A Memory-Bounded Best-First Beam Search and Its Application to Scheduling Halide Programs

Authors

  • Chao Gao Huawei Canada Research Centre
  • Jingwei Chen Huawei Canada Research Centre
  • Tong Mo Huawei Canada Research Centre
  • Tanvir Sajed Huawei Canada Research Centre
  • Shangling Jui Huawei Technologies, China
  • Min Qin Huawei Technologies, China
  • Laiyuan Gong Huawei Technologies, China
  • Wei Lu Huawei Canada Research Centre

Keywords:

Problem Solving Using Search, Analysis Of Search Algorithms, Real-time Search, Machine And Deep Learning In Search

Abstract

Beam search is a popular algorithm for solving real-world problems --- especially where search space is an enormously large tree but real-time solutions are most preferred. We present a memory-bounded best-first beam search (MB2FBS), which can be viewed as an improved and generalized version of standard beam search in trees. The algorithm takes three parameters --- in contrast to the singular parameter beam size in standard beam search. We discuss how to recover standard beam search and how to realize other search behavior by setting these three parameters correspondingly. In particular, we show that the principal version of MB2FBS can be thought as an algorithm whose search expense is similar or upper bounded by beam search of certain beam size; however it often finds better solutions as it decides the number of nodes to be searched each depth dynamically with respect cost landscape. We apply our algorithm for tensor program auto-scheduling in Halide, an important industrial problem that uses tree search for optimizing tensor program executions. We show that the principal variants of MB2FBS deliver better empirical results than the highly optimized beam search counterpart. Most importantly, it finds superior schedules while no more computation cost is used for search, which is highly desirable for real-time program compilation and optimization.

Downloads

Published

2022-07-17