HloEnv: A Graph Rewrite Environment for Deep Learning Compiler Optimization Research
Authors: Chin Yang Oh
Kunhao Zheng
Bingyi Kang
Xinyi Wan
Zhongwen Xu
Shuicheng YAN
Min Lin
Yangzihao Wang
Venue: Submitted to ICLR 2023 (2023)
Keywords: XLA
Compiler Optimization
Graph Rewrite
Reinforcement Learning Environment

Abstract: We introduce HloEnv, an environment based on Accelerated Linear Algebra (XLA) for deep learning (DL) compiler optimization research. HloEnv transforms all graph rewrites into a common representation, providing a flexible interface to control and modify existing graph optimization passes. In this representation, an XLA pass is converted into a set of sequential rewrite decisions, which control when and if the rewrites are applied. Along with HloEnv, we present a dataset with broad coverage of computation graphs drawn from modern real-world machine learning models. We select two XLA passes with the largest impact on the runtime of the compiled program, and explore the potential for further improvement over XLA in this decision space. We show that using simple heuristics for decision-making can achieve on-par or better performance than XLA. Using search algorithms further boosts performance. We intend for HloEnv and our dataset to be an open-source, community-driven effort that helps spur advances in DL compiler optimization research.

Under review as a conference paper at ICLR 2023
HLOENV: A GRAPH REWRITE ENVIRONMENT FOR DEEP LEARNING COMPILER OPTIMIZATION RESEARCH
Anonymous authors Paper under double-blind review
ABSTRACT
We introduce HloEnv, an environment based Accelerated Linear Algebra (XLA) for deep learning (DL) compiler optimization research. HloEnv transforms all graph rewrites into common representation, providing a flexible interface control and modify existing graph optimization passes. In this representation, an XLA pass is converted into a set of sequential rewrite decisions, which control when and if the rewrites are applied. Along with HloEnv, we present a dataset with broad coverage of computation graphs drawn from modern real-world machine learning models. We select two XLA passes with the largest impact on the runtime of the compiled program, and explore the potential for further improvement over XIA this decision space. We show that using simple heuristics for decision-making achieve on-par or better performance than XLA. Using search algorithms further boosts performance. We intend for HloEnv and our dataset to be open-source, community-dree driien effort that helps spur advances in DL compiler optimization research.
INTRODUCTION
Deep Learning (DL) models have been getting significantly larger and more computationally expen- sive (Thompson et al. 2020 As a result, computational efficiency is now increasingly important for the economic and technical viability, as well as the environmental sustainability of a DL project. DL compiler optimization is important for achieving this efficiency. DL compiler parses user-defined DL model code (usually written in Python) into high-level directed acyelic graph (DAG) that can then be optimized to run efficiently on DI. hardware through a sequence of sub-graph rewrite passes. Current production- readd DL compilers still heavily band-engineered, and require deep domain knowledge to create well-optimized results.
Great efforts have been made to alleviate the reliance human engineers. TASO (Jia et al. 2019c) is the most representative work on search-basee based DL compiler optimization. It automatically generates graph rewrites and searches for better optimization solutions a larger search space. However, the set of DL operators it considers contains only 12 operators, which does not generalize well to newly emerged DL. models. Recent works learning-based DL compiler optimization such as REGAL (Paliwal et al. (2020) and GO (Zhou et al.) (2020a) model a limited set of passes each with different representation. To our best knowledge, there has been no work that generalizes to all optimization passes with a common representation.
In short, at the current stage, research on DL compiler optimization is still facing the following challenges: First, due to their non-unified implementations, there is no systematic interface that has wide coverage optimization types. Second, most existing works focus on specific sets of passes. Third, current DI. compiler optimization benchmarks use either closed-source or small datasets with limited set of DI. models. The community has not yet centered its efforts build a publicly accessible dataset of real-world DL computation graphs.
We propose the following to address these challenges. First, we develop HloEnv, an environment for the optimization agent to inter-operate XLA (Leary & Wang 2017). production-quality cross- framework DL compiler. This environment provides a common representation for any type of graph rewrites. Second, present dataset with broad coverage of High-Level Operations (HLO) graphs drawn from real-world JAX-implemented machine learning code, extracted from a variety of open-source repositories GitHub (Table A.2), with spectrum spans through various domains. This

---

Under review as a conference paper at ICLR 2023
provides more representative dataset of workloads for DL compiler optimization research. Third, based on a thorough analysis of XLA optimization passes, we determine two XLA passes with the most significant impact the runtime of the compiled program. We explore using simple heuristics and search-based algorithms to further optimize these passes.
The design of HloEnv points to potential future where DL compiler engineers only need develop and maintain a simple set rewrite rules and leave the complicated heuristics to machine learning- generated optimization strategies that generalize both new DL models and new DL hardware.
SYSTEM DESIGN OF HL.ONV
2.1 XLA PRELIMINARIES
XLA compiles computation graphs in High-Level Operations (HLO) IR format into machine instruc- tions for different backends. As part of this compilation process, XLA runs series of passes to modify the HLO graph. The passes perform rewrites (using pattern matching and replacement) on the HLO graph to optimize the performance ensure the correctness of the graph. These passes can be composed in pipeline and recursively grouped in parent pipeline. These passes/pipelines are run sequentially in a fixed order and can be run either once or repeatedly in a loop until the pass no longer changes the HLO graph.
2.2 OVERVIEW OF HLOENV
Frontend Optimization Environment
Decision Making Agent
hlo.txt
Modified XLA Pipelines
parse
ON
Heuristics-bast Methods
HLO graph
Gry-mode
augmented HLO graph with altemative nodes
Actions
Search-based Methods
NO
OFF
converged?
Original XLA Pipelines
YES
Optimized HLO Module
Backend Optimizations
Figure 1: The HloEnv interaction loop.
HloEny aims to provide a flexible interface that allows for easy control of the XLA optimization passes and pipelines. Each pass and pipeline in HloEnv can be individually set to dry-mode to allow us to intercept and control the rewrites they perform.
As shown in Fig. T HloEnv's Python interface parses an HLO text file into HLO graph and loads it into the frontend optimization environment. A user-specified set of XIAA passes/pipelines is then applied to the HLO graph. HloEnv executes the original pass/pipeline directly if dry-mode is turned off, while it captures these rewrites without actually applying to the source graph when dry-mode is turned on. An augmented graph that contains both the source graph and all the rewrite opportunities is generated for the user. Using the augmented graph as an input, the user can develop various decision -aaiing agents to decide which rewrites apply. This process can be applied multiple times until the output HLO graphs stay unchanged (converge) or until a decision is made to end the optimization for that pass/pipeline. The user can then use XI.A's backend APIs to generate the final device code.
From the decision-making and control point of view, our system defines a Markov Decision Process (MDP) M = (S, P, R)) S stands for the state space, in our case, the augmented graph. From the state, the agent computes the action in the action space A that decides which rewrite rules to apply. P deseribes the transition function of the HloEnv, i.c., change of the graph when certain rewrite rules are applied. R is the reward generated from the decision, in our case, the improvement of runtime between the old and new graphs.
2

---

Under review as a conference paper at ICLR 2023
HloEnv allows users to design the action space at both macro (the ordering of passes and composition of passes in a pipeline) and a micro level (the ability to decide whether to apply individual rewrites from pass).
2.3 INTERCEPT REWRITES FROM XLA PIPELINES WITH DRY-MODE
Desirably, at each optimization step, the decision-making agent is provided with the source HLO graph, the matched sub sob-graphs of all rewrite rules, and the target sub-graphs that will replace the matched patterns. The actual rewriting of the graph only happens after the agent has made decisions on which rewrites apply.
However, existing XLA passes operate greedily by making an immediate replacement once a rewrite rule has its match. Hence is non-trivial refactor XLA into the above-described optimization iterations. Fortunately, we notice that all rewrites are carried out via only few core APIs that modify the graph. Therefore, by intercepting those graph modification APIs, we introduce a dry-mode to XLA. Instead of making immediate replacements, we save necessary information of all the rewrite opportunities under dry-mode while keeping the source graph unchanged. Given all the opportunities identified from the dry-mode, the agent makes decisions on which rewrites to apply, which get executed by the HloEnv afterward. The dry-mode general to most passes XLA and maximizes the reuse of XLA's existing code by introducing minimal code change. Dry-mode also brings additional positive side effect: parallel execution of passes becomes possible since the source graph read-only under dry-mode.
2.4 THE ALTERNATIVE GRAPH REPRESENTATION
Existing works such as REGAL Paliwal et al. 2019) and GO (Zhou et al. 2020a) use the source graph as the state and predict an action for each node of the graph. This state and action space work well for optimizations like device placement or execution priority. However, they are insufficient for complex graph rewrites like algebraic simplification.
D
XLA Dry Run
CD
Reroute
CD
CD
C
Augment the DAG
Prune
BC
BC
BC
B
A
Figure 2: The alternative graph-based optimization pipeline.
To generalize to any type of graph rewrite, and to enable joint consideration of multiple rewrite rules that potentially conflict with each other, introduce the representation shown in Fig. When a source graph is processed in pass with dry-mode enabled, we capture each rewrite opportunity instead applying them. We then augment the graph with the identified rewrite opportunities using special ternative instruction (yellow triangle), resulting in an alternative graph. This alternative graph serves the state for the agent. After decisions are made on which input path to take for each kAlternativ invennative instruction in the reroute step (details in Section 3), HloEnv applies pruning step to remove all unused nodes and perform additional clean-up.
Interfering Rewrites Commonly, two rewrites the same graph can interfere with each other, i.e., their matched pattern overlaps. Therefore, in existing XLA pipelines, rewrites happen sequentially in pre-defined order avoid race conditions. In pipeline, interfering rewrites are both inserted

---

Under review as a conference paper at ICLR 2023
alternatives. This enables our agent to make more knowledgeable decision based on all available opportunities. Two or more interfering rewrites can cause the resulting graph to violate the acyclic constraint. Hence introduce cycle detection function to reject all such alternatives.
Writing New Passes in XLA The generality of our alternative graph representation means that can easily enable dry-mode for any new pass we create, as long as it utilizes XLA's core APIs. This easily introduce modified versions XLA passes into HloEnv, with larger action spaces for an agent or heuristic to make the rewrite decisions. See our custom General Fusion pass Section 5.2 for an example.
important allows
2.5 CUSTOM HASH FUNCTION FOR HLO GRAPII
We require HLO graph hash function for de-duplicating the dataset uniquely labeling the state when performing a search over the state space. However, the existing hash implementation in XLA does not satisfy needs. It does not fully account for the graph topology the parameters specific to individual instruction types. This can potentially lead to higher rate hash collisions. At the same time, also is affected by irrelevant information such the ordering of the instrcciion/cooppuationns Hence, we developed HlooagHash, a more powerful hash that captures the graph topology and all instruction parameters allow for the unique identification of each graph. The details about its implementation can be found in Appendix
OPTIMIZATION STRATEGIES
Section 2.4 details how we organize alternative rewrites in the alternative graph (i.e. step 2 & 3 in Fig. [ and the Augment method in Algorithm| [bllow). In this section, we utilize HloEnv's ability to control individual rewrites explore alternative optimization strategies, including heuristic-based and search-bard based methods. These methods select one input path for each alternative node, as represented in the GenerateAction method in Algorithm below.
As shown Algorithm[1 I we frame the HLO graph optimization for a certain pass (identified as pass_id) as Markov Decision Process (MDP). Given a certain augmented graph & E S, the corresponding action space is A Ai An, where D is the number of kAlternative nodes in G.
Algorithm Frontend HLO Graph Optimization Pass
1: function OPTIMIZATION PAss(G", pass_id)
G (V, F)
2: Go < Augment(G", pass_id)
3: step +-
while (rsepp * G do
Loop while still having alternative nodes
a <- GenerateAction pass_id)
Action space .4 A
Ai x X A
6:
Gstep+1 ApplyAction(G*tin a)
Step 4 5 Fig.
7:
(1tee+1 + Augment(G pass_id)
8: step step
end while
10: return Ctstep
b The final optimized HLO graph
11: end function
3.1 HEURISTIC-BASER METHODS
Heuristic-based methods use human-designed rules for decisions on each alternative node in the alternative graph. We present simplistic pick-first heuristic. This baseline heuristic always takes the first choice available each kAlternativ ve instruction which leads to graph change (i.e., first alternative after the original sub-graph). This acts as baseline for other methods and the original XLAA pipeline.
4

---

Under review as a conference paper at ICLR 2023
3.2 SEARCH-BASE METHODS
Search-based methods explore multiple actions at each graph state while backtracking is allowed and determine an optimal decision sequence starting from G to an end graph G.¹ We exhibit two search-based methods: beam search (BS) and factorized Monte-Caie tree search (f-MCTS).
Beam Search For graph GH with action space A, beam search (BS) enumerates all next graph states {G ===1 and uses a runtime upper bound pruning rule to discard children violating Eq. in the search tree, equipping our beam search with adaptive bandwidth.
+( T(G").
Our beam search uses depth-firs strategy, implemented with a stack to prioritize search depth (details in Appendix C.2) We use large to approximate an exhaustive search at the cost of a large search
space. Therefore, beam search could only search on small graphs.
Factorized Monte-Carlo Tree Search To deal with graphs with arbitrary sizes, we base our search algorithm on the widely used (Silver et al. [016] 2018; (Schrittwieser et al.) 2020) Monte-Carlo tree search (MCTS). However, our computation graph optimization problem poses new challenges to existing MCTS methods. First, the optimization spaces differ for each computation graph as the action space varies from search node to search node. Second, the actions are naturally factorized in our problem, while the upper confidence bounds in MCTS are usually developed for flattened actions. Third, one search node might have multiple different parent nodes. Therefore, we propose factorized MCTS (MMMTS)) to address the above challenges.
MCTS maintains statistics for visit counts and action (Q) values, representing the expected cumulative reward for taking a specific action at a state. The key idea of f-MCTS is to replace the joint Q value with set of marginal Q values based on the number of kAl MAlterantiie nodes in the current computation graph. In this way, the space complexity for storing all Q values is linear to the number of kAl terantivi intive nodes, and different numbers of factors are allowed for different graphs. Thus, the tree search can operate on dynamic action spaces. In the simulation phase, each kAlternat can make decisions independently. However, in the backup phase, related marginal statisties will be updated to fulfill the marginalization operation (details in Appendix C.3). This way, our method can address the above three difficulties and work efficiently on arbitrarily large graphs.
node
ANALYSIS OF XLA OPTIMIZATION PASSES
The previous section assumes a certain pass is selected and layouts the optimization loop. In this section, our experiments lead us to the design of our action space, namely, which passes of interest.
We conduct pass analysis the optimization passes in XLA's frontend to determine their impact on runtime performance. Thanks to the flexible interface provided by HloEnv, we can easily reorder and disable any pass in our Python analysis script and evaluate its effect the resulting HLO graph's 's runtime. Without loss of generality, we only consider coimiizationffocusse passes (ignoring passes strictly for ensuring runtime correctness) and restrict our analysis NVIDIA GPUs, the most commonly used backend for existing DL compilers. From this analysis, we select optimization passes with the most significant impact to explore how changes in their rewrites can potentially improve performance over XLA's heuristics (Section 3.
Overview There are 222 passes in total, of which 143 passes operate on the HLO graph when compiled for GPU. We ignore correctness-critict passes and select 21 runtime/memory optimization- focused passes for our analysis. For each these passes, we utilize HloEnv to remove all instances of the pass type from the optimization pipeline and measure how this removal changes the runtime of the resulting HLO graph as compared to a fully optimized HLO graph with all optimization passes (see Table
Measurement of Performance Impact We analyze the impact on the performance of an optimiza- tion pass/pipeline on the HLO dataset from two perspectives: the proportion of the dataset affected by that pass (% Affected HLOs). and the average change in performance as a result of that pass (rumtime ratio w/ and w/o the pass). The results can be shown by few metrics presented in Table
5

---

Under review as a conference paper at ICLR 2023
Table 1: Analysis selected XIAA optimization passes. A higher runtime ratio indicates that the pass improved runtime since its absence in the optimization pipeline resulted in higher relative runtime.
% Affected HLOs
Runtime ratio w/ and w/o the pass
% %Impr. (<0.96)
Deggr. (>1.06)
Avg. Ratio (Changed)
Removed Pass/Pipeline
&Changed
Avg. Ratio
ZeroSizeaHooFlmminatto
7.46
0.32
0,46
1,000
1.004
A AlmphraieSimplifeer
36.85
0.92
5.91
1.012
1.033
DotMerger
5.31
0.03
0.36
1.001
1.019
Sortiimllifier
5.35
0.05
0.36
1.001
1.019
TupleSimplifier
5.73
0.08
0.38
1.001
1.021
WhileLoopSimplifer
7.20
0.07
2.24
1.012
1.165
HloCoustantFolding
11.24
0.14
0.63
1.002
1.018
ConditionalSimplifer
5.47
0.14
0.42
1.001
1.020
TransposeFolding
5.42
0.08
0.42
1.000
1.023
AllReduceFoldee
5,46
0.16
0.43
1.001
1.024
AllReduceReassociinn
5.48
0.13
0.42
1.001
1.023
Cudnn BiCaathee
5.54
0.18
0.46
1.001
1.025
5.46
0.15
0.44
1.001
1.023
CublasPadFor Pipeline
5.52
0.18
0,46
1.001
1.023
Gpu ReeuTrreeeccttttoonneeeee
5.71
0.21
0.52
1.001
1.024
Gemmeewriter
8.56
2.03
1.78
1.047
1.552
GemmBroadcan Froodcassoooddnggowiiii
5.54
0.20
1.001
1.024
Fusion Pipeline
49.22
0.25
44.61
1.579
2.178
AlGatherCombieer
5.49
0.28
0.46
1.001
1.019
AllReduccCommiier
5.57
0.26
0.53
1.001
1.023
ReduceSccaterCCmminnr
5.59
0.27
0.51
1.001
1.023
the percentage of graphs that have been transformed by the pass (%Changed) determined by comparing their HloDagHash;
the percentage of graphs that have improved/degraded performance (#llpt//%Degr.));
the average improvement in runtime they result in across all graphs (Avg. Ratio);
the average improvement in performance specifically for the graphs that change when the pass removed, i.e., graphs which the pass affects on (%Avg. Ratio Changed).
Some passes have significant impact on performance the graphs that they affect but only affect minimal number graphs (e.g., Whilel conileelooSSmmliieer)) Hence these passes have a lower average difference in the performance change. Due runtime noise, evaluate graph having improved performance when the relative runtime ratio against XLA (i.e., the runtime of that graph divided by the runtime of the fully XLA optimized graph) is less than 0.94, and degraded performance when is above 1.06 (see Appendix D.3 for more details on how we set these limits).
Passes/Pipelines of Significance There are two passes/pipelines which have the most significant impact on the HLO graphs on which they operate. Hence we choose to focus our experiments on these two passes/pipelines. These are the AlgebruicSimplifer pass and the Fusion pipeline (consisting of a variety of passes related instruction fusion). Of most significance is the Fusion pipeline, which affects the most significant percentage of HLO graphs and results in the largest performance improvement. AlgebraicSimplifer similarly affects large percentage of HLO graphs but results in more negligible general performance improvement (see Table
Insights Results from Table D show that these optimization passes do not always result in runtime improvement. For example, removing the GemmRewriter pass results in 2.1 1 of the HLO graphs showing more than 6% of runtime improvement. This applies even for trivial pass like HloConstantFolding which seems like it should always be applied. Our pass analysis found cases where removal of the HloConstantFolding pass resulted a final graph that ran approximately two times faster (see Fig. in Appendix E.. This demonstrates that there is much room for further optimization in many of the passes and pipelines, even at the macro level of deciding whether to run them given HI.O graph.

---

Under review as a conference paper at ICLR 2023
5 EXPERIMENTS
5.1 DATASETS
HLO graphs in our dataset are of high variance in terms of the number of instructions they contain. To better measure the performance various proposed methods, we developed tool that traverses full HLO graph and extracts random sub-graphs that contain a given range of instruction numbers. Applying this tool to the original HLO text files, we generate a sub-dataset with different ranges of instruction numbers: 10 to 20 (94332 sub-graphs) and 20 to 40 (31 18 sub-graphs). We refer these 2 sub-datasets as inst- 10-20 and inst-20-40. For each of the passes/pipelines analyzed (Algebraic further filter these sub-datasets to obtain experiment set for each
Simplification and Fusion), of these two passes. For each pass, we remove sub-graphs whose runtime correctness is affected by these passes. This allows to focus fully the optimization aspects of a pass. We also filter out optimization cannot conclude in a reasonable amount of
sub-graphs for which our search-heurise time. More details can be found in Appendix A.3
5.2 PASS SELECTION AND MODIFICATION
We select the Fusion pipeline of passes and the Algebraic Simplification pass to evaluate the perfor- mance of our alternative optimization strategies. These passes have the most significant impact on performance, shown in Table I
Algebraic Simplification Pass There are five separate Algebraic Simplification passes at different locations in the optimization pipeline. For our experiments, we selected the third Algebraic Simplifi- cation pass in the entire pipeline for optimization and disabled the other four Algebraic Simplification passes. This pass In contrast, the first Algebraic Simplification pass is located in smaller pipeline that is run multiple times in loop and has potential inter-dependency with these other passes. Selecting this pass for optimization results in a higher percentage of correctness issues; 2) The third Algebraic Simplification pass is run to convergence, i.e., runs multiple times until it no longer modifies the HLO graph. This gives us a larger space for optimization over multiple runs. To obtain a fair comparison, the same four passes were disabled in the XLA pipeline. The resulting pipeline was used obtain the reference results.
selected for two reasons: 1) It can be isolated from the passes before and after.
Fusion Pipeline XLA's Fusion Pipeline consists of various passes (e.g., MultoutputFusion. HorizontalFusion, etc.) that fuse different instructions and computations patterns. These passes are sequentially run fixed order. Additionally, they contain many hand-written heuristics that determine whether fusion should occur. In our alternative graph-based representation, however, the philosophy is to keep both alternatives and delay the heuristics to the routing step. Therefore, remove heuristic decisions from XLA's existing pipeline but keep only the rewrite rules. As testimony, we introduce a General Fusion pass that is heuristics free to replace the existing fusion pipeline. General Fusion has much shorter lines of code than the original fusion passes and provides a much larger search space.
5.3 EVALUATION AND METRICS
We note T(G) the runtime, given a graph G. Based on our experiment results on profiling of the runtime noise (Appendix D.3), we define T(G) = both pass/pipeline, we compute the runtime ratio of the final optimized graph w.r.t. XLA p = T (GGmetioa/TT(xxxt
For
and report its avg, max and min across the dataset. We also report the proportion of the graphs with performance compared against XLA with criteria. We statistically set < 0.94 (resp. p 1.06) as criteria for faster (resp. slower) than XLA to avoid false positives caused by noise (Appendix D.3 and Fig. 5) We identify Gmethod = GXLA using our custom Hl HlssahHash (Section (.5)) and report the proportion of graphs that end with equal hash. To further reduce the impact of noise on the maximum and minimum runtime ratios, we retest the graphs with the largest and smallest relative runtime ratios remove outliers.
7

---

Under review as a conference paper at ICLR 2023
Table Results on all sub-datasets of 2 passes. i-MCTS factorizd MCTS with uniform prior.
% of graphs of runtime ratio
Runtime ratio
Identical XI.A (equal hash)
Pass
Dataset
Method
W.r.t XLA
XIA
Faster 0.94)
Slower (> 1.06)
Avg.
Max.
Min.
pick-first
1.000
1.158
0.740
0.1
0.04
96.2 %
inst-10-20
BS
0.981
1.442
0.383
11.2 %
1.0 %
34.9 %
Alg-Simp
pick-first f-MCTS
1.000
1.197
0.883
0.19 %
0.25
92.2 %
imst-20-40
0.995
1.449
0.570
5.3 %
1.2 %
55.5 %
pick-first
0.997
1.709
0.679
3.1 %
0.3 %
26.4 %
inst-10-20
General
BS
0.986
1.252
0.430
5.2 %
0.3
11.7 %
Fusion
pick-first
0.989
3.313
0.738
13.4 %
1.9 %
11.6 %
insl-20-40
f-MCTS
0.992
1.755
0.349
14.5 %
7.7
4.9 %
5.4 RESULTS ANALYSIS
Finding exemplars of performance gains over XLA In cases where the alternative optimization strategy generated graph with faster performance, we performed additional evaluations, using T(G)= Sain..evvllutee (a0laatt(0).sssncc__aa 100).asyyn_tiimnnng confirm that their performance gains not due to noise. These cases are shown below as proof that an alternative optimization strategy can achieve significant performance gains over XLA.
5.4.1 HEURISTIC- BAUSED METHODS
On both inst-10- mo-t-10-20 and inst-20-40 (where graphs are not exhaustively scarchable), a pick-first heuristic serves strong baseline.
Algebraic Simplification Table 5.4 shows that the pick-first heuristic performs on par with XLA in the Algebraic Simplification pass, with more than 90 % of the final pick-first optimized graphs being identical to XLA's. This is expected as in the Algebraic Simplification case, pick-first after dry-run has high probability of recovering the behavior of original XLA.
General Fusion For the general fusion pass, the pick-first heuristic performs even better than XIAA (i.e. with average runtime ratio w.r.t. XLA of less than for both inst-10-20 and inst-20-40. In this case, we are comparing XLA's fusion pipeline, which contains variety of heuristics used to evaluate if fusion should be performed, with our modified general fusion pass, which removes most of these heuristics. Hence, the good performance of the pick-first heuristic on both sets shows that large portion of graphs can be optimized well despite ignoring XLA's hand-desggned heuristics and simply fusing everything.
5.4.2 SEARCI--BAS BASED METHOSS
In general, using search-based methods resulted in faster average runtime over XLA (as well as heuristic-based methods) for both the inst-10-20 and inst-20-40 sets for both the Fusion and Algebraic Simplification passes. We remove pruning from beam search on inst- 10-0 for exhaustive search.
Algebraic Simplification On inst-/0-20, using search-based methods to optimize the HIO graphs showed general reductions in their evaluation timings, with superior performance compared to the pick-first heuristic. On average, BS-optimized graphs were 1.9% faster than the equivalent XLA- optimized graphs, with the most optimized graphs performing up to 161.1% better (Fig. 7 On slightly worse than inst- inst-10-22 10-20, with f-MCTS being average 0.5 %
inst-20-40, performance faster than XLA. Nevertheless, the most optimized graphs were still significantly faster than XLA and ran up to 75.4% faster (Fig. g)
We examine the differences between the best-performing search-optimized and XLA-optimized HLO Graphs to understand the source of these performance gains. In these cases, performing certain
8

---

Under review as a conference paper at ICLR 2023
Algebraic Simplification rewrites either directly improve performance, allow for subsequent passes to make better optimizations (e.g., by allowing for later Fusion pass to fuse more instructions into a single kernel, or allowing different cuDNN call to be used). Exemplars of these cases can be found in Appendix F.T
General Fusion On inst-nn--20 using search-based methods showed general reductions in their evaluation timings. On average, best-first search optimized graphs were 1.42% faster than the equivalent XLA-optimized graphs, with the best graphs performing 132.5% faster (Fig. 13). On inst-20-40, J-MCTS has approximately equal performance to the pick-first heuristic. This is despite the best performing f f-MCTS graph running 186% faster than the XLA optimized equivalent (Fig. 14]. One reason for this could be that greedily performing as much fusion as possible is optimal in most cases.
As we did with the Algebraic Simplification results, we examine the best pest-performing -atthhoop graphs and compare them to XLA-optimized graphs to understand why they performed better. These cases can be split into two general types: 1) Trivial fusion cases, where the speed up happens due to the instructions being fused into a smaller number of kernels and 2) Non-trivia fusion cases where the number of kernels is the same or fewer, but changes to the topology of the resulting graph result in faster runtime. Exemplars of both these types can be found in Appendix F.2
RELATED WORKS
The efficient deployment of various DL models on diverse DL hardware is a challenging task (Jia
cet al., 2019a; Jouppi et al., 2017; Chen al.] 2014). It has driven development and research DL compilers such as TVVM (Chen et al.) 2018), XLA (Leary & Wang 2017). DLVM (Wei et al.
2017), nGraph (Cyphers et al. 2018), and Glow (Rotem et | 2018). Several third-party partt systems and
methods have also been proposed to enlarge the search space (Jia et al.) 2019b Wang et al. 2021 I).
include better optimization strategies (Looks al. 2017 Liu al. 2019; Zheng al. 2020), or providing new high-level IR (Roesch et al.) 2019). More recently, people have started exploring learning-based methods for DL compiler frontend optimization (Paliwal et al. 2020 Zhou al. 2020b). CompilerGym (Cummins et al. (022) and MLGO (Trofin et al. 2021)) are two works related to ours the traditional compiler optimization field. They provide Python environments for users to interact with a selected subset of LLVM optimization passes. TensNet (Zheng et al. 2021 is dataset for tensor programs that target lower-level optimization for tensor compilers.
7 DISCUSSIONS AND CONCLUSIONS
We develop HloEnv, the first (to our knowledge) general interface that inter-operates with a production- quality DL compiler. HloEny provides a common representation for graph rewrites. We utilize this common representation to test alternative optimization strategies for specific, high-impact passes on broad range of HLO graphs. We show that these alternative strategies can achieve on-par or better results than the DL compiler's ' native optimization strategy. Allowing the user to control the action space of the graph rewrites provides more generic and flexible (as result, more challenging) setup. This action space can be restricted (e.g., if users only select certain passes of interest) or, in principle, infinite (by defining and introducing new passes with novel graph rewrite rules). To supplement research with HloEnv, we also generate a large-scale HLO graph dataset to act as an ideal testbed for computation graph optimization.
We hope that HloEnv and this HLO graph dataset provide valuable tools to the community to spur progress developing DL compiler systems. Building on the baseline presented in this paper, we believe that a well-designed action space (system research) and a well-trained agent (MI. research) are both essential for this purpose. More specifically, we hope for HloEnv to enable DI. compiler development in the following directions: 1) open up opportunities for more types of decision-making agents that improve native optimization passes for existing DL compiler systems such as XLA; 2) less effort to introduce new passes as the heuristics will be replaced by the optimization agent, e.g.. General Fusion, and 3) lead to the development of learning- based policies to generate optimization strategies that generalize to new DL models running on new DL hardware.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Riyadh Baghdadi, Massinissa Merouani, Nohamed-Hichale Leghettas, Kamel Abdous, Taha Arbaoui, Karima Benatchba, and Saman P. Amarasinghe. A deep learning based cost model for automatic code optimization. In Alex Smola, Alex Dimakis, and lon Stoica (eds.), Proceedings of Machine Learning and Systems 2021, MLSys 2021, virtual, April 5-9. 2021. mlsys.org, 2021. URL https: httpss///cocccdnng..mmmm mlsys. org wttp://aae//e////a///////// 3def184ad 1408
Tianqi Chen, T. Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Q. Yan, Haichen Shen, M. Cowan, Leyuan Wang, Yuwei Hu, L. Ceze, Carlos Guestrin, and A. Krishnamurthy. Tvm: An automated end-to-end optimizing compiler for deep learning. 2018. URL.nttsss/ww..ennnticcch..... semanticscholar
//wp::
org/paper/df@1 3a1 45t 0400///0885555000444444333444430
Yunji Chen, Tao Luo, Shaoli Liu, Shijin Zhang, Liqiang He, Jia Wang, Ling Li, Tianshi Chen, Zhiwei Xu, Ninghui Sun, et al. Dadiannao: A machine-learning supercomputer. In 2014 47th Annual IEEE/ACM International Symposium on Microarchitecture, pp. 609-622 IEEE, 2014.
Chris Cummins, Bram Wasti, Jiadong Guo, Brandon Cui, Jason Ansel, Sahir Gomez, Somya Jain, Jia Liu, Olivier Teytaud, Benoit Steiner, et al. Compilergym: Robust, performant compiler optimization environments for ai research. In 2022 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), pp. 92-105, IEEE, 2022.
Scott Cyphers, Arjun K. Bansal, Anahita Bhiwandiwalla, Jayaram Bobba, Matthew Brookhart, Avijit Chakraborty, William Constable, Christian Convey, Leona Cook, Omar Kanawi, Robert Kimball, Jason Knight, Nikolay Korovaiko, Varun Kumar Vijay, Yixing Lao, Christopher R. Lishka, Jaikrishnan Menon, Jennifer Myers, Sandeep Aswath Narayana, Adam Procter, and Tristan J. Webb. Intel ngraph: An intermediate representation, compiler, and executor for deep learning. CoRR, abs//a0l...../... 808808058, 2018. URL
8888 08058
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
Zhe Jia, Blake Tillman, Marco Maggioni, and Daniele Paolo Scarpazza. Dissecting the graphcore ipu architecture via microbenchmarking arXiv preprint arXiv:1912.03441. 2,0,,,
Zhihao Jia, Oded Padon, James J. Thomas, Todd Warszawski, M. Zaharia, and A. Aiken. Taso: optimizing deep learning computation with automatic generation of graph substi- tutions. Proceedings of the 27th ACM Symposium on Operating Systems Principles, 2019b. doi: 10.1145/3341301 3359630.. URL http:: //www. www.seeennnncllarrrreeeee org/ /ttpppppeaaarrrppppppppppee/ /
Zhihao Jia, James Thomas, Todd Warszawski, Mingyu Gao, Matei Zaharia, and Alex Aiken. Opti- mizing dnn computation with relaxed graph substitutions. Proceedings of Machine Learning and Systems, 1:27-39, 2019c.
Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of tensor processing unit. In Proceedings of the 44th annual international symposium on computer architecture, pp. 1-12, 2017.
Rasmus Munk Larsen and Tatiana Shpeisman. Tensorflow graph optimizations, 2019. URL https //research. google/ e/ htte://pbs//uueee pub4b011/
Chris Leary and Todd Wang. Xla: Tensorflow, compiled, 2017.
Yizhi Liu, Yao Wang, Ruofei Yu, Mu Li. Vin Sharma, and Yida Wang. Optimizing cnn model infer- on cpus. In Dahlia Malkhi and Dan Tsafrir (eds.), USENIX Annual Technical Conference, pp. 1025-1040. USENIX Association, 2019. URL httt:://dbll uni seiitrrier.cm .edterdde/dae de/ cttp /usenix/ usenix2019 html #L 1UWYLSW19
10

---

Under review as a conference paper at ICLR 2023
Moshe Looks, Marcello Herreshoff, DeLesley Hutchins, and Peter Norvig. Deep learning with
dynamic computation graphs. In ICLR (Poster). OpenReview.net, 2017. URL http://dblp
Roman Novak, Lechao Xiao, Jiri Hron, Jaehoon Lee, Alexander A. Alemi, Jascha Sohl-Dickstein, and Samuel S. Schoenholz. Neural tangents: Fast and easy infinite neural networks in python. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https:///ppenrvview htt.://enweeeee..eeffooooo///oos nett:/orummid-kkkk3yyyrtt
Aditya Paliwal, Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, and Oriol Vinyals. Reinforced genetic algorithm learning for optimizing computation graphs. arXiv preprint MXXiv:990.002455 2019.
Aditya Paliwal, Felix Gimeno, Vinod Nair, Yujia Li, Miles Lubin, Pushmeet Kohli, and Oriol Vinyals. Reinforced genetic algorithm learning for optimizing computation graphs In International Conference on Learning Representations, 2020. URL. httpps////openreeeeeeeeeffrrrrre
[BYPP
Jared Roesch, Steven Lyubomirsky, Marisa Kirisame, Josh Pollock, Logan Weher; Ziheng Jiang, Tianqi Chen, Thierry Moreau, and Zachary Tatlock. Relay: A high-level IR for deep learning.
CoRR, abs/1904.08368, 2019. URL http:/pp///axxtoorr////baa/
083668
Nadav Rotem, Jordan Fix, Saleem Abdulrasool, Garret Catron, Summer Deng, Roman Dzhabarov, Nick Gibson, James Hegeman, Meghan Lele, Roman Levenstein, et al. Glow: Graph lowering compiler techniques for neural networks, arXiv preprint wwxiv:8805000900 2018.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, and David Silver. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 2020.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panncershelvam, Marc Lanctot, Sander Dicleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature, 529(7587):484 489, jan 2016. ISSN 0028-0836. doi: 10. 1038///ree 16661.
David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Simonyan, and Demis Hassabis. A general reinforcement learning algorithm that masters chess, shogi, and go through self-play. Science, 362(6419):1140 1144, 2018. doi: 10. 126sscience aar6404. URL httpps://ww. http://wwwwwwiiinccccc wrw.ci/ncc..rgg/o///ab/////////e
aar6404
Neil C. Thompson, Kristjan Greenewald, Keeheon Lee, and Gabriel F. Manso. The computational limits of deep learning, 2020. URL
.05558 cite arxiv:2007.0555: 005508
Mircea Trofin, Yundi Qian, Eugene Brevdo, Zinan Lin, Krzysztof Choromanski, and David Li. MLGO: a machine learning guided compiler optimizations framework. CoRR, 1ttp//2110040000 04808
2021. URL hnttps://arxiv.orr///eo
Haojie Wang, Jidong Zhai, Mingyu Gao, Zixuan Ma, Shizhi Tang, Liyan Zheng, Yuanzhi Li, Kaiyuan Rong. Yuanyong Chen, and Zhihao Jia. Pet: Optimizing tensor programs with partially equivalent transformations and automated corrections. 2021. URL httttt:////ww wttp:/w..mmmtttccooo.rrr.rrrrrr. org/
paper/ele 27 1ewee44ec16444488 160.e666666004422666666666666666
Richard Wei, Vikram S. Adve, and Lane Schwartz. DLVM: A modern compiler infrastructure for deep .00016
learning systems. CoRR,
2017. URL htttp://arxi..org//aaaaa/
Lianmin Zheng, Ruochen Liu, Junru Shao, Tianqi Chen, Joseph E. Gonzalez, Ion Stoica, and Ameer Haj Ali. 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 l), 2021. URL httt::///pppmeevieee.etttfrrrrrn

---

Under review as a conference paper at ICLR 2023
Zhen Zheng, Pengzhan Zhao, Guoping Long, Feiwen Zhu, Kai Zhu, Wenyi Zhao, Lansong Diao, Jun Yang. and Wei Lin. Fusionstitching; Boosting memory intensive computations for deep learning workloads. ArXiv, abs/2009 2020. URL Utt/https://////
https::///w// http:::/www..eemmnniccccoco...r
3B2bc@111Be7cc6555ebbb/4 04045556607774444444444444499119
Yanqi Zhou, Sudip Roy, Amirali Abdolrashidi, Daniel Wong, Peter Ma, Qiumin Xu, Hanxiao Liu, Phitchaya Phothilimtha, Shen Wang, Anna Goldie, et al. Transferable graph optimizers for ml compilers. Advances in Neural Information Processing Systems, 33:13844-13388 2020a.
Yanqi Zhou, Sudip Roy, AmirAli Abdolrashidi, Daniel Wong, Peter C. Ma, Qiumin Xu, Hanxiao Liu, Mangpo Phitchaya Phothilimtha, Shen Wang, Anna Goldie, Azalia Mirhoseini, and James Laudon. Transferable graph optimizers for ML compilers. CoRR, abs/2010. 12438, 2020b. URL
12438
12

---

Under review as a conference paper at ICLR 2023
A DATASET INFORMATION
The goal of our HLO dataset is twofold. First, we want to present large-scale dataset for training different computation graph optimization strategies. Second, we want the dataset to serve as an ideal bed against which people could measure the performance of arbitrary computation graph optimization strategies.
A.1 DATASET COLLECTION
We manually select a list of JAX implemented repositories from GitHub, and harvest the HI.O text files by setting the XLA DUMP TO flag while running the model. In this way, we dump all the unoptimized HLO graphs generated during JAX's lust-In-Time (JIT) compilation process. We then remove duplicate HLO text files by comparing hash using our HloDAGHash implementation (see Appendix [B for more details), and filter the resulting files to remove the very small ones which have minimal opportunities for optimization. After the above steps, we can guarantee three properties of HLO graphs in our dataset: 1) They all come from real-world deep learning models; 2) They all have different DAG and tensor shapes, and 3) They all provide at least some space for optimization opportunities. In total, we build dataset containing 40,711 HLO graphs from deep learning models defined in 26 distinguished GitHub repositories.
A.2 DATASET OVERVIEW AND ANALYSIS
208%
53%
73%
0:%
32%
200%
54%
5E%
641
74%
52%
60%
X0%
100%
45%
0%
A5%
33%
55%
16%
54%
och
62%
200%
54%
13%
2%
26%6
9%
22%
34%
22
20%
125
100
115
Besprett
Jedde
M
Aad
Figure 3: Distribution of the sizes of HLO graphs
Our dataset covers a broad range of network architectures, centering around modern, real-world models various domains. The GitHub repositories we pick include Vision Transformer (ViT) (Doso- vitskiy et al. 2021), Neural-Taangent (Novak et al.) 2020), MuZero (Schrittwieser et al. 2020), and more. Figure shsooss that most HLO graphs generated during JAX'sIIT compilation contain less than 1000 instructions. More than half of the repository set majorly contains dataset with less than 200 instructions. We also show a breakdown of top HLO opcodes that appear in our dataset in Figure
13

---

Under review as a conference paper at ICLR 2023
exponential
negate
1.0%%
custo I
broedrast
22.2%
concatinnate
18%
ped
trrapose
7.%%
divide
2.3%
idta
aod
2.0%%
del
3.4%
compare
4.4%
reshace
4.7%
select
muttply
7.4%
5.4%
Figure Distribution of the Ops our HI. dataset.
A.3 SUB-DATASET FILTERING
Algebraic Simplification The Algebraic Simplification pass is used for both optimization and correctness purposes. Hence it possible that modifications the rewrites applied by this pass result in graph that cannot be correctly evaluated on the backend. Of the 94332 inst-10-22 10-20 sub-graphs generated, we filtered out 34000 graphs for which the correctness of the graph was dependent on any graph rewrites performed by the Algebraic Simplification pass. Filtering inst-20-40 results in 1619 graphs
General Fusion Our implementation of a General Fusion pass allows almost any two instruc- tions/computations to be fused. Additionally, fusion result the cloning of computation. As result, the action space for General Fusion is significantly larger than for the Algebraic Simplification pass. To ensure search algorithms can finish in reasonable time, we empirically filter graphs from the sub-datasets with action space less than 10,000 size after the first optimization pass. Note that the number of instructions in graph significantly increase multiple rewrite operations during a search. This filtering leads to inst- inst-10-20 containing 28,507 graphs and inst-20-40 containing 2791 graphs.
B DETAILS ON THE IMPLEMENTATION OF THE HloDAGHash FUNCTION
The XLA HLO graph hash implementation is lacking two ways which increase the number of hash collisions: 1) simply hashes the instructions in the HLO graph in post-order, and does not recursively consider the structure and connections of each HLO instruction and computation in the HLO graph; 2) Instruction specific parameters (e.g. the size and stride of an HLO Convolution instruction) are not considered in the hash of each instruction well.
Our custom HI HLoDAGHash function builds upon XLA's hash implementation, but designed to more powerful hash that additionally accounts for graph topology and the parameters unique each
14

---

Under review as a conference paper at ICLR 2023
Table 3: Commit hashes and GitHub URLs of open-source repositories we used to generate our
dataset.
GitHub Repo Name
Commit Hash
URL
BayesNewton GANs-JAX
e3a7251
AaltoML
099111f
antibiotic-rrsistance
7634f9c
magi-1/anntiiiciicrrresstteee
bayeSDE
43511d2
xwinxu/bayeeDE
brax
730e05d
google/brax
continuation-jax
c145260
harshs8ff/continmattionnjja
cwvae-jax
d724112
julluskunzee/wwweeejaa
deeperwin
5c8d497
mdsuniviee/deeperwiin
deepmind-research
1642ae3
deepminnd//eeemiindrresaacc
dks
a92c184
deepmind/dks
dm_pix
6accc96
deepmind/dm_pix
efax
61ad838
NeilGirdhar/efax
efficientnet-jax
a65811f
rwightmmaneeffiiiiinttne----aa
flaxvision
d5b7e6c
rolandgwccFllaxvision
jax-bayes
b91432c
jameswuc/jax-bbaes
jax-enhance
3a3dd40
ssaacorrleyjjax-ehhhnce
jax-flows
26dce81
Chriswaitess/jax-flows
jax-gat
4d48e2b
gcucurull/jax-gat
jax-md
2b8754a
google/jax-md
jax-models
ae57505
Dershhanbessipannd//jjix-omelll
jax-rl
820cb5d
henry-prioor/jax-rl
jax-unirep
b8048db
ELArkk/jax-unirep
jax_muzero
b8ab362
jax_verify
e02aa65
deepmind/jax..eriffy
jaxfg
893c347
brentyi/jaxfg
jaxlie
65d6351
brentyi/jaxlie
jaxrl
1a86300
ikostrikov/jarl
jraph
36071d5
deepmind/jraph
maml_flax
a4a8819
gcucurull/maml_flan
medclip
ff9e5f6
Kaushalya//edcllip
metax
12f2478
tristanndeleu/mmeeax
mlp-gpt-jax
571ccf0
lucidrainns/mlppggtt-jea
neural-tangents
5bb274c
google/neural-tangenss
ott
f4dafa8
ott-jax/ott
revisiting_rrainoow
7cd2bc6
Johansanii//eevisttingggaaaaaai
scenic
978fea0
google-ressearch//sceic
siren-jax
0806e61
KeunwooPark/siren-jja
tree-math
46214b5
google/tree-math
Scale-Mixtures-
441 4d12
yuneg11/Scale-iiiiuuuss-------
of-NNGP
Toy-neural-
rasutt/Toy-neural-
3a1 d5d8
network-in-jax
network-in-jax
Information-fusion.
NuX
bd996dd
Lab-Umass/Nux
ASSNe0//Plysicc--infrmeed
PINN-JAX
561a8d9
neural-network-in--NN
Differenntiableenivvrree
jax_cosmo
52ca009
Initiative/jax.cosmo
15

---

Under review as a conference paper at ICLR 2023
instruction. This reduces the chance of hash collision when determining if graph has been before, or is identical to another graph.
Implementation The HloDAGHash algorithm walks the HLO graph, starting from the root instruc- tion, in Depth-First Traversal order. At each instruction, we take the original XLA hash of that instruction and additionally hash with two things. The first is the HloDAGHash of each operand of that instruction (unlike the XLA hash which hashes the shape of each operand of that instruction). In this way, we can ensure that two HIO graphs with the same post-order of instructions, but different structures, will not have the same hash. The second is that the XLA hash of each HLO instruction is further hashed with the attributes specific to that instruction opeode (e.g. the slice starts, limits, and strides of an HLO Slice instruction). This further decreases the chance of hash collision between two differing HLO graphs.
16

---

Under review as a conference paper at ICLR 2023
C DETAILS ON ALTERNATIVE OPTIMIZATION STRATEGIES
C.I NOTATIONS
We note G = EV the HLO computation graph parsed from HLO text file by our utility; & (V, F)) the alternative graph augmented from G, V = U [d].. with D KAlternative nodes. In reinforcement learning setup, & corresponds to the state. The action space given & (or G) is denoted as _A = A1 ........................... x AD. For all d = 1,..., D, Ad = {k + A E V] and |Aa = in_degree (d). We note an action = (ax) A We omit the superseript for A and note A when there is no ambiguity.
In our sequential decision problem, use superscript for G and to denote the search steps: Starting from a graph G we apply an action a to produce the next graph G²; we repeat the above for N steps to get the final graph G.V, We note T' the function for calculating the running time of graph G described in D.3 specifically, Eq. 7
C.2 BEAM SEARCH
We maintain stack S with no size limit and a global minimum runtime Tmin during the search, has one element G initially. At every search step t, we pop from the graph G*; we apply all possible actions E A to its alternative graph G to obtain new graphs {G runtime of each new graph and only push back into S with criteria
1 [/4
Ji=1i evaluate the
+( < az T(G*).
(2)
controls the pruning when runtime degradation happens. The number of new graphs being pushed back is capped by a pre-fixed expand budget. The search ends when the stack S is empty or a global timeout is triggered. Beam search becomes exhaustive when we push all new graphs into the stack (equivalent to setting a +00 with no expand budget) at every step.
When the search finishes, we extract the graph with runtime Tiin and its trajectory starting from G
Beam search can achieve optimization at the cost of large search space as the number of kAl MAlterrative nodes increases (e.g. for graph with 10 kalternative nodes with each node choices, the search space of size 210). Therefore, an exhaustive search is only feasible on graphs with considerably small number of alternatives.
C.3 FACTORIZED MCTS
Factorized MCTS (f-MCTS) maintains a search tree decide which action take to transit from Go to G++1 For a trajectory (G", ...., G.)), f-MCTS maintains N search trees, each with root node Go ..... G-- Without loss of generality, we present the algorithm for one search tree with root node Go.
We note (G. a) the state-action pair and G' the graph obtained after applying a on G. We define the reward function as follows:
(G),a,G)
T(G").
The action value function for G and (a1...., ap) is represented by Q(G, 01,,.,,...000 aD), which grows exponentially as the number of alternative nodes increases. To deal with the joint action space A, we propose to replace the joint Q with set of marginal Q value function Qa(G, ad) EagazalQ(G, (I,,,,.,.... ..... an)). d = 1, ..... D. Each Qa(G, aa) represents the expected value if only one action a, for d-th alternative node taken. In this way, we can select actions for different alternative nodes independently.
During the search procedure, we associate each search node with computation graph state G. Each search node maintains set of statistics (T) (Na(G, aa), Qa(G, aa), Pa(G, aa) !d=..... representing the running time T, marginal visit counts N. marginal action value Q and factorized policy P for each of the alternative vertices. For each action there is an edge (G, G') storing the transition information and the corresponding reward R. The search repeats the following three stages for given number of budgets.
17

---

Under review as a conference paper at ICLR 2023
Selection: We superscript & to denote the search depth in the tree. The root node is thus given by G All simulations start from the same root graph state G and finish when a leaf graph G' is achieved cycle is formed. For each time-step k along the search path, joint action a is obtained by selecting each o² according to the upper confidence bound (UCB) score described below:
VEAu Na(G,ba) 02 arg max Qa(G, ad) + Pa(G,as) 1 Na(G,
( (c1
01 + log EBg N(C,ba) 02 +1
(4)
Pa is prior policy while Qa accumulates knowledge from simulations. C1 and 02 are two hyperpa- rameters to trade off the relative importance of Pa and Qa. At the beginning of a search, UCB relies more the prior policy but gradually moves its attention to value statistics. In our experiments, we choose 01 = 1.25 and C2 = 19652 following AlphaGo (Silver et al. 2016).
Expansion happens when computation graph is visited for the first time in the search tree, ie., when a simulation terminates. Consider a terminal transition (G a--1_ G*), new node representing G will be created and added the search tree. Once prior policies d= for kAlternative nodes { and value function ty (G) obtain the value of are given. The node statistics will be initialized to (Na(G', aa) = 0, Qa(G", ad) = 0, Pa(G", ad) = PPla=i The running time is set to T = T( The reward for the current transition is also initialized R(G'-1 a--11 G*) T T Note that, for the expansion of the root note representing Go. there will not be a reward as transitions exist.
Expansion:
Backup: Each simulation generates search path {Go, G G'..... G² The statistics of nodes/graphs along this path need to be updated in reverse order. Let r" denote the reward for transition aa-1, G*), and y be the discounting factor. The (e - kt-step return estimation at k-th step is given by
8-1- GK=
Σ ==1
plt++ + ue
(5)
where of is the value for G For k = l,...,1, 1.0 we update the marginal statistics for each (G', (f))=t....) as follows:
Qa a Na(Gk, ,a)) Qa(C*, a²) + +
Ail Md=1.... I....,D;
Na(Gk,a) a
(6)
Na a = Na(G*, a0)) Ad = 1..... D.
However, the reward and value might have an arbitrary scale in our setting. We propose to normalize the values such that Q [0, to get stable calculation of the UCB score. To this end, keep
track of the minimum (Qmin) and maximum (Qmax) values observed in the search tree. A normalized
Q value is thus obtained by Q Que Qum When we calculate the UCB score in Eq. 4 we are actually using normalized Q instead of un-normalized ones in Eq.
Q-Qain
18

---

Under review as a conference paper at ICLR 2023
D EXPERIMENTS
D.1 HARDWARE AND SOFTWARE ENVIRONMENT
We empirically found that competing processes running on the same machine major source of noise in the runtime evaluation of a graph. To get the best estimation of runtime in a real-world environment, we directly evaluate the runtime of an HLO graph on a clean bare-metal GPU node with minimal other processes running. The GPU node has two AMD EPYC 7352 24-Core processors (with hyper-threading 96 cores), 512GB main memory, and eight 40GB memory NVIDIA A100 GPUs. All tests on Ubuntu 20.04 with CUDA 11.2, cuDNN 8.1.1, and TensorFlow 2.9.1.
D.2 XLA VERSION
HloEnv, along with all our experiments presented in this paper, was developed from the following version of XLA htttp://gggtha..cmm///errrr///// commit/ 0b07441177760ea555444442777f2822
D.3 PROFILING AN HLO GRAPH
To evaluate the effectiveness of any given optimization strategy. it is critical to get an accurate runtime oracle Ω Optimized, HLO O_Moo_leMeeee-- -> Runtime. To approximate oracle Q researchers in the community either build cost model or directly evaluate the runtime. The cost model is either learming-basee (i.e. trained from supervised dataset) (Baghdadi al. [222]] or rule-based. The latter requires large amount of engincering effort as it needs to predict the runtime without evaluation, e.g. Grappler (Larsen & Shpeisman 2019) for Tensorflow Graph. On the other hand, although direct runtime evaluation often suffers from real-world noises, given an environment with sufficient computing resources where noise can be controlled, it provides a way to do an accurate evaluation with minimum cost. In this paper, we use the direct runtime evaluation.
To profile the runtime of an HLO graph we need to obtain both the executable and parame- ters. We obtain the executable by calling the standard compiler provided by XLA while setting run uun_backed__ow _olly to prevent the re-invocation of HLO passes. For parameters, we randomly gener- ate (0, 1) for floating-poin parameters and fill const values for other types. A fixed random seed used to keep the parameters consistent across the optimization process so that we can verify the correctness of optimizations.
Reducing timing noise There is random variation in the evaluation timing of an HLO graph. Additionally, when an executable runs multiple times, the initial run is consistently much slower than subsequent runs of that executable. To reduce this noise in the evaluation timing, we evaluate the executable multiple times. The first three runs are treated as warm-up runs and are ignored, and the executable is then evaluated at least 10 additional times. Experiments showed that running the evaluation more than 10 times did not significantly reduce the variance in the final determined runtime. We then take the minimum of the timing results across all runs. We take the minimum of the results instead of the average due to the half-normal distribution of the timing.
Additionally, obtain three different measurements of the evaluation timing, with each being progressively more fine-grained:
full execution timing: The time measured from the moment the evaluation begins till when it concludes;
asynchronous evaluation timing: The time taken from the asynchronous dispatch of the computation to the moment it returns;
compute timing: The time spent in nanoseconds for the execution, without accounting for data transfer.
Experiments showed that the full execution timing measurement resulted in more evaluation timing noise, while the compute timing measurement was too fine-grat fine-grained and missed out some of the performance improvements as a result of memory-related optimizations. Hence, asynchronous evaluation timing is utilized the main timing metric for experiments.
19

---

Under review as a conference paper at ICLR 2023
Thus, the formula for obtaining the runtime formally reads:
T(G) = min miinGG.vvallatt(ll..ysss.tttiiid
Additionally, it was determined that noise was higher when both GPUs coupled to single NUMA node were utilized. Hence when obtaining our experimental results, we ensured that only a single GPU in each NUMA node was utilized (four out of eight GPUs on the bare-metal system total).
Profiling effects timing noise on evaluation of relative graph performance In our experiments, we have to frequently evaluate the relative performance of two HLO graphs, for instance in comparing whether the heuristic-based optimized graph performs better than XLA (Section or the relative change in performance when particular XLA optimization pass is removed from the full optimization pipeline as seen in Table
To determine what relative ratio be used determine with confidence that one HLO graph has faster run-time than another HLO graph, we profile the expected noise seen when evaluating the relative run-time ratio of two HLO graphs under the same conditions our experiments (i.e. only single GPU utilized per NUMA node). This is done by evaluating the same HLO graph twice and taking the ratio of the first runtime divided by the second runtime. This is repeated 500,000 times to obtain a distribution of the expected range in runtime ratios for a given HLO graph. We perform this profiling on 10 different graphs, spanning the range of runtimes (approximately 25000 ns to 1000000 ns)
in HLO dataset
From this distribution, we determine upper and lower bounds for the ratios, above/below which we can say with reasonable confidence that degradatiion/mprovenni This is evaluated by determining the ratios above and below which 99.9 of the data points lie. From our results, we can see that any run-time ratio < 0.94 and above 1.06 likely represents an actual change performance (Fig.
in run-time is not due to noise.
Non-empirical evaluation using HLO graph Cost Analysis The impact of optimization pass on an HLO graph can also be estimated by performing an HLO Cost Analysis on the resulting HLO graph, and seeing how the module changes in the metries of 1. number of FLOPs, 2. number of Transcendentals, and 3. Bytes accessed.
D.4 MODEL ARCHITECTURE AND RUNNING TIME
Beam Search The pruning factor in Eq. 1 is set to 20 cover most cases.
Factorized MCTS The simulation budget is set to 400. We decay to half (capped by 50) each step after an action is taken, and produce a new graph along the decision sequence. The budget decay is based the observation of the size shrink of action space after the first several steps. The value function is given by the Monte-Carlo evaluation. We launch 5 rollouts length 10 based on uniform sampling. We report the running time of f-MCTS with uniform prior for reference only, as most computation is done on the CPU side while only the environment and Monte-Carle evaluation require GPU. /MCTS with uniform prior takes 50 A 100 days insl-20-40.
20

---

Under review as a conference paper at ICLR 2023
Mean run-time: 24166 ns
Mean run-time: 25447 ns
3000
3000
0.962
0.944
1.039
2000
1.055
2000
1000
1000
0.9
1.0
1.1
1.2
Mean run-time: 28606 ns
Mean run-time: 29818
0.948
1500
1.057
1.041
1000
500
0.8
1.0
1.1
1.2
0.8
0.9
1.0
1.1
1.2
33514
run-time:
3000
0.969
2000
1.063
1.034
1000
1000
0
0.8
0.9
1.0
1.1
1.2
0.8
0.9
1.0
1.1
1,2
Mean run-time: 60626 ns
Mean run-time: 160943 ns
4000
0.971
0.987
1.026
4000
2000
2000
0.8
1.2
0.8
0.9
1.0
1.1
1.2
Mean run time: 586672 ns
Mean run timti 905120 ns
15000
0.996
0.992
10000
5000
5000
1.1
1.2
Figure 5: profile showing the distribution of the runtime ratio noise of different HLO graphs. The blue and green vertical lines encapsulate 99.9% of the ratios.
21

---

Under review as a conference paper at ICLR 2023
E
PASS ANALYSIS GRAPH EXAMPLES
XLA Optimized Graph without HioConstanlFolding Pass
XLA Optimized Graph
Figure 6: No HloConstantFolding pass runtime/XLA runtime 0.51 (96.1% % faster). This graph too large to display its details, but the overall structure both graphs can be seen to be visibly different. This example of how the removal of simple pass like HloConstantFolding can cause compounding differences in the final result as the other passes/pipelines are applied.
22

---

Under review as a conference paper at ICLR 2023
F SEARCH OPTIMIZED VS XLA OPTIMIZED GRAPH EXAMPLES
F.1 ALGEBRAIC SIMPLIFICATION
In these examples, the beam search and fMMCTS optimization strategies outperform XLA by removing specific graph rewrites that directly impact performance or allow for later passes to better optimize the graph.
Beam Search Optimized
XLA Optimized
ed
Figure 7: BS runtime/XI.A runtime 0.383 (161.1% faster). By not performing some optimizations
during the Algebraic Simplification pass. The reshape instruction before the final broadcast instruction does not get optimized out, making later fusion passes able to fully fuse the HLO graph into one computation.
23

---

Under review as a conference paper at ICLR 2023
f-MCTS Optimized
XLA Optimized
Computation wwwppetaiionneenncooooo whw..eeeeooooooobreeee.ooooo
Computation entry compatation
Parmuerter
Parsmeter
expression fusion kind KLoop
Purmaucter J
241160590222281111 (2.100))
cestams-cail cablas butal porsnn) Toel E4000, 1,3 2.1.0)
884 (641L:0550228.]11.22
hatch coppi-4tt0) (400) Iee bucch dita-n ())
Parameter [84 4850,28,28 22, (2,1
copa.1 [64] 48500 1J11 (0.2]
baxc, dimpo Lontrtatiue discam111
Parameter
reshapu,
Fesed cxpy medirs faskon Londnitoop
14030022,,8 (2.1100)
[64]] (61114000,23):1..u
H6ll 40100 1 1.20)
trauspose.2 dimensione- 10.2.11
brendeast.22 dimmensions- 10,21 20-4 14055028,
(64) 14050L2H, (1.2,00)
trampascal
Aireroketentn (0.2.1)
Holl 14000.20, (1.20)
multiply.
16(14450022228)112222
slioer (I] 14000) (8.283, (H1))) (1
164 140500 00.50..88.211
Niccett
reduce-4 Subcomputation: add dimansierie (1)
reslegpe.12 Fivi 14009 (UUU
(6-[1-4050.28] 1.00
operand 1641]
reskegee.1
264 2641100008.1112...0
resllage.t Li (66 Jit
slice.26
slice- [[1: [1:10005] (0-28] [o-I]:
364) 14009.26, 11,.1,12...
So schoomppttion add Bisl L4ber 28, 1 (111.0) epuecand (1-01106410000000 13.1)
operand Laoap
seloct
2 16-11 (04.9.2611
MiA 1-E000 THE 11
operand
2- L4040122.. 11
Vynamic-update-aee fi4j 14050, h8, 11 1.0) operand U- [84]144150226, (050028611 operand operand
sperand the Buqj Wodt
uperand operand (32)
sliee.22 25
slice- [1::] [1:1 8000] (0:28]] [0:1]
poddinggo) the Dud (11
fid] 14049, 26 11
ROOT
ROOT
Figure 8: FMCCTS runtime/XLA runtime = 0.57 (75.4% faster). The f-MCTS optimized graph preserves the map instruction and batch-gemm custom call.
24

---

Under review as a conference paper at ICLR 2023
Beam Search Optimized
XLA Optimized
Co carry
Compasater virn, wwegeseeeeeee tinee
BIG MIALASI
Cag Bor-
DOLL DER AAL!!
opmer & GRRTAND
ELAIT
BOOE
Figure 9: BS runtime/XLA runtime 0.71 (40.8% faster). this case, the beam search optimization results in different set of instructions and custom calls.
25

---

Under review as a conference paper at ICLR 2023
Beam Search Optimized
Computation collry_ www.:e/wlooooooooooo chsseclose chono clano clane
(7
22 11.0
loep hieton jLoopp
Parameter 1 +1.(6
Purometter (1
Parmindter x12/64.]]]11
offict dass (1), collapeod slice dans citz
(0, stat isde (0), indor veetor dan 01(042211.00)
dite (1), collapsed, sike dins shox sinowe
start indox inlleeeeeyeett Jin
reslope.2 6: 12111
rollpee55 [T (1-100
62264,2111.00 L.00
E2 32264,2111.00 1.!
broadensl.
broadeast 3 22 (2.44)
sebttrecD
(2(64/M22
add 22144442 (2.1.))
sperand 00.
ROOT
XLA Optimized
Compatation computaaions.llee
Parameter
Parsenter
Parameter
4nj
offict dinns(t), collopsed slice, dinss (0), aurt uxtex, sup- (9), index, vestor_ 11.2
oftset dane collapeed siice, dins 19), stunt index, _nap" 191, index, ventorddint din-t sloc. esspe: 1:21
Fused exprespor for fasten fiiion kind-kL.c
Parameier
Parameicr
D2
dimemommemlj) popp444-11.. 11
Almorsimenn 1 21:11 1.1
subrracLD 1216
B2164 S4 21 1.) operaed I (2646421 0,5
BOOT
Figure 10: BS runtime/XLA runtime 0.79 (26.6 % faster). By not performing optimizations
during the Algebraic Simplification pass, later fusion passes can fully fuse the HIOO graph into one computation.
26

---

Under review as a conference paper at ICLR 2023
F.2 FUSION
F.2.1
TRIVIAL EXAMPLES
In these examples, the beam search reduces the number of computations as compared to XLA by performing additional fusions of instructions and/or computations.
Beam Search Optimized
XLA Optimized
Computation entrr_cootrppeeeeeeeeeeeeeeeeeee entrr_computatiiocctorssse www.ppptiiiooooooo....o.oooooeee
Computation entry computation.clonn
Parameter 0
Parameter
Parameter
Parameter
[64[]
[64[]
f64[]
[64[]
Fuse expression for fusion.2
loop fusion
kind-kLoo kLoo
divide.18511
f64[]
Parameter 0 164[
Parameter
[64[
Fused expressin for fusion loop fusion
kind
divide.0
Parameter f64[]
164[
broadcast.3 dimensions-11 I [64[140500]
broadcast.0 dimensions- E [64[14050]
multiply.1
multiply.0 164 14050]
(65(0)
operand 0= constant.3 [64[14050]
operand 0 constant.3 f64[14050]
subtract.2
subtract.0
[64[[4050] operand 0- f64] 14050]
[64[14050] operand 0- 164 [14000]
ROOT
ROOT
Figure 11: runtime/XLA runtime 0.63 (58.7% faster). This is an example where the beam search finds more optimized case that trivially involves just performing additional instruction and computation fusions.
27

---

Under review as a conference paper at ICLR 2023
Optroad
Oplimtned
Competabo dww.cee.caaeeee cheaeheee
Ladittop
Fot take
Pauaaur FRIL
(A 2.18
Parseee
Pareeeer FA
Paneekrt
Fareeter
MISHLE
102 0oj16-11
openied cleane
Figure 12: BS runtime/XLA runtime 0.65 (53.8% faster). A second example where the beam search finds more optimized case that trivially involves just performing additional instruction and computation fusions.
28

---

Under review as a conference paper at ICLR 2023
F.2.2 NON-TRIVIAL EXAMPLES
In these examples, the beam search and F-MCTS optimized graphs outperform the XLA graphs despite fusing fewer instreciion/cooppuations.
Search
ONUN MANIAAL
THAZLE
DA CALURAA AAAAAAAA MAAIIII
Fand
Figure 13: BS run-time/XLA run-time 0.430 (132.5% faster).
PMCTS Optmieed
XLA Optimiedd
MEN LEREA
Nerta 11 SoUR
ALVOR L1713
Figure 14: f-CTTS runtime/XLA runtime = 0..349 (186.5% faster).
29

---

Under review as a conference paper at ICLR 2023
FMCTS Opömized
Opömized
Balle
Ball Bal-MN//p
MESDARI
LIB 1EA
SAU AULLIBUS LEA!
2.1)
AALLIIN LAMI
Figure 15: f-MCTS runtime/XLA runtime = 0.56 (78.6% faster).
MCTS Optimized
Optimized
Figure 16: f-MCTS runtime/XLA runtime 0.65 (53.8% faster).
30

---

Under review as a conference paper at ICLR 2023
f-CCTS Optimized
Por
and antoaat C.U
XLA Optimized
MUFCI
THE TE
pall pade Aupicc 11
Figure 17: f-MCTS runtime/XLA runtime 0.76 (31.6% faster). 31


Reviews:
[{"review_id": "T_h61cRLeCi", "reviewer": "Reviewer_UyBv", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper provides 1. a dataset of XLA-HLO graphs, 2. an environment to work with XLA in (HLOEnv), and 3. they explore the XLA optimization passes.\n\nThe dataset is gathered from 26 Jax repos, of varying sizes.\n\nHLOEnv provides 2 main additions beyond the basics: 1. A dry-mode for many of XLA's optimizations (i.e. the ability run the optimizations and collect information without modifying the graph), and 2. an alternate graph representation that provides an \"alternate graph\" representation of the optimization opportunity.\n\nThey also explore the optimizations done by XLA, and find two main optimizations that have a significant impact on performance - algebraic simplification and operator fusion. They explore various heuristics and find that it's possible to improve upon XLA's existing heuristics, even with fairly simple choices.\n\n", "summary_of_the_review": "See Conclusion.", "strength_and_weaknesses": "Overall, I found this paper to be an interesting read and contain lots of interesting analysis about the XLA compiler. That being said, I do feel that there is a substantial amount of overlap with prior work done on the XLA compiler, and so I am somewhat uncertain about the contribution this paper provides over prior work.\n\n### Prior work exploring learned heuristics over XLA graphs\n\nFor example, Paliwal et al. (2020) [1] provides a dataset of XLA graphs (although synthetic), Kaufman et al. (2021) [2]  explores a learned model for operator fusion, Yang et al. [3] explores equality saturation for graph rewrites, etc.\n\nThus, in my opinion, this statement in the abstract\n\n> We show that using simple heuristics for decision-making can achieve on-par or better performance than XLA. Using search algorithms further boosts performance.\n\nhas already been substantially explored in prior works. So, if the authors wish to make a substantial contribution on this front, I'd have liked to see more comparison of the search approaches used in this paper compared to prior work. Alternately, adapting prior approaches to the paper's infrastructure would also be a great contribution.\n\n### Dataset Quality\nI also have some worries about the dataset provided in the paper. For one, I'm not sure what the primary advantages are of this benchmark compared to TenSet (cited in the paper). Second, I'm somewhat worries that the repos chosen to be included in this benchmark are not ... representative benchmarks. In particular, many of them look to be somewhat toy projects that are not used in the community. In my spot-check of the repos, it seemed like perhaps 2/3 (or more) were repos that had less than 100 stars, and some of them (like https://github.com/rasutt/Toy-neural-network-in-jax) were merely toy repositories.\n\nI would not have an issue with the benchmarks chosen if they were merely to demonstrate performance improvements. However, as the dataset of this paper is meant to be a primary contribution, I think the inclusion of many low-traffic repositories is somewhat worrisome. (I would also have liked to see a split of where the graphs came from, since some of the repositories are mega-repositories like deepmind/deepmind-research). Benchmarks like MLPerf take great care to include \"popular\" and \"representative\" benchmarks, it worries me that, for example, I'm not sure I see a resnet implementation in any of the repositories provided.\n\nMore generally, I'm skeptical of approaches to benchmarks that provide a \"bag of graphs\". For a basic question, how many of these graphs are training benchmarks? How many of them are inference benchmarks? Although it's certainly *possible* to treat all of these graphs identically, the characteristics of inference graphs typically differ greatly from training graphs.\n\nThat being said, I do think the dataset improves meaningfully beyond existing XLA-HLO datasets, which range from proprietary to synthetic. However, the above concerns tamper my enthusiasm for this dataset.\n\n### Conclusion\n\nOverall, in my view, the main contributions listed by the paper are not strong enough for me to recommend acceptance. It has already been well established in the literature that XLA's heuristics can be improved through learned heuristics, and in my view the dataset contribution suffers from 1. lack of widely-used models, and 2. (related to 1) lack of more insight into the graphs and what they represent.\n\nThus, I think the primary contributions of this paper are: 1. A great investigation into what XLA heuristics matter, and 2. released code/data (as opposed to prior work which was mostly closed AFAIK). However, I think neither of these contributions are sufficiently technically novel, and I have several worries about the quality of the dataset.\n\n[1] https://arxiv.org/pdf/1905.02494.pdf\n\n[2] https://proceedings.mlsys.org/paper/2021/file/85d8ce590ad8981ca2c8286f79f59954-Paper.pdf\n\n[3] https://proceedings.mlsys.org/paper/2021/hash/65ded5353c5ee48d0b7d48c591b8f430-Abstract.html", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JCvptVk-XV", "reviewer": "Reviewer_ZQ7g", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper presents a system for extracting XLA graphs from the XLA compiler with annotations where the XLA compiler would have performed a rewrite that preserve the graph before and after the rewrite as alternatives. It looks into a system that applies different decisions that the default XLA rewrites by using beam search or monte-carlo tree search to choose between alternatives. It evaluates this system on a dataset of XLA programs gathered from GitHub repos.", "summary_of_the_review": "There are a number of factors here that suggest this paper needs more work before publication. It does not cite work like e-graphs which are very close to the representation used in the paper. Its results are also pretty weak, with only 1-2% overall speedups. Maybe this is due to averaging a lot of programs that are not effected by the optimizations, but the results are not presented in a way that makes it clear how the results would look if run on end-to-end programs, or whether improvements are due to a couple outliers. ", "strength_and_weaknesses": "Strengths\n-----------\nA fairly large dataset of XLA graphs as a benchmark for experimenting with optimizers is useful.\n\nThe description of how the system works was fairly easy to follow.\n\nWeaknesses\n--------------\n\n> Of the 94332 inst-10-20 sub-graphs generated, we filtered out 34000 graphs for which the correctness of the graph was not dependent on any graph rewrites performed by the Algebraic Simplification pass.\n\nThe approach of extracting rewrite rules from existing passing seems pretty fragile if ~2/3rds of the graphs produce incorrect results when it was changed. Similarly the fusion pass seems to filter out a similar proportion of the graphs because the method of searching would take too long to search. The paper needs to make this limitation much more clear, it is only mentioned in an appendix. Furthermore, it is not really a fair comparison to baseline XLA to exclude graphs it would compile fine. If averages are being reported, I would expect that the approach in the paper would have to include the numbers for the no-speedup cases where it cannot run, or would take too long to run.\n\nThe datasets are subsets of real world programs. Measuring performance improvements on a subset of a program is not necessarily representative of the whole program because time is not evenly spent in each operator.  I'd expect the best-improved subsets of programs to be much more improved than the program overall because maybe the subset chosen was precisely that part that would speed up a lot.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "7X1QDgLcPj0", "reviewer": "Reviewer_fXG4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Graph rewriting problem is an important problem in the field of deep learning compilers.\nThe paper notes that there is no good general testbed for deep learning compiler graph optimization research.\nThe paper develops and proposes a interface/framework for research.\nWith some preliminary results from the framework, the paper claims that it can provide more improvement compared to those of XLA.\nThe paper claims \"We intend for HloEnv and our dataset to be an open-source, community-driven effort that helps spur advances in DL compiler optimization research.\"", "summary_of_the_review": "The paper develops a framework for compiler optimization research, focusing on the graph rewrite problem.\nThe paper develops and proposes a interface/framework for research.\nWith some preliminary results from the framework, the paper claims that it can provide more improvement compared to those of XLA.\nThe paper claims \"We intend for HloEnv and our dataset to be an open-source, community-driven effort that helps spur advances in DL compiler optimization research.\"\n\nThe paper is very well written and makes a very clear contribution in the field of compiler optimization research.\nThis should have large impact when fully panned out. \nIn fact, considering the fact that the authors promised to open-source the efforts, the paper should be a big asset to the community. \n\nThe evaluation is rather preliminary.\nHowever, the main contribution of the paper comes from contributing a generalized framework for future development.\nEven such preliminary results suffice to show the potential benefits of the work.\n\nThe paper would be a nice addition to the program!", "strength_and_weaknesses": "+ Deep learning compiler can have major impact on the community (training + inference performance)\n+ Plans to open-source\n+ Clearly explains key ideas and results.\n\n- Only two examples which no clear evidence is provided about how good the baseline is.\n  (while reviewer agrees that XLA is good baseline, but is it for all HW and for all networks?)", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject