Merge pull request #1683 from CliMA/ck/compromise

Fix inference using promote_op in multiply_matrix_at_index

Author: charleskawczynski

.buildkite/pipeline.yml

@@ -630,7 +630,6 @@ steps:
       - label: "Unit: bidiag matrix row example (GPU)"
         key: gpu_compat_bidiag_matrix_row
         command: "julia --color=yes --project=test test/MatrixFields/gpu_compat_bidiag_matrix_row.jl"
-        soft_fail: true
         agents:
           slurm_gpus: 1
 

NEWS.md

@@ -4,7 +4,11 @@ ClimaCore.jl Release Notes
 main
 -------
 
+v0.13.4
+-------
+
 - ![][badge-🐛bugfix] We fixed some fieldvector broadcasting on Julia 1.9. PR [#1658](https://github.com/CliMA/ClimaCore.jl/pull/1658).
+- ![][badge-🚀performance] We fixed an inference failure with matrix field broadcasting. PR [#1683](https://github.com/CliMA/ClimaCore.jl/pull/1683).
 
 v0.13.3
 -------

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaCore"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 authors = ["CliMA Contributors <[email protected]>"]
-version = "0.13.3"
+version = "0.13.4"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/MatrixFields/matrix_multiplication.jl

@@ -270,29 +270,11 @@ function Operators.right_interior_idx(
     end
 end
 
-pick_inferred_type(
-    ::Type{Union{}},
-    ::Type{Y},
-) where {Y <: Geometry.LocalGeometry} = Y
-pick_inferred_type(
-    ::Type{X},
-    ::Type{Union{}},
-) where {X <: Geometry.LocalGeometry} = X
-pick_inferred_type(::Type{T}, ::Type{T}) where {T <: Geometry.LocalGeometry} = T
-pick_inferred_type(::Type{Union{}}, ::Type{Union{}}) =
-    error("Both LGs are not inferred")
-pick_inferred_type(::Type{X}, ::Type{Y}) where {X, Y} =
-    error("LGs do not match: X=$X, Y=$Y")
-
 function Operators.return_eltype(
     ::MultiplyColumnwiseBandMatrixField,
     matrix1,
     arg,
 )
-    # LG1 = local_geometry_type(typeof(axes(matrix1)))
-    # LG2 = local_geometry_type(typeof(axes(arg)))
-    # LG = pick_inferred_type(LG1, LG2)
-    # return Operators.return_eltype(op, matrix1, arg, LG)
     eltype(matrix1) <: BandMatrixRow || error(
         "The first argument of ⋅ must have elements of type BandMatrixRow, but \
          the given argument has elements of type $(eltype(matrix1))",
@@ -365,9 +347,8 @@ boundary_modified_ud(::BottomRightMatrixCorner, ud, column_space, i) =
 # matrix field broadcast expressions to take roughly 3 or 4 times longer to
 # evaluate, but this is less significant than the decrease in compilation time.
 function multiply_matrix_at_index(loc, space, idx, hidx, matrix1, arg, bc)
-    # lg = Geometry.LocalGeometry(space, idx, hidx)
-    # prod_type = Operators.return_eltype(⋅, matrix1, arg, typeof(lg))
-    prod_type = Operators.return_eltype(⋅, matrix1, arg)
+    lg = Geometry.LocalGeometry(space, idx, hidx)
+    prod_type = Operators.return_eltype(⋅, matrix1, arg, typeof(lg))
 
     column_space1 = column_axes(matrix1, space)
     ld1, ud1 = outer_diagonals(eltype(matrix1))

test/MatrixFields/operator_matrices.jl

@@ -302,7 +302,7 @@ end
             ᶠuvw,
             ᶜadvect(ᶠuvw, ᶜwinterp(ᶠscalar, ᶠrbias(ᶜlbias(ᶠinterp(ᶜnested))))),
         )),
-        time_ratio_limit = 20, # This case's ref function is fast on Buildkite.
+        time_ratio_limit = 25, # This case's ref function is fast on Buildkite.
         test_broken_with_cuda = true, # TODO: Fix this.
     )