Timezone: »

Tensor Program Optimization with Probabilistic Programs
Junru Shao · Xiyou Zhou · Siyuan Feng · Bohan Hou · Ruihang Lai · Hongyi Jin · Wuwei Lin · Masahiro Masuda · Cody Hao Yu · Tianqi Chen

Thu Dec 01 09:00 AM -- 11:00 AM (PST) @ Hall J #702

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.

Author Information

Junru Shao (OctoML)
Xiyou Zhou (OctoML)
Siyuan Feng (Shanghai Jiao Tong University)
Bohan Hou (School of Computer Science, Carnegie Mellon University)
Bohan Hou

I am a second-year Ph.D. student in Computer Science at Carnegie Mellon University, where I am very fortunate to be advised by Prof. Tianqi Chen. I am a member of Catalyst. My research interests lie in domain-specific system design and implementation, especially machine learning systems and deep learning compilers. Prior to coming to CMU, I received my Bachelor's degree in ACM Honors Class of Shanghai Jiao Tong University, where I worked under the supervision of Prof. Weinan Zhang and Prof. Yong Yu.

Ruihang Lai (Carnegie Mellon University)
Hongyi Jin (School of Computer Science, Carnegie Mellon University)
Wuwei Lin (OctoML)
Masahiro Masuda
Cody Hao Yu (Amazon Web Services)
Tianqi Chen (Carnegie Mellon University)

More from the Same Authors