Tensor Program Optimization with Probabilistic Programs
This addresses the need for efficient tensor program optimization in diverse deployment environments, enabling domain experts to grow search spaces modularly, though it is incremental relative to existing frameworks.
The paper tackles the problem of automatically optimizing tensor programs for deep learning by introducing MetaSchedule, a domain-specific probabilistic programming language that constructs a rich search space, resulting in a 48% speedup on end-to-end workloads.
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.