From 8d30be938df49c035493824dcfedd4bc12500d1f Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 16 Jul 2026 17:39:19 +0800 Subject: [PATCH 01/17] tci and trowexpand --- ...oas-implicit-tmp-materialization-design.md | 579 ++++++++++++++++++ include/PTO/Transforms/Passes.h | 2 + include/PTO/Transforms/Passes.td | 16 + lib/PTO/IR/PTO.cpp | 208 ++++++- lib/PTO/Transforms/CMakeLists.txt | 1 + .../Transforms/PTOMaterializeImplicitTmp.cpp | 290 +++++++++ test/lit/pto/easy_param_completion_emitc.pto | 4 +- test/lit/pto/tci_i16_emitc.pto | 4 +- .../pto/tci_implicit_tmp_level3_invalid.pto | 13 + .../pto/tci_implicit_tmp_materialization.pto | 25 + test/lit/pto/tci_tmp_contract_a3_invalid.pto | 27 + test/lit/pto/tci_ui16_emitc.pto | 4 +- test/lit/pto/tci_ui32_emitc.pto | 4 +- ...trowexpand_implicit_tmp_level3_invalid.pto | 24 + ...rowexpand_implicit_tmp_materialization.pto | 42 ++ .../trowexpand_tmp_contract_a3_invalid.pto | 33 + tools/ptoas/ptoas.cpp | 3 + 17 files changed, 1263 insertions(+), 16 deletions(-) create mode 100644 docs/designs/ptoas-implicit-tmp-materialization-design.md create mode 100644 lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp create mode 100644 test/lit/pto/tci_implicit_tmp_level3_invalid.pto create mode 100644 test/lit/pto/tci_implicit_tmp_materialization.pto create mode 100644 test/lit/pto/tci_tmp_contract_a3_invalid.pto create mode 100644 test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto create mode 100644 test/lit/pto/trowexpand_implicit_tmp_materialization.pto create mode 100644 test/lit/pto/trowexpand_tmp_contract_a3_invalid.pto diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md new file mode 100644 index 0000000000..40ebab49f3 --- /dev/null +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -0,0 +1,579 @@ +# PTOAS Implicit Tmp Materialization Design + +## 背景 + +PTOAS 前端 IR 中很多 tile op 的 `tmp` operand 是可选的。当前如果用户没有显式写 `tmp`,PTOAS 会继续 lowering 到后端不带 `tmp` 的 C++ 接口。 + +后续希望改成:前端仍允许省略 `tmp`,但 PTOAS 在内部为需要 tmp-aware 后端接口的 op 自动补充合法 tmp tile,并让这些 tmp tile 和其它 tile buffer 一起进入 memplan 做 local addr 规划。 + +本文档的目标是给所有类似 `pto.tci` 的 optional tmp op 提供整体改造方案。每个 op 再根据自己的后端 tmp 规格,补充 op-specific 的 tmp requirement、MemoryEffects、verifier 和测试。 + +## 目标 + +- 保持前端 IR 兼容:用户仍然可以写不带 `tmp` 的目标 op。 +- 在 memplan 之前补齐隐式 tmp,使 tmp 作为普通 local allocation root 参与地址规划。 +- 对需要 tmp-aware 后端接口的 op,EmitC lowering 统一走带 tmp 的 C++ overload,避免继续选择 no-tmp 后端接口。 +- 每个 op 的 tmp shape、dtype、address space、layout、容量等约束由该 op 的后端接口规格决定。 +- 不在 EmitC lowering 中临时分配 tmp 地址。 +- 不在 memplan 中特殊创建 tmp;memplan 只负责规划已经存在的 root。 + +## 非目标 + +- 本阶段不一次性覆盖所有 optional tmp op。 +- 本阶段不改变用户显式提供 tmp 的语义。 +- 本阶段不为 level3 自动分配 tmp 地址。 +- 本阶段不引入新的全局 workspace 规划。 +- 本阶段不把所有 op 的 tmp 规格抽象成完全统一的 shape;不同 op 可以有不同 tmp requirement。 + +## 总体方案 + +新增一个 IR 规范化 pass: + +```text +pto-materialize-implicit-tmp +``` + +该 pass 运行在 `PTOViewToMemref` 之后、`pto-plan-memory` 之前: + +```text +PTOViewToMemref + -> pto-materialize-implicit-tmp + -> pto-plan-memory + -> PTOResolveReservedBuffers + -> sync passes + -> PTOMaterializeTileHandles + -> PTOToEmitC +``` + +pass 的职责是扫描所有已纳入改造的目标 op。如果 op 没有 tmp operand,就根据该 op 的 `TmpRequirement` 在 op 前插入 `pto.alloc_tile(no addr)`,并重写原 op,使其显式携带 tmp。 + +抽象流程: + +```text +target_op(no tmp) + -> lookup TmpRequirement(target_op) + -> create pto.alloc_tile(no addr) tmp + -> rewrite target_op(with tmp) + -> memplan assigns addr to tmp + -> EmitC sees tmp and emits tmp-aware overload +``` + +`TmpRequirement` 至少应包含: + +```text +AddressSpace space; +Type elementType; +StaticShape or MinBytes requirement; +Layout/layout-family requirement; +uint64_t minBytes; +bool requireExplicitAtLevel3; +``` + +对自动生成的 tmp: + +- 使用 tile-native `pto.alloc_tile(no addr)`。 +- 不创建 `memref.alloc`。 +- 不创建 `pto.pointer_cast` / `pto.bind_tile`。 +- 不设置 `addr`,由 memplan 统一规划。 +- 尽量使用静态 full-valid shape,即 `v_row/v_col` 与 `rows/cols` 一致,不额外携带 `valid_row` / `valid_col` operand。 +- 定义位置必须支配目标 op。 +- 生命周期由后续 liveness/memplan 根据真实 use 计算。 + +## Memplan 接入 + +自动生成的 tmp 是 tile-native `pto.alloc_tile(no addr)`,因此复用当前 memplan 路径: + +```text +pto.alloc_tile(no addr) + -> local allocation root + -> legacy/modern memplan 分配 offset + -> pto.alloc_tile addr = ... +``` + +legacy memplan 和 modern memplan 都应把自动生成的 tmp 当成普通 local allocation root。memplan 不应该知道“这是某个 op 的隐式 tmp”,也不应该在内部临时创建 tmp。 + +memplan 侧需要依赖 op 的 MemoryEffects / semantic no-alias 信息保证正确复用: + +- tmp 如果是 scratch buffer,应通过 `Write(tmp)` 建模,使 scratch-output conflict 能禁止 tmp 与同 op output 错误复用。 +- 如果某个 op 的 tmp 与 output 不能 alias,但 tmp 不适合建模成 scratch write,则应在 semantic no-alias side table 中显式加入 `forbidAlias(tmp, output)`。 +- 每个 op 的专项改造必须说明 tmp 和 output、input 之间的 alias 约束。 + +## Level 行为 + +### level1 / level2 + +level1/level2 下 memplan 会运行,因此允许省略 tmp: + +```text +target_op(no tmp) + -> pto-materialize-implicit-tmp + -> pto.alloc_tile(no addr) tmp + -> pto-plan-memory 补 addr +``` + +用户显式提供 tmp 时,仍需满足该 op 的 tmp verifier 约束。level1/level2 下用户不应显式指定 local addr,地址由 memplan 统一规划。 + +### level3 + +level3 下用户显式管理 local 地址,memplan 通常跳过。因此不应自动创建无地址 tmp。 + +通用规则: + +```text +level3 + target_op(no tmp) => pass/verifier 报错 +``` + +用户在 level3 使用已纳入改造的目标 op 时,必须显式提供合法 tmp,并保证 tmp 自身带合法 local addr,或满足现有 level3 显式地址规则。 + +诊断信息示例: + +```text + requires explicit tmp when compiling at level3 because PlanMemory is skipped +``` + +## EmitC Lowering + +目标 op 的 EmitC lowering 应保持简单: + +- `op.getTmp()` 为空:生成 no-tmp C++ 调用,或者在该 op 改造完成后仅作为未经过 materialize pass 的兜底路径。 +- `op.getTmp()` 非空:生成带 tmp 的 C++ 调用。 + +引入 `pto-materialize-implicit-tmp` 后,level1/level2 的目标 op 在 EmitC 前都会携带 tmp,因此会自然走带 tmp 的 overload。 + +不建议在 PTOToEmitC 中补 tmp,原因: + +- EmitC 阶段已经错过 memplan。 +- 临时生成 tmp 无法获得 local addr。 +- 会绕过 liveness、sync 和 semantic no-alias 分析。 + +## TCI 针对性改造 + +本节描述 `pto.tci` 作为第一批目标 op 的具体落地规则。后续其它 optional tmp op 应新增类似小节,分别说明自己的 tmp 规格、pass 行为、MemoryEffects、verifier 和测试计划。 + +### TCI Tmp 约束 + +`pto.tci` 当前 ODS 已经支持可选 tmp: + +```td +Optional:$tmp +``` + +PTOToEmitC 也已经根据 `op.getTmp()` 选择带 tmp 或不带 tmp 的 C++ 调用。因此 TCI 改造不需要改 `pto.tci` 的 IR 语法,关键是保证进入 EmitC 前缺省 tmp 已经被显式 materialize。 + +TCI 后端 C++ 接口存在两类 overload: + +```cpp +TCI(dst, start) +TCI(dst, start, tmp) +``` + +A2/A3 上 no-tmp overload 可能走 scalar loop;带 tmp overload 才能走更优路径。A5 接受 tmp,但 tmp 可以作为兼容占位,不额外引入有效计算约束。 + +TCI tmp 不应要求固定 shape,应按 PTO-ISA 文档中的精细化 tmp 约束校验容量。PTOAS 对用户显式 tmp 和自动生成 tmp 采用同一组 A2/A3 合法性规则: + +```text +loc = vec +dtype = 4-byte type: f32 / i32 / ui32 +shape = static shape +layout = row_major +fractal = 512 +capacity = product(shape) * sizeof(dtype) +``` + +A2/A3 的最小容量由 dst 元素类型决定: + +```text +b32 dst: i32 / ui32 -> tmp capacity >= 768 bytes +b16 dst: i16 / ui16 -> tmp capacity >= 1792 bytes +``` + +其中 `shape` 可以是任意静态形状,只要总容量满足对应 dst 类型的最小容量。例如 b32 dst 可以使用 `1x192xf32`,b16 dst 可以使用 `1x448xf32`。`Tile` 是 PTO-ISA 文档中推荐的方便形状无关分配,容量为 2048 bytes (2KiB),可以同时覆盖 b32/b16。 + +A5 上 `tmp` Tile 被接受但不使用;A5 硬件直接使用 `vci` 向量指令,无需临时缓冲区。因此 A5 下 `pto.tci(no tmp)` 不需要自动 materialize tmp,用户显式传 tmp 时也不按 A2/A3 的容量规则校验。 + +### Pass 行为 + +对每个 `pto.tci`: + +- 如果已经有 `tmp`,pass 不修改。 +- A5 如果没有 `tmp`,pass 不修改;A5 后端直接使用 `vci`,不需要 tmp。 +- A2/A3 如果没有 `tmp`,且当前 build level 会运行 memplan,则自动补 tmp。 +- A2/A3 如果没有 `tmp`,但当前 level3 会跳过 memplan,则报错,要求用户显式提供带地址的 tmp。 + +重写前: + +```mlir +pto.tci ins(%s : i32) + outs(%dst : !pto.tile_buf) +``` + +A2/A3 重写后,以下以 b32 dst 自动生成 `f32 1x192` tmp 为例: + +```mlir +%tmp = pto.alloc_tile + : !pto.tile_buf + +pto.tci ins(%s, %tmp : i32, + !pto.tile_buf) + outs(%dst : !pto.tile_buf) +``` + +随后 memplan 会把 `%tmp` 当成普通 tile-native local allocation root,和其它 `pto.alloc_tile(no addr)` 一起规划 local address: + +```mlir +%tmp = pto.alloc_tile addr = %c4096_i64 + : !pto.tile_buf +``` + +TCI rewrite 需要保留原 op 的: + +- scalar operand `S`。 +- dst operand。 +- `descending` attr。 +- location。 +- 其它已有属性。 + +### MemoryEffects + +当前 `TCIOp::getEffects()` 只建模为: + +```text +Write(dst) +``` + +A2/A3 自动补 tmp 后,应改为: + +```text +Read(tmp) if tmp exists +Write(tmp) if tmp exists +Write(dst) +``` + +A5 上 tmp 被接受但不使用,因此不应把 tmp 建模为 Read/Write: + +```text +Write(dst) +``` + +原因: + +- liveness 需要看到 tmp 在 `pto.tci` 被使用。 +- sync pass 需要知道 `pto.tci` 会读 tmp 地址。 +- memplan 需要把 tmp 识别为 scratch buffer,避免 tmp 和同 op 的 dst 错误复用。 +- modern memplan 的 op semantic no-alias 和 root use 传播需要真实 use 信息。 + +如果 tmp 不被建模为 Read/Write,tmp 可能被认为没有 use 或不是 scratch,导致生命周期、复用或同步分析不准确。 + +TCI 也可以在 semantic no-alias side table 中显式加入: + +```text +op = pto.tci +forbidAlias(tmp, dst) +``` + +这不是 scratch conflict 生效的必要条件;A2/A3 只要 `TCIOp::getEffects()` 建模了 `Write(tmp)`,tmp 就会进入 scratch buffer conflict。但显式 side table 能防止未来有人调整 MemoryEffects 后破坏 tmp/dst no-alias 语义。 + +### Verifier + +`TCIOp::verify()` 应检查 tmp 是合法 tile buf,并满足后端接口容量约束: + +```text +如果 tmp 存在: + A5: tmp 被接受但不使用,不执行 A2/A3 tmp 容量校验。 + A2/A3: tmp 必须是 vec tile。 + A2/A3: tmp element type 必须是 4 字节类型(f32 / i32 / ui32)。 + A2/A3: tmp shape 必须是静态 shape。 + A2/A3: tmp layout 必须满足后端 TCI tmp 接口要求。 + A2/A3 b32 dst: tmp capacity 必须大于等于 768 bytes。 + A2/A3 b16 dst: tmp capacity 必须大于等于 1792 bytes。 +``` + +这里的关键是“容量满足接口约束”,而不是“shape 必须等于某个固定值”。TCI 按 dst 元素类型精细化检查: + +```text +// b32 dst: 768B 即可 +!pto.tile_buf // 合法 +!pto.tile_buf // 非法,容量不足 + +// b16 dst: 1792B 即可 +!pto.tile_buf // 合法 +!pto.tile_buf // 非法,容量不足 +``` + +自动生成 tmp 选择满足最小容量的 canonical shape: + +- b32 dst: `f32 1x192`。 +- b16 dst: `f32 1x448`。 +- A5: 不自动生成 tmp。 + +### 测试计划 + +#### lit:自动补 tmp + +新增用例: + +```text +test/lit/pto/tci_implicit_tmp_materialization.pto +``` + +检查: + +```text +CHECK: pto.alloc_tile +CHECK-SAME: dtype=f32 +CHECK: pto.tci ins(%{{.*}}, %{{.*}} +CHECK-NOT: memref.alloc +``` + +可以额外检查自动生成 shape:b32 dst 为 `f32 1x192`,b16 dst 为 `f32 1x448`。同时应覆盖 A5 下不自动生成 tmp。 + +#### lit:memplan 回写 addr + +检查 plan memory 后: + +```text +CHECK: pto.alloc_tile addr = +CHECK: pto.tci ins(%{{.*}}, %{{.*}} +``` + +legacy 和 modern 都应覆盖: + +```text +// RUN: ptoas --pto-level=level2 --plan-memory-impl=legacy ... +// RUN: ptoas --pto-level=level2 --plan-memory-impl=modern ... +``` + +#### lit:EmitC 走带 tmp overload + +检查 C++ 输出: + +```text +CHECK: TCI< +CHECK-SAME: Tile< +CHECK-SAME: float +CHECK: TCI{{.*}}({{.*}}, {{.*}}, {{.*}}) +``` + +#### lit:level3 负例 + +```text +level3 + pto.tci(no tmp) +``` + +期望: + +```text +expected-error {{pto.tci requires explicit tmp when compiling at level3}} +``` + +#### lit:verifier 负例 + +用户显式提供非法 tmp: + +- 非 vec space。 +- 非 f32 dtype。 +- dynamic shape。 +- layout 不满足 TCI tmp 接口约束。 +- A2/A3 b32 dst 的 tmp capacity 小于 768 bytes。 +- A2/A3 b16 dst 的 tmp capacity 小于 1792 bytes。 + +期望 verifier 报错。 + +## TROWEXPAND 二元 op 针对性改造 + +本节覆盖以下 row-expand 二元 op: + +```text +pto.trowexpandadd +pto.trowexpandsub +pto.trowexpandmul +pto.trowexpanddiv +pto.trowexpandmax +pto.trowexpandmin +``` + +这些 op 的 PTO-ISA 文档对 tmp 的描述一致:带 `TileDataTmp &tmp` 的 C++ overload 仅支持模式 1;A2/A3 上 tmp 用作行广播缓冲区;A5 接受 tmp 但不使用。 + +### RowExpand Tmp 约束 + +这些 op 有两种 row-broadcast 模式: + +- 模式 1:扩展操作数为 `ColMajor`,每行一个标量。带 tmp overload 仅支持该模式。 +- 模式 2:扩展操作数为 `RowMajor`,每行一个 32 字节块。该模式不需要 tmp,不应为了 tmp-aware overload 强行改写。 + +A2/A3 模式 1 下,tmp 作为 `vbrcb` 广播缓冲区使用。扩展操作数的每行标量会广播成一个 32 字节块;`vbrcb` repeat stride 为 8 个块,即 256 字节,每个 repeat 处理 8 行。 + +tmp 最小容量由 `R = dst.validRow` 决定: + +```text +if R < 256: + tmpBytes = ceil(R / 8) * 256 +else: + tmpBytes = 30 * 256 = 7680 +``` + +说明: + +- 当 `R >= 256` 时,后端按循环处理,每次循环最多 30 个 repeat,也就是 240 行;tmp 在循环间复用,因此每次循环只需要 7680 字节。 +- 一个紧凑的形状无关上界是 8KB,即 8192 字节。该上界可作为自动 materialize 的保守 canonical tmp 大小。 +- 不带 tmp 的 3 参数 overload 支持模式 1 和模式 2;对 A2/A3 的模式 1,后端使用内部 8KB 缓冲区 `TMP_UB_OFFSET`;模式 2 不需要广播缓冲区。 +- A5 硬件通过 `vlds` 广播模式原生支持行广播,tmp 被接口接受但不使用。 + +PTOAS 对用户显式 tmp 的合法性规则: + +```text +A2/A3: + op 必须是模式 1,才能使用显式 tmp。 + tmp 必须是 vec tile。 + tmp shape 必须静态可计算容量,或后续 verifier 能证明容量满足公式。 + tmp capacity >= min(ceil(R / 8) * 256, 7680)。 + +A5: + tmp 被接受但不使用,不执行 A2/A3 tmp 容量校验。 +``` + +### Pass 行为 + +对每个目标 row-expand 二元 op: + +- 如果已经有 `tmp`,pass 不修改,但 verifier 需要保证它只用于合法模式。 +- A5 如果没有 `tmp`,pass 不修改。 +- A2/A3 如果没有 `tmp`,且 op 是模式 1、当前 build level 会运行 memplan,则自动补 tmp。 +- A2/A3 如果没有 `tmp`,op 是模式 1、但当前 level3 会跳过 memplan,则报错,要求用户显式提供带地址的 tmp。 +- 模式 2 不需要 tmp;pass 不应自动补 tmp,也不应强制改成带 tmp overload。 + +自动补 tmp 的 canonical shape 建议采用形状无关上界: + +```mlir +%tmp = pto.alloc_tile + : !pto.tile_buf, + rows=1, cols=<8192 / sizeof(dst element type)>, + v_row=1, v_col=<8192 / sizeof(dst element type)>, + blayout=row_major, slayout=none_box, + fractal=512, pad=0> +``` + +默认使用 `dst element type` 作为 tmp element type,以贴合 row-expand 后端模板参数;如果后续确认某些后端实现允许更宽松的 tmp dtype,可在对应 op-specific verifier 中放宽。 + +这样不需要在 materialize pass 中依赖 `dst.validRow` 是否为静态值,也能覆盖 A2/A3 模式 1 的最大每轮 tmp 需求。后续如果希望节省 UB,可以在能静态证明 `R` 时生成更小 tmp: + +```text +tmpBytes = min(ceil(R / 8) * 256, 7680) +``` + +### MemoryEffects + +A2/A3 上这些 op 的 tmp 是广播 scratch buffer,应该建模为: + +```text +Read(non-tmp inputs) +Read(tmp) if tmp exists +Write(tmp) if tmp exists +Write(dst) +``` + +其中 `Write(tmp)` 用于让 memplan 的 scratch-output conflict 禁止 tmp 与同 op 的 `dst` 错误复用。 + +A5 上 tmp 被接受但不使用,因此不应把 tmp 建模为 Read/Write: + +```text +Read(non-tmp inputs) +Write(dst) +``` + +如果未来某个 row-expand op 的 MemoryEffects 不适合用 `Write(tmp)` 表达 scratch 语义,也应在 semantic no-alias side table 中显式加入: + +```text +op = pto.trowexpand* +forbidAlias(tmp, dst) +``` + +### Verifier + +这些 op 的 verifier 需要区分模式和 arch: + +```text +如果 tmp 存在: + A5: tmp 被接受但不使用,不执行 A2/A3 tmp 容量校验。 + A2/A3: op 必须是模式 1,即扩展操作数为 ColMajor 每行标量。 + A2/A3: tmp 必须是 vec tile。 + A2/A3: tmp capacity 必须满足 min(ceil(dst.validRow / 8) * 256, 7680)。 +``` + +模式识别规则沿用 ISA 文档: + +- `src0` 或 `src1` 中恰好一个与 `dst` 有相同 valid shape,该 operand 是全尺寸操作数。 +- 另一个 operand 是扩展操作数。 +- 扩展操作数为 `ColMajor` 且每行一个标量时是模式 1。 +- 扩展操作数为 `RowMajor` 且每行 `32 / sizeof(T)` 列时是模式 2。 + +如果 `dst.validRow` 是动态值,verifier 无法精确证明用户显式 tmp 是否足够小时,可以采用保守规则: + +- 用户显式 tmp 至少 8192 字节;或 +- 后续引入运行时/符号约束证明 tmp capacity 满足公式。 + +自动生成 tmp 建议先使用 8192 字节 canonical 上界,因此不会受动态 `dst.validRow` 影响。 + +### 测试计划 + +#### lit:自动补 tmp + +为至少一个代表 op 增加 A2/A3 模式 1 用例,例如 `pto.trowexpandadd(no tmp)`: + +```text +CHECK: pto.alloc_tile +CHECK: pto.trowexpandadd ins(%{{.*}}, %{{.*}}, %{{.*}} +CHECK-NOT: memref.alloc +``` + +同时检查 A5 下不自动生成 tmp。 + +#### lit:模式 2 不补 tmp + +构造 RowMajor 扩展操作数的模式 2 用例,确认 pass 不自动补 tmp,并继续走 no-tmp overload。 + +#### lit:memplan 回写 addr + +检查自动生成 tmp 在 plan memory 后带 `addr`: + +```text +CHECK: pto.alloc_tile addr = +CHECK: pto.trowexpand{{.*}} ins(%{{.*}}, %{{.*}}, %{{.*}} +``` + +legacy 和 modern 都应覆盖。 + +#### lit:level3 负例 + +A2/A3 level3 + 模式 1 + no tmp 应报错: + +```text +expected-error {{requires explicit tmp when compiling at level3}} +``` + +A5 level3 + no tmp 不应因为 tmp 缺失报错。 + +#### lit:verifier 负例 + +需要覆盖: + +- A2/A3 显式 tmp 用在模式 2,报错。 +- A2/A3 显式 tmp capacity 小于公式要求,报错。 +- A2/A3 动态 `dst.validRow` 且显式 tmp 小于 8192 字节,按保守规则报错。 +- A5 显式 tmp 不触发 A2/A3 容量校验。 + +## 后续扩展 + +后续新增其它 optional tmp op 时,需要补充一个 op-specific 小节,并明确: + +- op 的 tmp 后端接口规格。 +- 自动生成 tmp 的 canonical shape。 +- 用户显式 tmp 的 verifier 规则。 +- tmp 的 MemoryEffects。 +- tmp 与 input/output 的 alias 约束。 +- level3 下是否要求显式 tmp。 +- lit 覆盖自动补 tmp、memplan 回写 addr、EmitC overload、level3 负例和 verifier 负例。 diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index c1d9e82a14..a9b82c8e3f 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -75,6 +75,8 @@ createPlanMemoryModernPass(const PlanMemoryOptions &options); std::unique_ptr createPTORemoveRedundantBarrierPass(); std::unique_ptr createPTOValidateIntToPtrUsesPass(); std::unique_ptr createPTORematerializeFixpipeVectorQuantPass(); +std::unique_ptr +createPTOMaterializeImplicitTmpPass(bool requireExplicitTmp = false); std::unique_ptr createPTOResolveBufferSelectPass(); std::unique_ptr createInferPTOLayoutPass(); std::unique_ptr createPTOA5NormalizeTMovPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 7b08319dda..97206bb0fb 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -192,6 +192,22 @@ def PTORematerializeFixpipeVectorQuant let dependentDialects = ["mlir::pto::PTODialect", "mlir::func::FuncDialect"]; } +def PTOMaterializeImplicitTmp + : Pass<"pto-materialize-implicit-tmp", "func::FuncOp"> { + let summary = "Materialize implicit tmp tiles for PTO ops before memplan"; + let description = [{ + Rewrites PTO ops with optional tmp operands into explicit tmp forms when + the backend tmp-aware overload is required. The synthesized tmp is emitted + as tile-native `pto.alloc_tile(no addr)` so PlanMemory can assign its local + address together with other tile buffers. + }]; + let constructor = "mlir::pto::createPTOMaterializeImplicitTmpPass()"; + let dependentDialects = [ + "mlir::pto::PTODialect", + "mlir::func::FuncDialect" + ]; +} + def PlanMemory : Pass<"pto-plan-memory", "ModuleOp"> { let summary = "Plan memory for PTO Ops"; let constructor = "mlir::pto::createPlanMemoryPass()"; diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 86907e3255..90362862ab 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -2015,6 +2015,12 @@ static bool isRowMajorTileBuf(Type ty) { return tb && tb.getBLayoutValueI32() == static_cast(pto::BLayout::RowMajor); } +static bool isColMajorTileBuf(Type ty) { + auto tb = mlir::dyn_cast(ty); + return tb && tb.getBLayoutValueI32() == + static_cast(pto::BLayout::ColMajor); +} + static LogicalResult verifyRowReductionSrcLayout(Operation *op, Type ty, StringRef name) { if (failed(verifyTileBufCommon(op, ty, name))) @@ -5841,6 +5847,33 @@ LogicalResult pto::TCIOp::verify() { if (bw != 16 && bw != 32) return emitOpError("expects dst element type to be i16/i32"); + if (getTmp() && getTargetArch(getOperation()) != PTOArch::A5) { + auto tmpTy = mlir::dyn_cast(getTmp().getType()); + if (!tmpTy) + return emitOpError("expects tmp to be a tile buffer"); + auto tmpSpace = + mlir::dyn_cast_or_null(tmpTy.getMemorySpace()); + if (!tmpSpace || tmpSpace.getAddressSpace() != AddressSpace::VEC) + return emitOpError("expects tmp to be in vec address space"); + Type tmpElemTy = tmpTy.getElementType(); + if (!(tmpElemTy.isF32() || tmpElemTy.isInteger(32))) + return emitOpError("expects A2/A3 tmp element type to be a 4-byte type"); + if (tmpTy.getBLayoutValueI32() != static_cast(BLayout::RowMajor)) + return emitOpError("expects tmp blayout to be row_major"); + if (tmpTy.getSLayoutValueI32() != static_cast(SLayout::NoneBox)) + return emitOpError("expects tmp slayout to be none_box"); + if (tmpTy.getSFractalSizeI32() != 512) + return emitOpError("expects tmp fractal size to be 512"); + auto tmpBytes = getStaticByteSize(tmpTy); + if (!tmpBytes) + return emitOpError("expects tmp to have static byte size"); + uint64_t minTmpBytes = bw == 32 ? 768 : 1792; + if (*tmpBytes < minTmpBytes) + return emitOpError("expects A2/A3 tmp capacity to be at least ") + << minTmpBytes << " bytes for " << bw + << "-bit dst element type"; + } + auto sTy = mlir::dyn_cast(getOperand(0).getType()); if (!sTy) return emitOpError("expects S to be integer"); @@ -11723,6 +11756,112 @@ static FailureOr verifyTRowExpandBinaryCore(Operation *op, Type src0Ty, return getElemTy(src0Ty); } +enum class TRowExpandBinaryMode { + Unknown, + Mode1ColMajorScalar, + Mode2RowMajorBlock, +}; + +static bool validShapesCompatibleForTRowExpand(ArrayRef lhs, + ArrayRef rhs) { + if (lhs.size() != rhs.size()) + return false; + for (auto [l, r] : llvm::zip(lhs, rhs)) { + if (l != ShapedType::kDynamic && r != ShapedType::kDynamic && l != r) + return false; + } + return true; +} + +static TRowExpandBinaryMode classifyTRowExpandBinaryMode(Type src0Ty, + Type src1Ty, + Type dstTy) { + auto src0Valid = getValidShapeVec(src0Ty); + auto src1Valid = getValidShapeVec(src1Ty); + auto dstValid = getValidShapeVec(dstTy); + if (src0Valid.size() != 2 || src1Valid.size() != 2 || dstValid.size() != 2) + return TRowExpandBinaryMode::Unknown; + + Type expandedTy; + ArrayRef expandedValid; + if (validShapesCompatibleForTRowExpand(src0Valid, dstValid)) { + expandedTy = src1Ty; + expandedValid = src1Valid; + } else if (validShapesCompatibleForTRowExpand(src1Valid, dstValid)) { + expandedTy = src0Ty; + expandedValid = src0Valid; + } else { + return TRowExpandBinaryMode::Unknown; + } + + int64_t expandedCols = expandedValid[1]; + if (isColMajorTileBuf(expandedTy) && + (expandedCols == ShapedType::kDynamic || expandedCols == 1)) + return TRowExpandBinaryMode::Mode1ColMajorScalar; + + std::optional elemBytes = getElemBytes(getElemTy(dstTy)); + if (!elemBytes || *elemBytes == 0) + return TRowExpandBinaryMode::Unknown; + int64_t expectedMode2Cols = 32 / *elemBytes; + if (isRowMajorTileBuf(expandedTy) && + (expandedCols == ShapedType::kDynamic || + expandedCols == expectedMode2Cols)) + return TRowExpandBinaryMode::Mode2RowMajorBlock; + + return TRowExpandBinaryMode::Unknown; +} + +static int64_t getTRowExpandTmpMinBytes(int64_t dstValidRows) { + if (dstValidRows == ShapedType::kDynamic) + return 8192; + if (dstValidRows < 0) + return 8192; + if (dstValidRows < 256) + return ceilDivInt64(dstValidRows, 8) * 256; + return 30 * 256; +} + +static std::optional getStaticTileCapacityBytes(Type ty) { + auto numElems = getStaticNumElements(getShapeVec(ty)); + auto elemBytes = getElemBytes(getElemTy(ty)); + if (!numElems || !elemBytes) + return std::nullopt; + return *numElems * *elemBytes; +} + +static LogicalResult verifyTRowExpandImplicitTmpContract( + Operation *op, Type src0Ty, Type src1Ty, Type dstTy, Type tmpTy, + bool hasTmp, PTOArch targetArch) { + if (!hasTmp || targetArch == PTOArch::A5) + return success(); + + if (classifyTRowExpandBinaryMode(src0Ty, src1Ty, dstTy) != + TRowExpandBinaryMode::Mode1ColMajorScalar) { + return op->emitOpError( + "expects A2/A3 tmp-form trowexpand to use mode 1 " + "(ColMajor per-row scalar expanded operand)"); + } + + if (failed(verifyVecTileStorage(op, tmpTy, "tmp"))) + return failure(); + if (getElemTy(tmpTy) != getElemTy(dstTy)) + return op->emitOpError("expects tmp and dst to have the same element type"); + + auto dstValid = getValidShapeVec(dstTy); + if (dstValid.size() != 2) + return op->emitOpError("expects dst to have rank-2 valid_shape"); + int64_t minBytes = getTRowExpandTmpMinBytes(dstValid[0]); + std::optional tmpBytes = getStaticTileCapacityBytes(tmpTy); + if (!tmpBytes) + return op->emitOpError( + "expects A2/A3 trowexpand tmp capacity to be statically known"); + if (*tmpBytes < minBytes) + return op->emitOpError() + << "expects A2/A3 trowexpand tmp capacity to be at least " + << minBytes << " bytes, but got " << *tmpBytes << " bytes"; + return success(); +} + mlir::LogicalResult mlir::pto::TRowExpandDivOp::verify() { auto verifyByArch = [&](PTOArch targetArch) -> LogicalResult { Type src0Ty = getSrc0().getType(); @@ -11746,6 +11885,11 @@ mlir::LogicalResult mlir::pto::TRowExpandDivOp::verify() { } if (getPrecisionType() == pto::DivPrecision::HighPrecision && !getTmp()) return emitOpError("expects tmp when precisionType is high_precision"); + if (failed(verifyTRowExpandImplicitTmpContract( + getOperation(), src0Ty, src1Ty, dstTy, + getTmp() ? getTmp().getType() : Type{}, static_cast(getTmp()), + targetArch))) + return failure(); return mlir::success(); }; auto verifyA2A3 = [&]() -> LogicalResult { return verifyByArch(PTOArch::A3); }; @@ -11775,6 +11919,11 @@ mlir::LogicalResult mlir::pto::TRowExpandMulOp::verify() { return emitOpError( "expects A2/A3 trowexpandmul element type to be i16/i32/f16/f32"); } + if (failed(verifyTRowExpandImplicitTmpContract( + getOperation(), src0Ty, src1Ty, dstTy, + getTmp() ? getTmp().getType() : Type{}, static_cast(getTmp()), + targetArch))) + return failure(); return mlir::success(); }; auto verifyA2A3 = [&]() -> LogicalResult { return verifyByArch(PTOArch::A3); }; @@ -11804,6 +11953,11 @@ mlir::LogicalResult mlir::pto::TRowExpandSubOp::verify() { return emitOpError( "expects A2/A3 trowexpandsub element type to be i16/i32/f16/f32"); } + if (failed(verifyTRowExpandImplicitTmpContract( + getOperation(), src0Ty, src1Ty, dstTy, + getTmp() ? getTmp().getType() : Type{}, static_cast(getTmp()), + targetArch))) + return failure(); return mlir::success(); }; auto verifyA2A3 = [&]() -> LogicalResult { return verifyByArch(PTOArch::A3); }; @@ -11855,6 +12009,11 @@ mlir::LogicalResult mlir::pto::TRowExpandAddOp::verify() { if (src1Col != ShapedType::kDynamic && src1Col != 1) return emitOpError("expects non-row-major src1 valid_shape[1] to be 1"); } + if (failed(verifyTRowExpandImplicitTmpContract( + getOperation(), src0Ty, src1Ty, dstTy, + getTmp() ? getTmp().getType() : Type{}, static_cast(getTmp()), + targetArch))) + return failure(); return mlir::success(); }; auto verifyA2A3 = [&]() -> LogicalResult { return verifyByArch(PTOArch::A3); }; @@ -11866,6 +12025,7 @@ static LogicalResult verifyTRowExpandReduceLikeOp(Operation *op, Type src0Ty, Type src1Ty, Type dstTy, Type tmpTy, bool hasTmp, PTOArch targetArch, + bool enforceTmpContract, StringRef opName, bool allowIntegerTypes) { if (failed(verifyTileBufCommon(op, src0Ty, "src0")) || @@ -11984,14 +12144,23 @@ static LogicalResult verifyTRowExpandReduceLikeOp(Operation *op, Type src0Ty, // (A5 tmp-form invariant is checked earlier, before the empty-marker accept.) + auto verifyTmpContract = [&]() -> LogicalResult { + if (!enforceTmpContract) + return success(); + return verifyTRowExpandImplicitTmpContract(op, src0Ty, src1Ty, dstTy, + tmpTy, hasTmp, targetArch); + }; + if (src0MatchesDst) { if (succeeded(checkFullAndBroadcast(src0Ty, src0Valid, "src0", src1Ty, - src1Valid, "src1"))) + src1Valid, "src1")) && + succeeded(verifyTmpContract())) return success(); } if (src1MatchesDst) { if (succeeded(checkFullAndBroadcast(src1Ty, src1Valid, "src1", src0Ty, - src0Valid, "src0"))) + src0Valid, "src0")) && + succeeded(verifyTmpContract())) return success(); } @@ -12005,6 +12174,7 @@ mlir::LogicalResult mlir::pto::TRowExpandExpdifOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A3, + /*enforceTmpContract=*/false, "trowexpandexpdif", /*allowIntegerTypes=*/false); }; @@ -12013,6 +12183,7 @@ mlir::LogicalResult mlir::pto::TRowExpandExpdifOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A5, + /*enforceTmpContract=*/false, "trowexpandexpdif", /*allowIntegerTypes=*/false); }; @@ -12025,6 +12196,7 @@ mlir::LogicalResult mlir::pto::TRowExpandMaxOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A3, + /*enforceTmpContract=*/true, "trowexpandmax", /*allowIntegerTypes=*/true); }; @@ -12033,6 +12205,7 @@ mlir::LogicalResult mlir::pto::TRowExpandMaxOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A5, + /*enforceTmpContract=*/true, "trowexpandmax", /*allowIntegerTypes=*/true); }; @@ -12045,6 +12218,7 @@ mlir::LogicalResult mlir::pto::TRowExpandMinOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A3, + /*enforceTmpContract=*/true, "trowexpandmin", /*allowIntegerTypes=*/true); }; @@ -12053,6 +12227,7 @@ mlir::LogicalResult mlir::pto::TRowExpandMinOp::verify() { getSrc1().getType(), getDst().getType(), getTmp() ? getTmp().getType() : Type{}, (bool)getTmp(), PTOArch::A5, + /*enforceTmpContract=*/true, "trowexpandmin", /*allowIntegerTypes=*/true); }; @@ -14181,6 +14356,11 @@ PTO_DEFINE_UNARY_EFFECTS(TAndSOp, getSrcMutable(), getDstMutable()) // TCI: Write(dst) (generates sequence) void TCIOp::getEffects( SmallVectorImpl> &effects) { + if (auto tmp = getTmpMutable(); + !tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14449,8 +14629,10 @@ void TRowExpandDivOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14459,8 +14641,10 @@ void TRowExpandMulOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14469,8 +14653,10 @@ void TRowExpandSubOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14479,8 +14665,10 @@ void TRowExpandAddOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14499,8 +14687,10 @@ void TRowExpandMaxOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14509,8 +14699,10 @@ void TRowExpandMinOp::getEffects( PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); auto tmp = getTmpMutable(); - if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 75f2c25f85..5dbabf5e90 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -75,6 +75,7 @@ add_mlir_dialect_library(PTOTransforms InsertSync/PTOInsertSync.cpp PTOInjectBarrierAllSync.cpp InsertSync/InsertSyncDebug.cpp + PTOMaterializeImplicitTmp.cpp PTORematerializeFixpipeVectorQuant.cpp PTOValidateIntToPtrUses.cpp InsertTemplateAttributes.cpp diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp new file mode 100644 index 0000000000..abb14a89e2 --- /dev/null +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -0,0 +1,290 @@ +// 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. + +//===- PTOMaterializeImplicitTmp.cpp --------------------------------------===// + +#include "PTO/Transforms/Passes.h" + +#include "PTO/IR/PTO.h" +#include "PTO/IR/PTODialect.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" +#include "llvm/ADT/TypeSwitch.h" + +using namespace mlir; + +namespace { + +static pto::TileBufConfigAttr makeRowMajorNoneBoxConfig(MLIRContext *ctx) { + OpBuilder builder(ctx); + return pto::TileBufConfigAttr::get( + ctx, pto::BLayoutAttr::get(ctx, pto::BLayout::RowMajor), + pto::SLayoutAttr::get(ctx, pto::SLayout::NoneBox), + builder.getI32IntegerAttr(512), + pto::PadValueAttr::get(ctx, pto::PadValue::Null), + pto::CompactModeAttr::get(ctx, pto::CompactMode::Null)); +} + +static unsigned getTCIDstBitWidth(pto::TCIOp op) { + auto tileTy = dyn_cast(op.getDst().getType()); + if (!tileTy) + return 0; + auto elemTy = dyn_cast(tileTy.getElementType()); + if (!elemTy) + return 0; + return elemTy.getWidth(); +} + +static pto::TileBufType makeTCITmpType(MLIRContext *ctx, unsigned dstBitWidth) { + // PTO-ISA TCI A2/A3 vector path needs 768B for b32 dst and 1792B for + // b16 dst. Use an f32 1xN tmp with the exact minimum capacity. + int64_t cols = dstBitWidth == 16 ? 448 : 192; + return pto::TileBufType::get( + ctx, {1, cols}, Float32Type::get(ctx), + pto::AddressSpaceAttr::get(ctx, pto::AddressSpace::VEC), {1, cols}, + makeRowMajorNoneBoxConfig(ctx)); +} + +static std::optional getElemBytes(Type elemTy) { + unsigned bits = pto::getPTOStorageElemBitWidth(elemTy); + if (bits == 0 || bits % 8 != 0) + return std::nullopt; + return bits / 8; +} + +static SmallVector getValidShapeVec(Type ty) { + if (auto tileTy = dyn_cast(ty)) + return SmallVector(tileTy.getValidShape().begin(), + tileTy.getValidShape().end()); + return {}; +} + +static bool validShapesCompatible(ArrayRef lhs, + ArrayRef rhs) { + if (lhs.size() != rhs.size()) + return false; + for (auto [l, r] : llvm::zip(lhs, rhs)) { + if (l != ShapedType::kDynamic && r != ShapedType::kDynamic && l != r) + return false; + } + return true; +} + +static bool isRowMajorTile(Value value) { + auto tileTy = dyn_cast(value.getType()); + return tileTy && tileTy.getBLayoutValueI32() == + static_cast(pto::BLayout::RowMajor); +} + +static bool isColMajorTile(Value value) { + auto tileTy = dyn_cast(value.getType()); + return tileTy && tileTy.getBLayoutValueI32() == + static_cast(pto::BLayout::ColMajor); +} + +enum class RowExpandMode { + Unknown, + Mode1ColMajorScalar, + Mode2RowMajorBlock, +}; + +static RowExpandMode classifyTRowExpandBinaryMode(Value src0, Value src1, + Value dst) { + auto dstValid = getValidShapeVec(dst.getType()); + auto src0Valid = getValidShapeVec(src0.getType()); + auto src1Valid = getValidShapeVec(src1.getType()); + if (dstValid.size() != 2 || src0Valid.size() != 2 || src1Valid.size() != 2) + return RowExpandMode::Unknown; + + Value expanded; + ArrayRef expandedValid; + if (validShapesCompatible(src0Valid, dstValid)) { + expanded = src1; + expandedValid = src1Valid; + } else if (validShapesCompatible(src1Valid, dstValid)) { + expanded = src0; + expandedValid = src0Valid; + } else { + return RowExpandMode::Unknown; + } + + int64_t expandedCols = expandedValid[1]; + if (isColMajorTile(expanded) && + (expandedCols == ShapedType::kDynamic || expandedCols == 1)) + return RowExpandMode::Mode1ColMajorScalar; + + auto dstTileTy = dyn_cast(dst.getType()); + if (!dstTileTy) + return RowExpandMode::Unknown; + auto elemBytes = getElemBytes(dstTileTy.getElementType()); + if (!elemBytes || *elemBytes == 0) + return RowExpandMode::Unknown; + int64_t expectedMode2Cols = 32 / *elemBytes; + if (isRowMajorTile(expanded) && + (expandedCols == ShapedType::kDynamic || + expandedCols == expectedMode2Cols)) + return RowExpandMode::Mode2RowMajorBlock; + + return RowExpandMode::Unknown; +} + +static pto::TileBufType makeTRowExpandTmpType(MLIRContext *ctx, + pto::TileBufType dstTy) { + constexpr int64_t kTmpBytes = 8192; + std::optional elemBytes = getElemBytes(dstTy.getElementType()); + int64_t cols = elemBytes && *elemBytes > 0 ? kTmpBytes / *elemBytes : 2048; + return pto::TileBufType::get( + ctx, {1, cols}, dstTy.getElementType(), + pto::AddressSpaceAttr::get(ctx, pto::AddressSpace::VEC), {1, cols}, + makeRowMajorNoneBoxConfig(ctx)); +} + +static void replaceTRowExpandBinaryOpWithTmp(Operation *op, Value src0, + Value src1, Value tmp, Value dst) { + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands({src0, src1, tmp, dst}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1, 1})); + for (NamedAttribute attr : op->getAttrs()) { + if (attr.getName() == "operandSegmentSizes") + continue; + state.addAttribute(attr.getName(), attr.getValue()); + } + builder.create(state); + op->erase(); +} + +template +static LogicalResult materializeTRowExpandTmp(OpTy op, bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + if (pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5) + return success(); + + RowExpandMode mode = + classifyTRowExpandBinaryMode(op.getSrc0(), op.getSrc1(), op.getDst()); + if (mode != RowExpandMode::Mode1ColMajorScalar) + return success(); + + if (requireExplicitTmp) { + return op.emitOpError( + "requires explicit tmp for A2/A3 row-expand mode 1 when PlanMemory is skipped"); + } + + auto dstTy = dyn_cast(op.getDst().getType()); + if (!dstTy) + return op.emitOpError("expects tile_buf dst when materializing implicit tmp"); + + OpBuilder builder(op); + Value tmp = + builder + .create(op.getLoc(), + makeTRowExpandTmpType(ctx, dstTy), Value(), + Value(), Value()) + .getResult(); + replaceTRowExpandBinaryOpWithTmp(op.getOperation(), op.getSrc0(), op.getSrc1(), + tmp, op.getDst()); + return success(); +} + +struct PTOMaterializeImplicitTmpPass + : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PTOMaterializeImplicitTmpPass) + + PTOMaterializeImplicitTmpPass() = default; + explicit PTOMaterializeImplicitTmpPass(bool requireExplicitTmp) + : requireExplicitTmp(requireExplicitTmp) {} + + StringRef getArgument() const final { return "pto-materialize-implicit-tmp"; } + StringRef getDescription() const final { + return "Materialize implicit tmp tiles for PTO ops before memplan"; + } + + void runOnOperation() override { + func::FuncOp func = getOperation(); + MLIRContext *ctx = func.getContext(); + bool failed = false; + + SmallVector tciOps; + func.walk([&](pto::TCIOp op) { + if (!op.getTmp()) + tciOps.push_back(op); + }); + + for (pto::TCIOp op : tciOps) { + if (pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5) + continue; + + if (requireExplicitTmp) { + op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + failed = true; + continue; + } + + OpBuilder builder(op); + Location loc = op.getLoc(); + auto tmpType = makeTCITmpType(ctx, getTCIDstBitWidth(op)); + Value tmp = + builder.create(loc, tmpType, Value(), Value(), + Value()) + .getResult(); + + auto newOp = builder.create( + loc, TypeRange{}, op.getS(), tmp, op.getDst(), + op.getDescendingAttr()); + for (NamedAttribute attr : op->getAttrs()) { + if (attr.getName() == "operandSegmentSizes") + continue; + newOp->setAttr(attr.getName(), attr.getValue()); + } + op.erase(); + } + + SmallVector rowExpandOps; + func.walk([&](Operation *op) { + if (isa(op)) + rowExpandOps.push_back(op); + }); + + for (Operation *op : rowExpandOps) { + LogicalResult result = + llvm::TypeSwitch(op) + .Case( + [&](auto typedOp) { + return materializeTRowExpandTmp(typedOp, requireExplicitTmp, + ctx); + }) + .Default([](Operation *) { return success(); }); + if (mlir::failed(result)) + failed = true; + } + + if (failed) + signalPassFailure(); + } + +private: + bool requireExplicitTmp = false; +}; + +} // namespace + +std::unique_ptr +mlir::pto::createPTOMaterializeImplicitTmpPass(bool requireExplicitTmp) { + return std::make_unique(requireExplicitTmp); +} diff --git a/test/lit/pto/easy_param_completion_emitc.pto b/test/lit/pto/easy_param_completion_emitc.pto index 53636240bb..36c9f30b80 100644 --- a/test/lit/pto/easy_param_completion_emitc.pto +++ b/test/lit/pto/easy_param_completion_emitc.pto @@ -4,8 +4,8 @@ module { func.func @tci_with_tmp(%dst: !pto.partition_tensor_view<1x16xi16>) { %c0_i16 = arith.constant 0 : i16 %dst_tile = pto.alloc_tile : !pto.tile_buf - %tmp_tile = pto.alloc_tile : !pto.tile_buf - pto.tci ins(%c0_i16, %tmp_tile : i16, !pto.tile_buf) + %tmp_tile = pto.alloc_tile : !pto.tile_buf + pto.tci ins(%c0_i16, %tmp_tile : i16, !pto.tile_buf) outs(%dst_tile : !pto.tile_buf) pto.tstore ins(%dst_tile : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x16xi16>) {layout = #pto.layout, pto.inferred_layout = true} diff --git a/test/lit/pto/tci_i16_emitc.pto b/test/lit/pto/tci_i16_emitc.pto index cb7f1558c6..4fd5917dd8 100644 --- a/test/lit/pto/tci_i16_emitc.pto +++ b/test/lit/pto/tci_i16_emitc.pto @@ -10,5 +10,5 @@ module { } } -// A3: TCI, int16_t, 0>( -// A3-NOT: TCI, int32_t, 0>( +// A3: TCI, {{.*}}, int16_t, 0>({{.*}}, {{.*}}, {{.*}}) +// A3-NOT: TCI, {{.*}}, int32_t, 0>( diff --git a/test/lit/pto/tci_implicit_tmp_level3_invalid.pto b/test/lit/pto/tci_implicit_tmp_level3_invalid.pto new file mode 100644 index 0000000000..bdf628d5c1 --- /dev/null +++ b/test/lit/pto/tci_implicit_tmp_level3_invalid.pto @@ -0,0 +1,13 @@ +// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s + +module { + func.func @tci_implicit_tmp_level3() { + %c0_i32 = arith.constant 0 : i32 + %addr = arith.constant 0 : i64 + %tile = pto.alloc_tile addr = %addr : !pto.tile_buf + // CHECK: error: 'pto.tci' op requires explicit tmp when PlanMemory is skipped + pto.tci ins(%c0_i32 : i32) + outs(%tile : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/tci_implicit_tmp_materialization.pto b/test/lit/pto/tci_implicit_tmp_materialization.pto new file mode 100644 index 0000000000..78cec449e5 --- /dev/null +++ b/test/lit/pto/tci_implicit_tmp_materialization.pto @@ -0,0 +1,25 @@ +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --pto-level=level2 %s 2>&1 | FileCheck %s --check-prefix=CPP +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5IR + +module { + func.func @tci_implicit_tmp(%dst: memref<32xi32, #pto.address_space>) { + %c0_i32 = arith.constant 0 : i32 + %src = memref.reinterpret_cast %dst to offset: [0], sizes: [1, 32], strides: [32, 1] {layout = #pto.layout} : memref<32xi32, #pto.address_space> to memref<1x32xi32, strided<[32, 1], offset: ?>, #pto.address_space> + %tile = pto.alloc_tile : !pto.tile_buf + pto.tci ins(%c0_i32 : i32) + outs(%tile : !pto.tile_buf) + pto.tstore ins(%tile : !pto.tile_buf) + outs(%src : memref<1x32xi32, strided<[32, 1], offset: ?>, #pto.address_space>) {layout = #pto.layout, pto.inferred_layout = true} + return + } +} + +// IR: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// IR: pto.tci ins(%{{.*}}, %{{.*}} : i32, !pto.tile_buf) +// IR-NOT: memref.alloc + +// CPP: TCI<{{.*}}, {{.*}}, int32_t, 0>({{.*}}, {{.*}}, {{.*}}) + +// A5IR: pto.tci ins(%{{.*}} : i32) outs( +// A5IR-NOT: !pto.tile_buf diff --git a/test/lit/pto/tci_tmp_contract_a3_invalid.pto b/test/lit/pto/tci_tmp_contract_a3_invalid.pto new file mode 100644 index 0000000000..1694faf370 --- /dev/null +++ b/test/lit/pto/tci_tmp_contract_a3_invalid.pto @@ -0,0 +1,27 @@ +// RUN: not ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @tci_b32_tmp_too_small() { + %c0_i32 = arith.constant 0 : i32 + %dst = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + // CHECK: error: 'pto.tci' op expects A2/A3 tmp capacity to be at least 768 bytes for 32-bit dst element type + pto.tci ins(%c0_i32, %tmp : i32, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @tci_b16_tmp_too_small() { + %c0_i16 = arith.constant 0 : i16 + %dst = pto.alloc_tile + : !pto.tile_buf + %tmp = pto.alloc_tile + : !pto.tile_buf + // CHECK: error: 'pto.tci' op expects A2/A3 tmp capacity to be at least 1792 bytes for 16-bit dst element type + pto.tci ins(%c0_i16, %tmp : i16, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/tci_ui16_emitc.pto b/test/lit/pto/tci_ui16_emitc.pto index aba69bf2f0..4d435ae994 100644 --- a/test/lit/pto/tci_ui16_emitc.pto +++ b/test/lit/pto/tci_ui16_emitc.pto @@ -11,5 +11,5 @@ module { } } -// A3: TCI<{{.*}}, uint16_t, 0>( -// A3-NOT: TCI<{{.*}}, int16_t, 0>( +// A3: TCI<{{.*}}, {{.*}}, uint16_t, 0>({{.*}}, {{.*}}, {{.*}}) +// A3-NOT: TCI<{{.*}}, {{.*}}, int16_t, 0>( diff --git a/test/lit/pto/tci_ui32_emitc.pto b/test/lit/pto/tci_ui32_emitc.pto index ef6cc2a31c..beb7d241b7 100644 --- a/test/lit/pto/tci_ui32_emitc.pto +++ b/test/lit/pto/tci_ui32_emitc.pto @@ -11,5 +11,5 @@ module { } } -// A3: TCI<{{.*}}, uint32_t, 0>( -// A3-NOT: TCI<{{.*}}, int32_t, 0>( +// A3: TCI<{{.*}}, {{.*}}, uint32_t, 0>({{.*}}, {{.*}}, {{.*}}) +// A3-NOT: TCI<{{.*}}, {{.*}}, int32_t, 0>( diff --git a/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto b/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto new file mode 100644 index 0000000000..34f9d32a7c --- /dev/null +++ b/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto @@ -0,0 +1,24 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s + +module { + func.func @trowexpand_mode1_level3_requires_tmp() { + %addr0 = arith.constant 0 : i64 + %addr1 = arith.constant 1024 : i64 + %addr2 = arith.constant 2048 : i64 + %src0 = pto.alloc_tile addr = %addr0 : !pto.tile_buf + %src1 = pto.alloc_tile addr = %addr1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %addr2 : !pto.tile_buf + // CHECK: error: 'pto.trowexpandadd' op requires explicit tmp for A2/A3 row-expand mode 1 when PlanMemory is skipped + pto.trowexpandadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/trowexpand_implicit_tmp_materialization.pto b/test/lit/pto/trowexpand_implicit_tmp_materialization.pto new file mode 100644 index 0000000000..1ffe8abc36 --- /dev/null +++ b/test/lit/pto/trowexpand_implicit_tmp_materialization.pto @@ -0,0 +1,42 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 + +module { + func.func @trowexpand_mode1_add_materializes_tmp() { + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.trowexpandadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @trowexpand_mode2_add_keeps_no_tmp() { + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.trowexpandadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// A3-LABEL: func.func @trowexpand_mode1_add_materializes_tmp +// A3: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// A3: pto.trowexpandadd ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) +// A3-LABEL: func.func @trowexpand_mode2_add_keeps_no_tmp +// A3: pto.trowexpandadd ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3-NOT: !pto.tile_buf + +// A5-LABEL: func.func @trowexpand_mode1_add_materializes_tmp +// A5: pto.trowexpandadd ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5-LABEL: func.func @trowexpand_mode2_add_keeps_no_tmp +// A5-NOT: !pto.tile_buf diff --git a/test/lit/pto/trowexpand_tmp_contract_a3_invalid.pto b/test/lit/pto/trowexpand_tmp_contract_a3_invalid.pto new file mode 100644 index 0000000000..bea63f9b29 --- /dev/null +++ b/test/lit/pto/trowexpand_tmp_contract_a3_invalid.pto @@ -0,0 +1,33 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @trowexpand_mode2_rejects_explicit_tmp() { + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + // CHECK: error: 'pto.trowexpandadd' op expects A2/A3 tmp-form trowexpand to use mode 1 + pto.trowexpandadd ins(%src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @trowexpand_mode1_rejects_small_tmp() { + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + // CHECK: error: 'pto.trowexpandmax' op expects A2/A3 trowexpand tmp capacity to be at least 512 bytes, but got 256 bytes + pto.trowexpandmax ins(%src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 8fdaa31afa..7781a4281f 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -3460,6 +3460,9 @@ int mlir::pto::compilePTOASModule( pm.addNestedPass(pto::createPTOFusionRegionGenPass()); } + pm.addNestedPass( + pto::createPTOMaterializeImplicitTmpPass( + effectiveLevel == PTOBuildLevel::Level3)); pm.addNestedPass( pto::createPTORematerializeFixpipeVectorQuantPass()); From 17a59c23f98fda99e79d69ef28f0faa13002ee9e Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 17 Jul 2026 11:28:06 +0800 Subject: [PATCH 02/17] all rest ops --- ...oas-implicit-tmp-materialization-design.md | 295 ++++++ include/PTO/IR/PTOOps.td | 130 +-- lib/PTO/IR/PTO.cpp | 881 +++++++++++++++--- .../InsertSync/InsertSyncAnalysis.cpp | 49 + .../Transforms/PTOMaterializeImplicitTmp.cpp | 685 ++++++++++++++ lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 8 +- lib/PTO/Transforms/PTOToEmitC.cpp | 6 +- test/lit/pto/implicit_tmp_arg_reductions.pto | 31 + .../implicit_tmp_optional_level3_invalid.pto | 24 + ...licit_tmp_optional_ops_materialization.pto | 80 ++ .../implicit_tmp_remaining_level3_invalid.pto | 46 + ...icit_tmp_remaining_ops_materialization.pto | 42 + test/lit/pto/implicit_tmp_row_reductions.pto | 31 + .../pto/implicit_tmp_xor_materialization.pto | 27 + ...ssue533_loop_zero_trip_sync_regression.pto | 8 +- ...533_loop_zero_trip_sync_regression_gss.pto | 8 +- test/lit/pto/issue646_pipev_repeat_prune.pto | 16 +- test/lit/pto/tpow_fp_missing_tmp_invalid.pto | 11 +- 18 files changed, 2143 insertions(+), 235 deletions(-) create mode 100644 test/lit/pto/implicit_tmp_arg_reductions.pto create mode 100644 test/lit/pto/implicit_tmp_optional_level3_invalid.pto create mode 100644 test/lit/pto/implicit_tmp_optional_ops_materialization.pto create mode 100644 test/lit/pto/implicit_tmp_remaining_level3_invalid.pto create mode 100644 test/lit/pto/implicit_tmp_remaining_ops_materialization.pto create mode 100644 test/lit/pto/implicit_tmp_row_reductions.pto create mode 100644 test/lit/pto/implicit_tmp_xor_materialization.pto diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md index 40ebab49f3..719563b370 100644 --- a/docs/designs/ptoas-implicit-tmp-materialization-design.md +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -566,6 +566,301 @@ A5 level3 + no tmp 不应因为 tmp 缺失报错。 - A2/A3 动态 `dst.validRow` 且显式 tmp 小于 8192 字节,按保守规则报错。 - A5 显式 tmp 不触发 A2/A3 容量校验。 +## 批量 optional tmp op 分类设计 + +本节按 PTO-ISA tmp 行为对后续待改造 op 做分类。输入列表中的重复项只记录一次: + +```text +TCOLARGMAX, TCOLARGMIN, TROWARGMAX, TROWARGMIN, +TADDDEQRELU, TGATHER, TTRANS, TXOR, TXORS, TPRELU, TMRGSORT, +TROWPROD, TCOLSUM, TROWSUM, TROWMAX, TROWMIN, +TSEL, TSELS, TRSQRT, TPOW, TPOWS, TREM, TREMS, TCVT, +TSORT32, TQUANT +``` + +其中 `TADDDEQRELU` 对应 PTO-ISA 文档 `TAddDeqRelu_zh.md`;当前 PTOAS ODS 中未找到同名 op,先标记为待 IR 接入。 + +### 分类总览 + +| 分类 | Op / 模式 | 设计结论 | +| --- | --- | --- | +| A2/A3 使用 tmp,A5 接受但不使用 | `TCOLARGMAX`、`TCOLARGMIN`、`TROWARGMAX`、`TROWARGMIN`、`TGATHER`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS`、`TQUANT`、`TADDDEQRELU` | level1/2 在 A2/A3 生成真实 scratch;若 A5 C++ 签名仍要求 tmp,则生成不带 MemoryEffects 的 ABI placeholder。 | +| A2/A3 和 A5 都可能使用 tmp | `TCOLSUM(isBinary=true)`、`TSORT32` 非 32 对齐尾部、`TMRGSORT` 多列表归并 format2 | tmp 使用由 op 模式决定,不能只按 arch 判断。 | +| 条件性 tmp,不应无条件 materialize | `TTRANS`、`TCVT`、`TPOW`、`TPOWS`、`TRSQRT`、`TMRGSORT`、`TSORT32` | 需要先判断精度、dtype、layout、format 或尾部条件。 | +| 已从 mandatory tmp 改为 optional tmp | `TTRANS`、`TXOR`、`TXORS`、`TPRELU`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TROWARGMAX`、`TROWARGMIN`、`TCOLARGMAX`、`TCOLARGMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS` | ODS、parse/print、verifier、MemoryEffects、materialize 和 lowering 已接入。 | +| 当前 PTOAS IR 已有 optional tmp | `TGATHER`、`TCOLSUM`、`TRSQRT`、`TPOW`、`TPOWS`、`TSORT32`、`TQUANT` | 可直接纳入 `pto-materialize-implicit-tmp` 的后续实现。 | +| 当前 PTOAS IR 暂无对应 op | `TADDDEQRELU` | 需先完成 PTOAS IR 接入;`TCVT` 已新增 optional tmp operand。 | + +### 通用规则 + +- level1/level2:只有该 op 在当前 arch / 模式下实际需要 tmp,且 IR 允许省略 tmp 时,才自动 materialize `pto.alloc_tile(no addr)`。 +- level3:若该 op 在当前 arch / 模式下需要 tmp 且用户省略 tmp,则报错;若当前 arch / 模式不使用 tmp,则不强制补 tmp。 +- A5 仅接受但不使用 tmp 的 op:若后端存在 no-tmp overload,则不自动补 tmp;若 C++ 签名仍要求 tmp,则自动生成 ABI placeholder。placeholder 不建模 tmp 的 Read/Write,用户显式 tmp 也不按 A2/A3 容量规则校验。 +- tmp 是 scratch 的 op:MemoryEffects 需要建模为 `Read(tmp) + Write(tmp)`,或在 semantic no-alias side table 中显式加入 `forbidAlias(tmp, dst/output)`。 +- 原 mandatory tmp op 已统一改为 optional,并保持显式 tmp 文本格式兼容。 + +### Arg reduction 类 + +覆盖: + +```text +TCOLARGMAX, TCOLARGMIN, TROWARGMAX, TROWARGMIN +``` + +现状: + +- 这些 op 的 tmp 已改为 optional,并接入 `pto-materialize-implicit-tmp`。 +- A2/A3 使用 tmp;A5 接受 tmp 但不使用。 + +TCOLARGMAX / TCOLARGMIN: + +- tmp dtype 必须与 `src` 一致。 +- tmp 用于索引跟踪和当前比较值临时存储。 +- tmp 容量需要按 `tmpGapEles` 和输出模式计算:当 `srcValidCol >= elemPerRpt` 时,`tmpGapEles = elemPerRpt`;否则 `tmpGapEles = ceil(srcValidCol / elemPerBlock) * elemPerBlock`。 +- half + 纯索引模式是 tmp 使用量最大的组合;其它类型 / 模式下 tmp 中可能只需要区域 0,但自动 materialize 可先采用覆盖最大需求的保守形状。 + +TROWARGMAX / TROWARGMIN: + +- 仅索引模式在 A2/A3 可能不使用 tmp;值+索引模式和两阶段归约需要 tmp。 +- tmp 行数与 `src` 相同;每行 stride 按 PTO-ISA 文档公式计算。 +- 当前 PTOAS ODS 只有单输出索引模式;该模式仍需满足后端显式 tmp 参数签名,因此 level1/2 生成保守同形状 tmp。未来接入值+索引模式后,再按输出模式和归约阶段收紧容量。 + +MemoryEffects / alias: + +- A2/A3 实际使用 tmp 时建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dstIdx/dstVal)`。 +- tmp 不应与同 op 的输出 alias;如果 MemoryEffects 无法覆盖,应加入 `forbidAlias(tmp, dstIdx)` 和必要的 `forbidAlias(tmp, dstVal)`。 + +### Row reduction 类 + +覆盖: + +```text +TROWPROD, TROWSUM, TROWMAX, TROWMIN +``` + +现状: + +- 这些 op 的 tmp 已改为 optional,并接入 `pto-materialize-implicit-tmp`。 +- A2/A3 使用 tmp;A5 接受 tmp 但不使用。 + +tmp 规格: + +- tmp dtype 与 `src` / `dst` 一致。 +- 整数路径最小需要 1 个 vector block:`int32` 为 8 列,`int16` 为 16 列。 +- 浮点路径用于二叉树归约;安全默认形状可设为与 `src` 相同。 +- `TROWPROD` 的安全默认也可设为与 `src` 相同;最小需求为 1 行和 1 个 vector block。 + +Pass 行为: + +- A2/A3 level1/2:若 IR 已支持 optional tmp 且缺省 tmp,则自动生成 vec row-major none-box tmp。 +- A5:生成后端签名需要的 ABI placeholder,但不为其添加 tmp MemoryEffects。 +- level3:A2/A3 需要 tmp 时缺省 tmp 报错;A5 不强制 tmp。 + +MemoryEffects / alias: + +- A2/A3 建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dst)`。 +- tmp 与 `dst` 禁止 alias。 + +### Column sum 类 + +覆盖: + +```text +TCOLSUM +``` + +现状: + +- 当前 PTOAS IR 已支持 optional tmp。 +- no-tmp 形式表示顺序累加,不需要 tmp。 +- `isBinary=true` 时 A2/A3 和 A5 都使用 tmp 做二叉树累加。 + +tmp 规格: + +- tmp dtype 与 `src` / `dst` 一致。 +- tmp 为 vec row-major none-box tile。 +- `tmp.validCol >= src.validCol`。 +- `tmp.validRow >= ceil(src.validRow / 2)`。 + +Pass 行为: + +- 仅当 `isBinary=true` 且缺省 tmp 时自动 materialize。 +- `isBinary=false` 不自动补 tmp。 +- level3 下 `isBinary=true` 且缺省 tmp 报错。 + +MemoryEffects / alias: + +- `isBinary=true` 且 tmp 存在时,建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dst)`。 +- `isBinary=false` 且 tmp 缺省时,保持 no-tmp 顺序累加语义。 + +### Elementwise scratch 类 + +覆盖: + +```text +TXOR, TXORS, TPRELU, TREM, TREMS, TADDDEQRELU, TQUANT +``` + +现状: + +- `TXOR`、`TXORS`、`TPRELU`、`TREM`、`TREMS` 的 tmp 均已改为 optional 并接入 materialize pass。 +- `TQUANT` 当前 PTOAS IR 已支持 optional tmp。 +- `TADDDEQRELU` 当前 PTOAS ODS 中未找到同名 op,需先完成 IR 接入。 + +tmp 规格: + +- `TXOR` / `TXORS`:A2/A3 tmp dtype 与输入输出一致,row-major,容量覆盖 `dst` 有效区域;A5 不使用 tmp。 +- `TPRELU`:A2/A3 tmp dtype 为 `uint8_t`,row-major,`tmp.validRow > dst.validRow`,用于 mask buffer;A5 不使用 tmp。 +- `TREM`:A2/A3 tmp dtype 与 `dst` 一致,至少 2 行和 `dst.validCol` 列;A5 不使用 tmp。 +- `TREMS`:A2/A3 tmp dtype 与 `dst` 一致,至少 1 行和 `dst.validCol` 列;A5 不使用 tmp。 +- `TADDDEQRELU`:A2/A3 tmp dtype 为 `int32_t`,容量至少覆盖 `dst` 有效区域;A5 不使用 tmp。 +- `TQUANT`:A2/A3 tmp 为 FP32,形状与 `src` 同尺寸,用作 FP32 到 S32 转换中间结果;A5 不使用 tmp。 + +Pass 行为: + +- A2/A3 level1/2:缺省 tmp 且 IR 支持 optional tmp 时自动 materialize。 +- A5:`TXOR/TXORS` 因后端签名要求生成 ABI placeholder;其余 op 按各自后端是否存在 no-tmp overload 决定。 +- level3:A2/A3 缺省 tmp 报错;A5 不强制 tmp。 + +MemoryEffects / alias: + +- A2/A3 tmp 是 scratch 时建模 `Read(tmp) + Write(tmp)`。 +- tmp 与 `dst` 禁止 alias;`TXOR/TXORS/TPRELU/TREM/TREMS/TADDDEQRELU` 还应禁止 tmp 与同 op 输入错误 alias。 + +### Mask select 类 + +覆盖: + +```text +TSEL, TSELS +``` + +现状: + +- 当前 PTOAS IR 中 tmp 是 mandatory;隐式 tmp 支持前需要先改成 optional tmp。 +- A2/A3 使用 tmp;A5 接受 tmp 但不使用。 + +tmp 规格: + +- `TSEL`:tmp dtype 为 `uint32_t`,用于 mask buffer。16 位数据类型的 `cmpmaskLen = 4` 个 `uint32_t`;32 位数据类型的 `cmpmaskLen = 2` 个 `uint32_t`。 +- `TSELS`:tmp dtype 与 `src` 一致,至少 1 个元素,用于保存 scalar 和比较 mask。 + +Pass 行为: + +- A2/A3 level1/2:缺省 tmp 时自动 materialize。 +- A5:不自动补 tmp。 +- level3:A2/A3 缺省 tmp 报错;A5 不强制 tmp。 + +MemoryEffects / alias: + +- A2/A3 建模 `Read(mask/src) + Read(tmp) + Write(tmp) + Write(dst)`。 +- tmp 与 `dst` 禁止 alias。 + +### Data movement / layout 类 + +覆盖: + +```text +TGATHER, TTRANS, TCVT +``` + +TGATHER: + +- 当前 PTOAS IR 已支持 optional tmp。 +- index form:A2/A3 C++ API 需要 tmp;tmp dtype 与 indices dtype 一致,shape 覆盖 indices;A5 不使用 tmp。 +- compare form:A2/A3 tmp 是合并暂存缓冲区,包含 `cmpsTmp`、`indexTmp`、`cvtTmp` 三个区域;最小字节数按 PTO-ISA 文档公式计算;A5 不使用 tmp。 +- mask form 不使用 tmp,不应自动补 tmp。 +- A2/A3 level3 下 index / compare form 缺省 tmp 报错。 + +TTRANS: + +- 当前 PTOAS IR 中 tmp 已改为 optional,并接入 materialize pass。 +- tmp 只在满足高效转置路径条件时使用;scalar copy 和部分 layout 转换不需要 tmp。 +- 静态满足 stride 条件时生成与 src 同形状的保守 scratch;不满足并走 scalar copy 时生成 32 字节 ABI placeholder。 +- 只有真实 scratch 进入 `Read(tmp) + Write(tmp)` MemoryEffects;scalar-copy placeholder 不建模内存访问。 + +TCVT: + +- 当前 PTOAS IR 已新增 optional tmp operand,并完成 parser/printer、verifier、MemoryEffects、materialize 和 EmitC lowering。 +- A2/A3 仅在 `SaturationMode::OFF` 的 PyTorch 兼容非饱和窄化路径使用 tmp:`float -> int16`、`half -> int16`、`half -> int8`。 +- 其它转换不需要 tmp,不应自动 materialize。 +- tmp 按字节规划,容量使用 PTO-ISA 的 `tmpFloatToInt16Bytes`、`tmpHalfToInt16Bytes`、`tmpHalfToInt8Bytes` 公式。 +- level3 下上述三种路径缺省 tmp 报错;其它转换继续使用 no-tmp overload。 + +MemoryEffects / alias: + +- 只有实际使用 tmp 的路径建模 `Read(tmp) + Write(tmp)`。 +- tmp 与 `dst` 禁止 alias。 + +### Sort / merge 类 + +覆盖: + +```text +TSORT32, TMRGSORT +``` + +TSORT32: + +- 当前 PTOAS IR 已支持 optional tmp。 +- 3 参数形式适用于 `validCol` 已按 32 对齐的路径,不需要 tmp。 +- 4 参数形式用于非 32 对齐尾部,通过 tmp 保存填充后的行或尾块副本。 +- tmp dtype 与 `src` 一致;容量按 PTO-ISA `tmpSize` 公式计算,不能固定为 8KB。 +- level1/2 仅在能静态证明存在非 32 对齐尾部时自动补 tmp;否则保持 no-tmp。 +- level3 下非 32 对齐尾部缺省 tmp 报错。 + +TMRGSORT: + +- 当前 PTOAS IR 已支持 optional tmp;缺省 format2 使用显式 `no_tmp` 中间语法消除 src/tmp 个数歧义。 +- format1 单输入 block sort 不需要 tmp。 +- format2 多列表归并需要 tmp 和 executed list。 +- level1/2 对 format2 `no_tmp` 生成一行 row-major tmp,`tmp.cols = sum(src.cols)`;format1 不补 tmp。 +- level3 下 format2 `no_tmp` 报错。 + +MemoryEffects / alias: + +- 使用 tmp 的 sort / merge 路径建模 `Read(tmp) + Write(tmp)`。 +- tmp 与 `dst` 禁止 alias;`TMRGSORT` 还需考虑 `executed` output 的写 effect。 + +### Pow / rsqrt 类 + +覆盖: + +```text +TPOW, TPOWS, TRSQRT +``` + +TPOW / TPOWS: + +- 当前 PTOAS IR 已支持 optional tmp。 +- A2/A3 浮点路径使用 tmp;整数路径不使用 tmp。 +- A5 接受 tmp 但不使用。 +- tmp dtype 与 `dst` / `base` 一致,容量覆盖 `dst` 有效区域。 +- level1/2:A2/A3 浮点路径缺省 tmp 时自动 materialize;整数路径不补 tmp。 +- level3:A2/A3 浮点路径缺省 tmp 报错;整数路径不强制 tmp。 + +TRSQRT: + +- 当前 PTOAS IR 已支持 optional tmp。 +- no-tmp 默认实现不需要 tmp。 +- 带 tmp overload 当前仅作为 API 兼容 / 未来高精度路径保留,现阶段不自动 materialize。 +- 用户显式 tmp 时,A5 不按 A2/A3 scratch 容量规则校验。 + +MemoryEffects / alias: + +- `TPOW/TPOWS` 只有浮点 tmp-backed 路径建模 `Read(tmp) + Write(tmp)`。 +- `TRSQRT` 现阶段不因缺省 tmp 增加 MemoryEffects。 + +### 批量 op 测试策略 + +后续按类别实现时,每一类至少补以下 lit: + +- 自动补 tmp:level1/2 + A2/A3 + 缺省 tmp,检查 `pto.alloc_tile(no addr)` 经 memplan 后带 `addr`。 +- A5 不补 tmp:A5 + 缺省 tmp,检查保持 no-tmp 形态。 +- level3 负例:A2/A3 + 需要 tmp + 缺省 tmp 报错。 +- verifier 负例:显式 tmp dtype / shape / capacity / layout 不满足对应 op 规格时报错。 +- EmitC overload:需要 tmp 的路径最终走 tmp-aware C++ 调用;不需要 tmp 的路径不强行切换。 + ## 后续扩展 后续新增其它 optional tmp op 时,需要补充一个 op-specific 小节,并明确: diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 69f1e6c149..aa43df7569 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -861,16 +861,12 @@ def TTransOp : PTO_TOp<"ttrans", [ }]; let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); let results = (outs); - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let hasVerifier = 1; @@ -4070,7 +4066,7 @@ def TColArgMaxOp : PTO_TOp<"tcolargmax", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -4083,11 +4079,7 @@ def TColArgMaxOp : PTO_TOp<"tcolargmax", [ ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } }]; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; } def TColMinOp : PTO_TOp<"tcolmin", [ @@ -4127,7 +4119,7 @@ def TColArgMinOp : PTO_TOp<"tcolargmin", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -4140,11 +4132,7 @@ def TColArgMinOp : PTO_TOp<"tcolargmin", [ ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } }]; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; } def TColSumOp : PTO_TOp<"tcolsum", [ @@ -4214,6 +4202,7 @@ def TCvtOp : PTO_TOp<"tcvt", [ let arguments = (ins PTODpsType:$src, + Optional:$tmp, PTODpsType:$dst, DefaultValuedAttr:$rmode, DefaultValuedAttr:$sat_mode @@ -5149,6 +5138,7 @@ def TMrgSortOp: PTO_TOp<"tmrgsort", [ let extraClassDeclaration = [{ bool isFormat1() { return getSrcs().size() == 1u && getBlockLen() && getDsts().size() == 1u; } bool isFormat2() { return getSrcs().size() >= 2u && getSrcs().size() <= 4u && getTmp() && getDsts().size() == 1u && getExcuted(); } + bool isFormat2WithoutTmp() { return getSrcs().size() >= 2u && getSrcs().size() <= 4u && !getTmp() && !getBlockLen() && getDsts().size() == 1u && getExcuted(); } Value getSrc() { return getSrcs().front(); } Value getDst() { return getDsts().front(); } ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstsMutable(); } @@ -5565,7 +5555,7 @@ def TPReluOp: PTO_TOp<"tprelu", [ let arguments = (ins PTODpsType:$src0, PTODpsType:$src1, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -5573,11 +5563,7 @@ def TPReluOp: PTO_TOp<"tprelu", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src0 `,` $src1 `,` $tmp `:` qualified(type($src0)) `,` qualified(type($src1)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -5815,7 +5801,7 @@ def TRemOp: PTO_TOp<"trem", [ let arguments = (ins PTODpsType:$src0, PTODpsType:$src1, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst, DefaultValuedAttr:$precisionType ); @@ -5824,11 +5810,7 @@ def TRemOp: PTO_TOp<"trem", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src0 `,` $src1 `,` $tmp `:` qualified(type($src0)) `,` qualified(type($src1)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -5849,7 +5831,7 @@ def TRemSOp: PTO_TOp<"trems", [ let arguments = (ins PTODpsType:$src, ScalarType:$scalar, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -5857,11 +5839,7 @@ def TRemSOp: PTO_TOp<"trems", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $scalar `,` $tmp `:` qualified(type($src)) `,` type($scalar) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6154,7 +6132,7 @@ def TRowMaxOp: PTO_TOp<"trowmax", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6162,11 +6140,7 @@ def TRowMaxOp: PTO_TOp<"trowmax", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6183,7 +6157,7 @@ def TRowArgMaxOp: PTO_TOp<"trowargmax", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6191,11 +6165,7 @@ def TRowArgMaxOp: PTO_TOp<"trowargmax", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6215,7 +6185,7 @@ def TRowMinOp: PTO_TOp<"trowmin", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6223,11 +6193,7 @@ def TRowMinOp: PTO_TOp<"trowmin", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6244,7 +6210,7 @@ def TRowArgMinOp: PTO_TOp<"trowargmin", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6252,11 +6218,7 @@ def TRowArgMinOp: PTO_TOp<"trowargmin", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6276,7 +6238,7 @@ def TRowSumOp: PTO_TOp<"trowsum", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6284,11 +6246,7 @@ def TRowSumOp: PTO_TOp<"trowsum", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6305,7 +6263,7 @@ def TRowProdOp: PTO_TOp<"trowprod", [ let arguments = (ins PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6313,11 +6271,7 @@ def TRowProdOp: PTO_TOp<"trowprod", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $tmp `:` qualified(type($src)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6433,7 +6387,7 @@ def TSelOp: PTO_TOp<"tsel", [ PTODpsType:$mask, PTODpsType:$src0, PTODpsType:$src1, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6441,11 +6395,7 @@ def TSelOp: PTO_TOp<"tsel", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $mask `,` $src0 `,` $src1 `,` $tmp `:` qualified(type($mask)) `,` qualified(type($src0)) `,` qualified(type($src1)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6469,7 +6419,7 @@ def TSelSOp: PTO_TOp<"tsels", [ let arguments = (ins PTODpsType:$mask, PTODpsType:$src, - PTODpsType:$tmp, + Optional:$tmp, ScalarType:$scalar, PTODpsType:$dst ); @@ -6478,11 +6428,7 @@ def TSelSOp: PTO_TOp<"tsels", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $mask `,` $src `,` $tmp `,` $scalar `:` qualified(type($mask)) `,` qualified(type($src)) `,` qualified(type($tmp)) `,` type($scalar) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6857,7 +6803,7 @@ def TXorSOp: PTO_TOp<"txors", [ let arguments = (ins PTODpsType:$src, AnySignlessInteger:$scalar, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6865,11 +6811,7 @@ def TXorSOp: PTO_TOp<"txors", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src `,` $scalar `,` $tmp `:` qualified(type($src)) `,` type($scalar) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } @@ -6896,7 +6838,7 @@ def TXorOp: PTO_TOp<"txor", [ let arguments = (ins PTODpsType:$src0, PTODpsType:$src1, - PTODpsType:$tmp, + Optional:$tmp, PTODpsType:$dst ); @@ -6904,11 +6846,7 @@ def TXorOp: PTO_TOp<"txor", [ let hasVerifier = 1; - let assemblyFormat = [{ - `ins` `(` $src0 `,` $src1 `,` $tmp `:` qualified(type($src0)) `,` qualified(type($src1)) `,` qualified(type($tmp)) `)` - `outs` `(` $dst `:` qualified(type($dst) ) `)` - attr-dict - }]; + let hasCustomAssemblyFormat = 1; let extraClassDeclaration = [{ ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; } diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 90362862ab..984e0ad7d4 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -2238,6 +2238,25 @@ static LogicalResult verifyTColArgReductionOpA2A3(Operation *op, Type srcTy, return success(); } +static LogicalResult verifyTColArgReductionNoTmp(Operation *op, Type srcTy, + Type dstTy) { + if (failed(verifyNDStyleVecTile(op, srcTy, "src")) || + failed(verifyColArgReductionDstLayout(op, dstTy, "dst")) || + failed(verifyColReductionValidRegion(op, srcTy, dstTy, + /*requireNonZeroSrc=*/true))) + return failure(); + Type srcElemTy = getElemTy(srcTy); + unsigned srcElemBits = srcElemTy ? getPTOStorageElemBitWidth(srcElemTy) : 0; + if (!(mlir::isa(srcElemTy) && + (srcElemBits == 8 || srcElemBits == 16 || srcElemBits == 32))) + return op->emitOpError( + "expects src element type to be 1, 2, or 4 bytes wide"); + auto dstInt = dyn_cast(getElemTy(dstTy)); + if (!dstInt || dstInt.getWidth() != 32) + return op->emitOpError("expects dst element type to be i32 or ui32"); + return success(); +} + static LogicalResult verifyTColArgReductionOpA5(Operation *op, Type srcTy, Type tmpTy, Type dstTy) { if (failed(verifyNDStyleVecTile(op, srcTy, "src")) || @@ -2371,6 +2390,22 @@ static LogicalResult verifyTRowArgReductionOpA2A3(Operation *op, Type srcTy, return success(); } +static LogicalResult verifyTRowArgReductionNoTmp(Operation *op, Type srcTy, + Type dstTy) { + if (failed(verifyRowReductionSrcLayout(op, srcTy, "src")) || + failed(verifyRowReductionDstLayout(op, dstTy, "dst")) || + failed(verifyRowReductionValidRegion(op, srcTy, dstTy, + /*allowEmptyMarker=*/false))) + return failure(); + Type srcElem = getElemTy(srcTy); + if (!isSupportedRowReductionElemType(srcElem)) + return op->emitOpError("expects src element type to be i16/i32/f16/f32"); + auto dstInt = dyn_cast(getElemTy(dstTy)); + if (!dstInt || dstInt.getWidth() != 32) + return op->emitOpError("expects dst element type to be i32 or ui32"); + return success(); +} + static LogicalResult verifyTRowArgReductionOpA5(Operation *op, Type srcTy, Type tmpTy, Type dstTy) { if (failed(verifyRowReductionSrcLayout(op, srcTy, "src")) || @@ -6194,6 +6229,9 @@ LogicalResult pto::TColMaxOp::verify() { } LogicalResult pto::TColArgMaxOp::verify() { + if (!getTmp()) + return verifyTColArgReductionNoTmp(getOperation(), getSrc().getType(), + getDst().getType()); auto verifyA2A3 = [&]() -> LogicalResult { return verifyTColArgReductionOpA2A3(*this, getSrc().getType(), getTmp().getType(), getDst().getType()); @@ -6216,6 +6254,9 @@ LogicalResult pto::TColMinOp::verify() { } LogicalResult pto::TColArgMinOp::verify() { + if (!getTmp()) + return verifyTColArgReductionNoTmp(getOperation(), getSrc().getType(), + getDst().getType()); auto verifyA2A3 = [&]() -> LogicalResult { return verifyTColArgReductionOpA2A3(*this, getSrc().getType(), getTmp().getType(), getDst().getType()); @@ -6325,11 +6366,8 @@ LogicalResult pto::TColSumOp::verify() { return failure(); bool hasTmp = (bool)getTmp(); bool hasIsBinary = (bool)getIsBinaryAttr(); - if (hasTmp != hasIsBinary) { - if (hasTmp) - return emitOpError("tmp operand requires isBinary attribute"); - return emitOpError("isBinary attribute requires tmp operand"); - } + if (hasTmp && !hasIsBinary) + return emitOpError("tmp operand requires isBinary attribute"); if (getTmp()) { Type tmpTy = getTmp().getType(); if (failed(verifyNDStyleVecTile(*this, tmpTy, "tmp"))) @@ -6357,11 +6395,8 @@ LogicalResult pto::TColSumOp::verify() { return failure(); bool hasTmp = (bool)getTmp(); bool hasIsBinary = (bool)getIsBinaryAttr(); - if (hasTmp != hasIsBinary) { - if (hasTmp) - return emitOpError("tmp operand requires isBinary attribute"); - return emitOpError("isBinary attribute requires tmp operand"); - } + if (hasTmp && !hasIsBinary) + return emitOpError("tmp operand requires isBinary attribute"); if (getTmp()) { Type tmpTy = getTmp().getType(); if (failed(verifyNDStyleVecTile(*this, tmpTy, "tmp"))) @@ -6419,14 +6454,65 @@ llvm::LogicalResult mlir::pto::TCvtOp::verify() { return failure(); Type srcElem = getElemTy(srcTy); Type dstElem = getElemTy(dstTy); + auto needsTmp = [&]() { + if (getSatMode() != pto::SaturationMode::OFF) + return false; + return (srcElem.isF32() && dstElem.isInteger(16)) || + (srcElem.isF16() && + (dstElem.isInteger(16) || dstElem.isInteger(8))); + }; + auto verifyTmp = [&]() -> LogicalResult { + if (!getTmp()) + return success(); + Type tmpTy = getTmp().getType(); + if (failed(verifyVecTileCommon(*this, tmpTy, "tmp"))) + return failure(); + if (!needsTmp()) + return success(); + auto srcShape = getShapeVec(srcTy); + auto dstValid = getValidShapeVec(dstTy); + if (srcShape.size() != 2 || dstValid.size() != 2 || + llvm::is_contained(srcShape, ShapedType::kDynamic) || + llvm::is_contained(dstValid, ShapedType::kDynamic)) + return emitOpError( + "expects static src shape and dst valid_shape to verify tcvt tmp"); + int64_t rows = dstValid[0], cols = dstValid[1]; + int64_t requiredBytes = 0; + if (rows > 0 && cols > 0 && srcElem.isF32()) { + int64_t head = 4 * 64 * std::min(cols / 64, 255); + int64_t remainder = cols % 64; + int64_t tail = remainder == 0 + ? 0 + : 32 * ((std::min(rows, 255) - 1) * + (srcShape[1] / 8) + + llvm::divideCeil(remainder, int64_t{8})); + requiredBytes = std::max(head, tail); + } else if (cols > 0 && srcElem.isF16()) { + int64_t width = std::min(cols, 64); + int64_t halfToI16 = 32 * llvm::divideCeil(width, int64_t{8}); + int64_t halfToI8 = std::max( + halfToI16, + 128 + 32 * static_cast( + llvm::divideCeil(width, int64_t{16}))); + requiredBytes = dstElem.isInteger(8) ? halfToI8 : halfToI16; + } + auto tmpBytes = getStaticByteSize(tmpTy); + if (!tmpBytes || *tmpBytes < static_cast(requiredBytes)) + return emitOpError() + << "expects tcvt tmp capacity to be at least " << requiredBytes + << " bytes"; + return success(); + }; auto verifyA2A3 = [&]() -> LogicalResult { if (isPTOLowPrecisionType(srcElem) || isPTOLowPrecisionType(dstElem)) return emitOpError("expects A2/A3 tcvt low-precision element types to be unsupported"); - return success(); + return verifyTmp(); }; auto verifyA5 = [&]() -> LogicalResult { if (!isA5SupportedTCvtPair(srcElem, dstElem)) return emitOpError("expects A5 tcvt low-precision type pairs to match PTO-ISA support"); + if (getTmp() && failed(verifyVecTileCommon(*this, getTmp().getType(), "tmp"))) + return failure(); return success(); }; return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); @@ -7759,14 +7845,18 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { if (getAxisAttr()) return emitOpError("axis attribute must not be provided without maskPattern"); if (getCdst() || getKValue()) { - if (!getCdst() || !getKValue() || !getTmp()) - return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp"); + if (!getCdst() || !getKValue()) + return emitOpError("compare-form tgather expects dst, cdst, and kValue"); if (getIndices()) return emitOpError("compare-form tgather does not take indices"); + if (!getTmp()) + return success(); return verifyCompareForm(/*allowA5SrcTypes=*/false); } - if (!getIndices() || !getTmp()) - return emitOpError("index-form tgather expects both indices and tmp"); + if (!getIndices()) + return emitOpError("index-form tgather expects indices"); + if (!getTmp()) + return success(); return verifyIndexForm(/*allow16BitIndices=*/false, /*allowA5ElemTypes=*/false); }; @@ -7779,14 +7869,18 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { if (getAxisAttr()) return emitOpError("axis attribute must not be provided without maskPattern"); if (getCdst() || getKValue()) { - if (!getCdst() || !getKValue() || !getTmp()) - return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp"); + if (!getCdst() || !getKValue()) + return emitOpError("compare-form tgather expects dst, cdst, and kValue"); if (getIndices()) return emitOpError("compare-form tgather does not take indices"); + if (!getTmp()) + return success(); return verifyCompareForm(/*allowA5SrcTypes=*/true); } if (!getIndices()) return emitOpError("index-form tgather expects indices"); + if (!getTmp()) + return success(); return verifyIndexForm(/*allow16BitIndices=*/true, /*allowA5ElemTypes=*/true); }; @@ -9540,6 +9634,8 @@ LogicalResult MGatherOp::verify() { void mlir::pto::TCvtOp::print(OpAsmPrinter &p) { p << " ins(" << getSrc(); + if (getTmp()) + p << ", " << getTmp(); Builder builder(getContext()); NamedAttrList attrs; for (auto attr : (*this)->getAttrs()) { @@ -9551,18 +9647,28 @@ void mlir::pto::TCvtOp::print(OpAsmPrinter &p) { } p.printOptionalAttrDict(attrs.getAttrs()); p << " : " << getSrc().getType(); + if (getTmp()) + p << ", " << getTmp().getType(); p << ") outs(" << getDst() << " : " << getDst().getType() << ")"; } ParseResult mlir::pto::TCvtOp::parse(OpAsmParser &parser, OperationState &result) { - OpAsmParser::UnresolvedOperand src, dst; - Type srcTy, dstTy; + OpAsmParser::UnresolvedOperand src, tmp, dst; + Type srcTy, tmpTy, dstTy; + bool hasTmp = false; if (parser.parseKeyword("ins") || parser.parseLParen() || parser.parseOperand(src)) return failure(); + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(tmp)) + return failure(); + hasTmp = true; + } NamedAttrList attrs; if (parser.parseOptionalAttrDict(attrs) || parser.parseColonType(srcTy)) return failure(); + if (hasTmp && (parser.parseComma() || parser.parseType(tmpTy))) + return failure(); if (auto satmode = attrs.get("satmode")) { attrs.erase("satmode"); if (attrs.get("sat_mode")) @@ -9576,8 +9682,12 @@ ParseResult mlir::pto::TCvtOp::parse(OpAsmParser &parser, OperationState &result return failure(); if (parser.resolveOperand(src, srcTy, result.operands) || + (hasTmp && parser.resolveOperand(tmp, tmpTy, result.operands)) || parser.resolveOperand(dst, dstTy, result.operands)) return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr({1, hasTmp ? 1 : 0, 1})); return success(); } @@ -9586,13 +9696,17 @@ void mlir::pto::TMrgSortOp::print(OpAsmPrinter &p) { p << " ins(" << getSrc() << ", " << getBlockLen() << " : " << getSrc().getType() << ", " << getBlockLen().getType() << ") outs(" << getDst() << " : " << getDst().getType() << ")"; - } else if (isFormat2()) { + } else if (isFormat2() || isFormat2WithoutTmp()) { p << " ins("; llvm::interleaveComma(getSrcs(), p, [&](Value src) { p << src; }); - p << ", " << getTmp(); + if (getTmp()) + p << ", " << getTmp(); + else + p << " no_tmp"; p << " {exhausted = " << (getExhausted() ? "true" : "false") << "} : "; llvm::interleaveComma(getSrcs().getTypes(), p, [&](Type ty) { p << ty; }); - p << ", " << getTmp().getType(); + if (getTmp()) + p << ", " << getTmp().getType(); p << ") outs(" << getDst() << ", " << getExcuted() << " : " << getDst().getType() << ", " << getExcuted().getType() << ")"; } else { @@ -9637,10 +9751,15 @@ ParseResult mlir::pto::TMrgSortOp::parse(OpAsmParser &parser, OperationState &re return failure(); srcs.push_back(next); } - if (srcs.size() < 3 || srcs.size() > 5) - return parser.emitError(parser.getCurrentLocation(), - "tmrgsort format2 expects 2 to 4 src operands plus one tmp operand"); - OpAsmParser::UnresolvedOperand tmpOp = srcs.pop_back_val(); + bool noTmp = succeeded(parser.parseOptionalKeyword("no_tmp")); + if ((noTmp && (srcs.size() < 2 || srcs.size() > 4)) || + (!noTmp && (srcs.size() < 3 || srcs.size() > 5))) + return parser.emitError( + parser.getCurrentLocation(), + "tmrgsort format2 expects 2 to 4 src operands and optional no_tmp marker"); + OpAsmParser::UnresolvedOperand tmpOp; + if (!noTmp) + tmpOp = srcs.pop_back_val(); bool exhaustedVal = false; if (parser.parseOptionalLBrace().succeeded()) { if (parser.parseKeyword("exhausted") || parser.parseEqual()) @@ -9664,10 +9783,13 @@ ParseResult mlir::pto::TMrgSortOp::parse(OpAsmParser &parser, OperationState &re return failure(); srcTypes.push_back(nextTy); } - if (srcTypes.size() != srcs.size() + 1 || parser.parseRParen() || + if (srcTypes.size() != srcs.size() + (noTmp ? 0 : 1) || + parser.parseRParen() || parser.parseKeyword("outs") || parser.parseLParen()) return failure(); - Type tmpTy = srcTypes.pop_back_val(); + Type tmpTy; + if (!noTmp) + tmpTy = srcTypes.pop_back_val(); OpAsmParser::UnresolvedOperand dstOp, excutedOp; Type dstTy, excutedTy; if (parser.parseOperand(dstOp) || parser.parseComma() || parser.parseOperand(excutedOp) || @@ -9676,10 +9798,11 @@ ParseResult mlir::pto::TMrgSortOp::parse(OpAsmParser &parser, OperationState &re return failure(); result.addAttribute("operandSegmentSizes", parser.getBuilder().getDenseI32ArrayAttr( - {static_cast(srcs.size()), 0, 1, 1, 1})); + {static_cast(srcs.size()), 0, 1, + noTmp ? 0 : 1, 1})); if (parser.resolveOperands(srcs, srcTypes, parser.getCurrentLocation(), result.operands) || parser.resolveOperand(dstOp, dstTy, result.operands) || - parser.resolveOperand(tmpOp, tmpTy, result.operands) || + (!noTmp && parser.resolveOperand(tmpOp, tmpTy, result.operands)) || parser.resolveOperand(excutedOp, excutedTy, result.operands)) return failure(); if (parser.parseOptionalAttrDict(result.attributes)) @@ -9720,36 +9843,40 @@ mlir::LogicalResult mlir::pto::TMrgSortOp::verify() { } return mlir::success(); } - if (isFormat2()) { + if (isFormat2() || isFormat2WithoutTmp()) { for (Value v : getSrcs()) if (!isPTOShapedLike(v.getType())) return emitOpError() << "format2 expects PTO shaped-like type for each src"; if (getSrcs().size() < 2u || getSrcs().size() > 4u) return emitOpError() << "format2 expects 2 to 4 srcs"; - if (getDsts().size() != 1u || !getTmp() || !getExcuted()) - return emitOpError() << "format2 expects ins(srcs..., tmp), outs(dst), and excuted=vector"; + if (getDsts().size() != 1u || !getExcuted()) + return emitOpError() + << "format2 expects 2 to 4 srcs, one dst, and excuted=vector"; Type dstTy = getDst().getType(); - Type tmpTy = getTmp().getType(); - if (!isPTOShapedLike(dstTy) || !isPTOShapedLike(tmpTy)) + Type tmpTy = getTmp() ? getTmp().getType() : Type{}; + if (!isPTOShapedLike(dstTy) || + (tmpTy && !isPTOShapedLike(tmpTy))) return emitOpError() << "format2 dst/tmp must be PTO shaped-like"; auto excutedTy = mlir::dyn_cast(getExcuted().getType()); if (!excutedTy || excutedTy.getRank() != 1 || excutedTy.getNumElements() != 4 || !excutedTy.getElementType().isInteger(16)) return emitOpError() << "format2 excuted must be vector<4xi16>"; Type elemTy = getElemTy(dstTy); - if (elemTy != getElemTy(tmpTy)) + if (tmpTy && elemTy != getElemTy(tmpTy)) return emitOpError() << "format2 expects dst/tmp element types to match"; auto dstShape = getShapeVec(dstTy); - auto tmpShape = getShapeVec(tmpTy); - if (dstShape.size() != 2 || tmpShape.size() != 2) + auto tmpShape = tmpTy ? getShapeVec(tmpTy) : SmallVector{}; + if (dstShape.size() != 2 || (tmpTy && tmpShape.size() != 2)) return emitOpError() << "format2 expects dst/tmp to be rank-2 tile-shaped"; if ((dstShape[0] != mlir::ShapedType::kDynamic && dstShape[0] != 1) || - (tmpShape[0] != mlir::ShapedType::kDynamic && tmpShape[0] != 1)) + (tmpTy && tmpShape[0] != mlir::ShapedType::kDynamic && + tmpShape[0] != 1)) return emitOpError() << "format2 expects dst/tmp rows == 1"; - if (dstShape[1] != mlir::ShapedType::kDynamic && + if (tmpTy && dstShape[1] != mlir::ShapedType::kDynamic && tmpShape[1] != mlir::ShapedType::kDynamic && tmpShape[1] < dstShape[1]) return emitOpError() << "format2 expects tmp.cols >= dst.cols"; + int64_t requiredTmpCols = 0; for (Value src : getSrcs()) { Type srcTy = src.getType(); auto srcShape = getShapeVec(srcTy); @@ -9759,7 +9886,17 @@ mlir::LogicalResult mlir::pto::TMrgSortOp::verify() { return emitOpError() << "format2 expects src rows == 1"; if (getElemTy(srcTy) != elemTy) return emitOpError() << "format2 expects src/dst/tmp element types to match"; + if (srcShape[1] == mlir::ShapedType::kDynamic) + requiredTmpCols = mlir::ShapedType::kDynamic; + else if (requiredTmpCols != mlir::ShapedType::kDynamic) + requiredTmpCols += srcShape[1]; } + if (tmpTy && requiredTmpCols != mlir::ShapedType::kDynamic && + tmpShape[1] != mlir::ShapedType::kDynamic && + tmpShape[1] < requiredTmpCols) + return emitOpError() + << "format2 expects tmp.cols >= sum(src.cols) = " + << requiredTmpCols; return mlir::success(); } return emitOpError() << "tmrgsort expects format1 (1 src + blockLen + 1 dst) or " @@ -10275,16 +10412,17 @@ mlir::LogicalResult mlir::pto::TPReluOp::verify() { auto verifyCommon = [&]() -> FailureOr> { Type t0 = getSrc0().getType(); Type t1 = getSrc1().getType(); - Type tt = getTmp().getType(); + Type tt = getTmp() ? getTmp().getType() : Type{}; Type td = getDst().getType(); if (failed(verifyTileBufCommon(*this, t0, "src0")) || failed(verifyTileBufCommon(*this, t1, "src1")) || - failed(verifyTileBufCommon(*this, tt, "tmp")) || failed(verifyTileBufCommon(*this, td, "dst"))) return failure(); + if (tt && failed(verifyTileBufCommon(*this, tt, "tmp"))) + return failure(); - Type e0 = getElemTy(t0), e1 = getElemTy(t1), et = getElemTy(tt), ed = getElemTy(td); - if (!e0 || !e1 || !et || !ed) { + Type e0 = getElemTy(t0), e1 = getElemTy(t1), ed = getElemTy(td); + if (!e0 || !e1 || !ed) { emitOpError("failed to get element type for operands"); return failure(); } @@ -10317,6 +10455,8 @@ mlir::LogicalResult mlir::pto::TPReluOp::verify() { if (failed(tysOr)) return failure(); auto [t0, t1, tt, td] = *tysOr; + if (!tt) + return success(); Type tmpElem = getElemTy(tt); auto tmpIntTy = mlir::dyn_cast(tmpElem); if (!tmpIntTy || tmpIntTy.getWidth() != 8) @@ -10358,7 +10498,7 @@ mlir::LogicalResult mlir::pto::TPReluOp::verify() { (void)t0; (void)t1; (void)td; - if (failed(verifyVecTileCommon(*this, tt, "tmp"))) + if (tt && failed(verifyVecTileCommon(*this, tt, "tmp"))) return failure(); return success(); }; @@ -10859,11 +10999,9 @@ mlir::LogicalResult mlir::pto::TRemOp::verify() { Type src0Ty = getSrc0().getType(); Type src1Ty = getSrc1().getType(); - Type tmpTy = getTmp().getType(); Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, src0Ty, "src0")) || failed(verifyTileBufCommon(*this, src1Ty, "src1")) || - failed(verifyTileBufCommon(*this, tmpTy, "tmp")) || failed(verifyTileBufCommon(*this, dstTy, "dst"))) return failure(); if (failed(verifyTileBufSameElemType(*this, src0Ty, src1Ty, "src0", "src1")) || @@ -10875,11 +11013,30 @@ mlir::LogicalResult mlir::pto::TRemOp::verify() { !isRowMajorTileBuf(dstTy)) return emitOpError("expects src0, src1, and dst to use row-major layout"); auto dstValid = getValidShapeVec(dstTy); + + Type elem = getElemTy(src0Ty); + if (!getTmp()) { + auto verifyA2A3NoTmp = [&]() -> LogicalResult { + if (!(elem.isInteger(32) || elem.isF32())) + return emitOpError("expects A2/A3 trem element type to be i32/f32"); + return success(); + }; + auto verifyA5NoTmp = [&]() -> LogicalResult { + if (!(elem.isInteger(32) || elem.isInteger(16) || elem.isF16() || + elem.isF32())) + return emitOpError( + "expects A5 trem element type to be i32/i16/f16/f32"); + return success(); + }; + return dispatchVerifierByArch(getOperation(), verifyA2A3NoTmp, + verifyA5NoTmp); + } + Type tmpTy = getTmp().getType(); + if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); auto tmpValid = getValidShapeVec(tmpTy); if (dstValid.size() != 2 || tmpValid.size() != 2) return emitOpError("expects tmp and dst to be rank-2 tiles"); - - Type elem = getElemTy(src0Ty); auto verifyA2A3 = [&]() -> LogicalResult { if (failed(verifyVecTileCommon(*this, tmpTy, "tmp"))) return failure(); @@ -10914,11 +11071,9 @@ mlir::LogicalResult mlir::pto::TFModOp::verify() { mlir::LogicalResult mlir::pto::TRemSOp::verify() { Type ts = getSrc().getType(); - Type tt = getTmp().getType(); Type td = getDst().getType(); Type scalarTy = getScalar().getType(); if (failed(verifyTileBufCommon(*this, ts, "src")) || - failed(verifyTileBufCommon(*this, tt, "tmp")) || failed(verifyTileBufCommon(*this, td, "dst"))) return failure(); if (failed(verifyTileBufSameElemType(*this, ts, td, "src", "dst")) || @@ -10930,6 +11085,25 @@ mlir::LogicalResult mlir::pto::TRemSOp::verify() { if (scalarTy != elem) return emitOpError("expects scalar type to match the tile element type"); auto dstValid = getValidShapeVec(td); + if (!getTmp()) { + auto verifyA2A3NoTmp = [&]() -> LogicalResult { + if (!(elem.isInteger(32) || elem.isF32())) + return emitOpError("expects A2/A3 trems element type to be i32/f32"); + return success(); + }; + auto verifyA5NoTmp = [&]() -> LogicalResult { + if (!(elem.isInteger(32) || elem.isInteger(16) || elem.isF16() || + elem.isF32())) + return emitOpError( + "expects A5 trems element type to be i32/i16/f16/f32"); + return success(); + }; + return dispatchVerifierByArch(getOperation(), verifyA2A3NoTmp, + verifyA5NoTmp); + } + Type tt = getTmp().getType(); + if (failed(verifyTileBufCommon(*this, tt, "tmp"))) + return failure(); auto tmpValid = getValidShapeVec(tt); if (dstValid.size() != 2 || tmpValid.size() != 2) return emitOpError("expects tmp and dst to be rank-2 tiles"); @@ -11016,7 +11190,6 @@ mlir::LogicalResult mlir::pto::TPowOp::verify() { Type elem = getElemTy(baseTy); bool isIntElem = elem.isInteger(32) || elem.isInteger(16) || elem.isInteger(8); - bool isFpElem = elem.isF16() || elem.isF32() || elem.isBF16(); auto verifyA2A3 = [&]() -> LogicalResult { if (getPrecisionType() == pto::PowPrecision::HighPrecision) return emitOpError( @@ -11042,10 +11215,6 @@ mlir::LogicalResult mlir::pto::TPowOp::verify() { if (failed(dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5))) return failure(); - if (isFpElem && !getTmp()) - return emitOpError( - "expects tmp when element type is floating-point (required by the " - "floating-point pow lowering)"); if (isIntElem && getTmp()) return emitOpError( "does not accept tmp when element type is integer (the integer pow " @@ -11079,7 +11248,6 @@ mlir::LogicalResult mlir::pto::TPowSOp::verify() { // Same dtype matrix as TPowOp; see comment in TPowOp::verify. bool isIntElem = elem.isInteger(32) || elem.isInteger(16) || elem.isInteger(8); - bool isFpElem = elem.isF16() || elem.isF32() || elem.isBF16(); auto verifyA2A3 = [&]() -> LogicalResult { if (getPrecisionType() == pto::PowPrecision::HighPrecision) return emitOpError( @@ -11105,10 +11273,6 @@ mlir::LogicalResult mlir::pto::TPowSOp::verify() { if (failed(dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5))) return failure(); - if (isFpElem && !getTmp()) - return emitOpError( - "expects tmp when element type is floating-point (required by the " - "floating-point pow lowering)"); if (isIntElem && getTmp()) return emitOpError( "does not accept tmp when element type is integer (the integer pows " @@ -12235,8 +12399,267 @@ mlir::LogicalResult mlir::pto::TRowExpandMinOp::verify() { } +static ParseResult parseOptionalTmpRowReductionOp(OpAsmParser &parser, + OperationState &result) { + OpAsmParser::UnresolvedOperand src, tmp, dst; + Type srcTy, tmpTy, dstTy; + bool hasTmp = false; + + if (parser.parseKeyword("ins") || parser.parseLParen() || + parser.parseOperand(src)) + return failure(); + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(tmp)) + return failure(); + hasTmp = true; + } + if (parser.parseColonType(srcTy)) + return failure(); + if (hasTmp && (parser.parseComma() || parser.parseType(tmpTy))) + return failure(); + if (parser.parseRParen() || parser.parseKeyword("outs") || + parser.parseLParen() || parser.parseOperand(dst) || + parser.parseColonType(dstTy) || parser.parseRParen() || + parser.parseOptionalAttrDict(result.attributes)) + return failure(); + + if (parser.resolveOperand(src, srcTy, result.operands)) + return failure(); + if (hasTmp && parser.resolveOperand(tmp, tmpTy, result.operands)) + return failure(); + if (parser.resolveOperand(dst, dstTy, result.operands)) + return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr({1, hasTmp ? 1 : 0, 1})); + return success(); +} + +static void printOptionalTmpRowReductionOp(OpAsmPrinter &p, Operation *op, + Value src, Value tmp, Value dst) { + p << " ins(" << src; + if (tmp) + p << ", " << tmp; + p << " : " << src.getType(); + if (tmp) + p << ", " << tmp.getType(); + p << ") outs(" << dst << " : " << dst.getType() << ")"; + p.printOptionalAttrDict(op->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); +} + +static ParseResult parseOptionalTmpFixedDpsOp( + OpAsmParser &parser, OperationState &result, unsigned minInputs, + unsigned maxInputs, ArrayRef noTmpSegments, + ArrayRef withTmpSegments) { + SmallVector inputs; + SmallVector inputTypes; + OpAsmParser::UnresolvedOperand dst; + Type dstType; + if (parser.parseKeyword("ins") || parser.parseLParen()) + return failure(); + do { + inputs.emplace_back(); + if (parser.parseOperand(inputs.back())) + return failure(); + } while (succeeded(parser.parseOptionalComma())); + if (inputs.size() < minInputs || inputs.size() > maxInputs || + parser.parseColon()) + return failure(); + for (unsigned i = 0; i < inputs.size(); ++i) { + if (i && parser.parseComma()) + return failure(); + Type type; + if (parser.parseType(type)) + return failure(); + inputTypes.push_back(type); + } + if (parser.parseRParen() || parser.parseKeyword("outs") || + parser.parseLParen() || parser.parseOperand(dst) || + parser.parseColonType(dstType) || parser.parseRParen() || + parser.parseOptionalAttrDict(result.attributes) || + parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), + result.operands) || + parser.resolveOperand(dst, dstType, result.operands)) + return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr( + inputs.size() == minInputs ? noTmpSegments : withTmpSegments)); + return success(); +} + +static void printOptionalTmpFixedDpsOp(OpAsmPrinter &p, Operation *op, + ArrayRef inputs, Value dst) { + p << " ins("; + llvm::interleaveComma(inputs, p, [&](Value value) { p << value; }); + p << " : "; + llvm::interleaveComma(inputs, p, + [&](Value value) { p << value.getType(); }); + p << ") outs(" << dst << " : " << dst.getType() << ")"; + p.printOptionalAttrDict(op->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); +} + +ParseResult mlir::pto::TTransOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 1, 2, {1, 0, 1}, + {1, 1, 1}); +} +void mlir::pto::TTransOp::print(OpAsmPrinter &p) { + SmallVector inputs{getSrc()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TPReluOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 2, 3, {1, 1, 0, 1}, + {1, 1, 1, 1}); +} +void mlir::pto::TPReluOp::print(OpAsmPrinter &p) { + SmallVector inputs{getSrc0(), getSrc1()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TRemOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 2, 3, {1, 1, 0, 1}, + {1, 1, 1, 1}); +} +void mlir::pto::TRemOp::print(OpAsmPrinter &p) { + SmallVector inputs{getSrc0(), getSrc1()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TRemSOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 2, 3, {1, 1, 0, 1}, + {1, 1, 1, 1}); +} +void mlir::pto::TRemSOp::print(OpAsmPrinter &p) { + SmallVector inputs{getSrc(), getScalar()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TSelOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 3, 4, + {1, 1, 1, 0, 1}, {1, 1, 1, 1, 1}); +} +void mlir::pto::TSelOp::print(OpAsmPrinter &p) { + SmallVector inputs{getMask(), getSrc0(), getSrc1()}; + if (getTmp()) + inputs.push_back(getTmp()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TSelSOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpFixedDpsOp(parser, result, 3, 4, + {1, 1, 0, 1, 1}, {1, 1, 1, 1, 1}); +} +void mlir::pto::TSelSOp::print(OpAsmPrinter &p) { + SmallVector inputs{getMask(), getSrc()}; + if (getTmp()) + inputs.push_back(getTmp()); + inputs.push_back(getScalar()); + printOptionalTmpFixedDpsOp(p, getOperation(), inputs, getDst()); +} + +ParseResult mlir::pto::TColArgMaxOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TColArgMaxOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TColArgMinOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TColArgMinOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowMaxOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowMaxOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowArgMaxOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowArgMaxOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowMinOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowMinOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowArgMinOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowArgMinOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowSumOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowSumOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + +ParseResult mlir::pto::TRowProdOp::parse(OpAsmParser &parser, + OperationState &result) { + return parseOptionalTmpRowReductionOp(parser, result); +} + +void mlir::pto::TRowProdOp::print(OpAsmPrinter &p) { + printOptionalTmpRowReductionOp(p, getOperation(), getSrc(), getTmp(), + getDst()); +} + mlir::LogicalResult mlir::pto::TRowMaxOp::verify() { auto verifyByArch = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects element type to be i16/i32/f16/f32"); @@ -12245,6 +12668,9 @@ mlir::LogicalResult mlir::pto::TRowMaxOp::verify() { } mlir::LogicalResult mlir::pto::TRowArgMaxOp::verify() { + if (!getTmp()) + return verifyTRowArgReductionNoTmp(getOperation(), getSrc().getType(), + getDst().getType()); auto verifyA2A3 = [&]() -> LogicalResult { return verifyTRowArgReductionOpA2A3(*this, getSrc().getType(), getTmp().getType(), getDst().getType()); @@ -12260,6 +12686,10 @@ mlir::LogicalResult mlir::pto::TRowArgMaxOp::verify() { mlir::LogicalResult mlir::pto::TRowMinOp::verify() { auto verifyByArch = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects element type to be i16/i32/f16/f32"); @@ -12268,6 +12698,9 @@ mlir::LogicalResult mlir::pto::TRowMinOp::verify() { } mlir::LogicalResult mlir::pto::TRowArgMinOp::verify() { + if (!getTmp()) + return verifyTRowArgReductionNoTmp(getOperation(), getSrc().getType(), + getDst().getType()); auto verifyA2A3 = [&]() -> LogicalResult { return verifyTRowArgReductionOpA2A3(*this, getSrc().getType(), getTmp().getType(), getDst().getType()); @@ -12283,6 +12716,10 @@ mlir::LogicalResult mlir::pto::TRowArgMinOp::verify() { mlir::LogicalResult mlir::pto::TRowSumOp::verify() { auto verifyByArch = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects element type to be i16/i32/f16/f32"); @@ -12292,11 +12729,19 @@ mlir::LogicalResult mlir::pto::TRowSumOp::verify() { mlir::LogicalResult mlir::pto::TRowProdOp::verify() { auto verifyA2A3 = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects A2/A3 trowprod element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects A2/A3 trowprod element type to be i16/i32/f16/f32"); }; auto verifyA5 = [&]() -> LogicalResult { + if (!getTmp()) + return verifyTRowReductionNoTmpCommon( + *this, getSrc().getType(), getDst().getType(), + "expects A5 trowprod element type to be i16/i32/f16/f32"); return verifyTRowReductionWithTmpCommon( *this, getSrc().getType(), getTmp().getType(), getDst().getType(), "expects A5 trowprod element type to be i16/i32/f16/f32"); @@ -12555,16 +13000,17 @@ mlir::LogicalResult mlir::pto::TSelSOp::verify() { auto verifyCommon = [&]() -> FailureOr { Type tMask = getMask().getType(); Type tSrc = getSrc().getType(); - Type tTmp = getTmp().getType(); + Type tTmp = getTmp() ? getTmp().getType() : Type{}; Type tDst = getDst().getType(); if (failed(verifyTileBufCommon(*this, tMask, "mask")) || failed(verifyTileBufCommon(*this, tSrc, "src")) || - failed(verifyTileBufCommon(*this, tTmp, "tmp")) || failed(verifyTileBufCommon(*this, tDst, "dst"))) return failure(); + if (tTmp && failed(verifyTileBufCommon(*this, tTmp, "tmp"))) + return failure(); Type eMask = getElemTy(tMask), eSrc = getElemTy(tSrc); - Type eTmp = getElemTy(tTmp), eDst = getElemTy(tDst); - if (!eMask || !eSrc || !eTmp || !eDst) { + Type eDst = getElemTy(tDst); + if (!eMask || !eSrc || !eDst) { emitOpError("failed to get element type for operands"); return failure(); } @@ -12839,17 +13285,34 @@ mlir::LogicalResult mlir::pto::TSubSCOp::verify() { return emitOpError() << "expects src0, src1, and dst to have the same rank"; return mlir::success(); } +static bool ttransUsesTmp(Type srcTy, Type dstTy) { + auto srcShape = getShapeVec(srcTy); + auto dstShape = getShapeVec(dstTy); + unsigned elemBytes = getPTOStorageElemByteSize(getElemTy(srcTy)); + if (srcShape.size() != 2 || dstShape.size() != 2 || elemBytes == 0 || + llvm::is_contained(srcShape, ShapedType::kDynamic) || + llvm::is_contained(dstShape, ShapedType::kDynamic)) + return true; + int64_t rowStride = elemBytes == 1 ? 32 : 16; + int64_t elemPerBlock = 32 / elemBytes; + int64_t srcStride = srcShape[1]; + int64_t dstStride = dstShape[1]; + return dstStride % rowStride == 0 && srcStride % elemPerBlock == 0 && + srcStride / elemPerBlock <= 255; +} + mlir::LogicalResult mlir::pto::TTransOp::verify() { auto verifyA2A3 = [&]() -> LogicalResult { Type srcTy = getSrc().getType(); - Type tmpTy = getTmp().getType(); + Type tmpTy = getTmp() ? getTmp().getType() : Type{}; Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, srcTy, "src")) || - failed(verifyTileBufCommon(*this, tmpTy, "tmp")) || failed(verifyTileBufCommon(*this, dstTy, "dst"))) return failure(); + if (tmpTy && failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); Type srcElem = getElemTy(srcTy); - Type tmpElem = getElemTy(tmpTy); + Type tmpElem = tmpTy ? getElemTy(tmpTy) : srcElem; Type dstElem = getElemTy(dstTy); if (!srcElem || !tmpElem || !dstElem || srcElem != dstElem || srcElem != tmpElem) return emitOpError() << "expects src and dst to have the same element type"; @@ -12875,14 +13338,15 @@ mlir::LogicalResult mlir::pto::TTransOp::verify() { }; auto verifyA5 = [&]() -> LogicalResult { Type srcTy = getSrc().getType(); - Type tmpTy = getTmp().getType(); + Type tmpTy = getTmp() ? getTmp().getType() : Type{}; Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, srcTy, "src")) || - failed(verifyTileBufCommon(*this, tmpTy, "tmp")) || failed(verifyTileBufCommon(*this, dstTy, "dst"))) return failure(); + if (tmpTy && failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); Type srcElem = getElemTy(srcTy); - Type tmpElem = getElemTy(tmpTy); + Type tmpElem = tmpTy ? getElemTy(tmpTy) : srcElem; Type dstElem = getElemTy(dstTy); if (!srcElem || !tmpElem || !dstElem || srcElem != dstElem || srcElem != tmpElem) return emitOpError() << "expects src, tmp, and dst to have the same element type"; @@ -12920,6 +13384,100 @@ mlir::LogicalResult mlir::pto::TTransOp::verify() { return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } +ParseResult mlir::pto::TXorOp::parse(OpAsmParser &parser, + OperationState &result) { + OpAsmParser::UnresolvedOperand src0, src1, tmp, dst; + Type src0Ty, src1Ty, tmpTy, dstTy; + bool hasTmp = false; + if (parser.parseKeyword("ins") || parser.parseLParen() || + parser.parseOperand(src0) || parser.parseComma() || + parser.parseOperand(src1)) + return failure(); + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(tmp)) + return failure(); + hasTmp = true; + } + if (parser.parseColonType(src0Ty) || parser.parseComma() || + parser.parseType(src1Ty)) + return failure(); + if (hasTmp && (parser.parseComma() || parser.parseType(tmpTy))) + return failure(); + if (parser.parseRParen() || parser.parseKeyword("outs") || + parser.parseLParen() || parser.parseOperand(dst) || + parser.parseColonType(dstTy) || parser.parseRParen() || + parser.parseOptionalAttrDict(result.attributes)) + return failure(); + if (parser.resolveOperand(src0, src0Ty, result.operands) || + parser.resolveOperand(src1, src1Ty, result.operands) || + (hasTmp && parser.resolveOperand(tmp, tmpTy, result.operands)) || + parser.resolveOperand(dst, dstTy, result.operands)) + return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr({1, 1, hasTmp ? 1 : 0, 1})); + return success(); +} + +void mlir::pto::TXorOp::print(OpAsmPrinter &p) { + p << " ins(" << getSrc0() << ", " << getSrc1(); + if (getTmp()) + p << ", " << getTmp(); + p << " : " << getSrc0().getType() << ", " << getSrc1().getType(); + if (getTmp()) + p << ", " << getTmp().getType(); + p << ") outs(" << getDst() << " : " << getDst().getType() << ")"; + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); +} + +ParseResult mlir::pto::TXorSOp::parse(OpAsmParser &parser, + OperationState &result) { + OpAsmParser::UnresolvedOperand src, scalar, tmp, dst; + Type srcTy, scalarTy, tmpTy, dstTy; + bool hasTmp = false; + if (parser.parseKeyword("ins") || parser.parseLParen() || + parser.parseOperand(src) || parser.parseComma() || + parser.parseOperand(scalar)) + return failure(); + if (succeeded(parser.parseOptionalComma())) { + if (parser.parseOperand(tmp)) + return failure(); + hasTmp = true; + } + if (parser.parseColonType(srcTy) || parser.parseComma() || + parser.parseType(scalarTy)) + return failure(); + if (hasTmp && (parser.parseComma() || parser.parseType(tmpTy))) + return failure(); + if (parser.parseRParen() || parser.parseKeyword("outs") || + parser.parseLParen() || parser.parseOperand(dst) || + parser.parseColonType(dstTy) || parser.parseRParen() || + parser.parseOptionalAttrDict(result.attributes)) + return failure(); + if (parser.resolveOperand(src, srcTy, result.operands) || + parser.resolveOperand(scalar, scalarTy, result.operands) || + (hasTmp && parser.resolveOperand(tmp, tmpTy, result.operands)) || + parser.resolveOperand(dst, dstTy, result.operands)) + return failure(); + result.addAttribute( + "operandSegmentSizes", + parser.getBuilder().getDenseI32ArrayAttr({1, 1, hasTmp ? 1 : 0, 1})); + return success(); +} + +void mlir::pto::TXorSOp::print(OpAsmPrinter &p) { + p << " ins(" << getSrc() << ", " << getScalar(); + if (getTmp()) + p << ", " << getTmp(); + p << " : " << getSrc().getType() << ", " << getScalar().getType(); + if (getTmp()) + p << ", " << getTmp().getType(); + p << ") outs(" << getDst() << " : " << getDst().getType() << ")"; + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); +} + mlir::LogicalResult mlir::pto::TXorOp::verify() { auto verifyBase = [&]() -> FailureOr { return verifyMatchingRowMajorBinaryTileOpCommon( @@ -12931,16 +13489,20 @@ mlir::LogicalResult mlir::pto::TXorOp::verify() { FailureOr elemOr = verifyBase(); if (failed(elemOr)) return failure(); - Type tmpTy = getTmp().getType(); - if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) - return failure(); Type elem = *elemOr; - if (getElemTy(tmpTy) != elem) - return emitOpError("expects tmp to have the same element type as src0, src1, and dst"); - if (!isRowMajorTileBuf(tmpTy)) - return emitOpError("expects tmp to use row-major layout"); - if (failed(verifyTileBufSameValidShape(*this, tmpTy, getDst().getType(), "tmp", "dst"))) - return failure(); + if (getTmp()) { + Type tmpTy = getTmp().getType(); + if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); + if (getElemTy(tmpTy) != elem) + return emitOpError( + "expects tmp to have the same element type as src0, src1, and dst"); + if (!isRowMajorTileBuf(tmpTy)) + return emitOpError("expects tmp to use row-major layout"); + if (failed(verifyTileBufSameValidShape( + *this, tmpTy, getDst().getType(), "tmp", "dst"))) + return failure(); + } auto it = mlir::dyn_cast(elem); if (!it || (it.getWidth() != 8 && it.getWidth() != 16 && it.getWidth() != 32)) @@ -12975,14 +13537,17 @@ mlir::LogicalResult mlir::pto::TXorSOp::verify() { FailureOr elemOr = verifyCommon(); if (failed(elemOr)) return failure(); - Type tmpTy = getTmp().getType(); - if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) - return failure(); Type elem = *elemOr; - if (getElemTy(tmpTy) != elem) - return emitOpError("expects tmp to have the same element type as src and dst"); - if (!isRowMajorTileBuf(tmpTy)) - return emitOpError("expects tmp to use row-major layout"); + if (getTmp()) { + Type tmpTy = getTmp().getType(); + if (failed(verifyTileBufCommon(*this, tmpTy, "tmp"))) + return failure(); + if (getElemTy(tmpTy) != elem) + return emitOpError( + "expects tmp to have the same element type as src and dst"); + if (!isRowMajorTileBuf(tmpTy)) + return emitOpError("expects tmp to use row-major layout"); + } auto it = mlir::dyn_cast(elem); if (!it || (it.getWidth() != 8 && it.getWidth() != 16)) return emitOpError( @@ -14388,16 +14953,22 @@ PTO_DEFINE_UNARY_EFFECTS(TColProdOp, getSrcMutable(), getDstMutable()) void TColArgMaxOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TColArgMinOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14406,6 +14977,7 @@ void TColSumOp::getEffects( PTO_ADD_READ(getSrcMutable()); auto tmp = getTmpMutable(); if (!tmp.empty()) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); } PTO_ADD_WRITE(getDstMutable()); @@ -14414,6 +14986,11 @@ void TColSumOp::getEffects( void TCvtOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRandomOp::getEffects( @@ -14482,8 +15059,11 @@ void TGatherOp::getEffects( PTO_ADD_WRITE(cdst[0]); if (auto indices = getIndicesMutable(); !indices.empty()) PTO_ADD_READ(indices[0]); - if (auto tmp = getTmpMutable(); !tmp.empty()) + if (auto tmp = getTmpMutable(); + !tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14504,8 +15084,10 @@ void TMrgSortOp::getEffects( PTO_ADD_READ(opnd); } auto tmp = getTmpMutable(); - if (!tmp.empty()) + if (!tmp.empty()) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } for (auto &opnd : getDstsMutable()) { PTO_ADD_WRITE(opnd); } @@ -14552,8 +15134,11 @@ void TPReluOp::getEffects( // A5 pto-isa TPRELU implementation does not consume tmp; modeling tmp as a // write-only scratch on A5 incorrectly inflates local-memory planning and // can trigger false vec-overflow diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14565,8 +15150,10 @@ void TQuantOp::getEffects( if (!offsetRange.empty()) PTO_ADD_READ(offsetRange[0]); auto tmpRange = getTmpMutable(); - if (!tmpRange.empty() && getTargetArch(getOperation()) != PTOArch::A5) + if (!tmpRange.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmpRange[0]); PTO_ADD_WRITE(tmpRange[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14591,16 +15178,22 @@ void TRemOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrc0Mutable()); PTO_ADD_READ(getSrc1Mutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRemSOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14609,8 +15202,10 @@ void TPowOp::getEffects( PTO_ADD_READ(getBaseMutable()); PTO_ADD_READ(getExpMutable()); auto tmp = getTmpMutable(); - if (!tmp.empty()) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14618,8 +15213,10 @@ void TPowSOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); auto tmp = getTmpMutable(); - if (!tmp.empty()) + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } PTO_DEFINE_UNARY_EFFECTS(TRowExpandOp, getSrcMutable(), getDstMutable()) @@ -14710,8 +15307,11 @@ void TRowExpandMinOp::getEffects( void TRowMaxOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14721,16 +15321,22 @@ void TRowArgMaxOp::getEffects( // A5 lowering does not consume tmp for TROWARGMAX; modeling tmp as a // scratch write inflates local-memory planning and can trigger false // vec-overflow diagnostics, mirroring the fixed A5 TPRELU issue. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRowMinOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14740,32 +15346,43 @@ void TRowArgMinOp::getEffects( // A5 lowering does not consume tmp for TROWARGMIN; modeling tmp as a // scratch write inflates local-memory planning and can trigger false // vec-overflow diagnostics, mirroring the fixed A5 TPRELU issue. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRowSumOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRowProdOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } void TRsqrtOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); auto tmp = getTmpMutable(); - if (!tmp.empty()) + if (!tmp.empty()) { + PTO_ADD_READ(tmp[0]); PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14789,8 +15406,11 @@ void TSelOp::getEffects( // A5 lowering does not consume tmp for TSEL; modeling tmp as a scratch // write inflates local-memory planning and can trigger false vec-overflow // diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14802,8 +15422,11 @@ void TSelSOp::getEffects( // A5 lowering does not consume tmp for TSELS; modeling tmp as a scratch // write inflates local-memory planning and can trigger false vec-overflow // diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14836,8 +15459,11 @@ void TXorSOp::getEffects( // A5 lowering does not consume tmp for TXORS; modeling tmp as a scratch // write inflates local-memory planning and can trigger false vec-overflow // diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14849,8 +15475,11 @@ void TXorOp::getEffects( // A5 lowering does not consume tmp for TXOR; modeling tmp as a scratch // write inflates local-memory planning and can trigger false vec-overflow // diagnostics. - if (getTargetArch(getOperation()) != PTOArch::A5) - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && getTargetArch(getOperation()) != PTOArch::A5) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } @@ -14858,7 +15487,11 @@ void TXorOp::getEffects( void TTransOp::getEffects( SmallVectorImpl> &effects) { PTO_ADD_READ(getSrcMutable()); - PTO_ADD_WRITE(getTmpMutable()); + auto tmp = getTmpMutable(); + if (!tmp.empty() && ttransUsesTmp(getSrc().getType(), getDst().getType())) { + PTO_ADD_READ(tmp[0]); + PTO_ADD_WRITE(tmp[0]); + } PTO_ADD_WRITE(getDstMutable()); } diff --git a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp index 1709f10e4b..6ee5b20d26 100644 --- a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp +++ b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp @@ -19,6 +19,7 @@ #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Matchers.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/STLExtras.h" #include "llvm/Support/Casting.h" #include "llvm/Support/ErrorHandling.h" @@ -40,6 +41,44 @@ namespace { static constexpr uint64_t kVectorRegisterSizeInBytes = 256U; static constexpr unsigned kPipeVPruneMinRepeat = 16U; +static bool hasReadWriteScratchDependency( + Operation *op, const DepBaseMemInfoPairVec &dependencies) { + auto effectsOp = dyn_cast_or_null(op); + if (!effectsOp) + return false; + + llvm::DenseSet reads; + llvm::DenseSet writes; + SmallVector, 8> effects; + effectsOp.getEffects(effects); + for (const auto &effect : effects) { + Value value = effect.getValue(); + if (!value) + continue; + if (isa(effect.getEffect())) + reads.insert(value); + if (isa(effect.getEffect())) + writes.insert(value); + } + + ValueRange dpsInits; + if (auto ptoDpsOp = dyn_cast(op)) + dpsInits = ptoDpsOp.getDpsInits(); + else if (auto dpsOp = dyn_cast(op)) + dpsInits = dpsOp.getDpsInits(); + return llvm::any_of(writes, [&](Value value) { + if (!reads.contains(value) || llvm::is_contained(dpsInits, value)) + return false; + return llvm::any_of(dependencies, [&](const auto &dependency) { + auto matches = [&](const BaseMemInfo *info) { + return info && + (info->baseBuffer == value || info->rootBuffer == value); + }; + return matches(dependency.first) || matches(dependency.second); + }); + }); +} + struct RepeatAccessShape { SmallVector fullShape; SmallVector validShape; @@ -511,6 +550,16 @@ bool InsertSyncAnalysis::CanPrunePipeVBarrier( return false; } + // The same-access fast path only applies to a producer output consumed by + // the next op. A read/write non-DPS operand is scratch state; pruning its + // WAW dependency would allow two vector instructions to use it concurrently. + if (hasReadWriteScratchDependency(nowCompound->elementOp, + depBaseMemInfosVec) || + hasReadWriteScratchDependency(frontCompound->elementOp, + depBaseMemInfosVec)) { + return false; + } + // PIPE_V has a hardware-safe same-access chain case: exact same-access // dependencies from the producer result to the consumer source do not require // a vector-pipe barrier once the producer repeat is large enough. Keep the diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index abb14a89e2..8c0d60eef7 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -67,6 +67,64 @@ static SmallVector getValidShapeVec(Type ty) { return {}; } +static SmallVector getShapeVec(Type ty) { + if (auto tileTy = dyn_cast(ty)) + return SmallVector(tileTy.getShape().begin(), + tileTy.getShape().end()); + return {}; +} + +static int64_t ceilDiv(int64_t lhs, int64_t rhs) { + return (lhs + rhs - 1) / rhs; +} + +static bool hasDynamicDim(ArrayRef dims) { + return llvm::any_of(dims, [](int64_t dim) { + return dim == ShapedType::kDynamic; + }); +} + +static pto::TileBufType makeVecTmpType(MLIRContext *ctx, + ArrayRef shape, + Type elementType, + ArrayRef validShape) { + return pto::TileBufType::get( + ctx, shape, elementType, + pto::AddressSpaceAttr::get(ctx, pto::AddressSpace::VEC), validShape, + makeRowMajorNoneBoxConfig(ctx)); +} + +static FailureOr makeSameShapeTmpType(MLIRContext *ctx, + Value like, + Type elementType = {}) { + auto likeTy = dyn_cast(like.getType()); + if (!likeTy) + return failure(); + if (!elementType) + elementType = likeTy.getElementType(); + auto shape = getShapeVec(like.getType()); + auto validShape = getValidShapeVec(like.getType()); + if (shape.empty() || validShape.empty() || hasDynamicDim(shape) || + hasDynamicDim(validShape)) + return failure(); + return makeVecTmpType(ctx, shape, elementType, validShape); +} + +static FailureOr createAllocTmp(OpBuilder &builder, Location loc, + pto::TileBufType tmpType) { + return builder + .create(loc, tmpType, Value(), Value(), Value()) + .getResult(); +} + +static void copyAttrsExceptOperandSegments(Operation *from, OperationState &to) { + for (NamedAttribute attr : from->getAttrs()) { + if (attr.getName() == "operandSegmentSizes") + continue; + to.addAttribute(attr.getName(), attr.getValue()); + } +} + static bool validShapesCompatible(ArrayRef lhs, ArrayRef rhs) { if (lhs.size() != rhs.size()) @@ -197,6 +255,553 @@ static LogicalResult materializeTRowExpandTmp(OpTy op, bool requireExplicitTmp, return success(); } +static LogicalResult replaceTColSumWithTmp(pto::TColSumOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || !op.getIsBinary()) + return success(); + if (requireExplicitTmp) + return op.emitOpError( + "requires explicit tmp for binary tcolsum when PlanMemory is skipped"); + + auto srcTy = dyn_cast(op.getSrc().getType()); + if (!srcTy) + return op.emitOpError("expects tile_buf src when materializing implicit tmp"); + auto valid = getValidShapeVec(op.getSrc().getType()); + if (valid.size() != 2 || hasDynamicDim(valid)) + return op.emitOpError( + "requires static src valid_shape to materialize binary tcolsum tmp"); + + SmallVector tmpShape{ceilDiv(valid[0], 2), valid[1]}; + auto tmpType = makeVecTmpType(ctx, tmpShape, srcTy.getElementType(), tmpShape); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), *tmp, op.getDst()}); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTQuantWithTmp(pto::TQuantOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType( + ctx, op.getSrc(), Float32Type::get(ctx)); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf src to materialize implicit tquant tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + SmallVector operands{op.getSrc(), op.getFp()}; + if (op.getOffset()) + operands.push_back(op.getOffset()); + operands.push_back(*tmp); + operands.push_back(op.getDst()); + state.addOperands(operands); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr( + {1, 1, op.getOffset() ? 1 : 0, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static bool isFloatingPointTile(Value value) { + auto tileTy = dyn_cast(value.getType()); + return tileTy && isa(tileTy.getElementType()); +} + +static LogicalResult replaceTPowWithTmp(pto::TPowOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || !isFloatingPointTile(op.getDst())) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf dst to materialize implicit tpow tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getBase(), op.getExp(), op.getDst(), *tmp}); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTPowSWithTmp(pto::TPowSOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || !isFloatingPointTile(op.getDst())) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf dst to materialize implicit tpows tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), op.getScalar(), op.getDst(), *tmp}); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTGatherWithTmp(pto::TGatherOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || op.hasMaskForm()) + return success(); + if (!op.hasIndexForm() && !op.hasCompareForm()) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = failure(); + if (op.hasIndexForm()) { + tmpType = makeSameShapeTmpType(ctx, op.getIndices()); + } else { + auto srcTy = dyn_cast(op.getSrc().getType()); + auto dstTy = dyn_cast(op.getDst().getType()); + if (!srcTy || !dstTy) + return op.emitOpError( + "expects tile_buf operands when materializing compare-form tgather tmp"); + auto srcShape = getShapeVec(op.getSrc().getType()); + if (srcShape.size() != 2 || hasDynamicDim(srcShape)) + return op.emitOpError( + "requires static src shape to materialize compare-form tgather tmp"); + int64_t bytes = srcShape[0] * srcShape[1] * 4 + srcShape[0] * 4; + tmpType = makeVecTmpType(ctx, {1, bytes}, IntegerType::get(ctx, 8), + {1, bytes}); + } + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf indices/src to materialize implicit tgather tmp"); + + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + SmallVector operands{op.getSrc(), op.getDst()}; + if (op.getCdst()) + operands.push_back(op.getCdst()); + if (op.getIndices()) + operands.push_back(op.getIndices()); + operands.push_back(*tmp); + if (op.getKValue()) + operands.push_back(op.getKValue()); + state.addOperands(operands); + state.addAttribute( + "operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, op.getCdst() ? 1 : 0, + op.getIndices() ? 1 : 0, 1, + op.getKValue() ? 1 : 0})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTSort32WithTmp(pto::TSort32Op op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + auto valid = getValidShapeVec(op.getSrc().getType()); + if (valid.size() != 2 || valid[1] == ShapedType::kDynamic || + valid[1] % 32 == 0) + return success(); + if (requireExplicitTmp) + return op.emitOpError( + "requires explicit tmp for non-32-aligned tsort32 when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getSrc()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf src to materialize implicit tsort32 tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), op.getIdx(), *tmp, op.getDst()}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +template +static LogicalResult replaceRowReductionWithTmp(OpTy op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getSrc()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf src to materialize implicit row-reduction tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), *tmp, op.getDst()}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTXorWithTmp(pto::TXorOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf dst to materialize implicit txor tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc0(), op.getSrc1(), *tmp, op.getDst()}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceTXorSWithTmp(pto::TXorSOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp()) + return success(); + if (requireExplicitTmp) + return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); + FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + if (failed(tmpType)) + return op.emitOpError( + "requires static tile_buf dst to materialize implicit txors tmp"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); + if (failed(tmp)) + return failure(); + OperationState state(op.getLoc(), op->getName()); + state.addOperands({op.getSrc(), op.getScalar(), *tmp, op.getDst()}); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr({1, 1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + +static LogicalResult replaceFixedDpsOpWithTmp( + Operation *op, ArrayRef operands, pto::TileBufType tmpType, + ArrayRef operandSegments, bool requireExplicitTmp, + StringRef opName) { + if (requireExplicitTmp) + return op->emitOpError( + "requires explicit tmp when PlanMemory is skipped"); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op->getLoc(), tmpType); + if (failed(tmp)) + return failure(); + SmallVector finalOperands; + finalOperands.reserve(operands.size() + 1); + for (Value operand : operands) { + if (operand) + finalOperands.push_back(operand); + else + finalOperands.push_back(*tmp); + } + OperationState state(op->getLoc(), op->getName()); + state.addOperands(finalOperands); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr(operandSegments)); + copyAttrsExceptOperandSegments(op, state); + builder.create(state); + op->erase(); + (void)opName; + return success(); +} + +static FailureOr makeTPReluTmpType(MLIRContext *ctx, + Value dst) { + auto dstTy = dyn_cast(dst.getType()); + auto shape = getShapeVec(dst.getType()); + auto valid = getValidShapeVec(dst.getType()); + if (!dstTy || shape.size() != 2 || valid.size() != 2 || + hasDynamicDim(shape) || hasDynamicDim(valid)) + return failure(); + int64_t validCols = ceilDiv(valid[1], 8); + int64_t cols = std::max(32, ceilDiv(validCols, 32) * 32); + return makeVecTmpType(ctx, {valid[0] + 1, cols}, IntegerType::get(ctx, 8), + {valid[0], validCols}); +} + +static FailureOr makeRowsTmpType(MLIRContext *ctx, + Value dst, int64_t rows) { + auto dstTy = dyn_cast(dst.getType()); + auto shape = getShapeVec(dst.getType()); + auto valid = getValidShapeVec(dst.getType()); + if (!dstTy || shape.size() != 2 || valid.size() != 2 || + hasDynamicDim(shape) || hasDynamicDim(valid)) + return failure(); + return makeVecTmpType(ctx, {rows, shape[1]}, dstTy.getElementType(), + {rows, valid[1]}); +} + +static LogicalResult materializeFixedMandatoryTmp(Operation *op, + bool requireExplicitTmp, + MLIRContext *ctx) { + return llvm::TypeSwitch(op) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeTPReluTmpType(ctx, typedOp.getDst()); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf dst to materialize implicit tprelu tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getSrc0(), typedOp.getSrc1(), Value(), + typedOp.getDst()}, + *type, {1, 1, 1, 1}, requireExplicitTmp, "tprelu"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeRowsTmpType(ctx, typedOp.getDst(), 2); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf dst to materialize implicit trem tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getSrc0(), typedOp.getSrc1(), Value(), + typedOp.getDst()}, + *type, {1, 1, 1, 1}, requireExplicitTmp, "trem"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeRowsTmpType(ctx, typedOp.getDst(), 1); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf dst to materialize implicit trems tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getSrc(), typedOp.getScalar(), Value(), + typedOp.getDst()}, + *type, {1, 1, 1, 1}, requireExplicitTmp, "trems"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeVecTmpType(ctx, {1, 16}, IntegerType::get(ctx, 32), + {1, 16}); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getMask(), typedOp.getSrc0(), typedOp.getSrc1(), + Value(), typedOp.getDst()}, + type, {1, 1, 1, 1, 1}, requireExplicitTmp, "tsel"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto type = makeRowsTmpType(ctx, typedOp.getSrc(), 1); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf src to materialize implicit tsels tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getMask(), typedOp.getSrc(), Value(), + typedOp.getScalar(), typedOp.getDst()}, + *type, {1, 1, 1, 1, 1}, requireExplicitTmp, "tsels"); + }) + .Case([&](auto typedOp) -> LogicalResult { + if (typedOp.getTmp()) + return success(); + auto srcTy = dyn_cast(typedOp.getSrc().getType()); + auto dstTy = dyn_cast(typedOp.getDst().getType()); + auto srcShape = getShapeVec(typedOp.getSrc().getType()); + auto dstShape = getShapeVec(typedOp.getDst().getType()); + if (!srcTy || !dstTy || srcShape.size() != 2 || dstShape.size() != 2 || + hasDynamicDim(srcShape) || hasDynamicDim(dstShape)) + return typedOp.emitOpError( + "requires static tile_buf src to materialize implicit ttrans tmp"); + auto elemBytes = getElemBytes(srcTy.getElementType()); + if (!elemBytes) + return typedOp.emitOpError("failed to infer ttrans element size"); + int64_t rowStride = *elemBytes == 1 ? 32 : 16; + int64_t elemPerBlock = 32 / *elemBytes; + bool usesTmp = dstShape[1] % rowStride == 0 && + srcShape[1] % elemPerBlock == 0 && + srcShape[1] / elemPerBlock <= 255; + FailureOr type = makeSameShapeTmpType( + ctx, typedOp.getSrc()); + if (!usesTmp) + type = makeVecTmpType(ctx, {1, elemPerBlock}, + srcTy.getElementType(), {1, elemPerBlock}); + if (failed(type)) + return typedOp.emitOpError("failed to build implicit ttrans tmp"); + return replaceFixedDpsOpWithTmp( + op, {typedOp.getSrc(), Value(), typedOp.getDst()}, *type, + {1, 1, 1}, requireExplicitTmp, "ttrans"); + }) + .Default([](Operation *) { return success(); }); +} + +static bool tcvtNeedsTmp(pto::TCvtOp op) { + if (pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5 || + op.getSatMode() != pto::SaturationMode::OFF) + return false; + auto srcTy = dyn_cast(op.getSrc().getType()); + auto dstTy = dyn_cast(op.getDst().getType()); + if (!srcTy || !dstTy) + return false; + Type srcElem = srcTy.getElementType(); + Type dstElem = dstTy.getElementType(); + return (srcElem.isF32() && dstElem.isInteger(16)) || + (srcElem.isF16() && + (dstElem.isInteger(16) || dstElem.isInteger(8))); +} + +static FailureOr makeTCvtTmpType(MLIRContext *ctx, + pto::TCvtOp op) { + auto srcShape = getShapeVec(op.getSrc().getType()); + auto dstValid = getValidShapeVec(op.getDst().getType()); + auto srcTy = dyn_cast(op.getSrc().getType()); + auto dstTy = dyn_cast(op.getDst().getType()); + if (!srcTy || !dstTy || srcShape.size() != 2 || dstValid.size() != 2 || + hasDynamicDim(srcShape) || hasDynamicDim(dstValid)) + return failure(); + int64_t rows = dstValid[0], cols = dstValid[1]; + int64_t bytes = 0; + if (rows > 0 && cols > 0 && srcTy.getElementType().isF32()) { + int64_t head = 4 * 64 * std::min(cols / 64, 255); + int64_t remainder = cols % 64; + int64_t tail = remainder == 0 + ? 0 + : 32 * ((std::min(rows, 255) - 1) * + (srcShape[1] / 8) + + ceilDiv(remainder, 8)); + bytes = std::max(head, tail); + } else if (cols > 0 && srcTy.getElementType().isF16()) { + int64_t width = std::min(cols, 64); + int64_t halfToI16 = 32 * ceilDiv(width, 8); + int64_t halfToI8 = std::max(halfToI16, 128 + 32 * ceilDiv(width, 16)); + bytes = dstTy.getElementType().isInteger(8) ? halfToI8 : halfToI16; + } + int64_t allocatedBytes = std::max(32, ceilDiv(bytes, 32) * 32); + return makeVecTmpType(ctx, {1, allocatedBytes}, IntegerType::get(ctx, 8), + {1, allocatedBytes}); +} + +static LogicalResult materializeTCvtTmp(pto::TCvtOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (op.getTmp() || !tcvtNeedsTmp(op)) + return success(); + if (requireExplicitTmp) + return op.emitOpError( + "requires explicit tmp for non-saturating narrowing tcvt when PlanMemory is skipped"); + auto type = makeTCvtTmpType(ctx, op); + if (failed(type)) + return op.emitOpError( + "requires static tile_buf shapes to materialize implicit tcvt tmp"); + return replaceFixedDpsOpWithTmp(op.getOperation(), + {op.getSrc(), Value(), op.getDst()}, *type, + {1, 1, 1}, requireExplicitTmp, "tcvt"); +} + +static LogicalResult materializeTMrgSortTmp(pto::TMrgSortOp op, + bool requireExplicitTmp, + MLIRContext *ctx) { + if (!op.isFormat2WithoutTmp()) + return success(); + if (requireExplicitTmp) + return op.emitOpError( + "requires explicit tmp for tmrgsort format2 when PlanMemory is skipped"); + int64_t totalCols = 0; + Type elementType; + SmallVector operands; + for (Value src : op.getSrcs()) { + auto srcTy = dyn_cast(src.getType()); + auto shape = getShapeVec(src.getType()); + if (!srcTy || shape.size() != 2 || hasDynamicDim(shape)) + return op.emitOpError( + "requires static rank-2 tile_buf srcs to materialize tmrgsort tmp"); + if (!elementType) + elementType = srcTy.getElementType(); + totalCols += shape[1]; + operands.push_back(src); + } + if (!elementType || totalCols <= 0) + return op.emitOpError("failed to infer tmrgsort format2 tmp type"); + pto::TileBufType tmpType = + makeVecTmpType(ctx, {1, totalCols}, elementType, {1, totalCols}); + OpBuilder builder(op); + FailureOr tmp = createAllocTmp(builder, op.getLoc(), tmpType); + if (failed(tmp)) + return failure(); + OperationState state(op.getLoc(), op->getName()); + state.addOperands(operands); + state.addOperands(op.getDsts()); + state.addOperands(*tmp); + state.addOperands(op.getExcuted()); + state.addAttribute( + "operandSegmentSizes", + builder.getDenseI32ArrayAttr( + {static_cast(op.getSrcs().size()), 0, 1, 1, 1})); + copyAttrsExceptOperandSegments(op.getOperation(), state); + builder.create(state); + op.erase(); + return success(); +} + struct PTOMaterializeImplicitTmpPass : public PassWrapper> { @@ -274,6 +879,86 @@ struct PTOMaterializeImplicitTmpPass failed = true; } + SmallVector optionalTmpOps; + func.walk([&](Operation *op) { + if (isa(op)) + optionalTmpOps.push_back(op); + }); + + for (Operation *op : optionalTmpOps) { + LogicalResult result = + llvm::TypeSwitch(op) + .Case([&](auto typedOp) { + return replaceTColSumWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTGatherWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTQuantWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTPowWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTPowSWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTSort32WithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTXorWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return replaceTXorSWithTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return materializeTCvtTmp(typedOp, requireExplicitTmp, ctx); + }) + .Case([&](auto typedOp) { + return materializeTMrgSortTmp(typedOp, requireExplicitTmp, ctx); + }) + .Default([](Operation *) { return success(); }); + if (mlir::failed(result)) + failed = true; + } + + SmallVector rowReductionOps; + func.walk([&](Operation *op) { + if (isa(op)) + rowReductionOps.push_back(op); + }); + + for (Operation *op : rowReductionOps) { + LogicalResult result = + llvm::TypeSwitch(op) + .Case([&](auto typedOp) { + return replaceRowReductionWithTmp(typedOp, requireExplicitTmp, + ctx); + }) + .Default([](Operation *) { return success(); }); + if (mlir::failed(result)) + failed = true; + } + + SmallVector mandatoryTmpOps; + func.walk([&](Operation *op) { + if (isa(op)) + mandatoryTmpOps.push_back(op); + }); + for (Operation *op : mandatoryTmpOps) { + if (mlir::failed( + materializeFixedMandatoryTmp(op, requireExplicitTmp, ctx))) + failed = true; + } + if (failed) signalPassFailure(); } diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 6a92462fd4..1f669c556a 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -586,8 +586,14 @@ struct PlannerAnalysis { if (outputRoots.empty()) return; - for (Value scratch : getWrittenNonDpsOperands(op, dpsInits)) + for (Value scratch : getWrittenNonDpsOperands(op, dpsInits)) { addForbidAliasBetweenRoots(getRoots(scratch), outputRoots); + for (Value operand : op->getOperands()) { + if (operand == scratch || llvm::is_contained(dpsInits, operand)) + continue; + addForbidAliasBetweenRoots(getRoots(scratch), getRoots(operand)); + } + } } void recordInplacePolicyConflicts(Operation *op, ValueRange dpsInits) { diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 271f9e50a8..e0aedbd0b8 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -9386,7 +9386,11 @@ struct PTOCvtToEmitC : public OpConversionPattern { Value satModeVal = rewriter.create( loc, satModeTy, emitc::OpaqueAttr::get(ctx, satTok)); - SmallVector operands{dst, src, rmodeVal, satModeVal}; + SmallVector operands{dst, src}; + if (adaptor.getTmp()) + operands.push_back(peelUnrealized(adaptor.getTmp())); + operands.push_back(rmodeVal); + operands.push_back(satModeVal); rewriter.create( loc, TypeRange{}, "TCVT", diff --git a/test/lit/pto/implicit_tmp_arg_reductions.pto b/test/lit/pto/implicit_tmp_arg_reductions.pto new file mode 100644 index 0000000000..4f51f1da7c --- /dev/null +++ b/test/lit/pto/implicit_tmp_arg_reductions.pto @@ -0,0 +1,31 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR + +module { + func.func @implicit_arg_reduction_tmps() { + %src = pto.alloc_tile : !pto.tile_buf + %col_max = pto.alloc_tile : !pto.tile_buf + %col_min = pto.alloc_tile : !pto.tile_buf + %row_max = pto.alloc_tile : !pto.tile_buf + %row_min = pto.alloc_tile : !pto.tile_buf + pto.tcolargmax ins(%src : !pto.tile_buf) outs(%col_max : !pto.tile_buf) + pto.tcolargmin ins(%src : !pto.tile_buf) outs(%col_min : !pto.tile_buf) + pto.trowargmax ins(%src : !pto.tile_buf) outs(%row_max : !pto.tile_buf) + pto.trowargmin ins(%src : !pto.tile_buf) outs(%row_min : !pto.tile_buf) + return + } +} + +// IR-LABEL: func.func @implicit_arg_reduction_tmps +// IR: pto.tcolargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.tcolargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_optional_level3_invalid.pto b/test/lit/pto/implicit_tmp_optional_level3_invalid.pto new file mode 100644 index 0000000000..0f5951cda9 --- /dev/null +++ b/test/lit/pto/implicit_tmp_optional_level3_invalid.pto @@ -0,0 +1,24 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s + +module { + func.func @implicit_tpow_tmp_level3() { + %addr0 = arith.constant 0 : i64 + %addr1 = arith.constant 256 : i64 + %addr2 = arith.constant 512 : i64 + %base = pto.alloc_tile addr = %addr0 : !pto.tile_buf + %exp = pto.alloc_tile addr = %addr1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %addr2 : !pto.tile_buf + // CHECK: error: 'pto.tpow' op requires explicit tmp when PlanMemory is skipped + pto.tpow ins(%base, %exp : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/implicit_tmp_optional_ops_materialization.pto b/test/lit/pto/implicit_tmp_optional_ops_materialization.pto new file mode 100644 index 0000000000..7660a4bc88 --- /dev/null +++ b/test/lit/pto/implicit_tmp_optional_ops_materialization.pto @@ -0,0 +1,80 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 + +module { + func.func @implicit_tcolsum_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tcolsum ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) {isBinary = true} + return + } + + func.func @implicit_tgather_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %indices = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tgather ins(%src, %indices : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func @implicit_tquant_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tquant ins(%src, %fp : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) {quant_type = #pto} + return + } + + func.func @implicit_tpows_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + %scalar = arith.constant 2.0 : f32 + pto.tpows ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + return + } + + func.func @implicit_tsort32_tmp() { + %src = pto.alloc_tile : !pto.tile_buf + %idx = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tsort32 ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: func.func @implicit_tcolsum_tmp +// CHECK: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// CHECK: pto.tcolsum ins(%{{.*}}, %{{.*}} {{.*}}isBinary = true + +// CHECK-LABEL: func.func @implicit_tgather_tmp +// CHECK: pto.tgather ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + +// CHECK-LABEL: func.func @implicit_tquant_tmp +// CHECK: pto.tquant ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) outs(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) + +// CHECK-LABEL: func.func @implicit_tpows_tmp +// CHECK: pto.tpows ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, f32, !pto.tile_buf) + +// CHECK-LABEL: func.func @implicit_tsort32_tmp +// CHECK: pto.tsort32 ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + +// A5-LABEL: func.func @implicit_tgather_tmp +// A5: pto.tgather ins(%{{.*}}, %{{.*}}, %{{.*}} : +// A5-LABEL: func.func @implicit_tquant_tmp +// A5: pto.tquant ins(%{{.*}}, %{{.*}} : {{.*}}) outs(%{{.*}} : +// A5-NOT: pto.tquant ins({{.*}}) outs(%{{.*}}, %{{.*}} +// A5-LABEL: func.func @implicit_tpows_tmp +// A5: pto.tpows ins(%{{.*}}, %{{.*}}, %{{.*}} : diff --git a/test/lit/pto/implicit_tmp_remaining_level3_invalid.pto b/test/lit/pto/implicit_tmp_remaining_level3_invalid.pto new file mode 100644 index 0000000000..50666f4310 --- /dev/null +++ b/test/lit/pto/implicit_tmp_remaining_level3_invalid.pto @@ -0,0 +1,46 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s + +module { + func.func @tprelu_missing_tmp() { + %a0 = arith.constant 0 : i64 + %a1 = arith.constant 256 : i64 + %a2 = arith.constant 512 : i64 + %src0 = pto.alloc_tile addr = %a0 : !pto.tile_buf + %src1 = pto.alloc_tile addr = %a1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %a2 : !pto.tile_buf + pto.tprelu ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + + func.func @tcvt_missing_tmp() { + %a0 = arith.constant 0 : i64 + %a1 = arith.constant 256 : i64 + %src = pto.alloc_tile addr = %a0 : !pto.tile_buf + %dst = pto.alloc_tile addr = %a1 : !pto.tile_buf + pto.tcvt ins(%src {satmode = #pto} : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + + func.func @tmrgsort_missing_tmp(%executed : vector<4xi16>) { + %a0 = arith.constant 0 : i64 + %a1 = arith.constant 256 : i64 + %a2 = arith.constant 512 : i64 + %src0 = pto.alloc_tile addr = %a0 : !pto.tile_buf + %src1 = pto.alloc_tile addr = %a1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %a2 : !pto.tile_buf + pto.tmrgsort ins(%src0, %src1 no_tmp {exhausted = false} : !pto.tile_buf, !pto.tile_buf) outs(%dst, %executed : !pto.tile_buf, vector<4xi16>) + return + } +} + +// CHECK: error: 'pto.tprelu' op requires explicit tmp when PlanMemory is skipped +// CHECK: error: 'pto.tcvt' op requires explicit tmp for non-saturating narrowing tcvt when PlanMemory is skipped +// CHECK: error: 'pto.tmrgsort' op requires explicit tmp for tmrgsort format2 when PlanMemory is skipped diff --git a/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto b/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto new file mode 100644 index 0000000000..4f7aecd75e --- /dev/null +++ b/test/lit/pto/implicit_tmp_remaining_ops_materialization.pto @@ -0,0 +1,42 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --plan-memory-impl=modern --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @implicit_remaining_tmps(%executed : vector<4xi16>) { + %f32a = pto.alloc_tile : !pto.tile_buf + %f32b = pto.alloc_tile : !pto.tile_buf + %f32c = pto.alloc_tile : !pto.tile_buf + %mask = pto.alloc_tile : !pto.tile_buf + %i16dst = pto.alloc_tile : !pto.tile_buf + %scalar = arith.constant 3.0 : f32 + pto.tprelu ins(%f32a, %f32b : !pto.tile_buf, !pto.tile_buf) outs(%f32c : !pto.tile_buf) + pto.trem ins(%f32a, %f32b : !pto.tile_buf, !pto.tile_buf) outs(%f32c : !pto.tile_buf) + pto.trems ins(%f32a, %scalar : !pto.tile_buf, f32) outs(%f32c : !pto.tile_buf) + pto.tsel ins(%mask, %f32a, %f32b : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%f32c : !pto.tile_buf) + pto.tsels ins(%mask, %f32a, %scalar : !pto.tile_buf, !pto.tile_buf, f32) outs(%f32c : !pto.tile_buf) + pto.ttrans ins(%f32a : !pto.tile_buf) outs(%f32c : !pto.tile_buf) + pto.tcvt ins(%f32a {satmode = #pto} : !pto.tile_buf) outs(%i16dst : !pto.tile_buf) + %sort0 = pto.alloc_tile : !pto.tile_buf + %sort1 = pto.alloc_tile : !pto.tile_buf + %sortdst = pto.alloc_tile : !pto.tile_buf + pto.tmrgsort ins(%sort0, %sort1 no_tmp {exhausted = false} : !pto.tile_buf, !pto.tile_buf) outs(%sortdst, %executed : !pto.tile_buf, vector<4xi16>) + return + } +} + +// CHECK: pto.tprelu ins(%{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.trem ins(%{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.trems ins(%{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.tsel ins(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.tsels ins(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} +// CHECK: pto.ttrans ins(%{{.*}}, %{{.*}} +// CHECK: pto.tcvt ins(%{{.*}}, %{{.*}} +// CHECK: pto.tmrgsort ins(%{{.*}}, %{{.*}}, %{{.*}} {exhausted = false} diff --git a/test/lit/pto/implicit_tmp_row_reductions.pto b/test/lit/pto/implicit_tmp_row_reductions.pto new file mode 100644 index 0000000000..cfd3ade699 --- /dev/null +++ b/test/lit/pto/implicit_tmp_row_reductions.pto @@ -0,0 +1,31 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR + +module { + func.func @implicit_row_reduction_tmps() { + %src = pto.alloc_tile : !pto.tile_buf + %max = pto.alloc_tile : !pto.tile_buf + %min = pto.alloc_tile : !pto.tile_buf + %sum = pto.alloc_tile : !pto.tile_buf + %prod = pto.alloc_tile : !pto.tile_buf + pto.trowmax ins(%src : !pto.tile_buf) outs(%max : !pto.tile_buf) + pto.trowmin ins(%src : !pto.tile_buf) outs(%min : !pto.tile_buf) + pto.trowsum ins(%src : !pto.tile_buf) outs(%sum : !pto.tile_buf) + pto.trowprod ins(%src : !pto.tile_buf) outs(%prod : !pto.tile_buf) + return + } +} + +// IR-LABEL: func.func @implicit_row_reduction_tmps +// IR: pto.trowmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowsum ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// IR: pto.trowprod ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_xor_materialization.pto b/test/lit/pto/implicit_tmp_xor_materialization.pto new file mode 100644 index 0000000000..172d428c80 --- /dev/null +++ b/test/lit/pto/implicit_tmp_xor_materialization.pto @@ -0,0 +1,27 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR + +module { + func.func @implicit_xor_tmps() { + %scalar = arith.constant 3 : i16 + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %dst0 = pto.alloc_tile : !pto.tile_buf + %dst1 = pto.alloc_tile : !pto.tile_buf + pto.txor ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst0 : !pto.tile_buf) + pto.txors ins(%src0, %scalar : !pto.tile_buf, i16) outs(%dst1 : !pto.tile_buf) + return + } +} + +// IR-LABEL: func.func @implicit_xor_tmps +// IR: pto.txor ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) +// IR: pto.txors ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, i16, !pto.tile_buf) diff --git a/test/lit/pto/issue533_loop_zero_trip_sync_regression.pto b/test/lit/pto/issue533_loop_zero_trip_sync_regression.pto index 58e1c7b739..b213736de1 100644 --- a/test/lit/pto/issue533_loop_zero_trip_sync_regression.pto +++ b/test/lit/pto/issue533_loop_zero_trip_sync_regression.pto @@ -24,6 +24,7 @@ module attributes {pto.target_arch = "a2a3"} { %c128_i64 = arith.constant 128 : i64 %c160_i64 = arith.constant 160 : i64 %c4256_i64 = arith.constant 4256 : i64 + %c32768_i64 = arith.constant 32768 : i64 %c1024_index = arith.constant 1024 : index %c128_index = arith.constant 128 : index %c1_index = arith.constant 1 : index @@ -32,6 +33,7 @@ module attributes {pto.target_arch = "a2a3"} { %c0_index = arith.constant 0 : index %c8_index = arith.constant 8 : index %c16_index = arith.constant 16 : index + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf %all_oi_tmp__co_l0_rv_v1_view = pto.make_tensor_view %arg0, shape = [%c1024_index, %c128_index], strides = [%c128_index, %c1_index] {layout = #pto.layout}: !pto.tensor_view %all_cur_mi__co_l0_rv_v1_view = pto.make_tensor_view %arg1, shape = [%c512_index, %c1_index], strides = [%c1_index, %c512_index] {layout = #pto.layout}: !pto.tensor_view %all_cur_li__co_l0_rv_v1_view = pto.make_tensor_view %arg2, shape = [%c512_index, %c1_index], strides = [%c1_index, %c512_index] {layout = #pto.layout}: !pto.tensor_view @@ -97,9 +99,9 @@ module attributes {pto.target_arch = "a2a3"} { pto.tadd ins(%li__rm_a0_tmp_v19, %li__rm_a1_tmp_v20 : !pto.tile_buf, !pto.tile_buf) outs(%li__row_major_tmp_v21 : !pto.tile_buf) %3 = pto.alloc_tile addr = %c96_i64 : !pto.tile_buf %4 = pto.alloc_tile addr = %c160_i64 : !pto.tile_buf - pto.trowexpandmul ins(%oi__tile, %alpha__tile : !pto.tile_buf, !pto.tile_buf) outs(%4 : !pto.tile_buf) + pto.trowexpandmul ins(%oi__tile, %alpha__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%4 : !pto.tile_buf) %5 = pto.alloc_tile addr = %c10464_i64 : !pto.tile_buf - pto.trowexpandmul ins(%oi_tmp_valid__tile, %beta__tile : !pto.tile_buf, !pto.tile_buf) outs(%5 : !pto.tile_buf) + pto.trowexpandmul ins(%oi_tmp_valid__tile, %beta__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%5 : !pto.tile_buf) %6 = pto.alloc_tile addr = %c10464_i64 : !pto.tile_buf pto.tadd ins(%4, %5 : !pto.tile_buf, !pto.tile_buf) outs(%6 : !pto.tile_buf) %mi__ssa_v3 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf @@ -111,7 +113,7 @@ module attributes {pto.target_arch = "a2a3"} { pto.tmov ins(%6 : !pto.tile_buf) outs(%oi__tile_mv : !pto.tile_buf) } %ctx__tile = pto.alloc_tile addr = %c6304_i64 : !pto.tile_buf - pto.trowexpanddiv ins(%oi__tile, %li__tile : !pto.tile_buf, !pto.tile_buf) outs(%ctx__tile : !pto.tile_buf) + pto.trowexpanddiv ins(%oi__tile, %li__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%ctx__tile : !pto.tile_buf) %ctx_flat__tile = pto.alloc_tile addr = %c6304_i64 : !pto.tile_buf %ctx_flat_bf16__tile = pto.alloc_tile addr = %c4256_i64 : !pto.tile_buf pto.tcvt ins(%ctx_flat__tile{rmode = #pto} : !pto.tile_buf) outs(%ctx_flat_bf16__tile : !pto.tile_buf) diff --git a/test/lit/pto/issue533_loop_zero_trip_sync_regression_gss.pto b/test/lit/pto/issue533_loop_zero_trip_sync_regression_gss.pto index 56e75b410b..e8790c460f 100644 --- a/test/lit/pto/issue533_loop_zero_trip_sync_regression_gss.pto +++ b/test/lit/pto/issue533_loop_zero_trip_sync_regression_gss.pto @@ -24,6 +24,7 @@ module attributes {pto.target_arch = "a2a3"} { %c128_i64 = arith.constant 128 : i64 %c160_i64 = arith.constant 160 : i64 %c4256_i64 = arith.constant 4256 : i64 + %c32768_i64 = arith.constant 32768 : i64 %c1024_index = arith.constant 1024 : index %c128_index = arith.constant 128 : index %c1_index = arith.constant 1 : index @@ -32,6 +33,7 @@ module attributes {pto.target_arch = "a2a3"} { %c0_index = arith.constant 0 : index %c8_index = arith.constant 8 : index %c16_index = arith.constant 16 : index + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf %all_oi_tmp__co_l0_rv_v1_view = pto.make_tensor_view %arg0, shape = [%c1024_index, %c128_index], strides = [%c128_index, %c1_index] {layout = #pto.layout}: !pto.tensor_view %all_cur_mi__co_l0_rv_v1_view = pto.make_tensor_view %arg1, shape = [%c512_index, %c1_index], strides = [%c1_index, %c512_index] {layout = #pto.layout}: !pto.tensor_view %all_cur_li__co_l0_rv_v1_view = pto.make_tensor_view %arg2, shape = [%c512_index, %c1_index], strides = [%c1_index, %c512_index] {layout = #pto.layout}: !pto.tensor_view @@ -97,9 +99,9 @@ module attributes {pto.target_arch = "a2a3"} { pto.tadd ins(%li__rm_a0_tmp_v19, %li__rm_a1_tmp_v20 : !pto.tile_buf, !pto.tile_buf) outs(%li__row_major_tmp_v21 : !pto.tile_buf) %3 = pto.alloc_tile addr = %c96_i64 : !pto.tile_buf %4 = pto.alloc_tile addr = %c160_i64 : !pto.tile_buf - pto.trowexpandmul ins(%oi__tile, %alpha__tile : !pto.tile_buf, !pto.tile_buf) outs(%4 : !pto.tile_buf) + pto.trowexpandmul ins(%oi__tile, %alpha__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%4 : !pto.tile_buf) %5 = pto.alloc_tile addr = %c10464_i64 : !pto.tile_buf - pto.trowexpandmul ins(%oi_tmp_valid__tile, %beta__tile : !pto.tile_buf, !pto.tile_buf) outs(%5 : !pto.tile_buf) + pto.trowexpandmul ins(%oi_tmp_valid__tile, %beta__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%5 : !pto.tile_buf) %6 = pto.alloc_tile addr = %c10464_i64 : !pto.tile_buf pto.tadd ins(%4, %5 : !pto.tile_buf, !pto.tile_buf) outs(%6 : !pto.tile_buf) %mi__ssa_v3 = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf @@ -111,7 +113,7 @@ module attributes {pto.target_arch = "a2a3"} { pto.tmov ins(%6 : !pto.tile_buf) outs(%oi__tile_mv : !pto.tile_buf) } %ctx__tile = pto.alloc_tile addr = %c6304_i64 : !pto.tile_buf - pto.trowexpanddiv ins(%oi__tile, %li__tile : !pto.tile_buf, !pto.tile_buf) outs(%ctx__tile : !pto.tile_buf) + pto.trowexpanddiv ins(%oi__tile, %li__tile, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%ctx__tile : !pto.tile_buf) %ctx_flat__tile = pto.alloc_tile addr = %c6304_i64 : !pto.tile_buf %ctx_flat_bf16__tile = pto.alloc_tile addr = %c4256_i64 : !pto.tile_buf pto.tcvt ins(%ctx_flat__tile{rmode = #pto} : !pto.tile_buf) outs(%ctx_flat_bf16__tile : !pto.tile_buf) diff --git a/test/lit/pto/issue646_pipev_repeat_prune.pto b/test/lit/pto/issue646_pipev_repeat_prune.pto index d610b3326b..bb3584872b 100644 --- a/test/lit/pto/issue646_pipev_repeat_prune.pto +++ b/test/lit/pto/issue646_pipev_repeat_prune.pto @@ -9,11 +9,13 @@ module { %c0_i64 = arith.constant 0 : i64 %c4096_i64 = arith.constant 4096 : i64 %c8192_i64 = arith.constant 8192 : i64 + %c32768_i64 = arith.constant 32768 : i64 %acc = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf %src1 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf %scale = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf pto.tadd ins(%acc, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) - pto.trowexpanddiv ins(%acc, %scale : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) + pto.trowexpanddiv ins(%acc, %scale, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) return } @@ -25,11 +27,13 @@ module { %c0_i64 = arith.constant 0 : i64 %c4096_i64 = arith.constant 4096 : i64 %c12288_i64 = arith.constant 12288 : i64 + %c32768_i64 = arith.constant 32768 : i64 %acc = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf %src1 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf %scale = pto.alloc_tile addr = %c12288_i64 : !pto.tile_buf + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf pto.tadd ins(%acc, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) - pto.trowexpanddiv ins(%acc, %scale : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) + pto.trowexpanddiv ins(%acc, %scale, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) return } @@ -41,11 +45,13 @@ module { %c0_i64 = arith.constant 0 : i64 %c4096_i64 = arith.constant 4096 : i64 %c8192_i64 = arith.constant 8192 : i64 + %c32768_i64 = arith.constant 32768 : i64 %acc = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf %src1 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf %scale = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf pto.tadd ins(%acc, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) - pto.trowexpanddiv ins(%acc, %scale : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) + pto.trowexpanddiv ins(%acc, %scale, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) return } @@ -60,14 +66,16 @@ module { %c8192_i64 = arith.constant 8192 : i64 %c12288_i64 = arith.constant 12288 : i64 %c16384_i64 = arith.constant 16384 : i64 + %c32768_i64 = arith.constant 32768 : i64 %acc = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf %src1 = pto.alloc_tile addr = %c4096_i64 : !pto.tile_buf %scale = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf %tmp = pto.alloc_tile addr = %c12288_i64 : !pto.tile_buf %tmp_src = pto.alloc_tile addr = %c16384_i64 : !pto.tile_buf + %rowexpand_tmp = pto.alloc_tile addr = %c32768_i64 : !pto.tile_buf pto.tadd ins(%acc, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) pto.tadd ins(%tmp, %tmp_src : !pto.tile_buf, !pto.tile_buf) outs(%tmp : !pto.tile_buf) - pto.trowexpanddiv ins(%acc, %scale : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) + pto.trowexpanddiv ins(%acc, %scale, %rowexpand_tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) return } diff --git a/test/lit/pto/tpow_fp_missing_tmp_invalid.pto b/test/lit/pto/tpow_fp_missing_tmp_invalid.pto index 23d832af96..01813ab8ad 100644 --- a/test/lit/pto/tpow_fp_missing_tmp_invalid.pto +++ b/test/lit/pto/tpow_fp_missing_tmp_invalid.pto @@ -6,8 +6,9 @@ // 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. -// Floating-point tpow requires tmp scratch (used by PowF / TPowFloat). -// RUN: not ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s +// Floating-point tpow gets an implicit tmp before memplan. +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 module { func.func @tpow_fp_missing_tmp() { @@ -19,4 +20,8 @@ module { } } -// CHECK: expects tmp when element type is floating-point +// A3: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// A3: pto.tpow ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + +// A5: pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// A5: pto.tpow ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) From ebab68c69afc47b126377b60dc3028fbf5400847 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 4 Aug 2026 16:23:33 +0800 Subject: [PATCH 03/17] =?UTF-8?q?Fix=20rowexpand=5Ftile=5Fnative.pto=20RUN?= =?UTF-8?q?=20level=20(level3=20=E2=86=92=20level2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This test verifies tile-native IR preservation (!pto.tile_buf kept native, not lowered to memref), which is a level2 behavior. Running it at level3 invokes the full pipeline including PlanMemory skip, which triggers the new implicit-tmp verifier rejecting mode-1 col-major row-expand ops without explicit tmp (since implicit tmp cannot be materialized when PlanMemory is skipped at level3). At level2, PlanMemory runs and materializes implicit tmp for mode-1 row-expand ops correctly, matching the test's tile-native intent. --- test/lit/pto/rowexpand_tile_native.pto | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/lit/pto/rowexpand_tile_native.pto b/test/lit/pto/rowexpand_tile_native.pto index 61d9f7a4ec..a7523bae6b 100644 --- a/test/lit/pto/rowexpand_tile_native.pto +++ b/test/lit/pto/rowexpand_tile_native.pto @@ -6,8 +6,8 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE -// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC +// RUN: ptoas --pto-level=level2 --pto-arch=a5 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level2 --pto-arch=a3 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { func.func private @trowexpand_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { From 55a3f5e7ad8f3d21cc16347db4059e9811ae32c8 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 4 Aug 2026 17:18:17 +0800 Subject: [PATCH 04/17] Fix implicit-tmp tile-native lit tests after rebase onto main After rebasing onto main (which removed memref compatibility and added the optional tmp operand to tcvt), four lit tests had stale inputs or CHECK lines that no longer matched the generated IR: - tci_implicit_tmp_materialization.pto: rewrite inputs from memref to tile-native partition_tensor_view so the test parses on main. - tquant_no_implicit_tmp_a3.pto: this is a negative test expecting the implicit-tmp verifier to reject a dynamic-valid-shape tquant src. Use 'not ptoas' (the repo's standard negative-test idiom) so the pipe exit code reflects the expected failure, and match the emitted diagnostic instead of an unreachable func label. - cvt_tile_native.pto / tcvt_low_precision_a5_valid.pto: the optional tmp operand now prints 'operandSegmentSizes = array' in the tcvt IR; wildcard it in the CHECK lines. Verified on remote A3 (LLVM21): full check-pto = 1565 passed, 0 failed, 1 unsupported. --- test/lit/pto/cvt_tile_native.pto | 2 +- .../pto/tci_implicit_tmp_materialization.pto | 19 +++++++++++++++---- test/lit/pto/tcvt_low_precision_a5_valid.pto | 6 +++--- test/lit/pto/tquant_no_implicit_tmp_a3.pto | 18 +++++++++--------- 4 files changed, 28 insertions(+), 17 deletions(-) diff --git a/test/lit/pto/cvt_tile_native.pto b/test/lit/pto/cvt_tile_native.pto index af37bd6e13..c33ff11029 100644 --- a/test/lit/pto/cvt_tile_native.pto +++ b/test/lit/pto/cvt_tile_native.pto @@ -23,7 +23,7 @@ module { } // NATIVE-LABEL: func.func private @tcvt_arg( -// NATIVE: pto.tcvt ins(%arg0 {rmode = #pto, satmode = #pto} : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) +// NATIVE: pto.tcvt ins(%arg0 {{.*}}rmode = #pto, satmode = #pto{{.*}} : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) // NATIVE-NOT: memref< // EMITC-LABEL: tcvt_arg( diff --git a/test/lit/pto/tci_implicit_tmp_materialization.pto b/test/lit/pto/tci_implicit_tmp_materialization.pto index 78cec449e5..e8fbd41bcb 100644 --- a/test/lit/pto/tci_implicit_tmp_materialization.pto +++ b/test/lit/pto/tci_implicit_tmp_materialization.pto @@ -1,16 +1,27 @@ +// 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 PROGRAM 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. + +// TCI without an explicit tmp operand: ptoas should materialize an implicit +// tmp (a 1x192 fp32 buffer on A3) and pass it to the PTO-ISA TCI interface, +// then memory-plan all variables together with the tmp. + // RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR // RUN: ptoas --pto-arch=a3 --pto-level=level2 %s 2>&1 | FileCheck %s --check-prefix=CPP // RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5IR module { - func.func @tci_implicit_tmp(%dst: memref<32xi32, #pto.address_space>) { + func.func @tci_implicit_tmp(%dst: !pto.partition_tensor_view<1x32xi32>) { %c0_i32 = arith.constant 0 : i32 - %src = memref.reinterpret_cast %dst to offset: [0], sizes: [1, 32], strides: [32, 1] {layout = #pto.layout} : memref<32xi32, #pto.address_space> to memref<1x32xi32, strided<[32, 1], offset: ?>, #pto.address_space> - %tile = pto.alloc_tile : !pto.tile_buf + %tile = pto.declare_tile -> !pto.tile_buf pto.tci ins(%c0_i32 : i32) outs(%tile : !pto.tile_buf) pto.tstore ins(%tile : !pto.tile_buf) - outs(%src : memref<1x32xi32, strided<[32, 1], offset: ?>, #pto.address_space>) {layout = #pto.layout, pto.inferred_layout = true} + outs(%dst : !pto.partition_tensor_view<1x32xi32>) return } } diff --git a/test/lit/pto/tcvt_low_precision_a5_valid.pto b/test/lit/pto/tcvt_low_precision_a5_valid.pto index 15cc41a6da..aa9a8271b1 100644 --- a/test/lit/pto/tcvt_low_precision_a5_valid.pto +++ b/test/lit/pto/tcvt_low_precision_a5_valid.pto @@ -30,9 +30,9 @@ module { // CHECK: func.func @tcvt_low_precision_a5_valid() attributes {pto.kernel_kind = #pto.kernel_kind} // CHECK: pto.declare_tile -> !pto.tile_buf // CHECK: pto.declare_tile -> !pto.tile_buf -// CHECK: pto.tcvt ins(%{{.*}} {rmode = #pto, satmode = #pto} : !pto.tile_buf +// CHECK: pto.tcvt ins(%{{.*}} {{.*}}rmode = #pto, satmode = #pto{{.*}} : !pto.tile_buf // CHECK: outs(%{{.*}} : !pto.tile_buf) -// CHECK: pto.tcvt ins(%{{.*}} {rmode = #pto, satmode = #pto} : !pto.tile_buf +// CHECK: pto.tcvt ins(%{{.*}} {{.*}}rmode = #pto, satmode = #pto{{.*}} : !pto.tile_buf // CHECK: outs(%{{.*}} : !pto.tile_buf) -// CHECK: pto.tcvt ins(%{{.*}} {rmode = #pto, satmode = #pto} : !pto.tile_buf +// CHECK: pto.tcvt ins(%{{.*}} {{.*}}rmode = #pto, satmode = #pto{{.*}} : !pto.tile_buf // CHECK: outs(%{{.*}} : !pto.tile_buf) diff --git a/test/lit/pto/tquant_no_implicit_tmp_a3.pto b/test/lit/pto/tquant_no_implicit_tmp_a3.pto index be39a34079..02c5ba0d8d 100644 --- a/test/lit/pto/tquant_no_implicit_tmp_a3.pto +++ b/test/lit/pto/tquant_no_implicit_tmp_a3.pto @@ -2,11 +2,16 @@ // 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, +// THIS PROGRAM IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a3 %s -emit-pto-ir 2>&1 | FileCheck %s +// TQUANT with a dynamic-valid-shape src and no explicit tmp cannot have its +// implicit tmp materialized (the tmp type is derived from the static src +// shape). The implicit-tmp pass must reject it with a clear diagnostic +// instead of silently dropping the tmp. + +// RUN: not ptoas --pto-arch=a3 %s -emit-pto-ir 2>&1 | FileCheck %s module { func.func @tquant_no_implicit_tmp_a3(%valid_row: index, %valid_col: index) @@ -15,13 +20,6 @@ module { %fp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile valid_row = %valid_row valid_col = %valid_col : !pto.tile_buf - // CHECK-LABEL: func.func @tquant_no_implicit_tmp_a3 - // CHECK-SAME: (%[[ROW:arg[0-9]+]]: index, %[[COL:arg[0-9]+]]: index) - // CHECK: pto.alloc_tile{{.*}}valid_row = %[[ROW]] valid_col = %[[COL]]{{.*}} : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) @@ -29,3 +27,5 @@ module { return } } + +// CHECK: 'pto.tquant' op requires static tile_buf src to materialize implicit tquant tmp From 68e079d8c2c2188b37131ef702cc28562d1b889b Mon Sep 17 00:00:00 2001 From: FangRui Date: Wed, 5 Aug 2026 15:14:03 +0800 Subject: [PATCH 05/17] Fix ptodsl bindings and op printers for optional tmp operand Two CI failures after adding the optional tmp operand: 1. TypeError: trowmax() takes 2 positional arguments but 3 were given - The MLIR-generated Python bindings expose tmp as a keyword-only argument (tmp=None after *) for ops whose ODS declares Optional:$tmp. The ptodsl wrappers in _ops.py were still passing tmp positionally, which collided with the dst slot. - Fix trowsum/max/min/prod/argmax/argmin, tcolargmax/argmin, txor, txors to pass tmp=... as a keyword argument (matching tcolsum, tcvt, tsel, tci, tmrgsort, tgather which already did so). 2. docs_as_test cvt round-trip textual drift - After adding the optional tmp operand, MLIR auto-attaches an 'operandSegmentSizes' attribute to tcvt/tpow/tpows/tcolsum, which the custom assembly printers leaked via printOptionalAttrDict, breaking the parse->print round-trip stability checked by the docs_as_test harness. - Elide 'operandSegmentSizes' in TCvtOp/TPowOp/TPowSOp/TColSumOp printers (the repo's established idiom, already used by TRowExpand*/TXor/TGather/TSort32/TMrgSort). Verified on remote A3 (LLVM21): ptodsl_docs_as_test PASS, test_ptoas_frontend_verify PASS, full check-pto = 1565 passed, 0 failed, 1 unsupported. --- lib/PTO/IR/PTO.cpp | 13 ++++++++----- ptodsl/ptodsl/_ops.py | 20 ++++++++++---------- 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 984e0ad7d4..5d98065c50 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -6337,7 +6337,7 @@ void mlir::pto::TColSumOp::print(OpAsmPrinter &p) { // Format 2: ins(%src, %tmp {isBinary = ...}: type, type) outs(%dst : type) p << " ins(" << getSrc() << ", " << getTmp(); // Print isBinary attribute if present - SmallVector elidedAttrs; + SmallVector elidedAttrs = {"operandSegmentSizes"}; if (!getIsBinaryAttr() || getIsBinaryAttr().getValue() == false) { elidedAttrs.push_back("isBinary"); } @@ -6352,7 +6352,7 @@ void mlir::pto::TColSumOp::print(OpAsmPrinter &p) { // Print remaining attributes for format 1 (excluding isBinary) if (!getTmp()) { - SmallVector elidedAttrs = {"isBinary"}; + SmallVector elidedAttrs = {"isBinary", "operandSegmentSizes"}; p.printOptionalAttrDict((*this)->getAttrs(), elidedAttrs); } } @@ -9645,7 +9645,8 @@ void mlir::pto::TCvtOp::print(OpAsmPrinter &p) { } attrs.set(attr.getName(), attr.getValue()); } - p.printOptionalAttrDict(attrs.getAttrs()); + p.printOptionalAttrDict(attrs.getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); p << " : " << getSrc().getType(); if (getTmp()) p << ", " << getTmp().getType(); @@ -11710,7 +11711,8 @@ void mlir::pto::TPowOp::print(OpAsmPrinter &p) { p << ", " << getTmp().getType(); p << ")"; p << " outs(" << getDst() << " : " << getDst().getType() << ")"; - p.printOptionalAttrDict((*this)->getAttrs()); + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); } // TPOWS assembly format: @@ -11767,7 +11769,8 @@ void mlir::pto::TPowSOp::print(OpAsmPrinter &p) { p << ", " << getTmp().getType(); p << ")"; p << " outs(" << getDst() << " : " << getDst().getType() << ")"; - p.printOptionalAttrDict((*this)->getAttrs()); + p.printOptionalAttrDict((*this)->getAttrs(), + /*elidedAttrs=*/{"operandSegmentSizes"}); } static ParseResult parseTRowExpandBinaryLikeOp(OpAsmParser &parser, diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 914c39695b..5ec91c9a66 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -3490,8 +3490,8 @@ def trowsum(src, tmp, dst): """``pto.trowsum ins(src, tmp) outs(dst)``.""" _pto.trowsum( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3499,8 +3499,8 @@ def trowmax(src, tmp, dst): """``pto.trowmax ins(src, tmp) outs(dst)``.""" _pto.trowmax( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3508,8 +3508,8 @@ def trowmin(src, tmp, dst): """``pto.trowmin ins(src, tmp) outs(dst)``.""" _pto.trowmin( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3517,8 +3517,8 @@ def trowprod(src, tmp, dst): """``pto.trowprod ins(src, tmp) outs(dst)``.""" _pto.trowprod( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3526,8 +3526,8 @@ def trowargmax(src, tmp, dst): """``pto.trowargmax ins(src, tmp) outs(dst)``.""" _pto.trowargmax( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3535,8 +3535,8 @@ def trowargmin(src, tmp, dst): """``pto.trowargmin ins(src, tmp) outs(dst)``.""" _pto.trowargmin( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3578,8 +3578,8 @@ def tcolargmax(src, tmp, dst): """``pto.tcolargmax ins(src, tmp) outs(dst)``.""" _pto.tcolargmax( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -3587,8 +3587,8 @@ def tcolargmin(src, tmp, dst): """``pto.tcolargmin ins(src, tmp) outs(dst)``.""" _pto.tcolargmin( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -4118,8 +4118,8 @@ def txor(src0, src1, tmp, dst): _pto.txor( unwrap_surface_value(src0), unwrap_surface_value(src1), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) @@ -4128,8 +4128,8 @@ def txors(src, scalar, tmp, dst): _pto.txors( unwrap_surface_value(src), _coerce_tile_scalar_operand(src, scalar, context="txors"), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) From 110a9992dbd8e9009a5158032ddbb3ac5a5cf3d3 Mon Sep 17 00:00:00 2001 From: FangRui Date: Wed, 5 Aug 2026 16:01:03 +0800 Subject: [PATCH 06/17] Resolve TGather verifier conflict after rebase onto main main added a stricter TGather contract (PR #1080 'Add TGATHER indices and mask'): A2/A3 index-form and all compare-form tgather ops now require an explicit tmp ("expects both indices and tmp"), while A5 index-form deliberately emits WITHOUT tmp (TGATHER(src, indices, dst)). This conflicts with the implicit-tmp design, which materialized a tmp for any omitted-tmp tgather. Restore main's contract and stop materializing tgather tmp: - PTO.cpp TGatherOp::verify: revert to main's logic (A2/A3 needs tmp, A5 index-form lets verifyIndexForm handle the missing tmp). Remove the implicit-tmp early-return 'if (!getTmp()) return success()'. - PTOMaterializeImplicitTmp.cpp: drop TGatherOp from the optional-tmp dispatch and isa<...> list (tgather no longer gets an implicit tmp). Mark replaceTGatherWithTmp [[maybe_unused]] to keep the documented implementation around without a -Werror unused-function failure. - implicit_tmp_optional_ops_materialization.pto: drop the implicit_tgather_tmp test case (A2/A3 index-form now requires an explicit tmp; the implicit-tmp path no longer applies to tgather). Verified on remote A3 (LLVM21): full check-pto = 1646 passed, 0 failed, 1 unsupported; ptodsl_docs_as_test + frontend_verify PASS. --- lib/PTO/IR/PTO.cpp | 20 ++++++------------- .../Transforms/PTOMaterializeImplicitTmp.cpp | 11 ++++------ ...licit_tmp_optional_ops_materialization.pto | 14 ------------- 3 files changed, 10 insertions(+), 35 deletions(-) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 5d98065c50..202fefe1e3 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -7845,18 +7845,14 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { if (getAxisAttr()) return emitOpError("axis attribute must not be provided without maskPattern"); if (getCdst() || getKValue()) { - if (!getCdst() || !getKValue()) - return emitOpError("compare-form tgather expects dst, cdst, and kValue"); + if (!getCdst() || !getKValue() || !getTmp()) + return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp"); if (getIndices()) return emitOpError("compare-form tgather does not take indices"); - if (!getTmp()) - return success(); return verifyCompareForm(/*allowA5SrcTypes=*/false); } - if (!getIndices()) - return emitOpError("index-form tgather expects indices"); - if (!getTmp()) - return success(); + if (!getIndices() || !getTmp()) + return emitOpError("index-form tgather expects both indices and tmp"); return verifyIndexForm(/*allow16BitIndices=*/false, /*allowA5ElemTypes=*/false); }; @@ -7869,18 +7865,14 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() { if (getAxisAttr()) return emitOpError("axis attribute must not be provided without maskPattern"); if (getCdst() || getKValue()) { - if (!getCdst() || !getKValue()) - return emitOpError("compare-form tgather expects dst, cdst, and kValue"); + if (!getCdst() || !getKValue() || !getTmp()) + return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp"); if (getIndices()) return emitOpError("compare-form tgather does not take indices"); - if (!getTmp()) - return success(); return verifyCompareForm(/*allowA5SrcTypes=*/true); } if (!getIndices()) return emitOpError("index-form tgather expects indices"); - if (!getTmp()) - return success(); return verifyIndexForm(/*allow16BitIndices=*/true, /*allowA5ElemTypes=*/true); }; diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index 8c0d60eef7..d35c7147e3 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -376,9 +376,9 @@ static LogicalResult replaceTPowSWithTmp(pto::TPowSOp op, return success(); } -static LogicalResult replaceTGatherWithTmp(pto::TGatherOp op, - bool requireExplicitTmp, - MLIRContext *ctx) { +[[maybe_unused]] static LogicalResult +replaceTGatherWithTmp(pto::TGatherOp op, bool requireExplicitTmp, + MLIRContext *ctx) { if (op.getTmp() || op.hasMaskForm()) return success(); if (!op.hasIndexForm() && !op.hasCompareForm()) @@ -881,7 +881,7 @@ struct PTOMaterializeImplicitTmpPass SmallVector optionalTmpOps; func.walk([&](Operation *op) { - if (isa(op)) optionalTmpOps.push_back(op); @@ -893,9 +893,6 @@ struct PTOMaterializeImplicitTmpPass .Case([&](auto typedOp) { return replaceTColSumWithTmp(typedOp, requireExplicitTmp, ctx); }) - .Case([&](auto typedOp) { - return replaceTGatherWithTmp(typedOp, requireExplicitTmp, ctx); - }) .Case([&](auto typedOp) { return replaceTQuantWithTmp(typedOp, requireExplicitTmp, ctx); }) diff --git a/test/lit/pto/implicit_tmp_optional_ops_materialization.pto b/test/lit/pto/implicit_tmp_optional_ops_materialization.pto index 7660a4bc88..c7b4736c44 100644 --- a/test/lit/pto/implicit_tmp_optional_ops_materialization.pto +++ b/test/lit/pto/implicit_tmp_optional_ops_materialization.pto @@ -18,15 +18,6 @@ module { return } - func.func @implicit_tgather_tmp() { - %src = pto.alloc_tile : !pto.tile_buf - %indices = pto.alloc_tile : !pto.tile_buf - %dst = pto.alloc_tile : !pto.tile_buf - pto.tgather ins(%src, %indices : !pto.tile_buf, !pto.tile_buf) - outs(%dst : !pto.tile_buf) - return - } - func.func @implicit_tquant_tmp() { %src = pto.alloc_tile : !pto.tile_buf %fp = pto.alloc_tile : !pto.tile_buf @@ -59,9 +50,6 @@ module { // CHECK: pto.alloc_tile addr = {{.*}} : !pto.tile_buf // CHECK: pto.tcolsum ins(%{{.*}}, %{{.*}} {{.*}}isBinary = true -// CHECK-LABEL: func.func @implicit_tgather_tmp -// CHECK: pto.tgather ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) - // CHECK-LABEL: func.func @implicit_tquant_tmp // CHECK: pto.tquant ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) outs(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) @@ -71,8 +59,6 @@ module { // CHECK-LABEL: func.func @implicit_tsort32_tmp // CHECK: pto.tsort32 ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) -// A5-LABEL: func.func @implicit_tgather_tmp -// A5: pto.tgather ins(%{{.*}}, %{{.*}}, %{{.*}} : // A5-LABEL: func.func @implicit_tquant_tmp // A5: pto.tquant ins(%{{.*}}, %{{.*}} : {{.*}}) outs(%{{.*}} : // A5-NOT: pto.tquant ins({{.*}}) outs(%{{.*}}, %{{.*}} From 0c98de307194a931b41ec3890c227ab862b51adb Mon Sep 17 00:00:00 2001 From: FangRui Date: Wed, 5 Aug 2026 16:13:07 +0800 Subject: [PATCH 07/17] Update implicit-tmp design doc after main removed memref IR main removed the memref compatibility layer (PTOViewToMemref, PTOMaterializeTileHandles, memref.alloc, pto.pointer_cast / bind_tile ops) and tightened the TGather tmp contract (PR #1080). The design doc still referenced these, so refresh it: - Pipeline diagram: replace deleted PTOViewToMemref / PTOMaterializeTileHandles / PTOToEmitC with the actual current passes (PTOFusionRegionGen -> pto-materialize-implicit-tmp -> PTORematerializeFixpipeVectorQuant -> pto-plan-memory -> PTOResolveReservedBuffers -> sync passes -> PTOResolveBufferSelect -> EmitPTOManual). - Drop the 'memref.alloc' / 'pto.pointer_cast' / 'pto.bind_tile' bullets (those ops no longer exist) and the now-meaningless 'CHECK-NOT: memref.alloc' lit assertions; note that optional-tmp ops elide 'operandSegmentSizes' in the custom printer instead. - Rename stray PTOToEmitC references to EmitPTOManual (the real PTO->EmitC pass name). - TGATHER: mark it as NOT participating in implicit-tmp materialize. main requires an explicit tmp on A2/A3 index/compare-form and emits A5 index-form without tmp, so tgather is removed from the materialize dispatch; add a dedicated classification row and rewrite the TGATHER data-movement section accordingly. --- ...oas-implicit-tmp-materialization-design.md | 36 ++++++++++--------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md index 719563b370..c9c3c9cdec 100644 --- a/docs/designs/ptoas-implicit-tmp-materialization-design.md +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -33,16 +33,17 @@ PTOAS 前端 IR 中很多 tile op 的 `tmp` operand 是可选的。当前如果 pto-materialize-implicit-tmp ``` -该 pass 运行在 `PTOViewToMemref` 之后、`pto-plan-memory` 之前: +该 pass 运行在 fusion/调度 pass 之后、`pto-plan-memory` 之前: ```text -PTOViewToMemref +PTOFusionRegionGen -> pto-materialize-implicit-tmp - -> pto-plan-memory + -> PTORematerializeFixpipeVectorQuant + -> pto-plan-memory (level1/level2 only; skipped at level3) -> PTOResolveReservedBuffers - -> sync passes - -> PTOMaterializeTileHandles - -> PTOToEmitC + -> sync passes (InsertSync / GraphSyncSolver / BarrierAll ...) + -> PTOResolveBufferSelect + -> EmitPTOManual (PTO -> EmitC lowering) ``` pass 的职责是扫描所有已纳入改造的目标 op。如果 op 没有 tmp operand,就根据该 op 的 `TmpRequirement` 在 op 前插入 `pto.alloc_tile(no addr)`,并重写原 op,使其显式携带 tmp。 @@ -72,8 +73,6 @@ bool requireExplicitAtLevel3; 对自动生成的 tmp: - 使用 tile-native `pto.alloc_tile(no addr)`。 -- 不创建 `memref.alloc`。 -- 不创建 `pto.pointer_cast` / `pto.bind_tile`。 - 不设置 `addr`,由 memplan 统一规划。 - 尽量使用静态 full-valid shape,即 `v_row/v_col` 与 `rows/cols` 一致,不额外携带 `valid_row` / `valid_col` operand。 - 定义位置必须支配目标 op。 @@ -140,7 +139,7 @@ level3 + target_op(no tmp) => pass/verifier 报错 引入 `pto-materialize-implicit-tmp` 后,level1/level2 的目标 op 在 EmitC 前都会携带 tmp,因此会自然走带 tmp 的 overload。 -不建议在 PTOToEmitC 中补 tmp,原因: +不建议在 `EmitPTOManual`(PTO -> EmitC lowering)中补 tmp,原因: - EmitC 阶段已经错过 memplan。 - 临时生成 tmp 无法获得 local addr。 @@ -158,7 +157,7 @@ level3 + target_op(no tmp) => pass/verifier 报错 Optional:$tmp ``` -PTOToEmitC 也已经根据 `op.getTmp()` 选择带 tmp 或不带 tmp 的 C++ 调用。因此 TCI 改造不需要改 `pto.tci` 的 IR 语法,关键是保证进入 EmitC 前缺省 tmp 已经被显式 materialize。 +`EmitPTOManual`(PTO -> EmitC lowering)也已经根据 `op.getTmp()` 选择带 tmp 或不带 tmp 的 C++ 调用。因此 TCI 改造不需要改 `pto.tci` 的 IR 语法,关键是保证进入 EmitC 前缺省 tmp 已经被显式 materialize。 TCI 后端 C++ 接口存在两类 overload: @@ -326,9 +325,10 @@ test/lit/pto/tci_implicit_tmp_materialization.pto CHECK: pto.alloc_tile CHECK-SAME: dtype=f32 CHECK: pto.tci ins(%{{.*}}, %{{.*}} -CHECK-NOT: memref.alloc ``` +可选地检查 IR 不打印 `operandSegmentSizes`(可选 tmp 的自定义 assembly printer 已 elide 该属性,保证 round-trip 稳定)。 + 可以额外检查自动生成 shape:b32 dst 为 `f32 1x192`,b16 dst 为 `f32 1x448`。同时应覆盖 A5 下不自动生成 tmp。 #### lit:memplan 回写 addr @@ -527,7 +527,6 @@ forbidAlias(tmp, dst) ```text CHECK: pto.alloc_tile CHECK: pto.trowexpandadd ins(%{{.*}}, %{{.*}}, %{{.*}} -CHECK-NOT: memref.alloc ``` 同时检查 A5 下不自动生成 tmp。 @@ -584,11 +583,12 @@ TSORT32, TQUANT | 分类 | Op / 模式 | 设计结论 | | --- | --- | --- | -| A2/A3 使用 tmp,A5 接受但不使用 | `TCOLARGMAX`、`TCOLARGMIN`、`TROWARGMAX`、`TROWARGMIN`、`TGATHER`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS`、`TQUANT`、`TADDDEQRELU` | level1/2 在 A2/A3 生成真实 scratch;若 A5 C++ 签名仍要求 tmp,则生成不带 MemoryEffects 的 ABI placeholder。 | +| A2/A3 使用 tmp,A5 接受但不使用 | `TCOLARGMAX`、`TCOLARGMIN`、`TROWARGMAX`、`TROWARGMIN`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS`、`TQUANT`、`TADDDEQRELU` | level1/2 在 A2/A3 生成真实 scratch;若 A5 C++ 签名仍要求 tmp,则生成不带 MemoryEffects 的 ABI placeholder。 | | A2/A3 和 A5 都可能使用 tmp | `TCOLSUM(isBinary=true)`、`TSORT32` 非 32 对齐尾部、`TMRGSORT` 多列表归并 format2 | tmp 使用由 op 模式决定,不能只按 arch 判断。 | | 条件性 tmp,不应无条件 materialize | `TTRANS`、`TCVT`、`TPOW`、`TPOWS`、`TRSQRT`、`TMRGSORT`、`TSORT32` | 需要先判断精度、dtype、layout、format 或尾部条件。 | | 已从 mandatory tmp 改为 optional tmp | `TTRANS`、`TXOR`、`TXORS`、`TPRELU`、`TROWPROD`、`TROWSUM`、`TROWMAX`、`TROWMIN`、`TROWARGMAX`、`TROWARGMIN`、`TCOLARGMAX`、`TCOLARGMIN`、`TSEL`、`TSELS`、`TREM`、`TREMS` | ODS、parse/print、verifier、MemoryEffects、materialize 和 lowering 已接入。 | -| 当前 PTOAS IR 已有 optional tmp | `TGATHER`、`TCOLSUM`、`TRSQRT`、`TPOW`、`TPOWS`、`TSORT32`、`TQUANT` | 可直接纳入 `pto-materialize-implicit-tmp` 的后续实现。 | +| 当前 PTOAS IR 已有 optional tmp | `TCOLSUM`、`TRSQRT`、`TPOW`、`TPOWS`、`TSORT32`、`TQUANT` | 可直接纳入 `pto-materialize-implicit-tmp` 的后续实现。 | +| 已有 optional tmp 但不纳入 implicit-tmp materialize | `TGATHER` | main 分支要求 A2/A3 显式 tmp(verifier 拒绝省略),A5 index-form 设计为无 tmp;`replaceTGatherWithTmp` 实现保留但当前不在 dispatch 中启用。 | | 当前 PTOAS IR 暂无对应 op | `TADDDEQRELU` | 需先完成 PTOAS IR 接入;`TCVT` 已新增 optional tmp operand。 | ### 通用规则 @@ -766,11 +766,13 @@ TGATHER, TTRANS, TCVT TGATHER: -- 当前 PTOAS IR 已支持 optional tmp。 +- 当前 PTOAS IR 已支持 optional tmp,但 **tgather 不纳入 implicit-tmp materialize 范围**。 +- main 分支(PR #1080 "Add TGATHER indices and mask")对 tgather 的 tmp 契约更严:A2/A3 index-form 和所有 compare-form 都要求显式 `tmp`(verifier 报 `index-form tgather expects both indices and tmp` / `compare-form tgather expects dst, cdst, kValue, and tmp`);A5 index-form 设计为不带 tmp(emit `TGATHER(src, indices, dst)` 三参数)。 +- 因此 tgather 省略 tmp 时不由 `pto-materialize-implicit-tmp` 自动补齐,而是由 verifier 直接拒绝(A2/A3)或允许无 tmp(A5 index-form)。 - index form:A2/A3 C++ API 需要 tmp;tmp dtype 与 indices dtype 一致,shape 覆盖 indices;A5 不使用 tmp。 - compare form:A2/A3 tmp 是合并暂存缓冲区,包含 `cmpsTmp`、`indexTmp`、`cvtTmp` 三个区域;最小字节数按 PTO-ISA 文档公式计算;A5 不使用 tmp。 -- mask form 不使用 tmp,不应自动补 tmp。 -- A2/A3 level3 下 index / compare form 缺省 tmp 报错。 +- mask form 不使用 tmp。 +- A2/A3 index / compare form 必须显式提供 tmp;A5 index-form 不带 tmp。 TTRANS: From 1c2bcf5d8ddf9a00b565dec6e7bad4b6b8e73435 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 00:09:44 +0800 Subject: [PATCH 08/17] Fix tmp-aware sample bindings and rowexpand level3 --- ...oas-implicit-tmp-materialization-design.md | 2 +- .../Transforms/PTOMaterializeImplicitTmp.cpp | 8 +++-- .../pto/rowexpand_level3_no_tmp_preserved.pto | 29 +++++++++++++++++++ test/samples/Complex/mix_kernel.py | 2 +- test/samples/Rowmax/rowmax.py | 2 +- test/samples/Rowmin/rowmin.py | 2 +- test/samples/Rowprod/rowprod.py | 2 +- test/samples/Rowsum/rowsum.py | 2 +- test/samples/Sel/sel.py | 2 +- test/samples/Sels/sels.py | 4 +-- test/samples/Trans/trans.py | 2 +- test/samples/Xor/xor.py | 4 +-- test/samples/Xors/xors.py | 2 +- 13 files changed, 48 insertions(+), 15 deletions(-) create mode 100644 test/lit/pto/rowexpand_level3_no_tmp_preserved.pto diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md index c9c3c9cdec..aaad9a2276 100644 --- a/docs/designs/ptoas-implicit-tmp-materialization-design.md +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -443,7 +443,7 @@ A5: - 如果已经有 `tmp`,pass 不修改,但 verifier 需要保证它只用于合法模式。 - A5 如果没有 `tmp`,pass 不修改。 - A2/A3 如果没有 `tmp`,且 op 是模式 1、当前 build level 会运行 memplan,则自动补 tmp。 -- A2/A3 如果没有 `tmp`,op 是模式 1、但当前 level3 会跳过 memplan,则报错,要求用户显式提供带地址的 tmp。 +- A2/A3 如果没有 `tmp`,op 是模式 1、但当前 level3 会跳过 memplan,则保留 no-tmp overload,由后端使用内部 8KB `TMP_UB_OFFSET`,避免 pass 生成无地址 tmp。 - 模式 2 不需要 tmp;pass 不应自动补 tmp,也不应强制改成带 tmp overload。 自动补 tmp 的 canonical shape 建议采用形状无关上界: diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index d35c7147e3..138839984e 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -235,8 +235,12 @@ static LogicalResult materializeTRowExpandTmp(OpTy op, bool requireExplicitTmp, return success(); if (requireExplicitTmp) { - return op.emitOpError( - "requires explicit tmp for A2/A3 row-expand mode 1 when PlanMemory is skipped"); + // Row-expand is the one A2/A3 tmp-aware family whose no-tmp overload is + // still a valid backend contract: mode 1 falls back to pto-isa's internal + // 8KB TMP_UB_OFFSET scratch area, while mode 2 does not need tmp. Level3 + // inputs are already memory-planned by the frontend, so preserve their + // no-tmp form instead of creating an unaddressed alloc_tile. + return success(); } auto dstTy = dyn_cast(op.getDst().getType()); diff --git a/test/lit/pto/rowexpand_level3_no_tmp_preserved.pto b/test/lit/pto/rowexpand_level3_no_tmp_preserved.pto new file mode 100644 index 0000000000..ef1f6fa6ce --- /dev/null +++ b/test/lit/pto/rowexpand_level3_no_tmp_preserved.pto @@ -0,0 +1,29 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --pto-level=level3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module { + func.func @rowexpand_level3_no_tmp_preserved() { + %addr0 = arith.constant 0 : i64 + %addr1 = arith.constant 4096 : i64 + %addr2 = arith.constant 8192 : i64 + %src0 = pto.alloc_tile addr = %addr0 : !pto.tile_buf + %src1 = pto.alloc_tile addr = %addr1 : !pto.tile_buf + %dst = pto.alloc_tile addr = %addr2 : !pto.tile_buf + pto.trowexpandsub ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + pto.trowexpanddiv ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: func.func @rowexpand_level3_no_tmp_preserved +// CHECK: pto.trowexpandsub ins(%{{.*}}, %{{.*}} : +// CHECK-NOT: pto.trowexpandsub ins(%{{.*}}, %{{.*}}, %{{.*}} : +// CHECK: pto.trowexpanddiv ins(%{{.*}}, %{{.*}} : +// CHECK-NOT: pto.trowexpanddiv ins(%{{.*}}, %{{.*}}, %{{.*}} : diff --git a/test/samples/Complex/mix_kernel.py b/test/samples/Complex/mix_kernel.py index a7f4e8f6ba..031d0f2a07 100644 --- a/test/samples/Complex/mix_kernel.py +++ b/test/samples/Complex/mix_kernel.py @@ -132,7 +132,7 @@ def build(M=32, N=32, K=32, TM=32, TN=32, TK=32): pto.TLoadOp(None, sv_out, ubTile) # pto.trowmax ins(%src, %tmp) outs(%dst) - pto.TRowMaxOp(ubTile, ubTmpTile, ubReduceTile) + pto.TRowMaxOp(ubTile, ubReduceTile, tmp=ubTmpTile) pto.TStoreOp(None, ubReduceTile, sv_reduce) func.ReturnOp([]) diff --git a/test/samples/Rowmax/rowmax.py b/test/samples/Rowmax/rowmax.py index 30819eca8a..a63f8b05b5 100644 --- a/test/samples/Rowmax/rowmax.py +++ b/test/samples/Rowmax/rowmax.py @@ -67,7 +67,7 @@ def build(): pto.TLoadOp(None, sv0, tb0) # pto.trowmax ins(%src, %tmp) outs(%dst) - pto.TRowMaxOp(tb0, tb_tmp, tb1) + pto.TRowMaxOp(tb0, tb1, tmp=tb_tmp) # %8 = subview on output tensor_view sv1 = pto.PartitionViewOp(tile_view_32x1, tv1, offsets=[c0, c0], sizes=[c32, c1]).result diff --git a/test/samples/Rowmin/rowmin.py b/test/samples/Rowmin/rowmin.py index a51efdf81c..5840666982 100644 --- a/test/samples/Rowmin/rowmin.py +++ b/test/samples/Rowmin/rowmin.py @@ -60,7 +60,7 @@ def build(): tb1 = pto.AllocTileOp(tile_buf_32x1).result pto.TLoadOp(None, sv0, tb0) - pto.TRowMinOp(tb0, tb_tmp, tb1) + pto.TRowMinOp(tb0, tb1, tmp=tb_tmp) sv1 = pto.PartitionViewOp(tile_view_32x1, tv1, offsets=[c0, c0], sizes=[c32, c1]).result pto.TStoreOp(None, tb1, sv1) diff --git a/test/samples/Rowprod/rowprod.py b/test/samples/Rowprod/rowprod.py index 15e7c55734..c10ebe64e6 100644 --- a/test/samples/Rowprod/rowprod.py +++ b/test/samples/Rowprod/rowprod.py @@ -57,7 +57,7 @@ def build(): tb1 = pto.AllocTileOp(tile_buf_32x1).result pto.TLoadOp(None, sv0, tb0) - pto.TRowProdOp(tb0, tb_tmp, tb1) + pto.TRowProdOp(tb0, tb1, tmp=tb_tmp) sv1 = pto.PartitionViewOp(tile_view_32x1, tv1, offsets=[c0, c0], sizes=[c32, c1]).result pto.TStoreOp(None, tb1, sv1) diff --git a/test/samples/Rowsum/rowsum.py b/test/samples/Rowsum/rowsum.py index d7c5842f82..4a0e10b724 100644 --- a/test/samples/Rowsum/rowsum.py +++ b/test/samples/Rowsum/rowsum.py @@ -67,7 +67,7 @@ def build(): pto.TLoadOp(None, sv0, tb0) # result=None, valid_dims=[] # pto.trowsum ins(%src, %tmp) outs(%dst) - pto.TRowSumOp(tb0, tb_tmp, tb1) + pto.TRowSumOp(tb0, tb1, tmp=tb_tmp) # %8 = subview on output tensor_view sv1 = pto.PartitionViewOp(tile_view_32x1, tv1, offsets=[c0, c0], sizes=[c32, c1]).result diff --git a/test/samples/Sel/sel.py b/test/samples/Sel/sel.py index f4ae49b5ac..2d2a04dcdb 100644 --- a/test/samples/Sel/sel.py +++ b/test/samples/Sel/sel.py @@ -77,7 +77,7 @@ def build(): pto.TLoadOp(None, sv2, tb2) # result=None # pto.tsel ins(%mask,%src0,%src1,%tmp) outs(%dst) - pto.TSelOp(tb0, tb1, tb2, tb_tmp, tb3) + pto.TSelOp(tb0, tb1, tb2, tb3, tmp=tb_tmp) # %8 = subview on output tensor_view sv3 = pto.PartitionViewOp(tile_view_f32, tv3, offsets=[c0, c0], sizes=[c32, c32]).result diff --git a/test/samples/Sels/sels.py b/test/samples/Sels/sels.py index 1e10484e6b..f1e10bf5f5 100644 --- a/test/samples/Sels/sels.py +++ b/test/samples/Sels/sels.py @@ -65,8 +65,8 @@ def build(): pto.TLoadOp(None, sv0, tb0) # result=None pto.TLoadOp(None, sv1, tb1) # result=None - # TSELS(mask=tb0, src=tb1, tmp=tb2, scalar=c64) - pto.TSelSOp(tb0, tb1, tb2, c64, tb2) + # TSELS(mask=tb0, src=tb1, dst=tb2, tmp=tb2, scalar=c64) + pto.TSelSOp(tb0, tb1, c64, tb2, tmp=tb2) # %8 = subview on output tensor_view sv2 = pto.PartitionViewOp(tile_view_32, tv2, offsets=[c0, c0], sizes=[c32, c32]).result diff --git a/test/samples/Trans/trans.py b/test/samples/Trans/trans.py index 9b17b31b1e..e948b2cd5c 100644 --- a/test/samples/Trans/trans.py +++ b/test/samples/Trans/trans.py @@ -63,7 +63,7 @@ def build(): pto.TLoadOp(None, sv0, tb_src) # transpose: ttrans ins(%src, %tmp) outs(%dst) - pto.TTransOp(tb_src, tb_tmp, tb_dst) + pto.TTransOp(tb_src, tb_dst, tmp=tb_tmp) # output subview sv1 = pto.PartitionViewOp(tile_view_32, tv1, offsets=[c0, c0], sizes=[c32, c32]).result diff --git a/test/samples/Xor/xor.py b/test/samples/Xor/xor.py index 86acc61a46..2a8dd54b83 100644 --- a/test/samples/Xor/xor.py +++ b/test/samples/Xor/xor.py @@ -64,8 +64,8 @@ def build(): pto.TLoadOp(None, sv_src0, tb_src0) # result=None pto.TLoadOp(None, sv_src1, tb_src1) - pto.TXorOp(tb_src0, tb_src1, tb_tmp, tb_dst) - pto.TXorOp(tb_src0, tb_src1, tb_dst, tb_dst) + pto.TXorOp(tb_src0, tb_src1, tb_dst, tmp=tb_tmp) + pto.TXorOp(tb_src0, tb_src1, tb_dst, tmp=tb_dst) # output subview sv_dst = pto.PartitionViewOp(tile_view_32, tv_dst, offsets=[c0, c0], sizes=[c32, c32]).result diff --git a/test/samples/Xors/xors.py b/test/samples/Xors/xors.py index 3ef70a27ea..43bb176a3c 100644 --- a/test/samples/Xors/xors.py +++ b/test/samples/Xors/xors.py @@ -61,7 +61,7 @@ def build(): pto.TLoadOp(None, sv_src, tb_src) # result=None - pto.TXorSOp(tb_src, scale, tb_tmp, tb_dst) + pto.TXorSOp(tb_src, scale, tb_dst, tmp=tb_tmp) # output subview sv_dst = pto.PartitionViewOp(tile_view_32, tv_dst, offsets=[c0, c0], sizes=[c32, c32]).result From e2622b2db8cd20bad483499cb2dac299835e037d Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 08:52:59 +0800 Subject: [PATCH 09/17] Update rowexpand level3 tmp lit expectation --- ...lid.pto => trowexpandadd_level3_no_tmp_preserved.pto} | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) rename test/lit/pto/{trowexpand_implicit_tmp_level3_invalid.pto => trowexpandadd_level3_no_tmp_preserved.pto} (84%) diff --git a/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto b/test/lit/pto/trowexpandadd_level3_no_tmp_preserved.pto similarity index 84% rename from test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto rename to test/lit/pto/trowexpandadd_level3_no_tmp_preserved.pto index 34f9d32a7c..5c53792ae5 100644 --- a/test/lit/pto/trowexpand_implicit_tmp_level3_invalid.pto +++ b/test/lit/pto/trowexpandadd_level3_no_tmp_preserved.pto @@ -6,19 +6,22 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: not ptoas --pto-arch=a3 --pto-level=level3 %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch=a3 --pto-level=level3 --emit-pto-ir %s 2>&1 | FileCheck %s module { - func.func @trowexpand_mode1_level3_requires_tmp() { + func.func @trowexpand_mode1_level3_preserves_no_tmp() { %addr0 = arith.constant 0 : i64 %addr1 = arith.constant 1024 : i64 %addr2 = arith.constant 2048 : i64 %src0 = pto.alloc_tile addr = %addr0 : !pto.tile_buf %src1 = pto.alloc_tile addr = %addr1 : !pto.tile_buf %dst = pto.alloc_tile addr = %addr2 : !pto.tile_buf - // CHECK: error: 'pto.trowexpandadd' op requires explicit tmp for A2/A3 row-expand mode 1 when PlanMemory is skipped pto.trowexpandadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) return } } + +// CHECK-LABEL: func.func @trowexpand_mode1_level3_preserves_no_tmp +// CHECK: pto.trowexpandadd ins(%{{.*}}, %{{.*}} : +// CHECK-NOT: pto.trowexpandadd ins(%{{.*}}, %{{.*}}, %{{.*}} : From 64eeb54cf4f41f9a365b66660d6278d9e98da297 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 10:19:14 +0800 Subject: [PATCH 10/17] Fix tmrgsort smoke tmp tile size --- .../npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto index ed90162b9c..2131158a94 100644 --- a/test/tilelang_st/npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto +++ b/test/tilelang_st/npu/a5/src/st/smoke/testcase/tmrgsort/tmrgsort.pto @@ -107,7 +107,7 @@ func.func @TMRGSORT_f16_topk_1280_512(%src_ptr: !pto.ptr, %dst_ptr: !pto.pt %block1_tile = pto.alloc_tile : !pto.tile_buf %merge_tmp_tile = pto.alloc_tile - : !pto.tile_buf + : !pto.tile_buf %merge_dst_tile = pto.alloc_tile : !pto.tile_buf %ex_vec = arith.constant dense<0> : vector<4xi16> @@ -159,7 +159,7 @@ func.func @TMRGSORT_f16_topk_1280_512(%src_ptr: !pto.ptr, %dst_ptr: !pto.pt pto.tmrgsort ins(%block0_tile, %block1_tile, %merge_tmp_tile {exhausted = false} : !pto.tile_buf, !pto.tile_buf, - !pto.tile_buf) + !pto.tile_buf) outs(%merge_dst_tile, %ex_vec : !pto.tile_buf, vector<4xi16>) From 81bdb358967b379c602dc64bda7036da807ccb23 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 11:32:19 +0800 Subject: [PATCH 11/17] Fix PTODSL ttrans tmp builder call --- ptodsl/ptodsl/_ops.py | 4 ++-- ptodsl/tests/test_vector_cube_ops.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 5ec91c9a66..bf4110d623 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -3136,10 +3136,10 @@ def tmov(src, dst, *, mode=None): def ttrans(src, tmp, dst): """``pto.ttrans ins(src, tmp) outs(dst)`` – tile transpose (DPS).""" - _pto.ttrans( + _pto.TTransOp( unwrap_surface_value(src), - unwrap_surface_value(tmp), unwrap_surface_value(dst), + tmp=unwrap_surface_value(tmp), ) diff --git a/ptodsl/tests/test_vector_cube_ops.py b/ptodsl/tests/test_vector_cube_ops.py index ab134cb99e..acc4095989 100644 --- a/ptodsl/tests/test_vector_cube_ops.py +++ b/ptodsl/tests/test_vector_cube_ops.py @@ -764,6 +764,17 @@ def test_tile_row_reductions_expose_optional_tmp_and_synthesize_one(self): getattr(pto.tile, name)(src, dst, tmp=tmp) low_level_op.assert_called_once_with(src, tmp, dst) + def test_tile_transpose_wrapper_uses_tmp_keyword_builder(self): + src = object() + tmp = object() + dst = object() + + with patch.object(_ops, "unwrap_surface_value", side_effect=_identity), \ + patch.object(_ops._pto, "TTransOp") as ttrans_op: + pto.tile.transpose(src, tmp, dst) + + ttrans_op.assert_called_once_with(src, dst, tmp=tmp) + def test_tile_sort_gather_wrappers_call_low_level_ops(self): src = object() idx = object() From 657a4ba996aa5d63a13646c84bfc52ed728bb8a1 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 6 Aug 2026 23:23:44 +0800 Subject: [PATCH 12/17] Fix implicit tmp contracts --- lib/PTO/IR/PTO.cpp | 180 +++++++++++- .../Transforms/PTOMaterializeImplicitTmp.cpp | 267 ++++++++---------- test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto | 52 ++++ test/lit/pto/implicit_tmp_arg_reductions.pto | 19 +- test/lit/pto/implicit_tmp_row_reductions.pto | 19 +- .../pto/implicit_tmp_xor_materialization.pto | 13 +- .../pto/plan_memory_inplace_forbid_alias.pto | 4 +- test/lit/pto/select_tile_native.pto | 4 +- test/lit/pto/tsel_bf16.pto | 4 +- test/lit/pto/tsel_tmp_contract_a3_invalid.pto | 24 ++ 10 files changed, 398 insertions(+), 188 deletions(-) create mode 100644 test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto create mode 100644 test/lit/pto/tsel_tmp_contract_a3_invalid.pto diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 202fefe1e3..7c1f269e72 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -122,6 +122,9 @@ static bool isKnownZeroOrUnitExtent(int64_t value); static bool isByteIntegerType(Type ty); static LogicalResult verifyTileBufCommon(Operation *op, Type ty, StringRef name, bool allowLowPrecision = false); +static LogicalResult verifyTmpCapacityAtLeast(Operation *op, Type tmpTy, + uint64_t requiredBytes, + StringRef tmpName = "tmp"); namespace { struct PTOInlinerInterface : public DialectInlinerInterface { @@ -2144,11 +2147,17 @@ static LogicalResult verifyTRowReductionWithTmpCommon(Operation *op, Type srcTy, return failure(); if (getElemTy(srcTy) != getElemTy(dstTy)) return op->emitOpError("expects src and dst to have the same element type"); + if (getTargetArch(op) != PTOArch::A5 && + getElemTy(srcTy) != getElemTy(tmpTy)) + return op->emitOpError("expects A2/A3 tmp to have the same element type as src and dst"); if (failed(verifyRowReductionValidRegion(op, srcTy, dstTy, /*allowEmptyMarker=*/true))) return failure(); if (!isSupportedRowReductionElemType(getElemTy(srcTy))) return op->emitOpError(elemTypeError); + if (getTargetArch(op) != PTOArch::A5 && + failed(verifyTmpCapacityAtLeast(op, tmpTy, 32))) + return failure(); return success(); } @@ -2197,7 +2206,7 @@ static LogicalResult verifyTColArgTmpA2A3(Operation *op, Type srcTy, return failure(); if (hasExactKnownValidShape(srcTy, tmpTy)) - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); auto srcValid = getValidShapeVec(srcTy); auto tmpValid = getValidShapeVec(tmpTy); @@ -2214,7 +2223,7 @@ static LogicalResult verifyTColArgTmpA2A3(Operation *op, Type srcTy, << "expects A2/A3 tmp valid_shape[1] to be at least " << *minStride << " for src valid_shape[1] = " << srcValid[1]; } - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); } static LogicalResult verifyTColArgReductionOpA2A3(Operation *op, Type srcTy, @@ -2307,7 +2316,7 @@ static LogicalResult verifyTRowArgTmpA2A3(Operation *op, Type srcTy, return failure(); if (hasExactKnownValidShape(srcTy, tmpTy)) - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); auto srcShape = getShapeVec(srcTy); auto tmpShape = getShapeVec(tmpTy); @@ -2335,7 +2344,7 @@ static LogicalResult verifyTRowArgTmpA2A3(Operation *op, Type srcTy, return op->emitOpError() << "expects A2/A3 tmp DN layout to have valid_shape[0] >= " << (srcValid[0] * 2); - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); } if (!layout || *layout != pto::Layout::ND) @@ -2349,7 +2358,7 @@ static LogicalResult verifyTRowArgTmpA2A3(Operation *op, Type srcTy, if (tmpValid[1] != ShapedType::kDynamic && tmpValid[1] < 2) return op->emitOpError( "expects A2/A3 tmp valid_shape[1] to be at least 2 in the small-col ND path"); - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); } if (failed(verifyVecTileCommon(op, tmpTy, "tmp"))) @@ -2369,7 +2378,7 @@ static LogicalResult verifyTRowArgTmpA2A3(Operation *op, Type srcTy, << "expects A2/A3 tmp valid_shape[1] to be at least " << *minStride << " for src valid_shape[1] = " << srcValid[1]; } - return success(); + return verifyTmpCapacityAtLeast(op, tmpTy, 32); } static LogicalResult verifyTRowArgReductionOpA2A3(Operation *op, Type srcTy, @@ -4164,6 +4173,20 @@ static std::optional getStaticByteSize(Type ty) { return total; } +static LogicalResult verifyTmpCapacityAtLeast(Operation *op, Type tmpTy, + uint64_t requiredBytes, + StringRef tmpName) { + auto actualBytes = getStaticByteSize(tmpTy); + if (!actualBytes) + return op->emitOpError() + << "expects " << tmpName << " to have statically known byte capacity"; + if (*actualBytes < requiredBytes) + return op->emitOpError() + << "expects " << tmpName << " capacity to be at least " + << requiredBytes << " bytes, but got " << *actualBytes << " bytes"; + return success(); +} + static std::optional getPTOMemorySpaceEnum(Type ty) { if (auto ptr = dyn_cast(ty)) return ptr.getMemorySpace().getAddressSpace(); @@ -6376,6 +6399,19 @@ LogicalResult pto::TColSumOp::verify() { return emitOpError("expects src/tmp/dst element types to match"); if (failed(verifyTColSumTmpStride(*this, srcTy, tmpTy, getIsBinary()))) return failure(); + if (getIsBinary()) { + auto srcValid = getValidShapeVec(srcTy); + auto elemBytes = getElemByteSize(getElemTy(srcTy)); + if (srcValid.size() != 2 || srcValid[0] == ShapedType::kDynamic || + srcValid[1] == ShapedType::kDynamic || elemBytes == 0) + return emitOpError( + "expects static src valid_shape and element size to verify tcolsum tmp"); + uint64_t requiredBytes = + static_cast(ceilDivInt64(srcValid[0], 2)) * + static_cast(srcValid[1]) * elemBytes; + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, requiredBytes))) + return failure(); + } } if (getElemTy(srcTy) != getElemTy(dstTy)) return emitOpError("expects src/dst element types to match"); @@ -6405,6 +6441,19 @@ LogicalResult pto::TColSumOp::verify() { return emitOpError("expects src/tmp/dst element types to match"); if (failed(verifyTColSumTmpStride(*this, srcTy, tmpTy, getIsBinary()))) return failure(); + if (getIsBinary()) { + auto srcValid = getValidShapeVec(srcTy); + auto elemBytes = getElemByteSize(getElemTy(srcTy)); + if (srcValid.size() != 2 || srcValid[0] == ShapedType::kDynamic || + srcValid[1] == ShapedType::kDynamic || elemBytes == 0) + return emitOpError( + "expects static src valid_shape and element size to verify tcolsum tmp"); + uint64_t requiredBytes = + static_cast(ceilDivInt64(srcValid[0], 2)) * + static_cast(srcValid[1]) * elemBytes; + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, requiredBytes))) + return failure(); + } } if (getElemTy(srcTy) != getElemTy(dstTy)) return emitOpError("expects src/dst element types to match"); @@ -10473,6 +10522,18 @@ mlir::LogicalResult mlir::pto::TPReluOp::verify() { << "expects A2/A3 tmp valid_shape[1] to be at least ceil(dst valid_shape[1] / 8) (" << packedMaskCols << ")"; } + if (dstValid[0] == ShapedType::kDynamic || + dstValid[1] == ShapedType::kDynamic) + return emitOpError( + "expects A2/A3 tprelu dst valid_shape to be static when tmp is provided"); + int64_t packedCols = std::max( + 32, llvm::divideCeil(llvm::divideCeil(dstValid[1], int64_t{8}), + int64_t{32}) * + 32); + if (failed(verifyTmpCapacityAtLeast( + *this, tt, static_cast(dstValid[0] + 1) * + static_cast(packedCols)))) + return failure(); if (auto arch = getVerifierArchName(getOperation()); arch && arch->equals_insensitive("a3")) { if (getSrc0() == getSrc1() || getSrc0() == getTmp() || getSrc0() == getDst() || @@ -10740,7 +10801,11 @@ mlir::LogicalResult mlir::pto::TQuantOp::verify() { return emitOpError() << "expects A2/A3 tmp to have the same shape as src"; if (failed(verifyTileBufSameValidShape(*this, srcTy, tmpTy, "src", "tmp"))) return failure(); - return success(); + auto requiredBytes = getStaticByteSize(srcTy); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 tquant src shape to be static when tmp is provided"); + return verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes); }; if (getTmp() && failed(verifyA2A3Tmp(getTmp().getType()))) return failure(); @@ -11040,6 +11105,16 @@ mlir::LogicalResult mlir::pto::TRemOp::verify() { if (dstValid[1] != ShapedType::kDynamic && tmpValid[1] != ShapedType::kDynamic && tmpValid[1] < dstValid[1]) return emitOpError("expects A2/A3 tmp valid columns to cover dst valid columns"); + auto dstShape = getShapeVec(dstTy); + auto elemBytes = getElemByteSize(elem); + if (dstShape.size() != 2 || dstShape[1] == ShapedType::kDynamic || + elemBytes == 0) + return emitOpError( + "expects A2/A3 trem dst shape and element size to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast( + *this, tmpTy, static_cast(2) * + static_cast(dstShape[1]) * elemBytes))) + return failure(); if (!(elem.isInteger(32) || elem.isF32())) return emitOpError("expects A2/A3 trem element type to be i32/f32"); return success(); @@ -11110,6 +11185,15 @@ mlir::LogicalResult mlir::pto::TRemSOp::verify() { if (dstValid[1] != ShapedType::kDynamic && tmpValid[1] != ShapedType::kDynamic && tmpValid[1] < dstValid[1]) return emitOpError("expects A2/A3 tmp valid columns to cover dst valid columns"); + auto dstShape = getShapeVec(td); + auto elemBytes = getElemByteSize(elem); + if (dstShape.size() != 2 || dstShape[1] == ShapedType::kDynamic || + elemBytes == 0) + return emitOpError( + "expects A2/A3 trems dst shape and element size to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast( + *this, tt, static_cast(dstShape[1]) * elemBytes))) + return failure(); if (!(elem.isInteger(32) || elem.isF32())) return emitOpError("expects A2/A3 trems element type to be i32/f32"); return success(); @@ -11218,6 +11302,14 @@ mlir::LogicalResult mlir::pto::TPowOp::verify() { return failure(); if (failed(verifyTPowTmpShape(getOperation(), tmpTy, dstTy))) return failure(); + if (getTargetArch(getOperation()) != PTOArch::A5) { + auto requiredBytes = getStaticByteSize(dstTy); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 tpow dst shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes))) + return failure(); + } } return success(); } @@ -11276,6 +11368,14 @@ mlir::LogicalResult mlir::pto::TPowSOp::verify() { return failure(); if (failed(verifyTPowTmpShape(getOperation(), tmpTy, dstTy))) return failure(); + if (getTargetArch(getOperation()) != PTOArch::A5) { + auto requiredBytes = getStaticByteSize(dstTy); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 tpows dst shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes))) + return failure(); + } } return success(); } @@ -12934,6 +13034,9 @@ mlir::LogicalResult mlir::pto::TSelOp::verify() { failed(verifyTileBufCommon(*this, t1, "src1")) || failed(verifyTileBufCommon(*this, td, "dst"))) return failure(); + if (getTmp() && + failed(verifyVecTileCommon(*this, getTmp().getType(), "tmp"))) + return failure(); Type srcElem = getElemTy(t0); Type src1Elem = getElemTy(t1); @@ -12967,6 +13070,18 @@ mlir::LogicalResult mlir::pto::TSelOp::verify() { if (!ok) return emitOpError( "expects A2/A3 tsel src0, src1, and dst element type to be i16/i32/f16/bf16/f32"); + if (getTmp()) { + Type tmpTy = getTmp().getType(); + auto tmpElem = dyn_cast(getElemTy(tmpTy)); + if (!tmpElem || tmpElem.getWidth() != 32) + return emitOpError("expects A2/A3 tsel tmp element type to be i32"); + unsigned elemBits = getPTOStorageElemBitWidth(elem); + if (elemBits != 16 && elemBits != 32) + return emitOpError("expects A2/A3 tsel data element type to be 16 or 32 bits"); + uint64_t minBytes = elemBits == 16 ? 64 : 32; + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, minBytes))) + return failure(); + } return success(); }; @@ -13025,6 +13140,22 @@ mlir::LogicalResult mlir::pto::TSelSOp::verify() { if (!isRowMajorTileBuf(tSrc) || !isRowMajorTileBuf(tDst)) return emitOpError("expects src and dst to use row-major layout"); Type elem = *elemOr; + if (getTmp()) { + Type tmpTy = getTmp().getType(); + if (getElemTy(tmpTy) != elem) + return emitOpError("expects A2/A3 tsels tmp to have the same element type as src and dst"); + if (!isRowMajorTileBuf(tmpTy)) + return emitOpError("expects A2/A3 tsels tmp to use row-major layout"); + auto srcShape = getShapeVec(tSrc); + if (srcShape.size() != 2 || srcShape[1] == ShapedType::kDynamic) + return emitOpError( + "expects A2/A3 tsels src shape to be static when tmp is provided"); + auto elemBytes = getElemByteSize(elem); + if (elemBytes == 0 || + failed(verifyTmpCapacityAtLeast( + *this, tmpTy, static_cast(srcShape[1]) * elemBytes))) + return failure(); + } bool ok = elem.isF16() || elem.isF32(); if (auto it = mlir::dyn_cast(elem)) ok = (it.getWidth() == 16 || it.getWidth() == 32); @@ -13104,6 +13235,15 @@ mlir::LogicalResult mlir::pto::TSort32Op::verify() { if (getTmp() && failed(verifyVecTileCommon(*this, getTmp().getType(), "tmp"))) return failure(); + if (getTmp() && getTargetArch(getOperation()) != PTOArch::A5) { + auto requiredBytes = getStaticByteSize(srcTy); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 tsort32 src shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, getTmp().getType(), + *requiredBytes))) + return failure(); + } auto srcElem = getElemTy(srcTy); auto dstElem = getElemTy(dstTy); @@ -13329,6 +13469,18 @@ mlir::LogicalResult mlir::pto::TTransOp::verify() { }; if (!isAllowedWidthType(srcElem)) return emitOpError() << "expects transpose element type to match the supported set for its width"; + if (tmpTy) { + uint64_t requiredBytes = 32; + if (ttransUsesTmp(srcTy, dstTy)) { + auto srcBytes = getStaticByteSize(srcTy); + if (!srcBytes) + return emitOpError( + "expects A2/A3 transpose src shape to be static when tmp is used"); + requiredBytes = *srcBytes; + } + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, requiredBytes))) + return failure(); + } return mlir::success(); }; auto verifyA5 = [&]() -> LogicalResult { @@ -13359,6 +13511,8 @@ mlir::LogicalResult mlir::pto::TTransOp::verify() { }; if (!isAllowedWidthType(srcElem)) return emitOpError() << "expects transpose element type to match the supported set for its width"; + if (tmpTy && failed(verifyTmpCapacityAtLeast(*this, tmpTy, 32))) + return failure(); auto checkAlignedMajor = [&](Type ty, StringRef name) -> LogicalResult { auto tb = mlir::dyn_cast(ty); if (!tb) @@ -13497,6 +13651,12 @@ mlir::LogicalResult mlir::pto::TXorOp::verify() { if (failed(verifyTileBufSameValidShape( *this, tmpTy, getDst().getType(), "tmp", "dst"))) return failure(); + auto requiredBytes = getStaticByteSize(getDst().getType()); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 txor dst shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes))) + return failure(); } auto it = mlir::dyn_cast(elem); if (!it || (it.getWidth() != 8 && it.getWidth() != 16 && @@ -13542,6 +13702,12 @@ mlir::LogicalResult mlir::pto::TXorSOp::verify() { "expects tmp to have the same element type as src and dst"); if (!isRowMajorTileBuf(tmpTy)) return emitOpError("expects tmp to use row-major layout"); + auto requiredBytes = getStaticByteSize(getDst().getType()); + if (!requiredBytes) + return emitOpError( + "expects A2/A3 txors dst shape to be static when tmp is provided"); + if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, *requiredBytes))) + return failure(); } auto it = mlir::dyn_cast(elem); if (!it || (it.getWidth() != 8 && it.getWidth() != 16)) diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index 138839984e..b832503501 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -19,6 +19,8 @@ #include "mlir/Pass/Pass.h" #include "llvm/ADT/TypeSwitch.h" +#include + using namespace mlir; namespace { @@ -125,6 +127,21 @@ static void copyAttrsExceptOperandSegments(Operation *from, OperationState &to) } } +static void rebuildWithOperands(Operation *op, ArrayRef operands, + std::optional> segments) { + OpBuilder builder(op); + OperationState state(op->getLoc(), op->getName()); + state.addOperands(operands); + if (op->hasTrait()) { + assert(segments && "AttrSizedOperandSegments op must supply segments"); + state.addAttribute("operandSegmentSizes", + builder.getDenseI32ArrayAttr(*segments)); + } + copyAttrsExceptOperandSegments(op, state); + builder.create(state); + op->erase(); +} + static bool validShapesCompatible(ArrayRef lhs, ArrayRef rhs) { if (lhs.size() != rhs.size()) @@ -205,20 +222,23 @@ static pto::TileBufType makeTRowExpandTmpType(MLIRContext *ctx, makeRowMajorNoneBoxConfig(ctx)); } +static FailureOr makeA5PlaceholderTmpType( + MLIRContext *ctx, Value like, Type elementType = {}) { + auto likeTy = dyn_cast(like.getType()); + if (!likeTy) + return failure(); + if (!elementType) + elementType = likeTy.getElementType(); + auto elemBytes = getElemBytes(elementType); + if (!elemBytes || *elemBytes <= 0) + return failure(); + int64_t cols = std::max(1, 32 / *elemBytes); + return makeVecTmpType(ctx, {1, cols}, elementType, {1, cols}); +} + static void replaceTRowExpandBinaryOpWithTmp(Operation *op, Value src0, Value src1, Value tmp, Value dst) { - OpBuilder builder(op); - OperationState state(op->getLoc(), op->getName()); - state.addOperands({src0, src1, tmp, dst}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1, 1})); - for (NamedAttribute attr : op->getAttrs()) { - if (attr.getName() == "operandSegmentSizes") - continue; - state.addAttribute(attr.getName(), attr.getValue()); - } - builder.create(state); - op->erase(); + rebuildWithOperands(op, {src0, src1, tmp, dst}, ArrayRef{1, 1, 1, 1}); } template @@ -283,11 +303,8 @@ static LogicalResult replaceTColSumWithTmp(pto::TColSumOp op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), *tmp, op.getDst()}); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), {op.getSrc(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1}); return success(); } @@ -309,19 +326,13 @@ static LogicalResult replaceTQuantWithTmp(pto::TQuantOp op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); SmallVector operands{op.getSrc(), op.getFp()}; if (op.getOffset()) operands.push_back(op.getOffset()); operands.push_back(*tmp); operands.push_back(op.getDst()); - state.addOperands(operands); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr( - {1, 1, op.getOffset() ? 1 : 0, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), operands, + ArrayRef{1, 1, op.getOffset() ? 1 : 0, 1, 1}); return success(); } @@ -347,11 +358,9 @@ static LogicalResult replaceTPowWithTmp(pto::TPowOp op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getBase(), op.getExp(), op.getDst(), *tmp}); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), + {op.getBase(), op.getExp(), op.getDst(), *tmp}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -372,68 +381,9 @@ static LogicalResult replaceTPowSWithTmp(pto::TPowSOp op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), op.getScalar(), op.getDst(), *tmp}); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); - return success(); -} - -[[maybe_unused]] static LogicalResult -replaceTGatherWithTmp(pto::TGatherOp op, bool requireExplicitTmp, - MLIRContext *ctx) { - if (op.getTmp() || op.hasMaskForm()) - return success(); - if (!op.hasIndexForm() && !op.hasCompareForm()) - return success(); - if (requireExplicitTmp) - return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); - - FailureOr tmpType = failure(); - if (op.hasIndexForm()) { - tmpType = makeSameShapeTmpType(ctx, op.getIndices()); - } else { - auto srcTy = dyn_cast(op.getSrc().getType()); - auto dstTy = dyn_cast(op.getDst().getType()); - if (!srcTy || !dstTy) - return op.emitOpError( - "expects tile_buf operands when materializing compare-form tgather tmp"); - auto srcShape = getShapeVec(op.getSrc().getType()); - if (srcShape.size() != 2 || hasDynamicDim(srcShape)) - return op.emitOpError( - "requires static src shape to materialize compare-form tgather tmp"); - int64_t bytes = srcShape[0] * srcShape[1] * 4 + srcShape[0] * 4; - tmpType = makeVecTmpType(ctx, {1, bytes}, IntegerType::get(ctx, 8), - {1, bytes}); - } - if (failed(tmpType)) - return op.emitOpError( - "requires static tile_buf indices/src to materialize implicit tgather tmp"); - - OpBuilder builder(op); - FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); - if (failed(tmp)) - return failure(); - - OperationState state(op.getLoc(), op->getName()); - SmallVector operands{op.getSrc(), op.getDst()}; - if (op.getCdst()) - operands.push_back(op.getCdst()); - if (op.getIndices()) - operands.push_back(op.getIndices()); - operands.push_back(*tmp); - if (op.getKValue()) - operands.push_back(op.getKValue()); - state.addOperands(operands); - state.addAttribute( - "operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, op.getCdst() ? 1 : 0, - op.getIndices() ? 1 : 0, 1, - op.getKValue() ? 1 : 0})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), + {op.getSrc(), op.getScalar(), op.getDst(), *tmp}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -459,13 +409,8 @@ static LogicalResult replaceTSort32WithTmp(pto::TSort32Op op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), op.getIdx(), *tmp, op.getDst()}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), {op.getSrc(), op.getIdx(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -475,10 +420,14 @@ static LogicalResult replaceRowReductionWithTmp(OpTy op, MLIRContext *ctx) { if (op.getTmp()) return success(); - if (requireExplicitTmp) + + bool isA5 = pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5; + if (requireExplicitTmp && !isA5) return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); - FailureOr tmpType = makeSameShapeTmpType(ctx, op.getSrc()); + FailureOr tmpType = + isA5 ? makeA5PlaceholderTmpType(ctx, op.getSrc()) + : makeSameShapeTmpType(ctx, op.getSrc()); if (failed(tmpType)) return op.emitOpError( "requires static tile_buf src to materialize implicit row-reduction tmp"); @@ -487,13 +436,8 @@ static LogicalResult replaceRowReductionWithTmp(OpTy op, if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), *tmp, op.getDst()}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), {op.getSrc(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1}); return success(); } @@ -502,9 +446,12 @@ static LogicalResult replaceTXorWithTmp(pto::TXorOp op, MLIRContext *ctx) { if (op.getTmp()) return success(); - if (requireExplicitTmp) + bool isA5 = pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5; + if (requireExplicitTmp && !isA5) return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); - FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + FailureOr tmpType = + isA5 ? makeA5PlaceholderTmpType(ctx, op.getDst()) + : makeSameShapeTmpType(ctx, op.getDst()); if (failed(tmpType)) return op.emitOpError( "requires static tile_buf dst to materialize implicit txor tmp"); @@ -512,13 +459,9 @@ static LogicalResult replaceTXorWithTmp(pto::TXorOp op, FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc0(), op.getSrc1(), *tmp, op.getDst()}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), + {op.getSrc0(), op.getSrc1(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -527,9 +470,12 @@ static LogicalResult replaceTXorSWithTmp(pto::TXorSOp op, MLIRContext *ctx) { if (op.getTmp()) return success(); - if (requireExplicitTmp) + bool isA5 = pto::getTargetArch(op.getOperation()) == pto::PTOArch::A5; + if (requireExplicitTmp && !isA5) return op.emitOpError("requires explicit tmp when PlanMemory is skipped"); - FailureOr tmpType = makeSameShapeTmpType(ctx, op.getDst()); + FailureOr tmpType = + isA5 ? makeA5PlaceholderTmpType(ctx, op.getDst()) + : makeSameShapeTmpType(ctx, op.getDst()); if (failed(tmpType)) return op.emitOpError( "requires static tile_buf dst to materialize implicit txors tmp"); @@ -537,13 +483,9 @@ static LogicalResult replaceTXorSWithTmp(pto::TXorSOp op, FailureOr tmp = createAllocTmp(builder, op.getLoc(), *tmpType); if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands({op.getSrc(), op.getScalar(), *tmp, op.getDst()}); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr({1, 1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + rebuildWithOperands(op.getOperation(), + {op.getSrc(), op.getScalar(), *tmp, op.getDst()}, + ArrayRef{1, 1, 1, 1}); return success(); } @@ -566,13 +508,7 @@ static LogicalResult replaceFixedDpsOpWithTmp( else finalOperands.push_back(*tmp); } - OperationState state(op->getLoc(), op->getName()); - state.addOperands(finalOperands); - state.addAttribute("operandSegmentSizes", - builder.getDenseI32ArrayAttr(operandSegments)); - copyAttrsExceptOperandSegments(op, state); - builder.create(state); - op->erase(); + rebuildWithOperands(op, finalOperands, operandSegments); (void)opName; return success(); } @@ -610,60 +546,83 @@ static LogicalResult materializeFixedMandatoryTmp(Operation *op, .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeTPReluTmpType(ctx, typedOp.getDst()); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType( + ctx, typedOp.getDst(), + IntegerType::get(ctx, 8)) + : makeTPReluTmpType(ctx, typedOp.getDst()); if (failed(type)) return typedOp.emitOpError( "requires static tile_buf dst to materialize implicit tprelu tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getSrc0(), typedOp.getSrc1(), Value(), typedOp.getDst()}, - *type, {1, 1, 1, 1}, requireExplicitTmp, "tprelu"); + *type, {1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, + "tprelu"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeRowsTmpType(ctx, typedOp.getDst(), 2); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType(ctx, typedOp.getDst()) + : makeRowsTmpType(ctx, typedOp.getDst(), 2); if (failed(type)) return typedOp.emitOpError( "requires static tile_buf dst to materialize implicit trem tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getSrc0(), typedOp.getSrc1(), Value(), typedOp.getDst()}, - *type, {1, 1, 1, 1}, requireExplicitTmp, "trem"); + *type, {1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, "trem"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeRowsTmpType(ctx, typedOp.getDst(), 1); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType(ctx, typedOp.getDst()) + : makeRowsTmpType(ctx, typedOp.getDst(), 1); if (failed(type)) return typedOp.emitOpError( "requires static tile_buf dst to materialize implicit trems tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getSrc(), typedOp.getScalar(), Value(), typedOp.getDst()}, - *type, {1, 1, 1, 1}, requireExplicitTmp, "trems"); + *type, {1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, "trems"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeVecTmpType(ctx, {1, 16}, IntegerType::get(ctx, 32), - {1, 16}); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType( + ctx, typedOp.getDst(), + IntegerType::get(ctx, 32)) + : makeVecTmpType(ctx, {1, 16}, + IntegerType::get(ctx, 32), {1, 16}); + if (failed(type)) + return typedOp.emitOpError( + "requires static tile_buf dst to materialize implicit tsel tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getMask(), typedOp.getSrc0(), typedOp.getSrc1(), Value(), typedOp.getDst()}, - type, {1, 1, 1, 1, 1}, requireExplicitTmp, "tsel"); + *type, {1, 1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, "tsel"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); - auto type = makeRowsTmpType(ctx, typedOp.getSrc(), 1); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; + auto type = isA5 ? makeA5PlaceholderTmpType(ctx, typedOp.getSrc()) + : makeRowsTmpType(ctx, typedOp.getSrc(), 1); if (failed(type)) return typedOp.emitOpError( "requires static tile_buf src to materialize implicit tsels tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getMask(), typedOp.getSrc(), Value(), typedOp.getScalar(), typedOp.getDst()}, - *type, {1, 1, 1, 1, 1}, requireExplicitTmp, "tsels"); + *type, {1, 1, 1, 1, 1}, isA5 ? false : requireExplicitTmp, "tsels"); }) .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) @@ -791,18 +750,14 @@ static LogicalResult materializeTMrgSortTmp(pto::TMrgSortOp op, FailureOr tmp = createAllocTmp(builder, op.getLoc(), tmpType); if (failed(tmp)) return failure(); - OperationState state(op.getLoc(), op->getName()); - state.addOperands(operands); - state.addOperands(op.getDsts()); - state.addOperands(*tmp); - state.addOperands(op.getExcuted()); - state.addAttribute( - "operandSegmentSizes", - builder.getDenseI32ArrayAttr( - {static_cast(op.getSrcs().size()), 0, 1, 1, 1})); - copyAttrsExceptOperandSegments(op.getOperation(), state); - builder.create(state); - op.erase(); + SmallVector finalOperands; + finalOperands.append(operands.begin(), operands.end()); + finalOperands.append(op.getDsts().begin(), op.getDsts().end()); + finalOperands.push_back(*tmp); + finalOperands.push_back(op.getExcuted()); + rebuildWithOperands(op.getOperation(), finalOperands, + ArrayRef{static_cast(op.getSrcs().size()), + 0, 1, 1, 1}); return success(); } diff --git a/test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto b/test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto new file mode 100644 index 0000000000..3b66956054 --- /dev/null +++ b/test/lit/pto/implicit_tmp_a5_skip_no_tmp.pto @@ -0,0 +1,52 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: sed -E 's/ addr = %[A-Za-z0-9_]+//g' %s > %t.level2.pto && ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %t.level2.pto 2>&1 | FileCheck %s --check-prefix=A5 +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 + +module { + func.func @a5_skip_no_tmp_overloads() attributes {pto.kernel_kind = #pto.kernel_kind} { + %c0_i32 = arith.constant 0 : i32 + %a0 = arith.constant 0 : i64 + %a256 = arith.constant 256 : i64 + %a512 = arith.constant 512 : i64 + %a768 = arith.constant 768 : i64 + %a1024 = arith.constant 1024 : i64 + %a1280 = arith.constant 1280 : i64 + %a1536 = arith.constant 1536 : i64 + %a1792 = arith.constant 1792 : i64 + %a2048 = arith.constant 2048 : i64 + %seq = pto.alloc_tile addr = %a0 : !pto.tile_buf + pto.tci ins(%c0_i32 : i32) + outs(%seq : !pto.tile_buf) + + %cvt_src = pto.alloc_tile addr = %a256 : !pto.tile_buf + %cvt_dst = pto.alloc_tile addr = %a512 : !pto.tile_buf + pto.tcvt ins(%cvt_src : !pto.tile_buf) + outs(%cvt_dst : !pto.tile_buf) + + %q_src = pto.alloc_tile addr = %a768 : !pto.tile_buf + %q_fp = pto.alloc_tile addr = %a1024 : !pto.tile_buf + %q_dst = pto.alloc_tile addr = %a1280 : !pto.tile_buf + pto.tquant ins(%q_src, %q_fp : !pto.tile_buf, !pto.tile_buf) + outs(%q_dst : !pto.tile_buf) {quant_type = #pto} + + %re_src0 = pto.alloc_tile addr = %a1536 : !pto.tile_buf + %re_src1 = pto.alloc_tile addr = %a1792 : !pto.tile_buf + %re_dst = pto.alloc_tile addr = %a2048 : !pto.tile_buf + pto.trowexpandadd ins(%re_src0, %re_src1 : !pto.tile_buf, !pto.tile_buf) + outs(%re_dst : !pto.tile_buf) + return + } +} + +// A5-LABEL: func.func @a5_skip_no_tmp_overloads +// A5: pto.tci ins(%{{.*}} : i32) outs( +// A5: pto.tcvt ins(%{{.*}} {{.*}} : !pto.tile_buf) outs(%{{.*}} : !pto.tile_buf) +// A5: pto.tquant ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowexpandadd ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_arg_reductions.pto b/test/lit/pto/implicit_tmp_arg_reductions.pto index 4f51f1da7c..b144227d33 100644 --- a/test/lit/pto/implicit_tmp_arg_reductions.pto +++ b/test/lit/pto/implicit_tmp_arg_reductions.pto @@ -6,8 +6,8 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR -// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 module { func.func @implicit_arg_reduction_tmps() { @@ -24,8 +24,13 @@ module { } } -// IR-LABEL: func.func @implicit_arg_reduction_tmps -// IR: pto.tcolargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.tcolargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3-LABEL: func.func @implicit_arg_reduction_tmps +// A3: pto.tcolargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.tcolargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5-LABEL: func.func @implicit_arg_reduction_tmps +// A5: pto.tcolargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.tcolargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowargmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowargmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_row_reductions.pto b/test/lit/pto/implicit_tmp_row_reductions.pto index cfd3ade699..15c62e65b4 100644 --- a/test/lit/pto/implicit_tmp_row_reductions.pto +++ b/test/lit/pto/implicit_tmp_row_reductions.pto @@ -6,8 +6,8 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR -// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 module { func.func @implicit_row_reduction_tmps() { @@ -24,8 +24,13 @@ module { } } -// IR-LABEL: func.func @implicit_row_reduction_tmps -// IR: pto.trowmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowsum ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) -// IR: pto.trowprod ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3-LABEL: func.func @implicit_row_reduction_tmps +// A3: pto.trowmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowsum ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A3: pto.trowprod ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5-LABEL: func.func @implicit_row_reduction_tmps +// A5: pto.trowmax ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowmin ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowsum ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) +// A5: pto.trowprod ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf) diff --git a/test/lit/pto/implicit_tmp_xor_materialization.pto b/test/lit/pto/implicit_tmp_xor_materialization.pto index 172d428c80..d5e27a6dc4 100644 --- a/test/lit/pto/implicit_tmp_xor_materialization.pto +++ b/test/lit/pto/implicit_tmp_xor_materialization.pto @@ -6,8 +6,8 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR -// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 module { func.func @implicit_xor_tmps() { @@ -22,6 +22,9 @@ module { } } -// IR-LABEL: func.func @implicit_xor_tmps -// IR: pto.txor ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) -// IR: pto.txors ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, i16, !pto.tile_buf) +// A3-LABEL: func.func @implicit_xor_tmps +// A3: pto.txor ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) +// A3: pto.txors ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, i16, !pto.tile_buf) +// A5-LABEL: func.func @implicit_xor_tmps +// A5: pto.txor ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) +// A5: pto.txors ins(%{{.*}}, %{{.*}}, %{{.*}} : !pto.tile_buf, i16, !pto.tile_buf) diff --git a/test/lit/pto/plan_memory_inplace_forbid_alias.pto b/test/lit/pto/plan_memory_inplace_forbid_alias.pto index 92062d448e..8bfed61303 100644 --- a/test/lit/pto/plan_memory_inplace_forbid_alias.pto +++ b/test/lit/pto/plan_memory_inplace_forbid_alias.pto @@ -26,14 +26,14 @@ module attributes {"pto.target_arch" = "a3"} { %mask = pto.alloc_tile : !pto.tile_buf %src0 = pto.alloc_tile : !pto.tile_buf %src1 = pto.alloc_tile : !pto.tile_buf - %tmp = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, - !pto.tile_buf) + !pto.tile_buf) outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/pto/select_tile_native.pto b/test/lit/pto/select_tile_native.pto index a0ba395ba0..9e428bce07 100644 --- a/test/lit/pto/select_tile_native.pto +++ b/test/lit/pto/select_tile_native.pto @@ -10,8 +10,8 @@ // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { - func.func private @tsel_arg(%mask: !pto.tile_buf, %a: !pto.tile_buf, %b: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { - pto.tsel ins(%mask, %a, %b, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + func.func private @tsel_arg(%mask: !pto.tile_buf, %a: !pto.tile_buf, %b: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tsel ins(%mask, %a, %b, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) return } func.func private @tsels_arg(%mask: !pto.tile_buf, %src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: i16) { diff --git a/test/lit/pto/tsel_bf16.pto b/test/lit/pto/tsel_bf16.pto index f0dea4d774..cbc41cc577 100644 --- a/test/lit/pto/tsel_bf16.pto +++ b/test/lit/pto/tsel_bf16.pto @@ -14,10 +14,10 @@ module { %mask = pto.alloc_tile : !pto.tile_buf %src0 = pto.alloc_tile : !pto.tile_buf %src1 = pto.alloc_tile : !pto.tile_buf - %tmp = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/pto/tsel_tmp_contract_a3_invalid.pto b/test/lit/pto/tsel_tmp_contract_a3_invalid.pto new file mode 100644 index 0000000000..29b79e0ac4 --- /dev/null +++ b/test/lit/pto/tsel_tmp_contract_a3_invalid.pto @@ -0,0 +1,24 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module attributes {"pto.device-spec" = "Ascend910B3"} { + func.func @a3_tsel_tmp_too_small() attributes {pto.kernel_kind = #pto.kernel_kind} { + %mask = pto.alloc_tile : !pto.tile_buf + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK: error: 'pto.tsel' op expects tmp capacity to be at least 64 bytes, but got 32 bytes From 24da54612781003266cb2f30cda54d01b90ade7e Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 09:32:44 +0800 Subject: [PATCH 13/17] Relax tsel tmp verifier --- lib/PTO/IR/PTO.cpp | 5 ++--- test/lit/pto/tsel_tmp_contract_a3_4byte.pto | 25 +++++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) create mode 100644 test/lit/pto/tsel_tmp_contract_a3_4byte.pto diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 7c1f269e72..fa24568abc 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -13072,9 +13072,8 @@ mlir::LogicalResult mlir::pto::TSelOp::verify() { "expects A2/A3 tsel src0, src1, and dst element type to be i16/i32/f16/bf16/f32"); if (getTmp()) { Type tmpTy = getTmp().getType(); - auto tmpElem = dyn_cast(getElemTy(tmpTy)); - if (!tmpElem || tmpElem.getWidth() != 32) - return emitOpError("expects A2/A3 tsel tmp element type to be i32"); + if (getElemByteSize(getElemTy(tmpTy)) != 4) + return emitOpError("expects A2/A3 tsel tmp element type to be 4 bytes wide"); unsigned elemBits = getPTOStorageElemBitWidth(elem); if (elemBits != 16 && elemBits != 32) return emitOpError("expects A2/A3 tsel data element type to be 16 or 32 bits"); diff --git a/test/lit/pto/tsel_tmp_contract_a3_4byte.pto b/test/lit/pto/tsel_tmp_contract_a3_4byte.pto new file mode 100644 index 0000000000..09c894e950 --- /dev/null +++ b/test/lit/pto/tsel_tmp_contract_a3_4byte.pto @@ -0,0 +1,25 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s + +module attributes {"pto.device-spec" = "Ascend910B3"} { + func.func @a3_tsel_f32_tmp() attributes {pto.kernel_kind = #pto.kernel_kind} { + %mask = pto.alloc_tile : !pto.tile_buf + %src0 = pto.alloc_tile : !pto.tile_buf + %src1 = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// CHECK-LABEL: func.func @a3_tsel_f32_tmp +// CHECK: pto.tsel From 3df5f517ca8b01fa5f341b0071208bbba21746e4 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 09:40:38 +0800 Subject: [PATCH 14/17] Remove redundant tsel tmp lit --- test/lit/pto/tsel_tmp_contract_a3_4byte.pto | 25 --------------------- 1 file changed, 25 deletions(-) delete mode 100644 test/lit/pto/tsel_tmp_contract_a3_4byte.pto diff --git a/test/lit/pto/tsel_tmp_contract_a3_4byte.pto b/test/lit/pto/tsel_tmp_contract_a3_4byte.pto deleted file mode 100644 index 09c894e950..0000000000 --- a/test/lit/pto/tsel_tmp_contract_a3_4byte.pto +++ /dev/null @@ -1,25 +0,0 @@ -// Copyright (c) 2026 Huawei Technologies Co., Ltd. -// This program is free software, you can redistribute it and/or modify it under the terms and conditions of -// CANN Open Software License Agreement Version 2.0 (the "License"). -// Please refer to the License for details. You may not use this file except in compliance with the License. -// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, -// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. -// See LICENSE in the root of the software repository for the full text of the License. - -// RUN: ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s - -module attributes {"pto.device-spec" = "Ascend910B3"} { - func.func @a3_tsel_f32_tmp() attributes {pto.kernel_kind = #pto.kernel_kind} { - %mask = pto.alloc_tile : !pto.tile_buf - %src0 = pto.alloc_tile : !pto.tile_buf - %src1 = pto.alloc_tile : !pto.tile_buf - %tmp = pto.alloc_tile : !pto.tile_buf - %dst = pto.alloc_tile : !pto.tile_buf - pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) - outs(%dst : !pto.tile_buf) - return - } -} - -// CHECK-LABEL: func.func @a3_tsel_f32_tmp -// CHECK: pto.tsel From f9980f70c40c74b4a5bc69f74798e005f86ad411 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 11:44:09 +0800 Subject: [PATCH 15/17] Tighten tsel tmp capacity to backend usage --- lib/PTO/IR/PTO.cpp | 2 +- test/lit/pto/tsel_tmp_contract_a3_invalid.pto | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index fa24568abc..055d53f740 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -13077,7 +13077,7 @@ mlir::LogicalResult mlir::pto::TSelOp::verify() { unsigned elemBits = getPTOStorageElemBitWidth(elem); if (elemBits != 16 && elemBits != 32) return emitOpError("expects A2/A3 tsel data element type to be 16 or 32 bits"); - uint64_t minBytes = elemBits == 16 ? 64 : 32; + uint64_t minBytes = elemBits == 16 ? 16 : 8; if (failed(verifyTmpCapacityAtLeast(*this, tmpTy, minBytes))) return failure(); } diff --git a/test/lit/pto/tsel_tmp_contract_a3_invalid.pto b/test/lit/pto/tsel_tmp_contract_a3_invalid.pto index 29b79e0ac4..937c2d7e35 100644 --- a/test/lit/pto/tsel_tmp_contract_a3_invalid.pto +++ b/test/lit/pto/tsel_tmp_contract_a3_invalid.pto @@ -9,16 +9,15 @@ // RUN: not ptoas --pto-arch=a3 --emit-pto-ir %s 2>&1 | FileCheck %s module attributes {"pto.device-spec" = "Ascend910B3"} { - func.func @a3_tsel_tmp_too_small() attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @a3_tsel_tmp_too_small(%tmp: !pto.tile_buf) attributes {pto.kernel_kind = #pto.kernel_kind} { %mask = pto.alloc_tile : !pto.tile_buf %src0 = pto.alloc_tile : !pto.tile_buf %src1 = pto.alloc_tile : !pto.tile_buf - %tmp = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + pto.tsel ins(%mask, %src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) return } } -// CHECK: error: 'pto.tsel' op expects tmp capacity to be at least 64 bytes, but got 32 bytes +// CHECK: error: 'pto.tsel' op expects tmp capacity to be at least 16 bytes, but got 8 bytes From 51f858fb269e411a169eccf9d8e82b082601c33c Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 7 Aug 2026 16:08:23 +0800 Subject: [PATCH 16/17] Refresh implicit-tmp design for capacity and A5 level3 Align the doc with actual materialize/verifier behavior: reduction tmp capacity is checked from declared shape and the 32B floor is always satisfied because alloc_tile enforces 32B row alignment; A5 level3 synthesizes a no-effect 32B ABI placeholder instead of erroring. --- ...oas-implicit-tmp-materialization-design.md | 47 +++++++++++++------ 1 file changed, 32 insertions(+), 15 deletions(-) diff --git a/docs/designs/ptoas-implicit-tmp-materialization-design.md b/docs/designs/ptoas-implicit-tmp-materialization-design.md index aaad9a2276..8904bd9883 100644 --- a/docs/designs/ptoas-implicit-tmp-materialization-design.md +++ b/docs/designs/ptoas-implicit-tmp-materialization-design.md @@ -114,22 +114,31 @@ target_op(no tmp) ### level3 -level3 下用户显式管理 local 地址,memplan 通常跳过。因此不应自动创建无地址 tmp。 +level3 下用户显式管理 local 地址,memplan 通常跳过。pass 通过构造参数 `requireExplicitTmp` 判定 level3(`createPTOMaterializeImplicitTmpPass(effectiveLevel == Level3)`)。 -通用规则: +对 **A2/A3 实际使用 tmp 的 op**,level3 不自动创建无地址 tmp,缺省 tmp 直接报错: ```text -level3 + target_op(no tmp) => pass/verifier 报错 +level3 + A2/A3 target_op(no tmp) => pass 报错 ``` -用户在 level3 使用已纳入改造的目标 op 时,必须显式提供合法 tmp,并保证 tmp 自身带合法 local addr,或满足现有 level3 显式地址规则。 - -诊断信息示例: +实际诊断字符串(`PTOMaterializeImplicitTmp.cpp`): ```text - requires explicit tmp when compiling at level3 because PlanMemory is skipped + requires explicit tmp when PlanMemory is skipped ``` +个别 op 有更具体的变体,例如 binary tcolsum 为 `requires explicit tmp for binary tcolsum when PlanMemory is skipped`,非 32 对齐 tsort32 为 `requires explicit tmp for non-32-aligned tsort32 when PlanMemory is skipped`。 + +A2/A3 用户在 level3 使用这些 op 时,必须显式提供合法 tmp,并保证 tmp 自身带合法 local addr,或满足现有 level3 显式地址规则。 + +**A5 例外**:对 A5 只接受但不使用 tmp 的 op,pass 通过 `!isA5` 跳过上述 level3 报错: + +- 若后端 C++ 签名仍要求 tmp(row/arg reduction、TXOR/TXORS、TSEL/TSELS、TPRELU、TREM/TREMS 等),pass 即使在 level3 也自动生成固定 32 字节的 ABI placeholder(`makeA5PlaceholderTmpType`,形状 `{1, 32/sizeof(elem)}`)。该 placeholder 不建模 Read/Write MemoryEffects、不参与 memplan、无需回写地址,因此不违反 level3“不自动分配 tmp 地址”的非目标。 +- 若后端存在 no-tmp overload(TCI、TROWEXPAND*、TQUANT 等),A5 直接保持 no-tmp 形态,不补 placeholder。 + +因此当前实现下不存在 “A5 + level3 + 缺省 tmp” 报错的路径:要么 pass 自动补 placeholder,要么保持合法的 no-tmp 形态,verifier 的 no-tmp overload 也无条件接受。 + ## EmitC Lowering 目标 op 的 EmitC lowering 应保持简单: @@ -594,10 +603,11 @@ TSORT32, TQUANT ### 通用规则 - level1/level2:只有该 op 在当前 arch / 模式下实际需要 tmp,且 IR 允许省略 tmp 时,才自动 materialize `pto.alloc_tile(no addr)`。 -- level3:若该 op 在当前 arch / 模式下需要 tmp 且用户省略 tmp,则报错;若当前 arch / 模式不使用 tmp,则不强制补 tmp。 -- A5 仅接受但不使用 tmp 的 op:若后端存在 no-tmp overload,则不自动补 tmp;若 C++ 签名仍要求 tmp,则自动生成 ABI placeholder。placeholder 不建模 tmp 的 Read/Write,用户显式 tmp 也不按 A2/A3 容量规则校验。 +- level3:若该 op 在 A2/A3 当前模式下需要 tmp 且用户省略 tmp,则报错(`requires explicit tmp when PlanMemory is skipped`)。A5 见下一条,即使 level3 也不因缺省 tmp 报错。 +- A5 仅接受但不使用 tmp 的 op:若后端存在 no-tmp overload,则不自动补 tmp;若 C++ 签名仍要求 tmp,则自动生成固定 32 字节 ABI placeholder(level1/2/3 一致,通过 `!isA5` 绕过 level3 显式 tmp 检查)。placeholder 不建模 tmp 的 Read/Write(`getEffects` 以 `!tmp.empty() && arch != A5` 守卫)、不参与 memplan、无需地址,用户显式 tmp 也不按 A2/A3 容量规则校验。 - tmp 是 scratch 的 op:MemoryEffects 需要建模为 `Read(tmp) + Write(tmp)`,或在 semantic no-alias side table 中显式加入 `forbidAlias(tmp, dst/output)`。 - 原 mandatory tmp op 已统一改为 optional,并保持显式 tmp 文本格式兼容。 +- 容量校验统一走 `verifyTmpCapacityAtLeast`,按 tmp 的**声明 shape** × `sizeof(dtype)` 计算(`getStaticByteSize`,非 valid 区域),A5 分支不执行该校验。对 row/arg reduction 的 32 字节下限:因 `pto.alloc_tile` 已对 row-major none_box tile 强制行 `cols * sizeof(dtype)` 32 字节对齐,而 reduction `src` 必须是这类 tile,故任何合法 src 单行即 ≥32 字节,同形状 tmp 恒满足下限;materialize 无需为 sub-block src 额外兜底容量(<32 字节的合法 reduction src 无法构造)。 ### Arg reduction 类 @@ -625,9 +635,15 @@ TROWARGMAX / TROWARGMIN: - tmp 行数与 `src` 相同;每行 stride 按 PTO-ISA 文档公式计算。 - 当前 PTOAS ODS 只有单输出索引模式;该模式仍需满足后端显式 tmp 参数签名,因此 level1/2 生成保守同形状 tmp。未来接入值+索引模式后,再按输出模式和归约阶段收紧容量。 +容量校验: + +- A2/A3 verifier 在按 `tmpGapEles` / 布局(DN 1 列、ND 2 列、min stride 等)校验后,统一以 `verifyTmpCapacityAtLeast(op, tmp, 32)` 兜底;容量按声明 shape 计算,同 row reduction 一样被 `pto.alloc_tile` 的 32 字节行对齐恒满足。 +- A5 arg-reduction verifier(`verifyTColArgReductionOpA5` / `verifyTRowArgReductionOpA5`)不含任何容量校验。 + MemoryEffects / alias: - A2/A3 实际使用 tmp 时建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dstIdx/dstVal)`。 +- A5 placeholder 不建模 tmp 的 Read/Write(`getEffects` 以 `!tmp.empty() && arch != A5` 守卫)。 - tmp 不应与同 op 的输出 alias;如果 MemoryEffects 无法覆盖,应加入 `forbidAlias(tmp, dstIdx)` 和必要的 `forbidAlias(tmp, dstVal)`。 ### Row reduction 类 @@ -646,19 +662,20 @@ TROWPROD, TROWSUM, TROWMAX, TROWMIN tmp 规格: - tmp dtype 与 `src` / `dst` 一致。 -- 整数路径最小需要 1 个 vector block:`int32` 为 8 列,`int16` 为 16 列。 -- 浮点路径用于二叉树归约;安全默认形状可设为与 `src` 相同。 -- `TROWPROD` 的安全默认也可设为与 `src` 相同;最小需求为 1 行和 1 个 vector block。 +- ISA 最小需求为 1 行 1 个 vector block(32 字节):`int32` 为 8 列,`int16` 为 16 列;浮点二叉树归约同样以 1 个 block 为下限。 +- A2/A3 materialize 直接生成与 `src` 同形状的 tmp(`makeSameShapeTmpType`),不做逐 dtype 的特化裁剪;`TROWPROD` 亦同。 +- 容量校验:A2/A3 verifier 走 `verifyTmpCapacityAtLeast(op, tmp, 32)`,按声明 shape × `sizeof(dtype)` 计算。**该 32 字节下限对合法 IR 恒被满足**——reduction `src` 必须是 row-major none_box tile,`pto.alloc_tile` 已对其强制行 32 字节对齐,故同形状 tmp 单行即 ≥32 字节,无需为 sub-block src 特殊兜底。 Pass 行为: -- A2/A3 level1/2:若 IR 已支持 optional tmp 且缺省 tmp,则自动生成 vec row-major none-box tmp。 -- A5:生成后端签名需要的 ABI placeholder,但不为其添加 tmp MemoryEffects。 -- level3:A2/A3 需要 tmp 时缺省 tmp 报错;A5 不强制 tmp。 +- A2/A3 level1/2:若 IR 已支持 optional tmp 且缺省 tmp,则自动生成与 `src` 同形状的 vec row-major none-box tmp。 +- A5:生成后端签名需要的固定 32 字节 ABI placeholder,但不为其添加 tmp MemoryEffects,也不跑 32 字节容量校验。 +- level3:A2/A3 需要 tmp 时缺省 tmp 报错(`requires explicit tmp when PlanMemory is skipped`);A5 通过 `!isA5` 跳过报错,仍自动生成 ABI placeholder。 MemoryEffects / alias: - A2/A3 建模 `Read(src) + Read(tmp) + Write(tmp) + Write(dst)`。 +- A5 placeholder 不建模 tmp 的 Read/Write(`getEffects` 以 `!tmp.empty() && arch != A5` 守卫)。 - tmp 与 `dst` 禁止 alias。 ### Column sum 类 From a1aa09e510fa744fe48657a80fd2c817d1b74951 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 8 Aug 2026 11:01:26 +0800 Subject: [PATCH 17/17] Align A5 ttrans implicit tmp with other ops TTrans passed a bare requireExplicitTmp, causing A5+level3 to wrongly error on a missing tmp. The A5 backend ignores the ttrans tmp buffer, so materialize a no-address placeholder like the other A5 ops instead. --- .../Transforms/PTOMaterializeImplicitTmp.cpp | 11 ++++--- .../lit/pto/ttrans_implicit_tmp_a5_level3.pto | 30 +++++++++++++++++++ 2 files changed, 37 insertions(+), 4 deletions(-) create mode 100644 test/lit/pto/ttrans_implicit_tmp_a5_level3.pto diff --git a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp index b832503501..92139a09e6 100644 --- a/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp +++ b/lib/PTO/Transforms/PTOMaterializeImplicitTmp.cpp @@ -627,6 +627,8 @@ static LogicalResult materializeFixedMandatoryTmp(Operation *op, .Case([&](auto typedOp) -> LogicalResult { if (typedOp.getTmp()) return success(); + bool isA5 = + pto::getTargetArch(op) == pto::PTOArch::A5; auto srcTy = dyn_cast(typedOp.getSrc().getType()); auto dstTy = dyn_cast(typedOp.getDst().getType()); auto srcShape = getShapeVec(typedOp.getSrc().getType()); @@ -643,16 +645,17 @@ static LogicalResult materializeFixedMandatoryTmp(Operation *op, bool usesTmp = dstShape[1] % rowStride == 0 && srcShape[1] % elemPerBlock == 0 && srcShape[1] / elemPerBlock <= 255; - FailureOr type = makeSameShapeTmpType( - ctx, typedOp.getSrc()); - if (!usesTmp) + FailureOr type = + isA5 ? makeA5PlaceholderTmpType(ctx, typedOp.getSrc()) + : makeSameShapeTmpType(ctx, typedOp.getSrc()); + if (!isA5 && !usesTmp) type = makeVecTmpType(ctx, {1, elemPerBlock}, srcTy.getElementType(), {1, elemPerBlock}); if (failed(type)) return typedOp.emitOpError("failed to build implicit ttrans tmp"); return replaceFixedDpsOpWithTmp( op, {typedOp.getSrc(), Value(), typedOp.getDst()}, *type, - {1, 1, 1}, requireExplicitTmp, "ttrans"); + {1, 1, 1}, isA5 ? false : requireExplicitTmp, "ttrans"); }) .Default([](Operation *) { return success(); }); } diff --git a/test/lit/pto/ttrans_implicit_tmp_a5_level3.pto b/test/lit/pto/ttrans_implicit_tmp_a5_level3.pto new file mode 100644 index 0000000000..98ccbe2945 --- /dev/null +++ b/test/lit/pto/ttrans_implicit_tmp_a5_level3.pto @@ -0,0 +1,30 @@ +// 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. + +// On A5 the backend ignores the ttrans tmp buffer, so an implicit tmp must be +// materialized as a no-address placeholder at both level2 and level3 rather +// than rejected. level3 skips PlanMemory; the ttrans handler must not require +// an explicit tmp on A5 (consistent with the other A5 ops). + +// RUN: sed -E 's/ addr = %[A-Za-z0-9_]+//g' %s > %t.level2.pto && ptoas --pto-arch=a5 --pto-level=level2 --emit-pto-ir %t.level2.pto 2>&1 | FileCheck %s --check-prefix=A5 +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --emit-pto-ir %s 2>&1 | FileCheck %s --check-prefix=A5 + +module { + func.func @a5_ttrans_implicit_tmp() attributes {pto.kernel_kind = #pto.kernel_kind} { + %a0 = arith.constant 0 : i64 + %a2048 = arith.constant 2048 : i64 + %src = pto.alloc_tile addr = %a0 : !pto.tile_buf + %dst = pto.alloc_tile addr = %a2048 : !pto.tile_buf + pto.ttrans ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// A5-LABEL: func.func @a5_ttrans_implicit_tmp +// A5: pto.ttrans ins(%{{.*}}, %{{.*}} : !pto.tile_buf, !pto.tile_buf)