A Tensor Algebra Compiler for Sparse Differentiation
- Resource Type
- Conference
- Authors
- Shaikhha, Amir; Huot, Mathieu; Hashemian, Shideh
- Source
- 2024 IEEE/ACM International Symposium on Code Generation and Optimization (CGO) Code Generation and Optimization (CGO), 2024 IEEE/ACM International Symposium on. :1-12 Mar, 2024
- Subject
- Computing and Processing
Tensors
Codes
Dictionaries
Scalability
Pipelines
DSL
Optimization
Sparse Tensor Algebra
Automatic Differentiation
Semi-Ring Dictionaries
- Language
- ISSN
- 2643-2838
Sparse tensors are prevalent in many data-intensive applications. However, existing automatic differentiation (AD) frameworks are tailored towards dense tensors, which makes it a challenge to efficiently compute gradients through sparse tensor operations. This is due to irregular sparsity patterns that can result in substantial memory and computational overheads. We propose a novel framework that enables the efficient AD of sparse tensors. The key aspects of our work include a compilation pipeline leveraging two intermediate DSLs with AD-agnostic domain-specific optimizations followed by efficient C++ code generation. We showcase the effectiveness of our framework in terms of performance and scalability through extensive experimentation, outperforming state-of-the-art alternatives across a variety of synthetic and real-world datasets.