diff --git a/docs/designs/vpto-pipe-tilelib-lowering-design-zh.md b/docs/designs/vpto-pipe-tilelib-lowering-design-zh.md new file mode 100644 index 0000000000..2e87a48cd1 --- /dev/null +++ b/docs/designs/vpto-pipe-tilelib-lowering-design-zh.md @@ -0,0 +1,390 @@ +# VPTO Pipe PTODSL 与 TileLib 边界设计 + +**Issue:** [#966](https://github.com/hw-native-sys/PTOAS/issues/966) +**状态:** 编译器侧实现完成;TileLib 实现作为独立后续工作 +**目标:** A5 VPTO 后端 +**本次变更负责:** PTODSL pipe 表层校验、PTO IR、PTOAS pass、 +`ExpandTileOp` 契约生产、VPTO LLVM lowering 和聚焦的前端覆盖 +**本次不负责:** `ptodsl/tilelib/**`、TileLib template 测试,以及 +TileLib-ST/runtime 覆盖 + +## 1. 目的与范围 + +PTODSL 已可将表层 pipe 抽象 lower 为统一的 PTO pipe operation。VPTO 路径不能 +保留 EmitC 所使用的 C++ `TPipe` object,因为在进入 LLVM lowering 前,VPTO 需要 +在 PTO IR 中显式表示 FIFO 状态、同步和数据搬运。 + +本文记录当前已经落地的 PTOAS 编译器侧边界。它同时是编译器侧实现设计与 TileLib +owner 的交接契约。本文刻意不宣称:在没有相应 TileLib 改动的源树中,GM 或 local +split pipe 已可端到端运行。 + +稳定的前端到后端流程如下: + +```text +PTODSL Pipe object + -> frontend pipe operation + -> 统一 PTO pipe operation + -> 推导/校验 pipe nosplit 配置 + -> materialize PipeState 与结束 tdrain + -> 内存规划和 reserved-buffer 解析 + -> pipe metadata/candidate discovery + -> ExpandTileOp PipeSpec ABI + -> 外部 PTODSL TileLib template + -> VPTO LLVM lowering +``` + +## 2. 目标 + +- 保持已有 PTODSL Pipe API 及其逐 transaction 的 `split` 值。 +- 为 A5 VPTO 路径显式表示可变 FIFO state。 +- materialize terminal producer cleanup,且不从任意一条 `tpush` 推导其策略。 +- 在内存规划后,向 `ExpandTileOp` 提供完整、已解析的 pipe metadata 与资源 operand。 +- 在 A5 VPTO backend 之外,保持 EmitC 的 `TPipe` 路径不变。 +- 提供收敛的 ABI,使独立开发的 TileLib 可消费该信息,而不依赖 opaque + `!pto.pipe` value。 + +## 3. 非目标 + +- 替换既有 EmitC pipe 实现。 +- 新增一套并行的 PTODSL `pipe.push()` / `pipe.pop()` API,或重新设计已有 API。 + `pto.pipe.c2v(...)`、`pto.pipe.v2c(...)` 和 `pto.pipe.bidirectional(...)` + 已返回带有这些方法的 Pipe object。 +- 在本次变更中修改 `ptodsl/tilelib/**`。 +- 在 TileLib 实现和 TileLib-ST/simulator 覆盖合入前,宣称 GM/global 或 local + split 的运行时正确性已经完成。 +- 支持 A2/A3 VPTO pipe expansion、`v2c_ctrl`、fixpipe quantization 或 + phase/NZ conversion。 + +## 4. PTODSL Pipe 表层 + +### 4.1 已有公开 API + +已有 constructor 返回带方向信息的 Pipe object: + +```python +pipe = pto.pipe.c2v( + id=7, + gm_slot_tensor=slots, + slot_size=1024, +) +pipe.init_cube() +pipe.init_simd() +entry = pipe.alloc(split=1) +pipe.push(entry, split=1) +entry = pipe.pop(split=1) +pipe.free(entry, split=1) +``` + +`v2c(...)` 提供相同的 transaction 方法。`bidirectional(...)` 在执行带方向的 +transaction 前通过 `.c2v`、`.v2c` 暴露 endpoint。Pipe API 要求稳定、显式的 `id`。 + +公开方法直接 lower 到相应 frontend operation: + +| Pipe 方法 | lower 后的 frontend operation family | +|---|---| +| `alloc(split)` | `talloc_to_aiv` 或 `talloc_to_aic` | +| `push(entry, split)` | `tpush_to_aiv` 或 `tpush_to_aic` | +| `pop(split, ...)` | `tpop_from_aic` 或 `tpop_from_aiv` | +| `free(entry, split)` | `tfree_from_aic` 或 `tfree_from_aiv` | + +`split` 是 operation property,不是 Pipe-object property。因此同一个 +split-capable pipe 可以先执行 `split=1` transaction,再执行 `split=2` +transaction,但仍须满足既有 pipe configuration 校验。 + +### 4.2 GM slot-size 规则 + +对于 GM-entry pipe,只有 `nosplit=True` 时才能无歧义地推导 `slot_size`:此时一个 +`gm_slot_tensor` shape 就是一个完整 FIFO slot。对于 split-capable GM pipe,调用方 +必须显式传入完整 slot 的 `slot_size`。这避免把 split subregion 误认为 FIFO slot size。 + +### 4.3 Pipe 操作与 template 边界 + +PTODSL 源码和 frontend helper 使用公开的 `Pipe` object。其完整 operation surface +如下: + +| 类别 | 公开 surface | 契约 | +|---|---|---| +| constructor | `pto.pipe.c2v(...)`、`pto.pipe.v2c(...)`、`pto.pipe.bidirectional(...)` | 创建 C2V、V2C 或双向 logical pipe。`id` 必填,且始终是稳定的 pipe identity。 | +| initialization | `init_cube()`、`init_simd()` | 从 Cube 或 SIMD 侧初始化 pipe。双向 pipe 使用 root object 初始化。 | +| producer transaction | `alloc(split=0)` | 仅适用于 global-entry pipe;返回下一个 FIFO entry 的 descriptor。 | +| producer transaction | `push(entry, split=0)` | 将填充完成的 global entry 或 local tile 发布给 consumer。 | +| consumer transaction | `pop(split=0, result_type=None, valid_shape=None, valid_row=None, valid_col=None)` | 返回下一个 global-entry descriptor 或 local tile。local tile-entry pipe 必须提供 `result_type`;`valid_shape` 与 `valid_row` / `valid_col` 互斥。 | +| consumer transaction | `free(entry=None, split=0)` | 释放已消费的 FIFO entry。global-entry pipe 必须传入相应 `pop` 返回的 entry;local tile-entry pipe 可省略。 | +| read-only property | `id`、`slot_size`、`entry_type` | 暴露 compile-time identity、完整 logical slot 的 byte size,以及适用时的 global-entry descriptor type。 | + +`c2v` 和 `v2c` 只暴露方向合法的 transaction。双向 pipe 必须在调用 `alloc`、 +`push`、`pop` 或 `free` 前选择 `.c2v` 或 `.v2c`;其 root object 没有无歧义的 +transaction direction。`split` 是逐 transaction 的 compile-time value,不是可变的 +Pipe-object state:`0` 表示不切分,`1` 表示上下切分,`2` 表示左右切分。 + +这些是 PTODSL API,并非 TileLib template API。template expansion 时,template 不会 +收到 Python `Pipe` object,也不能调用其方法。PTOAS 通过第 7 节 ABI 提供每条 +operation 的 `PipeSpec`、有序 `PipeResources`、共享 `PipeState`、存在时的 entry,以及 +可选 AIV subblock value。TileLib owner 根据这些 ABI value 实现 transaction 的 FIFO +address、synchronization 和 counter 行为;本次变更既不新增 template-side Pipe wrapper, +也不修改 `ptodsl/tilelib/**`。 + +## 5. A5 VPTO 默认行为与兼容性 + +不存在 pipe expansion feature flag。当包含 frontend 或 unified pipe transaction 的 module +使用 `--pto-arch=a5 --pto-backend=vpto` 编译时,driver 会自动执行 pipe-specific validation、 +PipeState materialization、candidate discovery 和 expansion preparation。默认 TileLib backend +是 PTODSL,因此默认路径使用 PTODSL metadata 和 expansion。 + +PipeState materialization pass 只插入到 A5 VPTO backend。所有 EmitC 路径(包括 A5 +EmitC)仍使用既有 `TPipe` lowering,且不会 materialize PipeState 或 terminal `tdrain`。 +A2/A3 行为同样保持不变。 + +若所选 TileLib 尚未实现合法 pipe candidate,candidate discovery/expansion 会明确失败。 +PTOAS 不得静默回退到 C++ `TPipe` 或另一套 TileLib implementation。 + +## 6. PTO IR 契约 + +### 6.1 有状态 operation + +内部 `pto.talloc`、`pto.tpush`、`pto.tpop` 与 `pto.tfree` 均有可选的 +`pipe_state` operand。legacy source-authored IR 仍可不含该 operand,以保持既有 +EmitC 路径的 source compatibility。feature-owned VPTO IR 使用的类型固定为: + +```text +!pto.struct +``` + +verifier 拒绝其他 state shape。一旦某个 pipe 的任一 stateful user 带有 state,该 +pipe 的全部 `talloc`、`tpush`、`tpop`、`tfree` 和 `tdrain` user 都必须携带同一个 +SSA state。这样可以在 hand-authored internal IR 中拒绝部分 materialization 或各自 +分配 counter 的情况。`tpush`、`tpop` 和 `tfree` 使用 operand segment,使已有的可选 +entry/subblock operand 在 assembly form 中仍无歧义。 + +两个 field 的契约固定如下: + +| Field | 名称 | 含义 | +|---|---|---| +| 0 | `prod_index` | 下一个 producer FIFO position。 | +| 1 | `cons_index` | 下一个 consumer FIFO position。 | + +state 不包含 pipe handle、物理地址、`flag_base` 或 split mode。这些属于不可变的 +pipe configuration 或 runtime resource,会在 expansion 时单独提供。 + +### 6.2 结束 `tdrain` + +`pto.tdrain` 是只由 PipeState 路径插入的内部 operation: + +```text +pto.tdrain(%pipe, %state : !pto.pipe, !pto.struct) { split = <0|1> } +``` + +它表示此前由 EmitC `TPipe` object 生命周期执行的 producer-side cleanup。对于至少 +有一个 producer `tpush` 的 pipe,它在每个可达 `func.return` 前插入一次;没有 producer +的 pipe 不插入 drain。 + +`tdrain` 是 pipe-level cleanup。其 `split` 在 pipe configuration 解析后推导,不能从 +producer operation 推导: + +| 已解析 initializer `nosplit` | materialize 的 `tdrain.split` | +|---|---| +| `true` | `0` | +| `false` | `1` | + +因此 materialization pass 接受不同的 producer axis,例如 +`tpush(split=1)` 后再 `tpush(split=2)`。既有 infer/validate pass 仍负责拒绝非法的 +`split` / `nosplit` 组合。IR verifier 还要求 hand-authored `tdrain` 使用该精确推导 +的 split,并在 initializer 尚未解析 `nosplit` 时拒绝该 drain。 + +### 6.3 Materialization pass + +`pto-materialize-pipe-state` 在 `pto-infer-validate-pipe-init` 已解析 `nosplit` +之后,以每个 `func.func` 为单位运行。A5 VPTO driver 将其安排在 layout、fusion、 +memory planning 和 reserved-buffer resolution 之后、pipe-only candidate discovery 之前。 +这样既有 shared pass 继续消费原有 pipe IR,同时 pipe metadata 能同时获得 PipeState、 +已解析的 `flag_base` 与资源。 + +对于每个有 stateful user 的 initialized pipe,它会: + +1. 创建一个 `!pto.struct` 类型的 `pto.declare_struct`。 +2. 用 `pto.struct_set` 将两个 field 初始化为零。 +3. 将同一个 state 附加到该 pipe 的全部 `talloc`、`tpush`、`tpop` 与 `tfree` user。 +4. 标记 feature-owned pipe IR,供后续 expansion cleanup 使用。 +5. 只要该 pipe 有 producer,就插入 pipe-level `tdrain`。 + +initializer 必须支配每个需要插入 drain 的 return。pass 对不能表示的 nested/lifetime +形式给出诊断,而不是跨 region 移动 pipe handle。 + +## 7. Expansion ABI + +### 7.1 值的职责划分 + +`ExpandTileOp` 将 pipe 信息拆分为四个 logical operand: + +| Logical operand | 职责 | Physical helper argument | +|---|---|---| +| entry(存在时) | transaction 使用的 tile 或可变 GM descriptor | 是 | +| PipeSpec | 不可变配置和逐 operation 的 `split` | 否 | +| PipeResources | 有序 runtime address | 每个 resource 一个 argument | +| PipeState | 可变 producer/consumer counter | 一个 struct argument | + +`!pto.pipe` 不会成为 template helper argument。直接传递它会使 opaque pipe dependency +留在 VPTO lowering 中,违背 expansion 的目标。 + +对于 `tpush` 与 `tpop`,ABI 还包含 AIV subblock ID。原 operation 没有该 ID 时,PTOAS +会序列化值为零的 `i64` scalar,以保持 template logical signature 稳定。 + +### 7.2 PipeSpec producer 契约 + +在 `pto-resolve-reserved-buffers` 后,编译器从 pipe initializer 和单条 operation +推导 `PipeWireInfo`,并序列化如下 `pipe` operand: + +```json +{ + "kind": "pipe", + "init_kind": "l2g2l", + "dir_mask": 2, + "slot_size": 1024, + "slot_num": 8, + "local_slot_num": null, + "flag_base": 0, + "nosplit": false, + "split": 2, + "resource_names": ["gm_addr"] +} +``` + +`init_kind` 取值为 `l2l` 或 `l2g2l`。source direction encoding 仍是 PTO frontend +encoding(`1`、`2` 或 `3`);TileLib consumer 不得假设该 field 已被转成 +ISA-specific direction。 + +`split` 从每条 `talloc`、`tpush`、`tpop`、`tfree` 或 `tdrain` 复制而来,并纳入 +specialization identity,因此不同 split path 的 helper 不会错误复用。`tdrain` 因而 +拿到第 6.2 节确定的 pipe-level 值。 + +缺少 `flag_base`、无法解析 initializer、没有 PipeState,或使用不支持的 +`acc_push_epilogue`,都会在请求 helper 前报错。 + +### 7.3 Resources、state 和 entry + +`PipeResources` 按 `resource_names` 指定的精确顺序保存存在的资源: + +```text +gm_addr, local_addr, peer_local_addr +``` + +`l2l` pipe 可提供 local 和 peer-local resource。`l2g2l` pipe 还可提供 `gm_addr`; +不存在的 resource 不占用 helper argument。 + +`PipeState` 的序列化形式为: + +```json +{"kind": "pipe_state", "fields": ["i32", "i32"]} +``` + +已声明的 global entry 会序列化为 `pipe_entry`,而不是普通只读 view,因为 TileLib +可能需要重新绑定 caller-owned descriptor。其 helper ABI 是 `!pto.tensor_view<...>`。 +tile entry 保持原有 tile metadata。 + +### 7.4 Specialization 与 cleanup + +`ExpandTileOp` 使用 operation name、target、entry metadata、PipeSpec、resource +topology/type metadata、PipeState schema、subblock argument,以及 `kernel_kind` 等 +context attribute 构建 specialization key。它转发每条 operation 的 `split`,而不为 +整个 pipe 选择单一 split 值。 + +template expansion 成功后,PTOAS 只删除 use 已消失的 feature-owned unified pipe +operation 和 initializer。若 feature-owned pipe operation 仍存在,则 expansion 报错, +不回退到其他后端。 + +## 8. VPTO LLVM 边界 + +PipeState 必须跨 helper creation 和 inlining 存活,因此两个 VPTO LLVM emitter 都对 +受支持的 struct subset 进行 lowering: + +- `pto.declare_struct` lower 为 function-local storage; +- `pto.struct_get` lower 为 field address calculation 和 load; +- `pto.struct_set` lower 为 field address calculation 和 store。 + +pipe entry 路径还支持可变的 `!pto.tensor_view` descriptor: +`pto.declare_global`、`pto.tassign` 和 `pto.tensor_view_addr` 通过 descriptor +storage lower,不留下 memref bridge 或 VPTO pipeline 中的 +`unrealized_conversion_cast`。这是 TileLib ABI 的编译器支持,不是 FIFO address 或 +synchronization 语义的实现。 + +global pipe entry 可以直接使用 `pto.declare_global` 的结果,也可以使用从该声明 +经过一个或多个 `pto.tassign` 得到的结果。传给 pipe op 的仍是重绑定后的结果值。 + +## 9. TileLib 交接 + +以下工作明确属于 TileLib owner,不包含在此 PR: + +- 消费 `pipe`、`pipe_resources`、`pipe_state`、`pipe_entry` operand 的 metadata + class、renderer binding 与 candidate constraint; +- `talloc`、`tpush`、`tpop`、`tfree`、`tdrain` 的 A5 template; +- GM FIFO split address offset 与 local split/subblock constraint; +- operation-specific synchronization、FIFO index update 与 terminal drain 行为; +- daemon/unit coverage 以及 TileLib-ST 或 simulator 端到端覆盖。 + +TileLib 实现必须把 PipeSpec 视为不可变 configuration,按声明顺序消费 resource,并 +使用单条 operation 的 `split` 及可选 subblock 值。它不得从 `!pto.pipe` 恢复策略、 +重新推导 pipe-wide split,或将 `tdrain` 当作任意 producer transaction。 + +在该实现合入前,完整 pipeline 可能因没有合法 pipe template candidate,或 template +有意拒绝不支持配置而失败。这是预期的组件边界失败,不能据此认定 PTOAS pass 失败。 + +## 10. 诊断 + +编译器为下列边界违规提供可操作的失败信息: + +| 条件 | 诊断方向 | +|---|---| +| PipeState 类型非法 | 要求 `!pto.struct` | +| PipeState 关联不一致 | 要求同一 pipe 的全部 stateful user 共享一个 state | +| hand-authored `tdrain.split` 非法 | 要求使用由已解析 `nosplit` 推导的 split | +| drain 前 `nosplit` 未解析 | 要求在 materialization 前执行 pipe-init validation | +| drain lifetime 不支配 | 标识 initializer 与受影响 return | +| `flag_base` 未解析 | 要求 expansion 前执行 reserved-buffer resolution | +| pipe initializer/state 不可解析 | 标识对应 unified pipe operation | +| `acc_push_epilogue` 不支持 | 在 template invocation 前拒绝 | +| 没有 legal candidate 或遗留 feature-owned pipe IR | expansion 无回退地失败 | + +## 11. 验证与后续工作 + +### 11.1 已包含覆盖 + +PTODSL surface regression 覆盖显式 full-slot GM construction,以及两条独立的 +nonzero transaction(`split=1`、`split=2`)经过 `alloc`、`push`、`pop`、`free` 的 +路径。它检查每个 operation 收到自己的 split 值,而不是 pipe-wide default。 + +本次变更要求的本地检查为: + +```text +ninja -C build-local-vpto ptoas PTOPythonModules +ptodsl/tests/test_vector_cube_ops.py -v +``` + +当前 unit suite 通过 38 个测试,`git diff --check` 也通过。 + +额外的 focused compiler lit coverage 覆盖 PipeState materialization、PipeSpec RPC +payload(使用位于 `ptodsl/tilelib/**` 之外的 test-only mock daemon)、默认与显式 AIV +subblock operand、有序的 GM 和 local/peer-local resource list、可变 descriptor 的 +VPTO LLVM lowering、verifier 诊断、默认 A5 VPTO activation,以及 A5/A3 EmitC +compatibility。mock daemon 不代表 TileLib template 或 FIFO runtime 行为已经得到验证。 + +### 11.2 TileLib 必需的后续覆盖 + +TileLib 接入此 ABI 后,必须自行补充: + +- `split=0`、`1`、`2` 的 candidate selection 与 specialization separation; +- GM/local resource order、address offset 与 subblock handling; +- PipeState counter transition 与 terminal `tdrain` 行为; +- 不支持 direction/quantization 的诊断;以及 +- TileLib-ST/simulator 路径上的端到端 FIFO 正确性。 + +runtime suite 必须验证输出和同步行为。仅检查能编译或能渲染 IR,不能证明 +GM/local FIFO 语义正确。 + +## 12. 验收边界 + +当 frontend 保留 split、PipeState/`tdrain` materialization 遵循已解析的 `nosplit`、 +编译器产生本文档化的 expansion ABI,且 VPTO LLVM 可 lower 所需 state 和可变 +descriptor 形式时,PTOAS 侧工作完成。 + +只有独立的 TileLib 变更实现该交接契约,并以 TileLib-ST/simulator 覆盖证明支持配置 +后,完整功能才完成。PR 状态和 release claim 必须明确区分这两个里程碑。 diff --git a/docs/designs/vpto-pipe-tilelib-lowering-design.md b/docs/designs/vpto-pipe-tilelib-lowering-design.md new file mode 100644 index 0000000000..b165569002 --- /dev/null +++ b/docs/designs/vpto-pipe-tilelib-lowering-design.md @@ -0,0 +1,440 @@ +# VPTO Pipe PTODSL and TileLib Boundary Design + +**Issue:** [#966](https://github.com/hw-native-sys/PTOAS/issues/966) +**Status:** Compiler-side implementation complete; TileLib implementation is a +separate follow-up +**Target:** A5 VPTO backend +**This change owns:** PTODSL pipe surface validation, PTO IR, PTOAS passes, +`ExpandTileOp` contract production, VPTO LLVM lowering, and focused frontend +coverage +**Out of scope for this change:** `ptodsl/tilelib/**`, TileLib template tests, +and TileLib-ST/runtime coverage + +## 1. Purpose and Scope + +PTODSL already lowers its pipe surface to unified PTO pipe operations. The +VPTO path cannot retain the C++ `TPipe` object used by EmitC, because VPTO +emission needs explicit FIFO state, synchronization, and data movement in PTO +IR before LLVM lowering. + +This document records the implementation boundary now present in PTOAS. It is +both the design for the compiler-side implementation and the hand-off contract +for the TileLib owner. It deliberately does not claim that a source tree +without the corresponding TileLib change can execute GM or local split pipes +end to end. + +The stable frontend-to-backend flow is: + +```text +PTODSL Pipe object + -> frontend pipe operations + -> unified PTO pipe operations + -> infer/validate pipe nosplit configuration + -> materialize PipeState and terminal tdrain + -> memory planning and reserved-buffer resolution + -> pipe metadata/candidate discovery + -> ExpandTileOp PipeSpec ABI + -> external PTODSL TileLib template + -> VPTO LLVM lowering +``` + +## 2. Goals + +- Preserve the existing PTODSL Pipe API and its per-transaction `split` value. +- Make mutable FIFO state explicit for the A5 VPTO path. +- Materialize terminal producer cleanup without deriving its policy from an + arbitrary `tpush` operation. +- Supply `ExpandTileOp` with complete, resolved pipe metadata and resource + operands after memory planning. +- Keep the EmitC `TPipe` path unchanged outside the A5 VPTO backend. +- Provide a narrow ABI that an independently developed TileLib implementation + can consume without depending on opaque `!pto.pipe` values. + +## 3. Non-Goals + +- Replacing the existing EmitC pipe implementation. +- Adding a parallel PTODSL `pipe.push()` or `pipe.pop()` API, or redesigning + the existing API. `pto.pipe.c2v(...)`, `pto.pipe.v2c(...)`, and + `pto.pipe.bidirectional(...)` already return Pipe objects with these methods. +- Modifying `ptodsl/tilelib/**` in this change. +- Claiming GM/global or local split runtime correctness before the TileLib + implementation and TileLib-ST/simulator coverage are merged. +- Supporting A2/A3 VPTO pipe expansion, `v2c_ctrl`, fixpipe quantization, or + phase/NZ conversion in this implementation. + +## 4. PTODSL Pipe Surface + +### 4.1 Existing public API + +The existing constructors return a direction-aware Pipe object: + +```python +pipe = pto.pipe.c2v( + id=7, + gm_slot_tensor=slots, + slot_size=1024, +) +pipe.init_cube() +pipe.init_simd() +entry = pipe.alloc(split=1) +pipe.push(entry, split=1) +entry = pipe.pop(split=1) +pipe.free(entry, split=1) +``` + +`v2c(...)` has the same transaction methods. `bidirectional(...)` exposes +`.c2v` and `.v2c` endpoints before a direction-specific transaction is +performed. The Pipe API requires a stable explicit `id`. + +The public methods lower directly to the matching frontend operations: + +| Pipe method | Lowered frontend operation family | +|---|---| +| `alloc(split)` | `talloc_to_aiv` or `talloc_to_aic` | +| `push(entry, split)` | `tpush_to_aiv` or `tpush_to_aic` | +| `pop(split, ...)` | `tpop_from_aic` or `tpop_from_aiv` | +| `free(entry, split)` | `tfree_from_aic` or `tfree_from_aiv` | + +The `split` is an operation property, not a Pipe-object property. A +split-capable pipe may therefore issue one transaction with `split=1` and a +later transaction with `split=2`, subject to the normal pipe configuration +validation. + +### 4.2 GM slot-size rule + +For a GM-entry pipe, implicit `slot_size` inference is only unambiguous for +`nosplit=True`: in that case one `gm_slot_tensor` shape is one full FIFO slot. +For a split-capable GM pipe, the caller must provide an explicit full-slot +`slot_size`. This prevents a split subregion from being mistaken for the FIFO +slot size during frontend lowering. + +### 4.3 Pipe operations and the template boundary + +PTODSL source and frontend helpers use the public `Pipe` object. Its complete +operation surface is: + +| Category | Public surface | Contract | +|---|---|---| +| constructors | `pto.pipe.c2v(...)`, `pto.pipe.v2c(...)`, `pto.pipe.bidirectional(...)` | Create C2V, V2C, or bidirectional logical pipes. `id` is required and remains the stable pipe identity. | +| initialization | `init_cube()`, `init_simd()` | Initialize the pipe from the Cube or SIMD side. A bidirectional pipe uses its root object for initialization. | +| producer transaction | `alloc(split=0)` | Global-entry pipes only. Returns the descriptor for the next FIFO entry. | +| producer transaction | `push(entry, split=0)` | Publishes a filled global entry or local tile to the consumer. | +| consumer transaction | `pop(split=0, result_type=None, valid_shape=None, valid_row=None, valid_col=None)` | Returns the next global-entry descriptor or local tile. Local tile-entry pipes require `result_type`; `valid_shape` is mutually exclusive with `valid_row` / `valid_col`. | +| consumer transaction | `free(entry=None, split=0)` | Releases a consumed FIFO entry. Global-entry pipes require the entry returned by the corresponding `pop`; local tile-entry pipes may omit it. | +| read-only properties | `id`, `slot_size`, `entry_type` | Expose compile-time identity, full logical slot size in bytes, and the global-entry descriptor type when applicable. | + +`c2v` and `v2c` expose only directionally valid transactions. A bidirectional +pipe must select `.c2v` or `.v2c` before calling `alloc`, `push`, `pop`, or +`free`; its root object has no unambiguous transaction direction. `split` is a +per-transaction compile-time value, not mutable Pipe-object state. It is `0` +for no split, `1` for up/down split, and `2` for left/right split. + +These are PTODSL APIs, not TileLib template APIs. During expansion, a template +does not receive a Python `Pipe` object and cannot call its methods. PTOAS +passes the operation-specific `PipeSpec`, ordered `PipeResources`, shared +`PipeState`, the entry when present, and the optional AIV subblock value +through the ABI in Section 7. The TileLib owner implements the transaction's +FIFO address, synchronization, and counter behavior from those ABI values; +this change neither adds a template-side Pipe wrapper nor modifies +`ptodsl/tilelib/**`. + +## 5. A5 VPTO Default Behavior and Compatibility + +There is no pipe-expansion feature flag. When a module containing frontend or +unified pipe transactions is compiled with `--pto-arch=a5 --pto-backend=vpto`, +the driver automatically performs pipe-specific validation, PipeState +materialization, candidate discovery, and expansion preparation. The default +TileLib backend is PTODSL, so this default path uses PTODSL metadata and +expansion. + +The PipeState materialization pass is inserted only for the A5 VPTO backend. +All EmitC paths, including A5 EmitC, retain the existing `TPipe` lowering and +do not materialize PipeState or terminal `tdrain`. A2/A3 behavior is also +unchanged. + +If the selected TileLib has not implemented a legal pipe candidate, candidate +discovery/expansion fails explicitly. PTOAS must not silently fall back to C++ +`TPipe` or to another TileLib implementation. + +## 6. PTO IR Contract + +### 6.1 Stateful operations + +The internal `pto.talloc`, `pto.tpush`, `pto.tpop`, and `pto.tfree` operations +have an optional `pipe_state` operand. Legacy source-authored IR remains valid +without it so the existing EmitC path is source-compatible. Feature-owned VPTO +IR carries exactly this type: + +```text +!pto.struct +``` + +The verifier rejects another state shape. Once any stateful user of a pipe +has a state, every `talloc`, `tpush`, `tpop`, `tfree`, and `tdrain` user of +that pipe must carry the same SSA state. This prevents partial materialization +or independently allocated counters in hand-authored internal IR. `tpush`, +`tpop`, and `tfree` use operand segments so their existing optional +entry/subblock operands remain unambiguous in assembly form. + +The two fields have a fixed ownership contract: + +| Field | Name | Meaning | +|---|---|---| +| 0 | `prod_index` | Next producer FIFO position. | +| 1 | `cons_index` | Next consumer FIFO position. | + +The state contains no pipe handle, physical address, `flag_base`, or split +mode. Those are immutable pipe configuration or runtime resources and are +provided separately to expansion. + +### 6.2 Terminal `tdrain` + +`pto.tdrain` is an internal operation, inserted only by the PipeState path: + +```text +pto.tdrain(%pipe, %state : !pto.pipe, !pto.struct) { split = <0|1> } +``` + +It models producer-side cleanup formerly performed by the lifetime of the +EmitC `TPipe` object. It is generated once before each reachable `func.return` +for a pipe that has at least one producer `tpush`; a pipe without a producer +does not get a drain. + +`tdrain` is pipe-level cleanup. Its `split` is derived after pipe +configuration resolution, not from a producer operation: + +| Resolved initializer `nosplit` | Materialized `tdrain.split` | +|---|---| +| `true` | `0` | +| `false` | `1` | + +As a result, the materialization pass accepts different producer operation +axes, such as `tpush(split=1)` followed by `tpush(split=2)`. The existing +infer/validate pass remains responsible for rejecting illegal `split` and +`nosplit` combinations. The IR verifier additionally requires an authored +`tdrain` to use this exact derived split and rejects a drain when its +initializer has not resolved `nosplit` yet. + +### 6.3 Materialization pass + +`pto-materialize-pipe-state` runs per `func.func` after +`pto-infer-validate-pipe-init` has resolved `nosplit`. The A5 VPTO driver +schedules it after layout, fusion, memory planning, and reserved-buffer +resolution, immediately before pipe-only candidate discovery. This keeps the +existing shared passes on their pre-existing pipe IR while providing +PipeState, resolved `flag_base`, and resources together to pipe metadata. + +For each initialized pipe with stateful users, it: + +1. Creates one `pto.declare_struct` of type `!pto.struct`. +2. Initializes both fields to zero with `pto.struct_set`. +3. Attaches that same state to all `talloc`, `tpush`, `tpop`, and `tfree` users + of the pipe. +4. Marks feature-owned pipe IR for later expansion cleanup. +5. Inserts the pipe-level `tdrain` when the pipe has a producer. + +The initializer must dominate every return receiving the inserted drain. The +pass diagnoses a nested/lifetime form it cannot represent instead of moving a +pipe handle across regions. + +## 7. Expansion ABI + +### 7.1 Separation of values + +`ExpandTileOp` separates pipe information into four logical operands: + +| Logical operand | Role | Physical helper argument | +|---|---|---| +| entry, when present | tile or mutable GM descriptor used by the transaction | yes | +| PipeSpec | immutable configuration and per-operation `split` | no | +| PipeResources | ordered runtime addresses | one argument per resource | +| PipeState | mutable producer/consumer counters | one struct argument | + +`!pto.pipe` never becomes a template helper argument. Passing it through would +leave an opaque pipe dependency in VPTO lowering and defeat expansion. + +For `tpush` and `tpop`, the ABI additionally includes the AIV subblock ID. +When the original operation has no ID, PTOAS serializes a scalar `i64` value of +zero so the logical template signature is stable. + +### 7.2 PipeSpec producer contract + +After `pto-resolve-reserved-buffers`, the compiler derives `PipeWireInfo` from +the pipe initializer and the individual operation. It serializes a `pipe` +operand with this contract: + +```json +{ + "kind": "pipe", + "init_kind": "l2g2l", + "dir_mask": 2, + "slot_size": 1024, + "slot_num": 8, + "local_slot_num": null, + "flag_base": 0, + "nosplit": false, + "split": 2, + "resource_names": ["gm_addr"] +} +``` + +`init_kind` is `l2l` or `l2g2l`. The source direction encoding remains the +PTO frontend encoding (`1`, `2`, or `3`); a TileLib consumer must not assume +that this field has already been converted to an ISA-specific direction. + +`split` is copied from each `talloc`, `tpush`, `tpop`, `tfree`, or `tdrain`. +It is part of the specialization identity, so helpers for distinct split +paths cannot be incorrectly reused. `tdrain` consequently receives the +pipe-level value established in Section 6.2. + +Missing `flag_base`, an unresolved initializer, an absent PipeState, or an +unsupported `acc_push_epilogue` are diagnosed before a helper is requested. + +### 7.3 Resources, state, and entries + +`PipeResources` contains the present resources in the exact order specified +by `resource_names`: + +```text +gm_addr, local_addr, peer_local_addr +``` + +An `l2l` pipe may expose local and peer-local resources. An `l2g2l` pipe may +also expose `gm_addr`; absent resources do not consume a helper argument. + +`PipeState` is serialized as: + +```json +{"kind": "pipe_state", "fields": ["i32", "i32"]} +``` + +A declared global entry is serialized as `pipe_entry`, rather than a normal +read-only view, because TileLib may need to rebind its caller-owned descriptor. +Its helper ABI is `!pto.tensor_view<...>`. Tile entries retain their normal +tile metadata. + +### 7.4 Specialization and cleanup + +`ExpandTileOp` builds the specialization key from the operation name, target, +entry metadata, PipeSpec, resource topology/type metadata, PipeState schema, +subblock argument, and context attributes such as `kernel_kind`. It forwards +each operation's `split`; it does not select one split value for the whole +pipe. + +After successful template expansion, PTOAS removes only feature-owned unified +pipe operations and initializers whose uses have disappeared. A remaining +feature-owned pipe operation is an expansion error, not an opportunity to +fall back to another backend. + +## 8. VPTO LLVM Boundary + +PipeState survives helper creation and inlining, so both VPTO LLVM emitters +lower the supported struct subset: + +- `pto.declare_struct` becomes function-local storage; +- `pto.struct_get` becomes field address calculation and load; +- `pto.struct_set` becomes field address calculation and store. + +The pipe entry path also supports a mutable `!pto.tensor_view` descriptor: +`pto.declare_global`, `pto.tassign`, and `pto.tensor_view_addr` lower through +the descriptor storage rather than leaving a memref bridge or an +`unrealized_conversion_cast` in the VPTO pipeline. This is compiler support +for the TileLib ABI, not an implementation of FIFO address or synchronization +semantics. + +A global pipe entry may be a `pto.declare_global` result or a result reached +through one or more `pto.tassign` operations from that declaration. The +rebinding result remains the operand passed to the pipe operation. + +## 9. TileLib Handoff + +The following work intentionally belongs to the TileLib owner and is not +included in this PR: + +- metadata classes, renderer bindings, and candidate constraints consuming + `pipe`, `pipe_resources`, `pipe_state`, and `pipe_entry` operands; +- A5 templates for `talloc`, `tpush`, `tpop`, `tfree`, and `tdrain`; +- GM FIFO split address offsets and local split/subblock constraints; +- operation-specific synchronization, FIFO index updates, and terminal drain + behavior; +- daemon/unit coverage and TileLib-ST or simulator end-to-end coverage. + +The TileLib implementation must treat PipeSpec as immutable configuration, +consume resources in the declared order, and use the individual operation's +`split` together with the optional subblock value. It must not recover policy +from `!pto.pipe`, infer a new pipe-wide split value, or reinterpret `tdrain` +as an arbitrary producer transaction. + +Until that implementation is merged, a full pipeline may fail because no +legal pipe template candidate exists or because a template intentionally +rejects an unsupported configuration. Such a failure is expected at the +component boundary and does not demonstrate a PTOAS pass failure. + +## 10. Diagnostics + +The compiler provides actionable failures for these boundary violations: + +| Condition | Diagnostic direction | +|---|---| +| invalid PipeState type | require `!pto.struct` | +| inconsistent PipeState association | require every stateful user of one pipe to share one state | +| invalid authored `tdrain.split` | require the split derived from resolved `nosplit` | +| unresolved `nosplit` before drain | require pipe-init validation before materialization | +| non-dominating drain lifetime | identify initializer and affected return | +| unresolved `flag_base` | require reserved-buffer resolution before expansion | +| unresolved pipe initializer/state | identify the unified pipe operation | +| unsupported `acc_push_epilogue` | reject before template invocation | +| no legal candidate or leftover feature-owned pipe IR | fail expansion without fallback | + +## 11. Validation and Follow-up + +### 11.1 Included coverage + +The PTODSL surface regression covers explicit full-slot GM construction and +two independent nonzero transactions (`split=1` and `split=2`) through +`alloc`, `push`, `pop`, and `free`. It checks that each operation receives its +own split value rather than a pipe-wide default. + +The current compiler build and PTODSL unit suite are the required local checks +for this change: + +```text +ninja -C build-local-vpto ptoas PTOPythonModules +ptodsl/tests/test_vector_cube_ops.py -v +``` + +The current unit suite passes 38 tests. `git diff --check` also passes. + +Focused compiler lit coverage additionally verifies PipeState materialization, +the PipeSpec RPC payload (using a test-only mock daemon outside +`ptodsl/tilelib/**`), defaulted and explicit AIV subblock operands, ordered +GM and local/peer-local resource lists, VPTO LLVM lowering for mutable +descriptors, verifier diagnostics, default A5 VPTO activation, and A5/A3 +EmitC compatibility. It does not claim that the mock daemon validates a +TileLib template or FIFO runtime behavior. + +### 11.2 Required TileLib follow-up coverage + +Once TileLib consumes this ABI, it must add its own tests for: + +- candidate selection and specialization separation for `split=0`, `1`, and + `2`; +- GM and local resource order, address offsets, and subblock handling; +- PipeState counter transitions and terminal `tdrain` behavior; +- unsupported direction/quantization diagnostics; and +- end-to-end FIFO correctness on the TileLib-ST/simulator path. + +The runtime suite must verify output and synchronization behavior. A compile +or rendered-IR check alone is insufficient for GM/local FIFO semantics. + +## 12. Acceptance Boundary + +This PTOAS-side work is complete when the frontend preserves split values, +PipeState/`tdrain` materialization obeys resolved `nosplit`, the compiler +produces the documented expansion ABI, and VPTO LLVM can lower the required +state and mutable descriptor forms. + +The complete feature is complete only after the separate TileLib change +implements the hand-off contract and its TileLib-ST/simulator coverage proves +the supported configurations. The two milestones must remain distinct in PR +status and release claims. diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index f5f7938fc6..8b08aa8566 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -40,6 +40,14 @@ def TensorOrMemref : def PTODpsType : AnyTypeOf<[AnyRankedTensor, AnyMemRef, PartitionTensorViewType, TileBufType]>; +// TASSIGN also rebases a declared GlobalTensor descriptor. Keep this separate +// from PTODpsType: ordinary DPS operations must not accept a raw TensorViewType +// in place of their logical partition view. +def PTOTAssignType : + AnyTypeOf<[AnyRankedTensor, AnyMemRef, TensorViewType, + PartitionTensorViewType, TileBufType], + "TensorView, PartitionTensorView, TileBuf, Tensor, or MemRef">; + def PTOPipeEntryType : AnyTypeOf<[AnyRankedTensor, AnyMemRef, TensorViewType, TileBufType], "TensorView, TileBuf, Tensor, or MemRef">; @@ -2437,6 +2445,7 @@ def InitializeL2LPipeOp : PTO_Op<"initialize_l2l_pipe", [ } def TPushOp : PTO_TOp<"tpush", [ + AttrSizedOperandSegments, OpPipeInterface, DeclareOpInterfaceMethods ]> { @@ -2446,6 +2455,7 @@ def TPushOp : PTO_TOp<"tpush", [ PTOPipeEntryType:$tile, PipeType:$pipe_handle, Optional:$aiv_subblockid, + Optional:$pipe_state, I8Attr:$split ); @@ -2456,6 +2466,7 @@ def TPushOp : PTO_TOp<"tpush", [ `(` $tile `,` $pipe_handle `:` qualified(type($tile)) `,` qualified(type($pipe_handle)) `)` `{` `split` `=` $split `}` (`aiv_subblockid` `(` $aiv_subblockid^ `)`)? + (`pipe_state` `(` $pipe_state^ `:` qualified(type($pipe_state)) `)`)? attr-dict }]; @@ -2519,6 +2530,7 @@ def TAllocOp : PTO_TOp<"talloc", [ let arguments = (ins TensorViewType:$entry, PipeType:$pipe_handle, + Optional:$pipe_state, I8Attr:$split ); @@ -2527,7 +2539,9 @@ def TAllocOp : PTO_TOp<"talloc", [ let assemblyFormat = [{ `(` $entry `,` $pipe_handle `:` qualified(type($entry)) `,` qualified(type($pipe_handle)) `)` - `{` `split` `=` $split `}` attr-dict + `{` `split` `=` $split `}` + (`pipe_state` `(` $pipe_state^ `:` qualified(type($pipe_state)) `)`)? + attr-dict }]; let extraClassDeclaration = [{ @@ -2535,9 +2549,7 @@ def TAllocOp : PTO_TOp<"talloc", [ }]; } -def TAssignOp : PTO_Op<"tassign", [ - Pure - ]> { +def TAssignOp : PTO_Op<"tassign"> { let summary = "Rebind a declared tile handle to a runtime address"; let description = [{ Rebinds an existing tile handle to a new address. This op is intended for @@ -2546,10 +2558,10 @@ def TAssignOp : PTO_Op<"tassign", [ }]; let arguments = (ins - PTODpsType:$tile, + PTOTAssignType:$tile, I64:$addr ); - let results = (outs PTODpsType:$result); + let results = (outs PTOTAssignType:$result); let assemblyFormat = [{ $tile `,` $addr attr-dict `:` qualified(type($tile)) `->` qualified(type($result)) @@ -2790,6 +2802,7 @@ def DeclareTileMemRefOp : PTO_Op<"declare_tile_memref"> { def TPopOp : PTO_TOp<"tpop", [ PTO_DpsInitOpInterface, + AttrSizedOperandSegments, OpPipeInterface, DeclareOpInterfaceMethods ]> { @@ -2799,6 +2812,7 @@ def TPopOp : PTO_TOp<"tpop", [ PTOPipeEntryType:$tile, PipeType:$pipe_handle, Optional:$aiv_subblockid, + Optional:$pipe_state, I8Attr:$split ); @@ -2809,6 +2823,7 @@ def TPopOp : PTO_TOp<"tpop", [ `(` $tile `,` $pipe_handle `:` qualified(type($tile)) `,` qualified(type($pipe_handle)) `)` `{` `split` `=` $split `}` (`aiv_subblockid` `(` $aiv_subblockid^ `)`)? + (`pipe_state` `(` $pipe_state^ `:` qualified(type($pipe_state)) `)`)? attr-dict }]; @@ -2864,6 +2879,7 @@ def TPopOp : PTO_TOp<"tpop", [ } def TFreeOp : PTO_TOp<"tfree", [ + AttrSizedOperandSegments, OpPipeInterface, DeclareOpInterfaceMethods ]> { @@ -2872,6 +2888,7 @@ def TFreeOp : PTO_TOp<"tfree", [ let arguments = (ins Optional:$entry, PipeType:$pipe_handle, + Optional:$pipe_state, I8Attr:$split ); @@ -2884,6 +2901,35 @@ def TFreeOp : PTO_TOp<"tfree", [ }]; } +def TDrainOp : PTO_TOp<"tdrain", [ + OpPipeInterface, + DeclareOpInterfaceMethods +]> { + let summary = "Complete producer-side cleanup for a stateful unified pipe"; + let description = [{ + Internal operation materialized only by the VPTO PipeState path. It models + the terminal producer-side synchronization performed by the EmitC `TPipe` + object during destruction. + }]; + + let arguments = (ins + PipeType:$pipe_handle, + StructType:$pipe_state, + I8Attr:$split + ); + let results = (outs); + let hasVerifier = 1; + + let assemblyFormat = [{ + `(` $pipe_handle `,` $pipe_state `:` qualified(type($pipe_handle)) `,` qualified(type($pipe_state)) `)` + `{` `split` `=` $split `}` attr-dict + }]; + + let extraClassDeclaration = [{ + ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_MTE3; } + }]; +} + //===----------------------------------------------------------------------===// // Synchronization Ops //===----------------------------------------------------------------------===// diff --git a/include/PTO/IR/PTOTypeUtils.h b/include/PTO/IR/PTOTypeUtils.h index 67fcf8b3a2..e7a2f261f2 100644 --- a/include/PTO/IR/PTOTypeUtils.h +++ b/include/PTO/IR/PTOTypeUtils.h @@ -36,6 +36,10 @@ inline LogicalResult getPTOMemRefStridesAndOffset( return detail::getPTOMemRefStridesAndOffsetImpl(memTy, strides, offset, 0); } +/// Returns whether `value` is a pto.declare_global result or a result of one +/// or more pto.tassign operations rooted at that declaration. +bool isDeclaredGlobalOrTAssignResult(Value value); + bool isPTOFloat8Type(Type t); bool isPTOFloat8E4M3LikeType(Type t); bool isPTOFloat8E5M2LikeType(Type t); diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index 88c37ad3a5..812699461d 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -37,6 +37,7 @@ std::unique_ptr createLoweringSyncToPipePass(); std::unique_ptr createPTOAssignDefaultFrontendPipeIdPass(); std::unique_ptr createPTOLowerFrontendPipeOpsPass(); std::unique_ptr createPTOInferValidatePipeInitPass(); +std::unique_ptr createPTOMaterializePipeStatePass(); std::unique_ptr createPTOResolveReservedBuffersPass(); std::unique_ptr createPTOWrapFunctionsInSectionsPass(); std::unique_ptr createPTONormalizeUncoveredTileSectionsPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index d4366200fd..7eaa9f41ce 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -484,6 +484,25 @@ def PTOInferValidatePipeInit : Pass<"pto-infer-validate-pipe-init", "ModuleOp"> ]; } +def PTOMaterializePipeState : Pass<"pto-materialize-pipe-state", "func::FuncOp"> { + let summary = "Materialize explicit PipeState operands for unified pipe ops"; + let description = [{ + Creates one `!pto.struct` state object per internal pipe handle, + initializes its producer and consumer counters to zero, attaches it to + unified `talloc` / `tpush` / `tpop` / `tfree` operations, and inserts a + terminal `tdrain` for every pipe with producer pushes. The terminal drain + follows the resolved pipe-level `nosplit` configuration (`split = 0` when + true, otherwise `split = 1`), so producer operations on a split-capable + pipe may use different split axes. + }]; + let constructor = "mlir::pto::createPTOMaterializePipeStatePass()"; + let dependentDialects = [ + "mlir::arith::ArithDialect", + "mlir::pto::PTODialect", + "mlir::func::FuncDialect" + ]; +} + def PTOResolveReservedBuffers : Pass<"pto-resolve-reserved-buffers", "ModuleOp"> { let summary = "Resolve reserved local buffer addresses and peer pipe flag bases"; let description = [{ @@ -533,7 +552,13 @@ def InsertTemplateAttributes "PYTHONPATH root for PTODSL">, Option<"daemonHelperModule", "daemon-helper-module", "std::string", /*default=*/"\"ptodsl.tilelib.serving.helper\"", - "Python module used for daemon metadata RPC calls"> + "Python module used for daemon metadata RPC calls">, + Option<"pipeOnly", "pipe-only", "bool", + /*default=*/"false", + "Attach candidates only to unified internal pipe operations">, + Option<"skipPipeOps", "skip-pipe-ops", "bool", + /*default=*/"false", + "Exclude unified internal pipe operations from this metadata pass"> ]; } diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index c1cb5e9e61..89c07f8b97 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3589,11 +3589,16 @@ LogicalResult mlir::pto::SyncSetOp::verify() { } switch (getPipe().getPipe()) { case PIPE::PIPE_FIX: + case PIPE::PIPE_MTE1: + case PIPE::PIPE_MTE2: case PIPE::PIPE_MTE3: + case PIPE::PIPE_V: + case PIPE::PIPE_S: return success(); default: return emitOpError() - << "A5 sync.set expects pipe to be one of , "; + << "A5 sync.set expects pipe to be one of , , " + ", , , or "; } }; return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); @@ -3763,11 +3768,12 @@ LogicalResult mlir::pto::SyncWaitOp::verify() { case PIPE::PIPE_MTE2: case PIPE::PIPE_MTE3: case PIPE::PIPE_V: + case PIPE::PIPE_S: return success(); default: return emitOpError() << "A5 sync.wait expects pipe to be one of " ", , , " - ", "; + ", , "; } }; return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); @@ -17712,11 +17718,115 @@ LogicalResult InitializeL2LPipeOp::verify() { return success(); } +static LogicalResult verifyPipeStateOperand(Operation *op, Value state, + bool required = false) { + if (!state) { + if (required) + return op->emitOpError("requires pipe_state"); + return success(); + } + + auto stateType = dyn_cast(state.getType()); + if (!stateType || stateType.getNumFields() != 2 || + !stateType.getFieldType(0).isSignlessInteger(32) || + !stateType.getFieldType(1).isSignlessInteger(32)) { + return op->emitOpError() + << "pipe_state must have type !pto.struct"; + } + return success(); +} + +static bool isPipeStateOperation(Operation *op) { + return isa(op); +} + +static Value getPipeStateHandle(Operation *op) { + if (auto alloc = dyn_cast(op)) + return alloc.getPipeHandle(); + if (auto push = dyn_cast(op)) + return push.getPipeHandle(); + if (auto pop = dyn_cast(op)) + return pop.getPipeHandle(); + if (auto free = dyn_cast(op)) + return free.getPipeHandle(); + return cast(op).getPipeHandle(); +} + +static Value getPipeStateOperandForVerification(Operation *op) { + if (auto alloc = dyn_cast(op)) + return alloc.getPipeState(); + if (auto push = dyn_cast(op)) + return push.getPipeState(); + if (auto pop = dyn_cast(op)) + return pop.getPipeState(); + if (auto free = dyn_cast(op)) + return free.getPipeState(); + return cast(op).getPipeState(); +} + +static LogicalResult verifyPipeStateAssociation(Operation *op, + Value pipeHandle) { + Operation *firstStateUser = nullptr; + Value commonState; + bool hasMissingState = false; + + for (Operation *user : pipeHandle.getUsers()) { + if (!isPipeStateOperation(user) || getPipeStateHandle(user) != pipeHandle) + continue; + if (!firstStateUser) + firstStateUser = user; + + Value state = getPipeStateOperandForVerification(user); + if (!state) { + hasMissingState = true; + continue; + } + if (commonState && commonState != state) { + if (firstStateUser == op) + return op->emitOpError() + << "requires all stateful users of the same pipe to share " + "one pipe_state"; + return success(); + } + commonState = state; + } + + if (commonState && hasMissingState && firstStateUser == op) + return op->emitOpError() + << "requires every stateful user of a materialized pipe to carry " + "the shared pipe_state"; + return success(); +} + +static LogicalResult verifyTDrainSplit(Operation *op, Value pipeHandle, + int64_t split) { + BoolAttr noSplit; + if (auto init = pipeHandle.getDefiningOp()) + noSplit = init.getNosplitAttr(); + else if (auto init = pipeHandle.getDefiningOp()) + noSplit = init.getNosplitAttr(); + + if (!noSplit) + return op->emitOpError() + << "requires the pipe initializer to have resolved 'nosplit'"; + + const int64_t expectedSplit = noSplit.getValue() ? 0 : 1; + if (split != expectedSplit) + return op->emitOpError() << "requires split = " << expectedSplit + << " to match the pipe initializer's nosplit = " + << (noSplit.getValue() ? "true" : "false"); + return success(); +} + LogicalResult TPushOp::verify() { if (!isInsideSectionOrAttributedKernel(getOperation())) return emitOpError("must be inside pto.section.cube/vector or a kernel_kind function"); if (failed(verifyPipeHandleProducer(getOperation(), getPipeHandle()))) return failure(); + if (failed(verifyPipeStateOperand(getOperation(), getPipeState()))) + return failure(); + if (failed(verifyPipeStateAssociation(getOperation(), getPipeHandle()))) + return failure(); if (failed(verifySplitAttr(getOperation(), getSplit()))) return failure(); if (failed(verifyAivSubblockIdOperand(getOperation(), getAivSubblockid(), @@ -17736,6 +17846,10 @@ LogicalResult TAllocOp::verify() { return emitOpError("must be inside pto.section.cube/vector or a kernel_kind function"); if (failed(verifyPipeHandleProducer(getOperation(), getPipeHandle()))) return failure(); + if (failed(verifyPipeStateOperand(getOperation(), getPipeState()))) + return failure(); + if (failed(verifyPipeStateAssociation(getOperation(), getPipeHandle()))) + return failure(); if (failed(verifyTensorEntryMatchesInternalPipeInit( getOperation(), getPipeHandle(), getEntry().getType()))) return failure(); @@ -17747,6 +17861,10 @@ LogicalResult TPopOp::verify() { return emitOpError("must be inside pto.section.cube/vector or a kernel_kind function"); if (failed(verifyPipeHandleProducer(getOperation(), getPipeHandle()))) return failure(); + if (failed(verifyPipeStateOperand(getOperation(), getPipeState()))) + return failure(); + if (failed(verifyPipeStateAssociation(getOperation(), getPipeHandle()))) + return failure(); if (failed(verifySplitAttr(getOperation(), getSplit()))) return failure(); if (failed(verifyAivSubblockIdOperand(getOperation(), getAivSubblockid(), @@ -17767,6 +17885,10 @@ LogicalResult TFreeOp::verify() { return emitOpError("must be inside pto.section.cube/vector or a kernel_kind function"); if (failed(verifyPipeHandleProducer(getOperation(), getPipeHandle()))) return failure(); + if (failed(verifyPipeStateOperand(getOperation(), getPipeState()))) + return failure(); + if (failed(verifyPipeStateAssociation(getOperation(), getPipeHandle()))) + return failure(); if (getEntry() && failed(verifyTensorEntryMatchesInternalPipeInit( getOperation(), getPipeHandle(), getEntry().getType()))) @@ -17774,12 +17896,31 @@ LogicalResult TFreeOp::verify() { return verifySplitAttr(getOperation(), getSplit()); } +LogicalResult TDrainOp::verify() { + if (!isInsideSectionOrAttributedKernel(getOperation())) + return emitOpError( + "must be inside pto.section.cube/vector or a kernel_kind function"); + if (failed(verifyPipeHandleProducer(getOperation(), getPipeHandle()))) + return failure(); + if (failed(verifyPipeStateOperand(getOperation(), getPipeState(), + /*required=*/true))) + return failure(); + if (failed(verifyPipeStateAssociation(getOperation(), getPipeHandle()))) + return failure(); + if (failed(verifySplitAttr(getOperation(), getSplit()))) + return failure(); + return verifyTDrainSplit(getOperation(), getPipeHandle(), getSplit()); +} + ParseResult TFreeOp::parse(OpAsmParser &parser, OperationState &result) { OpAsmParser::UnresolvedOperand first; OpAsmParser::UnresolvedOperand pipe; Type firstTy; Type pipeTy; + OpAsmParser::UnresolvedOperand state; + Type stateTy; bool hasEntry = false; + bool hasState = false; if (parser.parseLParen() || parser.parseOperand(first)) return failure(); @@ -17802,15 +17943,30 @@ ParseResult TFreeOp::parse(OpAsmParser &parser, OperationState &result) { IntegerAttr splitAttr; if (parser.parseAttribute(splitAttr, parser.getBuilder().getI8Type(), "split", attrs) || - parser.parseRBrace() || parser.parseOptionalAttrDict(attrs)) + parser.parseRBrace()) + return failure(); + + if (succeeded(parser.parseOptionalKeyword("pipe_state"))) { + hasState = true; + if (parser.parseLParen() || parser.parseOperand(state) || + parser.parseColonType(stateTy) || parser.parseRParen()) + return failure(); + } + + if (parser.parseOptionalAttrDict(attrs)) return failure(); result.addAttributes(attrs); + result.addAttribute("operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr( + {hasEntry ? 1 : 0, 1, hasState ? 1 : 0})); if (hasEntry && parser.resolveOperand(first, firstTy, result.operands)) return failure(); if (parser.resolveOperand(pipe, pipeTy, result.operands)) return failure(); + if (hasState && parser.resolveOperand(state, stateTy, result.operands)) + return failure(); return success(); } @@ -17823,8 +17979,11 @@ void TFreeOp::print(OpAsmPrinter &p) { p << getPipeHandle() << " : " << getPipeHandle().getType(); } p << ") {split = " << static_cast(getSplit()) << "}"; + if (getPipeState()) + p << " pipe_state(" << getPipeState() << " : " + << getPipeState().getType() << ")"; p.printOptionalAttrDict((*this)->getAttrs(), - /*elidedAttrs=*/{"split"}); + /*elidedAttrs=*/{"split", "operandSegmentSizes"}); } static func::FuncOp getParentFunc(Operation *op) { @@ -18294,6 +18453,11 @@ void TPushOp::getEffects( auto aivSubblockId = getAivSubblockidMutable(); if (!aivSubblockId.empty()) addEffect(effects, &*aivSubblockId.begin(), MemoryEffects::Read::get()); + auto pipeState = getPipeStateMutable(); + if (!pipeState.empty()) { + addEffect(effects, &*pipeState.begin(), MemoryEffects::Read::get()); + addEffect(effects, &*pipeState.begin(), MemoryEffects::Write::get()); + } addEffect(effects, &getPipeHandleMutable(), MemoryEffects::Read::get()); addEffect(effects, &getPipeHandleMutable(), MemoryEffects::Write::get()); effects.emplace_back(MemoryEffects::Read::get(), @@ -18349,6 +18513,11 @@ void TAllocOp::getEffects( addEffect(effects, &getEntryMutable(), MemoryEffects::Write::get()); addEffect(effects, &getPipeHandleMutable(), MemoryEffects::Read::get()); addEffect(effects, &getPipeHandleMutable(), MemoryEffects::Write::get()); + auto pipeState = getPipeStateMutable(); + if (!pipeState.empty()) { + addEffect(effects, &*pipeState.begin(), MemoryEffects::Read::get()); + addEffect(effects, &*pipeState.begin(), MemoryEffects::Write::get()); + } } void TPopOp::getEffects( @@ -18359,6 +18528,11 @@ void TPopOp::getEffects( auto aivSubblockId = getAivSubblockidMutable(); if (!aivSubblockId.empty()) addEffect(effects, &*aivSubblockId.begin(), MemoryEffects::Read::get()); + auto pipeState = getPipeStateMutable(); + if (!pipeState.empty()) { + addEffect(effects, &*pipeState.begin(), MemoryEffects::Read::get()); + addEffect(effects, &*pipeState.begin(), MemoryEffects::Write::get()); + } addEffect(effects, &getTileMutable(), MemoryEffects::Write::get()); } @@ -18370,6 +18544,20 @@ void TFreeOp::getEffects( addEffect(effects, &*entry.begin(), MemoryEffects::Read::get()); addEffect(effects, &getPipeHandleMutable(), MemoryEffects::Read::get()); addEffect(effects, &getPipeHandleMutable(), MemoryEffects::Write::get()); + auto pipeState = getPipeStateMutable(); + if (!pipeState.empty()) { + addEffect(effects, &*pipeState.begin(), MemoryEffects::Read::get()); + addEffect(effects, &*pipeState.begin(), MemoryEffects::Write::get()); + } +} + +void TDrainOp::getEffects( + SmallVectorImpl> + &effects) { + addEffect(effects, &getPipeHandleMutable(), MemoryEffects::Read::get()); + addEffect(effects, &getPipeHandleMutable(), MemoryEffects::Write::get()); + addEffect(effects, &getPipeStateMutable(), MemoryEffects::Read::get()); + addEffect(effects, &getPipeStateMutable(), MemoryEffects::Write::get()); } void SetQuantScalarOp::getEffects( diff --git a/lib/PTO/IR/PTOTypeUtils.cpp b/lib/PTO/IR/PTOTypeUtils.cpp index 27851f9cad..6f76a127dd 100644 --- a/lib/PTO/IR/PTOTypeUtils.cpp +++ b/lib/PTO/IR/PTOTypeUtils.cpp @@ -17,6 +17,12 @@ namespace { constexpr unsigned kBitsPerByte = 8; } // namespace +bool mlir::pto::isDeclaredGlobalOrTAssignResult(Value value) { + while (auto assign = value.getDefiningOp()) + value = assign.getTile(); + return value.getDefiningOp() != nullptr; +} + bool mlir::pto::isPTOFloat8Type(Type t) { return isPTOFloat8E4M3LikeType(t) || isPTOFloat8E5M2LikeType(t); } diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 1c705df4b4..eebbc7e893 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -84,6 +84,7 @@ add_mlir_dialect_library(PTOTransforms PTOAssignDefaultFrontendPipeIdPass.cpp PTOLowerFrontendPipeOpsPass.cpp PTOInferValidatePipeInitPass.cpp + PTOMaterializePipeStatePass.cpp PTOResolveBufferSelect.cpp PTOResolveReservedBuffersPass.cpp SlotAffineAnalysis.cpp diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index dedc0267cb..e3333f2baa 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -31,6 +31,7 @@ #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Support/PythonExecutable.h" #include "PTO/Transforms/Passes.h" +#include "Utils.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" @@ -62,6 +63,7 @@ #include #include #include +#include extern "C" { extern char **environ; @@ -85,10 +87,11 @@ constexpr llvm::StringLiteral kCandidatesAttr = "candidates"; // ============================================================================ // OperandTypeInfo: describes one operand for template specialization. // -// Four kinds of operands: +// Eight kinds of operands: // Tile — from TileBufType. dtype + shape + memorySpace + config // all participate in the specialization key (SpecKey). -// View — from MemRefType (lowered PartitionTensorViewType). The element +// View — from MemRefType, TensorViewType, or PartitionTensorViewType. The +// element // dtype, shape, strides, memory space, and optional explicit layout // participate in SpecKey. PTODSL templates compile ViewSpec metadata // into helper bodies, so helpers with different view strides must not @@ -98,8 +101,22 @@ constexpr llvm::StringLiteral kCandidatesAttr = "candidates"; // distinguish auxiliary vector operands such as tmrgsort's // `excuted : vector<4xi16>`. // Scalar — from a scalar element type. Only dtype participates in SpecKey. +// Pipe — compile-time FIFO configuration only; never a helper argument. +// PipeResources — the initializer's runtime storage operands, flattened into +// physical helper arguments. +// PipeState — fixed !pto.struct cursor state. +// PipeEntry — mutable !pto.tensor_view descriptor for a GM pipe entry. // ============================================================================ -enum class OperandKind { Tile, View, Vector, Scalar }; +enum class OperandKind { + Tile, + View, + Vector, + Scalar, + Pipe, + PipeResources, + PipeState, + PipeEntry, +}; struct OperandTypeInfo { OperandKind kind = OperandKind::Tile; @@ -126,6 +143,20 @@ struct OperandTypeInfo { // --- Scalar-only --- std::optional scalarValue; + // --- Pipe-only --- + std::string pipeInitKind; + int64_t pipeDirMask = 0; + int64_t pipeSlotSize = 0; + int64_t pipeSlotNum = 0; + std::optional pipeLocalSlotNum; + int64_t pipeFlagBase = 0; + bool pipeNoSplit = false; + int64_t pipeSplit = 0; + std::vector pipeResourceNames; + + // --- PipeResources-only --- + std::vector pipeResources; + /// Equality for SpecKey caching — only compares fields relevant to each kind. bool operator==(const OperandTypeInfo &rhs) const { if (kind != rhs.kind || dtype != rhs.dtype) @@ -140,6 +171,19 @@ struct OperandTypeInfo { return vectorShape == rhs.vectorShape; if (kind == OperandKind::Scalar) return scalarValue == rhs.scalarValue; + if (kind == OperandKind::Pipe) + return pipeInitKind == rhs.pipeInitKind && + pipeDirMask == rhs.pipeDirMask && + pipeSlotSize == rhs.pipeSlotSize && + pipeSlotNum == rhs.pipeSlotNum && + pipeLocalSlotNum == rhs.pipeLocalSlotNum && + pipeFlagBase == rhs.pipeFlagBase && + pipeNoSplit == rhs.pipeNoSplit && pipeSplit == rhs.pipeSplit && + pipeResourceNames == rhs.pipeResourceNames; + if (kind == OperandKind::PipeResources) + return pipeResources == rhs.pipeResources; + if (kind == OperandKind::PipeState) + return true; return viewShape == rhs.viewShape && viewStrides == rhs.viewStrides && viewMemorySpace == rhs.viewMemorySpace && @@ -185,8 +229,17 @@ struct SpecKeyInfo : public llvm::DenseMapInfo { h = llvm::hash_combine(h, op.scalarValue.has_value()); if (op.scalarValue) h = llvm::hash_combine(h, *op.scalarValue); + } else if (op.kind == OperandKind::Pipe) { + h = llvm::hash_combine( + h, op.pipeInitKind, op.pipeDirMask, op.pipeSlotSize, + op.pipeSlotNum, op.pipeLocalSlotNum.has_value(), op.pipeFlagBase, + op.pipeNoSplit, op.pipeSplit); + if (op.pipeLocalSlotNum) + h = llvm::hash_combine(h, *op.pipeLocalSlotNum); + for (const std::string &name : op.pipeResourceNames) + h = llvm::hash_combine(h, name); } - if (op.kind == OperandKind::View) { + if (op.kind == OperandKind::View || op.kind == OperandKind::PipeEntry) { h = llvm::hash_combine(h, op.viewMemorySpace); for (int64_t d : op.viewShape) h = llvm::hash_combine(h, d); @@ -295,6 +348,12 @@ static std::string getMemorySpaceString(MemRefType mrTy) { return msAttr ? stringifyMemorySpace(msAttr.getAddressSpace()) : "gm"; } +static std::string getMemorySpaceString(pto::PartitionTensorViewType) { + return "gm"; +} + +static std::string getMemorySpaceString(pto::TensorViewType) { return "gm"; } + static std::string getBLayoutString(int32_t blayout) { if (blayout == static_cast(pto::BLayout::ColMajor)) return "col_major"; @@ -325,6 +384,11 @@ static std::optional resolveViewLayout(Value value) { Operation *def = value.getDefiningOp(); while (def) { + if (auto part = dyn_cast(def)) { + value = part.getSource(); + def = value.getDefiningOp(); + continue; + } if (auto layout = getLayoutAttrFromOp(def)) return layout; if (auto subview = dyn_cast(def)) { @@ -479,6 +543,23 @@ static std::string getTRandomRoundsString(pto::TRandomOp op) { static void appendOpContextAttrs( Operation *op, SmallVectorImpl> &attrs) { + if (isa(op)) { + auto func = op->getParentOfType(); + auto kernelKind = + func ? func->getAttrOfType( + pto::FunctionKernelKindAttr::name) + : pto::FunctionKernelKindAttr(); + if (!kernelKind) { + op->emitError( + "requires a cube or vector kernel_kind for TileLib pipe expansion"); + return; + } + attrs.emplace_back( + "kernel_kind", + kernelKind.getKernelKind() == pto::FunctionKernelKind::Cube ? "cube" + : "vector"); + } if (auto tcvt = dyn_cast(op)) { std::optional roundMode = getTCvtRoundModeString(tcvt); if (roundMode) @@ -622,6 +703,81 @@ static void populateViewShapeAndStrides(Value value, } } +static void populateTensorViewShapeAndStrides( + Value value, SmallVectorImpl &shape, + SmallVectorImpl &strides) { + if (!value) + return; + + if (auto make = value.getDefiningOp()) { + if (shape.empty()) { + if (auto viewTy = + dyn_cast(make.getResult().getType())) + shape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); + } + if (strides.empty()) { + strides.reserve(make.getStrides().size()); + for (Value strideValue : make.getStrides()) { + int64_t stride = ShapedType::kDynamic; + (void)getStaticIntFromValue(strideValue, stride); + strides.push_back(stride); + } + } + return; + } + + if (auto viewTy = dyn_cast(value.getType())) { + if (shape.empty()) + shape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); + } +} + +static void populatePTOViewShapeAndStrides( + Value value, SmallVectorImpl &shape, + SmallVectorImpl &strides) { + if (!value) + return; + + if (auto part = value.getDefiningOp()) { + if (shape.empty()) { + shape.reserve(part.getSizes().size()); + for (Value sizeValue : part.getSizes()) { + int64_t size = ShapedType::kDynamic; + (void)getStaticIntFromValue(sizeValue, size); + shape.push_back(size); + } + } + SmallVector sourceShape; + SmallVector sourceStrides; + populatePTOViewShapeAndStrides(part.getSource(), sourceShape, + sourceStrides); + if (strides.empty() && !sourceStrides.empty()) + strides = sourceStrides; + return; + } + + if (auto make = value.getDefiningOp()) { + if (shape.empty()) { + if (auto viewTy = dyn_cast(make.getResult().getType())) + shape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); + } + if (strides.empty()) { + strides.reserve(make.getStrides().size()); + for (Value strideValue : make.getStrides()) { + int64_t stride = ShapedType::kDynamic; + (void)getStaticIntFromValue(strideValue, stride); + strides.push_back(stride); + } + } + return; + } + + if (auto viewTy = dyn_cast(value.getType())) { + if (shape.empty()) + shape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); + } +} + static std::optional buildOperandTypeInfo(Value value) { Type ty = value.getType(); // Tile operand — from TileBufType. @@ -671,6 +827,34 @@ static std::optional buildOperandTypeInfo(Value value) { return info; } + if (auto viewTy = dyn_cast(ty)) { + OperandTypeInfo info; + info.kind = OperandKind::View; + info.dtype = getDtypeString(viewTy.getElementType()); + if (info.dtype.empty()) + return std::nullopt; + info.viewMemorySpace = getMemorySpaceString(viewTy); + info.viewLayout = resolveViewLayout(value); + populateTensorViewShapeAndStrides(value, info.viewShape, info.viewStrides); + if (info.viewShape.empty()) + info.viewShape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); + return info; + } + + if (auto viewTy = dyn_cast(ty)) { + OperandTypeInfo info; + info.kind = OperandKind::View; + info.dtype = getDtypeString(viewTy.getElementType()); + if (info.dtype.empty()) + return std::nullopt; + info.viewMemorySpace = getMemorySpaceString(viewTy); + info.viewLayout = resolveViewLayout(value); + populatePTOViewShapeAndStrides(value, info.viewShape, info.viewStrides); + if (info.viewShape.empty()) + info.viewShape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); + return info; + } + // Auxiliary vector operand — from builtin VectorType (e.g. vector<4xi16>). if (auto vecTy = dyn_cast(ty)) { OperandTypeInfo info; @@ -694,7 +878,201 @@ static std::optional buildOperandTypeInfo(Value value) { return info; } +static std::optional buildPipeEntryTypeInfo(Value value) { + if (!isa(value.getType())) + return buildOperandTypeInfo(value); + if (!pto::isDeclaredGlobalOrTAssignResult(value)) + return std::nullopt; + auto info = buildOperandTypeInfo(value); + if (!info) + return std::nullopt; + info->kind = OperandKind::PipeEntry; + return info; +} + +struct PipeWireInfo { + bool isL2G2L = false; + int64_t dirMask = 0; + int64_t slotSize = 0; + int64_t slotNum = 0; + std::optional localSlotNum; + int64_t flagBase = 0; + bool noSplit = false; + int64_t split = 0; + SmallVector, 3> resources; + Value state; + Value subblock; +}; + +static FailureOr getPipeWireInfo(Operation *op) { + Value pipe; + PipeWireInfo info; + if (auto push = dyn_cast(op)) { + pipe = push.getPipeHandle(); + info.state = push.getPipeState(); + info.subblock = push.getAivSubblockid(); + info.split = push.getSplit(); + } else if (auto pop = dyn_cast(op)) { + pipe = pop.getPipeHandle(); + info.state = pop.getPipeState(); + info.subblock = pop.getAivSubblockid(); + info.split = pop.getSplit(); + } else if (auto alloc = dyn_cast(op)) { + pipe = alloc.getPipeHandle(); + info.state = alloc.getPipeState(); + info.split = alloc.getSplit(); + } else if (auto free = dyn_cast(op)) { + pipe = free.getPipeHandle(); + info.state = free.getPipeState(); + info.split = free.getSplit(); + } else if (auto drain = dyn_cast(op)) { + pipe = drain.getPipeHandle(); + info.state = drain.getPipeState(); + info.split = drain.getSplit(); + } else { + return op->emitError("is not a unified internal pipe operation"); + } + + if (!info.state) + return op->emitError( + "requires a materialized !pto.struct pipe_state"); + + Operation *init = mlir::pto::getPipeInitDef(pipe); + if (!init) + return op->emitError("cannot resolve its pipe initializer"); + if (mlir::pto::getPipeInitAccPushEpilogue(init)) + return op->emitError( + "does not support acc_push_epilogue for TileLib pipe expansion"); + + auto addResource = [&](StringRef name, Value value) { + if (value) + info.resources.emplace_back(name.str(), value); + }; + if (auto l2l = dyn_cast(init)) { + if (!l2l.getFlagBaseAttr()) + return op->emitError( + "requires resolved pipe flag_base before TileLib expansion"); + info.dirMask = l2l.getDirMask(); + info.slotSize = l2l.getSlotSize(); + info.slotNum = l2l.getSlotNum(); + info.flagBase = l2l.getFlagBaseAttr().getInt(); + info.noSplit = l2l.getNosplitAttr() && l2l.getNosplitAttr().getValue(); + addResource("local_addr", l2l.getLocalAddr()); + addResource("peer_local_addr", l2l.getPeerLocalAddr()); + return info; + } + + auto l2g2l = dyn_cast(init); + if (!l2g2l) + return op->emitError("requires an l2l or l2g2l pipe initializer"); + if (!l2g2l.getFlagBaseAttr()) + return op->emitError( + "requires resolved pipe flag_base before TileLib expansion"); + info.isL2G2L = true; + info.dirMask = l2g2l.getDirMask(); + info.slotSize = l2g2l.getSlotSize(); + info.slotNum = l2g2l.getSlotNum(); + if (auto attr = l2g2l.getLocalSlotNumAttr()) + info.localSlotNum = attr.getInt(); + info.flagBase = l2g2l.getFlagBaseAttr().getInt(); + info.noSplit = l2g2l.getNosplitAttr() && l2g2l.getNosplitAttr().getValue(); + addResource("gm_addr", l2g2l.getGmAddr()); + addResource("local_addr", l2g2l.getLocalAddr()); + addResource("peer_local_addr", l2g2l.getPeerLocalAddr()); + return info; +} + +static OperandTypeInfo buildPipeTypeInfo(const PipeWireInfo &wire) { + OperandTypeInfo info; + info.kind = OperandKind::Pipe; + info.pipeInitKind = wire.isL2G2L ? "l2g2l" : "l2l"; + info.pipeDirMask = wire.dirMask; + info.pipeSlotSize = wire.slotSize; + info.pipeSlotNum = wire.slotNum; + info.pipeLocalSlotNum = wire.localSlotNum; + info.pipeFlagBase = wire.flagBase; + info.pipeNoSplit = wire.noSplit; + info.pipeSplit = wire.split; + for (const auto &resource : wire.resources) + info.pipeResourceNames.push_back(resource.first); + return info; +} + +static std::optional +buildPipeResourcesTypeInfo(const PipeWireInfo &wire) { + OperandTypeInfo info; + info.kind = OperandKind::PipeResources; + for (const auto &resource : wire.resources) { + auto resourceInfo = buildOperandTypeInfo(resource.second); + if (!resourceInfo) + return std::nullopt; + info.pipeResourceNames.push_back(resource.first); + info.pipeResources.push_back(*resourceInfo); + } + return info; +} + +static std::optional buildPipeSpecKey(Operation *op) { + FailureOr wire = getPipeWireInfo(op); + if (failed(wire)) + return std::nullopt; + + SpecKey key; + key.opName = getTileOpName(op).str(); + key.targetArch = getTargetArchString(op); + auto appendPipeEntry = [&](Value value) -> bool { + auto info = buildPipeEntryTypeInfo(value); + if (!info) + return false; + key.operands.push_back(*info); + return true; + }; + + if (auto push = dyn_cast(op)) { + if (!appendPipeEntry(push.getTile())) + return std::nullopt; + } else if (auto pop = dyn_cast(op)) { + if (!appendPipeEntry(pop.getTile())) + return std::nullopt; + } else if (auto alloc = dyn_cast(op)) { + if (!appendPipeEntry(alloc.getEntry())) + return std::nullopt; + } else if (auto free = dyn_cast(op)) { + if (free.getEntry() && !appendPipeEntry(free.getEntry())) + return std::nullopt; + } + + key.operands.push_back(buildPipeTypeInfo(*wire)); + auto resources = buildPipeResourcesTypeInfo(*wire); + if (!resources) + return std::nullopt; + key.operands.push_back(*resources); + OperandTypeInfo state; + state.kind = OperandKind::PipeState; + key.operands.push_back(state); + if (isa(op)) { + if (wire->subblock) { + auto subblock = buildOperandTypeInfo(wire->subblock); + if (!subblock) + return std::nullopt; + key.operands.push_back(*subblock); + } else { + OperandTypeInfo subblock; + subblock.kind = OperandKind::Scalar; + subblock.dtype = "i64"; + subblock.scalarValue = 0; + key.operands.push_back(subblock); + } + } + appendOpContextAttrs(op, key.contextAttrs); + return key; +} + static std::optional buildSpecKey(Operation *op) { + if (isa(op)) + return buildPipeSpecKey(op); + SpecKey key; key.opName = getTileOpName(op).str(); key.targetArch = getTargetArchString(op); @@ -775,14 +1153,9 @@ static void appendJsonDimArray(std::string &json, ArrayRef arr, json += "]"; } -static std::string buildOperandSpecsJson(const SpecKey &key) { - std::string json = "["; - for (size_t i = 0; i < key.operands.size(); ++i) { - const auto &op = key.operands[i]; - if (i > 0) - json += ","; - - if (op.kind == OperandKind::Tile) { +static void appendOperandSpecJson(std::string &json, + const OperandTypeInfo &op) { + if (op.kind == OperandKind::Tile) { json += "{\"kind\":\"tile\",\"dtype\":\"" + op.dtype + "\",\"shape\":"; appendJsonIntArray(json, op.tileShape); json += ",\"valid_shape\":"; @@ -799,11 +1172,13 @@ static std::string buildOperandSpecsJson(const SpecKey &key) { json += ",\"pad_value\":\"0x"; json += llvm::utohexstr(op.pad, /*LowerCase=*/false); json += "\"}}"; - continue; - } + return; + } - if (op.kind == OperandKind::View) { - json += "{\"kind\":\"view\",\"dtype\":\"" + op.dtype + "\",\"shape\":"; + if (op.kind == OperandKind::View || op.kind == OperandKind::PipeEntry) { + json += "{\"kind\":\""; + json += op.kind == OperandKind::PipeEntry ? "pipe_entry" : "view"; + json += "\",\"dtype\":\"" + op.dtype + "\",\"shape\":"; appendJsonDimArray(json, op.viewShape); if (!op.viewStrides.empty()) { json += ",\"strides\":["; @@ -824,23 +1199,75 @@ static std::string buildOperandSpecsJson(const SpecKey &key) { json += "\"}"; } json += "}"; - continue; - } + return; + } - if (op.kind == OperandKind::Vector) { + if (op.kind == OperandKind::Vector) { json += "{\"kind\":\"vector\",\"dtype\":\"" + op.dtype + "\",\"shape\":"; appendJsonIntArray(json, op.vectorShape); json += "}"; - continue; - } + return; + } - // Scalar + if (op.kind == OperandKind::Scalar) { json += "{\"kind\":\"scalar\",\"dtype\":\"" + op.dtype + "\""; if (op.scalarValue) { json += ",\"value\":"; json += std::to_string(*op.scalarValue); } json += "}"; + return; + } + + if (op.kind == OperandKind::Pipe) { + json += "{\"kind\":\"pipe\",\"init_kind\":\"" + + op.pipeInitKind + "\",\"dir_mask\":" + + std::to_string(op.pipeDirMask); + json += ",\"slot_size\":" + std::to_string(op.pipeSlotSize); + json += ",\"slot_num\":" + std::to_string(op.pipeSlotNum); + json += ",\"local_slot_num\":"; + json += op.pipeLocalSlotNum ? std::to_string(*op.pipeLocalSlotNum) : "null"; + json += ",\"flag_base\":" + std::to_string(op.pipeFlagBase); + json += ",\"nosplit\":"; + json += op.pipeNoSplit ? "true" : "false"; + json += ",\"split\":" + std::to_string(op.pipeSplit); + json += ",\"resource_names\":["; + for (auto [index, name] : llvm::enumerate(op.pipeResourceNames)) { + if (index != 0) + json += ","; + json += "\"" + name + "\""; + } + json += "]}"; + return; + } + + if (op.kind == OperandKind::PipeResources) { + json += "{\"kind\":\"pipe_resources\",\"names\":["; + for (auto [index, name] : llvm::enumerate(op.pipeResourceNames)) { + if (index != 0) + json += ","; + json += "\"" + name + "\""; + } + json += "]"; + json += ",\"values\":["; + for (auto [index, resource] : llvm::enumerate(op.pipeResources)) { + if (index != 0) + json += ","; + appendOperandSpecJson(json, resource); + } + json += "]}"; + return; + } + + json += "{\"kind\":\"pipe_state\",\"fields\":[\"i32\",\"i32\"]}"; +} + +static std::string buildOperandSpecsJson(const SpecKey &key) { + std::string json = "["; + for (size_t i = 0; i < key.operands.size(); ++i) { + if (i != 0) + json += ","; + appendOperandSpecJson(json, key.operands[i]); } json += "]"; return json; @@ -857,8 +1284,12 @@ static std::string buildUniqueFunctionBaseName(const SpecKey &key) { for (const auto &op : key.operands) { uniqueName += op.kind == OperandKind::Tile ? "_tile" : op.kind == OperandKind::View ? "_view" + : op.kind == OperandKind::PipeEntry ? "_pipe_entry" : op.kind == OperandKind::Vector ? "_vector" - : "_scalar"; + : op.kind == OperandKind::Scalar ? "_scalar" + : op.kind == OperandKind::Pipe ? "_pipe" + : op.kind == OperandKind::PipeResources ? "_pipe_resources" + : "_pipe_state"; uniqueName += "_" + op.dtype; if (op.kind == OperandKind::Tile) { for (int64_t d : op.tileShape) @@ -869,7 +1300,8 @@ static std::string buildUniqueFunctionBaseName(const SpecKey &key) { uniqueName += "_sl" + std::to_string(op.slayout); uniqueName += "_fr" + std::to_string(op.fractal); uniqueName += "_pd" + llvm::utohexstr(op.pad, /*LowerCase=*/false); - } else if (op.kind == OperandKind::View) { + } else if (op.kind == OperandKind::View || + op.kind == OperandKind::PipeEntry) { uniqueName += "_ms_" + op.viewMemorySpace; uniqueName += "_shape"; for (int64_t d : op.viewShape) @@ -884,6 +1316,25 @@ static std::string buildUniqueFunctionBaseName(const SpecKey &key) { uniqueName += "_" + std::to_string(d); } else if (op.kind == OperandKind::Scalar && op.scalarValue) { uniqueName += "_sv" + std::to_string(*op.scalarValue); + } else if (op.kind == OperandKind::Pipe) { + uniqueName += "_" + op.pipeInitKind; + uniqueName += "_d" + std::to_string(op.pipeDirMask); + uniqueName += "_ss" + std::to_string(op.pipeSlotSize); + uniqueName += "_sn" + std::to_string(op.pipeSlotNum); + uniqueName += "_fb" + std::to_string(op.pipeFlagBase); + uniqueName += "_ns" + std::to_string(op.pipeNoSplit); + uniqueName += "_sp" + std::to_string(op.pipeSplit); + } else if (op.kind == OperandKind::PipeResources) { + for (const auto &resource : op.pipeResources) { + uniqueName += "_" + resource.dtype; + if (resource.kind == OperandKind::Tile) { + for (int64_t dim : resource.tileShape) + uniqueName += "_" + dimSuffix(dim); + } else if (resource.kind == OperandKind::View) { + for (int64_t dim : resource.viewShape) + uniqueName += "_" + dimSuffix(dim); + } + } } } for (const auto &[attrName, attrValue] : key.contextAttrs) @@ -1361,25 +1812,110 @@ LogicalResult ExpandState::expandTileOpsInFunction(func::FuncOp func, return failure(); } - // Replace tile op with func.call. For view operands whose caller type - // (memref) differs from the template parameter type (tensor_view / - // partition_tensor_view), insert an unrealized_conversion_cast bridge. - // FoldTileBufIntrinsics will later resolve these casts. + // Replace the TileOp with a helper call. Pipe configuration is metadata + // only: the helper receives the initializer's storage values, PipeState, + // and (for push/pop) an explicit subblock id instead of !pto.pipe. builder.setInsertionPoint(op); SmallVector operands; auto fnArgTypes = dslFn.getArgumentTypes(); - for (unsigned i = 0; i < op->getNumOperands(); ++i) { - Value operand = op->getOperand(i); - if (i < fnArgTypes.size() && operand.getType() != fnArgTypes[i]) { + auto appendOperand = [&](Value operand) -> LogicalResult { + if (operands.size() >= fnArgTypes.size()) + return op->emitError( + "ExpandTileOp: generated helper has fewer physical arguments " + "than the pipe operation requires"); + Type expectedType = fnArgTypes[operands.size()]; + if (operand.getType() != expectedType) { operand = bridgeOperandToType(builder, op->getLoc(), operand, - fnArgTypes[i]); + expectedType); } operands.push_back(operand); + return success(); + }; + + const bool isPipeOp = isa(op); + if (!isPipeOp) { + for (Value operand : op->getOperands()) { + if (failed(appendOperand(operand))) + return failure(); + } + } else { + FailureOr wire = getPipeWireInfo(op); + if (failed(wire)) + return failure(); + if (auto push = dyn_cast(op)) { + if (failed(appendOperand(push.getTile()))) + return failure(); + } else if (auto pop = dyn_cast(op)) { + if (failed(appendOperand(pop.getTile()))) + return failure(); + } else if (auto alloc = dyn_cast(op)) { + if (failed(appendOperand(alloc.getEntry()))) + return failure(); + } else if (auto free = dyn_cast(op)) { + if (free.getEntry() && failed(appendOperand(free.getEntry()))) + return failure(); + } + + for (const auto &resource : wire->resources) { + if (failed(appendOperand(resource.second))) + return failure(); + } + if (failed(appendOperand(wire->state))) + return failure(); + + if (isa(op)) { + Value subblock = wire->subblock; + if (!subblock) + subblock = builder + .create(op->getLoc(), 0, 64) + .getResult(); + if (failed(appendOperand(subblock))) + return failure(); + } + } + if (operands.size() != fnArgTypes.size()) { + op->emitError("ExpandTileOp: generated helper physical argument count " + "does not match the expanded operation"); + return failure(); } builder.create(op->getLoc(), dslFn, operands); op->erase(); } + // The pipe handle is intentionally not passed through the TileLib helper. + // Once all owned pipe ops have become calls, erase the now-dead initializer + // so the VPTO path cannot retain !pto.pipe or EmitC's TPipe dependency. + SmallVector deadInitializers; + func.walk([&](Operation *candidate) { + if (!candidate->hasAttr("pto.pipe_tilelib_owned")) + return; + if (!isa(candidate)) + return; + if (llvm::all_of(candidate->getResults(), + [](Value result) { return result.use_empty(); })) + deadInitializers.push_back(candidate); + }); + for (Operation *initializer : deadInitializers) + initializer->erase(); + + bool residualPipe = false; + func.walk([&](Operation *candidate) { + if (candidate->hasAttr("pto.pipe_tilelib_owned") && + (isa(candidate) || + llvm::any_of(candidate->getResultTypes(), [](Type type) { + return isa(type); + }))) + residualPipe = true; + }); + if (residualPipe) { + func.emitError( + "ExpandTileOp left feature-owned !pto.pipe operations after expansion"); + return failure(); + } + return success(); } diff --git a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp index 2747850b70..b0c5a8e681 100644 --- a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp +++ b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp @@ -170,6 +170,27 @@ static std::pair findSetValidShapeOverride(Value tileBuf) { return {Value(), Value()}; } +static std::optional +findActiveTileAddress(Value tileBuf, Operation *user) { + Operation *producer = tileBuf.getDefiningOp(); + if (!producer || !user || producer->getBlock() != user->getBlock()) + return std::nullopt; + + pto::TAssignOp latestAssign; + for (Operation *candidate : tileBuf.getUsers()) { + auto assign = dyn_cast(candidate); + if (!assign || assign.getTile() != tileBuf || + assign->getBlock() != user->getBlock() || + !assign->isBeforeInBlock(user)) + continue; + if (!latestAssign || latestAssign->isBeforeInBlock(assign)) + latestAssign = assign; + } + if (!latestAssign) + return std::nullopt; + return latestAssign.getAddr(); +} + // Walk through value-conversion casts that preserve the underlying producer. // ExpandTileOp / PTOInlineLibCall and other passes may wrap tile_buf values in // UnrealizedConversionCastOp / arith.index_cast / memref.cast when bridging @@ -222,6 +243,17 @@ static std::optional resolveTileHandle(Value tileBuf, } tileBuf = unwrapBridgingCasts(tileBuf); + if (auto assign = tileBuf.getDefiningOp()) { + auto sourceInfo = resolveTileHandle(assign.getTile(), user); + if (!sourceInfo) + return std::nullopt; + + // TAssign rebinds the handle, so its result must use the new address even + // when the source was originally materialized from a memref. + return TileHandleInfo{Value(), assign.getAddr(), sourceInfo->validRow, + sourceInfo->validCol, sourceInfo->config}; + } + if (auto alloc = tileBuf.getDefiningOp()) { auto tileTy = dyn_cast(alloc.getResult().getType()); if (!tileTy) { @@ -229,7 +261,19 @@ static std::optional resolveTileHandle(Value tileBuf, "FoldTileBufIntrinsics: pto.alloc_tile must produce !pto.tile_buf"); return std::nullopt; } - return TileHandleInfo{Value(), alloc.getAddr(), alloc.getValidRow(), + Value addr = alloc.getAddr(); + if (!addr) { + auto reboundAddr = findActiveTileAddress(tileBuf, user); + if (!reboundAddr) { + user->emitError( + "FoldTileBufIntrinsics: pto.alloc_tile without an addr operand " + "must be rebound by a dominating pto.tassign before its " + "tile-buffer intrinsic use"); + return std::nullopt; + } + addr = *reboundAddr; + } + return TileHandleInfo{Value(), addr, alloc.getValidRow(), alloc.getValidCol(), tileTy.getConfigAttr()}; } @@ -239,6 +283,25 @@ static std::optional resolveTileHandle(Value tileBuf, materialize.getConfig()}; } + if (auto declare = tileBuf.getDefiningOp()) { + auto tileTy = dyn_cast(declare.getTile().getType()); + if (!tileTy) { + user->emitError( + "FoldTileBufIntrinsics: pto.declare_tile must produce !pto.tile_buf"); + return std::nullopt; + } + auto addr = findActiveTileAddress(tileBuf, user); + if (!addr) { + user->emitError( + "FoldTileBufIntrinsics: pto.declare_tile must be rebound by a " + "dominating pto.tassign before its tile-buffer intrinsic use"); + return std::nullopt; + } + auto [validRow, validCol] = findSetValidShapeOverride(tileBuf); + return TileHandleInfo{Value(), *addr, validRow, validCol, + tileTy.getConfigAttr()}; + } + if (auto reshape = tileBuf.getDefiningOp()) { auto sourceInfo = resolveTileHandle(reshape.getSrc(), user); if (!sourceInfo) @@ -261,7 +324,8 @@ static std::optional resolveTileHandle(Value tileBuf, user->emitError("FoldTileBufIntrinsics: expected tile_buf to be defined by " "the active materialized tile-handle bridge " - "(pto.alloc_tile, pto.materialize_tile, or pto.treshape, " + "(pto.alloc_tile, pto.materialize_tile, a rebound " + "pto.declare_tile, or pto.treshape, " "or a pto.fusion_region result that yields one of them)"); return std::nullopt; } @@ -314,8 +378,10 @@ static std::optional traceViewChain(Value tensorView, auto subviewOp = memrefVal.getDefiningOp(); if (!subviewOp) { + if (auto rcOp = memrefVal.getDefiningOp()) + return ViewChain{castOp, {}, rcOp, rcOp.getSource()}; user->emitError("FoldTileBufIntrinsics: expected memref to be defined by " - "memref.subview, got ") + "memref.subview or memref.reinterpret_cast, got ") << (memrefVal.getDefiningOp() ? memrefVal.getDefiningOp()->getName().getStringRef() : StringRef("block argument")); @@ -435,6 +501,56 @@ static Value computeLinearOffset(OpBuilder &builder, Location loc, return rcPart ? rcPart : svPart; } +struct PTOViewAddress { + Value basePtr; + Value linearOffset; +}; + +// PTODSL TileLib ViewSpec parameters are ABI-lowered as memrefs while their +// callers retain TensorView / PartitionTensorView values. Recover the original +// descriptor here so the inlined helper can use the same address-folding path +// as a mainline-lowered memref view. +static std::optional +tracePTOViewAddress(Value value, OpBuilder &builder, Location loc, + Operation *user) { + if (isa(value.getType())) { + auto bridge = value.getDefiningOp(); + if (!bridge || bridge.getNumOperands() != 1 || + !isa( + bridge.getOperand(0).getType())) + return std::nullopt; + value = bridge.getOperand(0); + } + + Value source = value; + SmallVector offsets; + if (auto part = value.getDefiningOp()) { + source = part.getSource(); + offsets.append(part.getOffsets().begin(), part.getOffsets().end()); + } else if (!isa(value.getType())) { + return std::nullopt; + } + + auto make = source.getDefiningOp(); + if (!make) { + user->emitError("FoldTileBufIntrinsics: expected PTODSL view bridge to " + "originate from pto.make_tensor_view"); + return std::nullopt; + } + if (offsets.size() != make.getStrides().size()) { + user->emitError("FoldTileBufIntrinsics: partition view rank does not " + "match its source tensor-view strides"); + return std::nullopt; + } + + SmallVector strides; + strides.append(make.getStrides().begin(), make.getStrides().end()); + return PTOViewAddress{make.getPtr(), + computeLinearOffset(builder, loc, + ArrayRef{}, offsets, + strides)}; +} + struct FoldTileBufIntrinsicsPass : public pto::impl::FoldTileBufIntrinsicsBase { using FoldTileBufIntrinsicsBase::FoldTileBufIntrinsicsBase; @@ -703,7 +819,9 @@ struct FoldTileBufIntrinsicsPass return signalPassFailure(); } - auto svTy = cast(chain->subview.getType()); + auto viewValue = chain->subview ? chain->subview.getResult() + : chain->reinterpretCast.getResult(); + auto svTy = cast(viewValue.getType()); if (dimIdx < 0 || dimIdx >= svTy.getRank()) { dimOp.emitError( "FoldTileBufIntrinsics: get_tensor_view_dim dim index out of " @@ -719,7 +837,9 @@ struct FoldTileBufIntrinsicsPass svTy.getDimSize(dimIdx)); } else { replacement = getValueOrCreateConstant( - builder, dimOp.getLoc(), chain->subview.getMixedSizes()[dimIdx]); + builder, dimOp.getLoc(), + chain->subview ? chain->subview.getMixedSizes()[dimIdx] + : chain->reinterpretCast.getMixedSizes()[dimIdx]); } dimOp.getResult().replaceAllUsesWith(replacement); @@ -739,7 +859,9 @@ struct FoldTileBufIntrinsicsPass return signalPassFailure(); } - auto svTy = cast(chain->subview.getType()); + auto viewValue = chain->subview ? chain->subview.getResult() + : chain->reinterpretCast.getResult(); + auto svTy = cast(viewValue.getType()); if (dimIdx < 0 || dimIdx >= svTy.getRank()) { strideOp.emitError( "FoldTileBufIntrinsics: get_tensor_view_stride dim index out of " @@ -751,7 +873,8 @@ struct FoldTileBufIntrinsicsPass Value replacement = computeResultStride( builder, strideOp.getLoc(), chain->reinterpretCast.getMixedStrides()[dimIdx], - chain->subview.getMixedStrides()[dimIdx]); + chain->subview ? chain->subview.getMixedStrides()[dimIdx] + : OpFoldResult(builder.getIndexAttr(1))); strideOp.getResult().replaceAllUsesWith(replacement); strideOp.erase(); @@ -760,6 +883,32 @@ struct FoldTileBufIntrinsicsPass if (shouldFoldAddrFamily(*mode)) { for (auto addrOp : tvAddrOps) { + builder.setInsertionPoint(addrOp); + if (auto directView = tracePTOViewAddress(addrOp.getSrc(), builder, + addrOp.getLoc(), addrOp)) { + auto resultPtrType = dyn_cast(addrOp.getDst().getType()); + if (!resultPtrType) { + addrOp.emitError( + "FoldTileBufIntrinsics: PTODSL view bridge requires " + "tensor_view_addr to produce !pto.ptr"); + return signalPassFailure(); + } + + Value basePtr = directView->basePtr; + if (basePtr.getType() != resultPtrType) + basePtr = builder.create(addrOp.getLoc(), + resultPtrType, basePtr); + Value replacement = + directView->linearOffset + ? builder.create(addrOp.getLoc(), + resultPtrType, basePtr, + directView->linearOffset) + : basePtr; + addrOp.getDst().replaceAllUsesWith(replacement); + addrOp.erase(); + continue; + } + auto chain = traceViewChain(addrOp.getSrc(), addrOp); if (!chain) return signalPassFailure(); @@ -783,11 +932,17 @@ struct FoldTileBufIntrinsicsPass return signalPassFailure(); } - Value linearOffset = - computeLinearOffset(builder, addrOp.getLoc(), - chain->reinterpretCast.getMixedOffsets(), - chain->subview.getMixedOffsets(), - chain->reinterpretCast.getMixedStrides()); + Value linearOffset = chain->subview + ? computeLinearOffset( + builder, addrOp.getLoc(), + chain->reinterpretCast.getMixedOffsets(), + chain->subview.getMixedOffsets(), + chain->reinterpretCast.getMixedStrides()) + : computeLinearOffset( + builder, addrOp.getLoc(), + chain->reinterpretCast.getMixedOffsets(), + ArrayRef{}, + chain->reinterpretCast.getMixedStrides()); Value basePtr = builder.create( addrOp.getLoc(), resultPtrType, chain->baseMemref); diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index 90a00ebc74..a9ffb86e2e 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/lib/PTO/Transforms/InsertTemplateAttributes.cpp @@ -10,6 +10,7 @@ #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Support/PythonExecutable.h" #include "PTO/Transforms/Passes.h" +#include "Utils.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" @@ -54,6 +55,11 @@ namespace { constexpr llvm::StringLiteral kCandidatesAttr = "candidates"; +static bool isUnifiedInternalPipeOperation(Operation *operation) { + return isa(operation); +} + struct CandidateMetadata { int64_t id; std::string name; @@ -151,6 +157,8 @@ static std::string getMemorySpaceString(pto::PartitionTensorViewType) { return "gm"; } +static std::string getMemorySpaceString(pto::TensorViewType) { return "gm"; } + static StringRef getBLayoutString(pto::BLayout layout) { return layout == pto::BLayout::ColMajor ? "col_major" : "row_major"; } @@ -481,6 +489,22 @@ static bool tryAppendPrecisionType( static void appendOpContextAttrs( Operation *op, SmallVectorImpl> &attrs) { + if (isUnifiedInternalPipeOperation(op)) { + auto func = op->getParentOfType(); + auto kernelKind = + func ? func->getAttrOfType( + pto::FunctionKernelKindAttr::name) + : pto::FunctionKernelKindAttr(); + if (!kernelKind) { + op->emitError( + "requires a cube or vector kernel_kind for TileLib pipe expansion"); + return; + } + attrs.emplace_back( + "kernel_kind", + kernelKind.getKernelKind() == pto::FunctionKernelKind::Cube ? "cube" + : "vector"); + } if (auto tcvt = dyn_cast(op)) { if (auto roundMode = getTCvtRoundModeString(tcvt)) attrs.emplace_back("round_mode", *roundMode); @@ -571,6 +595,35 @@ static void appendTileOperandSpecJson(std::string &json, json += "\"}}"; } +static std::optional getBoundTileType(Value value) { + auto bind = value.getDefiningOp(); + if (!bind) + return std::nullopt; + + auto memrefType = dyn_cast(value.getType()); + if (!memrefType || memrefType.getRank() != 2) + return std::nullopt; + + SmallVector validShape(memrefType.getShape().begin(), + memrefType.getShape().end()); + auto updateValidDim = [&](Value validDim, unsigned dimension) { + if (!validDim) + return; + int64_t staticValue = 0; + validShape[dimension] = getStaticIntFromValue(validDim, staticValue) + ? staticValue + : ShapedType::kDynamic; + }; + updateValidDim(bind.getValidRow(), 0); + updateValidDim(bind.getValidCol(), 1); + + return pto::TileBufType::get(value.getType().getContext(), + memrefType.getShape(), + memrefType.getElementType(), + memrefType.getMemorySpace(), validShape, + bind.getConfig()); +} + static void appendViewOperandSpecJson(std::string &json, Value operand, MemRefType memrefType) { std::string dtype = getDtypeString(memrefType.getElementType()); @@ -621,6 +674,33 @@ static void appendViewOperandSpecJson(std::string &json, Value operand, json += "}"; } +static void appendViewOperandSpecJson(std::string &json, Value operand, + pto::TensorViewType viewType, + StringRef kind = "view") { + std::string dtype = getDtypeString(viewType.getElementType()); + json += "{\"kind\":\"" + kind.str() + "\",\"dtype\":\"" + + dtype + "\",\"shape\":"; + SmallVector shape; + SmallVector strides; + populatePTOViewShapeAndStrides(operand, shape, strides); + if (shape.empty()) + shape.assign(viewType.getShape().begin(), viewType.getShape().end()); + appendJsonDimArray(json, shape); + if (!strides.empty()) { + json += ",\"strides\":"; + appendJsonDimArray(json, strides); + } + json += ",\"memory_space\":\""; + json += getMemorySpaceString(viewType); + json += "\""; + if (auto layout = getViewLayoutString(resolveViewLayout(operand))) { + json += ",\"config\":{\"layout\":\""; + json += *layout; + json += "\"}"; + } + json += "}"; +} + static void appendVectorOperandSpecJson(std::string &json, VectorType vectorType) { std::string dtype = getDtypeString(vectorType.getElementType()); @@ -640,63 +720,272 @@ static void appendScalarOperandSpecJson(std::string &json, Value operand) { json += "}"; } -static std::optional -buildOperandSpecsJson(Operation *operation) { - std::string json = "["; - for (auto [index, operand] : llvm::enumerate(operation->getOperands())) { +static LogicalResult appendOperandSpecJson(std::string &json, + Operation *operation, + Value operand) { + Type type = operand.getType(); + if (auto tileType = dyn_cast(type)) { + if (getDtypeString(tileType.getElementType()).empty()) + return operation->emitError( + "InsertTemplateAttributes encountered an unsupported tile dtype"); + appendTileOperandSpecJson(json, tileType); + return success(); + } + + if (auto memrefType = dyn_cast(type)) { + if (getDtypeString(memrefType.getElementType()).empty()) + return operation->emitError( + "InsertTemplateAttributes encountered an unsupported view dtype"); + appendViewOperandSpecJson(json, operand, memrefType); + return success(); + } + + if (auto viewType = dyn_cast(type)) { + if (getDtypeString(viewType.getElementType()).empty()) + return operation->emitError( + "InsertTemplateAttributes encountered an unsupported view dtype"); + appendViewOperandSpecJson(json, operand, viewType); + return success(); + } + + if (auto viewType = dyn_cast(type)) { + if (getDtypeString(viewType.getElementType()).empty()) + return operation->emitError( + "InsertTemplateAttributes encountered an unsupported view dtype"); + appendViewOperandSpecJson(json, operand, viewType); + return success(); + } + + if (auto vectorType = dyn_cast(type)) { + if (getDtypeString(vectorType.getElementType()).empty()) + return operation->emitError( + "InsertTemplateAttributes encountered an unsupported vector dtype"); + appendVectorOperandSpecJson(json, vectorType); + return success(); + } + + if (!getDtypeString(type).empty()) { + appendScalarOperandSpecJson(json, operand); + return success(); + } + + return operation->emitError( + "InsertTemplateAttributes encountered an unsupported operand type ") + << type; +} + +struct PipeWireInfo { + bool isL2G2L = false; + int64_t dirMask = 0; + int64_t slotSize = 0; + int64_t slotNum = 0; + std::optional localSlotNum; + int64_t flagBase = 0; + bool noSplit = false; + int64_t split = 0; + SmallVector> resources; +}; + +static FailureOr getPipeWireInfo(Operation *operation) { + Value pipe; + Value state; + int64_t split = 0; + if (auto push = dyn_cast(operation)) { + pipe = push.getPipeHandle(); + state = push.getPipeState(); + split = push.getSplit(); + } else if (auto pop = dyn_cast(operation)) { + pipe = pop.getPipeHandle(); + state = pop.getPipeState(); + split = pop.getSplit(); + } else if (auto alloc = dyn_cast(operation)) { + pipe = alloc.getPipeHandle(); + state = alloc.getPipeState(); + split = alloc.getSplit(); + } else if (auto free = dyn_cast(operation)) { + pipe = free.getPipeHandle(); + state = free.getPipeState(); + split = free.getSplit(); + } else if (auto drain = dyn_cast(operation)) { + pipe = drain.getPipeHandle(); + state = drain.getPipeState(); + split = drain.getSplit(); + } else { + return operation->emitError("is not a unified internal pipe operation"); + } + + if (!state) + return operation->emitError( + "requires a materialized !pto.struct pipe_state"); + + Operation *init = mlir::pto::getPipeInitDef(pipe); + if (!init) + return operation->emitError("cannot resolve its pipe initializer"); + if (mlir::pto::getPipeInitAccPushEpilogue(init)) + return operation->emitError( + "does not support acc_push_epilogue for TileLib pipe expansion"); + + PipeWireInfo info; + info.split = split; + auto addResource = [&](StringRef name, Value value) { + if (value) + info.resources.emplace_back(name.str(), value); + }; + if (auto l2l = dyn_cast(init)) { + if (!l2l.getFlagBaseAttr()) + return operation->emitError( + "requires resolved pipe flag_base before TileLib expansion"); + info.dirMask = l2l.getDirMask(); + info.slotSize = l2l.getSlotSize(); + info.slotNum = l2l.getSlotNum(); + info.flagBase = l2l.getFlagBaseAttr().getInt(); + info.noSplit = l2l.getNosplitAttr() && l2l.getNosplitAttr().getValue(); + addResource("local_addr", l2l.getLocalAddr()); + addResource("peer_local_addr", l2l.getPeerLocalAddr()); + return info; + } + + auto l2g2l = dyn_cast(init); + if (!l2g2l) + return operation->emitError("requires an l2l or l2g2l pipe initializer"); + if (!l2g2l.getFlagBaseAttr()) + return operation->emitError( + "requires resolved pipe flag_base before TileLib expansion"); + info.isL2G2L = true; + info.dirMask = l2g2l.getDirMask(); + info.slotSize = l2g2l.getSlotSize(); + info.slotNum = l2g2l.getSlotNum(); + if (auto attr = l2g2l.getLocalSlotNumAttr()) + info.localSlotNum = attr.getInt(); + info.flagBase = l2g2l.getFlagBaseAttr().getInt(); + info.noSplit = l2g2l.getNosplitAttr() && l2g2l.getNosplitAttr().getValue(); + addResource("gm_addr", l2g2l.getGmAddr()); + addResource("local_addr", l2g2l.getLocalAddr()); + addResource("peer_local_addr", l2g2l.getPeerLocalAddr()); + return info; +} + +static void appendPipeSpecJson(std::string &json, const PipeWireInfo &info) { + json += "{\"kind\":\"pipe\",\"init_kind\":\""; + json += info.isL2G2L ? "l2g2l" : "l2l"; + json += "\",\"dir_mask\":" + std::to_string(info.dirMask); + json += ",\"slot_size\":" + std::to_string(info.slotSize); + json += ",\"slot_num\":" + std::to_string(info.slotNum); + json += ",\"local_slot_num\":"; + json += info.localSlotNum ? std::to_string(*info.localSlotNum) : "null"; + json += ",\"flag_base\":" + std::to_string(info.flagBase); + json += ",\"nosplit\":"; + json += info.noSplit ? "true" : "false"; + json += ",\"split\":" + std::to_string(info.split); + json += ",\"resource_names\":["; + for (auto [index, resource] : llvm::enumerate(info.resources)) { if (index != 0) json += ","; + json += "\"" + resource.first + "\""; + } + json += "]}"; +} - Type type = operand.getType(); - if (auto tileType = dyn_cast(type)) { - if (getDtypeString(tileType.getElementType()).empty()) { - operation->emitError( - "InsertTemplateAttributes encountered an unsupported tile dtype"); - return std::nullopt; - } - appendTileOperandSpecJson(json, tileType); - continue; - } +static std::optional +buildPipeOperandSpecsJson(Operation *operation) { + FailureOr info = getPipeWireInfo(operation); + if (failed(info)) + return std::nullopt; - if (auto memrefType = dyn_cast(type)) { - if (getDtypeString(memrefType.getElementType()).empty()) { - operation->emitError( - "InsertTemplateAttributes encountered an unsupported view dtype"); - return std::nullopt; - } - appendViewOperandSpecJson(json, operand, memrefType); - continue; + std::string json = "["; + auto appendComma = [&]() { + if (json.size() != 1) + json += ","; + }; + auto appendPipeEntry = [&](Value value) -> LogicalResult { + appendComma(); + if (auto tileType = getBoundTileType(value)) { + if (getDtypeString(tileType->getElementType()).empty()) + return operation->emitError( + "InsertTemplateAttributes encountered an unsupported bound pipe tile dtype"); + appendTileOperandSpecJson(json, *tileType); + return success(); } + auto viewType = dyn_cast(value.getType()); + if (!viewType) + return appendOperandSpecJson(json, operation, value); + if (!pto::isDeclaredGlobalOrTAssignResult(value)) + return operation->emitError( + "requires a pto.declare_global tensor_view entry, optionally " + "rebound by pto.tassign, for pipe TileLib expansion"); + if (getDtypeString(viewType.getElementType()).empty()) + return operation->emitError( + "InsertTemplateAttributes encountered an unsupported pipe entry dtype"); + appendViewOperandSpecJson(json, value, viewType, "pipe_entry"); + return success(); + }; - if (auto viewType = dyn_cast(type)) { - if (getDtypeString(viewType.getElementType()).empty()) { - operation->emitError( - "InsertTemplateAttributes encountered an unsupported view dtype"); + if (auto push = dyn_cast(operation)) { + if (failed(appendPipeEntry(push.getTile()))) + return std::nullopt; + } else if (auto pop = dyn_cast(operation)) { + if (failed(appendPipeEntry(pop.getTile()))) + return std::nullopt; + } else if (auto alloc = dyn_cast(operation)) { + if (failed(appendPipeEntry(alloc.getEntry()))) + return std::nullopt; + } else if (auto free = dyn_cast(operation)) { + if (free.getEntry() && failed(appendPipeEntry(free.getEntry()))) + return std::nullopt; + } + + appendComma(); + appendPipeSpecJson(json, *info); + appendComma(); + json += "{\"kind\":\"pipe_resources\",\"names\":["; + for (auto [index, resource] : llvm::enumerate(info->resources)) { + if (index != 0) + json += ","; + json += "\"" + resource.first + "\""; + } + json += "],\"values\":["; + for (auto [index, resource] : llvm::enumerate(info->resources)) { + if (index != 0) + json += ","; + if (failed(appendOperandSpecJson(json, operation, resource.second))) + return std::nullopt; + } + json += "]}"; + appendComma(); + json += "{\"kind\":\"pipe_state\",\"fields\":[\"i32\",\"i32\"]}"; + + if (auto push = dyn_cast(operation)) { + appendComma(); + if (Value subblock = push.getAivSubblockid()) { + if (failed(appendOperandSpecJson(json, operation, subblock))) return std::nullopt; - } - appendViewOperandSpecJson(json, operand, viewType); - continue; + } else { + json += "{\"kind\":\"scalar\",\"dtype\":\"i64\",\"value\":0}"; } - - if (auto vectorType = dyn_cast(type)) { - if (getDtypeString(vectorType.getElementType()).empty()) { - operation->emitError( - "InsertTemplateAttributes encountered an unsupported vector dtype"); + } else if (auto pop = dyn_cast(operation)) { + appendComma(); + if (Value subblock = pop.getAivSubblockid()) { + if (failed(appendOperandSpecJson(json, operation, subblock))) return std::nullopt; - } - appendVectorOperandSpecJson(json, vectorType); - continue; + } else { + json += "{\"kind\":\"scalar\",\"dtype\":\"i64\",\"value\":0}"; } + } + json += "]"; + return json; +} - if (!getDtypeString(type).empty()) { - appendScalarOperandSpecJson(json, operand); - continue; - } +static std::optional +buildOperandSpecsJson(Operation *operation) { + if (isUnifiedInternalPipeOperation(operation)) + return buildPipeOperandSpecsJson(operation); - operation->emitError( - "InsertTemplateAttributes encountered an unsupported operand type ") - << type; - return std::nullopt; + std::string json = "["; + for (auto [index, operand] : llvm::enumerate(operation->getOperands())) { + if (index != 0) + json += ","; + if (failed(appendOperandSpecJson(json, operation, operand))) + return std::nullopt; } json += "]"; return json; @@ -948,6 +1237,11 @@ struct InsertTemplateAttributesPass module.walk([&](Operation *operation) { if (isa(operation)) return; + const bool isPipeOperation = isUnifiedInternalPipeOperation(operation); + if (skipPipeOps && isPipeOperation) + return; + if (pipeOnly != isPipeOperation) + return; if (isa(operation)) tileOperations.push_back(operation); }); diff --git a/lib/PTO/Transforms/PTOLowerFrontendPipeOpsPass.cpp b/lib/PTO/Transforms/PTOLowerFrontendPipeOpsPass.cpp index abe25d058c..21bf66b977 100644 --- a/lib/PTO/Transforms/PTOLowerFrontendPipeOpsPass.cpp +++ b/lib/PTO/Transforms/PTOLowerFrontendPipeOpsPass.cpp @@ -462,7 +462,7 @@ static LogicalResult lowerFrontendDataOps(func::FuncOp funcOp, alloc.getEntry().getType()); propagateGlobalTensorStrides(decl, handles.c2vSlotStrides, rewriter); rewriter.create(alloc.getLoc(), decl.getEntry(), - handles.c2vPipe, alloc.getSplitAttr()); + handles.c2vPipe, Value{}, alloc.getSplitAttr()); rewriter.replaceOp(alloc, decl.getEntry()); continue; } @@ -481,7 +481,7 @@ static LogicalResult lowerFrontendDataOps(func::FuncOp funcOp, alloc.getEntry().getType()); propagateGlobalTensorStrides(decl, handles.v2cSlotStrides, rewriter); rewriter.create(alloc.getLoc(), decl.getEntry(), - handles.v2cPipe, alloc.getSplitAttr()); + handles.v2cPipe, Value{}, alloc.getSplitAttr()); rewriter.replaceOp(alloc, decl.getEntry()); continue; } @@ -497,7 +497,7 @@ static LogicalResult lowerFrontendDataOps(func::FuncOp funcOp, return failure(); } rewriter.replaceOpWithNewOp(push, push.getTile(), handles.c2vPipe, - Value{}, push.getSplitAttr()); + Value{}, Value{}, push.getSplitAttr()); continue; } @@ -513,6 +513,7 @@ static LogicalResult lowerFrontendDataOps(func::FuncOp funcOp, } rewriter.replaceOpWithNewOp(push, push.getTile(), handles.v2cPipe, push.getAivSubblockid(), + Value{}, push.getSplitAttr()); continue; } @@ -543,7 +544,7 @@ static LogicalResult lowerFrontendDataOps(func::FuncOp funcOp, } } rewriter.create(pop.getLoc(), entry, handles.c2vPipe, - pop.getAivSubblockid(), pop.getSplitAttr()); + pop.getAivSubblockid(), Value{}, pop.getSplitAttr()); rewriter.replaceOp(pop, entry); continue; } @@ -574,7 +575,7 @@ static LogicalResult lowerFrontendDataOps(func::FuncOp funcOp, } } rewriter.create(pop.getLoc(), entry, handles.v2cPipe, - Value{}, pop.getSplitAttr()); + Value{}, Value{}, pop.getSplitAttr()); rewriter.replaceOp(pop, entry); continue; } @@ -591,6 +592,7 @@ static LogicalResult lowerFrontendDataOps(func::FuncOp funcOp, } rewriter.replaceOpWithNewOp(free, free.getEntry(), handles.c2vPipe, + Value{}, free.getSplitAttr()); continue; } @@ -607,6 +609,7 @@ static LogicalResult lowerFrontendDataOps(func::FuncOp funcOp, } rewriter.replaceOpWithNewOp(free, free.getEntry(), handles.v2cPipe, + Value{}, free.getSplitAttr()); } diff --git a/lib/PTO/Transforms/PTOMaterializePipeStatePass.cpp b/lib/PTO/Transforms/PTOMaterializePipeStatePass.cpp new file mode 100644 index 0000000000..1fd2cb9feb --- /dev/null +++ b/lib/PTO/Transforms/PTOMaterializePipeStatePass.cpp @@ -0,0 +1,195 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#include "PTO/IR/PTO.h" +#include "PTO/Transforms/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/Dominance.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Pass/Pass.h" + +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace pto { +namespace func = ::mlir::func; +#define GEN_PASS_DEF_PTOMATERIALIZEPIPESTATE +#include "PTO/Transforms/Passes.h.inc" +} // namespace pto +} // namespace mlir + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +constexpr llvm::StringLiteral kPipeTileLibOwnedAttr = + "pto.pipe_tilelib_owned"; + +static StructType getPipeStateType(MLIRContext *context) { + Builder builder(context); + return StructType::get(context, {builder.getI32Type(), builder.getI32Type()}); +} + +static bool isStatefulPipeOperation(Operation *op) { + return isa(op); +} + +static Value getPipeHandle(Operation *op) { + if (auto alloc = dyn_cast(op)) + return alloc.getPipeHandle(); + if (auto push = dyn_cast(op)) + return push.getPipeHandle(); + if (auto pop = dyn_cast(op)) + return pop.getPipeHandle(); + if (auto free = dyn_cast(op)) + return free.getPipeHandle(); + return {}; +} + +static Value getPipeState(Operation *op) { + if (auto alloc = dyn_cast(op)) + return alloc.getPipeState(); + if (auto push = dyn_cast(op)) + return push.getPipeState(); + if (auto pop = dyn_cast(op)) + return pop.getPipeState(); + if (auto free = dyn_cast(op)) + return free.getPipeState(); + return {}; +} + +static void setPipeState(Operation *op, Value state) { + if (auto alloc = dyn_cast(op)) { + alloc.getPipeStateMutable().assign(state); + return; + } + if (auto push = dyn_cast(op)) { + push.getPipeStateMutable().assign(state); + return; + } + if (auto pop = dyn_cast(op)) { + pop.getPipeStateMutable().assign(state); + return; + } + cast(op).getPipeStateMutable().assign(state); +} + +static bool isInternalPipeInitializer(Operation *op) { + return isa(op); +} + +static FailureOr getTerminalDrainSplit(Operation *initializer) { + BoolAttr noSplit; + if (auto l2l = dyn_cast(initializer)) + noSplit = l2l.getNosplitAttr(); + else if (auto l2g2l = dyn_cast(initializer)) + noSplit = l2g2l.getNosplitAttr(); + + if (!noSplit) { + initializer->emitOpError() + << "requires resolved 'nosplit' before terminal pipe drain " + "materialization"; + return failure(); + } + + // TPipe::~TPipe uses the pipe-level IsNoSplit configuration for cleanup, + // rather than the split value of any individual producer operation. + return noSplit.getValue() ? 0 : 1; +} + +struct PTOMaterializePipeStatePass + : public mlir::pto::impl::PTOMaterializePipeStateBase< + PTOMaterializePipeStatePass> { + void runOnOperation() override { + func::FuncOp funcOp = getOperation(); + MLIRContext *context = funcOp.getContext(); + DominanceInfo dominance(funcOp); + SmallVector initializers; + SmallVector returns; + + funcOp.walk([&](Operation *op) { + if (isInternalPipeInitializer(op)) + initializers.push_back(op); + if (auto returnOp = dyn_cast(op)) + returns.push_back(returnOp); + }); + + for (Operation *initializer : initializers) { + Value pipe = initializer->getResult(0); + SmallVector pipeUsers; + bool hasProducer = false; + for (Operation *user : pipe.getUsers()) { + if (user->getParentOfType() != funcOp || + !isStatefulPipeOperation(user)) + continue; + if (getPipeHandle(user) != pipe) + continue; + if (Value existingState = getPipeState(user)) { + user->emitOpError() + << "already has a pipe_state; " + "pto-materialize-pipe-state expects unmaterialized pipe IR"; + return signalPassFailure(); + } + pipeUsers.push_back(user); + hasProducer |= isa(user); + } + + if (pipeUsers.empty()) + continue; + + OpBuilder builder(initializer); + builder.setInsertionPointAfter(initializer); + StructType stateType = getPipeStateType(context); + auto state = builder.create(initializer->getLoc(), stateType).getS(); + Value zero = builder.create(initializer->getLoc(), 0, 32); + auto path0 = DenseI64ArrayAttr::get(context, {0}); + auto path1 = DenseI64ArrayAttr::get(context, {1}); + builder.create(initializer->getLoc(), state, path0, zero); + builder.create(initializer->getLoc(), state, path1, zero); + + for (Operation *user : pipeUsers) { + setPipeState(user, state); + user->setAttr(kPipeTileLibOwnedAttr, UnitAttr::get(context)); + } + initializer->setAttr(kPipeTileLibOwnedAttr, UnitAttr::get(context)); + + if (!hasProducer) + continue; + FailureOr drainSplit = getTerminalDrainSplit(initializer); + if (failed(drainSplit)) + return signalPassFailure(); + for (func::ReturnOp returnOp : returns) { + if (!dominance.dominates(initializer, returnOp)) { + initializer->emitOpError() + << "does not dominate a return that requires terminal pipe drain"; + return signalPassFailure(); + } + builder.setInsertionPoint(returnOp); + auto drain = + builder.create(returnOp.getLoc(), pipe, state, *drainSplit); + drain->setAttr(kPipeTileLibOwnedAttr, UnitAttr::get(context)); + } + } + } +}; + +} // namespace + +namespace mlir { +namespace pto { + +std::unique_ptr createPTOMaterializePipeStatePass() { + return std::make_unique(); +} + +} // namespace pto +} // namespace mlir diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index bfe0221c24..493bd5cbcd 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -135,7 +135,33 @@ static Type normalizeGEPElementTypeForLLVMLowering(Type type, return normalizePayloadTypeForLLVMLowering(type, builder); } +static Type convertVPTOType(Type type, Builder &builder); + +static Type convertVPTOStructStorageType(pto::StructType type, + Builder &builder) { + SmallVector fields; + fields.reserve(type.getNumFields()); + for (Type fieldType : type.getFieldTypes()) { + if (auto nested = dyn_cast(fieldType)) { + fields.push_back(convertVPTOStructStorageType(nested, builder)); + continue; + } + Type converted = convertVPTOType(fieldType, builder); + if (!converted) + return {}; + fields.push_back(converted); + } + return LLVM::LLVMStructType::getLiteral(builder.getContext(), fields); +} + static Type convertVPTOType(Type type, Builder &builder) { + if (auto structType = dyn_cast(type)) { + if (!convertVPTOStructStorageType(structType, builder)) + return {}; + return LLVM::LLVMPointerType::get(builder.getContext()); + } + if (isa(type)) + return LLVM::LLVMPointerType::get(builder.getContext()); if (auto vecType = dyn_cast(type)) { Type elementType = normalizePayloadTypeForLLVMLowering(vecType.getElementType(), builder); @@ -180,7 +206,8 @@ static unsigned getNaturalByteAlignment(Type type) { } static bool hasVPTOConvertibleType(Type type) { - return isa(type); + return isa(type); } static bool hasVPTOConvertibleType(TypeRange types) { @@ -9713,6 +9740,199 @@ class ConvertVPTOUnrealizedCastOp final } }; +static FailureOr getVPTOStructFieldAddress( + ConversionPatternRewriter &rewriter, Location loc, Value base, + pto::StructType rootType, ArrayRef path) { + if (path.empty()) + return failure(); + + auto pointerType = dyn_cast(base.getType()); + if (!pointerType) + return failure(); + + Builder builder(rewriter.getContext()); + Type storageType = convertVPTOStructStorageType(rootType, builder); + if (!storageType) + return failure(); + + Value cursor = base; + pto::StructType cursorType = rootType; + for (auto [depth, fieldIndex] : llvm::enumerate(path)) { + if (fieldIndex < 0 || + static_cast(fieldIndex) >= cursorType.getNumFields()) + return failure(); + + cursor = rewriter.create( + loc, pointerType, storageType, cursor, + ArrayRef{0, static_cast(fieldIndex)}); + + if (depth + 1 == path.size()) + return cursor; + + auto nestedType = dyn_cast( + cursorType.getFieldType(static_cast(fieldIndex))); + if (!nestedType) + return failure(); + cursorType = nestedType; + storageType = convertVPTOStructStorageType(cursorType, builder); + if (!storageType) + return failure(); + } + return failure(); +} + +class ConvertPtoDeclareStructOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::DeclareStructOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + (void)adaptor; + auto pointerType = dyn_cast( + getTypeConverter()->convertType(op.getS().getType())); + if (!pointerType) + return rewriter.notifyMatchFailure(op, + "failed to convert !pto.struct type"); + + Builder builder(rewriter.getContext()); + Type storageType = + convertVPTOStructStorageType(op.getS().getType(), builder); + if (!storageType) + return rewriter.notifyMatchFailure( + op, "failed to convert !pto.struct storage type"); + + Value one = rewriter.create( + op.getLoc(), rewriter.getI64Type(), rewriter.getI64IntegerAttr(1)); + rewriter.replaceOpWithNewOp(op, pointerType, storageType, + one, /*alignment=*/0); + return success(); + } +}; + +class ConvertPtoStructGetOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::StructGetOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type resultType = getTypeConverter()->convertType(op.getValue().getType()); + if (!resultType) + return rewriter.notifyMatchFailure(op, "failed to convert struct field type"); + + FailureOr address = getVPTOStructFieldAddress( + rewriter, op.getLoc(), adaptor.getS(), op.getS().getType(), op.getPath()); + if (failed(address)) + return rewriter.notifyMatchFailure(op, "failed to address struct field"); + + rewriter.replaceOpWithNewOp( + op, resultType, *address, getNaturalByteAlignment(resultType)); + return success(); + } +}; + +class ConvertPtoStructSetOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::StructSetOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr address = getVPTOStructFieldAddress( + rewriter, op.getLoc(), adaptor.getS(), op.getS().getType(), op.getPath()); + if (failed(address)) + return rewriter.notifyMatchFailure(op, "failed to address struct field"); + + rewriter.create( + op.getLoc(), adaptor.getValue(), *address, + getNaturalByteAlignment(adaptor.getValue().getType())); + rewriter.eraseOp(op); + return success(); + } +}; + +class ConvertPtoDeclareGlobalOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::DeclareGlobalOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + (void)adaptor; + auto descriptorType = dyn_cast( + getTypeConverter()->convertType(op.getEntry().getType())); + if (!descriptorType) + return rewriter.notifyMatchFailure( + op, "failed to convert declared global descriptor type"); + + Type storageType = LLVM::LLVMPointerType::get( + rewriter.getContext(), static_cast(pto::AddressSpace::GM)); + Value one = rewriter.create( + op.getLoc(), rewriter.getI64Type(), rewriter.getI64IntegerAttr(1)); + rewriter.replaceOpWithNewOp( + op, descriptorType, storageType, one, /*alignment=*/8); + return success(); + } +}; + +class ConvertPtoTAssignOp final : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::TAssignOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + if (!isa(op.getTile().getType())) + return failure(); + auto descriptorType = dyn_cast(adaptor.getTile().getType()); + if (!descriptorType) + return rewriter.notifyMatchFailure(op, + "expected a global descriptor pointer"); + + Value address = adaptor.getAddr(); + if (address.getType().isIndex()) + address = rewriter.create( + op.getLoc(), rewriter.getI64Type(), address); + if (!address.getType().isInteger(64)) + return rewriter.notifyMatchFailure(op, + "expected an i64 global descriptor address"); + + Type valueType = LLVM::LLVMPointerType::get( + rewriter.getContext(), static_cast(pto::AddressSpace::GM)); + Value pointer = + rewriter.create(op.getLoc(), valueType, address); + rewriter.create(op.getLoc(), pointer, adaptor.getTile(), + /*alignment=*/8); + rewriter.replaceOp(op, adaptor.getTile()); + return success(); + } +}; + +class ConvertPtoTensorViewAddrOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::TensorViewAddrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + if (!isa(op.getSrc().getType())) + return failure(); + Type resultType = getTypeConverter()->convertType(op.getDst().getType()); + if (!resultType) + return rewriter.notifyMatchFailure(op, + "failed to convert tensor view address type"); + rewriter.replaceOpWithNewOp(op, resultType, adaptor.getSrc(), + /*alignment=*/8); + return success(); + } +}; + class ConvertArithSelectOp final : public OpConversionPattern { public: ConvertArithSelectOp(TypeConverter &typeConverter, MLIRContext *context) @@ -10658,6 +10878,10 @@ static LogicalResult lowerVPTOTypes(ModuleOp module, llvm::raw_ostream &diagOS) }); populateVPTOStructuralTypePatterns(typeConverter, patterns, target); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); patterns.add(typeConverter, context); patterns.add fields; + fields.reserve(type.getNumFields()); + for (Type fieldType : type.getFieldTypes()) { + if (auto nested = dyn_cast(fieldType)) { + fields.push_back(convertVPTOStructStorageType(nested, builder)); + continue; + } + Type converted = convertVPTOType(fieldType, builder); + if (!converted) + return {}; + fields.push_back(converted); + } + return LLVM::LLVMStructType::getLiteral(builder.getContext(), fields); +} + static Type convertVPTOType(Type type, Builder &builder) { + if (auto structType = dyn_cast(type)) { + if (!convertVPTOStructStorageType(structType, builder)) + return {}; + return LLVM::LLVMPointerType::get(builder.getContext()); + } + if (isa(type)) + return LLVM::LLVMPointerType::get(builder.getContext()); if (auto vecType = dyn_cast(type)) { Type elementType = normalizePayloadTypeForLLVMLowering(vecType.getElementType(), builder); @@ -182,7 +208,8 @@ static unsigned getNaturalByteAlignment(Type type) { } static bool hasVPTOConvertibleType(Type type) { - return isa(type); + return isa(type); } static bool hasVPTOConvertibleType(TypeRange types) { @@ -10327,6 +10354,199 @@ class ConvertVPTOUnrealizedCastOp final } }; +static FailureOr getVPTOStructFieldAddress( + ConversionPatternRewriter &rewriter, Location loc, Value base, + pto::StructType rootType, ArrayRef path) { + if (path.empty()) + return failure(); + + auto pointerType = dyn_cast(base.getType()); + if (!pointerType) + return failure(); + + Builder builder(rewriter.getContext()); + Type storageType = convertVPTOStructStorageType(rootType, builder); + if (!storageType) + return failure(); + + Value cursor = base; + pto::StructType cursorType = rootType; + for (auto [depth, fieldIndex] : llvm::enumerate(path)) { + if (fieldIndex < 0 || + static_cast(fieldIndex) >= cursorType.getNumFields()) + return failure(); + + cursor = rewriter.create( + loc, pointerType, storageType, cursor, + ArrayRef{0, static_cast(fieldIndex)}); + + if (depth + 1 == path.size()) + return cursor; + + auto nestedType = dyn_cast( + cursorType.getFieldType(static_cast(fieldIndex))); + if (!nestedType) + return failure(); + cursorType = nestedType; + storageType = convertVPTOStructStorageType(cursorType, builder); + if (!storageType) + return failure(); + } + return failure(); +} + +class ConvertPtoDeclareStructOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::DeclareStructOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + (void)adaptor; + auto pointerType = dyn_cast( + getTypeConverter()->convertType(op.getS().getType())); + if (!pointerType) + return rewriter.notifyMatchFailure(op, + "failed to convert !pto.struct type"); + + Builder builder(rewriter.getContext()); + Type storageType = + convertVPTOStructStorageType(op.getS().getType(), builder); + if (!storageType) + return rewriter.notifyMatchFailure( + op, "failed to convert !pto.struct storage type"); + + Value one = rewriter.create( + op.getLoc(), rewriter.getI64Type(), rewriter.getI64IntegerAttr(1)); + rewriter.replaceOpWithNewOp(op, pointerType, storageType, + one, /*alignment=*/0); + return success(); + } +}; + +class ConvertPtoStructGetOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::StructGetOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type resultType = getTypeConverter()->convertType(op.getValue().getType()); + if (!resultType) + return rewriter.notifyMatchFailure(op, "failed to convert struct field type"); + + FailureOr address = getVPTOStructFieldAddress( + rewriter, op.getLoc(), adaptor.getS(), op.getS().getType(), op.getPath()); + if (failed(address)) + return rewriter.notifyMatchFailure(op, "failed to address struct field"); + + rewriter.replaceOpWithNewOp( + op, resultType, *address, getNaturalByteAlignment(resultType)); + return success(); + } +}; + +class ConvertPtoStructSetOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::StructSetOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + FailureOr address = getVPTOStructFieldAddress( + rewriter, op.getLoc(), adaptor.getS(), op.getS().getType(), op.getPath()); + if (failed(address)) + return rewriter.notifyMatchFailure(op, "failed to address struct field"); + + rewriter.create( + op.getLoc(), adaptor.getValue(), *address, + getNaturalByteAlignment(adaptor.getValue().getType())); + rewriter.eraseOp(op); + return success(); + } +}; + +class ConvertPtoDeclareGlobalOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::DeclareGlobalOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + (void)adaptor; + auto descriptorType = dyn_cast( + getTypeConverter()->convertType(op.getEntry().getType())); + if (!descriptorType) + return rewriter.notifyMatchFailure( + op, "failed to convert declared global descriptor type"); + + Type storageType = LLVM::LLVMPointerType::get( + rewriter.getContext(), static_cast(pto::AddressSpace::GM)); + Value one = rewriter.create( + op.getLoc(), rewriter.getI64Type(), rewriter.getI64IntegerAttr(1)); + rewriter.replaceOpWithNewOp( + op, descriptorType, storageType, one, /*alignment=*/8); + return success(); + } +}; + +class ConvertPtoTAssignOp final : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::TAssignOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + if (!isa(op.getTile().getType())) + return failure(); + auto descriptorType = dyn_cast(adaptor.getTile().getType()); + if (!descriptorType) + return rewriter.notifyMatchFailure(op, + "expected a global descriptor pointer"); + + Value address = adaptor.getAddr(); + if (address.getType().isIndex()) + address = rewriter.create( + op.getLoc(), rewriter.getI64Type(), address); + if (!address.getType().isInteger(64)) + return rewriter.notifyMatchFailure(op, + "expected an i64 global descriptor address"); + + Type valueType = LLVM::LLVMPointerType::get( + rewriter.getContext(), static_cast(pto::AddressSpace::GM)); + Value pointer = + rewriter.create(op.getLoc(), valueType, address); + rewriter.create(op.getLoc(), pointer, adaptor.getTile(), + /*alignment=*/8); + rewriter.replaceOp(op, adaptor.getTile()); + return success(); + } +}; + +class ConvertPtoTensorViewAddrOp final + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::TensorViewAddrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + if (!isa(op.getSrc().getType())) + return failure(); + Type resultType = getTypeConverter()->convertType(op.getDst().getType()); + if (!resultType) + return rewriter.notifyMatchFailure(op, + "failed to convert tensor view address type"); + rewriter.replaceOpWithNewOp(op, resultType, adaptor.getSrc(), + /*alignment=*/8); + return success(); + } +}; + class ConvertPtoTileBufAddrOp final : public OpConversionPattern { public: @@ -11437,6 +11657,10 @@ static LogicalResult lowerVPTOTypes(ModuleOp module, llvm::raw_ostream &diagOS) }); populateVPTOStructuralTypePatterns(typeConverter, patterns, target); + patterns.add(typeConverter, context); + patterns.add(typeConverter, context); patterns.add(typeConverter, context); diff --git a/ptodsl/docs/user_guide/10-sync-ops.md b/ptodsl/docs/user_guide/10-sync-ops.md index b16131d1cb..e82c21c979 100644 --- a/ptodsl/docs/user_guide/10-sync-ops.md +++ b/ptodsl/docs/user_guide/10-sync-ops.md @@ -403,7 +403,7 @@ pto.wait_cross_flag(pto.Pipe.FIX, 0) ### 10.5.2 Intra-block sync: `set_intra_flag`, `wait_intra_flag` -The Cube unit (matrix pipeline) has a dedicated synchronization channel separate from the standard pipe-flag mechanism used by MTE and Vector pipelines. `set_intra_flag` and `wait_intra_flag` synchronize Cube and Vector within the same block, ensuring that shared UB tile data is not accessed before the producer finishes. +The Cube unit (matrix pipeline) has a dedicated synchronization channel separate from the standard pipe-flag mechanism used by MTE and Vector pipelines. `set_intra_flag` and `wait_intra_flag` synchronize Cube, Vector, and MTE pipelines within the same block, ensuring that shared UB tile data is not accessed before the producer finishes. Cross-core flags use the public `0`-`7` event range. A5 intra-block sync uses a separate physical event range, described below. @@ -418,7 +418,7 @@ Unlike `wait_cross_flag`, `wait_intra_flag` only stalls the specified pipeline | Parameter | Type | Description | |-----------|------|-------------| -| `pipe` | `Pipe` | Trigger endpoint for the synchronization event. The public DSL accepts `Pipe.FIX` and `Pipe.MTE3` here. | +| `pipe` | `Pipe` | Trigger endpoint for the synchronization event. The public DSL accepts `Pipe.FIX`, `Pipe.MTE1`, `Pipe.MTE2`, `Pipe.MTE3`, `Pipe.V`, and `Pipe.S`. | | `event_id` | `int` or scalar expression | Physical event identifier (`0`–`31` for static IDs). | For A5 mixed C/V writeback code that must notify both AIV subblocks, emit two signals @@ -442,7 +442,7 @@ pto.set_intra_flag(pto.Pipe.MTE3, 0) | Parameter | Type | Description | |-----------|------|-------------| -| `pipe` | `Pipe` | Waiting endpoint for the synchronization event. The public DSL accepts `Pipe.FIX` and `Pipe.V` here. | +| `pipe` | `Pipe` | Waiting endpoint for the synchronization event. The public DSL accepts `Pipe.FIX`, `Pipe.MTE1`, `Pipe.MTE2`, `Pipe.MTE3`, `Pipe.V`, and `Pipe.S`. | | `event_id` | `int` or scalar expression | Physical event identifier to wait on (`0`–`31` for static IDs). | For A5 mixed C/V writeback code that must synchronize with both AIV subblocks, wait on both diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index a54c915d84..ab1fa03cd6 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -6014,7 +6014,7 @@ def set_intra_flag(pipe, event_id): _validate_sync_pipe( pipe, context="set_intra_flag(pipe, event_id)", - allowed=("PIPE_FIX", "PIPE_MTE3"), + allowed=("PIPE_FIX", "PIPE_MTE1", "PIPE_MTE2", "PIPE_MTE3", "PIPE_V", "PIPE_S"), ) event_operand = _sync_event_id_operand_in_range( event_id, @@ -6031,7 +6031,7 @@ def wait_intra_flag(pipe, event_id): _validate_sync_pipe( pipe, context="wait_intra_flag(pipe, event_id)", - allowed=("PIPE_FIX", "PIPE_V"), + allowed=("PIPE_FIX", "PIPE_MTE1", "PIPE_MTE2", "PIPE_MTE3", "PIPE_V", "PIPE_S"), ) event_operand = _sync_event_id_operand_in_range( event_id, diff --git a/ptodsl/ptodsl/_runtime/native_build.py b/ptodsl/ptodsl/_runtime/native_build.py index 310b44067e..077f151545 100644 --- a/ptodsl/ptodsl/_runtime/native_build.py +++ b/ptodsl/ptodsl/_runtime/native_build.py @@ -10,6 +10,7 @@ from __future__ import annotations import os +import shlex import subprocess from pathlib import Path @@ -47,6 +48,7 @@ def _run_ptoas( insert_sync: bool | None = None, backend: str | None = None, pto_level: str | None = None, + extra_args: tuple[str, ...] = (), ) -> None: ptoas = resolve_ptoas_binary() cmd = [ @@ -59,6 +61,7 @@ def _run_ptoas( cmd.append(f"--pto-level={pto_level}") if insert_sync is True: cmd.append("--enable-insert-sync") + cmd.extend(extra_args) cmd.extend([ "--enable-tile-op-expand", str(mlir_path), @@ -86,12 +89,17 @@ def _source_ptoas_overrides(module_spec) -> dict: return {"backend": module_spec.backend} +def _configured_ptoas_flags() -> tuple[str, ...]: + return tuple(shlex.split(os.environ.get("PTOAS_FLAGS", ""))) + + def _compile_config_text( *, module_spec, effective_insert_sync: bool, effective_pto_level: str | None, ptoas_overrides: dict, + ptoas_flags: tuple[str, ...], ) -> str: return "\n".join( [ @@ -101,6 +109,7 @@ def _compile_config_text( f"insert_sync={effective_insert_sync}", f"pto_level={effective_pto_level}", f"backend={ptoas_overrides.get('backend')}", + f"ptoas_flags={shlex.join(ptoas_flags)}", "enable_tile_op_expand=True", ] ) @@ -219,11 +228,13 @@ def build_native_library( ) effective_pto_level = _effective_pto_level(mode=module_spec.mode) ptoas_overrides = _source_ptoas_overrides(module_spec) + ptoas_flags = _configured_ptoas_flags() compile_config_text = _compile_config_text( module_spec=module_spec, effective_insert_sync=effective_insert_sync, effective_pto_level=effective_pto_level, ptoas_overrides=ptoas_overrides, + ptoas_flags=ptoas_flags, ) sim_mode = bool(os.environ.get("MSPROF_SIMULATOR_MODE")) link_config_text = "\n".join(runtime_library_flags(sim_mode=sim_mode)) @@ -247,6 +258,7 @@ def build_native_library( target_arch=module_spec.target_arch, insert_sync=effective_insert_sync, pto_level=effective_pto_level, + extra_args=ptoas_flags, **ptoas_overrides, ) diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py index 617fc05578..5d3d0f0b4e 100644 --- a/ptodsl/ptodsl/_tracing/session.py +++ b/ptodsl/ptodsl/_tracing/session.py @@ -834,11 +834,7 @@ def get_or_create_kernel_module_primary_function(self, spec: HelperFunctionSpec, self._attach_ptodsl_logical_name_attr(helper, spec.symbol_name) helper.attributes["sym_visibility"] = StringAttr.get("public") helper.attributes["pto.visibility"] = StringAttr.get("external") - if ( - module_spec.backend == "emitc" - and not module_spec.entry - and module_spec.kernel_kind in {"cube", "vector"} - ): + if not module_spec.entry and module_spec.kernel_kind in {"cube", "vector"}: helper.attributes["pto.kernel_kind"] = Attribute.parse( f"#pto.kernel_kind<{module_spec.kernel_kind}>" ) diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index d48f6e9ff3..a83cad38a6 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -11,7 +11,6 @@ from pathlib import Path import os import re -import sys from tempfile import TemporaryDirectory from importlib.util import module_from_spec, spec_from_file_location from typing import Optional @@ -4294,6 +4293,14 @@ def inline_source_backed_probe(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32): "func.func public @process_tile_module__ptodsl_" in kernel_module_call_text, "kernel-module callee definition should be materialized as a public ABI-specialized symbol", ) + kernel_module_primary_index = kernel_module_call_text.index( + "func.func public @process_tile_module__ptodsl_" + ) + expect( + "pto.kernel_kind = #pto.kernel_kind" + in kernel_module_call_text[kernel_module_primary_index : kernel_module_primary_index + 500], + "VPTO kernel-module primary definitions should preserve their vector kernel kind", + ) expect( 'pto.visibility = "external"' in kernel_module_call_text, "kernel-module ABI-specialized primary definitions should carry explicit external artifact visibility", @@ -4503,7 +4510,16 @@ def fake_artifacts(py_name, ir_function_name, specialization_key): manifest_path=cache_dir / "manifest.json", ) - def fake_run_ptoas(mlir_path, kernel_object, *, target_arch, insert_sync=None, backend=None, pto_level=None): + def fake_run_ptoas( + mlir_path, + kernel_object, + *, + target_arch, + insert_sync=None, + backend=None, + pto_level=None, + extra_args=(), + ): native_build_observations.append( { "mlir_path": mlir_path, @@ -4512,6 +4528,7 @@ def fake_run_ptoas(mlir_path, kernel_object, *, target_arch, insert_sync=None, b "insert_sync": insert_sync, "backend": backend, "pto_level": pto_level, + "extra_args": extra_args, "mlir_text": mlir_path.read_text(encoding="utf-8"), } ) @@ -4594,6 +4611,10 @@ def fake_link_shared_library(launch_object, kernel_object, shared_library, *, ke observation["pto_level"] == expected_pto_level, f"{label} native build should derive the PTOAS level from the authored mode", ) + expect( + observation["extra_args"] == (), + f"{label} native build should not add PTOAS options when PTOAS_FLAGS is unset", + ) expect( observation["mlir_text"] == compiled.mlir_text(), f"{label} native build should hand the backend-partitioned container MLIR to ptoas unchanged", @@ -4645,6 +4666,25 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None): "native build should still pass the shared PTOAS compile inputs and output path", ) ptoas_cmds.clear() + with mock.patch.dict(os.environ, {"PTOAS_FLAGS": "--tile-lib-backend=ptodsl"}), mock.patch.object( + native_build_runtime, "resolve_ptoas_binary", return_value=Path("/tmp/fake-ptoas") + ), mock.patch.object(native_build_runtime, "_run", side_effect=fake_run_ptoas_cmd): + configured_ptoas_flags = native_build_runtime._configured_ptoas_flags() + native_build_runtime._run_ptoas( + mlir_path, + kernel_object, + target_arch="a5", + extra_args=configured_ptoas_flags, + ) + expect( + ptoas_cmds[0].count("--tile-lib-backend=ptodsl") == 1, + "native build should forward PTOAS_FLAGS as individual PTOAS arguments", + ) + expect( + configured_ptoas_flags == ("--tile-lib-backend=ptodsl",), + "native build should retain PTOAS_FLAGS in its cache-keyed PTOAS arguments", + ) + ptoas_cmds.clear() with mock.patch.object(native_build_runtime, "resolve_ptoas_binary", return_value=Path("/tmp/fake-ptoas")), mock.patch.object( native_build_runtime, "_run", side_effect=fake_run_ptoas_cmd ): diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index 91d4011a88..f78688990a 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -9,7 +9,7 @@ import unittest import inspect from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock, call, patch import ptodsl._ops as _ops import ptodsl._pipe_namespace as _pipe_namespace @@ -704,8 +704,8 @@ def test_sync_facades_reject_illegal_pipe_endpoints(self): cases = [ (_ops.set_cross_flag, (pto.Pipe.V, 0), "set_cross_flag(pipe, event_id)", "", ""), (_ops.wait_cross_flag, (pto.Pipe.MTE3, 0), "wait_cross_flag(pipe, event_id)", "", ""), - (_ops.set_intra_flag, (pto.Pipe.V, 0), "set_intra_flag(pipe, event_id)", ", ", ""), - (_ops.wait_intra_flag, (pto.Pipe.MTE3, 0), "wait_intra_flag(pipe, event_id)", ", ", ""), + (_ops.set_intra_flag, (pto.Pipe.M, 0), "set_intra_flag(pipe, event_id)", ", , , , , ", ""), + (_ops.wait_intra_flag, (pto.Pipe.ALL, 0), "wait_intra_flag(pipe, event_id)", ", , , , , ", ""), ] with patch.object(_ops._pto, "sync_set") as sync_set_op, \ @@ -722,6 +722,29 @@ def test_sync_facades_reject_illegal_pipe_endpoints(self): sync_set_op.assert_not_called() sync_wait_op.assert_not_called() + def test_intra_sync_accepts_pipe_tilelib_endpoints(self): + supported_pipes = ( + pto.Pipe.FIX, + pto.Pipe.MTE1, + pto.Pipe.MTE2, + pto.Pipe.MTE3, + pto.Pipe.V, + pto.Pipe.S, + ) + + with patch.object(_ops, "_pipe_attr", side_effect=lambda pipe: f"pipe:{pipe}") as pipe_attr, \ + patch.object(_ops._pto, "sync_set") as sync_set_op, \ + patch.object(_ops._pto, "sync_wait") as sync_wait_op: + for pipe in supported_pipes: + with self.subTest(kind="set", pipe=pipe): + _ops.set_intra_flag(pipe, 31) + with self.subTest(kind="wait", pipe=pipe): + _ops.wait_intra_flag(pipe, 16) + + self.assertEqual(pipe_attr.call_count, 2 * len(supported_pipes)) + self.assertEqual(sync_set_op.call_count, len(supported_pipes)) + self.assertEqual(sync_wait_op.call_count, len(supported_pipes)) + def test_intra_sync_mixed_writeback_event_ranges(self): with patch.object(_ops, "_pipe_attr", side_effect=lambda pipe: f"pipe:{pipe}") as pipe_attr, \ patch.object(_ops._pto, "sync_set") as sync_set_op, \ @@ -756,21 +779,20 @@ def test_pipe_namespace_and_buffer_helpers_are_exposed(self): self.assertTrue(hasattr(pto, "gm_ptr"), "gm_ptr") - def test_global_pipe_methods_dispatch_to_matching_frontend_ops(self): - alloc_entry = object() - pop_entry = object() + def test_global_pipe_methods_preserve_nonzero_split_per_transaction(self): + alloc_entries = (object(), object()) + pop_entries = (object(), object()) with make_context(): gm_slot_type = _pipe_namespace._pto.TensorViewType.get([16, 16], F32Type.get()) gm_slot = SimpleNamespace(type=gm_slot_type) with patch.object(_pipe_namespace, "unwrap_surface_value", side_effect=_identity), \ - patch.object(_pipe_namespace, "wrap_surface_value", side_effect=_identity), \ - patch.object(_pipe_namespace, "_infer_unambiguous_global_slot_size", return_value=1024): + patch.object(_pipe_namespace, "wrap_surface_value", side_effect=_identity): pipe = pto.pipe.c2v( gm_slot_tensor=gm_slot, + slot_size=1024, id=7, - nosplit=True, ) self.assertEqual(pipe.id, 7) @@ -779,32 +801,58 @@ def test_global_pipe_methods_dispatch_to_matching_frontend_ops(self): with patch.object(_pipe_namespace._pto, "AicInitializePipeOp") as aic_init, \ patch.object(_pipe_namespace._pto, "AivInitializePipeOp") as aiv_init, \ - patch.object(_pipe_namespace._pto, "TAllocToAivOp", return_value=SimpleNamespace(result=alloc_entry)) as alloc_op, \ + patch.object(_pipe_namespace._pto, "TAllocToAivOp", side_effect=[ + SimpleNamespace(result=alloc_entries[0]), + SimpleNamespace(result=alloc_entries[1]), + ]) as alloc_op, \ patch.object(_pipe_namespace._pto, "TPushToAivOp") as push_op, \ - patch.object(_pipe_namespace._pto, "TPopFromAicOp", return_value=SimpleNamespace(result=pop_entry)) as pop_op, \ + patch.object(_pipe_namespace._pto, "TPopFromAicOp", side_effect=[ + SimpleNamespace(result=pop_entries[0]), + SimpleNamespace(result=pop_entries[1]), + ]) as pop_op, \ patch.object(_pipe_namespace._pto, "TFreeFromAicOp") as free_op, \ patch.object(_pipe_namespace, "wrap_surface_value", side_effect=_identity), \ patch.object(_pipe_namespace, "unwrap_surface_value", side_effect=_identity): pipe.init_cube() pipe.init_simd() - alloc_result = pipe.alloc() - pipe.push(alloc_result, split=1) - pop_result = pipe.pop() - pipe.free(pop_result, split=2) + alloc_up_down = pipe.alloc(split=1) + pipe.push(alloc_up_down, split=1) + alloc_left_right = pipe.alloc(split=2) + pipe.push(alloc_left_right, split=2) + pop_up_down = pipe.pop(split=1) + pipe.free(pop_up_down, split=1) + pop_left_right = pipe.pop(split=2) + pipe.free(pop_left_right, split=2) expected_init_kwargs = { "id": 7, - "nosplit": True, "gm_slot_tensor": gm_slot, } aic_init.assert_called_once_with(1, 1024, **expected_init_kwargs) aiv_init.assert_called_once_with(1, 1024, **expected_init_kwargs) - alloc_op.assert_called_once_with(gm_slot_type, 0, id=7) - push_op.assert_called_once_with(alloc_entry, 1, id=7) - pop_op.assert_called_once_with(gm_slot_type, 0, id=7) - free_op.assert_called_once_with(2, entry=pop_entry, id=7) - self.assertIs(alloc_result, alloc_entry) - self.assertIs(pop_result, pop_entry) + self.assertEqual( + alloc_op.call_args_list, + [call(gm_slot_type, 1, id=7), call(gm_slot_type, 2, id=7)], + ) + self.assertEqual( + push_op.call_args_list, + [call(alloc_entries[0], 1, id=7), call(alloc_entries[1], 2, id=7)], + ) + self.assertEqual( + pop_op.call_args_list, + [call(gm_slot_type, 1, id=7), call(gm_slot_type, 2, id=7)], + ) + self.assertEqual( + free_op.call_args_list, + [ + call(1, entry=pop_entries[0], id=7), + call(2, entry=pop_entries[1], id=7), + ], + ) + self.assertIs(alloc_up_down, alloc_entries[0]) + self.assertIs(alloc_left_right, alloc_entries[1]) + self.assertIs(pop_up_down, pop_entries[0]) + self.assertIs(pop_left_right, pop_entries[1]) def test_local_pipe_constructors_route_consumer_buffers(self): c2v_buf = object() diff --git a/test/lit/pto/pipe_state_verifier_invalid.pto b/test/lit/pto/pipe_state_verifier_invalid.pto new file mode 100644 index 0000000000..c219d3aac4 --- /dev/null +++ b/test/lit/pto/pipe_state_verifier_invalid.pto @@ -0,0 +1,74 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a5 %s 2>&1 | FileCheck %s + +module { + func.func @mismatched_pipe_state() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %slots = pto.declare_global -> !pto.tensor_view<16x16xf32> + %pipe = "pto.initialize_l2g2l_pipe"(%slots) <{ + dir_mask = 1 : i8, + slot_size = 1024 : i32, + slot_num = 8 : i32, + flag_base = 0 : i32, + nosplit = false, + operandSegmentSizes = array + }> : (!pto.tensor_view<16x16xf32>) -> !pto.pipe + %entry = pto.declare_global -> !pto.tensor_view<16x16xf32> + %state0 = pto.declare_struct -> !pto.struct + %state1 = pto.declare_struct -> !pto.struct + pto.talloc(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pipe_state(%state0 : !pto.struct) + pto.tpush(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pipe_state(%state1 : !pto.struct) + return + } + + func.func @mismatched_drain_split() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %slots = pto.declare_global -> !pto.tensor_view<16x16xf32> + %pipe = "pto.initialize_l2g2l_pipe"(%slots) <{ + dir_mask = 1 : i8, + slot_size = 1024 : i32, + slot_num = 8 : i32, + flag_base = 2 : i32, + nosplit = true, + operandSegmentSizes = array + }> : (!pto.tensor_view<16x16xf32>) -> !pto.pipe + %entry = pto.declare_global -> !pto.tensor_view<16x16xf32> + %state = pto.declare_struct -> !pto.struct + pto.tpush(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 0} + pipe_state(%state : !pto.struct) + pto.tdrain(%pipe, %state : !pto.pipe, !pto.struct) {split = 1} + return + } + + func.func @partial_pipe_state() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %slots = pto.declare_global -> !pto.tensor_view<16x16xf32> + %pipe = "pto.initialize_l2g2l_pipe"(%slots) <{ + dir_mask = 1 : i8, + slot_size = 1024 : i32, + slot_num = 8 : i32, + flag_base = 4 : i32, + nosplit = false, + operandSegmentSizes = array + }> : (!pto.tensor_view<16x16xf32>) -> !pto.pipe + %entry = pto.declare_global -> !pto.tensor_view<16x16xf32> + %state = pto.declare_struct -> !pto.struct + pto.talloc(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pipe_state(%state : !pto.struct) + pto.tpush(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + return + } +} + +// CHECK: error: 'pto.tpush' op requires all stateful users of the same pipe to share one pipe_state +// CHECK: error: 'pto.tdrain' op requires split = 0 to match the pipe initializer's nosplit = true +// CHECK: error: 'pto.tpush' op requires every stateful user of a materialized pipe to carry the shared pipe_state diff --git a/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/__init__.py b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/__init__.py new file mode 100644 index 0000000000..8433607418 --- /dev/null +++ b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. diff --git a/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/__init__.py b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/__init__.py new file mode 100644 index 0000000000..8433607418 --- /dev/null +++ b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. diff --git a/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/serving/__init__.py b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/serving/__init__.py new file mode 100644 index 0000000000..8433607418 --- /dev/null +++ b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/serving/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. diff --git a/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/serving/daemon.py b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/serving/daemon.py new file mode 100644 index 0000000000..b35e571a5e --- /dev/null +++ b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/serving/daemon.py @@ -0,0 +1,279 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""Test-only daemon for PipeSpec metadata and compiler expansion ABI.""" + +import argparse +import json +import os +import signal +import socket + + +def _recv_exactly(sock, length): + chunks = [] + while length: + chunk = sock.recv(length) + if not chunk: + raise ConnectionError("socket closed mid-message") + chunks.append(chunk) + length -= len(chunk) + return b"".join(chunks) + + +def _recv_message(sock): + length = int.from_bytes(_recv_exactly(sock, 4), byteorder="big") + return json.loads(_recv_exactly(sock, length).decode("utf-8")) + + +def _send_message(sock, payload): + encoded = json.dumps(payload).encode("utf-8") + sock.sendall(len(encoded).to_bytes(4, byteorder="big") + encoded) + + +def _expect_pipe_abi(params, request_counts): + target = params.get("target") + op = params.get("op") + operands = params.get("operand_specs") + if target != "a5" or not isinstance(operands, list): + raise ValueError("expected an A5 pipe metadata request") + + pipe_index = 0 if op == "pto.tdrain" else 1 + if len(operands) <= pipe_index + 2: + raise ValueError(f"incomplete PipeSpec operands for {op}") + + pipe = operands[pipe_index] + if pipe.get("init_kind") == "l2l": + _expect_local_pipe_abi(op, operands, pipe_index, pipe, request_counts) + return + + expected_splits = { + "pto.talloc": (1, 0), + "pto.tpush": (1, 2), + "pto.tpop": (1, 0), + "pto.tfree": (1, 0), + "pto.tdrain": (1,), + } + if op not in expected_splits: + raise ValueError(f"unexpected metadata operation: {op!r}") + + request_key = ("l2g2l", op) + request_index = request_counts.get(request_key, 0) + request_counts[request_key] = request_index + 1 + allowed_splits = expected_splits[op] + if request_index >= len(allowed_splits): + raise ValueError(f"unexpected duplicate metadata request for {op}") + + if op != "pto.tdrain" and operands[0].get("kind") != "pipe_entry": + raise ValueError(f"{op} must serialize its entry as pipe_entry") + + is_no_split_pipe = ( + op in ("pto.talloc", "pto.tpop", "pto.tfree") and request_index == 1 + ) + expected_pipe = { + "kind": "pipe", + "init_kind": "l2g2l", + "dir_mask": 1, + "slot_size": 1024, + "slot_num": 8, + "local_slot_num": None, + "flag_base": 4, + "nosplit": is_no_split_pipe, + "split": allowed_splits[request_index], + "resource_names": ["gm_addr"], + } + if pipe != expected_pipe: + raise ValueError(f"unexpected PipeSpec for {op}: {pipe!r}") + + resources = operands[pipe_index + 1] + if resources.get("kind") != "pipe_resources" or resources.get("names") != [ + "gm_addr" + ] or len(resources.get("values", ())) != 1: + raise ValueError(f"unexpected PipeResources for {op}: {resources!r}") + if operands[pipe_index + 2] != { + "kind": "pipe_state", + "fields": ["i32", "i32"], + }: + raise ValueError(f"missing PipeState for {op}") + + if op in ("pto.tpush", "pto.tpop"): + subblock = operands[pipe_index + 3] + if subblock != {"kind": "scalar", "dtype": "i64", "value": 0}: + raise ValueError(f"unexpected default subblock for {op}: {subblock!r}") + + +def _expect_local_pipe_abi(op, operands, pipe_index, pipe, request_counts): + if op not in ("pto.tpush", "pto.tpop", "pto.tdrain"): + raise ValueError(f"unexpected local pipe operation: {op!r}") + + request_key = ("l2l", op) + if request_counts.get(request_key, 0): + raise ValueError(f"unexpected duplicate local metadata request for {op}") + request_counts[request_key] = 1 + + if op != "pto.tdrain" and operands[0].get("kind") != "tile": + raise ValueError(f"{op} must serialize its local entry as tile") + + expected_pipe = { + "kind": "pipe", + "init_kind": "l2l", + "dir_mask": 3, + "slot_size": 512, + "slot_num": 4, + "local_slot_num": None, + "flag_base": 8, + "nosplit": False, + "split": 1, + "resource_names": ["local_addr", "peer_local_addr"], + } + if pipe != expected_pipe: + raise ValueError(f"unexpected local PipeSpec for {op}: {pipe!r}") + + resources = operands[pipe_index + 1] + expected_resources = { + "kind": "pipe_resources", + "names": ["local_addr", "peer_local_addr"], + "values": [ + {"kind": "scalar", "dtype": "i32", "value": 17}, + {"kind": "scalar", "dtype": "i32", "value": 29}, + ], + } + if resources != expected_resources: + raise ValueError(f"unexpected local PipeResources for {op}: {resources!r}") + if operands[pipe_index + 2] != { + "kind": "pipe_state", + "fields": ["i32", "i32"], + }: + raise ValueError(f"missing local PipeState for {op}") + + if op in ("pto.tpush", "pto.tpop"): + subblock = operands[pipe_index + 3] + if subblock != {"kind": "scalar", "dtype": "i64", "value": 7}: + raise ValueError(f"unexpected explicit subblock for {op}: {subblock!r}") + + +def _metadata_response(): + return { + "candidates": { + "mock_pipe_template": { + "id": 42, + "name": "mock_pipe_template", + "loop_depth": 0, + "is_post_update": False, + "has_tail": False, + } + } + } + + +def _mlir_type_for_spec(spec): + kind = spec.get("kind") + if kind in ("pipe_entry", "view"): + shape = spec.get("shape") + dtype = spec.get("dtype") + if not isinstance(shape, list) or not dtype: + raise ValueError(f"incomplete view spec: {spec!r}") + dims = "x".join("?" if dim == -1 else str(dim) for dim in shape) + return f"!pto.tensor_view<{dims}x{dtype}>" + if kind == "pipe_state": + return "!pto.struct" + if kind == "scalar" and spec.get("dtype"): + return spec["dtype"] + raise ValueError(f"unsupported mock helper operand: {spec!r}") + + +def _instantiate_response(params): + if params.get("candidate_id") != "mock_pipe_template": + raise ValueError(f"unexpected candidate: {params.get('candidate_id')!r}") + operands = params.get("operand_specs") + if params.get("target") != "a5" or not isinstance(operands, list): + raise ValueError("expected an A5 pipe instantiate request") + + argument_types = [] + for operand in operands: + kind = operand.get("kind") + if kind == "pipe": + continue + if kind == "pipe_resources": + argument_types.extend(_mlir_type_for_spec(value) + for value in operand.get("values", ())) + continue + argument_types.append(_mlir_type_for_spec(operand)) + + arguments = ", ".join( + f"%arg{index}: {type_name}" + for index, type_name in enumerate(argument_types) + ) + return ( + 'module attributes {pto.target_arch = "a5"} {\n' + f" func.func @mock_pipe_template({arguments}) " + "attributes {pto.tilelang.instance} {\n" + " return\n" + " }\n" + "}\n" + ) + + +def main(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--socket", required=True) + args = parser.parse_args(argv) + + try: + os.unlink(args.socket) + except FileNotFoundError: + pass + + stopping = False + + def stop(*_): + nonlocal stopping + stopping = True + + signal.signal(signal.SIGTERM, stop) + signal.signal(signal.SIGINT, stop) + request_counts = {} + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as server: + server.bind(args.socket) + server.listen() + server.settimeout(0.1) + try: + while not stopping: + try: + connection, _ = server.accept() + except socket.timeout: + continue + with connection: + try: + request = _recv_message(connection) + method = request.get("method") + params = request.get("params", {}) + if method == "get_metadata": + _expect_pipe_abi(params, request_counts) + result = _metadata_response() + elif method == "instantiate": + result = _instantiate_response(params) + else: + raise ValueError(f"unsupported method: {method!r}") + _send_message( + connection, + {"success": True, "result": result}, + ) + except Exception as error: + _send_message( + connection, + {"success": False, "error": f"{type(error).__name__}: {error}"}, + ) + finally: + try: + os.unlink(args.socket) + except FileNotFoundError: + pass + + +if __name__ == "__main__": + main() diff --git a/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/serving/helper.py b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/serving/helper.py new file mode 100644 index 0000000000..2f41a2bbf3 --- /dev/null +++ b/test/lit/vpto/Inputs/pipe_metadata_mock/ptodsl/tilelib/serving/helper.py @@ -0,0 +1,81 @@ +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +"""One-shot client for the test-only PipeSpec metadata and expansion daemon.""" + +import argparse +import json +import socket +import sys + + +def _recv_exactly(sock, length): + chunks = [] + while length: + chunk = sock.recv(length) + if not chunk: + raise ConnectionError("socket closed mid-message") + chunks.append(chunk) + length -= len(chunk) + return b"".join(chunks) + + +def _recv_message(sock): + length = int.from_bytes(_recv_exactly(sock, 4), byteorder="big") + return json.loads(_recv_exactly(sock, length).decode("utf-8")) + + +def _send_message(sock, payload): + encoded = json.dumps(payload).encode("utf-8") + sock.sendall(len(encoded).to_bytes(4, byteorder="big") + encoded) + + +def main(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--socket", required=True) + parser.add_argument("--target", required=True) + parser.add_argument("--op", required=True) + parser.add_argument("--operand-specs", required=True) + parser.add_argument("--context-attrs", default=None) + parser.add_argument( + "--method", + choices=("instantiate", "get_metadata"), + default="instantiate", + ) + parser.add_argument("--candidate-id", default=None) + args = parser.parse_args(argv) + + params = { + "target": args.target, + "op": args.op, + "operand_specs": json.loads(args.operand_specs), + "context_attrs": json.loads(args.context_attrs) + if args.context_attrs + else {}, + } + if args.method == "instantiate": + params["candidate_id"] = args.candidate_id + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(args.socket) + _send_message( + client, + { + "method": args.method, + "params": params, + }, + ) + response = _recv_message(client) + if not response.get("success"): + raise SystemExit(response.get("error", "metadata request failed")) + if args.method == "get_metadata": + sys.stdout.write(json.dumps(response["result"])) + else: + sys.stdout.write(response["result"]) + + +if __name__ == "__main__": + main() diff --git a/test/lit/vpto/fold_tile_buf_intrinsics_tassign.pto b/test/lit/vpto/fold_tile_buf_intrinsics_tassign.pto new file mode 100644 index 0000000000..6859c54080 --- /dev/null +++ b/test/lit/vpto/fold_tile_buf_intrinsics_tassign.pto @@ -0,0 +1,52 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Check the complete VPTO tile-op expansion path when a tile handle is rebound +// before reaching a template-generated tile_buf intrinsic. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --pto-level=level3 --emit-vpto %s -o - 2>&1 | FileCheck %s + +// CHECK-LABEL: func.func @tadd_tassign +// CHECK-NOT: pto.tadd ins +// CHECK: %[[REBOUND_ADDR:.*]] = arith.constant 4096 : i64 +// CHECK: pto.tassign %{{.*}}, %[[REBOUND_ADDR]] +// CHECK: pto.vecscope +// CHECK: pto.castptr %[[REBOUND_ADDR]] : i64 -> !pto.ptr +// CHECK: pto.vadd + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @tadd_tassign() { + %base_addr = arith.constant 0 : i64 + %rebound_addr = arith.constant 4096 : i64 + %a = pto.alloc_tile addr = %base_addr + : !pto.tile_buf + %b = pto.alloc_tile addr = %base_addr + : !pto.tile_buf + %dst = pto.alloc_tile addr = %base_addr + : !pto.tile_buf + %bound_a = pto.tassign %a, %rebound_addr + : !pto.tile_buf + -> !pto.tile_buf + + pto.tadd ins( + %bound_a, %b + : !pto.tile_buf, + !pto.tile_buf) + outs( + %dst + : !pto.tile_buf) + return + } +} diff --git a/test/lit/vpto/intra_block_sync_vpto_llvm.pto b/test/lit/vpto/intra_block_sync_vpto_llvm.pto index a2b7d001d3..d7f9ff9478 100644 --- a/test/lit/vpto/intra_block_sync_vpto_llvm.pto +++ b/test/lit/vpto/intra_block_sync_vpto_llvm.pto @@ -14,6 +14,8 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, 0 pto.sync.set , %event_id pto.sync.wait , %event_id + pto.sync.set , 16 + pto.sync.wait , %event_id return } } diff --git a/test/lit/vpto/pipe_descriptor_vpto_llvm.pto b/test/lit/vpto/pipe_descriptor_vpto_llvm.pto new file mode 100644 index 0000000000..97b26543ff --- /dev/null +++ b/test/lit/vpto/pipe_descriptor_vpto_llvm.pto @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ( mkdir -p %T && ptoas --pto-arch=a5 --pto-backend=vpto --pto-level=level3 %s -o %t --mlir-print-ir-after=convert-func-to-llvm 2>&1 || true ) | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @pipe_descriptor_lowering(%descriptor_addr: i64) + attributes {pto.kernel} { + %state = pto.declare_struct -> !pto.struct + %zero = arith.constant 0 : i32 + pto.struct_set %state[0], %zero : !pto.struct, i32 + pto.struct_set %state[1], %zero : !pto.struct, i32 + %counter = pto.struct_get %state[0] : !pto.struct -> i32 + + %entry = pto.declare_global -> !pto.tensor_view<16x16xi32> + %bound_entry = pto.tassign %entry, %descriptor_addr + : !pto.tensor_view<16x16xi32> -> !pto.tensor_view<16x16xi32> + %addr = pto.tensor_view_addr %bound_entry + : !pto.tensor_view<16x16xi32> -> !pto.ptr + %c0 = arith.constant 0 : index + pto.store_scalar %counter, %addr[%c0] : !pto.ptr, i32 + return + } +} + +// CHECK-LABEL: llvm.func @pipe_descriptor_lowering_mix_aiv +// CHECK: llvm.alloca {{.*}} x !llvm.struct<(i32, i32)> +// CHECK: llvm.getelementptr +// CHECK: llvm.store {{.*}} : i32, !llvm.ptr +// CHECK: llvm.load {{.*}} : !llvm.ptr -> i32 +// CHECK: llvm.alloca {{.*}} x !llvm.ptr<1> +// CHECK: llvm.inttoptr +// CHECK: llvm.store {{.*}} : !llvm.ptr<1>, !llvm.ptr +// CHECK: llvm.load {{.*}} : !llvm.ptr -> !llvm.ptr<1> +// CHECK-NOT: builtin.unrealized_conversion_cast diff --git a/test/lit/vpto/pipe_metadata_abi.pto b/test/lit/vpto/pipe_metadata_abi.pto new file mode 100644 index 0000000000..30959191df --- /dev/null +++ b/test/lit/vpto/pipe_metadata_abi.pto @@ -0,0 +1,58 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// The test-only daemon validates PipeSpec fields, ordered resources, the +// shared state operand, and the defaulted subblock argument before returning a +// candidate. It is deliberately outside ptodsl/tilelib so this remains a +// compiler ABI regression rather than a TileLib implementation test. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --pto-level=level3 --tile-lib-backend=ptodsl --ptodsl-pkg-path=%S/Inputs/pipe_metadata_mock --emit-pto-ir --mlir-print-ir-after=pto-insert-template-attributes %s -o /dev/null 2>&1 | FileCheck %s + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @pipe_metadata() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %slots = pto.declare_global -> !pto.tensor_view<16x16xf32> + %pipe = "pto.initialize_l2g2l_pipe"(%slots) <{ + dir_mask = 1 : i8, + slot_size = 1024 : i32, + slot_num = 8 : i32, + flag_base = 4 : i32, + operandSegmentSizes = array + }> : (!pto.tensor_view<16x16xf32>) -> !pto.pipe + %entry = pto.declare_global -> !pto.tensor_view<16x16xf32> + %entry_addr = arith.constant 0 : i64 + %bound_entry = pto.tassign %entry, %entry_addr + : !pto.tensor_view<16x16xf32> -> !pto.tensor_view<16x16xf32> + pto.talloc(%bound_entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pto.tpush(%bound_entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pto.tpush(%bound_entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 2} + pto.tpop(%bound_entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pto.tfree(%bound_entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + return + } + + func.func @pipe_metadata_explicit_subblock() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %local_addr = arith.constant 17 : i32 + %peer_local_addr = arith.constant 29 : i32 + %pipe = pto.initialize_l2l_pipe { + dir_mask = 3, + slot_size = 512, + slot_num = 4, + flag_base = 8, + nosplit = false + }(%local_addr : i32, %peer_local_addr : i32) -> !pto.pipe + %entry = pto.declare_tile -> !pto.tile_buf + %subblock = arith.constant 7 : i64 + pto.tpush(%entry, %pipe : !pto.tile_buf, !pto.pipe) {split = 1} aiv_subblockid(%subblock) + pto.tpop(%entry, %pipe : !pto.tile_buf, !pto.pipe) {split = 1} aiv_subblockid(%subblock) + return + } +} + +// CHECK-LABEL: func.func @pipe_metadata +// CHECK-COUNT-9: name = "mock_pipe_template" diff --git a/test/lit/vpto/pipe_state_materialization.pto b/test/lit/vpto/pipe_state_materialization.pto new file mode 100644 index 0000000000..047d38f7e7 --- /dev/null +++ b/test/lit/vpto/pipe_state_materialization.pto @@ -0,0 +1,91 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --tile-lib-backend=ptodsl --ptodsl-pkg-path=%S/Inputs/pipe_metadata_mock --emit-pto-ir --mlir-print-ir-after=pto-materialize-pipe-state %s -o /dev/null 2>&1 | FileCheck %s + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @state_with_producer() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %slots = pto.declare_global -> !pto.tensor_view<16x16xf32> + %pipe = "pto.initialize_l2g2l_pipe"(%slots) <{ + dir_mask = 1 : i8, + slot_size = 1024 : i32, + slot_num = 8 : i32, + flag_base = 4 : i32, + operandSegmentSizes = array + }> : (!pto.tensor_view<16x16xf32>) -> !pto.pipe + %entry = pto.declare_global -> !pto.tensor_view<16x16xf32> + pto.talloc(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pto.tpush(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pto.tpush(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 2} + pto.tpop(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pto.tfree(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + return + } + + func.func @state_without_producer() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %slots = pto.declare_global -> !pto.tensor_view<16x16xf32> + %pipe = "pto.initialize_l2g2l_pipe"(%slots) <{ + dir_mask = 1 : i8, + slot_size = 1024 : i32, + slot_num = 8 : i32, + flag_base = 4 : i32, + operandSegmentSizes = array + }> : (!pto.tensor_view<16x16xf32>) -> !pto.pipe + %entry = pto.declare_global -> !pto.tensor_view<16x16xf32> + pto.talloc(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 0} + pto.tpop(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 0} + pto.tfree(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 0} + return + } + + func.func @state_with_explicit_subblock() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %local_addr = arith.constant 17 : i32 + %peer_local_addr = arith.constant 29 : i32 + %pipe = pto.initialize_l2l_pipe { + dir_mask = 3, + slot_size = 512, + slot_num = 4, + flag_base = 8, + nosplit = false + }(%local_addr : i32, %peer_local_addr : i32) -> !pto.pipe + %entry = pto.declare_tile -> !pto.tile_buf + %subblock = arith.constant 7 : i64 + pto.tpush(%entry, %pipe : !pto.tile_buf, !pto.pipe) {split = 1} aiv_subblockid(%subblock) + pto.tpop(%entry, %pipe : !pto.tile_buf, !pto.pipe) {split = 1} aiv_subblockid(%subblock) + return + } +} + +// CHECK-LABEL: func.func @state_with_producer +// CHECK: %[[STATE:.*]] = pto.declare_struct -> !pto.struct +// CHECK: %[[ZERO:.*]] = arith.constant 0 : i32 +// CHECK: pto.struct_set %[[STATE]] [0], %[[ZERO]] : !pto.struct, i32 +// CHECK: pto.struct_set %[[STATE]] [1], %[[ZERO]] : !pto.struct, i32 +// CHECK: pto.talloc({{.*}}) {split = 1} pipe_state(%[[STATE]] : !pto.struct) +// CHECK: pto.tpush({{.*}}) {split = 1} pipe_state(%[[STATE]] : !pto.struct) +// CHECK: pto.tpush({{.*}}) {split = 2} pipe_state(%[[STATE]] : !pto.struct) +// CHECK: pto.tpop({{.*}}) {split = 1} pipe_state(%[[STATE]] : !pto.struct) +// CHECK: pto.tfree({{.*}}) {split = 1} pipe_state(%[[STATE]] : !pto.struct) +// CHECK: pto.tdrain({{.*}}, %[[STATE]] : !pto.pipe, !pto.struct) {split = 1} + +// CHECK-LABEL: func.func @state_without_producer +// CHECK: %[[NO_PUSH_STATE:.*]] = pto.declare_struct -> !pto.struct +// CHECK: pto.talloc({{.*}}) {split = 0} pipe_state(%[[NO_PUSH_STATE]] : !pto.struct) +// CHECK: pto.tpop({{.*}}) {split = 0} pipe_state(%[[NO_PUSH_STATE]] : !pto.struct) +// CHECK: pto.tfree({{.*}}) {split = 0} pipe_state(%[[NO_PUSH_STATE]] : !pto.struct) +// CHECK-NOT: pto.tdrain +// CHECK: return + +// CHECK-LABEL: func.func @state_with_explicit_subblock +// CHECK: %[[SUBBLOCK_STATE:.*]] = pto.declare_struct -> !pto.struct +// CHECK: pto.tpush({{.*}}) {split = 1} aiv_subblockid(%{{.*}}) pipe_state(%[[SUBBLOCK_STATE]] : !pto.struct) +// CHECK: pto.tpop({{.*}}) {split = 1} aiv_subblockid(%{{.*}}) pipe_state(%[[SUBBLOCK_STATE]] : !pto.struct) +// CHECK: pto.tdrain({{.*}}, %[[SUBBLOCK_STATE]] : !pto.pipe, !pto.struct) {split = 1} diff --git a/test/lit/vpto/pipe_tassign_tilelib_expand.pto b/test/lit/vpto/pipe_tassign_tilelib_expand.pto new file mode 100644 index 0000000000..77b31b4ad2 --- /dev/null +++ b/test/lit/vpto/pipe_tassign_tilelib_expand.pto @@ -0,0 +1,41 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// Run the complete VPTO TileLib path for a global pipe entry that is rebound by +// pto.tassign. The test-only daemon renders no-op helpers so it isolates the +// compiler's metadata, expansion, and helper-wiring contract from FIFO semantics. +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --pto-level=level3 --tile-lib-backend=ptodsl --ptodsl-pkg-path=%S/Inputs/pipe_metadata_mock --emit-vpto --mlir-print-ir-after=pto-expand-tile-op %s -o /dev/null 2>&1 | FileCheck %s + +// CHECK: IR Dump After ExpandTileOp +// CHECK-LABEL: func.func @pipe_tassign_tilelib_expand +// CHECK: %[[BOUND_ENTRY:.*]] = pto.tassign +// CHECK: call @__pto_tilelang_a5_talloc_pipe_entry{{.*}}(%[[BOUND_ENTRY]], +// CHECK: call @__pto_tilelang_a5_tpush_pipe_entry{{.*}}(%[[BOUND_ENTRY]], +// CHECK-NOT: pto.talloc +// CHECK-NOT: pto.tpush + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @pipe_tassign_tilelib_expand() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %slots = pto.declare_global -> !pto.tensor_view<16x16xf32> + %pipe = "pto.initialize_l2g2l_pipe"(%slots) <{ + dir_mask = 1 : i8, + slot_size = 1024 : i32, + slot_num = 8 : i32, + flag_base = 4 : i32, + operandSegmentSizes = array + }> : (!pto.tensor_view<16x16xf32>) -> !pto.pipe + %entry = pto.declare_global -> !pto.tensor_view<16x16xf32> + %entry_addr = arith.constant 0 : i64 + %bound_entry = pto.tassign %entry, %entry_addr + : !pto.tensor_view<16x16xf32> -> !pto.tensor_view<16x16xf32> + pto.talloc(%bound_entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + pto.tpush(%bound_entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 1} + return + } +} diff --git a/test/lit/vpto/pipe_tilelib_feature_gate.pto b/test/lit/vpto/pipe_tilelib_feature_gate.pto new file mode 100644 index 0000000000..8c06eb1014 --- /dev/null +++ b/test/lit/vpto/pipe_tilelib_feature_gate.pto @@ -0,0 +1,35 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// A5 EmitC retains the legacy pipe path: PipeState materialization and its +// terminal tdrain are only inserted for the A5 VPTO backend. +// RUN: ptoas --pto-arch=a5 --pto-backend=emitc %s 2>&1 | FileCheck %s --check-prefix=EMITC + +module { + func.func @a5_emitc_pipe_unchanged() + attributes {pto.kernel_kind = #pto.kernel_kind} { + %slots = pto.declare_global -> !pto.tensor_view<16x16xf32> + %pipe = "pto.initialize_l2g2l_pipe"(%slots) <{ + dir_mask = 1 : i8, + slot_size = 1024 : i32, + slot_num = 8 : i32, + flag_base = 0 : i32, + operandSegmentSizes = array + }> : (!pto.tensor_view<16x16xf32>) -> !pto.pipe + %entry = pto.declare_global -> !pto.tensor_view<16x16xf32> + pto.talloc(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 0} + pto.tpush(%entry, %pipe : !pto.tensor_view<16x16xf32>, !pto.pipe) {split = 0} + return + } +} + +// EMITC-LABEL: AICORE void a5_emitc_pipe_unchanged( +// EMITC: TPipe<0, Direction::DIR_C2V_GM, 1024, 8, 8, true> +// EMITC: TALLOC< +// EMITC: TPUSH< +// EMITC-NOT: TDrain diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index b650190a8e..daf436083a 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -744,6 +744,20 @@ static bool hasUnexpandedTileOps(ModuleOp module) { return found; } +static bool hasPipeTransactions(ModuleOp module) { + bool found = false; + module.walk([&](Operation *op) { + if (found) + return; + if (isa(op)) + found = true; + }); + return found; +} + using FunctionBlockArgHintMap = llvm::StringMap, 4>>; @@ -1118,6 +1132,9 @@ struct SerialFrontendPipeLoweringPass MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID( SerialFrontendPipeLoweringPass) + explicit SerialFrontendPipeLoweringPass(bool lowerNestedPipeFunctions) + : lowerNestedPipeFunctions(lowerNestedPipeFunctions) {} + void getDependentDialects(DialectRegistry ®istry) const override { registry.insert(); } @@ -1132,7 +1149,52 @@ struct SerialFrontendPipeLoweringPass // adaptor allows one function to be verified while another function is // still mutating its pipe ops. Keep these two small passes serial so every // verifier observes either the complete frontend or complete lowered form. - for (func::FuncOp funcOp : getOperation().getOps()) { + SmallVector functions; + if (lowerNestedPipeFunctions) { + getOperation().walk( + [&](func::FuncOp funcOp) { functions.push_back(funcOp); }); + } else { + for (func::FuncOp funcOp : getOperation().getOps()) + functions.push_back(funcOp); + } + for (func::FuncOp funcOp : functions) { + if (failed(runPipeline(functionPM, funcOp))) { + signalPassFailure(); + return; + } + } + } + +private: + bool lowerNestedPipeFunctions; +}; +} // namespace + +static std::unique_ptr +createSerialFrontendPipeLoweringPass(bool lowerNestedPipeFunctions) { + return std::make_unique( + lowerNestedPipeFunctions); +} + +namespace { +struct RecursivePipeStateMaterializationPass + : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID( + RecursivePipeStateMaterializationPass) + + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + } + + void runOnOperation() override { + OpPassManager functionPM(func::FuncOp::getOperationName()); + functionPM.addPass(pto::createPTOMaterializePipeStatePass()); + + SmallVector functions; + getOperation().walk( + [&](func::FuncOp funcOp) { functions.push_back(funcOp); }); + for (func::FuncOp funcOp : functions) { if (failed(runPipeline(functionPM, funcOp))) { signalPassFailure(); return; @@ -1142,8 +1204,8 @@ struct SerialFrontendPipeLoweringPass }; } // namespace -static std::unique_ptr createSerialFrontendPipeLoweringPass() { - return std::make_unique(); +static std::unique_ptr createRecursivePipeStateMaterializationPass() { + return std::make_unique(); } static void collectNonEntryBlocksInSourceOrder( @@ -2929,7 +2991,6 @@ int mlir::pto::compilePTOASModule( llvm::errs() << "Error: --enable-bufid_sync requires --pto-arch=a5.\n"; return 1; } - module->getOperation()->setAttr("pto.target_arch", mlir::StringAttr::get(module->getContext(), arch)); @@ -3076,7 +3137,12 @@ int mlir::pto::compilePTOASModule( } } - const bool hasTileOpsToExpand = hasUnexpandedTileOps(*module); + const bool enableA5VPTOPipeExpansion = + arch == "a5" && effectiveBackend == PTOBackend::VPTO && + hasPipeTransactions(*module); + + const bool hasTileOpsToExpand = + enableA5VPTOPipeExpansion || hasUnexpandedTileOps(*module); std::optional expandOptions; if (effectiveBackend == PTOBackend::VPTO && hasTileOpsToExpand && tileLibBackend == TileLibBackend::PTODSL) @@ -3109,7 +3175,8 @@ int mlir::pto::compilePTOASModule( // lifted to make it unconditional for all backends. if (effectiveBackend == PTOBackend::VPTO) pm.addNestedPass(pto::createPTOCanonicalizeIRPass()); - pm.addPass(createSerialFrontendPipeLoweringPass()); + pm.addPass( + createSerialFrontendPipeLoweringPass(enableA5VPTOPipeExpansion)); //pm.addNestedPass(pto::createPTOVerifyTFreePass()); pm.addPass(pto::createPTOInferValidatePipeInitPass()); pm.addNestedPass(pto::createLoweringSyncToPipePass()); @@ -3131,6 +3198,7 @@ int mlir::pto::compilePTOASModule( expandOptions->tileLibBackend == "ptodsl") { auto insertOptions = buildInsertTemplateAttributesOptions(*expandOptions); + insertOptions.skipPipeOps = enableA5VPTOPipeExpansion; pm.addPass( pto::createInsertTemplateAttributesPass(insertOptions)); } @@ -3170,6 +3238,19 @@ int mlir::pto::compilePTOASModule( pm.addPass(pto::createPlanMemoryPass(planMemoryOption)); } pm.addPass(pto::createPTOResolveReservedBuffersPass()); + // PipeState is needed only by pipe metadata and expansion. Materialize it + // after the shared layout/fusion/memory-planning prefix so those passes keep + // consuming the pre-existing pipe IR, but before pipe-only metadata needs + // the state and resolved pipe resources. + if (enableA5VPTOPipeExpansion) + pm.addPass(createRecursivePipeStateMaterializationPass()); + if (enableA5VPTOPipeExpansion && expandOptions && + expandOptions->tileLibBackend == "ptodsl") { + auto pipeInsertOptions = + buildInsertTemplateAttributesOptions(*expandOptions); + pipeInsertOptions.pipeOnly = true; + pm.addPass(pto::createInsertTemplateAttributesPass(pipeInsertOptions)); + } pm.addNestedPass(pto::createPTORemoveIdentityTMovPass()); // Conditionally add one automatic synchronization mode. Barrier-all is a