HloEnv

HloEnv is an environment based on Accelerated Linear Algebra (XLA) for deep learning 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.

HloEnv has the following major features:

• The alternative graph representation

  HloEnv frames the graph optimization problem as a sequential rewrite decision problem, serving as a single-player game engine. See always_pick_first and uniform_policy for examples of interacting with HloEnv and optimizing a HLO computation grpah.

• Python interface for full XLA optimization passes and pipelines

  HloEnv provides a Python interface to control and modify existing graph optimization passes. See general_fusion_pipeline for an example of our customized XLA optimization pipeline.

The design of HloEnv points to a potential future where deep learning compiler engineers only need to develop and maintain a simple set of rewrite rules. The complicated heuristics are left to machine learning-generated optimization strategies that generalize to both new deep learning models and new deep learning hardware.

Citations

HloEnv is described in details in the following paper:

@inproceedings{HloEnv2022,
  author = {Chin Yang Oh, Kunhao Zheng, Bingyi Kang, Xinyi Wan, Zhongwen Xu, Shuicheng Yan, Min Lin, Yangzihao Wang},
  title = {HloEnv: A Graph Rewrite Environment for Deep Learning Compiler Optimization Research},
  booktitle = {Workshop on ML for Systems at NeurIPS 2022},
  year = {2022},
  series = {NeurIPS '22},
  month = dec,
}

Docs

The full documentation for HloEnv can be found at this link.

Quick Start

Download the pre-built HloEnv wheel from the releases, and navigate to the folder that contains the .whl file and run:

python3 -m pip install hloenv-[x.y.z]-cp38-cp38-linux_x86_64.whl

To run examples, you need to install the following dependencies in your python package:

• TensorFlow==2.9.0
• NumPy

Build HloEnv

To build HloEnv and evaluate HLO graph runtime, you need to install the following dependencies:

• Bazel
• liblapack-dev
• CUDA

To clone HloEnv and install Tensorflow, run the following command:

git clone git@github.com:sail-sg/hloenv.git
cd hloenv
python -m venv hloenv-env
source hloenv-env/bin/activate
python3 -m pip install tensorflow==2.9.0

When met the minimum requirement, building HloEnv Python Wheels can be very simple.

In the root directory of HloEnv, run the following command:

make build

Will build the wheels file and put it under the folder dist.

make install

Will install the freshly built Python Wheels and all the dependencies.

HLO Dataset

You can find the unoptimized HLO text files harvested from JAX implemented repositories on GitHub here