核心发现
方法论
Prism通过sGraph实现张量程序的符号超优化，sGraph是一种符号化的层次表示法，可以紧凑地编码大类张量程序。Prism采用两级搜索策略：首先构建符号图表示程序族，然后将其实例化为具体实现。通过符号推理，Prism能够在搜索空间中结构化地剪枝，去除可证明的次优区域。其核心技术包括高效的符号图生成、通过e-graph重写进行等价验证，以及通过自动调优进行参数实例化。
关键结果
- 在五个常用的LLM工作负载上，Prism相较于最佳超优化器实现了最高2.2倍的加速，相较于最佳编译器方法实现了最高4.9倍的加速，同时将端到端优化时间减少了最高3.4倍。
- Prism通过符号推理有效地剪枝搜索空间，能够处理传统枚举方法无法解决的复杂优化问题。
- Prism保留了最优性保证，其剪枝过程是可靠的，不会排除最优解。
研究意义
Prism的研究意义在于其能够显著提升张量程序的执行效率，尤其是在现代机器学习工作负载中。通过符号超优化，Prism能够自动发现高效的张量程序实现，减少对手工设计规则的依赖。这不仅降低了支持新操作符或硬件目标的工程成本，还拓展了优化空间，捕捉到人类直觉难以发现的组合交互。
技术贡献
Prism的技术贡献在于其引入了sGraph这一符号化表示法，能够在不具体枚举候选程序的情况下组织搜索空间。相比于现有的枚举和采样方法，Prism能够探索更大的搜索空间，并发现更高质量的优化方案。此外，Prism通过符号推理实现了结构化剪枝，确保了最优性保证。
新颖性
Prism是首个针对张量程序的符号超优化器，其创新之处在于引入了sGraph这一符号化表示法，能够紧凑地编码程序族。与现有的基于枚举和采样的方法相比，Prism通过符号推理实现了更高效的搜索和剪枝。
局限性
- Prism在处理某些特定的复杂硬件约束时可能存在局限性，因为这些约束可能无法通过符号推理有效地表达和剪枝。
- Prism的性能表现依赖于符号图生成和等价验证的效率，这可能在某些情况下成为瓶颈。
- Prism的符号超优化方法在某些特定的张量程序上可能无法显著优于现有方法。
未来方向
未来的研究方向包括扩展Prism以支持更多类型的硬件约束，优化符号图生成和等价验证的效率，以及探索更多的符号推理技术以进一步提升优化效果。
AI 总览摘要
现代AI应用中，机器学习模型在GPU上的高效执行至关重要。然而，现有的优化方法通常依赖于手工设计的规则和调度模板，难以捕捉到复杂的组合交互。Prism通过引入sGraph这一符号化表示法，提供了一种新的符号超优化方法。sGraph能够紧凑地编码大类张量程序，通过符号推理实现结构化剪枝，去除搜索空间中的次优区域。
Prism的优化过程分为两级：首先构建符号图表示程序族，然后将其实例化为具体实现。通过这种方式，Prism能够在不具体枚举候选程序的情况下探索更大的搜索空间，并发现更高质量的优化方案。其核心技术包括高效的符号图生成、通过e-graph重写进行等价验证，以及通过自动调优进行参数实例化。
在实验中，Prism在五个常用的LLM工作负载上实现了显著的性能提升。相较于最佳超优化器，Prism实现了最高2.2倍的加速；相较于最佳编译器方法，Prism实现了最高4.9倍的加速。同时，Prism将端到端优化时间减少了最高3.4倍。这样的性能提升表明，Prism能够有效地处理现代机器学习工作负载中的复杂优化问题。
Prism的研究意义在于其能够显著提升张量程序的执行效率，减少对手工设计规则的依赖。这不仅降低了支持新操作符或硬件目标的工程成本，还拓展了优化空间，捕捉到人类直觉难以发现的组合交互。通过符号超优化，Prism为现代机器学习工作负载提供了一种高效的优化方法。
然而，Prism在处理某些特定的复杂硬件约束时可能存在局限性，其性能表现依赖于符号图生成和等价验证的效率。未来的研究方向包括扩展Prism以支持更多类型的硬件约束，优化符号图生成和等价验证的效率，以及探索更多的符号推理技术以进一步提升优化效果。
深度分析
研究背景
在现代机器学习系统中，张量程序通常以有向无环图（DAG）的形式表示，其中节点对应于张量操作符，边表示张量（即多维数组）。现有的机器学习系统通常通过领域专家手工设计的转换规则和调度模板来优化张量程序。然而，这种方法存在两个主要局限性：首先，支持新操作符或硬件目标需要大量的工程工作；其次，手工设计的规则只能探索有限的优化空间，难以捕捉到代数变换、数据布局和硬件特定调度决策之间的组合交互。超优化是一种自动发现快速张量程序的方法，不依赖于手工指定的规则。TASO通过自动生成图替换并应用于优化张量程序，开创了这一方法。Mirage通过统一表示法在GPU执行层次结构的多个层次上应用超优化，实现了协调的代数和调度变换。最近，大型语言模型（LLM）被用于生成优化的GPU内核，AlphaEvolve则提出了一种通用的超优化器，利用LLM引导搜索。然而，这些方法在优化张量程序方面的应用尚未被探索。
核心问题
现有的基于枚举和采样的超优化方法在处理大规模或深度嵌套的程序时存在扩展性瓶颈：候选程序的数量随着操作符和执行层次的增加而呈指数增长，使得穷举枚举对于大规模或深度嵌套的程序来说不切实际。基于学习的超优化器如AlphaEvolve使用学习的先验知识来引导搜索，能够探索更大的空间。然而，这些方法将优化空间视为基本无结构的，可能导致不稳定的搜索行为，并且对所探索程序空间的覆盖或完整性提供的保证有限。
核心创新
Prism是首个针对张量程序的符号超优化器。与通过暴力枚举或随机采样枚举具体候选程序不同，Prism将搜索空间组织为两级层次。在上层，它构建了一种称为sGraph的符号表示法，紧凑地编码了整个张量程序族；在下层，每个sGraph被实例化为多个具体实现。这种符号化的表述使Prism能够探索更大的搜索空间，并发现比现有的枚举和采样方法更高质量的优化方案。首先，sGraph将操作符语义、代数恒等式和硬件约束编码为符号表达式，使Prism能够在具体程序实现之前剪枝搜索空间中可证明的次优区域。这种结构化剪枝使得对穷举枚举无法解决的优化问题具有可扩展性。其次，Prism保留了最优性保证：剪枝过程是可靠的，不会排除最优解。这一特性从根本上将符号超优化与现有的基于采样的方法区分开来。
方法详解
- �� Prism通过sGraph实现张量程序的符号超优化。sGraph是一种符号化的层次表示法，可以紧凑地编码大类张量程序。
- �� Prism采用两级搜索策略：首先构建符号图表示程序族，然后将其实例化为具体实现。
- �� 通过符号推理，Prism能够在搜索空间中结构化地剪枝，去除可证明的次优区域。
- �� 其核心技术包括高效的符号图生成、通过e-graph重写进行等价验证，以及通过自动调优进行参数实例化。
- �� Prism的符号化表述使其能够探索更大的搜索空间，并发现更高质量的优化方案。
实验设计
Prism在五个常用的LLM工作负载上进行了评估，包括融合归一化-线性层、门控MLP和组查询注意力。实验中，Prism相较于最佳超优化器实现了最高2.2倍的加速，相较于最佳编译器方法实现了最高4.9倍的加速，同时将端到端优化时间减少了最高3.4倍。实验设计中，Prism通过符号推理有效地剪枝搜索空间，能够处理传统枚举方法无法解决的复杂优化问题。实验结果表明，Prism能够显著提升张量程序的执行效率，尤其是在现代机器学习工作负载中。
结果分析
在实验中，Prism在五个常用的LLM工作负载上实现了显著的性能提升。相较于最佳超优化器，Prism实现了最高2.2倍的加速；相较于最佳编译器方法，Prism实现了最高4.9倍的加速。同时，Prism将端到端优化时间减少了最高3.4倍。这样的性能提升表明，Prism能够有效地处理现代机器学习工作负载中的复杂优化问题。Prism通过符号推理有效地剪枝搜索空间，能够处理传统枚举方法无法解决的复杂优化问题。
应用场景
Prism的应用场景包括现代机器学习工作负载中的张量程序优化。通过符号超优化，Prism能够自动发现高效的张量程序实现，减少对手工设计规则的依赖。这不仅降低了支持新操作符或硬件目标的工程成本，还拓展了优化空间，捕捉到人类直觉难以发现的组合交互。Prism为现代机器学习工作负载提供了一种高效的优化方法，能够显著提升张量程序的执行效率。
局限与展望
Prism在处理某些特定的复杂硬件约束时可能存在局限性，因为这些约束可能无法通过符号推理有效地表达和剪枝。Prism的性能表现依赖于符号图生成和等价验证的效率，这可能在某些情况下成为瓶颈。Prism的符号超优化方法在某些特定的张量程序上可能无法显著优于现有方法。未来的研究方向包括扩展Prism以支持更多类型的硬件约束，优化符号图生成和等价验证的效率，以及探索更多的符号推理技术以进一步提升优化效果。
通俗解读 非专业人士也能看懂
想象一下你在厨房里做饭。你有一份食谱，告诉你需要哪些食材、怎么切、怎么煮。传统的方法就像是按照食谱一步步来，但有时候你可能会发现一些更快的方法，比如同时切菜和煮汤。Prism就像是一个聪明的厨师助手，它不仅能按照食谱来，还能自动发现这些更快的方法。它通过一种叫做sGraph的方式，把所有可能的做饭步骤都列出来，然后通过一些聪明的推理，去掉那些不太好的步骤。这样一来，它就能帮你找到最快、最好的做饭方法。Prism的厉害之处在于，它能处理非常复杂的食谱，甚至那些普通厨师都觉得难以应对的。通过这种方式，Prism能让你的厨房工作变得更高效，就像是为你的烹饪加了一个超级加速器。
简单解释 像给14岁少年讲一样
嘿，小伙伴们！你有没有想过，为什么有些游戏加载得特别快，而有些却要等很久？这就像是我们在用一种叫做Prism的超级工具来优化游戏的加载速度。想象一下，你在玩一个超大的游戏，里面有很多角色、场景和特效。Prism就像是一个聪明的游戏助手，它能帮你找到最快的加载方式。它会先把所有可能的加载方法都列出来，然后用一种叫做sGraph的魔法，把那些不太好的方法统统去掉。这样一来，游戏就能以最快的速度加载出来，让你不用再等得心急如焚。Prism的厉害之处在于，它能处理非常复杂的游戏场景，甚至那些普通工具都觉得难以应对的。通过这种方式，Prism能让你的游戏体验变得更流畅，就像是为你的游戏加了一个超级加速器。是不是很酷？
术语表
符号超优化 (Symbolic Superoptimization)
一种通过符号推理自动发现高效程序实现的方法，能够在不具体枚举候选程序的情况下组织搜索空间。
Prism通过符号超优化实现张量程序的高效执行。
sGraph
一种符号化的层次表示法，可以紧凑地编码大类张量程序，通过符号推理实现结构化剪枝。
Prism使用sGraph来表示和优化张量程序。
e-graph重写 (e-graph Rewriting)
一种通过等价类图进行程序等价验证的方法，能够在符号化表示中进行高效的等价检查。
Prism通过e-graph重写进行符号图的等价验证。
自动调优 (Auto-tuning)
一种通过自动化搜索找到最佳参数配置的方法，能够提高程序的执行效率。
Prism通过自动调优进行参数实例化以优化GPU内核。
张量程序 (Tensor Program)
一种用于表示多维数组计算的程序，通常在机器学习和深度学习中使用。
Prism针对张量程序进行符号超优化。
代数恒等式 (Algebraic Identity)
一种数学等式，表示两个表达式在所有情况下都是相等的，常用于优化计算。
Prism利用代数恒等式进行符号推理和剪枝。
硬件约束 (Hardware Constraint)
一种限制程序执行的硬件特性或条件，可能影响程序的性能和优化空间。
Prism在符号推理中考虑硬件约束以进行优化。
GPU内核 (GPU Kernel)
一种在GPU上执行的并行计算单元，通常用于加速大规模数据处理。
Prism通过符号超优化生成优化的GPU内核。
并行化参数 (Parallelization Parameter)
一种用于控制程序并行执行的参数，影响计算的分布和调度。
Prism通过符号化并行化参数进行优化。
等价验证 (Equivalence Verification)
一种验证两个程序在功能上是否等价的方法，确保优化过程的正确性。
Prism通过等价验证确保符号图的正确性。
开放问题 这项研究留下的未解疑问
- 1 如何在符号超优化中更好地表达和处理复杂的硬件约束？现有的方法可能无法有效地捕捉这些约束的细微差别，从而影响优化效果。
- 2 符号图生成和等价验证的效率如何进一步提升？在某些情况下，这些步骤可能成为性能瓶颈，限制了Prism的应用范围。
- 3 如何扩展Prism以支持更多类型的张量程序和操作符？现有的方法可能在某些特定的张量程序上无法显著优于现有方法。
- 4 符号推理技术如何进一步发展，以实现更高效的搜索和剪枝？现有的方法可能在某些情况下无法充分利用符号推理的潜力。
- 5 如何在符号超优化中更好地结合机器学习技术，以提高优化效果？现有的方法可能在某些情况下无法充分利用机器学习的优势。
应用场景
近期应用
机器学习模型优化
Prism可以用于优化机器学习模型中的张量程序，提高模型的执行效率，减少训练和推理时间。
GPU加速
通过生成优化的GPU内核，Prism可以显著提升大规模数据处理的速度，适用于科学计算和数据分析。
自动化软件优化
Prism可以用于自动化软件优化，减少对手工设计规则的依赖，提高软件的性能和可扩展性。
远期愿景
通用计算优化
Prism的符号超优化方法可以扩展到更广泛的计算领域，实现通用计算的自动化优化。
智能系统开发
通过结合机器学习技术，Prism可以用于开发更智能的系统，实现自动化的程序优化和性能提升。
原文摘要
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$.
参考文献 (20)
egg: Fast and extensible equality saturation
Max Willsey, Chandrakana Nandi, Y. Wang 等
AlphaEvolve: A coding agent for scientific and algorithmic discovery
Alexander Novikov, Ngân V˜u, Marvin Eisenberger 等
Mirage: A Multi-Level Superoptimizer for Tensor Programs
Mengdi Wu, Xinhao Cheng, Shengyu Liu 等
TASO: optimizing deep learning computation with automatic generation of graph substitutions
Zhihao Jia, Oded Padon, James J. Thomas 等
Ansor : Generating High-Performance Tensor Programs for Deep Learning
Lianmin Zheng, Chengfan Jia, Minmin Sun 等
Automatically scheduling halide image processing pipelines
Ravi Teja Mullapudi, Andrew Adams, Dillon Sharlet 等
Beyond Data and Model Parallelism for Deep Neural Networks
Zhihao Jia, M. Zaharia, A. Aiken
KernelFoundry: Hardware-aware evolutionary GPU kernel optimization
Nina Wiedemann, Quentin Leboutet, Michael Paulitsch 等
EINNET: Optimizing Tensor Programs with Derivation-Based Transformations
Liyan Zheng, Haojie Wang, Jidong Zhai 等
PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections
Haojie Wang, Jidong Zhai, Mingyu Gao 等
Unity: Accelerating DNN Training Through Joint Optimization of Algebraic Transformations and Parallelization
Colin Unger, Zhihao Jia, Wei Wu 等
NVIDIA Tensor Core Programmability, Performance & Precision
S. Markidis, Steven W. D. Chien, E. Laure 等
Superoptimizer: a look at the smallest program
H. Massalin
TVM: End-to-End Optimization Stack for Deep Learning
Tianqi Chen, T. Moreau, Ziheng Jiang 等
GraphPipe: Improving Performance and Scalability of DNN Training with Graph Pipeline Parallelism
Byungsoo Jeon, Mengdi Wu, Shiyi Cao 等
Astra: A Multi-Agent System for GPU Kernel Performance Optimization
Anjiang Wei, Tianran Sun, Yogesh Seenichamy 等
Transformer
Mukund R. Patel
Equality Saturation for Tensor Graph Superoptimization
Yicheng Yang, Mangpo Phitchaya Phothilimtha, Y. Wang 等
TensorFlow: A system for large-scale machine learning
Martín Abadi, P. Barham, Jianmin Chen 等