Tensor Program Optimization with Probabilistic Programs
Abstract
Automatic optimization for tensor programs becomes increasingly important as we deploy deep learning in various environments, and efﬁcient optimization relies on a rich search space and effective search.
Most existing efforts adopt a search space which lacks the ability to efﬁciently enable domain experts to grow the search space.
This paper introduces MetaSchedule, a domain-speciﬁc 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 ﬁnd 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.
Introduction
Deep learning has become pervasive in daily life.
From video understanding [28], natural language understanding [15] and recommendation system [29] to autonomous driving [21], different deep learning models are deployed on different hardware platforms and devices.
Deep learning frameworks usually rely on manually optimized libraries [13, 22] to accelerate deployment.
Engineers need to choose from many tensor programs that are logically equivalent but differ signiﬁcantly in performance due to memory access, threading, and the use of specialized hardware primitives.
The engineering effort required for tensor program optimization has become a signiﬁcant bottleneck for machine learning deployment with the growing number of models and hardware backends.
Initial Tensor Program e
Search Space S(e 0 )
Optimal Programe * Search Algorithm for i in range(1024): for j in range(1024): for k in range(1024): C[i, j] += A[i, k] * B[j, k] for i 0 , j 0 in grid(16, 8): for i 1 , j 1 in grid(8, 16): for k 0 in range(1024): for i 2 , j 2 in grid(8, 8): C[...] += ...
for i 0 , j 0 in grid(64, 8): for i 1 , j 1 in grid(4, 32): for k 0 in range(64): for i 2 , j 2 in grid(4, 4): for k 1 in range(16): C[...]
+= ...
Key
[Figure]
optimization
Examples
Automatic program optimization [10, 43, 1] is a recent sequence of efforts that aims to use machine learning to solve this problem.
There are two vital elements of automatic tensor program optimizations.
First, a search space is deﬁned to provide a possible set of equivalent tensor programs.
Then existing systems use learning-based search algorithms to ﬁnd an optimized tensor program in the search space with feedback from the deployment environment.
Most of the current approaches [10, 43, 1, 25, 4] use pre-deﬁned search spaces that effectively encode the domain knowledge of the authors once and focus on developing efﬁcient search algorithms.
While efﬁcient 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 [24] and hardware primitives [30] grow.
As a result, there is a strong need to enable easy customization and construction of the search space at scale by taking inputs from system engineers and domain experts.
Unfortunately, any change to search space construction currently requires surgical modiﬁcations to the automatic program optimization frameworks.
This research asks the following question: can we 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 an afﬁrmative answer to the question with two key observations.
First, we can parameterize an optimization search space by the initial program followed by a sequence of transformations on the program.
Next, using this parameterization, domain experts can then provide probabilistic choices that represent possible transformations after examining the program state.
These two observations lead to a simple yet powerful abstraction for search space construction through a domain-speciﬁc probabilistic language.
Finally, our framework composes multiple possible probabilistic transformations to form a 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 a learning-driven framework to ﬁnd optimized tensor programs speciﬁed by the search space constructed using our abstraction.
• We build an end-to-end system that can 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 on multiple hardware backends.
Experimental results show that our abstraction is expressive enough to cover the optimization space of a diverse set of tensor programs, delivering a competitive performance of popular deep learning models, and convenient to incorporate hardware-speciﬁc knowledge into the search space to outperform the state-of-the-art frameworks.
for i in range(1024): B[i] =
ReLU(A[i])
Initial tensor program: eparallel for i 0 in range(32):
for i 1 in range(8): i = i
* 32 + i 1 * 4 B[i : i + 4] =
ReLU(A[i : i + 4])
Equivalent optimizedprogram: e* ① for i 0 in range(32): for i 1 in range(8): for i 2 in range(4): i = i
* 32 + i 1 * 4 + i
B[i] =
ReLU(A[i]) ② ③ ③ Vectorize ② Parallelize ① Split
i, 32, 8, 4 ii
Transformation
Parameterization parameterized by: e 0 + ① parameterized by: e 0 + ①②③
Equivalent intermediate program: e
Parameterization
Equivalent Programs Induced by Parameterized Transformation
Background and Problem Overview
Figure 1 shows a typical workﬂow for tensor program optimization.
For a given program e 0 , a typical tensor program optimization framework will generate candidates from a pre-deﬁned search space
S ( e 0 ) containing semantically-equivalent programs.
Then the framework ﬁnds optimized tensor program e ⇤ 2 S ( e 0 ) with the minimum latency on the target hardware.
A typical search space S ( e 0 ) contains choices over threading, loop ordering, memory access, and hardware primitives.
Deﬁning the search space S ( e 0 ) for a wide range of tensor programs brings several challenges.
First, S ( e 0 ) is highly dependent on e 0 .
For example, S ( e 0 ) of a compute-intensive program (e.g., Dense ) needs to consider many more possible conﬁgurations than a communicationintensive program such as ReLU .
The space also differs signiﬁcantly in different hardware domains.
For example, S ( e 0 ) on CPU involves multi-core parallelism and vectorization, while S ( e 0 ) on GPU involves thread binding and tensorization.
Finally, as the hardware and model settings change, we need to bring in fresh domain experts to update the S ( e 0 ) to 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 to be able to build a 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 over 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 speciﬁed in our abstraction.
We will address the above goals in the following two sections.
Composable Search Space Parameterization
This section presents MetaSchedule, a probabilistic approach to search space parameterization.
3.1 Stochastic Search Space Construction
MetaSchedule constructs a search space S ( · ) with stochastic program transformations as the primitive building blocks.
Traditional program optimization can usually be represented by a sequence of transformations ⌧ , where at step i , the program e i  1 is transformed into a semantically-equivalent program e i , which ﬁnally leads to the optimized program e n .
MetaSchedule generalizes this idea by allowing further parameterization of each transformation step in ⌧ .
Taking Figure 2 as an example: e 0 is the initial program for the program B = ReLU ( A ) 1 .
In
MetaSchedule, transformation t 1 = Split is parameterized by a loop i and a sequence of integers indicating the loop extents after splitting; Similarly, transforms t 2 =
Parallelize and t 3 =
Vectorize are parameterized by loops respectively.
As a result, an optimized program e n is obtained by applying a sequence of parameterized transformations ⌧ to the initial program e 0 .
Accordingly, the search space S ( e 0 ) is composed of e 0 and all possible sequences of parameterized transformations.
[Figure]
' in grid( θ 4 , θ 5 ): D[...]
= ...
θ ReLU = θ ReLU =
Probabilistic
Stochastic Transformation
On the other hand, it could be less practical for practitioners to determine the best combination of the parameter values in transformations.
For instance, in Figure 2, it is usually efﬁcient to use 512-bit vectorization over the inner loop when AVX-512 2 vector instructions are available on Intel CPUs, or 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 efﬁciently deﬁne parameterized transformations without worrying about candidate values, MetaSchedule introduces random variables drawn from analysis, sampling .
Parameterized by random variables, a transformation naturally becomes stochastic, and the underlying probabilistic space reﬂects the space of possible transformations.
As illustrated in Figure 3, when creating Split transformations to tile loop i and j in the Dense operator, the tile sizes are drawn by random variables ✓ 0  3 deﬁned from Sample-Tile .
In this way, the Split transformation becomes stochastic.
Similarly, we use Sample-Compute-Location to enumerate valid loops in Dense after splitting for ReLU to fuse its computation.
In summary, 7 lines of MetaSchedule program covers a family of possible optimized tensor programs with stochastic transformations in its search space S ( e 0 ) , where e 0 is Dense-ReLU .
Notably, unlike orthogonal grid space in hyperparameter tuning, MetaSchedule captures longterm structural and arithmetic dependency between random variables and the tensor program e i being transformed.
As demonstrated on Step 2  in Figure 3, sampling distribution from
Sample-Compute-Location depends on the latest tensor program e 5 , 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 efﬁcient 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, we observe that it is straightforward to group a sub-sequence of transformations for a particular ﬁne-grained 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 efﬁciency.
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, a transformation module in MetaSchedule is deﬁned as either atomic stochastic transformation, or composition of
[Figure]
where analysis is done interactively to identify spatial (data parallel) and reduction loops, and then apply Split with the tiling factors drawn from Sample-Tile .
A ﬁnal Reorder organizes the loops into proper tiles.
Stochastic-Compose procedure C OMPOSE (modules, program) while locations are not exhausted do location next-location (program) m ⇠ Sample (modules) program m (location, program) end while end procedure
[Figure]
program analysis, sampling as well as smaller transformations.
Each transformation module can have a meaningful name so that it can be easily adopted by many workloads to hierarchically construct a search space.
Figure 4 shows hierarchical composition of transformation modules.
Speciﬁcally,
Multi-Level-Tiling interleaves program analysis on loops and the stochastic tiling of the loop structure and organizes the original tensor program into a 5-level tiling structure.
Notably, the transformation module is generic to tensor programs and thus could be applied to a variety of operators, including conv1d , conv3d , matmul , etc.
Figure 5 depicts an example of composing a search space with transformation modules.
In this simple example, we select a 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 a 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 of them can be covered by the MetaSchedule framework.
Domain speciﬁc languages for program transformations used by prior frameworks [32, 5, 9, 37] allow developers to easily optimize a program manually.
When there is 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 [10, 2, 25, 27] fully relies on developers to deﬁne a search space.
In
MetaSchedule, it means all random variables in a search space are deﬁned ahead of the transformations, so there is no interaction between program analysis and follow-up random sampling choices
The Probabilistic
[Figure]
Auto-scheduling [43, 45, 1, 19] requires developers to implement workload agnostic transformation rules.
MetaSchedule achieves the same programmability and functionality through speciﬁc probabilistic transformation modules that correspond to the search space generation rules.
Notably, all approaches mentioned above have important use-cases in tensor program optimizations, depending on how much domain knowledge we want to incorporate for a particular scenario.
By decoupling the search space construction from the search, we effectively build a single framework for all the use cases and enable further customization without surgical changes to the system.
Learning-driven Search
The last section provides a modular abstraction for search space.
We still need to do an effective search to ﬁnd an optimized program within the search space.
This section provides a generic learning-driven framework to ﬁnd an optimized program.
Objective formalization.
For a given probabilistic program e 0 , let us use ⌧ to denote the transformations performed on e 0 .
⌧ can be sampled from a prior distribution speciﬁed by the probabilistic program.
We deﬁne g ( e 0 , ⌧ ) to be the tensor program after applying transformation ⌧ to e 0 .
Let f ( e ) be the latency of the particular program e on the hardware environment.
We deﬁne a posterior probability of an optimized program as:
[Equation]
Intuitively, we want to assign a higher probability to the programs that perform well.
Our ﬁnal goal is to ﬁnd ⌧ ? = argmax ⌧ P ( ⌧ | e 0 ) 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.
Figure 6 demonstrates an example tracing process.
During program execution, the system records all samplings and transformations while ignoring control ﬂow and other constructs of the host language.
The resulting trace is a sequence of MetaSchedule primitives with only sampling and transformation instructions, which could be re-executed as a normal MetaSchedule program.
We can then continue to explore different sampling choices for a 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 workﬂow of our learning-driven framework.
The search algorithm ﬁrst samples the MetaSchedule program to obtain a collection of traces.
Then it continues to explore the space condition on the traces.
Notably, there is a signiﬁcantly higher cost measuring f ( e ) directly on the hardware, so we also incorporated a proxy cost model
ˆ f ( e ) , which is updated throughout the process, similar to previous efforts on tensor program optimization [10, 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.
[Figure]
Cost model.
Our approach allows extensive cost models, enabling us to supply those pre-trained from existing datasets [44].
We pick a tree-boosting-based cost model in ˆ f ( · ) by default and leverage a common set of features that are used in previous works [43].
Trace validation.
Importantly, invalid traces may show up as we propose updates.
Such a scenario can happen when some of the random variable choices go beyond the physical hardware limit or a variable that induces changes to the execution sequence.
Instead of enforcing a conservative proposal, we introduce a validator that validates the correctness of the trace.
The trace validation allows us to move around the space more freely while still ensuring the correctness of the sample outcome to be on the right support set.
Related Work
Tensor Program Transformations are proposed by many prior works, such as Halide [32], TVM [9],
Tiramisu [5] and TACO [23, 37].
Note that all the previous transformation languages are deterministic and cannot be directly used for search space construction, meaning that they have to introduce a separate programming model to express a search space.
This paper makes a simple but powerful generalization to domain-speciﬁc probabilistic language.
The resulting abstraction enables a uniﬁed approach to deterministic transformation and search space construction.
Black-box optimization has been adopted in high-performance computing libraries [16, 3].
Recent advances in automatic tensor program optimization brought a rich set of techniques to accelerate search through better cost modeling [10, 4, 34] and learning-based search [2, 25, 1, 19, 44], which could be incorporated into MetaSchedule search.
Different variations of pre-deﬁned search spaces have also been proposed that couple with the automatic tensor program optimization frameworks [10, 43, 1].
Polyhedral model [40, 6, 39] is one useful way to construct a rich pre-deﬁned search space.
This paper focuses on modular search space construction and provides orthogonal contributions to these prior works.
Probabilistic programming language is a powerful abstraction for incorporating domain knowledge and probabilistic inference.
There are many general-purpose probabilistic languages, such as
Church [17], Stan [8], Pyro [7], NumPyro [31], PyMC3 [35] and Edward [38].
This paper proposes a domain-speciﬁc probabilistic language for tensor program optimization with specializations such as tensor program analysis that would otherwise be opaque to the previous systems.
Our learningdriven search can be viewed as an application of previous works [42, 33, 41, 46] that use tracing to divide programs into subprograms with ﬁxed support.
We focus on the MAP inference problem where the posterior depends on an unknown cost function.
We solve the problem through a learned cost-model-driven evolutionary search over traces and with validation.
Automatic neural program synthesis [12, 18, 11] 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 be incorrect and use input-output pairs to ﬁlter out those erroring programs.
Our compiler approach requires us to ensure the correctness of all transformations.
However, some ideas like validation after creation might be reusable.
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 state-of-the-art optimization techniques?
To answer this question, we 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: 8d4f4dd73f ), which incorporates AutoTVM [10] and Ansor [43], is used as the state-of-the-art tensor program optimization system, and we pick the best among the two in each respective setups.
Full operators and hardware conﬁgurations are documented in Appendix A.2.
Figure 8 shows that, in all cases on CPU and GPU, MetaSchedule delivers performance comparable 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 a signiﬁcant 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 state-of-the-art works for end-to-end models?
Therefore, a series of experiments are conducted to compare MetaSchedule and TVM, including
BERT-Base [14], ResNet-50 [20], and MobileNet-v2 [36] on both CPU and GPU.
As shown in
Figure 9, MetaSchedule performance is on 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.
[Figure]
6.3 Search Space Composition and Hardware-Speciﬁc 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 it translate to performance?
We design an ablation study for transformation modules composition.
As indicated in Figure 10a, by progressively enriching the search space, the performance of optimized tensor programs consistently increases.
Composed with a hardware-speciﬁc module Use-Tensor-Core , MetaSchedule delivers signiﬁcantly 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 in Figure 10b.
Speciﬁcally, on BERT-large workloads, MetaSchedule with
Use-Tensor-Core delivers 48% speedup over TVM.
Notably, it took a graduate student only 2 days to 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 are in Appendix A.4.
Conclusion
This paper presents MetaSchedule, a programming model to describe search space construction in tensor program optimization.
Our method abstracts search space as a probabilistic language and enables ﬂexible incorporation of domain knowledge by allowing practitioners to implement customized probabilistic programs.
A learning-driven search algorithm is developed on top of the probabilistic language abstraction, which delivers competitive performance with state-of-the-art frameworks.
In the future, we will explore and modularize declarative API for various hardware environments.
Therefore, we will open-source our framework and hope it could enable broader collaboration between the machine learning deployment engineers and intelligent machine learning algorithms for tensor programs.
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.