GPU deployment using ROCmlink

IREE can accelerate model execution on AMD GPUs using ROCm.

Prerequisiteslink

In order to use ROCm to drive the GPU, you need to have a functional ROCm environment. It can be verified by the following steps:

rocm-smi | grep rocm

If rocm-smi does not exist, you will need to install the latest ROCm Toolkit SDK for Windows or Linux.

Get the IREE compilerlink

Download the compiler from a releaselink

Python packages are distributed through multiple channels. See the Python Bindings page for more details. The core iree-base-compiler package includes the ROCm compiler:

Stable releases Nightly releases Development packages

Stable release packages are published to PyPI.

python -m pip install iree-base-compiler

Nightly pre-releases are published on GitHub releases.

python -m pip install \
  --find-links https://iree.dev/pip-release-links.html \
  --upgrade --pre iree-base-compiler

Development packages are built at every commit and on pull requests, for limited configurations.

On Linux with Python 3.11, development packages can be installed into a Python venv using the build_tools/pkgci/setup_venv.py script:

# Install packages from a specific commit ref.
# See also the `--fetch-latest-main` and `--fetch-gh-workflow` options.
python ./build_tools/pkgci/setup_venv.py /tmp/.venv --fetch-git-ref=8230f41d
source /tmp/.venv/bin/activate

Tip

iree-compile and other tools are installed to your python module installation path. If you pip install with the user mode, it is under ${HOME}/.local/bin, or %APPDATA%\Python on Windows. You may want to include the path in your system's PATH environment variable:

export PATH=${HOME}/.local/bin:${PATH}

Build the compiler from sourcelink

Please make sure you have followed the Getting started page to build the IREE compiler, then enable the ROCm compiler target with the IREE_TARGET_BACKEND_ROCM option.

Tip

iree-compile will be built under the iree-build/tools/ directory. You may want to include this path in your system's PATH environment variable.

Get the IREE runtimelink

Next you will need to get an IREE runtime that includes the HIP HAL driver.

You can check for HIP support by looking for a matching driver and device:

$ iree-run-module --list_drivers

# ============================================================================
# Available HAL drivers
# ============================================================================
# Use --list_devices={driver name} to enumerate available devices.

        cuda: NVIDIA CUDA HAL driver (via dylib)
         hip: HIP HAL driver (via dylib)
  local-sync: Local execution using a lightweight inline synchronous queue
  local-task: Local execution using the IREE multithreading task system
      vulkan: Vulkan 1.x (dynamic)

$ iree-run-module --list_devices

  hip://GPU-00000000-1111-2222-3333-444444444444
  local-sync://
  local-task://
  vulkan://00000000-1111-2222-3333-444444444444

Download the runtime from a releaselink

Python packages are distributed through multiple channels. See the Python Bindings page for more details. The core iree-base-runtime package includes the HIP HAL driver:

Stable releases Nightly releases Development packages

Stable release packages are published to PyPI.

python -m pip install iree-base-runtime

Nightly pre-releases are published on GitHub releases.

python -m pip install \
  --find-links https://iree.dev/pip-release-links.html \
  --upgrade --pre iree-base-runtime

Development packages are built at every commit and on pull requests, for limited configurations.

On Linux with Python 3.11, development packages can be installed into a Python venv using the build_tools/pkgci/setup_venv.py script:

# Install packages from a specific commit ref.
# See also the `--fetch-latest-main` and `--fetch-gh-workflow` options.
python ./build_tools/pkgci/setup_venv.py /tmp/.venv --fetch-git-ref=8230f41d
source /tmp/.venv/bin/activate

Build the runtime from sourcelink

Please make sure you have followed the Getting started page to build IREE from source, then enable the HIP HAL driver with the IREE_HAL_DRIVER_HIP option.

Compile and run a program modellink

With the compiler and runtime ready, we can now compile programs and run them on GPUs.

Compile a programlink

The IREE compiler transforms a model into its final deployable format in several sequential steps. A model authored with Python in an ML framework should use the corresponding framework's import tool to convert into a format (i.e., MLIR) expected by the IREE compiler first.

Using a MobileNet model as an example, import using IREE's ONNX importer:

# Download the model you want to compile and run.
wget https://github.com/onnx/models/raw/refs/heads/main/validated/vision/classification/mobilenet/model/mobilenetv2-10.onnx

# Import to MLIR using IREE's ONNX importer.
pip install iree-base-compiler[onnx]
iree-import-onnx mobilenetv2-10.onnx --opset-version 17 -o mobilenetv2.mlir

Then run the following command to compile with the rocm target:

iree-compile \
    --iree-hal-target-backends=rocm \
    --iree-hip-target=<...> \
    mobilenetv2.mlir -o mobilenet_rocm.vmfb

Tip - HIP bitcode files

That IREE comes with bundled bitcode files, which are used for linking certain intrinsics on AMD GPUs. These will be used automatically or if the --iree-hip-bc-dir is empty. As additional support may be needed for different chips, users can use this flag to point to an explicit directory. For example, in ROCm installations on Linux, this is often found under /opt/rocm/amdgcn/bitcode.

Tip - HIP targets

A HIP target (iree-hip-target) matching the LLVM AMDGPU backend is needed to compile towards each GPU chip. Here is a table of commonly used architectures:

AMD GPU	SKU Name	Target Architecture	Architecture Code Name
AMD MI100	`mi100`	`gfx908`	`cdna1`
AMD MI210	`mi210`	`gfx90a`	`cdna2`
AMD MI250	`mi250`	`gfx90a`	`cdna2`
AMD MI300A	`mi300a`	`gfx942`	`cdna3`
AMD MI300X	`mi300x`	`gfx942`	`cdna3`
AMD MI308X	`mi308x`	`gfx942`	`cdna3`
AMD MI325X	`mi325x`	`gfx942`	`cdna3`
AMD RX7900XTX	`rx7900xtx`	`gfx1100`	`rdna3`
AMD RX7900XT	`rx7900xt`	`gfx1100`	`rdna3`
AMD PRO W7900	`w7900`	`gfx1100`	`rdna3`
AMD PRO W7800	`w7800`	`gfx1100`	`rdna3`
AMD RX7800XT	`rx7800xt`	`gfx1101`	`rdna3`
AMD RX7700XT	`rx7700xt`	`gfx1101`	`rdna3`
AMD PRO V710	`v710`	`gfx1101`	`rdna3`
AMD PRO W7700	`w7700`	`gfx1101`	`rdna3`

For a more comprehensive list of prior GPU generations, you can refer to the LLVM AMDGPU backend.

The iree-hip-target option support three schemes:

The exact GPU product (SKU), e.g., --iree-hip-target=mi300x. This allows the compiler to know about both the target architecture and about additional hardware details like the number of compute units. This extra information guides some compiler heuristics and allows for SKU-specific tuning specs.
The GPU architecture, as defined by LLVM, e.g., --iree-hip-target=gfx942. This scheme allows for architecture-specific tuning specs only.
The architecture code name, e.g., --iree-hip-target=cdna3. This scheme gets translated to closes matching GPU architecture under the hood.

We support for common code/SKU names without aiming to be exhaustive. If the ones you want are missing, please use the GPU architecture scheme (2.) as it is the most general.

Run a compiled programlink

To run the compiled program:

iree-run-module \
    --device=hip \
    --module=mobilenet_rocm.vmfb \
    --function=torch-jit-export \
    --input="1x3x224x224xf32=0"

The above assumes the exported function in the model is named torch-jit-export and it expects one 224x224 RGB image. We are feeding in an image with all 0 values here for brevity, see iree-run-module --help for the format to specify concrete values.