arXiv:2605.25954v1 [cs.LG] 25 May 2026
Step-TP: A Grounded, Step-Level Dataset with Chain-of-Thought Reasoning for LLM-Guided Tensor Program Optimization Mengfan Liu
Da Zheng
The University of Hong Kong Hong Kong, China [email protected]
Ant Group Hangzhou, China [email protected]
Junwei Su∗
Chuan Wu∗
University of Science and Technology of China Hefei, China [email protected]
The University of Hong Kong Hong Kong, China [email protected]
Abstract Despite the strong reasoning capabilities of large language models (LLMs), optimizing the execution efficiency of tensor programs remains challenging due to the need for precise, composable transformation decisions. Recent LLM-guided approaches frame tensor program optimization as an iterative decision process, but existing datasets provide only end-to-end optimized program pairs using token-inefficient representations, lacking verifiable step-level supervision and interpretability. As a result, LLMs struggle to make reliable single-step decisions in large combinatorial optimization spaces. We introduce Step-TP, a post-training dataset for tensor program optimization that provides grounded, atomic, step-level supervision with structured chain-of-thought (CoT) reasoning. Step-TP forms a closed reasoning loop over intermediate program states, enabling reliable multi-step optimization rather than outcome imitation. Its design is guided by four principles: (i) a token-efficient, verifiable intermediate representation (IR) that deterministically lowers to TVM TIR; (ii) atomic and composable optimization strategies that decompose complex trajectories into interpretable single-step decisions; (iii) structured CoT supervision coupled with explicit IR-to-IR state transitions; and (iv) strategy filtering to balance coverage while preventing shortcut exploitation. The dataset and implementation are available at a GitHub link https://github.com/LIUMENGFAN-gif/StepTP.
CCS Concepts • Computing methodologies → Artificial intelligence.
Keywords LLM-based Tensor Program Optimization, Tensor Program, Dataset, Chain-of-Thought Reasoning
1
Introduction
Background and Motivation. Efficient execution of deep neural networks on GPUs is fundamentally a problem of tensor program optimization, in which high-level mathematical operators must be lowered into kernels that effectively exploit the GPU’s massive parallelism and hierarchical memory system [1, 37, 38, 44, 51, 55]. Although modern GPUs provide extraordinary peak throughput, ∗ Corresponding authors.
realizing this performance in practice requires carefully coordinating computation, data movement, and parallel scheduling through techniques such as loop tiling, memory hierarchy selection, operator fusion, and thread–block binding. As model architectures become deeper, wider, and increasingly irregular, vanilla kernel implementations often lead to poor hardware utilization and excessive memory traffic, making performance highly sensitive to low-level code generation decisions [6, 17, 20, 22, 33, 36]. This challenge is further amplified in contemporary workloads involving large language models (LLMs) and foundation models [2, 25], where even minor tensor-program-level inefficiencies accumulate into substantial increases in end-to-end latency, energy consumption, and deployment cost at scale. Consequently, efficient GPU execution has emerged as a central enabling factor for scalable, cost-effective, and sustainable deep learning, motivating the development of principled and automated tensor optimization techniques that can systematically reason about GPU execution behavior. Motivated by these challenges, recent work has explored the use of LLMs as decision-makers in tensor program optimization, leveraging their capacity to reason over structured program representations and long-range dependencies in optimization sequences [4, 12, 16, 24, 26, 32, 35, 43, 47]. Tensor optimization entails navigating a vast, discrete, and highly nonconvex search space composed of different-level decisions (e.g., graph, operator, memory, math levels), where effective strategies often depend jointly on local program structure and global execution context [3, 7, 14, 18, 18, 39, 44, 53, 54, 56]. LLMs are particularly well suited to this setting because they naturally model optimization as a sequential decision process, integrating symbolic program information with high-level optimization intent to produce interpretable, step-wise transformations [4, 8, 12, 16, 35]. In contrast to black-box search methods or purely statistical cost models, LLM-based approaches offer the potential to incorporate domain knowledge, generalize across workloads and hardware backends, and reuse learned optimization patterns across programs. As tensor programs and accelerator architectures continue to grow in complexity, LLM-guided optimization provides a promising direction toward more flexible and data-efficient tensor optimization frameworks. Pressing Need for Post-Training Datasets. Despite recent progress, the practical effectiveness of LLM-based tensor optimization remains constrained without targeted post-training on domain-specific
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
data [8, 9, 43]. Although general-purpose LLMs exhibit strong reasoning and pattern recognition capabilities, they are not trained to manipulate low-level tensor programs under the strict correctness, hardware, and performance constraints required by real-world GPU execution. In the absence of post-training data that explicitly encodes these requirements, LLMs tend to fall back on superficial pattern matching or outcome imitation, resulting in brittle behavior, poor generalization, and limited interpretability [5, 43, 45, 47]. Meanwhile, reinforcement learning or search-based fine-tuning approaches are prohibitively expensive in this domain due to the high cost of compiling and benchmarking candidate programs [4, 24, 34]. Consequently, a carefully constructed dataset for post-training—one that provides grounded intermediate program states, atomic optimization actions, and verifiable state transitions—becomes essential for enabling data-efficient learning, stable reasoning, and systematic evaluation. Such datasets are critical for transforming LLMs from coarse heuristic generators into reliable optimization agents capable of reasoning about tensor programs at scale. Desired Properties of Post-Training Datasets. To be effective, such post-training datasets must do more than simply collect optimized programs—they must be explicitly designed to support step-wise reasoning and iterative decision making in tensor optimization. Tensor optimization is inherently an iterative decision process rather than a single-shot prediction task: it requires precise, step-wise reasoning over program representations, where each transformation must be valid, composable, and compatible with downstream steps [7, 15, 19, 40–42, 49–51]. Accordingly, effective post-training for LLM-based tensor optimization demands datasets that (i) expose step-level supervision instead of only final outcomes, (ii) provide faithfully grounded reasoning traces aligned with executable program transformations, and (iii) cover a diverse set of optimization strategies representative of real-world workloads. In addition, as the inference and reasoning capacity of LLMs is constrained by context length, post-training datasets should (iv) encode optimization processes in a context-efficient representation that supports effective reasoning within practical prompt-length limits. Limitations of Existing Datasets. Despite growing interest in this direction, existing datasets fail to jointly satisfy the desired requirements in several important respects. First, most existing datasets rely primarily on outcome-only supervision [28, 46–48, 52], providing only final high-performance tensor programs produced through complex compositions of multiple optimization strategies. Such supervision encourages LLMs to memorize surface patterns in optimized code rather than to internalize the underlying decisionmaking process, leading to weak reasoning ability, poor generalization to unseen programs, and limited capacity to explore novel strategy compositions. Second, existing datasets such as ConCuR [21] adopt low-level CUDA or Python code as the primary representation space for optimization. While expressive, these representations are verbose and poorly suited for compactly encoding optimization intent and intermediate program states, resulting in excessively long descriptions that hinder effective reasoning within the limited context length of LLMs [5, 23, 27]. Third, existing datasets (e.g., LOOPerSet [28], IR-OptSet [46], and ConCuR [21]) exhibit limited strategy diversity, focusing mainly on easily modularized transformations such as loop tiling while largely omitting more sophisticated mathematical optimizations (e.g., online softmax [29]).
Because strategy diversity directly determines the effective optimization search space, this narrow coverage constrains an LLM’s ability to reason about, compose, and generalize high-performance solutions for real-world tensor programs [10, 11, 13, 18, 44, 54]. Together, these limitations highlight a critical need for constructing datasets for LLM-based tensor program optimization that are representation-efficient, support grounded, step-level supervision, and enable reliable reasoning over a diverse set of optimization strategies. Contributions. To address these limitations, we introduce Step-TP, a post-training dataset for LLM-based tensor program optimization that provides grounded, atomic, step-level supervision with structured chain-of-thought (CoT) reasoning across diverse tensorprogram–level optimization strategies. A comparison between StepTP and existing datasets is shown in Table 1. The main contributions of this paper are twofold: (1) Design of an Effective Intermediate Representation (IR). We propose LEIR, a verifiable and token-efficient intermediate representation tailored for LLM-based tensor optimization. The LEIR supports compact and precise expression of intermediate program states, enables seamless application of atomic optimization strategies, and can be deterministically converted to TVM TIR, providing semantic grounding and correctness verification. (2) Construction of a Step-Level Post-Training Dataset. Building on the proposed IR, we construct Step-TP, a post-training dataset for LLM-based tensor optimization that incorporates: (i) a systematic decomposition of complex optimization trajectories into atomic, composable strategies, transforming a large and opaque search space into interpretable single-step decisions; (ii) structured CoT supervision that couples strategy-level rationale with explicit IR-to-IR state-transition mappings; and (iii) a strategy filtering mechanism based on preconditions, parameters, and synthesis depth to balance strategy distribution, ensure broad coverage, and prevent shortcut exploitation. We further conduct an extensive empirical study of multi-step optimization via step-level guidance. Our results show that Step-TP enables effective step-level guidance, empowering diverse search paradigms to achieve strong performance with remarkable efficiency. Our results demonstrate that this guidance allows models to generate executable, grounded transformations across a diverse set of strategies and can navigate long-horizon optimization trajectories across various GPU architectures.
2
Design of Intermediate Representation (IR)
This section presents the design of our IR for tensor-program-level transformations. We begin by examining why existing program representations are ill-suited for learning and reasoning about transformation logic, and distill a key structural insight from this analysis. This insight motivates a high-density loop–equation representation that covers full tensor program optimization space, which we formalize as LEIR and illustrate through a concrete matrix multiplication case study.
2.1
Limitations of Existing IR.
The reasoning capability of LLMs is constrained by finite context length, making representation efficiency a first-order concern. To
Step-TP Dataset for Tensor Program Optimization
Dataset LOOPerSet [28] TenSet [52] Tlp [48] TpuGraphs [31] IR-OptSet [46] ConCur [21] Step-TP (Ours)
Task Polyhedral compiler optimization Cost model Cost model Cost model Tensor program optimization Tensor program optimization Tensor program optimization
Target Platform CPU/GPU CPU/GPU CPU/GPU TPU CPU GPU GPU
Executable IR/Program × ⃝ ⃝ ⃝ ✓ ✓ ✓
CoT × × × × × ✓ ✓
Strategy-Driven × × × × × × ✓
Step-level Supervision × × × × × × ✓
Table 1: Representative Datasets for Tensor Programs. Our dataset Step-TP is the only post-training dataset for LLM-based tensor program optimization that provides grounded, atomic, step-level supervision with structured CoT across diverse optimization strategies. : loop part : equation part green: memory type; pink: loop notation; red: loop range; purple: data type; sky blue: loop index; orange: element-wise computation; black: variable extern "C" __global__ void __launch_bounds__(719) main_kernel(double* __restrict__ A, double* __restrict__ C, double* __restrict__ D) { for (int a = 0; a < 549; ++a) { for (int c = 0; c < 2252; ++c) { for (int d = 0; d < 2193; ++d) { if (d == 0) { D[(((((int)threadIdx.x) * 1236348) + (a * 2252)) + c)] = 0.0; } D[(((((int)threadIdx.x) * 1236348) + (a * 2252)) + c)] += (A[(((((int)threadIdx.x) * 1203957) + (a * 2193)) + d)] * C[((d * 2252) + c)]);}}}}
(a) Core part of CUDA def main(A: T.Buffer((719, 549, 2193), "float64"), C: T.Buffer((2193, 2252), "float64"), D: T.Buffer((719, 549, 2252), "float64")): T.func_attr({"tir.noalias": True}) # with T.block("root"): for tx in T.thread_binding(719, thread="threadIdx.x"): for a,c,d in T.grid(549,2252,2193): with T.block("D"): v_tx,v_c,v_a,v_d=T.axis.remap("SSSR",[tx,c,a,d]) T.reads(A[v_tx,v_a,v_d],C[v_d,v_c]) T.writes(D[v_tx,v_a,v_c]) with T.init(): D[v_tx,v_a,v_c]=T.float64(0) D[v_tx,v_a,v_c]=D[v_tx,v_a,v_c]+A[v_tx,v_a,v_d]*C[v_d,v_c]
(b) TIR B^{719}_{tx=0}L^{549}_{a=0}L^{2252}_{c=0}L^{2193}_{d=0} [D^{f64,g}_{tx,a,c}=D^{f64,g}_{tx,a,c}+A^{f64,g}_{tx,a,d}*C^{f64,g}_{d,c};];
(c) Our LEIR
Figure 1: Comparison of the same tensor program representation among CUDA, TVM TIR, and LEIR. As illustrated, our LEIR provides a more efficient representation for matrix multiplication (𝐷 = 𝐷 + 𝐴 × 𝐶) and its associated loop structure. enable effective modeling of program transformation logic, a tensor program representation should therefore be semantically dense and minimize entanglement with transformation-irrelevant details. Mainstream representations such as CUDA and TVM TIR [15] are executable and compiler-oriented, but they introduce substantial syntactic and structural noise—such as type annotations, controlflow scaffolding, and compiler-mandated boilerplate—that obscure the core transformation logic (e.g., loop restructuring and algebraic fusion). As a result, LLMs are overwhelmed with implementation artifacts that are orthogonal to tensor-program-level reasoning. We illustrate these limitations using a matrix multiplication example. CUDA. As an explicit, hardware-oriented imperative model, CUDA prioritizes fine-grained control over GPU execution, thereby causing high-level transformation logic to be scattered across fragmented, implementation-specific constructs. As illustrated in Figure 1(a), explicit type annotations in loop indices (e.g., int a,
((int)threadIdx.x)) embed formatting details within the loops and index arithmetic. Meanwhile, manual initialization via control flow (e.g., if (d == 0)) structurally separates the initialization of a reduction from its accumulation update. Although executed within the same loop nest, this separation breaks the structural coherence of the reduction, complicating the identification of the canonical matrix multiplication pattern (𝑒.𝑔., 𝐷 = 𝐷 +𝐴 ×𝐶) as a unified transformation unit for learning-based models. Furthermore, CUDA typically entangles the computation with micro-architectural designs such as memory bank-conflict avoidance. For instance, padding shared memory introduces intricate index offsets that obscure the logical iteration space, thereby introducing optimization concerns that are orthogonal to tensor-program-level transformations. TVM TIR. Compared to CUDA, TIR offers a more structured representation aligned with tensor-program-level optimizations. However, as a compiler-oriented IR, TIR imposes a heavy burden of declarative boilerplate. As shown in Figure 1(b), even a standard matrix multiplication is encased within extensive metadata (e.g., T.reads) and explicit axis-remapping mechanisms (e.g., T.axis.remap), which introduce substantial repetition without adding new semantic value to the algebraic computation. Moreover, the core 𝐷 = 𝐷 + 𝐴 × 𝐶 logic is buried under multiple layers of syntactic scaffolding, such as the nested T.block and T.init scopes. These constructs, while essential for compiler correctness, create a high degree of structural depth that weakens the visibility of the fundamental transformation intent for learning-based models. Therefore, neither CUDA nor TIR provides an efficient format for constructing transformation datasets.
2.2
Loop-Equation IR (LEIR).
To address these limitations, we examine the essential tensor-programlevel structure common to both CUDA and TIR. After abstracting away low-level execution details and compiler boilerplate, both representations reduce to two irreducible semantic components: (1) a loop structure defining the iteration space, and (2) an equation structure specifying the algebraic computation at each iteration, where these two components still cover the full tensor-programlevel optimization space detailed in Appendix B, in contrast to prior abstractions (e.g., EINNET [53]) that employ non-unified loop representations (hindering operator-level transformations like loop binding) and are restricted to summation-based computations. This observation motivates an IR that preserves only these two components, yielding a representation that is both semantically dense and suitable for step-level learning.
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
LEIR design principles. Building on this insight, we propose LEIR, a high-density representation that balances structural parsimony with the fidelity required to capture tensor-program-level transformations. Our design is guided by three core principles: (1) Irreducible semantic minimality. Unlike existing IRs that mandate extensive metadata for compiler analysis, LEIR distills the program representation into its minimal semantic components. It consolidates the fragmented constructs of CUDA and the multi-layered scaffolding of TIR into just two irreducible structures (i.e., nested loops and algebraic equations), thereby significantly enhancing the semantic density. This design ensures that the majority of tokens in the representation correspond directly to meaningful elements of the optimization space, rather than to syntactic overhead. (2) Explicit organization logic. Explicit organization logic. To maintain expressiveness, LEIR avoids the pitfall of excessive abstraction, such as representing programs solely as optimization parameters without the entire organizational logic of the computation. Instead, it explicitly preserves the structural hierarchy of execution. Specifically, the sequential order of loop descriptors and equations captures the execution flow. Meanwhile, the mapping of iteration spaces to logical execution levels (e.g., thread-block binding) is embedded within loop descriptors, and the assignment of data to memory hierarchies is encoded into tensor variables. This design ensures that the underlying computational pattern remains intact and reconstructible, allowing the learning-based models to reason about the spatial and temporal organization of the computation. (3) Parseable syntax. LEIR adopts a LaTeX-based syntax to represent the tensor programs, which leverages the prior knowledge of LLMs to enable direct parsing and reasoning. Together, these principles enable LEIR to provide a concise yet expressive enough representation to capture tensor-program-level optimizations, while remaining interpretable by LLMs. See Appendix A for the complete grammar definition. Case Study. As illustrated in Figure 1(c), we exemplify the design of LEIR through a matrix multiplication case. A typical tensor program in LEIR consists of one or more expressions separated by semicolons, with each expression comprising a loop part (in yellow) and an equation part (in green). To ensure brevity, LEIR employs implicit initialization to maintain a concise algebraic flow. Loop structure. The loop structure is represented by a main symbol indicating the loop type, with superscripts and subscripts specifying the loop index and iteration range. LEIR supports various loop types: serial loops (𝐿), parallel loops (𝑃), vectorized loops (𝑉 ), unrolled loops (𝑈 ), and thread/block-binding loops (𝐵). Notably, indices for binding loops are mapped to CUDA intrinsics: {𝑏𝑥, 𝑏𝑦, 𝑏𝑧} 719 defor blockIdx and {𝑡𝑥, 𝑡𝑦, 𝑡𝑧} for threadIdx. In this case, 𝐵𝑡𝑥=0 notes the outermost loop bound to threadIdx.x with a range of 549 represents a nested serial loop with a range of 549. 719, while 𝐿𝑎=0 Equation structure. The equation part specifies the computation performed under the given loop nest and consists of three elements: element-wise computation, tensor variables, and delimiters to define the computational scope and logical sequence. (1) Element-wise computation. We formalize the algebraic computations using a set of functional operators derived from TVM TIR,
including standard arithmetic (e.g., +, −), transcendental functions (e.g., exp, log), and conditional intrinsics (e.g., if_then_else), all applied in a purely element-wise manner. The example illustrates the core matric multiplication operation: 𝐷 = 𝐷 + 𝐴 × 𝐶. (2) Variable. Each variable is defined by an identity symbol (e.g., 𝐷), with subscripts for indices and superscripts for metadata. The metadata includes the data type (e.g., 𝑓 64 for float64) and memory hierarchy (e.g., 𝑔 for global, 𝑠 for shared, 𝑙 for local memory). For 𝑓 64,𝑔 example, 𝐷𝑡𝑥,𝑎,𝑐 denotes a double-precision tensor variable stored in global memory. (3) Delimiters. Two types of delimiters are employed to organize the program structure. Specifically, the square brackets ([ and ]) in the example bind the computation logic to the four-level loop nest. The internal semicolon (;) marks the completion of the expression, ensuring a clear logical sequence for operations. Implicit Initialization. To maintain an uninterrupted algebraic flow, LEIR employs implicit initialization for common reduction patterns. The identity element is automatically inferred from the operator: summations default to 0, products to 1, and extremes 𝑓 64,𝑔 (max/min) to ±∞. Consequently, the accumulator 𝐷𝑡𝑥,𝑎,𝑐 is initialized to 0 without requiring explicit code, maintaining a concise algebraic representation.
3
Dataset Construction
To construct a high-quality dataset for grounded, step-level supervision of tensor program transformation with reliable reasoning, we design a multi-stage pipeline as shown in Figure 2. Four stages are included: (i). The PyTorch-to-LEIR Translator captures the computation in a PyTorch program and then converts it into our LEIR. (ii). The One-step Strategy-driven transformation stage derives the applicable strategy set for each IR, applies strategies individually to perform single-step transformations, and generates a corresponding reasoning trace for each transformed LEIR to explain the optimization rationale. (iii). The transformed LEIR is then lowered to TVM TIR via an LEIR-to-TIR Translator, enabling execution with a mature compiler backend. (iv). Finally, the Verification and Filtering stage validates the correctness and semantic equivalence of transformed LEIRs, and applies a designed filtering mechanism to regulate the strategy distribution in the final dataset. While the translators provide the necessary infrastructure, the Transformation and Filtering stages constitute the core mechanisms, ensuring the reliability of reasoning traces and the data quality. This section is organized as follows: Sec. 3.1 outlines the overarching dataset composition; Sec. 3.2 details the transformation stage; Sec. 3.3 describes the specific data formats utilized for archival and training; and Sec. 3.4 covers verification and filtering mechanisms.
3.1
Dataset Composition
The dataset comprises source PyTorch programs constructed through a hierarchical approach. We first establish two fundamental building blocks: (1). single-operator programs, mainly based on KernelBench [30] level-1 dataset, covering computational backbones (e.g., matrix multiplication), nonlinear activations, lightweight operations (e.g., transpose), normalization and pooling, and common loss functions. (2) popular architectures, such as Matrix Multiplication and Softmax pipelines and attention modules (e.g., multi-head
Step-TP Dataset for Tensor Program Optimization
One-Step Strategy-Driven Transformation PyTorch Code PyTorch-to-LEIR Translator
Feasibility check Original LEIR
One-step transformation
CoT construction
One-step Applicable transformation strategy set
CoT construction
Verification and Filtering
Dataset
Filtered and verified dataset Filtering Verified dataset Verification
Dataset with LEIR pairs, strategy, and CoT
Chat-template data entry Single-answer version Multi-answer version
LEIRbased dataset
LEIR-to-TIR Translator
Dictionary-based data entry Original LEIR/TIR/CUDA Applied strategy Transformed LEIR/TIR/CUDA CoT
Figure 2: Pipeline of dataset construction. attention, multi-group attention). Based on these building blocks, we further construct composite programs by sampling and assembling 2-5 components from the aforementioned categories, covering the majority of KernelBench [30] level-2 programs and additional randomly composed cases. To further enhance data diversity, we randomize the input and output shapes (e.g., with dimension sizes up to 16,384), and vary data types (e.g., float16, float32, and float64). By instantiating the 189 distinct program types with these varied configurations, we ultimately produce a diverse dataset comprising 6,335 unique PyTorch programs, providing comprehensive coverage of representative tensor computation workloads. Based on the dataset composition, we employ a PyTorch-to-LEIR translator to convert these PyTorch programs into our LEIR. Since PyTorch operations abstract away low-level execution details (e.g., loop nesting), we align the underlying program structures of our IR with the corresponding vanilla TVM TIR implementations, ensuring correctness, semantic equivalence, and executability.
3.2
One-step Strategy-driven Transformation
This subsection details the one-step strategy-driven transformation stage, which shifts from traditional end-to-end mapping to step-level supervision with reliable reasoning. Two benefits are included: (i) step-level supervision guides the LLM through individual transformation, reducing learning difficulty and enabling generalization to different combinations of optimizations; (ii) the strategy-driven design ensures the independence and composability of each transformation strategy, avoiding the lack of interpretability in end-to-end optimizations. To implement this stage, we employ a three-phase transformation process: a feasibility check of strategy preconditions, a one-step transformation to generate target LEIR, and a CoT construction to trace the transformation logic. Feasibility Check. To ensure the validity of the generated candidate, each tensor program undergoes a feasibility check to identify applicable transformation strategies. We define nine essential preconditions, such as pattern-matching checks (e.g., identifying softmax for online softmax), with all preconditions and their corresponding strategy mappings in Appendix B. By checking these preconditions, we establish a set of feasible strategies for each LEIR. One-step Transformation. Based on the feasible strategy set for each original LEIR, we generate one-step transformed LEIRs via a decompose-modify-reassemble workflow to control the scope of modifications and improve reproducibility. The decomposition follows the structure of our LEIR, enabling strategies at different levels (i.e., graph, operator, memory, and mathematical) to modify only
the relevant components before reassembling them into a transformed LEIR. For example, the log simplification strategy works solely on the equations, while the loop reorder strategy acts only on the loops. When a strategy admits multiple valid outcomes (e.g., different loop split factors), we randomly sample one variant to enhance the diversity of transformations covered by each strategy. CoT Construction. After applying the strategy, a corresponding reasoning trace is synthesized to formalize the underlying transformation logic. As shown in Fig. 3(a), each reasoning trace comprises two components: (i) a brief description of the applied strategy to provide a high-level semantic anchor for the transformation, and (ii) an instance-specific explanation to delineate the targeted expressions and components (e.g., loop or equation segments) and document their reassembly into the resulting modified expressions. This structured CoT design yields several benefits for learning and interpretability. First, it improves the LLM’s understanding of program representations and enhances generalization to unseen LEIRs by explicitly constructing the reasoning with the underlying LEIR structure. Second, it facilitates the activation of tensor program optimization knowledge acquired during pretraining by presenting transformations in a strategy-centric and interpretable form. Third, the reasoning traces are directly derived from the actual transformation process rather than post-hoc rationales, and can be explicitly mapped to different-level modifications, ensuring full transparency and traceability. After transformation, the LEIR-to-TIR translator maps both the original and transformed LEIRs into TVM TIRs, leveraging the TVM backend to generate executable CUDA kernels for end-to-end performance evaluation.
3.3
Dataset Format
After the IR-to-TIR translation, we organize our data into two distinct formats: a comprehensive dictionary-based repository for archival purposes and a chat-template dataset for model training. Dictionary-based Dataset. In this format, each data entry contains the original and transformed versions of the LEIR, TIR, and CUDA code, along with the applied strategy and the corresponding CoT. While this dictionary supports various program representations for future research, this paper mainly focuses on our LEIR. Chat-template Dataset. Based on the dictionary repository, we construct a chat-template dataset, containing two specialized variants for training: (1).Single-answer format pairs an original LEIR
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
Give the following IR of Matmul3D: 'B^{719}_{tx=0}L^{549}_{a=0}L^{2252}_{c=0}L^{2193}_{d=0}[D^{f64,g}_{tx,a,c}=D^{f64,g} _{tx,a,c}+A^{f64,g}_{tx,a,d}*C^{f64,g}_{d,c};];', where the known variables are 'A' with the dtype torch.float64 and shape [719, 549, 2193], 'C' with the dtype torch.float64 and shape [2193, 2252], and 'D' with the dtype torch.float64 and shape [719, 549, 2252]. Do not change the names, shapes or dtypes of these known variables in the IR. Please give me a numerically equivalent transformed IR that produces exactly the same outputs for any floating‐point inputs (bitwise identical), and also provide applied strategy in this transformed IR. Return the answer only as a valid JSON object with the following keys:'transformed_IR', 'applied_strategies'. <think>The strategy loop reorder is used on the given IR to rearrange the nesting orders of loops within a loop nest. Under the given IR, the loops 'B^{719}_{tx=0}L^{549}_{a=0}L^{2252}_{c=0}L^{2193}_{d=0}' in the expression 'B^{719}_{tx=0}L^{549}_{a=0}L^{2252}_{c=0}L^{2193}_{d=0}[D^{f64,g}_{tx,a,c}=D^{f64,g} _{tx,a,c}+A^{f64,g}_{tx,a,d}*C^{f64,g}_{d,c};];' can be reordered as the loops 'B^{719}_{tx=0}L^{2252}_{c=0}L^{549}_{a=0}L^{2193}_{d=0}' in the expression 'B^{719}_{tx=0}L^{2252}_{c=0}L^{549}_{a=0}L^{2193}_{d=0}[D^{f64,g}_{tx,a,c}=D^{f64,g} _{tx,a,c}+A^{f64,g}_{tx,a,d}*C^{f64,g}{d,c};];'.</think> <answer>{'transformed_IR': 'B^{719}_{tx=0}L^{2252}_{c=0}L^{549}_{a=0}L^{2193}_{d=0}[D^{f64,g}_{tx,a,c}=D^{f64,g} _{tx,a,c}+A^{f64,g}_{tx,a,d}*C^{f64,g}_{d,c};];', 'applied_strategy': 'loop_reorder'}</ answer>
(a) Single-answer Give the following IR of Matmul3D: ...(the same IR information as (a))... Please give me at least 2 different numerically equivalent transformed IRs that produces exactly the same outputs for any floating‐point inputs (bitwise identical), and also provide applied strategy for each transformed IR. Return the answer list only as a valid JSON object, and each entry with the following keys: 'idx', 'transformed_IR', 'applied_strategies'. <think>These 2 transformed IRs can be individually analyzed as follows: 0. ...(analysis removed due to space limitation)... 1. ...(analysis removed due to space limitation)...</think> <answer>[{'idx': 0, 'transformed_IR': 'B^{719}_{tx=0}L^{9}_{h=0}L^{2}_{i=0}L^{61}_{a=0}L^{1126}_{c=0}L^{2193}_{d=0}[D^{f6 4,g}_{tx,h*61+a,i*1126+c}=D^{f64,g}_{tx,h*61+a,i*1126+c}+A^{f64,g}_{tx,h*61+a,d}*C^ {f64,g}_{d,i*1126+c};];', 'applied_strategy': 'loop_tiling'}, {'idx': 1, 'transformed_IR': 'B^{719}_{tx=0}L^{549}_{a=0}L^{4}_{h=0}L^{563}_{c=0}L^{2193}_{d=0}[D^{f64,g}_{tx,a,h *563+c}=D^{f64,g}_{tx,a,h*563+c}+A^{f64,g}_{tx,a,d}*C^{f64,g}_{d,h*563+c};];', 'applied_strategy': 'loop_split'}]</answer>
(b) Multiple-answer
Figure 3: Single-answer and multi-answer examples for Matmul3D in the chat-template dataset.
randomized input tensors and compare the outputs against the baseline. Only programs that exhibit consistent numerical equivalence across all trials are retained for the final dataset. Filtering. While the verification ensures functional correctness, it does not guarantee a high-quality distribution of transformation patterns. Without explicit control over strategy distribution, LLMs tend to disproportionately favor simplistic transformations. This bias stems from an asymmetry between simple and complex strategies. Simple strategies (e.g., log simplification) are frequently encountered during pre-training and can be activated with minimal supervision, resulting in lower predictive entropy and higher generation confidence. In contrast, complex strategies (e.g., loop split with index remapping and range recalculation) exhibit higher variability and structural diversity, which increases predictive uncertainty and causes the model to systematically avoid them. To mitigate this bias and encourage the LLMs to master sophisticated reasoning, we implement a difficulty-aware rebalancing approach. Strategy difficulty formulation. To operationalize this rebalancing, we formalize the difficulty of each strategy along three dimensions: • Preconditions (𝐾): the number of essential constraints identified during the feasibility check in Sec. 3.2; • Parameter modification (𝑃): the number of existing components modified in the original IR, covering six aspects (i.e., expressions, loop axes, range adjustments, equations, variables, and index calculation segments); • Synthesis depth (𝑆): the number of unique categories of newly introduced elements (including new variables, index calculations, and expressions). Each category contributes a fixed value of 1 to the depth score. These three dimensions reflect an ascending complexity, ranging from static constraint checking (𝐾) and structural modification (𝑃) to the generative synthesis of new logic (𝑆). Therefore, the final difficulty score is formulated as: difficulty score = 0.1𝐾 + 0.5(𝑃 − 1) + 𝑆
with one specific transformation; (2). Multi-answer format incorporates a randomized subset of multiple transformations for an original LEIR. This multi-answer design not only encourages diverse optimization reasoning but also accommodates various multi-step optimization scenarios, such as providing multiple candidates for node selection in beam search. As illustrated in Figure 3, both variants follow a standardized prompt-label architecture. The prompt integrates the original LEIR, essential metadata (e.g., program name, input/output shapes and data types), task specifications, and format requirements. The label comprises the CoT trace and the final answer. Specifically, in the multiple-answer variant, the CoT summarizes the number of transformed IRs and provides a numbered reasoning trace for each.
3.4
Verification and Filtering Mechanism
To guarantee dataset correctness and maintain a balanced strategy distribution, we implement a two-fold pipeline consisting of empirical verification and strategy-difficulty-aware filtering. Verification. To ensure the reliability of our dataset, we subject all original and transformed LEIRs to a rigorous verification process. For each program, we execute three independent trials using
(1)
where 𝑃 − 1 accounts for the baseline modification inherent in any transformation (e.g., selecting the target expression). Strategy balancing. Based on the calculated scores, we categorize strategies into three levels and apply differentiated filtering, detailed in Appendix B. For easy strategies (score< 1), we remove all instances from the multiple-answer dataset, whereas for the single-answer version, we retain only 20% of simplification-oriented strategies (e.g., log simplification) and a mere 4% of their inverse expansion counterpart (e.g., expand log simplification). These different ratios are based on our observation that LLMs can typically generalize to expansion tasks after mastering the corresponding simplification logic. For medium strategies (1 ≤score< 2.5), we solely cap their occurrences at 2,000 in the multi-answer version. Finally, all difficult strategies (score ≥ 2.5) are fully retained to maximize the LLM’s exposure to complex optimization logic. Dataset Summary. Following the verification and filtering stages, the 35,878 entries initially generated during the strategy-driven phase were refined to a final collection of 24,953 high-quality samples. The curated dataset includes 7,537 single-answer instances and 17,416 multi-answer instances. In the final distribution (treating each strategy in multi-answer samples independently), difficult
Token Number
Step-TP Dataset for Tensor Program Optimization
10000 7500 5000 2500 0
Model Qwen3-1.7B Qwen3-8B Qwen3-14B Qwen3-32B
Our LEIR
TIR
CUDA
Diff Pass 99% 100% 100% 100%
Build Pass 87% 95% 96% 96%
Exec Pass 87% 94% 96% 95%
Equal Pass 73% 88% 90% 92%
Table 2: Single-attempt pass rates on single-step tensor program transformations for models trained on StepTP.
Figure 4: Token consumption of 6335 tensor programs across LEIR, TIR, and CUDA. strategies account for the majority at 87.32%, while medium and easy strategies constitute 12.17%and 0.51%, respectively. This composition prioritizes high-complexity reasoning while retaining a minimal baseline to consolidate and activate the model’s existing knowledge of fundamental optimizations.
4
Evaluation
In this section, we present an experimental study of LEIR and the Step-TP dataset. Some technical details are deferred to the appendix. The goal of the experimental study is to empirically validate and answer the following main questions for our design. 1. Can LEIR achieve better token efficiency than CUDA and TIR? 2. Can Step-TP enable executable, grounded transformations across a diverse set of strategies? 3. Can Step-TP support long-horizon optimization and exhibit generality across different GPUs? Overview. The results answer these questions affirmatively, validating the effectiveness of LEIR and Step-TP.
4.1
Experimental Setup
Testbed. All experiments are conducted on a machine with 1536GB of host memory and eight NVIDIA H20-3e (140GB memory each). Unless otherwise specified, this platform serves as the default environment. To further evaluate hardware adaptability, we also perform some evaluations on a machine with eight NVIDIA A100 GPUs (80GB memory each) and 800GB of host memory. Models. The experiments are conducted on a range of Qwen3 models with different parameter scales, including Qwen3-1.7B, Qwen38B, Qwen3-14B, and Qwen3-32B. The training setups are detailed in Appendix C
4.2
Token Efficiency
To evaluate the context-efficiency of different IRs, we compare the token consumption of 6335 tensor programs across our LEIR, TVM TIR, and CUDA using the Qwen3 model’s AutoTokenizer. As illustrated in Figure 4, LEIR consistently exhibits the highest structural density. Specifically, the mean token count for LEIR is 499.3, which is significantly lower than that of TVM TIR (1244.2) and CUDA (2897.3). The disparity is even more pronounced at the upper tail of the distribution: while the longest CUDA kernel consumes over 11, 000 tokens and risks exhausting the effective context window of many models, the maximum length for the equivalent LEIR representation remains under 2, 900 tokens. These results demonstrate that LEIR reduces the average token footprint by approximately 60% compared to TIR and 83% compared to CUDA. This efficiency ensures that complex optimization processes can be encoded within prompt-length limits, allowing the
(a) Category
(b) Dfficulty
Figure 5: Number of strategies applied by LLMs of different sizes across 2248 test cases, grouped by (a) strategy category and (b) difficulty level. model to focus its reasoning capacity on strategy selection rather than parsing redundant syntax.
4.3
Single-step Transformations
We evaluate models trained on our dataset via single-step transformations, highlighting two key dataset properties: (i) enabling faithfully grounded, executable program transformations, and (ii) preliminarily supporting a diverse set of tensor-program-level strategies. Setup. We evaluate all models trained on our dataset across 2248 test cases, constructed from 180 distinct tensor programs by varying input/output data types and shapes. All test cases are strictly held out from training. Each model is prompted using our single-answer format, requiring it to autonomously select a valid transformation. Metrics. We employ four progressively stricter metrics to evaluate the transformations, where each success is counted only after passing three independent verification trials to ensure stability: (i) Different pass, to ensure the generated LEIR is syntactically modified; (ii) Build pass, to confirm the transformed LEIR is compilable into CUDA; (iii) Execute pass, to verify the successful execution; and (iv) Equal pass, to validate equivalence with the original program. Result Analysis. Table 2 demonstrates that models fine-tuned on our Step-TP dataset exhibit exceptional fidelity to executable and equivalent program transformations. Specifically, all models achieve near-perfect different pass rates, with Qwen3-32B reaching a 92% equal pass rate and even the 1.7B model achieving 73%. The consistently high success rates in build and execute passes suggest that our reasoning traces effectively guide the models to maintain functional correctness during complex IR modifications. To further examine the effect of repeated sampling, we also evaluate StepTP with more independent attempts in Appendix C. Qwen3-8B trained on our dataset improves from 88.08% equal Pass with a single attempt to 95.55% with 10 attempts and 96.22% with 16 attempts. This underscores the superior quality and robust generalization of the Step-TP dataset. We further analyze the diversity of strategies autonomously selected by the models. As shown in Figure 5(a), the category-wise distribution (graph, operator, memory, math) closely mirrors the underlying dataset ratio (8:9:5:21), confirming that models successfully cover all strategy categories proportionally. As shown in Figure 5(b), easy strategies occur least frequently, and medium and difficult strategies appear at comparable ratesindicating that the models do
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
Method Greedy Search Beam Search BFS Search DFS Search MCTS Chain Search Chain Search wo Parent Chain Search on A100
Avg. # Samples 35.91 103.17 28.27 28.11 17.41 16.63 15.95 16.85
Max. # Samples 41 114 31 31 21 21 21 21
Avg. Speedup 20.79 42.90 18.27 20.97 24.96 23.62 20.25 22.01
Median Speedup 1.78 4.57 2.11 1.68 4.54 2.05 4.96 3.62
Max. Speedup 193.55 561.82 173.13 242.62 171.9 286.92 96.14 265.35
Search Efficiency 0.58 0.60 0.65 0.75 1.43 1.42 1.63 1.31
Table 3: Performance comparison and search efficiency of various search algorithms guided by Qwen3-32B trained on Step-TP. Method Greedy Search Beam Search BFS Search DFS Search MCTS Chain Search Chain Search wo Parent Chain Search on A100
Avg. # Strategy 3.6 5.23 2.45 2.49 3.65 3.57 3.43 3.33
Max. # Strategy 8 10 4 4 9 7 7 9
Table 4: Number of strategies of multi-step optimization trajectories guided by Qwen3-32B trained on Step-TP. not avoid difficult strategies in favor of simpler alternatives. Moreover, we observe that the 32B model utilizes 31 unique strategies at all levels, while other scales cover 29, representing over 70% of the 43 total available strategies. This diverse coverage, aligned with our data distribution, confirms that the models have successfully mastered a broad spectrum of representative optimization patterns.
4.4
Multi-step Optimizations
We evaluate models trained on our dataset via multi-step transformations, highlighting three key dataset properties: (i) supporting long-horizon complex optimization by combining multiple strategies step by step, (ii) enabling efficient search through high-quality step-level supervision, and (iii) exhibiting generality across different GPU environments. We provide a detailed case study of the highest-performing test case in Appendix C. Setup. We evaluate Qwen3-32B, trained on Step-TP, across 100 distinct tensor programs strictly held out from training. The model is tasked with generating runtime-performance-optimized, equivalent LEIRs. To implement multi-step optimization, we deploy seven distinct search algorithms: Greedy Search, Breadth-First Search (BFS), Depth-First Search (DFS), Beam Search, Monte Carlo Tree Search (MCTS), and Chain-based Search (with/without parent nodes). The detailed setups are provided in Appendix C. Metrics. We evaluate results using four metrics: (i) # Samples: the total number of candidate LEIRs verified during the entire optimization process; (ii) Speedup: the runtime improvement of the transformed LEIR relative to the original LEIR , computed as Runtimeoriginal /Runtimetransformed ; (iii) # Strategies: the count of distinct strategy types applied along the final optimization trajectory, where each type is counted once regardless of different applications ; and (iv) Search efficiency: the average speedup achieved per verified sample, calculated as Average Speedup/# Samples. Result Analysis. We evaluate the ability of our dataset Step-TP across all search algorithms. Table 3 details the number of samples, the speedup, and the search efficiency, while Table 4 tracks
the structural complexity of the optimization trajectories via the number of strategies. The results illustrate key properties: (i) Supporting Long-horizon Complex Optimization. As shown in Table 3, across all search algorithms, the model trained on StepTP consistently achieves substantial performance improvements across all search paradigms. For instance, Beam Search achieves an average speedup of 42.90× and a peak speedup of 561.82×. Notably, even under the most restrictive search budgets (e.g., Chain-based Search without parent nodes), the model maintains a high median speedup of 4.96×. These results indicate performance gains are not driven by isolated outliers but demonstrate that our dataset captures high-quality optimization patterns, enabling algorithms to successfully compose sequences of transformations. Table 4 further reveals that optimal schedules require an average of 2.45 to 5.23 distinct strategy types, with a maximum depth of 10 (e.g., in Beam Search). This demonstrates the support of our dataset for diverse tensor-program-level strategies, which serve as the essential building blocks for composing these optimal trajectories. Together, these observations confirm that our dataset provides the necessary structural knowledge to navigate non-trivial, longhorizon optimization landscapes. (ii) Enabling Efficient Search. As shown in Table 3, the search efficiency ranging from 0.58 to 1.63 indicates that the majority of verified transformations contribute meaningfully to the final speedup. In particular, even search algorithms without backtracking or branching techniques achieve significant performance. For instance, Greedy Search and Chain-based Search variants achieve average speedups of 20.25× to 25.97× while requiring as few as 16 to 36 samples. The success of these short-insight paradigms suggests that Step-TP provides high-quality step-level supervision that enables the model to effectively identify high-potential transformation paths, reducing the cost of exhaustive trial and error. (iii) Exhibiting generality across different GPUs. As shown in Tables 3 and 4, Chain-based Search on A100 GPU maintains a strong speedup performance and complex trajectories up to 9 strategies, extending beyond the default H20-3e GPUs. This confirms that the optimization knowledge captured by Step-TP remains effective across different GPU generations.
4.5
Ablation Study
We conduct ablation studies to examine how three key Step-TP dataset designs affect the performance of trained models: the LEIR representation, structured CoT supervision, and strategy filtering. Setup. Four controlled dataset variants are constructed from the same source programs and transformation pipeline. Specifically, we train Qwen3-8B on: (i) TensorIR-based data without CoT but with
Step-TP Dataset for Tensor Program Optimization
Dataset Variant F TIR wo CoT F LEIR wo CoT UF LEIR with CoT Step-TP
Equal Pass 77% 83% 88% 88%
Difficult:Medium:Easy strategy Ratio 42:51:7 44:50:6 25:28:47 42:51:7
Table 5: Single-step transformation ablation results for Qwen3-8B trained on four dataset variants on 2248 test cases.
Dataset Variant F TIR wo CoT F LEIR w/o CoT UF LEIR with CoT Step-TP
Avg. #Samples 8.10 7.31 13.66 13.88
Avg. Speedup 0.72 0.75 8.14 12.10
Search Efficiency 0.08 0.10 0.59 0.87
Table 6: Multi-step optimization ablation results for Qwen38B trained on four dataset variants on 100 distinct test cases using chain-based search with parent nodes.
5
Conclusion
We introduce Step-TP, a step-level post-training dataset for LLMbased tensor program optimization that provides verifiable, compositional supervision for single-step transformation decisions. By combining a token-efficient intermediate representation (LEIR) with atomic strategy decomposition and deterministic equivalence checking, Step-TP enables models to reason about precise optimization steps, rather than relying on outcome-only shortcuts.
ACKNOWLEDGEMENT This work was supported in part by a collaborative research grant from Ant Group and grants from Hong Kong RGC under the contracts 17204423, 17205824, 17204625, C7004-22G (CRF), CRS_PolyU501/23, and T43-513/23-N (TRS).
References strategy filtering (F TIR wo CoT ); (ii) LEIR-based data without CoT but with strategy filtering (F LEIR wo CoT ); (iii) LEIR-based data with CoT but without strategy filtering (UF LEIR with CoT ); and (iv) the full Step-TP dataset with all three components enabled. Result Analysis. We compare the four dataset variants from two perspectives: single-step transformations, which measure transformation fidelity and strategy difficulty, and multi-step optimization, which measures optimization effectiveness and search efficiency. Table 5 reports single-step transformation performance in terms of Equal Pass and the difficulty distribution of the applied strategies, while Table 6 reports the average number of verified samples, average speedup, and search efficiency during multi-step optimization. These results lead to the following observations: (i) LEIR reduces representation-induced transformation errors. As shown in Table 5, replacing TensorIR with LEIR improves equal pass from 77% to 83% when CoT is removed and strategy filtering is kept. This suggests that LEIR helps the model generate semantically equivalent transformations by exposing the relevant loop and equation structures more directly, instead of requiring the model to reason through compiler-oriented TensorIR boilerplate. (ii) Structured CoT is critical for long-horizon optimization. In single-step transformation, Table 5 shows that adding CoT improves equal pass from 83% to 88%, indicating better semantic preservation. In multi-step optimization, Table 6 shows a much larger improvement, with average speedup increasing from 0.75 to 12.10 and search efficiency from 0.10 to 0.87. This suggests that structured CoT helps the model learn composable transformation logic, enabling effective optimization trajectories beyond locally valid IR edits. (iii) Strategy filtering reduces the bias toward simplistic transformations. As shown in Table 5, adding strategy filtering keeps equal pass unchanged at 88%, but changes the generated strategy distribution from 25:28:47 to 42:51:7 for difficult, medium, and easy strategies, respectively. This indicates that filtering suppresses the model’s preference for easy transformations without sacrificing single-step correctness. This shift becomes more meaningful in multi-step optimization, where Table 6 shows higher average speedup from 8.14 to 12.10 and higher search efficiency from 0.59 to 0.87.
[1] Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. 2016. { TensorFlow } : a system for { Large-Scale } machine learning. In 12th USENIX symposium on operating systems design and implementation (OSDI 16). 265–283. [2] Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. 2023. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 (2023). [3] Riyadh Baghdadi, Jessica Ray, Malek Ben Romdhane, Emanuele Del Sozzo, Abdurrahman Akkas, Yunming Zhang, Patricia Suriana, Shoaib Kamil, and Saman Amarasinghe. 2019. Tiramisu: A polyhedral compiler for expressing fast and portable code. In 2019 IEEE/ACM International Symposium on Code Generation and Optimization (CGO). IEEE, 193–205. [4] Carlo Baronio, Pietro Marsella, Ben Pan, Simon Guo, and Silas Alberti. 2025. Kevin: Multi-Turn RL for Generating CUDA Kernels. arXiv preprint arXiv:2507.11948 (2025). [5] Tyler A Chang and Benjamin K Bergen. 2024. Language model behavior: A comprehensive survey. Computational Linguistics 50, 1 (2024), 293–350. [6] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. 2018. { TVM } : An automated { End-to-End } optimizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18). 578–594. [7] Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang, Thierry Moreau, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. 2018. Learning to optimize tensor programs. Advances in Neural Information Processing Systems 31 (2018). [8] Chris Cummins, Volker Seeker, Dejan Grubisic, Mostafa Elhoushi, Youwei Liang, Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Kim Hazelwood, Gabriel Synnaeve, et al. 2023. Large language models for compiler optimization. arXiv preprint arXiv:2309.07062 (2023). [9] Chris Cummins, Volker Seeker, Dejan Grubisic, Baptiste Roziere, Jonas Gehring, Gabriel Synnaeve, and Hugh Leather. 2024. Meta large language model compiler: Foundation models of compiler optimization. arXiv preprint arXiv:2407.02524 (2024). [10] Tri Dao. 2023. Flashattention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691 (2023). [11] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in neural information processing systems 35 (2022), 16344–16359. [12] Juncheng Dong, Yang Yang, Tao Liu, Yang Wang, Feng Qi, Vahid Tarokh, Kaushik Rangadurai, and Shuang Yang. 2025. Stark: Strategic team of agents for refining kernels. arXiv preprint arXiv:2510.16996 (2025). [13] Jingzhi Fang, Yanyan Shen, Yue Wang, and Lei Chen. 2021. ETO: Accelerating optimization of DNN operators by high-performance tensor program reuse. Proceedings of the VLDB Endowment 15, 2 (2021), 183–195. [14] Pratik Fegade, Tianqi Chen, Phillip B Gibbons, and Todd C Mowry. 2024. ACRoBat: Optimizing auto-batching of dynamic deep learning at compile time. Proceedings of Machine Learning and Systems 6 (2024), 14–30. [15] Siyuan Feng, Bohan Hou, Hongyi Jin, Wuwei Lin, Junru Shao, Ruihang Lai, Zihao Ye, Lianmin Zheng, Cody Hao Yu, Yong Yu, et al. 2023. Tensorir: An abstraction for automatic tensorized program optimization. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2. 804–817.
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
[16] Junfeng Gong, Zhiyi Wei, Junying Chen, Cheng Liu, and Huawei Li. 2025. From large to small: Transferring cuda optimization expertise via reasoning graph. arXiv preprint arXiv:2510.19873 (2025). [17] Hanpeng Hu, Junwei Su, Juntao Zhao, Yanghua Peng, Yibo Zhu, Haibin Lin, and Chuan Wu. 2024. CDMPP: A device-model agnostic framework for latency prediction of tensor programs. In Proceedings of the Nineteenth European Conference on Computer Systems. 1054–1074. [18] Zhihao Jia, Oded Padon, James Thomas, Todd Warszawski, Matei Zaharia, and Alex Aiken. 2019. TASO: optimizing deep learning computation with automatic generation of graph substitutions. In Proceedings of the 27th ACM Symposium on Operating Systems Principles. 47–62. [19] Zhihao Jia, James Thomas, Todd Warszawski, Mingyu Gao, Matei Zaharia, and Alex Aiken. 2019. Optimizing DNN computation with relaxed graph substitutions. Proceedings of Machine Learning and Systems 1 (2019), 27–39. [20] Hyeonjin Kim, Sungwoo Ahn, Yunho Oh, Bogil Kim, Won Woo Ro, and William J Song. 2020. Duplo: Lifting redundant memory accesses of deep neural networks for gpu tensor cores. In 2020 53rd Annual IEEE/ACM International Symposium on Microarchitecture (MICRO). IEEE, 725–737. [21] Lingcheng Kong, Jiateng Wei, Hanzhang Shen, and Huan Wang. 2025. Concur: Conciseness makes state-of-the-art kernel generation. arXiv preprint arXiv:2510.07356 (2025). [22] Ao Li, Bojian Zheng, Gennady Pekhimenko, and Fan Long. 2022. Automatic horizontal fusion for GPU kernels. In 2022 IEEE/ACM International Symposium on Code Generation and Optimization (CGO). IEEE, 14–27. [23] Tianle Li, Ge Zhang, Quy Duc Do, Xiang Yue, and Wenhu Chen. 2024. Long-context llms struggle with long in-context learning. arXiv preprint arXiv:2404.02060 (2024). [24] Xiaoya Li, Xiaofei Sun, Albert Wang, Jiwei Li, and Chris Shum. 2025. Cuda-l1: Improving cuda optimization via contrastive reinforcement learning. arXiv preprint arXiv:2507.14111 (2025). [25] Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. 2024. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437 (2024). [26] Mengfan Liu, Wei Wang, and Chuan Wu. 2025. Optimizing distributed deployment of mixture-of-experts model inference in serverless computing. In Ieee infocom 2025-ieee conference on computer communications. IEEE, 1–10. [27] Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. 2024. Lost in the middle: How language models use long contexts. Transactions of the association for computational linguistics 12 (2024), 157–173. [28] Massinissa Merouani, Afif Boudaoud, and Riyadh Baghdadi. 2025. Looperset: A large-scale dataset for data-driven polyhedral compiler optimization. arXiv preprint arXiv:2510.10209 (2025). [29] Maxim Milakov and Natalia Gimelshein. 2018. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.02867 (2018). [30] Anne Ouyang, Simon Guo, Simran Arora, Alex L Zhang, William Hu, Christopher Ré, and Azalia Mirhoseini. 2025. KernelBench: Can LLMs write efficient GPU kernels?, 2025. URL https://arxiv. or g/abs/2502.10517 (2025). [31] Mangpo Phothilimthana, Sami Abu-El-Haija, Kaidi Cao, Bahare Fatemi, Michael Burrows, Charith Mendis, and Bryan Perozzi. 2023. Tpugraphs: A performance prediction dataset on large tensor computational graphs. Advances in Neural Information Processing Systems 36 (2023), 70355–70375. [32] Guicheng Qi, Junwei Su, Liqi Yang, Tao Li, Tingwen Xie, Yerui Sun, Yuchen Xie, and Chuan Wu. 2026. HetAuto: Cross-Cluster Auto-Parallelism for Heterogeneous Distributed Training. In Proceedings of the 21st European Conference on Computer Systems. 759–779. [33] Daniel Snider and Ruofan Liang. 2023. Operator fusion in XLA: analysis and evaluation. arXiv preprint arXiv:2301.13062 (2023). [34] Songqiao Su, Xiaofei Sun, Xiaoya Li, Albert Wang, Jiwei Li, and Chris Shum. 2025. CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning. arXiv preprint arXiv:2512.02551 (2025). [35] Annabelle Sujun Tang, Christopher Priebe, Rohan Mahapatra, Lianhui Qin, and Hadi Esmaeilzadeh. 2025. REASONING COMPILER: LLM-Guided Optimizations for Efficient Model Serving. In The Thirty-ninth Annual Conference on Neural Information Processing Systems. [36] Philippe Tillet, Hsiang-Tsung Kung, and David Cox. 2019. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages. 10–19. [37] Nicolas Vasilache, Oleksandr Zinenko, Theodoros Theodoridis, Priya Goyal, Zachary DeVito, William S Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. 2018. Tensor comprehensions: Framework-agnostic high-performance machine learning abstractions. arXiv preprint arXiv:1802.04730 (2018). [38] Vasily Volkov and James W Demmel. 2008. Benchmarking GPUs to tune dense linear algebra. In SC’08: Proceedings of the 2008 ACM/IEEE conference on Supercomputing. IEEE, 1–11.
[39] Haojie Wang, Jidong Zhai, Mingyu Gao, Zixuan Ma, Shizhi Tang, Liyan Zheng, Yuanzhi Li, Kaiyuan Rong, Yuanyong Chen, and Zhihao Jia. 2021. { PET } : Optimizing tensor programs with partially equivalent transformations and automated corrections. In 15th USENIX Symposium on Operating Systems Design and Implementation (OSDI 21). 37–54. [40] Lei Wang, Yu Cheng, Yining Shi, Zhengju Tang, Zhiwen Mo, Wenhao Xie, Lingxiao Ma, Yuqing Xia, Jilong Xue, Fan Yang, et al. 2025. TileLang: A Composable Tiled Programming Model for AI Systems. arXiv preprint arXiv:2504.17577 (2025). [41] Lei Wang, Lingxiao Ma, Shijie Cao, Quanlu Zhang, Jilong Xue, Yining Shi, Ningxin Zheng, Ziming Miao, Fan Yang, Ting Cao, et al. 2024. Ladder: Enabling efficient { Low-Precision } deep learning computing through hardware-aware tensor transformation. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24). 307–323. [42] Weiyang Wang, Moein Khazraee, Zhizhen Zhong, Zhijao Jia, Dheevatsa Mudigere, Ying Zhang, Anthony Kewitsch, and Manya Ghobadi. 2022. Topoopt: Optimizing the network topology for distributed dnn training. arXiv preprint arXiv:2202.00433 (2022). [43] Jiin Woo, Shaowei Zhu, Allen Nie, Zhen Jia, Yida Wang, and Youngsuk Park. 2025. Tritonrl: Training llms to think and code triton without cheating. arXiv preprint arXiv:2510.17891 (2025). [44] Mengdi Wu, Xinhao Cheng, Shengyu Liu, Chunan Shi, Jianan Ji, Man Kit Ao, Praveen Velliengiri, Xupeng Miao, Oded Padon, and Zhihao Jia. 2025. Mirage: A { Multi-Level } superoptimizer for tensor programs. In 19th USENIX Symposium on Operating Systems Design and Implementation (OSDI 25). 21–38. [45] Haofeng Xu, Junwei Su, Yukun Tian, Lansong Diao, Zhengping Qian, and Chuan Wu. 2026. GAC: Stabilizing Asynchronous RL Training for LLMs via Gradient Alignment Control. arXiv preprint arXiv:2603.01501 (2026). [46] Zi Yang, Lei Qiu, Fang Lyu, Ming Zhong, Zhilei Chai, Haojie Zhou, Huimin Cui, and Xiaobing Feng. [n. d.]. IR-OptSet: An Optimization-Sensitive Dataset for Advancing LLM-Based IR Optimizer. In The Thirty-ninth Annual Conference on Neural Information Processing Systems Datasets and Benchmarks Track. [47] Yi Zhai, Sijia Yang, Keyu Pan, Renwei Zhang, Shuo Liu, Chao Liu, Zichun Ye, Jianmin Ji, Jie Zhao, Yu Zhang, et al. 2024. Enabling Tensor Language Model to Assist in Generating { High-Performance } Tensor Programs for Deep Learning. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24). 289–305. [48] Yi Zhai, Yu Zhang, Shuo Liu, Xiaomeng Chu, Jie Peng, Jianmin Ji, and Yanyong Zhang. 2023. Tlp: A deep learning-based cost model for tensor program tuning. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2. 833–845. [49] Jie Zhao, Xiong Gao, Ruijie Xia, Zhaochuang Zhang, Deshi Chen, Lei Chen, Renwei Zhang, Zhen Geng, Bin Cheng, and Xuefeng Jin. 2022. Apollo: Automatic partition-based operator fusion through layer by layer optimization. Proceedings of Machine Learning and Systems 4 (2022), 1–19. [50] Jie Zhao, Bojie Li, Wang Nie, Zhen Geng, Renwei Zhang, Xiong Gao, Bin Cheng, Chen Wu, Yun Cheng, Zheng Li, et al. 2021. AKG: automatic kernel generation for neural processing units using polyhedral transformations. In Proceedings of the 42nd ACM SIGPLAN International Conference on Programming Language Design and Implementation. 1233–1248. [51] Lianmin Zheng, Chengfan Jia, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang, Danyang Zhuo, Koushik Sen, et al. 2020. Ansor: Generating { High-Performance } tensor programs for deep learning. In 14th USENIX symposium on operating systems design and implementation (OSDI 20). 863–879. [52] Lianmin Zheng, Ruochen Liu, Junru Shao, Tianqi Chen, Joseph E Gonzalez, Ion Stoica, and Ameer Haj Ali. 2021. Tenset: A large-scale program performance dataset for learned tensor compilers. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1). [53] Liyan Zheng, Haojie Wang, Jidong Zhai, Muyan Hu, Zixuan Ma, Tuowei Wang, Shuhong Huang, Xupeng Miao, Shizhi Tang, Kezhao Huang, et al. 2023. { EINNET } : Optimizing tensor programs with { Derivation-Based } transformations. In 17th USENIX Symposium on Operating Systems Design and Implementation (OSDI 23). 739–755. [54] Size Zheng, Siyuan Chen, Siyuan Gao, Liancheng Jia, Guangyu Sun, Runsheng Wang, and Yun Liang. 2023. Tileflow: A framework for modeling fusion dataflow via tree-based analysis. In Proceedings of the 56th Annual IEEE/ACM International Symposium on Microarchitecture. 1271–1288. [55] Yuchen Zhong, Junwei Su, Chuan Wu, and Minjie Wang. 2025. Heta: Distributed Training of Heterogeneous Graph Neural Networks. Proceedings of the VLDB Endowment 18, 9 (2025), 2790–2803. [56] Hongyu Zhu, Ruofan Wu, Yijia Diao, Shanbin Ke, Haoyu Li, Chen Zhang, Jilong Xue, Lingxiao Ma, Yuqing Xia, Wei Cui, et al. 2022. { ROLLER } : Fast and efficient tensor compilation for deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22). 233–248.
Step-TP Dataset for Tensor Program Optimization
A
Intermediate representation
Loop structure. The loop structure is represented by a main symbol indicating the loop type, with superscripts and subscripts specifying the loop index and iteration range. (1). Main symbol: Five loop notations represent specific loop types: serial loops (𝐿), parallel loops (𝑃), vectorized loops (𝑉 ), unrolled loops (𝑈 ), and thread/block-binding loops (𝐵). (2). Superscript and subscript: The numbers in the superscript 32 for a range of and subscript denote the loop range, such as 𝐿𝑎=0 [0, 32). The lowercase letters in the subscript indicate the loop index. Specifically, indices for thread/block-binding loops are categorized into six mappings to CUDA intrinsics: {𝑏𝑥, 𝑏𝑦, 𝑏𝑧} for blockIdx 6 and {𝑡𝑥, 𝑡𝑦, 𝑡𝑧} for threadIdx. For example, 𝐵𝑡𝑥=0 denotes a loop bound to threadIdx.x. Equation structure. The equation structure consists of three components: the element-wise computation, the variables to denote tensors, and delimiters to define the computational scope and logical sequence within the loop domain. Furthermore, we employ default initialization to maintain an uninterrupted algebraic flow. Element-wise computation. We formalize the algebraic computations using a set of functional operators derived from TVM TIR, including standard arithmetic (e.g., +, −), transcendental functions (e.g., exp, log), and conditional intrinsics (e.g., if_then_else), all applied in a purely element-wise manner. Variable. Each tensor variable is defined by a main symbol for its identity, with subscripts specifying indices and superscripts indicating data and memory type. (1). Main symbol: Tensor variables (input, output, and intermediate) are represented by an uppercase letter followed by optional lowercase letters (e.g., 𝐴, 𝐴𝑐, 𝐶𝑒𝑚). To avoid semantic ambiguity with loop types and existing compiler namespaces (such as the TVM T package), the symbols {𝐿, 𝑃, 𝑉 , 𝑈 , 𝐵,𝑇 } are reserved and excluded from variable naming. (2). Subscript: The subscript specifies the access indices using lowercase letters. To maintain semantic clarity and index integrity, certain characters are reserved: {𝑡, 𝑏, 𝑥, 𝑦, 𝑧} are dedicated to loop binding semantics, while 𝑒 is excluded to avoid syntactic confusion with the exponential operator (exp). (3). Superscript: The superscript encodes two mandatory attributes: the data type (e.g., 𝑓 32, 𝑖64, 𝑓 16 for FLOAT32, INT64, FLOAT16) and the memory type (𝑔 for global, 𝑙 for local, 𝑠 for shared memory). To ensure a fully specified execution state, both subscripts and 𝑓 32,𝑔 superscripts are mandatory for all tensor variables (e.g., 𝐴𝑐 𝑎,𝑐,𝑒 ). The only exception is for scalars, where subscripts are omitted as they lack associated iteration axes. Delimiter. Delimiters explicitly define the scope of computations and separate independent expressions within the IR. (1). Brackets: Square brackets [ and ] enclose one or more equations to define their computational scope, indicating that the enclosed logic is executed under the specified loop conditions. For 𝑓 32,𝑠 8 16 [𝐶 𝑓 32,𝑠 example, For example, 𝐵𝑡𝑥=0 𝐿𝑎=0 𝑡𝑥∗16+𝑎 = 𝐴𝑠𝑡𝑥∗16+𝑎 +1; ]; denotes the element-wise computation under a nested structure consisting of a thread-binding loop 𝑡𝑥 and a serial loop 𝑎.
(2). Semicolon: The semicolon ; acts as a separator for both independent equations and top-level expressions, where an expression comprises the loop nest and its enclosed equations. As illus4 [𝐶 𝑓 32,𝑔 = 𝐴 𝑓 32,𝑔 ∗ 3; 𝐷 𝑓 32,𝑔 = 𝐴 𝑓 32,𝑔 ∗ 2 + trated in the example, 𝐿𝑎=0 𝑎 𝑎 𝑎 𝑎 4 [𝐻 𝑓 32,𝑔 = 𝐴 𝑓 32,𝑔 ∗ 2; ];, the semicolon enforces a sequen1; ]; 𝐿𝑎=0 𝑎 𝑎 tial execution order: (i) within a loop body, it separates distinct equations (e.g., to obtain 𝐶 and 𝐷), ensuring their sequential evaluation within every iteration. (ii) At the expression level, it separates independent loop nests, indicating that the second loop (e.g., to calculate 𝐻 ) begins execution only after the entire preceding loop nest has completed. Default Initialization. To enhance brevity and semantic density, our IR employs implicit initialization for common reduction patterns, including summation, product, and extremum operations. The initial identity element is automatically inferred from the operator: summations default to 0, products to 1, and max or min reductions to −∞ or +∞, respectively.
B Strategy B.1 Tensor-Program-Level Strategies Scope of Tensor-Program-Level Optimization. Tensor-programlevel optimization focuses on the logical transformations of a program’s structure and algebraic representation, specifically targeting loop hierarchies, equation compositions, and tensor access patterns without altering computational semantics. While kernel- or hardware-level tuning manages physical execution details (e.g., memory bank conflicts, cache line utilization, and pipeline depths), this logical abstraction defines what is computed and how it is logically structured, independent of its implementation on specific hardware targets.
B.2
Strategy Categories
Motivation for Categorizing Strategies. We categorize tensorprogram-level optimizations into distinct strategies for two main reasons: (i). Facilitating step-level supervision. Effective supervision requires models to identify the semantic shift introduced at each transformation step. However, due to the vast tensor-program-level optimization space, each step often induces complex and entangled program changes, leading to supervision signals that are parse and semantically coarse. To ensure learnability, we partition the strategy space into a structured taxonomy. (ii). Handling strategy heterogeneity. Beyond mere partitioning, the optimization space benefits from organization into distinct semantic levels due to its inherent heterogeneity. Transformations at the tensor-program level act on diverse scopes (e.g., loop structure, algebraic formation), each requiring fundamentally different reasoning logic. Treating these as a flat space would obscure such structural differences, complicating the learning of disparate optimization logics. We therefore organize the strategy space by the program aspect being modified, ensuring each level represents a consistent semantic domain. Consequently, this stratification further allows complex end-to-end optimizations to be decomposed into interpretable and analyzable sequences of strategy steps.
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
Strategy operator fusion
Description and Examples it places consecutive equations from multiple similar loop nests into a single loop nest. 478 𝐿 478 [𝐸 𝑓 16,𝑔 = if_then_else(𝑡𝑥 >= For example, the consecutive expressions 𝐵𝑡𝑥=0 𝑎=0 𝑡𝑥,𝑎 𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
478 𝐿 478 [𝐷 𝑎, 𝐶𝑡𝑥,𝑎 , 0); ]; and 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = if_then_else((𝑡𝑥 >= 𝑎, 𝐴𝑡𝑥,𝑎 , 0); ]; has the 𝑎=0 478 𝐿 478 , and are fused into the expression 𝐵 478 𝐿 478 [𝐷 𝑓 16,𝑔 same loop nest 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥=0 𝑎=0 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 if_then_else((𝑡𝑥 >= 𝑎, 𝐴𝑡𝑥,𝑎 , 0); 𝐸𝑡𝑥,𝑎 = if_then_else((𝑡𝑥 >= 𝑎, 𝐶𝑡𝑥,𝑎 , 0); ];.
operator fission
=
it splits multiple equations in one loop nest into multiple separate loop nests. 975 𝐿 10081 [𝐶 𝑓 64,𝑠 = 1/(1 + 𝑒𝑥𝑝 (−𝐴 𝑓 64,𝑔 )); 𝐷 𝑓 64,𝑠 = 𝐴 𝑓 64,𝑔 ∗ 𝐶 𝑓 64,𝑠 ; ]; For example, 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
975 𝐿 10081 [𝐷 with multiple equations are split into 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 𝐴𝑡𝑥,𝑎 ∗ 𝐶𝑡𝑥,𝑎 ; ]; and 𝑎=0 975 𝐿 10081 [𝐶 𝑓 64,𝑔 = 1/(1 + 𝑒𝑥𝑝 (−𝐴 𝑓 64,𝑔 )); ]; as two expressions. 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0
compute inline
it merges related equations from multiple loop nests into one equation within a 175 𝐿 28272 [𝐶 𝑓 16,𝑔 = 𝑎𝑏𝑠 (𝐴 𝑓 16,𝑔 ); ]; single loop nest. For example, the expressions 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
175 𝐿 28272 [𝐹 175 28272 and 𝐵𝑡𝑥=0 𝑎=0 𝑡𝑥,0 = 𝐹𝑡𝑥,0 + 𝐶𝑡𝑥,𝑎 ; ]; has the same loop nest 𝐵𝑡𝑥=0 𝐿𝑎=0 and 𝑓 16,𝑔
𝑓 16,𝑔
175 𝐿 28272 [𝐹 related equations, and are merged into the expression 𝐵𝑡𝑥=0 𝑎=0 𝑡𝑥,0 = 𝐹𝑡𝑥,0 + 𝑓 16,𝑔
𝑎𝑏𝑠 (𝐴𝑡𝑥,𝑎 ); ];. it separates a merged equation in one loop nest into multiple equations under mul929 𝐿 10637 [𝐸 𝑓 32,𝑔 = 𝑒𝑥𝑝 (𝐷 𝑓 32,𝑔 − 𝐹 𝑓 32,𝑔 )/𝐺 𝑓 32,𝑔 ; ]; tiple loop nests. For example, 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥 𝑡𝑥 𝑎=0
expression splitting
𝑓 32,𝑔
929 𝐿 10637 [𝐽 with the merged equation are split into 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑔
= 𝑒𝑥𝑝 (𝐷𝑡𝑥,𝑎 −
𝑓 32,𝑔
929 𝐿 10637 [𝐸 𝐹𝑡𝑥 ); ]; 𝑎𝑛𝑑𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 𝐽𝑡𝑥,𝑎 /𝐺𝑡𝑥 ; ]; as two expressions 𝑎=0 it concatenates multiple input variables into one variable, merges the equations, and then splits the output variable to obtain multiple outputs. For example, for two sim99 𝐿 1167 𝐿 450 [𝐷 𝑓 32,𝑔 = 𝑒𝑥𝑝 (𝐴 𝑓 32,𝑔 ); ]; and 𝐵 99 𝐿 450 𝐿 365 [𝐸 𝑓 32,𝑔 = ilar expressions 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 𝑡𝑥=0 𝑎=0 𝑐=0 𝑡𝑥,𝑎,𝑐
tensor concat to fuse operators
𝑓 32,𝑔
𝑓 32,𝑔
99 𝐿 1167 𝐿 450 [𝐺 𝑒𝑥𝑝 (𝐶𝑡𝑥,𝑎,𝑐 ); ];, the input variables are contenated by 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 = 𝑎=0 𝑐=0 𝑓 32,𝑔 𝑓 32,𝑔 99 𝐿 450 𝐿 365 [𝐺 𝑓 32,𝑔 𝐴𝑡𝑥,𝑎,𝑐 ; ]; and 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑡𝑥,𝑎+1167,𝑐+450 = 𝐶𝑡𝑥,𝑎,𝑐 ; ];, the similar operations are 99 𝐿 1617 𝐿 815 [𝐻 𝑓 32,𝑔 = 𝑒𝑥𝑝 (𝐺 𝑓 32,𝑔 ); ];, and the output variable is split executed by 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 99 𝐿 1167 𝐿 450 [𝐷 𝑓 32,𝑔 = 𝐻 𝑓 32,𝑔 ; ]; and 𝐵 99 𝐿 450 𝐿 365 [𝐸 𝑓 32,𝑔 = into two outputs by 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 𝑡𝑥=0 𝑎=0 𝑐=0 𝑡𝑥,𝑎,𝑐 𝑓 32,𝑔 𝐻𝑡𝑥,𝑎+1167,𝑐+450 ; ];
tensor split to decouple operators
it splits an input variable into multiple variables, runs multiple equations, and then concatenates multiple output variables into one variable. For example, for 687 𝐿 217 [𝐶 𝑓 32,𝑔 = 𝑚𝑎𝑥 (0, 𝑚𝑖𝑛(1, (𝐴 𝑓 32,𝑔 + 3)/6)); ];, the input varithe expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥,𝑎 𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑔
57 𝐿 217 [𝐶 630 217 able is split into two inputs by 𝐵𝑡𝑥=0 𝑎=0 𝑡𝑥,𝑎 = 𝐹𝑡𝑥,𝑎 ; ]; and 𝐵𝑡𝑥=0 𝐿𝑎=0 [𝐸𝑡𝑥,𝑎 = 𝑓 32,𝑔 𝑓 32,𝑔 630 𝐿 217 [𝐶 𝑓 32,𝑔 𝐴𝑡𝑥+57,𝑎 ; ]; , 𝑡𝑤𝑜𝑠𝑖𝑚𝑖𝑙𝑎𝑟𝑜𝑝𝑒𝑟𝑎𝑡𝑖𝑜𝑛𝑠𝑎𝑟𝑒𝑒𝑥𝑒𝑐𝑢𝑡𝑒𝑑𝑏𝑦𝐵𝑡𝑥=0 𝑎=0 𝑡𝑥+57,𝑎 = 𝐺𝑡𝑥,𝑎 ; ]; and 57 𝐿 217 [𝐷 𝑓 32,𝑔 = 𝐴 𝑓 32,𝑔 ; ];, and two output variables are concatenated into one out𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 57 𝐿 217 [𝐹 𝑓 32,𝑔 = 𝑚𝑎𝑥 (0, 𝑚𝑖𝑛(1, (𝐷 𝑓 32,𝑔 + 3)/6)); ]; and 𝐵 630 𝐿 217 [𝐺 𝑓 32,𝑔 = put by 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥,𝑎 𝑡𝑥=0 𝑎=0 𝑓 32,𝑔 𝑚𝑎𝑥 (0, 𝑚𝑖𝑛(1, (𝐸𝑡𝑥,𝑎 + 3)/6)); ];.
common subexpression elimination
it computes duplicated expressions once and reuses the result to avoid redundant 𝑓 64,𝑔 calculations. For example, a common equation part 𝑒𝑥𝑝 (𝐴𝑡𝑥,𝑎 ) exists, so this part 𝑓 64,𝑔
expression reorder
𝑓 64,𝑔
917 𝐿 30201 [𝐹 can be computed in the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 𝑒𝑥𝑝 (𝐴𝑡𝑥,𝑎 ); ]; once. 𝑎=0 it rearranges the expressions. For example, two expressions 𝑓 64,𝑔 𝑓 64,𝑔 𝑓 64,𝑔 448 𝐿 2 𝐿 549 𝐿 549 𝐿 154 [𝑄 𝑓 64,𝑔 𝐵𝑡𝑥=0 = 𝑄 + 𝐾 ∗ 𝑂 ; ]; and 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑡𝑥,𝑎,𝑐,𝑑 𝑡𝑥,𝑎,𝑐,𝑑 𝑡𝑥,𝑎,𝑐,𝑓 𝑡𝑥,𝑎,𝑓 ,𝑑 1 𝐵𝑡𝑥=0 [𝑅 𝑓 64,𝑔 = 154; ]; can be reordered. Table 7: Graph-level atomic optimization strategies.
Four Categories of Strategies. Based on the semantic dimension primarily modified by a strategy, we categorize tensor-programlevel optimizations into four classes, as detailed in Table 7-Table 13 (1). Graph level. These optimizations act on the organization of multiple expressions. An example is the expression reorder strategy,
where independent expressions are permuted without violating data dependencies. (2). Operator level. These strategies target the loop–equation structure within a single expression. For example, the loop reorder
Step-TP Dataset for Tensor Program Optimization
strategy swaps the nesting order of LEIR iterators, such as trans719 𝐿 549 into 𝐿 549 𝐵 719 . forming 𝐵𝑡𝑥=0 𝑎=0 𝑎=0 𝑡𝑥=0 (3). Memory level. These optimizations primarily alter the logical storage and layouts of tensor variables. the storage scope strategy 𝑓 64,𝑔 rebinds a tensor 𝐸𝑎 in global memory to a local memory scope 𝑓 64,𝑙 𝐸𝑎 to optimize data proximity. (4). Mathematical level. These strategies directly rewrite algebraic formulations to change execution logic while preserving numerical semantics. For example, the online softmax strategy reformulates the global exponential sum into incremental update equations, decomposing a monolithic reduction into recursive algebraic steps within a loop. In total, we have identified 43 distinct strategies across these levels, comprising 8 graph-level, 9 operator-level, 5 memory-level, and 21 mathematical-level strategies. Among these, mathematicallevel strategies are particularly challenging to realize with modular or rule-based tools as they require global algebraic reasoning. LEIR addresses these challenges by providing a high-density, unified representation that abstracts these disparate dimensions into a standardized, step-level format.
B.3
Preconditions for Strategy Filtering
Given that strategies are constrained by specific program structures, we define nine essential preconditions as follows: (1). The pattern match check detects a specific computation pattern (e.g., softmax computation for online softmax strategy); (2). The dependency check verifies the existence of computation-order dependencies among expressions (e.g., for expression reorder strategy); (3). The operation identity check ensures that repeated computation operations exist (e.g., for common subexpression elimination); (4). The loop nest consistency check ensures that expressions share the same loop structure (e.g., for operator fusion strategy); (5). The equation count check ensures a sufficient number of equations exist (e.g., for operator fission strategy); (6). The loop axis count check verifies that the number of loop axes is sufficient (e.g., for loop reorder strategy); (7). The loop range factorization check ensures that the selected loop range can be split appropriately (e.g., for loop split strategy); (8). The reduction axis check prevents applying illegal strategies (e.g., loop binding) to reduction axes; (9). The intermediate variable check validates whether the tensor variables are inputs or outputs (e.g., for set storage scope strategy). (1) operator fusion: dependency check, loop nest consistency check; (2) operator fission: dependency check, equation count check; (3) compute inline: dependency check, loop nest consistency check; (4) expression splitting: n/a; (5) tensor concat to fuse operators: dependency check, operation identity check; (6) tensor split to decouple operators: n/a; (7) common subexpression elimination: operation identity check, loop nest consistency check; (8) expression reorder: dependency check ; (9) loop reorder: loop axis count check; (10) loop tiling: loop axis count check, loop range factorization check;
(11) loop split: loop range factorization check; (12) loop fusion: loop axis count check, reduction axis check; (13) loop unrolling: n/a; (14) loop parallelization: reduction axis check; (15) loop vectorization: reduction axis check; (16) loop binding: reduction axis check; (17) reduction factorization: reduction axis check; (18) cache read write: n/a; (19) layout transformation: n/a; (20) set storage scope: intermediate variable check; (21) set storage layout: intermediate variable check; (22) precompute indices: pattern match check; (23) factorization: pattern match check; (24) expand factorization: pattern match check; (25) cancellation: pattern match check; (26) expand cancellation: pattern match check; (27) apart: pattern match check; (28) together: pattern match check; (29) powsimp: pattern match check; (30) expand powsimp: pattern match check; (31) logsimp: pattern match check; (32) expand log: pattern match check; (33) collect: pattern match check; (34) expand collect: pattern match check; (35) partially equivalent then correct: n/a; (36) exponential split: pattern match check; (37) multiplicative split: n/a; (38) additive split: n/a; (39) normal loop max to prefix max: pattern match check; (40) normal loop summation on exp to prefix summation on exp: pattern match check; (41) online softmax: pattern match check; (42) flashattention wo tiling: pattern match check (43) normal matmul to prefix matmul based on online softmax: pattern match check.
B.4
Strategy Difficulty • Easy: operator fission, factorization, expand factorization, cancellation, expand cancellation, apart, together, powsimp, expand powsimp, logsimp, expand log, collect, expand collect; • Medium: operator fusion, compute inline, expression splitting, expression reorder, loop reorder, loop unrolling, loop parallelization, loop vectorization, loop binding, exponential split, ultiplicative split, additive split; • Dfficult:tensor concat to fuse operators, tensor split to decouple operators, common subexpression elimination, loop tiling, loop split, loop fusion, reduction factorization, cache read write, layout transformation, set storage scope , set storage layout, precompute indices, partially equivalent then correct, normal loop max to prefix max, normal loop summation on exp to prefix summation on exp, online softmax, flashattention wo tiling, normal matmul to prefix matmul based on online softmax;
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
C
Evaluation
Training. All models are fine-tuned using LoRA with a rank of 8, LoRA alpha set to 32, and a dropout rate of 0.05. LoRA adapters are applied to all linear layers. We use a learning rate of 1𝑒 − 4, a weight decay of 0.1, and a warmup ratio of 0.05. The batch size is set to 64 for Qwen3-1.7B and 16 for the larger models due to memory constraints. All models are trained for 3 epochs. LoRA fine-tuning is implemented using the ms-swift framework.
C.1
Single-step Transformation
Setup. We evaluate all models trained on our dataset across 2248 test cases, constructed from 180 distinct tensor programs by varying input/output data types and tensor shapes. All test cases are strictly held out from training. The workload of these test cases are detailed at the end of the section. Each model is tasked to generate a mathematically equivalent IR in a single step, following the prompt format of our dataset’s single-answer version. The temperature is set to 0 during inference. Crucially, we do not specify any strategies, requiring the model to autonomously select a valid transformation.
C.2
Multi-step Transformation
Setup. We evaluate Qwen3-32B, trained on our dataset, across 100 distinct tensor programs strictly held out from training. The temperature is set to 0.3 during inference. For accurate performance measurement, we employ TVM’s built-in time_evaluator to record execution time. Specifically, each kernel is executed for three warmup runs to mitigate transient hardware effects, followed by three measurement repetitions to ensure statistical stability, with the average latency reported. The model is tasked with generating mathematically equivalent, runtime-performance-optimized IRs to achieve higher speedups. Each test case prompt provides six key pieces of information: (i) the search algorithm, (ii) the current LEIR and related metadata, (iii) the exploration history if required by the search algorithm, with at most one parent node, (iv) the target GPU specifications, (v) all 43 potential strategies, and (vi) the task description. An example prompt is provided as follows: "Breadth-first-based optimization is used on a given IR to improve performance. Each IR is a state, and has a parent transformation and speedup performance. Give the current IR of Gemm Swish Divide Clamp Tanh Clamp: 728 𝐵 1243 𝐿 2022 [𝐶 𝑓 32,𝑔 = 𝐶 𝑓 32,𝑔 +𝐴 𝑓 32,𝑔 ∗𝐽 𝑓 32,𝑔 ; ]; 𝐵 728 𝐿 1243 [𝐶 𝑓 32,𝑔 𝐵𝑡𝑥=0 𝑡𝑥,𝑐 𝑡𝑥,𝑎 𝑡𝑥=0 𝑎=0 𝑏𝑥𝑎=0 𝑐=0 𝑡𝑥,𝑏𝑥𝑎 𝑡𝑥,𝑏𝑥𝑎 𝑏𝑥𝑎,𝑐 𝑓 32,𝑔
𝑓 32,𝑔
= 𝐶𝑡𝑥,𝑎 +𝐾𝑎
𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑔
728 𝐿 1243 [𝐹 ; ]; 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 0.5 ∗𝐶𝑡𝑥,𝑎 /(1 +𝑒𝑥𝑝 (−𝐶𝑡𝑥,𝑎 𝑎=0 𝑓 32,𝑔
𝑓 32,𝑔
728 𝐿 1243 [𝐺 728 1243 )); ]; 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 𝑚𝑖𝑛(𝑚𝑎𝑥 (𝐹𝑡𝑥,𝑎 , −1.0), 1.0); ]; 𝐵𝑡𝑥=0 𝐿𝑎=0 𝑎=0 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 728 𝐿 1243 [𝑁 𝑓 32,𝑔 = 𝑒𝑥𝑝 [𝑀𝑡𝑥,𝑎 = 𝑒𝑥𝑝 (𝐺𝑡𝑥,𝑎 ) − 𝑒𝑥𝑝 (−𝐺𝑡𝑥,𝑎 ); ]; 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑓 32,𝑔 𝑓 32,𝑔 728 𝐿 1243 [𝐻 𝑓 32,𝑔 = 𝑀 𝑓 32,𝑔 /𝑁 𝑓 32,𝑔 ; ]; 𝐵 728 (𝐺𝑡𝑥,𝑎 )+𝑒𝑥𝑝 (−𝐺𝑡𝑥,𝑎 ); ]; 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥=0 1243 [𝐼 𝑓 32,𝑔 = 𝑚𝑖𝑛(𝑚𝑎𝑥 (𝐻 𝑓 32,𝑔 , −1.0), 1.0); ];, where the known 𝐿𝑎=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎
variables are ’A’ with the dtype torch.float32 and shape [728, 2022], ’J’ with the dtype torch.float32 and shape [1243, 2022], ’K’ with the dtype torch.float32 and shape [1243], and ’I’ with the dtype torch.float32 and shape [728, 1243]. Do not change the names, shapes or dtypes of these known variables in the IR. History: The parent IR: same as the root IR, depth:0, speedup value: 1.
The speedup value of the current IR: 30.358607118641192, the depth is 1, and the current IR is obtained from the parent IR using the strategy ’loop_binding’. **Target hardware**: NVIDIA H20-3e GPU. CUDA binding rules: Loop axes bound to block (along x,y,z axis, max dimension value: 23 1 − 1, 65535, 65535) MUST be renamed with prefixes bx, by, bz and unique, respectively, followed by other unique lowercase letters. Loop axes bound to thread (along x,y,z axis, max dimension value: 1024, 1024, 64) MUST be renamed with prefixes tx, ty, tz, respectively, followed by other unique lowercase letters. Memory usage rules: Data indexed by block-level loops may be placed in shared (s) or global (g) memory. Data indexed by threadlevel loops may be placed in local (l), shared (s) or global (g) memory. The following strategies and any other mathematical strategies can be considered: operator fusion, operator fission, compute inline, expression splitting, tensor concat to fuse operators, tensor split to decouple operators, common subexpression elimination, expression reorder, loop reorder, loop tiling, loop split, loop fusion, loop unrolling, loop parallelization, loop vectorization, loop binding, reduction factorization, cache read write, layout transformation, set storage scope, set storage layout, precompute indices, factorization, expand factorization, cancellation,expand cancellation, apart, together, powsimp, expand powsimp, expand log, logsimp, collect, expand collect, partially equivalent then correct, normal loop max to prefix max, exponential split, multiplicative split, additive split, normal loop summation on exp to prefix summation on exp, online softmax, flashattention wo tiling, normal matmul to prefix matmul based on online softmax. **Task**: Please give me at least 2 different **numerically equivalent, runtime-performance-optimized** IRs that produce exactly the same outputs for any floating-point inputs (bitwise identical) to achieve higher speedup values (should be more than 1), and also provide applied strategy for each transformed IR. Return the answer list **only** as a valid JSON object, and each entry with the following keys: ’idx’, ’transformed_IR’, ’applied_strategies’. CRITICAL: 1. Before you suggest each new transformation, you MUST identify what has been changed in the current IR compared to the root IR. For each new optimization, you MUST build it ON TOP OF these existing changes, namely ON TOP OF the current IR. The strategies MUST be used on the current IR! You MUST compare your modified parts in each transformed IR with the current IR. If they are identical strings, your answer is WRONG. If the unmodified part in each transformed IR is different from the current IR, your answer is WRONG. 2.Don’t repeat the current or parent IRs! You MUST NOT revert to the parent IR: In particular, you are NOT allowed to apply any reverse or undo operation that reconstructs the current IR from its parent IR, including inverse transformations such as operator fusion <-> operator fission, loop tiling <-> loop fusion, loop split <-> loop fusion, apart <-> together, collect <-> expand collect, or similar reversals."
Step-TP Dataset for Tensor Program Optimization
Search Algorithm. To implement multi-step optimization, we deploy seven distinct search algorithms categorized by their exploration strategies: (i) Greedy Search, which generates multiple candidate transformations per step and selects the locally best one; (ii) Breadth-First Search (BFS) and (iii) Depth-First Search (DFS), representing exhaustive breadth and depth explorations; (iv) Beam Search, which generates multiple transformations per step and maintains the top-𝑘 candidates as the search frontier; (v) Monte Carlo Tree Search (MCTS), which balances exploration and exploitation using rollouts and value estimation to guide transformation selection; and (vi) Chain-based Search (with/without parent nodes), which sequentially refines transformations step by step, optionally considering the parent node. For greedy search, beam search, MCTS, and two chain-based search variants, the maximum number of iterations is set to 20. For DFS and BFS, the maximum depth is set to 20 andthe number of generated LEIRs for each node is set to 2. For greedy search, the breadth is set to 2. For beam search, the number of generated LEIRs for each node is set to 3, and the k value is set to 2 in order to maintain the top-k candidates. Case Study. We analyze a complex tensor program featuring ma457 𝐿 2265 𝐿 3520 trix multiplication, scaling, and residual addition(i.e., 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 457 2265 [𝐷𝑡𝑥,𝑎 = 𝐷𝑡𝑥,𝑎 + 𝐴𝑡𝑥,𝑐 ∗ 𝐼𝑎,𝑐 ; ]; 𝐵𝑡𝑥=0 𝐿𝑎=0 [𝐷𝑡𝑥,𝑎 = 𝐷𝑡𝑥,𝑎 + 𝑓 32,𝑔
457 𝐿 2265 [𝐸 𝑓 32,𝑔 = 𝐷 𝑓 32,𝑔 ∗ 𝐻 𝑓 32,𝑔 ; ]; 𝐵 457 𝐿 2265 [𝐹 𝑓 32,𝑔 = ; ]; 𝐵𝑡𝑥=0 𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥=0 𝑎=0 𝑓 32,𝑔 𝑓 32,𝑔 𝐸𝑡𝑥,𝑎 +𝐶𝑡𝑥,𝑎 ; ];), achieving a 561.82× speedup through a 15-step op-
𝐽𝑎
timization trajectory. The process begins with a layout transposition (Step 1, 8.34×) to align memory access. This is followed by intensive memory hierarchy and parallelism adjustments—including storage scope specification and loop binding—to reach 288.75× (Steps 2–6). Subsequent multi-level tiling and unrolling further optimize data locality and register pressure to 430.68× (Steps 7–12), concluding with fine-grained hardware binding to maximize GPU utiliza457 𝐿 3520 [𝑀 𝑓 32,𝑙 = tion (Steps 13–15). The final optimized LEIR is: 𝐵𝑡𝑥=0 𝑐,𝑡𝑥 𝑐=0 𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑙
𝑓 32,𝑔
𝑓 32,𝑔
457 𝐿 3520 [𝐾 457 3520 𝐴𝑡𝑥,𝑐 ; ]; 𝐵𝑡𝑥=0 𝑐,𝑡𝑥 = 𝑀𝑐,𝑡𝑥 ; ]; 𝐵𝑡𝑥=0𝑈𝑐=0 [𝑁𝑐∗457+𝑡𝑥 = 𝐾𝑐,𝑡𝑥 𝑐=0 𝑓 32,𝑠
𝑓 32,𝑠
𝑓 32,𝑔
𝑓 32,𝑔
457 𝐵 2265 𝑈 3520 [𝐷 457 ; ]; 𝐵𝑡𝑥=0 = 𝐷𝑡𝑥,𝑏𝑥𝑎 + 𝑁𝑐∗457+𝑡𝑥 ∗ 𝐼𝑏𝑥𝑎,𝑐 ; ]; 𝐵𝑡𝑥=0 𝑏𝑥𝑎=0 𝑐=0 𝑡𝑥,𝑏𝑥𝑎 2265 [𝐷 𝑓 32,𝑠 = 𝐷 𝑓 32,𝑠 + 𝐽 𝑓 32,𝑔 ; ]; 𝐵 457 𝐵 2265 [𝐸 𝑓 32,𝑔 = 𝐷 𝑓 32,𝑠 ∗ 𝐵𝑏𝑧𝑎=0 𝑡𝑥=0 𝑏𝑧𝑎=0 𝑡𝑥,𝑏𝑧𝑎 𝑡𝑥,𝑏𝑧𝑎 𝑡𝑥,𝑏𝑧𝑎 𝑏𝑧𝑎 𝑡𝑥,𝑏𝑧𝑎 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 457 𝐵 5 453 [𝐹 𝑓 32,𝑔 𝐻𝑏𝑧𝑎 ; ]; 𝐵𝑡𝑥=0 𝐿 = 𝐸 +𝐶 𝑏𝑥 𝑓 =0 𝑎=0 𝑡𝑥,𝑏𝑥 𝑓 ∗453+𝑎 𝑡𝑥,𝑏𝑥 𝑓 ∗453+𝑎 𝑡𝑥,𝑏𝑥 𝑓 ∗453+𝑎
; ];. Workload of testcases.
(1) InstanceNorm: the known variables are ’A’ with the dtype torch.float16 and shape [67, 27, 77, 10], ’D’ with the dtype torch.float16 and shape [27], ’E’ with the dtype torch.float16 and shape [27], and ’C’ with the dtype torch.float16 and shape [67, 27, 77, 10]. (2) Square matrix multiplication: the known variables are ’A’ with the dtype torch.float32 and shape [319, 319], ’C’ with the dtype torch.float32 and shape [319, 319], and ’D’ with the dtype torch.float32 and shape [319, 319]. (3) conv depthwise 2D square input square kernel: the known variables are ’A’ with the dtype torch.float64 and shape [64, 63, 49, 366], ’D’ with the dtype torch.float64 and shape [63, 1, 2, 2], and ’C’ with the dtype torch.float64 and shape [64, 63, 25, 184].
(4) Mean reduction over a dimension: the known variables are ’A’ with the dtype torch.float16 and shape [16, 256, 256], and ’C’ with the dtype torch.float16 and shape [16, 256]. (5) Gemm Sigmoid Sum LogSumExp: the known variables are ’A’ with the dtype torch.float32 and shape [128, 10], ’G’ with the dtype torch.float32 and shape [20, 10], ’H’ with the dtype torch.float32 and shape [20], and ’F’ with the dtype torch.float32 and shape []. (6) Product reduction over a dimension: the known variables are ’A’ with the dtype torch.float32 and shape [301, 3630, 744], and ’C’ with the dtype torch.float32 and shape [301, 744]. (7) Conv3d HardSwish ReLU Softmax Mean: the known variables are ’A’ with the dtype torch.float16 and shape [120, 2, 8, 22, 56], ’H’ with the dtype torch.float16 and shape [52, 2, 4, 4, 4], and ’G’ with the dtype torch.float16 and shape [120, 52]. (8) Matmul with small K dimension: the known variables are ’A’ with the dtype torch.float16 and shape [1024, 1024], ’C’ with the dtype torch.float16 and shape [1024, 32], and ’D’ with the dtype torch.float16 and shape [1024, 32]. (9) LeakyReLU: the known variables are ’A’ with the dtype torch.float16 and shape [931, 3862], and ’C’ with the dtype torch.float16 and shape [931, 3862]. (10) Matmul with transposed both: the known variables are ’A’ with the dtype torch.float64 and shape [120, 543], ’C’ with the dtype torch.float64 and shape [808, 120], and ’F’ with the dtype torch.float64 and shape [543, 808]. (11) Matmul Subtract Multiply ReLU: the known variables are ’A’ with the dtype torch.float16 and shape [738, 715], ’G’ with the dtype torch.float16 and shape [], ’H’ with the dtype torch.float16 and shape [], ’I’ with the dtype torch.float16 and shape [3251, 715], ’J’ with the dtype torch.float16 and shape [3251], and ’F’ with the dtype torch.float16 and shape [738, 3251]. (12) Gemm BiasAdd Hardtanh Mish GroupNorm: the known variables are ’A’ with the dtype torch.float32 and shape [924, 1220], ’H’ with the dtype torch.float32 and shape [16], ’I’ with the dtype torch.float32 and shape [16], ’J’ with the dtype torch.float32 and shape [16, 1220], ’K’ with the dtype torch.float32 and shape [16], ’M’ with the dtype torch.float32 and shape [16], ’N’ with the dtype torch.float32 and shape [16], and ’G’ with the dtype torch.float32 and shape [924, 16]. (13) Conv3d Softmax MaxPool MaxPool: the known variables are ’A’ with the dtype torch.float32 and shape [69, 38, 38, 6, 45], ’G’ with the dtype torch.float32 and shape [5, 38, 2, 2, 2], and ’F’ with the dtype torch.float32 and shape [69, 5, 37, 5, 44]. (14) GQA sum: the known variables are ’A’ with the dtype torch.float32 and shape [715, 105, 140], ’Ak’ with the dtype torch.bool and shape [1, 1, 105, 105], ’Am’ with the dtype torch.float32 and shape [140, 140], ’An’ with the dtype torch.float32 and shape [140], ’Ao’ with the dtype torch.float32 and shape [35, 140], ’Aq’ with the dtype torch.float32 and shape [35], ’Ar’ with the dtype torch.float32 and shape [35, 140], ’As’ with the dtype torch.float32 and shape [35], ’Aw’
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
with the dtype torch.float32 and shape [140, 140], ’Ax’ with the dtype torch.float32 and shape [140], and ’Aj’ with the dtype torch.float32 and shape [715, 1, 140]. (15) Matmul Mean Softmax: the known variables are ’A’ with the dtype torch.float32 and shape [128, 100], ’F’ with the dtype torch.float32 and shape [50, 100], ’G’ with the dtype torch.float32 and shape [50], and ’E’ with the dtype torch.float32 and shape [128, 1]. (16) Conv3d GroupNorm Mean: the known variables are ’A’ with the dtype torch.float16 and shape [115, 6, 9, 6, 48], ’F’ with the dtype torch.float16 and shape [16, 6, 1, 1, 1], ’G’ with the dtype torch.float16 and shape [16], ’H’ with the dtype torch.float16 and shape [16], and ’E’ with the dtype torch.float16 and shape [115]. (17) L1Norm: the known variables are ’A’ with the dtype torch.float32 and shape [24, 20767], and ’E’ with the dtype torch.float32 and shape [24, 20767]. (18) Gemm Sigmoid Scaling ResidualAdd: the known variables are ’A’ with the dtype torch.float16 and shape [275, 776], ’G’ with the dtype torch.float16 and shape [], ’H’ with the dtype torch.float16 and shape [776, 776], ’I’ with the dtype torch.float16 and shape [776], and ’F’ with the dtype torch.float16 and shape [275, 776]. (19) Matmul Min Subtract: the known variables are ’A’ with the dtype torch.float64 and shape [903, 763], ’F’ with the dtype torch.float64 and shape [], ’G’ with the dtype torch.float64 and shape [], ’H’ with the dtype torch.float64 and shape [2249, 763], ’I’ with the dtype torch.float64 and shape [2249], and ’E’ with the dtype torch.float64 and shape [903, 2249]. (20) Matmul Swish Scaling: the known variables are ’A’ with the dtype torch.float64 and shape [647, 1930], ’G’ with the dtype torch.float64 and shape [], ’H’ with the dtype torch.float64 and shape [2899, 1930], ’I’ with the dtype torch.float64 and shape [2899], and ’F’ with the dtype torch.float64 and shape [647, 2899]. (21) Gemm Swish Divide Clamp Tanh Clamp: the known variables are ’A’ with the dtype torch.float32 and shape [728, 2022], ’J’ with the dtype torch.float32 and shape [1243, 2022], ’K’ with the dtype torch.float32 and shape [1243], and ’I’ with the dtype torch.float32 and shape [728, 1243]. (22) CrossEntropyLoss: the known variables are ’A’ with the dtype torch.float16 and shape [603, 777], ’C’ with the dtype torch.int64 and shape [603], and ’D’ with the dtype torch.float16 and shape []. (23) BMM InstanceNorm Sum ResidualAdd Multiply: the known variables are ’A’ with the dtype torch.float64 and shape [270, 1026], ’C’ with the dtype torch.float64 and shape [270, 497, 2], ’M’ with the dtype torch.float64 and shape [497, 1026], ’N’ with the dtype torch.float64 and shape [497], ’O’ with the dtype torch.float64 and shape [497], ’Q’ with the dtype torch.float64 and shape [497], and ’K’ with the dtype torch.float64 and shape [270, 497, 2]. (24) Conv2d GroupNorm Scale MaxPool Clamp: the known variables are ’A’ with the dtype torch.float64 and shape [128, 3, 32, 32], ’H’ with the dtype torch.float64 and shape [16, 1, 1], ’I’ with the dtype torch.float64 and shape [16, 1, 1], ’J’ with the dtype torch.float64 and shape [16, 3, 3, 3], ’K’
with the dtype torch.float64 and shape [16], ’M’ with the dtype torch.float64 and shape [16], and ’G’ with the dtype torch.float64 and shape [128, 16, 15, 15]. (25) MHA max: the known variables are ’A’ with the dtype torch.float64 and shape [622, 5, 498], ’Ai’ with the dtype torch.bool and shape [1, 1, 5, 5], ’Aj’ with the dtype torch.float64 and shape [1494, 498], ’Ak’ with the dtype torch.float64 and shape [1494], ’Am’ with the dtype torch.float64 and shape [498, 498], ’An’ with the dtype torch.float64 and shape [498], and ’Ah’ with the dtype torch.float64 and shape [622, 1, 498]. (26) LogSoftmax the known variables are ’A’ with the dtype torch.float32 and shape [304, 12782], and ’C’ with the dtype torch.float32 and shape [304, 12782]. (27) LogSoftmax: the known variables are ’A’ with the dtype torch.float32 and shape [304, 12782], and ’C’ with the dtype torch.float32 and shape [304, 12782]. (28) Matrix scalar multiplication: the known variables are ’A’ with the dtype torch.float64 and shape [778, 2385], ’C’ with the dtype torch.float64 and shape [], and ’D’ with the dtype torch.float64 and shape [778, 2385]. (29) MSELoss: the known variables are ’A’ with the dtype torch.float16 and shape [426, 64], ’C’ with the dtype torch.float16 and shape [426, 64], and ’F’ with the dtype torch.float16 and shape []. (30) Swish: the known variables are ’A’ with the dtype torch.float16 and shape [32, 11421], and ’D’ with the dtype torch.float16 and shape [32, 11421]. (31) Gemm Subtract GlobalAvgPool LogSumExp GELU ResidualAdd: the known variables are ’A’ with the dtype torch.float64 and shape [10, 564], ’I’ with the dtype torch.float64 and shape [1313], ’J’ with the dtype torch.float64 and shape [1313], ’K’ with the dtype torch.float64 and shape [1313, 564], ’M’ with the dtype torch.float64 and shape [1313], and ’H’ with the dtype torch.float64 and shape [10, 564]. (32) Average Pooling 1D: the known variables are ’A’ with the dtype torch.float16 and shape [787, 1097, 481], and ’C’ with the dtype torch.float16 and shape [787, 1097, 475]. (33) HardSigmoid: the known variables are ’A’ with the dtype torch.float64 and shape [129, 7733], and ’C’ with the dtype torch.float64 and shape [129, 7733]. (34) Gemm Multiply LeakyReLU: the known variables are ’A’ with the dtype torch.float16 and shape [654, 1472], ’F’ with the dtype torch.float16 and shape [], ’G’ with the dtype torch.float16 and shape [738, 1472], ’H’ with the dtype torch.float16 and shape [738], and ’E’ with the dtype torch.float16 and shape [654, 738]. (35) MHA Gemm ReLU sum: the known variables are ’A’ with the dtype torch.float64 and shape [264, 37, 92], ’Aj’ with the dtype torch.bool and shape [1, 1, 37, 37], ’Ak’ with the dtype torch.float64 and shape [276, 92], ’Am’ with the dtype torch.float64 and shape [276], ’An’ with the dtype torch.float64 and shape [92, 92], ’Ao’ with the dtype torch.float64 and shape [92], ’Aq’ with the dtype torch.float64 and shape [935, 92], and ’Ai’ with the dtype torch.float64 and shape [264, 1, 935].
Step-TP Dataset for Tensor Program Optimization
(36) Max Pooling 1D: the known variables are ’A’ with the dtype torch.float16 and shape [16, 64, 128], and ’C’ with the dtype torch.float16 and shape [16, 64, 62]. (37) Average Pooling 2D: the known variables are ’A’ with the dtype torch.float16 and shape [19, 160, 945, 5], and ’C’ with the dtype torch.float16 and shape [19, 160, 236, 1]. (38) Max Pooling 2D: the known variables are ’A’ with the dtype torch.float16 and shape [16, 32, 128, 128], and ’C’ with the dtype torch.float16 and shape [16, 32, 64, 64]. (39) Conv3d GroupNorm Min Clamp: the known variables are ’A’ with the dtype torch.float32 and shape [79, 47, 59, 35, 5], ’G’ with the dtype torch.float32 and shape [], ’H’ with the dtype torch.float32 and shape [8, 47, 3, 3, 3], ’I’ with the dtype torch.float32 and shape [8], ’J’ with the dtype torch.float32 and shape [8], and ’F’ with the dtype torch.float32 and shape [79, 8, 57, 33, 3]. (40) Matmul MaxPool Sum Scale: the known variables are ’A’ with the dtype torch.float32 and shape [200, 3752], ’I’ with the dtype torch.float32 and shape [], ’J’ with the dtype torch.float32 and shape [2588, 3752], ’K’ with the dtype torch.float32 and shape [2588], and ’H’ with the dtype torch.float32 and shape [200]. (41) MinGPTNewGelu: the known variables are ’A’ with the dtype torch.float32 and shape [568, 11216], and ’J’ with the dtype torch.float32 and shape [568, 11216]. (42) FrobeniusNorm: the known variables are ’A’ with the dtype torch.float16 and shape [106, 29, 29, 89], and ’D’ with the dtype torch.float16 and shape [106, 29, 29, 89]. (43) TripletMarginLoss: the known variables are ’A’ with the dtype torch.float16 and shape [419, 1834], ’C’ with the dtype torch.float16 and shape [419, 1834], ’D’ with the dtype torch.float16 and shape [419, 1834], and ’E’ with the dtype torch.float16 and shape []. (44) MQA mean: the known variables are ’A’ with the dtype torch.float64 and shape [101, 63, 376], ’Af’ with the dtype torch.bool and shape [1, 1, 63, 63], ’Ag’ with the dtype torch.float64 and shape [376, 376], ’Ah’ with the dtype torch.float64 and shape [376], ’Ai’ with the dtype torch.float64 and shape [47, 376], ’Aj’ with the dtype torch.float64 and shape [47], ’Ak’ with the dtype torch.float64 and shape [47, 376], ’Am’ with the dtype torch.float64 and shape [47], ’An’ with the dtype torch.float64 and shape [376, 376], ’Ao’ with the dtype torch.float64 and shape [376], and ’Ae’ with the dtype torch.float64 and shape [101, 1, 376]. (45) Tanh: the known variables are ’A’ with the dtype torch.float32 and shape [776, 5482], and ’C’ with the dtype torch.float32 and shape [776, 5482]. (46) Softsign: the known variables are ’A’ with the dtype torch.float32 and shape [927, 11836], and ’E’ with the dtype torch.float32 and shape [927, 11836]. (47) Matrix Multiplication: the known variables are ’A’ with the dtype torch.float16 and shape [550, 54], ’C’ with the dtype torch.float16 and shape [54, 550], and ’D’ with the dtype torch.float16 and shape [550, 550]. (48) Matmul Mish Mish: the known variables are ’A’ with the dtype torch.float32 and shape [775, 177], ’F’ with the dtype torch.float32 and shape [2435, 177], ’G’ with the dtype
torch.float32 and shape [2435], and ’E’ with the dtype torch.float32 and shape [775, 2435]. (49) Gemm Scaling Hardtanh GELU:the known variables are ’A’ with the dtype torch.float16 and shape [919, 884], ’G’ with the dtype torch.float16 and shape [], ’H’ with the dtype torch.float16 and shape [1447, 884], ’I’ with the dtype torch.float16 and shape [1447], and ’F’ with the dtype torch.float16 and shape [919, 1447]. (50) Matmul Scaling ResidualAdd: the known variables are ’A’ with the dtype torch.float32 and shape [457, 3520], ’C’ with the dtype torch.float32 and shape [457, 2265], ’G’ with the dtype torch.float32 and shape [2265], ’H’ with the dtype torch.float32 and shape [2265], ’I’ with the dtype torch.float32 and shape [2265, 3520], ’J’ with the dtype torch.float32 and shape [2265], and ’F’ with the dtype torch.float32 and shape [457, 2265]. (51) Conv3d Multiply InstanceNorm Clamp Multiply Max: the known variables are ’A’ with the dtype torch.float32 and shape [89, 22, 50, 3, 28], ’C’ with the dtype torch.float32 and shape [], ’D’ with the dtype torch.float32 and shape [], ’M’ with the dtype torch.float32 and shape [3, 1, 1, 1], ’N’ with the dtype torch.float32 and shape [3, 1, 1, 1], ’O’ with the dtype torch.float32 and shape [3, 22, 1, 1, 1], ’Q’ with the dtype torch.float32 and shape [3], ’R’ with the dtype torch.float32 and shape [3], and ’K’ with the dtype torch.float32 and shape [89, 50, 3, 28]. (52) ELU: the known variables are ’A’ with the dtype torch.float16 and shape [784, 29541], and ’C’ with the dtype torch.float16 and shape [784, 29541]. (53) Matmul Swish Sum GroupNorm: the known variables are ’A’ with the dtype torch.float64 and shape [422, 297], ’H’ with the dtype torch.float64 and shape [64], ’I’ with the dtype torch.float64 and shape [64], ’J’ with the dtype torch.float64 and shape [64, 297], ’K’ with the dtype torch.float64 and shape [64], ’M’ with the dtype torch.float64 and shape [64], ’N’ with the dtype torch.float64 and shape [64], and ’G’ with the dtype torch.float64 and shape [422, 64]. (54) Conv2d InstanceNorm Divide: the known variables are ’A’ with the dtype torch.float16 and shape [28, 40, 223, 198], ’F’ with the dtype torch.float16 and shape [], ’G’ with the dtype torch.float16 and shape [12, 40, 1, 1], ’H’ with the dtype torch.float16 and shape [12], ’I’ with the dtype torch.float16 and shape [12], and ’E’ with the dtype torch.float16 and shape [28, 12, 223, 198]. (55) Conv2d Multiply LeakyReLU GELU: the known variables are ’A’ with the dtype torch.float32 and shape [87, 54, 446, 95], ’G’ with the dtype torch.float32 and shape [3, 1, 1], ’H’ with the dtype torch.float32 and shape [3, 1, 1], ’I’ with the dtype torch.float32 and shape [3, 54, 1, 1], and ’F’ with the dtype torch.float32 and shape [87, 3, 446, 95]. (56) Matmul BatchNorm BiasAdd Divide Swish: the known variables are ’A’ with the dtype torch.float64 and shape [800, 824], ’K’ with the dtype torch.float64 and shape [1], ’M’ with the dtype torch.float64 and shape [], ’N’ with the dtype torch.float64 and shape [1], ’O’ with the dtype torch.float64 and shape [926, 824], ’Q’ with the dtype torch.float64 and shape [926], ’R’ with the dtype torch.float64 and shape
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
[926], ’S’ with the dtype torch.float64 and shape [926], and ’J’ with the dtype torch.float64 and shape [800, 926]. (57) Matmul Sum Max AvgPool LogSumExp LogSumExp: the known variables are ’A’ with the dtype torch.float32 and shape [905, 3614], ’J’ with the dtype torch.float32 and shape [1475, 3614], ’K’ with the dtype torch.float32 and shape [1475], and ’I’ with the dtype torch.float32 and shape [905, 1]. (58) Matrix vector multiplication: the known variables are ’A’ with the dtype torch.float64 and shape [784, 778], ’C’ with the dtype torch.float64 and shape [778, 1], and ’D’ with the dtype torch.float64 and shape [784, 1]. (59) Conv2d Min Tanh Tanh: the known variables are ’A’ with the dtype torch.float64 and shape [216, 14, 40, 219], ’H’ with the dtype torch.float64 and shape [30, 14, 2, 2], and ’G’ with the dtype torch.float64 and shape [216, 1, 39, 218]. (60) Gemm GroupNorm Swish Multiply Swish: the known variables are ’A’ with the dtype torch.float64 and shape [581, 1021], ’J’ with the dtype torch.float64 and shape [128], ’K’ with the dtype torch.float64 and shape [128], ’M’ with the dtype torch.float64 and shape [128, 1021], ’N’ with the dtype torch.float64 and shape [128], ’O’ with the dtype torch.float64 and shape [128], ’Q’ with the dtype torch.float64 and shape [128], and ’I’ with the dtype torch.float64 and shape [581, 128]. (61) softmax: the known variables are ’A’ with the dtype torch.float64 and shape [604, 24802], and ’C’ with the dtype torch.float64 and shape [604, 24802]. (62) Gemm Scale BatchNorm: the known variables are ’A’ with the dtype torch.float16 and shape [515, 668], ’H’ with the dtype torch.float16 and shape [363], ’I’ with the dtype torch.float16 and shape [363], ’J’ with the dtype torch.float16 and shape [363, 668], ’K’ with the dtype torch.float16 and shape [363], ’M’ with the dtype torch.float16 and shape [363], ’N’ with the dtype torch.float16 and shape [363], and ’G’ with the dtype torch.float16 and shape [515, 363]. (63) Standard Matrix Multiplication :the known variables are ’A’ with the dtype torch.float16 and shape [268, 222], ’C’ with the dtype torch.float16 and shape [222, 2364], and ’D’ with the dtype torch.float16 and shape [268, 2364]. Standard matrix multiplication: the known variables are ’A’ with the dtype torch.float16 and shape [268, 222], ’C’ with the dtype torch.float16 and shape [222, 2364], and ’D’ with the dtype torch.float16 and shape [268, 2364]. (64) Gemm BatchNorm GELU GroupNorm Mean ReLU: the known variables are ’A’ with the dtype torch.float32 and shape [307, 2520], ’K’ with the dtype torch.float32 and shape [64, 2520], ’M’ with the dtype torch.float32 and shape [64], ’N’ with the dtype torch.float32 and shape [64], ’O’ with the dtype torch.float32 and shape [64], ’Q’ with the dtype torch.float32 and shape [64], ’R’ with the dtype torch.float32 and shape [64], and ’J’ with the dtype torch.float32 and shape [307, 1]. (65) Matmul with irregular shapes: the known variables are ’A’ with the dtype torch.float32 and shape [484, 2508], ’C’ with the dtype torch.float32 and shape [2508, 1224], and ’D’ with the dtype torch.float32 and shape [484, 1224].
(66) GQA mean: the known variables are ’A’ with the dtype torch.float64 and shape [77, 180, 112], ’Ak’ with the dtype torch.bool and shape [1, 1, 180, 180], ’Am’ with the dtype torch.float64 and shape [112, 112], ’An’ with the dtype torch.float64 and shape [112], ’Ao’ with the dtype torch.float64 and shape [56, 112], ’Aq’ with the dtype torch.float64 and shape [56], ’Ar’ with the dtype torch.float64 and shape [56, 112], ’As’ with the dtype torch.float64 and shape [56], ’Aw’ with the dtype torch.float64 and shape [112, 112], ’Ax’ with the dtype torch.float64 and shape [112], and ’Aj’ with the dtype torch.float64 and shape [77, 1, 112]. (67) MHA mean: the known variables are ’A’ with the dtype torch.float64 and shape [88, 189, 196], ’Ah’ with the dtype torch.bool and shape [1, 1, 189, 189], ’Ai’ with the dtype torch.float64 and shape [588, 196], ’Aj’ with the dtype torch.float64 and shape [588], ’Ak’ with the dtype torch.float64 and shape [196, 196], ’Am’ with the dtype torch.float64 and shape [196], and ’Ag’ with the dtype torch.float64 and shape [88, 1, 196]. (68) Matmul with transposed B: the known variables are ’A’ with the dtype torch.float16 and shape [1024, 512], ’C’ with the dtype torch.float16 and shape [768, 512], and ’E’ with the dtype torch.float16 and shape [1024, 768]. (69) Conv2d AvgPool Sigmoid Sum: the known variables are ’A’ with the dtype torch.float64 and shape [58, 33, 34, 29], ’G’ with the dtype torch.float64 and shape [50, 33, 1, 1], and ’F’ with the dtype torch.float64 and shape [58]. (70) Gemm ReLU Divide: the known variables are ’A’ with the dtype torch.float64 and shape [63, 1160], ’F’ with the dtype torch.float64 and shape [], ’G’ with the dtype torch.float64 and shape [1990, 1160], ’H’ with the dtype torch.float64 and shape [1990], and ’E’ with the dtype torch.float64 and shape [63, 1990]. (71) ReLU: the known variables are ’A’ with the dtype torch.float16 and shape [694, 17889], and ’C’ with the dtype torch.float16 and shape [694, 17889]. (72) Matmul Sigmoid Sum: the known variables are ’A’ with the dtype torch.float32 and shape [106, 3856], ’F’ with the dtype torch.float32 and shape [2265, 3856], ’G’ with the dtype torch.float32 and shape [2265], and ’E’ with the dtype torch.float32 and shape [106, 1]. (73) conv depthwise separable 2D: the known variables are ’A’ with the dtype torch.float16 and shape [16, 3, 256, 256], ’E’ with the dtype torch.float16 and shape [3, 1, 3, 3], ’F’ with the dtype torch.float16 and shape [64, 3, 1, 1], and ’D’ with the dtype torch.float16 and shape [16, 64, 254, 254]. (74) Conv2d GELU GlobalAvgPool: the known variables are ’A’ with the dtype torch.float16 and shape [96, 29, 241, 115], ’H’ with the dtype torch.float16 and shape [20, 29, 1, 1], and ’G’ with the dtype torch.float16 and shape [96, 20]. (75) GQA Gemm: the known variables are ’A’ with the dtype torch.float64 and shape [1006, 10, 288], ’Ak’ with the dtype torch.bool and shape [1, 1, 10, 10], ’Am’ with the dtype torch.float64 and shape [288, 288], ’An’ with the dtype torch.float64 and shape [288], ’Ao’ with the dtype torch.float64 and shape [18, 288], ’Aq’ with the dtype torch.float64 and shape [18], ’Ar’ with the dtype torch.float64 and shape [18,
Step-TP Dataset for Tensor Program Optimization
288], ’As’ with the dtype torch.float64 and shape [18], ’Aw’ with the dtype torch.float64 and shape [288, 288], ’Ax’ with the dtype torch.float64 and shape [288], ’Ay’ with the dtype torch.float64 and shape [297, 288], and ’Aj’ with the dtype torch.float64 and shape [1006, 10, 297]. (76) Matmul AvgPool GELU Scale Max: the known variables are ’A’ with the dtype torch.float16 and shape [845, 2795], ’K’ with the dtype torch.float16 and shape [], ’M’ with the dtype torch.float16 and shape [1148, 2795], ’N’ with the dtype torch.float16 and shape [1148], and ’J’ with the dtype torch.float16 and shape [845]. (77) Softplus: the known variables are ’A’ with the dtype torch.float32 and shape [966, 15929], and ’C’ with the dtype torch.float32 and shape [966, 15929]. (78) Max Pooling 3D: the known variables are ’A’ with the dtype torch.float16 and shape [16, 32, 64, 64, 64], and ’C’ with the dtype torch.float16 and shape [16, 32, 30, 30, 30]. (79) Matmul GELU Softmax: the known variables are ’A’ with the dtype torch.float32 and shape [34, 405], ’F’ with the dtype torch.float32 and shape [3679, 405], ’G’ with the dtype torch.float32 and shape [3679], and ’E’ with the dtype torch.float32 and shape [34, 3679]. (80) Matmul for upper triangular matrices: the known variables are ’A’ with the dtype torch.float16 and shape [894, 894], ’C’ with the dtype torch.float16 and shape [894, 894], and ’G’ with the dtype torch.float16 and shape [894, 894]. (81) Sigmoid: the known variables are ’A’ with the dtype torch.float32 and shape [701, 5484], and ’C’ with the dtype torch.float32 and shape [701, 5484]. (82) Max reduction over a dimension: the known variables are ’A’ with the dtype torch.float32 and shape [1182, 276, 82], and ’D’ with the dtype torch.float32 and shape [276, 82]. (83) HuberLoss: the known variables are ’A’ with the dtype torch.float16 and shape [228, 115], ’C’ with the dtype torch.float16 and shape [228, 115], and ’D’ with the dtype torch.float16 and shape []. (84) SELU: the known variables are ’A’ with the dtype torch.float32 and shape [264, 114], and ’C’ with the dtype torch.float32 and shape [264, 114]. (85) Gemm Divide Sum Scaling: the known variables are ’A’ with the dtype torch.float32 and shape [867, 1407], ’H’ with the dtype torch.float32 and shape [229, 1407], ’I’ with the dtype torch.float32 and shape [], and ’G’ with the dtype torch.float32 and shape [867, 1]. (86) Gemm GroupNorm Hardtanh: the known variables are ’A’ with the dtype torch.float64 and shape [576, 1440], ’I’ with the dtype torch.float64 and shape [8, 1440], ’J’ with the dtype torch.float64 and shape [8], ’K’ with the dtype torch.float64 and shape [8], ’M’ with the dtype torch.float64 and shape [8], and ’H’ with the dtype torch.float64 and shape [576, 8, 1, 2]. (87) Matmul GroupNorm LeakyReLU Sum: the known variables are ’A’ with the dtype torch.float32 and shape [132, 2982], ’G’ with the dtype torch.float32 and shape [1024, 2982], ’H’ with the dtype torch.float32 and shape [1024], ’I’ with the dtype torch.float32 and shape [1024], ’J’ with
the dtype torch.float32 and shape [1024], and ’F’ with the dtype torch.float32 and shape [132, 1024]. (88) Matmul with diagonal mat: the known variables are ’A’ with the dtype torch.float16 and shape [377], ’C’ with the dtype torch.float16 and shape [377, 291], and ’E’ with the dtype torch.float16 and shape [377, 291]. (89) Matmul Scale ResidualAdd Clamp LogSumExp Mish: the known variables are ’A’ with the dtype torch.float16 and shape [94, 2974], ’J’ with the dtype torch.float16 and shape [], ’K’ with the dtype torch.float16 and shape [862, 2974], ’M’ with the dtype torch.float16 and shape [862], and ’I’ with the dtype torch.float16 and shape [94, 1]. (90) Matmul Add Swish Tanh GELU Hardtanh: the known variables are ’A’ with the dtype torch.float16 and shape [244, 2994], ’J’ with the dtype torch.float16 and shape [38], ’K’ with the dtype torch.float16 and shape [38], ’M’ with the dtype torch.float16 and shape [38, 2994], ’N’ with the dtype torch.float16 and shape [38], and ’I’ with the dtype torch.float16 and shape [244, 38]. (91) GELU: the known variables are ’A’ with the dtype torch.float64 and shape [480, 15255], and ’C’ with the dtype torch.float64 and shape [480, 15255]. Min reduction over a dimension: the known variables are ’A’ with the dtype torch.float16 and shape [947, 817, 20], and ’D’ with the dtype torch.float16 and shape [817, 20]. (92) Matmul for lower triangular matrices: the known variables are ’A’ with the dtype torch.float32 and shape [935, 935], ’C’ with the dtype torch.float32 and shape [935, 935], and ’G’ with the dtype torch.float32 and shape [935, 935]. (93) L2Norm: the known variables are ’A’ with the dtype torch.float32 and shape [93, 12038], and ’D’ with the dtype torch.float32 and shape [93, 12038]. (94) Conv3d Scaling Tanh Multiply Sigmoid: the known variables are ’A’ with the dtype torch.float16 and shape [61, 62, 9, 54, 3], ’H’ with the dtype torch.float16 and shape [43, 1, 1, 1], ’I’ with the dtype torch.float16 and shape [43, 1, 1, 1], ’J’ with the dtype torch.float16 and shape [43, 1, 1, 1], ’K’ with the dtype torch.float16 and shape [43, 1, 1, 1], ’M’ with the dtype torch.float16 and shape [43, 62, 3, 3, 3], and ’G’ with the dtype torch.float16 and shape [61, 43, 7, 52, 1]. (95) Gemm GroupNorm Min BiasAdd: the known variables are ’A’ with the dtype torch.float32 and shape [76, 1833], ’K’ with the dtype torch.float32 and shape [1, 1, 1, 2], ’M’ with the dtype torch.float32 and shape [1, 1, 1, 2], ’N’ with the dtype torch.float32 and shape [1024, 1833], ’O’ with the dtype torch.float32 and shape [1024], ’Q’ with the dtype torch.float32 and shape [1024], ’R’ with the dtype torch.float32 and shape [1024], and ’J’ with the dtype torch.float32 and shape [76, 1, 1, 2]. (96) HardTanh: the known variables are ’A’ with the dtype torch.float16 and shape [603, 26865], and ’C’ with the dtype torch.float16 and shape [603, 26865]. (97) conv3d conv3d padding: the known variables are ’A’ with the dtype torch.float64 and shape [134, 4, 14, 2013, 9], ’C’ with the dtype torch.float64 and shape [134, 4, 14, 2013, 9], ’G’ with the dtype torch.float64 and shape [1, 4, 1, 2, 2], and
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
’F’ with the dtype torch.float64 and shape [134, 1, 8, 2018, 9]. (98) Conv3d Mish Tanh: the known variables are ’A’ with the dtype torch.float64 and shape [103, 40, 22, 8, 16], ’F’ with the dtype torch.float64 and shape [32, 40, 2, 2, 2], and ’E’ with the dtype torch.float64 and shape [103, 32, 8, 3, 6]. (99) Conv3d Min Softmax: the known variables are ’A’ with the dtype torch.float16 and shape [41, 5, 5, 3, 19], ’G’ with the
dtype torch.float16 and shape [36, 5, 3, 3, 3], and ’F’ with the dtype torch.float16 and shape [41, 3, 1, 17]. (100) Conv2d BatchNorm Scaling: the known variables are ’A’ with the dtype torch.float32 and shape [105, 55, 35, 144], ’H’ with the dtype torch.float32 and shape [], ’I’ with the dtype torch.float32 and shape [6, 55, 6, 6], ’J’ with the dtype torch.float32 and shape [6], ’K’ with the dtype torch.float32 and shape [6], and ’G’ with the dtype torch.float32 and shape [105, 6, 30, 139].
Step-TP Dataset for Tensor Program Optimization
Strategy loop reorder
Description and Example it rearranges the nesting orders of loops within a loop nest. For example, 55 𝐿 277 𝐿 4 𝐿 124 in the expression 𝐵 55 𝐿 277 𝐿 4 𝐿 124 [𝐴𝑑 𝑓 16,𝑔 the loops 𝐵𝑡𝑥=0 = 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 16,𝑔
55 𝐿 124 𝐿 4 𝐿 277 in the expression 𝐴𝑐𝑡𝑥,𝑎,𝑐,𝑑 ; ]; can be reordered as the loops 𝐵𝑡𝑥=0 𝑑=0 𝑐=0 𝑎=0 𝑓 16,𝑔
loop tiling
𝑓 16,𝑔
55 𝐿 124 𝐿 4 𝐿 277 [𝐴𝑑 𝐵𝑡𝑥=0 = 𝐴𝑐𝑡𝑥,𝑎,𝑐,𝑑 ; ];. 𝑑=0 𝑐=0 𝑎=0 𝑡𝑥,𝑎,𝑐,𝑑 The strategy loop tiling is used on the given IR to break two nested loops into four smaller loops (i.e., forming the tiles) within one loop nest. For exam255 𝐿 255 𝐿 255 in the expression 𝐵 255 𝐿 255 𝐿 255 [𝐷 𝑓 16,𝑔 = 𝐷 𝑓 16,𝑔 + ple, the loops 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑐=0 𝑡𝑥=0 𝑎=0 𝑐=0 𝑓 16,𝑔
𝑓 16,𝑔
15 𝐿 3 𝐵 17 𝐿 85 𝐿 255 in the expression 𝐴𝑡𝑥,𝑐 ∗ 𝐶𝑐,𝑎 ; ]; can be tiled as the loops 𝐿𝑔=0 ℎ=0 𝑡𝑥=0 𝑎=0 𝑐=0 𝑓 16,𝑔
loop split
𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
15 𝐿 3 𝐵 17 𝐿 85 𝐿 255 [𝐷 𝐿𝑔=0 = 𝐷𝑔∗17+𝑡𝑥,ℎ∗85+𝑎 + 𝐴𝑔∗17+𝑡𝑥,𝑐 ∗ 𝐶𝑐,ℎ∗85+𝑎 ; ];. ℎ=0 𝑡𝑥=0 𝑎=0 𝑐=0 𝑔∗17+𝑡𝑥,ℎ∗85+𝑎 it divides any loop within a loop nest into two nested loops. For example, the 16 𝐿 16384 in the expression 𝐵 16 𝐿 16384 [𝐸 𝑓 16,𝑔 = 𝐴 𝑓 16,𝑔 /𝐷 𝑓 16,𝑔 ; ]; can be loops 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥=0 𝑎=0 𝑓 16,𝑔
4 16384 in the expression 𝐿 4 𝐵 4 split as the loops 𝐿 4𝑓 =0 𝐵𝑡𝑥=0 𝐿𝑎=0 𝐿 16384 [𝐸 𝑓 ∗4+𝑡𝑥,𝑎 = 𝑓 =0 𝑡𝑥=0 𝑎=0 𝑓 16,𝑔
𝑓 16,𝑔
loop fusion
𝐴 𝑓 ∗4+𝑡𝑥,𝑎 /𝐷 𝑓 ∗4+𝑡𝑥,𝑎 ; ]; it combines multiple loops within a loop nest into one loop. For example, the loops 𝑓 32,𝑔 114 𝐿 491 in the expression 𝐿 2 𝐵 114 𝐿 491 [𝐸 𝑓 32,𝑔 𝐿 2𝑓 =0 𝐵𝑡𝑥=0 = 𝐶 𝑓 ∗114+𝑡𝑥,𝑎 ; ]; can be 𝑎=0 𝑓 =0 𝑡𝑥=0 𝑎=0 𝑎,𝑓 ∗114+𝑡𝑥
loop unrolling
228 𝐿 491 in the expression 𝐵 228 𝐿 491 [𝐸 fused as the loops 𝐵𝑡𝑥=0 𝑎=0 𝑡𝑥=0 𝑎=0 𝑎,𝑡𝑥 = 𝐶𝑡𝑥,𝑎 ; ];. it expands a loop body by computing its equation multiple times. For exam16 𝐿 32 𝐿 64 𝐿 64 𝐿 2 𝐿 2 ple, one loop axis in the loops 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑔=0 of the expression
𝑓 32,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝑓 32,𝑔
16 𝐿 32 𝐿 64 𝐿 64 𝐿 2 𝐿 2 [𝐶 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑔=0 𝑡𝑥,𝑎,𝑐,𝑑 = 𝑚𝑎𝑥 (𝐶𝑡𝑥,𝑎,𝑐,𝑑 , 𝐷𝑡𝑥,𝑎,𝑐∗2+𝑓 ∗3,𝑑∗2+𝑔∗3 ); ]; can be 16 𝑈 32 𝐿 64 𝐿 64 𝐿 2 𝐿 2 unrolled, then the loops become 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑔=0 of the expression
loop parallelization
16 𝑈 32 𝐿 64 𝐿 64 𝐿 2 𝐿 2 [𝐶 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑔=0 𝑡𝑥,𝑎,𝑐,𝑑 = 𝑚𝑎𝑥 (𝐶𝑡𝑥,𝑎,𝑐,𝑑 , 𝐷𝑡𝑥,𝑎,𝑐∗2+𝑓 ∗3,𝑑∗2+𝑔∗3 ); ];. it distributes iterations of a loop across multiple processing units. For 851 𝐿 36 𝐿 979 𝐿 36 example, one loop axis in the loops 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 of the expression 𝑓 32,𝑔
851 𝐿 36 𝐿 979 𝐿 36 [𝐸 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑
= if_then_else(1
<= 𝑐
<
978&2
<= 𝑑
<
𝑓 32,𝑔 851 𝐿 36 𝑃 979 𝐿 36 of 34, 𝐴𝑡𝑥,𝑎,𝑐 −1,𝑑 −2, 0); ]; can be parallel, then the loops become 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 32,𝑔 851 36 979 36 the expression 𝐵𝑡𝑥=0 𝐿𝑎=0 𝑃𝑐=0 𝐿𝑑=0 [𝐸𝑡𝑥,𝑎,𝑐,𝑑 = if_then_else(1 <= 𝑐 < 978&2 <= 𝑑 < 𝑓 32,𝑔 34, 𝐴𝑡𝑥,𝑎,𝑐 −1,𝑑 −2, 0); ];
loop vectorization
it transforms loop operations to use SIMD instructions and thereby processes multiple 178 𝐿 105 𝐿 1 data elements simultaneously. For example, one loop axis in the loops 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑓 64,𝑔 𝑓 64,𝑔 178 105 1 of the expression 𝐵𝑡𝑥=0 𝐿𝑎=0 𝐿𝑐=0 [𝐹𝑡𝑥,𝑎,𝑐,0 = 𝐸𝑡𝑥,𝑎,𝑐 ; ]; can be vectorized, then the loops 𝑓 64,𝑔
loop binding
𝑓 64,𝑔
178 𝐿 105 𝑉 1 of the expression 𝐵 178 𝐿 105 𝑉 1 [𝐹 become 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑡𝑥=0 𝑎=0 𝑐=0 𝑡𝑥,𝑎,𝑐,0 = 𝐸𝑡𝑥,𝑎,𝑐 ; ]; it maps loop iterations to specific GPU threads or blocks along [x, y, z] axes. For exam595 𝐿 71 𝐿 595 of the expression 𝐵 595 𝐿 71 𝐿 595 [𝐸 𝑓 16,𝑔 = ple, one loop axis in the loops 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑡𝑥=0 𝑎=0 𝑐=0 𝑡𝑥,𝑎 𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
595 𝐵 71 𝐸𝑡𝑥,𝑎 + 𝐷𝑡𝑥,𝑐 ∗ 𝐶𝑐,𝑎 ; ]; can be binded, then the loops become 𝐵𝑡𝑥=0 𝐿 595 of 𝑏𝑧𝑎=0 𝑐=0 𝑓 16,𝑔
reduction factorization
𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
595 𝐵 71 the expression 𝐵𝑡𝑥=0 𝐿 595 [𝐸 = 𝐸𝑡𝑥,𝑏𝑧𝑎 + 𝐷𝑡𝑥,𝑐 ∗ 𝐶𝑐,𝑏𝑧𝑎 ; ];. 𝑏𝑧𝑎=0 𝑐=0 𝑡𝑥,𝑏𝑧𝑎 it restructures a reduction expression into multiple reduction expressions by splitting the reduction loop axis and then adds the outputs. For example, 32 𝐿 4 𝐿 2 𝐿 1186 𝐿 5875 of the expression the reduction loop axis in the loops 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑓 32,𝑔
32 𝐿 4 𝐿 2 𝐿 1186 𝐿 5875 [𝐹 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑐=0 𝑑=0 𝑓 =0
𝑓 32,𝑔
= 𝐹𝑡𝑥,𝑎
𝑓 32,𝑔
+ 𝐴𝑡𝑥,𝑎∗2+𝑐,𝑑,𝑓 ; ]; can be split, so the 𝑓 32,𝑔
32 𝐿 4 𝐿 1186 𝐿 5875 of the expression 𝐵 32 𝐿 4 𝐿 1186 𝐿 5875 [𝐹 loops become 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑑=0 𝑓 =0 𝑡𝑥=0 𝑎=0 𝑑=0 𝑓 =0
=
𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 32 𝐿 4 𝐿 1 𝐿 1186 𝐿 5875 of the expression 𝐹𝑡𝑥,𝑎 + (𝐾𝑡𝑥,𝑎 + 𝑀𝑡𝑥,𝑎 ); ]; and 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 32 𝐿 4 𝐿 1 𝐿 1186 𝐿 5875 [𝐾 𝐵𝑡𝑥=0 = 𝐾 + 𝐴 ; ];. Then the outputs are added 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑡𝑥,𝑎∗2+𝑐,𝑑,𝑓 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 32 4 1 1186 5875 by 𝐵𝑡𝑥=0 𝐿𝑎=0 𝐿𝑐=0 𝐿𝑑=0 𝐿 𝑓 =0 [𝑀𝑡𝑥,𝑎 = 𝑀𝑡𝑥,𝑎 + 𝐴𝑡𝑥,𝑎∗2+𝑐+1,𝑑,𝑓 ; ];
Table 8: Operator-level atomic optimization strategies.
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
Strategy cache read/write
Description and Example it moves tensor variables between global (g), shared (s), and local (l) memory (i.e., 𝑓 32,𝑔 shown in the superscript of the variable). For example, the variable 𝐷𝑡𝑥,𝑎 in the 𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑙
413 𝐿 892 [𝐷 expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 𝑙𝑜𝑔(𝐴𝑡𝑥,𝑎 ); ]; can write from local memory as 𝐼𝑡𝑥,𝑎 𝑎=0
413 𝐿 892 [𝐷 𝑓 32,𝑔 = 𝐼 𝑓 32,𝑙 ; ]; and 𝐵 413 𝐿 892 [𝐼 𝑓 32,𝑙 = 𝑙𝑜𝑔(𝐴 𝑓 32,𝑔 ); ];. by two expressions 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥=0 𝑎=0 𝑡𝑥,𝑎
layout transformation
set storage scope
set storage layout
precompute indices
it changes the memory arrangement of tensor variables (i.e., shown in the subscript of the variable). For example, the memory layout of variable 𝐴 can be transformed 11 𝐵 77 𝐿 3182 [𝐸 𝑓 16,𝑔 = 𝐴 𝑓 16,𝑔 as the variable 𝐸 by the expression 𝐿𝑔=0 𝑔,𝑡𝑥,𝑐 𝑡𝑥=0 𝑐=0 𝑔∗77+𝑡𝑥,𝑐 ; ]; it directly sets intermediate tensor variables between global (g), shared (s), and local (l) memory (i.e., shown in the superscript of the variable). For example, the 𝑓 32,𝑔 𝑓 32,𝑠 intermediate variable 𝐸𝑡𝑥 can be set in shared memory as 𝐸𝑡𝑥 . it directly changes the memory arrangement of intermediate tensor variables (i.e., shown in the subscript of the variable). For example, the memory layout of the intermediate variable 𝐷 can be directly set by updating its subscripts (e.g., from 𝑡𝑥, 𝑎 to 𝑎 ∗ 311 + 𝑡𝑥). Note: all occurrences of D in the current IR are updated consistently. it precomputes indices of tensor variable (i.e., shown in the subscript of the variable) to store frequently used index expressions for reuse. For example, the in𝑓 32,𝑔 99 𝐿 26 𝐿 36 𝐿 253 𝐿 109 𝐿 3 𝐿 1 [𝐶 𝑓 32,𝑔 dices c*3+g in the expression 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑔=0 ℎ=0 𝑡𝑥,𝑎,𝑐,𝑑 = 𝐶𝑡𝑥,𝑎,𝑐,𝑑 + 𝑓 32,𝑔
𝑓 32,𝑔
𝑖64,𝑔
𝐸𝑡𝑥,𝑓 ,𝑐∗3+𝑔,𝑑∗3+ℎ ∗ 𝐷𝑎,𝑓 ,𝑔,ℎ ; ]; can be precomputed as the variable 𝐹𝑐,𝑔 by the expres36 𝐿 3 [𝐹 sion 𝐵𝑡𝑥=0 𝑔=0 𝑡𝑥,𝑔 = 𝑡𝑥 ∗ 3 + 𝑔; ]; 𝑖64,𝑔
Table 9: Memory-level atomic optimization strategies.
Step-TP Dataset for Tensor Program Optimization
Strategy factorization
Description and Example it decomposes a mathematical equation into products or sums of simpler com166 𝐿 13601 [𝐶 𝑓 64,𝑔 ponents. For example, the equation in the expression 𝐵𝑡𝑥=0 = 𝑡𝑥,𝑎 𝑎=0 𝑓 64,𝑔
𝑚𝑎𝑥 (0, 𝑚𝑖𝑛(1, (𝐴𝑡𝑥,𝑎 + 3)/6)); ]; can use the factorization strategy to transform 𝑓 64,𝑔
expand factorization
𝑓 16,𝑔
cancellation
𝑓 64,𝑔
166 𝐿 13601 [𝐶 as the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 𝑚𝑎𝑥 (0, 𝑚𝑖𝑛(1, 𝐴𝑡𝑥,𝑎 /6 + 1/2)); ];. 𝑎=0 it reconstructs a factored equation into its original full equation (i.e., the reverse pro7 3545 [𝐺 𝑓 16,𝑔 = cess of factorization). For example, the equation in the expression 𝐵𝑡𝑥=0 𝐿𝑎=0 𝑡𝑥 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 𝐺𝑡𝑥 + (𝐴𝑡𝑥,𝑎 −𝐷𝑡𝑥,𝑎 ) ∗∗2; ]; can use the expand factorization strategy to transform 𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
7 3545 [𝐺 as the expression 𝐵𝑡𝑥=0 𝐿𝑎=0 = 𝐺𝑡𝑥 +𝐴𝑡𝑥,𝑎 ∗ ∗2 − 2 ∗𝐴𝑡𝑥,𝑎 ∗ 𝐷𝑡𝑥,𝑎 + 𝐷𝑡𝑥,𝑎 ∗ 𝑡𝑥 ∗2; ]; it removes variables that offset each other to simplify the equation. For example, 𝑓 64,𝑔 𝑓 64,𝑔 128 𝐿 16 𝐿 30 𝐿 30 [𝐷 𝑓 64,𝑔 the equation in the expression 𝐵𝑡𝑥=0 = 𝐻𝑎 ∗ ((𝐶𝑡𝑥,𝑎,𝑐,𝑑 − 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝐾𝑡𝑥,𝑎 )/𝑠𝑞𝑟𝑡 (𝑁𝑡𝑥,𝑎 +1𝑒 −05))+𝐼𝑎
; ]; can use the cancellation strategy to transform
128 𝐿 16 𝐿 30 𝐿 30 [𝐷 𝑓 64,𝑔 = (𝐶 𝑓 64,𝑔 ∗ 𝐻 𝑓 64,𝑔 − 𝐻 𝑓 64,𝑔 ∗ 𝐾 𝑓 64,𝑔 + as the expression 𝐵𝑡𝑥=0 𝑎 𝑎 𝑡𝑥,𝑎 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 64,𝑔 𝑓 64,𝑔 𝑓 64,𝑔 𝐼𝑎 ∗ 𝑠𝑞𝑟𝑡 (𝑁𝑡𝑥,𝑎 + 1.0𝑒 − 5))/𝑠𝑞𝑟𝑡 (𝑁𝑡𝑥,𝑎 + 1.0𝑒 − 5); ];
expand cancellation
it reconstructs the canceled variables to recover the original equation (i.e., the reverse process of cancellation). For example, the equation in the expression 527 𝐿 11316 [𝐽 𝑓 64,𝑔 = 1 ∗ (𝐶 𝑓 64,𝑔 ∗ 𝐾 𝑓 64,𝑔 + 𝐶 𝑓 64,𝑔 ∗ 𝑀 𝑓 64,𝑔 )/𝑀 𝑓 64,𝑔 ; ]; can use the 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑓 64,𝑔
527 𝐿 11316 [𝐽 expand cancellation strategy to transform as the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑓 64,𝑔
apart
𝑓 64,𝑔
𝑓 16,𝑔
together
=
𝑓 64,𝑔
𝐶𝑡𝑥,𝑎 ∗ (1.0 + (𝐾𝑡𝑥,𝑎 /𝑀𝑡𝑥,𝑎 )); ]; it decomposes the rational fraction part into simpler partial fractions within one equation (i.e., the reverse process of together). For example, the equation in the 339 𝐿 1659 [𝐼 𝑓 16,𝑔 = (𝐾 𝑓 16,𝑔 + 𝑀 𝑓 16,𝑔 )/𝑀 𝑓 16,𝑔 ; ]; can use the apart strategy expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥,𝑎 𝑓 16,𝑔
𝑓 16,𝑔
339 𝐿 1659 [𝐼 to transform as the expression 𝐵𝑡𝑥=0 𝑎=0 𝑡𝑥,𝑎 = 1.0 + (𝐾𝑡𝑥,𝑎 /𝑀𝑡𝑥,𝑎 ); ];. it combines multiple fractions into a single fraction within one equation (i.e., the reverse process of apart). For example, the equation in the expression 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 21 𝐿 9036 𝐿 886 𝐿 9 [𝐶 𝑓 32,𝑔 𝐵𝑡𝑥=0 ∗ ((𝐴𝑡𝑥,𝑎,𝑐,𝑑 − 𝐺𝑡𝑥,𝑎 )/𝑠𝑞𝑟𝑡 (𝐼𝑡𝑥,𝑎 + 1𝑒 − 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 = 𝐷𝑎 𝑓 32,𝑔
05)) + 𝐸𝑎 ; ]; can use the together strategy to transform as the expression 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 21 𝐿 9036 𝐿 886 𝐿 9 [𝐶 𝑓 32,𝑔 𝐵𝑡𝑥=0 ∗ (𝐴𝑡𝑥,𝑎,𝑐,𝑑 − 𝐺𝑡𝑥,𝑎 ) + 𝐸𝑎 ∗ 𝑠𝑞𝑟𝑡 (𝐼𝑡𝑥,𝑎 + 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 = (𝐷𝑎 𝑓 32,𝑔
power simplification
1.0𝑒 − 5))/𝑠𝑞𝑟𝑡 (𝐼𝑡𝑥,𝑎 + 1.0𝑒 − 5); ];. it simplifies the equation by combining and reducing power operations. For example, 174 𝐿 3662 [𝐷 𝑓 32,𝑔 = 𝐴 𝑓 32,𝑔 ∗ (1/(1 + 𝑒𝑥𝑝 (−𝐴 𝑓 32,𝑔 ))); ]; the equation in the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑓 32,𝑔
174 𝐿 3662 [𝐷 can use the powsimp strategy to transform as the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 𝑎=0 𝑓 32,𝑔 𝑓 32,𝑔 𝐴𝑡𝑥,𝑎 /(1 + 𝑒𝑥𝑝 (−𝐴𝑡𝑥,𝑎 )); ];.
expand power simplification
it reverses a simplified power operation into separate power operations within one equation (i.e., the reverse process of power simplification). For example, the 693 𝐿 14522 [𝐷 𝑓 16,𝑔 = 𝐴 𝑓 16,𝑔 /(1 + 𝑒𝑥𝑝 (−𝐴 𝑓 16,𝑔 )); ]; can use equation in the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑓 16,𝑔
693 𝐿 14522 [𝐷 the expand powsimp strategy to transform as the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 𝑎=0 𝑓 16,𝑔
log simplification
𝑓 16,𝑔
𝐴𝑡𝑥,𝑎 ∗ (1/(1 + 𝑒𝑥𝑝 (−𝐴𝑡𝑥,𝑎 ))); ];. it simplifies logarithmic operations using log properties like product, quotient, and power rules within one equation. For example, the equation in the expres65 𝐿 883 𝐿 2399 [𝐹 𝑓 16,𝑔 = 𝑙𝑜𝑔(𝐴 𝑓 16,𝑔 ∗ 𝐶 𝑓 16,𝑔 ); ]; can use the logsimp strategy to sion 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 𝑓 16,𝑔
expand log simplification
𝑓 16,𝑔
𝑓 16,𝑔
65 𝐿 883 𝐿 2399 [𝐹 transform as the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 = 𝑙𝑜𝑔(𝐴𝑡𝑥,𝑎,𝑐 ) + 𝑙𝑜𝑔(𝐶𝑡𝑥,𝑎,𝑐 ); ];. 𝑎=0 𝑐=0 it reverses the simplified logarithm operation into separate log operations within one equation (i.e., the reverse process of log simplification). For example, the equa500 𝐿 644 𝐿 3730 [𝐹 𝑓 32,𝑔 = 𝑙𝑜𝑔(𝐴 𝑓 32,𝑔 ) + 𝑙𝑜𝑔(𝐶 𝑓 32,𝑔 ); ]; can use tion in the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 𝑓 32,𝑔
500 𝐿 644 𝐿 3730 [𝐹 the expand log strategy to transform as the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 = 𝑎=0 𝑐=0 𝑓 32,𝑔 𝑓 32,𝑔 𝑙𝑜𝑔(𝐴𝑡𝑥,𝑎,𝑐 ∗ 𝐶𝑡𝑥,𝑎,𝑐 ); ];.
Table 10: Mathematical-level atomic optimization strategies 1.
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
Strategy collect
Description and Example it combines like operations into a simplified operation in an equation. For example, the 264 𝐿 114 [𝐶 𝑓 32,𝑔 = 1.0507010221481323∗(𝑚𝑎𝑥 (0, 𝐴 𝑓 32,𝑔 )+ equation in the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥,𝑎 𝑓 32,𝑔
𝑚𝑖𝑛(0, 1.6732631921768188 ∗ (𝑒𝑥𝑝 (𝐴𝑡𝑥,𝑎 ) − 1))); ]; can use the collect strategy to 264 𝐿 114 [𝐶 𝑓 32,𝑔 = 1.0507010221481323∗𝑚𝑎𝑥 (0, 𝐴 𝑓 32,𝑔 )+ transform as the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥,𝑎 𝑓 32,𝑔 1.0507010221481323 ∗ 𝑚𝑖𝑛(0, 1.6732631921768188 ∗ (𝑒𝑥𝑝 (𝐴𝑡𝑥,𝑎 ) − 1 ∗ 1)); ];.
expand collect
it splits a combined like operation back into multiple operations in an equation (i.e., the reverse process of collect). For example, the equation in the expression 958 𝐿 1289 [𝐹 𝑓 32,𝑔 = 0.5 ∗ 𝐶 𝑓 32,𝑔 /(1 + 𝑒𝑥𝑝 (−𝐶 𝑓 32,𝑔 )); ]; can use the expand col𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑓 32,𝑔
𝑓 32,𝑔
958 𝐿 1289 [𝐹 lect strategy to transform as the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = (𝐶𝑡𝑥,𝑎 ∗ (1/(1 + 𝑎=0 𝑓 32,𝑔 𝑒𝑥𝑝 (−𝐶𝑡𝑥,𝑎 ))))/2.0; ];.
partially equivalent then correct
it first establishes partial equivalence of several similar expressions by concatenating the input variables, computing the fused expression, and splitting the output variables, as well as finally corrects differences to achieve full equivalence by adding another 274 𝐿 555 𝐿 1353 𝐿 417 𝐿 3 [𝐷 𝑓 16,𝑔 = expression. For example, for the expressions 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
𝐷𝑡𝑥,𝑎,𝑐 + 𝐻𝑡𝑥,𝑑,𝑐∗5+𝑓 ∗6 ∗ 𝐺𝑎,𝑑,𝑓 ; ]; 𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
274 𝐿 555 𝐿 1353 𝐿 417 𝐿 3 [𝐸 and 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑡𝑥,𝑎,𝑐
=
𝑓 16,𝑔
𝐸𝑡𝑥,𝑎,𝑐 + 𝐼𝑡𝑥,𝑑,𝑐∗5+𝑓 ∗6 ∗ 𝐺𝑎,𝑑,𝑓 ; ];, the inputs can be concatenated by 𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
274 𝐿 417 𝐿 6759 [𝐽 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 = 𝐴𝑡𝑥,𝑎,𝑐 ; 𝐽𝑡𝑥,𝑎,𝑐+6759 = 𝐶𝑡𝑥,𝑎,𝑐 ; ];, the similar operations are exe𝑎=0 𝑐=0 𝑓 16,𝑔
𝑓 16,𝑔
274 𝐿 417 𝐿 13534 [𝐾 cuted by 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 = if_then_else(8 <= 𝑐 < 13526, 𝐽𝑡𝑥,𝑎,𝑐 −8 , 0); ]; and 𝑎=0 𝑐=0 𝑓 16,𝑔 274 𝐿 555 𝐿 13532 𝐿 417 𝐿 3 [𝑀 𝑓 16,𝑔 = 𝑀 𝑓 16,𝑔 + 𝐾 𝑓 16,𝑔 𝐵𝑡𝑥=0 ∗ 𝐺𝑎,𝑑,𝑓 ; ];, then the output 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑡𝑥,𝑑,𝑐∗5+𝑓 ∗6 274 𝐿 555 𝐿 1353 [𝐷 𝑓 16,𝑔 = 𝑀 𝑓 16,𝑔 ; 𝐸 𝑓 16,𝑔 = 𝑀 𝑓 16,𝑔 can be split by 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 𝑡𝑥,𝑎,𝑐+12179 ; ];, and finally 𝑓 16,𝑔 274 555 2 417 𝐿 3 [𝐷 𝑓 16,𝑔 the results are corrected by 𝐵𝑡𝑥=0 𝐿𝑎=0 𝐿𝑐=0 [𝐷𝑡𝑥,𝑎,𝑐+1351 = 0; ]𝐿𝑑=0 𝑡𝑥,𝑎,𝑐+1351 = 𝑓 =0 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 274 𝐿 555 𝐿 2 [𝐸 𝑓 16,𝑔 = 𝐷𝑡𝑥,𝑎,𝑐+1351 + 𝐻𝑡𝑥,𝑑,(𝑐+1351)+𝑓 ∗6 ∗ 𝐺𝑎,𝑑,𝑓 ; ]; 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑡𝑥,𝑎,1−𝑐 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 417 3 0; ]𝐿𝑑=0 𝐿 𝑓 =0 [𝐸𝑡𝑥,𝑎,1−𝑐 = 𝐸𝑡𝑥,𝑎,1−𝑐 + 𝐼𝑡𝑥,𝑑,(1−𝑐 )+𝑓 ∗6 ∗ 𝐺𝑎,𝑑,𝑓 ; ];
exponential split
it decomposes an exponential operation into multiple factor operations by introducing an existing variable within one equation. For example, for the ex294 𝐿 32193 [𝐶 𝑓 16,𝑔 = 1/(1 + 𝑒𝑥𝑝 (−𝐴 𝑓 16,𝑔 )); ];, an exponential term pression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑡𝑥,𝑎 𝑎=0 𝑓 16,𝑔
𝑒𝑥𝑝 (if_then_else(𝑎 − 1 < 0, 0, 𝐴𝑡𝑥,𝑎−1 ) can be split from the original equation so 𝑓 16,𝑔
𝑓 16,𝑔
294 𝐿 32193 [𝐶 that the expression 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 = 1/(1 + 𝑒𝑥𝑝 (−𝐴𝑡𝑥,𝑎 − if_then_else(𝑎 − 1 < 𝑎=0 𝑓 16,𝑔
multiplicative split
𝑓 16,𝑔
0, 0, 𝐴𝑡𝑥,𝑎−1 )) ∗ 𝑒𝑥𝑝 (if_then_else(𝑎 − 1 < 0, 0, 𝐴𝑡𝑥,𝑎−1 ))); ]; is obtained. it decomposes one operation into multiple multiplicative factor operations by introducing an existing variable within one equation. For example, for the expression 445 𝐿 1 𝐿 424 𝐿 207 [𝐹 𝑓 16,𝑔 = 𝐾 𝑓 16,𝑔 /𝑀 𝑓 16,𝑔 ; ];, a multiplicative term ∗𝐸 𝑓 16,𝑔 can 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 𝑡𝑥,𝑎,𝑐,𝑑 𝑡𝑥,𝑎,𝑐,𝑑 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 16,𝑔
445 𝐿 1 𝐿 424 𝐿 207 [𝐹 be split from the original equation so that the expression 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 = 𝑓 16,𝑔
additive split
𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
(𝐾𝑡𝑥,𝑎,𝑐,𝑑 /𝐸𝑡𝑥,𝑎,𝑐,𝑑 ∗ 𝐸𝑡𝑥,𝑎,𝑐,𝑑 )/𝑀𝑡𝑥,𝑎,𝑐,𝑑 ; ]; is obtained. it decomposes one operation into multiple additive factor operations by introducing an existing variable within one equation. For example, for the expression 466 𝐿 1 𝐿 1 𝐿 2 [𝐽 𝑓 32,𝑔 = 𝐼 𝑓 32,𝑔 + 𝑀 𝑓 32,𝑔 ; ];, an additive term −𝐼 𝑓 32,𝑔 can be 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 𝑡𝑥,𝑎,𝑐,𝑑 0,𝑎,𝑐,𝑑 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 32,𝑔
466 𝐿 1 𝐿 1 𝐿 2 [𝐽 split from the original equation so that the expression 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 = 𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑔
𝐼𝑡𝑥,𝑎,𝑐,𝑑 + 𝐼𝑡𝑥,𝑎,𝑐,𝑑 − 𝐼𝑡𝑥,𝑎,𝑐,𝑑 + 𝑀0,𝑎,𝑐,𝑑 ; ]; is obtained. Table 11: Mathematical-level atomic optimization strategies 2.
Step-TP Dataset for Tensor Program Optimization
Strategy Description and Example normal loop to prefix loop for max opera- it transforms the maximum operations into online streaming operations where tion the current step is based on the previous step. For example, for the expres𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 156 𝐿 8 𝐿 274 𝐿 274 [𝐴𝑛 𝑓 16,𝑔 sion 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 = 𝐴𝑛𝑡𝑥,𝑎,𝑐 + 𝑒𝑥𝑝 (𝑌𝑡𝑥,𝑎,𝑐,𝑑 − 𝐴𝑚𝑡𝑥,𝑎,𝑐 ); ]; with 𝑎=0 𝑐=0 𝑑=0 the summation operation on the exponential term and its previous expression 156 𝐿 8 𝐿 274 𝐿 274 [𝐴𝑚 𝑓 16,𝑔 = 𝑚𝑎𝑥 (𝐴𝑚 𝑓 16,𝑔 , 𝑌 𝑓 16,𝑔 ); ]; with the maximum opera𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐 𝑡𝑥,𝑎,𝑐,𝑑 𝑎=0 𝑐=0 𝑑=0 tion, online streaming method can be used via updating the next step based on 156 𝐿 8 𝐿 274 𝐿 274 [𝐴𝑚 𝑓 16,𝑔 the previous step by 𝐵𝑡𝑥=0 = 𝑚𝑎𝑥 (if_then_else(𝑑 − 1 < 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
𝑓 16,𝑔
0, −𝑖𝑛𝑓 , 𝐴𝑚𝑡𝑥,𝑎,𝑐,𝑑 −1 ), 𝑌𝑡𝑥,𝑎,𝑐,𝑑 ); 𝐴𝑜𝑡𝑥,𝑎,𝑐,𝑑 = if_then_else(𝑑 − 1 < 0, 1, 𝐴𝑜𝑡𝑥,𝑎,𝑐,𝑑 −1 ) ∗ 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 𝑒𝑥𝑝 (if_then_else(𝑑 − 1 < 0, −𝑖𝑛𝑓 , 𝐴𝑚𝑡𝑥,𝑎,𝑐,𝑑 −1 ) − 𝐴𝑚𝑡𝑥,𝑎,𝑐,𝑑 ) + 𝑒𝑥𝑝 (𝑌𝑡𝑥,𝑎,𝑐,𝑑 𝑓 16,𝑔 156 𝐿 8 𝐿 274 [𝐴𝑛 𝑓 16,𝑔 = 𝐴𝑜 𝑓 16,𝑔 𝐴𝑚𝑡𝑥,𝑎,𝑐,𝑑 ); ]; and writing the output by 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 𝑡𝑥,𝑎,𝑐,273 ; ];.
normal loop to prefix loop for exponential operation
−
it transforms the summation of exponential operations into online streaming operations where the current step is based on the previous step. Note that this trick uses exponential cancellation to find the prefix relation. For example, for the expres933 𝐿 1 [𝐼 𝑓 16,𝑔 = 𝑚𝑎𝑥 (𝐼 𝑓 16,𝑔 , 𝐷 𝑓 16,𝑔 ); ]; with the maximum operation, online sion 𝐵𝑡𝑥=0 𝑡𝑥 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥 𝑓 16,𝑔
933 𝐿 1 [𝐾 streaming method can be used via initializing the output by 𝐵𝑡𝑥=0 𝑎=0 𝑡𝑥,𝑎
=
𝑓 16,𝑔 𝑓 16,𝑔 𝑚𝑎𝑥 (if_then_else(𝑎 − 1 < 0, −𝑖𝑛𝑓 , 𝐾𝑡𝑥,𝑎−1 ), 𝐷𝑡𝑥,𝑎 ); ]; and then updating the next 𝑓 16,𝑔 933 1 step based on the previous step by 𝐵𝑡𝑥=0 𝐿𝑎=0 [𝐾𝑡𝑥,𝑎 = 𝑚𝑎𝑥 (if_then_else(𝑎 − 1 < 𝑓 16,𝑔 𝑓 16,𝑔 0, −𝑖𝑛𝑓 , 𝐾𝑡𝑥,𝑎−1 ), 𝐷𝑡𝑥,𝑎 ); ];.
online softmax
it computes the softmax incrementally by updating the tensor variable step by step within a loop nested. For example, for the softmax expres𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 186 𝐿 1 [𝐼 𝑓 16,𝑔 186 𝐿 1 [𝐽 𝑓 16,𝑔 sions 𝐵𝑡𝑥=0 = 𝑚𝑎𝑥 (𝐼𝑡𝑥 , 𝐷𝑡𝑥,𝑎 ); ]; 𝐵𝑡𝑥=0 = 𝐽𝑡𝑥 + 𝑎=0 𝑡𝑥 𝑎=0 𝑡𝑥 𝑓 16,𝑔
𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 − 𝐼𝑡𝑥 )/𝐽𝑡𝑥 ; ];, the 𝑓 16,𝑔 186 𝐿 1 [𝐼 online softmax can be used by𝐵𝑡𝑥=0 = 𝑚𝑎𝑥 (if_then_else(𝑎 − 𝑎=0 𝑡𝑥,𝑎 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 1 < 0, −𝑖𝑛𝑓 , 𝐼𝑡𝑥,𝑎−1 ), 𝐷𝑡𝑥,𝑎 ); 𝐾𝑡𝑥,𝑎 = if_then_else(𝑎 − 1 < 0, 1, 𝐾𝑡𝑥,𝑎−1 ) ∗ 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 𝑓 16,𝑔 𝑒𝑥𝑝 (if_then_else(𝑎 − 1 < 0, −𝑖𝑛𝑓 , 𝐼𝑡𝑥,𝑎−1 ) − 𝐼𝑡𝑥,𝑎 ) + 𝑒𝑥𝑝 (𝐷𝑡𝑥,𝑎 − 𝐼𝑡𝑥,𝑎 ); ]; and the 186 𝐿 1 [𝐸 𝑓 16,𝑔 = 𝑒𝑥𝑝 (𝐷 𝑓 16,𝑔 − 𝐼 𝑓 16,𝑔 )/𝐾 𝑓 16,𝑔 ; ]; output can be written by 𝐵𝑡𝑥=0 𝑡𝑥,𝑎 𝑎=0 𝑡𝑥,𝑎 𝑡𝑥,0 𝑡𝑥,0
𝑒𝑥𝑝 (𝐷𝑡𝑥,𝑎
𝑓 16,𝑔
− 𝐼𝑡𝑥
𝑓 16,𝑔
186 𝐿 1 [𝐸 ); ]; 𝐵𝑡𝑥=0 𝑎=0 𝑡𝑥,𝑎
=
𝑒𝑥𝑝 (𝐷𝑡𝑥,𝑎
Table 12: Mathematical-level atomic optimization strategies 3.
Mengfan Liu, Da Zheng, Junwei Su, and Chuan Wu
Strategy flashattention without loop tiling
Description and Example it computes attention expressions in an online manner by incrementally calculating scaled dot-products and applying online softmax. For 29 𝐿 2 𝐿 3 𝐿 3 [𝐴𝑥 𝑓 64,𝑔 example, for the softmax expressions 𝐵𝑡𝑥=0 = 𝑡𝑥,𝑎,𝑐 𝑎=0 𝑐=0 𝑑=0 𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
29 𝐿 2 𝐿 3 𝐿 3 [𝐴𝑦 𝑚𝑎𝑥 (𝐴𝑥𝑡𝑥,𝑎,𝑐 , 𝐴𝑐𝑡𝑥,𝑎,𝑐,𝑑 ); ]; 𝐵𝑡𝑥=0 𝑡𝑥,𝑎,𝑐 = 𝐴𝑦𝑡𝑥,𝑎,𝑐 + 𝑒𝑥𝑝 (𝐴𝑐 𝑡𝑥,𝑎,𝑐,𝑑 − 𝑎=0 𝑐=0 𝑑=0 𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
𝑓 64,𝑔
29 𝐿 2 𝐿 3 𝐿 3 [𝐴𝑑 𝐴𝑥𝑡𝑥,𝑎,𝑐 ); ]; 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑
=
and
29 𝐿 2 𝐿 3 𝐿 139 𝐿 3 [𝐴𝑒 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑡𝑥,𝑎,𝑐,𝑑
the
matmul
expression
𝑒𝑥𝑝 (𝐴𝑐𝑡𝑥,𝑎,𝑐,𝑑 − 𝐴𝑥𝑡𝑥,𝑎,𝑐 )/𝐴𝑦𝑡𝑥,𝑎,𝑐 ; ]; 𝑓 64,𝑔
=
𝑓 64,𝑔 𝐴𝑒𝑡𝑥,𝑎,𝑐,𝑑
𝑓 64,𝑔 𝑓 64,𝑔 + 𝐴𝑑𝑡𝑥,𝑎,𝑐,𝑓 ∗ 𝑋𝑡𝑥,𝑎,𝑓 ,𝑑 ; ];, the flashattention without tiling can 29 𝐿 2 𝐿 3 𝐿 3 [𝐴𝑥 𝑓 64,𝑙 be used by 𝐵𝑡𝑥=0 = 𝑚𝑎𝑥 (if_then_else(𝑑 − 1 < 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 64,𝑙 𝑓 64,𝑔 𝑓 64,𝑙 𝑓 64,𝑙 0, −𝑖𝑛𝑓 , 𝐴𝑥𝑡𝑥,𝑎,𝑐,𝑑 −1 ), 𝐴𝑐𝑡𝑥,𝑎,𝑐,𝑑 ); 𝐴𝑧𝑡𝑥,𝑎,𝑐,𝑑 = if_then_else(𝑑 − 1 < 0, 1, 𝐴𝑧𝑡𝑥,𝑎,𝑐,𝑑 −1 ) ∗ 𝑓 64,𝑙 𝑓 64,𝑙 𝑓 64,𝑔 𝑒𝑥𝑝 (if_then_else(𝑑 − 1 < 0, −𝑖𝑛𝑓 , 𝐴𝑥𝑡𝑥,𝑎,𝑐,𝑑 −1 ) − 𝐴𝑥𝑡𝑥,𝑎,𝑐,𝑑 ) + 𝑒𝑥𝑝 (𝐴𝑐𝑡𝑥,𝑎,𝑐,𝑑 − 𝑓 64,𝑙 𝑓 64,𝑔 𝑓 64,𝑔 139 [𝐶𝑎 𝐴𝑥𝑡𝑥,𝑎,𝑐,𝑑 ); 𝐿𝑖=0 = if_then_else(𝑑 −1 < 0, 1, 𝐶𝑎𝑡𝑥,𝑎,𝑐,𝑖,𝑑 −1 ) ∗if_then_else(𝑑 − 𝑡𝑥,𝑎,𝑐,𝑖,𝑑 𝑓 64,𝑙 𝑓 64,𝑙 1 < 0, 1, 𝐴𝑧𝑡𝑥,𝑎,𝑐,𝑑 −1 ) ∗ 𝑒𝑥𝑝 (if_then_else(𝑑 − 1 < 0, −𝑖𝑛𝑓 , 𝐴𝑥𝑡𝑥,𝑎,𝑐,𝑑 −1 ) − 𝑓 64,𝑙 𝑓 64,𝑙 𝑓 64,𝑔 𝑓 64,𝑙 𝑓 64,𝑙 𝑓 64,𝑔 𝐴𝑥𝑡𝑥,𝑎,𝑐,𝑑 )/𝐴𝑧𝑡𝑥,𝑎,𝑐,𝑑 + 𝑒𝑥𝑝 (𝐴𝑐𝑡𝑥,𝑎,𝑐,𝑑 − 𝐴𝑥𝑡𝑥,𝑎,𝑐,𝑑 )/𝐴𝑧𝑡𝑥,𝑎,𝑐,𝑑 ∗ 𝑋𝑡𝑥,𝑎,𝑑,𝑖 ; ]; ]; and 𝑓 64,𝑔 𝑓 64,𝑔 29 𝐿 139 𝐿 2 𝐿 3 [𝐴𝑒 the output can be written by 𝐵𝑡𝑥=0 = 𝐶𝑎𝑡𝑥,𝑎,𝑐,𝑑,2 ; ];. 𝑑=0 𝑎=0 𝑐=0 𝑡𝑥,𝑎,𝑐,𝑑
normal loop to prefix loop for matmul op- it transforms a standard matrix multiplication into an online prefix computation eration based on online softmax based on online softmax applied in previous computations. For example, for the 110 𝐿 8 𝐿 2 𝐿 2 [𝐴𝑠 𝑓 32,𝑔 = 𝑚𝑎𝑥 (if_then_else(𝑑 − 1 < online softmax expressions 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑔
𝑓 32,𝑔
0, −𝑖𝑛𝑓 , 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 −1 ), 𝑌𝑡𝑥,𝑎,𝑐,𝑑 ); 𝐶𝑐𝑡𝑥,𝑎,𝑐,𝑑 = if_then_else(𝑑 − 1 < 0, 1, 𝐶𝑐𝑡𝑥,𝑎,𝑐,𝑑 −1 ) ∗ 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 < 0, −𝑖𝑛𝑓 , 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 −1 ) − 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 ) + 𝑒𝑥𝑝 (𝑌𝑡𝑥,𝑎,𝑐,𝑑 − 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 𝑓 32,𝑔 110 𝐿 8 𝐿 2 𝐿 2 [𝑍 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 ); ]; 𝐵𝑡𝑥=0 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 = 𝑒𝑥𝑝 (𝑌𝑡𝑥,𝑎,𝑐,𝑑 − 𝐴𝑠𝑡𝑥,𝑎,𝑐,1 )/𝐶𝑐 𝑡𝑥,𝑎,𝑐,1 ; ]; and 𝑓 32,𝑔 𝑓 32,𝑔 110 𝐿 8 𝐿 2 𝐿 50 𝐿 2 [𝐴𝑎 𝑓 32,𝑔 the matmul expression 𝐵𝑡𝑥=0 = 𝐴𝑎𝑡𝑥,𝑎,𝑐,𝑑 + 𝑍𝑡𝑥,𝑎,𝑐,𝑓 ∗ 𝑎=0 𝑐=0 𝑑=0 𝑓 =0 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 32,𝑔 110 𝐿 8 𝐿 2 𝐿 2 [𝐴𝑠 𝑓 32,𝑙 𝑁𝑡𝑥,𝑎,𝑓 ,𝑑 ; ];, the online method can be used by 𝐵𝑡𝑥=0 = 𝑎=0 𝑐=0 𝑑=0 𝑡𝑥,𝑎,𝑐,𝑑 𝑓 32,𝑙 𝑓 32,𝑔 𝑓 32,𝑙 𝑚𝑎𝑥 (if_then_else(𝑑 − 1 < 0, −𝑖𝑛𝑓 , 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 −1 ), 𝑌𝑡𝑥,𝑎,𝑐,𝑑 ); 𝐶𝑐𝑡𝑥,𝑎,𝑐,𝑑 = if_then_else(𝑑 − 𝑓 32,𝑙 𝑓 32,𝑙 𝑓 32,𝑙 1 < 0, 1, 𝐶𝑐𝑡𝑥,𝑎,𝑐,𝑑 −1 ) ∗ 𝑒𝑥𝑝 (if_then_else(𝑑 − 1 < 0, −𝑖𝑛𝑓 , 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 −1 ) − 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 ) + 𝑓 32,𝑔 𝑓 32,𝑙 𝑓 32,𝑔 50 [𝐶𝑑 𝑓 32,𝑔 𝑒𝑥𝑝 (𝑌𝑡𝑥,𝑎,𝑐,𝑑 − 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 ); 𝐿𝑖=0 = if_then_else(𝑑 − 1 < 0, 1, 𝐶𝑑𝑡𝑥,𝑎,𝑐,𝑖,𝑑 −1 ) ∗ 𝑡𝑥,𝑎,𝑐,𝑖,𝑑 𝑓 32,𝑙 𝑓 32,𝑙 if_then_else(𝑑−1 < 0, 1, 𝐶𝑐𝑡𝑥,𝑎,𝑐,𝑑 −1 )∗𝑒𝑥𝑝 (if_then_else(𝑑−1 < 0, −𝑖𝑛𝑓 , 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 −1 )− 𝑓 32,𝑙 𝑓 32,𝑙 𝑓 32,𝑔 𝑓 32,𝑙 𝑓 32,𝑙 𝑓 32,𝑔 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 )/𝐶𝑐𝑡𝑥,𝑎,𝑐,𝑑 + 𝑒𝑥𝑝 (𝑌𝑡𝑥,𝑎,𝑐,𝑑 − 𝐴𝑠𝑡𝑥,𝑎,𝑐,𝑑 )/𝐶𝑐𝑡𝑥,𝑎,𝑐,𝑑 ∗ 𝑁𝑡𝑥,𝑎,𝑑,𝑖 ; ]; ]; and the output 𝑓 32,𝑔 𝑓 32,𝑔 110 𝐿 50 𝐿 2 𝐿 8 [𝐴𝑎 can be written by 𝐵𝑡𝑥=0 = 𝐶𝑑𝑡𝑥,𝑎,𝑐,𝑑,1 ; ];. 𝑑=0 𝑐=0 𝑎=0 𝑡𝑥,𝑎,𝑐,𝑑
𝑒𝑥𝑝 (if_then_else(𝑑 − 1
Table 13: Mathematical-level atomic optimization strategies 4.