Add TSCATTER

include/PTO/IR/PTOOps.td

@@ -6629,6 +6629,7 @@ def TScatterOp: PTO_TOp<"tscatter", [
     PTODpsType:$src,
     PTODpsType:$dst,
     Optional<PTODpsType>:$indexes,
+    OptionalAttr<StrAttr>:$axis,
     OptionalAttr<PTO_MaskPatternAttr>:$maskPattern
   );
 
@@ -6640,6 +6641,7 @@ def TScatterOp: PTO_TOp<"tscatter", [
   let extraClassDeclaration = [{
     bool hasIndexForm() { return static_cast<bool>(getIndexes()); }
     bool hasMaskForm() { return static_cast<bool>(getMaskPatternAttr()); }
+    bool hasAxis() { return static_cast<bool>(getAxisAttr()); }
     ::mlir::pto::PIPE getPipe() {
       if (hasMaskForm())
         return ::mlir::pto::PIPE::PIPE_V;

lib/PTO/IR/PTO.cpp

@@ -1062,15 +1062,26 @@ ParseResult mlir::pto::TScatterOp::parse(OpAsmParser &parser,
     if (parser.parseKeyword("maskPattern") || parser.parseEqual())
       return failure();
     Attribute rawMaskAttr;
-    if (parser.parseAttribute(rawMaskAttr) || parser.parseRBrace())
+    if (parser.parseAttribute(rawMaskAttr))
       return failure();
     auto mp = llvm::dyn_cast<mlir::pto::MaskPatternAttr>(rawMaskAttr);
     if (!mp)
       return parser.emitError(parser.getCurrentLocation(),
                               "expected #pto.mask_pattern<Pxxxx> for maskPattern");
     result.addAttribute("maskPattern", mp);
     hasMask = true;
-    if (parser.parseColonType(srcTy) || parser.parseRParen())
+    if (parser.parseRBrace() || 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 {
     if (parser.parseOperand(indexes))
@@ -1103,21 +1114,23 @@ ParseResult mlir::pto::TScatterOp::parse(OpAsmParser &parser,
       parser.resolveOperand(dst, dstTy, result.operands) ||
       (hasIndexes && parser.resolveOperand(indexes, idxTy, result.operands)))
     return failure();
+
   return success();
 }
 
 void mlir::pto::TScatterOp::print(OpAsmPrinter &p) {
   p << " ins(" << getSrc() << ", ";
   if (getMaskPatternAttr()) {
-    p << "{maskPattern = " << getMaskPatternAttr() << "} : "
-      << getSrc().getType();
+    p << "{maskPattern = " << getMaskPatternAttr() << "} : " << getSrc().getType();
+    if (auto axisAttr = getAxisAttr())
+      p << ", " << axisAttr;
   } else {
     p << getIndexes() << " : " << getSrc().getType() << ", "
       << getIndexes().getType();
   }
   p << ") outs(" << getDst() << " : " << getDst().getType() << ")";
   p.printOptionalAttrDict((*this)->getAttrs(),
-                          /*elidedAttrs=*/{"maskPattern"});
+                          /*elidedAttrs=*/{"maskPattern", "axis"});
 }
 
 namespace {
@@ -12409,6 +12422,9 @@ mlir::LogicalResult mlir::pto::TScatterOp::verify() {
     return emitOpError(
         "expects exactly one of indexes operand or maskPattern attribute");
   }
+  if (hasIndexes && getAxisAttr()) {
+    return emitOpError("axis attribute must not be provided with indexes operand");
+  }
 
   auto isAllowedDataElem = [&](mlir::Type t) -> bool {
     if (t.isF16() || t.isF32() || t.isBF16()) return true;
@@ -12437,9 +12453,9 @@ mlir::LogicalResult mlir::pto::TScatterOp::verify() {
     Type ts = getSrc().getType();
     Type ti = getIndexes().getType();
     Type td = getDst().getType();
-    if (failed(verifyVecTileStorage(*this, ts, "src")) ||
-        failed(verifyVecTileStorage(*this, ti, "indexes")) ||
-        failed(verifyVecTileStorage(*this, td, "dst")))
+    if (failed(verifyVecTileCommon(*this, ts, "src")) ||
+        failed(verifyVecTileCommon(*this, ti, "indexes")) ||
+        failed(verifyVecTileCommon(*this, td, "dst")))
       return failure();
 
     Type srcElem = getElemTy(ts), dstElem = getElemTy(td), idxElem = getElemTy(ti);
@@ -12463,6 +12479,21 @@ mlir::LogicalResult mlir::pto::TScatterOp::verify() {
     unsigned expectedIdxBytes = (dataBytes == 1) ? 2 : dataBytes;
     if (idxBytes != expectedIdxBytes)
       return emitOpError("expects indexes element size to match the documented scatter rule");
+
+    auto srcValid = getValidShapeVec(ts);
+    auto idxValid = getValidShapeVec(ti);
+    auto dstValid = getValidShapeVec(td);
+    if (srcValid.size() != 2 || idxValid.size() != 2 || dstValid.size() != 2)
+      return emitOpError("expects src, indexes and dst to have rank-2 valid_shape");
+
+    for (unsigned d = 0; d < 2; ++d) {
+      if (srcValid[d] != ShapedType::kDynamic && idxValid[d] != ShapedType::kDynamic &&
+          srcValid[d] != idxValid[d])
+        return emitOpError("expects src and indexes to have the same valid_shape");
+      if (srcValid[d] != ShapedType::kDynamic && dstValid[d] != ShapedType::kDynamic &&
+          dstValid[d] < srcValid[d])
+        return emitOpError("expects dst valid_shape to be >= src valid_shape");
+    }
     return mlir::success();
   };
 
@@ -12488,16 +12519,31 @@ mlir::LogicalResult mlir::pto::TScatterOp::verify() {
     if (srcValid.size() != 2 || dstValid.size() != 2)
       return emitOpError("expects src and dst to have rank-2 valid_shape");
 
+    auto axisAttr = getAxisAttr();
+    if (!axisAttr)
+      return emitOpError("expects mask-pattern tscatter to provide axis attribute");
+    StringRef axisVal = axisAttr.getValue();
     auto mp = getMaskPatternAttr();
     if (!mp)
       return emitOpError("expects mask-pattern tscatter to provide maskPattern");
     const unsigned times = getMaskScatterTimes(mp);
-    if (srcValid[0] != ShapedType::kDynamic && dstValid[0] != ShapedType::kDynamic &&
-        srcValid[0] != dstValid[0])
-      return emitOpError("expects src and dst to have the same valid rows");
-    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");
+    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 &&
+          dstValid[1] != static_cast<int64_t>(srcValid[1] * times))
+        return emitOpError("expects dst valid cols to equal src 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 &&
+          dstValid[0] != static_cast<int64_t>(srcValid[0] * times))
+        return emitOpError("expects dst valid rows to equal src valid rows times the mask expansion factor for col direction");
+    } else {
+        return emitOpError("Invalid axis value, expected \"row\" or \"col\"");
+    }
 
     if (srcTB.getBLayoutValueI32() != static_cast<int32_t>(pto::BLayout::RowMajor) ||
         dstTB.getBLayoutValueI32() != static_cast<int32_t>(pto::BLayout::RowMajor))

lib/PTO/IR/VMI.cpp

@@ -2215,9 +2215,11 @@ LogicalResult VMIScatterOp::verify() {
     return failure();
 
   auto indexElementType = dyn_cast<IntegerType>(indicesType.getElementType());
-  if (!indexElementType || indexElementType.getWidth() != 32 ||
-      indexElementType.isSigned())
-    return emitOpError("requires signless or unsigned 32-bit integer indices");
+  if (!indexElementType || indexElementType.isSigned() ||
+      (indexElementType.getWidth() != 32 &&
+       indexElementType.getWidth() != 16))
+    return emitOpError(
+        "requires signless or unsigned 16-bit or 32-bit integer indices");
 
   if (failed(verifyAllSameVRegShapeAndLayout(getOperation(),
                                              {valueType, indicesType},
@@ -3217,9 +3219,11 @@ LogicalResult VMIVscatterOp::verify() {
 
   auto indexElementType =
       dyn_cast<IntegerType>(offsetsType.getElementType());
-  if (!indexElementType || indexElementType.getWidth() != 32 ||
-      indexElementType.isSigned())
-    return emitOpError("requires signless or unsigned 32-bit integer offsets");
+  if (!indexElementType || indexElementType.isSigned() ||
+      (indexElementType.getWidth() != 32 &&
+       indexElementType.getWidth() != 16))
+    return emitOpError(
+        "requires signless or unsigned 16-bit or 32-bit integer offsets");
 
   if (failed(verifyAllSameVRegShapeAndLayout(getOperation(),
                                              {valueType, offsetsType},

lib/PTO/IR/VPTO.cpp

@@ -6494,10 +6494,8 @@ LogicalResult VscatterOp::verify() {
   auto offsetsElemType = dyn_cast<IntegerType>(offsetsType.getElementType());
   if (!offsetsElemType)
     return emitOpError("offset vector must use integer element type");
-  if (offsetsElemType.getWidth() != 32)
-    return emitOpError("currently requires 32-bit offset vector elements");
-  if (offsetsType.getElementCount() != valueType.getElementCount())
-    return emitOpError("offset and value vectors must have the same element count");
+  if (offsetsElemType.getWidth() != 32 && offsetsElemType.getWidth() != 16)
+    return emitOpError("requires 16-bit or 32-bit offset vector elements");
   if (failed(verifyMaskTypeLike(*this, getMask().getType(), "mask type")))
     return failure();
   MemoryRole destinationRole = classifyMemoryRole(getDestination().getType());

lib/PTO/Transforms/ExpandTileOp.cpp

@@ -517,6 +517,16 @@ static void appendOpContextAttrs(
   if (auto tci = dyn_cast<pto::TCIOp>(op)) {
     attrs.emplace_back("descending", tci.getDescending() ? "true" : "false");
   }
+  if (auto tscatter = dyn_cast<pto::TScatterOp>(op)) {
+    if (auto maskPatternAttr = tscatter.getMaskPatternAttr()) {
+      attrs.emplace_back(
+          "mask_pattern",
+          stringifyMaskPattern(maskPatternAttr.getValue()).str());
+    }
+    if (auto axisAttr = tscatter.getAxisAttr()) {
+      attrs.emplace_back("axis_value", axisAttr.getValue().str());
+    }
+  }
   (void)(tryAppendPrecisionType<pto::TExpOp>(
              op, attrs, pto::ExpPrecision::HighPrecision) ||
          tryAppendPrecisionType<pto::TLogOp>(

lib/PTO/Transforms/InsertTemplateAttributes.cpp

@@ -516,6 +516,16 @@ static void appendOpContextAttrs(
       byte = byteAttr.getInt();
     attrs.emplace_back("byte", std::to_string(byte));
   }
+  if (auto tscatter = dyn_cast<pto::TScatterOp>(op)) {
+    if (auto maskPatternAttr = tscatter.getMaskPatternAttr()) {
+      attrs.emplace_back(
+          "mask_pattern",
+          stringifyMaskPattern(maskPatternAttr.getValue()).str());
+    }
+    if (auto axisAttr = tscatter.getAxisAttr()) {
+      attrs.emplace_back("axis_value", axisAttr.getValue().str());
+    }
+  }
   (void)(tryAppendPrecisionType<pto::TExpOp>(op, attrs) ||
          tryAppendPrecisionType<pto::TLogOp>(op, attrs) ||
          tryAppendPrecisionType<pto::TSqrtOp>(op, attrs) ||

lib/PTO/Transforms/PTOToEmitC.cpp

@@ -12010,10 +12010,14 @@ struct PTOScatterToEmitC : public OpConversionPattern<pto::TScatterOp> {
       auto targs = rewriter.getArrayAttr({
           emitc::OpaqueAttr::get(ctx, maskPatternTok(mp)),
       });
+      SmallVector<Value, 2> operands{dst, src};
+      SmallVector<Attribute, 2> args;
+      if (auto axisAttr = op.getAxisAttr())
+        args.push_back(emitc::OpaqueAttr::get(ctx, axisAttr.getValue().str()));
       rewriter.create<emitc::CallOpaqueOp>(
           loc, TypeRange{}, "TSCATTER",
-          /*args=*/ArrayAttr{}, /*templateArgs=*/targs,
-          /*operands=*/ValueRange{dst, src});
+          /*args=*/rewriter.getArrayAttr(args), /*templateArgs=*/targs,
+          /*operands=*/operands);
     } else {
       Value idx = peelUnrealized(adaptor.getIndexes());
       rewriter.create<emitc::CallOpaqueOp>(

lib/PTO/Transforms/VMIToVPTO.cpp

@@ -2055,15 +2055,24 @@ checkSupportedScatterShape(VMIScatterOp op, std::string *reason) {
     return fail("requires !pto.ptr destination because pto.vscatter is "
                 "pointer-only");
 
-  if (pto::getPTOStorageElemBitWidth(valueType.getElementType()) != 32)
-    return fail("currently requires 32-bit value element type so physical "
-                "index and value lane counts match pto.vscatter");
+  unsigned valueBits =
+      pto::getPTOStorageElemBitWidth(valueType.getElementType());
   auto indexElementType = dyn_cast<IntegerType>(indicesType.getElementType());
-  if (!indexElementType || indexElementType.getWidth() != 32 ||
-      indexElementType.isSigned())
-    return fail("requires signless or unsigned 32-bit indices");
-  if (maskType.getGranularity() != "b32")
-    return fail("requires b32 mask granularity");
+  if (!indexElementType || indexElementType.isSigned())
+    return fail("requires signless or unsigned integer indices");
+  bool isB8Scatter = valueBits == 8 && indexElementType.isUnsigned() &&
+                     indexElementType.getWidth() == 16 &&
+                     maskType.getGranularity() == "b16";
+  bool isB16Scatter = valueBits == 16 && indexElementType.isUnsigned() &&
+                      indexElementType.getWidth() == 16 &&
+                      maskType.getGranularity() == "b16";
+  bool isB32Scatter = valueBits == 32 && indexElementType.getWidth() == 32 &&
+                      maskType.getGranularity() == "b32";
+  if (!isB8Scatter && !isB16Scatter && !isB32Scatter)
+    return fail("requires either 32-bit values with 32-bit indices and b32 "
+                "mask, 16-bit values with unsigned 16-bit indices and b16 "
+                "mask, or 8-bit values with unsigned 16-bit indices and b16 "
+                "mask");
 
   FailureOr<int64_t> valueArity = getVMIPhysicalArity(valueType);
   FailureOr<int64_t> indicesArity = getVMIPhysicalArity(indicesType);

lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp

@@ -3798,6 +3798,33 @@ static FailureOr<StringRef> buildVscatterCallee(MLIRContext *context,
   return buildLaneTypedCallee(context, valueType, "vscatter", ".v300");
 }
 
+static FailureOr<Type> getVscatterOffsetsCarrierType(PatternRewriter &rewriter,
+                                                     Type valueType,
+                                                     Type offsetsType) {
+  Type elementType = getElementTypeFromVectorLike(valueType);
+  auto lanes = getElementCountFromVectorLike(valueType);
+  if (!elementType || !lanes || *lanes <= 0)
+    return failure();
+
+  Type carrierType = offsetsType;
+  unsigned elemBits = pto::getPTOStorageElemBitWidth(elementType);
+  if (elemBits == 16) {
+    if (*lanes % 2 != 0)
+      return failure();
+    carrierType = VectorType::get({*lanes / 2}, rewriter.getI32Type());
+  } else if (elemBits == 8) {
+    if (*lanes % 4 != 0)
+      return failure();
+    carrierType = VectorType::get({*lanes / 4}, rewriter.getI32Type());
+  }
+
+  std::optional<uint64_t> offsetsBits = getFixedVectorBitWidth(offsetsType);
+  std::optional<uint64_t> carrierBits = getFixedVectorBitWidth(carrierType);
+  if (!offsetsBits || !carrierBits || *offsetsBits != *carrierBits)
+    return failure();
+  return carrierType;
+}
+
 static FailureOr<StringRef> buildVaxpyCallee(MLIRContext *context,
                                              Type resultType) {
   return buildCANN900ModeTypedCallee(context, resultType, "vaxpy", "m");
@@ -7542,14 +7569,23 @@ class LowerVscatterOpPattern final
     if (failed(calleeName))
       return rewriter.notifyMatchFailure(op, "unsupported vscatter signature");
 
+    Value offsets = adaptor.getOffsets();
+    FailureOr<Type> offsetsCarrierType = getVscatterOffsetsCarrierType(
+        rewriter, op.getValue().getType(), offsets.getType());
+    if (failed(offsetsCarrierType))
+      return rewriter.notifyMatchFailure(op, "unsupported vscatter offsets carrier");
+    if (offsets.getType() != *offsetsCarrierType)
+      offsets = rewriter.create<LLVM::BitcastOp>(op.getLoc(), *offsetsCarrierType,
+                                                 offsets);
+
     auto funcType = rewriter.getFunctionType(
         TypeRange{adaptor.getValue().getType(), adaptor.getDestination().getType(),
-                  adaptor.getOffsets().getType(), adaptor.getMask().getType()},
+                  *offsetsCarrierType, adaptor.getMask().getType()},
         TypeRange{});
     rewriter.create<func::CallOp>(
         op.getLoc(), *calleeName, TypeRange{},
         ValueRange{adaptor.getValue(), adaptor.getDestination(),
-                   adaptor.getOffsets(), adaptor.getMask()});
+                   offsets, adaptor.getMask()});
     state.plannedDecls.push_back(PlannedDecl{calleeName->str(), funcType});
     rewriter.eraseOp(op);
     return success();