From b17eceeafa0961bcfdf0c70f4e3fe5ea888c5555 Mon Sep 17 00:00:00 2001 From: peanutchan Date: Mon, 3 Aug 2026 22:31:02 +0800 Subject: [PATCH 1/9] fix(ptodsl): coerce MLIR index bases to i32 for pto.vmi.vci MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dynamic loop IVs (TileLang T.serial / scf.for) arrive as MLIR index. VCI ODS/verify require an integer/float sreg element type; default index→i32 so dynamic bases lower to VCI Vd, Sn like Ascend S.vci. Co-authored-by: Cursor --- ptodsl/ptodsl/_vmi_namespace.py | 21 ++++++- ptodsl/tests/test_vmi_vci_dynamic_index.py | 73 ++++++++++++++++++++++ 2 files changed, 92 insertions(+), 2 deletions(-) create mode 100644 ptodsl/tests/test_vmi_vci_dynamic_index.py diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index 1c4695a0cb..fa8372d3e3 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -12,7 +12,17 @@ from collections.abc import Sequence from ptoas.mlir.dialects import pto as _pto -from ptoas.mlir.ir import BF16Type, F16Type, F32Type, Float8E4M3FNType, Float8E5M2Type, IntegerType, MemRefType, UnitAttr +from ptoas.mlir.ir import ( + BF16Type, + F16Type, + F32Type, + Float8E4M3FNType, + Float8E5M2Type, + IndexType, + IntegerType, + MemRefType, + UnitAttr, +) from ._scalar_coercion import coerce_scalar_to_type from ._surface_values import _coerce_index_value, _try_get_constant_index, unwrap_surface_value, wrap_surface_value @@ -269,7 +279,14 @@ def _derive_vci_result_type(base, size, *, context: str): f"{context} requires a typed scalar such as pto.i32(0) or " "pto.f32(0.0); plain Python scalars are ambiguous" ) - return _pto.VMIVRegType.get(size, raw_base.type) + elem_type = raw_base.type + # Dynamic loop indices (TileLang T.serial / scf.for IVs) are MLIR index. + # VCI requires an integer/float sreg element type; Ascend uses i32. + # Coerce index → signless i32 so pto.vmi.vci(dynamic_base) lowers to + # ``VCI Vd, Sn`` instead of failing ODS (index) or verify (i64). + if IndexType.isinstance(elem_type): + elem_type = IntegerType.get_signless(32) + return _pto.VMIVRegType.get(size, elem_type) def _derive_vmull_result_types(a, b, *, context: str): diff --git a/ptodsl/tests/test_vmi_vci_dynamic_index.py b/ptodsl/tests/test_vmi_vci_dynamic_index.py new file mode 100644 index 0000000000..25fd53f893 --- /dev/null +++ b/ptodsl/tests/test_vmi_vci_dynamic_index.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# 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. +"""Regression: pto.vmi.vci accepts a dynamic MLIR index base (loop IV). + +Dynamic loop indices must coerce to an i32 sreg so ODS/verify accept the op +and lowering emits ``VCI Vd, Sn`` (matches Ascend ``S.vci(T.int32(offset))``). + +Run: + python3 ptodsl/tests/test_vmi_vci_dynamic_index.py +""" + +from __future__ import annotations + +from ptodsl import pto + + +def expect(condition: bool, message: str) -> None: + if not condition: + raise AssertionError(message) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_vci_const_i32_probe(): + dst = pto.alloc_tile(shape=[1, 64], dtype=pto.i32) + offset = pto.const(0, dtype=pto.index) + idx = pto.vmi.vci(pto.i32(0), size=64) + pto.vmi.vstore(idx, dst.as_ptr(), offset) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_vci_dynamic_index_probe(): + """vci(pass_id * 64) where pass_id is an MLIR index loop IV.""" + dst = pto.alloc_tile(shape=[1, 128], dtype=pto.i32) + # AST rewrite turns range(...) into a dynamic index IV (scf.for). + for pass_id in range(2): + base = pass_id * 64 + idx = pto.vmi.vci(base, size=64) + pto.vmi.vstore(idx, dst.as_ptr(), base) + + +def main() -> None: + const_text = vmi_vci_const_i32_probe.compile().mlir_text() + expect("pto.vmi.vci" in const_text, "const probe must emit pto.vmi.vci") + expect( + ": i32 -> !pto.vmi.vreg" in const_text, + f"const probe must use i32 vci:\n{const_text[:800]}", + ) + + dyn_text = vmi_vci_dynamic_index_probe.compile().mlir_text() + expect("pto.vmi.vci" in dyn_text, "dynamic probe must emit pto.vmi.vci") + expect( + "index -> !pto.vmi.vreg" not in dyn_text, + "dynamic vci must not keep index as result element type", + ) + expect( + ": i32 -> !pto.vmi.vreg" in dyn_text, + f"dynamic probe must coerce index→i32 vci:\n{dyn_text[:1600]}", + ) + expect( + "arith.index_cast" in dyn_text or "index_cast" in dyn_text, + f"dynamic probe must index_cast before vci:\n{dyn_text[:1600]}", + ) + print("ptodsl_vmi_vci_dynamic_index: PASS") + + +if __name__ == "__main__": + main() From 4f82b6bd91b06d5dceb4a6906ecd6cff831a1e04 Mon Sep 17 00:00:00 2001 From: peanutchan Date: Mon, 3 Aug 2026 23:41:25 +0800 Subject: [PATCH 2/9] fix(vmi): lower grouped contiguous iota via vci(0)+vadds Dynamic group=2 (VL128) failed on camodel when contiguous iota used raw VCI with a register Sn base. Restore {group} on vci/iota, keep group-periodic lane offsets, and materialize contiguous chunks as vci(0)+vadds like deinterleaved/ASC so dynamic bases rematerialize. Co-authored-by: Cursor --- include/PTO/IR/VMIOps.td | 20 +++- lib/PTO/IR/VMI.cpp | 16 +++ lib/PTO/Transforms/VMILayoutRematerialize.cpp | 3 +- .../Transforms/VMILowerUnifiedToLegacy.cpp | 4 +- lib/PTO/Transforms/VMIToVPTO.cpp | 108 +++++++++++++++++- ptodsl/ptodsl/_vmi_namespace.py | 20 +++- ptodsl/tests/test_vmi_vci_dynamic_index.py | 27 +++++ test/lit/vmi_new/vmi_to_vpto_iota.pto | 27 +++-- test/lit/vmi_new/vmi_to_vpto_iota_group2.pto | 37 ++++++ test/lit/vmi_new/vmi_to_vpto_iota_tail.pto | 10 +- 10 files changed, 252 insertions(+), 20 deletions(-) create mode 100644 test/lit/vmi_new/vmi_to_vpto_iota_group2.pto diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index 29fbe5c6d5..731f4c5fa7 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -62,9 +62,19 @@ def VMIBroadcastOp : VMI_Op<"broadcast", [Pure]> { def VMIIotaOp : VMI_Op<"iota", [Pure]> { let summary = "Create a VMI logical index vector from a scalar base"; + let description = [{ + Without `{group}`, produces a contiguous ramp + `dst[i] = base + i` (ASC) over the full logical vector. + + With `{group = C}`, produces a **group-periodic** ramp: each of the C + groups of size `S = L / C` independently gets + `dst[g*S + j] = base + j` for `j in [0, S)`. Example for + `L=128, C=2, base=0`: `[0..63 | 0..63]`. + }]; let arguments = (ins AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, - OptionalAttr:$order + OptionalAttr:$order, + OptionalAttr:$group ); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; @@ -821,9 +831,15 @@ def VMIVbrcOp : VMI_Op<"vbrc", [Pure]> { def VMIVciOp : VMI_Op<"vci", [Pure]> { let summary = "Create a VMI logical index vector from a scalar base"; + let description = [{ + Unified form of `iota`. Without `{group}`, contiguous ramp over L lanes. + With `{group = C}`, group-periodic ramp: each group restarts at `base` + (see `iota`). Lowers to `iota` preserving `{group}`. + }]; let arguments = (ins AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, - OptionalAttr:$order + OptionalAttr:$order, + OptionalAttr:$group ); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index 01b72448b1..c2fd5a5ce1 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -955,6 +955,14 @@ LogicalResult VMIIotaOp::verify() { if (*order != "ASC" && *order != "DESC") return emitOpError("requires order to be ASC or DESC"); } + if (auto groupAttr = getGroupAttr()) { + int64_t numGroups = groupAttr.getInt(); + if (numGroups <= 0) + return emitOpError("requires group to be positive"); + if (resultType.getElementCount() % numGroups != 0) + return emitOpError("requires group to evenly divide result logical lane " + "count"); + } return success(); } @@ -2604,6 +2612,14 @@ LogicalResult VMIVciOp::verify() { if (*order != "ASC" && *order != "DESC") return emitOpError("requires order to be ASC or DESC"); } + if (auto groupAttr = getGroupAttr()) { + int64_t numGroups = groupAttr.getInt(); + if (numGroups <= 0) + return emitOpError("requires group to be positive"); + if (resultType.getElementCount() % numGroups != 0) + return emitOpError("requires group to evenly divide result logical lane " + "count"); + } return success(); } diff --git a/lib/PTO/Transforms/VMILayoutRematerialize.cpp b/lib/PTO/Transforms/VMILayoutRematerialize.cpp index b40a9b58d9..9318923314 100644 --- a/lib/PTO/Transforms/VMILayoutRematerialize.cpp +++ b/lib/PTO/Transforms/VMILayoutRematerialize.cpp @@ -242,7 +242,8 @@ static std::optional rematerializeDataProducer(Value value, if (auto iota = value.getDefiningOp()) return builder - .create(loc, resultType, iota.getBase(), iota.getOrderAttr()) + .create(loc, resultType, iota.getBase(), iota.getOrderAttr(), + iota.getGroupAttr()) .getResult(); return std::nullopt; diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index d2e498ff06..cc8d26d2b9 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -1189,7 +1189,7 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() { // ---- Category A: pure syntactic renames ---- if (auto vop = dyn_cast(op)) { - // vci -> iota + // vci -> iota (preserve optional {group}) builder.setInsertionPoint(op); StringAttr orderAttr; if (auto order = vop.getOrder()) @@ -1197,7 +1197,7 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() { Value result = builder .create(op->getLoc(), vop.getResult().getType(), - vop.getBase(), orderAttr) + vop.getBase(), orderAttr, vop.getGroupAttr()) .getResult(); vop.getResult().replaceAllUsesWith(result); op->erase(); diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 79dfa5a93a..a045f2133e 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -5207,12 +5207,47 @@ FailureOr createIotaContiguousChunk(Location loc, Type resultType, Value base, int64_t laneOffset, StringAttr orderAttr, PatternRewriter &rewriter) { + // Camodel / Bisheng llvm.hivm.vci honors immediate bases but currently + // ignores a non-zero *register* Sn base (dynamic sreg). Materialize + // absolute indices as vci(0) ± vadds/vdup(base), matching the + // deinterleaved iota path and Ascend's S.vci(0)+vadds pattern. + auto vregType = dyn_cast(resultType); + if (!vregType) + return failure(); + StringRef order = orderAttr ? orderAttr.getValue() : StringRef("ASC"); FailureOr chunkBase = createIotaChunkBase(loc, base, laneOffset, order, rewriter); if (failed(chunkBase)) return failure(); - return rewriter.create(loc, resultType, *chunkBase, orderAttr) + + FailureOr mask = createAllTrueMaskForVReg(loc, vregType, rewriter); + FailureOr zero = + createScalarOffsetConstant(loc, base.getType(), 0, rewriter); + if (failed(mask) || failed(zero)) + return failure(); + + Value local = + rewriter.create(loc, resultType, *zero, StringAttr{}).getResult(); + + // Fast-path: base+laneOffset folds to zero and ASC → plain vci(0). + if (order != "DESC") { + if (auto constBase = chunkBase->getDefiningOp()) { + if (auto intAttr = dyn_cast(constBase.getValue())) { + if (intAttr.getValue().isZero()) + return local; + } + } + return rewriter.create(loc, resultType, local, *chunkBase, *mask) + .getResult(); + } + + Value baseVector = + rewriter + .create(loc, resultType, *chunkBase, *mask, + /*position=*/nullptr) + .getResult(); + return rewriter.create(loc, resultType, baseVector, local, *mask) .getResult(); } @@ -5294,6 +5329,77 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { SmallVector results; results.reserve(resultTypes.size()); + // Optional {group = C}: group-periodic ramp. Each group of size + // S = L / C independently gets dst[g*S + j] = base + j. Contiguous + // physical part p therefore uses laneOffset = (p * lanesPerPart) % S + // (not p * lanesPerPart). Without {group}, keep the contiguous / + // deinterleaved absolute ramps below. + if (auto groupAttr = op.getGroupAttr()) { + int64_t numGroups = groupAttr.getInt(); + int64_t logicalLanes = resultVMIType.getElementCount(); + if (numGroups <= 0 || logicalLanes % numGroups != 0) + return rewriter.notifyMatchFailure( + op, "grouped iota requires group to divide logical lane count"); + int64_t groupSize = logicalLanes / numGroups; + if (groupSize % *lanesPerPart != 0) + return rewriter.notifyMatchFailure( + op, "grouped iota requires group_size to be a multiple of " + "physical lanes per part"); + + if (layout.isContiguous()) { + if (static_cast(resultTypes.size()) * *lanesPerPart != + logicalLanes) + return rewriter.notifyMatchFailure( + op, "grouped contiguous iota physical result count mismatch"); + for (auto [index, resultType] : llvm::enumerate(resultTypes)) { + if (!isa(resultType)) + return rewriter.notifyMatchFailure(op, "iota result must be vreg"); + int64_t laneOffset = + (static_cast(index) * *lanesPerPart) % groupSize; + FailureOr result = createIotaContiguousChunk( + op.getLoc(), resultType, *base, laneOffset, op.getOrderAttr(), + rewriter); + if (failed(result)) + return rewriter.notifyMatchFailure( + op, "failed to materialize grouped contiguous iota chunk"); + results.push_back(*result); + } + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } + + int64_t factor = layout.getFactor(); + if (factor != numGroups) + return rewriter.notifyMatchFailure( + op, "grouped deinterleaved iota requires layout factor == group"); + if (resultTypes.size() % factor != 0) + return rewriter.notifyMatchFailure( + op, "deinterleaved iota physical result count does not match " + "layout factor"); + int64_t chunksPerPart = resultTypes.size() / factor; + if (chunksPerPart * *lanesPerPart != groupSize) + return rewriter.notifyMatchFailure( + op, "grouped deinterleaved iota chunks-per-part mismatch"); + // Group-periodic deinterleaved: each part is one group and restarts + // at `base` (partOffset ignores the group index). + for (int64_t part = 0; part < factor; ++part) { + for (int64_t chunk = 0; chunk < chunksPerPart; ++chunk) { + Type resultType = resultTypes[part * chunksPerPart + chunk]; + FailureOr result = createIotaDeinterleavedChunk( + op.getLoc(), resultType, *base, /*factor=*/1, /*part=*/0, chunk, + *lanesPerPart, op.getOrderAttr(), rewriter); + if (failed(result)) + return rewriter.notifyMatchFailure( + op, "failed to materialize grouped deinterleaved iota chunk"); + results.push_back(*result); + } + } + replaceOpWithFlatConvertedValues(rewriter, op, results, + *this->getTypeConverter()); + return success(); + } + if (layout.isContiguous()) { for (auto [index, resultType] : llvm::enumerate(resultTypes)) { if (!isa(resultType)) diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index fa8372d3e3..3eb7fdf3bc 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -673,14 +673,26 @@ def vsstb(value, destination, offset, block_stride, mask, *, pmode=None, loc=Non ) @staticmethod - def vci(base, *, size, order=None, loc=None, ip=None): - result_type = _derive_vci_result_type(base, size, context="pto.vmi.vci(...)") + def vci(base, *, size, order=None, group=None, loc=None, ip=None): + context = "pto.vmi.vci(...)" + if group is not None: + if isinstance(group, bool) or not isinstance(group, int): + raise TypeError(f"{context} requires group to be a positive Python integer") + if group <= 0: + raise ValueError(f"{context} requires group to be positive, got {group!r}") + if size % group != 0: + raise ValueError( + f"{context} requires size divisible by group; got size={size!r}, group={group!r}" + ) + result_type = _derive_vci_result_type(base, size, context=context) base = coerce_scalar_to_type( base, - _vmi_element_type(result_type, context="pto.vmi.vci(...)"), + _vmi_element_type(result_type, context=context), context="pto.vmi.vci(base)", ) - return _call_value("vci", result_type, base, order=order, loc=loc, ip=ip) + return _call_value( + "vci", result_type, base, order=order, group=group, loc=loc, ip=ip + ) vadd = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vadd", lhs, rhs, mask, **kw)) vsub = staticmethod(lambda lhs, rhs, mask=None, **kw: _emit_binary("vsub", lhs, rhs, mask, **kw)) diff --git a/ptodsl/tests/test_vmi_vci_dynamic_index.py b/ptodsl/tests/test_vmi_vci_dynamic_index.py index 25fd53f893..a53e2abbfd 100644 --- a/ptodsl/tests/test_vmi_vci_dynamic_index.py +++ b/ptodsl/tests/test_vmi_vci_dynamic_index.py @@ -11,6 +11,8 @@ Dynamic loop indices must coerce to an i32 sreg so ODS/verify accept the op and lowering emits ``VCI Vd, Sn`` (matches Ascend ``S.vci(T.int32(offset))``). +Also covers ``group=2`` (VL128 group-periodic) with a dynamic base. + Run: python3 ptodsl/tests/test_vmi_vci_dynamic_index.py """ @@ -44,6 +46,16 @@ def vmi_vci_dynamic_index_probe(): pto.vmi.vstore(idx, dst.as_ptr(), base) +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_vci_dynamic_group2_probe(): + """Dynamic base + group=2 → VL128 group-periodic iota ([0..63|0..63]+base).""" + dst = pto.alloc_tile(shape=[1, 128], dtype=pto.i32) + for pass_id in range(2): + base = pass_id * 64 + idx = pto.vmi.vci(base, size=128, group=2) + pto.vmi.vstore(idx, dst.as_ptr(), pto.const(0, dtype=pto.index)) + + def main() -> None: const_text = vmi_vci_const_i32_probe.compile().mlir_text() expect("pto.vmi.vci" in const_text, "const probe must emit pto.vmi.vci") @@ -66,6 +78,21 @@ def main() -> None: "arith.index_cast" in dyn_text or "index_cast" in dyn_text, f"dynamic probe must index_cast before vci:\n{dyn_text[:1600]}", ) + + g2_text = vmi_vci_dynamic_group2_probe.compile().mlir_text() + expect("pto.vmi.vci" in g2_text, "group=2 probe must emit pto.vmi.vci") + expect( + "group = 2" in g2_text or "{group = 2" in g2_text, + f"group=2 probe must preserve group attr:\n{g2_text[:2000]}", + ) + expect( + ": i32 -> !pto.vmi.vreg" in g2_text, + f"group=2 probe must coerce index→i32 vci:\n{g2_text[:2000]}", + ) + expect( + "arith.index_cast" in g2_text or "index_cast" in g2_text, + f"group=2 probe must index_cast before vci:\n{g2_text[:2000]}", + ) print("ptodsl_vmi_vci_dynamic_index: PASS") diff --git a/test/lit/vmi_new/vmi_to_vpto_iota.pto b/test/lit/vmi_new/vmi_to_vpto_iota.pto index eebca50ef2..6b88ef8b6d 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota.pto @@ -61,20 +61,30 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_iota_i32_asc( -// CHECK: %[[P0:.*]] = pto.vci %arg0 : i32 -> !pto.vreg<64xi32> +// CHECK-SAME: %[[BASE:.*]]: i32 +// CHECK: %[[C0:.*]] = arith.constant 0 : i32 +// CHECK: %[[Z0:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> +// CHECK: %[[P0:.*]] = pto.vadds %[[Z0]], %[[BASE]] // CHECK: arith.constant 64 : i32 -// CHECK: arith.addi %arg0 -// CHECK: %[[P1:.*]] = pto.vci +// CHECK: %[[BASE64:.*]] = arith.addi %[[BASE]] +// CHECK: %[[Z1:.*]] = pto.vci {{.*}} : i32 -> !pto.vreg<64xi32> +// CHECK: %[[P1:.*]] = pto.vadds %[[Z1]], %[[BASE64]] // CHECK: return %[[P0]], %[[P1]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast // CHECK-LABEL: func.func @vmi_to_vpto_iota_i32_desc( -// CHECK: %[[P0:.*]] = pto.vci %arg0 {order = "DESC"} : i32 -> !pto.vreg<64xi32> +// CHECK-SAME: %[[BASE:.*]]: i32 +// CHECK: %[[C0:.*]] = arith.constant 0 : i32 +// CHECK: %[[Z0:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> +// CHECK: %[[DUP0:.*]] = pto.vdup %[[BASE]] +// CHECK: %[[P0:.*]] = pto.vsub %[[DUP0]], %[[Z0]] // CHECK: arith.constant 64 : i32 -// CHECK: arith.subi %arg0 -// CHECK: %[[P1:.*]] = pto.vci {{.*}} {order = "DESC"} +// CHECK: %[[BASE64:.*]] = arith.subi %[[BASE]] +// CHECK: %[[Z1:.*]] = pto.vci {{.*}} : i32 -> !pto.vreg<64xi32> +// CHECK: %[[DUP1:.*]] = pto.vdup %[[BASE64]] +// CHECK: %[[P1:.*]] = pto.vsub %[[DUP1]], %[[Z1]] // CHECK: return %[[P0]], %[[P1]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. @@ -95,7 +105,10 @@ module { // CHECK-NOT: unrealized_conversion_cast // CHECK-LABEL: func.func @vmi_to_vpto_iota_i16_asc( -// CHECK: %[[P16:.*]] = pto.vci %arg0 : i16 -> !pto.vreg<128xi16> +// CHECK-SAME: %[[BASE:.*]]: i16 +// CHECK: %[[C0:.*]] = arith.constant 0 : i16 +// CHECK: %[[Z:.*]] = pto.vci %[[C0]] : i16 -> !pto.vreg<128xi16> +// CHECK: %[[P16:.*]] = pto.vadds %[[Z]], %[[BASE]] // CHECK: return %[[P16]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto new file mode 100644 index 0000000000..1a2b34b63f --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto @@ -0,0 +1,37 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +// Contiguous VL128 {group=2} must restart both physical parts at the same +// base (group-periodic), materialized as vci(0)+vadds(base) for camodel. + +module { + func.func @vmi_to_vpto_iota_group2_i32(%base: i32) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) { + %value = pto.vmi.vci %base {group = 2 : i64} + : i32 -> !pto.vmi.vreg<128xi32, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%value) + : (!pto.vmi.vreg<128xi32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) + return %p0, %p1 : !pto.vreg<64xi32>, !pto.vreg<64xi32> + } +} + +// CHECK-LABEL: func.func @vmi_to_vpto_iota_group2_i32( +// CHECK-SAME: %[[BASE:.*]]: i32 +// CHECK: %[[C0:.*]] = arith.constant 0 : i32 +// CHECK: %[[Z0:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> +// CHECK: %[[P0:.*]] = pto.vadds %[[Z0]], %[[BASE]] +// CHECK: %[[C0B:.*]] = arith.constant 0 : i32 +// CHECK: %[[Z1:.*]] = pto.vci %[[C0B]] : i32 -> !pto.vreg<64xi32> +// CHECK: %[[P1:.*]] = pto.vadds %[[Z1]], %[[BASE]] +// CHECK: return %[[P0]], %[[P1]] +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. +// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_tail.pto b/test/lit/vmi_new/vmi_to_vpto_iota_tail.pto index 675f2a106d..7e01c65a3a 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota_tail.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota_tail.pto @@ -35,10 +35,14 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_iota_contiguous_tail( -// CHECK: %[[P0:.*]] = pto.vci %arg0 : i32 -> !pto.vreg<64xi32> +// CHECK-SAME: %[[BASE:.*]]: i32 +// CHECK: %[[C0:.*]] = arith.constant 0 : i32 +// CHECK: %[[Z0:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> +// CHECK: %[[P0:.*]] = pto.vadds %[[Z0]], %[[BASE]] // CHECK: arith.constant 64 : i32 -// CHECK: arith.addi %arg0 -// CHECK: %[[P1:.*]] = pto.vci +// CHECK: %[[BASE64:.*]] = arith.addi %[[BASE]] +// CHECK: %[[Z1:.*]] = pto.vci {{.*}} : i32 -> !pto.vreg<64xi32> +// CHECK: %[[P1:.*]] = pto.vadds %[[Z1]], %[[BASE64]] // CHECK: return %[[P0]], %[[P1]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. From da4a9ead97d346fa4139848cbe545a39753d8e56 Mon Sep 17 00:00:00 2001 From: peanutchan Date: Mon, 3 Aug 2026 23:59:02 +0800 Subject: [PATCH 3/9] fix(vmi): share one index vreg for identical group-periodic runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit group=2 is group-periodic [base..base+S | …], not a continuous 0..L-1 ramp. Physical parts with the same laneOffset reuse one iota chunk so VL128 group=2 does not re-emit duplicate vci/vadds. Co-authored-by: Cursor --- include/PTO/IR/VMIOps.td | 23 +++++++------ lib/PTO/Transforms/VMIToVPTO.cpp | 36 +++++++++++++------- test/lit/vmi_new/vmi_to_vpto_iota_group2.pto | 15 ++++---- 3 files changed, 44 insertions(+), 30 deletions(-) diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index 731f4c5fa7..abae1353d8 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -63,13 +63,15 @@ def VMIBroadcastOp : VMI_Op<"broadcast", [Pure]> { def VMIIotaOp : VMI_Op<"iota", [Pure]> { let summary = "Create a VMI logical index vector from a scalar base"; let description = [{ - Without `{group}`, produces a contiguous ramp - `dst[i] = base + i` (ASC) over the full logical vector. - - With `{group = C}`, produces a **group-periodic** ramp: each of the C - groups of size `S = L / C` independently gets - `dst[g*S + j] = base + j` for `j in [0, S)`. Example for - `L=128, C=2, base=0`: `[0..63 | 0..63]`. + Without `{group}`, produces a *continuous* ramp over the full logical + length L: `dst[i] = base + i` (ASC). For L=128 that is `base..base+127` + (physical parts `base+0..63` then `base+64..127`). + + With `{group = C}`, produces a **group-periodic** ramp (not continuous + across groups): each of the C groups of size `S = L / C` independently + gets `dst[g*S + j] = base + j`. Example `L=128, C=2, base=0`: + `[0..63 | 0..63]`. Both physical VL64 parts hold the same ramp and may + share one index register. Layout packing is orthogonal to this ramp shape. }]; let arguments = (ins AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, @@ -832,9 +834,10 @@ def VMIVbrcOp : VMI_Op<"vbrc", [Pure]> { def VMIVciOp : VMI_Op<"vci", [Pure]> { let summary = "Create a VMI logical index vector from a scalar base"; let description = [{ - Unified form of `iota`. Without `{group}`, contiguous ramp over L lanes. - With `{group = C}`, group-periodic ramp: each group restarts at `base` - (see `iota`). Lowers to `iota` preserving `{group}`. + Unified form of `iota`. Without `{group}`, continuous ramp over L lanes. + With `{group = C}`, group-periodic restart at `base` per group (see + `iota`); identical group runs share one physical index register after + lowering. Lowers to `iota` preserving `{group}`. }]; let arguments = (ins AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index a045f2133e..f3886ee0b2 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -5329,11 +5329,14 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { SmallVector results; results.reserve(resultTypes.size()); - // Optional {group = C}: group-periodic ramp. Each group of size - // S = L / C independently gets dst[g*S + j] = base + j. Contiguous - // physical part p therefore uses laneOffset = (p * lanesPerPart) % S - // (not p * lanesPerPart). Without {group}, keep the contiguous / - // deinterleaved absolute ramps below. + // Optional {group = C}: *group-periodic* ramp (not a continuous 0..L-1 + // ramp). Each of C groups of size S = L / C independently gets + // dst[g*S + j] = base + j. Example L=128,C=2 → [base..base+63 | + // base..base+63]. Layout only packs those logical lanes into + // physical parts; part p uses laneOffset = (p * lanesPerPart) % S. + // When several parts share the same offset (VL128 group=2 → both 0), + // materialize one physical index vreg and reuse it — same register for + // every identical group run. if (auto groupAttr = op.getGroupAttr()) { int64_t numGroups = groupAttr.getInt(); int64_t logicalLanes = resultVMIType.getElementCount(); @@ -5351,18 +5354,27 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { logicalLanes) return rewriter.notifyMatchFailure( op, "grouped contiguous iota physical result count mismatch"); + + // Cache by (resultType, laneOffset): identical group runs share one + // physical index register instead of re-emitting vci/vadds. + llvm::DenseMap, Value> sharedChunks; for (auto [index, resultType] : llvm::enumerate(resultTypes)) { if (!isa(resultType)) return rewriter.notifyMatchFailure(op, "iota result must be vreg"); int64_t laneOffset = (static_cast(index) * *lanesPerPart) % groupSize; - FailureOr result = createIotaContiguousChunk( - op.getLoc(), resultType, *base, laneOffset, op.getOrderAttr(), - rewriter); - if (failed(result)) - return rewriter.notifyMatchFailure( - op, "failed to materialize grouped contiguous iota chunk"); - results.push_back(*result); + auto key = std::make_pair(resultType, laneOffset); + auto it = sharedChunks.find(key); + if (it == sharedChunks.end()) { + FailureOr result = createIotaContiguousChunk( + op.getLoc(), resultType, *base, laneOffset, op.getOrderAttr(), + rewriter); + if (failed(result)) + return rewriter.notifyMatchFailure( + op, "failed to materialize grouped iota chunk"); + it = sharedChunks.try_emplace(key, *result).first; + } + results.push_back(it->second); } replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto index 1a2b34b63f..ebebedecea 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto @@ -8,8 +8,8 @@ // RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s -// Contiguous VL128 {group=2} must restart both physical parts at the same -// base (group-periodic), materialized as vci(0)+vadds(base) for camodel. +// VL128 {group=2} is group-periodic [base..base+63 | base..base+63], not a +// continuous 0..127 ramp. Both physical parts reuse one index vreg. module { func.func @vmi_to_vpto_iota_group2_i32(%base: i32) @@ -26,12 +26,11 @@ module { // CHECK-LABEL: func.func @vmi_to_vpto_iota_group2_i32( // CHECK-SAME: %[[BASE:.*]]: i32 // CHECK: %[[C0:.*]] = arith.constant 0 : i32 -// CHECK: %[[Z0:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> -// CHECK: %[[P0:.*]] = pto.vadds %[[Z0]], %[[BASE]] -// CHECK: %[[C0B:.*]] = arith.constant 0 : i32 -// CHECK: %[[Z1:.*]] = pto.vci %[[C0B]] : i32 -> !pto.vreg<64xi32> -// CHECK: %[[P1:.*]] = pto.vadds %[[Z1]], %[[BASE]] -// CHECK: return %[[P0]], %[[P1]] +// CHECK: %[[Z:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> +// CHECK: %[[IDX:.*]] = pto.vadds %[[Z]], %[[BASE]] +// CHECK-NOT: pto.vci +// CHECK-NOT: pto.vadds +// CHECK: return %[[IDX]], %[[IDX]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast From 3f39091f6ab1ba6c1c7fd6277cf6b4c111ce49f0 Mon Sep 17 00:00:00 2001 From: peanutchan Date: Tue, 4 Aug 2026 09:51:49 +0800 Subject: [PATCH 4/9] fix(vmi): restore raw vci(dyn) and share group=2 index VL Drop the vci(0)+vadds contiguous materialize path; contiguous iota is pto.vci(chunkBase) again, with group=2 reusing one SSA. Replace the topk-heavy share proof with lit expectations plus a simple vci(0)+vadds(1000)+vsts camodel example. Co-authored-by: Cursor --- lib/PTO/Transforms/VMIToVPTO.cpp | 45 +----- ptodsl/examples/vci_vadds_share_launch.py | 144 +++++++++++++++++++ ptodsl/tests/test_vmi_vci_dynamic_index.py | 7 +- test/lit/vmi_new/vmi_to_vpto_iota.pto | 27 +--- test/lit/vmi_new/vmi_to_vpto_iota_group2.pto | 8 +- test/lit/vmi_new/vmi_to_vpto_iota_tail.pto | 10 +- 6 files changed, 170 insertions(+), 71 deletions(-) create mode 100644 ptodsl/examples/vci_vadds_share_launch.py diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index f3886ee0b2..ef0466ded6 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -5207,47 +5207,15 @@ FailureOr createIotaContiguousChunk(Location loc, Type resultType, Value base, int64_t laneOffset, StringAttr orderAttr, PatternRewriter &rewriter) { - // Camodel / Bisheng llvm.hivm.vci honors immediate bases but currently - // ignores a non-zero *register* Sn base (dynamic sreg). Materialize - // absolute indices as vci(0) ± vadds/vdup(base), matching the - // deinterleaved iota path and Ascend's S.vci(0)+vadds pattern. - auto vregType = dyn_cast(resultType); - if (!vregType) - return failure(); - + // Contiguous iota is a direct VCI of the absolute chunk base (ASC + // `vci(offset_sreg)`). Group-periodic VL128 {group=2} shares one such VL64 + // result across both physical parts (see sharedChunks below). StringRef order = orderAttr ? orderAttr.getValue() : StringRef("ASC"); FailureOr chunkBase = createIotaChunkBase(loc, base, laneOffset, order, rewriter); if (failed(chunkBase)) return failure(); - - FailureOr mask = createAllTrueMaskForVReg(loc, vregType, rewriter); - FailureOr zero = - createScalarOffsetConstant(loc, base.getType(), 0, rewriter); - if (failed(mask) || failed(zero)) - return failure(); - - Value local = - rewriter.create(loc, resultType, *zero, StringAttr{}).getResult(); - - // Fast-path: base+laneOffset folds to zero and ASC → plain vci(0). - if (order != "DESC") { - if (auto constBase = chunkBase->getDefiningOp()) { - if (auto intAttr = dyn_cast(constBase.getValue())) { - if (intAttr.getValue().isZero()) - return local; - } - } - return rewriter.create(loc, resultType, local, *chunkBase, *mask) - .getResult(); - } - - Value baseVector = - rewriter - .create(loc, resultType, *chunkBase, *mask, - /*position=*/nullptr) - .getResult(); - return rewriter.create(loc, resultType, baseVector, local, *mask) + return rewriter.create(loc, resultType, *chunkBase, orderAttr) .getResult(); } @@ -5355,8 +5323,9 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { return rewriter.notifyMatchFailure( op, "grouped contiguous iota physical result count mismatch"); - // Cache by (resultType, laneOffset): identical group runs share one - // physical index register instead of re-emitting vci/vadds. + // Group-periodic VL128 {group=2}: one physical ramp shared by both + // parts → [base..base+63 | base..base+63]. Emit a single vci(dyn Sn) + // (or vci(chunkBase) when laneOffset≠0) and reuse that SSA value. llvm::DenseMap, Value> sharedChunks; for (auto [index, resultType] : llvm::enumerate(resultTypes)) { if (!isa(resultType)) diff --git a/ptodsl/examples/vci_vadds_share_launch.py b/ptodsl/examples/vci_vadds_share_launch.py new file mode 100644 index 0000000000..95afdf8832 --- /dev/null +++ b/ptodsl/examples/vci_vadds_share_launch.py @@ -0,0 +1,144 @@ +# 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. + +"""Shared ``vci`` + VL ``vadds`` + ``vsts`` probe (group=1 / group=2). + +Replaces topk-heavy validation for the dynamic contiguous iota share path: + + idx = vci(0, size=vl, group=G) # G=2 → one VL64 ramp reused as 0..63|0..63 + out = vadds(idx, 1000, mask) # VL-scale compute on that shared value + vstore(out) # → 1000..1063 (G=1) + # → 1000..1063|1000..1063 (G=2) + +Run under the CPU / camodel simulator: + + scripts/sim_dsl.sh --output /tmp/vci_vadds_g1_out \\ + ptodsl/examples/vci_vadds_share_launch.py -- --groups 1 + scripts/sim_dsl.sh --output /tmp/vci_vadds_g2_out \\ + ptodsl/examples/vci_vadds_share_launch.py -- --groups 2 +""" + +from __future__ import annotations + +import argparse +import sys +import time + +import numpy as np + +from ptodsl import pto + +_DEVICE = "npu:0" +PHYS_VL = 64 +ADD_SCALAR = 1000 +K_REMAT = 2 + + +def _expected(num_groups: int) -> np.ndarray: + ramp = np.arange(ADD_SCALAR, ADD_SCALAR + PHYS_VL, dtype=np.int32) + if num_groups == 1: + return ramp.reshape(1, PHYS_VL) + out = np.zeros((1, PHYS_VL * 2), dtype=np.int32) + out[0, :PHYS_VL] = ramp + out[0, PHYS_VL:] = ramp + return out + + +def _make_kernel(num_groups: int): + cols = PHYS_VL if num_groups == 1 else PHYS_VL * 2 + vl = PHYS_VL * num_groups + nbytes = cols * 4 + + @pto.jit( + name=f"vci_vadds_share_g{num_groups}", + kernel_kind="vector", + target="a5", + backend="vpto", + mode="explicit", + ) + def kernel(out_ptr: pto.ptr(pto.i32, "gm")): + ub = pto.alloc_tile(shape=[1, cols], dtype=pto.i32) + mask = pto.vmi.create_mask(vl, size=vl) + for _k in range(K_REMAT): + idx = pto.vmi.vci(0, size=vl, group=num_groups) + out_idx = pto.vmi.vadds(idx, pto.i32(ADD_SCALAR), mask) + pto.vmi.vstore(out_idx, ub.as_ptr(), pto.const(0, dtype=pto.index)) + pto.set_flag("V", "MTE3", event_id=0) + pto.wait_flag("V", "MTE3", event_id=0) + pto.mte_ub_gm( + ub.as_ptr(), + out_ptr, + nbytes, + nburst=(1, 0, 0), + ) + + return kernel + + +def init_torch_npu(): + import torch + import torch_npu # noqa: F401 + + torch.npu.config.allow_internal_format = False + torch_npu.npu.set_compile_mode(jit_compile=False) + torch.npu.set_device(_DEVICE) + return torch + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--groups", type=int, choices=(1, 2), default=2) + argv = [a for a in sys.argv[1:] if a != "--"] + args = ap.parse_args(argv) + + torch = init_torch_npu() + kernel = _make_kernel(args.groups) + expect = _expected(args.groups) + + t0 = time.perf_counter() + compiled = kernel.compile() + compile_s = time.perf_counter() - t0 + mlir = compiled.mlir_text() + if "pto.vmi.vci" not in mlir: + raise SystemExit("FAIL: expected pto.vmi.vci in frontend MLIR") + if "pto.vmi.vadds" not in mlir and "vadds" not in mlir: + raise SystemExit("FAIL: expected vadds in frontend MLIR") + + out = torch.zeros(expect.shape, dtype=torch.int32, device=_DEVICE) + stream = torch.npu.current_stream()._as_parameter_ # noqa: SLF001 + t0 = time.perf_counter() + compiled(out.data_ptr(), stream=stream) + torch.npu.synchronize() + launch_s = time.perf_counter() - t0 + + got = out.cpu().numpy() + if not np.array_equal(got, expect): + diff = np.argwhere(got != expect) + i, j = (int(x) for x in diff[0]) + print(f"FAIL groups={args.groups} num_mismatch={len(diff)}") + print(f" first bad lane={j} got={got[i, j]} expected={expect[i, j]}") + print(f" got={got[i].tolist()}") + print(f" exp={expect[i].tolist()}") + raise SystemExit(1) + + half = "" + if args.groups == 2: + half = ( + f" half0={got[0, :4].tolist()}" + f"|half1={got[0, PHYS_VL:PHYS_VL + 4].tolist()}" + ) + print( + f"PASS groups={args.groups} remat={K_REMAT} " + f"vci(0)+vadds({ADD_SCALAR})+vsts " + f"compile={compile_s:.2f}s launch={launch_s:.2f}s" + f"{half}" + ) + + +if __name__ == "__main__": + main() diff --git a/ptodsl/tests/test_vmi_vci_dynamic_index.py b/ptodsl/tests/test_vmi_vci_dynamic_index.py index a53e2abbfd..e4ca952101 100644 --- a/ptodsl/tests/test_vmi_vci_dynamic_index.py +++ b/ptodsl/tests/test_vmi_vci_dynamic_index.py @@ -11,7 +11,12 @@ Dynamic loop indices must coerce to an i32 sreg so ODS/verify accept the op and lowering emits ``VCI Vd, Sn`` (matches Ascend ``S.vci(T.int32(offset))``). -Also covers ``group=2`` (VL128 group-periodic) with a dynamic base. +Also covers ``group=2`` (VL128 group-periodic) with a dynamic base. Lit +``vmi_to_vpto_iota_group2.pto`` checks the share lowering: +``%idx = pto.vci %base`` then ``return %idx, %idx``. + +Camodel share+compute probe (not this unit test): + ``ptodsl/examples/vci_vadds_share_launch.py`` → ``vci(0)+vadds(1000)+vsts``. Run: python3 ptodsl/tests/test_vmi_vci_dynamic_index.py diff --git a/test/lit/vmi_new/vmi_to_vpto_iota.pto b/test/lit/vmi_new/vmi_to_vpto_iota.pto index 6b88ef8b6d..eebca50ef2 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota.pto @@ -61,30 +61,20 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_iota_i32_asc( -// CHECK-SAME: %[[BASE:.*]]: i32 -// CHECK: %[[C0:.*]] = arith.constant 0 : i32 -// CHECK: %[[Z0:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> -// CHECK: %[[P0:.*]] = pto.vadds %[[Z0]], %[[BASE]] +// CHECK: %[[P0:.*]] = pto.vci %arg0 : i32 -> !pto.vreg<64xi32> // CHECK: arith.constant 64 : i32 -// CHECK: %[[BASE64:.*]] = arith.addi %[[BASE]] -// CHECK: %[[Z1:.*]] = pto.vci {{.*}} : i32 -> !pto.vreg<64xi32> -// CHECK: %[[P1:.*]] = pto.vadds %[[Z1]], %[[BASE64]] +// CHECK: arith.addi %arg0 +// CHECK: %[[P1:.*]] = pto.vci // CHECK: return %[[P0]], %[[P1]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. // CHECK-NOT: unrealized_conversion_cast // CHECK-LABEL: func.func @vmi_to_vpto_iota_i32_desc( -// CHECK-SAME: %[[BASE:.*]]: i32 -// CHECK: %[[C0:.*]] = arith.constant 0 : i32 -// CHECK: %[[Z0:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> -// CHECK: %[[DUP0:.*]] = pto.vdup %[[BASE]] -// CHECK: %[[P0:.*]] = pto.vsub %[[DUP0]], %[[Z0]] +// CHECK: %[[P0:.*]] = pto.vci %arg0 {order = "DESC"} : i32 -> !pto.vreg<64xi32> // CHECK: arith.constant 64 : i32 -// CHECK: %[[BASE64:.*]] = arith.subi %[[BASE]] -// CHECK: %[[Z1:.*]] = pto.vci {{.*}} : i32 -> !pto.vreg<64xi32> -// CHECK: %[[DUP1:.*]] = pto.vdup %[[BASE64]] -// CHECK: %[[P1:.*]] = pto.vsub %[[DUP1]], %[[Z1]] +// CHECK: arith.subi %arg0 +// CHECK: %[[P1:.*]] = pto.vci {{.*}} {order = "DESC"} // CHECK: return %[[P0]], %[[P1]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. @@ -105,10 +95,7 @@ module { // CHECK-NOT: unrealized_conversion_cast // CHECK-LABEL: func.func @vmi_to_vpto_iota_i16_asc( -// CHECK-SAME: %[[BASE:.*]]: i16 -// CHECK: %[[C0:.*]] = arith.constant 0 : i16 -// CHECK: %[[Z:.*]] = pto.vci %[[C0]] : i16 -> !pto.vreg<128xi16> -// CHECK: %[[P16:.*]] = pto.vadds %[[Z]], %[[BASE]] +// CHECK: %[[P16:.*]] = pto.vci %arg0 : i16 -> !pto.vreg<128xi16> // CHECK: return %[[P16]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto index ebebedecea..6070e60420 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto @@ -8,8 +8,8 @@ // RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s -// VL128 {group=2} is group-periodic [base..base+63 | base..base+63], not a -// continuous 0..127 ramp. Both physical parts reuse one index vreg. +// VL128 {group=2} is group-periodic [base..base+63 | base..base+63]. +// One vci(dyn Sn); both physical parts reuse that single VL64 result. module { func.func @vmi_to_vpto_iota_group2_i32(%base: i32) @@ -25,9 +25,7 @@ module { // CHECK-LABEL: func.func @vmi_to_vpto_iota_group2_i32( // CHECK-SAME: %[[BASE:.*]]: i32 -// CHECK: %[[C0:.*]] = arith.constant 0 : i32 -// CHECK: %[[Z:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> -// CHECK: %[[IDX:.*]] = pto.vadds %[[Z]], %[[BASE]] +// CHECK: %[[IDX:.*]] = pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> // CHECK-NOT: pto.vci // CHECK-NOT: pto.vadds // CHECK: return %[[IDX]], %[[IDX]] diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_tail.pto b/test/lit/vmi_new/vmi_to_vpto_iota_tail.pto index 7e01c65a3a..675f2a106d 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota_tail.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota_tail.pto @@ -35,14 +35,10 @@ module { } // CHECK-LABEL: func.func @vmi_to_vpto_iota_contiguous_tail( -// CHECK-SAME: %[[BASE:.*]]: i32 -// CHECK: %[[C0:.*]] = arith.constant 0 : i32 -// CHECK: %[[Z0:.*]] = pto.vci %[[C0]] : i32 -> !pto.vreg<64xi32> -// CHECK: %[[P0:.*]] = pto.vadds %[[Z0]], %[[BASE]] +// CHECK: %[[P0:.*]] = pto.vci %arg0 : i32 -> !pto.vreg<64xi32> // CHECK: arith.constant 64 : i32 -// CHECK: %[[BASE64:.*]] = arith.addi %[[BASE]] -// CHECK: %[[Z1:.*]] = pto.vci {{.*}} : i32 -> !pto.vreg<64xi32> -// CHECK: %[[P1:.*]] = pto.vadds %[[Z1]], %[[BASE64]] +// CHECK: arith.addi %arg0 +// CHECK: %[[P1:.*]] = pto.vci // CHECK: return %[[P0]], %[[P1]] // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. From d8746033d9b7d7ee697968567558cd06cffd14be Mon Sep 17 00:00:00 2001 From: peanutchan Date: Tue, 4 Aug 2026 20:44:58 +0800 Subject: [PATCH 5/9] feat(vmi): sub-VL grouped iota + deint via ensure_layout Support group-periodic vci/iota when S < physVL (mask+vsel), and rewrite non-contiguous grouped results to contiguous iota + ensure_layout so layout infer handles interleave/deinterleave without a dedicated path. Co-authored-by: Cursor --- include/PTO/IR/VMIOps.td | 10 +- lib/PTO/IR/VMI.cpp | 14 ++ lib/PTO/Transforms/VMILayoutRematerialize.cpp | 11 +- .../Transforms/VMILowerUnifiedToLegacy.cpp | 74 ++++++- lib/PTO/Transforms/VMIToVPTO.cpp | 190 ++++++++++++------ ptodsl/ptodsl/_vmi_namespace.py | 37 ++++ ptodsl/tests/test_vmi_vci_dynamic_index.py | 56 ++++++ .../vmi_new/vmi_to_vpto_iota_group_deint.pto | 34 ++++ .../vmi_new/vmi_to_vpto_iota_group_subvl.pto | 79 ++++++++ 9 files changed, 437 insertions(+), 68 deletions(-) create mode 100644 test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_iota_group_subvl.pto diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index abae1353d8..829a689525 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -70,8 +70,11 @@ def VMIIotaOp : VMI_Op<"iota", [Pure]> { With `{group = C}`, produces a **group-periodic** ramp (not continuous across groups): each of the C groups of size `S = L / C` independently gets `dst[g*S + j] = base + j`. Example `L=128, C=2, base=0`: - `[0..63 | 0..63]`. Both physical VL64 parts hold the same ramp and may - share one index register. Layout packing is orthogonal to this ramp shape. + `[0..63 | 0..63]`. When `S` is a multiple of physical VL, parts that share + the same ramp may reuse one index register. When `S < VL` and `VL % S == 0` + (e.g. i32 `L=64, C=2` → `[0..31 | 0..31]` inside one VL64), lowering merges + masked `vci` runs via `vsel`. Grouped iota materializes as contiguous; + deinterleaved results go through contiguous iota + `ensure_layout`. }]; let arguments = (ins AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, @@ -837,7 +840,8 @@ def VMIVciOp : VMI_Op<"vci", [Pure]> { Unified form of `iota`. Without `{group}`, continuous ramp over L lanes. With `{group = C}`, group-periodic restart at `base` per group (see `iota`); identical group runs share one physical index register after - lowering. Lowers to `iota` preserving `{group}`. + lowering. Non-contiguous result layouts rewrite to contiguous iota + + `ensure_layout` in unified→legacy. Lowers to `iota` preserving `{group}`. }]; let arguments = (ins AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index c2fd5a5ce1..a9e7d3d656 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -962,6 +962,13 @@ LogicalResult VMIIotaOp::verify() { if (resultType.getElementCount() % numGroups != 0) return emitOpError("requires group to evenly divide result logical lane " "count"); + int64_t groupSize = resultType.getElementCount() / numGroups; + FailureOr lanesPerPart = getDataLanesPerPart(elementType); + if (succeeded(lanesPerPart) && groupSize % *lanesPerPart != 0 && + *lanesPerPart % groupSize != 0) + return emitOpError("requires group_size to divide or be a multiple of " + "physical lanes per part (") + << *lanesPerPart << ")"; } return success(); } @@ -2619,6 +2626,13 @@ LogicalResult VMIVciOp::verify() { if (resultType.getElementCount() % numGroups != 0) return emitOpError("requires group to evenly divide result logical lane " "count"); + int64_t groupSize = resultType.getElementCount() / numGroups; + FailureOr lanesPerPart = getDataLanesPerPart(elementType); + if (succeeded(lanesPerPart) && groupSize % *lanesPerPart != 0 && + *lanesPerPart % groupSize != 0) + return emitOpError("requires group_size to divide or be a multiple of " + "physical lanes per part (") + << *lanesPerPart << ")"; } return success(); } diff --git a/lib/PTO/Transforms/VMILayoutRematerialize.cpp b/lib/PTO/Transforms/VMILayoutRematerialize.cpp index 9318923314..35a3901784 100644 --- a/lib/PTO/Transforms/VMILayoutRematerialize.cpp +++ b/lib/PTO/Transforms/VMILayoutRematerialize.cpp @@ -240,11 +240,20 @@ static std::optional rematerializeDataProducer(Value value, return builder.create(loc, resultType, broadcast.getValue()) .getResult(); - if (auto iota = value.getDefiningOp()) + if (auto iota = value.getDefiningOp()) { + // Grouped iota only materializes as contiguous. Keep ensure_layout when + // the consumer wants a non-contiguous layout instead of rematerializing + // a grouped iota directly into deinterleaved/etc. + if (iota.getGroupAttr()) { + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!resultLayout || !resultLayout.isContiguous()) + return std::nullopt; + } return builder .create(loc, resultType, iota.getBase(), iota.getOrderAttr(), iota.getGroupAttr()) .getResult(); + } return std::nullopt; } diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index cc8d26d2b9..e920ed4e05 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -1189,17 +1189,48 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() { // ---- Category A: pure syntactic renames ---- if (auto vop = dyn_cast(op)) { - // vci -> iota (preserve optional {group}) + // vci -> iota (preserve optional {group}). + // Grouped iota only lowers as contiguous; non-contiguous results are + // rewritten to contiguous iota + ensure_layout so layout assignment / + // consumers can still request deinterleaved without a dedicated + // grouped-deint materialization. builder.setInsertionPoint(op); StringAttr orderAttr; if (auto order = vop.getOrder()) orderAttr = builder.getStringAttr(*order); - Value result = + + Type resultType = vop.getResult().getType(); + Value iotaResult; + if (vop.getGroupAttr()) { + if (auto vmiTy = dyn_cast(resultType)) { + VMILayoutAttr layout = vmiTy.getLayoutAttr(); + if (layout && !layout.isContiguous()) { + Type contigType = VMIVRegType::get( + op->getContext(), vmiTy.getElementCount(), + vmiTy.getElementType(), + VMILayoutAttr::getContiguous(op->getContext())); + Value contig = + builder + .create(op->getLoc(), contigType, vop.getBase(), + orderAttr, vop.getGroupAttr()) + .getResult(); + iotaResult = + builder + .create(op->getLoc(), vmiTy, contig) + .getResult(); + vop.getResult().replaceAllUsesWith(iotaResult); + op->erase(); + continue; + } + } + } + + iotaResult = builder - .create(op->getLoc(), vop.getResult().getType(), - vop.getBase(), orderAttr, vop.getGroupAttr()) + .create(op->getLoc(), resultType, vop.getBase(), + orderAttr, vop.getGroupAttr()) .getResult(); - vop.getResult().replaceAllUsesWith(result); + vop.getResult().replaceAllUsesWith(iotaResult); op->erase(); continue; } @@ -1582,6 +1613,39 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() { continue; } } + + // Direct legacy iota with {group} on a non-contiguous layout: rewrite to + // contiguous iota + ensure_layout (same contract as vci above). + SmallVector groupedNonContigIotas; + module.walk([&](VMIIotaOp iota) { + if (!iota.getGroupAttr()) + return; + auto vmiTy = dyn_cast(iota.getResult().getType()); + if (!vmiTy) + return; + VMILayoutAttr layout = vmiTy.getLayoutAttr(); + if (layout && !layout.isContiguous()) + groupedNonContigIotas.push_back(iota); + }); + for (VMIIotaOp iota : groupedNonContigIotas) { + if (!iota->getBlock()) + continue; + OpBuilder builder(iota); + auto vmiTy = cast(iota.getResult().getType()); + Type contigType = VMIVRegType::get( + iota.getContext(), vmiTy.getElementCount(), vmiTy.getElementType(), + VMILayoutAttr::getContiguous(iota.getContext())); + Value contig = + builder + .create(iota.getLoc(), contigType, iota.getBase(), + iota.getOrderAttr(), iota.getGroupAttr()) + .getResult(); + Value converted = + builder.create(iota.getLoc(), vmiTy, contig) + .getResult(); + iota.getResult().replaceAllUsesWith(converted); + iota.erase(); + } } std::unique_ptr mlir::pto::createVMILowerUnifiedToLegacyPass() { diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index ef0466ded6..32d0447eea 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -5219,6 +5219,91 @@ FailureOr createIotaContiguousChunk(Location loc, Type resultType, .getResult(); } +/// Pack group-periodic ramps inside one physical VL when S < physVL and +/// physVL % S == 0 (e.g. i32 L=64,group=2 → [base..base+31 | base..base+31]). +/// +/// Recipe (matches reduce/quant "mask + offset-by-S + VOR/Vsel"): +/// full = vci(base) // [base+0 .. base+VL) +/// for g in 0..G-1: +/// adj = full ∓ g*S // ASC: −; DESC: + +/// take lanes [g*S, (g+1)*S) from adj +/// +/// When S == physVL this is just `vci(base)` (single group fills the VL). +FailureOr createSubVLGroupPeriodicChunk(Location loc, Type resultType, + Value base, int64_t groupSize, + StringAttr orderAttr, + PatternRewriter &rewriter) { + auto vregType = dyn_cast(resultType); + if (!vregType) + return failure(); + + int64_t lanesPerPart = vregType.getElementCount(); + if (groupSize <= 0 || lanesPerPart % groupSize != 0) + return failure(); + + StringRef order = orderAttr ? orderAttr.getValue() : StringRef("ASC"); + FailureOr full = + createIotaContiguousChunk(loc, resultType, base, /*laneOffset=*/0, + orderAttr, rewriter); + FailureOr allMask = + createAllTrueMaskForVReg(loc, vregType, rewriter); + FailureOr maskType = + getMaskTypeForVReg(vregType, rewriter.getContext()); + FailureOr zeroScalar = + createScalarOffsetConstant(loc, base.getType(), 0, rewriter); + if (failed(full) || failed(allMask) || failed(maskType) || failed(zeroScalar)) + return failure(); + + int64_t groupsPerChunk = lanesPerPart / groupSize; + if (groupsPerChunk == 1) + return *full; + + Value result = rewriter + .create(loc, resultType, *zeroScalar, *allMask, + /*position=*/nullptr) + .getResult(); + for (int64_t localGroup = 0; localGroup < groupsPerChunk; ++localGroup) { + Value adjusted = *full; + if (localGroup != 0) { + int64_t delta = localGroup * groupSize; + FailureOr offsetScalar = + createScalarOffsetConstant(loc, base.getType(), delta, rewriter); + if (failed(offsetScalar)) + return failure(); + // ASC continuous is base+i; lane (g*S+j) holds base+g*S+j, want base+j + // → subtract g*S. DESC continuous is base-i; want base-j → add g*S. + if (order == "DESC") { + adjusted = rewriter + .create(loc, resultType, *full, *offsetScalar, + *allMask) + .getResult(); + } else { + Value negOffset = + isa(base.getType()) + ? rewriter + .create(loc, *offsetScalar) + .getResult() + : rewriter + .create(loc, *zeroScalar, *offsetScalar) + .getResult(); + adjusted = rewriter + .create(loc, resultType, *full, negOffset, + *allMask) + .getResult(); + } + } + FailureOr laneMask = + createLaneRangeMask(loc, *maskType, localGroup * groupSize, + (localGroup + 1) * groupSize, rewriter); + if (failed(laneMask)) + return failure(); + result = rewriter + .create(loc, resultType, adjusted, result, *laneMask) + .getResult(); + } + return result; +} + FailureOr createIotaDeinterleavedChunk(Location loc, Type resultType, Value base, int64_t factor, int64_t part, int64_t chunk, @@ -5300,11 +5385,12 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { // Optional {group = C}: *group-periodic* ramp (not a continuous 0..L-1 // ramp). Each of C groups of size S = L / C independently gets // dst[g*S + j] = base + j. Example L=128,C=2 → [base..base+63 | - // base..base+63]. Layout only packs those logical lanes into - // physical parts; part p uses laneOffset = (p * lanesPerPart) % S. - // When several parts share the same offset (VL128 group=2 → both 0), - // materialize one physical index vreg and reuse it — same register for - // every identical group run. + // base..base+63]. Only contiguous layout is lowered here; non-contiguous + // results are rewritten earlier to contiguous + ensure_layout. + // + // Contiguous materialization: + // * S % physVL == 0 → VCI chunk per distinct laneOffset (share parts). + // * physVL % S == 0 → sub-VL pack via mask + offset-by-S + vsel. if (auto groupAttr = op.getGroupAttr()) { int64_t numGroups = groupAttr.getInt(); int64_t logicalLanes = resultVMIType.getElementCount(); @@ -5312,69 +5398,55 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { return rewriter.notifyMatchFailure( op, "grouped iota requires group to divide logical lane count"); int64_t groupSize = logicalLanes / numGroups; - if (groupSize % *lanesPerPart != 0) + bool groupSizeMultipleOfPhys = groupSize % *lanesPerPart == 0; + bool physMultipleOfGroupSize = *lanesPerPart % groupSize == 0; + if (!groupSizeMultipleOfPhys && !physMultipleOfGroupSize) return rewriter.notifyMatchFailure( - op, "grouped iota requires group_size to be a multiple of " - "physical lanes per part"); + op, "grouped iota requires group_size to divide or be a multiple " + "of physical lanes per part"); - if (layout.isContiguous()) { - if (static_cast(resultTypes.size()) * *lanesPerPart != - logicalLanes) - return rewriter.notifyMatchFailure( - op, "grouped contiguous iota physical result count mismatch"); + if (!layout.isContiguous()) + return rewriter.notifyMatchFailure( + op, "grouped iota currently supports contiguous layout only; " + "ensure_layout to contiguous before lowering"); - // Group-periodic VL128 {group=2}: one physical ramp shared by both - // parts → [base..base+63 | base..base+63]. Emit a single vci(dyn Sn) - // (or vci(chunkBase) when laneOffset≠0) and reuse that SSA value. - llvm::DenseMap, Value> sharedChunks; - for (auto [index, resultType] : llvm::enumerate(resultTypes)) { - if (!isa(resultType)) - return rewriter.notifyMatchFailure(op, "iota result must be vreg"); - int64_t laneOffset = + if (static_cast(resultTypes.size()) * *lanesPerPart != + logicalLanes) + return rewriter.notifyMatchFailure( + op, "grouped contiguous iota physical result count mismatch"); + + // Under physVL % S == 0 every part holds the same in-VL pattern + // (lane j → base + j%S). Under S % physVL == 0 parts differ by + // laneOffset = (p * physVL) % S and are shared by that key. + llvm::DenseMap, Value> sharedChunks; + for (auto [index, resultType] : llvm::enumerate(resultTypes)) { + if (!isa(resultType)) + return rewriter.notifyMatchFailure(op, "iota result must be vreg"); + + int64_t laneOffset = 0; + if (groupSizeMultipleOfPhys) + laneOffset = (static_cast(index) * *lanesPerPart) % groupSize; - auto key = std::make_pair(resultType, laneOffset); - auto it = sharedChunks.find(key); - if (it == sharedChunks.end()) { - FailureOr result = createIotaContiguousChunk( - op.getLoc(), resultType, *base, laneOffset, op.getOrderAttr(), + + auto key = std::make_pair(resultType, laneOffset); + auto it = sharedChunks.find(key); + if (it == sharedChunks.end()) { + FailureOr result; + if (physMultipleOfGroupSize && groupSize < *lanesPerPart) { + result = createSubVLGroupPeriodicChunk( + op.getLoc(), resultType, *base, groupSize, op.getOrderAttr(), rewriter); - if (failed(result)) - return rewriter.notifyMatchFailure( - op, "failed to materialize grouped iota chunk"); - it = sharedChunks.try_emplace(key, *result).first; + } else { + result = createIotaContiguousChunk(op.getLoc(), resultType, *base, + laneOffset, op.getOrderAttr(), + rewriter); } - results.push_back(it->second); - } - replaceOpWithFlatConvertedValues(rewriter, op, results, - *this->getTypeConverter()); - return success(); - } - - int64_t factor = layout.getFactor(); - if (factor != numGroups) - return rewriter.notifyMatchFailure( - op, "grouped deinterleaved iota requires layout factor == group"); - if (resultTypes.size() % factor != 0) - return rewriter.notifyMatchFailure( - op, "deinterleaved iota physical result count does not match " - "layout factor"); - int64_t chunksPerPart = resultTypes.size() / factor; - if (chunksPerPart * *lanesPerPart != groupSize) - return rewriter.notifyMatchFailure( - op, "grouped deinterleaved iota chunks-per-part mismatch"); - // Group-periodic deinterleaved: each part is one group and restarts - // at `base` (partOffset ignores the group index). - for (int64_t part = 0; part < factor; ++part) { - for (int64_t chunk = 0; chunk < chunksPerPart; ++chunk) { - Type resultType = resultTypes[part * chunksPerPart + chunk]; - FailureOr result = createIotaDeinterleavedChunk( - op.getLoc(), resultType, *base, /*factor=*/1, /*part=*/0, chunk, - *lanesPerPart, op.getOrderAttr(), rewriter); if (failed(result)) return rewriter.notifyMatchFailure( - op, "failed to materialize grouped deinterleaved iota chunk"); - results.push_back(*result); + op, "failed to materialize grouped iota chunk"); + it = sharedChunks.try_emplace(key, *result).first; } + results.push_back(it->second); } replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter()); diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index 3eb7fdf3bc..8acd4ce147 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -289,6 +289,39 @@ def _derive_vci_result_type(base, size, *, context: str): return _pto.VMIVRegType.get(size, elem_type) +def _physical_lanes_per_part(elem_type, *, context: str) -> int | None: + """A5 256B physical VL lane count for VCI element types, else None.""" + if IntegerType.isinstance(elem_type): + width = IntegerType(elem_type).width + if width == 8: + return 256 + if width == 16: + return 128 + if width == 32: + return 64 + return None + # Float: f16→128, f32→64 (match getDataLanesPerPart). + name = str(elem_type) + if "f16" in name or "bf16" in name: + return 128 + if "f32" in name: + return 64 + return None + + +def _check_vci_group_tiles_phys_vl(elem_type, size, group, *, context: str) -> None: + group_size = size // group + phys = _physical_lanes_per_part(elem_type, context=context) + if phys is None: + return + if group_size % phys != 0 and phys % group_size != 0: + raise ValueError( + f"{context} requires group_size ({group_size}) to divide or be a " + f"multiple of physical lanes per part ({phys}) for element type " + f"{elem_type}" + ) + + def _derive_vmull_result_types(a, b, *, context: str): lhs_type = _as_vmi_vreg_type(_type_of(a), context=context) rhs_type = _as_vmi_vreg_type(_type_of(b), context=context) @@ -685,6 +718,10 @@ def vci(base, *, size, order=None, group=None, loc=None, ip=None): f"{context} requires size divisible by group; got size={size!r}, group={group!r}" ) result_type = _derive_vci_result_type(base, size, context=context) + if group is not None: + _check_vci_group_tiles_phys_vl( + result_type.element_type, size, group, context=context + ) base = coerce_scalar_to_type( base, _vmi_element_type(result_type, context=context), diff --git a/ptodsl/tests/test_vmi_vci_dynamic_index.py b/ptodsl/tests/test_vmi_vci_dynamic_index.py index e4ca952101..41517757ac 100644 --- a/ptodsl/tests/test_vmi_vci_dynamic_index.py +++ b/ptodsl/tests/test_vmi_vci_dynamic_index.py @@ -61,6 +61,22 @@ def vmi_vci_dynamic_group2_probe(): pto.vmi.vstore(idx, dst.as_ptr(), pto.const(0, dtype=pto.index)) +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_vci_subvl_i32_g2_probe(): + """Sub-VL: i32 size=64 group=2 → [0..31|0..31] in one physical VL.""" + dst = pto.alloc_tile(shape=[1, 64], dtype=pto.i32) + idx = pto.vmi.vci(pto.i32(0), size=64, group=2) + pto.vmi.vstore(idx, dst.as_ptr(), pto.const(0, dtype=pto.index)) + + +@pto.jit(target="a5", backend="vpto", mode="explicit") +def vmi_vci_subvl_i16_g2_probe(): + """Sub-VL: i16 size=128 group=2 → [0..63|0..63] in one physical VL.""" + dst = pto.alloc_tile(shape=[1, 128], dtype=pto.i16) + idx = pto.vmi.vci(pto.i16(0), size=128, group=2) + pto.vmi.vstore(idx, dst.as_ptr(), pto.const(0, dtype=pto.index)) + + def main() -> None: const_text = vmi_vci_const_i32_probe.compile().mlir_text() expect("pto.vmi.vci" in const_text, "const probe must emit pto.vmi.vci") @@ -98,6 +114,46 @@ def main() -> None: "arith.index_cast" in g2_text or "index_cast" in g2_text, f"group=2 probe must index_cast before vci:\n{g2_text[:2000]}", ) + + subvl32 = vmi_vci_subvl_i32_g2_probe.compile().mlir_text() + expect( + "group = 2" in subvl32 or "{group = 2" in subvl32, + f"sub-VL i32 g2 must preserve group:\n{subvl32[:1500]}", + ) + expect( + "!pto.vmi.vreg<64xi32" in subvl32, + f"sub-VL i32 g2 must be 64xi32:\n{subvl32[:1500]}", + ) + + subvl16 = vmi_vci_subvl_i16_g2_probe.compile().mlir_text() + expect( + "group = 2" in subvl16 or "{group = 2" in subvl16, + f"sub-VL i16 g2 must preserve group:\n{subvl16[:1500]}", + ) + expect( + "!pto.vmi.vreg<128xi16" in subvl16, + f"sub-VL i16 g2 must be 128xi16:\n{subvl16[:1500]}", + ) + + # Untileable: i32 size=48 group=2 → S=24 does not tile phys 64. + from ptodsl._vmi_namespace import _check_vci_group_tiles_phys_vl + from mlir.ir import Context, IntegerType + + with Context(): + i32 = IntegerType.get_signless(32) + try: + _check_vci_group_tiles_phys_vl( + i32, 48, 2, context="pto.vmi.vci(...)" + ) + raise AssertionError( + "expected ValueError for untileable group_size=24" + ) + except ValueError as err: + expect( + "physical lanes" in str(err), + f"untileable group must mention physical lanes, got: {err}", + ) + print("ptodsl_vmi_vci_dynamic_index: PASS") diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto new file mode 100644 index 0000000000..48a328a570 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto @@ -0,0 +1,34 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +// Grouped vci on a deinterleaved result is rewritten to contiguous group +// iota + ensure_layout, then lowered. No dedicated grouped-deint iota path. + +module { + func.func @vmi_to_vpto_iota_group2_deint(%base: i32) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) { + %value = pto.vmi.vci %base {group = 2 : i64} + : i32 -> !pto.vmi.vreg<128xi32, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%value) + : (!pto.vmi.vreg<128xi32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) + return %p0, %p1 : !pto.vreg<64xi32>, !pto.vreg<64xi32> + } +} + +// Contiguous group share first, then c→d ensure_layout (vdintlv). +// CHECK-LABEL: func.func @vmi_to_vpto_iota_group2_deint( +// CHECK-SAME: %[[BASE:.*]]: i32 +// CHECK: %[[IDX:.*]] = pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> +// CHECK-NOT: pto.vci +// CHECK: %[[E:.*]], %[[O:.*]] = pto.vdintlv %[[IDX]], %[[IDX]] +// CHECK: return %[[E]], %[[O]] +// CHECK-NOT: pto.vmi. +// CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group_subvl.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group_subvl.pto new file mode 100644 index 0000000000..b1092811d2 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group_subvl.pto @@ -0,0 +1,79 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +// Sub-VL group-periodic contiguous iota: S < physVL and physVL % S == 0. +// Materialize via vci(base) + per-group (∓ g*S) + lane-range vsel. +// +// i32 size=64 group=2 → [base..base+31 | base..base+31] in one VL64. +// i16 size=128 group=2 → [base..base+63 | base..base+63] in one VL128. +// i32 size=128 group=4 → two identical VL64 packs [b..b+31|b..b+31], shared. + +module { + func.func @vmi_to_vpto_iota_subvl_i32_g2(%base: i32) -> !pto.vreg<64xi32> { + %value = pto.vmi.vci %base {group = 2 : i64} + : i32 -> !pto.vmi.vreg<64xi32, #pto.vmi.layout> + %part = "pto.vmi.unpack"(%value) + : (!pto.vmi.vreg<64xi32, #pto.vmi.layout>) + -> !pto.vreg<64xi32> + return %part : !pto.vreg<64xi32> + } + + func.func @vmi_to_vpto_iota_subvl_i16_g2(%base: i16) -> !pto.vreg<128xi16> { + %value = pto.vmi.vci %base {group = 2 : i64} + : i16 -> !pto.vmi.vreg<128xi16, #pto.vmi.layout> + %part = "pto.vmi.unpack"(%value) + : (!pto.vmi.vreg<128xi16, #pto.vmi.layout>) + -> !pto.vreg<128xi16> + return %part : !pto.vreg<128xi16> + } + + func.func @vmi_to_vpto_iota_subvl_i32_g4_share(%base: i32) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) { + %value = pto.vmi.vci %base {group = 4 : i64} + : i32 -> !pto.vmi.vreg<128xi32, #pto.vmi.layout> + %p0, %p1 = "pto.vmi.unpack"(%value) + : (!pto.vmi.vreg<128xi32, #pto.vmi.layout>) + -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) + return %p0, %p1 : !pto.vreg<64xi32>, !pto.vreg<64xi32> + } +} + +// CHECK-LABEL: func.func @vmi_to_vpto_iota_subvl_i32_g2( +// CHECK-SAME: %[[BASE:.*]]: i32 +// CHECK: %[[FULL:.*]] = pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> +// CHECK: %[[ZERO:.*]] = arith.constant 0 : i32 +// CHECK: %[[INIT:.*]] = pto.vdup %[[ZERO]] +// CHECK: %[[MASK0:.*]] = pto.pset_b32 "PAT_VL32" +// CHECK: %[[G0:.*]] = pto.vsel %[[FULL]], %[[INIT]], %[[MASK0]] +// CHECK: %[[OFF:.*]] = arith.constant 32 : i32 +// CHECK: %[[NEGOFF:.*]] = arith.subi %[[ZERO]], %[[OFF]] +// CHECK: %[[ADJ:.*]] = pto.vadds %[[FULL]], %[[NEGOFF]] +// CHECK: pto.pnot +// CHECK: %[[OUT:.*]] = pto.vsel %[[ADJ]], %[[G0]] +// CHECK: return %[[OUT]] +// CHECK-NOT: pto.vmi. + +// CHECK-LABEL: func.func @vmi_to_vpto_iota_subvl_i16_g2( +// CHECK-SAME: %[[BASE:.*]]: i16 +// CHECK: %[[FULL:.*]] = pto.vci %[[BASE]] : i16 -> !pto.vreg<128xi16> +// CHECK: pto.pset_b16 "PAT_VL64" +// CHECK: pto.vsel +// CHECK: pto.vadds +// CHECK: pto.vsel +// CHECK-NOT: pto.vmi. + +// CHECK-LABEL: func.func @vmi_to_vpto_iota_subvl_i32_g4_share( +// CHECK-SAME: %[[BASE:.*]]: i32 +// CHECK: pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> +// CHECK: pto.vsel +// CHECK: pto.vadds +// CHECK: %[[OUT:.*]] = pto.vsel +// CHECK-NEXT: return %[[OUT]], %[[OUT]] +// CHECK-NOT: pto.vmi. From a5f56da58c1f0c95f88142a25d48b15f5f01f751 Mon Sep 17 00:00:00 2001 From: peanutchan Date: Tue, 4 Aug 2026 21:17:50 +0800 Subject: [PATCH 6/9] test(vmi): group iota lit via vadds/vsts/vcvt, not unpack Replace unpack-only checks with real consumers: shared group2 + vadds/vsts, more sub-VL group sizes (i32 g2/g4/g8, i16 g2/g4), and deint + sitofp vcvt. Co-authored-by: Cursor --- test/lit/vmi_new/vmi_to_vpto_iota_group2.pto | 36 +++-- .../vmi_new/vmi_to_vpto_iota_group_deint.pto | 30 ++-- .../vmi_new/vmi_to_vpto_iota_group_subvl.pto | 150 ++++++++++++------ 3 files changed, 144 insertions(+), 72 deletions(-) diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto index 6070e60420..234f166dc8 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group2.pto @@ -6,29 +6,35 @@ // 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s -// VL128 {group=2} is group-periodic [base..base+63 | base..base+63]. -// One vci(dyn Sn); both physical parts reuse that single VL64 result. +// VL128 {group=2}: group-periodic [base..base+63 | base..base+63]. +// Real consumers: shared phys vci → per-part vadds(+1000) → vsts. module { - func.func @vmi_to_vpto_iota_group2_i32(%base: i32) - -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) { - %value = pto.vmi.vci %base {group = 2 : i64} - : i32 -> !pto.vmi.vreg<128xi32, #pto.vmi.layout> - %p0, %p1 = "pto.vmi.unpack"(%value) - : (!pto.vmi.vreg<128xi32, #pto.vmi.layout>) - -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) - return %p0, %p1 : !pto.vreg<64xi32>, !pto.vreg<64xi32> + func.func @vmi_to_vpto_iota_group2_vadds_vsts( + %base: i32, %dst: !pto.ptr, %mask: !pto.vmi.mask<128xpred>) { + %c0 = arith.constant 0 : index + %idx = pto.vmi.vci %base {group = 2 : i64} + : i32 -> !pto.vmi.vreg<128xi32> + %c1000 = arith.constant 1000 : i32 + %out = pto.vmi.vadds %idx, %c1000, %mask + : !pto.vmi.vreg<128xi32>, i32, !pto.vmi.mask<128xpred> + -> !pto.vmi.vreg<128xi32> + pto.vmi.vstore %out, %dst[%c0] + : !pto.vmi.vreg<128xi32>, !pto.ptr + return } } -// CHECK-LABEL: func.func @vmi_to_vpto_iota_group2_i32( +// CHECK-LABEL: func.func @vmi_to_vpto_iota_group2_vadds_vsts( // CHECK-SAME: %[[BASE:.*]]: i32 // CHECK: %[[IDX:.*]] = pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> // CHECK-NOT: pto.vci -// CHECK-NOT: pto.vadds -// CHECK: return %[[IDX]], %[[IDX]] +// CHECK: %[[C1000:.*]] = arith.constant 1000 : i32 +// CHECK: pto.vadds %[[IDX]], %[[C1000]] +// CHECK: pto.vadds %[[IDX]], %[[C1000]] +// CHECK: pto.vsts +// CHECK: pto.vsts // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. -// CHECK-NOT: unrealized_conversion_cast diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto index 48a328a570..ba7cd2b3dd 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto @@ -8,27 +8,33 @@ // RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s -// Grouped vci on a deinterleaved result is rewritten to contiguous group -// iota + ensure_layout, then lowered. No dedicated grouped-deint iota path. +// Grouped vci on a deinterleaved result → contiguous group iota + +// ensure_layout (vdintlv), then sitofp vcvt + interleaved store. +// No unpack; vcvt/vstore are the real consumers. module { - func.func @vmi_to_vpto_iota_group2_deint(%base: i32) - -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) { - %value = pto.vmi.vci %base {group = 2 : i64} + func.func @vmi_to_vpto_iota_group2_deint_vcvt( + %base: i32, %dst: !pto.ptr) { + %c0 = arith.constant 0 : index + %idx = pto.vmi.vci %base {group = 2 : i64} : i32 -> !pto.vmi.vreg<128xi32, #pto.vmi.layout> - %p0, %p1 = "pto.vmi.unpack"(%value) - : (!pto.vmi.vreg<128xi32, #pto.vmi.layout>) - -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) - return %p0, %p1 : !pto.vreg<64xi32>, !pto.vreg<64xi32> + %fp = pto.vmi.vcvt %idx + : !pto.vmi.vreg<128xi32, #pto.vmi.layout> + -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> + pto.vmi.vstore %fp, %dst[%c0] + : !pto.vmi.vreg<128xf32, #pto.vmi.layout>, + !pto.ptr + return } } -// Contiguous group share first, then c→d ensure_layout (vdintlv). -// CHECK-LABEL: func.func @vmi_to_vpto_iota_group2_deint( +// CHECK-LABEL: func.func @vmi_to_vpto_iota_group2_deint_vcvt( // CHECK-SAME: %[[BASE:.*]]: i32 // CHECK: %[[IDX:.*]] = pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> // CHECK-NOT: pto.vci // CHECK: %[[E:.*]], %[[O:.*]] = pto.vdintlv %[[IDX]], %[[IDX]] -// CHECK: return %[[E]], %[[O]] +// CHECK: %[[FE:.*]] = pto.vcvt %[[E]] +// CHECK: %[[FO:.*]] = pto.vcvt %[[O]] +// CHECK: pto.vstsx2 %[[FE]], %[[FO]], {{.*}}, "INTLV_B32" // CHECK-NOT: pto.vmi. // CHECK-NOT: !pto.vmi. diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group_subvl.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group_subvl.pto index b1092811d2..d8827e5cfb 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota_group_subvl.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group_subvl.pto @@ -6,74 +6,134 @@ // 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-to-vpto | FileCheck %s // Sub-VL group-periodic contiguous iota: S < physVL and physVL % S == 0. -// Materialize via vci(base) + per-group (∓ g*S) + lane-range vsel. +// Materialize via vci(base) + per-group (∓ g*S) + lane-range vsel, then +// consume with vadds(+1000) + vsts (no unpack). // -// i32 size=64 group=2 → [base..base+31 | base..base+31] in one VL64. -// i16 size=128 group=2 → [base..base+63 | base..base+63] in one VL128. -// i32 size=128 group=4 → two identical VL64 packs [b..b+31|b..b+31], shared. +// i32 physVL=64: g2→S=32, g4→S=16, g8→S=8 +// i16 physVL=128: g2→S=64, g4→S=32 module { - func.func @vmi_to_vpto_iota_subvl_i32_g2(%base: i32) -> !pto.vreg<64xi32> { - %value = pto.vmi.vci %base {group = 2 : i64} - : i32 -> !pto.vmi.vreg<64xi32, #pto.vmi.layout> - %part = "pto.vmi.unpack"(%value) - : (!pto.vmi.vreg<64xi32, #pto.vmi.layout>) - -> !pto.vreg<64xi32> - return %part : !pto.vreg<64xi32> + func.func @vmi_to_vpto_iota_subvl_i32_g2( + %base: i32, %dst: !pto.ptr, %mask: !pto.vmi.mask<64xpred>) { + %c0 = arith.constant 0 : index + %idx = pto.vmi.vci %base {group = 2 : i64} + : i32 -> !pto.vmi.vreg<64xi32> + %c1000 = arith.constant 1000 : i32 + %out = pto.vmi.vadds %idx, %c1000, %mask + : !pto.vmi.vreg<64xi32>, i32, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xi32> + pto.vmi.vstore %out, %dst[%c0] + : !pto.vmi.vreg<64xi32>, !pto.ptr + return } - func.func @vmi_to_vpto_iota_subvl_i16_g2(%base: i16) -> !pto.vreg<128xi16> { - %value = pto.vmi.vci %base {group = 2 : i64} - : i16 -> !pto.vmi.vreg<128xi16, #pto.vmi.layout> - %part = "pto.vmi.unpack"(%value) - : (!pto.vmi.vreg<128xi16, #pto.vmi.layout>) - -> !pto.vreg<128xi16> - return %part : !pto.vreg<128xi16> + func.func @vmi_to_vpto_iota_subvl_i32_g4( + %base: i32, %dst: !pto.ptr, %mask: !pto.vmi.mask<64xpred>) { + %c0 = arith.constant 0 : index + %idx = pto.vmi.vci %base {group = 4 : i64} + : i32 -> !pto.vmi.vreg<64xi32> + %c1000 = arith.constant 1000 : i32 + %out = pto.vmi.vadds %idx, %c1000, %mask + : !pto.vmi.vreg<64xi32>, i32, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xi32> + pto.vmi.vstore %out, %dst[%c0] + : !pto.vmi.vreg<64xi32>, !pto.ptr + return } - func.func @vmi_to_vpto_iota_subvl_i32_g4_share(%base: i32) - -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) { - %value = pto.vmi.vci %base {group = 4 : i64} - : i32 -> !pto.vmi.vreg<128xi32, #pto.vmi.layout> - %p0, %p1 = "pto.vmi.unpack"(%value) - : (!pto.vmi.vreg<128xi32, #pto.vmi.layout>) - -> (!pto.vreg<64xi32>, !pto.vreg<64xi32>) - return %p0, %p1 : !pto.vreg<64xi32>, !pto.vreg<64xi32> + func.func @vmi_to_vpto_iota_subvl_i32_g8( + %base: i32, %dst: !pto.ptr, %mask: !pto.vmi.mask<64xpred>) { + %c0 = arith.constant 0 : index + %idx = pto.vmi.vci %base {group = 8 : i64} + : i32 -> !pto.vmi.vreg<64xi32> + %c1000 = arith.constant 1000 : i32 + %out = pto.vmi.vadds %idx, %c1000, %mask + : !pto.vmi.vreg<64xi32>, i32, !pto.vmi.mask<64xpred> + -> !pto.vmi.vreg<64xi32> + pto.vmi.vstore %out, %dst[%c0] + : !pto.vmi.vreg<64xi32>, !pto.ptr + return + } + + func.func @vmi_to_vpto_iota_subvl_i16_g2( + %base: i16, %dst: !pto.ptr, %mask: !pto.vmi.mask<128xpred>) { + %c0 = arith.constant 0 : index + %idx = pto.vmi.vci %base {group = 2 : i64} + : i16 -> !pto.vmi.vreg<128xi16> + %c1000 = arith.constant 1000 : i16 + %out = pto.vmi.vadds %idx, %c1000, %mask + : !pto.vmi.vreg<128xi16>, i16, !pto.vmi.mask<128xpred> + -> !pto.vmi.vreg<128xi16> + pto.vmi.vstore %out, %dst[%c0] + : !pto.vmi.vreg<128xi16>, !pto.ptr + return + } + + func.func @vmi_to_vpto_iota_subvl_i16_g4( + %base: i16, %dst: !pto.ptr, %mask: !pto.vmi.mask<128xpred>) { + %c0 = arith.constant 0 : index + %idx = pto.vmi.vci %base {group = 4 : i64} + : i16 -> !pto.vmi.vreg<128xi16> + %c1000 = arith.constant 1000 : i16 + %out = pto.vmi.vadds %idx, %c1000, %mask + : !pto.vmi.vreg<128xi16>, i16, !pto.vmi.mask<128xpred> + -> !pto.vmi.vreg<128xi16> + pto.vmi.vstore %out, %dst[%c0] + : !pto.vmi.vreg<128xi16>, !pto.ptr + return } } // CHECK-LABEL: func.func @vmi_to_vpto_iota_subvl_i32_g2( // CHECK-SAME: %[[BASE:.*]]: i32 // CHECK: %[[FULL:.*]] = pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> -// CHECK: %[[ZERO:.*]] = arith.constant 0 : i32 -// CHECK: %[[INIT:.*]] = pto.vdup %[[ZERO]] -// CHECK: %[[MASK0:.*]] = pto.pset_b32 "PAT_VL32" -// CHECK: %[[G0:.*]] = pto.vsel %[[FULL]], %[[INIT]], %[[MASK0]] -// CHECK: %[[OFF:.*]] = arith.constant 32 : i32 -// CHECK: %[[NEGOFF:.*]] = arith.subi %[[ZERO]], %[[OFF]] -// CHECK: %[[ADJ:.*]] = pto.vadds %[[FULL]], %[[NEGOFF]] -// CHECK: pto.pnot -// CHECK: %[[OUT:.*]] = pto.vsel %[[ADJ]], %[[G0]] -// CHECK: return %[[OUT]] +// CHECK: pto.pset_b32 "PAT_VL32" +// CHECK: pto.vsel +// CHECK: pto.vadds %[[FULL]] +// CHECK: %[[PACKED:.*]] = pto.vsel +// CHECK: %[[C1000:.*]] = arith.constant 1000 : i32 +// CHECK: pto.vadds %[[PACKED]], %[[C1000]] +// CHECK: pto.vsts +// CHECK-NOT: pto.vmi. + +// CHECK-LABEL: func.func @vmi_to_vpto_iota_subvl_i32_g4( +// CHECK-SAME: %[[BASE:.*]]: i32 +// CHECK: pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> +// CHECK-COUNT-3: pto.vsel +// CHECK: %[[C1000:.*]] = arith.constant 1000 : i32 +// CHECK: pto.vadds {{.*}}, %[[C1000]] +// CHECK: pto.vsts +// CHECK-NOT: pto.vmi. + +// CHECK-LABEL: func.func @vmi_to_vpto_iota_subvl_i32_g8( +// CHECK-SAME: %[[BASE:.*]]: i32 +// CHECK: pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> +// CHECK-COUNT-7: pto.vsel +// CHECK: %[[C1000:.*]] = arith.constant 1000 : i32 +// CHECK: pto.vadds {{.*}}, %[[C1000]] +// CHECK: pto.vsts // CHECK-NOT: pto.vmi. // CHECK-LABEL: func.func @vmi_to_vpto_iota_subvl_i16_g2( // CHECK-SAME: %[[BASE:.*]]: i16 -// CHECK: %[[FULL:.*]] = pto.vci %[[BASE]] : i16 -> !pto.vreg<128xi16> +// CHECK: pto.vci %[[BASE]] : i16 -> !pto.vreg<128xi16> // CHECK: pto.pset_b16 "PAT_VL64" // CHECK: pto.vsel // CHECK: pto.vadds // CHECK: pto.vsel +// CHECK: %[[C1000:.*]] = arith.constant 1000 : i16 +// CHECK: pto.vadds {{.*}}, %[[C1000]] +// CHECK: pto.vsts // CHECK-NOT: pto.vmi. -// CHECK-LABEL: func.func @vmi_to_vpto_iota_subvl_i32_g4_share( -// CHECK-SAME: %[[BASE:.*]]: i32 -// CHECK: pto.vci %[[BASE]] : i32 -> !pto.vreg<64xi32> -// CHECK: pto.vsel -// CHECK: pto.vadds -// CHECK: %[[OUT:.*]] = pto.vsel -// CHECK-NEXT: return %[[OUT]], %[[OUT]] +// CHECK-LABEL: func.func @vmi_to_vpto_iota_subvl_i16_g4( +// CHECK-SAME: %[[BASE:.*]]: i16 +// CHECK: pto.vci %[[BASE]] : i16 -> !pto.vreg<128xi16> +// CHECK-COUNT-3: pto.vsel +// CHECK: %[[C1000:.*]] = arith.constant 1000 : i16 +// CHECK: pto.vadds {{.*}}, %[[C1000]] +// CHECK: pto.vsts // CHECK-NOT: pto.vmi. From 163d1a2c2e47a598eef0044a3711381a6854c668 Mon Sep 17 00:00:00 2001 From: peanutchan Date: Tue, 4 Aug 2026 21:20:58 +0800 Subject: [PATCH 7/9] fix(ptodsl): use ptoas.mlir.ir in vci dynamic-index test Avoid mixing top-level mlir.ir bindings with ptoas MLIR types in the untileable group reject check. Co-authored-by: Cursor --- ptodsl/tests/test_vmi_vci_dynamic_index.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ptodsl/tests/test_vmi_vci_dynamic_index.py b/ptodsl/tests/test_vmi_vci_dynamic_index.py index 41517757ac..03affe4ecd 100644 --- a/ptodsl/tests/test_vmi_vci_dynamic_index.py +++ b/ptodsl/tests/test_vmi_vci_dynamic_index.py @@ -137,7 +137,7 @@ def main() -> None: # Untileable: i32 size=48 group=2 → S=24 does not tile phys 64. from ptodsl._vmi_namespace import _check_vci_group_tiles_phys_vl - from mlir.ir import Context, IntegerType + from ptoas.mlir.ir import Context, IntegerType with Context(): i32 = IntegerType.get_signless(32) From 760d9855714f2596a8cfa48ba55700ca8b2a68aa Mon Sep 17 00:00:00 2001 From: peanutchan Date: Wed, 5 Aug 2026 09:45:10 +0800 Subject: [PATCH 8/9] test(ptodsl): add native grouped VCI camodel probes Exercise full-VL sharing and sub-VL group restart semantics through vadds, vsts, and GM output so native regressions catch lowering gaps. Co-authored-by: Cursor --- ptodsl/examples/vci_subvl_group_launch.py | 143 ++++++++++++++++++++++ ptodsl/examples/vci_vadds_share_launch.py | 10 +- 2 files changed, 147 insertions(+), 6 deletions(-) create mode 100644 ptodsl/examples/vci_subvl_group_launch.py diff --git a/ptodsl/examples/vci_subvl_group_launch.py b/ptodsl/examples/vci_subvl_group_launch.py new file mode 100644 index 0000000000..06e91e5685 --- /dev/null +++ b/ptodsl/examples/vci_subvl_group_launch.py @@ -0,0 +1,143 @@ +# 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. + +"""Sub-VL group-periodic ``vci`` + ``vadds`` + ``vsts`` camodel probe. + +Cases (S < physVL, physVL % S == 0): + + i32 size=64 group=2 → [0..31|0..31] then +1000 + i32 size=64 group=4 → [0..15]x4 then +1000 + i16 size=128 group=2 → [0..63|0..63] then +1000 + +Run: + + scripts/sim_dsl.sh --output /tmp/vci_subvl_i32_g2_out \\ + ptodsl/examples/vci_subvl_group_launch.py -- --case i32_g2 + scripts/sim_dsl.sh --output /tmp/vci_subvl_i32_g4_out \\ + ptodsl/examples/vci_subvl_group_launch.py -- --case i32_g4 + scripts/sim_dsl.sh --output /tmp/vci_subvl_i16_g2_out \\ + ptodsl/examples/vci_subvl_group_launch.py -- --case i16_g2 +""" + +import argparse +import sys +import time + +import numpy as np + +from ptodsl import pto + +_DEVICE = "npu:0" +ADD_SCALAR = 1000 + +# name -> (dtype, size, group, np_dtype, elem_bytes) +_CASES = { + "i32_g2": (pto.i32, 64, 2, np.int32, 4), + "i32_g4": (pto.i32, 64, 4, np.int32, 4), + "i16_g2": (pto.i16, 128, 2, np.int16, 2), +} + + +def _expected(size: int, group: int, np_dtype) -> np.ndarray: + s = size // group + ramp = np.arange(ADD_SCALAR, ADD_SCALAR + s, dtype=np_dtype) + out = np.tile(ramp, group).reshape(1, size) + return out + + +def _make_kernel(case: str): + dtype, size, group, _np_dtype, elem_bytes = _CASES[case] + nbytes = size * elem_bytes + + @pto.jit( + name=f"vci_subvl_{case}", + kernel_kind="vector", + target="a5", + backend="vpto", + mode="explicit", + ) + def kernel(out_ptr: pto.ptr(dtype, "gm")): + ub = pto.castptr(pto.i64(0), pto.ptr(dtype, "ub")) + mask = pto.vmi.create_mask(size, size=size) + idx = pto.vmi.vci(dtype(0), size=size, group=group) + out_idx = pto.vmi.vadds(idx, dtype(ADD_SCALAR), mask) + pto.vmi.vstore(out_idx, ub, pto.const(0, dtype=pto.index)) + pto.set_flag("V", "MTE3", event_id=0) + pto.wait_flag("V", "MTE3", event_id=0) + pto.mte_ub_gm( + ub, + out_ptr, + nbytes, + nburst=(1, 0, 0), + ) + + return kernel + + +def init_torch_npu(): + import torch + import torch_npu # noqa: F401 + + torch.npu.config.allow_internal_format = False + torch_npu.npu.set_compile_mode(jit_compile=False) + torch.npu.set_device(_DEVICE) + return torch + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--case", choices=sorted(_CASES), default="i32_g2") + argv = [a for a in sys.argv[1:] if a != "--"] + args = ap.parse_args(argv) + + dtype, size, group, np_dtype, _ = _CASES[args.case] + torch = init_torch_npu() + kernel = _make_kernel(args.case) + expect = _expected(size, group, np_dtype) + + t0 = time.perf_counter() + compiled = kernel.compile() + compile_s = time.perf_counter() - t0 + mlir = compiled.mlir_text() + if "pto.vmi.vci" not in mlir: + raise SystemExit("FAIL: expected pto.vmi.vci in frontend MLIR") + if f"group = {group}" not in mlir and f"{{group = {group}" not in mlir: + raise SystemExit(f"FAIL: expected group={group} in frontend MLIR:\n{mlir[:1500]}") + + torch_dtype = torch.int32 if np_dtype == np.int32 else torch.int16 + out = torch.zeros(expect.shape, dtype=torch_dtype, device=_DEVICE) + stream = torch.npu.current_stream()._as_parameter_ # noqa: SLF001 + t0 = time.perf_counter() + compiled(out.data_ptr(), stream=stream) + torch.npu.synchronize() + launch_s = time.perf_counter() - t0 + + got = out.cpu().numpy() + if not np.array_equal(got, expect): + diff = np.argwhere(got != expect) + i, j = (int(x) for x in diff[0]) + s = size // group + print(f"FAIL case={args.case} num_mismatch={len(diff)}") + print(f" first bad lane={j} got={got[i, j]} expected={expect[i, j]}") + print(f" got groups={[got[0, g * s:(g + 1) * s].tolist() for g in range(group)]}") + print(f" exp groups={[expect[0, g * s:(g + 1) * s].tolist() for g in range(group)]}") + raise SystemExit(1) + + s = size // group + preview = " | ".join( + f"g{g}={got[0, g * s:g * s + min(4, s)].tolist()}" for g in range(group) + ) + print( + f"PASS case={args.case} size={size} group={group} " + f"vci(0)+vadds({ADD_SCALAR})+vsts " + f"compile={compile_s:.2f}s launch={launch_s:.2f}s {preview}" + ) + + +if __name__ == "__main__": + main() diff --git a/ptodsl/examples/vci_vadds_share_launch.py b/ptodsl/examples/vci_vadds_share_launch.py index 95afdf8832..3263d592bb 100644 --- a/ptodsl/examples/vci_vadds_share_launch.py +++ b/ptodsl/examples/vci_vadds_share_launch.py @@ -23,8 +23,6 @@ ptodsl/examples/vci_vadds_share_launch.py -- --groups 2 """ -from __future__ import annotations - import argparse import sys import time @@ -62,16 +60,16 @@ def _make_kernel(num_groups: int): mode="explicit", ) def kernel(out_ptr: pto.ptr(pto.i32, "gm")): - ub = pto.alloc_tile(shape=[1, cols], dtype=pto.i32) + ub = pto.castptr(pto.i64(0), pto.ptr(pto.i32, "ub")) mask = pto.vmi.create_mask(vl, size=vl) for _k in range(K_REMAT): - idx = pto.vmi.vci(0, size=vl, group=num_groups) + idx = pto.vmi.vci(pto.i32(0), size=vl, group=num_groups) out_idx = pto.vmi.vadds(idx, pto.i32(ADD_SCALAR), mask) - pto.vmi.vstore(out_idx, ub.as_ptr(), pto.const(0, dtype=pto.index)) + pto.vmi.vstore(out_idx, ub, pto.const(0, dtype=pto.index)) pto.set_flag("V", "MTE3", event_id=0) pto.wait_flag("V", "MTE3", event_id=0) pto.mte_ub_gm( - ub.as_ptr(), + ub, out_ptr, nbytes, nburst=(1, 0, 0), From dc668b96801c6dd3b775653fe5e4aec576d9a53e Mon Sep 17 00:00:00 2001 From: peanutchan Date: Wed, 5 Aug 2026 11:36:49 +0800 Subject: [PATCH 9/9] fix(vmi): preserve grouped iota layout contracts Lower grouped VCI through a contiguous-only internal producer so layout assignment materializes deinterleaved consumers, while normalizing group=1 to ordinary tail-capable iota. Co-authored-by: Cursor --- docs/isa/vmi-isa/02-index-gen.md | 16 ++++- include/PTO/IR/VMIOps.td | 30 +++++---- lib/PTO/IR/VMI.cpp | 61 +++++++++++------ lib/PTO/Transforms/VMILayoutAssignment.cpp | 6 ++ lib/PTO/Transforms/VMILayoutPropagation.cpp | 17 +++++ lib/PTO/Transforms/VMILayoutRematerialize.cpp | 22 ++++--- .../Transforms/VMILowerUnifiedToLegacy.cpp | 65 ++++++------------ lib/PTO/Transforms/VMIToVPTO.cpp | 16 +++-- .../14-vmi-virtual-instruction-set.md | 16 ++++- ptodsl/ptodsl/_vmi_namespace.py | 4 ++ ptodsl/tests/test_vmi_vci_dynamic_index.py | 25 ++++++- .../vmi_new/vmi_to_vpto_iota_group1_tail.pto | 66 +++++++++++++++++++ .../vmi_new/vmi_to_vpto_iota_group_deint.pto | 9 ++- .../vmi_to_vpto_iota_group_deint_assign.pto | 57 ++++++++++++++++ tools/ptoas/ptoas.cpp | 5 +- 15 files changed, 311 insertions(+), 104 deletions(-) create mode 100644 test/lit/vmi_new/vmi_to_vpto_iota_group1_tail.pto create mode 100644 test/lit/vmi_new/vmi_to_vpto_iota_group_deint_assign.pto diff --git a/docs/isa/vmi-isa/02-index-gen.md b/docs/isa/vmi-isa/02-index-gen.md index 185d0d279b..478ba78982 100644 --- a/docs/isa/vmi-isa/02-index-gen.md +++ b/docs/isa/vmi-isa/02-index-gen.md @@ -9,16 +9,27 @@ ## `pto.vmi.vci` -- **semantics:** Generate a per-lane index/counter vector from a single scalar base such as `[base, base±1, base±2, ...]`, lane `i` gets `base + i` (ASC) or `base - i` (DESC). It is the index source for `vgather`/`vscatter` offsets. +- **semantics:** Generate a per-lane index/counter vector from a single scalar base such as `[base, base±1, base±2, ...]`, lane `i` gets `base + i` (ASC) or `base - i` (DESC). It is the index source for `vgather`/`vscatter` offsets. ```c for (int i = 0; i < L; i++) dst[i] = base + (order == "ASC" ? i : -i); ``` + With `group=C>1`, each group of `S=L/C` lanes restarts the ramp: + + ```c + dst[g*S + j] = base + (order == "ASC" ? j : -j); + ``` + + `group=1` is normalized to ordinary continuous `iota`, so it has exactly the + same semantics and tail support as omitting `group`. Group-periodic iota is + an internal contiguous-only producer; layout assignment inserts + `ensure_layout` when a consumer requests a deinterleaved layout. + - **syntax:** ```mlir - %result = pto.vmi.vci %base {order = "ASC"} : T -> !pto.vmi.vreg + %result = pto.vmi.vci %base {order = "ASC", group = 2} : T -> !pto.vmi.vreg ``` - **operands:** @@ -37,6 +48,7 @@ | Attribute | Values | Default | Description | |---|---|---|---| | `order` | `"ASC"`, `"DESC"` | `"ASC"` | Index generation direction | + | `group` | positive integer | omitted | Number of equal groups. `1` is equivalent to omitted; values greater than one restart the ramp per group. | - **lowering to `pto.mi`:** ``` diff --git a/include/PTO/IR/VMIOps.td b/include/PTO/IR/VMIOps.td index 829a689525..36be32a063 100644 --- a/include/PTO/IR/VMIOps.td +++ b/include/PTO/IR/VMIOps.td @@ -63,23 +63,31 @@ def VMIBroadcastOp : VMI_Op<"broadcast", [Pure]> { def VMIIotaOp : VMI_Op<"iota", [Pure]> { let summary = "Create a VMI logical index vector from a scalar base"; let description = [{ - Without `{group}`, produces a *continuous* ramp over the full logical - length L: `dst[i] = base + i` (ASC). For L=128 that is `base..base+127` + Produces a continuous ramp over the full logical length L: + `dst[i] = base + i` (ASC). For L=128 that is `base..base+127` (physical parts `base+0..63` then `base+64..127`). + }]; + let arguments = (ins + AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, + OptionalAttr:$order + ); + let results = (outs VMI_VRegTypeConstraint:$result); + let hasVerifier = 1; + let assemblyFormat = "$base attr-dict `:` type($base) `->` type($result)"; +} - With `{group = C}`, produces a **group-periodic** ramp (not continuous - across groups): each of the C groups of size `S = L / C` independently - gets `dst[g*S + j] = base + j`. Example `L=128, C=2, base=0`: - `[0..63 | 0..63]`. When `S` is a multiple of physical VL, parts that share - the same ramp may reuse one index register. When `S < VL` and `VL % S == 0` - (e.g. i32 `L=64, C=2` → `[0..31 | 0..31]` inside one VL64), lowering merges - masked `vci` runs via `vsel`. Grouped iota materializes as contiguous; - deinterleaved results go through contiguous iota + `ensure_layout`. +def VMIGroupIotaOp : VMI_Op<"group_iota", [Pure]> { + let summary = "Create a contiguous group-periodic VMI index vector"; + let description = [{ + Internal legacy form for public `vci(..., group=C)` with C > 1. For logical + length L and group size S=L/C, produces + `dst[g*S + j] = base +/- j`. This producer is contiguous-only; consumers + requesting another layout receive an explicit `ensure_layout`. }]; let arguments = (ins AnyTypeOf<[AnyInteger, AnyFloat], "integer/float scalar">:$base, OptionalAttr:$order, - OptionalAttr:$group + I64Attr:$group ); let results = (outs VMI_VRegTypeConstraint:$result); let hasVerifier = 1; diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index a9e7d3d656..f3f9815fc4 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -955,21 +955,38 @@ LogicalResult VMIIotaOp::verify() { if (*order != "ASC" && *order != "DESC") return emitOpError("requires order to be ASC or DESC"); } - if (auto groupAttr = getGroupAttr()) { - int64_t numGroups = groupAttr.getInt(); - if (numGroups <= 0) - return emitOpError("requires group to be positive"); - if (resultType.getElementCount() % numGroups != 0) - return emitOpError("requires group to evenly divide result logical lane " - "count"); - int64_t groupSize = resultType.getElementCount() / numGroups; - FailureOr lanesPerPart = getDataLanesPerPart(elementType); - if (succeeded(lanesPerPart) && groupSize % *lanesPerPart != 0 && - *lanesPerPart % groupSize != 0) - return emitOpError("requires group_size to divide or be a multiple of " - "physical lanes per part (") - << *lanesPerPart << ")"; + return success(); +} + +LogicalResult VMIGroupIotaOp::verify() { + auto resultType = cast(getResult().getType()); + Type elementType = resultType.getElementType(); + if (!isVMIIotaElementType(elementType)) + return emitOpError("requires result element type to be integer 8/16/32 " + "or f16/f32"); + if (!isCompatibleScalarForSemanticType(elementType, getBase().getType())) + return emitOpError("requires base type to match result element type"); + if (std::optional order = getOrder()) { + if (*order != "ASC" && *order != "DESC") + return emitOpError("requires order to be ASC or DESC"); } + + int64_t numGroups = getGroupAttr().getInt(); + if (numGroups <= 1) + return emitOpError("requires group greater than one"); + if (resultType.getElementCount() % numGroups != 0) + return emitOpError("requires group to evenly divide result logical lane " + "count"); + int64_t groupSize = resultType.getElementCount() / numGroups; + FailureOr lanesPerPart = getDataLanesPerPart(elementType); + if (succeeded(lanesPerPart) && groupSize % *lanesPerPart != 0 && + *lanesPerPart % groupSize != 0) + return emitOpError("requires group_size to divide or be a multiple of " + "physical lanes per part (") + << *lanesPerPart << ")"; + if (VMILayoutAttr layout = resultType.getLayoutAttr(); + layout && !layout.isContiguous()) + return emitOpError("requires contiguous result layout"); return success(); } @@ -2626,13 +2643,15 @@ LogicalResult VMIVciOp::verify() { if (resultType.getElementCount() % numGroups != 0) return emitOpError("requires group to evenly divide result logical lane " "count"); - int64_t groupSize = resultType.getElementCount() / numGroups; - FailureOr lanesPerPart = getDataLanesPerPart(elementType); - if (succeeded(lanesPerPart) && groupSize % *lanesPerPart != 0 && - *lanesPerPart % groupSize != 0) - return emitOpError("requires group_size to divide or be a multiple of " - "physical lanes per part (") - << *lanesPerPart << ")"; + if (numGroups > 1) { + int64_t groupSize = resultType.getElementCount() / numGroups; + FailureOr lanesPerPart = getDataLanesPerPart(elementType); + if (succeeded(lanesPerPart) && groupSize % *lanesPerPart != 0 && + *lanesPerPart % groupSize != 0) + return emitOpError("requires group_size to divide or be a multiple of " + "physical lanes per part (") + << *lanesPerPart << ")"; + } } return success(); } diff --git a/lib/PTO/Transforms/VMILayoutAssignment.cpp b/lib/PTO/Transforms/VMILayoutAssignment.cpp index a48cae653c..5f2d3352f0 100644 --- a/lib/PTO/Transforms/VMILayoutAssignment.cpp +++ b/lib/PTO/Transforms/VMILayoutAssignment.cpp @@ -538,6 +538,12 @@ struct LayoutSolver { LogicalResult addConstraints() { WalkResult result = module.walk([&](Operation *op) -> WalkResult { + if (auto groupIota = dyn_cast(op)) { + if (failed(setNaturalLayout(groupIota.getResult(), + getContiguousLayout(), op))) + return WalkResult::interrupt(); + return WalkResult::advance(); + } if (auto maskAnd = dyn_cast(op)) { if (failed(uniteMask(maskAnd.getLhs(), maskAnd.getRhs(), op)) || failed(uniteMask(maskAnd.getLhs(), maskAnd.getResult(), op))) diff --git a/lib/PTO/Transforms/VMILayoutPropagation.cpp b/lib/PTO/Transforms/VMILayoutPropagation.cpp index fb72185c5f..bd2308f1bf 100644 --- a/lib/PTO/Transforms/VMILayoutPropagation.cpp +++ b/lib/PTO/Transforms/VMILayoutPropagation.cpp @@ -351,6 +351,20 @@ class VMIFreeResultLayoutTransfer final : public VMILayoutTransfer { } }; +class VMIContiguousResultLayoutTransfer final : public VMILayoutTransfer { +public: + FailureOr> + query(Operation *op, Value changedValue, VMILayoutAttr changedLayout, + const VMILayoutPropagator &propagator, + OpOperand *changedOperand) const override { + if (!isa(changedValue) || changedValue.getDefiningOp() != op || + !changedLayout.isContiguous()) + return failure(); + return makeSingleRelation(SmallVector{ + valueFact(changedValue, changedLayout)}); + } +}; + class VMILoadTransfer final : public VMILayoutTransfer { public: FailureOr> @@ -809,6 +823,7 @@ const VMILayoutTransfer *getTransfer(Operation *op) { static VMIBitcastTransfer bitcastTransfer; static VMIMaskGranularityCastTransfer maskGranularityCastTransfer; static VMIFreeResultLayoutTransfer freeResultLayoutTransfer; + static VMIContiguousResultLayoutTransfer contiguousResultLayoutTransfer; static VMILoadTransfer loadTransfer; static VMIDeinterleaveLoadTransfer deinterleaveLoadTransfer; static VMIGroupLoadTransfer groupLoadTransfer; @@ -826,6 +841,8 @@ const VMILayoutTransfer *getTransfer(Operation *op) { if (isa(op)) return &freeResultLayoutTransfer; + if (isa(op)) + return &contiguousResultLayoutTransfer; if (isa(op)) return &loadTransfer; if (isa(op)) diff --git a/lib/PTO/Transforms/VMILayoutRematerialize.cpp b/lib/PTO/Transforms/VMILayoutRematerialize.cpp index 35a3901784..1c4bf3f53d 100644 --- a/lib/PTO/Transforms/VMILayoutRematerialize.cpp +++ b/lib/PTO/Transforms/VMILayoutRematerialize.cpp @@ -241,17 +241,19 @@ static std::optional rematerializeDataProducer(Value value, .getResult(); if (auto iota = value.getDefiningOp()) { - // Grouped iota only materializes as contiguous. Keep ensure_layout when - // the consumer wants a non-contiguous layout instead of rematerializing - // a grouped iota directly into deinterleaved/etc. - if (iota.getGroupAttr()) { - VMILayoutAttr resultLayout = resultType.getLayoutAttr(); - if (!resultLayout || !resultLayout.isContiguous()) - return std::nullopt; - } return builder - .create(loc, resultType, iota.getBase(), iota.getOrderAttr(), - iota.getGroupAttr()) + .create(loc, resultType, iota.getBase(), iota.getOrderAttr()) + .getResult(); + } + + if (auto groupIota = value.getDefiningOp()) { + VMILayoutAttr resultLayout = resultType.getLayoutAttr(); + if (!resultLayout || !resultLayout.isContiguous()) + return std::nullopt; + return builder + .create(loc, resultType, groupIota.getBase(), + groupIota.getOrderAttr(), + groupIota.getGroupAttr()) .getResult(); } diff --git a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp index e920ed4e05..9a806ad674 100644 --- a/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp +++ b/lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp @@ -1189,19 +1189,16 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() { // ---- Category A: pure syntactic renames ---- if (auto vop = dyn_cast(op)) { - // vci -> iota (preserve optional {group}). - // Grouped iota only lowers as contiguous; non-contiguous results are - // rewritten to contiguous iota + ensure_layout so layout assignment / - // consumers can still request deinterleaved without a dedicated - // grouped-deint materialization. + // Public vci without grouping (or group=1) is ordinary continuous iota. + // group>1 lowers to the internal contiguous-only group_iota producer. builder.setInsertionPoint(op); StringAttr orderAttr; if (auto order = vop.getOrder()) orderAttr = builder.getStringAttr(*order); Type resultType = vop.getResult().getType(); - Value iotaResult; - if (vop.getGroupAttr()) { + IntegerAttr groupAttr = vop.getGroupAttr(); + if (groupAttr && groupAttr.getInt() > 1) { if (auto vmiTy = dyn_cast(resultType)) { VMILayoutAttr layout = vmiTy.getLayoutAttr(); if (layout && !layout.isContiguous()) { @@ -1211,26 +1208,34 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() { VMILayoutAttr::getContiguous(op->getContext())); Value contig = builder - .create(op->getLoc(), contigType, vop.getBase(), - orderAttr, vop.getGroupAttr()) + .create(op->getLoc(), contigType, + vop.getBase(), orderAttr, groupAttr) .getResult(); - iotaResult = + Value converted = builder .create(op->getLoc(), vmiTy, contig) .getResult(); - vop.getResult().replaceAllUsesWith(iotaResult); + vop.getResult().replaceAllUsesWith(converted); op->erase(); continue; } } + Value grouped = + builder + .create(op->getLoc(), resultType, vop.getBase(), + orderAttr, groupAttr) + .getResult(); + vop.getResult().replaceAllUsesWith(grouped); + op->erase(); + continue; } - iotaResult = + Value iota = builder .create(op->getLoc(), resultType, vop.getBase(), - orderAttr, vop.getGroupAttr()) + orderAttr) .getResult(); - vop.getResult().replaceAllUsesWith(iotaResult); + vop.getResult().replaceAllUsesWith(iota); op->erase(); continue; } @@ -1614,38 +1619,6 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() { } } - // Direct legacy iota with {group} on a non-contiguous layout: rewrite to - // contiguous iota + ensure_layout (same contract as vci above). - SmallVector groupedNonContigIotas; - module.walk([&](VMIIotaOp iota) { - if (!iota.getGroupAttr()) - return; - auto vmiTy = dyn_cast(iota.getResult().getType()); - if (!vmiTy) - return; - VMILayoutAttr layout = vmiTy.getLayoutAttr(); - if (layout && !layout.isContiguous()) - groupedNonContigIotas.push_back(iota); - }); - for (VMIIotaOp iota : groupedNonContigIotas) { - if (!iota->getBlock()) - continue; - OpBuilder builder(iota); - auto vmiTy = cast(iota.getResult().getType()); - Type contigType = VMIVRegType::get( - iota.getContext(), vmiTy.getElementCount(), vmiTy.getElementType(), - VMILayoutAttr::getContiguous(iota.getContext())); - Value contig = - builder - .create(iota.getLoc(), contigType, iota.getBase(), - iota.getOrderAttr(), iota.getGroupAttr()) - .getResult(); - Value converted = - builder.create(iota.getLoc(), vmiTy, contig) - .getResult(); - iota.getResult().replaceAllUsesWith(converted); - iota.erase(); - } } std::unique_ptr mlir::pto::createVMILowerUnifiedToLegacyPass() { diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index 32d0447eea..80c2299c37 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/lib/PTO/Transforms/VMIToVPTO.cpp @@ -5348,11 +5348,14 @@ FailureOr createIotaDeinterleavedChunk(Location loc, Type resultType, .getResult(); } -struct OneToNVMIIotaOpPattern : OpConversionPattern { - using OpConversionPattern::OpConversionPattern; +template +struct OneToNVMIIotaOpPattern : OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + using OneToNOpAdaptor = + typename OpConversionPattern::OneToNOpAdaptor; LogicalResult - matchAndRewrite(VMIIotaOp op, OneToNOpAdaptor adaptor, + matchAndRewrite(IotaOp op, OneToNOpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { auto resultVMIType = cast(op.getResult().getType()); VMILayoutAttr layout = resultVMIType.getLayoutAttr(); @@ -5391,8 +5394,8 @@ struct OneToNVMIIotaOpPattern : OpConversionPattern { // Contiguous materialization: // * S % physVL == 0 → VCI chunk per distinct laneOffset (share parts). // * physVL % S == 0 → sub-VL pack via mask + offset-by-S + vsel. - if (auto groupAttr = op.getGroupAttr()) { - int64_t numGroups = groupAttr.getInt(); + if constexpr (std::is_same_v) { + int64_t numGroups = op.getGroupAttr().getInt(); int64_t logicalLanes = resultVMIType.getElementCount(); if (numGroups <= 0 || logicalLanes % numGroups != 0) return rewriter.notifyMatchFailure( @@ -11627,7 +11630,8 @@ void populateVMIConversionPatterns( typeConverter, patterns.getContext()); patterns.add< OneToNVMIEnsureLayoutOpPattern, OneToNVMIEnsureMaskLayoutOpPattern, - OneToNVMIBroadcastOpPattern, OneToNVMIIotaOpPattern, + OneToNVMIBroadcastOpPattern, OneToNVMIIotaOpPattern, + OneToNVMIIotaOpPattern, OneToNVMIConstantOpPattern, OneToNVMIConstantMaskOpPattern, OneToNVMICreateMaskOpPattern, OneToNVMICreateGroupMaskOpPattern, OneToNVMIMaskBinaryOpPattern, diff --git a/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md index ddb210791c..5c89195475 100644 --- a/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md +++ b/ptodsl/docs/user_guide/14-vmi-virtual-instruction-set.md @@ -464,11 +464,14 @@ its `repeat_stride`. These instructions produce a new logical vector from a scalar seed — either as a lane-wise ramp or a uniform broadcast. -### `pto.vmi.vci(base, *, size, order=None) -> VRegType` +### `pto.vmi.vci(base, *, size, order=None, group=None) -> VRegType` **Description**: Builds a logical lane-wise index ramp starting from a scalar base value. Use it when you need an index vector for lane addressing, -gather/scatter offsets, or dynamic lane selection. +gather/scatter offsets, or dynamic lane selection. With `group=C`, the logical +vector is split into C equal groups and the ramp restarts from `base` in every +group. `group=1` is exactly equivalent to omitting `group`, including logical +tails that do not evenly tile the physical vector length. **Parameters**: @@ -477,6 +480,7 @@ gather/scatter offsets, or dynamic lane selection. | `base` | `ScalarType` | Typed scalar starting value for the ramp | | `size` | `int` | Logical lane count of the result vector | | `order` | `str` or `None` | Ramp order: `"ASC"` for ascending (default if omitted), or `"DESC"` for descending | +| `group` | `int` or `None` | Optional number of equal groups. Values greater than one produce a group-periodic ramp; `1` is equivalent to no grouping. | **Returns**: @@ -491,6 +495,14 @@ idx = pto.vmi.vci(pto.i32(0), size=64, order="ASC") out = pto.vmi.vselr(src, idx) ``` +```python +# [0..31 | 0..31] +idx = pto.vmi.vci(pto.i32(0), size=64, group=2) + +# Identical to an ungrouped 100-lane continuous ramp: [0..99]. +tail = pto.vmi.vci(pto.i32(0), size=100, group=1) +``` + **Constraints**: - `base` must already carry a scalar dtype. A plain Python literal like `0` is ambiguous, so use `pto.i32(0)`, `pto.i16(0)`, `pto.f16(0.0)`, or diff --git a/ptodsl/ptodsl/_vmi_namespace.py b/ptodsl/ptodsl/_vmi_namespace.py index 8acd4ce147..68a9c7f807 100644 --- a/ptodsl/ptodsl/_vmi_namespace.py +++ b/ptodsl/ptodsl/_vmi_namespace.py @@ -310,6 +310,10 @@ def _physical_lanes_per_part(elem_type, *, context: str) -> int | None: def _check_vci_group_tiles_phys_vl(elem_type, size, group, *, context: str) -> None: + # One group is exactly the ordinary continuous iota, including tails that + # do not tile physical VL (for example i32 size=100). + if group == 1: + return group_size = size // group phys = _physical_lanes_per_part(elem_type, context=context) if phys is None: diff --git a/ptodsl/tests/test_vmi_vci_dynamic_index.py b/ptodsl/tests/test_vmi_vci_dynamic_index.py index 03affe4ecd..39bac6cb71 100644 --- a/ptodsl/tests/test_vmi_vci_dynamic_index.py +++ b/ptodsl/tests/test_vmi_vci_dynamic_index.py @@ -154,8 +154,31 @@ def main() -> None: f"untileable group must mention physical lanes, got: {err}", ) - print("ptodsl_vmi_vci_dynamic_index: PASS") + # P2: group=1 is a single group → equivalent to ungrouped; legal even + # when size does not tile physical VL (same as ungrouped size=100). + _check_vci_group_tiles_phys_vl( + i32, 100, 1, context="pto.vmi.vci(...)" + ) + + @pto.jit(target="a5", backend="vpto", mode="explicit") + def vmi_vci_group1_tail_probe(): + dst = pto.alloc_tile(shape=[1, 128], dtype=pto.i32) + idx = pto.vmi.vci(pto.i32(0), size=100, group=1) + pto.vmi.vstore( + idx, dst.as_ptr(), pto.const(0, dtype=pto.index) + ) + + g1_tail = vmi_vci_group1_tail_probe.compile().mlir_text() + expect( + "pto.vmi.vci" in g1_tail, + f"group=1 size=100 must emit vci:\n{g1_tail[:1500]}", + ) + expect( + "!pto.vmi.vreg<100xi32" in g1_tail, + f"group=1 size=100 must keep logical length 100:\n{g1_tail[:1500]}", + ) + print("ptodsl_vmi_vci_dynamic_index: PASS") if __name__ == "__main__": main() diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group1_tail.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group1_tail.pto new file mode 100644 index 0000000000..b5cc6d4564 --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group1_tail.pto @@ -0,0 +1,66 @@ +// 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: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s + +// P2 regression: group=1 must be equivalent to omitting group. +// Ungrouped size=100 already works (vmi_to_vpto_iota_tail.pto). +// It normalizes to ordinary iota before grouped physical-tiling constraints, +// so the legal 100-lane tail remains the continuous ramp 0..99. + +module { + func.func @vmi_to_vpto_iota_group1_tail( + %base: i32, %offset: i32, + %mask: !pto.vmi.mask<100xb32, #pto.vmi.layout>, + %dst: !pto.ptr) { + %c0 = arith.constant 0 : index + %value = pto.vmi.vci %base {group = 1 : i64} + : i32 -> !pto.vmi.vreg<100xi32, #pto.vmi.layout> + %out = pto.vmi.vadds %value, %offset, %mask + : !pto.vmi.vreg<100xi32, #pto.vmi.layout>, i32, + !pto.vmi.mask<100xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<100xi32, #pto.vmi.layout> + pto.vmi.vstore %out, %dst[%c0] + : !pto.vmi.vreg<100xi32, #pto.vmi.layout>, + !pto.ptr + return + } + + func.func @vmi_to_vpto_iota_group1_full_vl( + %base: i32, %offset: i32, + %mask: !pto.vmi.mask<64xb32, #pto.vmi.layout>, + %dst: !pto.ptr) { + %c0 = arith.constant 0 : index + %value = pto.vmi.vci %base {group = 1 : i64} + : i32 -> !pto.vmi.vreg<64xi32, #pto.vmi.layout> + %out = pto.vmi.vadds %value, %offset, %mask + : !pto.vmi.vreg<64xi32, #pto.vmi.layout>, i32, + !pto.vmi.mask<64xb32, #pto.vmi.layout> + -> !pto.vmi.vreg<64xi32, #pto.vmi.layout> + pto.vmi.vstore %out, %dst[%c0] + : !pto.vmi.vreg<64xi32, #pto.vmi.layout>, + !pto.ptr + return + } +} + +// CHECK-LABEL: func.func @vmi_to_vpto_iota_group1_tail( +// Same as ungrouped contiguous tail: two physical vci chunks. +// CHECK: %[[P0:.*]] = pto.vci %arg0 : i32 -> !pto.vreg<64xi32> +// CHECK: arith.constant 64 : i32 +// CHECK: arith.addi %arg0 +// CHECK: %[[P1:.*]] = pto.vci +// CHECK: pto.vadds %[[P0]] +// CHECK: pto.vadds %[[P1]] +// CHECK: pto.vsts +// CHECK: pto.vsts + +// CHECK-LABEL: func.func @vmi_to_vpto_iota_group1_full_vl( +// CHECK: %[[P0:.*]] = pto.vci %arg0 : i32 -> !pto.vreg<64xi32> +// CHECK: pto.vadds %[[P0]] +// CHECK: pto.vsts diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto index ba7cd2b3dd..a5ed86d8d5 100644 --- a/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group_deint.pto @@ -8,9 +8,12 @@ // RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-to-vpto | FileCheck %s -// Grouped vci on a deinterleaved result → contiguous group iota + -// ensure_layout (vdintlv), then sitofp vcvt + interleaved store. -// No unpack; vcvt/vstore are the real consumers. +// Grouped vci on a *pre-annotated* deinterleaved result → contiguous group +// iota + ensure_layout (vdintlv), then sitofp vcvt + interleaved store. +// +// NOTE: this only covers the short-circuit rewrite when layout is already +// present at lower-unified time. Production surface IR has no layout yet — +// see vmi_to_vpto_iota_group_deint_assign.pto for the real assignment path. module { func.func @vmi_to_vpto_iota_group2_deint_vcvt( diff --git a/test/lit/vmi_new/vmi_to_vpto_iota_group_deint_assign.pto b/test/lit/vmi_new/vmi_to_vpto_iota_group_deint_assign.pto new file mode 100644 index 0000000000..d96e57c11e --- /dev/null +++ b/test/lit/vmi_new/vmi_to_vpto_iota_group_deint_assign.pto @@ -0,0 +1,57 @@ +// Copyright (c) 2026 Huawei Technologies Co., Ltd. +// This program is free software, you can redistribute it and/or modify it under the terms and conditions of +// CANN Open Software License Agreement Version 2.0 (the "License"). +// Please refer to the License for details. You may not use this file except in compliance with the License. +// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software repository for the full text of the License. + +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment | FileCheck %s --check-prefix=ASSIGN +// RUN: pto-test-opt %s -vmi-lower-unified-to-legacy -vmi-mask-granularity-assignment -vmi-layout-assignment -vmi-layout-rematerialize -vmi-to-vpto | FileCheck %s --check-prefix=LOWER + +// P1 regression: no-layout surface IR (production path). +// Widening f16→f32 naturally prefers deinterleaved=2. The internal group_iota +// remains contiguous and layout assignment materializes the consumer edge: +// group_iota(contiguous) → ensure_layout(deinterleaved=2) → consumers + +module { + func.func @grouped_vci_elem_with_widen_deint( + %base: f32, %src: !pto.vmi.vreg<128xf16>, %dst: !pto.ptr) { + %c0 = arith.constant 0 : index + %c128 = arith.constant 128 : index + %wide = pto.vmi.vcvt %src + : !pto.vmi.vreg<128xf16> -> !pto.vmi.vreg<128xf32> + %idx = pto.vmi.vci %base {group = 2 : i64} + : f32 -> !pto.vmi.vreg<128xf32> + %mask = pto.vmi.create_mask %c128 + : index -> !pto.vmi.mask<128xpred> + %sum = pto.vmi.vadd %wide, %idx, %mask + : !pto.vmi.vreg<128xf32>, !pto.vmi.vreg<128xf32>, + !pto.vmi.mask<128xpred> + -> !pto.vmi.vreg<128xf32> + pto.vmi.vstore %sum, %dst[%c0] + : !pto.vmi.vreg<128xf32>, !pto.ptr + return + } +} + +// ASSIGN-LABEL: func.func @grouped_vci_elem_with_widen_deint( +// ASSIGN: %[[WIDE:.*]] = pto.vmi.extf +// ASSIGN-SAME: -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> +// Grouped producer must remain contiguous; conversion is via ensure_layout. +// ASSIGN: %[[IOTA:.*]] = pto.vmi.group_iota +// ASSIGN-SAME: -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> +// ASSIGN: %[[IDX:.*]] = pto.vmi.ensure_layout %[[IOTA]] +// ASSIGN-SAME: -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> +// ASSIGN: pto.vmi.addf %[[WIDE]], %[[IDX]] +// ASSIGN-SAME: #pto.vmi.layout +// ASSIGN-NOT: pto.vmi.group_iota {{.*}} -> !pto.vmi.vreg<128xf32, #pto.vmi.layout> + +// LOWER-LABEL: func.func @grouped_vci_elem_with_widen_deint( +// Contiguous group=2 share: one physical vci, then layout conversion. +// LOWER: pto.vci +// LOWER: pto.vdintlv +// LOWER: pto.vadd +// LOWER: pto.vstsx2 +// LOWER-NOT: pto.vmi. +// LOWER-NOT: !pto.vmi. diff --git a/tools/ptoas/ptoas.cpp b/tools/ptoas/ptoas.cpp index 682c0ee25c..8fdaa31afa 100644 --- a/tools/ptoas/ptoas.cpp +++ b/tools/ptoas/ptoas.cpp @@ -3107,8 +3107,9 @@ static void appendVMISemanticPipeline(OpPassManager &pm) { // before any verifier, layout, or lowering pass sees them. pm.addNestedPass( pto::createVMINormalizeSignlessIntToUnsignedPass()); - // Expand unified VMI ops to legacy ops before layout assignment, - // so downstream passes only see legacy ops. + // Expand unified VMI ops before layout assignment so grouped vci becomes + // the contiguous-only legacy group_iota producer. Layout assignment can + // then materialize any consumer-requested non-contiguous use explicitly. pm.addPass(pto::createVMILowerUnifiedToLegacyPass()); pm.addPass(createCanonicalizerPass()); pm.addPass(pto::createVMILegalizeArithSelectPass());