Pure Tensor Program Rewriting via Access Patterns (Representation Pearl)
arXiv:2105.09377 · doi:10.1145/3460945.3464953
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 \textit{pure but high-level}, making low-level rewrites to hardware targets inexpressible, or \textit{low-level but impure}, hampering the use of term rewriting altogether. This paper introduces Glenside, a pure IR whose core abstraction -- the \textit{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 \texttt{im2col}. Glenside establishes a new foundation for exploring further term rewriting techniques in optimizing low-level tensor programs.
arXiv commentTo be published at MAPS 2021
References in corpus (17)
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- TensorFlow: A system for large-scale machine learning
- NVIDIA Tensor Core Programmability, Performance & Precision
- Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions
- Interstellar: Using Halide's Scheduling Language to Analyze DNN Accelerators
- Learning to Optimize Tensor Programs
- Survey and Benchmarking of Machine Learning Accelerators
- egg: Fast and Extensible Equality Saturation
- MLIR: A Compiler Infrastructure for the End of Moore's Law
- Kernel Operations on the GPU, with Autodiff, without Memory Overflows
- Synthesizing Structured CAD Models with Equality Saturation and Inverse Transformations
- Equality Saturation for Tensor Graph Superoptimization
- SPORES: Sum-Product Optimization via Relational Equality Saturation for Large Scale Linear Algebra
- A High-Performance Sparse Tensor Algebra Compiler in Multi-Level IR
- Bring Your Own Codegen to Deep Learning Compiler
- Bitcoin's future carbon footprint
- Fireiron: A Scheduling Language for High-Performance Linear Algebra on GPUs