zoomy_jax.symsolver.parallel

`zoomy_jax.symsolver.parallel`#

Distributed runs on JAX: the halo exchange and the global reductions of the runtime’s name table, one MPI rank per device. Used by zoomy_jax.symsolver.runtime when the solver gets comm=; the exchanges are mpi4jax collectives, differentiable in both directions.

Module Contents#

Classes#

no_halo

Context manager inside which halo is the identity, giving the rank’s own block of an operator (block-Jacobi preconditioner).

Functions#

`bind_device`	Bind this rank to one GPU of its node (round robin), before JAX starts if possible. No-op on the CPU, with one GPU, or when `CUDA_VISIBLE_DEVICES` is set.
`krylov`	The owned-cell mask `(nc,)` and the global sum (allreduce) a distributed Krylov solve needs.
`blocks`	The `halo` exchange and the global reductions of the runtime’s name table for this rank’s piece, each face and cell counted once.

API#

zoomy_jax.symsolver.parallel.bind_device(comm)#: Bind this rank to one GPU of its node (round robin), before JAX starts if possible. No-op on the CPU, with one GPU, or when CUDA_VISIBLE_DEVICES is set.

class zoomy_jax.symsolver.parallel.no_halo#: Context manager inside which halo is the identity, giving the rank’s own block of an operator (block-Jacobi preconditioner).

zoomy_jax.symsolver.parallel.krylov(part, nc)#

The owned-cell mask (nc,) and the global sum (allreduce) a distributed Krylov solve needs.

Returns

(mask, gsum); inside custom_linear_solve nothing is differentiated through gsum.

zoomy_jax.symsolver.parallel.blocks(part, mesh_np)#

The halo exchange and the global reductions of the runtime’s name table for this rank’s piece, each face and cell counted once.

Args

part: zoomy_core.mesh.distributed.Partition of this rank
mesh_np: the rank’s local mesh