概要

AI分野においてNVIDIAのCUDAエコシステムが広く普及している一方で、AMDはROCm (Radeon Open Compute) を通じて対抗している。ROCmは、AMDのGPUアクセラレーションおよびヘテロジニアスコンピューティングを支援するための開発プラットフォームである。本稿では、AMDの統合型GPU（APU）環境上にROCmを構築し、PyTorchやOpenCLを用いた動作検証を行うとともに、その裏側で動作するカーネルドライバの仕組みを分析する。

動作環境

OSはUbuntu 20.04.6 LTS (x86_64) を使用し、GPUにはAMD Ryzen 5 5600Gに内蔵されたRadeon Graphics（VEGAアーキテクチャ、VULKAN/OpenCL対応）を用いた。

ROCmのインストール手順

以下のコマンドを実行し、必要なパッケージのダウンロードとユーザーグループの設定を行い、ROCmをインストールする。

sudo apt update && sudo apt dist-upgrade
sudo apt-get install wget gnupg2 
sudo usermod -a -G video $LOGNAME
sudo usermod -a -G render $LOGNAME
echo 'ADD_EXTRA_GROUPS=1' | sudo tee -a /etc/adduser.conf
echo 'EXTRA_GROUPS=video' | sudo tee -a /etc/adduser.conf
echo 'EXTRA_GROUPS=render' | sudo tee -a /etc/adduser.conf
sudo wget https://repo.radeon.com/amdgpu-install/22.10/ubuntu/focal/amdgpu-install_22.10.50100-1_all.deb
sudo apt-get install ./amdgpu-install_22.10.50100-1_all.deb 
sudo amdgpu-install --usecase=dkms
amdgpu-install -y --usecase=rocm
echo 'export PATH=$PATH:/opt/rocm/bin:/opt/rocm/profiler/bin:/opt/rocm/opencl/bin' | sudo tee -a /etc/profile.d/rocm.sh

インストール完了後、コンパイラやデバッガなどの開発ツールが /opt/rocm 配下に配置される。

動作確認

rocm-smi コマンドでGPUのステータスを確認できる。また、rocminfo を実行すると、HSA (Heterogeneous System Architecture) に基づき、CPUとGPUがそれぞれ独立した「Agent」として認識されていることが確認できる。OpenCLのサポート状況は clinfo コマンドから確認可能である。

DockerコンテナでのPyTorch実行検証

ROCm公式のPyTorchイメージを起動し、GPUアクセラレーションが有効かどうかを確認する。起動時にはAMDGPUのデバイスノードである /dev/kfd と /dev/dri をコンテナに渡す必要がある。

sudo docker run -it -v $HOME:/data --privileged --rm \
  --device=/dev/kfd --device=/dev/dri \
  --group-add video --name pytorch_rocm rocm/pytorch:latest

コンテナ内でPythonインタプリタを起動し、以下を実行する。

import torch
print(torch.cuda.is_available())

結果が True となれば、AMDGPU上でのPyTorch実行環境が正常に構築されている。

KFDドライバとCUDA互換の仕組み

Docker起動時に渡した /dev/kfd は、AMD Kernel Fusion Driver (KFD) のキャラクタデバイスノードである。KFDはユーザースペースからAMD GPUを操作するためのインターフェースを提供し、Linuxカーネル内でCPUとGPU間の通信を高速に仲介する役割を担う。

PyTorchなどのフレームワークは元々CUDA APIを呼び出す設計になっているが、ROCm環境ではこれをどのように処理しているのか。ROCmは、HIP (Heterogeneous-Computing Interface for Portability) というレイヤーを通じてCUDA APIとの互換性を実現している。具体的には、CUDAのAPIインターフェースを模倣し、内部でHIP+ROCmの実装にマッピングするラッパーライブラリ (libcuda*.so) を提供している。これにより、CUDA向けに書かれたコードがAMD GPU上で透過的に実行可能となる。また、ソースコードレベルでは HIPIFY ツールを用いることで、CUDAコードをHIPコードに自動変換できる。

OpenCLによるカーネル実行テスト

ROCm上でのOpenCLの動作を検証するため、ベクトルの加算を行うカーネルを実装した。以下にホスト側のコードを示す。

#include <stdio.h>
#include <stdlib.h>
#include <CL/cl.h>

#define DATA_SIZE 256

const char* kernel_source = 
"__kernel void vec_add(__global const float* in_a, __global const float* in_b, __global float* out_c) {\n"
"    int idx = get_global_id(0);\n"
"    out_c[idx] = in_a[idx] + in_b[idx];\n"
"}\n";

int main() {
    cl_platform_id platform;
    cl_device_id device;
    cl_context context;
    cl_command_queue queue;
    cl_program program;
    cl_kernel kernel;
    cl_int err;

    // プラットフォームとデバイスの取得
    err = clGetPlatformIDs(1, &platform, NULL);
    err |= clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, &device, NULL);
    if (err != CL_SUCCESS) { return -1; }

    // コンテキストとコマンドキューの作成
    context = clCreateContext(NULL, 1, &device, NULL, NULL, &err);
    queue = clCreateCommandQueue(context, device, 0, &err);

    // プログラムのビルド
    program = clCreateProgramWithSource(context, 1, &kernel_source, NULL, &err);
    err = clBuildProgram(program, 1, &device, NULL, NULL, NULL);
    if (err != CL_SUCCESS) {
        char build_log[4096];
        clGetProgramBuildInfo(program, device, CL_PROGRAM_BUILD_LOG, sizeof(build_log), build_log, NULL);
        printf("Build Error:\n%s\n", build_log);
        return -1;
    }
    kernel = clCreateKernel(program, "vec_add", &err);

    // ホストデータの準備
    float host_a[DATA_SIZE], host_b[DATA_SIZE], host_c[DATA_SIZE];
    for (int i = 0; i < DATA_SIZE; i++) {
        host_a[i] = (float)i;
        host_b[i] = (float)i * 2.0f;
    }

    // バッファオブジェクトの作成
    cl_mem buf_a = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float) * DATA_SIZE, host_a, &err);
    cl_mem buf_b = clCreateBuffer(context, CL_MEM_READ_ONLY | CL_MEM_COPY_HOST_PTR, sizeof(float) * DATA_SIZE, host_b, &err);
    cl_mem buf_c = clCreateBuffer(context, CL_MEM_WRITE_ONLY, sizeof(float) * DATA_SIZE, NULL, &err);

    // カーネル引数の設定と実行
    err = clSetKernelArg(kernel, 0, sizeof(cl_mem), &buf_a);
    err |= clSetKernelArg(kernel, 1, sizeof(cl_mem), &buf_b);
    err |= clSetKernelArg(kernel, 2, sizeof(cl_mem), &buf_c);

    size_t global_size = DATA_SIZE;
    err = clEnqueueNDRangeKernel(queue, kernel, 1, NULL, &global_size, NULL, 0, NULL, NULL);

    // 結果の読み戻し
    clEnqueueReadBuffer(queue, buf_c, CL_TRUE, 0, sizeof(float) * DATA_SIZE, host_c, 0, NULL, NULL);

    // リソースの解放
    clReleaseMemObject(buf_a);
    clReleaseMemObject(buf_b);
    clReleaseMemObject(buf_c);
    clReleaseKernel(kernel);
    clReleaseProgram(program);
    clReleaseCommandQueue(queue);
    clReleaseContext(context);

    return 0;
}

このコードは、ROCmに含まれるClangコンパイラを用いてビルドする。

/opt/rocm/llvm/bin/clang -I/opt/rocm/opencl/include -L/opt/rocm/opencl/lib -lOpenCL ocl_vec_add.c -o ocl_vec_add

実行すると、内部でカーネルのビルドが行われ、GPU上でベクトルの加算が実行される。straceコマンドでシステムコールをトレースすると、/dev/kfd に対するioctl呼び出し（コマンドキューやメモリの割り当てなど）が行われていることを確認できる。

OpenMPオフロードの検証

ROCmはOpenMPを用いたGPUオフロードにも対応している。サンプルコード（veccopy.c）を用いて、OpenMPターゲット領域をAMD GPUにオフロードして実行できるか検証した。

/opt/rocm/llvm/bin/clang -O3 -fopenmp -fopenmp-targets=amdgcn-amd-amdhsa \
  -Xopenmp-target=amdgcn-amd-amdhsa -march=gfx908 veccopy.c -o veccopy

straceで追跡した結果、実行時にも /dev/kfd がオープンされ、HSAランタイムを介したメモリ割り当てやカーネル実行のためのioctlシステムコールが発行されていることが確認できた。これにより、OpenMPターゲットオフロード機能も、裏側ではKFDドライバを通じてAMD GPUと通信していることがわかる。