Learning Local Advantage Functions for Generalizable Graph Optimizations

Machine learning compilers rely on making optimized decisions in order to generate efficient code for a given computation graph. Many of these decision making processes can be formulated as graph optimization problems. The solution to these graph optimization problems is typically computed based on human designed heuristics. Learning/search-based methods have been recently investigated to improve upon or remove the need of human designed heuristics. However, existing methods that can reliably provide high quality solutions require iterative evaluations on real hardware. The evaluations can be costly especially for large graphs, making these methods infeasible to be deployed in production. To reduce or eliminate the evaluation cost, learning graph optimization strategies that can generalize across graphs is desirable. In this work, we propose learning local advantage functions for generalizable compiler graph optimizations. The learned model can be trained offline with supervised learning on massive amount of training data and then used to guide the search of optimal decisions on previously unseen graphs. We demonstrate the effectiveness of our approach on the operation fusion task and discuss several challenges we encountered in practice.