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

Commit

Permalink
Rewrite broadcast() using Base implementation and lifting semantics
Browse files Browse the repository at this point in the history 
Remove the custom implementation of broadcast(), and just call
the base method on the lifted method. For now, standard lifting
semantics are always used, even for logical operators and isnull/get.

Performance will probably be worse than before, but at least the
semantics will be correct. We can always re-implement the custom
and faster versions later when broadcast() has stabilized in Base
and Nullable support has settled.
  • Loading branch information
@nalimilan
nalimilan committed Nov 20, 2016
1 parent b7b50c3 commit 5de334f
Show file tree
Hide file tree
Showing 4 changed files with 103 additions and 203 deletions.
2 changes: 1 addition & 1 deletion REQUIRE
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
julia 0.4
julia 0.5
Compat 0.8.4
Reexport
262 changes: 86 additions & 176 deletions src/broadcast.jl
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
using Base: promote_eltype
using Base.Cartesian
using Base: _default_eltype, null_safe_op
using Compat

if VERSION >= v"0.6.0-dev.693"
using Base.Broadcast: check_broadcast_indices, broadcast_indices
else
Expand All @@ -8,198 +9,107 @@ else
const broadcast_indices = broadcast_shape
end

if VERSION >= v"0.5.0-dev+5189"
_to_shape(dims::Base.DimsOrInds) = map(_to_shape, dims)
_to_shape(r::Base.OneTo) = Int(last(r))
if !isdefined(Base.Broadcast, :ftype) # Julia < 0.6
ftype(f, A) = typeof(f)
ftype(f, A...) = typeof(a -> f(a...))
ftype(T::DataType, A) = Type{T}
ftype(T::DataType, A...) = Type{T}
else
_to_shape(x) = x
using Base.Broadcast: ftype
end

if VERSION < v"0.5.0-dev+5434"
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
if !isdefined(Base.Broadcast, :ziptype) # Julia < 0.6
if isdefined(Base, :Iterators)
using Base.Iterators: Zip2
else
using Base: Zip2
end
ziptype(A) = Tuple{eltype(A)}
ziptype(A, B) = Zip2{Tuple{eltype(A)}, Tuple{eltype(B)}}
@inline ziptype(A, B, C, D...) = Zip{Tuple{eltype(A)}, ziptype(B, C, D...)}
else
using Base.Broadcast: ziptype
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 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
@inline function broadcast_lift(f, x)
if null_safe_op(f, typeof(x))
return @compat Nullable(f(x.value), !isnull(x))
else
U = Core.Inference.return_type(f, Tuple{eltype(typeof(x))})
if isnull(x)
return Nullable{U}()
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
end
_F_
return Nullable(f(unsafe_get(x)))
end
end
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, map_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)
@inline function broadcast_lift(f, x1, x2)
if null_safe_op(f, typeof(x1), typeof(x2))
return @compat Nullable(
f(x1.value, x2.value), !(isnull(x1) | isnull(x2))
)
else
U = Core.Inference.return_type(
f, Tuple{eltype(typeof(x1)), eltype(typeof(x2))}
)
if isnull(x1) | isnull(x2)
return Nullable{U}()
else
return Nullable(f(unsafe_get(x1), unsafe_get(x2)))
end
# if 0 argument arrays, treat nullcheck as though it returns false
nargs >= 1 ? nullcheck : :(false)
end
end

@generated function Base.Broadcast._broadcast!{K,ID,XT,nargs}(f,
Z::NullableArray, keeps::K, Idefaults::ID, Xs::XT, ::Type{Val{nargs}}; lift=false)
nullcheck = _nullcheck(nargs)
quote
T = eltype(Z)
$(Expr(:meta, :noinline))
# destructure keeps and Xs tuples (common to both lifted and non-lifted broadcast)
@nexprs $nargs i->(keep_i = keeps[i])
@nexprs $nargs i->(Idefault_i = Idefaults[i])
if !lift
# destructure the keeps and As tuples
@nexprs $nargs i->(X_i = Xs[i])
@simd for I in CartesianRange(indices(Z))
# reverse-broadcast the indices
@nexprs $nargs i->(I_i = newindex(I, keep_i, Idefault_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)
@simd for I in CartesianRange(indices(Z))
# reverse-broadcast the indices
@nexprs $nargs i->(I_i = newindex(I, keep_i, Idefault_i))
if $nullcheck
# if any args are null, store null
@inbounds Z.isnull[I] = true
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
eltypes(x) = Tuple{eltype(x)}
eltypes(x, xs...) = Tuple{eltype(x), eltypes(xs...).parameters...}

"""
broadcast_lift(f, xs...)
Lift function `f`, passing it arguments `xs...`, using standard lifting semantics:
for a function call `f(xs...)`, return null if any `x` in `xs` is null; otherwise,
return `f` applied to values of `xs`.
"""
@inline function broadcast_lift(f, xs...)
if null_safe_op(f, map(typeof, xs)...)
# TODO: find a more efficient approach than mapreduce
# (i.e. one which gets lowedered to just isnull(x1) | isnull(x2) | ...)
return @compat Nullable(f(unsafe_get.(xs)...), !mapreduce(isnull, |, xs)
)
else
U = Core.Inference.return_type(f, eltypes(xs...))
# TODO: find a more efficient approach than mapreduce
# (i.e. one which gets lowedered to just isnull(x1) | isnull(x2) | ...)
if mapreduce(isnull, |, xs)
return Nullable{U}()
else
return Nullable(f(unsafe_get.(xs)...))
end
end
end

@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`.
""" ->
# Required to solve dispatch ambiguity between
# broadcast!(f, X::AbstractArray, x::Number...)
# broadcast!(f, Z::NullableArrays.NullableArray, Xs::NullableArrays.NullableArray...)
@inline Base.broadcast!(f, Z::NullableArray; lift=false) =
broadcast!(f, Z, Z; lift=lift)

@inline function Base.broadcast!(f, Z::NullableArray, Xs::NullableArray...;
lift=false)
nargs = length(Xs)
shape = indices(Z)
check_broadcast_indices(shape, Xs...)
keeps, Idefaults = map_newindexer(shape, Xs)
Base.Broadcast._broadcast!(f, Z, keeps, Idefaults, Xs, Val{nargs}; lift=lift)
return Z
"""
broadcast(f, As::NullableArray...)
Call `broadcast` with standard nullable lifting semantics and return a `NullableArray`.
Lifting means calling function `f` on the the values wrapped inside `Nullable` entries
of the input arrays, and returning null if any entry is missing.
"""
function Base.broadcast{N}(f, As::Vararg{NullableArray, N})
f2(x...) = broadcast_lift(f, x...)
T = _default_eltype(Base.Generator{ziptype(As...), ftype(f2, As...)})
if isleaftype(T) && !(T <: Nullable)
dest = similar(Array{eltype(T)}, broadcast_indices(As...))
else
dest = similar(NullableArray{eltype(T)}, broadcast_indices(As...))
end
invoke(broadcast!, Tuple{Function, AbstractArray, Vararg{AbstractArray, N}}, f2, dest, As...)
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
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, Xs::NullableArray...;lift::Bool=false)
return broadcast!(f, NullableArray(eltype(promote_eltype(Xs...)),
_to_shape(broadcast_indices(Xs...))),
Xs...; lift=lift)
function Base.broadcast!{N}(f, dest::AbstractArray, As::Vararg{NullableArray, N})
f2(x...) = broadcast_lift(f, x...)
invoke(broadcast!, Tuple{Function, AbstractArray, Vararg{AbstractArray, N}}, f2, dest, As...)
end

# broadcasted ops
Expand Down
5 changes: 2 additions & 3 deletions src/operators.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,7 @@ if VERSION < v"0.5.0-dev+5096"
end

"""
null_safe_op(f::Any, ::Type)::Bool
null_safe_op(f::Any, ::Type, ::Type)::Bool
null_safe_op(f::Any, ::Type...)::Bool
Returns whether an operation `f` can safely be applied to any value of the passed type(s).
Returns `false` by default.
Expand All @@ -38,7 +37,7 @@ always computing the result even for null values, a branch is avoided, which hel
vectorization.
"""
null_safe_op(f::Any, ::Type) = false
null_safe_op(f::Any, ::Type, ::Type) = false
null_safe_op(f::Any, ::Type, ::Type...) = false

typealias SafeSignedInts Union{Int128,Int16,Int32,Int64,Int8}
typealias SafeUnsignedInts Union{Bool,UInt128,UInt16,UInt32,UInt64,UInt8}
Expand Down
37 changes: 14 additions & 23 deletions test/broadcast.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,15 +31,10 @@ module TestBroadcast
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
f(x::Float64) = 5 * x
f(x::Float64, y::Float64) = x * y
f(x::Float64, y::Float64, z::Float64) = x * y * z

i = 1
for (dests, arrays, nullablearrays, mask) in
( ((C2, Z2), (A1, A2), (U1, U2), ()),
((C3, Z3), (A2, A3), (U2, U3), ()),
Expand All @@ -50,37 +45,25 @@ module TestBroadcast
((C3, Z3), (A1, A2, A3), (V1, V2, V3), (Q3,)),
)

# Base.broadcast!(f, B::NullableArray, As::NullableArray...; lift::Bool=false)
# Base.broadcast!(f, B::NullableArray, As::NullableArray...)
broadcast!(f, dests[1], arrays...)
broadcast!(f, dests[2], nullablearrays...)
@test isequal(dests[2], NullableArray(dests[1], mask...))

broadcast!(g, dests[1], arrays...)
broadcast!(g, dests[2], nullablearrays...; lift=true)
@test isequal(dests[2], NullableArray(dests[1], mask...))

# Base.broadcast(f, As::NullableArray...;lift::Bool=false)
# Base.broadcast(f, As::NullableArray...)
D = broadcast(f, arrays...)
X = broadcast(f, nullablearrays...)
@test isequal(X, NullableArray(D, mask...))

D = broadcast(g, arrays...)
X = broadcast(g, nullablearrays...; lift=true)
@test isequal(X, NullableArray(D, mask...))
end

# Base.broadcast!(f, X::NullableArray; lift::Bool=false)
# Base.broadcast!(f, X::NullableArray)
for (array, nullablearray, mask) in
( (A1, U1, ()), (A2, U2, ()), (A3, U3, ()),
(A1, V1, (M1,)), (A2, V2, (M2,)), (A3, V3, (M3,)),
)
broadcast!(f, array)
broadcast!(f, nullablearray)
@test isequal(nullablearray, NullableArray(array, mask...))

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

# test broadcasted arithmetic operators
Expand Down Expand Up @@ -111,5 +94,13 @@ module TestBroadcast
@test isequal(op(X1, Y), NullableArray(op(A, B), M))
end

A = rand(Bool, 100)
B = rand(Bool, 100)
M1 = rand(Bool, 100)
M2 = rand(Bool, 100)
X = NullableArray(A, M1)
Y = NullableArray(B, M2)
@test isequal(broadcast(&, X, Y), NullableArray(A & B, M1 | M2))
@test isequal(broadcast(|, X, Y), NullableArray(A | B, M1 | M2))

end # module

0 comments on commit 5de334f

Please sign in to comment.