Broadcast objects are not `instantiate`'d

[MasonProtter]

It seems like LazyBroadcast.jl is not using instantiate on the broadcast objects it creates. This can have some rather negative performance implications. Here's a quick comparison:

julia> using LazyBroadcast: @lazy_broadcasted

julia> using Base.Broadcast: Broadcasted, instantiate

julia> function foo(x)
           y = Broadcasted(+, (x, x))
           z = instantiate(Broadcasted(*, (2, x)))
           sum(z)
       end;

julia> function bar(x)
           y = @lazy_broadcasted x .+ x
           z = @lazy_broadcasted 2 .* y
           sum(z)
       end;

julia> let x = rand(10_000)
           print("Instantiated broadcast: "); @btime foo($x)
           print("LazyBroadcast:          "); @btime bar($x)
       end;
Instantiated broadcast:   683.796 ns (0 allocations: 0 bytes)
LazyBroadcast:            5.948 μs (0 allocations: 0 bytes)

[charleskawczynski]

Yeah, I had thought about this. Everywhere we use this we typically end up instantiating anyway, so maybe it's best that we just handle instantiating inside @lazy_broadcasted.

[MasonProtter]

Could also be an option to the @lazy_broadcasted macro.

[charleskawczynski]

True. I'm happy to change it, I don't think I'd consider it breaking to change the default behavior.

Also, I just had a look at DontMaterialize. The implementation is impressively small and simple.

I think the reason the LazyBroadcast implementation is complicated is because it basically was extracted from MultiBroadcastFusion.jl, which I think required certain things to be done at macro expansion time. One constraint I tried to satisfy was to avoid requiring our users to apply too many syntax changes when fusing expressions in ClimaAtmos, which has many @. annotations.

For simply extracting the broadcasted object, the DontMaterialize implementation seems sufficient.

[charleskawczynski]

It does look like instantiating is the main/only difference:

using LazyBroadcast: @lazy
using Base.Broadcast: Broadcasted, instantiate
using BenchmarkTools

function foo(x)
   y = Broadcasted(+, (x, x))
   z = instantiate(Broadcasted(*, (2, y)))
   sum(z)
end;

function bar(x)
   y = @lazy x .+ x
   z = @lazy 2 .* y
   sum(z)
end;

function baz(x)
   y = @lazy x .+ x
   z = instantiate(@lazy(2 .* y))
   sum(z)
end;

let x = rand(10_000)
   print("Instantiated broadcast:   "); @btime foo($x)
   print("LazyBroadcast:            "); @btime bar($x)
   print("LazyBroadcast+instantiate:"); @btime baz($x)
end

gives

Instantiated broadcast:     2.750 μs (0 allocations: 0 bytes)
LazyBroadcast:              9.958 μs (0 allocations: 0 bytes)
LazyBroadcast+instantiate:  2.750 μs (0 allocations: 0 bytes)

[charleskawczynski]

I'm happy to simplify our implementation, so long we can still support our existing syntax, which I think is possible.

[charleskawczynski]

Hm, I guess one "feature" that LazyBroadcast.jl supports, is that it can return the broadcasted expression. If we change the implementation to:

function _lazy_broadcasted end
const lazy = _lazy_broadcasted

struct LazyBroadcasted{T}
    value::T
end
Base.Broadcast.broadcasted(::typeof(_lazy_broadcasted), x) = LazyBroadcasted(x)
Base.Broadcast.materialize(x::LazyBroadcasted) = Base.Broadcast.instantiate(x.value)

import LazyBroadcast as LB
macro lazy_broadcasted(expr)
    Base.Broadcast.materialize(Base.Broadcast.broadcasted(LB.lazy, esc(expr)))
end

# Make shorter name available
macro lazy(expr)
    Base.Broadcast.materialize(Base.Broadcast.broadcasted(LB.lazy, esc(expr)))
end

then the lazy function works as expected, but @lazy still materializes. I think this is due to the order of macroexpansion (inner first).

[charleskawczynski]

I need to look around and see if we actually use rely on this anywhere.

[MasonProtter]

I think you just need to quote the output of that macro. e.g. something like

macro lazy_broadcasted(expr)
    :(Base.Broadcast.materialize(Base.Broadcast.broadcasted(LB.lazy, $(esc(expr)))))
end

[MasonProtter]

Oops, sorry one more change is required, you'd need to fuse the use of lazy here, e.g.

macro lazy_broadcasted(expr)
    :(LB.lazy.($(esc(expr))))
end

[charleskawczynski]

Oh, yeah, that would make sense, let me try this out

[charleskawczynski]

Hm, the expressions are not quite the same because what we currently have lowers the code. I'm not too concerned with this, because I don't think we actually use this.

But, some of the edge cases are failing with the simpler implementation:

julia> state = (; x = 1)
(x = 1,)

julia> bco = LB.@lazy_broadcasted @. state.x
1

julia> @test bco ==
             Base.Broadcast.broadcasted(identity, Base.getproperty(state, :x))
Test Failed at REPL[18]:1
  Expression: bco == Base.Broadcast.broadcasted(identity, Base.getproperty(state, :x))
   Evaluated: 1 == Base.Broadcast.Broadcasted(identity, (1,))

ERROR: There was an error during testing

I guess Julia is somehow evaluating/optimizing this? It's a bit odd that this would give a different result

[MasonProtter]

That's just because @. state.x isn't actually a broadcast expression.

julia> @macroexpand @. state.x
:(state.x)

[charleskawczynski]

Yeah, we treat it here as one, but it really doesn't need to be

[MasonProtter]

I think the old behaviour of turning it into broadcasted(identity, 1) is a bug, since it's not actually a lazy version of what @. does.

[charleskawczynski]

Yeah, I think we were trying to have @lazy_broadcasted always return a broadcasted object (https://github.com/CliMA/LazyBroadcast.jl/issues/8), but it might be better to just have it return the object in that case

[charleskawczynski]

I think it was previously erroring

[MasonProtter]

If you do want it to always return a Broadcasted object, I guess you could do

macro lazy_broadcasted(expr)
    quote
        res = LB.lazy.($(esc(expr)))
        if res isa Broadcast.Broadcasted
            res
        else
            broadcasted(identity, res)
        end
    end
end

[charleskawczynski]

I think I'm fine with returning the identity, it's not really an issue.

For the @lazy_broadcasted returning an expression, it's not really returning what I want:

julia> bce = LB.lazy_broadcasted_expression(:(@. a + b))
:(#= REPL[25]:1 =# @__dot__ a + b)

I'd like it to return

julia> bce = LB.lazy_broadcasted_expression(:(@. a + b))
:(Base.broadcasted(+, a, b))

or a Broadcasted object. I guess it maybe can't do the Broadcasted object, since that relies on dispatch, but Base.broadcasted should be possible

[charleskawczynski]

I think I've only used this feature once before.

I'll leave it in, but replace the core implementation.

[charleskawczynski]

Alright, I'm passing all of the tests now except for the in-place case: bco = LB.@lazy_broadcasted @. c += a + b, which I can't seem to catch in Base.Broadcast.materialize!(_, x::LazyBroadcasted) = Base.Broadcast.instantiate(x.value). Thoughts?

[charleskawczynski]

Basically, the test I would like to work is:

a = rand(3, 3)
b = rand(3, 3)
c = rand(3, 3)

bco = LB.@lazy_broadcasted @. c += a + b
@test bco == Base.broadcasted(+, c, Base.broadcasted(+, a, b))

[MasonProtter]

Hm, I think I'd call that a bug in LB, since that's not matching the semantics of broadcasting.

@. c += a + b

should be mutating c, not producing Base.broadcasted(+, c, Base.broadcasted(+, a, b)).

[MasonProtter]

disappointingly, I don't think there's a way for lazy.( a .= b) to get rid of the materialize! since there's an unconditional materialize!, so that case may need a macro.

[charleskawczynski]

Hm, I think I'd call that a bug in LB, since that's not matching the semantics of broadcasting.

@. c += a + b

should be mutating c, not producing Base.broadcasted(+, c, Base.broadcasted(+, a, b)).

Well, I'd argue it's not technically a bug, that's the broadcasted object that would be passed to materialize!(c, bc::Base.Broadcasted), but I agree that it doesn't match broadcasting semantics.

I don't think we strictly need this, but I'm also not sure I have the bandwidth to change the syntax everywhere we would need to make those changes.

If your dependence on DontMaterialize.jl is time-sensitive, I suggest continuing with registration.

In the mean time, I'm happy to update LB.jl to include instantiate calls.

[MasonProtter]

Well, I'd argue it's not technically a bug, that's the broadcasted object that would be passed to materialize!(c, bc::Base.Broadcasted), but I agree that it doesn't match broadcasting semantics.

Yeah, it's only a bug if you want it to match broadcast semantics. But I would usually hope that materialize(@lazy $expr) should typically be identical to just running $expr otherwise things become quite tricky to reason about.

If your dependence on DontMaterialize.jl is time-sensitive, I suggest continuing with registration.

I think I'll register it but put a not in the README pointing people to this package as well. It might be beneficial to have the two different choices out there for people.

In the mean time, I'm happy to update LB.jl to include instantiate calls.

Yeah, I think that should solve any performance issues and put the two on the same footing other than some syntactic / semantic differences.