hw-native-sys · frank-deng · Aug 4, 2026 · Aug 4, 2026 · Zhendong404 · Aug 4, 2026
diff --git a/docs/PTO_IR_manual.md b/docs/PTO_IR_manual.md
@@ -7474,7 +7474,7 @@ pto.tconcatidx ins(%src0, %src1, %idx0, %idx1 :
 
 **Summary:** Gathers elements from a source tile using one of three PTO-ISA-compatible forms:
 
-- index gather: `src + indices + tmp -> dst`
+- index gather: `src + indices[+ tmp] -> dst`
 - compare gather: `src + kValue + tmp -> dst + cdst`
 - mask-pattern gather: `src + maskPattern -> dst`
 
@@ -7500,7 +7500,7 @@ Mask form:
 | `dst` | `pto.tile_buf` | Main destination tile |
 | `cdst` | `Optional<pto.tile_buf>` | Secondary destination tile used only by compare form |
 | `indices` | `Optional<pto.tile_buf>` | Index tile used only by index form |
-| `tmp` | `Optional<pto.tile_buf>` | Temporary tile used by index form and compare form |
+| `tmp` | `Optional<pto.tile_buf>` | Temporary tile used by compare form; optionally used by index form on A2/A3 (required), A5 (optional) |
 | `kValue` | `Optional<scalar>` | Scalar compare value used only by compare form |
 | `maskPattern` | `Optional<MaskPatternAttr>` | Mask pattern used only by mask form |
 | `cmpMode` | `Optional<CmpModeAttr>` | Compare mode used only by compare form; defaults to `eq` when omitted |
@@ -7527,25 +7527,25 @@ pto.tgather ins(%src, %kValue, %tmp : !pto.tile_buf<...>, <scalar_type>, !pto.ti
            {cmpMode = #pto<cmp eq|gt>, offset = <i32>}
 
 // mask pattern
-pto.tgather ins(%src, {maskPattern = #pto.mask_pattern<Pxxxx>} : !pto.tile_buf<...>)
+pto.tgather ins(%src, {maskPattern = #pto.mask_pattern<Pxxxx>} : !pto.tile_buf<...>, "row")
            outs(%dst : !pto.tile_buf<...>)
 ```
 
 **Constraints & Verification:**
 
 - Exactly one of the following forms must be used:
-  - index form: `indices` and `tmp`
+  - index form: `indices` (A5: `tmp` is optional; A2/A3: `tmp` is required)
   - compare form: `kValue`, `tmp`, and `cdst`
   - mask form: `maskPattern`
 - **Index gather: implementation checks (A2/A3)**:
   - `src` and `dst` element types must match and be one of `i16/i32/f16/f32`.
   - `indices` element type must be `i32`.
-  - `tmp` element type must match `indices`.
+  - `tmp` is required; `tmp` element type must match `indices`.
   - `indices` and `tmp` must have the same valid shape.
 - **Index gather: implementation checks (A5)**:
   - `src` and `dst` element types must match and be one of `i8/i16/i32/f16/f32`, or a target-supported fp8 type (`f8E4M3*`/`f8E5M2*`).
   - `indices` element type must be `i16` or `i32`.
-  - PTO IR does not impose an extra tmp shape or valid-shape relation in the A5 index form.
+  - `tmp` is optional; PTO IR does not impose an extra tmp shape or valid-shape relation in the A5 index form.
 - **Compare gather: implementation checks (A2/A3)**:
   - `dst` element type must be `i32`.
   - `src` element type must be `f16/f32`, or `i32` when `cmpMode=eq`.
@@ -7583,7 +7583,7 @@ pto.tgather ins(%src, %k, %tmp : !pto.tile_buf<...>, f16, !pto.tile_buf<...>)
            {offset = 7 : i32}
 
 // mask pattern
-pto.tgather ins(%src, {maskPattern = #pto.mask_pattern<P1111>} : !pto.tile_buf<...>)
+pto.tgather ins(%src, {maskPattern = #pto.mask_pattern<P1111>} : !pto.tile_buf<...>, "row")
            outs(%dst : !pto.tile_buf<...>)
 ```
 
@@ -10427,7 +10427,8 @@ pto.comm.tbroadcast(%src, recv(%ping, %pong), group(%g0, %g1, %g2) :
 
 - `group` must be non-empty and all members must have identical types.
 - `dst` element type must match the group element type.
-- `ping` / `pong` must be local VEC tile-like values with matching element type.
+- `ping` / `pong` must be local VEC tile-like values with element type matching `dst`.
+- `root` must be a valid index into `group` (i.e. `root < group.size()`).
 
 **Examples:**
 

diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td
@@ -4869,6 +4869,7 @@ def TGatherOp : PTO_TOp<"tgather", [
     Optional<PTODpsType>:$tmp,
     Optional<ScalarType>:$kValue,
     OptionalAttr<PTO_MaskPatternAttr>:$maskPattern,
+    OptionalAttr<StrAttr>:$axis,
     OptionalAttr<PTO_CmpModeAttr>:$cmpMode,
     OptionalAttr<I32Attr>:$offset
   );
@@ -4882,6 +4883,7 @@ def TGatherOp : PTO_TOp<"tgather", [
     bool hasCompareForm() { return static_cast<bool>(getCdst()) || static_cast<bool>(getKValue()); }
     bool hasIndexForm() { return static_cast<bool>(getIndices()); }
     bool hasMaskForm() { return static_cast<bool>(getMaskPatternAttr()); }
+    bool hasAxis() { return static_cast<bool>(getAxisAttr()); }
     ::mlir::pto::PIPE getPipe() { return ::mlir::pto::PIPE::PIPE_V; }
     ::mlir::MutableOperandRange getDpsInitsMutable() {
       return ::mlir::MutableOperandRange(getOperation(), 1, 1 + (getCdst() ? 1 : 0));

diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp
@@ -851,7 +851,18 @@ ParseResult mlir::pto::TGatherOp::parse(OpAsmParser &parser, OperationState &res
     result.addAttribute("maskPattern", mp);
     hasMask = true;
 
-    if (parser.parseColonType(srcTy) || parser.parseRParen())
+    if (parser.parseColonType(srcTy))
+      return failure();
+    if (succeeded(parser.parseOptionalComma())) {
+      StringAttr axisAttr;
+      if (parser.parseAttribute(axisAttr))
+        return failure();
+      if (axisAttr.getValue() != "row" && axisAttr.getValue() != "col")
+        return parser.emitError(parser.getCurrentLocation(),
+                                "axis must be \"row\" or \"col\"");
+      result.addAttribute("axis", axisAttr);
+    }
+    if (parser.parseRParen())
       return failure();
   } else {
     OpAsmParser::UnresolvedOperand extra;
@@ -983,6 +994,8 @@ void mlir::pto::TGatherOp::print(OpAsmPrinter &p) {
   p << " ins(" << getSrc() << ", ";
   if (auto mp = getMaskPatternAttr()) {
     p << "{maskPattern = " << mp << "} : " << getSrc().getType();
+    if (auto axisAttr = getAxisAttr())
+      p << ", " << axisAttr;
   } else if (getCdst()) {
     p << getKValue();
     if (getTmp()) {
@@ -1012,10 +1025,10 @@ void mlir::pto::TGatherOp::print(OpAsmPrinter &p) {
 
   if (getMaskPatternAttr()) {
     p.printOptionalAttrDict((*this)->getAttrs(),
-                            /*elidedAttrs=*/{"maskPattern", "operandSegmentSizes"});
+                            /*elidedAttrs=*/{"maskPattern", "axis", "operandSegmentSizes"});
   } else {
     p.printOptionalAttrDict((*this)->getAttrs(),
-                            /*elidedAttrs=*/{"operandSegmentSizes"});
+                            /*elidedAttrs=*/{"axis", "operandSegmentSizes"});
   }
 }
 
@@ -7443,6 +7456,46 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() {
       return emitOpError("expects dst valid_shape[1] to equal dst cols");
     }
 
+    auto axisAttr = getAxisAttr();
+    if (!axisAttr)
+      return emitOpError("expects mask-pattern tgather to provide axis attribute");
+    StringRef axisVal = axisAttr.getValue();
+    auto mp = getMaskPatternAttr();
+    if (!mp)
+      return emitOpError("expects mask-pattern tgather to provide maskPattern");
+    auto getMaskGatherTimes = [](mlir::pto::MaskPatternAttr mp) -> unsigned {
+      switch (mp.getValue()) {
+      case mlir::pto::MaskPattern::P1111:
+        return 1;
+      case mlir::pto::MaskPattern::P0101:
+      case mlir::pto::MaskPattern::P1010:
+        return 2;
+      default:
+        return 4;
+      }
+    };
+    const unsigned times = getMaskGatherTimes(mp);
+    auto srcValid = getValidShapeVec(srcTy);
+    if (srcValid.size() != 2 || dstValid.size() != 2)
+      return emitOpError("expects src and dst to have rank-2 valid_shape");
+    if (axisVal == "row") {
+      if (srcValid[0] != ShapedType::kDynamic && dstValid[0] != ShapedType::kDynamic &&
+          dstValid[0] != srcValid[0])
+        return emitOpError("expects dst valid rows to equal src valid rows for row direction");
+      if (srcValid[1] != ShapedType::kDynamic && dstValid[1] != ShapedType::kDynamic &&
+          srcValid[1] != static_cast<int64_t>(dstValid[1] * times))
+        return emitOpError("expects src valid cols to equal dst valid cols times the mask expansion factor for row direction");
+    } else if (axisVal == "col") {
+      if (srcValid[1] != ShapedType::kDynamic && dstValid[1] != ShapedType::kDynamic &&
+          dstValid[1] != srcValid[1])
+        return emitOpError("expects dst valid cols to equal src valid cols for col direction");
+      if (srcValid[0] != ShapedType::kDynamic && dstValid[0] != ShapedType::kDynamic &&
+          srcValid[0] != static_cast<int64_t>(dstValid[0] * times))
+        return emitOpError("expects src valid rows to equal dst valid rows times the mask expansion factor for col direction");
+    } else {
+      return emitOpError("Invalid axis value, expected \"row\" or \"col\"");
+    }
+
     if (allowA5MaskTypes) {
       if (!(srcElemBytes == 1 || srcElemBytes == 2 || srcElemBytes == 4))
         return emitOpError("expects A5 mask-pattern gather element size to be 1, 2, or 4 bytes");
@@ -7460,12 +7513,15 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() {
     Type srcTy = getSrc().getType();
     Type dstTy = getDst().getType();
     Type idxTy = getIndices().getType();
-    Type tmpTy = getTmp().getType();
     if (failed(verifyTileBufCommon(*this, srcTy, "src", allowA5ElemTypes)) ||
         failed(verifyTileBufCommon(*this, dstTy, "dst", allowA5ElemTypes)) ||
-        failed(verifyTileBufCommon(*this, idxTy, "indices")) ||
-        failed(verifyTileBufCommon(*this, tmpTy, "tmp")))
+        failed(verifyTileBufCommon(*this, idxTy, "indices")))
       return failure();
+    if (getTmp()) {
+      Type tmpTy = getTmp().getType();
+      if (failed(verifyTileBufCommon(*this, tmpTy, "tmp")))
+        return failure();
+    }
 
     Type srcElem = getElemTy(srcTy);
     Type dstElem = getElemTy(dstTy);
@@ -7509,10 +7565,10 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() {
     }
 
     if (!allowA5ElemTypes) {
-      Type tmpElem = getElemTy(tmpTy);
+      Type tmpElem = getElemTy(getTmp().getType());
       if (tmpElem != idxElem)
         return emitOpError("expects tmp and indices to have the same element type");
-      if (failed(verifyTileBufSameValidShape(*this, idxTy, tmpTy, "indices", "tmp")))
+      if (failed(verifyTileBufSameValidShape(*this, idxTy, getTmp().getType(), "indices", "tmp")))
         return failure();
     }
     return success();
@@ -7575,6 +7631,8 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() {
         return emitOpError("mask-pattern tgather only allows src and dst operands");
       return verifyMaskForm(/*allowA5MaskTypes=*/false);
     }
+    if (getAxisAttr())
+      return emitOpError("axis attribute must not be provided without maskPattern");
     if (getCdst() || getKValue()) {
       if (!getCdst() || !getKValue() || !getTmp())
         return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp");
@@ -7593,15 +7651,17 @@ llvm::LogicalResult mlir::pto::TGatherOp::verify() {
         return emitOpError("mask-pattern tgather only allows src and dst operands");
       return verifyMaskForm(/*allowA5MaskTypes=*/true);
     }
+    if (getAxisAttr())
+      return emitOpError("axis attribute must not be provided without maskPattern");
     if (getCdst() || getKValue()) {
       if (!getCdst() || !getKValue() || !getTmp())
         return emitOpError("compare-form tgather expects dst, cdst, kValue, and tmp");
       if (getIndices())
         return emitOpError("compare-form tgather does not take indices");
       return verifyCompareForm(/*allowA5SrcTypes=*/true);
     }
-    if (!getIndices() || !getTmp())
-      return emitOpError("index-form tgather expects both indices and tmp");
+    if (!getIndices())
+      return emitOpError("index-form tgather expects indices");
     return verifyIndexForm(/*allow16BitIndices=*/true, /*allowA5ElemTypes=*/true);
   };
 

diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp
@@ -510,6 +510,9 @@ static void appendOpContextAttrs(
           "mask_pattern",
           stringifyMaskPattern(maskPatternAttr.getValue()).str());
     }
+    if (auto axisAttr = tgather.getAxisAttr()) {
+      attrs.emplace_back("axis_value", axisAttr.getValue().str());
+    }
   }
   if (auto ttri = dyn_cast<pto::TTriOp>(op)) {
     attrs.emplace_back("upper_or_lower", std::to_string(ttri.getUpperOrLower()));

diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp
@@ -511,6 +511,9 @@ static void appendOpContextAttrs(
           "mask_pattern",
           stringifyMaskPattern(maskPatternAttr.getValue()).str());
     }
+    if (auto axisAttr = tgather.getAxisAttr()) {
+      attrs.emplace_back("axis_value", axisAttr.getValue().str());
+    }
   }
   if (auto ttri = dyn_cast<pto::TTriOp>(op)) {
     attrs.emplace_back("upper_or_lower", std::to_string(ttri.getUpperOrLower()));

diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp
@@ -9944,15 +9944,17 @@ struct PTOGatherToEmitC : public OpConversionPattern<pto::TGatherOp> {
       return rewriter.notifyMatchFailure(op, (name + " must be emitc::OpaqueType (tile)").str());
     };
 
-    // Case 1: index-based TGATHER(dst, src0, indices, tmp)
+    // Case 1: index-based TGATHER(dst, src0, indices[, tmp])
     if (Value idx = adaptor.getIndices()) {
       idx = peelUnrealized(idx);
-      Value tmp = peelUnrealized(adaptor.getTmp());
+      SmallVector<Value, 4> operands{dst, src0, idx};
+      if (Value tmp = adaptor.getTmp())
+        operands.push_back(peelUnrealized(tmp));
 
       rewriter.create<emitc::CallOpaqueOp>(
           loc, TypeRange{}, "TGATHER",
           /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{},
-          /*operands=*/ValueRange{dst, src0, idx, tmp});
+          /*operands=*/operands);
 
       rewriter.eraseOp(op);
       return success();

diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp
@@ -7053,8 +7053,13 @@ class LowerVstsOpPattern final : public OpConversionPattern<pto::VstsOp> {
     Type elementType = getElementTypeFromVectorLike(op.getValue().getType());
     if (!elementType)
       return rewriter.notifyMatchFailure(op, "unsupported vsts element type");
+    Type offsetElementType = elementType;
+    if (auto ptrType = dyn_cast<pto::PtrType>(op.getDestination().getType()))
+      offsetElementType = ptrType.getElementType();
+    else if (auto memrefType = dyn_cast<BaseMemRefType>(op.getDestination().getType()))
+      offsetElementType = memrefType.getElementType();
     auto offsetBytes =
-        convertElementOffsetToBytes(op, adaptor.getOffset(), elementType);
+        convertElementOffsetToBytes(op, adaptor.getOffset(), offsetElementType);
     auto basePtr = dyn_cast<LLVM::LLVMPointerType>(adaptor.getDestination().getType());
     auto dist =
         parseStoreDistImmediate(op.getDist().value_or(""), elementType);

diff --git a/lib/TileOps/__init__.py b/lib/TileOps/__init__.py
@@ -48,6 +48,7 @@
     ("a5", "pto.tfillpad"): ".a5.tfillpad",
     ("a5", "pto.tfillpad_expand"): ".a5.tfillpad_expand",
     ("a5", "pto.tfillpad_inplace"): ".a5.tfillpad_inplace",
+    ("a5", "pto.tgather"): ".a5.tgather",
     ("a5", "pto.tgatherb"): ".a5.tgatherb",
     ("a5", "pto.tgemv"): ".a5.tgemv",
     ("a5", "pto.tgemv.acc"): ".a5.tgemv_acc",