Pure tensor program rewriting via access patterns (representation pearl)

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, layoutaware, 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.CCS Concepts: • Software and its engineering → Domain specific languages; • Computing methodologies → Machine learning; • Theory of computation → Equational logic and rewriting.

Verified tensor-program optimization via high-level scheduling rewrites

“…It is specifically derived from recent work on a functional tensor language for automatic differentiation [Bernstein et al 2020], which we have extended with features like reshape operators to express a richer space of implementation details relevant to performance optimization. Concurrent work on Glenside [Smith et al 2021] attempts to capture some similar implementation details by augmenting a functional tensor language with an algebra of łaccess patterns. ž…”

“…One popular approach is phrasing optimizations as transformations within a single source or intermediate language; the most common kind of transformation is a rewrite rule, a quantified term equality. There have been several success stories, including in Elevate [Fu et al 2021], Glenside [Smith et al 2021], and others [Kommrusch et al 2021;Steuwer et al 2015]. A typical such project presents rewrite rules as axioms and then studies engines to apply them effectively.…”

Specialized Accelerators and Compiler Flows: Replacing Accelerator APIs with a Formal Software/Hardware Interface

“…Compiler IR pattern Compiler IR pattern front-end and the code-generation capability provided by the TVM framework [14]. For instruction selection, it leverages the rewrite rules and the equality saturation engine provided by Glenside and egg [68,83]. The ILA-models of accelerators, the validation of IR-accelerator mappings, and compilation-results validation at the application level are powered by ILA-based methods in ILAng [35].…”

“…We leverage the egg library for equality saturation in our prototype [83]. First, the input program is translated from Relay to Glenside, a pure (side effectfree) tensor program representation that supports specifying rewrite rules for tensor programs [68]. Next, with both the compiler IR rewrites and IR-accelerator rewrites provided in Glenside, the equality saturation engine explores the space of possible rewrites as discussed in § 2.2.…”

“…For example, flexible matching revealed several offloads to Flex-ASR's linear layer in MobileNet-V2 by rewriting nn.dense to nn.dense followed by an add of a zero tensor. Also note that certain Glenside rewrites [68] that implement the im2col optimization [13] result in many more offloads to VTA in that table; this is due to 2D convolutions being rewritten into matrix multiplications. Hence, flexible matching allowed us to support 2D convolutions on VTA even though our prototype code generator did not explicitly implement 2D convolutions via VTA instructions.…”