☰
asc-devkit 算子原型注册:SetOpSelectFormat 动态推导 dtype 与 format 详解
2026/10/16 6:24:25 网站建设 项目流程
  • 人工智能
  • 深度学习
  • 算子库
  • CANN
  • Ascend

【免费下载链接】asc-devkit

本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言,原生支持C和C++标准规范,主要由类库和语言扩展层构成,提供多层级API,满足多维场景算子开发诉求。

项目地址:https://gitcode.com/cann/asc-devkit
点击查看免费下载

导读

在昇腾 AI 处理器算子开发中,算子原型(OpDef)通常需要在注册时静态声明输入输出所支持的数据类型(dtype)与数据格式(format)。但在算子输入 shape、属性随运行时变化的场景下,静态声明往往不够灵活。asc-devkit 提供的OpAICoreDef::SetOpSelectFormat接口允许开发者注册一个自定义推导回调函数,在算子融合阶段根据运行时算子描述动态推导输入输出支持的 dtype 与 format,配合DynamicFormatFlag(true)配置即可让算子原型彻底摆脱固定的 dtype/format 声明。读完本文,你将掌握SetOpSelectFormat的函数原型、回调函数签名与结果 JSON 格式、与DynamicFormatFlag的配合方式,以及完整可复制的注册示例。

接口定位:原型注册与管理体系中的一员

SetOpSelectFormat是OpAICoreDef(AI Core 算子定义)提供的注册接口之一。在 asc-devkit 的原型注册与管理体系中,算子原型通过继承OpDef的 C++ 类承载,类内通过this->AICore()返回的OpAICoreDef对象链式注册 Tiling、参数校验、格式推导等回调,相关接口还包括 SetTiling、SetCheckSupport、SetOpSupportInfo 等,完整列表见 OpAICoreDef 与 原型注册与管理总览。

功能说明

如果开发者需要自行推导算子输入输出所支持的数据类型与格式,可以实现推导回调函数,并通过SetOpSelectFormat接口进行注册。同时,需要将 DynamicFormatFlag 配置为true,则算子融合时框架会自动调用该推导函数进行数据类型与格式的设置,算子原型注册时无需再配置输入输出支持的数据类型与格式。

需要特别注意以下优先级规则:如果算子原型中已经通过DataType/Format注册过数据类型与格式,则以算子原型注册的数据类型与格式为准,即使注册了推导函数也不会执行。也就是说,SetOpSelectFormat的推导结果只有在原型未静态声明 dtype/format 时才会生效。

函数原型

OpAICoreDef &SetOpSelectFormat(optiling::OP_CHECK_FUNC func)

该接口是OpAICoreDef的成员函数,返回OpAICoreDef引用,支持链式调用。参数func即推导回调函数,其类型OP_CHECK_FUNC定义如下:

using OP_CHECK_FUNC = ge::graphStatus (*)(const ge::Operator &op, ge::AscendString &result);

该函数类型与 SetCheckSupport 等接口共用,均以ge::Operator作为算子描述入参,以ge::AscendString作为出参载体。

参数说明

参数输入/输出说明
func输入推导算子输入输出所支持数据类型与格式的函数,类型为optiling::OP_CHECK_FUNC

推导函数的入参是算子的描述,包括算子的输入、输出、属性等信息;出参为包含了当前算子输入输出支持的数据类型与格式列表的字符串。字符串采用 JSON 格式,样例如下:

{ "input0": {"name": "x","dtype": "float16,float32,int32","format": "ND,ND,ND"}, "input1": {"name": "y","dtype": "float16,float32,int32","format": "ND,ND,ND"}, "output0": {"name": "z","dtype": "float16,float32,int32","format": "ND,ND,ND"} }

对 JSON 字段逐一说明:

  • 顶层键:input0、input1、output0等,按索引编号描述第 0、1 个输入与第 0 个输出。
  • name:与该索引对应的输入/输出名称,必须与原型中Input("x")、Output("z")声明的名称一致。
  • dtype:支持的数据类型列表,多个取值用英文逗号分隔。
  • format:支持的数据格式列表,多个取值用英文逗号分隔。dtype与format列表按位置一一对应,例如"dtype": "float16,float32,int32"与"format": "ND,ND,ND"表示支持float16 + ND、float32 + ND、int32 + ND三组组合。

在实际推导函数中,还可以在每组描述中额外携带unknownshape_format字段(如"unknownshape_format": "ND,ND,ND"),用于声明未知 shape(动态 shape)场景下各 dtype 对应的 format,这在支持动态 shape 的算子中尤为常见。

返回值说明

返回OpAICoreDef算子定义对象(即this->AICore()返回的实例),支持继续链式调用其他注册接口。OpAICoreDef的完整接口列表请参考 OpAICoreDef。

约束说明

无特殊约束。不过从功能语义上需注意:推导函数的执行依赖DynamicFormatFlag(true),且当原型已静态注册 dtype/format 时推导结果不会生效(详见上文“功能说明”)。

调用示例

推导函数实现

如下展示了自定义 Add 算子推导函数的实现。示例以运行时第一个输入参数 shape 的第一个维度为判定条件:若x的第 0 维不大于 8,则支持float16、float32、int32三种数据类型(均为ND格式,并同时给出动态 shape 场景的unknownshape_format);否则仅支持int32。

static ge::graphStatus OpSelectFormat(const ge::Operator &op, ge::AscendString &result) { std::string resultJsonStr; // 如果本次执行第一个输入参数shape的第一个维度<=8,则支持更多的格式,否则仅支持int32 if (op.GetInputDesc(0).GetShape().GetDim(0) <= 8) { resultJsonStr = R"({ "input0": {"name": "x","dtype": "float16,float32,int32","format": "ND,ND,ND","unknownshape_format": "ND,ND,ND"}, "input1": {"name": "y","dtype": "float16,float32,int32","format": "ND,ND,ND","unknownshape_format": "ND,ND,ND"}, "output0": {"name": "z","dtype": "float16,float32,int32","format": "ND,ND,ND","unknownshape_format": "ND,ND,ND"} })"; } else { resultJsonStr = R"({ "input0": {"name": "x","dtype": "int32","format": "ND","unknownshape_format": "ND"}, "input1": {"name": "y","dtype": "int32","format": "ND","unknownshape_format": "ND"}, "output0": {"name": "z","dtype": "int32","format": "ND","unknownshape_format": "ND"} })"; } result = ge::AscendString(resultJsonStr.c_str()); return ge::GRAPH_SUCCESS; }

关键点说明:

  • 通过op.GetInputDesc(0).GetShape().GetDim(0)读取第 0 个输入描述的第 0 维大小,这是推导函数读取运行时信息的典型方式,ge::Operator中还提供读取输出描述、属性等信息的接口,均可用于推导逻辑。
  • 结果字符串以ge::AscendString写出,函数返回ge::GRAPH_SUCCESS表示推导成功。
  • 推导函数应为静态函数或位于命名空间中的普通函数,便于作为函数指针传入。

推导函数注册

推导函数的注册样例如下。自定义算子类AddCustom继承自OpDef,在构造函数中依次声明输入输出、注册SetInferShape,然后在AICore()链中注册 Tiling、TilingParse 与SetOpSelectFormat;随后创建OpAICoreConfig并将DynamicFormatFlag配置为true,最后通过AddConfig绑定到实际的 soc 版本:

class AddCustom : public OpDef { public: AddCustom(const char* name) : OpDef(name) { this->Input("x") .ParamType(REQUIRED); this->Input("y") .ParamType(REQUIRED); this->Output("z") .ParamType(REQUIRED); this->SetInferShape(ge::InferShape); this->AICore() .SetTiling(optiling::TilingFunc) .SetTilingParse(optiling::TilingPrepare) .SetOpSelectFormat(optiling::OpSelectFormat); OpAICoreConfig aicConfig; aicConfig.DynamicCompileStaticFlag(true) .DynamicFormatFlag(true) .DynamicRankSupportFlag(true) .DynamicShapeSupportFlag(true) .NeedCheckSupportFlag(false) .PrecisionReduceFlag(true); // 注意:soc_version请替换成实际的AI处理器型号 this->AICore().AddConfig("soc_version", aicConfig); } };

注册完成后,还需在算子原型文件末尾调用OP_ADD(AddCustom)完成算子注册,使编译工具能够识别该原型并生成 aclnn 接口与图模式原型相关代码。

与 DynamicFormatFlag 的配合关系

SetOpSelectFormat必须与 DynamicFormatFlag 配套使用才能生效:

  • DynamicFormatFlag(true):标识根据SetOpSelectFormat设置的函数自动推导算子输入输出支持的 dtype 与 format,此时无需在原型注册时配置固定的 dtype 与 format,算子融合时框架会自动调用推导函数。
  • DynamicFormatFlag(false)或不设置:不启用自动推导。

在OpAICoreConfig提供的全部配置开关中(DynamicCompileStaticFlag、DynamicRankSupportFlag、DynamicShapeSupportFlag、NeedCheckSupportFlag、PrecisionReduceFlag等,详见 OpAICoreConfig),DynamicFormatFlag专管 dtype/format 的自动推导,通常与DynamicShapeSupportFlag(true)、DynamicRankSupportFlag(true)一起开启,以支持动态 shape 场景下的完整推导能力。

仓库中的实际用法佐证

在 asc-devkit 仓库的算子原型注册测试用例 tests/tools/utils/opbuild/test_ops.cpp 中,AddTik2算子演示了该接口的典型组合用法:AICore()链中依次注册SetTiling、SetCheckSupport、SetOpSelectFormat、SetOpSupportInfo、SetOpSpecInfo、SetParamGeneralize,并在OpAICoreConfig中开启DynamicCompileStaticFlag(true)、DynamicFormatFlag(true)、DynamicRankSupportFlag(true)、DynamicShapeSupportFlag(true)等开关后通过AddConfig("ascend310p")/AddConfig("ascend910", aicConfig)绑定不同 AI 处理器:

this->AICore() .SetTiling(optiling::TilingTik2Add) .SetCheckSupport(optiling::check_op_support) .SetOpSelectFormat(optiling::op_select_format) .SetOpSupportInfo(optiling::get_op_support) .SetOpSpecInfo(optiling::get_op_specific_info) .SetParamGeneralize(optiling::generalize_config); OpAICoreConfig aicConfig; aicConfig.DynamicCompileStaticFlag(true) .DynamicFormatFlag(true) .DynamicRankSupportFlag(true) .DynamicShapeSupportFlag(true) .NeedCheckSupportFlag(true) .PrecisionReduceFlag(true); this->AICore().AddConfig("ascend310p"); this->AICore().AddConfig("ascend910", aicConfig);

该用例同时印证了两点:其一,SetOpSelectFormat的回调与SetCheckSupport等接口共用OP_CHECK_FUNC函数签名(ge::graphStatus (*)(const ge::Operator &op, ge::AscendString &result)),回调函数定义于optiling命名空间;其二,SetOpSelectFormat与DynamicFormatFlag(true)在真实工程中总是成对出现。

使用场景与设计建议

  • 按运行时条件动态收窄 dtype/format:如示例所示,当 shape 或属性满足特定条件时支持更丰富的数据类型组合,否则退回最小支持集合,可有效控制算子在不同输入下的编译与匹配范围。
  • 替代静态 dtype/format 声明:对于输入输出组合较多的算子,静态枚举全部组合会使原型膨胀且难以维护,动态推导可将支持集合的生成逻辑集中到单个回调函数中。相关静态声明方式(DataType/Format列表按位置对应)可参考 算子原型定义。
  • 与动态 shape 特性协同:推导结果中的unknownshape_format字段用于声明动态 shape 场景下的格式映射,建议与DynamicShapeSupportFlag(true)同时配置,保证动态 shape 与动态格式推导行为一致。

需要再次强调的是:该机制的本质是“按需推导、动态收口”,一旦原型中静态注册了 dtype/format,推导函数便不再执行。因此在使用SetOpSelectFormat时,应确保输入输出定义中不通过DataType/Format声明固定支持集合,让DynamicFormatFlag与推导回调成为 dtype/format 的唯一来源。

  • 人工智能
  • 深度学习
  • 算子库
  • CANN
  • Ascend

【免费下载链接】asc-devkit

本项目是CANN 推出的昇腾AI处理器专用的算子程序开发语言,原生支持C和C++标准规范,主要由类库和语言扩展层构成,提供多层级API,满足多维场景算子开发诉求。

项目地址:https://gitcode.com/cann/asc-devkit
点击查看免费下载

相关推荐

上一篇:为什么浏览器"另存为"总是失败?kage"先渲染、再剥离、后本地化"三阶段原理全解
下一篇:Transformer的灵魂:How to Train Your GPT用派对类比讲透多头注意力(含8步手算实例)

创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询