Timezone: »

 
HloEnv: A Graph Rewrite Environment for Deep Learning Compiler Optimization Research
Chin Yang Oh · Kunhao Zheng · Bingyi Kang · Xinyi Wan · Zhongwen Xu · Shuicheng Yan · Min Lin · Yangzihao Wang

@
in Machine Learning for Systems »
Event URL: https://openreview.net/forum?id=3gsBTrhZXr »

We introduce HloEnv, an environment based on Accelerated Linear Algebra (XLA) for deep learning (DL) compiler optimization research. HloEnv transforms all graph rewrites into a common representation, providing a flexible interface to control and modify existing graph optimization passes. In this representation, an XLA pass is converted into a set of sequential rewrite decisions, which control when and if the rewrites are applied. Along with HloEnv, we present a dataset with broad coverage of computation graphs drawn from modern real-world machine learning models. We select two XLA passes with the largest impact on the runtime of the compiled program, and explore the potential for further improvement over XLA in this decision space. We show that using simple heuristics for decision-making can achieve on-par or better performance than XLA. Using search algorithms further boosts performance. We intend for HloEnv and our dataset to be an open-source, community-driven effort that helps spur advances in DL compiler optimization research.

Author Information

Chin Yang Oh (Carnegie Mellon University)
Kunhao Zheng (Sea AI Lab)
Bingyi Kang (National University of Singapore)
Xinyi Wan (Zhejiang University)
Zhongwen Xu (Sea AI Lab)
Shuicheng Yan (Sea AI Lab)
Min Lin (Sea AI Lab)
Yangzihao Wang (Sea AI Lab)
Yangzihao Wang

Did CUDA programming model for graph computing for PhD. Ex-TensorFlow developer at Google Brain. Did large-scale training and GNN research at Tencent and WeChat. Current research interests include ML for system and system for ML.

More from the Same Authors