Skip to content

Commit db7b13d

Browse files
jimmychou0jimmychou0
committed
feat(ptodsl): redesign tileop simt/simd/cube subkernel surface
1 parent 7def542 commit db7b13d

69 files changed

Lines changed: 3228 additions & 245 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎docs/designs/ptodsl-redesign-of-simd-simt-cube-subkernel.md‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
2+
# 定稿设计方案（final）
3+
4+
## 1. 目标与用户模型
5+
6+
`pto.tileop` 统一 custom subkernel 标识（取代 `pto.cube`/`pto.simd` 作 subkernel 标识职责；`pto.simt` 专属 launched SIMT）。建模为 tile-level IR 上以 tile/tensorview/scalar 为 IO、带 phase 摘要的命名 helper + `func.call`，让 `PTOInsertSync`/`PTOPlanMemory` 当一等公民。用户零参数，**摘要全由后端 `PTOInferTileOpSummaryPass` 推导；canonical helper marker 统一收敛到 `pto.tileop.helper` unit attr。**
7+
8+
```python
9+
@pto.tileop
10+
def softmax(src_view: pto.TensorView, out_tile: pto.Tile, scratch_tile: pto.Tile,
11+
rows: pto.i32, cols: pto.i32):
12+
# caller 传入 scratch/out tile；body 不新建 tile buffer
13+
pto.tload(src_view, scratch_tile) # MTE
14+
m = pto.vmax(scratch_tile) # PIPE_V
15+
e = pto.vexp(pto.vsub(m, scratch_tile)) # PIPE_V
16+
s = pto.vsum(e) # PIPE_V
17+
r = pto.vdiv(e, s) # PIPE_V
18+
pto.tstore(out_tile, r) # MTE
19+
20+
@pto.jit
21+
def kernel(out, x):
22+
softmax(x, out, scratch, rows, cols)
23+
softmax(x, out, scratch, rows, cols) # 复用
24+
```
25+
26+
约束（编译期强制）：
27+
- IO 只允许 Tile/TensorView/PartitionTensorView/PTO scalar。**输出 tile/tensorview 全走 output operand + operand_effects=write/readwrite；func.call results MVP 只允许 scalar。**
28+
- **helper 内禁止 `alloc_tile`/`reserve_buffer`/`TAlloc`/任何需 PlanMemory 为 callee-local 规划的 op；内部 tile buffer 必须来自 caller operand。内部 vreg/mask/scalar 临时可存但不跨边界。**
29+
- body 允许 tload/tstore、vector ops、scalar(PIPE_S) ops、cube ops、`pipe_barrier` 同步。
30+
- 不允许 host tensor/TensorSpec/vreg/mask/pipe_handle 跨边界；不允许 SIMT-only op。
31+
- **tileop 不能在另一 tileop/section 内调用**（避免 inline 后嵌套 section，`VPTOSplitCVModule:113-118` 禁嵌套 section）。
32+
- **负例：tileop 只有 MTE/S/sync、无 vector/cube 主计算证据时报错。**
33+
- MVP 单主计算域（vector **或** cube）+ 多辅助 pipe；reject cube+vector 混算。多 phase 是 correctness 必需。
34+
35+
## 2. 关键后端事实（已核对，含 pipeline 实测）
36+
37+
- `PTOInsertSync`/`PTOPlanMemory` 均 `func.walk` 全递归进 region。
38+
- InsertSync 现有两条 PTODSL subkernel 路径：legacy `simd/cube` 兼容路径仍按 helper role→单 pipe、memory operand **保守建模为 read+write**；`tileop` 路径已读 `primary_domain/phases/operand_effects`，按 **non-empty boundary-effect phase** 拆多 `CompoundInstanceElement`。
39+
- **`CompoundInstanceElement`（`SyncCommon.h:334-341`）单 `kPipeValue` + 单组 def/use**；空 def/use 节点合法但不贡献跨边界依赖。
40+
- **`classifyTileOpByPipe`（`PTONormalizeUncoveredTileSections.cpp:252-254`）把 MTE1 归 Cube**；`inferSegmentKind` 对混合段报错。
41+
- **`normalizeFunction`（:740-764）**：对不带 kernel_kind 且不带 `subkernel_helper="tileop"` 的 func 会 `collectUncoveredTopLevelSegments`→`inferSegmentKind`，混合段 `emitSegmentInferenceError` 失败。`hasKnownKernelKindContext` 现已把 `subkernel_helper="tileop"` 视为已知上下文并直接跳过 NormalizeUncovered。
42+
- **`VPTOSplitCVModule`**：`hasSectionKind`（:58-83）检查 func 含 `SectionCubeOp`/`SectionVectorOp`；不带 section 的 split candidate 被 `eraseSectionSplitCandidatesWithoutSectionKind`（:170-175）擦除；`:135` 要求"must contain section"；`:113-118` 禁嵌套 section。
43+
- **实测 pipeline 顺序（`ptoas.cpp:1780-1900`）**：
44+
```
45+
preBackendPM: NormalizeUncoveredTileSections (1786)
46+
main pm: ... → PTOInferTileOpSummaryPass → PTOMaterializeTileOpSectionsPass
47+
→ PTOVerifyTileOpContractPass → ... → ViewToMemref
48+
→ PlanMemory → ResolveReservedBuffers
49+
→ VerifySubkernelPipeContract → InsertSync → ... →
50+
MaterializeTileHandles → InlineBackendHelpers
51+
```
52+
- MLIR attribute 不能引用 SSA value；PTO 无 ValueAttr 机制；custom attr 需在 `PTOAttrs.td` 注册（`PTO_Attr` 基类 :36）。
53+
- `alloc_tile`（`PTOOps.td:318`）、`reserve_buffer`（:1792）、`TAllocOp`（:2240，PlanMemory:478 处理）真实存在。
54+
55+
## 3. IR 载体：方案 B — named helper + `func.call` + 后端推导的 phase 摘要
56+
57+
复用 `func.call` + 命名 helper + callee 解析。**不复用 `kernel_kind`**，用后端推导的 `pto.tileop.primary_domain`。**不预套 section**，改 verifier + NormalizeUncovered 接受 tileop 裸 body。**前端只标 `pto.tileop.helper`，摘要全后端生成。**
58+
59+
```
60+
// 前端 trace 后：helper 只带统一 marker
61+
func.func @softmax(%src: !pto.tensorview<...>, %out: !pto.tile<...>, %scratch: !pto.tile<...>,
62+
%rows: i32, %cols: i32)
63+
{ pto.tileop.helper } {
64+
pto.tload %src, %scratch
65+
%m = pto.vmax %scratch
66+
...
67+
pto.tstore %out, %r
68+
return
69+
}
70+
71+
// PTOInferTileOpSummaryPass 后：补全摘要（真 MLIR attr 结构）
72+
func.func @softmax(%src, %out, %scratch, %rows, %cols)
73+
{ pto.tileop.helper,
74+
pto.tileop.primary_domain = #pto.kernel_kind<vector>,
75+
pto.tileop.phases = #array<#dict<{
76+
pipe = #pto.pipe<MTE1>, operand_uses = [0], operand_defs = [2], result_defs = []
77+
}, #dict<{
78+
pipe = #pto.pipe<V>, operand_uses = [2], operand_defs = [1], result_defs = []
79+
}, #dict<{
80+
pipe = #pto.pipe<MTE1>, operand_uses = [1], operand_defs = [1], result_defs = []
81+
}>>,
82+
pto.tileop.operand_effects = ["read", "readwrite", "readwrite", "read", "read"]
83+
} {
84+
...
85+
}
86+
87+
func.func @kernel(%out, %x, %scratch, %rows, %cols) {
88+
func.call @softmax(%x, %out, %scratch, %rows, %cols) : (...) -> ()
89+
func.call @softmax(%x, %out, %scratch, %rows, %cols) : (...) -> () // 复用
90+
}
91+
```
92+
93+
### phase attr schema（真 MLIR 结构，需在 PTOAttrs.td 注册）
94+
95+
- `pto.tileop.phases`: `ArrayAttr<DictionaryAttr>`，每 phase dict：
96+
- `pipe`: 复用现有 pipe 整数枚举（`PTOAttrs.td:213-227`）或新 `PipeAttr`（需注册），按 op `getPipe()` 推。
97+
- `operand_uses`: 整数 `ArrayAttr`（operand index 指向函数所有 operands；**InsertSync 只消费 memory-like operand**，scalar 可在 summary 供验证或忽略、不参与建图；可空）。
98+
- `operand_defs`: 整数 `ArrayAttr`（同上，可空）。
99+
- `result_defs`: 整数 `ArrayAttr`（**MVP 固定空或仅 scalar result，不参与 memory sync**；复杂语义后置）。
100+
- **effects 可空**：纯内部 phase 可全空，保留在 phases 用于校验/主域推导，**InsertSync 跳过不建 `CompoundInstanceElement`**。有 boundary effect 的 phase 才建节点。是否标 use/def 是 **policy 非 IR 不变量**。
101+
- `pto.tileop.operand_effects`：从 phases 非空 effects 派生（union use→read、def→write），unknown→readwrite；scalar 标 read 但不建图。
102+
- `pto.tileop.primary_domain`：主计算域 vector/cube（借用枚举值，不挂 `kernel_kind` attr）。
103+
- **去掉 `pipe_footprint`**：body pipe set 由 `phases.pipe` 集合表达。
104+
105+
### operand index 作用域（明确）
106+
107+
- index 指向**函数所有 operands**（含 scalar）。
108+
- InsertSync 只消费 **memory-like operand**（tile/tensorview）的 use/def 建 `CompoundInstanceElement`。
109+
- scalar operand 可在 summary 供验证或忽略；**不参与 def/use 建图**。
110+
111+
### 摘要属性职责划分
112+
113+
| 职责 | 由谁承担 |
114+
|---|---|
115+
| body 出现过的 pipe 集合 | `phases.pipe` 集合（verifier 校验 body op getPipe() ∈ 此集） |
116+
| caller 跨边界 sync 建模 | 有 boundary effect 的 phase（memory-like operand use/def 非空），InsertSync 为其建 `CompoundInstanceElement` |
117+
| 主计算域 | `primary_domain` |
118+
| 每 operand 副作用 | `operand_effects`（从非空 phase effects 派生；scalar 标 read 但不建图） |
119+
120+
> 不保留 `has_sync`：InsertSync 假设 helper 内部自管同步、caller 层只管跨边界。
121+
122+
### 输出/results 边界（MVP 硬约束）
123+
124+
- 输出 tile/tensorview 全走 output operand + operand_effects=write/readwrite。
125+
- **func.call results MVP 只允许 scalar**（alias handle 后置）。
126+
- **helper 内禁 `alloc_tile`/`reserve_buffer`/`TAlloc`**；内部 tile 必须来自 caller operand；内部 vreg/mask/scalar 临时不跨边界。
127+
128+
## 4. MVP 边界
129+
130+
- 单主计算域 + 多辅助 pipe；reject cube+vector 混算。
131+
- **多 phase correctness 必需**，pipe 按 `getPipe()` 推；effects 可空（非空才建 sync 节点），softmax V phase 标 use/def 是 policy。
132+
- MTE1/2/3/4、PIPE_S 归 phase、不参与 primary_domain 判定（tileop 专用规则，不改全局 `classifyTileOpByPipe`）。
133+
- 禁 helper-local tile allocation。
134+
- SIMT-only op 排除。**tileop 禁嵌套调用**。
135+
- **负例：tileop 只有 MTE/S/sync 无主计算证据报错。**
136+
137+
## 5. 改动点
138+
139+
### 前端（`_subkernels.py` + `_tracing/session.py`）
140+
141+
1. `_create_subkernel_section_op`：tileop **不预套 section**。
142+
2. helper 只附 `pto.tileop.helper`；**不写 primary_domain/phases/operand_effects**。
143+
3. helper 函数类型：输出全走 operand；results 只用于 scalar。
144+
4. 前端 public boundary 契约：保留 vreg/mask 不外逃；results 限 scalar；**禁 tileop 嵌套调用**。helper body 内 `alloc_tile/reserve_buffer/TAlloc` 等 helper-local 资源分配由后端 `PTOVerifyTileOpContractPass` 兜底拒绝。
145+
5. 装饰器无 `kind` 参数；`@pto.cube`/`@pto.simd` 别名（IR 层统一 "tileop"）。
146+
147+
### 后端
148+
149+
1. **verifier 改造**：tile op verifier 把带 `pto.tileop.helper` 的 func 当合法上下文；results 限 scalar；**拒绝 alloc_tile/reserve_buffer/TAlloc**；内部 vreg/mask/scalar 临时不跨边界。
150+
2. **`PTONormalizeUncoveredTileSections` 跳过 tileop**（P0）：`normalizeFunction`/`hasKnownKernelKindContext` 增条件——带 `pto.tileop.helper` 的 func 跳过，避免 preBackendPM:1786 扫到裸 body 混合段报错。
151+
3. **新增 `PTOInferTileOpSummaryPass`**：扫 helper body 推导 primary_domain + phases（pipe 按 `getPipe()`，effects 可空，operand index 指向所有 operand，memory-like 才建图）+ operand_effects（从非空 effects 派生）。tileop 专用 MTE/S 规则，不改全局 `classifyTileOpByPipe`。
152+
4. **新增 materialize pass**：按 primary_domain+phases 物化 `SectionCubeOp`/`SectionVectorOp`（只包 cube/vector 主段），MTE/S/sync 保持 top-level。**lit case 覆盖两类**：MTE+section.vector+MTE、MTE+section.cube+MTE，验证 `VPTOSplitCVModule`/EmitC/VPTO 接受（注意 `hasSectionKind`:58-83 要求 func 含 section，`eraseSectionSplitCandidatesWithoutSectionKind`:170-175 会擦除无 section 的 candidate）。
153+
5. **`UpdatePTODSLSubkernelCallInfo` 改造**：读 primary_domain+phases；按**有 boundary effect 的 phase**拆多 `CompoundInstanceElement`（空 effect phase 跳过，scalar operand 不建图）；memory operand 副作用从保守全 R+W 改读 operand_effects；支持 callsite scalar results 进依赖图。
154+
6. **新增 `PTOVerifyTileOpContractPass`**：校验 body op `getPipe()` ∈ phases pipe 集合、主域 pipe 与 primary_domain 一致、operand_effects == 非空 phases 派生、SIMT-only op 排除、cube+vector 混算 reject、results 限 scalar、tileop 无嵌套调用、**拒绝 alloc_tile/reserve_buffer/TAlloc**、**负例（只有 MTE/S/sync 无主计算证据报错）**。旧 `PTOVerifySubkernelPipeContractPass` 保留兼容 cube/simd。
155+
7. 主 pipeline 不再依赖 `PTOWrapFunctionsInSectionsPass` 为 tileop helper 自动套单段；tileop section 形成以 `PTOMaterializeTileOpSectionsPass` 为准。
156+
8. **`PTOInlineBackendHelpers`**：保证不丢围绕 call 的 sync ops；tileop 禁嵌套调用从源头避免 inline 后嵌套 section。
157+
158+
### pass 顺序（实测修正，P0）
159+
160+
```
161+
前端 trace → verifier(tileop 裸 body)
162+
→ [preBackendPM] NormalizeUncoveredTileSections (跳过 `pto.tileop.helper`) ← P0 必须跳过
163+
→ [main pm] PTOInferTileOpSummaryPass
164+
→ PTOMaterializeTileOpSectionsPass
165+
→ PTOVerifyTileOpContractPass
166+
→ ... → ViewToMemref → PlanMemory → ResolveReservedBuffers →
167+
VerifySubkernelPipeContract → InsertSync → ... →
168+
MaterializeTileHandles → InlineBackendHelpers
169+
```
170+
171+
### 可选后置（非 MVP）
172+
173+
- helper-local tile allocation 的 callsite clone/inline（放开 result 返回 tile / alloc_tile 等）。
174+
- alias handle result（放开 result 非 scalar / result_defs 复杂语义）。
175+
- phases def/use 细到 UB 子区域。
176+
- 内联 opt pass。
177+
178+
## 6. 当前落地状态
179+
180+
已落地并与本文主设计一致的部分：
181+
182+
1. `NormalizeUncoveredTileSections` 已把 `subkernel_helper="tileop"` helper 视为已知上下文并跳过预归一化。
183+
2. `PTOInferTileOpSummaryPass`、`PTOMaterializeTileOpSectionsPass`、`PTOVerifyTileOpContractPass` 已接入主 pipeline，且都位于 `PlanMemory` 之前。
184+
3. `UpdatePTODSLSubkernelCallInfo` 已能消费 tileop 摘要，按 phase 建模跨 helper 边界的 InsertSync 依赖；legacy `simd/cube` 兼容路径仍保留保守单-pipe 建模。
185+
4. tileop helper ABI 已收敛为 Tile/TensorView/PartitionTensorView/PTO scalar；`ptr` 仍为 SIMT-only。
186+
5. `@pto.tileop` / retained `@pto.simd` / `@pto.cube` 在 IR 层语义上已统一到 tileop helper role。
187+
188+
## 7. 仍待单独收敛的差异
189+
190+
1. **helper marker attr 仍是 legacy `pto.ptodsl.subkernel_helper = "tileop"`，尚未切到 canonical `pto.tileop.helper`。**
191+
设计目标是把“这是 tileop helper”这个内部标记也统一收敛到 `pto.tileop.*` 命名空间；当前实现与 lit 仍在读写旧 attr，下一步需要兼容迁移。
192+
2. **inline `with pto.tileop()` 仍有前端预套 section 的实现残留。**
193+
目标设计仍是 tileop helper 不预套 section、统一交后端 materialize；当前 inline 路径仍会先包 `SectionVectorOp`。
194+
3. **`pto.tileop.operand_effects` 的“无显式 boundary effect 时默认值”尚未与本文最终写法重新对齐。**
195+
本文目标写法仍按 `unknown→readwrite` 记录；当前实现会把无 boundary effect 的 operand 物化/校验为 `"read"`。这一点需要单独决策后再统一设计与实现。

‎include/PTO/Transforms/Passes.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,9 @@ std::unique_ptr<Pass> createPTOInferValidatePipeInitPass();
4040
std::unique_ptr<Pass> createPTOResolveReservedBuffersPass();
4141
std::unique_ptr<Pass> createPTOWrapFunctionsInSectionsPass();
4242
std::unique_ptr<Pass> createPTONormalizeUncoveredTileSectionsPass();
43+
std::unique_ptr<Pass> createPTOInferTileOpSummaryPass();
44+
std::unique_ptr<Pass> createPTOMaterializeTileOpSectionsPass();
45+
std::unique_ptr<Pass> createPTOVerifyTileOpContractPass();
4346
std::unique_ptr<Pass> createVPTOSplitCVModulePass();
4447
std::unique_ptr<Pass> createVPTONormalizeContainerPass();
4548
std::unique_ptr<Pass> createPTOVerifyTFreePass();

‎include/PTO/Transforms/Passes.td‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -312,6 +312,77 @@ def PTONormalizeUncoveredTileSections
312312
];
313313
}
314314

315+
def PTOInferTileOpSummary
316+
: Pass<"pto-infer-tileop-summary", "func::FuncOp"> {
317+
let summary = "Infer phase summary attributes for PTODSL tileop helpers";
318+
let description = [{
319+
Scans functions marked with `pto.ptodsl.subkernel_helper = "tileop"` and
320+
derives the backend-owned summary attributes:
321+
- `pto.tileop.primary_domain`
322+
- `pto.tileop.phases`
323+
- `pto.tileop.operand_effects`
324+
325+
This pass is intentionally summary-only. It does not materialize sections,
326+
inline helpers, or change InsertSync modeling.
327+
}];
328+
329+
let constructor = "mlir::pto::createPTOInferTileOpSummaryPass()";
330+
331+
let dependentDialects = [
332+
"mlir::func::FuncDialect",
333+
"mlir::pto::PTODialect"
334+
];
335+
}
336+
337+
def PTOMaterializeTileOpSections
338+
: Pass<"pto-materialize-tileop-sections", "func::FuncOp"> {
339+
let summary = "Materialize one primary PTO section for PTODSL tileop helpers";
340+
let description = [{
341+
Consumes backend-owned tileop summary attributes on functions marked with
342+
`pto.ptodsl.subkernel_helper = "tileop"` and wraps one contiguous
343+
primary-domain compute span in `pto.section.vector` or `pto.section.cube`.
344+
345+
The MVP implementation expects one contiguous primary compute span in the
346+
helper body. Leading and trailing MTE phases remain top-level so late
347+
helper inlining can expose VPTO section sugar to `vpto-split-cv-module`.
348+
}];
349+
350+
let constructor = "mlir::pto::createPTOMaterializeTileOpSectionsPass()";
351+
352+
let dependentDialects = [
353+
"mlir::func::FuncDialect",
354+
"mlir::pto::PTODialect"
355+
];
356+
}
357+
358+
def PTOVerifyTileOpContract
359+
: Pass<"pto-verify-tileop-contract", "func::FuncOp"> {
360+
let summary = "Verify backend-owned PTODSL tileop helper contracts";
361+
let description = [{
362+
Verifies functions marked with `pto.ptodsl.subkernel_helper = "tileop"`
363+
after summary inference and section materialization.
364+
365+
Current MVP contract:
366+
- results must be PTO scalar values only
367+
- helper-local `pto.alloc_tile`, `pto.reserve_buffer`, and `pto.talloc`
368+
are rejected
369+
- nested tileop helper calls are rejected
370+
- SIMT-only PTO ops are rejected
371+
- at least one primary compute op must exist, and all primary compute ops
372+
must belong to exactly one domain (`vector` or `cube`)
373+
- scalar/MTE/sync phases may coexist, but `pto.tileop.primary_domain`,
374+
`pto.tileop.phases`, and `pto.tileop.operand_effects` must remain
375+
consistent with the helper body
376+
}];
377+
378+
let constructor = "mlir::pto::createPTOVerifyTileOpContractPass()";
379+
380+
let dependentDialects = [
381+
"mlir::func::FuncDialect",
382+
"mlir::pto::PTODialect"
383+
];
384+
}
385+
315386
def VPTOSplitCVModule : Pass<"vpto-split-cv-module", "ModuleOp"> {
316387
let summary = "Split a VPTO module with cube/vector sections into kernel modules";
317388
let description = [{

0 commit comments

Comments
 (0)