Try using StaticOneTo instead of Val in unrolled_reduce

CliMA · Oct 14, 2024 · e713c65 · e713c65
1 parent 754271d
commit e713c65
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -24,20 +24,20 @@ steps:
       - "echo $$JULIA_DEPOT_PATH"
 
       - echo "--- Instantiate project"
-      - "julia --project -e 'using Pkg; Pkg.add(;url=\"https://github.com/CliMA/UnrolledUtilities.jl\", rev=\"dy/recursion_limit\")'"
+      - "julia --project -e 'using Pkg; Pkg.add(;url=\"https://github.com/CliMA/UnrolledUtilities.jl\", rev=\"dy/improve_compilation\")'"
       - "julia --project -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
       - "julia --project -e 'using Pkg; Pkg.precompile()'"
       - "julia --project -e 'using Pkg; Pkg.status()'"
 
       - echo "--- Instantiate examples"
-      - "julia --project=examples -e 'using Pkg; Pkg.add(;url=\"https://github.com/CliMA/UnrolledUtilities.jl\", rev=\"dy/recursion_limit\")'"
+      - "julia --project=examples -e 'using Pkg; Pkg.add(;url=\"https://github.com/CliMA/UnrolledUtilities.jl\", rev=\"dy/improve_compilation\")'"
       - "julia --project=examples -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
       - "julia --project=examples -e 'using Pkg; Pkg.precompile()'"
       - "julia --project=examples -e 'using CUDA; CUDA.precompile_runtime()'"
       - "julia --project=examples -e 'using Pkg; Pkg.status()'"
 
       - echo "--- Instantiate perf"
-      - "julia --project=perf -e 'using Pkg; Pkg.add(;url=\"https://github.com/CliMA/UnrolledUtilities.jl\", rev=\"dy/recursion_limit\")'"
+      - "julia --project=perf -e 'using Pkg; Pkg.add(;url=\"https://github.com/CliMA/UnrolledUtilities.jl\", rev=\"dy/improve_compilation\")'"
       - "julia --project=perf -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
       - "julia --project=perf -e 'using Pkg; Pkg.precompile()'"
       - "julia --project=perf -e 'using Pkg; Pkg.status()'"

diff --git a/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl b/src/parameterized_tendencies/gravity_wave_drag/non_orographic_gravity_wave.jl
@@ -566,8 +566,9 @@ function waveforcing_column_accumulate!(
         if level >= source_level - 1
             # check break condition for each gravity waves and calculate momentum flux of breaking gravity waves at each level
             # We use the unrolled_reduce function here because it performs better for parallel execution on the GPU, avoiding type instabilities.
+            n_values = StaticOneTo(nc)
             (mask, fm) =
-                unrolled_reduce(Val(nc), (mask, FT1(0.0))) do (mask, fm), (n)
+                unrolled_reduce(n_values, (mask, FT1(0.0))) do (mask, fm), (n)
                     if (mask[n]) == true
                         c_hat = c[n] - u_kp1 # c0mu
                         # f phase speed matches the wind speed, remove c(n) from the set of propagating waves.