Merge pull request #1881 from CliMA/ck/mat_field_opt_tests

Add JET tests to matrix field broadcasting

Author: charleskawczynski

test/MatrixFields/matrix_field_broadcasting.jl

@@ -24,7 +24,7 @@ results = Profile.fetch()
 Profile.clear()
 ProfileCanvas.html_file("flame.html", results)
 
-benchmark_getidx(bc)
+perf_getidx(bc)
 ```
 =#
 using Test

test/MatrixFields/matrix_field_test_utils.jl

@@ -316,7 +316,7 @@ end
 trunc_time(s::String) = count(',', s) > 1 ? join(split(s, ",")[1:2], ",") : s
 
 function get_getidx_args(bc)
-    space = axes(materialize(bc))
+    space = axes(bc)
     idx = if space.staggering isa Spaces.CellCenter
         10
     else
@@ -330,7 +330,8 @@ function get_getidx_args(bc)
     return (; space, bc, loc_l, loc_i, loc_r, idx, hidx)
 end
 
-function benchmark_getidx(bc)
+import JET
+function perf_getidx(bc; broken = false)
     (; space, bc, loc_l, loc_i, loc_r, idx, hidx) = get_getidx_args(bc)
     call_getidx(space, bc, loc_i, idx, hidx)
     call_getidx(space, bc, loc_l, idx, hidx)
@@ -345,6 +346,10 @@ function benchmark_getidx(bc)
     ber = time_and_units_str(
         BT.@belapsed call_getidx($space, $bc, $loc_r, $idx, $hidx)
     )
+
+    JET.@test_opt call_getidx(space, bc, loc_i, idx, hidx)
+    JET.@test_opt call_getidx(space, bc, loc_l, idx, hidx)
+    JET.@test_opt call_getidx(space, bc, loc_r, idx, hidx)
     @info "getidx times max(interior,left,right) = ($bei,$bel,$ber)"
     return nothing
 end

test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_1.jl

@@ -29,5 +29,5 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
     )
 
     test_opt && opt_test_field_broadcast(result, bc; ref_set_result!)
-    test_opt && !using_cuda && benchmark_getidx(bc)
+    test_opt && !using_cuda && perf_getidx(bc)
 end

test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_2.jl

@@ -34,5 +34,5 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
     )
 
     test_opt && opt_test_field_broadcast(result, bc; ref_set_result!)
-    test_opt && !using_cuda && benchmark_getidx(bc)
+    test_opt && !using_cuda && perf_getidx(bc)
 end

test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_3.jl

@@ -36,5 +36,5 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
     )
 
     test_opt && opt_test_field_broadcast(result, bc; ref_set_result!)
-    test_opt && !using_cuda && benchmark_getidx(bc)
+    test_opt && !using_cuda && perf_getidx(bc)
 end

test/MatrixFields/matrix_fields_broadcasting/test_non_scalar_4.jl

@@ -31,5 +31,5 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
     )
 
     test_opt && opt_test_field_broadcast(result, bc; ref_set_result!)
-    test_opt && !using_cuda && benchmark_getidx(bc)
+    test_opt && !using_cuda && perf_getidx(bc)
 end

test/MatrixFields/matrix_fields_broadcasting/test_scalar_1.jl

@@ -25,5 +25,5 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
         input_fields,
         using_cuda,
     )
-    test_opt && !using_cuda && benchmark_getidx(bc)
+    test_opt && !using_cuda && perf_getidx(bc)
 end

test/MatrixFields/matrix_fields_broadcasting/test_scalar_10.jl

@@ -59,5 +59,5 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
         ref_set_result!,
         using_cuda,
     )
-    test_opt && !using_cuda && benchmark_getidx(bc)
+    test_opt && !using_cuda && perf_getidx(bc)
 end

test/MatrixFields/matrix_fields_broadcasting/test_scalar_11.jl

@@ -55,5 +55,5 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
         ref_set_result!,
         using_cuda,
     )
-    test_opt && !using_cuda && benchmark_getidx(bc)
+    test_opt && !using_cuda && perf_getidx(bc)
 end

test/MatrixFields/matrix_fields_broadcasting/test_scalar_12.jl

@@ -56,5 +56,5 @@ test_opt = get(ENV, "BUILDKITE", "") == "true"
         ref_set_result!,
         using_cuda,
     )
-    test_opt && !using_cuda && benchmark_getidx(bc)
+    test_opt && !using_cuda && perf_getidx(bc)
 end