Tensor Program Optimization with Probabilistic Programs
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.