CANN opbase 的 OP_OUTSHAPE 宏详解:让 NonZero 类算子动态刷新输出 Shape 的实现原理与用法

CANN opbase 的 OP_OUTSHAPE 宏详解:让 NonZero 类算子动态刷新输出 Shape 的实现原理与用法 CANN opbase 的 OP_OUTSHAPE 宏详解让 NonZero 类算子动态刷新输出 Shape 的实现原理与用法【免费下载链接】opbase本项目是CANN算子库的基础框架库为算子提供公共依赖文件和基础调度能力。项目地址: https://gitcode.com/cann/opbase导读OP_OUTSHAPE是 CANN opbase 算子库中用于处理「依赖计算结果来确定输出 shape」的算子如 NonZero的关键宏它把 kernel 侧算出的输出 shape 以 aclTensor 形式带出并在 host 侧用它刷新对应输出 tensor 的 shape。本文从宏原型、参数语义、9 位 shape 存储格式与约束出发结合 opbase 仓库中 op_arg_def.h、outshape.cpp 等源码逐层拆解其底层实现与完整调用链并给出可直接复用的 host 侧 / kernel 侧示例。读完本文你将能够在自己的算子中正确声明、写入并刷新输出 shape理解动态 shape 场景下算子缓存的取舍。宏功能与适用场景针对需要计算结果来确定输出 shape 的算子如 NonZero 算子OP_OUTSHAPE宏用于存放此类算子输出 shape 的 aclTensor。这类算子的共同特征是输出 tensor 的维度个数与各维度取值无法在 shape 推导阶段静态确定必须等 kernel 真正执行后才知道。因此 opbase 提供了一条专门的「结果 shape 回传」通路host 侧预先分配一个outShapeTensorshape 大小按可能的最大输出规模预留kernel 执行时把算出的输出 shape 写入该 tensor 的 device 内存host 侧在算子执行完成后把 device 上的 shape 数据同步回来并刷新到对应的输出 tensor 上。OP_OUTSHAPE宏在整个 aclnn 算子实现中被归类为一种独立的算子参数类型与OP_INPUT、OP_OUTPUT、OP_ATTR、OP_WORKSPACE等并列。参考 ADD_TO_LAUNCHER_LIST_AICORE.md 中列出的关联接口清单它通常配合ADD_TO_LAUNCHER_LIST_AICORE等创建执行任务的宏一起使用是 AI Core 算子实现动态输出 shape 的标准入口。宏原型与参数说明OP_OUTSHAPE(x...)参数输入/输出说明x...输入包含两部分第一个参数是存放输出 tensor shape 的 aclTensor即 outShapeTensor第二个参数是需要更新输出 shape 的 tensor 索引。第一个参数的类型是aclTensor*第二个参数的类型是uint64_t。这一点可以从宏的源码定义直接得到印证op_arg_def.h 中#define OP_INPUT(x...) op::OpInput(std::make_tuple(x)) #define OP_OUTPUT(x...) op::OpOutput(std::make_tuple(x)) #define OP_ATTR(x...) op::OpAttr(std::make_tuple(x)) #define OP_WORKSPACE(x...) op::OpWorkspace(std::make_tuple(x)) #define OP_OUTSHAPE(x...) op::OpOutshape(std::tupleaclTensor*, uint64_t(x)) #define OP_OPTION(x...) op::OpOption(std::make_tuple(x)) #define OP_EMPTY_ARG op::EMPTY_OP_ARG #define OP_MODE(x...) op::OpMode(std::make_tuple(x))可以看到OP_OUTSHAPE与其他参数宏不同它显式将参数约束为std::tupleaclTensor*, uint64_t第一个成员是存放输出 shape 的 aclTensor 指针第二个成员是待刷新 shape 的输出 tensor 索引。在参数类型的枚举体系中outshape 也有独立的类型标识。op_arg_def.h 中定义了OP_OUTSHAPE_ARG 4并通过DEFINE_OP_ARG(OpOutshape, OP_OUTSHAPE_ARG)注册了对应的参数包装类OpOutshape。这意味着 outshape 参数会被框架作为一个独立的 OpArg 收集进OpArgContext后续可按OP_OUTSHAPE_ARG类型统一识别和分发。约束说明9 位 shape 存储格式OP_OUTSHAPE只支持存放一个输出 tensor shape 的 TensoroutShapeTensor其对应 shape 为(9 * 需要刷新的 tensor 个数,)每个输出 tensor 的 shape 占 9 位其中第 1 位表示维数rank剩下 8 位表示具体的每个维度的取值即最多支持 8 维输出。该约束与底层实现中的常量一一对应。outshape.cpp 中定义static constexpr size_t OP_OUTSHAPE_RANK 9; static constexpr size_t OP_OUTSHAPE_COUNT 2;OP_OUTSHAPE_RANK 9即每个输出 tensor 占用的 shape 槽位数OP_OUTSHAPE_COUNT 2则表示在 outshape 参数列表中每个刷新需求由「存放 shape 的 aclTensor 输出索引」两个元素组成。host 侧刷新逻辑正是按OP_OUTSHAPE_COUNT步长遍历、并按OP_OUTSHAPE_RANK偏移定位每个输出 tensor 的 shape 数据块。由此可以推导出两个实用结论若需要刷新 N 个输出 tensoroutShapeTensor 的元素个数应为9 * N每个输出 tensor 支持的最大维数为 8第 1 位存维数第 29 位存各维值超出 8 维的输出无法通过该机制表达。从实现看outShapeTensor 支持DT_FLOAT、DT_INT32、DT_INT64三种数据类型outshape.cpp文档示例使用的是DT_INT64。若传入其他数据类型RefreshOutputShape会打印Unsupported outshape dtype错误日志并返回ACLNN_ERR_INNER。底层实现device shape 回传 host 并刷新输出OP_OUTSHAPE的落地逻辑集中在 outshape.cpp 与 outshape.h 中整个流程可拆解为三步。1. 同步 device 数据到 hostSyncTensorDataoutshape.h按 outShapeTensor 的数据类型与元素个数计算字节数从BlockPool申请 host 内存然后调用aclrtMemcpy以ACL_MEMCPY_DEVICE_TO_HOST方向把 kernel 写入的 shape 数据拷贝到 host。若拷贝失败会释放内存并返回失败。2. 遍历 outshape 参数逐个解析 shape 数据块RefreshOutputShapeoutshape.cpp拿到同步回来的 host 内存后按数据类型分派处理if (dtype op::DataType::DT_INT64) { for (size_t i 1; i outputShape.count; i OP_OUTSHAPE_COUNT) { void* currShape PtrShift(shapeData, (i / OP_OUTSHAPE_COUNT * OP_OUTSHAPE_RANK * dtypeSize)); outputs.VisitAt(outputShape[i]-value, currShape { UpdateTensorShapeint64_t(idx, reinterpret_castaclTensor*(elem-pointer), currShape); }); } }注意循环从i 1开始、步长为 2每两个元素为一组第一个是 outShapeTensoroutputShape[0]只在首次用到第二个是待刷新输出 tensor 的索引存在outputShape[i]-value中。通过i / OP_OUTSHAPE_COUNT * OP_OUTSHAPE_RANK计算该输出对应的 shape 数据块在整体缓冲区中的偏移再通过outputs.VisitAt(索引, ...)定位到对应的输出 tensor 并刷新。3. 更新输出 tensor 的三套 shapeUpdateTensorShapeoutshape.h是真正修改输出 tensor 的地方。它先读取第 1 位作为维数dimNum随后按i * sizeof(T)依次偏移读取各维大小构造出新的gert::Shape最后一次性写入输出 tensor 的三套 shapeconst_castaclTensor*(arg)-SetStorageShape(newShape); const_castaclTensor*(arg)-SetOriginalShape(newShape); const_castaclTensor*(arg)-SetViewShape(newShape);也就是说刷新动作是「全量覆盖」——storage shape、original shape 与 view shape 被统一更新为 kernel 计算出的新 shape保证后续推理、内存分配与视图计算都基于真实 shape 进行。值得注意的细节当数据类型为int64_t时实现会对维数做兼容处理——若dimNum 0打印 warningdimNum should be positive否则执行dimNum 0xffffffff7fffffff位运算。源码注释说明这是为适配三类算子uint64 类型数据按 int64 解析时第 32 位置 0以及识别非正数场景避免误判。在完整调用链中的位置OP_OUTSHAPE不单独生效它经由创建执行任务的宏进入算子执行流程。以 AI Core 算子为例make_op_executor.h 中#define ADD_TO_LAUNCHER_LIST_AICORE(KERNEL_NAME, op_args...) \ ({ \ aclnnStatus addToLaunchRet; \ do { \ op::OpArgContext* opArgCtx GetOpArgContext(op_args); \ addToLaunchRet CreatAiCoreKernelLauncher(#KERNEL_NAME, \ KERNEL_NAME##OpTypeId(), \ executor, opArgCtx); \ } while (0); \ addToLaunchRet; \ })CreatAiCoreKernelLauncherop_executor.cpp中对 outshape 有两处关键处理关闭算子缓存当参数列表包含OP_OUTSHAPE_ARG时调用executor-AbandonCache(true)。原因很直接输出 shape 依赖 kernel 运行结果同一组输入每次运行产生的 shape 都可能不同因此这类算子不能走常规的缓存复用路径必须放弃缓存、保证每次都真实执行并回传 shape。参与构图将 outshape 参数传入BuildGraph(..., *args-GetOpArg(op::OP_OUTSHAPE_ARG))把刷新输出 shape 的依赖关系记录进执行图确保在 kernel 执行完成后、结果返回前完成 shape 刷新。从源码结构看可以推断整条链路为OP_OUTSHAPE收集参数 →GetOpArgContext打包成OpArgContext→CreatAiCoreKernelLauncher识别并关闭缓存、加入 launcher 队列 → kernel 执行并把 shape 写入 outShapeTensor →RefreshOutputShape同步回 host 并刷新输出 tensor 的三套 shape。opbase 中的INFER_SHAPE宏make_op_executor.h在调用ADD_TO_LAUNCHER_LIST_AICORE之前执行静态 shape 推导二者相互配合能静态推的由INFER_SHAPE推推不了的动态部分交给OP_OUTSHAPE在运行期补齐。调用示例完整可参考刷新单个输出 tensor// 表示算子将输出 tensor 的 shape 存放到 outShapeTensor 中并且用来更新 idx0 的输出 tensor 的 shape OP_OUTSHAPE({outShapeTensor, 0});刷新多个输出 tensorhost 侧当需要同时刷新多个输出 tensor 的 shape 时多个OP_OUTSHAPE复用同一个 outShapeTensor仅索引不同。本例中需要更新 idx0、3、4 三个输出 tensor 的 shape因此 outShapeTensor 的元素个数为9 * 3 27// host 侧 Shape outShapeShape{27}; auto outShapeTensor executor-AllocTensor(outShapeShape, DataType::DT_INT64, Format::FORMAT_ND); aclnnStatus ret ADD_TO_LAUNCHER_LIST_AICORE( xxx, OP_INPUT(...), OP_OUTPUT(...), OP_ATTR(...), OP_OUTSHAPE({outShapeTensor, 0}), OP_OUTSHAPE({outShapeTensor, 3}), OP_OUTSHAPE({outShapeTensor, 4}), );注意 outShapeTensor 的数据类型此处为DT_INT64必须与 kernel 侧写入的数据类型一致且落在DT_FLOAT/DT_INT32/DT_INT64三者之一否则刷新阶段会报Unsupported outshape dtype。kernel 侧写入 shapekernel 侧按 9 位一组的格式把计算结果写入 shape 缓冲区再进行数据搬运。以下示例展示了三个输出 tensor 的 shape 填充方式每个输出从9 * 输出序号的位置开始第 0 位写维数随后依次写各维取值// kernel 侧 __aicore__ inline void CopyOutShape(uint64_t dimNums1, uint64_t *dimNums2, uint64_t dimNums3) { LocalTensoruint64_t shapeTensor shapeBuf_.Getuint64_t(); shapeTensor.SetValue(0, 1); // 第一个输出tensor的维度信息 shapeTensor.SetValue(1, dimNums1); // 第一个输出tensor的第一维的shape值 shapeTensor.SetValue(9, 1); // 第二个输出tensor的维度信息 shapeTensor.SetValue(10, *(dimNums2)); // 第二个输出tensor的第一维的shape值 shapeTensor.SetValue(11, *(dimNums21)); // 第二个输出tensor的第二维的shape值 shapeTensor.SetValue(18, 1); // 第三个输出tensor的维度信息 shapeTensor.SetValue(19, dimNums3); // 第三个输出tensor的第一维的shape值 ... DataCopyPad(...); }结合约束说明可以确认第一个输出占位 08维数在 0各维值在 18第二个输出占位 9179 为维数10、11 为前两维值第三个输出占位 1826。若第二个输出的实际维数不止 2 维继续按 12、13……的顺序填充即可最多填充到第 8 位即偏移 16。测试用例佐证opbase 的单元测试中对OP_OUTSHAPE有直接覆盖。在 test_kernel_launch.cpp 等位置可以看到如下用法auto outshape OP_OUTSHAPE(tensorPtr4, 0);以及test_kernel_launch.cppauto outshape_arg OP_OUTSHAPE({outshape.get(), 0});两种写法分别对应宏定义中的两种传参形式直接传(tensorPtr4, 0)会经过std::tupleaclTensor*, uint64_t(x)的隐式包装而显式传{tensorPtr4, 0}则与文档示例一致。这些测试用例可作为验证自己算子中OP_OUTSHAPE用法是否规范的最直接参照。使用注意事项小结一个 outShapeTensor 只承载一份输出 shape 数据流约束说明明确「只支持存放一个输出 tensor shape 的 Tensor」多个输出需要刷新时通过多个OP_OUTSHAPE参数复用同一个 tensor按索引区分。容量必须按最大可能输出预留每个输出固定占 9 位1 位维数 8 位维度值需要刷新的输出个数为 N 时outShapeTensor 的 shape 必须为9 * N。当实际输出维度不足 8 维时剩余槽位无需填充host 侧只按实际维数读取。数据类型必须对齐host 分配 outShapeTensor 的 dtype 要与 kernel 写入的 dtype 一致且仅支持DT_FLOAT/DT_INT32/DT_INT64。缓存被强制放弃使用OP_OUTSHAPE的算子无法命中算子缓存每次调用都会真实执行并回传 shape在算子性能设计时应有所预期。执行顺序如果算子同时需要静态 shape 推导INFER_SHAPE需在ADD_TO_LAUNCHER_LIST_AICORE之前调用见 ADD_TO_LAUNCHER_LIST_AICORE.md 的约束说明动态部分则由OP_OUTSHAPE在运行期补齐。更多与OP_OUTSHAPE并列的参数宏OP_INPUT、OP_OUTPUT、OP_ATTR、OP_WORKSPACE等可参考 0_opdev_api_list.md 与 1_opdev_api_introduction.md 中的整体索引OP_OUTSHAPE的英文版本说明见 OP_OUTSHAPE.md。【免费下载链接】opbase本项目是CANN算子库的基础框架库为算子提供公共依赖文件和基础调度能力。项目地址: https://gitcode.com/cann/opbase创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考