fix(vmi): dynamic vci index + group-periodic share (raw vci)

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<L×T>
+  %result = pto.vmi.vci %base {order = "ASC", group = 2} : T -> !pto.vmi.vreg<L×T>
   ```
 - **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`:**
   ```

include/PTO/IR/VMIOps.td

@@ -62,6 +62,11 @@ 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 = [{
+    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<StrAttr>:$order
@@ -71,6 +76,24 @@ def VMIIotaOp : VMI_Op<"iota", [Pure]> {
   let assemblyFormat = "$base attr-dict `:` type($base) `->` type($result)";
 }
 
+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<StrAttr>:$order,
+    I64Attr:$group
+  );
+  let results = (outs VMI_VRegTypeConstraint:$result);
+  let hasVerifier = 1;
+  let assemblyFormat = "$base attr-dict `:` type($base) `->` type($result)";
+}
+
 def VMICreateMaskOp : VMI_Op<"create_mask", [Pure]> {
   let summary = "Create a VMI logical prefix predicate mask";
   let arguments = (ins Index:$active_lanes);
@@ -821,9 +844,17 @@ 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}`, 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. 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,
-    OptionalAttr<StrAttr>:$order
+    OptionalAttr<StrAttr>:$order,
+    OptionalAttr<I64Attr>:$group
   );
   let results = (outs VMI_VRegTypeConstraint:$result);
   let hasVerifier = 1;

lib/PTO/IR/VMI.cpp

@@ -958,6 +958,38 @@ LogicalResult VMIIotaOp::verify() {
   return success();
 }
 
+LogicalResult VMIGroupIotaOp::verify() {
+  auto resultType = cast<VMIVRegType>(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<StringRef> 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<int64_t> 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();
+}
+
 LogicalResult VMICreateMaskOp::verify() {
   return success();
 }
@@ -2604,6 +2636,23 @@ 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");
+    if (numGroups > 1) {
+      int64_t groupSize = resultType.getElementCount() / numGroups;
+      FailureOr<int64_t> 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();
 }
 

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<VMIGroupIotaOp>(op)) {
+        if (failed(setNaturalLayout(groupIota.getResult(),
+                                    getContiguousLayout(), op)))
+          return WalkResult::interrupt();
+        return WalkResult::advance();
+      }
       if (auto maskAnd = dyn_cast<VMIMaskAndOp>(op)) {
         if (failed(uniteMask(maskAnd.getLhs(), maskAnd.getRhs(), op)) ||
             failed(uniteMask(maskAnd.getLhs(), maskAnd.getResult(), op)))

lib/PTO/Transforms/VMILayoutPropagation.cpp

@@ -351,6 +351,20 @@ class VMIFreeResultLayoutTransfer final : public VMILayoutTransfer {
   }
 };
 
+class VMIContiguousResultLayoutTransfer final : public VMILayoutTransfer {
+public:
+  FailureOr<SmallVector<VMILayoutRelation, 4>>
+  query(Operation *op, Value changedValue, VMILayoutAttr changedLayout,
+        const VMILayoutPropagator &propagator,
+        OpOperand *changedOperand) const override {
+    if (!isa<OpResult>(changedValue) || changedValue.getDefiningOp() != op ||
+        !changedLayout.isContiguous())
+      return failure();
+    return makeSingleRelation(SmallVector<VMILayoutFact, 4>{
+        valueFact(changedValue, changedLayout)});
+  }
+};
+
 class VMILoadTransfer final : public VMILayoutTransfer {
 public:
   FailureOr<SmallVector<VMILayoutRelation, 4>>
@@ -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<VMIConstantOp, VMIBroadcastOp, VMIIotaOp, VMICreateMaskOp,
           VMICreateGroupMaskOp, VMIConstantMaskOp>(op))
     return &freeResultLayoutTransfer;
+  if (isa<VMIGroupIotaOp>(op))
+    return &contiguousResultLayoutTransfer;
   if (isa<VMILoadOp>(op))
     return &loadTransfer;
   if (isa<VMIDeinterleaveLoadOp>(op))

lib/PTO/Transforms/VMILayoutRematerialize.cpp

@@ -240,10 +240,22 @@ static std::optional<Value> rematerializeDataProducer(Value value,
     return builder.create<VMIBroadcastOp>(loc, resultType, broadcast.getValue())
         .getResult();
 
-  if (auto iota = value.getDefiningOp<VMIIotaOp>())
+  if (auto iota = value.getDefiningOp<VMIIotaOp>()) {
     return builder
         .create<VMIIotaOp>(loc, resultType, iota.getBase(), iota.getOrderAttr())
         .getResult();
+  }
+
+  if (auto groupIota = value.getDefiningOp<VMIGroupIotaOp>()) {
+    VMILayoutAttr resultLayout = resultType.getLayoutAttr();
+    if (!resultLayout || !resultLayout.isContiguous())
+      return std::nullopt;
+    return builder
+        .create<VMIGroupIotaOp>(loc, resultType, groupIota.getBase(),
+                                groupIota.getOrderAttr(),
+                                groupIota.getGroupAttr())
+        .getResult();
+  }
 
   return std::nullopt;
 }

lib/PTO/Transforms/VMILowerUnifiedToLegacy.cpp

@@ -1183,17 +1183,53 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() {
     // ---- Category A: pure syntactic renames ----
 
     if (auto vop = dyn_cast<VMIVciOp>(op)) {
-      // vci -> iota
+      // 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);
-      Value result =
+
+      Type resultType = vop.getResult().getType();
+      IntegerAttr groupAttr = vop.getGroupAttr();
+      if (groupAttr && groupAttr.getInt() > 1) {
+        if (auto vmiTy = dyn_cast<VMIVRegType>(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<VMIGroupIotaOp>(op->getLoc(), contigType,
+                                            vop.getBase(), orderAttr, groupAttr)
+                    .getResult();
+            Value converted =
+                builder
+                    .create<VMIEnsureLayoutOp>(op->getLoc(), vmiTy, contig)
+                    .getResult();
+            vop.getResult().replaceAllUsesWith(converted);
+            op->erase();
+            continue;
+          }
+        }
+        Value grouped =
+            builder
+                .create<VMIGroupIotaOp>(op->getLoc(), resultType, vop.getBase(),
+                                        orderAttr, groupAttr)
+                .getResult();
+        vop.getResult().replaceAllUsesWith(grouped);
+        op->erase();
+        continue;
+      }
+
+      Value iota =
           builder
-              .create<VMIIotaOp>(op->getLoc(), vop.getResult().getType(),
-                                 vop.getBase(), orderAttr)
+              .create<VMIIotaOp>(op->getLoc(), resultType, vop.getBase(),
+                                 orderAttr)
               .getResult();
-      vop.getResult().replaceAllUsesWith(result);
+      vop.getResult().replaceAllUsesWith(iota);
       op->erase();
       continue;
     }
@@ -1561,6 +1597,7 @@ void VMILowerUnifiedToLegacyPass::runOnOperation() {
       continue;
     }
   }
+
 }
 
 std::unique_ptr<Pass> mlir::pto::createVMILowerUnifiedToLegacyPass() {