A shared compilation stack for distributed-memory parallelism in stencil DSLs

Scientific communities build their own languages so physicists can write physics instead of loop nests, but each of those languages then has to grow a full compiler underneath it — parallel loops, message passing, vectorisation — and three separate stencil projects have built the same one three times with zero code shared; this work moves them onto the shared framework machine learning already uses, contributing the abstraction it was missing, which turns out to be distributed memory.

Why this problem is worth solving

The computation being targeted is narrow, common, and unusually well-suited to being automated.

Splitting the grid across machines is where the compiler earns its keep.

What people did before, and where it broke

• Write the parallel code by hand. What scientists did before these languages existed, and still do outside them. Where it broke. It puts cache blocking, vectorisation and message passing in front of people whose expertise is physics, and it produces code that has to be rewritten for each new machine.
• Build a domain-specific language with its own compiler. The approach all three projects here take, and it genuinely delivers — a physicist writes the equation and gets tuned parallel code. Where it broke. A large fraction of each compiler is generic: parallel directives, message passing, vectorisation, arithmetic simplification, loop transformations. All of it implemented three times, in three languages, with nothing shared.
• Accept the duplication as the cost of specialisation. The default position. Where it broke. It creates what the paper calls uncertainty around longevity and adoption at scale. Committing a national weather model to a language maintained by a small team is a real risk, and duplicated effort makes that team smaller than it needs to be.
• Use the extensible compiler framework that solved this for machine learning. The obvious move, and it works — deep learning libraries already share optimisation infrastructure through it while keeping distinct interfaces. Where it broke. It has abstractions for shared-memory parallelism, GPUs and vectorisation, but none for distributed memory. For supercomputing that is the missing piece, not a detail.
• Share code between stencil frameworks directly. Attempted by earlier work on common abstractions across stencil frameworks. Where it broke. Integrating across different implementation languages, without a standard interchange format, is exactly what makes sharing impractical in the first place.
• The gap. The framework and the need had both existed for a while. What was missing was the abstraction sitting between them — a way to say what data has to move where without saying which library sends it, so that the geometry of a domain decomposition becomes something a compiler can compute and optimise rather than something a person writes out.

How the shared stack lowers a stencil

The idea

Most of a domain-specific compiler is not about the domain, and that is where the waste is. The observation driving this work is deflating and correct. Each of these three projects presents a genuinely different interface — one lets you write symbolic mathematics in Python, one takes a stencil specification, one takes Fortran with conventions — and those interfaces matter enormously, because they are how each community talks to its own domain experts. But underneath, all three arrive at the same concept of a stencil, and then all three independently build parallel loop generation, message passing, vectorisation, common subexpression elimination and loop tiling. The proposal is to keep the interfaces and share everything below them, which is exactly the arrangement deep learning libraries already have. The consequence the paper cares about is not performance but survival: a modelling effort that will run for a decade is being asked to bet on a compiler maintained by a handful of people, and duplicated effort makes that bet worse than it needs to be. Building on a shared extensible framework rather than a private compiler is the same wager as the hardware accelerator representation work covered elsewhere on this site.

The missing abstraction was distributed memory, and it splits into geometry and mechanism. The shared framework already had layers for parallel loops, GPUs and vector operations, but supercomputing runs across machines and there was nothing for that. The contribution is two layers rather than one, and the split is the interesting part. The upper layer describes only geometry — a grid of participating machines, and for each exchange a rectangular region to receive together with a same-sized region to send and the relative direction of the neighbour involved. That is a complete description of a halo exchange with no mention of how a byte actually moves, and the paper states plainly that nothing in it is specific to the message passing library, so another communication library could be targeted instead. The lower layer mirrors that library directly and has been contributed upstream to the framework, so it is available to anyone. Making communication a first-class thing a compiler manipulates rather than a library call a programmer writes is the same move as the collective communication language work covered elsewhere on this site.

Constraining the upper layer is what makes the lower ones tractable, and joining the ecosystem buys capabilities nobody built. There is a quiet lesson in how the pieces fit. The stencil representation requires grid bounds to be known at compile time, which sounds like a limitation and means recompiling when a problem size changes — and the paper defends it on the grounds that constant-folding address computations reduces register pressure, which matters on GPUs. But it also notes something better: those known bounds greatly reduced the complexity of the decomposition layer built on top, so a restriction upstream paid for simplicity downstream. The other payoff is the one that makes the whole argument concrete. The Fortran-based project gained the ability to target field-programmable gate arrays, from unmodified source, purely by joining an ecosystem where somebody else had written that lowering — and because the stencil is described at a high level rather than as loops, the compiler could automatically restructure it into dataflow form, improving on the untransformed version by between 100 and 214 times. Automatically deriving a dataflow architecture from a high-level description is the same capability as the hierarchical dataflow synthesis work covered elsewhere on this site.

Turning a stencil into an exchange

Each project keeps everything above this point. The Fortran project routes anything it did not recognise as a stencil into a separate Fortran representation, which then calls into the stencil-generated code as ordinary functions, so partial adoption works.

Why it works

• Regular access patterns make the communication derivable. Every stencil update reads a fixed set of neighbouring offsets known at compile time → so the halo width and the exact slabs to exchange follow mechanically from scanning them → which turns the most error-prone part of distributed scientific code into a compiler pass.
• Separating geometry from mechanism keeps each layer simple. The upper layer says which rectangles move where and names no library → so it can be reasoned about, optimised and redundancy-eliminated on its own terms → and the lower layer only has to translate declarations into calls, buffers and packing.
• Explicit dataflow makes the optimisation nearly free. The representation records where every value is defined and used → so finding exchanges whose data has not changed needs no separate analysis infrastructure → and the conservative insert-everywhere pass can be cleaned up afterwards rather than being clever up front.
• A restriction upstream buys simplicity downstream. Requiring grid bounds at compile time forces recompilation when sizes change → but it lets address computations be constant-folded, which cuts register pressure on GPUs, and it substantially simplified the decomposition layer → so the constraint pays for itself twice.
• Joining an ecosystem gives you work you did not do. Both frontends now reach whatever targets the shared framework reaches → so the Fortran project acquired field-programmable gate array support from unmodified source without writing it → and the high-level stencil description let the compiler restructure into dataflow form automatically, worth between 100 and 214 times over the untransformed version.

What it costs (trade-offs and scope)

• It loses on the arithmetic-heavy kernels, and the paper says why. Against a mature baseline described honestly as highly efficient, the shared stack wins by up to 1.4 times on memory-bound diffusion kernels but runs at 0.6 to 0.8 times on the higher arithmetic intensity wave kernels. The stated cause is limited vectorisation in the lowered output, which the authors note is essential for stencils.
• Distributed scaling follows the trend but does not match it. The mature baseline supports diagonal exchanges in three dimensions, which the new decomposition does not, so it scales in the same shape while staying behind. Adding diagonal communication and overlapping communication with computation are named as future work.
• One benchmark runs at 0.62 times, for an infrastructure reason the authors identify precisely. A workload with eighteen separate stencil regions becomes eighteen separate parallel regions, because of how the shared framework lowers parallel loops, and profiling found the barrier spin-wait function to be the most expensive in the program. The conclusion names this lowering as the limitation the benchmarking exposed.
• The same lowering hurts on GPUs. Each parallel loop becomes a synchronous kernel launch with a device-wide wait at the end, so the same workload shows the same shortfall there.
• Grid bounds must be known at compile time. Changing the problem size means recompiling. Defended for good reasons, but it is a real constraint on how these compilers can be deployed.
• There is no high-level way to express boundary conditions. They are currently hand-encoded as conditionals over memory accesses, with a proper representation scheduled for future work — a notable omission given that boundary handling is a large part of real solver code.
• The decomposition handles rectangles only. Contiguous hyper-rectangular subsections, explicitly not sparse or irregularly shaped domains.
• Only two of the three compilers are actually integrated. The stencil representation comes from the third project and is extended here, but the prototype covers the other two; the third contributes the abstraction rather than being ported onto it.
• Comparisons are against what was available. The GPU comparison for one project is against its directive-based backend rather than its own hand-written accelerator backend, which required a commercial licence. The gate array results are acknowledged to fall short of the GPU, while matching what manual tuning achieves.
• This is a prototype on modified forks. The evaluation uses modified versions of all three codebases, and the message passing layer covers a subset of the older standard.

The evaluation

Check yourself

Sources

The paper is “A shared compilation stack for distributed-memory parallelism in stencil DSLs” by George Bisbas, Anton Lydike, Emilien Bauer, Nick Brown, Mathieu Fehr, Lawrence Mitchell, Gabriel Rodriguez-Canal, Maurice Jamieson, Paul H. J. Kelly, Michel Steuwer and Tobias Grosser, published at ASPLOS 2024 Volume 3, pages 38-56, DOI 10.1145/3620666.3651344. The first four authors contributed equally. It is open access under a Creative Commons Attribution licence, and the implementation is public.

This explainer is written from the full text, read from the authors’ copy on arXiv (2404.02218). Everything above comes from the paper’s own sections 2 through 8: the argument that the three stencil compilers share concepts while sharing no code, and that a large portion of each codebase addresses generic rather than domain-specific concerns; the description of the three projects and their distinct interfaces; the account of static single assignment and nested regions as the basis for interfacing domain-specific representations, and the use of a Python-native clone of the framework sharing its textual form; the section 4.1 changes to the stencil representation, including moving bounds into the types, generalising beyond three dimensions, adding a lowering suited to shared-memory parallelism, and the stated limitations of compile-time bounds and absent boundary condition modelling; the section 4.2 decomposition layer with its declarative exchange format, the pluggable decomposition strategy, the automatic insertion and subsequent dataflow-based elimination of exchanges, the note that nothing in it is specific to the message passing library, the observation that compile-time bounds simplified it, and its restriction to hyper-rectangular subsections; the section 4.3 message passing layer listing its supported operations and its contribution upstream; the section 5 integration of both frontends including the routing of unrecognised Fortran into a separate representation; and the section 6 evaluation on two supercomputers, a GPU and a gate array, covering the diffusion and wave results, the strong scaling comparison and its attribution to diagonal exchanges, the two advection benchmarks with the parallel-region diagnosis and the unified memory page fault finding, the gate array transformation results, and the conclusion naming the framework’s parallel loop lowerings as the limitation the benchmarking identified.