WACO: Learning Workload-Aware Co-optimization of the Format and Schedule of a Sparse Tensor Program

Jaeyeon Won, Charith Mendis, Joel S. Emer, Saman Amarasinghe

Research output: Chapter in Book/Report/Conference proceedingConference contribution

Abstract

In this paper, we present WACO, a novel method of co-optimizing the format and the schedule of a given sparsity pattern in a sparse tensor program. A core challenge in this paper is the design of a lightweight cost model that accurately predicts the runtime of a sparse tensor program by considering the sparsity pattern, the format, and the schedule. The key idea in addressing this is exploiting a sparse convolutional network to learn meaningful features of the sparsity pattern and embedding a coupled behavior between the format and the schedule using a specially designed schedule template. In addition, within the enormous search space of co-optimization, our novel search strategy, an approximate nearest neighbor search, efficiently and accurately retrieves the best format and schedule for a given sparsity pattern. We evaluated WACO for four different algorithms (SpMV, SpMM, SDDMM, and MTTKRP) on a CPU using 726 different sparsity patterns. Our experimental results showed that WACO outperformed four state-of-the-art baselines, Intel MKL, BestFormat, TACO with a default schedule, and ASpT. Compared to the best of four baselines, WACO achieved 1.43×, 1.18×, 1.14×, and 1.27× average speedups on SpMV, SpMM, SDDMM, and MTTKRP, respectively.

Original languageEnglish (US)
Title of host publicationASPLOS 2023 - Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems
EditorsTor M. Aamodt, Natalie Enright Jerger, Michael Swift
PublisherAssociation for Computing Machinery
Pages920-934
Number of pages15
ISBN (Electronic)9781450399166
DOIs
StatePublished - Jan 27 2023
Event28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, ASPLOS 2023 - Vancouver, Canada
Duration: Mar 25 2023Mar 29 2023

Publication series

NameInternational Conference on Architectural Support for Programming Languages and Operating Systems - ASPLOS
Volume2

Conference

Conference28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, ASPLOS 2023
Country/TerritoryCanada
CityVancouver
Period3/25/233/29/23

Keywords

  • Auto-Scheduling
  • Sparse Tensor
  • Tensor Compiler

ASJC Scopus subject areas

  • Software
  • Information Systems
  • Hardware and Architecture

Fingerprint

Dive into the research topics of 'WACO: Learning Workload-Aware Co-optimization of the Format and Schedule of a Sparse Tensor Program'. Together they form a unique fingerprint.

Cite this