Fix eager FD kernel bugs and replace op-specific isa checks with dispatch

Bug fixes (each reproduced on GPU before the fix):

- ext/cuda/operators_finite_difference.jl: read `mask.N[1]` under
  `CUDA.@allowscalar`. It is a one-element device array, so the masked eager
  path threw "Scalar indexing is disallowed" on its first call. Matches the
  three existing read sites in data_layouts_threadblock.jl.

- src/Operators/finitedifference.jl: restore
  `boundary_width(::DivergenceF2C, ::Extrapolate) = 1`. Without it the boundary
  centers fall in the interior window, so `op_matrix_first_row`/`_last_row` for
  `Extrapolate` were unreachable and the condition was silently ignored --
  `divf2c_extrapolate(u)` returned the plain interior stencil at the boundary.
  `test_op_matrix(DivergenceF2C, Extrapolate, ...)` could not catch this because
  both the matrix and operator paths degraded identically.

- ext/cuda/operators_fd_eager.jl: size the dynamic shared memory from the
  actual projected operand type (`cached_operand_type`) instead of the
  `AbstractTensor ? sizeof(eltype)*9 : sizeof` heuristic. `project_row2_for_mul`
  projects tensor leaves nested inside a `BandMatrixRow`, which the heuristic
  did not see: a `QuaddiagonalMatrixRow{Covariant1Vector{Float64}}` was budgeted
  32 bytes and written as 96, crashing the kernel with ERROR_ILLEGAL_ADDRESS.
  The new version mirrors `project_row2_for_mul`'s type logic, so it cannot
  drift from what the kernel writes.

- .buildkite/pipeline.yml: drop the stray `soft_fail: true` from the
  `agents:` block of unit_matrix_field_broadcasting_gpu_non_scalar_3. It is a
  step attribute, not an agent tag; there it became an agent-tag query that
  would leave the step unscheduled. The correct one three lines above stays.

Op-specific `isa` checks in generic machinery, now handled by dispatch:

- `should_call_left_boundary`/`should_call_right_boundary` no longer special-case
  `SetBoundaryOperator` + `NullBoundaryCondition`. Instead
  `boundary_width(::SetBoundaryOperator, ...)` returns 0 by default and 1 for the
  four value-fixing conditions, matching the idiom already used for
  GradientF2C/DivergenceF2C. This also fixes what the guard could not reach:
  `left_interior_window_idx`/`right_interior_window_idx` bypass
  `should_call_*`, so one-sided SetBoundaryOperators were needlessly widening
  the boundary window by a level. `should_call_right_boundary` now also uses the
  `boundary_condition` local it already binds.

- The `op_matrix.op isa Union{GradientF2C, DivergenceF2C} && bc isa SetValue`
  branch inside the generic `FDOperatorMatrix` `stencil_*_boundary` methods
  becomes a more specific method on `FDOperatorMatrix{<:InputFixingFDOperator}`.

- The `op isa SetBoundaryOperator` early return in the generic
  `OneArgFDOperator` StencilBroadcasted constructor becomes its own constructor
  method; the gradient/divergence adjoint ternaries become
  `adjoint_matrix_arg`/`adjoint_matrix_result`. The generic constructor now has
  no `isa` checks.

- Drop `split_bcs = op isa WeightedInterpolateC2F` from the `TwoArgFDOperator`
  constructor. Every other two-argument operator's conditions are linear, so
  `output_bcs` is already empty and `op_with_matrix_bcs` already returns the op
  unchanged; the flag was both a hardcoded op name and redundant.

`detect_ambiguities` is empty for both Operators and MatrixFields. Verified on
CPU and GPU: MatrixFields/operator_matrices.jl (CPU 275+41, GPU 207+25),
finitedifference/unit_column.jl (25 each), and
matrix_fields_broadcasting/test_non_scalar_2.jl (GPU 2) all pass.
test_non_scalar_3.jl remains red, but with a pre-existing InvalidIRError over
`AutoBroadcaster` row entries that reproduces with these changes stashed, so its
soft_fail is left in place.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BLMKMmk7JezBvagh3AybPQ

Author: imreddyTeja

.buildkite/pipeline.yml

@@ -1376,7 +1376,6 @@ steps:
         agents:
           slurm_gpus: 1
           slurm_mem: 10GB
-          soft_fail: true
 
       - label: "Unit: matrix field broadcasting (GPU)"
         key: unit_matrix_field_broadcasting_gpu_non_scalar_4

ext/cuda/operators_fd_eager.jl

@@ -38,22 +38,53 @@ fit and can always be cached.
 max_eager_shmem_per_thread(x) = 0
 max_eager_shmem_per_thread(bc::Union{Broadcasted, StencilBroadcasted}) =
     _max_eager_shmem_over_args(bc.args)
-function max_eager_shmem_per_thread(
+max_eager_shmem_per_thread(
     bc::StencilBroadcasted{S, <:MultiplyColumnwiseBandMatrixField},
-) where {S}
-    raw_arg_type = unsafe_eltype(bc.args[2])
-    # tensors may be projected to contain more components
-    # this is hacky and will sometimes overestimate the size of the projected row
-    modified_arg_size = raw_arg_type <: AbstractTensor ? sizeof(eltype(raw_arg_type)) * 9 : sizeof(raw_arg_type)
-    max(modified_arg_size, _max_eager_shmem_over_args(bc.args))
-end
+) where {S} =
+    max(sizeof(cached_operand_type(bc)), _max_eager_shmem_over_args(bc.args))
 
 _max_eager_shmem_over_args(::Tuple{}) = 0
 _max_eager_shmem_over_args(args::Tuple) = max(
     max_eager_shmem_per_thread(first(args)),
     _max_eager_shmem_over_args(Base.tail(args)),
 )
 
+"""
+    cached_operand_type(bc)
+
+The type that `calc_level_val` writes into shared memory for the multiplication `bc`,
+i.e. `typeof(project_row2_for_mul(mat1_row, mat2_row, hidx, mat2_space))`.
+
+The size cannot be read off the second operand directly, because
+`project_row2_for_mul` projects every tensor leaf of that operand onto the axis dual
+to the first operand's entries, which changes its size in either direction: a
+`Covariant1Vector` widens to a `Contravariant123Vector`, while a `Covariant12Vector`
+narrows to a `Contravariant1Vector`. For a matrix-matrix product those leaves are also
+nested inside a `BandMatrixRow`, so no property of the operand's outermost type
+bounds the projected size. Mirror `project_row2_for_mul`'s type-level logic instead
+and infer the projected type, so the buffer is always big enough for what the kernel
+writes into it.
+"""
+function cached_operand_type(bc)
+    mat1_type = unsafe_eltype(bc.args[1i32])
+    mat2_type = unsafe_eltype(bc.args[2i32])
+    mat1_et = mat1_type <: BandMatrixRow ? eltype(mat1_type) : mat1_type
+    project_onto =
+        ClimaCore.Geometry.recursively_find_dual_axes_for_projection(mat1_et)
+    isnothing(project_onto) && return mat2_type
+    lg_type = Spaces.local_geometry_type(typeof(axes(bc.args[2i32])))
+    projected_type = ClimaCore.Utilities.return_type(
+        recursively_project,
+        Tuple{Tuple{typeof(project_onto), lg_type}, mat2_type},
+    )
+    isconcretetype(projected_type) || error(
+        "Unable to size the eager finite difference kernel's shared memory: \
+         inference gave the non-concrete type $projected_type for the \
+         projection of a $mat2_type operand onto $project_onto",
+    )
+    return projected_type
+end
+
 
 ClimaCore.Utilities.unsafe_eltype(::CUDA.CuRefType{T}) where {T} = T
 

ext/cuda/operators_finite_difference.jl

@@ -94,7 +94,11 @@ function Base.copyto!(
         eager_shmem_per_thread = max_eager_shmem_per_thread(bc′)
         if !high_resolution
             #    32 <= n_face_levels * Ni <= 256
-            n_columns = mask isa NoMask ? Ni * Nj * Nh : mask.N[1]
+            # mask.N holds the active column count in a one-element device array;
+            # reading it on the host needs @allowscalar.
+            n_columns =
+                mask isa NoMask ? Ni * Nj * Nh :
+                CUDA.@allowscalar(mask.N[1])
             # 108 is the number of SMs in an A100. TODO: get this value from CUDA.jl to better optimize for different GPUs
             threads_dim_y = n_columns > 256 * 108 ? div(256, n_face_levels) : 1
             block_dim_x = div(n_columns, threads_dim_y, RoundUp)

src/MatrixFields/operator_matrices.jl

@@ -215,6 +215,47 @@ apply_boundary_operator(::Type{Style}, op, arg, axes, work) where {Style} =
         )
     end
 
+# A gradient operator matrix has vector entries and a divergence operator matrix has
+# covector entries, so for the plain `*` of the matrix multiply to produce a result of
+# the right rank, a gradient needs an adjoint on its argument and a divergence needs
+# one on its result.
+adjoint_matrix_arg(op, arg) = arg
+adjoint_matrix_arg(::Operators.GradientOperator, arg) =
+    Base.Broadcast.broadcasted(adjoint, arg)
+adjoint_matrix_result(op, result) = result
+adjoint_matrix_result(::Operators.DivergenceOperator, result) =
+    Base.Broadcast.broadcasted(adjoint, result)
+
+# Builds an ordinary StencilBroadcasted, without rewriting `op` into a matrix multiply.
+unconverted_stencil_broadcasted(
+    ::Type{Style},
+    op,
+    args::Args,
+    axes,
+    work::Work,
+) where {Style, Args, Work} = Operators.StencilBroadcasted{
+    Style,
+    typeof(op),
+    Args,
+    typeof(axes),
+    Work,
+}(
+    op,
+    args,
+    axes,
+    work,
+)
+
+# A SetBoundaryOperator has no operator matrix: it is what the conversions below use to
+# reapply the boundary conditions they strip out, so it is built verbatim.
+Operators.StencilBroadcasted{Style}(
+    op::Operators.SetBoundaryOperator,
+    args::Args,
+    axes::Spaces.AbstractSpace,
+    work::Work = nothing,
+) where {Style, Args, Work} =
+    unconverted_stencil_broadcasted(Style, op, args, axes, work)
+
 # Converts a broadcast over a one-argument operator, `op(arg)`, into the
 # equivalent operator matrix expression, `op_matrix() * arg`. Boundary conditions
 # that modify the operator's input or output are stripped from the matrix and
@@ -227,20 +268,6 @@ function Operators.StencilBroadcasted{Style}(
     axes::Spaces.AbstractSpace,
     work::Work = nothing,
 ) where {Style, Args, Work}
-    # SetBoundaryOperator is not represented by an operator matrix.
-    op isa Operators.SetBoundaryOperator && return Operators.StencilBroadcasted{
-        Style,
-        typeof(op),
-        Args,
-        typeof(axes),
-        Work,
-    }(
-        op,
-        args,
-        axes,
-        work,
-    )
-
     op_matrix = Base.Broadcast.broadcasted(
         FDOperatorMatrix(op_with_matrix_bcs(op)),
         Fields.local_geometry_field(operator_input_space(op, axes)),
@@ -256,14 +283,10 @@ function Operators.StencilBroadcasted{Style}(
             Base.axes(args[1]),
             nothing,
         )
-    arg =
-        op isa Operators.GradientOperator ?
-        Base.Broadcast.broadcasted(adjoint, arg) : arg
+    arg = adjoint_matrix_arg(op, arg)
 
     result = multiply_matrix_broadcasted(Style, op_matrix, arg, axes, work)
-    result =
-        op isa Operators.DivergenceOperator ?
-        Base.Broadcast.broadcasted(adjoint, result) : result
+    result = adjoint_matrix_result(op, result)
 
     bcs_out = output_bcs(op)
     return isempty(bcs_out) ? result :
@@ -278,23 +301,24 @@ end
 
 
 # Converts a broadcast over a two-argument operator, `op(weight, arg)`, into the
-# equivalent operator matrix expression, `op_matrix(weight) * arg`. Only
-# WeightedInterpolateC2F has boundary conditions (SetValue) that modify its
-# output; when present, they are stripped from the matrix and reapplied to the
-# result with a SetBoundaryOperator.
+# equivalent operator matrix expression, `op_matrix(weight) * arg`. As for one-argument
+# operators, boundary conditions that modify the output are stripped from the matrix and
+# reapplied to the result with a SetBoundaryOperator. In practice only
+# WeightedInterpolateC2F has such conditions (SetValue); every other two-argument
+# operator's conditions are linear, so `output_bcs` is empty for them and both
+# `op_with_matrix_bcs` and this function leave them untouched.
 function Operators.StencilBroadcasted{Style}(
     op::TwoArgFDOperator,
     args::Args,
     axes::Spaces.AbstractSpace,
     work::Work = nothing,
 ) where {Style, Args, Work}
-    split_bcs = op isa Operators.WeightedInterpolateC2F
-    matrix_op = split_bcs ? op_with_matrix_bcs(op) : op
-    op_matrix = Base.Broadcast.broadcasted(FDOperatorMatrix(matrix_op), args[1])
+    op_matrix =
+        Base.Broadcast.broadcasted(FDOperatorMatrix(op_with_matrix_bcs(op)), args[1])
 
     result = multiply_matrix_broadcasted(Style, op_matrix, args[2], axes, work)
 
-    bcs_out = split_bcs ? output_bcs(op) : (;)
+    bcs_out = output_bcs(op)
     return isempty(bcs_out) ? result :
            apply_boundary_operator(
         Style,
@@ -527,53 +551,78 @@ Operators.stencil_right_boundary(
 # constant contributed by the boundary value is zeroed out (see `has_affine_bc`).
 # For every operator except GradientF2C/DivergenceF2C, a value-fixing condition
 # prescribes the *output* at the boundary as a pure constant, so its linear part is
-# zero and the boundary row is empty (`rzero`). GradientF2C/DivergenceF2C with a
-# SetValue are the exception (as with `modifies_input`): the condition fixes an
-# *input* value, and the near-boundary output still depends linearly on the
-# adjacent interior input, so the row is the genuine boundary stencil (with the
-# fixed input's coefficient dropped) rather than zero.
-Base.@propagate_inbounds function Operators.stencil_left_boundary(
+# zero and the boundary row is empty (`rzero`).
+const ValueFixingBoundaryCondition = Union{
+    Operators.SetValue,
+    Operators.SetGradient,
+    Operators.SetDivergence,
+    Operators.SetCurl,
+}
+# The operators whose SetValue fixes an *input* rather than an output, and so keep a
+# genuine boundary stencil in the matrix (see `modifies_input`).
+const InputFixingFDOperator =
+    Union{Operators.GradientF2C, Operators.DivergenceF2C}
+
+Base.@propagate_inbounds Operators.stencil_left_boundary(
     op_matrix::FDOperatorMatrix,
-    bc::Union{
-        Operators.SetValue,
-        Operators.SetGradient,
-        Operators.SetDivergence,
-        Operators.SetCurl,
-    },
+    ::ValueFixingBoundaryCondition,
+    space,
+    idx,
+    hidx,
+    args...,
+) = rzero(Operators.return_eltype(op_matrix, args...))
+# Mirror of stencil_left_boundary above, for the right boundary.
+Base.@propagate_inbounds Operators.stencil_right_boundary(
+    op_matrix::FDOperatorMatrix,
+    ::ValueFixingBoundaryCondition,
+    space,
+    idx,
+    hidx,
+    args...,
+) = rzero(Operators.return_eltype(op_matrix, args...))
+
+# GradientF2C/DivergenceF2C with a SetValue are the exception (as with
+# `modifies_input`): the condition fixes an *input* value, and the near-boundary output
+# still depends linearly on the adjacent interior input, so the row is the genuine
+# boundary stencil (with the fixed input's coefficient dropped) rather than zero. The
+# last of `args` is the local geometry field that `op_matrix` was given, which the
+# underlying operator's row functions do not take.
+Base.@propagate_inbounds function Operators.stencil_left_boundary(
+    op_matrix::FDOperatorMatrix{<:InputFixingFDOperator},
+    bc::Operators.SetValue,
     space,
     idx,
     hidx,
     args...,
 )
-    if op_matrix.op isa Union{Operators.GradientF2C, Operators.DivergenceF2C} &&
-       bc isa Operators.SetValue
-        args′ = args[1:(end - 1)]
-        row = op_matrix_first_row(op_matrix.op, bc, space, idx, hidx, args′...)
-        return convert(Operators.return_eltype(op_matrix, args...), row)
-    end
-    rzero(Operators.return_eltype(op_matrix, args...))
+    row = op_matrix_first_row(
+        op_matrix.op,
+        bc,
+        space,
+        idx,
+        hidx,
+        args[1:(end - 1)]...,
+    )
+    return convert(Operators.return_eltype(op_matrix, args...), row)
 end
 # Mirror of stencil_left_boundary above, for the right boundary.
 Base.@propagate_inbounds function Operators.stencil_right_boundary(
-    op_matrix::FDOperatorMatrix,
-    bc::Union{
-        Operators.SetValue,
-        Operators.SetGradient,
-        Operators.SetDivergence,
-        Operators.SetCurl,
-    },
+    op_matrix::FDOperatorMatrix{<:InputFixingFDOperator},
+    bc::Operators.SetValue,
     space,
     idx,
     hidx,
     args...,
 )
-    if op_matrix.op isa Union{Operators.GradientF2C, Operators.DivergenceF2C} &&
-       bc isa Operators.SetValue
-        args′ = args[1:(end - 1)]
-        row = op_matrix_last_row(op_matrix.op, bc, space, idx, hidx, args′...)
-        return convert(Operators.return_eltype(op_matrix, args...), row)
-    end
-    rzero(Operators.return_eltype(op_matrix, args...))
+    row = op_matrix_last_row(
+        op_matrix.op,
+        bc,
+        space,
+        idx,
+        hidx,
+        args[1:(end - 1)]...,
+    )
+    return convert(Operators.return_eltype(op_matrix, args...), row)
 end
 
 ################################################################################

src/Operators/finitedifference.jl

@@ -1660,7 +1660,16 @@ Base.@propagate_inbounds stencil_interior(
     arg,
 ) = getidx(space, arg, idx, hidx)
 
-boundary_width(::SetBoundaryOperator, ::AbstractBoundaryCondition) = 1
+# A side with no boundary condition has no boundary row: the operator is the
+# identity there. Returning 0 (rather than gating on NullBoundaryCondition in the
+# generic windowing code) keeps one-sided SetBoundaryOperators out of the boundary
+# window entirely, so they never reach the NaN-producing NullBoundaryCondition
+# stencil methods.
+boundary_width(::SetBoundaryOperator, ::AbstractBoundaryCondition) = 0
+boundary_width(
+    ::SetBoundaryOperator,
+    ::Union{SetValue, SetGradient, SetCurl, SetDivergence},
+) = 1
 Base.@propagate_inbounds function stencil_left_boundary(
     ::SetBoundaryOperator,
     bc::Union{SetValue, SetGradient, SetCurl, SetDivergence},
@@ -1949,6 +1958,10 @@ stencil_interior_width(::DivergenceF2C, arg) = ((-half, half),)
 boundary_width(::DivergenceF2C, ::AbstractBoundaryCondition) = 0
 boundary_width(::DivergenceF2C, ::SetValue) = 1
 boundary_width(::DivergenceF2C, ::SetDivergence) = 1
+# Without this, `left_interior_idx`/`right_interior_idx` place the boundary
+# centers in the interior window, so `op_matrix_first_row`/`op_matrix_last_row`
+# for `Extrapolate` are never reached and the condition is silently ignored.
+boundary_width(::DivergenceF2C, ::Extrapolate) = 1
 
 # Extend `adapt_structure` for all boundary conditions containing a `val` field.
 function Adapt.adapt_structure(to, bc::AbstractBoundaryCondition)
@@ -2209,8 +2222,6 @@ end
     Topologies.isperiodic(space) && return false
     loc = left_boundary_window(space)
     boundary_condition = Operators.get_boundary(op, loc)
-    op isa SetBoundaryOperator && boundary_condition isa NullBoundaryCondition &&
-        return false
     return idx < Operators.left_interior_idx(
         space,
         op,
@@ -2223,12 +2234,10 @@ end
     Topologies.isperiodic(space) && return false
     loc = right_boundary_window(space)
     boundary_condition = Operators.get_boundary(op, loc)
-    op isa SetBoundaryOperator && boundary_condition isa NullBoundaryCondition &&
-        return false
     return idx > Operators.right_interior_idx(
         space,
         op,
-        Operators.get_boundary(op, loc),
+        boundary_condition,
         args...,
     )
 end