★ FEATURED ARTICLE 1. 两阶段调用到底拆开了什么1.1 从一次能跑但调度很重的经历说起做算子开发的朋友第一次接触 aclnn 接口时几乎都会产生同一个疑问为什么调一个算子要先跑一遍带GetWorkspaceSize后缀的接口再跑一遍真正的执行接口这不是白白多一次调用吗我刚接触 ops-nn 算子开发时也是这么想的甚至一度觉得这个设计多余。直到我负责的算子接入推理业务后发现小算子的端到端耗时里有相当一部分花在了 host 侧调度上每次调用都要根据输入张量的 shape、dtype、format 重新做一次算法选择、切分策略计算、参数打包而这些计算结果明明没有变化。两阶段调用就是针对这个问题设计的把算清楚怎么做和真正去做彻底分离开。一个算子要真正在设备上跑起来host 侧需要完成的准备工作远比想象中多。输入张量可能是各种 shape 和布局同样的算子内核要根据这些元信息决定采用哪种实现、数据怎么切块、中间结果放在哪里。aclnn 的第一阶段接口本质上是把这些准备工作全部前置做完后封装进一个不透明句柄第二阶段接口拿到句柄后几乎不再做任何解析直接把任务提交到指定 stream 上。1.2 查询阶段先把账算清楚再动手两阶段中的第一阶段也就是所有命名里带GetWorkspaceSize的接口名字直译过来是获取工作空间大小但它做的远不止算一个数字。它第一个职责是完成 host 侧的全部准备工作根据输入张量的描述信息推导输出 shape选择合适的算子实现路径计算 tiling 参数把最终要下发的指令打包好一并放进一个 executor 句柄里。第二个职责才是返回一个关键数字——workspaceSize即算子内核在执行过程中需要多少临时工作空间。为什么要单独把 workspace 大小算出来交给调用者因为很多算子内核在执行时确实需要一块额外的临时显存。以树状归并类算子为例多级归并在每一层都要保存中间结果这些中间结果的总大小和输入的数据量强相关只有根据当前 shape 精算才能得到确切数字。这块显存如果由算子内部偷偷分配必然涉及分配、释放职责和 stream 同步的问题调用者完全无法控制内存水位如果让调用者随便给一块又容易出现大小不足导致数据越界。两阶段设计的折中方案是框架负责精算大小调用者负责实际分配一块设备内存再把指针交给执行阶段使用。职责边界非常清晰。1.3 执行阶段拿着施工方案直接开工第二阶段接口一般就是算子同名接口本身它接收第一阶段返回的 executor、一块足够大的 workspace 指针以及调用者指定的 stream。这个接口内部不需要再做 shape 推导不需要再选算法只需要把 executor 里已经打包好的任务排进目标 stream剩下的交给设备侧异步执行。把执行拆成这两步最大的收益在于一次准备、多次执行。对固定的输入形状第一阶段可以只做一次之后的每次计算都直接复用那个 executor反复提交执行而当输入 shape 发生变化时只需要重新走一遍第一阶段第二阶段的执行代码完全不用改。这种模式对训练、推理这类高频小算子调用场景尤其友好省掉的是每一轮调用里重复的 host 侧计算开销。同时这个设计还让内存规划变得透明。调用者在第一阶段就知道当前这轮算子需要多大临时空间可以在算子执行前统一从内存池里取执行完统一还对整体显存水位心里有数。理解了这两个阶段各自的定位再看任何 ops-nn 算子暴露出来的 aclnn 接口都会觉得套路非常统一一个查询接口加一个执行接口中间夹着一块调用者负责管理的 workspace。2. GetWorkspaceSize 阶段实操元信息、返回值和易被忽略的边界2.1 入参全是元信息这阶段不碰数据第一阶段的入参需要特别留意这里传的是张量描述结构不是张量实际数据。以常见的逐元素算子为例第一阶段接口接收的是输入张量的描述、输出张量的描述外加两个出参指针——一个用于接收 workspaceSize一个用于接收 executor 句柄。这里有个容易忽略的坑第一阶段完全不会读取张量内容它只关心 shape、dtype、format 这些描述性信息。所以即使张量数据还没拷贝到设备上第一阶段也可以正常执行并返回成功。正确做法是在调用第一阶段之前先把输入数据准备好、确保它在设备侧可见再走查询和执行。如果反过来先把第一阶段跑完再去异步拷贝输入数据就可能出现执行阶段拿到的是旧数据或者脏数据的问题。另一个常见误解是有人以为 execGetWorkspaceSize 返回的 size 是固定的、只跟算子种类有关。实际上它跟输入 shape 强相关同一个算子处理 1x1 的输入和处理 1024x1024 的输入所需 workspace 可能差好几个数量级。开发时我习惯在第一阶段返回后打印一下 size这样既能验证查询逻辑也能在后续排查显存异常时作为参考基线。2.2 workspaceSize 为 0 时怎么处理不少算子确实不需要额外工作空间查询接口会返回 0。这个情况看起来省事但处理不当也会埋雷。首先要明确一点size 为 0 不代表执行接口不接收 workspace 指针部分算子的内核实现即使计算过程中用不到临时空间仍然会在入口处读取一下地址做判空或者对齐操作。稳妥的做法是把 workspace 指针统一分配哪怕是 size 为 0 的情况也通过平台提供的显存分配接口申请一个最小单位的空间。多花一次分配的开销可以忽略但换来的是执行接口永远不会因为空指针或非法地址而崩溃。我在团队里一直强调一个原则查询返回什么数字执行阶段就如实传什么数字不要自作主张地补零、对齐或放大。两阶段之间传递的所有参数包括 workspaceSize、executor都是内部契约的一部分。调用者一旦在中间环节改了参数破坏的就是内部封装好的预期排查起来非常痛苦。2.3 别把 size 当常数缓存有开发经验的人可能会想既然同一个算子对相似 shape 的 size 计算是确定的那我把第一次查询结果缓存下来后面直接复用不行吗这个想法方向对了但实操中很容易踩线。因为 size 不仅随 shape 变化还可能随算子的实现版本、输入 format、甚至是平台侧某个开关状态变化。算子库更新后同样的输入可能选用了不同的内核算法workspace 需求随之改变。如果调用端缓存了旧值场景一换就出问题。真正安全的缓存方式是缓存整个 executor 和配套 workspace 的分配结果并且在 exec 执行前用引用计数管理生命周期而不是只缓存一个数字。如果确实需要高频调用同一个 shape 的算子建议把第一阶段 分配 workspace这个组合体缓存起来每次执行直接复用 executor用完再统一释放。这样既避免了重复做 host 侧解析又不会因为缓存了一个过时的 size 而出错。这块在实操中收益非常明显我在后面第五节还会展开讲。3. 执行阶段与生命周期executor、workspace、stream 的三角关系3.1 executor 是第一阶段的核心产物第一阶段返回的 executor 句柄是整个两阶段机制里最容易被误用、也最容易被遗漏释放的对象。它内部封装了从张量描述推导出来的全部下发信息对调用者来说是一个不透明黑盒你不能修改它也不能窥探内容只能在执行接口里原样传回去。要理解 executor 的价值可以类比成外卖订单第一阶段是下单店家把食材、做法、配送路线全打包好了第二阶段是骑手取单配送不用重新核对订单。如果一个算子在循环里被调用一万次理想状态下第一阶段只做一次剩下九千九百九十九次都只是执行省下的 host 侧开销非常可观。不过 executor 是有生命周期的对象不是裸指针用完就丢。所有 ops-nn 算子的 aclnn 接口都提供了对应的销毁接口忘记调用就会造成 host 侧资源缓慢增长。我见过不少业务侧开发把执行接口封装成一个函数每次调用都重建 executor却忘了销毁跑一个长任务显存和内存双双上涨排查半天才定位到是句柄泄漏。3.2 workspace 的分配与对齐要求workspace 必须分配在设备侧全局内存上不能使用普通堆内存或者统一寻址内存。原因在于算子内核运行在设备上它通过设备地址直接访问这块临时空间任何 host 侧概念上的内存都替代不了。分配时最直接的方式是调用平台提供的显存分配接口。有一点值得注意平台分配接口本身通常已经做了地址对齐所以即使查询接口返回的 size 不是对齐值裸分配接口也能保证起始地址满足内核访问要求。但如果你自己搞了一个内存池从一块大显存里按偏移切小块就必须自己保证对齐否则内核访问未对齐地址轻则性能下降重则直接报错。另外我在实践中发现单一算子的临时空间需求往往不大但如果一个图里几十个算子串行执行每个算子都单独分配一块 workspace显存碎片和分配开销会被放大。更好的做法是维护一个按需增长的 workspace 内存池执行前统一取用、执行后统一归还这样整个流程的显存水位是平稳的。具体到 ops-nn 算子开发可以参考这个思路查询阶段拿到 size 后从池里取一块不小于 size 的内存执行完立即归还而不是持有到进程结束。3.3 异步语义与同步点的正确位置执行接口是异步的这是很多从 CPU 编程转过来的开发者最容易搞混的地方。调用执行接口返回成功只代表任务已经合法提交到指定 stream 上不代表算子已经跑完更不代表输出张量里的数据已经就绪。要拿到正确结果必须依靠 stream 上的同步机制。常见做法是调用同步接口阻塞等待该 stream 上所有任务完成或者给 stream 挂一个回调在回调里再读取输出。这里有一条重要经验不要在每个算子执行后都立刻同步那样等于把异步执行退化成同步执行性能会大打折扣。合理的做法是连续提交一批相互依赖的算子任务在最后一个关键输出前设置一个同步点。3.4 销毁顺序先释放执行资源再释放内存很多新手在资源释放顺序上栽过跟头。标准顺序应该是先确保任务执行完成同步或回调再释放 workspace 内存最后销毁 executor。如果反过来先销毁 executor 或者先释放 workspace而异步任务还在队列里等着用轻则报无效地址错误重则数据被覆盖成全随机数属于最难排查的一类故障。我习惯把这三类资源的生命周期全部收敛到同一个作用域里管理要么用 RAII 封装要么用统一的释放函数。这样做的好处是无论主流程怎么 return、怎么抛异常资源释放逻辑都不会漏。后面我会给一个可以直接照抄的模板。4. 一个可以直接照抄的标准调用模板4.1 模板代码两阶段合二为一把上面这些原则落到代码上我平时在 ops-nn 算子开发里使用的基础模板长这样// 第一阶段查询 workspace 大小并生成 executor aclnnStatus ret aclnnXxxGetWorkspaceSize( inputDesc, outputDesc, workspaceSize, executor); if (ret ! ACLNN_SUCCESS) { // 记录错误上下文方便排查 return ret; } // 分配 workspacesize 为 0 时也分配最小单位规避空指针路径 void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); if (ret ! ACL_SUCCESS) { aclnnDestroyExecutor(executor); return ret; } } // 第二阶段真正执行 ret aclnnXxx(executor, workspaceAddr, workspaceSize, stream); if (ret ! ACLNN_SUCCESS) { if (workspaceAddr ! nullptr) { aclrtFree(workspaceAddr); } aclnnDestroyExecutor(executor); return ret; } // 等待执行完成再统一释放资源 ret aclrtSynchronizeStream(stream); if (workspaceAddr ! nullptr) { aclrtFree(workspaceAddr); } aclnnDestroyExecutor(executor); return ret;这个模板是我在真实项目里反复使用的骨架它的核心思想是每一处可能提前返回的错误分支都必须保证已创建出来的资源被完整释放。你会发现这里没有任何一行是多余的——workspaceSize的判断、executor 的创建与销毁、stream 的同步每一样都是两阶段机制里绕不开的环节。4.2 模板里每个环节在防什么注释里的每一句话背后都有血泪教训。size 为 0 时也分配最小单位这一条就是因为我遇到过执行接口在这种情况下仍然会读取 workspace 指针的场景传入空指针直接导致非法地址访问。分配时选ACL_MEM_MALLOC_HUGE_FIRST是希望优先申请大块连续显存避免频繁触发碎片整理对于小 workspace 来说这个标志对性能几乎没有影响但对于反复分配的场景它能显著减少显存碎片。错误分支里的双重释放问题也值得强调。第一次分配失败时不需要释放 workspace因为地址还是空的但如果第一次分配成功了后面执行接口失败就需要同时释放 workspace 和执行器。这个模板在每个分支都检查了地址是否为空从根上避免了重复释放导致的二次崩溃。同步放在释放之前是为了保证异步队列里的任务确实用完了这块 workspace。不要小看这个顺序我见过有同事为了省一次同步直接释放 workspace结果下一个算子拿到的临时空间里全是上一个算子的残留数据跑出来的结果时对时错非常难查。4.3 从模板到工具函数每次写算子都复制这一大段模板效率并不高。实际开发中我把它封装成一个模板化的工具函数把两阶段调用收敛到一行调用里template typename GetSizeFn, typename ExecFn aclnnStatus InvokeAclnnTwoPhase( GetSizeFn getSize, // 第一阶段 lambda ExecFn exec, // 第二阶段 lambda void* stream, bool needSync true) { size_t workspaceSize 0; aclnnExecutor executor nullptr; auto ret getSize(workspaceSize, executor); if (ret ! ACLNN_SUCCESS) return ret; void* workspaceAddr nullptr; if (workspaceSize 0) { ret aclrtMalloc(workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); if (ret ! ACL_SUCCESS) { aclnnDestroyExecutor(executor); return ret; } } ret exec(executor, workspaceAddr, workspaceSize, stream); if (ret ! ACLNN_SUCCESS) { if (workspaceAddr) aclrtFree(workspaceAddr); aclnnDestroyExecutor(executor); return ret; } if (needSync) { ret aclrtSynchronizeStream(stream); } if (workspaceAddr) aclrtFree(workspaceAddr); aclnnDestroyExecutor(executor); return ret; }用这个工具函数后接入一个新算子只需要写两个 lambda一个负责查询、一个负责执行剩下的分配、同步、释放全部统一处理。我在负责的 ops-nn 算子接入工作中就是用这个工具把几十个算子的 host 侧调用代码全部统一化出问题的概率大幅下降。5. 高频故障排查从现象倒推根因的完整链路5.1 现象一换个 shape 程序就崩这是两阶段调用里最经典的一种故障。现象是用某个固定的输入 shape 反复验证都没问题一旦把输入张量的 shape 放大或者缩小程序立刻崩溃报的是设备侧内存非法访问。排查链路的第一步先确认是不是 workspace 不够用。很多人会在开发期把 workspaceSize 打印出来用固定值硬编码到分配逻辑里比如之前的算子都只要 4KB我就固定分配 4KB。一旦输入 shape 变大真实需要的 workspace 超过这个固定值内核写超边界崩溃就是必然的。正确做法是永远以查询接口返回的 size 为准不要自己估算。如果你确实怀疑查询结果有问题可以在分配时故意多分配一倍空间看崩溃是否消失——如果消失了说明问题出在 size 计算环节而不是内核逻辑。这一步隔离法能快速缩小排查范围。5.2 现象二多线程下偶发错误单线程完全正常另一种高频故障是并发场景下的偶发错误。两个线程同时执行同一个算子各自的查询与执行都对着同一块逻辑但时不时出现结果错误或者设备报错单线程压测又完全正常。这种问题十有八九出在 workspace 的共享使用上。如果两个线程从同一个内存池里取 workspace而内存池的取用逻辑没有做好线程安全或者两个线程拿到了同一块地址那么先执行的算子还没跑完后执行的算子就把这块空间覆写了。结果就是数据互相踩踏错误随机出现。排查时先看 workspace 的分配来源。如果是各自线程内部分配、互不共享那就排查 executor 是否被多个线程共用。executor 本身是否线程安全取决于平台实现稳妥做法是每个线程维护自己的 executor 和 workspace不要跨线程共享。我在团队里定过一条规矩两阶段产生的所有资源都归属当前线程绝不跨线程传递这基本能消灭这一类偶发问题。5.3 现象三显存和内存稳步上涨长跑任务里显存和 host 内存随着迭代次数稳步上涨哪怕每次迭代的算子调用看起来都正常。这类问题最容易被归因于平台泄漏实际上大概率是调用端没有正确销毁两阶段产生的资源。排查思路是先确认 executor 是否每次调用都被销毁。可以在迭代循环里加一个日志打印 executor 指针或者对象数量看它的数量是否持续增长。如果 executor 一直在增加说明执行接口跑完后销毁逻辑没有被走到常见原因是提前 return 的错误分支里漏写了销毁语句。另一个嫌疑对象是 workspace如果每次都重新分配而不释放显存水位同样会涨。修复方法是把我上面给的工具函数用起来把查询、执行、同步、释放全部包进去让资源生命周期跟函数作用域绑定。从根上杜绝某条分支忘记释放这种低级但极难查的问题。5.4 排查工具与套路排查两阶段调用问题我常用的套路是先看错误码再看报错位置。错误码能直接告诉你是 host 侧解析失败、设备侧执行失败还是参数非法。如果错误码指向设备侧再查 workspace 大小和地址是否合法。第二个套路是开平台提供的调试日志把调度信息和内核下发信息都打出来。两阶段调用里executor 内部到底选了什么算法、分配了多大 workspace日志里都能看到。我排查时最喜欢看的是查询返回的 size和实际分配的大小这两个数字一旦发现不一致问题基本就定位了。还有一个比较实用的方法把自己怀疑的环节分别隔离。比如固定 shape 连续跑一万次看是否稳定或者故意把 workspace 放大一倍看是否缓解。这种二分法的思路虽然朴素但对两阶段调用这种前后两个接口、中间一块内存的结构非常有效能快速区分是第一阶段算错了、还是第二阶段用错了。6. 效率翻倍把两阶段调用沉淀成开发习惯6.1 用 RAII 封装消除心智负担两阶段调用本身不难真正的开发负担在于那一大段单调的分配、同步、释放代码。我建议团队里所有 ops-nn 算子接入都走统一的 RAII 封装把 executor 和 workspace 的生命周期绑定到对象析构上让调用方完全不需要关心两阶段背后的资源管理细节。实际效果非常明显接入一个新算子的 host 侧代码从几十行模板收缩成几行业务逻辑出错概率也直线下降。开发人员可以把精力放在算子本身的 tiling 和内核优化上而不是反复跟内存分配做斗争。这比我见过的那种每个算子都复制一遍完整模板、改改函数名就上线的做法开发效率至少翻倍。6.2 开发期开启日志与统计我建议在开发阶段给两阶段调用加一层统计日志记录每个算子的调用次数、executor 复用次数、workspaceSize 分布。这些数据看起来不起眼实际上价值很大一是能帮你发现哪些算子在频繁重建 executor提示你可以把查询分配缓存起来二是能直观看到 workspace 的内存水位为显存规划提供依据。上线前再把日志关掉或者降级为 debug 级别不会给生产环境带来额外开销。这个习惯我保持了很久每次性能分析或者故障排查时这些日志都是第一手证据比事后猜要高效得多。6.3 用用例矩阵覆盖 shape 场景两阶段调用最容易翻车的地方在于 shape 变化所以我在开发时一定会准备一个 shape 用例矩阵覆盖极小 shape、极长 shape、多维广播 shape、以及连续调用场景。每次改完算子内核第一件事就是跑一遍这个矩阵确保查询接口返回的 workspaceSize 在各种 shape 下都能正确满足内核需求。这套用例矩阵还可以顺带验证异步语义在每个用例后随机加同步点确认结果一致再对比有同步和无同步两种模式下的输出基本能排除数据就绪顺序的隐患。工作量不大但能把两阶段调用相关的回归问题一次兜住。6.4 团队协作中的统一约定两阶段调用机制理解得不深很容易在团队里出现风格迥异的调用代码有人 src 错误分支忘释放有人 size 拍脑袋填有人把同步放在每个算子后拖累性能。我给团队定过几条硬约定统一用工具函数入口调用算子、所有资源归属调用线程、任何算子不允许缓存 size 数字、执行接口成功后按场景决定同步点位置。约定建立起来后code review 基本不用再为这些低级问题争论大家可以把时间花在更有价值的性能优化上。最后再分享一个我在实践中感触很深的点两阶段调用看似只是 aclnn 的一个接口设计但它背后计算与执行解耦、资源与任务分离的思想几乎是理解整个算子平台调度机制的钥匙。把这两阶段的边界吃透后续看内存池设计、算子融合、图调度都会顺畅很多。搞懂它绝对不亏。
阅读完成 · 觉得有帮助？
