Replace LazyBroadcast with custom (uninstantiate) version

CliMA · Jan 21, 2025 · 5bdf352 · 5bdf352
1 parent adcf4d4
commit 5bdf352
Show file tree

Hide file tree

Showing 32 changed files with 95 additions and 75 deletions.
diff --git a/benchmarks/scripts/benchmark_utils.jl b/benchmarks/scripts/benchmark_utils.jl
@@ -1,7 +1,7 @@
 # import CUDA
 import ClimaComms
 using BenchmarkTools, Dates
-using LazyBroadcast: @lazy_broadcast
+using ClimaCore: @_lazy
 
 """
     caller_name(@__FILE__)

diff --git a/benchmarks/scripts/index_swapping.jl b/benchmarks/scripts/index_swapping.jl
@@ -63,7 +63,7 @@ end;
 function custom_kernel_bc!(X, Y, us::UniversalSizesStatic; swap=0, printtb=false, nreps = 1, print_info = true, bm=nothing, n_trials=30)
     (; x1, x2, x3) = X
     (; y1) = Y
-    bc = @lazy_broadcast @. foo(x1, x2, x3)
+    bc = @_lazy @. foo(x1, x2, x3)
     @assert !(y1 isa Array)
     f = if swap==0
         custom_kernel_knl_bc_0swap!

diff --git a/benchmarks/scripts/indexing_and_static_ndranges.jl b/benchmarks/scripts/indexing_and_static_ndranges.jl
@@ -298,7 +298,7 @@ end;
 function custom_kernel_bc!(X, Y, us::AbstractUniversalSizes; printtb=false, use_pw=true, nreps = 1, bm=nothing, n_trials = 30)
     (; x1, x2, x3) = X
     (; y1) = Y
-    bc_base = @lazy_broadcast @. myadd(x1, x2, x3)
+    bc_base = @_lazy @. myadd(x1, x2, x3)
     bc = use_pw ? to_pointwise_bc(bc_base) : bc_base
     e = Inf
     if y1 isa Array

diff --git a/src/ClimaCore.jl b/src/ClimaCore.jl
@@ -4,6 +4,9 @@ using PkgVersion
 const VERSION = PkgVersion.@Version
 import ClimaComms
 
+# Temporary helpers to work around https://github.com/CliMA/ClimaCore.jl/issues/2146
+include("lazy_broadcast.jl")
+
 include("DebugOnly/DebugOnly.jl")
 include("interface.jl")
 include("devices.jl")

diff --git a/src/Fields/broadcast.jl b/src/Fields/broadcast.jl
@@ -218,7 +218,7 @@ _check_mismatched_spaces(space1, space2) =
     error("FusedMultiBroadcast spaces are not the same.")
 
 @noinline function error_mismatched_spaces(space1::Type, space2::Type)
-    error("Broacasted spaces are not the same.")
+    error("Broacasted spaces are not the same: \n$space1\n\n$space2.")
 end
 
 @inline function Base.Broadcast.broadcast_shape(

diff --git a/src/lazy_broadcast.jl b/src/lazy_broadcast.jl
@@ -0,0 +1,20 @@
+# Temporary helpers to work around https://github.com/CliMA/ClimaCore.jl/issues/2146
+using Base.Broadcast: materialize, instantiate
+import Base.Broadcast: broadcasted
+function _lazy_broadcast end
+struct LazyBroadcasted{T}
+    value::T
+end
+Base.Broadcast.broadcasted(::typeof(_lazy_broadcast), x) = LazyBroadcasted(x)
+# Cannot return instantiated object here, due to https://github.com/CliMA/ClimaCore.jl/issues/2146
+Base.materialize(x::LazyBroadcasted) = x.value
+macro _lazy_broadcast(expr)
+    return quote
+        _lazy_broadcast.($(esc(expr)))
+    end
+end
+macro _lazy(expr)
+    return quote
+        _lazy_broadcast.($(esc(expr)))
+    end
+end
diff --git a/test/DataLayouts/unit_has_uniform_datalayouts.jl b/test/DataLayouts/unit_has_uniform_datalayouts.jl
@@ -6,7 +6,7 @@ using Test
 using ClimaCore.DataLayouts
 import ClimaCore.Geometry
 import ClimaComms
-import LazyBroadcast: @lazy_broadcast
+import ClimaCore: @_lazy
 using StaticArrays
 import Random
 Random.seed!(1234)
@@ -30,9 +30,9 @@ Random.seed!(1234)
     data_VIJFH = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh)
     data_VIFH = VIFH{S}(ArrayType{FT}, zeros; Nv, Ni, Nh)
 
-    bc = @lazy_broadcast @. data_VIFH + data_VIFH
+    bc = @_lazy @. data_VIFH + data_VIFH
     @test DataLayouts.has_uniform_datalayouts(bc)
-    bc = @lazy_broadcast @. data_IJFH + data_VF
+    bc = @_lazy @. data_IJFH + data_VF
     @test !DataLayouts.has_uniform_datalayouts(bc)
 
     data_VIJFHᶜ = VIJFH{S}(ArrayType{FT}, zeros; Nv, Nij, Nh)
@@ -41,6 +41,6 @@ Random.seed!(1234)
     # This is not a valid broadcast expression,
     # but these two datalayouts can exist in a
     # valid broadcast expression (e.g., interpolation).
-    bc = @lazy_broadcast @. data_VIJFHᶜ + data_VIJFHᶠ
+    bc = @_lazy @. data_VIJFHᶜ + data_VIJFHᶠ
     @test DataLayouts.has_uniform_datalayouts(bc)
 end
diff --git a/test/DataLayouts/unit_linear_indexing.jl b/test/DataLayouts/unit_linear_indexing.jl
@@ -15,7 +15,7 @@ using ClimaComms
 ClimaComms.@import_required_backends
 import ClimaCore: Domains, Meshes, Topologies, Spaces, Fields, Operators
 import ClimaCore.DataLayouts: get_struct
-using LazyBroadcast: @lazy_broadcast
+using ClimaCore: @_lazy
 
 Random.seed!(1234)
 

diff --git a/test/DataLayouts/unit_non_extruded_broadcast.jl b/test/DataLayouts/unit_non_extruded_broadcast.jl
@@ -5,15 +5,15 @@ using Revise; include(joinpath("test", "DataLayouts", "unit_non_extruded_broadca
 using Test
 using ClimaComms
 import Base.Broadcast: broadcasted, instantiate
-using LazyBroadcast: @lazy_broadcast
+using ClimaCore: @_lazy
 using ClimaCore.DataLayouts
 using ClimaCore.Geometry
 using ClimaCore: Geometry, Domains, Topologies, Meshes, Spaces, Fields
 
 @testset "unit_non_extruded_broadcast" begin
     a = [1, 2, 3]
     b = [10, 20, 30]
-    bc = @lazy_broadcast @. a + b
+    bc = @_lazy @. a + b
     bc = instantiate(bc)
     @test !DataLayouts.isascalar(bc)
     bc = DataLayouts.to_non_extruded_broadcasted(bc)
@@ -32,7 +32,7 @@ end
     data = DataF{S}(ArrayType{FT}, zeros)
     data[] = 5.0
 
-    bc = @lazy_broadcast @. data + data
+    bc = @_lazy @. data + data
     bc = instantiate(bc)
     @test !DataLayouts.isascalar(bc)
     bc = DataLayouts.to_non_extruded_broadcasted(bc)
@@ -50,7 +50,7 @@ foo(a, b, c) = a
     data = DataF{S}(ArrayType{FT}, zeros)
     data_empty = similar(data, typeof(()))
 
-    bc = @lazy_broadcast @. foo(data_empty, data_empty, ())
+    bc = @_lazy @. foo(data_empty, data_empty, ())
     bc = instantiate(bc)
     @test !DataLayouts.isascalar(bc)
     bc = DataLayouts.to_non_extruded_broadcasted(bc)
@@ -86,7 +86,7 @@ end
 
 @testset "Conceptual test (to compare against Base)" begin
     foo(a, b) = a
-    bc = @lazy_broadcast @. foo((), ())
+    bc = @_lazy @. foo((), ())
     bc = instantiate(bc)
     @test_throws BoundsError bc[]
 end
diff --git a/test/MatrixFields/matrix_field_test_utils.jl b/test/MatrixFields/matrix_field_test_utils.jl
@@ -3,7 +3,7 @@ using JET
 import Dates
 import Random: seed!
 import Base.Broadcast: materialize, materialize!
-import LazyBroadcast: @lazy_broadcast
+import ClimaCore: @_lazy
 import BenchmarkTools as BT
 
 import ClimaComms

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_1.jl b/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_1.jl
@@ -11,7 +11,7 @@ end
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "matrix of covectors times matrix of vectors" begin
 
-    bc = @lazy_broadcast @. ᶜᶠmat_AC1 ⋅ ᶠᶜmat_C12
+    bc = @_lazy @. ᶜᶠmat_AC1 ⋅ ᶠᶜmat_C12
     result = materialize(bc)
 
     ref_set_result! =

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_2.jl b/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_2.jl
@@ -13,8 +13,7 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
                  of numbers times matrix of covectors times matrix of \
                  vectors" begin
 
-    bc =
-        @lazy_broadcast @. ᶜᶠmat_AC1 ⋅ ᶠᶜmat_C12 ⋅ ᶜᶠmat ⋅ ᶠᶜmat_AC1 ⋅ ᶜᶠmat_C12
+    bc = @_lazy @. ᶜᶠmat_AC1 ⋅ ᶠᶜmat_C12 ⋅ ᶜᶠmat ⋅ ᶠᶜmat_AC1 ⋅ ᶜᶠmat_C12
     result = materialize(bc)
 
     ref_set_result! =

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_3.jl b/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_3.jl
@@ -13,7 +13,7 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
                  and covectors times matrix of numbers and vectors times \
                  vector of numbers" begin
 
-    bc = @lazy_broadcast @. ᶜᶠmat_AC1_num ⋅ ᶠᶜmat_C12_AC1 ⋅ ᶜᶠmat_num_C12 ⋅ ᶠvec
+    bc = @_lazy @. ᶜᶠmat_AC1_num ⋅ ᶠᶜmat_C12_AC1 ⋅ ᶜᶠmat_num_C12 ⋅ ᶠvec
     result = materialize(bc)
 
     ref_set_result! =

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_4.jl b/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_4.jl
@@ -13,7 +13,7 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
                  times matrix of numbers times matrix of numbers times \
                  vector of nested values" begin
 
-    bc = @lazy_broadcast @. ᶜᶠmat_NT ⋅ ᶠᶜmat ⋅ ᶜᶠmat ⋅ ᶠᶜmat_NT ⋅ ᶜvec_NT
+    bc = @_lazy @. ᶜᶠmat_NT ⋅ ᶠᶜmat ⋅ ᶜᶠmat ⋅ ᶠᶜmat_NT ⋅ ᶜvec_NT
     result = materialize(bc)
 
     ref_set_result! =

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_5.jl b/test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_5.jl
@@ -11,7 +11,7 @@ end
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "matrix of vectors divided by scalar" begin
 
-    bc = @lazy_broadcast @. ᶜᶠmat_C12 / 2
+    bc = @_lazy @. ᶜᶠmat_C12 / 2
     result = materialize(bc)
 
     ref_set_result! =

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_1.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_1.jl
@@ -9,7 +9,7 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "diagonal matrix times vector" begin
-    bc = @lazy_broadcast @. ᶜᶜmat ⋅ ᶜvec
+    bc = @_lazy @. ᶜᶜmat ⋅ ᶜvec
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜvec)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_10.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_10.jl
@@ -10,7 +10,7 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "diagonal matrix times bi-diagonal matrix times \
                  tri-diagonal matrix times quad-diagonal matrix times \
                  vector, but with forced right-associativity" begin
-    bc = @lazy_broadcast @. ᶜᶜmat ⋅ (ᶜᶠmat ⋅ (ᶠᶠmat ⋅ (ᶠᶜmat ⋅ ᶜvec)))
+    bc = @_lazy @. ᶜᶜmat ⋅ (ᶜᶠmat ⋅ (ᶠᶠmat ⋅ (ᶠᶜmat ⋅ ᶜvec)))
     if using_cuda
         @test_throws invalid_ir_error materialize(bc)
         @warn "cuda is broken for this test, exiting."

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_11.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_11.jl
@@ -8,8 +8,7 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "linear combination of matrix products and LinearAlgebra.I" begin
-    bc =
-        @lazy_broadcast @. 2 * ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat + ᶠᶠmat ⋅ ᶠᶠmat / 3 - (4I,)
+    bc = @_lazy @. 2 * ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat + ᶠᶠmat ⋅ ᶠᶠmat / 3 - (4I,)
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜᶠmat, ᶠᶠmat, ᶠᶜmat)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_12.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_12.jl
@@ -9,8 +9,7 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "another linear combination of matrix products and \
                  LinearAlgebra.I" begin
-    bc = @lazy_broadcast @. ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat * 2 - (ᶠᶠmat / 3) ⋅ ᶠᶠmat +
-                       (4I,)
+    bc = @_lazy @. ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat * 2 - (ᶠᶠmat / 3) ⋅ ᶠᶠmat + (4I,)
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜᶠmat, ᶠᶠmat, ᶠᶜmat)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_13.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_13.jl
@@ -8,9 +8,8 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "matrix times linear combination" begin
-    bc = @lazy_broadcast @. ᶜᶠmat ⋅ (
-        2 * ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat + ᶠᶠmat ⋅ ᶠᶠmat / 3 - (4I,)
-    )
+    bc = @_lazy @. ᶜᶠmat ⋅
+              (2 * ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat + ᶠᶠmat ⋅ ᶠᶠmat / 3 - (4I,))
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜᶠmat, ᶠᶠmat, ᶠᶜmat)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_14.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_14.jl
@@ -8,9 +8,8 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "linear combination times another linear combination" begin
-    bc = @lazy_broadcast @. (
-        2 * ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat + ᶠᶠmat ⋅ ᶠᶠmat / 3 - (4I,)
-    ) ⋅ (ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat * 2 - (ᶠᶠmat / 3) ⋅ ᶠᶠmat + (4I,))
+    bc = @_lazy @. (2 * ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat + ᶠᶠmat ⋅ ᶠᶠmat / 3 - (4I,)) ⋅
+              (ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat * 2 - (ᶠᶠmat / 3) ⋅ ᶠᶠmat + (4I,))
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜᶠmat, ᶠᶠmat, ᶠᶜmat)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_15.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_15.jl
@@ -9,9 +9,11 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "matrix times matrix times linear combination times matrix \
                  times another linear combination times matrix" begin
-    bc = @lazy_broadcast @. ᶠᶜmat ⋅ ᶜᶠmat ⋅ (
-        2 * ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat + ᶠᶠmat ⋅ ᶠᶠmat / 3 - (4I,)
-    ) ⋅ ᶠᶠmat ⋅ (ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat * 2 - (ᶠᶠmat / 3) ⋅ ᶠᶠmat + (4I,)) ⋅ ᶠᶠmat
+    bc = @_lazy @. ᶠᶜmat ⋅ ᶜᶠmat ⋅
+              (2 * ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat + ᶠᶠmat ⋅ ᶠᶠmat / 3 - (4I,)) ⋅
+              ᶠᶠmat ⋅
+              (ᶠᶜmat ⋅ ᶜᶜmat ⋅ ᶜᶠmat * 2 - (ᶠᶠmat / 3) ⋅ ᶠᶠmat + (4I,)) ⋅
+              ᶠᶠmat
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜᶠmat, ᶠᶠmat, ᶠᶜmat)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_16.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_16.jl
@@ -8,9 +8,9 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "matrix constructions and multiplications" begin
-    bc = @lazy_broadcast @. BidiagonalMatrixRow(ᶜᶠmat ⋅ ᶠvec, ᶜᶜmat ⋅ ᶜvec) ⋅
-                       TridiagonalMatrixRow(ᶠvec, ᶠᶜmat ⋅ ᶜvec, 1) ⋅ ᶠᶠmat ⋅
-                       DiagonalMatrixRow(DiagonalMatrixRow(ᶠvec) ⋅ ᶠvec)
+    bc = @_lazy @. BidiagonalMatrixRow(ᶜᶠmat ⋅ ᶠvec, ᶜᶜmat ⋅ ᶜvec) ⋅
+              TridiagonalMatrixRow(ᶠvec, ᶠᶜmat ⋅ ᶜvec, 1) ⋅ ᶠᶠmat ⋅
+              DiagonalMatrixRow(DiagonalMatrixRow(ᶠvec) ⋅ ᶠvec)
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜᶠmat, ᶠᶠmat, ᶠᶜmat, ᶜvec, ᶠvec)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_2.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_2.jl
@@ -8,7 +8,7 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "tri-diagonal matrix times vector" begin
-    bc = @lazy_broadcast @. ᶠᶠmat ⋅ ᶠvec
+    bc = @_lazy @. ᶠᶠmat ⋅ ᶠvec
     result = materialize(bc)
 
     input_fields = (ᶠᶠmat, ᶠvec)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_3.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_3.jl
@@ -8,7 +8,7 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "quad-diagonal matrix times vector" begin
-    bc = @lazy_broadcast @. ᶠᶜmat ⋅ ᶜvec
+    bc = @_lazy @. ᶠᶜmat ⋅ ᶜvec
     result = materialize(bc)
 
     input_fields = (ᶠᶜmat, ᶜvec)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_4.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_4.jl
@@ -8,7 +8,7 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "diagonal matrix times bi-diagonal matrix" begin
-    bc = @lazy_broadcast @. ᶜᶜmat ⋅ ᶜᶠmat
+    bc = @_lazy @. ᶜᶜmat ⋅ ᶜᶠmat
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜᶠmat)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_5.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_5.jl
@@ -8,7 +8,7 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "tri-diagonal matrix times tri-diagonal matrix" begin
-    bc = @lazy_broadcast @. ᶠᶠmat ⋅ ᶠᶠmat
+    bc = @_lazy @. ᶠᶠmat ⋅ ᶠᶠmat
     result = materialize(bc)
 
     input_fields = (ᶠᶠmat,)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_6.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_6.jl
@@ -8,7 +8,7 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 #! format: on
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "quad-diagonal matrix times diagonal matrix" begin
-    bc = @lazy_broadcast @. ᶠᶜmat ⋅ ᶜᶜmat
+    bc = @_lazy @. ᶠᶜmat ⋅ ᶜᶜmat
     result = materialize(bc)
 
     input_fields = (ᶠᶜmat, ᶜᶜmat)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_7.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_7.jl
@@ -9,7 +9,7 @@ include(joinpath(pkgdir(ClimaCore),"test","MatrixFields","matrix_fields_broadcas
 test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "diagonal matrix times bi-diagonal matrix times \
                  tri-diagonal matrix times quad-diagonal matrix" begin
-    bc = @lazy_broadcast @. ᶜᶜmat ⋅ ᶜᶠmat ⋅ ᶠᶠmat ⋅ ᶠᶜmat
+    bc = @_lazy @. ᶜᶜmat ⋅ ᶜᶠmat ⋅ ᶠᶠmat ⋅ ᶠᶜmat
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜᶠmat, ᶠᶠmat, ᶠᶜmat)

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_8.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_8.jl
@@ -10,7 +10,7 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "diagonal matrix times bi-diagonal matrix times \
                  tri-diagonal matrix times quad-diagonal matrix, but with \
                  forced right-associativity" begin
-    bc = @lazy_broadcast @. ᶜᶜmat ⋅ (ᶜᶠmat ⋅ (ᶠᶠmat ⋅ ᶠᶜmat))
+    bc = @_lazy @. ᶜᶜmat ⋅ (ᶜᶠmat ⋅ (ᶠᶠmat ⋅ ᶠᶜmat))
     if using_cuda
         @test_throws invalid_ir_error materialize(bc)
         @warn "cuda is broken for this test, exiting."

diff --git a/test/MatrixFields/matrix_fields_broadcasting/test_scalar_9.jl b/test/MatrixFields/matrix_fields_broadcasting/test_scalar_9.jl
@@ -10,7 +10,7 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
 @testset "diagonal matrix times bi-diagonal matrix times \
                  tri-diagonal matrix times quad-diagonal matrix times \
                  vector" begin
-    bc = @lazy_broadcast @. ᶜᶜmat ⋅ ᶜᶠmat ⋅ ᶠᶠmat ⋅ ᶠᶜmat ⋅ ᶜvec
+    bc = @_lazy @. ᶜᶜmat ⋅ ᶜᶠmat ⋅ ᶠᶠmat ⋅ ᶠᶜmat ⋅ ᶜvec
     result = materialize(bc)
 
     input_fields = (ᶜᶜmat, ᶜᶠmat, ᶠᶠmat, ᶠᶜmat, ᶜvec)