diff --git a/docs/designs/mix-kernel-mix-backend-compile-flow.md b/docs/designs/mix-kernel-mix-backend-compile-flow.md
index 27c0d5c9b0..36b0671b69 100644
--- a/docs/designs/mix-kernel-mix-backend-compile-flow.md
+++ b/docs/designs/mix-kernel-mix-backend-compile-flow.md
@@ -133,7 +133,7 @@ For example:
```python
@pto.jit(target="a5", entry=False, backend="vpto", mode="explicit")
def scale_row(base_gm: pto.ptr(pto.f32, "gm"), row: pto.i32):
- with pto.simd():
+ with pto.tileop():
...
@pto.jit(target="a5", backend="emitc")
@@ -168,10 +168,9 @@ kernel. The Vector/Cube execution ownership is a PTOAS responsibility:
the VPTO backend.
This keeps the PTODSL programming model independent of the physical sectioning
-rules. PTODSL can still expose helper abstractions such as `@pto.simd`,
-`@pto.cube`, `with pto.simd():`, and `with pto.cube():`, but the design does
-not require users or the frontend to manually partition every operation into a
-final section.
+rules. PTODSL exposes logical helper abstractions such as `@pto.tileop` and
+`with pto.tileop():`, but the design does not require users or the frontend to
+manually partition every operation into a final section.
### PTODSL IR Codegen Shape
@@ -233,11 +232,10 @@ Python-only structure. This lowering records PTODSL helper structure and call
boundaries; it does not make PTODSL responsible for the final Vector/Cube
section partition.
-For `@pto.simd` / `@pto.cube` and inline `with pto.simd():` / `with pto.cube():`
-scopes, PTODSL:
+For `@pto.tileop` and inline `with pto.tileop():` scopes, PTODSL:
- outlines the subkernel body into a helper `func.func` when needed
-- marks the helper with `pto.ptodsl.subkernel_helper`
+- marks the helper with canonical `pto.tileop.helper`
- emits a helper call from the caller body
This is the PTODSL-side expression of a logical mixed kernel: the entry or
@@ -526,4 +524,5 @@ Use this order when debugging mixed compilation:
| `pto.aicore` | `func.func` | Legacy entry marker accepted for compatibility. |
| `pto.internal.non_entry` | `func.func` | Frontend/helper metadata; not used for current entry inference. |
| `pto.ptodsl.logical_name` | `func.func` | Source-level logical name used when assembling wrappers and peer references. |
-| `pto.ptodsl.subkernel_helper` | `func.func` | Frontend helper classification: `simd`, `cube`, or `simt`. |
+| `pto.tileop.helper` | `func.func` | Canonical tileop-style helper marker emitted for `@pto.tileop` helpers. |
+| `pto.ptodsl.subkernel_helper` | `func.func` | Legacy helper role marker retained for compatibility with older/manual IR. |
diff --git a/docs/designs/ptodsl-ast-preprocess-control-flow.md b/docs/designs/ptodsl-ast-preprocess-control-flow.md
index c67a1fd4df..5c555d39dd 100644
--- a/docs/designs/ptodsl-ast-preprocess-control-flow.md
+++ b/docs/designs/ptodsl-ast-preprocess-control-flow.md
@@ -19,8 +19,8 @@ static specialization readable.
## Goals
- Make native Python control-flow syntax usable by default for runtime control
- flow in `@pto.jit(...)` kernels and named `@pto.cube` / `@pto.simd` /
- `@pto.simt` sub-kernels.
+ flow in `@pto.jit(...)` kernels and named `@pto.tileop` / `@pto.simt`
+ sub-kernels.
- Use `ast_rewrite` as the public name for the source rewrite feature.
- Rewrite legal Python `if` / `for range(...)` into existing PTODSL
control-flow surfaces.
diff --git a/docs/designs/ptodsl-redesign-of-simd-simt-cube-subkernel.md b/docs/designs/ptodsl-redesign-of-simd-simt-cube-subkernel.md
new file mode 100644
index 0000000000..8eb94dd04c
--- /dev/null
+++ b/docs/designs/ptodsl-redesign-of-simd-simt-cube-subkernel.md
@@ -0,0 +1,346 @@
+# PTODSL TileOp 多 Phase 设计
+
+## 1. 背景和目标
+
+PTODSL 需要一种可复用的 tile 级 helper:调用者把已经规划好的片上
+`Tile` 传入,helper 用 MTE、SIMD、Cube 和必要时的 SIMT 微指令完成一段
+计算。一个真实的 helper 很少只属于一条 pipe。例如,一个 block 可能依次执行:
+
+```text
+MTE load -> SIMD normalize -> MTE move -> Cube matmul -> SIMD post-process
+```
+
+旧的 `simd` / `cube` subkernel 模型把一次 helper call 概括为单一角色。这样无法
+表达 helper 内不同阶段分别读取、写入了哪个边界 Tile,也无法正确处理多段
+Vector/Cube 计算、MTE 与计算域的归属,以及直接写在 helper 中的 SIMT 微指令。
+
+本设计以 `@pto.tileop` 作为唯一的 tile 级 helper surface。它的边界简单且稳定:
+只传递 `Tile` 和 PTO scalar。helper body 则可以包含多个不同 pipe 的微指令阶段。
+
+本文描述的是目标设计。当前实现仍保留 TensorView ABI、单一 `primary_domain` 和
+单个主计算 section 等过渡性行为;第 10 节列出差异和迁移顺序。
+
+## 2. 基础概念
+
+### 2.1 Kernel、helper、Tile 与 scalar
+
+- **kernel** 是 `@pto.jit` 定义的可执行入口,负责 tile 分配、整体调度和调用
+ helper。
+- **tileop helper** 是由 `@pto.tileop` 定义、可被 kernel 多次调用的命名函数。它
+ 不单独规划 Tile 生命周期,也不是设备执行入口。
+- **Tile** 是调用者已经拥有并交给 helper 使用的片上数据缓冲区。它是 tileop
+ 唯一的复合数据边界类型。
+- **PTO scalar** 是 `i32`、`f32` 等设备标量,可作为大小、步长、标志或计算参数。
+ scalar 不代表一块可由同步系统追踪的 Tile 数据。
+
+`@pto.simt` 保留为独立的 launched-SIMT public surface,可使用既有的 ptr ABI 和
+launch 语法。它不成为 tileop 的 ABI,也不能从 tileop 中被调用。本设计额外允许
+用户在 tileop body 中直接写 **SIMT 微指令**;这两件事不同。
+
+| surface | public ABI | 使用场景 |
+|---|---|---|
+| `@pto.tileop` | `Tile`、PTO scalar | 复用多 pipe 的 tile 级计算;body 可直接写 SIMT 微指令 |
+| `@pto.simt` | `ptr`、PTO scalar | 保留的 launched-SIMT helper;需要将 ptr 作为函数边界时使用 |
+
+因此,“tileop 允许 SIMT”只表示其 body 可以包含和自动 outline SIMT 微指令,
+不表示 `@pto.tileop` 接受 ptr,也不替换 `@pto.simt`。
+
+### 2.2 Pipe、phase 与 section
+
+- **pipe** 是一类执行资源,例如 MTE、Vector、Cube、SIMT 或 Scalar。
+- **phase** 是 helper 内一段有确定执行 pipe、边界 Tile 读写和顺序关系的操作。
+ 一个 helper 可以有任意多个 phase,同一 pipe 也可以多次出现。
+- **section** 是 PTOAS 为 VPTO 分模块和 lowering 物化的 IR 容器:
+ `pto.section.vector` 或 `pto.section.cube`。它不是用户选择的 helper role。
+
+phase 描述“这段操作做了什么以及依赖什么”;section 描述“这段 Vector/Cube 操作
+在后续 VPTO 流程中属于哪个计算模块”。MTE、SIMT、同步和控制流可构成 phase,
+但不自动变成 Vector/Cube section。
+
+## 3. `@pto.tileop` 的公共契约
+
+### 3.1 ABI
+
+`@pto.tileop` 的位置参数只允许:
+
+| 类型 | 用途 |
+|---|---|
+| `pto.Tile` | 输入、输出或 scratch Tile |
+| PTO scalar | 标量参数和标量结果 |
+
+不允许 `TensorView`、`PartitionTensorView`、`ptr`、`memref`、host tensor、
+`TensorSpec`、vreg、mask 或 pipe handle 跨 helper 边界。ptr 继续只属于独立
+`@pto.simt` 的 public ABI;这项 `@pto.simt` 契约不在本次 redesign 的收紧范围内。
+
+Tile 输出通过可写的 Tile 参数表达,不通过 function result 返回。function result
+仅允许 PTO scalar。这样 caller 始终拥有 Tile 的分配、别名和生命周期,PTOAS 也能
+将 helper 的读写明确映射到 call 的 Tile operands。
+
+```python
+# 设计示意:所有复合数据都通过 Tile 边界传递。
+@pto.tileop
+def fused_block(src: pto.Tile, weight: pto.Tile, dst: pto.Tile,
+ scratch: pto.Tile, rows: pto.i32, cols: pto.i32):
+ # body 见下一节。
+ ...
+
+@pto.jit
+def kernel(src, weight, dst):
+ # kernel 负责创建和规划 src_tile、weight_tile、dst_tile、scratch_tile。
+ fused_block(src_tile, weight_tile, dst_tile, scratch_tile, rows, cols)
+```
+
+### 3.2 Body 允许和禁止的内容
+
+tileop body 允许下列内容:
+
+- MTE、SIMD、Cube 和 SIMT 微指令;
+- 标量运算、地址派生、cast、`arith` 和 `scf` 等结构性操作;
+- 显式 pipe 同步和显式 SIMT 线程同步。
+
+tileop body 不允许下列内容:
+
+- 高层 TileOps,例如由自动调度或 Tile API 负责语义的 load/store/compute 操作;
+- `alloc_tile`、`reserve_buffer`、`TAlloc` 等需要为 callee 单独规划生命周期的
+ Tile 分配;scratch 和输出 Tile 必须由 caller 传入;
+- 对另一个 tileop helper 或 `@pto.simt` helper 的调用;
+- 将内部 vreg、mask、pointer 或 pipe handle 作为参数或结果泄漏到 helper 外。
+
+这里的限制不禁止在 body 内从 Tile 派生地址,也不禁止 SIMT 微指令使用这些内部
+地址。限制的是 public ABI:派生 pointer 的作用域不能超出当前 helper 或其自动
+生成的内部 SIMT entry。
+
+## 4. 用户可见的多 Phase 模型
+
+用户不为 helper 指定 `vector`、`cube` 或 `simt` role。用户按算法需要写微指令和
+显式同步;PTOAS 从 body 推导阶段。下面是概念流程,不表示可直接编译的 Python
+代码:
+
+```text
+Tile src, weight, dst, scratch
+ |
+ +-- MTE phase: src -> scratch
+ +-- SIMD phase: normalize(scratch)
+ +-- MTE phase: scratch -> Cube input buffer
+ +-- Cube phase: matmul(weight, scratch)
+ +-- SIMD phase: post-process -> dst
+ +-- SIMT phase: per-element fixup(dst)
+```
+
+同一 helper 可以有多个 Vector section、多个 Cube section,以及二者交错的 phase。
+例如 `Vector -> MTE -> Vector` 必须产生两个 Vector 计算 span,而不是为了保持
+“单一 primary span”把中间 MTE 包进 section。控制流也不改变这一原则:PTOAS 递归
+分析 `scf.for` 和 `scf.if` 的 region,在具体 region 内物化计算 span;不得因为顶层
+只有一个 `scf.for` 就跳过或包裹整个 loop。
+
+## 5. PTODSL tracing 与 PTOAS phase graph
+
+### 5.1 PTODSL tracing 的职责
+
+PTODSL tracing 是执行 Python helper 定义并记录 PTO IR 的过程。新设计中,它只做
+两件与 tileop 相关的事:
+
+1. 验证 Tile/Scalar public ABI,并记录带 `pto.tileop.helper` marker 的命名
+ `func.func`;
+2. 原样保留 body 中的微指令、结构化控制流和用户写出的同步。
+
+tracing 不预先选择 Vector 或 Cube,不写 `primary_domain`,不把 body 包进
+`pto.section.vector/cube`,也不预先填写 phase 或 operand effect 属性。inline
+`with pto.tileop():` 与命名 `@pto.tileop` 使用同一套 body 规则;inline 路径不因
+位置不同而硬编码一个 section kind。
+
+```text
+PTODSL tracing
+ -> func.func {pto.tileop.helper},Tile/Scalar ABI,原始微指令 body
+ -> PTOAS 推导 phase graph 并验证
+```
+
+### 5.2 Phase graph
+
+PTOAS 的 phase 推导 pass 扫描 helper body,读每个微指令自身声明的 `getPipe()`,
+并生成有序 phase graph。每个 phase 至少记录:
+
+- 执行 pipe;
+- 在源顺序和结构化控制流中的位置;
+- 读取和写入的 Tile 参数编号;
+- 与前后 phase 的数据依赖和控制依赖;
+- 用户显式同步形成的不可删除顺序边;
+- 对 MTE phase 而言,最终归属的 Vector 或 Cube 物理模块。
+
+读写摘要必须追溯到 helper 参数。例如 `tile[row, col:]`、地址计算、cast、
+`memref.subview`、`scf` 的 iter_arg/result 等会产生派生 SSA value;真正被微指令
+读写的可能是派生值而不是函数参数。PTOAS 沿这些透明派生关系追溯,最后记录
+“第几个 Tile 参数被读或写”。这样摘要既不把内部 pointer 误当作 ABI,也不会漏掉
+tile slice 的边界 effect。
+
+目标 IR 属性形态如下,字段名仅说明语义:
+
+```text
+func.func @fused_block(%src, %weight, %dst, %scratch, %rows, %cols)
+ {pto.tileop.helper,
+ pto.tileop.phases = [
+ {pipe = MTE2, tile_uses = [0], tile_defs = [3]},
+ {pipe = V, tile_uses = [3], tile_defs = [3]},
+ {pipe = MTE1, tile_uses = [3], tile_defs = [3], owner = CUBE},
+ {pipe = CUBE, tile_uses = [1, 3], tile_defs = [2]},
+ {pipe = V, tile_uses = [2], tile_defs = [2]},
+ {pipe = SIMT, tile_uses = [2], tile_defs = [2]}
+ ],
+ pto.tileop.operand_effects = [read, read, readwrite, readwrite, read, read]} {
+ ...
+}
+```
+
+`operand_effects` 是所有 phase 的 Tile effect 的并集,用于快速检查和 callsite
+建模;真正的顺序和逐阶段同步依据 `phases`。不再存在 `pto.tileop.primary_domain`。
+scalar 可以出现在参数列表和 phase 内,但不作为 Tile memory hazard 节点。
+
+## 6. 自动同步和用户显式同步
+
+PTOAS 负责普通数据 hazard:某个 phase 写入 Tile 后,后续不同 pipe 的 phase 或
+另一次 helper call 读取/写入同一 Tile 时,`InsertSync` 根据 phase graph 插入所需的
+pipe 同步。它既分析 helper body 内的常规生产者/消费者关系,也在 callsite 将摘要中
+的 Tile 参数编号映射回本次 `func.call` 的实际 Tile operand。
+
+```text
+helper A 的 phase 2 写 %scratch
+ +-----------------------+
+ v
+helper B 的 phase 0 读 %scratch
+
+InsertSync 根据写/读的 pipe 和顺序插入必要同步。
+```
+
+以下同步不能靠普通 Tile 读写关系可靠推导,必须由用户显式表达:
+
+- 算法规定的阶段边界、流水重叠和事件编号;
+- 对同一 Tile 的刻意并发访问或别名协议;
+- SIMT work-item 之间的 `syncthreads`、thread fence 等线程级同步。
+
+显式 pipe barrier、flag wait/set 和 SIMT 线程同步会成为 phase graph 的固定边。
+PTOAS 不删除它们;自动同步只补尚未覆盖的普通数据 hazard,且不能再插入等价的
+重复同步。
+
+## 7. TileOp 内的 SIMT 微指令
+
+tileop 允许直接包含 SIMT 微指令 phase,但这不等于允许调用 `@pto.simt` helper。
+PTOAS 会将可 outline 的连续 SIMT 微指令 span 变成内部实现细节:
+
+```text
+tileop 原始 body
+ -> pto.store_vfsimt_info(dim_z, dim_y, dim_x)
+ -> SIMT 微指令 span
+
+PTOAS SIMT outline
+ -> 内部 func.func {pto.simt_entry},capture 仅为 Tile/Scalar
+ -> pto.simt_launch(dim_x, dim_y, dim_z, captures...)
+```
+
+每个 SIMT span 前必须显式出现并支配该 span 的
+`pto.store_vfsimt_info(dim_z, dim_y, dim_x)`。该配置只消费给紧随其后的一个
+SIMT span;缺失、被多个 span 竞争、维度不合法或无法确定配置关联时均报错。显式
+维度避免 PTOAS 从 Tile shape 或循环结构猜测线程布局。
+
+outline 后的 `pto.simt_entry` 只捕获 Tile 和 scalar。SIMT body 内需要的 pointer
+从捕获 Tile 派生,不能被提升为 tileop 的函数参数。用户写出的 `syncthreads` 和
+其他 SIMT 同步原样保留在 entry 中。若一个 SIMT region 无法在保持结构化控制流和
+capture 规则的前提下 outline,PTOAS 必须诊断,而不能退化为普通 Vector 或忽略它。
+
+## 8. Section 物化、MTE 归属和 VPTO split
+
+### 8.1 Section materialization
+
+在同步摘要已经可用后,PTOAS 将每个最大 SIMD phase range 物化为
+`pto.section.vector`,将每个最大 Cube phase range 物化为 `pto.section.cube`。
+同一函数可以生成多个同类或异类 section。MTE、Scalar、SIMT 和同步保持在 section
+外,以便它们保留自己的 pipe 语义。
+
+```text
+MTE -> [section.vector] -> MTE -> [section.cube] -> [section.vector] -> SIMT
+```
+
+materialize pass 必须在嵌套 region 中处理 span,不能把混有 MTE、sync 或另一个
+计算域的整个 `scf.for` / `scf.if` 容器套成一个 Vector 或 Cube section。inline 后
+若出现同域 section 嵌套,由 section normalization 展开;跨域嵌套或无法确定语义的
+布局必须拒绝。
+
+### 8.2 MTE 归属
+
+MTE 本身不等于 Vector 或 Cube。PTOAS 根据 Tile 数据流把 MTE phase 归属到其唯一
+服务的 Vector 或 Cube 计算模块:例如 MTE 将数据搬到只被 Cube phase 消费的 Tile,
+该 MTE 归 Cube;计算结果经 MTE 写回且唯一来自 Vector phase,则归 Vector。
+
+一个 MTE phase 若同时服务两个计算域、数据流不足以确定归属,或其 Tile alias 使
+归属不唯一,必须报错。用户需要将数据移动拆开、使用独立 Tile,或写出能够消除
+歧义的算法结构;PTOAS 不能把同一条 MTE 指令复制到两个模块。
+
+### 8.3 `VPTOSplitCVModule` 与 `kernel_kind`
+
+`VPTOSplitCVModule` 以 section 和 MTE 归属为输入生成 Vector/Cube 模块。它必须:
+
+- 支持一个函数中的多个 Vector 和 Cube section;
+- 在 Vector clone 中保留全部 Vector section、归属 Vector 的 MTE,以及生成的
+ SIMT entry/launch;
+- 在 Cube clone 中保留全部 Cube section 和归属 Cube 的 MTE;
+- 只在 split clone 中删除另一域的内容,不能在用户显式要求单一 `kernel_kind`
+ 的模块中静默删除相反域的 phase。
+
+`pto.kernel_kind` 是 PTOAS 在模块分拆和 VPTO lowering 时使用的模块语义标记。
+它不是 `@pto.tileop` 的源级参数,也不能代替 phase 分析。显式单 kind 模块若仍含
+不属于该 kind 的 phase,应当诊断;否则错误会被隐藏到更晚的 lowering。
+
+## 9. 验证规则和预期诊断
+
+PTOAS 的 tileop contract verifier 至少应拒绝以下情况:
+
+- 非 Tile/Scalar ABI,或非 scalar function result;
+- 高层 TileOps、callee-local Tile allocation、tileop/helper 调用;
+- SIMT span 前没有唯一且有效的 `store_vfsimt_info`;
+- SIMT outline 需要捕获 ptr、vreg、mask 或 pipe handle;
+- 无法追溯到 Tile 参数的边界 memory effect;
+- 无法唯一归属到 Vector/Cube 的 MTE phase;
+- 无法物化的跨域/混合 section 结构;
+- 显式单 kind 模块中存在相反计算域的 phase。
+
+诊断应指出 helper、phase 和相关 Tile 参数,而不是只在最终 VPTO 或 LLVM lowering
+阶段报告“非法 section”。
+
+## 10. 当前实现差异和迁移计划
+
+| 项目 | 当前实现 | 目标设计 |
+|---|---|---|
+| tileop ABI | Tile、TensorView、PartitionTensorView、scalar | 仅 Tile、scalar |
+| 计算域 | 单一 `primary_domain` | 多 pipe、多 phase、无 primary domain |
+| SIMD/Cube | 只允许一个主计算域和主 span | 可有多个 Vector/Cube section 并交错 |
+| SIMT | tileop body 禁止 SIMT-only op | 允许直接写 SIMT 微指令并自动 outline |
+| 摘要 | phase + effect,服务单主域模型 | phase graph + Tile effect + MTE owner |
+| section | 只物化单个主计算 span | 物化所有最大 SIMD/Cube span |
+| split | 依赖单域化 section 形态 | 处理多 section、MTE owner 和 SIMT entry |
+| 同步 | 已开始按 phase 建 callsite 节点 | 覆盖多 phase body、跨 call Tile hazard 与显式边 |
+
+独立 `@pto.simt` 的 ptr ABI、launch syntax 和使用场景保持不变;本 redesign 仅修改
+`@pto.tileop` 的 ABI 及其 body 中的 SIMT 微指令处理方式。
+
+迁移按以下顺序进行:
+
+1. 收紧 PTODSL `@pto.tileop` ABI 到 Tile/Scalar,并补齐前端和后端负例;
+2. 让 tracing 只产生 raw tileop body,不写 `primary_domain` 或预套 section;
+3. 将 phase summary 扩展为无主域的多 phase graph,并完成 Tile 派生值追溯;
+4. 让 `InsertSync` 消费该 graph,先保证普通 Tile data hazard 和显式同步共存;
+5. 实现 tileop 内 SIMT span 的 launch 配置验证与 outline;
+6. 改造 section materialization、MTE ownership 和 `VPTOSplitCVModule`;
+7. 完成 VPTO/LLVM lowering 回归后,删除旧的单主域 tileop 路径和兼容属性。
+
+## 11. 回归测试
+
+目标实现至少需要以下覆盖:
+
+- Tile/Scalar ABI 正例,以及 TensorView、ptr、memref、vreg、mask 和非 scalar
+ result 的负例;
+- MTE+SIMD、MTE+Cube+SIMD、重复同域 section、控制流内 section;
+- Tile slice、address cast 和 `scf` iter_arg 派生值的 effect 追溯;
+- helper 内和跨 helper call 的自动 data-hazard sync,用户显式 barrier 保留且不重复;
+- tileop 内 SIMT 微指令、合法 launch 配置、线程同步、缺失配置和非法 capture;
+- MTE owner 推导成功、歧义 owner 诊断;
+- 多 section Vector/Cube split、SIMT 位于 Vector 侧、显式单 kind 冲突诊断;
+- 最终 VPTO 和 LLVM lowering 的端到端编译回归。
diff --git a/docs/designs/ptodsl-tileop-redesign-explainer.md b/docs/designs/ptodsl-tileop-redesign-explainer.md
new file mode 100644
index 0000000000..1da4bc2b02
--- /dev/null
+++ b/docs/designs/ptodsl-tileop-redesign-explainer.md
@@ -0,0 +1,954 @@
+# PTODSL TileOp Subkernel Redesign Explainer
+
+本文解释本次 PR 为什么要重做 tile-level
+subkernel 的 public surface,以及 PTODSL tracing、PTO IR、PTOAS 如何配合
+完成这次修改。
+
+更精确的最终 contract 见
+`docs/designs/ptodsl-redesign-of-simd-simt-cube-subkernel.md`。本文侧重背景铺垫
+和方案理解。
+
+## 1. 背景:PTODSL 在 PTOAS 中的位置
+
+PTOAS 的底层输入是 PTO IR。PTO IR 描述数据搬运、tile 计算、同步、内存规划
+和代码生成所需的信息。PTODSL 是 PTOAS 的 Python 前端:用户写 Python 函数,
+PTODSL tracing 把函数体记录成 PTO IR,然后交给 PTOAS 继续做 layout、memory、
+sync、EmitC 或 VPTO lowering。
+
+一个典型流程是:
+
+```text
+Python @pto.jit function
+ -> PTODSL tracing
+ -> PTO IR module
+ -> PTOAS
+ -> EmitC / VPTO output
+ -> runtime or simulator validation
+```
+
+## 2. 问题定义:旧 API 把 helper 名字和计算风格绑在一起
+
+issue 直接推动的是 PTODSL subkernel public surface 的收敛。修改前,PTODSL
+同时暴露 `@pto.simd`、`@pto.cube`、inline `with pto.simd():`、
+inline `with pto.cube():` 和 `@pto.simt`。
+
+这里的问题不是“名字太多”本身,而是 `simd` / `cube` 把两件不该耦合的事情绑
+在了一起:
+
+- 用户入口(public surface):用户要定义的是“一个 tile-level helper”,它的
+ 参数和返回值边界是什么、能不能含 load/compute/store 多个阶段。
+- 主计算风格:helper body 里真正的主计算是 Vector 还是 Cube。
+
+这两件事本应分开。一个 tile-level helper 的自然边界是 `Tile` / `TensorView`
+/ `PartitionTensorView` / PTO scalar;它的 body 可以是 vector-style,也可以是
+cube-style,还可能同时包含 MTE load/store、Scalar 控制和 sync(softmax 就是
+典型例子:MTE load → Vector compute → MTE store)。
+
+旧 API 要求用户先用 `@pto.simd` / `@pto.cube` 选主计算风格,于是 API 名字、
+参数契约、inline 路径和 PTOAS 后续分析都跟着 `simd` / `cube` 分叉成两套。但
+“是不是 vector 主计算”本该由 PTOAS 看 body 自己推导出来,而不是让用户提前
+用 decorator 名字声明。
+
+因此本次设计只保留两个用户入口:
+
+```text
+tile-level custom helper -> @pto.tileop
+launched SIMT helper -> @pto.simt
+```
+
+`@pto.simt` 不合并进 `@pto.tileop`,因为 SIMT helper 面向 launched SIMT 代码,
+允许 `pto.ptr(...)` 这类 raw pointer 边界;tileop helper 则是 tile/view/scalar
+层面的 helper。旧 `simd/cube` surface 不再作为兼容别名继续使用,而是诊断并提示
+迁移到 `@pto.tileop`。
+
+## 3. 一个 tile-level helper 要表达什么
+
+subkernel 是从主 kernel 里抽出来的可复用 helper。它通常不是独立 launched
+kernel,而是一个会被主 kernel 调用、之后再由 PTOAS inline 或拆分处理的函数。
+
+以 row-wise softmax 风格的 helper 为例,用户想表达的是:
+
+```python
+@pto.tileop
+def row_softmax(src: pto.TensorView,
+ dst: pto.TensorView,
+ scratch: pto.Tile,
+ rows: pto.i32,
+ cols: pto.i32):
+ pto.tile.load(src, scratch)
+ # vector compute on scratch
+ pto.tile.store(scratch, dst)
+
+@pto.jit
+def kernel(src, dst):
+ row_softmax(src, dst, scratch, rows, cols)
+```
+
+这个 helper 的 public ABI 是 `src/dst/scratch/rows/cols` 这些跨调用边界可见的
+值;helper body 里则有实际执行过程:
+
+```text
+MTE load -> Vector compute -> MTE store
+```
+
+
+## 4. 方案总览
+
+第 2 节已经说明用户入口的变化。这里从相同层级比较修改前后的职责分工,只关注
+tile-level helper;`@pto.simt` 在修改前后都保留独立入口,不在图中重复展示。
+
+修改前,用户必须先选择 `simd` 或 `cube`。PTODSL tracing 再把这个选择记录成
+分类标记,并按该分类生成对应的 section 形态:
+
+```mermaid
+flowchart LR
+ User["用户定义
tile-level helper"]
+ Surface["选择 @pto.simd
或 @pto.cube"]
+ Trace["PTODSL tracing
记录 simd/cube 分类标记
生成对应 section 形态"]
+ IR["PTO IR"]
+ PTOAS["PTOAS
消费前端给出的分类"]
+ Output["EmitC / VPTO output"]
+
+ User --> Surface --> Trace --> IR --> PTOAS --> Output
+```
+
+新方案在相同层级上的流程如下:
+
+```mermaid
+flowchart LR
+ User["用户定义
tile-level helper"]
+ Surface["统一使用 @pto.tileop"]
+ Trace["PTODSL tracing
记录 body 和
pto.tileop.helper"]
+ IR["PTO IR"]
+ PTOAS["PTOAS
从 body 推导分类、阶段和读写
再生成对应 section"]
+ Output["EmitC / VPTO output"]
+
+ User --> Surface --> Trace --> IR --> PTOAS --> Output
+```
+
+核心变化是分类职责发生了转移:修改前由用户选择 `simd/cube`,PTODSL tracing
+把这个选择直接写进 IR;修改后用户统一写 `@pto.tileop`,PTODSL tracing 只记录
+helper body,PTOAS 再从 body 推导主计算域、阶段和读写关系。
+
+因此新方案简化的是用户入口和 PTODSL tracing 路径。PTOAS 内部新增的分析和校验会在第 7 节展开。
+
+## 5. Public contract
+
+`@pto.tileop` 的核心含义是:它是 tile/view/scalar 层面的 helper,不是 raw
+pointer helper,也不是 launched SIMT kernel。
+
+### 5.1 参数能传什么
+
+`@pto.tileop` 参数只允许表达 tile-level 边界的数据:
+
+- `Tile`:片上 tile buffer。
+- `TensorView`:tensor 视图,常用作 load/store 的逻辑输入或输出边界。
+- `PartitionTensorView`:已经切分好的 tensor 分区。
+- PTO scalar:例如 rows、cols、scale、flag、loop bound。
+
+不允许把这些值作为 `@pto.tileop` 参数:
+
+- `pto.ptr(...)`:raw pointer 只给 `@pto.simt` 使用。
+- host tensor / `TensorSpec`:这是 host/kernel entry 层面的整 tensor 或规格描述,
+ 不是 tileop helper 的边界。
+- vector register、mask、pipe handle 等 helper 内部临时值或控制值。
+
+### 5.2 结果怎么传出来
+
+tile/view 结果通过参数传入的可写 tile/view 表达,而不是作为 Python 函数返回值
+返回。
+
+例如用户应写成“把结果写进 `out_tile` / `out_view`”:
+
+```python
+@pto.tileop
+def helper(src: pto.TensorView, out: pto.Tile):
+ ...
+ pto.tile.store(..., out)
+```
+
+不要把 tile/view 当作函数 result 返回:
+
+```python
+@pto.tileop
+def helper(src: pto.TensorView) -> pto.Tile: # not allowed
+ ...
+```
+
+当前版本只允许 scalar result。也就是说,helper 可以返回类似状态值或小标量,
+但主要数据输出仍然应该写入传进来的 tile/view。
+
+### 5.3 helper body 的边界
+
+`@pto.tileop` helper body 可以包含数据搬运、主计算、标量控制和同步,但有几个
+边界:
+
+- 当前 tileop contract 明确禁止 helper 在函数内部自己申请新的 tile buffer。
+ 需要的 scratch/output tile 应由 caller 创建后作为参数传进来。实现层面会拒绝
+ `alloc_tile`、`reserve_buffer`、`pto.talloc` 这类 helper-local tile allocation。
+ 这个限制只针对 tileop helper;普通 `@pto.jit` kernel 内仍可以使用
+ `pto.alloc_tile(...)`。
+- helper 必须有一个主计算域:Vector 或 Cube。MTE load/store、Scalar、sync 可以
+ 作为辅助阶段存在,但它们不决定主计算域。
+- 一个 helper 里不要同时混用 Vector 主计算和 Cube 主计算。当前实现要求一个
+ tileop helper 只表达一个主计算域。
+- tileop helper 不调用另一个 tileop helper。这样每个 helper 的读写摘要都是
+ 局部、清楚、可验证的;递归组合多个 helper summary 留给后续设计。
+
+## 6. 新 IR 形态
+
+PTODSL tracing 时不再提前写 `primary_domain`、`phases`、`operand_effects`,
+也不预套 `pto.section.vector/cube`。PTODSL tracing 只把 helper 标成
+canonical marker:
+
+```mlir
+// PTODSL tracing 刚完成后的 helper 形态。
+// attributes 里只有 pto.tileop.helper,没有 primary_domain/phases/effects。
+func.func private @softmax(%src: !pto.tensor_view<...>,
+ %dst: !pto.tensor_view<...>,
+ %scratch: !pto.tile_buf<...>,
+ %rows: i32,
+ %cols: i32)
+ attributes {pto.tileop.helper} {
+ // body 里先保留 tracing 得到的原始 op 顺序。
+ // 这里还没有 pto.section.vector 或 pto.section.cube。
+ pto.tload ... outs(%scratch) : ... // MTE load
+ pto.vadd ... outs(%scratch) : ... // Vector compute
+ pto.tstore ... ins(%scratch) : ... // MTE store
+ return
+}
+```
+
+
+PTOAS 后续再基于 helper body 推导这些属性:
+
+```mlir
+attributes {
+ pto.tileop.helper,
+ pto.tileop.primary_domain = #pto.kernel_kind,
+ pto.tileop.phases = [...],
+ pto.tileop.operand_effects = [...]
+}
+```
+
+## 7. lib/PTO 详细设计
+
+前面的章节说明了 public surface 和 PTODSL tracing 后的 IR 形态。本节继续说明
+`lib/PTO` 中负责摘要推导、section 物化、contract 校验、同步建模和 VPTO split
+的具体修改。
+
+本节涉及的 pass / 流程状态如下:
+
+| pass / 流程 | 状态 | 这次承担的职责 |
+|---|---|---|
+| helper marker 识别 | 修改 | 新增 canonical `pto.tileop.helper` marker,同时兼容旧 `pto.ptodsl.subkernel_helper = "tileop"` |
+| `PTOInferTileOpSummaryPass` | 新增 | 从 helper body 推导 `primary_domain`、`phases`、`operand_effects` |
+| `PTOMaterializeTileOpSectionsPass` | 新增 | 根据 summary 把主计算 span 包成 `pto.section.vector/cube` |
+| `PTOVerifyTileOpContractPass` | 新增 | 校验 tileop ABI、body 约束和 summary 是否 stale |
+| InsertSync 的 `PTOIRTranslator` | 修改 | helper call 不再按单个粗粒度节点建模,而是按 `pto.tileop.phases` 建同步节点 |
+| `VPTOSplitCVModule` | 修改 | 在已单域 VPTO module 里也展开同域 section、删除异域 section,避免 section 留到后续 lowering |
+
+### 7.1 Pipeline 位置
+
+修改前,从同一个 PTOAS pipeline 视角看,tile-level helper 更依赖 PTODSL tracing
+已经写好的角色标记和 section 形态。PTOAS 不会先从 helper body 推导
+`primary_domain`、`phases` 和边界 effect,后续同步建模也更接近按一个 helper
+call 整体处理:
+
+```mermaid
+flowchart LR
+ IR["PTO IR
func.func + 前端角色标记"]
+
+ subgraph PTOAS["PTOAS"]
+ Role["消费前端角色标记
simd / cube"]
+ Section["使用前端预套
vector/cube section"]
+ Sync["InsertSync PTOIRTranslator
按整个 helper call 建模"]
+ Inline["helper inline
后续 PTOAS passes"]
+ Split["VPTOSplitCVModule
处理 vector/cube section"]
+ end
+
+ Lower["EmitC / VPTO lowering"]
+
+ IR --> Role --> Section --> Sync --> Inline --> Split --> Lower
+```
+
+PTODSL tracing 的输出是 PTO IR。带 `pto.tileop.helper` marker 的 `func.func` 进入
+PTOAS 后,关键数据流如下:
+
+```mermaid
+flowchart LR
+ IR["PTO IR
func.func + pto.tileop.helper"]
+
+ subgraph PTOAS["PTOAS"]
+ Infer["PTOInferTileOpSummaryPass
推导 primary_domain / phases / effects"]
+ Materialize["PTOMaterializeTileOpSectionsPass
包主计算 span"]
+ Verify["PTOVerifyTileOpContractPass
校验 ABI / body / summary"]
+ Sync["InsertSync PTOIRTranslator
按 phase 建模 helper call"]
+ Inline["helper inline
后续 PTOAS passes"]
+ Split["VPTOSplitCVModule
处理 vector/cube section"]
+ end
+
+ Lower["EmitC / VPTO lowering"]
+
+ IR --> Infer --> Materialize --> Verify --> Sync --> Inline --> Split --> Lower
+```
+
+### 7.2 如何标记一个 tileop helper
+
+PTODSL tracing 生成 PTO IR 时,需要告诉 PTOAS:“这个函数不是普通函数,而是一个
+tileop helper”。做法是在函数上加一个没有额外取值的标签:
+
+```mlir
+attributes {pto.tileop.helper}
+```
+
+看到这个标签后,后续 PTOAS pass 才会对该函数执行 tileop 专用的摘要推导、
+section 物化和 contract 校验。普通函数没有这个标签,也就不会进入这些处理。
+
+修改前,PTODSL 使用一个字符串属性同时表示不同 helper 类型:
+
+```mlir
+pto.ptodsl.subkernel_helper = "simd"
+pto.ptodsl.subkernel_helper = "cube"
+pto.ptodsl.subkernel_helper = "tileop"
+```
+
+新生成的 IR 统一使用专门的 `pto.tileop.helper` 标签,不再把 tileop 塞进旧的
+`subkernel_helper` 字符串。为了让已有 IR 仍能通过 PTOAS,代码暂时也能识别旧的
+`pto.ptodsl.subkernel_helper = "tileop"` 写法。
+
+这两种形式的统一识别封装在 `include/PTO/IR/PTO.h` 中。后续 pass 只需要询问
+“这个函数是不是 tileop helper”,不需要分别处理新旧属性格式。
+
+### 7.3 Summary 属性设计
+
+`PTOInferTileOpSummaryPass` 在 tileop helper 函数上生成三类 PTOAS-owned 属性。
+下面是示意形态,省略了完整 MLIR attribute assembly 的细节:
+
+```mlir
+attributes {
+ pto.tileop.helper,
+ pto.tileop.primary_domain = #pto.kernel_kind,
+ pto.tileop.phases = [
+ {
+ pipe = #pto.pipe,
+ operand_uses = [0],
+ operand_defs = [2],
+ result_defs = []
+ },
+ {
+ pipe = #pto.pipe,
+ operand_uses = [2],
+ operand_defs = [2],
+ result_defs = []
+ },
+ {
+ pipe = #pto.pipe,
+ operand_uses = [2],
+ operand_defs = [1],
+ result_defs = []
+ }
+ ],
+ pto.tileop.operand_effects = ["read", "write", "readwrite"]
+}
+```
+
+这个例子可以对应:
+
+```text
+operand #0: src_view
+operand #1: dst_view
+operand #2: scratch_tile
+```
+
+每个属性的含义是:
+
+| 属性 | 粒度 | 用途 |
+|---|---|---|
+| `pto.tileop.primary_domain` | 整个 helper | 记录主计算域是 `vector` 还是 `cube`,用于 section materialization 和 summary 校验 |
+| `pto.tileop.phases` | helper body 内的 pipe phase | 记录每个 phase 的 pipe,以及这个 phase 读写了哪些 helper operand;InsertSync 用它把一个 helper call 拆成多个同步建模节点 |
+| `pto.tileop.operand_effects` | helper operand | 记录整个 helper 对每个 operand 的合并 effect,用于 contract 校验和 stale-summary 检查 |
+
+`result_defs` 当前保持为空,因为当前设计里 tile/view 输出不通过 function result
+返回;function result 只允许 scalar。
+
+### 7.4 PTOInferTileOpSummaryPass
+
+`PTOInferTileOpSummaryPass` 是 summary-only pass。它只写属性,不改 helper body,
+也不插 section。
+
+它的算法可以概括成:
+
+```text
+按源代码顺序递归遍历 helper body 里的每个 op:
+ pipe = classify(op) # 判断 op 属于哪个 pipe
+ if pipe 是 Vector 或 Cube 主计算:
+ 设置或校验 primary_domain
+ if pipe 和上一个 body op 不同:
+ 开始一个新的 phase
+ if op 带有 MemoryEffect:
+ 把 op 实际读写的值追溯回 helper 参数
+ 在当前 phase 记录这个参数被读(use)还是被写(def)
+ 合并得到整个 helper 对该参数的总体 read/write effect
+```
+
+pipe 分类分两层:
+
+- 优先通过 op 自己声明的 `getPipe()` 获取它所属执行管线,例如 MTE、Vector、
+ Cube。`getPipe()` 来自 op 实现的 `OpPipeInterface`。
+- 对部分没有明确 pipe interface 的 PTO op,用名字做 fallback 判断。例如
+ `pto.v*` 识别为 Vector pipe,`pto.mad*` 识别为 Cube pipe,MTE load/store
+ 识别为对应 MTE pipe,标量 load/store 或 predicate op 识别为 Scalar pipe。
+
+主计算域只由 Vector/Cube primary compute 决定:
+
+- `PIPE_V` / `PIPE_V2` -> `primary_domain = vector`
+- `PIPE_M` -> `primary_domain = cube`
+- MTE、Scalar、sync 不决定 `primary_domain`
+
+如果同一个 helper 里同时出现 Vector primary compute 和 Cube primary compute,
+summary inference 直接报错。这对应当前“单主计算域”的实现边界。
+
+边界 effect 的关键点是:op 实际读写的值可能是 helper 参数派生出来的内部临时值,
+例如 tile slice、address value、cast 后的值或 loop iter_arg。PTOAS 不能只记录
+这些内部临时值,因为 helper call 的外部边界只看得到函数参数。
+
+例如:
+
+```python
+@pto.tileop
+def helper(tile: pto.Tile, row: pto.i32):
+ x = tile[row, 0:]
+ v = pto.vlds(x)
+ pto.vsts(v, tile[row, 0:])
+```
+
+在 IR 里,`tile[row, 0:]` 可能会变成 `memref.subview`。`vlds/vsts` 的 memory
+effect 挂在这个 subview 上,但 summary 需要继续追溯,知道这个 subview 来自
+helper 参数 `tile`。最终 phase 摘要记录的是“读/写了参数 `tile`”,而不是
+“读/写了内部临时值 subview”。
+
+当前实现会把 subview、cast、reshape、tile/view address、loop iter_arg 等视为
+可透明追溯的中间值。追溯回 helper 参数后,只有 `Tile`、`TensorView`、
+`PartitionTensorView` 这类 tileop 边界值会进入 tileop 的公开 effect 摘要。
+`ptr` / `memref` 只在内部分析里用于兼容低层 IR 形态;它们不会因此变成 tileop
+的合法 public ABI,verifier 仍会拒绝 `ptr` / `memref` 作为 tileop 参数。
+
+### 7.5 PTOMaterializeTileOpSectionsPass
+
+`PTOMaterializeTileOpSectionsPass` 消费 summary,把 helper body 中的主计算段包进
+`pto.section.vector` 或 `pto.section.cube`。
+
+它不是简单把整个 helper body 包进 section。原因是 helper 可能长这样:
+
+```text
+MTE load
+Scalar loop setup
+Vector compute
+MTE store
+```
+
+正确的 IR 形态应该是:
+
+```text
+MTE load
+Scalar loop setup
+pto.section.vector {
+ Vector compute
+}
+MTE store
+```
+
+而不是:
+
+```text
+pto.section.vector {
+ MTE load
+ Scalar loop setup
+ Vector compute
+ MTE store
+}
+```
+
+materialize pass 的具体规则是:
+
+1. 只处理标记为 tileop helper 的函数。
+2. 如果 helper 里已经有 `pto.section.vector/cube`,不重复 materialize。
+3. 读取 `pto.tileop.primary_domain` 和 `pto.tileop.phases`,确认存在 primary
+ phase。
+4. 递归进入普通控制流 region,例如 `scf.for` / `scf.if`。
+5. 在每个 block 内收集 PTO body op,找到与 `primary_domain` 匹配的 primary
+ compute span。
+6. 当前实现要求这个 primary compute span 是连续的。如果写成
+ `Vector compute A -> Scalar/MTE/sync -> Vector compute B`,materialize pass
+ 会报错。
+7. 对 vector span,向前扩展包含产生 `mask` / `vreg` 这类 vector-scope local 值的
+ producer,避免这些局部值被留在 section 外。
+8. 把最终 span 包进 `pto.section.vector` 或 `pto.section.cube`,MTE/S/sync 留在
+ section 外。
+
+控制流中的主计算也需要被找到。这里的 helper entry block,指函数 body 最外层的
+代码块。用户把计算写在循环里时,最外层直接包含的是 `scf.for`,真正的 MTE 和
+Vector op 则在 `scf.for` 的循环体里:
+
+```mlir
+scf.for %i = %c0 to %rows step %c1 {
+ pto.tload ... // MTE load
+ pto.tadds ... // Vector compute
+ pto.tstore ... // MTE store
+}
+```
+
+如果 pass 只检查函数最外层直接包含的 op,它只能看到一个 `scf.for`,看不到循环
+体里的 `pto.tadds`,因而会错误地认为这个 helper 没有 Vector 主计算。
+
+当前实现会把 `scf.for` 当作包含内部代码的控制流容器,继续进入它的循环体检查。
+找到 `pto.tadds` 后,只在循环体内部给这段 Vector 主计算加 section:
+
+```mlir
+scf.for %i = %c0 to %rows step %c1 {
+ pto.tload ... // 仍在 section 外
+ pto.section.vector {
+ pto.tadds ...
+ }
+ pto.tstore ... // 仍在 section 外
+}
+```
+
+因此这里的“递归”只是指:从函数最外层进入 `scf.for` / `scf.if` 等控制流的内部
+代码块,继续寻找主计算段。pass 不会因为主计算写在循环或分支里,就把整个循环或
+分支都包进 Vector/Cube section。
+
+### 7.6 PTOVerifyTileOpContractPass
+
+`PTOVerifyTileOpContractPass` 是防线:它不负责推导新信息,而是确认 helper body 和
+summary 没有违反 tileop contract。
+
+它检查的内容包括:
+
+- 参数类型只能是 `TileBufType`、`TensorViewType`、`PartitionTensorViewType` 或
+ PTO scalar。
+- function result 只能是 scalar。
+- helper 必须至少包含一个 Vector 或 Cube primary compute op。
+- 一个 helper 不能混用 Vector primary compute 和 Cube primary compute。
+- helper 内不能出现 SIMT-only op,例如 thread id、launch、vote、shuffle、
+ thread fence 等。
+- helper 内不能分配 callee-local tile buffer,例如 `pto.alloc_tile`、
+ `pto.reserve_buffer`、`pto.talloc`。
+- tileop helper 不能调用另一个 tileop helper。
+- `pto.tileop.primary_domain`、`pto.tileop.phases`、
+ `pto.tileop.operand_effects` 必须和重新扫描 body 得到的结果一致。
+
+最后一条很重要:summary attr 是 PTOAS 后续 pass 的输入,如果 body 被改了但
+summary 没更新,InsertSync 或 section materialization 看到的就会是 stale 信息。
+verifier 通过重新推导一遍 summary 来发现这种不一致。
+
+### 7.7 InsertSync 如何消费 tileop summary
+
+InsertSync 的任务是根据不同 pipe 对同一份数据的读写关系插入同步。对于直接出现
+在 kernel 里的 PTO op,PTOAS 可以从 op 本身知道它属于哪个 pipe、读取什么、写入
+什么。但是看到一次 helper 调用时:
+
+```mlir
+func.call @row_softmax(%src, %dst, %scratch) : ...
+```
+
+这条 `func.call` 只列出了传给 helper 的三个值,看不出 helper 内部先做 MTE load、
+再做 Vector compute、最后做 MTE store。因此,PTOAS 需要读取 `@row_softmax`
+函数上的 `pto.tileop.phases`,用它补回调用语句中看不到的阶段和读写信息。
+
+一次 tileop helper 调用按下面的步骤处理:
+
+1. `PTOIRTranslator` 找到被调用的 helper,读取它的 `pto.tileop.phases`。
+2. 将摘要里的 helper 参数编号,对应到这次调用实际传入的值。
+3. 每个读写了 helper 边界参数的 phase,生成一个 InsertSync 建模节点。
+4. `InsertSyncAnalysis` 比较这些节点的 pipe 和读写关系,判断哪里需要同步。
+
+转换后,每个建模节点记录所属 pipe、读取的值和写入的值。实现中这个节点对应
+`CompoundInstanceElement`,但它不是新的 PTO IR op,只是 InsertSync 分析期间
+使用的一条阶段记录。
+
+如果一个 phase 只操作 helper 内部的临时值,没有读写任何 helper 参数,PTOAS
+不会为它建立调用边界节点,因为它不会与 helper 外部的操作产生数据依赖。
+
+#### 7.7.1 和旧同步建模的对比
+
+这次修改没有更换 InsertSync 判断同步依赖的算法,改变的是
+`PTOIRTranslator` 如何向它描述一次 helper 调用。
+
+旧 `simd/cube` helper:
+
+```mermaid
+flowchart TB
+ OldCall["func.call @helper(%src, %dst, %scratch)"]
+ OldRole["读取前端给出的
simd / cube 分类"]
+ OldOperands["收集调用时传入的
tile / view 等值"]
+ OldNode["把整个调用表示成 1 个节点
所有边界值保守地视为既读又写"]
+ OldHazard["InsertSyncAnalysis
判断同步依赖"]
+
+ OldCall --> OldRole --> OldOperands --> OldNode --> OldHazard
+```
+
+当前 `tileop` helper:
+
+```mermaid
+flowchart TB
+ NewCall["func.call @helper(%src, %dst, %scratch)"]
+ NewSummary["读取 helper 的
pto.tileop.phases"]
+ NewMap["把 helper 参数编号
对应到这次调用实际传入的值"]
+ NewNodes["按 phase 生成节点
分别记录 pipe、读取值和写入值"]
+ NewHazard["InsertSyncAnalysis
判断同步依赖"]
+
+ NewCall --> NewSummary --> NewMap --> NewNodes --> NewHazard
+```
+
+| 对比项 | 旧 `simd/cube` helper | 当前 `tileop` helper |
+|---|---|---|
+| 信息来源 | 前端给出的 `simd/cube` 分类和调用参数 | PTOAS 从 helper body 推导的 phase 摘要 |
+| 建模粒度 | 整个 helper 调用是一个节点 | 一个有边界读写的 phase 是一个节点 |
+| 读写关系 | 边界值保守地视为既读又写 | 分别记录每个 phase 实际读取和写入的参数 |
+| pipe 信息 | 整个调用只有一个粗粒度分类 | 每个 phase 保留自己的 MTE、Vector 或 Cube pipe |
+
+例如一个 helper body 是:
+
+```text
+MTE load(src -> scratch)
+Vector compute(scratch)
+MTE store(scratch -> dst)
+```
+
+旧方案把整个调用描述成一个节点,无法表达三个阶段各自属于哪个 pipe、读写哪个
+参数。当前方案则把它表示成三条阶段记录:
+
+```text
+phase 0: pipe = MTE, use src, def scratch
+phase 1: pipe = Vector, use scratch, def scratch
+phase 2: pipe = MTE, use scratch, def dst
+```
+
+当 kernel 中前后出现其他 PTO op 或 helper 调用时,InsertSync 可以用这些阶段记录
+判断它们是否在不同 pipe 上读写了同一个值,再决定是否插入同步;不再需要把整个
+helper 调用保守地当作一次不可拆分的读写。
+
+### 7.8 VPTOSplitCVModule 的 section rewrite
+
+section 的作用有明确的生命周期:计算域尚未确定时用于区分 Vector/Cube;计算域
+确定后就应被移除,因为 section 容器本身不对应后端指令。
+
+```mermaid
+flowchart LR
+ Input["包含 vector/cube section 的 module"]
+ Kind{"已有 pto.kernel_kind?"}
+ Split["按 section 拆成
Vector / Cube module"]
+ Rewrite["展开同域 section
删除异域 section"]
+ Ready["不再包含 section
进入 VPTO LLVM lowering"]
+
+ Input --> Kind
+ Kind -- "否" --> Split --> Rewrite --> Ready
+ Kind -- "是" --> Rewrite
+```
+
+`VPTOSplitCVModule` 的处理可以概括为:
+
+```text
+if module 已有 kernel_kind:
+ 展开同域 section,删除异域 section
+else:
+ 根据 vector/cube section 拆分 module
+ 在每个新 module 中展开同域 section,删除异域 section
+```
+
+主干原有实现会在“已有 `kernel_kind`”时直接返回,导致 helper inline 后的 section
+可能遗留到 LLVM lowering。本次修改补上该分支,并增加包含 `scf.for` 的回归测试。
+
+### 7.9 设计边界
+
+这套 `lib/PTO` 设计刻意把几件事分开:
+
+- summary inference 只负责“看懂 helper”,不改 IR body。
+- section materialization 只负责“把主计算 span 结构化”,不重新定义 ABI。
+- contract verifier 负责“发现非法 helper 或 stale summary”,不修正用户代码。
+- InsertSync 只消费 summary 建 dependency graph,不再重新进入 callee body 分析。
+- VPTO split 只处理 section sugar 和 module kind,不参与 tileop ABI 判断。
+
+这种拆分让每个 pass 的输入输出都比较明确:前一个 pass 产出的属性或 section,
+是后一个 pass 的显式输入;如果中途被改坏,verifier 会尽早报错。
+
+## 8. 相比旧方案新增或收紧的限制
+
+旧 `@pto.simd` / `@pto.cube` 路径主要依赖 frontend role 和预套 section。新
+`@pto.tileop` 路径把 helper 交给 PTOAS 推导和校验,因此也把一些边界变成了明确
+诊断。先从用户写 PTODSL 时能直接感知到的限制看:
+
+### 8.1 用户写 PTODSL 时会看到的限制
+
+| 用户写法 | 旧方案形态 | 当前 tileop 形态 |
+|---|---|---|
+| `@pto.simd` / `@pto.cube` | 作为 public decorator 使用 | 作为旧接口诊断报错,tile-level helper 统一写 `@pto.tileop` |
+| tileop 参数 | 旧 simd/cube 的参数边界不统一 | `@pto.tileop` 只允许 `Tile`、`TensorView`、`PartitionTensorView`、PTO scalar |
+| `pto.ptr(...)` 参数 | 容易被误认为 tile-level helper 也能接收 | 仍然只允许 `@pto.simt`;`@pto.tileop` 不能用 raw pointer 作为 public boundary |
+| tile/view 输出 | 容易被误写成 function result | 通过 output operand 传出;function result 只允许 scalar |
+| helper 内申请 tile | 普通 kernel 里可以 `pto.alloc_tile(...)` | `@pto.tileop` helper 内不能新建 helper-local tile buffer;scratch/output tile 由 caller 传入 |
+| helper body | 由 `simd` / `cube` 名字暗示主域 | body 必须真的包含 Vector 或 Cube primary compute,不能只有 MTE/Scalar/sync |
+| 混合 Vector/Cube | 旧名字无法清楚表达混合语义 | 同一个 `@pto.tileop` helper 不能同时写 Vector primary compute 和 Cube primary compute |
+| helper 互调 | 没有清晰的 phase summary 组合语义 | `@pto.tileop` helper 不能调用另一个 `@pto.tileop` helper |
+| SIMT 操作 | 容易和 tile-level helper 混用 | thread id、launch、vote、shuffle、thread fence 等 SIMT-only op 只放在 `@pto.simt` |
+
+### 8.2 PTOAS 内部额外收紧的检查
+
+还有一些限制用户不一定会在 Python 写法里直接看到,但会影响手写 IR、其他
+frontend,或后续 PTOAS pass 的输入:
+
+| 检查项 | 旧方案形态 | 当前 tileop 形态 |
+|---|---|---|
+| 后端 helper ABI | frontend / backend 边界规则容易分叉 | verifier 在 PTO IR 层拒绝 `ptr` / `memref` 作为 tileop helper 参数 |
+| summary 一致性 | 没有 `primary_domain/phases/effects` 这组属性 | verifier 会重新扫描 body,确认 summary 没有 stale |
+| 主计算 span | 旧路径常把整个 helper body 预套进一个 section | 当前只 materialize 主计算 span;同一个 block 内 primary compute 中间不能夹非 primary pipe |
+| sync 建模输入 | helper call 可被保守看成单个 read/write 节点 | InsertSync 按 `pto.tileop.phases` 建多个 phase 节点,summary 缺失或错误会被拒绝 |
+| section rewrite | section 可能留到后续 lowering | 已单域 VPTO module 里也会展开同域 section、删除异域 section |
+
+这些限制不是为了减少 tileop 能力,而是为了让当前版本的 ABI、summary、sync 和
+section 语义先保持可验证。后续如果要放开某一项,需要同时定义它对 summary、
+InsertSync 和 VPTO section rewrite 的影响。
+
+## 9. 方案同时覆盖的问题和收益
+
+前面第 7 节解释了 `lib/PTO` 如何落地这套设计,第 8 节列出了当前实现收紧的边界。
+本节再总结这套实现同时带来的收益和覆盖面。它们不是 issue 一开始逐条提出的问题,
+但统一到 `@pto.tileop` 后,这些相关边界可以一起收敛。
+
+### 9.1 tile-level helper 只保留一套 public contract
+
+旧 API 把 helper 名字直接命名成 `simd` 或 `cube`。这看起来直观:做 vector
+计算就写 `simd`,做 cube 计算就写 `cube`。但对 PTOAS 来说,这两类 helper
+共享同一个更重要的 contract:它们都是 tile-level helper,参数边界应该是
+tile/view/scalar,调用边界需要有 read/write effect,body 里也都可能包含
+MTE、主计算、Scalar/sync 等多个 phase。
+
+统一成 `@pto.tileop` 后,vector-style 和 cube-style 不再拥有两套 public
+boundary 规则、两套诊断文案、两套 inline/decorated 处理路径。区别保留在
+`primary_domain = vector/cube`,由 PTOAS 从 body 推导。
+
+### 9.2 phase summary 让跨 helper 调用的同步更清楚
+
+这里要区分两件事:一是 helper body 里实际出现了哪些 pipe(helper 真实做了
+什么),二是 PTOAS 在 helper 调用边界能看到的摘要粒度(PTOAS 处理连续 helper
+调用时能拿到多少内部信息)。例如面对:
+
+```text
+call @row_softmax(...)
+call @row_softmax(...)
+```
+
+这类连续 helper 调用,PTOAS 能看到多少内部读写信息,就取决于摘要粒度。
+
+旧 `simd`/`cube` 路径会把整个 helper call 压成一个粗粒度节点,例如“这是一个
+vector-style call”。这不表示 helper 里面真的只有 Vector pipe,而是表示 PTOAS
+只能把这个 call 当成一个整体来建模。它看不到这个 call 内部其实有
+`MTE load -> Vector compute -> MTE store` 三个阶段,也就无法分别知道:
+
+- MTE load 阶段读哪个输入 view、写哪个 scratch tile。
+- Vector compute 阶段读写哪个 tile。
+- MTE store 阶段读哪个 tile、写哪个输出 view。
+
+这就是“跨 call 建同步时丢掉阶段结构”的意思:不是 IR 里的操作消失了,而是
+PTOAS 在两个 helper call 之间判断依赖和插同步时,只拿到一个大粒度 call 摘要,
+拿不到每个 phase 的读写摘要。
+
+新的 `pto.tileop.phases` 用来表达这些 phase-level 信息。注意这不表示 MTE
+load/store 属于 Vector;它们仍是 MTE pipe,只是和 Vector 主计算共同构成一个
+`primary_domain = vector` 的 tileop helper。
+
+### 9.3 section 由 PTOAS 物化,避免 PTODSL tracing 过早判断
+
+section 是 PTO IR 里的结构化容器,用来告诉 PTOAS “这一段是 vector 主计算”
+或“这一段是 cube 主计算”:
+
+```mlir
+pto.section.vector {
+ // vector primary compute
+}
+```
+
+section 不是“整个 helper 的容器”。它只应该包主计算段。MTE load/store 是数据
+搬运,Scalar/sync 是辅助控制或同步;这些 op 通常应该留在主计算 section 外面。
+也就是说,一个典型 helper 更接近:
+
+```text
+MTE load
+pto.section.vector {
+ Vector compute
+}
+MTE store
+```
+
+新方案让 PTODSL tracing 只标记 `pto.tileop.helper`,不预套 section。PTOAS 先
+推导 summary,再递归进入控制流 body,找到连续的主计算 span 后只包这段,
+MTE/S/sync 保持在 section 外。这样 decorated helper 和 inline helper 也能走
+同一条 contract。
+
+### 9.4 tileop ABI 收敛到 Tile/View/Scalar,ptr 保留给 SIMT
+
+ABI 在这里指 helper 的参数和返回值 contract:哪些值允许跨 helper 边界,哪些
+值只能留在 helper 内部。对 tileop 来说,用户关心的是 tile-level 数据结构,而
+不是 raw pointer:
+
+- `Tile` 表示片上 tile buffer。
+- `TensorView` 表示一个 tensor 视图,可作为 load/store 的逻辑边界。
+- `PartitionTensorView` 表示已经切分好的 tensor 分区。
+- PTO scalar 表示 rows、cols、scale、loop bound 等标量参数。
+
+这些类型足够表达 tile-level helper 的自然边界。例如 softmax helper 可以读
+一个 `TensorView`,使用一个 `Tile` 做 scratch,再写一个 output view。
+
+`pto.ptr(...)` 不同。它是 raw pointer 风格的低层 ABI,适合 SIMT launched
+helper。raw pointer 只表达一段地址,不携带 tile/view 的结构化边界语义。因此
+tileop contract 收敛为:
+
+```text
+tileop: Tile / TensorView / PartitionTensorView / PTO scalar
+simt: pto.ptr(...) and SIMT-specific boundary
+```
+
+### 9.5 控制流和 tile slice 的建模更完整
+
+PTODSL 用户写 Python `for`/`if` 时,PTODSL tracing 会把它们记录成设备侧控制流,
+而不是在 Python 编译期展开。MLIR 里常见形式是 `scf.for` / `scf.if`。PTODSL
+用户写 `tile[row, col:]` 这类切片时,IR 里常见形式是 `memref.subview` 或 PTO
+自己的 view/tile address op。
+
+tracing 到 MLIR 后,它们会变成类似这样的结构:
+
+```mlir
+scf.for %i = %c0 to %rows step %c1 {
+ %slice = memref.subview %tile_addr[%i, 0] [1, 64] [1, 1]
+ %v = pto.vlds %slice[%c0] : ...
+ pto.vsts %v, %out_slice[%c0], %mask : ...
+}
+```
+
+对不熟 MLIR 的读者,可以把它理解成:
+
+- `scf.for` 是 IR 里的结构化 loop。
+- `memref.subview` 是从一个 tile/view 中取出一段 slice,例如第 `i` 行。
+- `vlds/vsts` 可能真正读写的是 subview 结果,而不是函数参数本身。
+
+这套方案要求 PTOAS 递归查看控制流 body,能看到 loop/if 内的主计算 op;同时
+识别由 helper 边界 operand 派生出的透明 view、tile address、cast、subview 等
+来源,能把 tile slice 的 effect 归回对应 helper 参数。这样 PTOAS 才能为常见
+row-wise/tile-slice helper 建出正确的跨 call 依赖。
+
+## 10. 用户迁移方式
+
+用户原来写:
+
+```python
+@pto.simd
+def helper(...):
+ ...
+
+@pto.cube
+def matmul_helper(...):
+ ...
+```
+
+现在统一写:
+
+```python
+@pto.tileop
+def helper(...):
+ ...
+```
+
+如果是 launched SIMT helper,仍然写:
+
+```python
+@pto.simt
+def helper(ptr: pto.ptr(pto.f32, "gm")):
+ ...
+```
+
+判断标准是:
+
+- 参数是 tile/view/scalar,body 是 tile-level load/compute/store:用
+ `@pto.tileop`。
+- 参数是 raw pointer,body 是 SIMT-style launched code:用 `@pto.simt`。
+
+旧代码里如果看到 `@pto.simd` 或 `@pto.cube`,迁移时不要按名字一比一替换成
+新的 Vector/Cube 专用 decorator,因为新设计没有这两个专用 decorator。统一改
+成 `@pto.tileop`,再让 PTOAS 从 body 推导 vector-style 或 cube-style 主域。
+
+## 11. 常见问题
+
+### MTE load/store 属于 vector 还是 cube?
+
+都不属于。MTE load/store 属于 MTE pipe,是数据搬运阶段。一个 helper 的
+`primary_domain` 由主计算 phase 决定:
+
+- `MTE load -> Vector compute -> MTE store` 的主域是 `vector`。
+- `MTE load -> Cube compute -> MTE store` 的主域是 `cube`。
+
+因此文档中说 “vector-style helper 包含 MTE phase” 时,不是在说 MTE 属于
+Vector,而是在说同一个 tileop helper 的执行过程可能包含多个 pipe。
+
+### 为什么 tileop 只能有一个 primary_domain?
+
+这是当前实现边界。一个 helper 如果同时包含 Vector 主计算和 Cube 主计算,PTOAS
+需要决定如何拆 section、如何拆 kernel module、如何在两个主域之间建同步。这些
+规则比单主域复杂很多。当前 contract 要求一个 tileop helper 只表达一个主计算
+域,MTE/S/sync 作为辅助 phase。
+
+### 为什么 `pto.ptr(...)` 不能传给 tileop?
+
+`pto.ptr(...)` 是低层 pointer ABI,适合 SIMT launched helper。tileop 的目标是
+让 helper 站在 tile/view/scalar 层面建模,这样 PTOAS 能基于结构化信息做
+memory、sync 和 section 处理。允许 raw pointer 进入 tileop ABI 会绕开这些
+结构化边界,所以仍保留为 SIMT-only。
+
+### 为什么不把所有 operand 都统一标成 readwrite?
+
+统一 readwrite 最保守,通常不容易漏同步,但会丢掉 phase 级别的真实依赖,
+可能引入不必要的同步,也让 PTOAS 无法判断哪些 phase 真的跨 helper 边界产生
+effect。`pto.tileop.phases` 的目的就是让 helper call 的同步模型更接近真实
+读写行为。
+
+### 为什么要禁止 tileop helper 调用另一个 tileop helper?
+
+如果 helper A 调 helper B,PTOAS 需要递归组合两个 helper 的 summary、phase 和
+operand effect,还要处理 inline 后可能产生的跨域 section 嵌套。当前版本先
+禁止 tileop helper 互调,保证每个 helper 的 summary 是局部可验证、可消费的。
+
+## 12. 测试覆盖
+
+本次修改对应的测试覆盖分几层:
+
+- PTODSL Python tests:验证 public API、legacy diagnostic、inline/decorated
+ tileop 行为。
+- PTO lit tests:验证 helper ABI、summary、contract verifier、memory planning
+ 和 sync。
+- VPTO lit tests:验证 section materialization、control-flow body、single-kind
+ module section rewrite。
+- DSL ST simulator:验证实际 PTODSL runtime/simulator 路径,例如
+ `vmulscvt.py`、`predicate_pack.py`、`cube_matrix_pipeline.py`。
+
+关键回归包括:
+
+- tileop 允许 `Tile` / `TensorView` / `PartitionTensorView` / scalar。
+- tileop 拒绝 `ptr`;SIMT 仍接受 `ptr`。legacy `simd/cube` public surface 会
+ 先被 PTODSL frontend 迁移诊断拒绝。
+- tile slice 经 `memref.subview` 后仍能被 summary effect 识别。
+- control-flow 内的 primary compute 能被 materialize 成 section。
+- mask/vreg local producer 不会逃出 vector scope。
+- 已单域 VPTO module 中的 section 会在 lowering 前被展开。
+
+## 13. 当前仍需注意的点
+
+第 8 节已经列出相比旧方案新增或收紧的限制。除此之外,使用当前实现时还要注意:
+
+- function result 只允许 scalar;tile/view 输出仍通过 output operand 表达。
+- 如果一个 operand 没有被任何 op 显式读写,它的 effect 默认记为 `read`。只有
+ 实现了 `MemoryEffectOpInterface`、且读写值能追溯回 helper 参数的 op,才会真正
+ 影响 phase summary 里的 use/def。
+
+这些边界不是因为 PTOAS/VPTO 永远无法支持,而是为了让当前 public contract 和
+sync/memory/codegen 行为先稳定下来。
diff --git a/include/PTO/IR/PTO.h b/include/PTO/IR/PTO.h
index e858212835..a4d76a0fa9 100644
--- a/include/PTO/IR/PTO.h
+++ b/include/PTO/IR/PTO.h
@@ -190,9 +190,36 @@ inline constexpr llvm::StringLiteral kPTOSimtMaxRegistersAttrName =
inline constexpr llvm::StringLiteral kPTOVisibilityAttrName = "pto.visibility";
inline constexpr llvm::StringLiteral kPTOVisibilityInternalValue = "internal";
inline constexpr llvm::StringLiteral kPTOVisibilityExternalValue = "external";
+inline constexpr llvm::StringLiteral kPTODSLSubkernelHelperAttrName =
+ "pto.ptodsl.subkernel_helper";
+inline constexpr llvm::StringLiteral kPTOTileOpHelperAttrName =
+ "pto.tileop.helper";
inline constexpr llvm::StringLiteral kPTODSLLogicalNameAttrName =
"pto.ptodsl.logical_name";
+/// Return the logical PTODSL helper role when present.
+///
+/// Canonical tileop helpers use the unit attr `pto.tileop.helper`. Legacy
+/// helper roles still use `pto.ptodsl.subkernel_helper = ""`.
+inline StringRef getPTODSLSubkernelHelperRole(::mlir::func::FuncOp func) {
+ if (!func)
+ return {};
+ if (func->hasAttrOfType(kPTOTileOpHelperAttrName))
+ return "tileop";
+ if (auto attr =
+ func->getAttrOfType(kPTODSLSubkernelHelperAttrName))
+ return attr.getValue();
+ return {};
+}
+
+inline bool hasPTODSLSubkernelHelperMarker(::mlir::func::FuncOp func) {
+ return !getPTODSLSubkernelHelperRole(func).empty();
+}
+
+inline bool isPTODSLTileOpHelper(::mlir::func::FuncOp func) {
+ return getPTODSLSubkernelHelperRole(func) == "tileop";
+}
+
/// Return the PTODSL logical function name when present, otherwise fall back to
/// the current symbol name. PTODSL uses this to mark ABI-specialized helper and
/// kernel-module symbols without relying on symbol-name parsing.
diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h
index 3de31a89bf..a955bca034 100644
--- a/include/PTO/Transforms/Passes.h
+++ b/include/PTO/Transforms/Passes.h
@@ -40,6 +40,9 @@ std::unique_ptr createPTOInferValidatePipeInitPass();
std::unique_ptr createPTOResolveReservedBuffersPass();
std::unique_ptr createPTOWrapFunctionsInSectionsPass();
std::unique_ptr createPTONormalizeUncoveredTileSectionsPass();
+std::unique_ptr createPTOInferTileOpSummaryPass();
+std::unique_ptr createPTOMaterializeTileOpSectionsPass();
+std::unique_ptr createPTOVerifyTileOpContractPass();
std::unique_ptr createVPTOSplitCVModulePass();
std::unique_ptr createVPTONormalizeContainerPass();
std::unique_ptr createPTOVerifyTFreePass();
diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td
index 04165718b9..ac4182c02c 100644
--- a/include/PTO/Transforms/Passes.td
+++ b/include/PTO/Transforms/Passes.td
@@ -340,6 +340,85 @@ def PTONormalizeUncoveredTileSections
];
}
+def PTOInferTileOpSummary
+ : Pass<"pto-infer-tileop-summary", "func::FuncOp"> {
+ let summary = "Infer phase summary attributes for PTODSL tileop helpers";
+ let description = [{
+ Scans functions marked with canonical `pto.tileop.helper` (while still
+ accepting legacy `pto.ptodsl.subkernel_helper = "tileop"`) and
+ derives the backend-owned summary attributes:
+ - `pto.tileop.primary_domain`
+ - `pto.tileop.phases`
+ - `pto.tileop.operand_effects`
+
+ This pass is intentionally summary-only. It does not materialize sections,
+ inline helpers, or change InsertSync modeling.
+ }];
+
+ let constructor = "mlir::pto::createPTOInferTileOpSummaryPass()";
+
+ let dependentDialects = [
+ "mlir::func::FuncDialect",
+ "mlir::pto::PTODialect"
+ ];
+}
+
+def PTOMaterializeTileOpSections
+ : Pass<"pto-materialize-tileop-sections", "func::FuncOp"> {
+ let summary = "Materialize one primary PTO section for PTODSL tileop helpers";
+ let description = [{
+ Consumes backend-owned tileop summary attributes on functions marked with
+ canonical `pto.tileop.helper` (while still accepting legacy
+ `pto.ptodsl.subkernel_helper = "tileop"`) and wraps one contiguous
+ primary-domain compute span in `pto.section.vector` or `pto.section.cube`.
+
+ Unlike `pto-normalize-uncovered-tile-sections`, this pass is helper-local
+ and runs after tileop phase summary inference. It uses those inferred phases
+ to preserve non-primary MTE/S/sync operations outside the materialized
+ primary-domain section.
+
+ The MVP implementation expects one contiguous primary compute span in the
+ helper body. Leading and trailing MTE phases remain top-level so late
+ helper inlining can expose VPTO section sugar to `vpto-split-cv-module`.
+ }];
+
+ let constructor = "mlir::pto::createPTOMaterializeTileOpSectionsPass()";
+
+ let dependentDialects = [
+ "mlir::func::FuncDialect",
+ "mlir::pto::PTODialect"
+ ];
+}
+
+def PTOVerifyTileOpContract
+ : Pass<"pto-verify-tileop-contract", "func::FuncOp"> {
+ let summary = "Verify backend-owned PTODSL tileop helper contracts";
+ let description = [{
+ Verifies functions marked with canonical `pto.tileop.helper` (while still
+ accepting legacy `pto.ptodsl.subkernel_helper = "tileop"`) after summary
+ inference and section materialization.
+
+ Current MVP contract:
+ - results must be PTO scalar values only
+ - helper-local `pto.alloc_tile`, `pto.reserve_buffer`, and `pto.talloc`
+ are rejected
+ - nested tileop helper calls are rejected
+ - SIMT-only PTO ops are rejected
+ - at least one primary compute op must exist, and all primary compute ops
+ must belong to exactly one domain (`vector` or `cube`)
+ - scalar/MTE/sync phases may coexist, but `pto.tileop.primary_domain`,
+ `pto.tileop.phases`, and `pto.tileop.operand_effects` must remain
+ consistent with the helper body
+ }];
+
+ let constructor = "mlir::pto::createPTOVerifyTileOpContractPass()";
+
+ let dependentDialects = [
+ "mlir::func::FuncDialect",
+ "mlir::pto::PTODialect"
+ ];
+}
+
def VPTOSplitCVModule : Pass<"vpto-split-cv-module", "ModuleOp"> {
let summary = "Split a VPTO module with cube/vector sections into kernel modules";
let description = [{
@@ -586,7 +665,8 @@ def PTOInlineBackendHelpers
let description = [{
Force-inlines backend helper functions that should not survive to backend-
specific lowering:
- - PTODSL subkernel helpers marked with `pto.ptodsl.subkernel_helper`
+ - PTODSL subkernel helpers marked with `pto.tileop.helper` or legacy
+ `pto.ptodsl.subkernel_helper`
This pass runs on the shared mainline before backend-specific pipelines so
both VPTO and EmitC consume the same helper-inlined IR.
@@ -622,10 +702,10 @@ def PTOVerifySubkernelPipeContract
: Pass<"pto-verify-subkernel-pipe-contract", "func::FuncOp"> {
let summary = "Verify PTODSL subkernel helpers stay within one InsertSync pipe contract";
let description = [{
- Verifies PTODSL subkernel helpers marked with `pto.ptodsl.subkernel_helper`
- and their same-module local call closure stay within one role-consistent PTO
- compute contract before InsertSync models the call boundary as one compound
- node.
+ Verifies legacy PTODSL subkernel helpers marked with
+ `pto.ptodsl.subkernel_helper` and their same-module local call closure stay
+ within one role-consistent PTO compute contract before InsertSync models
+ the call boundary as one compound node.
Current contract:
- `simd` helpers may only contain PTO ops on `PIPE_V`
diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp
index fca999855e..0ad9677a97 100644
--- a/lib/PTO/IR/PTO.cpp
+++ b/lib/PTO/IR/PTO.cpp
@@ -14569,9 +14569,15 @@ getEnclosingFunctionKernelKind(Operation *op) {
return kernelKindAttr.getKernelKind();
}
+static bool isInsideTileOpSubkernelHelper(Operation *op) {
+ auto funcOp = op->getParentOfType();
+ return pto::isPTODSLTileOpHelper(funcOp);
+}
+
static bool isInsideSectionOrAttributedKernel(Operation *op) {
return isInsideSectionCube(op) || isInsideSectionVector(op) ||
- getEnclosingFunctionKernelKind(op).has_value();
+ getEnclosingFunctionKernelKind(op).has_value() ||
+ isInsideTileOpSubkernelHelper(op);
}
static LogicalResult verifySplitAttr(Operation *op, int64_t split) {
@@ -14583,6 +14589,8 @@ static LogicalResult verifySplitAttr(Operation *op, int64_t split) {
static LogicalResult verifyFrontendKernelKind(Operation *op,
FunctionKernelKind expected,
StringRef kernelName) {
+ if (isInsideTileOpSubkernelHelper(op))
+ return success();
if (isInsideSectionCube(op)) {
if (expected == FunctionKernelKind::Cube)
return success();
@@ -15860,7 +15868,7 @@ LogicalResult InitializeL2LPipeOp::verify() {
LogicalResult TPushOp::verify() {
if (!isInsideSectionOrAttributedKernel(getOperation()))
- return emitOpError("must be inside pto.section.cube/vector or a kernel_kind function");
+ return emitOpError("must be inside pto.section.cube/vector, a kernel_kind function, or a tileop subkernel helper");
if (failed(verifyPipeHandleProducer(getOperation(), getPipeHandle())))
return failure();
if (failed(verifySplitAttr(getOperation(), getSplit())))
@@ -15876,7 +15884,7 @@ LogicalResult TPushOp::verify() {
LogicalResult TAllocOp::verify() {
if (!isInsideSectionOrAttributedKernel(getOperation()))
- return emitOpError("must be inside pto.section.cube/vector or a kernel_kind function");
+ return emitOpError("must be inside pto.section.cube/vector, a kernel_kind function, or a tileop subkernel helper");
if (failed(verifyPipeHandleProducer(getOperation(), getPipeHandle())))
return failure();
if (failed(verifyTensorEntryMatchesInternalPipeInit(
@@ -15887,7 +15895,7 @@ LogicalResult TAllocOp::verify() {
LogicalResult TPopOp::verify() {
if (!isInsideSectionOrAttributedKernel(getOperation()))
- return emitOpError("must be inside pto.section.cube/vector or a kernel_kind function");
+ return emitOpError("must be inside pto.section.cube/vector, a kernel_kind function, or a tileop subkernel helper");
if (failed(verifyPipeHandleProducer(getOperation(), getPipeHandle())))
return failure();
if (failed(verifySplitAttr(getOperation(), getSplit())))
@@ -15904,7 +15912,7 @@ LogicalResult TPopOp::verify() {
LogicalResult TFreeOp::verify() {
if (!isInsideSectionOrAttributedKernel(getOperation()))
- return emitOpError("must be inside pto.section.cube/vector or a kernel_kind function");
+ return emitOpError("must be inside pto.section.cube/vector, a kernel_kind function, or a tileop subkernel helper");
if (failed(verifyPipeHandleProducer(getOperation(), getPipeHandle())))
return failure();
if (getEntry() &&
diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt
index 3a6c04aed0..fb914c4585 100644
--- a/lib/PTO/Transforms/CMakeLists.txt
+++ b/lib/PTO/Transforms/CMakeLists.txt
@@ -69,6 +69,9 @@ add_mlir_dialect_library(PTOTransforms
SlotAffineAnalysis.cpp
PTOWrapFunctionsInSectionsPass.cpp
PTONormalizeUncoveredTileSections.cpp
+ PTOInferTileOpSummaryPass.cpp
+ PTOMaterializeTileOpSectionsPass.cpp
+ PTOVerifyTileOpContractPass.cpp
VPTONormalizeContainer.cpp
VPTOSplitCVModule.cpp
InsertSync/PTOIRTranslator.cpp
diff --git a/lib/PTO/Transforms/InsertSync/InsertSyncDebug.cpp b/lib/PTO/Transforms/InsertSync/InsertSyncDebug.cpp
index a01a50e1bc..72586626d5 100644
--- a/lib/PTO/Transforms/InsertSync/InsertSyncDebug.cpp
+++ b/lib/PTO/Transforms/InsertSync/InsertSyncDebug.cpp
@@ -14,6 +14,7 @@
#include "mlir/IR/AsmState.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/FormatVariadic.h"
+#include
using namespace mlir;
using namespace mlir::pto;
@@ -28,6 +29,8 @@ llvm::cl::opt insertSyncDebugLevelOpt(
"0=off, 1=phase, 2=syncir, 3=trace"),
llvm::cl::init(0));
+std::mutex insertSyncDebugDumpMutex;
+
} // namespace
unsigned mlir::pto::getInsertSyncDebugLevel() { return insertSyncDebugLevelOpt; }
@@ -321,16 +324,22 @@ void mlir::pto::dumpInsertSyncPhase(llvm::StringRef phase, const SyncIRs &syncIR
}
}
- os << "\n// === [PTOInsertSync Debug] " << phase << " === //\n";
- os << llvm::formatv("// nodes={0}, syncGroups={1}, activeOps={2} "
- "(set={3}, wait={4}, barrier={5}, blockSet={6}, "
- "blockWait={7}, blockAll={8})\n",
- syncIR.size(), syncOperations.size(), activeOps, setCnt,
- waitCnt, barrierCnt, blockSetCnt, blockWaitCnt,
- blockAllCnt);
+ std::string buffer;
+ llvm::raw_string_ostream bufferedOS(buffer);
+
+ bufferedOS << "\n// === [PTOInsertSync Debug] " << phase << " === //\n";
+ bufferedOS << llvm::formatv("// nodes={0}, syncGroups={1}, activeOps={2} "
+ "(set={3}, wait={4}, barrier={5}, blockSet={6}, "
+ "blockWait={7}, blockAll={8})\n",
+ syncIR.size(), syncOperations.size(), activeOps,
+ setCnt, waitCnt, barrierCnt, blockSetCnt,
+ blockWaitCnt, blockAllCnt);
if (level < static_cast(InsertSyncDebugLevel::SyncIR)) {
- os << "// ========================================= //\n";
+ bufferedOS << "// ========================================= //\n";
+ bufferedOS.flush();
+ std::lock_guard lock(insertSyncDebugDumpMutex);
+ os << buffer;
return;
}
@@ -340,6 +349,10 @@ void mlir::pto::dumpInsertSyncPhase(llvm::StringRef phase, const SyncIRs &syncIR
options.showMemInfo = showMemInfo;
options.showUselessSync = showMemInfo;
- dumpSyncIR(os, syncIR, opForPrinting, options, showMemInfo);
- os << "// ========================================= //\n";
+ dumpSyncIR(bufferedOS, syncIR, opForPrinting, options, showMemInfo);
+ bufferedOS << "// ========================================= //\n";
+ bufferedOS.flush();
+
+ std::lock_guard lock(insertSyncDebugDumpMutex);
+ os << buffer;
}
diff --git a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp
index cdf87a29b7..907813aa58 100644
--- a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp
+++ b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp
@@ -35,6 +35,12 @@ using namespace mlir::pto;
namespace {
+static constexpr llvm::StringLiteral kTileOpPrimaryDomainAttr =
+ "pto.tileop.primary_domain";
+static constexpr llvm::StringLiteral kTileOpPhasesAttr = "pto.tileop.phases";
+static constexpr llvm::StringLiteral kTileOpOperandEffectsAttr =
+ "pto.tileop.operand_effects";
+
constexpr size_t kTileRank2D = 2;
constexpr unsigned kStrideInlineCapacity = 4;
constexpr unsigned kMemoryEffectInlineCapacity = 4;
@@ -210,30 +216,104 @@ static func::FuncOp lookupPTODSLSubkernelHelper(func::CallOp callOp) {
auto callee = module.lookupSymbol(callOp.getCallee());
if (!callee)
return {};
- if (!callee->hasAttr("pto.ptodsl.subkernel_helper"))
+ if (!pto::hasPTODSLSubkernelHelperMarker(callee))
return {};
return callee;
}
+static StringRef getResolvedPTODSLSubkernelHelperRole(func::FuncOp callee) {
+ return pto::getPTODSLSubkernelHelperRole(callee);
+}
+
static std::optional
getPTODSLSubkernelHelperPipe(func::FuncOp callee) {
- auto roleAttr =
- callee->getAttrOfType("pto.ptodsl.subkernel_helper");
- if (!roleAttr)
+ StringRef role = getResolvedPTODSLSubkernelHelperRole(callee);
+ if (role.empty())
return std::nullopt;
return llvm::StringSwitch>(
- roleAttr.getValue())
+ role)
.Case("cube", pto::PipelineType::PIPE_M)
.Case("simd", pto::PipelineType::PIPE_V)
.Default(std::nullopt);
}
+static bool isTileOpSubkernelHelper(func::FuncOp callee) {
+ return pto::isPTODSLTileOpHelper(callee);
+}
+
static bool isPTODSLSubkernelMemoryOperand(Type type) {
return isa(type);
}
+static bool collectPTODSLTileOpCallOperands(func::CallOp callOp,
+ ArrayAttr operandIndices,
+ SmallVectorImpl &values) {
+ for (Attribute operandIndexAttr : operandIndices) {
+ auto indexAttr = dyn_cast(operandIndexAttr);
+ if (!indexAttr)
+ return false;
+
+ int64_t operandIndex = indexAttr.getInt();
+ if (operandIndex < 0 ||
+ operandIndex >= static_cast(callOp.getNumOperands()))
+ return false;
+
+ Value operand = callOp.getOperand(static_cast(operandIndex));
+ if (!isPTODSLSubkernelMemoryOperand(operand.getType()))
+ continue;
+ values.push_back(operand);
+ }
+ return true;
+}
+
+static bool getPTODSLTileOpCallPhases(func::CallOp callOp, func::FuncOp callee,
+ SmallVectorImpl &phases) {
+ auto primaryDomainAttr =
+ callee->getAttrOfType(kTileOpPrimaryDomainAttr);
+ auto phasesAttr = callee->getAttrOfType(kTileOpPhasesAttr);
+ auto operandEffectsAttr =
+ callee->getAttrOfType(kTileOpOperandEffectsAttr);
+ if (!primaryDomainAttr || !phasesAttr || !operandEffectsAttr)
+ return false;
+
+ if (operandEffectsAttr.size() != callOp.getNumOperands())
+ return false;
+
+ unsigned macroPhaseId = 0;
+ for (Attribute phaseAttr : phasesAttr) {
+ auto dictAttr = dyn_cast(phaseAttr);
+ auto pipeAttr =
+ dictAttr ? dyn_cast_or_null(dictAttr.get("pipe")) : PipeAttr();
+ auto usesAttr =
+ dictAttr ? dyn_cast_or_null(dictAttr.get("operand_uses"))
+ : ArrayAttr();
+ auto defsAttr =
+ dictAttr ? dyn_cast_or_null(dictAttr.get("operand_defs"))
+ : ArrayAttr();
+ auto resultsAttr =
+ dictAttr ? dyn_cast_or_null(dictAttr.get("result_defs"))
+ : ArrayAttr();
+ if (!dictAttr || !pipeAttr || !usesAttr || !defsAttr || !resultsAttr)
+ return false;
+
+ SyncMacroPhase phase;
+ phase.pipe = static_cast(pipeAttr.getPipe());
+ if (!collectPTODSLTileOpCallOperands(callOp, usesAttr, phase.useValues) ||
+ !collectPTODSLTileOpCallOperands(callOp, defsAttr, phase.defValues))
+ return false;
+
+ if (phase.useValues.empty() && phase.defValues.empty())
+ continue;
+
+ phase.phaseId = macroPhaseId++;
+ phases.push_back(std::move(phase));
+ }
+
+ return true;
+}
+
static pto::TCoreType getPTODSLSubkernelHelperCoreType(
pto::PipelineType pipe) {
return pipe == pto::PipelineType::PIPE_M ? pto::TCoreType::CUBE
@@ -749,6 +829,17 @@ void PTOIRTranslator::UpdatePTODSLSubkernelCallInfo(func::CallOp callOp) {
if (!callee)
return;
+ if (isTileOpSubkernelHelper(callee)) {
+ SmallVector phases;
+ if (!getPTODSLTileOpCallPhases(callOp, callee, phases))
+ return;
+ for (const auto &phase : phases) {
+ MakeMacroCompound(callOp, phase.pipe, ValueRange(phase.defValues),
+ ValueRange(phase.useValues), phase.phaseId);
+ }
+ return;
+ }
+
std::optional pipe = getPTODSLSubkernelHelperPipe(callee);
if (!pipe || *pipe == pto::PipelineType::PIPE_UNASSIGNED)
return;
diff --git a/lib/PTO/Transforms/PTOInferTileOpSummaryPass.cpp b/lib/PTO/Transforms/PTOInferTileOpSummaryPass.cpp
new file mode 100644
index 0000000000..2ee5965851
--- /dev/null
+++ b/lib/PTO/Transforms/PTOInferTileOpSummaryPass.cpp
@@ -0,0 +1,441 @@
+// 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/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallSet.h"
+
+namespace mlir {
+namespace pto {
+namespace func = ::mlir::func;
+#define GEN_PASS_DEF_PTOINFERTILEOPSUMMARY
+#include "PTO/Transforms/Passes.h.inc"
+} // namespace pto
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::pto;
+
+namespace {
+
+static constexpr llvm::StringLiteral kTileOpPrimaryDomainAttr =
+ "pto.tileop.primary_domain";
+static constexpr llvm::StringLiteral kTileOpPhasesAttr = "pto.tileop.phases";
+static constexpr llvm::StringLiteral kTileOpOperandEffectsAttr =
+ "pto.tileop.operand_effects";
+
+enum class BoundaryEffect : uint8_t {
+ None,
+ Read,
+ Write,
+ ReadWrite,
+};
+
+struct TileOpPhaseSummary {
+ PIPE pipe = PIPE::PIPE_UNASSIGNED;
+ llvm::SmallSet operandUses;
+ llvm::SmallSet operandDefs;
+};
+
+template
+static void walkTileOpBodyInSourceOrder(Block &block, CallbackT &&callback) {
+ for (Operation &op : block) {
+ if (op.hasTrait())
+ continue;
+ callback(&op);
+ if (op.getNumRegions() > 0 && op.hasTrait())
+ continue;
+ for (Region ®ion : op.getRegions())
+ for (Block &nestedBlock : region)
+ walkTileOpBodyInSourceOrder(nestedBlock, callback);
+ }
+}
+
+static bool isTileOpSubkernelHelper(func::FuncOp funcOp) {
+ return pto::isPTODSLTileOpHelper(funcOp);
+}
+
+static bool isMemoryLikeBoundaryType(Type type) {
+ return isa(type);
+}
+
+static bool isTileOpBodyOp(Operation *op) {
+ if (!op || isa(op))
+ return false;
+ if (op->getName().getDialectNamespace() != PTODialect::getDialectNamespace())
+ return false;
+ return true;
+}
+
+static std::optional getTileOpBodyPipe(Operation *op) {
+ if (!isTileOpBodyOp(op))
+ return std::nullopt;
+
+ if (auto pipeOp = dyn_cast(op)) {
+ PIPE pipe = pipeOp.getPipe();
+ if (pipe != PIPE::PIPE_UNASSIGNED)
+ return pipe;
+ }
+
+ if (isa(op))
+ return std::nullopt;
+
+ StringRef name = op->getName().getStringRef();
+ if (name.starts_with("pto.v"))
+ return PIPE::PIPE_V;
+ if (name.starts_with("pto.mad"))
+ return PIPE::PIPE_M;
+ if (name == "pto.plt_b8" || name == "pto.plt_b16" ||
+ name == "pto.plt_b32" || name == "pto.pltm_b8" ||
+ name == "pto.pltm_b16" || name == "pto.pltm_b32" ||
+ name == "pto.load" || name == "pto.store" || name == "pto.ldg" ||
+ name == "pto.stg")
+ return PIPE::PIPE_S;
+ if (name == "pto.copy_gm_to_ubuf" || name == "pto.mte_gm_ub" ||
+ name == "pto.mte_gm_l1" || name == "pto.mte_gm_l1_frac")
+ return PIPE::PIPE_MTE2;
+ if (name == "pto.mte_ub_gm" || name == "pto.mte_l0c_gm")
+ return PIPE::PIPE_MTE3;
+ if (name == "pto.mte_l1_l0a" || name == "pto.mte_l1_l0b" ||
+ name == "pto.mte_l1_l0a_mx" || name == "pto.mte_l1_l0b_mx")
+ return PIPE::PIPE_MTE1;
+ return std::nullopt;
+}
+
+static bool isMainVectorPipe(PIPE pipe) {
+ return pipe == PIPE::PIPE_V || pipe == PIPE::PIPE_V2;
+}
+
+static bool isMainCubePipe(PIPE pipe) {
+ return pipe == PIPE::PIPE_M;
+}
+
+static FunctionKernelKind getDomainForPipe(PIPE pipe) {
+ return isMainCubePipe(pipe) ? FunctionKernelKind::Cube
+ : FunctionKernelKind::Vector;
+}
+
+static BoundaryEffect joinEffect(BoundaryEffect oldEffect,
+ BoundaryEffect newEffect) {
+ if (oldEffect == BoundaryEffect::None)
+ return newEffect;
+ if (newEffect == BoundaryEffect::None || oldEffect == newEffect)
+ return oldEffect;
+ return BoundaryEffect::ReadWrite;
+}
+
+static StringRef stringifyBoundaryEffect(BoundaryEffect effect) {
+ switch (effect) {
+ case BoundaryEffect::None:
+ return "none";
+ case BoundaryEffect::Read:
+ return "read";
+ case BoundaryEffect::Write:
+ return "write";
+ case BoundaryEffect::ReadWrite:
+ return "readwrite";
+ }
+ llvm_unreachable("unexpected tileop boundary effect");
+}
+
+static LogicalResult buildOperandIndexMap(
+ func::FuncOp funcOp, llvm::DenseMap &operandIndex) {
+ if (funcOp.isDeclaration())
+ return success();
+ if (!funcOp.getBody().hasOneBlock())
+ return funcOp.emitOpError("tileop summary inference requires a single-block helper body");
+
+ Block &entry = funcOp.getBody().front();
+ for (auto [index, arg] : llvm::enumerate(entry.getArguments()))
+ operandIndex.try_emplace(arg, static_cast(index));
+ return success();
+}
+
+static void appendSortedI64Attrs(Builder &builder, llvm::SmallSet set,
+ SmallVectorImpl &attrs) {
+ SmallVector values(set.begin(), set.end());
+ llvm::sort(values);
+ for (int64_t value : values)
+ attrs.push_back(builder.getI64IntegerAttr(value));
+}
+
+static ArrayAttr getSortedI64ArrayAttr(Builder &builder,
+ llvm::SmallSet set) {
+ SmallVector attrs;
+ appendSortedI64Attrs(builder, set, attrs);
+ return builder.getArrayAttr(attrs);
+}
+
+static Value traceBoundaryOperandToHelperArg(Value value) {
+ int loopBound = 256;
+ while (value && loopBound-- > 0) {
+ if (auto arg = dyn_cast(value)) {
+ auto *parentOp = arg.getOwner()->getParentOp();
+ if (auto forOp = dyn_cast_or_null(parentOp)) {
+ if (arg.getArgNumber() > 0 &&
+ forOp.getInitArgs().size() >= arg.getArgNumber()) {
+ value = forOp.getInitArgs()[arg.getArgNumber() - 1];
+ continue;
+ }
+ }
+ return value;
+ }
+
+ Operation *def = value.getDefiningOp();
+ if (!def)
+ return value;
+
+ if (auto subview = dyn_cast(def)) {
+ value = subview.getSource();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ value = cast.getSource();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ value = cast.getSource();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ value = cast.getSource();
+ continue;
+ }
+ if (auto collapse = dyn_cast(def)) {
+ value = collapse.getSrc();
+ continue;
+ }
+ if (auto expand = dyn_cast(def)) {
+ value = expand.getSrc();
+ continue;
+ }
+ if (auto reshape = dyn_cast(def)) {
+ value = reshape.getSource();
+ continue;
+ }
+ if (auto transpose = dyn_cast(def)) {
+ value = transpose.getIn();
+ continue;
+ }
+ if (auto view = dyn_cast(def)) {
+ value = view.getViewSource();
+ continue;
+ }
+ if (auto tileBufAddr = dyn_cast(def)) {
+ value = tileBufAddr.getSrc();
+ continue;
+ }
+ if (auto tensorViewAddr = dyn_cast(def)) {
+ value = tensorViewAddr.getSrc();
+ continue;
+ }
+ if (auto bind = dyn_cast(def)) {
+ value = bind.getSource();
+ continue;
+ }
+ if (auto subview = dyn_cast(def)) {
+ value = subview.getSource();
+ continue;
+ }
+ if (auto bitcast = dyn_cast(def)) {
+ value = bitcast.getSrc();
+ continue;
+ }
+ if (auto reshape = dyn_cast(def)) {
+ value = reshape.getSrc();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ if (cast.getAddrs().empty())
+ return value;
+ value = cast.getAddrs().front();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ value = cast.getInput();
+ continue;
+ }
+ if (auto addPtr = dyn_cast(def)) {
+ value = addPtr.getPtr();
+ continue;
+ }
+ if (auto unrealized = dyn_cast(def)) {
+ if (unrealized.getInputs().empty())
+ return value;
+ if (auto result = dyn_cast(value)) {
+ unsigned resultNumber = result.getResultNumber();
+ if (resultNumber < unrealized.getInputs().size()) {
+ value = unrealized.getInputs()[resultNumber];
+ continue;
+ }
+ }
+ if (unrealized.getInputs().size() == 1) {
+ value = unrealized.getInputs().front();
+ continue;
+ }
+ return value;
+ }
+ if (auto forOp = dyn_cast(def)) {
+ if (auto result = dyn_cast(value)) {
+ unsigned resultNumber = result.getResultNumber();
+ if (resultNumber < forOp.getInitArgs().size()) {
+ value = forOp.getInitArgs()[resultNumber];
+ continue;
+ }
+ }
+ return value;
+ }
+ return value;
+ }
+ return value;
+}
+
+static void recordBoundaryEffects(
+ Operation *op, const llvm::DenseMap &operandIndex,
+ TileOpPhaseSummary &phase, SmallVectorImpl &operandEffects) {
+ auto effectInterface = dyn_cast(op);
+ if (!effectInterface)
+ return;
+
+ SmallVector, 8> effects;
+ effectInterface.getEffects(effects);
+ for (const auto &effect : effects) {
+ Value value = traceBoundaryOperandToHelperArg(effect.getValue());
+ if (!value)
+ continue;
+ auto it = operandIndex.find(value);
+ if (it == operandIndex.end())
+ continue;
+
+ int64_t index = it->second;
+ if (index < 0 ||
+ index >= static_cast(operandEffects.size()) ||
+ !isMemoryLikeBoundaryType(value.getType()))
+ continue;
+
+ BoundaryEffect boundaryEffect = BoundaryEffect::None;
+ if (isa(effect.getEffect())) {
+ phase.operandUses.insert(index);
+ boundaryEffect = BoundaryEffect::Read;
+ } else if (isa(effect.getEffect()) ||
+ isa(effect.getEffect()) ||
+ isa(effect.getEffect())) {
+ phase.operandDefs.insert(index);
+ boundaryEffect = BoundaryEffect::Write;
+ }
+
+ operandEffects[index] = joinEffect(operandEffects[index], boundaryEffect);
+ }
+}
+
+static Attribute buildPhaseAttr(Builder &builder,
+ const TileOpPhaseSummary &phase) {
+ NamedAttrList attrs;
+ attrs.append("pipe", PipeAttr::get(builder.getContext(), phase.pipe));
+ attrs.append("operand_uses", getSortedI64ArrayAttr(builder, phase.operandUses));
+ attrs.append("operand_defs", getSortedI64ArrayAttr(builder, phase.operandDefs));
+ attrs.append("result_defs", builder.getArrayAttr({}));
+ return builder.getDictionaryAttr(attrs);
+}
+
+static LogicalResult inferTileOpSummary(func::FuncOp funcOp) {
+ if (!isTileOpSubkernelHelper(funcOp))
+ return success();
+ if (funcOp.isDeclaration())
+ return success();
+
+ llvm::DenseMap operandIndex;
+ if (failed(buildOperandIndexMap(funcOp, operandIndex)))
+ return failure();
+
+ Builder builder(funcOp.getContext());
+ SmallVector operandEffects(
+ funcOp.getNumArguments(), BoundaryEffect::None);
+ SmallVector phases;
+ std::optional primaryDomain;
+
+ Block &entry = funcOp.getBody().front();
+ LogicalResult walkResult = success();
+ walkTileOpBodyInSourceOrder(entry, [&](Operation *op) {
+ if (failed(walkResult))
+ return;
+
+ std::optional maybePipe = getTileOpBodyPipe(op);
+ if (!maybePipe)
+ return;
+ PIPE pipe = *maybePipe;
+
+ if (isMainVectorPipe(pipe) || isMainCubePipe(pipe)) {
+ FunctionKernelKind domain = getDomainForPipe(pipe);
+ if (primaryDomain && *primaryDomain != domain) {
+ walkResult = op->emitError()
+ << "tileop helper mixes vector and cube primary compute "
+ "pipes; MVP supports exactly one primary domain";
+ return;
+ }
+ primaryDomain = domain;
+ }
+
+ if (phases.empty() || phases.back().pipe != pipe) {
+ TileOpPhaseSummary phase;
+ phase.pipe = pipe;
+ phases.push_back(std::move(phase));
+ }
+ recordBoundaryEffects(op, operandIndex, phases.back(), operandEffects);
+ });
+ if (failed(walkResult))
+ return failure();
+
+ if (!primaryDomain)
+ return success();
+
+ funcOp->setAttr(
+ kTileOpPrimaryDomainAttr,
+ FunctionKernelKindAttr::get(funcOp.getContext(), *primaryDomain));
+
+ SmallVector phaseAttrs;
+ phaseAttrs.reserve(phases.size());
+ for (const TileOpPhaseSummary &phase : phases)
+ phaseAttrs.push_back(buildPhaseAttr(builder, phase));
+ funcOp->setAttr(kTileOpPhasesAttr, builder.getArrayAttr(phaseAttrs));
+
+ SmallVector effectAttrs;
+ effectAttrs.reserve(operandEffects.size());
+ for (BoundaryEffect effect : operandEffects) {
+ if (effect == BoundaryEffect::None)
+ effect = BoundaryEffect::Read;
+ effectAttrs.push_back(builder.getStringAttr(stringifyBoundaryEffect(effect)));
+ }
+ funcOp->setAttr(kTileOpOperandEffectsAttr,
+ builder.getArrayAttr(effectAttrs));
+ return success();
+}
+
+struct PTOInferTileOpSummaryPass
+ : public mlir::pto::impl::PTOInferTileOpSummaryBase<
+ PTOInferTileOpSummaryPass> {
+ void runOnOperation() override {
+ if (failed(inferTileOpSummary(getOperation())))
+ signalPassFailure();
+ }
+};
+
+} // namespace
+
+std::unique_ptr mlir::pto::createPTOInferTileOpSummaryPass() {
+ return std::make_unique();
+}
diff --git a/lib/PTO/Transforms/PTOInstantiateAndInlineOpLib.cpp b/lib/PTO/Transforms/PTOInstantiateAndInlineOpLib.cpp
index c278bd196a..54ecfe02d2 100644
--- a/lib/PTO/Transforms/PTOInstantiateAndInlineOpLib.cpp
+++ b/lib/PTO/Transforms/PTOInstantiateAndInlineOpLib.cpp
@@ -50,7 +50,7 @@ static bool isTilelangInlineProcFunc(func::FuncOp fn) {
}
static bool isPTODSLSubkernelHelperFunc(func::FuncOp fn) {
- return fn->hasAttr("pto.ptodsl.subkernel_helper");
+ return pto::hasPTODSLSubkernelHelperMarker(fn);
}
static bool isTilelangTemplateFunc(func::FuncOp fn) {
diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp
index 10cf9da2bd..a3c54d01ba 100644
--- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp
+++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp
@@ -977,7 +977,7 @@ static bool isTileViewSemantics(StringAttr viewSemantics) {
}
static bool isPTODSLSubkernelHelper(func::FuncOp func) {
- return func->hasAttr("pto.ptodsl.subkernel_helper");
+ return pto::hasPTODSLSubkernelHelperMarker(func);
}
static std::optional>
diff --git a/lib/PTO/Transforms/PTOMaterializeTileOpSectionsPass.cpp b/lib/PTO/Transforms/PTOMaterializeTileOpSectionsPass.cpp
new file mode 100644
index 0000000000..8748040de3
--- /dev/null
+++ b/lib/PTO/Transforms/PTOMaterializeTileOpSectionsPass.cpp
@@ -0,0 +1,349 @@
+// 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/Func/IR/FuncOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+namespace pto {
+namespace func = ::mlir::func;
+#define GEN_PASS_DEF_PTOMATERIALIZETILEOPSECTIONS
+#include "PTO/Transforms/Passes.h.inc"
+} // namespace pto
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::pto;
+
+namespace {
+
+static constexpr llvm::StringLiteral kTileOpPrimaryDomainAttr =
+ "pto.tileop.primary_domain";
+static constexpr llvm::StringLiteral kTileOpPhasesAttr = "pto.tileop.phases";
+
+static bool isTileOpSubkernelHelper(func::FuncOp funcOp) {
+ return pto::isPTODSLTileOpHelper(funcOp);
+}
+
+static bool isTileOpBodyOp(Operation *op) {
+ if (!op || isa(op))
+ return false;
+ if (op->getName().getDialectNamespace() != PTODialect::getDialectNamespace())
+ return false;
+ return true;
+}
+
+static std::optional getTileOpBodyPipe(Operation *op) {
+ if (!isTileOpBodyOp(op))
+ return std::nullopt;
+
+ if (auto pipeOp = dyn_cast(op)) {
+ PIPE pipe = pipeOp.getPipe();
+ if (pipe != PIPE::PIPE_UNASSIGNED)
+ return pipe;
+ }
+
+ if (isa(op))
+ return std::nullopt;
+
+ StringRef name = op->getName().getStringRef();
+ if (name.starts_with("pto.v"))
+ return PIPE::PIPE_V;
+ if (name.starts_with("pto.mad"))
+ return PIPE::PIPE_M;
+ if (name == "pto.plt_b8" || name == "pto.plt_b16" ||
+ name == "pto.plt_b32" || name == "pto.pltm_b8" ||
+ name == "pto.pltm_b16" || name == "pto.pltm_b32" ||
+ name == "pto.load" || name == "pto.store" || name == "pto.ldg" ||
+ name == "pto.stg")
+ return PIPE::PIPE_S;
+ if (name == "pto.copy_gm_to_ubuf" || name == "pto.mte_gm_ub" ||
+ name == "pto.mte_gm_l1" || name == "pto.mte_gm_l1_frac")
+ return PIPE::PIPE_MTE2;
+ if (name == "pto.mte_ub_gm" || name == "pto.mte_l0c_gm")
+ return PIPE::PIPE_MTE3;
+ if (name == "pto.mte_l1_l0a" || name == "pto.mte_l1_l0b" ||
+ name == "pto.mte_l1_l0a_mx" || name == "pto.mte_l1_l0b_mx")
+ return PIPE::PIPE_MTE1;
+ return std::nullopt;
+}
+
+static bool hasExistingSection(func::FuncOp funcOp) {
+ bool found = false;
+ funcOp.walk([&](Operation *op) {
+ if (isa(op)) {
+ found = true;
+ return WalkResult::interrupt();
+ }
+ return WalkResult::advance();
+ });
+ return found;
+}
+
+static bool isPrimaryVectorPipe(PIPE pipe) {
+ return pipe == PIPE::PIPE_V || pipe == PIPE::PIPE_V2;
+}
+
+static bool isPrimaryCubePipe(PIPE pipe) {
+ return pipe == PIPE::PIPE_M;
+}
+
+static bool isPrimaryPipeForKind(PIPE pipe, FunctionKernelKind kind) {
+ switch (kind) {
+ case FunctionKernelKind::Vector:
+ return isPrimaryVectorPipe(pipe);
+ case FunctionKernelKind::Cube:
+ return isPrimaryCubePipe(pipe);
+ }
+ llvm_unreachable("unexpected kernel kind");
+}
+
+static FailureOr getPrimaryDomain(func::FuncOp funcOp) {
+ auto primaryAttr =
+ funcOp->getAttrOfType(kTileOpPrimaryDomainAttr);
+ if (!primaryAttr)
+ return funcOp.emitOpError("requires ")
+ << kTileOpPrimaryDomainAttr << " before tileop section materialization";
+ return primaryAttr.getKernelKind();
+}
+
+static FailureOr findFirstPrimaryPhase(func::FuncOp funcOp,
+ ArrayAttr phases,
+ FunctionKernelKind kind) {
+ for (auto [index, attr] : llvm::enumerate(phases)) {
+ auto dict = dyn_cast(attr);
+ auto pipeAttr = dict ? dyn_cast_or_null(dict.get("pipe")) : PipeAttr();
+ if (!dict || !pipeAttr) {
+ return funcOp.emitOpError("expects ")
+ << kTileOpPhasesAttr << " entries to carry a pipe attr";
+ }
+ if (isPrimaryPipeForKind(pipeAttr.getPipe(), kind))
+ return index;
+ }
+ return funcOp.emitOpError("requires at least one primary compute phase in ")
+ << kTileOpPhasesAttr;
+}
+
+static FailureOr findLastPrimaryPhase(func::FuncOp funcOp,
+ ArrayAttr phases,
+ FunctionKernelKind kind) {
+ for (int index = static_cast(phases.size()) - 1; index >= 0; --index) {
+ auto dict = dyn_cast(phases[index]);
+ auto pipeAttr = dict ? dyn_cast_or_null(dict.get("pipe")) : PipeAttr();
+ if (!dict || !pipeAttr) {
+ return funcOp.emitOpError("expects ")
+ << kTileOpPhasesAttr << " entries to carry a pipe attr";
+ }
+ if (isPrimaryPipeForKind(pipeAttr.getPipe(), kind))
+ return static_cast(index);
+ }
+ return funcOp.emitOpError("requires at least one primary compute phase in ")
+ << kTileOpPhasesAttr;
+}
+
+static SmallVector collectTileOpBodyOps(Block &block) {
+ SmallVector ops;
+ for (Operation &op : block.without_terminator()) {
+ if (!getTileOpBodyPipe(&op))
+ continue;
+ ops.push_back(&op);
+ }
+ return ops;
+}
+
+static FailureOr>>
+findPrimaryOpRange(func::FuncOp funcOp, ArrayRef ops,
+ FunctionKernelKind kind) {
+ int first = -1;
+ int last = -1;
+ for (auto [index, op] : llvm::enumerate(ops)) {
+ std::optional maybePipe = getTileOpBodyPipe(op);
+ if (!maybePipe)
+ continue;
+ PIPE pipe = *maybePipe;
+ if (!isPrimaryPipeForKind(pipe, kind))
+ continue;
+ if (first < 0)
+ first = static_cast(index);
+ last = static_cast(index);
+ }
+
+ if (first < 0 || last < 0)
+ return std::optional>();
+
+ for (int index = first; index <= last; ++index) {
+ std::optional maybePipe = getTileOpBodyPipe(ops[index]);
+ if (!maybePipe)
+ continue;
+ PIPE pipe = *maybePipe;
+ if (!isPrimaryPipeForKind(pipe, kind)) {
+ return ops[index]->emitError()
+ << "tileop primary compute span is not contiguous; MVP materializer "
+ "supports one contiguous primary-domain span";
+ }
+ }
+
+ return std::optional>(
+ std::make_pair(static_cast(first),
+ static_cast(last)));
+}
+
+static bool isVectorScopeLocalType(Type type) {
+ return isa(type);
+}
+
+static bool hasVectorScopeLocalResult(Operation *op) {
+ return llvm::any_of(op->getResultTypes(), isVectorScopeLocalType);
+}
+
+static Operation *expandStartForLocalProducers(Block &block, Operation *firstOp,
+ Operation *lastOp) {
+ Operation *expandedFirst = firstOp;
+ bool changed = true;
+ while (changed) {
+ changed = false;
+ auto expandedFirstIt = Block::iterator(expandedFirst);
+ auto afterLastIt = std::next(Block::iterator(lastOp));
+ for (auto it = expandedFirstIt; it != afterLastIt; ++it) {
+ Operation &op = *it;
+ for (Value operand : op.getOperands()) {
+ Operation *def = operand.getDefiningOp();
+ if (!def || def->getBlock() != &block || !def->isBeforeInBlock(expandedFirst))
+ continue;
+ if (!hasVectorScopeLocalResult(def))
+ continue;
+ expandedFirst = def;
+ changed = true;
+ break;
+ }
+ if (changed)
+ break;
+ }
+ }
+ return expandedFirst;
+}
+
+template
+static void wrapOperationRange(Block &block, Operation *firstOp,
+ Operation *lastOp) {
+ OpBuilder builder(firstOp);
+ auto sectionOp = builder.create(firstOp->getLoc());
+ if (!sectionOp.getBody().empty())
+ sectionOp.getBody().dropAllReferences();
+ while (!sectionOp.getBody().empty())
+ delete §ionOp.getBody().front();
+ sectionOp.getBody().push_back(new Block());
+ Block §ionBlock = sectionOp.getBody().front();
+
+ auto firstIt = Block::iterator(firstOp);
+ auto afterLastIt = std::next(Block::iterator(lastOp));
+ sectionBlock.getOperations().splice(sectionBlock.end(),
+ block.getOperations(), firstIt,
+ afterLastIt);
+}
+
+static LogicalResult materializePrimarySectionsInBlock(func::FuncOp funcOp,
+ Block &block,
+ FunctionKernelKind kind,
+ bool &materializedAny) {
+ SmallVector topLevelOps;
+ for (Operation &op : block.without_terminator())
+ topLevelOps.push_back(&op);
+
+ SmallVector bodyOps = collectTileOpBodyOps(block);
+ if (!bodyOps.empty()) {
+ FailureOr>> primaryRange =
+ findPrimaryOpRange(funcOp, bodyOps, kind);
+ if (failed(primaryRange))
+ return failure();
+
+ if (*primaryRange) {
+ Operation *firstPrimaryOp = bodyOps[(*primaryRange)->first];
+ Operation *lastPrimaryOp = bodyOps[(*primaryRange)->second];
+ firstPrimaryOp =
+ expandStartForLocalProducers(block, firstPrimaryOp, lastPrimaryOp);
+ switch (kind) {
+ case FunctionKernelKind::Vector:
+ wrapOperationRange(block, firstPrimaryOp,
+ lastPrimaryOp);
+ break;
+ case FunctionKernelKind::Cube:
+ wrapOperationRange(block, firstPrimaryOp, lastPrimaryOp);
+ break;
+ }
+ materializedAny = true;
+ }
+ }
+
+ for (Operation *op : topLevelOps) {
+ if (!op || getTileOpBodyPipe(op))
+ continue;
+ for (Region ®ion : op->getRegions())
+ for (Block &nestedBlock : region)
+ if (failed(materializePrimarySectionsInBlock(funcOp, nestedBlock, kind,
+ materializedAny)))
+ return failure();
+ }
+ return success();
+}
+
+static LogicalResult materializeTileOpSection(func::FuncOp funcOp) {
+ if (!isTileOpSubkernelHelper(funcOp) || funcOp.isDeclaration())
+ return success();
+
+ if (hasExistingSection(funcOp))
+ return success();
+
+ auto phases = funcOp->getAttrOfType(kTileOpPhasesAttr);
+ if (!phases || phases.empty())
+ return success();
+
+ FailureOr primaryDomain = getPrimaryDomain(funcOp);
+ if (failed(primaryDomain))
+ return failure();
+
+ FailureOr firstPhase =
+ findFirstPrimaryPhase(funcOp, phases, *primaryDomain);
+ if (failed(firstPhase))
+ return failure();
+ FailureOr lastPhase =
+ findLastPrimaryPhase(funcOp, phases, *primaryDomain);
+ if (failed(lastPhase))
+ return failure();
+ (void)firstPhase;
+ (void)lastPhase;
+
+ bool materializedAny = false;
+ if (failed(materializePrimarySectionsInBlock(funcOp, funcOp.getBody().front(),
+ *primaryDomain, materializedAny)))
+ return failure();
+ if (!materializedAny) {
+ return funcOp.emitOpError(
+ "requires at least one primary compute op in helper body");
+ }
+ return success();
+}
+
+struct PTOMaterializeTileOpSectionsPass
+ : public mlir::pto::impl::PTOMaterializeTileOpSectionsBase<
+ PTOMaterializeTileOpSectionsPass> {
+ void runOnOperation() override {
+ if (failed(materializeTileOpSection(getOperation())))
+ signalPassFailure();
+ }
+};
+
+} // namespace
+
+std::unique_ptr mlir::pto::createPTOMaterializeTileOpSectionsPass() {
+ return std::make_unique();
+}
diff --git a/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp b/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp
index 55b78f2863..8348b8aa8a 100644
--- a/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp
+++ b/lib/PTO/Transforms/PTONormalizeUncoveredTileSections.cpp
@@ -110,6 +110,10 @@ static bool hasExplicitFunctionKernelKind(func::FuncOp funcOp) {
funcOp->hasAttrOfType(FunctionKernelKindAttr::name);
}
+static bool isTileOpSubkernelHelper(func::FuncOp funcOp) {
+ return pto::isPTODSLTileOpHelper(funcOp);
+}
+
static bool isInsideKernelKindModule(func::FuncOp funcOp) {
if (!funcOp)
return false;
@@ -118,7 +122,9 @@ static bool isInsideKernelKindModule(func::FuncOp funcOp) {
}
static bool hasKnownKernelKindContext(func::FuncOp funcOp) {
- return isInsideKernelKindModule(funcOp) || hasExplicitFunctionKernelKind(funcOp);
+ return isInsideKernelKindModule(funcOp) ||
+ hasExplicitFunctionKernelKind(funcOp) ||
+ isTileOpSubkernelHelper(funcOp);
}
static std::optional getBufferAddressSpace(Type type) {
diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp
index 4a4f865885..36e7bd089b 100644
--- a/lib/PTO/Transforms/PTOPlanMemory.cpp
+++ b/lib/PTO/Transforms/PTOPlanMemory.cpp
@@ -473,6 +473,10 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) {
OpKillHandle(curOpInfo, live, op->getBlock());
} else if (auto storeOp = dyn_cast(op)) {
UpdateStoreOpInfo(curOpInfo, storeOp.getMemRef(), live);
+ } else if (isa(op)) {
+ UpdateOpGenInfo(curOpInfo, llvm::to_vector(op->getOperands()));
+ OpKillHandle(curOpInfo, live, op->getBlock());
} else if (auto ptoDpsOp = dyn_cast(op)) {
// PTO ops with destination (tile_buf, partition_view, etc.); no
// tensor/memref-only verification.
diff --git a/lib/PTO/Transforms/PTOVerifySubkernelPipeContractPass.cpp b/lib/PTO/Transforms/PTOVerifySubkernelPipeContractPass.cpp
index 8635497f64..f728b94164 100644
--- a/lib/PTO/Transforms/PTOVerifySubkernelPipeContractPass.cpp
+++ b/lib/PTO/Transforms/PTOVerifySubkernelPipeContractPass.cpp
@@ -26,14 +26,8 @@ using namespace mlir::pto;
namespace {
-static constexpr llvm::StringLiteral kPTODSLSubkernelHelperAttr =
- "pto.ptodsl.subkernel_helper";
-
static StringRef getSubkernelRole(func::FuncOp funcOp) {
- if (auto roleAttr =
- funcOp->getAttrOfType(kPTODSLSubkernelHelperAttr))
- return roleAttr.getValue();
- return {};
+ return pto::getPTODSLSubkernelHelperRole(funcOp);
}
static std::optional getExpectedPipeForRole(StringRef role) {
diff --git a/lib/PTO/Transforms/PTOVerifyTileOpContractPass.cpp b/lib/PTO/Transforms/PTOVerifyTileOpContractPass.cpp
new file mode 100644
index 0000000000..b3d4c3dbb0
--- /dev/null
+++ b/lib/PTO/Transforms/PTOVerifyTileOpContractPass.cpp
@@ -0,0 +1,632 @@
+// 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/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallSet.h"
+
+namespace mlir {
+namespace pto {
+namespace func = ::mlir::func;
+#define GEN_PASS_DEF_PTOVERIFYTILEOPCONTRACT
+#include "PTO/Transforms/Passes.h.inc"
+} // namespace pto
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::pto;
+
+namespace {
+
+static constexpr llvm::StringLiteral kTileOpPrimaryDomainAttr =
+ "pto.tileop.primary_domain";
+static constexpr llvm::StringLiteral kTileOpPhasesAttr = "pto.tileop.phases";
+static constexpr llvm::StringLiteral kTileOpOperandEffectsAttr =
+ "pto.tileop.operand_effects";
+
+enum class BoundaryEffect : uint8_t {
+ None,
+ Read,
+ Write,
+ ReadWrite,
+};
+
+struct TileOpPhaseSummary {
+ PIPE pipe = PIPE::PIPE_UNASSIGNED;
+ llvm::SmallSet operandUses;
+ llvm::SmallSet operandDefs;
+};
+
+template
+static void walkTileOpBodyInSourceOrder(Block &block, CallbackT &&callback) {
+ for (Operation &op : block) {
+ if (op.hasTrait())
+ continue;
+ callback(&op);
+ if (op.getNumRegions() > 0 && op.hasTrait())
+ continue;
+ for (Region ®ion : op.getRegions())
+ for (Block &nestedBlock : region)
+ walkTileOpBodyInSourceOrder(nestedBlock, callback);
+ }
+}
+
+static bool isTileOpSubkernelHelper(func::FuncOp funcOp) {
+ return pto::isPTODSLTileOpHelper(funcOp);
+}
+
+static func::CallOp getTransparentWrapperCall(func::FuncOp funcOp) {
+ if (!funcOp || funcOp.isDeclaration() || !funcOp.getBody().hasOneBlock())
+ return nullptr;
+
+ Block &entryBlock = funcOp.getBody().front();
+ func::CallOp callOp;
+ func::ReturnOp returnOp;
+ for (Operation &op : entryBlock.getOperations()) {
+ if (auto ret = dyn_cast(op)) {
+ returnOp = ret;
+ continue;
+ }
+ if (callOp)
+ return nullptr;
+ callOp = dyn_cast(op);
+ if (!callOp)
+ return nullptr;
+ }
+
+ if (!callOp || !returnOp)
+ return nullptr;
+ if (returnOp.getNumOperands() != callOp.getNumResults())
+ return nullptr;
+ for (auto [returned, forwarded] :
+ llvm::zip(returnOp.getOperands(), callOp.getResults())) {
+ if (returned != forwarded)
+ return nullptr;
+ }
+ return callOp;
+}
+
+static bool isMemoryLikeBoundaryType(Type type) {
+ return isa(type);
+}
+
+static bool isTileOpHelperBoundaryType(Type type) {
+ return isa(type);
+}
+
+static bool isTileOpScalarType(Type type) {
+ return isa(type);
+}
+
+static bool isMainVectorPipe(PIPE pipe) {
+ return pipe == PIPE::PIPE_V || pipe == PIPE::PIPE_V2;
+}
+
+static bool isMainCubePipe(PIPE pipe) {
+ return pipe == PIPE::PIPE_M;
+}
+
+static bool isTileOpBodyOp(Operation *op) {
+ if (!op || isa(op))
+ return false;
+ if (op->getName().getDialectNamespace() != PTODialect::getDialectNamespace())
+ return false;
+ return true;
+}
+
+static std::optional getTileOpBodyPipe(Operation *op) {
+ if (!isTileOpBodyOp(op))
+ return std::nullopt;
+
+ if (auto pipeOp = dyn_cast(op)) {
+ PIPE pipe = pipeOp.getPipe();
+ if (pipe != PIPE::PIPE_UNASSIGNED)
+ return pipe;
+ }
+
+ if (isa(op))
+ return std::nullopt;
+
+ StringRef name = op->getName().getStringRef();
+ if (name.starts_with("pto.v"))
+ return PIPE::PIPE_V;
+ if (name.starts_with("pto.mad"))
+ return PIPE::PIPE_M;
+ if (name == "pto.plt_b8" || name == "pto.plt_b16" ||
+ name == "pto.plt_b32" || name == "pto.pltm_b8" ||
+ name == "pto.pltm_b16" || name == "pto.pltm_b32" ||
+ name == "pto.load" || name == "pto.store" || name == "pto.ldg" ||
+ name == "pto.stg")
+ return PIPE::PIPE_S;
+ if (name == "pto.copy_gm_to_ubuf" || name == "pto.mte_gm_ub" ||
+ name == "pto.mte_gm_l1" || name == "pto.mte_gm_l1_frac")
+ return PIPE::PIPE_MTE2;
+ if (name == "pto.mte_ub_gm" || name == "pto.mte_l0c_gm")
+ return PIPE::PIPE_MTE3;
+ if (name == "pto.mte_l1_l0a" || name == "pto.mte_l1_l0b" ||
+ name == "pto.mte_l1_l0a_mx" || name == "pto.mte_l1_l0b_mx")
+ return PIPE::PIPE_MTE1;
+ return std::nullopt;
+}
+
+static FunctionKernelKind getDomainForPipe(PIPE pipe) {
+ return isMainCubePipe(pipe) ? FunctionKernelKind::Cube
+ : FunctionKernelKind::Vector;
+}
+
+static StringRef stringifyBoundaryEffect(BoundaryEffect effect) {
+ switch (effect) {
+ case BoundaryEffect::None:
+ return "none";
+ case BoundaryEffect::Read:
+ return "read";
+ case BoundaryEffect::Write:
+ return "write";
+ case BoundaryEffect::ReadWrite:
+ return "readwrite";
+ }
+ llvm_unreachable("unexpected tileop boundary effect");
+}
+
+static BoundaryEffect joinEffect(BoundaryEffect oldEffect,
+ BoundaryEffect newEffect) {
+ if (oldEffect == BoundaryEffect::None)
+ return newEffect;
+ if (newEffect == BoundaryEffect::None || oldEffect == newEffect)
+ return oldEffect;
+ return BoundaryEffect::ReadWrite;
+}
+
+static bool isSIMTOnlyPTOOp(Operation *op) {
+ return isa(op);
+}
+
+static LogicalResult buildOperandIndexMap(
+ func::FuncOp funcOp, llvm::DenseMap &operandIndex) {
+ if (funcOp.isDeclaration())
+ return success();
+ if (!funcOp.getBody().hasOneBlock())
+ return funcOp.emitOpError(
+ "tileop contract verification requires a single-block helper body");
+
+ for (auto [index, arg] :
+ llvm::enumerate(funcOp.getBody().front().getArguments()))
+ operandIndex.try_emplace(arg, static_cast(index));
+ return success();
+}
+
+static Value traceBoundaryOperandToHelperArg(Value value) {
+ int loopBound = 256;
+ while (value && loopBound-- > 0) {
+ if (auto arg = dyn_cast(value)) {
+ auto *parentOp = arg.getOwner()->getParentOp();
+ if (auto forOp = dyn_cast_or_null(parentOp)) {
+ if (arg.getArgNumber() > 0 &&
+ forOp.getInitArgs().size() >= arg.getArgNumber()) {
+ value = forOp.getInitArgs()[arg.getArgNumber() - 1];
+ continue;
+ }
+ }
+ return value;
+ }
+
+ Operation *def = value.getDefiningOp();
+ if (!def)
+ return value;
+
+ if (auto subview = dyn_cast(def)) {
+ value = subview.getSource();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ value = cast.getSource();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ value = cast.getSource();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ value = cast.getSource();
+ continue;
+ }
+ if (auto collapse = dyn_cast(def)) {
+ value = collapse.getSrc();
+ continue;
+ }
+ if (auto expand = dyn_cast(def)) {
+ value = expand.getSrc();
+ continue;
+ }
+ if (auto reshape = dyn_cast(def)) {
+ value = reshape.getSource();
+ continue;
+ }
+ if (auto transpose = dyn_cast(def)) {
+ value = transpose.getIn();
+ continue;
+ }
+ if (auto view = dyn_cast(def)) {
+ value = view.getViewSource();
+ continue;
+ }
+ if (auto tileBufAddr = dyn_cast(def)) {
+ value = tileBufAddr.getSrc();
+ continue;
+ }
+ if (auto tensorViewAddr = dyn_cast(def)) {
+ value = tensorViewAddr.getSrc();
+ continue;
+ }
+ if (auto bind = dyn_cast(def)) {
+ value = bind.getSource();
+ continue;
+ }
+ if (auto subview = dyn_cast(def)) {
+ value = subview.getSource();
+ continue;
+ }
+ if (auto bitcast = dyn_cast(def)) {
+ value = bitcast.getSrc();
+ continue;
+ }
+ if (auto reshape = dyn_cast(def)) {
+ value = reshape.getSrc();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ if (cast.getAddrs().empty())
+ return value;
+ value = cast.getAddrs().front();
+ continue;
+ }
+ if (auto cast = dyn_cast(def)) {
+ value = cast.getInput();
+ continue;
+ }
+ if (auto addPtr = dyn_cast(def)) {
+ value = addPtr.getPtr();
+ continue;
+ }
+ if (auto unrealized = dyn_cast(def)) {
+ if (unrealized.getInputs().empty())
+ return value;
+ if (auto result = dyn_cast(value)) {
+ unsigned resultNumber = result.getResultNumber();
+ if (resultNumber < unrealized.getInputs().size()) {
+ value = unrealized.getInputs()[resultNumber];
+ continue;
+ }
+ }
+ if (unrealized.getInputs().size() == 1) {
+ value = unrealized.getInputs().front();
+ continue;
+ }
+ return value;
+ }
+ if (auto forOp = dyn_cast(def)) {
+ if (auto result = dyn_cast(value)) {
+ unsigned resultNumber = result.getResultNumber();
+ if (resultNumber < forOp.getInitArgs().size()) {
+ value = forOp.getInitArgs()[resultNumber];
+ continue;
+ }
+ }
+ return value;
+ }
+ return value;
+ }
+ return value;
+}
+
+static void recordBoundaryEffects(
+ Operation *op, const llvm::DenseMap &operandIndex,
+ TileOpPhaseSummary &phase, SmallVectorImpl &operandEffects) {
+ auto effectInterface = dyn_cast(op);
+ if (!effectInterface)
+ return;
+
+ SmallVector, 8> effects;
+ effectInterface.getEffects(effects);
+ for (const auto &effect : effects) {
+ Value value = traceBoundaryOperandToHelperArg(effect.getValue());
+ if (!value)
+ continue;
+
+ auto it = operandIndex.find(value);
+ if (it == operandIndex.end())
+ continue;
+
+ int64_t index = it->second;
+ if (index < 0 || index >= static_cast(operandEffects.size()) ||
+ !isMemoryLikeBoundaryType(value.getType()))
+ continue;
+
+ BoundaryEffect boundaryEffect = BoundaryEffect::None;
+ if (isa(effect.getEffect())) {
+ phase.operandUses.insert(index);
+ boundaryEffect = BoundaryEffect::Read;
+ } else if (isa(effect.getEffect()) ||
+ isa(effect.getEffect()) ||
+ isa(effect.getEffect())) {
+ phase.operandDefs.insert(index);
+ boundaryEffect = BoundaryEffect::Write;
+ }
+
+ operandEffects[index] = joinEffect(operandEffects[index], boundaryEffect);
+ }
+}
+
+static LogicalResult verifyScalarResults(func::FuncOp funcOp) {
+ for (Type resultType : funcOp.getResultTypes()) {
+ if (isTileOpScalarType(resultType))
+ continue;
+ return funcOp.emitOpError()
+ << "tileop helper results are limited to PTO scalar values in the "
+ "MVP, but found result type "
+ << resultType;
+ }
+ return success();
+}
+
+static LogicalResult verifyArgumentBoundaryTypes(func::FuncOp funcOp) {
+ for (Type argType : funcOp.getArgumentTypes()) {
+ if (isTileOpHelperBoundaryType(argType) || isTileOpScalarType(argType))
+ continue;
+ return funcOp.emitOpError()
+ << "tileop helper arguments must be Tile/TensorView/"
+ "PartitionTensorView or PTO scalar values, but found "
+ << argType;
+ }
+ return success();
+}
+
+static LogicalResult verifySummaryAttrs(func::FuncOp funcOp,
+ std::optional inferredDomain,
+ ArrayRef inferredPhases,
+ ArrayRef inferredEffects) {
+ if (!inferredDomain)
+ return funcOp.emitOpError(
+ "requires at least one vector or cube primary compute op; helpers with "
+ "only MTE/scalar/sync phases are rejected");
+
+ auto primaryAttr =
+ funcOp->getAttrOfType(kTileOpPrimaryDomainAttr);
+ if (!primaryAttr)
+ return funcOp.emitOpError("requires ")
+ << kTileOpPrimaryDomainAttr << " before tileop contract verification";
+
+ if (primaryAttr.getKernelKind() != *inferredDomain) {
+ return funcOp.emitOpError("has stale ")
+ << kTileOpPrimaryDomainAttr << ": inferred primary domain is #pto.kernel_kind<"
+ << (*inferredDomain == FunctionKernelKind::Cube ? "cube" : "vector")
+ << ">";
+ }
+
+ auto phasesAttr = funcOp->getAttrOfType(kTileOpPhasesAttr);
+ if (!phasesAttr)
+ return funcOp.emitOpError("requires ") << kTileOpPhasesAttr;
+ if (phasesAttr.size() != inferredPhases.size()) {
+ return funcOp.emitOpError("has stale ")
+ << kTileOpPhasesAttr << ": inferred " << inferredPhases.size()
+ << " phase(s), but attribute stores " << phasesAttr.size();
+ }
+
+ auto effectsAttr = funcOp->getAttrOfType(kTileOpOperandEffectsAttr);
+ if (!effectsAttr)
+ return funcOp.emitOpError("requires ") << kTileOpOperandEffectsAttr;
+ if (effectsAttr.size() != inferredEffects.size()) {
+ return funcOp.emitOpError("has stale ")
+ << kTileOpOperandEffectsAttr << ": inferred "
+ << inferredEffects.size() << " operand effect(s), but attribute stores "
+ << effectsAttr.size();
+ }
+
+ for (unsigned index = 0; index < effectsAttr.size(); ++index) {
+ auto strAttr = dyn_cast(effectsAttr[index]);
+ if (!strAttr)
+ return funcOp.emitOpError("expects ")
+ << kTileOpOperandEffectsAttr
+ << " entries to be string attributes";
+
+ BoundaryEffect expected = inferredEffects[index];
+ if (expected == BoundaryEffect::None)
+ expected = BoundaryEffect::Read;
+ if (strAttr.getValue() != stringifyBoundaryEffect(expected)) {
+ return funcOp.emitOpError("has stale ")
+ << kTileOpOperandEffectsAttr << " at operand #" << index
+ << ": inferred \"" << stringifyBoundaryEffect(expected)
+ << "\", but attribute stores \"" << strAttr.getValue() << "\"";
+ }
+ }
+
+ for (unsigned index = 0; index < phasesAttr.size(); ++index) {
+ auto dict = dyn_cast(phasesAttr[index]);
+ if (!dict)
+ return funcOp.emitOpError("expects ")
+ << kTileOpPhasesAttr << " entries to be dictionary attributes";
+
+ auto pipeAttr = dyn_cast_or_null(dict.get("pipe"));
+ auto usesAttr = dyn_cast_or_null(dict.get("operand_uses"));
+ auto defsAttr = dyn_cast_or_null(dict.get("operand_defs"));
+ auto resultsAttr = dyn_cast_or_null(dict.get("result_defs"));
+ if (!pipeAttr || !usesAttr || !defsAttr || !resultsAttr) {
+ return funcOp.emitOpError("expects ")
+ << kTileOpPhasesAttr
+ << " entries to carry pipe/operand_uses/operand_defs/result_defs";
+ }
+
+ const TileOpPhaseSummary &expected = inferredPhases[index];
+ if (pipeAttr.getPipe() != expected.pipe) {
+ return funcOp.emitOpError("has stale ")
+ << kTileOpPhasesAttr << " at phase #" << index
+ << ": inferred pipe " << stringifyPIPE(expected.pipe)
+ << ", but attribute stores " << stringifyPIPE(pipeAttr.getPipe());
+ }
+
+ auto verifyIndexSet = [&](ArrayAttr values, llvm::SmallSet set,
+ StringRef fieldName) -> LogicalResult {
+ if (values.size() != set.size()) {
+ return funcOp.emitOpError("has stale ")
+ << kTileOpPhasesAttr << " at phase #" << index << " field '"
+ << fieldName << "'";
+ }
+ llvm::SmallSet actual;
+ for (Attribute valueAttr : values) {
+ auto intAttr = dyn_cast(valueAttr);
+ if (!intAttr)
+ return funcOp.emitOpError("expects ") << fieldName
+ << " indices to be integers";
+ actual.insert(intAttr.getInt());
+ }
+ if (actual != set) {
+ return funcOp.emitOpError("has stale ")
+ << kTileOpPhasesAttr << " at phase #" << index << " field '"
+ << fieldName << "'";
+ }
+ return success();
+ };
+
+ if (failed(verifyIndexSet(usesAttr, expected.operandUses, "operand_uses")) ||
+ failed(verifyIndexSet(defsAttr, expected.operandDefs, "operand_defs")))
+ return failure();
+ if (!resultsAttr.empty()) {
+ return funcOp.emitOpError("expects ")
+ << kTileOpPhasesAttr
+ << " result_defs to remain empty in the current MVP";
+ }
+ }
+
+ return success();
+}
+
+static LogicalResult verifyTileOpHelper(func::FuncOp funcOp) {
+ if (!isTileOpSubkernelHelper(funcOp) || funcOp.isDeclaration())
+ return success();
+
+ if (failed(verifyScalarResults(funcOp)) ||
+ failed(verifyArgumentBoundaryTypes(funcOp)))
+ return failure();
+
+ if (auto wrapperCall = getTransparentWrapperCall(funcOp)) {
+ auto callee = SymbolTable::lookupNearestSymbolFrom(
+ funcOp, wrapperCall.getCalleeAttr());
+ if (callee && callee != funcOp && isTileOpSubkernelHelper(callee))
+ return success();
+ }
+
+ llvm::DenseMap operandIndex;
+ if (failed(buildOperandIndexMap(funcOp, operandIndex)))
+ return failure();
+
+ SmallVector operandEffects(funcOp.getNumArguments(),
+ BoundaryEffect::None);
+ SmallVector phases;
+ std::optional primaryDomain;
+
+ Block &entry = funcOp.getBody().front();
+ LogicalResult walkResult = success();
+ walkTileOpBodyInSourceOrder(entry, [&](Operation *op) {
+ if (failed(walkResult))
+ return;
+
+ if (isa(op)) {
+ walkResult = op->emitError("is not allowed inside a tileop helper; "
+ "tileop helpers must not allocate "
+ "helper-local tile or reserved-buffer state");
+ return;
+ }
+
+ if (isSIMTOnlyPTOOp(op)) {
+ walkResult =
+ op->emitError("is SIMT-only and cannot appear inside a tileop helper");
+ return;
+ }
+
+ if (auto callOp = dyn_cast(op)) {
+ if (callOp->getParentOfType() != funcOp)
+ return;
+ auto module = funcOp->getParentOfType();
+ auto callee = module ? module.lookupSymbol(callOp.getCallee())
+ : func::FuncOp();
+ if (callee && isTileOpSubkernelHelper(callee)) {
+ InFlightDiagnostic diag =
+ callOp.emitOpError("cannot call tileop helper @");
+ diag << callee.getSymName()
+ << " from another tileop helper; nested tileop calls are "
+ "rejected";
+ walkResult = failure();
+ return;
+ }
+ return;
+ }
+
+ if (op->getName().getDialectNamespace() != PTODialect::getDialectNamespace())
+ return;
+
+ std::optional maybePipe = getTileOpBodyPipe(op);
+ if (!maybePipe)
+ return;
+ PIPE pipe = *maybePipe;
+
+ if (isMainVectorPipe(pipe) || isMainCubePipe(pipe)) {
+ FunctionKernelKind domain = getDomainForPipe(pipe);
+ if (primaryDomain && *primaryDomain != domain) {
+ walkResult = op->emitError(
+ "tileop helper mixes vector and cube primary compute pipes; MVP "
+ "supports exactly one primary domain");
+ return;
+ }
+ primaryDomain = domain;
+ }
+
+ if (phases.empty() || phases.back().pipe != pipe) {
+ TileOpPhaseSummary phase;
+ phase.pipe = pipe;
+ phases.push_back(std::move(phase));
+ }
+ recordBoundaryEffects(op, operandIndex, phases.back(), operandEffects);
+ });
+
+ if (failed(walkResult))
+ return failure();
+
+ return verifySummaryAttrs(funcOp, primaryDomain, phases, operandEffects);
+}
+
+struct PTOVerifyTileOpContractPass
+ : public mlir::pto::impl::PTOVerifyTileOpContractBase<
+ PTOVerifyTileOpContractPass> {
+ void runOnOperation() override {
+ if (failed(verifyTileOpHelper(getOperation())))
+ signalPassFailure();
+ }
+};
+
+} // namespace
+
+std::unique_ptr mlir::pto::createPTOVerifyTileOpContractPass() {
+ return std::make_unique();
+}
diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp
index e86a0a3ad8..8777ce77d4 100644
--- a/lib/PTO/Transforms/PTOViewToMemref.cpp
+++ b/lib/PTO/Transforms/PTOViewToMemref.cpp
@@ -747,6 +747,24 @@ static Type convertTileBufTypeToMemRef(mlir::pto::TileBufType tbTy) {
tbTy.getMemorySpace());
}
+static Type convertTensorViewLikeTypeToMemRef(ArrayRef shape,
+ Type elementType,
+ MLIRContext *ctx,
+ bool forceDynamicShape) {
+ SmallVector resultShape;
+ if (forceDynamicShape) {
+ resultShape.assign(shape.size(), ShapedType::kDynamic);
+ } else {
+ resultShape.assign(shape.begin(), shape.end());
+ }
+
+ SmallVector dynStrides(shape.size(), ShapedType::kDynamic);
+ auto layoutAttr =
+ StridedLayoutAttr::get(ctx, ShapedType::kDynamic, dynStrides);
+ auto gmSpace = AddressSpaceAttr::get(ctx, AddressSpace::GM);
+ return MemRefType::get(resultShape, elementType, layoutAttr, gmSpace);
+}
+
static Type convertPTOTypeToMemRef(Type t) {
// 1. 处理 !pto.ptr
if (auto pty = dyn_cast(t)) {
@@ -764,11 +782,15 @@ static Type convertPTOTypeToMemRef(Type t) {
if (auto mtbTy = dyn_cast(t))
return convertTileBufTypeToMemRef(mtbTy.getSlotType());
if (auto tvTy = dyn_cast(t))
- return MemRefType::get(tvTy.getShape(), tvTy.getElementType(),
- MemRefLayoutAttrInterface(), Attribute());
+ return convertTensorViewLikeTypeToMemRef(tvTy.getShape(),
+ tvTy.getElementType(),
+ tvTy.getContext(),
+ /*forceDynamicShape=*/true);
if (auto partTy = dyn_cast(t))
- return MemRefType::get(partTy.getShape(), partTy.getElementType(),
- MemRefLayoutAttrInterface(), Attribute());
+ return convertTensorViewLikeTypeToMemRef(partTy.getShape(),
+ partTy.getElementType(),
+ partTy.getContext(),
+ /*forceDynamicShape=*/false);
// 其他类型透传
return t;
}
diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp
index 7ce735669f..4fc209e3f3 100644
--- a/lib/PTO/Transforms/Utils.cpp
+++ b/lib/PTO/Transforms/Utils.cpp
@@ -117,6 +117,10 @@ std::optional> getOperationAliasInfo(Operation *op) {
return std::make_pair(toBufferOp.getBuffer(), toBufferOp.getTensor());
} else if (auto toTensorOp = dyn_cast(op)) {
return std::make_pair(toTensorOp.getResult(), toTensorOp.getOperand());
+ } else if (auto tileBufAddrOp = dyn_cast(op)) {
+ return std::make_pair(tileBufAddrOp.getDst(), tileBufAddrOp.getSrc());
+ } else if (auto tensorViewAddrOp = dyn_cast(op)) {
+ return std::make_pair(tensorViewAddrOp.getDst(), tensorViewAddrOp.getSrc());
}
return std::nullopt;
}
diff --git a/lib/PTO/Transforms/VPTOSplitCVModule.cpp b/lib/PTO/Transforms/VPTOSplitCVModule.cpp
index 622a6e9375..581736120f 100644
--- a/lib/PTO/Transforms/VPTOSplitCVModule.cpp
+++ b/lib/PTO/Transforms/VPTOSplitCVModule.cpp
@@ -223,7 +223,12 @@ static ModuleOp cloneModuleForKind(ModuleOp source, FunctionKernelKind kind,
}
static LogicalResult splitCVModule(ModuleOp module) {
- if (hasKernelKind(module) || hasKernelKindChildModule(module))
+ if (auto kernelKind = module->getAttrOfType(
+ FunctionKernelKindAttr::name)) {
+ rewriteSectionsForKind(module, kernelKind.getKernelKind());
+ return success();
+ }
+ if (hasKernelKindChildModule(module))
return success();
if (!hasCVSections(module))
return success();
diff --git a/ptodsl/README.md b/ptodsl/README.md
index d7182aa2b3..8d638f9abb 100644
--- a/ptodsl/README.md
+++ b/ptodsl/README.md
@@ -276,7 +276,9 @@ The user guide under `ptodsl/docs/user_guide/` is the canonical PTODSL API
reference. This README keeps only a compact map of the public surface:
- `@pto.jit`: the only host-visible kernel entry
-- `@pto.cube`, `@pto.simd`, `@pto.simt`: hardware-unit sub-kernels
+- `@pto.tileop`: custom tile-op helper surface for vector-style and cube-style sub-kernels
+- `@pto.simt`: SIMT helper surface with launch dimensions
+- legacy `@pto.cube`, `@pto.simd`: removed public aliases; PTODSL now diagnoses them and directs users to `@pto.tileop`
- `pto.ptr(...)` + runtime PTO scalar annotations: public entry ABI
- `pto.make_tensor_view(...)`, `pto.partition_view(...)`, `pto.alloc_tile(...)`:
core data-model builders
diff --git a/ptodsl/docs/user_guide/01-introduction.md b/ptodsl/docs/user_guide/01-introduction.md
index 7435052cd2..a1f7d0dca4 100644
--- a/ptodsl/docs/user_guide/01-introduction.md
+++ b/ptodsl/docs/user_guide/01-introduction.md
@@ -64,8 +64,7 @@ Python Wrapper L0 user-facing wrapper (NumPy, torch-npu, pure Pyth
│ └─ backend="emitc" EmitC backend, mode="auto" only
├─ Tile Ops tile.load, tile.store, tile.add, ...
├─ MTE Ops mte_load / mte_store / mte_gm_ub / ...
- ├─ @pto.cube matrix products (mad, mte_l1_l0a, mte_l0c_ub, ...)
- ├─ @pto.simd row-wise vector math (vlds, vadd, vexp, vsts, ...)
+ ├─ @pto.tileop custom tile-op helpers (vector- and cube-style bodies)
└─ @pto.simt scalar-like compute (lds, sts, pointwise blends, ...)
```
@@ -247,21 +246,25 @@ micro-instruction surface — MTE ops, explicit sync, and pointer-level
control — so you can mix tile operations with hand-authored instructions in
the same kernel.
-### Sub-kernels — `@pto.cube` / `@pto.simd` / `@pto.simt`
+### Sub-kernels — `@pto.tileop` / `@pto.simt`
These are hardware-bound compute sub-kernels, each mapped to a specific NPU compute unit:
-- **`@pto.cube`** consumes UB tiles and explicit cube-local scratch (LEFT, RIGHT, ACC, BIAS). Typical operations: `mad`, `mte_l1_l0a`, `mte_l1_l0b`, `mte_l0c_ub`.
-
-- **`@pto.simd`** operates on vector registers (`vreg`). Typical operations: `vlds`, `vadd`, `vexp`, `vcgmax`, `vsts`. Vector registers never cross the simd function boundary — persistent state is written back to UB tiles.
+- **`@pto.tileop`** is the custom tile-op surface for both vector-style and cube-style helper bodies that operate on UB tiles and hardware-local scratch. Typical operations range from `vlds` / `vadd` / `vexp` / `vsts` to `mte_l1_l0a` / `mad` / `mte_l0c_ub`.
- **`@pto.simt`** is a scalar-programmable processor group that executes scalar instructions across many work-items in parallel. Typical operations: `lds`, `sts`, scalar arithmetic and comparison. Well-suited for per-element tile walks, boundary metadata, and pointwise blends.
-Each can be invoked as a named decorated function (`@pto.cube` /
-`@pto.simd` / `@pto.simt`) or inline as a context manager
-(`with pto.cube():`, `with pto.simd():`, `with pto.simt():`).
+Tile-op helpers can be invoked as named decorated functions (`@pto.tileop`) or
+inline context managers (`with pto.tileop():`). SIMT helpers use `@pto.simt`
+and launch dimensions via `helper[x, y, z](...)` or `pto.simt_launch(...)`.
+
+Legacy `@pto.cube` / `@pto.simd` and `with pto.cube():` / `with pto.simd():`
+are no longer supported public surfaces. PTODSL raises a diagnostic directing
+users to `@pto.tileop`.
-The boundary contract is strict: vreg values do not escape a simd kernel, cube-local state does not leak into UB, and data crosses layer boundaries only through UB-backed tiles or typed UB pointers.
+The boundary contract is strict: transient `vreg` values do not escape a
+tileop helper, cube-local state does not leak into UB, and data crosses layer
+boundaries only through UB-backed tiles or typed UB pointers.
## 1.3 Tracing execution model
@@ -295,15 +298,16 @@ GM boundary.
**Explicit orchestration path** stages the current K and V blocks with
`mte_load`, issues `pipe_barrier(Pipe.ALL)` at phase boundaries, then
-sequences four sub-kernel calls: `qk_matmul` (cube),
-`online_softmax_rows` (simd), `pv_matmul` (cube), `blend_output_rows` (simt).
+sequences four sub-kernel calls: `qk_matmul` (cube-style tileop),
+`online_softmax_rows` (vector-style tileop), `pv_matmul` (cube-style tileop),
+`blend_output_rows` (simt).
-**`@pto.cube`** performs `mte_l1_l0a` / `mte_l1_l0b` / `mad` /
-`mte_l0c_ub` for both QK^T and P@V products.
+**Cube-style `@pto.tileop` bodies** perform `mte_l1_l0a` / `mte_l1_l0b` /
+`mad` / `mte_l0c_ub` for both QK^T and P@V products.
-**`@pto.simd`** implements the online softmax update: per-row max, exp, sum,
-and alpha/beta computation using vector ops (`vlds`, `vcgmax`, `vexp`,
-`vcgadd`, `vsts`).
+**Vector-style `@pto.tileop` bodies** implement the online softmax update:
+per-row max, exp, sum, and alpha/beta computation using vector ops (`vlds`,
+`vcgmax`, `vexp`, `vcgadd`, `vsts`).
**`@pto.simt`** blends the old and new output accumulators with per-element
`lds`/`sts` and scalar arithmetic.
@@ -325,7 +329,7 @@ Chapter 11 walks through this example in full detail.
|---------|-------|
| 1 | Introduction (this chapter) |
| 2 | Quick Start — a minimal working kernel |
-| 3 | Kernel entries, kernel modules, and sub-kernels: `@pto.jit(entry=True/False, backend=...)`, `@pto.cube`, `@pto.simd`, `@pto.simt` |
+| 3 | Kernel entries, kernel modules, and sub-kernels: `@pto.jit(entry=True/False, backend=...)`, `@pto.tileop`, `@pto.simt` |
| 4 | Type system and buffer management: scalars, tiles, views, allocation |
| 5 | Control flow: trace-time Python vs device-side `pto.for_` / `pto.if_` |
| 6 | Scalar and pointer operations |
diff --git a/ptodsl/docs/user_guide/02-quick-start.md b/ptodsl/docs/user_guide/02-quick-start.md
index ba46bae7ba..534c1ebd10 100644
--- a/ptodsl/docs/user_guide/02-quick-start.md
+++ b/ptodsl/docs/user_guide/02-quick-start.md
@@ -197,8 +197,8 @@ switch that kernel to `mode="explicit"`:
```python
-# SIMD sub-kernel — vector instructions on individual rows.
-@pto.simd
+# Tile-op helper — vector instructions on individual rows.
+@pto.tileop
def add_rows(a_tile: pto.Tile, b_tile: pto.Tile, o_tile: pto.Tile,
rows: pto.index, cols: pto.index):
VEC = pto.elements_per_vreg(pto.f32)
@@ -211,7 +211,7 @@ def add_rows(a_tile: pto.Tile, b_tile: pto.Tile, o_tile: pto.Tile,
o_vec = pto.vadd(a_vec, b_vec, mask)
pto.vsts(o_vec, o_tile[r, c:], mask)
-# Single kernel entry in explicit mode — micro-instruction staging plus SIMD sub-kernel.
+# Single kernel entry in explicit mode — micro-instruction staging plus a tile-op helper.
@pto.jit(target="a5", mode="explicit")
def vec_add_micro(
A_ptr: pto.ptr(pto.f32, "gm"),
@@ -252,11 +252,11 @@ def vec_add_micro(
loops over blocks, and directly authors the micro-instruction schedule for
each block.
-- **`@pto.simd` sub-kernel**: the top-level kernel calls a SIMD sub-kernel
+- **`@pto.tileop` sub-kernel**: the top-level kernel calls a vector custom OP
for the row-wise vector work while keeping instruction staging in the
explicit entry body.
-- **Inside `@pto.simd`**: the outer Python `for range(...)` iterates over rows,
+- **Inside `@pto.tileop`**: the outer Python `for range(...)` iterates over rows,
and the inner Python `for range(...)` iterates over column chunks of the hardware vector width
(`elements_per_vreg`). Each iteration loads a vector-width slice into a
`vreg`, does the addition under a mask (for tail elements), and stores the
diff --git a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md
index a7a566624d..f295c58038 100644
--- a/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md
+++ b/ptodsl/docs/user_guide/03-kernel-entry-and-subkernels.md
@@ -2,11 +2,12 @@
PTODSL provides one kernel decorator (`@pto.jit`) with two roles
(`entry=True` / `entry=False`), two compilation backends (`vpto` / `emitc`),
-and three compute-unit sub-kernel decorators (`@pto.cube`, `@pto.simd`,
-`@pto.simt`), plus matching context managers for inline use. This chapter covers
-the `@pto.jit` entry and module contracts, the two programming models, the two
-compilation backends, sub-kernel reference, parameter contracts, and boundary
-constraints.
+custom OP sub-kernel decorators (`@pto.tileop` and `@pto.simt`), plus
+matching context managers for inline use. Legacy `@pto.cube` / `@pto.simd`
+and `with pto.cube():` / `with pto.simd():` now raise migration diagnostics
+that direct users to `@pto.tileop`. This chapter covers the `@pto.jit` entry
+and module contracts, the two programming models, the two compilation
+backends, sub-kernel reference, parameter contracts, and boundary constraints.
## 3.1 `@pto.jit` — roles, backends, and modes
@@ -21,8 +22,8 @@ Decorator overview:
mode="auto" tile-first authoring, compiler-managed staging (default)
mode="explicit" micro-instruction authoring, user-managed staging
-@pto.cube Cube-unit matrix sub-kernel
-@pto.simd SIMD-unit vector sub-kernel
+@pto.tileop Custom tile-op sub-kernel (vector- or cube-style body)
+@pto.cube / @pto.simd Legacy sub-kernel aliases; use @pto.tileop
@pto.simt SIMT-unit scalar sub-kernel
```
@@ -56,9 +57,10 @@ build level and enables sync insertion by default, while `mode="explicit"` uses
manual-address, user-managed staging contract of explicit kernels.
`@pto.jit` owns compilation (tracing + lowering), caching, and — for
-`entry=True` — runtime launch binding. The compute-unit decorators
-(`@pto.cube`, `@pto.simd`, `@pto.simt`) define sub-kernels that are called from
-within `@pto.jit` bodies.
+`entry=True` — runtime launch binding. The supported sub-kernel decorators
+(`@pto.tileop` and `@pto.simt`) define helpers that are called from within
+`@pto.jit` bodies. Legacy `@pto.cube` / `@pto.simd` are diagnosed at the
+front-end boundary and should be migrated to `@pto.tileop`.
## 3.2 `entry=True` — host-launchable kernel entry
@@ -433,9 +435,9 @@ Module bodies follow the same AST rewrite rules as `@pto.jit(entry=True)`
(see Chapter 5). In the default `mode="auto"`, Python `for` / `if` are
rewritten to device-side control flow, and the compiler handles hardware
section placement automatically — you can write `vlds` / `vadd` / `vsts`
-directly in the module body without an explicit `with pto.simd():`. In
+directly in the module body without an explicit `with pto.tileop():`. In
`mode="explicit"`, you must manage hardware sections yourself with
-`with pto.simd():`, `with pto.cube():`, or `with pto.simt():`.
+`with pto.tileop():` or `with pto.simt():`.
### Interface protocol
@@ -572,7 +574,7 @@ there are two kinds of helpers:
- **Plain Python helpers** for code organization, repeated index math,
partition construction, and orchestration that should stay in the caller's
context.
-- **Sub-kernels** (`@pto.cube`, `@pto.simd`, `@pto.simt`) when the helper must
+- **Sub-kernels** (`@pto.tileop`, `@pto.simt`) when the helper must
run on a specific hardware unit or use unit-local value categories such as
`vreg` or cube-local scratch.
@@ -582,8 +584,9 @@ backend. They are traced as part of the enclosing `@pto.jit` specialization
and therefore inherit the caller's context.
Use a sub-kernel when the helper's semantics belong to a specific unit:
-vector register math on SIMD, matrix instructions on Cube, or scalar-thread
-work on SIMT. Sub-kernels are the only public way to express that boundary.
+vector register math or cube instructions via `@pto.tileop`, or scalar-thread
+work via `@pto.simt`. Sub-kernels are the only public way to express that
+boundary.
Named sub-kernels and plain nested helpers both use the same default AST
rewrite behavior when they are traced from a compiled specialization.
@@ -591,18 +594,18 @@ rewrite behavior when they are traced from a compiled specialization.
Sub-kernels are the mechanism for custom compute in PTODSL — when Tile Ops
cover your needs, you don't need one; when they don't, a sub-kernel gives you
direct access to the hardware unit. In auto mode, a sub-kernel's parameters
-are restricted to `Tile` and PTO scalar types — the compiler owns staging and
-sync. In explicit mode, sub-kernels may also accept `PartitionTensorView` and
-`pto.ptr` parameters, matching the richer type surface available there.
-This richer pointer surface belongs to the **in-kernel orchestration and
-sub-kernel boundary**, not to the public `@pto.jit` host entry ABI.
-Section 3.3 covers each sub-kernel decorator in detail.
+follow the decorator's role-specific ABI — the compiler still owns staging and
+sync for tileop-style helpers. `@pto.tileop` accepts `Tile`, `TensorView`,
+`PartitionTensorView`, and PTO scalar parameters. `@pto.simt` additionally
+accepts typed `pto.ptr(...)` parameters. These richer in-kernel boundary types
+do not change the public `@pto.jit` host entry ABI. Section 3.3 covers each
+sub-kernel decorator in detail.
### Module vs sub-kernel
**Module or sub-kernel?** A simple rule:
-- Logic that **must run on a specific hardware unit** (Cube, SIMD, or SIMT)
- and operates on tiles → use a sub-kernel (`@pto.cube`, `@pto.simd`, `@pto.simt`).
+- Logic that **must run as a custom tile op or SIMT helper**
+ and operates on tiles → use a sub-kernel (`@pto.tileop` or `@pto.simt`).
- General device-side code organisation — allocating tiles, partitioning GM
views, calling sub-kernels, mixing backends → use a kernel module
(`@pto.jit(entry=False)`).
@@ -611,11 +614,11 @@ Modules **can** call sub-kernels (they are callable from both entries and
modules). Sub-kernels **cannot** call modules — data crosses the sub-kernel
boundary only through UB tiles, not through nested function calls.
-| | `@pto.jit(entry=False)` module | `@pto.simd` / `@pto.simt` / `@pto.cube` |
+| | `@pto.jit(entry=False)` module | `@pto.tileop` / `@pto.simt` |
|---|---|---|
| Positioning | General device-side function | **Custom tile op** — hardware-bound compute primitive |
| Scope | Orchestration, tile allocation, data movement, sub-kernel dispatch | Single-hardware-unit compute logic |
-| ABI | **C ABI: ptr + PTO scalars only**. Tile/TensorView/PartitionTensorView cannot cross the function boundary. Caller passes `tile.as_ptr()`; module constructs local tiles internally | **Tile + PTO scalars**. In/out via mutable Tile parameters. `@pto.simt` additionally accepts typed UB pointers |
+| ABI | **C ABI: ptr + PTO scalars only**. Tile/TensorView/PartitionTensorView cannot cross the function boundary. Caller passes `tile.as_ptr()`; module constructs local tiles internally | Role-specific PTODSL helper ABI. `@pto.tileop`: `Tile`, `TensorView`, `PartitionTensorView`, PTO scalars. `@pto.simt`: `Tile`, typed `pto.ptr(...)`, PTO scalars |
| Backend | VPTO or EmitC | Always VPTO |
| Compilation | Compiled as a separate child module, linked automatically | Outlined as a helper function inside the owning caller/module |
| Callable from | Entries and other modules | Entries and modules |
@@ -651,10 +654,11 @@ gain access to MTE ops, explicit synchronization, and pointer manipulation.
When you need precise control over individual instructions and phase ordering,
you can drop below the tile abstraction without leaving the `@pto.jit` entry.
-The richer type surface also applies to sub-kernels: in auto mode, a
-sub-kernel's parameters are restricted to `Tile` and PTO scalar types; in
-explicit mode they may also accept `PartitionTensorView` and `pto.ptr`,
-matching the types available in the enclosing orchestration code.
+Explicit mode broadens the orchestration code you can write inside `@pto.jit`
+and `@pto.jit(entry=False)` bodies. Sub-kernel ABIs themselves remain
+role-specific rather than mode-specific: `@pto.tileop` accepts `Tile`,
+`TensorView`, `PartitionTensorView`, and PTO scalars, while `@pto.simt`
+accepts `Tile`, typed `pto.ptr(...)`, and PTO scalars.
```python
@@ -701,7 +705,7 @@ def process_block(q_tile, k_part, v_part, k_tile, v_tile,
nburst=(rows, ub_row_stride, gm_row_stride))
```
-Sub-kernel calls and inline sub-kernel scopes (`with pto.simd():`, etc.) work
+Sub-kernel calls and inline sub-kernel scopes (`with pto.tileop():`, etc.) work
identically in both modes.
### Choosing between modes
@@ -823,11 +827,16 @@ built-in ops don't cover your needs.
**Sub-kernels are custom tile ops.** Their contract is strict:
-- **Inputs**: `Tile` references and PTO scalars (`pto.i32`, `pto.f32`, ...).
- Data arrives from UB via tile handles; the sub-kernel does not own GM
- addressing or DMA orchestration.
-- **Outputs**: written back to UB tiles. Sub-kernels have no return values —
- results are communicated by writing to mutable `Tile` parameters.
+- **Inputs**: role-specific PTODSL boundary values. `@pto.tileop` accepts
+ `Tile`, `TensorView`, `PartitionTensorView`, and PTO scalars (`pto.i32`,
+ `pto.f32`, ...). `@pto.simt` accepts `Tile`, typed `pto.ptr(...)`, and PTO
+ scalars. Data still flows through device-side boundary objects; sub-kernels
+ do not define a host-visible ABI.
+- **Outputs**: written back to UB tiles. Mutable `Tile` parameters remain the
+ primary output path. In the current MVP, decorated tileop-style helpers may
+ additionally return PTO scalar values through `func.call` results, but
+ `Tile`, `TensorView`, `PartitionTensorView`, `vreg`, and mask values still do
+ not cross the helper boundary as results.
- **No cross-boundary vreg**: vector registers (`vreg`) and cube-local state
(LEFT, RIGHT, ACC) are private to the sub-kernel body and never escape.
@@ -838,32 +847,34 @@ When to use a sub-kernel vs a kernel module:
views, or mix backends → use an `@pto.jit(entry=False)` kernel module
instead. Modules can call sub-kernels, but sub-kernels cannot call modules.
-Sub-kernels are decorated with `@pto.cube`, `@pto.simd`, or `@pto.simt`.
-PTODSL lowers both surface forms to real helper `func.func` bodies instead of
-flattening them directly into the surrounding caller. They can be authored in
-two ways:
+Sub-kernels are decorated with `@pto.tileop` or `@pto.simt`. PTODSL lowers
+them to real helper `func.func` bodies instead of flattening them directly
+into the surrounding caller. Legacy `@pto.cube` / `@pto.simd` and inline
+`with pto.cube():` / `with pto.simd():` are removed public surfaces; PTODSL
+diagnoses them and directs users to `@pto.tileop`. Sub-kernels can be authored
+in two ways:
1. **As decorated functions** — reusable, named sub-kernels called from
`@pto.jit` entries and modules.
-2. **As context managers** (`with pto.cube():`, etc.) — inline blocks for
+2. **As context managers** (`with pto.tileop():`, etc.) — inline blocks for
one-off snippets (see Section 3.8).
Named sub-kernel decorators use the same default AST rewrite model as
`@pto.jit`: supported Python `if` and `for range(...)` statements lower to
device-side control flow.
-### 3.7.1 `@pto.cube` — Cube unit (matrix operations)
+### 3.7.1 Cube-style `@pto.tileop` — matrix operations
-**Role**: `@pto.cube` is the custom tile op for matrix multiplication on the
-Cube unit. It consumes UB-resident tiles and explicit cube-local scratch
-buffers. All parameters are `Tile` references — the caller owns tile
-allocation, and the sub-kernel only expresses the compute dataflow.
+**Role**: `@pto.tileop` is also the custom tile-op surface for matrix
+multiplication on the Cube unit. It consumes UB-resident tiles and explicit
+cube-local scratch buffers. All parameters are `Tile` references — the caller
+owns tile allocation, and the sub-kernel only expresses the compute dataflow.
-**Signature**: `@pto.cube(fn=None, *, name=None, target="a5")`
+**Signature**: `@pto.tileop(fn=None, *, name=None, target="a5")`
```python
-@pto.cube
+@pto.tileop
def my_cube_kernel(
input_tile: pto.Tile, # UB tile (source data)
output_tile: pto.Tile, # UB tile (destination)
@@ -882,7 +893,7 @@ allocated with the appropriate `memory_space` (e.g., `pto.MemorySpace.LEFT`,
```python
-@pto.cube
+@pto.tileop
def qk_matmul(
q_tile: pto.Tile,
k_tile: pto.Tile,
@@ -905,29 +916,31 @@ Cube-local state (LEFT, RIGHT, ACC, BIAS) never leaks into UB — it is the
caller's responsibility to allocate scratch buffers and pass them in
explicitly.
-**Lowering model**: a decorated `@pto.cube` function becomes one reusable
-helper function inside the owning PTODSL child module. Each callsite lowers to
-`func.call` of that helper; the helper body itself contains the `pto.section.cube`
-region.
+**Lowering model**: a decorated cube-style `@pto.tileop` function becomes one
+reusable helper function inside the owning PTODSL child module. Each callsite
+lowers to `func.call` of that helper. The helper body stays as a naked
+`tileop` helper; backend passes infer the primary domain and materialize
+`pto.section.cube` later in the PTOAS pipeline.
**Invocation modes**: can be called from `@pto.jit` in either mode, or authored
-as an anonymous inline helper with `with pto.cube():` (Section 3.8).
+as an anonymous inline helper with `with pto.tileop():` (Section 3.8).
-### 3.7.2 `@pto.simd` — SIMD unit (vector operations)
+### 3.7.2 `@pto.tileop` — custom tile op (vector operations)
-**Role**: `@pto.simd` is the custom tile op for row-wise vector compute on
-the SIMD unit. It operates on vector registers (`vreg`) loaded from UB tiles
-and stores results back to UB tiles. Parameters are `Tile` references and PTO
-scalars — the sub-kernel reads tile data, computes on vector hardware, and
-writes results back through mutable tile parameters. Vector registers are
-local to the function and never cross its boundary.
+**Role**: `@pto.tileop` is the recommended custom tile-op surface for
+row-wise vector compute. It operates on vector registers (`vreg`) loaded from UB tiles
+and stores results back to UB tiles. Parameters may be `Tile`,
+`TensorView`, `PartitionTensorView`, and PTO scalars — the sub-kernel reads
+device-side descriptors, computes on vector hardware, and writes results back
+through mutable tile parameters. Vector registers are local to the function
+and never cross its boundary.
-**Signature**: `@pto.simd(fn=None, *, name=None, target="a5")`
+**Signature**: `@pto.tileop(fn=None, *, name=None, target="a5")`
-
+
```python
-@pto.simd
-def my_simd_kernel(
+@pto.tileop
+def my_tileop_kernel(
input_tile: pto.Tile, # UB tile
output_tile: pto.Tile, # UB tile
rows: pto.i32, # PTO scalar
@@ -936,19 +949,19 @@ def my_simd_kernel(
return
```
-Parameters are UB `Tile` references and PTO scalar values (`pto.i32`,
-`pto.f32`, etc.). Scalar parameters may come from `lds` reads or compile-time
-constants.
+Parameters are device-side `Tile` / `TensorView` / `PartitionTensorView`
+references and PTO scalar values (`pto.i32`, `pto.f32`, etc.). Scalar
+parameters may come from `lds` reads or compile-time constants.
-This interface contract is enforced unconditionally. A decorated `@pto.simd`
-function does not gain extra pointer-style ABI forms in explicit mode; if you
-need a broader boundary, use `@pto.jit(entry=False)` instead.
+This interface contract is enforced unconditionally. A decorated `@pto.tileop`
+function does not accept `pto.ptr(...)`; typed pointers remain SIMT-only. If
+you need a broader C-style ABI, use `@pto.jit(entry=False)` instead.
**Typical body**:
-
+
```python
-@pto.simd
+@pto.tileop
def add_rows(a_tile: pto.Tile, b_tile: pto.Tile, o_tile: pto.Tile,
rows: pto.i32, cols: pto.i32):
VEC = pto.elements_per_vreg(pto.f32)
@@ -963,21 +976,27 @@ def add_rows(a_tile: pto.Tile, b_tile: pto.Tile, o_tile: pto.Tile,
```
The boundary contract: `vreg` values (`a_vec`, `b_vec`, `o_vec`) are local to
-the function. The only way to persist data across a `@pto.simd` call is to
+the function. The only way to persist data across a `@pto.tileop` call is to
write it back to a UB tile via `vsts` (or `psts`, etc.).
-**Lowering model**: a decorated `@pto.simd` function becomes one reusable
+**Lowering model**: a decorated `@pto.tileop` function becomes one reusable
helper function inside the owning PTODSL child module. Each callsite lowers to
-`func.call` of that helper; the helper body itself contains the `pto.section.vector`
-region.
+`func.call` of that helper. The helper body is emitted as a naked `tileop`
+helper with `pto.tileop.helper`; backend passes infer the
+primary domain and materialize `pto.section.vector` later in the PTOAS
+pipeline.
**Invocation modes**: can be called from `@pto.jit` in either mode, or authored
-as an anonymous inline helper with `with pto.simd():` (Section 3.8).
+as an anonymous inline helper with `with pto.tileop():` (Section 3.8).
+
+PTODSL no longer exposes a separate `@pto.simd` public surface. Vector-style
+helpers use `@pto.tileop`, and the backend infers `pto.section.vector` from
+the helper body later in the PTOAS pipeline.
### 3.7.3 `@pto.simt` — SIMT unit (scalar-parallel operations)
-**Role**: `@pto.simt` is the custom tile op for per-element scalar-parallel
-compute on the SIMT unit. SIMT (Single Instruction, Multiple Threads) is a
+**Role**: `@pto.simt` is the SIMT helper surface for per-element scalar-parallel
+compute. SIMT (Single Instruction, Multiple Threads) is a
programming model where you write instructions in scalar syntax
(`scalar.load`, `scalar.store`, `a + b`), and the hardware executes them in
parallel across many threads — analogous to how a GPU SM runs a CUDA kernel.
@@ -1106,16 +1125,18 @@ Specific SIMT micro-op APIs are documented in Chapter 13.
## 3.8 Inline context manager syntax
-In addition to the decorator form, each sub-kernel unit provides a context
-manager: `with pto.cube():`, `with pto.simd():`, and `with pto.simt():`. These
-open one-off anonymous sub-kernel bodies without requiring a separate named
-Python function. Inline scopes are supported in top-level `@pto.jit` bodies.
+In addition to the decorator form, custom tile-op helpers can be written with
+`with pto.tileop():`, and SIMT helpers can be written with `with pto.simt():`.
+These open one-off anonymous sub-kernel bodies without requiring a separate
+named Python function. Inline scopes are supported in top-level `@pto.jit`
+bodies. Legacy `with pto.cube():` / `with pto.simd():` are diagnosed and
+should be rewritten to `with pto.tileop():`.
### Syntax
```python
-with pto.simd():
+with pto.tileop():
a_vec = pto.vlds(a_tile[r, c:])
b_vec = pto.vlds(b_tile[r, c:])
o_vec = pto.vadd(a_vec, b_vec, mask)
@@ -1133,7 +1154,7 @@ with pto.simt():
```python
-with pto.cube():
+with pto.tileop():
pto.mte_l1_l0a(q_tile.as_ptr(), q_l0a.as_ptr(), m, k)
pto.mte_l1_l0b(k_tile.as_ptr(), k_l0b.as_ptr(), k, n, transpose=True)
pto.mad(q_l0a.as_ptr(), k_l0b.as_ptr(), s_acc.as_ptr(), m, n, k)
@@ -1146,10 +1167,12 @@ with pto.cube():
unit.
- On block exit, PTODSL outlines the block into one anonymous helper
`func.func` and replaces the original region with a `func.call`.
-- `with pto.simd():` and `with pto.cube():` preserve their `pto.section.vector`
- / `pto.section.cube` bodies inside the outlined helper.
+- Named/decorated `@pto.tileop` helpers and inline `with pto.tileop():`
+ scopes both lower as naked `pto.tileop.helper` bodies. PTOAS infers the
+ primary domain and materializes `pto.section.vector` /
+ `pto.section.cube` later in the pipeline.
- `with pto.simt():` preserves its scalar body inside one outlined
- `pto.simt_entry` helper, and the caller emits `pto.store_vfsimt_info`.
+ `pto.simt_entry` helper, and the caller emits `pto.simt_launch`.
- Values defined inside the inline sub-kernel cannot escape the block directly.
Use Tiles, typed pointers, or other mutable references to communicate results
back to the caller.
@@ -1171,10 +1194,10 @@ The two forms can be freely mixed in the same `@pto.jit` body.
## 3.9 Boundary contracts
-**Sub-kernels are custom tile ops.** Their I/O contract is strict: data enters
-via `Tile` handles and PTO scalars; results exit by writing to mutable `Tile`
-parameters. `TensorView` and `PartitionTensorView` belong to the orchestration
-layer and are NOT accepted by sub-kernels.
+**Sub-kernels are custom tile ops.** Their I/O contract is strict: tileop-style
+helpers accept `Tile`, `TensorView`, `PartitionTensorView`, and PTO scalars;
+SIMT helpers accept `Tile`, typed `pto.ptr(...)`, and PTO scalars. Results
+still exit by writing to mutable references or returning PTO scalar values only.
**Modules use the C ABI.** Module boundaries (`entry=False`) are real function
calls — only `pto.ptr` and PTO scalars can cross. `Tile`, `TensorView`, and
@@ -1184,13 +1207,13 @@ calls — only `pto.ptr` and PTO scalars can cross. `Tile`, `TensorView`, and
|----------|---------|
| Host → `@pto.jit(entry=True)` | explicit GM pointers + runtime scalars |
| Entry / module → `@pto.jit(entry=False)` module | **`pto.ptr` + PTO scalars only** (C ABI). Caller passes `tile.as_ptr()`; module constructs local tiles internally |
-| Entry / module → sub-kernel (`auto` mode) | **`Tile` + PTO scalars only**. Compiler handles staging + sync |
-| Entry / module → sub-kernel (`explicit` mode) | `Tile`, `PartitionTensorView`, `pto.ptr`, PTO scalars |
-| `@pto.jit` → `with pto.{cube,simd,simt}:` | Captured `Tile` / ptr / scalar values from enclosing scope |
+| Entry / module → `@pto.tileop` | `Tile`, `TensorView`, `PartitionTensorView`, PTO scalars |
+| Entry / module → `@pto.simt` | `Tile`, typed `pto.ptr(...)`, PTO scalars |
+| `@pto.jit` → `with pto.{tileop,simt}:` | Captured `Tile` / ptr / scalar values from enclosing scope |
| Sub-kernel → sub-kernel | Not allowed (go through UB tiles via the caller) |
| Sub-kernel → module | Not allowed (sub-kernels cannot call out) |
| Inline sub-kernel → caller | No direct SSA return path; write through Tile / ptr / mutable references |
-| `@pto.simd` → caller | Only via `vsts`/`psts` to UB tiles; `vreg` cannot escape |
+| `@pto.tileop` → caller | Write results through mutable Tile / view operands; transient `vreg` and cube-local state cannot escape |
| Cube-local → UB | Only via `mte_l0c_ub`; LEFT/RIGHT/ACC/BIAS are private |
| `entry=False` module → caller | No return values; data crosses only via mutable references |
diff --git a/ptodsl/docs/user_guide/05-control-flow.md b/ptodsl/docs/user_guide/05-control-flow.md
index 4ea132257b..32df78726d 100644
--- a/ptodsl/docs/user_guide/05-control-flow.md
+++ b/ptodsl/docs/user_guide/05-control-flow.md
@@ -4,7 +4,7 @@ PTODSL uses a **tracing** compilation model. When you call `kernel.compile(...)`
This has one critical implication for how you write loops and branches:
-- **Python native `for`/`if`** is rewritten to device-side control flow by default in `@pto.jit` bodies and named `@pto.cube` / `@pto.simd` / `@pto.simt` sub-kernels. A `for i in range(rows)` loop records a device loop, and a runtime `if` records both branches.
+- **Python native `for`/`if`** is rewritten to device-side control flow by default in `@pto.jit` bodies and named `@pto.tileop` / `@pto.simt` sub-kernels. A `for i in range(rows)` loop records a device loop, and a runtime `if` records both branches.
- **`pto.const_expr` / `pto.static_range`** keep compile-time Python behavior when you want trace-time specialization or unrolling.
- **`pto.for_` / `pto.if_`** produce device-side control flow. The loop bound or branch condition can be a runtime value, and the hardware will execute the loop or take the branch dynamically.
@@ -465,7 +465,7 @@ def debug_kernel(*, BLOCK: pto.const_expr = 4):
pto.pipe_barrier(pto.Pipe.ALL)
-@pto.simd(ast_rewrite=False)
+@pto.tileop(ast_rewrite=False)
def debug_simd_helper():
if pto.const_expr(True):
pto.pipe_barrier(pto.Pipe.ALL)
diff --git a/ptodsl/docs/user_guide/07-data-movement-ops.md b/ptodsl/docs/user_guide/07-data-movement-ops.md
index cabe44abdf..e49aa40b3c 100644
--- a/ptodsl/docs/user_guide/07-data-movement-ops.md
+++ b/ptodsl/docs/user_guide/07-data-movement-ops.md
@@ -307,7 +307,7 @@ def process_block(k_part, v_part, k_tile, v_tile, o_tile, o_part,
## 7.3 Vector loads (simd)
-Inside `@pto.simd`, data moves between UB tiles and vector registers (`vreg`). Vector loads read a contiguous chunk of a tile row into a `vreg`; the chunk size equals the hardware vector width for the element type (e.g., 64 elements for `f32`, 128 for `f16`).
+Inside `@pto.tileop`, data moves between UB tiles and vector registers (`vreg`). Vector loads read a contiguous chunk of a tile row into a `vreg`; the chunk size equals the hardware vector width for the element type (e.g., 64 elements for `f32`, 128 for `f16`).
### Tile-index syntax
@@ -824,7 +824,9 @@ pto.vstas(align, ub_dst_f32, pto.const(64))
## 7.5 Cube data movement (cube)
-Inside `@pto.cube`, data flows through a hierarchy of private buffers: GM → L1 (cbuf) → L0A/L0B (operand buffers) → L0C (accumulator) → UB or back to GM.
+Inside a cube-style `@pto.tileop` helper, data flows through a hierarchy of
+private buffers: GM → L1 (cbuf) → L0A/L0B (operand buffers) → L0C
+(accumulator) → UB or back to GM.
### Staging: GM → L1 and L1 → UB
@@ -1065,11 +1067,12 @@ pto.mte_l0c_ub(acc, ub, 16, 32, 16, 32, split=pto.SplitMode.N) # split N
### Typical cube dataflow in a matmul
-A full cube matmul (`@pto.cube`) follows this dataflow pattern:
+A full cube matmul (a cube-style `@pto.tileop` helper) follows this dataflow
+pattern:
```python
-@pto.cube
+@pto.tileop
def qk_matmul(
q_tile: pto.Tile,
k_tile: pto.Tile,
@@ -1226,7 +1229,7 @@ Cube (producer) side:
```python
-@pto.cube
+@pto.tileop
def producer(src_tile: pto.Tile):
c2v.init_cube()
entry = c2v.alloc(split=0)
@@ -1239,7 +1242,7 @@ Vector (consumer) side:
```python
-@pto.simd
+@pto.tileop
def consumer(dst_tile: pto.Tile):
c2v.init_simd()
entry = c2v.pop(split=0)
@@ -1270,7 +1273,7 @@ Vector (producer) side:
```python
-@pto.simd
+@pto.tileop
def producer(src_tile: pto.Tile):
v2c.init_simd()
entry = v2c.alloc(split=0)
@@ -1282,7 +1285,7 @@ Cube (consumer) side:
```python
-@pto.cube
+@pto.tileop
def consumer(dst_tile: pto.Tile):
v2c.init_cube()
entry = v2c.pop(split=0)
@@ -1320,7 +1323,7 @@ Cube (producer) transaction:
```python
-@pto.cube
+@pto.tileop
def producer(src_tile: pto.Tile):
c2v_peer.init_cube()
c2v_peer.push(src_tile, split=0)
@@ -1330,7 +1333,7 @@ Vector (consumer) transaction:
```python
-@pto.simd
+@pto.tileop
def consumer(dst_tile: pto.Tile):
c2v.init_simd()
tile = c2v.pop(result_type=dst_tile, split=0)
@@ -1385,7 +1388,7 @@ def cube_producer(
offsets=[0, 0], sizes=[16, 16])
a_tile = pto.alloc_tile(shape=[16, 16], dtype=pto.f32)
- with pto.cube():
+ with pto.tileop():
pto.tile.load(a_part, a_tile)
c2v.init_cube()
entry = c2v.alloc(split=0)
@@ -1412,7 +1415,7 @@ def vector_consumer(
pto.make_tensor_view(dst, shape=[16, 16], strides=[16, 1]),
offsets=[0, 0], sizes=[16, 16])
- with pto.simd():
+ with pto.tileop():
c2v.init_simd()
entry = c2v.pop(split=0)
entry_part = pto.partition_view(entry, offsets=[0, 0], sizes=[16, 16])
diff --git a/ptodsl/docs/user_guide/08-compute-operations.md b/ptodsl/docs/user_guide/08-compute-operations.md
index ac4e32a3a9..8f5b3f623a 100644
--- a/ptodsl/docs/user_guide/08-compute-operations.md
+++ b/ptodsl/docs/user_guide/08-compute-operations.md
@@ -1254,9 +1254,9 @@ pto.tile.gemv_mx_bias(lhs_l0a_mx, lhs_scale, rhs_l0b_mx, rhs_scale, bias_tile, a
---
-## 8.2 Vector compute (L3 — `@pto.simd`)
+## 8.2 Vector compute (L3 — `@pto.tileop`)
-Vector compute ops operate on `VRegType` values inside `@pto.simd` sub-kernels. Every vector op takes a `MaskType` predicate that gates which lanes participate; masked-off lanes produce an unspecified result (use the result only where the mask is true, or feed it to a masked store).
+Vector compute ops operate on `VRegType` values inside `@pto.tileop` sub-kernels. Every vector op takes a `MaskType` predicate that gates which lanes participate; masked-off lanes produce an unspecified result (use the result only where the mask is true, or feed it to a masked store).
All vector ops in this section follow the pattern established in Section 7.3 for tile-index and pointer-form addressing. The signatures below use the vector-register form — tile-index forms load into `vreg` first, then compute.
@@ -1605,7 +1605,7 @@ exp_f16_odd = pto.vmulscvt(exp_f32_odd, 1.0, mask, rnd=pto.VcvtRoundMode.A, par
### 8.2.7 Vector type conversion and packing
-These ops change the element type or layout of vector registers. They are distinct from the tile-level `tile.cvt` — they operate on `VRegType` values inside `@pto.simd` and are the explicit micro-op counterparts to higher-level conversion helpers.
+These ops change the element type or layout of vector registers. They are distinct from the tile-level `tile.cvt` — they operate on `VRegType` values inside `@pto.tileop` and are the explicit micro-op counterparts to higher-level conversion helpers.
#### `pto.vcvt(src: VRegType, to_dtype: DType, mask: MaskType, *, rnd: VcvtRoundMode | None = None, sat: VcvtSatMode | None = None, part: VcvtPartMode | None = None) -> VRegType`
@@ -1720,7 +1720,7 @@ packed_high = pto.vpack(vec_i32, pto.VPackPart.HIGHER) # upper 64 lanes -> 128
---
-## 8.3 Cube compute (L3 — `@pto.cube`)
+## 8.3 Cube compute (L3 — cube-style `@pto.tileop`)
The Cube unit performs matrix multiplication. Its operands are typed pointers into cube-local buffers — L0A (left operand), L0B (right operand), L0C (accumulator), and BIAS. Cube data movement (`mte_l1_l0a`, `mte_l1_l0b`, `mte_l0c_ub`, etc.) was covered in Section 7.5; this section covers the compute instruction itself.
@@ -1824,7 +1824,7 @@ A full cube matmul follows a three-stage pattern: stage operands into L0A/L0B, c
```python
-@pto.cube
+@pto.tileop
def qk_matmul(
q_tile: pto.Tile,
k_tile: pto.Tile,
diff --git a/ptodsl/docs/user_guide/10-sync-ops.md b/ptodsl/docs/user_guide/10-sync-ops.md
index 42b3089400..c0d649f01f 100644
--- a/ptodsl/docs/user_guide/10-sync-ops.md
+++ b/ptodsl/docs/user_guide/10-sync-ops.md
@@ -437,7 +437,7 @@ Where do sync operations belong in PTODSL's public entry model?
|---------|---------------------|
| `@pto.jit(mode="auto")` | Users can write sync explicitly when needed. PTOAS also provides an `--enable-insert-sync` option that auto-inserts `set_flag`/`wait_flag` pairs based on op-to-pipe mapping. |
| `@pto.jit(mode="explicit")` | The compiler does not insert sync — the user is fully responsible. Place `set_flag`/`wait_flag` between MTE and compute, `mem_bar` between compute phases, `pipe_barrier` at phase boundaries. |
-| Shared `@pto.cube` / `@pto.simd` / `@pto.simt` helpers | Cross-pipeline ordering is provided by the surrounding `@pto.jit` schedule. Helpers may still use `mem_bar` for intra-pipeline ordering when UB addresses alias. |
+| Shared `@pto.tileop` / `@pto.simt` helpers | Cross-pipeline ordering is provided by the surrounding `@pto.jit` schedule. Helpers may still use `mem_bar` for intra-pipeline ordering when UB addresses alias. |
**Rule of thumb**: in `mode="auto"`, think in tiles and let the compiler handle
orchestration. In `mode="explicit"`, think in micro-instructions and place the
diff --git a/ptodsl/docs/user_guide/11-flash-attention-walkthrough.md b/ptodsl/docs/user_guide/11-flash-attention-walkthrough.md
index ddd2fd42d9..1cac1de4b5 100644
--- a/ptodsl/docs/user_guide/11-flash-attention-walkthrough.md
+++ b/ptodsl/docs/user_guide/11-flash-attention-walkthrough.md
@@ -11,8 +11,8 @@ flash_attention(...) L0 user-facing wrapper
└─ @pto.jit(entry=True, mode="explicit") flash_attention_kernel
├─ Tile Ops tile.load / tile.store at the GM↔UB boundary
├─ explicit orchestration mte_load / pipe_barrier / pointer sequencing
- ├─ @pto.cube qk_matmul / pv_matmul
- ├─ @pto.simd online_softmax_rows
+ ├─ @pto.tileop qk_matmul / pv_matmul (cube-style)
+ ├─ @pto.tileop online_softmax_rows (vector-style)
└─ @pto.simt materialize_tile_bounds / blend_output_rows
```
@@ -438,13 +438,13 @@ The simt sub-kernel blends the old output accumulator with the new PV contributi
Each `pipe_barrier(Pipe.ALL)` between phases is explicit in the orchestration body. This is intentional: at the orchestration boundary, the user controls pipeline ordering. Auto mode may still use synchronization primitives where needed, but it does so around compiler-managed tile staging rather than user-authored instruction scheduling.
-## 11.5 Cube sub-kernel — `@pto.cube`
+## 11.5 Cube-style tileop sub-kernel
### `qk_matmul` — `S = Q @ K^T`
```python
-@pto.cube
+@pto.tileop
def qk_matmul(
q_mat: pto.Tile,
k_mat: pto.Tile,
@@ -476,7 +476,7 @@ The cube kernel does not allocate scratch — the caller (top-level kernel) owns
```python
-@pto.cube
+@pto.tileop
def pv_matmul(
p_mat: pto.Tile,
v_mat: pto.Tile,
@@ -497,10 +497,10 @@ def pv_matmul(
Structurally identical to `qk_matmul`, but without transposition and with different input/output tiles. The scratch tiles `p_l0a`, `v_l0b`, and `pv_acc` are reused across KV blocks — the caller (top-level kernel) allocates them once.
-## 11.6 SIMD sub-kernel — online softmax
+## 11.6 Tile-op sub-kernel — online softmax
```python
-@pto.simd
+@pto.tileop
def online_softmax_rows(
s_tile: pto.Tile,
p_tile: pto.Tile,
@@ -682,6 +682,6 @@ After all KV blocks: the top-level kernel issues `tile.store(o_final_tile, o_par
**Tile-level boundary vs micro-instruction boundary**: `tile.load`/`tile.store` are the tile-atomic surface used in auto mode and at the top-level tile boundary of this sketch. `mte_load` appears in explicit orchestration, authored as individual pointer-based instructions. The abstraction split is auto mode as tile-centric authoring, explicit mode as user-ordered orchestration.
-**No vreg across sub-kernel boundaries**: vector registers are local to each `@pto.simd` kernel. Data crosses sub-kernel boundaries through UB tiles — the boundary contract is enforced by the type system.
+**No vreg across sub-kernel boundaries**: vector registers are local to each `@pto.tileop` kernel. Data crosses sub-kernel boundaries through UB tiles — the boundary contract is enforced by the type system.
-**Invocation flexibility**: This sketch uses the explicit `@pto.jit(entry=True, mode="explicit")` path for full micro-instruction control. The same named sub-kernels can also be reused from `@pto.jit(mode="auto")` when the body stays within the auto-mode contract, or written inline as context managers (`with pto.simd():`, etc.). The orchestration logic could be extracted into `@pto.jit(entry=False)` kernel modules for reuse across multiple entry kernels. See Chapter 3 for details.
+**Invocation flexibility**: This sketch uses the explicit `@pto.jit(entry=True, mode="explicit")` path for full micro-instruction control. The same named sub-kernels can also be reused from `@pto.jit(mode="auto")` when the body stays within the auto-mode contract, or written inline as context managers (`with pto.tileop():`, etc.). The orchestration logic could be extracted into `@pto.jit(entry=False)` kernel modules for reuse across multiple entry kernels. See Chapter 3 for details.
diff --git a/ptodsl/docs/user_guide/12-additional-examples.md b/ptodsl/docs/user_guide/12-additional-examples.md
index 6e870c9a72..65a1958377 100644
--- a/ptodsl/docs/user_guide/12-additional-examples.md
+++ b/ptodsl/docs/user_guide/12-additional-examples.md
@@ -76,11 +76,11 @@ When a data dimension is not evenly divisible by the tile size or the hardware v
### 12.2.1 Tail handling in a SIMD kernel
-Below is a self-contained `@pto.simd` kernel that adds two tiles row by row, handling column tails with `make_mask`:
+Below is a self-contained `@pto.tileop` kernel that adds two tiles row by row, handling column tails with `make_mask`:
-
+
```python
-@pto.simd
+@pto.tileop
def add_rows_with_tail(a_tile: pto.Tile, b_tile: pto.Tile, o_tile: pto.Tile,
rows: pto.i32, cols: pto.i32):
VEC = pto.elements_per_vreg(pto.f32) # 64 for f32
@@ -160,13 +160,16 @@ def vec_add_with_tail(
## 12.3 GEMM: matrix multiplication on the Cube unit
-This example demonstrates a complete GEMM kernel: `C = A @ B` where A is `[M, K]` and B is `[K, N]`. It uses `@pto.jit` for tile allocation and loop scheduling, and `@pto.cube` for the actual matrix multiply.
+This example demonstrates a complete GEMM kernel: `C = A @ B` where A is
+`[M, K]` and B is `[K, N]`. It uses `@pto.jit` for tile allocation and loop
+scheduling, and a cube-style `@pto.tileop` helper for the actual matrix
+multiply.
### 12.3.1 Cube sub-kernel
```python
-@pto.cube
+@pto.tileop
def gemm_tile(a_mat: pto.Tile, b_mat: pto.Tile, o_tile: pto.Tile,
a_l0a: pto.Tile, b_l0b: pto.Tile, o_acc: pto.Tile):
m = a_mat.valid_shape[0]
@@ -380,10 +383,10 @@ def online_layernorm(
|------|-----|
| Whole-kernel orchestration, GM↔UB boundary | `@pto.jit` |
| Tile-level data movement | `tile.load` / `tile.store` |
-| Custom row-wise vector math | `@pto.simd` |
+| Custom row-wise vector math | `@pto.tileop` |
| Custom per-element logic | `@pto.simt` |
-| Matrix multiply | `@pto.cube` |
+| Matrix multiply | cube-style `@pto.tileop` |
| Micro-instruction-level control | `mode="explicit"` |
-| Inline compute for quick prototyping | `with pto.simd():` etc. |
+| Inline compute for quick prototyping | `with pto.tileop():` etc. |
-**Respect boundary contracts.** Vregs don't cross `@pto.simd` boundaries. Cube-local state doesn't leak into UB. Tile Ops and MTE Ops belong to different programming models — use Tile Ops in `mode="auto"`, and micro-instructions in `mode="explicit"`.
+**Respect boundary contracts.** Vregs don't cross `@pto.tileop` boundaries. Cube-local state doesn't leak into UB. Tile Ops and MTE Ops belong to different programming models — use Tile Ops in `mode="auto"`, and micro-instructions in `mode="explicit"`.
diff --git a/ptodsl/examples/dynamic_softmax_launch.py b/ptodsl/examples/dynamic_softmax_launch.py
index 92abd2faa3..59d4662996 100644
--- a/ptodsl/examples/dynamic_softmax_launch.py
+++ b/ptodsl/examples/dynamic_softmax_launch.py
@@ -92,7 +92,7 @@ def dynamic_softmax(
pto.set_flag("MTE2", "V", event_id=0)
pto.wait_flag("MTE2", "V", event_id=0)
- with pto.simd():
+ with pto.tileop():
remaining_rows = runtime_rows
for row_base in range(0, runtime_rows, lane_num):
active_rows, remaining_after_pack = pto.make_mask(pto.f32, remaining_rows)
diff --git a/ptodsl/examples/fa_dn_ptodsl.py b/ptodsl/examples/fa_dn_ptodsl.py
index 90f271745e..c4dd3f2a18 100644
--- a/ptodsl/examples/fa_dn_ptodsl.py
+++ b/ptodsl/examples/fa_dn_ptodsl.py
@@ -181,7 +181,7 @@ def __call__(self):
self.sync.free()
-@pto.cube
+@pto.tileop
def qk_matmul_stage(
qMatTile: pto.Tile,
kMatTile: pto.Tile,
@@ -200,7 +200,7 @@ def qk_matmul_stage(
pto.mte_l0c_ub(qkAccTile.as_ptr(), qkVecTileSub.as_ptr(), rows, cols, cols, cols, 0)
-@pto.cube
+@pto.tileop
def pv_matmul_stage(
pMatTileSub: pto.Tile,
vMatTile: pto.Tile,
diff --git a/ptodsl/examples/fast_inverse_dense_launch.py b/ptodsl/examples/fast_inverse_dense_launch.py
index c26a3f553d..8750db556a 100644
--- a/ptodsl/examples/fast_inverse_dense_launch.py
+++ b/ptodsl/examples/fast_inverse_dense_launch.py
@@ -94,7 +94,7 @@ def fast_inverse_dense_f32(
pto.set_flag("MTE2", "V", event_id=0)
pto.wait_flag("MTE2", "V", event_id=0)
- with pto.simd():
+ with pto.tileop():
active, _ = pto.make_mask(pto.f32, batch_i32)
d00 = pto.vlds(in_tile[0, 0:])
d01 = pto.vlds(in_tile[1, 0:])
diff --git a/ptodsl/examples/flash_attention/gu.py b/ptodsl/examples/flash_attention/gu.py
index ebca669034..679c485d99 100644
--- a/ptodsl/examples/flash_attention/gu.py
+++ b/ptodsl/examples/flash_attention/gu.py
@@ -13,7 +13,7 @@
- ``fa_gu_init_vpto_kernel`` / ``fa_gu_update_vpto_kernel`` are ptr-ABI VPTO
child modules.
-- ``fa_gu_init_vpto`` / ``fa_gu_update_vpto`` are Tile-ABI ``@pto.simd``
+- ``fa_gu_init_vpto`` / ``fa_gu_update_vpto`` are Tile-ABI ``@pto.tileop``
adapters for callers such as ``flash_attention_vf_fusion.py``.
- ``fa_gu_init_vpto_validate`` / ``fa_gu_update_vpto_validate`` are host-visible
launch wrappers for standalone validation.
@@ -105,7 +105,7 @@ def fa_gu_update_vpto_kernel(
pto.vsts(out_vec, o_ptr, row_base + col, mask, dist="NORM_B32")
-@pto.simd
+@pto.tileop
def fa_gu_init_vpto(
pv_tile: pto.Tile,
o_tile: pto.Tile,
@@ -119,7 +119,7 @@ def fa_gu_init_vpto(
)
-@pto.simd
+@pto.tileop
def fa_gu_update_vpto(
o_tile: pto.Tile,
pv_tile: pto.Tile,
diff --git a/ptodsl/examples/flash_attention/softmax.py b/ptodsl/examples/flash_attention/softmax.py
index e32f4709e2..da02c6bd6a 100644
--- a/ptodsl/examples/flash_attention/softmax.py
+++ b/ptodsl/examples/flash_attention/softmax.py
@@ -14,13 +14,13 @@
- ``fa_softmax_init_vpto_kernel`` / ``fa_softmax_update_vpto_kernel``:
ptr-ABI VPTO child modules intended to become separate backend objects
- ``fa_softmax_init_vpto`` / ``fa_softmax_update_vpto``:
- Tile-ABI ``@pto.simd`` adapters that materialize ``as_ptr()`` internally
+ Tile-ABI ``@pto.tileop`` adapters that materialize ``as_ptr()`` internally
- ``fa_softmax_vpto_probe``: minimal entry wrapper for compile-only inspection
The intended structure is:
- auto-mode callers only see Tile arguments
-- the ``@pto.simd`` adapter bridges Tile -> ptr
+- the ``@pto.tileop`` adapter bridges Tile -> ptr
- the explicit VPTO kernel module owns the micro-instruction body
"""
@@ -225,7 +225,7 @@ def fa_softmax_update_vpto_kernel(
pto.vsts(exp_scale, exp_scale_ptr, row, one32, dist="1PT_B32")
-@pto.simd
+@pto.tileop
def fa_softmax_init_vpto(
qk: pto.Tile,
p_nz: pto.Tile,
@@ -245,7 +245,7 @@ def fa_softmax_init_vpto(
)
-@pto.simd
+@pto.tileop
def fa_softmax_update_vpto(
qk: pto.Tile,
p_nz: pto.Tile,
diff --git a/ptodsl/examples/flash_attention_sketch.py b/ptodsl/examples/flash_attention_sketch.py
index a2d0019332..53f70fcfe0 100644
--- a/ptodsl/examples/flash_attention_sketch.py
+++ b/ptodsl/examples/flash_attention_sketch.py
@@ -16,8 +16,8 @@
└─ flash_attention_kernel (@pto.jit, mode="explicit")
├─ Tile Ops tile.load / tile.store at the GM↔UB boundary
├─ explicit orchestration mte_load / pipe_barrier / pointer sequencing
- ├─ @pto.cube matrix products (QK^T and P@V)
- ├─ @pto.simd row-wise online softmax
+ ├─ @pto.tileop matrix products (QK^T and P@V, cube-style)
+ ├─ @pto.tileop row-wise online softmax (vector-style)
└─ @pto.simt scalar metadata and output blending
Design rules illustrated here:
@@ -39,9 +39,10 @@
such as ``mte_load`` are used instead of tile ops where needed.
``mte_load`` / ``mte_store`` accept partitions and tiles directly,
deriving strides and burst sizes from the type metadata.
-6. ``simd`` / ``simt`` / ``cube`` are hardware boundaries. They do not expose
- vreg values across the function boundary. Data crosses the boundary through
- UB-backed tiles or typed UB pointers only.
+6. ``tileop`` / ``simt`` are the public helper boundaries. ``@pto.tileop``
+ covers both vector-style and cube-style helper bodies, but transient
+ hardware-local values still do not cross the function boundary. Data crosses
+ the boundary through UB-backed tiles or typed UB pointers only.
7. Named sub-kernels are reusable wherever their parameter contract is
satisfied. This sketch uses the explicit ``@pto.jit(mode="explicit")`` path
because it needs user-ordered DMA and phase barriers; smaller kernels can
@@ -361,11 +362,11 @@ def flash_attention_kernel(
# Boundary contract:
# - Tile arguments are UB-backed or cube-local buffers carrying addressable
# storage.
-# - No vector register escapes a simd function.
+# - No vector register escapes a tile-op helper.
# - No implicit global-memory access happens inside these kernels.
-@pto.cube
+@pto.tileop
def qk_matmul(
q_mat: pto.Tile, # MAT, [Br, dim]
k_mat: pto.Tile, # MAT, [Bc, dim]
@@ -392,7 +393,7 @@ def qk_matmul(
pto.mte_l0c_ub(s_acc.as_ptr(), s_tile.as_ptr(), m, n, n, n, 0)
-@pto.cube
+@pto.tileop
def pv_matmul(
p_mat: pto.Tile, # MAT, [Br, Bc]
v_mat: pto.Tile, # MAT, [Bc, dim]
@@ -417,7 +418,7 @@ def pv_matmul(
pto.mte_l0c_ub(pv_acc.as_ptr(), pv_tile.as_ptr(), m, n, n, n, 0)
-@pto.simd
+@pto.tileop
def online_softmax_rows(
s_tile: pto.Tile, # UB, [Br, Bc]
p_tile: pto.Tile, # UB, [Br, Bc], output
@@ -663,7 +664,7 @@ def kv_block_process(
# │ │
# │ Key idea: one place owns the "how this block runs on hardware" story. │
# ├──────────────────────────────────────────────────────────────────────────┤
-# │ @pto.cube Matrix-product kernels │
+# │ @pto.tileop Matrix-product kernels (cube-style bodies) │
# │ │
# │ qk_matmul: Q @ K^T │
# │ pv_matmul: P @ V │
@@ -671,7 +672,7 @@ def kv_block_process(
# │ │
# │ Key idea: UB tiles are inputs/outputs; cube-local state is explicit. │
# ├──────────────────────────────────────────────────────────────────────────┤
-# │ @pto.simd Row-wise vector math │
+# │ @pto.tileop Row-wise vector math │
# │ │
# │ online_softmax_rows │
# │ vreg stays local; persistent state is written back to UB tiles │
diff --git a/ptodsl/examples/flash_attention_softmax_launch.py b/ptodsl/examples/flash_attention_softmax_launch.py
index a4a2fc51ef..532a4f7ac1 100644
--- a/ptodsl/examples/flash_attention_softmax_launch.py
+++ b/ptodsl/examples/flash_attention_softmax_launch.py
@@ -105,7 +105,7 @@ def kernel(
pto.set_flag("MTE2", "V", event_id=0)
pto.wait_flag("MTE2", "V", event_id=0)
- with pto.simd():
+ with pto.tileop():
remaining_rows = runtime_rows
for row_base in range(0, runtime_rows, lane_num):
active_rows, remaining_after_pack = pto.make_mask(pto.f32, remaining_rows)
diff --git a/ptodsl/examples/inverse_block_inversion_launch.py b/ptodsl/examples/inverse_block_inversion_launch.py
index c6658b9125..faef2f8a12 100644
--- a/ptodsl/examples/inverse_block_inversion_launch.py
+++ b/ptodsl/examples/inverse_block_inversion_launch.py
@@ -93,7 +93,7 @@ def inverse_block_inversion_f32(
pto.set_flag("MTE2", "V", event_id=0)
pto.wait_flag("MTE2", "V", event_id=0)
- with pto.simd():
+ with pto.tileop():
active, _ = pto.make_mask(pto.f32, batch_i32)
d00 = pto.vlds(in_tile[0, 0:])
d10 = pto.vlds(in_tile[2, 0:])
diff --git a/ptodsl/examples/mixed_backend_kernel_module.py b/ptodsl/examples/mixed_backend_kernel_module.py
index 59361886d3..7c672cc757 100644
--- a/ptodsl/examples/mixed_backend_kernel_module.py
+++ b/ptodsl/examples/mixed_backend_kernel_module.py
@@ -44,7 +44,7 @@ def scale_row_kernel_module(
base_gm: pto.ptr(pto.f32, "gm"),
row: pto.i32,
):
- with pto.simd():
+ with pto.tileop():
c0_i64 = pto.const(0, dtype=pto.i64)
row_offset = row * _ROW_ELEMS
row_gm = pto.addptr(base_gm, row_offset)
diff --git a/ptodsl/examples/softmax_dsl.py b/ptodsl/examples/softmax_dsl.py
index 2b0f79e87a..e5f566b40e 100644
--- a/ptodsl/examples/softmax_dsl.py
+++ b/ptodsl/examples/softmax_dsl.py
@@ -71,7 +71,7 @@ def kernel(
pto.set_flag("MTE2", "V", event_id=0)
pto.wait_flag("MTE2", "V", event_id=0)
- with pto.simd():
+ with pto.tileop():
remaining_rows = runtime_rows
for row_base in range(0, runtime_rows, packed_rows):
active_rows, remaining_rows = pto.make_mask(pto.f32, remaining_rows)
diff --git a/ptodsl/examples/tadd_dsl.py b/ptodsl/examples/tadd_dsl.py
index 3248643a52..5c7deaf527 100644
--- a/ptodsl/examples/tadd_dsl.py
+++ b/ptodsl/examples/tadd_dsl.py
@@ -18,7 +18,7 @@
%c0_i64 = arith.constant 0 : i64 # pto.const(0, dtype=pto.int64)
%c16 = arith.constant 16 : index # pto.const(16, dtype=pto.index)
…
- pto.simd { # with pto.simd():
+ pto.tileop { # with pto.tileop():
%0 = pto.castptr %c4096_i64 … # pto.castptr(c4096_i64, …)
scf.for %arg0 = %c0 to %c16 … { # for i in range(c0, c16, c1):
%mask, _ = pto.plt_b32 … # pto.plt_b32(c64_i32)
@@ -43,7 +43,7 @@ def TADD():
c64_i32 = pto.const(64, dtype=pto.int32)
c64 = pto.const(64)
- with pto.simd():
+ with pto.tileop():
ptr_f32_ub = pto.ptr(pto.float32, "ub")
vf32 = pto.vreg_type(64, pto.float32)
ptr_src = pto.castptr(c4096_i64, ptr_f32_ub)
diff --git a/ptodsl/examples/tilelang_codegen.py b/ptodsl/examples/tilelang_codegen.py
index 86949c9c88..aed5a344e2 100644
--- a/ptodsl/examples/tilelang_codegen.py
+++ b/ptodsl/examples/tilelang_codegen.py
@@ -63,7 +63,7 @@ def _tilelang_generated_body(
pto.set_flag("MTE2", "V", event_id=iter % 2)
pto.wait_flag("MTE2", "V", event_id=iter % 2)
pto.wait_flag("MTE3", "V", event_id=iter % 2)
- with pto.simd():
+ with pto.tileop():
mask_cnt = 8192
with pto.for_(0, 128, step=1) as i:
mask = pto.pset_b32("PAT_ALL")
@@ -157,7 +157,7 @@ def _tilelang_generated_body_small(A, B, C):
pto.set_flag("MTE2", "V", event_id=iter % 2)
pto.wait_flag("MTE2", "V", event_id=iter % 2)
pto.wait_flag("MTE3", "V", event_id=iter % 2)
- with pto.simd():
+ with pto.tileop():
with pto.for_(0, 2, step=1) as i:
mask = pto.pset_b32("PAT_ALL")
r0 = pto.vlds(
diff --git a/ptodsl/ptodsl/_diagnostics.py b/ptodsl/ptodsl/_diagnostics.py
index 92226e000d..957b0e8429 100644
--- a/ptodsl/ptodsl/_diagnostics.py
+++ b/ptodsl/ptodsl/_diagnostics.py
@@ -359,14 +359,25 @@ def inline_subkernel_value_escape_error(role: str, type_text: str) -> RuntimeErr
)
-def simd_value_escape_error(type_text: str) -> RuntimeError:
+def simd_value_escape_error(type_text: str, *, surface: str = "@pto.simd") -> RuntimeError:
"""Return one diagnostic for transient SIMD values escaping a simd subkernel boundary."""
return RuntimeError(
- f"@pto.simd cannot return transient SIMD values across the subkernel boundary "
+ f"{surface} cannot return transient SIMD values across the subkernel boundary "
f"(got {type_text}). Write the value back to a Tile/UB buffer instead."
)
+def subkernel_return_boundary_error(role: str, observed: object) -> TypeError:
+ """Return one diagnostic for unsupported PTODSL subkernel return values."""
+ return TypeError(
+ f"@pto.{role} return values must be PTO scalar values or tuples/lists/dicts of PTO scalar values. "
+ f"Got {observed!r}. Return Tile/TensorView/PartitionTensorView/ptr data through explicit "
+ "subkernel operands instead."
+ )
+
+
+
+
def tile_row_alignment_error(*, shape, dtype, row_bytes: int, required_alignment: int) -> TypeError:
"""Return one diagnostic for authored tile shapes violating row-byte alignment."""
return TypeError(
@@ -433,19 +444,28 @@ def invalid_jit_backend_error(
)
+def legacy_subkernel_surface_error(surface: str) -> TypeError:
+ """Return one diagnostic for removed legacy subkernel public surfaces."""
+ return TypeError(
+ f"{surface} is a legacy PTODSL subkernel surface and is no longer supported. "
+ "Use @pto.tileop for named custom-op helpers, or inline custom-op code with "
+ "`with pto.tileop():`. Keep @pto.simt only for launched SIMT helpers."
+ )
+
+
def unsupported_public_surface_error(name: str) -> AttributeError:
"""Return one diagnostic for unsupported names on the public ``pto`` surface."""
hints = {
"ukernel": (
'Use @pto.jit(mode="explicit") for explicit DMA orchestration, and call or inline '
- "@pto.simd/@pto.simt/@pto.cube directly from that kernel."
+ "@pto.tileop helpers directly from that kernel."
),
"tile_buf_type": (
"Use pto.alloc_tile(shape=..., dtype=..., memory_space=..., valid_shape=..., addr=...) "
"to author tiles, and keep explicit tile-type construction inside internal implementation code only."
),
"vecscope": (
- "Use @pto.simd for named SIMD helpers, or inline SIMD code with `with pto.simd():`."
+ "Use @pto.tileop for named custom OP helpers, or inline custom OP code with `with pto.tileop():`."
),
"as_ptr": (
"Use tile.as_ptr(), view.as_ptr(), or partition.as_ptr() on the authored object itself "
@@ -505,6 +525,7 @@ def unsupported_public_surface_error(name: str) -> AttributeError:
"jit_source_file_error",
"invalid_jit_mode_error",
"invalid_jit_backend_error",
+ "legacy_subkernel_surface_error",
"jit_legacy_tensor_spec_helper_error",
"native_python_control_flow_error",
"simd_value_escape_error",
diff --git a/ptodsl/ptodsl/_subkernels.py b/ptodsl/ptodsl/_subkernels.py
index 1777177074..2cb228351d 100644
--- a/ptodsl/ptodsl/_subkernels.py
+++ b/ptodsl/ptodsl/_subkernels.py
@@ -17,17 +17,19 @@
from ._diagnostics import (
illegal_inline_subkernel_placement_error,
illegal_subkernel_placement_error,
+ legacy_subkernel_surface_error,
simd_value_escape_error,
subkernel_argument_type_error,
subkernel_host_tensor_boundary_error,
subkernel_illegal_annotation_error,
subkernel_illegal_parameter_kind_error,
subkernel_missing_annotation_error,
+ subkernel_return_boundary_error,
subkernel_signature_boundary_error,
)
from ._ast_rewrite import rewrite_jit_function
from ._host_tensors import TensorSpec, looks_like_host_tensor
-from ._surface_types import Tile
+from ._surface_types import PartitionTensorView, TensorView, Tile
from ._surface_values import unwrap_surface_value
from ._tracing import current_runtime, current_session
from ._types import (
@@ -59,9 +61,11 @@
ui32,
ui64,
)
+from mlir.ir import FloatType, IndexType, IntegerType
class KernelRole(str, Enum):
+ TILEOP = "tileop"
CUBE = "cube"
SIMD = "simd"
SIMT = "simt"
@@ -188,11 +192,15 @@ def _validate_invocation(self, *args, **kwargs) -> None:
raise subkernel_host_tensor_boundary_error(self.spec.role.value, name)
def _validate_result(self, result) -> None:
- if self.spec.role != KernelRole.SIMD:
+ if self.spec.role in {KernelRole.TILEOP, KernelRole.SIMD}:
+ escaped_type = _find_transient_simd_escape(result)
+ if escaped_type is not None:
+ raise simd_value_escape_error(escaped_type, surface=f"@pto.{self.spec.role.value}")
+ _validate_subkernel_scalar_result(self.spec.role.value, result)
+ return
+ if self.spec.role == KernelRole.CUBE:
+ _validate_subkernel_scalar_result(self.spec.role.value, result)
return
- escaped_type = _find_transient_simd_escape(result)
- if escaped_type is not None:
- raise simd_value_escape_error(escaped_type)
class _SimtLaunchTemplate:
@@ -233,6 +241,38 @@ def _find_transient_simd_escape(value):
return None
+def _is_scalar_result_type(type_obj) -> bool:
+ return (
+ IndexType.isinstance(type_obj)
+ or IntegerType.isinstance(type_obj)
+ or FloatType.isinstance(type_obj)
+ )
+
+
+def _validate_subkernel_scalar_result(role: str, value) -> None:
+ if value is None:
+ return
+ if isinstance(value, tuple):
+ for item in value:
+ _validate_subkernel_scalar_result(role, item)
+ return
+ if isinstance(value, list):
+ for item in value:
+ _validate_subkernel_scalar_result(role, item)
+ return
+ if isinstance(value, dict):
+ for item in value.values():
+ _validate_subkernel_scalar_result(role, item)
+ return
+ raw_value = unwrap_surface_value(value)
+ type_obj = getattr(raw_value, "type", None)
+ if type_obj is None:
+ raise subkernel_return_boundary_error(role, type(value).__name__)
+ if _is_scalar_result_type(type_obj):
+ return
+ raise subkernel_return_boundary_error(role, str(type_obj))
+
+
def _is_supported_runtime_scalar_annotation(annotation) -> bool:
return (
isinstance(annotation, _DType)
@@ -292,6 +332,10 @@ def _normalize_subkernel_annotation(annotation):
text = annotation.strip()
if text in {"Tile", "pto.Tile"}:
return Tile
+ if text in {"TensorView", "pto.TensorView"}:
+ return TensorView
+ if text in {"PartitionTensorView", "pto.PartitionTensorView"}:
+ return PartitionTensorView
if text in _POSTPONED_DTYPE_ANNOTATIONS:
return _POSTPONED_DTYPE_ANNOTATIONS[text]
if text.startswith("pto.ptr(") and text.endswith(")"):
@@ -301,11 +345,17 @@ def _normalize_subkernel_annotation(annotation):
return annotation
+def _allows_view_annotations(role: KernelRole) -> bool:
+ return role in {KernelRole.TILEOP, KernelRole.SIMD}
+
+
def _is_supported_subkernel_annotation(role: KernelRole, annotation) -> bool:
if annotation is Tile:
return True
if role == KernelRole.CUBE:
return False
+ if _allows_view_annotations(role) and annotation in {TensorView, PartitionTensorView}:
+ return True
if _is_supported_runtime_scalar_annotation(annotation):
return True
if role == KernelRole.SIMT and isinstance(annotation, _PtrDescriptor):
@@ -316,8 +366,11 @@ def _is_supported_subkernel_annotation(role: KernelRole, annotation) -> bool:
def _expected_subkernel_annotation_summary(role: KernelRole) -> str:
if role == KernelRole.CUBE:
return "pto.Tile parameters only"
- if role == KernelRole.SIMD:
- return "pto.Tile parameters plus PTO scalar annotations such as pto.i32/pto.f32"
+ if role in {KernelRole.TILEOP, KernelRole.SIMD}:
+ return (
+ "pto.Tile / pto.TensorView / pto.PartitionTensorView parameters plus PTO scalar "
+ "annotations such as pto.i32/pto.f32"
+ )
return "pto.Tile parameters, typed pto.ptr(...) values, and PTO scalar annotations"
@@ -347,12 +400,47 @@ def _is_runtime_scalar_value(value) -> bool:
)
+def _is_tensor_view_value(value) -> bool:
+ raw_value = unwrap_surface_value(value)
+ type_obj = getattr(raw_value, "type", None)
+ if type_obj is None:
+ return False
+ return str(type_obj).startswith("!pto.tensor_view<")
+
+
+def _is_partition_tensor_view_value(value) -> bool:
+ raw_value = unwrap_surface_value(value)
+ type_obj = getattr(raw_value, "type", None)
+ if type_obj is None:
+ return False
+ return str(type_obj).startswith("!pto.partition_tensor_view<")
+
+
def _normalize_subkernel_argument(role: KernelRole, name: str, annotation, value):
if annotation is Tile:
if isinstance(value, Tile):
return value
raise subkernel_argument_type_error(role.value, name, "a pto.Tile value", type(value).__name__)
+ if annotation is TensorView:
+ if _allows_view_annotations(role) and isinstance(value, TensorView) and _is_tensor_view_value(value):
+ return value
+ raise subkernel_argument_type_error(role.value, name, "a pto.TensorView value", type(value).__name__)
+
+ if annotation is PartitionTensorView:
+ if (
+ _allows_view_annotations(role)
+ and isinstance(value, PartitionTensorView)
+ and _is_partition_tensor_view_value(value)
+ ):
+ return value
+ raise subkernel_argument_type_error(
+ role.value,
+ name,
+ "a pto.PartitionTensorView value",
+ type(value).__name__,
+ )
+
if _is_supported_runtime_scalar_annotation(annotation):
if isinstance(value, (bool, int, float)):
from ._ops import const
@@ -415,6 +503,8 @@ def __init__(
self._session_cm = None
def __call__(self, fn):
+ if self._role in {KernelRole.CUBE, KernelRole.SIMD}:
+ raise legacy_subkernel_surface_error(f"@pto.{self._role.value}")
return SubkernelTemplate(
SubkernelSpec(
role=self._role,
@@ -428,6 +518,8 @@ def __call__(self, fn):
)
def __enter__(self):
+ if self._role in {KernelRole.CUBE, KernelRole.SIMD}:
+ raise legacy_subkernel_surface_error(f"with pto.{self._role.value}()")
if self._role == KernelRole.SIMT and (
self._simt_max_threads is not None or self._simt_max_regs is not None
):
@@ -505,6 +597,10 @@ def _decorate_subkernel(
)
+def tileop(fn=None, *, name: str | None = None, target: str = "a5", ast_rewrite: bool = True):
+ return _decorate_subkernel(KernelRole.TILEOP, fn, name=name, target=target, ast_rewrite=ast_rewrite)
+
+
def cube(fn=None, *, name: str | None = None, target: str = "a5", ast_rewrite: bool = True):
return _decorate_subkernel(KernelRole.CUBE, fn, name=name, target=target, ast_rewrite=ast_rewrite)
@@ -551,6 +647,7 @@ def simt(
"KernelRole",
"SubkernelSpec",
"SubkernelTemplate",
+ "tileop",
"cube",
"simd",
"simt",
diff --git a/ptodsl/ptodsl/_tracing/runtime.py b/ptodsl/ptodsl/_tracing/runtime.py
index e11740c98a..bb41b0b684 100644
--- a/ptodsl/ptodsl/_tracing/runtime.py
+++ b/ptodsl/ptodsl/_tracing/runtime.py
@@ -63,7 +63,7 @@ def finalize_session(self, session):
def dispatch_subkernel_call(self, subkernel, *args, **kwargs):
"""Dispatch a decorated PTODSL subkernel call in the active trace."""
session = require_active_session(f"@pto.{subkernel.spec.role.value}")
- if subkernel.spec.role.value in {"cube", "simd", "simt"}:
+ if subkernel.spec.role.value in {"tileop", "cube", "simd", "simt"}:
return session.lower_helper_subkernel(subkernel, *args, **kwargs)
return subkernel.emit_body(*args, **kwargs)
diff --git a/ptodsl/ptodsl/_tracing/session.py b/ptodsl/ptodsl/_tracing/session.py
index 8f81836996..6545dee66d 100644
--- a/ptodsl/ptodsl/_tracing/session.py
+++ b/ptodsl/ptodsl/_tracing/session.py
@@ -36,6 +36,7 @@
IntegerType,
Operation,
StringAttr,
+ TypeAttr,
UnitAttr,
)
@@ -154,6 +155,7 @@ def __init__(self, module_spec, module, entry_function):
self._carry_loop_stack = []
self._inline_subkernel_counter = 0
self._escaped_inline_values: dict[object, tuple[str, str]] = {}
+ self._helper_result_templates: dict[tuple[str, tuple], object] = {}
@property
def current_function(self):
@@ -232,16 +234,15 @@ def _next_inline_subkernel_symbol(self, base_symbol_name: str) -> str:
return f"{base_symbol_name}_{suffix}"
def _create_subkernel_section_op(self, role: str):
- if role == "simd":
- return _pto.SectionVectorOp()
- if role == "cube":
- return _pto.SectionCubeOp()
return None
+ def _canonical_helper_role(self, role: str) -> str:
+ if role in {"tileop", "simd", "cube"}:
+ return "tileop"
+ return role
+
def _create_inline_subkernel_wrapper(self, role: str):
- wrapper_op = None
- if self._subkernel_section_policy(role) != "function_kind":
- wrapper_op = self._create_subkernel_section_op(role)
+ wrapper_op = self._create_subkernel_section_op(role)
if wrapper_op is None:
wrapper_op = _pto.VecScopeOp()
body_block = wrapper_op.body.blocks.append()
@@ -272,28 +273,13 @@ def _subkernel_section_policy(self, role: str) -> str:
def _subkernel_helper_attributes(self, role: str) -> tuple[tuple[str, object], ...]:
attrs: list[tuple[str, object]] = []
- if role in {"simd", "cube"}:
- attrs.append(("pto.ptodsl.subkernel_helper", StringAttr.get(role)))
- if self._subkernel_section_policy(role) == "function_kind":
- attrs.append(
- (
- "pto.kernel_kind",
- Attribute.parse(
- f"#pto.kernel_kind<{self._subkernel_role_kernel_kind(role)}>"
- ),
- )
- )
+ helper_role = self._canonical_helper_role(role)
+ if helper_role == "tileop":
+ attrs.append(("pto.tileop.helper", UnitAttr.get()))
if role == "simt":
attrs.append(("pto.simt_entry", UnitAttr.get()))
return tuple(attrs)
- def _emit_simt_helper_launch_metadata(self) -> None:
- i32 = IntegerType.get_signless(32)
- dim_z = arith.ConstantOp(i32, 1).result
- dim_y = arith.ConstantOp(i32, 1).result
- dim_x = arith.ConstantOp(i32, 1).result
- _pto.StoreVfSimtInfoOp(dim_z, dim_y, dim_x)
-
def _erase_attached_op(self, op_view) -> None:
parent = op_view.operation.parent
if parent is not None:
@@ -419,6 +405,37 @@ def _note_escaped_inline_values(self, values, *, role: str) -> None:
for value in values:
self._escaped_inline_values[value] = (role, str(value.type))
+ def _flatten_helper_result_templates(self, value) -> tuple:
+ if value is None:
+ return ()
+ if isinstance(value, tuple):
+ flattened = []
+ for item in value:
+ flattened.extend(self._flatten_helper_result_templates(item))
+ return tuple(flattened)
+ if isinstance(value, list):
+ flattened = []
+ for item in value:
+ flattened.extend(self._flatten_helper_result_templates(item))
+ return tuple(flattened)
+ if isinstance(value, dict):
+ flattened = []
+ for item in value.values():
+ flattened.extend(self._flatten_helper_result_templates(item))
+ return tuple(flattened)
+ return (value,)
+
+ def _wrap_helper_call_results(self, template, results_iter):
+ if template is None:
+ return None
+ if isinstance(template, tuple):
+ return tuple(self._wrap_helper_call_results(item, results_iter) for item in template)
+ if isinstance(template, list):
+ return [self._wrap_helper_call_results(item, results_iter) for item in template]
+ if isinstance(template, dict):
+ return {name: self._wrap_helper_call_results(item, results_iter) for name, item in template.items()}
+ return wrap_like_surface_value(template, next(results_iter))
+
def _remap_captured_operands(self, root_ops, capture_mapping) -> None:
for op_view in self._walk_op_tree(root_ops):
operands = op_view.operation.operands
@@ -429,11 +446,7 @@ def _remap_captured_operands(self, root_ops, capture_mapping) -> None:
def _outline_inline_subkernel(self, outline_frame: InlineSubkernelOutlineFrame) -> None:
role = outline_frame.trace_frame.role
- section_policy = self._subkernel_section_policy(role)
- if role in {"simd", "cube"} and section_policy != "function_kind":
- root_ops = (outline_frame.wrapper_op,)
- else:
- root_ops = tuple(outline_frame.body_block.operations)
+ root_ops = tuple(outline_frame.body_block.operations)
defined_values = self._collect_defined_values(root_ops)
captures = self._collect_capture_values(root_ops)
@@ -453,23 +466,20 @@ def _outline_inline_subkernel(self, outline_frame: InlineSubkernelOutlineFrame)
with InsertionPoint(outline_frame.wrapper_op.operation):
if role == "simt":
- self._emit_simt_helper_launch_metadata()
- func.CallOp(helper_fn, list(captures))
+ self._emit_simt_launch_call(helper_fn, captures, dims=(1, 1, 1))
+ else:
+ func.CallOp(helper_fn, list(captures))
entry_block = helper_fn.add_entry_block()
with InsertionPoint(entry_block):
terminator = func.ReturnOp([])
return_anchor = terminator.operation.opview
- if role in {"simd", "cube"} and section_policy != "function_kind":
- outline_frame.wrapper_op.move_before(return_anchor)
- outlined_roots = (outline_frame.wrapper_op,)
- else:
- body_ops = tuple(outline_frame.body_block.operations)
- for op_view in body_ops:
- op_view.move_before(return_anchor)
- outline_frame.wrapper_op.operation.erase()
- outlined_roots = body_ops
+ body_ops = tuple(outline_frame.body_block.operations)
+ for op_view in body_ops:
+ op_view.move_before(return_anchor)
+ outline_frame.wrapper_op.operation.erase()
+ outlined_roots = body_ops
capture_mapping = dict(zip(captures, entry_block.arguments))
self._remap_captured_operands(outlined_roots, capture_mapping)
@@ -483,11 +493,13 @@ def lower_helper_subkernel(self, subkernel, *args, **kwargs):
arg_templates = tuple(args)
arg_types = tuple(unwrap_surface_value(arg).type for arg in arg_templates)
owner_symbol_name = self.current_function_owner_symbol_name
+ result_template = None
helper_spec = HelperFunctionSpec(
symbol_name=subkernel.spec.symbol_name,
arg_types=arg_types,
attributes=self._subkernel_helper_attributes(subkernel.spec.role.value),
)
+ helper_cache_key = (owner_symbol_name, helper_spec.cache_key())
helper_fn, created = self.get_or_create_helper_function(
helper_spec,
owner_symbol_name=owner_symbol_name,
@@ -505,10 +517,21 @@ def lower_helper_subkernel(self, subkernel, *args, **kwargs):
InsertionPoint(entry_block),
):
with self.enter_subkernel(subkernel):
- subkernel.emit_body(*wrapped_args, **kwargs)
- func.ReturnOp([])
+ result_template = subkernel.emit_body(*wrapped_args, **kwargs)
+ flat_results = self._flatten_helper_result_templates(result_template)
+ result_types = [unwrap_surface_value(value).type for value in flat_results]
+ helper_fn.operation.attributes["function_type"] = TypeAttr.get(
+ func.FunctionType.get(list(arg_types), result_types)
+ )
+ func.ReturnOp([unwrap_surface_value(value) for value in flat_results])
+ self._helper_result_templates[helper_cache_key] = result_template
+ else:
+ result_template = self._helper_result_templates.get(helper_cache_key)
- func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_templates])
+ call_op = func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_templates])
+ if result_template is None:
+ return None
+ return self._wrap_helper_call_results(result_template, iter(call_op.results))
def begin_carry_loop(self, start, stop, step, state_items):
"""Materialize one authored ``pto.for_(...).carry(...)`` loop body."""
@@ -536,12 +559,14 @@ def lower_simt_helper_subkernel(self, subkernel, *args, **kwargs):
"""Lower one ``@pto.simt`` call through a dedicated helper function."""
helper_fn, arg_templates = self._get_or_create_simt_helper_function(subkernel, *args, **kwargs)
- self._emit_simt_helper_launch_metadata()
- func.CallOp(helper_fn, [unwrap_surface_value(arg) for arg in arg_templates])
+ self._emit_simt_launch_call(helper_fn, arg_templates, dims=(1, 1, 1))
def lower_simt_launch_subkernel(self, subkernel, *args, dims, **kwargs):
"""Lower one explicit ``pto.simt_launch`` call through a SIMT helper."""
helper_fn, arg_templates = self._get_or_create_simt_helper_function(subkernel, *args, **kwargs)
+ self._emit_simt_launch_call(helper_fn, arg_templates, dims=dims)
+
+ def _emit_simt_launch_call(self, helper_fn, arg_templates, *, dims) -> None:
dim_x, dim_y, dim_z = _coerce_simt_launch_dims(dims)
Operation.create(
"pto.simt_launch",
diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py
index 9b6caffa4a..2050323532 100644
--- a/ptodsl/ptodsl/pto.py
+++ b/ptodsl/ptodsl/pto.py
@@ -145,7 +145,7 @@
# ── Decorator ─────────────────────────────────────────────────────────────────
from ._jit import jit, KernelHandle, merge_jit_modules # noqa: F401
-from ._subkernels import cube, simd, simt # noqa: F401
+from ._subkernels import tileop, cube, simd, simt # noqa: F401
from ._pipe_namespace import pipe # noqa: F401
# ── Shorthand dtype aliases ───────────────────────────────────────────────────
diff --git a/ptodsl/tests/support/docs_fragment_fixtures.py b/ptodsl/tests/support/docs_fragment_fixtures.py
index db833700c2..7f8496e654 100644
--- a/ptodsl/tests/support/docs_fragment_fixtures.py
+++ b/ptodsl/tests/support/docs_fragment_fixtures.py
@@ -558,7 +558,7 @@ def tail_vector_pattern_probe(*, BLOCK: pto.const_expr = 128):
{SNIPPET_PLACEHOLDER}
"""
),
- "tail.simd_helper": _fixture(
+ "tail.tileop_helper": _fixture(
f"""
{SNIPPET_PLACEHOLDER}
@@ -607,12 +607,12 @@ def kernel_entry_explicit_signature_probe(
),
"kernel_entry.explicit_body": _fixture(
f"""
- @pto.cube
+ @pto.tileop
def qk_matmul(q_tile: pto.Tile, k_tile: pto.Tile, s_tile: pto.Tile):
return
- @pto.simd
+ @pto.tileop
def online_softmax(s_tile: pto.Tile, o_tile: pto.Tile, rows: pto.i32, cols: pto.i32):
return
@@ -685,7 +685,7 @@ def kernel_entry_cube_signature_probe(
my_cube_kernel(input_tile, output_tile, left_scratch, right_scratch, acc_scratch)
"""
),
- "kernel_entry.simd_signature": _fixture(
+ "kernel_entry.tileop_signature": _fixture(
f"""
{SNIPPET_PLACEHOLDER}
@@ -694,10 +694,10 @@ def kernel_entry_cube_signature_probe(
def kernel_entry_simd_signature_probe(*, BLOCK: pto.const_expr = 128):
input_tile = pto.alloc_tile(shape=[1, BLOCK], dtype=pto.f32)
output_tile = pto.alloc_tile(shape=[1, BLOCK], dtype=pto.f32)
- my_simd_kernel(input_tile, output_tile, pto.const(1, dtype=pto.i32), pto.const(BLOCK, dtype=pto.i32))
+ my_tileop_kernel(input_tile, output_tile, pto.const(1, dtype=pto.i32), pto.const(BLOCK, dtype=pto.i32))
"""
),
- "kernel_entry.simd_body": _fixture(
+ "kernel_entry.tileop_body": _fixture(
f"""
{SNIPPET_PLACEHOLDER}
@@ -972,12 +972,12 @@ def data_movement_explicit_dma_probe(
),
"sync_ops.flag_pattern_explicit": _fixture(
f"""
- @pto.cube
+ @pto.tileop
def qk_matmul(q_tile: pto.Tile, k_tile: pto.Tile, p_tile: pto.Tile):
return
- @pto.cube
+ @pto.tileop
def pv_matmul(p_tile: pto.Tile, v_tile: pto.Tile, o_tile: pto.Tile):
return
@@ -1021,17 +1021,17 @@ def sync_ops_flag_pattern_explicit_probe(
),
"sync_ops.phase_barrier_explicit": _fixture(
f"""
- @pto.cube
+ @pto.tileop
def qk_matmul(q_tile: pto.Tile, k_tile: pto.Tile, s_tile: pto.Tile):
return
- @pto.simd
+ @pto.tileop
def online_softmax(s_tile: pto.Tile, p_tile: pto.Tile, rows: pto.i32, cols: pto.i32):
return
- @pto.cube
+ @pto.tileop
def pv_matmul(p_tile: pto.Tile, v_tile: pto.Tile, pv_tile: pto.Tile):
return
@@ -1134,7 +1134,7 @@ def data_movement_tile_slice_1d_probe(
),
"data_movement.cube_helper": _fixture(
f"""
- @pto.cube
+ @pto.tileop
def qk_matmul(
q_tile: pto.Tile,
k_tile: pto.Tile,
@@ -1164,7 +1164,7 @@ def data_movement_cube_helper_probe(
),
"compute_ops.vector_compute": _fixture(
f"""
- @pto.simd
+ @pto.tileop
def compute_ops_vector_helper(inp_tile: pto.Tile, out_tile: pto.Tile, row: pto.index):
col_mask = pto.make_mask(pto.f32, pto.const(16, dtype=pto.i32))
s_row = pto.vlds(inp_tile[row, 0:])
@@ -1513,7 +1513,7 @@ def flash_attention_l1_loop_body_probe(
),
"flash_attention.explicit_phase": _fixture(
f"""
- @pto.cube
+ @pto.tileop
def qk_matmul(
q_mat: pto.Tile,
k_mat: pto.Tile,
@@ -1525,7 +1525,7 @@ def qk_matmul(
return
- @pto.simd
+ @pto.tileop
def online_softmax_rows(
s_tile: pto.Tile,
p_tile: pto.Tile,
@@ -1542,7 +1542,7 @@ def online_softmax_rows(
return
- @pto.cube
+ @pto.tileop
def pv_matmul(
p_mat: pto.Tile,
v_mat: pto.Tile,
@@ -1722,7 +1722,7 @@ def flash_attention_inline_simt_scope_probe(*, BLOCK_Q: pto.const_expr = 16, BLO
),
"flash_attention.online_softmax_loop": _fixture(
f"""
- @pto.simd
+ @pto.tileop
def flash_attention_online_softmax_loop_helper(
s_tile: pto.Tile,
p_tile: pto.Tile,
@@ -1761,7 +1761,7 @@ def flash_attention_online_softmax_loop_probe(*, BLOCK: pto.const_expr = 16):
),
"flash_attention.online_softmax_compute": _fixture(
f"""
- @pto.simd
+ @pto.tileop
def flash_attention_online_softmax_compute_helper(
s_tile: pto.Tile,
p_tile: pto.Tile,
@@ -1805,7 +1805,7 @@ def flash_attention_online_softmax_compute_probe(*, BLOCK: pto.const_expr = 16):
),
"flash_attention.online_softmax_store": _fixture(
f"""
- @pto.simd
+ @pto.tileop
def flash_attention_online_softmax_store_helper(
s_tile: pto.Tile,
p_tile: pto.Tile,
@@ -1890,7 +1890,7 @@ def flash_attention_simt_blend_probe(*, BLOCK: pto.const_expr = 8):
),
"gemm.cube_helper": _fixture(
f"""
- @pto.cube
+ @pto.tileop
def gemm_tile(
a_mat: pto.Tile,
b_mat: pto.Tile,
@@ -1915,7 +1915,7 @@ def gemm_tile_probe(*, BLOCK_M: pto.const_expr = 64, BLOCK_K: pto.const_expr = 6
),
"gemm.jit_kernel": _fixture(
f"""
- @pto.cube
+ @pto.tileop
def gemm_tile(
a_mat: pto.Tile,
b_mat: pto.Tile,
diff --git a/ptodsl/tests/test_ast_rewrite_example_ir.py b/ptodsl/tests/test_ast_rewrite_example_ir.py
index 89f869126f..9219337359 100644
--- a/ptodsl/tests/test_ast_rewrite_example_ir.py
+++ b/ptodsl/tests/test_ast_rewrite_example_ir.py
@@ -88,7 +88,7 @@ def kernel():
c64_i32 = pto.const(64, dtype=pto.int32)
c64 = pto.const(64)
- with pto.simd():
+ with pto.tileop():
ptr_f32_ub = pto.ptr(pto.float32, "ub")
vf32 = pto.vreg_type(64, pto.float32)
ptr_src = pto.castptr(c4096_i64, ptr_f32_ub)
@@ -155,7 +155,7 @@ def kernel(
pto.set_flag("MTE2", "V", event_id=0)
pto.wait_flag("MTE2", "V", event_id=0)
- with pto.simd():
+ with pto.tileop():
row_loop = pto.for_(0, runtime_rows, step=packed_rows).carry(remained=runtime_rows)
with row_loop:
row_base = row_loop.iv
@@ -263,7 +263,7 @@ def kernel(
pto.set_flag("MTE2", "V", event_id=0)
pto.wait_flag("MTE2", "V", event_id=0)
- with pto.simd():
+ with pto.tileop():
row_loop = pto.for_(0, runtime_rows, step=lane_num).carry(remained=runtime_rows)
with row_loop:
row_base = row_loop.iv
diff --git a/ptodsl/tests/test_flash_attention_demo_compile.py b/ptodsl/tests/test_flash_attention_demo_compile.py
index ef4d5ee234..046f3ebb01 100644
--- a/ptodsl/tests/test_flash_attention_demo_compile.py
+++ b/ptodsl/tests/test_flash_attention_demo_compile.py
@@ -69,7 +69,7 @@ def main() -> None:
"flash attention wrapper compile should encode the VPTO backend directly on the child module",
)
expect("func.func @materialize_tile_bounds" in wrapper_text, "wrapper compile should emit the SIMT helper function")
- expect("pto.store_vfsimt_info" in wrapper_text, "wrapper compile should materialize SIMT caller metadata setup")
+ expect("pto.simt_launch" in wrapper_text, "wrapper compile should materialize SIMT launch ops")
expect("pto.barrier " in wrapper_text, "demo phase boundaries should lower to pipe_barrier(Pipe.ALL)")
compiled = demo.flash_attention_kernel.compile(
@@ -102,7 +102,7 @@ def main() -> None:
"direct compile should encode the VPTO backend directly on the child module",
)
expect("!pto.tile_buf= 2, "wrapper compile plus explicit compile should populate at least two cached specializations")
diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py
index 557c4b3a2a..4b53c3741f 100644
--- a/ptodsl/tests/test_jit_compile.py
+++ b/ptodsl/tests/test_jit_compile.py
@@ -256,7 +256,7 @@ def process_tile_module(
rows: pto.i32,
cols: pto.i32,
):
- with pto.simd():
+ with pto.tileop():
vec = pto.elements_per_vreg(pto.f32)
initial_remained = cols
with pto.for_(0, rows, step=1) as r:
@@ -278,7 +278,7 @@ def explicit_vpto_kernel_module(
o_tile: pto.Tile,
cols: pto.i32,
):
- with pto.simd():
+ with pto.tileop():
remained = cols
vec = pto.elements_per_vreg(pto.f32)
loop = pto.for_(0, cols, step=vec).carry(remained=remained)
@@ -296,7 +296,7 @@ def process_row_ptr_kernel_module(
dst_gm: pto.ptr(pto.f32, "gm"),
row: pto.i32,
):
- with pto.simd():
+ with pto.tileop():
c0_i64 = pto.const(0, dtype=pto.i64)
row_offset = row * 16
src_row = pto.addptr(src_gm, row_offset)
@@ -379,12 +379,15 @@ def emitc_entry_calls_vpto_kernel_module_probe(
process_row_ptr_kernel_module(A_ptr, O_ptr, row)
-@pto.simd
+@pto.tileop
def emitc_vpto_kernel_module_callsite_simd_helper(
src_tile: pto.Tile,
dst_tile: pto.Tile,
cols: pto.i32,
):
+ mask, _ = pto.make_mask(pto.f32, cols)
+ vec = pto.vlds(src_tile[0, 0:])
+ pto.vsts(vec, dst_tile[0, 0:], mask)
explicit_vpto_kernel_module(src_tile, dst_tile, cols)
@@ -666,57 +669,126 @@ def tile_surface_window_matmul_probe():
INLINE_SUBKERNEL_SCOPE_OBSERVATIONS = []
-@pto.simd
-def nested_simd_probe():
+@pto.tileop
+def nested_tileop_probe():
session = current_session()
frame = session.current_subkernel
SUBKERNEL_OBSERVATIONS.append((frame.role, frame.symbol_name, session.subkernel_stack_depth))
-@pto.cube
+@pto.tileop
def top_level_cube_probe():
session = current_session()
frame = session.current_subkernel
SUBKERNEL_OBSERVATIONS.append((frame.role, frame.symbol_name, session.subkernel_stack_depth))
-@pto.simd
+@pto.tileop
+def top_level_tileop_probe():
+ session = current_session()
+ frame = session.current_subkernel
+ SUBKERNEL_OBSERVATIONS.append((frame.role, frame.symbol_name, session.subkernel_stack_depth))
+
+
+@pto.tileop
def top_level_simd_probe():
session = current_session()
frame = session.current_subkernel
SUBKERNEL_OBSERVATIONS.append((frame.role, frame.symbol_name, session.subkernel_stack_depth))
-@pto.simd
-def explicit_vector_simd_probe():
+@pto.tileop
+def explicit_vector_tileop_probe():
pto.pipe_barrier(pto.Pipe.ALL)
-@pto.cube
-def explicit_vector_cube_probe():
- pto.pipe_barrier(pto.Pipe.ALL)
+@pto.tileop
+def vector_scalar_return_probe(inp_tile: pto.Tile, out_tile: pto.Tile, cols: pto.i32):
+ col_mask = pto.make_mask(pto.f32, cols)
+ row = pto.const(0)
+ s_row = pto.vlds(inp_tile[row, 0:])
+ pto.vsts(s_row, out_tile[row, 0:], col_mask)
+ return cols
+
+
+@pto.tileop
+def tileop_view_boundary_probe(
+ src_view: pto.TensorView,
+ src_part: pto.PartitionTensorView,
+ out_tile: pto.Tile,
+ rows: pto.i32,
+):
+ _ = src_view
+ pto.tile.load(src_part, out_tile)
+ return rows
+
+
+@pto.tileop
+def simd_view_boundary_probe(
+ src_view: pto.TensorView,
+ src_part: pto.PartitionTensorView,
+ out_tile: pto.Tile,
+ rows: pto.i32,
+):
+ _ = src_view
+ pto.tile.load(src_part, out_tile)
+ return rows
+
+
+@pto.jit(target="a5")
+def scalar_return_subkernel_lowering_probe():
+ inp_tile = pto.alloc_tile(shape=[1, 16], dtype=pto.f32, valid_shape=[1, 16])
+ out_tile = pto.alloc_tile(shape=[1, 16], dtype=pto.f32, valid_shape=[1, 16])
+ stats_tile = pto.alloc_tile(shape=[1, 8], dtype=pto.i32, valid_shape=[1, 2])
+ cols = pto.const(16, dtype=pto.i32)
+ returned_cols = vector_scalar_return_probe(inp_tile, out_tile, cols)
+ scalar.store(returned_cols, stats_tile.as_ptr(), 0)
+
+
+@pto.jit(target="a5")
+def tileop_view_boundary_entry_probe(
+ A_ptr: pto.ptr(pto.f32, "gm"),
+ rows: pto.i32,
+ cols: pto.i32,
+):
+ src_view = pto.make_tensor_view(A_ptr, shape=[rows, cols], strides=[cols, 1])
+ src_part = pto.partition_view(src_view, offsets=[0, 0], sizes=[rows, cols])
+ out_tile = pto.alloc_tile(shape=[1, 16], dtype=pto.f32, valid_shape=[rows, cols])
+ stats_tile = pto.alloc_tile(shape=[1, 8], dtype=pto.i32, valid_shape=[1, 1])
+ returned_rows = tileop_view_boundary_probe(src_view, src_part, out_tile, rows)
+ scalar.store(returned_rows, stats_tile.as_ptr(), 0)
+
+
+@pto.jit(target="a5")
+def simd_view_boundary_entry_probe(
+ A_ptr: pto.ptr(pto.f32, "gm"),
+ rows: pto.i32,
+ cols: pto.i32,
+):
+ src_view = pto.make_tensor_view(A_ptr, shape=[rows, cols], strides=[cols, 1])
+ src_part = pto.partition_view(src_view, offsets=[0, 0], sizes=[rows, cols])
+ out_tile = pto.alloc_tile(shape=[1, 16], dtype=pto.f32, valid_shape=[rows, cols])
+ stats_tile = pto.alloc_tile(shape=[1, 8], dtype=pto.i32, valid_shape=[1, 1])
+ returned_rows = simd_view_boundary_probe(src_view, src_part, out_tile, rows)
+ scalar.store(returned_rows, stats_tile.as_ptr(), 0)
@pto.jit(target="a5")
def shared_subkernel_lowering_probe(*, TRACE_TOKEN: pto.const_expr = 0):
top_level_cube_probe()
+ top_level_tileop_probe()
top_level_simd_probe()
- nested_simd_probe()
+ nested_tileop_probe()
@pto.jit(target="a5", kernel_kind="vector")
-def explicit_vector_calls_simd_probe(*, TRACE_TOKEN: pto.const_expr = 0):
- explicit_vector_simd_probe()
+def explicit_vector_calls_tileop_probe(*, TRACE_TOKEN: pto.const_expr = 0):
+ explicit_vector_tileop_probe()
@pto.jit(target="a5", kernel_kind="vector")
-def explicit_vector_calls_cube_probe(*, TRACE_TOKEN: pto.const_expr = 0):
- explicit_vector_cube_probe()
-
-
-@pto.jit(target="a5", kernel_kind="vector")
-def explicit_vector_inline_simd_probe(*, TRACE_TOKEN: pto.const_expr = 0):
- with pto.simd():
+def explicit_vector_inline_tileop_probe(*, TRACE_TOKEN: pto.const_expr = 0):
+ with pto.tileop():
pto.pipe_barrier(pto.Pipe.ALL)
@@ -729,11 +801,15 @@ def inline_subkernel_scope_probe(*, TRACE_TOKEN: pto.const_expr = 0):
frame = session.current_subkernel
INLINE_SUBKERNEL_SCOPE_OBSERVATIONS.append((frame.role, frame.symbol_name, session.subkernel_stack_depth))
scalar.store(0, meta_tile.as_ptr() + 0)
- with pto.simd():
+ with pto.tileop():
frame = session.current_subkernel
INLINE_SUBKERNEL_SCOPE_OBSERVATIONS.append((frame.role, frame.symbol_name, session.subkernel_stack_depth))
pto.pipe_barrier(pto.Pipe.ALL)
- with pto.cube():
+ with pto.tileop():
+ frame = session.current_subkernel
+ INLINE_SUBKERNEL_SCOPE_OBSERVATIONS.append((frame.role, frame.symbol_name, session.subkernel_stack_depth))
+ pto.pipe_barrier(pto.Pipe.ALL)
+ with pto.tileop():
frame = session.current_subkernel
INLINE_SUBKERNEL_SCOPE_OBSERVATIONS.append((frame.role, frame.symbol_name, session.subkernel_stack_depth))
pto.pipe_barrier(pto.Pipe.ALL)
@@ -896,7 +972,7 @@ def simt_invalid_atomic_signedness_probe(gm: pto.ptr(pto.f32, "gm")):
pto.atomic_add(gm, value, signedness="signed")
-@pto.simd
+@pto.tileop
def ast_subkernel_runtime_for_helper(rows: pto.i32):
for row in range(0, rows, 1):
_ = row
@@ -1314,7 +1390,7 @@ def ast_signature_closure_default_kernel(*, BLOCK: pto.const_expr = limit):
def make_ast_rebound_subkernel_probe():
limit = 2
- @pto.simd
+ @pto.tileop
def helper():
for _ in pto.static_range(limit):
pto.pipe_barrier(pto.Pipe.ALL)
@@ -1352,7 +1428,7 @@ def make_sourceless_subkernel_entry():
namespace = {"pto": pto}
exec(
"""
-@pto.simd
+@pto.tileop
def sourceless_subkernel_helper():
if True:
pto.pipe_barrier(pto.Pipe.ALL)
@@ -1498,7 +1574,7 @@ def host_runtime_scalar_entry_probe(
pto.tile.store(o_tile, o_part)
-@pto.simd
+@pto.tileop
def tile_slice_vector_probe(inp_tile: pto.Tile, out_tile: pto.Tile, row: pto.index):
mask, _ = pto.plt_b32(pto.const(64, dtype=pto.i32))
vec = pto.vlds(inp_tile[row, 0:])
@@ -1688,7 +1764,7 @@ def shared_index_coercion_probe():
pto.wait_flag(pto.Pipe.V, pto.Pipe.MTE2, event_id=limit)
-@pto.simd
+@pto.tileop
def public_vector_surface_probe(inp_tile: pto.Tile, out_tile: pto.Tile, stats_tile: pto.Tile):
col_mask = pto.make_mask(pto.f32, pto.const(16, dtype=pto.i32))
row = pto.const(0)
@@ -1702,7 +1778,7 @@ def public_vector_surface_probe(inp_tile: pto.Tile, out_tile: pto.Tile, stats_ti
scalar.store(row_sum, stats_tile[row, 1])
-@pto.cube
+@pto.tileop
def public_cube_surface_probe(
lhs_tile: pto.Tile,
rhs_tile: pto.Tile,
@@ -1793,7 +1869,7 @@ def public_cube_surface_probe(
pto.mte_l0c_ub(acc_tile.as_ptr(), out_tile.as_ptr(), m, n, n, n, split=pto.SplitMode.M, layout="nz2nd")
-@pto.cube
+@pto.tileop
def public_cube_tile_mx_probe(
mat_lhs: pto.Tile,
mat_lhs_scale: pto.Tile,
@@ -3625,8 +3701,9 @@ def inline_source_backed_probe(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32):
"mixed-backend EmitC entry should keep its top-level tile load/store path alongside the kernel-module call",
)
expect(
- mixed_backend_text.count("pto.section.vector {") == 1,
- "before PTOAS inferred normalization, the mixed-backend PTODSL IR should only carry the helper-authored explicit vector section",
+ mixed_backend_text.count("pto.section.vector {") == 0
+ and mixed_backend_text.count("pto.section.cube {") == 0,
+ "before PTOAS inferred normalization, the mixed-backend PTODSL IR should stay naked and let PTOAS infer missing sections later",
)
expect(
"pto.tload" in mixed_backend_text
@@ -3666,19 +3743,23 @@ def inline_source_backed_probe(ptr: pto.ptr(pto.f32, "gm"), rows: pto.i32):
decorated_mixed_backend_text = emitc_entry_calls_vpto_kernel_module_via_decorated_simd_probe.compile().mlir_text()
expect_parse_roundtrip_and_verify(
decorated_mixed_backend_text,
- "emitc entry calling vpto kernel-module through @pto.simd specialization",
+ "emitc entry calling vpto kernel-module through @pto.tileop specialization",
)
+ decorated_helper_body = decorated_mixed_backend_text.split(
+ "func.func @emitc_vpto_kernel_module_callsite_simd_helper__ptodsl_",
+ 1,
+ )[1].split("func.func private @explicit_vpto_kernel_module__ptodsl_", 1)[0]
expect(
re.search(
r"call @emitc_vpto_kernel_module_callsite_simd_helper__ptodsl_[0-9a-f]+"
r"\(%[a-zA-Z0-9_]+, %[a-zA-Z0-9_]+, %[a-zA-Z0-9_]+\)",
decorated_mixed_backend_text,
) is not None,
- "@pto.simd helper callsites should lower to helper function calls in the caller body",
+ "@pto.tileop helper callsites should lower to helper function calls in the caller body",
)
expect(
- "pto.section.vector {" in decorated_mixed_backend_text,
- "the outlined @pto.simd helper body should still materialize one vector section",
+ "pto.section.vector {" not in decorated_helper_body,
+ "decorated @pto.tileop helper bodies should now stay naked in PTODSL IR and rely on later PTOAS section materialization",
)
multi_abi_compiled = entry_calls_kernel_module_multiple_abi_probe.compile()
multi_abi_text = multi_abi_compiled.mlir_text()
@@ -4114,49 +4195,109 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None):
expect_parse_roundtrip_and_verify(shared_subkernel_text, "shared subkernel lowering specialization")
expect(
SUBKERNEL_OBSERVATIONS == [
- ("cube", "top_level_cube_probe", 1),
- ("simd", "top_level_simd_probe", 1),
- ("simd", "nested_simd_probe", 1),
+ ("tileop", "top_level_cube_probe", 1),
+ ("tileop", "top_level_tileop_probe", 1),
+ ("tileop", "top_level_simd_probe", 1),
+ ("tileop", "nested_tileop_probe", 1),
],
f"unexpected shared subkernel lowering observations: {SUBKERNEL_OBSERVATIONS!r}",
)
expect(
re.search(r"call @top_level_cube_probe__ptodsl_[0-9a-f]+\(\)", shared_subkernel_text) is not None
+ and re.search(r"call @top_level_tileop_probe__ptodsl_[0-9a-f]+\(\)", shared_subkernel_text) is not None
and re.search(r"call @top_level_simd_probe__ptodsl_[0-9a-f]+\(\)", shared_subkernel_text) is not None
- and re.search(r"call @nested_simd_probe__ptodsl_[0-9a-f]+\(\)", shared_subkernel_text) is not None,
- "@pto.cube/@pto.simd decorated subkernels should lower to helper calls in the caller body",
+ and re.search(r"call @nested_tileop_probe__ptodsl_[0-9a-f]+\(\)", shared_subkernel_text) is not None,
+ "@pto.tileop decorated subkernels should lower to helper calls in the caller body",
+ )
+ expect(
+ shared_subkernel_text.count("pto.tileop.helper") == 4
+ and 'pto.ptodsl.subkernel_helper = "tileop"' not in shared_subkernel_text
+ and 'pto.ptodsl.subkernel_helper = "simd"' not in shared_subkernel_text
+ and 'pto.ptodsl.subkernel_helper = "cube"' not in shared_subkernel_text,
+ "decorated @pto.tileop helpers should canonicalize to the tileop backend helper role",
)
expect(
- shared_subkernel_text.count("pto.section.vector {") == 2 and "pto.section.cube {" in shared_subkernel_text,
- "outlined decorated helper bodies should still preserve their PTO unit sections",
+ "pto.section.vector {" not in shared_subkernel_text
+ and "pto.section.cube {" not in shared_subkernel_text,
+ "decorated @pto.tileop helpers should now lower as naked tileop helpers without pre-materialized sections",
)
- explicit_vector_simd_text = explicit_vector_calls_simd_probe.compile(TRACE_TOKEN=1).mlir_text()
- expect_parse_roundtrip_and_verify(
- explicit_vector_simd_text,
- "explicit vector jit calling simd subkernel specialization",
+ scalar_return_subkernel_text = scalar_return_subkernel_lowering_probe.compile().mlir_text()
+ expect_parse_roundtrip_and_verify(scalar_return_subkernel_text, "scalar return subkernel lowering specialization")
+ expect(
+ re.search(r"func\.func @vector_scalar_return_probe__ptodsl_[0-9a-f]+\([^)]*\) -> i32", scalar_return_subkernel_text)
+ is not None,
+ "decorated tileop helpers that return PTO scalar values should materialize scalar helper result types",
)
expect(
- "pto.kernel_kind = #pto.kernel_kind" in explicit_vector_simd_text
- and "pto.section.vector {" not in explicit_vector_simd_text,
- "same-kind @pto.simd helpers inside explicit vector kernels should use function/kernel kind metadata without redundant sections",
+ "pto.vlds" in scalar_return_subkernel_text
+ and "pto.vsts" in scalar_return_subkernel_text
+ and re.search(r"%[a-zA-Z0-9_]+ = call @vector_scalar_return_probe__ptodsl_[0-9a-f]+\([^)]*\) : \([^)]*\) -> i32", scalar_return_subkernel_text)
+ is not None,
+ "decorated tileop helper callsites should consume scalar func.call results while preserving primary vector compute in the helper body",
)
- expect_raises(
- RuntimeError,
- lambda: explicit_vector_calls_cube_probe.compile(TRACE_TOKEN=1).mlir_text(),
- "@pto.cube cannot be lowered inside an explicit @pto.jit(kernel_kind='vector')",
+ expect(
+ "pto.store " in scalar_return_subkernel_text,
+ "scalar helper return values should remain usable by later PTODSL scalar stores",
+ )
+
+ tileop_view_boundary_text = tileop_view_boundary_entry_probe.compile().mlir_text()
+ expect_parse_roundtrip_and_verify(tileop_view_boundary_text, "tileop view boundary specialization")
+ expect(
+ re.search(
+ r"func\.func @tileop_view_boundary_probe__ptodsl_[0-9a-f]+\([^)]*!pto\.tensor_view<[^)]*!pto\.partition_tensor_view<[^)]*!pto\.tile_buf<",
+ tileop_view_boundary_text,
+ ) is not None,
+ "decorated tileop helpers should accept TensorView and PartitionTensorView formals in their lowered helper signature",
+ )
+ expect(
+ "pto.tload" in tileop_view_boundary_text
+ and re.search(
+ r"%[a-zA-Z0-9_]+ = call @tileop_view_boundary_probe__ptodsl_[0-9a-f]+\([^)]*\) : \([^)]*!pto\.tensor_view<[^)]*!pto\.partition_tensor_view<[^)]*\) -> i32",
+ tileop_view_boundary_text,
+ ) is not None,
+ "tileop callsites should pass TensorView/PartitionTensorView operands through the helper ABI and preserve scalar returns",
)
- explicit_vector_inline_simd_text = explicit_vector_inline_simd_probe.compile(
+
+ simd_view_boundary_text = simd_view_boundary_entry_probe.compile().mlir_text()
+ expect_parse_roundtrip_and_verify(simd_view_boundary_text, "simd view boundary specialization")
+ expect(
+ re.search(
+ r"func\.func @simd_view_boundary_probe__ptodsl_[0-9a-f]+\([^)]*!pto\.tensor_view<[^)]*!pto\.partition_tensor_view<[^)]*!pto\.tile_buf<",
+ simd_view_boundary_text,
+ ) is not None,
+ "tileop helpers should accept TensorView and PartitionTensorView formals in the same lowered helper ABI",
+ )
+ expect(
+ "pto.tload" in simd_view_boundary_text
+ and re.search(
+ r"%[a-zA-Z0-9_]+ = call @simd_view_boundary_probe__ptodsl_[0-9a-f]+\([^)]*\) : \([^)]*!pto\.tensor_view<[^)]*!pto\.partition_tensor_view<[^)]*\) -> i32",
+ simd_view_boundary_text,
+ ) is not None,
+ "tileop callsites should pass TensorView/PartitionTensorView operands through the same helper ABI and preserve scalar returns",
+ )
+
+ explicit_vector_tileop_text = explicit_vector_calls_tileop_probe.compile(TRACE_TOKEN=1).mlir_text()
+ expect_parse_roundtrip_and_verify(
+ explicit_vector_tileop_text,
+ "explicit vector jit calling tileop subkernel specialization",
+ )
+ expect(
+ "pto.kernel_kind = #pto.kernel_kind" in explicit_vector_tileop_text
+ and "pto.section.vector {" not in explicit_vector_tileop_text,
+ "same-kind @pto.tileop helpers inside explicit vector kernels should use function/kernel kind metadata without redundant sections",
+ )
+ explicit_vector_inline_tileop_text = explicit_vector_inline_tileop_probe.compile(
TRACE_TOKEN=1
).mlir_text()
expect_parse_roundtrip_and_verify(
- explicit_vector_inline_simd_text,
- "explicit vector jit calling inline simd specialization",
+ explicit_vector_inline_tileop_text,
+ "explicit vector jit calling inline tileop specialization",
)
expect(
- "pto.kernel_kind = #pto.kernel_kind" in explicit_vector_inline_simd_text
- and "pto.section.vector {" not in explicit_vector_inline_simd_text,
- "same-kind inline pto.simd() scopes inside explicit vector kernels should avoid redundant sections",
+ "pto.kernel_kind = #pto.kernel_kind" in explicit_vector_inline_tileop_text
+ and "pto.section.vector {" not in explicit_vector_inline_tileop_text,
+ "same-kind inline pto.tileop() scopes inside explicit vector kernels should avoid redundant sections",
)
INLINE_SUBKERNEL_SCOPE_OBSERVATIONS.clear()
@@ -4165,42 +4306,41 @@ def fake_run_ptoas_cmd(cmd, *, cwd=None):
expect(
INLINE_SUBKERNEL_SCOPE_OBSERVATIONS == [
("simt", "inline_simt", 1),
- ("simd", "inline_simd", 1),
- ("cube", "inline_cube", 1),
+ ("tileop", "inline_tileop", 1),
+ ("tileop", "inline_tileop", 1),
+ ("tileop", "inline_tileop", 1),
],
f"unexpected inline subkernel scope observations: {INLINE_SUBKERNEL_SCOPE_OBSERVATIONS!r}",
)
expect(
- inline_subkernel_scope_text.count("pto.store_vfsimt_info") == 1,
- "inline pto.simt() should materialize one caller-side store_vfsimt_info before the helper call",
+ re.search(r"pto\.simt_launch @inline_simt_[0-9]+__ptodsl_[0-9a-f]+<<<", inline_subkernel_scope_text) is not None,
+ "inline pto.simt() should materialize one caller-side pto.simt_launch",
+ )
+ expect(
+ re.search(r"call @inline_tileop_[0-9]+__ptodsl_[0-9a-f]+\([^\\n]*\)", inline_subkernel_scope_text) is not None
+ and len(re.findall(r"call @inline_tileop_[0-9]+__ptodsl_[0-9a-f]+\([^\\n]*\)", inline_subkernel_scope_text)) == 3,
+ "inline pto.tileop() scopes should each lower to one helper call",
)
expect(
- re.search(r"call @inline_simt_[0-9]+__ptodsl_[0-9a-f]+\([^\\n]*\)", inline_subkernel_scope_text) is not None
- and re.search(r"call @inline_simd_[0-9]+__ptodsl_[0-9a-f]+\([^\\n]*\)", inline_subkernel_scope_text) is not None
- and re.search(r"call @inline_cube_[0-9]+__ptodsl_[0-9a-f]+\([^\\n]*\)", inline_subkernel_scope_text) is not None,
- "inline pto.simt()/pto.simd()/pto.cube() scopes should each lower to one helper call",
+ inline_subkernel_scope_text.count("pto.tileop.helper") == 3
+ and 'pto.ptodsl.subkernel_helper = "tileop"' not in inline_subkernel_scope_text
+ and 'pto.ptodsl.subkernel_helper = "simd"' not in inline_subkernel_scope_text
+ and 'pto.ptodsl.subkernel_helper = "cube"' not in inline_subkernel_scope_text,
+ "outlined inline tileop helpers should canonicalize to the tileop backend helper role",
)
expect(
inline_subkernel_scope_text.count("pto.barrier ") >= 2
- and "pto.section.vector {" in inline_subkernel_scope_text
- and "pto.section.cube {" in inline_subkernel_scope_text
+ and "pto.section.vector {" not in inline_subkernel_scope_text
+ and "pto.section.cube {" not in inline_subkernel_scope_text
and "pto.store" in inline_subkernel_scope_text,
- "outlined inline helpers should preserve the authored SIMD/Cube sections and SIMT scalar ops",
+ "outlined inline helpers should lower as naked tileop bodies while preserving SIMT scalar ops",
)
simt_text = simt_helper_lowering_probe.compile(TRACE_TOKEN=1).mlir_text()
expect_parse_roundtrip_and_verify(simt_text, "simt helper lowering specialization")
expect(
- simt_text.count("pto.store_vfsimt_info") == 2,
- "each @pto.simt callsite should materialize a caller-side store_vfsimt_info",
- )
- expect(
- re.search(r"call @simt_tid_probe__simt_\d+\(\)", simt_text) is not None,
- "each @pto.simt callsite should lower to a func.call of the helper symbol",
- )
- expect(
- len(re.findall(r"call @simt_tid_probe__simt_\d+\(\)", simt_text)) == 2,
- "both @pto.simt callsites should call the same helper specialization",
+ len(re.findall(r"pto\.simt_launch @simt_tid_probe__simt_\d+<<<", simt_text)) == 2,
+ "each @pto.simt callsite should materialize a caller-side pto.simt_launch",
)
expect(
len(
@@ -4418,11 +4558,11 @@ def _enter_inline_simt_with_resource_attr():
)
expect(
ast_subkernel_runtime_for_text.count("scf.for") == 1,
- "@pto.simd helper should rewrite Python range(...) loops into runtime scf.for",
+ "@pto.tileop helper should rewrite Python range(...) loops into runtime scf.for",
)
expect(
"pto.barrier " in ast_subkernel_runtime_for_text,
- "rewritten @pto.simd helper body should lower inside the caller trace",
+ "rewritten @pto.tileop helper body should lower inside the caller trace",
)
carry_text = carry_loop_lowering_probe.compile(BLOCK=32).mlir_text()
@@ -4919,8 +5059,8 @@ def _enter_inline_simt_with_resource_attr():
simt_pointer_offset_text = simt_pointer_offset_probe.compile().mlir_text()
expect_parse_roundtrip_and_verify(simt_pointer_offset_text, "simt pointer offset specialization")
expect(
- re.search(r"call @simt_pointer_offset_helper__simt_\d+", simt_pointer_offset_text) is not None,
- "@pto.simt pointer helper should lower to a helper func.call",
+ re.search(r"pto\.simt_launch @simt_pointer_offset_helper__simt_\d+<<<", simt_pointer_offset_text) is not None,
+ "@pto.simt pointer helper should lower to a caller-side pto.simt_launch",
)
expect(
re.search(r"pto\.store %c9_i32, %(?:arg0|\d+)\[%c1(?:_\d+)?\]", simt_pointer_offset_text) is not None,
diff --git a/ptodsl/tests/test_ptoas_frontend_verify.py b/ptodsl/tests/test_ptoas_frontend_verify.py
index 0a25026476..c7072e4252 100644
--- a/ptodsl/tests/test_ptoas_frontend_verify.py
+++ b/ptodsl/tests/test_ptoas_frontend_verify.py
@@ -243,7 +243,7 @@ def process_row_ptr_kernel_module(
dst_gm: pto.ptr(pto.f32, "gm"),
row: pto.i32,
):
- with pto.simd():
+ with pto.tileop():
c0_i64 = pto.const(0, dtype=pto.i64)
row_offset = row * 16
src_row = pto.addptr(src_gm, row_offset)
@@ -472,8 +472,10 @@ def main() -> None:
expect(
"func.func public @scale_row_kernel_module__ptodsl_" in example_vpto_child
and 'pto.visibility = "external"' in example_vpto_child
- and "pto.section.vector {" in example_vpto_child,
- "mixed_backend_kernel_module.py VPTO child should expose a public helper definition with explicit vector authoring, matching the vector-helper side of mixed-external-vadd",
+ and "func.func @inline_tileop_0__ptodsl_" in example_vpto_child
+ and "pto.tileop.helper" in example_vpto_child
+ and "pto.section.vector {" not in example_vpto_child,
+ "mixed_backend_kernel_module.py VPTO child should expose a public wrapper plus a naked tileop helper body, leaving vector section materialization to later PTOAS normalization",
)
example_frontend_texts = run_ptoas_frontend_verify(
diff --git a/ptodsl/tests/test_subkernel_diagnostics.py b/ptodsl/tests/test_subkernel_diagnostics.py
index 34b209aa9a..ddd45330fb 100644
--- a/ptodsl/tests/test_subkernel_diagnostics.py
+++ b/ptodsl/tests/test_subkernel_diagnostics.py
@@ -28,27 +28,35 @@ def expect_raises(callback, exc_type, *message_fragments: str) -> None:
def define_bad_subkernel_signature_probe():
- @pto.simd
+ @pto.tileop
def bad_tensor_formal(A: TensorSpec(rank=2, dtype=pto.f32)):
pto.pipe_barrier(pto.Pipe.ALL)
return bad_tensor_formal
-def define_illegal_simd_ptr_signature_probe():
+def define_legacy_simd_surface_probe():
@pto.simd
- def bad_ptr_formal(meta_ptr: pto.ptr(pto.i32, pto.MemorySpace.UB)):
+ def legacy_simd_probe(tile: pto.Tile):
pto.pipe_barrier(pto.Pipe.ALL)
- return bad_ptr_formal
+ return legacy_simd_probe
-def define_illegal_cube_scalar_signature_probe():
+def define_legacy_cube_surface_probe():
@pto.cube
- def bad_cube_formal(tile: pto.Tile, cols: pto.i32):
+ def legacy_cube_probe(tile: pto.Tile):
pto.pipe_barrier(pto.Pipe.ALL)
- return bad_cube_formal
+ return legacy_cube_probe
+
+
+def define_illegal_tileop_ptr_signature_probe():
+ @pto.tileop
+ def bad_ptr_formal(meta_ptr: pto.ptr(pto.i32, pto.MemorySpace.UB)):
+ pto.pipe_barrier(pto.Pipe.ALL)
+
+ return bad_ptr_formal
def define_removed_ukernel_surface_probe():
@@ -71,7 +79,7 @@ def bad_mode_probe():
return bad_mode_probe
-@pto.simd
+@pto.tileop
def host_tensor_operand_probe(tensor: pto.Tile):
pto.pipe_barrier(pto.Pipe.ALL)
@@ -89,7 +97,7 @@ def nested_simt_probe():
pto.get_tid_x()
-@pto.simd
+@pto.tileop
def illegal_simt_placement_probe():
nested_simt_probe()
@@ -99,7 +107,7 @@ def nested_simt_from_simd_entry(*, TRACE_TOKEN: pto.const_expr = 0):
illegal_simt_placement_probe()
-@pto.simd
+@pto.tileop
def illegal_inline_simt_placement_probe():
with pto.simt():
pto.get_tid_x()
@@ -110,26 +118,37 @@ def nested_inline_simt_from_simd_entry(*, TRACE_TOKEN: pto.const_expr = 0):
illegal_inline_simt_placement_probe()
-@pto.simd
-def simd_value_escape_probe():
+@pto.tileop
+def tileop_value_escape_probe():
return pto.pset_b32("PAT_ALL")
@pto.jit(target="a5")
def simd_value_escape_entry(*, TRACE_TOKEN: pto.const_expr = 0):
- simd_value_escape_probe()
+ tileop_value_escape_probe()
-@pto.simd
+@pto.tileop
def tile_only_probe(inp_tile: pto.Tile):
pto.pipe_barrier(pto.Pipe.ALL)
+@pto.tileop
+def invalid_tileop_return_probe(inp_tile: pto.Tile):
+ return inp_tile
+
+
@pto.jit(target="a5")
def illegal_subkernel_callsite_entry(A_ptr: pto.ptr(pto.f32, "gm")):
tile_only_probe(A_ptr)
+@pto.jit(target="a5")
+def invalid_tileop_return_entry():
+ meta_tile = pto.alloc_tile(shape=[1, 8], dtype=pto.i32, valid_shape=[1, 1])
+ invalid_tileop_return_probe(meta_tile)
+
+
@pto.jit(target="a5", mode="explicit")
def inline_simt_value_escape_entry():
meta_tile = pto.alloc_tile(shape=[1, 8], dtype=pto.i32, valid_shape=[1, 1])
@@ -144,7 +163,7 @@ def main() -> None:
AttributeError,
"pto.ukernel is not a supported PTODSL public interface",
'@pto.jit(mode="explicit")',
- "@pto.simd/@pto.simt/@pto.cube",
+ "@pto.tileop",
)
expect_raises(
define_removed_tensor_spec_surface_probe,
@@ -170,21 +189,26 @@ def main() -> None:
expect_raises(
define_bad_subkernel_signature_probe,
TypeError,
- "@pto.simd parameter 'A' cannot be annotated with pto.tensor_spec(...)",
+ "@pto.tileop parameter 'A' cannot be annotated with pto.tensor_spec(...)",
"@pto.jit positional parameters",
)
expect_raises(
- define_illegal_simd_ptr_signature_probe,
+ define_legacy_simd_surface_probe,
TypeError,
- "@pto.simd parameter 'meta_ptr' uses unsupported subkernel annotation",
- "pto.Tile parameters plus PTO scalar annotations",
- "@pto.jit(entry=False)",
+ "@pto.simd is a legacy PTODSL subkernel surface",
+ "Use @pto.tileop",
)
expect_raises(
- define_illegal_cube_scalar_signature_probe,
+ define_legacy_cube_surface_probe,
TypeError,
- "@pto.cube parameter 'cols' uses unsupported subkernel annotation",
- "pto.Tile parameters only",
+ "@pto.cube is a legacy PTODSL subkernel surface",
+ "Use @pto.tileop",
+ )
+ expect_raises(
+ define_illegal_tileop_ptr_signature_probe,
+ TypeError,
+ "@pto.tileop parameter 'meta_ptr' uses unsupported subkernel annotation",
+ "pto.Tile / pto.TensorView / pto.PartitionTensorView parameters plus PTO scalar annotations",
"@pto.jit(entry=False)",
)
expect_raises(
@@ -198,28 +222,34 @@ def main() -> None:
nested_simt_from_simd_entry.compile,
RuntimeError,
"@pto.simt helper materialization is only supported from the top-level @pto.jit body",
- "inside @pto.simd",
+ "inside @pto.tileop",
)
expect_raises(
nested_inline_simt_from_simd_entry.compile,
RuntimeError,
"inline pto.simt() may only be used from the top-level @pto.jit body",
- "inside @pto.simd",
+ "inside @pto.tileop",
)
expect_raises(
simd_value_escape_entry.compile,
RuntimeError,
- "@pto.simd cannot return transient SIMD values",
+ "@pto.tileop cannot return transient SIMD values",
"!pto.mask",
"Write the value back to a Tile/UB buffer instead",
)
expect_raises(
illegal_subkernel_callsite_entry.compile,
TypeError,
- "@pto.simd argument 'inp_tile' violates the declared subkernel interface",
+ "@pto.tileop argument 'inp_tile' violates the declared subkernel interface",
"Expected a pto.Tile value",
"either pass a legal PTODSL boundary value or remove the subkernel decorator",
)
+ expect_raises(
+ invalid_tileop_return_entry.compile,
+ TypeError,
+ "@pto.tileop return values must be PTO scalar values",
+ "Return Tile/TensorView/PartitionTensorView/ptr data through explicit subkernel operands instead",
+ )
expect_raises(
inline_simt_value_escape_entry.compile,
RuntimeError,
diff --git a/test/dsl-st/vmulscvt.py b/test/dsl-st/vmulscvt.py
index 2c348d80f1..e5ad9bb742 100644
--- a/test/dsl-st/vmulscvt.py
+++ b/test/dsl-st/vmulscvt.py
@@ -79,7 +79,7 @@ def vmulscvt_pack_kernel(
pto.set_flag("MTE2", "V", event_id=0)
pto.wait_flag("MTE2", "V", event_id=0)
- with pto.simd():
+ with pto.tileop():
mask32 = pto.pset_b32(pto.MaskPattern.ALL)
mask16 = pto.pset_b16(pto.MaskPattern.ALL)
diff --git a/test/lit/pto/plan_memory_ptodsl_tileop_helper_vlds_vsts.pto b/test/lit/pto/plan_memory_ptodsl_tileop_helper_vlds_vsts.pto
new file mode 100644
index 0000000000..bd45c71a89
--- /dev/null
+++ b/test/lit/pto/plan_memory_ptodsl_tileop_helper_vlds_vsts.pto
@@ -0,0 +1,55 @@
+// 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.
+
+// Guards PlanMemory on PTODSL-style tileop helpers that bridge tile operands
+// through pto.tile_buf_addr into VPTO vector load/store ops.
+// RUN: ptoas --pto-arch=a5 --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s
+
+module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} {
+ func.func @plan_memory_ptodsl_tileop_helper_vlds_vsts() attributes {pto.entry} {
+ %src = pto.alloc_tile : !pto.tile_buf