Fix dss inference

CliMA · Jun 16, 2023 · ba5c556 · ba5c556
1 parent af8b6c7
commit ba5c556
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/src/Spaces/dss.jl b/src/Spaces/dss.jl
@@ -836,7 +836,7 @@ function dss_local_vertices!(
             sum_data = mapreduce(
                 ⊞,
                 vertex;
-                init = RecursiveApply.rmap(zero, slab(perimeter_data, 1, 1)[1]),
+                init = RecursiveApply.rzero(eltype(slab(perimeter_data, 1, 1))),
             ) do (lidx, vert)
                 ip = Topologies.perimeter_vertex_node_index(vert)
                 perimeter_slab = slab(perimeter_data, level, lidx)
@@ -906,9 +906,8 @@ function dss_local_ghost!(
                 sum_data = mapreduce(
                     ⊞,
                     vertex;
-                    init = RecursiveApply.rmap(
-                        zero,
-                        slab(perimeter_data, 1, 1)[1],
+                    init = RecursiveApply.rzero(
+                        eltype(slab(perimeter_data, 1, 1)),
                     ),
                 ) do (isghost, idx, vert)
                     ip = Topologies.perimeter_vertex_node_index(vert)

diff --git a/test/Operators/spectralelement/benchmark_ops.jl b/test/Operators/spectralelement/benchmark_ops.jl
@@ -138,14 +138,14 @@ using JET
         p = @allocated kernel_complicated_field_dss!(kernel_args)
         @test p == 0
         p = @allocated kernel_complicated_field2_dss!(kernel_args)
-        @test_broken p == 0
+        @test p == 0
         # Inference tests
         JET.@test_opt kernel_scalar_dss!(kernel_args)
         JET.@test_opt kernel_vector_dss!(kernel_args)
         JET.@test_opt kernel_field_dss!(kernel_args)
         JET.@test_opt kernel_ntuple_field_dss!(kernel_args)
         JET.@test_opt kernel_ntuple_floats_dss!(kernel_args)
         JET.@test_opt kernel_complicated_field_dss!(kernel_args)
-        # JET.@test_opt kernel_complicated_field2_dss!(kernel_args) # fails
+        JET.@test_opt kernel_complicated_field2_dss!(kernel_args)
     end
 end