☰
asc-devkit 算子开发:使用 SetCheckSupport 注册算子参数校验回调,实现 AI Core / AI CPU 支持性分流
2026/10/16 2:56:23 网站建设 项目流程
  • 人工智能
  • 深度学习
  • 算子库
  • CANN
  • Ascend

【免费下载链接】asc-devkit

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

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

导读

在 CANN 昇腾算子开发(asc-devkit)中,算子参数校验(CheckSupport)是算子融合阶段决定"这个算子到底该由 AI Core 执行,还是回退到 AI CPU 算子库执行"的关键机制。本文基于 SetCheckSupport 接口文档,完整讲解SetCheckSupport的函数原型、回调函数OP_CHECK_FUNC的签名约定、出参 JSON 格式、与其配套的 NeedCheckSupportFlag 开关,并结合仓库中的算子原型注册示例给出可直接复用的实现与注册代码。读完本文,你将掌握在自定义算子中自行实现参数校验回调、正确配置 OpAICoreConfig、并理解 AI Core 与 AI CPU 支持性分流底层逻辑的完整能力。

一、为什么需要"算子参数校验":融合阶段的执行路径分流

在昇腾算子编译与图融合流程中,框架需要回答一个核心问题:针对当前这一组具体的算子输入(data type、shape、format 等),本算子是否支持在 AI Core 上执行?

默认情况下,算子原型中的输入输出、属性等信息是静态声明的,而实际图上的输入往往千变万化(不同 dtype、不同 shape、不同维度)。为了在融合阶段精确判定支持性,asc-devkit 提供了"算子参数校验回调"机制:

  • 开发者实现一个参数校验函数,通过SetCheckSupport注册到算子的 AI Core 定义(OpAICoreDef)上;
  • 同时将 NeedCheckSupportFlag 配置为true;
  • 此后,算子编译和融合阶段会调用注册的校验函数对算子参数进行校验:
校验通过 → 代表 AI Core 支持此算子参数 → 选择 AI Core 上相应的算子执行 校验失败 → 继续查询 AI CPU 算子库 → 回退到 AI CPU 执行

这一"AI Core 优先、AI CPU 兜底"的执行路径,让算子在不同输入形态下都能找到可用的实现,同时尽可能优先利用 AI Core 的高算力。从仓库源码看,AI Core 侧算子实现中同样存在类似的支持性判断逻辑,例如 impl/adv_api/detail/matmul/scheduler/base/scheduler_base.h 中的CheckSupportTrianMatmul()就是 Matmul 三角形态场景下对输入 shape 的支持性检查。

二、接口速览:函数原型与返回值

SetCheckSupport是OpAICoreDef类的成员方法,用于注册算子参数校验函数,声明如下:

OpAICoreDef &SetCheckSupport(optiling::OP_CHECK_FUNC func)
  • 参数func(输入):算子参数校验函数,类型OP_CHECK_FUNC定义如下:

    using OP_CHECK_FUNC = ge::graphStatus (*)(const ge::Operator &op, ge::AscendString &result);
  • 返回值:OpAICoreDef引用。OpAICoreDef是算子在 AI Core 上相关实现信息的载体,除SetCheckSupport外还提供 SetTiling、SetOpSelectFormat、AddConfig 等注册接口(详见 OpAICoreDef 总览)。返回引用支持链式调用,例如this->AICore().SetTiling(...).SetCheckSupport(...)。

  • 约束说明:无特殊约束。

三、回调函数契约:入参、出参与返回码

OP_CHECK_FUNC类型的回调函数是整条校验链的核心,其契约包含三个方面:

1. 入参:算子描述ge::Operator &op

op携带当前算子的完整描述,包括算子的输入、输出、属性等信息。在校验函数中,可以通过op.GetInputDesc(i)、op.GetOutputDesc(i)等接口获取对应描述,再进一步读取 shape、dtype 等细节。例如:

op.GetInputDesc(0).GetShape().GetDim(0)

用于获取第 0 个输入描述中 shape 的第 0 维大小,这正是文档调用示例中校验的目标。

2. 出参:校验结果字符串ge::AscendString &result

出参是一个字符串,包含校验返回码和原因,格式固定为 JSON:

{"ret_code": "1","reason": "your reason"}
字段含义
ret_code校验结果返回码,示例中"1"表示支持、"0"表示不支持(具体取值语义以实际框架实现为准,回调中须与校验结论保持一致)
reason校验结果的原因描述,便于日志排查与框架侧定位

在校验函数末尾,需要将该 JSON 字符串写入result:

result = ge::AscendString(resultJsonStr.c_str());

3. 返回值:校验结论ge::graphStatus

  • 校验成功(AI Core 支持该参数):返回ge::GRAPH_SUCCESS;
  • 校验失败(AI Core 不支持):返回ge::GRAPH_FAILED。

返回值与result中的ret_code必须保持语义一致:返回GRAPH_SUCCESS时ret_code置为支持标识,返回GRAPH_FAILED时置为不支持标识。

四、配套开关:NeedCheckSupportFlag

仅有回调函数还不够,SetCheckSupport必须与 NeedCheckSupportFlag 配套使用,校验回调才会被框架真正调用:

OpAICoreConfig &NeedCheckSupportFlag(bool flag)
  • flag = true:标识是否在算子融合阶段调用算子参数校验函数进行data type 与 shape的校验。配置为true时,框架会调用通过SetCheckSupport设置的算子参数校验函数,检查算子是否支持指定输入——此场景下必须自行实现算子参数校验的回调函数;
  • flag = false:表示不需要进行校验。

需要注意的是,该开关的默认值为false。根据 AddConfig.md 默认配置表,当通过AddConfig(const char *soc)单参形式注册 AI 处理器型号、不显式传入OpAICoreConfig时,NeedCheckSupportFlag会被配置为默认值false(同表其他配置项DynamicCompileStaticFlag、DynamicFormatFlag、DynamicRankSupportFlag、DynamicShapeSupportFlag、PrecisionReduceFlag默认值均为true)。因此,只要期望融合阶段执行参数校验,就必须显式构造OpAICoreConfig并把NeedCheckSupportFlag(true)写入配置。

NeedCheckSupportFlag是OpAICoreConfig的成员方法,返回OpAICoreConfig &以支持链式配置(类总览见 OpAICoreConfig)。

五、完整实战:自定义 Add 算子的校验实现与注册

文档调用示例给出了一个完整的自定义AddCustom算子场景:仅支持第一个输入x的 shape 第一个维度为 8,其余 shape 一律判为不支持。

步骤 1:实现参数校验函数CheckSupported

static ge::graphStatus CheckSupported(const ge::Operator &op, ge::AscendString &result) { std::string resultJsonStr; // 仅支持第一个输入参数shape的第一个维度为8,其他shape不支持 if (op.GetInputDesc(0).GetShape().GetDim(0) == 8) { resultJsonStr = R"({"ret_code": "1","reason": "x.dim[0] is 8"})"; result = ge::AscendString(resultJsonStr.c_str()); return ge::GRAPH_SUCCESS; } resultJsonStr = R"({"ret_code": "0","reason": "xxx"})"; result = ge::AscendString(resultJsonStr.c_str()); return ge::GRAPH_FAILED; }

实现要点:

  • 通过op.GetInputDesc(0).GetShape().GetDim(0)读取输入描述与 shape 信息,按需扩展可继续校验 dtype、format、维度数等;
  • 两种分支都要同时写result(JSON 出参)并返回对应的ge::GRAPH_SUCCESS/ge::GRAPH_FAILED;
  • 该函数使用ge::命名空间下的类型,定义位置应与算子原型注册代码保持一致的命名空间组织。

步骤 2:在算子原型类中注册回调并配置开关

AddCustom继承自OpDef,在构造函数中依次完成输入输出声明、shape 推导注册,以及 AI Core 相关函数的注册:

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) .SetCheckSupport(optiling::CheckSupported); OpAICoreConfig aicConfig; aicConfig.DynamicCompileStaticFlag(true) .DynamicFormatFlag(true) .DynamicRankSupportFlag(true) .DynamicShapeSupportFlag(true) .NeedCheckSupportFlag(true) .PrecisionReduceFlag(true); // 注意:soc_version请替换成实际的AI处理器型号 this->AICore().AddConfig("soc_version", aicConfig); } };

注册要点:

  • 链式注册:AICore()返回OpAICoreDef,SetCheckSupport(optiling::CheckSupported)与SetTiling、SetTilingParse、SetOpSelectFormat一起完成 AI Core 侧实现函数的挂载;
  • 显式构造配置:OpAICoreConfig aicConfig后通过链式方法逐项打开动态能力开关,其中NeedCheckSupportFlag(true)是触发融合阶段参数校验的关键;
  • 双参AddConfig:AddConfig("soc_version", aicConfig)将配置与具体 AI 处理器型号绑定,soc_version需替换为实际型号(可参考算子工程编译配置CMakePresets.json中的ASCEND_COMPUTE_UNIT字段取值,该字段在工程创建时由 msOpGen 自动生成)。

步骤 3:通过 OP_ADD 完成原型注册

算子原型类定义完成后,需要调用 OP_ADD 接口 注册算子类型,框架才能正确识别、编译和执行该算子:

OP_ADD(AddCustom);

六、仓库示例佐证:真实工程中的原型注册形态

上述注册模式在仓库的 ACL 自定义算子示例中可以得到印证。例如 examples/01_simd_cpp_api/02_features/99_acl_based/00_acl_compilation/custom_op/op_host/add_custom/add_custom_host.cpp 中,AddCustom同样是OpDef的子类,构造函数内通过this->Input("x").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);

同目录下的 tiling 下沉示例 add_custom_tiling_sink.cpp 还展示了多型号注册形态:

this->AICore().SetTiling(optiling::AddCustomSinkTilingFunc).AddConfig("ascend910b").AddConfig("ascend910_93");

可以看到,仓库示例与文档样例在骨架上是完全一致的:OpDef派生类 → 声明输入输出与 shape 推导 →AICore()链式注册 AI Core 函数 →AddConfig绑定型号 →OP_ADD完成注册。若需要为某型号启用融合阶段的参数校验,只需在AddConfig处改用双参形式并传入带NeedCheckSupportFlag(true)的OpAICoreConfig即可。

七、小结

SetCheckSupport与NeedCheckSupportFlag共同构成了 asc-devkit 中算子融合阶段支持性校验的完整闭环:

  1. 注册:通过OpAICoreDef::SetCheckSupport(OP_CHECK_FUNC)挂载自定义校验回调,回调依据ge::Operator中的输入输出与属性信息判定支持性,以{"ret_code": "...", "reason": "..."}JSON 字符串输出结果,并返回GRAPH_SUCCESS/GRAPH_FAILED;
  2. 开关:通过OpAICoreConfig::NeedCheckSupportFlag(true)显式开启融合阶段校验(默认值为false,单参AddConfig不会自动开启);
  3. 分流:校验通过走 AI Core 算子执行,校验失败自动回退查询 AI CPU 算子库执行,保证算子在不同输入形态下的可用性。

对于需要精确控制"哪些 dtype/shape 才允许上 AI Core"的自定义算子开发者而言,这套机制是算子原型定义(原型注册与管理)中不可忽视的一环。本文涉及的完整接口族(SetTiling、SetOpSelectFormat、AddConfig等)可继续在 OpAICoreDef 与 OpAICoreConfig 文档中查阅,仓库示例代码可参考 examples/01_simd_cpp_api/02_features/99_acl_based/00_acl_compilation/custom_op 目录。

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

【免费下载链接】asc-devkit

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

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

相关推荐

上一篇:OpenDesign 中的 Shopify 设计系统实战:暗色优先的影院式电商界面设计指南
下一篇:Apache SkyWalking 浏览器监控(Browser Agent)接入指南:Client JS 上报与 receiver-browser 配置详解

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

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

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

立即咨询