> 本文只介绍通用算子开发流程，示例为普通逐元素运算，不包含任何未结束项目（博主的先导杯还没结束）的真实算子、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. 推荐的目录结构 一个最小扩展可以组织为： ```text 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 为例： ```text 输入x：DCU tensor、FP16、contiguous scale：标量float bias：标量float 输出：与x同shape、同dtype、同device ``` C++ binding 应在 launch 前检查： ```cpp 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骨架 以下代码只演示工程形式： ```cpp #include #include __global__ void affine_half_kernel( const __half* input, __half* output, float scale, float bias, long numel) { long index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; if (index >= numel) return; float value = __half2float(input[index]); output[index] = __float2half(value * scale + bias); } ``` launch函数： ```cpp 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((numel + threads - 1) / threads); hipLaunchKernelGGL( affine_half_kernel, dim3(blocks), dim3(threads), 0, 0, reinterpret_cast(input.data_ptr()), reinterpret_cast<__half*>(output.data_ptr()), scale, bias, numel); } ``` 生产实现还应使用当前PyTorch stream，并在需要时检查 launch error。上例省略了版本相关接口，避免把某一环境的写法误当作通用标准。 ## 5. C++ Binding只做契约和调度 ```cpp #include #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(scale), static_cast(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在线编译与加载 ```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`，可以使用新的扩展名或干净缓存目录验证。不要在生产程序里每次启动都删除缓存，这会强制重复编译。 提交源码时通常不应携带： ```text .torch_extensions/ *.so *.o __pycache__/ *.pyc ``` 预编译二进制还可能与评测机的PyTorch ABI和运行时不匹配。 ## 8. 三层验证缺一不可 ### 8.1 正确性 ```python 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 微基准 ```python 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 ``` reference和candidate必须使用相同输入、同步方式和重复次数。 ### 8.3 端到端 即使kernel微基准更快，也要放回完整应用检查： - 是否增加输入转换； - 是否破坏原有融合； - 是否增加常驻显存； - 是否引入同步； - 是否在真实shape上被调用； - 业务指标是否保持。 ## 9. 数值语义比公式等价更严格 设备优化中常见的误区是“公式相同，所以结果相同”。实际还要考虑： - FP16与FP32中间值； - reduction顺序； - FMA与分步乘加； - 原地写入时机； - 广播和舍入点； - NaN、Inf和有符号零。 如果原路径在两个操作之间写回FP16，融合kernel在FP32中连续计算就可能改变结果。是否允许这种差异，应由业务正确性门禁决定，而不是凭肉眼判断公式。 ## 10. 常见编译和运行错误 ### 10.1 找不到头文件 确认当前Python使用的PyTorch与编译环境一致，并检查扩展编译日志中的include目录。 ### 10.2 架构不匹配 确认 `rocminfo` 显示的目标架构与 `--offload-arch` 一致。不要把其他机器生成的code object直接复制过来。 ### 10.3 `undefined symbol` 通常是ABI或动态库版本不一致。使用 `ldd` 检查扩展实际链接到的库，并确认没有混用多个DTK环境。 ### 10.4 非法访存 优先检查：边界条件、元素数类型、tensor连续性、dtype、device、stream和输出buffer大小。先用极小shape复现，再扩大输入。 ## 11. 工程化发布清单 - loader明确检查HIP环境和设备架构； - binding检查device、dtype、shape和contiguous； - kernel覆盖非整除边界； - 使用当前框架stream； - 有reference正确性测试； - 有错误输入测试； - 有微基准和端到端验证； - 编译发生在正式计时边界之外； - 提交只包含源码，不包含缓存和二进制； - README写清环境、编译链和运行方法； - 不在公开文章中披露未结束项目的真实kernel和调参数据。 ## 参考与依赖说明 本文示例使用 PyTorch C++ Extension、HIP Runtime 和通用 C++17 接口，代码仅用于展示工程流程。实际项目应以对应版本的PyTorch、DTK和设备官方文档为准。文章不包含任何项目专有算子、模型、权重、shape或性能数据。