From bed323cfff4fc17d64761635f51e92f71bd6e215 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 28 Jul 2026 19:45:16 +0800 Subject: [PATCH 1/6] docs: design for aligning Layout::NZ inference with pto-isa 5D form Design proposal for issue #527. PTOAS currently infers Layout::ND for the canonical pto-isa 2D NZ 5D view ([1, cols/C0, rows/16, 16, C0]) and infers Layout::NZ for a contiguous ND view that happens to match an off-by-one-dim pattern. The doc analyses the root cause, states why shape/stride alone cannot distinguish NZ from a contiguous ND 5D view, proposes a single shared inference utility with an explicit-layout-wins policy plus a guarded canonical pattern rule, and reports the measured blast radius over the lit corpus. No behavior change in this commit (docs only). --- docs/designs/nz-layout-inference-alignment.md | 363 ++++++++++++++++++ 1 file changed, 363 insertions(+) create mode 100644 docs/designs/nz-layout-inference-alignment.md diff --git a/docs/designs/nz-layout-inference-alignment.md b/docs/designs/nz-layout-inference-alignment.md new file mode 100644 index 000000000..c046f193d --- /dev/null +++ b/docs/designs/nz-layout-inference-alignment.md @@ -0,0 +1,363 @@ + + +# PTOAS `Layout::NZ` 推断与 pto-isa 标准五维形对齐设计 + +关联 issue: [#527](https://github.com/hw-native-sys/PTOAS/issues/527) —— PTOAS 支持 pto-isa 中的 `Layout::NZ` 推断。 + +本文只讨论 GM 侧 `GlobalTensor` 的 layout 推断（`tensor_view` / `partition_tensor_view` → `pto::Layout`），不涉及 tile 侧 `blayout/slayout`、也不涉及 fixpipe 的 `NZ2ND/NZ2DN/NZ2NZ` 转换。 + +--- + +## 1. 事实基准：pto-isa 的 NZ 五维标准形 + +pto-isa 侧二维逻辑 tensor 的 NZ 五维表示由 `TileShape2D` / `BaseShape2D` 定义 +（`include/pto/common/pto_tile.hpp`，常量见 `include/pto/common/constants.hpp`： +`C0_SIZE_BYTE = 32`、`FRACTAL_NZ_ROW = 16`）： + +``` +C0 = 32 / sizeof(T) +shape = [1, cols / C0, rows / 16, 16, C0] +stride = [rows * cols, rows * C0, 16 * C0, C0, 1] +``` + +即 5 个维度的语义固定为： + +| 维度 | 含义 | 取值 | +|---|---|---| +| `d0` | batch（二维场景恒为 1） | `1` | +| `d1` | 列分块数 `n1` | `cols / C0` | +| `d2` | 行分块数 `m1` | `rows / 16` | +| `d3` | fractal 内行数 | `16`（`FRACTAL_NZ_ROW`） | +| `d4` | fractal 内列数 `C0` | `32 / sizeof(T)` | + +fp32 且 `(rows, cols) = (128, 64)`（`C0 = 8`）： + +``` +shape = [1, 8, 8, 16, 8] +stride = [8192, 1024, 128, 8, 1] +``` + +**关键结构不变量**（后文规则全部由它推出）： + +``` +shape[3] == 16 +shape[4] == C0 <=> shape[4] * sizeof(T) == 32 +shape[3] * shape[4] * sizeof(T) == 512 // 一个 fractal = 512B +stride[4] == 1 +stride[3] == C0 == shape[4] +stride[2] == 16 * C0 == shape[3] * shape[4] +stride[1] == shape[2] * stride[2] +stride[0] == shape[1] * stride[1] +``` + +## 2. PTOAS 现状 + +### 2.1 同一条规则在三处重复实现 + +| 位置 | 函数 | 用途 | +|---|---|---| +| `lib/PTO/Transforms/InferPTOLayout.cpp:174` | `inferNZLayout()` | `pto-infer-layout` pass，给 `make_tensor_view` / `reinterpret_cast` / `subview` / `tload` / `tstore` 打 `layout` 属性 | +| `lib/PTO/IR/PTO.cpp:1836` | `inferLayout()`（`getLogicalViewLayout` 调用） | verifier 侧判断逻辑 layout（如 mgather/mscatter 的 ND 约束，`lib/PTO/IR/PTO.cpp:4366`） | +| `lib/PTO/Transforms/PTOToEmitC.cpp:4350` | `inferFallbackGlobalTensorLayout()` | EmitC 兜底：`layout` 属性缺失时重新推断 | + +三份实现的 NZ 判定条件完全一致（右对齐到 5 维后）： + +```cpp +shape[2] == 16 +shape[2] * shape[3] * elemBytes == 512 +stride[4] == 1 +stride[3] == shape[4] +``` + +另有 `lib/PTO/Transforms/PTOCanonicalizeIR.cpp:88` 附近的 rank2 → rank5 规范化， +与 `rightAlignTo5D()` / `buildGlobalTensorShapeAndStride()` 共用同一套 padding 规则， +一并纳入"单一实现"的收敛范围。 + +### 2.2 现规则 = 标准形规则"错位一维" + +把现规则与第 1 节不变量并排看： + +| 条件 | 现规则 | pto-isa 标准形 | 结论 | +|---|---|---|---| +| fractal 行数 | `shape[2] == 16` | `shape[3] == 16` | **错位一维** | +| fractal 字节 | `shape[2] * shape[3] * eb == 512` | `shape[3] * shape[4] * eb == 512` | **错位一维** | +| 最内连续 | `stride[4] == 1` | `stride[4] == 1` | 一致 | +| C0 连续 | `stride[3] == shape[4]` | `stride[3] == shape[4]` | 一致 | +| fractal 跨度 | 未检查 | `stride[2] == shape[3] * shape[4]` | 缺失 | +| 分块跨度 | 未检查 | `stride[1] == shape[2] * stride[2]` | 缺失 | +| batch 跨度 | 未检查 | `stride[0] == shape[1] * stride[1]` | 缺失 | + +两条 stride 条件本来就是对的，**只有两条 align 条件整体错了一维**，再叠加"外层 +stride 完全不校验"，于是同时产生了漏判（标准形判成 ND）和误判（连续 ND 判成 NZ）。 + +### 2.3 复现（本地实测，非推演） + +用不依赖 python binding 的等价 `.pto` 复现 issue 的两个用例： + +```mlir +// 用例①：pto-isa canonical 2D NZ 五维形 +%view = pto.make_tensor_view %dst, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + : !pto.tensor_view<1x8x8x16x8xf32> +%part = pto.partition_view %view, + offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c8, %c8, %c16, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> -> !pto.partition_tensor_view<1x8x8x16x8xf32> +%tile = pto.alloc_tile : !pto.tile_buf +pto.tstore ins(%tile : ...) outs(%part : !pto.partition_tensor_view<1x8x8x16x8xf32>) +``` + +`ptoas --pto-arch=a5` 实际输出： + +```cpp +// ① canonical NZ 形 -> 被判成 ND（错） +GlobalTensor, pto::Stride<8192, 1024, 128, 8, 1>, + pto::Layout::ND> + +// ② 迎合现规则的形状 -> 被判成 NZ（本身是连续 ND 数据） +GlobalTensor, pto::Stride<2048, 2048, 128, 16, 1>, + pto::Layout::NZ> +``` + +**第三个症状（issue 未提，但更致命）**：用户显式标注也救不回来。给用例①的 +`make_tensor_view` 加上 `{layout = #pto.layout}`： + +``` +error: layout mismatch: user-specified layout=nz but inferred=nd +``` + +来源是 `lib/PTO/Transforms/InferPTOLayout.cpp:275` 的 `verifyOrSetLayoutAttr()`： +推断结果被当作"真值"去否决用户显式声明。也就是说目前**既推不出 NZ，也没有逃生通道**。 + +## 3. 一个必须先承认的约束 + +把第 1 节的 stride 展开： + +``` +stride[2] = shape[3] * stride[3] +stride[1] = shape[2] * stride[2] +stride[0] = shape[1] * stride[1] +``` + +这正是五维**连续（行主序累积积）**的定义。结论： + +> **NZ 五维标准形，在 shape/stride 数值上与"同 shape 的连续 ND 五维视图"完全相同， +> 无法区分。** + +差异只存在于"这 5 个维度分别代表什么"这个语义层面，不存在于内存 pattern 层面。 +由此推出两条设计原则： + +- **P1：layout 应该被"携带"，而不是被"猜"。** 显式 `layout` 属性必须是权威来源， + 推断只是缺省兜底。 +- **P2：纯 pattern 推断必须带消歧门槛。** 否则任何"末两维恰好是 `(16, 32/sizeof(T))` + 的连续视图"都会被升级成 NZ —— 包括最常见的 `16 x C0` 二维小 tile。 + +P2 不是理论担忧，第 5 节有实测数据。 + +## 4. 设计方案 + +### 4.1 收敛为唯一实现 + +新增 `include/PTO/IR/PTOLayoutUtils.h` + `lib/PTO/IR/PTOLayoutUtils.cpp`，导出： + +```cpp +namespace mlir::pto { + +// C0 = 32 / elemBytes；elemBytes 不能整除 32 时返回 nullopt（sub-byte / packed 类型） +std::optional getNZC0Elems(unsigned elemBytes); + +// 结构必要条件：shape/stride 是否是 pto-isa NZ 五维标准形 +bool isNZCompatible5D(ArrayRef shape5D, ArrayRef stride5D, + unsigned elemBytes); + +struct LayoutInferOptions { + bool allowNZFromPatternOnly = true; // 无外部证据时是否允许纯 pattern 判 NZ + std::optional preferredMinor2D; // 现有 ND/DN 歧义消解入口 +}; + +std::optional inferLayout5D(ArrayRef shape5D, + ArrayRef stride5D, + unsigned elemBytes, + const LayoutInferOptions &opts, + bool *isMinor2DAmbiguous = nullptr); +} // namespace mlir::pto +``` + +`InferPTOLayout.cpp` / `PTO.cpp` / `PTOToEmitC.cpp` 三处改为调用同一实现， +删除各自的私有副本。这一步是纯重构（NFC），单独成 PR，便于回归定位。 + +### 4.2 `isNZCompatible5D`：结构必要条件 + +```cpp +bool isNZCompatible5D(shape, stride, elemBytes) { + auto c0 = getNZC0Elems(elemBytes); // 32 % elemBytes != 0 -> false + if (!c0) return false; + return shape[3] == 16 && shape[4] == *c0 + && stride[4] == 1 + && stride[3] == *c0 + && stride[2] == 16 * *c0 + && stride[1] == shape[2] * stride[2] + && stride[0] == shape[1] * stride[1]; +} +``` + +用途有两个，且**两个用途门槛不同**（这是本设计的核心）： + +- 校验显式 `layout = nz` 是否自洽 —— 只看这个谓词； +- 无显式 layout 时的推断 —— 还要过 4.3 的门槛。 + +### 4.3 推断门槛：只认"不可能是二维 ND 视图"的形状 + +```cpp +bool inferNZFromPattern(shape, stride, elemBytes) { + return isNZCompatible5D(shape, stride, elemBytes) + && shape[0] == 1 // 二维 NZ，batch 维恒为 1 + && (shape[1] > 1 || shape[2] > 1); // 至少有一个分块维非退化 +} +``` + +两个附加条件的理由： + +- `shape[0] == 1`：issue 要求的是**二维** NZ；`batch > 1` 的形状在 pto-isa 侧没有 + 对应的 `TileShape2D`，留给显式标注，避免把连续五维 ND 张量整片吃掉。 +- `shape[1] > 1 || shape[2] > 1`：`[1,1,1,16,C0]` 是 rank2 视图规范化后的标准产物 + （`PTOCanonicalizeIR`），也是最常见的 `16 x C0` 向量 tile。此时 NZ 与 ND 的字节 + 排布**完全等价**（单 fractal），判成 NZ 没有任何收益，却会改变生成的 C++ 模板参数、 + 进而影响 pto-isa 侧的重载选择。实测这条门槛消除了 100% 的新增误判（第 5 节）。 + +保留 `LayoutInferOptions::allowNZFromPatternOnly`，是为了给"上层已经有更强证据"的 +调用点（例如 mgather/mscatter 的 ND-only 校验路径）一个关掉 pattern 推断的开关。 + +### 4.4 显式 layout 优先 + +`verifyOrSetLayoutAttr()` 的语义调整为： + +| 显式属性 | 结构自洽性 | 现行为 | 新行为 | +|---|---|---|---| +| `nz` | `isNZCompatible5D == true` | **报错** | 接受，保留 `nz`，不打 `pto.inferred_layout` | +| `nz` | `isNZCompatible5D == false` | 报错 | 报错，但错误信息升级为"哪一条不变量不满足" | +| `nd`/`dn` | 与推断不同 | 现有 minor-2D 歧义豁免 | 不变 | + +这条直接解决 2.3 的第三个症状：即使 pattern 推断因门槛保守而不升级 NZ，用户/前端 +也永远有一条显式通道。诊断信息形如： + +``` +error: layout mismatch: user-specified layout=nz but shape/stride is not an NZ + 5D form: expected shape[3]==16 (got 8), shape[4]==8 (C0 for f32, got 16) +``` + +### 4.5 `partition_view` / `memref.subview` 的 NZ 传播 + +现状：`InferPTOLayout.cpp` 的 subview 分支无条件继承源 layout。对 NZ 这是不安全的 +—— 在 fractal 内部切分后，结果已经不是合法 NZ。规则改为： + +1. 源为 NZ 时，只有当切分满足以下条件才继承 NZ： + - `d3`/`d4` 维保持完整（offset 为 0、size 等于源 size）； + - `d1`/`d2` 维的 size 可缩小，但 stride 保持不变； + - offset 在 `d1`/`d2` 上按整块对齐。 +2. 不满足时不 silently 退回 ND，而是发 `emitError`（"NZ view cannot be partitioned + inside a fractal"），避免生成静默错码。 + +issue 明确提到 `partition_tensor_view`，这条属于本次范围内。 + +### 4.6 动态形状（阶段二） + +pto-isa 的 `TileShape2D`/`BaseShape2D` 允许 `rows`/`cols` 为 `DYNAMIC`，此时 +`shape[1]`、`shape[2]`、`stride[0]`、`stride[1]` 动态，而结构维 +（`shape[3]`、`shape[4]`、`stride[2..4]`）仍是静态常量。因此可以在 +"结构维静态 + 分块维动态"时仍判定 NZ。当前实现要求全部 const-fold，直接放弃推断 +（`getStaticShapeAndStride()` 返回 false），属于可独立推进的增强，放到阶段二。 + +### 4.7 与旧规则的兼容策略 + +| 方案 | 说明 | 评价 | +|---|---|---| +| A. 直接替换 | 只保留标准形规则 | **推荐**。旧规则命中的非标准形本质是误判，继续保留会持续生成错误的 `GlobalTensor` 模板参数 | +| B. 并集（新规则 OR 旧规则） | 兼容一切现有行为 | 不推荐：把"连续 ND 判成 NZ"固化成契约 | +| C. 替换 + 一个版本的过渡开关 | 加 `--pto-legacy-nz-infer`（默认关） | 若下游有存量依赖再启用；本次实测影响面极小（第 5 节），倾向不引入 | + +推荐 A，若 review 中出现下游存量用例再降级到 C。 + +## 5. 影响面实测 + +方法：用当前 `ptoas` 对 `test/lit/pto/*.pto`（476）+ `test/lit/tile_fusion/*.pto`（42） +分别按 `--pto-arch=a5` 和 `--pto-arch=a3` 跑一遍，抽取全部生成的 +`GlobalTensor, Stride<...>, Layout::X>`（含 `using GTShape_*` 别名形式）， +再离线对同一批 shape/stride 分别套用旧规则、标准形规则、带门槛的标准形规则。 + +样本：242 个文件产生输出，共 **1558** 个 `GlobalTensor` 实例化点， +**64** 组去重后的 `(elem, shape, stride)`。当前 layout 分布：ND 1007 / NZ 529 / DN 22。 + +| 项 | 结果 | +|---|---| +| 当前判为 NZ 的去重形状 | 5 组 | +| 其中符合 pto-isa 标准形（改后仍为 NZ） | 2 组：`half`/`bfloat16_t` `[1,1,16,16,16]` `stride=[4096,4096,256,16,1]` | +| 其中由显式属性设置（不走推断，不受影响） | 1 组：`tinsert_a5_vec_mat_mode_lowering` 的 `half [1,1,1,32,32]` | +| **改后会失去 NZ** | 2 组：`int8_t [1,1,16,32,16]`、`int64_t [1,1,16,4,1]` —— 全部只出现在 `test/lit/pto/globaltensor_layout_bytewidth_emitc.pto` | +| 不带门槛时新增 NZ | 4 组 / 80 个点 / 38 个文件，**全部**是 `[1,1,1,16,C0]` 退化单 fractal（`float`/`half`/`float8_e4m3_t`/`hifloat8_t`） | +| **带门槛（4.3）新增 NZ** | **0** | + +两点解读： + +1. 门槛条款把误判从 80 个点压到 0，验证了 P2 的必要性。 +2. 唯一需要改期望值的测试是 `globaltensor_layout_bytewidth_emitc.pto` + （`comm.tbroadcast` 的 i8/i64 用例）。这两组形状按 pto-isa 定义本来就不是 NZ + （i8 的 `C0 = 32`，标准形应为 `[1, cols/32, rows/16, 16, 32]`），当前期望值本身 + 固化了 2.2 的错位 bug。改测试前需确认 `comm.tbroadcast` 侧对 `Layout` 是否敏感。 + +另需交叉验证：`PTO.cpp:4366` 的 mgather/mscatter "mem partition view 必须是 ND"校验 +走的是同一套推断。规则收紧后，理论上存在"原本判 ND 的五维 mem view 变成 NZ 导致新 +报错"的可能；带门槛后本地 lit 语料未观察到，但需要在实现 PR 里补一条针对性用例。 + +## 6. 实施拆分 + +| PR | 内容 | 风险 | +|---|---|---| +| P1 | 抽出 `PTOLayoutUtils`，三处调用点收敛，行为完全不变（NFC） | 低 | +| P2 | 标准形规则 + 4.3 门槛 + 4.7 方案 A；更新 `globaltensor_layout_bytewidth_emitc.pto` | 中 | +| P3 | 4.4 显式 layout 优先 + 诊断信息细化 | 低 | +| P4 | 4.5 partition/subview NZ 传播校验 | 中 | +| P5 | 4.6 动态形状支持（可选） | 中 | + +跨层同步检查（按 `.claude/rules/cross-layer-sync.md`）：本设计不改 ODS 算子签名， +`PTO_LayoutAttr` 已存在 `nz`；需要同步的是 IR verifier（`PTO.cpp`）、pass +（`InferPTOLayout.cpp`）、EmitC（`PTOToEmitC.cpp`）、以及 python 侧 +`make_tensor_view` 的 `layout` 传参样例与文档。 + +## 7. 测试计划 + +新增 `test/lit/pto/issue527_nz_canonical_view_infer.pto`（a5 + a3 双 RUN），覆盖： + +1. **正例**：canonical `[1,8,8,16,8]` fp32 → `pto::Layout::NZ`； +2. **正例**：`half` `[1,4,8,16,16]`、`int8_t` `[1,2,8,16,32]` → NZ（验证 C0 随 dtype 变化）； +3. **负例**：issue 用例② `[4,1,16,8,16]` → `ND`（旧规则误判被修掉）； +4. **负例**：`[1,1,1,16,8]` 单 fractal → 保持 `ND`（门槛条款，防回归 80 个点）； +5. **负例**：stride 被打断（如 `stride[1]` 非 `shape[2]*stride[2]`）→ `ND`； +6. **显式通道**：canonical 形 + `{layout = #pto.layout}` → 不再报 + `layout mismatch`，且输出 NZ； +7. **显式冲突**：`[4,1,16,8,16]` + `{layout = #pto.layout}` → 报错且信息指明 + 是哪条不变量不满足； +8. **partition**：对 NZ 视图在 `d1` 上整块切分 → 继承 NZ；在 `d3` 内部切分 → 报错。 + +E2E：`test/tilelang_st/npu/a5` 下补一条 NZ store 用例（若板卡资源允许）， +用 `nz_store_probe.py` 的逻辑做数据比对，确认生成的 C++ 与 pto-isa 语义一致。 + +## 8. 待确认问题 + +1. **是否接受 4.7 方案 A**（直接替换旧规则）？影响面只有一个 lit 测试的 2 组期望值。 +2. **`batch > 1` 的 NZ**（`shape[0] > 1`）是否需要推断？当前设计只在显式标注时接受。 +3. **sub-byte / packed 类型**（`int4`、`float4_e1m2x2`）：`32 % elemBytes` 语义如何定义？ + 当前设计直接不推断 NZ，需 pto-isa 侧确认是否存在这类 NZ 用法。 +4. `comm.tbroadcast` 路径对 `GlobalTensor` 的 `Layout` 模板参数是否敏感？ + 决定 `globaltensor_layout_bytewidth_emitc.pto` 的期望值怎么改。 +5. `mgather`/`mscatter` 的 "mem 必须 ND" 约束，在 mem 确实是 NZ 五维时是否应放开， + 还是维持报错（本设计维持现状，仅补测试）。 From c5623c338cbda1c3a351514cfc62d5368ac4462d Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Tue, 28 Jul 2026 20:08:15 +0800 Subject: [PATCH 2/6] docs: finalize NZ layout implementation design --- docs/designs/nz-layout-inference-alignment.md | 452 ++++++++---------- 1 file changed, 197 insertions(+), 255 deletions(-) diff --git a/docs/designs/nz-layout-inference-alignment.md b/docs/designs/nz-layout-inference-alignment.md index c046f193d..5ce41060e 100644 --- a/docs/designs/nz-layout-inference-alignment.md +++ b/docs/designs/nz-layout-inference-alignment.md @@ -8,67 +8,73 @@ INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A See LICENSE in the root of the software repository for the full text of the License. --> -# PTOAS `Layout::NZ` 推断与 pto-isa 标准五维形对齐设计 +# PTOAS `Layout::NZ` 推断与 pto-isa 标准五维形对齐 -关联 issue: [#527](https://github.com/hw-native-sys/PTOAS/issues/527) —— PTOAS 支持 pto-isa 中的 `Layout::NZ` 推断。 +关联 issue: [#527](https://github.com/hw-native-sys/PTOAS/issues/527)。 -本文只讨论 GM 侧 `GlobalTensor` 的 layout 推断（`tensor_view` / `partition_tensor_view` → `pto::Layout`），不涉及 tile 侧 `blayout/slayout`、也不涉及 fixpipe 的 `NZ2ND/NZ2DN/NZ2NZ` 转换。 +本文定义 PR #1027 的最终实现约束。设计、实现、测试和文档都在同一个 PR 中完成； +PR 合入后关闭 #527。 ---- +范围仅包含 GM 侧 `GlobalTensor` 的 layout 解析与传播 +（`tensor_view` / `partition_tensor_view` / memref view → `pto::Layout`），不改变 tile +侧 `blayout/slayout`，也不改变 fixpipe 的 `NZ2ND/NZ2DN/NZ2NZ` 转换语义。 -## 1. 事实基准：pto-isa 的 NZ 五维标准形 +## 1. pto-isa 的 NZ 五维标准形 -pto-isa 侧二维逻辑 tensor 的 NZ 五维表示由 `TileShape2D` / `BaseShape2D` 定义 -（`include/pto/common/pto_tile.hpp`，常量见 `include/pto/common/constants.hpp`： -`C0_SIZE_BYTE = 32`、`FRACTAL_NZ_ROW = 16`）： +pto-isa 使用 `TileShape2D` / `BaseShape2D` 表示二维 NZ tensor。常量 +`C0_SIZE_BYTE = 32`、`FRACTAL_NZ_ROW = 16`，标准根视图为： -``` +```text C0 = 32 / sizeof(T) shape = [1, cols / C0, rows / 16, 16, C0] stride = [rows * cols, rows * C0, 16 * C0, C0, 1] ``` -即 5 个维度的语义固定为： +五个维度的固定语义如下： -| 维度 | 含义 | 取值 | +| 维度 | 含义 | 标准根视图取值 | |---|---|---| -| `d0` | batch（二维场景恒为 1） | `1` | +| `d0` | 二维场景的占位维 | `1` | | `d1` | 列分块数 `n1` | `cols / C0` | | `d2` | 行分块数 `m1` | `rows / 16` | -| `d3` | fractal 内行数 | `16`（`FRACTAL_NZ_ROW`） | -| `d4` | fractal 内列数 `C0` | `32 / sizeof(T)` | +| `d3` | fractal 内行数 | `16` | +| `d4` | fractal 内列数 | `C0` | -fp32 且 `(rows, cols) = (128, 64)`（`C0 = 8`）： +例如 fp32 `(rows, cols) = (128, 64)` 时： -``` +```text shape = [1, 8, 8, 16, 8] stride = [8192, 1024, 128, 8, 1] ``` -**关键结构不变量**（后文规则全部由它推出）： +标准根视图满足： -``` +```text shape[3] == 16 -shape[4] == C0 <=> shape[4] * sizeof(T) == 32 -shape[3] * shape[4] * sizeof(T) == 512 // 一个 fractal = 512B +shape[4] == C0 stride[4] == 1 -stride[3] == C0 == shape[4] -stride[2] == 16 * C0 == shape[3] * shape[4] +stride[3] == C0 +stride[2] == 16 * C0 stride[1] == shape[2] * stride[2] stride[0] == shape[1] * stride[1] ``` -## 2. PTOAS 现状 +这里的 element 指 PTOAS/pto-isa C++ 类型的存储元素。对于 +`float4_e1m2x2_t` / `float4_e2m1x2_t`，一个存储元素是一字节 packed pair， +shape 的列数也按 packed pair 计数，因此 `sizeof(T) == 1`、`C0 == 32`。只有存储 +字节数为 0，或 32 不能被存储字节数整除时，才不能形成 NZ C0。 -### 2.1 同一条规则在三处重复实现 +## 2. 当前问题 -| 位置 | 函数 | 用途 | -|---|---|---| -| `lib/PTO/Transforms/InferPTOLayout.cpp:174` | `inferNZLayout()` | `pto-infer-layout` pass，给 `make_tensor_view` / `reinterpret_cast` / `subview` / `tload` / `tstore` 打 `layout` 属性 | -| `lib/PTO/IR/PTO.cpp:1836` | `inferLayout()`（`getLogicalViewLayout` 调用） | verifier 侧判断逻辑 layout（如 mgather/mscatter 的 ND 约束，`lib/PTO/IR/PTO.cpp:4366`） | -| `lib/PTO/Transforms/PTOToEmitC.cpp:4350` | `inferFallbackGlobalTensorLayout()` | EmitC 兜底：`layout` 属性缺失时重新推断 | +PTOAS 当前在三处重复实现 layout 推断： -三份实现的 NZ 判定条件完全一致（右对齐到 5 维后）： +| 位置 | 用途 | +|---|---| +| `lib/PTO/Transforms/InferPTOLayout.cpp` | 给 view、memref view 和 load/store 附加 layout | +| `lib/PTO/IR/PTO.cpp` | verifier 获取逻辑 layout | +| `lib/PTO/Transforms/PTOToEmitC.cpp` | EmitC 缺少 layout 属性时兜底 | + +当前 NZ 条件为： ```cpp shape[2] == 16 @@ -77,287 +83,223 @@ stride[4] == 1 stride[3] == shape[4] ``` -另有 `lib/PTO/Transforms/PTOCanonicalizeIR.cpp:88` 附近的 rank2 → rank5 规范化， -与 `rightAlignTo5D()` / `buildGlobalTensorShapeAndStride()` 共用同一套 padding 规则， -一并纳入"单一实现"的收敛范围。 - -### 2.2 现规则 = 标准形规则"错位一维" - -把现规则与第 1 节不变量并排看： - -| 条件 | 现规则 | pto-isa 标准形 | 结论 | -|---|---|---|---| -| fractal 行数 | `shape[2] == 16` | `shape[3] == 16` | **错位一维** | -| fractal 字节 | `shape[2] * shape[3] * eb == 512` | `shape[3] * shape[4] * eb == 512` | **错位一维** | -| 最内连续 | `stride[4] == 1` | `stride[4] == 1` | 一致 | -| C0 连续 | `stride[3] == shape[4]` | `stride[3] == shape[4]` | 一致 | -| fractal 跨度 | 未检查 | `stride[2] == shape[3] * shape[4]` | 缺失 | -| 分块跨度 | 未检查 | `stride[1] == shape[2] * stride[2]` | 缺失 | -| batch 跨度 | 未检查 | `stride[0] == shape[1] * stride[1]` | 缺失 | - -两条 stride 条件本来就是对的，**只有两条 align 条件整体错了一维**，再叠加"外层 -stride 完全不校验"，于是同时产生了漏判（标准形判成 ND）和误判（连续 ND 判成 NZ）。 - -### 2.3 复现（本地实测，非推演） - -用不依赖 python binding 的等价 `.pto` 复现 issue 的两个用例： - -```mlir -// 用例①：pto-isa canonical 2D NZ 五维形 -%view = pto.make_tensor_view %dst, - shape = [%c1, %c8, %c8, %c16, %c8], - strides = [%c8192, %c1024, %c128, %c8, %c1] - : !pto.tensor_view<1x8x8x16x8xf32> -%part = pto.partition_view %view, - offsets = [%c0, %c0, %c0, %c0, %c0], sizes = [%c1, %c8, %c8, %c16, %c8] - : !pto.tensor_view<1x8x8x16x8xf32> -> !pto.partition_tensor_view<1x8x8x16x8xf32> -%tile = pto.alloc_tile : !pto.tile_buf -pto.tstore ins(%tile : ...) outs(%part : !pto.partition_tensor_view<1x8x8x16x8xf32>) -``` +它相对 pto-isa 标准形错位一维，并且没有验证 `stride[2..0]`。结果是： -`ptoas --pto-arch=a5` 实际输出： +- canonical `[1,8,8,16,8]` fp32 被判成 ND； +- `[4,1,16,8,16]` 这类连续 ND 形状被误判成 NZ； +- 显式 `{layout = #pto.layout}` 仍会被错误的 pattern 推断否决。 -```cpp -// ① canonical NZ 形 -> 被判成 ND（错） -GlobalTensor, pto::Stride<8192, 1024, 128, 8, 1>, - pto::Layout::ND> +## 3. 必须保留的语义边界 -// ② 迎合现规则的形状 -> 被判成 NZ（本身是连续 ND 数据） -GlobalTensor, pto::Stride<2048, 2048, 128, 16, 1>, - pto::Layout::NZ> -``` +canonical NZ 根视图的数值 stride 同时满足连续五维 ND 的累积积关系。因此仅凭 +shape/stride 无法证明用户的逻辑语义一定是 NZ。 -**第三个症状（issue 未提，但更致命）**：用户显式标注也救不回来。给用例①的 -`make_tensor_view` 加上 `{layout = #pto.layout}`： +本 PR 采用一条确定的解析顺序： -``` -error: layout mismatch: user-specified layout=nz but inferred=nd -``` +1. **显式 layout 是权威来源。** 显式 `ND`、`DN`、`NZ` 都不会被 pattern 推断覆盖； +2. **派生 view 继承已经解析的源 layout。** NZ 子视图需要额外通过 fractal 边界校验； +3. **只有无显式属性、无源 layout 的根视图才走 pattern 推断；** +4. pattern 推断把满足 canonical NZ 规则的非退化二维五维形约定为 NZ，其他情况继续 + 使用现有 ND/DN 推断。 -来源是 `lib/PTO/Transforms/InferPTOLayout.cpp:275` 的 `verifyOrSetLayoutAttr()`： -推断结果被当作"真值"去否决用户显式声明。也就是说目前**既推不出 NZ，也没有逃生通道**。 +第 4 条是一项 PTOAS 输入约定，不声称它能从数值上区分所有连续 5D ND。用户需要表达 +同 shape/stride 的 ND 时，必须能够通过显式 `layout = nd` 覆盖该约定。 -## 3. 一个必须先承认的约束 +## 4. 唯一实现 -把第 1 节的 stride 展开： +### 4.1 共享 layout 工具 -``` -stride[2] = shape[3] * stride[3] -stride[1] = shape[2] * stride[2] -stride[0] = shape[1] * stride[1] -``` +新增 `include/PTO/IR/PTOLayoutUtils.h` 和 +`lib/PTO/IR/PTOLayoutUtils.cpp`，集中实现： -这正是五维**连续（行主序累积积）**的定义。结论： +```cpp +namespace mlir::pto { -> **NZ 五维标准形，在 shape/stride 数值上与"同 shape 的连续 ND 五维视图"完全相同， -> 无法区分。** +std::optional +getNZC0StorageElems(unsigned storageElemBytes); -差异只存在于"这 5 个维度分别代表什么"这个语义层面，不存在于内存 pattern 层面。 -由此推出两条设计原则： +bool hasNZInnerStructure5D(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes); -- **P1：layout 应该被"携带"，而不是被"猜"。** 显式 `layout` 属性必须是权威来源， - 推断只是缺省兜底。 -- **P2：纯 pattern 推断必须带消歧门槛。** 否则任何"末两维恰好是 `(16, 32/sizeof(T))` - 的连续视图"都会被升级成 NZ —— 包括最常见的 `16 x C0` 二维小 tile。 +bool isNZViewCompatible5D(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes); -P2 不是理论担忧，第 5 节有实测数据。 +bool isCanonicalNZRoot5D(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes); -## 4. 设计方案 +std::optional +inferLayout5D(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes, + std::optional preferredMinor2D = std::nullopt, + bool *isMinor2DAmbiguous = nullptr); -### 4.1 收敛为唯一实现 +} // namespace mlir::pto +``` -新增 `include/PTO/IR/PTOLayoutUtils.h` + `lib/PTO/IR/PTOLayoutUtils.cpp`，导出： +`InferPTOLayout.cpp`、`PTO.cpp`、`PTOToEmitC.cpp` 都调用这套实现，删除各自的私有 +NZ/ND/DN 判定副本。调用点必须先解析显式属性和源 view 属性，只有确实缺失 layout 时 +才调用 `inferLayout5D()`。 -```cpp -namespace mlir::pto { +### 4.2 根视图与子视图使用不同谓词 -// C0 = 32 / elemBytes；elemBytes 不能整除 32 时返回 nullopt（sub-byte / packed 类型） -std::optional getNZC0Elems(unsigned elemBytes); +`isCanonicalNZRoot5D()` 只用于无 layout 根视图的 pattern 推断，要求完整的紧密 +canonical stride： -// 结构必要条件：shape/stride 是否是 pto-isa NZ 五维标准形 -bool isNZCompatible5D(ArrayRef shape5D, ArrayRef stride5D, - unsigned elemBytes); +```cpp +hasNZInnerStructure5D(shape, stride, bytes) + && stride[1] == shape[2] * stride[2] + && stride[0] == shape[1] * stride[1] +``` -struct LayoutInferOptions { - bool allowNZFromPatternOnly = true; // 无外部证据时是否允许纯 pattern 判 NZ - std::optional preferredMinor2D; // 现有 ND/DN 歧义消解入口 -}; +`hasNZInnerStructure5D()` 只描述不能被切开的内部 fractal： -std::optional inferLayout5D(ArrayRef shape5D, - ArrayRef stride5D, - unsigned elemBytes, - const LayoutInferOptions &opts, - bool *isMinor2DAmbiguous = nullptr); -} // namespace mlir::pto +```cpp +shape[3] == 16 +shape[4] == C0 +stride[4] == 1 +stride[3] == C0 +stride[2] == 16 * C0 ``` -`InferPTOLayout.cpp` / `PTO.cpp` / `PTOToEmitC.cpp` 三处改为调用同一实现， -删除各自的私有副本。这一步是纯重构（NFC），单独成 PR，便于回归定位。 - -### 4.2 `isNZCompatible5D`：结构必要条件 +`isNZViewCompatible5D()` 在内部结构之上验证当前 view 的外层跨度： ```cpp -bool isNZCompatible5D(shape, stride, elemBytes) { - auto c0 = getNZC0Elems(elemBytes); // 32 % elemBytes != 0 -> false - if (!c0) return false; - return shape[3] == 16 && shape[4] == *c0 - && stride[4] == 1 - && stride[3] == *c0 - && stride[2] == 16 * *c0 - && stride[1] == shape[2] * stride[2] - && stride[0] == shape[1] * stride[1]; -} +hasNZInnerStructure5D(shape, stride, bytes) + && shape[0] == 1 + && stride[1] >= shape[2] * stride[2] + && stride[1] % C0 == 0 + && stride[0] >= shape[1] * stride[1] + && stride[0] % C0 == 0 ``` -用途有两个，且**两个用途门槛不同**（这是本设计的核心）： +它允许 `stride[0]` / `stride[1]` 包含按 C0 对齐的外层 gap，但不允许负 gap、重叠 +或破坏 32B block 对齐。`isCanonicalNZRoot5D()` 等价于 +`isNZViewCompatible5D()` 再要求两个外层 stride 都等于当前 shape 的紧密跨度。 + +不能使用根视图的外层紧密 stride 等式验证 partition 结果。合法子视图缩小 `d1` +或 `d2` 后会保留父 stride，此时 `stride[0]` 或 `stride[1]` 可以大于当前子视图的 +紧密跨度；pto-isa 的 NZ load/store 使用这个差值表示外层 gap。 -- 校验显式 `layout = nz` 是否自洽 —— 只看这个谓词； -- 无显式 layout 时的推断 —— 还要过 4.3 的门槛。 +### 4.3 无属性根视图的 NZ 约定 -### 4.3 推断门槛：只认"不可能是二维 ND 视图"的形状 +无显式 layout 且没有可继承源 layout 时，按以下唯一规则推断 NZ： ```cpp -bool inferNZFromPattern(shape, stride, elemBytes) { - return isNZCompatible5D(shape, stride, elemBytes) - && shape[0] == 1 // 二维 NZ，batch 维恒为 1 - && (shape[1] > 1 || shape[2] > 1); // 至少有一个分块维非退化 -} +isCanonicalNZRoot5D(shape, stride, storageElemBytes) + && shape[0] == 1 + && (shape[1] > 1 || shape[2] > 1) ``` -两个附加条件的理由： +`[1,1,1,16,C0]` 单 fractal 不自动升级为 NZ，因为它是 rank2 +`[16,C0]` ND 视图规范化后的常见形状。用户确实需要单 fractal NZ 时使用显式 +`layout = nz`。 -- `shape[0] == 1`：issue 要求的是**二维** NZ；`batch > 1` 的形状在 pto-isa 侧没有 - 对应的 `TileShape2D`，留给显式标注，避免把连续五维 ND 张量整片吃掉。 -- `shape[1] > 1 || shape[2] > 1`：`[1,1,1,16,C0]` 是 rank2 视图规范化后的标准产物 - （`PTOCanonicalizeIR`），也是最常见的 `16 x C0` 向量 tile。此时 NZ 与 ND 的字节 - 排布**完全等价**（单 fractal），判成 NZ 没有任何收益，却会改变生成的 C++ 模板参数、 - 进而影响 pto-isa 侧的重载选择。实测这条门槛消除了 100% 的新增误判（第 5 节）。 +显式 `layout = nd` 必须覆盖上述 pattern 约定。例如 fp32： -保留 `LayoutInferOptions::allowNZFromPatternOnly`，是为了给"上层已经有更强证据"的 -调用点（例如 mgather/mscatter 的 ND-only 校验路径）一个关掉 pattern 推断的开关。 +```text +shape = [1, 2, 1, 16, 8] +stride = [256, 128, 128, 8, 1] +layout = nd +``` -### 4.4 显式 layout 优先 +该 shape/stride 同时符合 canonical NZ 数值形式和连续 5D ND，PTOAS 必须保留用户 +声明的 ND，不得报 layout mismatch。 -`verifyOrSetLayoutAttr()` 的语义调整为： +### 4.4 显式 layout 校验 -| 显式属性 | 结构自洽性 | 现行为 | 新行为 | -|---|---|---|---| -| `nz` | `isNZCompatible5D == true` | **报错** | 接受，保留 `nz`，不打 `pto.inferred_layout` | -| `nz` | `isNZCompatible5D == false` | 报错 | 报错，但错误信息升级为"哪一条不变量不满足" | -| `nd`/`dn` | 与推断不同 | 现有 minor-2D 歧义豁免 | 不变 | +显式属性不再与 pattern 推断结果比较，而是按其自身语义校验： -这条直接解决 2.3 的第三个症状：即使 pattern 推断因门槛保守而不升级 NZ，用户/前端 -也永远有一条显式通道。诊断信息形如： +- 显式 `NZ` 根视图必须满足 NZ 内部 fractal 结构；canonical 紧密 stride 和合法的 + 外层 block gap 都可接受，统一使用 `isNZViewCompatible5D()` 校验； +- 显式 `ND` / `DN` 沿用现有各自的布局校验和 minor-2D 歧义处理； +- 校验失败时报告具体的 shape/stride 不变量，不回退到其他 layout； +- 显式属性不附加 `pto.inferred_layout`。 -``` -error: layout mismatch: user-specified layout=nz but shape/stride is not an NZ - 5D form: expected shape[3]==16 (got 8), shape[4]==8 (C0 for f32, got 16) -``` +### 4.5 NZ partition/subview 传播 -### 4.5 `partition_view` / `memref.subview` 的 NZ 传播 +源 layout 为 NZ 时，`partition_view` 和 `memref.subview` 按源 view 校验： -现状：`InferPTOLayout.cpp` 的 subview 分支无条件继承源 layout。对 NZ 这是不安全的 -—— 在 fractal 内部切分后，结果已经不是合法 NZ。规则改为： +- `d3`、`d4` 必须完整保留：offset 为 0，size 等于源 size； +- `d1`、`d2` 可以缩小，结果保留源 stride； +- `d1`、`d2` 的一个坐标单位本身就是一个完整 block，因此任意合法整数 offset + 都是 block 边界，不再增加含义不清的二次对齐条件； +- 二维 NZ 的 `d0` 保持 `offset = 0`、`size = 1`； +- 校验通过后直接继承 NZ，不再用 `isCanonicalNZRoot5D()` 重新推断； +- 在 `d3`/`d4` 内部切分时发出错误，不静默改成 ND。 -1. 源为 NZ 时，只有当切分满足以下条件才继承 NZ： - - `d3`/`d4` 维保持完整（offset 为 0、size 等于源 size）； - - `d1`/`d2` 维的 size 可缩小，但 stride 保持不变； - - offset 在 `d1`/`d2` 上按整块对齐。 -2. 不满足时不 silently 退回 ND，而是发 `emitError`（"NZ view cannot be partitioned - inside a fractal"），避免生成静默错码。 +动态 `d1`/`d2` offset 和 size 不改变 fractal 内部结构，可以传播 NZ。`d3`/`d4` +是否完整必须能够在编译期证明，否则报错。 -issue 明确提到 `partition_tensor_view`，这条属于本次范围内。 +### 4.6 动态 shape/stride -### 4.6 动态形状（阶段二） +显式或继承的 layout 在动态 shape 下继续携带。校验所有编译期已知的结构维，不使用 +未知值否决显式 NZ。 -pto-isa 的 `TileShape2D`/`BaseShape2D` 允许 `rows`/`cols` 为 `DYNAMIC`，此时 -`shape[1]`、`shape[2]`、`stride[0]`、`stride[1]` 动态，而结构维 -（`shape[3]`、`shape[4]`、`stride[2..4]`）仍是静态常量。因此可以在 -"结构维静态 + 分块维动态"时仍判定 NZ。当前实现要求全部 const-fold，直接放弃推断 -（`getStaticShapeAndStride()` 返回 false），属于可独立推进的增强，放到阶段二。 +无属性根视图只有在 NZ 内部结构和非退化门槛都能在编译期证明时才自动推断 NZ；无法 +证明时不猜测，沿用 ND/DN fallback，调用方应提供显式 layout。 -### 4.7 与旧规则的兼容策略 +### 4.7 旧规则处理 -| 方案 | 说明 | 评价 | -|---|---|---| -| A. 直接替换 | 只保留标准形规则 | **推荐**。旧规则命中的非标准形本质是误判，继续保留会持续生成错误的 `GlobalTensor` 模板参数 | -| B. 并集（新规则 OR 旧规则） | 兼容一切现有行为 | 不推荐：把"连续 ND 判成 NZ"固化成契约 | -| C. 替换 + 一个版本的过渡开关 | 加 `--pto-legacy-nz-infer`（默认关） | 若下游有存量依赖再启用；本次实测影响面极小（第 5 节），倾向不引入 | +直接删除旧的错位 NZ 判定，只保留本设计定义的标准规则。不保留旧规则并集，不增加 +legacy 开关。旧规则命中的非标准形属于误判，继续兼容会固定错误的 +`GlobalTensor<..., Layout::NZ>` 模板参数。 -推荐 A，若 review 中出现下游存量用例再降级到 C。 +## 5. 影响面 -## 5. 影响面实测 +对 `test/lit/pto/*.pto` 和 `test/lit/tile_fusion/*.pto` 的现有样本扫描得到： -方法：用当前 `ptoas` 对 `test/lit/pto/*.pto`（476）+ `test/lit/tile_fusion/*.pto`（42） -分别按 `--pto-arch=a5` 和 `--pto-arch=a3` 跑一遍，抽取全部生成的 -`GlobalTensor, Stride<...>, Layout::X>`（含 `using GTShape_*` 别名形式）， -再离线对同一批 shape/stride 分别套用旧规则、标准形规则、带门槛的标准形规则。 +- 242 个文件生成 1558 个 `GlobalTensor` 实例化点； +- 当前 NZ 去重形状中，两组符合 pto-isa 标准形； +- 两组只存在于 `globaltensor_layout_bytewidth_emitc.pto` 的旧错位形状将改回 ND； +- 不加非退化门槛时，38 个文件中的 80 个单-fractal ND 视图会被升级成 NZ； +- 加门槛后，当前语料没有新增 NZ。 -样本：242 个文件产生输出，共 **1558** 个 `GlobalTensor` 实例化点， -**64** 组去重后的 `(elem, shape, stride)`。当前 layout 分布：ND 1007 / NZ 529 / DN 22。 +“当前语料没有新增 NZ”仅是回归影响数据，不代表 shape/stride 已经从理论上完成 ND/NZ +消歧。因此实现仍必须保证所有显式 layout 优先，并增加“canonical 数值形式 + +显式 ND”的反例测试。 -| 项 | 结果 | -|---|---| -| 当前判为 NZ 的去重形状 | 5 组 | -| 其中符合 pto-isa 标准形（改后仍为 NZ） | 2 组：`half`/`bfloat16_t` `[1,1,16,16,16]` `stride=[4096,4096,256,16,1]` | -| 其中由显式属性设置（不走推断，不受影响） | 1 组：`tinsert_a5_vec_mat_mode_lowering` 的 `half [1,1,1,32,32]` | -| **改后会失去 NZ** | 2 组：`int8_t [1,1,16,32,16]`、`int64_t [1,1,16,4,1]` —— 全部只出现在 `test/lit/pto/globaltensor_layout_bytewidth_emitc.pto` | -| 不带门槛时新增 NZ | 4 组 / 80 个点 / 38 个文件，**全部**是 `[1,1,1,16,C0]` 退化单 fractal（`float`/`half`/`float8_e4m3_t`/`hifloat8_t`） | -| **带门槛（4.3）新增 NZ** | **0** | +`mgather`/`mscatter` 的 ND-only 校验也必须使用同一解析顺序：显式/继承 layout +优先，缺失时才 pattern 推断。真实 NZ 仍按现有约束报错；显式 ND 不得被 heuristic +改成 NZ。 -两点解读： +## 6. 本 PR 的实现范围 -1. 门槛条款把误判从 80 个点压到 0，验证了 P2 的必要性。 -2. 唯一需要改期望值的测试是 `globaltensor_layout_bytewidth_emitc.pto` - （`comm.tbroadcast` 的 i8/i64 用例）。这两组形状按 pto-isa 定义本来就不是 NZ - （i8 的 `C0 = 32`，标准形应为 `[1, cols/32, rows/16, 16, 32]`），当前期望值本身 - 固化了 2.2 的错位 bug。改测试前需确认 `comm.tbroadcast` 侧对 `Layout` 是否敏感。 +PR #1027 在现有设计提交之后继续完成以下内容： -另需交叉验证：`PTO.cpp:4366` 的 mgather/mscatter "mem partition view 必须是 ND"校验 -走的是同一套推断。规则收紧后，理论上存在"原本判 ND 的五维 mem view 变成 NZ 导致新 -报错"的可能；带门槛后本地 lit 语料未观察到，但需要在实现 PR 里补一条针对性用例。 +- 新增共享 `PTOLayoutUtils`； +- 同步 IR verifier、layout pass 和 EmitC fallback； +- 实现显式 layout 全优先； +- 替换旧 NZ 规则； +- 实现 NZ partition/subview 的 source-relative 校验和传播； +- 补齐静态、动态、packed storage element 和跨调用点一致性测试； +- 更新用户文档及 Python `make_tensor_view` 的显式 layout 示例。 -## 6. 实施拆分 +不再拆分成多个后续 PR。上述实现和验证全部完成后将 #1027 转为 ready；合入该 PR +即关闭 #527。 -| PR | 内容 | 风险 | -|---|---|---| -| P1 | 抽出 `PTOLayoutUtils`，三处调用点收敛，行为完全不变（NFC） | 低 | -| P2 | 标准形规则 + 4.3 门槛 + 4.7 方案 A；更新 `globaltensor_layout_bytewidth_emitc.pto` | 中 | -| P3 | 4.4 显式 layout 优先 + 诊断信息细化 | 低 | -| P4 | 4.5 partition/subview NZ 传播校验 | 中 | -| P5 | 4.6 动态形状支持（可选） | 中 | - -跨层同步检查（按 `.claude/rules/cross-layer-sync.md`）：本设计不改 ODS 算子签名， -`PTO_LayoutAttr` 已存在 `nz`；需要同步的是 IR verifier（`PTO.cpp`）、pass -（`InferPTOLayout.cpp`）、EmitC（`PTOToEmitC.cpp`）、以及 python 侧 -`make_tensor_view` 的 `layout` 传参样例与文档。 - -## 7. 测试计划 - -新增 `test/lit/pto/issue527_nz_canonical_view_infer.pto`（a5 + a3 双 RUN），覆盖： - -1. **正例**：canonical `[1,8,8,16,8]` fp32 → `pto::Layout::NZ`； -2. **正例**：`half` `[1,4,8,16,16]`、`int8_t` `[1,2,8,16,32]` → NZ（验证 C0 随 dtype 变化）； -3. **负例**：issue 用例② `[4,1,16,8,16]` → `ND`（旧规则误判被修掉）； -4. **负例**：`[1,1,1,16,8]` 单 fractal → 保持 `ND`（门槛条款，防回归 80 个点）； -5. **负例**：stride 被打断（如 `stride[1]` 非 `shape[2]*stride[2]`）→ `ND`； -6. **显式通道**：canonical 形 + `{layout = #pto.layout}` → 不再报 - `layout mismatch`，且输出 NZ； -7. **显式冲突**：`[4,1,16,8,16]` + `{layout = #pto.layout}` → 报错且信息指明 - 是哪条不变量不满足； -8. **partition**：对 NZ 视图在 `d1` 上整块切分 → 继承 NZ；在 `d3` 内部切分 → 报错。 - -E2E：`test/tilelang_st/npu/a5` 下补一条 NZ store 用例（若板卡资源允许）， -用 `nz_store_probe.py` 的逻辑做数据比对，确认生成的 C++ 与 pto-isa 语义一致。 - -## 8. 待确认问题 - -1. **是否接受 4.7 方案 A**（直接替换旧规则）？影响面只有一个 lit 测试的 2 组期望值。 -2. **`batch > 1` 的 NZ**（`shape[0] > 1`）是否需要推断？当前设计只在显式标注时接受。 -3. **sub-byte / packed 类型**（`int4`、`float4_e1m2x2`）：`32 % elemBytes` 语义如何定义？ - 当前设计直接不推断 NZ，需 pto-isa 侧确认是否存在这类 NZ 用法。 -4. `comm.tbroadcast` 路径对 `GlobalTensor` 的 `Layout` 模板参数是否敏感？ - 决定 `globaltensor_layout_bytewidth_emitc.pto` 的期望值怎么改。 -5. `mgather`/`mscatter` 的 "mem 必须 ND" 约束，在 mem 确实是 NZ 五维时是否应放开， - 还是维持报错（本设计维持现状，仅补测试）。 +## 7. 测试要求 + +新增 `test/lit/pto/issue527_nz_canonical_view_infer.pto`，至少覆盖： + +1. fp32 canonical `[1,8,8,16,8]` 自动推断 NZ； +2. half `[1,4,8,16,16]`、int8 `[1,2,8,16,32]` 自动推断 NZ； +3. 旧错位形状 `[4,1,16,8,16]` 推断 ND； +4. 单 fractal `[1,1,1,16,C0]` 无属性时保持 ND； +5. 单 fractal 显式 NZ 时保留 NZ； +6. canonical 数值形式 `[1,2,1,16,8]` 显式 ND 时保留 ND； +7. canonical NZ 显式 NZ 时不再出现 layout mismatch； +8. 非法 NZ 内部 shape/stride 显式 NZ 时给出精确诊断； +9. `d1`、`d2` 分别缩小的 partition 保留父 stride 并继承 NZ； +10. `d3` 或 `d4` 内部切分时报错； +11. 嵌套 partition 继续从已经解析的源 layout 传播； +12. FP4 packed-pair 按一字节 storage element 计算 C0； +13. verifier、pass 和 EmitC 对同一 view 得到一致 layout； +14. mgather/mscatter 的 ND-only 路径不覆盖显式 ND。 + +同时更新 `globaltensor_layout_bytewidth_emitc.pto` 中固化旧错位规则的期望值，并增加 +A3/A5 编译覆盖。具备板卡资源时运行一条 NZ load/store E2E 数据比对，确认生成的 +`GlobalTensor` shape、stride、layout 与 pto-isa 行为一致。 From adc4a91f5fcf4e54b6e9515ade06b69f93cbdeda Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Wed, 29 Jul 2026 09:18:47 +0800 Subject: [PATCH 3/6] docs: add NZ layout resolution examples --- docs/designs/nz-layout-inference-alignment.md | 122 +++++++++++++++++- 1 file changed, 119 insertions(+), 3 deletions(-) diff --git a/docs/designs/nz-layout-inference-alignment.md b/docs/designs/nz-layout-inference-alignment.md index 5ce41060e..72377eb06 100644 --- a/docs/designs/nz-layout-inference-alignment.md +++ b/docs/designs/nz-layout-inference-alignment.md @@ -248,7 +248,123 @@ layout = nd legacy 开关。旧规则命中的非标准形属于误判，继续兼容会固定错误的 `GlobalTensor<..., Layout::NZ>` 模板参数。 -## 5. 影响面 +## 5. 示例 + +以下示例省略 `arith.constant` 定义，常量名与数值一致。 + +### 5.1 canonical NZ 自动推断 + +输入是无显式 layout 的 fp32 canonical NZ 根视图： + +```mlir +%view = pto.make_tensor_view %dst, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + : !pto.tensor_view<1x8x8x16x8xf32> + +%part = pto.partition_view %view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c8, %c8, %c16, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x8x8x16x8xf32> +``` + +解析顺序中没有显式属性和源 layout，因此根视图进入 pattern 推断。它满足 +`isCanonicalNZRoot5D()` 和非退化门槛，预期生成： + +```cpp +GlobalTensor< + float, + pto::Shape<1, 8, 8, 16, 8>, + pto::Stride<8192, 1024, 128, 8, 1>, + pto::Layout::NZ> +``` + +`make_tensor_view` 会得到 `layout = #pto.layout` 和 +`pto.inferred_layout = true`，full-size partition 继承 NZ。 + +### 5.2 相同数值形式由显式 ND 覆盖 + +下面的连续 5D ND 在 shape/stride 数值上也满足 canonical NZ pattern，但用户显式 +声明其逻辑语义为 ND： + +```mlir +%view = pto.make_tensor_view %dst, + shape = [%c1, %c2, %c1, %c16, %c8], + strides = [%c256, %c128, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x2x1x16x8xf32> +``` + +显式属性在 pattern 推断之前解析，预期生成： + +```cpp +GlobalTensor< + float, + pto::Shape<1, 2, 1, 16, 8>, + pto::Stride<256, 128, 128, 8, 1>, + pto::Layout::ND> +``` + +该 view 不附加 `pto.inferred_layout`，也不得出现 +`layout mismatch: user-specified layout=nd but inferred=nz`。 + +### 5.3 NZ partition 的合法和非法切分 + +以 5.1 的 `%view` 为源，在 `d1` 上从第 2 个列块开始取 4 个块： + +```mlir +%d1_part = pto.partition_view %view, + offsets = [%c0, %c2, %c0, %c0, %c0], + sizes = [%c1, %c4, %c8, %c16, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x4x8x16x8xf32> +``` + +结果基址前移 `2 * stride[1] = 2048` 个 fp32 storage element，保留父 stride： + +```cpp +GlobalTensor< + float, + pto::Shape<1, 4, 8, 16, 8>, + pto::Stride<8192, 1024, 128, 8, 1>, + pto::Layout::NZ> +``` + +此时 `stride[0] != shape[1] * stride[1]`，但这是合法的外层 gap，不能用 canonical +根视图等式把它拒绝或重新推断成 ND。 + +同理，在 `d2` 上取 4 个行块也合法： + +```mlir +%d2_part = pto.partition_view %view, + offsets = [%c0, %c0, %c2, %c0, %c0], + sizes = [%c1, %c8, %c4, %c16, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x8x4x16x8xf32> +``` + +它的基址前移 `2 * stride[2] = 256` 个 storage element，stride 仍为 +`[8192,1024,128,8,1]`，layout 继续是 NZ。 + +在 `d3` 或 `d4` 内切分会破坏一个完整 fractal。例如： + +```mlir +%bad = pto.partition_view %view, + offsets = [%c0, %c0, %c0, %c1, %c0], + sizes = [%c1, %c8, %c8, %c15, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x8x8x15x8xf32> +``` + +预期在 verifier/layout pass 阶段失败，不生成 ND fallback： + +```text +error: NZ view cannot be partitioned inside a fractal: + d3 must keep offset=0 and size=16 (got offset=1, size=15) +``` + +## 6. 影响面 对 `test/lit/pto/*.pto` 和 `test/lit/tile_fusion/*.pto` 的现有样本扫描得到： @@ -266,7 +382,7 @@ legacy 开关。旧规则命中的非标准形属于误判，继续兼容会固 优先，缺失时才 pattern 推断。真实 NZ 仍按现有约束报错；显式 ND 不得被 heuristic 改成 NZ。 -## 6. 本 PR 的实现范围 +## 7. 本 PR 的实现范围 PR #1027 在现有设计提交之后继续完成以下内容： @@ -281,7 +397,7 @@ PR #1027 在现有设计提交之后继续完成以下内容： 不再拆分成多个后续 PR。上述实现和验证全部完成后将 #1027 转为 ready；合入该 PR 即关闭 #527。 -## 7. 测试要求 +## 8. 测试要求 新增 `test/lit/pto/issue527_nz_canonical_view_infer.pto`，至少覆盖： From 3c03c76a490b1a09e4748e38ba0ea72c382cfdf1 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Wed, 29 Jul 2026 10:47:23 +0800 Subject: [PATCH 4/6] feat: align NZ layout inference with pto-isa --- docs/PTO_IR_manual.md | 43 +- docs/designs/nz-layout-inference-alignment.md | 24 +- include/PTO/IR/PTOLayoutUtils.h | 57 +++ include/PTO/IR/PTOOps.td | 6 +- lib/PTO/IR/CMakeLists.txt | 1 + lib/PTO/IR/PTO.cpp | 243 +++++---- lib/PTO/IR/PTOLayoutUtils.cpp | 309 ++++++++++++ lib/PTO/Transforms/InferPTOLayout.cpp | 436 +++++++++------- lib/PTO/Transforms/PTOToEmitC.cpp | 25 +- lib/PTO/Transforms/PTOViewToMemref.cpp | 9 + ptodsl/docs/user_guide/01-introduction.md | 25 + .../globaltensor_layout_bytewidth_emitc.pto | 5 +- .../pto/issue527_nz_canonical_view_infer.pto | 472 ++++++++++++++++++ 13 files changed, 1329 insertions(+), 326 deletions(-) create mode 100644 include/PTO/IR/PTOLayoutUtils.h create mode 100644 lib/PTO/IR/PTOLayoutUtils.cpp create mode 100644 test/lit/pto/issue527_nz_canonical_view_infer.pto diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index 6903ff9c5..fc159f7d3 100644 --- a/docs/PTO_IR_manual.md +++ b/docs/PTO_IR_manual.md @@ -489,6 +489,19 @@ Global tensor layout inference for [`tensor_view` (Section 2.3)](#23-ptotensor_v **Attribute syntax:** `#pto.layout` +Resolution follows one order: an explicit `layout` is authoritative; a +partition inherits its source layout; only an unannotated root view is +pattern-inferred. For a storage element of `sizeof(T)` bytes, canonical NZ uses +`C0 = 32 / sizeof(T)` and: + +```text +shape = [1, n1, m1, 16, C0] +stride = [n1*m1*16*C0, m1*16*C0, 16*C0, C0, 1] +``` + +An unannotated single-fractal `[1,1,1,16,C0]` view remains ND. Use an explicit +`#pto.layout` when its logical meaning is NZ. + --- ## 4. Operations Reference @@ -678,7 +691,7 @@ This operation defines the physical "base" and stride rules for global memory. I | `ptr` | `AnyType` | Source pointer | | `shape` | `Variadic` | Dynamic shape dimensions | | `strides` | `Variadic` | Dynamic strides | -| `layout` | `LayoutAttr` (optional) | ND/DN/NZ layout hint | +| `layout` | `LayoutAttr` (optional) | Authoritative ND/DN/NZ logical layout | **Results:** `!pto.tensor_view<...>` @@ -687,14 +700,18 @@ This operation defines the physical "base" and stride rules for global memory. I - The operation has a custom verifier that checks: - `ptr` must be `!pto.ptr<...>` and its element type must match the result element type - `shape` and `strides` operand counts must match the tensor_view rank - - If `layout` is provided with static shapes/strides, it must be consistent with inferred layout + - Explicit ND/DN are preserved and are not replaced by shape/stride pattern + inference + - Explicit NZ must keep rank 5, `shape[0] == 1`, the complete + `[16,C0]` inner fractal, and compatible C0-aligned outer strides - `pto.inttoptr` results cannot feed `pto.make_tensor_view`. Tensor views must be constructed from a source pointer that already carries the desired element type. **Notes:** -- Stride patterns may allow the compiler to infer hardware layout hints (e.g., `layout = nz`) to guide later DMA operations. +- An unannotated non-degenerate canonical NZ five-dimensional shape is inferred + as NZ. Other unannotated views use the existing ND/DN fallback. **Hardware Mapping:** @@ -706,6 +723,16 @@ This operation defines the physical "base" and stride rules for global memory. I %tv = pto.make_tensor_view %ptr, shape = [%m, %n], strides = [%s0, %s1] : !pto.tensor_view ``` +For an explicit fp32 NZ view (`C0 = 8`): + +```mlir +%tv = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x8x8x16x8xf32> +``` + --- ##### `pto.get_tensor_view_dim` - Get Tensor View Dimension Size @@ -784,7 +811,8 @@ in elements rather than bytes. ##### `pto.partition_view` - Partition Tensor View -**Summary:** Creates a logical window on a tensor_view using offsets and sizes, producing a `partition_tensor_view`. +**Summary:** Creates a logical window on a tensor_view or partition_tensor_view +using offsets and sizes, producing a `partition_tensor_view`. **Semantics:** @@ -798,7 +826,7 @@ This op captures both static and dynamic shapes. It represents a logical slice w | Name | Type | Description | |------|------|-------------| -| `source` | `TensorViewType` | Input tensor view | +| `source` | `TensorViewType` or `PartitionTensorViewType` | Input root or partition view | | `offsets` | `Variadic` | Dynamic offsets | | `sizes` | `Variadic` | Dynamic sizes | @@ -807,10 +835,15 @@ This op captures both static and dynamic shapes. It represents a logical slice w **Constraints & Verification:** - `offsets`/`sizes` counts must match the rank of `source` +- An NZ partition may shrink or dynamically offset/size only `d1` and `d2` +- NZ `d0`, `d3`, and `d4` must be statically proven complete; slicing inside + `d3`/`d4` is rejected instead of falling back to ND **Notes:** - Pointer arithmetic is modeled as `BasePtr + Offset`, and the logical shape is determined by `sizes`. +- Nested partitions inherit the already-resolved source layout and preserve the + root strides. **Hardware Mapping:** diff --git a/docs/designs/nz-layout-inference-alignment.md b/docs/designs/nz-layout-inference-alignment.md index 72377eb06..28b35948c 100644 --- a/docs/designs/nz-layout-inference-alignment.md +++ b/docs/designs/nz-layout-inference-alignment.md @@ -130,6 +130,21 @@ bool isCanonicalNZRoot5D(ArrayRef shape5D, ArrayRef stride5D, unsigned storageElemBytes); +std::optional +getNZViewCompatibilityError(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes); + +std::optional +getNZSubviewCompatibilityError(ArrayRef sourceShape5D, + ArrayRef offsets5D, + ArrayRef sizes5D); + +bool isLayoutCompatible5D(Layout layout, + ArrayRef shape, + ArrayRef stride, + unsigned storageElemBytes); + std::optional inferLayout5D(ArrayRef shape5D, ArrayRef stride5D, @@ -215,7 +230,9 @@ layout = nd - 显式 `NZ` 根视图必须满足 NZ 内部 fractal 结构；canonical 紧密 stride 和合法的 外层 block gap 都可接受，统一使用 `isNZViewCompatible5D()` 校验； -- 显式 `ND` / `DN` 沿用现有各自的布局校验和 minor-2D 歧义处理； +- 显式 `ND` / `DN` 保持为逻辑布局的权威声明。它们不承诺当前 view 是紧密 + minor-2D 存储，因此不使用 dense stride 递推拒绝已有的 interleaved/gapped + view；minor-2D 歧义处理只用于无显式 layout 的推断； - 校验失败时报告具体的 shape/stride 不变量，不回退到其他 layout； - 显式属性不附加 `pto.inferred_layout`。 @@ -229,6 +246,8 @@ layout = nd 都是 block 边界，不再增加含义不清的二次对齐条件； - 二维 NZ 的 `d0` 保持 `offset = 0`、`size = 1`； - 校验通过后直接继承 NZ，不再用 `isCanonicalNZRoot5D()` 重新推断； +- `partition_view` 可以继续以 `partition_tensor_view` 为源，嵌套切分逐级沿源链 + 继承同一个 layout； - 在 `d3`/`d4` 内部切分时发出错误，不静默改成 ND。 动态 `d1`/`d2` offset 和 size 不改变 fractal 内部结构，可以传播 NZ。`d3`/`d4` @@ -370,7 +389,8 @@ error: NZ view cannot be partitioned inside a fractal: - 242 个文件生成 1558 个 `GlobalTensor` 实例化点； - 当前 NZ 去重形状中，两组符合 pto-isa 标准形； -- 两组只存在于 `globaltensor_layout_bytewidth_emitc.pto` 的旧错位形状将改回 ND； +- 两组只存在于 `globaltensor_layout_bytewidth_emitc.pto` 的旧错位形状不再是 NZ： + int8 样例回到 ND，末维为 1 的 int64 样例按既有 minor-2D 规则回到 DN； - 不加非退化门槛时，38 个文件中的 80 个单-fractal ND 视图会被升级成 NZ； - 加门槛后，当前语料没有新增 NZ。 diff --git a/include/PTO/IR/PTOLayoutUtils.h b/include/PTO/IR/PTOLayoutUtils.h new file mode 100644 index 000000000..9ad44d8cc --- /dev/null +++ b/include/PTO/IR/PTOLayoutUtils.h @@ -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. + +#ifndef PTO_IR_PTOLAYOUTUTILS_H +#define PTO_IR_PTOLAYOUTUTILS_H + +#include "PTO/IR/PTO.h" +#include "llvm/ADT/ArrayRef.h" + +#include +#include + +namespace mlir::pto { + +inline constexpr unsigned kPTOLayoutRank = 5; + +std::optional getNZC0StorageElems(unsigned storageElemBytes); + +bool hasNZInnerStructure5D(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes); + +bool isNZViewCompatible5D(ArrayRef shape5D, ArrayRef stride5D, + unsigned storageElemBytes); + +bool isCanonicalNZRoot5D(ArrayRef shape5D, ArrayRef stride5D, + unsigned storageElemBytes); + +std::optional +getNZViewCompatibilityError(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes); + +std::optional +getNZSubviewCompatibilityError(ArrayRef sourceShape5D, + ArrayRef offsets5D, + ArrayRef sizes5D); + +bool isLayoutCompatible5D(Layout layout, ArrayRef shape, + ArrayRef stride, unsigned storageElemBytes); + +std::optional +inferLayout5D(ArrayRef shape, ArrayRef stride, + unsigned storageElemBytes, + std::optional preferredMinor2D = std::nullopt, + bool *isMinor2DAmbiguous = nullptr); + +} // namespace mlir::pto + +#endif // PTO_IR_PTOLAYOUTUTILS_H diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index f5f7938fc..b6f3641dd 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -50,6 +50,10 @@ def PtrOrMemRef : def TileBufOrMemRef : AnyTypeOf<[TileBufType, AnyMemRef], "TileBuf or MemRef">; +def TensorViewOrPartitionView : + AnyTypeOf<[TensorViewType, PartitionTensorViewType], + "TensorView or PartitionTensorView">; + def ScalarPtrOrMemRef : TypeConstraint< CPred<"::mlir::pto::isScalarPtrOrMemRef($_self)">, @@ -257,7 +261,7 @@ def PartitionViewOp : PTO_Op<"partition_view", [AttrSizedOperandSegments]> { }]; let arguments = (ins - TensorViewType:$source, // 输入: 物理大底座 (MakeTensorViewOp 的结果) + TensorViewOrPartitionView:$source, Variadic:$offsets, // 动态 offsets Variadic:$sizes // 动态 sizes ); diff --git a/lib/PTO/IR/CMakeLists.txt b/lib/PTO/IR/CMakeLists.txt index fe99fff0f..b942f89ce 100644 --- a/lib/PTO/IR/CMakeLists.txt +++ b/lib/PTO/IR/CMakeLists.txt @@ -18,6 +18,7 @@ add_mlir_dialect_library(PTOIR VMI.cpp VPTOUbOps.cpp PTOAttrs.cpp + PTOLayoutUtils.cpp PTOSyncUtils.cpp PTOTypeDefs.cpp PTOTypeUtils.cpp diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index af6dd2859..b15485b9b 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -10,9 +10,10 @@ //===----------------------------------------------------------------------===// #include "PTO/IR/PTO.h" +#include "PTO/IR/PTOLayoutUtils.h" #include "PTO/IR/PTOMultiBuffer.h" -#include "PTO/IR/PTOTypeUtils.h" #include "PTO/IR/PTOSyncUtils.h" +#include "PTO/IR/PTOTypeUtils.h" #include "mlir/AsmParser/AsmParser.h" #include "mlir/Conversion/LLVMCommon/TypeConverter.h" @@ -1478,12 +1479,21 @@ static std::optional getConstIndexValue(Value v) { } static FailureOr -inferPartitionViewResultTypeFromSizes(mlir::pto::TensorViewType sourceType, - ValueRange sizes) { - if (!sourceType) +inferPartitionViewResultTypeFromSizes(Type sourceType, ValueRange sizes) { + int64_t sourceRank = 0; + Type elementType; + if (auto tensorView = dyn_cast(sourceType)) { + sourceRank = tensorView.getRank(); + elementType = tensorView.getElementType(); + } else if (auto partitionView = + dyn_cast(sourceType)) { + sourceRank = partitionView.getRank(); + elementType = partitionView.getElementType(); + } else { return failure(); + } - if ((int64_t)sizes.size() != sourceType.getRank()) + if ((int64_t)sizes.size() != sourceRank) return failure(); SmallVector shape; @@ -1496,8 +1506,8 @@ inferPartitionViewResultTypeFromSizes(mlir::pto::TensorViewType sourceType, shape.push_back(ShapedType::kDynamic); } - return mlir::pto::PartitionTensorViewType::get( - sourceType.getContext(), shape, sourceType.getElementType()); + return mlir::pto::PartitionTensorViewType::get(sourceType.getContext(), shape, + elementType); } ParseResult mlir::pto::PartitionViewOp::parse(OpAsmParser &parser, @@ -1547,8 +1557,8 @@ ParseResult mlir::pto::PartitionViewOp::parse(OpAsmParser &parser, ValueRange allOperands(result.operands); ValueRange sizeOperands = allOperands.slice(1 + offsets.size(), sizes.size()); - auto inferredResultType = inferPartitionViewResultTypeFromSizes( - dyn_cast(sourceTy), sizeOperands); + auto inferredResultType = + inferPartitionViewResultTypeFromSizes(sourceTy, sizeOperands); if (failed(inferredResultType)) { return parser.emitError(parser.getCurrentLocation(), "failed to infer pto.partition_view result type"); @@ -1568,8 +1578,8 @@ void mlir::pto::PartitionViewOp::print(OpAsmPrinter &printer) { /*elidedAttrs=*/{"operandSegmentSizes"}); printer << " : " << getSource().getType(); - auto inferredResultType = inferPartitionViewResultTypeFromSizes( - dyn_cast(getSource().getType()), getSizes()); + auto inferredResultType = + inferPartitionViewResultTypeFromSizes(getSource().getType(), getSizes()); if (succeeded(inferredResultType) && *inferredResultType == getResult().getType()) return; @@ -1812,69 +1822,6 @@ static bool isSupportedGatherElemTypeA5(Type ty) { return false; } -static bool isStaticLayoutInt(int64_t value) { - return value != ShapedType::kDynamic && value >= 0; -} - -static std::optional multiplyLayoutInts(int64_t lhs, int64_t rhs) { - int64_t product = 0; - if (llvm::MulOverflow(lhs, rhs, product)) - return std::nullopt; - return product; -} - -static std::optional -inferLayout(ArrayRef shape, ArrayRef strides, - unsigned elemBytes) { - if (shape.size() != strides.size() || elemBytes == 0) - return std::nullopt; - if (llvm::any_of(shape, [](int64_t dim) { return !isStaticLayoutInt(dim); }) || - llvm::any_of(strides, - [](int64_t stride) { return !isStaticLayoutInt(stride); })) - return std::nullopt; - - // NZ / fractal: rank>=5, check middle dims (sh3/sh4/sh5 per spec) - if (shape.size() >= 5) { - int64_t sh3 = shape[2], sh4 = shape[3], sh5 = shape[4]; - int64_t st4 = strides[3], st5 = strides[4]; - auto sh3TimesSh4 = multiplyLayoutInts(sh3, sh4); - auto fractalBytes = - sh3TimesSh4 - ? multiplyLayoutInts(*sh3TimesSh4, static_cast(elemBytes)) - : std::nullopt; - bool alignMatch = (sh3 == 16) && fractalBytes && (*fractalBytes == 512); - bool strideMatch = (st5 == 1) && (st4 == sh5); - if (alignMatch && strideMatch) - return mlir::pto::Layout::NZ; - } - - // ND: row-major contiguous - bool isRowMajor = true; - for (int i = 0, e = (int)shape.size() - 1; i < e; ++i) { - auto expectedStride = multiplyLayoutInts(strides[i + 1], shape[i + 1]); - if (!expectedStride || strides[i] != *expectedStride) { - isRowMajor = false; - break; - } - } - if (isRowMajor && strides.back() == 1) - return mlir::pto::Layout::ND; - - // DN: col-major - bool isColMajor = true; - for (int i = 0, e = (int)shape.size() - 1; i < e; ++i) { - auto expectedStride = multiplyLayoutInts(strides[i], shape[i]); - if (!expectedStride || strides[i + 1] != *expectedStride) { - isColMajor = false; - break; - } - } - if (isColMajor && strides.front() == 1) - return mlir::pto::Layout::DN; - - return mlir::pto::Layout::ND; // fallback -} - static std::optional getLogicalViewLayout(Value value) { if (!value) return std::nullopt; @@ -1903,11 +1850,45 @@ static std::optional getLogicalViewLayout(Value value) { return std::nullopt; strides.push_back(*cst); } - return inferLayout(shape, strides, getElemByteSize(tvTy.getElementType())); + return inferLayout5D(shape, strides, + getElemByteSize(tvTy.getElementType())); } return std::nullopt; } +static bool getLogicalViewShape(Value value, SmallVectorImpl &shape) { + if (auto make = value.getDefiningOp()) { + auto type = dyn_cast(make.getResult().getType()); + if (!type) + return false; + shape.assign(type.getShape().begin(), type.getShape().end()); + for (auto [index, operand] : llvm::enumerate(make.getShape())) { + if (shape[index] == ShapedType::kDynamic) { + if (auto constant = getConstIndexValue(operand)) + shape[index] = *constant; + } + } + return true; + } + + if (auto partition = value.getDefiningOp()) { + auto type = + dyn_cast(partition.getResult().getType()); + if (!type) + return false; + shape.assign(type.getShape().begin(), type.getShape().end()); + for (auto [index, operand] : llvm::enumerate(partition.getSizes())) { + if (shape[index] == ShapedType::kDynamic) { + if (auto constant = getConstIndexValue(operand)) + shape[index] = *constant; + } + } + return true; + } + + return false; +} + static std::optional getTileBufLogicalLayout(pto::TileBufType type) { if (!type) return std::nullopt; @@ -2467,60 +2448,73 @@ LogicalResult mlir::pto::MakeTensorViewOp::verify() { return emitOpError() << "shape/strides operand counts must match tensor_view rank=" << rank; - // Detect dynamic shape/stride. - bool hasDynamicShape = llvm::any_of(tvTy.getShape(), [](int64_t v) { - return v == ShapedType::kDynamic; - }); - bool hasDynamicStride = llvm::any_of(getStrides(), [](Value s) { - return !getConstIndexValue(s).has_value(); - }); - auto layoutAttr = getLayoutAttr(); - - // 1) Dynamic shape/stride without explicit layout: warn and keep going. - if ((hasDynamicShape || hasDynamicStride) && !layoutAttr) { + if (!layoutAttr) return success(); + + SmallVector shapeInts; + shapeInts.reserve(rank); + for (int64_t index = 0; index < rank; ++index) { + int64_t dim = tvTy.getShape()[index]; + if (dim == ShapedType::kDynamic) { + if (auto constant = getConstIndexValue(getShape()[index])) + dim = *constant; + } + shapeInts.push_back(dim); } - // 2) Static shape/stride with explicit layout: verify correctness. - bool allStaticStride = true; SmallVector strideInts; - strideInts.reserve(getStrides().size()); + strideInts.reserve(rank); for (Value s : getStrides()) { auto val = getConstIndexValue(s); - if (!val) { - allStaticStride = false; - break; - } - strideInts.push_back(*val); + strideInts.push_back(val.value_or(ShapedType::kDynamic)); } - bool allStaticShape = - llvm::none_of(tvTy.getShape(), [](int64_t v) { return v == ShapedType::kDynamic; }); - - if (layoutAttr && allStaticShape && allStaticStride) { - SmallVector shapeInts(tvTy.getShape().begin(), tvTy.getShape().end()); - if (auto inferred = inferLayout(shapeInts, strideInts, - getElemByteSize(tvTy.getElementType()))) { - (void)inferred; + Layout layout = layoutAttr.getLayout(); + unsigned storageElemBytes = getElemByteSize(tvTy.getElementType()); + if (!isLayoutCompatible5D(layout, shapeInts, strideInts, storageElemBytes)) { + if (layout == Layout::NZ) { + if (rank != static_cast(kPTOLayoutRank)) + return emitOpError() + << "user-specified layout=nz requires a rank-5 view, got rank " + << rank; + auto error = + getNZViewCompatibilityError(shapeInts, strideInts, storageElemBytes); + return emitOpError() + << "user-specified layout=nz is incompatible with shape/stride: " + << error.value_or("unknown NZ layout mismatch"); } + return emitOpError() << "user-specified layout=" << stringifyLayout(layout) + << " is incompatible with the view shape/stride"; } return success(); } LogicalResult mlir::pto::PartitionViewOp::verify() { - auto srcTy = dyn_cast(getSource().getType()); auto resTy = dyn_cast(getResult().getType()); - if (!srcTy || !resTy) - return emitOpError("expects tensor_view source and partition_tensor_view result"); + if (!resTy) + return emitOpError("expects partition_tensor_view result"); + + Type srcElementType; + int64_t srcRank = 0; + if (auto tensorView = + dyn_cast(getSource().getType())) { + srcElementType = tensorView.getElementType(); + srcRank = tensorView.getRank(); + } else if (auto partitionView = dyn_cast( + getSource().getType())) { + srcElementType = partitionView.getElementType(); + srcRank = partitionView.getRank(); + } else { + return emitOpError("expects tensor_view or partition_tensor_view source"); + } - if (srcTy.getElementType() != resTy.getElementType()) - return emitOpError() << "element type mismatch between source and result: src=" - << srcTy.getElementType() << " result=" - << resTy.getElementType(); + if (srcElementType != resTy.getElementType()) + return emitOpError() + << "element type mismatch between source and result: src=" + << srcElementType << " result=" << resTy.getElementType(); - int64_t srcRank = srcTy.getRank(); if ((int64_t)getOffsets().size() != srcRank) return emitOpError() << "offset count (" << getOffsets().size() << ") must match source rank (" << srcRank << ")"; @@ -2529,7 +2523,20 @@ LogicalResult mlir::pto::PartitionViewOp::verify() { return emitOpError() << "size count (" << getSizes().size() << ") must match source rank (" << srcRank << ")"; - ArrayRef srcShape = srcTy.getShape(); + SmallVector logicalSourceShape; + if (!getLogicalViewShape(getSource(), logicalSourceShape)) { + if (auto tensorView = + dyn_cast(getSource().getType())) + logicalSourceShape.assign(tensorView.getShape().begin(), + tensorView.getShape().end()); + else { + auto partitionView = + cast(getSource().getType()); + logicalSourceShape.assign(partitionView.getShape().begin(), + partitionView.getShape().end()); + } + } + ArrayRef srcShape = logicalSourceShape; ArrayRef resShape = resTy.getShape(); bool sameRank = resTy.getRank() == srcRank; @@ -2567,6 +2574,20 @@ LogicalResult mlir::pto::PartitionViewOp::verify() { << ") exceeds static source dim (" << srcDim << ")"; } + if (getLogicalViewLayout(getSource()) == Layout::NZ) { + SmallVector offsets; + SmallVector sizes; + offsets.reserve(srcRank); + sizes.reserve(srcRank); + for (Value offset : getOffsets()) + offsets.push_back( + getConstIndexValue(offset).value_or(ShapedType::kDynamic)); + for (Value size : getSizes()) + sizes.push_back(getConstIndexValue(size).value_or(ShapedType::kDynamic)); + if (auto error = getNZSubviewCompatibilityError(srcShape, offsets, sizes)) + return emitOpError(*error); + } + return success(); } diff --git a/lib/PTO/IR/PTOLayoutUtils.cpp b/lib/PTO/IR/PTOLayoutUtils.cpp new file mode 100644 index 000000000..d52b44919 --- /dev/null +++ b/lib/PTO/IR/PTOLayoutUtils.cpp @@ -0,0 +1,309 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under +// the terms and conditions of CANN Open Software License Agreement Version 2.0 +// (the "License"). Please refer to the License for details. You may not use +// this file except in compliance with the License. THIS SOFTWARE IS PROVIDED ON +// AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS +// FOR A PARTICULAR PURPOSE. See LICENSE in the root of the software repository +// for the full text of the License. + +#include "PTO/IR/PTOLayoutUtils.h" + +#include "mlir/IR/BuiltinTypes.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/Twine.h" +#include "llvm/Support/MathExtras.h" + +using namespace mlir; +using namespace mlir::pto; + +namespace { + +constexpr int64_t kNZBlockBytes = 32; +constexpr int64_t kNZInnerRows = 16; + +struct ShapeStride5D { + SmallVector shape; + SmallVector stride; +}; + +static bool isDynamic(int64_t value) { return value == ShapedType::kDynamic; } + +static bool isKnownInvalid(int64_t value) { + return value < 0 && !isDynamic(value); +} + +static bool matchesKnown(int64_t actual, int64_t expected) { + return isDynamic(actual) || actual == expected; +} + +static std::optional checkedMultiply(int64_t lhs, int64_t rhs) { + if (isDynamic(lhs) || isDynamic(rhs)) + return ShapedType::kDynamic; + int64_t product = 0; + if (llvm::MulOverflow(lhs, rhs, product)) + return std::nullopt; + return product; +} + +static std::optional rightAlignTo5D(ArrayRef shape, + ArrayRef stride) { + if (shape.empty() || shape.size() != stride.size() || + shape.size() > kPTOLayoutRank) + return std::nullopt; + if (llvm::any_of(shape, isKnownInvalid) || + llvm::any_of(stride, isKnownInvalid)) + return std::nullopt; + + ShapeStride5D result; + result.shape.assign(kPTOLayoutRank, 1); + result.stride.assign(kPTOLayoutRank, 1); + + const int64_t shift = kPTOLayoutRank - shape.size(); + for (auto [index, value] : llvm::enumerate(shape)) + result.shape[shift + index] = value; + for (auto [index, value] : llvm::enumerate(stride)) + result.stride[shift + index] = value; + + for (int64_t index = shift - 1; index >= 0; --index) { + auto value = + checkedMultiply(result.shape[index + 1], result.stride[index + 1]); + if (!value) + return std::nullopt; + result.stride[index] = *value; + } + return result; +} + +static bool allStatic(ArrayRef values) { + return llvm::none_of(values, isDynamic); +} + +static bool matchesNDMinor2D(int64_t rows, int64_t cols, int64_t rowStride, + int64_t colStride) { + if (!isDynamic(cols) && cols != 1 && !isDynamic(colStride) && colStride != 1) + return false; + if (!isDynamic(rows) && rows == 1) + return true; + if (!isDynamic(cols) && cols == 1) + return isDynamic(rowStride) || rowStride == 1; + if (isDynamic(rowStride) || isDynamic(cols)) + return true; + return rowStride == cols; +} + +static bool matchesDNMinor2D(int64_t rows, int64_t cols, int64_t rowStride, + int64_t colStride) { + if (!isDynamic(rows) && rows != 1 && !isDynamic(rowStride) && rowStride != 1) + return false; + if (!isDynamic(cols) && cols == 1) + return true; + if (!isDynamic(rows) && rows == 1) + return isDynamic(colStride) || colStride == 1; + if (isDynamic(colStride) || isDynamic(rows)) + return true; + return colStride == rows; +} + +static std::optional +inferMinor2DLayout(int64_t rows, int64_t cols, int64_t rowStride, + int64_t colStride, std::optional preferredMinor2D, + bool *isMinor2DAmbiguous) { + const bool nd = matchesNDMinor2D(rows, cols, rowStride, colStride); + const bool dn = matchesDNMinor2D(rows, cols, rowStride, colStride); + if (!nd && !dn) + return Layout::ND; + if (nd && dn) { + if (isMinor2DAmbiguous) + *isMinor2DAmbiguous = true; + if (preferredMinor2D && + (*preferredMinor2D == Layout::ND || *preferredMinor2D == Layout::DN)) + return *preferredMinor2D; + return (!isDynamic(cols) && cols == 1 && (isDynamic(rows) || rows != 1)) + ? Layout::DN + : Layout::ND; + } + return dn ? Layout::DN : Layout::ND; +} + +static std::string expectedActual(StringRef name, int64_t expected, + int64_t actual) { + return ("expected " + name + "==" + Twine(expected) + " (got " + + Twine(actual) + ")") + .str(); +} + +} // namespace + +std::optional +mlir::pto::getNZC0StorageElems(unsigned storageElemBytes) { + if (storageElemBytes == 0 || kNZBlockBytes % storageElemBytes != 0) + return std::nullopt; + return kNZBlockBytes / storageElemBytes; +} + +bool mlir::pto::hasNZInnerStructure5D(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes) { + if (shape5D.size() != kPTOLayoutRank || stride5D.size() != kPTOLayoutRank) + return false; + auto c0 = getNZC0StorageElems(storageElemBytes); + if (!c0) + return false; + return matchesKnown(shape5D[3], kNZInnerRows) && + matchesKnown(shape5D[4], *c0) && matchesKnown(stride5D[4], 1) && + matchesKnown(stride5D[3], *c0) && + matchesKnown(stride5D[2], kNZInnerRows * *c0); +} + +std::optional +mlir::pto::getNZViewCompatibilityError(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes) { + if (shape5D.size() != kPTOLayoutRank || stride5D.size() != kPTOLayoutRank) + return "NZ layout requires rank-5 shape and stride"; + + auto c0 = getNZC0StorageElems(storageElemBytes); + if (!c0) + return ("NZ layout requires a non-zero storage element byte size that " + "divides 32, got " + + Twine(storageElemBytes)) + .str(); + + if (!matchesKnown(shape5D[0], 1)) + return expectedActual("shape[0]", 1, shape5D[0]); + if (!matchesKnown(shape5D[3], kNZInnerRows)) + return expectedActual("shape[3]", kNZInnerRows, shape5D[3]); + if (!matchesKnown(shape5D[4], *c0)) + return expectedActual("shape[4]", *c0, shape5D[4]); + if (!matchesKnown(stride5D[4], 1)) + return expectedActual("stride[4]", 1, stride5D[4]); + if (!matchesKnown(stride5D[3], *c0)) + return expectedActual("stride[3]", *c0, stride5D[3]); + if (!matchesKnown(stride5D[2], kNZInnerRows * *c0)) + return expectedActual("stride[2]", kNZInnerRows * *c0, stride5D[2]); + + if (!isDynamic(stride5D[1])) { + if (stride5D[1] % *c0 != 0) + return ("expected stride[1] to be C0-aligned (" + Twine(*c0) + "), got " + + Twine(stride5D[1])) + .str(); + auto minimum = checkedMultiply(shape5D[2], stride5D[2]); + if (!minimum) + return "shape[2] * stride[2] overflows"; + if (!isDynamic(*minimum) && stride5D[1] < *minimum) + return ("expected stride[1] >= shape[2] * stride[2] (" + Twine(*minimum) + + "), got " + Twine(stride5D[1])) + .str(); + } + + if (!isDynamic(stride5D[0])) { + if (stride5D[0] % *c0 != 0) + return ("expected stride[0] to be C0-aligned (" + Twine(*c0) + "), got " + + Twine(stride5D[0])) + .str(); + auto minimum = checkedMultiply(shape5D[1], stride5D[1]); + if (!minimum) + return "shape[1] * stride[1] overflows"; + if (!isDynamic(*minimum) && stride5D[0] < *minimum) + return ("expected stride[0] >= shape[1] * stride[1] (" + Twine(*minimum) + + "), got " + Twine(stride5D[0])) + .str(); + } + + return std::nullopt; +} + +bool mlir::pto::isNZViewCompatible5D(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes) { + return !getNZViewCompatibilityError(shape5D, stride5D, storageElemBytes); +} + +std::optional +mlir::pto::getNZSubviewCompatibilityError(ArrayRef sourceShape5D, + ArrayRef offsets5D, + ArrayRef sizes5D) { + if (sourceShape5D.size() != kPTOLayoutRank || + offsets5D.size() != kPTOLayoutRank || sizes5D.size() != kPTOLayoutRank) + return "NZ partition requires rank-5 source, offsets, and sizes"; + + for (int64_t index : {0, 3, 4}) { + int64_t expectedSize = sourceShape5D[index]; + if (isDynamic(offsets5D[index]) || isDynamic(sizes5D[index]) || + isDynamic(expectedSize)) + return ("NZ partition must prove that d" + Twine(index) + + " remains complete at compile time") + .str(); + if (offsets5D[index] != 0 || sizes5D[index] != expectedSize) + return ("NZ view cannot be partitioned inside a fractal: d" + + Twine(index) + " must keep offset=0 and size=" + + Twine(expectedSize) + " (got offset=" + Twine(offsets5D[index]) + + ", size=" + Twine(sizes5D[index]) + ")") + .str(); + } + return std::nullopt; +} + +bool mlir::pto::isCanonicalNZRoot5D(ArrayRef shape5D, + ArrayRef stride5D, + unsigned storageElemBytes) { + if (!allStatic(shape5D) || !allStatic(stride5D) || + !isNZViewCompatible5D(shape5D, stride5D, storageElemBytes)) + return false; + auto stride1 = checkedMultiply(shape5D[2], stride5D[2]); + auto stride0 = checkedMultiply(shape5D[1], stride5D[1]); + return stride1 && stride0 && stride5D[1] == *stride1 && + stride5D[0] == *stride0; +} + +bool mlir::pto::isLayoutCompatible5D(Layout layout, ArrayRef shape, + ArrayRef stride, + unsigned storageElemBytes) { + auto padded = rightAlignTo5D(shape, stride); + if (!padded) + return false; + switch (layout) { + case Layout::NZ: + if (shape.size() != kPTOLayoutRank) + return false; + return isNZViewCompatible5D(padded->shape, padded->stride, + storageElemBytes); + case Layout::ND: + case Layout::DN: + // ND/DN are logical GlobalTensor interpretations, not compactness + // promises. Existing inputs legitimately use interleaved and padded + // outer views whose minor strides do not satisfy a dense 2-D recurrence. + // Their explicit layout remains authoritative. + return true; + case Layout::MX_A_ZZ: + case Layout::MX_B_NN: + return true; + } + return false; +} + +std::optional +mlir::pto::inferLayout5D(ArrayRef shape, ArrayRef stride, + unsigned storageElemBytes, + std::optional preferredMinor2D, + bool *isMinor2DAmbiguous) { + if (isMinor2DAmbiguous) + *isMinor2DAmbiguous = false; + if (storageElemBytes == 0) + return std::nullopt; + auto padded = rightAlignTo5D(shape, stride); + if (!padded || !allStatic(padded->shape) || !allStatic(padded->stride)) + return std::nullopt; + + if (shape.size() == kPTOLayoutRank && + isCanonicalNZRoot5D(padded->shape, padded->stride, storageElemBytes) && + padded->shape[0] == 1 && (padded->shape[1] > 1 || padded->shape[2] > 1)) + return Layout::NZ; + + return inferMinor2DLayout(padded->shape[3], padded->shape[4], + padded->stride[3], padded->stride[4], + preferredMinor2D, isMinor2DAmbiguous); +} diff --git a/lib/PTO/Transforms/InferPTOLayout.cpp b/lib/PTO/Transforms/InferPTOLayout.cpp index 572995a87..aaca6dd94 100644 --- a/lib/PTO/Transforms/InferPTOLayout.cpp +++ b/lib/PTO/Transforms/InferPTOLayout.cpp @@ -21,6 +21,7 @@ //===----------------------------------------------------------------------===// #include "PTO/IR/PTO.h" +#include "PTO/IR/PTOLayoutUtils.h" #include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/Passes.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -44,12 +45,6 @@ namespace { static constexpr llvm::StringLiteral kLayoutAttrName = "layout"; static constexpr llvm::StringLiteral kInferredLayoutAttrName = "pto.inferred_layout"; -static constexpr unsigned kPaddedLayoutRank = 5; -static constexpr int64_t kUnitExtent = 1; -static constexpr int64_t kNZInnerRows = 16; -static constexpr int64_t kNZFractalBytes = 512; - -using LayoutRankVector = SmallVector; static std::optional getConstInt(Value v) { if (auto c = v.getDefiningOp()) @@ -99,120 +94,6 @@ static bool isGlobalMemRef(MemRefType ty) { return true; } -struct ShapeStride5D { - LayoutRankVector shape; - LayoutRankVector stride; -}; - -static bool isMinor2DLayout(Layout layout) { - return layout == Layout::ND || layout == Layout::DN; -} - -static std::optional rightAlignTo5D(ArrayRef shape, - ArrayRef stride) { - if (shape.size() != stride.size()) - return std::nullopt; - if (shape.size() > kPaddedLayoutRank) - return std::nullopt; - - ShapeStride5D out; - out.shape.assign(kPaddedLayoutRank, kUnitExtent); - out.stride.assign(kPaddedLayoutRank, kUnitExtent); - - const int rank = static_cast(shape.size()); - const int shift = static_cast(kPaddedLayoutRank) - rank; - for (int i = 0; i < rank; ++i) { - out.shape[shift + i] = shape[i]; - out.stride[shift + i] = stride[i]; - } - - // Derive the padded leading strides with the same rule used in EmitC: - // stride[i] = shape[i+1] * stride[i+1]. - for (int i = shift - 1; i >= 0; --i) - out.stride[i] = out.shape[i + 1] * out.stride[i + 1]; - - return out; -} - -static bool matchesNDMinor2D(int64_t rows, int64_t cols, int64_t rowStride, - int64_t colStride) { - if (cols != 1 && colStride != 1) - return false; - if (rows == 1) - return true; - return cols == 1 ? rowStride == 1 : rowStride == cols; -} - -static bool matchesDNMinor2D(int64_t rows, int64_t cols, int64_t rowStride, - int64_t colStride) { - if (rows != 1 && rowStride != 1) - return false; - if (cols == 1) - return true; - return rows == 1 ? colStride == 1 : colStride == rows; -} - -static std::optional inferMinor2DLayout( - int64_t rows, int64_t cols, int64_t rowStride, int64_t colStride, - std::optional preferredMinor2D, bool *isMinor2DAmbiguous) { - const bool nd = matchesNDMinor2D(rows, cols, rowStride, colStride); - const bool dn = matchesDNMinor2D(rows, cols, rowStride, colStride); - if (!nd && !dn) - return Layout::ND; - if (nd && dn) { - if (isMinor2DAmbiguous) - *isMinor2DAmbiguous = true; - if (preferredMinor2D && - (*preferredMinor2D == Layout::ND || *preferredMinor2D == Layout::DN)) { - return *preferredMinor2D; - } - return (cols == 1 && rows != 1) ? Layout::DN : Layout::ND; - } - return dn ? Layout::DN : Layout::ND; -} - -static std::optional inferNZLayout(ArrayRef shape, - ArrayRef stride, - unsigned elemBytes) { - int64_t sh3 = shape[2]; - int64_t sh4 = shape[3]; - int64_t sh5 = shape[4]; - int64_t st4 = stride[3]; - int64_t st5 = stride[4]; - bool alignMatch = - (sh3 == kNZInnerRows) && - (sh3 * sh4 * static_cast(elemBytes) == kNZFractalBytes); - bool strideMatch = (st5 == kUnitExtent) && (st4 == sh5); - if (alignMatch && strideMatch) - return Layout::NZ; - return std::nullopt; -} - -static std::optional inferLayout5D(ArrayRef shape, - ArrayRef strides, - unsigned elemBytes, - std::optional preferredMinor2D = - std::nullopt, - bool *isMinor2DAmbiguous = nullptr) { - if (shape.size() != strides.size() || elemBytes == 0) - return std::nullopt; - if (isMinor2DAmbiguous) - *isMinor2DAmbiguous = false; - auto padded = rightAlignTo5D(shape, strides); - if (!padded) - return std::nullopt; - - if (auto nz = inferNZLayout(padded->shape, padded->stride, elemBytes)) - return nz; - - const int64_t rows = padded->shape[3]; - const int64_t cols = padded->shape[4]; - const int64_t rowStride = padded->stride[3]; - const int64_t colStride = padded->stride[4]; - return inferMinor2DLayout(rows, cols, rowStride, colStride, - preferredMinor2D, isMinor2DAmbiguous); -} - static std::optional tileBLayoutToGlobalLayout(Type tileLikeTy) { auto tbTy = dyn_cast(tileLikeTy); if (!tbTy) @@ -261,22 +142,50 @@ static void setLayoutAttr(Operation *op, Layout layout, bool inferred) { } template -static void verifyOrSetLayoutAttr(Operation *op, - std::optional inferred, - SignalFailureFn signalFailure, - bool isMinor2DAmbiguous = false) { +static bool verifyExistingLayoutAttr(Operation *op, ArrayRef shape, + ArrayRef strides, + unsigned storageElemBytes, + SignalFailureFn signalFailure) { auto existing = op->getAttrOfType(kLayoutAttrName); - if (existing) { - if (inferred && existing.getLayout() != *inferred) { - if (isMinor2DAmbiguous && isMinor2DLayout(existing.getLayout()) && - isMinor2DLayout(*inferred)) { - return; + if (!existing) + return true; + + Layout layout = existing.getLayout(); + if (!isLayoutCompatible5D(layout, shape, strides, storageElemBytes)) { + if (layout == Layout::NZ) { + if (shape.size() == kPTOLayoutRank && strides.size() == kPTOLayoutRank) { + auto error = + getNZViewCompatibilityError(shape, strides, storageElemBytes); + op->emitError() << "user-specified layout=nz is incompatible with " + "shape/stride: " + << error.value_or("unknown NZ layout mismatch"); + } else { + op->emitError() + << "user-specified layout=nz requires a rank-5 view, got rank " + << shape.size(); } - op->emitError() << "layout mismatch: user-specified layout=" - << stringifyLayout(existing.getLayout()) - << " but inferred=" << stringifyLayout(*inferred); - signalFailure(); + } else { + op->emitError() << "user-specified layout=" << stringifyLayout(layout) + << " is incompatible with the view shape/stride"; } + signalFailure(); + return false; + } + + if (!op->getAttrOfType(kInferredLayoutAttrName)) + op->removeAttr(kInferredLayoutAttrName); + return true; +} + +template +static void resolveOrInferLayoutAttr(Operation *op, ArrayRef shape, + ArrayRef strides, + unsigned storageElemBytes, + std::optional inferred, + SignalFailureFn signalFailure) { + if (op->getAttrOfType(kLayoutAttrName)) { + (void)verifyExistingLayoutAttr(op, shape, strides, storageElemBytes, + signalFailure); return; } setLayoutAttr(op, inferred.value_or(Layout::ND), /*inferred=*/true); @@ -284,7 +193,7 @@ static void verifyOrSetLayoutAttr(Operation *op, static std::optional inferFromStaticMemRefTy(MemRefType mrTy) { if (!mrTy.hasStaticShape() || mrTy.getRank() == 0 || - mrTy.getRank() > kPaddedLayoutRank) + mrTy.getRank() > kPTOLayoutRank) return std::nullopt; SmallVector strideInts; int64_t offset = ShapedType::kDynamic; @@ -454,26 +363,29 @@ static void reconcileAmbiguousTensorViewLayout(MakeTensorViewOp op, setLayoutAttr(op.getOperation(), *pref.preferred, /*inferred=*/true); } -static bool getStaticShapeAndStride(MakeTensorViewOp op, - SmallVectorImpl &shape, - SmallVectorImpl &strides) { +static bool getShapeAndStride(MakeTensorViewOp op, + SmallVectorImpl &shape, + SmallVectorImpl &strides, + bool &allStatic) { auto tvTy = dyn_cast(op.getResult().getType()); if (!tvTy) return false; const size_t rank = op.getShape().size(); - if (rank == 0 || rank > kPaddedLayoutRank) + if (rank == 0 || rank > kPTOLayoutRank) return false; + allStatic = true; shape.clear(); shape.reserve(rank); for (size_t i = 0; i < rank; ++i) { int64_t dim = tvTy.getShape()[i]; if (dim == ShapedType::kDynamic) { auto v = getConstInt(op.getShape()[i]); - if (!v) - return false; - dim = *v; + if (v) + dim = *v; + else + allStatic = false; } shape.push_back(dim); } @@ -482,27 +394,35 @@ static bool getStaticShapeAndStride(MakeTensorViewOp op, strides.reserve(rank); for (Value s : op.getStrides()) { auto v = getConstInt(s); - if (!v) - return false; - strides.push_back(*v); + if (v) { + strides.push_back(*v); + } else { + strides.push_back(ShapedType::kDynamic); + allStatic = false; + } } return true; } -template -static bool getConstFoldResults(ArrayRef values, - SmallVectorImpl &result, - SignalFailureFn signalFailure, - Operation *op) { +static bool getStaticShapeAndStride(MakeTensorViewOp op, + SmallVectorImpl &shape, + SmallVectorImpl &strides) { + bool allStatic = false; + return getShapeAndStride(op, shape, strides, allStatic) && allStatic; +} + +static bool getFoldResults(ArrayRef values, + SmallVectorImpl &result, bool &allStatic) { result.clear(); result.reserve(values.size()); for (OpFoldResult value : values) { auto folded = getConstInt(value); - if (!folded) { - verifyOrSetLayoutAttr(op, std::nullopt, signalFailure); - return false; + if (folded) { + result.push_back(*folded); + } else { + result.push_back(ShapedType::kDynamic); + allStatic = false; } - result.push_back(*folded); } return true; } @@ -524,26 +444,133 @@ static ResolvedLayoutInfo resolveLayoutFromViewValue(Value v) { def = v.getDefiningOp(); continue; } + if (auto subview = dyn_cast(def)) { + v = subview.getSource(); + def = v.getDefiningOp(); + continue; + } + if (auto reinterpret = dyn_cast(def)) { + v = reinterpret.getSource(); + def = v.getDefiningOp(); + continue; + } + if (auto cast = dyn_cast(def)) { + v = cast.getSource(); + def = v.getDefiningOp(); + continue; + } break; } return info; } +static void getValuesOrDynamic(ValueRange values, + SmallVectorImpl &result) { + result.clear(); + result.reserve(values.size()); + for (Value value : values) + result.push_back(getConstInt(value).value_or(ShapedType::kDynamic)); +} + +static void getFoldResultsOrDynamic(ArrayRef values, + SmallVectorImpl &result) { + result.clear(); + result.reserve(values.size()); + for (OpFoldResult value : values) + result.push_back(getConstInt(value).value_or(ShapedType::kDynamic)); +} + +static bool getResolvedViewShape(Value value, SmallVectorImpl &shape) { + if (auto make = value.getDefiningOp()) { + auto type = dyn_cast(make.getResult().getType()); + if (!type) + return false; + shape.assign(type.getShape().begin(), type.getShape().end()); + for (auto [index, operand] : llvm::enumerate(make.getShape())) { + if (shape[index] == ShapedType::kDynamic) { + if (auto constant = getConstInt(operand)) + shape[index] = *constant; + } + } + return true; + } + + if (auto partition = value.getDefiningOp()) { + auto type = + dyn_cast(partition.getResult().getType()); + if (!type) + return false; + shape.assign(type.getShape().begin(), type.getShape().end()); + for (auto [index, operand] : llvm::enumerate(partition.getSizes())) { + if (shape[index] == ShapedType::kDynamic) { + if (auto constant = getConstInt(operand)) + shape[index] = *constant; + } + } + return true; + } + + return false; +} + +static bool verifyNZPartitionView(PartitionViewOp op) { + SmallVector sourceShape; + if (!getResolvedViewShape(op.getSource(), sourceShape)) + return false; + SmallVector offsets; + SmallVector sizes; + getValuesOrDynamic(op.getOffsets(), offsets); + getValuesOrDynamic(op.getSizes(), sizes); + if (auto error = + getNZSubviewCompatibilityError(sourceShape, offsets, sizes)) { + op.emitError(*error); + return false; + } + return true; +} + +static bool verifyNZMemRefSubview(memref::SubViewOp op) { + auto sourceType = dyn_cast(op.getSource().getType()); + auto resultType = dyn_cast(op.getType()); + if (!sourceType || !resultType || + sourceType.getRank() != static_cast(kPTOLayoutRank) || + resultType.getRank() != sourceType.getRank()) { + op.emitError("NZ subview requires a rank-preserving rank-5 view"); + return false; + } + + SmallVector offsets; + SmallVector sizes; + getFoldResultsOrDynamic(op.getMixedOffsets(), offsets); + getFoldResultsOrDynamic(op.getMixedSizes(), sizes); + if (auto error = getNZSubviewCompatibilityError(sourceType.getShape(), + offsets, sizes)) { + op.emitError(*error); + return false; + } + return true; +} + template static void inferMakeTensorViewLayoutAttr(MakeTensorViewOp op, SignalFailureFn signalFailure) { SmallVector shape; SmallVector strides; - if (!getStaticShapeAndStride(op, shape, strides)) { - verifyOrSetLayoutAttr(op.getOperation(), std::nullopt, signalFailure); + bool allStatic = false; + if (!getShapeAndStride(op, shape, strides, allStatic)) return; - } bool isAmbiguous = false; - auto inferred = inferMakeTensorViewLayout(op, shape, strides, isAmbiguous); - verifyOrSetLayoutAttr(op.getOperation(), inferred, signalFailure, - isAmbiguous); - if (isAmbiguous) + std::optional inferred; + if (allStatic) + inferred = inferMakeTensorViewLayout(op, shape, strides, isAmbiguous); + resolveOrInferLayoutAttr( + op.getOperation(), shape, strides, + elemByteSize( + cast(op.getResult().getType()).getElementType()), + inferred, signalFailure); + if (allStatic && isAmbiguous && + op->getAttrOfType(kInferredLayoutAttrName)) reconcileAmbiguousTensorViewLayout(op, shape); } @@ -555,26 +582,24 @@ static void inferReinterpretCastLayoutAttr(memref::ReinterpretCastOp op, return; const size_t rank = op.getMixedSizes().size(); - if (rank == 0 || rank > kPaddedLayoutRank) { - verifyOrSetLayoutAttr(op.getOperation(), std::nullopt, signalFailure); + if (rank == 0 || rank > kPTOLayoutRank) return; - } SmallVector shape; SmallVector strides; - if (!getConstFoldResults(op.getMixedSizes(), shape, signalFailure, - op.getOperation()) || - !getConstFoldResults(op.getMixedStrides(), strides, signalFailure, - op.getOperation())) { - return; - } + bool allStatic = true; + (void)getFoldResults(op.getMixedSizes(), shape, allStatic); + (void)getFoldResults(op.getMixedStrides(), strides, allStatic); bool isMinor2DAmbiguous = false; - auto inferred = - inferLayout5D(shape, strides, elemByteSize(mrTy.getElementType()), - std::nullopt, &isMinor2DAmbiguous); - verifyOrSetLayoutAttr(op.getOperation(), inferred, signalFailure, - isMinor2DAmbiguous); + std::optional inferred; + if (allStatic) + inferred = + inferLayout5D(shape, strides, elemByteSize(mrTy.getElementType()), + std::nullopt, &isMinor2DAmbiguous); + resolveOrInferLayoutAttr(op.getOperation(), shape, strides, + elemByteSize(mrTy.getElementType()), inferred, + signalFailure); } struct InferPTOLayoutPass @@ -597,32 +622,69 @@ struct InferPTOLayoutPass }); // ------------------------------------------------------------------ - // 2) memref.reinterpret_cast (lowered from make_tensor_view) + // 2) pto.partition_view: validate against the resolved source layout. + // The logical inheritance is resolved through the source chain and later + // materialized on the lowered memref.subview. Avoid adding derived + // attributes here so the source remains the single authority. + // ------------------------------------------------------------------ + func.walk([&](PartitionViewOp op) { + auto sourceInfo = resolveLayoutFromViewValue(op.getSource()); + if (!sourceInfo.layout) + return; + if (*sourceInfo.layout == Layout::NZ && !verifyNZPartitionView(op)) { + signalPassFailure(); + return; + } + if (auto existing = op->getAttrOfType(kLayoutAttrName); + existing && existing.getLayout() != *sourceInfo.layout) { + op.emitError() << "partition layout=" + << stringifyLayout(existing.getLayout()) + << " does not match source layout=" + << stringifyLayout(*sourceInfo.layout); + signalPassFailure(); + return; + } + }); + + // ------------------------------------------------------------------ + // 3) memref.reinterpret_cast (lowered from make_tensor_view) // ------------------------------------------------------------------ func.walk([&](memref::ReinterpretCastOp op) { inferReinterpretCastLayoutAttr(op, [this] { signalPassFailure(); }); }); // ------------------------------------------------------------------ - // 3) memref.subview: layout is preserved from the source view + // 4) memref.subview: preserve layout only across a legal derived view. // ------------------------------------------------------------------ func.walk([&](memref::SubViewOp op) { auto resTy = dyn_cast(op.getType()); if (!resTy || !isGlobalMemRef(resTy)) return; - if (op->getAttrOfType(kLayoutAttrName)) + auto sourceInfo = resolveLayoutFromViewValue(op.getSource()); + auto existing = op->getAttrOfType(kLayoutAttrName); + if (existing && sourceInfo.layout && + existing.getLayout() != *sourceInfo.layout) { + op.emitError() << "subview layout=" + << stringifyLayout(existing.getLayout()) + << " does not match source layout=" + << stringifyLayout(*sourceInfo.layout); + signalPassFailure(); return; + } - if (Operation *def = op.getSource().getDefiningOp()) { - if (auto srcLayout = def->getAttrOfType(kLayoutAttrName)) { - op->setAttr(kLayoutAttrName, srcLayout); - if (auto inferred = - def->getAttrOfType(kInferredLayoutAttrName)) { - op->setAttr(kInferredLayoutAttrName, inferred); - } + if (existing || sourceInfo.layout) { + Layout layout = existing ? existing.getLayout() : *sourceInfo.layout; + bool inferred = + existing ? static_cast( + op->getAttrOfType(kInferredLayoutAttrName)) + : sourceInfo.inferred; + if (layout == Layout::NZ && !verifyNZMemRefSubview(op)) { + signalPassFailure(); return; } + setLayoutAttr(op.getOperation(), layout, inferred); + return; } // Fallback: if source memref type is fully static, infer from it. @@ -645,12 +707,16 @@ struct InferPTOLayoutPass auto inferred = inferLayout5D(srcTy.getShape(), strideInts, elemByteSize(srcTy.getElementType())); + if (inferred == Layout::NZ && !verifyNZMemRefSubview(op)) { + signalPassFailure(); + return; + } setLayoutAttr(op.getOperation(), inferred.value_or(Layout::ND), /*inferred=*/true); }); // ------------------------------------------------------------------ - // 4) pto.tload / pto.tstore: attach layout for static GM memrefs so EmitC + // 5) pto.tload / pto.tstore: attach layout for static GM memrefs so EmitC // doesn't need to infer again in buildGlobalTensorFromMemref(). // ------------------------------------------------------------------ func.walk([&](pto::TLoadOp op) { diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 02b2d84f0..e99615e0f 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -16,8 +16,9 @@ #include #include "PTO/IR/PTO.h" -#include "PTO/IR/PTOTypeUtils.h" +#include "PTO/IR/PTOLayoutUtils.h" #include "PTO/IR/PTOSyncUtils.h" +#include "PTO/IR/PTOTypeUtils.h" #include "PTO/Transforms/MemoryConsistencyAttrs.h" #include "PTO/Transforms/Passes.h" #include "Utils.h" @@ -4350,25 +4351,9 @@ static emitc::OpaqueType getGlobalTensorOpaqueTypeFromShape( static std::string inferFallbackGlobalTensorLayout(ArrayRef shape5D, ArrayRef stride5D, Type elemTy) { - int elemBytes = getGlobalTensorElementBytes(elemTy); - if (elemBytes == 0) - return "pto::Layout::ND"; - if (shape5D[2] == 16 && multiplyOrDynamic(shape5D[2], shape5D[3]) * elemBytes == 512 && - stride5D[4] == 1 && stride5D[3] == shape5D[4]) { - return "pto::Layout::NZ"; - } - - bool isRowMajor = stride5D[4] == 1; - for (int i = 3; i >= 0 && isRowMajor; --i) - isRowMajor = stride5D[i] == multiplyOrDynamic(stride5D[i + 1], shape5D[i + 1]); - - bool isColMajor = stride5D[0] == 1; - for (int i = 0; i < 4 && isColMajor; ++i) - isColMajor = stride5D[i + 1] == multiplyOrDynamic(stride5D[i], shape5D[i]); - - if (isColMajor) - return "pto::Layout::DN"; - return isRowMajor ? "pto::Layout::ND" : "pto::Layout::ND"; + auto layout = + inferLayout5D(shape5D, stride5D, getGlobalTensorElementBytes(elemTy)); + return layoutToEmitCString(layout.value_or(Layout::ND)); } static std::string resolveGlobalTensorLayout(Operation *anchor, Value basePtr, diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 46da6fb59..1398d1545 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -1213,6 +1213,8 @@ static LogicalResult lowerPtrToIntOps(func::FuncOp func, MLIRContext *ctx) { rc->setAttr("pto.addptr_trace", rewriter.getUnitAttr()); if (auto layoutAttr = op.getLayoutAttr()) rc->setAttr("layout", layoutAttr); + if (auto inferredAttr = op->getAttrOfType("pto.inferred_layout")) + rc->setAttr("pto.inferred_layout", inferredAttr); rewriter.replaceOp(op, rc.getResult()); } return success(); @@ -1380,6 +1382,9 @@ static LogicalResult lowerPartitionViewOps(func::FuncOp func, MLIRContext *ctx) if (Operation *srcDef = src.getDefiningOp()) { if (auto layoutAttr = srcDef->getAttrOfType("layout")) sv->setAttr("layout", layoutAttr); + if (auto inferredAttr = + srcDef->getAttrOfType("pto.inferred_layout")) + sv->setAttr("pto.inferred_layout", inferredAttr); } rewriter.replaceOp(op, sv.getResult()); } @@ -2192,6 +2197,10 @@ struct PTOViewToMemrefPass if (auto layoutAttr = op.getLayoutAttr()) { rc->setAttr("layout", layoutAttr); } + if (auto inferredAttr = + op->getAttrOfType("pto.inferred_layout")) { + rc->setAttr("pto.inferred_layout", inferredAttr); + } rewriter.replaceOp(op, rc.getResult()); } diff --git a/ptodsl/docs/user_guide/01-introduction.md b/ptodsl/docs/user_guide/01-introduction.md index 677e23de0..65196123f 100644 --- a/ptodsl/docs/user_guide/01-introduction.md +++ b/ptodsl/docs/user_guide/01-introduction.md @@ -157,6 +157,31 @@ contract is also owned here: the runtime grid (e.g., `batch * heads` blocks) is declared at the call site, and block/subblock indices are queried via `pto.get_block_idx()` and friends. +`make_tensor_view(..., layout=...)` accepts `"ND"`, `"DN"`, and `"NZ"` +(case-insensitive). Use it when the logical layout cannot be distinguished from +shape/stride alone. For example, fp32 canonical NZ has `C0 = 32 / 4 = 8`: + + +```python +from ptodsl import pto + + +@pto.jit(target="a5") +def explicit_nz_view(ptr: pto.ptr(pto.f32, "gm")): + view = pto.make_tensor_view( + ptr, + shape=[1, 8, 8, 16, 8], + strides=[8192, 1024, 128, 8, 1], + layout="NZ", + ) + _ = view + return +``` + +An explicit layout is preserved by PTOAS. Without one, only a non-degenerate +canonical NZ five-dimensional root is inferred as NZ; a single-fractal +`[1,1,1,16,C0]` view remains ND. + #### `entry=False` — kernel modules A kernel module is a device-side function that entries and other modules call diff --git a/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto b/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto index 8fbc3109e..4b40e070b 100644 --- a/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto +++ b/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto @@ -7,6 +7,7 @@ // See LICENSE in the root of the software repository for the full text of the License. // RUN: ptoas --pto-arch=a3 %s 2>&1 | FileCheck %s --check-prefix=A3 +// RUN: ptoas --pto-arch=a5 %s 2>&1 | FileCheck %s --check-prefix=A3 module { func.func @comm_globaltensor_bf16_nz( @@ -40,9 +41,9 @@ module { // 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: constexpr pto::Layout [[I8_LAYOUT:GT.*_layout]] = pto::Layout::ND; // 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: constexpr pto::Layout [[I64_LAYOUT:GT.*_layout]] = pto::Layout::DN; // A3: using [[I64_GT:GT.*]] = GlobalTensor; diff --git a/test/lit/pto/issue527_nz_canonical_view_infer.pto b/test/lit/pto/issue527_nz_canonical_view_infer.pto new file mode 100644 index 000000000..584fed0d9 --- /dev/null +++ b/test/lit/pto/issue527_nz_canonical_view_infer.pto @@ -0,0 +1,472 @@ +// 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: ptoas --pto-arch=a3 --emit-pto-ir --mlir-print-ir-after=pto-infer-layout %t/positive.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=INFER +// RUN: ptoas --pto-arch=a5 --emit-pto-ir --mlir-print-ir-after=pto-infer-layout %t/positive.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=INFER +// RUN: ptoas --pto-arch=a5 --emit-pto-ir --mlir-print-ir-after=pto-view-to-memref %t/positive.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=LOWER +// RUN: ptoas --pto-arch=a3 %t/emitc.pto 2>&1 | FileCheck %s --check-prefix=EMITC +// RUN: ptoas --pto-arch=a5 %t/emitc.pto 2>&1 | FileCheck %s --check-prefix=EMITC +// RUN: not ptoas --pto-arch=a5 %t/invalid_inner_shape.pto 2>&1 | FileCheck %s --check-prefix=BAD-SHAPE +// RUN: not ptoas --pto-arch=a5 %t/invalid_rank.pto 2>&1 | FileCheck %s --check-prefix=BAD-RANK +// RUN: not ptoas --pto-arch=a5 %t/invalid_inner_stride.pto 2>&1 | FileCheck %s --check-prefix=BAD-STRIDE +// RUN: not ptoas --pto-arch=a5 %t/invalid_outer_gap.pto 2>&1 | FileCheck %s --check-prefix=BAD-GAP +// RUN: not ptoas --pto-arch=a5 %t/invalid_d3_partition.pto 2>&1 | FileCheck %s --check-prefix=BAD-D3 +// RUN: not ptoas --pto-arch=a5 %t/invalid_d4_partition.pto 2>&1 | FileCheck %s --check-prefix=BAD-D4 +// RUN: not ptoas --pto-arch=a5 %t/dynamic_inner_partition.pto 2>&1 | FileCheck %s --check-prefix=DYNAMIC-INNER +// RUN: not ptoas --pto-arch=a5 %t/invalid_memref_subview.pto 2>&1 | FileCheck %s --check-prefix=BAD-MEMREF + +// INFER-DAG: pto.make_tensor_view {{.*}} shape = {{.*}} {layout = #pto.layout, pto.inferred_layout = true} : !pto.tensor_view<1x8x8x16x8xf32> +// INFER-DAG: pto.make_tensor_view {{.*}} shape = {{.*}} {layout = #pto.layout, pto.inferred_layout = true} : !pto.tensor_view<1x4x8x16x16xf16> +// INFER-DAG: pto.make_tensor_view {{.*}} shape = {{.*}} {layout = #pto.layout, pto.inferred_layout = true} : !pto.tensor_view<1x2x8x16x32xi8> +// INFER-DAG: pto.make_tensor_view {{.*}} shape = {{.*}} {layout = #pto.layout, pto.inferred_layout = true} : !pto.tensor_view<1x2x8x16x32x!pto.f4E1M2x2> +// INFER-DAG: pto.make_tensor_view {{.*}} shape = {{.*}} {layout = #pto.layout, pto.inferred_layout = true} : !pto.tensor_view<8x8x16x8xf32> +// INFER-DAG: pto.make_tensor_view {{.*}} shape = {{.*}} {layout = #pto.layout, pto.inferred_layout = true} : !pto.tensor_view<4x1x16x8x16xf32> +// INFER-DAG: pto.make_tensor_view {{.*}} shape = {{.*}} {layout = #pto.layout, pto.inferred_layout = true} : !pto.tensor_view<1x1x1x16x8xf32> +// INFER-DAG: pto.make_tensor_view {{.*}} {layout = #pto.layout} : !pto.tensor_view<1x1x1x16x8xf32> +// INFER-DAG: pto.make_tensor_view {{.*}} {layout = #pto.layout} : !pto.tensor_view<1x2x1x16x8xf32> +// INFER-DAG: pto.make_tensor_view {{.*}} {layout = #pto.layout} : !pto.tensor_view<1x8x8x16x8xf32> +// INFER-DAG: pto.make_tensor_view {{.*}} {layout = #pto.layout} : !pto.tensor_view<1x4x8x16x8xf32> +// INFER-DAG: pto.partition_view {{.*}} : !pto.partition_tensor_view<1x4x8x16x8xf32> +// INFER-DAG: pto.make_tensor_view {{.*}} {layout = #pto.layout} : !pto.tensor_view<1x?x?x16x8xf32> +// INFER-DAG: pto.mgather +// INFER-DAG: pto.mscatter +// INFER-DAG: memref.subview {{.*}} {layout = #pto.layout} + +// LOWER-DAG: memref.reinterpret_cast {{.*}} {layout = #pto.layout, pto.inferred_layout = true} +// LOWER-DAG: memref.subview {{.*}} {layout = #pto.layout, pto.inferred_layout = true} +// LOWER-DAG: memref.subview {{.*}} {layout = #pto.layout, pto.inferred_layout = true} +// LOWER-DAG: memref.subview {{.*}} {layout = #pto.layout, pto.inferred_layout = true} +// LOWER-DAG: memref.reinterpret_cast {{.*}} {layout = #pto.layout} +// LOWER-DAG: memref.subview {{.*}} {layout = #pto.layout} + +// EMITC: pto::Shape<1, 8, 8, 16, 8> +// EMITC: pto::Stride<8192, 1024, 128, 8, 1> +// EMITC: GlobalTensor, pto::Stride<8192, 1024, 128, 8, 1>, pto::Layout::NZ> + +// BAD-SHAPE: user-specified layout=nz is incompatible with shape/stride: expected shape[3]==16 (got 8) +// BAD-RANK: user-specified layout=nz requires a rank-5 view, got rank 4 +// BAD-STRIDE: user-specified layout=nz is incompatible with shape/stride: expected stride[2]==128 (got 64) +// BAD-GAP: user-specified layout=nz is incompatible with shape/stride: expected stride[1] >= shape[2] * stride[2] (1024), got 512 +// BAD-D3: NZ view cannot be partitioned inside a fractal: d3 must keep offset=0 and size=16 (got offset=1, size=15) +// BAD-D4: NZ view cannot be partitioned inside a fractal: d4 must keep offset=0 and size=8 (got offset=1, size=7) +// DYNAMIC-INNER: NZ partition must prove that d3 remains complete at compile time +// BAD-MEMREF: NZ view cannot be partitioned inside a fractal: d4 must keep offset=0 and size=8 (got offset=1, size=7) + +//--- positive.pto +module { + func.func @canonical_roots( + %fp32: !pto.ptr, %fp16: !pto.ptr, %i8: !pto.ptr, + %fp4: !pto.ptr) { + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c4 = arith.constant 4 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c128 = arith.constant 128 : index + %c256 = arith.constant 256 : index + %c512 = arith.constant 512 : index + %c1024 = arith.constant 1024 : index + %c2048 = arith.constant 2048 : index + %c4096 = arith.constant 4096 : index + %c8192 = arith.constant 8192 : index + + %fp32_nz = pto.make_tensor_view %fp32, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + : !pto.tensor_view<1x8x8x16x8xf32> + %fp16_nz = pto.make_tensor_view %fp16, + shape = [%c1, %c4, %c8, %c16, %c16], + strides = [%c8192, %c2048, %c256, %c16, %c1] + : !pto.tensor_view<1x4x8x16x16xf16> + %i8_nz = pto.make_tensor_view %i8, + shape = [%c1, %c2, %c8, %c16, %c32], + strides = [%c8192, %c4096, %c512, %c32, %c1] + : !pto.tensor_view<1x2x8x16x32xi8> + %fp4_nz = pto.make_tensor_view %fp4, + shape = [%c1, %c2, %c8, %c16, %c32], + strides = [%c8192, %c4096, %c512, %c32, %c1] + : !pto.tensor_view<1x2x8x16x32x!pto.f4E1M2x2> + %rank4_nd = pto.make_tensor_view %fp32, + shape = [%c8, %c8, %c16, %c8], + strides = [%c1024, %c128, %c8, %c1] + : !pto.tensor_view<8x8x16x8xf32> + %legacy_nd = pto.make_tensor_view %fp32, + shape = [%c4, %c1, %c16, %c8, %c16], + strides = [%c2048, %c2048, %c128, %c16, %c1] + : !pto.tensor_view<4x1x16x8x16xf32> + %single_nd = pto.make_tensor_view %fp32, + shape = [%c1, %c1, %c1, %c16, %c8], + strides = [%c128, %c128, %c128, %c8, %c1] + : !pto.tensor_view<1x1x1x16x8xf32> + return + } + + func.func @explicit_layouts(%ptr: !pto.ptr) { + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c4 = arith.constant 4 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %c8192 = arith.constant 8192 : index + + %single_nz = pto.make_tensor_view %ptr, + shape = [%c1, %c1, %c1, %c16, %c8], + strides = [%c128, %c128, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x1x1x16x8xf32> + %canonical_nd = pto.make_tensor_view %ptr, + shape = [%c1, %c2, %c1, %c16, %c8], + strides = [%c256, %c128, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x2x1x16x8xf32> + %canonical_nz = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x8x8x16x8xf32> + %gapped_nz = pto.make_tensor_view %ptr, + shape = [%c1, %c4, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x4x8x16x8xf32> + return + } + + func.func @partition_and_nested_propagation(%ptr: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c4 = arith.constant 4 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c1024 = arith.constant 1024 : index + %c8192 = arith.constant 8192 : index + + %root = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + : !pto.tensor_view<1x8x8x16x8xf32> + %d1 = pto.partition_view %root, + offsets = [%c0, %c2, %c0, %c0, %c0], + sizes = [%c1, %c4, %c8, %c16, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x4x8x16x8xf32> + %d2 = pto.partition_view %root, + offsets = [%c0, %c0, %c2, %c0, %c0], + sizes = [%c1, %c8, %c4, %c16, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x8x4x16x8xf32> + %nested = pto.partition_view %d1, + offsets = [%c0, %c1, %c2, %c0, %c0], + sizes = [%c1, %c2, %c4, %c16, %c8] + : !pto.partition_tensor_view<1x4x8x16x8xf32> + -> !pto.partition_tensor_view<1x2x4x16x8xf32> + return + } + + func.func @dynamic_outer_partition( + %ptr: !pto.ptr, %d1: index, %d2: index, + %stride0: index, %stride1: index, + %offset1: index, %offset2: index, + %size1: index, %size2: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + + %root = pto.make_tensor_view %ptr, + shape = [%c1, %d1, %d2, %c16, %c8], + strides = [%stride0, %stride1, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x?x?x16x8xf32> + %part = pto.partition_view %root, + offsets = [%c0, %offset1, %offset2, %c0, %c0], + sizes = [%c1, %size1, %size2, %c16, %c8] + : !pto.tensor_view<1x?x?x16x8xf32> + -> !pto.partition_tensor_view<1x?x?x16x8xf32> + return + } + + func.func @mgather_explicit_nd(%ptr: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c128 = arith.constant 128 : index + %c256 = arith.constant 256 : index + %view = pto.make_tensor_view %ptr, + shape = [%c1, %c2, %c1, %c16, %c8], + strides = [%c256, %c128, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x2x1x16x8xf32> + %part = pto.partition_view %view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c2, %c1, %c16, %c8] + : !pto.tensor_view<1x2x1x16x8xf32> + -> !pto.partition_tensor_view<1x2x1x16x8xf32> + %idx = pto.alloc_tile + : !pto.tile_buf + %dst = pto.alloc_tile + : !pto.tile_buf + pto.mgather + ins(%part, %idx : !pto.partition_tensor_view<1x2x1x16x8xf32>, !pto.tile_buf) + outs(%dst : !pto.tile_buf) + {coalesce = #pto} + return + } + + func.func @mscatter_explicit_nd(%ptr: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c256 = arith.constant 256 : index + %view = pto.make_tensor_view %ptr, + shape = [%c1, %c2, %c1, %c16, %c8], + strides = [%c256, %c128, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x2x1x16x8xf32> + %part = pto.partition_view %view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c2, %c1, %c16, %c8] + : !pto.tensor_view<1x2x1x16x8xf32> + -> !pto.partition_tensor_view<1x2x1x16x8xf32> + %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(%part : !pto.partition_tensor_view<1x2x1x16x8xf32>) + return + } + + func.func @memref_nz_subview( + %base: memref>) { + %root = memref.reinterpret_cast %base to + offset: [0], + sizes: [1, 8, 8, 16, 8], + strides: [8192, 1024, 128, 8, 1] + {layout = #pto.layout} + : memref> + to memref<1x8x8x16x8xf32, strided<[8192, 1024, 128, 8, 1]>, #pto.address_space> + %part = memref.subview %root[0, 2, 0, 0, 0] [1, 4, 8, 16, 8] [1, 1, 1, 1, 1] + : memref<1x8x8x16x8xf32, strided<[8192, 1024, 128, 8, 1]>, #pto.address_space> + to memref<1x4x8x16x8xf32, strided<[8192, 1024, 128, 8, 1], offset: 2048>, #pto.address_space> + %tile = pto.alloc_tile + : !pto.tile_buf + pto.tload + ins(%part : memref<1x4x8x16x8xf32, strided<[8192, 1024, 128, 8, 1], offset: 2048>, #pto.address_space>) + outs(%tile : !pto.tile_buf) + return + } +} + +//--- emitc.pto +module { + func.func @canonical_nz_emitc(%ptr: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c1024 = arith.constant 1024 : index + %c8192 = arith.constant 8192 : index + %root = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + : !pto.tensor_view<1x8x8x16x8xf32> + %part = pto.partition_view %root, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c8, %c8, %c16, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x8x8x16x8xf32> + %tile = pto.alloc_tile + : !pto.tile_buf + pto.tload + ins(%part : !pto.partition_tensor_view<1x8x8x16x8xf32>) + outs(%tile : !pto.tile_buf) + return + } +} + +//--- invalid_inner_shape.pto +module { + func.func @invalid_inner_shape(%ptr: !pto.ptr) { + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c64 = arith.constant 64 : index + %c512 = arith.constant 512 : index + %c4096 = arith.constant 4096 : index + %bad = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c8, %c8], + strides = [%c4096, %c512, %c64, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x8x8x8x8xf32> + return + } +} + +//--- invalid_rank.pto +module { + func.func @invalid_rank(%ptr: !pto.ptr) { + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c1024 = arith.constant 1024 : index + %bad = pto.make_tensor_view %ptr, + shape = [%c8, %c8, %c16, %c8], + strides = [%c1024, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<8x8x16x8xf32> + return + } +} + +//--- invalid_inner_stride.pto +module { + func.func @invalid_inner_stride(%ptr: !pto.ptr) { + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c64 = arith.constant 64 : index + %c512 = arith.constant 512 : index + %c4096 = arith.constant 4096 : index + %bad = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c4096, %c512, %c64, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x8x8x16x8xf32> + return + } +} + +//--- invalid_outer_gap.pto +module { + func.func @invalid_outer_gap(%ptr: !pto.ptr) { + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c512 = arith.constant 512 : index + %c4096 = arith.constant 4096 : index + %bad = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c4096, %c512, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x8x8x16x8xf32> + return + } +} + +//--- invalid_d3_partition.pto +module { + func.func @invalid_d3_partition(%ptr: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c15 = arith.constant 15 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c1024 = arith.constant 1024 : index + %c8192 = arith.constant 8192 : index + %root = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + : !pto.tensor_view<1x8x8x16x8xf32> + %bad = pto.partition_view %root, + offsets = [%c0, %c0, %c0, %c1, %c0], + sizes = [%c1, %c8, %c8, %c15, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x8x8x15x8xf32> + return + } +} + +//--- invalid_d4_partition.pto +module { + func.func @invalid_d4_partition(%ptr: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c7 = arith.constant 7 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c1024 = arith.constant 1024 : index + %c8192 = arith.constant 8192 : index + %root = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + : !pto.tensor_view<1x8x8x16x8xf32> + %bad = pto.partition_view %root, + offsets = [%c0, %c0, %c0, %c0, %c1], + sizes = [%c1, %c8, %c8, %c16, %c7] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x8x8x16x7xf32> + return + } +} + +//--- dynamic_inner_partition.pto +module { + func.func @dynamic_inner_partition( + %ptr: !pto.ptr, %innerOffset: index, %innerSize: index) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c1024 = arith.constant 1024 : index + %c8192 = arith.constant 8192 : index + %root = pto.make_tensor_view %ptr, + shape = [%c1, %c8, %c8, %c16, %c8], + strides = [%c8192, %c1024, %c128, %c8, %c1] + : !pto.tensor_view<1x8x8x16x8xf32> + %bad = pto.partition_view %root, + offsets = [%c0, %c0, %c0, %innerOffset, %c0], + sizes = [%c1, %c8, %c8, %innerSize, %c8] + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x8x8x?x8xf32> + return + } +} + +//--- invalid_memref_subview.pto +module { + func.func @invalid_memref_subview( + %base: memref>) { + %root = memref.reinterpret_cast %base to + offset: [0], + sizes: [1, 8, 8, 16, 8], + strides: [8192, 1024, 128, 8, 1] + {layout = #pto.layout} + : memref> + to memref<1x8x8x16x8xf32, strided<[8192, 1024, 128, 8, 1]>, #pto.address_space> + %bad = memref.subview %root[0, 0, 0, 0, 1] [1, 8, 8, 16, 7] [1, 1, 1, 1, 1] + : memref<1x8x8x16x8xf32, strided<[8192, 1024, 128, 8, 1]>, #pto.address_space> + to memref<1x8x8x16x7xf32, strided<[8192, 1024, 128, 8, 1], offset: 1>, #pto.address_space> + %tile = pto.alloc_tile + : !pto.tile_buf + pto.tload + ins(%bad : memref<1x8x8x16x7xf32, strided<[8192, 1024, 128, 8, 1], offset: 1>, #pto.address_space>) + outs(%tile : !pto.tile_buf) + return + } +} From af8b6adbab002e889cf47862a8d01ec29b8c194c Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Wed, 29 Jul 2026 11:46:23 +0800 Subject: [PATCH 5/6] test: align CI checks with vector fence lowering --- include/PTO/IR/PTOLayoutUtils.h | 14 ++++++-------- lib/PTO/IR/PTOLayoutUtils.cpp | 14 ++++++-------- test/lit/pto/issue872_tput_tnotify_release.pto | 18 ++++++------------ .../pto/signal_payload_cache_consistency.pto | 18 ++++++------------ 4 files changed, 24 insertions(+), 40 deletions(-) diff --git a/include/PTO/IR/PTOLayoutUtils.h b/include/PTO/IR/PTOLayoutUtils.h index 9ad44d8cc..bfc3577aa 100644 --- a/include/PTO/IR/PTOLayoutUtils.h +++ b/include/PTO/IR/PTOLayoutUtils.h @@ -1,12 +1,10 @@ // 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. +// 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. #ifndef PTO_IR_PTOLAYOUTUTILS_H #define PTO_IR_PTOLAYOUTUTILS_H diff --git a/lib/PTO/IR/PTOLayoutUtils.cpp b/lib/PTO/IR/PTOLayoutUtils.cpp index d52b44919..6220b375e 100644 --- a/lib/PTO/IR/PTOLayoutUtils.cpp +++ b/lib/PTO/IR/PTOLayoutUtils.cpp @@ -1,12 +1,10 @@ // 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. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. #include "PTO/IR/PTOLayoutUtils.h" diff --git a/test/lit/pto/issue872_tput_tnotify_release.pto b/test/lit/pto/issue872_tput_tnotify_release.pto index ecd19f65c..677da4762 100644 --- a/test/lit/pto/issue872_tput_tnotify_release.pto +++ b/test/lit/pto/issue872_tput_tnotify_release.pto @@ -3,14 +3,14 @@ // 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. +// 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 issue #872 under the explicit memory-consistency contract: // pto.comm.tput/pto.comm.tbroadcast can issue MTE3 GM stores internally. A // following TNotify publishes a cross-rank signal, so the producer side must -// execute a GM fence before the signal. The fence lowering conservatively -// drains MTE2, MTE3, and FIX before dsb(DSB_DDR). +// execute a GM fence before the signal. In a vector kernel, the fence lowering +// drains all vector pipes with PIPE_ALL before dsb(DSB_DDR). // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s @@ -149,9 +149,7 @@ module { // CHECK-LABEL: AICORE void tput_tnotify_release( // CHECK: pto::comm::TPUT( // CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK: pipe_barrier(PIPE_MTE2); -// CHECK-NEXT: pipe_barrier(PIPE_MTE3); -// CHECK-NEXT: pipe_barrier(PIPE_FIX); +// CHECK: pipe_barrier(PIPE_ALL); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( @@ -159,17 +157,13 @@ module { // CHECK: pto::comm::TPUT( // CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE // CHECK-NEXT: pipe_barrier(PIPE_ALL); -// CHECK-NEXT: pipe_barrier(PIPE_MTE2); -// CHECK-NEXT: pipe_barrier(PIPE_MTE3); -// CHECK-NEXT: pipe_barrier(PIPE_FIX); +// CHECK-NEXT: pipe_barrier(PIPE_ALL); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( // CHECK-LABEL: AICORE void tbroadcast_tnotify_release( // CHECK: pto::comm::TBROADCAST( // CHECK-NOT: PTOAS__DCCI_SINGLE_CACHE_LINE -// CHECK: pipe_barrier(PIPE_MTE2); -// CHECK-NEXT: pipe_barrier(PIPE_MTE3); -// CHECK-NEXT: pipe_barrier(PIPE_FIX); +// CHECK: pipe_barrier(PIPE_ALL); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( diff --git a/test/lit/pto/signal_payload_cache_consistency.pto b/test/lit/pto/signal_payload_cache_consistency.pto index 3a6fa9319..cf629d86d 100644 --- a/test/lit/pto/signal_payload_cache_consistency.pto +++ b/test/lit/pto/signal_payload_cache_consistency.pto @@ -3,13 +3,13 @@ // 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. +// 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. // Explicit signal/payload memory-consistency IR lowering. PTOAS no longer runs // the historical MemoryConsistency analysis pass, so this test only checks the -// public CMO/fence ops. The GM fence lowering conservatively drains MTE2, -// MTE3, and FIX before dsb(DSB_DDR). +// public CMO/fence ops. In a vector kernel, the GM fence lowering drains all +// vector pipes with PIPE_ALL before dsb(DSB_DDR). // RUN: ptoas --pto-arch=a3 %s -o - 2>&1 | FileCheck %s @@ -119,9 +119,7 @@ module { } // CHECK-LABEL: AICORE void explicit_gm_fence( -// CHECK: pipe_barrier(PIPE_MTE2); -// CHECK-NEXT: pipe_barrier(PIPE_MTE3); -// CHECK-NEXT: pipe_barrier(PIPE_FIX); +// CHECK: pipe_barrier(PIPE_ALL); // CHECK-NEXT: dsb(DSB_DDR); // CHECK-LABEL: AICORE void whole_cache_cmo( @@ -133,9 +131,7 @@ module { // CHECK-LABEL: AICORE void scalar_release_sequence( // CHECK: {{.*}}[{{.*}}] = {{.*}}; // CHECK: PTOAS__DCCI_SINGLE_CACHE_LINE({{.*}}); -// CHECK-NEXT: pipe_barrier(PIPE_MTE2); -// CHECK-NEXT: pipe_barrier(PIPE_MTE3); -// CHECK-NEXT: pipe_barrier(PIPE_FIX); +// CHECK-NEXT: pipe_barrier(PIPE_ALL); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( @@ -146,8 +142,6 @@ module { // CHECK-LABEL: AICORE void mte_release_sequence( // CHECK: TSTORE( -// CHECK: pipe_barrier(PIPE_MTE2); -// CHECK-NEXT: pipe_barrier(PIPE_MTE3); -// CHECK-NEXT: pipe_barrier(PIPE_FIX); +// CHECK: pipe_barrier(PIPE_ALL); // CHECK-NEXT: dsb(DSB_DDR); // CHECK: pto::comm::TNOTIFY( From 51a086033a45a9245882b74a3e4435c467938165 Mon Sep 17 00:00:00 2001 From: hecrereed <821896444@qq.com> Date: Wed, 29 Jul 2026 12:36:49 +0800 Subject: [PATCH 6/6] fix: preserve authoritative non-NZ layouts --- lib/PTO/IR/PTOLayoutUtils.cpp | 15 +++++++-------- .../pto/issue527_nz_canonical_view_infer.pto | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 8 deletions(-) diff --git a/lib/PTO/IR/PTOLayoutUtils.cpp b/lib/PTO/IR/PTOLayoutUtils.cpp index 6220b375e..38e26cb2c 100644 --- a/lib/PTO/IR/PTOLayoutUtils.cpp +++ b/lib/PTO/IR/PTOLayoutUtils.cpp @@ -260,15 +260,7 @@ bool mlir::pto::isCanonicalNZRoot5D(ArrayRef shape5D, bool mlir::pto::isLayoutCompatible5D(Layout layout, ArrayRef shape, ArrayRef stride, unsigned storageElemBytes) { - auto padded = rightAlignTo5D(shape, stride); - if (!padded) - return false; switch (layout) { - case Layout::NZ: - if (shape.size() != kPTOLayoutRank) - return false; - return isNZViewCompatible5D(padded->shape, padded->stride, - storageElemBytes); case Layout::ND: case Layout::DN: // ND/DN are logical GlobalTensor interpretations, not compactness @@ -279,6 +271,13 @@ bool mlir::pto::isLayoutCompatible5D(Layout layout, ArrayRef shape, case Layout::MX_A_ZZ: case Layout::MX_B_NN: return true; + case Layout::NZ: { + auto padded = rightAlignTo5D(shape, stride); + if (!padded || shape.size() != kPTOLayoutRank) + return false; + return isNZViewCompatible5D(padded->shape, padded->stride, + storageElemBytes); + } } return false; } diff --git a/test/lit/pto/issue527_nz_canonical_view_infer.pto b/test/lit/pto/issue527_nz_canonical_view_infer.pto index 584fed0d9..b70706671 100644 --- a/test/lit/pto/issue527_nz_canonical_view_infer.pto +++ b/test/lit/pto/issue527_nz_canonical_view_infer.pto @@ -9,6 +9,7 @@ // RUN: split-file %s %t // RUN: ptoas --pto-arch=a3 --emit-pto-ir --mlir-print-ir-after=pto-infer-layout %t/positive.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=INFER // RUN: ptoas --pto-arch=a5 --emit-pto-ir --mlir-print-ir-after=pto-infer-layout %t/positive.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=INFER +// RUN: ptoas --pto-arch=a3 --emit-pto-ir --mlir-print-ir-after=pto-infer-layout %t/nd_fallback.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=ND-FALLBACK // RUN: ptoas --pto-arch=a5 --emit-pto-ir --mlir-print-ir-after=pto-view-to-memref %t/positive.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=LOWER // RUN: ptoas --pto-arch=a3 %t/emitc.pto 2>&1 | FileCheck %s --check-prefix=EMITC // RUN: ptoas --pto-arch=a5 %t/emitc.pto 2>&1 | FileCheck %s --check-prefix=EMITC @@ -38,6 +39,8 @@ // INFER-DAG: pto.mscatter // INFER-DAG: memref.subview {{.*}} {layout = #pto.layout} +// ND-FALLBACK: pto.make_tensor_view {{.*}} {layout = #pto.layout, pto.inferred_layout = true} : !pto.tensor_view<4x8x16xf32> + // LOWER-DAG: memref.reinterpret_cast {{.*}} {layout = #pto.layout, pto.inferred_layout = true} // LOWER-DAG: memref.subview {{.*}} {layout = #pto.layout, pto.inferred_layout = true} // LOWER-DAG: memref.subview {{.*}} {layout = #pto.layout, pto.inferred_layout = true} @@ -306,6 +309,22 @@ module { } } +//--- nd_fallback.pto +module { + func.func @nd_fallback_for_unknown_stride(%ptr: !pto.ptr) { + %c-1 = arith.constant -1 : index + %c4 = arith.constant 4 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %view = pto.make_tensor_view %ptr, + shape = [%c4, %c8, %c16], + strides = [%c128, %c16, %c-1] + : !pto.tensor_view<4x8x16xf32> + return + } +} + //--- invalid_inner_shape.pto module { func.func @invalid_inner_shape(%ptr: !pto.ptr) {