Sharded static data and domain integrals (v0.2.1)#

Added in version v0.2.1: Per-device materialisation of StaticArray(replication="shard") under ShardedStencilNode, and the domain_integral_fields() API for cross-device reduction of partial-sum outputs.

This page explains how to write a stencil node whose non-evolving arrays are sharded alongside state — needed when the array is big enough that replicating it per device would defeat the point of sharding — and how to declare outputs that are domain integrals over the lattice so the wrapper all-reduces them.

The motivating consumer is MIME’s IBLBMFluidNode: a D3Q19 LBM node carrying a (nx, ny, nz) pipe-wall mask and a (19, nx, ny, nz) “missing-link” bounce-back mask, both far too large to replicate on every device when the simulation is sharded across a pencil mesh.

The picture#

        flowchart LR
    subgraph "ShardedStencilNode.update()"
        A[state pytree<br/>and StaticArray] --> B[device_put<br/>NamedSharding per shard_axis]
        B --> C[shard_map]
        C --> D[halo_exchange<br/>state slab + static slab]
        D --> E[inner.update_padded<br/>static_padded=...<br/>shard_info=...]
        E --> F[strip halos<br/>or psum integrals]
        F --> G[next state]
    end

Three things happen per step:

  1. The wrapper takes each StaticArray(replication="shard") on the inner node and device_puts it with a NamedSharding whose PartitionSpec puts the matching mesh axis on the array’s shard_axis. This happens once, not per step: the placement is cached and re-checked against the inner node’s static_data on every call, so a node that hands back a different array object (a mesh rebuilt after a parameter write, a static_data_provider reconstruction, a replace_node) is picked up on the next update. What the check cannot see is a static rewritten in place. GraphManager.compile() drops the cached placement on every node (and on every node they wrap), so the usual recovery is a recompile; outside a graph, call wrapper.invalidate_static_cache() after such a rewrite. Note that a rewrite with no recompile is not picked up by a later step() on its own: the graph’s static-data check hashes shape and dtype, never contents.

  2. Inside shard_map, each device’s slab is halo-exchanged along the matching spatial axis. Between two shards the halo holds the neighbour’s cells. At the edges of the global grid it follows the wrapper’s boundary when that is "periodic": the cells beyond a periodic edge are the opposite edge’s cells, for a static array as much as for the state. Under "edge" and "zero" there is no cell beyond the edge, and the static’s halo repeats the edge cell. (The halo fill is a question about space, not time; until 0.4.0 statics were edge-filled under every mode “because they don’t evolve”, and a periodic node that read a static in its halo stepped about 3% off the unsharded node, with no error.)

  3. The inner node’s update_padded receives the padded slab via static_padded[<key>] and uses it like any other padded array.

Declaring a sharded static array#

The declaration lives on the node:

from maddening.core.node import SimulationNode
from maddening.core.static_data import StaticArray

class WallBouncebackLBM(SimulationNode):

    def __init__(self, name, timestep, *, nx, ny, nz, pipe_radius):
        super().__init__(name, timestep,
                          nx=nx, ny=ny, nz=nz, pipe_radius=pipe_radius)
        self._pipe_wall = _build_pipe_mask(nx, ny, nz, pipe_radius)

    def halo_width(self):
        # D3Q19: read ±1 neighbour on every spatial axis
        return {0: 1, 1: 1, 2: 1}

    @property
    def static_data(self):
        return {
            "pipe_wall": StaticArray(
                self._pipe_wall,
                replication="shard",
                shard_axis=0,           # shard along the x axis
            ),
        }

shard_axis is the array’s own axis — it must match one of the spatial axes the wrapping ShardedStencilNode actually shards (shard_axis ∈ axis_map.values()) and the node must declare a non-zero halo_width() entry on that axis (otherwise there’s no neighbour slab to exchange with). Both invariants are checked at ShardedStencilNode.__init__ time:

sharded = ShardedStencilNode(
    WallBouncebackLBM("lbm", 0.001, nx=128, ny=64, nz=64, pipe_radius=20),
    mesh=mesh,
    axis_map={"spatial_x": 0},   # shard axis 0 on mesh "spatial_x"
    boundary="periodic",
)
# Raises ValueError immediately if pipe_wall's shard_axis isn't in
# axis_map.values() (i.e. it's not on the sharded spatial axis), or
# if halo_width() has no entry for shard_axis.

Replicate-mode StaticArray entries are not affected — they stay closure-captured in full on every device, exactly as in v0.2.0.

Reading the sharded slab in update_padded#

The inner node’s update_padded gets a new keyword-only argument:

def update_padded(
    self, state_padded, boundary_inputs, dt,
    *, static_padded=None, shard_info=None,
) -> dict:
    """Receives a per-shard slab of pipe_wall via `static_padded`."""
    f_pad      = state_padded["f"]               # (nx_local+2, ny, nz, 19)
    wall_pad   = static_padded["pipe_wall"]      # (nx_local+2, ny, nz)
    # interior view (halo stripped) is wall_pad[1:-1, :, :]
    # cells outside the interior come from the neighbour shard's
    # boundary cells via halo_exchange
    ...

The keyword is None when there are no sharded static_data entries on the node, or when the node is run outside of ShardedStencilNode (e.g. single-device path). Default to the closure-captured full array in that case:

def update_padded(
    self, state_padded, boundary_inputs, dt,
    *, static_padded=None, shard_info=None,
):
    if static_padded is not None and "pipe_wall" in static_padded:
        wall = static_padded["pipe_wall"]
    else:
        # Unsharded path; the full wall is closure-captured on self.
        # Pad it the way the wrapper would: this node is wrapped with
        # boundary="periodic", so the halo beyond each x face is the
        # opposite face ("wrap"); a node wrapped with "edge" or "zero"
        # would use mode="edge".
        wall = jnp.pad(self._pipe_wall, [(1, 1), (0, 0), (0, 0)], mode="wrap")
    ...

Declaring domain-integral outputs#

A node that computes a jnp.sum-over-lattice output (a drag force, a total mass, a heat-flux integral) needs cross-device reduction under sharding. Declare which output keys are integrals:

def domain_integral_fields(self) -> set[str]:
    return {"drag_force", "drag_torque"}

and have update_padded compute the partial sum on the local slab — no psum in the node:

def update_padded(self, state_padded, boundary_inputs, dt,
                  *, static_padded=None, shard_info=None):
    ...
    force_partial = jnp.sum(
        per_cell_force * wall_pad[1:-1, :, :, None],
        axis=(0, 1, 2),
    )
    return {
        "f": f_new,
        "drag_force": force_partial,    # (3,) — partial sum on this shard
    }

ShardedStencilNode applies lax.psum(value, axis_name=tuple(mesh.axis_names)) to every key listed in domain_integral_fields() after update_padded returns, and sets the corresponding out_spec to P() (fully replicated post-psum). Returned integral values are floating-point — psum on integer types risks wrap on large meshes.

Important

Every key returned from update_padded must be either in state_fields() or in domain_integral_fields(). The wrapper defensively raises ValueError at trace time on an unknown key — it has no way to infer the partition spec for an unclassified output.

shard_info: when you’d rather recompute the mask#

Not every static array needs to be materialised. For analytic masks (a sphere, a cylinder, a coordinate range), it’s often cheaper to recompute per shard than to ship the full mask through jax.device_put once and halo_exchange it every step.

ShardedStencilNode populates a shard_info dict for the inner node:

def update_padded(self, state_padded, boundary_inputs, dt,
                  *, static_padded=None, shard_info=None):
    # shard_info = {0: (global_offset, local_extent)}
    #
    # global_offset is a TRACED jax scalar:
    #   lax.axis_index(mesh_axis) * local_extent
    # local_extent is a Python int: the grid's extent along the axis
    # over the devices on its mesh axis.
    if shard_info is not None and 0 in shard_info:
        offset, extent = shard_info[0]
        # offset usable in dynamic_slice; NOT in Python int slicing
        global_x = offset + jnp.arange(extent + 2)    # +2 for halos
        ...

shard_info is None when the node is run outside of ShardedStencilNode.

local_extent is worked out before the step is traced, from a grid field’s global shape: a state field that is not a domain integral, or a sharded StaticArray split along the axis when no state field has it. A domain integral carried in the state never sets it. Before 0.4.0 the first local field in sorted key order did, integral or not, and a HeatNode carrying a vector energy integral named to sort before temperature was told the wrong block size and never closed its right rod end (MADD-ANO-056). A node that reads shard_info whose grid fields are split into blocks of different sizes along one axis is refused: one (offset, extent) pair cannot describe both.

ShardedUnstructuredNode passes a different shard_info, for a different layout: {0: (offset, n_local_max), "n_local": n_owned}. Each shard’s block of owned cells is padded to n_local_max, the largest shard’s count, so a shard that owns fewer cells has padding rows at the end of its block; n_owned (a traced int32 scalar, added in 0.4.0) is its own count. A domain integral on that path masks the padding out:

_, n_local_max = shard_info[0]
owned = jnp.arange(n_local_max) < shard_info["n_local"]
total_partial = jnp.sum(jnp.where(owned, x_new[:n_local_max], 0.0))

Before 0.4.0 nothing told the node where the padding started, and a node that summed its block summed the padding into the integral (on a 7-cell, 2-device partition, 36 where the answer is 35).

Sharding policy is part of the JIT cache key#

SimulationNode.static_data_hash() already incorporates (key, shape, dtype, replication, shard_axis). Switching a StaticArray from replication="replicate" to replication="shard" between graph compiles invalidates the JIT trace — no extra machinery needed in user code. The shard_map cache inside ShardedStencilNode also keys on static_data_hash(), so identity changes (a user replacing self._mask between steps) invalidate cleanly even when shape and dtype are unchanged.

A static copied along a mesh axis, inside a loop (MADD-ANO-068)#

On jaxlib 0.10.2 to 0.11.2, XLA compiles one kernel pattern wrongly when the step is traced inside a loop. The step reads a sharded static in its halo, reads another array through a window at a shard_info offset, and the static is copied along a mesh axis of two or more devices. A mesh axis the axis_map leaves unused copies it, and so does the other axis of a pencil. The result is silently wrong on the first kind of axis and fails to compile on the second. jit alone is right. GraphManager therefore refuses any path that loops over such a step:

  • run_scan, run_scan_with_history, run_sweep, run_adaptive_scan and sysid.windowed_loss, gradients through them included;

  • every entry point, step() included, when the node is in a coupling group, because the group’s iteration loops over the update inside the step.

The RuntimeError names the node, the static and the mesh axis.

The decision comes from your update_padded itself. The graph traces it once, never compiled, and walks the jaxpr. Each of the following keeps a node out of the refusal:

  • read the static only in its interior, with a slice (or a literal dynamic_slice) inside [halo, n - halo) along its shard axis;

  • don’t window another array at a shard offset. A window into the static itself is fine: taking the static’s own columns at shard_info[1] is how the multi-device validation goals’ 2-D field reads a full-width static on a pencil;

  • build the mesh from the axes the node is sharded along only.

A loop you write yourself around a wrapped node, such as a lax.scan over its update, is not guarded. Check it against step().

What’s still TODO#

  • ~~Partial-axis psum.~~ Done in v0.4.0: override domain_integral_axes() to map an integral key to the mesh axes it is reduced over ({"drag": ("spatial_z",)}); the result keeps one leading dimension per unreduced mesh axis, sharded along it, and an empty tuple stacks the per-shard partial values without reducing. A key without an entry is still psum-med over the full mesh. A mesh axis that splits nothing (one a stencil wrapper’s axis_map leaves unused) is neither reduced nor stacked over: every device along it holds the same block. Both ShardedStencilNode and ShardedUnstructuredNode honour it (tests/cloud/multigpu/test_partial_axis_integrals.py).

  • Sharded static arrays without halos. v0.2.1 rejects a sharded StaticArray whose shard_axis doesn’t appear in node.halo_width(). A future relaxation could allow such arrays by skipping the halo exchange, but the semantics get subtle near slab boundaries.

See also#