Add support for `Broadcasted` objects #64

Sbozzolo · 2024-08-13T20:24:19Z

This PR adds support for specifying diagnostics with a compute function instead of compute!. compute(state, cache, time) can return either a Field or a Base.Broadcast.Broadcasted object (obtained with @lazy) and are now the recommended way to specify diagnostics. compute! is still supported to preserve existing code.

Closes #5

charleskawczynski · 2024-08-13T20:52:55Z

Could we instead use LazyBroadcasts @lazy to reduce code duplication, instead? That way we can also easily fuse expressions.

Sbozzolo · 2024-08-13T21:31:47Z

Yes, my goal is to eventually allow for kernel fusion (which is what I meant with "pathway for future improvement"). However, I don't know how to use LazyBroadcast for that, so I was going for something that would allow us to use it internally in the future, while starting to reduce code verbosity.

Is it just a matter of adding @lazy_broadcasted in front of $(esc(out)) .= $(esc(rhs2)) and explicitely calling materialize! when needed?

Maybe this is a good opportunity to add some user documentation to LazyBroadcast.jl.

charleskawczynski · 2024-08-14T00:43:53Z

Yes, my goal is to eventually allow for kernel fusion (which is what I meant with "pathway for future improvement"). However, I don't know how to use LazyBroadcast for that, so I was going for something that would allow us to use it internally in the future, while starting to reduce code verbosity.

Is it just a matter of adding @lazy_broadcasted in front of $(esc(out)) .= $(esc(rhs2)) and explicitely calling materialize! when needed?

Maybe this is a good opportunity to add some user documentation to LazyBroadcast.jl.

I'm happy to add more docs to LazyBroadcast, what do you think is missing?

Using LazyBroadcast, compute_ta (from the docs) could be re-written as:

using LazyBroadcast: @lazy
compute_ta(state, cache, time) = @lazy @. state.ta

and the result could later be used as Base.Broadcast.materialize!(out, bc) where bc is the result of compute_ta.

Sbozzolo · 2024-08-14T02:22:11Z

Thank you! I got something to work and I managed to cluster all the materialize! together in a for loop so that it should be easy to fuse them. Two questions:

Should the broadcasted expressions be re-computed every time, or is it enough to re-materialize them?
With eager evaluation, we need to copy the first time we call compute! because ClimaDiagnostics modifies the output (to do accumulation/reduction). Do we need to do this for lazy evaluation too?

As for documentation: I am looking at LazyBroadcasts.jl and all the docs I see is the readme which does not mention @lazy or compute_ta.

This is what I think could be improved in the readme:

the problem that the package solves is not clearly identified
it assumes that the reader has specialized knowledge of how Brocasted objects work and how to use them
the only example provided is a test case which comes with no explanation and components that are not pedagogically useful (all the Test related stuff)
it is not explained how one would use bce or bco
the difference between bce and bco is left to the reader to understand
the fact that the package has restrictions is not mentioned
it does not guide the reader to further documentation, so the way I find myself trying to figure things out is by browsing the testcases

Here are few things I tried to get the code to work:

julia> import LazyBroadcast: @lazy
julia> compute_ta(state, cache, time) = @lazy @. state.ta
ERROR: LoadError: Invalid expression given to materialize_args
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] materialize_args(expr::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/utils.jl:8
 [3] transform(e::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:30 [4] _lazy_broadcasted(expr::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:43 [5] var"@lazy"(__source__::LineNumberNode, __module__::Module, expr::Any)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:54in expression starting at REPL[17]:1

julia> compute_ta(state, cache, time) = @lazy @. state
ERROR: LoadError: type Symbol has no field args
Stacktrace:
 [1] getproperty(x::Symbol, f::Symbol)
   @ Base ./Base.jl:37
 [2] code_info(expr::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/code_lowered_single_expression.jl:10
 [3] code_lowered_single_expression(expr::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/code_lowered_single_expression.jl:12
 [4] transform(e::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:29 [5] _lazy_broadcasted(expr::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:43 [6] var"@lazy"(__source__::LineNumberNode, __module__::Module, expr::Any)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:54in expression starting at REPL[18]:1

I also found that these do not work:

julia> a, b = [1], [2]

julia> @lazy a .= b
Expr
  head: Symbol .=
  args: Array{Any}((2,))
    1: Symbol a
    2: Symbol b
dump(expr) = nothing
Expr
  head: Symbol .=
  args: Array{Any}((2,))
    1: Symbol a
    2: Symbol b
dump(expr) = nothing
ERROR: LoadError: Uncaught edge case
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] check_restrictions(expr::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/utils.jl:48
 [3] _lazy_broadcasted(expr::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:42 [4] var"@lazy"(__source__::LineNumberNode, __module__::Module, expr::Any)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:54in expression starting at REPL[12]:1

julia> @lazy @dot a = b
Expr
  head: Symbol macrocall
  args: Array{Any}((3,))
    1: Symbol @dot
    2: LineNumberNode
      line: Int64 1
      file: Symbol REPL[19]
    3: Expr
      head: Symbol =
      args: Array{Any}((2,))
        1: Symbol a
        2: Symbol b
dump(expr) = nothing
Expr
  head: Symbol macrocall
  args: Array{Any}((3,))
    1: Symbol @dot
    2: LineNumberNode
      line: Int64 1
      file: Symbol REPL[19]
    3: Expr
      head: Symbol =
      args: Array{Any}((2,))
        1: Symbol a
        2: Symbol b
dump(expr) = nothing
ERROR: LoadError: Uncaught edge case
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] check_restrictions(expr::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/utils.jl:48
 [3] _lazy_broadcasted(expr::Expr)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:42 [4] var"@lazy"(__source__::LineNumberNode, __module__::Module, expr::Any)
   @ LazyBroadcast ~/.julia/packages/LazyBroadcast/SVKU8/src/LazyBroadcast.jl:54in expression starting at REPL[19]:1

Eventually I found that it works with @. (it returns Base.Broadcast.Broadcasted(identity, ([2],)))

From here, I think I understand that the correct function to write is

compute_ta(out, state, cache, time) = @lazy @. out = state.ta

This seems to work in my integration test case.

charleskawczynski · 2024-08-14T15:49:40Z

Thank you! I got something to work and I managed to cluster all the materialize! together in a for loop so that it should be easy to fuse them. Two questions:

Should the broadcasted expressions be re-computed every time, or is it enough to re-materialize them?

Great question. That's safe to do so long the data passed into the broadcasted objects point to the same memory. That said, it's actually very cheap to do this because all it's doing is making an instance of a struct with a bunch of pointers, so I recommend we reconstruct the object every time to be safe.

With eager evaluation, we need to copy the first time we call compute! because ClimaDiagnostics modifies the output (to do accumulation/reduction). Do we need to do this for lazy evaluation too?

We may be able to use Base.Broadcast.materialize(bc) if we want an instance (for the first time), and then use Base.Broadcast.materialize!(out, bc) thereafter. I don't think we need to explicitly use copy, since I think Base.Broadcast.materialize will do that for us.

As for documentation: I am looking at LazyBroadcasts.jl and all the docs I see is the readme which does not mention @lazy or compute_ta.

Yeah, the compute_ta is implementation specific. I think we can provide some some simple example use-cases, though.

This is what I think could be improved in the readme:

the problem that the package solves is not clearly identified

it assumes that the reader has specialized knowledge of how Brocasted objects work and how to use them

the only example provided is a test case which comes with no explanation and components that are not pedagogically useful (all the Test related stuff)

it is not explained how one would use bce or bco

the difference between bce and bco is left to the reader to understand

the fact that the package has restrictions is not mentioned

it does not guide the reader to further documentation, so the way I find myself trying to figure things out is by browsing the testcases

Yeah, I agree on all of these points. LazyBroadcast.jl was born because it was originally a bunch of utilities in MultiBroadcastFusion.jl, and I realized that they were more generally useful, so I split those utilities off into a separate package. I hadn't fully thought out the use-cases, and if/what pieces could be made more flexible (like the broken examples you showed below), so I wasn't sure how to best document things. I'll open an issue with these points.

Here are few things I tried to get the code to work:...

julia> import LazyBroadcast: @lazy
julia> compute_ta(state, cache, time) = @lazy @. state.ta
ERROR: LoadError: Invalid expression given to materialize_args

julia> compute_ta(state, cache, time) = @lazy @. state
ERROR: LoadError: type Symbol has no field args
...

I also found that these do not work:

julia> a, b = [1], [2]

julia> @lazy a .= b
Expr
  head: Symbol .=
...
ERROR: LoadError: Uncaught edge case

Yeah, this is a bit technical I think, and it does seem that there are some sharp edges (I didn't expect that last example to fail). It might be easier to discuss some of this over zoom. One tricky aspect of this is that, for example, foo(x) could return a vector, in which case, we wouldn't want to assume that the user meant @lazy foo.(x). So, it allows for more flexibility (there are some such examples in the test suite, but not much discussion in the docs about this).

compute_ta(out, state, cache, time) = @lazy @. out = state.ta

Yeah, this will work, and that's how I'd expect it be used. There might be something we can do to remove needing out, but that should be fine for now.

src/clima_diagnostics.jl

Sbozzolo · 2024-08-14T22:21:18Z

Passing ClimaAtmos build: https://buildkite.com/clima/climaatmos-ci/builds/20044

I did not change the EDMF and radiation diagnostics because they are much more complex.

The allocation flame graphs are failing because I added more for loops.

charleskawczynski · 2024-08-15T15:00:23Z

Passing ClimaAtmos build: https://buildkite.com/clima/climaatmos-ci/builds/20044

I did not change the EDMF and radiation diagnostics because they are much more complex.

The allocation flame graphs are failing because I added more for loops.

Nice! Do we still need @assign? I thought that it wouldn't be necessary since @lazy changes things to be functional

Sbozzolo · 2024-08-15T15:19:09Z

Passing ClimaAtmos build: https://buildkite.com/clima/climaatmos-ci/builds/20044
I did not change the EDMF and radiation diagnostics because they are much more complex.
The allocation flame graphs are failing because I added more for loops.

Nice! Do we still need @assign? I thought that it wouldn't be necessary since @lazy changes things to be functional

I am not sure. Does @lazy support diagnostics like this

Fields.array2field(
            cache.radiation.rrtmgp_model.face_sw_flux_dn,
            axes(state.f),
        )

?

charleskawczynski · 2024-08-15T15:33:10Z

Fields.array2field(
            cache.radiation.rrtmgp_model.face_sw_flux_dn,
            axes(state.f),
        )

For something like that, we don't even need @lazy, since there's no broadcasting. That said, we can still make a broadcasted object so that it's similar to what @lazy returns:

x = Fields.array2field(
            cache.radiation.rrtmgp_model.face_sw_flux_dn,
            axes(state.f),
        )
bc = Base.broadcasted(identity, x)

charleskawczynski · 2024-08-15T15:34:29Z

In fact, any diagnostic that simply returns a field, could just wrap it in Base.broadcasted(identity, ::Fields.Field), and that way everything could be a broadcasted object that will work with Base.Broadcast.materialize!/Base.Broadcast.materialize

Sbozzolo · 2024-08-15T18:40:59Z

Waiting for
CliMA/LazyBroadcast.jl#8
CliMA/LazyBroadcast.jl#7

Sbozzolo · 2025-03-24T15:12:06Z

Passing Atmos build: https://buildkite.com/clima/climaatmos-ci/builds/23646

NEWS.md

src/DiagnosticVariables.jl

src/clima_diagnostics.jl

`LazyBroadcast.jl` provides a way to return an unevaluated function. This is useful in two cases: 1. reduce code verbosity to handle the `isnothing(out)` case 2. allow clustering all the broadcasted expressions in a single place In turn, 2. is useful because it is the first step in fusing different broadcasted calls. This commit adds support for such functions.

gemini-code-assist · 2025-03-27T00:13:17Z

Warning

You have reached your daily quota limit. As a reminder, free tier users are limited to 5 requests per day. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist · 2025-03-27T00:13:19Z

Warning

You have reached your daily quota limit. As a reminder, free tier users are limited to 5 requests per day. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist · 2025-03-27T00:13:21Z

Warning

You have reached your daily quota limit. As a reminder, free tier users are limited to 5 requests per day. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist · 2025-03-27T00:13:22Z

Warning

You have reached your daily quota limit. As a reminder, free tier users are limited to 5 requests per day. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist · 2025-03-27T00:13:42Z

Warning

You have reached your daily quota limit. As a reminder, free tier users are limited to 5 requests per day. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist · 2025-03-27T00:13:51Z

Warning

You have reached your daily quota limit. As a reminder, free tier users are limited to 5 requests per day. Please wait up to 24 hours and I will start processing your requests again!

charleskawczynski · 2025-03-27T00:21:02Z

src/clima_diagnostics.jl

+            out_field = variable.compute!(nothing, Y, p, t)
+        else
+            out_or_broadcasted = variable.compute(Y, p, t)
+            if out_or_broadcasted isa Base.Broadcast.Broadcasted


Suggested change

if out_or_broadcasted isa Base.Broadcast.Broadcasted

if out_or_broadcasted isa Base.AbstractBroadcasted

I think this may work more generally (e.g., in case the broadcasted is a ClimaCore stencil broadcasted)

charleskawczynski

I left one more comment (to use Base.AbstractBroadcasted instead of Base.Broadcast.Broadcasted, as it's a bit more general). I think we're good to merge after that. Thanks for taking this on, I think this will be a nice upgrade! 🚀

Sbozzolo force-pushed the gb/assign branch from 6e8e28c to c6b7525 Compare August 14, 2024 04:45

charleskawczynski mentioned this pull request Aug 14, 2024

Add and improve documentation CliMA/LazyBroadcast.jl#7

Closed

Sbozzolo force-pushed the gb/assign branch from c6b7525 to 47ace43 Compare August 14, 2024 16:03

charleskawczynski reviewed Aug 14, 2024

View reviewed changes

src/clima_diagnostics.jl Outdated Show resolved Hide resolved

charleskawczynski reviewed Aug 14, 2024

View reviewed changes

src/clima_diagnostics.jl Outdated Show resolved Hide resolved

Sbozzolo force-pushed the gb/assign branch 3 times, most recently from ba077ff to 34bf23f Compare August 14, 2024 17:59

Sbozzolo changed the title ~~Add @assign~~ Add @assign, add support for Broadcasted objects Aug 14, 2024

Sbozzolo force-pushed the gb/assign branch from 4c08bd8 to 7c448d3 Compare August 14, 2024 18:12

Sbozzolo requested a review from charleskawczynski August 14, 2024 22:21

Sbozzolo marked this pull request as draft September 9, 2024 15:32

Fix typo

20195df

Sbozzolo force-pushed the gb/assign branch 2 times, most recently from cd7e940 to 93ec15f Compare March 20, 2025 23:17

Sbozzolo mentioned this pull request Mar 20, 2025

Make some diagnostics lazy CliMA/ClimaAtmos.jl#3727

Draft

Sbozzolo changed the title ~~Add @assign, add support for Broadcasted objects~~ Add support for Broadcasted objects Mar 20, 2025

Sbozzolo force-pushed the gb/assign branch 2 times, most recently from 9e5b5bb to 540752b Compare March 21, 2025 17:51

Sbozzolo requested review from charleskawczynski and removed request for charleskawczynski March 24, 2025 15:12