Prism: Symbolic Superoptimization of Tensor Programs

2026-04-16

AI Summary

Existing tensor program optimization approaches struggle to balance performance and search efficiency, lacking systematic exploration capabilities over large-scale program spaces. This work proposes Prism—the first symbolic superoptimizer for tensor programs—which introduces a novel sGraph symbolic graph representation to encode families of programs and employs a two-stage search framework. In the first stage, Prism performs provably correct pruning based on semantic equivalences, algebraic identities, and hardware constraints; in the second stage, it drives parameter instantiation through automated tuning. Experimental results demonstrate that Prism achieves up to 2.2× speedup over the best existing superoptimizer and outperforms conventional compiler-based methods by up to 4.9× across five mainstream LLM workloads, while reducing end-to-end optimization time by as much as 3.4×.

Technology Category

• Search and Optimization: Sampling/Simulation-based Search
• Machine Learning: Optimization
• Constraint Satisfaction and Optimization: Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines

Abstract

This paper presents Prism, the first symbolic superoptimizer for tensor programs. The key idea is sGraph, a symbolic, hierarchical representation that compactly encodes large classes of tensor programs by symbolically representing some execution parameters. Prism organizes optimization as a two-level search: it constructs symbolic graphs that represent families of programs, and then instantiates them into concrete implementations. This formulation enables structured pruning of provably suboptimal regions of the search space using symbolic reasoning over operator semantics, algebraic identities, and hardware constraints. We develop techniques for efficient symbolic graph generation, equivalence verification via e-graph rewriting, and parameter instantiation through auto-tuning. Together, these components allow Prism to bridge the rigor of exhaustive search with the scalability required for modern ML workloads. Evaluation on five commonly used LLM workloads shows that Prism achieves up to $2.2\times$ speedup over best superoptimizers and $4.9\times$ over best compiler-based approaches, while reducing end-to-end optimization time by up to $3.4\times$.

Problem

• superoptimization
• tensor programs
• symbolic reasoning
• program optimization
• machine learning

Innovation

• symbolic superoptimization
• tensor programs
• sGraph
• e-graph rewriting
• auto-tuning