Learning to Optimize Halide with Tree Search and Random Programs

 Andrew Adams   Karima Ma   Luke Anderson   Riyadh Baghdadi   Tzu-Mao Li   Michaël Gharbi 
Facebook AI Research UC Berkeley MIT CSAIL MIT CSAIL MIT CSAIL Adobe Research
 Benoit Steiner   Steven Johnson   Kayvon Fatahalian   Frédo Durand   Jonathan Ragan-Kelley 
Facebook AI Research Google Stanford University MIT CSAIL UC Berkeley


Abstract

We present a new algorithm to automatically schedule Halide programs for high-performance image processing and deep learning. We significantly improve upon the performance of previous methods, which considered a limited subset of schedules. We define a parameterization of possible schedules much larger than prior methods and use a variant of beam search to search over it. The search optimizes runtime predicted by a cost model based on a combination of new derived features and machine learning. We train the cost model by generating and featurizing hundreds of thousands of random programs and schedules. We show that this approach operates effectively with or without autotuning. It produces schedules which are on average almost twice as fast as the existing Halide autoscheduler without autotuning, or more than twice as fast with, and is the first automatic scheduling algorithm to significantly outperform human experts on average.

Publication

Andrew Adams, Karima Ma, Luke Anderson, Riyadh Baghdadi, Tzu-Mao Li, Michaël Gharbi,
Benoit Steiner, Steven Johnson, Kayvon Fatahalian, Frédo Durand, Jonathan Ragan-Kelley.
Learning to Optimize Halide with Tree Search and Random Programs
ACM Transactions on Graphics 38(4) (Proceedings of ACM SIGGRAPH 2019)

halide_autoscheduler_2019.pdf

Code

The core autoscheduler has been integrated into Halide main here: https://github.com/halide/Halide/tree/main/src/autoschedulers/adams2019

For full details of the experiments and benchmark applications, see this branch in the Halide repository: https://github.com/halide/Halide/tree/standalone_autoscheduler

Acknowledgements

This work was partially funded by Toyota, as well as the NSF/Intel Partnership on Computer Assisted Programming for Heterogeneous Architectures through grant CCF-1723445.