Prism: 텐서 프로그램의 상징적 슈퍼 최적화

Prism: Symbolic Superoptimization of Tensor Programs

본 논문에서는 텐서 프로그램의 첫 번째 상징적 슈퍼 최적화 도구인 Prism을 소개합니다. 핵심 아이디어는 sGraph로, 이는 몇 가지 실행 파라미터를 상징적으로 표현하여 텐서 프로그램의 광범위한 클래스를 간결하게 표현하는 상징적이고 계층적인 표현 방식입니다. Prism은 최적화를 두 단계의 검색으로 구성합니다. 먼저 프로그램 패밀리를 나타내는 상징적 그래프를 구성하고, 그 다음 이를 구체적인 구현으로 인스턴스화합니다. 이러한 접근 방식은 연산 의미, 대수적 항등식 및 하드웨어 제약 조건에 대한 상징적 추론을 통해 검색 공간의 증명적으로 최적이 아닌 영역을 구조적으로 제거할 수 있도록 합니다. 본 논문에서는 효율적인 상징적 그래프 생성, e-그래프 재작성을 통한 동등성 검증, 그리고 자동 튜닝을 통한 파라미터 인스턴스화 기술을 개발합니다. 이러한 구성 요소들을 결합하여 Prism은 완전 탐색의 엄격함과 최신 머신 러닝 워크로드에 필요한 확장성을 모두 갖추도록 합니다. 널리 사용되는 5가지 LLM 워크로드에 대한 평가 결과, Prism은 최고의 슈퍼 최적화 도구에 비해 최대 2.2배의 속도 향상을, 그리고 최고의 컴파일러 기반 접근 방식에 비해 최대 4.9배의 속도 향상을 달성했으며, 전체 최적화 시간을 최대 3.4배까지 단축했습니다.

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