DCU 自定义算子工程化：从 HIP 源码、PyTorch Binding 到在线加载

weixin_47594597

591人浏览 · 2026-07-16 23:01:59

weixin_47594597 · 2026-07-16 23:01:59 发布

本文只介绍通用算子开发流程，示例为普通逐元素运算，不包含任何未结束项目（博主的先导杯还没结束）的真实算子、shape、调度参数或性能结果。

摘要

写出一个能运行的 HIP kernel 并不难，难的是把它变成可维护的 PyTorch 扩展：输入检查要完整、目标架构要明确、编译缓存要可控、错误要能定位，还要有正确性和端到端性能门禁。

本文以一个通用 y = x × scale + bias 逐元素算子为例，说明 DCU 自定义算子的工程结构、C++ binding、HIP kernel、在线编译、缓存管理和测试方法。示例只演示基础流程，不对应任何实际项目的优化实现。

关键词：DCU、HIP、PyTorch Extension、C++ Binding、在线编译、算子开发

1. 什么时候值得写自定义算子

自定义kernel有长期成本：需要维护dtype、shape、设备、连续性、编译器版本和数值行为。满足以下条件时才值得考虑：

Profile确认它在端到端关键路径；
shape或dtype稳定，专用实现能利用先验；
多个相邻算子反复读写同一个大tensor；
现有后端没有自动融合；
替换后的正确性边界可清楚定义；
微基准收益足以覆盖launch、转换和编译成本。

如果只是单个通用 GEMM 或卷积，优先使用框架和数学库。自定义算子更适合稳定的逐元素融合、布局转换和小型专用reduction。

2. 推荐的目录结构

一个最小扩展可以组织为：

my_dcu_op/
├── __init__.py
├── loader.py
├── binding.cpp
├── op.h
└── op_kernel.hip

职责分离：

文件	职责
`loader.py`	架构检查、编译参数、缓存目录、Python接口
`binding.cpp`	Tensor契约检查、pybind导出、调用HIP入口
`op.h`	C++声明
`op_kernel.hip`	kernel和launch配置

不要把所有逻辑写在一个 .hip 文件中。输入契约和设备实现分开后，错误更容易定位。

3. 先定义Tensor契约

以逐元素 affine 为例：

输入x：DCU tensor、FP16、contiguous
scale：标量float
bias：标量float
输出：与x同shape、同dtype、同device

C++ binding 应在 launch 前检查：

TORCH_CHECK(x.is_cuda(), "x must be on accelerator");
TORCH_CHECK(x.scalar_type() == at::kHalf, "x must be FP16");
TORCH_CHECK(x.is_contiguous(), "x must be contiguous");

如果支持多输入，还要检查device一致、元素数匹配和禁止非法alias。把错误挡在C++入口，比让kernel非法访存安全得多。

4. 一个通用HIP kernel骨架

以下代码只演示工程形式：

#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>

__global__ void affine_half_kernel(
    const __half* input,
    __half* output,
    float scale,
    float bias,
    long numel) {
  long index = static_cast<long>(blockIdx.x) * blockDim.x + threadIdx.x;
  if (index >= numel) return;

  float value = __half2float(input[index]);
  output[index] = __float2half(value * scale + bias);
}

launch函数：

void launch_affine_half(
    const at::Tensor& input,
    at::Tensor& output,
    float scale,
    float bias) {
  constexpr int threads = 256;
  const auto numel = input.numel();
  const int blocks = static_cast<int>((numel + threads - 1) / threads);

  hipLaunchKernelGGL(
      affine_half_kernel,
      dim3(blocks), dim3(threads), 0, 0,
      reinterpret_cast<const __half*>(input.data_ptr<at::Half>()),
      reinterpret_cast<__half*>(output.data_ptr<at::Half>()),
      scale, bias, numel);
}

生产实现还应使用当前PyTorch stream，并在需要时检查 launch error。上例省略了版本相关接口，避免把某一环境的写法误当作通用标准。

5. C++ Binding只做契约和调度

#include <torch/extension.h>
#include "op.h"

torch::Tensor affine_half(
    const torch::Tensor& input,
    double scale,
    double bias) {
  TORCH_CHECK(input.is_cuda(), "input must be on accelerator");
  TORCH_CHECK(input.scalar_type() == at::kHalf, "input must be FP16");
  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");

  auto output = torch::empty_like(input);
  launch_affine_half(input, output,
                     static_cast<float>(scale),
                     static_cast<float>(bias));
  return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
  module.def("affine_half", &affine_half, "FP16 affine example");
}

binding层不应偷偷执行dtype转换或 .contiguous()。隐式转换会产生额外tensor，使微基准看起来正常、端到端却变慢。调用方需要转换时，应在Python层显式完成并纳入测量。

6. Python在线编译与加载

from pathlib import Path
import torch
from torch.utils.cpp_extension import load


def load_extension():
    if torch.version.hip is None:
        raise RuntimeError("This extension requires a HIP PyTorch build")

    root = Path(__file__).resolve().parent
    return load(
        name="my_dcu_op_ext",
        sources=[
            str(root / "binding.cpp"),
            str(root / "op_kernel.hip"),
        ],
        extra_cflags=["-O3", "-std=c++17"],
        extra_cuda_cflags=[
            "-O3",
            "--offload-arch=gfx936",
        ],
        verbose=True,
    )

虽然接口参数名是 extra_cuda_cflags，HIP版PyTorch通常会将其传递给HIP编译链。是否支持、具体参数如何解释，应以当前PyTorch/DTK版本为准。

在线编译适用于不能预装二进制扩展、但允许在目标机编译源码的环境。它能确保目标架构一致，代价是首次启动时间较长。

7. 编译缓存如何管理

PyTorch extension会缓存编译产物。缓存键至少受以下因素影响：

扩展名称；
源码内容；
编译参数；
Python与PyTorch版本；
编译器和运行时；
目标架构。

开发时如果修改源码却仍加载旧 .so，可以使用新的扩展名或干净缓存目录验证。不要在生产程序里每次启动都删除缓存，这会强制重复编译。

提交源码时通常不应携带：

.torch_extensions/
*.so
*.o
__pycache__/
*.pyc

预编译二进制还可能与评测机的PyTorch ABI和运行时不匹配。

8. 三层验证缺一不可

8.1 正确性

reference = input * scale + bias
actual = extension.affine_half(input, scale, bias)

torch.testing.assert_close(
    actual,
    reference,
    rtol=1e-3,
    atol=1e-3,
)

测试应覆盖：

空tensor和小tensor；
元素数不是block整数倍；
多维连续tensor；
极值、零和负数；
错误dtype、CPU输入、非连续输入；
多次调用和不同stream。

8.2 微基准

def bench(fn, iterations=100):
    for _ in range(10):
        fn()
    torch.cuda.synchronize()

    start = time.perf_counter()
    for _ in range(iterations):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - start) / iterations