jaxDecomp: JAX Library for 3D Domain Decomposition and Parallel FFTs

Important Version 0.2.0 includes a pure JAX backend that no longer requires MPI. For multi-node runs, MPI and NCCL backends are still available through cuDecomp.

JAX reimplementation and bindings for NVIDIA's cuDecomp library (Romero et al. 2022), enabling multi-node parallel FFTs and halo exchanges directly in low-level NCCL/CUDA-Aware MPI from your JAX code.

Important Starting from version 0.2.8, jaxDecomp supports JAX's Shardy partitioner, which can be activated via jax.config.update('jax_use_shardy_partitioner', True). This partitioner is enabled by default in JAX 0.7.x and later versions. Shardy support is an internal implementation change and users should not expect any behavioral differences outside of what the JAX sharding mechanism provides, as explained in the JAX Shardy migration documentation.

---

Usage

Below is a simple code snippet illustrating how to perform a 3D FFT on a distributed 3D array, followed by a halo exchange. For demonstration purposes, we force 8 CPU devices via environment variables:

```python import os os.environ["XLAFLAGS"] = "--xlaforcehostplatformdevicecount=8" os.environ["JAXPLATFORMNAME"] = "cpu"

import jax from jax.sharding import Mesh, PartitionSpec as P, NamedSharding, AxisType import jaxdecomp

Create a 2x4 mesh of devices on CPU

pdims = (2, 4) mesh = jax.makemesh(pdims, axisnames=('x', 'y') , axis_types=(AxisType.Auto, AxisType.Auto)) sharding = NamedSharding(mesh, P('x', 'y'))

Create a random 3D array and enforce sharding

a = jax.random.normal(jax.random.PRNGKey(0), (1024, 1024, 1024)) a = jax.lax.withshardingconstraint(a, sharding)

Parallel FFTs

karray = jaxdecomp.fft.pfft3d(a) recarray = jaxdecomp.fft.pifft3d(a)

Parallel halo exchange

exchanged = jaxdecomp.haloexchange(a, haloextents=(16, 16), halo_periods=(True, True)) ```

All these functions are JIT-compatible and support automatic differentiation (with some caveats).

See also: - Basic Usage - Distributed LPT Example

Important Multi-node FFTs work with both JAX and cuDecomp backends\ For CPU with JAX, Multi-node is supported starting JAX v0.5.1 (with gloo backend)

---

Running on an HPC Cluster

On HPC clusters (e.g., Jean Zay, Perlmutter), you typically launch your script with: bash srun python your_script.py or bash mpirun -n 8 python your_script.py

See the Slurm README and template script for more details.

---

Using cuDecomp (MPI and NCCL)

For other features, compile and install with cuDecomp enabled as described in install:

```python import jaxdecomp

Optionally select communication backends (defaults to NCCL)

jaxdecomp.config.update('halocommbackend', jaxdecomp.HALOCOMMMPI) jaxdecomp.config.update('transposecommbackend', jaxdecomp.TRANSPOSECOMMMPI_A2A)

Then specify 'backend="cudecomp"' in your FFT or halo calls:

karray = jaxdecomp.fft.pfft3d(globalarray, backend='cudecomp') recarray = jaxdecomp.fft.pifft3d(karray, backend='cudecomp') exchangedarray = jaxdecomp.haloexchange( paddedarray, haloextents=(16, 16), haloperiods=(True, True), backend='cudecomp' ) ```

Install

1. Pure JAX Version (Easy / Recommended)

jaxDecomp is on PyPI:

1. Install the appropriate JAX wheel:
   • GPU: bash pip install --upgrade "jax[cuda]"
   • CPU: bash pip install --upgrade "jax[cpu]"
2. Install jaxdecomp: bash pip install jaxdecomp

This setup uses the pure-JAX backend—no MPI required.

2. JAX + cuDecomp Backend (Advanced)

If you need to use MPI instead of NCCL for GPU, you can build from GitHub with cuDecomp enabled. This requires the NVIDIA HPC SDK. Ensure nvc, nvc++, and nvcc are in your PATH, CUDA, MPI, and NCCL shared libraries are on LD_LIBRARY_PATH, and set CC=nvc and CXX=nvc++ before building.

bash pip install -U pip pip install git+https://github.com/DifferentiableUniverseInitiative/jaxDecomp -Ccmake.define.JD_CUDECOMP_BACKEND=ON

Alternatively, clone the repository locally and install from your checkout:

bash git clone https://github.com/DifferentiableUniverseInitiative/jaxDecomp.git --recursive cd jaxDecomp pip install -U pip pip install . -Ccmake.define.JD_CUDECOMP_BACKEND=ON

• If CMake cannot find NVHPC, set: bash export CMAKE_PREFIX_PATH=$CMAKE_PREFIX_PATH:$NVCOMPILERS/$NVARCH/22.9/cmake and then install again.

If jax complains about incompatiliby with CuSparse or any other library, the easiest way to solve this is by installing jax localy by running pip install jax[cuda-local] and then installing jaxDecomp with cuDecomp support.

---

Machine-Specific Notes

IDRIS Jean Zay HPE SGI 8600 supercomputer

As of February 2026, loading modules in this exact order works:

```bash module load nvidia-compilers/25.1 cuda/12.6.3 openmpi/4.1.6-cuda nccl/2.26.2-1-cuda cudnn cmake

Install JAX

pip install --upgrade "jax[cuda-local]"

Install jaxDecomp with cuDecomp

export CMAKEPREFIXPATH=$NVHPCROOT/cmake # sometimes needed pip install git+https://github.com/DifferentiableUniverseInitiative/jaxDecomp -Ccmake.define.JDCUDECOMP_BACKEND=ON ```

Note: If using only the pure-JAX backend, you do not need NVHPC.

Important for JeanZay users Make sure to load the correct architucture module before loading the nvidia-compilers module. For example for A100 you need to load module load arch/a100 first. You also need to set the CXXFLAGS to export CXXFLAGS="-tp=zen2 -noswitcherror" if you are using the H100 or A100 partition or if you are using AMD CPUs in general. More info in Jean Zay documentation.

NERSC Perlmutter HPE Cray EX supercomputer

As of November 2022:

```bash module load PrgEnv-nvhpc python export CRAYACCELTARGET=nvidia80

Install JAX

pip install --upgrade "jax[cuda]"

Install jaxDecomp w/ cuDecomp

export CMAKEPREFIXPATH=/opt/nvidia/hpcsdk/Linuxx8664/22.5/cmake pip install git+https://github.com/DifferentiableUniverseInitiative/jaxDecomp -Ccmake.define.JDCUDECOMP_BACKEND=ON ```

---

Backend Configuration (cuDecomp Only)

By default, cuDecomp uses NCCL for inter-device communication. You can customize this at runtime:

```python import jaxdecomp

Choose MPI or NVSHMEM for halo and transpose ops

jaxdecomp.config.update('transposecommbackend', jaxdecomp.TRANSPOSECOMMMPIA2A) jaxdecomp.config.update('halocommbackend', jaxdecomp.HALOCOMM_MPI) ```

This can also be managed via environment variables, as described in the docs.

---

Autotune Computational Mesh (cuDecomp Only)

The cuDecomp library can autotune the partition layout to maximize performance:

```python automesh = jaxdecomp.autotune(shape=[512,512,512])

'automesh' is an optimized partition layout.

You can then create a JAX Sharding spec from this:

from jax.sharding import PositionalSharding sharding = PositionalSharding(automesh) ```

---

License: This project is licensed under the MIT License.

For more details, see the examples directory and the documentation. Contributions and issues are welcome!