Try #1436:

.buildkite/pipeline.yml

@@ -557,6 +557,16 @@ steps:
           slurm_gpus: 1
           slurm_mem: 40GB
 
+      - label: "Unit: field matrices (CPU)"
+        key: unit_field_matrices_cpu
+        command: "julia --color=yes --check-bounds=yes --project=test test/MatrixFields/field_matrices.jl"
+
+      - label: "Unit: field matrices (GPU)"
+        key: unit_field_matrices_gpu
+        command: "julia --color=yes --project=test test/MatrixFields/field_matrices.jl"
+        agents:
+          slurm_gpus: 1
+
   - group: "Unit: Hypsography"
     steps:
 

docs/src/matrix_fields.md

@@ -26,6 +26,26 @@ MultiplyColumnwiseBandMatrixField
 operator_matrix
 ```
 
+## Matrices of Fields
+
+```@docs
+FieldName
+@name
+FieldNameDict
+```
+
+# Linear Solvers
+
+```@docs
+FieldMatrixSolverAlgorithm
+FieldMatrixSolver
+field_matrix_solve!
+BlockDiagonalSolve
+BlockLowerTriangularSolve
+SchurComplementSolve
+ApproximateFactorizationSolve
+```
+
 ## Internals
 
 ```@docs
@@ -39,6 +59,9 @@ matrix_shape
 column_axes
 AbstractLazyOperator
 replace_lazy_operator
+FieldNameTree
+FieldNameSet
+field_vector_view
 ```
 
 ## Utilities

src/MatrixFields/MatrixFields.jl

@@ -27,38 +27,36 @@ for them:
 """
 module MatrixFields
 
-import CUDA: @allowscalar
-import LinearAlgebra: UniformScaling, Adjoint, AdjointAbsVec
+import CUDA
+import LinearAlgebra: I, UniformScaling, Adjoint, AdjointAbsVec, mul!, inv
 import StaticArrays: SMatrix, SVector
 import BandedMatrices: BandedMatrix, band, _BandedMatrix
+import ClimaComms
 
 import ..Utilities: PlusHalf, half
 import ..RecursiveApply:
     rmap, rmaptype, rpromote_type, rzero, rconvert, radd, rsub, rmul, rdiv
+import ..RecursiveApply: ⊠, ⊞, ⊟
 import ..DataLayouts: AbstractData
 import ..Geometry
 import ..Spaces
 import ..Fields
 import ..Operators
 
-export ⋅
 export DiagonalMatrixRow,
     BidiagonalMatrixRow,
     TridiagonalMatrixRow,
     QuaddiagonalMatrixRow,
     PentadiagonalMatrixRow
+export FieldVectorKeys, FieldVectorView, FieldVectorViewBroadcasted
+export FieldMatrixKeys, FieldMatrix, FieldMatrixBroadcasted
+export @name, ⋅, FieldMatrixSolver, field_matrix_solve!
 
 # Types that are teated as single values when using matrix fields.
 const SingleValue =
     Union{Number, Geometry.AxisTensor, Geometry.AdjointAxisTensor}
 
 include("band_matrix_row.jl")
-include("rmul_with_projection.jl")
-include("matrix_shape.jl")
-include("matrix_multiplication.jl")
-include("lazy_operators.jl")
-include("operator_matrices.jl")
-include("field2arrays.jl")
 
 const ColumnwiseBandMatrixField{V, S} = Fields.Field{
     V,
@@ -71,6 +69,19 @@ const ColumnwiseBandMatrixField{V, S} = Fields.Field{
     },
 }
 
+include("rmul_with_projection.jl")
+include("matrix_shape.jl")
+include("matrix_multiplication.jl")
+include("lazy_operators.jl")
+include("operator_matrices.jl")
+include("field2arrays.jl")
+include("unrolled_functions.jl")
+include("field_name.jl")
+include("field_name_set.jl")
+include("field_name_dict.jl")
+include("field_matrix_solver.jl")
+include("single_field_solver.jl")
+
 function Base.show(io::IO, field::ColumnwiseBandMatrixField)
     print(io, eltype(field), "-valued Field")
     if eltype(eltype(field)) <: Number
@@ -82,7 +93,10 @@ function Base.show(io::IO, field::ColumnwiseBandMatrixField)
         end
         column_field = Fields.column(field, 1, 1, 1)
         io = IOContext(io, :compact => true, :limit => true)
-        @allowscalar Base.print_array(io, column_field2array_view(column_field))
+        CUDA.@allowscalar Base.print_array(
+            io,
+            column_field2array_view(column_field),
+        )
     else
         # When a BandedMatrix with non-number entries is printed, it currently
         # either prints in an illegible format (e.g., if it has AxisTensor or

src/MatrixFields/band_matrix_row.jl

@@ -135,3 +135,9 @@ Base.:*(value::SingleValue, row::BandMatrixRow) =
 
 Base.:/(row::BandMatrixRow, value::Number) =
     map(entry -> rdiv(entry, value), row)
+
+inv(row::DiagonalMatrixRow) = DiagonalMatrixRow(inv(row[0]))
+inv(::BandMatrixRow{ld, bw}) where {ld, bw} = error(
+    "The inverse of a matrix with $bw diagonals is (usually) a dense matrix, \
+     so it cannot be represented using BandMatrixRows",
+)

src/MatrixFields/field2arrays.jl

@@ -53,11 +53,11 @@ function column_field2array(field::Fields.FiniteDifferenceField)
             last_row = matrix_d < n_cols - n_rows ? n_rows : n_cols - matrix_d
 
             diagonal_data_view = view(diagonal_data, first_row:last_row)
-            @allowscalar copyto!(matrix_diagonal, diagonal_data_view)
+            CUDA.@allowscalar copyto!(matrix_diagonal, diagonal_data_view)
         end
         return matrix
     else # field represents a vector
-        return @allowscalar Array(column_field2array_view(field))
+        return CUDA.@allowscalar Array(column_field2array_view(field))
     end
 end