In plain words: A compiler lets programmers separately describe how tensor data is laid out and how the math is split across CPUs and GPUs, then generates code. Its matrix multiply matches hand-tuned code on 256 nodes and beats other systems 1.8x to 3.7x on larger tensor operations.
Abstract · DISTAL: The Distributed Tensor Algebra Compiler
We introduce DISTAL, a compiler for dense tensor algebra that targets modern distributed and heterogeneous systems. DISTAL lets users independently describe how tensors and computation map onto target machines through separate format and scheduling languages. The combination of choices for data and computation distribution creates a large design space that includes many algorithms from both the past (e.g., Cannon's algorithm) and the present (e.g., COSMA). DISTAL compiles a tensor algebra domain specific language to a distributed task-based runtime system and supports nodes with multi-core CPUs and multiple GPUs. Code generated by DISTAL is competitive with optimized codes for matrix multiply on 256 nodes of the Lassen supercomputer and outperforms existing systems by between 1.8x to 3.7x (with a 45.7x outlier) on higher order tensor operations.
Rohan Yadav, Alex Aiken, Fredrik Kjolstad
arXiv:2203.08069 · cs.PL, cs.DC · submitted Mar 15, 2022 · updated Mar 17, 2022
abstract · pdf · html
I'm an infrastructure engineer, mostly focused on databases, data pipelines, ML infra for the past 10-15 years. Even when designing homogenous compute clusters, I had to dig in and understand compiler-level implementations in MLIR and LLVM. I'm not a compiler expert by any measure, but know just enough to be dangerous and curious about (safely) scheduling computations across a pool of volunteers machines. Seems especially important to chew on now, with training of foundational LLM weights costing 7-9 figures.