Skip to content

[Feature] Move plan memory / insert sync from memref world to tile_buf world #535

@Zhendong404

Description

@Zhendong404

Suggested label: enhancement

Summary

当前 PTOAS pipeline 将 plan memoryinsert sync 放在 View2Memref 之后执行:

[tile_buf world]
  (op fusion plan -> gen fusion region -> ...)
-> View2Memref
-> [memref world]
  (plan memory -> insert sync -> expand tileop -> ...)

这个 pass 边界在早期帮助我们复用 memref-world 的实现,但现在已经开始限制 tile-level 能力演进。我们希望将 plan memoryinsert sync 以及与之强相关的大部分中高层 pass 前移到 tile_buf world,把 View2Memref 明确收敛为高层 Tile IR 到低层 memref IR 的边界。

目标 pipeline 形态:

[tile_buf world]
  (plan memory -> insert sync -> op fusion -> expand tileop)
-> View2Memref
-> [memref world]
  (low level optimizations -> backend lowering -> ...)

这不是单纯的 pass 重排,而是为了让同步、融合和 TileOp 展开都发生在语义最完整的表示层上。

Motivation / use case

1. 当前 pass 边界直接阻塞更高等级的 op fusion

我们希望在 tile_buf world 做 op fusion,因为这个阶段保留了最完整的 tile 语义信息,包括:

  • tile 的布局和 config
  • valid shape
  • tile-level producer / consumer 关系
  • 更贴近上层语义的 alias / buffer 使用形态

level 3 的 op fusion 实际上依赖 insert sync 的结果,因为融合不能跨越 sync point。当前 insert sync 位于 memref world,意味着 tile_buf world 在做融合决策时看不到最终同步边界,只能停留在不依赖 sync 结果的 level 2 fusion。

换句话说,现有 pipeline 的问题不是“还没实现 level 3 fusion”,而是“当前 IR 分层方式让 level 3 fusion 无法在正确的位置实现”。

2. 当前 pipeline 让 expand tileop 之前发生了一次不必要的类型世界往返

ExpandTileOp 的模板函数接口本身就是 Tile / tile_buf。仓库里的 pass 描述也已经明确:

  • ExpandTileOp 生成的模板函数使用 tile_buf 参数
  • MemrefToTileBuf 的存在,是为了在 PTOViewToMemref + PlanMemory + InsertSync 之后重新恢复 tile_buf 语义,再继续做 Tile-level lowering

这意味着当前 pipeline 实际上在 expand tileop 前经历了:

tile_buf -> View2Memref -> memref -> MemrefToTileBuf -> tile_buf

这条链路带来几个问题:

  • 需要额外维护一条 memref -> tile_buf 的恢复路径,只是为了回到 TileOp lib 所需的接口形态
  • 中间经过一次语义降级后再恢复,增加了类型、metadata、anchor 和 verifier 一致性的维护成本
  • 新增 TileOp / TileBuf 能力时,需要同时考虑 lowering 和 recovery 两侧,增加系统脆弱性

如果 plan memory / insert sync 本来就是中高层语义 pass,那么让它们停留在 tile_buf world,可以避免这次“先降级再恢复”的往返。

3. 现在的分层方式已经开始影响 PTOAS 的功能完备性和稳定性

上述两个问题并不是孤立的工程不便,而是在持续影响 PTOAS 的能力边界:

  • 功能完备性:sync-aware fusion、更多 tile-level 优化和更自然的 TileOp 展开都会受到限制
  • 稳定性:额外的 IR world 往返意味着更多 bridge、更多隐式约束和更多恢复失败的可能性

随着 tile-level 功能继续扩展,把中高层 pass 长期放在 memref world 的成本只会越来越高。

Proposed API / behavior

建议将 tile_buf worldmemref world 的职责边界重新定义为:

  • tile_buf world:负责高层语义相关的 memory planning、sync insertion、fusion、tileop expansion
  • memref world:负责低层 canonicalization、backend-oriented optimization 和最终 lowering

具体上,希望逐步达成以下结果:

  1. plan memorytile_buf world 上工作,直接消费 Tile-level buffer / alias 信息,而不是等到 View2Memref 之后再规划。
  2. insert synctile_buf world 上工作,使 sync point 成为 fusion 和后续高层优化可见的边界。
  3. op fusion 在拿到 sync 信息后执行,能够支持真正受 sync point 约束的更高等级融合。
  4. expand tileop 直接在 tile_buf world 完成,避免为了调用 TileOp lib 再从 memref 恢复回 tile_buf
  5. View2Memref 作为更清晰的阶段边界,只承接已经完成高层决策的 Tile IR。

如果一次性迁移过大,也可以接受分阶段演进,但最终目标应当是让 MemrefToTileBuf 不再承担“为高层 pass 回补语义”的核心职责。

Additional context

仓库里已经存在一些信号,说明当前 pass 分层正在承受额外复杂度:

  • ExpandTileOp 的描述明确指出模板函数使用 tile_buf 参数,并直接以 tile_buf 作为调用接口
  • MemrefToTileBuf 的描述明确指出它是在 PTOViewToMemref + PlanMemory + InsertSync 之后恢复 tile_buf,以便后续 Tile->Vector lowering 继续基于 tile_buf 语义工作
  • PTOViewToMemref 还需要专门插入 pto.bind_tile 作为锚点,供后续恢复 tile_buf 类型使用

这些都说明当前 pipeline 已经不是“自然地在 memref world 完成后续流程”,而是“先降到 memref world,再想办法恢复出高层 Tile 语义继续工作”。

建议把这个 issue 作为 pipeline 分层调整的总入口,后续再拆分出更细的 implementation issue,例如:

  • tile_buf-world memory planning 的表示和地址回填方式
  • tile_buf-world sync analysis / sync insertion
  • sync-aware level 3 fusion
  • ExpandTileOp 前移后的 pipeline 收敛
  • MemrefToTileBuf 的降级、收缩或移除

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    Status

    Backlog

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions