Composing Distributed Computations Through Task and Kernel Fusion

Rohan Yadav, Shiv Sundram, Wonchan Lee, Michael Garland, Michael Bauer, Alex Aiken and 1 other

Drop-in distributed libraries (like a NumPy that automatically runs across many GPUs) let you compose programs freely, but the price is that each library call becomes its own batch of tasks and throwaway temporary arrays, so end-to-end code runs far slower than a hand-tuned version; this paper (Stanford and NVIDIA — Rohan Yadav, Shiv Sundram, and colleagues) presents Diffuse, a system that dynamically fuses tasks and their kernels across library boundaries in a distributed, task-based runtime, using a scale-free intermediate representation so the analysis cost does not grow with the machine size, and speeds up unmodified applications by 1.86 times on average on up to 128 GPUs.

The argument in one glance

Problem

Composable distributed libraries are convenient, but slow end-to-end

A popular way to scale scientific Python to a cluster is a drop-in library: you write ordinary NumPy or SciPy code, and a library like cuPyNumeric or Legate Sparse transparently runs it across many GPUs while keeping the same meaning. These libraries work by translating each operation into a stream of tasks (small user-defined functions, whose bodies are called kernels) issued to a task-based runtime system, which figures out the parallelism, communication, and synchronization for you. This is wonderful for composition: libraries can be written independently and then freely combined, because the runtime handles scheduling and keeps the distributed data coherent. But the same abstraction that makes composition easy hurts performance. Because each library operation is decomposed into its own tasks, a composed program allocates a temporary distributed array for every intermediate result and runs a separate pass over the data for each operation. That means poor data locality and far more memory traffic than a program organized by hand, where the operations would be fused together into a single pass. The convenience of composing libraries comes at a real end-to-end cost.
What people did

Fuse across library boundaries the way shared-memory systems do

The natural fix is fusion: combine the tasks and kernels of adjacent operations so intermediate results stay in registers or local memory instead of being written out to temporary arrays. Prior systems such as Weld and Split Annotations already do exactly this across library boundaries — but only for shared-memory libraries, where all the data lives in one address space and any operation can read any data cheaply.
Why it fell short

Distributed memory makes fusion depend on partitioning and aliasing

Those shared-memory techniques do not transfer to the distributed setting, because in distributed memory whether two operations can be fused is a subtle, data-dependent question. When data is spread across nodes, reading or writing shared data may require communication between nodes, and whether it does depends on how the distributed arrays are partitioned and whether they are aliasing views of the same underlying data. For example, a sequence of element-wise operations on two distributed arrays may or may not be fusible depending on whether the arrays are aliases of the same distributed data — fusing operations that would force one processor to read another processor's data is not allowed, because it hides communication that must actually happen. So a distributed fusion system needs an analysis that reasons about partitioning and aliasing to decide what is safe to fuse. The hard part is that a naive version of this analysis scales with the number of processors — it would examine each piece of distributed data on each node — which is exactly what you cannot afford on a large machine. What was missing was a way to fuse tasks and kernels across library boundaries in distributed memory, with an analysis whose cost does not grow with the size of the machine.
This paper's idea

A scale-free IR that makes distributed fusion analysis cheap

Diffuse dynamically fuses tasks and kernels in distributed, task-based runtimes, transparently achieving the optimizations of hand-tuned programs. Its foundation is a scale-free intermediate representation of distributed computation: it models data as stores (distributed arrays) with first-class partitions, and computation as a stream of index tasks over a launch domain, each task annotated with which partition of which store it reads, writes, or reduces. The representation is scale-free because a partition is described by a compact rule (a tiling with a shape, offset, and projection) rather than by listing every processor's piece, so the IR and the analyses on it do not grow with the number of processors. On top of this IR, Diffuse buffers the incoming task stream into a window and runs a dynamic task-fusion analysis that finds the longest prefix of tasks it can safely fuse: tasks are fusible when their dependencies are at most point-wise (each processor can proceed without talking to any other), which Diffuse checks with four constraints on partitions and privileges instead of materializing the full dependence map. It fuses that prefix into one task, eliminates the temporary arrays that are now internal to it, and hands the fused task to a JIT compiler built on MLIR that fuses the nested loops of the constituent kernels and generates CPU or GPU code. Because it reasons only about task privileges and partitions, Diffuse works across function and library boundaries: it accelerates unmodified cuPyNumeric and Legate Sparse applications by 1.86 times on average (and up to 10.7 times) on as many as 128 GPUs, matching or beating hand-optimized and MPI-based code.

Why this problem is worth solving

The abstraction that makes distributed libraries composable is the same one that makes composed programs slow, so bridging that gap has broad payoff.

Modern distributed scientific software increasingly relies on drop-in libraries that let you keep writing simple sequential-looking code while it quietly runs across a whole cluster. They achieve this by lowering each operation to a stream of tasks handed to a task-based runtime system, which extracts parallelism and inserts the necessary communication automatically. This architecture is what lets independently written libraries compose freely — you can mix a dense-array library and a sparse-linear-algebra library in one program, and the runtime keeps their distributed data coherent. The catch is granularity: because every library call becomes its own group of tasks operating on its own temporary distributed arrays, a program built by composing calls does far more memory traffic and far more separate passes over the data than the same computation written as one fused loop. A hand-tuned expert would fuse those operations to keep intermediates in fast memory, but the whole point of drop-in libraries is that users are not hand-tuning. So the value of closing this gap is large: it would let ordinary, composed high-level programs run at hand-tuned speed without asking the programmer to break the abstraction.

What makes the distributed version genuinely hard is that legality of fusion is a data-layout question, not a syntactic one.

In a single shared memory, deciding whether two operations can be fused is relatively easy, because any operation can read any data at roughly uniform cost. In distributed memory, the same two operations might be perfectly fusible or might require moving data between nodes, and which one it is depends entirely on how the arrays are partitioned across the machine and whether they are aliasing views of the same underlying data. Fusing two tasks is only safe when every processor can run its piece of both tasks using only the data it already holds — that is, when the dependency between the tasks is at most point-wise. If instead one processor would need a neighbor’s data (as in a stencil, where each cell reads its neighbors), fusing the write back into the next operation would silently bury communication that has to happen, giving wrong results. So a distributed fusion system must analyze partitions and aliasing to prove that no cross-processor communication is hidden. And it must do this analysis without its cost growing with the number of processors, because on a large machine, an analysis that inspects each node’s data individually is itself a bottleneck. Getting both — correctness under arbitrary partitioning and aliasing, and an analysis that stays cheap at scale — is the core difficulty.

What people did before, and where it broke

Compose libraries and accept the overhead. Let each operation become its own tasks and temporary arrays, and rely on the runtime to schedule them. Where it broke. Every intermediate result is a temporary distributed array and every operation is a separate pass over the data, so data locality is poor and composed programs run far slower than hand-tuned ones.
Fuse across library boundaries, shared-memory style (Weld, Split Annotations). Combine adjacent operations across libraries to keep intermediates in fast memory. Where it broke. These techniques assume shared memory; they have no model of how distributed data is partitioned, so they cannot tell when fusing two operations would hide required cross-node communication.
Materialize the full dependence map to check fusibility. For each candidate pair of tasks, compute exactly which point tasks depend on which. Where it broke. Materializing the dependence map scales with the number of processors, so the safety check itself becomes a bottleneck on large machines — the opposite of what a distributed system needs.
Rely on machine-learning-style graph compilers (Jax, PyTorch). Use structured program graphs and domain-specific fusion. Where it broke. These accept only restricted, structured representations and domain knowledge; they cannot express general programs with complex aliasing and mutation, like the fluid-dynamics and shallow-water simulations Diffuse targets.
The gap. No system fused tasks and kernels across library boundaries in distributed memory with an analysis whose cost is independent of the machine size, so composed distributed programs could not reach hand-tuned performance transparently.

How Diffuse works

Turning a stream of library tasks into fused, compiled work

↓ →
new
↓ →
new
↓ →
new
↓ →
new
↓ →
new

The idea

Fuse operations across library and function boundaries to remove the intermediates. Diffuse’s central act is to take a chain of separately-issued operations and fuse them into one, so the temporary arrays between them vanish and the data is touched in a single pass. Merging a sequence of independently-written operations into fused kernels so intermediate results never leave fast memory is the same undertaking the work on fusion-centric compilation for sparse deep learning (FuseFlow) covered elsewhere on this site pursues, making fusion the organizing principle of the compiler rather than an afterthought. Diffuse does this dynamically for distributed task-based programs, fusing across the boundaries between library calls that the programmer never has to break by hand.

Choose a representation whose cost does not grow with the machine. The reason Diffuse can analyze fusion on large clusters is its scale-free intermediate representation: partitions are described by compact rules (a tiling shape, offset, and projection) rather than by enumerating every processor’s piece, so the IR and its analyses stay the same size whether the machine has eight processors or thousands. Designing an intermediate representation carefully so that the important optimizations become cheap and tractable to compute on it is the same principle the work on an IR for optimized tensor computations on GPUs (Graphene) covered elsewhere on this site is built on, letting the representation itself carry the structure the analysis needs. Diffuse’s IR is what makes a correctness-preserving distributed fusion analysis affordable at scale.

Optimize a live stream of tasks that a runtime is already issuing. Rather than compiling a fixed program ahead of time, Diffuse sits in the middle of a running task-based system, buffering the stream of tasks the libraries emit and rewriting it on the fly. Working dynamically on the sequence of tasks a task-based runtime produces at execution time, and analyzing that stream to find optimization opportunities, is the same idea the work on automatic tracing in task-based runtime systems covered elsewhere on this site develops, treating the emitted task stream as the thing to analyze and transform. Diffuse applies this to fusion, finding fusible prefixes in the task window and replacing them before they reach the underlying runtime.

Following the five-point stencil through Diffuse

Fusing a distributed stencil computation, one window at a time

A trace, drawn from the paper's motivating example, of how Diffuse optimizes a multi-GPU cuPyNumeric five-point stencil.

Step 1 of 5

The library issues a stream of tasks with temporaries

The inner loop adds five shifted views of the grid and scales the result, then writes it back. cuPyNumeric lowers this to a stream of separate ADD tasks (each allocating a temporary distributed array), a MULT task, and a COPY task.

Every intermediate (t1, t2, t3, avg) is its own distributed allocation and its own pass.

Why it works

A scale-free IR keeps the analysis cheap at scale. A distributed fusion check could inspect each processor’s data → but Diffuse describes partitions by compact rules instead of per-processor lists → so the IR and its analyses stay independent of the number of processors, and fusion analysis does not become the bottleneck.
The point-wise dependence test captures exactly what is safe. Fusing tasks that need remote data would hide required communication → but Diffuse only fuses tasks whose dependencies are at most point-wise → so every processor can run the fused task using only data it already holds, and correctness is preserved.
Reasoning over privileges and partitions works across libraries. Library-specific fusion cannot cross into another library → but Diffuse looks only at task privileges (read, write, reduce) and partitions → so it finds fusion opportunities across function and library boundaries that the libraries could not see individually.
Eliminating temporaries is where the speed comes from. Fusing tasks alone still leaves intermediate arrays → but Diffuse demotes arrays internal to a fused task to task-local storage → so distributed allocations and their memory traffic disappear and arithmetic intensity rises.
Pairing task fusion with a kernel JIT compounds the win. A fused task still contains several separate loops → but the MLIR compiler fuses those nested loops into a single pass → so data is reused across the original operations and the result approaches hand-tuned code.

What it costs (trade-offs and scope)

Libraries must target Diffuse’s IR and expose their kernels. The transparency for end users is bought by work at the library level: cuPyNumeric and Legate Sparse were modified to emit Diffuse’s IR instead of talking to the runtime directly, and to expose their task bodies in MLIR for the JIT compiler. New libraries need the same integration.
The fusion analysis is sound but not complete. Diffuse’s constraints guarantee it never fuses something unsafe, but they are conservative: fusion opportunities that would require application-specific knowledge are out of scope, so some legal fusions are missed.
JIT compilation adds overhead. Fusing and compiling kernels at runtime costs time; the paper measures this separately and uses memoization to replay analysis and code generation across isomorphic task streams, but compilation is still a cost that must be amortized.
Benefit depends on how much fusible work the program has. Applications dominated by a single expensive operation gain little — dense Jacobi iteration, whose cost is a matrix-vector multiply, sees roughly 0.93 to 1.08 times (essentially no change), showing the analysis does not hurt but also cannot help where there is nothing to fuse.
It targets task-based runtimes with aliased distributed data. Diffuse is designed as a middle layer over a task-based system (Legion) with first-class partitions and privileges; programs that are explicitly parallel with manual communication (for example MPI code) require different analyses and are not the target.

The evaluation

Running composed distributed programs, with and without Diffuse

Property	Unmodified libraries	Hand-tuned / MPI baseline	Diffuse
Intermediate temporaries
Fusion across library boundaries
Analysis cost vs machine size
Application changes required
End-to-end speedup (up to 128 GPUs)

A comparison of running composed distributed programs. Diffuse eliminates temporaries and fuses across library boundaries automatically, with a fusion analysis whose cost is independent of the number of processors, and no application changes. It accelerates unmodified cuPyNumeric and Legate Sparse applications by 1.86 times on average (0.93 to 10.7 times) on up to 128 GPUs, matches or exceeds hand-optimized code (1.23 times over already hand-optimized), and beats the MPI-based PETSc library by 1.4 times on average.Underlined values explain themselves when clicked.

Where this stops working

What to keep in mind

This explainer is based on the paper’s full text (the authors’ arXiv copy of the ASPLOS 2025 paper). The system described is the paper’s: Diffuse dynamically performs task and kernel fusion in distributed, task-based runtime systems. Its key component is a scale-free intermediate representation of distributed computation — stores (distributed arrays), first-class partitions, and index tasks over launch domains with read, write, and reduce privileges — whose size and analysis cost are independent of the number of processors. Diffuse buffers the task stream into a window, uses a dynamic dependence analysis with four fusion constraints (launch-domain equivalence, true-, anti-, and reduction-dependence) to find a prefix of tasks whose dependencies are at most point-wise, fuses that prefix into one task, eliminates now-internal temporary arrays via a split reference-counting scheme, and pairs task fusion with an MLIR-based JIT compiler that fuses the kernels’ nested loops and generates CPU or GPU code, memoizing analyses across isomorphic task streams. The paper reports that Diffuse is a general enough target for two real libraries (cuPyNumeric and Legate Sparse), accelerates unmodified applications by 1.86 times on average (geo-mean) and between 0.93 and 10.7 times on up to 128 GPUs, matches or exceeds PETSc (1.4 times geo-mean), and finds opportunities missed by developers (1.23 times over hand-optimized code). Exact algorithms, the correctness argument, benchmark details, and full results should be read from the paper.

Background

Why composition and performance keep pulling against each other in distributed computing

It helps to see Diffuse within a long-running tension in high-performance computing: the abstractions that make software composable are usually the ones that leave performance behind. Layered libraries, clean interfaces, and drop-in replacements let people build large systems quickly and reuse each other’s work, but each abstraction boundary tends to force data to be materialized, copied, or moved when a monolithic hand-written version would have kept it in place. In distributed settings the stakes are highest, because crossing a boundary can mean not just a wasted memory pass but actual communication across the network, which is the scarcest resource. Much of the progress in this area comes from finding ways to preserve the programmer’s clean, composed view while a compiler or runtime quietly reorganizes the work underneath — fusing operations, overlapping phases, and eliminating the intermediate data that abstraction boundaries create. The effort to reduce the cost of crossing these boundaries by overlapping and interleaving distributed work is the same terrain the work on overlapping communication with dependent computation covered elsewhere on this site occupies, attacking the composition-versus-performance gap from the communication side. Diffuse attacks it from the fusion side, showing that with the right scale-free representation, a runtime can fuse across library boundaries in distributed memory and give composed programs the performance of code an expert would have written by hand.

If you remember only three things

1. Drop-in distributed libraries (a NumPy or SciPy that transparently scales to many GPUs) are built on task-based runtimes and compose cleanly, but each operation becomes its own group of tasks with its own temporary distributed arrays, so composed programs have poor locality and run far slower than hand-tuned code. Prior cross-library fusion worked only in shared memory, because in distributed memory whether two operations can be fused depends on how their data is partitioned and aliased — and the safety analysis must not blow up with the number of processors.
2. Diffuse dynamically fuses tasks and kernels across library boundaries using a scale-free intermediate representation: data is modeled as stores with compact-rule partitions, and computation as a stream of index tasks with read, write, and reduce privileges, so the IR and its analyses are independent of the machine size. It buffers the task stream, finds the longest prefix whose dependencies are at most point-wise (via four fusion constraints, without materializing the full dependence map), fuses it into one task, eliminates the now-internal temporaries, and JIT-compiles the fused kernels with MLIR into a single pass over the data.
3. On unmodified cuPyNumeric and Legate Sparse applications, Diffuse achieves a 1.86 times average speedup (0.93 to 10.7 times) on up to 128 GPUs, matches or exceeds hand-optimized code (1.23 times over hand-tuned) and beats the MPI-based PETSc by 1.4 times on average. The trade-offs: libraries must be integrated to emit Diffuse's IR and expose kernels in MLIR, the analysis is sound but conservative (misses fusions needing app knowledge), JIT compilation adds overhead, and programs with little fusible work (dense Jacobi) see essentially no change.

Check yourself

Check your understanding

1. Why do composed distributed programs built from drop-in libraries run slower than hand-tuned versions?
2. What makes deciding whether two operations can be fused harder in distributed memory than in shared memory?
3. What does it mean that Diffuse's intermediate representation is "scale-free," and why does it matter?
4. Predict: one application is a trivially parallel benchmark that is a long chain of element-wise operations (Black-Scholes), and another is dominated by a single dense matrix-vector multiply (Jacobi iteration). How will Diffuse affect each?

Sources

The paper is “Composing Distributed Computations Through Task and Kernel Fusion” by Rohan Yadav, Shiv Sundram, Wonchan Lee, Michael Garland, Michael Bauer, Alex Aiken, and Fredrik Kjolstad (Stanford University and NVIDIA), published at ASPLOS 2025, DOI 10.1145/3669940.3707216. It is licensed CC BY, with an author copy available at arXiv:2406.18109.

This explainer is written from the paper’s full text (the authors’ arXiv copy). Diffuse dynamically performs task and kernel fusion in distributed, task-based runtime systems. Its key component is an intermediate representation of distributed computation that makes the analyses needed for fusing distributed tasks scalable — a scale-free representation whose size is independent of the target machine, modeling data as stores with first-class partitions and computation as a stream of index tasks with read, write, and reduce privileges. Diffuse buffers the task stream into a window, runs a dynamic dependence analysis using four fusion constraints to find a prefix of tasks with at most point-wise dependencies, fuses that prefix into a single task, eliminates temporary arrays that become internal to it (via a split reference-counting scheme), and pairs task fusion with an MLIR-based JIT compiler that fuses the kernels within fused tasks and generates CPU or GPU code, memoizing analyses across isomorphic task streams. The paper shows the IR is general enough to be a target for two real libraries (cuPyNumeric and Legate Sparse), letting Diffuse find optimizations across function and library boundaries; it accelerates unmodified applications by 1.86 times on average (geo-mean) and between 0.93 and 10.7 times on up to 128 GPUs, enables high-level Python to match or exceed the MPI-based PETSc (1.4 times geo-mean), and finds opportunities missed by developers (1.23 times over hand-optimized code). Numbers are the paper’s own; exact algorithms, the correctness argument, and full results should be confirmed in the paper.