Taichi on ROCm

2025-08-03

5 min read time

Applies to Linux

Taichi is an open-source, imperative, and parallel programming language designed for high-performance numerical computation. Embedded in Python, it leverages just-in-time (JIT) compilation frameworks such as LLVM to accelerate compute-intensive Python code by compiling it to native GPU or CPU instructions.

For hardware, software, and third-party framework compatibility between ROCm and Taichi, see the following resources:

• System requirements (Linux)

• Taichi compatibility

Note

Taichi is supported on ROCm 6.3.2.

Install Taichi

To install Taichi on ROCm, you have the following options:

• Use the prebuilt Docker image (recommended)

• Use a wheels package

• Build your own docker image

Use a prebuilt Docker image with Taichi pre-installed

Docker is the recommended method to set up a Taichi environment, and it avoids potential installation issues. The tested, prebuilt image includes Taichi, Python, ROCm, and other dependencies.

1. Pull the Docker image

   docker pull rocm/taichi:taichi-1.8.0b1_rocm6.3.2_ubuntu22.04_py3.10.12
   

2. Launch and connect to the container

   docker run -it -d --network=host --device=/dev/kfd --device=/dev/dri --ipc=host --shm-size 64G \
   --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v $(pwd):/taichi_dir \
   --name rocm_taichi rocm/taichi:taichi-1.8.0b1_rocm6.3.2_ubuntu22.04_py3.10.12
   

Use a wheels package

The Taichi .whl packages are hosted on the AMD PyPI repository. Instead of manually downloading the files, you can simply install Taichi using pip with the provided URL. This command will automatically download and install the appropriate .whl file.

pip install amd-taichi==1.8.0b1 --extra-index-url=https://pypi.amd.com/simple
sudo apt-get update
sudo apt-get install -y lld

Build your own Docker image

If you prefer to use the ROCm Ubuntu image, or already have a ROCm Ubuntu container, follow these steps to install Taichi in the container.

1. Pull the ROCm Ubuntu Docker image

   docker pull rocm/dev-ubuntu-22.04:6.3.2
   

2. Launch the Docker container

   docker run -it -d --network=host --device=/dev/kfd --device=/dev/dri --ipc=host --shm-size 64G \
   --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v $(pwd):/taichi_dir \
   --name rocm_taichi rocm/dev-ubuntu-22.04:6.3.2 /bin/bash
   

3. Inside the running container, install build dependencies

   sudo apt-get update && apt-get install -y --no-install-recommends \
   git wget vim \
   freeglut3-dev libglfw3-dev libglm-dev libglu1-mesa-dev \
   libjpeg-dev liblz4-dev libpng-dev libssl-dev \
   libwayland-dev libx11-xcb-dev libxcb-dri3-dev libxcb-ewmh-dev \
   libxcb-keysyms1-dev libxcb-randr0-dev libxcursor-dev libxi-dev \
   libxinerama-dev libxrandr-dev libzstd-dev \
   python3-pip cmake pybind11-dev ca-certificates \
   llvm-15 clang-15 lld-15 \
   && apt-get clean && rm -rf /var/lib/apt/lists/*
   

4. Add LLVM 15 to PATH

   export LLVM_DIR=/usr/lib/llvm-15
   export PATH=${LLVM_DIR}/bin:$PATH
   

5. Clone the ROCm/taichi repository with the desired branch

   cd <working-directory>
   git clone --recursive https://github.com/ROCm/taichi -b amd-release/v1.8.0b1
   cd taichi
   

6. Build the Taichi wheel

   export GPU_TARGETS=gfx90a
   export TAICHI_CMAKE_ARGS="-DTI_WITH_VULKAN=OFF -DTI_WITH_OPENGL=OFF -DTI_BUILD_TESTS=ON -DTI_BUILD_EXAMPLES=OFF -DCMAKE_PREFIX_PATH=${LLVM_DIR}/lib/cmake -DCMAKE_CXX_COMPILER=${LLVM_DIR}/bin/clang++ -DTI_WITH_AMDGPU=ON -DTI_WITH_CUDA=OFF -DTI_AMDGPU_ARCHS=${GPU_TARGETS}"
   
   ./build.py
   

7. Install the Taichi .whl file

   pip3 install <taichi-src-dir>/dist/taichi*.whl
   

Test the Taichi installation

Clone the ROCm/taichi repository

sudo apt-get update
sudo apt-get install -y git
git clone --recursive https://github.com/ROCm/taichi -b amd-release/v1.8.0b1

To test the Taichi installation, run the laplace example in the source code:

python3 taichi/python/taichi/examples/algorithm/laplace.py

laplace example output:

[Taichi] version 1.8.0, llvm 15.0.0, commit f7911653, linux, python 3.10.12
[Taichi] Starting on arch=amdgpu
0.0
4.0
0.0
0.0
4.0
0.0
0.0
4.0
0.0
0.0

Run a Taichi example

Several examples have been collected in the ROCm/taichi_examples repository. This repository contains a Dockerfile for building a container and installing Taichi. It contains a README.md file that guides the building of the container using the Dockerfile and running the examples.

1. Clone the repository:

   git clone https://github.com/ROCm/taichi_examples.git
   

2. If you have Taichi installed in your environment, you do not need to build an additional container. Simply install the example dependencies:

   pip3 install pillow
   pip3 install --no-cache-dir torch --index-url https://download.pytorch.org/whl/nightly/rocm6.3/
   

3. Run the examples scripts from the taichi_examples directory:

   cd taichi_examples
   ./run_demos.sh
   ./run_algorithm_graph_examples.sh
   

Refer to the AMD ROCm blog to search for Taichi examples and best practices to optimize your workflows on AMD GPUs.