PRONTOLab · glwagner · Mar 6, 2026 · Mar 6, 2026 · Mar 6, 2026 · Mar 6, 2026
diff --git a/.github/workflows/CompileOrRun.yml b/.github/workflows/CompileOrRun.yml
@@ -48,6 +48,8 @@ jobs:
   simulation:
     name: Julia ${{ inputs.julia_version }} - ${{ inputs.sim_type }} - ${{ inputs.os }} - ${{ inputs.xla_runtime }} - ${{ inputs.compile_or_run }}
     runs-on: ${{ inputs.os }}
+    env:
+      JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
     strategy:
       fail-fast: false
     steps:
@@ -131,7 +133,7 @@ jobs:
         run: |
           OK="true"
           for OP in all-to-all all-gather all-reduce; do
-              COLLECTIVES=$(find . -name 'optimised_${{ inputs.sim_type }}*.xla' -exec grep --with-filename --line-number --extended-regexp "${OP}" '{}' \;)
+              COLLECTIVES=$(find . -name 'optimised_${{ inputs.sim_type }}*.xla' -exec grep --with-filename --line-number --extended-regexp "${OP}\(" '{}' \;)
               NUM_COLLECTIVES=$(echo "${COLLECTIVES}" | wc -l | xargs)
               if [[ -n "${COLLECTIVES}" ]]; then
                   MSG="There are ${NUM_COLLECTIVES} remaining ${OP} operations"

diff --git a/Project.toml b/Project.toml
@@ -8,17 +8,19 @@ ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
 BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
-ClimaOcean = "0376089a-ecfe-4b0e-a64f-9c555d74d754"
+ClimaSeaIce = "6ba0ff68-24e6-4315-936c-2e99227c95a4"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 HDF5_jll = "0234f1f7-429e-5d53-9886-15a909be8d59"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+NumericalEarth = "904d977b-046a-4731-8b86-9235c0d1ef02"
 Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reactant = "3c362404-f566-11ee-1572-e11a4b42c853"
+ReactantCore = "a3311ec8-5e00-46d5-b541-4f83e724a433"
 Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
 SeawaterPolynomials = "d496a93d-167e-4197-9f49-d3af4ff8fe40"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
@@ -27,18 +29,28 @@ Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 ArgParse = "1.2.0"
 BFloat16s = "0.6.1"
 CUDA = "5.7.1"
-ClimaOcean = "0.5.10"
+ClimaSeaIce = "=0.4.5"
 Dates = "1"
 HDF5_jll = "=1.14.3"
 KernelAbstractions = "0.9.34"
 MPI = "0.20.22"
-Oceananigans = "=0.96.26"
+NumericalEarth = "0.2.2"
+Oceananigans = "0.106.1"
 PrecompileTools = "1.2.1"
 Printf = "1"
 Profile = "1"
+ReactantCore = "0.1.18"
 Reactant = "0.2.247"
 SeawaterPolynomials = "0.3.9"
 Serialization = "1"
 
+[sources.Oceananigans]
+url = "https://github.com/CliMA/Oceananigans.jl.git"
+rev = "pb/step-fix"
+
+[sources.NumericalEarth]
+url = "https://github.com/NumericalEarth/NumericalEarth.jl.git"
+rev = "glw/manual-iterate"
+
 [preferences.Reactant]
 xla_runtime = "IFRT"
diff --git a/correctness/correctness_baroclinic_instability_simulation_run.jl b/correctness/correctness_baroclinic_instability_simulation_run.jl
@@ -9,6 +9,7 @@ const parsed_args = GordonBell25.parse_baroclinic_instability_args(;
 
 default_float_type = GordonBell25.float_type_from_args(parsed_args)
 Oceananigans.defaults.FloatType = default_float_type
+using CUDA
 using Reactant
 
 throw_error = true
@@ -40,6 +41,8 @@ vmodel = GordonBell25.baroclinic_instability_model(varch, Nx, Ny, Nz; model_kw..
 ui = 1e-3 .* rand(size(vmodel.velocities.u)...)
 vi = 1e-3 .* rand(size(vmodel.velocities.v)...)
 set!(vmodel, u=ui, v=vi)
+GordonBell25.zero_tendencies!(rmodel)
+GordonBell25.zero_tendencies!(vmodel)
 GordonBell25.sync_states!(rmodel, vmodel)
 
 @info "At the beginning:"
@@ -57,6 +60,9 @@ Oceananigans.TimeSteppers.update_state!(vmodel)
 GordonBell25.compare_states(rmodel, vmodel; include_halos, throw_error, rtol, atol)
 
 GordonBell25.sync_states!(rmodel, vmodel)
+GordonBell25.zero_tendencies!(rmodel)
+GordonBell25.zero_tendencies!(vmodel)
+# compile_options = CompileOptions(; sync=true, raise=true, strip_llvm_debuginfo=true, strip=:all)
 rfirst! = @compile compile_options=compile_options GordonBell25.first_time_step!(rmodel)
 @showtime rfirst!(rmodel)
 @showtime GordonBell25.first_time_step!(vmodel)

diff --git a/correctness/correctness_sharded_baroclinic_instability_simulation_run.jl b/correctness/correctness_sharded_baroclinic_instability_simulation_run.jl
@@ -1,5 +1,6 @@
 using GordonBell25
 using Oceananigans
+using Reactant
 
 const parsed_args = GordonBell25.parse_baroclinic_instability_args(;
     grid_x_default = 64,
@@ -9,7 +10,6 @@ const parsed_args = GordonBell25.parse_baroclinic_instability_args(;
 
 default_float_type = GordonBell25.float_type_from_args(parsed_args)
 Oceananigans.defaults.FloatType = default_float_type
-using Reactant
 
 if !GordonBell25.is_distributed_env_present()
     using MPI
@@ -64,6 +64,8 @@ end
 ui = 1e-3 .* rand(size(vmodel.velocities.u)...)
 vi = 1e-3 .* rand(size(vmodel.velocities.v)...)
 set!(vmodel, u=ui, v=vi)
+GordonBell25.zero_tendencies!(rmodel)
+GordonBell25.zero_tendencies!(vmodel)
 GordonBell25.sync_states!(rmodel, vmodel)
 
 compile_options = CompileOptions(; sync=true, raise=true, strip_llvm_debuginfo=true, strip=["enzymexla.kernel_call", "(::Reactant.Compiler.LLVMFunc", "ka_with_reactant", "(::KernelAbstractions.Kernel", "var\"#_launch!;_launch!"], multifloat=GordonBell25.multifloat_from_args(parsed_args))
@@ -81,6 +83,8 @@ Oceananigans.TimeSteppers.update_state!(vmodel)
 GordonBell25.compare_states(rmodel, vmodel; include_halos, throw_error, rtol, atol)
 
 GordonBell25.sync_states!(rmodel, vmodel)
+GordonBell25.zero_tendencies!(rmodel)
+GordonBell25.zero_tendencies!(vmodel)
 rfirst! = @compile compile_options=compile_options GordonBell25.first_time_step!(rmodel)
 @showtime rfirst!(rmodel)
 @showtime GordonBell25.first_time_step!(vmodel)

diff --git a/ext/PrecompileAB2StepF32.jl b/ext/PrecompileAB2StepF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin

diff --git a/ext/PrecompileCachePreviousTendenciesF32.jl b/ext/PrecompileCachePreviousTendenciesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
@@ -39,8 +38,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels:
 
     @compile_workload begin
         model = GordonBell25.baroclinic_instability_model(arch; resolution=4, Δt=60, Nz=4, grid_type=:simple_lat_lon)
-        Δt = ConcreteRNumber(60.0)
-        compiled! = @compile GordonBell25.correct_velocities_and_cache_previous_tendencies_workload!(model, Δt)
+        compiled! = @compile GordonBell25.cache_previous_tendencies_workload!(model)
     end
 
     # Reset float type

diff --git a/ext/PrecompileComputeAuxiliariesF32.jl b/ext/PrecompileComputeAuxiliariesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin

diff --git a/ext/PrecompileComputeBoundaryTendenciesF32.jl b/ext/PrecompileComputeBoundaryTendenciesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin
@@ -39,7 +38,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels:
 
     @compile_workload begin
         model = GordonBell25.baroclinic_instability_model(arch; resolution=4, Δt=60, Nz=4, grid_type=:simple_lat_lon)
-        compiled! = @compile GordonBell25.compute_boundary_tendencies_workload!(model)
+        compiled! = @compile GordonBell25.compute_interior_tracer_tendencies_workload!(model)
     end
 
     # Reset float type

diff --git a/ext/PrecompileComputeInteriorMomentumTendenciesF32.jl b/ext/PrecompileComputeInteriorMomentumTendenciesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin

diff --git a/ext/PrecompileComputeInteriorTracerTendenciesF32.jl b/ext/PrecompileComputeInteriorTracerTendenciesF32.jl
@@ -11,24 +11,23 @@ using Oceananigans: fields, prognostic_fields
 using Oceananigans.BoundaryConditions:
     fill_halo_regions!
 
-using Oceananigans.Fields:
-    tupled_fill_halo_regions!
-
 using Oceananigans.Models.HydrostaticFreeSurfaceModels:
     mask_immersed_model_fields!,
-    compute_auxiliaries!,
-    compute_tendencies!
+    compute_closure_fields!,
+    compute_momentum_tendencies!,
+    compute_tracer_tendencies!
 
 # For reference
 
-    * mask_immersed_model_fields!(model, grid)
-    * tupled_fill_halo_regions!(prognostic_fields(model), grid, model.clock, fields(model))
-    * compute_auxiliaries!(model)
-    * fill_halo_regions!(model.diffusivity_fields; only_local_halos=true)
-    * compute_tendencies!(model, callbacks)
+    * mask_immersed_model_fields!(model)
+    * fill_halo_regions!((u, v), model.clock, fields(model))
+    * fill_halo_regions!(tracers, model.clock, fields(model))
+    * compute_closure_fields!(model.closure_fields, model.closure, model, ...)
+    * fill_halo_regions!(model.closure_fields; only_local_halos=true)
+    * compute_momentum_tendencies!(model, callbacks)
     * ab2_step!(model, Δt)
-    * tupled_fill_halo_regions!(prognostic_fields(model), model.grid, model.clock, fields(model))
-    * correct_velocities_and_cache_previous_tendencies!(model, Δt)
+    * fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
+    * cache_previous_tendencies!(model)
 =#
 
 @setup_workload begin