深度解读昇腾 CANN 小 shape 算子计算优化技术，进一步减少调度开销

摘要：Host 调度模式下，GE 将模型中算子的执行单元划分为 Host CPU 执行与 Device（昇腾 AI 处理器）执行两大类。

本文分享自华为云社区《深度解读昇腾CANN小shape算子计算优化技术，进一步减少调度开销》，作者：昇腾 CANN。

GE（Graph Engine）将模型的调度分为 Host 调度与下沉调度两种模式。经过上期的介绍我们知道，在模型为静态 shape 时，由于其输入 tensor shape 固定不变，在编译时就能确定所有算子的输入输出 shape，并能提前完成模型级内存编排、tiling 计算等 Host 调度工作，因此采用模型下沉调度方式可以将整个模型下沉到 Device 侧执行，从而提升模型调度性能。

与之对应的，在模型为动态 shape 的情况下，由于输入 tensor shape 不确定，需要在上一个算子完成 shape 推导后，才能确定下一个算子的输入 shape 等信息，因此无法将整个模型下沉执行，只能采用 Host 调度模式。

1 Host 调度简介及优化背景

所谓 Host 调度，是指模型的调度主体位于 Host CPU，由 CPU 完成逐算子调度。一个算子的调度任务为 kernel 执行准备必要参数，通常包含 shape 推导、tiling、内存分配、launch 等。

Host 调度模式下，GE 将模型中算子的执行单元划分为 Host CPU 执行与 Device（昇腾 AI 处理器）执行两大类。对于卷积、MatMul 等对算力要求高的算子，会被划分到 Device 执行；而由于 shape 信息在 Host CPU 维护，Shape、Reshape 等算子更适合被划分到 Host CPU 执行；除此之外，还有一些算子，在 shape 较小时，计算量也很小，调度开销往往大于算子的实际计算开销，就需要考虑如何尽可能减少调度开销带来的性能影响。

图 1 网络拓扑片段

图 1 是一段网络拓扑片段示例，按照一般的调度机制，Gather、Concat 算子会下沉到 Device 侧计算，Shape、Unsqueeze、Reshape 算子在 Host 侧计算。其执行时序如图 2 所示，模型 E2E 执行耗时除了包含算子计算的时间外，还包含 Host 与 Device 之间的数据拷贝、算子下沉调度、Stream 同步等开销，整体执行 E2E 耗时在毫秒级别。

图 2 优化前执行时序

而对于小 shape（如 shape size 小于 8）的 Gather、Concat，算子本身在 Host 侧 CPU 的计算开销上仅微秒级别，与 Device 侧计算的性能相差无几。此时下发带来的额外开销就显得比较明显。针对上述这种 shape 较小且输入 Tensor 内存在 Host 的场景，GE 识别将这部分算子保留在 Host 侧执行，可有效减少调度开销带来的性能影响。

2 小 shape 算子计算优化实现

在图编译流程执行到引擎选择之后，GE 选择在 Host 侧执行的算子并将其作为锚点，然后向后递归查找计算数据个数小于 8 的算子，并将这些算子的执行引擎修改为 Host CPU。针对图 1 所示的网络片段，假设 shape 算子的输出的 shape size 小于 8，则 Gather、Concat 算子的执行引擎都会被刷新成 Host CPU。优化后执行时序如图 3 所示，此时模型执行只有算子计算带来的开销，经测试约为 10 微秒（3ms –> 10us），显著的提高了 E2E 执行性能。

图 3 优化前后执行前后时序对比

3 优化效果

以 LLaMA2 大语言推理模型为例，符合上述执行引擎刷新的算子有 Pack、Gather、Concat 等约 650+个，刷新前模型 E2E 耗时约 1.062S，刷新后执行时间优化到了 1.009S，吞吐提升 5%。

4 更多介绍

GE 小 shape 算子计算优化技术的相关介绍就到这里，欢迎大家关注后续技术分享。如需获取更多学习资源请登录昇腾社区。

点击关注，第一时间了解华为云新鲜技术~