zoomy_jax.fvm.halo_exchange_jax module

zoomy_jax.fvm.halo_exchange_jax module#

JAX-native halo exchange for SPMD-sharded FV solvers.

Pattern copied from JAX’s cloud_tpu_colabs/Wave_Equation.ipynb and used in production by JAX-Fluids / Autodesk XLB. No MPI dependency; just jax.lax.ppermute along a named axis inside shard_map.

Storage layout#

Each device holds a padded local Q of shape (n_var, n_local + 2*halo):

[ halo_left | n_local owned cells | halo_right ]

|← halo →| |← halo →|

halo_exchange_inplace pulls the owned edge slab of width halo from each side, ``ppermute``s it to the neighbor, and writes it into the neighbor’s halo slab on the opposite side. Devices at the domain boundary receive zeros in their outer halo (free boundary); the inline BC evaluator on the solver side overwrites those slots with the BC-evaluated face value, so the same kernel runs everywhere — no Python branching inside the JIT trace.

The cell axis is always the last axis of Q ((n_var, n_cells)).

Composition with the existing solver#

Already wired (see tests/unit/zoomy_jax/):

  • test_halo_exchange.py — bit-correct halo on 2/4 devices.

  • test_spmd_advection.py — bit-identity scalar advection on {2,4} devices × {16,32} cells vs single-device.

  • test_partition_jax.py — partition_1d_contiguous chops a global MeshJAX into per-partition padded slabs with shifted face indices.

  • test_spmd_solver_integration.py — existing ConstantReconstruction composes with SPMD shard_map when handed a per-partition mesh; bit-identity vs single-device.

Open work to lift this into the full HyperbolicSolver:

  1. Rebuild LSQ stencils per partition (LSQMUSCLReconstructionJAX consumes mesh.lsq_gradQ / mesh.lsq_neighbors / mesh.lsq_boundary_face_neighbors — partition_jax leaves these empty. Re-run LSQMesh._build_lsq_stencil on the padded coordinate system of each partition.)

  2. Drop the periodic-wrap demo path and use the inline BC kernel at global boundaries (rank-0 left face + rank-(N-1) right face). The partition’s boundary_face_* arrays already carry exactly these faces; the existing flux operator’s BC fori_loop runs over them per shard.

  3. Wrap HyperbolicSolver.step in shard_map(... in_specs= P(None, "cells"), out_specs=P(None, "cells")) with the halo exchange called inside the body before the flux operator.

  4. For multi-process deployment: jax.distributed.initialize() at process startup (SLURM auto-detection). Dev loop runs in one process via XLA_FLAGS=--xla_force_host_platform_device_count=N.

zoomy_jax.fvm.halo_exchange_jax.halo_exchange_inplace(Q_pad, halo, axis_name, n_devices)#

[Surrounding layout] Refill the halo slabs of Q_pad with neighbor data via lax.ppermute.

Q_pad layout: [halo_left | n_local owned | halo_right].

Used by the demo SPMD advection / partition tests. For solver- side LSQ-MUSCL composition (where the existing reconstruction iterates jnp.arange(n_inner_cells) from index 0), use halo_exchange_owned_first() instead — its layout matches LSQMesh’s convention (ghosts at the end).

zoomy_jax.fvm.halo_exchange_jax.halo_exchange_owned_first(Q_pad, n_local, halo, axis_name, n_devices)#

[Owned-first layout, matches LSQMesh convention] Refill the halo slabs of Q_pad via lax.ppermute.

Layout:

[  owned (n_local)  |  left_halo (halo)  |  right_halo (halo)  ]
   indices 0..n_local-1   n_local..n_local+halo-1   ...

The owned region is at the front so the existing per-cell reconstruction (jnp.arange(n_inner_cells) over owned cells) works directly. Halo cells live at the end of the array, mirroring how LSQMesh stores BC ghost cells at indices [n_inner_cells .. n_cells).

Send pattern:

  • MY leftmost slab [0:halo] → fills LEFT neighbor’s right_halo.

  • MY rightmost slab [n_local-halo:n_local] → fills RIGHT neighbor’s left_halo.

Receive (after ppermute):

  • MY left_halo = LEFT neighbor’s rightmost owned slab.

  • MY right_halo = RIGHT neighbor’s leftmost owned slab.

At the global domain boundary the corresponding halo slab contains zeros — the inline BC kernel on the solver side overwrites them.

Contents