Merge pull request JuliaLang#3100 from toivoh/staged_bsxfun

Staged bsxfun and other broadcast operations

base/broadcast.jl

@@ -0,0 +1,245 @@
+module Broadcast
+
+using ..Meta.quot
+import Base.(.+), Base.(.-), Base.(.*), Base.(./) 
+export broadcast, broadcast!, broadcast_function, broadcast!_function
+export broadcast_getindex, broadcast_setindex!
+
+
+## Broadcasting utilities ##
+
+# Calculate the broadcast shape of the arguments, or error if incompatible
+broadcast_shape() = ()
+function broadcast_shape(As::AbstractArray...)
+ nd = ndims(As[1])
+ for i = 2:length(As)
+ nd = max(nd, ndims(As[i]))
+ end
+ bshape = ones(Int, nd)
+ for A in As
+ for d = 1:ndims(A)
+ n = size(A, d)
+ if n != 1
+ if bshape[d] == 1
+ bshape[d] = n
+ elseif bshape[d] != n
+ error("arrays cannot be broadcast to a common size")
+ end
+ end
+ end
+ end
+ return tuple(bshape...)
+end
+
+# Check that all arguments are broadcast compatible with shape
+function check_broadcast_shape(shape::Dims, As::AbstractArray...)
+ for A in As
+ if ndims(A) > length(shape)
+ error("cannot broadcast array to have fewer dimensions")
+ end
+ for k in 1:ndims(A)
+ n, nA = shape[k], size(A, k)
+ if n != nA != 1
+ error("array cannot be broadcast to match destination")
+ end
+ end
+ end
+end
+
+# Calculate strides as will be used by the generated inner loops
+function calc_loop_strides(shape::Dims, As::AbstractArray...)
+ # squeeze out singleton dimensions in shape
+ dims = Array(Int, 0)
+ loopshape = Array(Int, 0)
+ nd = length(shape)
+ sizehint(dims, nd)
+ sizehint(loopshape, nd)
+ for i = 1:nd
+ s = shape[i]
+ if s != 1
+ push!(dims, i)
+ push!(loopshape, s)
+ end
+ end
+ nd = length(loopshape)
+
+ strides = [(size(A, d) > 1 ? stride(A, d) : 0) for A in As, d in dims]
+ # convert from regular strides to loop strides
+ for k=(nd-1):-1:1, a=1:length(As)
+ strides[a, k+1] -= strides[a, k]*loopshape[k]
+ end
+
+ tuple(loopshape...), strides
+end
+
+function broadcast_args(shape::Dims, As::(Array...))
+ loopshape, strides = calc_loop_strides(shape, As...)
+ (loopshape, As, ones(Int, length(As)), strides)
+end
+function broadcast_args(shape::Dims, As::(StridedArray...))
+ loopshape, strides = calc_loop_strides(shape, As...)
+ nA = length(As)
+ offs = Array(Int, nA)
+ baseAs = Array(Array, nA)
+ for (k, A) in enumerate(As)
+ offs[k],baseAs[k] = isa(A,SubArray) ? (A.first_index,A.parent) : (1,A)
+ end
+ (loopshape, tuple(baseAs...), offs, strides)
+end
+
+
+## Generation of inner loop instances ##
+
+function code_inner_loop(fname::Symbol, extra_args::Vector, initial,
+ innermost::Function, narrays::Int, nd::Int)
+ Asyms = [gensym("A$a") for a=1:narrays]
+ indsyms = [gensym("k$a") for a=1:narrays]
+ axissyms = [gensym("i$d") for d=1:nd]
+ sizesyms = [gensym("n$d") for d=1:nd]
+ stridesyms = [gensym("s$(a)_$d") for a=1:narrays, d=1:nd]
+
+ loop = innermost([:($arr[$ind]) for (arr, ind) in zip(Asyms, indsyms)]...)
+ for (d, (axis, n)) in enumerate(zip(axissyms, sizesyms))
+ loop = :(
+ for $axis=1:$n
+ $loop
+ $([:($ind += $(stridesyms[a, d]))
+ for (a, ind) in enumerate(indsyms)]...)
+ end
+ )
+ end
+
+ @gensym shape arrays offsets strides
+ quote
+ function $fname($shape::NTuple{$nd, Int},
+ $arrays::NTuple{$narrays, StridedArray},
+ $offsets::Vector{Int},
+ $strides::Matrix{Int}, $(extra_args...))
+ @assert size($strides) == ($narrays, $nd)
+ ($(sizesyms...),) = $shape
+ $([:(if $n==0; return; end) for n in sizesyms]...)
+ ($(Asyms...), ) = $arrays
+ ($(stridesyms...),) = $strides
+ ($(indsyms...), ) = $offsets
+ $initial
+ $loop
+ end
+ end
+end
+
+
+## Generation of inner loop staged functions ##
+
+function code_inner(fname::Symbol, extra_args::Vector, initial,
+ innermost::Function)
+ quote
+ function $fname(shape::(Int...), arrays::(StridedArray...), 
+ offsets::Vector{Int}, strides::Matrix{Int},
+ $(extra_args...))
+ f = eval(code_inner_loop($(quot(fname)), $(quot(extra_args)),
+ $(quot(initial)), $(quot(innermost)),
+ length(arrays), length(shape)))
+ f(shape, arrays, offsets, strides, $(extra_args...))
+ end
+ end
+end
+
+code_foreach_inner(fname::Symbol, extra_args::Vector, innermost::Function) =
+ code_inner(fname, extra_args, quote end, innermost)
+
+function code_map!_inner(fname::Symbol, dest, extra_args::Vector, 
+ innermost::Function)
+ @gensym k
+ code_inner(fname, {dest, extra_args...}, :($k=1),
+ (els...)->quote
+ $dest[$k] = $(innermost(:($dest[$k]), els...))
+ $k += 1
+ end)
+end
+
+
+## (Generation of) complete broadcast functions ##
+
+function code_broadcast(fname::Symbol, op)
+ inner! = gensym("$(fname)_inner!")
+ innerdef = code_map!_inner(inner!, :(result::Array), [],
+ (dest, els...) -> :( $op($(els...)) ))
+ quote
+ $innerdef
+ $fname() = $op()
+ function $fname(As::StridedArray...)
+ shape = broadcast_shape(As...)
+ result = Array(promote_type([eltype(A) for A in As]...), shape)
+ $inner!(broadcast_args(shape, As)..., result)
+ result
+ end 
+ end
+end
+
+function code_broadcast!(fname::Symbol, op)
+ inner! = gensym("$(fname)!_inner!")
+ innerdef = code_foreach_inner(inner!, [],
+ (dest, els...) -> :( $dest=$op($(els...)) ))
+ quote
+ $innerdef
+ function $fname(dest::StridedArray, As::StridedArray...)
+ shape = size(dest)
+ check_broadcast_shape(shape, As...)
+ $inner!(broadcast_args(shape, tuple(dest, As...))...)
+ dest
+ end 
+ end
+end
+
+eval(code_map!_inner(:broadcast_getindex_inner!,
+ :(result::Array), [:(A::AbstractArray)],
+ (dest, inds...) -> :( A[$(inds...)] )))
+function broadcast_getindex(A::AbstractArray,
+ ind1::StridedArray{Int},
+ inds::StridedArray{Int}...)
+ inds = tuple(ind1, inds...)
+ shape = broadcast_shape(inds...)
+ result = Array(eltype(A), shape)
+ broadcast_getindex_inner!(broadcast_args(shape, inds)..., result, A)
+ result
+end
+
+eval(code_foreach_inner(:broadcast_setindex!_inner!, [:(A::AbstractArray)],
+ (x, inds...)->:( A[$(inds...)] = $x )))
+function broadcast_setindex!(A::AbstractArray, X::StridedArray,
+ ind1::StridedArray{Int}, 
+ inds::StridedArray{Int}...)
+ Xinds = tuple(X, ind1, inds...)
+ shape = broadcast_shape(Xinds...)
+ broadcast_setindex!_inner!(broadcast_args(shape, Xinds)..., A)
+ Xinds[1]
+end
+
+
+## actual functions for broadcast and broadcast! ##
+
+for (fname, op) in {(:.+, +), (:.-, -), (:.*, *), (:./, /)}
+ eval(code_broadcast(fname, quot(op)))
+end
+
+broadcastfuns = (Function=>Function)[]
+function broadcast_function(op::Function)
+ (haskey(broadcastfuns, op) ? broadcastfuns[op] :
+ (broadcastfuns[op] = eval(code_broadcast(gensym("broadcast_$(op)"), 
+ quot(op)))))
+end
+broadcast(op::Function) = op()
+broadcast(op::Function, As::StridedArray...) = broadcast_function(op)(As...)
+
+broadcast!funs = (Function=>Function)[]
+function broadcast!_function(op::Function)
+ (haskey(broadcast!funs, op) ? broadcast!funs[op] :
+ (broadcast!funs[op] = eval(code_broadcast!(gensym("broadcast!_$(op)"), 
+ quot(op)))))
+end
+function broadcast!(op::Function, dest::StridedArray, As::StridedArray...)
+ broadcast!_function(op)(dest, As...)
+end
+
+
+end # module

base/exports.jl

@@ -452,7 +452,12 @@ export
 # arrays
  mapslices,
  reducedim,
- bsxfun,
+ broadcast,
+ broadcast!,
+ broadcast_function,
+ broadcast!_function,
+ broadcast_getindex,
+ broadcast_setindex!,
  cartesianmap,
  cat,
  cell,

base/sysimg.jl

@@ -164,6 +164,8 @@ push!(I18n.CALLBACKS, Help.clear_cache)
 include("sparse.jl")
 include("linalg.jl")
 importall .LinAlg
+include("broadcast.jl")
+importall .Broadcast
 
 # signal processing
 include("fftw.jl")

doc/manual/arrays.rst

@@ -329,13 +329,13 @@ vector to the size of the matrix::
  0.848333 1.66714 1.3262 
  1.26743 1.77988 1.13859
 
-This is wasteful when dimensions get large, so Julia offers the
-MATLAB-inspired ``bsxfun``, which expands singleton dimensions in
+This is wasteful when dimensions get large, so Julia offers
+``broadcast``, which expands singleton dimensions in
 array arguments to match the corresponding dimension in the other
-array without using extra memory, and applies the given binary
-function::
+array without using extra memory, and applies the given
+function elementwise::
 
- julia> bsxfun(+, a, A)
+ julia> broadcast(+, a, A)
  2x3 Float64 Array:
  0.848333 1.66714 1.3262 
  1.26743 1.77988 1.13859
@@ -344,11 +344,13 @@ function::
  1x2 Float64 Array:
  0.629799 0.754948
 
- julia> bsxfun(+, a, b)
+ julia> broadcast(+, a, b)
  2x2 Float64 Array:
  1.31849 1.44364
  1.56107 1.68622
 
+Elementwise operators such as ``.+`` and ``.*`` perform broadcasting if necessary. There is also a ``broadcast!`` function to specify an explicit destination, and ``broadcast_getindex`` and ``broadcast_setindex!`` that broadcast the indices before indexing.
+
 Implementation
 --------------
 

doc/stdlib/base.rst

@@ -2233,28 +2233,48 @@ Mathematical operators and functions
 
 All mathematical operations and functions are supported for arrays
 
-.. function:: bsxfun(fn, A, B[, C...])
+.. function:: broadcast(f, As...)
 
- Apply binary function ``fn`` to two or more arrays, with singleton dimensions expanded.
+ Broadcasts the arrays ``As`` to a common size by expanding singleton dimensions, and returns an array of the results ``f(as...)`` for each position.
+
+.. function:: broadcast!(f, dest, As...)
+
+ Like ``broadcast``, but store the result in the ``dest`` array.
+
+.. function:: broadcast_function(f)
+
+ Returns a function ``broadcast_f`` such that ``broadcast_function(f)(As...) === broadcast(f, As...)``. Most useful in the form ``const broadcast_f = broadcast_function(f)``.
+
+.. function:: broadcast!_function(f)
+
+ Like ``broadcast_function``, but for ``broadcast!``.
 
 Indexing, Assignment, and Concatenation
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. function:: getindex(A, ind)
+.. function:: getindex(A, inds...)
 
- Returns a subset of array ``A`` as specified by ``ind``, which may be an ``Int``, a ``Range``, or a ``Vector``.
+ Returns a subset of array ``A`` as specified by ``inds``, where each ``ind`` may be an ``Int``, a ``Range``, or a ``Vector``.
 
-.. function:: sub(A, ind)
+.. function:: sub(A, inds...)
 
- Returns a SubArray, which stores the input ``A`` and ``ind`` rather than computing the result immediately. Calling ``getindex`` on a SubArray computes the indices on the fly.
+ Returns a SubArray, which stores the input ``A`` and ``inds`` rather than computing the result immediately. Calling ``getindex`` on a SubArray computes the indices on the fly.
 
 .. function:: slicedim(A, d, i)
 
  Return all the data of ``A`` where the index for dimension ``d`` equals ``i``. Equivalent to ``A[:,:,...,i,:,:,...]`` where ``i`` is in position ``d``.
 
-.. function:: setindex!(A, X, ind)
+.. function:: setindex!(A, X, inds...)
+
+ Store values from array ``X`` within some subset of ``A`` as specified by ``inds``.
+
+.. function:: broadcast_getindex(A, inds...)
+
+ Broadcasts the ``inds`` arrays to a common size like ``broadcast``, and returns an array of the results ``A[ks...]``, where ``ks`` goes over the positions in the broadcast.
+
+.. function:: broadcast_setindex!(A, X, inds...)
 
- Store values from array ``X`` within some subset of ``A`` as specified by ``ind``.
+ Broadcasts the ``X`` and ``inds`` arrays to a common size and stores the value from each position in ``X`` at the indices given by the same positions in ``inds``.
 
 .. function:: cat(dim, A...)
 

test/Makefile

@@ -5,7 +5,7 @@ TESTS = all core keywordargs numbers strings unicode corelib hashing \
  remote iostring arrayops linalg blas fft dsp sparse bitarray \
  random math functional bigint sorting statistics spawn parallel \
  arpack bigfloat file git pkg pkg2 suitesparse complex version \
- pollfd mpfr
+ pollfd mpfr broadcast
 
 $(TESTS) ::
  $(QUIET_JULIA) $(call spawn,$(JULIA_EXECUTABLE)) ./runtests.jl $@