diff --git a/ReleaseNotes.md b/ReleaseNotes.md index 79882f63e0..5d5d16d871 100644 --- a/ReleaseNotes.md +++ b/ReleaseNotes.md @@ -7,6 +7,12 @@ ## 变更摘要 - PTOAS 首次发布 +## 未发布变更 +- GlobalTensor layout 推断现在识别 canonical rank-5 NZ view。对 + `pto.mgather` / `pto.mscatter`，若无 layout 的 shape/stride 同时命中该 NZ + 约定，而内核语义实际是 ND，需要在 `pto.make_tensor_view` 上显式写 + `layout = #pto.layout`（PTODSL 使用 `layout="ND"`）。 + ## 概述 PTOAS（PTO Assembler & Optimizer）是面向 PTO Bytecode 的编译器工具链，基于 LLVM/MLIR LLVM21 VPTO 分支 `vpto-dev/llvm-project:feature-vpto-llvm21` 构建。它提供 PTO Dialect 的定义、解析、验证、优化与代码生成能力，并输出可调用 `pto-isa` 的 C++ 代码。 diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md index 6903ff9c5f..fc159f7d33 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 new file mode 100644 index 0000000000..1c4c827787 --- /dev/null +++ b/docs/designs/nz-layout-inference-alignment.md @@ -0,0 +1,433 @@ + + +# PTOAS `Layout::NZ` 推断与 pto-isa 标准五维形对齐 + +关联 issue: [#527](https://github.com/hw-native-sys/PTOAS/issues/527)。 + +本文定义 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 五维标准形 + +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] +``` + +五个维度的固定语义如下： + +| 维度 | 含义 | 标准根视图取值 | +|---|---|---| +| `d0` | 二维场景的占位维 | `1` | +| `d1` | 列分块数 `n1` | `cols / C0` | +| `d2` | 行分块数 `m1` | `rows / 16` | +| `d3` | fractal 内行数 | `16` | +| `d4` | fractal 内列数 | `C0` | + +例如 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 +stride[4] == 1 +stride[3] == C0 +stride[2] == 16 * C0 +stride[1] == shape[2] * stride[2] +stride[0] == shape[1] * stride[1] +``` + +这里的 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. 当前问题 + +PTOAS 当前在三处重复实现 layout 推断： + +| 位置 | 用途 | +|---|---| +| `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 +shape[2] * shape[3] * elemBytes == 512 +stride[4] == 1 +stride[3] == shape[4] +``` + +它相对 pto-isa 标准形错位一维，并且没有验证 `stride[2..0]`。结果是： + +- canonical `[1,8,8,16,8]` fp32 被判成 ND； +- `[4,1,16,8,16]` 这类连续 ND 形状被误判成 NZ； +- 显式 `{layout = #pto.layout}` 仍会被错误的 pattern 推断否决。 + +## 3. 必须保留的语义边界 + +canonical NZ 根视图的数值 stride 同时满足连续五维 ND 的累积积关系。因此仅凭 +shape/stride 无法证明用户的逻辑语义一定是 NZ。 + +本 PR 采用一条确定的解析顺序： + +1. **显式 layout 是权威来源。** 显式 `ND`、`DN`、`NZ` 都不会被 pattern 推断覆盖； +2. **派生 view 继承已经解析的源 layout。** NZ 子视图需要额外通过 fractal 边界校验； +3. **只有无显式属性、无源 layout 的根视图才走 pattern 推断；** +4. pattern 推断把满足 canonical NZ 规则的非退化二维五维形约定为 NZ，其他情况继续 + 使用现有 ND/DN 推断。 + +第 4 条是一项 PTOAS 输入约定，不声称它能从数值上区分所有连续 5D ND。用户需要表达 +同 shape/stride 的 ND 时，必须能够通过显式 `layout = nd` 覆盖该约定。 + +## 4. 唯一实现 + +### 4.1 共享 layout 工具 + +新增 `include/PTO/IR/PTOLayoutUtils.h` 和 +`lib/PTO/IR/PTOLayoutUtils.cpp`，集中实现： + +```cpp +namespace mlir::pto { + +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, + unsigned storageElemBytes, + std::optional preferredMinor2D = std::nullopt, + bool *isMinor2DAmbiguous = nullptr); + +} // namespace mlir::pto +``` + +`getNZC0StorageElems()`、`isNZViewCompatible5D()` 和 +`isCanonicalNZRoot5D()` 仅是该实现文件内部的组合谓词，不作为跨 pass API 导出。 + +`InferPTOLayout.cpp`、`PTO.cpp`、`PTOToEmitC.cpp` 都调用这套实现，删除各自的私有 +NZ/ND/DN 判定副本。调用点必须先解析显式属性和源 view 属性，只有确实缺失 layout 时 +才调用 `inferLayout5D()`。 + +### 4.2 根视图与子视图使用不同谓词 + +`isCanonicalNZRoot5D()` 只用于无 layout 根视图的 pattern 推断，要求完整的紧密 +canonical stride： + +```cpp +NZInnerStructure(shape, stride, bytes) + && stride[1] == shape[2] * stride[2] + && stride[0] == shape[1] * stride[1] +``` + +其中 `NZInnerStructure` 是下列不能被切开的内部 fractal 条件的伪代码： + +```cpp +shape[3] == 16 +shape[4] == C0 +stride[4] == 1 +stride[3] == C0 +stride[2] == 16 * C0 +``` + +`isNZViewCompatible5D()` 在内部结构之上验证当前 view 的外层跨度： + +```cpp +NZInnerStructure(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。 + +### 4.3 无属性根视图的 NZ 约定 + +无显式 layout 且没有可继承源 layout 时，按以下唯一规则推断 NZ： + +```cpp +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`。 + +显式 `layout = nd` 必须覆盖上述 pattern 约定。例如 fp32： + +```text +shape = [1, 2, 1, 16, 8] +stride = [256, 128, 128, 8, 1] +layout = nd +``` + +该 shape/stride 同时符合 canonical NZ 数值形式和连续 5D ND，PTOAS 必须保留用户 +声明的 ND，不得报 layout mismatch。 + +### 4.4 显式 layout 校验 + +显式属性不再与 pattern 推断结果比较，而是按其自身语义校验： + +- 显式 `NZ` 根视图必须满足 NZ 内部 fractal 结构；canonical 紧密 stride 和合法的 + 外层 block gap 都可接受，统一使用 `isNZViewCompatible5D()` 校验； +- 显式 `ND` / `DN` 保持为逻辑布局的权威声明。它们不承诺当前 view 是紧密 + minor-2D 存储，因此不使用 dense stride 递推拒绝已有的 interleaved/gapped + view；minor-2D 歧义处理只用于无显式 layout 的推断； +- 校验失败时报告具体的 shape/stride 不变量，不回退到其他 layout； +- 显式属性不附加 `pto.inferred_layout`。 + +### 4.5 NZ partition/subview 传播 + +源 layout 为 NZ 时，`partition_view` 和 `memref.subview` 按源 view 校验： + +- `d3`、`d4` 必须完整保留：offset 为 0，size 等于源 size； +- `d1`、`d2` 可以缩小，结果保留源 stride； +- `d1`、`d2` 的一个坐标单位本身就是一个完整 block，因此任意合法整数 offset + 都是 block 边界，不再增加含义不清的二次对齐条件； +- 二维 NZ 的 `d0` 保持 `offset = 0`、`size = 1`； +- 校验通过后直接继承 NZ，不再用 `isCanonicalNZRoot5D()` 重新推断； +- `partition_view` 可以继续以 `partition_tensor_view` 为源，嵌套切分逐级沿源链 + 继承同一个 layout；lowering 必须同时保留根 view 的真实 stride 和累计 offset； +- EmitC 无法从源类型或源 view 链解析精确 stride 时必须报错，不能用连续 + row-major stride 猜测； +- 在 `d3`/`d4` 内部切分时发出错误，不静默改成 ND。 + +动态 `d1`/`d2` offset 和 size 不改变 fractal 内部结构，可以传播 NZ。`d3`/`d4` +是否完整必须能够在编译期证明，否则报错。 + +### 4.6 动态 shape/stride + +显式或继承的 layout 在动态 shape 下继续携带。校验所有编译期已知的结构维，不使用 +未知值否决显式 NZ。 + +无属性根视图只有在 NZ 内部结构和非退化门槛都能在编译期证明时才自动推断 NZ；无法 +证明时不猜测，沿用 ND/DN fallback，调用方应提供显式 layout。 + +### 4.7 旧规则处理 + +直接删除旧的错位 NZ 判定，只保留本设计定义的标准规则。不保留旧规则并集，不增加 +legacy 开关。旧规则命中的非标准形属于误判，继续兼容会固定错误的 +`GlobalTensor<..., Layout::NZ>` 模板参数。 + +## 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` 的现有样本扫描得到： + +- 242 个文件生成 1558 个 `GlobalTensor` 实例化点； +- 当前 NZ 去重形状中，两组符合 pto-isa 标准形； +- 两组只存在于 `globaltensor_layout_bytewidth_emitc.pto` 的旧错位形状不再是 NZ： + int8 样例回到 ND；末维为 1 时 ND/DN 的寻址等价，未提供消费者偏好时保持 + 兼容默认 ND； +- 不加非退化门槛时，38 个文件中的 80 个单-fractal ND 视图会被升级成 NZ； +- 加门槛后，当前语料没有新增 NZ。 + +“当前语料没有新增 NZ”仅是回归影响数据，不代表 shape/stride 已经从理论上完成 ND/NZ +消歧。因此实现仍必须保证所有显式 layout 优先，并增加“canonical 数值形式 + +显式 ND”的反例测试。 + +`mgather`/`mscatter` 的 ND-only 校验也必须使用同一解析顺序：显式/继承 layout +优先，缺失时才 pattern 推断。真实 NZ 仍按现有约束报错；显式 ND 不得被 heuristic +改成 NZ。 + +## 7. 本 PR 的实现范围 + +PR #1027 在现有设计提交之后继续完成以下内容： + +- 新增共享 `PTOLayoutUtils`； +- 同步 IR verifier、layout pass 和 EmitC fallback； +- 实现显式 layout 全优先； +- 替换旧 NZ 规则； +- 实现 NZ partition/subview 的 source-relative 校验和传播； +- 补齐静态、动态、packed storage element 和跨调用点一致性测试； +- 更新用户文档及 Python `make_tensor_view` 的显式 layout 示例。 + +不再拆分成多个后续 PR。上述实现和验证全部完成后将 #1027 转为 ready；合入该 PR +即关闭 #527。 + +## 8. 测试要求 + +新增 `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； +15. 带外层 gap 的根 view 经过嵌套 partition 后仍保留原 stride，并累计两级 offset。 + +同时更新 `globaltensor_layout_bytewidth_emitc.pto` 中固化旧错位规则的期望值，并增加 +A3/A5 编译覆盖。具备板卡资源时运行一条 NZ load/store E2E 数据比对，确认生成的 +`GlobalTensor` shape、stride、layout 与 pto-isa 行为一致。 diff --git a/include/PTO/IR/PTOLayoutUtils.h b/include/PTO/IR/PTOLayoutUtils.h new file mode 100644 index 0000000000..38f37dc08c --- /dev/null +++ b/include/PTO/IR/PTOLayoutUtils.h @@ -0,0 +1,43 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +#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 +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 32ee5a8f28..c276ec193c 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 fe99fff0fb..b942f89ceb 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 af6dd28592..b15485b9bf 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 0000000000..0dd8d8fb65 --- /dev/null +++ b/lib/PTO/IR/PTOLayoutUtils.cpp @@ -0,0 +1,293 @@ +// 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; + // Both interpretations address the same elements when a minor extent is + // one. Keep the historical row-major default unless a consumer supplies + // an explicit preference. + return 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 + +static std::optional +getNZC0StorageElems(unsigned storageElemBytes) { + if (storageElemBytes == 0 || kNZBlockBytes % storageElemBytes != 0) + return std::nullopt; + return kNZBlockBytes / storageElemBytes; +} + +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; +} + +static bool 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; +} + +static bool 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) { + switch (layout) { + 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; + case Layout::NZ: { + auto padded = rightAlignTo5D(shape, stride); + if (!padded || shape.size() != kPTOLayoutRank) + return false; + return isNZViewCompatible5D(padded->shape, padded->stride, + storageElemBytes); + } + } + 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 572995a870..14a6017b60 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,29 +394,36 @@ 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 void 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; } static ResolvedLayoutInfo resolveLayoutFromViewValue(Value v) { @@ -524,26 +443,136 @@ 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)) { + op.emitError( + "cannot resolve the source shape needed to validate NZ partition"); + 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 +584,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; + getFoldResults(op.getMixedSizes(), shape, allStatic); + 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 +624,70 @@ 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; + auto existingInferred = + op->getAttrOfType(kInferredLayoutAttrName); + bool inferred = + existing ? (existingInferred && existingInferred.getValue()) + : 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 +710,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/LowerPTOToUBufOps.cpp b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp index f9fe13866d..0e7f6d0cdc 100644 --- a/lib/PTO/Transforms/LowerPTOToUBufOps.cpp +++ b/lib/PTO/Transforms/LowerPTOToUBufOps.cpp @@ -1240,46 +1240,63 @@ struct LowerPTOToUBufOpsPass return !strides.empty() && strides.back() == 1; } - static FailureOr extractDmaViewInfo(pto::TLoadOp op) { - auto pvOp = op.getSrc().getDefiningOp(); - if (pvOp) { - auto mtvOp = pvOp.getSource().getDefiningOp(); - if (!mtvOp) - return failure(); - DmaViewInfo info; - info.gmPtr = mtvOp.getPtr(); - info.sizes.assign(pvOp.getSizes().begin(), pvOp.getSizes().end()); - info.strides.assign(mtvOp.getStrides().begin(), - mtvOp.getStrides().end()); - if (info.strides.empty() || - !matchPattern(info.strides.back(), m_One())) { - op.emitError("A2/A3 DMA lowering requires a unit innermost stride"); + static FailureOr + extractDirectDmaViewInfo(Value view, Operation *consumer) { + SmallVector partitions; + Value root = view; + while (auto partition = root.getDefiningOp()) { + partitions.push_back(partition); + root = partition.getSource(); + } + if (partitions.empty()) + return failure(); + + auto makeView = root.getDefiningOp(); + if (!makeView) + return failure(); + + DmaViewInfo info; + info.gmPtr = makeView.getPtr(); + info.sizes.assign(partitions.front().getSizes().begin(), + partitions.front().getSizes().end()); + info.strides.assign(makeView.getStrides().begin(), + makeView.getStrides().end()); + if (info.strides.empty() || + !matchPattern(info.strides.back(), m_One())) { + consumer->emitError( + "A2/A3 DMA lowering requires a unit innermost stride"); + return failure(); + } + + info.offsets.assign(partitions.front().getOffsets().begin(), + partitions.front().getOffsets().end()); + OpBuilder builder(consumer); + for (pto::PartitionViewOp partition : + llvm::drop_begin(partitions)) { + if (partition.getOffsets().size() != info.offsets.size() || + partition.getSizes().size() != info.sizes.size()) { + consumer->emitError( + "A2/A3 DMA lowering requires rank-preserving nested partitions"); return failure(); } - info.offsets.assign(pvOp.getOffsets().begin(), pvOp.getOffsets().end()); - return info; + for (auto [index, offset] : + llvm::enumerate(partition.getOffsets())) { + info.offsets[index] = builder.create( + consumer->getLoc(), info.offsets[index], offset); + } } + return info; + } + + static FailureOr extractDmaViewInfo(pto::TLoadOp op) { + if (op.getSrc().getDefiningOp()) + return extractDirectDmaViewInfo(op.getSrc(), op.getOperation()); return extractDmaMemRefViewInfo(op.getLoc(), op.getSrc(), op.getContext()); } static FailureOr extractDmaViewInfo(pto::TStoreOp op) { - auto pvOp = op.getDst().getDefiningOp(); - if (pvOp) { - auto mtvOp = pvOp.getSource().getDefiningOp(); - if (!mtvOp) - return failure(); - DmaViewInfo info; - info.gmPtr = mtvOp.getPtr(); - info.sizes.assign(pvOp.getSizes().begin(), pvOp.getSizes().end()); - info.strides.assign(mtvOp.getStrides().begin(), mtvOp.getStrides().end()); - if (info.strides.empty() || - !matchPattern(info.strides.back(), m_One())) { - op.emitError("A2/A3 DMA lowering requires a unit innermost stride"); - return failure(); - } - info.offsets.assign(pvOp.getOffsets().begin(), pvOp.getOffsets().end()); - return info; - } + if (op.getDst().getDefiningOp()) + return extractDirectDmaViewInfo(op.getDst(), op.getOperation()); return extractDmaMemRefViewInfo(op.getLoc(), op.getDst(), op.getContext()); } diff --git a/lib/PTO/Transforms/PTOCanonicalizeIR.cpp b/lib/PTO/Transforms/PTOCanonicalizeIR.cpp index c4495d0d14..28a1f102b7 100644 --- a/lib/PTO/Transforms/PTOCanonicalizeIR.cpp +++ b/lib/PTO/Transforms/PTOCanonicalizeIR.cpp @@ -220,19 +220,29 @@ static LogicalResult rewriteMakeTensorView(MakeTensorViewOp op, static LogicalResult rewritePartitionView(PartitionViewOp op, IRRewriter &rewriter) { - auto sourceType = dyn_cast(op.getSource().getType()); auto resultType = dyn_cast(op.getResult().getType()); - if (!sourceType || !resultType) + if (!resultType) return success(); + int64_t sourceRank = 0; + if (auto sourceType = + dyn_cast(op.getSource().getType())) { + sourceRank = sourceType.getRank(); + } else if (auto sourceType = + dyn_cast( + op.getSource().getType())) { + sourceRank = sourceType.getRank(); + } else { + return success(); + } if (op.getOffsets().size() != kLogicalRank2 || op.getSizes().size() != kLogicalRank2) return success(); - if (sourceType.getRank() != kCanonicalRank5) + if (sourceRank != kCanonicalRank5) return op.emitOpError( "rank-2 partition_tensor_view normalization expects canonical rank-5 " - "source tensor_view"); + "source view"); rewriter.setInsertionPoint(op); Value zero = getOrCreateIndexConstant(rewriter, op.getLoc(), 0); diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 33e5e66439..75869180b5 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" @@ -90,6 +91,9 @@ static std::string getGlobalTensorTypeStringFromShapeAndStrides( static emitc::OpaqueType getGlobalTensorOpaqueTypeFromShape( MLIRContext *ctx, Type elemTy, ArrayRef shape, StringRef layoutEnum = "pto::Layout::ND"); +static Value materializeGlobalTensorDataPointer( + ConversionPatternRewriter &rewriter, Location loc, Value value, + Type sourceType); static const char *addrSpaceQualifier(pto::AddressSpace as) { switch (as) { @@ -3784,6 +3788,74 @@ struct SubviewToEmitCPattern : public OpConversionPattern { return std::nullopt; } + LogicalResult + resolveSourceStrides(Value source, PatternRewriter &rewriter, + SmallVectorImpl &strides) const { + auto sourceType = dyn_cast(source.getType()); + if (!sourceType) + return failure(); + int64_t rank = sourceType.getRank(); + + if (auto reinterpretCast = + source.getDefiningOp()) { + auto mixedStrides = reinterpretCast.getMixedStrides(); + if (mixedStrides.size() != static_cast(rank)) + return failure(); + strides.assign(mixedStrides.begin(), mixedStrides.end()); + return success(); + } + + if (auto subview = source.getDefiningOp()) { + SmallVector parentStrides; + if (failed(resolveSourceStrides(subview.getSource(), rewriter, + parentStrides))) + return failure(); + auto steps = subview.getMixedStrides(); + if (parentStrides.size() != static_cast(rank) || + steps.size() != static_cast(rank)) + return failure(); + + strides.reserve(rank); + for (auto [parentStride, step] : + llvm::zip_equal(parentStrides, steps)) { + auto parentStatic = extractStaticInt(parentStride); + auto stepStatic = extractStaticInt(step); + if (parentStatic && stepStatic) { + strides.push_back( + rewriter.getIndexAttr((*parentStatic) * (*stepStatic))); + continue; + } + if (stepStatic && *stepStatic == 1) { + strides.push_back(parentStride); + continue; + } + if (parentStatic && *parentStatic == 1) { + strides.push_back(step); + continue; + } + return failure(); + } + return success(); + } + + if (auto cast = source.getDefiningOp()) + return resolveSourceStrides(cast.getSource(), rewriter, strides); + + SmallVector typeStrides; + int64_t offset = ShapedType::kDynamic; + if (failed(mlir::pto::getPTOMemRefStridesAndOffset( + sourceType, typeStrides, offset)) || + typeStrides.size() != static_cast(rank) || + llvm::any_of(typeStrides, [](int64_t stride) { + return stride == ShapedType::kDynamic; + })) + return failure(); + + for (int64_t stride : typeStrides) + strides.push_back(rewriter.getIndexAttr(stride)); + return success(); + } + LogicalResult matchAndRewrite(memref::SubViewOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { auto loc = op.getLoc(); @@ -3860,36 +3932,11 @@ struct SubviewToEmitCPattern : public OpConversionPattern { // 1. 获取 Source 的 Strides (支持动态 Stride 收集) SmallVector sourceStrides; - if (auto rc = op.getSource().getDefiningOp()) { - sourceStrides = rc.getMixedStrides(); - } else { - SmallVector strideInts; - int64_t offset = ShapedType::kDynamic; - bool useTypeStrides = - succeeded(mlir::pto::getPTOMemRefStridesAndOffset( - srcType, strideInts, offset)); - (void)offset; - if (useTypeStrides) { - for (int64_t s : strideInts) { - if (s == ShapedType::kDynamic) - useTypeStrides = false; - } - } - if (useTypeStrides) { - for (int64_t s : strideInts) { - sourceStrides.push_back(rewriter.getIndexAttr(s)); - } - } else { - // Fallback: Compact Layout - auto shape = srcType.getShape(); - int64_t current = 1; - sourceStrides.resize(rank); - for (int i = rank - 1; i >= 0; --i) { - sourceStrides[i] = rewriter.getIndexAttr(current); - if (shape[i] != ShapedType::kDynamic) current *= shape[i]; - } - } - } + if (failed(resolveSourceStrides(op.getSource(), rewriter, + sourceStrides))) + return rewriter.notifyMatchFailure( + op, "cannot resolve exact source strides; refusing to assume a " + "compact layout"); // 2. 计算运行时 Offset auto staticOffsets = op.getStaticOffsets(); @@ -3923,7 +3970,12 @@ struct SubviewToEmitCPattern : public OpConversionPattern { // NOTE: Some toolchains may materialize kernel pointer params as `void*` even // when the underlying element type is i16. Pointer arithmetic on `void*` // is ill-formed in C++, so we explicitly cast to a typed pointer for i16. - Value sourcePtr = adaptor.getSource(); + Value convertedSource = adaptor.getSource(); + if (auto cast = + convertedSource.getDefiningOp()) + convertedSource = cast.getOperand(0); + Value sourcePtr = materializeGlobalTensorDataPointer( + rewriter, loc, convertedSource, op.getSource().getType()); Value tileCandidate = sourcePtr; if (auto castOp = sourcePtr.getDefiningOp()) { tileCandidate = castOp.getOperand(); @@ -4118,46 +4170,10 @@ struct SubviewToEmitCPattern : public OpConversionPattern { layoutEnum = specialScaleSpec->layoutEnum; } else if (resolvedLayout) { layoutEnum = layoutToEmitCString(*resolvedLayout); - } else { - bool allStatic = - llvm::all_of(finalShape, [](int64_t value) { return value != -1; }) && - llvm::all_of(finalStride, [](int64_t value) { return value != -1; }); - - int layoutTag = 0; // ND - auto elemBytes = 4; // default float - if (elemTypeStr.find("half") != std::string::npos || - elemTypeStr.find("f16") != std::string::npos || - elemTypeStr.find("bf16") != std::string::npos) - elemBytes = 2; - else if (elemTypeStr.find("double") != std::string::npos || - elemTypeStr.find("f64") != std::string::npos) - elemBytes = 8; - - if (allStatic) { - if (finalShape[2] == 16 && - finalShape[2] * finalShape[3] * elemBytes == 512 && - finalStride[4] == 1 && finalStride[3] == finalShape[4]) { - layoutTag = 2; // NZ - } else { - bool isRow = finalStride[4] == 1; - for (int i = 3; i >= 0; --i) - isRow &= (finalStride[i] == - multiplyOrDynamic(finalStride[i + 1], finalShape[i + 1])); - bool isCol = finalStride[0] == 1; - for (int i = 0; i < 4; ++i) - isCol &= (finalStride[i + 1] == - multiplyOrDynamic(finalStride[i], finalShape[i])); - if (isCol) - layoutTag = 1; // DN - else - layoutTag = isRow ? 0 : 0; // fallback ND - } - } - - if (layoutTag == 1) - layoutEnum = "pto::Layout::DN"; - else if (layoutTag == 2) - layoutEnum = "pto::Layout::NZ"; + } else if (auto inferred = inferLayout5D( + shapeParamsVec, strideTemplateVec, + getPTOStorageElemByteSize(resTy.getElementType()))) { + layoutEnum = layoutToEmitCString(*inferred); } // GlobalTensor takes a Layout non-type template parameter; directly use the // enum constant. @@ -4347,37 +4363,21 @@ static emitc::OpaqueType getGlobalTensorOpaqueTypeFromShape( ctx, getGlobalTensorTypeStringFromShape(elemTy, shape, layoutEnum)); } -static std::string inferFallbackGlobalTensorLayout(ArrayRef shape5D, - ArrayRef stride5D, +static std::string inferFallbackGlobalTensorLayout(ArrayRef shape, + ArrayRef strides, 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(shape, strides, getGlobalTensorElementBytes(elemTy)); + return layoutToEmitCString(layout.value_or(Layout::ND)); } static std::string resolveGlobalTensorLayout(Operation *anchor, Value basePtr, - ArrayRef shape5D, - ArrayRef stride5D, + ArrayRef shape, + ArrayRef strides, Type elemTy) { if (auto layout = resolveLayoutForGlobalTensor(anchor, basePtr)) return layoutToEmitCString(*layout); - return inferFallbackGlobalTensorLayout(shape5D, stride5D, elemTy); + return inferFallbackGlobalTensorLayout(shape, strides, elemTy); } struct GlobalTensorTypeNames { @@ -4439,7 +4439,7 @@ static Value buildGlobalTensorFromMemref(ConversionPatternRewriter &rewriter, rewriter.create( loc, "using " + names.strideTypeName + " = pto::Stride<" + joinIntTemplateParams(stride5D) + ">;"); - layoutEnum = resolveGlobalTensorLayout(anchor, basePtr, shape5D, stride5D, + layoutEnum = resolveGlobalTensorLayout(anchor, basePtr, shape, strides, mrTy.getElementType()); } @@ -4490,6 +4490,12 @@ static Value maybeWrapGlobalMemrefAsGlobalTensor( if (!isGlobal) return loweredValue; + Type loweredType = loweredValue.getType(); + if (auto lvalueTy = dyn_cast(loweredType)) + loweredType = lvalueTy.getValueType(); + if (isEmitCGlobalTensorLikeType(loweredType)) + return loweredValue; + if (Value gt = buildGlobalTensorFromMemref(rewriter, loc, loweredValue, mrTy, anchor, tag)) return gt; @@ -8300,9 +8306,8 @@ static bool parseIntegerTemplateList(StringRef token, StringRef marker, } static LogicalResult getStaticTensorViewStrides( - Value source, Value convertedSource, pto::TensorViewType sourceType, + Value source, Value convertedSource, int64_t rank, SmallVectorImpl &strides) { - int64_t rank = sourceType.getRank(); strides.clear(); if (auto makeView = source.getDefiningOp()) { @@ -8329,9 +8334,7 @@ static LogicalResult getStaticTensorViewStrides( } } - auto fallback = buildRowMajorStrides(sourceType.getShape()); - strides.append(fallback.begin(), fallback.end()); - return success(); + return failure(); } struct PTOPartitionViewToEmitC @@ -8342,14 +8345,25 @@ struct PTOPartitionViewToEmitC LogicalResult matchAndRewrite(mlir::pto::PartitionViewOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { - auto srcTy = dyn_cast(op.getSource().getType()); auto resTy = dyn_cast(op.getResult().getType()); - if (!srcTy || !resTy) + Type srcElemTy; + int64_t srcRank = 0; + if (auto srcTy = dyn_cast(op.getSource().getType())) { + srcElemTy = srcTy.getElementType(); + srcRank = srcTy.getRank(); + } else if (auto srcTy = + dyn_cast( + op.getSource().getType())) { + srcElemTy = srcTy.getElementType(); + srcRank = srcTy.getRank(); + } + if (!srcElemTy || !resTy) return rewriter.notifyMatchFailure( - op, "expected tensor_view source and partition_tensor_view result"); + op, "expected tensor_view or partition_tensor_view source and " + "partition_tensor_view result"); - if (op.getOffsets().size() != static_cast(srcTy.getRank()) || - op.getSizes().size() != static_cast(srcTy.getRank())) + if (op.getOffsets().size() != static_cast(srcRank) || + op.getSizes().size() != static_cast(srcRank)) return rewriter.notifyMatchFailure(op, "rank mismatch"); for (auto [idx, value] : llvm::enumerate(op.getSizes())) { @@ -8365,9 +8379,10 @@ struct PTOPartitionViewToEmitC SmallVector srcStrides; if (failed(getStaticTensorViewStrides(op.getSource(), adaptor.getSource(), - srcTy, srcStrides))) + srcRank, srcStrides))) return rewriter.notifyMatchFailure( - op, "partition_view requires static source strides"); + op, "cannot resolve exact partition source strides; refusing to " + "assume a compact layout"); int64_t staticLinearOffset = 0; SmallVector> dynamicOffsetTerms; for (auto [idx, values] : @@ -8388,15 +8403,14 @@ struct PTOPartitionViewToEmitC } auto *ctx = rewriter.getContext(); - std::string elemTypeStr = getElemTypeStringForGT(srcTy.getElementType()); + std::string elemTypeStr = getElemTypeStringForGT(srcElemTy); auto ptrTy = emitc::PointerType::get( emitc::OpaqueType::get(ctx, "__gm__ " + elemTypeStr)); Value src = peelUnrealized(adaptor.getSource()); - auto data = rewriter - .create( - op.getLoc(), ptrTy, "PTOAS__GLOBAL_TENSOR_DATA", - ArrayAttr{}, ArrayAttr{}, ValueRange{src}) - .getResult(0); + Value data = materializeGlobalTensorDataPointer( + rewriter, op.getLoc(), src, op.getSource().getType()); + if (data.getType() != ptrTy) + data = rewriter.create(op.getLoc(), ptrTy, data); Value ptr = data; if (!dynamicOffsetTerms.empty()) { Type indexTy = emitc::OpaqueType::get(ctx, "int64_t"); @@ -8689,10 +8703,15 @@ struct ReinterpretCastToEmitC : public OpConversionPattern constantOffset = + mixedOffsets.empty() ? std::nullopt + : getConstantIntValue(mixedOffsets.front()); + const bool isZeroOffset = constantOffset && *constantOffset == 0; // GM: keep pointer arithmetic. if (isGm) { - if (!offsetVal) { + if (!offsetVal || (isZeroOffset && !emitAddPtrTrace)) { rewriter.replaceOp(op, source); return success(); } @@ -8815,7 +8834,7 @@ struct ReinterpretCastToEmitC : public OpConversionPattern(loc, u64Ty, offU64).getResult(); diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index 46da6fb596..596b8f7c01 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -1167,7 +1167,8 @@ static LogicalResult lowerPtrToIntOps(func::FuncOp func, MLIRContext *ctx) { Location loc = op.getLoc(); Value baseBuf = op.getOperand(0); - OpFoldResult off0 = rewriter.getIndexAttr(0); + OpFoldResult off0 = + rewriter.create(loc, 0).getResult(); bool foldedAddPtr = false; { Value cur = baseBuf; @@ -1193,11 +1194,16 @@ static LogicalResult lowerPtrToIntOps(func::FuncOp func, MLIRContext *ctx) { size_t rank = op.getShape().size(); int64_t dyn = ShapedType::kDynamic; - SmallVector dynStrides(rank, dyn); + SmallVector resultShape(rank, dyn); + SmallVector resultStrides(rank, dyn); + for (auto [index, value] : llvm::enumerate(op.getStrides())) { + int64_t constant = 0; + if (getConstIndexValue(value, constant)) + resultStrides[index] = constant; + } auto layout = - StridedLayoutAttr::get(ctx, /*offset=*/dyn, /*strides=*/dynStrides); - SmallVector dynShape(rank, dyn); - auto mrTy = MemRefType::get(dynShape, baseMr.getElementType(), layout, + StridedLayoutAttr::get(ctx, dyn, resultStrides); + auto mrTy = MemRefType::get(resultShape, baseMr.getElementType(), layout, baseMr.getMemorySpace()); SmallInlineVector sizes; @@ -1205,7 +1211,7 @@ static LogicalResult lowerPtrToIntOps(func::FuncOp func, MLIRContext *ctx) { sizes.push_back(ensureIndex(rewriter, loc, value, op)); SmallInlineVector strides; for (Value value : op.getStrides()) - strides.push_back(ensureIndex(rewriter, loc, value, op)); + appendMixedIndex(rewriter, loc, value, op, strides); auto rc = rewriter.create(loc, mrTy, baseBuf, off0, sizes, strides); @@ -1213,6 +1219,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(); @@ -1369,8 +1377,61 @@ static LogicalResult lowerPartitionViewOps(func::FuncOp func, MLIRContext *ctx) } int64_t dyn = ShapedType::kDynamic; - SmallVector dynStrides(rank, dyn); - auto layout = StridedLayoutAttr::get(ctx, dyn, dynStrides); + SmallVector resultStrides(rank, dyn); + int64_t sourceOffset = dyn; + SmallVector sourceTypeStrides; + if (succeeded(getPTOMemRefStridesAndOffset( + srcMrTy, sourceTypeStrides, sourceOffset)) && + sourceTypeStrides.size() == static_cast(rank)) + resultStrides.assign(sourceTypeStrides.begin(), + sourceTypeStrides.end()); + if (auto reinterpretCast = + src.getDefiningOp()) { + auto mixedStrides = reinterpretCast.getMixedStrides(); + if (mixedStrides.size() == static_cast(rank)) { + for (auto [idx, stride] : llvm::enumerate(mixedStrides)) { + if (auto attr = dyn_cast(stride)) { + if (auto intAttr = dyn_cast(attr)) { + resultStrides[idx] = intAttr.getInt(); + continue; + } + } + if (auto value = dyn_cast(stride)) { + int64_t constant = 0; + if (getConstIndexValue(value, constant)) + resultStrides[idx] = constant; + } + } + } + } + int64_t resultOffset = sourceOffset; + if (resultOffset != dyn) { + for (auto [offset, stride] : + llvm::zip_equal(mixedOffsets, resultStrides)) { + int64_t offsetValue = 0; + if (auto attr = dyn_cast(offset)) { + auto integerAttr = dyn_cast(attr); + if (!integerAttr) { + resultOffset = dyn; + break; + } + offsetValue = integerAttr.getInt(); + } else { + auto value = dyn_cast(offset); + if (!value || !getConstIndexValue(value, offsetValue)) { + resultOffset = dyn; + break; + } + } + if (stride == dyn) { + resultOffset = dyn; + break; + } + resultOffset += offsetValue * stride; + } + } + auto layout = + StridedLayoutAttr::get(ctx, resultOffset, resultStrides); auto resTy = MemRefType::get(staticSizes, srcMrTy.getElementType(), layout, srcMrTy.getMemorySpace()); @@ -1380,6 +1441,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()); } @@ -2135,7 +2199,8 @@ struct PTOViewToMemrefPass Location loc = op.getLoc(); Value baseBuf = op.getOperand(0); - OpFoldResult off0 = rewriter.getIndexAttr(0); + OpFoldResult off0 = + rewriter.create(loc, 0).getResult(); // Fold pto.addptr chains into the view base to avoid nested reinterpret_cast. bool foldedAddPtr = false; @@ -2165,24 +2230,29 @@ struct PTOViewToMemrefPass // [修复] 获取动态 Rank (根据 shape 输入的数量) size_t rank = op.getShape().size(); - // Construct target type with dynamic offset/strides + // Preserve every static stride operand in the memref type. Keep shape + // and offset dynamic so the view's SSA metadata remains authoritative. Type elemTy = baseMr.getElementType(); int64_t dyn = ShapedType::kDynamic; - - // [修复] 构建 N 维 Strided Layout - // strides 数组长度必须等于 rank - SmallVector dynStrides(rank, dyn); - auto layout = StridedLayoutAttr::get(ctx, /*offset=*/dyn, /*strides=*/dynStrides); - - // [修复] 构建 N 维 Shape - SmallVector dynShape(rank, dyn); - auto mrTy = MemRefType::get(dynShape, elemTy, layout, baseMr.getMemorySpace()); + SmallVector resultShape(rank, dyn); + SmallVector resultStrides(rank, dyn); + for (auto [index, value] : llvm::enumerate(op.getStrides())) { + int64_t constant = 0; + if (getConstIndexValue(value, constant)) + resultStrides[index] = constant; + } + auto layout = + StridedLayoutAttr::get(ctx, dyn, resultStrides); + auto mrTy = MemRefType::get(resultShape, elemTy, layout, + baseMr.getMemorySpace()); SmallInlineVector sizes; - for (Value v : op.getShape()) sizes.push_back(ensureIndex(rewriter, loc, v, op)); + for (Value value : op.getShape()) + sizes.push_back(ensureIndex(rewriter, loc, value, op)); SmallInlineVector strides; - for (Value v : op.getStrides()) strides.push_back(ensureIndex(rewriter, loc, v, op)); + for (Value value : op.getStrides()) + appendMixedIndex(rewriter, loc, value, op, strides); auto rc = rewriter.create( loc, mrTy, baseBuf, off0, sizes, strides); @@ -2192,6 +2262,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 677e23de01..65196123fa 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 8fbc3109e5..a366da8e6a 100644 --- a/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto +++ b/test/lit/pto/globaltensor_layout_bytewidth_emitc.pto @@ -7,11 +7,14 @@ // 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( %src: memref<16x16x16xbf16, #pto.address_space>, %peer0: memref<16x16x16xbf16, #pto.address_space>) { + %loaded = pto.alloc_tile : !pto.tile_buf + pto.tload ins(%src : memref<16x16x16xbf16, #pto.address_space>) outs(%loaded : !pto.tile_buf) %ping = pto.alloc_tile : !pto.tile_buf pto.comm.tbroadcast(%src, recv(%ping), group(%peer0) : memref<16x16x16xbf16, #pto.address_space>, !pto.tile_buf, memref<16x16x16xbf16, #pto.address_space>) {root = 0 : i32} return @@ -36,13 +39,17 @@ module { // A3: using [[BF16_SHAPE:GTShape.*]] = pto::Shape<1, 1, 16, 16, 16>; // A3: using [[BF16_STRIDE:GTStride.*]] = pto::Stride<4096, 4096, 256, 16, 1>; -// A3: constexpr pto::Layout [[BF16_LAYOUT:GT.*_layout]] = pto::Layout::NZ; +// A3: constexpr pto::Layout [[BF16_LAYOUT:GT.*_layout]] = pto::Layout::ND; // A3: using [[BF16_GT:GT.*]] = GlobalTensor; +// A3: using [[BF16_COMM_SHAPE:GTShape.*]] = pto::Shape<1, 1, 16, 16, 16>; +// A3: using [[BF16_COMM_STRIDE:GTStride.*]] = pto::Stride<4096, 4096, 256, 16, 1>; +// A3: constexpr pto::Layout [[BF16_COMM_LAYOUT:GT.*_layout]] = pto::Layout::ND; +// A3: using [[BF16_COMM_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::ND; // 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 0000000000..d2c37e0458 --- /dev/null +++ b/test/lit/pto/issue527_nz_canonical_view_infer.pto @@ -0,0 +1,632 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. +// +// 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 +// RUN: ptoas --pto-arch=a5 --emit-pto-ir --mlir-print-ir-after=pto-view-to-memref %t/emitc.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=GAPPED-LOWER +// 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 +// RUN: not ptoas --pto-arch=a5 %t/partition_layout_mismatch.pto 2>&1 | FileCheck %s --check-prefix=PARTITION-MISMATCH +// RUN: not ptoas --pto-arch=a5 %t/subview_layout_mismatch.pto 2>&1 | FileCheck %s --check-prefix=SUBVIEW-MISMATCH +// RUN: ptoas --pto-arch=a5 --emit-pto-ir --mlir-print-ir-after=pto-infer-layout %t/explicit_false_inferred_attr.pto -o /dev/null 2>&1 | FileCheck %s --check-prefix=FALSE-INFERRED + +// 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} + +// 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} +// 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> +// EMITC-LABEL: AICORE void gapped_nested_partition_emitc( +// EMITC: GlobalTensor, pto::Stride<16384, 2048, 128, 8, 1>, pto::Layout::NZ> +// EMITC: PTOAS__GLOBAL_TENSOR_DATA +// EMITC: pto::Shape<1, 2, 8, 16, 8> +// EMITC: pto::Stride<16384, 2048, 128, 8, 1> +// EMITC: GlobalTensor, pto::Stride<16384, 2048, 128, 8, 1>, pto::Layout::NZ> + +// GAPPED-LOWER-DAG: memref.reinterpret_cast {{.*}} to memref{{.*}}> +// GAPPED-LOWER-DAG: memref.subview {{.*}} : memref<1x3x8x16x8xf32, strided<[16384, 2048, 128, 8, 1], offset: ?>{{.*}}> to memref<1x2x8x16x8xf32, strided<[16384, 2048, 128, 8, 1], offset: ?>{{.*}}> + +// 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) +// PARTITION-MISMATCH: partition layout=nd does not match source layout=nz +// SUBVIEW-MISMATCH: subview layout=nd does not match source layout=nz +// FALSE-INFERRED: memref.subview {{.*}} {layout = #pto.layout} +// FALSE-INFERRED-NOT: pto.inferred_layout + +//--- 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.inferred_layout = true} + : !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 + } + + func.func @gapped_nested_partition_emitc(%ptr: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c2 = arith.constant 2 : index + %c3 = arith.constant 3 : index + %c4 = arith.constant 4 : index + %c8 = arith.constant 8 : index + %c16 = arith.constant 16 : index + %c128 = arith.constant 128 : index + %c2048 = arith.constant 2048 : index + %c16384 = arith.constant 16384 : index + %root = pto.make_tensor_view %ptr, + shape = [%c1, %c4, %c8, %c16, %c8], + strides = [%c16384, %c2048, %c128, %c8, %c1] + {layout = #pto.layout} + : !pto.tensor_view<1x4x8x16x8xf32> + %outer = pto.partition_view %root, + offsets = [%c0, %c1, %c0, %c0, %c0], + sizes = [%c1, %c3, %c8, %c16, %c8] + : !pto.tensor_view<1x4x8x16x8xf32> + -> !pto.partition_tensor_view<1x3x8x16x8xf32> + %nested = pto.partition_view %outer, + offsets = [%c0, %c1, %c0, %c0, %c0], + sizes = [%c1, %c2, %c8, %c16, %c8] + : !pto.partition_tensor_view<1x3x8x16x8xf32> + -> !pto.partition_tensor_view<1x2x8x16x8xf32> + %tile = pto.alloc_tile + : !pto.tile_buf + pto.tload + ins(%nested : !pto.partition_tensor_view<1x2x8x16x8xf32>) + outs(%tile : !pto.tile_buf) + return + } +} + +//--- 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) { + %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 + } +} + +//--- partition_layout_mismatch.pto +module { + func.func @partition_layout_mismatch(%ptr: !pto.ptr) { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : 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] + {layout = #pto.layout} + : !pto.tensor_view<1x8x8x16x8xf32> + %part = pto.partition_view %root, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c4, %c8, %c16, %c8] + {layout = #pto.layout} + : !pto.tensor_view<1x8x8x16x8xf32> + -> !pto.partition_tensor_view<1x4x8x16x8xf32> + return + } +} + +//--- subview_layout_mismatch.pto +module { + func.func @subview_layout_mismatch( + %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, 0, 0, 0, 0] + [1, 4, 8, 16, 8] [1, 1, 1, 1, 1] + {layout = #pto.layout} + : memref<1x8x8x16x8xf32, + strided<[8192, 1024, 128, 8, 1]>, + #pto.address_space> + to memref<1x4x8x16x8xf32, + strided<[8192, 1024, 128, 8, 1]>, + #pto.address_space> + %tile = pto.alloc_tile + : !pto.tile_buf + pto.tload + ins(%part : memref<1x4x8x16x8xf32, + strided<[8192, 1024, 128, 8, 1]>, + #pto.address_space>) + outs(%tile : !pto.tile_buf) + return + } +} + +//--- explicit_false_inferred_attr.pto +module { + func.func @explicit_false_inferred_attr( + %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, 0, 0, 0, 0] + [1, 4, 8, 16, 8] [1, 1, 1, 1, 1] + {layout = #pto.layout, pto.inferred_layout = false} + : memref<1x8x8x16x8xf32, + strided<[8192, 1024, 128, 8, 1]>, + #pto.address_space> + to memref<1x4x8x16x8xf32, + strided<[8192, 1024, 128, 8, 1]>, + #pto.address_space> + %tile = pto.alloc_tile + : !pto.tile_buf + pto.tload + ins(%part : memref<1x4x8x16x8xf32, + strided<[8192, 1024, 128, 8, 1]>, + #pto.address_space>) + outs(%tile : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/issue783_canonicalize_rank2_nd_views.pto b/test/lit/pto/issue783_canonicalize_rank2_nd_views.pto index 1593d63053..aaec3af387 100644 --- a/test/lit/pto/issue783_canonicalize_rank2_nd_views.pto +++ b/test/lit/pto/issue783_canonicalize_rank2_nd_views.pto @@ -10,13 +10,14 @@ // correct cumulative-product strides and that the layout attribute is // preserved after rank-5 padding. -// RUN: ptoas --pto-arch=a5 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-canonicalize-ir %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR +// RUN: ptoas --pto-arch=a3 --pto-backend=vpto --emit-pto-ir --mlir-print-ir-after=pto-canonicalize-ir %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=IR -module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { +module attributes {pto.target_arch = "a3", pto.kernel_kind = #pto.kernel_kind} { func.func @canonicalize_rank2_views(%src: !pto.ptr, %dst: !pto.ptr) attributes {pto.kernel_kind = #pto.kernel_kind} { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index %c16 = arith.constant 16 : index + %c256 = arith.constant 256 : index %c512 = arith.constant 512 : index %c8192 = arith.constant 8192 : index @@ -25,15 +26,18 @@ module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} : !pto.tensor_view %src_part = pto.partition_view %src_view, offsets = [%c0, %c512], sizes = [%c16, %c512] : !pto.tensor_view -> !pto.partition_tensor_view<16x512xbf16> %dst_part = pto.partition_view %dst_view, offsets = [%c0, %c512], sizes = [%c16, %c512] : !pto.tensor_view -> !pto.partition_tensor_view<16x512xbf16> - %tile = pto.declare_tile -> !pto.tile_buf - pto.tload ins(%src_part : !pto.partition_tensor_view<16x512xbf16>) outs(%tile : !pto.tile_buf) + %nested_src = pto.partition_view %src_part, offsets = [%c0, %c0], sizes = [%c16, %c256] : !pto.partition_tensor_view<16x512xbf16> -> !pto.partition_tensor_view<16x256xbf16> + %tile = pto.declare_tile -> !pto.tile_buf + pto.tload ins(%nested_src : !pto.partition_tensor_view<16x256xbf16>) outs(%tile : !pto.tile_buf) pto.section.vector { } return } } -// IR: pto.make_tensor_view {{.*}} shape = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}], strides = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}] {layout = #pto.layout} : !pto.tensor_view<1x1x1x?x?xbf16> -// IR: pto.partition_view {{.*}} offsets = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}], sizes = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}] : !pto.tensor_view<1x1x1x?x?xbf16> -// IR: !pto.partition_tensor_view<1x1x1x16x512xbf16> +// IR-DAG: pto.make_tensor_view {{.*}} shape = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}], strides = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}] {layout = #pto.layout} : !pto.tensor_view<1x1x1x?x?xbf16> +// IR-DAG: pto.partition_view {{.*}} offsets = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}], sizes = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}] : !pto.tensor_view<1x1x1x?x?xbf16> +// IR-DAG: pto.partition_view {{.*}} offsets = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}], sizes = [{{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^,]+}}, {{%[^]]+}}] : !pto.partition_tensor_view<1x1x1x16x512xbf16> +// IR-DAG: pto.tload ins({{.*}} : !pto.partition_tensor_view<1x1x1x16x256xbf16>) // IR-NOT: !pto.partition_tensor_view<16x512xbf16> +// IR-NOT: !pto.partition_tensor_view<16x256xbf16> diff --git a/tools/ptobc/testdata/tstore_fp_v0_roundtrip.pto b/tools/ptobc/testdata/tstore_fp_v0_roundtrip.pto index dab24f2324..5d955c2e29 100644 --- a/tools/ptobc/testdata/tstore_fp_v0_roundtrip.pto +++ b/tools/ptobc/testdata/tstore_fp_v0_roundtrip.pto @@ -10,12 +10,14 @@ module { func.func @tstore_fp_v0(%dst: !pto.ptr) { %c0 = arith.constant 0 : index %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index %c32 = arith.constant 32 : index %dst_tv = pto.make_tensor_view %dst, shape = [%c32, %c32], strides = [%c32, %c1] : !pto.tensor_view<32x32xi8> %dst_part = pto.partition_view %dst_tv, offsets = [%c0, %c0], sizes = [%c32, %c32] : !pto.tensor_view<32x32xi8> -> !pto.partition_tensor_view<32x32xi8> + %nested_dst = pto.partition_view %dst_part, offsets = [%c0, %c0], sizes = [%c16, %c32] : !pto.partition_tensor_view<32x32xi8> -> !pto.partition_tensor_view<16x32xi8> %acc_tile = pto.alloc_tile : !pto.tile_buf %fp_tile = pto.alloc_tile : !pto.tile_buf - pto.tstore_fp ins(%acc_tile, %fp_tile : !pto.tile_buf, !pto.tile_buf) outs(%dst_part : !pto.partition_tensor_view<32x32xi8>) + pto.tstore_fp ins(%acc_tile, %fp_tile : !pto.tile_buf, !pto.tile_buf) outs(%nested_dst : !pto.partition_tensor_view<16x32xi8>) return } } diff --git a/tools/ptobc/tests/tstore_fp_v0_encode.sh b/tools/ptobc/tests/tstore_fp_v0_encode.sh index 7d30d739a2..196055f227 100755 --- a/tools/ptobc/tests/tstore_fp_v0_encode.sh +++ b/tools/ptobc/tests/tstore_fp_v0_encode.sh @@ -32,4 +32,7 @@ ROUNDTRIP="${OUT_DIR}/tstore_fp_v0_roundtrip.roundtrip.pto" "${PTOBC_BIN}" decode "${BC}" -o "${ROUNDTRIP}" grep -F "pto.tstore_fp ins(" "${ROUNDTRIP}" >/dev/null +grep -F "pto.partition_view" "${ROUNDTRIP}" | + grep -F ": !pto.partition_tensor_view<32x32xi8>" >/dev/null +grep -F "!pto.partition_tensor_view<16x32xi8>" "${ROUNDTRIP}" >/dev/null grep -F "!pto.partition_tensor_view<32x32xi8>" "${ROUNDTRIP}" >/dev/null