Guided Tensor Lifting

TL;DR

STAGG addresses the challenge of lifting legacy tensor code to tensor-domain-specific languages by learning the search space from large language models and applying an enumerative template-based synthesis guided by a probabilistic grammar. It combines LLM-generated candidate solutions with a template grammar, dimension prediction from static analysis and LLMs, and two search strategies (Top-Down and Bottom-Up A*) to efficiently navigate the template space. Validation uses IO-based checks followed by bounded model checking in MLIR/CBMC with rational datatypes to ensure equivalence to the original code. The approach achieves $99\%$ lifting accuracy on a large dense tensor benchmark with an average time of $3.19$ seconds and outperforms state-of-the-art solvers like C2TACO and Tenspiler, demonstrating the value of learning the search space rather than hard-coding heuristics. This work suggests a scalable pathway for automatic tensor-code lifting across DSLs and motivates broader application of LLM-guided synthesis in domain-specific code generation.

Abstract

Domain-specific languages (DSLs) for machine learning are revolutionizing the speed and efficiency of machine learning workloads as they enable users easy access to high-performance compiler optimizations and accelerators. However, to take advantage of these capabilities, a user must first translate their legacy code from the language it is currently written in, into the new DSL. The process of automatically lifting code into these DSLs has been identified by several recent works, which propose program synthesis as a solution. However, synthesis is expensive and struggles to scale without carefully designed and hard-wired heuristics. In this paper, we present an approach for lifting that combines an enumerative synthesis approach with a Large Language Model used to automatically learn the domain-specific heuristics for program lifting, in the form of a probabilistic grammar. Our approach outperforms the state-of-the-art tools in this area, despite only using learned heuristics.

Figures (12)

• Figure 1: Overview of STAGG. We query the LLM to provide 10 possible solutions in TACO that are equivalent to the input code C. Based on the LLM response, we build a probabilistic grammar and enumerate the space of template programs described by the said grammar. We validate a candidate using I/O examples, and if it passes all tests, we proceed to verification to prove equivalence with the original C implementation. The input code is also analyzed to predict the dimensionality of the left-hand side tensor in the solution.
• Figure 2: A C implementation of $\sum_{i=0}^{N-1} \texttt{Mat1}(f \times N + i) \cdot \texttt{Mat2}(i)$. The result is a dot product between the $f$-th row of Mat1 and vector Mat2. The equivalent synthesized TACO expression is a(i) = b(i,j) $*$ c(j).
• Figure 3: A probabilistic context-free grammar template.
• Figure 4: Expression standardization. We omit part of the derivation for brevity.
• Figure 5: The grammar for TACO expression in Extended Backus–Naur form, defining the syntax for tensor expressions, identifiers, constants, and basic arithmetic expressions. The $^\texttt{*}$ symbol denotes Kleene star, indicates zero or more repetitions of the preceding element, while $^\texttt{+}$ denotes Kleene plus, requires one or more occurrences.

Theorems & Definitions (7)

• definition 1: Context-Free Grammar, CFG
• definition 2: Weighted Context-Free Grammar, CFG
• definition 3: Probabilistic Context-Free Grammar, pCFG
• definition 4: Templates
• definition 5: Dimension list
• definition 6: Derivations
• Example 6.1