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

本文只介绍通用算子开发流程示例为普通逐元素运算不包含任何未结束项目博主的先导杯还没结束的真实算子、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.cppTensor契约检查、pybind导出、调用HIP入口op.hC声明op_kernel.hipkernel和launch配置不要把所有逻辑写在一个.hip文件中。输入契约和设备实现分开后错误更容易定位。3. 先定义Tensor契约以逐元素 affine 为例输入xDCU tensor、FP16、contiguous scale标量float bias标量float 输出与x同shape、同dtype、同deviceC 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骨架以下代码只演示工程形式#includehip/hip_runtime.h#includehip/hip_fp16.h__global__voidaffine_half_kernel(const__half*input,__half*output,floatscale,floatbias,longnumel){longindexstatic_castlong(blockIdx.x)*blockDim.xthreadIdx.x;if(indexnumel)return;floatvalue__half2float(input[index]);output[index]__float2half(value*scalebias);}launch函数voidlaunch_affine_half(constat::Tensorinput,at::Tensoroutput,floatscale,floatbias){constexprintthreads256;constautonumelinput.numel();constintblocksstatic_castint((numelthreads-1)/threads);hipLaunchKernelGGL(affine_half_kernel,dim3(blocks),dim3(threads),0,0,reinterpret_castconst__half*(input.data_ptrat::Half()),reinterpret_cast__half*(output.data_ptrat::Half()),scale,bias,numel);}生产实现还应使用当前PyTorch stream并在需要时检查 launch error。上例省略了版本相关接口避免把某一环境的写法误当作通用标准。5. C Binding只做契约和调度#includetorch/extension.h#includeop.htorch::Tensoraffine_half(consttorch::Tensorinput,doublescale,doublebias){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);autooutputtorch::empty_like(input);launch_affine_half(input,output,static_castfloat(scale),static_castfloat(bias));returnoutput;}PYBIND11_MODULE(TORCH_EXTENSION_NAME,module){module.def(affine_half,affine_half,FP16 affine example);}binding层不应偷偷执行dtype转换或.contiguous()。隐式转换会产生额外tensor使微基准看起来正常、端到端却变慢。调用方需要转换时应在Python层显式完成并纳入测量。6. Python在线编译与加载frompathlibimportPathimporttorchfromtorch.utils.cpp_extensionimportloaddefload_extension():iftorch.version.hipisNone:raiseRuntimeError(This extension requires a HIP PyTorch build)rootPath(__file__).resolve().parentreturnload(namemy_dcu_op_ext,sources[str(root/binding.cpp),str(root/op_kernel.hip),],extra_cflags[-O3,-stdc17],extra_cuda_cflags[-O3,--offload-archgfx936,],verboseTrue,)虽然接口参数名是extra_cuda_cflagsHIP版PyTorch通常会将其传递给HIP编译链。是否支持、具体参数如何解释应以当前PyTorch/DTK版本为准。在线编译适用于不能预装二进制扩展、但允许在目标机编译源码的环境。它能确保目标架构一致代价是首次启动时间较长。7. 编译缓存如何管理PyTorch extension会缓存编译产物。缓存键至少受以下因素影响扩展名称源码内容编译参数Python与PyTorch版本编译器和运行时目标架构。开发时如果修改源码却仍加载旧.so可以使用新的扩展名或干净缓存目录验证。不要在生产程序里每次启动都删除缓存这会强制重复编译。提交源码时通常不应携带.torch_extensions/ *.so *.o __pycache__/ *.pyc预编译二进制还可能与评测机的PyTorch ABI和运行时不匹配。8. 三层验证缺一不可8.1 正确性referenceinput*scalebias actualextension.affine_half(input,scale,bias)torch.testing.assert_close(actual,reference,rtol1e-3,atol1e-3,)测试应覆盖空tensor和小tensor元素数不是block整数倍多维连续tensor极值、零和负数错误dtype、CPU输入、非连续输入多次调用和不同stream。8.2 微基准defbench(fn,iterations100):for_inrange(10):fn()torch.cuda.synchronize()starttime.perf_counter()for_inrange(iterations):fn()torch.cuda.synchronize()return(time.perf_counter()-start)/iterationsreference和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.3undefined symbol通常是ABI或动态库版本不一致。使用ldd检查扩展实际链接到的库并确认没有混用多个DTK环境。10.4 非法访存优先检查边界条件、元素数类型、tensor连续性、dtype、device、stream和输出buffer大小。先用极小shape复现再扩大输入。11. 工程化发布清单loader明确检查HIP环境和设备架构binding检查device、dtype、shape和contiguouskernel覆盖非整除边界使用当前框架stream有reference正确性测试有错误输入测试有微基准和端到端验证编译发生在正式计时边界之外提交只包含源码不包含缓存和二进制README写清环境、编译链和运行方法不在公开文章中披露未结束项目的真实kernel和调参数据。参考与依赖说明本文示例使用 PyTorch C Extension、HIP Runtime 和通用 C17 接口代码仅用于展示工程流程。实际项目应以对应版本的PyTorch、DTK和设备官方文档为准。文章不包含任何项目专有算子、模型、权重、shape或性能数据。