CANN-opbase-昇腾NPU算子开发的基础设施为什么这么重要

CANN-opbase-昇腾NPU算子开发的基础设施为什么这么重要所有 CANN AOL 算子仓库的底层都依赖 opbase。它不提供任何算子实现提供的是算子注册、编译、调度的基础设施。如果你要写自定义 Ascend C 算子opbase 是绕不过去的第一步。opbase 提供了什么组件功能算子开发者的感知算子注册框架定义算子的名称、输入输出、属性IMPL_OP(MyOp).Inputs(...)Tiling 框架管理分块参数的传递和存储TilingFunc回调编译工具链把 Ascend C 编译成 NPU 可执行文件op_build.sh运行时接口算子与 CANN 运行时的交互GetTensorAddr,GetTilingData通用数据结构Tensor 描述、Shape、FormatTensorDesc,GeShape这些组件都是 header-only 或预编译库算子代码#include就能用。算子注册流程每个算子必须在 opbase 的注册框架里声明CANN 运行时才能找到它// op_host/my_op/my_op.cpp#includeop_host.hnamespacege{IMPL_OP(MyOp).Inputs({x,y})// 输入 tensor 名称.Outputs({z})// 输出 tensor 名称.Attr(alpha,GeAttr::FLOAT)// 算子属性.Tiling(MyOpTiling);// Tiling 回调函数}注册信息告诉 CANN 运行时这个算子有几个输入/输出属性是什么类型Tiling 策略由哪个函数计算Tiling 框架Tiling 是昇腾NPU算子开发最核心的概念。opbase 的 Tiling 框架把分块参数的生成和传递标准化了// Tiling 函数根据输入 shape 计算分块参数ge::graphStatusMyOpTiling(constge::Operatorop,TilingContext*ctx){autox_shapectx-GetInputShape(0);// 获取输入 shapeint64_ttotal_elemsx_shape.GetShapeSize();// 计算分块参数int64_ttile_m128;int64_ttile_n128;int64_ttiles(total_elemstile_m*tile_n-1)/(tile_m*tile_n);// 写入 Tiling 数据ctx-SetTilingKey(1);// Tiling 策略编号ctx-SetTilingData(tile_m,tile_m);ctx-SetTilingData(tile_n,tile_n);ctx-SetTilingData(tiles,tiles);returnge::GRAPH_SUCCESS;}Tiling 函数跑在 CPU 上不在 NPU 上在算子执行前调用。计算出的分块参数通过 CANN 运行时传递给 NPU 上的 kernel。编译工具链opbase 提供了算子编译的标准流程# 编译单个算子cdop_kernel/my_op ascendc--chipAscend910B2--outputbuild my_op_kernel.cpp# 编译整个算子包bashbuild.sh编译产物是.so文件放在 CANN 的算子目录下。torch_npu 启动时扫描这些.so并注册。--chip参数指定目标芯片型号。不同型号的 Cube 单元大小不同910B2 的 Cube 是 16×16910B1 是 16×16 但频率不同编译器会针对具体型号优化指令调度。运行时接口NPU 上的 kernel 通过 opbase 的运行时接口获取输入数据和 Tiling 参数// op_kernel/my_op/my_op_kernel.cpp#includekernel_operator.hclassMyOpKernel{public:__aicore__inlinevoidInit(GM_ADDR x,GM_ADDR y,GM_ADDR z,TilingData*tiling){// 从 opbase 运行时获取 tensor 地址x_gm.SetGlobalBuffer((__gm__ half*)x);y_gm.SetGlobalBuffer((__gm__ half*)y);z_gm.SetGlobalBuffer((__gm__ half*)z);// 从 Tiling 数据获取分块参数tile_m_tiling-tile_m;tile_n_tiling-tile_n;}private:int64_ttile_m_,tile_n_;GlobalTensorhalfx_gm,y_gm,z_gm;};GM_ADDR是 opbase 定义的全局内存地址类型。TilingData是 opbase 的 Tiling 数据结构由 Tiling 函数填充、运行时传递。opbase 是 CANN 算子生态的地基。理解了注册框架、Tiling 框架、编译工具链和运行时接口这四个组件写自定义算子就不懵了。仓库在这里https://atomgit.com/cann/opbase