Taichi Lang on ROCm installation

Taichi Lang on ROCm installation#

2025-12-09

5 min read time

Applies to Linux

System requirements#

To use Taichi Lang 1.8.0b2, you need the following prerequisites:

ROCm version: 7.0.0 (recommended)
Operating system: Ubuntu 22.04, 24.04
GPU platform: AMD Instinct™ MI355X, MI325X, MI300X, MI250X, MI210
Python: 3.12.3, 3.10.12

Install Taichi#

To install Taichi Lang 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 Lang pre-installed#

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

Pull the Docker image:
Ubuntu 24.04
docker pull rocm/taichi:taichi-1.8.0b2_rocm7.0.0_ubuntu24.04_py3.12.3
See rocm/taichi:taichi-1.8.0b2_rocm7.0.0_ubuntu24.04_py3.12.3 on Docker Hub.
Ubuntu 22.04
docker pull rocm/taichi:taichi-1.8.0b2_rocm7.0.0_ubuntu22.04_py3.10.12
See rocm/taichi:taichi-1.8.0b2_rocm7.0.0_ubuntu22.04_py3.10.12 on Docker Hub.

Launch and connect to the container:

Ubuntu 24.04

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

Ubuntu 22.04

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

Use a wheels package#

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

pip install amd-taichi==1.8.0b2 --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.

Pull the ROCm Ubuntu Docker image:
Ubuntu 24.04
docker pull rocm/dev-ubuntu-24.04:7.0-complete
See rocm/dev-ubuntu-24.04:7.0-complete on Docker Hub.
Ubuntu 22.04
docker pull rocm/dev-ubuntu-22.04:7.0-complete
See rocm/dev-ubuntu-22.04:7.0-complete on Docker Hub.

Launch the Docker container:

Ubuntu 24.04

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

Ubuntu 22.04

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

Inside the running container, install build dependencies:

apt-get update && apt-get install -y --no-install-recommends \
   git wget vim git 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 python3-venv rocm-llvm-dev \
   gdb python3-dbg

Setup LLVM:

export LLVM_VERSION=20
export LLVM_PATH=/usr/lib/llvm-${LLVM_VERSION}
export PATH=${LLVM_PATH}/bin:$PATH

wget https://apt.llvm.org/llvm.sh \
   && chmod +x llvm.sh \
   && apt-get update && apt-get install -y \
   lsb-release software-properties-common gnupg \
   && ./llvm.sh ${LLVM_VERSION} llvm clang lld

Clone the ROCm/taichi repository with the desired branch:

cd /taichi_dir
git clone --recursive https://github.com/ROCm/taichi -b release/v1.8.0b2
cd taichi

Build the Taichi Lang wheel:

export GPU_TARGETS=gfx950,gfx942,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_PATH}/lib/cmake -DCMAKE_CXX_COMPILER=${LLVM_PATH}/bin/clang++ -DTI_WITH_AMDGPU=ON -DTI_WITH_CUDA=OFF -DTI_AMDGPU_ARCHS=${GPU_TARGETS} -DUSE_LLD=ON -DTI_WITH_LLVM=ON"

cd /taichi_dir/taichi/external/spdlog \
&& git apply /taichi_dir/taichi/spdlog_fmt.patch \
&& cd /taichi_dir/taichi \
&& ./build.py

Install the Taichi Lang .whl file:

python3 -m pip config set global.break-system-packages true \
&& python3 -m pip install /taichi_dir/taichi/dist/taichi*.whl

Test the Taichi Lang 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.0b2

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

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

Example output using laplace:

[Taichi] version 1.8.0, llvm 15.0.0, commit f7911653, linux, python python 3.12.3
[Taichi] Starting on arch=amdgpu
0
0
0
0
0
0
0
0
0
0

Run a Taichi Lang example#

A set of examples is available to help you get started. See run a Taichi Lang example for more details.