Pure Tensor Program Rewriting via Access Patterns

Abstract

Tensor kernels in machine learning (ML) often correspond to pure mathematical expressions, making term rewriting an attractive strategy for optimization and mapping to specialized hardware accelerators. However, existing ML intermediate representations (IRs) tend to either be pure but high-level, making low-level rewrites to hardware targets inexpressible, or low-level but impure, hampering the use of term rewriting altogether.

This paper introduces Glenside, a pure IR whose core abstraction — the access pattern — enables low-level, layout-aware, hardware-centric program rewrites. We demonstrate how term rewriting in Glenside can be used to map program fragments to hardware accelerator invocations and automatically discover classic data layout transformations like im2col. Glenside establishes a new foundation for exploring further term rewriting techniques in optimizing low-level tensor programs.

Talk

MAPS 2021 talk by Gus Smith.

BibTeX

@inproceedings{2021-maps-glenside,
  title     = {Pure Tensor Program Rewriting via Access Patterns (Representation Pearl)},
  author    = {Smith, Gus Henry and Liu, Andrew and Lyubomirsky, Steven and Davidson, Scott and McMahan, Joseph and Taylor, Michael and Ceze, Luis and Tatlock, Zachary},
  series    = {MAPS 2021},
  booktitle = {Proceedings of the 5th ACM SIGPLAN International Symposium on Machine Programming},
  date      = {2021},
  url       = {https://doi.org/10.1145/3460945.3464953},
  doi       = {10.1145/3460945.3464953},
  publisher = {Association for Computing Machinery},
}

📝 publications index