CANN opbase InferShape工具集：动态shape推导与Broadcast/Reduce实战

1. 从一次形状推导报错说起InferShape 到底在算什么做算子开发的同学大概率都遇到过这种场景算子写完单测跑起来日志里蹦出一句InferShape failed: dim num mismatch或者broadcast shape error然后你盯着那几行 shape 打印反复看明明觉得逻辑没问题可它就是过不去。更让人头疼的是有些算子在编译期根本拿不到完整的 shape 信息——输入张量的秩rank是未知的维度值也是动态的这时候 InferShape 该怎么写这就是 CANN opbase 里 InferShape 公共工具集要解决的核心问题。它不是一个单一函数而是一组围绕形状推导构建的工具方法覆盖了未知秩处理、Broadcast 形状推导、Elewise 逐元素推导、Reduce 归约推导这几大类典型场景。你写自定义算子的时候只要继承 opbase 提供的基类、调用对应的工具方法就能把大部分形状推导逻辑标准化不用每次从零手写。我先把话说在前面这篇内容适合已经接触过 CANN 算子开发、至少写过一个自定义算子的同学。如果你还没写过算子建议先去看算子注册和 Tiling 的基础流程否则直接看 InferShape 的细节会有点懵。但如果你已经踩过形状推导的坑那接下来的内容应该能帮你省下不少调试时间。InferShape 的本质是什么一句话在算子执行之前根据输入张量的 shape 和算子属性推导出输出张量的 shape。听起来简单但魔鬼在细节里。比如两个张量相加形状分别是[2, 3, 4]和[3, 4]输出是什么再比如一个 ReduceSum输入[2, 3, 4]axis1keepdimsfalse输出又是什么这些规则如果每个算子都手写一遍代码重复不说还容易出错。opbase 的工具集就是把这些通用规则沉淀下来让算子开发者直接复用。2. 未知秩场景下 InferShape 的应对策略2.1 什么是未知秩为什么它让形状推导变复杂在静态图编译场景里大部分时候输入张量的 shape 是完整的rank 和每一维的大小都确定。但实际工程中你会遇到几种不完整的情况rank 未知只知道是个张量但不知道是几维的。比如某些动态 shape 的输入编译期只能确定 dtyperank 要等到运行时才知道。dim 未知rank 确定但某一维的大小是动态的比如 batch 维度。rank 和 dim 都未知最极端的情况编译期几乎拿不到形状信息。未知秩之所以麻烦是因为很多形状推导规则本身就依赖 rank。比如 Broadcast 推导需要对齐维度Reduce 推导需要判断 axis 是否在合法范围内这些操作的前提都是你知道 rank 是多少。如果 rank 未知传统的推导逻辑直接就没法执行。2.2 opbase 对未知秩的处理方式opbase 的思路是在 rank 未知时不强行推导具体形状而是输出一个未知形状的占位同时保证推导过程不报错。具体来说工具集会判断输入的 shape 是否完整如果不完整就返回一个标记为未知的 shape 对象让后续流程知道这里需要运行时再确定。这里有个关键点未知秩不等于放弃推导。有些信息在 rank 未知时仍然可以推导出来。比如输出和输入的 rank 关系如果算子本身不改变 rank比如 Elewise 类算子那即使输入 rank 未知输出的 rank 也等于输入的 rank。opbase 的工具方法会尽量保留这些可推导的信息而不是一刀切地全部标记为未知。我在实际项目里遇到过这样一个案例一个自定义的逐元素激活算子输入 shape 在编译期是[-1, -1, 256]前两维动态第三维固定。如果用朴素的推导逻辑看到有-1就直接放弃那输出的 shape 也没法确定。但实际上逐元素算子的输出 shape 和输入完全一致所以即使有动态维度输出也应该是[-1, -1, 256]。opbase 的 Elewise 推导工具就是干这个的——它不关心具体维度值是多少只关心输入输出的形状对应关系。2.3 写未知秩推导时的几个实操要点第一不要假设 rank 一定存在。你在写 InferShape 函数的时候第一件事应该是检查输入的 shape 是否完整。opbase 提供了判断接口不要自己去猜。第二区分未知和零。有些同学看到动态维度就填 0这是错的。0 是一个合法的维度值比如空张量和未知完全是两回事。opbase 用专门的标记来表示未知你在推导过程中也要保持这个区分。第三未知秩下的输出要尽量保留已知信息。比如输入是[-1, 3, 256]经过一个不改变形状的算子输出也应该是[-1, 3, 256]而不是全部变成未知。这样后续算子还能继续利用已知的维度信息做优化。提示未知秩处理的核心原则是能推则推不能推则标记未知而不是遇到动态维度就全部放弃。这个原则在写自定义算子时非常关键。3. Broadcast 形状推导对齐规则与常见陷阱3.1 Broadcast 的基本规则回顾Broadcast 是形状推导里最常见也最容易出错的场景。基本规则大家都懂从最后一维开始对齐维度相同或其中一个为 1 时可以广播否则报错。但实际写代码的时候细节问题一大堆。先看一个典型例子输入 A输入 B输出[2, 3, 4][3, 4][2, 3, 4][2, 1, 4][1, 3, 1][2, 3, 4][5, 1][1, 6][5, 6][2, 3, 4][2, 3]报错最后一行为什么报错因为从最后一维对齐后A 的第三维是 4B 的第二维是 34 和 3 既不相等也不为 1无法广播。这个规则看起来简单但实际推导时rank 不一致的对齐处理、动态维度的传播、多输入广播的合并每一步都有坑。3.2 opbase 的 Broadcast 推导工具怎么用opbase 提供的 Broadcast 推导工具核心接口接收一组输入 shape返回广播后的输出 shape。它的内部逻辑大致是这样的找出所有输入中最大的 rank。把所有输入的 rank 补齐到最大 rank在前面补 1。从最后一维开始逐维比较按照广播规则计算输出维度。如果某一维无法广播返回失败。这个流程本身不复杂但 opbase 在实现上做了几个增强支持动态维度如果某一维是未知的输出也标记为未知但不直接报错。支持多输入不限于两个输入N 个输入的广播也能处理。错误信息友好广播失败时会明确指出是哪一维不匹配方便定位问题。我在用这个工具的时候最大的感受是不用自己处理 rank 对齐了。以前手写广播逻辑光是补齐 rank 那一段就要写十几行还容易漏掉边界情况。现在直接调工具传入输入 shape 列表拿输出 shape干净利落。3.3 Broadcast 推导中的三个高频坑坑一rank 对齐方向搞反。广播是从最后一维开始对齐的所以补齐 rank 是在前面补 1不是在后面补。我见过有同学在 shape 末尾补 1结果[3, 4]和[2, 3, 4]推导出[3, 4, 1]和[2, 3, 4]广播成[3, 4, 4]完全错了。坑二动态维度和 1 的优先级。如果一维是动态的未知另一维是 1输出应该是什么答案是动态维度。因为 1 可以广播到任意大小包括动态大小。但如果一维是动态的另一维是 3输出也是动态的——因为动态维度可能等于 3也可能不等于编译期无法确定。坑三多输入广播的顺序。三个及以上输入做广播时不要两两依次广播而应该一次性对齐所有输入。两两广播可能会丢失信息比如[2, 1, 4]、[1, 3, 1]、[1, 1, 5]三个输入两两广播的结果可能和一次性广播不一致。注意Broadcast 推导失败时优先检查 rank 对齐方向和各维的匹配情况90% 的错误都出在这两个地方。4. Elewise 与 Reduce 的形状推导逻辑4.1 Elewise 推导看似简单实则暗藏细节Elewise逐元素算子的形状推导规则很直观输出 shape 等于输入 shape。比如 ReLU、Sigmoid、Add同 shape 时这些算子输入输出形状完全一致。但实际场景中Elewise 算子往往伴随着 Broadcast。比如 Add 算子两个输入 shape 不同时需要先广播再逐元素相加。所以 opbase 的 Elewise 推导工具通常会结合 Broadcast 逻辑先对所有输入做广播得到统一的 shape然后输出这个 shape。这里有个容易忽略的点Elewise 算子的输出 shape 不一定等于第一个输入的 shape。如果第一个输入是[3, 4]第二个输入是[2, 3, 4]广播后的输出是[2, 3, 4]而不是[3, 4]。我见过有同学在写自定义 Elewise 算子时直接拿第一个输入的 shape 当输出结果遇到广播场景就挂了。另外Elewise 推导还要处理标量输入的情况。标量的 shape 是空的rank0和任意 shape 广播都得到对方的 shape。opbase 的工具会正确处理这种情况但你自己写的时候要注意别把空 shape 当成错误。4.2 Reduce 推导axis、keepdims 与动态 shape 的三角关系Reduce 类算子ReduceSum、ReduceMax、ReduceMean 等的形状推导比 Elewise 复杂得多因为它涉及三个变量输入 shape、axis、keepdims。推导规则可以总结为如果keepdimstrue被归约的维度变成 1其他维度不变。如果keepdimsfalse被归约的维度直接去掉。axis 可以是单个整数、整数列表也可以是负数从后往前数。举个例子输入[2, 3, 4, 5]axis[1, 3]keepdimsfalse被归约的是第 1 维大小 3和第 3 维大小 5。去掉这两维后输出是[2, 4]。如果 keepdimstrue输出是[2, 1, 4, 1]。看起来规则清晰但实际写推导时有几个坑坑一axis 为负数时的转换。axis-1 表示最后一维需要转换成 rank-1。如果 rank 是动态的这个转换就没法做需要特殊处理。坑二axis 为空列表。有些框架里axis[] 表示对所有维度归约输出是标量keepdimsfalse或全 1 的 shapekeepdimstrue。但有些实现里 axis[] 表示不归约输出等于输入。这个语义差异要在算子定义时明确。坑三动态维度下的 Reduce。如果被归约的维度是动态的输出维度怎么处理如果 keepdimsfalse这一维直接去掉没问题。如果 keepdimstrue这一维变成 1也没问题。但如果 axis 本身依赖动态 rank那就麻烦了——你根本不知道要归约哪一维。opbase 的 Reduce 推导工具对前两种情况都有处理第三种情况会返回未知形状等运行时再确定。4.3 三类推导工具的选型对照推导类型适用场景核心逻辑动态 shape 支持Broadcast多输入形状对齐从最后一维对齐逐维取最大值支持动态维传播Elewise逐元素运算输出等于广播后的输入 shape支持保留已知维度Reduce归约运算根据 axis 和 keepdims 变换维度部分支持axis 依赖 rank 时返回未知选型的时候先判断算子属于哪一类。如果是多输入且形状可能不同用 Broadcast如果是单输入逐元素变换用 Elewise如果涉及维度归约用 Reduce。有些算子可能同时涉及多种比如 Softmax 既有 Reduce 又有 Elewise那就组合使用。5. 把工具集用起来从算子注册到推导验证5.1 在自定义算子中接入 InferShape 工具假设你要写一个自定义的 Add 算子支持广播。InferShape 函数的骨架大概是这样// 伪代码示意具体接口以 opbase 实际 API 为准 ge::graphStatus InferShapeForAdd(gert::InferShapeContext* context) { // 1. 获取输入 shape auto input_shape_0 context-GetInputShape(0); auto input_shape_1 context-GetInputShape(1); // 2. 调用 Broadcast 推导工具 auto output_shape opbase::BroadcastInferShape({input_shape_0, input_shape_1}); // 3. 检查推导结果 if (output_shape nullptr) { return ge::GRAPH_FAILED; } // 4. 设置输出 shape context-GetOutputShape(0)-SetShape(*output_shape); return ge::GRAPH_SUCCESS; }关键点在于第二步不要自己手写广播逻辑直接调 opbase 的工具。工具内部已经处理了 rank 对齐、动态维度、错误检查这些细节。对于 Reduce 算子调用方式类似只是需要额外传入 axis 和 keepdims 属性auto output_shape opbase::ReduceInferShape(input_shape, axis, keepdims);5.2 推导结果的验证方法写完 InferShape 之后怎么验证它是对的我通常用三种方式方式一单测覆盖典型 shape。构造几组有代表性的输入 shape包括同 shape、可广播、不可广播、动态维度等情况检查输出是否符合预期。方式二和框架内置算子对比。如果你的算子和框架已有的某个算子语义一致可以用相同的输入跑一遍对比输出 shape 是否一致。方式三构造边界 case。比如 rank0 的标量、rank1 的向量、含 1 的维度、含动态维度的 shape这些边界情况最容易暴露问题。我自己的习惯是每写一个 InferShape 函数至少准备 10 组以上的测试 shape。别嫌麻烦形状推导的 bug 一旦漏到运行时排查成本比写单测高得多。5.3 几个提升推导健壮性的经验经验一永远不要信任输入的 shape 是完整的。即使你觉得这个算子只会接收静态 shape也要在推导函数里加未知秩的判断。工程实践中动态 shape 的出现频率比你想的高。经验二错误信息要带上下文。推导失败时不要只返回 FAILED要把输入 shape、axis、keepdims 这些关键信息打到日志里。不然出了问题你只能靠猜。经验三动态维度的传播要一致。如果输入有动态维度输出也要正确标记动态维度。不要一会儿用 -1 表示一会儿用特殊标记表示统一用 opbase 的约定。经验四组合算子的推导要分步验证。如果一个算子同时涉及 Broadcast 和 Reduce先验证 Broadcast 部分的输出再验证 Reduce 部分的输出不要一步到位。6. 动态 shape 场景下的推导边界与取舍6.1 什么时候该返回未知什么时候该报错这是动态 shape 推导里最需要判断力的问题。我的原则是能确定的信息尽量保留即使 rank 未知如果算子不改变 rank输出的 rank 也应该等于输入的 rank。无法确定的信息标记未知不要瞎猜也不要填默认值。确定非法的输入直接报错比如广播时两维既不相等也不为 1这是确定的错误应该报错而不是返回未知。区分未知和非法很重要。未知是信息不足非法是逻辑矛盾。前者应该容忍后者应该拒绝。6.2 动态 rank 对 Reduce 推导的影响Reduce 算子在动态 rank 下最麻烦。假设输入 rank 未知axis1你根本不知道第 1 维是否存在。这时候有两种处理策略策略一返回未知形状。保守做法等运行时再推导。缺点是编译期优化机会减少。策略二假设 rank 足够大。如果 axis1假设 rank 至少为 2按这个假设推导。如果运行时发现 rank 不够再报错。这种策略激进一些但能给编译期提供更多信息。opbase 的工具默认采用策略一但提供了接口让你根据算子语义选择策略二。选哪种取决于你的算子对动态 shape 的容忍度。6.3 一个真实的调试案例我之前写过一个自定义的归一化算子涉及 Reduce 和 Elewise 的组合。静态 shape 下一切正常但一上动态 shape 就报错。排查了半天发现问题是Reduce 之后的输出 shape 里被归约的维度变成了 1keepdimstrue但后续的 Elewise 推导没有正确处理这个含 1 的维度导致广播时对齐出错。修复方法很简单在 Elewise 推导之前先把 Reduce 输出的 shape 正确设置好确保含 1 的维度被保留。这个案例告诉我组合算子的推导顺序很重要前一步的输出必须完全正确后一步才能基于它推导。7. 写在最后一些个人体会InferShape 这个环节说难不难说简单也不简单。规则本身是清晰的但实际工程中的动态 shape、多输入广播、组合算子这些场景会让问题复杂度成倍上升。opbase 的公共工具集最大的价值就是把这些通用逻辑沉淀下来让你不用重复造轮子。我的建议是先把工具集的接口和语义搞清楚再动手写自己的推导逻辑。很多坑工具集已经帮你填过了你没必要再踩一遍。另外动态 shape 的测试一定要做足静态 shape 下跑通不代表动态 shape 下没问题。最后分享一个小技巧如果你不确定某个 shape 推导结果对不对可以先用框架内置的同类算子跑一遍把输入输出 shape 打印出来作为参考基准。这个方法在调试复杂算子时特别管用比对着文档推半天快得多。