33module Broadcast
44
55using Base. Cartesian
6- using Base: promote_eltype_op, linearindices, tail, OneTo, to_shape,
6+ using Base: promote_eltype_op, _default_eltype, linearindices, tail, OneTo, to_shape,
77 _msk_end, unsafe_bitgetindex, bitcache_chunks, bitcache_size, dumpbitcache
88import Base: .+ , .- , .* , ./ , .\ , .// , .== , .< , .!= , .<= , .÷ , .% , .<< , .>> , .^
99import Base: broadcast
@@ -16,7 +16,7 @@ export broadcast_getindex, broadcast_setindex!
1616broadcast (f) = f ()
1717@inline broadcast (f, x:: Number... ) = f (x... )
1818@inline broadcast {N} (f, t:: NTuple{N} , ts:: Vararg{NTuple{N}} ) = map (f, t, ts... )
19- @inline broadcast (f, As:: AbstractArray... ) = broadcast_t (f, promote_eltype_op (f, As ... ) , As... )
19+ @inline broadcast (f, As:: AbstractArray... ) = broadcast_c (f, Array , As... )
2020
2121# special cases for "X .= ..." (broadcast!) assignments
2222broadcast! (:: typeof (identity), X:: AbstractArray , x:: Number ) = fill! (X, x)
@@ -127,14 +127,14 @@ Base.@propagate_inbounds _broadcast_getindex(::Any, A, I) = A[I]
127127# # Broadcasting core
128128# nargs encodes the number of As arguments (which matches the number
129129# of keeps). The first two type parameters are to ensure specialization.
130- @generated function _broadcast! {K,ID,AT,nargs} (f, B:: AbstractArray , keeps:: K , Idefaults:: ID , As:: AT , :: Type{Val{nargs}} )
130+ @generated function _broadcast! {K,ID,AT,nargs} (f, B:: AbstractArray , keeps:: K , Idefaults:: ID , As:: AT , :: Type{Val{nargs}} , iter )
131131 quote
132132 $ (Expr (:meta , :noinline ))
133133 # destructure the keeps and As tuples
134134 @nexprs $ nargs i-> (A_i = As[i])
135135 @nexprs $ nargs i-> (keep_i = keeps[i])
136136 @nexprs $ nargs i-> (Idefault_i = Idefaults[i])
137- @simd for I in CartesianRange ( indices (B))
137+ @simd for I in iter
138138 # reverse-broadcast the indices
139139 @nexprs $ nargs i-> (I_i = newindex (I, keep_i, Idefault_i))
140140 # extract array values
148148
149149# For BitArray outputs, we cache the result in a "small" Vector{Bool},
150150# and then copy in chunks into the output
151- @generated function _broadcast! {K,ID,AT,nargs} (f, B:: BitArray , keeps:: K , Idefaults:: ID , As:: AT , :: Type{Val{nargs}} )
151+ @generated function _broadcast! {K,ID,AT,nargs} (f, B:: BitArray , keeps:: K , Idefaults:: ID , As:: AT , :: Type{Val{nargs}} , iter )
152152 quote
153153 $ (Expr (:meta , :noinline ))
154154 # destructure the keeps and As tuples
159159 Bc = B. chunks
160160 ind = 1
161161 cind = 1
162- @simd for I in CartesianRange ( indices (B))
162+ @simd for I in iter
163163 # reverse-broadcast the indices
164164 @nexprs $ nargs i-> (I_i = newindex (I, keep_i, Idefault_i))
165165 # extract array values
@@ -193,12 +193,12 @@ as in `broadcast!(f, A, A, B)` to perform `A[:] = broadcast(f, A, B)`.
193193 shape = indices (B)
194194 check_broadcast_indices (shape, As... )
195195 keeps, Idefaults = map_newindexer (shape, As)
196- _broadcast! (f, B, keeps, Idefaults, As, Val{nargs})
197- B
196+ iter = CartesianRange (shape)
197+ _broadcast! (f, B, keeps, Idefaults, As, Val{nargs}, iter)
198+ return B
198199end
199200
200201# broadcast with computed element type
201-
202202@generated function _broadcast! {K,ID,AT,nargs} (f, B:: AbstractArray , keeps:: K , Idefaults:: ID , As:: AT , :: Type{Val{nargs}} , iter, st, count)
203203 quote
204204 $ (Expr (:meta , :noinline ))
233233 end
234234end
235235
236- function broadcast_t (f, :: Type{Any} , As... )
237- shape = broadcast_indices (As... )
238- iter = CartesianRange (shape)
239- if isempty (iter)
240- return similar (Array{Any}, shape)
241- end
236+ # broadcast methods that dispatch on the type found by inference
237+ function broadcast_t (f, :: Type{Any} , shape, iter, As... )
242238 nargs = length (As)
243239 keeps, Idefaults = map_newindexer (shape, As)
244240 st = start (iter)
@@ -248,17 +244,46 @@ function broadcast_t(f, ::Type{Any}, As...)
248244 B[I] = val
249245 return _broadcast! (f, B, keeps, Idefaults, As, Val{nargs}, iter, st, 1 )
250246end
247+ @inline function broadcast_t (f, T, shape, iter, As... )
248+ B = similar (Array{T}, shape)
249+ nargs = length (As)
250+ keeps, Idefaults = map_newindexer (shape, As)
251+ _broadcast! (f, B, keeps, Idefaults, As, Val{nargs}, iter)
252+ return B
253+ end
251254
252- @inline broadcast_t (f, T, As... ) = broadcast! (f, similar (Array{T}, broadcast_indices (As... )), As... )
253-
255+ # broadcast method that uses inference to find the type, but preserves abstract
256+ # container types when possible (used by binary elementwise operators)
257+ @inline broadcast_elwise_op (f, As... ) =
258+ broadcast! (f, similar (Array{promote_eltype_op (f, As... )}, broadcast_indices (As... )), As... )
259+
260+ ftype (f, A) = typeof (a -> f (a))
261+ ftype (f, A... ) = typeof (a -> f (a... ))
262+ ftype (T:: DataType , A) = Type{T}
263+ ftype (T:: DataType , A... ) = Type{T}
264+ ziptype (A) = Tuple{eltype (A)}
265+ ziptype (A, B) = Iterators. Zip2{Tuple{eltype (A)}, Tuple{eltype (B)}}
266+ @inline ziptype (A, B, C, D... ) = Iterators. Zip{Tuple{eltype (A)}, ziptype (B, C, D... )}
267+
268+ # broadcast methods that dispatch on the type of the final container
269+ @inline function broadcast_c (f, :: Type{Array} , As... )
270+ T = _default_eltype (Base. Generator{ziptype (As... ), ftype (f, As... )})
271+ shape = broadcast_indices (As... )
272+ iter = CartesianRange (shape)
273+ if isleaftype (T)
274+ return broadcast_t (f, T, shape, iter, As... )
275+ end
276+ if isempty (iter)
277+ return similar (Array{T}, shape)
278+ end
279+ return broadcast_t (f, Any, shape, iter, As... )
280+ end
254281function broadcast_c (f, :: Type{Tuple} , As... )
255282 shape = broadcast_indices (As... )
256- check_broadcast_indices (shape, As... )
257283 n = length (shape[1 ])
258284 return ntuple (k-> f ((_broadcast_getindex (A, k) for A in As). .. ), n)
259285end
260286@inline broadcast_c (f, :: Type{Any} , a... ) = f (a... )
261- @inline broadcast_c (f, :: Type{Array} , As... ) = broadcast_t (f, promote_eltype_op (f, As... ), As... )
262287
263288"""
264289 broadcast(f, As...)
@@ -441,10 +466,10 @@ end
441466# # elementwise operators ##
442467
443468for op in (:÷ , :% , :<< , :>> , :- , :/ , :\ , :// , :^ )
444- @eval $ (Symbol (:., op))(A:: AbstractArray , B:: AbstractArray ) = broadcast ($ op, A, B)
469+ @eval $ (Symbol (:., op))(A:: AbstractArray , B:: AbstractArray ) = broadcast_elwise_op ($ op, A, B)
445470end
446- .+ (As:: AbstractArray... ) = broadcast (+ , As... )
447- .* (As:: AbstractArray... ) = broadcast (* , As... )
471+ .+ (As:: AbstractArray... ) = broadcast_elwise_op (+ , As... )
472+ .* (As:: AbstractArray... ) = broadcast_elwise_op (* , As... )
448473
449474# ## element-wise comparison operators returning BitArray ##
450475
0 commit comments