A Compilation-Based Approach to Performant Reduction and Redistribution Collective Communication Algorithms

Andreas Jocksch (ETH Zurich / CSCS), C. Nicole Avans, Riley Shipley, Anthony Skjellum (Tennessee Technological University) | IJHPCA, Special Issue Paper, 2026, Vol. 40(2), pp. 219–239 | DOI: 10.1177/10943420251363423

Expanded from the authors' PPAM 2024 paper (LNCS 15579, pp. 273–286). Code: github.com/eth-cscs/ext_mpi_collectives. Library name: ext_mpi.

Problem

As core counts per node keep rising (128 cores on dual-socket AMD EPYC, 288 cores across four sockets on NVIDIA Grace-Hopper), intra-node shared-memory collective communication becomes a central bottleneck, and vendor MPI libraries (MPICH, Open MPI) underperform for medium and large messages on high-core-count, multi-socket nodes. Many applications repeatedly invoke identical collectives — repeated allreduce in Neko, distributed AI training, and CP2K — which motivates paying a complex one-time optimisation cost that is then amortised over many executions. The classical constant-radix-two collective algorithms do not map cleanly onto hierarchical multi-socket hardware.

Core Insight

Treat a collective as something to be compiled: derive a hardware-tuned, multi-radix factorisation once in a setup phase, translate it through an intermediate colour-coded assembler into executable bytecode, and then run that bytecode many times. This cleanly separates deriving complex factorisations from executing the collective, mapping naturally onto the MPI 4.0 persistent interface, while a blocking interface with lazy initialisation and 'wisdom'-file caching serves applications with varying message sizes. Flexible (multi-radix) factors are mapped to the hardware hierarchy (quadrant → socket → node) so the largest data moves on the highest-bandwidth links, and barriers are integrated into the algorithm as counters.

Method

Allreduce is expressed as a reduce_scatter stage followed by an allgather stage, OR as a reduce stage followed by a broadcast stage, with different factors at each stage. Four shared-memory algorithms are used: recursive exchange (butterfly, non-commutative-capable), cyclic shift (commutative-only, dissemination variant), the tree algorithm (reduce+broadcast, avoids extra memory copies before/after inter-node communication), and the new cyclic copy-in reduction (commutative-only, combines copy-in with reduction). Factorisations are written as signed integer sets — negative for reduce_scatter, positive for allgather (e.g., {−4, −2, 2, 4}). Inter-node allreduce runs as reduce_scatter, a nested allreduce (implemented as allgather-with-reduction), then allgather; with multiple tasks per node it splits into shared-memory and inter-node phases. Constraints bound the factor search: split pieces stay at or above the 64-byte cache line, per-step message size should fit in cache, and inter-node factors relate to port count (factor = ports + 1). The compiler emits assembler (SMEMCPY / SREDUCE, IRECV / ISEND / WAITALL, SET/WAIT_NODE_BARRIER, MEMORY_FENCE) then bytecode; on GPU, consecutive copy/reduce operations are fused into a single CUDA kernel (a barrier, send, or receive breaks fusion). CPU shared buffers use XPMEM; GPU uses CUDA IPC. The library sits atop any MPI 4.0 library via the PMPI profiler hook, with fallback to the native routine.

Experimental Setup

Headline Quantitative Results

• Intra-node CPU allreduce: up to about half an order of magnitude faster than persistent and blocking MPICH / Open MPI on dual-socket AMD EPYC, and almost an order of magnitude faster on four-socket NVIDIA Grace-Hopper, for medium and large messages.
• Prime task counts hurt: with 127 tasks on EPYC, only {127} is possible for recursive exchange, so cyclic shift is used and performance drops to merely comparable with MPICH/Open MPI; on a single Grace socket there is no cyclic-shift penalty.
• GPU allreduce: on one Grace-Hopper node ext_mpi matches NCCL when NCCL's 'blocking' flag is enabled (both reduce to a single CUDA kernel per task) and beats HPE MPI for short messages; for long messages on four nodes ext_mpi is almost identical to NCCL and faster than HPE MPI.
• Initialisation cost: up to four orders of magnitude slower than a single short-message execution; whole-node init is two orders of magnitude costlier than per-socket init (MPICH/Open MPI init is far cheaper).
• CP2K QS/H2O-512: 2.5% overall speedup — allreduce time halved, reduce_scatter_block time cut to 1/20 vs MPICH (reduce_scatter_block 0.09834 s → 0.004390 s per call; allreduce 2.949·10⁻⁵ s → 1.410·10⁻⁵ s).
• CP2K QS_mp2_rpa: 1.5% overall speedup — allreduce again halved (2.808·10⁻⁵ s → 1.959·10⁻⁵ s), limited by the moderate reduce_scatter_block gain (0.2959 s → 0.2562 s).
• Combinatorial factor space (Table 1): recursive exchange grows modestly (128 tasks → 64 options), cyclic shift explodes (128 → 1378; 288 → 5421).

Limitations

• Initialisation is expensive (up to four orders of magnitude over a single short-message execution; whole-node two orders over per-socket) because of disadvantageous algorithmic complexity in the init routine.
• Cyclic shift and cyclic copy-in support commutative operators only.
• Prime task counts force a poor {prime} factorisation and degrade performance, most visibly on multi-socket AMD EPYC.
• The blocking interface is limited to message sizes needing no padding (or an extra memory copy when combined with inter-node communication).
• Reduce is only slightly faster than references because the reduction runs in a buffer with a final copy to root (the MPI standard forbids using non-root buffers directly).
• Not all CP2K MPI datatypes/operators are implemented, so the MPICH fallback reduces average speedups (especially for reduce_scatter_block); on Grace the vendor MPI used CMA not XPMEM due to current limitations.

Open Problems

1. Replacing empirical parameter (factorisation) selection with a machine-learning-guided heuristic for extensibility and performance portability, complementing the compiled approach.
2. Extending the algorithms to sparse data (common in deep learning) via data compression or novel locality-aware algorithms.
3. Reducing the expensive initialisation cost, whose complexity is currently unfavourable at whole-node scale relative to MPICH/Open MPI.
4. An MPI_Info extension permitting use of non-root buffers, which would unlock further reduce speedups.

Note on NCCL Tuning

The paper's central lever — choosing a multi-radix factorisation per hardware tier (quadrant → socket → node) and a message split whose smallest piece stays at the cache-line size — is structurally the same decision NCCL exposes through algorithm, protocol, and chunkSize selection. The optimal factor set is strongly hardware- and message-size-dependent (Table 3: {−16,−4,−2,2,4,16} on dual-socket EPYC versus {−9,−8,8,9} on single-socket Grace), and a wrong choice such as the forced {127} prime factorisation collapses performance — evidence that per-collective configuration search, not a fixed default, is what recovers the speedup. The authors' explicit future-work goal of replacing empirical parameter selection with a machine-learning-guided heuristic is the same problem statement as learned NCCL configuration tuning, applied to a compiled MPI stack.