Prism: Symbolic Superoptimization of Tensor Programs

Paper: 2604.15272 Authors: Mengdi Wu, Xiaoyu Jiang, Oded Padon, Zhihao Jia Categories: cs.PL, cs.AI, cs.LG

The Gap

Existing tensor program optimizers fall into two camps. Compiler-based approaches (TVM, XLA) use heuristic search with hand-crafted rules — fast but miss optimal implementations. Superoptimizers (MetaFlow, TASO) exhaustively enumerate concrete programs — thorough but explode combinatorially for real workloads. The gap: no one bridges exhaustive rigor with practical scalability for modern LLM-scale tensor computations.

Problem: Optimize tensor ops for LLMs
         |
         v
    Assumption: Can represent program families symbolically
         |
         +---> Method: Two-level search (symbolic + concrete)
         |              with provable pruning
         |
         +---> Evidence: 2.2x over superoptimizers, 4.9x over compilers
         |               on 5 LLM workloads
         |
         v
    Conclusion: Symbolic reasoning enables exhaustive search at scale

The Increment

One sentence: Before Prism, you chose between fast-but-incomplete heuristics or slow-but-thorough enumeration; after Prism, symbolic reasoning lets you exhaustively search while pruning provably bad paths.

Core Mechanism

Prism operates in two levels. First, it generates symbolic graphs (sGraphs) where execution parameters like tile sizes and loop orders are symbolic variables, not concrete values. Each sGraph represents an entire family of implementations. Second, it instantiates these symbolic graphs into concrete programs by assigning specific values to parameters.

The key innovation is structured pruning during symbolic generation. Prism encodes operator semantics, algebraic identities (like matrix multiplication associativity), and hardware constraints (memory hierarchy, parallelism limits) as symbolic rules. When exploring the search space, it can prove entire families of programs are suboptimal without instantiating them. For example, if a symbolic graph violates memory bandwidth constraints for all possible tile sizes, Prism discards the whole family.

Input Tensor Program
         |
         v
    [Symbolic Graph Generator]
         |
    (generates sGraphs with symbolic params)
         |
         +---> sGraph_1 (tile=T, order=O)
         |     sGraph_2 (tile=T', order=O')
         |     ...
         |
         v
    [Symbolic Pruning Engine]
    (uses operator semantics + algebra + HW constraints)
         |
    (prunes provably suboptimal families)
         |
         v
    Surviving sGraphs
         |
         v
    [Parameter Instantiation]
    (auto-tuning: T=32, O=ijk)
         |
         v
    Concrete Optimized Programs

Think of Prism like a chess engine with opening books. Traditional superoptimizers play out every possible game move-by-move (concrete enumeration) — exhausting. Compilers use opening heuristics but stop analyzing mid-game (incomplete). Prism maintains an opening book of symbolic positions (sGraphs) where each entry represents thousands of similar games. It prunes entire branches by proving “all games starting with this opening lose to this defense” (symbolic reasoning), then only plays out the promising openings with specific move sequences (instantiation). The opening book is the symbolic layer; the actual games played are the concrete layer.

Key Concepts

• Symbolic Graph (sGraph): Instead of representing one program with concrete tile size 32 and loop order ijk, an sGraph represents all programs with tile size T (symbolic variable) and loop order O (symbolic variable). It’s a template with holes. The power: you can reason about properties that hold for all values of T and O without trying each combination. Example: “For any T > memory_size, this sGraph will thrash cache” — proven once, eliminates infinite concrete programs.

• Two-Level Search: Most optimizers search a flat space of concrete programs. Prism searches a hierarchical space: first find good symbolic templates (families), then find good parameter values within each family. This is like first choosing “what kind of algorithm” (symbolic: divide-and-conquer vs dynamic programming) before tuning “how to implement it” (concrete: recursion depth, memoization table size). The levels interact: symbolic pruning uses hardware constraints that depend on parameters, but proves properties that hold across parameter ranges.

• E-graph Rewriting for Equivalence: When Prism generates two sGraphs, it needs to know if they’re equivalent (compute the same result). Traditional approaches execute both and compare outputs — expensive and only checks one parameter assignment. Prism uses e-graphs (equality graphs) that compactly represent equivalent expressions. It rewrites sGraphs using algebraic identities (A×B×C = A×(B×C)) and checks if they land in the same equivalence class. This proves equivalence symbolically for all parameter values at once.

Framework Shift

Before (mainstream approach):        After (this paper):

Concrete Program Space               Symbolic Family Space
                                            |
  P1  P2  P3  P4  P5  ...           sG1 ----+---- sG2
   \   |   |   /   /                 |             |
    \  |   |  /   /              (T=32,O=ijk)  (T=64,O=ikj)
     \ |   | /   /                   |             |
      Enumerate                     P1,P2        P3,P4
      & Test                          |             |
        |                          Test          Test
        v                             |             |
   Best Program                       +------+------+
                                             |
                                             v
                                        Best Program

(Try every program)              (Try families, prune symbolically,
                                  instantiate survivors)

From exhaustive concrete enumeration to hierarchical symbolic search, the core shift is: reason about program families before instantiating individuals.

Expert Assessment

Problem choice: Real gap. LLM inference is bottlenecked by tensor ops, and the optimization space genuinely explodes (tile sizes × loop orders × fusion strategies × parallelization schemes). Compilers leave 2-5× performance on the table; superoptimizers timeout on realistic workloads. This sits at the intersection of PL verification techniques and ML systems — timely and well-motivated.

Method maturity: Clever synthesis of existing ideas (symbolic execution, e-graphs, auto-tuning) rather than a single novel insight. The contribution is engineering them together effectively. The symbolic pruning rules are domain-specific (tensor algebra, memory hierarchy) — not a general framework. Could simpler approaches work? Possibly: learned cost models might prune similarly without symbolic reasoning, but the paper doesn’t compare against ML-guided search.

Experimental integrity: Baselines are fair (MetaFlow, TASO, TVM, XLA). Numbers look solid: 2.2× speedup is meaningful but not shocking given the exhaustive search. Red flag: only 5 workloads tested, all from LLM domain (attention, MLP, LayerNorm). Generalization to CNNs, graph neural nets, or scientific computing is unclear. The 3.4× optimization time reduction is impressive but measured on a single hardware setup (A100 GPU) — would be stronger with CPU/TPU results.

Writing quality: Section 3 (sGraph representation) is dense — the formalism obscures intuition. A running example traced through symbolic generation → pruning → instantiation would clarify. Section 5 (evaluation) front-loads speedup numbers but buries the search space reduction analysis (Table 3) — that’s the real story. The related work undersells how much this borrows from program synthesis (CEGIS, sketch-based synthesis).

Verdict: weak accept — Solid engineering contribution with real speedups, but incremental conceptually and narrow experimental scope. Would be stronger with broader workload coverage and ablation studies on which pruning rules matter most.

Takeaways

Hierarchical search with symbolic pruning: When your optimization space has structure (families of similar solutions), don’t search flat. Build a symbolic layer that represents families, prune at that level, then instantiate. Applies beyond tensor programs — database query optimization, circuit synthesis, configuration tuning all have this structure.

E-graphs for equivalence checking: If you’re generating many program variants and need to deduplicate, e-graphs are faster than execution-based comparison. They prove equivalence symbolically using rewrite rules. Steal this for any domain with algebraic identities (arithmetic, logic, relational algebra).

Two-phase auto-tuning: Separate “what to tune” (symbolic structure) from “how to tune it” (parameter values). Prism’s symbolic layer identifies promising structures; auto-tuning only runs on those. This avoids wasting tuning budget on bad structures. Useful for hyperparameter optimization, compiler flag tuning, system configuration.

Don’t over-generalize: Prism’s pruning rules are hand-crafted for tensor ops and GPU memory hierarchy. The paper doesn’t claim a general symbolic optimizer. If you’re in a different domain, you’ll need domain-specific rules. The framework (symbolic representation + pruning + instantiation) transfers; the rules don’t.