Polyhedral Specification and Code Generation of Sparse Tensor Contraction with Co-iteration

Tuowen Zhao, Tobi Popoola, Mary Hall, Catherine Olschanowsky, Michelle Strout

Research output: Contribution to journalArticlepeer-review

2 Scopus citations

Abstract

This article presents a code generator for sparse tensor contraction computations. It leverages a mathematical representation of loop nest computations in the sparse polyhedral framework (SPF), which extends the polyhedral model to support non-Affine computations, such as those that arise in sparse tensors. SPF is extended to perform layout specification, optimization, and code generation of sparse tensor code: (1) We develop a polyhedral layout specification that decouples iteration spaces for layout and computation; and (2) we develop efficient co-iteration of sparse tensors by combining polyhedra scanning over the layout of one sparse tensor with the synthesis of code to find corresponding elements in other tensors through an SMT solver.We compare the generated code with that produced by a state-of-The-Art tensor compiler, TACO. We achieve on average 1.63× faster parallel performance than TACO on sparse-sparse co-iteration and describe how to improve that to 2.72× average speedup by switching the find algorithms. We also demonstrate that decoupling iteration spaces of layout and computation enables additional layout and computation combinations to be supported.

Original languageEnglish (US)
Article number16
JournalACM Transactions on Architecture and Code Optimization
Volume20
Issue number1
DOIs
StatePublished - Dec 16 2022
Externally publishedYes

Keywords

  • Data layout
  • code synthesis
  • index array properties
  • polyhedral compilation
  • sparse tensor contraction
  • uninterpreted functions

ASJC Scopus subject areas

  • Software
  • Information Systems
  • Hardware and Architecture

Fingerprint

Dive into the research topics of 'Polyhedral Specification and Code Generation of Sparse Tensor Contraction with Co-iteration'. Together they form a unique fingerprint.

Cite this