Avoiding distributed fill_halo_regions! in dynamics by extending the halos (#88)

simone-silvestri · web-flow · commit cc0d33ad6383 · 2025-11-06T21:49:20.000+01:00
* this should work?

* launch this

* remove this

* go like this

* make it work

* try it out

* let's go!

* Update sea_ice_time_stepping.jl

* Update sea_ice_model.jl

* Optimize halo region filling for momentum equations

Refactor halo region filling to only process local halos.

* go ahead

* add this

* go back

* Update sea_ice_model.jl

* assume

* Update runtests.jl

* Adjust halo size calculation in sea ice model

* add a test

* separate tests

* remove ext

* ensure the correct group

* fix distributed tests

* add the dynamics

* fix tests

* sea ice model

* another bugfix

* add stuff

* add Oceananigans

* fix keyword

* fix these tests

* fix stuff

* add with _halo

* fix
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -1,54 +1,150 @@
 name: CI
+
 on:
-  push:
-    branches: [main]
-    tags: ["*"]
   pull_request:
+    paths:
+      - '.github/workflows/ci.yml'
+      - 'src/**'
+      - 'test/**'
+      - 'Project.toml'
+  push:
+    branches:
+      - main
+    tags: '*'
+    paths:
+      - '.github/workflows/ci.yml'
+      - 'src/**'
+      - 'test/**'
+      - 'Project.toml'
+
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
 jobs:
-  test:
-    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+  checkpointing_tests:
+    name: Checkpointing Tests - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 120
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.10'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v5
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          TEST_GROUP: "checkpointing"
+
+  distributed:
+    name: Distributed Tests - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 120
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.10'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v5
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          TEST_GROUP: "distributed"
+
+  advection:
+    name: Sea Ice Advection - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
     runs-on: ${{ matrix.os }}
+    timeout-minutes: 120
     strategy:
       fail-fast: false
       matrix:
         version:
           - '1.10'
-          - '1.12'
         os:
           - ubuntu-latest
         arch:
           - x64
-          - x86
-        include:
-          # test macOS and Windows with latest Julia only
-          - os: macOS-latest
-            arch: x64
-            version: '1.10'
-          - os: windows-latest
-            arch: x64
-            version: '1.10'
-          - os: windows-latest
-            arch: x86
-            version: '1.10'
     steps:
-      - uses: actions/checkout@v4
-      - uses: julia-actions/setup-julia@latest
+      - uses: actions/checkout@v5
+      - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - uses: actions/cache@v3
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
         env:
-          cache-name: cache-artifacts
+          TEST_GROUP: "advection"
+
+  timestepping:
+    name: Time Stepping - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 120
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.10'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v5
+      - uses: julia-actions/setup-julia@v2
         with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-test-${{ env.cache-name }}-
-            ${{ runner.os }}-test-
-            ${{ runner.os }}-
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
-      - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v2
+        env:
+          TEST_GROUP: "timestepping"
+
+  enthalpy:
+    name: Enthalpy Model - Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 120
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.10'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v5
+      - uses: julia-actions/setup-julia@v2
         with:
-          file: lcov.info
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          TEST_GROUP: "enthalpy"
diff --git a/src/SeaIceDynamics/sea_ice_momentum_equations.jl b/src/SeaIceDynamics/sea_ice_momentum_equations.jl
@@ -59,7 +59,7 @@ function SeaIceMomentumEquation(grid;
                                 top_momentum_stress    = nothing,
                                 bottom_momentum_stress = nothing,
                                 free_drift = nothing,
-                                solver = SplitExplicitSolver(150),
+                                solver = SplitExplicitSolver(grid; substeps=150),
                                 minimum_concentration = 1e-3,
                                 minimum_mass = 1.0)
 
diff --git a/src/SeaIceDynamics/split_explicit_momentum_equations.jl b/src/SeaIceDynamics/split_explicit_momentum_equations.jl
@@ -1,25 +1,41 @@
 using Oceananigans.Grids: AbstractGrid, architecture, halo_size
-using Oceananigans.BoundaryConditions: fill_halo_regions!, fill_halo_size, fill_halo_offset
 using Oceananigans.Utils: configure_kernel
 using Oceananigans.Architectures: convert_to_device
 using Oceananigans.Fields: instantiated_location, boundary_conditions
+using Oceananigans.DistributedComputations: DistributedGrid
 using Oceananigans.ImmersedBoundaries: peripheral_node
+using Oceananigans.BoundaryConditions: fill_halo_regions!, fill_halo_size, fill_halo_offset
+
+using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces: split_explicit_kernel_size
 
-struct SplitExplicitSolver 
-    substeps :: Int
+struct SplitExplicitSolver{I, K}
+    substeps :: I
+    kernel_parameters :: K
 end
 
 """
-    SplitExplicitSolver(; substeps=120)
+    SplitExplicitSolver(grid; substeps=120)
 
 Creates a `SplitExplicitSolver` that controls the dynamical evolution of sea-ice momentum
 by subcycling `substeps` times in between each ice_thermodynamics / tracer advection time step.
 
 The default number of substeps is 120.
 """
-SplitExplicitSolver(; substeps=120) = SplitExplicitSolver(substeps)
+SplitExplicitSolver(grid; substeps=120) = SplitExplicitSolver(substeps, :xy)
+
+function SplitExplicitSolver(grid::DistributedGrid; substeps=120) 
+    Nx, Ny, _ = size(grid)
+    Hx, Hy, _ = halo_size(grid)
+    TX, TY, _ = topology(grid)
+    kernel_sizes = map(split_explicit_kernel_size, (TX, TY), (Nx, Ny), (Hx, Hy))
+    return SplitExplicitSolver(substeps, KernelParameters(kernel_sizes...))
+end
+
+# When no grid is provided, we assume a serial grid with default kernel parameters
+SplitExplicitSolver(; substeps=120) = SplitExplicitSolver(substeps, :xy)
 
 const SplitExplicitMomentumEquation = SeaIceMomentumEquation{<:SplitExplicitSolver}
+const ExtendedSplitExplicitMomentumEquation = SeaIceMomentumEquation{<:SplitExplicitSolver{<:Any, <:KernelParameters}}
 
 """
     time_step_momentum!(model, rheology::AbstractExplicitRheology, Δt)
@@ -58,10 +74,10 @@ function time_step_momentum!(model, dynamics::SplitExplicitMomentumEquation, Δt
                          ℵ = model.ice_concentration, 
                          ρ = model.ice_density))
 
-    active_cells_map = Oceananigans.Grids.get_active_column_map(grid)
+    params = dynamics.solver.kernel_parameters
 
-    u_velocity_kernel!, _ = configure_kernel(arch, grid, :xy, _u_velocity_step!; active_cells_map)
-    v_velocity_kernel!, _ = configure_kernel(arch, grid, :xy, _v_velocity_step!; active_cells_map)
+    u_velocity_kernel!, _ = configure_kernel(arch, grid, params, _u_velocity_step!)
+    v_velocity_kernel!, _ = configure_kernel(arch, grid, params, _v_velocity_step!)
 
     substeps = dynamics.solver.substeps
     initialize_rheology!(model, dynamics.rheology)
@@ -101,6 +117,9 @@ function time_step_momentum!(model, dynamics::SplitExplicitMomentumEquation, Δt
         converted_stresses_args = convert_to_device(arch, stresses_args)
 
         for substep in 1 : substeps
+            fill_halo_regions!(converted_u_halo...; only_local_halos = true)
+            fill_halo_regions!(converted_v_halo...; only_local_halos = true)
+            
             # Compute stresses! depending on the particular rheology implementation
             compute_stresses!(dynamics, converted_stresses_args...)
 
@@ -115,9 +134,6 @@ function time_step_momentum!(model, dynamics::SplitExplicitMomentumEquation, Δt
                 v_velocity_kernel!(converted_v_args...)
                 u_velocity_kernel!(converted_u_args...)
             end
-
-            fill_halo_regions!(converted_u_halo...)
-            fill_halo_regions!(converted_v_halo...)
         end
     end
 
diff --git a/src/sea_ice_model.jl b/src/sea_ice_model.jl
@@ -4,13 +4,18 @@ using Oceananigans.TimeSteppers: TimeStepper
 using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions
 using Oceananigans: tupleit, tracernames
 using Oceananigans.Forcings: model_forcing
+using Oceananigans.Grids: halo_size, topology, with_halo
+using Oceananigans.Grids: LeftConnected, RightConnected, FullyConnected
 
+using ClimaSeaIce.SeaIceDynamics: ExtendedSplitExplicitMomentumEquation
 using ClimaSeaIce.SeaIceThermodynamics: PrescribedTemperature
 using ClimaSeaIce.SeaIceThermodynamics.HeatBoundaryConditions: flux_summary
 
 @inline instantiate(T::DataType) = T()
 @inline instantiate(T) = T
 
+const ConnectedTopology = Union{LeftConnected, RightConnected, FullyConnected}
+
 struct SeaIceModel{GR, TD, D, TS, CL, U, T, IT, IC, ID, CT, STF, A, F, Arch} <: AbstractModel{TS, Arch}
     architecture :: Arch
     grid :: GR
@@ -72,8 +77,24 @@ function SeaIceModel(grid;
     boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, field_names)
 
     if isnothing(velocities)
-        u = Field{Face, Center, Nothing}(grid, boundary_conditions=boundary_conditions.u)
-        v = Field{Center, Face, Nothing}(grid, boundary_conditions=boundary_conditions.v)
+
+        # Extend the halos for the velocity fields if the dynamics is
+        # an extended split explicit momentum equation
+        if dynamics isa ExtendedSplitExplicitMomentumEquation
+            old_halos = halo_size(grid)
+            Nsubsteps = length(dynamics.solver.substeps)
+            TX, TY    = topology(grid)
+            Hx = TX() isa ConnectedTopology ? Nsubsteps + old_halos[1] : old_halos[1]
+            Hy = TY() isa ConnectedTopology ? Nsubsteps + old_halos[2] : old_halos[2]
+
+            new_halos = (Hx, Hy, old_halos[3])
+            velocity_grid = with_halo(new_halos, grid)
+        else
+            velocity_grid = grid
+        end
+
+        u = Field{Face, Center, Nothing}(velocity_grid, boundary_conditions=boundary_conditions.u)
+        v = Field{Center, Face, Nothing}(velocity_grid, boundary_conditions=boundary_conditions.v)
         velocities = (; u, v)
     end
 
diff --git a/src/sea_ice_time_stepping.jl b/src/sea_ice_time_stepping.jl
@@ -83,10 +83,9 @@ function update_state!(model::SIM)
     
     foreach(prognostic_fields(model)) do field
         mask_immersed_field_xy!(field, k=size(model.grid, 3))
+        fill_halo_regions!(field, model.clock, fields(model))
     end
 
-    fill_halo_regions!(prognostic_fields(model), model.clock, fields(model))
-
     update_model_field_time_series!(model, model.clock)
 
     return nothing
@@ -104,4 +103,4 @@ function update_model_field_time_series!(model::SeaIceModel, clock::Clock)
     end
 
     return nothing
-end
+end
diff --git a/test/distributed_tests_utils.jl b/test/distributed_tests_utils.jl
@@ -0,0 +1,57 @@
+using JLD2
+using MPI
+using Oceananigans
+using Oceananigans.DistributedComputations: reconstruct_global_field, reconstruct_global_grid
+using Oceananigans.Units
+using ClimaSeaIce
+using ClimaSeaIce.SeaIceDynamics: SeaIceMomentumEquation, ElastoViscoPlasticRheology, SemiImplicitStress
+
+# Run the distributed grid simulation and save down reconstructed results
+function run_distributed_sea_ice(arch, filename)
+    distributed_grid = RectilinearGrid(arch; 
+                                       size = (100, 100, 1), 
+                                       x = (-10kilometers, 10kilometers), 
+                                       y = (-10kilometers, 10kilometers), 
+                                       z = (-1, 0), 
+                                       halo = (5, 5, 5))
+
+    model = run_distributed_simulation(distributed_grid)
+
+    u = reconstruct_global_field(model.velocities.u)
+    v = reconstruct_global_field(model.velocities.v)
+
+    if arch.local_rank == 0
+        jldsave(filename; u = Array(interior(u, :, :, 1)),
+                          v = Array(interior(v, :, :, 1)))
+    end
+
+    MPI.Barrier(MPI.COMM_WORLD)
+    MPI.Finalize()
+
+    return nothing
+end
+
+# Just a random simulation on a rectilinear grid
+function run_distributed_simulation(grid)
+
+    τᵤ = 0.01
+    τᵥ = 0.01
+    τₒ = SemiImplicitStress()
+
+    # We use an elasto-visco-plastic rheology and WENO seventh order 
+    # for advection of h and ℵ
+    dynamics = SeaIceMomentumEquation(grid; 
+                                      top_momentum_stress = (u=τᵤ, v=τᵥ),
+                                      bottom_momentum_stress = τₒ, 
+                                      rheology = ElastoViscoPlasticRheology(),
+                                      solver = SplitExplicitSolver(grid, substeps=50))
+
+    model = SeaIceModel(grid; ice_thermodynamics = nothing, dynamics, advection = WENO(order=7))
+
+    for N in 1:100
+        time_step!(model, 1minutes)
+    end
+
+    return model
+end
+
diff --git a/test/runtests.jl b/test/runtests.jl
diff --git a/test/test_distributed_sea_ice.jl b/test/test_distributed_sea_ice.jl
