[Paper Review] Stardust: Compiling Sparse Tensor Algebra to a Reconfigurable Dataflow Architecture

Olivia Hsu, Alexander Rucker+3 more|arXiv (Cornell University)|Nov 7, 2022

Parallel Computing and Optimization Techniques4 citations

TL;DR

Stardust is the first compiler that automatically translates sparse tensor algebra expressions into reconfigurable dataflow architectures (RDAs), using a novel scheduling and data representation language to map sparse computations efficiently. It achieves 138× speedup over CPU kernels and 41× over GPU kernels by leveraging parallel-patterns rewriting and on-chip memory optimization on the Capstan RDA.

ABSTRACT

We introduce Stardust, a compiler that compiles sparse tensor algebra to reconfigurable dataflow architectures (RDAs). Stardust introduces new user-provided data representation and scheduling language constructs for mapping to resource-constrained accelerated architectures. Stardust uses the information provided by these constructs to determine on-chip memory placement and to lower to the Capstan RDA through a parallel-patterns rewrite system that targets the Spatial programming model. The Stardust compiler is implemented as a new compilation path inside the TACO open-source system. Using cycle-accurate simulation, we demonstrate that Stardust can generate more Capstan tensor operations than its authors had implemented and that it results in 138$ imes$ better performance than generated CPU kernels and 41$ imes$ better performance than generated GPU kernels.

Motivation & Objective

• To address the performance and energy efficiency gap in accelerating the long-tail of sparse tensor algebra workloads beyond what CPUs and GPUs can achieve.
• To make reconfigurable dataflow architectures (RDAs) accessible to data scientists by enabling high-level compilation from sparse tensor expressions.
• To provide a separation of concerns between algorithm, data representation, and scheduling for efficient mapping to sparse accelerators.
• To demonstrate that compiler-optimized mappings to RDAs can outperform hand-optimized CPU and GPU implementations.

Proposed method

• Introduces a new data representation language to express accelerator-specific tensor placement on physical memories in RDAs.
• Designs a scheduling language to map transformed sparse iteration spaces to sparse hardware components like data combiners and compute units.
• Employs an algorithm to bind abstract data structure descriptions to specific on-chip memories, minimizing data movement.
• Uses a parallel-patterns rewrite system to lower sparse tensor algebra expressions into the declarative-sparse Spatial programming model for RDAs.
• Extends the open-source TACO compiler with a new compilation path targeting Capstan, a reconfigurable dataflow architecture.
• Leverages cycle-accurate simulation to evaluate performance against CPU and GPU kernels.

Experimental results

Research questions

• RQ1 Can a high-level compiler map arbitrary sparse tensor algebra expressions to reconfigurable dataflow architectures with performance surpassing CPUs and GPUs?
• RQ2 How can scheduling and data representation abstractions be designed to enable efficient memory placement and dataflow mapping on sparse RDAs?
• RQ3 Can the separation of algorithm, data representation, and schedule enable new, optimized mappings beyond what was previously hand-coded for RDAs?
• RQ4 To what extent can compiler optimizations improve arithmetic intensity and reduce redundant computation in sparse tensor kernels on RDAs?

Key findings

• Stardust successfully compiles a broader set of Capstan tensor operations than the original authors had hand-implemented.
• The system achieves 138× higher performance than generated CPU kernels and 41× higher than generated GPU kernels in cycle-accurate simulation.
• The proposed scheduling and data representation languages enable the synthesis of optimized algorithms that outperform existing hand-optimized implementations.
• The separation of algorithm, data representation, and scheduling allows for independent optimization and enables new, previously untargeted sparse tensor operations on Capstan.
• The compiler's lowering pipeline via parallel-patterns rewrite successfully translates high-level tensor expressions into efficient Spatial code for the RDA.