From 6a4c2339d4774c7650fee2f3374d65ba43381429 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 9 Jul 2026 15:48:55 +0800 Subject: [PATCH 01/91] Add dual memory planner support --- .../ptoas-dual-memory-planner-design.md | 188 +++++ include/PTO/IR/PTOAttrs.td | 13 - include/PTO/Transforms/Passes.h | 6 +- include/PTO/Transforms/Passes.td | 43 +- lib/PTO/Transforms/AllocToPointerCast.cpp | 99 +-- lib/PTO/Transforms/CMakeLists.txt | 1 + lib/PTO/Transforms/PTOPlanMemory.cpp | 158 +++- lib/PTO/Transforms/PTOPlanMemory.h | 5 + lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 701 ++++++++++++++++++ lib/PTO/Transforms/PTOViewToMemref.cpp | 232 ++---- lib/PTO/Transforms/Utils.cpp | 7 + ptodsl/tests/test_docs_as_test.py | 6 +- ptodsl/tests/test_ptoas_frontend_verify.py | 7 +- .../pto/alloc_tile_addr_alignment_invalid.pto | 1 + .../pto/alloc_tile_addr_level12_reject.pto | 21 + ...alloc_tile_plan_memory_no_memref_alloc.pto | 35 + .../multi_tile_get_const_slot_lowering.pto | 4 +- .../pto/multi_tile_level3_explicit_addr.pto | 13 +- .../pto/multi_tile_level3_missing_addr.pto | 6 +- .../plan_memory_bind_tile_alias_liveness.pto | 6 +- .../pto/plan_memory_for_iter_args_yield.pto | 1 + .../plan_memory_fragmentation_hole_fit.pto | 1 + .../plan_memory_fragmentation_two_holes.pto | 1 + test/lit/pto/plan_memory_if_in_loop.pto | 5 +- test/lit/pto/plan_memory_if_yield.pto | 1 + test/lit/pto/plan_memory_loop_in_if.pto | 2 +- .../plan_memory_loop_no_reuse_outer_live.pto | 1 + test/lit/pto/plan_memory_nested_loops.pto | 1 + test/lit/pto/plan_memory_no_reuse_overlap.pto | 1 + .../pto/plan_memory_order_by_size_noreuse.pto | 6 +- .../pto/plan_memory_order_by_size_reuse.pto | 6 +- .../pto/plan_memory_peak_8_overlapping.pto | 2 +- .../pto/plan_memory_peak_exact_capacity.pto | 2 +- ...erve_buffer_manual_reject_nested_alloc.pto | 3 +- .../pto/plan_memory_reserve_buffer_prefix.pto | 1 + test/lit/pto/plan_memory_reuse_sequential.pto | 1 + ...an_memory_spec_level0_no_reuse_overlap.pto | 31 + .../lit/pto/plan_memory_spec_level0_reuse.pto | 31 + ...eserve_buffer_level12_auto_base_reject.pto | 25 + ...erve_buffer_level12_manual_base_reject.pto | 25 + .../pto/reserve_buffer_level3_auto_reject.pto | 23 + .../pto/reserve_buffer_level3_manual_base.pto | 27 + tools/ptoas/ptoas.cpp | 128 +++- 43 files changed, 1571 insertions(+), 306 deletions(-) create mode 100644 docs/designs/ptoas-dual-memory-planner-design.md create mode 100644 lib/PTO/Transforms/PTOPlanMemoryModern.cpp create mode 100644 test/lit/pto/alloc_tile_addr_level12_reject.pto create mode 100644 test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto create mode 100644 test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto create mode 100644 test/lit/pto/plan_memory_spec_level0_reuse.pto create mode 100644 test/lit/pto/reserve_buffer_level12_auto_base_reject.pto create mode 100644 test/lit/pto/reserve_buffer_level12_manual_base_reject.pto create mode 100644 test/lit/pto/reserve_buffer_level3_auto_reject.pto create mode 100644 test/lit/pto/reserve_buffer_level3_manual_base.pto diff --git a/docs/designs/ptoas-dual-memory-planner-design.md b/docs/designs/ptoas-dual-memory-planner-design.md new file mode 100644 index 0000000000..0c0711e78b --- /dev/null +++ b/docs/designs/ptoas-dual-memory-planner-design.md @@ -0,0 +1,188 @@ +# PTOAS 双内存规划器设计说明 + +## 背景 + +PTOAS 之前曾临时删除旧版 `pto-plan-memory` pass，以便让 level1/level2 在没有完整 memplan 的情况下继续编译。同时，`pto.alloc_tile` 的 lowering 路径也从旧的 `pto.alloc_tile -> memref.alloc -> pto.pointer_cast -> pto.bind_tile` 调整为 tile-native 路径：无显式地址的 `pto.alloc_tile` 保留到 memplan，由 memplan 直接补充常量 `addr`，后续 `pto.t*` tile op 继续消费 `!pto.tile_buf`。 + +随着旧 memplan 恢复，需要同时满足两个需求： + +- 默认行为保持稳定，继续使用旧 memplan，避免影响现有 pipeline 和测试基线。 +- 保留重写后的 modern memplan，便于继续开发 alias-aware、SPEC_LEVEL_0 复用等新策略。 + +因此，本次提交将旧版和新版 memplan 拆成两个实现，由 `tools/ptoas/ptoas.cpp` 通过 CLI 参数选择使用哪一个。 + +## 目标 + +本次改动的目标如下： + +- 默认使用 legacy memplan，保证既有行为和兼容性。 +- 通过 `--plan-memory-impl=modern` 显式启用 modern memplan。 +- `pto-plan-memory` 仍保持 module 级 pass，pass 内部遍历 `func::FuncOp` 做规划。 +- `pto.alloc_tile(no addr)` 不再经过 `memref.alloc` 中间路径，而是由所选 memplan 直接补 `addr`。 +- 尽量减少对 legacy memplan 原有代码的结构性改动，只补充 tile-native alloc 所需的最小适配。 +- 将原先删除的 memplan lit/sample 用例恢复，并让现有 `plan_memory_*.pto` lit 同时覆盖 legacy 和 modern。 + +## 非目标 + +本次改动不试图完成 modern memplan 的全部设计目标： + +- 不实现 SPEC_LEVEL_1 / SPEC_LEVEL_2 的完整投机复用策略。 +- 不替换 legacy memplan 的 pipeline conflict、double-buffer 相关逻辑。 +- 不改变 level3 显式地址模式的规则。 +- 不把 `reserve_buffer` 手动 base 支持扩展到 level1/level2。 + +## 设计方案 + +### 文件组织 + +旧 memplan 仍保留在原文件： + +```text +lib/PTO/Transforms/PTOPlanMemory.cpp +lib/PTO/Transforms/PTOPlanMemory.h +``` + +modern memplan 放在独立文件： + +```text +lib/PTO/Transforms/PTOPlanMemoryModern.cpp +``` + +`PTOPlanMemory.cpp` 继续定义默认 pass factory： + +```cpp +createPlanMemoryPass(const PlanMemoryOptions &options = {}) +``` + +`PTOPlanMemoryModern.cpp` 定义 modern factory： + +```cpp +createPlanMemoryModernPass(const PlanMemoryOptions &options) +``` + +`lib/PTO/Transforms/CMakeLists.txt` 同时编译两个实现文件。 + +### Pass 级别 + +`pto-plan-memory` 保持 module 级 pass： + +```td +def PlanMemory : Pass<"pto-plan-memory", "ModuleOp"> +``` + +legacy 和 modern 都在 module pass 的 `runOnOperation()` 中遍历 module 内的 `func::FuncOp`，分别对每个函数做静态 local memory planning。这样保留旧 memplan 的 module-level 调用形态，也避免在 pass manager pipeline 中暴露两个不同的 nested func pass。 + +### 实现选择 + +实现选择不放在 `PTOPlanMemory.cpp` 内部，而是在 `tools/ptoas/ptoas.cpp` 中完成： + +```cpp +if (planMemoryImpl == "legacy") { + pm.addPass(pto::createPlanMemoryPass(planMemoryOptions)); +} else if (planMemoryImpl == "modern") { + pm.addPass(pto::createPlanMemoryModernPass(planMemoryOptions)); +} +``` + +CLI 参数为： + +```text +--plan-memory-impl=legacy # 默认 +--plan-memory-impl=modern +``` + +这样 `PTOPlanMemory.cpp` 可以继续表示 legacy memplan 本体，而不是变成 dispatcher。 + +### `memMode` + +`PlanMemoryOptions::memMode` 当前仍由 `ptoas.cpp` 固定设置为： + +```cpp +planMemoryOptions.memMode = "local"; +``` + +含义是规划 local memory buffer。legacy memplan 内部仍保留 `MemPlanMode` 枚举，用于尽量少改旧代码： + +```cpp +enum class MemPlanMode { + LOCAL_MEM_PLAN, + GLOBAL_WORKSPACE_PLAN, +}; +``` + +当前 PTOAS 主 pipeline 只使用 `"local"`。`GLOBAL_WORKSPACE_PLAN` 是 legacy 代码的历史模式保留，不作为当前 level1/level2 local memplan 的正式入口。 + +### `pto.alloc_tile(no addr)` 处理 + +`PTOViewToMemref` 不再把无地址的 `pto.alloc_tile` 降成 `memref.alloc`。因此，两个 memplan 都需要直接处理 tile-native allocation root： + +```mlir +%tile = pto.alloc_tile : !pto.tile_buf<...> +``` + +memplan 完成规划后直接补常量地址： + +```mlir +%c0_i64 = arith.constant 0 : i64 +%tile = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<...> +``` + +legacy memplan 为此增加了最小适配： + +- liveness 收集阶段识别 `pto.alloc_tile(no addr)`。 +- `BufferInfo` 支持从 `TileBufType` 计算 shape、element type 和字节数。 +- materialize 阶段新增 `pto.alloc_tile -> pto.alloc_tile addr` 的 rewrite。 +- `pto.tile_buf_addr` 被建模为读取 tile source，避免被未知 local-buffer op 检查误判。 + +### `memref.alloc` 处理 + +非 tile-native 的 `memref.alloc` root 仍沿用 legacy/modern 各自的既有 materialization 方式： + +```mlir +memref.alloc -> pto.pointer_cast(...) +``` + +也就是说，本次改动只要求 `pto.alloc_tile` 不再经过 `memref.alloc` 中间链路；并不禁止普通 memref root 继续使用 `pto.pointer_cast` 表达规划地址。 + +## 测试方案 + +测试覆盖分为三类。 + +### 恢复旧 memplan 测试 + +恢复 `test/samples/planmemory` 下之前删除的 sample case，覆盖 loop、if、nested loop、fragmentation、peak capacity、reuse 等场景。 + +### lit 双实现覆盖 + +现有 `test/lit/pto/plan_memory_*.pto` 增加 modern RUN： + +```mlir +// RUN: ptoas ... | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern ... | FileCheck %s +``` + +`order_by_size` 用例同时覆盖 legacy/modern 的默认顺序和 `--plan-memory-order-by-size` 顺序。负例也同时覆盖 legacy/modern。 + +### 回归命令 + +本次提交使用以下命令验证： + +```bash +cmake --build build --target ptoas -j8 + +PATH=/Users/fangrui/workspace/huawei/llvm21-workspace/llvm-project/llvm/build-assert/bin:$PATH \ + /Users/fangrui/workspace/huawei/llvm21-workspace/llvm-project/llvm/build-assert/bin/llvm-lit \ + -sv build/test/lit \ + --filter 'plan_memory|reserve_buffer|alloc_tile_addr|alloc_tile_plan_memory_no_memref_alloc|multi_tile_get_const_slot_lowering|multi_tile_.*planmem' + +ctest --test-dir build --output-on-failure -L PTODSL +``` + +同时对恢复的 `test/samples/planmemory/*.py` 做逐个生成和 `ptoas --emit-pto-ir` 编译验证。 + +## 后续工作 + +- modern memplan 后续继续补齐 SPEC_LEVEL_1 / SPEC_LEVEL_2。 +- 根据 modern 策略差异，为必要用例拆分 legacy/modern 的独立 FileCheck 前缀。 +- 如果确认 global workspace planning 不再需要，可单独清理 legacy 内部的 `GLOBAL_WORKSPACE_PLAN` 历史模式。 +- 继续减少 legacy 文件中的非必要 diff，保持旧实现可对照、可回退。 diff --git a/include/PTO/IR/PTOAttrs.td b/include/PTO/IR/PTOAttrs.td index b87e1486e6..da4a7237b5 100644 --- a/include/PTO/IR/PTOAttrs.td +++ b/include/PTO/IR/PTOAttrs.td @@ -600,19 +600,6 @@ def PTO_EventAttr : PTO_Attr<"Event", "event"> { }]; } -//===----------------------------------------------------------------------===// -// Plan Memory Mode -//===----------------------------------------------------------------------===// - -def PTO_LOCALMEMPLAN : I32EnumAttrCase<"LOCAL_MEM_PLAN", 0>; -def PTO_GLOBALWORKSPACEPLAN : I32EnumAttrCase<"GLOBAL_WORKSPACE_PLAN", 1>; - -def PTO_MemPlanModeEnum : PTO_I32Enum< - "MemPlanMode", "Mem Plan Mode", [ - PTO_LOCALMEMPLAN, - PTO_GLOBALWORKSPACEPLAN -]>; - //===----------------------------------------------------------------------===// // Pad Mode for LoadOp //===----------------------------------------------------------------------===// diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index 3555735a06..d88b918c95 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -67,10 +67,10 @@ std::unique_ptr createConvertToPTOOpPass(); /// PTO Ops. std::unique_ptr createInferPTOMemScopePass(); -/// Create a pass to plan memory. std::unique_ptr -createPlanMemoryPass(const PlanMemoryOptions &planMemoryOption = {}); - +createPlanMemoryPass(const PlanMemoryOptions &options = {}); +std::unique_ptr +createPlanMemoryModernPass(const PlanMemoryOptions &options); std::unique_ptr createPTORemoveRedundantBarrierPass(); std::unique_ptr createPTOViewToMemrefPass(); std::unique_ptr createPTOValidateIntToPtrUsesPass(); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 00b57ee697..2c782033cb 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -193,37 +193,26 @@ def PTORematerializeFixpipeVectorQuant } def PlanMemory : Pass<"pto-plan-memory", "ModuleOp"> { - let summary = "Plan memory for PTO Ops"; + let summary = "Plan memory for PTO Ops"; let constructor = "mlir::pto::createPlanMemoryPass()"; - let dependentDialects = ["pto::PTODialect", ]; + let options = [ - Option<"memMode", "mem-plan-mode", "pto::MemPlanMode", - "pto::MemPlanMode::LOCAL_MEM_PLAN", - "plan mem mode (default is LOCAL_MEM_PLAN)", - [{::llvm::cl::values( - clEnumValN(pto::MemPlanMode::LOCAL_MEM_PLAN, "local-mem-plan", - "plan mem mode is for memref.alloc"), - clEnumValN( - pto::MemPlanMode::GLOBAL_WORKSPACE_PLAN, - "global-work-space-plan", - "plan mem mode is for memref_ext.alloc_workspace"))}]>, - Option<"enableGlobalReuse", "enable-global-workspace-reuse", "bool", - /*default=*/"false", - "Enable global workspace reuse ,default : false">, - Option<"enablePrintMemoryAllocatedSize", "enable-print-memory-allocated-size", "bool", - /*default=*/"false", - "Enable print memory allocated size, default : false">, - Option<"restrictInplaceAsISA", "restrict-inplace-as-isa", "bool", - /*default=*/"false", - "restrict memory inplace as isa, default : false">, - Option<"orderBySize", "order-by-size", "bool", - /*default=*/"false", - "Process buffers largest-first (first-fit-decreasing order) during " - "local memory planning instead of the default DMA-first order. " - "Decreasing-size order packs heterogeneous-size buffers tighter " - "(matches XLA/TVM/SOMAS). default : false">, + Option<"memMode", "mem-mode", "std::string", "\"local\"", + "Planning mode. Supported value today: local">, + Option<"orderBySize", "order-by-size", "bool", "false", + "Plan larger buffers first inside one AddressSpace before applying " + "SPEC_LEVEL_0 first-fit reuse"> + ]; + + let dependentDialects = [ + "mlir::pto::PTODialect", + "mlir::memref::MemRefDialect", + "mlir::arith::ArithDialect", + "mlir::func::FuncDialect", + "mlir::scf::SCFDialect" ]; } + def PTOLoweringSyncToPipe : Pass<"pto-lowering-sync-to-pipe", "func::FuncOp"> { let summary = "Lower high-level sync ops to low-level pipe ops"; let description = [{ diff --git a/lib/PTO/Transforms/AllocToPointerCast.cpp b/lib/PTO/Transforms/AllocToPointerCast.cpp index f560458e5f..925993e8ec 100644 --- a/lib/PTO/Transforms/AllocToPointerCast.cpp +++ b/lib/PTO/Transforms/AllocToPointerCast.cpp @@ -10,16 +10,9 @@ //===----------------------------------------------------------------------===// #include "AllocToPointerCast.h" +#include "PTO/IR/PTOMultiBuffer.h" #include "PTO/IR/PTOTypeUtils.h" -#include "PTO/Transforms/Passes.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" - -namespace mlir { -#define GEN_PASS_DEF_ALLOCTOPOINTERCAST -#include "PTO/Transforms/Passes.h.inc" - -} // namespace mlir using namespace mlir; using namespace mlir::pto; @@ -30,62 +23,77 @@ namespace { constexpr uint64_t kDefaultAllocAlignmentBytes = 4096; constexpr size_t kDynamicValidShapeRank = 2; -static TileBufConfigAttr inferBindTileConfig(memref::AllocOp op) { +static uint64_t alignUpToDefault(uint64_t value) { + return ((value + kDefaultAllocAlignmentBytes - 1) / + kDefaultAllocAlignmentBytes) * + kDefaultAllocAlignmentBytes; +} + +static TileBufConfigAttr inferBindTileConfig(Value root, Operation *diagOp) { TileBufConfigAttr configAttr; - for (Operation *user : op.getResult().getUsers()) { + for (Operation *user : root.getUsers()) { auto bind = dyn_cast(user); - if (!bind || bind.getSource() != op.getResult()) + if (!bind || bind.getSource() != root) continue; if (!configAttr) { configAttr = bind.getConfigAttr(); continue; } if (configAttr != bind.getConfigAttr()) { - op.emitWarning("alloc has multiple bind_tile users with different configs; " - "using the first one"); + diagOp->emitWarning( + "alloc has multiple bind_tile users with different configs; " + "using the first one"); break; } } return configAttr; } -static SmallVector getAllocatedOffsets(memref::AllocOp op, - BaseMemRefType memRefType, - const DenseMap> &buffer2Offsets, - uint64_t &fallbackNextOffset) { - auto iter = buffer2Offsets.find(op.getResult()); +static SmallVector getAllocatedOffsets( + Value root, Operation *allocLikeOp, BaseMemRefType memRefType, + const DenseMap> &buffer2Offsets, + uint64_t &fallbackNextOffset) { + auto iter = buffer2Offsets.find(root); SmallVector offsets; if (iter != buffer2Offsets.end()) offsets = iter->second; - if (offsets.empty()) { - // Estimate buffer size (best-effort). Most PTO tile buffers are 32x32 and - // naturally align to 4096 bytes. - uint64_t bytes = kDefaultAllocAlignmentBytes; - if (auto memrefTy = dyn_cast(memRefType)) { - uint64_t elemBytes = getPTOStorageElemByteSize(memrefTy.getElementType()); - if (elemBytes != 0) { - uint64_t numel = 1; - bool allStatic = true; - for (int64_t d : memrefTy.getShape()) { - if (d == ShapedType::kDynamic) { - allStatic = false; - break; - } - numel *= static_cast(d); + if (!offsets.empty()) + return offsets; + + // Estimate buffer size (best-effort). Most PTO tile buffers are 32x32 and + // naturally align to 4096 bytes. + uint64_t bytes = kDefaultAllocAlignmentBytes; + if (auto memrefTy = dyn_cast(memRefType)) { + uint64_t elemBytes = getPTOStorageElemByteSize(memrefTy.getElementType()); + if (elemBytes != 0) { + uint64_t numel = 1; + bool allStatic = true; + for (int64_t d : memrefTy.getShape()) { + if (d == ShapedType::kDynamic) { + allStatic = false; + break; } - if (allStatic && numel != 0) - bytes = numel * elemBytes; + numel *= static_cast(d); } + if (allStatic && numel != 0) + bytes = numel * elemBytes; } - uint64_t stride = ((bytes + kDefaultAllocAlignmentBytes - 1) / - kDefaultAllocAlignmentBytes) * - kDefaultAllocAlignmentBytes; - uint64_t off = fallbackNextOffset; - fallbackNextOffset += - std::max(stride, kDefaultAllocAlignmentBytes); - offsets.push_back(off); } + + uint64_t slotCount = 1; + if (auto attr = + allocLikeOp->getAttrOfType(pto::kPtoMultiBufferAttrName)) { + slotCount = attr.getValue().getZExtValue(); + } + + uint64_t stride = alignUpToDefault(bytes); + uint64_t off = fallbackNextOffset; + for (uint64_t i = 0; i < slotCount; ++i) + offsets.push_back(off + i * stride); + + fallbackNextOffset += + std::max(stride * slotCount, kDefaultAllocAlignmentBytes); return offsets; } @@ -106,9 +114,10 @@ static std::pair getDynamicValidShapeValues(memref::AllocOp op) { LogicalResult MemrefAllocaOpToPointerCastOpPattern::matchAndRewrite( memref::AllocOp op, PatternRewriter &rewriter) const { const auto ¤tMemRefType = cast(op.getType()); - TileBufConfigAttr configAttr = inferBindTileConfig(op); - SmallVector offsets = getAllocatedOffsets( - op, currentMemRefType, buffer2Offsets, fallbackNextOffset); + TileBufConfigAttr configAttr = inferBindTileConfig(op.getResult(), op); + SmallVector offsets = + getAllocatedOffsets(op.getResult(), op, currentMemRefType, + buffer2Offsets, fallbackNextOffset); SmallVector addrs; addrs.reserve(offsets.size()); for (uint64_t offset : offsets) { diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index ccd135000a..92956a1a3d 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -55,6 +55,7 @@ add_mlir_dialect_library(PTOTransforms AllocToPointerCast.cpp InferPTOMemScope.cpp PTOPlanMemory.cpp + PTOPlanMemoryModern.cpp PTORemoveRedundantBarrier.cpp InferPTOLayout.cpp PTOA5NormalizeTMovPass.cpp diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index 9d672de706..46dc31f90a 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -13,6 +13,7 @@ #include "PTO/IR/PTOMultiBuffer.h" #include "PTO/IR/PTOTypeUtils.h" +#include "Utils.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/GPU/IR/GPUDialect.h" @@ -424,6 +425,32 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { // op and is consumed later by PTOResolveBufferSelect / sync. UpdateBufferAlias(slotOp.getResult(), slotOp.getSource()); return WalkResult::advance(); + } else if (isLocalMemPlan() && dyn_cast(op)) { + auto allocTileOp = cast(op); + if (allocTileOp.getAddr()) { + return WalkResult::advance(); + } + auto memorySpaceAttr = GetBufferSpaceAttr(allocTileOp.getResult()); + if (!isLocalBuffer(memorySpaceAttr)) { + allocTileOp.emitError("Alloc tile buffer not at local space"); + return WalkResult::interrupt(); + } + if (auto attr = allocTileOp->getAttrOfType( + mlir::pto::kPtoMultiBufferAttrName)) { + uint64_t n = attr.getValue().getZExtValue(); + if (n < mlir::pto::kPtoMultiBufferMinNum || + n > mlir::pto::kPtoMultiBufferMaxNum) { + allocTileOp.emitError() + << "pto.multi_buffer must be in [" + << mlir::pto::kPtoMultiBufferMinNum << ", " + << mlir::pto::kPtoMultiBufferMaxNum << "] (got " << n << ")"; + return WalkResult::interrupt(); + } + buffer2MultiNum[allocTileOp.getResult()] = static_cast(n); + } + UpdateOpBufferInfo(op, op->getResults()); + UpdateOperandGenInfo(curOpInfo, allocTileOp.getResult()); + return WalkResult::advance(); } else if (isLocalMemPlan() && dyn_cast(op)) { auto allocOp = cast(op); if (failed(CheckLocalBufferAllocOp(op))) { @@ -499,6 +526,9 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { UpdateBufferAlias(selectOp.getResult(), selectOp.getTrueValue(), true); UpdateBufferAlias(selectOp.getResult(), selectOp.getFalseValue(), true); OpKillHandle(curOpInfo, live, op->getBlock()); + } else if (auto tileBufAddrOp = dyn_cast(op)) { + UpdateOpGenInfo(curOpInfo, ValueRange{tileBufAddrOp.getSrc()}); + OpKillHandle(curOpInfo, live, op->getBlock()); } else if (auto callOp = dyn_cast(op)) { UpdateOpGenInfo(curOpInfo, llvm::to_vector(callOp->getOperands())); OpKillHandle(curOpInfo, live, op->getBlock()); @@ -914,14 +944,23 @@ BufferInfo MemLivenessAnalysis::GetBufferInfo(Operation *op, Value operand, bufferInfo.operation = op; bufferInfo.bufferScope = bufferScope; // get buffer size, now for static shape - Value traceValue = tracebackMemRef(operand); - auto memRefType = cast(traceValue.getType()); - bufferInfo.bufferType = memRefType.getElementType(); + ArrayRef shape; + Type elementType; + if (auto tileType = dyn_cast(operand.getType())) { + shape = tileType.getShape(); + elementType = tileType.getElementType(); + } else { + Value traceValue = tracebackMemRef(operand); + auto memRefType = cast(traceValue.getType()); + shape = memRefType.getShape(); + elementType = memRefType.getElementType(); + } + bufferInfo.bufferType = elementType; std::optional totalStaticSize = - getStaticTotalSize(memRefType.getShape()); + getStaticTotalSize(shape); if (!totalStaticSize.has_value()) llvm::report_fatal_error("failed to obtain buffer static shape size"); - unsigned elemBytes = getPTOStorageElemByteSize(memRefType.getElementType()); + unsigned elemBytes = getPTOStorageElemByteSize(elementType); if (elemBytes == 0) llvm::report_fatal_error("failed to obtain buffer element byte size"); bufferInfo.constBits = @@ -2423,8 +2462,69 @@ PlanRecord MemPlan::RollbackOutline(PlanRecHis &history, } namespace { + +class LegacyAllocTileOpAddPlannedAddressPattern + : public OpRewritePattern { +public: + explicit LegacyAllocTileOpAddPlannedAddressPattern( + MLIRContext *context, + DenseMap> buffer2Offsets) + : OpRewritePattern(context), + buffer2Offsets(std::move(buffer2Offsets)) {} + + LogicalResult matchAndRewrite(pto::AllocTileOp op, + PatternRewriter &rewriter) const override { + if (op.getAddr()) + return failure(); + + auto tileType = dyn_cast(op.getResult().getType()); + if (!tileType) + return failure(); + + auto it = buffer2Offsets.find(op.getResult()); + if (it == buffer2Offsets.end() || it->second.empty()) + return failure(); + + if (it->second.size() != 1) { + return rewriter.notifyMatchFailure( + op, "single alloc_tile root expects exactly one planned address"); + } + + Value addr = rewriter.create(op.getLoc(), + it->second.front(), 64); + auto planned = rewriter.create( + op.getLoc(), tileType, addr, + op.getValidRow() ? op.getValidRow() : Value(), + op.getValidCol() ? op.getValidCol() : Value()); + for (NamedAttribute attr : op->getAttrs()) { + if (attr.getName().getValue() == "operandSegmentSizes") + continue; + planned->setAttr(attr.getName(), attr.getValue()); + } + + rewriter.replaceOp(op, planned.getResult()); + return success(); + } + +private: + DenseMap> buffer2Offsets; +}; + +static FailureOr parseLegacyMemPlanMode(func::FuncOp func, + llvm::StringRef memMode) { + if (memMode.equals_insensitive("local") || + memMode.equals_insensitive("local-mem-plan")) + return MemPlanMode::LOCAL_MEM_PLAN; + if (memMode.equals_insensitive("global-work-space-plan")) + return MemPlanMode::GLOBAL_WORKSPACE_PLAN; + func.emitError("unsupported mem-mode '") + << memMode << "'; only 'local' is supported by the PTOAS pipeline"; + return failure(); +} + struct PlanMemoryPass : public mlir::pto::impl::PlanMemoryBase { public: + PlanMemoryPass() = default; explicit PlanMemoryPass(const mlir::pto::PlanMemoryOptions &planMemoryOption) : PlanMemoryBase(planMemoryOption) {} @@ -2432,11 +2532,13 @@ struct PlanMemoryPass : public mlir::pto::impl::PlanMemoryBase { private: void populateBufferAddressToAllocOp( - RewritePatternSet &patterns, + RewritePatternSet &patterns, MemPlanMode mode, DenseMap> buffer2Offsets) { - if (this->memMode == MemPlanMode::LOCAL_MEM_PLAN) { + if (mode == MemPlanMode::LOCAL_MEM_PLAN) { patterns.add(patterns.getContext(), buffer2Offsets); + patterns.add( + patterns.getContext(), buffer2Offsets); } } }; @@ -2445,12 +2547,17 @@ struct PlanMemoryPass : public mlir::pto::impl::PlanMemoryBase { void PlanMemoryPass::runOnOperation() { ModuleOp moduleOp = getOperation(); for (auto funcOp : moduleOp.getOps()) { + auto parsedMode = parseLegacyMemPlanMode(funcOp, this->memMode); + if (failed(parsedMode)) + return signalPassFailure(); + MemPlanMode mode = *parsedMode; + ReserveBufferPlans reservePlans; - if (this->memMode == MemPlanMode::LOCAL_MEM_PLAN && + if (mode == MemPlanMode::LOCAL_MEM_PLAN && failed(analyzeReserveBufferPlans(funcOp, reservePlans))) { return signalPassFailure(); } - if (this->memMode == MemPlanMode::LOCAL_MEM_PLAN) { + if (mode == MemPlanMode::LOCAL_MEM_PLAN) { for (ReserveBufferPlan &reservePlan : reservePlans) { if (reservePlan.mode != ReserveBufferMode::Manual) continue; @@ -2461,12 +2568,14 @@ void PlanMemoryPass::runOnOperation() { } } - MemLivenessAnalysis memLiveness(funcOp, this->memMode); + MemLivenessAnalysis memLiveness(funcOp, mode); memLiveness.build(); - MemPlan memPlan(this->memMode, this->enableGlobalReuse, - this->enablePrintMemoryAllocatedSize, - this->restrictInplaceAsISA, this->orderBySize); + constexpr bool enableGlobalReuse = false; + constexpr bool enablePrintMemoryAllocatedSize = false; + constexpr bool restrictInplaceAsISA = false; + MemPlan memPlan(mode, enableGlobalReuse, enablePrintMemoryAllocatedSize, + restrictInplaceAsISA, this->orderBySize); if (failed(memPlan.InitMemSpecsFromModule(funcOp))) { return signalPassFailure(); } @@ -2485,21 +2594,32 @@ void PlanMemoryPass::runOnOperation() { // Keep reserve_buffer allocation outside the core MemPlan algorithm: // normal local buffers are planned first, then reserve_buffer claims one // aligned hole in its target address space. - if (this->memMode == MemPlanMode::LOCAL_MEM_PLAN && - failed(assignAutoReserveBufferBases(reservePlans, memLiveness.bufferInfos, - memPlan.GetBuffer2Offsets()))) { + if (mode == MemPlanMode::LOCAL_MEM_PLAN && + failed(assignAutoReserveBufferBases( + reservePlans, memLiveness.bufferInfos, memPlan.GetBuffer2Offsets()))) { return signalPassFailure(); } RewritePatternSet patterns(&getContext()); - populateBufferAddressToAllocOp(patterns, memPlan.GetBuffer2Offsets()); + populateBufferAddressToAllocOp(patterns, mode, memPlan.GetBuffer2Offsets()); if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) { return signalPassFailure(); } + + bool hasUnplannedAllocTile = false; + funcOp.walk([&](pto::AllocTileOp op) { + if (op.getAddr()) + return; + op.emitError( + "PTOPlanMemory failed to assign an address to pto.alloc_tile"); + hasUnplannedAllocTile = true; + }); + if (hasUnplannedAllocTile) + return signalPassFailure(); } } std::unique_ptr -mlir::pto::createPlanMemoryPass(const PlanMemoryOptions &planMemoryOption) { - return std::make_unique(planMemoryOption); +mlir::pto::createPlanMemoryPass(const PlanMemoryOptions &options) { + return std::make_unique(options); } diff --git a/lib/PTO/Transforms/PTOPlanMemory.h b/lib/PTO/Transforms/PTOPlanMemory.h index 81f362be25..d6e2a37381 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.h +++ b/lib/PTO/Transforms/PTOPlanMemory.h @@ -43,6 +43,11 @@ enum BufferStatus { UNDEFFINED = 0, DEFFINED, GENED, KILLED }; /// Pair of inplace Value. using ValuePair = std::pair; +enum class MemPlanMode { + LOCAL_MEM_PLAN, + GLOBAL_WORKSPACE_PLAN, +}; + /// Result status after plan memory. enum PlanStatus { PLAN_SUCCESS = 0, diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp new file mode 100644 index 0000000000..2e8d6cd3c2 --- /dev/null +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -0,0 +1,701 @@ +// 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. + +//===- PTOPlanMemoryModern.cpp - modern static local memory planner -------===// + +#include "AllocToPointerCast.h" +#include "PTO/IR/PTOMultiBuffer.h" +#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/Transforms/Passes.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallSet.h" + +using namespace mlir; +using namespace mlir::pto; + +namespace mlir::pto { +LogicalResult runModernPlanMemory(func::FuncOp func, llvm::StringRef memMode, + bool orderBySize); +} // namespace mlir::pto + +namespace { + +struct MemSpec { + uint64_t capacityBytes = 0; + uint64_t alignmentBytes = 1; +}; + +struct RootInfo { + Value root; + Operation *defOp = nullptr; + AddressSpace space = AddressSpace::Zero; + uint64_t slotBytes = 0; + uint64_t totalBytes = 0; + uint64_t alignmentBytes = 1; + uint64_t slotCount = 1; + unsigned allocIndex = 0; + unsigned freeIndex = 0; + unsigned stableOrder = 0; + SmallVector offsets; +}; + +using RootList = SmallVector; + +static uint64_t alignUp(uint64_t value, uint64_t align) { + if (align == 0) + return value; + return ((value + align - 1) / align) * align; +} + +static std::optional getBufferAddressSpace(Type type) { + if (auto tileType = dyn_cast(type)) { + if (auto attr = + dyn_cast_or_null(tileType.getMemorySpace())) + return attr.getAddressSpace(); + return std::nullopt; + } + + if (auto memRefType = dyn_cast(type)) { + if (auto attr = + dyn_cast_or_null(memRefType.getMemorySpace())) + return attr.getAddressSpace(); + } + return std::nullopt; +} + +static bool isPlannableLocalSpace(std::optional space) { + return space && *space != AddressSpace::GM && *space != AddressSpace::Zero; +} + +static MemSpec getMemSpec(PTOArch arch, AddressSpace space) { + switch (space) { + case AddressSpace::VEC: + return {arch == PTOArch::A5 ? 253952ull : 196608ull, 256}; + case AddressSpace::MAT: + return {524288ull, 256}; + case AddressSpace::LEFT: + case AddressSpace::RIGHT: + return {65536ull, 4096}; + case AddressSpace::ACC: + return {arch == PTOArch::A5 ? 262144ull : 131072ull, 4096}; + case AddressSpace::BIAS: + case AddressSpace::SCALING: + return {524288ull, 256}; + case AddressSpace::GM: + case AddressSpace::Zero: + break; + } + return {}; +} + +static FailureOr computeStaticBufferBytes(Value value) { + ArrayRef shape; + Type elementType; + + if (auto tileType = dyn_cast(value.getType())) { + shape = tileType.getShape(); + elementType = tileType.getElementType(); + } else if (auto memRefType = dyn_cast(value.getType())) { + shape = memRefType.getShape(); + elementType = memRefType.getElementType(); + } else { + return failure(); + } + + uint64_t elemBytes = getPTOStorageElemByteSize(elementType); + if (elemBytes == 0) + return failure(); + + uint64_t numel = 1; + for (int64_t dim : shape) { + if (dim == ShapedType::kDynamic) + return failure(); + numel *= static_cast(dim); + } + return numel * elemBytes; +} + +static void appendUniqueRoot(RootList &roots, Value root) { + if (llvm::is_contained(roots, root)) + return; + roots.push_back(root); +} + +static RootList unionRoots(const RootList &lhs, const RootList &rhs) { + RootList result = lhs; + for (Value root : rhs) + appendUniqueRoot(result, root); + return result; +} + +struct PlannerAnalysis { + func::FuncOp func; + DenseMap valueToRoots; + SmallVector linearOps; + DenseMap opToIndex; + SmallVector roots; + DenseMap rootIndexByValue; + DenseMap semanticConflictMap; + bool failed = false; + + explicit PlannerAnalysis(func::FuncOp func) : func(func) {} + + RootList getRoots(Value value) const { + auto it = valueToRoots.find(value); + if (it == valueToRoots.end()) + return {}; + return it->second; + } + + void setRoots(Value value, const RootList &roots) { + if (roots.empty()) + return; + valueToRoots[value] = roots; + } + + void addRoot(Value value, Operation *defOp) { + if (rootIndexByValue.count(value)) + return; + + auto space = getBufferAddressSpace(value.getType()); + if (!isPlannableLocalSpace(space)) + return; + + auto bytesOr = computeStaticBufferBytes(value); + if (mlir::failed(bytesOr)) { + defOp->emitError("requires a static local buffer shape and known element " + "byte size for memory planning"); + failed = true; + return; + } + + uint64_t slotCount = 1; + if (auto attr = + defOp->getAttrOfType(pto::kPtoMultiBufferAttrName)) + slotCount = attr.getValue().getZExtValue(); + MemSpec spec = getMemSpec(getTargetArch(func), *space); + uint64_t slotBytes = alignUp(*bytesOr, spec.alignmentBytes); + + RootInfo info; + info.root = value; + info.defOp = defOp; + info.space = *space; + info.slotBytes = slotBytes; + info.totalBytes = slotBytes * slotCount; + info.alignmentBytes = spec.alignmentBytes; + info.slotCount = slotCount; + info.stableOrder = roots.size(); + roots.push_back(info); + rootIndexByValue[value] = roots.size() - 1; + valueToRoots[value] = RootList{value}; + } + + void markUse(Value value, unsigned index) { + for (Value root : getRoots(value)) { + auto found = rootIndexByValue.find(root); + if (found == rootIndexByValue.end()) + continue; + RootInfo &info = roots[found->second]; + if (index < info.allocIndex) + info.allocIndex = index; + info.freeIndex = std::max(info.freeIndex, index); + } + } + + void addSemanticConflict(Value a, Value b) { + if (a == b) + return; + if (!rootIndexByValue.count(a) || !rootIndexByValue.count(b)) + return; + appendUniqueRoot(semanticConflictMap[a], b); + appendUniqueRoot(semanticConflictMap[b], a); + } + + bool hasSemanticConflict(Value a, Value b) const { + if (a == b) + return false; + auto it = semanticConflictMap.find(a); + if (it == semanticConflictMap.end()) + return false; + return llvm::is_contained(it->second, b); + } + + void recordLocalOpConflicts(Operation *op) { + RootList operandRoots; + RootList resultRoots; + for (Value operand : op->getOperands()) + operandRoots = unionRoots(operandRoots, getRoots(operand)); + for (Value result : op->getResults()) + resultRoots = unionRoots(resultRoots, getRoots(result)); + + if (operandRoots.empty() || resultRoots.empty()) + return; + for (Value lhs : operandRoots) { + for (Value rhs : resultRoots) + addSemanticConflict(lhs, rhs); + } + } + + void seedForIterArgAliases(scf::ForOp forOp) { + if (forOp.getRegion().empty()) + return; + Block &body = forOp.getRegion().front(); + for (auto [iterArg, initArg] : + llvm::zip(body.getArguments().drop_front(1), forOp.getInitArgs())) { + setRoots(iterArg, getRoots(initArg)); + } + } + + void walkRegion(Region ®ion) { + for (Block &block : region) { + for (Operation &opRef : block) { + Operation *op = &opRef; + unsigned index = linearOps.size(); + linearOps.push_back(op); + opToIndex[op] = index; + + if (auto allocTile = dyn_cast(op)) { + if (!allocTile.getAddr()) { + addRoot(allocTile.getResult(), op); + if (failed) + return; + auto found = rootIndexByValue.find(allocTile.getResult()); + if (found != rootIndexByValue.end()) { + roots[found->second].allocIndex = index; + roots[found->second].freeIndex = index; + } + } + } else if (auto alloc = dyn_cast(op)) { + addRoot(alloc.getResult(), op); + if (failed) + return; + auto found = rootIndexByValue.find(alloc.getResult()); + if (found != rootIndexByValue.end()) { + roots[found->second].allocIndex = index; + roots[found->second].freeIndex = index; + } + } + + if (auto bind = dyn_cast(op)) { + setRoots(bind.getResult(), getRoots(bind.getSource())); + } else if (auto slotMarker = dyn_cast(op)) { + setRoots(slotMarker.getResult(), getRoots(slotMarker.getSource())); + } else if (auto select = dyn_cast(op)) { + setRoots(select.getResult(), + unionRoots(getRoots(select.getTrueValue()), + getRoots(select.getFalseValue()))); + } else if (auto castOp = dyn_cast(op)) { + setRoots(castOp.getResult(), getRoots(castOp.getSource())); + } else if (auto subview = dyn_cast(op)) { + setRoots(subview.getResult(), getRoots(subview.getSource())); + } else if (auto reinterpret = dyn_cast(op)) { + setRoots(reinterpret.getResult(), getRoots(reinterpret.getSource())); + } else if (auto forOp = dyn_cast(op)) { + seedForIterArgAliases(forOp); + } + + for (Value operand : op->getOperands()) + markUse(operand, index); + + for (Region &nested : op->getRegions()) { + walkRegion(nested); + if (failed) + return; + } + + if (auto ifOp = dyn_cast(op)) { + if (ifOp.getNumResults() != 0) { + auto thenYield = + cast(ifOp.thenBlock()->getTerminator()); + auto elseYield = + cast(ifOp.elseBlock()->getTerminator()); + for (auto [result, thenVal, elseVal] : + llvm::zip(ifOp.getResults(), thenYield.getResults(), + elseYield.getResults())) { + setRoots(result, + unionRoots(getRoots(thenVal), getRoots(elseVal))); + } + } + } else if (auto forOp = dyn_cast(op)) { + auto yieldOp = cast(forOp.getBody()->getTerminator()); + for (auto [result, initArg, yielded] : + llvm::zip(forOp.getResults(), forOp.getInitArgs(), + yieldOp.getResults())) { + setRoots(result, unionRoots(getRoots(initArg), getRoots(yielded))); + } + } + + recordLocalOpConflicts(op); + } + } + } +}; + +static bool lifetimesOverlap(const RootInfo &lhs, const RootInfo &rhs) { + return !(lhs.freeIndex < rhs.allocIndex || rhs.freeIndex < lhs.allocIndex); +} + +static bool intervalsOverlap(uint64_t lhsOffset, uint64_t lhsSize, + uint64_t rhsOffset, uint64_t rhsSize) { + return lhsOffset < rhsOffset + rhsSize && rhsOffset < lhsOffset + lhsSize; +} + +static SmallVector buildSlotOffsets(uint64_t base, uint64_t slotBytes, + uint64_t slotCount) { + SmallVector offsets; + offsets.reserve(slotCount); + for (uint64_t slot = 0; slot < slotCount; ++slot) + offsets.push_back(base + slot * slotBytes); + return offsets; +} + +static FailureOr planReserveBufferBase( + pto::ReserveBufferOp reserveOp, const MemSpec &spec, + SmallVectorImpl> &occupied) { + uint64_t sizeBytes = + alignUp(static_cast(reserveOp.getSize()), spec.alignmentBytes); + llvm::sort(occupied, [](const auto &lhs, const auto &rhs) { + return lhs.first < rhs.first; + }); + + SmallVector> merged; + for (const auto &interval : occupied) { + if (merged.empty() || interval.first > merged.back().second) { + merged.push_back(interval); + continue; + } + merged.back().second = std::max(merged.back().second, interval.second); + } + + uint64_t cursor = 0; + for (const auto &interval : merged) { + cursor = alignUp(cursor, spec.alignmentBytes); + if (cursor + sizeBytes <= interval.first) + return cursor; + cursor = std::max(cursor, interval.second); + } + cursor = alignUp(cursor, spec.alignmentBytes); + if (cursor + sizeBytes > spec.capacityBytes) + return failure(); + occupied.push_back({cursor, cursor + sizeBytes}); + return cursor; +} + +static LogicalResult +validateManualReserveBufferBase(pto::ReserveBufferOp reserveOp, + const MemSpec &spec) { + auto baseAttr = reserveOp.getBaseAttr(); + if (!baseAttr) + return reserveOp.emitError("expects 'base' when 'auto' is false"); + + int64_t signedBase = baseAttr.getInt(); + if (signedBase < 0) + return reserveOp.emitError( + "expects 'base' to be non-negative when present"); + + uint64_t base = static_cast(signedBase); + if (base % spec.alignmentBytes != 0) { + return reserveOp.emitError("expects 'base' to be aligned to ") + << spec.alignmentBytes << " bytes for " + << stringifyEnum(reserveOp.getLocation().getAddressSpace()); + } + + uint64_t size = static_cast(reserveOp.getSize()); + if (base > spec.capacityBytes || size > spec.capacityBytes - base) { + return reserveOp.emitError("reserved range exceeds ") + << stringifyEnum(reserveOp.getLocation().getAddressSpace()) + << " capacity: base " << base << " + size " << size << " > " + << spec.capacityBytes << " bytes"; + } + + return success(); +} + +static LogicalResult +validateReserveBufferForPlanMemory(pto::ReserveBufferOp reserveOp, + PTOArch arch) { + MemSpec spec = getMemSpec(arch, reserveOp.getLocation().getAddressSpace()); + + if (reserveOp.getAutoAlloc()) { + if (reserveOp.getBaseAttr()) { + return reserveOp.emitError( + "unexpected pre-populated 'base' on auto reserve_buffer before " + "pto-plan-memory; omit 'base' and let pto-plan-memory assign it"); + } + return success(); + } + + if (mlir::failed(validateManualReserveBufferBase(reserveOp, spec))) + return failure(); + + return reserveOp.emitError( + "pto.reserve_buffer with explicit 'base' (auto = false) is not " + "supported in PlanMemory; use --pto-level=level3 or set auto = true"); +} + +class AllocTileOpAddPlannedAddressPattern + : public OpRewritePattern { +public: + explicit AllocTileOpAddPlannedAddressPattern( + MLIRContext *context, + DenseMap> buffer2Offsets) + : OpRewritePattern(context), + buffer2Offsets(std::move(buffer2Offsets)) {} + + LogicalResult matchAndRewrite(pto::AllocTileOp op, + PatternRewriter &rewriter) const override { + if (op.getAddr()) + return failure(); + + auto tileType = dyn_cast(op.getResult().getType()); + if (!tileType) + return failure(); + + auto it = buffer2Offsets.find(op.getResult()); + if (it == buffer2Offsets.end() || it->second.empty()) + return failure(); + + if (it->second.size() != 1) { + return rewriter.notifyMatchFailure( + op, "single alloc_tile root expects exactly one planned address"); + } + + Value addr = rewriter.create(op.getLoc(), + it->second.front(), 64); + auto planned = rewriter.create( + op.getLoc(), tileType, addr, + op.getValidRow() ? op.getValidRow() : Value(), + op.getValidCol() ? op.getValidCol() : Value()); + for (NamedAttribute attr : op->getAttrs()) { + if (attr.getName().getValue() == "operandSegmentSizes") + continue; + planned->setAttr(attr.getName(), attr.getValue()); + } + + rewriter.replaceOp(op, planned.getResult()); + return success(); + } + +private: + DenseMap> buffer2Offsets; +}; + +static LogicalResult materializePlannedOffsets( + func::FuncOp func, DenseMap> buffer2Offsets) { + RewritePatternSet patterns(func.getContext()); + patterns.add(patterns.getContext(), + buffer2Offsets); + patterns.add(patterns.getContext(), + buffer2Offsets); + if (mlir::failed(applyPatternsGreedily(func, std::move(patterns)))) + return failure(); + return success(); +} + +} // namespace + +LogicalResult mlir::pto::runModernPlanMemory(func::FuncOp func, + llvm::StringRef memMode, + bool orderBySize) { + if (!memMode.equals_insensitive("local")) { + func.emitError("unsupported mem-mode '") + << memMode << "'; only 'local' is currently implemented"; + return failure(); + } + + PlannerAnalysis analysis(func); + analysis.walkRegion(func.getBody()); + if (analysis.failed) { + return failure(); + } + + DenseMap> buffer2Offsets; + llvm::MapVector> rootsBySpace; + for (RootInfo &info : analysis.roots) + rootsBySpace[info.space].push_back(&info); + + for (auto &entry : rootsBySpace) { + AddressSpace space = entry.first; + SmallVector &roots = entry.second; + MemSpec spec = getMemSpec(getTargetArch(func), space); + + llvm::stable_sort(roots, [&](const RootInfo *lhs, const RootInfo *rhs) { + if (orderBySize && lhs->totalBytes != rhs->totalBytes) + return lhs->totalBytes > rhs->totalBytes; + if (lhs->allocIndex != rhs->allocIndex) + return lhs->allocIndex < rhs->allocIndex; + return lhs->stableOrder < rhs->stableOrder; + }); + + SmallVector planned; + uint64_t scopeRequiredBytes = 0; + for (RootInfo *info : roots) { + llvm::SmallSet candidateSet; + candidateSet.insert(0); + for (RootInfo *other : planned) { + candidateSet.insert(other->offsets.front()); + candidateSet.insert( + alignUp(other->offsets.front() + other->totalBytes, + info->alignmentBytes)); + } + + SmallVector candidates(candidateSet.begin(), + candidateSet.end()); + llvm::sort(candidates); + + std::optional chosen; + for (uint64_t candidate : candidates) { + candidate = alignUp(candidate, info->alignmentBytes); + if (candidate + info->totalBytes > spec.capacityBytes) + continue; + + bool conflict = false; + for (RootInfo *other : planned) { + if (!intervalsOverlap(candidate, info->totalBytes, + other->offsets.front(), other->totalBytes)) + continue; + if (lifetimesOverlap(*info, *other) || + analysis.hasSemanticConflict(info->root, other->root)) { + conflict = true; + break; + } + } + if (!conflict) { + chosen = candidate; + break; + } + } + + if (!chosen) { + uint64_t fallback = 0; + for (RootInfo *other : planned) + fallback = + std::max(fallback, other->offsets.front() + other->totalBytes); + fallback = alignUp(fallback, info->alignmentBytes); + chosen = fallback; + } + + info->offsets = + buildSlotOffsets(*chosen, info->slotBytes, info->slotCount); + scopeRequiredBytes = + std::max(scopeRequiredBytes, *chosen + info->totalBytes); + planned.push_back(info); + } + + if (scopeRequiredBytes > spec.capacityBytes) { + func.emitError() << stringifyEnum(space) << " overflow, requires " + << (scopeRequiredBytes * 8) << " bits while " + << (spec.capacityBytes * 8) << " bits available"; + return failure(); + } + + for (RootInfo *info : roots) + buffer2Offsets[info->root] = info->offsets; + } + + DenseMap>> + occupiedBySpace; + for (const RootInfo &info : analysis.roots) { + if (info.offsets.empty()) + continue; + occupiedBySpace[info.space].push_back( + {info.offsets.front(), info.offsets.front() + info.totalBytes}); + } + + MLIRContext *ctx = func.getContext(); + PTOArch arch = getTargetArch(func); + func.walk([&](pto::ReserveBufferOp reserveOp) -> WalkResult { + if (mlir::failed(validateReserveBufferForPlanMemory(reserveOp, arch))) { + analysis.failed = true; + return WalkResult::interrupt(); + } + + AddressSpace space = reserveOp.getLocation().getAddressSpace(); + MemSpec spec = getMemSpec(arch, space); + auto baseOr = + planReserveBufferBase(reserveOp, spec, occupiedBySpace[space]); + if (mlir::failed(baseOr)) { + reserveOp.emitError("failed to reserve a local memory hole for " + "reserve_buffer"); + analysis.failed = true; + return WalkResult::interrupt(); + } + + reserveOp->setAttr("base", + IntegerAttr::get(IntegerType::get(ctx, 32), + static_cast(*baseOr))); + return WalkResult::advance(); + }); + + if (analysis.failed || + mlir::failed(materializePlannedOffsets(func, buffer2Offsets))) { + return failure(); + } + + bool hasUnplannedAllocTile = false; + func.walk([&](pto::AllocTileOp op) { + if (op.getAddr()) + return; + op.emitError( + "PTOPlanMemory failed to assign an address to pto.alloc_tile"); + hasUnplannedAllocTile = true; + }); + if (hasUnplannedAllocTile) + return failure(); + return success(); +} + +namespace { +struct PlanMemoryModernPass + : public PassWrapper> { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PlanMemoryModernPass) + + PlanMemoryModernPass() = default; + explicit PlanMemoryModernPass(const PlanMemoryOptions &options) + : memMode(options.memMode), orderBySize(options.orderBySize) {} + + StringRef getArgument() const final { return "pto-plan-memory"; } + StringRef getDescription() const final { + return "Plan local memory using the modern PTO planner"; + } + StringRef getName() const final { return "PlanMemoryModern"; } + + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + registry.insert(); + } + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + for (func::FuncOp funcOp : moduleOp.getOps()) { + if (failed(mlir::pto::runModernPlanMemory(funcOp, memMode, + orderBySize))) { + signalPassFailure(); + return; + } + } + } + +private: + std::string memMode = "local"; + bool orderBySize = false; +}; +} // namespace + +std::unique_ptr +mlir::pto::createPlanMemoryModernPass(const PlanMemoryOptions &options) { + return std::make_unique(options); +} diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index bd0cc4e228..fd3fd3edce 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -65,6 +65,18 @@ static void markForceDynamicValidShape(Operation *op, bool force, static Type convertPTOTypeToMemRef(Type t); +static bool hasTileNativeAllocTileRoot(func::FuncOp func) { + bool found = false; + auto result = func.walk([&](pto::AllocTileOp op) { + if (!isa(op.getResult().getType())) + return WalkResult::advance(); + found = true; + return WalkResult::interrupt(); + }); + (void)result; + return found; +} + constexpr size_t kTileRank2D = 2; constexpr size_t kRowDimensionIndex = 0; constexpr size_t kColumnDimensionIndex = 1; @@ -939,63 +951,6 @@ static void markForceDynamicValidShape(Operation *op, bool force, func.setFunctionType(FunctionType::get(ctx, newInputs, newResults)); } -[[maybe_unused]] static LogicalResult lowerAllocTileOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector allocTiles; - func.walk([&](mlir::pto::AllocTileOp op) { allocTiles.push_back(op); }); - - for (auto op : allocTiles) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - auto tbTy = dyn_cast(op.getResult().getType()); - if (!tbTy) - continue; - - SmallInlineVector shape(tbTy.getShape().begin(), - tbTy.getShape().end()); - Type elemTy = tbTy.getElementType(); - SmallVector strides = buildTileMemRefStrides(tbTy); - - auto targetLayout = - StridedLayoutAttr::get(ctx, ShapedType::kDynamic, strides); - auto targetType = - MemRefType::get(shape, elemTy, targetLayout, tbTy.getMemorySpace()); - - Value vRow = op.getValidRow(); - Value vCol = op.getValidCol(); - materializeStaticValidDims(rewriter, loc, tbTy, vRow, vCol); - - auto configAttr = tbTy.getConfigAttr(); - if (!configAttr) - configAttr = pto::TileBufConfigAttr::getDefault(ctx); - - if (Value addr = op.getAddr()) { - auto pc = rewriter.create( - loc, targetType, ValueRange{addr}, vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(pc, tbTy.hasDynamicValid(), ctx); - auto bindOp = rewriter.create( - loc, targetType, pc.getResult(), vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - rewriter.replaceOp(op, bindOp.getResult()); - continue; - } - - auto allocLayout = StridedLayoutAttr::get(ctx, 0, strides); - auto allocType = - MemRefType::get(shape, elemTy, allocLayout, tbTy.getMemorySpace()); - Value alloc = rewriter.create(loc, allocType); - auto bindOp = rewriter.create( - loc, targetType, alloc, vRow ? vRow : Value(), vCol ? vCol : Value(), - configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - rewriter.replaceOp(op, bindOp.getResult()); - } - return success(); -} - [[maybe_unused]] static LogicalResult lowerDeclareTileOps(func::FuncOp func, MLIRContext *ctx) { DefaultInlineVector declaredTiles; func.walk([&](mlir::pto::DeclareTileOp op) { declaredTiles.push_back(op); }); @@ -1250,6 +1205,40 @@ static LogicalResult lowerPtrLikeTileBufAddrOps(func::FuncOp func, return success(); } +static LogicalResult bridgeMemRefOperandsToExternalPtrCallees(ModuleOp mod, + func::FuncOp func, + MLIRContext *ctx) { + SmallVector callOps; + func.walk([&](func::CallOp callOp) { callOps.push_back(callOp); }); + + for (auto callOp : callOps) { + auto callee = mod.lookupSymbol(callOp.getCallee()); + if (!callee || !callee.isExternal()) + continue; + + FunctionType calleeType = callee.getFunctionType(); + if (calleeType.getNumInputs() != callOp.getNumOperands()) { + callOp.emitError("callee signature does not match call operands"); + return failure(); + } + + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(callOp); + for (auto [index, expectedType] : + llvm::enumerate(calleeType.getInputs())) { + Value operand = callOp.getOperand(index); + if (!isa(expectedType) || + !isa(operand.getType())) + continue; + + auto castOp = rewriter.create( + callOp.getLoc(), expectedType, operand); + callOp->setOperand(index, castOp.getResult()); + } + } + return success(); +} + [[maybe_unused]] static LogicalResult lowerTensorViewDimOps(func::FuncOp func, MLIRContext *ctx) { DefaultInlineVector tvDims; func.walk([&](mlir::pto::GetTensorViewDimOp op) { tvDims.push_back(op); }); @@ -1613,6 +1602,12 @@ struct PTOViewToMemrefPass return; } + // Keep external declarations in the authored ABI. PTOViewToMemref does + // not rewrite func.call operands for external/private helpers, so changing + // only the callee type can leave pointer-like call users mismatched. + if (func.isExternal()) + continue; + auto fnTy = func.getFunctionType(); // ------------------------------------------------------------------ @@ -1625,7 +1620,6 @@ struct PTOViewToMemrefPass for (Type t : fnTy.getResults()) newResults.push_back(convertPTOTypeToMemRef(t)); func.setFunctionType(FunctionType::get(ctx, newInputs, newResults)); - if (func.isExternal()) continue; Block &entry = func.front(); @@ -1636,6 +1630,11 @@ struct PTOViewToMemrefPass } } + if (failed(bridgeMemRefOperandsToExternalPtrCallees(mod, func, ctx))) { + signalPassFailure(); + return; + } + // ------------------------------------------------------------------ // Stage 0.20: lower pto.inttoptr result types to GM memrefs. // ------------------------------------------------------------------ @@ -1696,97 +1695,12 @@ struct PTOViewToMemrefPass } // ------------------------------------------------------------------ - // Stage 0.5: lower pto.alloc_tile -> memref.alloc + pto.bind_tile + // Stage 0.5: keep pto.alloc_tile as a tile-native allocation root. + // + // Explicit-address alloc_tile already carries the physical address for + // level3. No-address alloc_tile is assigned an address by PTOPlanMemory. + // Neither case needs a pointer_cast + bind_tile detour here. // ------------------------------------------------------------------ - DefaultInlineVector allocTiles; - func.walk([&](mlir::pto::AllocTileOp op) { allocTiles.push_back(op); }); - - for (auto op : allocTiles) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - auto tbTy = dyn_cast(op.getResult().getType()); - if (!tbTy) continue; - - // 1. 获取 Shape 和 ElementType - SmallInlineVector shape(tbTy.getShape().begin(), tbTy.getShape().end()); - Type elemTy = tbTy.getElementType(); - - // 2. 计算 Strides (layout-aware when possible) - SmallVector strides; - TileLayoutInfo info; - if (computeTileLayoutInfo(tbTy.getConfigAttr(), elemTy, shape, info)) { - strides = {info.rowStride, info.colStride}; - } else { - strides.resize(shape.size()); - int64_t s = 1; - for (int i = (int)shape.size() - 1; i >= 0; --i) { - strides[i] = s; - if (shape[i] != ShapedType::kDynamic) s *= shape[i]; - } - } - - // 3. 构造 [BindTile 输出] 的动态类型 (Offset: ?) - // 这必须与 convertPTOTypeToMemRef 返回的类型一致，以便与 Subview 兼容 - auto targetLayout = - StridedLayoutAttr::get(ctx, ShapedType::kDynamic, strides); // offset = ? - auto targetType = - MemRefType::get(shape, elemTy, targetLayout, tbTy.getMemorySpace()); - - // 4. Preserve tile valid dims (v_row / v_col). - // - // `pto.alloc_tile` encodes the valid shape in the result TileBufType - // (e.g. acc tile may be rows=16 but v_row=1). The alloc op itself does - // not necessarily carry explicit operands for static valid dims, so we - // must materialize them from the type to keep them through - // tile_buf -> memref lowering. - // - // For dynamically valid tiles (validShape == [-1, -1]), preserve the - // runtime operands if present. - Value vRow = op.getValidRow(); - Value vCol = op.getValidCol(); - // TileBuf valid dims use a negative sentinel (e.g. '?' / -1), which is - // distinct from MLIR's ShapedType::kDynamic (INT64_MIN). Treat any - // negative value as dynamic here. - materializeStaticValidDims(rewriter, loc, tbTy, vRow, vCol); - - // 5. 获取 Config (保持不变) - auto configAttr = tbTy.getConfigAttr(); - if (!configAttr) configAttr = pto::TileBufConfigAttr::getDefault(ctx); - - // 6. If alloc_tile provides an explicit address, keep the original - // pointer_cast lowering intact and additionally rebind through - // pto.bind_tile. PointerCastOp continues to carry the tile metadata - // used by existing lowering paths, while BindTileOp provides the - // unified anchor EmitC uses to recover tile_buf information. - if (Value addr = op.getAddr()) { - auto pc = rewriter.create( - loc, targetType, ValueRange{addr}, vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(pc, tbTy.hasDynamicValid(), ctx); - auto bindOp = rewriter.create( - loc, targetType, pc.getResult(), vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - rewriter.replaceOp(op, bindOp.getResult()); - continue; - } - - // 7. Otherwise allocate a concrete memref buffer and bind tile. - // memref.alloc 要求明确的 layout，不能是动态 offset。 - auto allocLayout = StridedLayoutAttr::get(ctx, 0, strides); // offset = 0 - auto allocType = MemRefType::get(shape, elemTy, allocLayout, tbTy.getMemorySpace()); - Value alloc = rewriter.create(loc, allocType); - - // BindTileOp 的 Builder 会自动处理空的 Value，将其视为静态维度 - auto bindOp = rewriter.create( - loc, targetType, alloc, vRow ? vRow : Value(), vCol ? vCol : Value(), - configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - - rewriter.replaceOp(op, bindOp.getResult()); - } // ------------------------------------------------------------------ // Stage 0.6: lower pto.alloc_multi_tile / pto.multi_tile_get @@ -1797,13 +1711,13 @@ struct PTOViewToMemrefPass // %m = pto.bind_tile %a, ... : memref -> memref // and replace all uses of the multi_tile_buf SSA with %m. The N-way // physical fan-out lives on the `pto.multi_buffer` attr and is - // materialized later by PTOPlanMemory. + // materialized later by the fallback address resolver. // // multi_tile_get consumes that memref and wraps it in // %slot = pto.slot_marker %m[%k] : memref -> memref // %t = pto.bind_tile %slot, ... : memref -> memref // The slot_marker carries the user-supplied slot SSA forward so that - // PlanMemory / sync / EnableBufferSelect can identify which physical + // sync analysis and buffer-select lowering can identify which physical // slot this use refers to. // // Ordering: alloc_multi_tile must be lowered before multi_tile_get so @@ -1873,13 +1787,12 @@ struct PTOViewToMemrefPass if (!configAttr) configAttr = pto::TileBufConfigAttr::getDefault(ctx); - // Level3 (caller-owned memory): an explicit base address is given and - // PlanMemory is skipped. Lay the N slots out contiguously as + // Level3 (caller-owned memory): an explicit base address is given. + // Lay the N slots out contiguously as // [addr, addr + slotBytes, ..., addr + (N-1)*slotBytes] and emit the - // multi-address `pto.pointer_cast` alloc anchor directly -- the same - // shape PlanMemory produces in the default pipeline (addrs are kept in - // slot order so PTOResolveBufferSelect can pick addrs[slot]). Sync and - // buffer-select then run unchanged. + // multi-address `pto.pointer_cast` alloc anchor directly (addrs are + // kept in slot order so PTOResolveBufferSelect can pick addrs[slot]). + // Sync and buffer-select then run unchanged. if (Value addr = op.getAddr()) { uint32_t n = mtbTy.getCount(); int64_t elemBytes = getElemBytes(elemTy); @@ -1927,7 +1840,7 @@ struct PTOViewToMemrefPass } // memref.alloc with N-slot annotation. The actual N-way address - // expansion happens in PlanMemory. + // expansion happens in PTOPlanMemory. auto allocLayout = StridedLayoutAttr::get(ctx, 0, strides); auto allocType = MemRefType::get(shape, elemTy, allocLayout, tbTy.getMemorySpace()); @@ -2406,6 +2319,7 @@ struct PTOViewToMemrefPass // Stage 3: Rewrite Compute Ops // [关键] 全面使用 op->getOperand(i) 避免 Typed Accessor Crash // ------------------------------------------------------------------ + if (!hasTileNativeAllocTileRoot(func)) { // --- TLoadOp [Src, Dst] --- DefaultInlineVector loads; @@ -4280,6 +4194,9 @@ struct PTOViewToMemrefPass auto srcTy = dyn_cast(src.getType()); auto tmpTy = tmp ? dyn_cast(tmp.getType()) : MemRefType(); if (!srcTy || (tmp && !tmpTy)) { + if (isa(src.getType()) && + (!tmp || isa(tmp.getType()))) + continue; op.emitError("ins/outs are not memref yet"); signalPassFailure(); return; @@ -4293,6 +4210,7 @@ struct PTOViewToMemrefPass dyn_cast_or_null( op.getProperties().printFormat)); } + } // ------------------------------------------------------------------ // Stage 4: Reconcile control-flow result types diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index 44b6d1005f..7c8de83043 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -131,6 +131,13 @@ void setBaseMemRefTypeScope(Value val, AddressSpaceAttr targetMemScope) { std::optional GetBufferSpaceAttr(Value operand) { + if (auto tileTy = dyn_cast(operand.getType())) { + if (auto memorySpaceAttr = + dyn_cast_or_null(tileTy.getMemorySpace())) + return memorySpaceAttr; + return std::nullopt; + } + if (!llvm::isa(operand.getType())) { return std::nullopt; } diff --git a/ptodsl/tests/test_docs_as_test.py b/ptodsl/tests/test_docs_as_test.py index 6520d2f296..9d1feda740 100644 --- a/ptodsl/tests/test_docs_as_test.py +++ b/ptodsl/tests/test_docs_as_test.py @@ -67,7 +67,7 @@ class DocTestDirective: symbol: Optional[str] = None compile_kwargs: Optional[dict[str, object]] = None fixture: Optional[str] = None - files: Optional[dict[str, str]] | None = None + files: Optional[dict[str, str]] = None @dataclass(frozen=True) @@ -324,7 +324,7 @@ def parse_test_directive(block: MarkdownCodeBlock) -> DocTestDirective: return DocTestDirective(mode=mode) -def _write_directive_files(snippet_dir: Path, files: dict[str, str] | None) -> None: +def _write_directive_files(snippet_dir: Path, files: Optional[dict[str, str]]) -> None: if not files: return for relative_path, text in files.items(): @@ -351,7 +351,7 @@ def execute_source( symbol: Optional[str] = None, *, extra_namespace: Optional[dict[str, object]] = None, - source_dir: Path | None = None, + source_dir: Optional[Path] = None, ) -> dict[str, object]: source_file = block.path if source_dir is None else source_dir / "case.py" namespace: dict[str, object] = { diff --git a/ptodsl/tests/test_ptoas_frontend_verify.py b/ptodsl/tests/test_ptoas_frontend_verify.py index 57592a8355..47d4f52320 100644 --- a/ptodsl/tests/test_ptoas_frontend_verify.py +++ b/ptodsl/tests/test_ptoas_frontend_verify.py @@ -451,12 +451,13 @@ def main() -> None: "ptr_like_tile_buf_addr_probe frontend verification output should preserve the kernel symbol", ) expect( - "memref address view during PTOViewToMemref", + "pto.alloc_tile addr" in ptr_like_addr_frontend_text + and "pto.tile_buf_addr" in ptr_like_addr_frontend_text, + "ptr-like tile_buf_addr lowering should keep the tile-native address path until PTOPlanMemory materializes alloc_tile addr", ) expect( "call @consume" in ptr_like_addr_frontend_text, - "ptr-like tile_buf_addr lowering should preserve call users after converting pointer-like operands", + "ptr-like tile_buf_addr lowering should preserve call users without half-converting pointer-like operands", ) example_mlir_text = emit_example_mlir(mixed_backend_example) diff --git a/test/lit/pto/alloc_tile_addr_alignment_invalid.pto b/test/lit/pto/alloc_tile_addr_alignment_invalid.pto index 2a94e206e0..a211d08e2f 100644 --- a/test/lit/pto/alloc_tile_addr_alignment_invalid.pto +++ b/test/lit/pto/alloc_tile_addr_alignment_invalid.pto @@ -1,4 +1,5 @@ // RUN: not ptoas --pto-level=level3 %s 2>&1 | FileCheck %s +// RUN: not ptoas --pto-level=level2 %s 2>&1 | FileCheck %s module { func.func @unaligned_alloc_tile_addr() { diff --git a/test/lit/pto/alloc_tile_addr_level12_reject.pto b/test/lit/pto/alloc_tile_addr_level12_reject.pto new file mode 100644 index 0000000000..c40ff6985e --- /dev/null +++ b/test/lit/pto/alloc_tile_addr_level12_reject.pto @@ -0,0 +1,21 @@ +// 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-level=level1 %s 2>&1 | FileCheck %s +// RUN: not ptoas --pto-level=level2 %s 2>&1 | FileCheck %s + +module { + func.func @reject_explicit_addr_below_level3() { + %addr = arith.constant 0 : i64 + %tile = pto.alloc_tile addr = %addr + : !pto.tile_buf + return + } +} + +// CHECK: error: unexpected 'addr' operand: only supported when --pto-level=level3 diff --git a/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto b/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto new file mode 100644 index 0000000000..55a835120e --- /dev/null +++ b/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto @@ -0,0 +1,35 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level2 --pto-arch=a3 --emit-pto-ir \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=VIEW +// RUN: ptoas --pto-level=level2 --pto-arch=a3 --emit-pto-ir \ +// RUN: --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN +// RUN: ptoas --pto-level=level2 --pto-arch=a3 %s >/dev/null + +module attributes {"pto.target_arch" = "a3"} { + func.func @alloc_tile_plan_memory_no_memref_alloc() { + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tmov ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// VIEW-LABEL: func.func @alloc_tile_plan_memory_no_memref_alloc +// VIEW: pto.alloc_tile +// VIEW: pto.alloc_tile +// VIEW-NOT: memref.alloc + +// PLAN-LABEL: func.func @alloc_tile_plan_memory_no_memref_alloc +// PLAN: %[[ADDR:.*]] = arith.constant 0 : i64 +// PLAN-NEXT: {{%.*}} = pto.alloc_tile addr = %[[ADDR]] +// PLAN-NOT: pto.pointer_cast +// PLAN-NOT: pto.bind_tile +// PLAN-NOT: memref.alloc diff --git a/test/lit/pto/multi_tile_get_const_slot_lowering.pto b/test/lit/pto/multi_tile_get_const_slot_lowering.pto index e8f88edc62..65c6c3430f 100644 --- a/test/lit/pto/multi_tile_get_const_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_const_slot_lowering.pto @@ -14,7 +14,7 @@ // 1. PTOViewToMemref lowers `alloc_multi_tile` to `memref.alloc` with the // `pto.multi_buffer = N : i32` attribute, and `multi_tile_get [%k]` to // `pto.slot_marker` carrying the constant slot SSA. -// 2. PTOPlanMemory reserves N physical addresses and emits a multi- +// 2. PTOPlanMemory assigns planned addresses and emits a multi- // address `pto.pointer_cast(addr0, addr1)`. // 3. PTOResolveBufferSelect lowers each `slot_marker` to a single- // address `pto.pointer_cast(addrK)` -- one per slot use. @@ -50,7 +50,7 @@ module { // CHECK: pto.slot_marker %{{.*}}[%c0 // CHECK: pto.slot_marker %{{.*}}[%c1 -// Stage 2 -- PlanMemory: two physical slots reserved -> 2-addr pointer_cast. +// Stage 2 -- PTOPlanMemory: two physical slots reserved -> 2-addr pointer_cast. // PLAN: IR Dump After PlanMemory // PLAN-LABEL: func.func @const_slot_two_buffers // PLAN: pto.pointer_cast({{.*}}, {{.*}}) diff --git a/test/lit/pto/multi_tile_level3_explicit_addr.pto b/test/lit/pto/multi_tile_level3_explicit_addr.pto index 9efaade82e..37c1828ea8 100644 --- a/test/lit/pto/multi_tile_level3_explicit_addr.pto +++ b/test/lit/pto/multi_tile_level3_explicit_addr.pto @@ -8,15 +8,14 @@ // RUN: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s // RUN: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT -// RUN: not ptoas --pto-level=level2 --pto-arch=a5 %s 2>&1 | FileCheck %s --check-prefix=LV2ERR +// RUN: not ptoas --pto-level=level1 --pto-arch=a5 %s 2>&1 | FileCheck %s --check-prefix=LVERR +// RUN: not ptoas --pto-level=level2 --pto-arch=a5 %s 2>&1 | FileCheck %s --check-prefix=LVERR // Under --pto-level=level3 the caller owns local memory, PTOPlanMemory does not // run, and `pto.alloc_multi_tile` carries an explicit base `addr`. The view -// lowering must fan the base address out into the multi-address -// `pto.pointer_cast` anchor PlanMemory would otherwise produce -- slot k at -// `addr + k * slotBytes` (slotBytes = 16*16*2 = 512 here) -- so sync and -// PTOResolveBufferSelect keep working. Without this the slots would all alias -// the same memref (silent single-buffer downgrade). +// lowering fans it out into a multi-address `pto.pointer_cast` anchor. Slot k is +// placed at `addr + k * slotBytes` (slotBytes = 16*16*2 = 512 here), so sync and +// PTOResolveBufferSelect keep working without automatic address planning. module { func.func @l3_two_buffers( @@ -63,4 +62,4 @@ module { // SELECT-NOT: pto.slot_marker // Outside level3 an explicit addr on alloc_multi_tile is rejected. -// LV2ERR: unexpected 'addr' operand on pto.alloc_multi_tile: only supported when --pto-level=level3 +// LVERR: unexpected 'addr' operand on pto.alloc_multi_tile: only supported when --pto-level=level3 diff --git a/test/lit/pto/multi_tile_level3_missing_addr.pto b/test/lit/pto/multi_tile_level3_missing_addr.pto index 9784b9db9d..f581df3ebd 100644 --- a/test/lit/pto/multi_tile_level3_missing_addr.pto +++ b/test/lit/pto/multi_tile_level3_missing_addr.pto @@ -8,9 +8,9 @@ // RUN: not ptoas --pto-level=level3 --pto-arch=a5 %s 2>&1 | FileCheck %s -// In level3 PTOPlanMemory is skipped, so multi-buffer slots have no physical -// addresses unless the caller provides a base `addr`. Reject the op instead of -// silently collapsing every slot onto one memref. +// In level3 automatic address planning is unavailable, so multi-buffer slots +// have no physical addresses unless the caller provides a base `addr`. Reject +// the op instead of silently collapsing every slot onto one memref. module { func.func @l3_missing_addr( diff --git a/test/lit/pto/plan_memory_bind_tile_alias_liveness.pto b/test/lit/pto/plan_memory_bind_tile_alias_liveness.pto index 32af387daf..ff075d063d 100644 --- a/test/lit/pto/plan_memory_bind_tile_alias_liveness.pto +++ b/test/lit/pto/plan_memory_bind_tile_alias_liveness.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @bind_tile_alias_liveness(%arg0: memref<16x16x16xf16, #pto.address_space>, @@ -30,6 +31,5 @@ module { // CHECK-NOT: memref.alloc // CHECK-DAG: %c0_i64 = arith.constant 0 : i64 // CHECK-DAG: %c8192_i64 = arith.constant 8192 : i64 -// CHECK-DAG: pto.pointer_cast(%c0_i64) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK-DAG: pto.pointer_cast(%c8192_i64) {{.*}} : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> - +// CHECK-DAG: pto.pointer_cast(%c0_i64){{.*}} : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> +// CHECK-DAG: pto.pointer_cast(%c8192_i64){{.*}} : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> diff --git a/test/lit/pto/plan_memory_for_iter_args_yield.pto b/test/lit/pto/plan_memory_for_iter_args_yield.pto index bb85772ad4..6f1578aac7 100644 --- a/test/lit/pto/plan_memory_for_iter_args_yield.pto +++ b/test/lit/pto/plan_memory_for_iter_args_yield.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @for_iter_args_yield(%arg0: memref<16x16x16xf16, #pto.address_space>, diff --git a/test/lit/pto/plan_memory_fragmentation_hole_fit.pto b/test/lit/pto/plan_memory_fragmentation_hole_fit.pto index aac3aa7640..39bb3b8494 100644 --- a/test/lit/pto/plan_memory_fragmentation_hole_fit.pto +++ b/test/lit/pto/plan_memory_fragmentation_hole_fit.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @fragmentation_hole_fit(%arg0: memref<16x16x16xf16, #pto.address_space>, diff --git a/test/lit/pto/plan_memory_fragmentation_two_holes.pto b/test/lit/pto/plan_memory_fragmentation_two_holes.pto index 04f7437e41..4a5c8be91a 100644 --- a/test/lit/pto/plan_memory_fragmentation_two_holes.pto +++ b/test/lit/pto/plan_memory_fragmentation_two_holes.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @fragmentation_two_holes(%arg0: memref<16x16x16xf16, #pto.address_space>, diff --git a/test/lit/pto/plan_memory_if_in_loop.pto b/test/lit/pto/plan_memory_if_in_loop.pto index f831840b20..1d902eedce 100644 --- a/test/lit/pto/plan_memory_if_in_loop.pto +++ b/test/lit/pto/plan_memory_if_in_loop.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @if_in_loop(%arg0: memref<16x16x16xf16, #pto.address_space>, @@ -31,7 +32,7 @@ module { // CHECK: func.func @if_in_loop // CHECK-NOT: memref.alloc // CHECK-DAG: %c0_i64 = arith.constant 0 : i64 -// CHECK-DAG: %c8192_i64 = arith.constant 8192 : i64 // CHECK: scf.for // CHECK: scf.if -// CHECK: pto.pointer_cast +// CHECK: pto.pointer_cast(%c0_i64) +// CHECK: pto.pointer_cast(%c0_i64) diff --git a/test/lit/pto/plan_memory_if_yield.pto b/test/lit/pto/plan_memory_if_yield.pto index 522464fc8e..d3f1e4a578 100644 --- a/test/lit/pto/plan_memory_if_yield.pto +++ b/test/lit/pto/plan_memory_if_yield.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @if_yield(%arg0: memref<16x16x16xf16, #pto.address_space>, diff --git a/test/lit/pto/plan_memory_loop_in_if.pto b/test/lit/pto/plan_memory_loop_in_if.pto index 558401ec1c..2b4e06be93 100644 --- a/test/lit/pto/plan_memory_loop_in_if.pto +++ b/test/lit/pto/plan_memory_loop_in_if.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @loop_in_if(%arg0: memref<16x16x16xf16, #pto.address_space>, @@ -34,4 +35,3 @@ module { // CHECK: scf.for // CHECK: } else { // CHECK: pto.pointer_cast - diff --git a/test/lit/pto/plan_memory_loop_no_reuse_outer_live.pto b/test/lit/pto/plan_memory_loop_no_reuse_outer_live.pto index 8e122bffeb..82131666c6 100644 --- a/test/lit/pto/plan_memory_loop_no_reuse_outer_live.pto +++ b/test/lit/pto/plan_memory_loop_no_reuse_outer_live.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @loop_outer_live(%arg0: memref<16x16x16xf16, #pto.address_space>, diff --git a/test/lit/pto/plan_memory_nested_loops.pto b/test/lit/pto/plan_memory_nested_loops.pto index 9bfb14ed07..f4a6710512 100644 --- a/test/lit/pto/plan_memory_nested_loops.pto +++ b/test/lit/pto/plan_memory_nested_loops.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @nested_loops(%arg0: memref<16x16x16xf16, #pto.address_space>, diff --git a/test/lit/pto/plan_memory_no_reuse_overlap.pto b/test/lit/pto/plan_memory_no_reuse_overlap.pto index 5a5d50d2ad..8e04a9e601 100644 --- a/test/lit/pto/plan_memory_no_reuse_overlap.pto +++ b/test/lit/pto/plan_memory_no_reuse_overlap.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @no_reuse_overlap(%arg0: memref<16x16x16xf16, #pto.address_space>, diff --git a/test/lit/pto/plan_memory_order_by_size_noreuse.pto b/test/lit/pto/plan_memory_order_by_size_noreuse.pto index 6e9c2d3cfa..fbfdfd6898 100644 --- a/test/lit/pto/plan_memory_order_by_size_noreuse.pto +++ b/test/lit/pto/plan_memory_order_by_size_noreuse.pto @@ -8,7 +8,9 @@ // contract that the option means the same thing whether or not reuse kicks in. // // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT +// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT // RUN: ptoas --pto-arch=a3 --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE +// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE module { func.func @order_by_size_noreuse(%src_ptr: !pto.ptr, %idx_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.kernel} { @@ -65,5 +67,5 @@ module { // On the no-reuse fast path the largest tile (dst, 1x8192xf32) gets offset 0 // ONLY with order-by-size; the default order leaves offset 0 to the smaller // first-generated input tile. -// BYSIZE: pto.pointer_cast(%c0_i64){{.*}} : memref<1x8192xf32 -// DEFAULT-NOT: pto.pointer_cast(%c0_i64){{.*}} : memref<1x8192xf32 +// BYSIZE: pto.alloc_tile addr = %c0_i64 : !pto.tile_buf +// DEFAULT-NOT: pto.alloc_tile addr = %c0_i64 : !pto.tile_buf diff --git a/test/lit/pto/plan_memory_order_by_size_reuse.pto b/test/lit/pto/plan_memory_order_by_size_reuse.pto index 05541cfd07..5bf11d33b5 100644 --- a/test/lit/pto/plan_memory_order_by_size_reuse.pto +++ b/test/lit/pto/plan_memory_order_by_size_reuse.pto @@ -6,7 +6,9 @@ // (first-fit-decreasing) the largest tile is allocated first and gets offset 0. // // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT +// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT // RUN: ptoas --pto-arch=a3 --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE +// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE module { func.func @order_by_size_reuse(%src_ptr: !pto.ptr, %idx_ptr: !pto.ptr, %dst_ptr: !pto.ptr) attributes {pto.kernel} { @@ -62,5 +64,5 @@ module { // The largest tile (dst, 1x32768xf32) is placed at offset 0 ONLY with // order-by-size; the default order leaves offset 0 to a smaller input tile. -// BYSIZE: pto.pointer_cast(%c0_i64){{.*}} : memref<1x32768xf32 -// DEFAULT-NOT: pto.pointer_cast(%c0_i64){{.*}} : memref<1x32768xf32 +// BYSIZE: pto.alloc_tile addr = %c0_i64 : !pto.tile_buf +// DEFAULT-NOT: pto.alloc_tile addr = %c0_i64 : !pto.tile_buf diff --git a/test/lit/pto/plan_memory_peak_8_overlapping.pto b/test/lit/pto/plan_memory_peak_8_overlapping.pto index 48d46f4537..b42f8fd997 100644 --- a/test/lit/pto/plan_memory_peak_8_overlapping.pto +++ b/test/lit/pto/plan_memory_peak_8_overlapping.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @peak_8_overlapping(%arg0: memref<16x16x16xf16, #pto.address_space>, @@ -56,4 +57,3 @@ module { // 8 live buffers implies a max offset of 7*8192 = 57344 bytes. // CHECK: %[[O57344:.*]] = arith.constant 57344 : i64 // CHECK: pto.pointer_cast(%[[O57344]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> - diff --git a/test/lit/pto/plan_memory_peak_exact_capacity.pto b/test/lit/pto/plan_memory_peak_exact_capacity.pto index 4ef8148e1a..8ecb828a99 100644 --- a/test/lit/pto/plan_memory_peak_exact_capacity.pto +++ b/test/lit/pto/plan_memory_peak_exact_capacity.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @peak_exact_capacity(%arg0: memref<16x16x16xf16, #pto.address_space>, @@ -138,4 +139,3 @@ module { // 24 live buffers implies a max offset of 23*8192 = 188416 bytes. // CHECK: %[[O188416:.*]] = arith.constant 188416 : i64 // CHECK: pto.pointer_cast(%[[O188416]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> - diff --git a/test/lit/pto/plan_memory_reserve_buffer_manual_reject_nested_alloc.pto b/test/lit/pto/plan_memory_reserve_buffer_manual_reject_nested_alloc.pto index 06ae763576..b89bb80022 100644 --- a/test/lit/pto/plan_memory_reserve_buffer_manual_reject_nested_alloc.pto +++ b/test/lit/pto/plan_memory_reserve_buffer_manual_reject_nested_alloc.pto @@ -1,4 +1,5 @@ // RUN: not ptoas %s 2>&1 1>/dev/null | FileCheck %s +// RUN: not ptoas --plan-memory-impl=modern %s 2>&1 1>/dev/null | FileCheck %s module { func.func @manual_nested_alloc(%cond: i1) { @@ -22,4 +23,4 @@ module { } } -// CHECK: error: 'pto.reserve_buffer' op pto.reserve_buffer with explicit 'base' (auto = false) is not supported in PlanMemory; use --pto-level=level3 or set auto = true +// CHECK: error: pto.reserve_buffer with explicit 'base' (auto = false) is not supported when --pto-level=level1 or level2; use --pto-level=level3 or set auto = true diff --git a/test/lit/pto/plan_memory_reserve_buffer_prefix.pto b/test/lit/pto/plan_memory_reserve_buffer_prefix.pto index fd4a577eb4..a5bcc2d249 100644 --- a/test/lit/pto/plan_memory_reserve_buffer_prefix.pto +++ b/test/lit/pto/plan_memory_reserve_buffer_prefix.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @reserve_prefix(%arg0: memref<16x16x16xf16, #pto.address_space>, diff --git a/test/lit/pto/plan_memory_reuse_sequential.pto b/test/lit/pto/plan_memory_reuse_sequential.pto index 6ed1a49599..425a3b399f 100644 --- a/test/lit/pto/plan_memory_reuse_sequential.pto +++ b/test/lit/pto/plan_memory_reuse_sequential.pto @@ -1,4 +1,5 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s module { func.func @reuse_sequential(%arg0: memref<16x16x16xf16, #pto.address_space>, diff --git a/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto b/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto new file mode 100644 index 0000000000..fe520f08bc --- /dev/null +++ b/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.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-level=level2 --pto-arch=a3 --plan-memory-impl=modern --emit-pto-ir \ +// RUN: --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s + +module attributes {"pto.target_arch" = "a3"} { + func.func @spec_level0_no_reuse_overlap( + %gm : memref<32x32xf16, #pto.address_space>) { + %src = memref.alloc() : memref<32x32xf16, #pto.address_space> + %dst = memref.alloc() : memref<32x32xf16, #pto.address_space> + + pto.tload ins(%gm : memref<32x32xf16, #pto.address_space>) + outs(%src : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} + pto.tmov ins(%src : memref<32x32xf16, #pto.address_space>) + outs(%dst : memref<32x32xf16, #pto.address_space>) + return + } +} + +// CHECK-LABEL: func.func @spec_level0_no_reuse_overlap +// CHECK-DAG: %[[ADDR0:.*]] = arith.constant 0 : i64 +// CHECK-DAG: %[[ADDR1:.*]] = arith.constant 2048 : i64 +// CHECK: pto.pointer_cast(%[[ADDR0]]) +// CHECK: pto.pointer_cast(%[[ADDR1]]) +// CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_spec_level0_reuse.pto b/test/lit/pto/plan_memory_spec_level0_reuse.pto new file mode 100644 index 0000000000..b55d457086 --- /dev/null +++ b/test/lit/pto/plan_memory_spec_level0_reuse.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-level=level2 --pto-arch=a3 --plan-memory-impl=modern --emit-pto-ir \ +// RUN: --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s + +module attributes {"pto.target_arch" = "a3"} { + func.func @spec_level0_reuse( + %gm0 : memref<32x32xf16, #pto.address_space>, + %gm1 : memref<32x32xf16, #pto.address_space>) { + %buf0 = memref.alloc() : memref<32x32xf16, #pto.address_space> + pto.tload ins(%gm0 : memref<32x32xf16, #pto.address_space>) + outs(%buf0 : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} + + %buf1 = memref.alloc() : memref<32x32xf16, #pto.address_space> + pto.tload ins(%gm1 : memref<32x32xf16, #pto.address_space>) + outs(%buf1 : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} + return + } +} + +// CHECK-LABEL: func.func @spec_level0_reuse +// CHECK: %[[ADDR0:.*]] = arith.constant 0 : i64 +// CHECK: %[[BUF0:.*]] = pto.pointer_cast(%[[ADDR0]]) +// CHECK: %[[BUF1:.*]] = pto.pointer_cast(%[[ADDR0]]) +// CHECK-NOT: memref.alloc diff --git a/test/lit/pto/reserve_buffer_level12_auto_base_reject.pto b/test/lit/pto/reserve_buffer_level12_auto_base_reject.pto new file mode 100644 index 0000000000..cc0c462343 --- /dev/null +++ b/test/lit/pto/reserve_buffer_level12_auto_base_reject.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: not ptoas --pto-level=level1 --pto-arch=a3 %s 2>&1 | FileCheck %s +// RUN: not ptoas --pto-level=level2 --pto-arch=a3 %s 2>&1 | FileCheck %s + +module { + func.func @auto_base_reject() { + %fifo = pto.reserve_buffer { + name = "fifo", + size = 8192, + location = #pto.address_space, + auto = true, + base = 0 + } -> i32 + return + } +} + +// CHECK: error: unexpected 'base' on auto reserve_buffer diff --git a/test/lit/pto/reserve_buffer_level12_manual_base_reject.pto b/test/lit/pto/reserve_buffer_level12_manual_base_reject.pto new file mode 100644 index 0000000000..3f1069f574 --- /dev/null +++ b/test/lit/pto/reserve_buffer_level12_manual_base_reject.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: not ptoas --pto-level=level1 --pto-arch=a3 %s 2>&1 | FileCheck %s +// RUN: not ptoas --pto-level=level2 --pto-arch=a3 %s 2>&1 | FileCheck %s + +module { + func.func @manual_base_reject() { + %fifo = pto.reserve_buffer { + name = "fifo", + size = 8192, + location = #pto.address_space, + auto = false, + base = 0 + } -> i32 + return + } +} + +// CHECK: error: pto.reserve_buffer with explicit 'base' (auto = false) is not supported when --pto-level=level1 or level2 diff --git a/test/lit/pto/reserve_buffer_level3_auto_reject.pto b/test/lit/pto/reserve_buffer_level3_auto_reject.pto new file mode 100644 index 0000000000..cb258ef978 --- /dev/null +++ b/test/lit/pto/reserve_buffer_level3_auto_reject.pto @@ -0,0 +1,23 @@ +// 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-level=level3 --pto-arch=a3 %s 2>&1 | FileCheck %s + +module { + func.func @level3_auto_reject() { + %fifo = pto.reserve_buffer { + name = "fifo", + size = 8192, + location = #pto.address_space, + auto = true + } -> i32 + return + } +} + +// CHECK: error: pto.reserve_buffer requires 'auto = false' and explicit 'base' when --pto-level=level3 diff --git a/test/lit/pto/reserve_buffer_level3_manual_base.pto b/test/lit/pto/reserve_buffer_level3_manual_base.pto new file mode 100644 index 0000000000..0164b8474a --- /dev/null +++ b/test/lit/pto/reserve_buffer_level3_manual_base.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-level=level3 --pto-arch=a3 --emit-pto-ir \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s + +module { + func.func @manual_base_level3() { + %fifo = pto.reserve_buffer { + name = "fifo", + size = 8192, + location = #pto.address_space, + auto = false, + base = 0 + } -> i32 + return + } +} + +// CHECK-LABEL: func.func @manual_base_level3 +// CHECK: arith.constant 0 : i32 +// CHECK-NOT: pto.reserve_buffer diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 2c62a9005c..4b77f7faa7 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -324,13 +324,18 @@ static llvm::cl::opt enableInsertSync("enable-insert-sync", llvm::cl::desc("Enable automatic synchronization insertion pass"), llvm::cl::init(false)); -static llvm::cl::opt planMemoryOrderBySize( +[[maybe_unused]] static llvm::cl::opt planMemoryOrderBySize( "plan-memory-order-by-size", - llvm::cl::desc("PlanMemory: allocate buffers largest-first " - "(first-fit-decreasing) instead of the default DMA-first " - "order"), + llvm::cl::desc("Plan larger local buffers first inside one AddressSpace " + "before applying the basic SPEC_LEVEL_0 reuse strategy"), llvm::cl::init(false)); +static llvm::cl::opt planMemoryImpl( + "plan-memory-impl", + llvm::cl::desc("Select local memory planner implementation: legacy or " + "modern"), + llvm::cl::init("legacy")); + static llvm::cl::opt enableBufidSync( "enable-bufid_sync", llvm::cl::desc("Enable A5 buffer-id synchronization insertion pass"), @@ -628,6 +633,97 @@ static bool parseBuildLevel(llvm::StringRef levelStr, PTOBuildLevel &out) { return false; } +struct ReserveBufferMemSpec { + uint64_t capacityBytes = 0; + uint64_t alignmentBytes = 1; +}; + +static ReserveBufferMemSpec getReserveBufferMemSpec(PTOArch arch, + AddressSpace space) { + switch (space) { + case AddressSpace::VEC: + return {arch == PTOArch::A5 ? 253952ull : 196608ull, 256}; + case AddressSpace::MAT: + return {524288ull, 256}; + case AddressSpace::LEFT: + case AddressSpace::RIGHT: + case AddressSpace::ACC: + case AddressSpace::BIAS: + case AddressSpace::SCALING: + case AddressSpace::GM: + case AddressSpace::Zero: + break; + } + return {}; +} + +static LogicalResult validateReserveBufferBase(pto::ReserveBufferOp op, + PTOArch arch) { + auto baseAttr = op.getBaseAttr(); + if (!baseAttr) + return op.emitError("expects explicit 'base'"); + + int64_t signedBase = baseAttr.getInt(); + if (signedBase < 0) + return op.emitError("expects 'base' to be non-negative when present"); + + ReserveBufferMemSpec spec = + getReserveBufferMemSpec(arch, op.getLocation().getAddressSpace()); + uint64_t base = static_cast(signedBase); + if (base % spec.alignmentBytes != 0) { + return op.emitError("expects 'base' to be aligned to ") + << spec.alignmentBytes << " bytes for " + << stringifyEnum(op.getLocation().getAddressSpace()); + } + + uint64_t size = static_cast(op.getSize()); + if (base > spec.capacityBytes || size > spec.capacityBytes - base) { + return op.emitError("reserved range exceeds ") + << stringifyEnum(op.getLocation().getAddressSpace()) + << " capacity: base " << base << " + size " << size + << " > " << spec.capacityBytes << " bytes"; + } + + return success(); +} + +static bool validateReserveBufferLevelRules(ModuleOp module, + PTOBuildLevel level) { + bool failed = false; + PTOArch arch = getTargetArch(module); + module.walk([&](pto::ReserveBufferOp op) { + if (level != PTOBuildLevel::Level3) { + if (op.getAutoAlloc()) { + if (op.getBaseAttr()) { + op.emitError("unexpected 'base' on auto reserve_buffer: " + "level1/level2 assign it in pto-plan-memory"); + failed = true; + } + return; + } + + if (op.getBaseAttr()) + (void)validateReserveBufferBase(op, arch); + op.emitError("pto.reserve_buffer with explicit 'base' (auto = false) is " + "not supported when --pto-level=level1 or level2; use " + "--pto-level=level3 or set auto = true"); + failed = true; + return; + } + + if (op.getAutoAlloc() || !op.getBaseAttr()) { + op.emitError("pto.reserve_buffer requires 'auto = false' and explicit " + "'base' when --pto-level=level3"); + failed = true; + return; + } + + if (mlir::failed(validateReserveBufferBase(op, arch))) + failed = true; + }); + return !failed; +} + static constexpr llvm::StringLiteral kAutoSyncTailPolicyBarrierAll = "barrier_all"; static constexpr llvm::StringLiteral kAutoSyncTailPolicyMte3ToSEvent0 = @@ -2909,6 +3005,9 @@ int mlir::pto::compilePTOASModule( return 1; } + if (!validateReserveBufferLevelRules(*module, effectiveLevel)) + return 1; + { PassManager preBackendPM(module->getContext()); preBackendPM.enableVerifier(); @@ -3002,12 +3101,18 @@ int mlir::pto::compilePTOASModule( pto::createPTORematerializeFixpipeVectorQuantPass()); if (effectiveLevel != PTOBuildLevel::Level3) { - PlanMemoryOptions planMemoryOption; - planMemoryOption.memMode = MemPlanMode::LOCAL_MEM_PLAN; - planMemoryOption.enableGlobalReuse = false; - planMemoryOption.enablePrintMemoryAllocatedSize = false; - planMemoryOption.orderBySize = planMemoryOrderBySize; - pm.addPass(pto::createPlanMemoryPass(planMemoryOption)); + pto::PlanMemoryOptions planMemoryOptions; + planMemoryOptions.memMode = "local"; + planMemoryOptions.orderBySize = planMemoryOrderBySize; + if (planMemoryImpl == "legacy") { + pm.addPass(pto::createPlanMemoryPass(planMemoryOptions)); + } else if (planMemoryImpl == "modern") { + pm.addPass(pto::createPlanMemoryModernPass(planMemoryOptions)); + } else { + llvm::errs() << "Error: invalid --plan-memory-impl='" << planMemoryImpl + << "', expected 'legacy' or 'modern'.\n"; + return 1; + } } pm.addPass(pto::createPTOResolveReservedBuffersPass()); pm.addNestedPass(pto::createPTORemoveIdentityTMovPass()); @@ -3039,7 +3144,8 @@ int mlir::pto::compilePTOASModule( // Materialize per-slot single-address `pto.pointer_cast` (constant slot) // or an `arith.select` chain (dynamic slot). The multi-address cast - // produced by PlanMemory survives as the alloc anchor. + // produced by explicit-address lowering or PTOPlanMemory survives as + // the alloc anchor. pm.addPass(pto::createPTOResolveBufferSelectPass()); if (effectiveBackend == PTOBackend::EmitC) pm.addPass(createNarrowUnusedMultiResultProvenancePass()); From 6f00019f33e5ee7b09a53a3b1590af098a22862c Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 13 Jul 2026 17:22:41 +0800 Subject: [PATCH 02/91] Implement modern memplan five-gate reuse checks --- ...est-first-fit-five-gates-memplan-design.md | 655 ++++++++++++++++ .../InsertSync/MemoryDependentAnalyzer.cpp | 60 +- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 708 ++++++++++++++---- ...lan_memory_five_gates_lifetime_overlap.pto | 40 + .../pto/plan_memory_five_gates_phi_family.pto | 43 ++ .../pto/plan_memory_inplace_forbid_alias.pto | 60 ++ .../pto/plan_memory_order_by_size_noreuse.pto | 2 +- .../pto/plan_memory_order_by_size_reuse.pto | 2 +- tools/ptoas/ptoas.cpp | 13 +- 9 files changed, 1421 insertions(+), 162 deletions(-) create mode 100644 docs/designs/ptoas-largest-first-fit-five-gates-memplan-design.md create mode 100644 test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto create mode 100644 test/lit/pto/plan_memory_five_gates_phi_family.pto create mode 100644 test/lit/pto/plan_memory_inplace_forbid_alias.pto diff --git a/docs/designs/ptoas-largest-first-fit-five-gates-memplan-design.md b/docs/designs/ptoas-largest-first-fit-five-gates-memplan-design.md new file mode 100644 index 0000000000..39e865910d --- /dev/null +++ b/docs/designs/ptoas-largest-first-fit-five-gates-memplan-design.md @@ -0,0 +1,655 @@ +# PTOAS Largest-First-Fit 与五道冲突闸门内存规划设计 + +## 总体方案 + +PTOAS 当前保留两套 local memory planner： + +- legacy memplan：默认启用，保留旧版 SPEC_LEVEL_0/1/2 三级投机复用，内存不够时回滚的策略。 +- modern memplan：通过 `--plan-memory-impl=modern` 显式启用，当前实现Largest-First-Fit 与四道冲突闸门内存规划设计 + +PTOAS 当前已经具备 largest-first 风格的规划开关：`--plan-memory-order-by-size`。该选项会让 planner 在同一 AddressSpace 内优先处理更大的 buffer。当用户显式选择 `--plan-memory-impl=modern` 且未显式指定 `--plan-memory-order-by-size` 时，modern memplan 默认开启该排序；legacy memplan 仍保持默认关闭。 + +当前 `pto.alloc_tile` 的 lowering 路径保持 tile-native： + +```text +pto.alloc_tile(no addr) + -> PTOViewToMemref 透传，不转换成 memref.alloc + -> pto-plan-memory 收集为 local allocation root + -> modern/legacy memplan 按 level 校验并规划 local addr + -> 直接给 pto.alloc_tile 补常量 addr + -> 后续 pto.t* tile op 继续使用 !pto.tile_buf 形态 +``` + +用户显式写 `pto.alloc_tile addr` 时，level 语义仍由 memplan 校验：level1/level2 禁止显式 addr，level3 要求显式 local addr。也就是说，`pto.alloc_tile` 不再通过 `memref.alloc -> pto.pointer_cast -> pto.bind_tile` 这条中间链路表达地址；只有非 `alloc_tile` 的 memref/address root 仍可在 materialize 阶段使用 `pto.pointer_cast + pto.bind_tile` 表达规划后的地址与 tile metadata。 + +## 目标 + +本设计目标如下： + +- 复用 PTOAS 已有 `--plan-memory-order-by-size` largest-first 能力。 +- 将复用判定收敛为统一的 `canShare(a, b)` 谓词。 +- 引入统一的扁平冲突闸门模型，避免不安全复用。 +- 当前实现范围不包含 PyPTO 闸门 4；PTOAS 显式 multibuffer 已覆盖 ping-pong slot 分离语义。 +- 保持按 AddressSpace 独立规划。 +- 保持不回滚、不降级的 deterministic 策略。 +- 保持 `pto.alloc_tile(no addr)` tile-native 路径：memplan 直接补 `addr`。 +- 保持 legacy memplan 默认行为不变。 + +## 非目标 + +本设计不包含以下内容： + +- 不修改 legacy memplan 的 StorageEntry / SPEC_LEVEL_1 / SPEC_LEVEL_2 逻辑。 +- 不在 legacy memplan 中实现五道闸门。 +- 不恢复旧版回滚式投机规划。 +- 不把 modern memplan 作为默认实现。 +- 不在本阶段实现跨函数、跨 module 的全局内存规划。 + +## 核心思想 + +将每个 plannable local allocation root 抽象成一个待装箱 item。这里的 item 对应当前 modern memplan 中的 `RootInfo`，不是所有 local SSA value： + +`pto.bind_tile`、`pto.slot_marker`、`memref.subview`、`memref.cast` 等 alias/view value 不会成为新的 item，而是通过 `valueToRoots` 归属到已有 root。四道闸门需要的额外约束数据不放进 `RootInfo`，而是由 `ConflictFacts` 这类 side table 提供。 + + +## 算法流程 + +### 1. Root 收集 + +modern memplan 继续收集以下 local root： + +- `pto.alloc_tile(no addr)`。 +- plannable local `memref.alloc`。 +- 后续如有其它明确的 local address root，可扩展到同一 root 表。 + +不参与 local memplan 的 root： + +- GM space。 +- `AddressSpace::Zero`。 +- 已经由 level3 显式指定地址的 local `pto.alloc_tile addr`。 +- 非静态 shape 或无法计算 element byte size 的 root。 + +### 2. LifetimeInterval + +每个 root 生成一个 `LifetimeInterval`： + +```text + Value root; + Operation *defOp; + AddressSpace space; + uint64_t slotBytes; + uint64_t totalBytes; + uint64_t alignmentBytes; + uint64_t slotCount; + unsigned allocIndex; + unsigned freeIndex; + unsigned stableOrder; + SmallVector offsets; +``` + +现代 planner 中已有线性化 walk 和 root alias 传播，本设计在此基础上补齐： + +- loop-aware lifetime extension。 +- branch / yield / iter_arg alias family 信息。 +- per-root semantic metadata。 + +### 3. 按 AddressSpace 分桶 + +复用只允许发生在同一 AddressSpace 内： + +```text +Vec 只和 Vec 复用 +Mat 只和 Mat 复用 +Left/Right/Acc/Bias/Scaling 各自独立 +GM 不参与 local memplan +``` + +分桶是硬约束，在进入装箱前完成。 + +### 4. Largest-First-Fit 排序与装箱 + +PTOAS 已有 `--plan-memory-order-by-size` 对应的 largest-first 排序能力。打开该选项后，每个 AddressSpace 桶内排序： + +```text +sizeBytes 降序 +sizeBytes 相同则按 defIndex 升序 +defIndex 相同则按 stableOrder 升序 +``` + +这样先让大 buffer 更早参与规划，再把生命周期不冲突的小 buffer 放入可复用位置，避免 definition-order greedy 只能“小复用先出现的大”的单向问题。 + +将一个物理 local buffer 抽象成一个 bin： + +```text +ReuseGroup + representative root + members + slot size bytes + address space + chosen offset +``` + +因此，本设计不是重新新增一个独立的 largest-first 开关，而是要求四道闸门的 `canShare` 判定接入现有 `--plan-memory-order-by-size` 路径。 + +特点： + +- 不回滚。 +- 不降级。 +- 装进第一个可容纳 group 后立即停止。 +- 后续容量超限直接报错，不为了适配容量放松闸门。 + +## 四道禁止内存复用的闸门 + +四道闸门是扁平 AND 关系，任意一道闸门失败，两个 root 就不能进入同一个 `ReuseGroup`，也就不能复用同一个 base offset。 + +### 闸门 1：生命周期 + +**目的：** 禁止两个运行时可能同时存活的 local root 复用同一块物理地址。 + +基础规则： + +```text +如果 a 和 b 生命周期重叠，则不能复用。 +如果 a.lastUse == b.def 或 b.lastUse == a.def，则认为 touching，可继续检查其它闸门。 +``` + +PTOAS 现有 `lifetimesOverlap` 可扩展为： + +```text +overlap = !(a.lastUse <= b.def || b.lastUse <= a.def) +``` + +`<=` 表示 touching 不算生命周期冲突。 + +**适用场景 sample：普通直线代码中的临时 buffer 复用。** + +```text +%a = alloc vec[1024] +use(%a) // %a.lastUse + +%b = alloc vec[1024] // %b.def +use(%b) +``` + +如果 `%a.lastUse <= %b.def`，闸门 1 允许 `%a` 和 `%b` 进入后续闸门。只要其它闸门也通过，二者可以复用同一个 offset： + +```text +%a.offset = 0 +%b.offset = 0 +``` + +反例： + +```text +%a = alloc vec[1024] +%b = alloc vec[1024] +use(%a) +use(%b) +``` + +此时 `%a` 和 `%b` 的生命周期重叠，闸门 1 直接失败，必须分配不同 offset。 + +### 闸门 2：phi family + +**目的：** 处理互斥控制流分支，避免过度保守的 liveness extension 阻止本来安全的复用。 + +场景： + +```text +%r = scf.if %cond -> tile { + scf.yield %a +} else { + scf.yield %b +} +``` + +`%a` 和 `%b` 在程序序生命周期上可能被扩展到 `%r` 的最后使用点，看起来重叠；但运行时两个分支互斥，不会同时存在。因此同一 phi family 的 yield source 可以豁免闸门 1 的生命周期冲突。 + +需要收集： + +```text +phiFamilyIds[root] = set +``` + +规则： + +- 同 family 且无其它非 family 真重叠成员：允许共享。 +- 如果混入外部 live root 或非 family 真重叠 root：不豁免。 + +实现建议： + +```text +Gate2_LifetimeAndPhi(a, b): + if !lifetimeOverlap(a, b): + return true + + if samePhiFamily(a, b): + return true + + return false +``` + +因此闸门 2 在实现上可以嵌入闸门 1，也可以保留独立函数名但由闸门 1 调用。 + +**适用场景 sample：if/else 分支内 local root 互斥复用。** + +```text +%r = scf.if %cond -> tile { + %then_buf = alloc vec[1024] + produce(%then_buf) + scf.yield %then_buf +} else { + %else_buf = alloc vec[1024] + produce(%else_buf) + scf.yield %else_buf +} + +consume(%r) +``` + +如果简单地把 `%then_buf` 和 `%else_buf` 的生命周期都延伸到 `consume(%r)`，二者看起来重叠，闸门 1 会失败。但运行时 then/else 互斥，不会同时分配这两个 branch-local root。因此它们属于同一个 phi family 时，可以共享同一 offset： + +```text +%then_buf.offset = 0 +%else_buf.offset = 0 +``` + +反例： + +```text +%outer = alloc vec[1024] + +%r = scf.if %cond -> tile { + %then_buf = alloc vec[1024] + use(%outer) + scf.yield %then_buf +} else { + %else_buf = alloc vec[1024] + use(%outer) + scf.yield %else_buf +} +``` + +`%then_buf` 和 `%else_buf` 之间可以因为 phi family 互斥而复用；但它们不能借这个豁免和 `%outer` 这种外部 live root 错误复用。 + + +### 闸门 3：target-specific load/tpop hazard + +**目的：** 表达特定 target 上的局部硬件/后端 hazard。它不是通用生命周期问题，而是某些 load-derived buffer 与 consumes-tpop writer 在 touching 点复用会触发目标相关错误。 + +参考 PyPTO 的 Ascend910B split-AIV hazard。PTOAS 中设计为 target-gated 闸门： + +```text +如果目标架构/后端不存在该 hazard，则闸门恒通过。 +如果存在该 hazard，则禁止 load-derived buffer 与 consumes-split-tpop 的 writer 在 touching 点复用。 +``` + +当前 PTOAS modern memplan 的启用条件： +A5 上该闸门恒通过。A3 上也只有识别到 split tpop 派生值时才可能触发。 + +#### 设计原理 + +这道闸门保护的是一种普通静态生命周期分析看不到的 target-specific touching 复用风险。普通 memplan 只看 allocation root 的 def/use 区间：如果一个 load-derived buffer 的最后一次使用正好等于另一个 writer output 的写入点，生命周期上属于 touching，通常可以复用同一个 offset。但在特定 target 上，`tload/tprefetch` 产生的 load-derived buffer 与 consuming split-tpop 的 writer 可能在后端流水或指令调度中存在更细粒度的运行时 overlap；如果二者复用同一地址，writer 的写入可能破坏 load-derived input 在该 target 上仍然需要保持稳定的内容。 + +因此该闸门不把问题抽象成通用“生命周期重叠”，而是显式识别一类 target hazard： + +```text +load-derived input 的 last use +touching +consumes split-tpop 的 writer output 写入 +``` + +只有当目标架构确认存在该 hazard 时才禁止复用。当前 PTOAS 中 A3 开启该闸门，A5 恒通过，避免把 A3 的局部硬件/后端约束错误扩散成所有 target 的通用内存规划规则。 + +实现上使用 writer op index，而不是 allocation root 的 def index，原因是 hazard 发生在“某个 writer op 同时读取 load-derived input 和 split-tpop-derived value，并写出 DPS output”的这个操作点。allocation root 的 def 只表示 buffer 被声明出来，不一定等于真正发生写入和触发 target hazard 的位置；用 writer op index 可以把 touching 条件精确绑定到产生风险的 op 上。 + +相关事实拆开收集： + +```text +loadDerivedRoots: + 标记来自 tload/tprefetch 的 DPS dst root + +split-tpop-derived values: + 标记 split tpop 结果、split tpop tile operand 以及它们经过 alias/view op 派生出的 value + +tpopConsumerRoots: + 标记同时读取 load-derived root 和 split-tpop-derived value 的 DPS writer output root + +tpopConsumerWriteIndices: + 记录上述 writer op 的 op index，用于判断是否正好在 touching 点复用 +``` + +这种拆分可以避免过度保守：不是所有 load buffer 都禁止复用，也不是所有 tpop 相关 buffer 都禁止复用，只有“load-derived input 的最后使用点”碰到“consumes split-tpop writer 的写入点”时才失败。 + +当前 PTOAS modern memplan 对 DPS output root 使用 writer-def liveness：如果某个 DPS output 是 pure overwrite，即该 output 没有被 memory effects 标记为 Read，则它的 `allocIndex` 会从 `pto.alloc_tile` / `memref.alloc` 收缩到第一次有效 writer op。这样 load-derived input 的 `freeIndex` 与 writer output 的 `allocIndex` 可以形成 touching，是否允许复用继续由 target-specific load/tpop hazard 和 op semantic no-alias 闸门判断。若 output 需要旧值，例如 read-modify-write / accumulate 语义，则保持 allocation-start 的保守生命周期。 + +需要收集： + +```text +loadDerivedRoots +tpopConsumerRoots +tpopConsumerWriteIndices +targetHazardEnabled +``` + +判定： + +```text +input.root in loadDerivedRoots +&& writer.root in tpopConsumerRoots +&& input.lastUseIndex in tpopConsumerWriteIndices[writer.root] +``` + +双向检查： + +```text +hazard(a, b) || hazard(b, a) +``` + +**适用场景 sample：target-gated touching 复用禁止。** + +```text +%load_buf = alloc vec[1024] +load_or_tpop_producer(%load_buf) +last_use(%load_buf) + +%writer_dst = alloc vec[1024] +writer_consumes_tpop(...) outs(%writer_dst) +``` + +从普通生命周期看，`%load_buf.lastUse == %writer_dst.def`，属于 touching，可以复用。但如果 target 标记了： + +```text +%load_buf in loadDerivedRoots +%writer_dst in tpopConsumerRoots +tpopConsumerWriteIndices[%writer_dst] contains writer op index +targetHazardEnabled = true +``` + +则闸门 3 失败，二者不能复用。若目标没有该 hazard，闸门 3 恒通过，不影响复用。 + +当前实现的 fact 来源： + +```text +loadDerivedRoots: + pto.tload / pto.tprefetch 的 DPS dst root + +split tpop derived value: + A3 上 split != 0 的 pto.tpop_from_aic result + A3 上 split != 0 的 pto.tpop tile operand + 以及从这些 value 经过 bind_tile / slot_marker / cast / subview / select 等 alias/view op 派生出的 value + +tpopConsumerRoots: + 某个 DPS writer 同时读取 split tpop derived value 和 load-derived root 时， + 记录该 writer 的 DPS output root +``` + + +### 闸门 4：op semantic no-alias + +**目的：** 表达“从生命周期看可以复用，但从 op 语义看不能 alias”的约束。闸门 4 对应 PyPTO 的 inplace 机制：`not_inplace_safe()` 与 `forbid_output_alias(i)`，并覆盖 PTOAS legacy memplan 中已有的 scratch-output conflict。 + +在 modern memplan 中，闸门 4 不应再只叫笼统的 `semantic conflict`，而应建模为一张明确的 forbid-alias side table： + +```text +forbidAlias[root] = forbiddenRootSet +``` + +这里的 `root` 是 plannable local allocation root，而不是任意 SSA value。收集时需要先通过 `valueToRoots` 把 operand、view、alias value 归约到 root，再记录 root 与 root 之间的禁止复用关系。 + +#### 4.1 scratch-output conflict + +PTOAS legacy memplan 的 semantic conflict 主要就是 scratch-output conflict。modern memplan 应继续覆盖这类场景： + +```text +op implements PTO_DpsInitOpInterface +dpsInits = op outs(...) + +effects = MemoryEffectOpInterface::getEffects(op) +scratchOperands = + Write operand + && operand in op operands + && operand not in dpsInits + +for scratch in scratchOperands: + for dst in dpsInits: + forbidAlias[root(scratch)].insert(root(dst)) + forbidAlias[root(dst)].insert(root(scratch)) +``` + +**适用场景 sample：`tmp` scratch 不能和 output 复用。** + +```text +%src = alloc vec[1024] +%tmp = alloc vec[1024] +%dst = alloc vec[1024] + +pto.ttrans ins(%src, %tmp) outs(%dst) +``` + +`%tmp` 是 op 执行过程中的 scratch workspace，`%dst` 是最终 output。即使未来 liveness 把 `%dst` 看作在 op 上定义、从而让 `%tmp.lastUse == %dst.def` 成为 touching，二者也不能复用，否则 scratch 写入可能覆盖 output。闸门 3 记录： + +```text +forbidAlias[%tmp].insert(%dst) +forbidAlias[%dst].insert(%tmp) +``` + +#### 4.2 not_inplace_safe + +PyPTO 中 `not_inplace_safe()` 表示该 op 不能做 `src == dst` 的 inplace 执行。映射到 PTOAS 时，规则是： + +```text +if opPolicy.notInplaceSafe: + for operand in op operands excluding dpsInits: + for dst in dpsInits: + forbidAlias[root(operand)].insert(root(dst)) + forbidAlias[root(dst)].insert(root(operand)) +``` + +典型 op 包括： + +```text +pto.ttrans +pto.tgather +pto.tands / pto.tors / pto.txors +pto.tfillpad_expand +pto.tfmod / pto.tfmods +pto.trecip / pto.trsqrt +pto.trowmax / pto.trowmin / pto.trowsum / pto.trowprod +pto.trowargmax / pto.trowargmin +pto.tcolargmax / pto.tcolargmin +pto.tsort32 / pto.tmrgsort +``` + +其中 `pto.tands` / `pto.tors` / `pto.txors` 和 `pto.tfillpad_expand` 是 PTOAS 侧额外保守标记的 non-inplace-safe op。它们虽然不是 scratch-output conflict，但后端/ISA 语义没有明确承诺 input/output alias 安全，memplan 不应通过地址复用隐式把它们变成 inplace 执行。 + +**适用场景 sample：算法本身不支持 input/output alias。** + +```text +%x = alloc vec[1024] +%y = alloc vec[1024] + +pto.tfmod ins(%x, %rhs) outs(%y) +``` + +`tfmod` 的实现可能会在计算中间覆盖某个源值，但后续仍需要原始源值。因此 `%y` 不能复用 `%x` 的物理地址。即使生命周期或 touching 规则允许，也必须由闸门 3 禁止： + +```text +forbidAlias[%x].insert(%y) +``` + +#### 4.3 forbid_output_alias(i) + +PyPTO 中 `forbid_output_alias(i)` 表示 op 整体可以对某些 value operand 做 inplace，但 output 不能 alias 第 `i` 个特定 operand。PTOAS 中需要按 PTO IR 的 DPS operand 布局映射到具体 operand。 + +典型场景： + +```text +pto.tsel: + forbid mask + forbid tmp + +pto.trowexpand / pto.tcolexpand: + forbid broadcast source + +pto.trowexpand* / pto.tcolexpand*: + forbid row/column vector operand +``` + +**适用场景 sample：broadcast vector 不能被 output 覆盖。** + +```text +%row = alloc vec[16x1] +%dst = alloc vec[16x64] + +pto.trowexpand ins(%row) outs(%dst) +``` + +`%row` 会被重复读取并广播到 `%dst` 的多个位置。如果 `%dst` 复用 `%row` 的地址，写 output 的过程中可能覆盖后续仍要读取的 broadcast source。因此闸门 3 记录： + +```text +forbidAlias[%row].insert(%dst) +``` + +**适用场景 sample：select 的 mask/tmp 不能和 output alias。** + +```text +%mask = alloc vec[1024] +%tmp = alloc vec[1024] +%dst = alloc vec[1024] + +pto.tsel ins(%mask, %lhs, %rhs, %tmp) outs(%dst) +``` + +`%lhs` / `%rhs` 是否允许和 `%dst` inplace 取决于 op 语义；但 `%mask` 和 `%tmp` 是被 op 读取或作为 scratch 使用的特殊 operand，不能被 `%dst` 覆盖。闸门 3 记录： + +```text +forbidAlias[%mask].insert(%dst) +forbidAlias[%tmp].insert(%dst) +``` + +#### 4.4 闸门 4 判定 + +`canShare(a, b)` 中的闸门 3 是双向检查： + +```text +Gate4_OpSemanticNoAlias(a, b): + return !forbidAlias[a].contains(b) + && !forbidAlias[b].contains(a) +``` + +由于 PTOAS 中 `pto.bind_tile`、`pto.slot_marker`、`memref.subview` 等 view/alias value 不一定是 root，收集 forbid-alias 时必须先做 root 归约： + +```text +for value in op operands / dpsInits: + roots = valueToRoots[value] +``` + +如果一个 value 通过 alias closure 对应多个 root，则需要记录所有 root 组合。这样后续 `ReuseGroup` 只需要比较 root 与 root，不需要在装箱阶段重新理解每个 op 的 operand 语义。 + +## PTOAS 数据结构建议 + +### RootInfo + +```cpp +struct RootInfo { + Value root; + Operation *defOp = nullptr; + AddressSpace space = AddressSpace::Zero; + uint64_t slotBytes = 0; + uint64_t totalBytes = 0; + uint64_t alignmentBytes = 1; + uint64_t slotCount = 1; + unsigned allocIndex = 0; + unsigned freeIndex = 0; + unsigned stableOrder = 0; + SmallVector offsets; +}; +``` + +本设计保持当前 modern memplan 的 `RootInfo` 字段不变。`RootInfo` 只表达 local allocation root 的基础事实： + +```text +root identity +definition op +address space +slot size / total size / slot count +lifetime interval +stable order +planned offsets +``` + +四道闸门所需的额外事实不直接塞进 `RootInfo`，而是放在 `ConflictFacts` 这类 side table 中。这样可以避免 root 结构随着每个闸门膨胀，也便于后续按阶段启用或删除某个闸门。 + +### ReuseGroup + +```cpp +struct ReuseGroup { + Value representative; + SmallVector memberIndices; + AddressSpace space; + uint64_t slotSizeBytes; + uint64_t offsetBytes; +}; +``` + +### ConflictFacts + +```cpp +struct ConflictFacts { + DenseMap> forbidAlias; + DenseSet loadDerivedRoots; + DenseSet tpopConsumerRoots; + DenseMap> phiFamilyIds; + + // Reserved for future implicit pipeline lowering support. Not used by the + // current PTOAS design because explicit multibuffer owns ping-pong slot + // separation. + // DenseMap> pipelineMembership; + // DenseSet pipelineLoadRoots; +}; +``` + +### 说明：pipeline stage load conflict（预留，不纳入当前实现） + +- 当前不实现 PyPTO 闸门 4。 +- 不新增 pipeline stage load 复用负例。 +- 保留设计占位：若未来 PTOAS 引入隐式 pipeline lowering，并能稳定提供 `pipelineMembership[root] = (group, stage)`，再接入该闸门。 +- 显式 `pto.alloc_multi_tile count=N` 继续由 multibuffer slot 分配保证 ping-pong 正确性。 + +## 测试计划 + +### lit 测试 + +新增或扩展以下测试： + +- `plan_memory_order_by_size_*.pto` 继续作为 largest-first 覆盖。 +- `plan_memory_five_gates_lifetime_touching.pto` +- `plan_memory_five_gates_phi_family.pto` +- `plan_memory_five_gates_semantic_no_alias.pto` +- `plan_memory_five_gates_target_hazard.pto` +- 暂不新增 `plan_memory_five_gates_pipeline_load.pto`；闸门 4 当前为预留设计。 + +已有 `plan_memory_*.pto` 应继续保留 legacy + modern 双 RUN。 + +### 验证命令 + +```bash +cmake --build build --target ptoas -j8 + +PATH=/Users/fangrui/workspace/huawei/llvm21-workspace/llvm-project/llvm/build-assert/bin:$PATH \ + /Users/fangrui/workspace/huawei/llvm21-workspace/llvm-project/llvm/build-assert/bin/llvm-lit \ + -sv build/test/lit \ + --filter 'plan_memory' + +ctest --test-dir build --output-on-failure -L PTODSL +``` + +## 风险与注意事项 + +- `--plan-memory-order-by-size` 本身会改变 modern memplan 的 offset 分配顺序，测试应避免把 order-sensitive 预期错误地复用于默认路径。 +- 四道闸门是硬约束，不应为了容量不足而放松。 +- target hazard 需要先确认 PTOAS IR 中稳定的标记来源。 +- pipeline metadata 闸门当前不实现；如果未来启用，需要先定义稳定的 pipeline membership 来源。 +- phi family 豁免必须保守，不能让外部 live alias 借互斥分支错误复用。 +- legacy memplan 不应受该设计影响，默认行为保持不变。 diff --git a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp index 5d8a5e9789..c8906f60ef 100644 --- a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp +++ b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp @@ -15,6 +15,7 @@ #include "PTO/Transforms/InsertSync/InsertSyncDebug.h" #include "mlir/Interfaces/ViewLikeInterface.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Arith/IR/Arith.h" #include "llvm/Support/Debug.h" #define DEBUG_TYPE "pto-inject-sync" @@ -104,7 +105,64 @@ static Value GetRealRoot(Value v) { } return v; } - + +// Extract the base UB address carried by `rootBuffer` when it is the i64 +// address SSA operand of a `pto.pointer_cast` (single-address path). In that +// model `rootBuffer` is an `arith.constant` i64 whose value is the absolute UB +// offset, and `BaseMemInfo::baseAddresses` holds {0} (relative deltas). For +// the multi-address path `rootBuffer` is the `PointerCastOp` result itself +// (not an i64 constant), so this returns 0 and `baseAddresses` already holds +// the absolute slot offsets — making `rootBase + baseAddresses[i]` resolve to +// the absolute address in both models. +static uint64_t getRootBaseAddress(Value rootBuffer) { + if (!rootBuffer) + return 0; + if (auto cst = rootBuffer.getDefiningOp()) { + if (auto attr = dyn_cast(cst.getValue())) + return attr.getValue().getZExtValue(); + } + return 0; +} + +// Cross-root byte-range overlap check for local memory (UB/L1). +// +// `MemAlias` historically only checked byte-range overlap when the two +// `rootBuffer` SSA values were identical (same `alloc_tile`/`pointer_cast`). +// When PlanMemory reuses the same physical UB region for two *different* +// allocations, their `rootBuffer` values differ even though their byte ranges +// overlap — so `MemAlias` returned false and InsertSync silently dropped the +// cross-pipe hazard (issue #934: MTE3 tstore racing a V-pipe write into an +// overlapping-but-different-root buffer). +// +// This helper re-derives the absolute address as `getRootBaseAddress(root) + +// baseAddresses[i]` and performs the same `maxStart < minEnd` overlap test +// across every address pair, regardless of root identity. +static bool isLocalBufferOverlapCrossRoot(const BaseMemInfo *a, + const BaseMemInfo *b) { + // Conservative: if either side has no known address or size, assume alias. + if (a->baseAddresses.empty() || b->baseAddresses.empty()) + return true; + if (a->allocateSize == 0 || b->allocateSize == 0) + return true; + + uint64_t rootBaseA = getRootBaseAddress(a->rootBuffer); + uint64_t rootBaseB = getRootBaseAddress(b->rootBuffer); + + for (uint64_t addrA : a->baseAddresses) { + for (uint64_t addrB : b->baseAddresses) { + uint64_t aStart = rootBaseA + addrA; + uint64_t bStart = rootBaseB + addrB; + uint64_t aEnd = aStart + a->allocateSize; + uint64_t bEnd = bStart + b->allocateSize; + uint64_t maxStart = std::max(aStart, bStart); + uint64_t minEnd = std::min(aEnd, bEnd); + if (maxStart < minEnd) + return true; + } + } + return false; +} + bool MemoryDependentAnalyzer::DepBetween( const SmallVector &a, const SmallVector &b, diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 2e8d6cd3c2..9b56c7bff1 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -15,10 +15,11 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/MapVector.h" #include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallSet.h" using namespace mlir; using namespace mlir::pto; @@ -46,11 +47,35 @@ struct RootInfo { unsigned allocIndex = 0; unsigned freeIndex = 0; unsigned stableOrder = 0; + bool hasWriter = false; + bool hasUseBeforeFirstWrite = false; SmallVector offsets; }; using RootList = SmallVector; +struct ConflictFacts { + DenseMap forbidAlias; + DenseSet loadDerivedRoots; + DenseSet tpopConsumerRoots; + DenseMap> tpopConsumerWriteIndices; + DenseMap> phiFamilyIds; + bool targetHazardEnabled = false; +}; + +struct InplacePolicy { + bool notInplaceSafe = false; + SmallVector forbidOutputAliasOperands; +}; + +struct ReuseGroup { + AddressSpace space = AddressSpace::Zero; + uint64_t sizeBytes = 0; + uint64_t alignmentBytes = 1; + uint64_t offsetBytes = 0; + SmallVector members; +}; + static uint64_t alignUp(uint64_t value, uint64_t align) { if (align == 0) return value; @@ -138,17 +163,116 @@ static RootList unionRoots(const RootList &lhs, const RootList &rhs) { return result; } +static bool isOneOf(StringRef name, ArrayRef names) { + return llvm::is_contained(names, name); +} + +static InplacePolicy getInplacePolicy(Operation *op) { + StringRef name = op->getName().getStringRef(); + InplacePolicy policy; + + policy.notInplaceSafe = isOneOf( + name, { + "pto.tands", "pto.tfillpad_expand", "pto.tfmod", + "pto.tfmods", "pto.tgather", "pto.tmrgsort", + "pto.tors", "pto.trecip", "pto.trsqrt", + "pto.tsort32", "pto.ttrans", "pto.trowargmax", + "pto.trowargmin", "pto.trowmax", "pto.trowmin", + "pto.trowprod", "pto.trowsum", "pto.tcolargmax", + "pto.tcolargmin", "pto.txors", + }); + + if (name == "pto.tsel") { + policy.forbidOutputAliasOperands.push_back(0); // mask + policy.forbidOutputAliasOperands.push_back(3); // tmp + } + + if (isOneOf(name, { + "pto.trowexpand", + "pto.tcolexpand", + })) { + policy.forbidOutputAliasOperands.push_back(0); + } + + if (isOneOf(name, { + "pto.trowexpandadd", + "pto.trowexpanddiv", + "pto.trowexpandexpdif", + "pto.trowexpandmax", + "pto.trowexpandmin", + "pto.trowexpandmul", + "pto.trowexpandsub", + "pto.tcolexpandadd", + "pto.tcolexpanddiv", + "pto.tcolexpandexpdif", + "pto.tcolexpandmax", + "pto.tcolexpandmin", + "pto.tcolexpandmul", + "pto.tcolexpandsub", + })) { + policy.forbidOutputAliasOperands.push_back(1); + } + + return policy; +} + +static SmallVector getWrittenNonDpsOperands(Operation *op, + ValueRange dpsInits) { + SmallVector scratchOperands; + auto memEffect = dyn_cast(op); + if (!memEffect) + return scratchOperands; + + SmallVector, 8> effects; + memEffect.getEffects(effects); + for (const auto &effect : effects) { + if (!isa(effect.getEffect())) + continue; + Value value = effect.getValue(); + if (!value) + continue; + if (!llvm::is_contained(op->getOperands(), value)) + continue; + if (llvm::is_contained(dpsInits, value)) + continue; + if (!llvm::is_contained(scratchOperands, value)) + scratchOperands.push_back(value); + } + return scratchOperands; +} + +static bool hasReadEffectOnValue(Operation *op, Value value) { + auto memEffect = dyn_cast(op); + if (!memEffect) + return false; + + SmallVector, 8> effects; + memEffect.getEffects(effects); + for (const auto &effect : effects) { + if (!isa(effect.getEffect())) + continue; + if (effect.getValue() == value) + return true; + } + return false; +} + struct PlannerAnalysis { func::FuncOp func; DenseMap valueToRoots; + DenseSet splitTpopDerivedValues; SmallVector linearOps; DenseMap opToIndex; SmallVector roots; DenseMap rootIndexByValue; - DenseMap semanticConflictMap; + ConflictFacts facts; + unsigned nextPhiFamilyId = 0; bool failed = false; - explicit PlannerAnalysis(func::FuncOp func) : func(func) {} + explicit PlannerAnalysis(func::FuncOp func) : func(func) { + facts.targetHazardEnabled = + getTargetArch(func.getOperation()) == PTOArch::A3; + } RootList getRoots(Value value) const { auto it = valueToRoots.find(value); @@ -206,46 +330,242 @@ struct PlannerAnalysis { if (found == rootIndexByValue.end()) continue; RootInfo &info = roots[found->second]; + if (!info.hasWriter) + info.hasUseBeforeFirstWrite = true; if (index < info.allocIndex) info.allocIndex = index; info.freeIndex = std::max(info.freeIndex, index); } } - void addSemanticConflict(Value a, Value b) { + void markWrite(Value value, unsigned index, bool pureOverwrite) { + for (Value root : getRoots(value)) { + auto found = rootIndexByValue.find(root); + if (found == rootIndexByValue.end()) + continue; + RootInfo &info = roots[found->second]; + if (!info.hasWriter) { + if (pureOverwrite && !info.hasUseBeforeFirstWrite) + info.allocIndex = index; + info.hasWriter = true; + } + info.freeIndex = std::max(info.freeIndex, index); + } + } + + void addForbidAlias(Value a, Value b) { if (a == b) return; if (!rootIndexByValue.count(a) || !rootIndexByValue.count(b)) return; - appendUniqueRoot(semanticConflictMap[a], b); - appendUniqueRoot(semanticConflictMap[b], a); + appendUniqueRoot(facts.forbidAlias[a], b); + appendUniqueRoot(facts.forbidAlias[b], a); } - bool hasSemanticConflict(Value a, Value b) const { + bool hasForbidAlias(Value a, Value b) const { if (a == b) return false; - auto it = semanticConflictMap.find(a); - if (it == semanticConflictMap.end()) + auto it = facts.forbidAlias.find(a); + if (it == facts.forbidAlias.end()) return false; return llvm::is_contained(it->second, b); } - void recordLocalOpConflicts(Operation *op) { - RootList operandRoots; - RootList resultRoots; - for (Value operand : op->getOperands()) - operandRoots = unionRoots(operandRoots, getRoots(operand)); + void addForbidAliasBetweenRoots(const RootList &lhsRoots, + const RootList &rhsRoots) { + for (Value lhs : lhsRoots) + for (Value rhs : rhsRoots) + addForbidAlias(lhs, rhs); + } + + void markRoots(DenseSet &set, const RootList &roots) { + for (Value root : roots) + set.insert(root); + } + + bool rootsContain(const DenseSet &set, const RootList &roots) const { + return llvm::any_of(roots, [&](Value root) { return set.contains(root); }); + } + + bool isSplitTpopDerived(Value value) const { + return splitTpopDerivedValues.contains(value); + } + + bool operandsContainSplitTpopDerived(Operation *op) const { + return llvm::any_of(op->getOperands(), [&](Value operand) { + return isSplitTpopDerived(operand); + }); + } + + bool operandsContainLoadDerivedRoot(Operation *op) const { + return llvm::any_of(op->getOperands(), [&](Value operand) { + return rootsContain(facts.loadDerivedRoots, getRoots(operand)); + }); + } + + void propagateSplitTpopDerived(Value result, ValueRange sources) { + if (!result) + return; + if (llvm::any_of(sources, [&](Value source) { + return isSplitTpopDerived(source); + })) + splitTpopDerivedValues.insert(result); + } + + void propagateSplitTpopDerivedFromOperands(Operation *op) { + if (!operandsContainSplitTpopDerived(op)) + return; for (Value result : op->getResults()) - resultRoots = unionRoots(resultRoots, getRoots(result)); + splitTpopDerivedValues.insert(result); + } + + bool isRootDefinedInRegion(Value root, Region ®ion) const { + auto it = rootIndexByValue.find(root); + if (it == rootIndexByValue.end()) + return false; + + Operation *defOp = roots[it->second].defOp; + for (Operation *cur = defOp; cur; cur = cur->getParentOp()) { + if (cur->getParentRegion() == ®ion) + return true; + } + return false; + } + + void addPhiFamily(Value root, unsigned familyId) { + if (!rootIndexByValue.count(root)) + return; + SmallVector &families = facts.phiFamilyIds[root]; + if (!llvm::is_contained(families, familyId)) + families.push_back(familyId); + } + + void recordIfPhiFamilies(scf::IfOp ifOp) { + if (ifOp.getNumResults() == 0) + return; + + auto thenYield = cast(ifOp.thenBlock()->getTerminator()); + auto elseYield = cast(ifOp.elseBlock()->getTerminator()); + for (auto [thenVal, elseVal] : + llvm::zip(thenYield.getResults(), elseYield.getResults())) { + RootList roots = unionRoots(getRoots(thenVal), getRoots(elseVal)); + SmallVector branchLocalRoots; + for (Value root : roots) { + if (isRootDefinedInRegion(root, ifOp.getThenRegion()) || + isRootDefinedInRegion(root, ifOp.getElseRegion())) + branchLocalRoots.push_back(root); + } + if (branchLocalRoots.size() < 2) + continue; + + unsigned familyId = nextPhiFamilyId++; + for (Value root : branchLocalRoots) + addPhiFamily(root, familyId); + } + } + + void recordDpsScratchConflicts(Operation *op, ValueRange dpsInits) { + RootList outputRoots; + for (Value init : dpsInits) + outputRoots = unionRoots(outputRoots, getRoots(init)); + if (outputRoots.empty()) + return; + + for (Value scratch : getWrittenNonDpsOperands(op, dpsInits)) + addForbidAliasBetweenRoots(getRoots(scratch), outputRoots); + } + + void recordInplacePolicyConflicts(Operation *op, ValueRange dpsInits) { + RootList outputRoots; + for (Value init : dpsInits) + outputRoots = unionRoots(outputRoots, getRoots(init)); + if (outputRoots.empty()) + return; + + InplacePolicy policy = getInplacePolicy(op); + if (policy.notInplaceSafe) { + for (Value operand : op->getOperands()) { + if (llvm::is_contained(dpsInits, operand)) + continue; + addForbidAliasBetweenRoots(getRoots(operand), outputRoots); + } + } + + for (unsigned operandIndex : policy.forbidOutputAliasOperands) { + if (operandIndex >= op->getNumOperands()) + continue; + Value operand = op->getOperand(operandIndex); + if (llvm::is_contained(dpsInits, operand)) + continue; + addForbidAliasBetweenRoots(getRoots(operand), outputRoots); + } + } + + void recordDpsInplaceConflicts(Operation *op) { + auto dpsOp = dyn_cast(op); + if (!dpsOp) + return; - if (operandRoots.empty() || resultRoots.empty()) + ValueRange dpsInits = dpsOp.getDpsInits(); + recordDpsScratchConflicts(op, dpsInits); + recordInplacePolicyConflicts(op, dpsInits); + } + + void recordLoadDerivedRoots(Operation *op, ValueRange dpsInits) { + if (!isa(op)) + return; + + for (Value init : dpsInits) + markRoots(facts.loadDerivedRoots, getRoots(init)); + } + + void recordTpopConsumerRoots(Operation *op, ValueRange dpsInits, + unsigned opIndex) { + if (!facts.targetHazardEnabled) return; - for (Value lhs : operandRoots) { - for (Value rhs : resultRoots) - addSemanticConflict(lhs, rhs); + if (!operandsContainSplitTpopDerived(op) || !operandsContainLoadDerivedRoot(op)) + return; + + RootList outputRoots; + for (Value init : dpsInits) + outputRoots = unionRoots(outputRoots, getRoots(init)); + markRoots(facts.tpopConsumerRoots, outputRoots); + for (Value root : outputRoots) { + SmallVector &indices = facts.tpopConsumerWriteIndices[root]; + if (!llvm::is_contained(indices, opIndex)) + indices.push_back(opIndex); } } + void recordDpsTargetHazardFacts(Operation *op, unsigned opIndex) { + auto dpsOp = dyn_cast(op); + if (!dpsOp) + return; + + ValueRange dpsInits = dpsOp.getDpsInits(); + recordLoadDerivedRoots(op, dpsInits); + recordTpopConsumerRoots(op, dpsInits, opIndex); + } + + void recordSplitTpopDerivedValue(Operation *op) { + if (!facts.targetHazardEnabled) + return; + + if (auto pop = dyn_cast(op)) { + if (pop.getSplit() != 0) + splitTpopDerivedValues.insert(pop.getTile()); + return; + } + + if (auto pop = dyn_cast(op)) { + if (pop.getSplit() != 0) + splitTpopDerivedValues.insert(pop.getTile()); + return; + } + + propagateSplitTpopDerivedFromOperands(op); + } + void seedForIterArgAliases(scf::ForOp forOp) { if (forOp.getRegion().empty()) return; @@ -288,24 +608,47 @@ struct PlannerAnalysis { if (auto bind = dyn_cast(op)) { setRoots(bind.getResult(), getRoots(bind.getSource())); + propagateSplitTpopDerived(bind.getResult(), ValueRange{bind.getSource()}); } else if (auto slotMarker = dyn_cast(op)) { setRoots(slotMarker.getResult(), getRoots(slotMarker.getSource())); + propagateSplitTpopDerived(slotMarker.getResult(), + ValueRange{slotMarker.getSource()}); } else if (auto select = dyn_cast(op)) { setRoots(select.getResult(), unionRoots(getRoots(select.getTrueValue()), getRoots(select.getFalseValue()))); + propagateSplitTpopDerived(select.getResult(), + ValueRange{select.getTrueValue(), + select.getFalseValue()}); } else if (auto castOp = dyn_cast(op)) { setRoots(castOp.getResult(), getRoots(castOp.getSource())); + propagateSplitTpopDerived(castOp.getResult(), + ValueRange{castOp.getSource()}); } else if (auto subview = dyn_cast(op)) { setRoots(subview.getResult(), getRoots(subview.getSource())); + propagateSplitTpopDerived(subview.getResult(), + ValueRange{subview.getSource()}); } else if (auto reinterpret = dyn_cast(op)) { setRoots(reinterpret.getResult(), getRoots(reinterpret.getSource())); + propagateSplitTpopDerived(reinterpret.getResult(), + ValueRange{reinterpret.getSource()}); } else if (auto forOp = dyn_cast(op)) { seedForIterArgAliases(forOp); } - for (Value operand : op->getOperands()) + auto dpsOp = dyn_cast(op); + ValueRange dpsInits = dpsOp ? dpsOp.getDpsInits() : ValueRange(); + for (Value operand : op->getOperands()) { + if (llvm::is_contained(dpsInits, operand)) + continue; markUse(operand, index); + } + for (Value init : dpsInits) { + bool readsOldValue = hasReadEffectOnValue(op, init); + if (readsOldValue) + markUse(init, index); + markWrite(init, index, /*pureOverwrite=*/!readsOldValue); + } for (Region &nested : op->getRegions()) { walkRegion(nested); @@ -325,6 +668,7 @@ struct PlannerAnalysis { setRoots(result, unionRoots(getRoots(thenVal), getRoots(elseVal))); } + recordIfPhiFamilies(ifOp); } } else if (auto forOp = dyn_cast(op)) { auto yieldOp = cast(forOp.getBody()->getTerminator()); @@ -335,19 +679,85 @@ struct PlannerAnalysis { } } - recordLocalOpConflicts(op); + recordSplitTpopDerivedValue(op); + recordDpsTargetHazardFacts(op, index); + recordDpsInplaceConflicts(op); } } } }; -static bool lifetimesOverlap(const RootInfo &lhs, const RootInfo &rhs) { - return !(lhs.freeIndex < rhs.allocIndex || rhs.freeIndex < lhs.allocIndex); +static bool lifetimesStrictlyOverlap(const RootInfo &lhs, const RootInfo &rhs) { + return !(lhs.freeIndex <= rhs.allocIndex || rhs.freeIndex <= lhs.allocIndex); +} + +static bool sharePhiFamily(Value lhs, Value rhs, const ConflictFacts &facts) { + auto lhsIt = facts.phiFamilyIds.find(lhs); + auto rhsIt = facts.phiFamilyIds.find(rhs); + if (lhsIt == facts.phiFamilyIds.end() || rhsIt == facts.phiFamilyIds.end()) + return false; + + for (unsigned lhsFamily : lhsIt->second) { + if (llvm::is_contained(rhsIt->second, lhsFamily)) + return true; + } + return false; +} + +static bool gateLifetimeAndPhi(const RootInfo &lhs, const RootInfo &rhs, + const ConflictFacts &facts) { + if (!lifetimesStrictlyOverlap(lhs, rhs)) + return true; + + // Gate 5 is intentionally embedded in Gate 1: branch-local roots yielded by + // the same scf.if result are mutually exclusive at runtime, so their + // post-phi liveness extension does not by itself forbid reuse. + return sharePhiFamily(lhs.root, rhs.root, facts); } -static bool intervalsOverlap(uint64_t lhsOffset, uint64_t lhsSize, - uint64_t rhsOffset, uint64_t rhsSize) { - return lhsOffset < rhsOffset + rhsSize && rhsOffset < lhsOffset + lhsSize; +static bool hasTargetHazard(const RootInfo &input, const RootInfo &writer, + const ConflictFacts &facts) { + if (!facts.targetHazardEnabled) + return false; + + if (!facts.loadDerivedRoots.contains(input.root) || + !facts.tpopConsumerRoots.contains(writer.root)) + return false; + + auto found = facts.tpopConsumerWriteIndices.find(writer.root); + if (found == facts.tpopConsumerWriteIndices.end()) + return false; + + return llvm::is_contained(found->second, input.freeIndex); +} + +static bool gateTargetHazard(const RootInfo &lhs, const RootInfo &rhs, + const ConflictFacts &facts) { + return !hasTargetHazard(lhs, rhs, facts) && !hasTargetHazard(rhs, lhs, facts); +} + +static bool gateSemanticNoAlias(const RootInfo &lhs, const RootInfo &rhs, + const PlannerAnalysis &analysis) { + return !analysis.hasForbidAlias(lhs.root, rhs.root); +} + +static bool canShare(const RootInfo &lhs, const RootInfo &rhs, + const PlannerAnalysis &analysis) { + const ConflictFacts &facts = analysis.facts; + return gateLifetimeAndPhi(lhs, rhs, facts) && + gateTargetHazard(lhs, rhs, facts) && + gateSemanticNoAlias(lhs, rhs, analysis); +} + +static bool canJoinReuseGroup(const RootInfo &info, const ReuseGroup &group, + const PlannerAnalysis &analysis) { + if (info.space != group.space) + return false; + for (const RootInfo *member : group.members) { + if (!canShare(info, *member, analysis)) + return false; + } + return true; } static SmallVector buildSlotOffsets(uint64_t base, uint64_t slotBytes, @@ -507,152 +917,138 @@ static LogicalResult materializePlannedOffsets( LogicalResult mlir::pto::runModernPlanMemory(func::FuncOp func, llvm::StringRef memMode, bool orderBySize) { - if (!memMode.equals_insensitive("local")) { - func.emitError("unsupported mem-mode '") - << memMode << "'; only 'local' is currently implemented"; - return failure(); - } - - PlannerAnalysis analysis(func); - analysis.walkRegion(func.getBody()); - if (analysis.failed) { - return failure(); - } - - DenseMap> buffer2Offsets; - llvm::MapVector> rootsBySpace; - for (RootInfo &info : analysis.roots) - rootsBySpace[info.space].push_back(&info); - - for (auto &entry : rootsBySpace) { - AddressSpace space = entry.first; - SmallVector &roots = entry.second; - MemSpec spec = getMemSpec(getTargetArch(func), space); - - llvm::stable_sort(roots, [&](const RootInfo *lhs, const RootInfo *rhs) { - if (orderBySize && lhs->totalBytes != rhs->totalBytes) - return lhs->totalBytes > rhs->totalBytes; - if (lhs->allocIndex != rhs->allocIndex) - return lhs->allocIndex < rhs->allocIndex; - return lhs->stableOrder < rhs->stableOrder; - }); - - SmallVector planned; - uint64_t scopeRequiredBytes = 0; - for (RootInfo *info : roots) { - llvm::SmallSet candidateSet; - candidateSet.insert(0); - for (RootInfo *other : planned) { - candidateSet.insert(other->offsets.front()); - candidateSet.insert( - alignUp(other->offsets.front() + other->totalBytes, - info->alignmentBytes)); - } - - SmallVector candidates(candidateSet.begin(), - candidateSet.end()); - llvm::sort(candidates); + if (!memMode.equals_insensitive("local")) { + func.emitError("unsupported mem-mode '") + << memMode << "'; only 'local' is currently implemented"; + return failure(); + } - std::optional chosen; - for (uint64_t candidate : candidates) { - candidate = alignUp(candidate, info->alignmentBytes); - if (candidate + info->totalBytes > spec.capacityBytes) - continue; + PlannerAnalysis analysis(func); + analysis.walkRegion(func.getBody()); + if (analysis.failed) { + return failure(); + } - bool conflict = false; - for (RootInfo *other : planned) { - if (!intervalsOverlap(candidate, info->totalBytes, - other->offsets.front(), other->totalBytes)) - continue; - if (lifetimesOverlap(*info, *other) || - analysis.hasSemanticConflict(info->root, other->root)) { - conflict = true; - break; - } - } - if (!conflict) { - chosen = candidate; - break; - } - } + DenseMap> buffer2Offsets; + llvm::MapVector> rootsBySpace; + for (RootInfo &info : analysis.roots) + rootsBySpace[info.space].push_back(&info); + + for (auto &entry : rootsBySpace) { + AddressSpace space = entry.first; + SmallVector &roots = entry.second; + MemSpec spec = getMemSpec(getTargetArch(func), space); + + llvm::stable_sort(roots, [&](const RootInfo *lhs, const RootInfo *rhs) { + if (orderBySize && lhs->totalBytes != rhs->totalBytes) + return lhs->totalBytes > rhs->totalBytes; + if (lhs->allocIndex != rhs->allocIndex) + return lhs->allocIndex < rhs->allocIndex; + return lhs->stableOrder < rhs->stableOrder; + }); - if (!chosen) { - uint64_t fallback = 0; - for (RootInfo *other : planned) - fallback = - std::max(fallback, other->offsets.front() + other->totalBytes); - fallback = alignUp(fallback, info->alignmentBytes); - chosen = fallback; + SmallVector groups; + uint64_t scopeRequiredBytes = 0; + for (RootInfo *info : roots) { + ReuseGroup *chosen = nullptr; + for (ReuseGroup &group : groups) { + if (canJoinReuseGroup(*info, group, analysis)) { + chosen = &group; + break; } + } - info->offsets = - buildSlotOffsets(*chosen, info->slotBytes, info->slotCount); - scopeRequiredBytes = - std::max(scopeRequiredBytes, *chosen + info->totalBytes); - planned.push_back(info); + if (!chosen) { + ReuseGroup group; + group.space = space; + group.alignmentBytes = info->alignmentBytes; + groups.push_back(std::move(group)); + chosen = &groups.back(); } - if (scopeRequiredBytes > spec.capacityBytes) { - func.emitError() << stringifyEnum(space) << " overflow, requires " - << (scopeRequiredBytes * 8) << " bits while " - << (spec.capacityBytes * 8) << " bits available"; - return failure(); + chosen->members.push_back(info); + chosen->sizeBytes = std::max(chosen->sizeBytes, info->totalBytes); + chosen->alignmentBytes = + std::max(chosen->alignmentBytes, info->alignmentBytes); + } + + uint64_t cursor = 0; + for (ReuseGroup &group : groups) { + cursor = alignUp(cursor, group.alignmentBytes); + group.offsetBytes = cursor; + if (group.offsetBytes + group.sizeBytes > spec.capacityBytes) { + scopeRequiredBytes = group.offsetBytes + group.sizeBytes; + break; } - for (RootInfo *info : roots) - buffer2Offsets[info->root] = info->offsets; + for (RootInfo *member : group.members) { + member->offsets = buildSlotOffsets(group.offsetBytes, member->slotBytes, + member->slotCount); + } + scopeRequiredBytes = + std::max(scopeRequiredBytes, group.offsetBytes + group.sizeBytes); + cursor = group.offsetBytes + group.sizeBytes; } - DenseMap>> - occupiedBySpace; - for (const RootInfo &info : analysis.roots) { - if (info.offsets.empty()) - continue; - occupiedBySpace[info.space].push_back( - {info.offsets.front(), info.offsets.front() + info.totalBytes}); + if (scopeRequiredBytes > spec.capacityBytes) { + func.emitError() << stringifyEnum(space) << " overflow, requires " + << (scopeRequiredBytes * 8) << " bits while " + << (spec.capacityBytes * 8) << " bits available"; + return failure(); } - MLIRContext *ctx = func.getContext(); - PTOArch arch = getTargetArch(func); - func.walk([&](pto::ReserveBufferOp reserveOp) -> WalkResult { - if (mlir::failed(validateReserveBufferForPlanMemory(reserveOp, arch))) { - analysis.failed = true; - return WalkResult::interrupt(); - } + for (RootInfo *info : roots) + buffer2Offsets[info->root] = info->offsets; + } - AddressSpace space = reserveOp.getLocation().getAddressSpace(); - MemSpec spec = getMemSpec(arch, space); - auto baseOr = - planReserveBufferBase(reserveOp, spec, occupiedBySpace[space]); - if (mlir::failed(baseOr)) { - reserveOp.emitError("failed to reserve a local memory hole for " - "reserve_buffer"); - analysis.failed = true; - return WalkResult::interrupt(); - } + DenseMap>> + occupiedBySpace; + for (const RootInfo &info : analysis.roots) { + if (info.offsets.empty()) + continue; + occupiedBySpace[info.space].push_back( + {info.offsets.front(), info.offsets.front() + info.totalBytes}); + } - reserveOp->setAttr("base", - IntegerAttr::get(IntegerType::get(ctx, 32), - static_cast(*baseOr))); - return WalkResult::advance(); - }); + MLIRContext *ctx = func.getContext(); + PTOArch arch = getTargetArch(func); + func.walk([&](pto::ReserveBufferOp reserveOp) -> WalkResult { + if (mlir::failed(validateReserveBufferForPlanMemory(reserveOp, arch))) { + analysis.failed = true; + return WalkResult::interrupt(); + } - if (analysis.failed || - mlir::failed(materializePlannedOffsets(func, buffer2Offsets))) { - return failure(); + AddressSpace space = reserveOp.getLocation().getAddressSpace(); + MemSpec spec = getMemSpec(arch, space); + auto baseOr = + planReserveBufferBase(reserveOp, spec, occupiedBySpace[space]); + if (mlir::failed(baseOr)) { + reserveOp.emitError("failed to reserve a local memory hole for " + "reserve_buffer"); + analysis.failed = true; + return WalkResult::interrupt(); } - bool hasUnplannedAllocTile = false; - func.walk([&](pto::AllocTileOp op) { - if (op.getAddr()) - return; - op.emitError( - "PTOPlanMemory failed to assign an address to pto.alloc_tile"); - hasUnplannedAllocTile = true; - }); - if (hasUnplannedAllocTile) - return failure(); - return success(); + reserveOp->setAttr("base", IntegerAttr::get(IntegerType::get(ctx, 32), + static_cast(*baseOr))); + return WalkResult::advance(); + }); + + if (analysis.failed || + mlir::failed(materializePlannedOffsets(func, buffer2Offsets))) { + return failure(); + } + + bool hasUnplannedAllocTile = false; + func.walk([&](pto::AllocTileOp op) { + if (op.getAddr()) + return; + op.emitError("PTOPlanMemory failed to assign an address to pto.alloc_tile"); + hasUnplannedAllocTile = true; + }); + if (hasUnplannedAllocTile) + return failure(); + return success(); } namespace { @@ -681,8 +1077,8 @@ struct PlanMemoryModernPass void runOnOperation() override { ModuleOp moduleOp = getOperation(); for (func::FuncOp funcOp : moduleOp.getOps()) { - if (failed(mlir::pto::runModernPlanMemory(funcOp, memMode, - orderBySize))) { + if (failed( + mlir::pto::runModernPlanMemory(funcOp, memMode, orderBySize))) { signalPassFailure(); return; } diff --git a/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto b/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto new file mode 100644 index 0000000000..3bf501afe1 --- /dev/null +++ b/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ +// RUN: | FileCheck %s + +module attributes {"pto.target_arch" = "a3"} { + func.func @five_gates_lifetime_overlap( + %gm0 : memref<16x16xf16, #pto.address_space>, + %gm1 : memref<16x16xf16, #pto.address_space>, + %out0 : memref<16x16xf16, #pto.address_space>, + %out1 : memref<16x16xf16, #pto.address_space>) { + %buf0 = memref.alloc() : memref<16x16xf16, #pto.address_space> + %buf1 = memref.alloc() : memref<16x16xf16, #pto.address_space> + + pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) + outs(%buf0 : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + pto.tload ins(%gm1 : memref<16x16xf16, #pto.address_space>) + outs(%buf1 : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + + pto.tstore ins(%buf0 : memref<16x16xf16, #pto.address_space>) + outs(%out0 : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + pto.tstore ins(%buf1 : memref<16x16xf16, #pto.address_space>) + outs(%out1 : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + return + } +} + +// CHECK-LABEL: func.func @five_gates_lifetime_overlap +// CHECK-DAG: %[[ADDR0:.*]] = arith.constant 0 : i64 +// CHECK-DAG: %[[ADDR1:.*]] = arith.constant 512 : i64 +// CHECK: pto.pointer_cast(%[[ADDR0]]) +// CHECK: pto.pointer_cast(%[[ADDR1]]) +// CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_five_gates_phi_family.pto b/test/lit/pto/plan_memory_five_gates_phi_family.pto new file mode 100644 index 0000000000..ddf295fd80 --- /dev/null +++ b/test/lit/pto/plan_memory_five_gates_phi_family.pto @@ -0,0 +1,43 @@ +// 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-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ +// RUN: | FileCheck %s + +module attributes {"pto.target_arch" = "a3"} { + func.func @five_gates_phi_family( + %gm0 : memref<16x16xf16, #pto.address_space>, + %gm1 : memref<16x16xf16, #pto.address_space>, + %out : memref<16x16xf16, #pto.address_space>) { + %cond = arith.constant true + %merged = scf.if %cond -> (memref<16x16xf16, #pto.address_space>) { + %then = memref.alloc() : memref<16x16xf16, #pto.address_space> + pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) + outs(%then : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + scf.yield %then : memref<16x16xf16, #pto.address_space> + } else { + %else = memref.alloc() : memref<16x16xf16, #pto.address_space> + pto.tload ins(%gm1 : memref<16x16xf16, #pto.address_space>) + outs(%else : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + scf.yield %else : memref<16x16xf16, #pto.address_space> + } + + pto.tstore ins(%merged : memref<16x16xf16, #pto.address_space>) + outs(%out : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + return + } +} + +// CHECK-LABEL: func.func @five_gates_phi_family +// CHECK: %[[ADDR0:.*]] = arith.constant 0 : i64 +// CHECK: scf.if +// CHECK: pto.pointer_cast(%[[ADDR0]]) +// CHECK: } else { +// CHECK: pto.pointer_cast(%[[ADDR0]]) +// CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_inplace_forbid_alias.pto b/test/lit/pto/plan_memory_inplace_forbid_alias.pto new file mode 100644 index 0000000000..54e65a0fd9 --- /dev/null +++ b/test/lit/pto/plan_memory_inplace_forbid_alias.pto @@ -0,0 +1,60 @@ +// 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-level=level2 --pto-arch=a3 --plan-memory-impl=modern \ +// RUN: --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s 2>&1 \ +// RUN: | FileCheck %s + +module attributes {"pto.target_arch" = "a3"} { + func.func @inplace_forbid_alias_ttrans() { + %src = memref.alloc() : memref<16x16xf16, #pto.address_space> + %tmp = memref.alloc() : memref<16x16xf16, #pto.address_space> + %dst = memref.alloc() : memref<16x16xf16, #pto.address_space> + + pto.ttrans ins(%src, %tmp : memref<16x16xf16, #pto.address_space>, + memref<16x16xf16, #pto.address_space>) + outs(%dst : memref<16x16xf16, #pto.address_space>) + return + } + + func.func @inplace_forbid_alias_tsel() { + %mask = memref.alloc() : memref<16x16xi16, #pto.address_space> + %src0 = memref.alloc() : memref<16x16xf16, #pto.address_space> + %src1 = memref.alloc() : memref<16x16xf16, #pto.address_space> + %tmp = memref.alloc() : memref<16x16xui8, #pto.address_space> + %dst = memref.alloc() : memref<16x16xf16, #pto.address_space> + + pto.tsel ins(%mask, %src0, %src1, %tmp + : memref<16x16xi16, #pto.address_space>, + memref<16x16xf16, #pto.address_space>, + memref<16x16xf16, #pto.address_space>, + memref<16x16xui8, #pto.address_space>) + outs(%dst : memref<16x16xf16, #pto.address_space>) + return + } +} + +// CHECK-LABEL: func.func @inplace_forbid_alias_ttrans +// CHECK-DAG: %[[A0:.*]] = arith.constant 0 : i64 +// CHECK-DAG: %[[A512:.*]] = arith.constant 512 : i64 +// CHECK-DAG: %[[A1024:.*]] = arith.constant 1024 : i64 +// CHECK: pto.pointer_cast(%[[A0]]) +// CHECK: pto.pointer_cast(%[[A512]]) +// CHECK: pto.pointer_cast(%[[A1024]]) +// CHECK-LABEL: func.func @inplace_forbid_alias_tsel +// CHECK-DAG: %[[B0:.*]] = arith.constant 0 : i64 +// CHECK-DAG: %[[B512:.*]] = arith.constant 512 : i64 +// CHECK-DAG: %[[B1024:.*]] = arith.constant 1024 : i64 +// CHECK-DAG: %[[B1536:.*]] = arith.constant 1536 : i64 +// CHECK-DAG: %[[B2048:.*]] = arith.constant 2048 : i64 +// CHECK: pto.pointer_cast(%[[B0]]) +// CHECK: pto.pointer_cast(%[[B512]]) +// CHECK: pto.pointer_cast(%[[B1024]]) +// CHECK: pto.pointer_cast(%[[B2048]]) +// CHECK: pto.pointer_cast(%[[B1536]]) +// CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_order_by_size_noreuse.pto b/test/lit/pto/plan_memory_order_by_size_noreuse.pto index fbfdfd6898..88ef992031 100644 --- a/test/lit/pto/plan_memory_order_by_size_noreuse.pto +++ b/test/lit/pto/plan_memory_order_by_size_noreuse.pto @@ -8,7 +8,7 @@ // contract that the option means the same thing whether or not reuse kicks in. // // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT -// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT +// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE // RUN: ptoas --pto-arch=a3 --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE // RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE diff --git a/test/lit/pto/plan_memory_order_by_size_reuse.pto b/test/lit/pto/plan_memory_order_by_size_reuse.pto index 5bf11d33b5..b83ff45090 100644 --- a/test/lit/pto/plan_memory_order_by_size_reuse.pto +++ b/test/lit/pto/plan_memory_order_by_size_reuse.pto @@ -6,7 +6,7 @@ // (first-fit-decreasing) the largest tile is allocated first and gets offset 0. // // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT -// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=DEFAULT +// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE // RUN: ptoas --pto-arch=a3 --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE // RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --plan-memory-order-by-size --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=BYSIZE diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 4b77f7faa7..62aef35cc8 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -324,10 +324,12 @@ static llvm::cl::opt enableInsertSync("enable-insert-sync", llvm::cl::desc("Enable automatic synchronization insertion pass"), llvm::cl::init(false)); -[[maybe_unused]] static llvm::cl::opt planMemoryOrderBySize( +static llvm::cl::opt planMemoryOrderBySize( "plan-memory-order-by-size", llvm::cl::desc("Plan larger local buffers first inside one AddressSpace " - "before applying the basic SPEC_LEVEL_0 reuse strategy"), + "before applying the basic SPEC_LEVEL_0 reuse strategy. " + "Defaults to true when --plan-memory-impl=modern is " + "explicitly selected"), llvm::cl::init(false)); static llvm::cl::opt planMemoryImpl( @@ -3103,7 +3105,12 @@ int mlir::pto::compilePTOASModule( if (effectiveLevel != PTOBuildLevel::Level3) { pto::PlanMemoryOptions planMemoryOptions; planMemoryOptions.memMode = "local"; - planMemoryOptions.orderBySize = planMemoryOrderBySize; + bool effectivePlanMemoryOrderBySize = planMemoryOrderBySize; + if (planMemoryImpl == "modern" && + planMemoryOrderBySize.getNumOccurrences() == 0) { + effectivePlanMemoryOrderBySize = true; + } + planMemoryOptions.orderBySize = effectivePlanMemoryOrderBySize; if (planMemoryImpl == "legacy") { pm.addPass(pto::createPlanMemoryPass(planMemoryOptions)); } else if (planMemoryImpl == "modern") { From c789dd5e70b85f86fb7965b02c971f4fe29c5be3 Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 13 Jul 2026 20:59:12 +0800 Subject: [PATCH 03/91] Fix cross-root local overlap sync analysis build --- .../InsertSync/MemoryDependentAnalyzer.cpp | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp index c8906f60ef..48a413ea3b 100644 --- a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp +++ b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp @@ -254,8 +254,20 @@ bool MemoryDependentAnalyzer::MemAlias(const BaseMemInfo *a, if (isTraceEnabled()) llvm::errs() << " -> Mismatch. Real roots differ.\n"; } - - return false; + + // 2.3 Cross-root absolute-address overlap check. + // + // Roots genuinely differ (different alloc_tile / different pointer_cast + // address SSA). Historically MemAlias returned false here, but PlanMemory + // can reuse the same physical UB region for distinct allocations whose + // byte ranges overlap. Re-derive the absolute address and check overlap so + // a cross-pipe hazard (e.g. MTE3 tstore vs a V-pipe write into an + // overlapping region) is not silently dropped (issue #934). + bool crossRootOverlap = isLocalBufferOverlapCrossRoot(a, b); + if (isTraceEnabled()) + llvm::errs() << " -> Cross-root overlap check: " + << (crossRootOverlap ? "true" : "false") << "\n"; + return crossRootOverlap; } bool MemoryDependentAnalyzer::isGMBufferOverlap(const BaseMemInfo *a, From e2d8d08beaf73ee3ca47bc774cfb09ac7b9c4438 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 14 Jul 2026 12:38:38 +0800 Subject: [PATCH 04/91] Fix tile-native memplan CI regressions --- .../InsertSync/MemoryDependentAnalyzer.cpp | 39 +++++++- .../Transforms/PTOMaterializeTileHandles.cpp | 91 +++++++++++++++++++ lib/PTO/Transforms/PTOPlanMemory.cpp | 58 +++++++++++- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 62 +++++++++++++ lib/PTO/Transforms/PTOViewToMemref.cpp | 48 ++++++++++ lib/PTO/Transforms/Utils.cpp | 7 ++ test/lit/pto/a5_unused_tmp_vec_overflow.pto | 12 +-- .../pto/alloc_tile_low_precision_valid.pto | 4 +- .../pto/compact_left_blayout_parser_a3.pto | 2 +- .../pto/compact_left_blayout_parser_a5.pto | 2 +- .../emitc_tile_data_sink_after_tassign.pto | 11 +-- ...fixpipe_frontend_emitc_vector_remat_a5.pto | 6 +- ...pe_frontend_quant_state_preserve_ir_a5.pto | 4 +- test/lit/pto/graph_sync_solver_basic.pto | 2 +- .../pto/issue531_tprelu_vec_overflow_a5.pto | 12 +-- ...ov_treshape_square_dynamic_valid_shape.pto | 8 +- test/lit/pto/issue708_zero_valid_subview.pto | 6 +- .../issue708_zero_valid_subview_inferred.pto | 12 +-- .../pto/mgather_mscatter_a5_attrs_emitc.pto | 8 +- .../pto/mgather_mscatter_base_a3_emitc.pto | 8 +- .../pto/ptodsl_subkernel_call_autosync.pto | 7 +- .../pto/subview_bind_tile_preserve_stride.pto | 9 +- ..._column_slice_preserve_parent_shape_a5.pto | 3 +- ...w_col_major_noncompact_preserve_stride.pto | 3 +- ...w_col_major_row_plus_one_stride_offset.pto | 5 +- ...dynamic_offset_static_valid_regression.pto | 3 +- test/lit/pto/subview_explicit_result_type.pto | 2 +- .../subview_row_plus_one_stride_offset.pto | 5 +- test/lit/pto/subview_validshape_guard.pto | 8 +- ...shape_partial_parent_requires_explicit.pto | 2 +- test/lit/pto/tile_compact_mode_emitc.pto | 6 +- test/lit/pto/tpow_emitc.pto | 4 +- test/lit/pto/tpow_int_no_tmp.pto | 4 +- test/lit/pto/tpows_emitc.pto | 4 +- test/lit/pto/tpows_int_no_tmp.pto | 2 +- test/lit/pto/tquant_dynamic_tmp_valid_a3.pto | 8 +- test/lit/pto/trem_emitc.pto | 8 +- test/lit/pto/trems_emitc.pto | 4 +- ...explicit_dynamic_valid_shape_preserved.pto | 3 +- test/lit/pto/trsqrt_emitc.pto | 8 +- test/lit/pto/tsort32_emitc.pto | 8 +- ...p_fusion_adapter_placement_level2_tadd.pto | 9 +- ...p_fusion_adapter_placement_level3_tadd.pto | 14 +-- .../op_fusion_region_pipeline_level2.pto | 4 +- ...rmalize_uncovered_tile_sections_vector.pto | 2 +- 45 files changed, 399 insertions(+), 138 deletions(-) diff --git a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp index 48a413ea3b..2088f159db 100644 --- a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp +++ b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp @@ -124,6 +124,35 @@ static uint64_t getRootBaseAddress(Value rootBuffer) { return 0; } +static bool isIntegerConstant(Value value) { + if (!value) + return false; + auto cst = value.getDefiningOp(); + return cst && isa(cst.getValue()); +} + +static bool hasKnownLocalAbsoluteAddress(const BaseMemInfo *info) { + if (!info || info->baseAddresses.empty() || info->allocateSize == 0) + return false; + + Value root = info->rootBuffer; + if (!root) + return false; + + if (auto alloc = root.getDefiningOp()) + return isIntegerConstant(alloc.getAddr()); + + if (isIntegerConstant(root)) + return true; + + // Multi-address pointer_cast records absolute slot offsets in baseAddresses + // and uses the pointer_cast result as rootBuffer. + if (auto ptrCast = root.getDefiningOp()) + return ptrCast.getAddrs().size() > 1; + + return false; +} + // Cross-root byte-range overlap check for local memory (UB/L1). // // `MemAlias` historically only checked byte-range overlap when the two @@ -139,11 +168,11 @@ static uint64_t getRootBaseAddress(Value rootBuffer) { // across every address pair, regardless of root identity. static bool isLocalBufferOverlapCrossRoot(const BaseMemInfo *a, const BaseMemInfo *b) { - // Conservative: if either side has no known address or size, assume alias. - if (a->baseAddresses.empty() || b->baseAddresses.empty()) - return true; - if (a->allocateSize == 0 || b->allocateSize == 0) - return true; + // Cross-root checks are only meaningful after physical addresses have been + // materialized. For unknown-address roots keep the historical behavior: + // different roots are treated as independent. + if (!hasKnownLocalAbsoluteAddress(a) || !hasKnownLocalAbsoluteAddress(b)) + return false; uint64_t rootBaseA = getRootBaseAddress(a->rootBuffer); uint64_t rootBaseB = getRootBaseAddress(b->rootBuffer); diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index 10cf9da2bd..41adfed168 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -508,8 +508,50 @@ static Value computeSubviewAddress(memref::SubViewOp subview, return builder.create(loc, base, linearOffset); } +static Value computePTOSubViewAddress(pto::SubViewOp subview, + OpBuilder &builder, Location loc) { + Value base = computeExplicitAddress(subview.getSource(), builder, loc); + if (!base) + return Value(); + + auto sourceTy = dyn_cast(subview.getSource().getType()); + if (!sourceTy || sourceTy.getShape().size() < 2) + return Value(); + + unsigned elemBytes = getPTOStorageElemByteSize(sourceTy.getElementType()); + if (elemBytes == 0) + return Value(); + + int64_t rowStride = 0; + int64_t colStride = 0; + if (!getTilePointerStrides(sourceTy.getConfigAttr(), sourceTy.getElementType(), + sourceTy.getShape()[0], sourceTy.getShape()[1], + rowStride, colStride)) + return Value(); + + SmallVector strides = {rowStride, colStride}; + Value linearOffset; + for (auto [offset, stride] : + llvm::zip(subview.getOffsets(), ArrayRef(strides))) { + Value offsetI64 = ensureI64(offset, builder, loc); + if (!offsetI64) + return Value(); + linearOffset = + addI64(linearOffset, mulI64(offsetI64, stride, builder, loc), builder, + loc); + } + + if (!linearOffset) + return base; + linearOffset = mulI64(linearOffset, elemBytes, builder, loc); + return builder.create(loc, base, linearOffset); +} + static Value computeExplicitAddress(Value value, OpBuilder &builder, Location loc) { + if (auto alloc = value.getDefiningOp()) + return ensureI64(alloc.getAddr(), builder, loc); + if (auto bind = value.getDefiningOp()) return computeExplicitAddress(bind.getSource(), builder, loc); @@ -522,6 +564,9 @@ static Value computeExplicitAddress(Value value, OpBuilder &builder, if (auto subview = value.getDefiningOp()) return computeSubviewAddress(subview, builder, loc); + if (auto subview = value.getDefiningOp()) + return computePTOSubViewAddress(subview, builder, loc); + if (auto select = value.getDefiningOp()) { Value trueAddr = computeExplicitAddress(select.getTrueValue(), builder, loc); Value falseAddr = @@ -535,6 +580,11 @@ static Value computeExplicitAddress(Value value, OpBuilder &builder, if (auto cast = value.getDefiningOp()) return computeExplicitAddress(cast.getSource(), builder, loc); + if (auto cast = value.getDefiningOp()) { + if (cast.getNumOperands() == 1) + return computeExplicitAddress(cast.getOperand(0), builder, loc); + } + return Value(); } @@ -962,6 +1012,42 @@ static void copyMaterializedTileAttrs(ArrayRef attrs, to->setAttr(kForceDynamicValidShapeAttrName, attr); } +static LogicalResult +materializePTOSubViewOps(ModuleOp module, OpBuilder &builder, + DenseMap &tileHandles) { + SmallVector subviews; + module.walk([&](pto::SubViewOp op) { subviews.push_back(op); }); + + for (pto::SubViewOp op : subviews) { + auto tileTy = dyn_cast(op.getResult().getType()); + if (!tileTy) { + op.emitError("expected pto.subview result to be a tile_buf"); + return failure(); + } + + builder.setInsertionPoint(op); + Value addr = computeExplicitAddress(op.getResult(), builder, op.getLoc()); + if (!addr) { + op.emitError("cannot materialize pto.subview without a planned source " + "address"); + return failure(); + } + + auto alloc = builder.create( + op.getLoc(), tileTy, addr, + getAllocValidOperand(tileTy, op.getValidRow(), 0, builder, + op.getLoc()), + getAllocValidOperand(tileTy, op.getValidCol(), 1, builder, + op.getLoc())); + alloc->setAttr("pto.view_semantics", builder.getStringAttr("subview")); + tileHandles[op.getResult()] = alloc.getResult(); + op.replaceAllUsesWith(alloc.getResult()); + op.erase(); + } + + return success(); +} + static void updateResultTypesAfterMaterializingOperand(Operation *op, unsigned operandNo, Type tileTy) { @@ -1449,6 +1535,11 @@ struct PTOMaterializeTileHandlesPass return; } + if (failed(materializePTOSubViewOps(module, builder, tileHandles))) { + signalPassFailure(); + return; + } + for (Operation *anchor : anchors) { if (anchor->getNumResults() != 1) continue; diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index 46dc31f90a..4f08600d92 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -94,6 +94,54 @@ static LocalMemSpec getLocalMemSpec(Operation *op, AddressSpace as) { } } +static bool isNameIn(StringRef name, ArrayRef names) { + return llvm::is_contained(names, name); +} + +static bool isIgnoredA5TmpOperandUse(OpOperand &use) { + Operation *owner = use.getOwner(); + unsigned operandNo = use.getOperandNumber(); + StringRef name = owner->getName().getStringRef(); + + if (isNameIn(name, {"pto.trowargmax", "pto.trowargmin", "pto.trowmax", + "pto.trowmin", "pto.trowsum", "pto.trowprod"})) + return operandNo == 1; + + if (name == "pto.txors") + return operandNo == 2; + + if (isNameIn(name, {"pto.tprelu", "pto.txor", "pto.tsels", + "pto.trowexpand", "pto.tcolexpand", + "pto.trowexpandadd", "pto.trowexpanddiv", + "pto.trowexpandexpdif", "pto.trowexpandmax", + "pto.trowexpandmin", "pto.trowexpandmul", + "pto.trowexpandsub", "pto.tcolexpandadd", + "pto.tcolexpanddiv", "pto.tcolexpandexpdif", + "pto.tcolexpandmax", "pto.tcolexpandmin", + "pto.tcolexpandmul", "pto.tcolexpandsub"})) + return operandNo == 2; + + if (name == "pto.tsel") + return operandNo == 3; + + return false; +} + +static bool isA5IgnoredTmpAlloc(pto::AllocTileOp allocTile) { + if (getTargetArch(allocTile.getOperation()) != PTOArch::A5) + return false; + + Value value = allocTile.getResult(); + if (value.use_empty()) + return false; + + for (OpOperand &use : value.getUses()) { + if (!isIgnoredA5TmpOperandUse(use)) + return false; + } + return true; +} + static void collectStableValueOrder(Region ®ion, AsmState &asmState, DenseMap &stableValueKeys, @@ -430,6 +478,8 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { if (allocTileOp.getAddr()) { return WalkResult::advance(); } + if (isA5IgnoredTmpAlloc(allocTileOp)) + return WalkResult::advance(); auto memorySpaceAttr = GetBufferSpaceAttr(allocTileOp.getResult()); if (!isLocalBuffer(memorySpaceAttr)) { allocTileOp.emitError("Alloc tile buffer not at local space"); @@ -449,7 +499,6 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { buffer2MultiNum[allocTileOp.getResult()] = static_cast(n); } UpdateOpBufferInfo(op, op->getResults()); - UpdateOperandGenInfo(curOpInfo, allocTileOp.getResult()); return WalkResult::advance(); } else if (isLocalMemPlan() && dyn_cast(op)) { auto allocOp = cast(op); @@ -480,6 +529,7 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { OpKillHandle(curOpInfo, live, op->getBlock()); } else if (auto tprintOp = dyn_cast(op)) { // TPrintOp only reads from buffer, similar to LoadOp + UpdateOpGenInfo(curOpInfo, llvm::to_vector(op->getOperands())); OpKillHandle(curOpInfo, live, op->getBlock()); } else if (auto tgetvalOp = dyn_cast(op)) { (void)tgetvalOp; @@ -503,7 +553,7 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { } else if (auto ptoDpsOp = dyn_cast(op)) { // PTO ops with destination (tile_buf, partition_view, etc.); no // tensor/memref-only verification. - SmallVector genBuffers = llvm::to_vector(ptoDpsOp.getDpsInits()); + SmallVector genBuffers = llvm::to_vector(op->getOperands()); auto scratchBuffers = getScratchBuffersFromEffects( op, ptoDpsOp.getDpsInits(), stableValueOrder); genBuffers.append(scratchBuffers.begin(), scratchBuffers.end()); @@ -2610,6 +2660,10 @@ void PlanMemoryPass::runOnOperation() { funcOp.walk([&](pto::AllocTileOp op) { if (op.getAddr()) return; + if (op->use_empty()) + return; + if (isA5IgnoredTmpAlloc(op)) + return; op.emitError( "PTOPlanMemory failed to assign an address to pto.alloc_tile"); hasUnplannedAllocTile = true; diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 9b56c7bff1..7c2ed8f9fc 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -102,6 +102,54 @@ static bool isPlannableLocalSpace(std::optional space) { return space && *space != AddressSpace::GM && *space != AddressSpace::Zero; } +static bool isNameIn(StringRef name, ArrayRef names) { + return llvm::is_contained(names, name); +} + +static bool isIgnoredA5TmpOperandUse(OpOperand &use) { + Operation *owner = use.getOwner(); + unsigned operandNo = use.getOperandNumber(); + StringRef name = owner->getName().getStringRef(); + + if (isNameIn(name, {"pto.trowargmax", "pto.trowargmin", "pto.trowmax", + "pto.trowmin", "pto.trowsum", "pto.trowprod"})) + return operandNo == 1; + + if (name == "pto.txors") + return operandNo == 2; + + if (isNameIn(name, {"pto.tprelu", "pto.txor", "pto.tsels", + "pto.trowexpand", "pto.tcolexpand", + "pto.trowexpandadd", "pto.trowexpanddiv", + "pto.trowexpandexpdif", "pto.trowexpandmax", + "pto.trowexpandmin", "pto.trowexpandmul", + "pto.trowexpandsub", "pto.tcolexpandadd", + "pto.tcolexpanddiv", "pto.tcolexpandexpdif", + "pto.tcolexpandmax", "pto.tcolexpandmin", + "pto.tcolexpandmul", "pto.tcolexpandsub"})) + return operandNo == 2; + + if (name == "pto.tsel") + return operandNo == 3; + + return false; +} + +static bool isA5IgnoredTmpAlloc(pto::AllocTileOp allocTile) { + if (getTargetArch(allocTile.getOperation()) != PTOArch::A5) + return false; + + Value value = allocTile.getResult(); + if (value.use_empty()) + return false; + + for (OpOperand &use : value.getUses()) { + if (!isIgnoredA5TmpOperandUse(use)) + return false; + } + return true; +} + static MemSpec getMemSpec(PTOArch arch, AddressSpace space) { switch (space) { case AddressSpace::VEC: @@ -586,6 +634,8 @@ struct PlannerAnalysis { if (auto allocTile = dyn_cast(op)) { if (!allocTile.getAddr()) { + if (isA5IgnoredTmpAlloc(allocTile)) + continue; addRoot(allocTile.getResult(), op); if (failed) return; @@ -628,6 +678,14 @@ struct PlannerAnalysis { setRoots(subview.getResult(), getRoots(subview.getSource())); propagateSplitTpopDerived(subview.getResult(), ValueRange{subview.getSource()}); + } else if (auto subview = dyn_cast(op)) { + setRoots(subview.getResult(), getRoots(subview.getSource())); + propagateSplitTpopDerived(subview.getResult(), + ValueRange{subview.getSource()}); + } else if (auto reshape = dyn_cast(op)) { + setRoots(reshape.getResult(), getRoots(reshape.getSrc())); + propagateSplitTpopDerived(reshape.getResult(), + ValueRange{reshape.getSrc()}); } else if (auto reinterpret = dyn_cast(op)) { setRoots(reinterpret.getResult(), getRoots(reinterpret.getSource())); propagateSplitTpopDerived(reinterpret.getResult(), @@ -1043,6 +1101,10 @@ LogicalResult mlir::pto::runModernPlanMemory(func::FuncOp func, func.walk([&](pto::AllocTileOp op) { if (op.getAddr()) return; + if (op->use_empty()) + return; + if (isA5IgnoredTmpAlloc(op)) + return; op.emitError("PTOPlanMemory failed to assign an address to pto.alloc_tile"); hasUnplannedAllocTile = true; }); diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index fd3fd3edce..cb865879fb 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -1239,6 +1239,39 @@ static LogicalResult bridgeMemRefOperandsToExternalPtrCallees(ModuleOp mod, return success(); } +static LogicalResult bridgeCallOperandsToConvertedCallees(ModuleOp mod, + MLIRContext *ctx) { + SmallVector callOps; + mod.walk([&](func::CallOp callOp) { callOps.push_back(callOp); }); + + for (func::CallOp callOp : callOps) { + auto callee = mod.lookupSymbol(callOp.getCallee()); + if (!callee || callee.isExternal()) + continue; + + FunctionType calleeType = callee.getFunctionType(); + if (calleeType.getNumInputs() != callOp.getNumOperands()) { + callOp.emitError("callee signature does not match call operands"); + return failure(); + } + + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(callOp); + for (auto [index, expectedType] : + llvm::enumerate(calleeType.getInputs())) { + Value operand = callOp.getOperand(index); + if (operand.getType() == expectedType) + continue; + + auto bridge = rewriter.create( + callOp.getLoc(), expectedType, operand); + callOp->setOperand(index, bridge.getResult(0)); + } + } + + return success(); +} + [[maybe_unused]] static LogicalResult lowerTensorViewDimOps(func::FuncOp func, MLIRContext *ctx) { DefaultInlineVector tvDims; func.walk([&](mlir::pto::GetTensorViewDimOp op) { tvDims.push_back(op); }); @@ -1388,6 +1421,12 @@ static LogicalResult lowerSubViewOps(func::FuncOp func, MLIRContext *ctx) { Value src = op->getOperand(0); auto srcMrTy = dyn_cast(src.getType()); if (!srcMrTy) { + // Tile-native alloc_tile roots intentionally survive this pass so that + // PTOPlanMemory can assign their addresses. Defer tile_buf subviews until + // after planning, when PTOMaterializeTileHandles can materialize the + // adjusted address directly. + if (isa(src.getType())) + continue; op.emitError("pto.subview source must be lowered to memref first"); return failure(); } @@ -1555,6 +1594,8 @@ static LogicalResult lowerTileBufViewLikeOps(func::FuncOp func, MLIRContext *ctx op.emitError("treshape result must be tile_buf type"); return failure(); } + if (isa(op->getOperand(0).getType())) + continue; Value lowered = buildTileBufViewLikeValue(op, op->getOperand(0), tbTy, "treshape", ctx); if (!lowered) @@ -1571,6 +1612,8 @@ static LogicalResult lowerTileBufViewLikeOps(func::FuncOp func, MLIRContext *ctx op.emitError("bitcast result must be tile_buf type"); return failure(); } + if (isa(op->getOperand(0).getType())) + continue; Value lowered = buildTileBufViewLikeValue(op, op->getOperand(0), tbTy, "bitcast", ctx); if (!lowered) @@ -4231,6 +4274,11 @@ struct PTOViewToMemrefPass return; } } + + if (failed(bridgeCallOperandsToConvertedCallees(mod, ctx))) { + signalPassFailure(); + return; + } // Debug Output LLVM_DEBUG(llvm::dbgs() << mod.getOperation()); diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index 7c8de83043..6356cc7b83 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -155,6 +155,10 @@ std::optional GetBufferSpaceAttr(Value operand) { std::optional> getOperationAliasInfo(Operation *op) { if (auto subViewOp = dyn_cast(op)) { return std::make_pair(subViewOp.getResult(), subViewOp.getViewSource()); + } else if (auto subViewOp = dyn_cast(op)) { + return std::make_pair(subViewOp.getResult(), subViewOp.getSource()); + } else if (auto reshapeOp = dyn_cast(op)) { + return std::make_pair(reshapeOp.getResult(), reshapeOp.getSrc()); } else if (auto bindTileOp = dyn_cast(op)) { return std::make_pair(bindTileOp.getResult(), bindTileOp.getSource()); } else if (auto slotMarkerOp = dyn_cast(op)) { @@ -179,6 +183,9 @@ std::optional> getOperationAliasInfo(Operation *op) { return std::make_pair(reshapeOp.getResult(), reshapeOp.getViewSource()); } else if (auto castOp = dyn_cast(op)) { return std::make_pair(castOp.getResult(), castOp.getViewSource()); + } else if (auto castOp = dyn_cast(op)) { + if (castOp.getNumOperands() == 1 && castOp.getNumResults() == 1) + return std::make_pair(castOp.getResult(0), castOp.getOperand(0)); } else if (auto extractStridedMetadataOp = dyn_cast(op)) { return std::make_pair(extractStridedMetadataOp.getBaseBuffer(), diff --git a/test/lit/pto/a5_unused_tmp_vec_overflow.pto b/test/lit/pto/a5_unused_tmp_vec_overflow.pto index a07f1b9588..500112c2ec 100644 --- a/test/lit/pto/a5_unused_tmp_vec_overflow.pto +++ b/test/lit/pto/a5_unused_tmp_vec_overflow.pto @@ -73,12 +73,12 @@ module { } func.func @TROWEXPANDMAX_TMP_OVERFLOW() { - %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.trowexpandmax ins(%src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) - outs(%dst : !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.trowexpandmax ins(%src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } } diff --git a/test/lit/pto/alloc_tile_low_precision_valid.pto b/test/lit/pto/alloc_tile_low_precision_valid.pto index ffa8240309..f385fb0290 100644 --- a/test/lit/pto/alloc_tile_low_precision_valid.pto +++ b/test/lit/pto/alloc_tile_low_precision_valid.pto @@ -18,4 +18,6 @@ module { } // CHECK: func.func @alloc_tile_low_precision_valid() -// CHECK-NOT: pto.alloc_tile +// CHECK: pto.alloc_tile : !pto.tile_buf +// CHECK: pto.alloc_tile : !pto.tile_buf +// CHECK: pto.alloc_tile : !pto.tile_buf diff --git a/test/lit/pto/compact_left_blayout_parser_a3.pto b/test/lit/pto/compact_left_blayout_parser_a3.pto index b1c7314036..271cdcafd6 100644 --- a/test/lit/pto/compact_left_blayout_parser_a3.pto +++ b/test/lit/pto/compact_left_blayout_parser_a3.pto @@ -8,4 +8,4 @@ module attributes {"pto.target_arch" = "a3"} { } // CHECK-LABEL: func.func @compact_left_blayout_parser_a3() { -// CHECK: #pto.tile_buf_config, slayout=#pto.slayout, s_fractal_size=512, pad=#pto.pad_value, compact=#pto.compact_mode> +// CHECK: pto.alloc_tile : !pto.tile_buf diff --git a/test/lit/pto/compact_left_blayout_parser_a5.pto b/test/lit/pto/compact_left_blayout_parser_a5.pto index 4dad24114e..23a5c1b6e5 100644 --- a/test/lit/pto/compact_left_blayout_parser_a5.pto +++ b/test/lit/pto/compact_left_blayout_parser_a5.pto @@ -8,4 +8,4 @@ module attributes {"pto.target_arch" = "a5"} { } // CHECK-LABEL: func.func @compact_left_blayout_parser_a5() { -// CHECK: #pto.tile_buf_config, slayout=#pto.slayout, s_fractal_size=512, pad=#pto.pad_value, compact=#pto.compact_mode> +// CHECK: pto.alloc_tile : !pto.tile_buf diff --git a/test/lit/pto/emitc_tile_data_sink_after_tassign.pto b/test/lit/pto/emitc_tile_data_sink_after_tassign.pto index 70086a0c89..e8014da937 100644 --- a/test/lit/pto/emitc_tile_data_sink_after_tassign.pto +++ b/test/lit/pto/emitc_tile_data_sink_after_tassign.pto @@ -29,15 +29,6 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind [[ADDR_BASE:v[0-9]+]]; // CHECK: // pto: %0 // CHECK-NEXT: Tile [[ADDR_TILE:v[0-9]+]] = [[ADDR_BASE]]; -// CHECK-NEXT: TASSIGN([[ADDR_TILE]], v1); // CHECK: // pto: %1 -// CHECK-NEXT: Tile [[SINK_BASE:v[0-9]+]]; -// CHECK: // pto: %1 -// CHECK-NEXT: Tile [[SINK_TILE:v[0-9]+]] = [[SINK_BASE]]; -// CHECK: // pto: %1 -// CHECK-NEXT: __ubuf__ float* [[ADDR_PTR:v[0-9]+]] = [[ADDR_TILE]].data(); -// CHECK: // pto: %1 -// CHECK-NEXT: uint64_t [[ADDR_BITS:v[0-9]+]] = reinterpret_cast([[ADDR_PTR]]); -// CHECK-NEXT: TASSIGN([[SINK_TILE]], [[ADDR_BITS]]); -// CHECK: __ubuf__ float* [[SINK_PTR:v[0-9]+]] = [[SINK_TILE]].data(); +// CHECK-NEXT: __ubuf__ float* [[SINK_PTR:v[0-9]+]] = [[ADDR_TILE]].data(); // CHECK-NEXT: sink_ptr([[SINK_PTR]]); diff --git a/test/lit/pto/fixpipe_frontend_emitc_vector_remat_a5.pto b/test/lit/pto/fixpipe_frontend_emitc_vector_remat_a5.pto index 8af360bcc1..9feaa0fae8 100644 --- a/test/lit/pto/fixpipe_frontend_emitc_vector_remat_a5.pto +++ b/test/lit/pto/fixpipe_frontend_emitc_vector_remat_a5.pto @@ -120,9 +120,9 @@ module { // A5-NEXT: TPUSH<{{.*}}, Pipe0FixpipeConfig>( // PLAN-LABEL: func.func @cube_kernel() -// PLAN: pto.set_quant_vector(%{{.*}} : memref<1x32xf16{{.*}}#pto.address_space>) {id = 0} +// PLAN: pto.set_quant_vector(%{{.*}} : !pto.tile_buf) {id = 0} // PLAN-NEXT: pto.tpush( -// PLAN: pto.set_quant_vector(%{{.*}} : memref<1x32xf16{{.*}}#pto.address_space>) {id = 1} +// PLAN: pto.set_quant_vector(%{{.*}} : !pto.tile_buf) {id = 1} // PLAN-NEXT: pto.tpush( -// PLAN: pto.set_quant_vector(%{{.*}} : memref<1x32xf16{{.*}}#pto.address_space>) {id = 0} +// PLAN: pto.set_quant_vector(%{{.*}} : !pto.tile_buf) {id = 0} // PLAN-NEXT: pto.tpush( diff --git a/test/lit/pto/fixpipe_frontend_quant_state_preserve_ir_a5.pto b/test/lit/pto/fixpipe_frontend_quant_state_preserve_ir_a5.pto index 9ce0ff5db4..13e9c893ff 100644 --- a/test/lit/pto/fixpipe_frontend_quant_state_preserve_ir_a5.pto +++ b/test/lit/pto/fixpipe_frontend_quant_state_preserve_ir_a5.pto @@ -130,7 +130,7 @@ module { // LATE-NEXT: pto.tpush( // LATE: pto.set_quant_scalar(%{{.*}} : f32) {id = 0} // LATE-NEXT: pto.tpush( -// LATE: pto.set_quant_vector(%{{.*}} : memref<1x32xf16{{.*}}#pto.address_space>) {id = 1} +// LATE: pto.set_quant_vector(%{{.*}} : !pto.tile_buf) {id = 1} // LATE-NEXT: pto.tpush( -// LATE: pto.set_quant_vector(%{{.*}} : memref<1x32xf16{{.*}}#pto.address_space>) {id = 1} +// LATE: pto.set_quant_vector(%{{.*}} : !pto.tile_buf) {id = 1} // LATE-NEXT: pto.tpush( diff --git a/test/lit/pto/graph_sync_solver_basic.pto b/test/lit/pto/graph_sync_solver_basic.pto index 95ebf777e9..dd34184526 100644 --- a/test/lit/pto/graph_sync_solver_basic.pto +++ b/test/lit/pto/graph_sync_solver_basic.pto @@ -148,7 +148,7 @@ module { // CHECK-LABEL: func.func @event_id_fallback // CHECK-COUNT-9: {{pto\.set_flag\[, , \]}} - // CHECK: pto.tmuls ins(%4, %cst : + // CHECK: pto.tmuls ins(%{{[0-9]+}}, %cst : func.func @event_id_fallback(%arg0: !pto.ptr, %arg2: index, %arg3: index, %arg4: index, diff --git a/test/lit/pto/issue531_tprelu_vec_overflow_a5.pto b/test/lit/pto/issue531_tprelu_vec_overflow_a5.pto index c8b9fdcbf3..9c76c24bf8 100644 --- a/test/lit/pto/issue531_tprelu_vec_overflow_a5.pto +++ b/test/lit/pto/issue531_tprelu_vec_overflow_a5.pto @@ -16,13 +16,13 @@ module { func.func @TPRELU_OVERFLOW() { - %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 + %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.tprelu ins(%src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) - outs(%dst : !pto.tile_buf) + pto.tprelu ins(%src0, %src1, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } } diff --git a/test/lit/pto/issue686_a5_tmov_treshape_square_dynamic_valid_shape.pto b/test/lit/pto/issue686_a5_tmov_treshape_square_dynamic_valid_shape.pto index 73626db68f..8b41fc5a0e 100644 --- a/test/lit/pto/issue686_a5_tmov_treshape_square_dynamic_valid_shape.pto +++ b/test/lit/pto/issue686_a5_tmov_treshape_square_dynamic_valid_shape.pto @@ -29,10 +29,10 @@ module attributes {pto.target_arch = "a5"} { // CHECK-DAG: [[PART:%[A-Za-z0-9_]+]] = arith.constant 8 : index // CHECK-DAG: [[FULL:%[A-Za-z0-9_]+]] = arith.constant 16 : index -// CHECK: [[SRC_BASE:%[0-9]+]] = pto.bind_tile {{.*}}, [[PART]], [[FULL]] {{.*}}blayout=#pto.blayout -// CHECK: [[DST_BASE:%[0-9]+]] = pto.bind_tile {{.*}}, [[PART]], [[FULL]] {{.*}}blayout=#pto.blayout -// CHECK: [[SRC:%[0-9]+]] = pto.bind_tile [[SRC_BASE]], [[PART]], [[FULL]] {{.*}}blayout=#pto.blayout{{.*}}pto.view_semantics = "treshape" -// CHECK: [[DST:%[0-9]+]] = pto.bind_tile [[DST_BASE]], [[PART]], [[FULL]] {{.*}}blayout=#pto.blayout{{.*}}pto.view_semantics = "treshape" +// CHECK: [[SRC_BASE:%[0-9]+]] = pto.alloc_tile addr = {{.*}} valid_row = [[PART]] valid_col = [[FULL]] : !pto.tile_buf +// CHECK: [[DST_BASE:%[0-9]+]] = pto.alloc_tile addr = {{.*}} valid_row = [[PART]] valid_col = [[FULL]] : !pto.tile_buf +// CHECK: [[SRC:%[0-9]+]] = pto.treshape [[SRC_BASE]] : !pto.tile_buf -> !pto.tile_buf +// CHECK: [[DST:%[0-9]+]] = pto.treshape [[DST_BASE]] : !pto.tile_buf -> !pto.tile_buf // CHECK: [[SRC_ROW:%[A-Za-z0-9_]+]], [[SRC_COL:%[A-Za-z0-9_]+]] = pto.get_validshape [[SRC_BASE]] // CHECK-NEXT: pto.set_validshape [[SRC]], [[SRC_COL]], [[SRC_ROW]] // CHECK-NEXT: [[DST_ROW:%[A-Za-z0-9_]+]], [[DST_COL:%[A-Za-z0-9_]+]] = pto.get_validshape [[DST_BASE]] diff --git a/test/lit/pto/issue708_zero_valid_subview.pto b/test/lit/pto/issue708_zero_valid_subview.pto index 3aa8615180..1c70846404 100644 --- a/test/lit/pto/issue708_zero_valid_subview.pto +++ b/test/lit/pto/issue708_zero_valid_subview.pto @@ -46,11 +46,9 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { } // CHECK-LABEL: func.func @issue708_zero_valid_subview_row -// CHECK: %[[ROW:[0-9A-Za-z_]+]] = pto.bind_tile {{.*}}, %c0, %c128 -// CHECK-SAME: pto.view_semantics = "subview" +// CHECK: %[[ROW:[0-9A-Za-z_]+]] = pto.subview {{.*}} valid [%c0, %c128] // CHECK: pto.tstore ins(%[[ROW]] // CHECK-LABEL: func.func @issue708_zero_valid_subview_col -// CHECK: %[[COL:[0-9A-Za-z_]+]] = pto.bind_tile {{.*}}, %c5, %c0 -// CHECK-SAME: pto.view_semantics = "subview" +// CHECK: %[[COL:[0-9A-Za-z_]+]] = pto.subview {{.*}} valid [%c5, %c0] // CHECK: pto.tstore ins(%[[COL]] diff --git a/test/lit/pto/issue708_zero_valid_subview_inferred.pto b/test/lit/pto/issue708_zero_valid_subview_inferred.pto index 2abc8bae8f..38f1893b20 100644 --- a/test/lit/pto/issue708_zero_valid_subview_inferred.pto +++ b/test/lit/pto/issue708_zero_valid_subview_inferred.pto @@ -95,21 +95,17 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { // from the result type, i.e. 0 for an empty axis -- never the subview size. // CHECK-LABEL: func.func @issue708_inferred_empty_store -// CHECK: %[[BOTH:[0-9A-Za-z_]+]] = pto.bind_tile {{.*}}, %c0, %c0 -// CHECK-SAME: pto.view_semantics = "subview" +// CHECK: %[[BOTH:[0-9A-Za-z_]+]] = pto.subview {{.*}} -> !pto.tile_buf // CHECK: pto.tstore ins(%[[BOTH]] // CHECK-LABEL: func.func @issue708_inferred_empty_cvt -// CHECK: %[[CVTV:[0-9A-Za-z_]+]] = pto.bind_tile {{.*}}, %c0, %c0 -// CHECK-SAME: pto.view_semantics = "subview" +// CHECK: %[[CVTV:[0-9A-Za-z_]+]] = pto.subview {{.*}} -> !pto.tile_buf // CHECK: pto.tcvt ins(%[[CVTV]] // CHECK-LABEL: func.func @issue708_inferred_row_empty -// CHECK: %[[ROW:[0-9A-Za-z_]+]] = pto.bind_tile {{.*}}, %c0, %c128 -// CHECK-SAME: pto.view_semantics = "subview" +// CHECK: %[[ROW:[0-9A-Za-z_]+]] = pto.subview {{.*}} -> !pto.tile_buf // CHECK: pto.tstore ins(%[[ROW]] // CHECK-LABEL: func.func @issue708_inferred_partial_valid -// CHECK: %[[PARTIAL:[0-9A-Za-z_]+]] = pto.bind_tile {{.*}}, %c3, %c128 -// CHECK-SAME: pto.view_semantics = "subview" +// CHECK: %[[PARTIAL:[0-9A-Za-z_]+]] = pto.subview {{.*}} -> !pto.tile_buf // CHECK: pto.tstore ins(%[[PARTIAL]] diff --git a/test/lit/pto/mgather_mscatter_a5_attrs_emitc.pto b/test/lit/pto/mgather_mscatter_a5_attrs_emitc.pto index 7c50049378..e061e82159 100644 --- a/test/lit/pto/mgather_mscatter_a5_attrs_emitc.pto +++ b/test/lit/pto/mgather_mscatter_a5_attrs_emitc.pto @@ -88,12 +88,8 @@ module attributes {pto.target_arch = "a5"} { // CHECK-LABEL: AICORE void mgather_emitc // CHECK: TLOAD({{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); -// CHECK-DAG: __ubuf__ int32_t* [[MGATHER_IDX:[_A-Za-z][_A-Za-z0-9]*]] = {{.*}}; -// CHECK-DAG: __ubuf__ float* [[MGATHER_DST:[_A-Za-z][_A-Za-z0-9]*]] = {{.*}}; -// CHECK: MGATHER([[MGATHER_DST]], {{[_A-Za-z][_A-Za-z0-9]*}}, [[MGATHER_IDX]]); +// CHECK: MGATHER({{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); // CHECK-LABEL: AICORE void mscatter_emitc // CHECK: TLOAD({{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); // CHECK: TLOAD({{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); -// CHECK-DAG: __ubuf__ float* [[MSCATTER_SRC:[_A-Za-z][_A-Za-z0-9]*]] = {{.*}}; -// CHECK-DAG: __ubuf__ int32_t* [[MSCATTER_IDX:[_A-Za-z][_A-Za-z0-9]*]] = {{.*}}; -// CHECK: MSCATTER({{[_A-Za-z][_A-Za-z0-9]*}}, [[MSCATTER_SRC]], [[MSCATTER_IDX]]); +// CHECK: MSCATTER({{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); diff --git a/test/lit/pto/mgather_mscatter_base_a3_emitc.pto b/test/lit/pto/mgather_mscatter_base_a3_emitc.pto index ac413fa848..e407a8e9db 100644 --- a/test/lit/pto/mgather_mscatter_base_a3_emitc.pto +++ b/test/lit/pto/mgather_mscatter_base_a3_emitc.pto @@ -63,11 +63,7 @@ module attributes {pto.target_arch = "a3"} { } // CHECK-LABEL: AICORE void mgather_emitc_base( -// CHECK: __ubuf__ int32_t* [[MGATHER_IDX:[_A-Za-z][_A-Za-z0-9]*]] = {{[_A-Za-z][_A-Za-z0-9]*}}.data(); -// CHECK: __ubuf__ float* [[MGATHER_DST:[_A-Za-z][_A-Za-z0-9]*]] = {{[_A-Za-z][_A-Za-z0-9]*}}.data(); -// CHECK: MGATHER([[MGATHER_DST]], {{[_A-Za-z][_A-Za-z0-9]*}}, [[MGATHER_IDX]]); +// CHECK: MGATHER({{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); // CHECK-LABEL: AICORE void mscatter_emitc_base( -// CHECK: __ubuf__ float* [[MSCATTER_SRC:[_A-Za-z][_A-Za-z0-9]*]] = {{[_A-Za-z][_A-Za-z0-9]*}}.data(); -// CHECK: __ubuf__ int32_t* [[MSCATTER_IDX:[_A-Za-z][_A-Za-z0-9]*]] = {{[_A-Za-z][_A-Za-z0-9]*}}.data(); -// CHECK: MSCATTER({{[_A-Za-z][_A-Za-z0-9]*}}, [[MSCATTER_SRC]], [[MSCATTER_IDX]]); +// CHECK: MSCATTER({{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); diff --git a/test/lit/pto/ptodsl_subkernel_call_autosync.pto b/test/lit/pto/ptodsl_subkernel_call_autosync.pto index 798c67b3df..5d5532fb28 100644 --- a/test/lit/pto/ptodsl_subkernel_call_autosync.pto +++ b/test/lit/pto/ptodsl_subkernel_call_autosync.pto @@ -39,9 +39,8 @@ module { // CHECK-LABEL: AICORE void subkernel_call_autosync( // CHECK: TLOAD( -// CHECK: set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); -// CHECK: wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); +// CHECK: set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); // CHECK: TNEG( -// CHECK: set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); -// CHECK: wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); +// CHECK: pipe_barrier(PIPE_V); +// CHECK: wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); // CHECK: TSTORE( diff --git a/test/lit/pto/subview_bind_tile_preserve_stride.pto b/test/lit/pto/subview_bind_tile_preserve_stride.pto index c611c540e2..5b160d2bb3 100644 --- a/test/lit/pto/subview_bind_tile_preserve_stride.pto +++ b/test/lit/pto/subview_bind_tile_preserve_stride.pto @@ -39,13 +39,12 @@ module { } } -// CHECK-DAG: Tile -// CHECK-DAG: Tile -// CHECK-DAG: Tile -// CHECK-DAG: Tile +// CHECK-DAG: Tile +// CHECK-DAG: Tile +// CHECK-DAG: Tile +// CHECK-DAG: Tile // CHECK-DAG: TSTORE( // CHECK-DAG: TSTORE( // CHECK-DAG: TSTORE( // CHECK-DAG: TSTORE( -// CHECK-NOT: Tile // CHECK-NOT: copy_ubuf_to_gm_align_b32( diff --git a/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto b/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto index 1ecda2915c..cd72ab78ae 100644 --- a/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto +++ b/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto @@ -34,8 +34,7 @@ module { // IR: %[[ELEM_BYTES:.*]] = arith.constant 4 : i64 // IR: %[[BYTE_OFFSET:.*]] = arith.muli %[[ELEM_OFFSET]], %[[ELEM_BYTES]] : i64 // IR: %[[ADDR:.*]] = arith.addi %{{.*}}, %[[BYTE_OFFSET]] : i64 -// IR: pto.alloc_tile addr = %[[ADDR]] : !pto.tile_buf -// IR-NOT: pto.view_semantics +// IR: pto.alloc_tile addr = %[[ADDR]] {pto.view_semantics = "subview"} : !pto.tile_buf // IR-NOT: pto.materialize_tile // EMITC: Tile diff --git a/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto b/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto index ef374d41d1..9908f5361f 100644 --- a/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto +++ b/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto @@ -21,5 +21,4 @@ module { } } -// CHECK: Tile -// CHECK-NOT: Tile +// CHECK: Tile diff --git a/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto b/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto index b83bb4bdfd..9911993b05 100644 --- a/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto +++ b/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto @@ -21,6 +21,5 @@ module { } // ColMajor + RowPlusOne: major stride should be 17 (not 16). -// CHECK: pto.pointer_cast(%c0_i64) {{.*}} : memref<16x16xf32, strided<[1, 17]>, #pto.address_space> -// CHECK: compact=#pto.compact_mode -// CHECK: memref.subview {{.*}} to memref<8x8xf32, strided<[1, 17], offset: ?>, #pto.address_space> +// CHECK: pto.alloc_tile addr = %c0_i64 : !pto.tile_buf +// CHECK: pto.subview {{.*}} sizes [8, 8] : !pto.tile_buf -> !pto.tile_buf diff --git a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto index 9826fc294f..f2c1979946 100644 --- a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto +++ b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto @@ -24,5 +24,4 @@ module { } // CHECK: func.func @subview_dynamic_offset_static_valid_regression -// CHECK: memref.subview %{{.*}}[0, %{{.*}}] [1, 64] [1, 1] -// CHECK: pto.bind_tile %{{.*}}, %c1, %c64 +// CHECK: pto.subview %{{.*}}[%c0, %{{.*}}] sizes [1, 64] : !pto.tile_buf -> !pto.tile_buf diff --git a/test/lit/pto/subview_explicit_result_type.pto b/test/lit/pto/subview_explicit_result_type.pto index e6f74e5705..b76ce90a41 100644 --- a/test/lit/pto/subview_explicit_result_type.pto +++ b/test/lit/pto/subview_explicit_result_type.pto @@ -21,4 +21,4 @@ module { } } -// CHECK: Tile +// CHECK: Tile diff --git a/test/lit/pto/subview_row_plus_one_stride_offset.pto b/test/lit/pto/subview_row_plus_one_stride_offset.pto index c2587b3a6e..02e9b11ded 100644 --- a/test/lit/pto/subview_row_plus_one_stride_offset.pto +++ b/test/lit/pto/subview_row_plus_one_stride_offset.pto @@ -21,6 +21,5 @@ module { } // RowPlusOne: major stride should be 17 (not 16). -// CHECK: pto.pointer_cast(%c0_i64) {{.*}} : memref<16x16xf32, strided<[17, 1]>, #pto.address_space> -// CHECK: compact=#pto.compact_mode -// CHECK: memref.subview {{.*}} to memref<8x8xf32, strided<[17, 1], offset: ?>, #pto.address_space> +// CHECK: pto.alloc_tile addr = %c0_i64 : !pto.tile_buf +// CHECK: pto.subview {{.*}} sizes [8, 8] : !pto.tile_buf -> !pto.tile_buf diff --git a/test/lit/pto/subview_validshape_guard.pto b/test/lit/pto/subview_validshape_guard.pto index fafcdea660..52d0caa9e7 100644 --- a/test/lit/pto/subview_validshape_guard.pto +++ b/test/lit/pto/subview_validshape_guard.pto @@ -66,14 +66,14 @@ module { // Default non-compact subview keeps parent physical shape (4x16) but valid // defaults to subview shape (2x8), not clipped from parent valid shape. -// CHECK: Tile -// CHECK-NOT: Tile +// CHECK: Tile +// CHECK-NOT: Tile // Default row-major full-col subview can use the compact child physical shape (2x8). // CHECK: Tile // Explicit valid override takes effect. -// CHECK: Tile +// CHECK: Tile // Dynamic explicit valid keeps dynamic valid dims at lowering. -// CHECK: Tile +// CHECK: Tile diff --git a/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto b/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto index 973ef5ee60..df24d9cb57 100644 --- a/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto +++ b/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto @@ -15,5 +15,5 @@ module { } } -// CHECK: Tile +// CHECK: Tile // CHECK-NOT: error: diff --git a/test/lit/pto/tile_compact_mode_emitc.pto b/test/lit/pto/tile_compact_mode_emitc.pto index f03d81b79f..7c827be7e4 100644 --- a/test/lit/pto/tile_compact_mode_emitc.pto +++ b/test/lit/pto/tile_compact_mode_emitc.pto @@ -25,9 +25,9 @@ module { } } -// A3-DAG: memref.alloc() : memref<1x16xf16, strided<[16, 1]>, #pto.address_space> -// A3-DAG: memref.alloc() : memref<1x16xf16, strided<[16, 1]>, #pto.address_space> -// A3-DAG: memref.alloc() : memref<1x16xf16, strided<[17, 1]>, #pto.address_space> +// A3-DAG: pto.alloc_tile : !pto.tile_buf +// A3-DAG: pto.alloc_tile : !pto.tile_buf +// A3-DAG: pto.alloc_tile : !pto.tile_buf // A3: Tile [[DEFAULT:[_A-Za-z][_A-Za-z0-9]*]]; // A3: Tile [[COMPACT:[_A-Za-z][_A-Za-z0-9]*]]; // A3: Tile [[ROWP1:[_A-Za-z][_A-Za-z0-9]*]]; diff --git a/test/lit/pto/tpow_emitc.pto b/test/lit/pto/tpow_emitc.pto index 186ce1c57c..d252177805 100644 --- a/test/lit/pto/tpow_emitc.pto +++ b/test/lit/pto/tpow_emitc.pto @@ -19,6 +19,6 @@ module { } } -// A3: pto.tpow ins([[BASE:%[0-9]+]], [[EXP:%[0-9]+]], [[TMP:%[0-9]+]] : memref<1x16xf32 -// A3: outs([[DST:%[0-9]+]] : memref<1x16xf32 +// A3: pto.tpow ins([[BASE:%[0-9]+]], [[EXP:%[0-9]+]], [[TMP:%[0-9]+]] : !pto.tile_buf, diff --git a/test/lit/pto/trem_emitc.pto b/test/lit/pto/trem_emitc.pto index a5aa39f7ee..36736b0102 100644 --- a/test/lit/pto/trem_emitc.pto +++ b/test/lit/pto/trem_emitc.pto @@ -11,8 +11,8 @@ module { } } -// A3: pto.trem ins([[SRC0:%[0-9]+]], [[SRC1:%[0-9]+]], [[TMP:%[0-9]+]] : memref<1x16xf32 -// A3-SAME: memref<1x16xf32 -// A3-SAME: memref<2x16xf32 -// A3: outs([[DST:%[0-9]+]] : memref<1x16xf32 +// A3: pto.trem ins([[SRC0:%[0-9]+]], [[SRC1:%[0-9]+]], [[TMP:%[0-9]+]] : !pto.tile_buf{{.*}}pto.view_semantics = "treshape" +// CHECK: pto.alloc_tile addr = {{.*}} valid_row = [[PART]] valid_col = [[FULL]] : !pto.tile_buf +// CHECK: pto.treshape {{.*}} : !pto.tile_buf -> !pto.tile_buf diff --git a/test/lit/pto/trsqrt_emitc.pto b/test/lit/pto/trsqrt_emitc.pto index b9b1b55bf3..7fc664730b 100644 --- a/test/lit/pto/trsqrt_emitc.pto +++ b/test/lit/pto/trsqrt_emitc.pto @@ -13,9 +13,9 @@ module { } } -// A3: pto.trsqrt ins([[SRC0:%[0-9]+]] : memref<1x16xf16 -// A3: outs([[DST0:%[0-9]+]] : memref<1x16xf16 -// A3: pto.trsqrt ins([[SRC1:%[0-9]+]], [[TMP:%[0-9]+]] : memref<1x16xf16 -// A3: outs([[DST1:%[0-9]+]] : memref<1x16xf16 +// A3: pto.trsqrt ins([[SRC0:%[0-9]+]] : !pto.tile_buf // LEVEL2-PLAIN: pto.tadd ins( // LEVEL2-PLAIN-NEXT: pto.tmul ins( // LEVEL2-PLAIN: pto.tstore ins( @@ -83,7 +83,7 @@ module { // LEVEL2-SYNC-NEXT: pto.trowexpandmul // LEVEL2-SYNC-NEXT: pto.tadd // LEVEL2-SYNC: pto.yield( -// LEVEL2-SYNC: } {pto.fusion.group_id = 0 : i64} : memref<32x32xf32 +// LEVEL2-SYNC: } {pto.fusion.group_id = 0 : i64} : !pto.tile_buf // LEVEL2-SYNC: pto.tadd ins( // LEVEL2-SYNC-NEXT: pto.tmul ins( // LEVEL2-SYNC: pto.tstore ins( diff --git a/test/lit/vpto/normalize_uncovered_tile_sections_vector.pto b/test/lit/vpto/normalize_uncovered_tile_sections_vector.pto index 665d08b682..728c2a2e5c 100644 --- a/test/lit/vpto/normalize_uncovered_tile_sections_vector.pto +++ b/test/lit/vpto/normalize_uncovered_tile_sections_vector.pto @@ -39,7 +39,7 @@ module attributes {pto.target_arch = "a5"} { // CHECK: module attributes // CHECK-SAME: pto.backend = "vpto", pto.kernel_kind = #pto.kernel_kind, pto.target_arch = "a5" // CHECK: func.func @auto_vector_segment -// CHECK: pto.bind_tile +// CHECK: pto.alloc_tile addr = // CHECK: pto.tload // CHECK: pto.tadds // CHECK: pto.tstore From 73d966d26f19f3c2d68ad974cd0337aba5e869ec Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 14 Jul 2026 14:12:08 +0800 Subject: [PATCH 05/91] Avoid treating DPS outputs as ignored A5 tmp --- lib/PTO/Transforms/PTOPlanMemory.cpp | 8 ++++++++ lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index 4f08600d92..81656b7859 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -103,6 +103,14 @@ static bool isIgnoredA5TmpOperandUse(OpOperand &use) { unsigned operandNo = use.getOperandNumber(); StringRef name = owner->getName().getStringRef(); + if (auto dpsOp = dyn_cast(owner)) { + if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) + return false; + } else if (auto dpsOp = dyn_cast(owner)) { + if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) + return false; + } + if (isNameIn(name, {"pto.trowargmax", "pto.trowargmin", "pto.trowmax", "pto.trowmin", "pto.trowsum", "pto.trowprod"})) return operandNo == 1; diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 7c2ed8f9fc..1e9b30402d 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -111,6 +111,14 @@ static bool isIgnoredA5TmpOperandUse(OpOperand &use) { unsigned operandNo = use.getOperandNumber(); StringRef name = owner->getName().getStringRef(); + if (auto dpsOp = dyn_cast(owner)) { + if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) + return false; + } else if (auto dpsOp = dyn_cast(owner)) { + if (llvm::is_contained(dpsOp.getDpsInits(), use.get())) + return false; + } + if (isNameIn(name, {"pto.trowargmax", "pto.trowargmin", "pto.trowmax", "pto.trowmin", "pto.trowsum", "pto.trowprod"})) return operandNo == 1; From e0a7f1a6e74092326b007041447b30b2c6ceb5a7 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 14 Jul 2026 16:46:19 +0800 Subject: [PATCH 06/91] Fix tile subview materialization for strided windows --- .../Transforms/PTOMaterializeTileHandles.cpp | 26 +++++++++-- lib/PTO/Transforms/PTOToEmitC.cpp | 9 +++- .../pto/subview_bind_tile_preserve_stride.pto | 8 ++-- ..._column_slice_preserve_parent_shape_a5.pto | 2 +- ...w_col_major_compact_keeps_normal_shape.pto | 4 +- ...w_col_major_noncompact_preserve_stride.pto | 2 +- .../subview_compact_keeps_normal_shape.pto | 4 +- test/lit/pto/subview_explicit_result_type.pto | 2 +- .../lit/pto/subview_tmatmul_acc_slices_a5.pto | 2 +- test/lit/pto/subview_validshape_guard.pto | 10 ++--- ...shape_partial_parent_requires_explicit.pto | 2 +- ...ile_assemble_level3_subview_tmov_emitc.pto | 43 +++++++++++++++++++ .../pto/tinsert_level3_tile_native_sync.pto | 38 ++++++++++++++++ 13 files changed, 130 insertions(+), 22 deletions(-) create mode 100644 test/lit/pto/tile_assemble_level3_subview_tmov_emitc.pto create mode 100644 test/lit/pto/tinsert_level3_tile_native_sync.pto diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index 41adfed168..f3ad1810ef 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -998,6 +998,25 @@ static Value getAllocValidOperand(TileBufType tileTy, Value operand, return Value(); } +static TileBufType buildPTOSubViewPhysicalTileType(pto::SubViewOp op, + TileBufType resultTy) { + auto sourceTy = dyn_cast(op.getSource().getType()); + if (!sourceTy || sourceTy.getRank() != 2 || resultTy.getRank() != 2) + return resultTy; + + SmallVector physicalShape(sourceTy.getShape().begin(), + sourceTy.getShape().end()); + SmallVector validShape(resultTy.getValidShape().begin(), + resultTy.getValidShape().end()); + if (validShape.size() != 2) + validShape.assign(resultTy.getShape().begin(), resultTy.getShape().end()); + + return TileBufType::get(op.getContext(), physicalShape, + resultTy.getElementType(), + resultTy.getMemorySpace(), validShape, + resultTy.getConfigAttr()); +} + static Attribute getAttr(ArrayRef attrs, StringRef name) { for (NamedAttribute attr : attrs) { if (attr.getName().getValue() == name) @@ -1033,11 +1052,12 @@ materializePTOSubViewOps(ModuleOp module, OpBuilder &builder, return failure(); } + TileBufType physicalTileTy = buildPTOSubViewPhysicalTileType(op, tileTy); auto alloc = builder.create( - op.getLoc(), tileTy, addr, - getAllocValidOperand(tileTy, op.getValidRow(), 0, builder, + op.getLoc(), physicalTileTy, addr, + getAllocValidOperand(physicalTileTy, op.getValidRow(), 0, builder, op.getLoc()), - getAllocValidOperand(tileTy, op.getValidCol(), 1, builder, + getAllocValidOperand(physicalTileTy, op.getValidCol(), 1, builder, op.getLoc())); alloc->setAttr("pto.view_semantics", builder.getStringAttr("subview")); tileHandles[op.getResult()] = alloc.getResult(); diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index ab57d08e24..1d8ad818ae 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -1350,7 +1350,14 @@ static LogicalResult insertFixpipeConfigAliases(ModuleOp mop) { if (aliases.empty()) continue; - OpBuilder builder(&funcOp.front(), funcOp.front().begin()); + if (funcOp.empty()) { + funcOp.emitError("cannot insert fixpipe config aliases into an external " + "function"); + return failure(); + } + + OpBuilder builder(funcOp.getContext()); + builder.setInsertionPointToStart(&funcOp.front()); for (const auto &[pipeId, configTok] : aliases) { std::string line = "using " + buildFixpipeConfigAliasName(pipeId) + " = " + configTok + ";"; diff --git a/test/lit/pto/subview_bind_tile_preserve_stride.pto b/test/lit/pto/subview_bind_tile_preserve_stride.pto index 5b160d2bb3..9e2ca7168f 100644 --- a/test/lit/pto/subview_bind_tile_preserve_stride.pto +++ b/test/lit/pto/subview_bind_tile_preserve_stride.pto @@ -39,10 +39,10 @@ module { } } -// CHECK-DAG: Tile -// CHECK-DAG: Tile -// CHECK-DAG: Tile -// CHECK-DAG: Tile +// CHECK-DAG: Tile +// CHECK-DAG: Tile +// CHECK-DAG: Tile +// CHECK-DAG: Tile // CHECK-DAG: TSTORE( // CHECK-DAG: TSTORE( // CHECK-DAG: TSTORE( diff --git a/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto b/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto index cd72ab78ae..d4fd3668d2 100644 --- a/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto +++ b/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto @@ -34,7 +34,7 @@ module { // IR: %[[ELEM_BYTES:.*]] = arith.constant 4 : i64 // IR: %[[BYTE_OFFSET:.*]] = arith.muli %[[ELEM_OFFSET]], %[[ELEM_BYTES]] : i64 // IR: %[[ADDR:.*]] = arith.addi %{{.*}}, %[[BYTE_OFFSET]] : i64 -// IR: pto.alloc_tile addr = %[[ADDR]] {pto.view_semantics = "subview"} : !pto.tile_buf +// IR: pto.alloc_tile addr = %[[ADDR]] {pto.view_semantics = "subview"} : !pto.tile_buf // IR-NOT: pto.materialize_tile // EMITC: Tile diff --git a/test/lit/pto/subview_col_major_compact_keeps_normal_shape.pto b/test/lit/pto/subview_col_major_compact_keeps_normal_shape.pto index 4f635b6225..108c0b070a 100644 --- a/test/lit/pto/subview_col_major_compact_keeps_normal_shape.pto +++ b/test/lit/pto/subview_col_major_compact_keeps_normal_shape.pto @@ -20,5 +20,5 @@ module { } } -// CHECK: Tile -// CHECK-NOT: Tile +// CHECK: Tile +// CHECK-NOT: Tile diff --git a/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto b/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto index 9908f5361f..644bbe8386 100644 --- a/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto +++ b/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto @@ -21,4 +21,4 @@ module { } } -// CHECK: Tile +// CHECK: Tile diff --git a/test/lit/pto/subview_compact_keeps_normal_shape.pto b/test/lit/pto/subview_compact_keeps_normal_shape.pto index 82092c7745..ae18825a05 100644 --- a/test/lit/pto/subview_compact_keeps_normal_shape.pto +++ b/test/lit/pto/subview_compact_keeps_normal_shape.pto @@ -20,5 +20,5 @@ module { } } -// CHECK: Tile -// CHECK-NOT: Tile +// CHECK: Tile +// CHECK-NOT: Tile diff --git a/test/lit/pto/subview_explicit_result_type.pto b/test/lit/pto/subview_explicit_result_type.pto index b76ce90a41..e6f74e5705 100644 --- a/test/lit/pto/subview_explicit_result_type.pto +++ b/test/lit/pto/subview_explicit_result_type.pto @@ -21,4 +21,4 @@ module { } } -// CHECK: Tile +// CHECK: Tile diff --git a/test/lit/pto/subview_tmatmul_acc_slices_a5.pto b/test/lit/pto/subview_tmatmul_acc_slices_a5.pto index e2cf4f04c4..102c03a4cf 100644 --- a/test/lit/pto/subview_tmatmul_acc_slices_a5.pto +++ b/test/lit/pto/subview_tmatmul_acc_slices_a5.pto @@ -24,5 +24,5 @@ module { } } -// CHECK-COUNT-2: Tile +// CHECK-COUNT-2: Tile // CHECK-COUNT-2: TMATMUL( diff --git a/test/lit/pto/subview_validshape_guard.pto b/test/lit/pto/subview_validshape_guard.pto index 52d0caa9e7..2720df1937 100644 --- a/test/lit/pto/subview_validshape_guard.pto +++ b/test/lit/pto/subview_validshape_guard.pto @@ -66,14 +66,14 @@ module { // Default non-compact subview keeps parent physical shape (4x16) but valid // defaults to subview shape (2x8), not clipped from parent valid shape. -// CHECK: Tile +// CHECK: Tile // CHECK-NOT: Tile -// Default row-major full-col subview can use the compact child physical shape (2x8). -// CHECK: Tile +// Default row-major full-col subview keeps parent physical shape (4x8). +// CHECK: Tile // Explicit valid override takes effect. -// CHECK: Tile +// CHECK: Tile // Dynamic explicit valid keeps dynamic valid dims at lowering. -// CHECK: Tile +// CHECK: Tile diff --git a/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto b/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto index df24d9cb57..973ef5ee60 100644 --- a/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto +++ b/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto @@ -15,5 +15,5 @@ module { } } -// CHECK: Tile +// CHECK: Tile // CHECK-NOT: error: diff --git a/test/lit/pto/tile_assemble_level3_subview_tmov_emitc.pto b/test/lit/pto/tile_assemble_level3_subview_tmov_emitc.pto new file mode 100644 index 0000000000..85a92aff72 --- /dev/null +++ b/test/lit/pto/tile_assemble_level3_subview_tmov_emitc.pto @@ -0,0 +1,43 @@ +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --enable-insert-sync %s | FileCheck %s + +module attributes {"pto.target_arch" = "a5"} { + func.func @tile_assemble_level3_subview_tmov_emitc( + %x: memref<32x32xf32, #pto.address_space>, + %src: memref<32x16xf32, #pto.address_space>, + %out: memref<32x32xf32, #pto.address_space>) { + %c0 = arith.constant 0 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %addr0 = arith.constant 0 : i64 + %addr1 = arith.constant 4096 : i64 + + %tile_x = pto.alloc_tile addr = %addr0 : + !pto.tile_buf + %tile_src = pto.alloc_tile addr = %addr1 : + !pto.tile_buf + + pto.tload ins(%x : memref<32x32xf32, #pto.address_space>) + outs(%tile_x : !pto.tile_buf) + pto.tload ins(%src : memref<32x16xf32, #pto.address_space>) + outs(%tile_src : !pto.tile_buf) + %view = pto.subview %tile_x[%c0, %c16] sizes [32, 16] valid [%c32, %c16] : + !pto.tile_buf -> + !pto.tile_buf + pto.tmov ins(%tile_src : + !pto.tile_buf) + outs(%view : !pto.tile_buf) + pto.tstore ins(%tile_x : + !pto.tile_buf) + outs(%out : memref<32x32xf32, #pto.address_space>) + return + } +} + +// CHECK-LABEL: AICORE void tile_assemble_level3_subview_tmov_emitc +// CHECK: set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); +// CHECK: Tile>, + %src: memref<32x16xf32, #pto.address_space>, + %out: memref<32x32xf32, #pto.address_space>) { + %c0 = arith.constant 0 : index + %addr0 = arith.constant 0 : i64 + %addr1 = arith.constant 4096 : i64 + + %tile_x = pto.alloc_tile addr = %addr0 : + !pto.tile_buf + %tile_src = pto.alloc_tile addr = %addr1 : + !pto.tile_buf + + pto.tload ins(%x : memref<32x32xf32, #pto.address_space>) + outs(%tile_x : !pto.tile_buf) + pto.tload ins(%src : memref<32x16xf32, #pto.address_space>) + outs(%tile_src : !pto.tile_buf) + pto.tinsert ins(%tile_src, %c0, %c0 : + !pto.tile_buf, index, index) + outs(%tile_x : !pto.tile_buf) + pto.tstore ins(%tile_x : + !pto.tile_buf) + outs(%out : memref<32x32xf32, #pto.address_space>) + return + } +} + +// CHECK-LABEL: func.func @tinsert_level3_tile_native_sync +// CHECK: pto.tload +// CHECK: pto.tload +// CHECK: pto.set_flag[, , ] +// CHECK: pto.wait_flag[, , ] +// CHECK: pto.tinsert +// CHECK: pto.set_flag[, , ] +// CHECK: pto.wait_flag[, , ] From 9a16d4b46ed94225aa8a985ca086045959fd96dd Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 14 Jul 2026 17:50:42 +0800 Subject: [PATCH 07/91] Treat tile bitcast as memory-planning alias --- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 4 ++++ lib/PTO/Transforms/Utils.cpp | 2 ++ 2 files changed, 6 insertions(+) diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 1e9b30402d..04a555a6fb 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -690,6 +690,10 @@ struct PlannerAnalysis { setRoots(subview.getResult(), getRoots(subview.getSource())); propagateSplitTpopDerived(subview.getResult(), ValueRange{subview.getSource()}); + } else if (auto bitcast = dyn_cast(op)) { + setRoots(bitcast.getResult(), getRoots(bitcast.getSrc())); + propagateSplitTpopDerived(bitcast.getResult(), + ValueRange{bitcast.getSrc()}); } else if (auto reshape = dyn_cast(op)) { setRoots(reshape.getResult(), getRoots(reshape.getSrc())); propagateSplitTpopDerived(reshape.getResult(), diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index 6356cc7b83..2b34417c10 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -157,6 +157,8 @@ std::optional> getOperationAliasInfo(Operation *op) { return std::make_pair(subViewOp.getResult(), subViewOp.getViewSource()); } else if (auto subViewOp = dyn_cast(op)) { return std::make_pair(subViewOp.getResult(), subViewOp.getSource()); + } else if (auto bitcastOp = dyn_cast(op)) { + return std::make_pair(bitcastOp.getResult(), bitcastOp.getSrc()); } else if (auto reshapeOp = dyn_cast(op)) { return std::make_pair(reshapeOp.getResult(), reshapeOp.getSrc()); } else if (auto bindTileOp = dyn_cast(op)) { From ebc9a2995278b103eac360a74c23329f493ffc98 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 14 Jul 2026 20:45:13 +0800 Subject: [PATCH 08/91] Lower PTO castptr in EmitC backend --- lib/PTO/Transforms/PTOToEmitC.cpp | 46 ++++++++++++++++++++++++++++++- 1 file changed, 45 insertions(+), 1 deletion(-) diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 1d8ad818ae..30b8b7e292 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -4648,6 +4648,50 @@ static FailureOr buildSyncAllWorkspaceTileValue( //===----------------------------------------------------------------------===// // pto.pointer_cast lowering //===----------------------------------------------------------------------=== + +struct CastPtrConversion : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::CastPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type convertedResultType = + getTypeConverter()->convertType(op.getResult().getType()); + if (!convertedResultType) + return failure(); + + Value input = adaptor.getInput(); + if (input.getType() == convertedResultType) { + rewriter.replaceOp(op, input); + return success(); + } + + if (auto resultPtrTy = dyn_cast(op.getResult().getType())) { + std::string elemTok = getEmitCScalarTypeToken(resultPtrTy.getElementType()); + std::optional as = + getAddressSpaceOrGM(resultPtrTy.getMemorySpace()); + if (!as) + return rewriter.notifyMatchFailure(op, "unsupported ptr address space"); + + Value ptr = materializeAddressAsPointer(rewriter, op.getLoc(), input, *as, + elemTok); + if (ptr.getType() != convertedResultType) + ptr = rewriter.create(op.getLoc(), convertedResultType, ptr) + .getResult(); + rewriter.replaceOp(op, ptr); + return success(); + } + + if (emitc::isSupportedEmitCType(input.getType()) && + emitc::isSupportedEmitCType(convertedResultType)) { + rewriter.replaceOpWithNewOp(op, convertedResultType, input); + return success(); + } + + return rewriter.notifyMatchFailure(op, "unsupported castptr conversion"); + } +}; + struct PointerCastConversion : public OpConversionPattern { static bool getIndexConst(Value v, int64_t &out) { if (auto cst = v.getDefiningOp()) { @@ -14152,7 +14196,7 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add(typeConverter, ctx); + patterns.add(typeConverter, ctx); patterns.add Date: Thu, 16 Jul 2026 11:07:37 +0800 Subject: [PATCH 09/91] Fix EmitC castptr ptr/int lowering --- .../ptoas-dual-memory-planner-design.md | 188 ------------------ ...st-first-fit-four-gates-memplan-design.md} | 6 +- lib/PTO/Transforms/PTOToEmitC.cpp | 62 +++++- test/lit/pto/castptr_emitc_tile_and_int.pto | 28 +++ 4 files changed, 89 insertions(+), 195 deletions(-) delete mode 100644 docs/designs/ptoas-dual-memory-planner-design.md rename docs/designs/{ptoas-largest-first-fit-five-gates-memplan-design.md => ptoas-largest-first-fit-four-gates-memplan-design.md} (99%) create mode 100644 test/lit/pto/castptr_emitc_tile_and_int.pto diff --git a/docs/designs/ptoas-dual-memory-planner-design.md b/docs/designs/ptoas-dual-memory-planner-design.md deleted file mode 100644 index 0c0711e78b..0000000000 --- a/docs/designs/ptoas-dual-memory-planner-design.md +++ /dev/null @@ -1,188 +0,0 @@ -# PTOAS 双内存规划器设计说明 - -## 背景 - -PTOAS 之前曾临时删除旧版 `pto-plan-memory` pass，以便让 level1/level2 在没有完整 memplan 的情况下继续编译。同时，`pto.alloc_tile` 的 lowering 路径也从旧的 `pto.alloc_tile -> memref.alloc -> pto.pointer_cast -> pto.bind_tile` 调整为 tile-native 路径：无显式地址的 `pto.alloc_tile` 保留到 memplan，由 memplan 直接补充常量 `addr`，后续 `pto.t*` tile op 继续消费 `!pto.tile_buf`。 - -随着旧 memplan 恢复，需要同时满足两个需求： - -- 默认行为保持稳定，继续使用旧 memplan，避免影响现有 pipeline 和测试基线。 -- 保留重写后的 modern memplan，便于继续开发 alias-aware、SPEC_LEVEL_0 复用等新策略。 - -因此，本次提交将旧版和新版 memplan 拆成两个实现，由 `tools/ptoas/ptoas.cpp` 通过 CLI 参数选择使用哪一个。 - -## 目标 - -本次改动的目标如下： - -- 默认使用 legacy memplan，保证既有行为和兼容性。 -- 通过 `--plan-memory-impl=modern` 显式启用 modern memplan。 -- `pto-plan-memory` 仍保持 module 级 pass，pass 内部遍历 `func::FuncOp` 做规划。 -- `pto.alloc_tile(no addr)` 不再经过 `memref.alloc` 中间路径，而是由所选 memplan 直接补 `addr`。 -- 尽量减少对 legacy memplan 原有代码的结构性改动，只补充 tile-native alloc 所需的最小适配。 -- 将原先删除的 memplan lit/sample 用例恢复，并让现有 `plan_memory_*.pto` lit 同时覆盖 legacy 和 modern。 - -## 非目标 - -本次改动不试图完成 modern memplan 的全部设计目标： - -- 不实现 SPEC_LEVEL_1 / SPEC_LEVEL_2 的完整投机复用策略。 -- 不替换 legacy memplan 的 pipeline conflict、double-buffer 相关逻辑。 -- 不改变 level3 显式地址模式的规则。 -- 不把 `reserve_buffer` 手动 base 支持扩展到 level1/level2。 - -## 设计方案 - -### 文件组织 - -旧 memplan 仍保留在原文件： - -```text -lib/PTO/Transforms/PTOPlanMemory.cpp -lib/PTO/Transforms/PTOPlanMemory.h -``` - -modern memplan 放在独立文件： - -```text -lib/PTO/Transforms/PTOPlanMemoryModern.cpp -``` - -`PTOPlanMemory.cpp` 继续定义默认 pass factory： - -```cpp -createPlanMemoryPass(const PlanMemoryOptions &options = {}) -``` - -`PTOPlanMemoryModern.cpp` 定义 modern factory： - -```cpp -createPlanMemoryModernPass(const PlanMemoryOptions &options) -``` - -`lib/PTO/Transforms/CMakeLists.txt` 同时编译两个实现文件。 - -### Pass 级别 - -`pto-plan-memory` 保持 module 级 pass： - -```td -def PlanMemory : Pass<"pto-plan-memory", "ModuleOp"> -``` - -legacy 和 modern 都在 module pass 的 `runOnOperation()` 中遍历 module 内的 `func::FuncOp`，分别对每个函数做静态 local memory planning。这样保留旧 memplan 的 module-level 调用形态，也避免在 pass manager pipeline 中暴露两个不同的 nested func pass。 - -### 实现选择 - -实现选择不放在 `PTOPlanMemory.cpp` 内部，而是在 `tools/ptoas/ptoas.cpp` 中完成： - -```cpp -if (planMemoryImpl == "legacy") { - pm.addPass(pto::createPlanMemoryPass(planMemoryOptions)); -} else if (planMemoryImpl == "modern") { - pm.addPass(pto::createPlanMemoryModernPass(planMemoryOptions)); -} -``` - -CLI 参数为： - -```text ---plan-memory-impl=legacy # 默认 ---plan-memory-impl=modern -``` - -这样 `PTOPlanMemory.cpp` 可以继续表示 legacy memplan 本体，而不是变成 dispatcher。 - -### `memMode` - -`PlanMemoryOptions::memMode` 当前仍由 `ptoas.cpp` 固定设置为： - -```cpp -planMemoryOptions.memMode = "local"; -``` - -含义是规划 local memory buffer。legacy memplan 内部仍保留 `MemPlanMode` 枚举，用于尽量少改旧代码： - -```cpp -enum class MemPlanMode { - LOCAL_MEM_PLAN, - GLOBAL_WORKSPACE_PLAN, -}; -``` - -当前 PTOAS 主 pipeline 只使用 `"local"`。`GLOBAL_WORKSPACE_PLAN` 是 legacy 代码的历史模式保留，不作为当前 level1/level2 local memplan 的正式入口。 - -### `pto.alloc_tile(no addr)` 处理 - -`PTOViewToMemref` 不再把无地址的 `pto.alloc_tile` 降成 `memref.alloc`。因此，两个 memplan 都需要直接处理 tile-native allocation root： - -```mlir -%tile = pto.alloc_tile : !pto.tile_buf<...> -``` - -memplan 完成规划后直接补常量地址： - -```mlir -%c0_i64 = arith.constant 0 : i64 -%tile = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf<...> -``` - -legacy memplan 为此增加了最小适配： - -- liveness 收集阶段识别 `pto.alloc_tile(no addr)`。 -- `BufferInfo` 支持从 `TileBufType` 计算 shape、element type 和字节数。 -- materialize 阶段新增 `pto.alloc_tile -> pto.alloc_tile addr` 的 rewrite。 -- `pto.tile_buf_addr` 被建模为读取 tile source，避免被未知 local-buffer op 检查误判。 - -### `memref.alloc` 处理 - -非 tile-native 的 `memref.alloc` root 仍沿用 legacy/modern 各自的既有 materialization 方式： - -```mlir -memref.alloc -> pto.pointer_cast(...) -``` - -也就是说，本次改动只要求 `pto.alloc_tile` 不再经过 `memref.alloc` 中间链路；并不禁止普通 memref root 继续使用 `pto.pointer_cast` 表达规划地址。 - -## 测试方案 - -测试覆盖分为三类。 - -### 恢复旧 memplan 测试 - -恢复 `test/samples/planmemory` 下之前删除的 sample case，覆盖 loop、if、nested loop、fragmentation、peak capacity、reuse 等场景。 - -### lit 双实现覆盖 - -现有 `test/lit/pto/plan_memory_*.pto` 增加 modern RUN： - -```mlir -// RUN: ptoas ... | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern ... | FileCheck %s -``` - -`order_by_size` 用例同时覆盖 legacy/modern 的默认顺序和 `--plan-memory-order-by-size` 顺序。负例也同时覆盖 legacy/modern。 - -### 回归命令 - -本次提交使用以下命令验证： - -```bash -cmake --build build --target ptoas -j8 - -PATH=/Users/fangrui/workspace/huawei/llvm21-workspace/llvm-project/llvm/build-assert/bin:$PATH \ - /Users/fangrui/workspace/huawei/llvm21-workspace/llvm-project/llvm/build-assert/bin/llvm-lit \ - -sv build/test/lit \ - --filter 'plan_memory|reserve_buffer|alloc_tile_addr|alloc_tile_plan_memory_no_memref_alloc|multi_tile_get_const_slot_lowering|multi_tile_.*planmem' - -ctest --test-dir build --output-on-failure -L PTODSL -``` - -同时对恢复的 `test/samples/planmemory/*.py` 做逐个生成和 `ptoas --emit-pto-ir` 编译验证。 - -## 后续工作 - -- modern memplan 后续继续补齐 SPEC_LEVEL_1 / SPEC_LEVEL_2。 -- 根据 modern 策略差异，为必要用例拆分 legacy/modern 的独立 FileCheck 前缀。 -- 如果确认 global workspace planning 不再需要，可单独清理 legacy 内部的 `GLOBAL_WORKSPACE_PLAN` 历史模式。 -- 继续减少 legacy 文件中的非必要 diff，保持旧实现可对照、可回退。 diff --git a/docs/designs/ptoas-largest-first-fit-five-gates-memplan-design.md b/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md similarity index 99% rename from docs/designs/ptoas-largest-first-fit-five-gates-memplan-design.md rename to docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md index 39e865910d..3bd87e6b03 100644 --- a/docs/designs/ptoas-largest-first-fit-five-gates-memplan-design.md +++ b/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md @@ -1,4 +1,4 @@ -# PTOAS Largest-First-Fit 与五道冲突闸门内存规划设计 +# PTOAS Largest-First-Fit 与四道冲突闸门内存规划设计 ## 总体方案 @@ -40,7 +40,7 @@ pto.alloc_tile(no addr) 本设计不包含以下内容： - 不修改 legacy memplan 的 StorageEntry / SPEC_LEVEL_1 / SPEC_LEVEL_2 逻辑。 -- 不在 legacy memplan 中实现五道闸门。 +- 不在 legacy memplan 中实现四道闸门。 - 不恢复旧版回滚式投机规划。 - 不把 modern memplan 作为默认实现。 - 不在本阶段实现跨函数、跨 module 的全局内存规划。 @@ -530,7 +530,7 @@ forbidAlias[%tmp].insert(%dst) #### 4.4 闸门 4 判定 -`canShare(a, b)` 中的闸门 3 是双向检查： +`canShare(a, b)` 中的闸门 4 是双向检查： ```text Gate4_OpSemanticNoAlias(a, b): diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 30b8b7e292..4fbe445370 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -4661,8 +4661,9 @@ struct CastPtrConversion : public OpConversionPattern { return failure(); Value input = adaptor.getInput(); - if (input.getType() == convertedResultType) { - rewriter.replaceOp(op, input); + Value peeledInput = peelUnrealized(input); + if (peeledInput.getType() == convertedResultType) { + rewriter.replaceOp(op, peeledInput); return success(); } @@ -4673,8 +4674,18 @@ struct CastPtrConversion : public OpConversionPattern { if (!as) return rewriter.notifyMatchFailure(op, "unsupported ptr address space"); - Value ptr = materializeAddressAsPointer(rewriter, op.getLoc(), input, *as, - elemTok); + if (isEmitCTileLikeType(peeledInput.getType())) { + Value ptr = rewriter + .create( + op.getLoc(), convertedResultType, "PTOAS__TILE_DATA", + ArrayAttr{}, ArrayAttr{}, ValueRange{peeledInput}) + .getResult(0); + rewriter.replaceOp(op, ptr); + return success(); + } + + Value ptr = materializeAddressAsPointer(rewriter, op.getLoc(), peeledInput, + *as, elemTok); if (ptr.getType() != convertedResultType) ptr = rewriter.create(op.getLoc(), convertedResultType, ptr) .getResult(); @@ -4682,6 +4693,49 @@ struct CastPtrConversion : public OpConversionPattern { return success(); } + if (isa(op.getResult().getType()) && + emitc::isSupportedEmitCType(convertedResultType)) { + Value source = input; + if (!emitc::isSupportedEmitCType(source.getType())) { + if (auto inputPtrTy = dyn_cast(op.getInput().getType())) { + std::string elemTok = + getEmitCScalarTypeToken(inputPtrTy.getElementType()); + std::optional as = + getAddressSpaceOrGM(inputPtrTy.getMemorySpace()); + if (!as) + return rewriter.notifyMatchFailure(op, + "unsupported ptr address space"); + if (isEmitCTileLikeType(peeledInput.getType())) { + Type convertedInputType = + getTypeConverter()->convertType(op.getInput().getType()); + if (!convertedInputType) + return rewriter.notifyMatchFailure(op, + "failed to convert input ptr type"); + source = rewriter + .create( + op.getLoc(), convertedInputType, "PTOAS__TILE_DATA", + ArrayAttr{}, ArrayAttr{}, ValueRange{peeledInput}) + .getResult(0); + } else { + source = materializeAddressAsPointer(rewriter, op.getLoc(), + peeledInput, *as, elemTok); + } + } + } + if (!emitc::isSupportedEmitCType(source.getType())) + return rewriter.notifyMatchFailure(op, + "unsupported castptr integer source"); + auto templateArgs = rewriter.getArrayAttr( + {emitc::OpaqueAttr::get(rewriter.getContext(), + cast(convertedResultType) + .getValue())}); + auto cast = rewriter.create( + op.getLoc(), convertedResultType, "reinterpret_cast", ArrayAttr{}, + templateArgs, ValueRange{source}); + rewriter.replaceOp(op, cast.getResult(0)); + return success(); + } + if (emitc::isSupportedEmitCType(input.getType()) && emitc::isSupportedEmitCType(convertedResultType)) { rewriter.replaceOpWithNewOp(op, convertedResultType, input); diff --git a/test/lit/pto/castptr_emitc_tile_and_int.pto b/test/lit/pto/castptr_emitc_tile_and_int.pto new file mode 100644 index 0000000000..751c9c698b --- /dev/null +++ b/test/lit/pto/castptr_emitc_tile_and_int.pto @@ -0,0 +1,28 @@ +// 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 --cann-output-version=9.0.0 --pto-arch=a5 --pto-level=level3 %s -o - 2>&1 | FileCheck %s + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + func.func @castptr_ptr_to_int(%src: memref<1xf32, #pto.address_space>) + attributes {pto.kernel} { + func.call @helper(%src) : (memref<1xf32, #pto.address_space>) -> i64 + return + } + + func.func private @helper(%src: memref<1xf32, #pto.address_space>) -> i64 { + %ptr = builtin.unrealized_conversion_cast %src + : memref<1xf32, #pto.address_space> to !pto.ptr + %addr = pto.castptr %ptr : !pto.ptr -> i64 + return %addr : i64 + } +} + +// CHECK-LABEL: AICORE void castptr_ptr_to_int(__ubuf__ float* +// CHECK: helper( +// CHECK: reinterpret_cast From d3faad75bbb60137716bfa99ff6e92fac8c70ab0 Mon Sep 17 00:00:00 2001 From: FangRui Date: Thu, 16 Jul 2026 11:41:18 +0800 Subject: [PATCH 10/91] Fix flaky plan memory lit checks --- test/lit/pto/castptr_emitc_tile_and_int.pto | 3 ++- test/lit/pto/plan_memory_reused_tstore_sync_level2.pto | 10 ++++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/test/lit/pto/castptr_emitc_tile_and_int.pto b/test/lit/pto/castptr_emitc_tile_and_int.pto index 751c9c698b..4bd3f4d7ee 100644 --- a/test/lit/pto/castptr_emitc_tile_and_int.pto +++ b/test/lit/pto/castptr_emitc_tile_and_int.pto @@ -23,6 +23,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind // CHECK-LABEL: AICORE void castptr_ptr_to_int(__ubuf__ float* // CHECK: helper( -// CHECK: reinterpret_cast diff --git a/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto b/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto index 91fd32e8b4..06a47e100b 100644 --- a/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto +++ b/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto @@ -8,7 +8,7 @@ // End-to-end level2 coverage for PlanMemory reuse followed by InsertSync. The // source uses logical alloc_tile operations without explicit addresses, so the -// physical pointer_cast addresses checked below are produced by PlanMemory. +// physical alloc_tile addresses checked below are produced by PlanMemory. // `stored` occupies [8192, 16384), while `reused` occupies [8448, 16640). // // RUN: ptoas --pto-arch=a3 --pto-level=level2 --enable-insert-sync --emit-pto-ir --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=PLAN @@ -70,13 +70,11 @@ module attributes {pto.target_arch = "a2a3"} { } // PLAN-LABEL: func.func @plan_memory_reused_tstore_sync_level2( -// PLAN-DAG: %[[STORED_ADDR:[A-Za-z0-9_]+]] = arith.constant 8192 : i64 // PLAN-DAG: %[[REUSED_ADDR:[A-Za-z0-9_]+]] = arith.constant 8448 : i64 -// PLAN: %[[STORED_CAST:[0-9]+]] = pto.pointer_cast(%[[STORED_ADDR]]){{.*}} : memref<16x128xf32 -// PLAN-NEXT: %[[STORED:[0-9]+]] = pto.bind_tile %[[STORED_CAST]] +// PLAN-DAG: %[[STORED_ADDR:[A-Za-z0-9_]+]] = arith.constant 8192 : i64 +// PLAN: %[[STORED:[0-9]+]] = pto.alloc_tile addr = %[[STORED_ADDR]] : !pto.tile_buf // PLAN: pto.tstore ins(%[[STORED]] -// PLAN: %[[REUSED_CAST:[0-9]+]] = pto.pointer_cast(%[[REUSED_ADDR]]){{.*}} : memref<16x128xf32 -// PLAN-NEXT: %[[REUSED:[0-9]+]] = pto.bind_tile %[[REUSED_CAST]] +// PLAN: %[[REUSED:[0-9]+]] = pto.alloc_tile addr = %[[REUSED_ADDR]] : !pto.tile_buf // PLAN: pto.tmuls {{.*}} outs(%[[REUSED]] // SYNC-LABEL: AICORE void plan_memory_reused_tstore_sync_level2( From 8ffa703a9cb111f185af71734fa9cf914d022630 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 17 Jul 2026 18:59:20 +0800 Subject: [PATCH 11/91] Keep multi-tile buffers native through memory planning --- ...ptoas-tile-native-mainline-op-migration.md | 377 ++++++++++++++++++ include/PTO/IR/PTOMultiBuffer.h | 11 +- include/PTO/IR/PTOOps.td | 28 +- .../Transforms/InsertSync/PTOIRTranslator.h | 4 + include/PTO/Transforms/Passes.td | 27 +- lib/PTO/IR/PTO.cpp | 37 ++ .../Transforms/GraphSyncSolver/MemInfo.cpp | 67 ++++ .../SyncSolverIRTranslator.cpp | 7 +- .../Transforms/InsertSync/PTOIRTranslator.cpp | 74 +++- .../Transforms/PTOMaterializeTileHandles.cpp | 14 + lib/PTO/Transforms/PTOPlanMemory.cpp | 54 +++ lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 66 ++- lib/PTO/Transforms/PTOResolveBufferSelect.cpp | 105 +++++ lib/PTO/Transforms/PTOViewToMemref.cpp | 253 +----------- lib/PTO/Transforms/SlotAffineAnalysis.cpp | 2 + lib/PTO/Transforms/Utils.cpp | 8 + .../lit/pto/multi_tile_buf_n3_planmem_e2e.pto | 5 +- ...lti_tile_const_preload_dyn_loop_select.pto | 23 +- .../multi_tile_get_const_slot_lowering.pto | 40 +- .../pto/multi_tile_get_dyn_slot_lowering.pto | 27 +- .../pto/multi_tile_level3_explicit_addr.pto | 31 +- test/lit/pto/multi_tile_n4_planmem_e2e.pto | 28 +- test/lit/pto/multi_tile_no_loop_unroll.pto | 15 +- .../pto/multi_tile_reject_internal_addrs.pto | 23 ++ tools/ptoas/ptoas.cpp | 12 + 25 files changed, 975 insertions(+), 363 deletions(-) create mode 100644 docs/designs/ptoas-tile-native-mainline-op-migration.md create mode 100644 test/lit/pto/multi_tile_reject_internal_addrs.pto diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md new file mode 100644 index 0000000000..3ceb81502d --- /dev/null +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -0,0 +1,377 @@ +# PTOAS Tile-Native Mainline 按 Op 迁移设计 + +## 1. 背景 + +PTOAS 当前主 pipeline 在内存规划前后包含两个 ModuleOp pass： + +```text +PTOViewToMemref + -> implicit tmp materialization + -> legacy/modern memplan + -> InsertSync/GSS/BufidSync + -> PTOResolveBufferSelect + -> PTOMaterializeTileHandles + -> EmitC/VPTO +``` + +`PTOViewToMemref` 把函数 ABI、PTO view、部分 local tile 和计算 op 转成 +memref 形态；`PTOMaterializeTileHandles` 在规划和同步完成后重新构造 +`!pto.tile_buf` handle。普通 `pto.alloc_tile` 已经是例外：它保持 +tile-native，由 memplan 直接填写 `addr`。 + +本设计的目标是逐个 op 消除上述往返转换，最终删除： + +- `lib/PTO/Transforms/PTOViewToMemref.cpp` +- `lib/PTO/Transforms/PTOMaterializeTileHandles.cpp` +- `pto-view-to-memref` +- `pto-materialize-tile-handles` + +迁移过程中不增加用户可见的双 pipeline 开关。每个 op 在现有 pipeline 中 +先补齐 tile-native 分析和 lowering；全部 op 完成后一次性删除两个 pass。 + +## 2. 目标 IR + +```text +PTO tile/view IR + -> implicit tmp materialization（产生 pto.alloc_tile(no addr)） + -> memplan（直接填写 alloc_tile/alloc_multi_tile 地址） + -> sync（直接分析 PTO root、view、slot 和物理地址） + -> EmitC/VPTO（直接消费 PTO tile/view IR） +``` + +基本原则： + +1. local allocation root 使用 PTO allocation op 表达，不中转 `memref.alloc`。 +2. tile view 保持 `!pto.tile_buf`，view op 自身携带 shape/layout/offset 语义。 +3. GM tensor view 保持 PTO view op，直到 EmitC/VPTO backend lowering。 +4. memplan、InsertSync 和 backend 各只有一套算法主体，通过统一的 root/alias/address + 查询处理不同 PTO op。 +5. `pto.bind_tile`、`pto.declare_tile_memref`、`pto.slot_marker` 仅在完成对应 + tile-native 迁移后删除。 + +## 3. 每个 Op 的完成标准 + +每个 op 只有同时满足以下条件才算迁移完成： + +- verifier 和 parse/print 能表达目标 tile-native 语义； +- legacy 和 modern memplan 能识别其 root/alias/size/address； +- InsertSync、Graph Sync Solver 和 BufidSync 能恢复其物理内存访问； +- EmitC 和 VPTO 都能直接 lowering； +- level1/level2/level3 行为明确； +- A2/A3 和 A5 行为明确； +- 至少有一个正例和必要的 verifier/level 负例； +- 测试不再检查 `bind_tile`、`memref.subview` 等旧中间形态。 + +## 4. Allocation 与内部 Handle Op + +| Op | 当前行为 | Tile-native 目标 | Memplan | InsertSync | Backend 与测试 | +|---|---|---|---|---|---| +| `pto.alloc_tile` | `PTOViewToMemref` 已透传；memplan 填写 `addr` | 保持现状，作为单 slot local root | legacy/modern 均继续收集无地址 root；level3 校验显式地址 | 现有 InsertSync 已读取类型、大小和常量地址；GSS 需从 `alloc_tile addr` 构造 `PointerLikeInfo` | EmitC/VPTO 已有直接 lowering；保留 level1/2 自动地址和 level3 显式地址测试 | +| `pto.alloc_multi_tile` | 已保持 tile-native；memplan 写入内部 N-address 属性，level3 保留用户 base | op 本身作为 N-slot root，最终可由 backend 直接消费 | legacy/modern 已支持 `slotCount=N`、`slotBytes`、`totalBytes` 和每 slot offsets；禁止 sibling slot 互相 alias | InsertSync/GSS 已直接识别 multi root，不依赖 `slot_marker -> pointer_cast` | 当前由 `PTOResolveBufferSelect` 展开成带地址的 `alloc_tile`；已覆盖 constant/dynamic slot、loop/non-loop 测试 | +| `pto.multi_tile_get` | 已保持到 sync 和 `PTOResolveBufferSelect` | 最终继续保持到 backend | result 已 alias 到 multi root，并携带 slot expression | constant slot 只访问一个地址；dynamic slot 保守访问全部地址，并保留 dyn event-id 推导 | 当前 resolve pass 根据 slot 生成带单地址的 `alloc_tile`；不再生成 `slot_marker` | +| `pto.declare_tile` | 转成 `pto.declare_tile_memref + pto.bind_tile` | 保持 tile handle，表示地址由 `tpop` 等运行时操作绑定，不参与静态 local allocation | 不作为普通静态 root；需要单独标记 runtime-bound root | 建立 `declare_tile -> tpop/tassign -> tile users` 的 root/address 数据流 | EmitC/VPTO 直接 lowering；覆盖 tpush/tpop、分支和循环 | +| `pto.declare_tile_memref` | `PTOViewToMemref` 生成的内部占位 op | 所有 `declare_tile` 用户迁移后删除 | 删除 legacy/modern 特判 | 删除 InsertSync 的专用 root 注册 | 删除 EmitC pattern 和相关测试 | +| `pto.bind_tile` | memref 与 tile metadata 的桥接和 alias anchor | local tile 不再需要；只允许迁移期兼容，最终删除或限制为外部 ABI bridge | 先统一 `getAliasSource()`，切换后删除 bind 分支 | 先统一 alias tracing，切换后删除 bind 分支 | 删除 materialize/EmitC bind lowering 测试 | +| `pto.pointer_cast` | 表达 memref root 的物理地址或 multi-address root | 仅保留真正的 raw-address ABI bridge；普通 local allocation 改由 allocation op 持有地址 | 不再作为 `alloc_tile` 的 materialization 结果 | 若保留，sync 继续读取其地址；multi-buffer 不再必须依赖它 | 后端保留 raw-address lowering，删除 alloc fallback 用法 | +| `pto.slot_marker` | `multi_tile_get` 在 memref 层的内部 slot 标签 | `multi_tile_get` 直接承担 slot 标签语义，最终删除 | 删除 alias 特判 | 将 slot narrowing 和 dyn event-id 逻辑迁移到 `multi_tile_get` | 删除 `PTOResolveBufferSelect` 对 slot_marker 的入口 | +| `pto.materialize_tile` | 用于显式恢复 tile handle 的兼容 op | 若无独立用户则随 materialize pass 一并删除 | 无 | 无 | 搜索并删除残余 pattern/test | + +## 5. Tile View 与 Metadata Op + +| Op | 当前行为 | Tile-native 目标 | Memplan / Alias | InsertSync | Backend 与测试 | +|---|---|---|---|---|---| +| `pto.subview` | 通常转成 `memref.subview + pto.bind_tile`，之后再恢复 tile | 全程保留 `pto.subview` | result 与 source 同 root；按 layout/stride/offset 计算 byte range | 现有 InsertSync 已有精确 `pto.subview` 地址计算；GSS 和 BufidSync 补齐同等能力 | EmitC/VPTO 增加直接 lowering；覆盖 row/col-major、dynamic offset、valid shape | +| `pto.treshape` | 转成 memref view，materialize 阶段恢复 `treshape` | 保持原 op | result 与 source 完全 alias，size 不变 | InsertSync/GSS 增加 result-to-source alias 传播 | backend 已有 tile op 路径；覆盖动态 valid shape 和控制流 | +| `pto.bitcast` | 转成 memref view，materialize 阶段恢复 `bitcast` | 保持原 op | result 与 source alias；以 byte range 为准，校验总容量一致 | InsertSync/GSS 增加 alias 传播，不能按 element type 分裂 root | EmitC/VPTO 直接 lowering；覆盖 dtype 改变和相同容量 | +| `pto.set_validshape` | 在 memref 层通过 bind metadata 记录，materialize 后恢复 tile 类型 | 直接更新/产生 tile handle metadata，不改变物理 root | result 与 source alias，物理 size 使用 allocation shape，不使用 valid shape | 同一物理范围；repeat 分析读取更新后的 valid shape | 两后端直接消费；覆盖 if/for 和动态 valid shape | +| `pto.get_validshape` | materialize 后依赖 tile handle 读取 metadata | 直接读取 tile operand metadata | 不产生 root | 无内存 effect | 两后端直接 lowering；覆盖静态折叠和动态值 | +| `pto.tile_buf_addr` | ptr-like 形式转成线性 memref，后续 materialize 可能恢复 | 直接从 tile root/view 计算地址；返回 pointer-like PTO 类型 | 不产生新 root | 若结果参与 load/store，需要保留到原 root 的 provenance | EmitC/VPTO 直接生成地址表达式；覆盖 alloc、subview、multi slot | + +## 6. GM Tensor View 与地址 Op + +| Op | 当前行为 | Tile-native 目标 | Memplan / Sync | Backend 与测试 | +|---|---|---|---|---| +| `pto.make_tensor_view` | 转成 `memref.reinterpret_cast` | 保持 PTO tensor view 到 backend | 不参与 local memplan；InsertSync 已能把 result alias 到 pointer source | EmitC 增加完整 direct pattern；VPTO 保持 canonical rank/layout lowering；覆盖 addptr base、dynamic shape/stride | +| `pto.partition_view` | 转成 `memref.subview` | 保持 PTO partition view | 不参与 local memplan；sync result alias source，并根据 offset/size缩小 GM range | EmitC 已有部分 direct pattern，补齐全部 layout；VPTO 增加直接支持 | +| `pto.get_tensor_view_dim` | 转成 `memref.dim` | 直接读取 PTO view shape | 不影响 memplan/sync | backend 支持静态折叠和动态 dim | +| `pto.get_tensor_view_stride` | 转成 strided memref metadata | 直接读取 PTO view stride | 不影响 memplan/sync | backend 支持静态折叠和动态 stride | +| `pto.inttoptr` | 结果改成 GM memref，并限制用途 | 保持 PTO pointer-like value | 不参与 local memplan；sync 将其视作 GM provenance | 保留 restricted-use verifier；EmitC/VPTO 直接 lowering | +| `pto.ptrtoint` | 折叠 `addptr` 链并生成 byte offset | 保持 PTO op，或迁移到独立 address-canonicalization pass | 不影响 local memplan；不能丢失 GM provenance | backend 生成整数地址；覆盖 ptr、addptr、view base | +| `pto.addptr` | 折叠进 tensor view、scalar load/store 或 pipe init | 保持到独立 address canonicalization/backend | sync 需要把 result alias 到 base，并记录 offset | EmitC/VPTO 直接 lowering；保留非法 escape 校验 | +| `pto.castptr` | 作为 pointer/memref 适配 op | 保持 PTO pointer cast | result alias source，地址空间变化必须校验 | 两后端直接 lowering | +| `pto.load_scalar` | `addptr` offset 被折叠进 op | op 直接接受 base+offset，或 backend 统一折叠 | InsertSync 记录 GM read；不参与 local memplan | 覆盖 inttoptr/addptr 和动态 offset | +| `pto.store_scalar` | `addptr` offset 被折叠进 op | op 直接接受 base+offset，或 backend 统一折叠 | InsertSync 记录 GM write | 覆盖跨 pipe flush/sync | +| `pto.initialize_l2g2l_pipe` | `gm_addr` 上的 addptr 被提前折叠 | 迁移到独立 address canonicalization 或 op verifier/lowering | sync 保留 GM base provenance | EmitC/VPTO 覆盖动态地址 | + +## 7. Control Flow 与函数 ABI + +| Op | 当前行为 | Tile-native 目标 | 需要改动 | +|---|---|---|---| +| `func.func` | tile 参数/结果可能改成 memref，入口插入 `bind_tile` | authored tile ABI 保持不变 | backend helper、外部声明和 kernel ABI 明确区分；memplan 仍按函数独立规划 | +| `func.call` | 依赖 ViewToMemref 的签名桥接，materialize 后恢复 helper tile ABI | caller/callee 直接使用一致的 tile/pointer ABI | call verifier、inline helper 和 InsertSync call effects 需要统一；跨函数不做 local allocation 合并 | +| `scf.if` | pass 重建 result type 和两个 `scf.yield` | tile result 原样穿过分支 | memplan 和 sync 合并两个分支的 root family；backend 支持 tile result | +| `scf.for` | pass 重建 iter_arg/result/yield 类型 | tile loop-carried value 原样保持 | memplan 扩展 loop-carried root 生命周期；sync 建立 iter_arg/yield alias 和 back-edge dependency | +| `scf.yield` | 随外层控制流重建 operand type | 直接 yield tile value | 不能把 yield 当新 root；传播到外层 result/下一轮 iter_arg | +| `pto.fusion_region` | 协调 region result 与 `pto.yield` 类型 | 保持 tile-native region ABI | FusionAnalysis、memplan、sync 和 backend 共享同一 root family | +| `pto.yield` | materialize 阶段恢复 tile operand | 直接 yield tile value | 与 `scf.yield` 相同，传播 alias/liveness | + +## 8. 特殊后端/同步 Op + +| Op | 当前依赖 | 删除两个 pass 后的工作 | +|---|---|---| +| `pto.vlds` | materialize pass 为 view operand 恢复 tile address | backend 和 sync 直接从 tile/view root 计算地址；保持 Read/Write effects | +| `pto.vsts` | 同上 | 直接处理 tile view 地址和 post-update 结果 | +| `pto.vsstb` | 同上 | 直接处理 tile view 地址、packed stride 和 post-update result | +| `pto.mgather` | ViewToMemref 会在无 tile root 的旧路径重建 op | 保持 tile operands；sync macro model 直接读取 index/src/dst/scratch 的 effects | +| `pto.mscatter` | 同上 | 保持 tile operands；直接计算 GM/local ranges 和 coalesce 语义 | +| `pto.tassign` | ViewToMemref 协调 result type，materialize 用它恢复地址 | 明确为 tile handle 地址重绑定；result 类型始终等于 tile operand 类型 | memplan 不把 result 当新 allocation；sync 将其 alias 到 tile root并更新地址 provenance；backend 直接 lowering | + +## 9. Compute Op 逐项处理 + +以下 op 在 `PTOViewToMemref` 的 Stage 3 中主要为了适配 memref operand 而被 +重新创建。tile-native 主路径中不应重建它们。每个 op 的共同要求是： + +- 保持 authored `!pto.tile_buf` operand/result； +- memplan 只通过 MemoryEffects、DPS output 和 alias side table 使用它，不改 op； +- InsertSync 直接读取 MemoryEffects；tmp/scratch 必须保持正确的 Read/Write； +- EmitC/VPTO 的现有 tile lowering 必须在没有 materialize pass 的情况下通过； +- 每个 op 至少保留一个 end-to-end lit 或 sample。 + +| Op | 额外检查 | +|---|---| +| `pto.tload` | GM view 到 local tile 的 MTE2 write；dst 是 writer output，地址来自 `alloc_tile`/subview | +| `pto.tstore` | local tile 到 GM view 的 MTE3 read；检查 reused local addr 与后续 writer 的同步 | +| `pto.ttrans` | optional tmp 的条件 materialization、A5 placeholder 和 scratch conflict | +| `pto.texp` | precision attr 原样保留 | +| `pto.tmul` | DPS inplace 规则和 input/output alias | +| `pto.tmuls` | scalar operand 顺序和 precision attr | +| `pto.tadd` | DPS inplace 规则 | +| `pto.taddc` | carry/result operand 顺序和多输出语义 | +| `pto.tadds` | scalar operand和 valid-shape 校验 | +| `pto.taddsc` | scalar/carry 组合和多输出语义 | +| `pto.tmatmul` | LEFT/RIGHT/ACC 地址空间、tile config 和 role 不再依赖 materialize 推断 | +| `pto.tmatmul_acc` | ACC init/output 的 Read+Write 和 inplace 语义 | +| `pto.tmatmul_bias` | bias root/address space 和 scratch/output conflict | +| `pto.tmatmul_mx` | MX scale tile role 和低精度 packed shape | +| `pto.tmatmul_mx_acc` | MX scale + ACC init/output | +| `pto.tmatmul_mx_bias` | MX scale + bias root | +| `pto.tgemv` | vector/matrix role 和 dst writer lifetime | +| `pto.tgemv_acc` | accumulator init/output effects | +| `pto.tgemv_bias` | bias tile role | +| `pto.tgemv_mx` | MX scale role | +| `pto.tgemv_mx_acc` | MX scale + accumulator effects | +| `pto.tgemv_mx_bias` | MX scale + bias effects | +| `pto.tmov` | identity removal、view metadata 和 src/dst alias | +| `pto.tmov_fp` | fp/pre-quant optional operands和地址空间 | +| `pto.tabs` | unary Read(src)/Write(dst) | +| `pto.tand` | binary input/output effects | +| `pto.tands` | scalar form operand 顺序 | +| `pto.tor` | binary input/output effects | +| `pto.tors` | scalar form operand 顺序 | +| `pto.tnot` | unary effects | +| `pto.tneg` | unary effects | +| `pto.tcmp` | mask/output类型和 compare mode | +| `pto.tcmps` | scalar compare operand 顺序 | +| `pto.tconcat` | 多 input root 的 liveness | +| `pto.tconcatidx` | index/output effects | +| `pto.tci` | implicit tmp、A2/A3 scratch、A5 unused placeholder | +| `pto.tcolexpand` | broadcast source/output shape | +| `pto.tcolexpandmul` | binary broadcast effects | +| `pto.tcolexpandmax` | binary broadcast effects | +| `pto.tcolexpandmin` | binary broadcast effects | +| `pto.tcolmax` | reduction tmp/output effects | +| `pto.tcolmin` | reduction tmp/output effects | +| `pto.tcolsum` | `isBinary` 条件 tmp 和 arch 行为 | +| `pto.tcvt` | 条件 tmp、rmode/satmode operand 顺序 | +| `pto.tdiv` | precision attr 和 inplace policy | +| `pto.tdivs` | scalar 顺序和 precision attr | +| `pto.texpands` | shape扩展和 scalar operand | +| `pto.textract` | tile role、offset 和可选 pre-quant | +| `pto.textract_fp` | fp tile role和地址空间 | +| `pto.tinsert` | materialize pass 当前对 tile config 有特殊推断；目标是由 result type 完整携带 | +| `pto.tinsert_fp` | fp/pre-quant tile role | +| `pto.tfillpad` | src/dst alias 和 A5 MAT/PIPE 选择 | +| `pto.tfillpad_inplace` | same-SSA inplace 和 MemoryEffects | +| `pto.tsetval` | tile writer 和 result type | +| `pto.tgetval` | tile reader和 scalar result | +| `pto.tgather` | optional tmp、compare/index form 和 sync macro model | +| `pto.tgatherb` | mask/index buffer effects | +| `pto.tlog` | precision attr | +| `pto.tlrelu` | unary input/output effects | +| `pto.tmax` | binary inplace policy | +| `pto.tmaxs` | scalar form | +| `pto.tmin` | binary inplace policy | +| `pto.tmins` | scalar form | +| `pto.tquant` | 删除 ViewToMemref 内部 tmp 补全；统一由 `PTOMaterializeImplicitTmp` 负责 | +| `pto.tmrgsort` | format1/format2、optional tmp 和多 source roots | +| `pto.tpartadd` | partition参数和 dst writer | +| `pto.tpartmul` | partition参数和 dst writer | +| `pto.tprint` | optional format tmp、tile read 和无输出语义 | + +以下 op 当前没有出现在 Stage 3 的逐 op 重建列表中，通常已经保持 tile-native； +但 `PTOMaterializeTileHandles` 会泛化处理所有 `pto.t*` 的 memref operand，因此 +删除 pass 前仍需逐项确认它们不依赖 memref metadata 回溯： + +| Op | 额外检查 | +|---|---| +| `pto.trowexpandadd` | optional tmp、A2/A3 scratch 和 shared-tmp PIPE_V dependency | +| `pto.trowexpandsub` | optional tmp 和 scratch conflict | +| `pto.trowexpandmul` | optional tmp 和 inplace dst | +| `pto.trowexpanddiv` | optional tmp、precision attr 和 PIPE_V barrier pruning | +| `pto.trowexpandmax` | optional tmp 和 reduction/broadcast mode | +| `pto.trowexpandmin` | optional tmp 和 reduction/broadcast mode | +| `pto.trowexpandexpdif` | tmp effects 和 dst writer | +| `pto.tcolexpandadd` | broadcast input/output alias | +| `pto.tcolexpandsub` | broadcast input/output alias | +| `pto.tcolexpanddiv` | precision attr 和 broadcast alias | +| `pto.tcolexpandexpdif` | tmp/output effects | +| `pto.trowmax` | A2/A3 tmp、A5 placeholder 和 output lifetime | +| `pto.trowmin` | A2/A3 tmp、A5 placeholder 和 output lifetime | +| `pto.trowsum` | A2/A3 tmp、A5 placeholder 和 output lifetime | +| `pto.trowprod` | A2/A3 tmp、A5 placeholder 和 output lifetime | +| `pto.tcolargmax` | tmp capacity、value/index outputs 和多结果 root | +| `pto.tcolargmin` | tmp capacity、value/index outputs 和多结果 root | +| `pto.trowargmax` | index-only/value+index 模式和条件 tmp | +| `pto.trowargmin` | index-only/value+index 模式和条件 tmp | +| `pto.tprelu` | ui8 tmp、tmp/dst no-alias 和 A5 unused tmp | +| `pto.trem` | 2-row tmp、precision attr 和 scratch effects | +| `pto.trems` | 1-row tmp、scalar operand 和 scratch effects | +| `pto.tsel` | ui32 mask tmp、mask/src/dst effects | +| `pto.tsels` | one-row tmp、scalar select 和 A5 unused tmp | +| `pto.tpow` | A2/A3 floating tmp、integer no-tmp 和 A5 no-tmp | +| `pto.tpows` | A2/A3 floating tmp、scalar exponent 和 A5 no-tmp | +| `pto.trsqrt` | API compatibility tmp 不应触发自动 allocation | +| `pto.tsort32` | 非 32 对齐尾部的条件 tmp | +| `pto.txor` | tmp dtype/shape 和 tmp/output conflict | +| `pto.txors` | scalar operand、tmp dtype/shape 和 conflict | +| `pto.tpartmax` | partition参数和 dst writer | +| `pto.tpartmin` | partition参数和 dst writer | + +其它没有在两个 pass 中出现、且 operand/result 始终是 tile-native 的 `pto.t*` +op 默认不需要迁移实现；仍需通过完整 lit 确认它们没有间接依赖 `bind_tile`、 +`pointer_cast` 或 memref metadata 回溯。 + +## 10. InsertSync 专项改造 + +### 10.1 普通 InsertSync + +现有 `PTOIRTranslator` 已直接支持： + +- `pto.alloc_tile` +- `pto.make_tensor_view` +- `pto.partition_view` +- `pto.subview` +- `pto.bind_tile` +- `memref.subview` + +删除两个 pass 前的支持状态： + +1. `pto.treshape` 和 `pto.bitcast` 的 alias 传播。 +2. `pto.declare_tile` runtime-bound root。 +3. `pto.alloc_multi_tile` 的 N-address root：已完成。 +4. `pto.multi_tile_get` 的 constant/dynamic slot narrowing：已完成。 +5. tile 类型 helper argument 和 `func.call` effects。 +6. `scf.yield/iter_arg`、`pto.yield/fusion result` 的 root family 传播检查。 +7. `tile_buf_addr` 结果到原 tile root 的 provenance。 + +`CanPrunePipeVBarrier()` 已支持 `TileBufType`，不需要因删除两个 pass重写； +但新增 view/multi-buffer alias 后必须确认 dependency pair 仍能区分普通 +output-input RAW 和 scratch WAW/WAR。 + +### 10.2 Graph Sync Solver + +GSS 原先主要从 `pto.pointer_cast` 读取物理地址。`getMemInfo()` 的支持状态： + +- `pto.alloc_tile addr` -> 单地址 `PointerLikeInfo`； +- `pto.alloc_multi_tile addrs` -> N 地址 `PointerLikeInfo`：已完成； +- `pto.multi_tile_get` -> constant slot 单地址或 dynamic slot 全地址：已完成； +- `pto.subview` -> base address加 byte offset并缩小访问范围； +- `treshape/bitcast/set_validshape` -> 传播原 root。 + +否则 memplan 把两个不同 SSA allocation 复用到同一物理地址后，GSS 可能漏掉 +cross-root hazard。 + +### 10.3 BufidSync 与 Macro Model + +检查所有 unwrap helper，去掉对 `bind_tile/memref.subview` 的必需依赖，直接穿透： + +- `pto.subview` +- `pto.treshape` +- `pto.bitcast` +- `pto.multi_tile_get` +- `pto.set_validshape` + +## 11. 推荐迁移顺序 + +按照风险从低到高逐 op 处理： + +1. `pto.treshape` +2. `pto.bitcast` +3. `pto.set_validshape/get_validshape` +4. `pto.subview` +5. `pto.tile_buf_addr` +6. `pto.declare_tile` +7. `pto.alloc_multi_tile` +8. `pto.multi_tile_get` +9. `pto.make_tensor_view` +10. `pto.partition_view` +11. `pto.get_tensor_view_dim/get_tensor_view_stride` +12. `pto.inttoptr/ptrtoint/addptr/castptr` +13. `pto.load_scalar/store_scalar/initialize_l2g2l_pipe` +14. `scf.if/scf.for/scf.yield` +15. `pto.fusion_region/pto.yield` +16. `func.func/func.call/helper ABI` +17. `vlds/vsts/vsstb` +18. 逐项回归第 9 节 compute op +19. 删除 `bind_tile/declare_tile_memref/slot_marker` 兼容路径 +20. 从 `ptoas.cpp` 删除两个 pass，随后删除源文件和 Passes.td 注册 + +`alloc_multi_tile/multi_tile_get` 迁移已完成：地址规划、slot identity 和动态 +event id 均不再依赖 `PTOViewToMemref`。剩余 op 仍按上述顺序逐项迁移。 + +## 12. 测试矩阵 + +每批 op 至少运行： + +```text +level1 + legacy memplan +level1 + modern memplan +level2 + legacy memplan +level2 + modern memplan +level3 explicit address + +A3 + InsertSync +A3 + Graph Sync Solver +A5 + BufidSync + +EmitC backend +VPTO backend +``` + +必须重点恢复/改写以下现有测试类别： + +- `materialize_tile_handles_*` +- `multi_tile_*` +- `subview_*` +- `treshape_*` +- `ptr_int_cast.pto` +- PTODSL subkernel helper ABI +- tile fusion control-flow result +- plan memory reused-address sync +- implicit tmp end-to-end + +旧测试中检查 `memref.alloc`、`memref.subview`、`pto.bind_tile`、 +`pto.pointer_cast` 或 `pto.slot_marker` 的部分，应改为检查 tile-native op、规划后 +地址和最终 backend 输出。 + +## 13. 最终删除条件 + +只有满足以下条件才能删除两个 pass： + +- 主 pipeline 中不再产生需要恢复为 tile 的 local memref； +- legacy/modern memplan 都通过完整测试； +- InsertSync/GSS/BufidSync 都能直接分析 tile-native multi-buffer 和 view； +- EmitC/VPTO 不再依赖 bind/memref metadata 回溯； +- helper、控制流、fusion 和 multi-buffer 测试全部通过； +- `rg "PTOViewToMemref|PTOMaterializeTileHandles"` 只剩历史文档； +- `rg "BindTileOp|DeclareTileMemRefOp|SlotMarkerOp"` 不再存在主路径依赖。 diff --git a/include/PTO/IR/PTOMultiBuffer.h b/include/PTO/IR/PTOMultiBuffer.h index 53c89f6c3c..1b99a15108 100644 --- a/include/PTO/IR/PTOMultiBuffer.h +++ b/include/PTO/IR/PTOMultiBuffer.h @@ -9,9 +9,10 @@ //===- PTOMultiBuffer.h - Shared constants for multi-buffer ----*- C++ -*-===// // // Shared constants for the multi-buffer expression scheme: -// - `kPtoMultiBufferAttrName` is the memref-level attribute name written by -// PTOViewToMemref when lowering an `alloc_multi_tile` op. PlanMemory and -// downstream passes read it to reserve N physical slots. +// - `kPtoMultiBufferAttrName` is the legacy memref-level slot-count +// attribute. +// - `kPtoMultiBufferAddrsAttrName` is the internal address list written +// directly on tile-native `pto.alloc_multi_tile` by PlanMemory. // - `kPtoMultiBufferMaxNum` is the upper bound on the slot count N. It is // kept in lock-step with the InsertSync `MAX_MULTI_BUFFER_NUM`. // @@ -29,6 +30,10 @@ namespace pto { inline constexpr llvm::StringLiteral kPtoMultiBufferAttrName = "pto.multi_buffer"; +/// Internal DenseI64ArrayAttr containing one byte address per physical slot. +inline constexpr llvm::StringLiteral kPtoMultiBufferAddrsAttrName = + "pto.multi_buffer_addrs"; + /// Upper bound for N; must stay consistent with `MAX_MULTI_BUFFER_NUM` in /// insert-sync. inline constexpr unsigned kPtoMultiBufferMaxNum = 16; diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index c722d5406a..b5839380ca 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -343,19 +343,18 @@ def AllocTileOp : PTO_Op<"alloc_tile", [AttrSizedOperandSegments]> { //===----------------------------------------------------------------------===// // // `pto.alloc_multi_tile` declares an N-slot logical tile buffer. The actual -// physical allocation of N slots is performed by PTOPlanMemory once the op -// has been lowered to `memref.alloc {pto.multi_buffer = N : i32}` by -// PTOViewToMemref. +// physical allocation of N slots is performed directly by PTOPlanMemory. +// The op remains tile-native until PTOResolveBufferSelect materializes an +// addressed `pto.alloc_tile` for each selected slot. // // `pto.multi_tile_get` is the *only* way to consume a `multi_tile_buf`: it // returns a regular `tile_buf` view onto the chosen slot. The slot index // may be a constant or any SSA `index` value; PTOAS does not synthesize // `iv mod N` automatically -- the user expression IS the slot selector. // -// `pto.slot_marker` is an internal op materialized by PTOViewToMemref to -// thread the slot SSA from `multi_tile_get` through the memref layer down -// to PlanMemory / sync analysis. It is not intended for direct frontend -// use. +// `pto.slot_marker` remains as an internal compatibility op for legacy +// memref-based multi-buffer IR. PTOViewToMemref does not create it for +// tile-native `pto.multi_tile_get`. def AllocMultiTileOp : PTO_Op<"alloc_multi_tile", [AttrSizedOperandSegments]> { let summary = "Allocate an N-slot multi-buffer tile"; @@ -415,9 +414,9 @@ def MultiTileGetOp : PTO_Op<"multi_tile_get", [ sync analysis treats this use as a slot-indexed dynamic access; ptoas does NOT rewrite the user expression into `iv mod N`. - This op is metadata-only (no data movement); the lowering merely - annotates the underlying memref view with the slot index for PlanMemory - / sync / EnableBufferSelect to consume. + This op is metadata-only (no data movement). PlanMemory preserves it, + sync analysis consumes its slot index directly, and + PTOResolveBufferSelect replaces it with an addressed `pto.alloc_tile`. }]; let arguments = (ins @@ -1556,11 +1555,10 @@ def SlotMarkerOp : PTO_Op<"slot_marker", [ ]> { let summary = "Tag a memref view as referring to one slot of a multi_tile_buf"; let description = [{ - Internal op materialized by `PTOViewToMemref` while lowering - `pto.multi_tile_get`. It carries the slot SSA index through the memref - layer so that PlanMemory, sync analysis (InsertSync / GraphSyncSolver), - and the buffer-select lowering pass can identify which physical slot - this memref reference touches. + Internal compatibility op for legacy memref-based multi-buffer IR. It + carries a slot SSA index so sync analysis (InsertSync / GraphSyncSolver) + and the buffer-select lowering pass can identify which physical slot a + memref reference touches. The op is metadata-only (no data movement, no extra address arithmetic); its result memref aliases the source memref byte-for-byte. Frontends do diff --git a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h index 71b3f1b10e..b9c439ea5a 100644 --- a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h +++ b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h @@ -70,6 +70,7 @@ class PTOIRTranslator { // --- 内存/Alias 分析 --- void UpdateKernelArgMemInfo(); LogicalResult UpdateAllocTileOpMemInfo(pto::AllocTileOp op); + LogicalResult UpdateAllocMultiTileOpMemInfo(pto::AllocMultiTileOp op); LogicalResult UpdateDeclareGlobalOpMemInfo(pto::DeclareGlobalOp op); LogicalResult UpdateDeclareTileMemRefOpMemInfo(pto::DeclareTileMemRefOp op); LogicalResult UpdatePointerCastOpMemInfo(pto::PointerCastOp op); @@ -81,6 +82,9 @@ class PTOIRTranslator { void UpdateMemrefSubViewAliasBufferInfo(memref::SubViewOp op); void UpdateTileSubViewAliasBufferInfo(pto::SubViewOp op); void UpdateSlotMarkerAliasBufferInfo(pto::SlotMarkerOp op); + void UpdateMultiTileGetAliasBufferInfo(pto::MultiTileGetOp op); + void UpdateSlotSelectedAliasBufferInfo(Value result, Value source, + Value slot); // --- 控制流处理 (SCF) --- void UpdateForOpInfo(scf::ForOp forOp); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 2c782033cb..210f9d8160 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -741,25 +741,20 @@ def PTOMaterializeTileHandles : Pass<"pto-materialize-tile-handles", "ModuleOp"> } def PTOResolveBufferSelect : Pass<"pto-resolve-buffer-select", "ModuleOp"> { - let summary = "Lower pto.slot_marker to single-slot pointer_cast (multi-buffer)"; + let summary = "Resolve multi-buffer slot selections to addressed tile handles"; let description = [{ - Consumes `pto.slot_marker %src[%k]` ops written by PTOViewToMemref to - thread multi-buffer slot selection through the memref layer. + Primarily consumes tile-native `pto.multi_tile_get %src[%k]` operations. + The source `pto.alloc_multi_tile` carries either planner-assigned physical + slot addresses or a level3 base address. Each selection is replaced by a + `pto.alloc_tile` carrying the selected address; constant slots select one + address directly, while dynamic slots use an N-way `arith.select` chain. - The op is replaced by a fresh single-address `pto.pointer_cast` that - refers to the chosen physical slot: + The pass also retains compatibility with legacy memref-based + `pto.slot_marker` IR, where the selection is materialized as a + single-address `pto.pointer_cast`. - - constant slot `k`: a single `pto.pointer_cast(addrK)` is emitted at - the use site, using slot k from the underlying multi-address - pointer_cast created by PTOPlanMemory / AllocToPointerCast. - - dynamic slot `%k`: per-slot single-address `pto.pointer_cast`s are - created and an N-way `arith.select` chain picks one according to the - user-supplied SSA. This pass does NOT rewrite `%k` into `iv mod N`; - the frontend expression is the slot selector. - - The original multi-address `pto.pointer_cast` is kept in IR as the - "alloc anchor" so future sync passes can recognize the multi-buffer - geometry. + This pass does not rewrite `%k` into `iv mod N`; the frontend expression + remains the slot selector. }]; let constructor = "mlir::pto::createPTOResolveBufferSelectPass()"; diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 1b66634049..08d45505c7 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3036,6 +3036,43 @@ LogicalResult AllocMultiTileOp::verify() { << kPtoMultiBufferMaxNum << "] (got " << count << ")"; } + if (Attribute rawAddrs = (*this)->getAttr(kPtoMultiBufferAddrsAttrName)) { + auto addrs = dyn_cast(rawAddrs); + if (!addrs) + return emitOpError() << "expects internal '" + << kPtoMultiBufferAddrsAttrName + << "' to be a dense i64 array"; + if (getAddr()) + return emitOpError() << "cannot carry both base 'addr' and internal '" + << kPtoMultiBufferAddrsAttrName << "'"; + if (addrs.size() != count) + return emitOpError() << "expects " << count << " planned slot addresses, got " + << addrs.size(); + + uint64_t elemBytes = getPTOStorageElemByteSize(slotTy.getElementType()); + uint64_t slotBytes = elemBytes; + for (int64_t dim : slotTy.getShape()) { + if (dim == ShapedType::kDynamic) + return emitOpError( + "planned multi-buffer addresses require a static slot shape"); + slotBytes *= static_cast(dim); + } + for (auto [lhsIdx, lhs] : llvm::enumerate(addrs.asArrayRef())) { + if (lhs < 0) + return emitOpError("planned slot addresses must be non-negative"); + uint64_t lhsBegin = static_cast(lhs); + uint64_t lhsEnd = lhsBegin + slotBytes; + for (size_t rhsIdx = lhsIdx + 1; + rhsIdx < static_cast(addrs.size()); ++rhsIdx) { + uint64_t rhsBegin = static_cast(addrs[rhsIdx]); + uint64_t rhsEnd = rhsBegin + slotBytes; + if (std::max(lhsBegin, rhsBegin) < std::min(lhsEnd, rhsEnd)) + return emitOpError() << "planned slots " << lhsIdx << " and " + << rhsIdx << " overlap"; + } + } + } + return success(); } diff --git a/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp b/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp index ba67ba3990..b0adfb765d 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp @@ -11,6 +11,7 @@ #include "PTO/Transforms/GraphSyncSolver/MemInfo.h" #include "PTO/IR/PTO.h" +#include "PTO/IR/PTOMultiBuffer.h" #include "PTO/IR/PTOTypeUtils.h" #include "../Utils.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -26,6 +27,19 @@ using namespace pto::syncsolver; namespace mlir::pto::syncsolver { static std::optional getBufferBitSize(Value value) { + if (auto multiType = dyn_cast(value.getType())) { + auto slotType = multiType.getSlotType(); + auto bitWidth = getPTOStorageElemBitWidth(slotType.getElementType()); + if (bitWidth == 0) + return ShapedType::kDynamic; + int64_t numElements = 1; + for (int64_t dim : slotType.getShape()) { + if (dim == ShapedType::kDynamic) + return ShapedType::kDynamic; + numElements *= dim; + } + return numElements * bitWidth; + } auto shaped = dyn_cast(value.getType()); if (!shaped || !shaped.hasStaticShape()) { return ShapedType::kDynamic; @@ -72,6 +86,55 @@ PointerLikeInfo getPointerLikeInfo(pto::PointerCastOp pointerCastOp) { return pointerLikeInfo; } +static std::optional getConstantI64(Value value) { + IntegerAttr attr; + if (!value || !matchPattern(value, m_Constant(&attr))) + return std::nullopt; + return attr.getValue().getSExtValue(); +} + +static PointerLikeInfo getPointerLikeInfo(pto::AllocMultiTileOp alloc) { + PointerLikeInfo info(alloc); + info.allocateSize = getBufferBitSize(alloc.getResult()); + auto slotType = alloc.getResult().getType().getSlotType(); + if (auto space = dyn_cast_or_null( + slotType.getMemorySpace())) + info.addressSpace = space.getAddressSpace(); + + if (auto planned = alloc->getAttrOfType( + pto::kPtoMultiBufferAddrsAttrName)) { + for (int64_t address : planned.asArrayRef()) + info.addresses.push_back(address * pto::kBitsToByte); + } else if (auto base = getConstantI64(alloc.getAddr())) { + int64_t slotBits = info.allocateSize.value_or(ShapedType::kDynamic); + for (uint32_t slot = 0; slot < alloc.getResult().getType().getCount(); ++slot) + info.addresses.push_back(*base * pto::kBitsToByte + slot * slotBits); + } + if (auto loop = alloc->getParentOfType()) + info.parentLoop = loop; + return info; +} + +static MemInfo getMemInfoForMultiTileGet(pto::MultiTileGetOp get) { + auto alloc = get.getSource().getDefiningOp(); + if (!alloc) + return MemInfo(get.getResult(), isWorkSpaceFuncArgument(get.getResult())); + + PointerLikeInfo info = getPointerLikeInfo(alloc); + IntegerAttr slotAttr; + if (matchPattern(get.getSlot(), m_Constant(&slotAttr)) && + info.addresses.size() > 1) { + int64_t slot = slotAttr.getValue().getSExtValue(); + if (slot >= 0 && slot < static_cast(info.addresses.size())) { + int64_t address = info.addresses[static_cast(slot)]; + info.addresses.assign(1, address); + } + } + if (auto loop = get->getParentOfType()) + info.parentLoop = loop; + return MemInfo(get.getResult(), info); +} + // Walk back through metadata-only view ops (`pto.bind_tile`) to the // nearest `pto.pointer_cast`. Used to anchor slot_marker MemInfo on its // underlying multi-address alloc cast. @@ -143,6 +206,10 @@ MemInfo getMemInfo(Value val) { if (auto slotMarker = llvm::dyn_cast(defOp)) { return getMemInfoForSlotMarker(slotMarker); } + if (auto allocMulti = llvm::dyn_cast(defOp)) + return MemInfo(val, getPointerLikeInfo(allocMulti)); + if (auto multiGet = llvm::dyn_cast(defOp)) + return getMemInfoForMultiTileGet(multiGet); } return MemInfo(val, isWorkSpaceFuncArgument(val)); } diff --git a/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp b/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp index 38e8c602e0..9e9758e541 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp @@ -81,7 +81,7 @@ llvm::SmallVector IRTranslator::tracebackMemValsStep(Value val) { // `getOperationAliasInfo` path below would treat slot_marker as a // transparent view and let the trace fall through to the underlying // multi-address `pto.pointer_cast`, dropping the slot. - if (isa(defOp)) { + if (isa(defOp)) { return out; } @@ -131,8 +131,9 @@ llvm::SmallVector IRTranslator::tracebackMemVals(Value val) { // stop, `getOperationAliasInfo` would let the walk slip past slot_marker // and reach the underlying multi-address `pto.pointer_cast`, dropping // the slot index. - if (isa(defOp)) { + if (isa(defOp)) { leaves.insert(result); continue; } diff --git a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp index 7d19f50db0..e699d19a3b 100644 --- a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp +++ b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp @@ -12,6 +12,7 @@ // See LICENSE in the root of the software repository for the full text of the License. #include "PTO/Transforms/InsertSync/PTOIRTranslator.h" +#include "PTO/IR/PTOMultiBuffer.h" #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/InsertSync/SyncMacroModel.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -363,6 +364,11 @@ void PTOIRTranslator::RecursionIR(Region *region) { return WalkResult::interrupt(); } } + else if (auto allocMultiOp = dyn_cast(op)) { + if (failed(UpdateAllocMultiTileOpMemInfo(allocMultiOp))) { + return WalkResult::interrupt(); + } + } // 支持标准 memref.alloc else if (auto memAllocOp = dyn_cast(op)) { if (failed(UpdateMemrefAllocOpMemInfo(memAllocOp))) { @@ -402,6 +408,9 @@ void PTOIRTranslator::RecursionIR(Region *region) { else if (auto slotMarker = dyn_cast(op)) { UpdateSlotMarkerAliasBufferInfo(slotMarker); } + else if (auto multiGet = dyn_cast(op)) { + UpdateMultiTileGetAliasBufferInfo(multiGet); + } else if (auto castOp = dyn_cast(op)) { UpdateAliasBufferInfo(castOp.getResult(), castOp.getSource()); } @@ -517,6 +526,55 @@ LogicalResult PTOIRTranslator::UpdateAllocTileOpMemInfo(pto::AllocTileOp op) { return success(); } +LogicalResult +PTOIRTranslator::UpdateAllocMultiTileOpMemInfo(pto::AllocMultiTileOp op) { + Value result = op.getResult(); + auto multiType = op.getResult().getType(); + pto::TileBufType slotType = multiType.getSlotType(); + + uint64_t slotBytes = pto::getPTOStorageElemByteSize(slotType.getElementType()); + if (slotBytes == 0) + return failure(); + for (int64_t dim : slotType.getShape()) { + if (dim == ShapedType::kDynamic) + return op.emitError("requires a static slot shape for sync analysis"); + slotBytes *= static_cast(dim); + } + + pto::AddressSpace space = pto::AddressSpace::MAT; + if (auto attr = dyn_cast_or_null( + slotType.getMemorySpace())) + space = attr.getAddressSpace(); + + SmallVector addresses; + bool hasKnownAddresses = false; + if (auto planned = op->getAttrOfType( + pto::kPtoMultiBufferAddrsAttrName)) { + if (planned.size() != multiType.getCount()) + return op.emitError("planned address count does not match slot count"); + for (int64_t address : planned.asArrayRef()) + addresses.push_back(static_cast(address)); + hasKnownAddresses = true; + } else if (Value base = op.getAddr()) { + if (std::optional knownBase = getKnownPhysicalAddress(base)) { + for (uint32_t slot = 0; slot < multiType.getCount(); ++slot) + addresses.push_back(*knownBase + slot * slotBytes); + hasKnownAddresses = true; + } + } + + if (addresses.empty()) { + return op.emitError( + "requires planner-assigned slot addresses or a constant level3 base"); + } + + auto info = std::make_unique( + result, result, space, std::move(addresses), slotBytes, + hasKnownAddresses && isLocalAddressSpace(space)); + buffer2MemInfoMap_[result].emplace_back(info->clone()); + return success(); +} + LogicalResult PTOIRTranslator::UpdatePointerCastOpMemInfo(pto::PointerCastOp op) { Value res = op.getResult(); auto memRefType = dyn_cast(res.getType()); @@ -986,14 +1044,24 @@ void PTOIRTranslator::UpdateConservativeAliasBufferInfo(Value result, } void PTOIRTranslator::UpdateSlotMarkerAliasBufferInfo(pto::SlotMarkerOp op) { - Value result = op.getResult(); - Value source = op.getSource(); + UpdateSlotSelectedAliasBufferInfo(op.getResult(), op.getSource(), + op.getSlot()); +} + +void PTOIRTranslator::UpdateMultiTileGetAliasBufferInfo( + pto::MultiTileGetOp op) { + UpdateSlotSelectedAliasBufferInfo(op.getResult(), op.getSource(), + op.getSlot()); +} + +void PTOIRTranslator::UpdateSlotSelectedAliasBufferInfo(Value result, + Value source, + Value slot) { if (!result || !source) return; if (!buffer2MemInfoMap_.contains(source)) return; - Value slot = op.getSlot(); IntegerAttr constAttr; bool isConstSlot = matchPattern(slot, m_Constant(&constAttr)); int64_t constSlotIdx = isConstSlot ? constAttr.getValue().getSExtValue() : -1; diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index f3ad1810ef..2df6dc6a4d 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -1520,6 +1520,20 @@ struct PTOMaterializeTileHandlesPass ModuleOp module = getOperation(); MLIRContext *ctx = module.getContext(); + bool hasUnresolvedMultiTile = false; + module.walk([&](Operation *op) { + if (!isa(op)) + return; + op->emitError( + "must be resolved by pto-resolve-buffer-select before " + "pto-materialize-tile-handles"); + hasUnresolvedMultiTile = true; + }); + if (hasUnresolvedMultiTile) { + signalPassFailure(); + return; + } + OpBuilder builder(ctx); DenseMap tileHandles; DenseSet mustMaterialize; diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index 81656b7859..de8b9f84dd 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -508,6 +508,19 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { } UpdateOpBufferInfo(op, op->getResults()); return WalkResult::advance(); + } else if (isLocalMemPlan() && dyn_cast(op)) { + auto allocMultiOp = cast(op); + if (allocMultiOp.getAddr()) + return WalkResult::advance(); + auto memorySpaceAttr = GetBufferSpaceAttr(allocMultiOp.getResult()); + if (!isLocalBuffer(memorySpaceAttr)) { + allocMultiOp.emitError("Alloc multi tile buffer not at local space"); + return WalkResult::interrupt(); + } + uint32_t count = allocMultiOp.getResult().getType().getCount(); + buffer2MultiNum[allocMultiOp.getResult()] = count; + UpdateOpBufferInfo(op, op->getResults()); + return WalkResult::advance(); } else if (isLocalMemPlan() && dyn_cast(op)) { auto allocOp = cast(op); if (failed(CheckLocalBufferAllocOp(op))) { @@ -1007,6 +1020,9 @@ BufferInfo MemLivenessAnalysis::GetBufferInfo(Operation *op, Value operand, if (auto tileType = dyn_cast(operand.getType())) { shape = tileType.getShape(); elementType = tileType.getElementType(); + } else if (auto multiType = dyn_cast(operand.getType())) { + shape = multiType.getSlotType().getShape(); + elementType = multiType.getSlotType().getElementType(); } else { Value traceValue = tracebackMemRef(operand); auto memRefType = cast(traceValue.getType()); @@ -2568,6 +2584,42 @@ class LegacyAllocTileOpAddPlannedAddressPattern DenseMap> buffer2Offsets; }; +class LegacyAllocMultiTileOpAddPlannedAddressesPattern + : public OpRewritePattern { +public: + explicit LegacyAllocMultiTileOpAddPlannedAddressesPattern( + MLIRContext *context, + DenseMap> buffer2Offsets) + : OpRewritePattern(context), + buffer2Offsets(std::move(buffer2Offsets)) {} + + LogicalResult matchAndRewrite(pto::AllocMultiTileOp op, + PatternRewriter &rewriter) const override { + if (op.getAddr() || op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) + return failure(); + auto it = buffer2Offsets.find(op.getResult()); + if (it == buffer2Offsets.end() || it->second.empty()) + return failure(); + if (it->second.size() != op.getResult().getType().getCount()) { + return rewriter.notifyMatchFailure( + op, "planned address count does not match multi_tile_buf count"); + } + + SmallVector addrs; + addrs.reserve(it->second.size()); + for (uint64_t offset : it->second) + addrs.push_back(static_cast(offset)); + rewriter.modifyOpInPlace(op, [&] { + op->setAttr(pto::kPtoMultiBufferAddrsAttrName, + rewriter.getDenseI64ArrayAttr(addrs)); + }); + return success(); + } + +private: + DenseMap> buffer2Offsets; +}; + static FailureOr parseLegacyMemPlanMode(func::FuncOp func, llvm::StringRef memMode) { if (memMode.equals_insensitive("local") || @@ -2597,6 +2649,8 @@ struct PlanMemoryPass : public mlir::pto::impl::PlanMemoryBase { buffer2Offsets); patterns.add( patterns.getContext(), buffer2Offsets); + patterns.add( + patterns.getContext(), buffer2Offsets); } } }; diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 04a555a6fb..2fa4d9882d 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -90,6 +90,9 @@ static std::optional getBufferAddressSpace(Type type) { return std::nullopt; } + if (auto multiType = dyn_cast(type)) + return getBufferAddressSpace(multiType.getSlotType()); + if (auto memRefType = dyn_cast(type)) { if (auto attr = dyn_cast_or_null(memRefType.getMemorySpace())) @@ -186,6 +189,9 @@ static FailureOr computeStaticBufferBytes(Value value) { if (auto tileType = dyn_cast(value.getType())) { shape = tileType.getShape(); elementType = tileType.getElementType(); + } else if (auto multiType = dyn_cast(value.getType())) { + shape = multiType.getSlotType().getShape(); + elementType = multiType.getSlotType().getElementType(); } else if (auto memRefType = dyn_cast(value.getType())) { shape = memRefType.getShape(); elementType = memRefType.getElementType(); @@ -360,8 +366,10 @@ struct PlannerAnalysis { } uint64_t slotCount = 1; - if (auto attr = - defOp->getAttrOfType(pto::kPtoMultiBufferAttrName)) + if (auto multiType = dyn_cast(value.getType())) + slotCount = multiType.getCount(); + else if (auto attr = + defOp->getAttrOfType(pto::kPtoMultiBufferAttrName)) slotCount = attr.getValue().getZExtValue(); MemSpec spec = getMemSpec(getTargetArch(func), *space); uint64_t slotBytes = alignUp(*bytesOr, spec.alignmentBytes); @@ -653,6 +661,17 @@ struct PlannerAnalysis { roots[found->second].freeIndex = index; } } + } else if (auto allocMulti = dyn_cast(op)) { + if (!allocMulti.getAddr()) { + addRoot(allocMulti.getResult(), op); + if (failed) + return; + auto found = rootIndexByValue.find(allocMulti.getResult()); + if (found != rootIndexByValue.end()) { + roots[found->second].allocIndex = index; + roots[found->second].freeIndex = index; + } + } } else if (auto alloc = dyn_cast(op)) { addRoot(alloc.getResult(), op); if (failed) @@ -671,6 +690,10 @@ struct PlannerAnalysis { setRoots(slotMarker.getResult(), getRoots(slotMarker.getSource())); propagateSplitTpopDerived(slotMarker.getResult(), ValueRange{slotMarker.getSource()}); + } else if (auto multiGet = dyn_cast(op)) { + setRoots(multiGet.getResult(), getRoots(multiGet.getSource())); + propagateSplitTpopDerived(multiGet.getResult(), + ValueRange{multiGet.getSource()}); } else if (auto select = dyn_cast(op)) { setRoots(select.getResult(), unionRoots(getRoots(select.getTrueValue()), @@ -970,6 +993,43 @@ class AllocTileOpAddPlannedAddressPattern DenseMap> buffer2Offsets; }; +class AllocMultiTileOpAddPlannedAddressesPattern + : public OpRewritePattern { +public: + explicit AllocMultiTileOpAddPlannedAddressesPattern( + MLIRContext *context, + DenseMap> buffer2Offsets) + : OpRewritePattern(context), + buffer2Offsets(std::move(buffer2Offsets)) {} + + LogicalResult matchAndRewrite(pto::AllocMultiTileOp op, + PatternRewriter &rewriter) const override { + if (op.getAddr() || op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) + return failure(); + + auto it = buffer2Offsets.find(op.getResult()); + if (it == buffer2Offsets.end() || it->second.empty()) + return failure(); + if (it->second.size() != op.getResult().getType().getCount()) { + return rewriter.notifyMatchFailure( + op, "planned address count does not match multi_tile_buf count"); + } + + SmallVector addrs; + addrs.reserve(it->second.size()); + for (uint64_t offset : it->second) + addrs.push_back(static_cast(offset)); + rewriter.modifyOpInPlace(op, [&] { + op->setAttr(pto::kPtoMultiBufferAddrsAttrName, + rewriter.getDenseI64ArrayAttr(addrs)); + }); + return success(); + } + +private: + DenseMap> buffer2Offsets; +}; + static LogicalResult materializePlannedOffsets( func::FuncOp func, DenseMap> buffer2Offsets) { RewritePatternSet patterns(func.getContext()); @@ -977,6 +1037,8 @@ static LogicalResult materializePlannedOffsets( buffer2Offsets); patterns.add(patterns.getContext(), buffer2Offsets); + patterns.add( + patterns.getContext(), buffer2Offsets); if (mlir::failed(applyPatternsGreedily(func, std::move(patterns)))) return failure(); return success(); diff --git a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp index 858593e6d3..8bff4e5e37 100644 --- a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp +++ b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp @@ -30,6 +30,7 @@ #include "PTO/IR/PTO.h" #include "PTO/IR/PTOMultiBuffer.h" +#include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/Passes.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -54,6 +55,105 @@ using namespace mlir; namespace { +static FailureOr getStaticSlotBytes(pto::TileBufType slotType) { + uint64_t elemBytes = pto::getPTOStorageElemByteSize(slotType.getElementType()); + if (elemBytes == 0) + return failure(); + uint64_t bytes = elemBytes; + for (int64_t dim : slotType.getShape()) { + if (dim == ShapedType::kDynamic) + return failure(); + bytes *= static_cast(dim); + } + return bytes; +} + +static LogicalResult getMultiTileAddresses(pto::AllocMultiTileOp alloc, + IRRewriter &rewriter, + SmallVectorImpl &addrs) { + uint32_t count = alloc.getResult().getType().getCount(); + if (auto planned = alloc->getAttrOfType( + pto::kPtoMultiBufferAddrsAttrName)) { + if (planned.size() != count) + return alloc.emitError("planned address count does not match slot count"); + for (int64_t address : planned.asArrayRef()) + addrs.push_back(rewriter.create( + alloc.getLoc(), address, 64)); + return success(); + } + + Value base = alloc.getAddr(); + if (!base) + return alloc.emitError( + "has neither a level3 base address nor planner-assigned slot addresses"); + auto slotBytes = getStaticSlotBytes(alloc.getResult().getType().getSlotType()); + if (failed(slotBytes)) + return alloc.emitError( + "requires a static slot shape and known element byte size"); + + addrs.push_back(base); + for (uint32_t slot = 1; slot < count; ++slot) { + Value offset = rewriter.create( + alloc.getLoc(), static_cast(slot * *slotBytes), 64); + addrs.push_back( + rewriter.create(alloc.getLoc(), base, offset)); + } + return success(); +} + +static LogicalResult resolveTileNativeMultiGets(ModuleOp module, + MLIRContext *ctx) { + SmallVector gets; + module.walk([&](pto::MultiTileGetOp op) { gets.push_back(op); }); + + for (pto::MultiTileGetOp op : gets) { + auto alloc = op.getSource().getDefiningOp(); + if (!alloc) + return op.emitError( + "currently requires a direct pto.alloc_multi_tile source"); + + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(op); + SmallVector addrs; + if (failed(getMultiTileAddresses(alloc, rewriter, addrs))) + return failure(); + + Value selectedAddr; + IntegerAttr constSlotAttr; + if (matchPattern(op.getSlot(), m_Constant(&constSlotAttr))) { + int64_t slot = constSlotAttr.getValue().getSExtValue(); + if (slot < 0 || slot >= static_cast(addrs.size())) + return op.emitError("constant slot is outside planned address range"); + selectedAddr = addrs[static_cast(slot)]; + } else { + selectedAddr = addrs.front(); + for (uint32_t slot = 1; slot < addrs.size(); ++slot) { + Value slotValue = rewriter.create(op.getLoc(), slot); + Value matches = rewriter.create( + op.getLoc(), arith::CmpIPredicate::eq, op.getSlot(), slotValue); + selectedAddr = rewriter.create( + op.getLoc(), matches, addrs[slot], selectedAddr); + } + } + + auto slotHandle = rewriter.create( + op.getLoc(), op.getResult().getType(), selectedAddr, + alloc.getValidRow() ? alloc.getValidRow() : Value(), + alloc.getValidCol() ? alloc.getValidCol() : Value()); + rewriter.replaceOp(op, slotHandle.getResult()); + } + + SmallVector allocs; + module.walk([&](pto::AllocMultiTileOp op) { allocs.push_back(op); }); + for (pto::AllocMultiTileOp alloc : allocs) { + if (!alloc.getResult().use_empty()) + return alloc.emitError( + "has unsupported uses after resolving pto.multi_tile_get"); + alloc.erase(); + } + return success(); +} + /// Walk back through pure metadata ops (`pto.bind_tile`, `pto.slot_marker`) /// to find the root multi-address `pto.pointer_cast` that this view ties /// to. Returns nullptr if the chain does not terminate on a @@ -113,6 +213,11 @@ struct PTOResolveBufferSelectPass ModuleOp mod = getOperation(); MLIRContext *ctx = &getContext(); + if (failed(resolveTileNativeMultiGets(mod, ctx))) { + signalPassFailure(); + return; + } + SmallVector markers; mod.walk([&](pto::SlotMarkerOp op) { markers.push_back(op); }); if (markers.empty()) diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index cb865879fb..1da3c177aa 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -13,7 +13,6 @@ // metadata through binding ops and SSA backtracking. #include "PTO/IR/PTO.h" -#include "PTO/IR/PTOMultiBuffer.h" #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/Passes.h" @@ -65,13 +64,14 @@ static void markForceDynamicValidShape(Operation *op, bool force, static Type convertPTOTypeToMemRef(Type t); -static bool hasTileNativeAllocTileRoot(func::FuncOp func) { +static bool hasTileNativeAllocationRoot(func::FuncOp func) { bool found = false; - auto result = func.walk([&](pto::AllocTileOp op) { - if (!isa(op.getResult().getType())) - return WalkResult::advance(); - found = true; - return WalkResult::interrupt(); + auto result = func.walk([&](Operation *op) { + if (isa(op)) { + found = true; + return WalkResult::interrupt(); + } + return WalkResult::advance(); }); (void)result; return found; @@ -1746,240 +1746,11 @@ struct PTOViewToMemrefPass // ------------------------------------------------------------------ // ------------------------------------------------------------------ - // Stage 0.6: lower pto.alloc_multi_tile / pto.multi_tile_get - // - // alloc_multi_tile produces a `multi_tile_buf`. We lower - // it into: - // %a = memref.alloc() {pto.multi_buffer = N : i32} : memref - // %m = pto.bind_tile %a, ... : memref -> memref - // and replace all uses of the multi_tile_buf SSA with %m. The N-way - // physical fan-out lives on the `pto.multi_buffer` attr and is - // materialized later by the fallback address resolver. - // - // multi_tile_get consumes that memref and wraps it in - // %slot = pto.slot_marker %m[%k] : memref -> memref - // %t = pto.bind_tile %slot, ... : memref -> memref - // The slot_marker carries the user-supplied slot SSA forward so that - // sync analysis and buffer-select lowering can identify which physical - // slot this use refers to. - // - // Ordering: alloc_multi_tile must be lowered before multi_tile_get so - // that the get's source SSA has already been rewired to a memref. + // Stage 0.6: keep pto.alloc_multi_tile / pto.multi_tile_get tile-native. + // PlanMemory writes the physical slot addresses on alloc_multi_tile; + // sync consumes the slot selector directly from multi_tile_get, and + // PTOResolveBufferSelect materializes the selected alloc_tile handle. // ------------------------------------------------------------------ - SmallVector allocMultiTiles; - func.walk( - [&](mlir::pto::AllocMultiTileOp op) { allocMultiTiles.push_back(op); }); - - for (auto op : allocMultiTiles) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - auto mtbTy = op.getResult().getType(); - if (!mtbTy) - continue; - auto tbTy = mtbTy.getSlotType(); - if (!tbTy) - continue; - - SmallVector shape(tbTy.getShape().begin(), - tbTy.getShape().end()); - Type elemTy = tbTy.getElementType(); - - // Stride / layout reuse the same logic as alloc_tile. - SmallVector strides; - TileLayoutInfo info; - if (computeTileLayoutInfo(tbTy.getConfigAttr(), elemTy, shape, info)) { - strides = {info.rowStride, info.colStride}; - } else { - strides.resize(shape.size()); - int64_t s = 1; - for (int i = (int)shape.size() - 1; i >= 0; --i) { - strides[i] = s; - if (shape[i] != ShapedType::kDynamic) - s *= shape[i]; - } - } - - auto targetLayout = - StridedLayoutAttr::get(ctx, ShapedType::kDynamic, strides); - auto targetType = - MemRefType::get(shape, elemTy, targetLayout, tbTy.getMemorySpace()); - - Value vRow = op.getValidRow(); - Value vCol = op.getValidCol(); - ArrayRef validShape = tbTy.getValidShape(); - if (!tbTy.hasDynamicValid()) { - if (validShape.size() >= 1 && validShape[0] >= 0) { - vRow = rewriter - .create( - loc, rewriter.getIndexType(), - rewriter.getIndexAttr(validShape[0])) - .getResult(); - } - if (validShape.size() >= 2 && validShape[1] >= 0) { - vCol = rewriter - .create( - loc, rewriter.getIndexType(), - rewriter.getIndexAttr(validShape[1])) - .getResult(); - } - } - - auto configAttr = tbTy.getConfigAttr(); - if (!configAttr) - configAttr = pto::TileBufConfigAttr::getDefault(ctx); - - // Level3 (caller-owned memory): an explicit base address is given. - // Lay the N slots out contiguously as - // [addr, addr + slotBytes, ..., addr + (N-1)*slotBytes] and emit the - // multi-address `pto.pointer_cast` alloc anchor directly (addrs are - // kept in slot order so PTOResolveBufferSelect can pick addrs[slot]). - // Sync and buffer-select then run unchanged. - if (Value addr = op.getAddr()) { - uint32_t n = mtbTy.getCount(); - int64_t elemBytes = getElemBytes(elemTy); - int64_t numElems = 1; - bool staticShape = true; - for (int64_t d : shape) { - if (d == ShapedType::kDynamic) { - staticShape = false; - break; - } - numElems *= d; - } - if (!staticShape || elemBytes <= 0) { - op.emitError("pto.alloc_multi_tile with explicit addr requires a " - "static slot shape and known element byte size"); - signalPassFailure(); - return; - } - int64_t slotBytes = numElems * elemBytes; - - SmallVector slotAddrs; - slotAddrs.reserve(n); - slotAddrs.push_back(addr); - for (uint32_t slot = 1; slot < n; ++slot) { - Value off = rewriter.create( - loc, rewriter.getI64Type(), - rewriter.getI64IntegerAttr(static_cast(slot) * - slotBytes)); - slotAddrs.push_back( - rewriter.create(loc, addr, off).getResult()); - } - - auto pc = rewriter.create( - loc, targetType, slotAddrs, vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(pc, tbTy.hasDynamicValid(), ctx); - - auto bindOp = rewriter.create( - loc, targetType, pc.getResult(), vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - - rewriter.replaceOp(op, bindOp.getResult()); - continue; - } - - // memref.alloc with N-slot annotation. The actual N-way address - // expansion happens in PTOPlanMemory. - auto allocLayout = StridedLayoutAttr::get(ctx, 0, strides); - auto allocType = - MemRefType::get(shape, elemTy, allocLayout, tbTy.getMemorySpace()); - auto allocOp = rewriter.create(loc, allocType); - auto i32Ty = IntegerType::get(ctx, 32); - allocOp->setAttr( - mlir::pto::kPtoMultiBufferAttrName, - IntegerAttr::get(i32Ty, static_cast(mtbTy.getCount()))); - - auto bindOp = rewriter.create( - loc, targetType, allocOp.getResult(), vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - - rewriter.replaceOp(op, bindOp.getResult()); - } - - SmallVector multiTileGets; - func.walk( - [&](mlir::pto::MultiTileGetOp op) { multiTileGets.push_back(op); }); - - for (auto op : multiTileGets) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - // After alloc_multi_tile has been replaced, the source SSA value - // is a memref (no longer multi_tile_buf). We accept the raw Value. - Value srcMem = op->getOperand(0); - if (!llvm::isa(srcMem.getType())) { - op.emitError( - "multi_tile_get expects its source to have been lowered to " - "memref by alloc_multi_tile rewriting; got ") - << srcMem.getType(); - signalPassFailure(); - return; - } - auto srcMemTy = llvm::cast(srcMem.getType()); - - // Recover per-slot tile_buf info from the get op's result type - // (which still describes the per-slot tile shape, valid dims, - // config, etc.). - auto tbTy = llvm::dyn_cast(op.getResult().getType()); - if (!tbTy) { - op.emitError("multi_tile_get result must be `!pto.tile_buf<...>`"); - signalPassFailure(); - return; - } - - // The slot view aliases the source memref byte-for-byte. - auto slotMarker = rewriter.create( - loc, srcMemTy, srcMem, op.getSlot()); - - // Recover valid_row / valid_col from the per-slot tile type so - // downstream BindTile carries the same metadata as the alloc. - Value vRow; - Value vCol; - ArrayRef validShape = tbTy.getValidShape(); - if (!tbTy.hasDynamicValid()) { - if (validShape.size() >= 1 && validShape[0] >= 0) { - vRow = rewriter - .create( - loc, rewriter.getIndexType(), - rewriter.getIndexAttr(validShape[0])) - .getResult(); - } - if (validShape.size() >= 2 && validShape[1] >= 0) { - vCol = rewriter - .create( - loc, rewriter.getIndexType(), - rewriter.getIndexAttr(validShape[1])) - .getResult(); - } - } - - auto configAttr = tbTy.getConfigAttr(); - if (!configAttr) - configAttr = pto::TileBufConfigAttr::getDefault(ctx); - - // BindTile result type uses the same layout signature as in the - // alloc lowering, so subsequent subview / DMA ops see a consistent - // memref view. - SmallVector strides = buildTileMemRefStrides(tbTy); - auto targetLayout = - StridedLayoutAttr::get(ctx, ShapedType::kDynamic, strides); - auto targetType = - MemRefType::get(tbTy.getShape(), tbTy.getElementType(), targetLayout, - tbTy.getMemorySpace()); - - auto bindOp = rewriter.create( - loc, targetType, slotMarker.getResult(), vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - - rewriter.replaceOp(op, bindOp.getResult()); - } // ------------------------------------------------------------------ // Stage 0.75: lower pto.declare_tile -> pto.declare_tile_memref + @@ -2362,7 +2133,7 @@ struct PTOViewToMemrefPass // Stage 3: Rewrite Compute Ops // [关键] 全面使用 op->getOperand(i) 避免 Typed Accessor Crash // ------------------------------------------------------------------ - if (!hasTileNativeAllocTileRoot(func)) { + if (!hasTileNativeAllocationRoot(func)) { // --- TLoadOp [Src, Dst] --- DefaultInlineVector loads; diff --git a/lib/PTO/Transforms/SlotAffineAnalysis.cpp b/lib/PTO/Transforms/SlotAffineAnalysis.cpp index ab37697c6f..cf633b1795 100644 --- a/lib/PTO/Transforms/SlotAffineAnalysis.cpp +++ b/lib/PTO/Transforms/SlotAffineAnalysis.cpp @@ -30,6 +30,8 @@ Value findSlotMarkerExpr(Value v) { return {}; if (auto sm = dyn_cast(op)) return sm.getSlot(); + if (auto get = dyn_cast(op)) + return get.getSlot(); if (auto bind = dyn_cast(op)) { v = bind.getSource(); continue; diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index 2b34417c10..4978986158 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -137,6 +137,12 @@ std::optional GetBufferSpaceAttr(Value operand) { return memorySpaceAttr; return std::nullopt; } + if (auto multiTy = dyn_cast(operand.getType())) { + if (auto memorySpaceAttr = dyn_cast_or_null( + multiTy.getSlotType().getMemorySpace())) + return memorySpaceAttr; + return std::nullopt; + } if (!llvm::isa(operand.getType())) { return std::nullopt; @@ -168,6 +174,8 @@ std::optional> getOperationAliasInfo(Operation *op) { // physical slot of a multi-buffer alloc. From an alias-walking // standpoint it behaves like any other view-like op. return std::make_pair(slotMarkerOp.getResult(), slotMarkerOp.getSource()); + } else if (auto multiGetOp = dyn_cast(op)) { + return std::make_pair(multiGetOp.getResult(), multiGetOp.getSource()); } else if (auto extSliceOp = dyn_cast(op)) { return std::make_pair(extSliceOp.getResult(), extSliceOp.getSource()); } else if (auto collapseShapeOp = dyn_cast(op)) { diff --git a/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto b/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto index 76b9171330..7773a855d9 100644 --- a/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto +++ b/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto @@ -7,6 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s // Sanity for N == 3 to confirm `ExpandMultiBufferStorageEntry` produces 3 // physical slots and `UpdateBuffer2Offsets` emits them in slot order. @@ -39,4 +40,6 @@ module { // CHECK: IR Dump After PlanMemory // CHECK-LABEL: func.func @n3_three_slots -// CHECK: pto.pointer_cast({{[^,]*}}, {{[^,]*}}, {{[^,]*}}) +// CHECK: pto.alloc_multi_tile {pto.multi_buffer_addrs = array} +// CHECK: pto.multi_tile_get +// CHECK-NOT: pto.pointer_cast diff --git a/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto b/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto index c2a1b19d2d..93b2490ee7 100644 --- a/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto +++ b/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto @@ -45,21 +45,24 @@ module { // CHECK: IR Dump After PTOViewToMemref // CHECK-LABEL: func.func @const_preload_dyn_loop_select -// CHECK: memref.alloc() {pto.multi_buffer = 2 : i32} -// CHECK: pto.slot_marker %{{.*}}[%c0 +// CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile +// CHECK: pto.multi_tile_get %[[MULTI]][%c0 // CHECK: scf.for // CHECK: %[[IDX:.*]] = arith.remui %arg{{.*}}, %c2 -// CHECK: pto.slot_marker %{{.*}}[%[[IDX]] +// CHECK: pto.multi_tile_get %[[MULTI]][%[[IDX]] +// CHECK-NOT: memref.alloc +// CHECK-NOT: pto.slot_marker -// After ResolveBufferSelect, the const preload slot is a direct single-address -// pointer_cast, while the loop slot is chosen with arith.select. +// After ResolveBufferSelect, the const preload slot becomes a directly +// addressed alloc_tile, while the loop slot is chosen with arith.select. // SELECT: IR Dump After PTOResolveBufferSelect // SELECT-LABEL: func.func @const_preload_dyn_loop_select -// SELECT: %[[ANCHOR:.*]] = pto.pointer_cast(%[[ADDR0:.*]], %[[ADDR1:.*]]) -// SELECT-NOT: pto.slot_marker -// SELECT: pto.pointer_cast(%[[ADDR0]]) +// SELECT-NOT: pto.alloc_multi_tile +// SELECT-NOT: pto.multi_tile_get +// SELECT: pto.alloc_tile addr = %{{.*}} // SELECT: scf.for // SELECT: arith.select +// SELECT: pto.alloc_tile addr = %{{.*}} // MaterializeTileHandles must carry the dynamic selected buffer address into // alloc_tile, otherwise EmitC creates a tile without TASSIGN in the loop. @@ -69,8 +72,10 @@ module { // MAT-DAG: %[[ADDR1:[A-Za-z0-9_]+]] = arith.constant 512 : i64 // MAT: pto.alloc_tile addr = %[[ADDR0]] // MAT: scf.for +// MAT-DAG: %[[LOOP_ADDR0:[A-Za-z0-9_]+]] = arith.constant 0 : i64 +// MAT-DAG: %[[LOOP_ADDR1:[A-Za-z0-9_]+]] = arith.constant 512 : i64 // MAT: %[[COND:[A-Za-z0-9_]+]] = arith.cmpi eq -// MAT: %[[DYN_ADDR:[A-Za-z0-9_]+]] = arith.select %[[COND]], %[[ADDR1]], %[[ADDR0]] : i64 +// MAT: %[[DYN_ADDR:[A-Za-z0-9_]+]] = arith.select %[[COND]], %[[LOOP_ADDR1]], %[[LOOP_ADDR0]] : i64 // MAT: pto.alloc_tile addr = %[[DYN_ADDR]] // EMITC: AICORE void const_preload_dyn_loop_select diff --git a/test/lit/pto/multi_tile_get_const_slot_lowering.pto b/test/lit/pto/multi_tile_get_const_slot_lowering.pto index 65c6c3430f..52b5ca5f01 100644 --- a/test/lit/pto/multi_tile_get_const_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_const_slot_lowering.pto @@ -9,15 +9,13 @@ // RUN: ptoas --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN // RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT -// Verifies the full multi-buffer pipeline for constant slots: -// 1. PTOViewToMemref lowers `alloc_multi_tile` to `memref.alloc` with the -// `pto.multi_buffer = N : i32` attribute, and `multi_tile_get [%k]` to -// `pto.slot_marker` carrying the constant slot SSA. -// 2. PTOPlanMemory assigns planned addresses and emits a multi- -// address `pto.pointer_cast(addr0, addr1)`. -// 3. PTOResolveBufferSelect lowers each `slot_marker` to a single- -// address `pto.pointer_cast(addrK)` -- one per slot use. +// Verifies the tile-native multi-buffer pipeline for constant slots: +// 1. PTOViewToMemref preserves alloc_multi_tile and multi_tile_get. +// 2. PTOPlanMemory records the physical slot addresses on alloc_multi_tile. +// 3. PTOResolveBufferSelect replaces each selection with alloc_tile(addrK). module { func.func @const_slot_two_buffers( @@ -46,21 +44,23 @@ module { // CHECK: IR Dump After PTOViewToMemref // CHECK-LABEL: func.func @const_slot_two_buffers -// CHECK: memref.alloc() {pto.multi_buffer = 2 : i32} -// CHECK: pto.slot_marker %{{.*}}[%c0 -// CHECK: pto.slot_marker %{{.*}}[%c1 +// CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile +// CHECK: pto.multi_tile_get %[[MULTI]][%c0 +// CHECK: pto.multi_tile_get %[[MULTI]][%c1 +// CHECK-NOT: memref.alloc +// CHECK-NOT: pto.slot_marker -// Stage 2 -- PTOPlanMemory: two physical slots reserved -> 2-addr pointer_cast. +// Stage 2 -- PTOPlanMemory records two physical slot addresses. // PLAN: IR Dump After PlanMemory // PLAN-LABEL: func.func @const_slot_two_buffers -// PLAN: pto.pointer_cast({{.*}}, {{.*}}) -// PLAN: pto.slot_marker +// PLAN: pto.alloc_multi_tile {pto.multi_buffer_addrs = array} +// PLAN: pto.multi_tile_get +// PLAN-NOT: pto.pointer_cast -// Stage 3 -- ResolveBufferSelect: each slot_marker becomes a single-addr -// pointer_cast; original 2-addr pointer_cast remains as alloc anchor. +// Stage 3 -- ResolveBufferSelect materializes one addressed tile per use. // SELECT: IR Dump After PTOResolveBufferSelect // SELECT-LABEL: func.func @const_slot_two_buffers -// SELECT: %[[ANCHOR:.*]] = pto.pointer_cast(%{{.*}}, %{{.*}}) -// SELECT-NOT: pto.slot_marker -// SELECT: pto.pointer_cast(%[[ADDR0:.*]]) -// SELECT: pto.pointer_cast(%[[ADDR1:.*]]) +// SELECT-NOT: pto.alloc_multi_tile +// SELECT-NOT: pto.multi_tile_get +// SELECT: pto.alloc_tile addr = %{{.*}} +// SELECT: pto.alloc_tile addr = %{{.*}} diff --git a/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto b/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto index 31b4583b6f..14c055cd44 100644 --- a/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto @@ -9,11 +9,9 @@ // RUN: ptoas --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s // RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT -// Verifies that `multi_tile_get` with a dynamic slot SSA preserves the -// user-supplied slot expression all the way through `pto.slot_marker` and -// the final `arith.select` chain. The lowering must NOT rewrite the slot -// expression into `iv mod N` -- the frontend SSA is the source of truth -// for which slot this use refers to. +// Verifies that tile-native `multi_tile_get` preserves the user-supplied slot +// expression through the final `arith.select` chain. The lowering must not +// synthesize an implicit `iv mod N` expression. module { func.func @dyn_slot_prefetch( @@ -48,17 +46,22 @@ module { // CHECK: IR Dump After PTOViewToMemref // CHECK-LABEL: func.func @dyn_slot_prefetch -// CHECK: memref.alloc() {pto.multi_buffer = 2 : i32} +// CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile // CHECK: %[[NEXT:.*]] = arith.addi %{{.*}}, %c1 // CHECK-DAG: %[[CUR_IDX:.*]] = arith.remui %arg{{.*}}, %c2 // CHECK-DAG: %[[NEXT_IDX:.*]] = arith.remui %[[NEXT]], %c2 -// CHECK: pto.slot_marker %{{.*}}[%[[NEXT_IDX]] -// CHECK: pto.slot_marker %{{.*}}[%[[CUR_IDX]] +// CHECK: pto.multi_tile_get %[[MULTI]][%[[NEXT_IDX]] +// CHECK: pto.multi_tile_get %[[MULTI]][%[[CUR_IDX]] +// CHECK-NOT: memref.alloc +// CHECK-NOT: pto.slot_marker -// After ResolveBufferSelect: slot_marker -> per-slot pointer_cast + -// arith.select chain driven by the user SSA (not iv mod N). +// ResolveBufferSelect creates an addressed alloc_tile from an arith.select +// chain driven by the user SSA (not an implicit iv mod N). // SELECT: IR Dump After PTOResolveBufferSelect // SELECT-LABEL: func.func @dyn_slot_prefetch -// SELECT-NOT: pto.slot_marker -// SELECT: pto.pointer_cast(%{{.*}}, %{{.*}}) +// SELECT-NOT: pto.alloc_multi_tile +// SELECT-NOT: pto.multi_tile_get // SELECT: arith.select +// SELECT: pto.alloc_tile addr = %{{.*}} +// SELECT: arith.select +// SELECT: pto.alloc_tile addr = %{{.*}} diff --git a/test/lit/pto/multi_tile_level3_explicit_addr.pto b/test/lit/pto/multi_tile_level3_explicit_addr.pto index 37c1828ea8..6289855628 100644 --- a/test/lit/pto/multi_tile_level3_explicit_addr.pto +++ b/test/lit/pto/multi_tile_level3_explicit_addr.pto @@ -12,10 +12,8 @@ // RUN: not ptoas --pto-level=level2 --pto-arch=a5 %s 2>&1 | FileCheck %s --check-prefix=LVERR // Under --pto-level=level3 the caller owns local memory, PTOPlanMemory does not -// run, and `pto.alloc_multi_tile` carries an explicit base `addr`. The view -// lowering fans it out into a multi-address `pto.pointer_cast` anchor. Slot k is -// placed at `addr + k * slotBytes` (slotBytes = 16*16*2 = 512 here), so sync and -// PTOResolveBufferSelect keep working without automatic address planning. +// run, and `pto.alloc_multi_tile` carries an explicit base `addr`. Slot k is +// resolved as `addr + k * slotBytes` (slotBytes = 16*16*2 = 512 here). module { func.func @l3_two_buffers( @@ -43,23 +41,24 @@ module { } } -// PTOViewToMemref: explicit base address fans out into a 2-address -// pointer_cast; slot 1 = base + 512. No memref.alloc / multi_buffer attr. +// PTOViewToMemref preserves the tile-native allocation and selections. // CHECK: IR Dump After PTOViewToMemref // CHECK-LABEL: func.func @l3_two_buffers -// CHECK-NOT: pto.multi_buffer -// CHECK-DAG: %[[OFF1:.*]] = arith.constant 512 : i64 -// CHECK-DAG: %[[A1:.*]] = arith.addi %{{.*}}, %[[OFF1]] : i64 -// CHECK: pto.pointer_cast(%{{.*}}, %[[A1]]) -// CHECK: pto.slot_marker %{{.*}}[%c0 -// CHECK: pto.slot_marker %{{.*}}[%c1 +// CHECK: %[[BASE:.*]] = arith.constant 0 : i64 +// CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile addr = %[[BASE]] +// CHECK: pto.multi_tile_get %[[MULTI]][%c0 +// CHECK: pto.multi_tile_get %[[MULTI]][%c1 +// CHECK-NOT: pto.multi_buffer_addrs +// CHECK-NOT: memref.alloc -// ResolveBufferSelect: each slot_marker becomes a single-address pointer_cast; -// the 2-address anchor is preserved. +// ResolveBufferSelect materializes one addressed alloc_tile per selection. // SELECT: IR Dump After PTOResolveBufferSelect // SELECT-LABEL: func.func @l3_two_buffers -// SELECT: pto.pointer_cast(%{{.*}}, %{{.*}}) -// SELECT-NOT: pto.slot_marker +// SELECT-NOT: pto.alloc_multi_tile +// SELECT-NOT: pto.multi_tile_get +// SELECT: pto.alloc_tile addr = %{{.*}} +// SELECT: arith.addi +// SELECT: pto.alloc_tile addr = %{{.*}} // Outside level3 an explicit addr on alloc_multi_tile is rejected. // LVERR: unexpected 'addr' operand on pto.alloc_multi_tile: only supported when --pto-level=level3 diff --git a/test/lit/pto/multi_tile_n4_planmem_e2e.pto b/test/lit/pto/multi_tile_n4_planmem_e2e.pto index 361e06da7e..d44d97b3dd 100644 --- a/test/lit/pto/multi_tile_n4_planmem_e2e.pto +++ b/test/lit/pto/multi_tile_n4_planmem_e2e.pto @@ -8,11 +8,11 @@ // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN // RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN +// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT -// End-to-end test for the N == 4 multi-buffer path: PlanMemory must -// reserve 4 physical slots and emit a 4-address `pto.pointer_cast`. -// PTOResolveBufferSelect then lowers each constant-slot `pto.slot_marker` -// to a single-address pointer_cast picking the right slot. +// End-to-end test for N == 4: PlanMemory records four physical slot addresses, +// then PTOResolveBufferSelect materializes each selected tile handle. module { func.func @n4_rotate( @@ -50,18 +50,18 @@ module { } } -// PlanMemory: 4 physical addresses, one `pto.pointer_cast` with 4 i64 args. +// PlanMemory records four physical addresses on alloc_multi_tile. // PLAN: IR Dump After PlanMemory // PLAN-LABEL: func.func @n4_rotate -// PLAN: pto.pointer_cast({{[^,]*}}, {{[^,]*}}, {{[^,]*}}, {{[^,]*}}) +// PLAN: pto.alloc_multi_tile {pto.multi_buffer_addrs = array} +// PLAN-NOT: pto.pointer_cast -// PTOResolveBufferSelect: four single-addr casts (one per slot), no leftover -// slot_marker. The original 4-addr cast remains as the alloc anchor. +// PTOResolveBufferSelect: four addressed alloc_tile ops, no multi-buffer ops. // SELECT: IR Dump After PTOResolveBufferSelect // SELECT-LABEL: func.func @n4_rotate -// SELECT-NOT: pto.slot_marker -// SELECT: pto.pointer_cast({{[^,]*}}, {{[^,]*}}, {{[^,]*}}, {{[^,]*}}) -// SELECT: pto.pointer_cast(%{{[^,)]+}}) -// SELECT: pto.pointer_cast(%{{[^,)]+}}) -// SELECT: pto.pointer_cast(%{{[^,)]+}}) -// SELECT: pto.pointer_cast(%{{[^,)]+}}) +// SELECT-NOT: pto.alloc_multi_tile +// SELECT-NOT: pto.multi_tile_get +// SELECT: pto.alloc_tile addr = %{{.*}} +// SELECT: pto.alloc_tile addr = %{{.*}} +// SELECT: pto.alloc_tile addr = %{{.*}} +// SELECT: pto.alloc_tile addr = %{{.*}} diff --git a/test/lit/pto/multi_tile_no_loop_unroll.pto b/test/lit/pto/multi_tile_no_loop_unroll.pto index e7b4373e96..fe9c70ee80 100644 --- a/test/lit/pto/multi_tile_no_loop_unroll.pto +++ b/test/lit/pto/multi_tile_no_loop_unroll.pto @@ -8,11 +8,8 @@ // RUN: ptoas --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s -// Verifies that multi-buffer expression and slot selection work WITHOUT an -// enclosing `scf.for`. Two constant slots are tagged via `pto.slot_marker` -// at the memref layer and live independently of any loop induction -// variable. This is something the PR615 fully-automatic `iv mod N` path -// could not express. +// Verifies that tile-native multi-buffer selection works without an enclosing +// `scf.for`; the two constant slots do not depend on a loop induction value. module { func.func @unroll_two_slots( @@ -42,6 +39,8 @@ module { // CHECK: IR Dump After PTOViewToMemref // CHECK-LABEL: func.func @unroll_two_slots // CHECK-NOT: scf.for -// CHECK: memref.alloc() {pto.multi_buffer = 2 : i32} -// CHECK: pto.slot_marker %{{.*}}[%c0 -// CHECK: pto.slot_marker %{{.*}}[%c1 +// CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile +// CHECK: pto.multi_tile_get %[[MULTI]][%c0 +// CHECK: pto.multi_tile_get %[[MULTI]][%c1 +// CHECK-NOT: memref.alloc +// CHECK-NOT: pto.slot_marker diff --git a/test/lit/pto/multi_tile_reject_internal_addrs.pto b/test/lit/pto/multi_tile_reject_internal_addrs.pto new file mode 100644 index 0000000000..b547f8d48f --- /dev/null +++ b/test/lit/pto/multi_tile_reject_internal_addrs.pto @@ -0,0 +1,23 @@ +// 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 %s 2>&1 | FileCheck %s + +// Planner-generated slot addresses are an internal contract. Users must not +// bypass level validation or memory planning by attaching the attribute. + +module { + func.func @reject_internal_addrs() { + %multi = pto.alloc_multi_tile { + pto.multi_buffer_addrs = array + } : !pto.multi_tile_buf, count=2> + return + } +} + +// CHECK: attribute 'pto.multi_buffer_addrs' is reserved for pto-plan-memory diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 62aef35cc8..17fcfbb0a8 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -8,6 +8,7 @@ #include "ptoas.h" #include "PTO/IR/PTO.h" +#include "PTO/IR/PTOMultiBuffer.h" #include "PTO/Transforms/VPTOLLVMEmitter.h" #include "PTO/Transforms/Passes.h" #include "PTO/Transforms/BufferizableOpInterfaceImpl.h" @@ -2965,6 +2966,17 @@ int mlir::pto::compilePTOASModule( return 1; } + bool hasUserPlannedMultiAddrs = false; + module->walk([&](pto::AllocMultiTileOp op) { + if (!op->hasAttr(pto::kPtoMultiBufferAddrsAttrName)) + return; + op.emitError() << "attribute '" << pto::kPtoMultiBufferAddrsAttrName + << "' is reserved for pto-plan-memory"; + hasUserPlannedMultiAddrs = true; + }); + if (hasUserPlannedMultiAddrs) + return 1; + if (effectiveLevel == PTOBuildLevel::Level3) { // In level3 the caller owns local memory and PTOPlanMemory is skipped, so // every allocation must carry an explicit physical address. For From 3ab2d27e46aab411d55facc9f89249700a83ba67 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 17 Jul 2026 21:59:21 +0800 Subject: [PATCH 12/91] Keep declared tile handles native --- ...ptoas-tile-native-mainline-op-migration.md | 6 +- include/PTO/IR/PTOOps.td | 8 +- .../Transforms/InsertSync/PTOIRTranslator.h | 1 + .../Transforms/InsertSync/PTOIRTranslator.cpp | 34 ++++++++ lib/PTO/Transforms/PTOPlanMemory.cpp | 6 +- lib/PTO/Transforms/PTOToEmitC.cpp | 44 +++++++++- lib/PTO/Transforms/PTOViewToMemref.cpp | 86 +------------------ test/lit/pto/declare_tile_tile_native.pto | 50 +++++++++++ test/lit/pto/tcvt_low_precision_a5_valid.pto | 17 ++-- ...tload_tprefetch_low_precision_a5_valid.pto | 3 +- .../lit/pto/tstore_low_precision_a5_valid.pto | 9 +- 11 files changed, 157 insertions(+), 107 deletions(-) create mode 100644 test/lit/pto/declare_tile_tile_native.pto diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index 3ceb81502d..9d7155fe3e 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -69,8 +69,8 @@ PTO tile/view IR | `pto.alloc_tile` | `PTOViewToMemref` 已透传；memplan 填写 `addr` | 保持现状，作为单 slot local root | legacy/modern 均继续收集无地址 root；level3 校验显式地址 | 现有 InsertSync 已读取类型、大小和常量地址；GSS 需从 `alloc_tile addr` 构造 `PointerLikeInfo` | EmitC/VPTO 已有直接 lowering；保留 level1/2 自动地址和 level3 显式地址测试 | | `pto.alloc_multi_tile` | 已保持 tile-native；memplan 写入内部 N-address 属性，level3 保留用户 base | op 本身作为 N-slot root，最终可由 backend 直接消费 | legacy/modern 已支持 `slotCount=N`、`slotBytes`、`totalBytes` 和每 slot offsets；禁止 sibling slot 互相 alias | InsertSync/GSS 已直接识别 multi root，不依赖 `slot_marker -> pointer_cast` | 当前由 `PTOResolveBufferSelect` 展开成带地址的 `alloc_tile`；已覆盖 constant/dynamic slot、loop/non-loop 测试 | | `pto.multi_tile_get` | 已保持到 sync 和 `PTOResolveBufferSelect` | 最终继续保持到 backend | result 已 alias 到 multi root，并携带 slot expression | constant slot 只访问一个地址；dynamic slot 保守访问全部地址，并保留 dyn event-id 推导 | 当前 resolve pass 根据 slot 生成带单地址的 `alloc_tile`；不再生成 `slot_marker` | -| `pto.declare_tile` | 转成 `pto.declare_tile_memref + pto.bind_tile` | 保持 tile handle，表示地址由 `tpop` 等运行时操作绑定，不参与静态 local allocation | 不作为普通静态 root；需要单独标记 runtime-bound root | 建立 `declare_tile -> tpop/tassign -> tile users` 的 root/address 数据流 | EmitC/VPTO 直接 lowering；覆盖 tpush/tpop、分支和循环 | -| `pto.declare_tile_memref` | `PTOViewToMemref` 生成的内部占位 op | 所有 `declare_tile` 用户迁移后删除 | 删除 legacy/modern 特判 | 删除 InsertSync 的专用 root 注册 | 删除 EmitC pattern 和相关测试 | +| `pto.declare_tile` | 已保持 tile-native，不参与静态 local allocation | 保持现状，表示地址由 `tpop/tassign` 等运行时操作绑定 | legacy/modern 不收集为 allocation root | InsertSync 直接注册 symbolic root；GSS 按 declared SSA identity 建模 | EmitC 已直接声明 Tile；已覆盖 tpush/tpop 和 interleaved sync/GSS | +| `pto.declare_tile_memref` | 仅保留 legacy memref IR 兼容，不再由 `PTOViewToMemref` 生成 | 其余旧兼容用户清除后删除 | legacy 仅跳过，不分配 | InsertSync 暂保兼容入口 | 暂保 legacy EmitC pattern，最终随兼容路径删除 | | `pto.bind_tile` | memref 与 tile metadata 的桥接和 alias anchor | local tile 不再需要；只允许迁移期兼容，最终删除或限制为外部 ABI bridge | 先统一 `getAliasSource()`，切换后删除 bind 分支 | 先统一 alias tracing，切换后删除 bind 分支 | 删除 materialize/EmitC bind lowering 测试 | | `pto.pointer_cast` | 表达 memref root 的物理地址或 multi-address root | 仅保留真正的 raw-address ABI bridge；普通 local allocation 改由 allocation op 持有地址 | 不再作为 `alloc_tile` 的 materialization 结果 | 若保留，sync 继续读取其地址；multi-buffer 不再必须依赖它 | 后端保留 raw-address lowering，删除 alloc fallback 用法 | | `pto.slot_marker` | `multi_tile_get` 在 memref 层的内部 slot 标签 | `multi_tile_get` 直接承担 slot 标签语义，最终删除 | 删除 alias 特判 | 将 slot narrowing 和 dyn event-id 逻辑迁移到 `multi_tile_get` | 删除 `PTOResolveBufferSelect` 对 slot_marker 的入口 | @@ -267,7 +267,7 @@ op 默认不需要迁移实现；仍需通过完整 lit 确认它们没有间接 删除两个 pass 前的支持状态： 1. `pto.treshape` 和 `pto.bitcast` 的 alias 传播。 -2. `pto.declare_tile` runtime-bound root。 +2. `pto.declare_tile` runtime-bound root：已完成 EmitC、memplan、InsertSync/GSS 主路径迁移。 3. `pto.alloc_multi_tile` 的 N-address root：已完成。 4. `pto.multi_tile_get` 的 constant/dynamic slot narrowing：已完成。 5. tile 类型 helper argument 和 `func.call` effects。 diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index b5839380ca..75a6803f76 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -2632,12 +2632,10 @@ def LocalArraySetOp : PTO_Op<"local_array_set"> { } def DeclareTileMemRefOp : PTO_Op<"declare_tile_memref"> { - let summary = "Internal memref placeholder for a tile whose address is assigned later"; + let summary = "Legacy memref placeholder for a runtime-bound tile"; let description = [{ - Internal lowering op used by PTOViewToMemref. This op does not allocate - storage; it only provides a memref-typed SSA handle so later passes can - attach tile metadata through pto.bind_tile before the address is filled by - pipe operations such as pto.tpop. + Compatibility op for legacy memref-based IR. New PTOViewToMemref pipelines + preserve `pto.declare_tile` directly instead of producing this op. }]; let results = (outs AnyMemRef:$result); diff --git a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h index b9c439ea5a..0012f79f29 100644 --- a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h +++ b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h @@ -71,6 +71,7 @@ class PTOIRTranslator { void UpdateKernelArgMemInfo(); LogicalResult UpdateAllocTileOpMemInfo(pto::AllocTileOp op); LogicalResult UpdateAllocMultiTileOpMemInfo(pto::AllocMultiTileOp op); + LogicalResult UpdateDeclareTileOpMemInfo(pto::DeclareTileOp op); LogicalResult UpdateDeclareGlobalOpMemInfo(pto::DeclareGlobalOp op); LogicalResult UpdateDeclareTileMemRefOpMemInfo(pto::DeclareTileMemRefOp op); LogicalResult UpdatePointerCastOpMemInfo(pto::PointerCastOp op); diff --git a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp index e699d19a3b..e6aef631f1 100644 --- a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp +++ b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp @@ -375,6 +375,10 @@ void PTOIRTranslator::RecursionIR(Region *region) { return WalkResult::interrupt(); } } + else if (auto declareOp = dyn_cast(op)) { + if (failed(UpdateDeclareTileOpMemInfo(declareOp))) + return WalkResult::interrupt(); + } else if (auto declareOp = dyn_cast(op)) { if (failed(UpdateDeclareTileMemRefOpMemInfo(declareOp))) { return WalkResult::interrupt(); @@ -655,6 +659,36 @@ LogicalResult PTOIRTranslator::UpdatePointerCastOpMemInfo(pto::PointerCastOp op) return success(); } +LogicalResult +PTOIRTranslator::UpdateDeclareTileOpMemInfo(pto::DeclareTileOp op) { + Value result = op.getTile(); + auto tileType = dyn_cast(result.getType()); + if (!tileType) + return op.emitError("requires a tile_buf result for sync analysis"); + + uint64_t sizeInBytes = pto::getPTOStorageElemByteSize( + tileType.getElementType()); + if (sizeInBytes == 0) + return failure(); + for (int64_t dim : tileType.getShape()) { + if (dim == ShapedType::kDynamic) { + sizeInBytes = 0; + break; + } + sizeInBytes *= static_cast(dim); + } + + pto::AddressSpace space = pto::AddressSpace::MAT; + if (auto attr = dyn_cast_or_null( + tileType.getMemorySpace())) + space = attr.getAddressSpace(); + + auto info = std::make_unique( + result, result, space, SmallVector{0}, sizeInBytes); + buffer2MemInfoMap_[result].emplace_back(info->clone()); + return success(); +} + LogicalResult PTOIRTranslator::UpdateDeclareTileMemRefOpMemInfo(pto::DeclareTileMemRefOp op) { Value res = op.getResult(); diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index de8b9f84dd..cfd58e3481 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -464,10 +464,8 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { if (mayAliasOp.has_value()) { auto aliasPair = mayAliasOp.value(); UpdateBufferAlias(aliasPair.first, aliasPair.second); - } else if (isa(op)) { - // Internal placeholder for a tile whose runtime address is assigned by - // pipe operations such as tpop. This op does not allocate local storage - // and should not participate in memory planning. + } else if (isa(op)) { + // Runtime-bound tile handles do not allocate static local storage. return WalkResult::advance(); } else if (auto bindOp = dyn_cast(op)) { // BindTile result is only an alias of the source buffer. Treat every use diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 4fbe445370..d5f148c003 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -13168,7 +13168,8 @@ struct PTOAllocTileToEmitC static FailureOr createEmitCTileVariable(ConversionPatternRewriter &rewriter, Location loc, const TypeConverter *typeConverter, - pto::TileBufType tileTy) { + pto::TileBufType tileTy, + bool initializeDynamicValidToShape = false) { auto tileTypeString = getEmitCTileTypeString(tileTy); if (!tileTypeString) return failure(); @@ -13177,6 +13178,25 @@ createEmitCTileVariable(ConversionPatternRewriter &rewriter, Location loc, if (!convertedTy) convertedTy = emitc::OpaqueType::get(rewriter.getContext(), *tileTypeString); + if (initializeDynamicValidToShape && tileTy.hasDynamicValid()) { + auto shape = tileTy.getShape(); + if (shape.size() != 2 || llvm::is_contained(shape, ShapedType::kDynamic)) + return failure(); + Type i32Ty = emitc::OpaqueType::get(rewriter.getContext(), "int32_t"); + pto::BLayout blayout = getTileBufBLayoutValue(tileTy.getConfigAttr()); + SmallVector constructorArgs; + constructorArgs.push_back(makeEmitCIntConstant( + rewriter, loc, i32Ty, + renderTileTemplateDim(shape[0], tileTy.getElementType(), blayout, 0))); + constructorArgs.push_back(makeEmitCIntConstant( + rewriter, loc, i32Ty, + renderTileTemplateDim(shape[1], tileTy.getElementType(), blayout, 1))); + return rewriter + .create(loc, convertedTy, *tileTypeString, + ArrayAttr{}, ArrayAttr{}, constructorArgs) + .getResult(0); + } + Value tile = rewriter .create( loc, getEmitCVariableResultType(convertedTy), @@ -13185,6 +13205,27 @@ createEmitCTileVariable(ConversionPatternRewriter &rewriter, Location loc, return loadEmitCVariableIfNeeded(rewriter, loc, tile); } +struct PTODeclareTileToEmitC + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(pto::DeclareTileOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + (void)adaptor; + auto tileType = dyn_cast(op.getTile().getType()); + if (!tileType) + return rewriter.notifyMatchFailure(op, "expected a tile_buf result"); + FailureOr tile = createEmitCTileVariable( + rewriter, op.getLoc(), getTypeConverter(), tileType, + /*initializeDynamicValidToShape=*/true); + if (failed(tile)) + return rewriter.notifyMatchFailure( + op, "only rank-2 declare_tile handles can be converted to EmitC"); + rewriter.replaceOp(op, *tile); + return success(); + } +}; + struct PTOTReshapeToEmitC : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -14137,6 +14178,7 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, PTOArch targetArch) { patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); + patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 1da3c177aa..4e745b6232 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -67,7 +67,7 @@ static Type convertPTOTypeToMemRef(Type t); static bool hasTileNativeAllocationRoot(func::FuncOp func) { bool found = false; auto result = func.walk([&](Operation *op) { - if (isa(op)) { + if (isa(op)) { found = true; return WalkResult::interrupt(); } @@ -951,46 +951,6 @@ static void markForceDynamicValidShape(Operation *op, bool force, func.setFunctionType(FunctionType::get(ctx, newInputs, newResults)); } -[[maybe_unused]] static LogicalResult lowerDeclareTileOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector declaredTiles; - func.walk([&](mlir::pto::DeclareTileOp op) { declaredTiles.push_back(op); }); - - for (auto op : declaredTiles) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - auto tbTy = dyn_cast(op.getTile().getType()); - if (!tbTy) { - op.emitError("declare_tile result must be tile_buf type"); - return failure(); - } - - auto targetType = dyn_cast(convertPTOTypeToMemRef(tbTy)); - if (!targetType) { - op.emitError("failed to convert declare_tile result to memref type"); - return failure(); - } - - auto configAttr = tbTy.getConfigAttr(); - if (!configAttr) - configAttr = pto::TileBufConfigAttr::getDefault(ctx); - - Value vRow; - Value vCol; - materializeStaticValidDims(rewriter, loc, tbTy, vRow, vCol); - - auto declaredMemRef = - rewriter.create(loc, targetType); - auto bindOp = rewriter.create( - loc, targetType, declaredMemRef.getResult(), vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - rewriter.replaceOp(op, bindOp.getResult()); - } - return success(); -} - static Value castIndexToI64(IRRewriter &rewriter, Location loc, Value value) { Type i64Ty = rewriter.getI64Type(); if (value.getType() == i64Ty) @@ -1753,48 +1713,10 @@ struct PTOViewToMemrefPass // ------------------------------------------------------------------ // ------------------------------------------------------------------ - // Stage 0.75: lower pto.declare_tile -> pto.declare_tile_memref + - // pto.bind_tile + // Stage 0.75: keep pto.declare_tile tile-native. It is a symbolic handle + // whose address is assigned at runtime by tpop/tassign, not a static + // allocation root for PlanMemory. // ------------------------------------------------------------------ - DefaultInlineVector declaredTiles; - func.walk([&](mlir::pto::DeclareTileOp op) { declaredTiles.push_back(op); }); - - for (auto op : declaredTiles) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - auto tbTy = dyn_cast(op.getTile().getType()); - if (!tbTy) { - op.emitError("declare_tile result must be tile_buf type"); - signalPassFailure(); - return; - } - - auto targetType = dyn_cast(convertPTOTypeToMemRef(tbTy)); - if (!targetType) { - op.emitError("failed to convert declare_tile result to memref type"); - signalPassFailure(); - return; - } - - auto configAttr = tbTy.getConfigAttr(); - if (!configAttr) - configAttr = pto::TileBufConfigAttr::getDefault(ctx); - - Value vRow; - Value vCol; - materializeStaticValidDims(rewriter, loc, tbTy, vRow, vCol); - - auto declaredMemRef = - rewriter.create(loc, targetType); - auto bindOp = rewriter.create( - loc, targetType, declaredMemRef.getResult(), - vRow ? vRow : Value(), vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - - rewriter.replaceOp(op, bindOp.getResult()); - } // ------------------------------------------------------------------ // Stage 0.8: normalize pto.tassign result type to match tile operand diff --git a/test/lit/pto/declare_tile_tile_native.pto b/test/lit/pto/declare_tile_tile_native.pto new file mode 100644 index 0000000000..9a8a40ab50 --- /dev/null +++ b/test/lit/pto/declare_tile_tile_native.pto @@ -0,0 +1,50 @@ +// 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 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC + +module { + func.func @declare_tile_native( + %gm_slot_buffer: memref<256xf32, #pto.address_space>, + %consumer_buf: i32) attributes {pto.kernel_kind = #pto.kernel_kind} { + %pipe = pto.initialize_l2g2l_pipe { + dir_mask = 1, + slot_size = 1024, + slot_num = 8, + local_slot_num = 8, + flag_base = 0 + }(%gm_slot_buffer : memref<256xf32, #pto.address_space>, + %consumer_buf : i32) -> !pto.pipe + %tile = pto.declare_tile + -> !pto.tile_buf + pto.tpop(%tile, %pipe + : !pto.tile_buf, !pto.pipe) { + split = 1 + } + return + } +} + +// NATIVE-LABEL: func.func @declare_tile_native +// NATIVE: pto.declare_tile -> !pto.tile_buf +// NATIVE: pto.tpop({{.*}} : !pto.tile_buf +// NATIVE-NOT: pto.declare_tile_memref +// NATIVE-NOT: pto.bind_tile +// NATIVE-NOT: memref.alloc + +// EMITC-LABEL: AICORE void declare_tile_native +// EMITC: Tile} -// CHECK: pto.declare_tile_memref -> memref<16x8x!pto.f4E1M2x2 -// CHECK: pto.declare_tile_memref -> memref<16x16x!pto.hif8 -// CHECK: pto.tcvt ins(%1 {rmode = #pto, satmode = #pto} : memref<16x16xf32 -// CHECK: outs(%3 : memref<16x16xf8E4M3FN -// CHECK: pto.tcvt ins(%5 {rmode = #pto, satmode = #pto} : memref<16x16xbf16 -// CHECK: outs(%7 : memref<16x8x!pto.f4E1M2x2 -// CHECK: pto.tcvt ins(%9 {rmode = #pto, satmode = #pto} : memref<16x16x!pto.hif8 -// CHECK: outs(%11 : memref<16x16xf32 +// CHECK: pto.declare_tile -> !pto.tile_buf +// CHECK: pto.declare_tile -> !pto.tile_buf +// CHECK-NOT: pto.declare_tile_memref +// 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: outs(%{{.*}} : !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/tload_tprefetch_low_precision_a5_valid.pto b/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto index affb448722..0b5c48df87 100644 --- a/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto +++ b/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto @@ -24,6 +24,7 @@ module { } // CHECK: func.func @tload_tprefetch_low_precision_a5_valid(%arg0: memref<16x16xf8E4M3FN>, %arg1: memref<16x16x!pto.hif8>) -// CHECK: pto.declare_tile_memref -> memref<16x16x!pto.hif8 +// CHECK: pto.declare_tile -> !pto.tile_buf +// CHECK-NOT: pto.declare_tile_memref // CHECK: pto.tload ins(%arg0 : memref<16x16xf8E4M3FN>) outs( // CHECK: pto.tprefetch ins(%arg1 : memref<16x16x!pto.hif8>) outs( diff --git a/test/lit/pto/tstore_low_precision_a5_valid.pto b/test/lit/pto/tstore_low_precision_a5_valid.pto index b139645a04..520ef20e64 100644 --- a/test/lit/pto/tstore_low_precision_a5_valid.pto +++ b/test/lit/pto/tstore_low_precision_a5_valid.pto @@ -25,12 +25,15 @@ module { } // CHECK: func.func @tstore_low_precision_a5_valid(%arg0: memref<16x16xf8E4M3FN>, %arg1: memref<16x16x!pto.hif8>, %arg2: i64) -// CHECK: pto.bind_tile +// CHECK: %[[VEC:.*]] = pto.declare_tile -> !pto.tile_buf +// CHECK: %[[ACC:.*]] = pto.declare_tile -> !pto.tile_buf +// CHECK-NOT: pto.declare_tile_memref +// CHECK-NOT: pto.bind_tile // CHECK: pto.section.vector { -// CHECK: pto.tstore ins( +// CHECK: pto.tstore ins(%[[VEC]] : !pto.tile_buf) // CHECK: outs(%arg0 : memref<16x16xf8E4M3FN>) // CHECK: } // CHECK: pto.section.cube { -// CHECK: pto.tstore ins( +// CHECK: pto.tstore ins(%[[ACC]] : !pto.tile_buf, %arg2 : i64) // CHECK: outs(%arg1 : memref<16x16x!pto.hif8>) // CHECK: } From 9bf81ef01da9e2e2d930fab9903b23c1f0b0b436 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 17 Jul 2026 22:12:06 +0800 Subject: [PATCH 13/91] Keep tile reshape views native --- ...ptoas-tile-native-mainline-op-migration.md | 2 +- .../Transforms/InsertSync/PTOIRTranslator.cpp | 3 ++ lib/PTO/Transforms/PTOViewToMemref.cpp | 47 ++++++++++--------- test/lit/pto/treshape_tile_native_arg.pto | 41 ++++++++++++++++ 4 files changed, 70 insertions(+), 23 deletions(-) create mode 100644 test/lit/pto/treshape_tile_native_arg.pto diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index 9d7155fe3e..94aaab1771 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -81,7 +81,7 @@ PTO tile/view IR | Op | 当前行为 | Tile-native 目标 | Memplan / Alias | InsertSync | Backend 与测试 | |---|---|---|---|---|---| | `pto.subview` | 通常转成 `memref.subview + pto.bind_tile`，之后再恢复 tile | 全程保留 `pto.subview` | result 与 source 同 root；按 layout/stride/offset 计算 byte range | 现有 InsertSync 已有精确 `pto.subview` 地址计算；GSS 和 BufidSync 补齐同等能力 | EmitC/VPTO 增加直接 lowering；覆盖 row/col-major、dynamic offset、valid shape | -| `pto.treshape` | 转成 memref view，materialize 阶段恢复 `treshape` | 保持原 op | result 与 source 完全 alias，size 不变 | InsertSync/GSS 增加 result-to-source alias 传播 | backend 已有 tile op 路径；覆盖动态 valid shape 和控制流 | +| `pto.treshape` | 已保持 tile-native，不再生成 memref view 或 `bind_tile` | 保持原 op | result 与 source 完全 alias，size 不变；legacy/modern memplan 已传播 alias | InsertSync/GSS 已直接传播 result-to-source alias | EmitC 已直接 lowering；覆盖 tile 参数、动态 valid shape 和静态 valid shape | | `pto.bitcast` | 转成 memref view，materialize 阶段恢复 `bitcast` | 保持原 op | result 与 source alias；以 byte range 为准，校验总容量一致 | InsertSync/GSS 增加 alias 传播，不能按 element type 分裂 root | EmitC/VPTO 直接 lowering；覆盖 dtype 改变和相同容量 | | `pto.set_validshape` | 在 memref 层通过 bind metadata 记录，materialize 后恢复 tile 类型 | 直接更新/产生 tile handle metadata，不改变物理 root | result 与 source alias，物理 size 使用 allocation shape，不使用 valid shape | 同一物理范围；repeat 分析读取更新后的 valid shape | 两后端直接消费；覆盖 if/for 和动态 valid shape | | `pto.get_validshape` | materialize 后依赖 tile handle 读取 metadata | 直接读取 tile operand metadata | 不产生 root | 无内存 effect | 两后端直接 lowering；覆盖静态折叠和动态值 | diff --git a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp index e6aef631f1..a36dd89310 100644 --- a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp +++ b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp @@ -415,6 +415,9 @@ void PTOIRTranslator::RecursionIR(Region *region) { else if (auto multiGet = dyn_cast(op)) { UpdateMultiTileGetAliasBufferInfo(multiGet); } + else if (auto reshape = dyn_cast(op)) { + UpdateAliasBufferInfo(reshape.getResult(), reshape.getSrc()); + } else if (auto castOp = dyn_cast(op)) { UpdateAliasBufferInfo(castOp.getResult(), castOp.getSource()); } diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 4e745b6232..3253ec4618 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -77,6 +77,16 @@ static bool hasTileNativeAllocationRoot(func::FuncOp func) { return found; } +static bool hasMigratedTileNativeView(func::FuncOp func) { + bool found = false; + auto result = func.walk([&](pto::TReshapeOp) { + found = true; + return WalkResult::interrupt(); + }); + (void)result; + return found; +} + constexpr size_t kTileRank2D = 2; constexpr size_t kRowDimensionIndex = 0; constexpr size_t kColumnDimensionIndex = 1; @@ -1546,24 +1556,6 @@ static Value buildTileBufViewLikeValue(Operation *anchorOp, Value src, } static LogicalResult lowerTileBufViewLikeOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector reshapes; - func.walk([&](mlir::pto::TReshapeOp op) { reshapes.push_back(op); }); - for (auto op : reshapes) { - auto tbTy = dyn_cast(op.getResult().getType()); - if (!tbTy) { - op.emitError("treshape result must be tile_buf type"); - return failure(); - } - if (isa(op->getOperand(0).getType())) - continue; - Value lowered = buildTileBufViewLikeValue(op, op->getOperand(0), tbTy, - "treshape", ctx); - if (!lowered) - return failure(); - IRRewriter rewriter(ctx); - rewriter.replaceOp(op, lowered); - } - DefaultInlineVector bitcasts; func.walk([&](mlir::pto::BitcastOp op) { bitcasts.push_back(op); }); for (auto op : bitcasts) { @@ -1612,15 +1604,26 @@ struct PTOViewToMemrefPass continue; auto fnTy = func.getFunctionType(); + bool preserveTileABI = hasMigratedTileNativeView(func); + bool tileNativeMainline = + preserveTileABI || hasTileNativeAllocationRoot(func); // ------------------------------------------------------------------ // Stage 0.10: Rewrite Function Signature // ------------------------------------------------------------------ SmallVector newInputs; - for (Type t : fnTy.getInputs()) newInputs.push_back(convertPTOTypeToMemRef(t)); + for (Type t : fnTy.getInputs()) { + newInputs.push_back(preserveTileABI && isa(t) + ? t + : convertPTOTypeToMemRef(t)); + } SmallVector newResults; - for (Type t : fnTy.getResults()) newResults.push_back(convertPTOTypeToMemRef(t)); + for (Type t : fnTy.getResults()) { + newResults.push_back(preserveTileABI && isa(t) + ? t + : convertPTOTypeToMemRef(t)); + } func.setFunctionType(FunctionType::get(ctx, newInputs, newResults)); @@ -1673,7 +1676,7 @@ struct PTOViewToMemrefPass for (unsigned i = 0; i < entry.getNumArguments(); ++i) { Type origTy = fnTy.getInputs()[i]; auto tbTy = dyn_cast(origTy); - if (!tbTy) + if (!tbTy || newInputs[i] == origTy) continue; auto configAttr = tbTy.getConfigAttr(); @@ -2055,7 +2058,7 @@ struct PTOViewToMemrefPass // Stage 3: Rewrite Compute Ops // [关键] 全面使用 op->getOperand(i) 避免 Typed Accessor Crash // ------------------------------------------------------------------ - if (!hasTileNativeAllocationRoot(func)) { + if (!tileNativeMainline) { // --- TLoadOp [Src, Dst] --- DefaultInlineVector loads; diff --git a/test/lit/pto/treshape_tile_native_arg.pto b/test/lit/pto/treshape_tile_native_arg.pto new file mode 100644 index 0000000000..bbf649ad92 --- /dev/null +++ b/test/lit/pto/treshape_tile_native_arg.pto @@ -0,0 +1,41 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o - 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE + +module attributes {pto.target_arch = "a5"} { + func.func private @reshape_arg( + %src: !pto.tile_buf) + -> !pto.tile_buf { + %result = pto.treshape %src + : !pto.tile_buf + -> !pto.tile_buf + return %result : !pto.tile_buf + } +} + +// NATIVE-LABEL: func.func private @reshape_arg( +// NATIVE-SAME: !pto.tile_buf +// NATIVE: %[[RESULT:.*]] = pto.treshape %arg0 +// NATIVE-SAME: !pto.tile_buf +// NATIVE-SAME: -> !pto.tile_buf +// NATIVE: return %[[RESULT]] : !pto.tile_buf +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile From 72e6f1502bf4f92cbd31b3ee8ed0e3d7aba8c4a4 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 17 Jul 2026 22:29:19 +0800 Subject: [PATCH 14/91] Keep tile bitcast views native --- ...ptoas-tile-native-mainline-op-migration.md | 2 +- .../Transforms/InsertSync/PTOIRTranslator.cpp | 3 + lib/PTO/Transforms/PTOViewToMemref.cpp | 119 ++---------------- test/lit/pto/bitcast_tile_native.pto | 40 ++++++ 4 files changed, 51 insertions(+), 113 deletions(-) create mode 100644 test/lit/pto/bitcast_tile_native.pto diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index 94aaab1771..4a28170ea2 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -82,7 +82,7 @@ PTO tile/view IR |---|---|---|---|---|---| | `pto.subview` | 通常转成 `memref.subview + pto.bind_tile`，之后再恢复 tile | 全程保留 `pto.subview` | result 与 source 同 root；按 layout/stride/offset 计算 byte range | 现有 InsertSync 已有精确 `pto.subview` 地址计算；GSS 和 BufidSync 补齐同等能力 | EmitC/VPTO 增加直接 lowering；覆盖 row/col-major、dynamic offset、valid shape | | `pto.treshape` | 已保持 tile-native，不再生成 memref view 或 `bind_tile` | 保持原 op | result 与 source 完全 alias，size 不变；legacy/modern memplan 已传播 alias | InsertSync/GSS 已直接传播 result-to-source alias | EmitC 已直接 lowering；覆盖 tile 参数、动态 valid shape 和静态 valid shape | -| `pto.bitcast` | 转成 memref view，materialize 阶段恢复 `bitcast` | 保持原 op | result 与 source alias；以 byte range 为准，校验总容量一致 | InsertSync/GSS 增加 alias 传播，不能按 element type 分裂 root | EmitC/VPTO 直接 lowering；覆盖 dtype 改变和相同容量 | +| `pto.bitcast` | 已保持 tile-native，不再生成 memref view 或 `bind_tile` | 保持原 op | result 与 source alias；legacy/modern memplan 已按 byte range 传播同一 root，并校验总容量一致 | InsertSync/GSS 已直接传播 alias，不按 element type 分裂 root | EmitC 已直接 lowering；覆盖 tile 参数和同容量 dtype 改变 | | `pto.set_validshape` | 在 memref 层通过 bind metadata 记录，materialize 后恢复 tile 类型 | 直接更新/产生 tile handle metadata，不改变物理 root | result 与 source alias，物理 size 使用 allocation shape，不使用 valid shape | 同一物理范围；repeat 分析读取更新后的 valid shape | 两后端直接消费；覆盖 if/for 和动态 valid shape | | `pto.get_validshape` | materialize 后依赖 tile handle 读取 metadata | 直接读取 tile operand metadata | 不产生 root | 无内存 effect | 两后端直接 lowering；覆盖静态折叠和动态值 | | `pto.tile_buf_addr` | ptr-like 形式转成线性 memref，后续 materialize 可能恢复 | 直接从 tile root/view 计算地址；返回 pointer-like PTO 类型 | 不产生新 root | 若结果参与 load/store，需要保留到原 root 的 provenance | EmitC/VPTO 直接生成地址表达式；覆盖 alloc、subview、multi slot | diff --git a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp index a36dd89310..dd8355fccc 100644 --- a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp +++ b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp @@ -418,6 +418,9 @@ void PTOIRTranslator::RecursionIR(Region *region) { else if (auto reshape = dyn_cast(op)) { UpdateAliasBufferInfo(reshape.getResult(), reshape.getSrc()); } + else if (auto bitcast = dyn_cast(op)) { + UpdateAliasBufferInfo(bitcast.getResult(), bitcast.getSrc()); + } else if (auto castOp = dyn_cast(op)) { UpdateAliasBufferInfo(castOp.getResult(), castOp.getSource()); } diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 3253ec4618..a8bc184232 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -79,17 +79,18 @@ static bool hasTileNativeAllocationRoot(func::FuncOp func) { static bool hasMigratedTileNativeView(func::FuncOp func) { bool found = false; - auto result = func.walk([&](pto::TReshapeOp) { - found = true; - return WalkResult::interrupt(); + auto result = func.walk([&](Operation *op) { + if (isa(op)) { + found = true; + return WalkResult::interrupt(); + } + return WalkResult::advance(); }); (void)result; return found; } constexpr size_t kTileRank2D = 2; -constexpr size_t kRowDimensionIndex = 0; -constexpr size_t kColumnDimensionIndex = 1; constexpr unsigned kShapeVectorInlineCapacity = 4; constexpr unsigned kOperationVectorInlineCapacity = 8; @@ -258,27 +259,6 @@ static OpTy replaceOpWithClonedAttrs(IRRewriter &rewriter, Operation *op, return newOp; } -static Value resolveTileBufViewLikeSource(Value src) { - if (isa(src.getType())) - return src; - - if (auto regionResult = dyn_cast(src)) { - if (auto fusionRegion = - dyn_cast(regionResult.getOwner())) { - auto yieldOp = dyn_cast( - fusionRegion.getBody().front().getTerminator()); - if (!yieldOp) - return Value(); - unsigned resultIndex = regionResult.getResultNumber(); - if (resultIndex >= yieldOp.getNumOperands()) - return Value(); - return resolveTileBufViewLikeSource(yieldOp.getOperand(resultIndex)); - } - } - - return Value(); -} - static LogicalResult synthesizeMissingTQuantTmpOps(func::FuncOp func, MLIRContext *ctx) { if (mlir::pto::getTargetArch(func.getOperation()) == mlir::pto::PTOArch::A5) @@ -639,19 +619,6 @@ static SmallVector computeCompactStrides(ArrayRef shape) { return strides; } -static void materializeStaticValidDims(IRRewriter &rewriter, Location loc, - mlir::pto::TileBufType tbTy, Value &vRow, - Value &vCol) { - ArrayRef validShape = tbTy.getValidShape(); - if (tbTy.hasDynamicValid()) - return; - if (!validShape.empty() && validShape[kRowDimensionIndex] >= 0) - vRow = makeIndexConstant(rewriter, loc, validShape[kRowDimensionIndex]); - if (validShape.size() >= kTileRank2D && - validShape[kColumnDimensionIndex] >= 0) - vCol = makeIndexConstant(rewriter, loc, validShape[kColumnDimensionIndex]); -} - static bool checkMultipleOf(Operation *op, int64_t value, int64_t divisor, StringRef label) { if (divisor <= 0) { @@ -1508,74 +1475,6 @@ static LogicalResult lowerSubViewOps(func::FuncOp func, MLIRContext *ctx) { return success(); } -static Value buildTileBufViewLikeValue(Operation *anchorOp, Value src, - mlir::pto::TileBufType tbTy, - StringRef viewSemantics, - MLIRContext *ctx) { - Location loc = anchorOp->getLoc(); - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(anchorOp); - - src = resolveTileBufViewLikeSource(src); - auto srcMrTy = dyn_cast_or_null(src.getType()); - if (!srcMrTy) { - anchorOp->emitError("tile_buf view op src must be lowered to memref first"); - return Value(); - } - - auto targetType = dyn_cast(convertPTOTypeToMemRef(tbTy)); - if (!targetType) { - anchorOp->emitError("failed to convert tile_buf type to memref type"); - return Value(); - } - for (int64_t dim : targetType.getShape()) { - if (dim == ShapedType::kDynamic) { - anchorOp->emitError("dynamic shapes are not supported for tile_buf view ops"); - return Value(); - } - } - - Value parentVRow; - Value parentVCol; - lookupValidDims(src, parentVRow, parentVCol); - Value vRow = parentVRow; - Value vCol = parentVCol; - materializeStaticValidDims(rewriter, loc, tbTy, vRow, vCol); - - auto configAttr = tbTy.getConfigAttr(); - if (!configAttr) - configAttr = pto::TileBufConfigAttr::getDefault(ctx); - - auto bindOp = rewriter.create( - loc, targetType, src, vRow ? vRow : Value(), vCol ? vCol : Value(), - configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - if (!viewSemantics.empty()) - bindOp->setAttr("pto.view_semantics", rewriter.getStringAttr(viewSemantics)); - return bindOp.getResult(); -} - -static LogicalResult lowerTileBufViewLikeOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector bitcasts; - func.walk([&](mlir::pto::BitcastOp op) { bitcasts.push_back(op); }); - for (auto op : bitcasts) { - auto tbTy = dyn_cast(op.getResult().getType()); - if (!tbTy) { - op.emitError("bitcast result must be tile_buf type"); - return failure(); - } - if (isa(op->getOperand(0).getType())) - continue; - Value lowered = buildTileBufViewLikeValue(op, op->getOperand(0), tbTy, - "bitcast", ctx); - if (!lowered) - return failure(); - IRRewriter rewriter(ctx); - rewriter.replaceOp(op, lowered); - } - return success(); -} - // ============================================================================= // The Pass Implementation // ============================================================================= @@ -1861,12 +1760,8 @@ struct PTOViewToMemrefPass } // ------------------------------------------------------------------ - // Stage 1.4: Lower tile_buf view-like ops (treshape/bitcast) + // Stage 1.4: treshape/bitcast stay tile-native. // ------------------------------------------------------------------ - if (failed(lowerTileBufViewLikeOps(func, ctx))) { - signalPassFailure(); - return; - } if (failed(reconcileFusionRegionResultTypes(func))) { signalPassFailure(); return; diff --git a/test/lit/pto/bitcast_tile_native.pto b/test/lit/pto/bitcast_tile_native.pto new file mode 100644 index 0000000000..82aa068367 --- /dev/null +++ b/test/lit/pto/bitcast_tile_native.pto @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o - 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE + +module attributes {pto.target_arch = "a5"} { + func.func private @bitcast_arg( + %src: !pto.tile_buf) + -> !pto.tile_buf { + %result = pto.bitcast %src + : !pto.tile_buf + -> !pto.tile_buf + return %result : !pto.tile_buf + } +} + +// NATIVE-LABEL: func.func private @bitcast_arg( +// NATIVE-SAME: !pto.tile_buf +// NATIVE: %[[RESULT:.*]] = pto.bitcast %arg0 +// NATIVE-SAME: !pto.tile_buf -> !pto.tile_buf +// NATIVE: return %[[RESULT]] : !pto.tile_buf +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile From e9af3906281b1275b38293b09ab839f59a47cec1 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 17 Jul 2026 22:39:33 +0800 Subject: [PATCH 15/91] Keep tile valid-shape metadata native --- ...ptoas-tile-native-mainline-op-migration.md | 4 +- include/PTO/IR/PTOOps.td | 4 +- lib/PTO/IR/PTO.cpp | 58 ++++++------------- lib/PTO/Transforms/PTOToEmitC.cpp | 4 -- lib/PTO/Transforms/PTOViewToMemref.cpp | 31 +--------- test/lit/pto/validshape_tile_native.pto | 43 ++++++++++++++ 6 files changed, 68 insertions(+), 76 deletions(-) create mode 100644 test/lit/pto/validshape_tile_native.pto diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index 4a28170ea2..5cfbc362b7 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -83,8 +83,8 @@ PTO tile/view IR | `pto.subview` | 通常转成 `memref.subview + pto.bind_tile`，之后再恢复 tile | 全程保留 `pto.subview` | result 与 source 同 root；按 layout/stride/offset 计算 byte range | 现有 InsertSync 已有精确 `pto.subview` 地址计算；GSS 和 BufidSync 补齐同等能力 | EmitC/VPTO 增加直接 lowering；覆盖 row/col-major、dynamic offset、valid shape | | `pto.treshape` | 已保持 tile-native，不再生成 memref view 或 `bind_tile` | 保持原 op | result 与 source 完全 alias，size 不变；legacy/modern memplan 已传播 alias | InsertSync/GSS 已直接传播 result-to-source alias | EmitC 已直接 lowering；覆盖 tile 参数、动态 valid shape 和静态 valid shape | | `pto.bitcast` | 已保持 tile-native，不再生成 memref view 或 `bind_tile` | 保持原 op | result 与 source alias；legacy/modern memplan 已按 byte range 传播同一 root，并校验总容量一致 | InsertSync/GSS 已直接传播 alias，不按 element type 分裂 root | EmitC 已直接 lowering；覆盖 tile 参数和同容量 dtype 改变 | -| `pto.set_validshape` | 在 memref 层通过 bind metadata 记录，materialize 后恢复 tile 类型 | 直接更新/产生 tile handle metadata，不改变物理 root | result 与 source alias，物理 size 使用 allocation shape，不使用 valid shape | 同一物理范围；repeat 分析读取更新后的 valid shape | 两后端直接消费；覆盖 if/for 和动态 valid shape | -| `pto.get_validshape` | materialize 后依赖 tile handle 读取 metadata | 直接读取 tile operand metadata | 不产生 root | 无内存 effect | 两后端直接 lowering；覆盖静态折叠和动态值 | +| `pto.set_validshape` | 已保持 tile-native，输入已收紧为 `TileBufType` | 直接更新 tile handle metadata，不改变物理 root | 不产生新 root；物理 size 使用 allocation shape，不使用 valid shape | 对原 root 建模为 metadata Write | EmitC 已直接消费；覆盖 if 和动态 valid shape | +| `pto.get_validshape` | 已保持 tile-native，输入已收紧为 `TileBufType` | 直接读取 tile operand metadata | 不产生 root | 对原 root 建模为 metadata Read | EmitC 已直接 lowering；覆盖 tile 参数和动态值 | | `pto.tile_buf_addr` | ptr-like 形式转成线性 memref，后续 materialize 可能恢复 | 直接从 tile root/view 计算地址；返回 pointer-like PTO 类型 | 不产生新 root | 若结果参与 load/store，需要保留到原 root 的 provenance | EmitC/VPTO 直接生成地址表达式；覆盖 alloc、subview、multi slot | ## 6. GM Tensor View 与地址 Op diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 75a6803f76..60758484df 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -559,7 +559,7 @@ def SetValidShapeOp : PTO_Op<"set_validshape", [ }]; let arguments = (ins - TileBufOrMemRef:$source, + TileBufType:$source, Index:$valid_row, Index:$valid_col ); @@ -583,7 +583,7 @@ def GetValidShapeOp : PTO_Op<"get_validshape", [ }]; let arguments = (ins - TileBufOrMemRef:$source + TileBufType:$source ); let results = (outs diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 08d45505c7..95aac073b8 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -11149,9 +11149,6 @@ static std::optional getElemBytes(Type elemTy) { return mlir::isa(ty); } -static constexpr llvm::StringLiteral kLoweredSetValidShapeAttrName = - "__pto.lowered_set_validshape"; - static bool isLocallyBoundTileSource(Value value) { if (!value || isa(value)) return false; @@ -11191,32 +11188,22 @@ static std::optional getConstIndexLike(Value v) { mlir::LogicalResult mlir::pto::SetValidShapeOp::verify() { SmallVector shape; - if (auto srcTy = llvm::dyn_cast(getSource().getType())) { - if (srcTy.getRank() != 2) - return emitOpError("expects rank-2 tile_buf source"); + auto srcTy = getSource().getType(); + if (srcTy.getRank() != 2) + return emitOpError("expects rank-2 tile_buf source"); - ArrayRef validShape = srcTy.getValidShape(); - if (validShape.size() != 2) - return emitOpError("expects source validShape to be rank-2"); - if (!srcTy.hasDynamicValid()) - return emitOpError("expects source tile_buf to have dynamic validShape (?, ?)"); + ArrayRef validShape = srcTy.getValidShape(); + if (validShape.size() != 2) + return emitOpError("expects source validShape to be rank-2"); + if (!srcTy.hasDynamicValid()) + return emitOpError("expects source tile_buf to have dynamic validShape (?, ?)"); - shape.assign(srcTy.getShape().begin(), srcTy.getShape().end()); + shape.assign(srcTy.getShape().begin(), srcTy.getShape().end()); - if (!isLocallyBoundTileSource(getSource())) - return emitOpError( - "requires a locally bound tile source; function arguments/results " - "are unsupported"); - } else if (auto srcTy = llvm::dyn_cast(getSource().getType())) { - if (!(*this)->hasAttr(kLoweredSetValidShapeAttrName)) - return emitOpError( - "expects tile_buf source; memref source is only valid for the internal lowered form"); - if (srcTy.getRank() != 2) - return emitOpError("expects rank-2 memref source after tile lowering"); - shape.assign(srcTy.getShape().begin(), srcTy.getShape().end()); - } else { - return emitOpError("expects tile_buf source (or lowered memref source)"); - } + if (!isLocallyBoundTileSource(getSource())) + return emitOpError( + "requires a locally bound tile source; function arguments/results " + "are unsupported"); auto checkDim = [&](Value operand, unsigned dimIdx, StringRef dimName) -> LogicalResult { @@ -11243,19 +11230,12 @@ mlir::LogicalResult mlir::pto::SetValidShapeOp::verify() { } mlir::LogicalResult mlir::pto::GetValidShapeOp::verify() { - if (auto srcTy = llvm::dyn_cast(getSource().getType())) { - if (srcTy.getRank() != 2) - return emitOpError("expects rank-2 tile_buf source"); - if (srcTy.getValidShape().size() != 2) - return emitOpError("expects source validShape to be rank-2"); - return success(); - } - if (auto srcTy = llvm::dyn_cast(getSource().getType())) { - if (srcTy.getRank() != 2) - return emitOpError("expects rank-2 memref source after tile lowering"); - return success(); - } - return emitOpError("expects tile_buf source (or lowered memref source)"); + auto srcTy = getSource().getType(); + if (srcTy.getRank() != 2) + return emitOpError("expects rank-2 tile_buf source"); + if (srcTy.getValidShape().size() != 2) + return emitOpError("expects source validShape to be rank-2"); + return success(); } diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index d5f148c003..8b903cf656 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -200,10 +200,6 @@ static Location getIndexedNameHintLoc(Location fallbackLoc, unsigned index) { fallbackLoc); } -[[maybe_unused]] static constexpr llvm::StringLiteral kLoweredSetValidShapeAttrName = - "__pto.lowered_set_validshape"; -[[maybe_unused]] static constexpr llvm::StringLiteral kLoweredSetValidShapeConfigAttrName = - "__pto.lowered_set_validshape_config"; static constexpr llvm::StringLiteral kForceDynamicValidShapeAttrName = "__pto.force_dynamic_valid_shape"; static constexpr llvm::StringLiteral kGlobalTensorStridesAttrName = diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index a8bc184232..7e59f0812b 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -52,8 +52,6 @@ using namespace mlir; namespace mlir { namespace pto { -static constexpr llvm::StringLiteral kLoweredSetValidShapeAttrName = - "__pto.lowered_set_validshape"; static constexpr llvm::StringLiteral kForceDynamicValidShapeAttrName = "__pto.force_dynamic_valid_shape"; @@ -80,7 +78,8 @@ static bool hasTileNativeAllocationRoot(func::FuncOp func) { static bool hasMigratedTileNativeView(func::FuncOp func) { bool found = false; auto result = func.walk([&](Operation *op) { - if (isa(op)) { + if (isa(op)) { found = true; return WalkResult::interrupt(); } @@ -881,25 +880,6 @@ static LogicalResult reconcileFusionRegionResultTypes(func::FuncOp func) { return success(); } -static LogicalResult markLoweredSetValidShapeOps(func::FuncOp func, - MLIRContext *ctx) { - WalkResult result = func.walk([&](mlir::pto::SetValidShapeOp op) { - if (isa(op.getSource().getType())) { - if (!lookupConfig(op.getSource())) { - op.emitError( - "set_validshape requires a locally bound tile source; function " - "arguments/results are unsupported"); - return WalkResult::interrupt(); - } - op->setAttr(kLoweredSetValidShapeAttrName, UnitAttr::get(ctx)); - return WalkResult::advance(); - } - op->removeAttr(kLoweredSetValidShapeAttrName); - return WalkResult::advance(); - }); - return result.wasInterrupted() ? failure() : success(); -} - static void markForceDynamicValidShape(Operation *op, bool force, MLIRContext *ctx) { if (force) { @@ -3857,13 +3837,6 @@ struct PTOViewToMemrefPass signalPassFailure(); return; } - // Mark memref-form set_validshape only after control-flow result-type - // reconciliation. Values such as scf.if results can stay tile_buf until - // this late stage. - if (failed(markLoweredSetValidShapeOps(func, ctx))) { - signalPassFailure(); - return; - } } if (failed(bridgeCallOperandsToConvertedCallees(mod, ctx))) { diff --git a/test/lit/pto/validshape_tile_native.pto b/test/lit/pto/validshape_tile_native.pto new file mode 100644 index 0000000000..17b4609ddf --- /dev/null +++ b/test/lit/pto/validshape_tile_native.pto @@ -0,0 +1,43 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o - 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE + +module attributes {pto.target_arch = "a5"} { + func.func private @get_validshape_arg( + %src: !pto.tile_buf) { + %c0_i64 = arith.constant 0 : i64 + %c16 = arith.constant 16 : index + %local = pto.alloc_tile addr = %c0_i64 valid_row = %c16 valid_col = %c16 + : !pto.tile_buf + %row, %col = pto.get_validshape %src + : !pto.tile_buf + pto.set_validshape %local, %row, %col + : !pto.tile_buf + return + } +} + +// NATIVE-LABEL: func.func private @get_validshape_arg( +// NATIVE-SAME: !pto.tile_buf) +// NATIVE: %[[ROW:.*]], %[[COL:.*]] = pto.get_validshape %arg0 +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.set_validshape {{.*}}, %[[ROW]], %[[COL]] +// NATIVE: return +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile From 86a4205389a458aba57733cd1ae17546f345b4f9 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 17 Jul 2026 23:07:55 +0800 Subject: [PATCH 16/91] Keep tile subviews native through sync --- ...ptoas-tile-native-mainline-op-migration.md | 2 +- lib/PTO/IR/PTO.cpp | 13 +- .../Transforms/PTOMaterializeTileHandles.cpp | 61 ----- lib/PTO/Transforms/PTOResolveBufferSelect.cpp | 148 ++++++++++++ lib/PTO/Transforms/PTOViewToMemref.cpp | 228 +----------------- test/lit/pto/issue708_zero_valid_subview.pto | 4 +- .../issue708_zero_valid_subview_inferred.pto | 12 +- ...w_col_major_row_plus_one_stride_offset.pto | 4 +- ...dynamic_offset_static_valid_regression.pto | 10 +- .../subview_row_plus_one_stride_offset.pto | 4 +- 10 files changed, 186 insertions(+), 300 deletions(-) diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index 5cfbc362b7..468cf32fe8 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -80,7 +80,7 @@ PTO tile/view IR | Op | 当前行为 | Tile-native 目标 | Memplan / Alias | InsertSync | Backend 与测试 | |---|---|---|---|---|---| -| `pto.subview` | 通常转成 `memref.subview + pto.bind_tile`，之后再恢复 tile | 全程保留 `pto.subview` | result 与 source 同 root；按 layout/stride/offset 计算 byte range | 现有 InsertSync 已有精确 `pto.subview` 地址计算；GSS 和 BufidSync 补齐同等能力 | EmitC/VPTO 增加直接 lowering；覆盖 row/col-major、dynamic offset、valid shape | +| `pto.subview` | 已在 memplan 和 sync 主线保持 tile-native，不再生成 `memref.subview + pto.bind_tile` | 同步完成前保留 `pto.subview` | result 与 source 同 root；按 layout/stride/offset 计算 byte range | InsertSync 已精确计算地址范围；GSS 通过通用 view alias 回溯 source | `PTOResolveBufferSelect` 在同步后解析为带偏移地址的 `alloc_tile`，供 EmitC/VPTO 共用；覆盖 row/col-major、boxed、dynamic offset 和 valid shape | | `pto.treshape` | 已保持 tile-native，不再生成 memref view 或 `bind_tile` | 保持原 op | result 与 source 完全 alias，size 不变；legacy/modern memplan 已传播 alias | InsertSync/GSS 已直接传播 result-to-source alias | EmitC 已直接 lowering；覆盖 tile 参数、动态 valid shape 和静态 valid shape | | `pto.bitcast` | 已保持 tile-native，不再生成 memref view 或 `bind_tile` | 保持原 op | result 与 source alias；legacy/modern memplan 已按 byte range 传播同一 root，并校验总容量一致 | InsertSync/GSS 已直接传播 alias，不按 element type 分裂 root | EmitC 已直接 lowering；覆盖 tile 参数和同容量 dtype 改变 | | `pto.set_validshape` | 已保持 tile-native，输入已收紧为 `TileBufType` | 直接更新 tile handle metadata，不改变物理 root | 不产生新 root；物理 size 使用 allocation shape，不使用 valid shape | 对原 root 建模为 metadata Write | EmitC 已直接消费；覆盖 if 和动态 valid shape | diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 95aac073b8..a78d2dec07 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -6380,7 +6380,18 @@ llvm::LogicalResult mlir::pto::TCvtOp::verify() { if (failed(verifyTileBufCommon(*this, srcTy, "src", /*allowLowPrecision=*/true)) || failed(verifyTileBufCommon(*this, dstTy, "dst", /*allowLowPrecision=*/true))) return failure(); - if (failed(verifyTileBufSameLogicalExtent(*this, srcTy, dstTy, "src", "dst", + auto isResolvedSubview = [](Value value) { + auto alloc = value.getDefiningOp(); + auto semantics = + alloc ? alloc->getAttrOfType("pto.view_semantics") + : StringAttr(); + return semantics && semantics.getValue() == "subview"; + }; + // A resolved subview keeps its parent's physical shape so the generated + // Tile retains the parent stride. Its valid shape is the logical tcvt + // extent, so comparing physical shapes would reject a valid sliced tile. + if (!isResolvedSubview(getSrc()) && !isResolvedSubview(getDst()) && + failed(verifyTileBufSameLogicalExtent(*this, srcTy, dstTy, "src", "dst", /*compareValidShape=*/false))) return failure(); if (failed(verifyTileBufSameLogicalExtent(*this, srcTy, dstTy, "src", "dst", diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index 2df6dc6a4d..2a4c4bf1d1 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -998,25 +998,6 @@ static Value getAllocValidOperand(TileBufType tileTy, Value operand, return Value(); } -static TileBufType buildPTOSubViewPhysicalTileType(pto::SubViewOp op, - TileBufType resultTy) { - auto sourceTy = dyn_cast(op.getSource().getType()); - if (!sourceTy || sourceTy.getRank() != 2 || resultTy.getRank() != 2) - return resultTy; - - SmallVector physicalShape(sourceTy.getShape().begin(), - sourceTy.getShape().end()); - SmallVector validShape(resultTy.getValidShape().begin(), - resultTy.getValidShape().end()); - if (validShape.size() != 2) - validShape.assign(resultTy.getShape().begin(), resultTy.getShape().end()); - - return TileBufType::get(op.getContext(), physicalShape, - resultTy.getElementType(), - resultTy.getMemorySpace(), validShape, - resultTy.getConfigAttr()); -} - static Attribute getAttr(ArrayRef attrs, StringRef name) { for (NamedAttribute attr : attrs) { if (attr.getName().getValue() == name) @@ -1031,43 +1012,6 @@ static void copyMaterializedTileAttrs(ArrayRef attrs, to->setAttr(kForceDynamicValidShapeAttrName, attr); } -static LogicalResult -materializePTOSubViewOps(ModuleOp module, OpBuilder &builder, - DenseMap &tileHandles) { - SmallVector subviews; - module.walk([&](pto::SubViewOp op) { subviews.push_back(op); }); - - for (pto::SubViewOp op : subviews) { - auto tileTy = dyn_cast(op.getResult().getType()); - if (!tileTy) { - op.emitError("expected pto.subview result to be a tile_buf"); - return failure(); - } - - builder.setInsertionPoint(op); - Value addr = computeExplicitAddress(op.getResult(), builder, op.getLoc()); - if (!addr) { - op.emitError("cannot materialize pto.subview without a planned source " - "address"); - return failure(); - } - - TileBufType physicalTileTy = buildPTOSubViewPhysicalTileType(op, tileTy); - auto alloc = builder.create( - op.getLoc(), physicalTileTy, addr, - getAllocValidOperand(physicalTileTy, op.getValidRow(), 0, builder, - op.getLoc()), - getAllocValidOperand(physicalTileTy, op.getValidCol(), 1, builder, - op.getLoc())); - alloc->setAttr("pto.view_semantics", builder.getStringAttr("subview")); - tileHandles[op.getResult()] = alloc.getResult(); - op.replaceAllUsesWith(alloc.getResult()); - op.erase(); - } - - return success(); -} - static void updateResultTypesAfterMaterializingOperand(Operation *op, unsigned operandNo, Type tileTy) { @@ -1569,11 +1513,6 @@ struct PTOMaterializeTileHandlesPass return; } - if (failed(materializePTOSubViewOps(module, builder, tileHandles))) { - signalPassFailure(); - return; - } - for (Operation *anchor : anchors) { if (anchor->getNumResults() != 1) continue; diff --git a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp index 8bff4e5e37..d9742130d3 100644 --- a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp +++ b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp @@ -55,6 +55,150 @@ using namespace mlir; namespace { +static Value ensureI64(Value value, IRRewriter &rewriter, Location loc) { + if (!value) + return {}; + if (value.getType().isInteger(64)) + return value; + if (value.getType().isIndex()) + return rewriter.create(loc, rewriter.getI64Type(), value); + if (isa(value.getType())) + return rewriter.create(loc, rewriter.getI64Type(), value); + return {}; +} + +static bool getTilePointerStrides(pto::TileBufType type, int64_t &rowStride, + int64_t &colStride) { + auto shape = type.getShape(); + if (shape.size() != 2 || llvm::is_contained(shape, ShapedType::kDynamic)) + return false; + + auto config = type.getConfigAttr(); + int32_t bl = static_cast(config.getBLayout().getValue()); + int32_t sl = static_cast(config.getSLayout().getValue()); + if (sl == 0) { + bool rowPlusOne = + type.getCompactModeI32() == + static_cast(pto::CompactMode::RowPlusOne); + rowStride = bl == 1 ? 1 : shape[1] + (rowPlusOne ? 1 : 0); + colStride = bl == 1 ? shape[0] + (rowPlusOne ? 1 : 0) : 1; + return true; + } + + unsigned elemBytes = pto::getPTOStorageElemByteSize(type.getElementType()); + if (elemBytes == 0) + return false; + int64_t innerRows = 1; + int64_t innerCols = 1; + int32_t fractal = config.getSFractalSize().getInt(); + if (fractal == 1024) { + innerRows = 16; + innerCols = 16; + } else if (fractal == 32) { + innerRows = 16; + innerCols = 2; + } else if (fractal == 512 && sl == 1) { + innerRows = 16; + innerCols = 32 / elemBytes; + } else if (fractal == 512 && sl == 2) { + innerRows = 32 / elemBytes; + innerCols = 16; + } else { + return false; + } + + if (bl == 1) { + if (sl != 1) + return false; + rowStride = innerCols; + colStride = shape[0]; + } else { + rowStride = shape[1]; + colStride = innerRows; + } + return true; +} + +static Value computeTileAddress(Value value, IRRewriter &rewriter, + Location loc) { + if (auto alloc = value.getDefiningOp()) + return ensureI64(alloc.getAddr(), rewriter, loc); + if (auto subview = value.getDefiningOp()) { + Value base = computeTileAddress(subview.getSource(), rewriter, loc); + auto sourceType = subview.getSource().getType(); + int64_t rowStride = 0; + int64_t colStride = 0; + if (!base || !getTilePointerStrides(sourceType, rowStride, colStride) || + subview.getOffsets().size() != 2) + return {}; + Value row = ensureI64(subview.getOffsets()[0], rewriter, loc); + Value col = ensureI64(subview.getOffsets()[1], rewriter, loc); + if (!row || !col) + return {}; + Value rowScale = rewriter.create(loc, rowStride, 64); + Value colScale = rewriter.create(loc, colStride, 64); + row = rewriter.create(loc, row, rowScale); + col = rewriter.create(loc, col, colScale); + Value elements = rewriter.create(loc, row, col); + int64_t elemBytes = static_cast( + pto::getPTOStorageElemByteSize(sourceType.getElementType())); + if (elemBytes == 0) + return {}; + Value byteScale = rewriter.create(loc, elemBytes, 64); + Value bytes = rewriter.create(loc, elements, byteScale); + return rewriter.create(loc, base, bytes); + } + return {}; +} + +static pto::TileBufType getSubviewPhysicalType(pto::SubViewOp op) { + pto::TileBufType sourceType = op.getSource().getType(); + pto::TileBufType resultType = op.getResult().getType(); + return pto::TileBufType::get( + op.getContext(), sourceType.getShape(), resultType.getElementType(), + resultType.getMemorySpace(), resultType.getValidShape(), + resultType.getConfigAttr()); +} + +static Value getSubviewValidOperand(pto::SubViewOp op, + pto::TileBufType physicalType, + unsigned dim, IRRewriter &rewriter) { + Value operand = dim == 0 ? op.getValidRow() : op.getValidCol(); + ArrayRef validShape = physicalType.getValidShape(); + if (validShape.size() <= dim || validShape[dim] >= 0) + return {}; + if (operand) + return operand; + ArrayRef shape = physicalType.getShape(); + if (shape.size() > dim && shape[dim] != ShapedType::kDynamic) + return rewriter.create(op.getLoc(), shape[dim]); + return {}; +} + +static LogicalResult resolveTileNativeSubviews(ModuleOp module, + MLIRContext *ctx) { + SmallVector subviews; + module.walk([&](pto::SubViewOp op) { subviews.push_back(op); }); + for (pto::SubViewOp op : subviews) { + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(op); + Value addr = computeTileAddress(op.getResult(), rewriter, op.getLoc()); + // A tile function argument is a symbolic runtime-bound handle. Keep its + // subview tile-native; only planned local roots can be normalized to an + // addressed alloc_tile here. + if (!addr) + continue; + pto::TileBufType physicalType = getSubviewPhysicalType(op); + auto alloc = rewriter.create( + op.getLoc(), physicalType, addr, + getSubviewValidOperand(op, physicalType, 0, rewriter), + getSubviewValidOperand(op, physicalType, 1, rewriter)); + alloc->setAttr("pto.view_semantics", rewriter.getStringAttr("subview")); + rewriter.replaceOp(op, alloc.getResult()); + } + return success(); +} + static FailureOr getStaticSlotBytes(pto::TileBufType slotType) { uint64_t elemBytes = pto::getPTOStorageElemByteSize(slotType.getElementType()); if (elemBytes == 0) @@ -217,6 +361,10 @@ struct PTOResolveBufferSelectPass signalPassFailure(); return; } + if (failed(resolveTileNativeSubviews(mod, ctx))) { + signalPassFailure(); + return; + } SmallVector markers; mod.walk([&](pto::SlotMarkerOp op) { markers.push_back(op); }); diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 7e59f0812b..bcc0cfadee 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -79,7 +79,7 @@ static bool hasMigratedTileNativeView(func::FuncOp func) { bool found = false; auto result = func.walk([&](Operation *op) { if (isa(op)) { + pto::GetValidShapeOp, pto::SubViewOp>(op)) { found = true; return WalkResult::interrupt(); } @@ -133,58 +133,6 @@ using SmallInlineVector = SmallVector; template using DefaultInlineVector = SmallVector; -// ============================================================================= -// Helper: Metadata Backtracking (核心机制) -// ============================================================================= -// 从一个 MemRef Value 向上回溯，找到它绑定的 TileBufConfig。 -// 这解决了 "Type Erasure" 问题：memref 类型本身不包含 config，但 SSA 定义链包含。 -static mlir::pto::TileBufConfigAttr lookupConfig(Value v) { - // 1. 最直接的情况：它就是 bind_tile 的结果 - if (auto bind = v.getDefiningOp()) { - return bind.getConfig(); - } - // PointerCastOp can also carry tile metadata (used when alloc_tile specifies - // an explicit address). - if (auto pc = v.getDefiningOp()) { - if (auto cfg = pc.getConfig()) - return *cfg; - return {}; - } - - // 2. 穿透 View 操作 (SubView, Cast 等) 向上查找 - if (auto subview = v.getDefiningOp()) { - return lookupConfig(subview.getSource()); - } - if (auto cast = v.getDefiningOp()) { - return lookupConfig(cast.getSource()); - } - if (auto cast = v.getDefiningOp()) { - return lookupConfig(cast.getSource()); - } - if (auto slot = v.getDefiningOp()) { - return lookupConfig(slot.getSource()); - } - - // 3. pto.fusion_region result 本身不携带 config；作为兜底情况，沿着 - // pto.yield 回溯到 region 内真正的 tile handle/memref 定义链继续查找。 - if (auto regionResult = dyn_cast(v)) { - if (auto fusionRegion = - dyn_cast(regionResult.getOwner())) { - auto yieldOp = dyn_cast( - fusionRegion.getBody().front().getTerminator()); - if (!yieldOp) - return {}; - unsigned resultIndex = regionResult.getResultNumber(); - if (resultIndex >= yieldOp.getNumOperands()) - return {}; - return lookupConfig(yieldOp.getOperand(resultIndex)); - } - } - - // 如果追溯到 BlockArgument (函数参数) 或其他无法穿透的 Op，则返回空 - return {}; -} - // ============================================================================= // Helper: Valid dims backtracking (v_row / v_col) // ============================================================================= @@ -618,19 +566,6 @@ static SmallVector computeCompactStrides(ArrayRef shape) { return strides; } -static bool checkMultipleOf(Operation *op, int64_t value, int64_t divisor, - StringRef label) { - if (divisor <= 0) { - op->emitError("boxed layout requires positive divisor for ") << label; - return false; - } - if (value % divisor == 0) - return true; - op->emitError("boxed layout requires ") - << label << " multiple of " << divisor << ", got " << value; - return false; -} - // 确保 Value 是 Index 类型 static Value ensureIndex(IRRewriter &rewriter, Location loc, Value v, Operation *anchorOp) { @@ -680,31 +615,6 @@ static bool foldAddPtrChainIntoOffset(IRRewriter &rewriter, Location loc, return folded; } -static Value clampSubViewValidDim(IRRewriter &rewriter, Location loc, - Value explicitValid, int64_t size, - int64_t inferredValid, Operation *anchorOp) { - if (!explicitValid) { - // No explicit valid operand: take the valid extent the result type - // declares. For an ordinary subview this equals `size`; for an empty - // tail/no-op-replay tile the type carries 0, which must survive to - // bind_tile rather than being widened back to `size`. A dynamic declared - // extent is materialized via markForceDynamicValidShape, so fall back to - // `size` here. - int64_t fallback = inferredValid >= 0 ? inferredValid : size; - return rewriter.create(loc, fallback); - } - - Value sizeVal = rewriter.create(loc, size); - int64_t cst = 0; - if (getConstIndexValue(explicitValid, cst)) - return rewriter.create(loc, std::min(cst, size)); - - Value v = ensureIndex(rewriter, loc, explicitValid, anchorOp); - Value lt = rewriter.create(loc, arith::CmpIPredicate::slt, v, - sizeVal); - return rewriter.create(loc, lt, v, sizeVal); -} - [[maybe_unused]] static void dumpPretty(Operation *op, llvm::raw_ostream &os) { OpPrintingFlags flags; flags.useLocalScope(); @@ -1325,136 +1235,6 @@ static LogicalResult lowerPartitionViewOps(func::FuncOp func, MLIRContext *ctx) return success(); } -static LogicalResult lowerSubViewOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector subViews; - func.walk([&](mlir::pto::SubViewOp op) { subViews.push_back(op); }); - - for (auto op : subViews) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - auto resultTileTy = - dyn_cast(op.getResult().getType()); - Value src = op->getOperand(0); - auto srcMrTy = dyn_cast(src.getType()); - if (!srcMrTy) { - // Tile-native alloc_tile roots intentionally survive this pass so that - // PTOPlanMemory can assign their addresses. Defer tile_buf subviews until - // after planning, when PTOMaterializeTileHandles can materialize the - // adjusted address directly. - if (isa(src.getType())) - continue; - op.emitError("pto.subview source must be lowered to memref first"); - return failure(); - } - - ArrayAttr sizeAttr = op.getSizes(); - SmallVector staticSizes; - SmallVector mixedSizes; - for (Attribute attr : sizeAttr) { - int64_t size = cast(attr).getInt(); - staticSizes.push_back(size); - mixedSizes.push_back(rewriter.getIndexAttr(size)); - } - - SmallVector mixedOffsets; - for (Value offset : op.getOffsets()) { - appendMixedIndex(rewriter, loc, offset, op, mixedOffsets); - } - - auto configAttr = lookupConfig(src); - if (!configAttr) - configAttr = pto::TileBufConfigAttr::getDefault(ctx); - - TileLayoutInfo layoutInfo; - if (!computeTileLayoutInfo(configAttr, srcMrTy.getElementType(), - srcMrTy.getShape(), layoutInfo)) { - op.emitError("unsupported tile layout for pto.subview"); - return failure(); - } - - if (layoutInfo.boxed) { - if (staticSizes.size() != kTileRank2D || - op.getOffsets().size() != kTileRank2D) { - op.emitError("boxed layout subview expects 2D sizes/offsets"); - return failure(); - } - if (!checkMultipleOf(op, staticSizes[0], layoutInfo.innerRows, "row size") || - !checkMultipleOf(op, staticSizes[1], layoutInfo.innerCols, "col size")) { - return failure(); - } - - int64_t off0 = 0; - int64_t off1 = 0; - bool off0Const = getConstIndexValue(op.getOffsets()[0], off0); - bool off1Const = getConstIndexValue(op.getOffsets()[1], off1); - if (off0Const && - !checkMultipleOf(op, off0, layoutInfo.innerRows, "row offset")) { - return failure(); - } - if (off1Const && - !checkMultipleOf(op, off1, layoutInfo.innerCols, "col offset")) { - return failure(); - } - - } - - SmallVector srcStrides; - int64_t srcOffset = ShapedType::kDynamic; - if (failed(mlir::pto::getPTOMemRefStridesAndOffset(srcMrTy, srcStrides, - srcOffset))) - srcStrides = computeCompactStrides(srcMrTy.getShape()); - - // Keep parent physical shape + strides for bound tile semantics. - auto resultLayout = - StridedLayoutAttr::get(ctx, ShapedType::kDynamic, srcStrides); - auto parentShape = srcMrTy.getShape(); - auto resultMemRefType = - MemRefType::get(parentShape, srcMrTy.getElementType(), resultLayout, - srcMrTy.getMemorySpace()); - - // Intermediate memref.subview keeps logical subview size. - auto subViewMemRefType = - MemRefType::get(staticSizes, srcMrTy.getElementType(), resultLayout, - srcMrTy.getMemorySpace()); - - SmallVector mixedStrides(staticSizes.size(), - rewriter.getIndexAttr(1)); - auto sv = rewriter.create(loc, subViewMemRefType, src, - mixedOffsets, mixedSizes, - mixedStrides); - - // When a valid operand is omitted, fall back to the extent the result type - // declares (which the verifier pins to either `sizes` or an empty 0 marker) - // rather than the physical subview size, so a no-op-replay v_row/v_col=0 - // survives lowering. - ArrayRef resultValid = - resultTileTy ? resultTileTy.getValidShape() : ArrayRef{}; - auto inferredValidDim = [&](unsigned d) -> int64_t { - return d < resultValid.size() ? resultValid[d] : ShapedType::kDynamic; - }; - - Value vRow; - Value vCol; - if (!staticSizes.empty()) - vRow = clampSubViewValidDim(rewriter, loc, op.getValidRow(), - staticSizes[0], inferredValidDim(0), op); - if (staticSizes.size() > 1) - vCol = clampSubViewValidDim(rewriter, loc, op.getValidCol(), - staticSizes[1], inferredValidDim(1), op); - - auto bindOp = rewriter.create( - loc, resultMemRefType, sv.getResult(), vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(bindOp, - resultTileTy && resultTileTy.hasDynamicValid(), - ctx); - bindOp->setAttr("pto.view_semantics", rewriter.getStringAttr("subview")); - rewriter.replaceOp(op, bindOp.getResult()); - } - return success(); -} - // ============================================================================= // The Pass Implementation // ============================================================================= @@ -1732,12 +1512,8 @@ struct PTOViewToMemrefPass } // ------------------------------------------------------------------ - // Stage 1.35: Lower pto.subview -> memref.subview + pto.bind_tile + // Stage 1.35: pto.subview stays tile-native through memory and sync. // ------------------------------------------------------------------ - if (failed(lowerSubViewOps(func, ctx))) { - signalPassFailure(); - return; - } // ------------------------------------------------------------------ // Stage 1.4: treshape/bitcast stay tile-native. diff --git a/test/lit/pto/issue708_zero_valid_subview.pto b/test/lit/pto/issue708_zero_valid_subview.pto index 1c70846404..b4c7be1bc1 100644 --- a/test/lit/pto/issue708_zero_valid_subview.pto +++ b/test/lit/pto/issue708_zero_valid_subview.pto @@ -46,9 +46,9 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { } // CHECK-LABEL: func.func @issue708_zero_valid_subview_row -// CHECK: %[[ROW:[0-9A-Za-z_]+]] = pto.subview {{.*}} valid [%c0, %c128] +// CHECK: %[[ROW:[0-9A-Za-z_]+]] = pto.alloc_tile addr = %{{.*}} {pto.view_semantics = "subview"} : !pto.tile_buf // CHECK: pto.tstore ins(%[[ROW]] // CHECK-LABEL: func.func @issue708_zero_valid_subview_col -// CHECK: %[[COL:[0-9A-Za-z_]+]] = pto.subview {{.*}} valid [%c5, %c0] +// CHECK: %[[COL:[0-9A-Za-z_]+]] = pto.alloc_tile addr = %{{.*}} {pto.view_semantics = "subview"} : !pto.tile_buf // CHECK: pto.tstore ins(%[[COL]] diff --git a/test/lit/pto/issue708_zero_valid_subview_inferred.pto b/test/lit/pto/issue708_zero_valid_subview_inferred.pto index 38f1893b20..faa4df0930 100644 --- a/test/lit/pto/issue708_zero_valid_subview_inferred.pto +++ b/test/lit/pto/issue708_zero_valid_subview_inferred.pto @@ -91,21 +91,21 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { } } -// The bind_tile that replaces each pto.subview must carry valid operands taken -// from the result type, i.e. 0 for an empty axis -- never the subview size. +// The resolved subview handle must carry valid extents taken from the result +// type, i.e. 0 for an empty axis -- never the subview size. // CHECK-LABEL: func.func @issue708_inferred_empty_store -// CHECK: %[[BOTH:[0-9A-Za-z_]+]] = pto.subview {{.*}} -> !pto.tile_buf +// CHECK: %[[BOTH:[0-9A-Za-z_]+]] = pto.alloc_tile addr = %{{.*}} {pto.view_semantics = "subview"} : !pto.tile_buf // CHECK: pto.tstore ins(%[[BOTH]] // CHECK-LABEL: func.func @issue708_inferred_empty_cvt -// CHECK: %[[CVTV:[0-9A-Za-z_]+]] = pto.subview {{.*}} -> !pto.tile_buf +// CHECK: %[[CVTV:[0-9A-Za-z_]+]] = pto.alloc_tile addr = %{{.*}} {pto.view_semantics = "subview"} : !pto.tile_buf // CHECK: pto.tcvt ins(%[[CVTV]] // CHECK-LABEL: func.func @issue708_inferred_row_empty -// CHECK: %[[ROW:[0-9A-Za-z_]+]] = pto.subview {{.*}} -> !pto.tile_buf +// CHECK: %[[ROW:[0-9A-Za-z_]+]] = pto.alloc_tile addr = %{{.*}} {pto.view_semantics = "subview"} : !pto.tile_buf // CHECK: pto.tstore ins(%[[ROW]] // CHECK-LABEL: func.func @issue708_inferred_partial_valid -// CHECK: %[[PARTIAL:[0-9A-Za-z_]+]] = pto.subview {{.*}} -> !pto.tile_buf +// CHECK: %[[PARTIAL:[0-9A-Za-z_]+]] = pto.alloc_tile addr = %{{.*}} {pto.view_semantics = "subview"} : !pto.tile_buf // CHECK: pto.tstore ins(%[[PARTIAL]] diff --git a/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto b/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto index 9911993b05..847484b8ca 100644 --- a/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto +++ b/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto @@ -22,4 +22,6 @@ module { // ColMajor + RowPlusOne: major stride should be 17 (not 16). // CHECK: pto.alloc_tile addr = %c0_i64 : !pto.tile_buf -// CHECK: pto.subview {{.*}} sizes [8, 8] : !pto.tile_buf -> !pto.tile_buf +// CHECK: %c17_i64 = arith.constant 17 : i64 +// CHECK: arith.muli %{{.*}}, %c17_i64 : i64 +// CHECK: pto.alloc_tile addr = %{{.*}} {pto.view_semantics = "subview"} : !pto.tile_buf diff --git a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto index f2c1979946..f9f92a74b3 100644 --- a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto +++ b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto @@ -1,4 +1,5 @@ // RUN: ptoas --emit-pto-ir %s -o - | FileCheck %s +// RUN: ptoas %s --mlir-print-ir-after=pto-view-to-memref -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE module { func.func @subview_dynamic_offset_static_valid_regression( @@ -24,4 +25,11 @@ module { } // CHECK: func.func @subview_dynamic_offset_static_valid_regression -// CHECK: pto.subview %{{.*}}[%c0, %{{.*}}] sizes [1, 64] : !pto.tile_buf -> !pto.tile_buf +// CHECK: %[[BYTE_OFFSET:.*]] = arith.muli %{{.*}}, %c4_i64 : i64 +// CHECK: %[[ADDR:.*]] = arith.addi %c0_i64, %[[BYTE_OFFSET]] : i64 +// CHECK: pto.alloc_tile addr = %[[ADDR]] {pto.view_semantics = "subview"} : !pto.tile_buf + +// NATIVE: IR Dump After PTOViewToMemref +// NATIVE: pto.subview %{{.*}}[%c0, %{{.*}}] sizes [1, 64] : !pto.tile_buf -> !pto.tile_buf +// NATIVE-NOT: memref.subview +// NATIVE-NOT: pto.bind_tile diff --git a/test/lit/pto/subview_row_plus_one_stride_offset.pto b/test/lit/pto/subview_row_plus_one_stride_offset.pto index 02e9b11ded..80d1625663 100644 --- a/test/lit/pto/subview_row_plus_one_stride_offset.pto +++ b/test/lit/pto/subview_row_plus_one_stride_offset.pto @@ -22,4 +22,6 @@ module { // RowPlusOne: major stride should be 17 (not 16). // CHECK: pto.alloc_tile addr = %c0_i64 : !pto.tile_buf -// CHECK: pto.subview {{.*}} sizes [8, 8] : !pto.tile_buf -> !pto.tile_buf +// CHECK: %c17_i64 = arith.constant 17 : i64 +// CHECK: arith.muli %{{.*}}, %c17_i64 : i64 +// CHECK: pto.alloc_tile addr = %{{.*}} {pto.view_semantics = "subview"} : !pto.tile_buf From dfb80fb917af131a67409dcecfc6678a0730c623 Mon Sep 17 00:00:00 2001 From: FangRui Date: Fri, 17 Jul 2026 23:24:34 +0800 Subject: [PATCH 17/91] Keep tile address extraction native --- ...ptoas-tile-native-mainline-op-migration.md | 2 +- include/PTO/IR/VPTOOps.td | 13 +++---- lib/PTO/Transforms/PTOToEmitC.cpp | 13 ++----- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- ...lize_tile_handles_subkernel_helper_abi.pto | 13 ++++--- .../lit/pto/tile_buf_addr_tile_native_arg.pto | 39 +++++++++++++++++++ 6 files changed, 58 insertions(+), 24 deletions(-) create mode 100644 test/lit/pto/tile_buf_addr_tile_native_arg.pto diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index 468cf32fe8..5b0912e3bd 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -85,7 +85,7 @@ PTO tile/view IR | `pto.bitcast` | 已保持 tile-native，不再生成 memref view 或 `bind_tile` | 保持原 op | result 与 source alias；legacy/modern memplan 已按 byte range 传播同一 root，并校验总容量一致 | InsertSync/GSS 已直接传播 alias，不按 element type 分裂 root | EmitC 已直接 lowering；覆盖 tile 参数和同容量 dtype 改变 | | `pto.set_validshape` | 已保持 tile-native，输入已收紧为 `TileBufType` | 直接更新 tile handle metadata，不改变物理 root | 不产生新 root；物理 size 使用 allocation shape，不使用 valid shape | 对原 root 建模为 metadata Write | EmitC 已直接消费；覆盖 if 和动态 valid shape | | `pto.get_validshape` | 已保持 tile-native，输入已收紧为 `TileBufType` | 直接读取 tile operand metadata | 不产生 root | 对原 root 建模为 metadata Read | EmitC 已直接 lowering；覆盖 tile 参数和动态值 | -| `pto.tile_buf_addr` | ptr-like 形式转成线性 memref，后续 materialize 可能恢复 | 直接从 tile root/view 计算地址；返回 pointer-like PTO 类型 | 不产生新 root | 若结果参与 load/store，需要保留到原 root 的 provenance | EmitC/VPTO 直接生成地址表达式；覆盖 alloc、subview、multi slot | +| `pto.tile_buf_addr` | tile-native 输入已保持原 op；仅 legacy memref 输入仍走线性 memref 兼容路径 | 直接从 tile root/view 计算地址；返回 pointer-like PTO 类型 | 不产生新 root；legacy memplan 将 source 记录为 use | 包含该 op 的函数保持 tile ABI，地址结果保留 source provenance | EmitC 直接生成 `tile.data()`，VPTO pointer normalize 保持 typed pointer；覆盖 tile 参数和 alloc root | ## 6. GM Tensor View 与地址 Op diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index 757034284e..603a18dc51 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -220,11 +220,10 @@ def TensorViewAddrOp : PTO_Op<"tensor_view_addr", [Pure]> { typed PTO pointer, depending on the requested destination type. In authoring-form IR this op preserves the descriptor-style surface; - during view-to-memref lowering it collapses to the underlying memref value - or to a memref-derived pointer. + tile-native inputs remain unchanged through memory planning and sync. - This op may also accept a memref operand after earlier view lowering, in - which case it behaves as an identity marker and is removed by lowering. + This op may also accept a legacy memref operand, in which case it behaves + as an identity marker and is removed by compatibility lowering. }]; let arguments = (ins AnyTypeOf<[TensorViewType, PartitionTensorViewType, AnyMemRef], @@ -247,9 +246,9 @@ def TileBufAddrOp : PTO_Op<"tile_buf_addr", [Pure]> { shape and address space. This op is emitted by TileLang DSL templates and resolved by the - FoldTileBufIntrinsics pass after inlining. Hand-written `.pto` may also - use it directly on the memref result of `pto.bind_tile` / lowered - `pto.alloc_tile`. + FoldTileBufIntrinsics pass after inlining. Tile-native uses remain in PTO + IR through memory planning and sync. Hand-written legacy `.pto` may also + use it directly on the memref result of `pto.bind_tile`. }]; let arguments = (ins AnyTypeOf<[TileBufType, AnyMemRef], diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 8b903cf656..5f369b9054 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -13302,23 +13302,18 @@ struct PTOTileBufAddrToEmitC : public OpConversionPattern { LogicalResult matchAndRewrite(pto::TileBufAddrOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - auto ptrTy = dyn_cast(op.getResult().getType()); - if (!ptrTy) - return failure(); - Value src = peelUnrealized(adaptor.getSrc()); + Type dstTy = getTypeConverter()->convertType(op.getResult().getType()); + if (!dstTy) + return failure(); if (isEmitCTileLikeType(src.getType())) { rewriter.replaceOpWithNewOp( - op, - TypeRange{getTypeConverter()->convertType(op.getResult().getType())}, + op, TypeRange{dstTy}, "PTOAS__TILE_DATA", ArrayAttr{}, ArrayAttr{}, ValueRange{src}); return success(); } - Type dstTy = getTypeConverter()->convertType(op.getResult().getType()); - if (!dstTy) - return failure(); rewriter.replaceOpWithNewOp(op, dstTy, src); return success(); } diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index bcc0cfadee..835ec43814 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -79,7 +79,7 @@ static bool hasMigratedTileNativeView(func::FuncOp func) { bool found = false; auto result = func.walk([&](Operation *op) { if (isa(op)) { + pto::GetValidShapeOp, pto::SubViewOp, pto::TileBufAddrOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto b/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto index 0fe4c0e375..c5c62816a6 100644 --- a/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto +++ b/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto @@ -6,8 +6,10 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// Guards PTODSL subkernel helper arguments restored to tile_buf ABI. +// Guards PTODSL subkernel helper arguments staying in tile_buf ABI when the +// helper directly extracts their address views. // RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --pto-level=level3 %s -o /dev/null module { func.func @ptodsl_subkernel_helper_tile_abi( @@ -37,12 +39,11 @@ module { } // IR: IR Dump After PTOMaterializeTileHandles -// IR-LABEL: func.func private @ptodsl_subkernel_helper_tile_abi( +// IR-LABEL: func.func @ptodsl_subkernel_helper_tile_abi( // IR-SAME: %{{.*}}: !pto.tile_buf // IR-SAME: %{{.*}}: !pto.tile_buf // IR-SAME: %{{.*}}: f32 // IR-SAME: attributes -// IR: builtin.unrealized_conversion_cast %{{.*}} : !pto.tile_buf to memref<16x1xf32, strided<[1, 16], offset: ?>, #pto.address_space> -// IR: builtin.unrealized_conversion_cast %{{.*}} : !pto.tile_buf to memref<16x16xf16, strided<[16, 1], offset: ?>, #pto.address_space> -// IR: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -// IR: pto.tile_buf_addr %{{.*}} : !pto.tile_buf +// IR: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> memref<16x16xf16, #pto.address_space> +// IR: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> memref<16x1xf32, #pto.address_space> +// IR-NOT: builtin.unrealized_conversion_cast diff --git a/test/lit/pto/tile_buf_addr_tile_native_arg.pto b/test/lit/pto/tile_buf_addr_tile_native_arg.pto new file mode 100644 index 0000000000..76588d3ed7 --- /dev/null +++ b/test/lit/pto/tile_buf_addr_tile_native_arg.pto @@ -0,0 +1,39 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-level=level3 %s --mlir-print-ir-after=pto-view-to-memref -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a5 --pto-level=level3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @tile_buf_addr_tile_native_arg( + %tile: !pto.tile_buf) { + %ptr = pto.tile_buf_addr %tile + : !pto.tile_buf -> !pto.ptr + func.call @consume(%ptr) : (!pto.ptr) -> () + return + } + + func.func private @consume(!pto.ptr) +} + +// NATIVE: IR Dump After PTOViewToMemref +// NATIVE-LABEL: func.func @tile_buf_addr_tile_native_arg( +// NATIVE-SAME: %[[TILE:.*]]: !pto.tile_buf +// NATIVE: %[[PTR:.*]] = pto.tile_buf_addr %[[TILE]] : !pto.tile_buf -> !pto.ptr +// NATIVE: call @consume(%[[PTR]]) : (!pto.ptr) -> () +// NATIVE-NOT: memref.reinterpret_cast +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: AICORE void tile_buf_addr_tile_native_arg( +// EMITC: __ubuf__ float* {{.*}} = {{.*}}.data(); +// EMITC: consume( + +// VPTO-LABEL: func.func @tile_buf_addr_tile_native_arg( +// VPTO-SAME: !pto.tile_buf +// VPTO: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> !pto.ptr From c321132e60f521e90d88ddeb55af745aac209e45 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 10:59:01 +0800 Subject: [PATCH 18/91] Keep GM tensor views native through sync --- ...ptoas-tile-native-mainline-op-migration.md | 8 +- lib/PTO/IR/PTO.cpp | 18 +- lib/PTO/Transforms/PTOResolveBufferSelect.cpp | 170 ++++++++++++++++++ lib/PTO/Transforms/PTOViewToMemref.cpp | 13 +- .../gm_tensor_view_native_through_sync.pto | 47 +++++ 5 files changed, 243 insertions(+), 13 deletions(-) create mode 100644 test/lit/pto/gm_tensor_view_native_through_sync.pto diff --git a/docs/designs/ptoas-tile-native-mainline-op-migration.md b/docs/designs/ptoas-tile-native-mainline-op-migration.md index 5b0912e3bd..b66a7f5c68 100644 --- a/docs/designs/ptoas-tile-native-mainline-op-migration.md +++ b/docs/designs/ptoas-tile-native-mainline-op-migration.md @@ -91,10 +91,10 @@ PTO tile/view IR | Op | 当前行为 | Tile-native 目标 | Memplan / Sync | Backend 与测试 | |---|---|---|---|---| -| `pto.make_tensor_view` | 转成 `memref.reinterpret_cast` | 保持 PTO tensor view 到 backend | 不参与 local memplan；InsertSync 已能把 result alias 到 pointer source | EmitC 增加完整 direct pattern；VPTO 保持 canonical rank/layout lowering；覆盖 addptr base、dynamic shape/stride | -| `pto.partition_view` | 转成 `memref.subview` | 保持 PTO partition view | 不参与 local memplan；sync result alias source，并根据 offset/size缩小 GM range | EmitC 已有部分 direct pattern，补齐全部 layout；VPTO 增加直接支持 | -| `pto.get_tensor_view_dim` | 转成 `memref.dim` | 直接读取 PTO view shape | 不影响 memplan/sync | backend 支持静态折叠和动态 dim | -| `pto.get_tensor_view_stride` | 转成 strided memref metadata | 直接读取 PTO view stride | 不影响 memplan/sync | backend 支持静态折叠和动态 stride | +| `pto.make_tensor_view` | 已保持 PTO 形态穿过 `PTOViewToMemref`、memplan 和 sync | 最终保持到 backend | 不参与 local memplan；InsertSync 已把 result alias 到 pointer source | 当前在 `PTOResolveBufferSelect` 与完整 GM view 链一起转成 `memref.reinterpret_cast`，复用成熟 backend lowering | +| `pto.partition_view` | 已保持 PTO 形态穿过 memplan 和 sync | 最终保持到 backend | 不参与 local memplan；sync result alias source，并根据 offset/size 缩小 GM range | memref-backed source 在 `PTOResolveBufferSelect` 转成 `memref.subview`；`declare_global` 等运行时 source 保持 PTO op 走已有直接 EmitC lowering | +| `pto.get_tensor_view_dim` | 已保持 PTO 形态穿过 memplan 和 sync | 直接读取 PTO view shape | 不影响 memplan/sync | 当前在 `PTOResolveBufferSelect` 随所属 view 转成 `memref.dim` | +| `pto.get_tensor_view_stride` | 已保持 PTO 形态穿过 memplan 和 sync | 直接读取 PTO view stride | 不影响 memplan/sync | 当前在 `PTOResolveBufferSelect` 随所属 view 转成 strided memref metadata | | `pto.inttoptr` | 结果改成 GM memref，并限制用途 | 保持 PTO pointer-like value | 不参与 local memplan；sync 将其视作 GM provenance | 保留 restricted-use verifier；EmitC/VPTO 直接 lowering | | `pto.ptrtoint` | 折叠 `addptr` 链并生成 byte offset | 保持 PTO op，或迁移到独立 address-canonicalization pass | 不影响 local memplan；不能丢失 GM provenance | backend 生成整数地址；覆盖 ptr、addptr、view base | | `pto.addptr` | 折叠进 tensor view、scalar load/store 或 pipe init | 保持到独立 address canonicalization/backend | sync 需要把 result alias 到 base，并记录 offset | EmitC/VPTO 直接 lowering；保留非法 escape 校验 | diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index a78d2dec07..30000d78e4 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -2408,13 +2408,19 @@ LogicalResult mlir::pto::MakeTensorViewOp::verify() { if (!tvTy) return emitOpError("result must be pto.tensor_view<...>"); - auto pty = dyn_cast(getPtr().getType()); - if (!pty) - return emitOpError("ptr operand must be !pto.ptr<...>"); + Type ptrElemTy; + if (auto pty = dyn_cast(getPtr().getType())) + ptrElemTy = pty.getElementType(); + else if (auto memrefTy = dyn_cast(getPtr().getType())) + ptrElemTy = memrefTy.getElementType(); + else + return emitOpError( + "ptr operand must be !pto.ptr<...> or a memref-backed pointer"); - if (pty.getElementType() != tvTy.getElementType()) - return emitOpError() << "ptr element type must match tensor_view element type, but got ptr=" - << pty.getElementType() << " view=" << tvTy.getElementType(); + if (ptrElemTy != tvTy.getElementType()) + return emitOpError() << "ptr element type must match tensor_view element " + "type, but got ptr=" + << ptrElemTy << " view=" << tvTy.getElementType(); int64_t rank = tvTy.getRank(); diff --git a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp index d9742130d3..e6ef8ae590 100644 --- a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp +++ b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp @@ -36,6 +36,7 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/Matchers.h" @@ -199,6 +200,171 @@ static LogicalResult resolveTileNativeSubviews(ModuleOp module, return success(); } +static LogicalResult reconcileSCFViewResultTypes(ModuleOp module) { + SmallVector ifOps; + module.walk([&](scf::IfOp op) { ifOps.push_back(op); }); + for (scf::IfOp op : llvm::reverse(ifOps)) { + if (op.getNumResults() == 0) + continue; + auto thenYield = dyn_cast(op.thenBlock()->getTerminator()); + auto elseYield = dyn_cast(op.elseBlock()->getTerminator()); + if (!thenYield || !elseYield || + thenYield.getNumOperands() != op.getNumResults() || + elseYield.getNumOperands() != op.getNumResults()) + return op.emitError("cannot reconcile scf.if view result types"); + for (unsigned i = 0; i < op.getNumResults(); ++i) { + Type thenType = thenYield.getOperand(i).getType(); + if (thenType != elseYield.getOperand(i).getType()) + return op.emitError("scf.if branches yield different view types"); + op.getResult(i).setType(thenType); + } + } + return success(); +} + +static LogicalResult resolveGMTensorViews(ModuleOp module, MLIRContext *ctx) { + SmallVector makeViews; + module.walk([&](pto::MakeTensorViewOp op) { makeViews.push_back(op); }); + DenseMap loweredViews; + for (pto::MakeTensorViewOp op : makeViews) { + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(op); + Value base = op.getPtr(); + OpFoldResult offset = rewriter.getIndexAttr(0); + Value totalOffset; + while (auto add = base.getDefiningOp()) { + Value term = add.getOffset(); + totalOffset = + totalOffset + ? rewriter.create(op.getLoc(), totalOffset, term) + : term; + base = add.getPtr(); + } + if (totalOffset) + offset = totalOffset; + + auto baseType = dyn_cast(base.getType()); + if (!baseType) + continue; + int64_t rank = static_cast(op.getShape().size()); + SmallVector dynamicShape(rank, ShapedType::kDynamic); + SmallVector dynamicStrides(rank, ShapedType::kDynamic); + auto layout = + StridedLayoutAttr::get(ctx, ShapedType::kDynamic, dynamicStrides); + auto resultType = MemRefType::get(dynamicShape, baseType.getElementType(), + layout, baseType.getMemorySpace()); + SmallVector sizes(op.getShape().begin(), op.getShape().end()); + SmallVector strides(op.getStrides().begin(), + op.getStrides().end()); + auto view = rewriter.create( + op.getLoc(), resultType, base, offset, sizes, strides); + if (totalOffset) + view->setAttr("pto.addptr_trace", rewriter.getUnitAttr()); + if (auto layoutAttr = op.getLayoutAttr()) + view->setAttr("layout", layoutAttr); + loweredViews[op.getResult()] = view.getResult(); + } + + SmallVector dimOps; + module.walk([&](pto::GetTensorViewDimOp op) { dimOps.push_back(op); }); + for (pto::GetTensorViewDimOp op : dimOps) { + Value view = op.getTensorView(); + if (Value lowered = loweredViews.lookup(view)) + view = lowered; + if (!isa(view.getType())) + continue; + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(op); + rewriter.replaceOpWithNewOp(op, view, op.getDimIndex()); + } + + SmallVector partitions; + module.walk([&](pto::PartitionViewOp op) { partitions.push_back(op); }); + for (pto::PartitionViewOp op : partitions) { + Value source = op.getSource(); + if (Value lowered = loweredViews.lookup(source)) + source = lowered; + auto sourceType = dyn_cast(source.getType()); + if (!sourceType) + continue; + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(op); + SmallVector staticSizes; + SmallVector sizes; + for (Value size : op.getSizes()) { + IntegerAttr attr; + if (matchPattern(size, m_Constant(&attr))) { + int64_t value = attr.getValue().getSExtValue(); + staticSizes.push_back(value); + sizes.push_back(rewriter.getIndexAttr(value)); + } else { + staticSizes.push_back(ShapedType::kDynamic); + sizes.push_back(size); + } + } + SmallVector offsets(op.getOffsets().begin(), + op.getOffsets().end()); + SmallVector strides(sourceType.getRank(), + rewriter.getIndexAttr(1)); + SmallVector dynamicStrides(sourceType.getRank(), + ShapedType::kDynamic); + auto layout = + StridedLayoutAttr::get(ctx, ShapedType::kDynamic, dynamicStrides); + auto resultType = MemRefType::get(staticSizes, sourceType.getElementType(), + layout, sourceType.getMemorySpace()); + auto subview = rewriter.create( + op.getLoc(), resultType, source, offsets, sizes, strides); + if (Operation *def = source.getDefiningOp()) + if (auto attr = def->getAttrOfType("layout")) + subview->setAttr("layout", attr); + rewriter.replaceOp(op, subview.getResult()); + } + + SmallVector strideOps; + module.walk([&](pto::GetTensorViewStrideOp op) { strideOps.push_back(op); }); + for (pto::GetTensorViewStrideOp op : strideOps) { + Value view = op.getTensorView(); + if (Value lowered = loweredViews.lookup(view)) + view = lowered; + auto type = dyn_cast(view.getType()); + if (!type) + continue; + IntegerAttr dimAttr; + if (!matchPattern(op.getDimIndex(), m_Constant(&dimAttr))) + return op.emitError( + "get_tensor_view_stride expects a constant dim index"); + int64_t dim = dimAttr.getValue().getSExtValue(); + if (dim < 0 || dim >= type.getRank()) + return op.emitError("get_tensor_view_stride dim index is out of bounds"); + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(op); + SmallVector staticStrides; + int64_t staticOffset = ShapedType::kDynamic; + if (succeeded(pto::getPTOMemRefStridesAndOffset(type, staticStrides, + staticOffset)) && + staticStrides[dim] != ShapedType::kDynamic) { + rewriter.replaceOpWithNewOp(op, + staticStrides[dim]); + continue; + } + auto metadata = + rewriter.create(op.getLoc(), view); + rewriter.replaceOp(op, metadata.getStrides()[dim]); + } + + // Replace roots only after all strongly typed PTO view users have been + // removed. This avoids constructing an intermediate partition_view whose + // source has already changed from tensor_view to memref. + for (pto::MakeTensorViewOp op : makeViews) { + Value lowered = loweredViews.lookup(op.getResult()); + if (!lowered) + continue; + IRRewriter rewriter(ctx); + rewriter.replaceOp(op, lowered); + } + return reconcileSCFViewResultTypes(module); +} + static FailureOr getStaticSlotBytes(pto::TileBufType slotType) { uint64_t elemBytes = pto::getPTOStorageElemByteSize(slotType.getElementType()); if (elemBytes == 0) @@ -365,6 +531,10 @@ struct PTOResolveBufferSelectPass signalPassFailure(); return; } + if (failed(resolveGMTensorViews(mod, ctx))) { + signalPassFailure(); + return; + } SmallVector markers; mod.walk([&](pto::SlotMarkerOp op) { markers.push_back(op); }); diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 835ec43814..bceecfb6ab 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -79,7 +79,9 @@ static bool hasMigratedTileNativeView(func::FuncOp func) { bool found = false; auto result = func.walk([&](Operation *op) { if (isa(op)) { + pto::GetValidShapeOp, pto::SubViewOp, pto::TileBufAddrOp, + pto::MakeTensorViewOp, pto::PartitionViewOp, + pto::GetTensorViewDimOp, pto::GetTensorViewStrideOp>(op)) { found = true; return WalkResult::interrupt(); } @@ -1412,8 +1414,13 @@ struct PTOViewToMemrefPass // ------------------------------------------------------------------ // Stage 1: Lower pto.make_tensor_view -> memref.reinterpret_cast // ------------------------------------------------------------------ + // GM view ops stay authored until PTOResolveBufferSelect when a + // function contains the migrated view family. Legacy functions keep + // the established memref lowering below. DefaultInlineVector makeViews; func.walk([&](mlir::pto::MakeTensorViewOp op) { makeViews.push_back(op); }); + if (preserveTileABI) + makeViews.clear(); for (auto op : makeViews) { IRRewriter rewriter(ctx); @@ -1506,7 +1513,7 @@ struct PTOViewToMemrefPass // ------------------------------------------------------------------ // Stage 1.3: Lower pto.partition_view -> memref.subview // ------------------------------------------------------------------ - if (failed(lowerPartitionViewOps(func, ctx))) { + if (!preserveTileABI && failed(lowerPartitionViewOps(func, ctx))) { signalPassFailure(); return; } @@ -1521,7 +1528,6 @@ struct PTOViewToMemrefPass if (failed(reconcileFusionRegionResultTypes(func))) { signalPassFailure(); return; - } // ------------------------------------------------------------------ // Stage 1.5: Lower pto.get_tensor_view_stride -> strided memref metadata @@ -1566,6 +1572,7 @@ struct PTOViewToMemrefPass rewriter.create(loc, view); rewriter.replaceOp(op, metadata.getStrides()[dimIndex]); } + } // ------------------------------------------------------------------ // Stage 1.6: Fold pto.addptr chains into load/store_scalar. diff --git a/test/lit/pto/gm_tensor_view_native_through_sync.pto b/test/lit/pto/gm_tensor_view_native_through_sync.pto new file mode 100644 index 0000000000..26223f478d --- /dev/null +++ b/test/lit/pto/gm_tensor_view_native_through_sync.pto @@ -0,0 +1,47 @@ +// 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 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-buffer-select %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=RESOLVED + +module { + func.func @gm_tensor_view_native_through_sync(%ptr: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %view = pto.make_tensor_view %ptr, shape = [%c32, %c16], + strides = [%c16, %c1] : !pto.tensor_view + %dim = pto.get_tensor_view_dim %view, %c0 + : !pto.tensor_view -> index + %stride = pto.get_tensor_view_stride %view, %c1 + : !pto.tensor_view -> index + %part = pto.partition_view %view, offsets = [%c0, %stride], + sizes = [%dim, %c16] + : !pto.tensor_view -> !pto.partition_tensor_view<32x16xf32> + return + } +} + +// NATIVE: IR Dump After PTOViewToMemref +// NATIVE: pto.make_tensor_view +// NATIVE: pto.get_tensor_view_dim +// NATIVE: pto.get_tensor_view_stride +// NATIVE: pto.partition_view +// NATIVE-NOT: memref.reinterpret_cast +// NATIVE-NOT: memref.subview + +// RESOLVED: IR Dump After PTOResolveBufferSelect +// RESOLVED: memref.reinterpret_cast +// RESOLVED: memref.dim +// RESOLVED: memref.subview +// RESOLVED: memref.extract_strided_metadata +// RESOLVED-NOT: pto.make_tensor_view +// RESOLVED-NOT: pto.get_tensor_view_dim +// RESOLVED-NOT: pto.get_tensor_view_stride +// RESOLVED-NOT: pto.partition_view From 3bd0577307549eba12e1b2b28099bd12eed9851b Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 12:08:43 +0800 Subject: [PATCH 19/91] Keep simple unary ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 7 ++- .../pto/ptodsl_subkernel_call_autosync.pto | 7 ++- test/lit/pto/unary_compute_tile_native.pto | 57 +++++++++++++++++++ 3 files changed, 66 insertions(+), 5 deletions(-) create mode 100644 test/lit/pto/unary_compute_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index bceecfb6ab..f8a13272a4 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -75,13 +75,14 @@ static bool hasTileNativeAllocationRoot(func::FuncOp func) { return found; } -static bool hasMigratedTileNativeView(func::FuncOp func) { +static bool hasMigratedTileNativeOp(func::FuncOp func) { bool found = false; auto result = func.walk([&](Operation *op) { if (isa(op)) { + pto::GetTensorViewDimOp, pto::GetTensorViewStrideOp, + pto::TAbsOp, pto::TNegOp, pto::TNotOp>(op)) { found = true; return WalkResult::interrupt(); } @@ -1265,7 +1266,7 @@ struct PTOViewToMemrefPass continue; auto fnTy = func.getFunctionType(); - bool preserveTileABI = hasMigratedTileNativeView(func); + bool preserveTileABI = hasMigratedTileNativeOp(func); bool tileNativeMainline = preserveTileABI || hasTileNativeAllocationRoot(func); diff --git a/test/lit/pto/ptodsl_subkernel_call_autosync.pto b/test/lit/pto/ptodsl_subkernel_call_autosync.pto index 5d5532fb28..8cd65ca69a 100644 --- a/test/lit/pto/ptodsl_subkernel_call_autosync.pto +++ b/test/lit/pto/ptodsl_subkernel_call_autosync.pto @@ -39,8 +39,11 @@ module { // CHECK-LABEL: AICORE void subkernel_call_autosync( // CHECK: TLOAD( -// CHECK: set_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); +// CHECK: set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); +// CHECK: wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0); // CHECK: TNEG( // CHECK: pipe_barrier(PIPE_V); -// CHECK: wait_flag(PIPE_MTE2, PIPE_MTE3, EVENT_ID0); +// CHECK: TNEG( +// CHECK: set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); +// CHECK: wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0); // CHECK: TSTORE( diff --git a/test/lit/pto/unary_compute_tile_native.pto b/test/lit/pto/unary_compute_tile_native.pto new file mode 100644 index 0000000000..d73a152dde --- /dev/null +++ b/test/lit/pto/unary_compute_tile_native.pto @@ -0,0 +1,57 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tabs_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tabs ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tneg_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tneg ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tnot_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tnot ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tabs_arg( +// NATIVE-SAME: %arg0: !pto.tile_buf +// NATIVE-SAME: %arg1: !pto.tile_buf +// NATIVE: pto.tabs ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tneg_arg( +// NATIVE: pto.tneg ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tnot_arg( +// NATIVE: pto.tnot ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tabs_arg( +// EMITC: TABS( +// EMITC-LABEL: tneg_arg( +// EMITC: TNEG( +// EMITC-LABEL: tnot_arg( +// EMITC: TNOT( From 142422de03d68aa951bef88e50c7e3fbb293cea6 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 12:38:00 +0800 Subject: [PATCH 20/91] Keep transcendental unary ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- .../pto/transcendental_unary_tile_native.pto | 47 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/transcendental_unary_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index f8a13272a4..434265f912 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -82,7 +82,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::GetValidShapeOp, pto::SubViewOp, pto::TileBufAddrOp, pto::MakeTensorViewOp, pto::PartitionViewOp, pto::GetTensorViewDimOp, pto::GetTensorViewStrideOp, - pto::TAbsOp, pto::TNegOp, pto::TNotOp>(op)) { + pto::TAbsOp, pto::TNegOp, pto::TNotOp, pto::TExpOp, + pto::TLogOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/transcendental_unary_tile_native.pto b/test/lit/pto/transcendental_unary_tile_native.pto new file mode 100644 index 0000000000..37e90e4527 --- /dev/null +++ b/test/lit/pto/transcendental_unary_tile_native.pto @@ -0,0 +1,47 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @texp_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.texp ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {precisionType = #pto} + return + } + + func.func private @tlog_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tlog ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {precisionType = #pto} + return + } +} + +// NATIVE-LABEL: func.func private @texp_arg( +// NATIVE-SAME: %arg0: !pto.tile_buf +// NATIVE-SAME: %arg1: !pto.tile_buf +// NATIVE: pto.texp ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) {precisionType = #pto} +// NATIVE-LABEL: func.func private @tlog_arg( +// NATIVE: pto.tlog ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) {precisionType = #pto} +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: texp_arg( +// EMITC: TEXP( +// EMITC-LABEL: tlog_arg( +// EMITC: TLOG( From d6978803992660a2568ebaa8a63025941c63a970 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 12:46:34 +0800 Subject: [PATCH 21/91] Keep TLRelu tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/tlrelu_tile_native.pto | 35 ++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/tlrelu_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 434265f912..a20b10773a 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -83,7 +83,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::MakeTensorViewOp, pto::PartitionViewOp, pto::GetTensorViewDimOp, pto::GetTensorViewStrideOp, pto::TAbsOp, pto::TNegOp, pto::TNotOp, pto::TExpOp, - pto::TLogOp>(op)) { + pto::TLogOp, pto::TLReluOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/tlrelu_tile_native.pto b/test/lit/pto/tlrelu_tile_native.pto new file mode 100644 index 0000000000..9324743b06 --- /dev/null +++ b/test/lit/pto/tlrelu_tile_native.pto @@ -0,0 +1,35 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tlrelu_arg( + %src: !pto.tile_buf, + %slope: f32, + %dst: !pto.tile_buf) { + pto.tlrelu ins(%src, %slope : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tlrelu_arg( +// NATIVE-SAME: %arg0: !pto.tile_buf +// NATIVE-SAME: %arg1: f32 +// NATIVE-SAME: %arg2: !pto.tile_buf +// NATIVE: pto.tlrelu ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tlrelu_arg( +// EMITC: TLRELU( From a40aa1b7a5d5b937f7075892c679eca3cefbddb4 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 12:55:23 +0800 Subject: [PATCH 22/91] Keep TMax and TMin tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/minmax_tile_native.pto | 54 ++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/minmax_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index a20b10773a..ade640fd5a 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -83,7 +83,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::MakeTensorViewOp, pto::PartitionViewOp, pto::GetTensorViewDimOp, pto::GetTensorViewStrideOp, pto::TAbsOp, pto::TNegOp, pto::TNotOp, pto::TExpOp, - pto::TLogOp, pto::TLReluOp>(op)) { + pto::TLogOp, pto::TLReluOp, pto::TMaxOp, pto::TMinOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/minmax_tile_native.pto b/test/lit/pto/minmax_tile_native.pto new file mode 100644 index 0000000000..8b9db4389e --- /dev/null +++ b/test/lit/pto/minmax_tile_native.pto @@ -0,0 +1,54 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tmax_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tmax ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tmax ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%src0 : !pto.tile_buf) + return + } + + func.func private @tmin_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tmin ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tmin ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%src0 : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tmax_arg( +// NATIVE-SAME: %arg0: !pto.tile_buf +// NATIVE-SAME: %arg1: !pto.tile_buf +// NATIVE-SAME: %arg2: !pto.tile_buf +// NATIVE: pto.tmax ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tmax ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg0 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tmin_arg( +// NATIVE: pto.tmin ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tmin ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg0 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tmax_arg( +// EMITC-COUNT-2: TMAX( +// EMITC-LABEL: tmin_arg( +// EMITC-COUNT-2: TMIN( From 9fbd807311459fbfda088570cdbce6222246a10a Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 13:03:06 +0800 Subject: [PATCH 23/91] Keep TMaxS and TMinS tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/minmax_scalar_tile_native.pto | 54 ++++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/minmax_scalar_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index ade640fd5a..538e87891b 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -83,7 +83,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::MakeTensorViewOp, pto::PartitionViewOp, pto::GetTensorViewDimOp, pto::GetTensorViewStrideOp, pto::TAbsOp, pto::TNegOp, pto::TNotOp, pto::TExpOp, - pto::TLogOp, pto::TLReluOp, pto::TMaxOp, pto::TMinOp>(op)) { + pto::TLogOp, pto::TLReluOp, pto::TMaxOp, pto::TMinOp, + pto::TMaxSOp, pto::TMinSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/minmax_scalar_tile_native.pto b/test/lit/pto/minmax_scalar_tile_native.pto new file mode 100644 index 0000000000..7f053bb34c --- /dev/null +++ b/test/lit/pto/minmax_scalar_tile_native.pto @@ -0,0 +1,54 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tmaxs_arg( + %src: !pto.tile_buf, + %scalar: f32, + %dst: !pto.tile_buf) { + pto.tmaxs ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + pto.tmaxs ins(%src, %scalar : !pto.tile_buf, f32) + outs(%src : !pto.tile_buf) + return + } + + func.func private @tmins_arg( + %src: !pto.tile_buf, + %scalar: f32, + %dst: !pto.tile_buf) { + pto.tmins ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + pto.tmins ins(%src, %scalar : !pto.tile_buf, f32) + outs(%src : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tmaxs_arg( +// NATIVE-SAME: %arg0: !pto.tile_buf +// NATIVE-SAME: %arg1: f32 +// NATIVE-SAME: %arg2: !pto.tile_buf +// NATIVE: pto.tmaxs ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tmaxs ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg0 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tmins_arg( +// NATIVE: pto.tmins ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tmins ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg0 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tmaxs_arg( +// EMITC-COUNT-2: TMAXS( +// EMITC-LABEL: tmins_arg( +// EMITC-COUNT-2: TMINS( From 949fd3e09c9141cfc7853689a116a594c3bf00e6 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 13:18:16 +0800 Subject: [PATCH 24/91] Keep TAdd and TMul tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/addmul_tile_native.pto | 54 ++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/addmul_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 538e87891b..646bd54aa3 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -84,7 +84,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::GetTensorViewDimOp, pto::GetTensorViewStrideOp, pto::TAbsOp, pto::TNegOp, pto::TNotOp, pto::TExpOp, pto::TLogOp, pto::TLReluOp, pto::TMaxOp, pto::TMinOp, - pto::TMaxSOp, pto::TMinSOp>(op)) { + pto::TMaxSOp, pto::TMinSOp, pto::TAddOp, pto::TMulOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/addmul_tile_native.pto b/test/lit/pto/addmul_tile_native.pto new file mode 100644 index 0000000000..a26a71029e --- /dev/null +++ b/test/lit/pto/addmul_tile_native.pto @@ -0,0 +1,54 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tadd_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%src0 : !pto.tile_buf) + return + } + + func.func private @tmul_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tmul ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tmul ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%src0 : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tadd_arg( +// NATIVE-SAME: %arg0: !pto.tile_buf +// NATIVE-SAME: %arg1: !pto.tile_buf +// NATIVE-SAME: %arg2: !pto.tile_buf +// NATIVE: pto.tadd ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tadd ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg0 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tmul_arg( +// NATIVE: pto.tmul ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tmul ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg0 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tadd_arg( +// EMITC-COUNT-2: TADD( +// EMITC-LABEL: tmul_arg( +// EMITC-COUNT-2: TMUL( From c9f3234ee3ba39a9248aaf57aa140baa1b5d1747 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 13:23:42 +0800 Subject: [PATCH 25/91] Keep TAddS and TMulS tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/addmul_scalar_tile_native.pto | 54 ++++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/addmul_scalar_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 646bd54aa3..f2db46b11e 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -84,7 +84,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::GetTensorViewDimOp, pto::GetTensorViewStrideOp, pto::TAbsOp, pto::TNegOp, pto::TNotOp, pto::TExpOp, pto::TLogOp, pto::TLReluOp, pto::TMaxOp, pto::TMinOp, - pto::TMaxSOp, pto::TMinSOp, pto::TAddOp, pto::TMulOp>(op)) { + pto::TMaxSOp, pto::TMinSOp, pto::TAddOp, pto::TMulOp, + pto::TAddSOp, pto::TMulSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/addmul_scalar_tile_native.pto b/test/lit/pto/addmul_scalar_tile_native.pto new file mode 100644 index 0000000000..65f0a83039 --- /dev/null +++ b/test/lit/pto/addmul_scalar_tile_native.pto @@ -0,0 +1,54 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tadds_arg( + %src: !pto.tile_buf, + %scalar: f32, + %dst: !pto.tile_buf) { + pto.tadds ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + pto.tadds ins(%src, %scalar : !pto.tile_buf, f32) + outs(%src : !pto.tile_buf) + return + } + + func.func private @tmuls_arg( + %src: !pto.tile_buf, + %scalar: f32, + %dst: !pto.tile_buf) { + pto.tmuls ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + pto.tmuls ins(%src, %scalar : !pto.tile_buf, f32) + outs(%src : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tadds_arg( +// NATIVE-SAME: %arg0: !pto.tile_buf +// NATIVE-SAME: %arg1: f32 +// NATIVE-SAME: %arg2: !pto.tile_buf +// NATIVE: pto.tadds ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tadds ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg0 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tmuls_arg( +// NATIVE: pto.tmuls ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tmuls ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg0 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tadds_arg( +// EMITC-COUNT-2: TADDS( +// EMITC-LABEL: tmuls_arg( +// EMITC-COUNT-2: TMULS( From 415886d51da318d3976f6fa65dd74e4f307be118 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 13:33:30 +0800 Subject: [PATCH 26/91] Make VPTO cache control checks order independent --- test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto | 36 +++++++++++++----------- 1 file changed, 19 insertions(+), 17 deletions(-) diff --git a/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto b/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto index f91b6a0389..a0381aa7ae 100644 --- a/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto +++ b/test/lit/vpto/mte_ub_gm_l2_cache_ctl.pto @@ -6,8 +6,10 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-before=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ROUNDTRIP -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=EXPAND +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-before=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ROUNDTRIP-CTL +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-before=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ROUNDTRIP-DEFAULT +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=EXPAND-CTL +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-expand-wrapper-ops %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=EXPAND-DEFAULT module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @mte_ub_gm_l2_cache_ctl(%out: !pto.ptr) attributes {pto.aicore} { @@ -36,19 +38,19 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, !pto.ptr, i64, i64, i64, i64, i64{{$}} -// ROUNDTRIP-LABEL: func.func @mte_ub_gm_l2_cache_ctl_default -// ROUNDTRIP: pto.mte_ub_gm %{{[^,]+}}, %arg0, %{{[^ ]+}} nburst(%{{[^,]+}}, %{{[^,]+}}, %{{[^)]+}}) -// ROUNDTRIP-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64{{$}} +// ROUNDTRIP-CTL-LABEL: func.func @mte_ub_gm_l2_cache_ctl( +// ROUNDTRIP-CTL: %[[L2:.*]] = arith.constant 5 : i64 +// ROUNDTRIP-CTL: pto.mte_ub_gm %{{[^,]+}}, %arg0, %{{[^ ]+}} nburst(%{{[^,]+}}, %{{[^,]+}}, %{{[^)]+}}) l2_cache_ctl(%[[L2]]) +// ROUNDTRIP-CTL-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64{{$}} +// ROUNDTRIP-DEFAULT-LABEL: func.func @mte_ub_gm_l2_cache_ctl_default( +// ROUNDTRIP-DEFAULT: pto.mte_ub_gm %{{[^,]+}}, %arg0, %{{[^ ]+}} nburst(%{{[^,]+}}, %{{[^,]+}}, %{{[^)]+}}) +// ROUNDTRIP-DEFAULT-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64{{$}} -// EXPAND-LABEL: func.func @mte_ub_gm_l2_cache_ctl -// EXPAND: %[[L2:.*]] = arith.constant 5 : i64 -// EXPAND: pto.copy_ubuf_to_gm %{{[^,]+}}, %arg0, %{{[^,]+}}, %{{[^,]+}}, %{{[^,]+}}, %[[L2]], %{{[^,]+}}, %{{[^:]+}} -// EXPAND-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 -// EXPAND-LABEL: func.func @mte_ub_gm_l2_cache_ctl_default -// EXPAND: %[[ZERO:.*]] = arith.constant 0 : i64 -// EXPAND: pto.copy_ubuf_to_gm %{{[^,]+}}, %arg0, %[[ZERO]], %{{[^,]+}}, %{{[^,]+}}, %[[ZERO]], %{{[^,]+}}, %{{[^:]+}} -// EXPAND-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 +// EXPAND-CTL-LABEL: func.func @mte_ub_gm_l2_cache_ctl( +// EXPAND-CTL: %[[L2:.*]] = arith.constant 5 : i64 +// EXPAND-CTL: pto.copy_ubuf_to_gm %{{[^,]+}}, %arg0, %{{[^,]+}}, %{{[^,]+}}, %{{[^,]+}}, %[[L2]], %{{[^,]+}}, %{{[^:]+}} +// EXPAND-CTL-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 +// EXPAND-DEFAULT-LABEL: func.func @mte_ub_gm_l2_cache_ctl_default( +// EXPAND-DEFAULT: %[[ZERO:.*]] = arith.constant 0 : i64 +// EXPAND-DEFAULT: pto.copy_ubuf_to_gm %{{[^,]+}}, %arg0, %[[ZERO]], %{{[^,]+}}, %{{[^,]+}}, %[[ZERO]], %{{[^,]+}}, %{{[^:]+}} +// EXPAND-DEFAULT-SAME: : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 From dfcbe1bcb38537e6c3b6682b84351cecb25eadc7 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 13:37:35 +0800 Subject: [PATCH 27/91] Keep TAnd and TOr tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/logic_binary_tile_native.pto | 51 +++++++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/logic_binary_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index f2db46b11e..99d55ad864 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -85,7 +85,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TAbsOp, pto::TNegOp, pto::TNotOp, pto::TExpOp, pto::TLogOp, pto::TLReluOp, pto::TMaxOp, pto::TMinOp, pto::TMaxSOp, pto::TMinSOp, pto::TAddOp, pto::TMulOp, - pto::TAddSOp, pto::TMulSOp>(op)) { + pto::TAddSOp, pto::TMulSOp, pto::TAndOp, pto::TOrOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/logic_binary_tile_native.pto b/test/lit/pto/logic_binary_tile_native.pto new file mode 100644 index 0000000000..147bfda9ae --- /dev/null +++ b/test/lit/pto/logic_binary_tile_native.pto @@ -0,0 +1,51 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tand_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tand ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tand ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%src0 : !pto.tile_buf) + return + } + + func.func private @tor_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tor ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tor ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%src0 : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tand_arg( +// NATIVE: pto.tand ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tand ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg0 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tor_arg( +// NATIVE: pto.tor ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE: pto.tor ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg0 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tand_arg( +// EMITC-COUNT-2: TAND( +// EMITC-LABEL: tor_arg( +// EMITC-COUNT-2: TOR( From af9e591567c03e48071bd7399ded20f230afd48f Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 13:43:53 +0800 Subject: [PATCH 28/91] Keep TAndS and TOrS tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/logic_scalar_tile_native.pto | 45 +++++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/logic_scalar_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 99d55ad864..3cc0651ed0 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -85,7 +85,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TAbsOp, pto::TNegOp, pto::TNotOp, pto::TExpOp, pto::TLogOp, pto::TLReluOp, pto::TMaxOp, pto::TMinOp, pto::TMaxSOp, pto::TMinSOp, pto::TAddOp, pto::TMulOp, - pto::TAddSOp, pto::TMulSOp, pto::TAndOp, pto::TOrOp>(op)) { + pto::TAddSOp, pto::TMulSOp, pto::TAndOp, pto::TOrOp, + pto::TAndSOp, pto::TOrSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/logic_scalar_tile_native.pto b/test/lit/pto/logic_scalar_tile_native.pto new file mode 100644 index 0000000000..dc6c457f68 --- /dev/null +++ b/test/lit/pto/logic_scalar_tile_native.pto @@ -0,0 +1,45 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tands_arg( + %src: !pto.tile_buf, + %scalar: i16, + %dst: !pto.tile_buf) { + pto.tands ins(%src, %scalar : !pto.tile_buf, i16) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tors_arg( + %src: !pto.tile_buf, + %scalar: i16, + %dst: !pto.tile_buf) { + pto.tors ins(%src, %scalar : !pto.tile_buf, i16) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tands_arg( +// NATIVE: pto.tands ins(%arg0, %arg1 : !pto.tile_buf, i16) outs(%arg2 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tors_arg( +// NATIVE: pto.tors ins(%arg0, %arg1 : !pto.tile_buf, i16) outs(%arg2 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tands_arg( +// EMITC: TANDS( +// EMITC-LABEL: tors_arg( +// EMITC: TORS( From d277e8ce745456e923401629d0c369346a6b61b9 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 13:50:19 +0800 Subject: [PATCH 29/91] Keep TCmp and TCmpS tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/compare_tile_native.pto | 45 ++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/compare_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 3cc0651ed0..2eab81cfc9 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -86,7 +86,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TLogOp, pto::TLReluOp, pto::TMaxOp, pto::TMinOp, pto::TMaxSOp, pto::TMinSOp, pto::TAddOp, pto::TMulOp, pto::TAddSOp, pto::TMulSOp, pto::TAndOp, pto::TOrOp, - pto::TAndSOp, pto::TOrSOp>(op)) { + pto::TAndSOp, pto::TOrSOp, pto::TCmpOp, pto::TCmpSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/compare_tile_native.pto b/test/lit/pto/compare_tile_native.pto new file mode 100644 index 0000000000..2728956cf8 --- /dev/null +++ b/test/lit/pto/compare_tile_native.pto @@ -0,0 +1,45 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tcmp_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcmp ins(%src0, %src1 {cmpMode = #pto} : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tcmps_arg( + %src: !pto.tile_buf, + %scalar: f32, + %dst: !pto.tile_buf) { + pto.tcmps ins(%src, %scalar {cmpMode = #pto} : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tcmp_arg( +// NATIVE: pto.tcmp ins(%arg0, %arg1 {cmpMode = #pto} : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tcmps_arg( +// NATIVE: pto.tcmps ins(%arg0, %arg1 {cmpMode = #pto} : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tcmp_arg( +// EMITC: TCMP( +// EMITC-LABEL: tcmps_arg( +// EMITC: TCMPS( From cb70d49499a5461a437422ff635a9540b1c3dbd5 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:01:09 +0800 Subject: [PATCH 30/91] Keep TDiv and TDivS tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/div_tile_native.pto | 61 ++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/div_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 2eab81cfc9..98188d4fe2 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -86,7 +86,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TLogOp, pto::TLReluOp, pto::TMaxOp, pto::TMinOp, pto::TMaxSOp, pto::TMinSOp, pto::TAddOp, pto::TMulOp, pto::TAddSOp, pto::TMulSOp, pto::TAndOp, pto::TOrOp, - pto::TAndSOp, pto::TOrSOp, pto::TCmpOp, pto::TCmpSOp>(op)) { + pto::TAndSOp, pto::TOrSOp, pto::TCmpOp, pto::TCmpSOp, + pto::TDivOp, pto::TDivSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/div_tile_native.pto b/test/lit/pto/div_tile_native.pto new file mode 100644 index 0000000000..dc606c661d --- /dev/null +++ b/test/lit/pto/div_tile_native.pto @@ -0,0 +1,61 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tdiv_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tdiv ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {precisionType = #pto} + return + } + + func.func private @tdivs_tile_scalar_arg( + %src: !pto.tile_buf, + %scalar: f32, + %dst: !pto.tile_buf) { + pto.tdivs ins(%src, %scalar : !pto.tile_buf, f32) + outs(%dst : !pto.tile_buf) + {precisionType = #pto} + return + } + + func.func private @tdivs_scalar_tile_arg( + %scalar: f32, + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tdivs ins(%scalar, %src : f32, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {precisionType = #pto} + return + } +} + +// NATIVE-LABEL: func.func private @tdiv_arg( +// NATIVE: pto.tdiv ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) {precisionType = #pto} +// NATIVE-LABEL: func.func private @tdivs_tile_scalar_arg( +// NATIVE: pto.tdivs ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) {precisionType = #pto} +// NATIVE-LABEL: func.func private @tdivs_scalar_tile_arg( +// NATIVE: pto.tdivs ins(%arg0, %arg1 : f32, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) {precisionType = #pto} +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tdiv_arg( +// EMITC: TDIV( +// EMITC-LABEL: tdivs_tile_scalar_arg( +// EMITC: TDIVS( +// EMITC-LABEL: tdivs_scalar_tile_arg( +// EMITC: TDIVS( From 0a493abcfaf565604e85f8a97207ceeccf71cc17 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:06:41 +0800 Subject: [PATCH 31/91] Keep TConcat ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/concat_tile_native.pto | 47 ++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/concat_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 98188d4fe2..912a17cc7c 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -87,7 +87,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TMaxSOp, pto::TMinSOp, pto::TAddOp, pto::TMulOp, pto::TAddSOp, pto::TMulSOp, pto::TAndOp, pto::TOrOp, pto::TAndSOp, pto::TOrSOp, pto::TCmpOp, pto::TCmpSOp, - pto::TDivOp, pto::TDivSOp>(op)) { + pto::TDivOp, pto::TDivSOp, pto::TConcatOp, + pto::TConcatidxOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/concat_tile_native.pto b/test/lit/pto/concat_tile_native.pto new file mode 100644 index 0000000000..588116a7bb --- /dev/null +++ b/test/lit/pto/concat_tile_native.pto @@ -0,0 +1,47 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tconcat_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tconcat ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tconcatidx_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %idx0: !pto.tile_buf, + %idx1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tconcatidx ins(%src0, %src1, %idx0, %idx1 : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tconcat_arg( +// NATIVE: pto.tconcat ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tconcatidx_arg( +// NATIVE: pto.tconcatidx ins(%arg0, %arg1, %arg2, %arg3 : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%arg4 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tconcat_arg( +// EMITC: TCONCAT( +// EMITC-LABEL: tconcatidx_arg( +// EMITC: TCONCAT( From 5e6424b721fea52895f74574adb7d741bd5e38d4 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:12:38 +0800 Subject: [PATCH 32/91] Keep TPart arithmetic ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/part_arithmetic_tile_native.pto | 45 ++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/part_arithmetic_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 912a17cc7c..1b7810ec43 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -88,7 +88,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TAddSOp, pto::TMulSOp, pto::TAndOp, pto::TOrOp, pto::TAndSOp, pto::TOrSOp, pto::TCmpOp, pto::TCmpSOp, pto::TDivOp, pto::TDivSOp, pto::TConcatOp, - pto::TConcatidxOp>(op)) { + pto::TConcatidxOp, pto::TPartAddOp, pto::TPartMulOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/part_arithmetic_tile_native.pto b/test/lit/pto/part_arithmetic_tile_native.pto new file mode 100644 index 0000000000..e21d0b3be0 --- /dev/null +++ b/test/lit/pto/part_arithmetic_tile_native.pto @@ -0,0 +1,45 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tpartadd_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tpartadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tpartmul_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tpartmul ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tpartadd_arg( +// NATIVE: pto.tpartadd ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tpartmul_arg( +// NATIVE: pto.tpartmul ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tpartadd_arg( +// EMITC: TPARTADD( +// EMITC-LABEL: tpartmul_arg( +// EMITC: TPARTMUL( From 91beee997af9b229e32040de3348b6b26c5a3b4c Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:18:16 +0800 Subject: [PATCH 33/91] Keep TColExpand ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 4 +- test/lit/pto/colexpand_tile_native.pto | 70 ++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/colexpand_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 1b7810ec43..b2df8b85b4 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -88,7 +88,9 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TAddSOp, pto::TMulSOp, pto::TAndOp, pto::TOrOp, pto::TAndSOp, pto::TOrSOp, pto::TCmpOp, pto::TCmpSOp, pto::TDivOp, pto::TDivSOp, pto::TConcatOp, - pto::TConcatidxOp, pto::TPartAddOp, pto::TPartMulOp>(op)) { + pto::TConcatidxOp, pto::TPartAddOp, pto::TPartMulOp, + pto::TColExpandOp, pto::TColExpandMulOp, pto::TColExpandMaxOp, + pto::TColExpandMinOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/colexpand_tile_native.pto b/test/lit/pto/colexpand_tile_native.pto new file mode 100644 index 0000000000..75ef9fb8d4 --- /dev/null +++ b/test/lit/pto/colexpand_tile_native.pto @@ -0,0 +1,70 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tcolexpand_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolexpand ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tcolexpandmul_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolexpandmul ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tcolexpandmax_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolexpandmax ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tcolexpandmin_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolexpandmin ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tcolexpand_arg( +// NATIVE: pto.tcolexpand ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tcolexpandmul_arg( +// NATIVE: pto.tcolexpandmul ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tcolexpandmax_arg( +// NATIVE: pto.tcolexpandmax ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tcolexpandmin_arg( +// NATIVE: pto.tcolexpandmin ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tcolexpand_arg( +// EMITC: TCOLEXPAND( +// EMITC-LABEL: tcolexpandmul_arg( +// EMITC: TCOLEXPANDMUL( +// EMITC-LABEL: tcolexpandmax_arg( +// EMITC: TCOLEXPANDMAX( +// EMITC-LABEL: tcolexpandmin_arg( +// EMITC: TCOLEXPANDMIN( From 961cfda283292f658f7e43cec97ed7a2b0320c0f Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:24:01 +0800 Subject: [PATCH 34/91] Keep TExpands tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/expands_tile_native.pto | 31 ++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/expands_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index b2df8b85b4..c8835c0a88 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -90,7 +90,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TDivOp, pto::TDivSOp, pto::TConcatOp, pto::TConcatidxOp, pto::TPartAddOp, pto::TPartMulOp, pto::TColExpandOp, pto::TColExpandMulOp, pto::TColExpandMaxOp, - pto::TColExpandMinOp>(op)) { + pto::TColExpandMinOp, pto::TExpandsOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/expands_tile_native.pto b/test/lit/pto/expands_tile_native.pto new file mode 100644 index 0000000000..680ec473c5 --- /dev/null +++ b/test/lit/pto/expands_tile_native.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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @texpands_arg( + %scalar: f32, + %dst: !pto.tile_buf) { + pto.texpands ins(%scalar : f32) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @texpands_arg( +// NATIVE: pto.texpands ins(%arg0 : f32) outs(%arg1 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: texpands_arg( +// EMITC: TEXPANDS( From 78c5aadeeeb9f78adea42ffc3c3215b435909ad9 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:29:50 +0800 Subject: [PATCH 35/91] Keep tile scalar access ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- .../pto/tile_scalar_access_tile_native.pto | 46 +++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/tile_scalar_access_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index c8835c0a88..13401f54fc 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -90,7 +90,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TDivOp, pto::TDivSOp, pto::TConcatOp, pto::TConcatidxOp, pto::TPartAddOp, pto::TPartMulOp, pto::TColExpandOp, pto::TColExpandMulOp, pto::TColExpandMaxOp, - pto::TColExpandMinOp, pto::TExpandsOp>(op)) { + pto::TColExpandMinOp, pto::TExpandsOp, pto::TSetValOp, + pto::TGetValOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/tile_scalar_access_tile_native.pto b/test/lit/pto/tile_scalar_access_tile_native.pto new file mode 100644 index 0000000000..d201f2c382 --- /dev/null +++ b/test/lit/pto/tile_scalar_access_tile_native.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: ptoas --pto-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tsetval_arg( + %dst: !pto.tile_buf, + %offset: index, + %val: f32) { + pto.tsetval ins(%offset, %val : index, f32) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tgetval_arg( + %src: !pto.tile_buf, + %offset: index) -> f32 { + %val = pto.tgetval ins(%src, %offset : !pto.tile_buf, index) + outs : f32 + return %val : f32 + } +} + +// NATIVE-LABEL: func.func private @tsetval_arg( +// NATIVE: pto.tsetval ins(%arg1, %arg2 : index, f32) outs(%arg0 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tgetval_arg( +// NATIVE-SAME: -> f32 +// NATIVE: %0 = pto.tgetval ins(%arg0, %arg1 : !pto.tile_buf, index) outs : f32 +// NATIVE: return %0 : f32 +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tsetval_arg( +// EMITC: .SetValue( +// EMITC-LABEL: tgetval_arg( +// EMITC: .GetValue( From cf8a9aa6837c8e0e21e31add9af7e26156eea513 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:35:32 +0800 Subject: [PATCH 36/91] Keep TExtract and TInsert tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/extract_insert_tile_native.pto | 47 +++++++++++++++++++++ 2 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/extract_insert_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 13401f54fc..699d7294ba 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -91,7 +91,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TConcatidxOp, pto::TPartAddOp, pto::TPartMulOp, pto::TColExpandOp, pto::TColExpandMulOp, pto::TColExpandMaxOp, pto::TColExpandMinOp, pto::TExpandsOp, pto::TSetValOp, - pto::TGetValOp>(op)) { + pto::TGetValOp, pto::TExtractOp, pto::TInsertOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/extract_insert_tile_native.pto b/test/lit/pto/extract_insert_tile_native.pto new file mode 100644 index 0000000000..86a1d8ecfc --- /dev/null +++ b/test/lit/pto/extract_insert_tile_native.pto @@ -0,0 +1,47 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @textract_arg( + %src: !pto.tile_buf, + %row: index, + %col: index, + %dst: !pto.tile_buf) { + pto.textract ins(%src, %row, %col : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tinsert_arg( + %src: !pto.tile_buf, + %row: index, + %col: index, + %dst: !pto.tile_buf) { + pto.tinsert ins(%src, %row, %col : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @textract_arg( +// NATIVE: pto.textract ins(%arg0, %arg1, %arg2 : !pto.tile_buf, index, index) outs(%arg3 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tinsert_arg( +// NATIVE: pto.tinsert ins(%arg0, %arg1, %arg2 : !pto.tile_buf, index, index) outs(%arg3 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: textract_arg( +// EMITC: TEXTRACT( +// EMITC-LABEL: tinsert_arg( +// EMITC: TINSERT( From 32ac9ef9bcde477d0997656b9f6d79d33c27ebf4 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:40:50 +0800 Subject: [PATCH 37/91] Keep TExtractFP and TInsertFP tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- .../lit/pto/extract_insert_fp_tile_native.pto | 47 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/extract_insert_fp_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 699d7294ba..5fb10f0eb3 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -91,7 +91,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TConcatidxOp, pto::TPartAddOp, pto::TPartMulOp, pto::TColExpandOp, pto::TColExpandMulOp, pto::TColExpandMaxOp, pto::TColExpandMinOp, pto::TExpandsOp, pto::TSetValOp, - pto::TGetValOp, pto::TExtractOp, pto::TInsertOp>(op)) { + pto::TGetValOp, pto::TExtractOp, pto::TInsertOp, + pto::TExtractFPOp, pto::TInsertFPOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/extract_insert_fp_tile_native.pto b/test/lit/pto/extract_insert_fp_tile_native.pto new file mode 100644 index 0000000000..e3c740b1ed --- /dev/null +++ b/test/lit/pto/extract_insert_fp_tile_native.pto @@ -0,0 +1,47 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @textract_fp_arg( + %src: !pto.tile_buf, + %fp: !pto.tile_buf, + %row: index, %col: index, + %dst: !pto.tile_buf) { + pto.textract_fp ins(%src, %fp, %row, %col : !pto.tile_buf, !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tinsert_fp_arg( + %src: !pto.tile_buf, + %fp: !pto.tile_buf, + %row: index, %col: index, + %dst: !pto.tile_buf) { + pto.tinsert_fp ins(%src, %fp, %row, %col : !pto.tile_buf, !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @textract_fp_arg( +// NATIVE: pto.textract_fp ins(%arg0, %arg1, %arg2, %arg3 +// NATIVE-LABEL: func.func private @tinsert_fp_arg( +// NATIVE: pto.tinsert_fp ins(%arg0, %arg1, %arg2, %arg3 +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: textract_fp_arg( +// EMITC: TEXTRACT_FP( +// EMITC-LABEL: tinsert_fp_arg( +// EMITC: TINSERT_FP( From 0d9372b8ae73d3b4187c60a3822b5da4a106b764 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:46:41 +0800 Subject: [PATCH 38/91] Keep carry add ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/add_carry_tile_native.pto | 48 ++++++++++++++++++++++++++ 2 files changed, 50 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/add_carry_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 5fb10f0eb3..721f9eca14 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -92,7 +92,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TColExpandOp, pto::TColExpandMulOp, pto::TColExpandMaxOp, pto::TColExpandMinOp, pto::TExpandsOp, pto::TSetValOp, pto::TGetValOp, pto::TExtractOp, pto::TInsertOp, - pto::TExtractFPOp, pto::TInsertFPOp>(op)) { + pto::TExtractFPOp, pto::TInsertFPOp, pto::TAddCOp, + pto::TAddSCOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/add_carry_tile_native.pto b/test/lit/pto/add_carry_tile_native.pto new file mode 100644 index 0000000000..6b5d337527 --- /dev/null +++ b/test/lit/pto/add_carry_tile_native.pto @@ -0,0 +1,48 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @taddc_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %src2: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.taddc ins(%src0, %src1, %src2 : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @taddsc_arg( + %src0: !pto.tile_buf, + %scalar: f32, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.taddsc ins(%src0, %scalar, %src1 : !pto.tile_buf, f32, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @taddc_arg( +// NATIVE: pto.taddc ins(%arg0, %arg1, %arg2 : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%arg3 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @taddsc_arg( +// NATIVE: pto.taddsc ins(%arg0, %arg1, %arg2 : !pto.tile_buf, f32, !pto.tile_buf) outs(%arg3 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: taddc_arg( +// EMITC-COUNT-2: TADD( +// EMITC-LABEL: taddsc_arg( +// EMITC: TADDS( +// EMITC: TADD( From ca65b371e607816418c55351462d233d69317360 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:53:21 +0800 Subject: [PATCH 39/91] Keep TMovFP tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/tmov_fp_tile_native.pto | 32 ++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/tmov_fp_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 721f9eca14..b62111b280 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -93,7 +93,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TColExpandMinOp, pto::TExpandsOp, pto::TSetValOp, pto::TGetValOp, pto::TExtractOp, pto::TInsertOp, pto::TExtractFPOp, pto::TInsertFPOp, pto::TAddCOp, - pto::TAddSCOp>(op)) { + pto::TAddSCOp, pto::TMovFPOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/tmov_fp_tile_native.pto b/test/lit/pto/tmov_fp_tile_native.pto new file mode 100644 index 0000000000..4e2822406c --- /dev/null +++ b/test/lit/pto/tmov_fp_tile_native.pto @@ -0,0 +1,32 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tmov_fp_arg( + %src: !pto.tile_buf, + %fp: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tmov.fp ins(%src, %fp : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tmov_fp_arg( +// NATIVE: pto.tmov.fp ins(%arg0, %arg1 +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tmov_fp_arg( +// EMITC: TMOV_FP< From 018e097a492fe75a66a3f561b54c32b86d8f2f6a Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 14:59:02 +0800 Subject: [PATCH 40/91] Keep TFillPad ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/fillpad_tile_native.pto | 41 ++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/fillpad_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index b62111b280..97907a5826 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -93,7 +93,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TColExpandMinOp, pto::TExpandsOp, pto::TSetValOp, pto::TGetValOp, pto::TExtractOp, pto::TInsertOp, pto::TExtractFPOp, pto::TInsertFPOp, pto::TAddCOp, - pto::TAddSCOp, pto::TMovFPOp>(op)) { + pto::TAddSCOp, pto::TMovFPOp, pto::TFillPadOp, + pto::TFillPadInplaceOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/fillpad_tile_native.pto b/test/lit/pto/fillpad_tile_native.pto new file mode 100644 index 0000000000..081127f5ca --- /dev/null +++ b/test/lit/pto/fillpad_tile_native.pto @@ -0,0 +1,41 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tfillpad_arg( + %tile: !pto.tile_buf) { + pto.tfillpad ins(%tile : !pto.tile_buf) + outs(%tile : !pto.tile_buf) + return + } + + func.func private @tfillpad_inplace_arg( + %tile: !pto.tile_buf) { + pto.tfillpad_inplace ins(%tile : !pto.tile_buf) + outs(%tile : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tfillpad_arg( +// NATIVE: pto.tfillpad ins(%arg0 +// NATIVE-LABEL: func.func private @tfillpad_inplace_arg( +// NATIVE: pto.tfillpad_inplace ins(%arg0 +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tfillpad_arg( +// EMITC: TFILLPAD( +// EMITC-LABEL: tfillpad_inplace_arg( +// EMITC: TFILLPAD_INPLACE( From a73aafa42a2adb0a6086983ac99b5103ed67f4da Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 15:05:15 +0800 Subject: [PATCH 41/91] Keep TLoad and TStore local tiles native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/load_store_tile_native.pto | 42 +++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/load_store_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 97907a5826..b40008c314 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -94,7 +94,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TGetValOp, pto::TExtractOp, pto::TInsertOp, pto::TExtractFPOp, pto::TInsertFPOp, pto::TAddCOp, pto::TAddSCOp, pto::TMovFPOp, pto::TFillPadOp, - pto::TFillPadInplaceOp>(op)) { + pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/load_store_tile_native.pto b/test/lit/pto/load_store_tile_native.pto new file mode 100644 index 0000000000..ad73bd6915 --- /dev/null +++ b/test/lit/pto/load_store_tile_native.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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tload_arg( + %src: !pto.partition_tensor_view<16x16xf32>, + %dst: !pto.tile_buf) { + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tstore_arg( + %src: !pto.tile_buf, + %dst: !pto.partition_tensor_view<16x16xf32>) { + pto.tstore ins(%src : !pto.tile_buf) + outs(%dst : !pto.partition_tensor_view<16x16xf32>) + return + } +} + +// NATIVE-LABEL: func.func private @tload_arg( +// NATIVE: pto.tload ins(%arg0 : memref<16x16xf32>) outs(%arg1 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tstore_arg( +// NATIVE: pto.tstore ins(%arg0 : !pto.tile_buf) outs(%arg1 : memref<16x16xf32>) +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tload_arg( +// EMITC: TLOAD( +// EMITC-LABEL: tstore_arg( +// EMITC: TSTORE( From 8115ef40d602e277257ec67c82ea79642675c65e Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 15:10:59 +0800 Subject: [PATCH 42/91] Keep TColMax and TColMin tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/col_minmax_tile_native.pto | 43 +++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/col_minmax_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index b40008c314..f01ceeef6e 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -94,7 +94,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TGetValOp, pto::TExtractOp, pto::TInsertOp, pto::TExtractFPOp, pto::TInsertFPOp, pto::TAddCOp, pto::TAddSCOp, pto::TMovFPOp, pto::TFillPadOp, - pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp>(op)) { + pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp, + pto::TColMaxOp, pto::TColMinOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/col_minmax_tile_native.pto b/test/lit/pto/col_minmax_tile_native.pto new file mode 100644 index 0000000000..5be86f1bf5 --- /dev/null +++ b/test/lit/pto/col_minmax_tile_native.pto @@ -0,0 +1,43 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tcolmax_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolmax ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tcolmin_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolmin ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tcolmax_arg( +// NATIVE: pto.tcolmax ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tcolmin_arg( +// NATIVE: pto.tcolmin ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tcolmax_arg( +// EMITC: TCOLMAX( +// EMITC-LABEL: tcolmin_arg( +// EMITC: TCOLMIN( From 9a5b53d3847a8d1fefae82b223cb414504b2d9d9 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 15:16:35 +0800 Subject: [PATCH 43/91] Keep TColSum tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/colsum_tile_native.pto | 44 ++++++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/colsum_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index f01ceeef6e..cc7e178287 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -95,7 +95,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TExtractFPOp, pto::TInsertFPOp, pto::TAddCOp, pto::TAddSCOp, pto::TMovFPOp, pto::TFillPadOp, pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp, - pto::TColMaxOp, pto::TColMinOp>(op)) { + pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/colsum_tile_native.pto b/test/lit/pto/colsum_tile_native.pto new file mode 100644 index 0000000000..252428a442 --- /dev/null +++ b/test/lit/pto/colsum_tile_native.pto @@ -0,0 +1,44 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tcolsum_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolsum ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tcolsum_binary_arg( + %src: !pto.tile_buf, + %tmp: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolsum ins(%src, %tmp {isBinary = true} : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tcolsum_arg( +// NATIVE: pto.tcolsum ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tcolsum_binary_arg( +// NATIVE: pto.tcolsum ins(%arg0, %arg1 {isBinary = true} : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tcolsum_arg( +// EMITC: TCOLSUM( +// EMITC-LABEL: tcolsum_binary_arg( +// EMITC: TCOLSUM( From 88ab4c52d000e15ff8efae47f2d9625781a7a128 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 15:22:58 +0800 Subject: [PATCH 44/91] Keep TCvt tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 ++- test/lit/pto/cvt_tile_native.pto | 31 ++++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/cvt_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index cc7e178287..1b0eafebba 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -95,7 +95,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TExtractFPOp, pto::TInsertFPOp, pto::TAddCOp, pto::TAddSCOp, pto::TMovFPOp, pto::TFillPadOp, pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp, - pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp>(op)) { + pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp, + pto::TCvtOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/cvt_tile_native.pto b/test/lit/pto/cvt_tile_native.pto new file mode 100644 index 0000000000..f711431ce6 --- /dev/null +++ b/test/lit/pto/cvt_tile_native.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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tcvt_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcvt ins(%src {rmode = #pto, satmode = #pto} : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// 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-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tcvt_arg( +// EMITC: TCVT( From 738a59e4860f432a329caaa695b35ba0d3a48954 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 15:29:26 +0800 Subject: [PATCH 45/91] Keep gather ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/gather_tile_native.pto | 58 ++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/gather_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 1b0eafebba..2acb7127dc 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -96,7 +96,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TAddSCOp, pto::TMovFPOp, pto::TFillPadOp, pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp, pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp, - pto::TCvtOp>(op)) { + pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/gather_tile_native.pto b/test/lit/pto/gather_tile_native.pto new file mode 100644 index 0000000000..579e00f151 --- /dev/null +++ b/test/lit/pto/gather_tile_native.pto @@ -0,0 +1,58 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tgather_mask_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tgather ins(%src, {maskPattern = #pto.mask_pattern} : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tgather_index_arg( + %src: !pto.tile_buf, + %indices: !pto.tile_buf, + %tmp: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tgather ins(%src, %indices, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tgatherb_arg( + %src: !pto.tile_buf, + %offsets: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tgatherb ins(%src, %offsets : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tgather_mask_arg( +// NATIVE: pto.tgather +// NATIVE-LABEL: func.func private @tgather_index_arg( +// NATIVE: pto.tgather +// NATIVE-LABEL: func.func private @tgatherb_arg( +// NATIVE: pto.tgatherb ins(%arg0, %arg1 +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tgather_mask_arg( +// EMITC: TGATHER< +// EMITC-LABEL: tgather_index_arg( +// EMITC: TGATHER( +// EMITC-LABEL: tgatherb_arg( +// EMITC: TGATHERB( From 7f9f998a2d100440c7cdafbe45db2186436dab6b Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 15:40:42 +0800 Subject: [PATCH 46/91] Keep TCI tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 2acb7127dc..a9fcfe3fcc 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -96,7 +96,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TAddSCOp, pto::TMovFPOp, pto::TFillPadOp, pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp, pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp, - pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp>(op)) { + pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp, pto::TCIOp>(op)) { found = true; return WalkResult::interrupt(); } From ba29931cfd5c40aa327e2a7d961edfbef8569173 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 15:47:09 +0800 Subject: [PATCH 47/91] Keep TPrint tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/print_tile_native.pto | 40 ++++++++++++++++++++++++++ 2 files changed, 42 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/print_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index a9fcfe3fcc..404e2a9bca 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -96,7 +96,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TAddSCOp, pto::TMovFPOp, pto::TFillPadOp, pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp, pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp, - pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp, pto::TCIOp>(op)) { + pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp, pto::TCIOp, + pto::TPrintOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/print_tile_native.pto b/test/lit/pto/print_tile_native.pto new file mode 100644 index 0000000000..84247d2a5f --- /dev/null +++ b/test/lit/pto/print_tile_native.pto @@ -0,0 +1,40 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tprint_arg( + %src: !pto.tile_buf) { + pto.tprint ins(%src : !pto.tile_buf) + return + } + + func.func private @tprint_tmp_arg( + %src: !pto.tile_buf, + %tmp: !pto.partition_tensor_view<16x16xf32>) { + pto.tprint ins(%src, %tmp : !pto.tile_buf, !pto.partition_tensor_view<16x16xf32>) + {printFormat = #pto} + return + } +} + +// NATIVE-LABEL: func.func private @tprint_arg( +// NATIVE: pto.tprint ins(%arg0 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tprint_tmp_arg( +// NATIVE: pto.tprint ins(%arg0, %arg1 : !pto.tile_buf, memref<16x16xf32>) +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tprint_arg( +// EMITC: TPRINT( +// EMITC-LABEL: tprint_tmp_arg( +// EMITC: TPRINT( From d83449509dfb41f5e554f596ed20db28e647888b Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 16:30:18 +0800 Subject: [PATCH 48/91] Keep TQuant tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 404e2a9bca..de5872a2ed 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -97,7 +97,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp, pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp, pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp, pto::TCIOp, - pto::TPrintOp>(op)) { + pto::TPrintOp, pto::TQuantOp>(op)) { found = true; return WalkResult::interrupt(); } From 9b218cd0eadd753ea3861ef83962a57b4296014d Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 16:33:45 +0800 Subject: [PATCH 49/91] Keep TTrans tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/trans_tile_native.pto | 33 ++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/trans_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index de5872a2ed..55390fbbcf 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -97,7 +97,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp, pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp, pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp, pto::TCIOp, - pto::TPrintOp, pto::TQuantOp>(op)) { + pto::TPrintOp, pto::TQuantOp, pto::TTransOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/trans_tile_native.pto b/test/lit/pto/trans_tile_native.pto new file mode 100644 index 0000000000..af88a74480 --- /dev/null +++ b/test/lit/pto/trans_tile_native.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: ptoas --pto-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @ttrans_arg( + %src: !pto.tile_buf, + %tmp: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.ttrans ins(%src, %tmp : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @ttrans_arg +// NATIVE-SAME: (%[[SRC:[^:]+]]: !pto.tile_buf, %[[TMP:[^:]+]]: !pto.tile_buf, %[[DST:[^:]+]]: !pto.tile_buf) +// NATIVE: pto.ttrans ins(%[[SRC]], %[[TMP]] : !pto.tile_buf, !pto.tile_buf) outs(%[[DST]] : !pto.tile_buf) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: ttrans_arg( +// EMITC: TTRANS( From 143ec495e1cdf044d8f096f0c1f654082ffb4e65 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 16:36:35 +0800 Subject: [PATCH 50/91] Keep TMrgSort tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/mrgsort_tile_native.pto | 47 ++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/mrgsort_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 55390fbbcf..cf8b07328c 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -97,7 +97,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TFillPadInplaceOp, pto::TLoadOp, pto::TStoreOp, pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp, pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp, pto::TCIOp, - pto::TPrintOp, pto::TQuantOp, pto::TTransOp>(op)) { + pto::TPrintOp, pto::TQuantOp, pto::TTransOp, + pto::TMrgSortOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/mrgsort_tile_native.pto b/test/lit/pto/mrgsort_tile_native.pto new file mode 100644 index 0000000000..5fa7f960c4 --- /dev/null +++ b/test/lit/pto/mrgsort_tile_native.pto @@ -0,0 +1,47 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tmrgsort_format1_arg( + %src: !pto.tile_buf, + %dst: !pto.tile_buf, + %block_len: i32) { + pto.tmrgsort ins(%src, %block_len : !pto.tile_buf, i32) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tmrgsort_format2_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %tmp: !pto.tile_buf, + %dst: !pto.tile_buf, + %executed: vector<4xi16>) { + pto.tmrgsort ins(%src0, %src1, %tmp {exhausted = true} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) + outs(%dst, %executed : !pto.tile_buf, vector<4xi16>) + return + } +} + +// NATIVE-LABEL: func.func private @tmrgsort_format1_arg +// NATIVE: pto.tmrgsort ins(%{{.*}}, %{{.*}} : !pto.tile_buf, i32) outs(%{{.*}} : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tmrgsort_format2_arg +// NATIVE: pto.tmrgsort ins(%{{.*}}, %{{.*}}, %{{.*}} {exhausted = true} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%{{.*}}, %{{.*}} : !pto.tile_buf, vector<4xi16>) +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tmrgsort_format1_arg( +// EMITC: TMRGSORT( +// EMITC-LABEL: tmrgsort_format2_arg( +// EMITC: TMRGSORT<{{.*}}, true>( From 5673c6ef1a31a93e60f885b28855c28f88eb50b7 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 16:40:10 +0800 Subject: [PATCH 51/91] Keep TAssign tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/tassign_tile_native_arg.pto | 32 ++++++++++++++++++++++++ 2 files changed, 33 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/tassign_tile_native_arg.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index cf8b07328c..6aa7e51934 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -98,7 +98,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp, pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp, pto::TCIOp, pto::TPrintOp, pto::TQuantOp, pto::TTransOp, - pto::TMrgSortOp>(op)) { + pto::TMrgSortOp, pto::TAssignOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/tassign_tile_native_arg.pto b/test/lit/pto/tassign_tile_native_arg.pto new file mode 100644 index 0000000000..29b945d619 --- /dev/null +++ b/test/lit/pto/tassign_tile_native_arg.pto @@ -0,0 +1,32 @@ +// 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-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tassign_arg(%tile: !pto.tile_buf) + -> !pto.tile_buf { + %addr = arith.constant 0 : i64 + %rebound = pto.tassign %tile, %addr + : !pto.tile_buf -> !pto.tile_buf + return %rebound : !pto.tile_buf + } +} + +// NATIVE-LABEL: func.func private @tassign_arg +// NATIVE-SAME: (%[[TILE:.*]]: !pto.tile_buf) +// NATIVE: %{{.*}} = pto.tassign %[[TILE]], %{{.*}} : !pto.tile_buf -> !pto.tile_buf +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tassign_arg( +// EMITC: TASSIGN( From 5874cfa898bae37011a5c2211b5ff0a04e8b9b22 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 16:43:01 +0800 Subject: [PATCH 52/91] Keep MGather and MScatter tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/mgather_mscatter_tile_native.pto | 52 +++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/mgather_mscatter_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 6aa7e51934..4e467f25c8 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -98,7 +98,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TColMaxOp, pto::TColMinOp, pto::TColSumOp, pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp, pto::TCIOp, pto::TPrintOp, pto::TQuantOp, pto::TTransOp, - pto::TMrgSortOp, pto::TAssignOp>(op)) { + pto::TMrgSortOp, pto::TAssignOp, pto::MGatherOp, + pto::MScatterOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/mgather_mscatter_tile_native.pto b/test/lit/pto/mgather_mscatter_tile_native.pto new file mode 100644 index 0000000000..72be822671 --- /dev/null +++ b/test/lit/pto/mgather_mscatter_tile_native.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: ptoas --pto-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @mgather_arg( + %mem: memref<1x1x1x8x16xf16, #pto.address_space>, + %idx: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.mgather ins(%mem, %idx : memref<1x1x1x8x16xf16, #pto.address_space>, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {coalesce = #pto, gatherOob = #pto} + return + } + + func.func private @mscatter_arg( + %src: !pto.tile_buf, + %idx: !pto.tile_buf, + %mem: memref<1x1x1x8x16xf16, #pto.address_space>) { + pto.mscatter ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) + outs(%mem : memref<1x1x1x8x16xf16, #pto.address_space>) + {coalesce = #pto, scatterOob = #pto} + return + } +} + +// NATIVE-LABEL: func.func private @mgather_arg +// NATIVE-SAME: memref<1x1x1x8x16xf16, #pto.address_space> +// NATIVE-SAME: !pto.tile_buf +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.mgather +// NATIVE-LABEL: func.func private @mscatter_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE-SAME: !pto.tile_buf +// NATIVE-SAME: memref<1x1x1x8x16xf16, #pto.address_space> +// NATIVE: pto.mscatter +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: mgather_arg( +// EMITC: MGATHER( +// EMITC-LABEL: mscatter_arg( +// EMITC: MSCATTER( From bd19835f6c48dd74054ae8a2dbd88de314db768c Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 16:47:13 +0800 Subject: [PATCH 53/91] Keep subtraction ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/sub_tile_native.pto | 58 ++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/sub_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 4e467f25c8..369129b520 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -99,7 +99,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TCvtOp, pto::TGatherOp, pto::TGatherBOp, pto::TCIOp, pto::TPrintOp, pto::TQuantOp, pto::TTransOp, pto::TMrgSortOp, pto::TAssignOp, pto::MGatherOp, - pto::MScatterOp>(op)) { + pto::MScatterOp, pto::TSubOp, pto::TSubSOp, pto::TSubCOp, + pto::TSubSCOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/sub_tile_native.pto b/test/lit/pto/sub_tile_native.pto new file mode 100644 index 0000000000..676fa74573 --- /dev/null +++ b/test/lit/pto/sub_tile_native.pto @@ -0,0 +1,58 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level3 --pto-arch=a5 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tsub_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tsub ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tsubs_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: i16) { + pto.tsubs ins(%src, %scalar : !pto.tile_buf, i16) outs(%dst : !pto.tile_buf) + return + } + func.func private @tsubc_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %c: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tsubc ins(%a, %b, %c : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tsubsc_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: i16) { + pto.tsubsc ins(%a, %scalar, %b : !pto.tile_buf, i16, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: func.func private @tsub_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tsub +// NATIVE-LABEL: func.func private @tsubs_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tsubs +// NATIVE-LABEL: func.func private @tsubc_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tsubc +// NATIVE-LABEL: func.func private @tsubsc_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tsubsc +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tsub_arg( +// EMITC: TSUB( +// EMITC-LABEL: tsubs_arg( +// EMITC: TSUBS( +// EMITC-LABEL: tsubc_arg( +// EMITC: TSUB( +// EMITC: TADD( +// EMITC-LABEL: tsubsc_arg( +// EMITC: TSUBS( +// EMITC: TADD( From a6689002d0e8498676ae8a4fa27bfc397058e089 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 16:49:57 +0800 Subject: [PATCH 54/91] Keep basic floating-point ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/basic_float_tile_native.pto | 62 ++++++++++++++++++++++++ 2 files changed, 64 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/basic_float_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 369129b520..44d96ca1c4 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -100,7 +100,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TPrintOp, pto::TQuantOp, pto::TTransOp, pto::TMrgSortOp, pto::TAssignOp, pto::MGatherOp, pto::MScatterOp, pto::TSubOp, pto::TSubSOp, pto::TSubCOp, - pto::TSubSCOp>(op)) { + pto::TSubSCOp, pto::TReluOp, pto::TRecipOp, pto::TSqrtOp, + pto::TFModOp, pto::TFModSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/basic_float_tile_native.pto b/test/lit/pto/basic_float_tile_native.pto new file mode 100644 index 0000000000..35de5643af --- /dev/null +++ b/test/lit/pto/basic_float_tile_native.pto @@ -0,0 +1,62 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @trelu_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trelu ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trecip_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trecip ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tsqrt_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tsqrt ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tfmod_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tfmod ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tfmods_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: f32) { + pto.tfmods ins(%src, %scalar : !pto.tile_buf, f32) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @trelu_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.trelu +// NATIVE-LABEL: @trecip_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.trecip +// NATIVE-LABEL: @tsqrt_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tsqrt +// NATIVE-LABEL: @tfmod_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tfmod +// NATIVE-LABEL: @tfmods_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tfmods +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: trelu_arg( +// EMITC: TRELU( +// EMITC-LABEL: trecip_arg( +// EMITC: TRECIP( +// EMITC-LABEL: tsqrt_arg( +// EMITC: TSQRT( +// EMITC-LABEL: tfmod_arg( +// EMITC: TFMOD( +// EMITC-LABEL: tfmods_arg( +// EMITC: TFMODS( From d0716680eb7b52be094ca76133981fe41d5bb6cc Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 16:53:48 +0800 Subject: [PATCH 55/91] Keep shift ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/shift_tile_native.pto | 53 ++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/shift_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 44d96ca1c4..2f8d4292d9 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -101,7 +101,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TMrgSortOp, pto::TAssignOp, pto::MGatherOp, pto::MScatterOp, pto::TSubOp, pto::TSubSOp, pto::TSubCOp, pto::TSubSCOp, pto::TReluOp, pto::TRecipOp, pto::TSqrtOp, - pto::TFModOp, pto::TFModSOp>(op)) { + pto::TFModOp, pto::TFModSOp, pto::TShlOp, pto::TShlSOp, + pto::TShrOp, pto::TShrSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/shift_tile_native.pto b/test/lit/pto/shift_tile_native.pto new file mode 100644 index 0000000000..0470930f7f --- /dev/null +++ b/test/lit/pto/shift_tile_native.pto @@ -0,0 +1,53 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @tshl_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tshl ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tshr_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tshr ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tshls_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf, %shift: i16) { + pto.tshls ins(%src, %shift : !pto.tile_buf, i16) outs(%dst : !pto.tile_buf) + return + } + func.func private @tshrs_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf, %shift: i16) { + pto.tshrs ins(%src, %shift : !pto.tile_buf, i16) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tshl_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tshl +// NATIVE-LABEL: @tshr_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tshr +// NATIVE-LABEL: @tshls_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tshls +// NATIVE-LABEL: @tshrs_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tshrs +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tshl_arg( +// EMITC: TSHL( +// EMITC-LABEL: tshr_arg( +// EMITC: TSHR( +// EMITC-LABEL: tshls_arg( +// EMITC: TSHLS( +// EMITC-LABEL: tshrs_arg( +// EMITC: TSHRS( From 7a030f073d1677eec53849bc3e675af288ff0fad Mon Sep 17 00:00:00 2001 From: FangRui Date: Sat, 18 Jul 2026 16:56:26 +0800 Subject: [PATCH 56/91] Keep TPartMax and TPartMin tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 ++- test/lit/pto/part_arithmetic_tile_native.pto | 26 ++++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 2f8d4292d9..f9ace44310 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -102,7 +102,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::MScatterOp, pto::TSubOp, pto::TSubSOp, pto::TSubCOp, pto::TSubSCOp, pto::TReluOp, pto::TRecipOp, pto::TSqrtOp, pto::TFModOp, pto::TFModSOp, pto::TShlOp, pto::TShlSOp, - pto::TShrOp, pto::TShrSOp>(op)) { + pto::TShrOp, pto::TShrSOp, pto::TPartMaxOp, + pto::TPartMinOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/part_arithmetic_tile_native.pto b/test/lit/pto/part_arithmetic_tile_native.pto index e21d0b3be0..2282b5573e 100644 --- a/test/lit/pto/part_arithmetic_tile_native.pto +++ b/test/lit/pto/part_arithmetic_tile_native.pto @@ -30,12 +30,34 @@ module { outs(%dst : !pto.tile_buf) return } + + func.func private @tpartmax_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tpartmax ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } + + func.func private @tpartmin_arg( + %src0: !pto.tile_buf, + %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tpartmin ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + return + } } // NATIVE-LABEL: func.func private @tpartadd_arg( // NATIVE: pto.tpartadd ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) // NATIVE-LABEL: func.func private @tpartmul_arg( // NATIVE: pto.tpartmul ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tpartmax_arg( +// NATIVE: pto.tpartmax ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tpartmin_arg( +// NATIVE: pto.tpartmin ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) // NATIVE-NOT: memref< // NATIVE-NOT: pto.bind_tile @@ -43,3 +65,7 @@ module { // EMITC: TPARTADD( // EMITC-LABEL: tpartmul_arg( // EMITC: TPARTMUL( +// EMITC-LABEL: tpartmax_arg( +// EMITC: TPARTMAX( +// EMITC-LABEL: tpartmin_arg( +// EMITC: TPARTMIN( From 8ea20356dae250fdfd35df20e8397d27ea88d5a1 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 20:38:37 +0800 Subject: [PATCH 57/91] Keep remaining TColExpand ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/colexpand_tile_native.pto | 41 ++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index f9ace44310..20dbed7c17 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -103,7 +103,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TSubSCOp, pto::TReluOp, pto::TRecipOp, pto::TSqrtOp, pto::TFModOp, pto::TFModSOp, pto::TShlOp, pto::TShlSOp, pto::TShrOp, pto::TShrSOp, pto::TPartMaxOp, - pto::TPartMinOp>(op)) { + pto::TPartMinOp, pto::TColExpandAddOp, pto::TColExpandSubOp, + pto::TColExpandDivOp, pto::TColExpandExpdifOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/colexpand_tile_native.pto b/test/lit/pto/colexpand_tile_native.pto index 75ef9fb8d4..a30a107fdd 100644 --- a/test/lit/pto/colexpand_tile_native.pto +++ b/test/lit/pto/colexpand_tile_native.pto @@ -47,6 +47,31 @@ module { outs(%dst : !pto.tile_buf) return } + + func.func private @tcolexpandadd_arg( + %src0: !pto.tile_buf, %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolexpandadd ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tcolexpandsub_arg( + %src0: !pto.tile_buf, %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolexpandsub ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tcolexpanddiv_arg( + %src0: !pto.tile_buf, %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolexpanddiv ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tcolexpandexpdif_arg( + %src0: !pto.tile_buf, %src1: !pto.tile_buf, + %dst: !pto.tile_buf) { + pto.tcolexpandexpdif ins(%src0, %src1 : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } } // NATIVE-LABEL: func.func private @tcolexpand_arg( @@ -57,6 +82,14 @@ module { // NATIVE: pto.tcolexpandmax ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) // NATIVE-LABEL: func.func private @tcolexpandmin_arg( // NATIVE: pto.tcolexpandmin ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) +// NATIVE-LABEL: func.func private @tcolexpandadd_arg( +// NATIVE: pto.tcolexpandadd +// NATIVE-LABEL: func.func private @tcolexpandsub_arg( +// NATIVE: pto.tcolexpandsub +// NATIVE-LABEL: func.func private @tcolexpanddiv_arg( +// NATIVE: pto.tcolexpanddiv +// NATIVE-LABEL: func.func private @tcolexpandexpdif_arg( +// NATIVE: pto.tcolexpandexpdif // NATIVE-NOT: memref< // NATIVE-NOT: pto.bind_tile @@ -68,3 +101,11 @@ module { // EMITC: TCOLEXPANDMAX( // EMITC-LABEL: tcolexpandmin_arg( // EMITC: TCOLEXPANDMIN( +// EMITC-LABEL: tcolexpandadd_arg( +// EMITC: TCOLEXPANDADD( +// EMITC-LABEL: tcolexpandsub_arg( +// EMITC: TCOLEXPANDSUB( +// EMITC-LABEL: tcolexpanddiv_arg( +// EMITC: TCOLEXPANDDIV( +// EMITC-LABEL: tcolexpandexpdif_arg( +// EMITC: TCOLEXPANDEXPDIF( From f853e9d7e3e96f9ced3b2ab46e77a84abcb23c88 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 20:43:59 +0800 Subject: [PATCH 58/91] Keep TRowExpand ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 6 +- test/lit/pto/rowexpand_tile_native.pto | 82 ++++++++++++++++++++++++++ 2 files changed, 87 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/rowexpand_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 20dbed7c17..b524ddae46 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -104,7 +104,11 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TFModOp, pto::TFModSOp, pto::TShlOp, pto::TShlSOp, pto::TShrOp, pto::TShrSOp, pto::TPartMaxOp, pto::TPartMinOp, pto::TColExpandAddOp, pto::TColExpandSubOp, - pto::TColExpandDivOp, pto::TColExpandExpdifOp>(op)) { + pto::TColExpandDivOp, pto::TColExpandExpdifOp, + pto::TRowExpandOp, pto::TRowExpandAddOp, + pto::TRowExpandSubOp, pto::TRowExpandMulOp, + pto::TRowExpandDivOp, pto::TRowExpandMaxOp, + pto::TRowExpandMinOp, pto::TRowExpandExpdifOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/rowexpand_tile_native.pto b/test/lit/pto/rowexpand_tile_native.pto new file mode 100644 index 0000000000..1eb14eaeaf --- /dev/null +++ b/test/lit/pto/rowexpand_tile_native.pto @@ -0,0 +1,82 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @trowexpand_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowexpand ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowexpandadd_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowexpandadd ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowexpandsub_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowexpandsub ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowexpandmul_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowexpandmul ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowexpanddiv_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowexpanddiv ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowexpandmax_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowexpandmax ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowexpandmin_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowexpandmin ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowexpandexpdif_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowexpandexpdif ins(%a, %b : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @trowexpand_arg +// NATIVE-SAME: !pto.tile_buf< +// NATIVE: pto.trowexpand +// NATIVE-LABEL: @trowexpandadd_arg +// NATIVE: pto.trowexpandadd +// NATIVE-LABEL: @trowexpandsub_arg +// NATIVE: pto.trowexpandsub +// NATIVE-LABEL: @trowexpandmul_arg +// NATIVE: pto.trowexpandmul +// NATIVE-LABEL: @trowexpanddiv_arg +// NATIVE: pto.trowexpanddiv +// NATIVE-LABEL: @trowexpandmax_arg +// NATIVE: pto.trowexpandmax +// NATIVE-LABEL: @trowexpandmin_arg +// NATIVE: pto.trowexpandmin +// NATIVE-LABEL: @trowexpandexpdif_arg +// NATIVE: pto.trowexpandexpdif +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: trowexpand_arg( +// EMITC: TROWEXPAND( +// EMITC-LABEL: trowexpandadd_arg( +// EMITC: TROWEXPANDADD( +// EMITC-LABEL: trowexpandsub_arg( +// EMITC: TROWEXPANDSUB( +// EMITC-LABEL: trowexpandmul_arg( +// EMITC: TROWEXPANDMUL( +// EMITC-LABEL: trowexpanddiv_arg( +// EMITC: TROWEXPANDDIV( +// EMITC-LABEL: trowexpandmax_arg( +// EMITC: TROWEXPANDMAX( +// EMITC-LABEL: trowexpandmin_arg( +// EMITC: TROWEXPANDMIN( +// EMITC-LABEL: trowexpandexpdif_arg( +// EMITC: TROWEXPANDEXPDIF( From b20873690479a5105de32553812952d91e9d20e9 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 20:47:02 +0800 Subject: [PATCH 59/91] Keep reduction ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 4 +- test/lit/pto/reduction_tile_native.pto | 57 ++++++++++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/reduction_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index b524ddae46..a80e38ed3e 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -108,7 +108,9 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TRowExpandOp, pto::TRowExpandAddOp, pto::TRowExpandSubOp, pto::TRowExpandMulOp, pto::TRowExpandDivOp, pto::TRowExpandMaxOp, - pto::TRowExpandMinOp, pto::TRowExpandExpdifOp>(op)) { + pto::TRowExpandMinOp, pto::TRowExpandExpdifOp, + pto::TRowMaxOp, pto::TRowMinOp, pto::TRowSumOp, + pto::TRowProdOp, pto::TColProdOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/reduction_tile_native.pto b/test/lit/pto/reduction_tile_native.pto new file mode 100644 index 0000000000..068644d33e --- /dev/null +++ b/test/lit/pto/reduction_tile_native.pto @@ -0,0 +1,57 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @trowmax_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowmax ins(%src, %tmp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowmin_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowmin ins(%src, %tmp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowsum_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowsum ins(%src, %tmp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowprod_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowprod ins(%src, %tmp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tcolprod_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tcolprod ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @trowmax_arg +// NATIVE: pto.trowmax +// NATIVE-LABEL: @trowmin_arg +// NATIVE: pto.trowmin +// NATIVE-LABEL: @trowsum_arg +// NATIVE: pto.trowsum +// NATIVE-LABEL: @trowprod_arg +// NATIVE: pto.trowprod +// NATIVE-LABEL: @tcolprod_arg +// NATIVE: pto.tcolprod +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: trowmax_arg( +// EMITC: TROWMAX( +// EMITC-LABEL: trowmin_arg( +// EMITC: TROWMIN( +// EMITC-LABEL: trowsum_arg( +// EMITC: TROWSUM( +// EMITC-LABEL: trowprod_arg( +// EMITC: TROWPROD( +// EMITC-LABEL: tcolprod_arg( +// EMITC: TCOLPROD( From 2259172c1bfcd506bf24ba3b111929730790c556 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 20:50:23 +0800 Subject: [PATCH 60/91] Keep arg-reduction ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 4 +- test/lit/pto/arg_reduction_tile_native.pto | 65 ++++++++++++++++++++++ 2 files changed, 68 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/arg_reduction_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index a80e38ed3e..10416caccd 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -110,7 +110,9 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TRowExpandDivOp, pto::TRowExpandMaxOp, pto::TRowExpandMinOp, pto::TRowExpandExpdifOp, pto::TRowMaxOp, pto::TRowMinOp, pto::TRowSumOp, - pto::TRowProdOp, pto::TColProdOp>(op)) { + pto::TRowProdOp, pto::TColProdOp, pto::TColArgMaxOp, + pto::TColArgMinOp, pto::TRowArgMaxOp, pto::TRowArgMinOp, + pto::TPartArgMaxOp, pto::TPartArgMinOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/arg_reduction_tile_native.pto b/test/lit/pto/arg_reduction_tile_native.pto new file mode 100644 index 0000000000..3a54887913 --- /dev/null +++ b/test/lit/pto/arg_reduction_tile_native.pto @@ -0,0 +1,65 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @tcolargmax_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tcolargmax ins(%src, %tmp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tcolargmin_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tcolargmin ins(%src, %tmp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowargmax_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowargmax ins(%src, %tmp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trowargmin_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trowargmin ins(%src, %tmp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tpartargmax_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %ai: !pto.tile_buf, %bi: !pto.tile_buf, %dst: !pto.tile_buf, %di: !pto.tile_buf) { + pto.tpartargmax ins(%a, %b, %ai, %bi : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst, %di : !pto.tile_buf, !pto.tile_buf) + return + } + func.func private @tpartargmin_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %ai: !pto.tile_buf, %bi: !pto.tile_buf, %dst: !pto.tile_buf, %di: !pto.tile_buf) { + pto.tpartargmin ins(%a, %b, %ai, %bi : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst, %di : !pto.tile_buf, !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tcolargmax_arg +// NATIVE: pto.tcolargmax +// NATIVE-LABEL: @tcolargmin_arg +// NATIVE: pto.tcolargmin +// NATIVE-LABEL: @trowargmax_arg +// NATIVE: pto.trowargmax +// NATIVE-LABEL: @trowargmin_arg +// NATIVE: pto.trowargmin +// NATIVE-LABEL: @tpartargmax_arg +// NATIVE: pto.tpartargmax +// NATIVE-LABEL: @tpartargmin_arg +// NATIVE: pto.tpartargmin +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tcolargmax_arg( +// EMITC: TCOLARGMAX( +// EMITC-LABEL: tcolargmin_arg( +// EMITC: TCOLARGMIN( +// EMITC-LABEL: trowargmax_arg( +// EMITC: TROWARGMAX( +// EMITC-LABEL: trowargmin_arg( +// EMITC: TROWARGMIN( +// EMITC-LABEL: tpartargmax_arg( +// EMITC: TPARTARGMAX( +// EMITC-LABEL: tpartargmin_arg( +// EMITC: TPARTARGMIN( From e2f682dcf676233e0fb0c6f0500da70a6c8c6202 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 20:53:22 +0800 Subject: [PATCH 61/91] Keep pow and rsqrt ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/pow_rsqrt_tile_native.pto | 41 ++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/pow_rsqrt_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 10416caccd..7a26506d2c 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -112,7 +112,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TRowMaxOp, pto::TRowMinOp, pto::TRowSumOp, pto::TRowProdOp, pto::TColProdOp, pto::TColArgMaxOp, pto::TColArgMinOp, pto::TRowArgMaxOp, pto::TRowArgMinOp, - pto::TPartArgMaxOp, pto::TPartArgMinOp>(op)) { + pto::TPartArgMaxOp, pto::TPartArgMinOp, pto::TPowOp, + pto::TPowSOp, pto::TRsqrtOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/pow_rsqrt_tile_native.pto b/test/lit/pto/pow_rsqrt_tile_native.pto new file mode 100644 index 0000000000..b2a99514e6 --- /dev/null +++ b/test/lit/pto/pow_rsqrt_tile_native.pto @@ -0,0 +1,41 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @tpow_arg(%base: !pto.tile_buf, %exp: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tpow ins(%base, %exp, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tpows_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: f32) { + pto.tpows ins(%src, %scalar, %tmp : !pto.tile_buf, f32, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trsqrt_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trsqrt ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tpow_arg +// NATIVE: pto.tpow +// NATIVE-LABEL: @tpows_arg +// NATIVE: pto.tpows +// NATIVE-LABEL: @trsqrt_arg +// NATIVE: pto.trsqrt +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tpow_arg( +// EMITC: TPOW( +// EMITC-LABEL: tpows_arg( +// EMITC: TPOWS( +// EMITC-LABEL: trsqrt_arg( +// EMITC: TRSQRT( From e3e619abd4753666d9d48fbbde8fc16e2d786b95 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 20:56:16 +0800 Subject: [PATCH 62/91] Keep scratch elementwise ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- .../pto/scratch_elementwise_tile_native.pto | 41 +++++++++++++++++++ 2 files changed, 43 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/scratch_elementwise_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 7a26506d2c..7cc86c9219 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -113,7 +113,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TRowProdOp, pto::TColProdOp, pto::TColArgMaxOp, pto::TColArgMinOp, pto::TRowArgMaxOp, pto::TRowArgMinOp, pto::TPartArgMaxOp, pto::TPartArgMinOp, pto::TPowOp, - pto::TPowSOp, pto::TRsqrtOp>(op)) { + pto::TPowSOp, pto::TRsqrtOp, pto::TPReluOp, + pto::TRemOp, pto::TRemSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/scratch_elementwise_tile_native.pto b/test/lit/pto/scratch_elementwise_tile_native.pto new file mode 100644 index 0000000000..af45853b52 --- /dev/null +++ b/test/lit/pto/scratch_elementwise_tile_native.pto @@ -0,0 +1,41 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @tprelu_arg(%src: !pto.tile_buf, %slope: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tprelu ins(%src, %slope, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trem_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.trem ins(%a, %b, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trems_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: f32) { + pto.trems ins(%src, %scalar, %tmp : !pto.tile_buf, f32, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tprelu_arg +// NATIVE: pto.tprelu +// NATIVE-LABEL: @trem_arg +// NATIVE: pto.trem +// NATIVE-LABEL: @trems_arg +// NATIVE: pto.trems +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tprelu_arg( +// EMITC: TPRELU( +// EMITC-LABEL: trem_arg( +// EMITC: TREM( +// EMITC-LABEL: trems_arg( +// EMITC: TREMS( From 1dfd65de4a2233ea1317497e61b2be79abef2692 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 20:59:14 +0800 Subject: [PATCH 63/91] Keep select ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 ++- test/lit/pto/select_tile_native.pto | 33 ++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/select_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 7cc86c9219..f180aed5b5 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -114,7 +114,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TColArgMinOp, pto::TRowArgMaxOp, pto::TRowArgMinOp, pto::TPartArgMaxOp, pto::TPartArgMinOp, pto::TPowOp, pto::TPowSOp, pto::TRsqrtOp, pto::TPReluOp, - pto::TRemOp, pto::TRemSOp>(op)) { + pto::TRemOp, pto::TRemSOp, pto::TSelOp, + pto::TSelSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/select_tile_native.pto b/test/lit/pto/select_tile_native.pto new file mode 100644 index 0000000000..c79fde6bc4 --- /dev/null +++ b/test/lit/pto/select_tile_native.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: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +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) + return + } + func.func private @tsels_arg(%mask: !pto.tile_buf, %src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: i16) { + pto.tsels ins(%mask, %src, %tmp, %scalar : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, i16) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tsel_arg +// NATIVE: pto.tsel +// NATIVE-LABEL: @tsels_arg +// NATIVE: pto.tsels +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tsel_arg( +// EMITC: TSEL( +// EMITC-LABEL: tsels_arg( +// EMITC: TSELS( From fa6fe5313e5167cbcd5593c427cbb76ae22e17f5 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 21:02:03 +0800 Subject: [PATCH 64/91] Keep XOR ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/xor_tile_native.pto | 33 ++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/xor_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index f180aed5b5..42e6293266 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -115,7 +115,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TPartArgMaxOp, pto::TPartArgMinOp, pto::TPowOp, pto::TPowSOp, pto::TRsqrtOp, pto::TPReluOp, pto::TRemOp, pto::TRemSOp, pto::TSelOp, - pto::TSelSOp>(op)) { + pto::TSelSOp, pto::TXorOp, pto::TXorSOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/xor_tile_native.pto b/test/lit/pto/xor_tile_native.pto new file mode 100644 index 0000000000..c728792ef4 --- /dev/null +++ b/test/lit/pto/xor_tile_native.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: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @txor_arg(%a: !pto.tile_buf, %b: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.txor ins(%a, %b, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @txors_arg(%src: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: i16) { + pto.txors ins(%src, %scalar, %tmp : !pto.tile_buf, i16, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @txor_arg +// NATIVE: pto.txor +// NATIVE-LABEL: @txors_arg +// NATIVE: pto.txors +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: txor_arg( +// EMITC: TXOR( +// EMITC-LABEL: txors_arg( +// EMITC: TXORS( From 6c583dbec1a7e24b5da4043e01199b4e789fa017 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 21:04:50 +0800 Subject: [PATCH 65/91] Keep TSort32 tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 ++- test/lit/pto/sort32_tile_native.pto | 33 ++++++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/sort32_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 42e6293266..ca21f8cb68 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -115,7 +115,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TPartArgMaxOp, pto::TPartArgMinOp, pto::TPowOp, pto::TPowSOp, pto::TRsqrtOp, pto::TPReluOp, pto::TRemOp, pto::TRemSOp, pto::TSelOp, - pto::TSelSOp, pto::TXorOp, pto::TXorSOp>(op)) { + pto::TSelSOp, pto::TXorOp, pto::TXorSOp, + pto::TSort32Op>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/sort32_tile_native.pto b/test/lit/pto/sort32_tile_native.pto new file mode 100644 index 0000000000..d3d0648a85 --- /dev/null +++ b/test/lit/pto/sort32_tile_native.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: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @tsort32_arg(%src: !pto.tile_buf, %idx: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tsort32 ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tsort32_tmp_arg(%src: !pto.tile_buf, %idx: !pto.tile_buf, %tmp: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tsort32 ins(%src, %idx, %tmp : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tsort32_arg +// NATIVE: pto.tsort32 +// NATIVE-LABEL: @tsort32_tmp_arg +// NATIVE: pto.tsort32 +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tsort32_arg( +// EMITC: TSORT32( +// EMITC-LABEL: tsort32_tmp_arg( +// EMITC: TSORT32( From efa96b3c8d7b5166db00e6f9c8bfcb0afe922eb0 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 21:08:30 +0800 Subject: [PATCH 66/91] Keep TQuantMx tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/quant_mx_tile_native.pto | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/quant_mx_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index ca21f8cb68..a37ad4f924 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -116,7 +116,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TPowSOp, pto::TRsqrtOp, pto::TPReluOp, pto::TRemOp, pto::TRemSOp, pto::TSelOp, pto::TSelSOp, pto::TXorOp, pto::TXorSOp, - pto::TSort32Op>(op)) { + pto::TSort32Op, pto::TQuantMxOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/quant_mx_tile_native.pto b/test/lit/pto/quant_mx_tile_native.pto new file mode 100644 index 0000000000..2d154a36d5 --- /dev/null +++ b/test/lit/pto/quant_mx_tile_native.pto @@ -0,0 +1,26 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tquant_mx_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf, %exp: !pto.tile_buf, %max: !pto.tile_buf, %scaling: !pto.tile_buf) { + pto.tquant.mx ins(%src : !pto.tile_buf) outs(%dst, %exp, %max, %scaling : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) {quant_type = #pto} + return + } +} + +// NATIVE-LABEL: @tquant_mx_arg +// NATIVE-SAME: !pto.tile_buf< +// NATIVE: pto.tquant.mx +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tquant_mx_arg( +// EMITC: TQUANT Date: Sun, 19 Jul 2026 21:12:06 +0800 Subject: [PATCH 67/91] Keep TAxpy and TDequant tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 ++- test/lit/pto/axpy_dequant_tile_native.pto | 33 +++++++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/axpy_dequant_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index a37ad4f924..fce2a066c2 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -116,7 +116,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TPowSOp, pto::TRsqrtOp, pto::TPReluOp, pto::TRemOp, pto::TRemSOp, pto::TSelOp, pto::TSelSOp, pto::TXorOp, pto::TXorSOp, - pto::TSort32Op, pto::TQuantMxOp>(op)) { + pto::TSort32Op, pto::TQuantMxOp, pto::TAxpyOp, + pto::TDequantOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/axpy_dequant_tile_native.pto b/test/lit/pto/axpy_dequant_tile_native.pto new file mode 100644 index 0000000000..d43dbc27da --- /dev/null +++ b/test/lit/pto/axpy_dequant_tile_native.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: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %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 + +module { + func.func private @taxpy_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf, %scalar: f16) { + pto.taxpy ins(%src, %scalar : !pto.tile_buf, f16) outs(%dst : !pto.tile_buf) + return + } + func.func private @tdequant_arg(%src: !pto.tile_buf, %scale: !pto.tile_buf, %offset: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tdequant ins(%src, %scale, %offset : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @taxpy_arg +// NATIVE: pto.taxpy +// NATIVE-LABEL: @tdequant_arg +// NATIVE: pto.tdequant +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: taxpy_arg( +// EMITC: TAXPY( +// EMITC-LABEL: tdequant_arg( +// EMITC: TDEQUANT< From bd65abe789879fbc524eb1f396946d019473d06f Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 21:15:44 +0800 Subject: [PATCH 68/91] Keep tile movement metadata ops native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- .../lit/pto/movement_metadata_tile_native.pto | 33 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/movement_metadata_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index fce2a066c2..0b3bec1602 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -117,7 +117,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TRemOp, pto::TRemSOp, pto::TSelOp, pto::TSelSOp, pto::TXorOp, pto::TXorSOp, pto::TSort32Op, pto::TQuantMxOp, pto::TAxpyOp, - pto::TDequantOp>(op)) { + pto::TDequantOp, pto::TFillPadExpandOp, + pto::TGetScaleAddrOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/movement_metadata_tile_native.pto b/test/lit/pto/movement_metadata_tile_native.pto new file mode 100644 index 0000000000..4c0e74eee2 --- /dev/null +++ b/test/lit/pto/movement_metadata_tile_native.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: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tfillpad_expand_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tfillpad_expand ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @tget_scale_addr_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tget_scale_addr ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tfillpad_expand_arg +// NATIVE: pto.tfillpad_expand +// NATIVE-LABEL: @tget_scale_addr_arg +// NATIVE: pto.tget_scale_addr +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tfillpad_expand_arg( +// EMITC: TFILLPAD_EXPAND( +// EMITC-LABEL: tget_scale_addr_arg( +// EMITC: GetScaleAddr From 18e32aefa0ae99152521692a29d264979371e914 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 21:18:28 +0800 Subject: [PATCH 69/91] Keep TMov tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/mov_tile_native.pto | 26 ++++++++++++++++++++++++++ 2 files changed, 27 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/mov_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 0b3bec1602..826bfcd3c0 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -118,7 +118,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TSelSOp, pto::TXorOp, pto::TXorSOp, pto::TSort32Op, pto::TQuantMxOp, pto::TAxpyOp, pto::TDequantOp, pto::TFillPadExpandOp, - pto::TGetScaleAddrOp>(op)) { + pto::TGetScaleAddrOp, pto::TMovOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/mov_tile_native.pto b/test/lit/pto/mov_tile_native.pto new file mode 100644 index 0000000000..b8b77e809e --- /dev/null +++ b/test/lit/pto/mov_tile_native.pto @@ -0,0 +1,26 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tmov_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tmov ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tmov_arg +// NATIVE-SAME: !pto.tile_buf +// NATIVE: pto.tmov +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tmov_arg( +// EMITC: TMOV( From 920f18d9a71adb047e01f41772c0a17753bf3d94 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 21:22:14 +0800 Subject: [PATCH 70/91] Keep TStoreFP tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 ++- test/lit/pto/store_fp_tile_native.pto | 27 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/store_fp_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 826bfcd3c0..5b41d636f4 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -118,7 +118,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TSelSOp, pto::TXorOp, pto::TXorSOp, pto::TSort32Op, pto::TQuantMxOp, pto::TAxpyOp, pto::TDequantOp, pto::TFillPadExpandOp, - pto::TGetScaleAddrOp, pto::TMovOp>(op)) { + pto::TGetScaleAddrOp, pto::TMovOp, + pto::TStoreFPOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/store_fp_tile_native.pto b/test/lit/pto/store_fp_tile_native.pto new file mode 100644 index 0000000000..72b30a402d --- /dev/null +++ b/test/lit/pto/store_fp_tile_native.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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tstore_fp_arg(%src: !pto.tile_buf, %fp: !pto.tile_buf, %dst: memref<1x32xf16, #pto.address_space>) { + pto.tstore_fp ins(%src, %fp : !pto.tile_buf, !pto.tile_buf) outs(%dst : memref<1x32xf16, #pto.address_space>) + return + } +} + +// NATIVE-LABEL: @tstore_fp_arg +// NATIVE-SAME: !pto.tile_buf> +// NATIVE: pto.tstore_fp +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: tstore_fp_arg( +// EMITC: TSTORE_FP( From cc8e53915a1082218c19add3cec3845bf404fcee Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 21:26:09 +0800 Subject: [PATCH 71/91] Keep vector specialty ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 +- test/lit/pto/vector_specialty_tile_native.pto | 49 +++++++++++++++++++ 2 files changed, 51 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/vector_specialty_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 5b41d636f4..2b4528681e 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -119,7 +119,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TSort32Op, pto::TQuantMxOp, pto::TAxpyOp, pto::TDequantOp, pto::TFillPadExpandOp, pto::TGetScaleAddrOp, pto::TMovOp, - pto::TStoreFPOp>(op)) { + pto::TStoreFPOp, pto::THistogramOp, pto::TRandomOp, + pto::TScatterOp, pto::TTriOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/vector_specialty_tile_native.pto b/test/lit/pto/vector_specialty_tile_native.pto new file mode 100644 index 0000000000..601c6ce412 --- /dev/null +++ b/test/lit/pto/vector_specialty_tile_native.pto @@ -0,0 +1,49 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC + +module { + func.func private @thistogram_arg(%src: !pto.tile_buf, %idx: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.thistogram ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @trandom_arg(%dst: !pto.tile_buf, %k0: i32, %k1: i32, %c0: i32, %c1: i32, %c2: i32, %c3: i32) { + pto.trandom ins(%k0, %k1, %c0, %c1, %c2, %c3 {rounds = 7 : i32} : i32, i32, i32, i32, i32, i32) outs(%dst : !pto.tile_buf) + return + } + func.func private @tscatter_arg(%src: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tscatter ins(%src, {maskPattern = #pto.mask_pattern} : !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } + func.func private @ttri_arg(%dst: !pto.tile_buf, %diag: i32) { + pto.ttri ins(%diag {upperOrLower = 1 : i32} : i32) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @thistogram_arg +// NATIVE: pto.thistogram +// NATIVE-LABEL: @trandom_arg +// NATIVE: pto.trandom +// NATIVE-LABEL: @tscatter_arg +// NATIVE: pto.tscatter +// NATIVE-LABEL: @ttri_arg +// NATIVE: pto.ttri +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: thistogram_arg( +// EMITC: THISTOGRAM< +// EMITC-LABEL: trandom_arg( +// EMITC: PTOAS__TRANDOM<7>( +// EMITC-LABEL: tscatter_arg( +// EMITC: TSCATTER( +// EMITC-LABEL: ttri_arg( +// EMITC: TTRI< From d399f3681de4c1bebad5b49b907f98c23c6887c7 Mon Sep 17 00:00:00 2001 From: FangRui Date: Sun, 19 Jul 2026 21:29:45 +0800 Subject: [PATCH 72/91] Keep matmul ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 4 +++- test/lit/pto/matmul_tile_native.pto | 28 ++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/matmul_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 2b4528681e..886c19458f 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -120,7 +120,9 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TDequantOp, pto::TFillPadExpandOp, pto::TGetScaleAddrOp, pto::TMovOp, pto::TStoreFPOp, pto::THistogramOp, pto::TRandomOp, - pto::TScatterOp, pto::TTriOp>(op)) { + pto::TScatterOp, pto::TTriOp, pto::TMatmulOp, + pto::TMatmulAccOp, pto::TMatmulBiasOp, pto::TMatmulMxOp, + pto::TMatmulMxAccOp, pto::TMatmulMxBiasOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/matmul_tile_native.pto b/test/lit/pto/matmul_tile_native.pto new file mode 100644 index 0000000000..ab70cf51de --- /dev/null +++ b/test/lit/pto/matmul_tile_native.pto @@ -0,0 +1,28 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tmatmul_arg(%lhs: !pto.tile_buf, %rhs: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tmatmul ins(%lhs, %rhs : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tmatmul_arg +// NATIVE-SAME: !pto.tile_buf Date: Mon, 20 Jul 2026 09:25:44 +0800 Subject: [PATCH 73/91] Fix matmul tile-native type checks --- test/lit/pto/matmul_tile_native.pto | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/lit/pto/matmul_tile_native.pto b/test/lit/pto/matmul_tile_native.pto index ab70cf51de..574ff3af1c 100644 --- a/test/lit/pto/matmul_tile_native.pto +++ b/test/lit/pto/matmul_tile_native.pto @@ -17,9 +17,9 @@ module { } // NATIVE-LABEL: @tmatmul_arg -// NATIVE-SAME: !pto.tile_buf Date: Mon, 20 Jul 2026 09:29:11 +0800 Subject: [PATCH 74/91] Keep gemv ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 4 +++- test/lit/pto/gemv_tile_native.pto | 28 ++++++++++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/gemv_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 886c19458f..dfb440009d 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -122,7 +122,9 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TStoreFPOp, pto::THistogramOp, pto::TRandomOp, pto::TScatterOp, pto::TTriOp, pto::TMatmulOp, pto::TMatmulAccOp, pto::TMatmulBiasOp, pto::TMatmulMxOp, - pto::TMatmulMxAccOp, pto::TMatmulMxBiasOp>(op)) { + pto::TMatmulMxAccOp, pto::TMatmulMxBiasOp, pto::TGemvOp, + pto::TGemvAccOp, pto::TGemvBiasOp, pto::TGemvMxOp, + pto::TGemvMxAccOp, pto::TGemvMxBiasOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/gemv_tile_native.pto b/test/lit/pto/gemv_tile_native.pto new file mode 100644 index 0000000000..4e506374dc --- /dev/null +++ b/test/lit/pto/gemv_tile_native.pto @@ -0,0 +1,28 @@ +// 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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tgemv_arg(%lhs: !pto.tile_buf, %rhs: !pto.tile_buf, %dst: !pto.tile_buf) { + pto.tgemv ins(%lhs, %rhs : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tgemv_arg +// NATIVE-SAME: !pto.tile_buf Date: Mon, 20 Jul 2026 10:14:50 +0800 Subject: [PATCH 75/91] Keep TPrefetch tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 ++- test/lit/pto/tprefetch_tile_native.pto | 27 ++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/tprefetch_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index dfb440009d..07d9c8745b 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -124,7 +124,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TMatmulAccOp, pto::TMatmulBiasOp, pto::TMatmulMxOp, pto::TMatmulMxAccOp, pto::TMatmulMxBiasOp, pto::TGemvOp, pto::TGemvAccOp, pto::TGemvBiasOp, pto::TGemvMxOp, - pto::TGemvMxAccOp, pto::TGemvMxBiasOp>(op)) { + pto::TGemvMxAccOp, pto::TGemvMxBiasOp, + pto::TPrefetchOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/tprefetch_tile_native.pto b/test/lit/pto/tprefetch_tile_native.pto new file mode 100644 index 0000000000..e630565932 --- /dev/null +++ b/test/lit/pto/tprefetch_tile_native.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-level=level3 --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 %s | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tprefetch_arg(%src: memref<16x16xf16, #pto.address_space>, %dst: !pto.tile_buf) { + pto.tprefetch ins(%src : memref<16x16xf16, #pto.address_space>) + outs(%dst : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tprefetch_arg +// NATIVE-SAME: %{{.*}}: memref<16x16xf16, #pto.address_space> +// NATIVE-SAME: %{{.*}}: !pto.tile_buf Date: Mon, 20 Jul 2026 10:19:22 +0800 Subject: [PATCH 76/91] Keep pipe queue ops tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 3 ++- test/lit/pto/pipe_tile_native.pto | 35 ++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/pipe_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 07d9c8745b..65c8918f3f 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -125,7 +125,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TMatmulMxAccOp, pto::TMatmulMxBiasOp, pto::TGemvOp, pto::TGemvAccOp, pto::TGemvBiasOp, pto::TGemvMxOp, pto::TGemvMxAccOp, pto::TGemvMxBiasOp, - pto::TPrefetchOp>(op)) { + pto::TPrefetchOp, pto::TPushOp, pto::TPopOp, + pto::TFreeOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/pipe_tile_native.pto b/test/lit/pto/pipe_tile_native.pto new file mode 100644 index 0000000000..91fb05b3e0 --- /dev/null +++ b/test/lit/pto/pipe_tile_native.pto @@ -0,0 +1,35 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE + +module { + func.func private @push_arg(%gm: memref<256xf32, #pto.address_space>, %consumer: i32, %tile: !pto.tile_buf) attributes {pto.kernel_kind = #pto.kernel_kind} { + %pipe = pto.initialize_l2g2l_pipe {dir_mask = 1, slot_size = 1024, slot_num = 8, local_slot_num = 8, flag_base = 0}(%gm : memref<256xf32, #pto.address_space>, %consumer : i32) -> !pto.pipe + pto.tpush(%tile, %pipe : !pto.tile_buf, !pto.pipe) {split = 0} + return + } + + func.func private @pop_free_arg(%gm: memref<256xf32, #pto.address_space>, %consumer: i32, %tile: !pto.tile_buf) attributes {pto.kernel_kind = #pto.kernel_kind} { + %pipe = pto.initialize_l2g2l_pipe {dir_mask = 1, slot_size = 1024, slot_num = 8, local_slot_num = 8, flag_base = 0}(%gm : memref<256xf32, #pto.address_space>, %consumer : i32) -> !pto.pipe + pto.tpop(%tile, %pipe : !pto.tile_buf, !pto.pipe) {split = 1} + pto.tfree(%tile, %pipe : !pto.tile_buf, !pto.pipe) {split = 1} + return + } +} + +// NATIVE-LABEL: @push_arg +// NATIVE-SAME: !pto.tile_buf Date: Mon, 20 Jul 2026 10:21:26 +0800 Subject: [PATCH 77/91] Keep TSync tile-native --- lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- test/lit/pto/tsync_tile_native.pto | 25 +++++++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) create mode 100644 test/lit/pto/tsync_tile_native.pto diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 65c8918f3f..c2cafb11e8 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -126,7 +126,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TGemvAccOp, pto::TGemvBiasOp, pto::TGemvMxOp, pto::TGemvMxAccOp, pto::TGemvMxBiasOp, pto::TPrefetchOp, pto::TPushOp, pto::TPopOp, - pto::TFreeOp>(op)) { + pto::TFreeOp, pto::TSyncOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/tsync_tile_native.pto b/test/lit/pto/tsync_tile_native.pto new file mode 100644 index 0000000000..1a8404a830 --- /dev/null +++ b/test/lit/pto/tsync_tile_native.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-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a5 %s | FileCheck %s --check-prefix=EMITC + +module { + func.func private @tsync_tile_arg(%event: !pto.tile_buf) { + pto.tsync ins(%event : !pto.tile_buf) + return + } +} + +// NATIVE-LABEL: @tsync_tile_arg +// NATIVE-SAME: !pto.tile_buf Date: Mon, 20 Jul 2026 12:22:53 +0800 Subject: [PATCH 78/91] Support tensor views directly in backends --- include/PTO/IR/PTOOps.td | 6 +- lib/PTO/Transforms/ExpandTileOp.cpp | 50 +++ lib/PTO/Transforms/FoldTileBufIntrinsics.cpp | 229 +++++++++++++- lib/PTO/Transforms/PTOResolveBufferSelect.cpp | 174 ----------- lib/PTO/Transforms/PTOToEmitC.cpp | 285 ++++++++++++++++-- lib/PTO/Transforms/PTOViewToMemref.cpp | 25 +- .../lit/pto/gm_tensor_view_direct_backend.pto | 92 ++++++ .../gm_tensor_view_native_through_sync.pto | 19 +- .../pto/inttoptr_make_tensor_view_invalid.pto | 6 +- .../pto/issue481_addptr_gm_slot_buffer.pto | 8 +- .../issue481_addptr_gm_slot_buffer_gss.pto | 8 +- .../tpush_tpop_globaltensor_frontend_a3.pto | 6 +- .../tpush_tpop_globaltensor_internal_a3.pto | 4 +- ...push_tpop_globaltensor_loop_nosplit_a3.pto | 4 +- ...h_tpop_globaltensor_split_half_slot_a3.pto | 2 +- 15 files changed, 677 insertions(+), 241 deletions(-) create mode 100644 test/lit/pto/gm_tensor_view_direct_backend.pto diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 60758484df..a1f6dd9026 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -223,7 +223,8 @@ def StoreScalarOp : PTO_Op<"store_scalar", [ }]; } -def MakeTensorViewOp : PTO_Op<"make_tensor_view", [AttrSizedOperandSegments]> { +def MakeTensorViewOp : PTO_Op<"make_tensor_view", [AttrSizedOperandSegments, + Pure]> { let summary = "Wrap a pointer as a tensor_view descriptor (no allocation, no copy)."; let arguments = (ins @@ -242,7 +243,8 @@ def MakeTensorViewOp : PTO_Op<"make_tensor_view", [AttrSizedOperandSegments]> { // ============================================================================= // PartitionViewOp // ============================================================================= -def PartitionViewOp : PTO_Op<"partition_view", [AttrSizedOperandSegments]> { +def PartitionViewOp : PTO_Op<"partition_view", [AttrSizedOperandSegments, + Pure]> { let summary = "Partition a tensor view into a smaller logical view (logical slicing)."; let description = [{ Captures a specific calculation region from a large view. diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 3c5b83c9d0..e72fe04607 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -341,6 +341,11 @@ static std::optional resolveViewLayout(Value value) { def = value.getDefiningOp(); continue; } + if (auto partition = dyn_cast(def)) { + value = partition.getSource(); + def = value.getDefiningOp(); + continue; + } break; } return std::nullopt; @@ -580,6 +585,35 @@ static void populateViewShapeAndStrides(Value value, if (!value) return; + if (auto partition = value.getDefiningOp()) { + populateViewShapeAndStrides(partition.getSource(), shape, strides); + SmallVector partitionShape; + partitionShape.reserve(partition.getSizes().size()); + for (Value size : partition.getSizes()) { + int64_t staticSize = ShapedType::kDynamic; + (void)getStaticIntFromValue(size, staticSize); + partitionShape.push_back(staticSize); + } + shape = std::move(partitionShape); + return; + } + + if (auto makeView = value.getDefiningOp()) { + shape.clear(); + strides.clear(); + for (Value dim : makeView.getShape()) { + int64_t staticDim = ShapedType::kDynamic; + (void)getStaticIntFromValue(dim, staticDim); + shape.push_back(staticDim); + } + for (Value stride : makeView.getStrides()) { + int64_t staticStride = ShapedType::kDynamic; + (void)getStaticIntFromValue(stride, staticStride); + strides.push_back(staticStride); + } + return; + } + if (auto subview = value.getDefiningOp()) { populateViewShapeAndStrides(subview.getSource(), shape, strides); SmallVector subviewShape; @@ -670,6 +704,22 @@ static std::optional buildOperandTypeInfo(Value value) { return info; } + if (auto viewTy = dyn_cast(ty)) { + OperandTypeInfo info; + info.kind = OperandKind::View; + info.dtype = getDtypeString(viewTy.getElementType()); + if (info.dtype.empty()) + return std::nullopt; + info.viewMemorySpace = "gm"; + info.viewLayout = resolveViewLayout(value); + populateViewShapeAndStrides(value, info.viewShape, info.viewStrides); + if (info.viewShape.empty()) + info.viewShape.assign(viewTy.getShape().begin(), viewTy.getShape().end()); + if (info.viewStrides.empty()) + info.viewStrides.assign(viewTy.getRank(), ShapedType::kDynamic); + return info; + } + // Auxiliary vector operand — from builtin VectorType (e.g. vector<4xi16>). if (auto vecTy = dyn_cast(ty)) { OperandTypeInfo info; diff --git a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp index b6064832be..af5e49d85b 100644 --- a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp +++ b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp @@ -40,6 +40,7 @@ #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Pass/Pass.h" @@ -432,6 +433,187 @@ static Value computeLinearOffset(OpBuilder &builder, Location loc, return rcPart ? rcPart : svPart; } +static Value unwrapPTOViewBridge(Value value) { + while (auto cast = value.getDefiningOp()) { + if (cast.getNumOperands() != 1 || cast.getNumResults() != 1) + break; + value = cast.getOperand(0); + } + return value; +} + +enum class PTOViewProjectionKind { Dim, Stride, Address }; + +static Value projectSCFIfViewResult(Value view, PTOViewProjectionKind kind, + int64_t dim, pto::PtrType resultPtrType, + OpBuilder &builder, Operation *user); + +static Value resolvePTOViewDim(Value view, int64_t dim, OpBuilder &builder, + Operation *user) { + view = unwrapPTOViewBridge(view); + if (Value projected = projectSCFIfViewResult( + view, PTOViewProjectionKind::Dim, dim, {}, builder, user)) + return projected; + if (auto partition = view.getDefiningOp()) { + if (dim < 0 || dim >= static_cast(partition.getSizes().size())) + return {}; + return partition.getSizes()[dim]; + } + if (auto makeView = view.getDefiningOp()) { + if (dim < 0 || dim >= static_cast(makeView.getShape().size())) + return {}; + return makeView.getShape()[dim]; + } + (void)builder; + (void)user; + return {}; +} + +static Value resolvePTOViewStride(Value view, int64_t dim, OpBuilder &builder, + Operation *user) { + view = unwrapPTOViewBridge(view); + if (Value projected = projectSCFIfViewResult( + view, PTOViewProjectionKind::Stride, dim, {}, builder, user)) + return projected; + if (auto partition = view.getDefiningOp()) + return resolvePTOViewStride(partition.getSource(), dim, builder, user); + if (auto makeView = view.getDefiningOp()) { + if (dim < 0 || dim >= static_cast(makeView.getStrides().size())) + return {}; + return makeView.getStrides()[dim]; + } + return {}; +} + +static Value resolvePTOViewAddress(Value view, pto::PtrType resultType, + OpBuilder &builder, Operation *user) { + view = unwrapPTOViewBridge(view); + if (Value projected = projectSCFIfViewResult( + view, PTOViewProjectionKind::Address, 0, resultType, builder, user)) + return projected; + if (auto makeView = view.getDefiningOp()) { + Value ptr = makeView.getPtr(); + if (ptr.getType() == resultType) + return ptr; + return builder.create(user->getLoc(), resultType, ptr); + } + auto partition = view.getDefiningOp(); + if (!partition) + return {}; + + Value linearOffset; + for (unsigned index = 0; index < partition.getOffsets().size(); ++index) { + Value stride = resolvePTOViewStride(partition.getSource(), index, builder, + user); + if (!stride) + return {}; + Value offset = partition.getOffsets()[index]; + Value term = builder.create(user->getLoc(), offset, stride); + linearOffset = linearOffset + ? builder.create(user->getLoc(), + linearOffset, term) + : term; + } + Value base = + resolvePTOViewAddress(partition.getSource(), resultType, builder, user); + if (!base) + return {}; + if (!linearOffset) + return base; + return builder.create(user->getLoc(), resultType, base, + linearOffset); +} + +static void cloneBlockWithoutTerminator(Block *source, Block *target, + IRMapping &mapping, + OpBuilder &builder) { + builder.setInsertionPointToStart(target); + for (Operation &op : source->without_terminator()) + builder.clone(op, mapping); +} + +static Value projectSCFIfViewResult(Value view, PTOViewProjectionKind kind, + int64_t dim, pto::PtrType resultPtrType, + OpBuilder &builder, Operation *user) { + auto result = dyn_cast(view); + auto ifOp = result ? dyn_cast(result.getOwner()) : scf::IfOp(); + if (!ifOp || ifOp.getNumRegions() != 2) + return {}; + if (kind == PTOViewProjectionKind::Address && !resultPtrType) + return {}; + auto thenYield = dyn_cast(ifOp.thenBlock()->getTerminator()); + auto elseYield = dyn_cast(ifOp.elseBlock()->getTerminator()); + unsigned resultIndex = result.getResultNumber(); + if (!thenYield || !elseYield || + resultIndex >= thenYield.getNumOperands() || + resultIndex >= elseYield.getNumOperands()) + return {}; + + Type projectionType = kind == PTOViewProjectionKind::Address + ? Type(resultPtrType) + : Type(builder.getIndexType()); + if (!projectionType) + return {}; + SmallVector resultTypes(ifOp.getResultTypes().begin(), + ifOp.getResultTypes().end()); + resultTypes.push_back(projectionType); + + OpBuilder ifBuilder(ifOp); + auto newIf = ifBuilder.create( + ifOp.getLoc(), resultTypes, ifOp.getCondition(), + /*addThenBlock=*/true, /*addElseBlock=*/true); + newIf->setAttrs(ifOp->getAttrs()); + + auto cloneBranch = [&](Block *source, Block *target, scf::YieldOp oldYield, + bool &failed) { + IRMapping mapping; + cloneBlockWithoutTerminator(source, target, mapping, ifBuilder); + SmallVector yields; + yields.reserve(oldYield.getNumOperands() + 1); + for (Value operand : oldYield.getOperands()) + yields.push_back(mapping.lookupOrDefault(operand)); + Value yieldedView = yields[resultIndex]; + ifBuilder.setInsertionPointToEnd(target); + Value projection; + switch (kind) { + case PTOViewProjectionKind::Dim: + projection = + resolvePTOViewDim(yieldedView, dim, ifBuilder, user); + break; + case PTOViewProjectionKind::Stride: + projection = + resolvePTOViewStride(yieldedView, dim, ifBuilder, user); + break; + case PTOViewProjectionKind::Address: + projection = resolvePTOViewAddress(yieldedView, resultPtrType, ifBuilder, + user); + break; + } + if (!projection) { + failed = true; + return; + } + yields.push_back(projection); + ifBuilder.create(oldYield.getLoc(), yields); + }; + + bool failed = false; + cloneBranch(ifOp.thenBlock(), newIf.thenBlock(), thenYield, failed); + cloneBranch(ifOp.elseBlock(), newIf.elseBlock(), elseYield, failed); + if (failed) { + newIf.erase(); + return {}; + } + + for (auto [oldResult, newResult] : + llvm::zip_equal(ifOp.getResults(), newIf.getResults().take_front( + ifOp.getNumResults()))) + oldResult.replaceAllUsesWith(newResult); + Value projection = newIf.getResult(newIf.getNumResults() - 1); + ifOp.erase(); + return projection; +} + struct FoldTileBufIntrinsicsPass : public pto::impl::FoldTileBufIntrinsicsBase { using FoldTileBufIntrinsicsBase::FoldTileBufIntrinsicsBase; @@ -686,10 +868,6 @@ struct FoldTileBufIntrinsicsPass } for (auto dimOp : tvDimOps) { - auto chain = traceViewChain(dimOp.getTensorView(), dimOp); - if (!chain) - return signalPassFailure(); - int64_t dimIdx = 0; if (!getConstIndexValue(dimOp.getDimIndex(), dimIdx)) { dimOp.emitError( @@ -698,6 +876,18 @@ struct FoldTileBufIntrinsicsPass return signalPassFailure(); } + builder.setInsertionPoint(dimOp); + if (Value direct = resolvePTOViewDim(dimOp.getTensorView(), dimIdx, + builder, dimOp)) { + dimOp.getResult().replaceAllUsesWith(direct); + dimOp.erase(); + continue; + } + + auto chain = traceViewChain(dimOp.getTensorView(), dimOp); + if (!chain) + return signalPassFailure(); + auto svTy = cast(chain->subview.getType()); if (dimIdx < 0 || dimIdx >= svTy.getRank()) { dimOp.emitError( @@ -722,10 +912,6 @@ struct FoldTileBufIntrinsicsPass } for (auto strideOp : tvStrideOps) { - auto chain = traceViewChain(strideOp.getTensorView(), strideOp); - if (!chain) - return signalPassFailure(); - int64_t dimIdx = 0; if (!getConstIndexValue(strideOp.getDimIndex(), dimIdx)) { strideOp.emitError( @@ -734,6 +920,18 @@ struct FoldTileBufIntrinsicsPass return signalPassFailure(); } + builder.setInsertionPoint(strideOp); + if (Value direct = resolvePTOViewStride( + strideOp.getTensorView(), dimIdx, builder, strideOp)) { + strideOp.getResult().replaceAllUsesWith(direct); + strideOp.erase(); + continue; + } + + auto chain = traceViewChain(strideOp.getTensorView(), strideOp); + if (!chain) + return signalPassFailure(); + auto svTy = cast(chain->subview.getType()); if (dimIdx < 0 || dimIdx >= svTy.getRank()) { strideOp.emitError( @@ -755,13 +953,22 @@ struct FoldTileBufIntrinsicsPass if (shouldFoldAddrFamily(*mode)) { for (auto addrOp : tvAddrOps) { + builder.setInsertionPoint(addrOp); + + auto resultPtrType = dyn_cast(addrOp.getDst().getType()); + if (resultPtrType) { + if (Value direct = resolvePTOViewAddress( + addrOp.getSrc(), resultPtrType, builder, addrOp)) { + addrOp.getDst().replaceAllUsesWith(direct); + addrOp.erase(); + continue; + } + } + auto chain = traceViewChain(addrOp.getSrc(), addrOp); if (!chain) return signalPassFailure(); - builder.setInsertionPoint(addrOp); - - auto resultPtrType = dyn_cast(addrOp.getDst().getType()); if (!resultPtrType) { if (auto resultMemrefType = dyn_cast(addrOp.getDst().getType())) { diff --git a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp index e6ef8ae590..d70e2d3e6a 100644 --- a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp +++ b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp @@ -34,11 +34,7 @@ #include "PTO/Transforms/Passes.h" #include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/IRMapping.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Pass/Pass.h" @@ -200,171 +196,6 @@ static LogicalResult resolveTileNativeSubviews(ModuleOp module, return success(); } -static LogicalResult reconcileSCFViewResultTypes(ModuleOp module) { - SmallVector ifOps; - module.walk([&](scf::IfOp op) { ifOps.push_back(op); }); - for (scf::IfOp op : llvm::reverse(ifOps)) { - if (op.getNumResults() == 0) - continue; - auto thenYield = dyn_cast(op.thenBlock()->getTerminator()); - auto elseYield = dyn_cast(op.elseBlock()->getTerminator()); - if (!thenYield || !elseYield || - thenYield.getNumOperands() != op.getNumResults() || - elseYield.getNumOperands() != op.getNumResults()) - return op.emitError("cannot reconcile scf.if view result types"); - for (unsigned i = 0; i < op.getNumResults(); ++i) { - Type thenType = thenYield.getOperand(i).getType(); - if (thenType != elseYield.getOperand(i).getType()) - return op.emitError("scf.if branches yield different view types"); - op.getResult(i).setType(thenType); - } - } - return success(); -} - -static LogicalResult resolveGMTensorViews(ModuleOp module, MLIRContext *ctx) { - SmallVector makeViews; - module.walk([&](pto::MakeTensorViewOp op) { makeViews.push_back(op); }); - DenseMap loweredViews; - for (pto::MakeTensorViewOp op : makeViews) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Value base = op.getPtr(); - OpFoldResult offset = rewriter.getIndexAttr(0); - Value totalOffset; - while (auto add = base.getDefiningOp()) { - Value term = add.getOffset(); - totalOffset = - totalOffset - ? rewriter.create(op.getLoc(), totalOffset, term) - : term; - base = add.getPtr(); - } - if (totalOffset) - offset = totalOffset; - - auto baseType = dyn_cast(base.getType()); - if (!baseType) - continue; - int64_t rank = static_cast(op.getShape().size()); - SmallVector dynamicShape(rank, ShapedType::kDynamic); - SmallVector dynamicStrides(rank, ShapedType::kDynamic); - auto layout = - StridedLayoutAttr::get(ctx, ShapedType::kDynamic, dynamicStrides); - auto resultType = MemRefType::get(dynamicShape, baseType.getElementType(), - layout, baseType.getMemorySpace()); - SmallVector sizes(op.getShape().begin(), op.getShape().end()); - SmallVector strides(op.getStrides().begin(), - op.getStrides().end()); - auto view = rewriter.create( - op.getLoc(), resultType, base, offset, sizes, strides); - if (totalOffset) - view->setAttr("pto.addptr_trace", rewriter.getUnitAttr()); - if (auto layoutAttr = op.getLayoutAttr()) - view->setAttr("layout", layoutAttr); - loweredViews[op.getResult()] = view.getResult(); - } - - SmallVector dimOps; - module.walk([&](pto::GetTensorViewDimOp op) { dimOps.push_back(op); }); - for (pto::GetTensorViewDimOp op : dimOps) { - Value view = op.getTensorView(); - if (Value lowered = loweredViews.lookup(view)) - view = lowered; - if (!isa(view.getType())) - continue; - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - rewriter.replaceOpWithNewOp(op, view, op.getDimIndex()); - } - - SmallVector partitions; - module.walk([&](pto::PartitionViewOp op) { partitions.push_back(op); }); - for (pto::PartitionViewOp op : partitions) { - Value source = op.getSource(); - if (Value lowered = loweredViews.lookup(source)) - source = lowered; - auto sourceType = dyn_cast(source.getType()); - if (!sourceType) - continue; - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - SmallVector staticSizes; - SmallVector sizes; - for (Value size : op.getSizes()) { - IntegerAttr attr; - if (matchPattern(size, m_Constant(&attr))) { - int64_t value = attr.getValue().getSExtValue(); - staticSizes.push_back(value); - sizes.push_back(rewriter.getIndexAttr(value)); - } else { - staticSizes.push_back(ShapedType::kDynamic); - sizes.push_back(size); - } - } - SmallVector offsets(op.getOffsets().begin(), - op.getOffsets().end()); - SmallVector strides(sourceType.getRank(), - rewriter.getIndexAttr(1)); - SmallVector dynamicStrides(sourceType.getRank(), - ShapedType::kDynamic); - auto layout = - StridedLayoutAttr::get(ctx, ShapedType::kDynamic, dynamicStrides); - auto resultType = MemRefType::get(staticSizes, sourceType.getElementType(), - layout, sourceType.getMemorySpace()); - auto subview = rewriter.create( - op.getLoc(), resultType, source, offsets, sizes, strides); - if (Operation *def = source.getDefiningOp()) - if (auto attr = def->getAttrOfType("layout")) - subview->setAttr("layout", attr); - rewriter.replaceOp(op, subview.getResult()); - } - - SmallVector strideOps; - module.walk([&](pto::GetTensorViewStrideOp op) { strideOps.push_back(op); }); - for (pto::GetTensorViewStrideOp op : strideOps) { - Value view = op.getTensorView(); - if (Value lowered = loweredViews.lookup(view)) - view = lowered; - auto type = dyn_cast(view.getType()); - if (!type) - continue; - IntegerAttr dimAttr; - if (!matchPattern(op.getDimIndex(), m_Constant(&dimAttr))) - return op.emitError( - "get_tensor_view_stride expects a constant dim index"); - int64_t dim = dimAttr.getValue().getSExtValue(); - if (dim < 0 || dim >= type.getRank()) - return op.emitError("get_tensor_view_stride dim index is out of bounds"); - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - SmallVector staticStrides; - int64_t staticOffset = ShapedType::kDynamic; - if (succeeded(pto::getPTOMemRefStridesAndOffset(type, staticStrides, - staticOffset)) && - staticStrides[dim] != ShapedType::kDynamic) { - rewriter.replaceOpWithNewOp(op, - staticStrides[dim]); - continue; - } - auto metadata = - rewriter.create(op.getLoc(), view); - rewriter.replaceOp(op, metadata.getStrides()[dim]); - } - - // Replace roots only after all strongly typed PTO view users have been - // removed. This avoids constructing an intermediate partition_view whose - // source has already changed from tensor_view to memref. - for (pto::MakeTensorViewOp op : makeViews) { - Value lowered = loweredViews.lookup(op.getResult()); - if (!lowered) - continue; - IRRewriter rewriter(ctx); - rewriter.replaceOp(op, lowered); - } - return reconcileSCFViewResultTypes(module); -} - static FailureOr getStaticSlotBytes(pto::TileBufType slotType) { uint64_t elemBytes = pto::getPTOStorageElemByteSize(slotType.getElementType()); if (elemBytes == 0) @@ -531,11 +362,6 @@ struct PTOResolveBufferSelectPass signalPassFailure(); return; } - if (failed(resolveGMTensorViews(mod, ctx))) { - signalPassFailure(); - return; - } - SmallVector markers; mod.walk([&](pto::SlotMarkerOp op) { markers.push_back(op); }); if (markers.empty()) diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 5f369b9054..a8fb32a927 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -79,14 +79,10 @@ static void buildGlobalTensorShapeAndStride(ArrayRef shape, SmallVectorImpl &stride5D); static std::string joinIntTemplateParams(ArrayRef values); static SmallVector buildRowMajorStrides(ArrayRef shape); -static std::string getGlobalTensorTypeStringFromShape(Type elemTy, - ArrayRef shape, - StringRef layoutEnum = - "pto::Layout::ND"); static std::string getGlobalTensorTypeStringFromShapeAndStrides( Type elemTy, ArrayRef shape, ArrayRef strides, StringRef layoutEnum = "pto::Layout::ND"); -static emitc::OpaqueType getGlobalTensorOpaqueTypeFromShape( +static emitc::OpaqueType getRuntimeGlobalTensorOpaqueType( MLIRContext *ctx, Type elemTy, ArrayRef shape, StringRef layoutEnum = "pto::Layout::ND"); @@ -900,13 +896,13 @@ class PTOToEmitCTypeConverter : public TypeConverter { }); addConversion([Ctx](pto::TensorViewType type) -> Type { - return getGlobalTensorOpaqueTypeFromShape( - Ctx, type.getElementType(), type.getShape()); + return getRuntimeGlobalTensorOpaqueType(Ctx, type.getElementType(), + type.getShape()); }); addConversion([Ctx](pto::PartitionTensorViewType type) -> Type { - return getGlobalTensorOpaqueTypeFromShape( - Ctx, type.getElementType(), type.getShape()); + return getRuntimeGlobalTensorOpaqueType(Ctx, type.getElementType(), + type.getShape()); }); addConversion([Ctx](pto::TileBufType type) -> std::optional { @@ -4221,14 +4217,6 @@ static SmallVector buildRowMajorStrides(ArrayRef shape) { return strides; } -static std::string getGlobalTensorTypeStringFromShape(Type elemTy, - ArrayRef shape, - StringRef layoutEnum) { - SmallVector strides = buildRowMajorStrides(shape); - return getGlobalTensorTypeStringFromShapeAndStrides(elemTy, shape, strides, - layoutEnum); -} - static std::string getGlobalTensorTypeStringFromShapeAndStrides( Type elemTy, ArrayRef shape, ArrayRef strides, StringRef layoutEnum) { @@ -4244,11 +4232,23 @@ static std::string getGlobalTensorTypeStringFromShapeAndStrides( strideType + ", " + layoutEnum.str() + ">"; } -static emitc::OpaqueType getGlobalTensorOpaqueTypeFromShape( +static emitc::OpaqueType getRuntimeGlobalTensorOpaqueType( MLIRContext *ctx, Type elemTy, ArrayRef shape, StringRef layoutEnum) { + SmallVector shape5D(5, 1); + SmallVector stride5D(5, -1); + int64_t shift = 5 - static_cast(shape.size()); + for (auto [index, dim] : llvm::enumerate(shape)) + shape5D[shift + static_cast(index)] = + ShapedType::isDynamic(dim) ? -1 : dim; + + std::string elemTypeStr = getElemTypeStringForGT(elemTy); + std::string shapeType = "pto::Shape<" + joinIntTemplateParams(shape5D) + ">"; + std::string strideType = + "pto::Stride<" + joinIntTemplateParams(stride5D) + ">"; return emitc::OpaqueType::get( - ctx, getGlobalTensorTypeStringFromShape(elemTy, shape, layoutEnum)); + ctx, "GlobalTensor<" + elemTypeStr + ", " + shapeType + ", " + + strideType + ", " + layoutEnum.str() + ">"); } static std::string inferFallbackGlobalTensorLayout(ArrayRef shape5D, @@ -4742,6 +4742,21 @@ struct CastPtrConversion : public OpConversionPattern { } }; +struct PTOAddPtrToEmitC : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(pto::AddPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type resultType = getTypeConverter()->convertType(op.getResult().getType()); + if (!resultType) + return rewriter.notifyMatchFailure(op, "failed to convert pointer type"); + Value ptr = peelUnrealized(adaptor.getPtr()); + Value offset = peelUnrealized(adaptor.getOffset()); + rewriter.replaceOpWithNewOp(op, resultType, ptr, offset); + return success(); + } +}; + struct PointerCastConversion : public OpConversionPattern { static bool getIndexConst(Value v, int64_t &out) { if (auto cst = v.getDefiningOp()) { @@ -8138,6 +8153,140 @@ static bool parseIntegerTemplateList(StringRef token, StringRef marker, return true; } +static Value castViewIndexToEmitC(ConversionPatternRewriter &rewriter, + Location loc, Value value) { + Type indexTy = emitc::OpaqueType::get(rewriter.getContext(), "int64_t"); + value = peelUnrealized(value); + if (value.getType() == indexTy) + return value; + return rewriter.create(loc, indexTy, value).getResult(); +} + +static Value makeViewIndexConstant(ConversionPatternRewriter &rewriter, + Location loc, int64_t value) { + return makeEmitCIntConstant( + rewriter, loc, emitc::OpaqueType::get(rewriter.getContext(), "int64_t"), + value); +} + +static Value getRuntimeGlobalTensorMetadata( + ConversionPatternRewriter &rewriter, Location loc, Value tensor, + Value logicalDim, int64_t rank, StringRef helper) { + Value dim = castViewIndexToEmitC(rewriter, loc, logicalDim); + int64_t shift = 5 - rank; + if (shift != 0) { + dim = rewriter + .create(loc, dim.getType(), dim, + makeViewIndexConstant(rewriter, loc, shift)) + .getResult(); + } + return rewriter + .create(loc, dim.getType(), helper, ArrayAttr{}, + ArrayAttr{}, ValueRange{tensor, dim}) + .getResult(0); +} + +static FailureOr buildRuntimeGlobalTensor( + ConversionPatternRewriter &rewriter, Location loc, Value ptr, Type elemTy, + ArrayRef staticShape, ValueRange runtimeShape, + ValueRange runtimeStrides, StringRef layoutEnum = "pto::Layout::ND") { + if (staticShape.size() > 5 || runtimeShape.size() != staticShape.size() || + runtimeStrides.size() != staticShape.size()) + return failure(); + + SmallVector shape5D(5, 1); + SmallVector stride5D(5, -1); + int64_t shift = 5 - static_cast(staticShape.size()); + for (auto [index, dim] : llvm::enumerate(staticShape)) + shape5D[shift + static_cast(index)] = + ShapedType::isDynamic(dim) ? -1 : dim; + + std::string shapeType = "pto::Shape<" + joinIntTemplateParams(shape5D) + ">"; + std::string strideType = + "pto::Stride<" + joinIntTemplateParams(stride5D) + ">"; + SmallVector shapeValues; + SmallVector strideValues; + for (int64_t dim = 0; dim < shift; ++dim) { + shapeValues.push_back(makeViewIndexConstant(rewriter, loc, 1)); + strideValues.push_back(makeViewIndexConstant(rewriter, loc, 1)); + } + for (Value value : runtimeShape) + shapeValues.push_back(castViewIndexToEmitC(rewriter, loc, value)); + for (Value value : runtimeStrides) + strideValues.push_back(castViewIndexToEmitC(rewriter, loc, value)); + + Value shape = rewriter + .create( + loc, emitc::OpaqueType::get(rewriter.getContext(), + shapeType), + shapeType, ArrayAttr{}, ArrayAttr{}, shapeValues) + .getResult(0); + Value stride = rewriter + .create( + loc, emitc::OpaqueType::get(rewriter.getContext(), + strideType), + strideType, ArrayAttr{}, ArrayAttr{}, strideValues) + .getResult(0); + auto resultType = getRuntimeGlobalTensorOpaqueType( + rewriter.getContext(), elemTy, staticShape, layoutEnum); + return rewriter + .create(loc, resultType, resultType.getValue(), + ArrayAttr{}, ArrayAttr{}, + ValueRange{ptr, shape, stride}) + .getResult(0); +} + +struct PTOMakeTensorViewToEmitC + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(mlir::pto::MakeTensorViewOp op, + OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto resultType = dyn_cast(op.getResult().getType()); + if (!resultType) + return rewriter.notifyMatchFailure(op, "expected tensor_view result"); + auto layout = op.getLayoutAttr() + ? layoutToEmitCString(op.getLayoutAttr().getLayout()) + : "pto::Layout::ND"; + auto result = buildRuntimeGlobalTensor( + rewriter, op.getLoc(), peelUnrealized(adaptor.getPtr()), + resultType.getElementType(), resultType.getShape(), adaptor.getShape(), + adaptor.getStrides(), layout); + if (failed(result)) + return rewriter.notifyMatchFailure( + op, "failed to build runtime GlobalTensor descriptor"); + rewriter.replaceOp(op, *result); + return success(); + } +}; + +template +struct PTOGetTensorViewMetadataToEmitC : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite( + OpTy op, typename OpTy::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type sourceType = op.getTensorView().getType(); + int64_t rank = 0; + if (auto type = dyn_cast(sourceType)) + rank = type.getRank(); + else if (auto type = dyn_cast(sourceType)) + rank = type.getRank(); + else + return rewriter.notifyMatchFailure(op, "expected PTO tensor view"); + + Value result = getRuntimeGlobalTensorMetadata( + rewriter, op.getLoc(), peelUnrealized(adaptor.getTensorView()), + adaptor.getDimIndex(), rank, + IsStride ? "PTOAS__GLOBAL_TENSOR_GET_STRIDE" + : "PTOAS__GLOBAL_TENSOR_GET_SHAPE"); + rewriter.replaceOp(op, result); + return success(); + } +}; + static LogicalResult getStaticTensorViewStrides( Value source, Value convertedSource, pto::TensorViewType sourceType, SmallVectorImpl &strides) { @@ -8175,6 +8324,73 @@ static LogicalResult getStaticTensorViewStrides( struct PTOPartitionViewToEmitC : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(mlir::pto::PartitionViewOp op, + OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto sourceType = dyn_cast(op.getSource().getType()); + auto resultType = + dyn_cast(op.getResult().getType()); + if (!sourceType || !resultType) + return rewriter.notifyMatchFailure(op, "expected PTO tensor view types"); + + Value source = peelUnrealized(adaptor.getSource()); + SmallVector sourceStrides; + sourceStrides.reserve(sourceType.getRank()); + for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) { + Value logicalDim = makeViewIndexConstant(rewriter, op.getLoc(), dim); + sourceStrides.push_back(getRuntimeGlobalTensorMetadata( + rewriter, op.getLoc(), source, logicalDim, sourceType.getRank(), + "PTOAS__GLOBAL_TENSOR_GET_STRIDE")); + } + + Value linearOffset = makeViewIndexConstant(rewriter, op.getLoc(), 0); + for (auto [offset, stride] : + llvm::zip(adaptor.getOffsets(), sourceStrides)) { + Value term = rewriter + .create( + op.getLoc(), linearOffset.getType(), + castViewIndexToEmitC(rewriter, op.getLoc(), offset), + stride) + .getResult(); + linearOffset = rewriter + .create(op.getLoc(), + linearOffset.getType(), + linearOffset, term) + .getResult(); + } + + std::string elemTypeStr = + getElemTypeStringForGT(sourceType.getElementType()); + auto ptrType = emitc::PointerType::get(emitc::OpaqueType::get( + rewriter.getContext(), "__gm__ " + elemTypeStr)); + Value data = rewriter + .create( + op.getLoc(), ptrType, "PTOAS__GLOBAL_TENSOR_DATA", + ArrayAttr{}, ArrayAttr{}, ValueRange{source}) + .getResult(0); + Value ptr = rewriter + .create(op.getLoc(), ptrType, data, + linearOffset) + .getResult(); + + auto layout = resolveLayoutForGlobalTensor(op.getOperation(), op.getSource()); + std::string layoutString = + layout ? layoutToEmitCString(*layout) : "pto::Layout::ND"; + auto result = buildRuntimeGlobalTensor( + rewriter, op.getLoc(), ptr, resultType.getElementType(), + resultType.getShape(), adaptor.getSizes(), sourceStrides, layoutString); + if (failed(result)) + return rewriter.notifyMatchFailure( + op, "failed to build partition GlobalTensor descriptor"); + rewriter.replaceOp(op, *result); + return success(); + } +}; + +struct PTOPartitionViewStaticToEmitC + : public OpConversionPattern { using OpConversionPattern< mlir::pto::PartitionViewOp>::OpConversionPattern; @@ -14283,7 +14499,8 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add(typeConverter, ctx); + patterns.add( + typeConverter, ctx); patterns.add(typeConverter, ctx, targetArch); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add(typeConverter, ctx); + patterns.add, + PTOGetTensorViewMetadataToEmitC>(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); @@ -14512,7 +14732,9 @@ struct EmitPTOManualPass needsEventIdArrayHelper = true; if (isa(op)) needsTRandomHelper = true; - if (isa(op)) + if (isa(op)) needsGlobalTensorDataHelper = true; if (isa decltype(tensor.data()) { return tensor.data(); } + +template +static AICORE inline int64_t PTOAS__GLOBAL_TENSOR_GET_SHAPE(Tensor &tensor, + int64_t dim) { + return tensor.GetShape(static_cast(dim)); +} + +template +static AICORE inline int64_t PTOAS__GLOBAL_TENSOR_GET_STRIDE(Tensor &tensor, + int64_t dim) { + return tensor.GetStride(static_cast(dim)); +} )cpp")); } if (needsEventIdArrayHelper) { @@ -14805,6 +15039,13 @@ static AICORE inline void PTOAS__DCCI_SINGLE_CACHE_LINE(Ptr ptr) { return; } + if (Type convertedOutTy = typeConverter.convertType(outTy); + convertedOutTy && convertedOutTy == inTy) { + output.replaceAllUsesWith(input); + castsToErase.push_back(cast); + return; + } + // IndexType is lowered to int64_t for EmitC. SCF structural conversion // can still materialize temporary index<->int64_t bridges; keeping them // as emitc.cast leaves illegal index-typed EmitC IR for LLVM 21's C++ diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index c2cafb11e8..d3a525e86b 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -1319,14 +1319,20 @@ struct PTOViewToMemrefPass // ------------------------------------------------------------------ SmallVector newInputs; for (Type t : fnTy.getInputs()) { - newInputs.push_back(preserveTileABI && isa(t) + newInputs.push_back(preserveTileABI && + isa(t) ? t : convertPTOTypeToMemRef(t)); } SmallVector newResults; for (Type t : fnTy.getResults()) { - newResults.push_back(preserveTileABI && isa(t) + newResults.push_back(preserveTileABI && + isa(t) ? t : convertPTOTypeToMemRef(t)); } @@ -1350,7 +1356,7 @@ struct PTOViewToMemrefPass // ------------------------------------------------------------------ // Stage 0.20: lower pto.inttoptr result types to GM memrefs. // ------------------------------------------------------------------ - if (failed(lowerIntToPtrOps(func, ctx))) { + if (!preserveTileABI && failed(lowerIntToPtrOps(func, ctx))) { signalPassFailure(); return; } @@ -1366,7 +1372,7 @@ struct PTOViewToMemrefPass // ------------------------------------------------------------------ // Stage 0.30: materialize pto.ptrtoint(addptr ...) byte offsets. // ------------------------------------------------------------------ - if (failed(lowerPtrToIntOps(func, ctx))) { + if (!preserveTileABI && failed(lowerPtrToIntOps(func, ctx))) { signalPassFailure(); return; } @@ -1752,7 +1758,16 @@ struct PTOViewToMemrefPass for (auto *op : addPtrs) { if (!op) continue; - op->emitError("addptr must feed make_tensor_view, initialize_l2g2l_pipe(gm_addr) or load/store_scalar for lowering"); + if (llvm::all_of(op->getUsers(), [op](Operation *user) { + if (isa(user)) + return true; + auto init = dyn_cast(user); + return init && init.getGmAddr() == op->getResult(0); + })) + continue; + op->emitError("addptr must feed make_tensor_view, " + "initialize_l2g2l_pipe(gm_addr) or load/store_scalar " + "for lowering"); signalPassFailure(); return; } diff --git a/test/lit/pto/gm_tensor_view_direct_backend.pto b/test/lit/pto/gm_tensor_view_direct_backend.pto new file mode 100644 index 0000000000..d161e392d1 --- /dev/null +++ b/test/lit/pto/gm_tensor_view_direct_backend.pto @@ -0,0 +1,92 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: ptoas --pto-arch=a5 --pto-level=level3 %s | FileCheck %s --check-prefix=EMITC +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s --check-prefix=VPTO + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @gm_tensor_view_direct_backend( + %base: !pto.ptr, %base_offset: index, + %rows: index, %cols: index, %row_stride: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c0_i64 = arith.constant 0 : i64 + %ptr = pto.addptr %base, %base_offset + : !pto.ptr -> !pto.ptr + %view = pto.make_tensor_view %ptr, shape = [%rows, %cols], + strides = [%row_stride, %c1] : !pto.tensor_view + %part = pto.partition_view %view, offsets = [%c0, %c0], + sizes = [%c16, %c16] + : !pto.tensor_view -> !pto.partition_tensor_view<16x16xf32> + %tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + pto.tload ins(%part : !pto.partition_tensor_view<16x16xf32>) + outs(%tile : !pto.tile_buf) + return + } + + func.func @gm_tensor_view_control_flow( + %base: !pto.ptr, %alternate: !pto.ptr, %cond: i1, + %rows: index, %cols: index, %row_stride: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c0_i64 = arith.constant 0 : i64 + %view = scf.if %cond -> (!pto.tensor_view) { + %then_view = pto.make_tensor_view %base, shape = [%rows, %cols], + strides = [%row_stride, %c1] : !pto.tensor_view + scf.yield %then_view : !pto.tensor_view + } else { + %else_view = pto.make_tensor_view %alternate, shape = [%rows, %cols], + strides = [%row_stride, %c1] : !pto.tensor_view + scf.yield %else_view : !pto.tensor_view + } + %dim = pto.get_tensor_view_dim %view, %c0 + : !pto.tensor_view -> index + %stride = pto.get_tensor_view_stride %view, %c1 + : !pto.tensor_view -> index + %part = pto.partition_view %view, offsets = [%c0, %stride], + sizes = [%dim, %c16] + : !pto.tensor_view -> !pto.partition_tensor_view<16x16xf32> + %tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + pto.tload ins(%part : !pto.partition_tensor_view<16x16xf32>) + outs(%tile : !pto.tile_buf) + return + } +} + +// EMITC-LABEL: gm_tensor_view_direct_backend(__gm__ float* +// EMITC: pto::Shape<1, 1, 1, -1, -1> +// EMITC: pto::Stride<-1, -1, -1, -1, -1> +// EMITC: PTOAS__GLOBAL_TENSOR_GET_STRIDE +// EMITC: TLOAD( +// EMITC-NOT: memref +// EMITC-LABEL: gm_tensor_view_control_flow( +// EMITC: if ( +// EMITC: PTOAS__GLOBAL_TENSOR_GET_SHAPE +// EMITC: PTOAS__GLOBAL_TENSOR_GET_STRIDE +// EMITC: TLOAD( + +// VPTO-LABEL: func.func @gm_tensor_view_direct_backend +// VPTO-NOT: pto.make_tensor_view +// VPTO-NOT: pto.partition_view +// VPTO-NOT: pto.get_tensor_view_dim +// VPTO-NOT: pto.get_tensor_view_stride +// VPTO: pto.copy_gm_to_ubuf +// VPTO-LABEL: func.func @gm_tensor_view_control_flow +// VPTO-NOT: pto.make_tensor_view +// VPTO-NOT: pto.partition_view +// VPTO: pto.copy_gm_to_ubuf diff --git a/test/lit/pto/gm_tensor_view_native_through_sync.pto b/test/lit/pto/gm_tensor_view_native_through_sync.pto index 26223f478d..36193cab32 100644 --- a/test/lit/pto/gm_tensor_view_native_through_sync.pto +++ b/test/lit/pto/gm_tensor_view_native_through_sync.pto @@ -10,7 +10,8 @@ // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-buffer-select %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=RESOLVED module { - func.func @gm_tensor_view_native_through_sync(%ptr: !pto.ptr) { + func.func @gm_tensor_view_native_through_sync(%ptr: !pto.ptr) + -> !pto.partition_tensor_view<32x16xf32> { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -24,7 +25,7 @@ module { %part = pto.partition_view %view, offsets = [%c0, %stride], sizes = [%dim, %c16] : !pto.tensor_view -> !pto.partition_tensor_view<32x16xf32> - return + return %part : !pto.partition_tensor_view<32x16xf32> } } @@ -37,11 +38,9 @@ module { // NATIVE-NOT: memref.subview // RESOLVED: IR Dump After PTOResolveBufferSelect -// RESOLVED: memref.reinterpret_cast -// RESOLVED: memref.dim -// RESOLVED: memref.subview -// RESOLVED: memref.extract_strided_metadata -// RESOLVED-NOT: pto.make_tensor_view -// RESOLVED-NOT: pto.get_tensor_view_dim -// RESOLVED-NOT: pto.get_tensor_view_stride -// RESOLVED-NOT: pto.partition_view +// RESOLVED: pto.make_tensor_view +// RESOLVED: pto.get_tensor_view_dim +// RESOLVED: pto.get_tensor_view_stride +// RESOLVED: pto.partition_view +// RESOLVED-NOT: memref.reinterpret_cast +// RESOLVED-NOT: memref.subview diff --git a/test/lit/pto/inttoptr_make_tensor_view_invalid.pto b/test/lit/pto/inttoptr_make_tensor_view_invalid.pto index 3c9e78c856..a35137f388 100644 --- a/test/lit/pto/inttoptr_make_tensor_view_invalid.pto +++ b/test/lit/pto/inttoptr_make_tensor_view_invalid.pto @@ -1,11 +1,13 @@ // RUN: not ptoas --pto-arch=a3 %s -o /dev/null 2>&1 | FileCheck %s module { - func.func @inttoptr_make_tensor_view_invalid(%addr: i64) attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @inttoptr_make_tensor_view_invalid(%addr: i64) + -> !pto.tensor_view + attributes {pto.kernel_kind = #pto.kernel_kind} { %c1 = arith.constant 1 : index %ptr = pto.inttoptr %addr : i64 -> !pto.ptr %view = pto.make_tensor_view %ptr, shape = [%c1], strides = [%c1] : !pto.tensor_view - return + return %view : !pto.tensor_view } } diff --git a/test/lit/pto/issue481_addptr_gm_slot_buffer.pto b/test/lit/pto/issue481_addptr_gm_slot_buffer.pto index 5bbe4a6281..5c2cc19c4c 100644 --- a/test/lit/pto/issue481_addptr_gm_slot_buffer.pto +++ b/test/lit/pto/issue481_addptr_gm_slot_buffer.pto @@ -36,11 +36,11 @@ module { // CHECK-LABEL: AICORE void cube_kernel(__gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}}, __gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}}) // CHECK: int64_t {{[_A-Za-z][_A-Za-z0-9]*}} = get_block_idx(); -// CHECK: __gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}} = {{[_A-Za-z][_A-Za-z0-9]*}} + {{[_A-Za-z][_A-Za-z0-9]*}}; -// CHECK: auto {{[_A-Za-z][_A-Za-z0-9]*}} = TPipe<0, Direction::DIR_C2V, 1024, 8, 8, false>( +// CHECK-DAG: {{[_A-Za-z][_A-Za-z0-9]*}} + +// CHECK-DAG: auto {{[_A-Za-z][_A-Za-z0-9]*}} = TPipe<0, Direction::DIR_C2V, 1024, 8, 8, false>( // CHECK-LABEL: AICORE void vector_kernel(__gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}}, __gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}}) // CHECK: int64_t {{[_A-Za-z][_A-Za-z0-9]*}} = get_block_idx(); -// CHECK: __gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}} = {{[_A-Za-z][_A-Za-z0-9]*}} + {{[_A-Za-z][_A-Za-z0-9]*}}; -// CHECK: auto {{[_A-Za-z][_A-Za-z0-9]*}} = TPipe<0, Direction::DIR_C2V, 1024, 8, 8, false>( +// CHECK-DAG: {{[_A-Za-z][_A-Za-z0-9]*}} + +// CHECK-DAG: auto {{[_A-Za-z][_A-Za-z0-9]*}} = TPipe<0, Direction::DIR_C2V, 1024, 8, 8, false>( // CHECK-NOT: pto.addptr diff --git a/test/lit/pto/issue481_addptr_gm_slot_buffer_gss.pto b/test/lit/pto/issue481_addptr_gm_slot_buffer_gss.pto index ffd82412ae..000ace1c83 100644 --- a/test/lit/pto/issue481_addptr_gm_slot_buffer_gss.pto +++ b/test/lit/pto/issue481_addptr_gm_slot_buffer_gss.pto @@ -36,11 +36,11 @@ module { // CHECK-LABEL: AICORE void cube_kernel(__gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}}, __gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}}) // CHECK: int64_t {{[_A-Za-z][_A-Za-z0-9]*}} = get_block_idx(); -// CHECK: __gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}} = {{[_A-Za-z][_A-Za-z0-9]*}} + {{[_A-Za-z][_A-Za-z0-9]*}}; -// CHECK: auto {{[_A-Za-z][_A-Za-z0-9]*}} = TPipe<0, Direction::DIR_C2V, 1024, 8, 8, false>( +// CHECK-DAG: {{[_A-Za-z][_A-Za-z0-9]*}} + +// CHECK-DAG: auto {{[_A-Za-z][_A-Za-z0-9]*}} = TPipe<0, Direction::DIR_C2V, 1024, 8, 8, false>( // CHECK-LABEL: AICORE void vector_kernel(__gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}}, __gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}}) // CHECK: int64_t {{[_A-Za-z][_A-Za-z0-9]*}} = get_block_idx(); -// CHECK: __gm__ float* {{[_A-Za-z][_A-Za-z0-9]*}} = {{[_A-Za-z][_A-Za-z0-9]*}} + {{[_A-Za-z][_A-Za-z0-9]*}}; -// CHECK: auto {{[_A-Za-z][_A-Za-z0-9]*}} = TPipe<0, Direction::DIR_C2V, 1024, 8, 8, false>( +// CHECK-DAG: {{[_A-Za-z][_A-Za-z0-9]*}} + +// CHECK-DAG: auto {{[_A-Za-z][_A-Za-z0-9]*}} = TPipe<0, Direction::DIR_C2V, 1024, 8, 8, false>( // CHECK-NOT: pto.addptr diff --git a/test/lit/pto/tpush_tpop_globaltensor_frontend_a3.pto b/test/lit/pto/tpush_tpop_globaltensor_frontend_a3.pto index d8973881e6..4afe53c020 100644 --- a/test/lit/pto/tpush_tpop_globaltensor_frontend_a3.pto +++ b/test/lit/pto/tpush_tpop_globaltensor_frontend_a3.pto @@ -65,14 +65,16 @@ module { // CHECK-LABEL: AICORE void cube_kernel // CHECK-SAME: (__gm__ float* [[CUBE_GM:[_A-Za-z][_A-Za-z0-9]*]], -// CHECK: TPipe<0, Direction::DIR_C2V, 1024, 8, 8, true>([[CUBE_GM]], {{.*}}, {{.*}}); +// CHECK: __gm__ float* [[CUBE_DATA:[_A-Za-z][_A-Za-z0-9]*]] = PTOAS__GLOBAL_TENSOR_DATA( +// CHECK: TPipe<0, Direction::DIR_C2V, 1024, 8, 8, true>([[CUBE_DATA]], {{.*}}, {{.*}}); // CHECK: GlobalTensor, pto::Stride<256, 256, 256, 16, 1>, pto::Layout::ND> {{[_A-Za-z][_A-Za-z0-9]*}}(nullptr); // CHECK: TALLOC, GlobalTensor, pto::Stride<256, 256, 256, 16, 1>, pto::Layout::ND>, TileSplitAxis::TILE_NO_SPLIT>({{.*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); // CHECK: TSTORE // CHECK: TPUSH, GlobalTensor, pto::Stride<256, 256, 256, 16, 1>, pto::Layout::ND>, TileSplitAxis::TILE_NO_SPLIT>({{.*}}, {{[_A-Za-z][_A-Za-z0-9]*}}); // CHECK-LABEL: AICORE void vector_kernel // CHECK-SAME: (__gm__ float* [[VEC_GM:[_A-Za-z][_A-Za-z0-9]*]], -// CHECK: TPipe<0, Direction::DIR_C2V, 1024, 8, 8, true>([[VEC_GM]], {{.*}}, {{.*}}); +// CHECK: __gm__ float* [[VEC_DATA:[_A-Za-z][_A-Za-z0-9]*]] = PTOAS__GLOBAL_TENSOR_DATA( +// CHECK: TPipe<0, Direction::DIR_C2V, 1024, 8, 8, true>([[VEC_DATA]], {{.*}}, {{.*}}); // CHECK: GlobalTensor, pto::Stride<256, 256, 256, 16, 1>, pto::Layout::ND> {{[_A-Za-z][_A-Za-z0-9]*}}(nullptr); // CHECK: TPOP, GlobalTensor, pto::Stride<256, 256, 256, 16, 1>, pto::Layout::ND>, TileSplitAxis::TILE_NO_SPLIT>( // CHECK: TLOAD diff --git a/test/lit/pto/tpush_tpop_globaltensor_internal_a3.pto b/test/lit/pto/tpush_tpop_globaltensor_internal_a3.pto index c56c19ed82..04e0089942 100644 --- a/test/lit/pto/tpush_tpop_globaltensor_internal_a3.pto +++ b/test/lit/pto/tpush_tpop_globaltensor_internal_a3.pto @@ -22,5 +22,5 @@ module { // CHECK-LABEL: AICORE void cube_globaltensor_internal( // CHECK: __gm__ float* [[GM_DATA:[_A-Za-z][_A-Za-z0-9]*]] = PTOAS__GLOBAL_TENSOR_DATA( // CHECK: TPipe<0, Direction::DIR_C2V, 1024, 8, 8, true>([[GM_DATA]], {{.*}}, {{.*}}); -// CHECK: TALLOC, GlobalTensor, pto::Stride<256, 256, 256, 16, 1>, pto::Layout::ND>, TileSplitAxis::TILE_NO_SPLIT>( -// CHECK: TPUSH, GlobalTensor, pto::Stride<256, 256, 256, 16, 1>, pto::Layout::ND>, TileSplitAxis::TILE_NO_SPLIT>( +// CHECK: TALLOC, GlobalTensor, pto::Stride<-1, -1, -1, -1, -1>, pto::Layout::ND>, TileSplitAxis::TILE_NO_SPLIT>( +// CHECK: TPUSH, GlobalTensor, pto::Stride<-1, -1, -1, -1, -1>, pto::Layout::ND>, TileSplitAxis::TILE_NO_SPLIT>( diff --git a/test/lit/pto/tpush_tpop_globaltensor_loop_nosplit_a3.pto b/test/lit/pto/tpush_tpop_globaltensor_loop_nosplit_a3.pto index 20b7f098f2..bccadc35bb 100644 --- a/test/lit/pto/tpush_tpop_globaltensor_loop_nosplit_a3.pto +++ b/test/lit/pto/tpush_tpop_globaltensor_loop_nosplit_a3.pto @@ -111,9 +111,9 @@ module { // CHECK: get_subblockid // CHECK: TPOP, GlobalTensor, pto::Stride<65536, 65536, 16384, 128, 1>, pto::Layout::ND>, TileSplitAxis::TILE_NO_SPLIT> // CHECK: for ( -// CHECK: GlobalTensor, pto::Stride<65536, 65536, 16384, 128, 1>, pto::Layout::ND> +// CHECK: GlobalTensor, pto::Stride<-1, -1, -1, -1, -1>, pto::Layout::ND> // CHECK: TLOAD // CHECK: TADDS -// CHECK: GlobalTensor, pto::Stride<512, 512, 128, 512, 1>, pto::Layout::ND> +// CHECK: GlobalTensor, pto::Stride<-1, -1, -1, -1, -1>, pto::Layout::ND> // CHECK: TSTORE // CHECK: TFREE, GlobalTensor, pto::Stride<65536, 65536, 16384, 128, 1>, pto::Layout::ND>, TileSplitAxis::TILE_NO_SPLIT> diff --git a/test/lit/pto/tpush_tpop_globaltensor_split_half_slot_a3.pto b/test/lit/pto/tpush_tpop_globaltensor_split_half_slot_a3.pto index dd3f30890c..b561807000 100644 --- a/test/lit/pto/tpush_tpop_globaltensor_split_half_slot_a3.pto +++ b/test/lit/pto/tpush_tpop_globaltensor_split_half_slot_a3.pto @@ -74,7 +74,7 @@ module { // CHECK-LABEL: AICORE void vector_kernel // CHECK: TPipe<0, Direction::DIR_C2V, 262144, 8, 8, false> // CHECK: TPOP, GlobalTensor, pto::Stride<65536, 65536, 16384, 128, 1>, pto::Layout::ND>, TileSplitAxis::TILE_UP_DOWN> -// CHECK: GlobalTensor, pto::Stride<65536, 65536, 16384, 128, 1>, pto::Layout::ND> +// CHECK: GlobalTensor, pto::Stride<-1, -1, -1, -1, -1>, pto::Layout::ND> // CHECK: TFREE, GlobalTensor, pto::Stride<65536, 65536, 16384, 128, 1>, pto::Layout::ND>, TileSplitAxis::TILE_UP_DOWN> // CHECK-LABEL: __global__ AICORE void call_both // CHECK: set_ffts_base_addr From 84ab543d3dc86c0a1acfd38c3c87b7ba6eed4281 Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 20 Jul 2026 12:42:27 +0800 Subject: [PATCH 79/91] Keep PTO address ops native through view lowering --- include/PTO/IR/PTOOps.td | 5 +- lib/PTO/Transforms/PTOViewToMemref.cpp | 443 +------------------- lib/PTO/Transforms/VPTOPtrNormalize.cpp | 47 ++- test/lit/pto/castptr_emitc_tile_and_int.pto | 5 + test/lit/pto/ptr_int_cast.pto | 44 +- 5 files changed, 76 insertions(+), 468 deletions(-) diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index a1f6dd9026..1e9fbf2ebb 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -110,9 +110,8 @@ def AddPtrOp : PTO_Op<"addptr", [ def PtrToIntOp : PTO_Op<"ptrtoint", [Pure]> { let summary = "Convert a global pointer to a byte address"; let description = [{ - Converts a global pointer value to an i64 byte address. If the source - pointer is produced by `pto.addptr`, lowering materializes the addptr - offset as an explicit byte offset before converting to EmitC. + Converts a global pointer value to an i64 byte address. Pointer arithmetic + such as `pto.addptr` remains explicit until backend lowering. }]; let arguments = (ins diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index d3a525e86b..490178db78 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -80,6 +80,8 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { auto result = func.walk([&](Operation *op) { if (isa()) { - folded = true; - Value off = ensureIndex(rewriter, loc, add.getOperand(1), add); - totalOffset = - totalOffset ? rewriter.create(loc, totalOffset, off) : off; - base = add.getOperand(0); - } - return folded; -} - [[maybe_unused]] static void dumpPretty(Operation *op, llvm::raw_ostream &os) { OpPrintingFlags flags; flags.useLocalScope(); @@ -865,188 +854,6 @@ static void markForceDynamicValidShape(Operation *op, bool force, func.setFunctionType(FunctionType::get(ctx, newInputs, newResults)); } -static Value castIndexToI64(IRRewriter &rewriter, Location loc, Value value) { - Type i64Ty = rewriter.getI64Type(); - if (value.getType() == i64Ty) - return value; - return rewriter.create(loc, i64Ty, value).getResult(); -} - -static FailureOr -materializePtrToIntAddPtrAddress(IRRewriter &rewriter, Location loc, - mlir::pto::PtrToIntOp anchor, Value source) { - SmallVector addPtrChain; - Value base = source; - while (auto add = base.getDefiningOp()) { - addPtrChain.push_back(add); - base = add.getOperand(0); - } - - if (addPtrChain.empty()) - return failure(); - - auto baseMemTy = dyn_cast(base.getType()); - if (!baseMemTy) { - anchor.emitOpError( - "pto.addptr source base could not be lowered to a GM memref"); - return failure(); - } - - Value byteAddress = rewriter.create( - loc, rewriter.getI64Type(), base); - for (auto add : addPtrChain) { - auto addPtrTy = dyn_cast(add.getResult().getType()); - if (!addPtrTy) { - anchor.emitOpError("requires pto.addptr source to have !pto.ptr result " - "type before byte-address lowering"); - return failure(); - } - - unsigned elemBytes = - mlir::pto::getPTOStorageElemByteSize(addPtrTy.getElementType()); - if (elemBytes == 0) { - anchor.emitOpError("cannot lower pto.addptr source with unknown element " - "byte size to a byte address"); - return failure(); - } - - Value byteOffset = castIndexToI64(rewriter, loc, add.getOffset()); - if (elemBytes != 1) { - Value elemBytesValue = - rewriter.create(loc, elemBytes, 64); - byteOffset = - rewriter.create(loc, byteOffset, elemBytesValue) - .getResult(); - } - byteAddress = - rewriter.create(loc, byteAddress, byteOffset).getResult(); - } - - return byteAddress; -} - -static LogicalResult lowerIntToPtrOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector intToPtrs; - func.walk([&](mlir::pto::IntToPtrOp op) { intToPtrs.push_back(op); }); - - for (auto op : intToPtrs) { - if (!isa(op.getResult().getType())) - continue; - - auto targetTy = - dyn_cast(convertPTOTypeToMemRef(op.getResult().getType())); - if (!targetTy) { - op.emitError("failed to convert inttoptr result to memref type"); - return failure(); - } - - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - auto lowered = - rewriter.create(op.getLoc(), targetTy, - op.getAddr()); - lowered->setAttrs(op->getAttrs()); - rewriter.replaceOp(op, lowered.getResult()); - } - - return success(); -} - -static LogicalResult lowerPtrToIntOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector ptrToInts; - func.walk([&](mlir::pto::PtrToIntOp op) { ptrToInts.push_back(op); }); - - for (auto op : ptrToInts) { - Value source = op.getPtr(); - if (source.getDefiningOp()) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - FailureOr byteAddress = - materializePtrToIntAddPtrAddress(rewriter, op.getLoc(), op, source); - if (failed(byteAddress)) - return failure(); - rewriter.replaceOp(op, *byteAddress); - continue; - } - - if (isa(source.getType())) - continue; - } - - DefaultInlineVector remaining; - func.walk([&](mlir::pto::PtrToIntOp op) { - if (isa(op.getPtr().getType())) - remaining.push_back(op); - }); - for (auto op : remaining) { - op.emitError("ptrtoint source could not be lowered to a GM memref"); - return failure(); - } - - return success(); -} - -[[maybe_unused]] static LogicalResult lowerMakeTensorViewOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector makeViews; - func.walk([&](mlir::pto::MakeTensorViewOp op) { makeViews.push_back(op); }); - - for (auto op : makeViews) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - Value baseBuf = op.getOperand(0); - OpFoldResult off0 = rewriter.getIndexAttr(0); - bool foldedAddPtr = false; - { - Value cur = baseBuf; - Value totalOffset; - while (auto add = cur.getDefiningOp()) { - foldedAddPtr = true; - Value off = ensureIndex(rewriter, loc, add.getOperand(1), add); - totalOffset = totalOffset ? rewriter.create(loc, totalOffset, off) - : off; - cur = add.getOperand(0); - } - if (cur != baseBuf) { - baseBuf = cur; - off0 = totalOffset ? OpFoldResult(totalOffset) : off0; - } - } - - auto baseMr = dyn_cast(baseBuf.getType()); - if (!baseMr) { - op.emitError("make_tensor_view base must be memref"); - return failure(); - } - - size_t rank = op.getShape().size(); - int64_t dyn = ShapedType::kDynamic; - SmallVector dynStrides(rank, dyn); - auto layout = - StridedLayoutAttr::get(ctx, /*offset=*/dyn, /*strides=*/dynStrides); - SmallVector dynShape(rank, dyn); - auto mrTy = MemRefType::get(dynShape, baseMr.getElementType(), layout, - baseMr.getMemorySpace()); - - SmallInlineVector sizes; - for (Value value : op.getShape()) - sizes.push_back(ensureIndex(rewriter, loc, value, op)); - SmallInlineVector strides; - for (Value value : op.getStrides()) - strides.push_back(ensureIndex(rewriter, loc, value, op)); - - auto rc = rewriter.create(loc, mrTy, baseBuf, off0, - sizes, strides); - if (foldedAddPtr) - rc->setAttr("pto.addptr_trace", rewriter.getUnitAttr()); - if (auto layoutAttr = op.getLayoutAttr()) - rc->setAttr("layout", layoutAttr); - rewriter.replaceOp(op, rc.getResult()); - } - return success(); -} - static LogicalResult lowerPtrLikeTileBufAddrOps(func::FuncOp func, MLIRContext *ctx) { SmallVector addrOps; @@ -1163,66 +970,6 @@ static LogicalResult bridgeCallOperandsToConvertedCallees(ModuleOp mod, return success(); } -[[maybe_unused]] static LogicalResult foldAddPtrIntoScalarOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector loadScalars; - func.walk([&](mlir::pto::LoadScalarOp op) { loadScalars.push_back(op); }); - for (auto op : loadScalars) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - Value base = op.getPtr(); - Value totalOffset = ensureIndex(rewriter, loc, op.getOffset(), op); - bool foldedAddPtr = foldAddPtrChainIntoOffset(rewriter, loc, base, totalOffset); - if (foldedAddPtr) { - auto newOp = - rewriter.create(loc, op.getValue().getType(), base, - totalOffset); - rewriter.replaceOp(op, newOp.getValue()); - } - } - - DefaultInlineVector storeScalars; - func.walk([&](mlir::pto::StoreScalarOp op) { storeScalars.push_back(op); }); - for (auto op : storeScalars) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - Value base = op.getPtr(); - Value totalOffset = ensureIndex(rewriter, loc, op.getOffset(), op); - bool foldedAddPtr = foldAddPtrChainIntoOffset(rewriter, loc, base, totalOffset); - if (foldedAddPtr) { - rewriter.create(loc, base, totalOffset, op.getValue()); - rewriter.eraseOp(op); - } - } - - DefaultInlineVector addPtrs; - func.walk([&](mlir::pto::AddPtrOp op) { addPtrs.push_back(op.getOperation()); }); - bool changed = true; - while (changed) { - changed = false; - for (auto &op : addPtrs) { - if (!op) - continue; - if (op->use_empty()) { - op->erase(); - op = nullptr; - changed = true; - } - } - } - for (Operation *op : addPtrs) { - if (!op) - continue; - op->emitError( - "addptr must feed make_tensor_view or load/store_scalar for lowering"); - return failure(); - } - return success(); -} - static LogicalResult lowerPartitionViewOps(func::FuncOp func, MLIRContext *ctx) { DefaultInlineVector partitionViews; func.walk([&](mlir::pto::PartitionViewOp op) { partitionViews.push_back(op); }); @@ -1353,14 +1100,6 @@ struct PTOViewToMemrefPass return; } - // ------------------------------------------------------------------ - // Stage 0.20: lower pto.inttoptr result types to GM memrefs. - // ------------------------------------------------------------------ - if (!preserveTileABI && failed(lowerIntToPtrOps(func, ctx))) { - signalPassFailure(); - return; - } - // ------------------------------------------------------------------ // Stage 0.25: synthesize missing A2/A3 tquant tmp tiles before type lowering. // ------------------------------------------------------------------ @@ -1369,14 +1108,6 @@ struct PTOViewToMemrefPass return; } - // ------------------------------------------------------------------ - // Stage 0.30: materialize pto.ptrtoint(addptr ...) byte offsets. - // ------------------------------------------------------------------ - if (!preserveTileABI && failed(lowerPtrToIntOps(func, ctx))) { - signalPassFailure(); - return; - } - // Stage 0.40 Insert pto.bind_tile for function args that were tile_buf. // ------------------------------------------------------------------ // Later materialization and intrinsic folding use BindTileOp as the @@ -1481,26 +1212,6 @@ struct PTOViewToMemrefPass Value baseBuf = op.getOperand(0); OpFoldResult off0 = rewriter.getIndexAttr(0); - // Fold pto.addptr chains into the view base to avoid nested reinterpret_cast. - bool foldedAddPtr = false; - { - Value cur = baseBuf; - Value totalOffset; - while (auto add = cur.getDefiningOp()) { - foldedAddPtr = true; - Value off = ensureIndex(rewriter, loc, add.getOperand(1), add); - if (totalOffset) - totalOffset = rewriter.create(loc, totalOffset, off); - else - totalOffset = off; - cur = add.getOperand(0); - } - if (cur != baseBuf) { - baseBuf = cur; - off0 = totalOffset ? OpFoldResult(totalOffset) : off0; - } - } - auto baseMr = dyn_cast(baseBuf.getType()); if (!baseMr) { op.emitError("make_tensor_view base must be memref"); signalPassFailure(); return; @@ -1530,9 +1241,6 @@ struct PTOViewToMemrefPass auto rc = rewriter.create( loc, mrTy, baseBuf, off0, sizes, strides); - if (foldedAddPtr) { - rc->setAttr("pto.addptr_trace", rewriter.getUnitAttr()); - } if (auto layoutAttr = op.getLayoutAttr()) { rc->setAttr("layout", layoutAttr); } @@ -1625,153 +1333,6 @@ struct PTOViewToMemrefPass } } - // ------------------------------------------------------------------ - // Stage 1.6: Fold pto.addptr chains into load/store_scalar. - // ------------------------------------------------------------------ - DefaultInlineVector loadScalars; - func.walk([&](mlir::pto::LoadScalarOp op) { loadScalars.push_back(op); }); - - for (auto op : loadScalars) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - Value base = op.getPtr(); - Value totalOffset = ensureIndex(rewriter, loc, op.getOffset(), op); - - bool foldedAddPtr = false; - while (auto add = base.getDefiningOp()) { - foldedAddPtr = true; - Value off = ensureIndex(rewriter, loc, add.getOperand(1), add); - if (totalOffset) - totalOffset = rewriter.create(loc, totalOffset, off); - else - totalOffset = off; - base = add.getOperand(0); - } - - if (foldedAddPtr) { - auto newOp = rewriter.create( - loc, op.getValue().getType(), base, totalOffset); - rewriter.replaceOp(op, newOp.getValue()); - } - } - - DefaultInlineVector storeScalars; - func.walk([&](mlir::pto::StoreScalarOp op) { storeScalars.push_back(op); }); - - for (auto op : storeScalars) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - Value base = op.getPtr(); - Value totalOffset = ensureIndex(rewriter, loc, op.getOffset(), op); - - bool foldedAddPtr = false; - while (auto add = base.getDefiningOp()) { - foldedAddPtr = true; - Value off = ensureIndex(rewriter, loc, add.getOperand(1), add); - if (totalOffset) - totalOffset = rewriter.create(loc, totalOffset, off); - else - totalOffset = off; - base = add.getOperand(0); - } - - if (foldedAddPtr) { - rewriter.create( - loc, base, totalOffset, op.getValue()); - rewriter.eraseOp(op); - } - } - - // ------------------------------------------------------------------ - // Stage 1.75: Fold addptr used by initialize_l2g2l_pipe(gm_addr). - // This keeps IR well-typed after function arguments are rewritten from - // !pto.ptr to memref. - // ------------------------------------------------------------------ - bool foldedPipeInitAddPtr = true; - while (foldedPipeInitAddPtr) { - foldedPipeInitAddPtr = false; - DefaultInlineVector addPtrsForPipeInit; - func.walk([&](mlir::pto::AddPtrOp op) { - bool eligible = !op->use_empty(); - for (Operation *user : op->getUsers()) { - auto init = dyn_cast(user); - if (!init || init.getGmAddr() != op->getResult(0)) { - eligible = false; - break; - } - } - if (eligible) - addPtrsForPipeInit.push_back(op); - }); - - for (auto op : addPtrsForPipeInit) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - Value base = op->getOperand(0); - Value totalOffset = ensureIndex(rewriter, loc, op->getOperand(1), op); - while (auto add = base.getDefiningOp()) { - Value off = ensureIndex(rewriter, loc, add->getOperand(1), add); - totalOffset = rewriter.create(loc, totalOffset, off); - base = add->getOperand(0); - } - - auto baseMrTy = dyn_cast(base.getType()); - if (!baseMrTy || baseMrTy.getRank() != 1) - continue; - - int64_t dyn = ShapedType::kDynamic; - auto layout = StridedLayoutAttr::get(ctx, dyn, {dyn}); - auto targetTy = MemRefType::get({dyn}, baseMrTy.getElementType(), layout, - baseMrTy.getMemorySpace()); - SmallVector sizes{rewriter.getIndexAttr(1)}; - SmallVector strides{rewriter.getIndexAttr(1)}; - auto rc = rewriter.create( - loc, targetTy, base, OpFoldResult(totalOffset), sizes, strides); - rc->setAttr("pto.addptr_trace", rewriter.getUnitAttr()); - rewriter.replaceOp(op, rc.getResult()); - foldedPipeInitAddPtr = true; - } - } - - // Clean up: addptr should be folded into make_tensor_view. - DefaultInlineVector addPtrs; - func.walk([&](mlir::pto::AddPtrOp op) { addPtrs.push_back(op.getOperation()); }); - bool changed = true; - while (changed) { - changed = false; - for (auto &op : addPtrs) { - if (!op) - continue; - if (op->use_empty()) { - op->erase(); - op = nullptr; - changed = true; - } - } - } - for (auto *op : addPtrs) { - if (!op) - continue; - if (llvm::all_of(op->getUsers(), [op](Operation *user) { - if (isa(user)) - return true; - auto init = dyn_cast(user); - return init && init.getGmAddr() == op->getResult(0); - })) - continue; - op->emitError("addptr must feed make_tensor_view, " - "initialize_l2g2l_pipe(gm_addr) or load/store_scalar " - "for lowering"); - signalPassFailure(); - return; - } - // ------------------------------------------------------------------ // Stage 3: Rewrite Compute Ops // [关键] 全面使用 op->getOperand(i) 避免 Typed Accessor Crash diff --git a/lib/PTO/Transforms/VPTOPtrNormalize.cpp b/lib/PTO/Transforms/VPTOPtrNormalize.cpp index 404c6634aa..c9e5928142 100644 --- a/lib/PTO/Transforms/VPTOPtrNormalize.cpp +++ b/lib/PTO/Transforms/VPTOPtrNormalize.cpp @@ -308,7 +308,8 @@ struct ConvertPointerCastToCastPtrPattern LogicalResult matchAndRewrite(pto::PointerCastOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - Type convertedType = getTypeConverter()->convertType(op.getResult().getType()); + Type convertedType = + getTypeConverter()->convertType(op.getResult().getType()); auto ptrType = dyn_cast(convertedType); if (!ptrType) return failure(); @@ -322,6 +323,43 @@ struct ConvertPointerCastToCastPtrPattern } }; +struct ConvertIntToPtrToCastPtrPattern + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::IntToPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type convertedType = + getTypeConverter()->convertType(op.getResult().getType()); + if (!isa(convertedType)) + return rewriter.notifyMatchFailure(op, "expected pointer result type"); + + rewriter.replaceOpWithNewOp(op, convertedType, + adaptor.getAddr()); + return success(); + } +}; + +struct ConvertPtrToIntToCastPtrPattern + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(pto::PtrToIntOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type convertedType = getTypeConverter()->convertType(op.getResult().getType()); + if (!isa(convertedType)) + return rewriter.notifyMatchFailure(op, "expected integer result type"); + if (!isa(adaptor.getPtr().getType())) + return rewriter.notifyMatchFailure(op, "expected pointer input type"); + + rewriter.replaceOpWithNewOp(op, convertedType, + adaptor.getPtr()); + return success(); + } +}; + struct ConvertCastPtrPattern : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -816,7 +854,8 @@ struct VPTOPtrNormalizePass scf::SCFDialect>(); target.addDynamicallyLegalDialect([](Operation *op) { return !isa(op); + pto::IntToPtrOp, pto::PtrToIntOp, pto::BindTileOp, + pto::VldsOp, pto::VstsOp, pto::VsstbOp>(op); }); target.addLegalOp(); target.addDynamicallyLegalOp([&](func::FuncOp op) { @@ -943,7 +982,9 @@ struct VPTOPtrNormalizePass populateCallOpTypeConversionPattern(patterns, typeConverter); populateReturnOpTypeConversionPattern(patterns, typeConverter); patterns.add&1 | FileCheck %s +// RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @castptr_ptr_to_int(%src: memref<1xf32, #pto.address_space>) @@ -27,3 +28,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind // CHECK-LABEL: AICORE void castptr_ptr_to_int(__ubuf__ float* // CHECK: helper( + +// NATIVE: IR Dump After PTOViewToMemref +// NATIVE-LABEL: func.func private @helper +// NATIVE: pto.castptr {{.*}} : !pto.ptr -> i64 diff --git a/test/lit/pto/ptr_int_cast.pto b/test/lit/pto/ptr_int_cast.pto index bba4f125dd..39eb5c78f5 100644 --- a/test/lit/pto/ptr_int_cast.pto +++ b/test/lit/pto/ptr_int_cast.pto @@ -1,7 +1,8 @@ // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=CPP +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO -module { +module attributes {pto.kernel_kind = #pto.kernel_kind} { func.func @ptr_int_cast_kernel(%src64: !pto.ptr, %dst32: !pto.ptr, %idx: index) attributes {pto.kernel_kind = #pto.kernel_kind} { %c0 = arith.constant 0 : index %src64_off = pto.addptr %src64, %idx : !pto.ptr -> !pto.ptr @@ -24,36 +25,37 @@ module { } // IR-LABEL: func.func @ptr_int_cast_kernel -// IR: pto.ptrtoint {{.*}} : memref> -> i64 -// IR: arith.index_cast {{.*}} : index to i64 -// IR: arith.constant 8 : i64 -// IR: arith.muli {{.*}} : i64 -// IR: arith.addi {{.*}} : i64 -// IR: pto.inttoptr {{.*}} : i64 -> memref> -// IR-NOT: !pto.ptr +// IR-SAME: %[[SRC:[^,]+]]: !pto.ptr +// IR: %[[OFF:.*]] = pto.addptr %[[SRC]], %{{.*}} : -> +// IR: %[[ADDR:.*]] = pto.ptrtoint %[[OFF]] : !pto.ptr -> i64 +// IR: %[[PTR:.*]] = pto.inttoptr %[[ADDR]] : i64 -> !pto.ptr +// IR-NOT: memref< // CPP-LABEL: AICORE void ptr_int_cast_kernel // CPP-SAME: (__gm__ uint64_t* [[SRC:[_A-Za-z][_A-Za-z0-9]*]], __gm__ uint32_t* [[DST:[_A-Za-z][_A-Za-z0-9]*]], int64_t [[IDX:[_A-Za-z][_A-Za-z0-9]*]]) -// CPP: int64_t [[ADDR:[_A-Za-z][_A-Za-z0-9]*]] = reinterpret_cast([[SRC]]); -// CPP: __gm__ uint32_t* [[SRC32:[_A-Za-z][_A-Za-z0-9]*]] = reinterpret_cast<__gm__ uint32_t*>({{.*}}[[ADDR]] +// CPP: int64_t [[ADDR:[_A-Za-z][_A-Za-z0-9]*]] = reinterpret_cast([[SRC]] + [[IDX]]); +// CPP: __gm__ uint32_t* [[SRC32:[_A-Za-z][_A-Za-z0-9]*]] = reinterpret_cast<__gm__ uint32_t*>([[ADDR]]); // CPP: uint32_t [[VAL:[_A-Za-z][_A-Za-z0-9]*]] = [[SRC32]][{{.*}}]; // CPP: [[DST]][[[IDX]]] = [[VAL]]; // IR-LABEL: func.func @ptrtoint_addptr_multi_consumer -// IR: pto.load_scalar {{.*}} : memref> -> ui64 -// IR: pto.ptrtoint {{.*}} : memref> -> i64 -// IR: arith.index_cast {{.*}} : index to i64 -// IR: arith.constant 8 : i64 -// IR: arith.muli {{.*}} : i64 -// IR: arith.addi {{.*}} : i64 -// IR-NOT: !pto.ptr +// IR-SAME: %[[SRC64:[^,]+]]: !pto.ptr +// IR: %[[OFF2:.*]] = pto.addptr %[[SRC64]], %{{.*}} : -> +// IR: pto.load_scalar %[[OFF2]]{{.*}} : !pto.ptr -> ui64 +// IR: pto.ptrtoint %[[OFF2]] : !pto.ptr -> i64 +// IR-NOT: memref< // CPP-LABEL: AICORE void ptrtoint_addptr_multi_consumer // CPP-SAME: (__gm__ uint64_t* [[SRC64:[_A-Za-z][_A-Za-z0-9]*]], __gm__ uint64_t* [[DST64:[_A-Za-z][_A-Za-z0-9]*]], __gm__ int64_t* [[DSTADDR:[_A-Za-z][_A-Za-z0-9]*]], int64_t [[IDX2:[_A-Za-z][_A-Za-z0-9]*]]) // CPP: const int64_t [[ZERO:[_A-Za-z][_A-Za-z0-9]*]] = 0; -// CPP: const int64_t [[ELEM_BYTES:[_A-Za-z][_A-Za-z0-9]*]] = 8; -// CPP: uint64_t [[A:[_A-Za-z][_A-Za-z0-9]*]] = [[SRC64]][[[IDX2]]]; -// CPP: int64_t [[ADDR2:[_A-Za-z][_A-Za-z0-9]*]] = reinterpret_cast([[SRC64]]); -// CPP: int64_t [[BYTE_ADDR:[_A-Za-z][_A-Za-z0-9]*]] = {{.*}}[[ADDR2]]{{.*}}[[IDX2]]{{.*}}[[ELEM_BYTES]] +// CPP: __gm__ uint64_t* [[OFF3:[_A-Za-z][_A-Za-z0-9]*]] = [[SRC64]] + [[IDX2]]; +// CPP: uint64_t [[A:[_A-Za-z][_A-Za-z0-9]*]] = [[OFF3]][[[ZERO]]]; +// CPP: int64_t [[BYTE_ADDR:[_A-Za-z][_A-Za-z0-9]*]] = reinterpret_cast([[OFF3]]); // CPP: [[DST64]][[[ZERO]]] = [[A]]; // CPP: [[DSTADDR]][[[ZERO]]] = [[BYTE_ADDR]]; + +// VPTO-LABEL: func.func @ptr_int_cast_kernel +// VPTO: %[[VPTO_OFF:.*]] = pto.addptr +// VPTO: %[[VPTO_ADDR:.*]] = pto.castptr %[[VPTO_OFF]] : !pto.ptr -> i64 +// VPTO: %{{.*}} = pto.castptr %[[VPTO_ADDR]] : i64 -> !pto.ptr +// VPTO-NOT: memref From 766e8764d401e9280797c98febb110c7ab20663f Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 20 Jul 2026 14:47:09 +0800 Subject: [PATCH 80/91] Keep SCF tile values native across control flow --- lib/PTO/Transforms/FoldTileBufIntrinsics.cpp | 35 ++- .../Transforms/PTOMaterializeTileHandles.cpp | 255 +----------------- lib/PTO/Transforms/PTOViewToMemref.cpp | 93 ------- ...alize_tile_handles_control_flow_result.pto | 44 +-- 4 files changed, 61 insertions(+), 366 deletions(-) diff --git a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp index af5e49d85b..4eac1b7fca 100644 --- a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp +++ b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp @@ -197,6 +197,16 @@ static Value unwrapBridgingCasts(Value v) { return v; } +static bool isSCFTileCarrier(Value value) { + if (auto result = dyn_cast(value)) { + return isa(result.getOwner()); + } + + auto blockArg = dyn_cast(value); + return blockArg && + isa_and_nonnull(blockArg.getOwner()->getParentOp()); +} + static std::optional resolveTileHandle(Value tileBuf, Operation *user) { if (auto regionResult = dyn_cast(tileBuf)) { @@ -675,24 +685,30 @@ struct FoldTileBufIntrinsicsPass if (!isa(gvsOp.getSource().getType())) continue; - auto handleInfo = resolveTileHandle(gvsOp.getSource(), gvsOp); - if (!handleInfo) - return signalPassFailure(); - builder.setInsertionPoint(gvsOp); auto tileTy = cast(gvsOp.getSource().getType()); auto validShape = tileTy.getValidShape(); - Value rowReplacement = handleInfo->validRow; + Value rowReplacement; if (!validShape.empty() && validShape[0] != ShapedType::kDynamic) rowReplacement = builder.create(gvsOp.getLoc(), validShape[0]); - Value colReplacement = handleInfo->validCol; + Value colReplacement; if (validShape.size() >= 2 && validShape[1] != ShapedType::kDynamic) colReplacement = builder.create(gvsOp.getLoc(), validShape[1]); + if (!rowReplacement || !colReplacement) { + auto handleInfo = resolveTileHandle(gvsOp.getSource(), gvsOp); + if (!handleInfo) + return signalPassFailure(); + if (!rowReplacement) + rowReplacement = handleInfo->validRow; + if (!colReplacement) + colReplacement = handleInfo->validCol; + } + if (!rowReplacement || !colReplacement) { gvsOp.emitError("FoldTileBufIntrinsics: pto.get_validshape could not " "resolve a concrete valid_row / valid_col"); @@ -711,6 +727,13 @@ struct FoldTileBufIntrinsicsPass // When the requested result type is already !pto.ptr<...>, cast from the // recovered memref instead of leaving tile_buf_addr in the IR. for (auto addrOp : addrOps) { + // An SCF result/iter_arg is already a runtime-selected tile handle. + // Keep tile_buf_addr attached to that handle; VPTO pointer + // normalization converts it directly without choosing one branch's + // allocation address here. + if (isSCFTileCarrier(addrOp.getSrc())) + continue; + auto handleInfo = resolveTileHandle(addrOp.getSrc(), addrOp); if (!handleInfo) return signalPassFailure(); diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index 2a4c4bf1d1..c2a4434d93 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -24,7 +24,6 @@ #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Builders.h" -#include "mlir/IR/IRMapping.h" #include "mlir/IR/PatternMatch.h" #include "mlir/Pass/Pass.h" @@ -671,240 +670,6 @@ static Value lookupMaterializedTileHandle( return it->second; } -static void cloneBlockWithoutTerminator(Block *from, Block *to, - IRMapping &mapping) { - OpBuilder builder(to, to->end()); - Operation *terminator = from->getTerminator(); - for (Operation &op : *from) { - if (&op == terminator) - break; - builder.clone(op, mapping); - } -} - -static bool isOperationNestedIn(Operation *op, Operation *ancestor) { - for (; op; op = op->getParentOp()) { - if (op == ancestor) - return true; - } - return false; -} - -static bool isValueOwnedByOperation(Value value, Operation *owner) { - if (auto result = dyn_cast(value)) - return isOperationNestedIn(result.getOwner(), owner); - if (auto arg = dyn_cast(value)) - return isOperationNestedIn(arg.getOwner()->getParentOp(), owner); - return false; -} - -static void eraseTileHandlesOwnedBy(Operation *owner, - DenseMap &tileHandles) { - SmallVector keysToErase; - for (auto &entry : tileHandles) { - if (isValueOwnedByOperation(entry.first, owner) || - isValueOwnedByOperation(entry.second, owner)) - keysToErase.push_back(entry.first); - } - for (Value key : keysToErase) - tileHandles.erase(key); -} - -static FailureOr -materializeSCFIfResults(ModuleOp module, DenseMap &tileHandles) { - bool changed = false; - - SmallVector ifOps; - module.walk([&](scf::IfOp ifOp) { ifOps.push_back(ifOp); }); - - for (scf::IfOp ifOp : llvm::reverse(ifOps)) { - if (ifOp.getNumResults() == 0) - continue; - - auto thenYield = dyn_cast(ifOp.thenBlock()->getTerminator()); - auto elseYield = dyn_cast(ifOp.elseBlock()->getTerminator()); - if (!thenYield || !elseYield) - continue; - - if (thenYield.getNumOperands() != ifOp.getNumResults() || - elseYield.getNumOperands() != ifOp.getNumResults()) { - ifOp.emitOpError("result count does not match branch yield operands"); - return failure(); - } - - SmallVector resultTypes(ifOp->getResultTypes()); - SmallVector thenYieldOperands(thenYield.getOperands().begin(), - thenYield.getOperands().end()); - SmallVector elseYieldOperands(elseYield.getOperands().begin(), - elseYield.getOperands().end()); - SmallVector materializedResults; - - for (auto [idx, result] : llvm::enumerate(ifOp.getResults())) { - if (!isLocalTileMemRef(result.getType())) - continue; - - Value thenTile = - lookupMaterializedTileHandle(thenYield.getOperand(idx), tileHandles); - Value elseTile = - lookupMaterializedTileHandle(elseYield.getOperand(idx), tileHandles); - if (!thenTile || !elseTile) - continue; - - if (thenTile.getType() != elseTile.getType()) { - ifOp.emitOpError() - << "cannot materialize tile result #" << idx - << " because branch tile types differ: " << thenTile.getType() - << " vs " << elseTile.getType(); - return failure(); - } - - Type tileTy = thenTile.getType(); - resultTypes[idx] = tileTy; - thenYieldOperands[idx] = thenTile; - elseYieldOperands[idx] = elseTile; - materializedResults.push_back(idx); - } - - if (materializedResults.empty()) - continue; - - OpBuilder builder(ifOp); - auto newIf = builder.create( - ifOp.getLoc(), TypeRange(resultTypes), ifOp.getCondition(), - /*addThenBlock=*/true, /*addElseBlock=*/true); - newIf->setAttrs(ifOp->getAttrs()); - - IRMapping thenMapping; - cloneBlockWithoutTerminator(ifOp.thenBlock(), newIf.thenBlock(), - thenMapping); - builder.setInsertionPointToEnd(newIf.thenBlock()); - for (Value &operand : thenYieldOperands) - operand = thenMapping.lookupOrDefault(operand); - builder.create(thenYield.getLoc(), thenYieldOperands); - - IRMapping elseMapping; - cloneBlockWithoutTerminator(ifOp.elseBlock(), newIf.elseBlock(), - elseMapping); - builder.setInsertionPointToEnd(newIf.elseBlock()); - for (Value &operand : elseYieldOperands) - operand = elseMapping.lookupOrDefault(operand); - builder.create(elseYield.getLoc(), elseYieldOperands); - - for (auto [oldResult, newResult] : - llvm::zip_equal(ifOp.getResults(), newIf.getResults())) - oldResult.replaceAllUsesWith(newResult); - - for (unsigned idx : materializedResults) { - tileHandles[newIf.getResult(idx)] = newIf.getResult(idx); - } - - eraseTileHandlesOwnedBy(ifOp, tileHandles); - ifOp.erase(); - changed = true; - } - - return changed; -} - -static FailureOr -materializeSCFForResults(ModuleOp module, DenseMap &tileHandles) { - bool changed = false; - - SmallVector forOps; - module.walk([&](scf::ForOp forOp) { forOps.push_back(forOp); }); - - for (scf::ForOp forOp : llvm::reverse(forOps)) { - if (forOp.getNumResults() == 0) - continue; - - auto yield = dyn_cast(forOp.getBody()->getTerminator()); - if (!yield) - continue; - - if (yield.getNumOperands() != forOp.getNumResults() || - forOp.getInitArgs().size() != forOp.getNumResults()) { - forOp.emitOpError("result count does not match iter/yield operands"); - return failure(); - } - - SmallVector initArgs(forOp.getInitArgs().begin(), - forOp.getInitArgs().end()); - SmallVector yieldOperands(yield.getOperands().begin(), - yield.getOperands().end()); - SmallVector materializedResults; - - for (auto [idx, result] : llvm::enumerate(forOp.getResults())) { - if (!isLocalTileMemRef(result.getType())) - continue; - - Value initTile = - lookupMaterializedTileHandle(forOp.getInitArgs()[idx], tileHandles); - if (!initTile) - continue; - - BlockArgument iterArg = forOp.getRegionIterArg(idx); - Value yieldValue = yield.getOperand(idx); - Value yieldTile = lookupMaterializedTileHandle(yieldValue, tileHandles); - bool yieldIsIterArg = !yieldTile && yieldValue == iterArg; - if (yieldIsIterArg) - yieldTile = iterArg; - if (!yieldTile) - continue; - - Type yieldTy = yieldIsIterArg ? initTile.getType() : yieldTile.getType(); - if (initTile.getType() != yieldTy) { - forOp.emitOpError() - << "cannot materialize tile result #" << idx - << " because init/yield tile types differ: " << initTile.getType() - << " vs " << yieldTy; - return failure(); - } - - initArgs[idx] = initTile; - yieldOperands[idx] = yieldTile; - materializedResults.push_back(idx); - } - - if (materializedResults.empty()) - continue; - - OpBuilder builder(forOp); - auto newFor = builder.create( - forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), - forOp.getStep(), initArgs); - newFor->setAttrs(forOp->getAttrs()); - - IRMapping mapping; - mapping.map(forOp.getInductionVar(), newFor.getInductionVar()); - for (auto [oldArg, newArg] : - llvm::zip_equal(forOp.getBody()->getArguments().drop_front(), - newFor.getBody()->getArguments().drop_front())) - mapping.map(oldArg, newArg); - - cloneBlockWithoutTerminator(forOp.getBody(), newFor.getBody(), mapping); - builder.setInsertionPointToEnd(newFor.getBody()); - for (Value &operand : yieldOperands) - operand = mapping.lookupOrDefault(operand); - builder.create(yield.getLoc(), yieldOperands); - - for (auto [oldResult, newResult] : - llvm::zip_equal(forOp.getResults(), newFor.getResults())) - oldResult.replaceAllUsesWith(newResult); - - for (unsigned idx : materializedResults) { - BlockArgument newIterArg = newFor.getRegionIterArg(idx); - tileHandles[newIterArg] = newIterArg; - tileHandles[newFor.getResult(idx)] = newFor.getResult(idx); - } - - eraseTileHandlesOwnedBy(forOp, tileHandles); - forOp.erase(); - changed = true; - } - - return changed; -} - static FailureOr materializeFusionRegionResults(ModuleOp module, DenseMap &tileHandles) { @@ -955,24 +720,12 @@ materializeFusionRegionResults(ModuleOp module, } static LogicalResult -materializeControlFlowTileResults(ModuleOp module, - DenseMap &tileHandles) { +materializeFusionRegionTileResults(ModuleOp module, + DenseMap &tileHandles) { bool changed = false; do { changed = false; - FailureOr ifChanged = - materializeSCFIfResults(module, tileHandles); - if (failed(ifChanged)) - return failure(); - changed |= *ifChanged; - - FailureOr forChanged = - materializeSCFForResults(module, tileHandles); - if (failed(forChanged)) - return failure(); - changed |= *forChanged; - FailureOr fusionChanged = materializeFusionRegionResults(module, tileHandles); if (failed(fusionChanged)) @@ -1528,7 +1281,7 @@ struct PTOMaterializeTileHandlesPass restoreMaterializedTileBufAddrOps(module, builder, ctx, tileHandles); - if (failed(materializeControlFlowTileResults(module, tileHandles))) { + if (failed(materializeFusionRegionTileResults(module, tileHandles))) { signalPassFailure(); return; } @@ -1589,7 +1342,7 @@ struct PTOMaterializeTileHandlesPass return; } - if (failed(materializeControlFlowTileResults(module, tileHandles))) { + if (failed(materializeFusionRegionTileResults(module, tileHandles))) { signalPassFailure(); return; } diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 490178db78..6a1800ad76 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -707,88 +707,6 @@ static Type convertPTOTypeToMemRef(Type t) { return t; } -// Ensure scf.if result types follow the rewritten yield operand types. -// PTOViewToMemref rewrites tile values to memref in branch bodies, but scf.if -// result types are not auto-updated by those op-local rewrites. -static LogicalResult reconcileSCFIfResultTypes(func::FuncOp func) { - DefaultInlineVector ifOps; - func.walk([&](scf::IfOp ifOp) { ifOps.push_back(ifOp); }); - - for (scf::IfOp ifOp : ifOps) { - if (ifOp.getNumResults() == 0) - continue; - - auto thenYield = dyn_cast(ifOp.thenBlock()->getTerminator()); - auto elseYield = dyn_cast(ifOp.elseBlock()->getTerminator()); - if (!thenYield || !elseYield) { - ifOp.emitError("result-bearing scf.if must end with scf.yield in both " - "then/else regions"); - return failure(); - } - - if (thenYield.getNumOperands() != ifOp.getNumResults() || - elseYield.getNumOperands() != ifOp.getNumResults()) { - ifOp.emitError("scf.if result count does not match yielded values"); - return failure(); - } - - for (unsigned i = 0; i < ifOp.getNumResults(); ++i) { - Type thenTy = thenYield.getOperand(i).getType(); - Type elseTy = elseYield.getOperand(i).getType(); - if (thenTy != elseTy) { - ifOp.emitError() << "scf.if branch yield type mismatch at result #" << i - << ": then=" << thenTy << ", else=" << elseTy; - return failure(); - } - - if (ifOp.getResult(i).getType() != thenTy) - ifOp.getResult(i).setType(thenTy); - } - } - - return success(); -} - -static LogicalResult reconcileSCFForResultTypes(func::FuncOp func) { - DefaultInlineVector forOps; - func.walk([&](scf::ForOp forOp) { forOps.push_back(forOp); }); - - for (scf::ForOp forOp : forOps) { - if (forOp.getNumResults() == 0) - continue; - - auto yield = dyn_cast(forOp.getBody()->getTerminator()); - if (!yield) { - forOp.emitError("result-bearing scf.for must end with scf.yield"); - return failure(); - } - - if (yield.getNumOperands() != forOp.getNumResults() || - forOp.getInitArgs().size() != forOp.getNumResults()) { - forOp.emitError("scf.for result count does not match iter/yield values"); - return failure(); - } - - for (unsigned i = 0; i < forOp.getNumResults(); ++i) { - Type initTy = forOp.getInitArgs()[i].getType(); - Type yieldTy = yield.getOperand(i).getType(); - if (initTy != yieldTy) { - forOp.emitError() << "scf.for init/yield type mismatch at result #" << i - << ": init=" << initTy << ", yield=" << yieldTy; - return failure(); - } - - BlockArgument iterArg = forOp.getRegionIterArg(i); - if (iterArg.getType() != initTy) - iterArg.setType(initTy); - if (forOp.getResult(i).getType() != initTy) - forOp.getResult(i).setType(initTy); - } - } - - return success(); -} - // Ensure pto.fusion_region result types follow the rewritten pto.yield operand // types. PTOViewToMemref rewrites region-local tile values to memref, but the // region result types are not auto-updated by those op-local rewrites. @@ -3230,17 +3148,6 @@ struct PTOViewToMemrefPass } } - // ------------------------------------------------------------------ - // Stage 4: Reconcile control-flow result types - // ------------------------------------------------------------------ - if (failed(reconcileSCFIfResultTypes(func))) { - signalPassFailure(); - return; - } - if (failed(reconcileSCFForResultTypes(func))) { - signalPassFailure(); - return; - } } if (failed(bridgeCallOperandsToConvertedCallees(mod, ctx))) { diff --git a/test/lit/pto/materialize_tile_handles_control_flow_result.pto b/test/lit/pto/materialize_tile_handles_control_flow_result.pto index 4fba96e31b..3548bd86c9 100644 --- a/test/lit/pto/materialize_tile_handles_control_flow_result.pto +++ b/test/lit/pto/materialize_tile_handles_control_flow_result.pto @@ -1,7 +1,9 @@ -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO -module { +module attributes {pto.kernel_kind = #pto.kernel_kind} { func.func @control_flow_result_tile_materialization(%cond: i1) { %src = pto.alloc_tile : !pto.tile_buf %add = pto.alloc_tile : !pto.tile_buf @@ -48,20 +50,22 @@ module { } } -// IR: %[[SEL:.*]] = scf.if {{.*}} -> (!pto.tile_buf) { -// IR: pto.tadd -// IR: scf.yield %{{.*}} : !pto.tile_buf -// IR: } else { -// IR: pto.tmul -// IR: scf.yield %{{.*}} : !pto.tile_buf -// IR: } -// IR: pto.tadd ins(%[[SEL]], %[[SEL]] -// IR-LABEL: func.func @loop_carried_tile_materialization -// IR: %[[RES:.*]] = scf.for {{.*}} -> (!pto.tile_buf) { -// IR: pto.tadd -// IR: scf.yield %{{.*}} : !pto.tile_buf -// IR: } -// IR: pto.tmul ins(%[[RES]], %[[RES]] +// NATIVE: IR Dump After PTOViewToMemref +// MATERIALIZED: IR Dump After PTOMaterializeTileHandles +// COMMON: %[[SEL:.*]] = scf.if {{.*}} -> (!pto.tile_buf) { +// COMMON: pto.tadd +// COMMON: scf.yield %{{.*}} : !pto.tile_buf +// COMMON: } else { +// COMMON: pto.tmul +// COMMON: scf.yield %{{.*}} : !pto.tile_buf +// COMMON: } +// COMMON: pto.tadd ins(%[[SEL]], %[[SEL]] +// COMMON-LABEL: func.func @loop_carried_tile_materialization +// COMMON: %[[RES:.*]] = scf.for {{.*}} -> (!pto.tile_buf) { +// COMMON: pto.tadd +// COMMON: scf.yield %{{.*}} : !pto.tile_buf +// COMMON: } +// COMMON: pto.tmul ins(%[[RES]], %[[RES]] // EMITC-LABEL: AICORE void control_flow_result_tile_materialization // EMITC: if ( @@ -73,3 +77,11 @@ module { // EMITC: for ( // EMITC: TADD( // EMITC: TMUL( + +// VPTO-LABEL: func.func @control_flow_result_tile_materialization +// VPTO: %[[SELECTED:.*]] = arith.select {{.*}} : !pto.tile_buf +// VPTO: pto.tile_buf_addr %[[SELECTED]] : !pto.tile_buf -> !pto.ptr +// VPTO-LABEL: func.func @loop_carried_tile_materialization +// VPTO: %[[RESULT:.*]] = scf.for {{.*}} -> (!pto.tile_buf) { +// VPTO: scf.yield {{.*}} : !pto.tile_buf +// VPTO: pto.tile_buf_addr %[[RESULT]] : !pto.tile_buf -> !pto.ptr From f9eefb842cce8117a47099486e8799ce298b6cc4 Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 20 Jul 2026 14:54:12 +0800 Subject: [PATCH 81/91] Keep fusion region tile values native --- .../Transforms/PTOMaterializeTileHandles.cpp | 76 ------------------- lib/PTO/Transforms/PTOViewToMemref.cpp | 42 ---------- 2 files changed, 118 deletions(-) diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index c2a4434d93..81dccf5579 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -670,72 +670,6 @@ static Value lookupMaterializedTileHandle( return it->second; } -static FailureOr -materializeFusionRegionResults(ModuleOp module, - DenseMap &tileHandles) { - bool changed = false; - - SmallVector fusionRegions; - module.walk([&](pto::FusionRegionOp fusionRegion) { - fusionRegions.push_back(fusionRegion); - }); - - for (pto::FusionRegionOp fusionRegion : llvm::reverse(fusionRegions)) { - if (fusionRegion.getNumResults() == 0) - continue; - - auto yield = - dyn_cast(fusionRegion.getBody().front().getTerminator()); - if (!yield) { - fusionRegion.emitOpError( - "result-bearing pto.fusion_region must terminate with pto.yield"); - return failure(); - } - if (yield.getNumOperands() != fusionRegion.getNumResults()) { - fusionRegion.emitOpError() - << "cannot materialize tile results because yield/result arity " - "mismatch: " - << yield.getNumOperands() << " vs " << fusionRegion.getNumResults(); - return failure(); - } - - for (auto [idx, result] : llvm::enumerate(fusionRegion.getResults())) { - if (!isLocalTileMemRef(result.getType())) - continue; - - Value yieldTile = - lookupMaterializedTileHandle(yield.getOperand(idx), tileHandles); - if (!yieldTile) - continue; - - Type tileTy = yieldTile.getType(); - yield->setOperand(idx, yieldTile); - result.setType(tileTy); - tileHandles[result] = result; - changed = true; - } - } - - return changed; -} - -static LogicalResult -materializeFusionRegionTileResults(ModuleOp module, - DenseMap &tileHandles) { - bool changed = false; - do { - changed = false; - - FailureOr fusionChanged = - materializeFusionRegionResults(module, tileHandles); - if (failed(fusionChanged)) - return failure(); - changed |= *fusionChanged; - } while (changed); - - return success(); -} - static Value getAllocValidOperand(TileBufType tileTy, Value operand, unsigned dim, OpBuilder &builder, Location loc) { @@ -1281,11 +1215,6 @@ struct PTOMaterializeTileHandlesPass restoreMaterializedTileBufAddrOps(module, builder, ctx, tileHandles); - if (failed(materializeFusionRegionTileResults(module, tileHandles))) { - signalPassFailure(); - return; - } - SmallVector, 32> operandsToRewrite; module.walk([&](Operation *op) { if (!shouldMaterializeOperand(op)) @@ -1342,11 +1271,6 @@ struct PTOMaterializeTileHandlesPass return; } - if (failed(materializeFusionRegionTileResults(module, tileHandles))) { - signalPassFailure(); - return; - } - bool erasedBind = true; while (erasedBind) { erasedBind = false; diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 6a1800ad76..4e97ad80ba 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -707,43 +707,6 @@ static Type convertPTOTypeToMemRef(Type t) { return t; } -// Ensure pto.fusion_region result types follow the rewritten pto.yield operand -// types. PTOViewToMemref rewrites region-local tile values to memref, but the -// region result types are not auto-updated by those op-local rewrites. -static LogicalResult reconcileFusionRegionResultTypes(func::FuncOp func) { - SmallVector fusionRegions; - func.walk([&](pto::FusionRegionOp fusionRegion) { - fusionRegions.push_back(fusionRegion); - }); - - for (pto::FusionRegionOp fusionRegion : fusionRegions) { - if (fusionRegion.getNumResults() == 0) - continue; - - auto yieldOp = - dyn_cast(fusionRegion.getBody().front().getTerminator()); - if (!yieldOp) { - fusionRegion.emitError("result-bearing pto.fusion_region must end with " - "pto.yield"); - return failure(); - } - - if (yieldOp.getNumOperands() != fusionRegion.getNumResults()) { - fusionRegion.emitError( - "pto.fusion_region result count does not match yielded values"); - return failure(); - } - - for (unsigned i = 0; i < fusionRegion.getNumResults(); ++i) { - Type yieldedTy = yieldOp.getOperand(i).getType(); - if (fusionRegion.getResult(i).getType() != yieldedTy) - fusionRegion.getResult(i).setType(yieldedTy); - } - } - - return success(); -} - static void markForceDynamicValidShape(Operation *op, bool force, MLIRContext *ctx) { if (force) { @@ -1202,10 +1165,6 @@ struct PTOViewToMemrefPass // ------------------------------------------------------------------ // Stage 1.4: treshape/bitcast stay tile-native. // ------------------------------------------------------------------ - if (failed(reconcileFusionRegionResultTypes(func))) { - signalPassFailure(); - return; - // ------------------------------------------------------------------ // Stage 1.5: Lower pto.get_tensor_view_stride -> strided memref metadata // ------------------------------------------------------------------ @@ -1249,7 +1208,6 @@ struct PTOViewToMemrefPass rewriter.create(loc, view); rewriter.replaceOp(op, metadata.getStrides()[dimIndex]); } - } // ------------------------------------------------------------------ // Stage 3: Rewrite Compute Ops From 7635046ed9815cead9f43055a6d1a898187e2b3a Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 20 Jul 2026 15:34:05 +0800 Subject: [PATCH 82/91] Keep function and call tile ABI native --- .../Transforms/PTOMaterializeTileHandles.cpp | 374 ------------------ lib/PTO/Transforms/PTOViewToMemref.cpp | 184 --------- ...lize_tile_handles_subkernel_helper_abi.pto | 33 +- ...ptodsl_subkernel_helper_tile_abi_emitc.pto | 73 ++-- 4 files changed, 49 insertions(+), 615 deletions(-) diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index 81dccf5579..387db6bd51 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -659,17 +659,6 @@ static void emitMissingExplicitAddressError(Operation *owner, Value value) { diag << "; unsupported defining op is '" << def->getName() << "'"; } -static Value lookupMaterializedTileHandle( - Value value, DenseMap &tileHandles) { - if (isa(value.getType())) - return value; - - auto it = tileHandles.find(value); - if (it == tileHandles.end()) - return Value(); - return it->second; -} - static Value getAllocValidOperand(TileBufType tileTy, Value operand, unsigned dim, OpBuilder &builder, Location loc) { @@ -713,356 +702,6 @@ static bool isTileViewSemantics(StringAttr viewSemantics) { viewSemantics.getValue() == "bitcast"); } -static bool isPTODSLSubkernelHelper(func::FuncOp func) { - return func->hasAttr("pto.ptodsl.subkernel_helper"); -} - -static std::optional> -inferSubkernelHelperTileSignature(func::FuncOp func, MLIRContext *ctx) { - if (func.isExternal() || func.empty() || !isPTODSLSubkernelHelper(func)) - return std::nullopt; - - Block &entry = func.front(); - SmallVector inputTypes(func.getFunctionType().getInputs().begin(), - func.getFunctionType().getInputs().end()); - bool changed = false; - - for (BlockArgument arg : entry.getArguments()) { - if (!isLocalTileMemRef(arg.getType())) - continue; - - BindTileOp bind; - for (Operation *user : arg.getUsers()) { - auto candidate = dyn_cast(user); - if (!candidate || candidate.getSource() != arg) - continue; - bind = candidate; - break; - } - if (!bind) - continue; - - auto memTy = dyn_cast(arg.getType()); - if (!memTy) - continue; - - TileHandleMetadata meta = getTileHandleMetadata(bind.getResult(), ctx); - inputTypes[arg.getArgNumber()] = buildTileTypeFromMemRef(memTy, meta, ctx); - changed = true; - } - - if (!changed) - return std::nullopt; - return inputTypes; -} - -static BindTileOp findEntryArgBindTile(BlockArgument arg) { - for (Operation *user : arg.getUsers()) { - auto bind = dyn_cast(user); - if (bind && bind.getSource() == arg) - return bind; - } - return {}; -} - -static Operation *findMaterializedTileAnchorForValue(Value value) { - if (Operation *def = value.getDefiningOp()) { - if (isMaterializedTileAnchor(def)) - return def; - } - return nullptr; -} - -static Value materializeValueForSubkernelCall( - func::CallOp call, unsigned operandNo, Type tileTy, OpBuilder &builder, - MLIRContext *ctx, DenseMap &tileHandles, - bool &failedMaterialization) { - Value oldValue = call.getOperand(operandNo); - if (oldValue.getType() == tileTy) - return oldValue; - - if (Value materialized = lookupMaterializedTileHandle(oldValue, tileHandles)) { - if (materialized.getType() == tileTy) - return materialized; - } - - auto memTy = dyn_cast(oldValue.getType()); - if (!memTy || !isLocalTileMemRef(memTy)) - return Value(); - - TileHandleMetadata meta = getTileHandleMetadata(oldValue, ctx); - auto expectedTileTy = dyn_cast(tileTy); - if (expectedTileTy) - meta.config = expectedTileTy.getConfigAttr(); - Type materializedTy = - expectedTileTy ? tileTy : Type(buildTileTypeFromMemRef(memTy, meta, ctx)); - - Operation *anchor = findMaterializedTileAnchorForValue(oldValue); - Location loc = anchor ? anchor->getLoc() : call.getLoc(); - if (anchor) - builder.setInsertionPointAfter(anchor); - else - builder.setInsertionPoint(call); - Value addr = computeExplicitAddress(oldValue, builder, loc); - if (!addr && isUnsupportedControlFlowAddress(oldValue)) { - emitMissingExplicitAddressError(call, oldValue); - failedMaterialization = true; - return Value(); - } - - auto alloc = builder.create( - loc, materializedTy, addr ? addr : Value(), - getAllocValidOperand(cast(materializedTy), meta.validRow, 0, - builder, loc), - getAllocValidOperand(cast(materializedTy), meta.validCol, 1, - builder, loc)); - copyMaterializedTileAttrs(meta.attrs, alloc); - Value materialized = alloc.getResult(); - tileHandles[oldValue] = materialized; - return materialized; -} - -static LogicalResult restorePTODSLSubkernelHelperTileABI( - ModuleOp module, OpBuilder &builder, MLIRContext *ctx, - DenseMap &tileHandles, bool &failedMaterialization) { - DenseMap> helperInputTypes; - - for (func::FuncOp func : module.getOps()) { - std::optional> maybeInputs = - inferSubkernelHelperTileSignature(func, ctx); - if (!maybeInputs) - continue; - - FunctionType fnTy = func.getFunctionType(); - if (maybeInputs->size() != fnTy.getNumInputs()) { - func.emitOpError("cannot restore PTODSL subkernel helper tile ABI: " - "inferred input arity does not match function type"); - return failure(); - } - - Block &entry = func.front(); - for (auto [arg, newTy] : llvm::zip(entry.getArguments(), *maybeInputs)) { - if (arg.getType() == newTy) - continue; - - BindTileOp bind = findEntryArgBindTile(arg); - Type oldTy = arg.getType(); - builder.setInsertionPointToStart(&entry); - auto backCast = builder.create( - func.getLoc(), oldTy, arg); - arg.setType(newTy); - arg.replaceAllUsesWith(backCast.getResult(0)); - backCast.getInputsMutable().assign(ValueRange{arg}); - if (bind) - tileHandles[bind.getResult()] = arg; - } - func.setFunctionType( - FunctionType::get(ctx, *maybeInputs, fnTy.getResults())); - func.setPrivate(); - helperInputTypes[func.getSymName()] = *maybeInputs; - } - - if (helperInputTypes.empty()) - return success(); - - SmallVector calls; - module.walk([&](func::CallOp call) { - if (helperInputTypes.contains(call.getCallee())) - calls.push_back(call); - }); - - for (func::CallOp call : calls) { - auto it = helperInputTypes.find(call.getCallee()); - if (it == helperInputTypes.end()) - continue; - ArrayRef inputTypes = it->second; - if (inputTypes.size() != call.getNumOperands()) { - call.emitOpError("cannot restore PTODSL subkernel helper tile ABI: " - "call operand arity does not match callee"); - return failure(); - } - - for (auto [idx, inputTy] : llvm::enumerate(inputTypes)) { - if (!isa(inputTy)) - continue; - - Value materialized = materializeValueForSubkernelCall( - call, idx, inputTy, builder, ctx, tileHandles, failedMaterialization); - if (!materialized) { - call.emitOpError("cannot restore PTODSL subkernel helper tile ABI for " - "operand #") - << idx << ": expected " << inputTy << " but found " - << call.getOperand(idx).getType(); - return failure(); - } - call->setOperand(idx, materialized); - } - } - - return success(); -} - -static bool isFlattenedTileBufAddrView(memref::ReinterpretCastOp cast) { - auto resultTy = dyn_cast(cast.getResult().getType()); - if (!resultTy || resultTy.getRank() != 1) - return false; - - auto sourceTy = dyn_cast(cast.getSource().getType()); - if (!sourceTy || !isLocalTileMemRef(sourceTy)) - return false; - - auto offset = cast.getMixedOffsets(); - if (offset.size() != 1) - return false; - if (auto attr = offset.front().dyn_cast()) { - auto intAttr = dyn_cast(attr); - if (!intAttr || intAttr.getInt() != 0) - return false; - } - - auto strides = cast.getMixedStrides(); - if (strides.size() != 1) - return false; - if (auto attr = strides.front().dyn_cast()) { - auto intAttr = dyn_cast(attr); - if (!intAttr || intAttr.getInt() != 1) - return false; - } - - return true; -} - -static bool canRetypeCalleeOperand(ModuleOp module, func::CallOp call, - unsigned operandNo, Type newTy) { - func::FuncOp callee = module.lookupSymbol(call.getCallee()); - if (!callee) - return false; - FunctionType fnTy = callee.getFunctionType(); - if (operandNo >= fnTy.getNumInputs()) - return false; - if (fnTy.getInput(operandNo) == newTy) - return true; - return callee.isExternal(); -} - -static bool canRetypeTileBufAddrDirectUse(Operation *owner, unsigned operandNo) { - if (auto vlds = dyn_cast(owner)) - return operandNo == vlds.getSourceMutable().getOperandNumber(); - if (auto vsts = dyn_cast(owner)) - return operandNo == vsts.getDestinationMutable().getOperandNumber(); - if (auto vsstb = dyn_cast(owner)) - return operandNo == vsstb.getDestinationMutable().getOperandNumber(); - return false; -} - -static void updateDirectUseResultTypesAfterPtrRetype(Operation *owner, - unsigned operandNo, - Type newTy) { - if (auto vlds = dyn_cast(owner)) { - if (operandNo == vlds.getSourceMutable().getOperandNumber() && - vlds.getUpdatedBase()) - vlds.getUpdatedBase().setType(newTy); - return; - } - - if (auto vsts = dyn_cast(owner)) { - if (operandNo == vsts.getDestinationMutable().getOperandNumber() && - vsts.getUpdatedBase()) - vsts.getUpdatedBase().setType(newTy); - return; - } - - if (auto vsstb = dyn_cast(owner)) { - if (operandNo == vsstb.getDestinationMutable().getOperandNumber() && - vsstb.getUpdatedBase()) - vsstb.getUpdatedBase().setType(newTy); - } -} - -static void retypeCalleeOperand(ModuleOp module, func::CallOp call, - unsigned operandNo, Type newTy, - MLIRContext *ctx) { - func::FuncOp callee = module.lookupSymbol(call.getCallee()); - if (!callee) - return; - - FunctionType fnTy = callee.getFunctionType(); - if (operandNo >= fnTy.getNumInputs() || fnTy.getInput(operandNo) == newTy) - return; - - SmallVector inputs(fnTy.getInputs().begin(), fnTy.getInputs().end()); - inputs[operandNo] = newTy; - callee.setFunctionType(FunctionType::get(ctx, inputs, fnTy.getResults())); -} - -static void restoreMaterializedTileBufAddrOps( - ModuleOp module, OpBuilder &builder, MLIRContext *ctx, - DenseMap &tileHandles) { - SmallVector flattenedViews; - module.walk([&](memref::ReinterpretCastOp cast) { - if (isFlattenedTileBufAddrView(cast)) - flattenedViews.push_back(cast); - }); - - for (memref::ReinterpretCastOp view : flattenedViews) { - Value tile = lookupMaterializedTileHandle(view.getSource(), tileHandles); - auto tileTy = tile ? dyn_cast(tile.getType()) : TileBufType(); - if (!tileTy) - continue; - - auto resultTy = cast(view.getResult().getType()); - if (resultTy.getElementType() != tileTy.getElementType()) - continue; - - auto tileSpace = - dyn_cast_or_null(tileTy.getMemorySpace()); - if (!tileSpace) - continue; - - auto ptrTy = - PtrType::get(ctx, resultTy.getElementType(), tileSpace); - - SmallVector, 4> callUses; - SmallVector, 4> directUses; - for (OpOperand &use : view.getResult().getUses()) { - auto call = dyn_cast(use.getOwner()); - if (call) { - unsigned operandNo = use.getOperandNumber(); - if (!canRetypeCalleeOperand(module, call, operandNo, ptrTy)) - continue; - callUses.push_back({call, operandNo}); - continue; - } - - Operation *owner = use.getOwner(); - unsigned operandNo = use.getOperandNumber(); - if (!canRetypeTileBufAddrDirectUse(owner, operandNo)) - continue; - directUses.push_back({owner, operandNo}); - } - - if (callUses.empty() && directUses.empty()) - continue; - - builder.setInsertionPoint(view); - auto addr = - builder.create(view.getLoc(), ptrTy, tile); - - for (auto [call, operandNo] : callUses) { - call->setOperand(operandNo, addr.getDst()); - retypeCalleeOperand(module, call, operandNo, ptrTy, ctx); - } - - for (auto [owner, operandNo] : directUses) { - owner->setOperand(operandNo, addr.getDst()); - updateDirectUseResultTypesAfterPtrRetype(owner, operandNo, ptrTy); - } - - if (view.use_empty()) - view.erase(); - } -} - static Value materializeAnchorResult(Operation *anchor, Value anchoredValue, OpBuilder &builder, MLIRContext *ctx, DenseMap &tileHandles, @@ -1189,17 +828,6 @@ struct PTOMaterializeTileHandlesPass mustMaterialize.insert(meta.source); } - if (failed(restorePTODSLSubkernelHelperTileABI( - module, builder, ctx, tileHandles, failedMaterialization))) { - signalPassFailure(); - return; - } - - if (failedMaterialization) { - signalPassFailure(); - return; - } - for (Operation *anchor : anchors) { if (anchor->getNumResults() != 1) continue; @@ -1213,8 +841,6 @@ struct PTOMaterializeTileHandlesPass return; } - restoreMaterializedTileBufAddrOps(module, builder, ctx, tileHandles); - SmallVector, 32> operandsToRewrite; module.walk([&](Operation *op) { if (!shouldMaterializeOperand(op)) diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 4e97ad80ba..a556fb2338 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -52,14 +52,8 @@ using namespace mlir; namespace mlir { namespace pto { -static constexpr llvm::StringLiteral kForceDynamicValidShapeAttrName = - "__pto.force_dynamic_valid_shape"; - namespace { -static void markForceDynamicValidShape(Operation *op, bool force, - MLIRContext *ctx); - static Type convertPTOTypeToMemRef(Type t); static bool hasTileNativeAllocationRoot(func::FuncOp func) { @@ -707,34 +701,6 @@ static Type convertPTOTypeToMemRef(Type t) { return t; } -static void markForceDynamicValidShape(Operation *op, bool force, - MLIRContext *ctx) { - if (force) { - op->setAttr(kForceDynamicValidShapeAttrName, UnitAttr::get(ctx)); - return; - } - op->removeAttr(kForceDynamicValidShapeAttrName); -} - -[[maybe_unused]] static void rewriteFunctionSignature(func::FuncOp func, MLIRContext *ctx) { - Block &entry = func.front(); - auto fnTy = func.getFunctionType(); - - SmallVector newInputs; - for (Type type : fnTy.getInputs()) - newInputs.push_back(convertPTOTypeToMemRef(type)); - - SmallVector newResults; - for (Type type : fnTy.getResults()) - newResults.push_back(convertPTOTypeToMemRef(type)); - - for (unsigned i = 0; i < entry.getNumArguments(); ++i) { - if (entry.getArgument(i).getType() != newInputs[i]) - entry.getArgument(i).setType(newInputs[i]); - } - func.setFunctionType(FunctionType::get(ctx, newInputs, newResults)); -} - static LogicalResult lowerPtrLikeTileBufAddrOps(func::FuncOp func, MLIRContext *ctx) { SmallVector addrOps; @@ -767,73 +733,6 @@ static LogicalResult lowerPtrLikeTileBufAddrOps(func::FuncOp func, return success(); } -static LogicalResult bridgeMemRefOperandsToExternalPtrCallees(ModuleOp mod, - func::FuncOp func, - MLIRContext *ctx) { - SmallVector callOps; - func.walk([&](func::CallOp callOp) { callOps.push_back(callOp); }); - - for (auto callOp : callOps) { - auto callee = mod.lookupSymbol(callOp.getCallee()); - if (!callee || !callee.isExternal()) - continue; - - FunctionType calleeType = callee.getFunctionType(); - if (calleeType.getNumInputs() != callOp.getNumOperands()) { - callOp.emitError("callee signature does not match call operands"); - return failure(); - } - - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(callOp); - for (auto [index, expectedType] : - llvm::enumerate(calleeType.getInputs())) { - Value operand = callOp.getOperand(index); - if (!isa(expectedType) || - !isa(operand.getType())) - continue; - - auto castOp = rewriter.create( - callOp.getLoc(), expectedType, operand); - callOp->setOperand(index, castOp.getResult()); - } - } - return success(); -} - -static LogicalResult bridgeCallOperandsToConvertedCallees(ModuleOp mod, - MLIRContext *ctx) { - SmallVector callOps; - mod.walk([&](func::CallOp callOp) { callOps.push_back(callOp); }); - - for (func::CallOp callOp : callOps) { - auto callee = mod.lookupSymbol(callOp.getCallee()); - if (!callee || callee.isExternal()) - continue; - - FunctionType calleeType = callee.getFunctionType(); - if (calleeType.getNumInputs() != callOp.getNumOperands()) { - callOp.emitError("callee signature does not match call operands"); - return failure(); - } - - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(callOp); - for (auto [index, expectedType] : - llvm::enumerate(calleeType.getInputs())) { - Value operand = callOp.getOperand(index); - if (operand.getType() == expectedType) - continue; - - auto bridge = rewriter.create( - callOp.getLoc(), expectedType, operand); - callOp->setOperand(index, bridge.getResult(0)); - } - } - - return success(); -} - [[maybe_unused]] static LogicalResult lowerTensorViewDimOps(func::FuncOp func, MLIRContext *ctx) { DefaultInlineVector tvDims; func.walk([&](mlir::pto::GetTensorViewDimOp op) { tvDims.push_back(op); }); @@ -931,56 +830,13 @@ struct PTOViewToMemrefPass return; } - // Keep external declarations in the authored ABI. PTOViewToMemref does - // not rewrite func.call operands for external/private helpers, so changing - // only the callee type can leave pointer-like call users mismatched. if (func.isExternal()) continue; - auto fnTy = func.getFunctionType(); bool preserveTileABI = hasMigratedTileNativeOp(func); bool tileNativeMainline = preserveTileABI || hasTileNativeAllocationRoot(func); - // ------------------------------------------------------------------ - // Stage 0.10: Rewrite Function Signature - // ------------------------------------------------------------------ - SmallVector newInputs; - for (Type t : fnTy.getInputs()) { - newInputs.push_back(preserveTileABI && - isa(t) - ? t - : convertPTOTypeToMemRef(t)); - } - - SmallVector newResults; - for (Type t : fnTy.getResults()) { - newResults.push_back(preserveTileABI && - isa(t) - ? t - : convertPTOTypeToMemRef(t)); - } - - func.setFunctionType(FunctionType::get(ctx, newInputs, newResults)); - - Block &entry = func.front(); - - // Update entry block arguments - for (unsigned i = 0; i < entry.getNumArguments(); ++i) { - if (entry.getArgument(i).getType() != newInputs[i]) { - entry.getArgument(i).setType(newInputs[i]); - } - } - - if (failed(bridgeMemRefOperandsToExternalPtrCallees(mod, func, ctx))) { - signalPassFailure(); - return; - } - // ------------------------------------------------------------------ // Stage 0.25: synthesize missing A2/A3 tquant tmp tiles before type lowering. // ------------------------------------------------------------------ @@ -989,41 +845,6 @@ struct PTOViewToMemrefPass return; } - // Stage 0.40 Insert pto.bind_tile for function args that were tile_buf. - // ------------------------------------------------------------------ - // Later materialization and intrinsic folding use BindTileOp as the - // anchor to recover tile metadata after the Stage-0 type rewrite. - { - IRRewriter rewriter(ctx); - // Insert after existing block args, before any existing ops. - rewriter.setInsertionPointToStart(&entry); - for (unsigned i = 0; i < entry.getNumArguments(); ++i) { - Type origTy = fnTy.getInputs()[i]; - auto tbTy = dyn_cast(origTy); - if (!tbTy || newInputs[i] == origTy) - continue; - - auto configAttr = tbTy.getConfigAttr(); - if (!configAttr) configAttr = pto::TileBufConfigAttr::getDefault(ctx); - - Value vRow, vCol; - auto vs = tbTy.getValidShape(); - if (vs.size() == 2) { - if (vs[0] != ShapedType::kDynamic) - vRow = rewriter.create(func.getLoc(), vs[0]); - if (vs[1] != ShapedType::kDynamic) - vCol = rewriter.create(func.getLoc(), vs[1]); - } - - auto bindOp = rewriter.create( - func.getLoc(), newInputs[i], entry.getArgument(i), - vRow ? vRow : Value(), vCol ? vCol : Value(), configAttr); - markForceDynamicValidShape(bindOp, tbTy.hasDynamicValid(), ctx); - - entry.getArgument(i).replaceAllUsesExcept(bindOp.getResult(), bindOp); - } - } - // ------------------------------------------------------------------ // Stage 0.5: keep pto.alloc_tile as a tile-native allocation root. // @@ -3108,11 +2929,6 @@ struct PTOViewToMemrefPass } - if (failed(bridgeCallOperandsToConvertedCallees(mod, ctx))) { - signalPassFailure(); - return; - } - // Debug Output LLVM_DEBUG(llvm::dbgs() << mod.getOperation()); } diff --git a/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto b/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto index c5c62816a6..450a5d4cc8 100644 --- a/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto +++ b/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto @@ -8,7 +8,8 @@ // Guards PTODSL subkernel helper arguments staying in tile_buf ABI when the // helper directly extracts their address views. -// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON // RUN: ptoas --pto-arch=a5 --pto-level=level3 %s -o /dev/null module { @@ -21,29 +22,27 @@ module { pto.section.vector { %0 = pto.tile_buf_addr %arg0 : !pto.tile_buf - -> memref<16x16xf16, #pto.address_space> + -> !pto.ptr %1 = pto.tile_buf_addr %arg1 : !pto.tile_buf - -> memref<16x1xf32, #pto.address_space> + -> !pto.ptr func.call @ptodsl_subkernel_helper_tile_abi_sink(%0, %1, %arg2) - : (memref<16x16xf16, #pto.address_space>, - memref<16x1xf32, #pto.address_space>, f32) -> () + : (!pto.ptr, !pto.ptr, f32) -> () } return } func.func private @ptodsl_subkernel_helper_tile_abi_sink( - memref<16x16xf16, #pto.address_space>, - memref<16x1xf32, #pto.address_space>, - f32) + !pto.ptr, !pto.ptr, f32) } -// IR: IR Dump After PTOMaterializeTileHandles -// IR-LABEL: func.func @ptodsl_subkernel_helper_tile_abi( -// IR-SAME: %{{.*}}: !pto.tile_buf -// IR-SAME: %{{.*}}: !pto.tile_buf -// IR-SAME: %{{.*}}: f32 -// IR-SAME: attributes -// IR: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> memref<16x16xf16, #pto.address_space> -// IR: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> memref<16x1xf32, #pto.address_space> -// IR-NOT: builtin.unrealized_conversion_cast +// NATIVE: IR Dump After PTOViewToMemref +// MATERIALIZED: IR Dump After PTOMaterializeTileHandles +// COMMON-LABEL: func.func @ptodsl_subkernel_helper_tile_abi( +// COMMON-SAME: %{{.*}}: !pto.tile_buf +// COMMON-SAME: %{{.*}}: !pto.tile_buf +// COMMON: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> !pto.ptr +// COMMON: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> !pto.ptr +// COMMON: func.call @ptodsl_subkernel_helper_tile_abi_sink +// COMMON-SAME: (!pto.ptr, !pto.ptr, f32) -> () +// COMMON-NOT: builtin.unrealized_conversion_cast diff --git a/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto b/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto index 21f3d0c14b..de967da6e3 100644 --- a/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto +++ b/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto @@ -7,73 +7,62 @@ // See LICENSE in the root of the software repository for the full text of the License. // Guards late PTODSL helper inlining after tile handle materialization. -// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON // RUN: ptoas --pto-arch=a5 --pto-level=level3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @ptodsl_subkernel_helper_tile_abi_emitc( - %src: memref>, - %scale: f32) { + %src: !pto.ptr, %scale: f32) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index %c0_i64 = arith.constant 0 : i64 pto.section.vector { - %gm = memref.reinterpret_cast %src to offset: [0], sizes: [%c1, %c16, %c16], - strides: [%c16, %c16, %c1] - : memref> - to memref, #pto.address_space> - %mem = pto.pointer_cast(%c0_i64) - : memref<16x16xf32, strided<[16, 1]>, #pto.address_space> - %tile = pto.bind_tile %mem, %c16, %c16 {config = #pto.tile_buf_config, slayout=#pto.slayout, s_fractal_size=512, pad=#pto.pad_value, compact=#pto.compact_mode>} - : memref<16x16xf32, strided<[16, 1]>, #pto.address_space> - -> memref<16x16xf32, strided<[16, 1]>, #pto.address_space> - %part = memref.subview %gm[0, 0, 0] [1, 16, 16] [1, 1, 1] - : memref, #pto.address_space> - to memref<1x16x16xf32, strided<[?, ?, ?], offset: ?>, #pto.address_space> - pto.tload ins(%part : memref<1x16x16xf32, strided<[?, ?, ?], offset: ?>, #pto.address_space>) - outs(%tile : memref<16x16xf32, strided<[16, 1]>, #pto.address_space>) + %gm = pto.make_tensor_view %src, shape = [%c16, %c16], + strides = [%c16, %c1] : !pto.tensor_view<16x16xf32> + %part = pto.partition_view %gm, offsets = [%c0, %c0], + sizes = [%c16, %c16] + : !pto.tensor_view<16x16xf32> + -> !pto.partition_tensor_view<16x16xf32> + %tile = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + pto.tload ins(%part : !pto.partition_tensor_view<16x16xf32>) + outs(%tile : !pto.tile_buf) func.call @ptodsl_helper(%tile, %scale) - : (memref<16x16xf32, strided<[16, 1]>, #pto.address_space>, f32) -> () + : (!pto.tile_buf, f32) -> () } return } func.func private @ptodsl_helper( - %arg0: memref<16x16xf32, strided<[16, 1]>, #pto.address_space>, + %arg0: !pto.tile_buf, %scale: f32) attributes { pto.ptodsl.subkernel_helper = "simd" } { pto.section.vector { - %c16 = arith.constant 16 : index - %tile = pto.bind_tile %arg0, %c16, %c16 {config = #pto.tile_buf_config, slayout=#pto.slayout, s_fractal_size=512, pad=#pto.pad_value, compact=#pto.compact_mode>} - : memref<16x16xf32, strided<[16, 1]>, #pto.address_space> - -> memref<16x16xf32, strided<[16, 1]>, #pto.address_space> - %ptr = memref.reinterpret_cast %tile to offset: [0], sizes: [256], - strides: [1] - : memref<16x16xf32, strided<[16, 1]>, #pto.address_space> - to memref> + %ptr = pto.tile_buf_addr %arg0 + : !pto.tile_buf -> !pto.ptr func.call @ptodsl_helper_sink(%ptr, %scale) - : (memref>, f32) -> () + : (!pto.ptr, f32) -> () } return } func.func private @ptodsl_helper_sink( - memref>, f32) + !pto.ptr, f32) } -// IR: IR Dump After PTOMaterializeTileHandles -// IR-LABEL: func.func @ptodsl_subkernel_helper_tile_abi_emitc -// IR: %[[ALLOC:.*]] = pto.alloc_tile addr = {{.*}} : !pto.tile_buf -// IR: pto.tload -// IR-SAME: outs(%[[ALLOC]] : !pto.tile_buf) -// IR: func.call @ptodsl_helper(%[[ALLOC]], {{.*}}) : (!pto.tile_buf, f32) -> () -// IR-LABEL: func.func private @ptodsl_helper( -// IR-SAME: !pto.tile_buf -// IR: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> !pto.ptr -// IR: func.call @ptodsl_helper_sink(%{{.*}}, {{.*}}) : (!pto.ptr, f32) -> () -// IR: func.func private @ptodsl_helper_sink(!pto.ptr, f32) +// NATIVE: IR Dump After PTOViewToMemref +// MATERIALIZED: IR Dump After PTOMaterializeTileHandles +// COMMON-LABEL: func.func @ptodsl_subkernel_helper_tile_abi_emitc +// COMMON: %[[ALLOC:.*]] = pto.alloc_tile addr = {{.*}} : !pto.tile_buf +// COMMON: func.call @ptodsl_helper(%[[ALLOC]], {{.*}}) : (!pto.tile_buf, f32) -> () +// COMMON-LABEL: func.func private @ptodsl_helper( +// COMMON-SAME: !pto.tile_buf +// COMMON: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> !pto.ptr +// COMMON: func.call @ptodsl_helper_sink(%{{.*}}, {{.*}}) : (!pto.ptr, f32) -> () +// COMMON: func.func private @ptodsl_helper_sink(!pto.ptr, f32) // EMITC-LABEL: extern "C" AICORE void ptodsl_helper_sink( // EMITC-LABEL: AICORE void ptodsl_subkernel_helper_tile_abi_emitc @@ -81,3 +70,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, f32) -> () +// VPTO: func.func private @ptodsl_helper_sink(!pto.ptr, f32) From 2bc85ecc83584c55141ed0c88bd32fbb76b5198f Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 20 Jul 2026 15:46:27 +0800 Subject: [PATCH 83/91] Keep async session scratch tile-native --- .../Transforms/PTOMaterializeTileHandles.cpp | 2 -- lib/PTO/Transforms/PTOValidateVPTOIR.cpp | 8 +++++ lib/PTO/Transforms/PTOViewToMemref.cpp | 2 +- .../pto/build_async_session_tile_native.pto | 36 +++++++++++++++++++ test/lit/vpto/async_session_unsupported.pto | 24 +++++++++++++ 5 files changed, 69 insertions(+), 3 deletions(-) create mode 100644 test/lit/pto/build_async_session_tile_native.pto create mode 100644 test/lit/vpto/async_session_unsupported.pto diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index 387db6bd51..72caf512a9 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -101,8 +101,6 @@ static bool shouldMaterializeOperand(Operation *owner) { return true; if (name == "pto.get_validshape") return true; - if (name == "pto.build_async_session") - return true; if (!name.consume_front("pto.")) return false; return name.starts_with("t"); diff --git a/lib/PTO/Transforms/PTOValidateVPTOIR.cpp b/lib/PTO/Transforms/PTOValidateVPTOIR.cpp index d414e4abcb..f0778bd287 100644 --- a/lib/PTO/Transforms/PTOValidateVPTOIR.cpp +++ b/lib/PTO/Transforms/PTOValidateVPTOIR.cpp @@ -1086,6 +1086,14 @@ class VPTOLegalityValidator { LogicalResult validateEmissionOperationSurface() { WalkResult walkResult = helper.getModule().walk([&](Operation *op) { + if (isa(op)) { + op->emitOpError() + << "is not supported by the VPTO backend; async SDMA session " + "operations currently require the EmitC backend"; + return WalkResult::interrupt(); + } + VPTOBufferAddressFamily family = VPTOLegalityHelper::classifyBufferAddressFamily(op); diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index a556fb2338..1853ea3a71 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -122,7 +122,7 @@ static bool hasMigratedTileNativeOp(func::FuncOp func) { pto::TGemvAccOp, pto::TGemvBiasOp, pto::TGemvMxOp, pto::TGemvMxAccOp, pto::TGemvMxBiasOp, pto::TPrefetchOp, pto::TPushOp, pto::TPopOp, - pto::TFreeOp, pto::TSyncOp>(op)) { + pto::TFreeOp, pto::TSyncOp, pto::BuildAsyncSessionOp>(op)) { found = true; return WalkResult::interrupt(); } diff --git a/test/lit/pto/build_async_session_tile_native.pto b/test/lit/pto/build_async_session_tile_native.pto new file mode 100644 index 0000000000..e4f92f08d2 --- /dev/null +++ b/test/lit/pto/build_async_session_tile_native.pto @@ -0,0 +1,36 @@ +// 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 \ +// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 --pto-level=level3 %s -o - 2>&1 \ +// RUN: | FileCheck %s --check-prefix=EMITC + +module { + func.func @build_async_session_tile_native( + %scratch: !pto.tile_buf, + %workspace: !pto.ptr) { + %session = pto.comm.build_async_session( + %scratch, %workspace : !pto.tile_buf, + !pto.ptr) -> !pto.async_session + return + } +} + +// NATIVE: IR Dump After PTOViewToMemref +// NATIVE-LABEL: func.func @build_async_session_tile_native( +// NATIVE-SAME: %[[SCRATCH:.*]]: !pto.tile_buf +// NATIVE-SAME: %[[WORKSPACE:.*]]: !pto.ptr +// NATIVE: pto.comm.build_async_session(%[[SCRATCH]], %[[WORKSPACE]] +// NATIVE-SAME: !pto.tile_buf, !pto.ptr +// NATIVE-NOT: memref< +// NATIVE-NOT: pto.bind_tile + +// EMITC-LABEL: AICORE void build_async_session_tile_native( +// EMITC: pto::comm::BuildAsyncSession( diff --git a/test/lit/vpto/async_session_unsupported.pto b/test/lit/vpto/async_session_unsupported.pto new file mode 100644 index 0000000000..b8d02db110 --- /dev/null +++ b/test/lit/vpto/async_session_unsupported.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=a5 --pto-level=level3 --pto-backend=vpto \ +// RUN: --emit-vpto %s -o - 2>&1 | FileCheck %s + +module attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @unsupported_async_session( + %scratch: !pto.tile_buf, + %workspace: !pto.ptr) attributes {pto.kernel} { + %session = pto.comm.build_async_session( + %scratch, %workspace : !pto.tile_buf, + !pto.ptr) -> !pto.async_session + return + } +} + +// CHECK: 'pto.comm.build_async_session' op is not supported by the VPTO backend +// CHECK-SAME: async SDMA session operations currently require the EmitC backend From 24e61b527fd9104fb4c7a5cd8f7a0b94ad5ea271 Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 20 Jul 2026 16:26:19 +0800 Subject: [PATCH 84/91] Remove retired memref tile bridge ops --- include/PTO/IR/PTOOps.td | 80 -------- .../Transforms/GraphSyncSolver/SyncSolverIR.h | 2 +- .../Transforms/InsertSync/PTOIRTranslator.h | 2 - include/PTO/Transforms/Passes.td | 12 +- include/PTO/Transforms/SlotAffineAnalysis.h | 8 +- lib/PTO/IR/PTO.cpp | 46 +---- lib/PTO/Transforms/FoldTileBufIntrinsics.cpp | 79 ++------ .../Transforms/GraphSyncSolver/MemInfo.cpp | 66 ------- .../Transforms/GraphSyncSolver/SyncSolver.cpp | 6 +- .../SyncSolverIRTranslator.cpp | 19 +- .../InsertSync/InsertSyncAnalysis.cpp | 8 +- .../Transforms/InsertSync/PTOIRTranslator.cpp | 62 +----- .../Transforms/PTOMaterializeTileHandles.cpp | 2 +- lib/PTO/Transforms/PTOPlanMemory.cpp | 19 +- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 4 - lib/PTO/Transforms/PTOResolveBufferSelect.cpp | 143 +------------- lib/PTO/Transforms/PTOToEmitC.cpp | 184 ------------------ lib/PTO/Transforms/PTOViewToMemref.cpp | 5 - lib/PTO/Transforms/SlotAffineAnalysis.cpp | 4 +- .../Transforms/TileFusion/FusionAnalysis.cpp | 7 - lib/PTO/Transforms/Utils.cpp | 5 - test/lit/pto/declare_tile_tile_native.pto | 1 - .../pto/materialize_tile_handles_emitc.pto | 1 - ...lti_tile_const_preload_dyn_loop_select.pto | 1 - .../multi_tile_const_slot_gss_compiles.pto | 5 +- .../multi_tile_get_const_slot_lowering.pto | 1 - .../pto/multi_tile_get_dyn_slot_lowering.pto | 1 - test/lit/pto/multi_tile_no_loop_unroll.pto | 1 - .../pto/multi_tile_two_buf_slot_binding.pto | 2 +- .../pto/retired_memref_bridge_ops_invalid.pto | 42 ++++ ..._column_slice_preserve_parent_shape_a5.pto | 1 - test/lit/pto/tcvt_low_precision_a5_valid.pto | 1 - ...tload_tprefetch_low_precision_a5_valid.pto | 1 - .../lit/pto/tstore_low_precision_a5_valid.pto | 1 - tools/ptoas/ptoas.cpp | 10 +- 35 files changed, 103 insertions(+), 729 deletions(-) create mode 100644 test/lit/pto/retired_memref_bridge_ops_invalid.pto diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 1e9fbf2ebb..2ace62496b 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -353,10 +353,6 @@ def AllocTileOp : PTO_Op<"alloc_tile", [AttrSizedOperandSegments]> { // may be a constant or any SSA `index` value; PTOAS does not synthesize // `iv mod N` automatically -- the user expression IS the slot selector. // -// `pto.slot_marker` remains as an internal compatibility op for legacy -// memref-based multi-buffer IR. PTOViewToMemref does not create it for -// tile-native `pto.multi_tile_get`. - def AllocMultiTileOp : PTO_Op<"alloc_multi_tile", [AttrSizedOperandSegments]> { let summary = "Allocate an N-slot multi-buffer tile"; let description = [{ @@ -470,34 +466,6 @@ def BindTileOp : PTO_Op<"bind_tile", [ }]; } -def MaterializeTileOp : PTO_Op<"materialize_tile", [ - Pure, - AttrSizedOperandSegments - ]> { - let summary = "Materializes a tile buffer handle from a planned memref"; - let description = [{ - Bridges the memref-based memory-planning/sync IR back into tile-buffer - handle IR immediately before EmitC. The source remains the planned memref - address, while the result carries the tile type, valid dims, and config - expected by tile operations and C++ tile codegen. - }]; - - let arguments = (ins - AnyMemRef:$source, - Optional:$valid_row, - Optional:$valid_col, - TileBufConfigAttr:$config - ); - - let results = (outs TileBufType:$result); - - let hasVerifier = 1; - - let assemblyFormat = [{ - $source (`,` $valid_row^ `,` $valid_col)? attr-dict `:` qualified(type($source)) `->` qualified(type($result)) - }]; -} - def SubViewOp : PTO_Op<"subview", [ Pure, ViewLikeOpInterface, @@ -1549,40 +1517,6 @@ def PointerCastOp : PTO_Op<"pointer_cast", [AttrSizedOperandSegments, Pure]> { }]; } -def SlotMarkerOp : PTO_Op<"slot_marker", [ - Pure, - ViewLikeOpInterface, - AllTypesMatch<["source", "result"]> - ]> { - let summary = "Tag a memref view as referring to one slot of a multi_tile_buf"; - let description = [{ - Internal compatibility op for legacy memref-based multi-buffer IR. It - carries a slot SSA index so sync analysis (InsertSync / GraphSyncSolver) - and the buffer-select lowering pass can identify which physical slot a - memref reference touches. - - The op is metadata-only (no data movement, no extra address arithmetic); - its result memref aliases the source memref byte-for-byte. Frontends do - not produce this op directly -- use `pto.multi_tile_get` instead. - }]; - - let arguments = (ins - AnyMemRef:$source, - Index:$slot - ); - - let results = (outs AnyMemRef:$result); - - let assemblyFormat = [{ - $source `[` $slot `]` attr-dict - `:` qualified(type($source)) `->` qualified(type($result)) - }]; - - let extraClassDeclaration = [{ - ::mlir::Value getViewSource() { return getSource(); } - }]; -} - // ============================================================================= // System/Runtime Query Ops // ============================================================================= @@ -2632,20 +2566,6 @@ def LocalArraySetOp : PTO_Op<"local_array_set"> { }]; } -def DeclareTileMemRefOp : PTO_Op<"declare_tile_memref"> { - let summary = "Legacy memref placeholder for a runtime-bound tile"; - let description = [{ - Compatibility op for legacy memref-based IR. New PTOViewToMemref pipelines - preserve `pto.declare_tile` directly instead of producing this op. - }]; - - let results = (outs AnyMemRef:$result); - - let assemblyFormat = [{ - attr-dict `->` qualified(type($result)) - }]; -} - def TPopOp : PTO_TOp<"tpop", [ PTO_DpsInitOpInterface, OpPipeInterface, diff --git a/include/PTO/Transforms/GraphSyncSolver/SyncSolverIR.h b/include/PTO/Transforms/GraphSyncSolver/SyncSolverIR.h index 44a3a82aff..38ee5050e0 100644 --- a/include/PTO/Transforms/GraphSyncSolver/SyncSolverIR.h +++ b/include/PTO/Transforms/GraphSyncSolver/SyncSolverIR.h @@ -451,7 +451,7 @@ class SetWaitOp : public SyncOp { bool checkLastIter{false}; // Slot SSA at this access site. Populated when the sync corresponds to a // multi-buffer back-edge dep produced by `pto.multi_tile_get` / lowered - // `pto.slot_marker`. When non-null and `eventIds.size() > 1`, codegen + // `pto.multi_tile_get`. When non-null and `eventIds.size() > 1`, codegen // emits `pto.set_flag_dyn` / `pto.wait_flag_dyn` with a runtime event id // selected by `slotSSAExpr % eventIds.size()` instead of fanning out // into N static `set_flag` / `wait_flag` pairs per iteration. diff --git a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h index 0012f79f29..f53af0ce67 100644 --- a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h +++ b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h @@ -73,7 +73,6 @@ class PTOIRTranslator { LogicalResult UpdateAllocMultiTileOpMemInfo(pto::AllocMultiTileOp op); LogicalResult UpdateDeclareTileOpMemInfo(pto::DeclareTileOp op); LogicalResult UpdateDeclareGlobalOpMemInfo(pto::DeclareGlobalOp op); - LogicalResult UpdateDeclareTileMemRefOpMemInfo(pto::DeclareTileMemRefOp op); LogicalResult UpdatePointerCastOpMemInfo(pto::PointerCastOp op); LogicalResult UpdateMemrefAllocOpMemInfo(memref::AllocOp op); @@ -82,7 +81,6 @@ class PTOIRTranslator { void UpdateConservativeAliasBufferInfo(Value result, Value source); void UpdateMemrefSubViewAliasBufferInfo(memref::SubViewOp op); void UpdateTileSubViewAliasBufferInfo(pto::SubViewOp op); - void UpdateSlotMarkerAliasBufferInfo(pto::SlotMarkerOp op); void UpdateMultiTileGetAliasBufferInfo(pto::MultiTileGetOp op); void UpdateSlotSelectedAliasBufferInfo(Value result, Value source, Value slot); diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 210f9d8160..8eb5755b04 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -723,11 +723,11 @@ def PTOValidateIntToPtrUses : Pass<"pto-validate-inttoptr-uses", "func::FuncOp"> } def PTOMaterializeTileHandles : Pass<"pto-materialize-tile-handles", "ModuleOp"> { - let summary = "Materialize tile_buf handles from planned memrefs before EmitC"; + let summary = "Restore tile_buf handles for remaining legacy memref values"; let description = [{ - Re-wraps local memref values used by tile operations into explicit - pto.materialize_tile ops. This keeps PlanMemory and InjectSync in memref IR - while making the tile handle boundary visible to EmitC lowering. + Re-wraps any remaining local memref values used by tile operations as + addressed `pto.alloc_tile` handles before backend lowering. Tile-native + operations already bypass this compatibility path. }]; let constructor = "mlir::pto::createPTOMaterializeTileHandlesPass()"; @@ -749,10 +749,6 @@ def PTOResolveBufferSelect : Pass<"pto-resolve-buffer-select", "ModuleOp"> { `pto.alloc_tile` carrying the selected address; constant slots select one address directly, while dynamic slots use an N-way `arith.select` chain. - The pass also retains compatibility with legacy memref-based - `pto.slot_marker` IR, where the selection is materialized as a - single-address `pto.pointer_cast`. - This pass does not rewrite `%k` into `iv mod N`; the frontend expression remains the slot selector. }]; diff --git a/include/PTO/Transforms/SlotAffineAnalysis.h b/include/PTO/Transforms/SlotAffineAnalysis.h index 4e265e77aa..303e04ade6 100644 --- a/include/PTO/Transforms/SlotAffineAnalysis.h +++ b/include/PTO/Transforms/SlotAffineAnalysis.h @@ -9,7 +9,7 @@ //===- SlotAffineAnalysis.h - Multi-buffer slot affine compare --*- C++ -*-===// // // Small affine helper used by the multi-buffer sync path. Both InsertSync -// and GraphSyncSolver consume it to decide, for two `pto.slot_marker` +// and GraphSyncSolver consume it to decide, for two `pto.multi_tile_get` // slot-index SSA expressions, whether they are provably equal modulo N, // provably disjoint modulo N, or indeterminate. The result lets sync // shrink event-id count or skip same-iter forward syncs entirely when @@ -37,9 +37,9 @@ enum class SlotRelation { }; /// Walk back through metadata-only ops (`pto.bind_tile`) to the nearest -/// `pto.slot_marker` and return its slot SSA value. Returns a null Value -/// if the chain does not pass through a slot_marker. -mlir::Value findSlotMarkerExpr(mlir::Value v); +/// `pto.multi_tile_get` and return its slot SSA value. Returns a null Value +/// if the chain does not pass through a multi_tile_get. +mlir::Value findMultiTileSlotExpr(mlir::Value v); /// Compare two slot SSA expressions modulo `N`. The analysis is /// intentionally narrow: it accepts the forms commonly produced by diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 30000d78e4..6cbbe7e952 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3125,49 +3125,6 @@ LogicalResult PointerCastOp::verify() { return success(); } -LogicalResult MaterializeTileOp::verify() { - auto sourceTy = cast(getSource().getType()); - auto resultTy = cast(getResult().getType()); - - if (sourceTy.getRank() != 2) - return emitOpError("source memref must be rank-2 to materialize a tile handle"); - if (resultTy.getRank() != 2) - return emitOpError("result tile_buf must be rank-2"); - if (failed(verifyTileBufLayoutConstraints(*this, resultTy, "result"))) - return failure(); - - auto viewSemantics = (*this)->getAttrOfType("pto.view_semantics"); - bool isSubview = viewSemantics && viewSemantics.getValue() == "subview"; - if (!isSubview && sourceTy.getShape() != resultTy.getShape()) - return emitOpError() << "source/result shape mismatch: source=" - << sourceTy << " result=" << resultTy; - - if (sourceTy.getElementType() != resultTy.getElementType()) - return emitOpError() << "source/result element type mismatch: source=" - << sourceTy.getElementType() - << " result=" << resultTy.getElementType(); - - if (sourceTy.getMemorySpace() != resultTy.getMemorySpace()) - return emitOpError() << "source/result memory space mismatch"; - - if (getConfig() != resultTy.getConfigAttr()) - return emitOpError("config attribute must match the result tile_buf config"); - - auto shape = resultTy.getShape(); - auto validShape = resultTy.getValidShape(); - if (validShape.size() != 2) - return emitOpError("result tile_buf must have rank-2 validShape"); - for (unsigned i = 0; i < 2; ++i) { - if (shape[i] != ShapedType::kDynamic && - validShape[i] != ShapedType::kDynamic && validShape[i] > shape[i]) { - return emitOpError() << "valid_shape[" << i << "] must be <= shape[" - << i << "]"; - } - } - - return success(); -} - LogicalResult TAssignOp::verify() { if (getTile().getType() != getResult().getType()) { return emitOpError("result type must match tile operand type"); @@ -11170,8 +11127,7 @@ static bool isLocallyBoundTileSource(Value value) { if (!value || isa(value)) return false; - if (isa( + if (isa( value.getDefiningOp())) return true; diff --git a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp index 4eac1b7fca..258b9b1fe9 100644 --- a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp +++ b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp @@ -22,9 +22,8 @@ // - pto.get_tensor_view_stride → extract dimension stride // // This pass resolves them against the concrete values at the call site. -// For tile_buf intrinsics, the active VPTO path folds against materialized tile -// handles produced by the shared tile-handle bridge (`pto.alloc_tile` or -// `pto.materialize_tile`). +// For tile_buf intrinsics, the active VPTO path folds against addressed +// `pto.alloc_tile` handles produced by the shared tile-handle bridge. // For tensor_view intrinsics, the pass traces through the full // unrealized_conversion_cast → memref.subview → memref.reinterpret_cast // chain to fold directly to constants or SSA operands from the @@ -151,7 +150,6 @@ static void eraseDeadViewChains(func::FuncOp func) { } struct TileHandleInfo { - Value sourceMemref; Value addr; Value validRow; Value validCol; @@ -237,16 +235,10 @@ static std::optional resolveTileHandle(Value tileBuf, "FoldTileBufIntrinsics: pto.alloc_tile must produce !pto.tile_buf"); return std::nullopt; } - return TileHandleInfo{Value(), alloc.getAddr(), alloc.getValidRow(), + return TileHandleInfo{alloc.getAddr(), alloc.getValidRow(), alloc.getValidCol(), tileTy.getConfigAttr()}; } - if (auto materialize = tileBuf.getDefiningOp()) { - return TileHandleInfo{materialize.getSource(), Value(), - materialize.getValidRow(), materialize.getValidCol(), - materialize.getConfig()}; - } - if (auto reshape = tileBuf.getDefiningOp()) { auto sourceInfo = resolveTileHandle(reshape.getSrc(), user); if (!sourceInfo) @@ -260,8 +252,7 @@ static std::optional resolveTileHandle(Value tileBuf, } auto [validRow, validCol] = findSetValidShapeOverride(tileBuf); - return TileHandleInfo{sourceInfo->sourceMemref, - sourceInfo->addr, + return TileHandleInfo{sourceInfo->addr, validRow ? validRow : sourceInfo->validRow, validCol ? validCol : sourceInfo->validCol, tileTy.getConfigAttr()}; @@ -269,7 +260,7 @@ static std::optional resolveTileHandle(Value tileBuf, user->emitError("FoldTileBufIntrinsics: expected tile_buf to be defined by " "the active materialized tile-handle bridge " - "(pto.alloc_tile, pto.materialize_tile, or pto.treshape, " + "(pto.alloc_tile or pto.treshape, " "or a pto.fusion_region result that yields one of them)"); return std::nullopt; } @@ -722,8 +713,7 @@ struct FoldTileBufIntrinsicsPass // Fold pto.tile_buf_addr by recovering the active materialized tile // handle contract: - // - pto.materialize_tile → use the source memref directly - // - pto.alloc_tile → rebuild a memref from the explicit addr + // - pto.alloc_tile → rebuild a memref from the explicit addr // When the requested result type is already !pto.ptr<...>, cast from the // recovered memref instead of leaving tile_buf_addr in the IR. for (auto addrOp : addrOps) { @@ -747,23 +737,6 @@ struct FoldTileBufIntrinsicsPass if (auto resultMemrefType = dyn_cast(addrOp.getDst().getType())) { - if (handleInfo->sourceMemref) { - Value srcMemref = handleInfo->sourceMemref; - if (!isa(srcMemref.getType())) { - addrOp.emitError("FoldTileBufIntrinsics: pto.materialize_tile " - "source is not a memref"); - return signalPassFailure(); - } - - // The declared tile_buf_addr result type may differ from the actual - // materialized source layout (e.g. plain shape vs. strided layout). - if (srcMemref.getType() != resultMemrefType) - addrOp.getDst().setType(cast(srcMemref.getType())); - addrOp.getDst().replaceAllUsesWith(srcMemref); - addrOp.erase(); - continue; - } - if (!handleInfo->addr) { addrOp.emitError("FoldTileBufIntrinsics: pto.alloc_tile used by " "tile_buf_addr must carry an addr operand on the " @@ -789,31 +762,21 @@ struct FoldTileBufIntrinsicsPass return signalPassFailure(); } - Value memrefValue; - if (handleInfo->sourceMemref) { - memrefValue = handleInfo->sourceMemref; - if (!isa(memrefValue.getType())) { - addrOp.emitError("FoldTileBufIntrinsics: pto.materialize_tile " - "source is not a memref"); - return signalPassFailure(); - } - } else { - if (!handleInfo->addr) { - addrOp.emitError("FoldTileBufIntrinsics: pto.alloc_tile used by " - "tile_buf_addr must carry an addr operand on the " - "VPTO path"); - return signalPassFailure(); - } - - builder.setInsertionPoint(addrOp); - auto canonicalMemrefType = getCanonicalMemRefTypeForTileBuf(tileTy); - memrefValue = builder.create( - addrOp.getLoc(), canonicalMemrefType, ValueRange{handleInfo->addr}, - handleInfo->validRow ? handleInfo->validRow : Value(), - handleInfo->validCol ? handleInfo->validCol : Value(), - handleInfo->config); + if (!handleInfo->addr) { + addrOp.emitError("FoldTileBufIntrinsics: pto.alloc_tile used by " + "tile_buf_addr must carry an addr operand on the " + "VPTO path"); + return signalPassFailure(); } + builder.setInsertionPoint(addrOp); + auto canonicalMemrefType = getCanonicalMemRefTypeForTileBuf(tileTy); + Value memrefValue = builder.create( + addrOp.getLoc(), canonicalMemrefType, ValueRange{handleInfo->addr}, + handleInfo->validRow ? handleInfo->validRow : Value(), + handleInfo->validCol ? handleInfo->validCol : Value(), + handleInfo->config); + builder.setInsertionPoint(addrOp); Value replacement = builder.create(addrOp.getLoc(), resultPtrType, @@ -1065,7 +1028,7 @@ struct FoldTileBufIntrinsicsPass op.erase(); // DCE tile-handle view / alloc ops left behind after valid-shape - // folding (treshape / materialize_tile / alloc_tile / bridging casts). + // folding (treshape / alloc_tile / bridging casts). bool tileDceChanged = true; while (tileDceChanged) { tileDceChanged = false; @@ -1073,7 +1036,7 @@ struct FoldTileBufIntrinsicsPass func.walk([&](Operation *op) { if (!op->use_empty()) return; - if (isa(op)) + if (isa(op)) deadTileOps.push_back(op); else if (auto castOp = dyn_cast(op)) { if (castOp.getNumOperands() == 1 && diff --git a/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp b/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp index b0adfb765d..c4d9765e8c 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp @@ -135,77 +135,11 @@ static MemInfo getMemInfoForMultiTileGet(pto::MultiTileGetOp get) { return MemInfo(get.getResult(), info); } -// Walk back through metadata-only view ops (`pto.bind_tile`) to the -// nearest `pto.pointer_cast`. Used to anchor slot_marker MemInfo on its -// underlying multi-address alloc cast. -static pto::PointerCastOp findUnderlyingPointerCast(Value v) { - int hops = 0; - while (v && hops++ < 32) { - Operation *op = v.getDefiningOp(); - if (!op) - return {}; - if (auto pc = llvm::dyn_cast(op)) - return pc; - if (auto bind = llvm::dyn_cast(op)) { - v = bind.getSource(); - continue; - } - return {}; - } - return {}; -} - -// Build a MemInfo for a `pto.slot_marker` use. For a constant slot K the -// MemInfo carries just slot K's physical address so two const-slot -// accesses on different slots come back as non-conflicting via the -// existing `PointerLikeInfo::checkConflict` byte-range overlap. For a -// dynamic slot the MemInfo carries all N physical addresses; downstream -// `checkMultiBufferEventIdInfo` then deduces N event ids using the -// `(i % N) == (j % N)` slot-skipping rule, which is exactly the -// multi-buffer prefetch pattern. -// -// Note on `parentLoop`: `getPointerLikeInfo` records the parent loop of -// the cast op, which is typically outside the multi-buffer scf.for (the -// alloc/cast lives at function scope). The multi-buffer geometry, though, -// is keyed by the loop that *uses* the slot. We override `parentLoop` -// with the slot_marker's enclosing LoopLikeOpInterface so -// `getMultiBufferLoop` finds the right anchor. -static MemInfo getMemInfoForSlotMarker(pto::SlotMarkerOp slotMarker) { - pto::PointerCastOp castOp = findUnderlyingPointerCast(slotMarker.getSource()); - if (!castOp) { - return MemInfo(slotMarker.getResult(), - isWorkSpaceFuncArgument(slotMarker.getResult())); - } - - PointerLikeInfo info = getPointerLikeInfo(castOp); - - IntegerAttr constSlotAttr; - if (matchPattern(slotMarker.getSlot(), m_Constant(&constSlotAttr)) && - info.addresses.size() > 1) { - int64_t slotIdx = constSlotAttr.getValue().getSExtValue(); - if (slotIdx >= 0 && slotIdx < static_cast(info.addresses.size())) { - int64_t picked = info.addresses[static_cast(slotIdx)]; - info.addresses.clear(); - info.addresses.push_back(picked); - } - } - - if (auto useLoop = - slotMarker->template getParentOfType()) { - info.parentLoop = useLoop; - } - - return MemInfo(slotMarker.getResult(), info); -} - MemInfo getMemInfo(Value val) { if (auto *defOp = val.getDefiningOp()) { if (auto pointerCastOp = llvm::dyn_cast(defOp)) { return MemInfo(val, getPointerLikeInfo(pointerCastOp)); } - if (auto slotMarker = llvm::dyn_cast(defOp)) { - return getMemInfoForSlotMarker(slotMarker); - } if (auto allocMulti = llvm::dyn_cast(defOp)) return MemInfo(val, getPointerLikeInfo(allocMulti)); if (auto multiGet = llvm::dyn_cast(defOp)) diff --git a/lib/PTO/Transforms/GraphSyncSolver/SyncSolver.cpp b/lib/PTO/Transforms/GraphSyncSolver/SyncSolver.cpp index ed717e5875..d9c2e17b05 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/SyncSolver.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/SyncSolver.cpp @@ -48,8 +48,8 @@ mlir::pto::SlotRelation compareMemInfoSlotSSA(const MemInfo &memInfo1, size_t n = std::max(memInfo1.getSz(), memInfo2.getSz()); if (n < 2 || n > std::numeric_limits::max()) return mlir::pto::SlotRelation::kUnknown; - Value slot1 = mlir::pto::findSlotMarkerExpr(memInfo1.value); - Value slot2 = mlir::pto::findSlotMarkerExpr(memInfo2.value); + Value slot1 = mlir::pto::findMultiTileSlotExpr(memInfo1.value); + Value slot2 = mlir::pto::findMultiTileSlotExpr(memInfo2.value); if (!slot1 || !slot2) return mlir::pto::SlotRelation::kUnknown; return mlir::pto::compareSlotSSA(slot1, slot2, static_cast(n)); @@ -507,7 +507,7 @@ Solver::getMultiBufferEventIdInfo(Occurrence *occ1, Occurrence *occ2, Value slotOp1, slotOp2; bool slotOp1Ambiguous = false, slotOp2Ambiguous = false; auto collectSlot = [](const MemInfo &mi, Value &slot, bool &ambiguous) { - Value s = mlir::pto::findSlotMarkerExpr(mi.value); + Value s = mlir::pto::findMultiTileSlotExpr(mi.value); if (!s) return; if (!slot) diff --git a/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp b/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp index 9e9758e541..97cb286297 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp @@ -76,12 +76,9 @@ llvm::SmallVector IRTranslator::tracebackMemValsStep(Value val) { out.push_back(whileOp.getYieldedValues()[resultNo]); } - // Stop the walk at `pto.slot_marker` so the multi-buffer slot index is - // preserved for `getMemInfo`. Without this special case, the generic - // `getOperationAliasInfo` path below would treat slot_marker as a - // transparent view and let the trace fall through to the underlying - // multi-address `pto.pointer_cast`, dropping the slot. - if (isa(defOp)) { + // Stop at `pto.multi_tile_get` so getMemInfo preserves the selected slot + // instead of tracing through the metadata-only view. + if (isa(defOp)) { return out; } @@ -126,14 +123,10 @@ llvm::SmallVector IRTranslator::tracebackMemVals(Value val) { if (!result) continue; Operation *defOp = result.getDefiningOp(); - // `pto.slot_marker` is a multi-buffer slot tag and stops traversal so - // `getMemInfo` can extract slot-narrowed addresses below. Without this - // stop, `getOperationAliasInfo` would let the walk slip past slot_marker - // and reach the underlying multi-address `pto.pointer_cast`, dropping - // the slot index. + // `pto.multi_tile_get` stops traversal so getMemInfo can extract the + // slot-narrowed physical addresses. if (isa(defOp)) { + tensor::EmptyOp, memref::AllocOp, pto::MultiTileGetOp>(defOp)) { leaves.insert(result); continue; } diff --git a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp index 2fd3a44526..723dbea214 100644 --- a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp +++ b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp @@ -454,8 +454,8 @@ static bool isForwardDepDroppableBySlotAffine(const BaseMemInfo *a, size_t n = std::max(aN, bN); if (n < 2) return false; - Value slotA = findSlotMarkerExpr(a->baseBuffer); - Value slotB = findSlotMarkerExpr(b->baseBuffer); + Value slotA = findMultiTileSlotExpr(a->baseBuffer); + Value slotB = findMultiTileSlotExpr(b->baseBuffer); if (!slotA || !slotB) return false; return compareSlotSSA(slotA, slotB, static_cast(n)) == @@ -629,9 +629,9 @@ void InsertSyncAnalysis::InsertSyncOperation( Value consumerSlot; for (auto &pair : depBaseMemInfosVec) { if (pair.second && pair.second->baseBuffer) - producerSlot = findSlotMarkerExpr(pair.second->baseBuffer); + producerSlot = findMultiTileSlotExpr(pair.second->baseBuffer); if (pair.first && pair.first->baseBuffer) - consumerSlot = findSlotMarkerExpr(pair.first->baseBuffer); + consumerSlot = findMultiTileSlotExpr(pair.first->baseBuffer); if (producerSlot && consumerSlot) break; } diff --git a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp index dd8355fccc..fbbfcab29f 100644 --- a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp +++ b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp @@ -379,11 +379,6 @@ void PTOIRTranslator::RecursionIR(Region *region) { if (failed(UpdateDeclareTileOpMemInfo(declareOp))) return WalkResult::interrupt(); } - else if (auto declareOp = dyn_cast(op)) { - if (failed(UpdateDeclareTileMemRefOpMemInfo(declareOp))) { - return WalkResult::interrupt(); - } - } else if (auto declareGlobalOp = dyn_cast(op)) { if (failed(UpdateDeclareGlobalOpMemInfo(declareGlobalOp))) { return WalkResult::interrupt(); @@ -409,9 +404,6 @@ void PTOIRTranslator::RecursionIR(Region *region) { else if (auto memrefSubView = dyn_cast(op)) { UpdateMemrefSubViewAliasBufferInfo(memrefSubView); } - else if (auto slotMarker = dyn_cast(op)) { - UpdateSlotMarkerAliasBufferInfo(slotMarker); - } else if (auto multiGet = dyn_cast(op)) { UpdateMultiTileGetAliasBufferInfo(multiGet); } @@ -633,13 +625,8 @@ LogicalResult PTOIRTranslator::UpdatePointerCastOpMemInfo(pto::PointerCastOp op) return success(); } - // Multi-address (multi-buffer) cast. Use the cast result as `rootBuffer` - // so every downstream `pto.slot_marker` from the same alloc shares one - // root, and populate `baseAddresses` with each slot's physical offset - // (extracted from the constant i64 operands emitted by - // AllocToPointerCast). `pto.slot_marker` then narrows or keeps these - // offsets according to its slot SSA; MemAlias's existing - // `isBufferAddressRangeOverlap` does the per-slot disambiguation. + // Multi-address cast. Use the cast result as `rootBuffer` and retain every + // constant physical offset for conservative range-overlap analysis. SmallVector slotOffsets; slotOffsets.reserve(op.getAddrs().size()); for (Value a : op.getAddrs()) { @@ -695,46 +682,6 @@ PTOIRTranslator::UpdateDeclareTileOpMemInfo(pto::DeclareTileOp op) { return success(); } -LogicalResult -PTOIRTranslator::UpdateDeclareTileMemRefOpMemInfo(pto::DeclareTileMemRefOp op) { - Value res = op.getResult(); - auto memRefType = dyn_cast(res.getType()); - if (!memRefType) - return failure(); - - uint64_t sizeInBytes = 0; - if (memRefType.hasStaticShape()) { - int64_t elemSize = static_cast( - pto::getPTOStorageElemByteSize(memRefType.getElementType())); - if (elemSize == 0) - return failure(); - - int64_t numElements = 1; - for (auto dim : memRefType.getShape()) - numElements *= dim; - sizeInBytes = numElements * elemSize; - } - - pto::AddressSpace space = pto::AddressSpace::MAT; - if (auto attr = memRefType.getMemorySpace()) { - if (auto ptoAttr = dyn_cast(attr)) - space = ptoAttr.getAddressSpace(); - } - - // declare_tile_memref is only a symbolic placeholder. Use its SSA result as - // both base/root so later bind_tile aliases and tpop consumers can be - // connected by InsertSync without inventing a fake allocation. - auto newMemInfo = std::make_unique( - res, - res, - space, - SmallVector{0}, - sizeInBytes); - - buffer2MemInfoMap_[res].emplace_back(newMemInfo->clone()); - return success(); -} - LogicalResult PTOIRTranslator::UpdateDeclareGlobalOpMemInfo(pto::DeclareGlobalOp op) { Value res = op.getEntry(); @@ -1083,11 +1030,6 @@ void PTOIRTranslator::UpdateConservativeAliasBufferInfo(Value result, resultMemInfoVec.emplace_back(parentInfo->clone(result)); } -void PTOIRTranslator::UpdateSlotMarkerAliasBufferInfo(pto::SlotMarkerOp op) { - UpdateSlotSelectedAliasBufferInfo(op.getResult(), op.getSource(), - op.getSlot()); -} - void PTOIRTranslator::UpdateMultiTileGetAliasBufferInfo( pto::MultiTileGetOp op) { UpdateSlotSelectedAliasBufferInfo(op.getResult(), op.getSource(), diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index 72caf512a9..1a91787481 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -93,7 +93,7 @@ static bool isLocalTileMemRef(Type type) { } static bool shouldMaterializeOperand(Operation *owner) { - if (isa(owner)) + if (isa(owner)) return false; StringRef name = owner->getName().getStringRef(); diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index cfd58e3481..8ee98ada83 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -464,7 +464,7 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { if (mayAliasOp.has_value()) { auto aliasPair = mayAliasOp.value(); UpdateBufferAlias(aliasPair.first, aliasPair.second); - } else if (isa(op)) { + } else if (isa(op)) { // Runtime-bound tile handles do not allocate static local storage. return WalkResult::advance(); } else if (auto bindOp = dyn_cast(op)) { @@ -472,13 +472,6 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { // of the result as a use of the source in liveness analysis. UpdateBufferAlias(bindOp.getResult(), bindOp.getSource()); return WalkResult::advance(); - } else if (auto slotOp = dyn_cast(op)) { - // SlotMarker is metadata-only: it tags which physical slot of a - // multi-buffer alloc this view refers to. From the planner's point of - // view its result aliases the source; the slot index travels with the - // op and is consumed later by PTOResolveBufferSelect / sync. - UpdateBufferAlias(slotOp.getResult(), slotOp.getSource()); - return WalkResult::advance(); } else if (isLocalMemPlan() && dyn_cast(op)) { auto allocTileOp = cast(op); if (allocTileOp.getAddr()) { @@ -808,15 +801,7 @@ bool MemLivenessAnalysis::isSkippableOp(Operation *op) const { // Call-like ops are still modeled explicitly. Only pure terminators and // dim queries are skipped here. // - // `pto.slot_marker` is a metadata-only view added by PTOViewToMemref to - // thread multi-buffer slot selection through the memref layer. Until - // PlanMemory acquires first-class multi-buffer support (the design's - // §5.2 work), treat it as a passthrough so the rest of the pipeline can - // still be exercised. The N-way physical fan-out lives on the - // `pto.multi_buffer` attr of the underlying `memref.alloc` and is a - // follow-up. - return isa(op); + return isa(op); } LogicalResult diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 2fa4d9882d..aaaa584b2b 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -686,10 +686,6 @@ struct PlannerAnalysis { if (auto bind = dyn_cast(op)) { setRoots(bind.getResult(), getRoots(bind.getSource())); propagateSplitTpopDerived(bind.getResult(), ValueRange{bind.getSource()}); - } else if (auto slotMarker = dyn_cast(op)) { - setRoots(slotMarker.getResult(), getRoots(slotMarker.getSource())); - propagateSplitTpopDerived(slotMarker.getResult(), - ValueRange{slotMarker.getSource()}); } else if (auto multiGet = dyn_cast(op)) { setRoots(multiGet.getResult(), getRoots(multiGet.getSource())); propagateSplitTpopDerived(multiGet.getResult(), diff --git a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp index d70e2d3e6a..9b0268f0d7 100644 --- a/lib/PTO/Transforms/PTOResolveBufferSelect.cpp +++ b/lib/PTO/Transforms/PTOResolveBufferSelect.cpp @@ -10,21 +10,10 @@ // // Lowering for multi-buffer slot selection. // -// Consumes `pto.slot_marker %src[%k] : memref<...>` ops written by -// PTOViewToMemref while lowering `pto.multi_tile_get`. By the time this pass -// runs, PTOPlanMemory has already converted the underlying `memref.alloc` to -// a multi-address `pto.pointer_cast(addr0, ..., addrN-1)`. This pass picks -// the right per-slot address(es) for each slot_marker use: -// -// * Constant slot k: emit a single-address `pto.pointer_cast(addrK)` at -// the use site and replace the slot_marker. -// * Dynamic slot %k: emit N single-address per-slot pointer_casts and -// pick one via an N-way `arith.select` chain. The user's SSA selects -// the slot -- this pass does NOT synthesize `iv mod N`. -// -// The original multi-address `pto.pointer_cast` is left in IR as the -// "alloc anchor" so future sync extensions can still see the multi-buffer -// geometry (e.g. for `set_flag_dyn` / `wait_flag_dyn` derivation). +// Consumes tile-native `pto.multi_tile_get` operations after memory planning. +// Constant slots become addressed `pto.alloc_tile` handles directly; dynamic +// slots select an address through an N-way `arith.select` chain. The user SSA +// remains the slot selector; this pass does not synthesize `iv mod N`. // //===----------------------------------------------------------------------===// @@ -295,56 +284,6 @@ static LogicalResult resolveTileNativeMultiGets(ModuleOp module, return success(); } -/// Walk back through pure metadata ops (`pto.bind_tile`, `pto.slot_marker`) -/// to find the root multi-address `pto.pointer_cast` that this view ties -/// to. Returns nullptr if the chain does not terminate on a -/// `pto.pointer_cast` -- in which case this slot_marker is not a multi- -/// buffer reference and should not be touched. -static pto::PointerCastOp lookupRootPointerCast(Value v) { - while (Operation *def = v.getDefiningOp()) { - if (auto pc = dyn_cast(def)) - return pc; - if (auto bind = dyn_cast(def)) { - v = bind.getSource(); - continue; - } - if (auto sm = dyn_cast(def)) { - // Nested slot_marker should not happen (verifier disallows nested - // multi_tile_get), but follow the chain defensively. - v = sm.getSource(); - continue; - } - return {}; - } - return {}; -} - -/// Lookup tile-buf config from the existing PointerCastOp's optional attr. -/// Returns nullptr if not set. -static Attribute getCastConfigAttr(pto::PointerCastOp root) { - auto cfg = root.getConfig(); - if (cfg.has_value()) - return *cfg; - return Attribute(); -} - -/// Create a fresh single-address pointer_cast that aliases slot `slotIdx` -/// of `root`. The result type matches `targetType`. `vRow` / `vCol` and -/// `config` are forwarded from the root. -static Value emitSlotPointerCast(IRRewriter &rewriter, Location loc, - pto::PointerCastOp root, uint32_t slotIdx, - Type targetType) { - auto rootAddrs = root.getAddrs(); - assert(slotIdx < rootAddrs.size() && "slot index out of range"); - Value vRow = root.getValidRow(); - Value vCol = root.getValidCol(); - Attribute cfg = getCastConfigAttr(root); - auto pc = rewriter.create( - loc, targetType, ValueRange{rootAddrs[slotIdx]}, - vRow ? vRow : Value(), vCol ? vCol : Value(), cfg); - return pc.getResult(); -} - struct PTOResolveBufferSelectPass : public mlir::pto::impl::PTOResolveBufferSelectBase< PTOResolveBufferSelectPass> { @@ -362,80 +301,6 @@ struct PTOResolveBufferSelectPass signalPassFailure(); return; } - SmallVector markers; - mod.walk([&](pto::SlotMarkerOp op) { markers.push_back(op); }); - if (markers.empty()) - return; - - for (auto op : markers) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - // Find the root multi-address pto.pointer_cast that this slot_marker - // refers to. If the chain does not land on one, the marker is not a - // multi-buffer reference; downgrade silently by forwarding source. - pto::PointerCastOp root = lookupRootPointerCast(op.getSource()); - if (!root) { - rewriter.replaceOp(op, op.getSource()); - continue; - } - - auto rootAddrs = root.getAddrs(); - uint32_t n = static_cast(rootAddrs.size()); - if (n < 2) { - // Single-address root: treat slot_marker as identity. - rewriter.replaceOp(op, op.getSource()); - continue; - } - if (n > mlir::pto::kPtoMultiBufferMaxNum) { - op.emitError() << "underlying pointer_cast has " << n - << " addresses, exceeds max " - << mlir::pto::kPtoMultiBufferMaxNum; - signalPassFailure(); - return; - } - - Type targetType = op.getResult().getType(); - - // Constant slot: emit a single-address pointer_cast for that slot. - IntegerAttr constSlotAttr; - if (matchPattern(op.getSlot(), m_Constant(&constSlotAttr))) { - int64_t slotI = constSlotAttr.getValue().getSExtValue(); - if (slotI < 0 || slotI >= static_cast(n)) { - op.emitError() << "constant slot " << slotI - << " is out of range for " - << n << " physical buffers"; - signalPassFailure(); - return; - } - Value picked = emitSlotPointerCast(rewriter, loc, root, - static_cast(slotI), - targetType); - rewriter.replaceOp(op, picked); - continue; - } - - // Dynamic slot: emit per-slot single-addr casts + N-way arith.select. - // The select chain uses the user-supplied SSA verbatim -- ptoas does - // NOT replace it with `iv mod N`. - SmallVector slotMems; - slotMems.reserve(n); - for (uint32_t i = 0; i < n; ++i) - slotMems.push_back( - emitSlotPointerCast(rewriter, loc, root, i, targetType)); - - Value selected = slotMems[0]; - Value slot = op.getSlot(); - for (uint32_t i = 1; i < n; ++i) { - Value iIdx = rewriter.create(loc, i); - Value isThis = rewriter.create( - loc, arith::CmpIPredicate::eq, slot, iIdx); - selected = rewriter.create(loc, isThis, slotMems[i], - selected); - } - rewriter.replaceOp(op, selected); - } } }; } // namespace diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index a8fb32a927..998ddc13d0 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -7857,25 +7857,6 @@ struct PTOSignalCommToEmitC : public OpConversionPattern { std::string callee; }; -struct PTODeclareTileMemRefToEmitC - : public OpConversionPattern { - using OpConversionPattern< - mlir::pto::DeclareTileMemRefOp>::OpConversionPattern; - - LogicalResult matchAndRewrite(mlir::pto::DeclareTileMemRefOp op, - OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - (void)adaptor; - Type convertedType = getTypeConverter()->convertType(op.getResult().getType()); - if (!convertedType) - return rewriter.notifyMatchFailure( - op, "failed to convert declare_tile_memref result type"); - rewriter.replaceOp(op, makeEmitCOpaqueConstant(rewriter, op.getLoc(), - convertedType, "nullptr")); - return success(); - } -}; - struct PTODeclareGlobalToEmitC : public OpConversionPattern { using OpConversionPattern< @@ -13156,14 +13137,6 @@ struct PTOBindTileToEmitC : public OpConversionPattern { return rewriter.create(loc, u64Ty, rawPtr).getResult(); }; - if (op.getSource().getDefiningOp()) { - FailureOr tileSpec = buildTileSpec(); - if (failed(tileSpec)) - return failure(); - rewriter.replaceOp(op, buildTileValue(*tileSpec)); - return success(); - } - Value tileCandidate = peelAllCasts(adaptor.getSource()); if (viewSemantics && viewSemantics.getValue() == "bitcast" && isTileLike(tileCandidate)) { @@ -13535,161 +13508,6 @@ struct PTOTileBufAddrToEmitC : public OpConversionPattern { } }; -struct PTOMaterializeTileToEmitC - : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - static bool isTileLike(Value v) { - auto ot = dyn_cast(v.getType()); - if (!ot) - return false; - StringRef s = ot.getValue(); - return s.contains("Tile<") || s.contains("ConvTile<"); - } - - LogicalResult matchAndRewrite(pto::MaterializeTileOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - Location loc = op.getLoc(); - MLIRContext *ctx = rewriter.getContext(); - auto tileTy = cast(op.getResult().getType()); - auto tileTypeString = getEmitCTileTypeString(tileTy); - if (!tileTypeString) - return rewriter.notifyMatchFailure( - op, "only rank-2 tile_buf handles can be materialized to EmitC"); - - Type convertedTy = getTypeConverter()->convertType(tileTy); - if (!convertedTy) - convertedTy = emitc::OpaqueType::get(ctx, *tileTypeString); - - Value source = peelUnrealized(adaptor.getSource()); - if (auto castOp = source.getDefiningOp()) - source = castOp.getOperand(); - - auto viewSemantics = op->getAttrOfType("pto.view_semantics"); - bool forceDynamicValid = op->hasAttr(kForceDynamicValidShapeAttrName); - bool isReshape = viewSemantics && viewSemantics.getValue() == "treshape"; - bool isSubview = viewSemantics && viewSemantics.getValue() == "subview"; - bool sourceIsDeclaredTile = - op.getSource().getDefiningOp(); - - auto createTileValue = [&]() -> Value { - SmallVector constructorArgs; - bool useConstructor = false; - pto::BLayout blayout = getTileBufBLayoutValue(tileTy.getConfigAttr()); - Type elemTy = tileTy.getElementType(); - auto shape = tileTy.getShape(); - auto validShape = tileTy.getValidShape(); - - auto makeCtorDimValue = [&](Value emitted, int64_t fallback) -> Value { - if (emitted) - return emitted; - return makeEmitCIntConstant( - rewriter, loc, emitc::OpaqueType::get(ctx, "int32_t"), fallback); - }; - auto maybeScaleDynamicValid = [&](Value emitted, int dimIdx) -> Value { - if (!emitted || !pto::isPTOFloat4PackedType(elemTy)) - return emitted; - int packedDim = blayout == pto::BLayout::ColMajor ? 0 : 1; - if (dimIdx != packedDim) - return emitted; - auto i32Ty = emitc::OpaqueType::get(ctx, "int32_t"); - Value two = makeEmitCIntConstant(rewriter, loc, i32Ty, 2); - return rewriter.create(loc, i32Ty, emitted, two).getResult(); - }; - auto fallbackDim = [&](int dimIdx) { - return renderTileTemplateDim(shape[dimIdx], elemTy, blayout, dimIdx); - }; - - if (forceDynamicValid) { - useConstructor = true; - constructorArgs.push_back(makeCtorDimValue( - maybeScaleDynamicValid(adaptor.getValidRow(), 0), fallbackDim(0))); - constructorArgs.push_back(makeCtorDimValue( - maybeScaleDynamicValid(adaptor.getValidCol(), 1), fallbackDim(1))); - } else { - if (validShape[0] == ShapedType::kDynamic) { - useConstructor = true; - constructorArgs.push_back(makeCtorDimValue( - maybeScaleDynamicValid(adaptor.getValidRow(), 0), fallbackDim(0))); - } - if (validShape[1] == ShapedType::kDynamic) { - useConstructor = true; - constructorArgs.push_back(makeCtorDimValue( - maybeScaleDynamicValid(adaptor.getValidCol(), 1), fallbackDim(1))); - } - } - - if (useConstructor) { - return rewriter - .create(loc, convertedTy, *tileTypeString, - ArrayAttr{}, ArrayAttr{}, - ValueRange(constructorArgs)) - .getResult(0); - } - - Value tile = rewriter - .create( - loc, getEmitCVariableResultType(convertedTy), - emitc::OpaqueAttr::get(ctx, "")) - .getResult(); - return loadEmitCVariableIfNeeded(rewriter, loc, tile); - }; - - if (!isSubview && !forceDynamicValid && isTileLike(source)) { - if (auto srcTy = dyn_cast(source.getType())) { - if (srcTy.getValue() == *tileTypeString) { - rewriter.replaceOp(op, source); - return success(); - } - } - } - - Value tile = createTileValue(); - if (sourceIsDeclaredTile) { - rewriter.replaceOp(op, tile); - return success(); - } - - if (isReshape && isTileLike(source)) { - rewriter.create(loc, TypeRange{}, "TRESHAPE", - ArrayAttr{}, ArrayAttr{}, - ValueRange{tile, source}); - rewriter.replaceOp(op, tile); - return success(); - } - - pto::AddressSpace as = pto::AddressSpace::GM; - if (auto asAttr = - dyn_cast_or_null(tileTy.getMemorySpace())) - as = asAttr.getAddressSpace(); - std::string elemTok = getEmitCScalarTypeToken(tileTy.getElementType()); - - Value rawPtr = source; - if (isTileLike(rawPtr)) - rawPtr = materializeTileDataValue(rewriter, loc, rawPtr, as, elemTok); - - auto u64Ty = emitc::OpaqueType::get(ctx, "uint64_t"); - Value addr = rawPtr; - if (isSetFFTsPointerLikeType(rawPtr.getType())) { - auto rcU64 = - rewriter.getArrayAttr({emitc::OpaqueAttr::get(ctx, "uint64_t")}); - addr = rewriter - .create(loc, u64Ty, "reinterpret_cast", - ArrayAttr{}, rcU64, - ValueRange{rawPtr}) - .getResult(0); - } else if (rawPtr.getType() != u64Ty) { - addr = rewriter.create(loc, u64Ty, rawPtr).getResult(); - } - - rewriter.create(loc, TypeRange{}, "TASSIGN", - ArrayAttr{}, ArrayAttr{}, - ValueRange{tile, addr}); - rewriter.replaceOp(op, tile); - return success(); - } -}; - // ============================================================================= // Arith CmpI -> EmitC Cmp // ============================================================================= @@ -14386,7 +14204,6 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); @@ -14619,7 +14436,6 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, typeConverter, ctx, "PTOAS__ASYNC_EVENT_TEST"); patterns.add(typeConverter, ctx, targetArch); patterns.add(typeConverter, ctx, targetArch); - patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add, diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 1853ea3a71..7908d30d30 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -207,11 +207,6 @@ static void lookupValidDims(Value v, Value &vRow, Value &vCol) { lookupValidDims(cast.getSource(), vRow, vCol); return; } - if (auto slot = v.getDefiningOp()) { - lookupValidDims(slot.getSource(), vRow, vCol); - return; - } - // pto.fusion_region result 不直接携带 valid dims；作为兜底情况，沿着 // pto.yield 回溯到 region 内真正的 tile handle/memref 定义链继续查找。 if (auto regionResult = dyn_cast(v)) { diff --git a/lib/PTO/Transforms/SlotAffineAnalysis.cpp b/lib/PTO/Transforms/SlotAffineAnalysis.cpp index cf633b1795..ddc7847060 100644 --- a/lib/PTO/Transforms/SlotAffineAnalysis.cpp +++ b/lib/PTO/Transforms/SlotAffineAnalysis.cpp @@ -22,14 +22,12 @@ using namespace mlir; namespace mlir { namespace pto { -Value findSlotMarkerExpr(Value v) { +Value findMultiTileSlotExpr(Value v) { int hops = 0; while (v && hops++ < 32) { Operation *op = v.getDefiningOp(); if (!op) return {}; - if (auto sm = dyn_cast(op)) - return sm.getSlot(); if (auto get = dyn_cast(op)) return get.getSlot(); if (auto bind = dyn_cast(op)) { diff --git a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp index c134ee9c4b..0c23b29cab 100644 --- a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp +++ b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp @@ -324,13 +324,6 @@ static void bindExplicitValidDims(ShapeConstraintSolver &solver, dims.cols, bind.getValidCol()); return; } - if (auto materialize = value.getDefiningOp()) { - bindDimToValue(solver, symbolDimByValue, canonicalByValue, signatureMap, - dims.rows, materialize.getValidRow()); - bindDimToValue(solver, symbolDimByValue, canonicalByValue, signatureMap, - dims.cols, materialize.getValidCol()); - return; - } if (auto subview = value.getDefiningOp()) { bindDimToValue(solver, symbolDimByValue, canonicalByValue, signatureMap, dims.rows, subview.getValidRow()); diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index 4978986158..2f4c24a331 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -169,11 +169,6 @@ std::optional> getOperationAliasInfo(Operation *op) { return std::make_pair(reshapeOp.getResult(), reshapeOp.getSrc()); } else if (auto bindTileOp = dyn_cast(op)) { return std::make_pair(bindTileOp.getResult(), bindTileOp.getSource()); - } else if (auto slotMarkerOp = dyn_cast(op)) { - // `pto.slot_marker` is a metadata-only view that tags a memref with the - // physical slot of a multi-buffer alloc. From an alias-walking - // standpoint it behaves like any other view-like op. - return std::make_pair(slotMarkerOp.getResult(), slotMarkerOp.getSource()); } else if (auto multiGetOp = dyn_cast(op)) { return std::make_pair(multiGetOp.getResult(), multiGetOp.getSource()); } else if (auto extSliceOp = dyn_cast(op)) { diff --git a/test/lit/pto/declare_tile_tile_native.pto b/test/lit/pto/declare_tile_tile_native.pto index 9a8a40ab50..0a5d5f6bf9 100644 --- a/test/lit/pto/declare_tile_tile_native.pto +++ b/test/lit/pto/declare_tile_tile_native.pto @@ -41,7 +41,6 @@ module { // NATIVE-LABEL: func.func @declare_tile_native // NATIVE: pto.declare_tile -> !pto.tile_buf // NATIVE: pto.tpop({{.*}} : !pto.tile_buf -// NATIVE-NOT: pto.declare_tile_memref // NATIVE-NOT: pto.bind_tile // NATIVE-NOT: memref.alloc diff --git a/test/lit/pto/materialize_tile_handles_emitc.pto b/test/lit/pto/materialize_tile_handles_emitc.pto index adf1e27029..58cce54f14 100644 --- a/test/lit/pto/materialize_tile_handles_emitc.pto +++ b/test/lit/pto/materialize_tile_handles_emitc.pto @@ -16,7 +16,6 @@ module { // IR: IR Dump After PTOMaterializeTileHandles // IR: pto.alloc_tile addr = {{.*}} : !pto.tile_buf // IR: pto.alloc_tile addr = {{.*}} : !pto.tile_buf -// IR-NOT: pto.materialize_tile // EMITC-LABEL: AICORE void materialize_tile_handles_emitc() // EMITC: Tile diff --git a/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto b/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto index 93b2490ee7..225922f5d7 100644 --- a/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto +++ b/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto @@ -51,7 +51,6 @@ module { // CHECK: %[[IDX:.*]] = arith.remui %arg{{.*}}, %c2 // CHECK: pto.multi_tile_get %[[MULTI]][%[[IDX]] // CHECK-NOT: memref.alloc -// CHECK-NOT: pto.slot_marker // After ResolveBufferSelect, the const preload slot becomes a directly // addressed alloc_tile, while the loop slot is chosen with arith.select. diff --git a/test/lit/pto/multi_tile_const_slot_gss_compiles.pto b/test/lit/pto/multi_tile_const_slot_gss_compiles.pto index 1dbfb17087..e9ba742136 100644 --- a/test/lit/pto/multi_tile_const_slot_gss_compiles.pto +++ b/test/lit/pto/multi_tile_const_slot_gss_compiles.pto @@ -10,9 +10,8 @@ // Smoke test: a constant-slot multi-buffer program compiles cleanly through // the GraphSyncSolver pipeline as well. GSS uses its own IR translator -// distinct from InsertSync's, so this guards against regressions when -// adding new view-like ops (here `pto.bind_tile` and `pto.slot_marker`) -// into the shared `getOperationAliasInfo` helper. We don't FileCheck +// distinct from InsertSync's, so this guards the tile-native +// `pto.multi_tile_get` alias path. We don't FileCheck // flag-emission specifics: a follow-up will teach GSS multi-buffer // event-id allocation, after which we can lock in optimization patterns. diff --git a/test/lit/pto/multi_tile_get_const_slot_lowering.pto b/test/lit/pto/multi_tile_get_const_slot_lowering.pto index 52b5ca5f01..a411d88b85 100644 --- a/test/lit/pto/multi_tile_get_const_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_const_slot_lowering.pto @@ -48,7 +48,6 @@ module { // CHECK: pto.multi_tile_get %[[MULTI]][%c0 // CHECK: pto.multi_tile_get %[[MULTI]][%c1 // CHECK-NOT: memref.alloc -// CHECK-NOT: pto.slot_marker // Stage 2 -- PTOPlanMemory records two physical slot addresses. // PLAN: IR Dump After PlanMemory diff --git a/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto b/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto index 14c055cd44..e96049174d 100644 --- a/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto @@ -53,7 +53,6 @@ module { // CHECK: pto.multi_tile_get %[[MULTI]][%[[NEXT_IDX]] // CHECK: pto.multi_tile_get %[[MULTI]][%[[CUR_IDX]] // CHECK-NOT: memref.alloc -// CHECK-NOT: pto.slot_marker // ResolveBufferSelect creates an addressed alloc_tile from an arith.select // chain driven by the user SSA (not an implicit iv mod N). diff --git a/test/lit/pto/multi_tile_no_loop_unroll.pto b/test/lit/pto/multi_tile_no_loop_unroll.pto index fe9c70ee80..b659bfc7fd 100644 --- a/test/lit/pto/multi_tile_no_loop_unroll.pto +++ b/test/lit/pto/multi_tile_no_loop_unroll.pto @@ -43,4 +43,3 @@ module { // CHECK: pto.multi_tile_get %[[MULTI]][%c0 // CHECK: pto.multi_tile_get %[[MULTI]][%c1 // CHECK-NOT: memref.alloc -// CHECK-NOT: pto.slot_marker diff --git a/test/lit/pto/multi_tile_two_buf_slot_binding.pto b/test/lit/pto/multi_tile_two_buf_slot_binding.pto index f5c45e587f..3be33e95c1 100644 --- a/test/lit/pto/multi_tile_two_buf_slot_binding.pto +++ b/test/lit/pto/multi_tile_two_buf_slot_binding.pto @@ -18,7 +18,7 @@ // Because the producing `tadd` touches *two* multi-buffers at different // phases, the dyn event lane for B's hazard must be keyed off B's slot // (i%2), NOT the tadd's first read operand A (which is (i+1)%2). The earlier -// code grabbed the op's first slot_marker memref and would have keyed the +// code grabbed the op's first multi_tile_get and would have keyed the // wait off A's slot -- selecting the wrong event lane and letting the real // cross-iter dependency on B through. diff --git a/test/lit/pto/retired_memref_bridge_ops_invalid.pto b/test/lit/pto/retired_memref_bridge_ops_invalid.pto new file mode 100644 index 0000000000..30fbc6158c --- /dev/null +++ b/test/lit/pto/retired_memref_bridge_ops_invalid.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: split-file %s %t +// RUN: not ptoas %t/materialize.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=MATERIALIZE +// RUN: not ptoas %t/declare.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=DECLARE +// RUN: not ptoas %t/slot.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=SLOT + +//--- materialize.pto +module { + func.func @materialize() { + %0 = "pto.materialize_tile"() : () -> i32 + return + } +} + +// MATERIALIZE: error: unregistered operation 'pto.materialize_tile' found in dialect ('pto') + +//--- declare.pto +module { + func.func @declare() { + %0 = "pto.declare_tile_memref"() : () -> i32 + return + } +} + +// DECLARE: error: unregistered operation 'pto.declare_tile_memref' found in dialect ('pto') + +//--- slot.pto +module { + func.func @slot() { + %0 = "pto.slot_marker"() : () -> i32 + return + } +} + +// SLOT: error: unregistered operation 'pto.slot_marker' found in dialect ('pto') diff --git a/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto b/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto index d4fd3668d2..ad9fb57710 100644 --- a/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto +++ b/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto @@ -35,7 +35,6 @@ module { // IR: %[[BYTE_OFFSET:.*]] = arith.muli %[[ELEM_OFFSET]], %[[ELEM_BYTES]] : i64 // IR: %[[ADDR:.*]] = arith.addi %{{.*}}, %[[BYTE_OFFSET]] : i64 // IR: pto.alloc_tile addr = %[[ADDR]] {pto.view_semantics = "subview"} : !pto.tile_buf -// IR-NOT: pto.materialize_tile // EMITC: Tile // EMITC-NOT: Tile diff --git a/test/lit/pto/tcvt_low_precision_a5_valid.pto b/test/lit/pto/tcvt_low_precision_a5_valid.pto index 94c02082cd..15cc41a6da 100644 --- a/test/lit/pto/tcvt_low_precision_a5_valid.pto +++ b/test/lit/pto/tcvt_low_precision_a5_valid.pto @@ -30,7 +30,6 @@ 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-NOT: pto.declare_tile_memref // 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 diff --git a/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto b/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto index 0b5c48df87..84fbf30052 100644 --- a/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto +++ b/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto @@ -25,6 +25,5 @@ module { // CHECK: func.func @tload_tprefetch_low_precision_a5_valid(%arg0: memref<16x16xf8E4M3FN>, %arg1: memref<16x16x!pto.hif8>) // CHECK: pto.declare_tile -> !pto.tile_buf -// CHECK-NOT: pto.declare_tile_memref // CHECK: pto.tload ins(%arg0 : memref<16x16xf8E4M3FN>) outs( // CHECK: pto.tprefetch ins(%arg1 : memref<16x16x!pto.hif8>) outs( diff --git a/test/lit/pto/tstore_low_precision_a5_valid.pto b/test/lit/pto/tstore_low_precision_a5_valid.pto index 520ef20e64..df0099391d 100644 --- a/test/lit/pto/tstore_low_precision_a5_valid.pto +++ b/test/lit/pto/tstore_low_precision_a5_valid.pto @@ -27,7 +27,6 @@ module { // CHECK: func.func @tstore_low_precision_a5_valid(%arg0: memref<16x16xf8E4M3FN>, %arg1: memref<16x16x!pto.hif8>, %arg2: i64) // CHECK: %[[VEC:.*]] = pto.declare_tile -> !pto.tile_buf // CHECK: %[[ACC:.*]] = pto.declare_tile -> !pto.tile_buf -// CHECK-NOT: pto.declare_tile_memref // CHECK-NOT: pto.bind_tile // CHECK: pto.section.vector { // CHECK: pto.tstore ins(%[[VEC]] : !pto.tile_buf) diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 17fcfbb0a8..b6343a34ca 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -3139,8 +3139,8 @@ int mlir::pto::compilePTOASModule( // Conditionally add one automatic synchronization mode. Barrier-all is a // conservative standalone pass; InsertSync and GraphSyncSolver are set/wait // solvers. Sync runs BEFORE PTOResolveBufferSelect so it sees per-use - // `pto.slot_marker` ops and can keep multi-buffer slot identity (const slot - // K vs slot K' or dynamic slot) for the alias / event-id analysis. + // `pto.multi_tile_get` operations and keeps their slot identity for alias + // and event-id analysis. // solvers, while BufidSync is A5-only get_buf/rls_buf synchronization. pm.addNestedPass( pto::createPTOVerifySubkernelPipeContractPass()); @@ -3161,10 +3161,8 @@ int mlir::pto::compilePTOASModule( pto::createPTOGraphSyncSolverPass(graphSyncOpts)); } - // Materialize per-slot single-address `pto.pointer_cast` (constant slot) - // or an `arith.select` chain (dynamic slot). The multi-address cast - // produced by explicit-address lowering or PTOPlanMemory survives as - // the alloc anchor. + // Materialize each `pto.multi_tile_get` as an addressed `pto.alloc_tile`; + // dynamic selections use an `arith.select` chain over planned addresses. pm.addPass(pto::createPTOResolveBufferSelectPass()); if (effectiveBackend == PTOBackend::EmitC) pm.addPass(createNarrowUnusedMultiResultProvenancePass()); From 1af4ed048587673d5141c768911bf3720f8520a1 Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 20 Jul 2026 16:55:52 +0800 Subject: [PATCH 85/91] Remove retired bind tile bridge --- include/PTO/IR/PTOOps.td | 30 - include/PTO/IR/VPTOOps.td | 2 +- include/PTO/Transforms/Passes.td | 2 +- include/PTO/Transforms/SlotAffineAnalysis.h | 3 +- lib/PTO/IR/PTO.cpp | 27 +- lib/PTO/IR/VPTO.cpp | 6 +- lib/PTO/Transforms/AllocToPointerCast.cpp | 23 +- .../InsertSync/InsertSyncAnalysis.cpp | 20 - .../Transforms/InsertSync/PTOIRTranslator.cpp | 3 - .../Transforms/InsertSync/SyncMacroModel.cpp | 2 - .../Transforms/PTOMaterializeTileHandles.cpp | 36 +- lib/PTO/Transforms/PTOPlanMemory.cpp | 5 - lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 5 +- lib/PTO/Transforms/PTOToEmitC.cpp | 648 +----------------- lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp | 2 +- lib/PTO/Transforms/PTOValidateVPTOIR.cpp | 2 +- lib/PTO/Transforms/PTOViewToMemref.cpp | 7 +- lib/PTO/Transforms/SlotAffineAnalysis.cpp | 4 - .../Transforms/TileFusion/FusionAnalysis.cpp | 7 - .../TileFusion/PTOFusionLoadStoreElision.cpp | 44 -- .../TileFusion/PTOLowLevelLoopFusion.cpp | 2 - lib/PTO/Transforms/Utils.cpp | 4 - .../Transforms/VPTOBufferMaterialization.cpp | 4 - lib/PTO/Transforms/VPTOPtrNormalize.cpp | 37 +- test/lit/pto/add_carry_tile_native.pto | 1 - test/lit/pto/addmul_scalar_tile_native.pto | 1 - test/lit/pto/addmul_tile_native.pto | 1 - ...alloc_tile_plan_memory_no_memref_alloc.pto | 1 - test/lit/pto/arg_reduction_tile_native.pto | 1 - test/lit/pto/axpy_dequant_tile_native.pto | 1 - test/lit/pto/basic_float_tile_native.pto | 1 - .../bias_tile_subview_emitc_pointer_type.pto | 7 +- test/lit/pto/bitcast_tile_native.pto | 1 - .../pto/build_async_session_tile_native.pto | 1 - test/lit/pto/col_minmax_tile_native.pto | 1 - test/lit/pto/colexpand_tile_native.pto | 1 - test/lit/pto/colsum_tile_native.pto | 1 - test/lit/pto/compare_tile_native.pto | 1 - test/lit/pto/concat_tile_native.pto | 1 - test/lit/pto/cvt_tile_native.pto | 1 - test/lit/pto/declare_tile_tile_native.pto | 1 - test/lit/pto/div_tile_native.pto | 1 - test/lit/pto/expands_tile_native.pto | 1 - .../lit/pto/extract_insert_fp_tile_native.pto | 1 - test/lit/pto/extract_insert_tile_native.pto | 1 - test/lit/pto/fillpad_tile_native.pto | 1 - test/lit/pto/gather_tile_native.pto | 1 - test/lit/pto/gemv_tile_native.pto | 1 - .../issue708_zero_valid_subview_inferred.pto | 4 +- test/lit/pto/load_store_tile_native.pto | 1 - test/lit/pto/logic_binary_tile_native.pto | 1 - test/lit/pto/logic_scalar_tile_native.pto | 1 - test/lit/pto/matmul_tile_native.pto | 1 - test/lit/pto/mgather_mscatter_tile_native.pto | 1 - test/lit/pto/minmax_scalar_tile_native.pto | 1 - test/lit/pto/minmax_tile_native.pto | 1 - test/lit/pto/mov_tile_native.pto | 1 - .../lit/pto/movement_metadata_tile_native.pto | 1 - test/lit/pto/mrgsort_tile_native.pto | 1 - test/lit/pto/part_arithmetic_tile_native.pto | 1 - test/lit/pto/pipe_tile_native.pto | 2 - .../plan_memory_bind_tile_alias_liveness.pto | 35 - test/lit/pto/pow_rsqrt_tile_native.pto | 1 - test/lit/pto/print_tile_native.pto | 1 - test/lit/pto/quant_mx_tile_native.pto | 1 - test/lit/pto/reduction_tile_native.pto | 1 - .../pto/retired_memref_bridge_ops_invalid.pto | 11 + test/lit/pto/rowexpand_tile_native.pto | 1 - .../pto/scratch_elementwise_tile_native.pto | 1 - test/lit/pto/select_tile_native.pto | 1 - test/lit/pto/shift_tile_native.pto | 1 - test/lit/pto/sort32_tile_native.pto | 1 - test/lit/pto/store_fp_tile_native.pto | 1 - test/lit/pto/sub_tile_native.pto | 1 - ...dynamic_offset_static_valid_regression.pto | 1 - ...> subview_tile_native_preserve_stride.pto} | 2 +- test/lit/pto/tassign_tile_native_arg.pto | 1 - .../lit/pto/tile_buf_addr_tile_native_arg.pto | 1 - .../pto/tile_scalar_access_tile_native.pto | 1 - test/lit/pto/tlrelu_tile_native.pto | 1 - test/lit/pto/tmov_fp_tile_native.pto | 1 - test/lit/pto/tprefetch_tile_native.pto | 1 - test/lit/pto/trans_tile_native.pto | 1 - .../pto/transcendental_unary_tile_native.pto | 1 - test/lit/pto/treshape_tile_native_arg.pto | 1 - .../lit/pto/tstore_low_precision_a5_valid.pto | 1 - test/lit/pto/tsync_tile_native.pto | 1 - test/lit/pto/unary_compute_tile_native.pto | 1 - test/lit/pto/validshape_tile_native.pto | 1 - test/lit/pto/vector_specialty_tile_native.pto | 1 - test/lit/pto/xor_tile_native.pto | 1 - test/lit/vpto/fold_tile_buf_intrinsics.pto | 1 - ...veness.py => plan_memory_tile_liveness.py} | 2 +- test/samples/runop.sh | 2 +- 94 files changed, 38 insertions(+), 1014 deletions(-) delete mode 100644 test/lit/pto/plan_memory_bind_tile_alias_liveness.pto rename test/lit/pto/{subview_bind_tile_preserve_stride.pto => subview_tile_native_preserve_stride.pto} (98%) rename test/samples/planmemory/{plan_memory_bind_tile_alias_liveness.py => plan_memory_tile_liveness.py} (96%) diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 2ace62496b..b56c6bb8e1 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -436,36 +436,6 @@ def MultiTileGetOp : PTO_Op<"multi_tile_get", [ } -// ============================================================================ -// BindTileOp: 将 Config 和 Valid Dims 绑定到 MemRef 上 -// ============================================================================ -def BindTileOp : PTO_Op<"bind_tile", [ - Pure, - AttrSizedOperandSegments - // 允许输入 offset:0 -> 输出 offset:? - ]> { - let summary = "Binds metadata and implicitly casts layout"; - let description = [{ - Wraps a memref with PTO metadata (valid dimensions and config). - }]; - - // [修改] 以前是 Variadic:$valid_dims - // 现在改为明确的 Optional Row/Col，与 PointerCastOp 保持一致 - let arguments = (ins - AnyMemRef:$source, - Optional:$valid_row, - Optional:$valid_col, - TileBufConfigAttr:$config - ); - - let results = (outs AnyMemRef:$result); - - // [修改] assemblyFormat: 去掉 valid_dims，改为 ($row, $col)? - let assemblyFormat = [{ - $source (`,` $valid_row^ `,` $valid_col)? attr-dict `:` qualified(type($source)) `->` qualified(type($result)) - }]; -} - def SubViewOp : PTO_Op<"subview", [ Pure, ViewLikeOpInterface, diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index 603a18dc51..9419565165 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -248,7 +248,7 @@ def TileBufAddrOp : PTO_Op<"tile_buf_addr", [Pure]> { This op is emitted by TileLang DSL templates and resolved by the FoldTileBufIntrinsics pass after inlining. Tile-native uses remain in PTO IR through memory planning and sync. Hand-written legacy `.pto` may also - use it directly on the memref result of `pto.bind_tile`. + use it directly on an addressed memref. }]; let arguments = (ins AnyTypeOf<[TileBufType, AnyMemRef], diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 8eb5755b04..420d540df2 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -873,7 +873,7 @@ def PTOVPTOPtrBoundary rewrites eligible memref function arguments to same-space `!pto.ptr`, rejects memref function results, canonicalizes supported body-level VPTO buffer-like ops to ptr-form, and drops dead boundary/view scaffold such as - trivial `pto.castptr`, `pto.bind_tile`, `memref.subview`, + trivial `pto.castptr`, `memref.subview`, `memref.reinterpret_cast`, and `memref.memory_space_cast` once they become unused. }]; diff --git a/include/PTO/Transforms/SlotAffineAnalysis.h b/include/PTO/Transforms/SlotAffineAnalysis.h index 303e04ade6..3b99bc8814 100644 --- a/include/PTO/Transforms/SlotAffineAnalysis.h +++ b/include/PTO/Transforms/SlotAffineAnalysis.h @@ -36,8 +36,7 @@ enum class SlotRelation { kUnknown, // can neither prove equal nor disjoint }; -/// Walk back through metadata-only ops (`pto.bind_tile`) to the nearest -/// `pto.multi_tile_get` and return its slot SSA value. Returns a null Value +/// Return the slot SSA value carried by `pto.multi_tile_get`. Returns null /// if the chain does not pass through a multi_tile_get. mlir::Value findMultiTileSlotExpr(mlir::Value v); diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 6cbbe7e952..b0128fbe6c 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -495,8 +495,6 @@ static bool shouldBypassDecodedMemrefVerifier(Operation *op) { for (Value operand : op->getOperands()) { if (isa(operand.getType())) return true; - if (operand.getDefiningOp()) - return true; } return false; } @@ -3836,26 +3834,10 @@ static SmallVector getLogicalTileExtentVec(Type ty, return dims; } -static int64_t getConstantIndexOrDynamic(Value value) { - if (!value) - return ShapedType::kDynamic; - if (auto cst = value.getDefiningOp()) - return cst.value(); - if (auto cst = value.getDefiningOp()) - return cst.value(); - return ShapedType::kDynamic; -} - static SmallVector getValidShapeVec(Value value) { if (!value) return {}; auto valid = getValidShapeVec(value.getType()); - if (auto bind = value.getDefiningOp()) { - if (valid.size() >= 1 && bind.getValidRow()) - valid[0] = getConstantIndexOrDynamic(bind.getValidRow()); - if (valid.size() >= 2 && bind.getValidCol()) - valid[1] = getConstantIndexOrDynamic(bind.getValidCol()); - } return valid; } @@ -11127,7 +11109,7 @@ static bool isLocallyBoundTileSource(Value value) { if (!value || isa(value)) return false; - if (isa( + if (isa( value.getDefiningOp())) return true; @@ -12503,9 +12485,7 @@ mlir::LogicalResult mlir::pto::TStoreFPOp::verify() { auto shouldBypassDecoded = [&]() -> bool { Value src = getSrc(); Value fp = getFp(); - return isa(src.getType()) || isa(fp.getType()) || - src.getDefiningOp() || - fp.getDefiningOp(); + return isa(src.getType()) || isa(fp.getType()); }; auto verifySrcDtypeAlways = [&]() -> LogicalResult { @@ -13995,8 +13975,7 @@ mlir::LogicalResult mlir::pto::SubViewOp::verify() { return emitOpError("expects result to have rank-2 valid_shape"); // With the valid operand omitted, the result type is authoritative for the // valid extent: accept any static value in [0, size] (this subsumes both the - // full-size default and the v=0 no-op-replay empty marker). Lowering derives - // the bind_tile valid operand from this type. A dynamic result valid still + // full-size default and the v=0 no-op-replay empty marker). A dynamic result valid still // requires an explicit operand to supply the runtime extent, so it stays // rejected on this path. bool rowInferred = !getValidRow() && dstValid[0] != ShapedType::kDynamic && diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index 11dea4fd23..bf4aaeb14e 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -4960,10 +4960,8 @@ LogicalResult TileBufAddrOp::verify() { srcMemorySpace = srcTileType.getMemorySpace(); srcRank = static_cast(srcTileType.getShape().size()); } else if (auto srcMemRefType = dyn_cast(getSrc().getType())) { - // PTOViewToMemref may lower tile_buf producers to memref + pto.bind_tile - // before the shared materialization bridge restores tile handles. - // Hand-written pto.tile_buf_addr may therefore temporarily see a tile-bound - // memref operand in that intermediate form. + // Hand-written legacy IR may use tile_buf_addr directly on an addressed + // memref before the shared materialization bridge restores tile handles. elementType = srcMemRefType.getElementType(); srcMemorySpace = srcMemRefType.getMemorySpace(); srcRank = srcMemRefType.getRank(); diff --git a/lib/PTO/Transforms/AllocToPointerCast.cpp b/lib/PTO/Transforms/AllocToPointerCast.cpp index 925993e8ec..667481964b 100644 --- a/lib/PTO/Transforms/AllocToPointerCast.cpp +++ b/lib/PTO/Transforms/AllocToPointerCast.cpp @@ -29,26 +29,6 @@ static uint64_t alignUpToDefault(uint64_t value) { kDefaultAllocAlignmentBytes; } -static TileBufConfigAttr inferBindTileConfig(Value root, Operation *diagOp) { - TileBufConfigAttr configAttr; - for (Operation *user : root.getUsers()) { - auto bind = dyn_cast(user); - if (!bind || bind.getSource() != root) - continue; - if (!configAttr) { - configAttr = bind.getConfigAttr(); - continue; - } - if (configAttr != bind.getConfigAttr()) { - diagOp->emitWarning( - "alloc has multiple bind_tile users with different configs; " - "using the first one"); - break; - } - } - return configAttr; -} - static SmallVector getAllocatedOffsets( Value root, Operation *allocLikeOp, BaseMemRefType memRefType, const DenseMap> &buffer2Offsets, @@ -114,7 +94,6 @@ static std::pair getDynamicValidShapeValues(memref::AllocOp op) { LogicalResult MemrefAllocaOpToPointerCastOpPattern::matchAndRewrite( memref::AllocOp op, PatternRewriter &rewriter) const { const auto ¤tMemRefType = cast(op.getType()); - TileBufConfigAttr configAttr = inferBindTileConfig(op.getResult(), op); SmallVector offsets = getAllocatedOffsets(op.getResult(), op, currentMemRefType, buffer2Offsets, fallbackNextOffset); @@ -129,7 +108,7 @@ LogicalResult MemrefAllocaOpToPointerCastOpPattern::matchAndRewrite( auto [vRow, vCol] = getDynamicValidShapeValues(op); auto ptoPointerCastOp = rewriter.create( op.getLoc(), currentMemRefType, ValueRange(addrs), vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); + vCol ? vCol : Value(), Attribute()); rewriter.replaceOp(op, ptoPointerCastOp->getResults()); return success(); diff --git a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp index 723dbea214..eac1305f22 100644 --- a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp +++ b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp @@ -75,31 +75,11 @@ static std::optional getKnownRepeatAccessShapeFromType(Type t return std::nullopt; } -static std::optional getConstantIndex(Value value) { - if (!value) return std::nullopt; - APInt intValue; - if (!matchPattern(value, m_ConstantInt(&intValue))) return std::nullopt; - return intValue.getSExtValue(); -} - static std::optional getKnownRepeatAccessShape(Value access) { if (!access) return std::nullopt; auto shape = getKnownRepeatAccessShapeFromType(access.getType()); if (!shape) return std::nullopt; - if (auto bind = access.getDefiningOp()) { - auto row = getConstantIndex(bind.getValidRow()); - auto col = getConstantIndex(bind.getValidCol()); - if (row && col) { - if (*row < 0 || *col < 0 || *row > shape->fullShape[0] || - *col > shape->fullShape[1]) - return std::nullopt; - shape->validShape = SmallVector{*row, *col}; - } else if (bind.getValidRow() || bind.getValidCol()) { - return std::nullopt; - } - } - return shape; } diff --git a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp index fbbfcab29f..eccdd1303a 100644 --- a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp +++ b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp @@ -392,9 +392,6 @@ void PTOIRTranslator::RecursionIR(Region *region) { else if (auto makeViewOp = dyn_cast(op)) { UpdateAliasBufferInfo(makeViewOp.getResult(), makeViewOp.getPtr()); } - else if (auto bindTileOp = dyn_cast(op)) { - UpdateAliasBufferInfo(bindTileOp.getResult(), bindTileOp.getSource()); - } else if (auto subViewOp = dyn_cast(op)) { UpdateAliasBufferInfo(subViewOp.getResult(), subViewOp.getSource()); } diff --git a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp index 14ecc4d453..86ad3f99a7 100644 --- a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp +++ b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp @@ -239,8 +239,6 @@ static std::optional getConstantIndex(Value value) { } static std::optional> lookupMGatherValidDims(Value value) { - if (auto bind = value.getDefiningOp()) - return SmallVector{bind.getValidRow(), bind.getValidCol()}; if (auto pc = value.getDefiningOp()) return SmallVector{pc.getValidRow(), pc.getValidCol()}; if (auto subview = value.getDefiningOp()) diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp index 1a91787481..ec8f7201e8 100644 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp @@ -93,7 +93,7 @@ static bool isLocalTileMemRef(Type type) { } static bool shouldMaterializeOperand(Operation *owner) { - if (isa(owner)) + if (isa(owner)) return false; StringRef name = owner->getName().getStringRef(); @@ -238,16 +238,6 @@ static TileHandleMetadata getTileHandleMetadata(Value value, meta.source = value; meta.config = TileBufConfigAttr::getDefault(ctx); - if (auto bind = value.getDefiningOp()) { - meta.source = bind.getSource(); - meta.validRow = bind.getValidRow(); - meta.validCol = bind.getValidCol(); - meta.config = bind.getConfig(); - meta.explicitConfig = true; - copyTileHandleAttrs(bind, meta.attrs); - return meta; - } - if (auto cast = value.getDefiningOp()) { meta.validRow = cast.getValidRow(); meta.validCol = cast.getValidCol(); @@ -404,7 +394,7 @@ static TileBufType buildTileTypeFromMemRef(MemRefType memTy, } static bool isMaterializedTileAnchor(Operation *op) { - return isa(op); + return isa(op); } static Value makeI64Constant(OpBuilder &builder, Location loc, int64_t value) { @@ -549,9 +539,6 @@ static Value computeExplicitAddress(Value value, OpBuilder &builder, if (auto alloc = value.getDefiningOp()) return ensureI64(alloc.getAddr(), builder, loc); - if (auto bind = value.getDefiningOp()) - return computeExplicitAddress(bind.getSource(), builder, loc); - if (auto cast = value.getDefiningOp()) { if (cast.getAddrs().empty()) return Value(); @@ -596,11 +583,6 @@ static bool isControlFlowAddressProducer(Operation *op) { static Value peelAddressSource(Value value) { while (true) { - if (auto bind = value.getDefiningOp()) { - value = bind.getSource(); - continue; - } - if (auto subview = value.getDefiningOp()) { value = subview.getSource(); continue; @@ -895,20 +877,6 @@ struct PTOMaterializeTileHandlesPass return; } - bool erasedBind = true; - while (erasedBind) { - erasedBind = false; - SmallVector deadBinds; - module.walk([&](BindTileOp op) { - if (op.getResult().use_empty()) - deadBinds.push_back(op); - }); - for (Operation *op : deadBinds) { - op->erase(); - erasedBind = true; - } - } - } }; diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index 8ee98ada83..daa804e043 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -467,11 +467,6 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { } else if (isa(op)) { // Runtime-bound tile handles do not allocate static local storage. return WalkResult::advance(); - } else if (auto bindOp = dyn_cast(op)) { - // BindTile result is only an alias of the source buffer. Treat every use - // of the result as a use of the source in liveness analysis. - UpdateBufferAlias(bindOp.getResult(), bindOp.getSource()); - return WalkResult::advance(); } else if (isLocalMemPlan() && dyn_cast(op)) { auto allocTileOp = cast(op); if (allocTileOp.getAddr()) { diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index aaaa584b2b..d1eb6de846 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -683,10 +683,7 @@ struct PlannerAnalysis { } } - if (auto bind = dyn_cast(op)) { - setRoots(bind.getResult(), getRoots(bind.getSource())); - propagateSplitTpopDerived(bind.getResult(), ValueRange{bind.getSource()}); - } else if (auto multiGet = dyn_cast(op)) { + if (auto multiGet = dyn_cast(op)) { setRoots(multiGet.getResult(), getRoots(multiGet.getSource())); propagateSplitTpopDerived(multiGet.getResult(), ValueRange{multiGet.getSource()}); diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 998ddc13d0..4c5af18f12 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -522,9 +522,7 @@ getSpecialScaleGlobalTensorTypeSpecForTileValue(Value dstValue, return std::nullopt; std::optional configAttr; - if (auto bind = dstValue.getDefiningOp()) { - configAttr = bind.getConfigAttr(); - } else if (auto cast = dstValue.getDefiningOp()) { + if (auto cast = dstValue.getDefiningOp()) { if (auto config = cast.getConfig()) configAttr = config; } @@ -4525,9 +4523,7 @@ static FailureOr buildAsyncScratchTileValue( auto *ctx = rewriter.getContext(); pto::TileBufConfigAttr configAttr = pto::TileBufConfigAttr::getDefault(ctx); - if (auto bind = originalScratch.getDefiningOp()) { - configAttr = bind.getConfig(); - } else if (auto cast = originalScratch.getDefiningOp()) { + if (auto cast = originalScratch.getDefiningOp()) { if (auto config = cast.getConfig()) configAttr = *config; } @@ -4596,9 +4592,7 @@ static FailureOr buildSyncAllWorkspaceTileValue( auto *ctx = rewriter.getContext(); pto::TileBufConfigAttr configAttr = pto::TileBufConfigAttr::getDefault(ctx); - if (auto bind = originalWorkspace.getDefiningOp()) { - configAttr = bind.getConfig(); - } else if (auto cast = originalWorkspace.getDefiningOp()) { + if (auto cast = originalWorkspace.getDefiningOp()) { if (auto config = cast.getConfig()) configAttr = *config; } @@ -8592,7 +8586,7 @@ struct PTOTPushToEmitC : public OpConversionPattern { if (failed(pipeTok)) return rewriter.notifyMatchFailure(op, "failed to resolve pipe token"); // Read the tile type token from the already-converted OpaqueType, which - // preserves the exact layout produced by BindTileOp / PointerCastOp EmitC. + // preserves the exact layout produced by PointerCastOp EmitC. Value convertedTile = peelUnrealized(adaptor.getTile()); auto tileTok = getPipeDataTypeToken(convertedTile); if (failed(tileTok)) @@ -12613,639 +12607,6 @@ struct PTOTrapOpToEmitC : public OpConversionPattern { } }; -// ============================================================================= -// 2. BindTileOp Lowering (FIX: Trace back to physical address) -// ============================================================================= -struct PTOBindTileToEmitC : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - struct TileBuildSpec { - std::string tileTypeStr; - bool useConstructor = false; - SmallVector constructorArgs; - }; - - enum class MxValidAlignPolicy { - None, - AlignRowToShape, - AlignColToShape, - }; - - static bool getIndexConst(Value v, int64_t &out) { - if (!v) - return false; - if (auto cst = v.getDefiningOp()) { - if (auto ia = dyn_cast(cst.getValue())) { - out = getIntegerAttrSignedValue(ia); - return true; - } - } - return false; - } - - static bool getTilePointerStrides(pto::TileBufConfigAttr configAttr, - Type elemTy, int64_t rows, int64_t cols, - int64_t &rowStride, - int64_t &colStride) { - if (rows == ShapedType::kDynamic || cols == ShapedType::kDynamic) - return false; - - int32_t blVal = 0; - if (auto blAttr = dyn_cast(configAttr.getBLayout())) - blVal = static_cast(blAttr.getValue()); - else if (auto intAttr = dyn_cast(configAttr.getBLayout())) - blVal = static_cast(getIntegerAttrSignedValue(intAttr)); - - int32_t slVal = 0; - if (auto slAttr = dyn_cast(configAttr.getSLayout())) - slVal = static_cast(slAttr.getValue()); - else if (auto intAttr = dyn_cast(configAttr.getSLayout())) - slVal = static_cast(getIntegerAttrSignedValue(intAttr)); - - bool boxed = slVal != 0; - int64_t innerRows = 1; - int64_t innerCols = 1; - if (boxed) { - int32_t fractal = 512; - if (auto frAttr = dyn_cast(configAttr.getSFractalSize())) - fractal = static_cast(getIntegerAttrSignedValue(frAttr)); - - unsigned elemBytes = pto::getPTOStorageElemByteSize(elemTy); - if (elemBytes == 0) - return false; - - switch (fractal) { - case 1024: - innerRows = 16; - innerCols = 16; - break; - case 32: - innerRows = 16; - innerCols = 2; - break; - case 512: - if (slVal == 1) { - innerRows = 16; - innerCols = 32 / elemBytes; - } else if (slVal == 2) { - innerRows = 32 / elemBytes; - innerCols = 16; - } else { - return false; - } - break; - default: - return false; - } - if (innerRows <= 0 || innerCols <= 0) - return false; - } - - if (!boxed) { - if (blVal == 1) { - rowStride = 1; - colStride = rows; - } else { - rowStride = cols; - colStride = 1; - } - return true; - } - - if (blVal == 1) { - if (slVal != 1) - return false; - rowStride = innerCols; - colStride = rows; - return true; - } - - rowStride = cols; - colStride = innerRows; - return true; - } - - LogicalResult matchAndRewrite(pto::BindTileOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op.getLoc(); - auto *ctx = rewriter.getContext(); - auto configAttr = op.getConfigAttr(); - auto viewSemantics = op->getAttrOfType("pto.view_semantics"); - bool isSubView = viewSemantics && viewSemantics.getValue() == "subview"; - - auto peelAllCasts = [](Value v) { - while (auto castOp = v.getDefiningOp()) - v = castOp.getOperand(0); - if (auto castOp = v.getDefiningOp()) - v = castOp.getOperand(); - return v; - }; - auto isTileLike = [](Value v) -> bool { - auto ot = dyn_cast(v.getType()); - if (!ot) - return false; - StringRef s = ot.getValue(); - return s.contains("Tile<") || s.contains("ConvTile<"); - }; - auto buildTileSpec = [&]() -> FailureOr { - auto resMrTy = dyn_cast(op.getType()); - if (!resMrTy) - return failure(); - - const char *roleTok = "TileType::Vec"; - pto::AddressSpace resultAS = pto::AddressSpace::GM; - Type elemTy = resMrTy.getElementType(); - if (auto asAttr = - dyn_cast_or_null(resMrTy.getMemorySpace())) { - resultAS = asAttr.getAddressSpace(); - switch (asAttr.getAddressSpace()) { - case pto::AddressSpace::VEC: - roleTok = "TileType::Vec"; - break; - case pto::AddressSpace::MAT: - roleTok = "TileType::Mat"; - break; - case pto::AddressSpace::LEFT: - roleTok = "TileType::Left"; - break; - case pto::AddressSpace::RIGHT: - roleTok = "TileType::Right"; - break; - case pto::AddressSpace::ACC: - roleTok = "TileType::Acc"; - break; - case pto::AddressSpace::BIAS: - roleTok = "TileType::Bias"; - break; - case pto::AddressSpace::SCALING: - roleTok = scalingRoleToken(elemTy, configAttr); - break; - case pto::AddressSpace::GM: - case pto::AddressSpace::Zero: - roleTok = "TileType::Vec"; - break; - } - } - - Type emitElemTy = getTypeConverter()->convertType(elemTy); - if (!emitElemTy) - return failure(); - auto emitElemOpaque = dyn_cast(emitElemTy); - if (!emitElemOpaque) - return failure(); - std::string elemTypeStr = emitElemOpaque.getValue().str(); - - if (resMrTy.getRank() < 2) - return failure(); - int64_t rows = resMrTy.getDimSize(0); - int64_t cols = resMrTy.getDimSize(1); - if (rows == ShapedType::kDynamic || cols == ShapedType::kDynamic) - return failure(); - - auto inferMxValidAlignPolicy = [&]() -> MxValidAlignPolicy { - auto matchTMatmulMxLike = [&](auto user) -> MxValidAlignPolicy { - if (resultAS == pto::AddressSpace::LEFT && - peelAllCasts(user.getA()) == op.getResult()) - return MxValidAlignPolicy::AlignColToShape; - if (resultAS == pto::AddressSpace::RIGHT && - peelAllCasts(user.getB()) == op.getResult()) - return MxValidAlignPolicy::AlignRowToShape; - return MxValidAlignPolicy::None; - }; - auto matchTGemvMxLike = [&](auto user) -> MxValidAlignPolicy { - if (resultAS == pto::AddressSpace::RIGHT && - peelAllCasts(user.getB()) == op.getResult()) - return MxValidAlignPolicy::AlignRowToShape; - return MxValidAlignPolicy::None; - }; - - auto walkUsers = [&](auto &&self, Value value) -> MxValidAlignPolicy { - for (Operation *user : value.getUsers()) { - if (auto unrealized = dyn_cast(user)) { - for (Value result : unrealized.getResults()) { - if (auto policy = self(self, result); - policy != MxValidAlignPolicy::None) - return policy; - } - continue; - } - if (auto cast = dyn_cast(user)) { - if (auto policy = self(self, cast.getResult()); - policy != MxValidAlignPolicy::None) - return policy; - continue; - } - - if (auto tmatmul = dyn_cast(user)) { - auto policy = matchTMatmulMxLike(tmatmul); - if (policy != MxValidAlignPolicy::None) - return policy; - continue; - } - if (auto tmatmulAcc = dyn_cast(user)) { - auto policy = matchTMatmulMxLike(tmatmulAcc); - if (policy != MxValidAlignPolicy::None) - return policy; - continue; - } - if (auto tmatmulBias = dyn_cast(user)) { - auto policy = matchTMatmulMxLike(tmatmulBias); - if (policy != MxValidAlignPolicy::None) - return policy; - continue; - } - if (auto tgemv = dyn_cast(user)) { - auto policy = matchTGemvMxLike(tgemv); - if (policy != MxValidAlignPolicy::None) - return policy; - continue; - } - if (auto tgemvAcc = dyn_cast(user)) { - auto policy = matchTGemvMxLike(tgemvAcc); - if (policy != MxValidAlignPolicy::None) - return policy; - continue; - } - if (auto tgemvBias = dyn_cast(user)) { - auto policy = matchTGemvMxLike(tgemvBias); - if (policy != MxValidAlignPolicy::None) - return policy; - continue; - } - } - return MxValidAlignPolicy::None; - }; - return walkUsers(walkUsers, op.getResult()); - }; - - MxValidAlignPolicy mxValidAlignPolicy = inferMxValidAlignPolicy(); - - std::string blTok = "BLayout::RowMajor"; - if (auto blAttr = dyn_cast(configAttr.getBLayout())) { - if (static_cast(blAttr.getValue()) == 1) - blTok = "BLayout::ColMajor"; - } - pto::BLayout blayout = getTileBufBLayoutValue(configAttr); - - if (isSubView) { - auto subMrTy = dyn_cast(op.getSource().getType()); - auto subViewOp = op.getSource().getDefiningOp(); - if (subMrTy && subMrTy.getRank() >= 2 && subViewOp) { - int64_t subRows = subMrTy.getDimSize(0); - int64_t subCols = subMrTy.getDimSize(1); - SmallVector inheritedStrides; - int64_t inheritedOffset = ShapedType::kDynamic; - - if (!pto::isPTOFloat4PackedType(elemTy) && - subRows != ShapedType::kDynamic && - subCols != ShapedType::kDynamic && - succeeded(mlir::pto::getPTOMemRefStridesAndOffset( - subMrTy, inheritedStrides, inheritedOffset)) && - inheritedStrides.size() >= 2) { - int64_t childRowStride = 0; - int64_t childColStride = 0; - bool sameStrides = getTilePointerStrides( - configAttr, elemTy, subRows, subCols, childRowStride, - childColStride); - sameStrides = sameStrides && - inheritedStrides[0] == childRowStride && - inheritedStrides[1] == childColStride; - if (sameStrides) { - rows = subRows; - cols = subCols; - } - } - } - } - - std::string slTok = "SLayout::NoneBox"; - if (auto slAttr = dyn_cast(configAttr.getSLayout())) { - int32_t slVal = static_cast(slAttr.getValue()); - slTok = (slVal == 1) ? "SLayout::RowMajor" - : (slVal == 2) ? "SLayout::ColMajor" - : "SLayout::NoneBox"; - } - - int32_t fractal = 512; - if (auto frAttr = dyn_cast(configAttr.getSFractalSize())) - fractal = static_cast(getIntegerAttrSignedValue(frAttr)); - - std::string padTok = "PadValue::Null"; - if (auto padAttr = dyn_cast(configAttr.getPad())) { - switch (static_cast(padAttr.getValue())) { - case 1: - padTok = "PadValue::Zero"; - break; - case 2: - padTok = "PadValue::Max"; - break; - case 3: - padTok = "PadValue::Min"; - break; - default: - padTok = "PadValue::Null"; - break; - } - } - - std::string compactTok = "CompactMode::Null"; - if (auto compactAttr = dyn_cast(configAttr.getCompactMode())) { - switch (static_cast(compactAttr.getValue())) { - case 1: - compactTok = "CompactMode::Normal"; - break; - case 2: - compactTok = "CompactMode::RowPlusOne"; - break; - default: - compactTok = "CompactMode::Null"; - break; - } - } - - std::string vrowTok, vcolTok; - bool useConstructor = false; - bool rowIsDynamic = false; - bool colIsDynamic = false; - SmallVector constructorArgs; - - Value vRow = op.getValidRow(); - Value vCol = op.getValidCol(); - Value vRowEmitC = adaptor.getValidRow(); - Value vColEmitC = adaptor.getValidCol(); - bool forceDynamicValid = op->hasAttr(kForceDynamicValidShapeAttrName); - int64_t cRow = 0, cCol = 0; - bool rowIsConst = vRow && getIndexConst(vRow, cRow); - bool colIsConst = vCol && getIndexConst(vCol, cCol); - - auto makeCtorDimValue = [&](Value emitted, int64_t fallback) -> Value { - if (emitted) - return emitted; - return makeEmitCIntConstant( - rewriter, loc, emitc::OpaqueType::get(ctx, "int32_t"), fallback); - }; - auto maybeScaleDynamicValid = [&](Value emitted, int dimIdx) -> Value { - if (!emitted || !pto::isPTOFloat4PackedType(elemTy)) - return emitted; - int packedDim = blayout == pto::BLayout::ColMajor ? 0 : 1; - if (dimIdx != packedDim) - return emitted; - auto i32Ty = emitc::OpaqueType::get(ctx, "int32_t"); - Value two = makeEmitCIntConstant(rewriter, loc, i32Ty, 2); - return rewriter.create(loc, i32Ty, emitted, two).getResult(); - }; - - if (forceDynamicValid) { - vrowTok = "-1"; - vcolTok = "-1"; - useConstructor = true; - constructorArgs.push_back( - makeCtorDimValue(maybeScaleDynamicValid(vRowEmitC, 0), - renderTileTemplateDim(rowIsConst ? cRow : rows, - elemTy, blayout, 0))); - constructorArgs.push_back( - makeCtorDimValue(maybeScaleDynamicValid(vColEmitC, 1), - renderTileTemplateDim(colIsConst ? cCol : cols, - elemTy, blayout, 1))); - } else { - if (rowIsConst) { - vrowTok = std::to_string( - renderTileTemplateDim(cRow, elemTy, blayout, 0)); - } else if (vRow) { - vrowTok = "-1"; - rowIsDynamic = true; - useConstructor = true; - } else { - vrowTok = std::to_string( - renderTileTemplateDim(rows, elemTy, blayout, 0)); - } - - if (colIsConst) { - vcolTok = std::to_string( - renderTileTemplateDim(cCol, elemTy, blayout, 1)); - } else if (vCol) { - vcolTok = "-1"; - colIsDynamic = true; - useConstructor = true; - } else { - vcolTok = std::to_string( - renderTileTemplateDim(cols, elemTy, blayout, 1)); - } - - if (useConstructor) { - if (rowIsDynamic && vRowEmitC) - constructorArgs.push_back(maybeScaleDynamicValid(vRowEmitC, 0)); - if (colIsDynamic && vColEmitC) - constructorArgs.push_back(maybeScaleDynamicValid(vColEmitC, 1)); - } - } - - if (!forceDynamicValid) { - bool alignedByRole = false; - if (isLowPrecisionCubeOperandType(elemTy)) { - if (resultAS == pto::AddressSpace::LEFT && rows != 1 && !vcolTok.empty() && - vcolTok != "-1") { - vcolTok = std::to_string( - renderTileTemplateDim(cols, elemTy, blayout, 1)); - alignedByRole = true; - } else if (resultAS == pto::AddressSpace::RIGHT && !vrowTok.empty() && - vrowTok != "-1") { - vrowTok = std::to_string( - renderTileTemplateDim(rows, elemTy, blayout, 0)); - alignedByRole = true; - } - } - if (!alignedByRole) { - if (mxValidAlignPolicy == MxValidAlignPolicy::AlignRowToShape) - vrowTok = std::to_string( - renderTileTemplateDim(rows, elemTy, blayout, 0)); - if (mxValidAlignPolicy == MxValidAlignPolicy::AlignColToShape) - vcolTok = std::to_string( - renderTileTemplateDim(cols, elemTy, blayout, 1)); - } - } - - std::string tileTypeStr = std::string("Tile<") + roleTok + ", " + - elemTypeStr + ", " + - std::to_string(renderTileTemplateDim( - rows, elemTy, blayout, 0)) + - ", " + - std::to_string(renderTileTemplateDim( - cols, elemTy, blayout, 1)) + - ", " + blTok + - ", " + vrowTok + ", " + vcolTok + ", " + slTok + - ", " + std::to_string(fractal) + ", " + padTok + - ", " + compactTok + - ">"; - return TileBuildSpec{tileTypeStr, useConstructor, constructorArgs}; - }; - - auto buildTileValue = [&](const TileBuildSpec &spec, - bool forceDeclaration = false) -> Value { - auto tileType = emitc::OpaqueType::get(ctx, spec.tileTypeStr); - if (spec.useConstructor && !forceDeclaration) { - return rewriter - .create(loc, tileType, spec.tileTypeStr, - ArrayAttr{}, ArrayAttr{}, - ValueRange(spec.constructorArgs)) - .getResult(0); - } - - Value tile = rewriter - .create( - loc, getEmitCVariableResultType(tileType), - emitc::OpaqueAttr::get(ctx, "")) - .getResult(); - return loadEmitCVariableIfNeeded(rewriter, loc, tile); - }; - - auto emitElemTypeToString = [&](Type elemTy) -> std::string { - return getEmitCScalarTypeToken(elemTy); - }; - - auto buildIntegralAddress = [&](Value sourceValue) -> FailureOr { - auto u64Ty = emitc::OpaqueType::get(ctx, "uint64_t"); - auto rcU64 = - rewriter.getArrayAttr({emitc::OpaqueAttr::get(ctx, "uint64_t")}); - - Value rawPtr = sourceValue; - if (auto ot = dyn_cast(sourceValue.getType())) { - StringRef tyStr = ot.getValue(); - if (tyStr.contains("Tile<") || tyStr.contains("ConvTile<")) { - auto srcMrTy = dyn_cast(op.getSource().getType()); - if (!srcMrTy) - return failure(); - std::string elemTok = emitElemTypeToString(srcMrTy.getElementType()); - pto::AddressSpace as = pto::AddressSpace::GM; - if (auto asAttr = - dyn_cast_or_null(srcMrTy.getMemorySpace())) - as = asAttr.getAddressSpace(); - rawPtr = materializeTileDataValue(rewriter, loc, sourceValue, as, - elemTok); - } - } - - if (isSetFFTsPointerLikeType(rawPtr.getType())) { - return rewriter - .create(loc, u64Ty, "reinterpret_cast", - ArrayAttr{}, rcU64, ValueRange{rawPtr}) - .getResult(0); - } - - if (rawPtr.getType() == u64Ty) - return rawPtr; - return rewriter.create(loc, u64Ty, rawPtr).getResult(); - }; - - Value tileCandidate = peelAllCasts(adaptor.getSource()); - if (viewSemantics && viewSemantics.getValue() == "bitcast" && - isTileLike(tileCandidate)) { - FailureOr tileSpec = buildTileSpec(); - if (failed(tileSpec)) - return failure(); - Value dstTile = buildTileValue(*tileSpec); - FailureOr addr = buildIntegralAddress(tileCandidate); - if (failed(addr)) - return failure(); - - rewriter.create(loc, TypeRange{}, "TASSIGN", - ArrayAttr{}, ArrayAttr{}, - ValueRange{dstTile, *addr}); - rewriter.replaceOp(op, dstTile); - return success(); - } - - if (viewSemantics && viewSemantics.getValue() == "treshape" && - isTileLike(tileCandidate)) { - FailureOr tileSpec = buildTileSpec(); - if (failed(tileSpec)) - return failure(); - Value dstTile = buildTileValue(*tileSpec, /*forceDeclaration=*/true); - - rewriter.create(loc, TypeRange{}, "TRESHAPE", - ArrayAttr{}, ArrayAttr{}, - ValueRange{dstTile, tileCandidate}); - rewriter.replaceOp(op, dstTile); - return success(); - } - - // Subview origins are kept distinct from generic tile rebinding: - // even when source/destination C++ tile types match, subview may carry - // shifted base address semantics and should materialize a fresh handle. - if (isSubView) { - FailureOr tileSpec = buildTileSpec(); - if (failed(tileSpec)) - return failure(); - Value dstTile = buildTileValue(*tileSpec); - FailureOr addr = buildIntegralAddress(tileCandidate); - if (failed(addr)) - return failure(); - - rewriter.create(loc, TypeRange{}, "TASSIGN", - ArrayAttr{}, ArrayAttr{}, - ValueRange{dstTile, *addr}); - rewriter.replaceOp(op, dstTile); - return success(); - } - - // Generic tile-to-tile rebind path: preserve the same backing storage and - // rebuild a sibling tile with updated metadata/valid dims. - if (isTileLike(tileCandidate)) { - FailureOr tileSpec = buildTileSpec(); - if (failed(tileSpec)) - return failure(); - - if (!tileSpec->useConstructor) { - if (auto srcTy = dyn_cast(tileCandidate.getType())) { - if (srcTy.getValue() == tileSpec->tileTypeStr) { - rewriter.replaceOp(op, tileCandidate); - return success(); - } - } - } - - Value dstTile = buildTileValue(*tileSpec); - FailureOr addr = buildIntegralAddress(tileCandidate); - if (failed(addr)) - return failure(); - - rewriter.create(loc, TypeRange{}, "TASSIGN", - ArrayAttr{}, ArrayAttr{}, - ValueRange{dstTile, *addr}); - rewriter.replaceOp(op, dstTile); - return success(); - } - - SmallVector physAddrs; - Value source = op.getSource(); - - while (auto castOp = source.getDefiningOp()) - source = castOp.getOperand(0); - - if (auto upstreamCast = source.getDefiningOp()) { - auto upstreamOperands = upstreamCast.getAddrs(); - physAddrs.append(upstreamOperands.begin(), upstreamOperands.end()); - } else { - physAddrs.push_back(adaptor.getSource()); - } - - Value vRow = op.getValidRow(); - Value vCol = op.getValidCol(); - - auto newCast = rewriter.create( - loc, op.getType(), physAddrs, vRow ? vRow : Value(), - vCol ? vCol : Value(), configAttr); - if (viewSemantics) - newCast->setAttr("pto.view_semantics", viewSemantics); - if (op->hasAttr(kForceDynamicValidShapeAttrName)) - newCast->setAttr(kForceDynamicValidShapeAttrName, - op->getAttr(kForceDynamicValidShapeAttrName)); - rewriter.replaceOp(op, newCast.getResult()); - - return success(); - } -}; - struct PTOAllocTileToEmitC : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -14205,7 +13566,6 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx, "pto.set_flag_dyn", "set_flag"); diff --git a/lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp b/lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp index 636d82a64d..71256793aa 100644 --- a/lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp +++ b/lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp @@ -63,7 +63,7 @@ static LogicalResult eraseDeadVPTOMemRefScaffold(ModuleOp module) { if (!op->use_empty()) return; - if (isa(op)) deadOps.push_back(op); }); diff --git a/lib/PTO/Transforms/PTOValidateVPTOIR.cpp b/lib/PTO/Transforms/PTOValidateVPTOIR.cpp index f0778bd287..e6fafa37a8 100644 --- a/lib/PTO/Transforms/PTOValidateVPTOIR.cpp +++ b/lib/PTO/Transforms/PTOValidateVPTOIR.cpp @@ -343,7 +343,7 @@ class VPTOLegalityHelper { } static bool isResidualEmissionScaffold(Operation *op) { - return isa(op) || isTrivialEmissionCastPtr(op); } diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 7908d30d30..9120207880 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -185,11 +185,6 @@ static void lookupValidDims(Value v, Value &vRow, Value &vCol) { vCol = alloc.getValidCol(); return; } - if (auto bind = v.getDefiningOp()) { - vRow = bind.getValidRow(); - vCol = bind.getValidCol(); - return; - } if (auto pc = v.getDefiningOp()) { vRow = pc.getValidRow(); vCol = pc.getValidCol(); @@ -845,7 +840,7 @@ struct PTOViewToMemrefPass // // Explicit-address alloc_tile already carries the physical address for // level3. No-address alloc_tile is assigned an address by PTOPlanMemory. - // Neither case needs a pointer_cast + bind_tile detour here. + // Neither case needs a pointer_cast detour here. // ------------------------------------------------------------------ // ------------------------------------------------------------------ diff --git a/lib/PTO/Transforms/SlotAffineAnalysis.cpp b/lib/PTO/Transforms/SlotAffineAnalysis.cpp index ddc7847060..e24d26a2dd 100644 --- a/lib/PTO/Transforms/SlotAffineAnalysis.cpp +++ b/lib/PTO/Transforms/SlotAffineAnalysis.cpp @@ -30,10 +30,6 @@ Value findMultiTileSlotExpr(Value v) { return {}; if (auto get = dyn_cast(op)) return get.getSlot(); - if (auto bind = dyn_cast(op)) { - v = bind.getSource(); - continue; - } return {}; } return {}; diff --git a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp index 0c23b29cab..3d867822ab 100644 --- a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp +++ b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp @@ -317,13 +317,6 @@ static void bindExplicitValidDims(ShapeConstraintSolver &solver, dims.cols, alloc.getValidCol()); return; } - if (auto bind = value.getDefiningOp()) { - bindDimToValue(solver, symbolDimByValue, canonicalByValue, signatureMap, - dims.rows, bind.getValidRow()); - bindDimToValue(solver, symbolDimByValue, canonicalByValue, signatureMap, - dims.cols, bind.getValidCol()); - return; - } if (auto subview = value.getDefiningOp()) { bindDimToValue(solver, symbolDimByValue, canonicalByValue, signatureMap, dims.rows, subview.getValidRow()); diff --git a/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp b/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp index 35ad9fc966..9c52282a85 100644 --- a/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOFusionLoadStoreElision.cpp @@ -128,10 +128,6 @@ static Value getCanonicalTrackedValue(Value value) { if (!def) break; - if (auto bind = dyn_cast(def)) { - value = bind.getSource(); - continue; - } if (auto tileBufAddr = dyn_cast(def)) { value = tileBufAddr.getSrc(); continue; @@ -200,43 +196,6 @@ static Value getCanonicalTrackedValue(Value value) { return value; } -static bool normalizeFusionRegionYieldFrontier(pto::FusionRegionOp fusionRegion) { - Block &body = fusionRegion.getBody().front(); - auto yieldOp = dyn_cast(body.getTerminator()); - if (!yieldOp) - return false; - - bool changed = false; - for (auto [index, yielded] : llvm::enumerate(yieldOp.getValues())) { - auto bind = yielded.getDefiningOp(); - if (!bind) - continue; - - Value normalized = bind.getSource(); - if (!normalized || normalized == yielded) - continue; - - Value regionResult = fusionRegion.getResult(index); - Type originalResultType = regionResult.getType(); - - yieldOp->setOperand(index, normalized); - if (regionResult.getType() != normalized.getType()) - regionResult.setType(normalized.getType()); - - if (originalResultType != normalized.getType() && !regionResult.use_empty()) { - OpBuilder builder(fusionRegion); - builder.setInsertionPointAfter(fusionRegion); - auto rebound = builder.create( - bind.getLoc(), originalResultType, regionResult, bind.getValidRow(), - bind.getValidCol(), bind.getConfig()); - rebound->setAttrs(bind->getAttrDictionary()); - regionResult.replaceAllUsesExcept(rebound.getResult(), rebound); - } - changed = true; - } - return changed; -} - static Operation *getTopLevelAncestorInBlock(Operation *op, Block *block) { for (Operation *cur = op; cur; cur = cur->getParentOp()) if (cur->getBlock() == block) @@ -551,9 +510,6 @@ struct PTOFusionLoadStoreElisionPass return; bool changed = false; - func.walk([&](pto::FusionRegionOp fusionRegion) { - changed |= normalizeFusionRegionYieldFrontier(fusionRegion); - }); llvm::DenseMap regionContexts; func.walk([&](pto::FusionRegionOp fusionRegion) { diff --git a/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp b/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp index 58d68fcbf3..2784ecf4a2 100644 --- a/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp @@ -159,8 +159,6 @@ static Value traceAliasRootOneStep(Value value) { return cast.getSource(); if (auto transpose = dyn_cast(def)) return transpose.getIn(); - if (auto bind = dyn_cast(def)) - return bind.getSource(); if (auto subview = dyn_cast(def)) return subview.getSource(); if (auto bitcast = dyn_cast(def)) diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index 2f4c24a331..2ac03c8d39 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -167,8 +167,6 @@ std::optional> getOperationAliasInfo(Operation *op) { return std::make_pair(bitcastOp.getResult(), bitcastOp.getSrc()); } else if (auto reshapeOp = dyn_cast(op)) { return std::make_pair(reshapeOp.getResult(), reshapeOp.getSrc()); - } else if (auto bindTileOp = dyn_cast(op)) { - return std::make_pair(bindTileOp.getResult(), bindTileOp.getSource()); } else if (auto multiGetOp = dyn_cast(op)) { return std::make_pair(multiGetOp.getResult(), multiGetOp.getSource()); } else if (auto extSliceOp = dyn_cast(op)) { @@ -249,8 +247,6 @@ Value tracebackImpl(Value memrefVal) { } else if (auto op = dyn_cast(def)) { // trace back memref.alloc support scf.for result = op.getInitArgs()[cast(memrefVal).getResultNumber()]; - } else if (auto op = dyn_cast(def)) { - result = op.getSource(); } if (result) { diff --git a/lib/PTO/Transforms/VPTOBufferMaterialization.cpp b/lib/PTO/Transforms/VPTOBufferMaterialization.cpp index a12fc87718..cc0a965345 100644 --- a/lib/PTO/Transforms/VPTOBufferMaterialization.cpp +++ b/lib/PTO/Transforms/VPTOBufferMaterialization.cpp @@ -68,10 +68,6 @@ Value materializeBufferPointer(Value value, Type elementType, if (value.getType() == ptrType) return value; - if (auto bind = value.getDefiningOp()) - return materializeBufferPointer(bind.getSource(), elementType, memorySpace, - rewriter, loc); - if (auto cast = value.getDefiningOp()) { if (cast.getAddrs().empty()) return {}; diff --git a/lib/PTO/Transforms/VPTOPtrNormalize.cpp b/lib/PTO/Transforms/VPTOPtrNormalize.cpp index c9e5928142..7ea765a967 100644 --- a/lib/PTO/Transforms/VPTOPtrNormalize.cpp +++ b/lib/PTO/Transforms/VPTOPtrNormalize.cpp @@ -113,7 +113,7 @@ static LogicalResult computeSubviewElementOffset(memref::SubViewOp op, if (failed(mlir::pto::getPTOMemRefStridesAndOffset(sourceType, strides, baseOffset))) return failure(); - // The SSA source already names the base address after bind_tile/pointer_cast + // The SSA source already names the base address after pointer_cast // normalization. A dynamic memref layout offset here is metadata we can // ignore for ptr normalization and model as zero. if (baseOffset == ShapedType::kDynamic) @@ -219,9 +219,6 @@ static Value materializeScalarAccessPtr(Value source, PatternRewriter &rewriter, return rewriter.create(loc, ptrType, basePtr, offset); } - if (auto bind = source.getDefiningOp()) - return materializeScalarAccessPtr(bind.getSource(), rewriter, loc); - if (auto pointerCast = source.getDefiningOp()) { if (pointerCast.getAddrs().empty()) return {}; @@ -392,31 +389,6 @@ struct ConvertCastPtrPattern : public OpConversionPattern { } }; -struct ConvertBindTileToPtrPattern : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(pto::BindTileOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - Type convertedType = getTypeConverter()->convertType(op.getResult().getType()); - auto ptrType = dyn_cast(convertedType); - if (!ptrType) - return failure(); - - Value ptr = - materializeSubviewInputPtr(adaptor.getSource(), rewriter, op.getLoc()); - if (!ptr) - return rewriter.notifyMatchFailure(op, - "failed to materialize bind_tile input ptr"); - - if (ptr.getType() != ptrType) - ptr = rewriter.create(op.getLoc(), ptrType, ptr); - - rewriter.replaceOp(op, ptr); - return success(); - } -}; - struct ConvertSubviewToAddPtrPattern : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -854,7 +826,7 @@ struct VPTOPtrNormalizePass scf::SCFDialect>(); target.addDynamicallyLegalDialect([](Operation *op) { return !isa(op); }); target.addLegalOp(); @@ -888,10 +860,6 @@ struct VPTOPtrNormalizePass return !isMemRefType(op.getInput().getType()) && !isMemRefType(op.getResult().getType()); }); - target.addDynamicallyLegalOp([&](pto::BindTileOp op) { - return op.getResult().getType() == - typeConverter.convertType(op.getResult().getType()); - }); target.addDynamicallyLegalOp( [&](pto::VldsOp op) { return isa(op.getSource().getType()) && @@ -985,7 +953,6 @@ struct VPTOPtrNormalizePass ConvertPointerCastToCastPtrPattern, ConvertIntToPtrToCastPtrPattern, ConvertPtrToIntToCastPtrPattern, ConvertCastPtrPattern, - ConvertBindTileToPtrPattern, ConvertSubviewToAddPtrPattern, ConvertVldsSubviewOperandPattern, ConvertVstsSubviewOperandPattern, ConvertVsstbSubviewOperandPattern, diff --git a/test/lit/pto/add_carry_tile_native.pto b/test/lit/pto/add_carry_tile_native.pto index 6b5d337527..7146fb03fe 100644 --- a/test/lit/pto/add_carry_tile_native.pto +++ b/test/lit/pto/add_carry_tile_native.pto @@ -39,7 +39,6 @@ module { // NATIVE-LABEL: func.func private @taddsc_arg( // NATIVE: pto.taddsc ins(%arg0, %arg1, %arg2 : !pto.tile_buf, f32, !pto.tile_buf) outs(%arg3 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: taddc_arg( // EMITC-COUNT-2: TADD( diff --git a/test/lit/pto/addmul_scalar_tile_native.pto b/test/lit/pto/addmul_scalar_tile_native.pto index 65f0a83039..ed75b9bf4c 100644 --- a/test/lit/pto/addmul_scalar_tile_native.pto +++ b/test/lit/pto/addmul_scalar_tile_native.pto @@ -46,7 +46,6 @@ module { // NATIVE: pto.tmuls ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) // NATIVE: pto.tmuls ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg0 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tadds_arg( // EMITC-COUNT-2: TADDS( diff --git a/test/lit/pto/addmul_tile_native.pto b/test/lit/pto/addmul_tile_native.pto index a26a71029e..eac26c4a74 100644 --- a/test/lit/pto/addmul_tile_native.pto +++ b/test/lit/pto/addmul_tile_native.pto @@ -46,7 +46,6 @@ module { // NATIVE: pto.tmul ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) // NATIVE: pto.tmul ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg0 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tadd_arg( // EMITC-COUNT-2: TADD( diff --git a/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto b/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto index 55a835120e..72437129f2 100644 --- a/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto +++ b/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto @@ -31,5 +31,4 @@ module attributes {"pto.target_arch" = "a3"} { // PLAN: %[[ADDR:.*]] = arith.constant 0 : i64 // PLAN-NEXT: {{%.*}} = pto.alloc_tile addr = %[[ADDR]] // PLAN-NOT: pto.pointer_cast -// PLAN-NOT: pto.bind_tile // PLAN-NOT: memref.alloc diff --git a/test/lit/pto/arg_reduction_tile_native.pto b/test/lit/pto/arg_reduction_tile_native.pto index 3a54887913..6f8e728f25 100644 --- a/test/lit/pto/arg_reduction_tile_native.pto +++ b/test/lit/pto/arg_reduction_tile_native.pto @@ -49,7 +49,6 @@ module { // NATIVE-LABEL: @tpartargmin_arg // NATIVE: pto.tpartargmin // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tcolargmax_arg( // EMITC: TCOLARGMAX( diff --git a/test/lit/pto/axpy_dequant_tile_native.pto b/test/lit/pto/axpy_dequant_tile_native.pto index d43dbc27da..9efd14142b 100644 --- a/test/lit/pto/axpy_dequant_tile_native.pto +++ b/test/lit/pto/axpy_dequant_tile_native.pto @@ -25,7 +25,6 @@ module { // NATIVE-LABEL: @tdequant_arg // NATIVE: pto.tdequant // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: taxpy_arg( // EMITC: TAXPY( diff --git a/test/lit/pto/basic_float_tile_native.pto b/test/lit/pto/basic_float_tile_native.pto index 35de5643af..4788a0812b 100644 --- a/test/lit/pto/basic_float_tile_native.pto +++ b/test/lit/pto/basic_float_tile_native.pto @@ -48,7 +48,6 @@ module { // NATIVE-SAME: !pto.tile_buf // NATIVE: pto.tfmods // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: trelu_arg( // EMITC: TRELU( diff --git a/test/lit/pto/bias_tile_subview_emitc_pointer_type.pto b/test/lit/pto/bias_tile_subview_emitc_pointer_type.pto index 4c3d698d9c..f381d7ce46 100644 --- a/test/lit/pto/bias_tile_subview_emitc_pointer_type.pto +++ b/test/lit/pto/bias_tile_subview_emitc_pointer_type.pto @@ -3,15 +3,10 @@ module { func.func @bias_tile_subview_emitc_pointer_type() attributes {pto.kernel_kind = #pto.kernel_kind} { - %c1 = arith.constant 1 : index %c32 = arith.constant 32 : index - %c64 = arith.constant 64 : index %bias = memref.alloc() : memref<1x64xf32, #pto.address_space> - %bias_view = pto.bind_tile %bias, %c1, %c64 - {config = #pto.tile_buf_config} - : memref<1x64xf32, #pto.address_space> -> memref<1x64xf32, #pto.address_space> - %bias_sub = memref.subview %bias_view[0, 32] [1, 32] [1, 1] + %bias_sub = memref.subview %bias[0, 32] [1, 32] [1, 1] : memref<1x64xf32, #pto.address_space> to memref<1x32xf32, strided<[64, 1], offset: 32>, #pto.address_space> call @consume_bias(%bias_sub) : (memref<1x32xf32, strided<[64, 1], offset: 32>, #pto.address_space>) -> () diff --git a/test/lit/pto/bitcast_tile_native.pto b/test/lit/pto/bitcast_tile_native.pto index 82aa068367..55581dc82b 100644 --- a/test/lit/pto/bitcast_tile_native.pto +++ b/test/lit/pto/bitcast_tile_native.pto @@ -37,4 +37,3 @@ module attributes {pto.target_arch = "a5"} { // NATIVE-SAME: !pto.tile_buf -> !pto.tile_buf // NATIVE: return %[[RESULT]] : !pto.tile_buf // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile diff --git a/test/lit/pto/build_async_session_tile_native.pto b/test/lit/pto/build_async_session_tile_native.pto index e4f92f08d2..9c43766337 100644 --- a/test/lit/pto/build_async_session_tile_native.pto +++ b/test/lit/pto/build_async_session_tile_native.pto @@ -30,7 +30,6 @@ module { // NATIVE: pto.comm.build_async_session(%[[SCRATCH]], %[[WORKSPACE]] // NATIVE-SAME: !pto.tile_buf, !pto.ptr // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: AICORE void build_async_session_tile_native( // EMITC: pto::comm::BuildAsyncSession( diff --git a/test/lit/pto/col_minmax_tile_native.pto b/test/lit/pto/col_minmax_tile_native.pto index 5be86f1bf5..24e7a0678e 100644 --- a/test/lit/pto/col_minmax_tile_native.pto +++ b/test/lit/pto/col_minmax_tile_native.pto @@ -35,7 +35,6 @@ module { // NATIVE-LABEL: func.func private @tcolmin_arg( // NATIVE: pto.tcolmin ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tcolmax_arg( // EMITC: TCOLMAX( diff --git a/test/lit/pto/colexpand_tile_native.pto b/test/lit/pto/colexpand_tile_native.pto index a30a107fdd..a252693183 100644 --- a/test/lit/pto/colexpand_tile_native.pto +++ b/test/lit/pto/colexpand_tile_native.pto @@ -91,7 +91,6 @@ module { // NATIVE-LABEL: func.func private @tcolexpandexpdif_arg( // NATIVE: pto.tcolexpandexpdif // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tcolexpand_arg( // EMITC: TCOLEXPAND( diff --git a/test/lit/pto/colsum_tile_native.pto b/test/lit/pto/colsum_tile_native.pto index 252428a442..4cb9c94017 100644 --- a/test/lit/pto/colsum_tile_native.pto +++ b/test/lit/pto/colsum_tile_native.pto @@ -36,7 +36,6 @@ module { // NATIVE-LABEL: func.func private @tcolsum_binary_arg( // NATIVE: pto.tcolsum ins(%arg0, %arg1 {isBinary = true} : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tcolsum_arg( // EMITC: TCOLSUM( diff --git a/test/lit/pto/compare_tile_native.pto b/test/lit/pto/compare_tile_native.pto index 2728956cf8..1b09451969 100644 --- a/test/lit/pto/compare_tile_native.pto +++ b/test/lit/pto/compare_tile_native.pto @@ -37,7 +37,6 @@ module { // NATIVE-LABEL: func.func private @tcmps_arg( // NATIVE: pto.tcmps ins(%arg0, %arg1 {cmpMode = #pto} : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tcmp_arg( // EMITC: TCMP( diff --git a/test/lit/pto/concat_tile_native.pto b/test/lit/pto/concat_tile_native.pto index 588116a7bb..81cf638414 100644 --- a/test/lit/pto/concat_tile_native.pto +++ b/test/lit/pto/concat_tile_native.pto @@ -39,7 +39,6 @@ module { // NATIVE-LABEL: func.func private @tconcatidx_arg( // NATIVE: pto.tconcatidx ins(%arg0, %arg1, %arg2, %arg3 : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%arg4 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tconcat_arg( // EMITC: TCONCAT( diff --git a/test/lit/pto/cvt_tile_native.pto b/test/lit/pto/cvt_tile_native.pto index f711431ce6..c65d57e1cf 100644 --- a/test/lit/pto/cvt_tile_native.pto +++ b/test/lit/pto/cvt_tile_native.pto @@ -25,7 +25,6 @@ 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-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tcvt_arg( // EMITC: TCVT( diff --git a/test/lit/pto/declare_tile_tile_native.pto b/test/lit/pto/declare_tile_tile_native.pto index 0a5d5f6bf9..cea347f13a 100644 --- a/test/lit/pto/declare_tile_tile_native.pto +++ b/test/lit/pto/declare_tile_tile_native.pto @@ -41,7 +41,6 @@ module { // NATIVE-LABEL: func.func @declare_tile_native // NATIVE: pto.declare_tile -> !pto.tile_buf // NATIVE: pto.tpop({{.*}} : !pto.tile_buf -// NATIVE-NOT: pto.bind_tile // NATIVE-NOT: memref.alloc // EMITC-LABEL: AICORE void declare_tile_native diff --git a/test/lit/pto/div_tile_native.pto b/test/lit/pto/div_tile_native.pto index dc606c661d..0868d36a9c 100644 --- a/test/lit/pto/div_tile_native.pto +++ b/test/lit/pto/div_tile_native.pto @@ -51,7 +51,6 @@ module { // NATIVE-LABEL: func.func private @tdivs_scalar_tile_arg( // NATIVE: pto.tdivs ins(%arg0, %arg1 : f32, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) {precisionType = #pto} // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tdiv_arg( // EMITC: TDIV( diff --git a/test/lit/pto/expands_tile_native.pto b/test/lit/pto/expands_tile_native.pto index 680ec473c5..1d16f49493 100644 --- a/test/lit/pto/expands_tile_native.pto +++ b/test/lit/pto/expands_tile_native.pto @@ -25,7 +25,6 @@ module { // NATIVE-LABEL: func.func private @texpands_arg( // NATIVE: pto.texpands ins(%arg0 : f32) outs(%arg1 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: texpands_arg( // EMITC: TEXPANDS( diff --git a/test/lit/pto/extract_insert_fp_tile_native.pto b/test/lit/pto/extract_insert_fp_tile_native.pto index e3c740b1ed..cbbeb7514e 100644 --- a/test/lit/pto/extract_insert_fp_tile_native.pto +++ b/test/lit/pto/extract_insert_fp_tile_native.pto @@ -39,7 +39,6 @@ module { // NATIVE-LABEL: func.func private @tinsert_fp_arg( // NATIVE: pto.tinsert_fp ins(%arg0, %arg1, %arg2, %arg3 // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: textract_fp_arg( // EMITC: TEXTRACT_FP( diff --git a/test/lit/pto/extract_insert_tile_native.pto b/test/lit/pto/extract_insert_tile_native.pto index 86a1d8ecfc..c5e4ba7fd2 100644 --- a/test/lit/pto/extract_insert_tile_native.pto +++ b/test/lit/pto/extract_insert_tile_native.pto @@ -39,7 +39,6 @@ module { // NATIVE-LABEL: func.func private @tinsert_arg( // NATIVE: pto.tinsert ins(%arg0, %arg1, %arg2 : !pto.tile_buf, index, index) outs(%arg3 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: textract_arg( // EMITC: TEXTRACT( diff --git a/test/lit/pto/fillpad_tile_native.pto b/test/lit/pto/fillpad_tile_native.pto index 081127f5ca..535d264920 100644 --- a/test/lit/pto/fillpad_tile_native.pto +++ b/test/lit/pto/fillpad_tile_native.pto @@ -33,7 +33,6 @@ module { // NATIVE-LABEL: func.func private @tfillpad_inplace_arg( // NATIVE: pto.tfillpad_inplace ins(%arg0 // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tfillpad_arg( // EMITC: TFILLPAD( diff --git a/test/lit/pto/gather_tile_native.pto b/test/lit/pto/gather_tile_native.pto index 579e00f151..8ebc4a33f8 100644 --- a/test/lit/pto/gather_tile_native.pto +++ b/test/lit/pto/gather_tile_native.pto @@ -48,7 +48,6 @@ module { // NATIVE-LABEL: func.func private @tgatherb_arg( // NATIVE: pto.tgatherb ins(%arg0, %arg1 // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tgather_mask_arg( // EMITC: TGATHER< diff --git a/test/lit/pto/gemv_tile_native.pto b/test/lit/pto/gemv_tile_native.pto index 4e506374dc..1134f71fc0 100644 --- a/test/lit/pto/gemv_tile_native.pto +++ b/test/lit/pto/gemv_tile_native.pto @@ -22,7 +22,6 @@ module { // NATIVE-SAME: !pto.tile_buf>) attributes {pto.kernel_kind = #pto.kernel_kind} { diff --git a/test/lit/pto/load_store_tile_native.pto b/test/lit/pto/load_store_tile_native.pto index ad73bd6915..2620bfc1a4 100644 --- a/test/lit/pto/load_store_tile_native.pto +++ b/test/lit/pto/load_store_tile_native.pto @@ -34,7 +34,6 @@ module { // NATIVE: pto.tload ins(%arg0 : memref<16x16xf32>) outs(%arg1 : !pto.tile_buf) // NATIVE-LABEL: func.func private @tstore_arg( // NATIVE: pto.tstore ins(%arg0 : !pto.tile_buf) outs(%arg1 : memref<16x16xf32>) -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tload_arg( // EMITC: TLOAD( diff --git a/test/lit/pto/logic_binary_tile_native.pto b/test/lit/pto/logic_binary_tile_native.pto index 147bfda9ae..7a29ee6e5e 100644 --- a/test/lit/pto/logic_binary_tile_native.pto +++ b/test/lit/pto/logic_binary_tile_native.pto @@ -43,7 +43,6 @@ module { // NATIVE: pto.tor ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) // NATIVE: pto.tor ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg0 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tand_arg( // EMITC-COUNT-2: TAND( diff --git a/test/lit/pto/logic_scalar_tile_native.pto b/test/lit/pto/logic_scalar_tile_native.pto index dc6c457f68..2d55a6ec0a 100644 --- a/test/lit/pto/logic_scalar_tile_native.pto +++ b/test/lit/pto/logic_scalar_tile_native.pto @@ -37,7 +37,6 @@ module { // NATIVE-LABEL: func.func private @tors_arg( // NATIVE: pto.tors ins(%arg0, %arg1 : !pto.tile_buf, i16) outs(%arg2 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tands_arg( // EMITC: TANDS( diff --git a/test/lit/pto/matmul_tile_native.pto b/test/lit/pto/matmul_tile_native.pto index 574ff3af1c..62745e4e36 100644 --- a/test/lit/pto/matmul_tile_native.pto +++ b/test/lit/pto/matmul_tile_native.pto @@ -22,7 +22,6 @@ module { // NATIVE-SAME: !pto.tile_buf // NATIVE-SAME: memref<1x1x1x8x16xf16, #pto.address_space> // NATIVE: pto.mscatter -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: mgather_arg( // EMITC: MGATHER( diff --git a/test/lit/pto/minmax_scalar_tile_native.pto b/test/lit/pto/minmax_scalar_tile_native.pto index 7f053bb34c..5b9c6a22ea 100644 --- a/test/lit/pto/minmax_scalar_tile_native.pto +++ b/test/lit/pto/minmax_scalar_tile_native.pto @@ -46,7 +46,6 @@ module { // NATIVE: pto.tmins ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) // NATIVE: pto.tmins ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg0 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tmaxs_arg( // EMITC-COUNT-2: TMAXS( diff --git a/test/lit/pto/minmax_tile_native.pto b/test/lit/pto/minmax_tile_native.pto index 8b9db4389e..deb676e216 100644 --- a/test/lit/pto/minmax_tile_native.pto +++ b/test/lit/pto/minmax_tile_native.pto @@ -46,7 +46,6 @@ module { // NATIVE: pto.tmin ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) // NATIVE: pto.tmin ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg0 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tmax_arg( // EMITC-COUNT-2: TMAX( diff --git a/test/lit/pto/mov_tile_native.pto b/test/lit/pto/mov_tile_native.pto index b8b77e809e..71ba65ad91 100644 --- a/test/lit/pto/mov_tile_native.pto +++ b/test/lit/pto/mov_tile_native.pto @@ -20,7 +20,6 @@ module { // NATIVE-SAME: !pto.tile_buf // NATIVE: pto.tmov // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tmov_arg( // EMITC: TMOV( diff --git a/test/lit/pto/movement_metadata_tile_native.pto b/test/lit/pto/movement_metadata_tile_native.pto index 4c0e74eee2..1b4999c9c1 100644 --- a/test/lit/pto/movement_metadata_tile_native.pto +++ b/test/lit/pto/movement_metadata_tile_native.pto @@ -25,7 +25,6 @@ module { // NATIVE-LABEL: @tget_scale_addr_arg // NATIVE: pto.tget_scale_addr // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tfillpad_expand_arg( // EMITC: TFILLPAD_EXPAND( diff --git a/test/lit/pto/mrgsort_tile_native.pto b/test/lit/pto/mrgsort_tile_native.pto index 5fa7f960c4..fcba49ef06 100644 --- a/test/lit/pto/mrgsort_tile_native.pto +++ b/test/lit/pto/mrgsort_tile_native.pto @@ -39,7 +39,6 @@ module { // NATIVE-LABEL: func.func private @tmrgsort_format2_arg // NATIVE: pto.tmrgsort ins(%{{.*}}, %{{.*}}, %{{.*}} {exhausted = true} : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%{{.*}}, %{{.*}} : !pto.tile_buf, vector<4xi16>) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tmrgsort_format1_arg( // EMITC: TMRGSORT( diff --git a/test/lit/pto/part_arithmetic_tile_native.pto b/test/lit/pto/part_arithmetic_tile_native.pto index 2282b5573e..085a223b9c 100644 --- a/test/lit/pto/part_arithmetic_tile_native.pto +++ b/test/lit/pto/part_arithmetic_tile_native.pto @@ -59,7 +59,6 @@ module { // NATIVE-LABEL: func.func private @tpartmin_arg( // NATIVE: pto.tpartmin ins(%arg0, %arg1 : !pto.tile_buf, !pto.tile_buf) outs(%arg2 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tpartadd_arg( // EMITC: TPARTADD( diff --git a/test/lit/pto/pipe_tile_native.pto b/test/lit/pto/pipe_tile_native.pto index 91fb05b3e0..c187386d6d 100644 --- a/test/lit/pto/pipe_tile_native.pto +++ b/test/lit/pto/pipe_tile_native.pto @@ -26,10 +26,8 @@ module { // NATIVE-LABEL: @push_arg // NATIVE-SAME: !pto.tile_buf&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @bind_tile_alias_liveness(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - %c16 = arith.constant 16 : index - - %a = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %av = pto.bind_tile %a, %c16, %c16 - {config = #pto.tile_buf_config} - : memref<16x16x16xf16, #pto.address_space> -> memref<16x16x16xf16, #pto.address_space> - - %b = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%b : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%b : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - // Using %av should keep %a live; %b must not reuse %a's offset. - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%av : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%av : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @bind_tile_alias_liveness -// CHECK-NOT: memref.alloc -// CHECK-DAG: %c0_i64 = arith.constant 0 : i64 -// CHECK-DAG: %c8192_i64 = arith.constant 8192 : i64 -// CHECK-DAG: pto.pointer_cast(%c0_i64){{.*}} : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK-DAG: pto.pointer_cast(%c8192_i64){{.*}} : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> diff --git a/test/lit/pto/pow_rsqrt_tile_native.pto b/test/lit/pto/pow_rsqrt_tile_native.pto index b2a99514e6..61306ba443 100644 --- a/test/lit/pto/pow_rsqrt_tile_native.pto +++ b/test/lit/pto/pow_rsqrt_tile_native.pto @@ -31,7 +31,6 @@ module { // NATIVE-LABEL: @trsqrt_arg // NATIVE: pto.trsqrt // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tpow_arg( // EMITC: TPOW( diff --git a/test/lit/pto/print_tile_native.pto b/test/lit/pto/print_tile_native.pto index 84247d2a5f..dec7f9d0ef 100644 --- a/test/lit/pto/print_tile_native.pto +++ b/test/lit/pto/print_tile_native.pto @@ -32,7 +32,6 @@ module { // NATIVE: pto.tprint ins(%arg0 : !pto.tile_buf) // NATIVE-LABEL: func.func private @tprint_tmp_arg( // NATIVE: pto.tprint ins(%arg0, %arg1 : !pto.tile_buf, memref<16x16xf32>) -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tprint_arg( // EMITC: TPRINT( diff --git a/test/lit/pto/quant_mx_tile_native.pto b/test/lit/pto/quant_mx_tile_native.pto index 2d154a36d5..5ee9681dd9 100644 --- a/test/lit/pto/quant_mx_tile_native.pto +++ b/test/lit/pto/quant_mx_tile_native.pto @@ -20,7 +20,6 @@ module { // NATIVE-SAME: !pto.tile_buf< // NATIVE: pto.tquant.mx // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tquant_mx_arg( // EMITC: TQUANT&1 | FileCheck %s --check-prefix=MATERIALIZE // RUN: not ptoas %t/declare.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=DECLARE // RUN: not ptoas %t/slot.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=SLOT +// RUN: not ptoas %t/bind.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=BIND //--- materialize.pto module { @@ -40,3 +41,13 @@ module { } // SLOT: error: unregistered operation 'pto.slot_marker' found in dialect ('pto') + +//--- bind.pto +module { + func.func @bind() { + %0 = "pto.bind_tile"() : () -> i32 + return + } +} + +// BIND: error: unregistered operation 'pto.bind_tile' found in dialect ('pto') diff --git a/test/lit/pto/rowexpand_tile_native.pto b/test/lit/pto/rowexpand_tile_native.pto index 1eb14eaeaf..26ed6b0735 100644 --- a/test/lit/pto/rowexpand_tile_native.pto +++ b/test/lit/pto/rowexpand_tile_native.pto @@ -62,7 +62,6 @@ module { // NATIVE-LABEL: @trowexpandexpdif_arg // NATIVE: pto.trowexpandexpdif // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: trowexpand_arg( // EMITC: TROWEXPAND( diff --git a/test/lit/pto/scratch_elementwise_tile_native.pto b/test/lit/pto/scratch_elementwise_tile_native.pto index af45853b52..32826f0de6 100644 --- a/test/lit/pto/scratch_elementwise_tile_native.pto +++ b/test/lit/pto/scratch_elementwise_tile_native.pto @@ -31,7 +31,6 @@ module { // NATIVE-LABEL: @trems_arg // NATIVE: pto.trems // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tprelu_arg( // EMITC: TPRELU( diff --git a/test/lit/pto/select_tile_native.pto b/test/lit/pto/select_tile_native.pto index c79fde6bc4..a0501b5fd9 100644 --- a/test/lit/pto/select_tile_native.pto +++ b/test/lit/pto/select_tile_native.pto @@ -25,7 +25,6 @@ module { // NATIVE-LABEL: @tsels_arg // NATIVE: pto.tsels // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tsel_arg( // EMITC: TSEL( diff --git a/test/lit/pto/shift_tile_native.pto b/test/lit/pto/shift_tile_native.pto index 0470930f7f..84a3dbabe2 100644 --- a/test/lit/pto/shift_tile_native.pto +++ b/test/lit/pto/shift_tile_native.pto @@ -41,7 +41,6 @@ module { // NATIVE-SAME: !pto.tile_buf // NATIVE: pto.tshrs // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tshl_arg( // EMITC: TSHL( diff --git a/test/lit/pto/sort32_tile_native.pto b/test/lit/pto/sort32_tile_native.pto index d3d0648a85..4740d2e397 100644 --- a/test/lit/pto/sort32_tile_native.pto +++ b/test/lit/pto/sort32_tile_native.pto @@ -25,7 +25,6 @@ module { // NATIVE-LABEL: @tsort32_tmp_arg // NATIVE: pto.tsort32 // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tsort32_arg( // EMITC: TSORT32( diff --git a/test/lit/pto/store_fp_tile_native.pto b/test/lit/pto/store_fp_tile_native.pto index 72b30a402d..26c7974d2a 100644 --- a/test/lit/pto/store_fp_tile_native.pto +++ b/test/lit/pto/store_fp_tile_native.pto @@ -21,7 +21,6 @@ module { // NATIVE-SAME: !pto.tile_buf> // NATIVE: pto.tstore_fp -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tstore_fp_arg( // EMITC: TSTORE_FP( diff --git a/test/lit/pto/sub_tile_native.pto b/test/lit/pto/sub_tile_native.pto index 676fa74573..adfbaa648a 100644 --- a/test/lit/pto/sub_tile_native.pto +++ b/test/lit/pto/sub_tile_native.pto @@ -44,7 +44,6 @@ module { // NATIVE-SAME: !pto.tile_buf // NATIVE: pto.tsubsc // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tsub_arg( // EMITC: TSUB( diff --git a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto index f9f92a74b3..1646298be1 100644 --- a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto +++ b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto @@ -32,4 +32,3 @@ module { // NATIVE: IR Dump After PTOViewToMemref // NATIVE: pto.subview %{{.*}}[%c0, %{{.*}}] sizes [1, 64] : !pto.tile_buf -> !pto.tile_buf // NATIVE-NOT: memref.subview -// NATIVE-NOT: pto.bind_tile diff --git a/test/lit/pto/subview_bind_tile_preserve_stride.pto b/test/lit/pto/subview_tile_native_preserve_stride.pto similarity index 98% rename from test/lit/pto/subview_bind_tile_preserve_stride.pto rename to test/lit/pto/subview_tile_native_preserve_stride.pto index 9e2ca7168f..2d3cb2a329 100644 --- a/test/lit/pto/subview_bind_tile_preserve_stride.pto +++ b/test/lit/pto/subview_tile_native_preserve_stride.pto @@ -1,7 +1,7 @@ // RUN: ptoas %s 2>&1 | FileCheck %s module { - func.func @subview_bind_tile_preserve_stride( + func.func @subview_tile_native_preserve_stride( %src: memref<16x16xf32, #pto.address_space>, %dst0: memref<8x8xf32, #pto.address_space>, %dst1: memref<8x8xf32, #pto.address_space>, diff --git a/test/lit/pto/tassign_tile_native_arg.pto b/test/lit/pto/tassign_tile_native_arg.pto index 29b945d619..442c522729 100644 --- a/test/lit/pto/tassign_tile_native_arg.pto +++ b/test/lit/pto/tassign_tile_native_arg.pto @@ -26,7 +26,6 @@ module { // NATIVE-SAME: (%[[TILE:.*]]: !pto.tile_buf) // NATIVE: %{{.*}} = pto.tassign %[[TILE]], %{{.*}} : !pto.tile_buf -> !pto.tile_buf // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tassign_arg( // EMITC: TASSIGN( diff --git a/test/lit/pto/tile_buf_addr_tile_native_arg.pto b/test/lit/pto/tile_buf_addr_tile_native_arg.pto index 76588d3ed7..9ee7e6e987 100644 --- a/test/lit/pto/tile_buf_addr_tile_native_arg.pto +++ b/test/lit/pto/tile_buf_addr_tile_native_arg.pto @@ -28,7 +28,6 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { // NATIVE: %[[PTR:.*]] = pto.tile_buf_addr %[[TILE]] : !pto.tile_buf -> !pto.ptr // NATIVE: call @consume(%[[PTR]]) : (!pto.ptr) -> () // NATIVE-NOT: memref.reinterpret_cast -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: AICORE void tile_buf_addr_tile_native_arg( // EMITC: __ubuf__ float* {{.*}} = {{.*}}.data(); diff --git a/test/lit/pto/tile_scalar_access_tile_native.pto b/test/lit/pto/tile_scalar_access_tile_native.pto index d201f2c382..9e4315946c 100644 --- a/test/lit/pto/tile_scalar_access_tile_native.pto +++ b/test/lit/pto/tile_scalar_access_tile_native.pto @@ -38,7 +38,6 @@ module { // NATIVE: %0 = pto.tgetval ins(%arg0, %arg1 : !pto.tile_buf, index) outs : f32 // NATIVE: return %0 : f32 // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tsetval_arg( // EMITC: .SetValue( diff --git a/test/lit/pto/tlrelu_tile_native.pto b/test/lit/pto/tlrelu_tile_native.pto index 9324743b06..cfbea9cd5c 100644 --- a/test/lit/pto/tlrelu_tile_native.pto +++ b/test/lit/pto/tlrelu_tile_native.pto @@ -29,7 +29,6 @@ module { // NATIVE-SAME: %arg2: !pto.tile_buf // NATIVE: pto.tlrelu ins(%arg0, %arg1 : !pto.tile_buf, f32) outs(%arg2 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tlrelu_arg( // EMITC: TLRELU( diff --git a/test/lit/pto/tmov_fp_tile_native.pto b/test/lit/pto/tmov_fp_tile_native.pto index 4e2822406c..1ebdb7912d 100644 --- a/test/lit/pto/tmov_fp_tile_native.pto +++ b/test/lit/pto/tmov_fp_tile_native.pto @@ -26,7 +26,6 @@ module { // NATIVE-LABEL: func.func private @tmov_fp_arg( // NATIVE: pto.tmov.fp ins(%arg0, %arg1 // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tmov_fp_arg( // EMITC: TMOV_FP< diff --git a/test/lit/pto/tprefetch_tile_native.pto b/test/lit/pto/tprefetch_tile_native.pto index e630565932..16aea07fdc 100644 --- a/test/lit/pto/tprefetch_tile_native.pto +++ b/test/lit/pto/tprefetch_tile_native.pto @@ -21,7 +21,6 @@ module { // NATIVE-SAME: %{{.*}}: memref<16x16xf16, #pto.address_space> // NATIVE-SAME: %{{.*}}: !pto.tile_buf, %[[TMP:[^:]+]]: !pto.tile_buf, %[[DST:[^:]+]]: !pto.tile_buf) // NATIVE: pto.ttrans ins(%[[SRC]], %[[TMP]] : !pto.tile_buf, !pto.tile_buf) outs(%[[DST]] : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: ttrans_arg( // EMITC: TTRANS( diff --git a/test/lit/pto/transcendental_unary_tile_native.pto b/test/lit/pto/transcendental_unary_tile_native.pto index 37e90e4527..233b910669 100644 --- a/test/lit/pto/transcendental_unary_tile_native.pto +++ b/test/lit/pto/transcendental_unary_tile_native.pto @@ -39,7 +39,6 @@ module { // NATIVE-LABEL: func.func private @tlog_arg( // NATIVE: pto.tlog ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.tile_buf) {precisionType = #pto} // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: texp_arg( // EMITC: TEXP( diff --git a/test/lit/pto/treshape_tile_native_arg.pto b/test/lit/pto/treshape_tile_native_arg.pto index bbf649ad92..0d56a70714 100644 --- a/test/lit/pto/treshape_tile_native_arg.pto +++ b/test/lit/pto/treshape_tile_native_arg.pto @@ -38,4 +38,3 @@ module attributes {pto.target_arch = "a5"} { // NATIVE-SAME: -> !pto.tile_buf // NATIVE: return %[[RESULT]] : !pto.tile_buf // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile diff --git a/test/lit/pto/tstore_low_precision_a5_valid.pto b/test/lit/pto/tstore_low_precision_a5_valid.pto index df0099391d..2dbcfc59d3 100644 --- a/test/lit/pto/tstore_low_precision_a5_valid.pto +++ b/test/lit/pto/tstore_low_precision_a5_valid.pto @@ -27,7 +27,6 @@ module { // CHECK: func.func @tstore_low_precision_a5_valid(%arg0: memref<16x16xf8E4M3FN>, %arg1: memref<16x16x!pto.hif8>, %arg2: i64) // CHECK: %[[VEC:.*]] = pto.declare_tile -> !pto.tile_buf // CHECK: %[[ACC:.*]] = pto.declare_tile -> !pto.tile_buf -// CHECK-NOT: pto.bind_tile // CHECK: pto.section.vector { // CHECK: pto.tstore ins(%[[VEC]] : !pto.tile_buf) // CHECK: outs(%arg0 : memref<16x16xf8E4M3FN>) diff --git a/test/lit/pto/tsync_tile_native.pto b/test/lit/pto/tsync_tile_native.pto index 1a8404a830..4233de5f25 100644 --- a/test/lit/pto/tsync_tile_native.pto +++ b/test/lit/pto/tsync_tile_native.pto @@ -19,7 +19,6 @@ module { // NATIVE-LABEL: @tsync_tile_arg // NATIVE-SAME: !pto.tile_buf) outs(%arg1 : !pto.tile_buf) // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: tabs_arg( // EMITC: TABS( diff --git a/test/lit/pto/validshape_tile_native.pto b/test/lit/pto/validshape_tile_native.pto index 17b4609ddf..487514c1a0 100644 --- a/test/lit/pto/validshape_tile_native.pto +++ b/test/lit/pto/validshape_tile_native.pto @@ -40,4 +40,3 @@ module attributes {pto.target_arch = "a5"} { // NATIVE: pto.set_validshape {{.*}}, %[[ROW]], %[[COL]] // NATIVE: return // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile diff --git a/test/lit/pto/vector_specialty_tile_native.pto b/test/lit/pto/vector_specialty_tile_native.pto index 601c6ce412..479a07bffd 100644 --- a/test/lit/pto/vector_specialty_tile_native.pto +++ b/test/lit/pto/vector_specialty_tile_native.pto @@ -37,7 +37,6 @@ module { // NATIVE-LABEL: @ttri_arg // NATIVE: pto.ttri // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: thistogram_arg( // EMITC: THISTOGRAM< diff --git a/test/lit/pto/xor_tile_native.pto b/test/lit/pto/xor_tile_native.pto index c728792ef4..d22269b0e2 100644 --- a/test/lit/pto/xor_tile_native.pto +++ b/test/lit/pto/xor_tile_native.pto @@ -25,7 +25,6 @@ module { // NATIVE-LABEL: @txors_arg // NATIVE: pto.txors // NATIVE-NOT: memref< -// NATIVE-NOT: pto.bind_tile // EMITC-LABEL: txor_arg( // EMITC: TXOR( diff --git a/test/lit/vpto/fold_tile_buf_intrinsics.pto b/test/lit/vpto/fold_tile_buf_intrinsics.pto index 0e1f719cec..1c9b59e249 100644 --- a/test/lit/vpto/fold_tile_buf_intrinsics.pto +++ b/test/lit/vpto/fold_tile_buf_intrinsics.pto @@ -37,7 +37,6 @@ // ADDR-NOT: pto.tile_buf_addr // ADDR-NOT: pto.tile_valid_rows // ADDR-NOT: pto.tile_valid_cols -// ADDR-NOT: pto.bind_tile // ADDR: pto.vsts module attributes {pto.kernel_kind = #pto.kernel_kind} { diff --git a/test/samples/planmemory/plan_memory_bind_tile_alias_liveness.py b/test/samples/planmemory/plan_memory_tile_liveness.py similarity index 96% rename from test/samples/planmemory/plan_memory_bind_tile_alias_liveness.py rename to test/samples/planmemory/plan_memory_tile_liveness.py index 5097c00afd..c0238920c0 100644 --- a/test/samples/planmemory/plan_memory_bind_tile_alias_liveness.py +++ b/test/samples/planmemory/plan_memory_tile_liveness.py @@ -9,7 +9,7 @@ PTO_IR = r""" module { - func.func @bind_tile_alias_liveness(%arg0: memref<16x256xf16, #pto.address_space>, + func.func @tile_liveness(%arg0: memref<16x256xf16, #pto.address_space>, %arg1: memref<16x256xf16, #pto.address_space>) { %c16 = arith.constant 16 : index diff --git a/test/samples/runop.sh b/test/samples/runop.sh index 9a48538bd5..e1ee947d94 100755 --- a/test/samples/runop.sh +++ b/test/samples/runop.sh @@ -901,7 +901,7 @@ process_one_dir() { fi # Regression guard for Issue #207: - # SSA `pto.treshape` (lowered into `pto.bind_tile`) must lower to a single + # SSA `pto.treshape` must lower to a single # `TRESHAPE(dst, src)` instead of an invalid Tile-to-pointer cast sequence. if [[ "$base" == "reshape" ]]; then if ! grep -Fq "TRESHAPE(" "$cpp"; then From aab6cb1dcf51faafd99ec128eb2679825f879363 Mon Sep 17 00:00:00 2001 From: FangRui Date: Mon, 20 Jul 2026 19:27:53 +0800 Subject: [PATCH 86/91] Remove legacy PTO view-to-memref lowering --- docs/PTO_IR_manual.md | 2 +- include/PTO/IR/PTOOps.td | 21 +- include/PTO/Transforms/Passes.h | 2 - include/PTO/Transforms/Passes.td | 17 - lib/PTO/IR/PTO.cpp | 21 +- lib/PTO/Transforms/CMakeLists.txt | 1 - .../Transforms/InsertSync/SyncMacroModel.cpp | 4 +- lib/PTO/Transforms/PTOCanonicalizeIR.cpp | 2 +- lib/PTO/Transforms/PTOToEmitC.cpp | 4 +- lib/PTO/Transforms/PTOViewToMemref.cpp | 2934 ----------------- test/lit/pto/add_carry_tile_native.pto | 2 +- test/lit/pto/addmul_scalar_tile_native.pto | 2 +- test/lit/pto/addmul_tile_native.pto | 2 +- ...alloc_tile_plan_memory_no_memref_alloc.pto | 2 +- test/lit/pto/arg_reduction_tile_native.pto | 2 +- test/lit/pto/axpy_dequant_tile_native.pto | 2 +- test/lit/pto/basic_float_tile_native.pto | 2 +- test/lit/pto/bitcast_tile_native.pto | 2 +- .../pto/build_async_session_tile_native.pto | 4 +- test/lit/pto/castptr_emitc_tile_and_int.pto | 4 +- test/lit/pto/col_minmax_tile_native.pto | 2 +- test/lit/pto/colexpand_tile_native.pto | 2 +- test/lit/pto/colsum_tile_native.pto | 2 +- .../pto/compact_left_blayout_parser_a3.pto | 2 +- .../pto/compact_left_blayout_parser_a5.pto | 2 +- test/lit/pto/compare_tile_native.pto | 2 +- test/lit/pto/concat_tile_native.pto | 2 +- test/lit/pto/cvt_tile_native.pto | 2 +- test/lit/pto/declare_tile_tile_native.pto | 2 +- test/lit/pto/div_tile_native.pto | 2 +- test/lit/pto/expands_tile_native.pto | 2 +- .../lit/pto/extract_insert_fp_tile_native.pto | 2 +- test/lit/pto/extract_insert_tile_native.pto | 2 +- test/lit/pto/fillpad_tile_native.pto | 2 +- test/lit/pto/gather_tile_native.pto | 2 +- test/lit/pto/gemv_tile_native.pto | 2 +- .../gm_tensor_view_native_through_sync.pto | 4 +- test/lit/pto/graph_sync_solver_basic.pto | 2 +- .../issue31_partition_view_parser_compat.pto | 12 +- ...ov_treshape_square_dynamic_valid_shape.pto | 2 +- test/lit/pto/load_store_tile_native.pto | 6 +- test/lit/pto/logic_binary_tile_native.pto | 2 +- test/lit/pto/logic_scalar_tile_native.pto | 2 +- ...alize_tile_handles_control_flow_result.pto | 4 +- ...lize_tile_handles_subkernel_helper_abi.pto | 4 +- test/lit/pto/matmul_tile_native.pto | 2 +- test/lit/pto/mgather_gm2l1_sync_a5.pto | 2 +- ...ub_a5_elem_1x1_memref_macro_sync_model.pto | 2 +- test/lit/pto/mgather_mscatter_tile_native.pto | 2 +- test/lit/pto/minmax_scalar_tile_native.pto | 2 +- test/lit/pto/minmax_tile_native.pto | 2 +- test/lit/pto/mov_tile_native.pto | 2 +- .../lit/pto/movement_metadata_tile_native.pto | 2 +- test/lit/pto/mrgsort_tile_native.pto | 2 +- .../pto/multi_tile_buf_type_parse_print.pto | 2 +- ...lti_tile_const_preload_dyn_loop_select.pto | 4 +- .../multi_tile_get_const_slot_lowering.pto | 6 +- .../pto/multi_tile_get_dyn_slot_lowering.pto | 4 +- .../pto/multi_tile_level3_explicit_addr.pto | 6 +- test/lit/pto/multi_tile_no_loop_unroll.pto | 4 +- test/lit/pto/part_arithmetic_tile_native.pto | 2 +- test/lit/pto/pipe_tile_native.pto | 2 +- test/lit/pto/pow_rsqrt_tile_native.pto | 2 +- test/lit/pto/print_tile_native.pto | 4 +- ...ptodsl_subkernel_helper_tile_abi_emitc.pto | 4 +- test/lit/pto/ptr_int_cast.pto | 2 +- test/lit/pto/quant_mx_tile_native.pto | 2 +- test/lit/pto/reduction_tile_native.pto | 2 +- test/lit/pto/rowexpand_tile_native.pto | 2 +- .../pto/scratch_elementwise_tile_native.pto | 2 +- test/lit/pto/select_tile_native.pto | 2 +- test/lit/pto/shift_tile_native.pto | 2 +- test/lit/pto/sort32_tile_native.pto | 2 +- test/lit/pto/store_fp_tile_native.pto | 2 +- test/lit/pto/sub_tile_native.pto | 2 +- ...dynamic_offset_static_valid_regression.pto | 4 +- test/lit/pto/tassign_reject_memref.pto | 23 + test/lit/pto/tassign_tile_native_arg.pto | 2 +- .../lit/pto/tile_buf_addr_tile_native_arg.pto | 4 +- test/lit/pto/tile_compact_mode_emitc.pto | 8 +- .../pto/tile_scalar_access_tile_native.pto | 2 +- test/lit/pto/tlrelu_tile_native.pto | 2 +- test/lit/pto/tmov_fp_tile_native.pto | 2 +- test/lit/pto/tpow_emitc.pto | 2 +- test/lit/pto/tpow_int_no_tmp.pto | 2 +- test/lit/pto/tpows_emitc.pto | 2 +- test/lit/pto/tpows_int_no_tmp.pto | 2 +- test/lit/pto/tprefetch_tile_native.pto | 2 +- ...d_a3.pto => tquant_no_implicit_tmp_a3.pto} | 7 +- test/lit/pto/trans_tile_native.pto | 2 +- .../pto/transcendental_unary_tile_native.pto | 2 +- test/lit/pto/trem_emitc.pto | 2 +- test/lit/pto/trems_emitc.pto | 2 +- ...explicit_dynamic_valid_shape_preserved.pto | 2 +- test/lit/pto/treshape_tile_native_arg.pto | 2 +- test/lit/pto/trsqrt_emitc.pto | 2 +- test/lit/pto/tsort32_emitc.pto | 2 +- test/lit/pto/tsync_tile_native.pto | 2 +- test/lit/pto/unary_compute_tile_native.pto | 2 +- test/lit/pto/validshape_tile_native.pto | 2 +- test/lit/pto/vector_specialty_tile_native.pto | 2 +- test/lit/pto/xor_tile_native.pto | 2 +- .../final_emitc_last_use_level2.pto | 2 +- ...p_fusion_adapter_placement_level2_tadd.pto | 2 +- ...p_fusion_adapter_placement_level3_tadd.pto | 4 +- .../op_fusion_region_pipeline_level2.pto | 2 +- tools/ptoas/ptoas.cpp | 2 - 107 files changed, 172 insertions(+), 3110 deletions(-) delete mode 100644 lib/PTO/Transforms/PTOViewToMemref.cpp create mode 100644 test/lit/pto/tassign_reject_memref.pto rename test/lit/pto/{tquant_dynamic_tmp_valid_a3.pto => tquant_no_implicit_tmp_a3.pto} (87%) diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index 776df4657a..511fcf3f5f 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -8503,7 +8503,7 @@ dst[i, j] = Quantize(src[i, j]; fp, quant_type) | `fp` | `pto.tile_buf` | Scaling parameter tile (`f32`) | | `offset` | `pto.tile_buf` | Optional asymmetric offset tile (`f32`, required for `INT8_ASYM`) | | `dst` | `pto.tile_buf` | Destination tile (`i8` for SYM, `ui8` for ASYM) | -| `tmp` | `pto.tile_buf` | Optional scratch tile. A2/A3 uses it for row-broadcast and fp32-to-s32 conversion scratch; PTOAS auto-synthesizes it when omitted. A5 accepts it as a placeholder but does not require it. | +| `tmp` | `pto.tile_buf` | Optional scratch tile. A2/A3 uses it for row-broadcast and fp32-to-s32 conversion scratch. When omitted, PTOAS preserves the no-tmp form and selects the 4-argument backend overload. A5 accepts it as a placeholder but does not require it. | **Attributes:** diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index b56c6bb8e1..5a0ae134b7 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -1377,8 +1377,8 @@ def TMovOp : PTO_TOp<"tmov", [ return as.getAddressSpace(); return std::nullopt; } - // Post PTOViewToMemref: tile_buf is erased to memref but memorySpace is - // preserved in memref's memorySpace attribute. + // Legacy imported IR may carry memref values with the address space in + // the memref memorySpace attribute. if (auto mr = llvm::dyn_cast<::mlir::MemRefType>(ty)) { if (auto ms = mr.getMemorySpace()) { if (auto as = @@ -1415,8 +1415,7 @@ def TMovOp : PTO_TOp<"tmov", [ // PTOAS uses SCALING for two distinct PTO-ISA destinations. A5 MX // scale tiles are ScaleLeft/ScaleRight loads into L0A/L0B and run on // MTE1. Regular Scaling tiles are FBUF loads (TMOV_M2S) and run on - // FIX. PTOViewToMemref preserves the element type but erases the tile - // role before sync insertion, so distinguish the MX-only f8E8M0 type. + // Distinguish the MX-only f8E8M0 type from regular scaling tiles. Type dstElemTy; if (auto tb = llvm::dyn_cast<::mlir::pto::TileBufType>(getDst().getType())) dstElemTy = tb.getElementType(); @@ -2399,7 +2398,7 @@ def TAllocOp : PTO_TOp<"talloc", [ }]; } -def TAssignOp : PTO_Op<"tassign", [Pure]> { +def TAssignOp : PTO_Op<"tassign", [Pure, AllTypesMatch<["tile", "result"]>]> { let summary = "Rebind a declared tile handle to a runtime address"; let description = [{ Rebinds an existing tile handle to a new address. This op is intended for @@ -3077,9 +3076,8 @@ def MGatherOp : PTO_TOp<"mgather", [ // MTE2 pipe on every arch, mirroring pto.tload. The GM -> UB (VEC) path // below keeps its A5 SIMT (PIPE_V) / A2A3 MTE2 selection. // - // Resolve the dst address space from BOTH the tile_buf and the lowered - // memref form: PTOViewToMemref rewrites dst into a memref<...,mat> before - // InsertSync / GraphSyncSolver run, and both consume getPipe() through the + // Resolve the dst address space from both tile_buf and legacy imported + // memref forms. InsertSync / GraphSyncSolver consume getPipe() through the // OpPipeInterface. A tile-only check would fall through to PIPE_V on A5. auto getAddressSpace = [](::mlir::Type ty) -> std::optional<::mlir::pto::AddressSpace> { @@ -5577,9 +5575,10 @@ def TQuantOp : PTO_TOp<"tquant", [ TQUANT_IMPL row-wise broadcast (the fp->s32 conversion reuses it). On pto-isa builds where the 4-arg TQUANT overload routes through the fixed TMP_UB_OFFSET scratch area, supplying tmp selects the tmp-aware 5-arg - overload and avoids the fixed-address collision. PTOAS auto-synthesizes - this scratch tile for A2/A3 only. On A5 the tmp operand is accepted as a - placeholder and does not add extra shape or memory-planning constraints. + overload and avoids the fixed-address collision. PTOAS preserves an omitted + tmp and lowers it through the 4-arg overload. On A5 the tmp operand is + accepted as a placeholder and does not add extra shape or memory-planning + constraints. DPS form (INT8_SYM): pto.tquant ins(%src, %fp : , ) diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index d88b918c95..4dada8a275 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -72,7 +72,6 @@ createPlanMemoryPass(const PlanMemoryOptions &options = {}); std::unique_ptr createPlanMemoryModernPass(const PlanMemoryOptions &options); std::unique_ptr createPTORemoveRedundantBarrierPass(); -std::unique_ptr createPTOViewToMemrefPass(); std::unique_ptr createPTOValidateIntToPtrUsesPass(); std::unique_ptr createPTORematerializeFixpipeVectorQuantPass(); std::unique_ptr createPTOMaterializeTileHandlesPass(); @@ -120,7 +119,6 @@ std::unique_ptr createPTOInlineLibCallPass(const PTOInlineLibCallOptions &options = {}); std::unique_ptr createPTOInlineBackendHelpersPass( const PTOInlineBackendHelpersOptions &options = {}); -void registerPTOViewToMemrefPass(); //===----------------------------------------------------------------------===// // Registration diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 420d540df2..49bf964833 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -689,23 +689,6 @@ def PTOVerifySubkernelPipeContract ]; } -def PTOViewToMemref : Pass<"pto-view-to-memref", "ModuleOp"> { - let summary = "Lower PTO views to memref with Metadata Binding"; - let description = [{ - Lower PTO tile/view operations to memref-based IR while preserving tile - metadata through binding ops and SSA backtracking. - }]; - - let constructor = "mlir::pto::createPTOViewToMemrefPass()"; - - let dependentDialects = [ - "mlir::pto::PTODialect", - "mlir::memref::MemRefDialect", - "mlir::arith::ArithDialect", - "mlir::func::FuncDialect" - ]; -} - def PTOValidateIntToPtrUses : Pass<"pto-validate-inttoptr-uses", "func::FuncOp"> { let summary = "Validate restricted uses of PTO inttoptr results"; let description = [{ diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index b0128fbe6c..7f8cd60f69 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -2997,8 +2997,8 @@ LogicalResult AllocMultiTileOp::verify() { return failure(); // Multi-buffer slots are placed at product(shape) * element_size byte - // intervals -- both the --pto-level=level3 lowering (PTOViewToMemref) and - // PTOPlanMemory size them that way. `row_plus_one` compaction inflates the + // intervals -- both level3 validation and PTOPlanMemory size them that way. + // `row_plus_one` compaction inflates the // major stride by one element per row, so the slot's physical strided // footprint exceeds product(shape) and adjacent slots would silently overlap // (data corruption). Reject it until the slot stride is derived from the true @@ -3128,16 +3128,13 @@ LogicalResult TAssignOp::verify() { return emitOpError("result type must match tile operand type"); } - Type tileType = getTile().getType(); - if (auto tileTy = dyn_cast(tileType)) { - if (failed(verifyConstantLocalAddress(getOperation(), getAddr(), - tileTy.getMemorySpace()))) - return failure(); - } else if (auto memRefTy = dyn_cast(tileType)) { - if (failed(verifyConstantLocalAddress(getOperation(), getAddr(), - memRefTy.getMemorySpace()))) - return failure(); - } + auto tileTy = dyn_cast(getTile().getType()); + if (!tileTy) + return emitOpError("expects tile operand and result to be !pto.tile_buf"); + + if (failed(verifyConstantLocalAddress(getOperation(), getAddr(), + tileTy.getMemorySpace()))) + return failure(); return success(); } diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index 92956a1a3d..e9d9eb7c86 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -40,7 +40,6 @@ add_mlir_dialect_library(PTOTransforms InsertSync/PTOInsertSync.cpp PTOInjectBarrierAllSync.cpp InsertSync/InsertSyncDebug.cpp - PTOViewToMemref.cpp PTORematerializeFixpipeVectorQuant.cpp PTOValidateIntToPtrUses.cpp InsertTemplateAttributes.cpp diff --git a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp index 86ad3f99a7..9db12e038c 100644 --- a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp +++ b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp @@ -196,8 +196,8 @@ std::optional getTGatherSyncMacroModel(pto::TGatherOp op) { } // Resolve the PTO address space of an MGatherOp operand type. Mirrors the -// getAddressSpace helper used by MGatherOp::getPipe() so the model matches the -// post-PTOViewToMemref (memref<...,mat>) form that InsertSync actually sees. +// getAddressSpace helper used by MGatherOp::getPipe() and retain compatibility +// with legacy imported memref<...,mat> IR. static std::optional getMGatherOperandAddressSpace(::mlir::Type ty) { if (auto tb = ::mlir::dyn_cast<::mlir::pto::TileBufType>(ty)) { diff --git a/lib/PTO/Transforms/PTOCanonicalizeIR.cpp b/lib/PTO/Transforms/PTOCanonicalizeIR.cpp index c4495d0d14..e85d47155f 100644 --- a/lib/PTO/Transforms/PTOCanonicalizeIR.cpp +++ b/lib/PTO/Transforms/PTOCanonicalizeIR.cpp @@ -34,7 +34,7 @@ namespace { // // This pass canonicalizes rank-2 TensorViewType / PartitionTensorViewType // into the right-aligned rank-5 form [1, 1, 1, R, C] used by all backends -// (A3, A5, VPTO EmitC codegen and the 5D memref rank in PTOViewToMemref). +// (A3, A5, and VPTO EmitC codegen). // // Ops that carry **rank-dependent operands** must be structurally rewritten // (their operand count or operand values change when rank changes): diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 4c5af18f12..cb88d45f9b 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -10047,9 +10047,7 @@ struct PTOFillPadToEmitC : public OpConversionPattern { ArrayAttr templateArgs{}; if (auto padValueAttr = op.getPadValueAttr()) { // The verifier only accepts explicit padValue for loc=mat tile-form - // tfillpad. PTOViewToMemref preserves that attribute after rewriting the - // operands to memref form, so lowering must trust the preserved semantic - // contract instead of re-checking the now-erased TileBufType. + // tfillpad, so lowering can trust the preserved semantic contract. templateArgs = rewriter.getArrayAttr( {emitc::OpaqueAttr::get(ctx, padValueTok(padValueAttr.getValue()))}); } diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp deleted file mode 100644 index 9120207880..0000000000 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ /dev/null @@ -1,2934 +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. - -//===- PTOViewToMemref.cpp ------------------------------------------------===// -//===----------------------------------------------------------------------===// -// -// Lower PTO tile/view operations to memref-based IR while preserving tile -// metadata through binding ops and SSA backtracking. - -#include "PTO/IR/PTO.h" -#include "PTO/IR/PTOTypeUtils.h" -#include "PTO/Transforms/Passes.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" - -#include "mlir/IR/AsmState.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/IR/SymbolTable.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassRegistry.h" - -namespace mlir { -namespace pto { -#define GEN_PASS_DEF_PTOVIEWTOMEMREF -#include "PTO/Transforms/Passes.h.inc" -} // namespace pto -} // namespace mlir - -#include "llvm/ADT/SmallVector.h" -#include "llvm/Support/raw_ostream.h" -#include "Utils.h" // 假设包含一些通用的工具函数 - -#include -#include -#include - -#define DEBUG_TYPE "pto-view-to-memref" - -using namespace mlir; - -namespace mlir { -namespace pto { - -namespace { - -static Type convertPTOTypeToMemRef(Type t); - -static bool hasTileNativeAllocationRoot(func::FuncOp func) { - bool found = false; - auto result = func.walk([&](Operation *op) { - if (isa(op)) { - found = true; - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - (void)result; - return found; -} - -static bool hasMigratedTileNativeOp(func::FuncOp func) { - bool found = false; - auto result = func.walk([&](Operation *op) { - if (isa(op)) { - found = true; - return WalkResult::interrupt(); - } - return WalkResult::advance(); - }); - (void)result; - return found; -} - -constexpr size_t kTileRank2D = 2; -constexpr unsigned kShapeVectorInlineCapacity = 4; -constexpr unsigned kOperationVectorInlineCapacity = 8; - -constexpr int64_t kElementBytes1 = 1; -constexpr int64_t kElementBytes2 = 2; -constexpr int64_t kElementBytes4 = 4; -constexpr int64_t kElementBytes8 = 8; -constexpr int64_t kElementBytes16 = 16; -constexpr int64_t kElementBytes32 = 32; - -constexpr int64_t kInnerExtent1 = 1; -constexpr int64_t kInnerExtent2 = 2; -constexpr int64_t kInnerExtent4 = 4; -constexpr int64_t kInnerExtent8 = 8; -constexpr int64_t kInnerExtent16 = 16; -constexpr int64_t kInnerExtent32 = 32; - -constexpr int32_t kFractalSize32 = 32; -constexpr int32_t kFractalSize512 = 512; -constexpr int32_t kFractalSize1024 = 1024; - -constexpr int32_t kBLayoutColMajor = - static_cast(BLayout::ColMajor); -constexpr int32_t kSLayoutNoneBox = - static_cast(SLayout::NoneBox); -constexpr int32_t kSLayoutRowMajor = - static_cast(SLayout::RowMajor); -constexpr int32_t kSLayoutColMajor = - static_cast(SLayout::ColMajor); -constexpr int32_t kCompactModeRowPlusOne = - static_cast(CompactMode::RowPlusOne); - -constexpr unsigned kThirdOperandIndex = 2; -constexpr unsigned kFourthOperandIndex = 3; -constexpr unsigned kFifthOperandIndex = 4; -constexpr unsigned kSixthOperandIndex = 5; - -template -using SmallInlineVector = SmallVector; - -template -using DefaultInlineVector = SmallVector; - -// ============================================================================= -// Helper: Valid dims backtracking (v_row / v_col) -// ============================================================================= -static void lookupValidDims(Value v, Value &vRow, Value &vCol) { - if (auto alloc = v.getDefiningOp()) { - vRow = alloc.getValidRow(); - vCol = alloc.getValidCol(); - return; - } - if (auto pc = v.getDefiningOp()) { - vRow = pc.getValidRow(); - vCol = pc.getValidCol(); - return; - } - if (auto subview = v.getDefiningOp()) { - lookupValidDims(subview.getSource(), vRow, vCol); - return; - } - if (auto cast = v.getDefiningOp()) { - lookupValidDims(cast.getSource(), vRow, vCol); - return; - } - if (auto cast = v.getDefiningOp()) { - lookupValidDims(cast.getSource(), vRow, vCol); - return; - } - // pto.fusion_region result 不直接携带 valid dims；作为兜底情况，沿着 - // pto.yield 回溯到 region 内真正的 tile handle/memref 定义链继续查找。 - if (auto regionResult = dyn_cast(v)) { - if (auto fusionRegion = - dyn_cast(regionResult.getOwner())) { - auto yieldOp = dyn_cast( - fusionRegion.getBody().front().getTerminator()); - if (!yieldOp) { - vRow = Value(); - vCol = Value(); - return; - } - unsigned resultIndex = regionResult.getResultNumber(); - if (resultIndex >= yieldOp.getNumOperands()) { - vRow = Value(); - vCol = Value(); - return; - } - lookupValidDims(yieldOp.getOperand(resultIndex), vRow, vCol); - return; - } - } - - vRow = Value(); - vCol = Value(); -} - -template -static OpTy replaceOpWithClonedAttrs(IRRewriter &rewriter, Operation *op, - Args &&...args) { - auto newOp = - rewriter.create(op->getLoc(), std::forward(args)...); - newOp->setAttrs(op->getAttrs()); - rewriter.replaceOp(op, newOp->getResults()); - return newOp; -} - -static LogicalResult synthesizeMissingTQuantTmpOps(func::FuncOp func, - MLIRContext *ctx) { - if (mlir::pto::getTargetArch(func.getOperation()) == mlir::pto::PTOArch::A5) - return success(); - - DefaultInlineVector quantOps; - func.walk([&](mlir::pto::TQuantOp op) { quantOps.push_back(op); }); - - for (auto op : quantOps) { - if (op.getTmp()) - continue; - - auto srcTy = dyn_cast(op.getSrc().getType()); - if (!srcTy) - continue; - - if (srcTy.getRank() != static_cast(kTileRank2D)) { - op.emitError("expects rank-2 src tile to synthesize tquant tmp"); - return failure(); - } - - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - SmallInlineVector shape(srcTy.getShape().begin(), - srcTy.getShape().end()); - SmallInlineVector validShape(srcTy.getValidShape().begin(), - srcTy.getValidShape().end()); - if (validShape.empty()) - validShape.assign(shape.begin(), shape.end()); - - auto tmpTy = mlir::pto::TileBufType::get( - ctx, shape, srcTy.getElementType(), srcTy.getMemorySpace(), validShape, - mlir::pto::TileBufConfigAttr::getDefault(ctx)); - Value validRow; - Value validCol; - lookupValidDims(op.getSrc(), validRow, validCol); - Value tmp = rewriter - .create(loc, tmpTy, Value(), - validRow, validCol) - .getResult(); - - rewriter.create( - loc, TypeRange{}, op.getSrc(), op.getFp(), op.getOffset(), tmp, - op.getDst(), op.getQuantTypeAttr()); - rewriter.eraseOp(op); - } - - return success(); -} - -// ============================================================================= -// Helper Functions for Layout Normalization -// ============================================================================= - -struct TileLayoutInfo { - int64_t rowStride = 1; - int64_t colStride = 1; - int64_t innerRows = 1; - int64_t innerCols = 1; - bool boxed = false; // slayout != NoneBox -}; - -struct TileLayoutConfig { - int32_t bLayout = 0; - int32_t sLayout = 0; - int32_t fractalSize = kFractalSize512; - int32_t compactMode = 0; -}; - -static int64_t getElemBytes(Type elemTy) { - unsigned bytes = getPTOStorageElemByteSize(elemTy); - return bytes == 0 ? -1 : static_cast(bytes); -} - -template -static bool readEnumAttrOrIntegerI32(Attribute attr, int32_t &out) { - if (auto enumAttr = dyn_cast(attr)) { - out = static_cast(enumAttr.getValue()); - return true; - } - if (auto intAttr = dyn_cast(attr)) { - out = static_cast(intAttr.getInt()); - return true; - } - return false; -} - -static bool readBLayoutI32(Attribute attr, int32_t &out) { - return readEnumAttrOrIntegerI32(attr, out); -} - -static bool readSLayoutI32(Attribute attr, int32_t &out) { - return readEnumAttrOrIntegerI32(attr, out); -} - -static bool readCompactModeI32(Attribute attr, int32_t &out) { - return readEnumAttrOrIntegerI32(attr, out); -} - -static Value peelIndexLikeCast(Value value) { - while (true) { - if (auto castOp = value.getDefiningOp()) { - value = castOp.getIn(); - continue; - } - if (auto extOp = value.getDefiningOp()) { - value = extOp.getIn(); - continue; - } - if (auto extOp = value.getDefiningOp()) { - value = extOp.getIn(); - continue; - } - if (auto truncOp = value.getDefiningOp()) { - value = truncOp.getIn(); - continue; - } - return value; - } -} - -static bool getConstIndexValue(Value value, int64_t &out) { - value = peelIndexLikeCast(value); - if (auto constIndex = value.getDefiningOp()) { - out = constIndex.value(); - return true; - } - if (auto constInt = value.getDefiningOp()) { - out = constInt.value(); - return true; - } - auto constOp = value.getDefiningOp(); - auto intAttr = - constOp ? dyn_cast(constOp.getValue()) : IntegerAttr(); - if (!intAttr) - return false; - out = intAttr.getInt(); - return true; -} - -static TileLayoutConfig getTileLayoutConfig(mlir::pto::TileBufConfigAttr cfg) { - TileLayoutConfig config; - (void)readBLayoutI32(cfg.getBLayout(), config.bLayout); - (void)readSLayoutI32(cfg.getSLayout(), config.sLayout); - if (auto attr = dyn_cast(cfg.getSFractalSize())) - config.fractalSize = static_cast(attr.getInt()); - (void)readCompactModeI32(cfg.getCompactMode(), config.compactMode); - return config; -} - -static bool getFractal512InnerExtent(int64_t elemBytes, int64_t &extent) { - switch (elemBytes) { - case kElementBytes1: - extent = kInnerExtent32; - return true; - case kElementBytes2: - extent = kInnerExtent16; - return true; - case kElementBytes4: - extent = kInnerExtent8; - return true; - case kElementBytes8: - extent = kInnerExtent4; - return true; - case kElementBytes16: - extent = kInnerExtent2; - return true; - case kElementBytes32: - extent = kInnerExtent1; - return true; - default: - return false; - } -} - -static bool computeBoxInnerShape(const TileLayoutConfig &config, Type elemTy, - TileLayoutInfo &info) { - info.boxed = config.sLayout != kSLayoutNoneBox; - if (!info.boxed) { - info.innerRows = kInnerExtent1; - info.innerCols = kInnerExtent1; - return true; - } - - int64_t elemBytes = getElemBytes(elemTy); - if (elemBytes <= 0) - return false; - - switch (config.fractalSize) { - case kFractalSize1024: - info.innerRows = kInnerExtent16; - info.innerCols = kInnerExtent16; - return true; - case kFractalSize32: - info.innerRows = kInnerExtent16; - info.innerCols = kInnerExtent2; - return true; - case kFractalSize512: - if (config.sLayout == kSLayoutRowMajor) { - info.innerRows = kInnerExtent16; - return getFractal512InnerExtent(elemBytes, info.innerCols); - } - if (config.sLayout == kSLayoutColMajor) { - if (!getFractal512InnerExtent(elemBytes, info.innerRows)) - return false; - info.innerCols = kInnerExtent16; - return true; - } - return false; - default: - return false; - } -} - -static bool computeTilePointerStrides(const TileLayoutConfig &config, - ArrayRef shape, - TileLayoutInfo &info) { - int64_t rows = shape[0]; - int64_t cols = shape[1]; - auto applyCompactToMajorStride = [&](int64_t majorStride) -> int64_t { - if (config.compactMode == kCompactModeRowPlusOne) - return majorStride + kInnerExtent1; - return majorStride; - }; - if (!info.boxed) { - if (config.bLayout == kBLayoutColMajor) { - info.rowStride = kInnerExtent1; - info.colStride = applyCompactToMajorStride(rows); - return true; - } - info.rowStride = applyCompactToMajorStride(cols); - info.colStride = kInnerExtent1; - return true; - } - - if (config.bLayout == kBLayoutColMajor) { - if (config.sLayout != kSLayoutRowMajor) - return false; - info.rowStride = info.innerCols; - info.colStride = applyCompactToMajorStride(rows); - return true; - } - - info.rowStride = applyCompactToMajorStride(cols); - info.colStride = info.innerRows; - return true; -} - -static bool computeTileLayoutInfo(mlir::pto::TileBufConfigAttr cfg, Type elemTy, - ArrayRef shape, - TileLayoutInfo &info) { - if (shape.size() != kTileRank2D || - llvm::is_contained(shape, ShapedType::kDynamic)) - return false; - - TileLayoutConfig config = getTileLayoutConfig(cfg); - return computeBoxInnerShape(config, elemTy, info) && - computeTilePointerStrides(config, shape, info); -} - -static void collectAffineAddTerms(AffineExpr root, - SmallVectorImpl &terms) { - SmallInlineVector pending{root}; - while (!pending.empty()) { - AffineExpr current = pending.pop_back_val(); - auto addExpr = llvm::dyn_cast(current); - if (!addExpr || addExpr.getKind() != AffineExprKind::Add) { - terms.push_back(current); - continue; - } - pending.push_back(addExpr.getRHS()); - pending.push_back(addExpr.getLHS()); - } -} - -static bool tryAssignAffineStride(AffineExpr expr, - MutableArrayRef strides) { - if (auto dim = llvm::dyn_cast(expr)) { - strides[dim.getPosition()] = 1; - return true; - } - - auto mulExpr = llvm::dyn_cast(expr); - if (!mulExpr || mulExpr.getKind() != AffineExprKind::Mul) - return false; - - auto assignStride = [&](AffineExpr dimExpr, - AffineExpr constantExpr) -> bool { - auto dim = llvm::dyn_cast(dimExpr); - auto constant = llvm::dyn_cast(constantExpr); - if (!dim || !constant) - return false; - strides[dim.getPosition()] = constant.getValue(); - return true; - }; - return assignStride(mulExpr.getLHS(), mulExpr.getRHS()) || - assignStride(mulExpr.getRHS(), mulExpr.getLHS()); -} - -[[maybe_unused]] static void decomposeStridedLayout(AffineMap map, - SmallVectorImpl &strides) { - strides.assign(map.getNumDims(), 0); - if (map.getNumResults() != 1) - return; - - SmallInlineVector terms; - collectAffineAddTerms(map.getResult(0), terms); - for (AffineExpr term : terms) - (void)tryAssignAffineStride(term, strides); -} - -static Value makeIndexConstant(IRRewriter &rewriter, Location loc, - int64_t value) { - return rewriter.create(loc, rewriter.getIndexType(), - rewriter.getIndexAttr(value)); -} - -static Value buildFlattenedPtrLikeMemRefView(IRRewriter &rewriter, Location loc, - Value sourceMemref, - MemRefType targetType, - Operation *anchorOp) { - auto sourceType = dyn_cast(sourceMemref.getType()); - if (!sourceType) { - anchorOp->emitError( - "tile_buf_addr ptr-like lowering expects the source to be memref-backed"); - return Value(); - } - - Value totalElems = makeIndexConstant(rewriter, loc, 1); - for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) { - Value extent; - if (sourceType.isDynamicDim(dim)) - extent = rewriter.create(loc, sourceMemref, dim); - else - extent = makeIndexConstant(rewriter, loc, sourceType.getDimSize(dim)); - totalElems = rewriter.create(loc, totalElems, extent); - } - - SmallVector sizes{totalElems}; - SmallVector strides{rewriter.getIndexAttr(1)}; - return rewriter - .create(loc, targetType, sourceMemref, - rewriter.getIndexAttr(0), sizes, - strides) - .getResult(); -} - -static SmallVector computeCompactStrides(ArrayRef shape) { - SmallVector strides(shape.size(), 1); - int64_t stride = 1; - for (int i = static_cast(shape.size()) - 1; i >= 0; --i) { - strides[i] = stride; - if (shape[i] != ShapedType::kDynamic) - stride *= shape[i]; - } - return strides; -} - -// 确保 Value 是 Index 类型 -static Value ensureIndex(IRRewriter &rewriter, Location loc, Value v, - Operation *anchorOp) { - if (v.getType().isIndex()) - return v; - if (isa(v.getType())) - return rewriter.create(loc, rewriter.getIndexType(), v); - if (anchorOp) - anchorOp->emitError() << "expected index or integer, but got " << v.getType(); - return Value(); -} - -static bool tryGetIndexAttrFromValue(IRRewriter &rewriter, Value v, - IntegerAttr &constAttr) { - if (auto cOp = v.getDefiningOp()) { - constAttr = rewriter.getIndexAttr(cOp.value()); - return true; - } - if (auto cInt = v.getDefiningOp()) { - constAttr = rewriter.getIndexAttr(cInt.value()); - return true; - } - return false; -} - -static void appendMixedIndex(IRRewriter &rewriter, Location loc, Value v, - Operation *anchorOp, - SmallVectorImpl &mixedVals) { - IntegerAttr constAttr; - if (tryGetIndexAttrFromValue(rewriter, v, constAttr)) { - mixedVals.push_back(constAttr); - return; - } - mixedVals.push_back(ensureIndex(rewriter, loc, v, anchorOp)); -} - -[[maybe_unused]] static void dumpPretty(Operation *op, llvm::raw_ostream &os) { - OpPrintingFlags flags; - flags.useLocalScope(); - AsmState state(op, flags); - op->print(os, state); - os << "\n"; - os.flush(); -} - -// ============================================================================= -// Type Converter Logic -// ============================================================================= - -static SmallVector buildTileMemRefStrides(mlir::pto::TileBufType tbTy) { - TileLayoutInfo info; - if (computeTileLayoutInfo(tbTy.getConfigAttr(), tbTy.getElementType(), - tbTy.getShape(), info)) { - return {info.rowStride, info.colStride}; - } - return computeCompactStrides(tbTy.getShape()); -} - -static Type convertTileBufTypeToMemRef(mlir::pto::TileBufType tbTy) { - auto layoutAttr = StridedLayoutAttr::get(tbTy.getContext(), - ShapedType::kDynamic, - buildTileMemRefStrides(tbTy)); - return MemRefType::get(tbTy.getShape(), tbTy.getElementType(), layoutAttr, - tbTy.getMemorySpace()); -} - -static Type convertPTOTypeToMemRef(Type t) { - // 1. 处理 !pto.ptr - if (auto pty = dyn_cast(t)) { - return MemRefType::get({ShapedType::kDynamic}, pty.getElementType(), - MemRefLayoutAttrInterface(), pty.getMemorySpace()); - } - - // 2. 处理 !pto.tile_buf<...> - if (auto tbTy = dyn_cast(t)) - return convertTileBufTypeToMemRef(tbTy); - // 3. !pto.multi_tile_buf: collapses to the slot memref shape; - // the N-slot fan-out lives on the `pto.multi_buffer` attr written by - // the alloc lowering, not in the type. This branch is defensive -- - // by design multi_tile_buf does not appear on function boundaries. - if (auto mtbTy = dyn_cast(t)) - return convertTileBufTypeToMemRef(mtbTy.getSlotType()); - if (auto tvTy = dyn_cast(t)) - return MemRefType::get(tvTy.getShape(), tvTy.getElementType(), - MemRefLayoutAttrInterface(), Attribute()); - if (auto partTy = dyn_cast(t)) - return MemRefType::get(partTy.getShape(), partTy.getElementType(), - MemRefLayoutAttrInterface(), Attribute()); - // 其他类型透传 - return t; -} - -static LogicalResult lowerPtrLikeTileBufAddrOps(func::FuncOp func, - MLIRContext *ctx) { - SmallVector addrOps; - func.walk([&](mlir::pto::TileBufAddrOp op) { addrOps.push_back(op); }); - - for (auto op : addrOps) { - if (!isa(op.getDst().getType())) - continue; - - auto targetType = - dyn_cast(convertPTOTypeToMemRef(op.getDst().getType())); - if (!targetType) { - op.emitError("failed to convert tile_buf_addr pointer result to memref"); - return failure(); - } - - Value source = op.getSrc(); - if (!isa(source.getType())) - continue; - - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Value replacement = - buildFlattenedPtrLikeMemRefView(rewriter, op.getLoc(), source, - targetType, op.getOperation()); - if (!replacement) - return failure(); - rewriter.replaceOp(op, replacement); - } - return success(); -} - -[[maybe_unused]] static LogicalResult lowerTensorViewDimOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector tvDims; - func.walk([&](mlir::pto::GetTensorViewDimOp op) { tvDims.push_back(op); }); - - for (auto op : tvDims) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Value view = op.getTensorView(); - auto mrTy = dyn_cast(view.getType()); - if (!mrTy) - continue; - Value dim = rewriter.create(op.getLoc(), view, op.getDimIndex()); - rewriter.replaceOp(op, dim); - } - return success(); -} - -static LogicalResult lowerPartitionViewOps(func::FuncOp func, MLIRContext *ctx) { - DefaultInlineVector partitionViews; - func.walk([&](mlir::pto::PartitionViewOp op) { partitionViews.push_back(op); }); - - for (auto op : partitionViews) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - Value src = op.getOperand(0); - auto srcMrTy = dyn_cast(src.getType()); - if (!srcMrTy) - continue; - int64_t rank = srcMrTy.getRank(); - - SmallVector staticSizes; - SmallVector mixedSizes; - for (Value size : op.getSizes()) { - IntegerAttr constAttr; - bool isStatic = false; - if (auto cOp = size.getDefiningOp()) { - constAttr = rewriter.getIndexAttr(cOp.value()); - isStatic = true; - } else if (auto cInt = size.getDefiningOp()) { - constAttr = rewriter.getIndexAttr(cInt.value()); - isStatic = true; - } - - if (isStatic) { - mixedSizes.push_back(constAttr); - staticSizes.push_back(constAttr.getInt()); - } else { - mixedSizes.push_back(ensureIndex(rewriter, loc, size, op)); - staticSizes.push_back(ShapedType::kDynamic); - } - } - - SmallVector mixedOffsets; - for (Value offset : op.getOffsets()) { - appendMixedIndex(rewriter, loc, offset, op, mixedOffsets); - } - - int64_t dyn = ShapedType::kDynamic; - SmallVector dynStrides(rank, dyn); - auto layout = StridedLayoutAttr::get(ctx, dyn, dynStrides); - auto resTy = MemRefType::get(staticSizes, srcMrTy.getElementType(), layout, - srcMrTy.getMemorySpace()); - - SmallVector mixedStrides(rank, rewriter.getIndexAttr(1)); - auto sv = rewriter.create(loc, resTy, src, mixedOffsets, - mixedSizes, mixedStrides); - if (Operation *srcDef = src.getDefiningOp()) { - if (auto layoutAttr = srcDef->getAttrOfType("layout")) - sv->setAttr("layout", layoutAttr); - } - rewriter.replaceOp(op, sv.getResult()); - } - return success(); -} - -// ============================================================================= -// The Pass Implementation -// ============================================================================= - -struct PTOViewToMemrefPass - : public mlir::pto::impl::PTOViewToMemrefBase { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PTOViewToMemrefPass) - - void runOnOperation() override { - ModuleOp mod = getOperation(); - MLIRContext *ctx = &getContext(); - - for (auto func : mod.getOps()) { - // ------------------------------------------------------------------ - // Stage 0: ensure inttoptr values remain scalar-load/store only. - // ------------------------------------------------------------------ - if (failed(validateIntToPtrUses(func))) { - signalPassFailure(); - return; - } - - if (func.isExternal()) - continue; - - bool preserveTileABI = hasMigratedTileNativeOp(func); - bool tileNativeMainline = - preserveTileABI || hasTileNativeAllocationRoot(func); - - // ------------------------------------------------------------------ - // Stage 0.25: synthesize missing A2/A3 tquant tmp tiles before type lowering. - // ------------------------------------------------------------------ - if (failed(synthesizeMissingTQuantTmpOps(func, ctx))) { - signalPassFailure(); - return; - } - - // ------------------------------------------------------------------ - // Stage 0.5: keep pto.alloc_tile as a tile-native allocation root. - // - // Explicit-address alloc_tile already carries the physical address for - // level3. No-address alloc_tile is assigned an address by PTOPlanMemory. - // Neither case needs a pointer_cast detour here. - // ------------------------------------------------------------------ - - // ------------------------------------------------------------------ - // Stage 0.6: keep pto.alloc_multi_tile / pto.multi_tile_get tile-native. - // PlanMemory writes the physical slot addresses on alloc_multi_tile; - // sync consumes the slot selector directly from multi_tile_get, and - // PTOResolveBufferSelect materializes the selected alloc_tile handle. - // ------------------------------------------------------------------ - - // ------------------------------------------------------------------ - // Stage 0.75: keep pto.declare_tile tile-native. It is a symbolic handle - // whose address is assigned at runtime by tpop/tassign, not a static - // allocation root for PlanMemory. - // ------------------------------------------------------------------ - - // ------------------------------------------------------------------ - // Stage 0.8: normalize pto.tassign result type to match tile operand - // after tile_buf -> memref lowering (required for verifier consistency). - // ------------------------------------------------------------------ - DefaultInlineVector tassignOps; - func.walk([&](mlir::pto::TAssignOp op) { tassignOps.push_back(op); }); - for (auto op : tassignOps) { - Type targetTy = op.getTile().getType(); - if (op.getResult().getType() == targetTy) - continue; - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - auto normalized = - rewriter.create(op.getLoc(), targetTy, op.getTile(), - op.getAddr()); - rewriter.replaceOp(op, normalized.getResult()); - } - - // ------------------------------------------------------------------ - // Stage 0.9: Lower ptr-like pto.tile_buf_addr -> memref. - // The original authoring surface is pointer-like, so shape information - // is intentionally discarded here and only the linear address view is - // preserved for later memref-based rewriting. - // ------------------------------------------------------------------ - if (failed(lowerPtrLikeTileBufAddrOps(func, ctx))) { - signalPassFailure(); - return; - } - - // ------------------------------------------------------------------ - // Stage 1: Lower pto.make_tensor_view -> memref.reinterpret_cast - // ------------------------------------------------------------------ - // GM view ops stay authored until PTOResolveBufferSelect when a - // function contains the migrated view family. Legacy functions keep - // the established memref lowering below. - DefaultInlineVector makeViews; - func.walk([&](mlir::pto::MakeTensorViewOp op) { makeViews.push_back(op); }); - if (preserveTileABI) - makeViews.clear(); - - for (auto op : makeViews) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - Value baseBuf = op.getOperand(0); - OpFoldResult off0 = rewriter.getIndexAttr(0); - - auto baseMr = dyn_cast(baseBuf.getType()); - if (!baseMr) { - op.emitError("make_tensor_view base must be memref"); signalPassFailure(); return; - } - - // [修复] 获取动态 Rank (根据 shape 输入的数量) - size_t rank = op.getShape().size(); - - // Construct target type with dynamic offset/strides - Type elemTy = baseMr.getElementType(); - int64_t dyn = ShapedType::kDynamic; - - // [修复] 构建 N 维 Strided Layout - // strides 数组长度必须等于 rank - SmallVector dynStrides(rank, dyn); - auto layout = StridedLayoutAttr::get(ctx, /*offset=*/dyn, /*strides=*/dynStrides); - - // [修复] 构建 N 维 Shape - SmallVector dynShape(rank, dyn); - auto mrTy = MemRefType::get(dynShape, elemTy, layout, baseMr.getMemorySpace()); - - SmallInlineVector sizes; - for (Value v : op.getShape()) sizes.push_back(ensureIndex(rewriter, loc, v, op)); - - SmallInlineVector strides; - for (Value v : op.getStrides()) strides.push_back(ensureIndex(rewriter, loc, v, op)); - - auto rc = rewriter.create( - loc, mrTy, baseBuf, off0, sizes, strides); - if (auto layoutAttr = op.getLayoutAttr()) { - rc->setAttr("layout", layoutAttr); - } - - rewriter.replaceOp(op, rc.getResult()); - } - - // ------------------------------------------------------------------ - // Stage 1.25: Lower pto.get_tensor_view_dim -> memref.dim - // ------------------------------------------------------------------ - DefaultInlineVector tvDims; - func.walk([&](mlir::pto::GetTensorViewDimOp op) { tvDims.push_back(op); }); - - for (auto op : tvDims) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - Value view = op.getTensorView(); - auto mrTy = dyn_cast(view.getType()); - if (!mrTy) - continue; // leave it to later passes if it hasn't been lowered yet - - Value dimIdx = op.getDimIndex(); - Value dim = rewriter.create(loc, view, dimIdx); - rewriter.replaceOp(op, dim); - } - - // ------------------------------------------------------------------ - // Stage 1.3: Lower pto.partition_view -> memref.subview - // ------------------------------------------------------------------ - if (!preserveTileABI && failed(lowerPartitionViewOps(func, ctx))) { - signalPassFailure(); - return; - } - - // ------------------------------------------------------------------ - // Stage 1.35: pto.subview stays tile-native through memory and sync. - // ------------------------------------------------------------------ - - // ------------------------------------------------------------------ - // Stage 1.4: treshape/bitcast stay tile-native. - // ------------------------------------------------------------------ - // ------------------------------------------------------------------ - // Stage 1.5: Lower pto.get_tensor_view_stride -> strided memref metadata - // ------------------------------------------------------------------ - SmallVector tvStrides; - func.walk([&](mlir::pto::GetTensorViewStrideOp op) { tvStrides.push_back(op); }); - - for (auto op : tvStrides) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Location loc = op.getLoc(); - - Value view = op.getTensorView(); - auto mrTy = dyn_cast(view.getType()); - if (!mrTy) - continue; // leave it to later passes if it hasn't been lowered yet - - int64_t dimIndex = 0; - if (!getConstIndexValue(op.getDimIndex(), dimIndex)) { - op.emitError("get_tensor_view_stride currently expects a constant dim index"); - signalPassFailure(); - return; - } - if (dimIndex < 0 || dimIndex >= mrTy.getRank()) { - op.emitError("get_tensor_view_stride dim index is out of bounds"); - signalPassFailure(); - return; - } - - SmallVector staticStrides; - int64_t offset = ShapedType::kDynamic; - if (succeeded(mlir::pto::getPTOMemRefStridesAndOffset( - mrTy, staticStrides, offset)) && - dimIndex < (int64_t)staticStrides.size() && - staticStrides[dimIndex] != ShapedType::kDynamic) { - rewriter.replaceOpWithNewOp( - op, staticStrides[dimIndex]); - continue; - } - - auto metadata = - rewriter.create(loc, view); - rewriter.replaceOp(op, metadata.getStrides()[dimIndex]); - } - - // ------------------------------------------------------------------ - // Stage 3: Rewrite Compute Ops - // [关键] 全面使用 op->getOperand(i) 避免 Typed Accessor Crash - // ------------------------------------------------------------------ - if (!tileNativeMainline) { - - // --- TLoadOp [Src, Dst] --- - DefaultInlineVector loads; - func.walk([&](mlir::pto::TLoadOp op) { loads.push_back(op); }); - for (auto op : loads) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op->getOperand(0); - Value dst = op->getOperand(1); - - auto newOp = - rewriter.create(op.getLoc(), TypeRange{}, src, dst); - newOp->setAttrs(op->getAttrs()); - rewriter.replaceOp(op, newOp->getResults()); - } - - // --- TStoreOp [Src, Dst] --- - DefaultInlineVector storeops; - func.walk([&](mlir::pto::TStoreOp op) { storeops.push_back(op); }); - for (auto op : storeops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op->getOperand(0); - Value dst = op->getOperand(1); - Value preQuant = op.getPreQuantScalar(); - - pto::TStoreOp newOp; - if (preQuant) { - newOp = rewriter.create(op.getLoc(), TypeRange{}, - src, dst, preQuant); - } else { - newOp = rewriter.create(op.getLoc(), TypeRange{}, - src, dst, Value{}); - } - newOp->setAttrs(op->getAttrs()); - rewriter.replaceOp(op, newOp->getResults()); - } - - // --- TTransOp [Src, Tmp, Dst] --- - DefaultInlineVector trans; - func.walk([&](mlir::pto::TTransOp op) { trans.push_back(op); }); - for (auto op : trans) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex)); - } - - // --- TExpOp [Src, Dst] --- - DefaultInlineVector exp; - func.walk([&](mlir::pto::TExpOp op) { exp.push_back(op); }); - for (auto op : exp) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs(rewriter, op, TypeRange{}, - op->getOperand(0), - op->getOperand(1), op.getPrecisionTypeAttr()); - } - - // --- TMulOp [Src, Scalar, Dst] --- - DefaultInlineVector mul; - func.walk([&](mlir::pto::TMulOp op) { mul.push_back(op); }); - for (auto op : mul) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, op->getOperand(0), op.getOperand(1), - op->getOperand(kThirdOperandIndex)); - } - - // --- TMulSOp [Src, Scalar, Dst] --- - DefaultInlineVector muls; - func.walk([&](mlir::pto::TMulSOp op) { muls.push_back(op); }); - for (auto op : muls) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, op->getOperand(0), op.getScalar(), - op->getOperand(kThirdOperandIndex)); - } - - // --- TAddOp [Src0, Src1, Dst] --- - DefaultInlineVector addops; - func.walk([&](mlir::pto::TAddOp op) { addops.push_back(op); }); - for (auto op : addops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex)); - } - - // --- TMatmulOp [Lhs, Rhs, Dst] (no optional bias in ODS) --- - DefaultInlineVector matmuls; - func.walk([&](mlir::pto::TMatmulOp op) { matmuls.push_back(op); }); - for (auto op : matmuls) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - Value lhs = op->getOperand(0); - Value rhs = op->getOperand(1); - Value dst = op->getOperand(kThirdOperandIndex); - - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, lhs, rhs, dst, op.getAccPhaseAttr()); - } - - // --- TMatmulAccOp [Acc, Lhs, Rhs, Dst] --- - DefaultInlineVector matmulAccs; - func.walk([&](mlir::pto::TMatmulAccOp op) { matmulAccs.push_back(op); }); - for (auto op : matmulAccs) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), op.getAccPhaseAttr()); - } - - // --- TMatmulBiasOp [Acc, Lhs, Rhs, Bias, Dst] --- - DefaultInlineVector matmulBiass; - func.walk([&](mlir::pto::TMatmulBiasOp op) { matmulBiass.push_back(op); }); - for (auto op : matmulBiass) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), op.getAccPhaseAttr()); - } - - // --- TMatmulMxOp--- - DefaultInlineVector matmulMxs; - func.walk([&](mlir::pto::TMatmulMxOp op) { matmulMxs.push_back(op); }); - for (auto op : matmulMxs) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), - op->getOperand(kFifthOperandIndex), op.getAccPhaseAttr()); - } - - // --- TMatmulMxAccOp --- - DefaultInlineVector matmulMxAccs; - func.walk([&](mlir::pto::TMatmulMxAccOp op) { matmulMxAccs.push_back(op); }); - for (auto op : matmulMxAccs) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), - op->getOperand(kFifthOperandIndex), - op->getOperand(kSixthOperandIndex), op.getAccPhaseAttr()); - } - - // --- TMatmulMxBiasOp --- - DefaultInlineVector matmulMxBiass; - func.walk([&](mlir::pto::TMatmulMxBiasOp op) { matmulMxBiass.push_back(op); }); - for (auto op : matmulMxBiass) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), - op->getOperand(kFifthOperandIndex), - op->getOperand(kSixthOperandIndex)); - } - - // --- TGemvOp [Lhs, Rhs, Dst] --- - DefaultInlineVector gemvs; - func.walk([&](mlir::pto::TGemvOp op) { gemvs.push_back(op); }); - for (auto op : gemvs) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value lhs = op->getOperand(0); - Value rhs = op->getOperand(1); - Value dst = op->getOperand(kThirdOperandIndex); - - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, lhs, rhs, dst, op.getAccPhaseAttr()); - } - - // --- TGemvAccOp [Acc, Lhs, Rhs, Dst] --- - DefaultInlineVector gemvAccs; - func.walk([&](mlir::pto::TGemvAccOp op) { gemvAccs.push_back(op); }); - for (auto op : gemvAccs) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), op.getAccPhaseAttr()); - } - - // --- TGemvBiasOp [Acc, Lhs, Rhs, Bias, Dst] --- - DefaultInlineVector gemvBiass; - func.walk([&](mlir::pto::TGemvBiasOp op) { gemvBiass.push_back(op); }); - for (auto op : gemvBiass) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), op.getAccPhaseAttr()); - } - - // --- TGemvMxOp [A, AScale, B, BScale, Dst] --- - DefaultInlineVector gemvMxs; - func.walk([&](mlir::pto::TGemvMxOp op) { gemvMxs.push_back(op); }); - for (auto op : gemvMxs) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), - op->getOperand(kFifthOperandIndex), op.getAccPhaseAttr()); - } - - // --- TGemvMxAccOp [CIn, A, AScale, B, BScale, Dst] --- - DefaultInlineVector gemvMxAccs; - func.walk([&](mlir::pto::TGemvMxAccOp op) { gemvMxAccs.push_back(op); }); - for (auto op : gemvMxAccs) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), - op->getOperand(kFifthOperandIndex), - op->getOperand(kSixthOperandIndex), op.getAccPhaseAttr()); - } - - // --- TGemvMxBiasOp [A, AScale, B, BScale, Bias, Dst] --- - DefaultInlineVector gemvMxBiass; - func.walk([&](mlir::pto::TGemvMxBiasOp op) { gemvMxBiass.push_back(op); }); - for (auto op : gemvMxBiass) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, - op->getOperand(0), op->getOperand(1), - op->getOperand(kThirdOperandIndex), - op->getOperand(kFourthOperandIndex), - op->getOperand(kFifthOperandIndex), - op->getOperand(kSixthOperandIndex)); - } - - // --- TMovOp [Src, Dst] --- - DefaultInlineVector movs; - func.walk([&](mlir::pto::TMovOp op) { movs.push_back(op); }); - for (auto op : movs) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, op.getSrc(), op.getDst(), op.getFp(), - op.getPreQuantScalar(), op.getAccToVecModeAttr(), - op.getReluPreModeAttr()); - } - - DefaultInlineVector abseops; - func.walk([&](mlir::pto::TAbsOp op) { abseops.push_back(op); }); - - for (auto op : abseops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - dst); - } - - DefaultInlineVector addcops; - func.walk([&](mlir::pto::TAddCOp op) { addcops.push_back(op); }); - - for (auto op : addcops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value src2 = op.getSrc2(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto src2Ty = dyn_cast(src2.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !src2Ty ||!dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src0, - src1, - src2, - dst); - } - - DefaultInlineVector addsops; - func.walk([&](mlir::pto::TAddSOp op) { addsops.push_back(op); }); - - for (auto op : addsops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value scalar = op.getScalar(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs(rewriter, op, TypeRange{}, src, - scalar, dst); - } - - DefaultInlineVector addscops; - func.walk([&](mlir::pto::TAddSCOp op) { addscops.push_back(op); }); - - for (auto op : addscops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value scalar = op.getScalar(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src0, - scalar, - src1, - dst); - } - - DefaultInlineVector andops; - func.walk([&](mlir::pto::TAndOp op) { andops.push_back(op); }); - - for (auto op : andops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src0, - src1, - dst); - } - - DefaultInlineVector concats; - func.walk([&](mlir::pto::TConcatOp op) { concats.push_back(op); }); - - for (auto op : concats) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src0, - src1, - dst); - } - - DefaultInlineVector concatIdxs; - func.walk([&](mlir::pto::TConcatidxOp op) { concatIdxs.push_back(op); }); - - IRRewriter rewriter(ctx); - for (auto op : concatIdxs) { - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value src0Idx = op.getSrc0Idx(); - Value src1Idx = op.getSrc1Idx(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto src0IdxTy = dyn_cast(src0Idx.getType()); - auto src1IdxTy = dyn_cast(src1Idx.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !src0IdxTy || !src1IdxTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src0, - src1, - src0Idx, - src1Idx, - dst); - } - - DefaultInlineVector andsops; - func.walk([&](mlir::pto::TAndSOp op) { andsops.push_back(op); }); - - for (auto op : andsops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value scalar = op.getScalar(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - scalar, - dst); - } - - DefaultInlineVector ciops; - func.walk([&](mlir::pto::TCIOp op) { ciops.push_back(op); }); - - for (auto op : ciops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value s = op->getOperand(0); - Value tmp = op.getTmp(); - Value dst = op.getDst(); - bool descending = op.getDescending(); - - auto sTy = dyn_cast(s.getType()); - auto dstTy = dyn_cast(dst.getType()); - auto tmpTy = tmp ? dyn_cast(tmp.getType()) : MemRefType{}; - if (!sTy || !dstTy || (tmp && !tmpTy)) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - s, - tmp, - dst, - descending); - } - - DefaultInlineVector cmpops; - func.walk([&](mlir::pto::TCmpOp op) { cmpops.push_back(op); }); - - for (auto op : cmpops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src0, - src1, - dst); - } - - DefaultInlineVector cmpsops; - func.walk([&](mlir::pto::TCmpSOp op) { cmpsops.push_back(op); }); - - for (auto op : cmpsops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value scalar = op.getScalar(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - auto scalarTy = scalar.getType(); - bool scalarOk = - isa(scalarTy); // ScalarType in ODS: int/float - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - if (!scalarOk) { - op.emitError("expects scalar to be an integer or float type"); - signalPassFailure(); - return; - } - - auto cmpMode = op.getCmpModeAttr(); - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - scalar, - cmpMode, - dst); - } - - DefaultInlineVector colexpand; - func.walk([&](mlir::pto::TColExpandOp op) { colexpand.push_back(op); }); - - for (auto op : colexpand) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if ( !srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - dst); - } - - DefaultInlineVector colmaxops; - func.walk([&](mlir::pto::TColMaxOp op) { colmaxops.push_back(op); }); - - for (auto op : colmaxops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if ( !srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - dst); - } - - DefaultInlineVector colminops; - func.walk([&](mlir::pto::TColMinOp op) { colminops.push_back(op); }); - - for (auto op : colminops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if ( !srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - dst); - } - - DefaultInlineVector colexpandmulops; - func.walk([&](mlir::pto::TColExpandMulOp op) { - colexpandmulops.push_back(op); - }); - - for (auto op : colexpandmulops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src0, - src1, - dst); - } - - DefaultInlineVector colexpandmaxops; - func.walk([&](mlir::pto::TColExpandMaxOp op) { - colexpandmaxops.push_back(op); - }); - - for (auto op : colexpandmaxops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src0, - src1, - dst); - } - - DefaultInlineVector colexpandminops; - func.walk([&](mlir::pto::TColExpandMinOp op) { - colexpandminops.push_back(op); - }); - - for (auto op : colexpandminops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src0, - src1, - dst); - } - - DefaultInlineVector colsumops; - func.walk([&](mlir::pto::TColSumOp op) { colsumops.push_back(op); }); - - for (auto op : colsumops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - Value tmp = op.getTmp(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("src/dst are not memref yet"); - signalPassFailure(); - return; - } - - // If tmp exists, it must have isBinary attribute - if (tmp) { - auto tmpTy = dyn_cast(tmp.getType()); - if (!tmpTy) { - op.emitError("tmp is not memref yet"); - signalPassFailure(); - return; - } - - // Get isBinary attribute (should exist if tmp exists) - BoolAttr isBinaryAttr = op.getIsBinaryAttr(); - if (!isBinaryAttr) { - isBinaryAttr = BoolAttr::get(ctx, false); - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - tmp, - dst, - isBinaryAttr); - } else { - // Format 1: no tmp, no isBinary - // Use generic builder to avoid adding default isBinary attribute - SmallVector operands = {src, dst}; - SmallVector attrs; - // Copy all attributes except isBinary - for (auto attr : op->getAttrs()) { - if (attr.getName() != "isBinary") { - attrs.push_back(attr); - } - } - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - operands, - attrs); - } - } - - DefaultInlineVector cvtops; - func.walk([&](mlir::pto::TCvtOp op) { cvtops.push_back(op); }); - - for (auto op : cvtops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - auto rmodeAttr = op.getRmodeAttr(); // PTO_RoundModeAttr - auto satModeAttr = op.getSatModeAttr(); - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - dst, - rmodeAttr, - satModeAttr); - } - - DefaultInlineVector divops; - func.walk([&](mlir::pto::TDivOp op) { divops.push_back(op); }); - - for (auto op : divops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs(rewriter, - op, - TypeRange{}, - src0, - src1, - dst, - op.getPrecisionTypeAttr()); - } - - DefaultInlineVector divsops; - func.walk([&](mlir::pto::TDivSOp op) { divsops.push_back(op); }); - - for (auto op : divsops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value scale = op.getScalar(); - Value dst = op.getDst(); - - // Check types - they might still be TileBufType or already converted to MemRefType - auto srcTy = dyn_cast(src.getType()); - auto srcTileTy = dyn_cast(src.getType()); - auto scaleTileTy = dyn_cast(scale.getType()); - auto dstTy = dyn_cast(dst.getType()); - auto dstTileTy = dyn_cast(dst.getType()); - - // Determine which operand is tile-like and which is scalar-like. - // Keep the original operand order (set by parser textual form). - // Check if src is memref/tensor/tile (not scalar) - bool srcIsMemref = (srcTy != nullptr || srcTileTy != nullptr || - isa(src.getType()) || - isa(src.getType())); - // Check if scale is memref/tensor/tile (not scalar) - bool scaleIsMemref = (isa(scale.getType()) || - scaleTileTy != nullptr || - isa(scale.getType()) || - isa(scale.getType())); - - // Type validation - ensure we have the right types - if (!srcIsMemref && !scaleIsMemref) { - op.emitError("at least one operand (src or scale) must be tile_buf or memref"); - signalPassFailure(); - return; - } - if (srcIsMemref && scaleIsMemref) { - op.emitError("exactly one operand (src or scale) must be tile_buf or memref, the other must be scalar"); - signalPassFailure(); - return; - } - - if (!dstTy && !dstTileTy) { - op.emitError("dst operand must be tile_buf or memref"); - signalPassFailure(); - return; - } - replaceOpWithClonedAttrs(rewriter, - op, - TypeRange{}, - src, - scale, - dst, - op.getPrecisionTypeAttr()); - } - - DefaultInlineVector expandsops; - func.walk([&](mlir::pto::TExpandsOp op) { expandsops.push_back(op); }); - - for (auto op : expandsops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value scalar = op.getScalar(); - Value dst = op.getDst(); - - auto dstTy = dyn_cast(dst.getType()); - if (!dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs(rewriter, op, TypeRange{}, - scalar, dst); - } - - DefaultInlineVector extractops; - func.walk([&](mlir::pto::TExtractOp op) { extractops.push_back(op); }); - - for (auto op : extractops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value indexRow = op.getIndexRow(); - Value indexCol = op.getIndexCol(); - Value preQuantScalar = op.getPreQuantScalar(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto indexRowTy = dyn_cast(indexRow.getType()); - auto indexColTy = dyn_cast(indexCol.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !indexRowTy || !indexColTy || !dstTy) { - op.emitError("ins/outs are not correct yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, src, indexRow, indexCol, - preQuantScalar, dst, op.getAccToVecModeAttr(), - op.getReluPreModeAttr()); - } - - DefaultInlineVector extractfpops; - func.walk([&](mlir::pto::TExtractFPOp op) { extractfpops.push_back(op); }); - for (auto op : extractfpops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value fp = op.getFp(); - Value indexRow = op.getIndexRow(); - Value indexCol = op.getIndexCol(); - Value dst = op.getDst(); - - if (!dyn_cast(src.getType()) || - !dyn_cast(fp.getType()) || - !dyn_cast(indexRow.getType()) || - !dyn_cast(indexCol.getType()) || - !dyn_cast(dst.getType())) { - op.emitError("ins/outs are not correct yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, src, fp, indexRow, indexCol, dst, - op.getAccToVecModeAttr(), op.getReluPreModeAttr()); - } - - DefaultInlineVector insertops; - func.walk([&](mlir::pto::TInsertOp op) { insertops.push_back(op); }); - for (auto op : insertops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value indexRow = op.getIndexRow(); - Value indexCol = op.getIndexCol(); - Value fp = op.getFp(); - Value preQuantScalar = op.getPreQuantScalar(); - Value dst = op.getDst(); - - if (!dyn_cast(src.getType()) || - !dyn_cast(indexRow.getType()) || - !dyn_cast(indexCol.getType()) || - !dyn_cast(dst.getType()) || - (fp && !dyn_cast(fp.getType()))) { - op.emitError("ins/outs are not correct yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, src, indexRow, indexCol, dst, fp, - preQuantScalar, op.getAccToVecModeAttr(), op.getReluPreModeAttr(), - op.getTinsertModeAttr()); - } - - DefaultInlineVector insertfpops; - func.walk([&](mlir::pto::TInsertFPOp op) { insertfpops.push_back(op); }); - for (auto op : insertfpops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value fp = op.getFp(); - Value indexRow = op.getIndexRow(); - Value indexCol = op.getIndexCol(); - Value dst = op.getDst(); - - if (!dyn_cast(src.getType()) || - !dyn_cast(fp.getType()) || - !dyn_cast(indexRow.getType()) || - !dyn_cast(indexCol.getType()) || - !dyn_cast(dst.getType())) { - op.emitError("ins/outs are not correct yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, op, TypeRange{}, src, fp, indexRow, indexCol, dst, - op.getAccToVecModeAttr(), op.getReluPreModeAttr()); - } - - DefaultInlineVector fillpadops; - func.walk([&](mlir::pto::TFillPadOp op) { fillpadops.push_back(op); }); - - for (auto op : fillpadops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - dst, - op.getPadValueAttr()); - } - - DefaultInlineVector fillpadInplaceOps; - func.walk( - [&](mlir::pto::TFillPadInplaceOp op) { fillpadInplaceOps.push_back(op); }); - - for (auto op : fillpadInplaceOps) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - dst); - } - - // --- TSetValOp [Dst, Offset, Val] --- - // Lower tile-world scalar write to memref-world SETVAL DPS op. - DefaultInlineVector tsetvalops; - func.walk([&](mlir::pto::TSetValOp op) { tsetvalops.push_back(op); }); - - for (auto op : tsetvalops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value dst = op.getDst(); - Value offset = op.getOffset(); - Value val = op.getVal(); - - auto dstTy = dyn_cast(dst.getType()); - if (!dstTy) { - op.emitError("dst is not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - dst, - offset, - val); - } - - // --- TGetValOp [Src, Offset] -> Scalar --- - // Lower tile-world scalar read to memref-world GETVAL DPS op. - DefaultInlineVector tgetvalops; - func.walk([&](mlir::pto::TGetValOp op) { tgetvalops.push_back(op); }); - - for (auto op : tgetvalops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value offset = op.getOffset(); - Type dstType = op.getDst().getType(); - - auto srcTy = dyn_cast(src.getType()); - if (!srcTy) { - op.emitError("src is not memref yet"); - signalPassFailure(); - return; - } - - auto newOp = rewriter.create( - op.getLoc(), - dstType, - src, - offset); - rewriter.replaceOp(op, newOp.getDst()); - } - - DefaultInlineVector gatherops; - func.walk([&](mlir::pto::TGatherOp op) { gatherops.push_back(op); }); - - for (auto op : gatherops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - Value cdst = op.getCdst(); - Value indices = op.getIndices(); - Value tmp = op.getTmp(); - Value kValue = op.getKValue(); - auto maskPattern = op.getMaskPatternAttr(); - auto cmpMode = op.getCmpModeAttr(); - auto offset = op.getOffsetAttr(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - if (maskPattern) { - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - dst, - /*cdst=*/Value(), - /*indices=*/Value(), - /*tmp=*/Value(), - /*kValue=*/Value(), - /*maskPattern=*/maskPattern, - /*cmpMode=*/pto::CmpModeAttr(), - /*offset=*/IntegerAttr()); - continue; - } - - if (cdst || kValue) { - auto cdstTy = dyn_cast(cdst.getType()); - auto tmpTy = dyn_cast(tmp.getType()); - if (!cdstTy || !tmpTy) { - op.emitError("compare-form tgather expects cdst/tmp to be memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - dst, - cdst, - /*indices=*/Value(), - tmp, - kValue, - /*maskPattern=*/pto::MaskPatternAttr(), - cmpMode, - offset); - continue; - } - - if (indices || tmp) { - auto indicesTy = dyn_cast(indices.getType()); - auto tmpTy = dyn_cast(tmp.getType()); - if (!indicesTy || !tmpTy) { - op.emitError("index-form tgather expects indices/tmp to be memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - dst, - /*cdst=*/Value(), - indices, - tmp, - /*kValue=*/Value(), - /*maskPattern=*/pto::MaskPatternAttr(), - /*cmpMode=*/pto::CmpModeAttr(), - /*offset=*/IntegerAttr()); - continue; - } - - op.emitError("expects tgather to be in mask, index+tmp, or compare+tmp form"); - signalPassFailure(); - return; - } - - DefaultInlineVector gatherbops; - func.walk([&](mlir::pto::TGatherBOp op) { gatherbops.push_back(op); }); - - for (auto op : gatherbops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value offsets = op.getOffsets(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto offsetsTy = dyn_cast(offsets.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !offsetsTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - offsets, - dst); - } - - DefaultInlineVector logops; - func.walk([&](mlir::pto::TLogOp op) { logops.push_back(op); }); - - for (auto op : logops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - auto newOp = rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - dst, - op.getPrecisionTypeAttr()); - newOp->setAttrs(op->getAttrs()); - } - - DefaultInlineVector lreluops; - func.walk([&](mlir::pto::TLReluOp op) { lreluops.push_back(op); }); - - for (auto op : lreluops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value slope = op.getSlope(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto slopeTy = dyn_cast(slope.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !slopeTy || !dstTy) { - op.emitError("ins/outs are not correct type yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - slope, - dst); - } - - DefaultInlineVector maxops; - func.walk([&](mlir::pto::TMaxOp op) { maxops.push_back(op); }); - - for (auto op : maxops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs(rewriter, - op, - TypeRange{}, - src0, - src1, - dst); - } - - DefaultInlineVector maxsops; - func.walk([&](mlir::pto::TMaxSOp op) { maxsops.push_back(op); }); - - for (auto op : maxsops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value scalar = op.getScalar(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - bool scalarIsScalar = isa(scalar.getType()); - if (!srcTy || !scalarIsScalar || !dstTy) { - op.emitError("expects src/dst to be memref and scalar to be integer/float"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs(rewriter, - op, - TypeRange{}, - src, - scalar, - dst); - } - - DefaultInlineVector minops; - func.walk([&](mlir::pto::TMinOp op) { minops.push_back(op); }); - - for (auto op : minops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs(rewriter, - op, - TypeRange{}, - src0, - src1, - dst); - } - - DefaultInlineVector minsops; - func.walk([&](mlir::pto::TMinSOp op) { minsops.push_back(op); }); - - for (auto op : minsops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value scalar = op.getScalar(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - bool scalarIsScalar = isa(scalar.getType()); - if (!srcTy || !scalarIsScalar || !dstTy) { - op.emitError("expects src/dst to be memref and scalar to be integer/float"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs(rewriter, - op, - TypeRange{}, - src, - scalar, - dst); - } - - DefaultInlineVector movfpops; - func.walk([&](mlir::pto::TMovFPOp op) { movfpops.push_back(op); }); - - for (auto op : movfpops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value fp = op.getFp(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto fpTy = dyn_cast(fp.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !fpTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - fp, - dst); - } - - DefaultInlineVector quantops; - func.walk([&](mlir::pto::TQuantOp op) { quantops.push_back(op); }); - - for (auto op : quantops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value fp = op.getFp(); - Value tmp = op.getTmp(); - Value offset = op.getOffset(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto fpTy = dyn_cast(fp.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !fpTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - if (tmp && !dyn_cast(tmp.getType())) { - op.emitError("tmp is not memref yet"); - signalPassFailure(); - return; - } - if (offset && !dyn_cast(offset.getType())) { - op.emitError("offset is not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - fp, - offset, - tmp, - dst, - op.getQuantTypeAttr()); - } - - DefaultInlineVector mrgsortops; - func.walk([&](mlir::pto::TMrgSortOp op) { mrgsortops.push_back(op); }); - - for (auto op : mrgsortops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - if (op.isFormat1()) { - Value src = op.getSrc(); - Value dst = op.getDst(); - Value blockLenVal = op.getBlockLen(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - auto newOp = rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - ValueRange{src}, - blockLenVal, - ValueRange{dst}, - Value() /*tmp*/, - Value() /*excuted*/, - op.getExhaustedAttr()); - newOp->setAttrs(op->getAttrs()); - } else if (op.isFormat2()) { - bool allMemRef = true; - for (Value v : op.getSrcs()) - if (!dyn_cast(v.getType())) { allMemRef = false; break; } - if (!allMemRef) { - op.emitError("format2 ins/outs are not memref yet"); - signalPassFailure(); - return; - } - if (op.getDsts().size() != 1u || !op.getTmp()) { - op.emitError("format2 expects outs(dst) and ins(tmp)"); - signalPassFailure(); - return; - } - - Value dst = op.getDst(); - Value tmp = op.getTmp(); - Value excuted = op.getExcuted(); - if (!dyn_cast(dst.getType()) || !dyn_cast(tmp.getType())) { - op.emitError("format2 dst/tmp must be memref"); - signalPassFailure(); - return; - } - if (!dyn_cast(excuted.getType())) { - op.emitError("format2 outs(excuted) must be vector"); - signalPassFailure(); - return; - } - - auto newOp = rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - op.getSrcs(), - Value() /*blockLen*/, - ValueRange{dst}, - tmp, - excuted, - op.getExhaustedAttr()); - newOp->setAttrs(op->getAttrs()); - } else { - op.emitError("tmrgsort must be format1 or format2"); - signalPassFailure(); - return; - } - } - - DefaultInlineVector negops; - func.walk([&](mlir::pto::TNegOp op) { negops.push_back(op); }); - - for (auto op : negops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - dst); - } - - DefaultInlineVector notops; - func.walk([&](mlir::pto::TNotOp op) { notops.push_back(op); }); - - for (auto op : notops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - dst); - } - - DefaultInlineVector orops; - func.walk([&](mlir::pto::TOrOp op) { orops.push_back(op); }); - - for (auto op : orops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src0, - src1, - dst); - } - - DefaultInlineVector orsops; - func.walk([&](mlir::pto::TOrSOp op) { orsops.push_back(op); }); - - for (auto op : orsops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value scalar = op.getScalar(); - Value dst = op.getDst(); - - auto srcTy = dyn_cast(src.getType()); - auto scalarTy = dyn_cast(scalar.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!srcTy || !scalarTy || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - replaceOpWithClonedAttrs( - rewriter, - op, - TypeRange{}, - src, - scalar, - dst); - } - - DefaultInlineVector partaddops; - func.walk([&](mlir::pto::TPartAddOp op) { partaddops.push_back(op); }); - - for (auto op : partaddops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - auto attrs = op->getAttrs(); - auto newOp = rewriter.replaceOpWithNewOp( - op, - src0, - src1, - dst); - newOp->setAttrs(attrs); - } - - DefaultInlineVector partmulops; - func.walk([&](mlir::pto::TPartMulOp op) { partmulops.push_back(op); }); - - for (auto op : partmulops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src0 = op.getSrc0(); - Value src1 = op.getSrc1(); - Value dst = op.getDst(); - - auto src0Ty = dyn_cast(src0.getType()); - auto src1Ty = dyn_cast(src1.getType()); - auto dstTy = dyn_cast(dst.getType()); - if (!src0Ty || !src1Ty || !dstTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - auto attrs = op->getAttrs(); - auto newOp = rewriter.replaceOpWithNewOp( - op, - src0, - src1, - dst); - newOp->setAttrs(attrs); - } - - DefaultInlineVector mgatherops; - func.walk([&](mlir::pto::MGatherOp op) { mgatherops.push_back(op); }); - - for (auto op : mgatherops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value dst = op.getDst(); - Value idx = op.getIdx(); - Value mem = op.getMem(); - - auto dstTy = dyn_cast(dst.getType()); - auto idxTy = dyn_cast(idx.getType()); - auto memTy = dyn_cast(mem.getType()); - if (!dstTy || !idxTy || !memTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - mem, - idx, - dst, - /*scratch=*/op.getScratch(), - op.getCoalesceAttr(), - op.getGatherOobAttr()); - } - - DefaultInlineVector mascatterops; - func.walk([&](mlir::pto::MScatterOp op) { mascatterops.push_back(op); }); - - for (auto op : mascatterops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value idx = op.getIdx(); - Value mem = op.getMem(); - - auto srcTy = dyn_cast(src.getType()); - auto idxTy = dyn_cast(idx.getType()); - auto memTy = dyn_cast(mem.getType()); - if (!srcTy || !idxTy || !memTy) { - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - idx, - mem, - op.getCoalesceAttr(), - op.getScatterAtomicOpAttr(), - op.getScatterOobAttr(), - op.getScatterConflictAttr()); - } - DefaultInlineVector printops; - func.walk([&](mlir::pto::TPrintOp op) { printops.push_back(op); }); - - for (auto op : printops) { - IRRewriter rewriter(ctx); - rewriter.setInsertionPoint(op); - - Value src = op.getSrc(); - Value tmp = op->getNumOperands() > 1 ? op->getOperand(1) : Value(); - - auto srcTy = dyn_cast(src.getType()); - auto tmpTy = tmp ? dyn_cast(tmp.getType()) : MemRefType(); - if (!srcTy || (tmp && !tmpTy)) { - if (isa(src.getType()) && - (!tmp || isa(tmp.getType()))) - continue; - op.emitError("ins/outs are not memref yet"); - signalPassFailure(); - return; - } - - rewriter.replaceOpWithNewOp( - op, - TypeRange{}, - src, - tmp, - dyn_cast_or_null( - op.getProperties().printFormat)); - } - } - - } - - // Debug Output - LLVM_DEBUG(llvm::dbgs() << mod.getOperation()); - } -}; - -} // namespace - -std::unique_ptr createPTOViewToMemrefPass() { - return std::make_unique(); -} - -} // namespace pto -} // namespace mlir diff --git a/test/lit/pto/add_carry_tile_native.pto b/test/lit/pto/add_carry_tile_native.pto index 7146fb03fe..5e6e3dcc02 100644 --- a/test/lit/pto/add_carry_tile_native.pto +++ b/test/lit/pto/add_carry_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/addmul_scalar_tile_native.pto b/test/lit/pto/addmul_scalar_tile_native.pto index ed75b9bf4c..815a75b3b9 100644 --- a/test/lit/pto/addmul_scalar_tile_native.pto +++ b/test/lit/pto/addmul_scalar_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/addmul_tile_native.pto b/test/lit/pto/addmul_tile_native.pto index eac26c4a74..c63bafb0ba 100644 --- a/test/lit/pto/addmul_tile_native.pto +++ b/test/lit/pto/addmul_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto b/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto index 72437129f2..eaa6423049 100644 --- a/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto +++ b/test/lit/pto/alloc_tile_plan_memory_no_memref_alloc.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level2 --pto-arch=a3 --emit-pto-ir \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=VIEW +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=VIEW // RUN: ptoas --pto-level=level2 --pto-arch=a3 --emit-pto-ir \ // RUN: --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN // RUN: ptoas --pto-level=level2 --pto-arch=a3 %s >/dev/null diff --git a/test/lit/pto/arg_reduction_tile_native.pto b/test/lit/pto/arg_reduction_tile_native.pto index 6f8e728f25..4c8c9cae6f 100644 --- a/test/lit/pto/arg_reduction_tile_native.pto +++ b/test/lit/pto/arg_reduction_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/axpy_dequant_tile_native.pto b/test/lit/pto/axpy_dequant_tile_native.pto index 9efd14142b..4be75d6c5b 100644 --- a/test/lit/pto/axpy_dequant_tile_native.pto +++ b/test/lit/pto/axpy_dequant_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/basic_float_tile_native.pto b/test/lit/pto/basic_float_tile_native.pto index 4788a0812b..26f4559e38 100644 --- a/test/lit/pto/basic_float_tile_native.pto +++ b/test/lit/pto/basic_float_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/bitcast_tile_native.pto b/test/lit/pto/bitcast_tile_native.pto index 55581dc82b..66c59ecd55 100644 --- a/test/lit/pto/bitcast_tile_native.pto +++ b/test/lit/pto/bitcast_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o - 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o - 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE module attributes {pto.target_arch = "a5"} { diff --git a/test/lit/pto/build_async_session_tile_native.pto b/test/lit/pto/build_async_session_tile_native.pto index 9c43766337..7fb6f327d8 100644 --- a/test/lit/pto/build_async_session_tile_native.pto +++ b/test/lit/pto/build_async_session_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-arch=a3 --pto-level=level3 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 --pto-level=level3 %s -o - 2>&1 \ // RUN: | FileCheck %s --check-prefix=EMITC @@ -23,7 +23,7 @@ module { } } -// NATIVE: IR Dump After PTOViewToMemref +// NATIVE: IR Dump After PTOResolveReservedBuffers // NATIVE-LABEL: func.func @build_async_session_tile_native( // NATIVE-SAME: %[[SCRATCH:.*]]: !pto.tile_buf // NATIVE-SAME: %[[WORKSPACE:.*]]: !pto.ptr diff --git a/test/lit/pto/castptr_emitc_tile_and_int.pto b/test/lit/pto/castptr_emitc_tile_and_int.pto index d496bade74..255700c48e 100644 --- a/test/lit/pto/castptr_emitc_tile_and_int.pto +++ b/test/lit/pto/castptr_emitc_tile_and_int.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-level=level3 %s -o - 2>&1 | FileCheck %s -// RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --cann-output-version=9.0.0 --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { func.func @castptr_ptr_to_int(%src: memref<1xf32, #pto.address_space>) @@ -29,6 +29,6 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind -> i64 diff --git a/test/lit/pto/col_minmax_tile_native.pto b/test/lit/pto/col_minmax_tile_native.pto index 24e7a0678e..f618c20e9a 100644 --- a/test/lit/pto/col_minmax_tile_native.pto +++ b/test/lit/pto/col_minmax_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/colexpand_tile_native.pto b/test/lit/pto/colexpand_tile_native.pto index a252693183..1b34eff01e 100644 --- a/test/lit/pto/colexpand_tile_native.pto +++ b/test/lit/pto/colexpand_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/colsum_tile_native.pto b/test/lit/pto/colsum_tile_native.pto index 4cb9c94017..193f96504a 100644 --- a/test/lit/pto/colsum_tile_native.pto +++ b/test/lit/pto/colsum_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/compact_left_blayout_parser_a3.pto b/test/lit/pto/compact_left_blayout_parser_a3.pto index 271cdcafd6..b64df93341 100644 --- a/test/lit/pto/compact_left_blayout_parser_a3.pto +++ b/test/lit/pto/compact_left_blayout_parser_a3.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch a3 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s module attributes {"pto.target_arch" = "a3"} { func.func @compact_left_blayout_parser_a3() { diff --git a/test/lit/pto/compact_left_blayout_parser_a5.pto b/test/lit/pto/compact_left_blayout_parser_a5.pto index 23a5c1b6e5..f34458bd1b 100644 --- a/test/lit/pto/compact_left_blayout_parser_a5.pto +++ b/test/lit/pto/compact_left_blayout_parser_a5.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch a5 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s +// RUN: ptoas --pto-arch a5 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s module attributes {"pto.target_arch" = "a5"} { func.func @compact_left_blayout_parser_a5() { diff --git a/test/lit/pto/compare_tile_native.pto b/test/lit/pto/compare_tile_native.pto index 1b09451969..62be2e6b42 100644 --- a/test/lit/pto/compare_tile_native.pto +++ b/test/lit/pto/compare_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/concat_tile_native.pto b/test/lit/pto/concat_tile_native.pto index 81cf638414..09f823b97a 100644 --- a/test/lit/pto/concat_tile_native.pto +++ b/test/lit/pto/concat_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/cvt_tile_native.pto b/test/lit/pto/cvt_tile_native.pto index c65d57e1cf..af37bd6e13 100644 --- a/test/lit/pto/cvt_tile_native.pto +++ b/test/lit/pto/cvt_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/declare_tile_tile_native.pto b/test/lit/pto/declare_tile_tile_native.pto index cea347f13a..bd2e8a6dab 100644 --- a/test/lit/pto/declare_tile_tile_native.pto +++ b/test/lit/pto/declare_tile_tile_native.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE diff --git a/test/lit/pto/div_tile_native.pto b/test/lit/pto/div_tile_native.pto index 0868d36a9c..61239a4547 100644 --- a/test/lit/pto/div_tile_native.pto +++ b/test/lit/pto/div_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/expands_tile_native.pto b/test/lit/pto/expands_tile_native.pto index 1d16f49493..e9e4a4a0a2 100644 --- a/test/lit/pto/expands_tile_native.pto +++ b/test/lit/pto/expands_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/extract_insert_fp_tile_native.pto b/test/lit/pto/extract_insert_fp_tile_native.pto index cbbeb7514e..aa4e3f3f45 100644 --- a/test/lit/pto/extract_insert_fp_tile_native.pto +++ b/test/lit/pto/extract_insert_fp_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/extract_insert_tile_native.pto b/test/lit/pto/extract_insert_tile_native.pto index c5e4ba7fd2..7074684aa0 100644 --- a/test/lit/pto/extract_insert_tile_native.pto +++ b/test/lit/pto/extract_insert_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/fillpad_tile_native.pto b/test/lit/pto/fillpad_tile_native.pto index 535d264920..dc24cacc0f 100644 --- a/test/lit/pto/fillpad_tile_native.pto +++ b/test/lit/pto/fillpad_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/gather_tile_native.pto b/test/lit/pto/gather_tile_native.pto index 8ebc4a33f8..7bb90c5e2e 100644 --- a/test/lit/pto/gather_tile_native.pto +++ b/test/lit/pto/gather_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/gemv_tile_native.pto b/test/lit/pto/gemv_tile_native.pto index 1134f71fc0..ab80df500d 100644 --- a/test/lit/pto/gemv_tile_native.pto +++ b/test/lit/pto/gemv_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/gm_tensor_view_native_through_sync.pto b/test/lit/pto/gm_tensor_view_native_through_sync.pto index 36193cab32..a1230f0f60 100644 --- a/test/lit/pto/gm_tensor_view_native_through_sync.pto +++ b/test/lit/pto/gm_tensor_view_native_through_sync.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-buffer-select %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=RESOLVED module { @@ -29,7 +29,7 @@ module { } } -// NATIVE: IR Dump After PTOViewToMemref +// NATIVE: IR Dump After PTOResolveReservedBuffers // NATIVE: pto.make_tensor_view // NATIVE: pto.get_tensor_view_dim // NATIVE: pto.get_tensor_view_stride diff --git a/test/lit/pto/graph_sync_solver_basic.pto b/test/lit/pto/graph_sync_solver_basic.pto index dd34184526..6eb2814bff 100644 --- a/test/lit/pto/graph_sync_solver_basic.pto +++ b/test/lit/pto/graph_sync_solver_basic.pto @@ -6,7 +6,7 @@ // directly instead of going through the CCEC backend renderer. // // All three functions live in the same module: ptoas's pipeline contains -// ModuleOp-scoped passes (pto-view-to-memref, pto-resolve-reserved-buffers) +// ModuleOp-scoped passes (pto-resolve-reserved-buffers, pto-resolve-reserved-buffers) // that won't visit nested modules introduced by `-split-input-file`. module { diff --git a/test/lit/pto/issue31_partition_view_parser_compat.pto b/test/lit/pto/issue31_partition_view_parser_compat.pto index f6f5bfef8b..afadd750de 100755 --- a/test/lit/pto/issue31_partition_view_parser_compat.pto +++ b/test/lit/pto/issue31_partition_view_parser_compat.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --mlir-print-ir-before=pto-view-to-memref %s 2>&1 | FileCheck %s +// RUN: ptoas --mlir-print-ir-before=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s module { func.func @new_format_static(%src: !pto.ptr) { @@ -19,7 +19,8 @@ module { return } - func.func @old_format_static(%src: !pto.ptr) { + func.func @old_format_static(%src: !pto.ptr) + -> !pto.partition_tensor_view<16x32xf32> { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index @@ -29,10 +30,11 @@ module { : !pto.tensor_view %sv = pto.partition_view %tv, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.tensor_view -> !pto.partition_tensor_view<16x32xf32> - return + return %sv : !pto.partition_tensor_view<16x32xf32> } - func.func @old_format_dynamic(%src: !pto.ptr, %m: index, %n: index) { + func.func @old_format_dynamic(%src: !pto.ptr, %m: index, %n: index) + -> !pto.partition_tensor_view { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c64 = arith.constant 64 : index @@ -41,7 +43,7 @@ module { : !pto.tensor_view %sv = pto.partition_view %tv, offsets = [%c0, %c0], sizes = [%m, %n] : !pto.tensor_view -> !pto.partition_tensor_view - return + return %sv : !pto.partition_tensor_view } } diff --git a/test/lit/pto/issue686_a5_tmov_treshape_square_dynamic_valid_shape.pto b/test/lit/pto/issue686_a5_tmov_treshape_square_dynamic_valid_shape.pto index 8b41fc5a0e..8c11c78d1c 100644 --- a/test/lit/pto/issue686_a5_tmov_treshape_square_dynamic_valid_shape.pto +++ b/test/lit/pto/issue686_a5_tmov_treshape_square_dynamic_valid_shape.pto @@ -1,5 +1,5 @@ // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o - 2>&1 | FileCheck %s +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o - 2>&1 | FileCheck %s module attributes {pto.target_arch = "a5"} { func.func @a5_tmov_treshape_square_dynamic_valid_shape() { diff --git a/test/lit/pto/load_store_tile_native.pto b/test/lit/pto/load_store_tile_native.pto index 2620bfc1a4..f2bfcc6009 100644 --- a/test/lit/pto/load_store_tile_native.pto +++ b/test/lit/pto/load_store_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC @@ -31,9 +31,9 @@ module { } // NATIVE-LABEL: func.func private @tload_arg( -// NATIVE: pto.tload ins(%arg0 : memref<16x16xf32>) outs(%arg1 : !pto.tile_buf) +// NATIVE: pto.tload ins(%arg0 : !pto.partition_tensor_view<16x16xf32>) outs(%arg1 : !pto.tile_buf) // NATIVE-LABEL: func.func private @tstore_arg( -// NATIVE: pto.tstore ins(%arg0 : !pto.tile_buf) outs(%arg1 : memref<16x16xf32>) +// NATIVE: pto.tstore ins(%arg0 : !pto.tile_buf) outs(%arg1 : !pto.partition_tensor_view<16x16xf32>) // EMITC-LABEL: tload_arg( // EMITC: TLOAD( diff --git a/test/lit/pto/logic_binary_tile_native.pto b/test/lit/pto/logic_binary_tile_native.pto index 7a29ee6e5e..81d889b7c7 100644 --- a/test/lit/pto/logic_binary_tile_native.pto +++ b/test/lit/pto/logic_binary_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/logic_scalar_tile_native.pto b/test/lit/pto/logic_scalar_tile_native.pto index 2d55a6ec0a..809269529a 100644 --- a/test/lit/pto/logic_scalar_tile_native.pto +++ b/test/lit/pto/logic_scalar_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/materialize_tile_handles_control_flow_result.pto b/test/lit/pto/materialize_tile_handles_control_flow_result.pto index 3548bd86c9..8987f119f3 100644 --- a/test/lit/pto/materialize_tile_handles_control_flow_result.pto +++ b/test/lit/pto/materialize_tile_handles_control_flow_result.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO @@ -50,7 +50,7 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { } } -// NATIVE: IR Dump After PTOViewToMemref +// NATIVE: IR Dump After PTOResolveReservedBuffers // MATERIALIZED: IR Dump After PTOMaterializeTileHandles // COMMON: %[[SEL:.*]] = scf.if {{.*}} -> (!pto.tile_buf) { // COMMON: pto.tadd diff --git a/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto b/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto index 450a5d4cc8..ae7bafd536 100644 --- a/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto +++ b/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto @@ -8,7 +8,7 @@ // Guards PTODSL subkernel helper arguments staying in tile_buf ABI when the // helper directly extracts their address views. -// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON // RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON // RUN: ptoas --pto-arch=a5 --pto-level=level3 %s -o /dev/null @@ -36,7 +36,7 @@ module { !pto.ptr, !pto.ptr, f32) } -// NATIVE: IR Dump After PTOViewToMemref +// NATIVE: IR Dump After PTOResolveReservedBuffers // MATERIALIZED: IR Dump After PTOMaterializeTileHandles // COMMON-LABEL: func.func @ptodsl_subkernel_helper_tile_abi( // COMMON-SAME: %{{.*}}: !pto.tile_buf diff --git a/test/lit/pto/matmul_tile_native.pto b/test/lit/pto/matmul_tile_native.pto index 62745e4e36..c72bfeaf3a 100644 --- a/test/lit/pto/matmul_tile_native.pto +++ b/test/lit/pto/matmul_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/mgather_gm2l1_sync_a5.pto b/test/lit/pto/mgather_gm2l1_sync_a5.pto index 91f51ff8ea..4dc59f9c85 100644 --- a/test/lit/pto/mgather_gm2l1_sync_a5.pto +++ b/test/lit/pto/mgather_gm2l1_sync_a5.pto @@ -6,7 +6,7 @@ // 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. -// Regression for the A5 GM->L1 mgather pipe. PTOViewToMemref rewrites the mat +// Regression for the A5 GM->L1 mgather pipe. PTOResolveReservedBuffers rewrites the mat // destination into a memref<...,mat> before InsertSync / GraphSyncSolver run, so // MGatherOp::getPipe() must resolve the MAT address space from the lowered // memref form too. If it only matched a TileBufType, A5 would fall through to diff --git a/test/lit/pto/mgather_gm2ub_a5_elem_1x1_memref_macro_sync_model.pto b/test/lit/pto/mgather_gm2ub_a5_elem_1x1_memref_macro_sync_model.pto index 3a7b79859b..ba146e0b64 100644 --- a/test/lit/pto/mgather_gm2ub_a5_elem_1x1_memref_macro_sync_model.pto +++ b/test/lit/pto/mgather_gm2ub_a5_elem_1x1_memref_macro_sync_model.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // A5 elem-coalesced 1x1 GM->UB mgather lowers to the scalar overload even when -// InsertSync sees the dst after PTOViewToMemref lowering. The macro model must +// InsertSync sees the dst after PTOResolveReservedBuffers lowering. The macro model must // recover the 1x1 shape from MemRefType instead of requiring TileBufType. // // RUN: ptoas --pto-arch=a5 --enable-insert-sync --pto-insert-sync-debug=2 %s -o - 2>&1 | FileCheck %s diff --git a/test/lit/pto/mgather_mscatter_tile_native.pto b/test/lit/pto/mgather_mscatter_tile_native.pto index 7ff3d63027..bbdd0173af 100644 --- a/test/lit/pto/mgather_mscatter_tile_native.pto +++ b/test/lit/pto/mgather_mscatter_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/minmax_scalar_tile_native.pto b/test/lit/pto/minmax_scalar_tile_native.pto index 5b9c6a22ea..2faf76bf44 100644 --- a/test/lit/pto/minmax_scalar_tile_native.pto +++ b/test/lit/pto/minmax_scalar_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/minmax_tile_native.pto b/test/lit/pto/minmax_tile_native.pto index deb676e216..4f79a51315 100644 --- a/test/lit/pto/minmax_tile_native.pto +++ b/test/lit/pto/minmax_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/mov_tile_native.pto b/test/lit/pto/mov_tile_native.pto index 71ba65ad91..f0828f3548 100644 --- a/test/lit/pto/mov_tile_native.pto +++ b/test/lit/pto/mov_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/movement_metadata_tile_native.pto b/test/lit/pto/movement_metadata_tile_native.pto index 1b4999c9c1..7120bb3178 100644 --- a/test/lit/pto/movement_metadata_tile_native.pto +++ b/test/lit/pto/movement_metadata_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/mrgsort_tile_native.pto b/test/lit/pto/mrgsort_tile_native.pto index fcba49ef06..23729a549b 100644 --- a/test/lit/pto/mrgsort_tile_native.pto +++ b/test/lit/pto/mrgsort_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/multi_tile_buf_type_parse_print.pto b/test/lit/pto/multi_tile_buf_type_parse_print.pto index 90b9f47e9c..195c346024 100644 --- a/test/lit/pto/multi_tile_buf_type_parse_print.pto +++ b/test/lit/pto/multi_tile_buf_type_parse_print.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-before=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --mlir-print-ir-before=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s // Verifies that both the compact and verbose forms of `!pto.multi_tile_buf` // parse and round-trip through the printer. The compact form (preferred) diff --git a/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto b/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto index 225922f5d7..7ed312799d 100644 --- a/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto +++ b/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s // RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=MAT // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC @@ -43,7 +43,7 @@ module { } } -// CHECK: IR Dump After PTOViewToMemref +// CHECK: IR Dump After PTOResolveReservedBuffers // CHECK-LABEL: func.func @const_preload_dyn_loop_select // CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile // CHECK: pto.multi_tile_get %[[MULTI]][%c0 diff --git a/test/lit/pto/multi_tile_get_const_slot_lowering.pto b/test/lit/pto/multi_tile_get_const_slot_lowering.pto index a411d88b85..e96fd6a7ac 100644 --- a/test/lit/pto/multi_tile_get_const_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_const_slot_lowering.pto @@ -6,14 +6,14 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s // RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN // RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT // RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=PLAN // RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT // Verifies the tile-native multi-buffer pipeline for constant slots: -// 1. PTOViewToMemref preserves alloc_multi_tile and multi_tile_get. +// 1. PTOResolveReservedBuffers preserves alloc_multi_tile and multi_tile_get. // 2. PTOPlanMemory records the physical slot addresses on alloc_multi_tile. // 3. PTOResolveBufferSelect replaces each selection with alloc_tile(addrK). @@ -42,7 +42,7 @@ module { } } -// CHECK: IR Dump After PTOViewToMemref +// CHECK: IR Dump After PTOResolveReservedBuffers // CHECK-LABEL: func.func @const_slot_two_buffers // CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile // CHECK: pto.multi_tile_get %[[MULTI]][%c0 diff --git a/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto b/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto index e96049174d..1ca5908c3c 100644 --- a/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s // RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT // Verifies that tile-native `multi_tile_get` preserves the user-supplied slot @@ -44,7 +44,7 @@ module { } } -// CHECK: IR Dump After PTOViewToMemref +// CHECK: IR Dump After PTOResolveReservedBuffers // CHECK-LABEL: func.func @dyn_slot_prefetch // CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile // CHECK: %[[NEXT:.*]] = arith.addi %{{.*}}, %c1 diff --git a/test/lit/pto/multi_tile_level3_explicit_addr.pto b/test/lit/pto/multi_tile_level3_explicit_addr.pto index 6289855628..6b5a06a9a4 100644 --- a/test/lit/pto/multi_tile_level3_explicit_addr.pto +++ b/test/lit/pto/multi_tile_level3_explicit_addr.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s // RUN: ptoas --pto-level=level3 --pto-arch=a5 --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT // RUN: not ptoas --pto-level=level1 --pto-arch=a5 %s 2>&1 | FileCheck %s --check-prefix=LVERR // RUN: not ptoas --pto-level=level2 --pto-arch=a5 %s 2>&1 | FileCheck %s --check-prefix=LVERR @@ -41,8 +41,8 @@ module { } } -// PTOViewToMemref preserves the tile-native allocation and selections. -// CHECK: IR Dump After PTOViewToMemref +// PTOResolveReservedBuffers preserves the tile-native allocation and selections. +// CHECK: IR Dump After PTOResolveReservedBuffers // CHECK-LABEL: func.func @l3_two_buffers // CHECK: %[[BASE:.*]] = arith.constant 0 : i64 // CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile addr = %[[BASE]] diff --git a/test/lit/pto/multi_tile_no_loop_unroll.pto b/test/lit/pto/multi_tile_no_loop_unroll.pto index b659bfc7fd..43682523bc 100644 --- a/test/lit/pto/multi_tile_no_loop_unroll.pto +++ b/test/lit/pto/multi_tile_no_loop_unroll.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 1>/dev/null | FileCheck %s +// RUN: ptoas --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s // Verifies that tile-native multi-buffer selection works without an enclosing // `scf.for`; the two constant slots do not depend on a loop induction value. @@ -36,7 +36,7 @@ module { } } -// CHECK: IR Dump After PTOViewToMemref +// CHECK: IR Dump After PTOResolveReservedBuffers // CHECK-LABEL: func.func @unroll_two_slots // CHECK-NOT: scf.for // CHECK: %[[MULTI:.*]] = pto.alloc_multi_tile diff --git a/test/lit/pto/part_arithmetic_tile_native.pto b/test/lit/pto/part_arithmetic_tile_native.pto index 085a223b9c..393162f4db 100644 --- a/test/lit/pto/part_arithmetic_tile_native.pto +++ b/test/lit/pto/part_arithmetic_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/pipe_tile_native.pto b/test/lit/pto/pipe_tile_native.pto index c187386d6d..8e04257b22 100644 --- a/test/lit/pto/pipe_tile_native.pto +++ b/test/lit/pto/pipe_tile_native.pto @@ -6,7 +6,7 @@ // 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=a3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE module { func.func private @push_arg(%gm: memref<256xf32, #pto.address_space>, %consumer: i32, %tile: !pto.tile_buf) attributes {pto.kernel_kind = #pto.kernel_kind} { diff --git a/test/lit/pto/pow_rsqrt_tile_native.pto b/test/lit/pto/pow_rsqrt_tile_native.pto index 61306ba443..7f05f55766 100644 --- a/test/lit/pto/pow_rsqrt_tile_native.pto +++ b/test/lit/pto/pow_rsqrt_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/print_tile_native.pto b/test/lit/pto/print_tile_native.pto index dec7f9d0ef..d33e50f7de 100644 --- a/test/lit/pto/print_tile_native.pto +++ b/test/lit/pto/print_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a5 %s \ // RUN: | FileCheck %s --check-prefix=EMITC @@ -31,7 +31,7 @@ module { // NATIVE-LABEL: func.func private @tprint_arg( // NATIVE: pto.tprint ins(%arg0 : !pto.tile_buf) // NATIVE-LABEL: func.func private @tprint_tmp_arg( -// NATIVE: pto.tprint ins(%arg0, %arg1 : !pto.tile_buf, memref<16x16xf32>) +// NATIVE: pto.tprint ins(%arg0, %arg1 : !pto.tile_buf, !pto.partition_tensor_view<16x16xf32>) // EMITC-LABEL: tprint_arg( // EMITC: TPRINT( diff --git a/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto b/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto index de967da6e3..3a5c05124f 100644 --- a/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto +++ b/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // Guards late PTODSL helper inlining after tile handle materialization. -// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON // RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON // RUN: ptoas --pto-arch=a5 --pto-level=level3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC // RUN: ptoas --pto-arch=a5 --pto-level=level3 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO @@ -53,7 +53,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind, f32) } -// NATIVE: IR Dump After PTOViewToMemref +// NATIVE: IR Dump After PTOResolveReservedBuffers // MATERIALIZED: IR Dump After PTOMaterializeTileHandles // COMMON-LABEL: func.func @ptodsl_subkernel_helper_tile_abi_emitc // COMMON: %[[ALLOC:.*]] = pto.alloc_tile addr = {{.*}} : !pto.tile_buf diff --git a/test/lit/pto/ptr_int_cast.pto b/test/lit/pto/ptr_int_cast.pto index 39eb5c78f5..ecd2da917f 100644 --- a/test/lit/pto/ptr_int_cast.pto +++ b/test/lit/pto/ptr_int_cast.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=CPP // RUN: ptoas --pto-arch=a3 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO diff --git a/test/lit/pto/quant_mx_tile_native.pto b/test/lit/pto/quant_mx_tile_native.pto index 5ee9681dd9..27d08760fd 100644 --- a/test/lit/pto/quant_mx_tile_native.pto +++ b/test/lit/pto/quant_mx_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/reduction_tile_native.pto b/test/lit/pto/reduction_tile_native.pto index b6c34126c2..4826d37448 100644 --- a/test/lit/pto/reduction_tile_native.pto +++ b/test/lit/pto/reduction_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/rowexpand_tile_native.pto b/test/lit/pto/rowexpand_tile_native.pto index 26ed6b0735..61d9f7a4ec 100644 --- a/test/lit/pto/rowexpand_tile_native.pto +++ b/test/lit/pto/rowexpand_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/scratch_elementwise_tile_native.pto b/test/lit/pto/scratch_elementwise_tile_native.pto index 32826f0de6..63f30f5a36 100644 --- a/test/lit/pto/scratch_elementwise_tile_native.pto +++ b/test/lit/pto/scratch_elementwise_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/select_tile_native.pto b/test/lit/pto/select_tile_native.pto index a0501b5fd9..a0ba395ba0 100644 --- a/test/lit/pto/select_tile_native.pto +++ b/test/lit/pto/select_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/shift_tile_native.pto b/test/lit/pto/shift_tile_native.pto index 84a3dbabe2..eb0fa6d5a2 100644 --- a/test/lit/pto/shift_tile_native.pto +++ b/test/lit/pto/shift_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/sort32_tile_native.pto b/test/lit/pto/sort32_tile_native.pto index 4740d2e397..0f06041f9a 100644 --- a/test/lit/pto/sort32_tile_native.pto +++ b/test/lit/pto/sort32_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/pto/store_fp_tile_native.pto b/test/lit/pto/store_fp_tile_native.pto index 26c7974d2a..ea25964b4c 100644 --- a/test/lit/pto/store_fp_tile_native.pto +++ b/test/lit/pto/store_fp_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/sub_tile_native.pto b/test/lit/pto/sub_tile_native.pto index adfbaa648a..febff85943 100644 --- a/test/lit/pto/sub_tile_native.pto +++ b/test/lit/pto/sub_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto index 1646298be1..7dde92792d 100644 --- a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto +++ b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto @@ -1,5 +1,5 @@ // RUN: ptoas --emit-pto-ir %s -o - | FileCheck %s -// RUN: ptoas %s --mlir-print-ir-after=pto-view-to-memref -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas %s --mlir-print-ir-after=pto-resolve-reserved-buffers -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE module { func.func @subview_dynamic_offset_static_valid_regression( @@ -29,6 +29,6 @@ module { // CHECK: %[[ADDR:.*]] = arith.addi %c0_i64, %[[BYTE_OFFSET]] : i64 // CHECK: pto.alloc_tile addr = %[[ADDR]] {pto.view_semantics = "subview"} : !pto.tile_buf -// NATIVE: IR Dump After PTOViewToMemref +// NATIVE: IR Dump After PTOResolveReservedBuffers // NATIVE: pto.subview %{{.*}}[%c0, %{{.*}}] sizes [1, 64] : !pto.tile_buf -> !pto.tile_buf // NATIVE-NOT: memref.subview diff --git a/test/lit/pto/tassign_reject_memref.pto b/test/lit/pto/tassign_reject_memref.pto new file mode 100644 index 0000000000..d730dbe89f --- /dev/null +++ b/test/lit/pto/tassign_reject_memref.pto @@ -0,0 +1,23 @@ +// 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-level=level3 --pto-arch=a5 %s -o /dev/null 2>&1 \ +// RUN: | FileCheck %s + +module { + func.func private @tassign_reject_memref() + -> memref<16x16xf16, #pto.address_space> { + %addr = arith.constant 0 : i64 + %tile = memref.alloc() : memref<16x16xf16, #pto.address_space> + // CHECK: error: 'pto.tassign' op expects tile operand and result to be !pto.tile_buf + %rebound = pto.tassign %tile, %addr + : memref<16x16xf16, #pto.address_space> + -> memref<16x16xf16, #pto.address_space> + return %rebound : memref<16x16xf16, #pto.address_space> + } +} diff --git a/test/lit/pto/tassign_tile_native_arg.pto b/test/lit/pto/tassign_tile_native_arg.pto index 442c522729..ce66e72d0b 100644 --- a/test/lit/pto/tassign_tile_native_arg.pto +++ b/test/lit/pto/tassign_tile_native_arg.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a5 %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/tile_buf_addr_tile_native_arg.pto b/test/lit/pto/tile_buf_addr_tile_native_arg.pto index 9ee7e6e987..052f30a282 100644 --- a/test/lit/pto/tile_buf_addr_tile_native_arg.pto +++ b/test/lit/pto/tile_buf_addr_tile_native_arg.pto @@ -6,7 +6,7 @@ // INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. // See LICENSE in the root of the software repository for the full text of the License. -// RUN: ptoas --pto-arch=a5 --pto-level=level3 %s --mlir-print-ir-after=pto-view-to-memref -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a5 --pto-level=level3 %s --mlir-print-ir-after=pto-resolve-reserved-buffers -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a5 --pto-level=level3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC // RUN: ptoas --pto-arch=a5 --pto-level=level3 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO @@ -22,7 +22,7 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { func.func private @consume(!pto.ptr) } -// NATIVE: IR Dump After PTOViewToMemref +// NATIVE: IR Dump After PTOResolveReservedBuffers // NATIVE-LABEL: func.func @tile_buf_addr_tile_native_arg( // NATIVE-SAME: %[[TILE:.*]]: !pto.tile_buf // NATIVE: %[[PTR:.*]] = pto.tile_buf_addr %[[TILE]] : !pto.tile_buf -> !pto.ptr diff --git a/test/lit/pto/tile_compact_mode_emitc.pto b/test/lit/pto/tile_compact_mode_emitc.pto index 7c827be7e4..24bd77b17a 100644 --- a/test/lit/pto/tile_compact_mode_emitc.pto +++ b/test/lit/pto/tile_compact_mode_emitc.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @tile_compact_mode_emitc() { @@ -25,9 +25,9 @@ module { } } -// A3-DAG: pto.alloc_tile : !pto.tile_buf -// A3-DAG: pto.alloc_tile : !pto.tile_buf -// A3-DAG: pto.alloc_tile : !pto.tile_buf +// A3-DAG: pto.alloc_tile{{.*}} : !pto.tile_buf +// A3-DAG: pto.alloc_tile{{.*}} : !pto.tile_buf +// A3-DAG: pto.alloc_tile{{.*}} : !pto.tile_buf // A3: Tile [[DEFAULT:[_A-Za-z][_A-Za-z0-9]*]]; // A3: Tile [[COMPACT:[_A-Za-z][_A-Za-z0-9]*]]; // A3: Tile [[ROWP1:[_A-Za-z][_A-Za-z0-9]*]]; diff --git a/test/lit/pto/tile_scalar_access_tile_native.pto b/test/lit/pto/tile_scalar_access_tile_native.pto index 9e4315946c..b33ee33a2d 100644 --- a/test/lit/pto/tile_scalar_access_tile_native.pto +++ b/test/lit/pto/tile_scalar_access_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/tlrelu_tile_native.pto b/test/lit/pto/tlrelu_tile_native.pto index cfbea9cd5c..5567936e4d 100644 --- a/test/lit/pto/tlrelu_tile_native.pto +++ b/test/lit/pto/tlrelu_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/tmov_fp_tile_native.pto b/test/lit/pto/tmov_fp_tile_native.pto index 1ebdb7912d..8cdcb77cca 100644 --- a/test/lit/pto/tmov_fp_tile_native.pto +++ b/test/lit/pto/tmov_fp_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/tpow_emitc.pto b/test/lit/pto/tpow_emitc.pto index d252177805..08ca27b23c 100644 --- a/test/lit/pto/tpow_emitc.pto +++ b/test/lit/pto/tpow_emitc.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @tpow_tmp_form() { diff --git a/test/lit/pto/tpow_int_no_tmp.pto b/test/lit/pto/tpow_int_no_tmp.pto index 19e2f237c8..9d3cbb7d10 100644 --- a/test/lit/pto/tpow_int_no_tmp.pto +++ b/test/lit/pto/tpow_int_no_tmp.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // Integer path: pto-isa TPowI does not consume tmp, so tmp must be omitted. -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @tpow_int_no_tmp() { diff --git a/test/lit/pto/tpows_emitc.pto b/test/lit/pto/tpows_emitc.pto index d7bc569989..a7d9c94f91 100644 --- a/test/lit/pto/tpows_emitc.pto +++ b/test/lit/pto/tpows_emitc.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @tpows_tmp_form() { diff --git a/test/lit/pto/tpows_int_no_tmp.pto b/test/lit/pto/tpows_int_no_tmp.pto index c2de201be2..0cbc452408 100644 --- a/test/lit/pto/tpows_int_no_tmp.pto +++ b/test/lit/pto/tpows_int_no_tmp.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @tpows_int_no_tmp() { diff --git a/test/lit/pto/tprefetch_tile_native.pto b/test/lit/pto/tprefetch_tile_native.pto index 16aea07fdc..0ea39b2449 100644 --- a/test/lit/pto/tprefetch_tile_native.pto +++ b/test/lit/pto/tprefetch_tile_native.pto @@ -6,7 +6,7 @@ // 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=a3 --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-level=level3 --pto-arch=a3 --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 %s | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/tquant_dynamic_tmp_valid_a3.pto b/test/lit/pto/tquant_no_implicit_tmp_a3.pto similarity index 87% rename from test/lit/pto/tquant_dynamic_tmp_valid_a3.pto rename to test/lit/pto/tquant_no_implicit_tmp_a3.pto index c2c83ca731..be39a34079 100644 --- a/test/lit/pto/tquant_dynamic_tmp_valid_a3.pto +++ b/test/lit/pto/tquant_no_implicit_tmp_a3.pto @@ -9,20 +9,19 @@ // RUN: ptoas --pto-arch=a3 %s -emit-pto-ir 2>&1 | FileCheck %s module { - func.func @tquant_dynamic_tmp_valid_a3(%valid_row: index, %valid_col: index) + func.func @tquant_no_implicit_tmp_a3(%valid_row: index, %valid_col: index) attributes {pto.kernel_kind = #pto.kernel_kind} { %src = pto.alloc_tile valid_row = %valid_row valid_col = %valid_col : !pto.tile_buf %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_dynamic_tmp_valid_a3 + // 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) diff --git a/test/lit/pto/trans_tile_native.pto b/test/lit/pto/trans_tile_native.pto index 174eb9bc47..c77201f6c0 100644 --- a/test/lit/pto/trans_tile_native.pto +++ b/test/lit/pto/trans_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/transcendental_unary_tile_native.pto b/test/lit/pto/transcendental_unary_tile_native.pto index 233b910669..b513717bc0 100644 --- a/test/lit/pto/transcendental_unary_tile_native.pto +++ b/test/lit/pto/transcendental_unary_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/trem_emitc.pto b/test/lit/pto/trem_emitc.pto index 36736b0102..5477c3d5f8 100644 --- a/test/lit/pto/trem_emitc.pto +++ b/test/lit/pto/trem_emitc.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @trem_tmp_form() { diff --git a/test/lit/pto/trems_emitc.pto b/test/lit/pto/trems_emitc.pto index a284156020..778f660641 100644 --- a/test/lit/pto/trems_emitc.pto +++ b/test/lit/pto/trems_emitc.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @trems_tmp_form() { diff --git a/test/lit/pto/treshape_explicit_dynamic_valid_shape_preserved.pto b/test/lit/pto/treshape_explicit_dynamic_valid_shape_preserved.pto index 841545ca45..acc803abc1 100644 --- a/test/lit/pto/treshape_explicit_dynamic_valid_shape_preserved.pto +++ b/test/lit/pto/treshape_explicit_dynamic_valid_shape_preserved.pto @@ -1,5 +1,5 @@ // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o - 2>&1 | FileCheck %s +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o - 2>&1 | FileCheck %s module { func.func @treshape_explicit_dynamic_valid_shape_preserved() { diff --git a/test/lit/pto/treshape_tile_native_arg.pto b/test/lit/pto/treshape_tile_native_arg.pto index 0d56a70714..b7c772957d 100644 --- a/test/lit/pto/treshape_tile_native_arg.pto +++ b/test/lit/pto/treshape_tile_native_arg.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o - 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o - 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE module attributes {pto.target_arch = "a5"} { diff --git a/test/lit/pto/trsqrt_emitc.pto b/test/lit/pto/trsqrt_emitc.pto index 7fc664730b..9dec9ee11d 100644 --- a/test/lit/pto/trsqrt_emitc.pto +++ b/test/lit/pto/trsqrt_emitc.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @trsqrt_forms() { diff --git a/test/lit/pto/tsort32_emitc.pto b/test/lit/pto/tsort32_emitc.pto index ebf8b96648..36a6c89e9a 100644 --- a/test/lit/pto/tsort32_emitc.pto +++ b/test/lit/pto/tsort32_emitc.pto @@ -6,7 +6,7 @@ // 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 --mlir-print-ir-after=pto-view-to-memref %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @tsort32_forms() { diff --git a/test/lit/pto/tsync_tile_native.pto b/test/lit/pto/tsync_tile_native.pto index 4233de5f25..f57318e27e 100644 --- a/test/lit/pto/tsync_tile_native.pto +++ b/test/lit/pto/tsync_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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=a5 %s | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/unary_compute_tile_native.pto b/test/lit/pto/unary_compute_tile_native.pto index 2b2269384c..4e53e0b891 100644 --- a/test/lit/pto/unary_compute_tile_native.pto +++ b/test/lit/pto/unary_compute_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o /dev/null 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-level=level3 --pto-arch=a3 --enable-insert-sync %s \ // RUN: | FileCheck %s --check-prefix=EMITC diff --git a/test/lit/pto/validshape_tile_native.pto b/test/lit/pto/validshape_tile_native.pto index 487514c1a0..ca9fca8c00 100644 --- a/test/lit/pto/validshape_tile_native.pto +++ b/test/lit/pto/validshape_tile_native.pto @@ -7,7 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-level=level3 --pto-arch=a5 \ -// RUN: --mlir-print-ir-after=pto-view-to-memref %s -o - 2>&1 \ +// RUN: --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o - 2>&1 \ // RUN: | FileCheck %s --check-prefix=NATIVE module attributes {pto.target_arch = "a5"} { diff --git a/test/lit/pto/vector_specialty_tile_native.pto b/test/lit/pto/vector_specialty_tile_native.pto index 479a07bffd..4685d719a5 100644 --- a/test/lit/pto/vector_specialty_tile_native.pto +++ b/test/lit/pto/vector_specialty_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/xor_tile_native.pto b/test/lit/pto/xor_tile_native.pto index d22269b0e2..221ea17e3a 100644 --- a/test/lit/pto/xor_tile_native.pto +++ b/test/lit/pto/xor_tile_native.pto @@ -6,7 +6,7 @@ // 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-view-to-memref %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// 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 module { diff --git a/test/lit/tile_fusion/final_emitc_last_use_level2.pto b/test/lit/tile_fusion/final_emitc_last_use_level2.pto index 029d5b5e23..57722044ee 100644 --- a/test/lit/tile_fusion/final_emitc_last_use_level2.pto +++ b/test/lit/tile_fusion/final_emitc_last_use_level2.pto @@ -11,7 +11,7 @@ // - the final generated C++ still emits output-first last_use slots as // [[pto::last_use(0, 1, 1)]] CALLEE(...) // -// RUN: { ptoas --pto-arch=a5 --pto-level=level2 --enable-op-fusion --emit-pto-ir %s --mlir-print-ir-after=pto-view-to-memref -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=IR +// RUN: { ptoas --pto-arch=a5 --pto-level=level2 --enable-op-fusion --emit-pto-ir %s --mlir-print-ir-after=pto-resolve-reserved-buffers -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=IR // RUN: ptoas --pto-arch=a5 --pto-level=level2 --enable-op-fusion %s -o - | FileCheck %s --check-prefix=CPP module { diff --git a/test/lit/tile_fusion/op_fusion_adapter_placement_level2_tadd.pto b/test/lit/tile_fusion/op_fusion_adapter_placement_level2_tadd.pto index 1d493ea705..139fef5b97 100644 --- a/test/lit/tile_fusion/op_fusion_adapter_placement_level2_tadd.pto +++ b/test/lit/tile_fusion/op_fusion_adapter_placement_level2_tadd.pto @@ -10,7 +10,7 @@ // shared memref lowering, optional sync insertion, and the non-fused baseline. // // RUN: ptoas --pto-arch=a5 --pto-level=level2 %s --emit-pto-ir -o - | FileCheck %s --check-prefix=NOFUSE -// RUN: { ptoas --pto-arch=a5 --pto-level=level2 --enable-op-fusion --emit-pto-ir %s --mlir-print-ir-after=pto-view-to-memref -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=SEAM +// RUN: { ptoas --pto-arch=a5 --pto-level=level2 --enable-op-fusion --emit-pto-ir %s --mlir-print-ir-after=pto-resolve-reserved-buffers -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=SEAM // RUN: { ptoas --pto-arch=a5 --pto-level=level2 --enable-op-fusion --enable-insert-sync --emit-pto-ir %s --mlir-print-ir-after=pto-insert-sync -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=SYNC module { diff --git a/test/lit/tile_fusion/op_fusion_adapter_placement_level3_tadd.pto b/test/lit/tile_fusion/op_fusion_adapter_placement_level3_tadd.pto index 9c948530f4..00646515db 100644 --- a/test/lit/tile_fusion/op_fusion_adapter_placement_level3_tadd.pto +++ b/test/lit/tile_fusion/op_fusion_adapter_placement_level3_tadd.pto @@ -11,7 +11,7 @@ // and stays non-fused when op fusion is disabled. // // RUN: ptoas --pto-arch=a5 --pto-level=level3 %s --emit-pto-ir -o - | FileCheck %s --check-prefix=NOFUSE -// RUN: { ptoas --pto-arch=a5 --pto-level=level3 --enable-op-fusion --emit-pto-ir %s --mlir-print-ir-after=pto-view-to-memref -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=SEAM +// RUN: { ptoas --pto-arch=a5 --pto-level=level3 --enable-op-fusion --emit-pto-ir %s --mlir-print-ir-after=pto-resolve-reserved-buffers -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=SEAM module { func.func @fusion_adapter_placement_level3_tadd(%dst_ptr: !pto.ptr) { @@ -60,7 +60,7 @@ module { // NOFUSE: pto.tstore ins( // NOFUSE: return -// SEAM: // -----// IR Dump After PTOViewToMemref (pto-view-to-memref) //----- // +// SEAM: // -----// IR Dump After PTOResolveReservedBuffers (pto-resolve-reserved-buffers) //----- // // SEAM-LABEL: func.func @fusion_adapter_placement_level3_tadd( // SEAM-NOT: pto.fusion_region // SEAM: pto.alloc_tile addr = %c0_i64 diff --git a/test/lit/tile_fusion/op_fusion_region_pipeline_level2.pto b/test/lit/tile_fusion/op_fusion_region_pipeline_level2.pto index dafa62885a..ab108a6100 100644 --- a/test/lit/tile_fusion/op_fusion_region_pipeline_level2.pto +++ b/test/lit/tile_fusion/op_fusion_region_pipeline_level2.pto @@ -10,7 +10,7 @@ // preservation through the level2 pipeline, with and without insert-sync. // // RUN: { ptoas --pto-backend=vpto --pto-arch=a5 --pto-level=level2 --enable-op-fusion --emit-vpto %s --mlir-print-ir-after=pto-fusion-region-gen -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=FORM -// RUN: { ptoas --pto-backend=vpto --pto-arch=a5 --pto-level=level2 --enable-op-fusion --emit-vpto %s --mlir-print-ir-after=pto-view-to-memref -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=LEVEL2-PLAIN +// RUN: { ptoas --pto-backend=vpto --pto-arch=a5 --pto-level=level2 --enable-op-fusion --emit-vpto %s --mlir-print-ir-after=pto-resolve-reserved-buffers -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=LEVEL2-PLAIN // RUN: { ptoas --pto-backend=vpto --pto-arch=a5 --pto-level=level2 --enable-op-fusion --enable-insert-sync --emit-vpto %s --mlir-print-ir-after=pto-insert-sync -o /dev/null 2>&1 || true; } | FileCheck %s --check-prefix=LEVEL2-SYNC module { diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index b6343a34ca..12c5d6ab97 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -123,7 +123,6 @@ void mlir::pto::registerPTOASDialects(DialectRegistry ®istry) { void mlir::pto::registerPTOASPassesAndCLOptions() { mlir::registerAllPasses(); mlir::pto::registerPTOPasses(); - mlir::pto::registerPTOViewToMemrefPass(); mlir::pto::registerPTOInlineLibCall(); mlir::pto::registerFoldTileBufIntrinsics(); mlir::pto::registerExpandTileOp(); @@ -3110,7 +3109,6 @@ int mlir::pto::compilePTOASModule( pm.addNestedPass(pto::createPTOFusionRegionGenPass()); } - pm.addPass(pto::createPTOViewToMemrefPass()); pm.addNestedPass( pto::createPTORematerializeFixpipeVectorQuantPass()); From b1920d875a9b1cee1e3ce7b1b25295fa58965cc7 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 21 Jul 2026 11:25:09 +0800 Subject: [PATCH 87/91] Remove remaining local memref tile compatibility paths --- docs/PTO_IR_manual.md | 25 +- include/PTO/Transforms/Passes.h | 1 - include/PTO/Transforms/Passes.td | 18 - lib/PTO/IR/PTO.cpp | 113 +-- lib/PTO/Transforms/AllocToPointerCast.cpp | 115 --- lib/PTO/Transforms/AllocToPointerCast.h | 67 -- lib/PTO/Transforms/CMakeLists.txt | 2 - lib/PTO/Transforms/OptMemPlanForPipeline.cpp | 28 +- lib/PTO/Transforms/OptMemPlanForPipeline.h | 11 +- .../Transforms/PTOMaterializeTileHandles.cpp | 890 ------------------ lib/PTO/Transforms/PTOPlanMemory.cpp | 108 +-- lib/PTO/Transforms/PTOPlanMemory.h | 20 +- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 34 - .../pto/alloc_tile_physical_overlap_sync.pto | 12 +- .../bias_tile_subview_emitc_pointer_type.pto | 25 - test/lit/pto/declare_tile_tile_native.pto | 2 +- test/lit/pto/easy_param_completion_emitc.pto | 12 +- ..._sync_solver_a3_dependent_mmad_barrier.pto | 10 +- .../graph_sync_solver_section_transparent.pto | 43 +- .../inject_barrier_all_sync_tpush_tpop.pto | 8 +- .../pto/issue157_64bit_view_offset_emitc.pto | 25 +- ...ue667_subview_disjoint_no_mte2_barrier.pto | 16 +- ...issue667_tload_overlap_no_mte2_barrier.pto | 8 +- test/lit/pto/issue708_zero_valid_subview.pto | 8 +- .../issue708_zero_valid_subview_inferred.pto | 12 +- .../pto/issue870_identity_tmov_pruning.pto | 30 +- .../low_precision_globaltensor_emitc_a5.pto | 12 +- ...alize_tile_handles_control_flow_result.pto | 4 +- .../pto/materialize_tile_handles_emitc.pto | 4 +- ...lize_tile_handles_subkernel_helper_abi.pto | 4 +- ...ub_a5_elem_1x1_memref_macro_sync_model.pto | 14 +- .../pto/multi_tile_affine_disjoint_slots.pto | 8 +- .../lit/pto/multi_tile_buf_n3_planmem_e2e.pto | 8 +- ...lti_tile_const_preload_dyn_loop_select.pto | 33 +- .../multi_tile_const_slot_disjoint_sync.pto | 8 +- .../multi_tile_const_slot_gss_compiles.pto | 8 +- .../multi_tile_dyn_slot_conservative_sync.pto | 4 +- .../multi_tile_get_const_slot_lowering.pto | 8 +- .../pto/multi_tile_get_dyn_slot_lowering.pto | 6 +- .../pto/multi_tile_level3_explicit_addr.pto | 8 +- .../pto/multi_tile_level3_missing_addr.pto | 4 +- test/lit/pto/multi_tile_n4_planmem_e2e.pto | 10 +- test/lit/pto/multi_tile_no_loop_unroll.pto | 8 +- .../pto/multi_tile_prefetch_dyn_event_id.pto | 8 +- .../pto/multi_tile_prefetch_gss_event_id.pto | 8 +- .../pto/multi_tile_preload_loop_set_wait.pto | 12 +- .../pto/multi_tile_two_buf_slot_binding.pto | 8 +- ...lti_tile_unknown_slot_gss_dyn_event_id.pto | 8 +- ...lan_memory_five_gates_lifetime_overlap.pto | 62 +- .../pto/plan_memory_five_gates_phi_family.pto | 52 +- .../pto/plan_memory_for_iter_args_yield.pto | 35 - .../plan_memory_fragmentation_hole_fit.pto | 152 --- .../plan_memory_fragmentation_two_holes.pto | 163 ---- test/lit/pto/plan_memory_if_in_loop.pto | 38 - test/lit/pto/plan_memory_if_yield.pto | 34 - .../pto/plan_memory_inplace_forbid_alias.pto | 49 +- test/lit/pto/plan_memory_loop_in_if.pto | 37 - .../plan_memory_loop_no_reuse_outer_live.pto | 41 - test/lit/pto/plan_memory_nested_loops.pto | 46 - test/lit/pto/plan_memory_no_reuse_overlap.pto | 31 - .../pto/plan_memory_peak_8_overlapping.pto | 59 -- .../pto/plan_memory_peak_exact_capacity.pto | 141 --- ...erve_buffer_manual_reject_nested_alloc.pto | 26 - .../pto/plan_memory_reserve_buffer_prefix.pto | 27 - test/lit/pto/plan_memory_reuse_sequential.pto | 203 ---- .../pto/plan_memory_reused_tstore_sync.pto | 59 +- .../plan_memory_reused_tstore_sync_level2.pto | 58 +- ...an_memory_spec_level0_no_reuse_overlap.pto | 31 +- .../lit/pto/plan_memory_spec_level0_reuse.pto | 41 +- ...ptodsl_subkernel_helper_tile_abi_emitc.pto | 4 +- ...el_pipe_contract_reject_transitive_mte.pto | 12 +- ..._column_slice_preserve_parent_shape_a5.pto | 4 +- ...w_col_major_compact_keeps_normal_shape.pto | 8 +- ...w_col_major_noncompact_preserve_stride.pto | 8 +- ...w_col_major_row_plus_one_stride_offset.pto | 8 +- .../subview_compact_keeps_normal_shape.pto | 8 +- ...dynamic_offset_static_valid_regression.pto | 4 +- test/lit/pto/subview_explicit_result_type.pto | 8 +- .../subview_row_plus_one_stride_offset.pto | 8 +- .../subview_tile_native_preserve_stride.pto | 20 +- test/lit/pto/subview_validshape_guard.pto | 16 +- ...shape_partial_parent_requires_explicit.pto | 4 +- test/lit/pto/tci_i16_emitc.pto | 6 +- test/lit/pto/tci_ui16_emitc.pto | 6 +- test/lit/pto/tci_ui32_emitc.pto | 6 +- test/lit/pto/tdequant.pto | 34 +- .../textract_a5_scaling_pipe_selection.pto | 4 +- ...textract_a5_scaling_pipe_selection_gss.pto | 4 +- .../pto/tfillpad_a5_mat_pipe_selection.pto | 4 +- ...ile_assemble_level3_subview_tmov_emitc.pto | 12 +- .../tinsert_a2a3_vec_vec_pipe_selection.pto | 20 +- ...insert_a2a3_vec_vec_pipe_selection_gss.pto | 20 +- test/lit/pto/tinsert_a3_pipe_selection.pto | 77 +- .../lit/pto/tinsert_a3_pipe_selection_gss.pto | 52 +- test/lit/pto/tinsert_a5_pipe_selection.pto | 77 +- .../lit/pto/tinsert_a5_pipe_selection_gss.pto | 52 +- .../pto/tinsert_a5_vec_mat_mode_lowering.pto | 8 +- .../lit/pto/tinsert_a5_vec_pipe_selection.pto | 20 +- .../pto/tinsert_a5_vec_pipe_selection_gss.pto | 20 +- .../pto/tinsert_level3_tile_native_sync.pto | 12 +- ...tload_tprefetch_low_precision_a5_valid.pto | 6 +- test/lit/pto/tmov_acc_mat_pipe_selection.pto | 49 +- .../pto/tmov_acc_mat_pipe_selection_gss.pto | 33 +- .../lit/pto/tmov_acc_to_vec_mode_a5_emitc.pto | 45 +- .../pto/tmov_acc_to_vec_mode_a5_emitc_gss.pto | 32 +- test/lit/pto/tmov_fbuf_pipe_selection.pto | 4 +- .../pto/tmov_mx_scale_pipe_selection_a5.pto | 8 +- ...efetch_low_precision_memref_invalid_a3.pto | 8 +- test/lit/pto/tquant_mxfp8_a5.pto | 64 +- .../pto/tscatter_indexed_macro_sync_model.pto | 29 +- .../lit/pto/tstore_low_precision_a5_valid.pto | 6 +- test/lit/vpto/expand_tile_op_tilelang.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tabs.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tand.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tcmp.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tdiv.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_texp.pto | 2 +- .../vpto/expand_tile_op_tilelang_texpand.pto | 2 +- .../vpto/expand_tile_op_tilelang_tfillpad.pto | 2 +- ...xpand_tile_op_tilelang_tfillpad_expand.pto | 2 +- ...pand_tile_op_tilelang_tfillpad_inplace.pto | 2 +- .../vpto/expand_tile_op_tilelang_tfmod.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tlog.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tmax.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tmin.pto | 2 +- .../vpto/expand_tile_op_tilelang_tmrgsort.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tmul.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tneg.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tnot.pto | 2 +- test/lit/vpto/expand_tile_op_tilelang_tor.pto | 2 +- .../vpto/expand_tile_op_tilelang_tpartadd.pto | 2 +- .../vpto/expand_tile_op_tilelang_tpartmax.pto | 2 +- .../vpto/expand_tile_op_tilelang_tpartmin.pto | 2 +- .../vpto/expand_tile_op_tilelang_tpartmul.pto | 2 +- .../vpto/expand_tile_op_tilelang_trecip.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_trem.pto | 2 +- .../expand_tile_op_tilelang_trowargmax.pto | 2 +- .../expand_tile_op_tilelang_trowargmin.pto | 2 +- .../vpto/expand_tile_op_tilelang_trowmax.pto | 2 +- .../vpto/expand_tile_op_tilelang_trowmin.pto | 2 +- .../vpto/expand_tile_op_tilelang_trowprod.pto | 2 +- .../vpto/expand_tile_op_tilelang_trowsum.pto | 2 +- .../vpto/expand_tile_op_tilelang_trsqrt.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tshl.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tshr.pto | 2 +- .../vpto/expand_tile_op_tilelang_tsort32.pto | 2 +- .../vpto/expand_tile_op_tilelang_tsqrt.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_tsub.pto | 2 +- .../lit/vpto/expand_tile_op_tilelang_txor.pto | 2 +- .../vpto/expand_tile_op_tlrelu_tilelang.pto | 2 +- .../vpto/expand_tile_op_tprelu_tilelang.pto | 2 +- .../vpto/expand_tile_op_trelu_tilelang.pto | 2 +- .../lit/vpto/expand_tile_op_tsel_tilelang.pto | 2 +- .../vpto/expand_tile_op_tsels_tilelang.pto | 2 +- ...le_sections_reject_residual_nested_mix.pto | 20 +- test/lit/vpto/tcolargmax.pto | 2 +- test/lit/vpto/tcolargmin.pto | 2 +- test/lit/vpto/tcolexpand_tilelang.pto | 2 +- test/lit/vpto/tcolexpandadd_tilelang.pto | 2 +- test/lit/vpto/tcolexpanddiv_tilelang.pto | 2 +- test/lit/vpto/tcolexpandexpdif_tilelang.pto | 2 +- test/lit/vpto/tcolexpandmax_tilelang.pto | 2 +- test/lit/vpto/tcolexpandmin_tilelang.pto | 2 +- test/lit/vpto/tcolexpandmul_tilelang.pto | 2 +- test/lit/vpto/tcolexpandsub_tilelang.pto | 2 +- test/lit/vpto/tcolmax.pto | 2 +- test/lit/vpto/tcolmin.pto | 2 +- test/lit/vpto/tcolprod.pto | 2 +- test/lit/vpto/tcolsum.pto | 2 +- test/lit/vpto/trowexpand_tile_op_expand.pto | 2 +- .../lit/vpto/trowexpandadd_tile_op_expand.pto | 2 +- .../lit/vpto/trowexpanddiv_tile_op_expand.pto | 2 +- .../vpto/trowexpandexpdif_tile_op_expand.pto | 2 +- .../lit/vpto/trowexpandmax_tile_op_expand.pto | 2 +- .../lit/vpto/trowexpandmin_tile_op_expand.pto | 2 +- .../lit/vpto/trowexpandmul_tile_op_expand.pto | 2 +- .../lit/vpto/trowexpandsub_tile_op_expand.pto | 2 +- tools/ptoas/ptoas.cpp | 7 +- 178 files changed, 876 insertions(+), 3334 deletions(-) delete mode 100644 lib/PTO/Transforms/AllocToPointerCast.cpp delete mode 100644 lib/PTO/Transforms/AllocToPointerCast.h delete mode 100644 lib/PTO/Transforms/PTOMaterializeTileHandles.cpp delete mode 100644 test/lit/pto/bias_tile_subview_emitc_pointer_type.pto delete mode 100644 test/lit/pto/plan_memory_for_iter_args_yield.pto delete mode 100644 test/lit/pto/plan_memory_fragmentation_hole_fit.pto delete mode 100644 test/lit/pto/plan_memory_fragmentation_two_holes.pto delete mode 100644 test/lit/pto/plan_memory_if_in_loop.pto delete mode 100644 test/lit/pto/plan_memory_if_yield.pto delete mode 100644 test/lit/pto/plan_memory_loop_in_if.pto delete mode 100644 test/lit/pto/plan_memory_loop_no_reuse_outer_live.pto delete mode 100644 test/lit/pto/plan_memory_nested_loops.pto delete mode 100644 test/lit/pto/plan_memory_no_reuse_overlap.pto delete mode 100644 test/lit/pto/plan_memory_peak_8_overlapping.pto delete mode 100644 test/lit/pto/plan_memory_peak_exact_capacity.pto delete mode 100644 test/lit/pto/plan_memory_reserve_buffer_manual_reject_nested_alloc.pto delete mode 100644 test/lit/pto/plan_memory_reserve_buffer_prefix.pto delete mode 100644 test/lit/pto/plan_memory_reuse_sequential.pto diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index 511fcf3f5f..b32ea5150f 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -216,11 +216,11 @@ selection): regular `tile_buf` that flows through every existing DMA / compute / view op unchanged -The N-way physical fan-out lives on the `pto.multi_buffer = N : i32` -attribute that PTOViewToMemref writes onto the lowered `memref.alloc`; -downstream passes (PlanMemory / InsertSync / GraphSyncSolver) consume that -attribute. The per-use slot index threaded through `pto.multi_tile_get` is -forwarded to the memref layer via the internal `pto.slot_marker` view op. +The slot count is part of `!pto.multi_tile_buf`. In level1/level2, +`PTOPlanMemory` plans all slots and records their physical addresses in the +internal `pto.multi_buffer_addrs` attribute. `PTOResolveBufferSelect`, +InsertSync, and GraphSyncSolver consume the tile-native allocation and slot +selection directly; no memref allocation is introduced. See `docs/designs/ptoas-multi-buffer-explicit-design.md` for the full design. @@ -244,9 +244,8 @@ compiler, not by PTO memory planning. - `T` is a scalar integer or float **Disjoint from tile-buf world.** Values of `!pto.local_array<...>` never -participate in `pto.pointer_cast`, `pto-plan-memory`, or -`AllocToPointerCast` rewrites — these passes match on `memref` / tile-buf -types and simply do not see this type. +participate in `pto-plan-memory`; the planner only collects tile allocation +roots and does not see this type. **Syntax:** ```mlir @@ -748,7 +747,7 @@ result = alloc_tile(base_addr, valid_row, valid_col) // operands are optional ##### `pto.alloc_multi_tile` - Allocate N-Slot Multi-Buffer Tile -**Summary:** Declares the lifetime of an N-slot multi-buffer tile. Each slot has the same `tile_buf` shape; only the underlying physical address differs. In the default pipeline (`--pto-level=level1|level2`) the N physical slots are reserved by `PTOPlanMemory` from the `pto.multi_buffer = N` attribute written onto the lowered `memref.alloc`. Under `--pto-level=level3` the caller owns local memory and PlanMemory does not run, so an explicit base `addr` operand is required (mirroring `pto.alloc_tile`); the lowering then fans the base out into the multi-address `pto.pointer_cast` PlanMemory would otherwise produce. +**Summary:** Declares the lifetime of an N-slot multi-buffer tile. Each slot has the same `tile_buf` shape; only the underlying physical address differs. In the default pipeline (`--pto-level=level1|level2`), `PTOPlanMemory` plans the tile-native allocation directly and records all slot addresses in `pto.multi_buffer_addrs`. Under `--pto-level=level3`, the caller owns local memory and must provide an explicit base `addr`; `PTOResolveBufferSelect` derives each selected slot address from that base. **Semantics:** @@ -1120,15 +1119,15 @@ op is modeled as writing the prefetched data into `dst`. | Name | Type | Description | |------|------|-------------| -| `src` | `pto.partition_tensor_view` or lowered GM memref | Source global view | -| `dst` | `pto.tile_buf` or lowered local memref | Destination local tile | +| `src` | `pto.partition_tensor_view` | Source global view | +| `dst` | `pto.tile_buf` | Destination local tile | **Results:** None. Writes into `dst` via DPS pattern. **Constraints & Verification:** -- `src` must be a partition view before lowering, or the corresponding lowered ranked memref form after `PTOViewToMemref`. -- `dst` must be a tile buffer before lowering, or the corresponding lowered ranked memref form after `PTOViewToMemref`. +- `src` must be a partition tensor view. +- `dst` must be a tile buffer. - `dst` must use `loc=vec` or `loc=mat`. - Static source extents must be positive when known; static destination valid extents must be non-negative when known. - `src` and `dst` element types must have the same element size in bytes. diff --git a/include/PTO/Transforms/Passes.h b/include/PTO/Transforms/Passes.h index 4dada8a275..70a79b98d4 100644 --- a/include/PTO/Transforms/Passes.h +++ b/include/PTO/Transforms/Passes.h @@ -74,7 +74,6 @@ createPlanMemoryModernPass(const PlanMemoryOptions &options); std::unique_ptr createPTORemoveRedundantBarrierPass(); std::unique_ptr createPTOValidateIntToPtrUsesPass(); std::unique_ptr createPTORematerializeFixpipeVectorQuantPass(); -std::unique_ptr createPTOMaterializeTileHandlesPass(); 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 49bf964833..3b6c3ad571 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -705,24 +705,6 @@ def PTOValidateIntToPtrUses : Pass<"pto-validate-inttoptr-uses", "func::FuncOp"> ]; } -def PTOMaterializeTileHandles : Pass<"pto-materialize-tile-handles", "ModuleOp"> { - let summary = "Restore tile_buf handles for remaining legacy memref values"; - let description = [{ - Re-wraps any remaining local memref values used by tile operations as - addressed `pto.alloc_tile` handles before backend lowering. Tile-native - operations already bypass this compatibility path. - }]; - - let constructor = "mlir::pto::createPTOMaterializeTileHandlesPass()"; - - let dependentDialects = [ - "mlir::pto::PTODialect", - "mlir::memref::MemRefDialect", - "mlir::arith::ArithDialect", - "mlir::func::FuncDialect" - ]; -} - def PTOResolveBufferSelect : Pass<"pto-resolve-buffer-select", "ModuleOp"> { let summary = "Resolve multi-buffer slot selections to addressed tile handles"; let description = [{ diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 7f8cd60f69..271c2bb5ec 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -489,16 +489,6 @@ static std::optional getVerifierArchName(Operation *op) { return std::nullopt; } -static bool shouldBypassDecodedMemrefVerifier(Operation *op) { - if (!op) - return false; - for (Value operand : op->getOperands()) { - if (isa(operand.getType())) - return true; - } - return false; -} - static SmallVector canonicalizeTileBufValidShape(ArrayRef validShape) { SmallVector canonical; canonical.reserve(validShape.size()); @@ -510,8 +500,6 @@ static SmallVector canonicalizeTileBufValidShape(ArrayRef v template static LogicalResult dispatchVerifierByArch(Operation *op, FnA2A3 &&verifyA2A3, FnA5 &&verifyA5) { - if (shouldBypassDecodedMemrefVerifier(op)) - return success(); switch (getVerifierTargetArch(op)) { case VerifierTargetArch::A2A3: return verifyA2A3(); @@ -3738,8 +3726,6 @@ LogicalResult TStoreOp::verify() { } LogicalResult pto::TAbsOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); if (failed(verifyVecTileCommon(*this, srcTy, "src")) || @@ -4413,14 +4399,8 @@ static LogicalResult verifyTileBufCommon(Operation *op, Type ty, StringRef name, if (!allowLowPrecision && isPTOLowPrecisionType(elemTy)) return op->emitOpError() << name << ": dtype " << elemTy << " is not supported by this op yet"; - } else if (auto mr = dyn_cast(ty)) { - if (mr.getRank() != 2) - return op->emitOpError() << "expects " << name << " to be a rank-2 memref"; - if (!allowLowPrecision && isPTOLowPrecisionType(mr.getElementType())) - return op->emitOpError() << name << ": dtype " << mr.getElementType() - << " is not supported by this op yet"; } else { - return op->emitOpError() << "expects " << name << " to be a !pto.tile_buf or rank-2 memref"; + return op->emitOpError() << "expects " << name << " to be a !pto.tile_buf"; } auto validShape = getValidShapeVec(ty); @@ -5328,8 +5308,6 @@ LogicalResult pto::TAddOp::verify() { } LogicalResult pto::TAddCOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type t0 = getSrc0().getType(); Type t1 = getSrc1().getType(); Type t2 = getSrc2().getType(); @@ -5414,8 +5392,6 @@ LogicalResult pto::TAxpyOp::verify() { } LogicalResult pto::TAddSCOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type ts0 = getSrc0().getType(); Type ts1 = getSrc1().getType(); Type td = getDst().getType(); @@ -5774,8 +5750,6 @@ void mlir::pto::TCIOp::print(OpAsmPrinter &p) { } LogicalResult pto::TCIOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, dstTy, "dst"))) return failure(); @@ -5806,8 +5780,6 @@ LogicalResult pto::TCIOp::verify() { } LogicalResult pto::TTriOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type dstTy = getDst().getType(); if (failed(verifyVecTileCommon(*this, dstTy, "dst"))) @@ -5968,8 +5940,6 @@ LogicalResult pto::TCmpSOp::verify() { return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } LogicalResult pto::TColExpandOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); if (failed(verifyNDStyleVecTile(*this, srcTy, "src")) || @@ -6315,8 +6285,6 @@ LogicalResult pto::TColProdOp::verify() { } llvm::LogicalResult mlir::pto::TCvtOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, srcTy, "src", /*allowLowPrecision=*/true)) || @@ -6359,8 +6327,6 @@ llvm::LogicalResult mlir::pto::TRandomOp::verify() { return emitOpError("trandom is only supported for A5 targets"); }; auto verifyA5 = [&]() -> LogicalResult { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, dstTy, "dst"))) @@ -7722,8 +7688,6 @@ mlir::LogicalResult mlir::pto::TGatherBOp::verify() { } mlir::LogicalResult mlir::pto::TLogOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); if (failed(verifyVecTileUnaryOp(*this, srcTy, dstTy, "src", "dst", @@ -8880,8 +8844,6 @@ LogicalResult TMatmulMxBiasOp::verify() { } // ---- TSetValOp ---- LogicalResult TSetValOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); // dst can be tile/tensor/tilebuf (PTODpsType). Keep checks minimal. if (auto shaped = dyn_cast(getDst().getType())) { if (shaped.getElementType() != getVal().getType()) @@ -8891,8 +8853,6 @@ LogicalResult TSetValOp::verify() { } // ---- TGetValOp ---- LogicalResult TGetValOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); if (!mlir::isa(srcTy)) return emitOpError("expects src to be tile_buf or memref type"); @@ -9146,8 +9106,6 @@ void mlir::pto::MScatterOp::print(OpAsmPrinter &p) { } LogicalResult MScatterOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type idxTy = getIdx().getType(); @@ -9388,8 +9346,6 @@ void mlir::pto::MGatherOp::print(OpAsmPrinter &p) { } LogicalResult MGatherOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type memTy = getMem().getType(); Type idxTy = getIdx().getType(); @@ -9606,8 +9562,6 @@ ParseResult mlir::pto::TMrgSortOp::parse(OpAsmParser &parser, OperationState &re } mlir::LogicalResult mlir::pto::TMrgSortOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (isFormat1()) { Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); @@ -9703,8 +9657,6 @@ mlir::LogicalResult mlir::pto::TMulSOp::verify() { } mlir::LogicalResult mlir::pto::TShlSOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); if (failed(verifyTileBufCommon(*this, srcTy, "src")) || @@ -10191,8 +10143,6 @@ mlir::LogicalResult mlir::pto::TPartMulOp::verify() { } mlir::LogicalResult mlir::pto::TPReluOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); auto verifyCommon = [&]() -> FailureOr> { Type t0 = getSrc0().getType(); Type t1 = getSrc1().getType(); @@ -10481,9 +10431,6 @@ mlir::LogicalResult mlir::pto::TQuantOp::verify() { if (failed(verifyStructural())) return failure(); - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); - auto verifyInt8Common = [&]() -> LogicalResult { Type srcTy = getSrc().getType(); Type fpTy = getFp().getType(); @@ -10573,8 +10520,6 @@ mlir::LogicalResult mlir::pto::TQuantMxOp::verify() { }; auto verifyA5 = [&]() -> LogicalResult { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); @@ -10715,9 +10660,6 @@ mlir::LogicalResult mlir::pto::TDequantOp::verify() { if (failed(verifyStructural())) return failure(); - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); - auto verifyCommon = [&]() -> LogicalResult { if (failed(verifyTileBufCommon(*this, getSrc().getType(), "src")) || failed(verifyTileBufCommon(*this, getScale().getType(), "scale")) || @@ -10743,8 +10685,6 @@ mlir::LogicalResult mlir::pto::TDequantOp::verify() { } mlir::LogicalResult mlir::pto::TRecipOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type ts = getSrc().getType(); Type td = getDst().getType(); if (failed(verifyVecTileUnaryOp(*this, ts, td, "src", "dst", @@ -10787,8 +10727,6 @@ mlir::LogicalResult mlir::pto::TReluOp::verify() { mlir::LogicalResult mlir::pto::TRemOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type src0Ty = getSrc0().getType(); Type src1Ty = getSrc1().getType(); @@ -10846,8 +10784,6 @@ mlir::LogicalResult mlir::pto::TFModOp::verify() { } mlir::LogicalResult mlir::pto::TRemSOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type ts = getSrc().getType(); Type tt = getTmp().getType(); Type td = getDst().getType(); @@ -10893,8 +10829,6 @@ mlir::LogicalResult mlir::pto::TRemSOp::verify() { } mlir::LogicalResult mlir::pto::TFModSOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); @@ -10934,8 +10868,6 @@ static LogicalResult verifyTPowTmpShape(Operation *op, Type tmpTy, Type dstTy) { } mlir::LogicalResult mlir::pto::TPowOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type baseTy = getBase().getType(); Type expTy = getExp().getType(); @@ -11000,8 +10932,6 @@ mlir::LogicalResult mlir::pto::TPowOp::verify() { } mlir::LogicalResult mlir::pto::TPowSOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); @@ -11192,8 +11122,6 @@ mlir::LogicalResult mlir::pto::GetValidShapeOp::verify() { mlir::LogicalResult mlir::pto::TReshapeOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type ts = getSrc().getType(); Type tr = getResult().getType(); auto srcTb = dyn_cast(ts); @@ -11235,8 +11163,6 @@ mlir::LogicalResult mlir::pto::TReshapeOp::verify() { } mlir::LogicalResult mlir::pto::BitcastOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); auto srcTy = llvm::dyn_cast(getSrc().getType()); auto dstTy = llvm::dyn_cast(getResult().getType()); if (!srcTy || !dstTy) @@ -12114,8 +12040,6 @@ mlir::LogicalResult mlir::pto::TRowProdOp::verify() { mlir::LogicalResult mlir::pto::TRsqrtOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type ts = getSrc().getType(); Type td = getDst().getType(); if (failed(verifyVecTileUnaryOp(*this, ts, td, "src", "dst", @@ -12430,8 +12354,6 @@ mlir::LogicalResult mlir::pto::TShrOp::verify() { mlir::LogicalResult mlir::pto::TSort32Op::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); Type idxTy = getIdx().getType(); @@ -12459,8 +12381,6 @@ mlir::LogicalResult mlir::pto::TSort32Op::verify() { mlir::LogicalResult mlir::pto::TSqrtOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type srcTy = getSrc().getType(); Type dstTy = getDst().getType(); if (failed(verifyVecTileUnaryOp(*this, srcTy, dstTy, "src", "dst", @@ -12597,8 +12517,6 @@ mlir::LogicalResult mlir::pto::TSubOp::verify() { mlir::LogicalResult mlir::pto::TSubCOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type src0Ty = getSrc0().getType(); Type src1Ty = getSrc1().getType(); Type src2Ty = getSrc2().getType(); @@ -12625,8 +12543,6 @@ mlir::LogicalResult mlir::pto::TSubSOp::verify() { mlir::LogicalResult mlir::pto::TSubSCOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); Type src0Ty = getSrc0().getType(); Type src1Ty = getSrc1().getType(); Type dstTy = getDst().getType(); @@ -12862,9 +12778,6 @@ mlir::LogicalResult mlir::pto::TPrintOp::verify() { dyn_cast_or_null(getProperties().printFormat); if (printFormatAttr && !tmp) return emitOpError() << "expects printFormat only when tmp is present"; - - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (auto tb = mlir::dyn_cast(srcType)) { auto elem = tb.getElementType(); if (!(elem.isF16() || elem.isF32() || @@ -13013,8 +12926,6 @@ LogicalResult mlir::pto::TGemvOp::verify() { } LogicalResult mlir::pto::TMatmulAccOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); auto verifyA2A3 = [&]() -> LogicalResult { if (failed(verifyAccTileCommon(*this, getAccIn().getType(), "acc_in")) || failed(verifyMatTileOperands(*this, getLhs().getType(), getRhs().getType(), @@ -13038,8 +12949,6 @@ LogicalResult mlir::pto::TMatmulAccOp::verify() { } LogicalResult mlir::pto::TGemvAccOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyAccTileCommon(*this, getAccIn().getType(), "acc_in")) || failed(verifyGemvTileOperands(*this, getLhs().getType(), getRhs().getType(), getDst().getType()))) @@ -13884,8 +13793,6 @@ mlir::LogicalResult mlir::pto::SimdTileToMemrefOp::verify() { } mlir::LogicalResult mlir::pto::SubViewOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); auto srcTy = llvm::dyn_cast(getSource().getType()); auto dstTy = llvm::dyn_cast(getResult().getType()); if (!srcTy || !dstTy) @@ -16360,8 +16267,6 @@ LogicalResult TGetAsyncOp::verify() { } LogicalResult TPutOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyCommGlobalLike(*this, getDst(), "dst")) || failed(verifyCommGlobalLike(*this, getSrc(), "src")) || failed(verifyCommStagingTileLike(*this, getPing(), "ping")) || @@ -16378,8 +16283,6 @@ LogicalResult TPutOp::verify() { } LogicalResult TGetOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyCommGlobalLike(*this, getDst(), "dst")) || failed(verifyCommGlobalLike(*this, getSrc(), "src")) || failed(verifyCommStagingTileLike(*this, getPing(), "ping")) || @@ -16396,8 +16299,6 @@ LogicalResult TGetOp::verify() { } LogicalResult TNotifyOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyCommSignalLike(*this, getSignal(), "signal"))) return failure(); auto valueTy = dyn_cast(getValue().getType()); @@ -16407,8 +16308,6 @@ LogicalResult TNotifyOp::verify() { } LogicalResult TWaitOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyCommSignalLike(*this, getSignal(), "signal"))) return failure(); auto cmpTy = dyn_cast(getCmpValue().getType()); @@ -16418,8 +16317,6 @@ LogicalResult TWaitOp::verify() { } LogicalResult TTestOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyCommSignalLike(*this, getSignal(), "signal"))) return failure(); auto cmpTy = dyn_cast(getCmpValue().getType()); @@ -16554,8 +16451,6 @@ LogicalResult SyncAllOp::verify() { } LogicalResult TBroadcastOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyCommGlobalLike(*this, getSrc(), "src")) || failed(verifyCommStagingTileLike(*this, getPing(), "ping")) || failed(verifyCommPingPongSameType(*this, getPing(), getPong(), "ping", @@ -16572,8 +16467,6 @@ LogicalResult TBroadcastOp::verify() { } LogicalResult CommTGatherOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyCommGlobalLike(*this, getDst(), "dst")) || failed(verifyCommStagingTileLike(*this, getPing(), "ping")) || failed(verifyCommPingPongSameType(*this, getPing(), getPong(), "ping", @@ -16590,8 +16483,6 @@ LogicalResult CommTGatherOp::verify() { } LogicalResult CommTScatterOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyCommGlobalLike(*this, getSrc(), "src")) || failed(verifyCommStagingTileLike(*this, getPing(), "ping")) || failed(verifyCommPingPongSameType(*this, getPing(), getPong(), "ping", @@ -16608,8 +16499,6 @@ LogicalResult CommTScatterOp::verify() { } LogicalResult TReduceOp::verify() { - if (shouldBypassDecodedMemrefVerifier(getOperation())) - return success(); if (failed(verifyCommGlobalLike(*this, getDst(), "dst")) || failed(verifyCommStagingTileLike(*this, getAcc(), "acc")) || failed(verifyCommStagingTileLike(*this, getRecvPing(), "recv_ping")) || diff --git a/lib/PTO/Transforms/AllocToPointerCast.cpp b/lib/PTO/Transforms/AllocToPointerCast.cpp deleted file mode 100644 index 667481964b..0000000000 --- a/lib/PTO/Transforms/AllocToPointerCast.cpp +++ /dev/null @@ -1,115 +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. - -//===- AllocToPointerCast.cpp - convert memref.AllocOp to pto.pointercastOp.// -//===----------------------------------------------------------------------===// - -#include "AllocToPointerCast.h" -#include "PTO/IR/PTOMultiBuffer.h" -#include "PTO/IR/PTOTypeUtils.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" - -using namespace mlir; -using namespace mlir::pto; - -namespace {} // namespace - -namespace { -constexpr uint64_t kDefaultAllocAlignmentBytes = 4096; -constexpr size_t kDynamicValidShapeRank = 2; - -static uint64_t alignUpToDefault(uint64_t value) { - return ((value + kDefaultAllocAlignmentBytes - 1) / - kDefaultAllocAlignmentBytes) * - kDefaultAllocAlignmentBytes; -} - -static SmallVector getAllocatedOffsets( - Value root, Operation *allocLikeOp, BaseMemRefType memRefType, - const DenseMap> &buffer2Offsets, - uint64_t &fallbackNextOffset) { - auto iter = buffer2Offsets.find(root); - SmallVector offsets; - if (iter != buffer2Offsets.end()) - offsets = iter->second; - - if (!offsets.empty()) - return offsets; - - // Estimate buffer size (best-effort). Most PTO tile buffers are 32x32 and - // naturally align to 4096 bytes. - uint64_t bytes = kDefaultAllocAlignmentBytes; - if (auto memrefTy = dyn_cast(memRefType)) { - uint64_t elemBytes = getPTOStorageElemByteSize(memrefTy.getElementType()); - if (elemBytes != 0) { - uint64_t numel = 1; - bool allStatic = true; - for (int64_t d : memrefTy.getShape()) { - if (d == ShapedType::kDynamic) { - allStatic = false; - break; - } - numel *= static_cast(d); - } - if (allStatic && numel != 0) - bytes = numel * elemBytes; - } - } - - uint64_t slotCount = 1; - if (auto attr = - allocLikeOp->getAttrOfType(pto::kPtoMultiBufferAttrName)) { - slotCount = attr.getValue().getZExtValue(); - } - - uint64_t stride = alignUpToDefault(bytes); - uint64_t off = fallbackNextOffset; - for (uint64_t i = 0; i < slotCount; ++i) - offsets.push_back(off + i * stride); - - fallbackNextOffset += - std::max(stride * slotCount, kDefaultAllocAlignmentBytes); - return offsets; -} - -static std::pair getDynamicValidShapeValues(memref::AllocOp op) { - Value vRow; - Value vCol; - auto dynSizes = op.getDynamicSizes(); - if (dynSizes.size() >= kDynamicValidShapeRank) { - vRow = dynSizes[0]; - vCol = dynSizes[1]; - } else if (dynSizes.size() == 1) { - vCol = dynSizes[0]; - } - return {vRow, vCol}; -} -} // namespace - -LogicalResult MemrefAllocaOpToPointerCastOpPattern::matchAndRewrite( - memref::AllocOp op, PatternRewriter &rewriter) const { - const auto ¤tMemRefType = cast(op.getType()); - SmallVector offsets = - getAllocatedOffsets(op.getResult(), op, currentMemRefType, - buffer2Offsets, fallbackNextOffset); - SmallVector addrs; - addrs.reserve(offsets.size()); - for (uint64_t offset : offsets) { - auto constantIntOffsetOp = - rewriter.create(op->getLoc(), offset, 64); - addrs.push_back(constantIntOffsetOp); - } - - auto [vRow, vCol] = getDynamicValidShapeValues(op); - auto ptoPointerCastOp = rewriter.create( - op.getLoc(), currentMemRefType, ValueRange(addrs), vRow ? vRow : Value(), - vCol ? vCol : Value(), Attribute()); - - rewriter.replaceOp(op, ptoPointerCastOp->getResults()); - return success(); -} diff --git a/lib/PTO/Transforms/AllocToPointerCast.h b/lib/PTO/Transforms/AllocToPointerCast.h deleted file mode 100644 index f22de97554..0000000000 --- a/lib/PTO/Transforms/AllocToPointerCast.h +++ /dev/null @@ -1,67 +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. - -//===- AllocToPointerCast.h --Convert memref.AllocOp to pto.pointercastOp-===// -#ifndef LLVM_PROJECT_ALLOCTOPOINTERCAST_H -#define LLVM_PROJECT_ALLOCTOPOINTERCAST_H -#include "PTO/IR/PTO.h" -#include "PTO/Transforms/Passes.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "llvm/ADT/SmallSet.h" - -namespace mlir { -namespace pto { -class MemrefAllocaOpToPointerCastOpPattern - : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - /// map from buffer to its allocated addresses - /// note: the buffer which does multibuffer n optimization will be allocated n - /// addresses. - DenseMap> buffer2Offsets; - mutable uint64_t fallbackNextOffset = 0; - - explicit MemrefAllocaOpToPointerCastOpPattern( - MLIRContext *context, - DenseMap> buffer2Offsets) - : OpRewritePattern(context), - buffer2Offsets(std::move(buffer2Offsets)) { - // Seed fallback offsets above any known planned offsets to reduce collisions. - constexpr uint64_t kAlign = 4096; - uint64_t maxOff = 0; - for (const auto &kv : this->buffer2Offsets) { - for (uint64_t off : kv.second) - maxOff = std::max(maxOff, off); - } - fallbackNextOffset = ((maxOff + kAlign - 1) / kAlign) * kAlign; - } - LogicalResult matchAndRewrite(memref::AllocOp op, - PatternRewriter &rewriter) const final; -}; - -// class UpdateWorkSpaceAllocaOpOffsetPattern -// : public OpRewritePattern { -// public: -// using OpRewritePattern< -// bishengir::memref_ext::AllocWorkspaceOp>::OpRewritePattern; - -// DenseMap> buffer2Offsets; - -// explicit UpdateWorkSpaceAllocaOpOffsetPattern( -// MLIRContext *context, -// DenseMap> buffer2Offsets) -// : OpRewritePattern(context), -// buffer2Offsets(buffer2Offsets) {} -// LogicalResult matchAndRewrite(bishengir::memref_ext::AllocWorkspaceOp op, -// PatternRewriter &rewriter) const final; -// }; -} // namespace pto -} // namespace mlir - -#endif // LLVM_PROJECT_ALLOCTOPOINTERCAST_H diff --git a/lib/PTO/Transforms/CMakeLists.txt b/lib/PTO/Transforms/CMakeLists.txt index e9d9eb7c86..2f416332eb 100644 --- a/lib/PTO/Transforms/CMakeLists.txt +++ b/lib/PTO/Transforms/CMakeLists.txt @@ -51,7 +51,6 @@ add_mlir_dialect_library(PTOTransforms CppPostprocess.cpp Utils.cpp OptMemPlanForPipeline.cpp - AllocToPointerCast.cpp InferPTOMemScope.cpp PTOPlanMemory.cpp PTOPlanMemoryModern.cpp @@ -60,7 +59,6 @@ add_mlir_dialect_library(PTOTransforms PTOA5NormalizeTMovPass.cpp PTORemoveIdentityTMov.cpp PTOCanonicalizeIR.cpp - PTOMaterializeTileHandles.cpp BufferizableOpInterfaceImpl.cpp ConvertToPTOOp.cpp PTOAssignDefaultFrontendPipeIdPass.cpp diff --git a/lib/PTO/Transforms/OptMemPlanForPipeline.cpp b/lib/PTO/Transforms/OptMemPlanForPipeline.cpp index 1a9162892f..ff6091ac81 100644 --- a/lib/PTO/Transforms/OptMemPlanForPipeline.cpp +++ b/lib/PTO/Transforms/OptMemPlanForPipeline.cpp @@ -15,17 +15,7 @@ using namespace mlir::detail; using namespace mlir::pto; void OptMemPlanForDma::build(func::FuncOp func) { - auto result = func->walk([&](Operation *op) { - if (auto loadOp = dyn_cast(op)) { - UpdateScalarBuffers(loadOp); - } else if (auto storeOp = dyn_cast(op)) { - UpdateScalarBuffers(storeOp); - } - return WalkResult::advance(); - }); - if (result == WalkResult::interrupt()) { - llvm_unreachable("OptMemPlanForLoop Traverse IR Failed! "); - } + (void)func; } void OptMemPlanForDma::UpdateDmaBuffers(SmallVector dpsOperand) { @@ -34,7 +24,7 @@ void OptMemPlanForDma::UpdateDmaBuffers(SmallVector dpsOperand) { if (!isLocalBuffer(memorySpaceAttr)) { continue; } - DmaBuffers.insert(tracebackMemRef(operand)); + DmaBuffers.insert(operand); } } @@ -75,18 +65,6 @@ bool OptMemPlanForDma::BufferPipeConflict(const Value buf1, return false; } -template -typename std::enable_if || - std::is_same_v, - void>::type -OptMemPlanForDma::UpdateScalarBuffers(OP op) { - auto memorySpaceAttr = GetBufferSpaceAttr(op.getMemRef()); - if (!isLocalBuffer(memorySpaceAttr)) { - return; - } - ScalarBuffers.insert(tracebackMemRef(op.getMemRef())); -} - bool OptMemPlanForDma::IsScalarBuffer(const Value buf) const { if (ScalarBuffers.empty()) { return false; @@ -105,6 +83,6 @@ void OptMemPlanForDma::UpdateScalarBuffersForLowerToLoops(Operation *op) { if (!isLocalBuffer(memorySpaceAttr)) { continue; } - ScalarBuffers.insert(tracebackMemRef(operand)); + ScalarBuffers.insert(operand); } } diff --git a/lib/PTO/Transforms/OptMemPlanForPipeline.h b/lib/PTO/Transforms/OptMemPlanForPipeline.h index 553e68c3a4..6475ded13c 100644 --- a/lib/PTO/Transforms/OptMemPlanForPipeline.h +++ b/lib/PTO/Transforms/OptMemPlanForPipeline.h @@ -10,11 +10,8 @@ #ifndef OPT_MEM_PLAN_FOR_PIPELINE_H #define OPT_MEM_PLAN_FOR_PIPELINE_H #include "PTO/IR/PTO.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "Utils.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" namespace mlir { namespace pto { @@ -38,12 +35,6 @@ class OptMemPlanForDma { /// Update the buffers for MTE2 and MTE3. void UpdateDmaBuffers(SmallVector dpsOperand); - template - typename std::enable_if || - std::is_same_v, - void>::type - UpdateScalarBuffers(OP op); - void UpdateScalarBuffersForLowerToLoops(Operation *operands); /// Buffer in MTE2 and MTE3. diff --git a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp b/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp deleted file mode 100644 index ec8f7201e8..0000000000 --- a/lib/PTO/Transforms/PTOMaterializeTileHandles.cpp +++ /dev/null @@ -1,890 +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. - -//===- PTOMaterializeTileHandles.cpp -------------------------------------===// -//===----------------------------------------------------------------------===// -// -// Reintroduce explicit tile_buf handles after memory planning/sync have used -// memref IR. EmitC can then lower tile operations from tile-typed operands -// instead of rediscovering tile metadata from every memref use. - -#include "PTO/IR/PTO.h" -#include "PTO/IR/PTOTypeUtils.h" -#include "PTO/Transforms/Passes.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Pass/Pass.h" - -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/DenseSet.h" - -#include -#include -#include - -namespace mlir { -namespace pto { -#define GEN_PASS_DEF_PTOMATERIALIZETILEHANDLES -#include "PTO/Transforms/Passes.h.inc" -} // namespace pto -} // namespace mlir - -using namespace mlir; - -namespace mlir { -namespace pto { -namespace { - -static constexpr llvm::StringLiteral kForceDynamicValidShapeAttrName = - "__pto.force_dynamic_valid_shape"; - -static int64_t getIntegerAttrSignedValue(IntegerAttr attr) { - const APInt &value = attr.getValue(); - return value.getBitWidth() == 0 ? 0 : value.getSExtValue(); -} - -struct TileHandleMetadata { - Value source; - Value validRow; - Value validCol; - TileBufConfigAttr config; - bool explicitConfig = false; - SmallVector attrs; -}; - -static bool isLocalTileMemRef(Type type) { - auto memTy = dyn_cast(type); - if (!memTy || memTy.getRank() != 2) - return false; - - auto asAttr = dyn_cast_or_null(memTy.getMemorySpace()); - if (!asAttr) - return false; - - switch (asAttr.getAddressSpace()) { - case AddressSpace::GM: - case AddressSpace::Zero: - return false; - case AddressSpace::VEC: - case AddressSpace::MAT: - case AddressSpace::LEFT: - case AddressSpace::RIGHT: - case AddressSpace::ACC: - case AddressSpace::BIAS: - case AddressSpace::SCALING: - return true; - } - return false; -} - -static bool shouldMaterializeOperand(Operation *owner) { - if (isa(owner)) - return false; - - StringRef name = owner->getName().getStringRef(); - if (name == "pto.set_validshape") - return true; - if (name == "pto.get_validshape") - return true; - if (!name.consume_front("pto.")) - return false; - return name.starts_with("t"); -} - -static bool shouldMaterializeYieldOperand(Operation *owner) { - return isa(owner); -} - -static bool hasStringAttr(ArrayRef attrs, StringRef name, - StringRef value) { - return llvm::any_of(attrs, [&](NamedAttribute attr) { - if (attr.getName().getValue() != name) - return false; - auto strAttr = dyn_cast(attr.getValue()); - return strAttr && strAttr.getValue() == value; - }); -} - -static bool hasAttr(ArrayRef attrs, StringRef name) { - return llvm::any_of(attrs, [&](NamedAttribute attr) { - return attr.getName().getValue() == name; - }); -} - -static int64_t getConstantIndexOrDynamic(Value value) { - if (!value) - return ShapedType::kDynamic; - if (auto cst = value.getDefiningOp()) - return cst.value(); - if (auto cst = value.getDefiningOp()) - return cst.value(); - return ShapedType::kDynamic; -} - -static void copyTileHandleAttrs(Operation *from, - SmallVectorImpl &attrs) { - StringRef names[] = {"pto.view_semantics", kForceDynamicValidShapeAttrName}; - for (StringRef name : names) { - if (Attribute attr = from->getAttr(name)) - attrs.push_back(NamedAttribute(StringAttr::get(from->getContext(), name), - attr)); - } -} - -static std::optional getAddressSpace(Type type) { - if (auto memTy = dyn_cast(type)) { - if (auto asAttr = - dyn_cast_or_null(memTy.getMemorySpace())) - return asAttr.getAddressSpace(); - return std::nullopt; - } - if (auto tileTy = dyn_cast(type)) { - if (auto asAttr = - dyn_cast_or_null(tileTy.getMemorySpace())) - return asAttr.getAddressSpace(); - } - return std::nullopt; -} - -static bool isA5Target(Operation *op) { - auto module = op->getParentOfType(); - if (!module) - return false; - - if (auto arch = module->getAttrOfType("pto.target_arch")) { - if (arch.getValue().equals_insensitive("a5")) - return true; - } - if (auto spec = module->getAttrOfType("pto.device-spec")) { - StringRef value = spec.getValue(); - if (value.starts_with("Ascend950") || value.starts_with("Ascend910_95")) - return true; - } - return false; -} - -static TileBufConfigAttr makeTileConfig(MLIRContext *ctx, BLayout bl, - SLayout sl) { - Builder builder(ctx); - return TileBufConfigAttr::get( - ctx, BLayoutAttr::get(ctx, bl), SLayoutAttr::get(ctx, sl), - builder.getI32IntegerAttr(512), PadValueAttr::get(ctx, PadValue::Null), - CompactModeAttr::get(ctx, CompactMode::Null)); -} - -static void inferConfigForMaterializedUse(Operation *owner, unsigned operandNo, - Type operandType, - TileHandleMetadata &meta, - MLIRContext *ctx) { - if (meta.explicitConfig) - return; - - auto colRow = [&]() { - return makeTileConfig(ctx, BLayout::ColMajor, SLayout::RowMajor); - }; - auto rowCol = [&]() { - return makeTileConfig(ctx, BLayout::RowMajor, SLayout::ColMajor); - }; - - if (isa(owner)) { - if (!isA5Target(owner)) - return; - if (operandNo == 0 || operandNo == 2) - meta.config = colRow(); - else if (operandNo == 1) - meta.config = rowCol(); - return; - } - - if (isa(owner)) { - if (!isA5Target(owner)) - return; - if (operandNo == 0 || operandNo == 1 || operandNo == 3) - meta.config = colRow(); - else if (operandNo == 2) - meta.config = rowCol(); - return; - } - - if (isa(owner)) { - if (operandNo != 0 && operandNo != 3) - return; - auto as = getAddressSpace(operandType); - if (!as) - return; - if (*as == AddressSpace::ACC || *as == AddressSpace::MAT) - meta.config = colRow(); - } -} - -static TileHandleMetadata getTileHandleMetadata(Value value, - MLIRContext *ctx) { - TileHandleMetadata meta; - meta.source = value; - meta.config = TileBufConfigAttr::getDefault(ctx); - - if (auto cast = value.getDefiningOp()) { - meta.validRow = cast.getValidRow(); - meta.validCol = cast.getValidCol(); - if (auto config = cast.getConfig()) { - meta.config = *config; - meta.explicitConfig = true; - } - copyTileHandleAttrs(cast, meta.attrs); - return meta; - } - - return meta; -} - -static bool getTilePointerStrides(TileBufConfigAttr configAttr, Type elemTy, - int64_t rows, int64_t cols, - int64_t &rowStride, int64_t &colStride) { - if (rows == ShapedType::kDynamic || cols == ShapedType::kDynamic) - return false; - - int32_t blVal = 0; - if (auto blAttr = dyn_cast(configAttr.getBLayout())) - blVal = static_cast(blAttr.getValue()); - else if (auto intAttr = dyn_cast(configAttr.getBLayout())) - blVal = static_cast(getIntegerAttrSignedValue(intAttr)); - - int32_t slVal = 0; - if (auto slAttr = dyn_cast(configAttr.getSLayout())) - slVal = static_cast(slAttr.getValue()); - else if (auto intAttr = dyn_cast(configAttr.getSLayout())) - slVal = static_cast(getIntegerAttrSignedValue(intAttr)); - - bool boxed = slVal != 0; - int64_t innerRows = 1; - int64_t innerCols = 1; - if (boxed) { - int32_t fractal = 512; - if (auto frAttr = dyn_cast(configAttr.getSFractalSize())) - fractal = static_cast(getIntegerAttrSignedValue(frAttr)); - - unsigned elemBytes = pto::getPTOStorageElemByteSize(elemTy); - if (elemBytes == 0) - return false; - - switch (fractal) { - case 1024: - innerRows = 16; - innerCols = 16; - break; - case 32: - innerRows = 16; - innerCols = 2; - break; - case 512: - if (slVal == 1) { - innerRows = 16; - innerCols = 32 / elemBytes; - } else if (slVal == 2) { - innerRows = 32 / elemBytes; - innerCols = 16; - } else { - return false; - } - break; - default: - return false; - } - if (innerRows <= 0 || innerCols <= 0) - return false; - } - - if (!boxed) { - if (blVal == 1) { - rowStride = 1; - colStride = rows; - } else { - rowStride = cols; - colStride = 1; - } - return true; - } - - if (blVal == 1) { - if (slVal != 1) - return false; - rowStride = innerCols; - colStride = rows; - return true; - } - - rowStride = cols; - colStride = innerRows; - return true; -} - -static SmallVector -getMaterializedTileShape(MemRefType memTy, const TileHandleMetadata &meta) { - SmallVector shape(memTy.getShape().begin(), - memTy.getShape().end()); - if (!hasStringAttr(meta.attrs, "pto.view_semantics", "subview")) - return shape; - - auto sourceMrTy = dyn_cast_or_null(meta.source.getType()); - if (!sourceMrTy || sourceMrTy.getRank() < 2 || - !meta.source.getDefiningOp()) - return shape; - - int64_t subRows = sourceMrTy.getDimSize(0); - int64_t subCols = sourceMrTy.getDimSize(1); - if (pto::isPTOFloat4PackedType(sourceMrTy.getElementType()) || - subRows == ShapedType::kDynamic || subCols == ShapedType::kDynamic) - return shape; - - SmallVector inheritedStrides; - int64_t inheritedOffset = ShapedType::kDynamic; - if (failed(mlir::pto::getPTOMemRefStridesAndOffset( - sourceMrTy, inheritedStrides, inheritedOffset)) || - inheritedStrides.size() < 2) - return shape; - - int64_t childRowStride = 0; - int64_t childColStride = 0; - if (!getTilePointerStrides(meta.config, sourceMrTy.getElementType(), subRows, - subCols, childRowStride, childColStride)) - return shape; - - if (inheritedStrides[0] == childRowStride && - inheritedStrides[1] == childColStride) { - shape[0] = subRows; - shape[1] = subCols; - } - - return shape; -} - -static TileBufType buildTileTypeFromMemRef(MemRefType memTy, - const TileHandleMetadata &meta, - MLIRContext *ctx) { - SmallVector shape = getMaterializedTileShape(memTy, meta); - SmallVector validShape(shape.begin(), shape.end()); - bool forceDynamic = hasAttr(meta.attrs, kForceDynamicValidShapeAttrName); - if (forceDynamic) { - validShape[0] = ShapedType::kDynamic; - validShape[1] = ShapedType::kDynamic; - } else { - if (meta.validRow) - validShape[0] = getConstantIndexOrDynamic(meta.validRow); - if (meta.validCol) - validShape[1] = getConstantIndexOrDynamic(meta.validCol); - } - - return TileBufType::get(ctx, shape, memTy.getElementType(), - memTy.getMemorySpace(), validShape, meta.config); -} - -static bool isMaterializedTileAnchor(Operation *op) { - return isa(op); -} - -static Value makeI64Constant(OpBuilder &builder, Location loc, int64_t value) { - return builder.create(loc, value, 64); -} - -static Value ensureI64(Value value, OpBuilder &builder, Location loc) { - if (!value) - return Value(); - - auto i64Ty = builder.getI64Type(); - if (value.getType() == i64Ty) - return value; - if (isa(value.getType())) - return builder.create(loc, i64Ty, value); - if (auto intTy = dyn_cast(value.getType())) { - if (intTy.getWidth() == 64) - return value; - if (intTy.getWidth() < 64) - return builder.create(loc, i64Ty, value); - return builder.create(loc, i64Ty, value); - } - return Value(); -} - -static Value materializeOffset(OpFoldResult ofr, OpBuilder &builder, - Location loc) { - if (isa(ofr)) { - Attribute attr = cast(ofr); - if (auto intAttr = dyn_cast(attr)) - return makeI64Constant(builder, loc, getIntegerAttrSignedValue(intAttr)); - return Value(); - } - return ensureI64(cast(ofr), builder, loc); -} - -static Value addI64(Value lhs, Value rhs, OpBuilder &builder, Location loc) { - if (!lhs) - return rhs; - if (!rhs) - return lhs; - return builder.create(loc, lhs, rhs); -} - -static Value mulI64(Value lhs, int64_t rhs, OpBuilder &builder, Location loc) { - if (!lhs) - return Value(); - if (rhs == 0) - return makeI64Constant(builder, loc, 0); - if (rhs == 1) - return lhs; - return builder.create(loc, lhs, - makeI64Constant(builder, loc, rhs)); -} - -static Value computeExplicitAddress(Value value, OpBuilder &builder, - Location loc); - -static Value computeSubviewAddress(memref::SubViewOp subview, - OpBuilder &builder, Location loc) { - Value base = computeExplicitAddress(subview.getSource(), builder, loc); - if (!base) - return Value(); - - auto sourceTy = dyn_cast(subview.getSource().getType()); - if (!sourceTy) - return Value(); - unsigned elemBytes = getPTOStorageElemByteSize(sourceTy.getElementType()); - if (elemBytes == 0) - return Value(); - - SmallVector sourceStrides; - int64_t sourceOffset = ShapedType::kDynamic; - if (failed(mlir::pto::getPTOMemRefStridesAndOffset( - sourceTy, sourceStrides, sourceOffset))) - return Value(); - - auto mixedOffsets = subview.getMixedOffsets(); - if (sourceStrides.size() < mixedOffsets.size()) - return Value(); - - Value linearOffset; - for (auto [offsetOfr, stride] : - llvm::zip_equal(mixedOffsets, ArrayRef(sourceStrides).take_front( - mixedOffsets.size()))) { - if (stride == ShapedType::kDynamic) - return Value(); - Value offset = materializeOffset(offsetOfr, builder, loc); - if (!offset) - return Value(); - linearOffset = addI64(linearOffset, mulI64(offset, stride, builder, loc), - builder, loc); - } - - if (!linearOffset) - return base; - linearOffset = mulI64(linearOffset, elemBytes, builder, loc); - return builder.create(loc, base, linearOffset); -} - -static Value computePTOSubViewAddress(pto::SubViewOp subview, - OpBuilder &builder, Location loc) { - Value base = computeExplicitAddress(subview.getSource(), builder, loc); - if (!base) - return Value(); - - auto sourceTy = dyn_cast(subview.getSource().getType()); - if (!sourceTy || sourceTy.getShape().size() < 2) - return Value(); - - unsigned elemBytes = getPTOStorageElemByteSize(sourceTy.getElementType()); - if (elemBytes == 0) - return Value(); - - int64_t rowStride = 0; - int64_t colStride = 0; - if (!getTilePointerStrides(sourceTy.getConfigAttr(), sourceTy.getElementType(), - sourceTy.getShape()[0], sourceTy.getShape()[1], - rowStride, colStride)) - return Value(); - - SmallVector strides = {rowStride, colStride}; - Value linearOffset; - for (auto [offset, stride] : - llvm::zip(subview.getOffsets(), ArrayRef(strides))) { - Value offsetI64 = ensureI64(offset, builder, loc); - if (!offsetI64) - return Value(); - linearOffset = - addI64(linearOffset, mulI64(offsetI64, stride, builder, loc), builder, - loc); - } - - if (!linearOffset) - return base; - linearOffset = mulI64(linearOffset, elemBytes, builder, loc); - return builder.create(loc, base, linearOffset); -} - -static Value computeExplicitAddress(Value value, OpBuilder &builder, - Location loc) { - if (auto alloc = value.getDefiningOp()) - return ensureI64(alloc.getAddr(), builder, loc); - - if (auto cast = value.getDefiningOp()) { - if (cast.getAddrs().empty()) - return Value(); - return ensureI64(cast.getAddrs().front(), builder, loc); - } - - if (auto subview = value.getDefiningOp()) - return computeSubviewAddress(subview, builder, loc); - - if (auto subview = value.getDefiningOp()) - return computePTOSubViewAddress(subview, builder, loc); - - if (auto select = value.getDefiningOp()) { - Value trueAddr = computeExplicitAddress(select.getTrueValue(), builder, loc); - Value falseAddr = - computeExplicitAddress(select.getFalseValue(), builder, loc); - if (!trueAddr || !falseAddr) - return Value(); - return builder.create(loc, select.getCondition(), - trueAddr, falseAddr); - } - - if (auto cast = value.getDefiningOp()) - return computeExplicitAddress(cast.getSource(), builder, loc); - - if (auto cast = value.getDefiningOp()) { - if (cast.getNumOperands() == 1) - return computeExplicitAddress(cast.getOperand(0), builder, loc); - } - - return Value(); -} - -static bool isControlFlowAddressProducer(Operation *op) { - if (!op) - return false; - - StringRef name = op->getName().getStringRef(); - return name == "scf.if" || name == "scf.for" || name == "scf.while" || - name == "scf.execute_region" || name == "scf.index_switch"; -} - -static Value peelAddressSource(Value value) { - while (true) { - if (auto subview = value.getDefiningOp()) { - value = subview.getSource(); - continue; - } - - if (auto cast = value.getDefiningOp()) { - value = cast.getSource(); - continue; - } - - return value; - } -} - -static bool isFunctionEntryBlockArgument(BlockArgument arg) { - Operation *parent = arg.getOwner()->getParentOp(); - auto func = dyn_cast_or_null(parent); - return func && arg.getOwner() == &func.getBody().front(); -} - -static bool isUnsupportedControlFlowAddress(Value value) { - value = peelAddressSource(value); - - if (auto arg = dyn_cast(value)) - return !isFunctionEntryBlockArgument(arg); - - return isControlFlowAddressProducer(value.getDefiningOp()); -} - -static void emitMissingExplicitAddressError(Operation *owner, Value value) { - value = peelAddressSource(value); - auto diag = owner->emitOpError() - << "cannot materialize tile handle for local memref because its " - "explicit byte address cannot be recovered"; - - if (isa(value)) { - diag << "; region block arguments and loop-carried memref values are " - "unsupported here"; - return; - } - - Operation *def = value.getDefiningOp(); - if (!def) { - diag << "; value has no defining op"; - return; - } - - if (isControlFlowAddressProducer(def)) { - diag << "; control-flow result '" << def->getName() - << "' cannot carry a local memref into tile materialization"; - return; - } - - diag << "; unsupported defining op is '" << def->getName() << "'"; -} - -static Value getAllocValidOperand(TileBufType tileTy, Value operand, - unsigned dim, OpBuilder &builder, - Location loc) { - auto validShape = tileTy.getValidShape(); - if (validShape.size() <= dim || validShape[dim] >= 0) - return Value(); - if (operand) - return operand; - - auto shape = tileTy.getShape(); - if (shape.size() > dim && shape[dim] != ShapedType::kDynamic) - return builder.create(loc, shape[dim]); - return Value(); -} - -static Attribute getAttr(ArrayRef attrs, StringRef name) { - for (NamedAttribute attr : attrs) { - if (attr.getName().getValue() == name) - return attr.getValue(); - } - return {}; -} - -static void copyMaterializedTileAttrs(ArrayRef attrs, - Operation *to) { - if (Attribute attr = getAttr(attrs, kForceDynamicValidShapeAttrName)) - to->setAttr(kForceDynamicValidShapeAttrName, attr); -} - -static void updateResultTypesAfterMaterializingOperand(Operation *op, - unsigned operandNo, - Type tileTy) { - if (auto tassign = dyn_cast(op)) { - if (operandNo == 0) - tassign.getResult().setType(tileTy); - } -} - -static bool isTileViewSemantics(StringAttr viewSemantics) { - return viewSemantics && (viewSemantics.getValue() == "treshape" || - viewSemantics.getValue() == "bitcast"); -} - -static Value materializeAnchorResult(Operation *anchor, Value anchoredValue, - OpBuilder &builder, MLIRContext *ctx, - DenseMap &tileHandles, - const DenseSet &mustMaterialize, - bool &failedMaterialization) { - auto memTy = dyn_cast(anchoredValue.getType()); - if (!memTy || !isLocalTileMemRef(memTy)) - return Value(); - - SmallVector usesToRewrite; - for (OpOperand &use : anchoredValue.getUses()) { - if (shouldMaterializeOperand(use.getOwner()) || - shouldMaterializeYieldOperand(use.getOwner())) - usesToRewrite.push_back(&use); - } - - TileHandleMetadata meta = getTileHandleMetadata(anchoredValue, ctx); - auto viewSemantics = dyn_cast_or_null( - getAttr(meta.attrs, "pto.view_semantics")); - bool isTileView = isTileViewSemantics(viewSemantics); - if (usesToRewrite.empty() && !isTileView && - !mustMaterialize.contains(anchoredValue)) - return Value(); - - if (Value existing = tileHandles.lookup(anchoredValue)) { - for (OpOperand *use : usesToRewrite) { - Operation *owner = use->getOwner(); - unsigned operandNo = use->getOperandNumber(); - use->set(existing); - updateResultTypesAfterMaterializingOperand(owner, operandNo, - existing.getType()); - } - return existing; - } - - for (OpOperand *use : usesToRewrite) - inferConfigForMaterializedUse(use->getOwner(), use->getOperandNumber(), - anchoredValue.getType(), meta, ctx); - auto tileTy = buildTileTypeFromMemRef(memTy, meta, ctx); - - builder.setInsertionPointAfter(anchor); - Value materialized; - Value sourceTile = meta.source ? tileHandles.lookup(meta.source) : Value(); - if (sourceTile && viewSemantics && - viewSemantics.getValue() == "treshape") { - materialized = - builder.create(anchor->getLoc(), tileTy, sourceTile) - .getResult(); - } else if (sourceTile && viewSemantics && - viewSemantics.getValue() == "bitcast") { - materialized = - builder.create(anchor->getLoc(), tileTy, sourceTile) - .getResult(); - } else { - Value addr = computeExplicitAddress(anchoredValue, builder, anchor->getLoc()); - if (!addr && isUnsupportedControlFlowAddress(anchoredValue)) { - emitMissingExplicitAddressError(anchor, anchoredValue); - failedMaterialization = true; - return Value(); - } - auto alloc = builder.create( - anchor->getLoc(), tileTy, addr ? addr : Value(), - getAllocValidOperand(tileTy, meta.validRow, 0, builder, - anchor->getLoc()), - getAllocValidOperand(tileTy, meta.validCol, 1, builder, - anchor->getLoc())); - copyMaterializedTileAttrs(meta.attrs, alloc); - materialized = alloc.getResult(); - } - - for (OpOperand *use : usesToRewrite) { - Operation *owner = use->getOwner(); - unsigned operandNo = use->getOperandNumber(); - use->set(materialized); - updateResultTypesAfterMaterializingOperand(owner, operandNo, tileTy); - } - - tileHandles[anchoredValue] = materialized; - return materialized; -} - -struct PTOMaterializeTileHandlesPass - : public impl::PTOMaterializeTileHandlesBase< - PTOMaterializeTileHandlesPass> { - void runOnOperation() override { - ModuleOp module = getOperation(); - MLIRContext *ctx = module.getContext(); - - bool hasUnresolvedMultiTile = false; - module.walk([&](Operation *op) { - if (!isa(op)) - return; - op->emitError( - "must be resolved by pto-resolve-buffer-select before " - "pto-materialize-tile-handles"); - hasUnresolvedMultiTile = true; - }); - if (hasUnresolvedMultiTile) { - signalPassFailure(); - return; - } - - OpBuilder builder(ctx); - DenseMap tileHandles; - DenseSet mustMaterialize; - bool failedMaterialization = false; - - SmallVector anchors; - module.walk([&](Operation *op) { - if (!isMaterializedTileAnchor(op)) - return; - anchors.push_back(op); - }); - for (Operation *anchor : anchors) { - if (anchor->getNumResults() != 1) - continue; - Value anchoredValue = anchor->getResult(0); - if (!isLocalTileMemRef(anchoredValue.getType())) - continue; - TileHandleMetadata meta = getTileHandleMetadata(anchoredValue, ctx); - auto viewSemantics = dyn_cast_or_null( - getAttr(meta.attrs, "pto.view_semantics")); - if (isTileViewSemantics(viewSemantics) && meta.source) - mustMaterialize.insert(meta.source); - } - - for (Operation *anchor : anchors) { - if (anchor->getNumResults() != 1) - continue; - materializeAnchorResult(anchor, anchor->getResult(0), builder, ctx, - tileHandles, mustMaterialize, - failedMaterialization); - } - - if (failedMaterialization) { - signalPassFailure(); - return; - } - - SmallVector, 32> operandsToRewrite; - module.walk([&](Operation *op) { - if (!shouldMaterializeOperand(op)) - return; - for (OpOperand &operand : op->getOpOperands()) { - if (isLocalTileMemRef(operand.get().getType())) - operandsToRewrite.push_back({op, operand.getOperandNumber()}); - } - }); - - for (auto [op, operandNo] : operandsToRewrite) { - Value oldValue = op->getOperand(operandNo); - if (!isa(oldValue.getType())) - continue; - if (op->getName().getStringRef() == "pto.tassign" && operandNo == 0) - continue; - - if (Value existing = tileHandles.lookup(oldValue)) { - op->setOperand(operandNo, existing); - updateResultTypesAfterMaterializingOperand(op, operandNo, - existing.getType()); - continue; - } - - auto memTy = cast(oldValue.getType()); - TileHandleMetadata meta = getTileHandleMetadata(oldValue, ctx); - inferConfigForMaterializedUse(op, operandNo, oldValue.getType(), meta, - ctx); - auto tileTy = buildTileTypeFromMemRef(memTy, meta, ctx); - - builder.setInsertionPoint(op); - Value materialized; - Value addr = computeExplicitAddress(oldValue, builder, op->getLoc()); - if (!addr && isUnsupportedControlFlowAddress(oldValue)) { - emitMissingExplicitAddressError(op, oldValue); - failedMaterialization = true; - continue; - } - auto alloc = builder.create( - op->getLoc(), tileTy, addr ? addr : Value(), - getAllocValidOperand(tileTy, meta.validRow, 0, builder, - op->getLoc()), - getAllocValidOperand(tileTy, meta.validCol, 1, builder, - op->getLoc())); - copyMaterializedTileAttrs(meta.attrs, alloc); - materialized = alloc.getResult(); - tileHandles[oldValue] = materialized; - op->setOperand(operandNo, materialized); - updateResultTypesAfterMaterializingOperand(op, operandNo, tileTy); - } - - if (failedMaterialization) { - signalPassFailure(); - return; - } - - } -}; - -} // namespace - -std::unique_ptr createPTOMaterializeTileHandlesPass() { - return std::make_unique(); -} - -} // namespace pto -} // namespace mlir diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index daa804e043..a822324cc2 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -20,7 +20,6 @@ #include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/IR/AsmState.h" #include "mlir/Transforms/DialectConversion.h" -#include "AllocToPointerCast.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "llvm/Support/Debug.h" @@ -507,33 +506,6 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { buffer2MultiNum[allocMultiOp.getResult()] = count; UpdateOpBufferInfo(op, op->getResults()); return WalkResult::advance(); - } else if (isLocalMemPlan() && dyn_cast(op)) { - auto allocOp = cast(op); - if (failed(CheckLocalBufferAllocOp(op))) { - return WalkResult::interrupt(); - } - // Pick up the multi-buffer slot count when present. Range checking - // mirrors the type-level verifier so a malformed memref alloc (e.g. - // from a hand-written test) still gets a clear diagnostic. The - // sibling expansion in `ExpandMultiBufferStorageEntry` supports any - // N in the legal range. - if (auto attr = allocOp->getAttrOfType( - mlir::pto::kPtoMultiBufferAttrName)) { - uint64_t n = attr.getValue().getZExtValue(); - if (n < mlir::pto::kPtoMultiBufferMinNum || - n > mlir::pto::kPtoMultiBufferMaxNum) { - allocOp.emitError() - << "pto.multi_buffer must be in [" - << mlir::pto::kPtoMultiBufferMinNum << ", " - << mlir::pto::kPtoMultiBufferMaxNum << "] (got " << n << ")"; - return WalkResult::interrupt(); - } - buffer2MultiNum[allocOp.getResult()] = static_cast(n); - } - UpdateOpBufferInfo(op, op->getResults()); - return WalkResult::advance(); - } else if (auto loadOp = dyn_cast(op)) { - OpKillHandle(curOpInfo, live, op->getBlock()); } else if (auto tprintOp = dyn_cast(op)) { // TPrintOp only reads from buffer, similar to LoadOp UpdateOpGenInfo(curOpInfo, llvm::to_vector(op->getOperands())); @@ -555,11 +527,8 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { // buffer for liveness, but the scalar row/col results are not buffers. UpdateOpGenInfo(curOpInfo, ValueRange{op->getOperand(0)}); OpKillHandle(curOpInfo, live, op->getBlock()); - } else if (auto storeOp = dyn_cast(op)) { - UpdateStoreOpInfo(curOpInfo, storeOp.getMemRef(), live); } else if (auto ptoDpsOp = dyn_cast(op)) { - // PTO ops with destination (tile_buf, partition_view, etc.); no - // tensor/memref-only verification. + // PTO ops with destination (tile_buf, partition_view, etc.). SmallVector genBuffers = llvm::to_vector(op->getOperands()); auto scratchBuffers = getScratchBuffersFromEffects( op, ptoDpsOp.getDpsInits(), stableValueOrder); @@ -572,10 +541,7 @@ void MemLivenessAnalysis::RecursionIR(Region *region, Liveness live) { } OpKillHandle(curOpInfo, live, op->getBlock()); } else if (auto dstStyleOp = dyn_cast(op)) { - // Process the operation of pto instructions as follows: - // pto.hir.copy ins(%0 : memref<16xf16, #pto.address_space>) - // outs(%1 : memref<16xxf16, #pto.address_space>) - // need to handle kill buffer. + // Preserve generic destination-style aliasing for non-tile tensors. UpdateInitAndResAlias(dstStyleOp); UpdateOpGenInfo(curOpInfo, llvm::to_vector(dstStyleOp.getDpsInits())); OpKillHandle(curOpInfo, live, op->getBlock()); @@ -662,14 +628,7 @@ void MemLivenessAnalysis::UpdateForOpBufferAlias(scf::ForOp forOp) { } void MemLivenessAnalysis::RecursiveForOp(scf::ForOp forOp, Liveness live) { - // Process the operation of ForOp as follows: - // alloca %allocA - // %0 = scf.for %arg4 = %c0 to %c1024 step %c128 iter_args(%arg5 = %4)-> - // (memref<16x16x16xf16, #pto.address_space>): - // def(allocA) - // ... - // scf.yield %alloc0 : memref<16xf16,#pto.address_space> - // need to handle kill buffer. + // Model loop-carried tile handles as aliases of the yielded roots. auto forBeginSeq = UpdateLinearOperation(forOp.getOperation()); UpdateOpGenInfo(forBeginSeq, GetLiveBuffersInLoop(forOp, live)); UpdateForOpInitArgsAlias(forOp); @@ -705,11 +664,7 @@ void MemLivenessAnalysis::UpdateIfOpBufferAlias(scf::IfOp ifOp, } void MemLivenessAnalysis::RecursiveIfOp(scf::IfOp ifOp, Liveness live) { - // Process the operation of IfOp as follows: - // %0 = scf.if %cond -> (memref<16xf16, #pto.address_space>) - // scf.yield %alloc0: memref<16xf16, #pto.address_space> - // else: - // scf.yield %alloc1 : memref<16xf16, #pto.address_space> + // Join the tile roots yielded by each branch into the if result. UpdateLinearOperation(ifOp.getOperation()); RecursionIR(&ifOp.getThenRegion(), live); auto curIfElse = UpdateLinearOperation(ifOp.getOperation()); @@ -779,24 +734,11 @@ SmallVector MemLivenessAnalysis::GetLiveBuffersInLoop(scf::ForOp forOp, return allocBeforeLoopBuffers; } -LogicalResult -MemLivenessAnalysis::CheckLocalBufferAllocOp(Operation *op) const { - auto allocOp = dyn_cast(op); - if (!allocOp) - return op->emitError("must be alloc op"), failure(); - auto memorySpaceAttr = GetBufferSpaceAttr(allocOp.getResult()); - if (isLocalBuffer(memorySpaceAttr)) { - return success(); - } - allocOp.getOperation()->emitError("Alloc buffer not at UB space! "); - return failure(); -} - bool MemLivenessAnalysis::isSkippableOp(Operation *op) const { // Call-like ops are still modeled explicitly. Only pure terminators and // dim queries are skipped here. // - return isa(op); + return isa(op); } LogicalResult @@ -832,8 +774,7 @@ void MemLivenessAnalysis::UpdateBufferAlias(Value buffer, Value aliasBuffer, buffer2AliasVec[buf] = clonedAliasSet; } - // mark the alias buffer as ignoring Inplace if it is not generated by - // memref.alloc. + // Mark aliases that must not participate in inplace merging. auto it = bufferInfos.find(aliasBuffer); if (isIgnoreInplace && it != bufferInfos.end()) { it->second.ignoreInplace = true; @@ -859,17 +800,6 @@ SetVector MemLivenessAnalysis::GetAliasBuffers(Value aliasBuffer) { return {}; } -void MemLivenessAnalysis::UpdateStoreOpInfo(OpInfo *opInfo, - const Value storeValue, - Liveness live) { - // The src of memref store may also serve as a gen buffer. - SmallVector storeValues; - storeValues.push_back(storeValue); - UpdateOpGenInfo(opInfo, storeValues); - // Collect kill buffers corresponding to operation. - OpKillHandle(opInfo, live, opInfo->operation->getBlock()); -} - void MemLivenessAnalysis::UpdateOpBufferInfo(Operation *op, const ValueRange &results) { for (const Value &operand : results) { @@ -1002,10 +932,7 @@ BufferInfo MemLivenessAnalysis::GetBufferInfo(Operation *op, Value operand, shape = multiType.getSlotType().getShape(); elementType = multiType.getSlotType().getElementType(); } else { - Value traceValue = tracebackMemRef(operand); - auto memRefType = cast(traceValue.getType()); - shape = memRefType.getShape(); - elementType = memRefType.getElementType(); + llvm_unreachable("local memory planner expects tile buffer roots"); } bufferInfo.bufferType = elementType; std::optional totalStaticSize = @@ -1679,13 +1606,8 @@ void MemPlan::ReportMemLifeDebugInfo(StorageEntry *rootStorageEntry) { } void MemPlan::MemLifeDebugInfo(StorageEntry *storageEntry) { - for (auto &buffer : storageEntry->inplaceBuffers) { - if (buffer.getDefiningOp()) { - if (auto allocOp = dyn_cast(buffer.getDefiningOp())) { - LDBG("Buffer : " << allocOp.getResult() << "\n"); - } - } - } + for (auto &buffer : storageEntry->inplaceBuffers) + LDBG("Buffer : " << buffer << "\n"); for (auto &bufferLife : storageEntry->bufferLifeVec) { (void)bufferLife; LDBG("bufferLife : " @@ -1696,14 +1618,8 @@ void MemPlan::MemLifeDebugInfo(StorageEntry *storageEntry) { } void MemPlan::ReportCurEntryDebugInfo(const StorageEntry *curEntry) { - for (auto &buffer : curEntry->inplaceBuffers) { - if (buffer.getDefiningOp()) { - if (auto allocOp = dyn_cast(buffer.getDefiningOp())) { - LDBG("buffer : "); - LDBG(allocOp.getResult()); - } - } - } + for (auto &buffer : curEntry->inplaceBuffers) + LDBG("buffer : " << buffer); } StorageEntry * @@ -2623,8 +2539,6 @@ struct PlanMemoryPass : public mlir::pto::impl::PlanMemoryBase { RewritePatternSet &patterns, MemPlanMode mode, DenseMap> buffer2Offsets) { if (mode == MemPlanMode::LOCAL_MEM_PLAN) { - patterns.add(patterns.getContext(), - buffer2Offsets); patterns.add( patterns.getContext(), buffer2Offsets); patterns.add( diff --git a/lib/PTO/Transforms/PTOPlanMemory.h b/lib/PTO/Transforms/PTOPlanMemory.h index d6e2a37381..65bcdc6317 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.h +++ b/lib/PTO/Transforms/PTOPlanMemory.h @@ -17,7 +17,6 @@ #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Analysis/Liveness.h" #include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "llvm/ADT/SmallSet.h" #include @@ -383,13 +382,7 @@ class MemLivenessAnalysis { /// Determine whether the current operation can be skipped. bool isSkippableOp(Operation *op) const; - /// Update store op information. - void UpdateStoreOpInfo(OpInfo *opInfo, const Value storeValue, Liveness live); - - /// Check if it is local buffer with memory space - LogicalResult CheckLocalBufferAllocOp(Operation *op) const; - - /// kill buffer handle. + /// Kill a buffer handle at the current operation. void OpKillHandle(OpInfo *opInfo, Liveness live, Block *block); /// Process kill buffer based on the result live of op. @@ -406,8 +399,7 @@ class MemLivenessAnalysis { void GenerateBufferLife(); /// initialize the buffers that must be inplaced together - /// namely, the alias buffers of memref.alloc, - /// e.g. for iter arg and for yield. + /// namely, tile buffer aliases such as iter args and yields. void InitializeInplacePairList(); /// Record semantic non-reuse pairs for buffers that may be used @@ -511,10 +503,10 @@ class MemPlan { /// Post-plan sanity check for local memory overflow. bool RecordOverflowIfAny(); - /// Prepare the memref.alloc plan. + /// Prepare the local tile buffer plan. PlanStatus PlanLocalMemAddress(); - /// Prepare the memrefExt.alloc_workspace plan. + /// Prepare the global workspace plan. PlanStatus PlanWorkSpaceMemAddress(); /// merge all storage entry to the first storage entry for WorkSpaceArg. @@ -699,7 +691,7 @@ class MemPlan { /// Report tensor life time for debug. void MemLifeDebugInfo(StorageEntry *storageEntry); - /// Report tensor which is defined by memref allco. + /// Report the allocation root represented by this entry. void ReportCurEntryDebugInfo(const StorageEntry *curEntry); /// Report tensor allocate info. @@ -733,7 +725,7 @@ class MemPlan { /// Whether to adopt a split strategy. bool splitOutline{false}; - /// map from memref buffer to plan memory address. + /// Map from local allocation root to planned addresses. DenseMap> buffer2Offsets; /// map from each scope to its root StorageEntry. diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index d1eb6de846..068096d6d7 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -8,12 +8,10 @@ //===- PTOPlanMemoryModern.cpp - modern static local memory planner -------===// -#include "AllocToPointerCast.h" #include "PTO/IR/PTOMultiBuffer.h" #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/Passes.h" #include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Interfaces/SideEffectInterfaces.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" @@ -93,11 +91,6 @@ static std::optional getBufferAddressSpace(Type type) { if (auto multiType = dyn_cast(type)) return getBufferAddressSpace(multiType.getSlotType()); - if (auto memRefType = dyn_cast(type)) { - if (auto attr = - dyn_cast_or_null(memRefType.getMemorySpace())) - return attr.getAddressSpace(); - } return std::nullopt; } @@ -192,9 +185,6 @@ static FailureOr computeStaticBufferBytes(Value value) { } else if (auto multiType = dyn_cast(value.getType())) { shape = multiType.getSlotType().getShape(); elementType = multiType.getSlotType().getElementType(); - } else if (auto memRefType = dyn_cast(value.getType())) { - shape = memRefType.getShape(); - elementType = memRefType.getElementType(); } else { return failure(); } @@ -672,15 +662,6 @@ struct PlannerAnalysis { roots[found->second].freeIndex = index; } } - } else if (auto alloc = dyn_cast(op)) { - addRoot(alloc.getResult(), op); - if (failed) - return; - auto found = rootIndexByValue.find(alloc.getResult()); - if (found != rootIndexByValue.end()) { - roots[found->second].allocIndex = index; - roots[found->second].freeIndex = index; - } } if (auto multiGet = dyn_cast(op)) { @@ -694,14 +675,6 @@ struct PlannerAnalysis { propagateSplitTpopDerived(select.getResult(), ValueRange{select.getTrueValue(), select.getFalseValue()}); - } else if (auto castOp = dyn_cast(op)) { - setRoots(castOp.getResult(), getRoots(castOp.getSource())); - propagateSplitTpopDerived(castOp.getResult(), - ValueRange{castOp.getSource()}); - } else if (auto subview = dyn_cast(op)) { - setRoots(subview.getResult(), getRoots(subview.getSource())); - propagateSplitTpopDerived(subview.getResult(), - ValueRange{subview.getSource()}); } else if (auto subview = dyn_cast(op)) { setRoots(subview.getResult(), getRoots(subview.getSource())); propagateSplitTpopDerived(subview.getResult(), @@ -714,10 +687,6 @@ struct PlannerAnalysis { setRoots(reshape.getResult(), getRoots(reshape.getSrc())); propagateSplitTpopDerived(reshape.getResult(), ValueRange{reshape.getSrc()}); - } else if (auto reinterpret = dyn_cast(op)) { - setRoots(reinterpret.getResult(), getRoots(reinterpret.getSource())); - propagateSplitTpopDerived(reinterpret.getResult(), - ValueRange{reinterpret.getSource()}); } else if (auto forOp = dyn_cast(op)) { seedForIterArgAliases(forOp); } @@ -1026,8 +995,6 @@ class AllocMultiTileOpAddPlannedAddressesPattern static LogicalResult materializePlannedOffsets( func::FuncOp func, DenseMap> buffer2Offsets) { RewritePatternSet patterns(func.getContext()); - patterns.add(patterns.getContext(), - buffer2Offsets); patterns.add(patterns.getContext(), buffer2Offsets); patterns.add( @@ -1197,7 +1164,6 @@ struct PlanMemoryModernPass void getDependentDialects(DialectRegistry ®istry) const override { registry.insert(); - registry.insert(); registry.insert(); registry.insert(); registry.insert(); diff --git a/test/lit/pto/alloc_tile_physical_overlap_sync.pto b/test/lit/pto/alloc_tile_physical_overlap_sync.pto index 8e6cc457d9..7e5ff662d0 100644 --- a/test/lit/pto/alloc_tile_physical_overlap_sync.pto +++ b/test/lit/pto/alloc_tile_physical_overlap_sync.pto @@ -14,9 +14,9 @@ module attributes {pto.target_arch = "a2a3"} { func.func @alloc_tile_physical_overlap_sync( - %src0: memref<16x128xf32, #pto.address_space>, - %src1: memref<16x128xf32, #pto.address_space>, - %dst: memref<16x128xf32, #pto.address_space>) + %src0: !pto.partition_tensor_view<16x128xf32>, + %src1: !pto.partition_tensor_view<16x128xf32>, + %dst: !pto.partition_tensor_view<16x128xf32>) attributes {pto.kernel_kind = #pto.kernel_kind} { %c1280_i64 = arith.constant 1280 : i64 %c8192_i64 = arith.constant 8192 : i64 @@ -30,14 +30,14 @@ module attributes {pto.target_arch = "a2a3"} { %reused = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf - pto.tload ins(%src0 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src0 : !pto.partition_tensor_view<16x128xf32>) outs(%input : !pto.tile_buf) - pto.tload ins(%src1 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src1 : !pto.partition_tensor_view<16x128xf32>) outs(%stored : !pto.tile_buf) pto.tmuls ins(%input, %one : !pto.tile_buf, f32) outs(%stored : !pto.tile_buf) pto.tstore ins(%stored : !pto.tile_buf) - outs(%dst : memref<16x128xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<16x128xf32>) pto.tmuls ins(%input, %one : !pto.tile_buf, f32) outs(%reused : !pto.tile_buf) return diff --git a/test/lit/pto/bias_tile_subview_emitc_pointer_type.pto b/test/lit/pto/bias_tile_subview_emitc_pointer_type.pto deleted file mode 100644 index f381d7ce46..0000000000 --- a/test/lit/pto/bias_tile_subview_emitc_pointer_type.pto +++ /dev/null @@ -1,25 +0,0 @@ -// RUN: ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s - -module { - func.func @bias_tile_subview_emitc_pointer_type() - attributes {pto.kernel_kind = #pto.kernel_kind} { - %c32 = arith.constant 32 : index - - %bias = memref.alloc() : memref<1x64xf32, #pto.address_space> - %bias_sub = memref.subview %bias[0, 32] [1, 32] [1, 1] - : memref<1x64xf32, #pto.address_space> - to memref<1x32xf32, strided<[64, 1], offset: 32>, #pto.address_space> - call @consume_bias(%bias_sub) : (memref<1x32xf32, strided<[64, 1], offset: 32>, #pto.address_space>) -> () - return - } - - func.func private @consume_bias(%arg0: memref<1x32xf32, strided<[64, 1], offset: 32>, #pto.address_space>) -} - -// CHECK-LABEL: AICORE void bias_tile_subview_emitc_pointer_type() -// CHECK-DAG: const int64_t [[ONE:[_A-Za-z][_A-Za-z0-9]*]] = 1; -// CHECK-DAG: const int64_t [[THIRTYTWO:[_A-Za-z][_A-Za-z0-9]*]] = 32; -// CHECK: __gm__ float* [[BIAS_BASE:[_A-Za-z][_A-Za-z0-9]*]] = reinterpret_cast<__gm__ float*>({{[_A-Za-z][_A-Za-z0-9]*}}); -// CHECK: __gm__ float* [[BIAS_SUB:[_A-Za-z][_A-Za-z0-9]*]] = {{.*}}[[BIAS_BASE]]{{.*}}[[THIRTYTWO]] * [[ONE]]{{.*}}; -// CHECK: consume_bias([[BIAS_SUB]]); -// CHECK-NOT: __gm__ float** diff --git a/test/lit/pto/declare_tile_tile_native.pto b/test/lit/pto/declare_tile_tile_native.pto index bd2e8a6dab..b365375ee6 100644 --- a/test/lit/pto/declare_tile_tile_native.pto +++ b/test/lit/pto/declare_tile_tile_native.pto @@ -9,7 +9,7 @@ // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-before=pto-inline-backend-helpers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NATIVE // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC module { diff --git a/test/lit/pto/easy_param_completion_emitc.pto b/test/lit/pto/easy_param_completion_emitc.pto index 95c3cde76a..53636240bb 100644 --- a/test/lit/pto/easy_param_completion_emitc.pto +++ b/test/lit/pto/easy_param_completion_emitc.pto @@ -1,14 +1,14 @@ // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s --check-prefix=A3 module { - func.func @tci_with_tmp(%dst: memref<1x16xi16, #pto.address_space>) { + 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) outs(%dst_tile : !pto.tile_buf) pto.tstore ins(%dst_tile : !pto.tile_buf) - outs(%dst : memref<1x16xi16, #pto.address_space>) {layout = #pto.layout, pto.inferred_layout = true} + outs(%dst : !pto.partition_tensor_view<1x16xi16>) {layout = #pto.layout, pto.inferred_layout = true} return } @@ -23,21 +23,21 @@ module { } func.func @mgather_with_explicit_attrs_a3( - %mem: memref<1x1x1x8x16xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x8x16xf16>) { %idx = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.mgather ins(%mem, %idx : memref<1x1x1x8x16xf16, #pto.address_space>, !pto.tile_buf) + pto.mgather ins(%mem, %idx : !pto.partition_tensor_view<1x1x1x8x16xf16>, !pto.tile_buf) outs(%dst : !pto.tile_buf) {coalesce = #pto, gatherOob = #pto} return } func.func @mscatter_with_explicit_attrs_a3( - %mem: memref<1x1x1x8x16xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x8x16xf16>) { %src = pto.alloc_tile : !pto.tile_buf %idx = pto.alloc_tile : !pto.tile_buf pto.mscatter ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) - outs(%mem : memref<1x1x1x8x16xf16, #pto.address_space>) + outs(%mem : !pto.partition_tensor_view<1x1x1x8x16xf16>) {coalesce = #pto, scatterAtomicOp = #pto, scatterOob = #pto} diff --git a/test/lit/pto/graph_sync_solver_a3_dependent_mmad_barrier.pto b/test/lit/pto/graph_sync_solver_a3_dependent_mmad_barrier.pto index de05f06c47..aa061bf18a 100644 --- a/test/lit/pto/graph_sync_solver_a3_dependent_mmad_barrier.pto +++ b/test/lit/pto/graph_sync_solver_a3_dependent_mmad_barrier.pto @@ -10,11 +10,11 @@ module { // A5: pto.tmatmul // A5-NEXT: pto.tmatmul.acc func.func @dependent_mmad_same_acc( - %lhs: memref<32x32xf16, #pto.address_space>, - %rhs: memref<32x32xf16, #pto.address_space>, - %acc: memref<32x32xf32, #pto.address_space>) { - pto.tmatmul ins(%lhs, %rhs : memref<32x32xf16, #pto.address_space>, memref<32x32xf16, #pto.address_space>) outs(%acc : memref<32x32xf32, #pto.address_space>) - pto.tmatmul.acc ins(%acc, %lhs, %rhs : memref<32x32xf32, #pto.address_space>, memref<32x32xf16, #pto.address_space>, memref<32x32xf16, #pto.address_space>) outs(%acc : memref<32x32xf32, #pto.address_space>) + %lhs: !pto.tile_buf, + %rhs: !pto.tile_buf, + %acc: !pto.tile_buf) { + pto.tmatmul ins(%lhs, %rhs : !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) + pto.tmatmul.acc ins(%acc, %lhs, %rhs : !pto.tile_buf, !pto.tile_buf, !pto.tile_buf) outs(%acc : !pto.tile_buf) return } } diff --git a/test/lit/pto/graph_sync_solver_section_transparent.pto b/test/lit/pto/graph_sync_solver_section_transparent.pto index 91900aa6e2..9f95ef9804 100644 --- a/test/lit/pto/graph_sync_solver_section_transparent.pto +++ b/test/lit/pto/graph_sync_solver_section_transparent.pto @@ -10,30 +10,31 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK: pto.set_flag[, // CHECK: pto.wait_flag[, // CHECK: pto.tinsert - func.func @sectioned_cube_pipeline(%a: memref<32x32xf32, #pto.address_space>, - %b: memref<32x32xf32, #pto.address_space>) { + func.func @sectioned_cube_pipeline( + %a: !pto.partition_tensor_view<32x32xf32>, + %b: !pto.partition_tensor_view<32x32xf32>) { pto.section.cube { %c0 = arith.constant 0 : index - %a_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %a_left = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_right = memref.alloc() : memref<32x32xf32, #pto.address_space> - %src_acc = memref.alloc() : memref<32x32xf32, #pto.address_space> - %dst_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> + %a_mat = pto.alloc_tile : !pto.tile_buf + %b_mat = pto.alloc_tile : !pto.tile_buf + %a_left = pto.alloc_tile : !pto.tile_buf + %b_right = pto.alloc_tile : !pto.tile_buf + %src_acc = pto.alloc_tile : !pto.tile_buf + %dst_mat = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%a : memref<32x32xf32, #pto.address_space>) - outs(%a_mat : memref<32x32xf32, #pto.address_space>) - pto.tload ins(%b : memref<32x32xf32, #pto.address_space>) - outs(%b_mat : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%a_mat : memref<32x32xf32, #pto.address_space>) - outs(%a_left : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%b_mat : memref<32x32xf32, #pto.address_space>) - outs(%b_right : memref<32x32xf32, #pto.address_space>) - pto.tmatmul ins(%a_left, %b_right : memref<32x32xf32, #pto.address_space>, - memref<32x32xf32, #pto.address_space>) - outs(%src_acc : memref<32x32xf32, #pto.address_space>) - pto.tinsert ins(%src_acc, %c0, %c0 : memref<32x32xf32, #pto.address_space>, index, index) - outs(%dst_mat : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%a : !pto.partition_tensor_view<32x32xf32>) + outs(%a_mat : !pto.tile_buf) + pto.tload ins(%b : !pto.partition_tensor_view<32x32xf32>) + outs(%b_mat : !pto.tile_buf) + pto.tmov ins(%a_mat : !pto.tile_buf) + outs(%a_left : !pto.tile_buf) + pto.tmov ins(%b_mat : !pto.tile_buf) + outs(%b_right : !pto.tile_buf) + pto.tmatmul ins(%a_left, %b_right : !pto.tile_buf, + !pto.tile_buf) + outs(%src_acc : !pto.tile_buf) + pto.tinsert ins(%src_acc, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst_mat : !pto.tile_buf) } return } diff --git a/test/lit/pto/inject_barrier_all_sync_tpush_tpop.pto b/test/lit/pto/inject_barrier_all_sync_tpush_tpop.pto index 7e527aad35..6f9d66d77f 100644 --- a/test/lit/pto/inject_barrier_all_sync_tpush_tpop.pto +++ b/test/lit/pto/inject_barrier_all_sync_tpush_tpop.pto @@ -2,13 +2,13 @@ // RUN: ptoas --pto-arch=a3 --enable-inject-barrier-all-sync %s 2>&1 | FileCheck %s --check-prefix=CPP module { - func.func @regular_pipe_memory_ops(%src: memref<16x16xf32, #pto.address_space>) attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @regular_pipe_memory_ops(%src: !pto.partition_tensor_view<16x16xf32>) attributes {pto.kernel_kind = #pto.kernel_kind} { // CHECK-LABEL: func.func @regular_pipe_memory_ops %mat = pto.alloc_tile : !pto.tile_buf %left = pto.alloc_tile : !pto.tile_buf // CHECK: pto.barrier // CHECK-NEXT: pto.tload - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%mat : !pto.tile_buf) // CHECK: pto.barrier // CHECK-NEXT: pto.tmov @@ -19,13 +19,13 @@ module { return } - func.func @manual_pipe_all_barrier_not_duplicated(%src: memref<16x16xf32, #pto.address_space>) attributes {pto.kernel_kind = #pto.kernel_kind} { + func.func @manual_pipe_all_barrier_not_duplicated(%src: !pto.partition_tensor_view<16x16xf32>) attributes {pto.kernel_kind = #pto.kernel_kind} { // CHECK-LABEL: func.func @manual_pipe_all_barrier_not_duplicated // CHECK: pto.barrier // CHECK-NEXT: pto.tload %mat = pto.alloc_tile : !pto.tile_buf pto.barrier - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%mat : !pto.tile_buf) // CHECK: pto.barrier // CHECK-NEXT: return diff --git a/test/lit/pto/issue157_64bit_view_offset_emitc.pto b/test/lit/pto/issue157_64bit_view_offset_emitc.pto index b749376013..e59e4f03be 100644 --- a/test/lit/pto/issue157_64bit_view_offset_emitc.pto +++ b/test/lit/pto/issue157_64bit_view_offset_emitc.pto @@ -9,19 +9,6 @@ // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s --check-prefix=A3 module { - func.func @issue157_subview_remote_offset_i8( - %src: memref>, %off: index) { - %c64 = arith.constant 64 : index - %sub = memref.subview %src[%off] [%c64] [1] - : memref> - to memref, #pto.address_space> - %tile = pto.alloc_tile - : !pto.tile_buf - pto.tload ins(%sub : memref, #pto.address_space>) - outs(%tile : !pto.tile_buf) - return - } - func.func @issue157_partition_remote_offset_i8(%src: !pto.ptr, %off: index) { %c1 = arith.constant 1 : index %c64 = arith.constant 64 : index @@ -38,18 +25,12 @@ module { } } -// A3-LABEL: AICORE void issue157_subview_remote_offset_i8 -// A3-SAME: (__gm__ int8_t* [[SUB_SRC:[_A-Za-z][_A-Za-z0-9]*]], int64_t [[SUB_OFF:[_A-Za-z][_A-Za-z0-9]*]]) -// A3-DAG: const int64_t [[SUB_ONE:[_A-Za-z][_A-Za-z0-9]*]] = 1; -// A3-DAG: const int64_t [[SUB_ZERO:[_A-Za-z][_A-Za-z0-9]*]] = 0; -// A3-NOT: unsigned -// A3: GlobalTensor>, - %src1: memref<16x128xf32, #pto.address_space>, - %src2: memref<16x128xf32, #pto.address_space>, - %src3: memref<16x128xf32, #pto.address_space>) { + %src0: !pto.partition_tensor_view<16x128xf32>, + %src1: !pto.partition_tensor_view<16x128xf32>, + %src2: !pto.partition_tensor_view<16x128xf32>, + %src3: !pto.partition_tensor_view<16x128xf32>) { %c0 = arith.constant 0 : index %c128 = arith.constant 128 : index %c256 = arith.constant 256 : index @@ -28,13 +28,13 @@ module { !pto.tile_buf -> !pto.tile_buf - pto.tload ins(%src0 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src0 : !pto.partition_tensor_view<16x128xf32>) outs(%s0 : !pto.tile_buf) - pto.tload ins(%src1 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src1 : !pto.partition_tensor_view<16x128xf32>) outs(%s1 : !pto.tile_buf) - pto.tload ins(%src2 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src2 : !pto.partition_tensor_view<16x128xf32>) outs(%s2 : !pto.tile_buf) - pto.tload ins(%src3 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src3 : !pto.partition_tensor_view<16x128xf32>) outs(%s3 : !pto.tile_buf) return } diff --git a/test/lit/pto/issue667_tload_overlap_no_mte2_barrier.pto b/test/lit/pto/issue667_tload_overlap_no_mte2_barrier.pto index fb51009558..324fd1c3f9 100644 --- a/test/lit/pto/issue667_tload_overlap_no_mte2_barrier.pto +++ b/test/lit/pto/issue667_tload_overlap_no_mte2_barrier.pto @@ -2,8 +2,8 @@ module { func.func @issue667_tload_overlap_does_not_need_mte2_barrier( - %src0: memref<16x128xf32, #pto.address_space>, - %src1: memref<16x128xf32, #pto.address_space>) { + %src0: !pto.partition_tensor_view<16x128xf32>, + %src1: !pto.partition_tensor_view<16x128xf32>) { %c0 = arith.constant 0 : index %c64 = arith.constant 64 : index %c0_i64 = arith.constant 0 : i64 @@ -18,9 +18,9 @@ module { !pto.tile_buf -> !pto.tile_buf - pto.tload ins(%src0 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src0 : !pto.partition_tensor_view<16x128xf32>) outs(%s0 : !pto.tile_buf) - pto.tload ins(%src1 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src1 : !pto.partition_tensor_view<16x128xf32>) outs(%s1 : !pto.tile_buf) return } diff --git a/test/lit/pto/issue708_zero_valid_subview.pto b/test/lit/pto/issue708_zero_valid_subview.pto index b4c7be1bc1..6b3794c29a 100644 --- a/test/lit/pto/issue708_zero_valid_subview.pto +++ b/test/lit/pto/issue708_zero_valid_subview.pto @@ -10,7 +10,7 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { func.func @issue708_zero_valid_subview_row( - %dst: memref<5x128xf32, #pto.address_space>) + %dst: !pto.partition_tensor_view<5x128xf32>) attributes {pto.kernel_kind = #pto.kernel_kind} { %c0 = arith.constant 0 : index %c128 = arith.constant 128 : index @@ -22,12 +22,12 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { -> !pto.tile_buf pto.tstore ins(%tail : !pto.tile_buf) - outs(%dst : memref<5x128xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<5x128xf32>) return } func.func @issue708_zero_valid_subview_col( - %dst: memref<5x128xf32, #pto.address_space>) + %dst: !pto.partition_tensor_view<5x128xf32>) attributes {pto.kernel_kind = #pto.kernel_kind} { %c0 = arith.constant 0 : index %c5 = arith.constant 5 : index @@ -40,7 +40,7 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { -> !pto.tile_buf pto.tstore ins(%tail : !pto.tile_buf) - outs(%dst : memref<5x128xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<5x128xf32>) return } } diff --git a/test/lit/pto/issue708_zero_valid_subview_inferred.pto b/test/lit/pto/issue708_zero_valid_subview_inferred.pto index 8ba7d64b79..736d754ad5 100644 --- a/test/lit/pto/issue708_zero_valid_subview_inferred.pto +++ b/test/lit/pto/issue708_zero_valid_subview_inferred.pto @@ -19,7 +19,7 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { // Both axes empty: trims a 16-row Q_HEAD_PAD tile to a 5x128 view whose valid // region is fully empty on the replay lane, then stores it. func.func @issue708_inferred_empty_store( - %dst: memref<5x128xf32, #pto.address_space>) + %dst: !pto.partition_tensor_view<5x128xf32>) attributes {pto.kernel_kind = #pto.kernel_kind} { %c0 = arith.constant 0 : index @@ -30,7 +30,7 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { -> !pto.tile_buf pto.tstore ins(%tail : !pto.tile_buf) - outs(%dst : memref<5x128xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<5x128xf32>) return } @@ -56,7 +56,7 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { // full valid extent (v_col=128). Lowering must thread 0 for the row and 128 // for the column -- not widen the row, not zero the column. func.func @issue708_inferred_row_empty( - %dst: memref<5x128xf32, #pto.address_space>) + %dst: !pto.partition_tensor_view<5x128xf32>) attributes {pto.kernel_kind = #pto.kernel_kind} { %c0 = arith.constant 0 : index @@ -67,7 +67,7 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { -> !pto.tile_buf pto.tstore ins(%tail : !pto.tile_buf) - outs(%dst : memref<5x128xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<5x128xf32>) return } @@ -75,7 +75,7 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { // non-size static valid (v_row=3). The result type is authoritative, so // lowering threads 3 (not the subview size 5) into the tile handle. func.func @issue708_inferred_partial_valid( - %dst: memref<5x128xf32, #pto.address_space>) + %dst: !pto.partition_tensor_view<3x128xf32>) attributes {pto.kernel_kind = #pto.kernel_kind} { %c0 = arith.constant 0 : index @@ -86,7 +86,7 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { -> !pto.tile_buf pto.tstore ins(%tail : !pto.tile_buf) - outs(%dst : memref<5x128xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<3x128xf32>) return } } diff --git a/test/lit/pto/issue870_identity_tmov_pruning.pto b/test/lit/pto/issue870_identity_tmov_pruning.pto index 2b75517e9c..69bcb8172a 100644 --- a/test/lit/pto/issue870_identity_tmov_pruning.pto +++ b/test/lit/pto/issue870_identity_tmov_pruning.pto @@ -23,13 +23,13 @@ module attributes {pto.target_arch = "a2a3", pto.kernel_kind = #pto.kernel_kind< func.func @live_dst_same_addr_keeps_tmov() { %c0_i64 = arith.constant 0 : i64 %c64_i64 = arith.constant 64 : i64 - %src = pto.pointer_cast(%c0_i64) : memref<1x16xf32, #pto.address_space> - %dst = pto.pointer_cast(%c0_i64) : memref<1x16xf32, #pto.address_space> - %out = pto.pointer_cast(%c64_i64) : memref<1x16xf32, #pto.address_space> - pto.tmov ins(%src : memref<1x16xf32, #pto.address_space>) - outs(%dst : memref<1x16xf32, #pto.address_space>) - pto.tmov ins(%dst : memref<1x16xf32, #pto.address_space>) - outs(%out : memref<1x16xf32, #pto.address_space>) + %src = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %dst = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %out = pto.alloc_tile addr = %c64_i64 : !pto.tile_buf + pto.tmov ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tmov ins(%dst : !pto.tile_buf) + outs(%out : !pto.tile_buf) return } @@ -50,10 +50,10 @@ module attributes {pto.target_arch = "a2a3", pto.kernel_kind = #pto.kernel_kind< %c0_i64 = arith.constant 0 : i64 %c0_idx = arith.constant 0 : index %c0_alt = arith.index_cast %c0_idx : index to i64 - %src = pto.pointer_cast(%c0_i64) : memref<1x16xf32, #pto.address_space> - %dst = pto.pointer_cast(%c0_alt) : memref<1x16xf32, #pto.address_space> - pto.tmov ins(%src : memref<1x16xf32, #pto.address_space>) - outs(%dst : memref<1x16xf32, #pto.address_space>) + %src = pto.alloc_tile addr = %c0_i64 : !pto.tile_buf + %dst = pto.alloc_tile addr = %c0_alt : !pto.tile_buf + pto.tmov ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } @@ -87,10 +87,10 @@ module attributes {pto.target_arch = "a2a3", pto.kernel_kind = #pto.kernel_kind< } func.func @non_identity_tmov_dynamic_addr(%addr0: i64, %addr1: i64) { - %src = pto.pointer_cast(%addr0) : memref<1x16xf32, #pto.address_space> - %dst = pto.pointer_cast(%addr1) : memref<1x16xf32, #pto.address_space> - pto.tmov ins(%src : memref<1x16xf32, #pto.address_space>) - outs(%dst : memref<1x16xf32, #pto.address_space>) + %src = pto.alloc_tile addr = %addr0 : !pto.tile_buf + %dst = pto.alloc_tile addr = %addr1 : !pto.tile_buf + pto.tmov ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } } diff --git a/test/lit/pto/low_precision_globaltensor_emitc_a5.pto b/test/lit/pto/low_precision_globaltensor_emitc_a5.pto index 7bfbf6c742..f7b691ad4d 100644 --- a/test/lit/pto/low_precision_globaltensor_emitc_a5.pto +++ b/test/lit/pto/low_precision_globaltensor_emitc_a5.pto @@ -10,19 +10,19 @@ module { func.func @low_precision_globaltensor_emitc_a5( - %src_hif8: memref<16x32x!pto.hif8, #pto.address_space>, - %src_f8: memref<16x32xf8E4M3FN, #pto.address_space>, - %dst_f4: memref<16x32x!pto.f4E1M2x2, #pto.address_space>) { + %src_hif8: !pto.partition_tensor_view<16x32x!pto.hif8>, + %src_f8: !pto.partition_tensor_view<16x32xf8E4M3FN>, + %dst_f4: !pto.partition_tensor_view<16x32x!pto.f4E1M2x2>) { %tile_hif8 = pto.alloc_tile : !pto.tile_buf %tile_f8 = pto.alloc_tile : !pto.tile_buf %tile_f4 = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src_hif8 : memref<16x32x!pto.hif8, #pto.address_space>) + pto.tload ins(%src_hif8 : !pto.partition_tensor_view<16x32x!pto.hif8>) outs(%tile_hif8 : !pto.tile_buf) - pto.tprefetch ins(%src_f8 : memref<16x32xf8E4M3FN, #pto.address_space>) + pto.tprefetch ins(%src_f8 : !pto.partition_tensor_view<16x32xf8E4M3FN>) outs(%tile_f8 : !pto.tile_buf) pto.tstore ins(%tile_f4 : !pto.tile_buf) - outs(%dst_f4 : memref<16x32x!pto.f4E1M2x2, #pto.address_space>) + outs(%dst_f4 : !pto.partition_tensor_view<16x32x!pto.f4E1M2x2>) return } } diff --git a/test/lit/pto/materialize_tile_handles_control_flow_result.pto b/test/lit/pto/materialize_tile_handles_control_flow_result.pto index 8987f119f3..bbfef684dc 100644 --- a/test/lit/pto/materialize_tile_handles_control_flow_result.pto +++ b/test/lit/pto/materialize_tile_handles_control_flow_result.pto @@ -1,5 +1,5 @@ // RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-before=pto-inline-backend-helpers %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO @@ -51,7 +51,7 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { } // NATIVE: IR Dump After PTOResolveReservedBuffers -// MATERIALIZED: IR Dump After PTOMaterializeTileHandles +// MATERIALIZED: IR Dump Before PTOInlineBackendHelpers // COMMON: %[[SEL:.*]] = scf.if {{.*}} -> (!pto.tile_buf) { // COMMON: pto.tadd // COMMON: scf.yield %{{.*}} : !pto.tile_buf diff --git a/test/lit/pto/materialize_tile_handles_emitc.pto b/test/lit/pto/materialize_tile_handles_emitc.pto index 58cce54f14..e9a2ce5606 100644 --- a/test/lit/pto/materialize_tile_handles_emitc.pto +++ b/test/lit/pto/materialize_tile_handles_emitc.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-before=pto-inline-backend-helpers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC module { @@ -13,7 +13,7 @@ module { } } -// IR: IR Dump After PTOMaterializeTileHandles +// IR: IR Dump Before PTOInlineBackendHelpers // IR: pto.alloc_tile addr = {{.*}} : !pto.tile_buf // IR: pto.alloc_tile addr = {{.*}} : !pto.tile_buf diff --git a/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto b/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto index ae7bafd536..9e311eab4c 100644 --- a/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto +++ b/test/lit/pto/materialize_tile_handles_subkernel_helper_abi.pto @@ -9,7 +9,7 @@ // Guards PTODSL subkernel helper arguments staying in tile_buf ABI when the // helper directly extracts their address views. // RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON -// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-before=pto-inline-backend-helpers %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON // RUN: ptoas --pto-arch=a5 --pto-level=level3 %s -o /dev/null module { @@ -37,7 +37,7 @@ module { } // NATIVE: IR Dump After PTOResolveReservedBuffers -// MATERIALIZED: IR Dump After PTOMaterializeTileHandles +// MATERIALIZED: IR Dump Before PTOInlineBackendHelpers // COMMON-LABEL: func.func @ptodsl_subkernel_helper_tile_abi( // COMMON-SAME: %{{.*}}: !pto.tile_buf // COMMON-SAME: %{{.*}}: !pto.tile_buf diff --git a/test/lit/pto/mgather_gm2ub_a5_elem_1x1_memref_macro_sync_model.pto b/test/lit/pto/mgather_gm2ub_a5_elem_1x1_memref_macro_sync_model.pto index ba146e0b64..57aed3750f 100644 --- a/test/lit/pto/mgather_gm2ub_a5_elem_1x1_memref_macro_sync_model.pto +++ b/test/lit/pto/mgather_gm2ub_a5_elem_1x1_memref_macro_sync_model.pto @@ -7,20 +7,18 @@ // See LICENSE in the root of the software repository for the full text of the License. // A5 elem-coalesced 1x1 GM->UB mgather lowers to the scalar overload even when -// InsertSync sees the dst after PTOResolveReservedBuffers lowering. The macro model must -// recover the 1x1 shape from MemRefType instead of requiring TileBufType. +// InsertSync sees the tile-native dst after PTOResolveReservedBuffers. // // RUN: ptoas --pto-arch=a5 --enable-insert-sync --pto-insert-sync-debug=2 %s -o - 2>&1 | FileCheck %s module attributes {pto.target_arch = "a5"} { func.func @mgather_elem_1x1_memref_model( - %mem: memref<1xf32, #pto.address_space>, - %idx: memref<1x1xi32, #pto.address_space>) { - %dst = memref.alloc() : memref<1x1xf32, #pto.address_space> + %mem: !pto.partition_tensor_view<1xf32>, + %idx: !pto.tile_buf) { + %dst = pto.alloc_tile : !pto.tile_buf pto.mgather ins(%mem, %idx - : memref<1xf32, #pto.address_space>, - memref<1x1xi32, #pto.address_space>) - outs(%dst : memref<1x1xf32, #pto.address_space>) + : !pto.partition_tensor_view<1xf32>, !pto.tile_buf) + outs(%dst : !pto.tile_buf) {coalesce = #pto} return } diff --git a/test/lit/pto/multi_tile_affine_disjoint_slots.pto b/test/lit/pto/multi_tile_affine_disjoint_slots.pto index 03806f6688..15c3e28a6c 100644 --- a/test/lit/pto/multi_tile_affine_disjoint_slots.pto +++ b/test/lit/pto/multi_tile_affine_disjoint_slots.pto @@ -23,8 +23,8 @@ module { func.func @prefetch_disjoint_slots( - %gm : memref<16x16xf16, #pto.address_space>, - %dst : memref<16x16xf16, #pto.address_space>, + %gm : !pto.partition_tensor_view<16x16xf16>, + %dst : !pto.partition_tensor_view<16x16xf16>, %n : index) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -42,7 +42,7 @@ module { %prod = pto.multi_tile_get %mb[%nxt_idx] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%prod : !pto.tile_buf) // Consumer (MTE3): tstore reads slot iv%2. @@ -50,7 +50,7 @@ module { : !pto.multi_tile_buf -> !pto.tile_buf pto.tstore ins(%cons : !pto.tile_buf) - outs(%dst : memref<16x16xf16, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<16x16xf16>) } return } diff --git a/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto b/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto index 7773a855d9..e13b59271f 100644 --- a/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto +++ b/test/lit/pto/multi_tile_buf_n3_planmem_e2e.pto @@ -13,7 +13,7 @@ // physical slots and `UpdateBuffer2Offsets` emits them in slot order. module { - func.func @n3_three_slots(%gm : memref<16x16xf16, #pto.address_space>) { + func.func @n3_three_slots(%gm : !pto.partition_tensor_view<16x16xf16>) { %mb = pto.alloc_multi_tile : !pto.multi_tile_buf %c0 = arith.constant 0 : index @@ -28,11 +28,11 @@ module { %s2 = pto.multi_tile_get %mb[%c2] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s0 : !pto.tile_buf) - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s1 : !pto.tile_buf) - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s2 : !pto.tile_buf) return } diff --git a/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto b/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto index 7ed312799d..4a09d00ba2 100644 --- a/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto +++ b/test/lit/pto/multi_tile_const_preload_dyn_loop_select.pto @@ -8,7 +8,7 @@ // RUN: ptoas --mlir-print-ir-after=pto-resolve-reserved-buffers %s 2>&1 1>/dev/null | FileCheck %s // RUN: ptoas --mlir-print-ir-after=pto-resolve-buffer-select %s 2>&1 1>/dev/null | FileCheck %s --check-prefix=SELECT -// RUN: ptoas --pto-arch=a3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=MAT +// RUN: ptoas --pto-arch=a3 --mlir-print-ir-before=pto-inline-backend-helpers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=MAT // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC // Mixed slot-select example: loop-external preload uses a constant slot id, @@ -16,10 +16,21 @@ module { func.func @const_preload_dyn_loop_select( - %gm : memref<16x16xf16, #pto.address_space>, %n : index) { + %gm : !pto.ptr, %n : index) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index + %c16 = arith.constant 16 : index + + %gm_view = pto.make_tensor_view %gm, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c16, %c16, %c16, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %gm_part = pto.partition_view %gm_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> + -> !pto.partition_tensor_view<1x1x1x16x16xf16> %mb = pto.alloc_multi_tile : !pto.multi_tile_buf @@ -28,7 +39,7 @@ module { %pre = pto.multi_tile_get %mb[%c0] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) outs(%pre : !pto.tile_buf) scf.for %i = %c0 to %n step %c1 { @@ -36,7 +47,7 @@ module { %slot = pto.multi_tile_get %mb[%idx] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) outs(%slot : !pto.tile_buf) } return @@ -63,18 +74,16 @@ module { // SELECT: arith.select // SELECT: pto.alloc_tile addr = %{{.*}} -// MaterializeTileHandles must carry the dynamic selected buffer address into -// alloc_tile, otherwise EmitC creates a tile without TASSIGN in the loop. -// MAT: IR Dump After PTOMaterializeTileHandles +// ResolveBufferSelect carries the dynamic selected buffer address directly in +// alloc_tile, so EmitC can emit TASSIGN in the loop without a memref bridge. +// MAT: IR Dump Before PTOInlineBackendHelpers // MAT-LABEL: func.func @const_preload_dyn_loop_select -// MAT-DAG: %[[ADDR0:[A-Za-z0-9_]+]] = arith.constant 0 : i64 -// MAT-DAG: %[[ADDR1:[A-Za-z0-9_]+]] = arith.constant 512 : i64 +// MAT: %[[ADDR0:[A-Za-z0-9_]+]] = arith.constant 0 : i64 // MAT: pto.alloc_tile addr = %[[ADDR0]] // MAT: scf.for -// MAT-DAG: %[[LOOP_ADDR0:[A-Za-z0-9_]+]] = arith.constant 0 : i64 -// MAT-DAG: %[[LOOP_ADDR1:[A-Za-z0-9_]+]] = arith.constant 512 : i64 +// MAT: %[[ADDR1:[A-Za-z0-9_]+]] = arith.constant 512 : i64 // MAT: %[[COND:[A-Za-z0-9_]+]] = arith.cmpi eq -// MAT: %[[DYN_ADDR:[A-Za-z0-9_]+]] = arith.select %[[COND]], %[[LOOP_ADDR1]], %[[LOOP_ADDR0]] : i64 +// MAT: %[[DYN_ADDR:[A-Za-z0-9_]+]] = arith.select %[[COND]], %[[ADDR1]], %[[ADDR0]] : i64 // MAT: pto.alloc_tile addr = %[[DYN_ADDR]] // EMITC: AICORE void const_preload_dyn_loop_select diff --git a/test/lit/pto/multi_tile_const_slot_disjoint_sync.pto b/test/lit/pto/multi_tile_const_slot_disjoint_sync.pto index 16412acec6..3c97d15503 100644 --- a/test/lit/pto/multi_tile_const_slot_disjoint_sync.pto +++ b/test/lit/pto/multi_tile_const_slot_disjoint_sync.pto @@ -18,8 +18,8 @@ module { func.func @const_slot_no_inter_sync( - %gm0 : memref<16x16xf16, #pto.address_space>, - %gm1 : memref<16x16xf16, #pto.address_space>) { + %gm0 : !pto.partition_tensor_view<16x16xf16>, + %gm1 : !pto.partition_tensor_view<16x16xf16>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -33,9 +33,9 @@ module { : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm0 : !pto.partition_tensor_view<16x16xf16>) outs(%s0 : !pto.tile_buf) - pto.tload ins(%gm1 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm1 : !pto.partition_tensor_view<16x16xf16>) outs(%s1 : !pto.tile_buf) return } diff --git a/test/lit/pto/multi_tile_const_slot_gss_compiles.pto b/test/lit/pto/multi_tile_const_slot_gss_compiles.pto index e9ba742136..a2ae682b51 100644 --- a/test/lit/pto/multi_tile_const_slot_gss_compiles.pto +++ b/test/lit/pto/multi_tile_const_slot_gss_compiles.pto @@ -17,8 +17,8 @@ module { func.func @const_slot_gss_compiles( - %gm0 : memref<16x16xf16, #pto.address_space>, - %gm1 : memref<16x16xf16, #pto.address_space>) { + %gm0 : !pto.partition_tensor_view<16x16xf16>, + %gm1 : !pto.partition_tensor_view<16x16xf16>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -32,9 +32,9 @@ module { : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm0 : !pto.partition_tensor_view<16x16xf16>) outs(%s0 : !pto.tile_buf) - pto.tload ins(%gm1 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm1 : !pto.partition_tensor_view<16x16xf16>) outs(%s1 : !pto.tile_buf) return } diff --git a/test/lit/pto/multi_tile_dyn_slot_conservative_sync.pto b/test/lit/pto/multi_tile_dyn_slot_conservative_sync.pto index ea2020dca4..0dd8e36ae0 100644 --- a/test/lit/pto/multi_tile_dyn_slot_conservative_sync.pto +++ b/test/lit/pto/multi_tile_dyn_slot_conservative_sync.pto @@ -19,7 +19,7 @@ module { func.func @dyn_slot_compiles( - %gm : memref<16x16xf16, #pto.address_space>, %n : index) { + %gm : !pto.partition_tensor_view<16x16xf16>, %n : index) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index @@ -32,7 +32,7 @@ module { %s = pto.multi_tile_get %mb[%idx] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s : !pto.tile_buf) } return diff --git a/test/lit/pto/multi_tile_get_const_slot_lowering.pto b/test/lit/pto/multi_tile_get_const_slot_lowering.pto index e96fd6a7ac..50b41802e2 100644 --- a/test/lit/pto/multi_tile_get_const_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_const_slot_lowering.pto @@ -19,8 +19,8 @@ module { func.func @const_slot_two_buffers( - %gm0 : memref<16x16xf16, #pto.address_space>, - %gm1 : memref<16x16xf16, #pto.address_space>) { + %gm0 : !pto.partition_tensor_view<16x16xf16>, + %gm1 : !pto.partition_tensor_view<16x16xf16>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -34,9 +34,9 @@ module { : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm0 : !pto.partition_tensor_view<16x16xf16>) outs(%s0 : !pto.tile_buf) - pto.tload ins(%gm1 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm1 : !pto.partition_tensor_view<16x16xf16>) outs(%s1 : !pto.tile_buf) return } diff --git a/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto b/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto index 1ca5908c3c..62dbfa1875 100644 --- a/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_dyn_slot_lowering.pto @@ -15,7 +15,7 @@ module { func.func @dyn_slot_prefetch( - %gm : memref<16x16xf16, #pto.address_space>, %n : index) { + %gm : !pto.partition_tensor_view<16x16xf16>, %n : index) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index @@ -31,13 +31,13 @@ module { %s_next = pto.multi_tile_get %mb[%next_idx] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s_next : !pto.tile_buf) %s_cur = pto.multi_tile_get %mb[%cur_idx] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s_cur : !pto.tile_buf) } return diff --git a/test/lit/pto/multi_tile_level3_explicit_addr.pto b/test/lit/pto/multi_tile_level3_explicit_addr.pto index 6b5a06a9a4..a10fd2ceab 100644 --- a/test/lit/pto/multi_tile_level3_explicit_addr.pto +++ b/test/lit/pto/multi_tile_level3_explicit_addr.pto @@ -17,8 +17,8 @@ module { func.func @l3_two_buffers( - %gm0 : memref<16x16xf16, #pto.address_space>, - %gm1 : memref<16x16xf16, #pto.address_space>) attributes {pto.entry} { + %gm0 : !pto.partition_tensor_view<16x16xf16>, + %gm1 : !pto.partition_tensor_view<16x16xf16>) attributes {pto.entry} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %base = arith.constant 0 : i64 @@ -33,9 +33,9 @@ module { : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm0 : !pto.partition_tensor_view<16x16xf16>) outs(%s0 : !pto.tile_buf) - pto.tload ins(%gm1 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm1 : !pto.partition_tensor_view<16x16xf16>) outs(%s1 : !pto.tile_buf) return } diff --git a/test/lit/pto/multi_tile_level3_missing_addr.pto b/test/lit/pto/multi_tile_level3_missing_addr.pto index f581df3ebd..9bdddb02dd 100644 --- a/test/lit/pto/multi_tile_level3_missing_addr.pto +++ b/test/lit/pto/multi_tile_level3_missing_addr.pto @@ -14,7 +14,7 @@ module { func.func @l3_missing_addr( - %gm0 : memref<16x16xf16, #pto.address_space>) attributes {pto.entry} { + %gm0 : !pto.partition_tensor_view<16x16xf16>) attributes {pto.entry} { %c0 = arith.constant 0 : index %mb = pto.alloc_multi_tile @@ -24,7 +24,7 @@ module { : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm0 : !pto.partition_tensor_view<16x16xf16>) outs(%s0 : !pto.tile_buf) return } diff --git a/test/lit/pto/multi_tile_n4_planmem_e2e.pto b/test/lit/pto/multi_tile_n4_planmem_e2e.pto index d44d97b3dd..5667c623e3 100644 --- a/test/lit/pto/multi_tile_n4_planmem_e2e.pto +++ b/test/lit/pto/multi_tile_n4_planmem_e2e.pto @@ -16,7 +16,7 @@ module { func.func @n4_rotate( - %gm : memref<16x16xf16, #pto.address_space>) { + %gm : !pto.partition_tensor_view<16x16xf16>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index @@ -38,13 +38,13 @@ module { : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s0 : !pto.tile_buf) - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s1 : !pto.tile_buf) - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s2 : !pto.tile_buf) - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%s3 : !pto.tile_buf) return } diff --git a/test/lit/pto/multi_tile_no_loop_unroll.pto b/test/lit/pto/multi_tile_no_loop_unroll.pto index 43682523bc..1329211265 100644 --- a/test/lit/pto/multi_tile_no_loop_unroll.pto +++ b/test/lit/pto/multi_tile_no_loop_unroll.pto @@ -13,8 +13,8 @@ module { func.func @unroll_two_slots( - %gm0 : memref<16x16xf16, #pto.address_space>, - %gm1 : memref<16x16xf16, #pto.address_space>) { + %gm0 : !pto.partition_tensor_view<16x16xf16>, + %gm1 : !pto.partition_tensor_view<16x16xf16>) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -24,13 +24,13 @@ module { %s0 = pto.multi_tile_get %mb[%c0] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm0 : !pto.partition_tensor_view<16x16xf16>) outs(%s0 : !pto.tile_buf) %s1 = pto.multi_tile_get %mb[%c1] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm1 : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm1 : !pto.partition_tensor_view<16x16xf16>) outs(%s1 : !pto.tile_buf) return } diff --git a/test/lit/pto/multi_tile_prefetch_dyn_event_id.pto b/test/lit/pto/multi_tile_prefetch_dyn_event_id.pto index 4f00bb59bb..dfe9c8ab97 100644 --- a/test/lit/pto/multi_tile_prefetch_dyn_event_id.pto +++ b/test/lit/pto/multi_tile_prefetch_dyn_event_id.pto @@ -22,8 +22,8 @@ module { func.func @prefetch_dyn_eid( - %gm : memref<16x16xf16, #pto.address_space>, - %dst : memref<16x16xf16, #pto.address_space>, + %gm : !pto.partition_tensor_view<16x16xf16>, + %dst : !pto.partition_tensor_view<16x16xf16>, %n : index) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -38,12 +38,12 @@ module { %prod = pto.multi_tile_get %mb[%idx] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%prod : !pto.tile_buf) // Consumer: V op on dynamic slot. The MTE2 -> V dep on a // multi-buffer alloc should lower to a dyn event-id pair. pto.tstore ins(%prod : !pto.tile_buf) - outs(%dst : memref<16x16xf16, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<16x16xf16>) } return } diff --git a/test/lit/pto/multi_tile_prefetch_gss_event_id.pto b/test/lit/pto/multi_tile_prefetch_gss_event_id.pto index 62188120a4..c051333806 100644 --- a/test/lit/pto/multi_tile_prefetch_gss_event_id.pto +++ b/test/lit/pto/multi_tile_prefetch_gss_event_id.pto @@ -18,8 +18,8 @@ module { func.func @prefetch_gss( - %gm : memref<16x16xf16, #pto.address_space>, - %dst : memref<16x16xf16, #pto.address_space>, + %gm : !pto.partition_tensor_view<16x16xf16>, + %dst : !pto.partition_tensor_view<16x16xf16>, %n : index) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -33,10 +33,10 @@ module { %prod = pto.multi_tile_get %mb[%idx] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%prod : !pto.tile_buf) pto.tstore ins(%prod : !pto.tile_buf) - outs(%dst : memref<16x16xf16, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<16x16xf16>) } return } diff --git a/test/lit/pto/multi_tile_preload_loop_set_wait.pto b/test/lit/pto/multi_tile_preload_loop_set_wait.pto index 8a8f4cce15..db95a22200 100644 --- a/test/lit/pto/multi_tile_preload_loop_set_wait.pto +++ b/test/lit/pto/multi_tile_preload_loop_set_wait.pto @@ -15,8 +15,8 @@ module { func.func @preload_and_loop_set_wait( - %gm : memref<16x16xf16, #pto.address_space>, - %dst : memref<16x16xf16, #pto.address_space>, + %gm : !pto.partition_tensor_view<16x16xf16>, + %dst : !pto.partition_tensor_view<16x16xf16>, %n : index) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -29,20 +29,20 @@ module { %pre = pto.multi_tile_get %mb[%c0] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%pre : !pto.tile_buf) pto.tstore ins(%pre : !pto.tile_buf) - outs(%dst : memref<16x16xf16, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<16x16xf16>) scf.for %i = %c0 to %n step %c1 { %idx = arith.remui %i, %c2 : index %slot = pto.multi_tile_get %mb[%idx] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%slot : !pto.tile_buf) pto.tstore ins(%slot : !pto.tile_buf) - outs(%dst : memref<16x16xf16, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<16x16xf16>) } return } diff --git a/test/lit/pto/multi_tile_two_buf_slot_binding.pto b/test/lit/pto/multi_tile_two_buf_slot_binding.pto index 3be33e95c1..47c8264030 100644 --- a/test/lit/pto/multi_tile_two_buf_slot_binding.pto +++ b/test/lit/pto/multi_tile_two_buf_slot_binding.pto @@ -23,8 +23,8 @@ // cross-iter dependency on B through. module { - func.func @two_mb(%gm : memref<16x16xf16, #pto.address_space>, - %dst : memref<16x16xf16, #pto.address_space>, + func.func @two_mb(%gm : !pto.partition_tensor_view<16x16xf16>, + %dst : !pto.partition_tensor_view<16x16xf16>, %n : index) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -45,13 +45,13 @@ module { : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%sA : !pto.tile_buf) pto.tadd ins(%sA, %sA : !pto.tile_buf, !pto.tile_buf) outs(%sB : !pto.tile_buf) pto.tstore ins(%sB : !pto.tile_buf) - outs(%dst : memref<16x16xf16, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<16x16xf16>) } return } diff --git a/test/lit/pto/multi_tile_unknown_slot_gss_dyn_event_id.pto b/test/lit/pto/multi_tile_unknown_slot_gss_dyn_event_id.pto index 7ea560bb43..29bd0cb239 100644 --- a/test/lit/pto/multi_tile_unknown_slot_gss_dyn_event_id.pto +++ b/test/lit/pto/multi_tile_unknown_slot_gss_dyn_event_id.pto @@ -14,8 +14,8 @@ module { func.func @unknown_slot_gss( - %gm : memref<16x16xf16, #pto.address_space>, - %dst : memref<16x16xf16, #pto.address_space>, + %gm : !pto.partition_tensor_view<16x16xf16>, + %dst : !pto.partition_tensor_view<16x16xf16>, %n : index) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index @@ -32,14 +32,14 @@ module { %prod = pto.multi_tile_get %mb[%prod_idx] : !pto.multi_tile_buf -> !pto.tile_buf - pto.tload ins(%gm : memref<16x16xf16, #pto.address_space>) + pto.tload ins(%gm : !pto.partition_tensor_view<16x16xf16>) outs(%prod : !pto.tile_buf) %cons = pto.multi_tile_get %mb[%cons_idx] : !pto.multi_tile_buf -> !pto.tile_buf pto.tstore ins(%cons : !pto.tile_buf) - outs(%dst : memref<16x16xf16, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<16x16xf16>) } return } diff --git a/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto b/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto index 3bf501afe1..3bc17eb8ec 100644 --- a/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto +++ b/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto @@ -12,22 +12,51 @@ module attributes {"pto.target_arch" = "a3"} { func.func @five_gates_lifetime_overlap( - %gm0 : memref<16x16xf16, #pto.address_space>, - %gm1 : memref<16x16xf16, #pto.address_space>, - %out0 : memref<16x16xf16, #pto.address_space>, - %out1 : memref<16x16xf16, #pto.address_space>) { - %buf0 = memref.alloc() : memref<16x16xf16, #pto.address_space> - %buf1 = memref.alloc() : memref<16x16xf16, #pto.address_space> + %gm0 : !pto.ptr, %gm1 : !pto.ptr, + %out0 : !pto.ptr, %out1 : !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %gm0_view = pto.make_tensor_view %gm0, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c16, %c16, %c16, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %gm1_view = pto.make_tensor_view %gm1, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c16, %c16, %c16, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %out0_view = pto.make_tensor_view %out0, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c16, %c16, %c16, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %out1_view = pto.make_tensor_view %out1, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c16, %c16, %c16, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %gm0_part = pto.partition_view %gm0_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %gm1_part = pto.partition_view %gm1_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %out0_part = pto.partition_view %out0_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %out1_part = pto.partition_view %out1_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %buf0 = pto.alloc_tile : !pto.tile_buf + %buf1 = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) - outs(%buf0 : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} - pto.tload ins(%gm1 : memref<16x16xf16, #pto.address_space>) - outs(%buf1 : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + pto.tload ins(%gm0_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) + outs(%buf0 : !pto.tile_buf) {layout = #pto.layout} + pto.tload ins(%gm1_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) + outs(%buf1 : !pto.tile_buf) {layout = #pto.layout} - pto.tstore ins(%buf0 : memref<16x16xf16, #pto.address_space>) - outs(%out0 : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} - pto.tstore ins(%buf1 : memref<16x16xf16, #pto.address_space>) - outs(%out1 : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + pto.tstore ins(%buf0 : !pto.tile_buf) + outs(%out0_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) {layout = #pto.layout} + pto.tstore ins(%buf1 : !pto.tile_buf) + outs(%out1_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) {layout = #pto.layout} return } } @@ -35,6 +64,7 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK-LABEL: func.func @five_gates_lifetime_overlap // CHECK-DAG: %[[ADDR0:.*]] = arith.constant 0 : i64 // CHECK-DAG: %[[ADDR1:.*]] = arith.constant 512 : i64 -// CHECK: pto.pointer_cast(%[[ADDR0]]) -// CHECK: pto.pointer_cast(%[[ADDR1]]) +// CHECK: pto.alloc_tile addr = %[[ADDR0]] +// CHECK: pto.alloc_tile addr = %[[ADDR1]] +// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_five_gates_phi_family.pto b/test/lit/pto/plan_memory_five_gates_phi_family.pto index ddf295fd80..bbec280795 100644 --- a/test/lit/pto/plan_memory_five_gates_phi_family.pto +++ b/test/lit/pto/plan_memory_five_gates_phi_family.pto @@ -12,24 +12,43 @@ module attributes {"pto.target_arch" = "a3"} { func.func @five_gates_phi_family( - %gm0 : memref<16x16xf16, #pto.address_space>, - %gm1 : memref<16x16xf16, #pto.address_space>, - %out : memref<16x16xf16, #pto.address_space>) { + %gm0 : !pto.ptr, %gm1 : !pto.ptr, %out : !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %gm0_view = pto.make_tensor_view %gm0, + shape = [%c1, %c1, %c1, %c16, %c16], strides = [%c16, %c16, %c16, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %gm1_view = pto.make_tensor_view %gm1, + shape = [%c1, %c1, %c1, %c16, %c16], strides = [%c16, %c16, %c16, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %out_view = pto.make_tensor_view %out, + shape = [%c1, %c1, %c1, %c16, %c16], strides = [%c16, %c16, %c16, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %gm0_part = pto.partition_view %gm0_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %gm1_part = pto.partition_view %gm1_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> %cond = arith.constant true - %merged = scf.if %cond -> (memref<16x16xf16, #pto.address_space>) { - %then = memref.alloc() : memref<16x16xf16, #pto.address_space> - pto.tload ins(%gm0 : memref<16x16xf16, #pto.address_space>) - outs(%then : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} - scf.yield %then : memref<16x16xf16, #pto.address_space> + %merged = scf.if %cond -> (!pto.tile_buf) { + %then = pto.alloc_tile : !pto.tile_buf + pto.tload ins(%gm0_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) + outs(%then : !pto.tile_buf) {layout = #pto.layout} + scf.yield %then : !pto.tile_buf } else { - %else = memref.alloc() : memref<16x16xf16, #pto.address_space> - pto.tload ins(%gm1 : memref<16x16xf16, #pto.address_space>) - outs(%else : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} - scf.yield %else : memref<16x16xf16, #pto.address_space> + %else = pto.alloc_tile : !pto.tile_buf + pto.tload ins(%gm1_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) + outs(%else : !pto.tile_buf) {layout = #pto.layout} + scf.yield %else : !pto.tile_buf } - pto.tstore ins(%merged : memref<16x16xf16, #pto.address_space>) - outs(%out : memref<16x16xf16, #pto.address_space>) {layout = #pto.layout} + pto.tstore ins(%merged : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) {layout = #pto.layout} return } } @@ -37,7 +56,8 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK-LABEL: func.func @five_gates_phi_family // CHECK: %[[ADDR0:.*]] = arith.constant 0 : i64 // CHECK: scf.if -// CHECK: pto.pointer_cast(%[[ADDR0]]) +// CHECK: pto.alloc_tile addr = %[[ADDR0]] // CHECK: } else { -// CHECK: pto.pointer_cast(%[[ADDR0]]) +// CHECK: pto.alloc_tile addr = %[[ADDR0]] +// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_for_iter_args_yield.pto b/test/lit/pto/plan_memory_for_iter_args_yield.pto deleted file mode 100644 index 6f1578aac7..0000000000 --- a/test/lit/pto/plan_memory_for_iter_args_yield.pto +++ /dev/null @@ -1,35 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @for_iter_args_yield(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c2 = arith.constant 2 : index - scf.for %i = %c0 to %c2 step %c1 { - // Two allocs inside the loop body exercise liveness within loops and - // per-iteration memory planning for overlapping lifetimes. - %a = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %b = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%a : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%b : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%a : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%b : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @for_iter_args_yield -// CHECK-NOT: memref.alloc -// CHECK: scf.for -// After matching the loop header, require at least one planned buffer inside -// the loop body (the inner alloc becomes a pointer_cast). -// CHECK: pto.pointer_cast -// CHECK: pto.pointer_cast diff --git a/test/lit/pto/plan_memory_fragmentation_hole_fit.pto b/test/lit/pto/plan_memory_fragmentation_hole_fit.pto deleted file mode 100644 index 39bb3b8494..0000000000 --- a/test/lit/pto/plan_memory_fragmentation_hole_fit.pto +++ /dev/null @@ -1,152 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @fragmentation_hole_fit(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - // Force a tight situation: - // - Keep 23 UB buffers live across the function (23 * 8192 = 188416 B). - // - Only 1 UB slot remains. Two short-lived buffers must reuse that slot. - %k0 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k1 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k2 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k3 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k4 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k5 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k6 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k7 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k8 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k9 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k10 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k11 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k12 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k13 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k14 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k15 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k16 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k17 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k18 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k19 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k20 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k21 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k22 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - - // Touch all k-buffers early so their lifetimes start before the temps. - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k0 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k1 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k2 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k3 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k4 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k5 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k6 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k7 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k8 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k9 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k10 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k11 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k12 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k13 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k14 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k15 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k16 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k17 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k18 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k19 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k20 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k21 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k22 : memref<16x16x16xf16, #pto.address_space>) - - %t0 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%t0 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%t0 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %t1 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%t1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%t1 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - // Keep all k-buffers live until the end. - pto.tstore ins(%k0 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k1 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k2 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k3 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k4 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k5 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k6 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k7 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k8 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k9 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k10 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k11 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k12 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k13 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k14 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k15 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k16 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k17 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k18 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k19 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k20 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k21 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k22 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @fragmentation_hole_fit -// CHECK-NOT: memref.alloc -// With 23 live UB buffers, there is exactly one remaining 8192B slot. The two -// short-lived buffers must reuse that slot (offset 23*8192 = 188416). -// CHECK-DAG: %[[O188416:.*]] = arith.constant 188416 : i64 -// CHECK-DAG: pto.pointer_cast(%[[O188416]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK-DAG: pto.pointer_cast(%[[O188416]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> diff --git a/test/lit/pto/plan_memory_fragmentation_two_holes.pto b/test/lit/pto/plan_memory_fragmentation_two_holes.pto deleted file mode 100644 index 4a5c8be91a..0000000000 --- a/test/lit/pto/plan_memory_fragmentation_two_holes.pto +++ /dev/null @@ -1,163 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @fragmentation_two_holes(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - // Tight + overlap: - // - Keep 22 UB buffers live (22 * 8192 = 180224 B), leaving 2 free slots. - // - Allocate 2 short-lived buffers with overlapping lifetimes twice. - // The two free slots are at offsets 22*8192 and 23*8192. - %k0 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k1 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k2 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k3 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k4 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k5 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k6 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k7 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k8 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k9 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k10 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k11 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k12 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k13 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k14 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k15 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k16 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k17 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k18 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k19 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k20 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %k21 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - - // Touch all k-buffers early so their lifetimes start before the temps. - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k0 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k1 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k2 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k3 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k4 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k5 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k6 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k7 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k8 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k9 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k10 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k11 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k12 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k13 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k14 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k15 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k16 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k17 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k18 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k19 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k20 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%k21 : memref<16x16x16xf16, #pto.address_space>) - - // Stage 1: two overlapping temps. - %a0 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %b0 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%a0 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%b0 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%a0 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%b0 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - // Stage 2: two overlapping temps again, which should reuse the same two slots. - %a1 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %b1 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%a1 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%b1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%a1 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%b1 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - // Keep all k-buffers live until the end. - pto.tstore ins(%k0 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k1 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k2 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k3 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k4 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k5 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k6 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k7 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k8 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k9 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k10 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k11 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k12 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k13 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k14 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k15 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k16 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k17 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k18 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k19 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k20 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%k21 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @fragmentation_two_holes -// CHECK-NOT: memref.alloc -// With 22 live UB buffers, there are exactly two remaining 8192B slots at -// offsets 22*8192 and 23*8192, reused across the two stages. -// CHECK-DAG: %[[O180224:.*]] = arith.constant 180224 : i64 -// CHECK-DAG: %[[O188416:.*]] = arith.constant 188416 : i64 -// CHECK-DAG: pto.pointer_cast(%[[O180224]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK-DAG: pto.pointer_cast(%[[O180224]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK-DAG: pto.pointer_cast(%[[O188416]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK-DAG: pto.pointer_cast(%[[O188416]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> diff --git a/test/lit/pto/plan_memory_if_in_loop.pto b/test/lit/pto/plan_memory_if_in_loop.pto deleted file mode 100644 index 1d902eedce..0000000000 --- a/test/lit/pto/plan_memory_if_in_loop.pto +++ /dev/null @@ -1,38 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @if_in_loop(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c2 = arith.constant 2 : index - - scf.for %i = %c0 to %c2 step %c1 { - %is0 = arith.cmpi eq, %i, %c0 : index - scf.if %is0 { - %a = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%a : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%a : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } else { - %b = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%b : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%b : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } - } - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @if_in_loop -// CHECK-NOT: memref.alloc -// CHECK-DAG: %c0_i64 = arith.constant 0 : i64 -// CHECK: scf.for -// CHECK: scf.if -// CHECK: pto.pointer_cast(%c0_i64) -// CHECK: pto.pointer_cast(%c0_i64) diff --git a/test/lit/pto/plan_memory_if_yield.pto b/test/lit/pto/plan_memory_if_yield.pto deleted file mode 100644 index d3f1e4a578..0000000000 --- a/test/lit/pto/plan_memory_if_yield.pto +++ /dev/null @@ -1,34 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @if_yield(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - %cond = arith.constant true - // Use scf.if control-flow without yielding a memref (the current emitc - // lowering can't handle memref-typed scf.if results), but still ensure - // PlanMemory rewrites allocs inside both branches. - scf.if %cond { - %then = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%then : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%then : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } else { - %els = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%els : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%els : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @if_yield -// CHECK-NOT: memref.alloc -// CHECK: scf.if -// CHECK: pto.pointer_cast -// CHECK: } else { -// CHECK: pto.pointer_cast diff --git a/test/lit/pto/plan_memory_inplace_forbid_alias.pto b/test/lit/pto/plan_memory_inplace_forbid_alias.pto index 54e65a0fd9..de21408ebc 100644 --- a/test/lit/pto/plan_memory_inplace_forbid_alias.pto +++ b/test/lit/pto/plan_memory_inplace_forbid_alias.pto @@ -12,29 +12,29 @@ module attributes {"pto.target_arch" = "a3"} { func.func @inplace_forbid_alias_ttrans() { - %src = memref.alloc() : memref<16x16xf16, #pto.address_space> - %tmp = memref.alloc() : memref<16x16xf16, #pto.address_space> - %dst = memref.alloc() : memref<16x16xf16, #pto.address_space> + %src = pto.alloc_tile : !pto.tile_buf + %tmp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf - pto.ttrans ins(%src, %tmp : memref<16x16xf16, #pto.address_space>, - memref<16x16xf16, #pto.address_space>) - outs(%dst : memref<16x16xf16, #pto.address_space>) + pto.ttrans ins(%src, %tmp : !pto.tile_buf, + !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } func.func @inplace_forbid_alias_tsel() { - %mask = memref.alloc() : memref<16x16xi16, #pto.address_space> - %src0 = memref.alloc() : memref<16x16xf16, #pto.address_space> - %src1 = memref.alloc() : memref<16x16xf16, #pto.address_space> - %tmp = memref.alloc() : memref<16x16xui8, #pto.address_space> - %dst = memref.alloc() : memref<16x16xf16, #pto.address_space> + %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 - : memref<16x16xi16, #pto.address_space>, - memref<16x16xf16, #pto.address_space>, - memref<16x16xf16, #pto.address_space>, - memref<16x16xui8, #pto.address_space>) - outs(%dst : memref<16x16xf16, #pto.address_space>) + : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } } @@ -43,18 +43,19 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK-DAG: %[[A0:.*]] = arith.constant 0 : i64 // CHECK-DAG: %[[A512:.*]] = arith.constant 512 : i64 // CHECK-DAG: %[[A1024:.*]] = arith.constant 1024 : i64 -// CHECK: pto.pointer_cast(%[[A0]]) -// CHECK: pto.pointer_cast(%[[A512]]) -// CHECK: pto.pointer_cast(%[[A1024]]) +// CHECK: pto.alloc_tile addr = %[[A0]] +// CHECK: pto.alloc_tile addr = %[[A512]] +// CHECK: pto.alloc_tile addr = %[[A1024]] // CHECK-LABEL: func.func @inplace_forbid_alias_tsel // CHECK-DAG: %[[B0:.*]] = arith.constant 0 : i64 // CHECK-DAG: %[[B512:.*]] = arith.constant 512 : i64 // CHECK-DAG: %[[B1024:.*]] = arith.constant 1024 : i64 // CHECK-DAG: %[[B1536:.*]] = arith.constant 1536 : i64 // CHECK-DAG: %[[B2048:.*]] = arith.constant 2048 : i64 -// CHECK: pto.pointer_cast(%[[B0]]) -// CHECK: pto.pointer_cast(%[[B512]]) -// CHECK: pto.pointer_cast(%[[B1024]]) -// CHECK: pto.pointer_cast(%[[B2048]]) -// CHECK: pto.pointer_cast(%[[B1536]]) +// CHECK: pto.alloc_tile addr = %[[B0]] +// CHECK: pto.alloc_tile addr = %[[B512]] +// CHECK: pto.alloc_tile addr = %[[B1024]] +// CHECK: pto.alloc_tile addr = %[[B2048]] +// CHECK: pto.alloc_tile addr = %[[B1536]] +// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_loop_in_if.pto b/test/lit/pto/plan_memory_loop_in_if.pto deleted file mode 100644 index 2b4e06be93..0000000000 --- a/test/lit/pto/plan_memory_loop_in_if.pto +++ /dev/null @@ -1,37 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @loop_in_if(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - %true = arith.constant true - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c2 = arith.constant 2 : index - - scf.if %true { - scf.for %i = %c0 to %c2 step %c1 { - %a = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%a : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%a : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } - } else { - %b = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%b : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%b : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @loop_in_if -// CHECK-NOT: memref.alloc -// CHECK: scf.if -// CHECK: scf.for -// CHECK: } else { -// CHECK: pto.pointer_cast diff --git a/test/lit/pto/plan_memory_loop_no_reuse_outer_live.pto b/test/lit/pto/plan_memory_loop_no_reuse_outer_live.pto deleted file mode 100644 index 82131666c6..0000000000 --- a/test/lit/pto/plan_memory_loop_no_reuse_outer_live.pto +++ /dev/null @@ -1,41 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @loop_outer_live(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c4 = arith.constant 4 : index - - // A buffer that remains live across the loop (used after the loop). - %outer = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%outer : memref<16x16x16xf16, #pto.address_space>) - - // A loop-local buffer used inside the loop. - %inner = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - scf.for %i = %c0 to %c4 step %c1 { - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%inner : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%inner : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } - - // Use %outer after the loop to keep it live across the loop. - pto.tstore ins(%outer : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @loop_outer_live -// CHECK-NOT: memref.alloc -// Expect a loop, and two planned buffers at distinct offsets. -// CHECK-DAG: %c0_i64 = arith.constant 0 : i64 -// CHECK-DAG: %c8192_i64 = arith.constant 8192 : i64 -// CHECK-DAG: pto.pointer_cast(%c0_i64) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK-DAG: pto.pointer_cast(%c8192_i64) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK: scf.for diff --git a/test/lit/pto/plan_memory_nested_loops.pto b/test/lit/pto/plan_memory_nested_loops.pto deleted file mode 100644 index f4a6710512..0000000000 --- a/test/lit/pto/plan_memory_nested_loops.pto +++ /dev/null @@ -1,46 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @nested_loops(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %c2 = arith.constant 2 : index - - // Outer buffer stays live across both loops (used after). - %outer = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%outer : memref<16x16x16xf16, #pto.address_space>) - - scf.for %i = %c0 to %c2 step %c1 { - // Buffer allocated inside the outer loop. - %mid = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%mid : memref<16x16x16xf16, #pto.address_space>) - scf.for %j = %c0 to %c2 step %c1 { - // A temp buffer in the inner loop. - %tmp = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%tmp : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%tmp : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } - pto.tstore ins(%mid : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - } - - pto.tstore ins(%outer : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @nested_loops -// CHECK-NOT: memref.alloc -// CHECK-DAG: %c0_i64 = arith.constant 0 : i64 -// CHECK-DAG: %c8192_i64 = arith.constant 8192 : i64 -// CHECK: scf.for -// CHECK: scf.for -// CHECK: pto.pointer_cast diff --git a/test/lit/pto/plan_memory_no_reuse_overlap.pto b/test/lit/pto/plan_memory_no_reuse_overlap.pto deleted file mode 100644 index 8e04a9e601..0000000000 --- a/test/lit/pto/plan_memory_no_reuse_overlap.pto +++ /dev/null @@ -1,31 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @no_reuse_overlap(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - %ub0 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %ub1 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - - // Make lifetimes overlap by using both buffers after both are created. - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub0 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub0 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub1 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @no_reuse_overlap -// CHECK-NOT: memref.alloc -// With overlapping lifetimes, offsets must differ. -// CHECK-DAG: %c0_i64 = arith.constant 0 : i64 -// CHECK-DAG: %c8192_i64 = arith.constant 8192 : i64 -// CHECK-DAG: pto.pointer_cast(%c0_i64) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK-DAG: pto.pointer_cast(%c8192_i64) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> diff --git a/test/lit/pto/plan_memory_peak_8_overlapping.pto b/test/lit/pto/plan_memory_peak_8_overlapping.pto deleted file mode 100644 index b42f8fd997..0000000000 --- a/test/lit/pto/plan_memory_peak_8_overlapping.pto +++ /dev/null @@ -1,59 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @peak_8_overlapping(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - // Peak liveness: 8 buffers live at once. - %u0 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u1 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u2 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u3 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u4 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u5 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u6 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u7 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u0 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u1 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u2 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u3 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u4 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u5 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u6 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u7 : memref<16x16x16xf16, #pto.address_space>) - - pto.tstore ins(%u0 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u1 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u2 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u3 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u4 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u5 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u6 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u7 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @peak_8_overlapping -// CHECK-NOT: memref.alloc -// 8 live buffers implies a max offset of 7*8192 = 57344 bytes. -// CHECK: %[[O57344:.*]] = arith.constant 57344 : i64 -// CHECK: pto.pointer_cast(%[[O57344]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> diff --git a/test/lit/pto/plan_memory_peak_exact_capacity.pto b/test/lit/pto/plan_memory_peak_exact_capacity.pto deleted file mode 100644 index 8ecb828a99..0000000000 --- a/test/lit/pto/plan_memory_peak_exact_capacity.pto +++ /dev/null @@ -1,141 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @peak_exact_capacity(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - // Default UB size is 1572864 bits (196608 bytes). Each buffer here is - // 16*16*16*f16 = 8192 bytes. 24 buffers live at once should fit exactly: - // 24 * 8192 = 196608 bytes. - %u0 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u1 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u2 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u3 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u4 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u5 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u6 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u7 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u8 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u9 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u10 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u11 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u12 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u13 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u14 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u15 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u16 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u17 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u18 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u19 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u20 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u21 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u22 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - %u23 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u0 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u1 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u2 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u3 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u4 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u5 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u6 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u7 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u8 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u9 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u10 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u11 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u12 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u13 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u14 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u15 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u16 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u17 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u18 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u19 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u20 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u21 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u22 : memref<16x16x16xf16, #pto.address_space>) - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%u23 : memref<16x16x16xf16, #pto.address_space>) - - pto.tstore ins(%u0 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u1 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u2 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u3 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u4 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u5 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u6 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u7 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u8 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u9 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u10 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u11 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u12 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u13 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u14 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u15 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u16 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u17 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u18 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u19 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u20 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u21 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u22 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%u23 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @peak_exact_capacity -// CHECK-NOT: memref.alloc -// 24 live buffers implies a max offset of 23*8192 = 188416 bytes. -// CHECK: %[[O188416:.*]] = arith.constant 188416 : i64 -// CHECK: pto.pointer_cast(%[[O188416]]) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> diff --git a/test/lit/pto/plan_memory_reserve_buffer_manual_reject_nested_alloc.pto b/test/lit/pto/plan_memory_reserve_buffer_manual_reject_nested_alloc.pto deleted file mode 100644 index b89bb80022..0000000000 --- a/test/lit/pto/plan_memory_reserve_buffer_manual_reject_nested_alloc.pto +++ /dev/null @@ -1,26 +0,0 @@ -// RUN: not ptoas %s 2>&1 1>/dev/null | FileCheck %s -// RUN: not ptoas --plan-memory-impl=modern %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @manual_nested_alloc(%cond: i1) { - %fifo = pto.reserve_buffer { - name = "fifo", - size = 8192, - location = #pto.address_space, - auto = false, - base = 0 - } -> i32 - - %ub = scf.if %cond -> (memref<16x16x16xf16, #pto.address_space>) { - %true_ub = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - scf.yield %true_ub : memref<16x16x16xf16, #pto.address_space> - } else { - %false_ub = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - scf.yield %false_ub : memref<16x16x16xf16, #pto.address_space> - } - memref.dealloc %ub : memref<16x16x16xf16, #pto.address_space> - return - } -} - -// CHECK: error: pto.reserve_buffer with explicit 'base' (auto = false) is not supported when --pto-level=level1 or level2; use --pto-level=level3 or set auto = true diff --git a/test/lit/pto/plan_memory_reserve_buffer_prefix.pto b/test/lit/pto/plan_memory_reserve_buffer_prefix.pto deleted file mode 100644 index a5bcc2d249..0000000000 --- a/test/lit/pto/plan_memory_reserve_buffer_prefix.pto +++ /dev/null @@ -1,27 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @reserve_prefix(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - %fifo = pto.reserve_buffer { - name = "fifo", - size = 8192, - location = #pto.address_space, - auto = true - } -> i32 - - %ub = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - return - } -} - -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @reserve_prefix( -// CHECK: %c0_i64 = arith.constant 0 : i64 -// CHECK: %{{.*}} = pto.reserve_buffer{name = "fifo", size = 8192, location = , auto = true, base = 8192} -> i32 -// CHECK: %{{.*}} = pto.pointer_cast(%c0_i64) : memref<16x16x16xf16, #pto.address_space> diff --git a/test/lit/pto/plan_memory_reuse_sequential.pto b/test/lit/pto/plan_memory_reuse_sequential.pto deleted file mode 100644 index 425a3b399f..0000000000 --- a/test/lit/pto/plan_memory_reuse_sequential.pto +++ /dev/null @@ -1,203 +0,0 @@ -// RUN: ptoas --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s -// RUN: ptoas --plan-memory-impl=modern --mlir-print-ir-after=pto-plan-memory %s 2>&1 1>/dev/null | FileCheck %s - -module { - func.func @reuse_sequential(%arg0: memref<16x16x16xf16, #pto.address_space>, - %arg1: memref<16x16x16xf16, #pto.address_space>) { - // Force reuse: - // UB capacity (default) is 1572864 bits (196608 bytes). Each buffer here is - // 16*16*16*f16 = 8192 bytes. Allocating 30 such buffers exceeds UB capacity - // unless memory reuse is applied. - %ub0 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub0 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub0 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub1 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub1 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub1 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub2 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub2 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub2 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub3 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub3 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub3 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub4 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub4 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub4 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub5 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub5 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub5 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub6 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub6 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub6 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub7 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub7 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub7 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub8 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub8 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub8 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub9 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub9 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub9 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub10 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub10 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub10 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub11 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub11 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub11 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub12 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub12 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub12 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub13 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub13 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub13 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub14 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub14 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub14 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub15 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub15 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub15 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub16 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub16 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub16 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub17 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub17 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub17 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub18 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub18 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub18 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub19 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub19 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub19 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub20 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub20 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub20 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub21 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub21 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub21 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub22 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub22 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub22 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub23 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub23 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub23 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub24 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub24 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub24 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub25 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub25 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub25 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub26 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub26 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub26 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub27 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub27 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub27 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub28 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub28 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub28 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - %ub29 = memref.alloc() : memref<16x16x16xf16, #pto.address_space> - pto.tload ins(%arg0 : memref<16x16x16xf16, #pto.address_space>) - outs(%ub29 : memref<16x16x16xf16, #pto.address_space>) - pto.tstore ins(%ub29 : memref<16x16x16xf16, #pto.address_space>) - outs(%arg1 : memref<16x16x16xf16, #pto.address_space>) - - return - } -} - -// Anchor checks after the PlanMemory debug marker (ptoas prints the module -// before and after planning). -// CHECK: IR Dump After PlanMemory -// CHECK: func.func @reuse_sequential -// CHECK-NOT: memref.alloc -// Expect at least two distinct allocations to reuse offset 0. -// CHECK: %c0_i64 = arith.constant 0 : i64 -// CHECK: %[[BUF0:.*]] = pto.pointer_cast(%c0_i64) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> -// CHECK: %[[BUF1:.*]] = pto.pointer_cast(%c0_i64) : memref<16x16x16xf16, #pto.address_space<{{vec|ub}}>> diff --git a/test/lit/pto/plan_memory_reused_tstore_sync.pto b/test/lit/pto/plan_memory_reused_tstore_sync.pto index bece344319..831aa98488 100644 --- a/test/lit/pto/plan_memory_reused_tstore_sync.pto +++ b/test/lit/pto/plan_memory_reused_tstore_sync.pto @@ -6,8 +6,7 @@ // 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. -// PlanMemory lowers local allocations to pointer_casts before InsertSync. This -// test models the lowered level2 form directly: `stored` occupies +// This test models explicit level3 tile addresses directly: `stored` occupies // [1280, 9472) and `reused` occupies [8192, 16384). `stored` dies at tstore, // but MTE3 still reads its physical range asynchronously, so the later vector // write needs an MTE3 -> V handshake. @@ -16,32 +15,50 @@ module attributes {pto.target_arch = "a2a3"} { func.func @plan_memory_reused_tstore_sync( - %src0: memref<16x128xf32, #pto.address_space>, - %src1: memref<16x128xf32, #pto.address_space>, - %dst: memref<16x128xf32, #pto.address_space>) + %src0: !pto.ptr, %src1: !pto.ptr, %dst: !pto.ptr) attributes {pto.kernel_kind = #pto.kernel_kind} { %c1280_i64 = arith.constant 1280 : i64 %c8192_i64 = arith.constant 8192 : i64 %c16384_i64 = arith.constant 16384 : i64 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index %one = arith.constant 1.000000e+00 : f32 - %input = pto.pointer_cast(%c16384_i64) - : memref<16x128xf32, #pto.address_space> - %stored = pto.pointer_cast(%c1280_i64) - : memref<16x128xf32, #pto.address_space> - %reused = pto.pointer_cast(%c8192_i64) - : memref<16x128xf32, #pto.address_space> + %src0_view = pto.make_tensor_view %src0, + shape = [%c1, %c1, %c1, %c16, %c128], strides = [%c128, %c128, %c128, %c128, %c1] + : !pto.tensor_view<1x1x1x16x128xf32> + %src1_view = pto.make_tensor_view %src1, + shape = [%c1, %c1, %c1, %c16, %c128], strides = [%c128, %c128, %c128, %c128, %c1] + : !pto.tensor_view<1x1x1x16x128xf32> + %dst_view = pto.make_tensor_view %dst, + shape = [%c1, %c1, %c1, %c16, %c128], strides = [%c128, %c128, %c128, %c128, %c1] + : !pto.tensor_view<1x1x1x16x128xf32> + %src0_part = pto.partition_view %src0_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c128] + : !pto.tensor_view<1x1x1x16x128xf32> -> !pto.partition_tensor_view<1x1x1x16x128xf32> + %src1_part = pto.partition_view %src1_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c128] + : !pto.tensor_view<1x1x1x16x128xf32> -> !pto.partition_tensor_view<1x1x1x16x128xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c128] + : !pto.tensor_view<1x1x1x16x128xf32> -> !pto.partition_tensor_view<1x1x1x16x128xf32> - pto.tload ins(%src0 : memref<16x128xf32, #pto.address_space>) - outs(%input : memref<16x128xf32, #pto.address_space>) - pto.tload ins(%src1 : memref<16x128xf32, #pto.address_space>) - outs(%stored : memref<16x128xf32, #pto.address_space>) - pto.tmuls ins(%input, %one : memref<16x128xf32, #pto.address_space>, f32) - outs(%stored : memref<16x128xf32, #pto.address_space>) - pto.tstore ins(%stored : memref<16x128xf32, #pto.address_space>) - outs(%dst : memref<16x128xf32, #pto.address_space>) - pto.tmuls ins(%input, %one : memref<16x128xf32, #pto.address_space>, f32) - outs(%reused : memref<16x128xf32, #pto.address_space>) + %input = pto.alloc_tile addr = %c16384_i64 : !pto.tile_buf + %stored = pto.alloc_tile addr = %c1280_i64 : !pto.tile_buf + %reused = pto.alloc_tile addr = %c8192_i64 : !pto.tile_buf + + pto.tload ins(%src0_part : !pto.partition_tensor_view<1x1x1x16x128xf32>) + outs(%input : !pto.tile_buf) + pto.tload ins(%src1_part : !pto.partition_tensor_view<1x1x1x16x128xf32>) + outs(%stored : !pto.tile_buf) + pto.tmuls ins(%input, %one : !pto.tile_buf, f32) + outs(%stored : !pto.tile_buf) + pto.tstore ins(%stored : !pto.tile_buf) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x128xf32>) + pto.tmuls ins(%input, %one : !pto.tile_buf, f32) + outs(%reused : !pto.tile_buf) return } } diff --git a/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto b/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto index 06a47e100b..5dffba7a56 100644 --- a/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto +++ b/test/lit/pto/plan_memory_reused_tstore_sync_level2.pto @@ -16,37 +16,75 @@ module attributes {pto.target_arch = "a2a3"} { func.func @plan_memory_reused_tstore_sync_level2( - %pressure_src: memref<1x45056xf32, #pto.address_space>, - %pressure_dst: memref<1x45056xf32, #pto.address_space>, - %src0: memref<16x128xf32, #pto.address_space>, - %src1: memref<16x128xf32, #pto.address_space>, - %dst: memref<16x128xf32, #pto.address_space>, + %pressure_src: !pto.ptr, %pressure_dst: !pto.ptr, + %src0: !pto.ptr, %src1: !pto.ptr, %dst: !pto.ptr, %blocker_rows: index, %blocker_cols: index) attributes {pto.kernel_kind = #pto.kernel_kind} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c45056 = arith.constant 45056 : index %one = arith.constant 1.000000e+00 : f32 + %pressure_src_view = pto.make_tensor_view %pressure_src, + shape = [%c1, %c1, %c1, %c1, %c45056], + strides = [%c45056, %c45056, %c45056, %c45056, %c1] + : !pto.tensor_view<1x1x1x1x45056xf32> + %pressure_dst_view = pto.make_tensor_view %pressure_dst, + shape = [%c1, %c1, %c1, %c1, %c45056], + strides = [%c45056, %c45056, %c45056, %c45056, %c1] + : !pto.tensor_view<1x1x1x1x45056xf32> + %src0_view = pto.make_tensor_view %src0, + shape = [%c1, %c1, %c1, %c16, %c128], + strides = [%c128, %c128, %c128, %c128, %c1] + : !pto.tensor_view<1x1x1x16x128xf32> + %src1_view = pto.make_tensor_view %src1, + shape = [%c1, %c1, %c1, %c16, %c128], + strides = [%c128, %c128, %c128, %c128, %c1] + : !pto.tensor_view<1x1x1x16x128xf32> + %dst_view = pto.make_tensor_view %dst, + shape = [%c1, %c1, %c1, %c16, %c128], + strides = [%c128, %c128, %c128, %c128, %c1] + : !pto.tensor_view<1x1x1x16x128xf32> + %pressure_src_part = pto.partition_view %pressure_src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c1, %c45056] + : !pto.tensor_view<1x1x1x1x45056xf32> -> !pto.partition_tensor_view<1x1x1x1x45056xf32> + %pressure_dst_part = pto.partition_view %pressure_dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c1, %c45056] + : !pto.tensor_view<1x1x1x1x45056xf32> -> !pto.partition_tensor_view<1x1x1x1x45056xf32> + %src0_part = pto.partition_view %src0_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c128] + : !pto.tensor_view<1x1x1x16x128xf32> -> !pto.partition_tensor_view<1x1x1x16x128xf32> + %src1_part = pto.partition_view %src1_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c128] + : !pto.tensor_view<1x1x1x16x128xf32> -> !pto.partition_tensor_view<1x1x1x16x128xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c1, %c1, %c16, %c128] + : !pto.tensor_view<1x1x1x16x128xf32> -> !pto.partition_tensor_view<1x1x1x16x128xf32> + // Force PlanMemory's reuse path: 180224 bytes here plus the target tiles // and the 256-byte blocker exceed the default 196608-byte UB capacity, // while their actual live ranges remain non-overlapping enough to fit. %pressure = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%pressure_src : memref<1x45056xf32, #pto.address_space>) + pto.tload ins(%pressure_src_part : !pto.partition_tensor_view<1x1x1x1x45056xf32>) outs(%pressure : !pto.tile_buf) pto.tstore ins(%pressure : !pto.tile_buf) - outs(%pressure_dst : memref<1x45056xf32, #pto.address_space>) + outs(%pressure_dst_part : !pto.partition_tensor_view<1x1x1x1x45056xf32>) %input = pto.alloc_tile : !pto.tile_buf %stored = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src0 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src0_part : !pto.partition_tensor_view<1x1x1x16x128xf32>) outs(%input : !pto.tile_buf) - pto.tload ins(%src1 : memref<16x128xf32, #pto.address_space>) + pto.tload ins(%src1_part : !pto.partition_tensor_view<1x1x1x16x128xf32>) outs(%stored : !pto.tile_buf) pto.tmuls ins(%input, %one : !pto.tile_buf, f32) outs(%stored : !pto.tile_buf) pto.tstore ins(%stored : !pto.tile_buf) - outs(%dst : memref<16x128xf32, #pto.address_space>) + outs(%dst_part : !pto.partition_tensor_view<1x1x1x16x128xf32>) %blocker = pto.alloc_tile valid_row = %blocker_rows valid_col = %blocker_cols : !pto.tile_buf %reused = pto.alloc_tile : diff --git a/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto b/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto index fe520f08bc..abce67d01f 100644 --- a/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto +++ b/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto @@ -11,14 +11,26 @@ module attributes {"pto.target_arch" = "a3"} { func.func @spec_level0_no_reuse_overlap( - %gm : memref<32x32xf16, #pto.address_space>) { - %src = memref.alloc() : memref<32x32xf16, #pto.address_space> - %dst = memref.alloc() : memref<32x32xf16, #pto.address_space> + %gm : !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %gm_view = pto.make_tensor_view %gm, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c32, %c32, %c32, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + %gm_part = pto.partition_view %gm_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> + -> !pto.partition_tensor_view<1x1x1x32x32xf16> + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%gm : memref<32x32xf16, #pto.address_space>) - outs(%src : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} - pto.tmov ins(%src : memref<32x32xf16, #pto.address_space>) - outs(%dst : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%gm_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + outs(%src : !pto.tile_buf) {layout = #pto.layout} + pto.tmov ins(%src : !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } } @@ -26,6 +38,7 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK-LABEL: func.func @spec_level0_no_reuse_overlap // CHECK-DAG: %[[ADDR0:.*]] = arith.constant 0 : i64 // CHECK-DAG: %[[ADDR1:.*]] = arith.constant 2048 : i64 -// CHECK: pto.pointer_cast(%[[ADDR0]]) -// CHECK: pto.pointer_cast(%[[ADDR1]]) +// CHECK: pto.alloc_tile addr = %[[ADDR0]] +// CHECK: pto.alloc_tile addr = %[[ADDR1]] +// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_spec_level0_reuse.pto b/test/lit/pto/plan_memory_spec_level0_reuse.pto index b55d457086..7fcbdebfd5 100644 --- a/test/lit/pto/plan_memory_spec_level0_reuse.pto +++ b/test/lit/pto/plan_memory_spec_level0_reuse.pto @@ -11,21 +11,42 @@ module attributes {"pto.target_arch" = "a3"} { func.func @spec_level0_reuse( - %gm0 : memref<32x32xf16, #pto.address_space>, - %gm1 : memref<32x32xf16, #pto.address_space>) { - %buf0 = memref.alloc() : memref<32x32xf16, #pto.address_space> - pto.tload ins(%gm0 : memref<32x32xf16, #pto.address_space>) - outs(%buf0 : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} + %gm0 : !pto.ptr, %gm1 : !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c32 = arith.constant 32 : index + %gm0_view = pto.make_tensor_view %gm0, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c32, %c32, %c32, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + %gm1_view = pto.make_tensor_view %gm1, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c32, %c32, %c32, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + %gm0_part = pto.partition_view %gm0_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> + -> !pto.partition_tensor_view<1x1x1x32x32xf16> + %gm1_part = pto.partition_view %gm1_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> + -> !pto.partition_tensor_view<1x1x1x32x32xf16> + %buf0 = pto.alloc_tile : !pto.tile_buf + pto.tload ins(%gm0_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + outs(%buf0 : !pto.tile_buf) {layout = #pto.layout} - %buf1 = memref.alloc() : memref<32x32xf16, #pto.address_space> - pto.tload ins(%gm1 : memref<32x32xf16, #pto.address_space>) - outs(%buf1 : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} + %buf1 = pto.alloc_tile : !pto.tile_buf + pto.tload ins(%gm1_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + outs(%buf1 : !pto.tile_buf) {layout = #pto.layout} return } } // CHECK-LABEL: func.func @spec_level0_reuse // CHECK: %[[ADDR0:.*]] = arith.constant 0 : i64 -// CHECK: %[[BUF0:.*]] = pto.pointer_cast(%[[ADDR0]]) -// CHECK: %[[BUF1:.*]] = pto.pointer_cast(%[[ADDR0]]) +// CHECK: %[[BUF0:.*]] = pto.alloc_tile addr = %[[ADDR0]] +// CHECK: %[[BUF1:.*]] = pto.alloc_tile addr = %[[ADDR0]] +// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto b/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto index 3a5c05124f..609ff8db15 100644 --- a/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto +++ b/test/lit/pto/ptodsl_subkernel_helper_tile_abi_emitc.pto @@ -8,7 +8,7 @@ // Guards late PTODSL helper inlining after tile handle materialization. // RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-resolve-reserved-buffers %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=NATIVE,COMMON -// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON +// RUN: ptoas --pto-arch=a5 --pto-level=level3 --mlir-print-ir-before=pto-inline-backend-helpers %s -o /dev/null 2>&1 | FileCheck %s --check-prefixes=MATERIALIZED,COMMON // RUN: ptoas --pto-arch=a5 --pto-level=level3 %s -o - 2>&1 | FileCheck %s --check-prefix=EMITC // RUN: ptoas --pto-arch=a5 --pto-level=level3 --pto-backend=vpto --emit-vpto %s -o - 2>&1 | FileCheck %s --check-prefix=VPTO @@ -54,7 +54,7 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind // COMMON: func.call @ptodsl_helper(%[[ALLOC]], {{.*}}) : (!pto.tile_buf, f32) -> () diff --git a/test/lit/pto/ptodsl_subkernel_pipe_contract_reject_transitive_mte.pto b/test/lit/pto/ptodsl_subkernel_pipe_contract_reject_transitive_mte.pto index a253bcae99..0248217f76 100644 --- a/test/lit/pto/ptodsl_subkernel_pipe_contract_reject_transitive_mte.pto +++ b/test/lit/pto/ptodsl_subkernel_pipe_contract_reject_transitive_mte.pto @@ -9,22 +9,22 @@ // RUN: not ptoas --pto-arch=a5 --emit-pto-ir %s 2>&1 | FileCheck %s module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { - func.func @kernel(%src: memref<16x16xf32, #pto.address_space>) { - func.call @simd_helper(%src) : (memref<16x16xf32, #pto.address_space>) -> () + func.func @kernel(%src: !pto.partition_tensor_view<16x16xf32>) { + func.call @simd_helper(%src) : (!pto.partition_tensor_view<16x16xf32>) -> () func.return } - func.func private @simd_helper(%src: memref<16x16xf32, #pto.address_space>) attributes {pto.ptodsl.subkernel_helper = "simd"} { + func.func private @simd_helper(%src: !pto.partition_tensor_view<16x16xf32>) attributes {pto.ptodsl.subkernel_helper = "simd"} { pto.section.vector { - func.call @local_helper(%src) : (memref<16x16xf32, #pto.address_space>) -> () + func.call @local_helper(%src) : (!pto.partition_tensor_view<16x16xf32>) -> () } func.return } - func.func private @local_helper(%src: memref<16x16xf32, #pto.address_space>) { + func.func private @local_helper(%src: !pto.partition_tensor_view<16x16xf32>) { pto.section.vector { %tile = pto.declare_tile -> !pto.tile_buf - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) outs(%tile : !pto.tile_buf) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%tile : !pto.tile_buf) } func.return } diff --git a/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto b/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto index ad9fb57710..96d7516763 100644 --- a/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto +++ b/test/lit/pto/subview_boxed_row_major_column_slice_preserve_parent_shape_a5.pto @@ -1,4 +1,4 @@ -// RUN: ptoas --pto-arch=a5 --mlir-print-ir-after=pto-materialize-tile-handles %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a5 --mlir-print-ir-before=pto-inline-backend-helpers %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR // RUN: ptoas --pto-arch=a5 %s 2>&1 | FileCheck %s --check-prefix=EMITC module { @@ -29,7 +29,7 @@ module { } } -// IR: IR Dump After PTOMaterializeTileHandles +// IR: IR Dump Before PTOInlineBackendHelpers // IR: %[[ELEM_OFFSET:.*]] = arith.addi %{{.*}}, %{{.*}} : i64 // IR: %[[ELEM_BYTES:.*]] = arith.constant 4 : i64 // IR: %[[BYTE_OFFSET:.*]] = arith.muli %[[ELEM_OFFSET]], %[[ELEM_BYTES]] : i64 diff --git a/test/lit/pto/subview_col_major_compact_keeps_normal_shape.pto b/test/lit/pto/subview_col_major_compact_keeps_normal_shape.pto index 108c0b070a..841d44d287 100644 --- a/test/lit/pto/subview_col_major_compact_keeps_normal_shape.pto +++ b/test/lit/pto/subview_col_major_compact_keeps_normal_shape.pto @@ -2,12 +2,12 @@ module { func.func @subview_col_major_compact_keeps_normal_shape( - %src: memref<16x16xf32, #pto.address_space>, - %dst: memref<16x8xf32, #pto.address_space>) { + %src: !pto.partition_tensor_view<16x16xf32>, + %dst: !pto.partition_tensor_view<16x8xf32>) { %c0 = arith.constant 0 : index %tile = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%tile : !pto.tile_buf) %s = pto.subview %tile[%c0, %c0] sizes [16, 8] : @@ -15,7 +15,7 @@ module { -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<16x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<16x8xf32>) return } } diff --git a/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto b/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto index 644bbe8386..b1269251c0 100644 --- a/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto +++ b/test/lit/pto/subview_col_major_noncompact_preserve_stride.pto @@ -2,13 +2,13 @@ module { func.func @subview_col_major_noncompact_preserve_stride( - %src: memref<16x16xf32, #pto.address_space>, - %dst: memref<8x8xf32, #pto.address_space>) { + %src: !pto.partition_tensor_view<16x16xf32>, + %dst: !pto.partition_tensor_view<8x8xf32>) { %c0 = arith.constant 0 : index %c8 = arith.constant 8 : index %tile = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%tile : !pto.tile_buf) %s = pto.subview %tile[%c8, %c0] sizes [8, 8] : @@ -16,7 +16,7 @@ module { -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<8x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<8x8xf32>) return } } diff --git a/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto b/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto index 847484b8ca..e980326bdc 100644 --- a/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto +++ b/test/lit/pto/subview_col_major_row_plus_one_stride_offset.pto @@ -2,12 +2,12 @@ module { func.func @subview_col_major_row_plus_one_stride_offset( - %src: memref<16x16xf32, #pto.address_space>, - %dst: memref<8x8xf32, #pto.address_space>) { + %src: !pto.partition_tensor_view<16x16xf32>, + %dst: !pto.partition_tensor_view<8x8xf32>) { %c8 = arith.constant 8 : index %tile = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%tile : !pto.tile_buf) %s = pto.subview %tile[%c8, %c8] sizes [8, 8] : @@ -15,7 +15,7 @@ module { -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<8x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<8x8xf32>) return } } diff --git a/test/lit/pto/subview_compact_keeps_normal_shape.pto b/test/lit/pto/subview_compact_keeps_normal_shape.pto index ae18825a05..7d8abb6e06 100644 --- a/test/lit/pto/subview_compact_keeps_normal_shape.pto +++ b/test/lit/pto/subview_compact_keeps_normal_shape.pto @@ -2,12 +2,12 @@ module { func.func @subview_compact_keeps_normal_shape( - %src: memref<16x16xf32, #pto.address_space>, - %dst: memref<8x16xf32, #pto.address_space>) { + %src: !pto.partition_tensor_view<16x16xf32>, + %dst: !pto.partition_tensor_view<8x16xf32>) { %c0 = arith.constant 0 : index %tile = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%tile : !pto.tile_buf) %s0 = pto.subview %tile[%c0, %c0] sizes [8, 16] : @@ -15,7 +15,7 @@ module { -> !pto.tile_buf pto.tstore ins(%s0 : !pto.tile_buf) - outs(%dst : memref<8x16xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<8x16xf32>) return } } diff --git a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto index 7dde92792d..65c43d9cc2 100644 --- a/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto +++ b/test/lit/pto/subview_dynamic_offset_static_valid_regression.pto @@ -4,7 +4,7 @@ module { func.func @subview_dynamic_offset_static_valid_regression( %arg0: index, - %dst: memref<1x64xf32, #pto.address_space>) { + %dst: !pto.partition_tensor_view<1x64xf32>) { %c0 = arith.constant 0 : index %c64 = arith.constant 64 : index %off = arith.muli %arg0, %c64 : index @@ -19,7 +19,7 @@ module { -> !pto.tile_buf pto.tstore ins(%sub : !pto.tile_buf) - outs(%dst : memref<1x64xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<1x64xf32>) return } } diff --git a/test/lit/pto/subview_explicit_result_type.pto b/test/lit/pto/subview_explicit_result_type.pto index e6f74e5705..3fba213221 100644 --- a/test/lit/pto/subview_explicit_result_type.pto +++ b/test/lit/pto/subview_explicit_result_type.pto @@ -2,13 +2,13 @@ module { func.func @subview_explicit_result_type( - %src: memref<16x16xf32, #pto.address_space>, - %dst: memref<8x8xf32, #pto.address_space>) { + %src: !pto.partition_tensor_view<16x16xf32>, + %dst: !pto.partition_tensor_view<8x8xf32>) { %c0 = arith.constant 0 : index %c8 = arith.constant 8 : index %tile = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%tile : !pto.tile_buf) %s = pto.subview %tile[%c8, %c0] sizes [8, 8] : @@ -16,7 +16,7 @@ module { -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<8x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<8x8xf32>) return } } diff --git a/test/lit/pto/subview_row_plus_one_stride_offset.pto b/test/lit/pto/subview_row_plus_one_stride_offset.pto index 80d1625663..44e379f5a4 100644 --- a/test/lit/pto/subview_row_plus_one_stride_offset.pto +++ b/test/lit/pto/subview_row_plus_one_stride_offset.pto @@ -2,12 +2,12 @@ module { func.func @subview_row_plus_one_stride_offset( - %src: memref<16x16xf32, #pto.address_space>, - %dst: memref<8x8xf32, #pto.address_space>) { + %src: !pto.partition_tensor_view<16x16xf32>, + %dst: !pto.partition_tensor_view<8x8xf32>) { %c8 = arith.constant 8 : index %tile = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%tile : !pto.tile_buf) %s = pto.subview %tile[%c8, %c8] sizes [8, 8] : @@ -15,7 +15,7 @@ module { -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<8x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<8x8xf32>) return } } diff --git a/test/lit/pto/subview_tile_native_preserve_stride.pto b/test/lit/pto/subview_tile_native_preserve_stride.pto index 2d3cb2a329..d44357ce38 100644 --- a/test/lit/pto/subview_tile_native_preserve_stride.pto +++ b/test/lit/pto/subview_tile_native_preserve_stride.pto @@ -2,16 +2,16 @@ module { func.func @subview_tile_native_preserve_stride( - %src: memref<16x16xf32, #pto.address_space>, - %dst0: memref<8x8xf32, #pto.address_space>, - %dst1: memref<8x8xf32, #pto.address_space>, - %dst2: memref<8x8xf32, #pto.address_space>, - %dst3: memref<8x8xf32, #pto.address_space>) { + %src: !pto.partition_tensor_view<16x16xf32>, + %dst0: !pto.partition_tensor_view<8x8xf32>, + %dst1: !pto.partition_tensor_view<8x8xf32>, + %dst2: !pto.partition_tensor_view<8x8xf32>, + %dst3: !pto.partition_tensor_view<8x8xf32>) { %c0 = arith.constant 0 : index %c8 = arith.constant 8 : index %tile = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) outs(%tile : !pto.tile_buf) %s00 = pto.subview %tile[%c0, %c0] sizes [8, 8] : @@ -28,13 +28,13 @@ module { -> !pto.tile_buf pto.tstore ins(%s00 : !pto.tile_buf) - outs(%dst0 : memref<8x8xf32, #pto.address_space>) + outs(%dst0 : !pto.partition_tensor_view<8x8xf32>) pto.tstore ins(%s01 : !pto.tile_buf) - outs(%dst1 : memref<8x8xf32, #pto.address_space>) + outs(%dst1 : !pto.partition_tensor_view<8x8xf32>) pto.tstore ins(%s10 : !pto.tile_buf) - outs(%dst2 : memref<8x8xf32, #pto.address_space>) + outs(%dst2 : !pto.partition_tensor_view<8x8xf32>) pto.tstore ins(%s11 : !pto.tile_buf) - outs(%dst3 : memref<8x8xf32, #pto.address_space>) + outs(%dst3 : !pto.partition_tensor_view<8x8xf32>) return } } diff --git a/test/lit/pto/subview_validshape_guard.pto b/test/lit/pto/subview_validshape_guard.pto index 2720df1937..c8a6853103 100644 --- a/test/lit/pto/subview_validshape_guard.pto +++ b/test/lit/pto/subview_validshape_guard.pto @@ -2,7 +2,7 @@ module { func.func @subview_validshape_noncompact_default( - %dst: memref<2x8xf32, #pto.address_space>) { + %dst: !pto.partition_tensor_view<2x8xf32>) { %c2 = arith.constant 2 : index %c8 = arith.constant 8 : index @@ -12,12 +12,12 @@ module { -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<2x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<2x8xf32>) return } func.func @subview_validshape_compact_default( - %dst: memref<2x8xf32, #pto.address_space>) { + %dst: !pto.partition_tensor_view<2x8xf32>) { %c0 = arith.constant 0 : index %c2 = arith.constant 2 : index @@ -27,12 +27,12 @@ module { -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<2x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<2x8xf32>) return } func.func @subview_validshape_explicit_override( - %dst: memref<2x8xf32, #pto.address_space>) { + %dst: !pto.partition_tensor_view<1x1xf32>) { %c1 = arith.constant 1 : index %c2 = arith.constant 2 : index %c8 = arith.constant 8 : index @@ -43,13 +43,13 @@ module { -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<2x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<1x1xf32>) return } func.func @subview_validshape_explicit_dynamic( %vr: index, %vc: index, - %dst: memref<2x8xf32, #pto.address_space>) { + %dst: !pto.partition_tensor_view) { %c2 = arith.constant 2 : index %c8 = arith.constant 8 : index @@ -59,7 +59,7 @@ module { -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<2x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view) return } } diff --git a/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto b/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto index 973ef5ee60..68c9faa390 100644 --- a/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto +++ b/test/lit/pto/subview_validshape_partial_parent_requires_explicit.pto @@ -2,7 +2,7 @@ module { func.func @subview_implicit_valid_partial_parent_allowed( - %dst: memref<2x8xf32, #pto.address_space>) { + %dst: !pto.partition_tensor_view<2x8xf32>) { %c2 = arith.constant 2 : index %c8 = arith.constant 8 : index %tile = pto.alloc_tile : !pto.tile_buf @@ -10,7 +10,7 @@ module { !pto.tile_buf -> !pto.tile_buf pto.tstore ins(%s : !pto.tile_buf) - outs(%dst : memref<2x8xf32, #pto.address_space>) + outs(%dst : !pto.partition_tensor_view<2x8xf32>) return } } diff --git a/test/lit/pto/tci_i16_emitc.pto b/test/lit/pto/tci_i16_emitc.pto index b4b072427c..cb7f1558c6 100644 --- a/test/lit/pto/tci_i16_emitc.pto +++ b/test/lit/pto/tci_i16_emitc.pto @@ -1,13 +1,11 @@ // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s --check-prefix=A3 module { - func.func @tci_i16_kernel(%dst: memref<16xi16, #pto.address_space>) { + func.func @tci_i16_kernel(%dst: !pto.partition_tensor_view<1x16xi16>) { %c0_i16 = arith.constant 0 : i16 - %c0 = arith.constant 0 : index - %src = memref.reinterpret_cast %dst to offset: [0], sizes: [1, 16], strides: [16, 1] {layout = #pto.layout} : memref<16xi16, #pto.address_space> to memref<1x16xi16, strided<[16, 1], offset: ?>, #pto.address_space> %tile = pto.alloc_tile : !pto.tile_buf pto.tci ins(%c0_i16 : i16) outs(%tile : !pto.tile_buf) - pto.tstore ins(%tile : !pto.tile_buf) outs(%src : memref<1x16xi16, strided<[16, 1], offset: ?>, #pto.address_space>) {layout = #pto.layout, pto.inferred_layout = true} + pto.tstore ins(%tile : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x16xi16>) {layout = #pto.layout, pto.inferred_layout = true} return } } diff --git a/test/lit/pto/tci_ui16_emitc.pto b/test/lit/pto/tci_ui16_emitc.pto index 17c7eec7f1..aba69bf2f0 100644 --- a/test/lit/pto/tci_ui16_emitc.pto +++ b/test/lit/pto/tci_ui16_emitc.pto @@ -1,14 +1,12 @@ // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s --check-prefix=A3 module { - func.func @tci_ui16_kernel(%dst: memref<16xui16, #pto.address_space>) { + func.func @tci_ui16_kernel(%dst: !pto.partition_tensor_view<1x16xui16>) { %c7_i16 = arith.constant 7 : i16 %c7_ui16 = builtin.unrealized_conversion_cast %c7_i16 : i16 to ui16 - %c0 = arith.constant 0 : index - %src = memref.reinterpret_cast %dst to offset: [0], sizes: [1, 16], strides: [16, 1] {layout = #pto.layout} : memref<16xui16, #pto.address_space> to memref<1x16xui16, strided<[16, 1], offset: ?>, #pto.address_space> %tile = pto.alloc_tile : !pto.tile_buf pto.tci ins(%c7_ui16 : ui16) outs(%tile : !pto.tile_buf) - pto.tstore ins(%tile : !pto.tile_buf) outs(%src : memref<1x16xui16, strided<[16, 1], offset: ?>, #pto.address_space>) {layout = #pto.layout, pto.inferred_layout = true} + pto.tstore ins(%tile : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x16xui16>) {layout = #pto.layout, pto.inferred_layout = true} return } } diff --git a/test/lit/pto/tci_ui32_emitc.pto b/test/lit/pto/tci_ui32_emitc.pto index e36175500f..ef6cc2a31c 100644 --- a/test/lit/pto/tci_ui32_emitc.pto +++ b/test/lit/pto/tci_ui32_emitc.pto @@ -1,14 +1,12 @@ // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s --check-prefix=A3 module { - func.func @tci_ui32_kernel(%dst: memref<32xui32, #pto.address_space>) { + func.func @tci_ui32_kernel(%dst: !pto.partition_tensor_view<1x32xui32>) { %c5_i32 = arith.constant 5 : i32 %c5_ui32 = builtin.unrealized_conversion_cast %c5_i32 : i32 to ui32 - %c0 = arith.constant 0 : index - %src = memref.reinterpret_cast %dst to offset: [0], sizes: [1, 32], strides: [32, 1] {layout = #pto.layout} : memref<32xui32, #pto.address_space> to memref<1x32xui32, strided<[32, 1], offset: ?>, #pto.address_space> %tile = pto.alloc_tile : !pto.tile_buf pto.tci ins(%c5_ui32 : ui32) outs(%tile : !pto.tile_buf) - pto.tstore ins(%tile : !pto.tile_buf) outs(%src : memref<1x32xui32, strided<[32, 1], offset: ?>, #pto.address_space>) {layout = #pto.layout, pto.inferred_layout = true} + pto.tstore ins(%tile : !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xui32>) {layout = #pto.layout, pto.inferred_layout = true} return } } diff --git a/test/lit/pto/tdequant.pto b/test/lit/pto/tdequant.pto index 7322d7f9bf..aaefc3bbf8 100644 --- a/test/lit/pto/tdequant.pto +++ b/test/lit/pto/tdequant.pto @@ -3,29 +3,29 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { func.func @tdequant_i16() attributes {pto.kernel_kind = #pto.kernel_kind} { - %src = memref.alloc() : memref<16x16xi16, #pto.address_space> - %scale = memref.alloc() : memref<16x8xf32, #pto.address_space> - %offset = memref.alloc() : memref<16x8xf32, #pto.address_space> - %dst = memref.alloc() : memref<16x16xf32, #pto.address_space> + %src = pto.alloc_tile : !pto.tile_buf + %scale = pto.alloc_tile : !pto.tile_buf + %offset = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf - pto.tdequant ins(%src, %scale, %offset : memref<16x16xi16, #pto.address_space>, - memref<16x8xf32, #pto.address_space>, - memref<16x8xf32, #pto.address_space>) - outs(%dst : memref<16x16xf32, #pto.address_space>) + pto.tdequant ins(%src, %scale, %offset : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } func.func @tdequant_i8() attributes {pto.kernel_kind = #pto.kernel_kind} { - %src = memref.alloc() : memref<16x32xi8, #pto.address_space> - %scale = memref.alloc() : memref<16x8xf32, #pto.address_space> - %offset = memref.alloc() : memref<16x8xf32, #pto.address_space> - %dst = memref.alloc() : memref<16x32xf32, #pto.address_space> - - pto.tdequant ins(%src, %scale, %offset : memref<16x32xi8, #pto.address_space>, - memref<16x8xf32, #pto.address_space>, - memref<16x8xf32, #pto.address_space>) - outs(%dst : memref<16x32xf32, #pto.address_space>) + %src = pto.alloc_tile : !pto.tile_buf + %scale = pto.alloc_tile : !pto.tile_buf + %offset = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + + pto.tdequant ins(%src, %scale, %offset : !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) + outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/pto/textract_a5_scaling_pipe_selection.pto b/test/lit/pto/textract_a5_scaling_pipe_selection.pto index 2fc6b462f5..4532a7428e 100644 --- a/test/lit/pto/textract_a5_scaling_pipe_selection.pto +++ b/test/lit/pto/textract_a5_scaling_pipe_selection.pto @@ -3,12 +3,12 @@ module attributes {"pto.target_arch" = "a5"} { // A5: textract mat->scaling follows TEXTRACT_M2LR data-movement path and // must be treated as PIPE_MTE1 in sync insertion. - func.func @textract_mat_scaling_sync(%in: memref<32x32xf16, #pto.address_space>) { + func.func @textract_mat_scaling_sync(%in: !pto.partition_tensor_view<32x32xf16>) { %c0 = arith.constant 0 : index %src = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%in : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%in : !pto.partition_tensor_view<32x32xf16>) outs(%src : !pto.tile_buf) pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) outs(%dst : !pto.tile_buf) diff --git a/test/lit/pto/textract_a5_scaling_pipe_selection_gss.pto b/test/lit/pto/textract_a5_scaling_pipe_selection_gss.pto index 92ea332623..1540f1f1eb 100644 --- a/test/lit/pto/textract_a5_scaling_pipe_selection_gss.pto +++ b/test/lit/pto/textract_a5_scaling_pipe_selection_gss.pto @@ -3,12 +3,12 @@ module attributes {"pto.target_arch" = "a5"} { // A5: textract mat->scaling follows TEXTRACT_M2LR data-movement path and // must be treated as PIPE_MTE1 in sync insertion. - func.func @textract_mat_scaling_sync(%in: memref<32x32xf16, #pto.address_space>) { + func.func @textract_mat_scaling_sync(%in: !pto.partition_tensor_view<32x32xf16>) { %c0 = arith.constant 0 : index %src = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%in : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%in : !pto.partition_tensor_view<32x32xf16>) outs(%src : !pto.tile_buf) pto.textract ins(%src, %c0, %c0 : !pto.tile_buf, index, index) outs(%dst : !pto.tile_buf) diff --git a/test/lit/pto/tfillpad_a5_mat_pipe_selection.pto b/test/lit/pto/tfillpad_a5_mat_pipe_selection.pto index 866559923c..46f5ee14a9 100644 --- a/test/lit/pto/tfillpad_a5_mat_pipe_selection.pto +++ b/test/lit/pto/tfillpad_a5_mat_pipe_selection.pto @@ -3,7 +3,7 @@ module attributes {"pto.target_arch" = "a5"} { // A5: mat->mat TFILLPAD follows the cube data-movement path and must not be // treated as PIPE_V, otherwise sync insertion emits invalid V pipe events. - func.func @tfillpad_mat_sync(%in: memref<32x16xf16, #pto.address_space>) { + func.func @tfillpad_mat_sync(%in: !pto.partition_tensor_view<32x16xf16>) { %src = pto.alloc_tile : !pto.tile_buf @@ -11,7 +11,7 @@ module attributes {"pto.target_arch" = "a5"} { : !pto.tile_buf - pto.tload ins(%in : memref<32x16xf16, #pto.address_space>) + pto.tload ins(%in : !pto.partition_tensor_view<32x16xf16>) outs(%src : !pto.tile_buf) pto.tfillpad ins(%src : !pto.tile_buf>, - %src: memref<32x16xf32, #pto.address_space>, - %out: memref<32x32xf32, #pto.address_space>) { + %x: !pto.partition_tensor_view<32x32xf32>, + %src: !pto.partition_tensor_view<32x16xf32>, + %out: !pto.partition_tensor_view<32x32xf32>) { %c0 = arith.constant 0 : index %c16 = arith.constant 16 : index %c32 = arith.constant 32 : index @@ -16,9 +16,9 @@ module attributes {"pto.target_arch" = "a5"} { %tile_src = pto.alloc_tile addr = %addr1 : !pto.tile_buf - pto.tload ins(%x : memref<32x32xf32, #pto.address_space>) + pto.tload ins(%x : !pto.partition_tensor_view<32x32xf32>) outs(%tile_x : !pto.tile_buf) - pto.tload ins(%src : memref<32x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<32x16xf32>) outs(%tile_src : !pto.tile_buf) %view = pto.subview %tile_x[%c0, %c16] sizes [32, 16] valid [%c32, %c16] : !pto.tile_buf -> @@ -28,7 +28,7 @@ module attributes {"pto.target_arch" = "a5"} { outs(%view : !pto.tile_buf) pto.tstore ins(%tile_x : !pto.tile_buf) - outs(%out : memref<32x32xf32, #pto.address_space>) + outs(%out : !pto.partition_tensor_view<32x32xf32>) return } } diff --git a/test/lit/pto/tinsert_a2a3_vec_vec_pipe_selection.pto b/test/lit/pto/tinsert_a2a3_vec_vec_pipe_selection.pto index 20f1a6051d..e7a89ee59c 100644 --- a/test/lit/pto/tinsert_a2a3_vec_vec_pipe_selection.pto +++ b/test/lit/pto/tinsert_a2a3_vec_vec_pipe_selection.pto @@ -1,18 +1,18 @@ // RUN: ptoas --pto-arch=a3 --enable-insert-sync %s | FileCheck %s module attributes {"pto.target_arch" = "a3"} { - func.func @tinsert_a2a3_vec_vec_pipeline(%a: memref<32x32xf16, #pto.address_space>) { + func.func @tinsert_a2a3_vec_vec_pipeline(%a: !pto.partition_tensor_view<32x32xf16>) { %c0 = arith.constant 0 : index - %src_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> - %dst_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> + %src_vec = pto.alloc_tile : !pto.tile_buf + %dst_vec = pto.alloc_tile : !pto.tile_buf + %out_vec = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%a : memref<32x32xf16, #pto.address_space>) - outs(%src_vec : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} - pto.tinsert ins(%src_vec, %c0, %c0 : memref<32x32xf16, #pto.address_space>, index, index) - outs(%dst_vec : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_vec : memref<32x32xf16, #pto.address_space>) - outs(%out_vec : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%a : !pto.partition_tensor_view<32x32xf16>) + outs(%src_vec : !pto.tile_buf) {layout = #pto.layout} + pto.tinsert ins(%src_vec, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst_vec : !pto.tile_buf) + pto.tmov ins(%dst_vec : !pto.tile_buf) + outs(%out_vec : !pto.tile_buf) return } } diff --git a/test/lit/pto/tinsert_a2a3_vec_vec_pipe_selection_gss.pto b/test/lit/pto/tinsert_a2a3_vec_vec_pipe_selection_gss.pto index a3bf500c44..87cdef0842 100644 --- a/test/lit/pto/tinsert_a2a3_vec_vec_pipe_selection_gss.pto +++ b/test/lit/pto/tinsert_a2a3_vec_vec_pipe_selection_gss.pto @@ -1,18 +1,18 @@ // RUN: ptoas --pto-arch=a3 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %s | FileCheck %s module attributes {"pto.target_arch" = "a3"} { - func.func @tinsert_a2a3_vec_vec_pipeline(%a: memref<32x32xf16, #pto.address_space>) { + func.func @tinsert_a2a3_vec_vec_pipeline(%a: !pto.partition_tensor_view<32x32xf16>) { %c0 = arith.constant 0 : index - %src_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> - %dst_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> + %src_vec = pto.alloc_tile : !pto.tile_buf + %dst_vec = pto.alloc_tile : !pto.tile_buf + %out_vec = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%a : memref<32x32xf16, #pto.address_space>) - outs(%src_vec : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} - pto.tinsert ins(%src_vec, %c0, %c0 : memref<32x32xf16, #pto.address_space>, index, index) - outs(%dst_vec : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_vec : memref<32x32xf16, #pto.address_space>) - outs(%out_vec : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%a : !pto.partition_tensor_view<32x32xf16>) + outs(%src_vec : !pto.tile_buf) {layout = #pto.layout} + pto.tinsert ins(%src_vec, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst_vec : !pto.tile_buf) + pto.tmov ins(%dst_vec : !pto.tile_buf) + outs(%out_vec : !pto.tile_buf) return } } diff --git a/test/lit/pto/tinsert_a3_pipe_selection.pto b/test/lit/pto/tinsert_a3_pipe_selection.pto index efab133cea..6037a501f6 100644 --- a/test/lit/pto/tinsert_a3_pipe_selection.pto +++ b/test/lit/pto/tinsert_a3_pipe_selection.pto @@ -2,50 +2,51 @@ module attributes {"pto.target_arch" = "a3"} { // A3: acc->mat TINSERT uses PIPE_FIX. - func.func @tinsert_acc_mat_pipeline_a3(%a: memref<32x32xf32, #pto.address_space>, - %b: memref<32x32xf32, #pto.address_space>, - %i: memref<32x32xf16, #pto.address_space>, - %out: memref<32x32xf32, #pto.address_space>) { + func.func @tinsert_acc_mat_pipeline_a3( + %a: !pto.partition_tensor_view<32x32xf32>, + %b: !pto.partition_tensor_view<32x32xf32>, + %i: !pto.partition_tensor_view<32x32xf16>, + %out: !pto.partition_tensor_view<32x32xf32>) { %c0 = arith.constant 0 : index - %a_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %a_left = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_right = memref.alloc() : memref<32x32xf32, #pto.address_space> - %src_acc = memref.alloc() : memref<32x32xf32, #pto.address_space> + %a_mat = pto.alloc_tile : !pto.tile_buf + %b_mat = pto.alloc_tile : !pto.tile_buf + %a_left = pto.alloc_tile : !pto.tile_buf + %b_right = pto.alloc_tile : !pto.tile_buf + %src_acc = pto.alloc_tile : !pto.tile_buf - %dst_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %i_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_left = memref.alloc() : memref<32x32xf16, #pto.address_space> - %i_right = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_acc = memref.alloc() : memref<32x32xf32, #pto.address_space> + %dst_mat = pto.alloc_tile : !pto.tile_buf + %i_mat = pto.alloc_tile : !pto.tile_buf + %out_left = pto.alloc_tile : !pto.tile_buf + %i_right = pto.alloc_tile : !pto.tile_buf + %out_acc = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%a : memref<32x32xf32, #pto.address_space>) - outs(%a_mat : memref<32x32xf32, #pto.address_space>) - pto.tload ins(%b : memref<32x32xf32, #pto.address_space>) - outs(%b_mat : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%a_mat : memref<32x32xf32, #pto.address_space>) - outs(%a_left : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%b_mat : memref<32x32xf32, #pto.address_space>) - outs(%b_right : memref<32x32xf32, #pto.address_space>) - pto.tmatmul ins(%a_left, %b_right : memref<32x32xf32, #pto.address_space>, - memref<32x32xf32, #pto.address_space>) - outs(%src_acc : memref<32x32xf32, #pto.address_space>) + pto.tload ins(%a : !pto.partition_tensor_view<32x32xf32>) + outs(%a_mat : !pto.tile_buf) + pto.tload ins(%b : !pto.partition_tensor_view<32x32xf32>) + outs(%b_mat : !pto.tile_buf) + pto.tmov ins(%a_mat : !pto.tile_buf) + outs(%a_left : !pto.tile_buf) + pto.tmov ins(%b_mat : !pto.tile_buf) + outs(%b_right : !pto.tile_buf) + pto.tmatmul ins(%a_left, %b_right : !pto.tile_buf, + !pto.tile_buf) + outs(%src_acc : !pto.tile_buf) - pto.tinsert ins(%src_acc, %c0, %c0 : memref<32x32xf32, #pto.address_space>, index, index) - outs(%dst_mat : memref<32x32xf16, #pto.address_space>) + pto.tinsert ins(%src_acc, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst_mat : !pto.tile_buf) - pto.tload ins(%i : memref<32x32xf16, #pto.address_space>) - outs(%i_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_mat : memref<32x32xf16, #pto.address_space>) - outs(%out_left : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%i_mat : memref<32x32xf16, #pto.address_space>) - outs(%i_right : memref<32x32xf16, #pto.address_space>) - pto.tmatmul ins(%out_left, %i_right : memref<32x32xf16, #pto.address_space>, - memref<32x32xf16, #pto.address_space>) - outs(%out_acc : memref<32x32xf32, #pto.address_space>) - pto.tstore ins(%out_acc : memref<32x32xf32, #pto.address_space>) - outs(%out : memref<32x32xf32, #pto.address_space>) + pto.tload ins(%i : !pto.partition_tensor_view<32x32xf16>) + outs(%i_mat : !pto.tile_buf) + pto.tmov ins(%dst_mat : !pto.tile_buf) + outs(%out_left : !pto.tile_buf) + pto.tmov ins(%i_mat : !pto.tile_buf) + outs(%i_right : !pto.tile_buf) + pto.tmatmul ins(%out_left, %i_right : !pto.tile_buf, + !pto.tile_buf) + outs(%out_acc : !pto.tile_buf) + pto.tstore ins(%out_acc : !pto.tile_buf) + outs(%out : !pto.partition_tensor_view<32x32xf32>) return } } diff --git a/test/lit/pto/tinsert_a3_pipe_selection_gss.pto b/test/lit/pto/tinsert_a3_pipe_selection_gss.pto index dc1209c5b9..1453f9fe30 100644 --- a/test/lit/pto/tinsert_a3_pipe_selection_gss.pto +++ b/test/lit/pto/tinsert_a3_pipe_selection_gss.pto @@ -1,54 +1,4 @@ -// RUN: ptoas --pto-arch a3 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %s | FileCheck %s - -module attributes {"pto.target_arch" = "a3"} { - // A3: acc->mat TINSERT uses PIPE_FIX. - func.func @tinsert_acc_mat_pipeline_a3(%a: memref<32x32xf32, #pto.address_space>, - %b: memref<32x32xf32, #pto.address_space>, - %i: memref<32x32xf16, #pto.address_space>, - %out: memref<32x32xf32, #pto.address_space>) { - %c0 = arith.constant 0 : index - - %a_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %a_left = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_right = memref.alloc() : memref<32x32xf32, #pto.address_space> - %src_acc = memref.alloc() : memref<32x32xf32, #pto.address_space> - - %dst_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %i_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_left = memref.alloc() : memref<32x32xf16, #pto.address_space> - %i_right = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_acc = memref.alloc() : memref<32x32xf32, #pto.address_space> - - pto.tload ins(%a : memref<32x32xf32, #pto.address_space>) - outs(%a_mat : memref<32x32xf32, #pto.address_space>) - pto.tload ins(%b : memref<32x32xf32, #pto.address_space>) - outs(%b_mat : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%a_mat : memref<32x32xf32, #pto.address_space>) - outs(%a_left : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%b_mat : memref<32x32xf32, #pto.address_space>) - outs(%b_right : memref<32x32xf32, #pto.address_space>) - pto.tmatmul ins(%a_left, %b_right : memref<32x32xf32, #pto.address_space>, - memref<32x32xf32, #pto.address_space>) - outs(%src_acc : memref<32x32xf32, #pto.address_space>) - - pto.tinsert ins(%src_acc, %c0, %c0 : memref<32x32xf32, #pto.address_space>, index, index) - outs(%dst_mat : memref<32x32xf16, #pto.address_space>) - - pto.tload ins(%i : memref<32x32xf16, #pto.address_space>) - outs(%i_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_mat : memref<32x32xf16, #pto.address_space>) - outs(%out_left : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%i_mat : memref<32x32xf16, #pto.address_space>) - outs(%i_right : memref<32x32xf16, #pto.address_space>) - pto.tmatmul ins(%out_left, %i_right : memref<32x32xf16, #pto.address_space>, - memref<32x32xf16, #pto.address_space>) - outs(%out_acc : memref<32x32xf32, #pto.address_space>) - pto.tstore ins(%out_acc : memref<32x32xf32, #pto.address_space>) - outs(%out : memref<32x32xf32, #pto.address_space>) - return - } -} +// RUN: ptoas --pto-arch a3 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %S/tinsert_a3_pipe_selection.pto | FileCheck %s // CHECK-LABEL: AICORE void tinsert_acc_mat_pipeline_a3( // CHECK: set_flag(PIPE_M, PIPE_FIX, EVENT_ID{{[0-9]+}}); diff --git a/test/lit/pto/tinsert_a5_pipe_selection.pto b/test/lit/pto/tinsert_a5_pipe_selection.pto index 75f48acaad..c8ca4d3821 100644 --- a/test/lit/pto/tinsert_a5_pipe_selection.pto +++ b/test/lit/pto/tinsert_a5_pipe_selection.pto @@ -2,50 +2,51 @@ module attributes {"pto.device-spec" = "Ascend950"} { // A5: acc->mat TINSERT must use PIPE_FIX (regular A2M path). - func.func @tinsert_acc_mat_pipeline(%a: memref<32x32xf32, #pto.address_space>, - %b: memref<32x32xf32, #pto.address_space>, - %i: memref<32x32xf16, #pto.address_space>, - %out: memref<32x32xf32, #pto.address_space>) { + func.func @tinsert_acc_mat_pipeline( + %a: !pto.partition_tensor_view<32x32xf32>, + %b: !pto.partition_tensor_view<32x32xf32>, + %i: !pto.partition_tensor_view<32x32xf16>, + %out: !pto.partition_tensor_view<32x32xf32>) { %c0 = arith.constant 0 : index - %a_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %a_left = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_right = memref.alloc() : memref<32x32xf32, #pto.address_space> - %src_acc = memref.alloc() : memref<32x32xf32, #pto.address_space> + %a_mat = pto.alloc_tile : !pto.tile_buf + %b_mat = pto.alloc_tile : !pto.tile_buf + %a_left = pto.alloc_tile : !pto.tile_buf + %b_right = pto.alloc_tile : !pto.tile_buf + %src_acc = pto.alloc_tile : !pto.tile_buf - %dst_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %i_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_left = memref.alloc() : memref<32x32xf16, #pto.address_space> - %i_right = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_acc = memref.alloc() : memref<32x32xf32, #pto.address_space> + %dst_mat = pto.alloc_tile : !pto.tile_buf + %i_mat = pto.alloc_tile : !pto.tile_buf + %out_left = pto.alloc_tile : !pto.tile_buf + %i_right = pto.alloc_tile : !pto.tile_buf + %out_acc = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%a : memref<32x32xf32, #pto.address_space>) - outs(%a_mat : memref<32x32xf32, #pto.address_space>) - pto.tload ins(%b : memref<32x32xf32, #pto.address_space>) - outs(%b_mat : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%a_mat : memref<32x32xf32, #pto.address_space>) - outs(%a_left : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%b_mat : memref<32x32xf32, #pto.address_space>) - outs(%b_right : memref<32x32xf32, #pto.address_space>) - pto.tmatmul ins(%a_left, %b_right : memref<32x32xf32, #pto.address_space>, - memref<32x32xf32, #pto.address_space>) - outs(%src_acc : memref<32x32xf32, #pto.address_space>) + pto.tload ins(%a : !pto.partition_tensor_view<32x32xf32>) + outs(%a_mat : !pto.tile_buf) + pto.tload ins(%b : !pto.partition_tensor_view<32x32xf32>) + outs(%b_mat : !pto.tile_buf) + pto.tmov ins(%a_mat : !pto.tile_buf) + outs(%a_left : !pto.tile_buf) + pto.tmov ins(%b_mat : !pto.tile_buf) + outs(%b_right : !pto.tile_buf) + pto.tmatmul ins(%a_left, %b_right : !pto.tile_buf, + !pto.tile_buf) + outs(%src_acc : !pto.tile_buf) - pto.tinsert ins(%src_acc, %c0, %c0 : memref<32x32xf32, #pto.address_space>, index, index) - outs(%dst_mat : memref<32x32xf16, #pto.address_space>) + pto.tinsert ins(%src_acc, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst_mat : !pto.tile_buf) - pto.tload ins(%i : memref<32x32xf16, #pto.address_space>) - outs(%i_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_mat : memref<32x32xf16, #pto.address_space>) - outs(%out_left : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%i_mat : memref<32x32xf16, #pto.address_space>) - outs(%i_right : memref<32x32xf16, #pto.address_space>) - pto.tmatmul ins(%out_left, %i_right : memref<32x32xf16, #pto.address_space>, - memref<32x32xf16, #pto.address_space>) - outs(%out_acc : memref<32x32xf32, #pto.address_space>) - pto.tstore ins(%out_acc : memref<32x32xf32, #pto.address_space>) - outs(%out : memref<32x32xf32, #pto.address_space>) + pto.tload ins(%i : !pto.partition_tensor_view<32x32xf16>) + outs(%i_mat : !pto.tile_buf) + pto.tmov ins(%dst_mat : !pto.tile_buf) + outs(%out_left : !pto.tile_buf) + pto.tmov ins(%i_mat : !pto.tile_buf) + outs(%i_right : !pto.tile_buf) + pto.tmatmul ins(%out_left, %i_right : !pto.tile_buf, + !pto.tile_buf) + outs(%out_acc : !pto.tile_buf) + pto.tstore ins(%out_acc : !pto.tile_buf) + outs(%out : !pto.partition_tensor_view<32x32xf32>) return } } diff --git a/test/lit/pto/tinsert_a5_pipe_selection_gss.pto b/test/lit/pto/tinsert_a5_pipe_selection_gss.pto index 986499d675..3b3f6c0504 100644 --- a/test/lit/pto/tinsert_a5_pipe_selection_gss.pto +++ b/test/lit/pto/tinsert_a5_pipe_selection_gss.pto @@ -1,54 +1,4 @@ -// RUN: ptoas --pto-arch a5 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %s | FileCheck %s - -module attributes {"pto.device-spec" = "Ascend950"} { - // A5: acc->mat TINSERT must use PIPE_FIX (regular A2M path). - func.func @tinsert_acc_mat_pipeline(%a: memref<32x32xf32, #pto.address_space>, - %b: memref<32x32xf32, #pto.address_space>, - %i: memref<32x32xf16, #pto.address_space>, - %out: memref<32x32xf32, #pto.address_space>) { - %c0 = arith.constant 0 : index - - %a_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_mat = memref.alloc() : memref<32x32xf32, #pto.address_space> - %a_left = memref.alloc() : memref<32x32xf32, #pto.address_space> - %b_right = memref.alloc() : memref<32x32xf32, #pto.address_space> - %src_acc = memref.alloc() : memref<32x32xf32, #pto.address_space> - - %dst_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %i_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_left = memref.alloc() : memref<32x32xf16, #pto.address_space> - %i_right = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_acc = memref.alloc() : memref<32x32xf32, #pto.address_space> - - pto.tload ins(%a : memref<32x32xf32, #pto.address_space>) - outs(%a_mat : memref<32x32xf32, #pto.address_space>) - pto.tload ins(%b : memref<32x32xf32, #pto.address_space>) - outs(%b_mat : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%a_mat : memref<32x32xf32, #pto.address_space>) - outs(%a_left : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%b_mat : memref<32x32xf32, #pto.address_space>) - outs(%b_right : memref<32x32xf32, #pto.address_space>) - pto.tmatmul ins(%a_left, %b_right : memref<32x32xf32, #pto.address_space>, - memref<32x32xf32, #pto.address_space>) - outs(%src_acc : memref<32x32xf32, #pto.address_space>) - - pto.tinsert ins(%src_acc, %c0, %c0 : memref<32x32xf32, #pto.address_space>, index, index) - outs(%dst_mat : memref<32x32xf16, #pto.address_space>) - - pto.tload ins(%i : memref<32x32xf16, #pto.address_space>) - outs(%i_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_mat : memref<32x32xf16, #pto.address_space>) - outs(%out_left : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%i_mat : memref<32x32xf16, #pto.address_space>) - outs(%i_right : memref<32x32xf16, #pto.address_space>) - pto.tmatmul ins(%out_left, %i_right : memref<32x32xf16, #pto.address_space>, - memref<32x32xf16, #pto.address_space>) - outs(%out_acc : memref<32x32xf32, #pto.address_space>) - pto.tstore ins(%out_acc : memref<32x32xf32, #pto.address_space>) - outs(%out : memref<32x32xf32, #pto.address_space>) - return - } -} +// RUN: ptoas --pto-arch a5 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %S/tinsert_a5_pipe_selection.pto | FileCheck %s // CHECK-LABEL: AICORE void tinsert_acc_mat_pipeline( // CHECK: set_flag(PIPE_M, PIPE_FIX, EVENT_ID{{[0-9]+}}); diff --git a/test/lit/pto/tinsert_a5_vec_mat_mode_lowering.pto b/test/lit/pto/tinsert_a5_vec_mat_mode_lowering.pto index 61aec5ca96..114c12f1fe 100644 --- a/test/lit/pto/tinsert_a5_vec_mat_mode_lowering.pto +++ b/test/lit/pto/tinsert_a5_vec_mat_mode_lowering.pto @@ -1,22 +1,22 @@ // RUN: ptoas --pto-arch=a5 %s | FileCheck %s module attributes {"pto.target_arch" = "a5"} { - func.func @tinsert_vec_mat_nd(%src_gm: memref<32x32xf16, #pto.address_space>) { + func.func @tinsert_vec_mat_nd(%src_gm: !pto.partition_tensor_view<32x32xf16>) { %c0 = arith.constant 0 : index %src = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src_gm : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%src_gm : !pto.partition_tensor_view<32x32xf16>) outs(%src : !pto.tile_buf) {layout = #pto.layout} pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index) outs(%dst : !pto.tile_buf) return } - func.func @tinsert_vec_mat_nz(%src_gm: memref<32x32xf16, #pto.address_space>) { + func.func @tinsert_vec_mat_nz(%src_gm: !pto.partition_tensor_view<32x32xf16>) { %c0 = arith.constant 0 : index %src = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src_gm : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%src_gm : !pto.partition_tensor_view<32x32xf16>) outs(%src : !pto.tile_buf) {layout = #pto.layout} pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index) outs(%dst : !pto.tile_buf) diff --git a/test/lit/pto/tinsert_a5_vec_pipe_selection.pto b/test/lit/pto/tinsert_a5_vec_pipe_selection.pto index 7ac4de225e..11a1626835 100644 --- a/test/lit/pto/tinsert_a5_vec_pipe_selection.pto +++ b/test/lit/pto/tinsert_a5_vec_pipe_selection.pto @@ -2,18 +2,18 @@ module attributes {"pto.device-spec" = "Ascend950"} { // A5: vec->vec TINSERT keeps PIPE_V path and no explicit mode template. - func.func @tinsert_vec_vec_pipeline(%a: memref<32x32xf16, #pto.address_space>) { + func.func @tinsert_vec_vec_pipeline(%a: !pto.partition_tensor_view<32x32xf16>) { %c0 = arith.constant 0 : index - %src_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> - %dst_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> + %src_vec = pto.alloc_tile : !pto.tile_buf + %dst_vec = pto.alloc_tile : !pto.tile_buf + %out_vec = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%a : memref<32x32xf16, #pto.address_space>) - outs(%src_vec : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} - pto.tinsert ins(%src_vec, %c0, %c0 : memref<32x32xf16, #pto.address_space>, index, index) - outs(%dst_vec : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_vec : memref<32x32xf16, #pto.address_space>) - outs(%out_vec : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%a : !pto.partition_tensor_view<32x32xf16>) + outs(%src_vec : !pto.tile_buf) {layout = #pto.layout} + pto.tinsert ins(%src_vec, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst_vec : !pto.tile_buf) + pto.tmov ins(%dst_vec : !pto.tile_buf) + outs(%out_vec : !pto.tile_buf) return } } diff --git a/test/lit/pto/tinsert_a5_vec_pipe_selection_gss.pto b/test/lit/pto/tinsert_a5_vec_pipe_selection_gss.pto index 2d6b9063bf..ade805ae73 100644 --- a/test/lit/pto/tinsert_a5_vec_pipe_selection_gss.pto +++ b/test/lit/pto/tinsert_a5_vec_pipe_selection_gss.pto @@ -2,18 +2,18 @@ module attributes {"pto.device-spec" = "Ascend950"} { // A5: vec->vec TINSERT keeps PIPE_V path and no explicit mode template. - func.func @tinsert_vec_vec_pipeline(%a: memref<32x32xf16, #pto.address_space>) { + func.func @tinsert_vec_vec_pipeline(%a: !pto.partition_tensor_view<32x32xf16>) { %c0 = arith.constant 0 : index - %src_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> - %dst_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> - %out_vec = memref.alloc() : memref<32x32xf16, #pto.address_space> + %src_vec = pto.alloc_tile : !pto.tile_buf + %dst_vec = pto.alloc_tile : !pto.tile_buf + %out_vec = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%a : memref<32x32xf16, #pto.address_space>) - outs(%src_vec : memref<32x32xf16, #pto.address_space>) {layout = #pto.layout} - pto.tinsert ins(%src_vec, %c0, %c0 : memref<32x32xf16, #pto.address_space>, index, index) - outs(%dst_vec : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_vec : memref<32x32xf16, #pto.address_space>) - outs(%out_vec : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%a : !pto.partition_tensor_view<32x32xf16>) + outs(%src_vec : !pto.tile_buf) {layout = #pto.layout} + pto.tinsert ins(%src_vec, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst_vec : !pto.tile_buf) + pto.tmov ins(%dst_vec : !pto.tile_buf) + outs(%out_vec : !pto.tile_buf) return } } diff --git a/test/lit/pto/tinsert_level3_tile_native_sync.pto b/test/lit/pto/tinsert_level3_tile_native_sync.pto index 255c4ab7e0..a3089ab526 100644 --- a/test/lit/pto/tinsert_level3_tile_native_sync.pto +++ b/test/lit/pto/tinsert_level3_tile_native_sync.pto @@ -2,9 +2,9 @@ module attributes {"pto.target_arch" = "a5"} { func.func @tinsert_level3_tile_native_sync( - %x: memref<32x32xf32, #pto.address_space>, - %src: memref<32x16xf32, #pto.address_space>, - %out: memref<32x32xf32, #pto.address_space>) { + %x: !pto.partition_tensor_view<32x32xf32>, + %src: !pto.partition_tensor_view<32x16xf32>, + %out: !pto.partition_tensor_view<32x32xf32>) { %c0 = arith.constant 0 : index %addr0 = arith.constant 0 : i64 %addr1 = arith.constant 4096 : i64 @@ -14,16 +14,16 @@ module attributes {"pto.target_arch" = "a5"} { %tile_src = pto.alloc_tile addr = %addr1 : !pto.tile_buf - pto.tload ins(%x : memref<32x32xf32, #pto.address_space>) + pto.tload ins(%x : !pto.partition_tensor_view<32x32xf32>) outs(%tile_x : !pto.tile_buf) - pto.tload ins(%src : memref<32x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<32x16xf32>) outs(%tile_src : !pto.tile_buf) pto.tinsert ins(%tile_src, %c0, %c0 : !pto.tile_buf, index, index) outs(%tile_x : !pto.tile_buf) pto.tstore ins(%tile_x : !pto.tile_buf) - outs(%out : memref<32x32xf32, #pto.address_space>) + outs(%out : !pto.partition_tensor_view<32x32xf32>) return } } diff --git a/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto b/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto index 84fbf30052..6259007a18 100644 --- a/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto +++ b/test/lit/pto/tload_tprefetch_low_precision_a5_valid.pto @@ -23,7 +23,7 @@ module { } } -// CHECK: func.func @tload_tprefetch_low_precision_a5_valid(%arg0: memref<16x16xf8E4M3FN>, %arg1: memref<16x16x!pto.hif8>) +// CHECK: func.func @tload_tprefetch_low_precision_a5_valid(%arg0: !pto.partition_tensor_view<16x16xf8E4M3FN>, %arg1: !pto.partition_tensor_view<16x16x!pto.hif8>) // CHECK: pto.declare_tile -> !pto.tile_buf -// CHECK: pto.tload ins(%arg0 : memref<16x16xf8E4M3FN>) outs( -// CHECK: pto.tprefetch ins(%arg1 : memref<16x16x!pto.hif8>) outs( +// CHECK: pto.tload ins(%arg0 : !pto.partition_tensor_view<16x16xf8E4M3FN>) outs( +// CHECK: pto.tprefetch ins(%arg1 : !pto.partition_tensor_view<16x16x!pto.hif8>) outs( diff --git a/test/lit/pto/tmov_acc_mat_pipe_selection.pto b/test/lit/pto/tmov_acc_mat_pipe_selection.pto index b44604e712..3a079960fe 100644 --- a/test/lit/pto/tmov_acc_mat_pipe_selection.pto +++ b/test/lit/pto/tmov_acc_mat_pipe_selection.pto @@ -2,31 +2,32 @@ module attributes {"pto.device-spec" = "Ascend950"} { // A5: tmov acc->mat should be treated as PIPE_FIX for sync insertion. - func.func @tmov_acc_mat_pipeline(%a: memref<32x32xf16, #pto.address_space>, - %b: memref<32x32xf16, #pto.address_space>) { - %a_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %b_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %a_left = memref.alloc() : memref<32x32xf16, #pto.address_space> - %b_right = memref.alloc() : memref<32x32xf16, #pto.address_space> - %acc = memref.alloc() : memref<32x32xf32, #pto.address_space> - %dst_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %dst_left = memref.alloc() : memref<32x32xf16, #pto.address_space> + func.func @tmov_acc_mat_pipeline( + %a: !pto.partition_tensor_view<32x32xf16>, + %b: !pto.partition_tensor_view<32x32xf16>) { + %a_mat = pto.alloc_tile : !pto.tile_buf + %b_mat = pto.alloc_tile : !pto.tile_buf + %a_left = pto.alloc_tile : !pto.tile_buf + %b_right = pto.alloc_tile : !pto.tile_buf + %acc = pto.alloc_tile : !pto.tile_buf + %dst_mat = pto.alloc_tile : !pto.tile_buf + %dst_left = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%a : memref<32x32xf16, #pto.address_space>) - outs(%a_mat : memref<32x32xf16, #pto.address_space>) - pto.tload ins(%b : memref<32x32xf16, #pto.address_space>) - outs(%b_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%a_mat : memref<32x32xf16, #pto.address_space>) - outs(%a_left : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%b_mat : memref<32x32xf16, #pto.address_space>) - outs(%b_right : memref<32x32xf16, #pto.address_space>) - pto.tmatmul ins(%a_left, %b_right : memref<32x32xf16, #pto.address_space>, - memref<32x32xf16, #pto.address_space>) - outs(%acc : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%acc : memref<32x32xf32, #pto.address_space>) - outs(%dst_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_mat : memref<32x32xf16, #pto.address_space>) - outs(%dst_left : memref<32x32xf16, #pto.address_space>) + pto.tload ins(%a : !pto.partition_tensor_view<32x32xf16>) + outs(%a_mat : !pto.tile_buf) + pto.tload ins(%b : !pto.partition_tensor_view<32x32xf16>) + outs(%b_mat : !pto.tile_buf) + pto.tmov ins(%a_mat : !pto.tile_buf) + outs(%a_left : !pto.tile_buf) + pto.tmov ins(%b_mat : !pto.tile_buf) + outs(%b_right : !pto.tile_buf) + pto.tmatmul ins(%a_left, %b_right : !pto.tile_buf, + !pto.tile_buf) + outs(%acc : !pto.tile_buf) + pto.tmov ins(%acc : !pto.tile_buf) + outs(%dst_mat : !pto.tile_buf) + pto.tmov ins(%dst_mat : !pto.tile_buf) + outs(%dst_left : !pto.tile_buf) return } } diff --git a/test/lit/pto/tmov_acc_mat_pipe_selection_gss.pto b/test/lit/pto/tmov_acc_mat_pipe_selection_gss.pto index e09ed36f17..c64fc9641a 100644 --- a/test/lit/pto/tmov_acc_mat_pipe_selection_gss.pto +++ b/test/lit/pto/tmov_acc_mat_pipe_selection_gss.pto @@ -1,35 +1,4 @@ -// RUN: ptoas --pto-arch a5 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %s | FileCheck %s - -module attributes {"pto.device-spec" = "Ascend950"} { - // A5: tmov acc->mat should be treated as PIPE_FIX for sync insertion. - func.func @tmov_acc_mat_pipeline(%a: memref<32x32xf16, #pto.address_space>, - %b: memref<32x32xf16, #pto.address_space>) { - %a_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %b_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %a_left = memref.alloc() : memref<32x32xf16, #pto.address_space> - %b_right = memref.alloc() : memref<32x32xf16, #pto.address_space> - %acc = memref.alloc() : memref<32x32xf32, #pto.address_space> - %dst_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %dst_left = memref.alloc() : memref<32x32xf16, #pto.address_space> - - pto.tload ins(%a : memref<32x32xf16, #pto.address_space>) - outs(%a_mat : memref<32x32xf16, #pto.address_space>) - pto.tload ins(%b : memref<32x32xf16, #pto.address_space>) - outs(%b_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%a_mat : memref<32x32xf16, #pto.address_space>) - outs(%a_left : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%b_mat : memref<32x32xf16, #pto.address_space>) - outs(%b_right : memref<32x32xf16, #pto.address_space>) - pto.tmatmul ins(%a_left, %b_right : memref<32x32xf16, #pto.address_space>, - memref<32x32xf16, #pto.address_space>) - outs(%acc : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%acc : memref<32x32xf32, #pto.address_space>) - outs(%dst_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%dst_mat : memref<32x32xf16, #pto.address_space>) - outs(%dst_left : memref<32x32xf16, #pto.address_space>) - return - } -} +// RUN: ptoas --pto-arch a5 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %S/tmov_acc_mat_pipe_selection.pto | FileCheck %s // CHECK-LABEL: AICORE void tmov_acc_mat_pipeline( // CHECK: set_flag(PIPE_M, PIPE_FIX, EVENT_ID{{[0-9]+}}); diff --git a/test/lit/pto/tmov_acc_to_vec_mode_a5_emitc.pto b/test/lit/pto/tmov_acc_to_vec_mode_a5_emitc.pto index bb6fc1ce5d..058dae6b1f 100644 --- a/test/lit/pto/tmov_acc_to_vec_mode_a5_emitc.pto +++ b/test/lit/pto/tmov_acc_to_vec_mode_a5_emitc.pto @@ -1,29 +1,30 @@ // RUN: ptoas --pto-arch a5 --enable-insert-sync %s | FileCheck %s --check-prefix=A5 // RUN: ptoas --pto-arch a3 %s | FileCheck %s --check-prefix=A3 -module attributes {"pto.device-spec" = "Ascend950"} { - func.func @tmov_acc_to_vec_mode_a5(%a: memref<32x32xf16, #pto.address_space>, - %b: memref<32x32xf16, #pto.address_space>) { - %a_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %b_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %a_left = memref.alloc() : memref<32x32xf16, #pto.address_space> - %b_right = memref.alloc() : memref<32x32xf16, #pto.address_space> - %acc = memref.alloc() : memref<32x32xf32, #pto.address_space> - %dst_vec = memref.alloc() : memref<32x32xf32, #pto.address_space> +module { + func.func @tmov_acc_to_vec_mode_a5( + %a: !pto.partition_tensor_view<32x32xf16>, + %b: !pto.partition_tensor_view<32x32xf16>) { + %a_mat = pto.alloc_tile : !pto.tile_buf + %b_mat = pto.alloc_tile : !pto.tile_buf + %a_left = pto.alloc_tile : !pto.tile_buf + %b_right = pto.alloc_tile : !pto.tile_buf + %acc = pto.alloc_tile : !pto.tile_buf + %dst_vec = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%a : memref<32x32xf16, #pto.address_space>) - outs(%a_mat : memref<32x32xf16, #pto.address_space>) - pto.tload ins(%b : memref<32x32xf16, #pto.address_space>) - outs(%b_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%a_mat : memref<32x32xf16, #pto.address_space>) - outs(%a_left : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%b_mat : memref<32x32xf16, #pto.address_space>) - outs(%b_right : memref<32x32xf16, #pto.address_space>) - pto.tmatmul ins(%a_left, %b_right : memref<32x32xf16, #pto.address_space>, - memref<32x32xf16, #pto.address_space>) - outs(%acc : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%acc : memref<32x32xf32, #pto.address_space>) - outs(%dst_vec : memref<32x32xf32, #pto.address_space>) + pto.tload ins(%a : !pto.partition_tensor_view<32x32xf16>) + outs(%a_mat : !pto.tile_buf) + pto.tload ins(%b : !pto.partition_tensor_view<32x32xf16>) + outs(%b_mat : !pto.tile_buf) + pto.tmov ins(%a_mat : !pto.tile_buf) + outs(%a_left : !pto.tile_buf) + pto.tmov ins(%b_mat : !pto.tile_buf) + outs(%b_right : !pto.tile_buf) + pto.tmatmul ins(%a_left, %b_right : !pto.tile_buf, + !pto.tile_buf) + outs(%acc : !pto.tile_buf) + pto.tmov ins(%acc : !pto.tile_buf) + outs(%dst_vec : !pto.tile_buf) {accToVecMode = #pto.acc_to_vec_mode} return } diff --git a/test/lit/pto/tmov_acc_to_vec_mode_a5_emitc_gss.pto b/test/lit/pto/tmov_acc_to_vec_mode_a5_emitc_gss.pto index c63fa14d26..be89c7ba1e 100644 --- a/test/lit/pto/tmov_acc_to_vec_mode_a5_emitc_gss.pto +++ b/test/lit/pto/tmov_acc_to_vec_mode_a5_emitc_gss.pto @@ -1,33 +1,5 @@ -// RUN: ptoas --pto-arch a5 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %s | FileCheck %s --check-prefix=A5 -// RUN: ptoas --pto-arch a3 %s | FileCheck %s --check-prefix=A3 - -module attributes {"pto.device-spec" = "Ascend950"} { - func.func @tmov_acc_to_vec_mode_a5(%a: memref<32x32xf16, #pto.address_space>, - %b: memref<32x32xf16, #pto.address_space>) { - %a_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %b_mat = memref.alloc() : memref<32x32xf16, #pto.address_space> - %a_left = memref.alloc() : memref<32x32xf16, #pto.address_space> - %b_right = memref.alloc() : memref<32x32xf16, #pto.address_space> - %acc = memref.alloc() : memref<32x32xf32, #pto.address_space> - %dst_vec = memref.alloc() : memref<32x32xf32, #pto.address_space> - - pto.tload ins(%a : memref<32x32xf16, #pto.address_space>) - outs(%a_mat : memref<32x32xf16, #pto.address_space>) - pto.tload ins(%b : memref<32x32xf16, #pto.address_space>) - outs(%b_mat : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%a_mat : memref<32x32xf16, #pto.address_space>) - outs(%a_left : memref<32x32xf16, #pto.address_space>) - pto.tmov ins(%b_mat : memref<32x32xf16, #pto.address_space>) - outs(%b_right : memref<32x32xf16, #pto.address_space>) - pto.tmatmul ins(%a_left, %b_right : memref<32x32xf16, #pto.address_space>, - memref<32x32xf16, #pto.address_space>) - outs(%acc : memref<32x32xf32, #pto.address_space>) - pto.tmov ins(%acc : memref<32x32xf32, #pto.address_space>) - outs(%dst_vec : memref<32x32xf32, #pto.address_space>) - {accToVecMode = #pto.acc_to_vec_mode} - return - } -} +// RUN: ptoas --pto-arch a5 --enable-graph-sync-solver --graph-sync-solver-event-id-max=64 %S/tmov_acc_to_vec_mode_a5_emitc.pto | FileCheck %s --check-prefix=A5 +// RUN: ptoas --pto-arch a3 %S/tmov_acc_to_vec_mode_a5_emitc.pto | FileCheck %s --check-prefix=A3 // A5-LABEL: AICORE void tmov_acc_to_vec_mode_a5( // A5: set_flag(PIPE_M, PIPE_FIX, EVENT_ID{{[0-9]+}}); diff --git a/test/lit/pto/tmov_fbuf_pipe_selection.pto b/test/lit/pto/tmov_fbuf_pipe_selection.pto index 7f0c1956e8..a2b06e13cf 100644 --- a/test/lit/pto/tmov_fbuf_pipe_selection.pto +++ b/test/lit/pto/tmov_fbuf_pipe_selection.pto @@ -13,11 +13,11 @@ module { func.func @tmov_fbuf_pipe_selection( - %input: memref<1x16xui64, #pto.address_space>) { + %input: !pto.partition_tensor_view<1x16xui64>) { %src = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%input : memref<1x16xui64, #pto.address_space>) + pto.tload ins(%input : !pto.partition_tensor_view<1x16xui64>) outs(%src : !pto.tile_buf) pto.tmov ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) diff --git a/test/lit/pto/tmov_mx_scale_pipe_selection_a5.pto b/test/lit/pto/tmov_mx_scale_pipe_selection_a5.pto index ddf9db8473..da6cea07ab 100644 --- a/test/lit/pto/tmov_mx_scale_pipe_selection_a5.pto +++ b/test/lit/pto/tmov_mx_scale_pipe_selection_a5.pto @@ -11,11 +11,11 @@ module { func.func @tmov_mx_scale_left_pipe_selection( - %input: memref<1x4x!pto.f8E8M0, #pto.address_space>) { + %input: !pto.partition_tensor_view<1x4x!pto.f8E8M0>) { %src = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%input : memref<1x4x!pto.f8E8M0, #pto.address_space>) + pto.tload ins(%input : !pto.partition_tensor_view<1x4x!pto.f8E8M0>) outs(%src : !pto.tile_buf) pto.tmov ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) @@ -23,11 +23,11 @@ module { } func.func @tmov_mx_scale_right_pipe_selection( - %input: memref<4x16x!pto.f8E8M0, #pto.address_space>) { + %input: !pto.partition_tensor_view<4x16x!pto.f8E8M0>) { %src = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%input : memref<4x16x!pto.f8E8M0, #pto.address_space>) + pto.tload ins(%input : !pto.partition_tensor_view<4x16x!pto.f8E8M0>) outs(%src : !pto.tile_buf) pto.tmov ins(%src : !pto.tile_buf) outs(%dst : !pto.tile_buf) diff --git a/test/lit/pto/tprefetch_low_precision_memref_invalid_a3.pto b/test/lit/pto/tprefetch_low_precision_memref_invalid_a3.pto index a89562de3d..da5034390d 100644 --- a/test/lit/pto/tprefetch_low_precision_memref_invalid_a3.pto +++ b/test/lit/pto/tprefetch_low_precision_memref_invalid_a3.pto @@ -9,10 +9,10 @@ // RUN: not ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s func.func @tprefetch_low_precision_memref_invalid_a3( - %src: memref<16x16x!pto.hif8, #pto.address_space>, - %dst: memref<16x16xi8, #pto.address_space>) { - pto.tprefetch ins(%src : memref<16x16x!pto.hif8, #pto.address_space>) - outs(%dst : memref<16x16xi8, #pto.address_space>) + %src: !pto.partition_tensor_view<16x16x!pto.hif8>, + %dst: !pto.tile_buf) { + pto.tprefetch ins(%src : !pto.partition_tensor_view<16x16x!pto.hif8>) + outs(%dst : !pto.tile_buf) return } diff --git a/test/lit/pto/tquant_mxfp8_a5.pto b/test/lit/pto/tquant_mxfp8_a5.pto index 4c47b853f8..e8233e342a 100644 --- a/test/lit/pto/tquant_mxfp8_a5.pto +++ b/test/lit/pto/tquant_mxfp8_a5.pto @@ -10,55 +10,55 @@ module attributes {"pto.device-spec" = "Ascend910B3"} { func.func @tquant_mxfp8() attributes {pto.kernel_kind = #pto.kernel_kind} { - %src = memref.alloc() : memref<32x32xf32, #pto.address_space> - %exp = memref.alloc() : memref<1x32xui8, #pto.address_space> - %max = memref.alloc() : memref<1x32xf32, #pto.address_space> - %scaling = memref.alloc() : memref<1x32xf32, #pto.address_space> - %dst = memref.alloc() : memref<32x32xui8, #pto.address_space> + %src = pto.alloc_tile : !pto.tile_buf + %exp = pto.alloc_tile : !pto.tile_buf + %max = pto.alloc_tile : !pto.tile_buf + %scaling = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf - pto.tquant.mx ins(%src : memref<32x32xf32, #pto.address_space>) + pto.tquant.mx ins(%src : !pto.tile_buf) outs(%dst, %exp, %max, %scaling : - memref<32x32xui8, #pto.address_space>, - memref<1x32xui8, #pto.address_space>, - memref<1x32xf32, #pto.address_space>, - memref<1x32xf32, #pto.address_space>) + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) {quant_type = #pto} return } func.func @tquant_mxfp4_nv() attributes {pto.kernel_kind = #pto.kernel_kind} { - %src = memref.alloc() : memref<32x64xbf16, #pto.address_space> - %exp = memref.alloc() : memref<1x64xui8, #pto.address_space> - %max = memref.alloc() : memref<1x64xbf16, #pto.address_space> - %scaling = memref.alloc() : memref<1x64xbf16, #pto.address_space> - %dst = memref.alloc() : memref<32x32x!pto.f4E2M1x2, #pto.address_space> + %src = pto.alloc_tile : !pto.tile_buf + %exp = pto.alloc_tile : !pto.tile_buf + %max = pto.alloc_tile : !pto.tile_buf + %scaling = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf - pto.tquant.mx ins(%src : memref<32x64xbf16, #pto.address_space>) + pto.tquant.mx ins(%src : !pto.tile_buf) outs(%dst, %exp, %max, %scaling : - memref<32x32x!pto.f4E2M1x2, #pto.address_space>, - memref<1x64xui8, #pto.address_space>, - memref<1x64xbf16, #pto.address_space>, - memref<1x64xbf16, #pto.address_space>) + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) {quant_type = #pto, quantScaleAlg = #pto} return } func.func @tquant_mxfp8_nz() attributes {pto.kernel_kind = #pto.kernel_kind} { - %src = memref.alloc() : memref<32x32xf32, #pto.address_space> - %exp = memref.alloc() : memref<1x32xui8, #pto.address_space> - %max = memref.alloc() : memref<1x32xf32, #pto.address_space> - %scaling = memref.alloc() : memref<1x32xf32, #pto.address_space> - %exp_zz = memref.alloc() : memref<1x32xui8, #pto.address_space> - %dst = memref.alloc() : memref<32x32xui8, #pto.address_space> + %src = pto.alloc_tile : !pto.tile_buf + %exp = pto.alloc_tile : !pto.tile_buf + %max = pto.alloc_tile : !pto.tile_buf + %scaling = pto.alloc_tile : !pto.tile_buf + %exp_zz = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf - pto.tquant.mx ins(%src : memref<32x32xf32, #pto.address_space>) + pto.tquant.mx ins(%src : !pto.tile_buf) outs(%dst, %exp, %max, %scaling, %exp_zz : - memref<32x32xui8, #pto.address_space>, - memref<1x32xui8, #pto.address_space>, - memref<1x32xf32, #pto.address_space>, - memref<1x32xf32, #pto.address_space>, - memref<1x32xui8, #pto.address_space>) + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf, + !pto.tile_buf) {quant_type = #pto, storeMode = #pto} return diff --git a/test/lit/pto/tscatter_indexed_macro_sync_model.pto b/test/lit/pto/tscatter_indexed_macro_sync_model.pto index 30f8c8f3cc..5b894fd132 100644 --- a/test/lit/pto/tscatter_indexed_macro_sync_model.pto +++ b/test/lit/pto/tscatter_indexed_macro_sync_model.pto @@ -1,21 +1,22 @@ // RUN: ptoas --pto-arch=a3 --enable-insert-sync --pto-insert-sync-debug=2 %s -o - 2>&1 | FileCheck %s module { - func.func @tscatter_indexed_macro_model(%src_gm: memref<1x32xf16, #pto.address_space>, - %idx_gm: memref<1x32xi16, #pto.address_space>) { - %src = memref.alloc() : memref<1x32xf16, #pto.address_space> - %idx = memref.alloc() : memref<1x32xi16, #pto.address_space> - %dst = memref.alloc() : memref<1x32xf16, #pto.address_space> - %out = memref.alloc() : memref<1x32xf16, #pto.address_space> + func.func @tscatter_indexed_macro_model( + %src_gm: !pto.partition_tensor_view<1x32xf16>, + %idx_gm: !pto.partition_tensor_view<1x32xi16>) { + %src = pto.alloc_tile : !pto.tile_buf + %idx = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + %out = pto.alloc_tile : !pto.tile_buf - pto.tload ins(%src_gm : memref<1x32xf16, #pto.address_space>) - outs(%src : memref<1x32xf16, #pto.address_space>) {layout = #pto.layout} - pto.tload ins(%idx_gm : memref<1x32xi16, #pto.address_space>) - outs(%idx : memref<1x32xi16, #pto.address_space>) {layout = #pto.layout} - pto.tscatter ins(%src, %idx : memref<1x32xf16, #pto.address_space>, memref<1x32xi16, #pto.address_space>) - outs(%dst : memref<1x32xf16, #pto.address_space>) - pto.tmov ins(%dst : memref<1x32xf16, #pto.address_space>) - outs(%out : memref<1x32xf16, #pto.address_space>) + pto.tload ins(%src_gm : !pto.partition_tensor_view<1x32xf16>) + outs(%src : !pto.tile_buf) {layout = #pto.layout} + pto.tload ins(%idx_gm : !pto.partition_tensor_view<1x32xi16>) + outs(%idx : !pto.tile_buf) {layout = #pto.layout} + pto.tscatter ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + pto.tmov ins(%dst : !pto.tile_buf) + outs(%out : !pto.tile_buf) return } } diff --git a/test/lit/pto/tstore_low_precision_a5_valid.pto b/test/lit/pto/tstore_low_precision_a5_valid.pto index 2dbcfc59d3..a5d5d58524 100644 --- a/test/lit/pto/tstore_low_precision_a5_valid.pto +++ b/test/lit/pto/tstore_low_precision_a5_valid.pto @@ -24,14 +24,14 @@ module { } } -// CHECK: func.func @tstore_low_precision_a5_valid(%arg0: memref<16x16xf8E4M3FN>, %arg1: memref<16x16x!pto.hif8>, %arg2: i64) +// CHECK: func.func @tstore_low_precision_a5_valid(%arg0: !pto.partition_tensor_view<16x16xf8E4M3FN>, %arg1: !pto.partition_tensor_view<16x16x!pto.hif8>, %arg2: i64) // CHECK: %[[VEC:.*]] = pto.declare_tile -> !pto.tile_buf // CHECK: %[[ACC:.*]] = pto.declare_tile -> !pto.tile_buf // CHECK: pto.section.vector { // CHECK: pto.tstore ins(%[[VEC]] : !pto.tile_buf) -// CHECK: outs(%arg0 : memref<16x16xf8E4M3FN>) +// CHECK: outs(%arg0 : !pto.partition_tensor_view<16x16xf8E4M3FN>) // CHECK: } // CHECK: pto.section.cube { // CHECK: pto.tstore ins(%[[ACC]] : !pto.tile_buf, %arg2 : i64) -// CHECK: outs(%arg1 : memref<16x16x!pto.hif8>) +// CHECK: outs(%arg1 : !pto.partition_tensor_view<16x16x!pto.hif8>) // CHECK: } diff --git a/test/lit/vpto/expand_tile_op_tilelang.pto b/test/lit/vpto/expand_tile_op_tilelang.pto index cef4e22bb9..cdda474a56 100644 --- a/test/lit/vpto/expand_tile_op_tilelang.pto +++ b/test/lit/vpto/expand_tile_op_tilelang.pto @@ -10,7 +10,7 @@ // expands pto.tadd via the default TileLang Python DSL template // lib/TileOps/tadd_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tabs.pto b/test/lit/vpto/expand_tile_op_tilelang_tabs.pto index 40d19d663c..6ffd2e35f1 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tabs.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tabs.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tand.pto b/test/lit/vpto/expand_tile_op_tilelang_tand.pto index d59c146000..a2d51ec318 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tand.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tand.pto @@ -10,7 +10,7 @@ // expands pto.tand via the default TileLang Python DSL template // lib/TileOps/tand_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tcmp.pto b/test/lit/vpto/expand_tile_op_tilelang_tcmp.pto index 4acfda6920..192cd09318 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tcmp.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tcmp.pto @@ -10,7 +10,7 @@ // expands pto.tcmp via the default TileLang Python DSL template // lib/TileOps/tcmp_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tdiv.pto b/test/lit/vpto/expand_tile_op_tilelang_tdiv.pto index 2811e8316c..bfc5ed09f4 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tdiv.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tdiv.pto @@ -10,7 +10,7 @@ // expands pto.tdiv via the default TileLang Python DSL template // lib/TileOps/tdiv_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_texp.pto b/test/lit/vpto/expand_tile_op_tilelang_texp.pto index 6243cd09c0..1109f0f84a 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_texp.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_texp.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_texpand.pto b/test/lit/vpto/expand_tile_op_tilelang_texpand.pto index 0915e1f2d1..a4b1796ce0 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_texpand.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_texpand.pto @@ -10,7 +10,7 @@ // expands pto.texpands via the default TileLang Python DSL template // lib/TileOps/texpands_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto index 7dab0b04c2..14613b738c 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad.pto @@ -10,7 +10,7 @@ // expands pto.tfillpad via the default TileLang Python DSL template // lib/TileOps/tfillpad_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto index e975231566..7809de61df 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_expand.pto @@ -10,7 +10,7 @@ // expands pto.tfillpad_expand via the default TileLang Python DSL template // lib/TileOps/tfillpad_expand_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto index d3e7b00e10..68a58e8cf5 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfillpad_inplace.pto @@ -10,7 +10,7 @@ // expands pto.tfillpad (inplace mode) via the default TileLang Python DSL template // lib/TileOps/tfillpad_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tfmod.pto b/test/lit/vpto/expand_tile_op_tilelang_tfmod.pto index 4cfd8b4510..c2e04900f1 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tfmod.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tfmod.pto @@ -10,7 +10,7 @@ // expands pto.tfmod via the default TileLang Python DSL template // lib/TileOps/tfmod_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tlog.pto b/test/lit/vpto/expand_tile_op_tilelang_tlog.pto index d2656a4097..b588b77c1a 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tlog.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tlog.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tmax.pto b/test/lit/vpto/expand_tile_op_tilelang_tmax.pto index 4ec289f0e0..0ef8a20ba3 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tmax.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tmax.pto @@ -10,7 +10,7 @@ // expands pto.tmax via the default TileLang Python DSL template // lib/TileOps/tmax_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tmin.pto b/test/lit/vpto/expand_tile_op_tilelang_tmin.pto index 621edc55fe..b9abda842b 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tmin.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tmin.pto @@ -10,7 +10,7 @@ // expands pto.tmin via the default TileLang Python DSL template // lib/TileOps/tmin_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tmrgsort.pto b/test/lit/vpto/expand_tile_op_tilelang_tmrgsort.pto index 1d50ea4c7e..f606774ba3 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tmrgsort.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tmrgsort.pto @@ -10,7 +10,7 @@ // expands pto.tmrgsort format1 (single input list) via the TileLang Python DSL template // lib/TileOps/tmrgsort_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tmul.pto b/test/lit/vpto/expand_tile_op_tilelang_tmul.pto index de37d60390..e365ff4026 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tmul.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tmul.pto @@ -10,7 +10,7 @@ // expands pto.tmul via the default TileLang Python DSL template // lib/TileOps/tmul_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tneg.pto b/test/lit/vpto/expand_tile_op_tilelang_tneg.pto index e932bb6afc..a1c285b3ea 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tneg.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tneg.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tnot.pto b/test/lit/vpto/expand_tile_op_tilelang_tnot.pto index 6e15b0e394..be5dbf4508 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tnot.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tnot.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tor.pto b/test/lit/vpto/expand_tile_op_tilelang_tor.pto index d12f3e739d..df214ab4f4 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tor.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tor.pto @@ -10,7 +10,7 @@ // expands pto.tor via the default TileLang Python DSL template // lib/TileOps/tor_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tpartadd.pto b/test/lit/vpto/expand_tile_op_tilelang_tpartadd.pto index 8f6883c34c..d4c5616ffc 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tpartadd.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tpartadd.pto @@ -9,7 +9,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // expands pto.tpartadd/tpartmul/tpartmax/tpartmin via TileLang Python DSL templates. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tpartmax.pto b/test/lit/vpto/expand_tile_op_tilelang_tpartmax.pto index 522d5f3b28..34c19b573f 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tpartmax.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tpartmax.pto @@ -9,7 +9,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // expands pto.tpartadd/tpartmul/tpartmax/tpartmin via TileLang Python DSL templates. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tpartmin.pto b/test/lit/vpto/expand_tile_op_tilelang_tpartmin.pto index 7980d7cfe5..46d91d472f 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tpartmin.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tpartmin.pto @@ -9,7 +9,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // expands pto.tpartadd/tpartmul/tpartmax/tpartmin via TileLang Python DSL templates. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tpartmul.pto b/test/lit/vpto/expand_tile_op_tilelang_tpartmul.pto index 53c140baef..295bac5a46 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tpartmul.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tpartmul.pto @@ -9,7 +9,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // expands pto.tpartadd/tpartmul/tpartmax/tpartmin via TileLang Python DSL templates. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_trecip.pto b/test/lit/vpto/expand_tile_op_tilelang_trecip.pto index 8e082462a7..22e33cce73 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_trecip.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_trecip.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_trem.pto b/test/lit/vpto/expand_tile_op_tilelang_trem.pto index 2aeca7f559..790e9cd8ed 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_trem.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_trem.pto @@ -10,7 +10,7 @@ // expands pto.trem via the default TileLang Python DSL template // lib/TileOps/trem_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_trowargmax.pto b/test/lit/vpto/expand_tile_op_tilelang_trowargmax.pto index 8d0de3bd5d..123c09eb27 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_trowargmax.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_trowargmax.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s // After the full tile-op-expand path on the VPTO backend, the original diff --git a/test/lit/vpto/expand_tile_op_tilelang_trowargmin.pto b/test/lit/vpto/expand_tile_op_tilelang_trowargmin.pto index 88556082d2..ff23eb56d4 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_trowargmin.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_trowargmin.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s // After the full tile-op-expand path on the VPTO backend, the original diff --git a/test/lit/vpto/expand_tile_op_tilelang_trowmax.pto b/test/lit/vpto/expand_tile_op_tilelang_trowmax.pto index 8be5003928..7fed90a350 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_trowmax.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_trowmax.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s // After the full tile-op-expand path on the VPTO backend, the original diff --git a/test/lit/vpto/expand_tile_op_tilelang_trowmin.pto b/test/lit/vpto/expand_tile_op_tilelang_trowmin.pto index f9bcb767a5..23d2b3ae55 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_trowmin.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_trowmin.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s // After the full tile-op-expand path on the VPTO backend, the original diff --git a/test/lit/vpto/expand_tile_op_tilelang_trowprod.pto b/test/lit/vpto/expand_tile_op_tilelang_trowprod.pto index 9fcb891126..2c5ee7f957 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_trowprod.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_trowprod.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s // After the full tile-op-expand path on the VPTO backend, the original diff --git a/test/lit/vpto/expand_tile_op_tilelang_trowsum.pto b/test/lit/vpto/expand_tile_op_tilelang_trowsum.pto index b5b15f6742..b6b67d0249 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_trowsum.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_trowsum.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s // After the full tile-op-expand path on the VPTO backend, the original diff --git a/test/lit/vpto/expand_tile_op_tilelang_trsqrt.pto b/test/lit/vpto/expand_tile_op_tilelang_trsqrt.pto index 6564102bd4..e9b21e32d1 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_trsqrt.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_trsqrt.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tshl.pto b/test/lit/vpto/expand_tile_op_tilelang_tshl.pto index a05fa95ec3..6d5bda4200 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tshl.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tshl.pto @@ -10,7 +10,7 @@ // expands pto.tshl via the default TileLang Python DSL template // lib/TileOps/tshl_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tshr.pto b/test/lit/vpto/expand_tile_op_tilelang_tshr.pto index 6ce58add51..1164daa218 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tshr.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tshr.pto @@ -10,7 +10,7 @@ // expands pto.tshr via the default TileLang Python DSL template // lib/TileOps/tshr_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tsort32.pto b/test/lit/vpto/expand_tile_op_tilelang_tsort32.pto index 906b72fb4b..1617b17ae8 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tsort32.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tsort32.pto @@ -10,7 +10,7 @@ // expands pto.tsort32 via the TileLang Python DSL template // lib/TileOps/tsort32_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tsqrt.pto b/test/lit/vpto/expand_tile_op_tilelang_tsqrt.pto index 6672237523..4f2778877d 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tsqrt.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tsqrt.pto @@ -8,7 +8,7 @@ // Test that ExpandTileOp + InlineLibCall + FoldTileBufIntrinsics pipeline // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_tsub.pto b/test/lit/vpto/expand_tile_op_tilelang_tsub.pto index 7f53712feb..56655353dc 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_tsub.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_tsub.pto @@ -10,7 +10,7 @@ // expands pto.tsub via the default TileLang Python DSL template // lib/TileOps/tsub_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tilelang_txor.pto b/test/lit/vpto/expand_tile_op_tilelang_txor.pto index 5441bf62c3..9870d48437 100644 --- a/test/lit/vpto/expand_tile_op_tilelang_txor.pto +++ b/test/lit/vpto/expand_tile_op_tilelang_txor.pto @@ -10,7 +10,7 @@ // expands pto.txor via the default TileLang Python DSL template // lib/TileOps/txor_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tlrelu_tilelang.pto b/test/lit/vpto/expand_tile_op_tlrelu_tilelang.pto index c5b1b4dc97..6f541748c0 100644 --- a/test/lit/vpto/expand_tile_op_tlrelu_tilelang.pto +++ b/test/lit/vpto/expand_tile_op_tlrelu_tilelang.pto @@ -10,7 +10,7 @@ // expands pto.tlrelu via the default TileLang Python DSL template // lib/TileOps/tlrelu_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tprelu_tilelang.pto b/test/lit/vpto/expand_tile_op_tprelu_tilelang.pto index 91d3810f00..d23a95baad 100644 --- a/test/lit/vpto/expand_tile_op_tprelu_tilelang.pto +++ b/test/lit/vpto/expand_tile_op_tprelu_tilelang.pto @@ -10,7 +10,7 @@ // expands pto.tprelu via the default TileLang Python DSL template // lib/TileOps/tprelu_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_trelu_tilelang.pto b/test/lit/vpto/expand_tile_op_trelu_tilelang.pto index c11c8b3598..702b1e8366 100644 --- a/test/lit/vpto/expand_tile_op_trelu_tilelang.pto +++ b/test/lit/vpto/expand_tile_op_trelu_tilelang.pto @@ -10,7 +10,7 @@ // expands pto.trelu via the default TileLang Python DSL template // lib/TileOps/trelu_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tsel_tilelang.pto b/test/lit/vpto/expand_tile_op_tsel_tilelang.pto index 42e045bd28..37fd3b5115 100644 --- a/test/lit/vpto/expand_tile_op_tsel_tilelang.pto +++ b/test/lit/vpto/expand_tile_op_tsel_tilelang.pto @@ -10,7 +10,7 @@ // expands pto.tsel via the default TileLang Python DSL template // lib/TileOps/tsel_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/expand_tile_op_tsels_tilelang.pto b/test/lit/vpto/expand_tile_op_tsels_tilelang.pto index 71efabde1b..84fdb92bc4 100644 --- a/test/lit/vpto/expand_tile_op_tsels_tilelang.pto +++ b/test/lit/vpto/expand_tile_op_tsels_tilelang.pto @@ -10,7 +10,7 @@ // expands pto.tsels via the default TileLang Python DSL template // lib/TileOps/tsels_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/normalize_uncovered_tile_sections_reject_residual_nested_mix.pto b/test/lit/vpto/normalize_uncovered_tile_sections_reject_residual_nested_mix.pto index 1ae41f30c7..45327f19e9 100644 --- a/test/lit/vpto/normalize_uncovered_tile_sections_reject_residual_nested_mix.pto +++ b/test/lit/vpto/normalize_uncovered_tile_sections_reject_residual_nested_mix.pto @@ -12,23 +12,23 @@ module attributes {pto.target_arch = "a5"} { func.func @reject_nested_section_with_sibling_uncovered_tiles( - %src: memref<16x16xf32, #pto.address_space>, - %dst: memref<16x16xf32, #pto.address_space>) attributes {pto.kernel} { + %src: !pto.partition_tensor_view<16x16xf32>, + %dst: !pto.partition_tensor_view<16x16xf32>) attributes {pto.kernel} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index - %tile0 = memref.alloc() : memref<16x16xf32, #pto.address_space> - %tile1 = memref.alloc() : memref<16x16xf32, #pto.address_space> + %tile0 = pto.alloc_tile : !pto.tile_buf + %tile1 = pto.alloc_tile : !pto.tile_buf %scalar = arith.constant 1.0 : f32 scf.for %i = %c0 to %c1 step %c1 { pto.section.vector { - pto.tload ins(%src : memref<16x16xf32, #pto.address_space>) - outs(%tile0 : memref<16x16xf32, #pto.address_space>) + pto.tload ins(%src : !pto.partition_tensor_view<16x16xf32>) + outs(%tile0 : !pto.tile_buf) } - pto.tadds ins(%tile0, %scalar : memref<16x16xf32, #pto.address_space>, f32) - outs(%tile1 : memref<16x16xf32, #pto.address_space>) - pto.tstore ins(%tile1 : memref<16x16xf32, #pto.address_space>) - outs(%dst : memref<16x16xf32, #pto.address_space>) + pto.tadds ins(%tile0, %scalar : !pto.tile_buf, f32) + outs(%tile1 : !pto.tile_buf) + pto.tstore ins(%tile1 : !pto.tile_buf) + outs(%dst : !pto.partition_tensor_view<16x16xf32>) } return } diff --git a/test/lit/vpto/tcolargmax.pto b/test/lit/vpto/tcolargmax.pto index 2862fb7586..9aa908a012 100644 --- a/test/lit/vpto/tcolargmax.pto +++ b/test/lit/vpto/tcolargmax.pto @@ -10,7 +10,7 @@ // expands pto.tcolargmax via the TileLang Python DSL template // lib/TileOps/tcolargmax_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolargmin.pto b/test/lit/vpto/tcolargmin.pto index 4b703f1386..1742150698 100644 --- a/test/lit/vpto/tcolargmin.pto +++ b/test/lit/vpto/tcolargmin.pto @@ -10,7 +10,7 @@ // expands pto.tcolargmin via the TileLang Python DSL template // lib/TileOps/tcolargmin_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolexpand_tilelang.pto b/test/lit/vpto/tcolexpand_tilelang.pto index a35feeb1d8..d276cc82c4 100644 --- a/test/lit/vpto/tcolexpand_tilelang.pto +++ b/test/lit/vpto/tcolexpand_tilelang.pto @@ -2,7 +2,7 @@ // expands pto.tcolexpand via the default TileLang Python DSL template // lib/TileOps/tcolexpand_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolexpandadd_tilelang.pto b/test/lit/vpto/tcolexpandadd_tilelang.pto index a5ed9a21cd..385bf7daf8 100644 --- a/test/lit/vpto/tcolexpandadd_tilelang.pto +++ b/test/lit/vpto/tcolexpandadd_tilelang.pto @@ -2,7 +2,7 @@ // expands pto.tcolexpandadd via the default TileLang Python DSL template // lib/TileOps/tcolexpandadd_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolexpanddiv_tilelang.pto b/test/lit/vpto/tcolexpanddiv_tilelang.pto index 3e2b32905a..84420d2502 100644 --- a/test/lit/vpto/tcolexpanddiv_tilelang.pto +++ b/test/lit/vpto/tcolexpanddiv_tilelang.pto @@ -2,7 +2,7 @@ // expands pto.tcolexpanddiv via the default TileLang Python DSL template // lib/TileOps/tcolexpanddiv_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolexpandexpdif_tilelang.pto b/test/lit/vpto/tcolexpandexpdif_tilelang.pto index 6e20622521..a5900e4939 100644 --- a/test/lit/vpto/tcolexpandexpdif_tilelang.pto +++ b/test/lit/vpto/tcolexpandexpdif_tilelang.pto @@ -2,7 +2,7 @@ // expands pto.tcolexpandexpdif via the default TileLang Python DSL template // lib/TileOps/tcolexpandexpdif_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolexpandmax_tilelang.pto b/test/lit/vpto/tcolexpandmax_tilelang.pto index 9d0fb20ec3..7e8cf8af81 100644 --- a/test/lit/vpto/tcolexpandmax_tilelang.pto +++ b/test/lit/vpto/tcolexpandmax_tilelang.pto @@ -2,7 +2,7 @@ // expands pto.tcolexpandmax via the default TileLang Python DSL template // lib/TileOps/tcolexpandmax_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolexpandmin_tilelang.pto b/test/lit/vpto/tcolexpandmin_tilelang.pto index 269307986b..b0c89668a4 100644 --- a/test/lit/vpto/tcolexpandmin_tilelang.pto +++ b/test/lit/vpto/tcolexpandmin_tilelang.pto @@ -2,7 +2,7 @@ // expands pto.tcolexpandmin via the default TileLang Python DSL template // lib/TileOps/tcolexpandmin_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolexpandmul_tilelang.pto b/test/lit/vpto/tcolexpandmul_tilelang.pto index 3e3d148a22..c23060d170 100644 --- a/test/lit/vpto/tcolexpandmul_tilelang.pto +++ b/test/lit/vpto/tcolexpandmul_tilelang.pto @@ -2,7 +2,7 @@ // expands pto.tcolexpandmul via the default TileLang Python DSL template // lib/TileOps/tcolexpandmul_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolexpandsub_tilelang.pto b/test/lit/vpto/tcolexpandsub_tilelang.pto index 49bfd7314f..0de8fa5932 100644 --- a/test/lit/vpto/tcolexpandsub_tilelang.pto +++ b/test/lit/vpto/tcolexpandsub_tilelang.pto @@ -2,7 +2,7 @@ // expands pto.tcolexpandsub via the default TileLang Python DSL template // lib/TileOps/tcolexpandsub_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolmax.pto b/test/lit/vpto/tcolmax.pto index 3d6af696e9..aa31a21521 100644 --- a/test/lit/vpto/tcolmax.pto +++ b/test/lit/vpto/tcolmax.pto @@ -10,7 +10,7 @@ // expands pto.tcolmax via the TileLang Python DSL template // lib/TileOps/tcolmax_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolmin.pto b/test/lit/vpto/tcolmin.pto index 83e6a621ab..695da2fecc 100644 --- a/test/lit/vpto/tcolmin.pto +++ b/test/lit/vpto/tcolmin.pto @@ -10,7 +10,7 @@ // expands pto.tcolmin via the TileLang Python DSL template // lib/TileOps/tcolmin_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolprod.pto b/test/lit/vpto/tcolprod.pto index af31f0595a..f7b3e02943 100644 --- a/test/lit/vpto/tcolprod.pto +++ b/test/lit/vpto/tcolprod.pto @@ -10,7 +10,7 @@ // expands pto.tcolprod via the TileLang Python DSL template // lib/TileOps/tcolprod_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/tcolsum.pto b/test/lit/vpto/tcolsum.pto index ee470c1dc7..d255a768f7 100644 --- a/test/lit/vpto/tcolsum.pto +++ b/test/lit/vpto/tcolsum.pto @@ -10,7 +10,7 @@ // expands pto.tcolsum via the TileLang Python DSL template // lib/TileOps/tcolsum_template.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/trowexpand_tile_op_expand.pto b/test/lit/vpto/trowexpand_tile_op_expand.pto index 54db4ff81c..479b35c45f 100644 --- a/test/lit/vpto/trowexpand_tile_op_expand.pto +++ b/test/lit/vpto/trowexpand_tile_op_expand.pto @@ -10,7 +10,7 @@ // expands pto.trowexpand via the TileLang Python DSL template // lib/TileOps/trowexpand.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/trowexpandadd_tile_op_expand.pto b/test/lit/vpto/trowexpandadd_tile_op_expand.pto index 208a3f55b6..2658ee6eff 100644 --- a/test/lit/vpto/trowexpandadd_tile_op_expand.pto +++ b/test/lit/vpto/trowexpandadd_tile_op_expand.pto @@ -10,7 +10,7 @@ // expands pto.trowexpandadd via the TileLang Python DSL template // lib/TileOps/trowexpandadd.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/trowexpanddiv_tile_op_expand.pto b/test/lit/vpto/trowexpanddiv_tile_op_expand.pto index 115dd538ad..a70b55ee86 100644 --- a/test/lit/vpto/trowexpanddiv_tile_op_expand.pto +++ b/test/lit/vpto/trowexpanddiv_tile_op_expand.pto @@ -10,7 +10,7 @@ // expands pto.trowexpanddiv via the TileLang Python DSL template // lib/TileOps/trowexpanddiv.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/trowexpandexpdif_tile_op_expand.pto b/test/lit/vpto/trowexpandexpdif_tile_op_expand.pto index b45529c99d..4e848079f9 100644 --- a/test/lit/vpto/trowexpandexpdif_tile_op_expand.pto +++ b/test/lit/vpto/trowexpandexpdif_tile_op_expand.pto @@ -10,7 +10,7 @@ // expands pto.trowexpandexpdif via the TileLang Python DSL template // lib/TileOps/trowexpandexpdif.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/trowexpandmax_tile_op_expand.pto b/test/lit/vpto/trowexpandmax_tile_op_expand.pto index 599d32ff1b..7e175642d8 100644 --- a/test/lit/vpto/trowexpandmax_tile_op_expand.pto +++ b/test/lit/vpto/trowexpandmax_tile_op_expand.pto @@ -10,7 +10,7 @@ // expands pto.trowexpandmax via the TileLang Python DSL template // lib/TileOps/trowexpandmax.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/trowexpandmin_tile_op_expand.pto b/test/lit/vpto/trowexpandmin_tile_op_expand.pto index e1b00b44af..e42646d4bf 100644 --- a/test/lit/vpto/trowexpandmin_tile_op_expand.pto +++ b/test/lit/vpto/trowexpandmin_tile_op_expand.pto @@ -10,7 +10,7 @@ // expands pto.trowexpandmin via the TileLang Python DSL template // lib/TileOps/trowexpandmin.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/trowexpandmul_tile_op_expand.pto b/test/lit/vpto/trowexpandmul_tile_op_expand.pto index 7e50bf2f84..2830d198e3 100644 --- a/test/lit/vpto/trowexpandmul_tile_op_expand.pto +++ b/test/lit/vpto/trowexpandmul_tile_op_expand.pto @@ -10,7 +10,7 @@ // expands pto.trowexpandmul via the TileLang Python DSL template // lib/TileOps/trowexpandmul.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/test/lit/vpto/trowexpandsub_tile_op_expand.pto b/test/lit/vpto/trowexpandsub_tile_op_expand.pto index a668612c37..e05a8b9f67 100644 --- a/test/lit/vpto/trowexpandsub_tile_op_expand.pto +++ b/test/lit/vpto/trowexpandsub_tile_op_expand.pto @@ -10,7 +10,7 @@ // expands pto.trowexpandsub via the TileLang Python DSL template // lib/TileOps/trowexpandsub.py. // -// Pipeline: PTOMaterializeTileHandles -> ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics +// Pipeline: ExpandTileOp -> InlineLibCall -> FoldTileBufIntrinsics // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --enable-tile-op-expand %s -o - 2>/dev/null | FileCheck %s diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 12c5d6ab97..463a382810 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -3181,13 +3181,8 @@ int mlir::pto::compilePTOASModule( return 0; } - // Reintroduce tile-native handles once on the shared mainline so both - // backends consume the same post-planning seam IR. - pm.addPass(pto::createPTOMaterializeTileHandlesPass()); pm.addPass(createCSEPass()); - // Inline PTODSL backend helpers only after the shared mainline has - // materialized tile-native handles, so helper arguments are restored to the - // tile_buf ABI before qk.as_ptr()-style bridges are cloned into callers. + // PTODSL backend helpers already use the tile-native ABI. pm.addPass(pto::createPTOInlineBackendHelpersPass()); if (effectiveBackend == PTOBackend::EmitC) pm.addPass(createNarrowUnusedMultiResultProvenancePass()); From fddb3605ac87adf4f4a52c6fac9cda647cd77f0c Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 21 Jul 2026 12:34:05 +0800 Subject: [PATCH 88/91] Remove memref compatibility from GM interfaces --- docs/PTO_IR_manual.md | 78 ++-- include/PTO/IR/PTO.h | 1 - include/PTO/IR/PTOOps.td | 137 ++++--- lib/PTO/IR/PTO.cpp | 342 ++++-------------- lib/PTO/Transforms/PTOToEmitC.cpp | 149 ++------ test/lit/pto/async_put_get_emitc.pto | 22 +- test/lit/pto/async_put_invalid_non_1d.pto | 12 +- .../pto/easy_param_completion_a5_emitc.pto | 8 +- .../globaltensor_layout_bytewidth_emitc.pto | 47 +-- ...ather_gm2l1_elem_scratch_reuse_sync_a5.pto | 24 +- test/lit/pto/mgather_gm2l1_elem_sync_a5.pto | 12 +- test/lit/pto/mgather_gm2l1_emitc.pto | 20 +- test/lit/pto/mgather_gm2l1_sync_a5.pto | 8 +- ...explicit_row_coalesce_validshape_emitc.pto | 8 +- ...gather_mscatter_low_precision_a5_emitc.pto | 8 +- test/lit/pto/mgather_mscatter_tile_native.pto | 12 +- test/lit/pto/store_fp_tile_native.pto | 6 +- test/lit/pto/syncall_emitc.pto | 8 +- test/lit/pto/syncall_invalid_missing_ub.pto | 4 +- test/lit/pto/tassign_reject_memref.pto | 3 +- test/lit/pto/tprefetch_async_emitc.pto | 4 +- .../pto/tprefetch_async_invalid_non_1d.pto | 6 +- test/lit/pto/tprefetch_tile_native.pto | 6 +- test/lit/pto/tsync_emitc.pto | 6 +- 24 files changed, 311 insertions(+), 620 deletions(-) diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index b32ea5150f..17347cd946 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -1167,7 +1167,7 @@ Lowering maps `%ctx` to `pto::PrefetchAsyncContext`, emits | Name | Type | Description | |------|------|-------------| -| `src` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Source GM region to prefetch | +| `src` | `pto.tensor_view` / `pto.partition_tensor_view` | Source GM region to prefetch | | `ctx` | `!pto.prefetch_async_context` | Explicit PTO prefetch async context | **Results:** @@ -1189,7 +1189,7 @@ Lowering maps `%ctx` to `pto::PrefetchAsyncContext`, emits -> !pto.prefetch_async_context %event = pto.tprefetch_async( %src, %ctx - : memref<128xf32, #pto.address_space>, + : !pto.partition_tensor_view<128xf32>, !pto.prefetch_async_context) -> !pto.async_event %session = pto.get_prefetch_async_session %ctx @@ -7592,7 +7592,7 @@ elem mode: dst[i, j] = mem[idx[i, j]] | Name | Type | Default | Description | |------|------|---------|-------------| -| `mem` | `!pto.partition_tensor_view<...>` / GM memref | `NA` | Global source table | +| `mem` | `!pto.partition_tensor_view<...>` | `NA` | Global source table | | `idx` | `pto.tile_buf` | `NA` | Index tile | | `dst` | `pto.tile_buf` | `NA` | Destination VEC tile | | `coalesce` | `#pto` | required | Explicit coalesce mode (`row` / `elem`) | @@ -7616,7 +7616,6 @@ elem mode: dst[i, j] = mem[idx[i, j]] - Element mode: `idx valid_shape == dst valid_shape`. - Row mode: `idx valid_shape` may be `[1, dst.valid_row]` or `[dst.valid_row, 1]`. - The `[1, R]` row-mode variant uses `row_major`; the `[R, 1]` row-mode variant uses `col_major`. - - If `mem` is a rank-5 static GM memref, it must satisfy `<1, 1, 1, Rows, RowWidth>`. - **Out-of-bounds mode** - Default `gatherOob = undefined` lowers to the default `MGATHER(dst, mem, idx)` overload. @@ -7634,15 +7633,15 @@ elem mode: dst[i, j] = mem[idx[i, j]] **Basic Example:** ```mlir -pto.mgather ins(%mem, %idx : memref<...>, !pto.tile_buf<...>) +pto.mgather ins(%mem, %idx : !pto.partition_tensor_view<...>, !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) {coalesce = #pto} -pto.mgather ins(%mem, %idx : memref<...>, !pto.tile_buf<...>) +pto.mgather ins(%mem, %idx : !pto.partition_tensor_view<...>, !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) {coalesce = #pto} -pto.mgather ins(%mem, %idx : memref<...>, !pto.tile_buf<...>) +pto.mgather ins(%mem, %idx : !pto.partition_tensor_view<...>, !pto.tile_buf<...>) outs(%dst : !pto.tile_buf<...>) {coalesce = #pto, gatherOob = #pto} ``` @@ -7657,7 +7656,7 @@ when `dst` is a `loc=mat` tile. - **`dst`** — `loc=mat`, `blayout=col_major`, `slayout=row_major`, `fractal=512` (NZ). Padded `cols` must be a multiple of `C0 = 32 / sizeof(elem)` and padded `rows` a multiple of `16` (`FRACTAL_NZ_ROW`). -- **`idx`** — supplied as a **GM tensor** (`memref` / `partition_tensor_view`) +- **`idx`** — supplied as a GM `partition_tensor_view` of `i32`, **not** a UB tile: on A5 the cube core that issues the L1 transfer cannot read AIV's UB. - **`coalesce`** — must be **explicit** (`row` or `elem`); there is no UB index @@ -7672,8 +7671,8 @@ when `dst` is a `loc=mat` tile. ```mlir // Row: no scratch, idx is a GM [1, R] tensor. -pto.mgather ins(%mem, %idx : memref<1x1x1x64x32xf16, #pto.address_space>, - memref<1x1x1x1x32xi32, #pto.address_space>) +pto.mgather ins(%mem, %idx : !pto.partition_tensor_view<1x1x1x64x32xf16>, + !pto.partition_tensor_view<1x1x1x1x32xi32>) outs(%dst : !pto.tile_buf) @@ -7681,9 +7680,9 @@ pto.mgather ins(%mem, %idx : memref<1x1x1x64x32xf16, #pto.address_space>, // Elem: GM scratch workspace staged in NZ layout before the bulk copy. pto.mgather ins(%mem, %idx, %scratch - : memref<1x1x1x64x32xf16, #pto.address_space>, - memref<1x1x1x32x32xi32, #pto.address_space>, - memref<1x1x1x32x32xf16, #pto.address_space>) + : !pto.partition_tensor_view<1x1x1x64x32xf16>, + !pto.partition_tensor_view<1x1x1x32x32xi32>, + !pto.partition_tensor_view<1x1x1x32x32xf16>) outs(%dst : !pto.tile_buf) @@ -7709,7 +7708,7 @@ elem mode: mem[idx[i, j]] = src[i, j] |------|------|---------|-------------| | `src` | `pto.tile_buf` | `NA` | Source VEC tile | | `idx` | `pto.tile_buf` | `NA` | Index tile | -| `mem` | `!pto.partition_tensor_view<...>` / GM memref | `NA` | Global destination table | +| `mem` | `!pto.partition_tensor_view<...>` | `NA` | Global destination table | | `coalesce` | `#pto` | inferred | Explicit coalesce mode (`row` / `elem`) | | `scatterAtomicOp` | `#pto` | `none` | Atomic mode (`none/add/max/min`) | | `scatterOob` | `#pto` | `undefined` | Out-of-bounds mode (`undefined/skip/clamp/wrap`) | @@ -7733,7 +7732,6 @@ elem mode: mem[idx[i, j]] = src[i, j] - Element mode: `idx valid_shape == src valid_shape`. - Row mode: `idx valid_shape` may be `[1, src.valid_row]` or `[src.valid_row, 1]`. - The `[1, R]` row-mode variant uses `row_major`; the `[R, 1]` row-mode variant uses `col_major`. - - If `mem` is a rank-5 static GM memref, it must satisfy `<1, 1, 1, Rows, RowWidth>`. - **Atomic modes** - Default `scatterAtomicOp = none` lowers to the default `MSCATTER(mem, src, idx)` overload. @@ -7757,19 +7755,19 @@ elem mode: mem[idx[i, j]] = src[i, j] ```mlir pto.mscatter ins(%src, %idx : !pto.tile_buf<...>, !pto.tile_buf<...>) - outs(%mem : memref<...>) + outs(%mem : !pto.partition_tensor_view<...>) pto.mscatter ins(%src, %idx : !pto.tile_buf<...>, !pto.tile_buf<...>) - outs(%mem : memref<...>) + outs(%mem : !pto.partition_tensor_view<...>) {scatterAtomicOp = #pto} pto.mscatter ins(%src, %idx : !pto.tile_buf<...>, !pto.tile_buf<...>) - outs(%mem : memref<...>) + outs(%mem : !pto.partition_tensor_view<...>) {scatterAtomicOp = #pto, scatterOob = #pto} pto.mscatter ins(%src, %idx : !pto.tile_buf<...>, !pto.tile_buf<...>) - outs(%mem : memref<...>) + outs(%mem : !pto.partition_tensor_view<...>) {coalesce = #pto, scatterConflict = #pto} ``` @@ -8801,9 +8799,9 @@ pto.wait_event [#pto.pipe_event_type, #pto.pipe_event_type` | `hard` or `soft` | | `core_type` | `#pto.sync_core_type<...>` | `aiv_only`, `aic_only`, or `mix` | @@ -8817,8 +8815,8 @@ pto.wait_event [#pto.pipe_event_type, #pto.pipe_event_type, #pto.pipe_event_type, mode = #pto.sync_all_mode, core_type = #pto.sync_core_type -} : (memref<64xi32, #pto.address_space>, - memref<64xi32, #pto.address_space>, +} : (!pto.partition_tensor_view<64xi32>, + !pto.tile_buf, i32) -> () "pto.syncall"() { @@ -9938,8 +9936,8 @@ This section documents PTO communication primitives. PTOAS currently exposes: | Name | Type | Description | |------|------|-------------| -| `scratch` | `pto.tile_buf` / local memref | Local scratch/staging buffer used by the async runtime | -| `workspace` | `!pto.ptr<...>` / GM memref | Global workspace backing the async session | +| `scratch` | `pto.tile_buf` | Local scratch/staging buffer used by the async runtime | +| `workspace` | `!pto.ptr<...>` | Global workspace backing the async session | | `sync_id` | optional `i32` attr | Session synchronization ID | | `block_bytes` | optional `i64` attr | Communication block size in bytes | | `comm_block_offset` | optional `i64` attr | Per-block GM offset in bytes | @@ -9950,8 +9948,8 @@ This section documents PTO communication primitives. PTOAS currently exposes: **Constraints & Verification:** -- `scratch` must be tile-like local storage. -- `workspace` must be a GM pointer/memref. +- `scratch` must be a local tile buffer. +- `workspace` must be a GM pointer. - Optional attrs are forwarded as session configuration and must use the declared integer types. **Basic Example:** @@ -9970,8 +9968,8 @@ This section documents PTO communication primitives. PTOAS currently exposes: | Name | Type | Description | |------|------|-------------| -| `dst` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Remote destination buffer | -| `src` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Local source buffer | +| `dst` | `pto.tensor_view` / `pto.partition_tensor_view` | Remote destination buffer | +| `src` | `pto.tensor_view` / `pto.partition_tensor_view` | Local source buffer | | `session` | `!pto.async_session` | Async DMA session | **Results:** `!pto.async_event` @@ -9998,8 +9996,8 @@ This section documents PTO communication primitives. PTOAS currently exposes: | Name | Type | Description | |------|------|-------------| -| `dst` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Local destination buffer | -| `src` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Remote source buffer | +| `dst` | `pto.tensor_view` / `pto.partition_tensor_view` | Local destination buffer | +| `src` | `pto.tensor_view` / `pto.partition_tensor_view` | Remote source buffer | | `session` | `!pto.async_session` | Async DMA session | **Results:** `!pto.async_event` @@ -10051,8 +10049,8 @@ This section documents PTO communication primitives. PTOAS currently exposes: | Name | Type | Description | |------|------|-------------| -| `dst` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Remote destination buffer | -| `src` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Local source buffer | +| `dst` | `pto.tensor_view` / `pto.partition_tensor_view` | Remote destination buffer | +| `src` | `pto.tensor_view` / `pto.partition_tensor_view` | Local source buffer | | `buf` | `buf(%ping)` or `buf(%ping, %pong)` | Staging bundle: one or two local VEC tiles | | `atomicType` | `#pto` | Atomic mode, e.g. `atomic_none` or `atomic_add` | @@ -10082,10 +10080,10 @@ pto.comm.tput(%dst, %src, buf(%ping, %pong) : !pto.partition_tensor_view<128xf32 | Name | Type | Description | |------|------|-------------| -| `dst` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Local destination buffer | -| `src` | GM memref / `pto.tensor_view` / `pto.partition_tensor_view` | Remote source buffer | -| `ping` | `pto.tile_buf` / local VEC memref | Required staging tile (wrapped in `buf(%ping)`) | -| `pong` | `pto.tile_buf` / local VEC memref | Optional second staging tile (`buf(%ping, %pong)`) | +| `dst` | `pto.tensor_view` / `pto.partition_tensor_view` | Local destination buffer | +| `src` | `pto.tensor_view` / `pto.partition_tensor_view` | Remote source buffer | +| `ping` | `pto.tile_buf` | Required staging tile (wrapped in `buf(%ping)`) | +| `pong` | `pto.tile_buf` | Optional second staging tile (`buf(%ping, %pong)`) | **Constraints & Verification:** diff --git a/include/PTO/IR/PTO.h b/include/PTO/IR/PTO.h index e858212835..e535a5d79a 100644 --- a/include/PTO/IR/PTO.h +++ b/include/PTO/IR/PTO.h @@ -141,7 +141,6 @@ inline constexpr char kPTOTargetArchAttrName[] = "pto.target_arch"; AddressSpaceAttr getPTOAddressSpaceAttr(Type type); /// Return true if type is a ptr/memref in GM address space (or default). -bool isScalarPtrOrMemRef(Type type); enum class PTOArch { A3, diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 5a0ae134b7..6862877560 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -38,23 +38,23 @@ def TensorOrMemref : AnyTypeOf<[AnyMemRef, AnyRankedTensor], "Tensor or Memref">; def PTODpsType : - AnyTypeOf<[AnyRankedTensor, AnyMemRef, PartitionTensorViewType, TileBufType]>; + AnyTypeOf<[AnyRankedTensor, PartitionTensorViewType, TileBufType]>; def PTOPipeEntryType : - AnyTypeOf<[AnyRankedTensor, AnyMemRef, TensorViewType, TileBufType], - "TensorView, TileBuf, Tensor, or MemRef">; + AnyTypeOf<[AnyRankedTensor, TensorViewType, TileBufType], + "TensorView, TileBuf, or Tensor">; +def PTOCommType : + AnyTypeOf<[AnyRankedTensor, TensorViewType, PartitionTensorViewType, + TileBufType], "TensorView, PartitionTensorView, TileBuf, or Tensor">; + +// VPTO address extraction may still materialize a low-level memref result. def PtrOrMemRef : AnyTypeOf<[PtrType, AnyMemRef], "Ptr or MemRef">; def TileBufOrMemRef : AnyTypeOf<[TileBufType, AnyMemRef], "TileBuf or MemRef">; -def ScalarPtrOrMemRef : - TypeConstraint< - CPred<"::mlir::pto::isScalarPtrOrMemRef($_self)">, - "Ptr or GM MemRef">; - def ScalarType : AnyTypeOf<[AnySignlessInteger, AnyFloat], "numeric (integer/float)">; @@ -115,7 +115,7 @@ def PtrToIntOp : PTO_Op<"ptrtoint", [Pure]> { }]; let arguments = (ins - PtrOrMemRef:$ptr + PtrType:$ptr ); let results = (outs @@ -144,7 +144,7 @@ def IntToPtrOp : PTO_Op<"inttoptr", [Pure]> { ); let results = (outs - PtrOrMemRef:$result + PtrType:$result ); let hasVerifier = 1; @@ -163,7 +163,6 @@ def CastPtrOp : PTO_Op<"castptr", [Pure]> { - integer -> !pto.ptr - !pto.ptr -> integer - !pto.ptr -> !pto.ptr - - memref<..., space> -> !pto.ptr (extract the aligned base ptr) Pointer-to-pointer casts must stay within the same PTO memory space. Cross space casts such as gm <-> ub are rejected by the verifier. @@ -189,7 +188,7 @@ def LoadScalarOp : PTO_Op<"load_scalar", [ let summary = "Load a single scalar element from a pointer at offset."; let arguments = (ins - ScalarPtrOrMemRef:$ptr, + PtrType:$ptr, Index:$offset ); @@ -208,7 +207,7 @@ def StoreScalarOp : PTO_Op<"store_scalar", [ let summary = "Store a single scalar element to a pointer at offset."; let arguments = (ins - ScalarPtrOrMemRef:$ptr, + PtrType:$ptr, Index:$offset, AnyType:$value ); @@ -227,7 +226,7 @@ def MakeTensorViewOp : PTO_Op<"make_tensor_view", [AttrSizedOperandSegments, let summary = "Wrap a pointer as a tensor_view descriptor (no allocation, no copy)."; let arguments = (ins - AnyType:$ptr, + PtrType:$ptr, Variadic:$shape, Variadic:$strides, OptionalAttr:$layout @@ -266,27 +265,24 @@ def PartitionViewOp : PTO_Op<"partition_view", [AttrSizedOperandSegments, let hasCustomAssemblyFormat = 1; } -// Helper: tensor_view or memref (after lowering tensor_view to memref). -def TensorViewLikeOrMemRef : - AnyTypeOf<[TensorViewType, PartitionTensorViewType, AnyMemRef], - "TensorView, PartitionTensorView, or MemRef">; +// Helper for either whole-tensor or partitioned GM views. +def TensorViewLike : + AnyTypeOf<[TensorViewType, PartitionTensorViewType], + "TensorView or PartitionTensorView">; -// Get the size of a dimension of a tensor_view or its lowered memref view. +// Get the size of a dimension of a tensor_view. // Result type: Index (use arith.index_cast if i32 is needed). def GetTensorViewDimOp : PTO_Op<"get_tensor_view_dim", [Pure]> { let summary = "Get the size of a dimension of a tensor_view."; let description = [{ Returns the size of the given dimension of a logical tensor view. - This op accepts either !pto.tensor_view or the memref it is lowered to. - IR (tensor_view form): + This op accepts !pto.tensor_view or !pto.partition_tensor_view. + IR: %dim_size = pto.get_tensor_view_dim %tv, %dim_index : !pto.tensor_view, index -> index - IR (memref form, after lowering): - %dim_size = pto.get_tensor_view_dim %mr, %dim_index - : memref<...>, index -> index }]; let arguments = (ins - TensorViewLikeOrMemRef:$tensor_view, + TensorViewLike:$tensor_view, Index:$dim_index ); let results = (outs Index:$result); @@ -302,11 +298,11 @@ def GetTensorViewStrideOp : PTO_Op<"get_tensor_view_stride", [Pure]> { let summary = "Get the stride of a dimension of a tensor_view."; let description = [{ Returns the stride, measured in elements, of the given dimension of a - logical tensor view. This op accepts either !pto.tensor_view or the memref - it is lowered to. + logical tensor view. This op accepts !pto.tensor_view or + !pto.partition_tensor_view. }]; let arguments = (ins - TensorViewLikeOrMemRef:$tensor_view, + TensorViewLike:$tensor_view, Index:$dim_index ); let results = (outs Index:$result); @@ -727,7 +723,7 @@ def MakePrefetchAsyncContextOp : PTO_Op<"make_prefetch_async_context"> { }]; let arguments = (ins - ScalarPtrOrMemRef:$workspace + PtrType:$workspace ); let results = (outs @@ -1972,8 +1968,8 @@ def BuildAsyncSessionOp : PTO_Op<"comm.build_async_session", [ let summary = "Build an async DMA session handle for TPUT_ASYNC/TGET_ASYNC"; let arguments = (ins - TileBufOrMemRef:$scratch, - ScalarPtrOrMemRef:$workspace, + TileBufType:$scratch, + PtrType:$workspace, OptionalAttr:$sync_id, OptionalAttr:$block_bytes, OptionalAttr:$comm_block_offset, @@ -1985,7 +1981,7 @@ def BuildAsyncSessionOp : PTO_Op<"comm.build_async_session", [ let hasVerifier = 1; let assemblyFormat = [{ - `(` $scratch `,` $workspace `:` qualified(type($scratch)) `,` type($workspace) `)` + `(` $scratch `,` $workspace `:` qualified(type($scratch)) `,` qualified(type($workspace)) `)` attr-dict `->` qualified(type($session)) }]; } @@ -1996,8 +1992,8 @@ def TPutAsyncOp : PTO_Op<"comm.tput_async", [ let summary = "Asynchronous remote write from local GM to remote GM"; let arguments = (ins - PTODpsType:$dst, - PTODpsType:$src, + PTOCommType:$dst, + PTOCommType:$src, AsyncSessionType:$session ); @@ -2017,8 +2013,8 @@ def TGetAsyncOp : PTO_Op<"comm.tget_async", [ let summary = "Asynchronous remote read from remote GM to local GM"; let arguments = (ins - PTODpsType:$dst, - PTODpsType:$src, + PTOCommType:$dst, + PTOCommType:$src, AsyncSessionType:$session ); @@ -2075,10 +2071,10 @@ def TPutOp : PTO_Op<"comm.tput", [ ]> { let summary = "Synchronous remote write from local GM to remote GM"; let arguments = (ins - PTODpsType:$dst, - PTODpsType:$src, - PTODpsType:$ping, - Optional:$pong, + PTOCommType:$dst, + PTOCommType:$src, + PTOCommType:$ping, + Optional:$pong, DefaultValuedAttr:$atomicType ); let results = (outs); @@ -2095,10 +2091,10 @@ def TGetOp : PTO_Op<"comm.tget", [ ]> { let summary = "Synchronous remote read from remote GM to local GM"; let arguments = (ins - PTODpsType:$dst, - PTODpsType:$src, - PTODpsType:$ping, - Optional:$pong + PTOCommType:$dst, + PTOCommType:$src, + PTOCommType:$ping, + Optional:$pong ); let results = (outs); let hasVerifier = 1; @@ -2114,7 +2110,7 @@ def TNotifyOp : PTO_Op<"comm.tnotify", [ ]> { let summary = "Send a signal notification to remote GM"; let arguments = (ins - PTODpsType:$signal, + PTOCommType:$signal, AnySignlessInteger:$value, PTO_NotifyOpAttr:$notifyOp ); @@ -2131,7 +2127,7 @@ def TWaitOp : PTO_Op<"comm.twait", [ ]> { let summary = "Block until signal(s) satisfy a comparison"; let arguments = (ins - PTODpsType:$signal, + PTOCommType:$signal, AnySignlessInteger:$cmpValue, PTO_WaitCmpAttr:$cmp ); @@ -2148,7 +2144,7 @@ def TTestOp : PTO_Op<"comm.ttest", [ ]> { let summary = "Non-blocking signal comparison test"; let arguments = (ins - PTODpsType:$signal, + PTOCommType:$signal, AnySignlessInteger:$cmpValue, PTO_WaitCmpAttr:$cmp ); @@ -2166,10 +2162,10 @@ def TBroadcastOp : PTO_Op<"comm.tbroadcast", [ ]> { let summary = "Broadcast local GM data to all group members"; let arguments = (ins - PTODpsType:$src, - PTODpsType:$ping, - Optional:$pong, - Variadic:$group, + PTOCommType:$src, + PTOCommType:$ping, + Optional:$pong, + Variadic:$group, I32Attr:$root ); let results = (outs); @@ -2183,10 +2179,10 @@ def CommTGatherOp : PTO_Op<"comm.tgather", [ ]> { let summary = "Gather remote GM data from a communication group"; let arguments = (ins - PTODpsType:$dst, - PTODpsType:$ping, - Optional:$pong, - Variadic:$group, + PTOCommType:$dst, + PTOCommType:$ping, + Optional:$pong, + Variadic:$group, I32Attr:$root ); let results = (outs); @@ -2200,10 +2196,10 @@ def CommTScatterOp : PTO_Op<"comm.tscatter", [ ]> { let summary = "Scatter local GM data to a communication group"; let arguments = (ins - PTODpsType:$src, - PTODpsType:$ping, - Optional:$pong, - Variadic:$group, + PTOCommType:$src, + PTOCommType:$ping, + Optional:$pong, + Variadic:$group, I32Attr:$root ); let results = (outs); @@ -2217,11 +2213,11 @@ def TReduceOp : PTO_Op<"comm.treduce", [ ]> { let summary = "Reduce remote GM data from a communication group"; let arguments = (ins - PTODpsType:$dst, - PTODpsType:$acc, - PTODpsType:$recvPing, - Optional:$recvPong, - Variadic:$group, + PTOCommType:$dst, + PTOCommType:$acc, + PTOCommType:$recvPing, + Optional:$recvPong, + Variadic:$group, PTO_ReduceOpAttr:$reduceOp, I32Attr:$root ); @@ -3076,9 +3072,8 @@ def MGatherOp : PTO_TOp<"mgather", [ // MTE2 pipe on every arch, mirroring pto.tload. The GM -> UB (VEC) path // below keeps its A5 SIMT (PIPE_V) / A2A3 MTE2 selection. // - // Resolve the dst address space from both tile_buf and legacy imported - // memref forms. InsertSync / GraphSyncSolver consume getPipe() through the - // OpPipeInterface. A tile-only check would fall through to PIPE_V on A5. + // Resolve the destination address space for InsertSync and + // GraphSyncSolver through the OpPipeInterface. auto getAddressSpace = [](::mlir::Type ty) -> std::optional<::mlir::pto::AddressSpace> { if (auto tb = ::mlir::dyn_cast<::mlir::pto::TileBufType>(ty)) { @@ -3087,14 +3082,6 @@ def MGatherOp : PTO_TOp<"mgather", [ return as.getAddressSpace(); return std::nullopt; } - if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty)) { - if (auto ms = mr.getMemorySpace()) { - if (auto as = - ::mlir::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) - return as.getAddressSpace(); - } - return std::nullopt; - } return std::nullopt; }; if (auto as = getAddressSpace(getDst().getType()); diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 271c2bb5ec..53b7e20c40 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -2394,14 +2394,10 @@ LogicalResult mlir::pto::MakeTensorViewOp::verify() { if (!tvTy) return emitOpError("result must be pto.tensor_view<...>"); - Type ptrElemTy; - if (auto pty = dyn_cast(getPtr().getType())) - ptrElemTy = pty.getElementType(); - else if (auto memrefTy = dyn_cast(getPtr().getType())) - ptrElemTy = memrefTy.getElementType(); - else - return emitOpError( - "ptr operand must be !pto.ptr<...> or a memref-backed pointer"); + auto ptrTy = dyn_cast(getPtr().getType()); + if (!ptrTy) + return emitOpError("ptr operand must be !pto.ptr<...>"); + Type ptrElemTy = ptrTy.getElementType(); if (ptrElemTy != tvTy.getElementType()) return emitOpError() << "ptr element type must match tensor_view element " @@ -2535,32 +2531,9 @@ LogicalResult mlir::pto::AddPtrOp::verify() { return success(); } -static LogicalResult verifyPtrLikeForAddressCast(Operation *op, Type type, - StringRef name) { - if (isa(type)) - return success(); - - auto memTy = dyn_cast(type); - if (!memTy) - return op->emitOpError() - << "expects " << name << " to be !pto.ptr<...> or a GM memref"; - - if (memTy.getRank() != 1) - return op->emitOpError() - << "expects lowered memref " << name << " to be rank-1"; - - if (!isGmAddressSpaceAttr(memTy.getMemorySpace())) - return op->emitOpError() - << "expects lowered memref " << name << " to use GM address space"; - - return success(); -} - static Type getPointerLikeElementType(Type type) { if (auto ptrTy = dyn_cast(type)) return ptrTy.getElementType(); - if (auto memTy = dyn_cast(type)) - return memTy.getElementType(); return Type(); } @@ -2586,8 +2559,9 @@ LogicalResult mlir::pto::PtrToIntOp::verify() { if (!intTy || intTy.getWidth() != 64) return emitOpError("result must be i64"); - return verifyPtrLikeForAddressCast(getOperation(), getPtr().getType(), - "ptr operand"); + if (!isa(getPtr().getType())) + return emitOpError("ptr operand must be !pto.ptr<...>"); + return success(); } LogicalResult mlir::pto::IntToPtrOp::verify() { @@ -2595,9 +2569,8 @@ LogicalResult mlir::pto::IntToPtrOp::verify() { if (!addrTy || addrTy.getWidth() != 64) return emitOpError("address operand must be i64"); - if (failed(verifyPtrLikeForAddressCast(getOperation(), getResult().getType(), - "result"))) - return failure(); + if (!isa(getResult().getType())) + return emitOpError("result must be !pto.ptr<...>"); Type dstElem = getPointerLikeElementType(getResult().getType()); if (!isEmitCSupportedScalarType(dstElem)) @@ -2664,7 +2637,8 @@ LogicalResult mlir::pto::CastPtrOp::verify() { inputMemRefType.getMemorySpace()); auto resultSpace = resultPtrType.getMemorySpace(); if (memrefSpace && memrefSpace != resultSpace) - return emitOpError("memref-to-ptr cast must stay within the same PTO memory space"); + return emitOpError( + "memref-to-ptr cast must stay within the same PTO memory space"); } if (inputPtrType && resultPtrType && @@ -2708,14 +2682,6 @@ AddressSpaceAttr mlir::pto::getPTOAddressSpaceAttr(Type type) { return scopeAttr; } -bool mlir::pto::isScalarPtrOrMemRef(Type type) { - if (auto pty = dyn_cast(type)) - return static_cast(pty); - if (auto memTy = dyn_cast(type)) - return isGmAddressSpaceAttr(memTy.getMemorySpace()); - return false; -} - bool mlir::pto::hasExplicitPTOEntryAttr(func::FuncOp func) { return func && (func->hasAttrOfType(kPTOEntryAttrName) || func->hasAttrOfType(kLegacyHACCEntryAttrName) || @@ -3235,52 +3201,32 @@ LogicalResult TPrefetchOp::verify() { Type srcElem; Type dstElem; - if (auto srcPart = dyn_cast(srcTy)) { - auto srcShape = srcPart.getShape(); - for (unsigned i = 0; i < srcShape.size(); ++i) { - if (srcShape[i] != ShapedType::kDynamic && srcShape[i] <= 0) - return emitOpError() << "expects src shape[" << i << "] to be positive"; - } - srcElem = srcPart.getElementType(); - } else if (auto srcMr = dyn_cast(srcTy)) { - if (!srcMr.hasRank()) - return emitOpError("expects src memref to be ranked"); - for (int64_t dim : srcMr.getShape()) { - if (dim != ShapedType::kDynamic && dim <= 0) - return emitOpError("expects src memref shape to be positive"); - } - srcElem = srcMr.getElementType(); - } else { - return emitOpError("expects src to be !pto.partition_tensor_view or memref"); + auto srcPart = dyn_cast(srcTy); + if (!srcPart) + return emitOpError("expects src to be !pto.partition_tensor_view"); + auto srcShape = srcPart.getShape(); + for (unsigned i = 0; i < srcShape.size(); ++i) { + if (srcShape[i] != ShapedType::kDynamic && srcShape[i] <= 0) + return emitOpError() << "expects src shape[" << i << "] to be positive"; } + srcElem = srcPart.getElementType(); - if (auto dstTile = dyn_cast(dstTy)) { - if (failed(verifyTileBufCommon(*this, dstTile, "dst", allowLowPrecision))) - return failure(); - auto dstValid = dstTile.getValidShape(); - for (unsigned i = 0; i < dstValid.size(); ++i) { - if (dstValid[i] != ShapedType::kDynamic && dstValid[i] < 0) - return emitOpError() << "expects dst valid_shape[" << i - << "] to be non-negative"; - } - auto dstSpace = getPTOMemorySpaceEnum(dstTile); - if (!dstSpace || (*dstSpace != pto::AddressSpace::VEC && - *dstSpace != pto::AddressSpace::MAT)) - return emitOpError("expects dst to use loc=vec or loc=mat"); - dstElem = dstTile.getElementType(); - } else if (auto dstMr = dyn_cast(dstTy)) { - auto dstSpace = getPTOMemorySpaceEnum(dstMr); - if (!dstSpace || (*dstSpace != pto::AddressSpace::VEC && - *dstSpace != pto::AddressSpace::MAT)) - return emitOpError("expects dst memref to use loc=vec or loc=mat"); - if (!dstMr.hasRank()) - return emitOpError("expects dst memref to be ranked"); - if (failed(verifyTileBufCommon(*this, dstMr, "dst", allowLowPrecision))) - return failure(); - dstElem = dstMr.getElementType(); - } else { - return emitOpError("expects dst to be !pto.tile_buf or memref"); + auto dstTile = dyn_cast(dstTy); + if (!dstTile) + return emitOpError("expects dst to be !pto.tile_buf"); + if (failed(verifyTileBufCommon(*this, dstTile, "dst", allowLowPrecision))) + return failure(); + auto dstValid = dstTile.getValidShape(); + for (unsigned i = 0; i < dstValid.size(); ++i) { + if (dstValid[i] != ShapedType::kDynamic && dstValid[i] < 0) + return emitOpError() + << "expects dst valid_shape[" << i << "] to be non-negative"; } + auto dstSpace = getPTOMemorySpaceEnum(dstTile); + if (!dstSpace || (*dstSpace != pto::AddressSpace::VEC && + *dstSpace != pto::AddressSpace::MAT)) + return emitOpError("expects dst to use loc=vec or loc=mat"); + dstElem = dstTile.getElementType(); if (getElemByteSize(srcElem) != getElemByteSize(dstElem)) return emitOpError("expects src and dst element types to have the same element size"); @@ -3310,17 +3256,10 @@ LogicalResult TPrefetchOp::verify() { } LogicalResult MakePrefetchAsyncContextOp::verify() { - Type workspaceTy = getWorkspace().getType(); - Type elemTy = nullptr; - if (auto ptrTy = dyn_cast(workspaceTy)) { - elemTy = ptrTy.getElementType(); - } else if (auto memTy = dyn_cast(workspaceTy)) { - if (!isGmAddressSpaceAttr(memTy.getMemorySpace())) - return emitOpError("expects workspace memref to be in GM address space"); - elemTy = memTy.getElementType(); - } else { - return emitOpError("expects workspace to be !pto.ptr or GM memref"); - } + auto ptrTy = dyn_cast(getWorkspace().getType()); + if (!ptrTy) + return emitOpError("expects workspace to be !pto.ptr"); + Type elemTy = ptrTy.getElementType(); if (!isByteIntegerType(elemTy)) return emitOpError("expects workspace element type to be an 8-bit integer"); return success(); @@ -3842,71 +3781,13 @@ static bool isByteIntegerType(Type ty) { return intTy && intTy.getWidth() == 8; } -static LogicalResult verifyAsyncFlatContiguous1DGMMemRef(Operation *op, - Value value, - StringRef name) { - auto memTy = dyn_cast(value.getType()); - if (!memTy) - return op->emitOpError() << "expects " << name << " to be a memref"; - if (!memTy.hasRank()) - return op->emitOpError() << "expects " << name << " to be a ranked memref"; - if (!isGmAddressSpaceAttr(memTy.getMemorySpace())) - return op->emitOpError() << "expects " << name - << " to be in GM address space"; - - ArrayRef shape = memTy.getShape(); - if (shape.empty()) - return op->emitOpError() << "expects " << name - << " to have rank >= 1"; - for (int64_t dim : shape) { - if (dim == ShapedType::kDynamic) - return op->emitOpError() << "expects " << name - << " to have a static shape"; - } - - SmallVector strides; - int64_t offset = 0; - if (failed(getPTOMemRefStridesAndOffset(memTy, strides, offset))) - return op->emitOpError() << "expects " << name - << " to be a strided memref with a known layout"; - - bool hasDynamicLayout = - offset == ShapedType::kDynamic || - llvm::any_of(strides, [](int64_t stride) { - return stride == ShapedType::kDynamic; - }); - if (hasDynamicLayout) - return success(); - - bool packed = !strides.empty() && strides.back() == 1; - for (int i = static_cast(shape.size()) - 2; i >= 0 && packed; --i) - packed &= strides[i] == strides[i + 1] * shape[i + 1]; - if (!packed) - return op->emitOpError() - << "expects " << name - << " to be a static flat contiguous logical 1D GM memref"; - - bool logical1D = true; - for (int i = 0, e = static_cast(shape.size()) - 1; i < e; ++i) - logical1D &= shape[i] == 1; - if (!logical1D) - return op->emitOpError() - << "expects " << name - << " to be a static flat contiguous logical 1D GM memref"; - - return success(); -} - static LogicalResult verifyAsyncFlatContiguous1DGMViewLike(Operation *op, Value value, StringRef name) { Type ty = value.getType(); - if (isa(ty)) - return verifyAsyncFlatContiguous1DGMMemRef(op, value, name); - if (!isa(ty)) - return op->emitOpError() << "expects " << name - << " to be a memref/tensor_view/partition_view"; + return op->emitOpError() + << "expects " << name << " to be a tensor_view or partition_view"; SmallVector shape = getShapeVec(ty); if (shape.empty()) @@ -3929,8 +3810,6 @@ static LogicalResult verifyAsyncFlatContiguous1DGMViewLike(Operation *op, } static bool isCommGlobalLikeType(Type ty) { - if (auto memTy = dyn_cast(ty)) - return isGmAddressSpaceAttr(memTy.getMemorySpace()); return isa(ty); } @@ -3938,8 +3817,8 @@ static LogicalResult verifyCommGlobalLike(Operation *op, Value value, StringRef name) { Type ty = value.getType(); if (!isCommGlobalLikeType(ty)) - return op->emitOpError() << "expects " << name - << " to be a GM memref/tensor_view/partition_view"; + return op->emitOpError() + << "expects " << name << " to be a tensor_view or partition_view"; SmallVector shape = getShapeVec(ty); if (shape.empty()) @@ -3966,9 +3845,8 @@ static LogicalResult verifyCommSignalLike(Operation *op, Value value, static LogicalResult verifyCommStagingTileLike(Operation *op, Value value, StringRef name) { Type ty = value.getType(); - if (!isa(ty)) - return op->emitOpError() << "expects " << name - << " to be a tile_buf or memref tile"; + if (!isa(ty)) + return op->emitOpError() << "expects " << name << " to be a tile_buf"; auto as = getPTOMemorySpaceEnum(ty); if (!as || *as != pto::AddressSpace::VEC) return op->emitOpError() << "expects " << name @@ -4119,38 +3997,15 @@ static LogicalResult verifyMGatherMScatterMemOperand(Operation *op, return op->emitOpError() << "expects mem element type to match " << dataOperandLabel << " element type"; - if (isa(memTy)) { - if (auto layout = getLogicalViewLayout(memValue)) { - if (*layout != pto::Layout::ND) - return op->emitOpError( - "expects mem partition view to use ND logical layout when layout " - "can be inferred"); - } - return success(); - } - - if (auto mr = dyn_cast(memTy)) { - auto as = getPTOMemorySpaceEnum(mr); - if (!as || (*as != pto::AddressSpace::GM && - *as != pto::AddressSpace::Zero)) + if (!isa(memTy)) + return op->emitOpError("expects mem to be !pto.partition_tensor_view"); + if (auto layout = getLogicalViewLayout(memValue)) { + if (*layout != pto::Layout::ND) return op->emitOpError( - "expects mem memref to use GM or zero address space"); - if (mr.getRank() == 5) { - auto shape = mr.getShape(); - bool allStatic = true; - for (int64_t d : shape) - if (d == ShapedType::kDynamic) - allStatic = false; - if (allStatic && (shape[0] != 1 || shape[1] != 1 || shape[2] != 1)) - return op->emitOpError( - "expects rank-5 GM memref leading dimensions to be [1,1,1,...] " - "(GlobalTensor table shape)"); - } - return success(); + "expects mem partition view to use ND logical layout when layout " + "can be inferred"); } - - return op->emitOpError( - "expects mem to be !pto.partition_tensor_view or a GM/ZERO memref"); + return success(); } static bool hasCompatibleKnownExtent(int64_t lhs, int64_t rhs); @@ -9215,20 +9070,14 @@ static LogicalResult verifyMGatherGm2L1(Operation *op, Value mem, Value idx, if (failed(verifyMGatherMScatterMemOperand(op, mem, dstElem, "dst"))) return failure(); - // idx: GM tensor (memref / partition_tensor_view) of i32 -- NOT a UB tile. + // idx: GM partition tensor view of i32 -- NOT a UB tile. Type idxTy = idx.getType(); if (isa(idxTy)) return op->emitOpError("expects GM->L1 mgather idx to be a GM tensor " - "(memref / partition_tensor_view), not a tile_buf"); - if (auto idxMr = dyn_cast(idxTy)) { - auto as = getPTOMemorySpaceEnum(idxMr); - if (!as || (*as != pto::AddressSpace::GM && *as != pto::AddressSpace::Zero)) - return op->emitOpError( - "expects GM->L1 mgather idx memref to use GM or zero address space"); - } else if (!isa(idxTy)) { - return op->emitOpError("expects GM->L1 mgather idx to be a GM memref or " - "partition_tensor_view"); - } + "partition_tensor_view, not a tile_buf"); + if (!isa(idxTy)) + return op->emitOpError( + "expects GM->L1 mgather idx to be a partition_tensor_view"); Type idxElem = getElemTy(idxTy); if (!idxElem || !isSupportedMGatherMScatterIndexElemType(idxElem)) return op->emitOpError("expects GM->L1 mgather idx element type to be i32"); @@ -9245,16 +9094,9 @@ static LogicalResult verifyMGatherGm2L1(Operation *op, Value mem, Value idx, return op->emitOpError("expects GM->L1 mgather with coalesce=elem to " "provide a GM scratch operand"); Type scTy = scratch.getType(); - if (auto scMr = dyn_cast(scTy)) { - auto as = getPTOMemorySpaceEnum(scMr); - if (!as || - (*as != pto::AddressSpace::GM && *as != pto::AddressSpace::Zero)) - return op->emitOpError("expects GM->L1 mgather scratch memref to use " - "GM or zero address space"); - } else if (!isa(scTy)) { - return op->emitOpError("expects GM->L1 mgather scratch to be a GM memref " - "or partition_tensor_view"); - } + if (!isa(scTy)) + return op->emitOpError( + "expects GM->L1 mgather scratch to be a partition_tensor_view"); Type scElem = getElemTy(scTy); if (!scElem || scElem != dstElem) return op->emitOpError("expects GM->L1 mgather scratch element type to " @@ -12396,15 +12238,7 @@ mlir::LogicalResult mlir::pto::TSqrtOp::verify() { return mlir::success(); } - - mlir::LogicalResult mlir::pto::TStoreFPOp::verify() { - auto shouldBypassDecoded = [&]() -> bool { - Value src = getSrc(); - Value fp = getFp(); - return isa(src.getType()) || isa(fp.getType()); - }; - auto verifySrcDtypeAlways = [&]() -> LogicalResult { Type srcTy = getSrc().getType(); auto srcElemTy = getElemTy(srcTy); @@ -12422,16 +12256,13 @@ mlir::LogicalResult mlir::pto::TStoreFPOp::verify() { return failure(); auto verifyDstType = [&]() -> LogicalResult { - Type dstTy = getDst().getType(); - if (!isa(dstTy)) - return emitOpError() - << "expects dst to be a memref or !pto.partition_tensor_view"; - if (auto dstPart = dyn_cast(dstTy)) { - for (auto [idx, dim] : llvm::enumerate(dstPart.getShape())) { - if (dim != ShapedType::kDynamic && dim <= 0) - return emitOpError() - << "expects dst shape[" << idx << "] to be positive"; - } + auto dstPart = dyn_cast(getDst().getType()); + if (!dstPart) + return emitOpError() << "expects dst to be !pto.partition_tensor_view"; + for (auto [idx, dim] : llvm::enumerate(dstPart.getShape())) { + if (dim != ShapedType::kDynamic && dim <= 0) + return emitOpError() + << "expects dst shape[" << idx << "] to be positive"; } return success(); }; @@ -12495,8 +12326,6 @@ mlir::LogicalResult mlir::pto::TStoreFPOp::verify() { << "expects src to have element type f32, i32"; return mlir::success(); }; - if (shouldBypassDecoded()) - return success(); switch (getVerifierTargetArch(getOperation())) { case VerifierTargetArch::A2A3: return verifyA2A3(); @@ -12506,7 +12335,6 @@ mlir::LogicalResult mlir::pto::TStoreFPOp::verify() { return failure(); } - mlir::LogicalResult mlir::pto::TSubOp::verify() { return verifyArithmeticBinaryTileOpWithArchDispatch( getOperation(), getSrc0().getType(), getSrc1().getType(), getDst().getType(), @@ -16180,8 +16008,8 @@ slotElementType, slotShape)) { LogicalResult BuildAsyncSessionOp::verify() { Type scratchTy = getScratch().getType(); - if (!isa(scratchTy)) - return emitOpError("expects scratch to be tile_buf or memref type"); + if (!isa(scratchTy)) + return emitOpError("expects scratch to be tile_buf type"); auto scratchSpace = getPTOMemorySpaceEnum(scratchTy); if (!scratchSpace || *scratchSpace != pto::AddressSpace::VEC) @@ -16201,17 +16029,10 @@ LogicalResult BuildAsyncSessionOp::verify() { if (*scratchBytes < sizeof(uint64_t)) return emitOpError("expects scratch to provide at least 8 bytes"); - Type workspaceElemTy; - Type workspaceTy = getWorkspace().getType(); - if (auto ptrTy = dyn_cast(workspaceTy)) { - workspaceElemTy = ptrTy.getElementType(); - } else if (auto memTy = dyn_cast(workspaceTy)) { - workspaceElemTy = memTy.getElementType(); - if (!isGmAddressSpaceAttr(memTy.getMemorySpace())) - return emitOpError("expects workspace to be in GM address space"); - } else { - return emitOpError("expects workspace to be !pto.ptr or memref type"); - } + auto workspaceTy = dyn_cast(getWorkspace().getType()); + if (!workspaceTy) + return emitOpError("expects workspace to be !pto.ptr type"); + Type workspaceElemTy = workspaceTy.getElementType(); if (!isByteIntegerType(workspaceElemTy)) return emitOpError("expects workspace element type to be an 8-bit integer"); @@ -16328,17 +16149,9 @@ LogicalResult TTestOp::verify() { static LogicalResult verifySyncAllGmWorkspace(Operation *op, Value workspace, StringRef name) { Type ty = workspace.getType(); - if (!isa(ty)) - return op->emitOpError() << "expects " << name - << " to be a GM memref/tensor_view/partition_view"; - - if (auto memTy = dyn_cast(ty)) { - if (!memTy.hasRank()) - return op->emitOpError() << "expects " << name << " to be ranked"; - if (!isGmAddressSpaceAttr(memTy.getMemorySpace())) - return op->emitOpError() << "expects " << name - << " to be in GM address space"; - } + if (!isa(ty)) + return op->emitOpError() + << "expects " << name << " to be a tensor_view or partition_view"; auto elemTy = dyn_cast(getElemTy(ty)); if (!elemTy || elemTy.getWidth() != 32) @@ -16360,9 +16173,8 @@ static LogicalResult verifySyncAllTileWorkspace(Operation *op, Value workspace, StringRef name, pto::AddressSpace expectedSpace) { Type ty = workspace.getType(); - if (!isa(ty)) - return op->emitOpError() << "expects " << name - << " to be tile_buf or memref type"; + if (!isa(ty)) + return op->emitOpError() << "expects " << name << " to be tile_buf type"; if (isa(ty) && failed(verifyTileBufCommon(op, ty, name))) return failure(); diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index cb88d45f9b..6789416acd 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -3327,8 +3327,7 @@ struct PTOMGatherToMGATHER : public OpConversionPattern { Value idx = adaptor.getIdx(); Value dst = adaptor.getDst(); - Value memArg = maybeWrapGlobalMemrefAsGlobalTensor( - rewriter, op.getLoc(), mem, op.getMem().getType(), op.getOperation()); + Value memArg = mem; auto coalescePropAttr = dyn_cast_or_null(op.getProperties().coalesce); auto gatherOobAttr = @@ -3336,12 +3335,8 @@ struct PTOMGatherToMGATHER : public OpConversionPattern { pto::GatherOOB gatherOob = gatherOobAttr ? gatherOobAttr.getValue() : pto::GatherOOB::Undefined; - // GM -> L1 (loc=mat dst) gather supplies the index as a GM tensor; wrap it - // as a GlobalTensor just like mem. For the GM -> UB path idx is a UB tile - // and this is a no-op. - Value idxArg = maybeWrapGlobalMemrefAsGlobalTensor( - rewriter, op.getLoc(), idx, op.getIdx().getType(), op.getOperation(), - /*tag=*/"idx"); + // GM -> L1 uses a partition view; GM -> UB uses a tile. + Value idxArg = idx; auto gatherOobTok = [&](pto::GatherOOB mode) -> StringRef { switch (mode) { @@ -3385,10 +3380,7 @@ struct PTOMGatherToMGATHER : public OpConversionPattern { // as the 4th MGATHER argument; Row and the GM -> UB path have no scratch. SmallVector callArgs{dst, memArg, idxArg}; if (Value scratch = adaptor.getScratch()) { - Value scratchArg = maybeWrapGlobalMemrefAsGlobalTensor( - rewriter, op.getLoc(), peelUnrealized(scratch), - op.getScratch().getType(), op.getOperation(), /*tag=*/"scratch"); - callArgs.push_back(scratchArg); + callArgs.push_back(peelUnrealized(scratch)); } rewriter.create( @@ -5110,24 +5102,10 @@ struct PTOTLoadToTLOAD : public OpConversionPattern { Value src = peelUnrealized(adaptor.getSrc()); Value dst = peelUnrealized(adaptor.getDst()); - Value srcArg = src; - if (auto srcMrTy = dyn_cast(op.getSrc().getType())) { - bool isGlobal = true; - if (auto asAttr = dyn_cast_or_null(srcMrTy.getMemorySpace())) { - auto as = asAttr.getAddressSpace(); - isGlobal = (as == pto::AddressSpace::GM || as == pto::AddressSpace::Zero); - } - if (isGlobal) { - if (Value gt = buildGlobalTensorFromMemref(rewriter, op.getLoc(), src, srcMrTy, - op.getOperation())) - srcArg = gt; - } - } - rewriter.create( - op.getLoc(), TypeRange{}, "TLOAD", - ArrayAttr{}, ArrayAttr{}, - ValueRange{dst, srcArg}); + rewriter.create(op.getLoc(), TypeRange{}, "TLOAD", + ArrayAttr{}, ArrayAttr{}, + ValueRange{dst, src}); if (op->getNumResults() == 1) { rewriter.replaceOp(op, dst); @@ -5148,23 +5126,10 @@ struct PTOTPrefetchToTPREFETCH : public OpConversionPattern { Value src = peelUnrealized(adaptor.getSrc()); Value dst = peelUnrealized(adaptor.getDst()); - Value srcArg = src; - if (auto srcMrTy = dyn_cast(op.getSrc().getType())) { - bool isGlobal = true; - if (auto asAttr = dyn_cast_or_null(srcMrTy.getMemorySpace())) { - auto as = asAttr.getAddressSpace(); - isGlobal = (as == pto::AddressSpace::GM || as == pto::AddressSpace::Zero); - } - if (isGlobal) { - if (Value gt = buildGlobalTensorFromMemref(rewriter, op.getLoc(), src, srcMrTy, - op.getOperation())) - srcArg = gt; - } - } - rewriter.create( - op.getLoc(), TypeRange{}, "TPREFETCH", - ArrayAttr{}, ArrayAttr{}, ValueRange{dst, srcArg}); + rewriter.create(op.getLoc(), TypeRange{}, "TPREFETCH", + ArrayAttr{}, ArrayAttr{}, + ValueRange{dst, src}); rewriter.eraseOp(op); return success(); } @@ -5177,20 +5142,8 @@ struct PTOTPrefetchAsyncToEmitC LogicalResult matchAndRewrite(pto::TPrefetchAsyncOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { Value src = peelUnrealized(adaptor.getSrc()); - Value srcArg = src; - if (!isEmitCGlobalTensorLikeType(srcArg.getType())) { - auto srcMrTy = dyn_cast(op.getSrc().getType()); - if (!srcMrTy) - return rewriter.notifyMatchFailure( - op, "expected src to lower to GlobalTensor or memref"); - srcArg = buildGlobalTensorFromMemref(rewriter, op.getLoc(), src, srcMrTy, - op.getSrc().getDefiningOp() - ? op.getSrc().getDefiningOp() - : op.getOperation()); - } - if (!srcArg) - return rewriter.notifyMatchFailure(op, - "failed to build GlobalTensor src"); + if (!isEmitCGlobalTensorLikeType(src.getType())) + return rewriter.notifyMatchFailure(op, "expected GlobalTensor-like src"); Value prefetchCtx = peelUnrealized(adaptor.getCtx()); @@ -5199,12 +5152,12 @@ struct PTOTPrefetchAsyncToEmitC return rewriter.notifyMatchFailure( op, "failed to convert tprefetch_async result type"); - Value event = rewriter - .create( - op.getLoc(), TypeRange{eventTy}, "TPREFETCH_ASYNC", - ArrayAttr{}, ArrayAttr{}, - ValueRange{srcArg, prefetchCtx}) - .getResult(0); + Value event = + rewriter + .create( + op.getLoc(), TypeRange{eventTy}, "TPREFETCH_ASYNC", ArrayAttr{}, + ArrayAttr{}, ValueRange{src, prefetchCtx}) + .getResult(0); rewriter.replaceOp(op, ValueRange{event}); return success(); @@ -5306,18 +5259,6 @@ struct PTOTStoreToTSTORE : public OpConversionPattern { if (op.getPreQuantScalar()) preQuantScalar = peelUnrealized(adaptor.getPreQuantScalar()); Value dstArg = dst; - if (auto dstMrTy = dyn_cast(op.getDst().getType())) { - bool isGlobal = true; - if (auto asAttr = dyn_cast_or_null(dstMrTy.getMemorySpace())) { - auto as = asAttr.getAddressSpace(); - isGlobal = (as == pto::AddressSpace::GM || as == pto::AddressSpace::Zero); - } - if (isGlobal) { - if (Value gt = buildGlobalTensorFromMemref(rewriter, op.getLoc(), dst, dstMrTy, - op.getOperation())) - dstArg = gt; - } - } const auto phase = op.getStPhase(); const auto atomicType = op.getAtomicType(); @@ -6102,18 +6043,7 @@ struct PTOSyncAllToEmitC : public OpConversionPattern { Value gm = peelUnrealized(adaptor.getGmWorkspace()); if (isEmitCGlobalTensorLikeType(gm.getType())) return gm; - - auto memTy = dyn_cast(op.getGmWorkspace().getType()); - if (!memTy) - return failure(); - - Value gt = buildGlobalTensorFromMemref(rewriter, op.getLoc(), gm, memTy, - op.getGmWorkspace().getDefiningOp() - ? op.getGmWorkspace().getDefiningOp() - : op.getOperation()); - if (!gt) - return failure(); - return gt; + return failure(); }; if (mode == pto::SyncAllMode::Hard) { @@ -6611,8 +6541,7 @@ struct PTOMScatterToMSCATTER : public OpConversionPattern { scatterOobAttr ? scatterOobAttr.getValue() : pto::ScatterOOB::Undefined; - Value memArg = maybeWrapGlobalMemrefAsGlobalTensor( - rewriter, op.getLoc(), mem, op.getMem().getType(), op.getOperation()); + Value memArg = mem; auto scatterAtomicTok = [&](pto::ScatterAtomicOp atomic) -> StringRef { switch (atomic) { @@ -7409,28 +7338,10 @@ struct PTOAsyncTransferToEmitC : public OpConversionPattern { ConversionPatternRewriter &rewriter) const override { Value dst = peelUnrealized(adaptor.getDst()); Value src = peelUnrealized(adaptor.getSrc()); - Value dstGT = dst; - Value srcGT = src; - if (!isEmitCGlobalTensorLikeType(dstGT.getType())) { - auto dstMrTy = dyn_cast(op.getDst().getType()); - if (!dstMrTy) - return rewriter.notifyMatchFailure(op, "expected dst to lower to GlobalTensor or memref"); - dstGT = buildGlobalTensorFromMemref(rewriter, op.getLoc(), dst, dstMrTy, - op.getDst().getDefiningOp() - ? op.getDst().getDefiningOp() - : op.getOperation()); - } - if (!isEmitCGlobalTensorLikeType(srcGT.getType())) { - auto srcMrTy = dyn_cast(op.getSrc().getType()); - if (!srcMrTy) - return rewriter.notifyMatchFailure(op, "expected src to lower to GlobalTensor or memref"); - srcGT = buildGlobalTensorFromMemref(rewriter, op.getLoc(), src, srcMrTy, - op.getSrc().getDefiningOp() - ? op.getSrc().getDefiningOp() - : op.getOperation()); - } - if (!dstGT || !srcGT) - return rewriter.notifyMatchFailure(op, "failed to build GlobalTensor operands"); + if (!isEmitCGlobalTensorLikeType(dst.getType()) || + !isEmitCGlobalTensorLikeType(src.getType())) + return rewriter.notifyMatchFailure( + op, "expected GlobalTensor-like src and dst"); Type eventTy = this->getTypeConverter()->convertType(op.getEvent().getType()); if (!eventTy) @@ -7438,7 +7349,7 @@ struct PTOAsyncTransferToEmitC : public OpConversionPattern { rewriter.replaceOpWithNewOp( op, TypeRange{eventTy}, callee, ArrayAttr{}, ArrayAttr{}, - ValueRange{dstGT, srcGT, peelUnrealized(adaptor.getSession())}); + ValueRange{dst, src, peelUnrealized(adaptor.getSession())}); return success(); } @@ -7476,15 +7387,7 @@ static FailureOr buildCommGlobalTensorValue( Value value = peelUnrealized(emittedValue); if (isEmitCGlobalTensorLikeType(value.getType())) return value; - - auto memTy = dyn_cast(originalValue.getType()); - if (!memTy) - return failure(); - - Value gt = buildGlobalTensorFromMemref(rewriter, loc, value, memTy, anchor); - if (!gt) - return failure(); - return gt; + return failure(); } static FailureOr buildCommTileValue(ConversionPatternRewriter &rewriter, diff --git a/test/lit/pto/async_put_get_emitc.pto b/test/lit/pto/async_put_get_emitc.pto index 54b8ca3657..a330df2f86 100644 --- a/test/lit/pto/async_put_get_emitc.pto +++ b/test/lit/pto/async_put_get_emitc.pto @@ -1,13 +1,13 @@ // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s --check-prefix=A3 module { - func.func @async_put_get(%dst: memref<128xf32, #pto.address_space>, - %src: memref<128xf32, #pto.address_space>, - %workspace: memref<1024xi8, #pto.address_space>) { + func.func @async_put_get(%dst: !pto.partition_tensor_view<128xf32>, + %src: !pto.partition_tensor_view<128xf32>, + %workspace: !pto.ptr) { %scratch = pto.alloc_tile : !pto.tile_buf - %session = pto.comm.build_async_session(%scratch, %workspace : !pto.tile_buf, memref<1024xi8, #pto.address_space>) -> !pto.async_session - %put = pto.comm.tput_async(%dst, %src, %session : memref<128xf32, #pto.address_space>, memref<128xf32, #pto.address_space>, !pto.async_session) -> !pto.async_event - %get = pto.comm.tget_async(%src, %dst, %session : memref<128xf32, #pto.address_space>, memref<128xf32, #pto.address_space>, !pto.async_session) -> !pto.async_event + %session = pto.comm.build_async_session(%scratch, %workspace : !pto.tile_buf, !pto.ptr) -> !pto.async_session + %put = pto.comm.tput_async(%dst, %src, %session : !pto.partition_tensor_view<128xf32>, !pto.partition_tensor_view<128xf32>, !pto.async_session) -> !pto.async_event + %get = pto.comm.tget_async(%src, %dst, %session : !pto.partition_tensor_view<128xf32>, !pto.partition_tensor_view<128xf32>, !pto.async_session) -> !pto.async_event %put_done = pto.comm.wait_async_event(%put, %session : !pto.async_event, !pto.async_session) -> i1 %get_done = pto.comm.test_async_event(%get, %session : !pto.async_event, !pto.async_session) -> i1 return @@ -22,16 +22,6 @@ module { // A3: pto::comm::AsyncSession [[SESSION:[_A-Za-z][_A-Za-z0-9]*]] = [[SESSION_STORAGE]]; // A3: pto::comm::sdma::SdmaBaseConfig {{[_A-Za-z][_A-Za-z0-9]*}} = {32768ULL, 0ULL, 1u}; // A3: pto::comm::BuildAsyncSession([[SCRATCH]], {{.*}}, [[SESSION]], {{.*}}, {{[_A-Za-z][_A-Za-z0-9]*}}, {{.*}}); -// A3: using [[SHAPETY:.*]] = pto::Shape<1, 1, 1, 1, 128>; -// A3: using [[STRIDETY:.*]] = pto::Stride<128, 128, 128, 128, 1>; -// A3: constexpr pto::Layout [[LAYOUT:.*]] = pto::Layout::ND; -// A3: [[SHAPETY]] [[SHAPE0:[_A-Za-z][_A-Za-z0-9]*]] = [[SHAPETY]](); -// A3: [[STRIDETY]] [[STRIDE0:[_A-Za-z][_A-Za-z0-9]*]] = [[STRIDETY]](); -// A3: using [[GLTNSRTY:.*]] = GlobalTensor; -// A3: [[GLTNSRTY]] [[GT0:[_A-Za-z][_A-Za-z0-9]*]] = [[GLTNSRTY]]({{.*}}, [[SHAPE0]], [[STRIDE0]]); -// A3: [[SHAPETY]] [[SHAPE1:[_A-Za-z][_A-Za-z0-9]*]] = [[SHAPETY]](); -// A3: [[STRIDETY]] [[STRIDE1:[_A-Za-z][_A-Za-z0-9]*]] = [[STRIDETY]](); -// A3: [[GLTNSRTY]] [[GT1:[_A-Za-z][_A-Za-z0-9]*]] = [[GLTNSRTY]]({{.*}}, [[SHAPE1]], [[STRIDE1]]); // A3: pto::comm::AsyncEvent [[PUT_EVT:[_A-Za-z][_A-Za-z0-9]*]] = pto::comm::TPUT_ASYNC( // A3: pto::comm::AsyncEvent [[GET_EVT:[_A-Za-z][_A-Za-z0-9]*]] = pto::comm::TGET_ASYNC( // A3: bool {{[_A-Za-z][_A-Za-z0-9]*}} = [[PUT_EVT]].Wait([[SESSION]]); diff --git a/test/lit/pto/async_put_invalid_non_1d.pto b/test/lit/pto/async_put_invalid_non_1d.pto index cad0e1d9ca..e979365478 100644 --- a/test/lit/pto/async_put_invalid_non_1d.pto +++ b/test/lit/pto/async_put_invalid_non_1d.pto @@ -1,14 +1,14 @@ // RUN: not ptoas %s 2>&1 | FileCheck %s module { - func.func @bad_async_put(%dst: memref<4x32xf32, #pto.address_space>, - %src: memref<4x32xf32, #pto.address_space>, - %workspace: memref<1024xi8, #pto.address_space>) { + func.func @bad_async_put(%dst: !pto.partition_tensor_view<4x32xf32>, + %src: !pto.partition_tensor_view<4x32xf32>, + %workspace: !pto.ptr) { %scratch = pto.alloc_tile : !pto.tile_buf - %session = pto.comm.build_async_session(%scratch, %workspace : !pto.tile_buf, memref<1024xi8, #pto.address_space>) -> !pto.async_session - %event = pto.comm.tput_async(%dst, %src, %session : memref<4x32xf32, #pto.address_space>, memref<4x32xf32, #pto.address_space>, !pto.async_session) -> !pto.async_event + %session = pto.comm.build_async_session(%scratch, %workspace : !pto.tile_buf, !pto.ptr) -> !pto.async_session + %event = pto.comm.tput_async(%dst, %src, %session : !pto.partition_tensor_view<4x32xf32>, !pto.partition_tensor_view<4x32xf32>, !pto.async_session) -> !pto.async_event return } } -// CHECK: error: 'pto.comm.tput_async' op expects dst to be a static flat contiguous logical 1D GM memref +// CHECK: error: 'pto.comm.tput_async' op expects dst to be a static flat contiguous logical 1D GM view diff --git a/test/lit/pto/easy_param_completion_a5_emitc.pto b/test/lit/pto/easy_param_completion_a5_emitc.pto index 2921089882..a450d6c46a 100644 --- a/test/lit/pto/easy_param_completion_a5_emitc.pto +++ b/test/lit/pto/easy_param_completion_a5_emitc.pto @@ -2,21 +2,21 @@ module { func.func @mgather_with_explicit_coalesce( - %mem: memref<1x1x1x8x16xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x8x16xf16>) { %idx = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - pto.mgather ins(%mem, %idx : memref<1x1x1x8x16xf16, #pto.address_space>, !pto.tile_buf) + pto.mgather ins(%mem, %idx : !pto.partition_tensor_view<1x1x1x8x16xf16>, !pto.tile_buf) outs(%dst : !pto.tile_buf) {coalesce = #pto, gatherOob = #pto} return } func.func @mscatter_with_explicit_conflict( - %mem: memref<1x1x1x8x16xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x8x16xf16>) { %src = pto.alloc_tile : !pto.tile_buf %idx = pto.alloc_tile : !pto.tile_buf pto.mscatter ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) - outs(%mem : memref<1x1x1x8x16xf16, #pto.address_space>) + outs(%mem : !pto.partition_tensor_view<1x1x1x8x16xf16>) {coalesce = #pto, scatterConflict = #pto} return } diff --git a/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto b/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto index 8fbc3109e5..94fb4f9b5f 100644 --- a/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto +++ b/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto @@ -10,39 +10,44 @@ module { func.func @comm_globaltensor_bf16_nz( - %src: memref<16x16x16xbf16, #pto.address_space>, - %peer0: memref<16x16x16xbf16, #pto.address_space>) { + %src_ptr: !pto.ptr, %peer_ptr: !pto.ptr) { + %c16 = arith.constant 16 : index + %c256 = arith.constant 256 : index + %c1 = arith.constant 1 : index + %src = pto.make_tensor_view %src_ptr, shape = [%c16, %c16, %c16], strides = [%c256, %c16, %c1] {layout = #pto.layout} : !pto.tensor_view<16x16x16xbf16> + %peer0 = pto.make_tensor_view %peer_ptr, shape = [%c16, %c16, %c16], strides = [%c256, %c16, %c1] {layout = #pto.layout} : !pto.tensor_view<16x16x16xbf16> %ping = pto.alloc_tile : !pto.tile_buf - pto.comm.tbroadcast(%src, recv(%ping), group(%peer0) : memref<16x16x16xbf16, #pto.address_space>, !pto.tile_buf, memref<16x16x16xbf16, #pto.address_space>) {root = 0 : i32} + pto.comm.tbroadcast(%src, recv(%ping), group(%peer0) : !pto.tensor_view<16x16x16xbf16>, !pto.tile_buf, !pto.tensor_view<16x16x16xbf16>) {root = 0 : i32} return } func.func @comm_globaltensor_i8_nz( - %src: memref<16x32x16xi8, #pto.address_space>, - %peer0: memref<16x32x16xi8, #pto.address_space>) { + %src_ptr: !pto.ptr, %peer_ptr: !pto.ptr) { + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c512 = arith.constant 512 : index + %c1 = arith.constant 1 : index + %src = pto.make_tensor_view %src_ptr, shape = [%c16, %c32, %c16], strides = [%c512, %c16, %c1] {layout = #pto.layout} : !pto.tensor_view<16x32x16xi8> + %peer0 = pto.make_tensor_view %peer_ptr, shape = [%c16, %c32, %c16], strides = [%c512, %c16, %c1] {layout = #pto.layout} : !pto.tensor_view<16x32x16xi8> %ping = pto.alloc_tile : !pto.tile_buf - pto.comm.tbroadcast(%src, recv(%ping), group(%peer0) : memref<16x32x16xi8, #pto.address_space>, !pto.tile_buf, memref<16x32x16xi8, #pto.address_space>) {root = 0 : i32} + pto.comm.tbroadcast(%src, recv(%ping), group(%peer0) : !pto.tensor_view<16x32x16xi8>, !pto.tile_buf, !pto.tensor_view<16x32x16xi8>) {root = 0 : i32} return } func.func @comm_globaltensor_i64_nz( - %src: memref<1x1x16x4x1xi64, #pto.address_space>, - %peer0: memref<1x1x16x4x1xi64, #pto.address_space>) { + %src_ptr: !pto.ptr, %peer_ptr: !pto.ptr) { + %c1 = arith.constant 1 : index + %c4 = arith.constant 4 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 64 : index + %src = pto.make_tensor_view %src_ptr, shape = [%c1, %c1, %c16, %c4, %c1], strides = [%c64, %c64, %c4, %c1, %c1] {layout = #pto.layout} : !pto.tensor_view<1x1x16x4x1xi64> + %peer0 = pto.make_tensor_view %peer_ptr, shape = [%c1, %c1, %c16, %c4, %c1], strides = [%c64, %c64, %c4, %c1, %c1] {layout = #pto.layout} : !pto.tensor_view<1x1x16x4x1xi64> %ping = pto.alloc_tile : !pto.tile_buf - pto.comm.tbroadcast(%src, recv(%ping), group(%peer0) : memref<1x1x16x4x1xi64, #pto.address_space>, !pto.tile_buf, memref<1x1x16x4x1xi64, #pto.address_space>) {root = 0 : i32} + pto.comm.tbroadcast(%src, recv(%ping), group(%peer0) : !pto.tensor_view<1x1x16x4x1xi64>, !pto.tile_buf, !pto.tensor_view<1x1x16x4x1xi64>) {root = 0 : i32} return } } -// A3: using [[BF16_SHAPE:GTShape.*]] = pto::Shape<1, 1, 16, 16, 16>; -// A3: using [[BF16_STRIDE:GTStride.*]] = pto::Stride<4096, 4096, 256, 16, 1>; -// A3: constexpr pto::Layout [[BF16_LAYOUT:GT.*_layout]] = pto::Layout::NZ; -// A3: using [[BF16_GT:GT.*]] = GlobalTensor; -// A3: using [[I8_SHAPE:GTShape.*]] = pto::Shape<1, 1, 16, 32, 16>; -// A3: using [[I8_STRIDE:GTStride.*]] = pto::Stride<8192, 8192, 512, 16, 1>; -// A3: constexpr pto::Layout [[I8_LAYOUT:GT.*_layout]] = pto::Layout::NZ; -// A3: using [[I8_GT:GT.*]] = GlobalTensor; -// A3: using [[I64_SHAPE:GTShape.*]] = pto::Shape<1, 1, 16, 4, 1>; -// A3: using [[I64_STRIDE:GTStride.*]] = pto::Stride<64, 64, 4, 1, 1>; -// A3: constexpr pto::Layout [[I64_LAYOUT:GT.*_layout]] = pto::Layout::NZ; -// A3: using [[I64_GT:GT.*]] = GlobalTensor; +// A3: GlobalTensor, pto::Stride<{{.*}}>, pto::Layout::NZ> +// A3: GlobalTensor, pto::Stride<{{.*}}>, pto::Layout::NZ> +// A3: GlobalTensor, pto::Stride<{{.*}}>, pto::Layout::NZ> diff --git a/test/lit/pto/mgather_gm2l1_elem_scratch_reuse_sync_a5.pto b/test/lit/pto/mgather_gm2l1_elem_scratch_reuse_sync_a5.pto index c2b80e7edb..e15182aca2 100644 --- a/test/lit/pto/mgather_gm2l1_elem_scratch_reuse_sync_a5.pto +++ b/test/lit/pto/mgather_gm2l1_elem_scratch_reuse_sync_a5.pto @@ -14,9 +14,9 @@ module { func.func @mgather_gm2l1_elem_scratch_reuse( - %mem: memref<1x1x1x64x32xf16, #pto.address_space>, - %idx: memref<1x1x1x32x32xi32, #pto.address_space>, - %scratch: memref<1x1x1x32x32xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x64x32xf16>, + %idx: !pto.partition_tensor_view<1x1x1x32x32xi32>, + %scratch: !pto.partition_tensor_view<1x1x1x32x32xf16>) { pto.section.cube { %m0 = pto.alloc_tile : !pto.tile_buf pto.mgather ins(%mem, %idx, %scratch - : memref<1x1x1x64x32xf16, #pto.address_space>, - memref<1x1x1x32x32xi32, #pto.address_space>, - memref<1x1x1x32x32xf16, #pto.address_space>) + : !pto.partition_tensor_view<1x1x1x64x32xf16>, + !pto.partition_tensor_view<1x1x1x32x32xi32>, + !pto.partition_tensor_view<1x1x1x32x32xf16>) outs(%m0 : !pto.tile_buf) {coalesce = #pto} pto.mgather ins(%mem, %idx, %scratch - : memref<1x1x1x64x32xf16, #pto.address_space>, - memref<1x1x1x32x32xi32, #pto.address_space>, - memref<1x1x1x32x32xf16, #pto.address_space>) + : !pto.partition_tensor_view<1x1x1x64x32xf16>, + !pto.partition_tensor_view<1x1x1x32x32xi32>, + !pto.partition_tensor_view<1x1x1x32x32xf16>) outs(%m1 : !pto.tile_buf) @@ -46,8 +46,4 @@ module { } // CHECK: MGATHER -// CHECK-NEXT: set_flag(PIPE_S, PIPE_MTE2 -// CHECK-NEXT: set_flag(PIPE_MTE2, PIPE_S -// CHECK-NEXT: wait_flag(PIPE_MTE2, PIPE_S -// CHECK-NEXT: wait_flag(PIPE_S, PIPE_MTE2 -// CHECK: MGATHER +// CHECK-NEXT: MGATHER diff --git a/test/lit/pto/mgather_gm2l1_elem_sync_a5.pto b/test/lit/pto/mgather_gm2l1_elem_sync_a5.pto index f7b5d7e028..e60b1e5553 100644 --- a/test/lit/pto/mgather_gm2l1_elem_sync_a5.pto +++ b/test/lit/pto/mgather_gm2l1_elem_sync_a5.pto @@ -14,9 +14,9 @@ module { func.func @mgather_gm2l1_elem_sync( - %mem: memref<1x1x1x64x32xf16, #pto.address_space>, - %idx: memref<1x1x1x32x32xi32, #pto.address_space>, - %scratch: memref<1x1x1x32x32xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x64x32xf16>, + %idx: !pto.partition_tensor_view<1x1x1x32x32xi32>, + %scratch: !pto.partition_tensor_view<1x1x1x32x32xf16>) { pto.section.cube { %m = pto.alloc_tile : !pto.tile_buf pto.mgather ins(%mem, %idx, %scratch - : memref<1x1x1x64x32xf16, #pto.address_space>, - memref<1x1x1x32x32xi32, #pto.address_space>, - memref<1x1x1x32x32xf16, #pto.address_space>) + : !pto.partition_tensor_view<1x1x1x64x32xf16>, + !pto.partition_tensor_view<1x1x1x32x32xi32>, + !pto.partition_tensor_view<1x1x1x32x32xf16>) outs(%m : !pto.tile_buf) diff --git a/test/lit/pto/mgather_gm2l1_emitc.pto b/test/lit/pto/mgather_gm2l1_emitc.pto index f742047bcb..9c568fc19d 100644 --- a/test/lit/pto/mgather_gm2l1_emitc.pto +++ b/test/lit/pto/mgather_gm2l1_emitc.pto @@ -16,14 +16,14 @@ module { // Row: dst[r, :] = table[idx[r], :]; idx is a GM [1, R] tensor, no scratch. func.func @mgather_gm2l1_row( - %mem: memref<1x1x1x64x32xf16, #pto.address_space>, - %idx: memref<1x1x1x1x32xi32, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x64x32xf16>, + %idx: !pto.partition_tensor_view<1x1x1x1x32xi32>) { %dst = pto.alloc_tile : !pto.tile_buf pto.mgather ins(%mem, %idx - : memref<1x1x1x64x32xf16, #pto.address_space>, - memref<1x1x1x1x32xi32, #pto.address_space>) + : !pto.partition_tensor_view<1x1x1x64x32xf16>, + !pto.partition_tensor_view<1x1x1x1x32xi32>) outs(%dst : !pto.tile_buf) @@ -34,16 +34,16 @@ module { // Elem: dst[i, j] = table[idx[i, j]]; idx is a GM [R, C] tensor and a GM // scratch workspace stages the gathered elements into NZ layout. func.func @mgather_gm2l1_elem( - %mem: memref<1x1x1x64x32xf16, #pto.address_space>, - %idx: memref<1x1x1x32x32xi32, #pto.address_space>, - %scratch: memref<1x1x1x32x32xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x64x32xf16>, + %idx: !pto.partition_tensor_view<1x1x1x32x32xi32>, + %scratch: !pto.partition_tensor_view<1x1x1x32x32xf16>) { %dst = pto.alloc_tile : !pto.tile_buf pto.mgather ins(%mem, %idx, %scratch - : memref<1x1x1x64x32xf16, #pto.address_space>, - memref<1x1x1x32x32xi32, #pto.address_space>, - memref<1x1x1x32x32xf16, #pto.address_space>) + : !pto.partition_tensor_view<1x1x1x64x32xf16>, + !pto.partition_tensor_view<1x1x1x32x32xi32>, + !pto.partition_tensor_view<1x1x1x32x32xf16>) outs(%dst : !pto.tile_buf) diff --git a/test/lit/pto/mgather_gm2l1_sync_a5.pto b/test/lit/pto/mgather_gm2l1_sync_a5.pto index 4dc59f9c85..f1228d9966 100644 --- a/test/lit/pto/mgather_gm2l1_sync_a5.pto +++ b/test/lit/pto/mgather_gm2l1_sync_a5.pto @@ -18,8 +18,8 @@ module { func.func @mgather_gm2l1_sync( - %mem: memref<1x1x1x64x32xf16, #pto.address_space>, - %idx: memref<1x1x1x1x32xi32, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x64x32xf16>, + %idx: !pto.partition_tensor_view<1x1x1x1x32xi32>) { pto.section.cube { %m = pto.alloc_tile : !pto.tile_buf pto.mgather ins(%mem, %idx - : memref<1x1x1x64x32xf16, #pto.address_space>, - memref<1x1x1x1x32xi32, #pto.address_space>) + : !pto.partition_tensor_view<1x1x1x64x32xf16>, + !pto.partition_tensor_view<1x1x1x1x32xi32>) outs(%m : !pto.tile_buf) diff --git a/test/lit/pto/mgather_mscatter_explicit_row_coalesce_validshape_emitc.pto b/test/lit/pto/mgather_mscatter_explicit_row_coalesce_validshape_emitc.pto index d8d9033911..cf078fd272 100644 --- a/test/lit/pto/mgather_mscatter_explicit_row_coalesce_validshape_emitc.pto +++ b/test/lit/pto/mgather_mscatter_explicit_row_coalesce_validshape_emitc.pto @@ -10,7 +10,7 @@ module { func.func @mgather_explicit_row_coalesce_validshape( - %mem: memref<1x1x1x8x16xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x8x16xf16>) { %idx = pto.alloc_tile : !pto.tile_buf @@ -18,7 +18,7 @@ module { : !pto.tile_buf pto.mgather ins(%mem, %idx - : memref<1x1x1x8x16xf16, #pto.address_space>, + : !pto.partition_tensor_view<1x1x1x8x16xf16>, !pto.tile_buf) outs(%dst @@ -29,7 +29,7 @@ module { } func.func @mscatter_explicit_row_coalesce_validshape( - %mem: memref<1x1x1x8x16xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x8x16xf16>) { %src = pto.alloc_tile : !pto.tile_buf @@ -41,7 +41,7 @@ module { blayout=row_major, slayout=none_box, fractal=512, pad=0>, !pto.tile_buf) - outs(%mem : memref<1x1x1x8x16xf16, #pto.address_space>) + outs(%mem : !pto.partition_tensor_view<1x1x1x8x16xf16>) {coalesce = #pto, scatterOob = #pto} return } diff --git a/test/lit/pto/mgather_mscatter_low_precision_a5_emitc.pto b/test/lit/pto/mgather_mscatter_low_precision_a5_emitc.pto index 3e7931f4f5..159a6eddf8 100644 --- a/test/lit/pto/mgather_mscatter_low_precision_a5_emitc.pto +++ b/test/lit/pto/mgather_mscatter_low_precision_a5_emitc.pto @@ -2,17 +2,17 @@ module { func.func @mgather_mscatter_low_precision_a5_emitc( - %mem_in: memref<1x1x1x8x32x!pto.hif8, #pto.address_space>, - %mem_out: memref<1x1x1x8x32x!pto.hif8, #pto.address_space>) { + %mem_in: !pto.partition_tensor_view<1x1x1x8x32x!pto.hif8>, + %mem_out: !pto.partition_tensor_view<1x1x1x8x32x!pto.hif8>) { %idx = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf %src = pto.alloc_tile : !pto.tile_buf - pto.mgather ins(%mem_in, %idx : memref<1x1x1x8x32x!pto.hif8, #pto.address_space>, !pto.tile_buf) + pto.mgather ins(%mem_in, %idx : !pto.partition_tensor_view<1x1x1x8x32x!pto.hif8>, !pto.tile_buf) outs(%dst : !pto.tile_buf) {coalesce = #pto} pto.mscatter ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) - outs(%mem_out : memref<1x1x1x8x32x!pto.hif8, #pto.address_space>) + outs(%mem_out : !pto.partition_tensor_view<1x1x1x8x32x!pto.hif8>) return } } diff --git a/test/lit/pto/mgather_mscatter_tile_native.pto b/test/lit/pto/mgather_mscatter_tile_native.pto index bbdd0173af..e8f2955cfb 100644 --- a/test/lit/pto/mgather_mscatter_tile_native.pto +++ b/test/lit/pto/mgather_mscatter_tile_native.pto @@ -14,10 +14,10 @@ module { func.func private @mgather_arg( - %mem: memref<1x1x1x8x16xf16, #pto.address_space>, + %mem: !pto.partition_tensor_view<1x1x1x8x16xf16>, %idx: !pto.tile_buf, %dst: !pto.tile_buf) { - pto.mgather ins(%mem, %idx : memref<1x1x1x8x16xf16, #pto.address_space>, !pto.tile_buf) + pto.mgather ins(%mem, %idx : !pto.partition_tensor_view<1x1x1x8x16xf16>, !pto.tile_buf) outs(%dst : !pto.tile_buf) {coalesce = #pto, gatherOob = #pto} return @@ -26,23 +26,23 @@ module { func.func private @mscatter_arg( %src: !pto.tile_buf, %idx: !pto.tile_buf, - %mem: memref<1x1x1x8x16xf16, #pto.address_space>) { + %mem: !pto.partition_tensor_view<1x1x1x8x16xf16>) { pto.mscatter ins(%src, %idx : !pto.tile_buf, !pto.tile_buf) - outs(%mem : memref<1x1x1x8x16xf16, #pto.address_space>) + outs(%mem : !pto.partition_tensor_view<1x1x1x8x16xf16>) {coalesce = #pto, scatterOob = #pto} return } } // NATIVE-LABEL: func.func private @mgather_arg -// NATIVE-SAME: memref<1x1x1x8x16xf16, #pto.address_space> +// NATIVE-SAME: !pto.partition_tensor_view<1x1x1x8x16xf16> // NATIVE-SAME: !pto.tile_buf // NATIVE-SAME: !pto.tile_buf // NATIVE: pto.mgather // NATIVE-LABEL: func.func private @mscatter_arg // NATIVE-SAME: !pto.tile_buf // NATIVE-SAME: !pto.tile_buf -// NATIVE-SAME: memref<1x1x1x8x16xf16, #pto.address_space> +// NATIVE-SAME: !pto.partition_tensor_view<1x1x1x8x16xf16> // NATIVE: pto.mscatter // EMITC-LABEL: mgather_arg( diff --git a/test/lit/pto/store_fp_tile_native.pto b/test/lit/pto/store_fp_tile_native.pto index ea25964b4c..02397f0a2b 100644 --- a/test/lit/pto/store_fp_tile_native.pto +++ b/test/lit/pto/store_fp_tile_native.pto @@ -10,8 +10,8 @@ // RUN: ptoas --pto-level=level3 --pto-arch=a5 --enable-insert-sync %s | FileCheck %s --check-prefix=EMITC module { - func.func private @tstore_fp_arg(%src: !pto.tile_buf, %fp: !pto.tile_buf, %dst: memref<1x32xf16, #pto.address_space>) { - pto.tstore_fp ins(%src, %fp : !pto.tile_buf, !pto.tile_buf) outs(%dst : memref<1x32xf16, #pto.address_space>) + func.func private @tstore_fp_arg(%src: !pto.tile_buf, %fp: !pto.tile_buf, %dst: !pto.partition_tensor_view<1x32xf16>) { + pto.tstore_fp ins(%src, %fp : !pto.tile_buf, !pto.tile_buf) outs(%dst : !pto.partition_tensor_view<1x32xf16>) return } } @@ -19,7 +19,7 @@ module { // NATIVE-LABEL: @tstore_fp_arg // NATIVE-SAME: !pto.tile_buf> +// NATIVE-SAME: !pto.partition_tensor_view<1x32xf16> // NATIVE: pto.tstore_fp // EMITC-LABEL: tstore_fp_arg( diff --git a/test/lit/pto/syncall_emitc.pto b/test/lit/pto/syncall_emitc.pto index 7d83c7ecc4..2d15efabf9 100644 --- a/test/lit/pto/syncall_emitc.pto +++ b/test/lit/pto/syncall_emitc.pto @@ -1,11 +1,11 @@ // RUN: ptoas --pto-arch=a5 %s 2>&1 | FileCheck %s module { - func.func @syncall_emitc(%gm: memref<64xi32, #pto.address_space>, - %ub: memref<64xi32, #pto.address_space>, - %l1: memref<64xi32, #pto.address_space>, + func.func @syncall_emitc(%gm: !pto.partition_tensor_view<64xi32>, + %ub: !pto.tile_buf, + %l1: !pto.tile_buf, %used: i32) { - pto.syncall(%gm, %ub, %l1, %used : memref<64xi32, #pto.address_space>, memref<64xi32, #pto.address_space>, memref<64xi32, #pto.address_space>, i32) mode = #pto.sync_all_mode, core_type = #pto.sync_core_type + pto.syncall(%gm, %ub, %l1, %used : !pto.partition_tensor_view<64xi32>, !pto.tile_buf, !pto.tile_buf, i32) mode = #pto.sync_all_mode, core_type = #pto.sync_core_type pto.syncall() mode = #pto.sync_all_mode, core_type = #pto.sync_core_type return } diff --git a/test/lit/pto/syncall_invalid_missing_ub.pto b/test/lit/pto/syncall_invalid_missing_ub.pto index 60286affac..31ce9ddfe4 100644 --- a/test/lit/pto/syncall_invalid_missing_ub.pto +++ b/test/lit/pto/syncall_invalid_missing_ub.pto @@ -1,8 +1,8 @@ // RUN: not ptoas %s 2>&1 | FileCheck %s module { - func.func @bad_syncall(%gm: memref<64xi32, #pto.address_space>) { - pto.syncall(%gm : memref<64xi32, #pto.address_space>) mode = #pto.sync_all_mode, core_type = #pto.sync_core_type + func.func @bad_syncall(%gm: !pto.partition_tensor_view<64xi32>) { + pto.syncall(%gm : !pto.partition_tensor_view<64xi32>) mode = #pto.sync_all_mode, core_type = #pto.sync_core_type return } } diff --git a/test/lit/pto/tassign_reject_memref.pto b/test/lit/pto/tassign_reject_memref.pto index d730dbe89f..9a13b29cf7 100644 --- a/test/lit/pto/tassign_reject_memref.pto +++ b/test/lit/pto/tassign_reject_memref.pto @@ -14,7 +14,8 @@ module { -> memref<16x16xf16, #pto.address_space> { %addr = arith.constant 0 : i64 %tile = memref.alloc() : memref<16x16xf16, #pto.address_space> - // CHECK: error: 'pto.tassign' op expects tile operand and result to be !pto.tile_buf + // CHECK: error: 'pto.tassign' op operand #0 must be + // CHECK-SAME: but got 'memref<16x16xf16, #pto.address_space>' %rebound = pto.tassign %tile, %addr : memref<16x16xf16, #pto.address_space> -> memref<16x16xf16, #pto.address_space> diff --git a/test/lit/pto/tprefetch_async_emitc.pto b/test/lit/pto/tprefetch_async_emitc.pto index 020aa96bf5..5536152da0 100644 --- a/test/lit/pto/tprefetch_async_emitc.pto +++ b/test/lit/pto/tprefetch_async_emitc.pto @@ -1,10 +1,10 @@ // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s module { - func.func @tprefetch_async_emitc(%src: memref<128xf32, #pto.address_space>, + func.func @tprefetch_async_emitc(%src: !pto.partition_tensor_view<128xf32>, %workspace: !pto.ptr) { %ctx = pto.make_prefetch_async_context(%workspace : !pto.ptr) -> !pto.prefetch_async_context - %event = pto.tprefetch_async(%src, %ctx : memref<128xf32, #pto.address_space>, !pto.prefetch_async_context) -> !pto.async_event + %event = pto.tprefetch_async(%src, %ctx : !pto.partition_tensor_view<128xf32>, !pto.prefetch_async_context) -> !pto.async_event %session = pto.get_prefetch_async_session %ctx : !pto.prefetch_async_context -> !pto.async_session %done = pto.comm.wait_async_event(%event, %session : !pto.async_event, !pto.async_session) -> i1 return diff --git a/test/lit/pto/tprefetch_async_invalid_non_1d.pto b/test/lit/pto/tprefetch_async_invalid_non_1d.pto index 27e17f290b..5ac9052d00 100644 --- a/test/lit/pto/tprefetch_async_invalid_non_1d.pto +++ b/test/lit/pto/tprefetch_async_invalid_non_1d.pto @@ -1,12 +1,12 @@ // RUN: not ptoas %s 2>&1 | FileCheck %s module { - func.func @bad_tprefetch_async(%src: memref<4x32xf32, #pto.address_space>, + func.func @bad_tprefetch_async(%src: !pto.partition_tensor_view<4x32xf32>, %workspace: !pto.ptr) { %ctx = pto.make_prefetch_async_context(%workspace : !pto.ptr) -> !pto.prefetch_async_context - %event = pto.tprefetch_async(%src, %ctx : memref<4x32xf32, #pto.address_space>, !pto.prefetch_async_context) -> !pto.async_event + %event = pto.tprefetch_async(%src, %ctx : !pto.partition_tensor_view<4x32xf32>, !pto.prefetch_async_context) -> !pto.async_event return } } -// CHECK: error: 'pto.tprefetch_async' op expects src to be a static flat contiguous logical 1D GM memref +// CHECK: error: 'pto.tprefetch_async' op expects src to be a static flat contiguous logical 1D GM view diff --git a/test/lit/pto/tprefetch_tile_native.pto b/test/lit/pto/tprefetch_tile_native.pto index 0ea39b2449..d503940f5e 100644 --- a/test/lit/pto/tprefetch_tile_native.pto +++ b/test/lit/pto/tprefetch_tile_native.pto @@ -10,15 +10,15 @@ // RUN: ptoas --pto-level=level3 --pto-arch=a3 %s | FileCheck %s --check-prefix=EMITC module { - func.func private @tprefetch_arg(%src: memref<16x16xf16, #pto.address_space>, %dst: !pto.tile_buf) { - pto.tprefetch ins(%src : memref<16x16xf16, #pto.address_space>) + func.func private @tprefetch_arg(%src: !pto.partition_tensor_view<16x16xf16>, %dst: !pto.tile_buf) { + pto.tprefetch ins(%src : !pto.partition_tensor_view<16x16xf16>) outs(%dst : !pto.tile_buf) return } } // NATIVE-LABEL: @tprefetch_arg -// NATIVE-SAME: %{{.*}}: memref<16x16xf16, #pto.address_space> +// NATIVE-SAME: %{{.*}}: !pto.partition_tensor_view<16x16xf16> // NATIVE-SAME: %{{.*}}: !pto.tile_buf>, - %src1: memref<32x32xf32, #pto.address_space>) { - pto.tsync ins(%src0, %src1 : memref<32x32xf32, #pto.address_space>, memref<32x32xf32, #pto.address_space>) + func.func @tsync_emitc(%src0: !pto.partition_tensor_view<32x32xf32>, + %src1: !pto.partition_tensor_view<32x32xf32>) { + pto.tsync ins(%src0, %src1 : !pto.partition_tensor_view<32x32xf32>, !pto.partition_tensor_view<32x32xf32>) return } From 6821e4c5b59abefd6243c913c9cb78e2270ec54f Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 21 Jul 2026 17:36:46 +0800 Subject: [PATCH 89/91] Remove pto.pointer_cast op Drop the pto.pointer_cast ODS definition, verifier, EmitC lowering, InsertSync/GSS metadata handling, and related VPTO compatibility hooks. Keep tile_buf_addr as the VPTO bridge until vpto-ptr-normalize, and update focused lit coverage for the new tile-native path. --- docs/PTO_IR_manual.md | 2 +- docs/designs/ptoas-auto-sync-design.md | 2 +- ...est-first-fit-four-gates-memplan-design.md | 2 +- include/PTO/IR/PTOOps.td | 47 +-- include/PTO/IR/PTOTypeDefs.td | 2 +- .../PTO/Transforms/GraphSyncSolver/MemInfo.h | 2 - .../Transforms/InsertSync/PTOIRTranslator.h | 1 - lib/PTO/IR/PTO.cpp | 18 +- lib/PTO/Transforms/FoldTileBufIntrinsics.cpp | 37 +- .../Transforms/GraphSyncSolver/MemInfo.cpp | 20 - .../SyncSolverIRTranslator.cpp | 2 +- .../InsertSync/InsertSyncAnalysis.cpp | 2 +- .../InsertSync/MemoryDependentAnalyzer.cpp | 18 +- .../Transforms/InsertSync/PTOIRTranslator.cpp | 82 +--- .../Transforms/InsertSync/SyncMacroModel.cpp | 2 - lib/PTO/Transforms/PTORemoveIdentityTMov.cpp | 5 - lib/PTO/Transforms/PTOToEmitC.cpp | 392 +----------------- lib/PTO/Transforms/PTOVPTOPtrBoundary.cpp | 4 +- .../Transforms/TileFusion/FusionAnalysis.cpp | 6 - .../TileFusion/PTOLowLevelLoopFusion.cpp | 5 +- .../Transforms/VPTOBufferMaterialization.cpp | 7 - lib/PTO/Transforms/VPTOPtrNormalize.cpp | 39 +- ...alloc_tile_plan_memory_no_memref_alloc.pto | 1 - .../pto/issue870_identity_tmov_pruning.pto | 4 +- .../lit/pto/multi_tile_buf_n3_planmem_e2e.pto | 1 - .../multi_tile_const_slot_disjoint_sync.pto | 2 +- .../multi_tile_get_const_slot_lowering.pto | 1 - test/lit/pto/multi_tile_n4_planmem_e2e.pto | 1 - ...lan_memory_five_gates_lifetime_overlap.pto | 1 - .../pto/plan_memory_five_gates_phi_family.pto | 1 - .../pto/plan_memory_inplace_forbid_alias.pto | 1 - ...an_memory_spec_level0_no_reuse_overlap.pto | 1 - .../lit/pto/plan_memory_spec_level0_reuse.pto | 1 - ...pointer_cast_addr_alignment_gm_ignored.pto | 9 - .../pointer_cast_addr_alignment_invalid.pto | 11 - .../pto/retired_memref_bridge_ops_invalid.pto | 11 + .../op_fusion_region_pipeline_level3.pto | 1 - test/lit/vpto/fold_tile_buf_intrinsics.pto | 29 +- .../issue_791_mte_ub_gm_ptr_normalize.pto | 13 +- tools/ptoas/ptoas.cpp | 4 +- 40 files changed, 74 insertions(+), 716 deletions(-) delete mode 100644 test/lit/pto/pointer_cast_addr_alignment_gm_ignored.pto delete mode 100644 test/lit/pto/pointer_cast_addr_alignment_invalid.pto diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index 17347cd946..c42abc0a49 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -10319,7 +10319,7 @@ pto.comm.treduce(%dst, %acc, recv(%ping, %pong), group(%g0, %g1, %g2) : Minimum support for **C++ stack-local statically-shaped arrays of scalars** — suitable for small auxiliary buffers in host scalar code. Disjoint from the tile-buf world: these values do not participate in PTO memory planning or -`pto.pointer_cast`, and their underlying address is decided by the host C++ +tile address materialization, and their underlying address is decided by the host C++ compiler. Naming and asm style mirror the `eventid_array` triad. Operates on the [`!pto.local_array<...>`](#26-ptolocal_arrayd1-x-d2-x--x-dk-x-t) type. See Section 2.6 for type-level constraints. diff --git a/docs/designs/ptoas-auto-sync-design.md b/docs/designs/ptoas-auto-sync-design.md index b8f5c0a2d5..dd9e4fc126 100644 --- a/docs/designs/ptoas-auto-sync-design.md +++ b/docs/designs/ptoas-auto-sync-design.md @@ -199,7 +199,7 @@ MLIR op，最终生成 `pto::SetFlagOp` / `pto::WaitFlagOp` / `pto::BarrierOp` - 主要逻辑： 1. `UpdateKernelArgMemInfo()`：把 kernel 参数登记为 GM root buffer。 2. `RecursionIR(&func.getBody())`：前序遍历 region： - - `pto::AllocTileOp` / `DeclareTileMemRefOp` / `PointerCastOp` / + - `pto::AllocTileOp` / `DeclareTileMemRefOp` / `memref::AllocOp` 经 `Update*MemInfo` 写入 `buffer2MemInfoMap_`。 - View / Subview / Cast / Mov 调 `UpdateAliasBufferInfo(result, source)`， 把派生 buffer 的 `BaseMemInfo` 链回到原 root。 diff --git a/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md b/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md index 3bd87e6b03..cd9fc80d85 100644 --- a/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md +++ b/docs/designs/ptoas-largest-first-fit-four-gates-memplan-design.md @@ -20,7 +20,7 @@ pto.alloc_tile(no addr) -> 后续 pto.t* tile op 继续使用 !pto.tile_buf 形态 ``` -用户显式写 `pto.alloc_tile addr` 时，level 语义仍由 memplan 校验：level1/level2 禁止显式 addr，level3 要求显式 local addr。也就是说，`pto.alloc_tile` 不再通过 `memref.alloc -> pto.pointer_cast -> pto.bind_tile` 这条中间链路表达地址；只有非 `alloc_tile` 的 memref/address root 仍可在 materialize 阶段使用 `pto.pointer_cast + pto.bind_tile` 表达规划后的地址与 tile metadata。 +用户显式写 `pto.alloc_tile addr` 时，level 语义仍由 memplan 校验：level1/level2 禁止显式 addr，level3 要求显式 local addr。也就是说，`pto.alloc_tile` 不再通过 `memref.alloc -> pto.bind_tile` 这类中间链路表达地址；普通 local tile allocation 由 memplan 直接把常量地址回写到 `pto.alloc_tile addr`。历史 `pto.pointer_cast` bridge 已删除，不再作为 memplan materialize 结果。 ## 目标 diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 6862877560..6f08bea192 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -1437,51 +1437,6 @@ def TMovOp : PTO_TOp<"tmov", [ } -//===----------------------------------------------------------------------===// -// Pointer Cast Op (existing) -//===----------------------------------------------------------------------===// - -def PointerCastOp : PTO_Op<"pointer_cast", [AttrSizedOperandSegments, Pure]> { - let summary = "Casts an integer address to a MemRef with optional valid dims"; - - // 参数定义 (保持 Optional) - let arguments = (ins - Variadic:$addrs, - Optional:$valid_row, - Optional:$valid_col, - OptionalAttr:$config - ); - - let results = (outs Res:$result); - - // Assembly Format (去掉了 []) - let assemblyFormat = [{ - `(` $addrs `)` ($valid_row^ `,` $valid_col)? attr-dict `:` qualified(type($result)) - }]; - - // Builder (修复 addOperand -> addOperands) - let builders = [ - OpBuilder<(ins "Type":$result, "ValueRange":$addrs, "Value":$vRow, "Value":$vCol, "Attribute":$config), [{ - $_state.addTypes(result); - $_state.addOperands(addrs); - // [关键修复] addOperand -> addOperands - if (vRow) $_state.addOperands(vRow); - if (vCol) $_state.addOperands(vCol); - if (config) $_state.addAttribute("config", config); - - int32_t addrsSize = addrs.size(); - int32_t vRowSize = vRow ? 1 : 0; - int32_t vColSize = vCol ? 1 : 0; - $_state.addAttribute("operandSegmentSizes", - $_builder.getDenseI32ArrayAttr({addrsSize, vRowSize, vColSize})); - }]> - ]; - - let extraClassDeclaration = [{ - ::mlir::LogicalResult verify(); - }]; -} - // ============================================================================= // System/Runtime Query Ops // ============================================================================= @@ -2468,7 +2423,7 @@ def EventIdArraySetOp : PTO_Op<"eventid_array_set"> { // Mirrors the `eventid_array` triad (declare / get / set) for a generic // statically-shaped scalar stack array. Disjoint from the tile-buf / scratch // memory world: the array's address is decided by the host C++ compiler at -// codegen time, never by PTO memory planning or `pto.pointer_cast`. +// codegen time, never by PTO memory planning. def DeclareLocalArrayOp : PTO_Op<"declare_local_array"> { let summary = "Declare a stack-local statically-shaped array"; diff --git a/include/PTO/IR/PTOTypeDefs.td b/include/PTO/IR/PTOTypeDefs.td index 69003cf5a4..b98f8e7363 100644 --- a/include/PTO/IR/PTOTypeDefs.td +++ b/include/PTO/IR/PTOTypeDefs.td @@ -280,7 +280,7 @@ def EventIdArrayType : TypeDef { // Statically-shaped C++ stack-local array. Lowers to a plain `T a[D1][D2]...;` // declaration in the emitted C++. Element type is restricted to scalars // (integer / float). This type does NOT participate in PTO memory planning -// or pto.pointer_cast lowering — the underlying address is decided by the +// lowering; the underlying address is decided by the // host C++ compiler, not by PTOAS. def LocalArrayType : TypeDef { let mnemonic = "local_array"; diff --git a/include/PTO/Transforms/GraphSyncSolver/MemInfo.h b/include/PTO/Transforms/GraphSyncSolver/MemInfo.h index 7e74779dae..6cc42ff223 100644 --- a/include/PTO/Transforms/GraphSyncSolver/MemInfo.h +++ b/include/PTO/Transforms/GraphSyncSolver/MemInfo.h @@ -70,8 +70,6 @@ struct MemInfo { llvm::SmallVector getAddresses(const llvm::SmallVector &addrs); -PointerLikeInfo getPointerLikeInfo(pto::PointerCastOp pointerCastOp); - MemInfo getMemInfo(Value val); MemInfo getMemInfo(const llvm::SmallVector &addrs); diff --git a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h index f53af0ce67..8fd6edd5df 100644 --- a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h +++ b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h @@ -73,7 +73,6 @@ class PTOIRTranslator { LogicalResult UpdateAllocMultiTileOpMemInfo(pto::AllocMultiTileOp op); LogicalResult UpdateDeclareTileOpMemInfo(pto::DeclareTileOp op); LogicalResult UpdateDeclareGlobalOpMemInfo(pto::DeclareGlobalOp op); - LogicalResult UpdatePointerCastOpMemInfo(pto::PointerCastOp op); LogicalResult UpdateMemrefAllocOpMemInfo(memref::AllocOp op); // 处理 View/Alias (MakeTensorView, Subview, Mov) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 53b7e20c40..98dd5920d6 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3062,21 +3062,6 @@ LogicalResult MultiTileGetOp::verify() { return success(); } -LogicalResult PointerCastOp::verify() { - auto memRefTy = dyn_cast(getResult().getType()); - if (!memRefTy) - return emitOpError("result must be a memref type"); - - for (auto [idx, addr] : llvm::enumerate(getAddrs())) { - if (failed(verifyConstantLocalAddress(getOperation(), addr, - memRefTy.getMemorySpace(), - static_cast(idx)))) - return failure(); - } - - return success(); -} - LogicalResult TAssignOp::verify() { if (getTile().getType() != getResult().getType()) { return emitOpError("result type must match tile operand type"); @@ -10878,8 +10863,7 @@ static bool isLocallyBoundTileSource(Value value) { if (!value || isa(value)) return false; - if (isa( - value.getDefiningOp())) + if (isa(value.getDefiningOp())) return true; if (auto bitcast = value.getDefiningOp()) diff --git a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp index 258b9b1fe9..bbe2e11b2f 100644 --- a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp +++ b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp @@ -265,11 +265,6 @@ static std::optional resolveTileHandle(Value tileBuf, return std::nullopt; } -static MemRefType getCanonicalMemRefTypeForTileBuf(pto::TileBufType tileTy) { - return MemRefType::get(tileTy.getShape(), tileTy.getElementType(), - AffineMap(), tileTy.getMemorySpace()); -} - struct ViewChain { UnrealizedConversionCastOp cast; memref::SubViewOp subview; @@ -735,25 +730,8 @@ struct FoldTileBufIntrinsicsPass return signalPassFailure(); } - if (auto resultMemrefType = - dyn_cast(addrOp.getDst().getType())) { - if (!handleInfo->addr) { - addrOp.emitError("FoldTileBufIntrinsics: pto.alloc_tile used by " - "tile_buf_addr must carry an addr operand on the " - "VPTO path"); - return signalPassFailure(); - } - - builder.setInsertionPoint(addrOp); - Value replacement = builder.create( - addrOp.getLoc(), resultMemrefType, ValueRange{handleInfo->addr}, - handleInfo->validRow ? handleInfo->validRow : Value(), - handleInfo->validCol ? handleInfo->validCol : Value(), - handleInfo->config); - addrOp.getDst().replaceAllUsesWith(replacement); - addrOp.erase(); + if (isa(addrOp.getDst().getType())) continue; - } auto resultPtrType = dyn_cast(addrOp.getDst().getType()); if (!resultPtrType) { @@ -770,17 +748,8 @@ struct FoldTileBufIntrinsicsPass } builder.setInsertionPoint(addrOp); - auto canonicalMemrefType = getCanonicalMemRefTypeForTileBuf(tileTy); - Value memrefValue = builder.create( - addrOp.getLoc(), canonicalMemrefType, ValueRange{handleInfo->addr}, - handleInfo->validRow ? handleInfo->validRow : Value(), - handleInfo->validCol ? handleInfo->validCol : Value(), - handleInfo->config); - - builder.setInsertionPoint(addrOp); - Value replacement = - builder.create(addrOp.getLoc(), resultPtrType, - memrefValue); + Value replacement = builder.create( + addrOp.getLoc(), resultPtrType, handleInfo->addr); addrOp.getDst().replaceAllUsesWith(replacement); addrOp.erase(); } diff --git a/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp b/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp index c4d9765e8c..a73f4781b0 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/MemInfo.cpp @@ -69,23 +69,6 @@ llvm::SmallVector getAddresses(const llvm::SmallVector &addrs) { return offsets; } -PointerLikeInfo getPointerLikeInfo(pto::PointerCastOp pointerCastOp) { - PointerLikeInfo pointerLikeInfo(pointerCastOp); - pointerLikeInfo.addresses = getAddresses(pointerCastOp.getAddrs()); - pointerLikeInfo.allocateSize = getBufferBitSize(pointerCastOp.getResult()); - if (!pointerLikeInfo.allocateSize.has_value()) { - pointerCastOp.emitError("unknown buffer size"); - llvm_unreachable("unknown buffer size"); - } - if (auto spaceAttr = GetBufferSpaceAttr(pointerCastOp.getResult())) { - pointerLikeInfo.addressSpace = spaceAttr->getAddressSpace(); - } - if (auto parentLoop = pointerCastOp->getParentOfType()) { - pointerLikeInfo.parentLoop = parentLoop; - } - return pointerLikeInfo; -} - static std::optional getConstantI64(Value value) { IntegerAttr attr; if (!value || !matchPattern(value, m_Constant(&attr))) @@ -137,9 +120,6 @@ static MemInfo getMemInfoForMultiTileGet(pto::MultiTileGetOp get) { MemInfo getMemInfo(Value val) { if (auto *defOp = val.getDefiningOp()) { - if (auto pointerCastOp = llvm::dyn_cast(defOp)) { - return MemInfo(val, getPointerLikeInfo(pointerCastOp)); - } if (auto allocMulti = llvm::dyn_cast(defOp)) return MemInfo(val, getPointerLikeInfo(allocMulti)); if (auto multiGet = llvm::dyn_cast(defOp)) diff --git a/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp b/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp index 97cb286297..94bdad6dbf 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp @@ -125,7 +125,7 @@ llvm::SmallVector IRTranslator::tracebackMemVals(Value val) { Operation *defOp = result.getDefiningOp(); // `pto.multi_tile_get` stops traversal so getMemInfo can extract the // slot-narrowed physical addresses. - if (isa(defOp)) { leaves.insert(result); continue; diff --git a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp index eac1305f22..dc34ffc2c1 100644 --- a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp +++ b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp @@ -767,7 +767,7 @@ void InsertSyncAnalysis::InsertLastPipeAll() { // ============================================================================== bool InsertSyncAnalysis::IsMemAllocOp(Operation *op) const { - return isa(op) || isa(op); + return isa(op); } SmallVector InsertSyncAnalysis::GetMemInfoBuffers( diff --git a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp index 2088f159db..247e5bd96c 100644 --- a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp +++ b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp @@ -106,14 +106,9 @@ static Value GetRealRoot(Value v) { return v; } -// Extract the base UB address carried by `rootBuffer` when it is the i64 -// address SSA operand of a `pto.pointer_cast` (single-address path). In that -// model `rootBuffer` is an `arith.constant` i64 whose value is the absolute UB -// offset, and `BaseMemInfo::baseAddresses` holds {0} (relative deltas). For -// the multi-address path `rootBuffer` is the `PointerCastOp` result itself -// (not an i64 constant), so this returns 0 and `baseAddresses` already holds -// the absolute slot offsets — making `rootBase + baseAddresses[i]` resolve to -// the absolute address in both models. +// Extract an absolute base address when a root is represented by an integer +// constant. Tile roots otherwise carry their planned addresses in +// `BaseMemInfo::baseAddresses`. static uint64_t getRootBaseAddress(Value rootBuffer) { if (!rootBuffer) return 0; @@ -145,18 +140,13 @@ static bool hasKnownLocalAbsoluteAddress(const BaseMemInfo *info) { if (isIntegerConstant(root)) return true; - // Multi-address pointer_cast records absolute slot offsets in baseAddresses - // and uses the pointer_cast result as rootBuffer. - if (auto ptrCast = root.getDefiningOp()) - return ptrCast.getAddrs().size() > 1; - return false; } // Cross-root byte-range overlap check for local memory (UB/L1). // // `MemAlias` historically only checked byte-range overlap when the two -// `rootBuffer` SSA values were identical (same `alloc_tile`/`pointer_cast`). +// `rootBuffer` SSA values were identical (same `alloc_tile`). // When PlanMemory reuses the same physical UB region for two *different* // allocations, their `rootBuffer` values differ even though their byte ranges // overlap — so `MemAlias` returned false and InsertSync silently dropped the diff --git a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp index eccdd1303a..05f6d3e412 100644 --- a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp +++ b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp @@ -384,10 +384,6 @@ void PTOIRTranslator::RecursionIR(Region *region) { return WalkResult::interrupt(); } } - else if (auto castOp = dyn_cast(op)) { - if (failed(UpdatePointerCastOpMemInfo(castOp))) return WalkResult::interrupt(); - } - // --- Case B: 别名/视图操作 --- else if (auto makeViewOp = dyn_cast(op)) { UpdateAliasBufferInfo(makeViewOp.getResult(), makeViewOp.getPtr()); @@ -574,81 +570,6 @@ PTOIRTranslator::UpdateAllocMultiTileOpMemInfo(pto::AllocMultiTileOp op) { return success(); } -LogicalResult PTOIRTranslator::UpdatePointerCastOpMemInfo(pto::PointerCastOp op) { - Value res = op.getResult(); - auto memRefType = dyn_cast(res.getType()); - if (!memRefType) return failure(); - - if (op.getAddrs().empty()) { - return op.emitError("PointerCast must have at least one address operand"); - } - - uint64_t sizeInBytes = 0; - if (memRefType.hasStaticShape()) { - int64_t elemSize = static_cast( - pto::getPTOStorageElemByteSize(memRefType.getElementType())); - if (elemSize == 0) - return failure(); - int64_t numElements = 1; - for (auto dim : memRefType.getShape()) numElements *= dim; - sizeInBytes = numElements * elemSize; - } - - pto::AddressSpace space = pto::AddressSpace::GM; - if (auto attr = memRefType.getMemorySpace()) { - if (auto ptoAttr = dyn_cast(attr)) { - space = ptoAttr.getAddressSpace(); - } - } - - if (op.getAddrs().size() == 1) { - // Keep the address SSA as the root so dynamic pointer casts retain their - // historical identity-based behavior. For a static local cast, also keep - // its byte address for cross-root physical range analysis; downstream view - // ops accumulate their delta into baseAddresses[0]. - Value rootSrc = op.getAddrs().front(); - SmallVector baseAddresses{0}; - bool hasKnownPhysicalAddresses = false; - if (isLocalAddressSpace(space)) { - if (std::optional address = getKnownPhysicalAddress(rootSrc)) { - baseAddresses[0] = *address; - hasKnownPhysicalAddresses = true; - } - } - auto newMemInfo = std::make_unique( - res, rootSrc, space, std::move(baseAddresses), sizeInBytes, - hasKnownPhysicalAddresses); - buffer2MemInfoMap_[res].emplace_back(newMemInfo->clone()); - return success(); - } - - // Multi-address cast. Use the cast result as `rootBuffer` and retain every - // constant physical offset for conservative range-overlap analysis. - SmallVector slotOffsets; - slotOffsets.reserve(op.getAddrs().size()); - for (Value a : op.getAddrs()) { - std::optional address = getKnownPhysicalAddress(a); - if (!address) { - // Non-constant slot address: fall back to single-address semantics - // with the first operand as rootBuffer so existing non-multi-buffer - // codepaths that happen to feed non-constant i64s keep their - // historical behavior. - Value rootSrc = op.getAddrs().front(); - auto newMemInfo = std::make_unique( - res, rootSrc, space, SmallVector{0}, sizeInBytes); - buffer2MemInfoMap_[res].emplace_back(newMemInfo->clone()); - return success(); - } - slotOffsets.push_back(*address); - } - - auto newMemInfo = std::make_unique( - res, res, space, std::move(slotOffsets), sizeInBytes, - isLocalAddressSpace(space)); - buffer2MemInfoMap_[res].emplace_back(newMemInfo->clone()); - return success(); -} - LogicalResult PTOIRTranslator::UpdateDeclareTileOpMemInfo(pto::DeclareTileOp op) { Value result = op.getTile(); @@ -1049,8 +970,7 @@ void PTOIRTranslator::UpdateSlotSelectedAliasBufferInfo(Value result, for (auto &parentInfo : buffer2MemInfoMap_[source]) { auto newInfo = parentInfo->clone(result); // Multi-buffer parent: `baseAddresses` lists every physical slot's - // offset (populated by `UpdatePointerCastOpMemInfo` for multi-address - // casts). For a constant slot index, narrow `baseAddresses` to just + // planner-assigned offset. For a constant slot index, narrow it to just // that one slot so MemAlias's range-overlap check returns false when // two const-slot uses pick different slots. For a dynamic slot index, // keep all addresses -- the runtime SSA could resolve to any slot, so diff --git a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp index 9db12e038c..89a911bcc6 100644 --- a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp +++ b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp @@ -239,8 +239,6 @@ static std::optional getConstantIndex(Value value) { } static std::optional> lookupMGatherValidDims(Value value) { - if (auto pc = value.getDefiningOp()) - return SmallVector{pc.getValidRow(), pc.getValidCol()}; if (auto subview = value.getDefiningOp()) return lookupMGatherValidDims(subview.getSource()); if (auto cast = value.getDefiningOp()) diff --git a/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp b/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp index b2b42b75df..330b52820b 100644 --- a/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp +++ b/lib/PTO/Transforms/PTORemoveIdentityTMov.cpp @@ -145,11 +145,6 @@ tryGetConcreteRootAddress(const BaseMemInfo *info) { if (auto alloc = dyn_cast(defOp)) return tryEvalI64Constant(alloc.getAddr()); - if (auto pointerCast = dyn_cast(defOp)) { - if (pointerCast.getAddrs().size() == 1) - return tryEvalI64Constant(pointerCast.getAddrs().front()); - } - return std::nullopt; } diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 6789416acd..86608934f2 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -196,8 +196,6 @@ static Location getIndexedNameHintLoc(Location fallbackLoc, unsigned index) { fallbackLoc); } -static constexpr llvm::StringLiteral kForceDynamicValidShapeAttrName = - "__pto.force_dynamic_valid_shape"; static constexpr llvm::StringLiteral kGlobalTensorStridesAttrName = "__pto.globaltensor_strides"; static constexpr llvm::StringLiteral kPipePeerOwnerFuncAttrName = @@ -513,27 +511,19 @@ getSpecialScaleGlobalTensorTypeSpecForTileValue(Value dstValue, Type elemTy) { dstValue = peelUnrealized(dstValue); - auto dstMrTy = dyn_cast(dstValue.getType()); - if (!dstMrTy) + auto dstTileTy = dyn_cast(dstValue.getType()); + if (!dstTileTy) return std::nullopt; - auto dstSpace = dyn_cast_or_null(dstMrTy.getMemorySpace()); + auto dstSpace = dyn_cast_or_null( + dstTileTy.getMemorySpace()); if (!dstSpace || dstSpace.getAddressSpace() != pto::AddressSpace::MAT) return std::nullopt; - std::optional configAttr; - if (auto cast = dstValue.getDefiningOp()) { - if (auto config = cast.getConfig()) - configAttr = config; - } - if (!configAttr) - return std::nullopt; - - ArrayRef effectiveShape = shape; - if (dstMrTy.getRank() == 2 && dstMrTy.hasStaticShape()) - effectiveShape = dstMrTy.getShape(); - - auto config = *configAttr; + ArrayRef effectiveShape = dstTileTy.getShape(); + if (effectiveShape.empty()) + effectiveShape = shape; + auto config = dstTileTy.getConfigAttr(); if (!isF8E8M0ElemType(elemTy)) return std::nullopt; if (effectiveShape.size() != 2) @@ -4515,11 +4505,6 @@ static FailureOr buildAsyncScratchTileValue( auto *ctx = rewriter.getContext(); pto::TileBufConfigAttr configAttr = pto::TileBufConfigAttr::getDefault(ctx); - if (auto cast = originalScratch.getDefiningOp()) { - if (auto config = cast.getConfig()) - configAttr = *config; - } - int32_t fractal = 512; if (auto frAttr = dyn_cast(configAttr.getSFractalSize())) fractal = static_cast(getIntegerAttrSignedValue(frAttr)); @@ -4584,11 +4569,6 @@ static FailureOr buildSyncAllWorkspaceTileValue( auto *ctx = rewriter.getContext(); pto::TileBufConfigAttr configAttr = pto::TileBufConfigAttr::getDefault(ctx); - if (auto cast = originalWorkspace.getDefiningOp()) { - if (auto config = cast.getConfig()) - configAttr = *config; - } - Attribute memorySpace = memTy.getMemorySpace(); if (!memorySpace) return failure(); @@ -4628,7 +4608,7 @@ static FailureOr buildSyncAllWorkspaceTileValue( } //===----------------------------------------------------------------------===// -// pto.pointer_cast lowering +// PTO pointer lowering //===----------------------------------------------------------------------=== struct CastPtrConversion : public OpConversionPattern { @@ -4743,355 +4723,6 @@ struct PTOAddPtrToEmitC : public OpConversionPattern { } }; -struct PointerCastConversion : public OpConversionPattern { - static bool getIndexConst(Value v, int64_t &out) { - if (auto cst = v.getDefiningOp()) { - if (auto ia = dyn_cast(cst.getValue())) { - out = getIntegerAttrSignedValue(ia); - return true; - } - } - return false; - } - - using OpConversionPattern::OpConversionPattern; - - enum class TileRole { Vec, Mat, Left, Right, Acc, Bias, Scaling }; - - static void collectUserOpsThroughCasts(Value v, SmallVectorImpl &out) { - for (Operation *u : v.getUsers()) { - if (auto castOp = dyn_cast(u)) { - for (Value r : castOp.getResults()) - collectUserOpsThroughCasts(r, out); - continue; - } - out.push_back(u); - } - } - - static Value peelUnrealized(Value v) { - while (auto castOp = v.getDefiningOp()) { - v = castOp.getOperand(0); - } - return v; - } - - static TileRole inferRole(pto::PointerCastOp op) { - // 1. 优先检查 AddressSpace - if (auto memRefTy = dyn_cast(op.getType())) { - Attribute memorySpace = memRefTy.getMemorySpace(); - if (auto ptoAttr = dyn_cast_or_null(memorySpace)) { - switch (ptoAttr.getAddressSpace()) { - case pto::AddressSpace::LEFT: return TileRole::Left; - case pto::AddressSpace::RIGHT: return TileRole::Right; - case pto::AddressSpace::ACC: return TileRole::Acc; - case pto::AddressSpace::BIAS: return TileRole::Bias; - case pto::AddressSpace::MAT: return TileRole::Mat; - case pto::AddressSpace::SCALING: return TileRole::Scaling; - default: break; - } - } - } - - // 2. 通过 Usage 推导 (Fallback) - SmallVector users; - collectUserOpsThroughCasts(op.getResult(), users); - - for (Operation *user : users) { - if (auto mm = dyn_cast(user)) { - if (mm.getDst() && peelUnrealized(mm.getDst()) == op.getResult()) return TileRole::Acc; - if (peelUnrealized(mm.getLhs()) == op.getResult()) return TileRole::Left; - if (peelUnrealized(mm.getRhs()) == op.getResult()) return TileRole::Right; - } - if (auto mmacc = dyn_cast(user)) { - if (mmacc.getDst() && peelUnrealized(mmacc.getDst()) == op.getResult()) return TileRole::Acc; - if (peelUnrealized(mmacc.getAccIn()) == op.getResult()) return TileRole::Acc; - if (peelUnrealized(mmacc.getLhs()) == op.getResult()) return TileRole::Left; - if (peelUnrealized(mmacc.getRhs()) == op.getResult()) return TileRole::Right; - } - if (auto tgemv = dyn_cast(user)) { - if (peelUnrealized(tgemv.getAScale()) == op.getResult()) return TileRole::Left; - if (peelUnrealized(tgemv.getBScale()) == op.getResult()) return TileRole::Right; - } - if (auto tgemvAcc = dyn_cast(user)) { - if (peelUnrealized(tgemvAcc.getAScale()) == op.getResult()) return TileRole::Left; - if (peelUnrealized(tgemvAcc.getBScale()) == op.getResult()) return TileRole::Right; - } - if (auto tgemvBias = dyn_cast(user)) { - if (peelUnrealized(tgemvBias.getAScale()) == op.getResult()) return TileRole::Left; - if (peelUnrealized(tgemvBias.getBScale()) == op.getResult()) return TileRole::Right; - } - if (auto tmatmulMx = dyn_cast(user)) { - if (peelUnrealized(tmatmulMx.getAScale()) == op.getResult()) return TileRole::Left; - if (peelUnrealized(tmatmulMx.getBScale()) == op.getResult()) return TileRole::Right; - } - if (auto tmatmulMxAcc = dyn_cast(user)) { - if (peelUnrealized(tmatmulMxAcc.getAScale()) == op.getResult()) return TileRole::Left; - if (peelUnrealized(tmatmulMxAcc.getBScale()) == op.getResult()) return TileRole::Right; - } - if (auto tmatmulMxBias = dyn_cast(user)) { - if (peelUnrealized(tmatmulMxBias.getAScale()) == op.getResult()) return TileRole::Left; - if (peelUnrealized(tmatmulMxBias.getBScale()) == op.getResult()) return TileRole::Right; - } - } - - return TileRole::Vec; - } - - // [新增] 辅助函数：判断 Value 是否源自 arith.constant - static bool isConstant(Value v, int64_t &outVal) { - if (!v) return false; - if (auto cst = v.getDefiningOp()) { - if (auto attr = dyn_cast(cst.getValue())) { - outVal = getIntegerAttrSignedValue(attr); - return true; - } - } - return false; - } - - LogicalResult matchAndRewrite(pto::PointerCastOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op.getLoc(); - auto *ctx = rewriter.getContext(); - auto selfType = mlir::cast(op.getType()); - ArrayRef shape = selfType.getShape(); - Type elemType = selfType.getElementType(); - - // 1. 推导 Tile Role - TileRole role = inferRole(op); - - // 2. 类型字符串生成 (elemTypeStr, dimStr) - std::string elemTypeStr = getEmitCScalarTypeToken(elemType); - - std::string dimStr; - pto::BLayout blayout = pto::BLayout::RowMajor; - auto dimToString = [&](int64_t dim, const char *symbol, - int dimIdx) -> std::string { - if (dim == ShapedType::kDynamic) - return std::string(symbol); - return std::to_string(renderTileTemplateDim(dim, elemType, blayout, - dimIdx)); - }; - - // 3. Role Token - const char *roleTok = "TileType::Vec"; - switch (role) { - case TileRole::Left: roleTok = "TileType::Left"; break; - case TileRole::Right: roleTok = "TileType::Right"; break; - case TileRole::Acc: roleTok = "TileType::Acc"; break; - case TileRole::Bias: roleTok = "TileType::Bias"; break; - case TileRole::Mat: roleTok = "TileType::Mat"; break; - case TileRole::Vec: roleTok = "TileType::Vec"; break; - case TileRole::Scaling: - if (auto configOpt = op.getConfig()) - roleTok = scalingRoleToken(elemType, *configOpt); - else - roleTok = "TileType::Scaling"; - break; - } - - // 4. Config & Layout (support BLayoutAttr/SLayoutAttr/PadValueAttr after namespace change) - std::string layoutParams = "BLayout::RowMajor"; - std::string extraParams = ""; - if (auto configOpt = op.getConfig()) { - auto config = *configOpt; - int32_t blVal = 0; - if (auto attr = dyn_cast(config.getBLayout())) - blVal = static_cast(attr.getValue()); - - if (blVal == 1) layoutParams = "BLayout::ColMajor"; - blayout = blVal == 1 ? pto::BLayout::ColMajor : pto::BLayout::RowMajor; - - int32_t slVal = 0; - if (auto attr = dyn_cast(config.getSLayout())) - slVal = static_cast(attr.getValue()); - - std::string slStr = (slVal == 1) ? "SLayout::RowMajor" : (slVal == 2) ? "SLayout::ColMajor" : "SLayout::NoneBox"; - - int32_t frVal = 0; - if (auto attr = dyn_cast(config.getSFractalSize())) - frVal = static_cast(getIntegerAttrSignedValue(attr)); - - int32_t padVal = 0; - if (auto attr = dyn_cast(config.getPad())) - padVal = static_cast(attr.getValue()); - - std::string padStr = "PadValue::Null"; - switch (padVal) { - case 1: padStr = "PadValue::Zero"; break; - case 2: padStr = "PadValue::Max"; break; - case 3: padStr = "PadValue::Min"; break; - } - - int32_t compactVal = 0; - if (auto attr = dyn_cast(config.getCompactMode())) - compactVal = static_cast(attr.getValue()); - - std::string compactStr = "CompactMode::Null"; - switch (compactVal) { - case 1: compactStr = "CompactMode::Normal"; break; - case 2: compactStr = "CompactMode::RowPlusOne"; break; - } - - if (!slStr.empty()) { - extraParams += ", " + slStr + ", " + std::to_string(frVal) + ", " + - padStr + ", " + compactStr; - } - } else { - extraParams = ", SLayout::NoneBox, 512, PadValue::Null, CompactMode::Null"; - } - - if (role == TileRole::Left) - dimStr = dimToString(shape[0], "M", 0) + ", " + - dimToString(shape[1], "K", 1); - else if (role == TileRole::Right) - dimStr = dimToString(shape[0], "K", 0) + ", " + - dimToString(shape[1], "N", 1); - else if (role == TileRole::Bias) - dimStr = "1, " + dimToString(shape[1], "N", 1); - else - dimStr = dimToString(shape[0], "M", 0) + ", " + - dimToString(shape[1], "N", 1); - - // [核心修改] Valid Dims 处理逻辑 (支持混合静态/动态) - std::string vrowTok, vcolTok; - bool useConstructor = false; - - bool rowIsDynamic = false; - bool colIsDynamic = false; - - SmallVector constructorArgs; - - Value vRow = op.getValidRow(); - Value vCol = op.getValidCol(); - Value vRowEmitC = adaptor.getValidRow(); - Value vColEmitC = adaptor.getValidCol(); - bool forceDynamicValid = op->hasAttr(kForceDynamicValidShapeAttrName); - - int64_t cRow = 0, cCol = 0; - bool rowIsConst = vRow && isConstant(vRow, cRow); - bool colIsConst = vCol && isConstant(vCol, cCol); - - auto makeCtorDimValue = [&](Value emitted, int64_t fallback) -> Value { - if (emitted) - return emitted; - return makeEmitCIntConstant( - rewriter, loc, emitc::OpaqueType::get(ctx, "int32_t"), fallback); - }; - auto maybeScaleDynamicValid = [&](Value emitted, int dimIdx) -> Value { - if (!emitted || !pto::isPTOFloat4PackedType(elemType)) - return emitted; - int packedDim = blayout == pto::BLayout::ColMajor ? 0 : 1; - if (dimIdx != packedDim) - return emitted; - auto i32Ty = emitc::OpaqueType::get(ctx, "int32_t"); - Value two = makeEmitCIntConstant(rewriter, loc, i32Ty, 2); - return rewriter.create(loc, i32Ty, emitted, two).getResult(); - }; - - if (forceDynamicValid) { - vrowTok = "-1"; - vcolTok = "-1"; - useConstructor = true; - constructorArgs.push_back( - makeCtorDimValue(maybeScaleDynamicValid(vRowEmitC, 0), - renderTileTemplateDim(rowIsConst ? cRow : shape[0], - elemType, blayout, 0))); - constructorArgs.push_back( - makeCtorDimValue(maybeScaleDynamicValid(vColEmitC, 1), - renderTileTemplateDim(colIsConst ? cCol : shape[1], - elemType, blayout, 1))); - } else { - if (rowIsConst) { - vrowTok = std::to_string( - renderTileTemplateDim(cRow, elemType, blayout, 0)); - } else if (vRow) { - vrowTok = "-1"; - rowIsDynamic = true; - useConstructor = true; - } else { - vrowTok = std::to_string( - renderTileTemplateDim(shape[0], elemType, blayout, 0)); - } - - if (colIsConst) { - vcolTok = std::to_string( - renderTileTemplateDim(cCol, elemType, blayout, 1)); - } else if (vCol) { - vcolTok = "-1"; - colIsDynamic = true; - useConstructor = true; - } else { - vcolTok = std::to_string( - renderTileTemplateDim(shape[1], elemType, blayout, 1)); - } - - if (useConstructor) { - if (rowIsDynamic && vRowEmitC) - constructorArgs.push_back(maybeScaleDynamicValid(vRowEmitC, 0)); - if (colIsDynamic && vColEmitC) - constructorArgs.push_back(maybeScaleDynamicValid(vColEmitC, 1)); - } - } - - // 5. 生成 Tile 类型字符串 - std::string tileTypeStr = - std::string("Tile<") + roleTok + ", " + elemTypeStr + ", " + dimStr + ", " + - layoutParams + ", " + vrowTok + ", " + vcolTok + extraParams + ">"; - - auto tileType = emitc::OpaqueType::get(ctx, tileTypeStr); - Value resultValue; - - if (useConstructor) { - // 使用 CallOpaqueOp 生成构造函数调用 (Tile v = Tile(...)) - auto ctorOp = rewriter.create( - loc, - tileType, // Result Type - tileTypeStr, // Callee Name (类名) - ArrayAttr{}, // args - ArrayAttr{}, // template_args - ValueRange(constructorArgs) // operands - ); - resultValue = ctorOp.getResult(0); - } else { - // 静态情况 (Tile v;) - auto varOp = rewriter.create( - loc, - getEmitCVariableResultType(tileType), - emitc::OpaqueAttr::get(ctx, "") - ); - resultValue = loadEmitCVariableIfNeeded(rewriter, loc, varOp.getResult()); - } - - // TASSIGN: pto-isa expects an integral address. - Value addr = adaptor.getAddrs()[0]; - if (isa(addr.getType()) || - (isa(addr.getType()) && - cast(addr.getType()).getValue().ends_with("*"))) { - auto u64Ty = emitc::OpaqueType::get(ctx, "uint64_t"); - auto rcU64 = rewriter.getArrayAttr({emitc::OpaqueAttr::get(ctx, "uint64_t")}); - addr = rewriter.create( - loc, u64Ty, "reinterpret_cast", - /*args=*/ArrayAttr{}, /*templateArgs=*/rcU64, - /*operands=*/ValueRange{addr}) - .getResult(0); - } - - rewriter.create( - loc, TypeRange{}, "TASSIGN", - ArrayAttr{}, ArrayAttr{}, - ValueRange{resultValue, addr}); - - rewriter.replaceOp(op, resultValue); - return success(); - } -}; - -//===----------------------------------------------------------------------===// -// pto.load_dps / pto.store_dps lowering (FIX: keep optional result) -//===----------------------------------------------------------------------=== - struct PTOTLoadToTLOAD : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -8489,7 +8120,7 @@ struct PTOTPushToEmitC : public OpConversionPattern { if (failed(pipeTok)) return rewriter.notifyMatchFailure(op, "failed to resolve pipe token"); // Read the tile type token from the already-converted OpaqueType, which - // preserves the exact layout produced by PointerCastOp EmitC. + // preserves the exact tile layout produced during EmitC conversion. Value convertedTile = peelUnrealized(adaptor.getTile()); auto tileTok = getPipeDataTypeToken(convertedTile); if (failed(tileTok)) @@ -13577,8 +13208,7 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); - patterns.add( - typeConverter, ctx); + patterns.add(typeConverter, ctx); patterns.adduse_empty()) return; - if (isa(op)) + if (isa(op)) deadOps.push_back(op); }); diff --git a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp index 3d867822ab..44f50acb2e 100644 --- a/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp +++ b/lib/PTO/Transforms/TileFusion/FusionAnalysis.cpp @@ -73,12 +73,6 @@ static bool isShapeComputableOp(Operation *op) { opName == "arith.divui" || opName == "arith.index_cast"; } - // pto.pointer_cast is a pure, regionless op used in interstage setup; - // it is not directly a shape computation but may appear in valid-shape - // expression trees after lowering. - if (opName == "pto.pointer_cast") - return true; - return false; } diff --git a/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp b/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp index 2784ecf4a2..42cbe80493 100644 --- a/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp +++ b/lib/PTO/Transforms/TileFusion/PTOLowLevelLoopFusion.cpp @@ -81,10 +81,7 @@ static bool isInterstageSetupOp(Operation *op) { if (isPureNoRegionOp(op)) return true; - // Tile-native buffers lower to backend memrefs through pto.pointer_cast - // between adjacent stage loops. Treat these address materializations as - // stage-boundary-transparent so loop-run collection can keep walking. - return isa(op); + return false; } static bool areEquivalentOperations(Operation *lhs, Operation *rhs) { diff --git a/lib/PTO/Transforms/VPTOBufferMaterialization.cpp b/lib/PTO/Transforms/VPTOBufferMaterialization.cpp index cc0a965345..33abcc91e4 100644 --- a/lib/PTO/Transforms/VPTOBufferMaterialization.cpp +++ b/lib/PTO/Transforms/VPTOBufferMaterialization.cpp @@ -68,13 +68,6 @@ Value materializeBufferPointer(Value value, Type elementType, if (value.getType() == ptrType) return value; - if (auto cast = value.getDefiningOp()) { - if (cast.getAddrs().empty()) - return {}; - return rewriter.create(loc, ptrType, cast.getAddrs().front()) - .getResult(); - } - Value memrefValue = materializeTileBufferView(value, rewriter, loc); auto memrefType = dyn_cast_or_null(memrefValue.getType()); if (!memrefValue || !memrefType) diff --git a/lib/PTO/Transforms/VPTOPtrNormalize.cpp b/lib/PTO/Transforms/VPTOPtrNormalize.cpp index 7ea765a967..dd03c7c4c3 100644 --- a/lib/PTO/Transforms/VPTOPtrNormalize.cpp +++ b/lib/PTO/Transforms/VPTOPtrNormalize.cpp @@ -219,15 +219,6 @@ static Value materializeScalarAccessPtr(Value source, PatternRewriter &rewriter, return rewriter.create(loc, ptrType, basePtr, offset); } - if (auto pointerCast = source.getDefiningOp()) { - if (pointerCast.getAddrs().empty()) - return {}; - Value addr = pointerCast.getAddrs().front(); - if (isa(addr.getType())) - return addr; - return materializeScalarAccessPtr(addr, rewriter, loc); - } - // Restrict normalization to memref views that already sit on top of a ptr-like // boundary bridge. Materializing fresh memref -> ptr casts here would leave // illegal pto.castptr(memref) behind in this pass. @@ -298,28 +289,6 @@ struct ConvertTileBufAddrToPtrPattern } }; -struct ConvertPointerCastToCastPtrPattern - : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(pto::PointerCastOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - Type convertedType = - getTypeConverter()->convertType(op.getResult().getType()); - auto ptrType = dyn_cast(convertedType); - if (!ptrType) - return failure(); - - if (adaptor.getAddrs().empty()) - return rewriter.notifyMatchFailure(op, "expected at least one address"); - - rewriter.replaceOpWithNewOp(op, ptrType, - adaptor.getAddrs().front()); - return success(); - } -}; - struct ConvertIntToPtrToCastPtrPattern : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -825,7 +794,7 @@ struct VPTOPtrNormalizePass target.addLegalDialect(); target.addDynamicallyLegalDialect([](Operation *op) { - return !isa(op); }); @@ -851,11 +820,6 @@ struct VPTOPtrNormalizePass return op.getDst().getType() == typeConverter.convertType(op.getDst().getType()); }); - target.addDynamicallyLegalOp( - [&](pto::PointerCastOp op) { - return op.getResult().getType() == - typeConverter.convertType(op.getResult().getType()); - }); target.addDynamicallyLegalOp([&](pto::CastPtrOp op) { return !isMemRefType(op.getInput().getType()) && !isMemRefType(op.getResult().getType()); @@ -950,7 +914,6 @@ struct VPTOPtrNormalizePass populateCallOpTypeConversionPattern(patterns, typeConverter); populateReturnOpTypeConversionPattern(patterns, typeConverter); patterns.add} // CHECK: pto.multi_tile_get -// CHECK-NOT: pto.pointer_cast diff --git a/test/lit/pto/multi_tile_const_slot_disjoint_sync.pto b/test/lit/pto/multi_tile_const_slot_disjoint_sync.pto index 3c97d15503..66afd2a2d2 100644 --- a/test/lit/pto/multi_tile_const_slot_disjoint_sync.pto +++ b/test/lit/pto/multi_tile_const_slot_disjoint_sync.pto @@ -10,7 +10,7 @@ // Confirms that two constant-slot uses of a `multi_tile_buf` are recognised // as disjoint by InsertSync: the two MTE2 loads land on distinct physical -// slots (different `pto.pointer_cast` rootBuffer i64 constants), so no +// slots (different planner-assigned physical addresses), so no // inter-slot `set_flag` / `wait_flag` is needed between them. Only the // pipeline tail barrier remains. Without slot-aware buffer planning this // would otherwise emit redundant MTE2-MTE2 synchronization between the diff --git a/test/lit/pto/multi_tile_get_const_slot_lowering.pto b/test/lit/pto/multi_tile_get_const_slot_lowering.pto index 50b41802e2..cb6bd1f8b8 100644 --- a/test/lit/pto/multi_tile_get_const_slot_lowering.pto +++ b/test/lit/pto/multi_tile_get_const_slot_lowering.pto @@ -54,7 +54,6 @@ module { // PLAN-LABEL: func.func @const_slot_two_buffers // PLAN: pto.alloc_multi_tile {pto.multi_buffer_addrs = array} // PLAN: pto.multi_tile_get -// PLAN-NOT: pto.pointer_cast // Stage 3 -- ResolveBufferSelect materializes one addressed tile per use. // SELECT: IR Dump After PTOResolveBufferSelect diff --git a/test/lit/pto/multi_tile_n4_planmem_e2e.pto b/test/lit/pto/multi_tile_n4_planmem_e2e.pto index 5667c623e3..13244bc2a6 100644 --- a/test/lit/pto/multi_tile_n4_planmem_e2e.pto +++ b/test/lit/pto/multi_tile_n4_planmem_e2e.pto @@ -54,7 +54,6 @@ module { // PLAN: IR Dump After PlanMemory // PLAN-LABEL: func.func @n4_rotate // PLAN: pto.alloc_multi_tile {pto.multi_buffer_addrs = array} -// PLAN-NOT: pto.pointer_cast // PTOResolveBufferSelect: four addressed alloc_tile ops, no multi-buffer ops. // SELECT: IR Dump After PTOResolveBufferSelect diff --git a/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto b/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto index 3bc17eb8ec..1f0ddd9ce0 100644 --- a/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto +++ b/test/lit/pto/plan_memory_five_gates_lifetime_overlap.pto @@ -66,5 +66,4 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK-DAG: %[[ADDR1:.*]] = arith.constant 512 : i64 // CHECK: pto.alloc_tile addr = %[[ADDR0]] // CHECK: pto.alloc_tile addr = %[[ADDR1]] -// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_five_gates_phi_family.pto b/test/lit/pto/plan_memory_five_gates_phi_family.pto index bbec280795..5007e8874f 100644 --- a/test/lit/pto/plan_memory_five_gates_phi_family.pto +++ b/test/lit/pto/plan_memory_five_gates_phi_family.pto @@ -59,5 +59,4 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK: pto.alloc_tile addr = %[[ADDR0]] // CHECK: } else { // CHECK: pto.alloc_tile addr = %[[ADDR0]] -// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_inplace_forbid_alias.pto b/test/lit/pto/plan_memory_inplace_forbid_alias.pto index de21408ebc..eab357c125 100644 --- a/test/lit/pto/plan_memory_inplace_forbid_alias.pto +++ b/test/lit/pto/plan_memory_inplace_forbid_alias.pto @@ -57,5 +57,4 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK: pto.alloc_tile addr = %[[B1024]] // CHECK: pto.alloc_tile addr = %[[B2048]] // CHECK: pto.alloc_tile addr = %[[B1536]] -// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto b/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto index abce67d01f..aea758ab72 100644 --- a/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto +++ b/test/lit/pto/plan_memory_spec_level0_no_reuse_overlap.pto @@ -40,5 +40,4 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK-DAG: %[[ADDR1:.*]] = arith.constant 2048 : i64 // CHECK: pto.alloc_tile addr = %[[ADDR0]] // CHECK: pto.alloc_tile addr = %[[ADDR1]] -// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/plan_memory_spec_level0_reuse.pto b/test/lit/pto/plan_memory_spec_level0_reuse.pto index 7fcbdebfd5..33a6db3d08 100644 --- a/test/lit/pto/plan_memory_spec_level0_reuse.pto +++ b/test/lit/pto/plan_memory_spec_level0_reuse.pto @@ -48,5 +48,4 @@ module attributes {"pto.target_arch" = "a3"} { // CHECK: %[[ADDR0:.*]] = arith.constant 0 : i64 // CHECK: %[[BUF0:.*]] = pto.alloc_tile addr = %[[ADDR0]] // CHECK: %[[BUF1:.*]] = pto.alloc_tile addr = %[[ADDR0]] -// CHECK-NOT: pto.pointer_cast // CHECK-NOT: memref.alloc diff --git a/test/lit/pto/pointer_cast_addr_alignment_gm_ignored.pto b/test/lit/pto/pointer_cast_addr_alignment_gm_ignored.pto deleted file mode 100644 index e5b1591cb7..0000000000 --- a/test/lit/pto/pointer_cast_addr_alignment_gm_ignored.pto +++ /dev/null @@ -1,9 +0,0 @@ -// RUN: ptoas --pto-level=level3 %s >/dev/null - -module { - func.func @gm_pointer_cast_addr_is_ignored() { - %neg = arith.constant -1 : i64 - %buf = pto.pointer_cast(%neg) : memref<32x64xf16, #pto.address_space> - return - } -} diff --git a/test/lit/pto/pointer_cast_addr_alignment_invalid.pto b/test/lit/pto/pointer_cast_addr_alignment_invalid.pto deleted file mode 100644 index 881bb6581e..0000000000 --- a/test/lit/pto/pointer_cast_addr_alignment_invalid.pto +++ /dev/null @@ -1,11 +0,0 @@ -// RUN: not ptoas --pto-level=level3 %s 2>&1 | FileCheck %s - -module { - func.func @unaligned_pointer_cast_addr() { - %bad = arith.constant 70660 : i64 - // CHECK: pto.pointer_cast - // CHECK: addr[0] must be aligned to 32 bytes for local tile memory space, got 70660 - %buf = pto.pointer_cast(%bad) : memref<32x64xf16, #pto.address_space> - return - } -} diff --git a/test/lit/pto/retired_memref_bridge_ops_invalid.pto b/test/lit/pto/retired_memref_bridge_ops_invalid.pto index 87e0cd1cec..82bde7b255 100644 --- a/test/lit/pto/retired_memref_bridge_ops_invalid.pto +++ b/test/lit/pto/retired_memref_bridge_ops_invalid.pto @@ -11,6 +11,7 @@ // RUN: not ptoas %t/declare.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=DECLARE // RUN: not ptoas %t/slot.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=SLOT // RUN: not ptoas %t/bind.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=BIND +// RUN: not ptoas %t/pointer-cast.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=POINTER-CAST //--- materialize.pto module { @@ -51,3 +52,13 @@ module { } // BIND: error: unregistered operation 'pto.bind_tile' found in dialect ('pto') + +//--- pointer-cast.pto +module { + func.func @pointer_cast() { + %0 = "pto.pointer_cast"() : () -> i32 + return + } +} + +// POINTER-CAST: error: unregistered operation 'pto.pointer_cast' found in dialect ('pto') diff --git a/test/lit/tile_fusion/op_fusion_region_pipeline_level3.pto b/test/lit/tile_fusion/op_fusion_region_pipeline_level3.pto index e67a74f1ec..56c0f74d29 100644 --- a/test/lit/tile_fusion/op_fusion_region_pipeline_level3.pto +++ b/test/lit/tile_fusion/op_fusion_region_pipeline_level3.pto @@ -61,7 +61,6 @@ module { // CHECK: pto.alloc_tile addr = %c12288_i64 // CHECK: pto.alloc_tile addr = %c16384_i64 // CHECK: pto.fusion_region { -// CHECK-NOT: pto.pointer_cast // CHECK: pto.alloc_tile addr = %c0_i64 // CHECK: pto.alloc_tile addr = %c4096_i64 // CHECK: pto.alloc_tile addr = %c8192_i64 diff --git a/test/lit/vpto/fold_tile_buf_intrinsics.pto b/test/lit/vpto/fold_tile_buf_intrinsics.pto index 1c9b59e249..b52453e749 100644 --- a/test/lit/vpto/fold_tile_buf_intrinsics.pto +++ b/test/lit/vpto/fold_tile_buf_intrinsics.pto @@ -10,6 +10,7 @@ // // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-fold-tile-buf-intrinsics %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=SHAPE // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=pto-fold-tile-buf-intrinsics %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=ADDR +// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-ptr-normalize %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NORMALIZED // // After the first shape-only FoldTileBufIntrinsics: // - tile_valid_rows / tile_valid_cols should be gone @@ -25,20 +26,36 @@ // SHAPE: pto.vsts // After the later addr-only FoldTileBufIntrinsics: -// - tile_buf_addr should also be gone -// - the expanded body should use concrete memrefs/constants +// - tile_buf_addr should still carry the planned tile address +// - VPTO pointer normalization will rewrite the memref bridge to pto.ptr later // ADDR-LABEL: func.func @TADD -// ADDR: pto.pointer_cast( +// ADDR: pto.tile_buf_addr {{.*}} -> memref<16x64xf32, #pto.address_space> // ADDR: pto.vlds -// ADDR: pto.pointer_cast( +// ADDR: pto.tile_buf_addr {{.*}} -> memref<16x64xf32, #pto.address_space> // ADDR: pto.vlds // ADDR: pto.vadd -// ADDR: pto.pointer_cast( -// ADDR-NOT: pto.tile_buf_addr +// ADDR: pto.tile_buf_addr {{.*}} -> memref<16x64xf32, #pto.address_space> +// ADDR-NOT: pto.pointer_cast // ADDR-NOT: pto.tile_valid_rows // ADDR-NOT: pto.tile_valid_cols // ADDR: pto.vsts +// After VPTO pointer normalization: +// - tile_buf_addr returns !pto.ptr directly +// - memref.subview has been converted to pto.addptr +// NORMALIZED-LABEL: func.func @TADD +// NORMALIZED: pto.tile_buf_addr {{.*}} -> !pto.ptr +// NORMALIZED: pto.addptr +// NORMALIZED: pto.vlds +// NORMALIZED: pto.tile_buf_addr {{.*}} -> !pto.ptr +// NORMALIZED: pto.addptr +// NORMALIZED: pto.vlds +// NORMALIZED: pto.vadd +// NORMALIZED: pto.tile_buf_addr {{.*}} -> !pto.ptr +// NORMALIZED: pto.addptr +// NORMALIZED-NOT: pto.pointer_cast +// NORMALIZED: pto.vsts + module attributes {pto.kernel_kind = #pto.kernel_kind} { func.func @TADD() { %a = pto.alloc_tile diff --git a/test/lit/vpto/issue_791_mte_ub_gm_ptr_normalize.pto b/test/lit/vpto/issue_791_mte_ub_gm_ptr_normalize.pto index d9a9f664d2..a4e3883e02 100644 --- a/test/lit/vpto/issue_791_mte_ub_gm_ptr_normalize.pto +++ b/test/lit/vpto/issue_791_mte_ub_gm_ptr_normalize.pto @@ -8,10 +8,8 @@ // Regression guard for issue 791. // -// Guard item 1: VPTOPtrNormalize must rewrite memref-backed `pto.mte_ub_gm` -// operands into ptr-form operands before wrapper expansion. If this regresses, -// the op reaches later passes with a memref source and eventually reintroduces -// an illegal `pto.castptr` from memref to !pto.ptr. +// Guard item 1: VPTOPtrNormalize preserves the pointer-form `pto.mte_ub_gm` +// operands before wrapper expansion. // RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-vpto --mlir-print-ir-after=vpto-ptr-normalize %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=NORMALIZE // // Guard item 2: after `vpto-expand-wrapper-ops`, the expanded DMA store must @@ -24,11 +22,10 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind> + %src = pto.castptr %c0_i64 : i64 -> !pto.ptr pto.mte_ub_gm %src, %out_gm, %c32_i64 nburst(%c16_i64, %c32_i64, %c32_i64) - : memref<16x16xf16, #pto.address_space>, !pto.ptr, i64, i64, i64, i64 + : !pto.ptr, !pto.ptr, i64, i64, i64, i64 return } } @@ -44,5 +41,5 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind) attributes {pto.aicore} { // EXPAND: %[[SRC:.*]] = pto.castptr %{{.*}} : i64 -> !pto.ptr -// EXPAND-NOT: pto.castptr %{{.*}} : memref<16x16xf16, #pto.address_space> -> !pto.ptr +// EXPAND-NOT: memref // EXPAND: pto.copy_ubuf_to_gm %[[SRC]], %arg0, %{{.*}} : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 463a382810..0d2243dfe8 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -2979,9 +2979,7 @@ int mlir::pto::compilePTOASModule( if (effectiveLevel == PTOBuildLevel::Level3) { // In level3 the caller owns local memory and PTOPlanMemory is skipped, so // every allocation must carry an explicit physical address. For - // multi-buffer, `addr` is the base of the contiguous N-slot region; the - // alloc lowering fans it out into the multi-address `pto.pointer_cast` - // PlanMemory would otherwise produce. + // multi-buffer, `addr` is the base of the contiguous N-slot region. bool missing = false; module->walk([&](pto::AllocTileOp op) { if (!op.getAddr()) { From 1688ae49185d4fa7afd3592af4550c3c1001b450 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 21 Jul 2026 19:07:14 +0800 Subject: [PATCH 90/91] Remove memref compatibility from sync analysis Drop InsertSync and GraphSyncSolver handling for memref alloc/view/load/store compatibility paths so sync analysis only models PTO tile, pointer, and view roots. Update the auto-sync design note and keep existing tile-native sync coverage passing. --- docs/designs/ptoas-auto-sync-design.md | 6 +- .../Transforms/GraphSyncSolver/SyncSolverIR.h | 2 +- .../GraphSyncSolver/SyncSolverIRTranslator.h | 3 - .../InsertSync/InsertSyncAnalysis.h | 3 - .../Transforms/InsertSync/PTOIRTranslator.h | 2 - .../PTO/Transforms/InsertSync/SyncCommon.h | 2 +- .../Transforms/GraphSyncSolver/SyncSolver.cpp | 4 +- .../SyncSolverIRTranslator.cpp | 35 +- .../InsertSync/InsertSyncAnalysis.cpp | 30 +- .../InsertSync/MemoryDependentAnalyzer.cpp | 45 +-- .../Transforms/InsertSync/PTOIRTranslator.cpp | 307 +----------------- .../Transforms/InsertSync/SyncMacroModel.cpp | 21 +- 12 files changed, 22 insertions(+), 438 deletions(-) diff --git a/docs/designs/ptoas-auto-sync-design.md b/docs/designs/ptoas-auto-sync-design.md index dd9e4fc126..7b41826cfd 100644 --- a/docs/designs/ptoas-auto-sync-design.md +++ b/docs/designs/ptoas-auto-sync-design.md @@ -84,7 +84,7 @@ | 字段 | 类型 | 作用 | | --- | --- | --- | | `baseBuffer` | `Value` | 当前 op 直接看到的 SSA buffer（可能是 view/cast 链顶端） | -| `rootBuffer` | `Value` | 静态可知的最根缓冲区（`alloc_tile` / kernel arg / `memref.alloc`） | +| `rootBuffer` | `Value` | 静态可知的最根缓冲区（`alloc_tile` / `alloc_multi_tile` / kernel ptr arg） | | `scope` | `pto::AddressSpace` | 地址空间（GM/MAT/VEC/ACC/LEFT/RIGHT 等） | | `baseAddresses` | `SmallVector` | 已知的偏移列表，配合 `allocateSize` 做精确区间重叠 | | `allocateSize` | `uint64_t` | 字节大小 | @@ -199,8 +199,8 @@ MLIR op，最终生成 `pto::SetFlagOp` / `pto::WaitFlagOp` / `pto::BarrierOp` - 主要逻辑： 1. `UpdateKernelArgMemInfo()`：把 kernel 参数登记为 GM root buffer。 2. `RecursionIR(&func.getBody())`：前序遍历 region： - - `pto::AllocTileOp` / `DeclareTileMemRefOp` / - `memref::AllocOp` 经 `Update*MemInfo` 写入 `buffer2MemInfoMap_`。 + - `pto::AllocTileOp` / `pto::AllocMultiTileOp` / `pto::DeclareTileOp` + 经 `Update*MemInfo` 写入 `buffer2MemInfoMap_`。 - View / Subview / Cast / Mov 调 `UpdateAliasBufferInfo(result, source)`， 把派生 buffer 的 `BaseMemInfo` 链回到原 root。 - `scf::ForOp` / `WhileOp` / `IfOp` / `YieldOp` 经 `UpdateForOpInfo` / diff --git a/include/PTO/Transforms/GraphSyncSolver/SyncSolverIR.h b/include/PTO/Transforms/GraphSyncSolver/SyncSolverIR.h index 38ee5050e0..14b6d2b191 100644 --- a/include/PTO/Transforms/GraphSyncSolver/SyncSolverIR.h +++ b/include/PTO/Transforms/GraphSyncSolver/SyncSolverIR.h @@ -73,7 +73,7 @@ struct EventIdInfo { // (== rwOp1 / rwOp2 of the conflict). Captured from the actual hazard // MemInfo pair in `getMultiBufferEventIdInfo` so dyn-event-id codegen keys // the event lane off the buffer that drove the event-id allocation, not - // the op's first `slot_marker` memref (an op may touch several + // the op's first slot-bearing value (an op may touch several // multi-buffers at different slots/phases). Null when the slot is absent // or ambiguous (more than one distinct slot in this pair's conflicts), // which makes codegen fall back to the safe N-static `set_flag` / diff --git a/include/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.h b/include/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.h index d986802b1b..096c3a5e47 100644 --- a/include/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.h +++ b/include/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.h @@ -88,9 +88,6 @@ class IRTranslator { std::pair, llvm::SmallVector> getReadWriteMemoryOps(Operation *op); - template - std::unique_ptr getLoadStoreOp(OP op, OperationBase *parentOp); - std::unique_ptr getPipeInterfaceOp(pto::OpPipeInterface op, OperationBase *parentOp); diff --git a/include/PTO/Transforms/InsertSync/InsertSyncAnalysis.h b/include/PTO/Transforms/InsertSync/InsertSyncAnalysis.h index f9e98a0270..cf6f62be6b 100644 --- a/include/PTO/Transforms/InsertSync/InsertSyncAnalysis.h +++ b/include/PTO/Transforms/InsertSync/InsertSyncAnalysis.h @@ -171,9 +171,6 @@ class InsertSyncAnalysis { /// 辅助函数：获取所有涉及的 Buffer (用于 LCA 计算，虽然现在简化了，保留接口) SmallVector GetMemInfoBuffers(const DepBaseMemInfoPairVec &depBaseMemInfosVec); - /// 判断 buffer 是否是 Alloc 类操作 (用于溯源) - bool IsMemAllocOp(Operation *op) const; - /// 判断两个操作是否构成GM的读写冲突 bool IsGMHazard(const CompoundInstanceElement *nowCompound, const CompoundInstanceElement *frontCompound) const; diff --git a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h index 8fd6edd5df..2135beead3 100644 --- a/include/PTO/Transforms/InsertSync/PTOIRTranslator.h +++ b/include/PTO/Transforms/InsertSync/PTOIRTranslator.h @@ -73,12 +73,10 @@ class PTOIRTranslator { LogicalResult UpdateAllocMultiTileOpMemInfo(pto::AllocMultiTileOp op); LogicalResult UpdateDeclareTileOpMemInfo(pto::DeclareTileOp op); LogicalResult UpdateDeclareGlobalOpMemInfo(pto::DeclareGlobalOp op); - LogicalResult UpdateMemrefAllocOpMemInfo(memref::AllocOp op); // 处理 View/Alias (MakeTensorView, Subview, Mov) void UpdateAliasBufferInfo(Value result, Value source); void UpdateConservativeAliasBufferInfo(Value result, Value source); - void UpdateMemrefSubViewAliasBufferInfo(memref::SubViewOp op); void UpdateTileSubViewAliasBufferInfo(pto::SubViewOp op); void UpdateMultiTileGetAliasBufferInfo(pto::MultiTileGetOp op); void UpdateSlotSelectedAliasBufferInfo(Value result, Value source, diff --git a/include/PTO/Transforms/InsertSync/SyncCommon.h b/include/PTO/Transforms/InsertSync/SyncCommon.h index 7b1970a634..49ec77fe86 100644 --- a/include/PTO/Transforms/InsertSync/SyncCommon.h +++ b/include/PTO/Transforms/InsertSync/SyncCommon.h @@ -100,7 +100,7 @@ struct BaseMemInfo { pto::AddressSpace scope; SmallVector baseAddresses; // 用于 Offset 分析 uint64_t allocateSize; - // PlanMemory materializes static local allocations as pointer_cast constants. + // PlanMemory materializes static local allocations as physical addresses. // This distinguishes their physical addresses from root-relative offsets. bool hasKnownPhysicalAddresses; diff --git a/lib/PTO/Transforms/GraphSyncSolver/SyncSolver.cpp b/lib/PTO/Transforms/GraphSyncSolver/SyncSolver.cpp index d9c2e17b05..36bd537806 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/SyncSolver.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/SyncSolver.cpp @@ -500,7 +500,7 @@ Solver::getMultiBufferEventIdInfo(Occurrence *occ1, Occurrence *occ2, // Capture the slot SSA of the conflicting multi-buffer per side, so dyn // event-id codegen keys the event lane off the *hazard* buffer rather than - // the op's first slot_marker memref. `slotOpN` is the slot as accessed by + // the op's first slot-bearing value. `slotOpN` is the slot as accessed by // rwOpN; ambiguity (more than one distinct slot participating in this // pair's conflicts) collapses to null, making codegen fall back to the // safe N-static fanout. @@ -2387,7 +2387,7 @@ SyncBeforeAfterMap Solver::getBeforeAfterSyncMaps() { // codegen can lower into `pto.set_flag_dyn` / `pto.wait_flag_dyn`. // The slots were captured in `getMultiBufferEventIdInfo` from the // actual conflicting MemInfo pair, so the dyn event lane is indexed by - // the *hazard* buffer's slot -- not the op's first slot_marker memref. + // the *hazard* buffer's slot -- not the op's first slot-bearing value. // `slotExprOp1/Op2` are keyed to op1(==rwOp1)/op2(==rwOp2); the set // side may land on either op depending on the program order resolved // by `getSetWaitOcc`, so map by op identity. A null slot (absent or diff --git a/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp b/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp index 94bdad6dbf..f35da80dd4 100644 --- a/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp +++ b/lib/PTO/Transforms/GraphSyncSolver/SyncSolverIRTranslator.cpp @@ -13,10 +13,8 @@ #include "PTO/IR/PTO.h" #include "../Utils.h" -#include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" #include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Interfaces/DestinationStyleOpInterface.h" @@ -125,8 +123,8 @@ llvm::SmallVector IRTranslator::tracebackMemVals(Value val) { Operation *defOp = result.getDefiningOp(); // `pto.multi_tile_get` stops traversal so getMemInfo can extract the // slot-narrowed physical addresses. - if (isa(defOp)) { + if (isa(defOp)) { leaves.insert(result); continue; } @@ -173,23 +171,6 @@ IRTranslator::getReadWriteMemoryOps(Operation *op) { return {std::move(reads), std::move(writes)}; } -template -std::unique_ptr -IRTranslator::getLoadStoreOp(OP loadStoreOp, OperationBase *parentOp) { - auto pipe = pto::PIPE::PIPE_S; - llvm::SmallVector reads; - llvm::SmallVector writes; - if constexpr (std::is_same_v || - std::is_same_v) { - reads = getMemoryOps({loadStoreOp.getMemRef()}); - } else { - writes = getMemoryOps({loadStoreOp.getMemRef()}); - } - return std::make_unique( - loadStoreOp.getOperation(), parentOp, TCoreType::CUBE_OR_VECTOR, pipe, - pipe, reads, writes); -} - std::unique_ptr IRTranslator::getPipeInterfaceOp(pto::OpPipeInterface op, OperationBase *parentOp) { @@ -349,18 +330,6 @@ void IRTranslator::translateBlockIntoScope(Block &block, Scope *parScope, if (auto pipeOp = dyn_cast(op)) { if (auto rw = getPipeInterfaceOp(pipeOp, parScope)) parScope->body.push_back(std::move(rw)); - } else if (auto storeOp = dyn_cast(op)) { - if (auto rw = getLoadStoreOp(storeOp, parScope)) - parScope->body.push_back(std::move(rw)); - } else if (auto loadOp = dyn_cast(op)) { - if (auto rw = getLoadStoreOp(loadOp, parScope)) - parScope->body.push_back(std::move(rw)); - } else if (auto storeOp = dyn_cast(op)) { - if (auto rw = getLoadStoreOp(storeOp, parScope)) - parScope->body.push_back(std::move(rw)); - } else if (auto loadOp = dyn_cast(op)) { - if (auto rw = getLoadStoreOp(loadOp, parScope)) - parScope->body.push_back(std::move(rw)); } else if (auto extractOp = dyn_cast(op)) { if (auto rw = getTensorExtractOp(extractOp, parScope)) parScope->body.push_back(std::move(rw)); diff --git a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp index dc34ffc2c1..a0720fc94a 100644 --- a/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp +++ b/lib/PTO/Transforms/InsertSync/InsertSyncAnalysis.cpp @@ -16,7 +16,6 @@ #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/InsertSync/SyncCommon.h" #include "PTO/Transforms/SlotAffineAnalysis.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/Matchers.h" @@ -63,15 +62,6 @@ static std::optional getKnownRepeatAccessShapeFromType(Type t tileTy.getElementType()}; } - if (auto memRefTy = dyn_cast(ty)) { - if (!memRefTy.hasStaticShape() || memRefTy.getRank() != 2) - return std::nullopt; - auto shape = memRefTy.getShape(); - return RepeatAccessShape{SmallVector{shape[0], shape[1]}, - SmallVector{shape[0], shape[1]}, - memRefTy.getElementType()}; - } - return std::nullopt; } @@ -92,19 +82,6 @@ static std::optional getKnownBLayout(Type ty) { return BLayout::ColMajor; } - if (auto memRefTy = dyn_cast(ty)) { - SmallVector strides; - int64_t offset = 0; - if (failed(mlir::pto::getPTOMemRefStridesAndOffset(memRefTy, strides, - offset)) || - strides.size() != 2) { - return std::nullopt; - } - ArrayRef shape = memRefTy.getShape(); - if (strides[1] == 1 && strides[0] == shape[1]) return BLayout::RowMajor; - if (strides[0] == 1 && strides[1] == shape[0]) return BLayout::ColMajor; - } - return std::nullopt; } @@ -618,8 +595,7 @@ void InsertSyncAnalysis::InsertSyncOperation( if (!producerSlot || !consumerSlot) { // No slot SSA threaded through -- fall back to single event id. // This keeps non-multi-buffer codepaths untouched even if their - // baseAddresses happen to have multiple entries for some other - // reason (e.g. memref subview). + // baseAddresses happen to have multiple entries for another reason. eventIdNum = 1; } else { setOp->slotSSAExpr = producerSlot; @@ -766,10 +742,6 @@ void InsertSyncAnalysis::InsertLastPipeAll() { // 7. Helpers // ============================================================================== -bool InsertSyncAnalysis::IsMemAllocOp(Operation *op) const { - return isa(op); -} - SmallVector InsertSyncAnalysis::GetMemInfoBuffers( const DepBaseMemInfoPairVec &depBaseMemInfosVec) { llvm::DenseSet touchedBuffer; diff --git a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp index 247e5bd96c..fba700f283 100644 --- a/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp +++ b/lib/PTO/Transforms/InsertSync/MemoryDependentAnalyzer.cpp @@ -13,8 +13,6 @@ #include "PTO/Transforms/InsertSync/MemoryDependentAnalyzer.h" #include "PTO/Transforms/InsertSync/InsertSyncDebug.h" -#include "mlir/Interfaces/ViewLikeInterface.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "llvm/Support/Debug.h" @@ -64,39 +62,6 @@ static Value GetRealRoot(Value v) { break; } - if (auto op = dyn_cast(defOp)) { - if (trace) - llvm::errs() << " -> Hit CollapseShapeOp. Peel off.\n"; - v = op.getSrc(); - continue; - } - if (auto op = dyn_cast(defOp)) { - if (trace) - llvm::errs() << " -> Hit ExpandShapeOp. Peel off.\n"; - v = op.getSrc(); - continue; - } - if (auto op = dyn_cast(defOp)) { - if (trace) - llvm::errs() << " -> Hit ViewOp. Peel off.\n"; - v = op.getSource(); - continue; - } - if (auto view = dyn_cast(defOp)) { - if (trace) - llvm::errs() << " -> Hit ViewLikeInterface. Peel off.\n"; - v = view.getViewSource(); - continue; - } - if (auto cast = dyn_cast(defOp)) { - v = cast.getSource(); - continue; - } - if (auto reCast = dyn_cast(defOp)) { - v = reCast.getSource(); - continue; - } - if (trace) { llvm::errs() << " -> Hit Alloc/Other [" << defOp->getName() << "]. Stop.\n"; @@ -231,9 +196,9 @@ bool MemoryDependentAnalyzer::MemAlias(const BaseMemInfo *a, } // 2. Local Memory (UB/L1) - // PTOPlanMemory turns each allocation into a distinct pointer_cast. Once an + // PTOPlanMemory writes physical addresses back to allocation roots. Once an // async MTE3 store has consumed the source SSA, a later allocation can reuse - // the same physical range with a different pointer_cast root. Compare those + // the same physical range with a different alloc_tile root. Compare those // ranges directly when both sides carry known physical local addresses. if (a->hasKnownPhysicalAddresses && b->hasKnownPhysicalAddresses) { if (isTraceEnabled()) @@ -276,9 +241,9 @@ bool MemoryDependentAnalyzer::MemAlias(const BaseMemInfo *a, // 2.3 Cross-root absolute-address overlap check. // - // Roots genuinely differ (different alloc_tile / different pointer_cast - // address SSA). Historically MemAlias returned false here, but PlanMemory - // can reuse the same physical UB region for distinct allocations whose + // Roots genuinely differ (different alloc_tile/address SSA). Historically + // MemAlias returned false here, but PlanMemory can reuse the same physical + // UB region for distinct allocations whose // byte ranges overlap. Re-derive the absolute address and check overlap so // a cross-pipe hazard (e.g. MTE3 tstore vs a V-pipe write into an // overlapping region) is not silently dropped (issue #934). diff --git a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp index 05f6d3e412..a0b1c3e7eb 100644 --- a/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp +++ b/lib/PTO/Transforms/InsertSync/PTOIRTranslator.cpp @@ -16,14 +16,11 @@ #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/InsertSync/SyncMacroModel.h" #include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "llvm/Support/Debug.h" #include "mlir/IR/AsmState.h" -#include "llvm/Support/FormatVariadic.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/IR/Matchers.h" #include "llvm/ADT/StringSwitch.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/FormatVariadic.h" // [P0 新增] 引入副作用接口和 PTO 接口 #include "mlir/Interfaces/SideEffectInterfaces.h" @@ -37,10 +34,8 @@ using namespace mlir::pto; namespace { constexpr size_t kTileRank2D = 2; -constexpr unsigned kStrideInlineCapacity = 4; constexpr unsigned kMemoryEffectInlineCapacity = 4; -using StrideVector = SmallVector; using MemoryEffectVector = SmallVector, kMemoryEffectInlineCapacity>; @@ -158,62 +153,6 @@ getPtoSubViewBaseAddresses(pto::SubViewOp op, pto::TileBufType sourceType, return addresses; } -static std::optional> -getMemrefSubViewBaseAddresses(memref::SubViewOp op, MemRefType sourceType, - int64_t elemBytes) { - if (!sourceType.hasStaticShape() || sourceType.getRank() != kTileRank2D) - return std::nullopt; - - SmallVector strides; - int64_t baseOffset = ShapedType::kDynamic; - if (failed(mlir::pto::getPTOMemRefStridesAndOffset(sourceType, strides, - baseOffset)) || - strides.size() != 2 || - llvm::is_contained(strides, ShapedType::kDynamic)) - return std::nullopt; - - ArrayRef staticOffsets = op.getStaticOffsets(); - ArrayRef staticSizes = op.getStaticSizes(); - ArrayRef staticSubViewStrides = op.getStaticStrides(); - if (staticOffsets.size() != 2 || staticSizes.size() != 2 || - staticSubViewStrides.size() != 2) - return std::nullopt; - if (llvm::is_contained(staticOffsets, ShapedType::kDynamic) || - llvm::is_contained(staticSizes, ShapedType::kDynamic) || - llvm::is_contained(staticSubViewStrides, ShapedType::kDynamic)) - return std::nullopt; - if (staticSubViewStrides[0] != 1 || staticSubViewStrides[1] != 1) - return std::nullopt; - - int64_t rowOffset = staticOffsets[0]; - int64_t colOffset = staticOffsets[1]; - int64_t rowSize = staticSizes[0]; - int64_t colSize = staticSizes[1]; - if (rowOffset < 0 || colOffset < 0 || rowSize <= 0 || colSize <= 0) - return std::nullopt; - - SmallVector addresses; - if (strides[1] == 1) { - addresses.reserve(static_cast(rowSize)); - for (int64_t row = 0; row < rowSize; ++row) { - int64_t elemOffset = - (rowOffset + row) * strides[0] + colOffset * strides[1]; - addresses.push_back(static_cast(elemOffset * elemBytes)); - } - } else if (strides[0] == 1) { - addresses.reserve(static_cast(colSize)); - for (int64_t col = 0; col < colSize; ++col) { - int64_t elemOffset = - rowOffset * strides[0] + (colOffset + col) * strides[1]; - addresses.push_back(static_cast(elemOffset * elemBytes)); - } - } else { - return std::nullopt; - } - - return addresses; -} - namespace { static func::FuncOp lookupPTODSLSubkernelHelper(func::CallOp callOp) { @@ -243,7 +182,7 @@ getPTODSLSubkernelHelperPipe(func::FuncOp callee) { } static bool isPTODSLSubkernelMemoryOperand(Type type) { - return isa(type); } @@ -255,69 +194,6 @@ static pto::TCoreType getPTODSLSubkernelHelperCoreType( } // namespace -// [辅助函数] 尝试从 Operation 中计算相对于 Source 的字节偏移量和新大小 -// 返回值: pair -// 如果无法计算静态值，返回 {-1, -1} 表示这是动态的 -static std::pair getStaticOffsetAndSize(Operation *op, Value src) { - auto srcType = dyn_cast(src.getType()); - if (!srcType) return {0, 0}; - - int64_t elemSize = - static_cast(pto::getPTOStorageElemByteSize(srcType.getElementType())); - if (elemSize == 0) - return {-1, -1}; - - // === Case 1: memref.subview === - if (auto subView = dyn_cast(op)) { - int64_t baseOffset; - StrideVector strides; - if (failed(mlir::pto::getPTOMemRefStridesAndOffset(srcType, strides, - baseOffset))) { - return {-1, -1}; - } - - int64_t newSize = 1; - for (int64_t s : subView.getStaticSizes()) { - if (s == ShapedType::kDynamic) return {-1, -1}; - newSize *= s; - } - newSize *= elemSize; - - int64_t totalOffset = 0; - auto staticOffsets = subView.getStaticOffsets(); - - if (staticOffsets.empty()) return {-1, -1}; - if (staticOffsets.size() > strides.size()) return {-1, -1}; - - for (size_t i = 0; i < staticOffsets.size(); ++i) { - int64_t off = staticOffsets[i]; - if (off == ShapedType::kDynamic) return {-1, -1}; - - int64_t stride = 1; - if (i < strides.size() && strides[i] != ShapedType::kDynamic) { - stride = strides[i]; - } else { - return {-1, -1}; - } - - totalOffset += off * stride; - } - - return {totalOffset * elemSize, newSize}; - } - - // === Case 2: memref.reinterpret_cast === - if (auto castOp = dyn_cast(op)) { - auto staticOffsets = castOp.getStaticOffsets(); - if (staticOffsets.empty() || staticOffsets[0] == ShapedType::kDynamic) { - return {0, 0}; - } - return {staticOffsets[0] * elemSize, 0}; - } - - return {0, 0}; -} - // ============================================================================ // 1. 构建入口 // ============================================================================ @@ -336,7 +212,7 @@ void PTOIRTranslator::UpdateKernelArgMemInfo() { Value funcArg = func_.getArgument(i); Type argType = funcArg.getType(); - if (!isa(argType) && !isa(argType)) { + if (!isa(argType)) { continue; } @@ -369,12 +245,6 @@ void PTOIRTranslator::RecursionIR(Region *region) { return WalkResult::interrupt(); } } - // 支持标准 memref.alloc - else if (auto memAllocOp = dyn_cast(op)) { - if (failed(UpdateMemrefAllocOpMemInfo(memAllocOp))) { - return WalkResult::interrupt(); - } - } else if (auto declareOp = dyn_cast(op)) { if (failed(UpdateDeclareTileOpMemInfo(declareOp))) return WalkResult::interrupt(); @@ -394,9 +264,6 @@ void PTOIRTranslator::RecursionIR(Region *region) { else if (auto subViewOp = dyn_cast(op)) { UpdateTileSubViewAliasBufferInfo(subViewOp); } - else if (auto memrefSubView = dyn_cast(op)) { - UpdateMemrefSubViewAliasBufferInfo(memrefSubView); - } else if (auto multiGet = dyn_cast(op)) { UpdateMultiTileGetAliasBufferInfo(multiGet); } @@ -406,31 +273,6 @@ void PTOIRTranslator::RecursionIR(Region *region) { else if (auto bitcast = dyn_cast(op)) { UpdateAliasBufferInfo(bitcast.getResult(), bitcast.getSrc()); } - else if (auto castOp = dyn_cast(op)) { - UpdateAliasBufferInfo(castOp.getResult(), castOp.getSource()); - } - // [Fix] 添加 CollapseShape 和 ExpandShape 的支持 - else if (auto collapseOp = dyn_cast(op)) { - UpdateAliasBufferInfo(collapseOp.getResult(), collapseOp.getSrc()); - } - else if (auto expandOp = dyn_cast(op)) { - UpdateAliasBufferInfo(expandOp.getResult(), expandOp.getSrc()); - } - // arith.select on memref values: emitted by `PTOResolveBufferSelect` - // for the dynamic-slot path of `pto.multi_tile_get`. The result - // conservatively aliases both branches so sync analysis cannot miss a - // dependency that arises through either possible slot. This preserves - // correctness for runtime slot indices; finer-grained "slot - // expressions are disjoint" reasoning is a future affine-analysis - // upgrade and would let sync emit fewer flags in prefetch idioms. - else if (auto selectOp = dyn_cast(op)) { - if (isa(selectOp.getResult().getType())) { - UpdateConservativeAliasBufferInfo(selectOp.getResult(), - selectOp.getTrueValue()); - UpdateConservativeAliasBufferInfo(selectOp.getResult(), - selectOp.getFalseValue()); - } - } // --- Case C: 控制流 (SCF) --- else if (auto forOp = dyn_cast(op)) { @@ -908,34 +750,9 @@ void PTOIRTranslator::UpdateAliasBufferInfo(Value result, Value source) { if (!result || !source) return; if (!buffer2MemInfoMap_.contains(source)) return; - int64_t deltaOffset = 0; - int64_t newSize = -1; - - if (auto op = result.getDefiningOp()) { - auto info = getStaticOffsetAndSize(op, source); - if (info.first != -1) { - deltaOffset = info.first; - if (info.second > 0) newSize = info.second; - } - } - auto &resultMemInfoVec = buffer2MemInfoMap_[result]; - - for (auto &parentInfo : buffer2MemInfoMap_[source]) { - auto newInfo = parentInfo->clone(result); - - if (!newInfo->baseAddresses.empty()) { - newInfo->baseAddresses[0] += deltaOffset; - } else { - newInfo->baseAddresses.push_back(deltaOffset); - } - - if (newSize > 0) { - newInfo->allocateSize = newSize; - } - - resultMemInfoVec.emplace_back(std::move(newInfo)); - } + for (auto &parentInfo : buffer2MemInfoMap_[source]) + resultMemInfoVec.emplace_back(parentInfo->clone(result)); } void PTOIRTranslator::UpdateConservativeAliasBufferInfo(Value result, @@ -987,72 +804,6 @@ void PTOIRTranslator::UpdateSlotSelectedAliasBufferInfo(Value result, } } -void PTOIRTranslator::UpdateMemrefSubViewAliasBufferInfo(memref::SubViewOp op) { - Value result = op.getResult(); - Value source = op.getSource(); - if (!result || !source) return; - if (!buffer2MemInfoMap_.contains(source)) return; - - auto sourceType = dyn_cast(source.getType()); - if (!sourceType) { - UpdateConservativeAliasBufferInfo(result, source); - return; - } - - unsigned elemBytes = pto::getPTOStorageElemByteSize(sourceType.getElementType()); - auto subViewAddresses = - elemBytes == 0 ? std::nullopt - : getMemrefSubViewBaseAddresses( - op, sourceType, static_cast(elemBytes)); - if (!subViewAddresses || subViewAddresses->empty()) { - UpdateConservativeAliasBufferInfo(result, source); - return; - } - - SmallVector strides; - int64_t baseOffset = ShapedType::kDynamic; - if (failed(mlir::pto::getPTOMemRefStridesAndOffset(sourceType, strides, - baseOffset)) || - strides.size() != 2) { - UpdateConservativeAliasBufferInfo(result, source); - return; - } - - ArrayRef staticSizes = op.getStaticSizes(); - uint64_t segmentSize = 0; - if (strides[1] == 1) { - segmentSize = static_cast( - staticSizes[1] * static_cast(elemBytes)); - } else if (strides[0] == 1) { - segmentSize = static_cast( - staticSizes[0] * static_cast(elemBytes)); - } else { - UpdateConservativeAliasBufferInfo(result, source); - return; - } - - for (auto &parentInfo : buffer2MemInfoMap_[source]) { - if (!parentInfo || parentInfo->baseAddresses.size() != 1 || - parentInfo->allocateSize == 0) { - UpdateConservativeAliasBufferInfo(result, source); - return; - } - } - - auto &resultMemInfoVec = buffer2MemInfoMap_[result]; - for (auto &parentInfo : buffer2MemInfoMap_[source]) { - auto newInfo = parentInfo->clone(result); - SmallVector addresses; - addresses.reserve(subViewAddresses->size()); - uint64_t parentBase = parentInfo->baseAddresses[0]; - for (uint64_t offset : *subViewAddresses) - addresses.push_back(parentBase + offset); - newInfo->baseAddresses = std::move(addresses); - newInfo->allocateSize = segmentSize; - resultMemInfoVec.emplace_back(std::move(newInfo)); - } -} - void PTOIRTranslator::UpdateTileSubViewAliasBufferInfo(pto::SubViewOp op) { Value result = op.getResult(); Value source = op.getSource(); @@ -1106,52 +857,6 @@ void PTOIRTranslator::UpdateTileSubViewAliasBufferInfo(pto::SubViewOp op) { } } -// ============================================================================ -// 实现 UpdateMemrefAllocOpMemInfo -// ============================================================================ -LogicalResult PTOIRTranslator::UpdateMemrefAllocOpMemInfo(memref::AllocOp op) { - Value res = op.getResult(); - auto memRefType = dyn_cast(res.getType()); - if (!memRefType) return failure(); - - // 1. 计算大小 (Bytes) - uint64_t sizeInBytes = 0; - if (memRefType.hasStaticShape()) { - int64_t elemSize = static_cast( - pto::getPTOStorageElemByteSize(memRefType.getElementType())); - if (elemSize == 0) - return failure(); - - int64_t numElements = 1; - for (auto dim : memRefType.getShape()) numElements *= dim; - sizeInBytes = numElements * elemSize; - } - - // 2. 解析地址空间 (Scope) - // 默认视为 MAT/UB (Local Memory)，这是 alloc 的常见用途 - // 如果有显式属性，则覆盖 - pto::AddressSpace space = pto::AddressSpace::MAT; - - if (auto attr = memRefType.getMemorySpace()) { - if (auto ptoAttr = dyn_cast(attr)) { - space = ptoAttr.getAddressSpace(); - } - } - - // 3. 注册 Buffer 信息 - // 对于 alloc，它自己就是 Root - auto newMemInfo = std::make_unique( - res, // baseBuffer - res, // rootBuffer (Self is root) - space, - SmallVector{0}, // Base Addresses (Offset 0) - sizeInBytes - ); - - buffer2MemInfoMap_[res].emplace_back(newMemInfo->clone()); - return success(); -} - void PTOIRTranslator::UpdateDefUseVec(ValueRange values, SmallVector &vec) { for (Value v : values) { if (buffer2MemInfoMap_.contains(v)) { diff --git a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp index 89a911bcc6..eeaed5eb61 100644 --- a/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp +++ b/lib/PTO/Transforms/InsertSync/SyncMacroModel.cpp @@ -8,7 +8,6 @@ #include "PTO/Transforms/InsertSync/SyncMacroModel.h" #include "PTO/IR/PTO.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/IR/Matchers.h" using namespace mlir; @@ -195,9 +194,7 @@ std::optional getTGatherSyncMacroModel(pto::TGatherOp op) { return model; } -// Resolve the PTO address space of an MGatherOp operand type. Mirrors the -// getAddressSpace helper used by MGatherOp::getPipe() and retain compatibility -// with legacy imported memref<...,mat> IR. +// Resolve the PTO address space of an MGatherOp tile operand type. static std::optional getMGatherOperandAddressSpace(::mlir::Type ty) { if (auto tb = ::mlir::dyn_cast<::mlir::pto::TileBufType>(ty)) { @@ -206,13 +203,6 @@ getMGatherOperandAddressSpace(::mlir::Type ty) { return as.getAddressSpace(); return std::nullopt; } - if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty)) { - if (auto ms = mr.getMemorySpace()) { - if (auto as = ::mlir::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) - return as.getAddressSpace(); - } - return std::nullopt; - } return std::nullopt; } @@ -223,8 +213,6 @@ getMGatherOperandShapeFromType(::mlir::Type ty, bool useValidShape = true) { useValidShape ? tb.getValidShape() : tb.getShape(); return SmallVector(shape.begin(), shape.end()); } - if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty)) - return SmallVector(mr.getShape().begin(), mr.getShape().end()); return std::nullopt; } @@ -239,13 +227,6 @@ static std::optional getConstantIndex(Value value) { } static std::optional> lookupMGatherValidDims(Value value) { - if (auto subview = value.getDefiningOp()) - return lookupMGatherValidDims(subview.getSource()); - if (auto cast = value.getDefiningOp()) - return lookupMGatherValidDims(cast.getSource()); - if (auto cast = value.getDefiningOp()) - return lookupMGatherValidDims(cast.getSource()); - if (auto regionResult = dyn_cast(value)) { if (auto fusionRegion = dyn_cast(regionResult.getOwner())) { auto yieldOp = From bab26e3b6647fe73af7dd6a15e5389960c6a6829 Mon Sep 17 00:00:00 2001 From: FangRui Date: Tue, 21 Jul 2026 21:01:59 +0800 Subject: [PATCH 91/91] Remove memref entry points from PTO address views --- include/PTO/IR/PTOOps.td | 109 +---- include/PTO/IR/VPTOOps.td | 24 +- include/PTO/Transforms/Passes.td | 5 +- lib/PTO/IR/PTO.cpp | 393 +++--------------- lib/PTO/IR/VPTO.cpp | 51 +-- lib/PTO/Transforms/FoldTileBufIntrinsics.cpp | 25 +- lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp | 183 -------- ptodsl/ptodsl/_ops.py | 31 +- ptodsl/ptodsl/_surface_values.py | 108 +---- ptodsl/ptodsl/tilelib/_render_runtime.py | 9 +- ptodsl/ptodsl/tilelib/metadata.py | 25 +- .../fixtures/tadd_a5_8x64_f32.golden.mlir | 20 +- ptodsl/tests/test_jit_compile.py | 15 +- ptodsl/tests/test_tilelib_catalog.py | 12 - ptodsl/tests/test_tilelib_constraints.py | 2 +- ptodsl/tests/test_tilelib_daemon.py | 2 +- ptodsl/tests/test_tilelib_elementwise.py | 2 +- ptodsl/tests/test_tilelib_render.py | 7 +- .../cmo_cacheinvalid_single_line_invalid.pto | 2 +- .../pto/issue481_addptr_gm_slot_buffer.pto | 4 +- .../issue481_addptr_gm_slot_buffer_gss.pto | 4 +- ...e622_v_mte2_eventid_overlap_reproducer.pto | 4 +- test/lit/pto/ptr_int_cast.pto | 8 +- .../lit/pto/tile_buf_addr_tile_native_arg.pto | 4 +- ...h_tpop_globaltensor_split_half_slot_a3.pto | 4 +- test/lit/vpto/fold_tile_buf_intrinsics.pto | 24 +- ...o_tilelang_inline_soft_divmod_fastpath.pto | 10 +- 27 files changed, 153 insertions(+), 934 deletions(-) diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 6f08bea192..23573a1917 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -34,9 +34,6 @@ include "mlir/Interfaces/ViewLikeInterface.td" // Types //===----------------------------------------------------------------------===// -def TensorOrMemref : - AnyTypeOf<[AnyMemRef, AnyRankedTensor], "Tensor or Memref">; - def PTODpsType : AnyTypeOf<[AnyRankedTensor, PartitionTensorViewType, TileBufType]>; @@ -48,13 +45,6 @@ def PTOCommType : AnyTypeOf<[AnyRankedTensor, TensorViewType, PartitionTensorViewType, TileBufType], "TensorView, PartitionTensorView, TileBuf, or Tensor">; -// VPTO address extraction may still materialize a low-level memref result. -def PtrOrMemRef : - AnyTypeOf<[PtrType, AnyMemRef], "Ptr or MemRef">; - -def TileBufOrMemRef : - AnyTypeOf<[TileBufType, AnyMemRef], "TileBuf or MemRef">; - def ScalarType : AnyTypeOf<[AnySignlessInteger, AnyFloat], "numeric (integer/float)">; @@ -829,13 +819,6 @@ def TStoreOp: PTO_TOp<"tstore", [ return as.getAddressSpace() == ::mlir::pto::AddressSpace::ACC; return false; } - if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty)) { - if (auto ms = mr.getMemorySpace()) { - if (auto as = ::mlir::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) - return as.getAddressSpace() == ::mlir::pto::AddressSpace::ACC; - } - return false; - } return false; }; if (isAcc(getSrc().getType())) @@ -883,21 +866,6 @@ def TTransOp : PTO_TOp<"ttrans", [ } }]; } -//===----------------------------------------------------------------------===// -// TMATMUL_BIAS / TMATMUL_MX family -// - DPS: AnyMemRef -// - TOp: PTODpsType -// - All in ins/outs form -//===----------------------------------------------------------------------===// - -// ------------------------- -// DPS versions (AnyMemRef) -// ------------------------- - -// ------------------------- -// Tile-world TOp versions (PTODpsType) -// ------------------------- - def TMatmulBiasOp : PTO_TOp<"tmatmul.bias", [ PTO_DpsInitOpInterface, OpPipeInterface, @@ -1366,23 +1334,12 @@ def TMovOp : PTO_TOp<"tmov", [ // while UB->UB copies are vector-pipe operations. auto getASFromType = [](Type ty) -> std::optional<::mlir::pto::AddressSpace> { - // Pre-lowering: tile_buf carries the address space in its memorySpace. if (auto tb = llvm::dyn_cast<::mlir::pto::TileBufType>(ty)) { if (auto as = llvm::dyn_cast_or_null<::mlir::pto::AddressSpaceAttr>( tb.getMemorySpace())) return as.getAddressSpace(); return std::nullopt; } - // Legacy imported IR may carry memref values with the address space in - // the memref memorySpace attribute. - if (auto mr = llvm::dyn_cast<::mlir::MemRefType>(ty)) { - if (auto ms = mr.getMemorySpace()) { - if (auto as = - llvm::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) - return as.getAddressSpace(); - } - return std::nullopt; - } return std::nullopt; }; @@ -1415,8 +1372,6 @@ def TMovOp : PTO_TOp<"tmov", [ Type dstElemTy; if (auto tb = llvm::dyn_cast<::mlir::pto::TileBufType>(getDst().getType())) dstElemTy = tb.getElementType(); - else if (auto mr = llvm::dyn_cast<::mlir::MemRefType>(getDst().getType())) - dstElemTy = mr.getElementType(); if (dstElemTy && llvm::isa<::mlir::pto::F8E8M0Type>(dstElemTy)) return ::mlir::pto::PIPE::PIPE_MTE1; return ::mlir::pto::PIPE::PIPE_FIX; @@ -2282,13 +2237,6 @@ def TPushOp : PTO_TOp<"tpush", [ return as.getAddressSpace(); return std::nullopt; } - if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty)) { - if (auto ms = mr.getMemorySpace()) { - if (auto as = ::mlir::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) - return as.getAddressSpace(); - } - return std::nullopt; - } return std::nullopt; }; @@ -2519,13 +2467,6 @@ def TPopOp : PTO_TOp<"tpop", [ return as.getAddressSpace(); return std::nullopt; } - if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty)) { - if (auto ms = mr.getMemorySpace()) { - if (auto as = ::mlir::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) - return as.getAddressSpace(); - } - return std::nullopt; - } return std::nullopt; }; @@ -2878,32 +2819,6 @@ def SyncAllOp : PTO_Op<"syncall", [AttrSizedOperandSegments]> { let hasCustomAssemblyFormat = 1; } -//===----------------------------------------------------------------------===// -// SIMD Bridge Ops -//===----------------------------------------------------------------------===// - -def SimdTileToMemrefOp : PTO_Op<"simd.tile_to_memref", [Pure]> { - let summary = "Bridge cast from tile_buf to memref in OP-Lib bodies."; - let description = [{ - This op is the canonical bridge marker for OP-Lib templates to expose a - memref view from tile-like values while keeping external ABI as - !pto.tile_buf. - In tile_buf world, src is !pto.tile_buf and dst is the corresponding - memref bridge type. - After memref-world lowering, src may already be memref and this op remains - as a marker for backend lowering (EmitC) to materialize tile data access. - }]; - - let arguments = (ins TileBufOrMemRef:$src); - let results = (outs AnyMemRef:$dst); - - let hasVerifier = 1; - - let assemblyFormat = [{ - $src attr-dict `:` qualified(type($src)) `to` type($dst) - }]; -} - // --------------------------------------------------------------------------- // TileBuf intrinsics — used in TileLang DSL-generated template functions. // These ops extract memref address and valid shape from tile_buf parameters. @@ -2950,7 +2865,7 @@ def TileValidColsOp : PTO_Op<"tile_valid_cols", [Pure]> { def SetFFTsOp : PTO_Op<"set_ffts", [MemoryEffects<[MemRead, MemWrite]>]> { let summary = "Set FFTS/flags pointer for runtime (side-effecting)."; - let arguments = (ins AnyMemRef:$ffts); + let arguments = (ins PtrType:$ffts); let results = (outs); let assemblyFormat = "$ffts attr-dict `:` type($ffts)"; @@ -4448,14 +4363,6 @@ def TExtractOp : PTO_TOp<"textract", [ return as.getAddressSpace(); return std::nullopt; } - if (auto mr = llvm::dyn_cast<::mlir::MemRefType>(ty)) { - if (auto ms = mr.getMemorySpace()) { - if (auto as = - llvm::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) - return as.getAddressSpace(); - } - return std::nullopt; - } return std::nullopt; }; @@ -4594,13 +4501,6 @@ def TInsertOp : PTO_TOp<"tinsert", [ return as.getAddressSpace(); return std::nullopt; } - if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty)) { - if (auto ms = mr.getMemorySpace()) { - if (auto as = ::mlir::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) - return as.getAddressSpace(); - } - return std::nullopt; - } return std::nullopt; }; @@ -4688,13 +4588,6 @@ def TFillPadOp : PTO_TOp<"tfillpad", [ return as.getAddressSpace(); return std::nullopt; } - if (auto mr = ::mlir::dyn_cast<::mlir::MemRefType>(ty)) { - if (auto ms = mr.getMemorySpace()) { - if (auto as = ::mlir::dyn_cast<::mlir::pto::AddressSpaceAttr>(ms)) - return as.getAddressSpace(); - } - return std::nullopt; - } return std::nullopt; }; diff --git a/include/PTO/IR/VPTOOps.td b/include/PTO/IR/VPTOOps.td index 9419565165..6bfd81f08f 100644 --- a/include/PTO/IR/VPTOOps.td +++ b/include/PTO/IR/VPTOOps.td @@ -216,19 +216,15 @@ def TensorViewAddrOp : PTO_Op<"tensor_view_addr", [Pure]> { let summary = "Extract address from a tensor view."; let description = [{ Returns the address view carried by a `tensor_view` or - `partition_tensor_view` value. The result may be either a memref view or a - typed PTO pointer, depending on the requested destination type. + `partition_tensor_view` value as a typed PTO pointer. In authoring-form IR this op preserves the descriptor-style surface; tile-native inputs remain unchanged through memory planning and sync. - - This op may also accept a legacy memref operand, in which case it behaves - as an identity marker and is removed by compatibility lowering. }]; - let arguments = (ins AnyTypeOf<[TensorViewType, PartitionTensorViewType, AnyMemRef], - "TensorViewLike or MemRef">:$src); - let results = (outs PtrOrMemRef:$dst); + let arguments = (ins AnyTypeOf<[TensorViewType, PartitionTensorViewType], + "TensorViewLike">:$src); + let results = (outs PtrType:$dst); let hasVerifier = 1; @@ -241,19 +237,15 @@ def TileBufAddrOp : PTO_Op<"tile_buf_addr", [Pure]> { let summary = "Extract address from a tile_buf."; let description = [{ Returns the address view of the data region of a `tile_buf` value. - The result may be either a memref view or a typed PTO pointer, depending - on the requested destination type. Memref results use the tile's static - shape and address space. + The result is a typed PTO pointer. This op is emitted by TileLang DSL templates and resolved by the FoldTileBufIntrinsics pass after inlining. Tile-native uses remain in PTO - IR through memory planning and sync. Hand-written legacy `.pto` may also - use it directly on an addressed memref. + IR through memory planning and sync. }]; - let arguments = (ins AnyTypeOf<[TileBufType, AnyMemRef], - "tile_buf or tile-bound memref">:$src); - let results = (outs PtrOrMemRef:$dst); + let arguments = (ins TileBufType:$src); + let results = (outs PtrType:$dst); let hasVerifier = 1; diff --git a/include/PTO/Transforms/Passes.td b/include/PTO/Transforms/Passes.td index 3b6c3ad571..2634b2a7f5 100644 --- a/include/PTO/Transforms/Passes.td +++ b/include/PTO/Transforms/Passes.td @@ -549,9 +549,8 @@ def FoldTileBufIntrinsics : Pass<"pto-fold-tile-buf-intrinsics", "mlir::func::Fu structured-view intrinsics whose operands are now bound to concrete values. This pass resolves them: - - pto.tile_buf_addr → replaced by the memref recovered from the active - materialized tile handle, or by pto.castptr when the requested result - type is !pto.ptr + - pto.tile_buf_addr → replaced by pto.castptr from the active materialized + tile handle address - pto.tile_valid_rows → folded to arith.constant if v_row is static, or replaced with the dynamic valid_row operand carried by the materialized tile handle diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 98dd5920d6..09f8290f9f 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -235,7 +235,7 @@ static ParseResult parseI32LiteralAttr(OpAsmParser &parser, IntegerAttr &attr); } static int64_t getPTOTypeRank(Type type) { - // 1. 处理标准的 MLIR 类型 (MemRef, Tensor, Vector) + // 1. 处理标准的 MLIR 类型 (Tensor, Vector) if (auto shapedTy = dyn_cast(type)) { if (shapedTy.hasRank()) return shapedTy.getRank(); @@ -306,16 +306,6 @@ func::FuncOp mlir::pto::lookupPeerFuncAcrossContainer(Operation *op, return {}; } -static bool isGmAddressSpaceAttr(Attribute memorySpace) { - if (!memorySpace) - return true; - if (auto addr = mlir::dyn_cast(memorySpace)) - return addr.getAddressSpace() == pto::AddressSpace::GM; - if (auto intAttr = mlir::dyn_cast(memorySpace)) - return intAttr.getInt() == 0; - return false; -} - static bool isA5DeviceSpec(StringRef spec) { return spec.starts_with("Ascend950") || spec.starts_with("Ascend910_95"); } @@ -1902,8 +1892,6 @@ static LogicalResult verifyRowReductionSrcLayout(Operation *op, Type ty, if (tb.getBLayoutValueI32() != static_cast(pto::BLayout::RowMajor)) return op->emitOpError() << "expects " << name << " to use the row_major blayout"; } - if (auto mr = dyn_cast(ty)) - (void)mr; if (auto tb = dyn_cast(ty)) { if (tb.getSLayoutValueI32() != static_cast(pto::SLayout::NoneBox)) return op->emitOpError() << "expects " << name @@ -1936,8 +1924,6 @@ static LogicalResult verifyRowReductionDstLayout(Operation *op, Type ty, return op->emitOpError() << "expects " << name << " to use the row_major or col_major blayout"; } - if (auto mr = dyn_cast(ty)) - (void)mr; if (auto tb = dyn_cast(ty)) { auto layout = getTileBufLogicalLayout(tb); if (layout && *layout == pto::Layout::DN) { @@ -2617,30 +2603,17 @@ LogicalResult mlir::pto::CastPtrOp::verify() { auto inputPtrType = dyn_cast(inputType); auto resultPtrType = dyn_cast(resultType); - auto inputMemRefType = dyn_cast(inputType); bool inputIsInteger = isa(inputType); bool resultIsInteger = isa(resultType); - if (!inputPtrType && !inputMemRefType && !inputIsInteger) - return emitOpError("input must be an integer, memref, or !pto.ptr<...>"); + if (!inputPtrType && !inputIsInteger) + return emitOpError("input must be an integer or !pto.ptr<...>"); if (!resultPtrType && !resultIsInteger) return emitOpError("result must be an integer or !pto.ptr<...>"); if (inputIsInteger && resultIsInteger) return emitOpError("integer-to-integer cast is not a ptr cast"); - if (inputMemRefType && resultIsInteger) - return emitOpError("memref-to-integer cast is unsupported"); - - if (inputMemRefType && resultPtrType) { - auto memrefSpace = dyn_cast_or_null( - inputMemRefType.getMemorySpace()); - auto resultSpace = resultPtrType.getMemorySpace(); - if (memrefSpace && memrefSpace != resultSpace) - return emitOpError( - "memref-to-ptr cast must stay within the same PTO memory space"); - } - if (inputPtrType && resultPtrType && inputPtrType.getMemorySpace() != resultPtrType.getMemorySpace()) { return emitOpError("ptr-to-ptr cast must stay within the same PTO memory space"); @@ -2673,13 +2646,7 @@ void PTODialect::initialize() { AddressSpaceAttr mlir::pto::getPTOAddressSpaceAttr(Type type) { if (auto ptrType = dyn_cast(type)) return ptrType.getMemorySpace(); - auto memRefType = dyn_cast(type); - if (!memRefType) - return {}; - auto scopeAttr = dyn_cast(memRefType.getMemorySpace()); - if (!scopeAttr) - return {}; - return scopeAttr; + return {}; } bool mlir::pto::hasExplicitPTOEntryAttr(func::FuncOp func) { @@ -2763,69 +2730,8 @@ void mlir::pto::annotatePTOEntryFunctions(ModuleOp module) { //===----------------------------------------------------------------------===// // PTO Load/Store/Addf (non-DPS polymorphic) verification + inference. -// - If operands are memref/tensor: verify strictly. -// - Otherwise (tile_view/tile etc): accept (so old IR can still parse). //===----------------------------------------------------------------------===// -[[maybe_unused]] static LogicalResult verifyMemrefToTensorLoad(Operation *op, Value src, Value res) { - auto mr = dyn_cast(src.getType()); - auto rt = dyn_cast(res.getType()); - if (!mr) - return success(); // non-memref case: don't block old IR - if (!rt) - return op->emitOpError("when src is memref, result must be ranked tensor"); - - if (mr.getElementType() != rt.getElementType()) - return op->emitOpError() << "memref/tensor element type mismatch: memref=" - << mr.getElementType() << " tensor=" << rt.getElementType(); - - if (mr.getRank() != rt.getRank()) - return op->emitOpError() << "rank mismatch: memref rank=" << mr.getRank() - << " tensor rank=" << rt.getRank(); - - if (mr.hasStaticShape()) { - if (!rt.hasStaticShape()) - return op->emitOpError("memref has static shape but result tensor is not static"); - if (mr.getShape() != rt.getShape()) - return op->emitOpError() << "shape mismatch: memref=" << mr << " tensor=" << rt; - } else { - // For dynamic memref dims: if tensor dim is static, allow it; if it's dynamic too, also fine. - // We only reject when a memref static dim conflicts with tensor static dim. - for (int64_t i = 0; i < mr.getRank(); ++i) { - int64_t md = mr.getDimSize(i); - int64_t td = rt.getDimSize(i); - if (md != ShapedType::kDynamic && td != ShapedType::kDynamic && md != td) - return op->emitOpError() << "dim mismatch at " << i << ": memref=" << md << " tensor=" << td; - } - } - return success(); -} - -[[maybe_unused]] static LogicalResult verifyMemrefTensorStore(Operation *op, Value dst, Value src) { - auto mr = dyn_cast(dst.getType()); - if (!mr) - return success(); // non-memref case: old tile IR allowed - auto rt = dyn_cast(src.getType()); - if (!rt) - return op->emitOpError("when dst is memref, src must be ranked tensor"); - - if (mr.getElementType() != rt.getElementType()) - return op->emitOpError() << "memref/tensor element type mismatch: memref=" - << mr.getElementType() << " tensor=" << rt.getElementType(); - - if (mr.getRank() != rt.getRank()) - return op->emitOpError() << "rank mismatch: memref rank=" << mr.getRank() - << " tensor rank=" << rt.getRank(); - - for (int64_t i = 0; i < mr.getRank(); ++i) { - int64_t md = mr.getDimSize(i); - int64_t td = rt.getDimSize(i); - if (md != ShapedType::kDynamic && td != ShapedType::kDynamic && md != td) - return op->emitOpError() << "dim mismatch at " << i << ": memref=" << md << " tensor=" << td; - } - return success(); -} - static std::optional getLocalAddressAlignmentBytes(Attribute memorySpace) { auto addrSpace = dyn_cast_or_null(memorySpace); @@ -3258,11 +3164,12 @@ LogicalResult TPrefetchAsyncOp::verify() { } LogicalResult mlir::pto::SetFFTsOp::verify() { - auto mr = llvm::dyn_cast(getFfts().getType()); - if (!mr) - return emitOpError("expects a memref operand"); + auto ptrTy = llvm::dyn_cast(getFfts().getType()); + if (!ptrTy) + return emitOpError("expects a !pto.ptr operand"); - if (!mr.getElementType().isInteger(64) && !mr.getElementType().isInteger(8)) + if (!ptrTy.getElementType().isInteger(64) && + !ptrTy.getElementType().isInteger(8)) return emitOpError("expects element type i64 (or i8)"); return mlir::success(); @@ -3662,8 +3569,6 @@ LogicalResult pto::TAbsOp::verify() { Type elemTy; if (auto tb = dyn_cast(srcTy)) elemTy = tb.getElementType(); - else if (auto mr = dyn_cast(srcTy)) - elemTy = mr.getElementType(); if (!(elemTy.isF16() || elemTy.isF32())) return emitOpError() << "expects element type to be f16 or f32"; @@ -3672,17 +3577,15 @@ LogicalResult pto::TAbsOp::verify() { // PTO.cpp static bool isPTOShapedLike(Type ty) { - return mlir::isa(ty); } static bool isTileLikeType(Type ty) { - return isa(ty); + return isa(ty); } static Type getElemTy(Type ty) { - if (auto mr = mlir::dyn_cast(ty)) return mr.getElementType(); if (auto tt = mlir::dyn_cast(ty)) return tt.getElementType(); if (auto tv = mlir::dyn_cast(ty)) return tv.getElementType(); if (auto tb = mlir::dyn_cast(ty)) return tb.getElementType(); @@ -3692,8 +3595,6 @@ static Type getElemTy(Type ty) { static SmallVector getShapeVec(Type ty) { SmallVector s; - if (auto mr = mlir::dyn_cast(ty)) - return SmallVector(mr.getShape().begin(), mr.getShape().end()); if (auto tt = mlir::dyn_cast(ty)) return SmallVector(tt.getShape().begin(), tt.getShape().end()); if (auto tv = mlir::dyn_cast(ty)) @@ -3906,12 +3807,6 @@ static std::optional getPTOMemorySpaceEnum(Type ty) { return as.getAddressSpace(); return std::nullopt; } - if (auto mr = dyn_cast(ty)) { - if (auto as = dyn_cast_or_null(mr.getMemorySpace())) - return as.getAddressSpace(); - if (!mr.getMemorySpace()) - return pto::AddressSpace::GM; - } return std::nullopt; } @@ -4261,7 +4156,7 @@ static LogicalResult verifyTileBufSameElemType(Operation *op, Type lhs, Type rhs StringRef rhsName) { if (!isTileLikeType(lhs) || !isTileLikeType(rhs)) return op->emitOpError() << "expects " << lhsName << " and " << rhsName - << " to be !pto.tile_buf or memref"; + << " to be !pto.tile_buf"; if (getElemTy(lhs) != getElemTy(rhs)) return op->emitOpError() << "expects " << lhsName << " and " << rhsName << " to have the same element type"; @@ -5155,7 +5050,7 @@ LogicalResult pto::TAddCOp::verify() { if (!isPTOShapedLike(t0) || !isPTOShapedLike(t1) || !isPTOShapedLike(t2) || !isPTOShapedLike(td)) - return emitOpError("expects src0/src1/src2/dst to be memref/tile_buf types"); + return emitOpError("expects src0/src1/src2/dst to be PTO shaped-like types"); auto s0 = getShapeVec(t0); auto s1 = getShapeVec(t1); @@ -6229,7 +6124,7 @@ LogicalResult mlir::pto::TDivOp::verify() { mlir::LogicalResult mlir::pto::TDivSOp::verify() { auto isTileLike = [](Type ty) -> bool { - return isa(ty); }; auto isScalarLike = [](Type ty) -> bool { @@ -7244,12 +7139,11 @@ mlir::LogicalResult mlir::pto::TFillPadOp::verify() { auto dstSpace = getPTOMemorySpaceEnum(getDst().getType()); if (!dstSpace || *dstSpace != pto::AddressSpace::MAT) return emitOpError("expects padValue attribute only for loc=mat tfillpad"); - if (auto dstTileTy = dyn_cast(getDst().getType())) { - if (dstTileTy.getPadValueI32() != static_cast(padValueAttr.getValue())) - return emitOpError("expects padValue attribute to match dst tile pad configuration"); - } else if (!isa(getDst().getType())) { - return emitOpError("expects dst to be tile_buf or memref when padValue is specified"); - } + auto dstTileTy = dyn_cast(getDst().getType()); + if (!dstTileTy) + return emitOpError("expects dst to be tile_buf when padValue is specified"); + if (dstTileTy.getPadValueI32() != static_cast(padValueAttr.getValue())) + return emitOpError("expects padValue attribute to match dst tile pad configuration"); } return success(); @@ -7801,20 +7695,20 @@ mlir::LogicalResult mlir::pto::TMovFPOp::verify() { } // 辅助函数：获取 Rank，支持 ShapedType 和 PTO TileTypes static int64_t getRankHelper(Type t) { - if (auto s = dyn_cast(t)) return s.getRank(); + if (auto s = dyn_cast(t)) return s.getRank(); if (auto tile = dyn_cast(t)) return tile.getRank(); if (auto view = dyn_cast(t)) return view.getRank(); return -1; } static LogicalResult verifyMatmulLike(Operation *op, Type aTy, Type bTy, Type dstTy, bool checkRank = true) { - // 1. 检查类型 (ShapedType 或 Tile 类型) - bool aValid = isa(aTy); - bool bValid = isa(bTy); - bool dValid = isa(dstTy); + // 1. 检查类型 (Tensor 或 Tile 类型) + bool aValid = isa(aTy); + bool bValid = isa(bTy); + bool dValid = isa(dstTy); if (!aValid || !bValid || !dValid) - return op->emitOpError("expects inputs/outputs to be shaped types or PTO tile types"); + return op->emitOpError("expects inputs/outputs to be tensors or PTO tile types"); if (checkRank) { int64_t aRank = getRankHelper(aTy); @@ -7837,12 +7731,8 @@ LogicalResult LoadScalarOp::verify() { Type elemTy; if (auto pty = dyn_cast(ptrTy)) { elemTy = pty.getElementType(); - } else if (auto memTy = dyn_cast(ptrTy)) { - elemTy = memTy.getElementType(); - if (!isGmAddressSpaceAttr(memTy.getMemorySpace())) - return emitOpError() << "scalar load only supports GM address space pointers"; } else { - return emitOpError("expects ptr to be !pto.ptr or memref type"); + return emitOpError("expects ptr to be !pto.ptr type"); } if (getValue().getType() != elemTy) @@ -7856,12 +7746,8 @@ LogicalResult StoreScalarOp::verify() { Type elemTy; if (auto pty = dyn_cast(ptrTy)) { elemTy = pty.getElementType(); - } else if (auto memTy = dyn_cast(ptrTy)) { - elemTy = memTy.getElementType(); - if (!isGmAddressSpaceAttr(memTy.getMemorySpace())) - return emitOpError() << "scalar store only supports GM address space pointers"; } else { - return emitOpError("expects ptr to be !pto.ptr or memref type"); + return emitOpError("expects ptr to be !pto.ptr type"); } if (getValue().getType() != elemTy) @@ -7878,10 +7764,6 @@ static bool isGmOrDefaultAddressSpace(pto::AddressSpace space) { static bool isGmOrDefaultCmoAddressType(Type type) { if (auto ptrTy = dyn_cast(type)) return isGmOrDefaultAddressSpace(ptrTy.getMemorySpace().getAddressSpace()); - if (auto memTy = dyn_cast(type)) { - auto spaceAttr = dyn_cast_or_null(memTy.getMemorySpace()); - return !spaceAttr || isGmOrDefaultAddressSpace(spaceAttr.getAddressSpace()); - } if (isa(type)) return true; return false; @@ -7935,7 +7817,7 @@ LogicalResult CmoCacheInvalidOp::verify() { if (Value addr = getAddr()) { if (!isGmOrDefaultCmoAddressType(addr.getType())) - return emitOpError("single_cache_line address expects a GM pointer, GM memref, or GM tensor view"); + return emitOpError("single_cache_line address expects a GM pointer or GM tensor view"); } return success(); @@ -8694,14 +8576,11 @@ LogicalResult TSetValOp::verify() { // ---- TGetValOp ---- LogicalResult TGetValOp::verify() { Type srcTy = getSrc().getType(); - if (!mlir::isa(srcTy)) - return emitOpError("expects src to be tile_buf or memref type"); + if (!mlir::isa(srcTy)) + return emitOpError("expects src to be tile_buf type"); // Memory space must be vec (Ascend does not support getval from MAT etc.). - Attribute memSpace = - isa(srcTy) - ? cast(srcTy).getMemorySpace() - : cast(srcTy).getMemorySpace(); + Attribute memSpace = cast(srcTy).getMemorySpace(); auto addrSpaceAttr = dyn_cast_or_null(memSpace); if (!addrSpaceAttr || addrSpaceAttr.getAddressSpace() != pto::AddressSpace::VEC) { @@ -9707,7 +9586,7 @@ static FailureOr verifyPTOShapedBinarySameElemAndShape(Operation *op, if (!isPTOShapedLike(src0Ty) || !isPTOShapedLike(src1Ty) || !isPTOShapedLike(dstTy)) return op->emitOpError( - "expects src0/src1/dst to be memref/tensor/tile_buf/tile_view types"), + "expects src0/src1/dst to be tensor/tile_buf/tile_view types"), failure(); Type e0 = getElemTy(src0Ty), e1 = getElemTy(src1Ty), ed = getElemTy(dstTy); if (!e0 || !e1 || !ed) @@ -10855,10 +10734,6 @@ static std::optional getElemBytes(Type elemTy) { return std::nullopt; } -[[maybe_unused]] static bool isTileBufOrMemref(Type ty) { - return mlir::isa(ty); -} - static bool isLocallyBoundTileSource(Value value) { if (!value || isa(value)) return false; @@ -12615,10 +12490,9 @@ mlir::LogicalResult mlir::pto::TPrintOp::verify() { } if (tmp) return emitOpError() << "expects tmp only when src is a tile_buf"; - if (mlir::dyn_cast(srcType) || - mlir::dyn_cast(srcType)) + if (mlir::dyn_cast(srcType)) return mlir::success(); - return emitOpError() << "expects tile_buf, memref, or partition_tensor_view for src"; + return emitOpError() << "expects tile_buf or partition_tensor_view for src"; } @@ -12626,11 +12500,11 @@ mlir::LogicalResult mlir::pto::TPrintOp::verify() { [[maybe_unused]] static LogicalResult verifyMatmulCommon(Operation *op, Value lhs, Value rhs, Value biasOpt, Type maybeDstElemTy, Type maybeResultElemTy) { - // ---- case A: tensor/memref (ShapedType) ---- - if (auto lhsTy = dyn_cast(lhs.getType())) { - auto rhsTy = dyn_cast(rhs.getType()); - if (!rhsTy || !lhsTy.hasRank() || !rhsTy.hasRank()) - return op->emitOpError("expects lhs and rhs to be ranked tensors or memrefs"); + // ---- case A: tensor ---- + if (auto lhsTy = dyn_cast(lhs.getType())) { + auto rhsTy = dyn_cast(rhs.getType()); + if (!rhsTy) + return op->emitOpError("expects lhs and rhs to be ranked tensors"); if (lhsTy.getElementType() != rhsTy.getElementType()) return op->emitOpError() @@ -12638,9 +12512,9 @@ mlir::LogicalResult mlir::pto::TPrintOp::verify() { << lhsTy.getElementType() << " rhs=" << rhsTy.getElementType(); if (biasOpt) { - auto biasTy = dyn_cast(biasOpt.getType()); - if (!biasTy || !biasTy.hasRank()) - return op->emitOpError("expects bias to be a ranked tensor or memref"); + auto biasTy = dyn_cast(biasOpt.getType()); + if (!biasTy) + return op->emitOpError("expects bias to be a ranked tensor"); if (biasTy.getElementType() != lhsTy.getElementType()) return op->emitOpError() << "expects bias to have the same element type as lhs and rhs, but got bias=" @@ -12664,7 +12538,7 @@ mlir::LogicalResult mlir::pto::TPrintOp::verify() { auto lhsTile = dyn_cast(lhs.getType()); auto rhsTile = dyn_cast(rhs.getType()); if (!lhsTile || !rhsTile) - return op->emitOpError("expects lhs and rhs to be ranked tensors, memrefs, or !pto.tile"); + return op->emitOpError("expects lhs and rhs to be ranked tensors or !pto.tile"); if (lhsTile.getElementType() != rhsTile.getElementType()) return op->emitOpError() << "expects lhs and rhs tiles to have the same element type, but got lhs=" @@ -12804,9 +12678,9 @@ LogicalResult mlir::pto::TGemvAccOp::verify() { if (operands.size() < 2) return RankedTensorType(); - auto lhsTy = dyn_cast(operands[0].getType()); - auto rhsTy = dyn_cast(operands[1].getType()); - if (!lhsTy || !rhsTy || !lhsTy.hasRank() || !rhsTy.hasRank()) + auto lhsTy = dyn_cast(operands[0].getType()); + auto rhsTy = dyn_cast(operands[1].getType()); + if (!lhsTy || !rhsTy) return RankedTensorType(); Type elemTy = lhsTy.getElementType(); @@ -12814,10 +12688,6 @@ LogicalResult mlir::pto::TGemvAccOp::verify() { if (operands.size() >= 3) { if (auto biasRT = dyn_cast(operands[2].getType())) return RankedTensorType::get(biasRT.getShape(), elemTy); - if (auto biasMR = dyn_cast(operands[2].getType())) { - if (biasMR.hasStaticShape()) - return RankedTensorType::get(biasMR.getShape(), elemTy); - } } if (lhsTy.getRank() >= 2 && rhsTy.getRank() >= 2) { @@ -13459,151 +13329,6 @@ static LogicalResult computeInnerShape(TileBufConfigAttr cfg, Type elemTy, return failure(); } -static LogicalResult -computeExpectedTileBufMemrefStrides(TileBufType tileTy, - SmallVectorImpl &expectedStrides) { - if (tileTy.getRank() != 2) - return failure(); - - ArrayRef shape = tileTy.getShape(); - if (shape.size() != 2) - return failure(); - if (shape[0] == ShapedType::kDynamic || shape[1] == ShapedType::kDynamic) - return failure(); - - auto cfg = tileTy.getConfigAttr(); - if (!cfg) - cfg = TileBufConfigAttr::getDefault(tileTy.getContext()); - - int64_t innerRows = 1, innerCols = 1; - bool boxed = false; - int32_t bl = 0, sl = 0; - if (failed(computeInnerShape(cfg, tileTy.getElementType(), innerRows, innerCols, - boxed, bl, sl))) - return failure(); - - expectedStrides.clear(); - if (!boxed) { - if (bl == 1) { - expectedStrides.push_back(1); - expectedStrides.push_back(shape[0]); - } else { - expectedStrides.push_back(shape[1]); - expectedStrides.push_back(1); - } - return success(); - } - - if (bl == 1) { - if (sl != 1) - return failure(); - expectedStrides.push_back(innerCols); - expectedStrides.push_back(shape[0]); - return success(); - } - - expectedStrides.push_back(shape[1]); - expectedStrides.push_back(innerRows); - return success(); -} - -mlir::LogicalResult mlir::pto::SimdTileToMemrefOp::verify() { - auto memTy = dyn_cast(getDst().getType()); - if (!memTy) - return emitOpError("expects result to be memref"); - - Type srcTy = getSrc().getType(); - if (auto tileTy = dyn_cast(srcTy)) { - if (memTy.getElementType() != tileTy.getElementType()) - return emitOpError( - "expects memref element type to match tile_buf element type"); - - if (memTy.getMemorySpace() != tileTy.getMemorySpace()) - return emitOpError( - "expects memref memory space to match tile_buf memory space"); - - if (memTy.getRank() != tileTy.getRank()) - return emitOpError("expects memref rank to match tile_buf rank"); - - ArrayRef tileShape = tileTy.getShape(); - ArrayRef validShape = tileTy.getValidShape(); - ArrayRef memShape = memTy.getShape(); - if (tileShape.size() != memShape.size()) - return emitOpError( - "expects memref shape rank to match tile_buf shape rank"); - - if (validShape.size() != memShape.size()) - return emitOpError( - "expects tile_buf valid shape rank to match memref shape rank"); - - for (unsigned i = 0; i < validShape.size(); ++i) { - int64_t expect = validShape[i]; - if (expect < 0) { - if (memShape[i] >= 0 && memShape[i] != tileShape[i]) { - return emitOpError() - << "expects memref dim " << i - << " to be dynamic or match physical tile dim " << tileShape[i] - << " because tile_buf valid dim is ?"; - } - continue; - } - - if (memShape[i] != expect) { - return emitOpError() << "expects memref dim " << i - << " to match tile_buf valid dim; got " - << memShape[i] << ", expected " << expect; - } - } - - SmallVector expectedStrides; - if (failed(computeExpectedTileBufMemrefStrides(tileTy, expectedStrides))) - return emitOpError("cannot infer expected strides from tile_buf layout"); - - SmallVector memStrides; - int64_t memOffset = ShapedType::kDynamic; - if (failed(getPTOMemRefStridesAndOffset(memTy, memStrides, memOffset))) - return emitOpError("expects memref to use strided layout"); - if (memOffset != 0) - return emitOpError("expects memref offset to be 0"); - if (memStrides.size() != expectedStrides.size()) - return emitOpError("expects memref stride rank to match tile_buf rank"); - for (unsigned i = 0; i < expectedStrides.size(); ++i) { - if (memStrides[i] != expectedStrides[i]) { - return emitOpError() - << "expects memref strides to match tile_buf layout; got " - << memStrides[i] << " at dim " << i << ", expected " - << expectedStrides[i]; - } - } - return success(); - } - - auto srcMemTy = dyn_cast(srcTy); - if (!srcMemTy) - return emitOpError("expects src to be !pto.tile_buf or memref"); - - if (srcMemTy.getElementType() != memTy.getElementType()) - return emitOpError("expects src/result memref element types to match"); - - if (srcMemTy.getMemorySpace() != memTy.getMemorySpace()) - return emitOpError("expects src/result memref memory spaces to match"); - - if (srcMemTy.getRank() != memTy.getRank()) - return emitOpError("expects src/result memref ranks to match"); - - ArrayRef srcShape = srcMemTy.getShape(); - ArrayRef dstShape = memTy.getShape(); - for (unsigned i = 0; i < srcShape.size(); ++i) { - if (srcShape[i] >= 0 && dstShape[i] >= 0 && srcShape[i] != dstShape[i]) { - return emitOpError() - << "expects compatible src/result memref shapes; dim " << i - << " mismatches (" << srcShape[i] << " vs " << dstShape[i] << ")"; - } - } - - return success(); -} - mlir::LogicalResult mlir::pto::SubViewOp::verify() { auto srcTy = llvm::dyn_cast(getSource().getType()); auto dstTy = llvm::dyn_cast(getResult().getType()); @@ -13750,17 +13475,6 @@ using namespace mlir::pto; // Helper Functions // ============================================================================= -[[maybe_unused]] static AddressSpace getAddressSpace(Value val) { - auto type = llvm::dyn_cast(val.getType()); - if (!type) return AddressSpace::Zero; // Default - - // 假设你的 AddressSpaceAttr 存储在 MemRef 的 memorySpace 中 - // 需要根据你的 getPTOAddressSpaceAttr 实现来调整 - auto attr = llvm::dyn_cast_or_null(type.getMemorySpace()); - if (attr) return attr.getAddressSpace(); - return AddressSpace::Zero; -} - // ============================================================================= // Side Effects Implementation // ============================================================================= @@ -15899,11 +15613,6 @@ static bool getTensorLikeElementAndShape(Type ty, Type &elementType, shape = tvTy.getShape(); return true; } - if (auto memrefTy = dyn_cast(ty)) { - elementType = memrefTy.getElementType(); - shape = memrefTy.getShape(); - return true; - } return false; } @@ -15932,7 +15641,7 @@ static LogicalResult verifyTensorEntryMatchesInternalPipeInit(Operation *op, slotElementType, slotShape)) { return op->emitOpError() << "expects !pto.tensor_view pipe entry to use " - "pto.initialize_l2g2l_pipe gm_addr with tensor/memref slot type"; + "pto.initialize_l2g2l_pipe gm_addr with tensor_view slot type"; } if (slotElementType != entryViewTy.getElementType()) { @@ -16952,15 +16661,12 @@ LogicalResult SetQuantVectorOp::verify() { Operation *initOp = *initOr; pto::AccPushEpilogueAttr accPushEpilogue; - bool isLoweredProducerPipe = false; if (auto aicInit = dyn_cast(initOp)) { accPushEpilogue = aicInit.getAccPushEpilogueAttr(); } else if (auto l2lInit = dyn_cast(initOp)) { accPushEpilogue = l2lInit.getAccPushEpilogueAttr(); - isLoweredProducerPipe = true; } else if (auto l2g2lInit = dyn_cast(initOp)) { accPushEpilogue = l2g2lInit.getAccPushEpilogueAttr(); - isLoweredProducerPipe = true; } else { return emitOpError() << "expects 'id' = " << getId() @@ -16984,10 +16690,7 @@ LogicalResult SetQuantVectorOp::verify() { } Type scalingTy = getScalingTile().getType(); - if (!isa(scalingTy)) { - return emitOpError("expects 'scaling_tile' to be a tile type"); - } - if (!isa(scalingTy) && !isLoweredProducerPipe) { + if (!isa(scalingTy)) { return emitOpError("expects 'scaling_tile' to be a tile type"); } auto scalingSpace = getPTOMemorySpaceEnum(scalingTy); diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index bf4aaeb14e..d258b9d83a 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -4905,42 +4905,20 @@ LogicalResult TensorViewAddrOp::verify() { Type dstType = getDst().getType(); Type elementType; - int64_t expectedRank = -1; auto gmSpace = pto::AddressSpaceAttr::get(getContext(), pto::AddressSpace::GM); if (auto tvType = dyn_cast(srcType)) { elementType = tvType.getElementType(); - expectedRank = tvType.getRank(); } else if (auto partType = dyn_cast(srcType)) { elementType = partType.getElementType(); - expectedRank = partType.getRank(); - } else if (auto memrefType = dyn_cast(srcType)) { - elementType = memrefType.getElementType(); - expectedRank = memrefType.getRank(); - auto srcSpace = dyn_cast_or_null(memrefType.getMemorySpace()); - if (srcSpace && srcSpace != gmSpace) - return emitOpError("memref source must stay in gm memory space"); } else { return emitOpError( - "source must be a tensor_view, partition_tensor_view, or memref"); - } - - if (auto dstMemRefType = dyn_cast(dstType)) { - if (dstMemRefType.getElementType() != elementType) - return emitOpError( - "memref result element type must match source element type"); - if (dstMemRefType.getRank() != expectedRank) - return emitOpError("memref result rank must match source rank"); - auto dstSpace = - dyn_cast_or_null(dstMemRefType.getMemorySpace()); - if (dstSpace && dstSpace != gmSpace) - return emitOpError("memref result must stay in gm memory space"); - return success(); + "source must be a tensor_view or partition_tensor_view"); } auto dstPtrType = dyn_cast(dstType); if (!dstPtrType) - return emitOpError("result must be a memref or !pto.ptr<...>"); + return emitOpError("result must be !pto.ptr<...>"); if (dstPtrType.getElementType() != elementType) return emitOpError( "pointer result element type must match source element type"); @@ -4953,40 +4931,19 @@ LogicalResult TileBufAddrOp::verify() { Type dstType = getDst().getType(); Type elementType; Attribute srcMemorySpace; - int64_t srcRank = 0; if (auto srcTileType = dyn_cast(getSrc().getType())) { elementType = srcTileType.getElementType(); srcMemorySpace = srcTileType.getMemorySpace(); - srcRank = static_cast(srcTileType.getShape().size()); - } else if (auto srcMemRefType = dyn_cast(getSrc().getType())) { - // Hand-written legacy IR may use tile_buf_addr directly on an addressed - // memref before the shared materialization bridge restores tile handles. - elementType = srcMemRefType.getElementType(); - srcMemorySpace = srcMemRefType.getMemorySpace(); - srcRank = srcMemRefType.getRank(); } else { - return emitOpError("source must be a !pto.tile_buf<...> or memref"); + return emitOpError("source must be a !pto.tile_buf<...>"); } auto srcSpace = dyn_cast_or_null(srcMemorySpace); - if (auto dstMemRefType = dyn_cast(dstType)) { - if (dstMemRefType.getElementType() != elementType) - return emitOpError( - "memref result element type must match tile element type"); - if (dstMemRefType.getRank() != srcRank) - return emitOpError("memref result rank must match tile rank"); - auto dstSpace = - dyn_cast_or_null(dstMemRefType.getMemorySpace()); - if (srcSpace && dstSpace && srcSpace != dstSpace) - return emitOpError("memref result must stay within the tile memory space"); - return success(); - } - auto dstPtrType = dyn_cast(dstType); if (!dstPtrType) - return emitOpError("result must be a memref or !pto.ptr<...>"); + return emitOpError("result must be !pto.ptr<...>"); if (dstPtrType.getElementType() != elementType) return emitOpError( "pointer result element type must match tile element type"); diff --git a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp index bbe2e11b2f..4ba0f8445e 100644 --- a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp +++ b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp @@ -12,12 +12,12 @@ // structured-view intrinsics that reference template parameters: // // tile_buf family: -// - pto.tile_buf_addr → extract memref address from tile_buf +// - pto.tile_buf_addr → extract pointer address from tile_buf // - pto.tile_valid_rows → extract valid row count // - pto.tile_valid_cols → extract valid column count // // tensor_view family: -// - pto.tensor_view_addr → extract memref/ptr from tensor_view +// - pto.tensor_view_addr → extract pointer from tensor_view // - pto.get_tensor_view_dim → extract dimension size // - pto.get_tensor_view_stride → extract dimension stride // @@ -708,9 +708,7 @@ struct FoldTileBufIntrinsicsPass // Fold pto.tile_buf_addr by recovering the active materialized tile // handle contract: - // - pto.alloc_tile → rebuild a memref from the explicit addr - // When the requested result type is already !pto.ptr<...>, cast from the - // recovered memref instead of leaving tile_buf_addr in the IR. + // - pto.alloc_tile → cast the explicit addr to the requested pointer. for (auto addrOp : addrOps) { // An SCF result/iter_arg is already a runtime-selected tile handle. // Keep tile_buf_addr attached to that handle; VPTO pointer @@ -730,13 +728,10 @@ struct FoldTileBufIntrinsicsPass return signalPassFailure(); } - if (isa(addrOp.getDst().getType())) - continue; - auto resultPtrType = dyn_cast(addrOp.getDst().getType()); if (!resultPtrType) { addrOp.emitError("FoldTileBufIntrinsics: tile_buf_addr result must " - "be memref or !pto.ptr"); + "be !pto.ptr"); return signalPassFailure(); } @@ -925,18 +920,8 @@ struct FoldTileBufIntrinsicsPass return signalPassFailure(); if (!resultPtrType) { - if (auto resultMemrefType = - dyn_cast(addrOp.getDst().getType())) { - Value base = chain->baseMemref; - if (base.getType() != resultMemrefType) - addrOp.getDst().setType(cast(base.getType())); - addrOp.getDst().replaceAllUsesWith(base); - addrOp.erase(); - continue; - } addrOp.emitError( - "FoldTileBufIntrinsics: tensor_view_addr result must be memref or " - "!pto.ptr"); + "FoldTileBufIntrinsics: tensor_view_addr result must be !pto.ptr"); return signalPassFailure(); } diff --git a/lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp b/lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp index 42f6309020..dde973d6c7 100644 --- a/lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp +++ b/lib/PTO/Transforms/PTOLowerToOpLibCalls.cpp @@ -15,192 +15,9 @@ using namespace mlir; -namespace { - -static int64_t getElemBytes(Type elemTy) { - if (auto intTy = dyn_cast(elemTy)) - return (intTy.getWidth() + 7) / 8; - if (auto floatTy = dyn_cast(elemTy)) - return (floatTy.getWidth() + 7) / 8; - return -1; -} - -static bool readBLayoutI32(Attribute attr, int32_t &out) { - if (auto intAttr = dyn_cast(attr)) { - out = static_cast(intAttr.getInt()); - return true; - } - return false; -} - -static bool readSLayoutI32(Attribute attr, int32_t &out) { - if (auto intAttr = dyn_cast(attr)) { - out = static_cast(intAttr.getInt()); - return true; - } - return false; -} - -static FailureOr inferSimdBridgeMemRefType(pto::TileBufType tileTy, - MLIRContext *ctx) { - if (tileTy.getRank() != 2) - return failure(); - - ArrayRef physicalShape = tileTy.getShape(); - if (physicalShape.size() != 2) - return failure(); - if (physicalShape[0] == ShapedType::kDynamic || - physicalShape[1] == ShapedType::kDynamic) - return failure(); - - SmallVector memShape(physicalShape.begin(), physicalShape.end()); - ArrayRef validShape = tileTy.getValidShape(); - if (validShape.size() == memShape.size()) { - for (unsigned i = 0; i < validShape.size(); ++i) - memShape[i] = validShape[i] < 0 ? physicalShape[i] : validShape[i]; - } - - auto cfg = tileTy.getConfigAttr(); - if (!cfg) - cfg = pto::TileBufConfigAttr::getDefault(ctx); - - int32_t bl = 0; - int32_t sl = 0; - int32_t fr = 512; - (void)readBLayoutI32(cfg.getBLayout(), bl); - (void)readSLayoutI32(cfg.getSLayout(), sl); - if (auto attr = dyn_cast(cfg.getSFractalSize())) - fr = static_cast(attr.getInt()); - - int64_t innerRows = 1; - int64_t innerCols = 1; - if (sl != 0) { - int64_t elemBytes = getElemBytes(tileTy.getElementType()); - if (elemBytes <= 0) - return failure(); - if (fr == 1024) { - innerRows = 16; - innerCols = 16; - } else if (fr == 32) { - innerRows = 16; - innerCols = 2; - } else if (fr == 512) { - if (sl == 1) { - innerRows = 16; - innerCols = 32 / elemBytes; - } else if (sl == 2) { - innerRows = 32 / elemBytes; - innerCols = 16; - } else { - return failure(); - } - } else { - return failure(); - } - } - - SmallVector strides; - if (sl == 0) { - if (bl == 1) { - strides.push_back(1); - strides.push_back(physicalShape[0]); - } else { - strides.push_back(physicalShape[1]); - strides.push_back(1); - } - } else if (bl == 1) { - if (sl != 1) - return failure(); - strides.push_back(innerCols); - strides.push_back(physicalShape[0]); - } else { - strides.push_back(physicalShape[1]); - strides.push_back(innerRows); - } - - auto layout = StridedLayoutAttr::get(ctx, /*offset=*/0, strides); - return MemRefType::get(memShape, tileTy.getElementType(), layout, - tileTy.getMemorySpace()); -} - -static bool areIntegerCarrierTypesCompatible(Type lhs, Type rhs) { - auto lhsInt = dyn_cast(lhs); - auto rhsInt = dyn_cast(rhs); - if (!lhsInt || !rhsInt) - return false; - return lhsInt.getWidth() == rhsInt.getWidth(); -} - -static bool canRemapSimdBridgeViaCarrierCast(MemRefType actualTy, - MemRefType templateTy) { - if (actualTy.getRank() != templateTy.getRank()) - return false; - if (actualTy.getMemorySpace() != templateTy.getMemorySpace()) - return false; - return areIntegerCarrierTypesCompatible(actualTy.getElementType(), - templateTy.getElementType()); -} - -static MemRefType remapMemRefToTemplateCarrier(MemRefType actualTy, - MemRefType templateTy) { - return MemRefType::get(actualTy.getShape(), templateTy.getElementType(), - actualTy.getLayout(), actualTy.getMemorySpace()); -} - -} // namespace - FailureOr mlir::pto::tryCloneOpLibInlineBridgeOp(OpBuilder &builder, Operation &op, IRMapping &mapping) { - if (auto bridge = dyn_cast(&op)) { - Value mappedSrc = mapping.lookupOrNull(bridge.getSrc()); - if (!mappedSrc) - return failure(); - - auto templateMemTy = dyn_cast(bridge.getDst().getType()); - if (auto mappedTileTy = dyn_cast(mappedSrc.getType())) { - FailureOr inferredTyOr = - inferSimdBridgeMemRefType(mappedTileTy, builder.getContext()); - if (failed(inferredTyOr)) - return failure(); - - auto inferredTy = *inferredTyOr; - auto newBridge = builder.create( - bridge.getLoc(), inferredTy, mappedSrc); - if (templateMemTy && inferredTy != templateMemTy && - canRemapSimdBridgeViaCarrierCast(inferredTy, templateMemTy)) { - auto carrierTy = remapMemRefToTemplateCarrier(inferredTy, templateMemTy); - auto cast = builder.create( - bridge.getLoc(), TypeRange{carrierTy}, ValueRange{newBridge.getDst()}); - mapping.map(bridge.getDst(), cast.getResult(0)); - } else { - mapping.map(bridge.getDst(), newBridge.getDst()); - } - return true; - } - - auto mappedMemTy = dyn_cast(mappedSrc.getType()); - auto dstMemTy = templateMemTy; - if (!mappedMemTy || !dstMemTy) - return failure(); - if (mappedMemTy.getRank() != dstMemTy.getRank()) - return failure(); - - auto newBridge = builder.create( - bridge.getLoc(), mappedMemTy, mappedSrc); - if (mappedMemTy.getElementType() == dstMemTy.getElementType()) { - mapping.map(bridge.getDst(), newBridge.getDst()); - return true; - } - if (!canRemapSimdBridgeViaCarrierCast(mappedMemTy, dstMemTy)) - return failure(); - auto carrierTy = remapMemRefToTemplateCarrier(mappedMemTy, dstMemTy); - auto cast = builder.create( - bridge.getLoc(), TypeRange{carrierTy}, ValueRange{newBridge.getDst()}); - mapping.map(bridge.getDst(), cast.getResult(0)); - return true; - } - if (auto cast = dyn_cast(&op)) { if (cast->getNumOperands() != 1 || cast->getNumResults() != 1) return failure(); diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 30a77f71e7..409a96d488 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -262,7 +262,7 @@ def vlds(src_ptr, offset=None, result_vreg_type=None, *, dist=None, post_update= post_mode = _normalize_post_update_mode(post_update, context="vlds(..., post_update=...)") if isinstance(src_ptr, TileSliceValue): if offset is not None or result_vreg_type is not None: - raise TypeError("vlds(tile[row, col:]) infers its memref slice and vreg type; do not pass offset/result_vreg_type") + raise TypeError("vlds(tile[row, col:]) infers its pointer slice and vreg type; do not pass offset/result_vreg_type") if post_mode != "NO_POST_UPDATE": raise TypeError("vlds(tile[...], post_update=...) only supports post_update=PostUpdate.OFF; use the pointer form for stateful loads") kwargs = {} @@ -272,13 +272,14 @@ def vlds(src_ptr, offset=None, result_vreg_type=None, *, dist=None, post_update= allowed=_VLOAD_DIST_TOKENS, context="vlds(..., dist)", ) - raw_source = unwrap_surface_value(src_ptr) + source, source_offset = _tile_slice_address(src_ptr) + raw_source = unwrap_surface_value(source) return wrap_surface_value( _pto.VldsOp( _infer_vreg_type_from_tile_slice(src_ptr), None, raw_source, - _index_zero(), + source_offset, **kwargs, ).result ) @@ -386,11 +387,12 @@ def vldsx2(source, offset_or_dist, dist=None): if dist is not None: raise TypeError("vldsx2(tile[row, col:], dist) does not accept a separate offset argument") result_type = _infer_vreg_type_from_tile_slice(source) + source, source_offset = _tile_slice_address(source) op = _pto.Vldsx2Op( result_type, result_type, unwrap_surface_value(source), - _index_zero(), + source_offset, _normalize_dist_token( offset_or_dist, allowed=_DEINTERLEAVE_DIST_TOKENS, @@ -847,12 +849,13 @@ def vsts(val, dst_ptr, offset, mask=None, *, dist=None, post_update="OFF"): allowed=_VSTORE_DIST_TOKENS, context="vsts(..., dist)", ) - raw_destination = unwrap_surface_value(dst_ptr) + destination, destination_offset = _tile_slice_address(dst_ptr) + raw_destination = unwrap_surface_value(destination) _pto.VstsOp( None, unwrap_surface_value(val), raw_destination, - _index_zero(), + destination_offset, unwrap_surface_value(offset), **kwargs, ) @@ -904,11 +907,12 @@ def vstsx2(low, high, dst_ptr, offset_or_dist, dist_or_mask=None, mask=None): if isinstance(dst_ptr, TileSliceValue): if mask is not None: raise TypeError("vstsx2(low, high, tile[row, col:], dist, mask) does not accept a separate offset argument") + destination, destination_offset = _tile_slice_address(dst_ptr) _pto.Vstsx2Op( unwrap_surface_value(low), unwrap_surface_value(high), - unwrap_surface_value(dst_ptr), - _index_zero(), + unwrap_surface_value(destination), + destination_offset, _normalize_dist_token( offset_or_dist, allowed=_INTERLEAVE_DIST_TOKENS, @@ -995,6 +999,8 @@ def vsldb(source, block_stride, repeat_stride, mask): if isinstance(source, TileSliceValue) else _infer_vreg_type_from_address_source(source) ) + if isinstance(source, TileSliceValue): + source = _tile_slice_ptr(source) return wrap_surface_value( _pto.VsldbOp( result_type, @@ -3803,9 +3809,14 @@ def _tile_slice_ptr(tile_slice: TileSliceValue): return addptr(base_ptr, _coerce_index(linear_offset, context="tile slice pointer lowering")) +def _tile_slice_address(tile_slice: TileSliceValue): + base_ptr = emit_as_ptr(tile_slice.tile) + linear_offset = _tile_slice_linear_offset(tile_slice) + return base_ptr, _coerce_index(linear_offset, context="tile slice pointer lowering") + + def _infer_vreg_type_from_tile_slice(tile_slice: TileSliceValue): - memref_type = MemRefType(tile_slice.type) - elem_type = memref_type.element_type + elem_type = infer_tile_element_type(tile_slice.tile) lanes = _elements_per_vreg(elem_type) return _resolve(vreg_type(lanes, elem_type)) diff --git a/ptodsl/ptodsl/_surface_values.py b/ptodsl/ptodsl/_surface_values.py index 94374333f3..5b07b1effd 100644 --- a/ptodsl/ptodsl/_surface_values.py +++ b/ptodsl/ptodsl/_surface_values.py @@ -23,9 +23,8 @@ from ._types import _normalize_address_space, _resolve, ptr from mlir.dialects import arith -from mlir.dialects import memref from mlir.dialects import pto as _pto -from mlir.ir import IndexType, IntegerAttr, IntegerType, MemRefType, ShapedType, StridedLayoutAttr, Type, VectorType +from mlir.ir import IndexType, IntegerAttr, IntegerType, MemRefType, ShapedType, Type, VectorType def _validate_surface_value_access(value): @@ -347,7 +346,7 @@ def __iter__(self): class AddressValue(_SurfaceValue): - """Author-facing address view backed by either a PTO ptr or a memref.""" + """Author-facing address view backed by a PTO ptr.""" def __add__(self, offset): return AddressOffsetValue(self, offset) @@ -422,7 +421,7 @@ def __iter__(self): class TileSliceValue(_SurfaceValue): - """Author-facing memref view produced by `tile[row, col:]` style indexing.""" + """Author-facing tile slice descriptor produced by `tile[row, col:]` indexing.""" def __init__(self, value, *, tile: "TileValue", offsets, shape): super().__init__(value) @@ -962,55 +961,7 @@ def _materialize_tile_slice(tile: TileValue, key): def _build_tile_slice_view(tile: TileValue, *, raw_offsets, shape): - base_memref = _emit_tile_memref(tile) - base_type = MemRefType(base_memref.type) - rank = len(base_type.shape) - offset_operands, static_offsets = _split_dynamic_index_operands(raw_offsets) - shape_operands, static_shape = _split_dynamic_index_operands(shape) - base_strides, base_offset = base_type.get_strides_and_offset() - if rank == 1: - slice_type = _make_strided_memref_type( - [_static_extent_if_known(shape[0])], - base_type.element_type, - [base_strides[0]], - base_type.memory_space, - offset=_compose_static_subview_offset(base_offset, base_strides, raw_offsets), - ) - slice_value = memref.SubViewOp( - slice_type, - base_memref, - offset_operands, - shape_operands, - [], - static_offsets, - static_shape, - [1], - ).result - return TileSliceValue(slice_value, tile=tile, offsets=tuple(raw_offsets), shape=shape) - - slice_type = _make_strided_memref_type( - [_static_extent_if_known(shape[0])], - base_type.element_type, - [base_strides[1]], - base_type.memory_space, - offset=_compose_static_subview_offset(base_offset, base_strides, raw_offsets), - ) - slice_value = memref.SubViewOp( - slice_type, - base_memref, - offset_operands, - shape_operands, - [], - static_offsets, - [1, static_shape[0]], - [1, 1], - ).result - return TileSliceValue(slice_value, tile=tile, offsets=tuple(raw_offsets), shape=shape) - - -def _emit_tile_memref(tile: TileValue): - memref_type = infer_memref_type_from_surface_value(tile) - return _pto.TileBufAddrOp(memref_type, tile.value).result + return TileSliceValue(tile.value, tile=tile, offsets=tuple(raw_offsets), shape=shape) def _dynamic_extent(static_dim, start): @@ -1019,57 +970,6 @@ def _dynamic_extent(static_dim, start): return arith.SubIOp(_index_const(static_dim), _coerce_index_value(start)).result -def _static_extent_if_known(extent): - return extent if _is_python_index_literal(extent) else ShapedType.get_dynamic_size() - - -def _static_index_attr(value): - return value if _is_python_index_literal(value) else ShapedType.get_dynamic_size() - - -def _split_dynamic_index_operands(values): - operands = [] - static_attrs = [] - for value in values: - if _is_python_index_literal(value): - static_attrs.append(value) - else: - operands.append(_coerce_index_value(value)) - static_attrs.append(ShapedType.get_dynamic_size()) - return operands, static_attrs - - -def _make_strided_memref_type_with_offset(shape, element_type, strides, memory_space, *, offset): - return MemRefType.get( - list(shape), - element_type, - StridedLayoutAttr.get(offset, list(strides)), - memory_space, - ) - - -def _make_strided_memref_type(shape, element_type, strides, memory_space, *, offset=ShapedType.get_dynamic_size()): - return _make_strided_memref_type_with_offset( - shape, - element_type, - strides, - memory_space, - offset=offset, - ) - - -def _compose_static_subview_offset(base_offset, base_strides, raw_offsets): - if base_offset == ShapedType.get_dynamic_size(): - return ShapedType.get_dynamic_size() - - linear_offset = base_offset - for stride, authored_offset in zip(base_strides, raw_offsets): - if not _is_python_index_literal(authored_offset): - return ShapedType.get_dynamic_size() - linear_offset += stride * authored_offset - return linear_offset - - def _mul_index(lhs, rhs): lhs = _normalize_index(lhs) rhs = _normalize_index(rhs) diff --git a/ptodsl/ptodsl/tilelib/_render_runtime.py b/ptodsl/ptodsl/tilelib/_render_runtime.py index fc9bd8aa8e..dbc6d1e399 100644 --- a/ptodsl/ptodsl/tilelib/_render_runtime.py +++ b/ptodsl/ptodsl/tilelib/_render_runtime.py @@ -8,9 +8,9 @@ """ TileLib render runtime. -Traces a tilelang-style template body into a standalone ``func.func`` whose MLIR is on -par with the legacy tilelang-dsl render (``tile_buf_addr`` -> memref, ``memref.subview``, -``pto.vlds/vadd/vsts``, dynamic ``pto.tile_valid_rows/cols``). +Traces a tilelang-style template body into a standalone ``func.func`` whose MLIR +uses tile-native handles plus pointer-style ``pto.tile_buf_addr`` / +``pto.vlds/vadd/vsts`` and dynamic ``pto.tile_valid_rows/cols``. Control flow is handled by the engine's AST rewrite (``rewrite_jit_function``): plain ``for x in range(...)`` in the template body is rewritten at trace time to @@ -54,8 +54,7 @@ class _TemplateTile(TileValue): static ``v_row/v_col`` carried in the tile_buf type). Metadata (shape/dtype/memory_space) is supplied from the ``TileSpec`` because a raw - entry-block ``tile_buf`` type is not introspectable by ``parse_tile_type_metadata``; - supplying it explicitly takes the fast path in ``infer_memref_type_from_surface_value``. + entry-block ``tile_buf`` type is not introspectable by ``parse_tile_type_metadata``. """ def __init__(self, value, spec: TileSpec): diff --git a/ptodsl/ptodsl/tilelib/metadata.py b/ptodsl/ptodsl/tilelib/metadata.py index 9ebaef8877..b3e8b32d93 100644 --- a/ptodsl/ptodsl/tilelib/metadata.py +++ b/ptodsl/ptodsl/tilelib/metadata.py @@ -33,6 +33,7 @@ si32 as _si32, si64 as _si64, tile_buf_type as _tile_buf_type, + tensor_view_type_from_dims as _tensor_view_type_from_dims, ui8 as _ui8, ui16 as _ui16, ui32 as _ui32, @@ -197,7 +198,7 @@ def mlir_type(self): @dataclass(frozen=True) class ViewSpec: - """Concrete specialization of one view/memref TileOp operand.""" + """Concrete specialization of one tensor-view TileOp operand.""" shape: tuple dtype: ScalarType @@ -206,11 +207,8 @@ class ViewSpec: layout: str | None = None def mlir_type(self): - dims = "x".join("?" if dim is None else str(dim) for dim in self.shape) - addr_space = _memref_address_space_token(self.memory_space) - elem = _resolve(scalar_descriptor(self.dtype)) - return Type.parse( - f"memref<{dims}x{elem}, #pto.address_space<{addr_space}>>" + return _resolve( + _tensor_view_type_from_dims(self.shape, scalar_descriptor(self.dtype)) ) @@ -227,21 +225,6 @@ def mlir_type(self): return Type.parse(f"vector<{dims}x{elem}>") -def _memref_address_space_token(value: str) -> str: - aliases = { - "ub": "vec", - "vec": "vec", - "gm": "gm", - "mat": "mat", - "left": "left", - "right": "right", - "acc": "acc", - "bias": "bias", - "scaling": "scaling", - } - return aliases.get(str(value), str(value)) - - @dataclass(frozen=True) class TemplateMetadata: """Hard constraints + selection hints for one registered template version.""" diff --git a/ptodsl/tests/fixtures/tadd_a5_8x64_f32.golden.mlir b/ptodsl/tests/fixtures/tadd_a5_8x64_f32.golden.mlir index 61a8e57c24..21aa1ccb63 100644 --- a/ptodsl/tests/fixtures/tadd_a5_8x64_f32.golden.mlir +++ b/ptodsl/tests/fixtures/tadd_a5_8x64_f32.golden.mlir @@ -11,25 +11,21 @@ %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c64 = arith.constant 64 : index - %tmp_0 = pto.tile_buf_addr %arg0 : !pto.tile_buf -> memref<8x64xf32, #pto.address_space> - %tmp_1 = pto.tile_buf_addr %arg1 : !pto.tile_buf -> memref<8x64xf32, #pto.address_space> - %tmp_2 = pto.tile_buf_addr %arg2 : !pto.tile_buf -> memref<8x64xf32, #pto.address_space> + %tmp_0 = pto.tile_buf_addr %arg0 : !pto.tile_buf -> !pto.ptr + %tmp_1 = pto.tile_buf_addr %arg1 : !pto.tile_buf -> !pto.ptr + %tmp_2 = pto.tile_buf_addr %arg2 : !pto.tile_buf -> !pto.ptr %valid_rows_1 = pto.tile_valid_rows %arg2 : !pto.tile_buf -> index %valid_cols_2 = pto.tile_valid_cols %arg2 : !pto.tile_buf -> index scf.for %row_3 = %c0 to %valid_rows_1 step %c1 { %tmp_3 = arith.index_cast %valid_cols_2 : index to i32 %remained_11:1 = scf.for %col_5 = %c0 to %valid_cols_2 step %c64 iter_args(%remained_iter_0 = %tmp_3) -> (i32) { %mask_6, %remained_7 = pto.plt_b32 %remained_iter_0 : i32 -> !pto.mask, i32 - %tmp_4 = arith.subi %c64, %col_5 : index - %tmp_5 = memref.subview %tmp_0[%row_3, %col_5] [%c1, %tmp_4] [%c1, %c1] : memref<8x64xf32, #pto.address_space> to memref, #pto.address_space> - %lhs_8 = pto.vlds %tmp_5[%c0] : memref, #pto.address_space> -> !pto.vreg<64xf32> - %tmp_6 = arith.subi %c64, %col_5 : index - %tmp_7 = memref.subview %tmp_1[%row_3, %col_5] [%c1, %tmp_6] [%c1, %c1] : memref<8x64xf32, #pto.address_space> to memref, #pto.address_space> - %rhs_9 = pto.vlds %tmp_7[%c0] : memref, #pto.address_space> -> !pto.vreg<64xf32> + %row_off_4 = arith.muli %row_3, %c64 : index + %linear_5 = arith.addi %row_off_4, %col_5 : index + %lhs_8 = pto.vlds %tmp_0[%linear_5] : !pto.ptr -> !pto.vreg<64xf32> + %rhs_9 = pto.vlds %tmp_1[%linear_5] : !pto.ptr -> !pto.vreg<64xf32> %summed_10 = pto.vadd %lhs_8, %rhs_9, %mask_6 : !pto.vreg<64xf32>, !pto.vreg<64xf32>, !pto.mask -> !pto.vreg<64xf32> - %tmp_8 = arith.subi %c64, %col_5 : index - %tmp_9 = memref.subview %tmp_2[%row_3, %col_5] [%c1, %tmp_8] [%c1, %c1] : memref<8x64xf32, #pto.address_space> to memref, #pto.address_space> - pto.vsts %summed_10, %tmp_9[%c0], %mask_6 : !pto.vreg<64xf32>, memref, #pto.address_space>, !pto.mask + pto.vsts %summed_10, %tmp_2[%linear_5], %mask_6 : !pto.vreg<64xf32>, !pto.ptr, !pto.mask scf.yield %remained_7 : i32 } } diff --git a/ptodsl/tests/test_jit_compile.py b/ptodsl/tests/test_jit_compile.py index 1d565a402f..e6265428fb 100644 --- a/ptodsl/tests/test_jit_compile.py +++ b/ptodsl/tests/test_jit_compile.py @@ -5135,21 +5135,20 @@ def _enter_inline_simt_with_resource_attr(): tile_slice_text = tile_slice_surface_probe.compile(BLOCK=128).mlir_text() expect_parse_roundtrip_and_verify(tile_slice_text, "tile slice surface specialization") - expect("memref.subview" in tile_slice_text, "tile[row, col:] should lower through memref.subview") - expect("memref.collapse_shape" not in tile_slice_text, "2D tile[row, col:] should lower directly to a rank-reduced memref view") - expect("pto.tile_buf_addr" in tile_slice_text, "tile[row, col:] should materialize a memref tile address view") + expect("memref.subview" not in tile_slice_text, "tile[row, col:] should no longer lower through memref.subview") + expect("pto.tile_buf_addr" in tile_slice_text, "tile[row, col:] should materialize a tile address pointer") expect( - "pto.vlds" in tile_slice_text and "memref<128xf32, strided<[1], offset: ?>, #pto.address_space>" in tile_slice_text, - "vlds(tile[row, col:]) should lower against the memref slice view", + "pto.vlds" in tile_slice_text and "!pto.ptr" in tile_slice_text, + "vlds(tile[row, col:]) should lower against the pointer slice view", ) expect( - "pto.vsts" in tile_slice_text and "memref<128xf32, strided<[1], offset: ?>, #pto.address_space>" in tile_slice_text, - "vsts(vec, tile[row, col:], mask) should lower against the memref slice view", + "pto.vsts" in tile_slice_text and "!pto.ptr" in tile_slice_text, + "vsts(vec, tile[row, col:], mask) should lower against the pointer slice view", ) tile_slice_1d_text = tile_slice_1d_surface_probe.compile(BLOCK=128).mlir_text() expect_parse_roundtrip_and_verify(tile_slice_1d_text, "1D tile slice surface specialization") - expect("memref.subview" in tile_slice_1d_text, "tile[start:] should lower through memref.subview") + expect("memref.subview" not in tile_slice_1d_text, "tile[start:] should lower without memref.subview") expect("pto.vldas" in tile_slice_1d_text, "vldas(tile[start:]) should lower against the 1D slice view") expect("pto.vldus" in tile_slice_1d_text, "vldus(tile[start:], align) should lower against the 1D slice view") expect("pto.vsts" in tile_slice_1d_text, "vsts(vec, tile[start:], mask) should lower against the 1D slice view") diff --git a/ptodsl/tests/test_tilelib_catalog.py b/ptodsl/tests/test_tilelib_catalog.py index c027e989e3..709e6ff709 100644 --- a/ptodsl/tests/test_tilelib_catalog.py +++ b/ptodsl/tests/test_tilelib_catalog.py @@ -253,7 +253,6 @@ SPECIAL_VALID_SHAPES[(_op, "src1")] = (1, 64) SHARED_RENDERED_OPS = ( "pto.tile_buf_addr", - "memref.subview", "scf.for", "pto.vsts", "pto.tilelang.instance", @@ -265,12 +264,6 @@ OPS_WITHOUT_VECTOR_STORE = {"pto.tcmp", "pto.tcmps", "pto.tsort32"} OPS_WITHOUT_VECTOR_STORE = OPS_WITHOUT_VECTOR_STORE | {"pto.tload", "pto.tstore", "pto.tstore_fp", "pto.textract_fp"} OPS_WITHOUT_VECTOR_STORE = OPS_WITHOUT_VECTOR_STORE | CUBE_OPS -OPS_WITHOUT_MEMREF_SUBVIEW = {"pto.tcmps", "pto.tsort32"} -OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | {"pto.texpands", "pto.tdivs", "pto.tfillpad_inplace"} -OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | {"pto.tload", "pto.tstore", "pto.tstore_fp", "pto.textract_fp"} -OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | ROW_REDUCTIONS -OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | ARG_COLUMN_REDUCTIONS -OPS_WITHOUT_MEMREF_SUBVIEW = OPS_WITHOUT_MEMREF_SUBVIEW | CUBE_OPS OPS_WITHOUT_LOOP = {"pto.tmrgsort"} OPS_WITHOUT_LOOP = OPS_WITHOUT_LOOP | {"pto.tstore_fp", "pto.textract_fp"} OPS_WITHOUT_LOOP = OPS_WITHOUT_LOOP | CUBE_OPS @@ -481,11 +474,6 @@ def test_each_catalog_entry_selects_and_renders(self): shared_op for shared_op in shared_ops if shared_op != "pto.vsts" ) - if op in OPS_WITHOUT_MEMREF_SUBVIEW: - shared_ops = tuple( - shared_op for shared_op in shared_ops - if shared_op != "memref.subview" - ) if op in OPS_WITHOUT_LOOP: shared_ops = tuple( shared_op for shared_op in shared_ops diff --git a/ptodsl/tests/test_tilelib_constraints.py b/ptodsl/tests/test_tilelib_constraints.py index c94bcd87ec..d68ca8f00a 100644 --- a/ptodsl/tests/test_tilelib_constraints.py +++ b/ptodsl/tests/test_tilelib_constraints.py @@ -47,7 +47,7 @@ def test_rejected_when_not_none_box(self): def test_legal_colmax_renders_structured_mlir(self): chosen = select("pto.tcolmax", "a5", _specs()) mlir = chosen.specialize(**_specs()).mlir_text() - for op in ("pto.tile_valid_rows", "memref.subview", "scf.for", "iter_args", + for op in ("pto.tile_valid_rows", "!pto.ptr", "scf.for", "iter_args", "pto.vmax", "pto.vsts", "pto.tilelang.instance"): self.assertIn(op, mlir) self.assertNotIn("pto.castptr", mlir) diff --git a/ptodsl/tests/test_tilelib_daemon.py b/ptodsl/tests/test_tilelib_daemon.py index ae16a168f1..5d565d7a9c 100644 --- a/ptodsl/tests/test_tilelib_daemon.py +++ b/ptodsl/tests/test_tilelib_daemon.py @@ -92,7 +92,7 @@ def test_instantiate_named_candidate_returns_structured_mlir(self): self.assertIn(f"func.func @{TADD}", mlir) for operation in ( "pto.tile_buf_addr", - "memref.subview", + "!pto.ptr", "pto.vlds", "pto.vadd", "pto.vsts", diff --git a/ptodsl/tests/test_tilelib_elementwise.py b/ptodsl/tests/test_tilelib_elementwise.py index 3b138ba9e4..29a95eafdc 100644 --- a/ptodsl/tests/test_tilelib_elementwise.py +++ b/ptodsl/tests/test_tilelib_elementwise.py @@ -22,7 +22,7 @@ } # Structured abstraction every elementwise template must preserve. -SHARED_OPS = ["pto.tile_buf_addr", "memref.subview", "scf.for", "iter_args", +SHARED_OPS = ["pto.tile_buf_addr", "!pto.ptr", "scf.for", "iter_args", "pto.plt_b32", "pto.vlds", "pto.vsts", "pto.tilelang.instance"] diff --git a/ptodsl/tests/test_tilelib_render.py b/ptodsl/tests/test_tilelib_render.py index d1405ef0a6..95305d1666 100644 --- a/ptodsl/tests/test_tilelib_render.py +++ b/ptodsl/tests/test_tilelib_render.py @@ -9,7 +9,7 @@ Asserts the rendered MLIR is *on par* with the tilelang golden (fixtures/tadd_a5_8x64_f32.golden.mlir): same structured abstraction, not byte-identical -(ptodsl differs in SSA naming, constant hoisting, index-vs-i32 carry, subview rank). +(ptodsl differs in SSA naming, constant hoisting, index-vs-i32 carry, ptr typing). """ import unittest @@ -23,13 +23,12 @@ # The structured abstraction the migration must preserve (see golden fixture). REQUIRED_OPS = [ "pto.tile_buf_addr", - "memref<8x64xf32, #pto.address_space>", + "!pto.ptr", "pto.tile_valid_rows", "pto.tile_valid_cols", "scf.for", "iter_args", # the inner loop carries `remained` (AST-rewrite .carry path) "pto.plt_b32", - "memref.subview", "pto.vlds", "pto.vadd", "pto.vsts", @@ -63,7 +62,7 @@ def test_func_is_a_tilelang_instance(self): def test_golden_fixture_uses_same_abstraction(self): self.assertTrue(FIXTURE.exists(), f"missing golden fixture {FIXTURE}") golden = FIXTURE.read_text(encoding="utf-8") - for op in ("pto.tile_buf_addr", "memref.subview", "pto.vlds", "pto.vadd", "pto.vsts", "pto.plt_b32"): + for op in ("pto.tile_buf_addr", "!pto.ptr", "pto.vlds", "pto.vadd", "pto.vsts", "pto.plt_b32"): self.assertIn(op, golden) diff --git a/test/lit/pto/cmo_cacheinvalid_single_line_invalid.pto b/test/lit/pto/cmo_cacheinvalid_single_line_invalid.pto index 82a866280d..70ce7b4c95 100644 --- a/test/lit/pto/cmo_cacheinvalid_single_line_invalid.pto +++ b/test/lit/pto/cmo_cacheinvalid_single_line_invalid.pto @@ -16,4 +16,4 @@ module { } } -// CHECK: single_cache_line address expects a GM pointer, GM memref, or GM tensor view +// CHECK: single_cache_line address expects a GM pointer or GM tensor view diff --git a/test/lit/pto/issue481_addptr_gm_slot_buffer.pto b/test/lit/pto/issue481_addptr_gm_slot_buffer.pto index 5c2cc19c4c..e5ca35d13d 100644 --- a/test/lit/pto/issue481_addptr_gm_slot_buffer.pto +++ b/test/lit/pto/issue481_addptr_gm_slot_buffer.pto @@ -26,8 +26,8 @@ module { return } - func.func @call_both(%arg0: memref<256xi64>, %arg1: !pto.ptr, %arg2: !pto.ptr, %arg3: !pto.ptr) attributes {pto.entry} { - pto.set_ffts %arg0 : memref<256xi64> + func.func @call_both(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr, %arg3: !pto.ptr) attributes {pto.entry} { + pto.set_ffts %arg0 : !pto.ptr call @cube_kernel(%arg1, %arg2) : (!pto.ptr, !pto.ptr) -> () call @vector_kernel(%arg1, %arg3) : (!pto.ptr, !pto.ptr) -> () return diff --git a/test/lit/pto/issue481_addptr_gm_slot_buffer_gss.pto b/test/lit/pto/issue481_addptr_gm_slot_buffer_gss.pto index 000ace1c83..60709a3126 100644 --- a/test/lit/pto/issue481_addptr_gm_slot_buffer_gss.pto +++ b/test/lit/pto/issue481_addptr_gm_slot_buffer_gss.pto @@ -26,8 +26,8 @@ module { return } - func.func @call_both(%arg0: memref<256xi64>, %arg1: !pto.ptr, %arg2: !pto.ptr, %arg3: !pto.ptr) attributes {pto.entry} { - pto.set_ffts %arg0 : memref<256xi64> + func.func @call_both(%arg0: !pto.ptr, %arg1: !pto.ptr, %arg2: !pto.ptr, %arg3: !pto.ptr) attributes {pto.entry} { + pto.set_ffts %arg0 : !pto.ptr call @cube_kernel(%arg1, %arg2) : (!pto.ptr, !pto.ptr) -> () call @vector_kernel(%arg1, %arg3) : (!pto.ptr, !pto.ptr) -> () return diff --git a/test/lit/pto/issue622_v_mte2_eventid_overlap_reproducer.pto b/test/lit/pto/issue622_v_mte2_eventid_overlap_reproducer.pto index 62d85aeeea..ce71d4c170 100644 --- a/test/lit/pto/issue622_v_mte2_eventid_overlap_reproducer.pto +++ b/test/lit/pto/issue622_v_mte2_eventid_overlap_reproducer.pto @@ -183,9 +183,9 @@ module { return } - func.func @call_both(%ffts: memref<256xi64>, %gm0: !pto.ptr, %gm1: !pto.ptr, %gm2: !pto.ptr) + func.func @call_both(%ffts: !pto.ptr, %gm0: !pto.ptr, %gm1: !pto.ptr, %gm2: !pto.ptr) attributes {pto.entry} { - pto.set_ffts %ffts : memref<256xi64> + pto.set_ffts %ffts : !pto.ptr call @cube_kernel(%gm0, %gm1, %gm2) : (!pto.ptr, !pto.ptr, !pto.ptr) -> () call @vector_kernel(%gm0, %gm1, %gm2) : (!pto.ptr, !pto.ptr, !pto.ptr) -> () return diff --git a/test/lit/pto/ptr_int_cast.pto b/test/lit/pto/ptr_int_cast.pto index ecd2da917f..bfe9783403 100644 --- a/test/lit/pto/ptr_int_cast.pto +++ b/test/lit/pto/ptr_int_cast.pto @@ -27,8 +27,8 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { // IR-LABEL: func.func @ptr_int_cast_kernel // IR-SAME: %[[SRC:[^,]+]]: !pto.ptr // IR: %[[OFF:.*]] = pto.addptr %[[SRC]], %{{.*}} : -> -// IR: %[[ADDR:.*]] = pto.ptrtoint %[[OFF]] : !pto.ptr -> i64 -// IR: %[[PTR:.*]] = pto.inttoptr %[[ADDR]] : i64 -> !pto.ptr +// IR: %[[ADDR:.*]] = pto.ptrtoint %[[OFF]] : -> i64 +// IR: %[[PTR:.*]] = pto.inttoptr %[[ADDR]] : i64 -> // IR-NOT: memref< // CPP-LABEL: AICORE void ptr_int_cast_kernel @@ -41,8 +41,8 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { // IR-LABEL: func.func @ptrtoint_addptr_multi_consumer // IR-SAME: %[[SRC64:[^,]+]]: !pto.ptr // IR: %[[OFF2:.*]] = pto.addptr %[[SRC64]], %{{.*}} : -> -// IR: pto.load_scalar %[[OFF2]]{{.*}} : !pto.ptr -> ui64 -// IR: pto.ptrtoint %[[OFF2]] : !pto.ptr -> i64 +// IR: pto.load_scalar %[[OFF2]]{{.*}} : -> ui64 +// IR: pto.ptrtoint %[[OFF2]] : -> i64 // IR-NOT: memref< // CPP-LABEL: AICORE void ptrtoint_addptr_multi_consumer diff --git a/test/lit/pto/tile_buf_addr_tile_native_arg.pto b/test/lit/pto/tile_buf_addr_tile_native_arg.pto index 052f30a282..096b2e5d37 100644 --- a/test/lit/pto/tile_buf_addr_tile_native_arg.pto +++ b/test/lit/pto/tile_buf_addr_tile_native_arg.pto @@ -25,7 +25,7 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { // NATIVE: IR Dump After PTOResolveReservedBuffers // NATIVE-LABEL: func.func @tile_buf_addr_tile_native_arg( // NATIVE-SAME: %[[TILE:.*]]: !pto.tile_buf -// NATIVE: %[[PTR:.*]] = pto.tile_buf_addr %[[TILE]] : !pto.tile_buf -> !pto.ptr +// NATIVE: %[[PTR:.*]] = pto.tile_buf_addr %[[TILE]] : !pto.tile_buf -> {{.*}} // NATIVE: call @consume(%[[PTR]]) : (!pto.ptr) -> () // NATIVE-NOT: memref.reinterpret_cast @@ -35,4 +35,4 @@ module attributes {pto.kernel_kind = #pto.kernel_kind} { // VPTO-LABEL: func.func @tile_buf_addr_tile_native_arg( // VPTO-SAME: !pto.tile_buf -// VPTO: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> !pto.ptr +// VPTO: pto.tile_buf_addr %{{.*}} : !pto.tile_buf -> {{.*}} diff --git a/test/lit/pto/tpush_tpop_globaltensor_split_half_slot_a3.pto b/test/lit/pto/tpush_tpop_globaltensor_split_half_slot_a3.pto index b561807000..0902cd3285 100644 --- a/test/lit/pto/tpush_tpop_globaltensor_split_half_slot_a3.pto +++ b/test/lit/pto/tpush_tpop_globaltensor_split_half_slot_a3.pto @@ -12,9 +12,9 @@ // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s module { - func.func @call_both(%ffts : memref<256xi64>, %fifo_mem : !pto.ptr) + func.func @call_both(%ffts : !pto.ptr, %fifo_mem : !pto.ptr) attributes {pto.entry} { - pto.set_ffts %ffts : memref<256xi64> + pto.set_ffts %ffts : !pto.ptr func.call @cube_kernel(%fifo_mem) : (!pto.ptr) -> () func.call @vector_kernel(%fifo_mem) : (!pto.ptr) -> () return diff --git a/test/lit/vpto/fold_tile_buf_intrinsics.pto b/test/lit/vpto/fold_tile_buf_intrinsics.pto index b52453e749..7b23584bca 100644 --- a/test/lit/vpto/fold_tile_buf_intrinsics.pto +++ b/test/lit/vpto/fold_tile_buf_intrinsics.pto @@ -26,33 +26,31 @@ // SHAPE: pto.vsts // After the later addr-only FoldTileBufIntrinsics: -// - tile_buf_addr should still carry the planned tile address -// - VPTO pointer normalization will rewrite the memref bridge to pto.ptr later +// - tile_buf_addr should still carry the planned tile address as pto.ptr // ADDR-LABEL: func.func @TADD -// ADDR: pto.tile_buf_addr {{.*}} -> memref<16x64xf32, #pto.address_space> +// ADDR: pto.tile_buf_addr {{.*}} -> {{.*}} // ADDR: pto.vlds -// ADDR: pto.tile_buf_addr {{.*}} -> memref<16x64xf32, #pto.address_space> +// ADDR: pto.tile_buf_addr {{.*}} -> {{.*}} // ADDR: pto.vlds // ADDR: pto.vadd -// ADDR: pto.tile_buf_addr {{.*}} -> memref<16x64xf32, #pto.address_space> +// ADDR: pto.tile_buf_addr {{.*}} -> {{.*}} // ADDR-NOT: pto.pointer_cast // ADDR-NOT: pto.tile_valid_rows // ADDR-NOT: pto.tile_valid_cols // ADDR: pto.vsts // After VPTO pointer normalization: -// - tile_buf_addr returns !pto.ptr directly -// - memref.subview has been converted to pto.addptr +// - tile_buf_addr has been folded to concrete pto.castptr addresses +// - tile-slice addressing is carried by the vlds/vsts offset operand // NORMALIZED-LABEL: func.func @TADD -// NORMALIZED: pto.tile_buf_addr {{.*}} -> !pto.ptr -// NORMALIZED: pto.addptr +// NORMALIZED: pto.castptr +// NORMALIZED: arith.muli // NORMALIZED: pto.vlds -// NORMALIZED: pto.tile_buf_addr {{.*}} -> !pto.ptr -// NORMALIZED: pto.addptr +// NORMALIZED: pto.castptr // NORMALIZED: pto.vlds // NORMALIZED: pto.vadd -// NORMALIZED: pto.tile_buf_addr {{.*}} -> !pto.ptr -// NORMALIZED: pto.addptr +// NORMALIZED: pto.castptr +// NORMALIZED-NOT: pto.tile_buf_addr // NORMALIZED-NOT: pto.pointer_cast // NORMALIZED: pto.vsts diff --git a/test/vpto_tilelang_inline_soft_divmod_fastpath.pto b/test/vpto_tilelang_inline_soft_divmod_fastpath.pto index a6ea117484..473a78496f 100644 --- a/test/vpto_tilelang_inline_soft_divmod_fastpath.pto +++ b/test/vpto_tilelang_inline_soft_divmod_fastpath.pto @@ -25,15 +25,15 @@ module attributes {pto.target_arch = "a5"} { func.func @kernel(%arg0: !pto.tile_buf, %arg1: !pto.tile_buf) attributes { pto.tilelang.instance } { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index - %tmp_0 = pto.tile_buf_addr %arg0 : !pto.tile_buf -> memref<8x16xi16, #pto.address_space> - %tmp_1 = pto.tile_buf_addr %arg1 : !pto.tile_buf -> memref<8x16xi16, #pto.address_space> + %tmp_0 = pto.tile_buf_addr %arg0 : !pto.tile_buf -> !pto.ptr + %tmp_1 = pto.tile_buf_addr %arg1 : !pto.tile_buf -> !pto.ptr pto.vecscope { %mask_0 = pto.pset_b16 "PAT_ALL" : !pto.mask - %vec_1 = pto.vlds %tmp_1[%c0] : memref<8x16xi16, #pto.address_space> -> !pto.vreg<128xi16> + %vec_1 = pto.vlds %tmp_1[%c0] : !pto.ptr -> !pto.vreg<128xi16> %q_59 = func.call @__tl_inline__tl_soft_vdiv_2(%vec_1, %vec_1, %mask_0) : (!pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.mask) -> !pto.vreg<128xi16> %r_81 = func.call @__tl_inline__tl_soft_vmod_4(%vec_1, %vec_1, %mask_0) : (!pto.vreg<128xi16>, !pto.vreg<128xi16>, !pto.mask) -> !pto.vreg<128xi16> - pto.vsts %q_59, %tmp_0[%c0], %mask_0 : !pto.vreg<128xi16>, memref<8x16xi16, #pto.address_space>, !pto.mask - pto.vsts %r_81, %tmp_0[%c1], %mask_0 : !pto.vreg<128xi16>, memref<8x16xi16, #pto.address_space>, !pto.mask + pto.vsts %q_59, %tmp_0[%c0], %mask_0 : !pto.vreg<128xi16>, !pto.ptr, !pto.mask + pto.vsts %r_81, %tmp_0[%c1], %mask_0 : !pto.vreg<128xi16>, !pto.ptr, !pto.mask } return }