Skip to content
This repository has been archived by the owner on May 4, 2019. It is now read-only.

Commit

Permalink
Merge pull request #123 from JuliaStats/broadcast
Browse files Browse the repository at this point in the history 
Fix broadcast for 0.5
  • Loading branch information
@quinnj
quinnj authored Jun 23, 2016
2 parents 1616e96 + 4dec5b9 commit 61398f3
Show file tree
Hide file tree
Showing 2 changed files with 161 additions and 80 deletions.
235 changes: 157 additions & 78 deletions src/broadcast.jl
View file
Original file line number Diff line number Diff line change
@@ -1,105 +1,184 @@
using Base.Broadcast: check_broadcast_shape
# using Base.Cartesian

function gen_nullcheck(narrays::Int, nd::Int)
e_nullcheck = macroexpand(:( @nref $nd isnull_1 d->j_d_1 ))
for k = 2:narrays
isnull = Symbol("isnull_$k")
j_d_k = Symbol("j_d_$k")
e_isnull_k = macroexpand(:( @nref $nd $(isnull) d->$(j_d_k) ))
e_nullcheck = Expr(:||, e_nullcheck, e_isnull_k)
using Base.Cartesian

if VERSION < v"0.5.0-dev+4724"
function gen_nullcheck(narrays::Int, nd::Int)
e_nullcheck = macroexpand(:( @nref $nd isnull_1 d->j_d_1 ))
for k = 2:narrays
isnull = Symbol("isnull_$k")
j_d_k = Symbol("j_d_$k")
e_isnull_k = macroexpand(:( @nref $nd $(isnull) d->$(j_d_k) ))
e_nullcheck = Expr(:||, e_nullcheck, e_isnull_k)
end
return e_nullcheck
end
return e_nullcheck
end

function gen_broadcast_body(nd::Int, narrays::Int, f, lift::Bool)
F = Expr(:quote, f)
e_nullcheck = gen_nullcheck(narrays, nd)
if lift
return quote
# set up aliases to facilitate subsequent Base.Cartesian magic
B_isnull = B.isnull
@nexprs $narrays k->(values_k = A_k.values)
@nexprs $narrays k->(isnull_k = A_k.isnull)
# check size
@assert ndims(B) == $nd
@ncall $narrays Base.Broadcast.check_broadcast_shape size(B) k->A_k
# main loops
@nloops($nd, i, B,
d->(@nexprs $narrays k->(j_d_k = size(A_k, d) == 1 ? 1 : i_d)), # pre
begin # body
if $e_nullcheck
@inbounds (@nref $nd B_isnull i) = true
else
@nexprs $narrays k->(@inbounds v_k = @nref $nd values_k d->j_d_k)
@inbounds (@nref $nd B i) = (@ncall $narrays $F v)
function gen_broadcast_body(nd::Int, narrays::Int, f, lift::Bool)
F = Expr(:quote, f)
e_nullcheck = gen_nullcheck(narrays, nd)
if lift
return quote
# set up aliases to facilitate subsequent Base.Cartesian magic
B_isnull = B.isnull
@nexprs $narrays k->(values_k = A_k.values)
@nexprs $narrays k->(isnull_k = A_k.isnull)
# check size
@assert ndims(B) == $nd
@ncall $narrays Base.Broadcast.check_broadcast_shape size(B) k->A_k
# main loops
@nloops($nd, i, B,
d->(@nexprs $narrays k->(j_d_k = size(A_k, d) == 1 ? 1 : i_d)), # pre
begin # body
if $e_nullcheck
@inbounds (@nref $nd B_isnull i) = true
else
@nexprs $narrays k->(@inbounds v_k = @nref $nd values_k d->j_d_k)
@inbounds (@nref $nd B i) = (@ncall $narrays $F v)
end
end
end
)
)
end
else
return Base.Broadcast.gen_broadcast_body_cartesian(nd, narrays, f)
end
else
return Base.Broadcast.gen_broadcast_body_cartesian(nd, narrays, f)
end
end

function gen_broadcast_function(nd::Int, narrays::Int, f, lift::Bool)
As = [Symbol("A_"*string(i)) for i = 1:narrays]
body = gen_broadcast_body(nd, narrays, f, lift)
@eval let
local _F_
function _F_(B, $(As...))
$body
function gen_broadcast_function(nd::Int, narrays::Int, f, lift::Bool)
As = [Symbol("A_"*string(i)) for i = 1:narrays]
body = gen_broadcast_body(nd, narrays, f, lift)
@eval let
local _F_
function _F_(B, $(As...))
$body
end
_F_
end
_F_
end
end

function Base.broadcast!(f, X::NullableArray; lift::Bool=false)
broadcast!(f, X, X; lift=lift)
end
function Base.broadcast!(f, X::NullableArray; lift::Bool=false)
broadcast!(f, X, X; lift=lift)
end

@eval let cache = Dict{Any, Dict{Bool, Dict{Int, Dict{Int, Any}}}}()
@doc """
`broadcast!(f, B::NullableArray, As::NullableArray...; lift::Bool=false)`
This method implements the same behavior as that of `broadcast!` when called on
regular `Array` arguments. It also includes the `lift` keyword argument, which
when set to true will lift `f` over the entries of the `As`.
Lifting is disabled by default. Note that this method's signature specifies
the destination `B` array as well as the source `As` arrays as all
`NullableArray`s. Thus, calling `broadcast!` on a arguments consisting
of both `Array`s and `NullableArray`s will fall back to the implementation
of `broadcast!` in `base/broadcast.jl`.
""" ->
function Base.broadcast!(f, B::NullableArray, As::NullableArray...; lift::Bool=false)
nd = ndims(B)
narrays = length(As)

cache_f = Base.@get! cache f Dict{Bool, Dict{Int, Dict{Int, Any}}}()
cache_lift = Base.@get! cache_f lift Dict{Int, Dict{Int, Any}}()
cache_f_na = Base.@get! cache_lift narrays Dict{Int, Any}()
func = Base.@get! cache_f_na nd gen_broadcast_function(nd, narrays, f, lift)

func(B, As...)
return B
end
end # let cache
else
using Base.Broadcast: newindexer, newindex

function _nullcheck(nargs)
nullcheck = :(isnull_1[I_1])
for i in 2:nargs
sym_isnull = Symbol("isnull_$i")
sym_idx = Symbol("I_$i")
nullcheck = Expr(:||, :($sym_isnull[$sym_idx]), nullcheck)
end
# if 0 argument arrays, treat nullcheck as though it returns false
nargs >= 1 ? nullcheck : :(false)
end

@generated function Base.Broadcast._broadcast!{M,XT,nargs}(f,
Z::NullableArray, indexmaps::M, Xs::XT, ::Type{Val{nargs}}; lift=false)
nullcheck = _nullcheck(nargs)
quote
T = eltype(Z)
$(Expr(:meta, :noinline))
if !lift
# destructure the indexmaps and As tuples
@nexprs $nargs i->(X_i = Xs[i])
@nexprs $nargs i->(imap_i = indexmaps[i])
@simd for I in CartesianRange(indices(Z))
# reverse-broadcast the indices
@nexprs $nargs i->(I_i = newindex(I, imap_i))
# extract array values
@nexprs $nargs i->(@inbounds val_i = X_i[I_i])
# call the function and store the result
@inbounds Z[I] = @ncall $nargs f val
end
else
# destructure the indexmaps and Xs tuples
@nexprs $nargs i->(values_i = Xs[i].values)
@nexprs $nargs i->(isnull_i = Xs[i].isnull)
@nexprs $nargs i->(imap_i = indexmaps[i])
@simd for I in CartesianRange(indices(Z))
# reverse-broadcast the indices
@nexprs $nargs i->(I_i = newindex(I, imap_i))
if $nullcheck
# if any args are null, store null
@inbounds Z[I] = Nullable{T}()
else
# extract array values
@nexprs $nargs i->(@inbounds val_i = values_i[I_i])
# call the function and store the result
@inbounds Z[I] = @ncall $nargs f val
end
end
end
end
end

@eval let cache = Dict{Any, Dict{Bool, Dict{Int, Dict{Int, Any}}}}()
@doc """
`broadcast!(f, B::NullableArray, As::NullableArray...; lift::Bool=false)`
This method implements the same behavior as that of `broadcast!` when called on
regular `Array` arguments. It also includes the `lift` keyword argument, which
when set to true will lift `f` over the entries of the `As`. Lifting is
disabled by default. Note that this method's signature specifies the destination
`B` array as well as the source `As` arrays as all `NullableArray`s.
Thus, calling `broadcast!` on a arguments consisting of both `Array`s and
`NullableArray`s will fall back to the implementation of `broadcast!` in
`base/broadcast.jl`.
""" ->
function Base.broadcast!(f, B::NullableArray, As::NullableArray...; lift::Bool=false)
nd = ndims(B)
narrays = length(As)

cache_f = Base.@get! cache f Dict{Bool, Dict{Int, Dict{Int, Any}}}()
cache_lift = Base.@get! cache_f lift Dict{Int, Dict{Int, Any}}()
cache_f_na = Base.@get! cache_lift narrays Dict{Int, Any}()
func = Base.@get! cache_f_na nd gen_broadcast_function(nd, narrays, f, lift)
This method implements the same behavior as that of `broadcast!` when called
on regular `Array` arguments. It also includes the `lift` keyword argument,
which when set to true will lift `f` over the entries of the `As`.
func(B, As...)
return B
Lifting is disabled by default. Note that this method's signature specifies
the destination `B` array as well as the source `As` arrays as all
`NullableArray`s. Thus, calling `broadcast!` on a arguments consisting of
both `Array`s and `NullableArray`s will fall back to the implementation of
`broadcast!` in `base/broadcast.jl`.
""" ->
@inline function Base.broadcast!(f, Z::NullableArray, Xs::NullableArray...;
lift=false)
nargs = length(Xs)
sz = size(Z)
check_broadcast_shape(sz, Xs...)
mapindex = map(x->newindexer(sz, x), Xs)
Base.Broadcast._broadcast!(f, Z, mapindex, Xs, Val{nargs}; lift=lift)
Z
end
end # let cache
end

@doc """
`broadcast(f, As::NullableArray...;lift::Bool=false)`
This method implements the same behavior as that of `broadcast` when called on
regular `Array` arguments. It also includes the `lift` keyword argument, which
when set to true will lift `f` over the entries of the `As`. Lifting is
disabled by default. Note that this method's signature specifies the source
`As` arrays as all `NullableArray`s. Thus, calling `broadcast!` on a arguments
consisting of both `Array`s and `NullableArray`s will fall back to the
when set to true will lift `f` over the entries of the `As`.
Lifting is disabled by default. Note that this method's signature specifies the
source `As` arrays as all `NullableArray`s. Thus, calling `broadcast!` on
arguments consisting of both `Array`s and `NullableArray`s will fall back to the
implementation of `broadcast` in `base/broadcast.jl`.
""" ->
function Base.broadcast(f, As::NullableArray...;lift::Bool=false)
return broadcast!(f, NullableArray(eltype(Base.promote_eltype(As...)),
Base.Broadcast.broadcast_shape(As...)),
As...; lift=lift)
@inline function Base.broadcast(f, Xs::NullableArray...;lift::Bool=false)
return broadcast!(f, NullableArray(eltype(Base.promote_eltype(Xs...)),
Base.Broadcast.broadcast_shape(Xs...)),
Xs...; lift=lift)
end

# broadcasted ops
Expand Down
6 changes: 4 additions & 2 deletions test/broadcast.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,9 +30,11 @@ module TestBroadcast
Z2 = NullableArray(Float64, 10, dims...)
Z3 = NullableArray(Float64, 10, [dims; i]...)

f() = 5
f(x) = Nullable(5) * x
f(x, y) = x + y
f(x, y, z) = x + y + z
g() = 5
g(x::Float64) = 5 * x
g(x::Float64, y::Float64) = x * y
g(x::Float64, y::Float64, z::Float64) = x * y * z
Expand Down Expand Up @@ -72,11 +74,11 @@ module TestBroadcast
( (A1, U1, ()), (A2, U2, ()), (A3, U3, ()),
(A1, V1, (M1,)), (A2, V2, (M2,)), (A3, V3, (M3,)),
)
map!(g, array)
broadcast!(f, array)
broadcast!(f, nullablearray)
@test isequal(nullablearray, NullableArray(array, mask...))

map!(g, array)
broadcast!(g, array)
broadcast!(g, nullablearray; lift=true)
@test isequal(nullablearray, NullableArray(array, mask...))
end
Expand Down

0 comments on commit 61398f3

Please sign in to comment.