asc-devkit 算子原型注册 OpDef 接口全解析:输入输出、Shape 推导与多引擎配置
人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载导读OpDef是 CANN 昇腾 AI 处理器算子程序开发语言asc-devkit中承载算子原型定义的核心基类自定义算子工程通过继承OpDef的 C 类在构造函数中声明算子的输入、输出、属性注册 Shape/DataType 推导函数并关联 Tiling 实现与各计算引擎AICore、AICPU、HostCPU的配置信息。本文以 OpDef 接口文档 为骨架逐项解析其下 13 个注册接口的函数原型、参数语义、约束与调用示例并结合仓库中的 AddCustom 示例工程与《算子原型定义》指南给出源码级佐证。读完本文你将掌握如何用OpDef完整描述一个自定义算子的对外接口与支持范围并理解图模式推导函数、Fallback 下发与 format 匹配模式等进阶配置。一、OpDef 在算子原型注册体系中的定位算子原型描述了算子的输入输出、属性等信息以及算子在 AI 处理器上的相关实现信息并关联 Tiling 等函数是框架识别、编译和执行算子的依据。在 asc-devkit 的原型注册与管理体系中见 prototype_register_management.mdOpDef位于承上启下的位置注册入口算子原型通过自定义算子类承载该类继承自OpDef完成原型定义后调用 OP_ADD 接口传入算子类型自定义类名完成注册参数描述载体Input/Output返回 OpParamDefAttr返回 OpAttrDef二者可继续链式配置参数类型、dtype、format 等细节引擎配置载体AICore/AICPU/HostCPU分别返回 OpAICoreDef、OpAICPUDef、OpHostCPUDef用于配置各引擎的执行信息MC2则返回 OpMC2Def。仓库中的 AddCustom 样例工程给出了最小可运行的原型定义形态见 add_custom_host.cppnamespace ops { class AddCustom : public OpDef { public: explicit AddCustom(const char* name) : OpDef(name) { this-Input(x).ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); this-Input(y).ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); this-Output(z).ParamType(REQUIRED).DataType({ge::DT_FLOAT16}).Format({ge::FORMAT_ND}); this-SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType); this-AICore().SetTiling(optiling::TilingFunc).AddConfig(ascend910b); } }; OP_ADD(AddCustom); } // namespace ops从中可以看到 OpDef 的三类职责声明输入输出Input/OutputOpParamDef链、注册图模式推导函数SetInferShape/SetInferDataType、配置 AICore 执行信息SetTiling/AddConfig。下文逐项展开。二、输入、输出与属性注册Input / Output / Attr2.1 Input注册算子输入调用Input接口注册一个算子输入返回OpParamDef结构后续通过该结构继续配置该输入的参数类型、dtype、format 等信息。OpParamDef Input(const char *name)参数输入/输出说明name输入算子输入名称返回值算子参数定义OpParamDef实例。约束参数注册的顺序需要与算子核函数Kernel入口函数一致——因为注册顺序会直接影响 aclnn 接口参数顺序、tiling 侧取参顺序以及核函数Kernel入口参数顺序这是自定义算子开发中最容易踩坑的隐性约定之一。2.2 Output注册算子输出OpParamDef Output(const char *name)参数输入/输出说明name输入算子输出名称返回值算子参数定义OpParamDef实例。约束与Input相同注册顺序必须与核函数Kernel入口函数一致。2.3 输入输出的链式配置OpParamDef 关键方法Input/Output返回的OpParamDef支持链式调用详见 OpParamDef常用方法如下方法语义取值示例ParamType定义参数类型OPTIONAL可选、REQUIRED必选、DYNAMIC动态输入DataType定义参数数据类型与 Format 按位置一一对应ge::DT_FLOAT、ge::DT_FLOAT16、ge::DT_INT32等DataTypeList以集合方式定义支持的数据类型同上Format定义参数数据格式与 DataType 按位置一一对应ge::FORMAT_ND、ge::FORMAT_NCHW、ge::FORMAT_NHWC等FormatList以集合方式定义支持的 format同上Follow指定当前参数的 datatype、format、shape 跟随前面已定义的输入Follow(x1)要点同一个输入或输出支持多组 dtype/format 组合时DataType与Format列表需按位置一一对应。例如以下写法表达DT_FLOAT16 FORMAT_ND、DT_FLOAT FORMAT_ND、DT_INT32 FORMAT_ND三组组合this-Input(x) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_INT32}) .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});当输出完全跟随某个输入时优先使用Follow表达输出与输入一致的关系通常无需再编写 shape 推导函数Follow与SetInferShape不能混用详见 算子原型定义。2.4 Attr注册算子属性参数当某个参数不参与核函数Kernel侧计算时可以将其注册为算子属性参数如维度、阈值、缩放系数、开关等配置项OpAttrDef Attr(const char *name)参数输入/输出说明name输入算子属性名称返回值算子属性定义OpAttrDef具体接口见 OpAttrDef如AttrType(REQUIRED/OPTIONAL)、Bool/Float/Int等类型与默认值接口。属性声明后Host 侧 Tiling 函数可通过context-GetAttrs()获取属性值若核函数Kernel侧需要属性参与计算则通过 TilingData 传递。约束重要Attr属性名不能与以下 Python 关键字及内置变量名相同否则会导致未定义错误常见 Python 关键字and、as、assert、break、class、continue、def、del、elif、else、except、finally、for、from、global、if、import、in、is、lambda、not、or、pass、raise、return、try、while、with、yield、False、None、True、nonlocal、arg内置变量名__inputs__、__outputs__、__attrs__、options、bisheng、bisheng_path、tikcpp_path、impl_mode、custom_compile_options、custom_all_compile_options、soc_version、soc_short、custom_compile_options_soc、custom_all_compile_options_soc、origin_func_name、ascendc_src_dir_ex、ascendc_src_dir、ascendc_src_file、src、op_type、code_channel、op_info、compile_op、get_code_channel、result、isinstance、attr、get_current_build_config、_build_args、get_dtype_fmt_options、shutil、os、get_kernel_source、ascendc_api_version_header_path、ascendc_api_version_file、ascendc_api_version、re。从上述名单可以看出asc-devkit 的编译链路会生成 Python 侧的算子构建代码属性名会进入该生成代码的命名空间因此必须避开这些保留名称。三、图模式推导函数注册SetInferShape / SetInferShapeRange / SetInferDataType仅通过 aclnn 单算子接口调用算子时不需要注册输出 shape 和 dtype 推导函数但算子需要支持图模式构图、图编译、图执行时必须通过以下三个接口注册推导函数。三个接口均返回OpDef引用支持链式调用。3.1 SetInferShape注册 Shape 推导函数OpDef SetInferShape(gert::OpImplRegisterV2::InferShapeKernelFunc func)参数输入/输出说明func输入Shape 推导函数类型定义为using InferShapeKernelFunc UINT32 (*)(InferShapeContext *);AddCustom 示例中的推导函数实现见 add_custom_host.cpp展示了如何读取输入 shape 并设置输出 shapestatic graphStatus InferShape(gert::InferShapeContext* context) { const gert::Shape* x1_shape context-GetInputShape(0); gert::Shape* y_shape context-GetOutputShape(0); *y_shape *x1_shape; return GRAPH_SUCCESS; }3.2 SetInferShapeRange注册 ShapeRange 推导函数OpDef SetInferShapeRange(gert::OpImplRegisterV2::InferShapeRangeKernelFunc func)参数输入/输出说明func输入ShapeRange 推导函数类型定义为using InferShapeRangeKernelFunc UINT32 (*)(InferShapeRangeContext *);该接口用于动态 Shape 场景下推导输出 shape 的取值范围。3.3 SetInferDataType注册 DataType 推导函数OpDef SetInferDataType(gert::OpImplRegisterV2::InferDataTypeKernelFunc func)参数输入/输出说明func输入DataType 推导函数类型定义为using InferDataTypeKernelFunc UINT32 (*)(InferDataTypeContext *);AddCustom 中的实现通过输入 dtype 直接推导输出 dtypestatic graphStatus InferDataType(gert::InferDataTypeContext* context) { const auto inputDataType context-GetInputDataType(0); context-SetOutputDataType(0, inputDataType); return ge::GRAPH_SUCCESS; }组合注册图模式下三个推导函数可按需组合注册例如this-SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);。注意若注册了SetInferShape则需由推导函数完成所有输出的 shape 推导此时不宜再对单个输出使用Follow。四、计算引擎配置AICore / AICPU / HostCPU4.1 AICoreAICore 配置信息OpAICoreDef AICore(void)无参数。用于设置 AICore 的配置信息包括Tiling 处理回调函数、能力检查回调函数、基础的配置信息等。返回 OpAICoreDef其常用方法方法功能SetTiling注册 Tiling 函数原型固定为using TilingKernelFunc UINT32 (*)(TilingContext *);由框架调用并传入TilingContext可在其上获取输入输出的 Shape 指针等SetCheckSupport注册能力检查回调函数AddConfig注册算子支持的 AI 处理器型号以及 OpAICoreConfig 信息典型用法仓库示例this-AICore().SetTiling(optiling::TilingFunc).AddConfig(ascend910b);。其中SetTiling只负责建立原型与 Tiling 函数之间的调用关系TilingFunc内部读取输入 shape 等信息计算运行参数并写入 TilingData。AddConfig 默认配置AddConfig有两种重载——void AddConfig(const char *soc)与void AddConfig(const char *soc, OpAICoreConfig aicore_config)。soc填写规则参考算子工程目录下编译配置项文件CMakePresets.json中的ASCEND_COMPUTE_UNIT字段该字段在使用 msOpGen 创建工程时自动生成。不传入aicore_config时OpAICoreConfig部分参数采用默认值配置参数说明默认值DynamicCompileStaticFlag标识算子实现是否支持入图时的静态 Shape 编译trueDynamicFormatFlag标识是否根据 SetOpSelectFormat 设置的函数自动推导算子输入输出支持的 dtype 和 formattrueDynamicRankSupportFlag标识算子是否支持 dynamicRank动态维度trueDynamicShapeSupportFlag标识算子是否支持入图时的动态 Shape 场景trueNeedCheckSupportFlag标识是否在算子融合阶段调用算子参数校验函数进行 dtype 与 Shape 校验falsePrecisionReduceFlag用于 ATC 模型转换或网络调测时控制算子的精度模式trueOpAICoreConfig中未配置的输入输出继承OpDef中的定义重新配置的同名输入输出覆盖通用定义详见 AddConfig 与 算子原型定义。4.2 AICPUAI CPU 配置信息OpAICPUDef AICPU(void)无参数。用于设置 AI CPU 的配置信息包括算子输入输出的名称、支持的数据类型、数据格式等信息。返回 OpAICPUDef其子接口包括Engine引擎类型、FlagPartial、ComputeCost、FlagAsync、OpkernelLib、KernelSo、FunctionName、UserDefined、ExtendCfgInfo等用于描述 AICPU 算子的实现载体与调度属性。4.3 HostCPUHost CPU 配置信息OpHostCPUDef HostCPU(void)无参数。用于设置 Host CPU 的配置信息包括算子输入输出的name、支持的 dtype、format等信息。返回 OpHostCPUDef子接口与OpAICPUDef结构一致Engine、FlagPartial、ComputeCost、FlagAsync、OpKernelLib、KernelSo、FunctionName、UserDefined、ExtendCfgInfo。五、通算融合算子注册MC2OpMC2Def MC2(void)无参数。调用该接口将当前算子注册为通算融合算子返回 OpMC2Def 结构后续可通过该结构配置通信域名称HcclGroup、HcclServerType等见 OpMC2Def 简介。兼容性约束重要基于旧版本 CANN 包不支持 MC2 特性生成的自定义算子工程无法兼容 MC2 接口。在使用非当前版本 CANN 包生成的自定义算子工程时需特别注意。可通过查看自定义算子工程下cmake/util/ascendc_impl_build.py中有无_build_mc2_ctx字段来确认当前工程是否支持该特性如果未找到该字段则需要重新生成自定义算子工程以启用 MC2 特性。六、动态图 Fallback 下发EnableFallBackOpDef EnableFallBack(void)无参数。启用 fallback 配置后将自动生成一个 fallback 函数并注册给 GEGraph Engine。该 fallback 函数的核心逻辑是把 GE 的输入、输出及属性转换为 aclnn 单算子 API 所需的参数格式随后调用 aclnn 接口。动态图场景下GE 可直接调用 fallback 函数函数中调用了 aclnn 接口从而简化调度流程。约束算子需要注册并实现InferShape函数算子需要注册并实现InferDataType函数。调用示例完整展示 OpDef 构造与注册流程class AddCustom : public OpDef { public: AddCustom(const char* name) : OpDef(name) { this-Input(x).ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); this-Input(y).ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}); this-Output(z).ParamType(REQUIRED).DataType({ge::DT_FLOAT}).Format({ge::FORMAT_ND}) this-AICore().AddConfig(ascendxxx); this-SetInferShape(ge::InferShapeFunc); this-SetInferDataType(ge::InferDataTypeFunc); this-EnableFallBack(); } }; OP_ADD(AddCustom);七、原型注释与格式匹配Comment / FormatMatchMode7.1 Comment设置算子分组与原型注释OpDef Comment(CommentSection section, const char *comment)用于设置算子分组信息以及算子原型注释算子简述、算子约束等在自动生成算子原型头文件时同步生成算子原型注释。基于 OpDef 算子原型定义自定义算子工程可以实现自动生成图模式场景使用的算子原型头文件开发者可基于生成的算子原型进行构图、图编译、图执行等操作生成的注释还有助于自动生成算子原型的文档说明内置 CANN 算子使用较多开发者可按需使用。section通过CommentSection类指定注释功能支持以下取值取值含义CATEGORYcomment 内容为算子分组名称BRIEFcomment 内容为算子brief注释即算子功能简述CONSTRAINTScomment 内容为算子Attention Constraints注释即该算子的约束RESTRICTIONScomment 内容为算子Restrictions注释当前属于试验参数不推荐使用SEEcomment 内容为算子see注释可表示该算子的相关算子THIRDPARTYFWKCOMPATcomment 内容为算子参考的第三方算子约束使用CATEGORY设置算子分组名称时会对应生成同名的代码文件若文件名过长编译时可能超过 tar 打包的文件名称长度限制导致报错具体排查方法见 算子工程编译时出现文件名过长报错。调用示例展示各类注释的组合使用AddCustomComment(const char* name) : OpDef(name) { this-Comment(CommentSection::CATEGORY, catg); // 算子分组 this-Comment(CommentSection::BRIEF, Brief cmt) // BRIEF注释 .Comment(CommentSection::CONSTRAINTS, Constraints cmt1) // CONSTRAINTS注释 .Comment(CommentSection::CONSTRAINTS, Constraints cmt2); this-Comment(CommentSection::RESTRICTIONS, Restrictions cmt1) // RESTRICTIONS注释 .Comment(CommentSection::RESTRICTIONS, Restrictions cmt2) .Comment(CommentSection::THIRDPARTYFWKCOMPAT, Third-party framework compatibility cmt1) // THIRDPARTYFWKCOMPAT注释 .Comment(CommentSection::THIRDPARTYFWKCOMPAT, Third-party framework compatibility cmt2) .Comment(CommentSection::SEE, See cmt1) // SEE注释 .Comment(CommentSection::SEE, See cmt2); this-Input(x) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT, ge::DT_INT32}) .FormatList({ge::FORMAT_ND}); this-Input(y) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT, ge::DT_INT32}) .FormatList({ge::FORMAT_ND}); this-Output(z) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT, ge::DT_INT32}) .FormatList({ge::FORMAT_ND}); this-AICore() .SetTiling(optiling::TilingFunc); this-AICore().AddConfig(ascendxxx); }注意示例中Comment支持链式多次调用同一 section 可追加多条注释。7.2 FormatMatchMode设置 format 匹配模式OpDef FormatMatchMode(FormatCheckOption option)用于设置输入输出 tensor 的 format 匹配模式。option为FormatCheckOption枚举支持取值行为DEFAULT对 NCHW/NHWC/DHWCN/NCDHW/NCL 格式的输入输出转成 ND 格式进行处理STRICT对数据格式严格区分针对 NCHW/NHWC/DHWCN/NCDHW/NCL 格式aclnn 框架侧不做转换处理约束不调用该接口时默认将 NCHW/NHWC/DHWCN/NCDHW/NCL 格式的输入输出转成 ND 格式处理即等效于DEFAULT。调用示例下面示例中算子 AddCustom 输入x只支持 format 为 NCHW输入y只支持 format 为 NHWC此时必须配置FormatMatchMode(FormatCheckOption::STRICT)如果不配置aclnn 框架会将二者转成 ND 格式传给算子 tiling导致 format 语义丢失AddCustom(const char* name) : OpDef(name) { this-Input(x) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT}) .FormatList({ge::FORMAT_NCHW}); this-Input(y) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT}) .FormatList({ge::FORMAT_NHWC}); this-Output(z) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT}) .FormatList({ge::FORMAT_ND}); this-AICore().SetTiling(optiling::TilingFunc); this-AICore().AddConfig(ascendxxx); this-FormatMatchMode(FormatCheckOption::STRICT); }八、实战整合一个完整的 OpDef 原型定义综合以上接口一个覆盖输入输出、属性、图模式推导、AICore 配置与注释的完整算子原型定义如下融合了 EnableFallBack 与 FormatMatchMode 的示例写法namespace ge { static graphStatus InferShape(gert::InferShapeContext* context); // 输出 shape 推导 static graphStatus InferDataType(gert::InferDataTypeContext* context); // 输出 dtype 推导 } // namespace ge namespace ops { class AddCustom : public OpDef { public: explicit AddCustom(const char* name) : OpDef(name) { this-Comment(CommentSection::CATEGORY, MathOps); this-Comment(CommentSection::BRIEF, Element-wise add: z x y); this-Comment(CommentSection::CONSTRAINTS, x/y/z shapes must be identical); this-Input(x) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16}) .Format({ge::FORMAT_ND}); this-Input(y) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16}) .Format({ge::FORMAT_ND}); this-Output(z) .ParamType(REQUIRED) .DataType({ge::DT_FLOAT16}) .Format({ge::FORMAT_ND}); this-Attr(alpha).AttrType(OPTIONAL).Float(1.0); // 不参与 Kernel 计算的属性 this-SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType); this-AICore().SetTiling(optiling::TilingFunc).AddConfig(ascend910b); this-EnableFallBack(); // 动态图场景下发 aclnn this-FormatMatchMode(FormatCheckOption::STRICT); // 严格匹配 format } }; OP_ADD(AddCustom); } // namespace ops编写时的关键检查点Input/Output注册顺序与核函数Kernel入口参数顺序严格一致Attr属性名避开 Python 关键字与内置变量名见 2.4 节清单支持图模式时必须注册SetInferShape/SetInferDataType启用EnableFallBack时二者缺一不可AddConfig的soc取值与CMakePresets.json的ASCEND_COMPUTE_UNIT一致多组 dtype/format 组合时DataType与Format列表按位置一一对应多个输入输出 format 不一致时通过FormatMatchMode(STRICT)阻止 aclnn 框架的 ND 转换。九、延伸阅读原型注册与管理总览OpDef 与其他注册接口OP_ADD、REGISTER_OP_AICORE_CONFIG、OpParamDef、OpAttrDef、OpAICoreDef、OpAICPUDef、OpHostCPUDef、OpAICoreConfig、OpMC2Def的关系总览算子原型定义从算子功能设计到 OpDef 代码落地的完整流程含 OpType 命名规范大驼峰、不与内置算子重名与optype_collector工具使用说明AddCustom 示例工程仓库内可直接编译运行的最小原型定义参考算子工程编译时出现文件名过长报错CATEGORY分组名称导致打包报错的排查方法。赞分享人工智能深度学习算子库CANNAscend【免费下载链接】asc-devkit本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言原生支持C和C标准规范主要由类库和语言扩展层构成提供多层级API满足多维场景算子开发诉求。项目地址https://gitcode.com/cann/asc-devkit点击查看免费下载相关推荐asc-devkit 算子原型注册OpDef::FormatMatchMode 接口详解与 format 匹配模式配置实战asc devkit 算子原型注册OpDef::FormatMatchMode 接口详解与 format 匹配模式配置实战 FormatMatchMode 是人工智能深度学习算子库CANNAscendasc-devkit 算子原型注册OpAICPUDef::Engine 接口详解与 AICPU 算子引擎配置实战asc devkit 算子原型注册OpAICPUDef::Engine 接口详解与 AICPU 算子引擎配置实战 本篇技术指南围绕 asc devkitCA人工智能深度学习算子库CANNAscendasc-devkit 算子原型注册OpDef::HostCPU() 设置 Host CPU 算子配置信息详解asc devkit 算子原型注册OpDef::HostCPU 设置 Host CPU 算子配置信息详解 在 CANN 昇腾 AI 处理器的算子开发语言 as人工智能深度学习算子库CANNAscend上一篇qwen-code cua-driver macOS 后台驱动指南不抢占前台的 GUI 自动化契约与实践下一篇Linux 内核动态 DMA 映射完全指南从地址模型到驱动实践DMA API HOWTO创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考