把 CUDA 算子搬到 ROCm，最常被问的一句话是“值不值？”
我们先把结论放这儿：同一套 PyTorch 扩展，在 MI210 上打开 hipGraph 后，端到端吞吐比 A100 高 18%，显存占用低 7%。
迁移成本呢？5 个文件、不到 120 行 diff，CI 编译时间从 11 min 降到 8 min——因为 hipcc 并行度比 nvcc 更激进。
下面把全过程拆成 5 步，每一步都给出可复制的代码或命令，照着敲就能跑通。

2. 语法差异速查表：30 秒定位改写点

CUDA 写法	ROCm 写法	备注
__global__ void foo()	__global__ void foo()	关键字兼容，直接留
threadIdx.x	hipThreadIdx_x	宏统一小写，加 hip 前缀
__syncthreads()	__syncthreads()	同名，无需改
cudaMalloc	hipMalloc	前缀替换
cudaMemcpyDeviceToHost	hipMemcpyDeviceToHost	前缀替换
cublasHandle_t	rocblas_handle	类型名不同，注意大小写
nvcc –arch=sm_80	hipcc –offload-arch=gfx908	架构代号换 gfx 系列

把上表贴到 IDE 侧边栏，全局搜索-替换 5 分钟就能扫完。

3. 第 1 步：宏开关隔离平台差异

// common.hpp
#pragma once
#if defined(__HIPCC__)
  #define DEVICE __device__
  #define HOST   __host__
  #define THREAD_IDX_X (hipThreadIdx_x)
#else
  #define DEVICE __device__
  #define HOST   __host__
  #define THREAD_IDX_X (threadIdx.x)
#endif

以后所有 .cu / .hip 文件只要 #include "common.hpp"，就能一份代码双平台编译，告别 #ifdef 漫天飞。

4. 第 2 步：把 nvcc 换成 hipcc，CMakeLists 这样写

cmake_minimum_required(VERSION 3.18)
project(my_extension LANGUAGES CXX)

find_package(Torch REQUIRED)

# 关键：自动检测 GPU 架构
if(DEFINED ENV{ROCM_PATH})
  set(CMAKE_HIP_COMPILER "${ROCM_PATH}/bin/hipcc")
  set(HIP_HIP_PLATFORM hcc)
  # gfx908=MI210, gfx90a=MI250, 按需追加
  list(APPEND CMAKE_HIP_ARCHITECTURES gfx908)
  enable_language(HIP)
  set_source_files_properties(
      my_kernel.hip
      PROPERTIES LANGUAGE HIP)
else()
  list(APPEND CMAKE_CUDA_ARCHITECTURES 80)
  enable_language(CUDA)
  set_source_files_properties(
      my_kernel.cu
      PROPERTIES LANGUAGE CUDA)
endif()

add_library(my_ops SHARED
    my_kernel.${CMAKE_HIP_COMPILER_EXTENSION}
    my_ops.cpp
)
target_link_libraries(my_ops "${TORCH_LIBRARIES}")
set_property(TARGET my_ops PROPERTY CXX_STANDARD 17)

把文件后缀改成 .hip 只是约定，hipcc 照样能编译 .cu，但后缀区分方便 CI 做缓存哈希。

5. 第 3 步：本地编译 & 打包 .so

# 假设已安装 rocm/dev 镜像
export ROCM_PATH=/opt/rocm
pip install torch torchvision --index-url https://download.pytorch.org/whl/rocm5.6
mkdir build && cd build
cmake .. -DCMAKE_PREFIX_PATH=$(python -c "import torch; print(torch.utils.cmake_prefix_path)")
make -j$(nproc)

编译完得到 libmy_ops.so，python -c "import torch; torch.ops.load_library('./libmy_ops.so')" 无报错即成功。

6. 第 4 步：pytest 端到端验证

# test_ops.py
import torch, time, pytest

torch.ops.load_library('build/libmy_ops.so')

@pytest.mark.parametrize('N', [128, 256, 512])
def test_custom_gemm(benchmark, N):
    a = torch.randn(N, N, device='cuda', dtype=torch.half)
    b = torch.randn(N, N, device='cuda', dtype=torch.half)
    # warmup
    torch.ops.my_ops.custom_gemm(a, b)
    torch.cuda.synchronize()

    def run():
        return torch.ops.my_ops.custom_gemm(a, b)

    benchmark(run)

跑 pytest test_ops.py --benchmark-json=result.json 就能拿到 latency 分布，后续画曲线直接喂给 matplotlib。

7. 第 5 步：打开 ROCm 独占的 hipGraph，白捡 18%

CUDA 有 cudaGraph，ROCm 对应的是 hipGraph，接口几乎 1:1。
在 PyTorch C++ 扩展里，只要在第一次执行时把 kernel 序列捕获下来，后面直接 replay：

static hipGraph_t g = nullptr;
static hipGraphExec_t gExec = nullptr;

void launch_custom_gemm(...) {
    if (g == nullptr) {
        hipStreamCaptureBegin(&g);
        // 这里调你的 kernel
        hipStreamCaptureEnd(&g);
        hipGraphInstantiate(&gExec, g, nullptr, nullptr, 0);
    }
    hipGraphLaunch(gExec, stream);
}

图 2：关闭 vs 打开 hipGraph 的吞吐对比（batch=64，序列长度 512）

打开后，MI210 在 200 step 的均值吞吐从 2.34 k seq/s 提到 2.76 k seq/s，A100 同期是 2.34 k，差距就是这么赤裸裸。

8. 踩坑实录：3 个暗坑与快速自检清单

1. 符号未找到：hipcc 默认把 device 函数 inline 掉，加 -fgpu-rdc 让 relocatable device code 生成，否则 __device__ 函数在 .so 里找不到定义。
2. 精度漂移：rocBLAS 的 rocblas_gemm_ex 默认用 f32_r 累加，要显式指定 compute_type = rocblas_compute_f16 才能和 cublas 对齐。
3. 显存泄漏：hipGraph 捕获时会把 malloc 一并录进来，replay 阶段会反复 alloc。解决方式：在捕获前预分配 workspace，把指针作为 kernel 参数传进去，别让内存节点入图。

自检清单（CI 加一行就能过）：

hipcc --version | grep -q "HIP 5" || exit 1
rocminfo | grep -q "gfx908" || exit 1
python -c "import torch; print(torch.version.hip)" || exit 1

9. 云端一键镜像：魔搭创空间 3 分钟拉起

懒得本地装驱动？ModelScope 社区已提供 rocm5.6_pytorch2.1 镜像，启动命令：

docker run -it --device=/dev/kfd \
  -v $PWD:/workspace \
  registry.cn-hangzhou.aliyuncs.com/modelscope/rocm5.6_pytorch2.1:latest \
  bash -c "cd /workspace && python setup.py install && pytest"

图 3：魔搭容器内 htop 看到的 64 核 Zen3 + 64 GB HBM

容器里 /opt/rocm/bin/hipcc 已配好，CMake 能自动搜到路径，CI 直接复用这条命令，零本地依赖。

10. 小结：把“能跑”变“跑得爽”

迁移不是目的，榨干硬件才是。
5 步下来，我们拿到了：

• 一份双平台宏隔离的源码

• 一条 cmake 命令打出 .so

• 一套 pytest 基准

• 一张 hipGraph 开关就能提速 18%

  图 4：最终吞吐曲线汇总（MI210 开 hipGraph 后全程压在 A100 上方）

下次再有新算法，先在 MI210 上跑通，再决定要不要回 CUDA——毕竟，省下来的卡时费是真金白银。

---

立即加入AI开发者计划，免费领取 100 小时算力
添加微信小助手 csdn-01 还可额外领取「Openclaw 实战秘籍」