Add support for unevaluated compute! functions

`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.

Author: Sbozzolo

.buildkite/Project.toml

@@ -12,6 +12,7 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
+LazyBroadcast = "9dccce8e-a116-406d-9fcc-a88ed4f510c8"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"

NEWS.md

@@ -110,6 +110,7 @@ v0.2.4
 
 - Add `EveryCalendarDtSchedule` for schedules with calendar periods.
 
+
 v0.2.3
 -------
 

Project.toml

@@ -24,6 +24,7 @@ ClimaUtilities = "0.1.22"
 Dates = "1"
 Documenter = "1"
 ExplicitImports = "1.6"
+LazyBroadcast = "0.1.4"
 JuliaFormatter = "1"
 NCDatasets = "0.14"
 OrderedCollections = "1.4"
@@ -41,10 +42,11 @@ ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
+LazyBroadcast = "9dccce8e-a116-406d-9fcc-a88ed4f510c8"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 ProfileCanvas = "efd6af41-a80b-495e-886c-e51b0c7d77a3"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "BenchmarkTools", "ClimaTimeSteppers", "Documenter", "ExplicitImports", "JuliaFormatter", "Profile", "ProfileCanvas", "SafeTestsets", "Test"]
+test = ["Aqua", "BenchmarkTools", "ClimaTimeSteppers", "Documenter", "ExplicitImports", "JuliaFormatter", "LazyBroadcast", "Profile", "ProfileCanvas", "SafeTestsets", "Test"]

docs/src/developer_guide.md

@@ -246,6 +246,27 @@ add_diagnostic_variable!(
 )
 ```
 
+When writing compute functions, consider using [`@assign`](@ref) to simplify
+your code and, when possible, making your expression lazy with
+[LazyBroadcast.jl](https://github.com/CliMA/LazyBroadcast.jl) to further improve
+clarity and performance. To do that, add `LazyBroadcast` to your dependencies
+and import `@lazy`. The previous example would look like:
+
+```julia
+###
+# Density (3d)
+###
+add_diagnostic_variable!(
+    short_name = "rhoa",
+    long_name = "Air Density",
+    standard_name = "air_density",
+    units = "kg m^-3",
+    compute! = (out, state, cache, time) -> begin
+        return @lazy @. state.c.ρ
+    end,
+)
+```
+
 It is a good idea to put safeguards in place to ensure that your users will not
 be allowed to call diagnostics that do not make sense for the simulation they
 are running. If your package has a notion of `Model` that is stored in `p`, you
@@ -269,11 +290,7 @@ function compute_hus!(
     time,
     moisture_model::T,
 ) where {T <: Union{EquilMoistModel, NonEquilMoistModel}}
-    if isnothing(out)
-        return state.c.ρq_tot ./ state.c.ρ
-    else
-        out .= state.c.ρq_tot ./ state.c.ρ
-    end
+    @assign out state.c.ρq_tot ./ state.c.ρ
 end
 
 add_diagnostic_variable!(

docs/src/index.md

@@ -24,5 +24,6 @@ the Developer guide page.
 - Allow users to define arbitrary new diagnostics;
 - Trigger diagnostics on arbitrary conditions;
 - Save output to HDF5 or NetCDF files, or a dictionary in memory;
-
+- Work with lazy expressions (such as the ones produced by
+  [LazyBroadcast.jl](https://github.com/CliMA/LazyBroadcast.jl)).
 

docs/src/user_guide.md

@@ -58,7 +58,8 @@ var = DiagnosticVariable(;
 `compute_ta!` is the key function here. It determines how the variable should be
 computed from the `state`, `cache`, and `time` of the simulation. Typically,
 these are packaged within an `integrator` object (e.g., `state = integrator.u`
-or `integrator.Y`).
+or `integrator.Y`). `copy` is needed because we want to return a new area of
+memory (without `copy`, `ClimaDiagnostics` might end up modifying `state.ta`).
 
 `compute_ta!` takes another argument, `out`. `out` is an area of memory managed
 by `ClimaDiagnostics` that is used to reduce the number of allocations needed
@@ -67,9 +68,18 @@ of memory is allocated and filled with the value (this is when `out` is
 `nothing`). All the subsequent times, the same space is overwritten, leading to
 much better performance. You should follow this pattern in all your diagnostics.
 
-> Note, in the future, we hope to improve this rather clumsy way to write
-> diagnostics. Hopefully, at some point you will just have to write something like
-> `state.ta` and not worry about the `out` at all.
+`ClimaDiagnostics` supports working with unevaluated expressions represented by
+`Base.Broadcast.Broadcasted` objects, such as the ones produced with
+[LazyBroadcast.jl](https://github.com/CliMA/LazyBroadcast.jl). Using
+`LazyBroadcast.jl`, the snippet above can be rewritten as
+```julia
+import LazyBroadcast: @lazy
+
+function compute_ta!(state, cache, time)
+    return @lazy @. state.ta
+end
+```
+Using lazy expressions can lead to improved performance and clearer code.
 
 A `DiagnosticVariable` defines what a variable is and how to compute it, but
 does not specify when to compute/output it. For that, we need

src/clima_diagnostics.jl

@@ -13,7 +13,14 @@ include("reduction_identities.jl")
 A struct that contains the scheduled diagnostics, ancillary data and areas of memory needed
 to store and accumulate results.
 """
-struct DiagnosticsHandler{SD, V <: Vector{Int}, STORAGE, ACC <: Dict, COUNT}
+struct DiagnosticsHandler{
+    SD,
+    V <: Vector{Int},
+    STORAGE,
+    ACC <: Dict,
+    COUNT,
+    BROAD,
+}
     """An iterable with the `ScheduledDiagnostic`s that are scheduled."""
     scheduled_diagnostics::SD
 
@@ -31,6 +38,11 @@ struct DiagnosticsHandler{SD, V <: Vector{Int}, STORAGE, ACC <: Dict, COUNT}
     """Container holding a counter that tracks how many times the given
     diagnostics was computed from the last time it was output to disk."""
     counters::COUNT
+
+    """Dictionary that maps a given `ScheduledDiagnostic` to a Base.Broadcast.Broadcasted
+    object. This is used to allow lazy evaluation of expressions, which can lead to reduced
+    code verbosity and improved performance."""
+    broadcasted_expressions::BROAD
 end
 
 """
@@ -78,6 +90,10 @@ function DiagnosticsHandler(scheduled_diagnostics, Y, p, t; dt = nothing)
     counters = Int[]
     scheduled_diagnostics_keys = Int[]
 
+    # broadcasted_expressions contains the un-evaluated expression, which are then moved to
+    # storage
+    broadcasted_expressions = Dict()
+
     # NOTE: unique requires isequal and hash to both be implemented. We don't
     # really want to do that. So, we roll our own unique. This is O(N^2) but it
     # is run only once, so it should be fine.
@@ -119,11 +135,31 @@ function DiagnosticsHandler(scheduled_diagnostics, Y, p, t; dt = nothing)
         isa_time_reduction = !isnothing(diag.reduction_time_func)
 
         # The first time we call compute! we use its return value. All the subsequent times
-        # (in the callbacks), we will write the result in place. We call copy to acquire ownership
-        # of the data in case compute! returned a reference.
-        push!(storage, copy(variable.compute!(nothing, Y, p, t)))
-        push!(counters, 1)
+        # (in the callbacks), we will write the result in place.
+        #
+        # ClimaDiagnostics supports LazyBroadcast.jl. In this case, the return value of
+        # `compute!` is a `Base.Broadcast.Broadcasted` and we have to manually materialize
+        # the result. When using LazyBroadcast.jl, compute! does not need to have the `out`
+        # argument, so we have to check how many arguments to pass.
+
+        three_args = (Y, p, t)
 
+        full_args =
+            hasmethod(variable.compute!, typeof.(three_args)) ? three_args :
+            (nothing, three_args...)
+
+        out_or_broadcasted_expr = variable.compute!(full_args...)
+
+        if out_or_broadcasted_expr isa Base.Broadcast.Broadcasted
+            broadcasted_expressions[diag] = out_or_broadcasted_expr
+            push!(storage, Base.Broadcast.materialize(out_or_broadcasted_expr))
+        else
+            # We call copy to acquire ownership of the data in case compute! returned a
+            # reference.
+            push!(storage, copy(out_or_broadcasted_expr))
+        end
+
+        push!(counters, 1)
         # If it is not a reduction, call the output writer as well
         if !isa_time_reduction
             interpolate_field!(diag.output_writer, storage[i], diag, Y, p, t)
@@ -148,6 +184,7 @@ function DiagnosticsHandler(scheduled_diagnostics, Y, p, t; dt = nothing)
         storage,
         accumulators,
         counters,
+        broadcasted_expressions,
     )
 end
 
@@ -175,6 +212,10 @@ function orchestrate_diagnostics(
     active_output = Bool[]
     active_sync = Bool[]
 
+    # Used in the compute loop
+    three_args_type =
+        (typeof(integrator.u), typeof(integrator.p), typeof(integrator.t))
+
     for diag in scheduled_diagnostics
         push!(active_compute, diag.compute_schedule_func(integrator))
         push!(active_output, diag.output_schedule_func(integrator))
@@ -186,13 +227,52 @@ function orchestrate_diagnostics(
         active_compute[diag_index] || continue
         diag = scheduled_diagnostics[diag_index]
 
-        diag.variable.compute!(
+        diagnostic_handler.counters[diag_index] += 1
+
+        # ClimaDiagnostics supports LazyBroadcast.jl. When used, the return value of
+        # `compute!` is a `Base.Broadcast.Broadcasted`. We materialize the output to
+        # diagnostic_handler.storage[diag] in the next for loop. If the output is not a
+        # `Base.Broadcast.Broadcasted`, we don't have to do anything because compute! will
+        # already update diagnostic_handler.storage[diag]. Here, too, we have to check if
+        # compute! is the three or four argument variant.
+        #
+        # Here we use a more verbose way of writing this to avoid working with tuples with
+        # very complex types (that might increase compilation costs).
+
+        if hasmethod(diag.variable.compute!, three_args_type)
+            out_or_broadcasted_expr =
+                diag.variable.compute!(integrator.u, integrator.p, integrator.t)
+        else
+            out_or_broadcasted_expr = diag.variable.compute!(
+                diagnostic_handler.storage[diag_index],
+                integrator.u,
+                integrator.p,
+                integrator.t,
+            )
+        end
+
+        if out_or_broadcasted_expr isa Base.Broadcast.Broadcasted
+            diagnostic_handler.broadcasted_expressions[diag] =
+                out_or_broadcasted_expr
+        end
+    end
+
+    # Evaluate the lazy compute (aka, materialize everything)
+    for diag_index in 1:length(scheduled_diagnostics)
+        active_compute[diag_index] || continue
+        diag = scheduled_diagnostics[diag_index]
+        haskey(diagnostic_handler.broadcasted_expressions, diag) || continue
+
+        Base.Broadcast.materialize!(
             diagnostic_handler.storage[diag_index],
-            integrator.u,
-            integrator.p,
-            integrator.t,
+            diagnostic_handler.broadcasted_expressions[diag],
         )
-        diagnostic_handler.counters[diag_index] += 1
+    end
+
+    # Process possible time reductions (now we have evaluated storage[diag])
+    for diag_index in 1:length(scheduled_diagnostics)
+        active_compute[diag_index] || continue
+        diag = scheduled_diagnostics[diag_index]
 
         isa_time_reduction = !isnothing(diag.reduction_time_func)
         if isa_time_reduction

test/integration_test.jl

@@ -14,6 +14,8 @@ import ClimaComms
     ClimaComms.@import_required_backends
 end
 
+import LazyBroadcast: @lazy
+
 const context = ClimaComms.context()
 ClimaComms.init(context)
 
@@ -52,12 +54,32 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
         end
     end
 
+    function compute_my_var_lazy!(out, u, p, t)
+        return @lazy @. out = u.my_var
+    end
+
+    function compute_my_var_lazy_three_args!(u, p, t)
+        return @lazy @. u.my_var
+    end
+
     simple_var = ClimaDiagnostics.DiagnosticVariable(;
         compute! = compute_my_var!,
         short_name = "YO",
         long_name = "YO YO",
     )
 
+    simple_var_lazy = ClimaDiagnostics.DiagnosticVariable(;
+        compute! = compute_my_var_lazy!,
+        short_name = "YO LAZY",
+        long_name = "YO YO LAZY",
+    )
+
+    simple_var_lazy_three = ClimaDiagnostics.DiagnosticVariable(;
+        compute! = compute_my_var_lazy_three_args!,
+        short_name = "YO LAZY THREE",
+        long_name = "YO YO LAZY THREE",
+    )
+
     average_diagnostic = ClimaDiagnostics.ScheduledDiagnostic(
         variable = simple_var,
         output_writer = nc_writer,
@@ -69,6 +91,14 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
         variable = simple_var,
         output_writer = nc_writer,
     )
+    inst_diagnostic_lazy = ClimaDiagnostics.ScheduledDiagnostic(
+        variable = simple_var_lazy,
+        output_writer = nc_writer,
+    )
+    inst_diagnostic_lazy_three = ClimaDiagnostics.ScheduledDiagnostic(
+        variable = simple_var_lazy_three,
+        output_writer = nc_writer,
+    )
     inst_every3s_diagnostic = ClimaDiagnostics.ScheduledDiagnostic(
         variable = simple_var,
         output_writer = nc_writer,
@@ -92,6 +122,8 @@ function setup_integrator(output_dir; context, more_compute_diagnostics = 0)
     scheduled_diagnostics = [
         average_diagnostic,
         inst_diagnostic,
+        inst_diagnostic_lazy,
+        inst_diagnostic_lazy_three,
         inst_diagnostic_h5,
         inst_every3s_diagnostic,
         inst_every3s_diagnostic_another,