Prism: Symbolic Superoptimization of Tensor Programs

In modern machine learning systems, tensor programs are typically represented as directed acyclic graphs (DAGs), where nodes correspond to tensor operators and edges denote tensors (i.e., multi-dimensional arrays). Existing ML systems generally optimize tensor programs through transformation rules and scheduling templates manually designed by domain experts. However, this approach has two fundamental limitations: first, it requires substantial engineering effort to support new operators or hardware targets; second, manually designed rules explore only a limited portion of the optimization space, failing to capture the combinatorial interactions among algebraic transformations, data layouts, and hardware-specific scheduling decisions. Superoptimization is an approach to automatically discover fast tensor programs without relying on manually specified rules. TASO pioneered this approach by automatically generating graph substitutions and applying them to optimize tensor programs. Mirage applies superoptimization across multiple levels of the GPU execution hierarchy through a unified representation, enabling coordinated algebraic and scheduling transformations. Recently, large language models (LLMs) have been used to generate optimized GPU kernels, and AlphaEvolve has been proposed as a general-purpose superoptimizer leveraging LLMs to guide the search. However, their applicability to optimizing tensor programs has not been explored yet.

Existing enumeration and sampling-based superoptimization methods face scalability bottlenecks when dealing with large or deeply nested programs: the number of candidate programs grows combinatorially with the number of operators and levels in the execution hierarchy, making exhaustive enumeration impractical. Learning-based superoptimizers like AlphaEvolve use learned priors to guide the search, enabling exploration over substantially larger spaces. However, these methods treat the optimization landscape as largely unstructured, which can lead to unstable search behavior and provides limited guarantees on coverage or completeness of the explored program space.

Prism is the first symbolic superoptimizer for tensor programs. Rather than enumerating concrete candidate programs via brute-force enumeration or stochastic sampling, Prism organizes the search space into a two-level hierarchy. At the upper level, it constructs a representation called sGraph, which compactly encodes entire families of tensor programs; at the lower level, each sGraph is instantiated into many concrete implementations. This symbolic formulation enables Prism to explore substantially larger search spaces and discover higher-quality optimizations compared to prior enumeration- and sampling-based approaches for two key reasons. First, sGraph encodes operator semantics, algebraic identities, and hardware constraints as symbolic expressions, allowing Prism to prune provably suboptimal regions of the search space before materializing concrete programs. This structured pruning enables scalability to optimization problems that are intractable for exhaustive enumeration. Second, Prism preserves optimality guarantees: the pruning process is sound and does not eliminate optimal solutions. This property fundamentally distinguishes symbolic superoptimization from prior sampling-based approaches.

• �� Prism employs sGraph, a symbolic hierarchical representation, to achieve superoptimization of tensor programs. sGraph compactly encodes large classes of tensor programs by symbolically representing execution parameters.

• �� Prism organizes optimization as a two-level search: constructing symbolic graphs representing program families and instantiating them into concrete implementations.

• �� Through symbolic reasoning, Prism achieves structured pruning, eliminating suboptimal regions in the search space.

• �� Key techniques include efficient symbolic graph generation, equivalence verification via e-graph rewriting, and parameter instantiation through auto-tuning.

• �� Prism's symbolic formulation allows it to explore larger search spaces and discover higher-quality optimizations.

Prism was evaluated on five commonly used LLM workloads, including fused normalization-linear layers, gated MLPs, and group-query attention. In these experiments, Prism achieved up to 2.2x speedup over the best superoptimizers and 4.9x over the best compiler-based approaches, while reducing end-to-end optimization time by up to 3.4x. The experimental design involved symbolic reasoning to effectively prune the search space, enabling Prism to handle complex optimization problems that traditional enumeration methods cannot solve. The results demonstrate that Prism significantly enhances the execution efficiency of tensor programs, particularly in modern ML workloads.

In experiments, Prism demonstrated significant performance improvements on five commonly used LLM workloads. Compared to the best superoptimizers, Prism achieved up to 2.2x speedup; compared to the best compiler-based approaches, Prism achieved up to 4.9x speedup. Additionally, Prism reduced end-to-end optimization time by up to 3.4x. These performance gains indicate that Prism effectively handles complex optimization problems in modern ML workloads. Through symbolic reasoning, Prism effectively prunes the search space, handling complex optimization problems that traditional enumeration methods cannot.

Prism's application scenarios include tensor program optimization in modern ML workloads. By leveraging symbolic superoptimization, Prism automatically discovers efficient tensor program implementations, reducing reliance on manually designed rules. This decreases engineering effort for supporting new operators or hardware targets and expands the optimization space to capture combinatorial interactions that human intuition might miss. Prism provides an efficient optimization method for modern ML workloads, significantly enhancing the execution efficiency of tensor programs.

Prism may have limitations in handling specific complex hardware constraints, as these constraints might not be effectively expressed and pruned through symbolic reasoning. Prism's performance depends on the efficiency of symbolic graph generation and equivalence verification, which may become bottlenecks in certain scenarios. Prism's symbolic superoptimization approach may not significantly outperform existing methods on certain specific tensor programs. Future directions include extending Prism to support more types of hardware constraints, optimizing the efficiency of symbolic graph generation and equivalence verification, and exploring more symbolic reasoning techniques to further enhance optimization effectiveness.