Using CuPy on AMD GPU (experimental)

CuPy has an experimental support for AMD GPU (ROCm).

Requirements

The following ROCm libraries are required:

$ sudo apt install hipblas hipsparse rocsparse rocrand rocthrust rocsolver rocfft hipcub rocprim rccl

Building CuPy for ROCm

Currently, you need to build CuPy from source to run on AMD GPU.

$ export HCC_AMDGPU_TARGET=gfx900  # This value should be changed based on your GPU
$ export __HIP_PLATFORM_HCC__
$ export CUPY_INSTALL_USE_HIP=1
$ pip install cupy

Note that HCC_AMDGPU_TARGET must be set to the ISA name supported by your GPU. Run rocminfo and use the value displayed in Name: line (e.g., gfx900).

You may also need to set ROCM_HOME (e.g., ROCM_HOME=/opt/rocm).