vgather support i8/ui8 -> i16/ui16

docs/isa/vmi-isa/07-sfu.md

@@ -413,7 +413,7 @@
 
 ### `pto.vmi.vgather`
 
-- **semantics:** Indexed gather from UB at B32 granularity. For each active
+- **semantics:** Indexed gather from UB at B32/B16 granularity. For each active
   lane `i`, load `src[offsets[i]]`.
 
   ```c
@@ -423,7 +423,17 @@
 
 - **syntax:**
   ```mlir
-  %g = pto.vmi.vgather %src, %offsets, %mask : !pto.ptr<T, ub>, !pto.vmi.vreg<L×i32>, !pto.vmi.mask<L> -> !pto.vmi.vreg<L×T>
+  // B32 path
+  %g = pto.vmi.vgather %src, %offsets, %mask
+      : !pto.ptr<T, ub>, !pto.vmi.vreg<L×i32>, !pto.vmi.mask<L×b32> -> !pto.vmi.vreg<L×T>   // T in {i32,ui32,f32}
+
+  // B16 path
+  %g = pto.vmi.vgather %src, %offsets, %mask
+      : !pto.ptr<T, ub>, !pto.vmi.vreg<L×ui16>, !pto.vmi.mask<L×b16> -> !pto.vmi.vreg<L×T>   // T in {i16,ui16,f16,bf16}
+  %g = pto.vmi.vgather %src, %offsets, %mask
+      : !pto.ptr<i8, ub>, !pto.vmi.vreg<L×ui16>, !pto.vmi.mask<L×b16> -> !pto.vmi.vreg<L×i16>
+  %g = pto.vmi.vgather %src, %offsets, %mask
+      : !pto.ptr<ui8, ub>, !pto.vmi.vreg<L×ui16>, !pto.vmi.mask<L×b16> -> !pto.vmi.vreg<L×ui16>
   ```
 - **operands:**
 
@@ -435,12 +445,10 @@
 
 - **results:** `!pto.vmi.vreg<L×T>`
 - **attributes:** `pmode`
-- **datatypes:** `i8`–`i32`, `f16`, `bf16`, `f32`
-- **lowering to `pto.mi`:**
-  ```
-  K × pto.vgather2
-  ```
-  `#mi = K`, `dep = 1`, util data-dependent.
+- **datatypes:** B32 -- `i32`/`ui32`/`f32`; B16 -- `i16`/`ui16`/`f16`/`bf16`,
+  plus `i8`/`ui8` -> `i16`/`ui16` zero-extension.
+- **lowering:** B16 -> `K × pto.vgather2`; B32 -> `K × pto.vgather2_bc`.
+  A statically all-active mask omits the trailing `vsel`.
 
 ### `pto.vmi.vgatherb`
 

include/PTO/IR/VMIOps.td

@@ -710,6 +710,13 @@ def VMIMaskedLoadOp : VMI_Op<"masked_load", [DeclareOpInterfaceMethods<MemoryEff
 
 def VMIGatherOp : VMI_Op<"gather", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "VMI logical masked indexed gather with passthrough lanes";
+  let description = [{
+    Gathers elements from UB memory using an index vector.
+    i8/ui8 sources may be gathered into matching i16/ui16 results. The
+    8-to-16 promotion is a zero-extension: one byte is read per lane and the
+    upper 8 bits are zero-filled, regardless of whether the source element
+    type is signed (i8) or unsigned (ui8). Sign-extension is not supported.
+  }];
   let arguments = (ins PtrOrMemRef:$source,
                        VMI_VRegTypeConstraint:$indices,
                        VMI_MaskTypeConstraint:$mask,
@@ -1584,9 +1591,13 @@ def VMIVdintlvOp : VMI_Op<"vdintlv"> {
 }
 
 def VMIVgatherOp : VMI_Op<"vgather", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
-  let summary = "VMI logical masked indexed gather (B32 granularity)";
+  let summary = "VMI logical masked indexed gather";
   let description = [{
     Gathers elements from UB memory using an index vector.
+    i8/ui8 sources may be gathered into matching i16/ui16 results. The
+    8-to-16 promotion is a zero-extension: one byte is read per lane and the
+    upper 8 bits are zero-filled, regardless of whether the source element
+    type is signed (i8) or unsigned (ui8). Sign-extension is not supported.
     Inactive lanes are controlled by pmode (no explicit passthru operand).
 
     pmode controls inactive-lane behavior:

include/PTO/IR/VPTOOps.td

@@ -2717,6 +2717,12 @@ def PTO_Vmrgsort4Op : PTO_VectorMicroOp<"vmrgsort4"> {
 def PTO_Vgather2Op : PTO_VectorMicroOp<"vgather2", [
     DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
   ]> {
+  let description = [{
+    Masked indexed gather from a UB-backed source. For 8-bit integer sources
+    the result is a 16-bit integer (i8/ui8 -> i16/ui16) produced by
+    zero-extension (one byte per lane, upper 8 bits zero-filled).
+    Sign-extension is not performed.
+  }];
   let arguments = (ins
     PTO_BufferType:$source,
     PTO_VectorType:$offsets,

lib/PTO/IR/VMI.cpp

@@ -562,6 +562,72 @@ static LogicalResult verifyMemoryElementMatches(Operation *op, Type memoryType,
   return success();
 }
 
+// 8->16 gather promotion is a zero-extension (unsigned) operation. signless
+// i8/i16 are accepted and treated as unsigned bytes; sign-extension is not
+// supported (see VMIVgatherOp / Vgather2Op description).
+static bool isVMI8To16GatherPair(Type sourceElemType, Type resultElemType) {
+  auto srcInt = dyn_cast<IntegerType>(sourceElemType);
+  auto resInt = dyn_cast<IntegerType>(resultElemType);
+  if (!srcInt || !resInt ||
+      srcInt.getWidth() != mlir::pto::kValue8 ||
+      resInt.getWidth() != mlir::pto::kValue16) {
+    return false;
+  }
+  if (srcInt.isUnsigned()) {
+    return resInt.isUnsigned();
+  }
+  return !resInt.isUnsigned();
+}
+
+static LogicalResult verifyGatherMemoryElementMatches(
+    Operation *op, Type memoryType, VMIVRegType dataType, StringRef role) {
+  Type memoryElementType = getMemoryElementType(memoryType);
+  if (!memoryElementType) {
+    return success();
+  }
+  if (memoryElementType == dataType.getElementType()) {
+    return success();
+  }
+  if (isVMI8To16GatherPair(memoryElementType, dataType.getElementType())) {
+    return success();
+  }
+  return op->emitOpError()
+         << "requires memory " << role
+         << " element type to match VMI data element type"
+            " or be an 8-bit integer promoted to a matching 16-bit integer";
+}
+
+static bool isSameWidth16BitGatherPair(Type sourceElemType,
+                                       Type resultElemType) {
+  // Existing VMI f16/bf16 path: same-width 16-bit float gather.
+  if (sourceElemType == resultElemType &&
+      (sourceElemType.isF16() || sourceElemType.isBF16())) {
+    return true;
+  }
+  auto srcInt = dyn_cast<IntegerType>(sourceElemType);
+  auto resInt = dyn_cast<IntegerType>(resultElemType);
+  // Existing VMI ui16/i16 path: same-width 16-bit integer gather with matching
+  // integer semantics (signless i16 / i16 is accepted as the non-unsigned side).
+  if (!srcInt || !resInt ||
+      srcInt.getWidth() != mlir::pto::kValue16 ||
+      resInt.getWidth() != mlir::pto::kValue16) {
+    return false;
+  }
+  if (srcInt.isUnsigned()) {
+    return resInt.isUnsigned();
+  }
+  return !resInt.isUnsigned();
+}
+
+static bool isSupported16BitGatherResult(Type sourceElemType,
+                                         Type resultElemType) {
+  // New 8 -> 16 path: i8/ui8 -> i16/ui16 with matching integer semantics.
+  if (isVMI8To16GatherPair(sourceElemType, resultElemType)) {
+    return true;
+  }
+  return isSameWidth16BitGatherPair(sourceElemType, resultElemType);
+}
+
 static LogicalResult verifyContiguousIfLayoutAssigned(Operation *op,
                                                       VMIVRegType type,
                                                       StringRef role) {
@@ -2582,7 +2648,7 @@ LogicalResult VMIGatherOp::verify() {
   auto maskType = cast<VMIMaskType>(getMask().getType());
   auto passthruType = cast<VMIVRegType>(getPassthru().getType());
   auto resultType = cast<VMIVRegType>(getResult().getType());
-  if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(),
+  if (failed(verifyGatherMemoryElementMatches(getOperation(), getSource().getType(),
                                         resultType, "source"))) {
     return failure();
   }
@@ -2609,13 +2675,13 @@ LogicalResult VMIGatherOp::verify() {
     return failure();
   }
 
-  auto resultIntegerType = dyn_cast<IntegerType>(resultType.getElementType());
   if (indexElementType.getWidth() == mlir::pto::kValue16 &&
-      (!resultIntegerType || !resultIntegerType.isUnsigned() ||
-       resultIntegerType.getWidth() != mlir::pto::kValue16)) {
+      !isSupported16BitGatherResult(
+          getMemoryElementType(getSource().getType()),
+          resultType.getElementType())) {
     return emitOpError(
-        "requires ui16 result and passthru element type when using ui16 "
-        "indices");
+        "requires i16/ui16/f16/bf16 result and passthru element type when "
+        "using ui16 indices, or i8/ui8 -> i16/ui16 integer promotion");
   }
   return verifyMaskMatchesData(getOperation(), maskType, resultType);
 }
@@ -4022,7 +4088,7 @@ LogicalResult VMIVgatherOp::verify() {
     return failure();
   }
 
-  if (failed(verifyMemoryElementMatches(getOperation(), getSource().getType(),
+  if (failed(verifyGatherMemoryElementMatches(getOperation(), getSource().getType(),
                                         resultType, "source"))) {
     return failure();
   }
@@ -4044,16 +4110,15 @@ LogicalResult VMIVgatherOp::verify() {
     return failure();
   }
 
-  // 16-bit offsets only address the ui16 gather path (pto.vgather2 / b16 mask),
-  // which requires a ui16 result element type. Reject other 16-bit-offset
-  // results here so the error surfaces at the vgather op rather than later in
-  // the legacy gather it lowers to.
-  auto resultIntegerType = dyn_cast<IntegerType>(resultType.getElementType());
+  // 16-bit offsets address the pto.vgather2 / b16 mask path.  It supports
+  // i16/ui16/f16/bf16 same-width gather and i8/ui8 -> i16/ui16 promotion.
   if (indexElementType.getWidth() == mlir::pto::kValue16 &&
-      (!resultIntegerType || !resultIntegerType.isUnsigned() ||
-       resultIntegerType.getWidth() != mlir::pto::kValue16)) {
+      !isSupported16BitGatherResult(
+          getMemoryElementType(getSource().getType()),
+          resultType.getElementType())) {
     return emitOpError(
-        "requires ui16 result element type when using ui16 offsets");
+        "requires i16/ui16/f16/bf16 result element type when using ui16 "
+        "offsets, or i8/ui8 -> i16/ui16 integer promotion");
   }
 
   if (auto pmode = getPmode()) {

lib/PTO/Transforms/VMIToVPTO.cpp

@@ -1941,14 +1941,37 @@ checkSupportedGatherShape(VMIGatherOp op, std::string *reason) {
   auto indexElementType = dyn_cast<IntegerType>(indicesType.getElementType());
   if (!indexElementType || indexElementType.isSigned())
     return fail("requires signless or unsigned integer indices");
-  bool isU16Gather = resultBits == 16 && indexElementType.isUnsigned() &&
-                     indexElementType.getWidth() == 16 &&
-                     maskType.getGranularity() == "b16";
+  Type sourceElemType = getMemoryElementType(op.getSource().getType());
+  auto sourceInt = dyn_cast<IntegerType>(sourceElemType);
+  auto resultInt = dyn_cast<IntegerType>(resultType.getElementType());
+  bool isB8To16Gather =
+      resultBits == 16 && sourceInt && resultInt &&
+      sourceInt.getWidth() == mlir::pto::kValue8 &&
+      resultInt.getWidth() == mlir::pto::kValue16 &&
+      indexElementType.isUnsigned() &&
+      indexElementType.getWidth() == 16 &&
+      maskType.getGranularity() == "b16" &&
+      ((sourceInt.isUnsigned() && resultInt.isUnsigned()) ||
+       (!sourceInt.isUnsigned() && !resultInt.isUnsigned()));
+  bool isSameWidth16Gather =
+      resultBits == 16 && indexElementType.isUnsigned() &&
+      indexElementType.getWidth() == 16 &&
+      maskType.getGranularity() == "b16" &&
+      ((sourceInt && resultInt &&
+        sourceInt.getWidth() == mlir::pto::kValue16 &&
+        resultInt.getWidth() == mlir::pto::kValue16 &&
+        ((sourceInt.isUnsigned() && resultInt.isUnsigned()) ||
+         (!sourceInt.isUnsigned() && !resultInt.isUnsigned()))) ||
+       (sourceElemType == resultType.getElementType() &&
+        (sourceElemType.isF16() || sourceElemType.isBF16())));
+  bool isB16Gather = isSameWidth16Gather || isB8To16Gather;
   bool isB32Gather = resultBits == 32 && indexElementType.getWidth() == 32 &&
                      maskType.getGranularity() == "b32";
-  if (!isU16Gather && !isB32Gather)
+  if (!isB16Gather && !isB32Gather) {
     return fail("requires either 32-bit results with 32-bit indices and b32 "
-                "mask, or ui16 results with ui16 indices and b16 mask");
+                "mask, or ui16/i16/f16/bf16 results with ui16 indices and "
+                "b16 mask (including i8/ui8 -> i16/ui16 promotion)");
+  }
 
   FailureOr<int64_t> resultArity = getVMIPhysicalArity(resultType);
   FailureOr<int64_t> indicesArity = getVMIPhysicalArity(indicesType);
@@ -1962,7 +1985,16 @@ checkSupportedGatherShape(VMIGatherOp op, std::string *reason) {
     return fail("requires result, indices, passthru, and mask to have the "
                 "same physical arity");
 
-  if (isB32Gather) {
+  // Each pto.vgather2/pto.vgather2_bc emits one physical vector register per
+  // result part. The ISA has a hard limit of four physical registers per
+  // pto.vmi instruction, so reject anything above that instead of silently
+  // lowering to five or more physical gathers.
+  if (*resultArity > mlir::pto::kValue4) {
+    return fail("gather exceeds the 4 physical register limit per VMI "
+                "instruction");
+  }
+
+  if (isB32Gather || (isB16Gather && *resultArity != 1)) {
     std::string resultReason;
     std::string indicesReason;
     std::string passthruReason;
@@ -1978,8 +2010,6 @@ checkSupportedGatherShape(VMIGatherOp op, std::string *reason) {
                   passthruReason);
     if (failed(checkFullVMIPhysicalChunks(maskType, &maskReason)))
       return fail(Twine("mask requires full physical chunks; ") + maskReason);
-  } else if (*resultArity != 1) {
-    return fail("ui16 gather currently supports one physical chunk");
   }
 
   return success();
@@ -7048,6 +7078,13 @@ struct OneToNVMIGatherOpPattern : OneToNOpConversionPattern<VMIGatherOp> {
         indicesParts.size() != passthruParts.size() ||
         indicesParts.size() != resultTypes.size())
       return rewriter.notifyMatchFailure(op, "gather physical arity mismatch");
+    // Static all-active masks select gathered[0] for every lane, so the
+    // trailing vsel is a semantic no-op. Skip it and keep gathered directly.
+    // Non-static masks still take the original gather + vsel path.
+    auto resultVMIType = cast<VMIVRegType>(op.getResult().getType());
+    bool allActive = isStaticAllActiveMask(op.getMask(),
+                                           resultVMIType.getElementCount());
+
 
     SmallVector<Value> results;
     results.reserve(resultTypes.size());
@@ -7069,10 +7106,14 @@ struct OneToNVMIGatherOpPattern : OneToNOpConversionPattern<VMIGatherOp> {
                                  .create<Vgather2BcOp>(op.getLoc(), resultType,
                                                        *source, indices, mask)
                                  .getResult();
-      results.push_back(
-          rewriter
-              .create<VselOp>(op.getLoc(), resultType, gathered, passthru, mask)
-              .getResult());
+      if (allActive) {
+        results.push_back(gathered);
+      } else {
+        results.push_back(
+            rewriter
+                .create<VselOp>(op.getLoc(), resultType, gathered, passthru, mask)
+                .getResult());
+      }
     }
 
     replaceOpWithFlatConvertedValues(rewriter, op, results, *this->getTypeConverter());
@@ -13394,9 +13435,10 @@ verifySupportedVMIToVPTOOps(ModuleOp module,
         return WalkResult::advance();
       gather.emitError()
           << kVMIDiagUnsupportedPrefix
-          << "pto.vmi.gather lowers through pto.vgather2_bc + pto.vsel only "
+          << "pto.vmi.gather lowers through pto.vgather2/pto.vgather2_bc + pto.vsel only "
              "for UB pointer sources, contiguous full physical chunks, "
-             "32-bit result elements, i32 indices, and b32 masks ("
+             "ui16/i16/f16/bf16 results with ui16 indices and b16 masks, "
+             "or 32-bit results with i32 indices and b32 masks ("
           << reason << ")";
       return WalkResult::interrupt();
     }

lib/PTO/Transforms/VPTOMaskSimplify.cpp

@@ -52,11 +52,27 @@ struct SimplifyAllTruePredicateReorder : public OpRewritePattern<OpTy> {
   }
 };
 
+struct SimplifyVselAllTrueMask : public OpRewritePattern<VselOp> {
+  using OpRewritePattern<VselOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(VselOp op,
+                                PatternRewriter &rewriter) const override {
+    // vsel(src0, src1, mask) selects src0 for all-true masks.
+    if (!isAllTrueMask(op.getMask())) {
+      return failure();
+    }
+
+    rewriter.replaceOp(op, op.getSrc0());
+    return success();
+  }
+};
+
 struct VPTOMaskSimplifyPass
     : public pto::impl::VPTOMaskSimplifyBase<VPTOMaskSimplifyPass> {
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
-    patterns.add<SimplifyAllTruePredicateReorder<PintlvB8Op>,
+    patterns.add<SimplifyVselAllTrueMask,
+                 SimplifyAllTruePredicateReorder<PintlvB8Op>,
                  SimplifyAllTruePredicateReorder<PintlvB16Op>,
                  SimplifyAllTruePredicateReorder<PintlvB32Op>,
                  SimplifyAllTruePredicateReorder<PdintlvB8Op>,

ptodsl/ptodsl/_vmi_namespace.py

@@ -432,10 +432,21 @@ def _derive_hist_result_type(acc, *, context: str):
 
 def _derive_vgather_result_type(source, offsets, *, context: str):
     offsets_type = _as_vmi_vreg_type(_type_of(offsets), context=context)
-    result_type = _pointer_element_type(_type_of(source), context=context)
+    result_elem_type = _pointer_element_type(_type_of(source), context=context)
+    # 8-bit integer sources use the B16 gather promotion path (i8 -> i16,
+    # ui8 -> ui16, si8 -> si16).
+    if IntegerType.isinstance(result_elem_type):
+        source_int_type = IntegerType(result_elem_type)
+        if source_int_type.width == 8:
+            if source_int_type.is_unsigned:
+                result_elem_type = IntegerType.get_unsigned(16)
+            elif source_int_type.is_signed:
+                result_elem_type = IntegerType.get_signed(16)
+            else:
+                result_elem_type = IntegerType.get_signless(16)
     return _pto.VMIVRegType.get(
         offsets_type.element_count,
-        result_type,
+        result_elem_type,
         layout=offsets_type.layout,
     )