intercept column spectral broadcast

Author: nefrathenrici

ext/cuda/operators_spectral_element.jl

@@ -35,17 +35,6 @@ function Base.copyto!(
     mask = DataLayouts.NoMask(),
 )
     space = axes(out)
-    # Column spaces don't have horizontal spectral element structure,
-    # so operators with empty axes just return zero without launching a kernel
-    if space isa Spaces.FiniteDifferenceSpace
-        fill!(
-            parent(Fields.field_values(out)),
-            zero(eltype(parent(Fields.field_values(out)))),
-        )
-        Operators.call_post_op_callback() &&
-            Operators.post_op_callback(out, out, sbc)
-        return out
-    end
     us = UniversalSize(Fields.field_values(out))
     # executed
     p = spectral_partition(us)

src/Operators/spectralelement.jl

@@ -130,6 +130,12 @@ function Base.Broadcast.instantiate(sbc::SpectralBroadcasted)
             Base.Broadcast.check_broadcast_axes(axes, args...)
         end
     end
+    # For FiniteDifferenceSpace, fall back to pointwise broadcasting
+    if axes isa Spaces.FiniteDifferenceSpace
+        return Base.Broadcast.instantiate(
+            Base.Broadcast.Broadcasted{Fields.FieldStyle}(op, args, axes)
+        )
+    end
     # If we've already instantiated, then we need to strip the type parameters,
     # for example, `Divergence{()}(axes)`.
     op = unionall_type(typeof(op)){()}(axes)
@@ -149,6 +155,12 @@ function Base.Broadcast.instantiate(
         axes = bc.axes
         Base.Broadcast.check_broadcast_axes(axes, args...)
     end
+    # For FiniteDifferenceSpace, fall back to pointwise broadcasting
+    if axes isa Spaces.FiniteDifferenceSpace
+        return Base.Broadcast.instantiate(
+            Base.Broadcast.Broadcasted{Fields.FieldStyle}(bc.f, args, axes)
+        )
+    end
     Style = AbstractSpectralStyle(ClimaComms.device(axes))
     return Base.Broadcast.Broadcasted{Style}(bc.f, args, axes)
 end