Register the lazy-composition test in CI and cover the lazy path

Adds CPU and CUDA Buildkite steps for the new test, an assertion using
LazyBroadcast.lazy alongside the explicit instantiate, and trims the
style-rule comment.

Author: haakon-e

.buildkite/pipeline.yml

@@ -705,6 +705,11 @@ steps:
         retry: *retry_policy
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Operators/spectralelement/rectilinear.jl"
 
+      - label: "Unit: lazy composition"
+        key: unit_lazy_composition
+        retry: *retry_policy
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Operators/spectralelement/lazy_composition.jl"
+
       - label: "Unit: diffusion2d"
         key: unit_diffusion2d
         retry: *retry_policy
@@ -926,6 +931,15 @@ steps:
         agents:
           slurm_gpus: 1
 
+      - label: "Unit: lazy composition cuda"
+        key: unit_lazy_composition_cuda
+        retry: *retry_policy
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Operators/spectralelement/lazy_composition.jl"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+
       - label: "Unit: operators on levels and extruded spaces"
         key: unit_cuda_operators_example
         retry: *retry_policy

src/Operators/spectralelement.jl

@@ -507,10 +507,8 @@ Base.Broadcast.BroadcastStyle(
     ::Fields.AbstractFieldStyle,
 ) = style
 
-# The device-resolved style assigned by `instantiate` takes precedence over the
-# generic construction-time style, so an instantiated spectral broadcast can be
-# used as an argument to a new operator broadcast. See the finite-difference
-# analogue for `ColumnStencilStyle` in the CUDA extension.
+# The device-resolved style assigned by `instantiate` takes precedence over
+# the generic construction-time style.
 Base.Broadcast.BroadcastStyle(
     ::SpectralStyle,
     style::AbstractSpectralStyle,

test/Operators/spectralelement/lazy_composition.jl

@@ -12,6 +12,7 @@ import ClimaCore:
     Quadratures
 import ClimaCore.Operators as O
 import Base.Broadcast as BB
+import LazyBroadcast: lazy
 using LinearAlgebra, IntervalSets
 
 @testset "spectral broadcast style combination" begin
@@ -48,4 +49,7 @@ end
     )
     out = @. wdiv(K * ∇f)
     @test parent(out) ≈ 2 .* parent(ref)
+    ∇f_lazy = @. lazy(grad(f) + grad(f))
+    out_lazy = @. wdiv(K * ∇f_lazy)
+    @test parent(out_lazy) ≈ 2 .* parent(ref)
 end