算子融合

本文档从 kernel 组织方式、中间数据搬运和片上资源使用等角度,说明 PTO Tile Lib 中与算子融合相关的考虑因素。

本文档重点讨论能够通过 PTO 编程模型直接表达的融合机会。

1. PTO 中的算子融合

在 PTO 语境下,算子融合通常指:在合法且有收益的前提下,把多个阶段的工作组织在同一个 kernel 或同一个 tile 级计算流程中,从而减少不必要的中间 GM 读写。

常见目标包括:

  • 减少中间结果的 GM 读取和写回
  • 提高片上数据复用
  • 改善 load、transform、compute、store 等阶段的稳态衔接

这种理解与 docs/coding/opt.md 中的整体优化思路是一致的。

2. Fusion characteristics

The following characteristics are typical of fusion in PTO kernels:

  • 在一个 kernel 中组合多个逻辑阶段,通常可以减少中间内存流量
  • 是否适合融合,取决于 tile 布局、valid region 约束、backend 支持情况以及可用片上资源
  • 融合只有在保持正确性并真正改善瓶颈时才有意义
  • 仓库中的一些高性能 kernel,本身就会把多个阶段组织在同一个 kernel 中,而不是把每个中间结果都写回 GM

相比之下,把固定加速比、未公开的编译器融合 pass 或虚构 fused API 写成 PTO 的既定能力,就不够严谨。

3. 在开发者编写的 kernel 中如何理解融合

在本仓库中,更合适的理解方式是:融合首先是一种 kernel 组织方式,而不是一个已经自动保证存在的编译器特性。

也就是说,开发者可以:

  • 在同一个 tile 级 kernel 中保留多个有依赖的操作阶段
  • 在中间结果能够留在片上时,避免把它们写回 GM
  • 让中间 tile 直接进入后续阶段,而不是被拆成多个独立 kernel

例如,一个按行归一化或 attention 相关的 kernel,就可能自然地把以下阶段放在同一个 kernel 中:

  • load
  • reduction
  • elementwise transform
  • normalization
  • store

4. 融合时真正需要考虑的问题

4.1 中间结果的存储代价

如果某个中间结果原本需要写回 GM,然后很快又被下一阶段读回,那么融合通常就有吸引力。

如果中间 tile 可以一直保留在片上并直接供下一阶段使用,GM 流量就可能明显下降。

4.2 片上资源限制

融合并不是零成本的。

一个融合后的 kernel 可能需要更多:

  • tile buffer
  • 临时 tile
  • 同步边
  • 布局转换

如果融合后显著增加了片上存储压力,结果可能反而更差。

4.3 指令合法性与布局约束

即使一个数学表达式在概念上很容易融合,PTO 实现仍然必须满足真实的指令约束:

  • tile type 是否满足指令要求
  • 布局是否对参与的操作都合法
  • valid region 处理是否仍然正确
  • 相关指令是否在目标 backend 上可用

因此,每次尝试融合时,都应结合 docs/isa/include/README.md 进行核对。

4.4 围绕瓶颈做判断

是否融合,应围绕瓶颈来判断。

如果一个 kernel 的主要瓶颈是 GM 流量,那么减少中间结果写回通常会更有帮助。

如果主要瓶颈是 compute 或 transform 成本,那么更有效的改进手段可能是调整 tiling、重叠方式或布局选择。

5. 与流水线重叠的关系

融合与流水线重叠相关,但并不相同:

  • 融合:减少不必要的中间结果落地和阶段拆分
  • 流水线重叠:让仍然保留的阶段之间更高效地重叠执行

一个融合后的 kernel 依然需要良好的流水线组织。很多情况下,更好的结果来自于同时做到:

  • 减少中间数据的 GM 往返
  • 对剩余阶段使用更合理的缓冲和同步结构

可进一步参考:

  • docs/coding/pipeline-parallel.md
  • docs/coding/Event.md
  • docs/coding/opt.md

6. 本文档的适用边界

除非仓库代码或正式文档明确给出依据,否则以下说法都不够严谨:

  • 把某个固定内核启动开销量写成 PTO 的通用常数
  • 直接写出“100% 缓存命中”这类绝对结论
  • 在没有 shape 和 backend 证据的情况下直接写“3× 加速”
  • 使用不属于 PTO 公共 intrinsic 的虚构 API 或占位式指令
  • 把未正式公开的自动融合能力写成当前稳定编译器保证

这些表述在非正式讨论中可以帮助说明直觉,但不适合作为严格的仓库文档内容。

7. 融合流程

一个更实用的流程是:

  1. 先从正确的非融合或最小融合 kernel 出发
  2. 确认中间 GM 流量是否真的是瓶颈
  3. 只保留那些既合法又有收益的组合阶段
  4. 重新检查 tile 约束、backend 支持和同步关系
  5. 先验证正确性,再在代表性 shape 上比较性能

这样做可以让融合决策建立在实际收益上,而不是只依赖直觉。

8. 结语

在 PTO Tile Lib 中,更准确的算子融合表述应当是:

  • 它是一种减少中间 GM 流量的方式
  • 它首先是一种 kernel 组织与优化手段
  • 它受到 tile 合法性、backend 支持情况和片上资源约束

相比把推测性的 fused API、固定加速比或未公开的编译器自动化能力写成稳定接口,这样的描述更符合当前仓库事实,也更适合作为严谨文档。