Prism: Symbolic Superoptimization of Tensor Programs

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$.

• The paper introduces the sGraph representation that encodes equivalence classes of tensor programs for efficient symbolic pruning.
• It integrates symbolic search with mapping instantiation, achieving up to 3.4× faster optimization compared to traditional methods.
• Empirical results show kernel performance improvements up to 4.9× over state-of-the-art approaches in LLM-centric workloads.

Symbolic Superoptimization of Tensor Programs: The Prism (SIGMA) Framework

Introduction and Motivation

Tensor program optimization underpins efficient execution of deep neural networks, especially on modern GPU architectures. Traditional approaches—manual graph rewrite rules, kernel libraries, or template-based scheduling—suffer from high engineering overhead and poor coverage as DNN model and hardware complexity escalate. Enumeration-based superoptimizers (e.g., Mirage, TASO) automate search for equivalent and faster tensor programs, but scalability is bottlenecked by the combinatorial explosion of candidate programs. Sampling-based approaches guided by LLMs or evolutionary strategies can explore larger spaces, but lack structural reasoning and coverage guarantees.

The "Prism: Symbolic Superoptimization of Tensor Programs" paper (SIGMA) (2604.15272) introduces a fundamentally different approach: symbolic superoptimization. The core innovation is the sGraph representation, which encodes broad equivalence classes of tensor programs with operator semantics, transformations, and hardware mappings as symbolic variables. This enables structured, sound pruning across the optimization space, and program synthesis at a scale previously unattainable for ML workloads (such as LLM kernels and fused operator blocks).

sGraph: Symbolic, Hierarchical Representation

The sGraph abstraction advances from Mirage's μGraph by lifting grid/block/loop dimensions and tensor-to-thread mappings from concrete values to symbolic variables. Each sGraph encodes:

• Operator algebra and execution hierarchy (kernel/block/thread graphs);
• Symbolic mappings (Boolean variables for each tensor/data/parallelization dimension combination);
• Symbolic parallelization parameters (grid/block/loop sizes).

This methodology compresses the representation of vast families of concrete tensor programs, and supports symbolic reasoning: pruning and constraint propagation can eliminate entire infeasible subspaces without enumerating specific schedule assignments.

Symbolic Search, Pruning, and Instantiation Pipeline

SIGMA's superoptimization workflow comprises several key innovations:

Symbolic Graph Generation

The generator constructs sGraphs by incrementally adding operators and postponing the assignment of concrete mappings. Weakly valid partial graphs are pruned by two mechanisms:

1. Symbolic Dimension Matching: Shape compatibility constraints are expressed as equations over symbolic mapping variables; only fully congruent symbolic expressions are retained.
2. Expression-Guided Pruning: Additional pruning is achieved by abstracting all parameters to unit sizes (i.e., identities with respect to $\mathbf{d}$), then performing subexpression checks, reducing the search cost while maintaining feasible coverage.

Mapping Instantiation

Only after candidate sGraphs pass symbolic pruning are concrete mappings enumerated, enforcing linear constraints from graph structure and dimension matching. Symmetry-breaking further compresses the search space by considering only canonical mapping permutations, eliminating redundant verification effort.

sGraph Verification

Functional equivalence with the input program is verified over symbolic parallelization parameters via e-graphs [egg], using a carefully designed axiom set that covers associative algebraic identities, commutation/cancellation of parallelization ops, fusion/splitting, and compound parallel reductions. The system leverages efficient equality saturation for scalable equivalence testing, compensating for the lack of completeness in the axiom set with extensive random testing as a fallback.

Parameter Instantiation and Autotuning

Verified sGraphs are lowered to parameterized kernel templates. Concrete parallelization parameters are chosen by random sampling and GPU profiling. This approach parallelizes compilation and avoids the iterative bottleneck of evolutionary search, covering more of the space given fixed compilation resources.

Empirical Results and Numerical Highlights

Evaluation focuses on five LLM-centric fused kernel workloads across RMSNorm, GLU MLPs, SwiGLU, and Attention variants. Across multiple input shapes, SIGMA consistently identifies kernels that outperform prior superoptimization and auto-tuning approaches.

Figure 2: Kernel performance and end-to-end optimization time for SIGMA, Mirage, and TVM (Ansor) across five tensor workloads.

Key quantitative results:

• Kernel performance: Up to $2.2\times$ faster than Mirage (state-of-the-art superoptimizer) and $4.9\times$ over TVM compiler-based search in several fused attention workloads.
• Optimization time: Up to $3.4\times$ reduction versus Mirage, particularly pronounced in workloads like RMSNorm-MLP where concrete enumeration-based search becomes intractable.
• Graph diversity: SIGMA discovers significantly more unique fused program graphs than Mirage, especially in high-dimensional attention where exhaustive concrete enumeration is infeasible.

Furthermore, sGraph symbolic search times (search-only, not counting parameter tuning) are orders-of-magnitude lower than concrete search—often sub-second for moderate workloads, compared to tens or hundreds of seconds for enumeration.

Theoretical and Practical Implications

The adoption of symbolic encodings at the graph and hardware mapping level introduces several key implications:

1. Optimization Scalability: The decoupling of structural graph search from mapping and parameter enumeration supports superoptimization for architectures and workloads previously unattainable due to combinatorial explosion.
2. Optimality Guarantees: The soundness of symbolic pruning and verification ensures that optimal implementations are never prematurely pruned (unlike sampling-based search).
3. Verification Costs: The symbolic approach shifts the bottleneck from search to equivalence checking. However, with e-graph based rewriting and axiomatization, verification scales effectively for practically relevant operator sets.

These characteristics position sGraph-based superoptimization as a promising foundation for new ML compilation systems, especially as attention, normalization, and custom fusion operators continue to grow in complexity and diversity across AI hardware.

Limitations and Future Directions

Despite the advances, several open problems remain:

• Axiom Completeness: While the current axiom set is extensive, it is not formally complete for all equivalences expressible by general tensor program superoptimization. Certain deep or obscure fusion opportunities may remain undiscovered.
• Cost Model Integration: Current parameter instantiation uses random sampling. Integrating learned cost models or hybrid statistical approaches (as in TVM) could further reduce autotuning latency.
• Heterogeneous and Custom Hardware: Adapting symbolic superoptimization to support emerging heterogeneous architectures with more intricate memory hierarchies, communication, and exotic ops (e.g., tensor cores, custom accelerators) will require significant extensions to the operator and mapping models.

Conclusion

SIGMA establishes sGraphs as a powerful abstraction for symbolic superoptimization of tensor programs, enabling both exponential compression of the candidate program space and structured verification and pruning. Empirical results demonstrate clear superiority over prior state-of-the-art both in kernel performance and search efficiency, particularly for fused operators in LLM kernels. Symbolic superoptimization opens a viable path to fully automatic generation of high-efficiency GPU kernels for the next generation of large-scale models and diverse hardware contexts, motivating further efforts in axiom discovery, cost modeling, and generalization to hybrid compute environments.