DISTAL: The Distributed Tensor Algebra Compiler
arXiv:2203.08069 · doi:10.1145/3519939.3523437
Abstract
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.
References in corpus (1)
Cited by in corpus (8)
- The Sparse Abstract Machine
- TAMM: Tensor Algebra for Many-body Methods
- Minimum Cost Loop Nests for Contraction of a Sparse Tensor with a Tensor Network
- Full Version: (De/Re)-Composition of Data-Parallel Computations via Multi-Dimensional Homomorphisms
- Generating coupled cluster code for modern distributed memory tensor software
- Scaling Deep Learning Computation over the Inter-Core Connected Intelligence Processor with T10
- PartIR: Composing SPMD Partitioning Strategies for Machine Learning
- Fast Kronecker Matrix-Matrix Multiplication on GPUs