Fix edge case, apply formatter

src/MatrixFields/field_matrix_solver.jl

@@ -247,6 +247,17 @@ function check_field_matrix_solver(::BlockDiagonalSolve, _, A, b)
     end
 end
 
+# TODO: we can remove the uniform_vertical_levels
+# limitation while still using static shared memory
+# once Nv is in the type space.
+function uniform_vertical_levels(x, names)
+    _, _, _, Nv1, _ = size(Fields.field_values(x[first(names)]))
+    return all(Base.tail(names)) do name
+        _, _, _, Nv, _ = size(Fields.field_values(x[name]))
+        Nv == Nv1
+    end
+end
+
 NVTX.@annotate function run_field_matrix_solver!(
     ::BlockDiagonalSolve,
     cache,
@@ -256,7 +267,8 @@ NVTX.@annotate function run_field_matrix_solver!(
 )
     names = matrix_row_keys(keys(A))
     if length(names) == 1 ||
-       all(name -> A[name, name] isa UniformScaling, names.values)
+       all(name -> A[name, name] isa UniformScaling, names.values) ||
+       !uniform_vertical_levels(x, names.values)
         foreach(names) do name
             single_field_solve!(cache[name], x[name], A[name, name], b[name])
         end

test/MatrixFields/field_matrix_solvers.jl

@@ -58,8 +58,10 @@ function test_field_matrix_solver(; test_name, alg, A, b, use_rel_error = false)
             AnyFrameModule(MatrixFields.KrylovKit),
             AnyFrameModule(Base.CoreLogging),
         )
-        using_cuda || @test_opt ignored_modules = ignored FieldMatrixSolver(alg, A, b)
-        using_cuda || @test_opt ignored_modules = ignored field_matrix_solve!(args...)
+        using_cuda ||
+            @test_opt ignored_modules = ignored FieldMatrixSolver(alg, A, b)
+        using_cuda ||
+            @test_opt ignored_modules = ignored field_matrix_solve!(args...)
         @test_opt ignored_modules = ignored field_matrix_mul!(b, A, x)
 
         using_cuda || @test @allocated(field_matrix_solve!(args...)) == 0