From 8a203fc91f1e358acd0f8819a9d2dadb2b4ea028 Mon Sep 17 00:00:00 2001 From: frank-deng <13382084679@163.com> Date: Tue, 28 Jul 2026 15:11:19 +0800 Subject: [PATCH 1/2] vscatter add non-32b type support --- lib/PTO/IR/VMI.cpp | 16 +++++--- lib/PTO/IR/VPTO.cpp | 6 +-- lib/PTO/Transforms/VMIToVPTO.cpp | 25 ++++++++---- lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp | 40 ++++++++++++++++++- lib/PTO/Transforms/VPTOLLVMEmitter.cpp | 40 ++++++++++++++++++- tilelang-dsl/python/tilelang_dsl/semantic.py | 5 ++- 6 files changed, 108 insertions(+), 24 deletions(-) diff --git a/lib/PTO/IR/VMI.cpp b/lib/PTO/IR/VMI.cpp index 6767c03870..0dfb613105 100644 --- a/lib/PTO/IR/VMI.cpp +++ b/lib/PTO/IR/VMI.cpp @@ -2215,9 +2215,11 @@ LogicalResult VMIScatterOp::verify() { return failure(); auto indexElementType = dyn_cast(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(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}, diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index 5a2fc23a68..d3aeb6c3ab 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -6494,10 +6494,8 @@ LogicalResult VscatterOp::verify() { auto offsetsElemType = dyn_cast(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()); diff --git a/lib/PTO/Transforms/VMIToVPTO.cpp b/lib/PTO/Transforms/VMIToVPTO.cpp index b76bbd3251..28053eac3a 100644 --- a/lib/PTO/Transforms/VMIToVPTO.cpp +++ b/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(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 valueArity = getVMIPhysicalArity(valueType); FailureOr indicesArity = getVMIPhysicalArity(indicesType); diff --git a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp index c432b313c2..facfd835ae 100644 --- a/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOCANN900LLVMEmitter.cpp @@ -3798,6 +3798,33 @@ static FailureOr buildVscatterCallee(MLIRContext *context, return buildLaneTypedCallee(context, valueType, "vscatter", ".v300"); } +static FailureOr 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 offsetsBits = getFixedVectorBitWidth(offsetsType); + std::optional carrierBits = getFixedVectorBitWidth(carrierType); + if (!offsetsBits || !carrierBits || *offsetsBits != *carrierBits) + return failure(); + return carrierType; +} + static FailureOr 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 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(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( 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(); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 27cf34c93f..9449e6212d 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -3837,6 +3837,33 @@ static FailureOr buildVscatterCallee(MLIRContext *context, return buildLaneTypedCallee(context, valueType, "vscatter", ".v300"); } +static FailureOr 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 offsetsBits = getFixedVectorBitWidth(offsetsType); + std::optional carrierBits = getFixedVectorBitWidth(carrierType); + if (!offsetsBits || !carrierBits || *offsetsBits != *carrierBits) + return failure(); + return carrierType; +} + static FailureOr buildVaxpyCallee(MLIRContext *context, Type resultType) { return buildLaneTypedCallee(context, resultType, "vaxpy", ".m"); @@ -8141,14 +8168,23 @@ class LowerVscatterOpPattern final if (failed(calleeName)) return rewriter.notifyMatchFailure(op, "unsupported vscatter signature"); + Value offsets = adaptor.getOffsets(); + FailureOr 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(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( 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(); diff --git a/tilelang-dsl/python/tilelang_dsl/semantic.py b/tilelang-dsl/python/tilelang_dsl/semantic.py index 12186df658..9af965e8ed 100644 --- a/tilelang-dsl/python/tilelang_dsl/semantic.py +++ b/tilelang-dsl/python/tilelang_dsl/semantic.py @@ -1608,8 +1608,9 @@ def _analyze_vector_store_stmt( offsets_type = self._require_vreg_expr(offsets, "pto.vscatter offsets") if not is_integer_dtype(offsets_type.element_dtype): raise TypeError("pto.vscatter offsets must use an integer vector type in TileLang DSL v1") - if integer_bitwidth(offsets_type.element_dtype) != 32: - raise TypeError("pto.vscatter currently requires i32 offset vectors in TileLang DSL v1") + offsets_bw = integer_bitwidth(offsets_type.element_dtype) + if offsets_bw not in (16, 32): + raise TypeError("pto.vscatter requires i16 or i32 offset vectors in TileLang DSL v1") if value_type.lanes != offsets_type.lanes: raise TypeError("pto.vscatter value and offsets must use the same lane count in TileLang DSL v1") self._require_matching_vector_pointer(value_type, destination.type, "pto.vscatter") From 75412affe8465f5a189c8189c12e192a4da0cfa6 Mon Sep 17 00:00:00 2001 From: frank-deng <13382084679@163.com> Date: Tue, 21 Jul 2026 19:43:12 +0800 Subject: [PATCH 2/2] Add TSCATTER --- include/PTO/IR/PTOOps.td | 2 + lib/PTO/IR/PTO.cpp | 74 +++- lib/PTO/Transforms/ExpandTileOp.cpp | 10 + .../Transforms/InsertTemplateAttributes.cpp | 10 + lib/PTO/Transforms/PTOToEmitC.cpp | 8 +- .../docs/user_guide/08-compute-operations.md | 75 +++- ptodsl/ptodsl/_ops.py | 19 +- ptodsl/ptodsl/_tile_namespace.py | 1 + ptodsl/ptodsl/pto.py | 2 +- ptodsl/ptodsl/tilelib/templates/__init__.py | 1 + .../ptodsl/tilelib/templates/a5/tscatter.py | 227 ++++++++++++ python/pto/dialects/pto.py | 3 + .../lit/pto/tscatter_maskpattern_a5_emitc.pto | 6 +- test/lit/pto/tscatter_maskpattern_emitc.pto | 6 +- ...atter_maskpattern_missing_mode_invalid.pto | 2 +- test/tilelib-st/a5/tscatter/case.py | 340 ++++++++++++++++++ .../tscatter_maskpattern_v0_roundtrip.pto | 2 +- 17 files changed, 761 insertions(+), 27 deletions(-) create mode 100644 ptodsl/ptodsl/tilelib/templates/a5/tscatter.py create mode 100644 test/tilelib-st/a5/tscatter/case.py diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index f5f7938fc6..bee35b431d 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -6629,6 +6629,7 @@ def TScatterOp: PTO_TOp<"tscatter", [ PTODpsType:$src, PTODpsType:$dst, Optional:$indexes, + OptionalAttr:$axis, OptionalAttr:$maskPattern ); @@ -6640,6 +6641,7 @@ def TScatterOp: PTO_TOp<"tscatter", [ let extraClassDeclaration = [{ bool hasIndexForm() { return static_cast(getIndexes()); } bool hasMaskForm() { return static_cast(getMaskPatternAttr()); } + bool hasAxis() { return static_cast(getAxisAttr()); } ::mlir::pto::PIPE getPipe() { if (hasMaskForm()) return ::mlir::pto::PIPE::PIPE_V; diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index af6dd28592..1f1dffdbc7 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -1062,7 +1062,7 @@ 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(rawMaskAttr); if (!mp) @@ -1070,7 +1070,18 @@ ParseResult mlir::pto::TScatterOp::parse(OpAsmParser &parser, "expected #pto.mask_pattern 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(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(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(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(pto::BLayout::RowMajor) || dstTB.getBLayoutValueI32() != static_cast(pto::BLayout::RowMajor)) diff --git a/lib/PTO/Transforms/ExpandTileOp.cpp b/lib/PTO/Transforms/ExpandTileOp.cpp index 5410383e34..1e765769bf 100644 --- a/lib/PTO/Transforms/ExpandTileOp.cpp +++ b/lib/PTO/Transforms/ExpandTileOp.cpp @@ -517,6 +517,16 @@ static void appendOpContextAttrs( if (auto tci = dyn_cast(op)) { attrs.emplace_back("descending", tci.getDescending() ? "true" : "false"); } + if (auto tscatter = dyn_cast(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( op, attrs, pto::ExpPrecision::HighPrecision) || tryAppendPrecisionType( diff --git a/lib/PTO/Transforms/InsertTemplateAttributes.cpp b/lib/PTO/Transforms/InsertTemplateAttributes.cpp index 7b0f9b28e1..006ead855e 100644 --- a/lib/PTO/Transforms/InsertTemplateAttributes.cpp +++ b/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(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(op, attrs) || tryAppendPrecisionType(op, attrs) || tryAppendPrecisionType(op, attrs) || diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 02b2d84f0a..4bfb0aa6f0 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -12010,10 +12010,14 @@ struct PTOScatterToEmitC : public OpConversionPattern { auto targs = rewriter.getArrayAttr({ emitc::OpaqueAttr::get(ctx, maskPatternTok(mp)), }); + SmallVector operands{dst, src}; + SmallVector args; + if (auto axisAttr = op.getAxisAttr()) + args.push_back(emitc::OpaqueAttr::get(ctx, axisAttr.getValue().str())); rewriter.create( 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( diff --git a/ptodsl/docs/user_guide/08-compute-operations.md b/ptodsl/docs/user_guide/08-compute-operations.md index 9825b0edb3..08b5567f40 100644 --- a/ptodsl/docs/user_guide/08-compute-operations.md +++ b/ptodsl/docs/user_guide/08-compute-operations.md @@ -283,6 +283,79 @@ namespace. --- +#### `pto.tile.scatter(src: Tile, dst: Tile, *, indexes: Tile | None = None, axis: str | None = None, mask_pattern: str | None = None) -> None` + +**Description**: Scatters elements from `src` into `dst`. Two modes are available: + +**Index mode** (pass `indexes`, omit `mask_pattern`): Each element `src[i, j]` is written to `dst` at the column offset specified by `indexes[i, j]`. The destination tile is zero-initialized before scattering. Uses `pto.vscatter` under the hood. + +**Mask-pattern mode** (pass `mask_pattern` and `axis`, omit `indexes`): Elements from `src` are scattered into `dst` with a regular spacing pattern controlled by `mask_pattern` and `axis`. The destination tile is zero-initialized before scattering. + +- `axis="row"`: source elements are scattered across columns within each row, interleaved with zeros according to the mask pattern. +- `axis="col"`: source elements are scattered across rows within each column, placed at strided row positions. + +Supported mask patterns: + +| Pattern | Row semantics (elements placed at column multiples) | Column semantics (stride, start) | +|---------|------------------------------------------------------|----------------------------------| +| `P1111` | Direct copy (no interleaving) | Direct copy (stride=1, start=0) | +| `P0101` | Every 2nd col, starting at 1 | stride=2, start=1 | +| `P1010` | Every 2nd col, starting at 0 | stride=2, start=0 | +| `P0001` | Every 4th col, starting at 3 | stride=4, start=3 | +| `P0010` | Every 4th col, starting at 2 | stride=4, start=2 | +| `P0100` | Every 4th col, starting at 1 | stride=4, start=1 | +| `P1000` | Every 4th col, starting at 0 | stride=4, start=0 | + +**Parameters**: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `src` | `Tile` | Source tile containing data to scatter | +| `dst` | `Tile` | Destination tile (zero-initialized, then receives scattered data) | +| `indexes` | `Tile | None` | Index tile (`i32` or `ui32` dtype) specifying per-element column offsets (index mode) | +| `axis` | `str | None` | Scatter direction: `"row"` or `"col"` (mask-pattern mode) | +| `mask_pattern` | `str | None` | Spacing pattern: `"P1111"`, `"P0101"`, `"P1010"`, `"P0001"`, `"P0010"`, `"P0100"`, or `"P1000"` (mask-pattern mode) | + +**Returns**: None (writes to `dst`). + +**Constraints**: + +- A5 target only. +- Supported element types: `i8`, `i16`, `i32`, `ui8`, `ui16`, `ui32`, `f16`, `bf16`, `f32`. +- Index mode: `indexes` must have `i32` or `ui32` dtype and the same shape as `src`. +- Mask-pattern mode: exactly one of `axis` and `mask_pattern` must be provided together; `indexes` must not be set. +- `dst` must use row-major layout in UB memory space. +- Runs on `PIPE_V` (vector pipe). + +**Example** — index-mode scatter: + +```python +src_tile = pto.alloc_tile(shape=[4, 32], dtype=pto.f32) +dst_tile = pto.alloc_tile(shape=[4, 32], dtype=pto.f32) +idx_tile = pto.alloc_tile(shape=[4, 32], dtype=pto.i32) +pto.tile.scatter(src_tile, dst_tile, indexes=idx_tile) +``` + +**Example** — mask-pattern scatter along rows: + +```python +src_tile = pto.alloc_tile(shape=[4, 32], dtype=pto.f16) +dst_tile = pto.alloc_tile(shape=[4, 64], dtype=pto.f16) +pto.tile.scatter(src_tile, dst_tile, axis="row", mask_pattern="P0101") +``` + +**Example** — mask-pattern scatter along columns: + +```python +src_tile = pto.alloc_tile(shape=[4, 32], dtype=pto.f32) +dst_tile = pto.alloc_tile(shape=[16, 32], dtype=pto.f32) +pto.tile.scatter(src_tile, dst_tile, axis="col", mask_pattern="P0010") +``` + +The low-level alias `pto.tscatter` is also available when a kernel needs to bypass the `pto.tile` namespace. + +--- + ### 8.1.7 Broadcast and expansion Expansion ops take a narrow source (scalar, row vector, or column vector) and broadcast it to a full tile shape. They are useful for applying per-row or per-column coefficients to a tile. @@ -1453,7 +1526,7 @@ pto.tile.store(dst_tile, out_view) | Activation | `tile.relu`, `tile.lrelu` | | Row reductions | `tile.rowsum`, `tile.rowmax`, `tile.rowmin`, `tile.rowprod`, `tile.rowargmax`, `tile.rowargmin` | | Column reductions | `tile.colsum`, `tile.colmax`, `tile.colmin`, `tile.colprod` | -| Sort/gather | `tile.sort32`, `tile.mrgsort`, `tile.gather` | +| Sort/gather/scatter | `tile.sort32`, `tile.mrgsort`, `tile.gather`, `tile.scatter` | | Broadcast | `tile.expands`, `tile.rowexpand`, `tile.colexpand` | | Row-expand arith | `tile.rowexpandadd`, `tile.rowexpandsub`, `tile.rowexpandmul`, `tile.rowexpanddiv`, `tile.rowexpandmax`, `tile.rowexpandmin`, `tile.rowexpandexpdif` | | Col-expand arith | `tile.colexpandadd`, `tile.colexpandsub`, `tile.colexpandmul`, `tile.colexpanddiv`, `tile.colexpandmax`, `tile.colexpandmin`, `tile.colexpandexpdif` | diff --git a/ptodsl/ptodsl/_ops.py b/ptodsl/ptodsl/_ops.py index 062933443e..0642452e4a 100644 --- a/ptodsl/ptodsl/_ops.py +++ b/ptodsl/ptodsl/_ops.py @@ -3788,6 +3788,23 @@ def tci(start, dst, *, tmp=None, descending=False): ) +def tscatter(src, dst, *, indexes=None, axis=None, mask_pattern=None): + """``pto.tscatter`` tile scatter wrapper.""" + if indexes is not None and (axis is not None or mask_pattern is not None): + raise ValueError("indexes and axis/mask_pattern cannot be provided together") + if indexes is None and (axis is None or mask_pattern is None): + raise ValueError(f"non indexes mode must provide both axis and mask_pattern") + if axis is not None and axis not in ("row", "col"): + raise ValueError(f"axis must be 'row' or 'col', got {axis!r}") + _pto.tscatter( + unwrap_surface_value(src), + unwrap_surface_value(dst), + indexes=None if indexes is None else unwrap_surface_value(indexes), + axis=None if axis is None else axis, + mask_pattern=None if mask_pattern is None else _tile_mask_pattern_attr(mask_pattern), + ) + + def tsel(mask, src0, src1, dst, *, tmp=None): """``pto.tsel ins(mask, src0, src1, tmp) outs(dst)`` with synthesized scratch when omitted.""" resolved_tmp = tmp if tmp is not None else _resolve_selection_tmp(dst, tmp, context="tsel") @@ -6262,7 +6279,7 @@ def import_reserved_buffer(name, *, peer_func): "texpands", "treshape", "trowexpand", "tcolexpand", "trowexpandadd", "trowexpandsub", "trowexpandmul", "trowexpanddiv", "trowexpandmax", "trowexpandmin", "trowexpandexpdif", "tcolexpandadd", "tcolexpandsub", "tcolexpandmul", "tcolexpanddiv", "tcolexpandmax", "tcolexpandmin", "tcolexpandexpdif", - "tsort32", "tmrgsort", "tgather", + "tsort32", "tmrgsort", "tgather", "tscatter", "tsel", "tsels", "tcvt", "tnot", "tand", "tands", "tor", "tors", "txor", "txors", "tshl", "tshls", "tshr", "tshrs", "tpartadd", "tpartmul", "tpartmax", "tpartmin", diff --git a/ptodsl/ptodsl/_tile_namespace.py b/ptodsl/ptodsl/_tile_namespace.py index 0ceb3c534b..26cbaeafef 100644 --- a/ptodsl/ptodsl/_tile_namespace.py +++ b/ptodsl/ptodsl/_tile_namespace.py @@ -179,6 +179,7 @@ def rowargmin(src, dst, *, tmp=None): mrgsort = staticmethod(_ops.tmrgsort) gather = staticmethod(_ops.tgather) gatherb = staticmethod(_ops.tgatherb) + scatter = staticmethod(_ops.tscatter) tri = staticmethod(_ops.ttri) histogram = staticmethod(_ops.tthistogram) diff --git a/ptodsl/ptodsl/pto.py b/ptodsl/ptodsl/pto.py index 0598fb1453..aa3208f616 100644 --- a/ptodsl/ptodsl/pto.py +++ b/ptodsl/ptodsl/pto.py @@ -118,7 +118,7 @@ vsel, make_tensor_view, partition_view, alloc_buffer, alloc_tile, - tsort32, tmrgsort, tgather, + tsort32, tmrgsort, tgather, tscatter, mte_load, mte_store, mte_gm_ub, mte_ub_gm, mte_ub_ub, mte_ub_l1, mte_gm_l1, mte_l1_ub, mte_gm_l1_frac, mte_l1_bt, mte_l1_fb, mem_bar, mte_l1_l0a, mte_l1_l0b, mte_l1_l0a_mx, mte_l1_l0b_mx, diff --git a/ptodsl/ptodsl/tilelib/templates/__init__.py b/ptodsl/ptodsl/tilelib/templates/__init__.py index eee2fa0ae0..7105c77de8 100644 --- a/ptodsl/ptodsl/tilelib/templates/__init__.py +++ b/ptodsl/ptodsl/tilelib/templates/__init__.py @@ -109,6 +109,7 @@ ("a5", "pto.trowmin"): ".a5.trowmin", ("a5", "pto.trowprod"): ".a5.trowprod", ("a5", "pto.trowsum"): ".a5.trowsum", + ("a5", "pto.tscatter"): ".a5.tscatter", ("a5", "pto.tsel"): ".a5.tsel", ("a5", "pto.tsels"): ".a5.tsels", ("a5", "pto.tshl"): ".a5.tshl", diff --git a/ptodsl/ptodsl/tilelib/templates/a5/tscatter.py b/ptodsl/ptodsl/tilelib/templates/a5/tscatter.py new file mode 100644 index 0000000000..01c291c90d --- /dev/null +++ b/ptodsl/ptodsl/tilelib/templates/a5/tscatter.py @@ -0,0 +1,227 @@ +# 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. +"""PTODSL TileLib templates for ``pto.tscatter``.""" + +from ptodsl import pto, scalar +import ptodsl.tilelib as tilelib + +from ._common import NUMERIC_DTYPES +from ._row_arg import _scalar_literal + + +def _init_ub_buffer(dst: pto.Tile): + dtype = dst.dtype + tile_rows, tile_cols = dst.shape + lanes = pto.elements_per_vreg(dtype) + zeros = pto.vbr(_scalar_literal(dtype, 0)) + mask, _ = pto.make_mask(dtype, lanes) + for row in range(0, tile_rows, 1): + for col in range(0, tile_cols, lanes): + pto.vsts(zeros, dst[row, col:], mask) + pto.mem_bar(pto.BarrierType.VST_VST) + + +_SCATTER_MASK_DTYPES = [(dtype, dtype) for dtype in NUMERIC_DTYPES] + +_SCATTER_INDEX_DTYPES = [ + (dtype, dtype, idx) + for dtype in NUMERIC_DTYPES + for idx in ( + ("i16", "ui16") + if dtype in ("i8", "ui8", "i16", "ui16", "f16", "bf16") + else ("i32", "ui32") + ) +] + + +def _static_valid_dim(tile, index): + static_valid_shape = getattr(tile, "_template_static_valid_shape", None) + if static_valid_shape is None: + static_valid_shape = getattr(tile, "static_valid_shape", None) + if static_valid_shape is not None and static_valid_shape[index] is not None: + return static_valid_shape[index] + return tile.valid_shape[index] + + +def _template_tscatter_mask_row(src: pto.Tile, dst: pto.Tile, interleave_args): + dtype = dst.dtype + valid_rows, valid_cols = src.valid_shape + lanes = pto.elements_per_vreg(dtype) + times = 1 << len(interleave_args) + dst_valid_col = valid_cols * times + zeros = pto.vbr(_scalar_literal(dtype, 0)) + for row in range(0, valid_rows, 1): + py_rem = dst_valid_col + for col in range(0, valid_cols, lanes): + src_reg = pto.vlds(src[row, col:]) + if not interleave_args: + mask, _ = pto.make_mask(dtype, py_rem) + pto.vsts(src_reg, dst[row, col:], mask) + py_rem -= lanes + elif len(interleave_args) == 1: + if interleave_args[0]: + reg0, reg1 = pto.vintlv(src_reg, zeros) + else: + reg0, reg1 = pto.vintlv(zeros, src_reg) + mask, _ = pto.make_mask(dtype, py_rem) + pto.vsts(reg0, dst[row, col * times:], mask) + py_rem -= lanes + if py_rem > 0: + mask, _ = pto.make_mask(dtype, py_rem) + pto.vsts(reg1, dst[row, col * times + lanes:], mask) + py_rem -= lanes + else: + if interleave_args[0]: + tmp0, tmp1 = pto.vintlv(src_reg, zeros) + else: + tmp0, tmp1 = pto.vintlv(zeros, src_reg) + if interleave_args[1]: + reg0, reg1 = pto.vintlv(tmp0, zeros) + reg2, reg3 = pto.vintlv(tmp1, zeros) + else: + reg0, reg1 = pto.vintlv(zeros, tmp0) + reg2, reg3 = pto.vintlv(zeros, tmp1) + mask, _ = pto.make_mask(dtype, py_rem) + pto.vsts(reg0, dst[row, col * times:], mask) + py_rem -= lanes + if py_rem > 0: + mask, _ = pto.make_mask(dtype, py_rem) + pto.vsts(reg1, dst[row, col * times + lanes:], mask) + py_rem -= lanes + if py_rem > 0: + mask, _ = pto.make_mask(dtype, py_rem) + pto.vsts(reg2, dst[row, col * times + lanes * 2:], mask) + py_rem -= lanes + if py_rem > 0: + mask, _ = pto.make_mask(dtype, py_rem) + pto.vsts(reg3, dst[row, col * times + lanes * 3:], mask) + py_rem -= lanes + + +def _no_mask_pattern(mask_pattern=True): + return mask_pattern is True + + +def _axis_is_row(axis_value, **_): + return axis_value == "row" + + +def _axis_is_col(axis_value, **_): + return axis_value == "col" + + +@tilelib.tile_template( + op="pto.tscatter", + target="a5", + name="template_tscatter_index", + dtypes=_SCATTER_INDEX_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + layouts=("row_major",), + loop_depth=2, + is_post_update=False, + constraints=(_no_mask_pattern,), + id=0, +) +def template_tscatter_index(src: pto.Tile, dst: pto.Tile, idx: pto.Tile): + _init_ub_buffer(dst) + valid_rows, valid_cols = dst.valid_shape + dtype = dst.dtype + idx_dtype = idx.dtype + dst_ptr = dst.as_ptr() + valid_rows, valid_cols = dst.valid_shape + lanes = pto.elements_per_vreg(idx_dtype) + src_dist = "UNPK_B8" if pto.bytewidth(dtype) == 1 else None + for row in range(0, valid_rows, 1): + remained = valid_cols + for col in range(0, valid_cols, lanes): + idx_mask, remained = pto.make_mask(idx_dtype, remained) + src_reg = pto.vlds(src[row, col:], dist=src_dist) + idx_reg = pto.vlds(idx[row, col:]) + pto.vscatter(src_reg, dst_ptr, idx_reg, idx_mask) + + +def _template_tscatter_mask_col(src: pto.Tile, dst: pto.Tile, start, stride): + dtype = dst.dtype + valid_rows, valid_cols = src.valid_shape + lanes = pto.elements_per_vreg(dtype) + for row in range(0, valid_rows, 1): + remained = valid_cols + row_dst = row * stride + start + for col in range(0, valid_cols, lanes): + src_reg = pto.vlds(src[row, col:]) + mask, remained = pto.make_mask(dtype, remained) + pto.vsts(src_reg, dst[row_dst, col:], mask) + + +_MASK_PATTERN_TO_INTERLEAVE = { + "P1111": (), + "P0101": (True,), + "P1010": (False,), + "P0001": (True, True), + "P0010": (True, False), + "P0100": (False, True), + "P1000": (False, False), +} + +_MASK_PATTERN_TO_STRIDE = { + "P1111": (0, 1), + "P0101": (0, 2), + "P1010": (1, 2), + "P0001": (0, 4), + "P0010": (1, 4), + "P0100": (2, 4), + "P1000": (3, 4), +} + + +@tilelib.tile_template( + op="pto.tscatter", + target="a5", + name="template_tscatter_mask_row", + dtypes=_SCATTER_MASK_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + layouts=("row_major",), + loop_depth=2, + is_post_update=False, + tags=("scatter", "mask"), + constraints=(_axis_is_row,), + id=1, +) +def template_tscatter_mask_row(src: pto.Tile, dst: pto.Tile): + mask_pattern = pto.get_op_attr("mask_pattern", "P1111") + if mask_pattern != "P1111": + _init_ub_buffer(dst) + interleave_args = _MASK_PATTERN_TO_INTERLEAVE[mask_pattern] + _template_tscatter_mask_row(src, dst, interleave_args) + + +@tilelib.tile_template( + op="pto.tscatter", + target="a5", + name="template_tscatter_mask_col", + dtypes=_SCATTER_MASK_DTYPES, + iteration_axis="none", + op_engine="vector", + op_class="other", + layouts=("row_major",), + loop_depth=2, + is_post_update=False, + tags=("scatter", "mask"), + constraints=(_axis_is_col,), + id=2, +) +def template_tscatter_mask_col(src: pto.Tile, dst: pto.Tile): + mask_pattern = pto.get_op_attr("mask_pattern", "P1111") + if mask_pattern != "P1111": + _init_ub_buffer(dst) + start, stride = _MASK_PATTERN_TO_STRIDE[mask_pattern] + _template_tscatter_mask_col(src, dst, start, stride) diff --git a/python/pto/dialects/pto.py b/python/pto/dialects/pto.py index 3e66f0bdb1..975f060982 100644 --- a/python/pto/dialects/pto.py +++ b/python/pto/dialects/pto.py @@ -964,6 +964,7 @@ def __init__( dst=_TSCATTER_UNSET, indexes=_TSCATTER_UNSET, maskPattern=_TSCATTER_UNSET, + axis=None, loc=None, ip=None, ): @@ -1035,6 +1036,8 @@ def _matches_src_type(value): kwargs["indexes"] = indexes if maskPattern is not _TSCATTER_UNSET: kwargs["maskPattern"] = maskPattern + if axis is not None: + kwargs["axis"] = axis super().__init__(src, **kwargs, loc=loc, ip=ip) diff --git a/test/lit/pto/tscatter_maskpattern_a5_emitc.pto b/test/lit/pto/tscatter_maskpattern_a5_emitc.pto index 351465515e..48a2e00e97 100644 --- a/test/lit/pto/tscatter_maskpattern_a5_emitc.pto +++ b/test/lit/pto/tscatter_maskpattern_a5_emitc.pto @@ -2,9 +2,9 @@ module attributes {pto.target_arch = "a5"} { func.func @tscatter_maskpattern_a5() { - %src = pto.alloc_tile : !pto.tile_buf - %dst = pto.alloc_tile : !pto.tile_buf - pto.tscatter ins(%src, {maskPattern = #pto.mask_pattern} : !pto.tile_buf) outs(%dst : !pto.tile_buf) + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tscatter ins(%src, {maskPattern = #pto.mask_pattern} : !pto.tile_buf, "row") outs(%dst : !pto.tile_buf) func.return } } diff --git a/test/lit/pto/tscatter_maskpattern_emitc.pto b/test/lit/pto/tscatter_maskpattern_emitc.pto index 05314de55c..08064f7440 100644 --- a/test/lit/pto/tscatter_maskpattern_emitc.pto +++ b/test/lit/pto/tscatter_maskpattern_emitc.pto @@ -2,9 +2,9 @@ module attributes {pto.target_arch = "a2a3"} { func.func @tscatter_maskpattern() { - %src = pto.alloc_tile : !pto.tile_buf - %dst = pto.alloc_tile : !pto.tile_buf - pto.tscatter ins(%src, {maskPattern = #pto.mask_pattern} : !pto.tile_buf) outs(%dst : !pto.tile_buf) + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tscatter ins(%src, {maskPattern = #pto.mask_pattern} : !pto.tile_buf, "row") outs(%dst : !pto.tile_buf) func.return } } diff --git a/test/lit/pto/tscatter_maskpattern_missing_mode_invalid.pto b/test/lit/pto/tscatter_maskpattern_missing_mode_invalid.pto index d9478cdca1..335dbd4074 100644 --- a/test/lit/pto/tscatter_maskpattern_missing_mode_invalid.pto +++ b/test/lit/pto/tscatter_maskpattern_missing_mode_invalid.pto @@ -4,7 +4,7 @@ module attributes {pto.target_arch = "a2a3"} { func.func @tscatter_missing_mode() { %src = pto.alloc_tile : !pto.tile_buf %dst = pto.alloc_tile : !pto.tile_buf - "pto.tscatter"(%src, %dst) : (!pto.tile_buf, !pto.tile_buf) -> () + "pto.tscatter"(%src, %dst) <{}> : (!pto.tile_buf, !pto.tile_buf) -> () func.return } } diff --git a/test/tilelib-st/a5/tscatter/case.py b/test/tilelib-st/a5/tscatter/case.py new file mode 100644 index 0000000000..d0a76ec544 --- /dev/null +++ b/test/tilelib-st/a5/tscatter/case.py @@ -0,0 +1,340 @@ +#!/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. + +# PTODSL rewrite of pto-isa/tests/npu/a5/src/st/testcase/tscatter. +# +# TSCATTER scatters elements from a source tile into a destination tile. +# Two modes are supported: +# - Indexed scatter: each element of src is written to dst at the position +# specified by the corresponding index value. +# - Masked scatter: src elements are placed into dst at positions determined +# by a mask pattern (e.g., P0101 interleaves with zeros at 2x expansion). +# Both row and col directions are tested. + +from pathlib import Path +import sys +import zlib + +import numpy as np + +if __package__ in {None, ""}: + sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from common import auto_main, golden_output_case +from ptodsl import pto + + +PTO_TO_NP_DTYPE = { + pto.f32: np.float32, + pto.f16: np.float16, + pto.i32: np.int32, + pto.i16: np.int16, + pto.i8: np.int8, + pto.ui32: np.uint32, + pto.ui16: np.uint16, + pto.ui8: np.uint8, +} + + +def npy_dtype(pto_type) -> np.dtype: + return PTO_TO_NP_DTYPE[pto_type] + + +# --- Indexed scatter cases --- +# (name, src_dtype, idx_dtype, src_shape, idx_shape) +INDEX_CASES = [ + ("i8_uint16_256x32_256x32", pto.i8, pto.ui16, (256, 32), (256, 32)), + ("ui8_int16_256x32_256x32", pto.ui8, pto.i16, (256, 32), (256, 32)), + ("i16_uint16_15x192_15x192", pto.i16, pto.ui16, (15, 192), (15, 192)), + ("ui16_int16_63x64_63x64", pto.ui16, pto.i16, (63, 64), (63, 64)), + ("f16_uint16_63x64_63x64", pto.f16, pto.ui16, (63, 64), (63, 64)), + ("bf16_int16_32x64_32x64", pto.bf16, pto.i16, (32, 64), (32, 64)), + ("i32_uint32_31x128_31x128", pto.i32, pto.ui32, (31, 128), (31, 128)), + ("ui32_int32_32x64_32x64", pto.ui32, pto.i32, (32, 64), (32, 64)), + ("f32_uint32_7x448_7x448", pto.f32, pto.ui32, (7, 448), (7, 448)), + ("f32_uint32_32x64_32x64", pto.f32, pto.ui32, (32, 64), (32, 64)), +] + +# --- Masked scatter cases (axis=0, row) --- +# Row: scatter along columns within each row. +# src(R, C_src) -> dst(R, C_dst), C_dst >= C_src +# (name, dtype, src_shape, dst_shape, pattern, axis) +MASK_CASES_ROW = [ + ("mask_row_f16_16x64_16x64_P1111", pto.f16, (16, 64), (16, 64), "P1111", "row"), + ("mask_row_f32_16x64_16x64_P1111", pto.f32, (16, 64), (16, 64), "P1111", "row"), + ("mask_row_i32_16x64_16x64_P1111", pto.i32, (16, 64), (16, 64), "P1111", "row"), + ("mask_row_i16_16x64_16x64_P1111", pto.i16, (16, 64), (16, 64), "P1111", "row"), + ("mask_row_f16_16x64_16x128_P1010", pto.f16, (16, 64), (16, 128), "P1010", "row"), + ("mask_row_f16_16x64_16x128_P0101", pto.f16, (16, 64), (16, 128), "P0101", "row"), + ("mask_row_f32_16x64_16x128_P1010", pto.f32, (16, 64), (16, 128), "P1010", "row"), + ("mask_row_f32_16x64_16x128_P0101", pto.f32, (16, 64), (16, 128), "P0101", "row"), + ("mask_row_i32_16x64_16x128_P1010", pto.i32, (16, 64), (16, 128), "P1010", "row"), + ("mask_row_i32_16x64_16x128_P0101", pto.i32, (16, 64), (16, 128), "P0101", "row"), + ("mask_row_i16_16x64_16x128_P1010", pto.i16, (16, 64), (16, 128), "P1010", "row"), + ("mask_row_i16_16x64_16x128_P0101", pto.i16, (16, 64), (16, 128), "P0101", "row"), + ("mask_row_f16_16x64_16x256_P1000", pto.f16, (16, 64), (16, 256), "P1000", "row"), + ("mask_row_f16_16x64_16x256_P0100", pto.f16, (16, 64), (16, 256), "P0100", "row"), + ("mask_row_f16_16x64_16x256_P0010", pto.f16, (16, 64), (16, 256), "P0010", "row"), + ("mask_row_f16_16x64_16x256_P0001", pto.f16, (16, 64), (16, 256), "P0001", "row"), + ("mask_row_f32_16x64_16x256_P1000", pto.f32, (16, 64), (16, 256), "P1000", "row"), + ("mask_row_f32_16x64_16x256_P0100", pto.f32, (16, 64), (16, 256), "P0100", "row"), + ("mask_row_f32_16x64_16x256_P0010", pto.f32, (16, 64), (16, 256), "P0010", "row"), + ("mask_row_f32_16x64_16x256_P0001", pto.f32, (16, 64), (16, 256), "P0001", "row"), + ("mask_row_i32_16x64_16x256_P1000", pto.i32, (16, 64), (16, 256), "P1000", "row"), + ("mask_row_i32_16x64_16x256_P0100", pto.i32, (16, 64), (16, 256), "P0100", "row"), + ("mask_row_i32_16x64_16x256_P0010", pto.i32, (16, 64), (16, 256), "P0010", "row"), + ("mask_row_i32_16x64_16x256_P0001", pto.i32, (16, 64), (16, 256), "P0001", "row"), + ("mask_row_i16_16x64_16x256_P1000", pto.i16, (16, 64), (16, 256), "P1000", "row"), + ("mask_row_i16_16x64_16x256_P0100", pto.i16, (16, 64), (16, 256), "P0100", "row"), + ("mask_row_i16_16x64_16x256_P0010", pto.i16, (16, 64), (16, 256), "P0010", "row"), + ("mask_row_i16_16x64_16x256_P0001", pto.i16, (16, 64), (16, 256), "P0001", "row"), +] + +# --- Masked scatter cases (axis=1, col) --- +# Col: scatter across rows, keeping columns intact. +# src(R_src, C) -> dst(R_dst, C), R_dst >= R_src +# (name, dtype, src_shape, dst_shape, pattern, axis) +MASK_CASES_COL = [ + ("mask_col_f16_16x64_16x64_P1111", pto.f16, (16, 64), (16, 64), "P1111", "col"), + ("mask_col_f32_16x64_16x64_P1111", pto.f32, (16, 64), (16, 64), "P1111", "col"), + ("mask_col_i32_16x64_16x64_P1111", pto.i32, (16, 64), (16, 64), "P1111", "col"), + ("mask_col_i16_16x64_16x64_P1111", pto.i16, (16, 64), (16, 64), "P1111", "col"), + ("mask_col_f16_16x64_32x64_P1010", pto.f16, (16, 64), (32, 64), "P1010", "col"), + ("mask_col_f16_16x64_32x64_P0101", pto.f16, (16, 64), (32, 64), "P0101", "col"), + ("mask_col_f32_16x64_32x64_P1010", pto.f32, (16, 64), (32, 64), "P1010", "col"), + ("mask_col_f32_16x64_32x64_P0101", pto.f32, (16, 64), (32, 64), "P0101", "col"), + ("mask_col_i32_16x64_32x64_P1010", pto.i32, (16, 64), (32, 64), "P1010", "col"), + ("mask_col_i32_16x64_32x64_P0101", pto.i32, (16, 64), (32, 64), "P0101", "col"), + ("mask_col_i16_16x64_32x64_P1010", pto.i16, (16, 64), (32, 64), "P1010", "col"), + ("mask_col_i16_16x64_32x64_P0101", pto.i16, (16, 64), (32, 64), "P0101", "col"), + ("mask_col_f16_16x64_64x64_P1000", pto.f16, (16, 64), (64, 64), "P1000", "col"), + ("mask_col_f16_16x64_64x64_P0100", pto.f16, (16, 64), (64, 64), "P0100", "col"), + ("mask_col_f16_16x64_64x64_P0010", pto.f16, (16, 64), (64, 64), "P0010", "col"), + ("mask_col_f16_16x64_64x64_P0001", pto.f16, (16, 64), (64, 64), "P0001", "col"), + ("mask_col_f32_16x64_64x64_P1000", pto.f32, (16, 64), (64, 64), "P1000", "col"), + ("mask_col_f32_16x64_64x64_P0100", pto.f32, (16, 64), (64, 64), "P0100", "col"), + ("mask_col_f32_16x64_64x64_P0010", pto.f32, (16, 64), (64, 64), "P0010", "col"), + ("mask_col_f32_16x64_64x64_P0001", pto.f32, (16, 64), (64, 64), "P0001", "col"), + ("mask_col_i32_16x64_64x64_P1000", pto.i32, (16, 64), (64, 64), "P1000", "col"), + ("mask_col_i32_16x64_64x64_P0100", pto.i32, (16, 64), (64, 64), "P0100", "col"), + ("mask_col_i32_16x64_64x64_P0010", pto.i32, (16, 64), (64, 64), "P0010", "col"), + ("mask_col_i32_16x64_64x64_P0001", pto.i32, (16, 64), (64, 64), "P0001", "col"), + ("mask_col_i16_16x64_64x64_P1000", pto.i16, (16, 64), (64, 64), "P1000", "col"), + ("mask_col_i16_16x64_64x64_P0100", pto.i16, (16, 64), (64, 64), "P0100", "col"), + ("mask_col_i16_16x64_64x64_P0010", pto.i16, (16, 64), (64, 64), "P0010", "col"), + ("mask_col_i16_16x64_64x64_P0001", pto.i16, (16, 64), (64, 64), "P0001", "col"), +] + +MASK_CASES = MASK_CASES_ROW + MASK_CASES_COL + + +# --------------------------------------------------------------------------- +# Kernel bodies +# --------------------------------------------------------------------------- + +def _tscatter_index_body(src_ptr, idx_ptr, dst_ptr, *, src_rows, src_cols, + idx_rows, idx_cols, dtype, idx_dtype): + src_view = pto.make_tensor_view(src_ptr, shape=[src_rows, src_cols], + strides=[src_cols, 1]) + idx_view = pto.make_tensor_view(idx_ptr, shape=[idx_rows, idx_cols], + strides=[idx_cols, 1]) + dst_view = pto.make_tensor_view(dst_ptr, shape=[src_rows, src_cols], + strides=[src_cols, 1]) + + src_tile = pto.alloc_tile(shape=[src_rows, src_cols], dtype=dtype) + idx_tile = pto.alloc_tile(shape=[idx_rows, idx_cols], dtype=idx_dtype) + dst_tile = pto.alloc_tile(shape=[src_rows, src_cols], dtype=dtype) + + pto.tile.load(src_view, src_tile) + pto.tile.load(idx_view, idx_tile) + pto.tile.scatter(src_tile, dst_tile, indexes=idx_tile) + pto.tile.store(dst_tile, dst_view) + + +def _tscatter_mask_body(src_ptr, dst_ptr, *, src_rows, src_cols, + dst_rows, dst_cols, dtype, pattern, axis): + src_view = pto.make_tensor_view(src_ptr, shape=[src_rows, src_cols], + strides=[src_cols, 1]) + dst_view = pto.make_tensor_view(dst_ptr, shape=[dst_rows, dst_cols], + strides=[dst_cols, 1]) + + src_tile = pto.alloc_tile(shape=[src_rows, src_cols], dtype=dtype) + dst_tile = pto.alloc_tile(shape=[dst_rows, dst_cols], dtype=dtype) + + pto.tile.load(src_view, src_tile) + pto.tile.scatter(src_tile, dst_tile, mask_pattern=pattern, + axis=axis) + pto.tile.store(dst_tile, dst_view) + + +# --------------------------------------------------------------------------- +# One @pto.jit kernel per case variant +# --------------------------------------------------------------------------- + +_index_kernels = {} +for _name, _src_dtype, _idx_dtype, _src_shape, _idx_shape in INDEX_CASES: + _sr, _sc = _src_shape + _ir, _ic = _idx_shape + + def _make_idx(sr=_sr, sc=_sc, ir=_ir, ic=_ic, src_dtype=_src_dtype, + idx_dtype=_idx_dtype, kernel_name=f"tscatter_{_name}"): + @pto.jit(name=kernel_name, target="a5") + def _kernel( + src_ptr: pto.ptr(src_dtype, "gm"), + idx_ptr: pto.ptr(idx_dtype, "gm"), + dst_ptr: pto.ptr(src_dtype, "gm"), + ): + _tscatter_index_body( + src_ptr, idx_ptr, dst_ptr, + src_rows=sr, src_cols=sc, idx_rows=ir, idx_cols=ic, + dtype=src_dtype, idx_dtype=idx_dtype, + ) + return _kernel + + _index_kernels[_name] = _make_idx() + + +_mask_kernels = {} +for _name, _dtype, _src_shape, _dst_shape, _pattern, _axis in MASK_CASES: + _sr, _sc = _src_shape + _dr, _dc = _dst_shape + + def _make_mask(sr=_sr, sc=_sc, dr=_dr, dc=_dc, dtype=_dtype, + pattern=_pattern, axis=_axis, + kernel_name=f"tscatter_{_name}"): + @pto.jit(name=kernel_name, target="a5") + def _kernel( + src_ptr: pto.ptr(dtype, "gm"), + dst_ptr: pto.ptr(dtype, "gm"), + ): + _tscatter_mask_body( + src_ptr, dst_ptr, + src_rows=sr, src_cols=sc, dst_rows=dr, dst_cols=dc, + dtype=dtype, pattern=pattern, axis=axis, + ) + return _kernel + + _mask_kernels[_name] = _make_mask() + + +# --------------------------------------------------------------------------- +# Input generators and golden functions +# --------------------------------------------------------------------------- + +def _scatter_index_golden(src, indices): + dst = np.zeros_like(src).flatten() + flat_src = src.flatten() + flat_idx = indices.flatten() + for i in range(len(flat_idx)): + dst[int(flat_idx[i])] = flat_src[i] + return dst.reshape(src.shape) + + +_ROW_PATTERN_PARAMS = { + "P0101": (0, 2), + "P1010": (1, 2), + "P0001": (0, 4), + "P0010": (1, 4), + "P0100": (2, 4), + "P1000": (3, 4), + "P1111": (0, 1), +} + +_COL_PATTERN_PARAMS = { + "P0101": (0, 2), + "P1010": (1, 2), + "P0001": (0, 4), + "P0010": (1, 4), + "P0100": (2, 4), + "P1000": (3, 4), + "P1111": (0, 1), +} + + +def _scatter_mask_row_golden(src, dst_rows, dst_cols, pattern): + offset, stride = _ROW_PATTERN_PARAMS[pattern] + dst = np.zeros((dst_rows, dst_cols), dtype=src.dtype) + src_rows, src_cols = src.shape + for i in range(src_rows): + for j in range(src_cols): + dst[i, j * stride + offset] = src[i, j] + return dst + + +def _scatter_mask_col_golden(src, dst_rows, dst_cols, pattern): + start, stride = _COL_PATTERN_PARAMS[pattern] + dst = np.zeros((dst_rows, dst_cols), dtype=src.dtype) + src_rows, src_cols = src.shape + for i in range(src_rows): + dst_row = i * stride + start + for j in range(src_cols): + dst[dst_row, j] = src[i, j] + return dst + + +def _make_index_inputs(name, src_dtype, idx_dtype, src_shape, idx_shape): + np_dt = npy_dtype(src_dtype) + np_idx_dt = npy_dtype(idx_dtype) + rng = np.random.RandomState(zlib.crc32(name.encode())) + + src = rng.uniform(0, 100, src_shape).astype(np_dt) + raw_idx = rng.randint(0, 2, idx_shape).astype(np_idx_dt) + cols = src_shape[1] + indices = np.empty_like(raw_idx) + for row in range(raw_idx.shape[0]): + for col in range(raw_idx.shape[1]): + indices[row, col] = raw_idx[row, col] * cols + col + + return [src, indices] + + +def _make_mask_inputs(name, dtype, src_shape): + np_dt = npy_dtype(dtype) + rng = np.random.RandomState(zlib.crc32(name.encode())) + src = rng.uniform(0, 100, src_shape).astype(np_dt) + return [src] + + +# --------------------------------------------------------------------------- +# Build CASES +# --------------------------------------------------------------------------- + +CASES = [] + +for _name, _src_dtype, _idx_dtype, _src_shape, _idx_shape in INDEX_CASES: + CASES.append( + golden_output_case( + "tscatter_" + _name, + _index_kernels[_name], + inputs=lambda _n=_name, _sd=_src_dtype, _id=_idx_dtype, _ss=_src_shape, _is=_idx_shape: + _make_index_inputs(_n, _sd, _id, _ss, _is), + expected=lambda src, idx: _scatter_index_golden(src, idx), + rtol=1e-6, + atol=1e-6, + ) + ) + +for _name, _dtype, _src_shape, _dst_shape, _pattern, _axis in MASK_CASES: + _golden = (_scatter_mask_row_golden if _axis == "row" + else _scatter_mask_col_golden) + CASES.append( + golden_output_case( + "tscatter_" + _name, + _mask_kernels[_name], + inputs=lambda _n=_name, _d=_dtype, _ss=_src_shape: + _make_mask_inputs(_n, _d, _ss), + expected=lambda src, _dr=_dst_shape[0], _dc=_dst_shape[1], _p=_pattern, _g=_golden: + _g(src, _dr, _dc, _p), + rtol=1e-6, + atol=1e-6, + ) + ) + + +auto_main(globals()) diff --git a/tools/ptobc/testdata/tscatter_maskpattern_v0_roundtrip.pto b/tools/ptobc/testdata/tscatter_maskpattern_v0_roundtrip.pto index 1728349063..1169890c4a 100644 --- a/tools/ptobc/testdata/tscatter_maskpattern_v0_roundtrip.pto +++ b/tools/ptobc/testdata/tscatter_maskpattern_v0_roundtrip.pto @@ -2,7 +2,7 @@ module attributes {pto.target_arch = "a2a3"} { func.func @tscatter_maskpattern_v0_roundtrip() { %src_mask = pto.alloc_tile : !pto.tile_buf %dst_mask = pto.alloc_tile : !pto.tile_buf - pto.tscatter ins(%src_mask, {maskPattern = #pto.mask_pattern} : !pto.tile_buf) outs(%dst_mask : !pto.tile_buf) + pto.tscatter ins(%src_mask, {maskPattern = #pto.mask_pattern} : !pto.tile_buf, "row") outs(%dst_mask : !pto.tile_buf) %src_idx = pto.alloc_tile : !pto.tile_buf %dst_idx = pto.alloc_tile : !pto.tile_buf