Tensor Program Optimization with Probabilistic Programs
Authors: Junru Shao
Xiyou Zhou
Siyuan Feng
Bohan Hou
Ruihang Lai
Hongyi Jin
Wuwei Lin
Masahiro Masuda
Cody Hao Yu
Tianqi Chen
Venue: NeurIPS 2022 Accept (2022)
Keywords: Tensor Program Optimization
Deep Learning Deployment
Machine Learning Compilation
Probabilistic Programming

Abstract: Automatic optimization for tensor programs becomes increasingly important as we deploy deep learning in various environments, and efficient optimization relies on a rich search space and effective search. Most existing efforts adopt a search space which lacks the ability to efficiently enable domain experts to grow the search space. This paper introduces MetaSchedule, a domain-specific probabilistic programming language abstraction to construct a rich search space of tensor programs. Our abstraction allows domain experts to analyze the program, and easily propose stochastic choices in a modular way to compose program transformation accordingly. We also build an end-to-end learning-driven framework to find an optimized program for a given search space. Experimental results show that MetaSchedule can cover the search space used in the state-of-the-art tensor program optimization frameworks in a modular way. Additionally, it empowers domain experts to conveniently grow the search space and modularly enhance the system, which brings 48% speedup on end-to-end deep learning workloads.

Tensor Program Optimization with Probabilistic Programs
Junru Shao OctoML jshao@octoml.as
Xiyou Zhou OctoML xiyou@octoml. ai
Siyuan Feng Shanghai Jiao Tong University hzfengsyāsjtu. edu. cn
Bohan Hou Carnegie Mellon University bohanhou@cs. cmu
Ruihang Lai
Hongyi Jin Carnegie Mellon University longyijecs cmu edu
Carnegie Mellon University ruihangl@cs cmu. edu
Wuwei Lin OctoML
Masahiro Masuda
Cody Hao Yu
OctoML. mmasuda@octoml.c ai
Amazon Web Services
wlin@octoml.at ai
hyuz@amazon.cmm com
Tianqi Chen
Carnegie Mellon University, OctoML tqchen@cmu edu tqchen@octoml.. ai
Abstract
Automatic optimization for tensor programs becomes increasingly important we deploy deep learning in various environments, and efficient optimization relies rich search space and effective search. Most existing efforts adopt search space which lacks the ability to efficiently enable domain experts to grow the search space. This paper introduces MetaSchedule, domain-specific probabilis- tic programming language abstraction to construct a rich search space of tensor programs. Our abstraction allows domain experts analyze the program, and easily propose stochastic choices in modular way to compose program transfor- mation accordingly. We also build an end-to-end learning- driven framework find an optimized program for given search space. Experimental results show that MetaSchedule can cover the search space used the state-of oofte--o-th-arr tensor program optimization frameworks in a modular way. Additionally. it empowers domain experts to conveniently grow the search space and modularly enhance the system, which brings 48% speedup on end-to-end deep learning workloads.
Introduction
Deep learning has become pervasive in daily life. From video understanding 228 natural language
understanding [15], and recommendation system (2 autonomous driving 21. different deep
learning models are deployed on different hardware platforms and devices. Deep learning frameworks
usually rely on manually optimized libraries [3] p accelerate deployment. Engineers need to choose from many tensor programs that are logically equivalent but differ significantly in performance due to memory access, threading, and the use of specialized hardware primitives. The enginecring effort required for tensor program optimization has become significant bottleneck for machine learning deployment with the growing number models and hardware backends.
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Key elements in automatic tensor program optimization
of Fiesst frggrrmin
Stag
Optinal Progun Pogane
Dxapples
megelte94): (agge12824))
Far fas Ie artefit, fit for Au 2a grickly for by recet(heea))
for for grittee, atr for to is prietsy soy for in resgeddos "an los in graits, di for raagetials ([oc)]
for in (orge(18240)
far lo gri4ts clood
Figure 1: Automatic tensor program optimization contains two key elements: the search space S(cv) and the search algorithm that finds the optimal tensor program e*. The search space usually incorporates choices over loop transformation, vectorization, threading patterns, and hardware acceleration.
Automatic program optimization 10 43
is a recent sequence of efforts that aims to use machine vital elements of automatic tensor program optimizations.
learning solve this problem. There First, search space is defined to provide a possible set of equivalent tensor programs. Then existing systems use learning-based search algorithms to find an optimized tensor program in the search space with feedback from the deployment environment. Most of the current approaches [T0] 43 II 255 use pre-defined search spaces that effectively encode the domain knowledge of the authors once and focus on developing efficient search algorithms.
While efficient search is essential, the search space itself fundamentally limits the best possible performance search algorithms can get. To construct a good search space, domain experts have to make numerous choices over loop transformation, vectorization, threading patterns, and hardware acceleration. Additionally, the best search space itself evolves as new tensor program optimization techniques and hardware primitives 30 grow, As result, there strong need enable easy customization and construction of the search space at scale by taking inputs from system engineers and domain experts. Unfortunately, any change search space construction currently requires surgical modifications to the automatic program optimization frameworks.
This research asks the following question:
decouple the search space construction from search and provide an adequate abstraction for domain experts and the learning system to collaborate on search space construction? We give affirmative answer the question with key observations. First, we can parameterize an optimization search space by the initial program followed by sequence of transformations
the program. Next, using this parameterization, domain experts then provide probabilistic choices that represent possible transformations after examining the program state. These two observations lead a simple yet powerful abstraction for search space construction through a domain-specific probabilistic language. Finally, our framework composes multiple possible probabilistic transformations form rich search space. We make the following contributions:
We introduce a simple yet powerful probabilistic language abstraction for tensor program search space construction.
We build learning-driven framework to find optimized tensor programs specified by the search space constructed using our abstraction.
We build an end-to-end system that take prior knowledge from domain experts to construct optimization search space to optimize deep learning deployment on multiple platforms.
Our end-to-end system can easily expand search space that matches previous approaches without surgical changes and achieve comparable performance multiple hardware backends. Experimental results show that our abstraction is expressive enough to cover the optimization space of diverse set of tensor programs, delivering competitive performance of popular deep learning models, and convenient to incorporate hardware-specific knowledge into the search space to outperform the state-of-the-art frameworks.

---

Parameterization
Equivalent Programs Induced by Parameterized Transformation
Initial tensor program: e.
Equivalent intermediate program: R2 1y rangel32):
Equivalent optimized program: e* paraltel, for to in rangel32); for da in renget01: i = la + 32 + K
for rance ( [[i] (e]U(A(1))
for by in range(8):
For reriee(d:: (1::)
Transformation
Parametorization
1. 32 1,
A[s 4]
Splitt
B[4] ReLu)A(())
RoLUCA[1 1 41)
10 Vectorize
parameterized by: 6,,
parameterized by: %, +DE
Figure 2: Parameterizing programs with the initial program and sequence of transformations. Tensor program é1 is parameterized by the initial program co. plus Step D, which is further parameterized by loop and resulting loop extents 32, 8, 4, respectively.
2 Background and Problem Overview
Figure shows a typical workflow for tensor program optimization. For given program én, a typical tensor program optimization framework will generate candidates from pre-defined search space S(eo) containing mannanically-equivlllee programs. Then the framework finds optimized tensor program e* E S(ea) with the minimum latency on the target hardware.
typical search space S(en) contains choices over threading, loop ordering, memory access, and hardware primitives. Defining the search space S(eo) for wide range of tensor programs brings several challenges. First, S(e)) is highly dependent on co- For example, S(ea) of a compute-intent intensive program (e.g., Dense) needs consider many more possible configurations than communication- intensive program such as ReLU. The space also differs significantly in different hardware domains. For example, (eo) on CPU involves multi-core parallelism and vectorization, while S(eo) GPU involves thread binding and tensorization. Finally, as the hardware and model settings change, we need bring fresh domain experts update the S(ev) leverage the latest improvements.
This paper aims to provide a programmable abstraction to construct S(-) in a composable and modular way. Our key goals are listed as follows: Expressiveness. We need be able to build rich search space that covers the optimization programs that domain experts will write. Modularity, Tensor program optimization likely will involve inputs from multiple domain experts different periods. Therefore, we need to be able to combine prior knowledge in a composable and modular way. Designed for learning. We need to build a generic learning-based framework to enable diverse variations of the cost model and search for search space specified in our abstraction. We will address the above goals in the following two sections.
3 Composable Search Space Parameterization
This section presents MetaSchedule, a probabilistic approach to search space parameterization.
3.1 Stochastic Search Space Construction
MetaSchedule constructs search space S(.) with stochastic program transformations the primitive building blocks. Traditional program optimization can usually be represented by a sequence of transformations т, where step the program 4-1 is transformed into semantically-equival program e,, which finally leads to the optimized program en. MetaSchedule generalizes this idea by allowing further parameterization of each transformation step T.
Taking Figure 2 as an example: 6n is the initial program for the program B = ReLU( (A) MetaSchedule, transformation t Split is parameterized by a loop i and sequence of integers indicating the loop extents after splitting; Similarly, transforms t/ = Parallelize and f3 = Vectorize are parameterized by loops respectively. As result, an optimized program e, is obtained by applying a sequence of parameterized transformations r to the initial program eg. Accordingly, the search space S(ep) is composed of en and all possible sequences of parameterized transformations.
In
In practice, we ingest models from PyToorch//TnssoFHo///AXX See Appendix A.6 for details.

---

Probabiiiticc Program
Stochastic Tra Transformation
Bens: Ie reprtsiii
Apply !
kall
Frobabillutte-- Proggnntt- * Leep tatiing for benis
Hu ii Saepl E-llell parreadt
porr in rergettele
supuss
Ai Sec.ueeelee Jaledke pet las de Splitta Ihus
for 27 in rergettt2) for 1" renge22991s nfoel
for AG 3 le praett, Edtt Hied
St Ji Splaed,, Ris ayl
a- Iuadin
pridkt, fut I- griets, Clud
1.0 g000 24ti0ee Poos!
da Ja in gradto, Bit for far in grodie, fit
[ heali
for corpette))
Clo.el
2f rarge15221: for rengel2301: ofo.e
2', pristeu, Blovs
Figure 3: The MetaSchedule probabilistic language. The language contains two key elements: (l) sampling of random variables; (2) program transformation based on random variables. An example execution instance: Step ①: Draw tile sizes of and then organize the loops into two-level tiling structure. Step @: Decide where to fuse the ReLU operator.
On the other hand, it could be less practical for practitioners determine the best combination of the parameter values in transformations. For instance, in Figure 2] it is usually efficient to use 512-bit vectorization over the inner loop when AVX-512 vector instructions available on Intel CPUs, other vectorization lengths may lead to better performance, otherwise. Therefore, deep knowledge of the target hardware is mandatory to enumerate plausible parameter combinations to control the search space size while covering the optimal program.
To let practitioners efficiently define parameterized transformations without worrying about candidate values, MetaSchedule introduces random variables drawn from analysis, sampling. Parameterized by random variables, transformation naturally becomes stochastic, and the underlying probabilistic space reflects the space of possible transformations.
when creating Split transformations to tile loop i and j in the Dense operator, the tile sizes are drawn by random variables 6--3 defined from Sample-Tile. In this way, the Split transformation becomes stochastic. Similarly, we use Homple-Compute Locattion to enumerate valid loops in Dense after splitting for ReLU to fuse its computation. summary, lines of MetaSchedule program covers a family of possible optimized tensor programs with stochastic transformations in its search space S(eo), where 6o is Dense-ReLU.
As illustrated in Figure 3,
Notably, unlike orthogonal grid space in hyperparameter tuning, MetaSchedule captures long- term structural and arithmetic dependency between random variables and the tensor program e being transformed. As demonstrated on Step @ in Figure 3] sampling distribution from -ompplte Location depends on the latest tensor program C5, whose structure depends on all previous random variables.
3.2 Modular Search Space Composition
Although the search space constructed by stochastic transformations proposed in the previous subsection is efficient and is capable of covering the optimal tensor program, it is hard for other developers to learn how the search space is constructed by reading a long sequence of transformations. It makes transformations designed for a workload hard to be reused by other workloads. Meanwhile, observe that it is straightforward to group sub-sequence of transformations for particular fine-graine optimization. For example, some transformations implement multi-level tiling for better memory locality in compute-intensive operators like Conv2d and Dense; some other transformations are used to fold/inline elementwise operations such as activation functions into their predecessors or successors for better memory bandwidth efficiency.
To improve the usability and make MetaSchedule more practical, we introduce transformation module. Just like the convolutional module with Conv2D, BiasAdd and ReLU in ResNet, transformation module in MetaSchedule is defined as either atomic stochastic transformation, or composition of
A single X86 instruction that performs the computation 512 bits vector in CPU cycle.

---

Transformation Module
Example: Execution of the Transformation Module
car matcs www---ilnnggleg reat Luttlapl::
riles verp wqll Disrd (aserts] der tit, talallosstinn loep Leep, tit Listionit: 64) suplan stettee, prttslestssiasaak) tiled_lops splsctlos,, 500
Tooefomed
far ia rangetsazl: for in negelzze): for a rangatitic Cool
rargetill grasti,
zezeite Jey, rited twest
yredirrction
for leegurests
"F RCGARELL Seapoar BiCCeee (0, (2"1) <tt reduttie, teaocit
Cool
Figure 4: Transformation modules. A transformation module consists of tensor program analysis, sampling, and stochastic transformations. The figure uses MelliiLLeve--Te where analysis is done interactively to identify spatial (data parallel) and reduction loops, and then apply Split with the tiling factors drawn from Sample- -ill.. A final Reorder organizes the loops into proper tiles.
as an example.
Transformation Module
Eechastic-Compore
Oenaa knawledge InnorCano
Multi-Lvve
procedure ComPosE(modules. program) while locations are not exhausted do
Skc
Use Tensor-Core
location t next-location ( modules)
Anto-inliee
Seath
m ~ Sample
program +- m(location, program)
Cross Thread Reduction
end while end procedure
Praperse
Figure 5: Left: An example algorithm to compose transformation modules. A sequence of trans- formation modules is composed together into single transformation module. Right: Hierarchical composition of transformation modules gives generic search space. In this example, a hardware- specific module Use- Tene-r-no---ooe -orre composed together with other generic modules into a module that generates search space for any tensor program.
program analysis, sampling as well as smaller transformations. Each transformation module can have meaningful name that it can be easily adopted by many workloads to hierarchically construct search space.
Figure 4 shows hierarchical composition of transformation modules. Specifically, Multi-Level-Tiiimg interleaves program analysis on loops and the stochastic tiling of the loop structure and organizes the original tensor program into 5-level tiling structure. Notably, the transformation module generic to tensor programs and thus could be applied to a variety of operators, including convid, conv3d, matmul, etc.
Figure 5 depicts example of composing search space with transformation modules. In this simple example, we select set of transformation modules, which are implemented in advance by practitioners with prior domain knowledge, and apply them to every available location in the tensor program to form search space. Consequently, the formed search space covers common optimizations on diverse hardware.
3.3 Relation to Existing Tensor Program Optimization Methods
In this subsection, we discuss prior approaches for automatic tensor program optimization and illustrate that many them can be covered by the MetaSchedule framework.
Domain specific languages for program transformations used by prior frameworks 32 allow developers to easily optimize program manually. When there no random variable sampling in the program, MetaSchedule reduces to a DSL for deterministic program transformations and achieves the same functionality.
Template-guided auto-tuning 100
fully relies on developers to define a search space. In MetaSchedule, it means all random variables in search space are defined ahead of the transforma- tions, so there no interaction between program analysis and follow-up random sampling choices conditioned on the program state.
25
27
5

---

The Probsbildte Program
def WultinlaweleTiiingg pest: List [tmpll::
Bieet, Nstt, 11sti1]
Srotecr Leop, 11le_ Listlattle (6) Sinpte-Tieellag parts (entseeen()5))
Inttill
thr & Senpte-matats, parten2) On Bi Sempla-mataty partungi tplasia, the dil splant), Boorderlle. Ju dur 307
"*een 5, tile siprtale__dd, tilel lospal:
for 1 in toop_rest resti
Torptii: torpls, 233
traced instuciion
redrrdacti loopiat: tale, torpis, (3]1)
Resrder (ies))
Figure 6: Execution tracing in MetaSchedule. The probabilistic language on the left defines the entire search space S(eo). Tracing the program exccution across different runs leads to set of linearized probabilistic programs on the right. Only sampling and transformation instructions are traced, while all other constructs and control flow in the host language is ignored.
19 requires developers to implement workload agnostic transforma- tion rules. MetaSchedule achieves the same programmability and functionality through specific probabilistic transformation modules that correspond to the search space generation rules.
Auto-scheduling +
Notably, all approaches mentioned above have important use-cases in tensor program optimizations, depending on how much domain knowledge we want to incorporate for particular scenario. By decoupling the search space construction from the search, we effectively build single framework for all the use cases and enable further customization without surgical changes to the system.
4 Learning-driven Search
The last section provides a modular abstraction for search space. We still need to do an effective search to find an optimized program within the search space. This section provides generic learning-driven framework to find an optimized program.
Objective formalization. For given probabilistic program cu. let us use T to denote the transfor- mations performed on cg. can be sampled from a prior distribution specified by the probabilistic program. We define y(eo, r) to be the tensor program after applying transformation [(e) be the latency of the particular program e on the hardware environment. We define posterior probability of an optimized program as:
ep- Let
P(r eo) x -((((at)))) P(T)-
Intuitively, we want to assign higher probability to the programs that perform well. Our final goal to find 7* argmax, P(r co) that maximizes the posterior through maximum a posteriori estimation (MAP) estimation.
Execution tracing. To enable domain experts to express their knowledge via transformations modules productively, we embed MetaSchedule in Python. We introduced execution tracing to reduce the cost of repetitive re-execution of the Python program. Figuree demonstrates example tracing process. During program execution, the system records all samplings and transformations while ignoring control flow and other constructs of the host language. The resulting trace is sequence of MetaSchedule primitives with only sampling and transformation instructions, which could be re-executed as a normal MetaSchedule program. We can then continue explore different sampling choices for given collection of initial traces. Conceptually, this is equivalent to dividing up our support set and then sampling the program condition on the execution sequence of the program.
End-to-end search. Figure 7 shows the overall workflow of our learning-driven framework. The search algorithm first samples the MetaSchedule program obtain collection of traces. Then continues to explore the space condition on the traces. Notably, there is a significantly higher cost measuring fee directly on the hardware, so we also incorporated proxy cost model f(e), which is updated throughout the process, similar previous efforts on tensor program optimization [T 43 At each iteration, we adopt an evolutionary search algorithm that proposes a new variant of the trace by mutating the random variables, then accept or reject the proposal based on the cost model. While evolutionary search could be viewed as parallel chain MCMC, we also made our system modular enough to incorporate other ways to select the probabilistic choices, such as those through Bayesian optimization and reinforcement learning.

---

Learning driven
Transformation Module
liing
Intill tonsor
reject
progrom
sample
program
RandoosAnnnol
Serect
Hardhware ((e)
Figure 7: Learning-drin driven search. Based on the traces of MetaSchedule execution, candidate tensor programs are proposed by mutating sampling decisions in traces, among which the invalid ones are rejected by the validator. In every iteration, proposed candidates are accepted or rejected via annealed Netropolis-Hastinge with prediction from a learned cost model f. while the cost model incrementally updated according f, the measured latency of tensor programs real hardware.
Cost model. Our approach allows extensive cost models, enabling us to supply those pre-trained from existing datasets 44 We pick a prre-boosting-base cost model in f(-) by default and leverage a common of features that are used in previous works 43
Trace validation. Importantly, invalid traces may show up as we propose updates. Such a scenario happen when some of the random variable choices go beyond the physical hardware limit variable that induces changes to the execution sequence. Instead of enforcing a conservative proposal, introduce validator that validates the correctness of the trace. The trace validation allows to move around the space more freely while still ensuring the correctness of the sample outcome to be on the right support set.
5 Related Work
Tensor Program Transformations are proposed by many prior works, such as Halide 32 Tiramisu [5] and TACO 23 E Note that all the previous transformation languages are deterministic and cannot be directly used for search space construction, meaning that they have introduce separate programming model to express a search space. This paper makes a simple but powerful generalization to domain-specifie probabilistic language. The resulting abstraction enables unified approach to deterministic transformation and search space construction.
TVM M
Black-box optimization has been adopted in high-performance computing libraries 16 3
Recent
advances automatic tensor program optimization brought a rich set techniques to accelerate and learning-basse search [2]
search through better cost modeling 110
44 which could be incorporated into MetaSchedule search. Different variations of pre-defined search spaces have also been proposed that couple with the automatic tensor program optimization frameworks [CC
4 3E
25
19
43 M. Polyhedral model 400 39 is one useful way to construct rich pre-defined search space. This paper focuses modular search space construction and provides orthogonal contributions to these prior works.
Probabilistic programming language is powerful abstraction for incorporating domain knowl- edge and probabilistic inference. There are many general-purpose probabílistic languages, such as Church [77. Stan [8) Pyro 77. NumPyro BT,, PyMC3 (55) and Edward [3)) This paper proposes
speeiffc probabilistic language for tensor program optimization with specializations such
as tensor program analysis that would otherwise be opaque to the previous systems. Our learning- driven search can be viewed as an application previous works
All 46 that use tracing to divide programs into subprograms with fixed support. We focus the MAP inference problem where the posterior depends on an unknown cost function. We solve the problem through a learned cost-model-driver evolutionary search over traces and with validation.
[1) 33
Automatic neural program synthesis
18 has seen large progress recently. Alphacode 26 builds a system that can output creative and sound solutions to problems that require deeper-level human understanding. These approaches generate abstract syntax trees (ASTs) that can incorrect and use input-output pairs to filter out those erroring programs. Our compiler approach requires us to ensure the correctness of transformations. However, some ideas like validation after creation might be reusable.
7

---

PyTorch
TVM
MetaSchedu e
GPU
Normalized Porformance
CiD
C2D
DEP
GMM
GRP
T2D
CBR
TBG
NRM
SFM
GPU
C1D
C2D
C3D
DEP
DIL
GMM
GRP
T2D
CBR
TBG
NRM
SFM
Figure 8: Operator- and subgraph-level performance. MetaSchedule always achieves similar or better performance compared with TVM (Ansor), and exceeds Py PyTorch significantly in most workloads whose operators are carefully hand-optimized.
PyTorch
TVM
MetaSchedule
1.0
Nomallzzd Performance
0.5
BERT-Base
ResNet50
MoblleNetV2
BERT-Base
ResNet50
MobileNetVZ
Figure 9: Optimizing end-to-end deep learning models. performance compared TVM (Ansor) all the models.
MetaSchedule reaches close or better
Experiments
6.1 Expressiveness to Cover Common Optimization Techniques
This section aims to answer the following question: Is MetaSchedule expressive enough to capture the search space of the latt-of-the-art optimization techniques? To answer this question, evaluate our work on a diverse set of operators and subgraphs extracted from popular deep learning models, including variants of convolution, dense, and normalization.
As baselines, PyTorch (v1.11.0) results are provided to compare performance with vendor libraries; TVM (commit: 8d4f4dd7311 which incorporates AutoTVM [10] and Ansor 43), is used as the state-of-thed tensor program optimization system, and we pick the best among the two in each respective setups. Full operators and hardware configurations are documented in Appendix A.2
Figure 8 shows that, in all cases on CPU and GPU, MetaSchedule delivers performance compa- rable with or even better than TVM, from which we could infer that MetaSchedule could express optimization techniques comparable to TVM on diverse workloads. Additionally, in most of the cases, MetaSchedule outperforms PyTorch by significant margin except for SFM, which is highly optimized manually in PyTorch.
6.2 Optimizing End-to-End Deep Learning Models
Operator performance does not always translate directly to full model optimization. Therefore, this section is dedicated to answering the following question: Can MetaSchedule deliver competitive performance with tfate- works for end-to-end models?
Therefore, series of experiments are conducted to compare MetaSchedule and TVM, including BERT- Base 14] ResNet-50 20).. and MobileNet-v2 (6 on both CPU and GPU. As shown in Figure MetaSchedule performance is parity with TVM, while surpassing PyTorch in all cases, which indicates that MetaSchedule framework delivers end-to-end performance. Additionally, tuning time is provided in Appendix A.5

---

*Mut--uvve Ting
Pylorrh
TVM
MetaSchedule
Normelized
100
1000
Normall and Performance
BERT Large
(a) Performance with different search spaces.
(b) BERT- Large Performance.
Figure 10: Left: Search space composition conducted on representative subgraph of BERT called fused-dense. Gradually composing more transformation modules, the space covers more optimized programs. Right:
Introduction of hardware specifcc Tese-Tenor-oee module, composed with existing search space, brought 48% speedup over TVM (AutoTVM).
6.3 Search Space Composition and Hardware-Speciffc Modules
Besides performance parity with existing work, in this section, we demonstrate the extra value of modular search space composition by answering the following question: How convenient is it to compose transformation modules, and how does translate to performance?
We design an ablation study for transformation modules composition. As indicated in Figure I0a by progressively enriching the search space, the performance of optimized tensor programs consis- tently increases. Composed with hardware-specific module Use-Tensor-Core delivers significantly better performance compared with generic search space. The performance gain, brought by search space composition with customized rules, does translate to end-to-end model performance, as shown Figure 10b Specifically, on BERT-large workloads, MetaSchedule with Use-Tensor-Core delivers 48% speedup over TVM.
MetaSchedule
Notably, it took graduate student only days craft the 82-line Use-Tensor-Core module in Python (see supplementary materials), which provides strong evidence of the convenience of customization and composition, More details Appendix A.4
Conclusion
This paper presents MetaSchedule, a programming model to describe search space construction tensor program optimization. Our method abstracts search space as a probabilistic language and enables flexible incorporation of domain knowledge by allowing practitioners to implement customized probabilistic programs. A cearning-driven search algorithm developed on top of the probabilistic language abstraction, which delivers competitive performance with state-of-ther of-the-art frameworks. the future, we will explore and modularize declarative API for various hardware environments. Therefore, collaboration between the machine learning deployment engineers and intelligent machine learning algorithms for tensor programs.
will open-source framework and hope it could enable broader
Acknowledgement
This work is supported in part by a gift from Oppo. We would like to thank Josh Fromm, Denise Kutnick and Sunghyun Park from OctoML for helpful discussion and support of the work.

---

References
[1] Andrew Adams, Karima Ma, Luke Anderson, Riyadh Baghdadi, Tzu-Mao Li, Michaël Gharbi, Benoit Steiner, Steven Johnson. Kayvon Fatahalian, Frédo Durand, et al. Learning to optimize halide with tree search and random programs. ACM Transactions on Graphics (TOG), 38(4):1-
12, 2019.
[2] Byung Hoon Ahn, Prannoy Pilligundla, Amir Yazdanbakhsh, and Hadi Esmacilzadeh. Chameleon: Adaptive code optimization for expedited deep neural network compilation. arXiv preprint arXiv:2001. 008743, 2020.
[3] Collaboration ATLAS, Torsten Åkesson, Paula Eerola, Vincent Hedberg, Göran Jarlskog, Björn Lundberg, Ulf Mjörmmark, Oxana Smirnova, Sverker Almehed, et al. Atlas computing: technical design report. 2005,
[4] Riyadh Baghdadi, Massinissa Merouani, Mohamed-Hichhe Leghettas, Kamel Abdous, Taha Arbaoui, Karima Benatchba, et al. A deep learning based cost model for automatic code optimization. Proceedings of Machine Learning and Systems, 3:181-195 2021.
[5] Riyadh Baghdadi, Jessica Ray, Malek Ben Romdhane, Emanuele Del Sozzo, Abdurrahman Akkas, Yunming Zhang, Patricia Suriana, Shoaib Kamil, and Saman Amarasinghe. Tiramisu: A polyhedral compiler for expressing fast and portable code. In 2019 IEEE/ACM International Symposium on Code Generation and Optimization (CGO), pages 193-205. IEEF, 2019.
[6] Riyadh Baghdadi, Jessica Ray, Malek Ben Romdhane, Emanuele Del Sozzo, Abdurrahman Akkas, Yunming Zhang, Patricia Suriana, Shoaib Kamil, and Saman Amarasinghe. Tiramisu: A polyhedral compiler for expressing fast and portable code. In Proceedings of the 2019 IEEE/ACM International Symposium Code Generation and Optimization, CGO 2019, page 193-205. IEEE Press, 2019.
[7] Eli Bingham, Jonathan P Chen, Martin Jankowiak, Fritz Obermeyer, Neeraj Pradhan, Theofanis Karaletsos, Rohit Singh, Paul Szerlip, Paul Horsfall, and Noah D Goodman. Pyro: Deep universal probabilistic programming. The Journal of Machine Learning Research, 20(1):973
978, 2019.
[8] Bob Carpenter, Andrew Gelman, Matthew D Hoffman, Daniel Lee, Ben Goodrich, Michael Betancourt, Marcus Brubaker, Jiqiang Guo, Peter Li, and Allen Riddell. Stan: A probabilistic programming language. Journal of statistical software, 76(1), 2017.
[9] Tianqi Chen, Thierry Moreau, Ziheng Jiang, Lianmin Zheng, Eddie Yan, Haichen Shen, Meghan Cowan, Leyuan Wang, Yuwei Hu, Luis Ceze, et al. [TVM): An automated (End-to-End optimizing compiler for deep learning. In 13th USENIX Symposium on Operating Systems Design and Implementation (OSDI 18), pages 578-594, 2018.
[10] Tianqi Chen, Lianmin Zheng, Eddie Yan, Ziheng Jiang, Thierry Morcau, Luis Ceze, Carlos Guestrin, and Arvind Krishnamurthy. Learning to optimize tensor programs. Advances in Neural Information Processing Systems, 31, 2018.
[11] Xinyun Chen, Chang Liu, and Dawn Song. Execution- neural program synthesis. In International Conference Learning Representations, 2019.
[12] Xinyun Chen, Dawn Song, and Yuandong Tian. Latent execution for neural program synthesis beyond domain-specifie languages. Advances in Neural Information Processing Systems, 34, 2021.
[13] Sharan Chetlur, Cliff Woolley, Philippe Vandermersch, Jonathan Cohen, John Tran, Bryan Catanzaro, and Evan Shelhamer. cudnn: Efficient primitives for deep learning, arXiv preprint 1rXiv:1410.0759, 2014.
[14] Jacob Devlin, Ming-Wai Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 1rX0v:1810..0 04005, 2018.
[15] Jacob Devlin, Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre- Pre-training of deep bidirectional transformers for language understanding. Proceedings of the 2019 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
10

---

[16] Matteo Frigo and Steven G Johnson. Fftw: An adaptive software architecture for the fft. In Proceedings of the 1998 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP'98 98 (Cat. No. 98CH36181), volume 3, pages 1381-1384. IEEE, 1998.
[17] Noah Goodman, Vikash Mansinghka, Daniel M Roy, Keith Bonawitz, and Joshua B Tenenbaum. Church: language for generative models. arXiv preprini arXiv: 1206.225, 2012.
[18] Kavi Gupta, Peter Ebert Christensen, Xinyun Chen, and Dawn Song. Synthesize, execute and debug: Learning to repair for neural program synthesis. Advances in Neural Information Processing Systems, 33: 13:11766-16655 2020.
[19] Ameer Haj-Ali, Hasan Genc, Qijing Huang, William Moses, John Wawrzynek, Krste Asanović, and Ion Stoica. Protuner: tuning programs with monte carlo tree search. arXiv preprini arXiv:2005. 13685, 2020.
[20] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and paltern recognition, pages 770-778, 2016.
[21] Yu Huang and Yue Chen. Autonomous driving with deep learning: A survey of state-of-art technologies. arXiv preprint arXiv:2006.0609), 2020.
[22] Intel. Intel@ math kernel library for deep learning networks, 2017.
[23] Fredrik Kjolstad, Shoaib Kamil, Stephen Chou, David Lugato, and Saman Amarasinghe. The tensor algebra compiler. Proceedings of the ACM on Programming Languages, 1(00PSLA))1- 29, 2017.
[24] Andrew Lavin and Scott Gray. Fast algorithms for convolutional neural networks. In 2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016, pages 4013-4021, 2016.
[25] Menghao Li, Minjia Zhang, Chi Wang, and Mingqin Li. Adatune: Adaptive tensor program compilation made efficient. Advances in Neural Information Processing Systems, 33:14807-
14819, 2020.
[26] Yujia Li, David Choi, Junyoung Chung, Nate Kushman, Julian Schrittwieser, Rémi Leblond, Tom Eccles, James Keeling, Felix Gimeno, Agustin Dal Lago, et al. Competition code generation with alphacode. arXiv preprint arXiv:2203.07814 2022.
[27] Yizhi Liu. Yao Wang, Ruofei Yu, Mu Li, Vin Sharma, and Yida Wang. Optimizing (CNN) model inference on {CPUs}. In 2019 USENIX Annual Technical Conference (USENIX ATC 19). pages 1025-1040, 2019.
[28] Ze Liu, Jia Ning, Yue Cao, Yixuan Wei, Zheng Zhang, Stephen Lin, and Han Hu. Video swin transformer. arXiv preprint arXiv:2106. 1.1000...2.. 13230, 2021.
[29] Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sun- daraman, Jongsoo Park, Xiaodong Wang. Udit Gupta, Carole-Jean Wu, Alisson G. Azzolini, Dmytro Dzhulgakov, Andrey Mallevich, Ilia Cherniavskii, Yinghai Lu, Raghuraman Krish- namoorthi, Ansha Yu, Volodymyr Kondratenko, Stephanie Pereira, Xianjie Chen, Wenlin Chen, Vijay Rao, Bill Jia, Liang Xiong, and Misha Smelyanskiy. Deep learning recommendation model for personalization and recommendation systems. CoRR, abbs/1906.0009) 2019.
[30] Nvidia. Nvidia tensor cores, 2017.
[31] Du Phan, Neeraj Pradhan, and Martin Jankowiak, Composable effects for flexible and acceler- ated probabilistic programming in numpyro. arXiv preprint arXiv:/9/2.1 1554, 2019.
[32] Jonathan Regan-Kelley, Connelly Barnes, Andrew Adams, Sylvain Paris, Frédo Durand, and Saman Amarasinghe. Halide: a language and compiler for optimizing parallelism, locality, and recomputation in image processing pipelines. Acm Sigplan Notices, 48(6):519-530, 2013.
[33] Daniel Ritchie, Andreas Stuhlmüller, and Noah Goodman. C3: Lightweight incremental- ized meme for probabilistic programs using continuations and callsite caching. In Artificial Intelligence and Statistics, pages 28-37. PMLR, 2016.
[34] Jaehun Ryu and Hyojin Sung. Metatune: Meta-learning based cost model for fast and efficient auto-tuning frameworks. arXiv preprint arXiv:2102.041 041999 2021.
[35] John Salvatier, Thomas Wiecki, and Christopher Fonnesbeck. Probabilistic programming in python using pymc3. PeerJ Computer Science, 2:e55, 2016.

---

[36] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang Chich Chen. Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4510-4520, 2018.
[37] Ryan Senanayake, Changwan Hong, Ziheng Wang, Amalee Wilson, Stephen Chou, Shoaib Kamil, Saman Amarasinghe, and Fredrik Kjolstad. A sparse iteration space transformation framework for sparse tensor algebra. Proc. ACM Program. Lang,, 4(0OPSLA), November
2020.
[38] Dustin Tran. Matthew D. Hoffman, Dave Moore, Christopher Suter, Srinivas Vasudevan, Alexey Radul, Matthew Johnson, and Rif A. Saurous. Simple, distributed, and accelerated probabilistic programming. In Neural Information Processing Systems, 2018.
[39] Nicolas Vasilache, Cédric Bastoul, and Albert Cohen. Polyhedral code generation in the real world. In International Conference on Compiler Construction, pages 185-201. Springer, 2006.
[40] Nicolas Vasilache, Oleksandr Zinenko. Theodoros Theodoridis, Priya Goyal, Zachary De Vito, William S Moses, Sven Verdoolaege, Andrew Adams, and Albert Cohen. Tensor comprehen- sions: Framework- aggnostic high-performance machine learning abstractions. arXiv preprint arXiv:1802.04430, 2018.
[41] David Wingate, Andreas Stuhlmüller, and Noah Goodman. Lightweight implementations of probabilistic programming languages via transformational compilation. In Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics, pages 770-778.
JMLR Workshop and Conference Proceedings, 2011.
[42] Lingfeng Yang. Patrick Hanrahan, and Noah Goodman. Generating efficient meme kernels from probabilistic programs. In Artificial Intelligence and Statistics, pages 1068-1077 PMLR, 2014.
[43] Lianmin Zheng, Chengfan Jia, Minmin Sun, Zhao Wu, Cody Hao Yu, Ameer Haj-Ali, Yida Wang, Jun Yang. Danyang Zhuo, Koushik Sen, et al. Ansor: Generating (Higg-Performance) tensor programs for deep learning. In 14th USENIX Symposium on Operating Systems Design and Implementation (OSDI 20), pages 863-879, 2020.
[44] Lianmin Zheng, Ruochen Liu, Junru Shao, Tianqi Chen. Joseph E Gonzalez, lon Stoica, and Ameer Haj Ali. Tenset: A large-scale program performance dataset for learned tensor compilers. In Thirty-fiith Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021.
[45] Size Zheng. Yun Liang, Shuo Wang. Renze Chen, and Kaiwen Sheng. Flextensor: An automatic schedule exploration and optimization framework for tensor computation on heterogeneous system. In Proceedings of the Twenty-Fifhh International Conference on Architectural Support for Programming Languages and Operating Systems, pages 859-873, 2020.
[46] Yuan Zhou, Hongseok Yang, Yee whye Tch, and Tom Rainforth. Divide, conquer, and combine: new inference strategy for probabilistic programs with stochastic support. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 1 11534 11545. PMLR, 13-18 Jul 2020.
12

---

Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on how to answer these questions. For each question, change the default [TODO] to [Yes], [No], , or [N/A] - You are strongly encouraged to include a justification your answer, either by referencing the appropriate section of your paper or providing brief inline description. For example:
Did you include the license to the code and datasets? [Yes] See Section gen_ jon_inst.
Did you include the license to the code and datasets? [No] The code and the data are proprietary.
Did you include the license to the code and datasets? [N/A]
Please do not modify the questions and only use the provided macros for your answers. Note that the Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/amswers below.
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See section
(c) Did you discuss any potential negative societal impacts of your work? [No] To the best of our knowledge, there is no potential negative societal impact of our work, given the only focus is to accelerate existing machine learning models.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A] Our work does not include theoretical results.
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper- imental results (either in the supplemental material or as URL)? [No] We will not include the URL for codebase for anonymity, and will release the link after the review process.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] Operator configurations and hyperparameters for evolutionary search are shown the Appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)? [No] The variance of running time of tensor programs is negligible across several runs.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We use the codes of TVM and PyTorch which are both cited.
(b) Did you mention the license of the assets? [No] TVM under Apache 2.0 license.
PyTorch is under Modified BSD license.
(c) Did you include any new assets either in the supplemental material or as a URL? [No] We did not use any new assets.
(d) Did you discuss whether and how consent using/curating? [N/A] We did not use data obtained from other people.
obtained from people whose data you're
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A] We did not use data contains personally identifiable information or offensive content.
13

---

5. If you used crowdsourcing conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A] We did not use any crowdsourcing conduct any research with human subjects.
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
14


Reviews:
[{"review_id": "gaiV3wXWnkY", "reviewer": "Reviewer_AXec", "summary": "Program optimization relies on rewriting a given program into an equivalent, but more efficient version of itself. These equivalent programs are searched within hand-engineered search spaces which are typically hard to customize. This paper introduces SpaceCraft, a domain-specific probabilistic-programming-language-abstraction which enables easy customization and construction of tensor program search spaces. SpaceCraft achieves this goal by providing modularity for composing program transformations. Further, the authors combine SpaceCraft with MAP estimation for automatic tensor programs optimization. SpaceCraft can cover the optimization search space used in many prior work and optimization with SpaceCraft delivers similar performance to prior work. On growing the search-space with a hardware specific module, the optimization yields a significant improvement over prior work.", "questions": "Can the authors can address the weaknesses raised in the previous section? \n\nSpecifically:\n \n1. Can they address the issues with the text of the paper?\n2. Can the address the drawbacks of their experimental design? \n3. Can the authors explain why certain important experiments have not been conducted?", "limitations": "No. The authors have not discussed limitations of this approach. The authors can discuss trade-offs of using MAP estimation vs. techniques such as reinforcement learning or simulated annealing used in prior work. The authors can also highlight the cost of learning a new DSL for crafting search spaces in SpaceCraft.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths:\n\n1. The program tackles an very useful problem - construction of the program search space. Typically,  construction of the \"right\" search space requires many trial and errors. SpaceCraft can significantly help reduce the complexity of this process.\n2. Clarity: The figures effectively convey the intended message. \n3. The use of probabilistic programs for tensor program optimization is novel.\n4. SpaceCraft covers the prior methods under a single framework.\n\n# Weaknesses:\n\n1. **Presentation**:  The presentation of the paper requires further refinement. Apart from plenty of minor grammatical mistakes (specified below), the paper also consist of many confusing sentences: \n\t1. The paper fails to specify what exactly are the hurdles with search space construction/customization in prior work. The paper simply states previous methods \"... requires surgical modification...\".  What are the surgical modifications? Why are they costly/infeasible?  \n\t2. In Line 44-52, paper states that \" These two observations lead to a simple yet powerful abstraction ...\" but do not to specify what exactly this abstraction is.  \n\t3. In Line 124-126, the paper states that SpaceCraft captures long-term structural arithmatic dependency between random variables and tensor program. The paper would benefit from specifically state why SpaceCraft achieves this. \n\t4. The presentation in Section 3.3 is very confusing. In this section, the reader needs to be aware of the authors intentions to understand some sentences. For example,  in Lines 163-165, what is implied by \"it\" is not clear. \n\n2.  **Experiments**: There are multiple drawbacks with the experiments section in the paper:\n\t1. A central benefit of SpaceCraft is that it simplifies the construction and modification of the optimization search spaces. This benefit however, and not been evaluated at all. This benefit should be evaluated by comparing the human efforts required for introducing new modules in prior frameworks versus the effort required for the same in SpaceCraft. \n\t2. The paper states (Line 278) that it takes 2 days to introduce a new module, but what is the reader supposed to compare this number with? What is the time it would take for introducing similar modules in prior work? \n\t3. The paper does not compare their optimization results to newer optimization techniques such as AutoTVM [1] and CHAMELEON [2] making it unclear how well this framework fares w.r.t. the current state-of-the-art. \n\t4. Comparisons in Figure 10 (b) are unfair as the search space of SpaceCraft and TVM are different. Are there fundamental limitations to TVM which make expanding the search space with operations from \"Use-Tensor-Core\" module infeasible? \n\t5. All the plots in Figure 8, 9, and 10  present results with normalized performance without clarifying what is the normalized performance measures. Further, the paper would benefit if more interpret-able statistics such as wall clock time or TFLOPS are used.\n\t6. The paper does not contain any measurement of the time it takes to perform the proposed automatic optimization. This is important as a high optimization cost would make the usage of SpaceCraft infeasible for many situations.   \n\n## Minor Weaknesses\n\n\n1. In the conclusions, the paper attempt to say they will open-source the framework. However, in the checklist, they mention the code and data is proprietary. In that case, how will they eventually open source the framework?\n2.   In the checklist, the authors state that limitation of the work are described in section 7. However, there are no sentences that discuss any limitations of the work in the entire paper.  \n\n\n## Errors: \n\nLine 18: has become\n\nLine 70 : Then, the framework finds\n\nLine 83: optimization programs\n\nLine 84: Therefore/Hence, we need\n\nLine 87: build a generic\n\nLine 288: We will open-source our framework \n\nLine 481: in the Appendix. \n\n\n## References:\n\n\n[1]  Learning to Optimize Tensor Programs, Tianqi Chen et al., NeurIPS 2018\n\n[2] Chameleon: Adaptive Code Optimization for Expedited Deep Neural Network Compilation, Byung Hoon Ahn et al., ICLR 2020", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "L9WGw3QQc3d", "reviewer": "Reviewer_AcTZ", "summary": "The paper outlines an approach to streamlining the process of tensor program optimization for different hardware substrates, by integrating a prior over optimization approaches. The latency of the compiled program (or a prediction thereof) is added to the prior probability of the optimization transforms as a potential function for random walk metropolis hastings, to explore alternative optimization approaches. Adding hardware specific optimizations is feasible and unlocks performance improvements on modern DL models, relative to TVM and PyTorch.", "questions": "How easily is the new system integrated with existing ML frameworks? Can we easily write code in PyTorch or JAX that calls into this system? Does optimization need to be re-run with every process startup, or is there a cache of previous work maintained?\n\nHow important is the prior, and how is it specified? None of the examples suggest the user can provide the prior probability of selecting any given categorical decision or relative probabilities of different transformation modules being applicable. So, is a uniform prior over categories assumed?\nDo you need to scale the prior relative to the runtime \"energy\"?\n\nIt seems like some discussion of Bayesopt / RL based approaches would be in order. I am surprised if noone has tried RL / contextual bandits in this space, but even if that's the case it's worth calling out as future work.\n\nIt seems like the system supports only categorical/discrete decisions. Are there cases in compiler optimization where a continuous random variable is useful? If we had such a case would it be easily supported by SpaceCraft?\n\nDo the programs produced yield interesting numerical issues? Sometimes, reordering arithmetic operations such as sum reductions can result in catastrophic cancellations.", "limitations": "Main limitation discussed is invalid output programs (OOM, etc), and validation being a viable approach to rectify this.\nIt'd be interesting to learn how long optimization takes, possibly see curves of \"best runtime by # steps\".", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Originality\nThe approach seems to be a novel one, and contributes to an important field of compiler optimizations for ML models. These systems facilitate performant deployment of models from training devices (GPU, high end CPU) to inference endpoints (mobile CPU, GPU). Previous work in this area could include Tensor Comprehensions, TVM, Halide, and beyond.\n\nQuality\nThe \"user experience\" of specifying kernels in the probabilistic DSL looks a little bit awkward. An upside, as a python-embedded traced DSL, we can interleave host language control flow to study conditions such as spatial-vs-reduction. A downside is that the resulting kernel code may need to be specialized to particular tensor shapes and ranks. It also appears (but I could have misunderstood?) that transformation modules are specific to tensor programs, which could limit the applicability to performance-critical code segments. The decision to use MCMC to sample this space is an interesting one, because I would expect the outcomes to be highly nonlinear relative to the decisions being sampled. A non-smooth posterior is a nightmare for most MCMC samplers, and I'd wonder whether there's any value to MCMC relative to some kind of prior-regularized evolutionary approach (or is that basically how we are viewing the MH work here?). Do you need to scale the prior relative to the runtime \"energy\"?\n\nClarity\nThis is a \"systems\" type of paper, where observing code and before/after different optimizations is probably an effective mode of presentation, but is necessarily limited by the page limit. The \"kinds of things\" that can be done with SpaceCraft are reasonably well presented, but I did not feel that the \"way those things are done\" was very clearly communicated. Also unclear was the time spent exploring for optimizations relative to a deterministic system like TVM.\n\nSignificance\n~2x speedup relative to PyTorch on BERT-Large seems significant, but without very-convenient tie-in to major ML systems may not be very usable by the community. The claim on 278-80 about ease of adding a novel TC-specific module are enticing, but probably deserve to be fleshed out further in the body of the paper rather than supplementary.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "IV3j1vSrMJz", "reviewer": "Reviewer_zFyW", "summary": "SpaceCraft is developed as an improvement over previous TVM IR such that it can unlock two key features: (1) not having to define the template like initial AutoTVM and its variants, and (2) achieving same programmability as Auto Scheduling such as Ansor. Overall, the paper devise a new abstraction with probabilistic language abstractions and search algorithms that improve upon previous black-box optimizations.\n\nFirst of all, it is important to emphasize that the paper makes a novel and timely contribution to the code-generation for tensor programs. There has been two major directions in code-generation. One being the black-box optimization building on TVM or TVM-like templates. Another being the auto-scheduling techniques that aim to minimize human intervention in devising the templates. The abstraction introduced in the paper initiates a third thrust of probabilistic programs. It seems also important to note that it builds on TVM.\n\nThe paper enables a stochastic search space construction which the contributions are straight-forward. Especially, Figure 3 clearly demonstrates an example. However, it would assist readers better if there could be some table that enumerates the additional methods used in the probabilistic program such as Sample-Tile and so on. Also, it would be good to add an example that contrasts different methods for tensor program generation.\n\nOne issue that I have with the paper is that unlike AutoTVM for instance, it does not outline the optimization details. For example, it is not intuitive as how to incorporate pre-trained cost-models from datasets such as TenSet. It seems that lack of these information in this respect leaves readers with lots of questions. How do we incorporate pre-trained cost-models? Is the “Learning-driven search” part exactly equivalent to the AutoTVM and its variants?, In that case doesn’t that make transfer learning more difficult as the search spaces become more heterogeneous? While these are just some hypothetical questions that I enumerated, adding more details to the Section 4 would prevent the confusion.\n\nRegarding evaluation, it seems that the largest benefits come from (1) multi-level tiling and (2) use tensor cores. However, reasons behind this nor insights were shared with the readers. This seems to make it difficult for readers to gain more insights as to how to optimize the search.\n\nLastly, the paper is focused on CPUs and GPUs. Considering the explosion of number of accelerators in the recent decade, it would be interesting to see how this can be incorporated into compilation for these diverse set of accelerators.\n\nIn sum, despite the above issues, the paper makes in important contribution that would enable interesting works on search and optimization. To this end, I would like to recommend the paper for a Weak Accept. However, with more concreteness to the paper addressing the issues above, I would be willing to reevaluate.", "questions": "* Can you fortify the Section 4?\n* Would this new abstraction require changes over prior works on black-box optimization given the probabilistic programs would render much diverse set of search spaces?\n* How can transfer learning in AutoTVM be enabled? Also, datasets such as TenSet?\n* How would the Figure 8 look like with some fused operations such as Conv + Pool + Relu? Could you share some results if possible with some of the common patterns in DNNs?", "limitations": "NA", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "+ timely contribution of probabilistic programming to tensor program compilation\n+ abstraction proposed in Section 3 can instigate interesting research in optimization that would benefit the readers.\n\n- unclear on details in Section 4 about learning\n- unclear on how it relates to prior works in more detail (not just high level statements in 3.3)", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}, {"review_id": "EVJeBl0CmRt", "reviewer": "Reviewer_UWHJ", "summary": "This paper introduces SpaceCraft, which improves the efficiency of designing tensor programs’ schedulers via establishing the search space in a composable and modular way. The key of SpaceCraft is to use the probabilistic programming method to describe the search space. Designers can improve their efficiency by combining SpaceCraft’s reusable program transformation modules with prior domain knowledge. The paper also introduces a learning-driven searching framework that can search efficiently under the guidance of a learned cost model. The experiments demonstrate that SpaceCraft achieves comparable optimization ability against TVM. ", "questions": "Please address the weaknesses listed above. Besides, in Figure 10(b), does TVM use Tensor Core as well? ", "limitations": "Supporting tensor computation unit is critical while challenging for existing exploration-based approached. Both FlexTensor and Ansor cannot support Tensor Core in a convenient manner. The readers would expect to see more explanation on the Use-Tensor-Core module.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This paper proposes a new programming model, which abstracts the search space as a probabilistic language and allows incorporation of domain knowledge, for tensor program optimization. The proposed method is technically sound. It seems that prior wok such as AutoTVM and Ansor can be easily implemented by using SpaceCraft, and SpaceCraft works well for domain-specific accelerators such as Tensor Core. The related work, including template-guided auto-tuning and auto-scheduling, is adequately cited.\n\nThis work is evaluated in terms of both expressiveness and optimization performance. Although it achieves comparable performance against TVM, it would be quite helpful to compare the performance with state-of-the-art approaches including FlexTensor and Ansor. Moreover, as SpaceCraft can provide a better search space, I suggest to directly compare the quality of produced search spaces against previous work, for example, ratio of high-performance candidate programs or performance difference of candidate programs, etc.\n\nThe advantage of the proposed learning-driven search over existing AutoTVM is unclear, which employs the cost model as well. I would like to see more ablation study on the effectiveness of the proposed learning framework.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept