Hi,
I am upgrading a package of mine to v0.7. So far the trickiest bit is updating the broadcasting code to the new broadcasting facilities of v0.7. I have read the relevant manual pages, plus other bits here and there, but still have some questions.
The idea is that I have a type, called AugmentedState
that wraps two AbstractArrays
, x
and q
, i.e.
struct AugmentedState{X, Q}
x::X
q::Q
end
This type is part of the internals of the package and there is no risk that objects of this type collide with other objects. The usage case is that objects of this type are used for in-place operations, exclusively involving other AugmentedState
objects and optionally Number
s, as in a .= b .+ 2.0*c
.
Objects of this type do not support indexing: the only aim is to have arithmetic operations such as this one forwarded equally to the two wrapped arrays. This is what I have so far
# extract parts
@inline _state(x::AugmentedState) = x.x
@inline _quad(x::AugmentedState) = x.q
@inline _state(x::Base.RefValue) = _state(getindex(x))
@inline _quad(x::Base.RefValue) = _quad(getindex(x))
@inline _state(x) = x
@inline _quad(x) = x
@inline function Base.Broadcast.materialize!(dest::AugmentedState, bc::Base.Broadcast.Broadcasted)
bcf = Base.Broadcast.flatten(bc)
return __broadcast(bcf.f, dest, bcf.args...)
end
@generated function __broadcast(f, dest, args...)
quote
$(Expr(:meta, :inline))
Base.Broadcast.broadcast!(f,
_state(dest),
map(_state, args)...)
Base.Broadcast.broadcast!(f,
_quad(dest),
map(_quad, args)...)
return dest
end
end
This seems to work, but it incurs in additional allocations, plus a warning. Demo:
a = AugmentedState([1, 2], [3])
b = AugmentedState([1, 2], [3])
c = AugmentedState([0, 0], [0])
foo(a, b, c) = (c .= a .+ b .+ 2.0.*a; c)
@time foo(a, b, c)
@time foo(a, b, c)
@time foo(a, b, c)
@time foo(a, b, c)
gives
ā Warning: broadcast will default to iterating over its arguments in the future. Wrap arguments of
ā type `x::Flows.AugmentedState{Array{Int64,1},Array{Int64,1}}` with `Ref(x)` to ensure they broadcast as "scalar" elements.
ā caller = ip:0x0
ā @ Core :-1
1.569558 seconds (3.24 M allocations: 156.342 MiB, 9.89% gc time)
0.001696 seconds (852 allocations: 61.859 KiB)
0.001987 seconds (852 allocations: 61.859 KiB)
0.001628 seconds (852 allocations: 61.859 KiB)
In v0.6, I achieved this with this code
@generated function Base.Broadcast.broadcast!(f, dest::AugmentedState, args::Vararg{Any})
quote
$(Expr(:meta, :inline))
broadcast!(f, _state(dest), map(_state, args)...)
broadcast!(f, _quad(dest), map(_quad, args)...)
return dest
end
end
and had no extra allocations.
Any help is appreciated!