diff --git a/.gitignore b/.gitignore index 3357b79f60..bd6d22aec6 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ build/ build_plain/ build_plan/ install/ +install-optllvm/ # TileLang ST standalone build outputs (see temp_docs/standalone_st.md) test/tilelang_st/npu/a5/src/st/build/ diff --git a/include/PTO/IR/PTOAttrs.td b/include/PTO/IR/PTOAttrs.td index c156d96892..8c5658a9b5 100644 --- a/include/PTO/IR/PTOAttrs.td +++ b/include/PTO/IR/PTOAttrs.td @@ -862,6 +862,16 @@ def PTO_AccToVecModeAttr : PTO_Attr<"AccToVecMode", "acc_to_vec_mode"> { let parameters = (ins EnumParameter:$value); let assemblyFormat = "`<` params `>`"; } + +def PTO_TInsertMode_Enum : PTO_I32Enum<"TInsertMode", "TINSERT split mode (A5)", [ + I32EnumAttrCase<"SPLIT2", 2, "split2">, + I32EnumAttrCase<"SPLIT4", 3, "split4"> +]>; + +def PTO_TInsertModeAttr : PTO_Attr<"TInsertMode", "tinsert_mode"> { + let parameters = (ins EnumParameter:$value); + let assemblyFormat = "`<` params `>`"; +} //===----------------------------------------------------------------------===// // MaskPattern //===----------------------------------------------------------------------===// diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index c4351b6949..83fdbc81b7 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -3996,6 +3996,7 @@ def TExtractFPOp : PTO_TOp<"textract_fp", [ } def TInsertOp : PTO_TOp<"tinsert", [ + AttrSizedOperandSegments, PTO_DpsInitOpInterface, OpPipeInterface, DeclareOpInterfaceMethods @@ -4006,7 +4007,12 @@ def TInsertOp : PTO_TOp<"tinsert", [ PTODpsType:$src, Index:$indexRow, Index:$indexCol, - PTODpsType:$dst + PTODpsType:$dst, + Optional:$fp, + Optional:$preQuantScalar, + OptionalAttr:$accToVecMode, + DefaultValuedAttr:$reluPreMode, + OptionalAttr:$tinsertMode ); let results = (outs); @@ -4014,7 +4020,10 @@ def TInsertOp : PTO_TOp<"tinsert", [ let hasVerifier = 1; let assemblyFormat = [{ - `ins` `(` $src `,` $indexRow `,` $indexCol `:` qualified(type($src)) `,` type($indexRow) `,` type($indexCol) `)` + `ins` `(` $src `,` $indexRow `,` $indexCol `:` qualified(type($src)) `,` type($indexRow) `,` type($indexCol) + (`,` $fp^ `:` qualified(type($fp)))? + (`,` $preQuantScalar^ `:` type($preQuantScalar))? + `)` `outs` `(` $dst `:` qualified(type($dst) ) `)` attr-dict }]; @@ -4023,7 +4032,7 @@ def TInsertOp : PTO_TOp<"tinsert", [ // TINSERT runs on different pipes depending on target/path: // - Vec(UB)->Vec(UB): PIPE_V. // - A5 Vec(UB)->Mat(L1): PIPE_MTE3 (custom UB->L1 copy path). - // - Acc->Mat (A2/A3/A5 regular path): PIPE_FIX. + // - Acc->Mat / Acc->Vec (A2/A3/A5 regular path): PIPE_FIX. ::mlir::pto::PIPE getPipe() { auto isA5Target = [&]() -> bool { auto moduleOp = getOperation()->getParentOfType<::mlir::ModuleOp>(); @@ -4073,8 +4082,12 @@ def TInsertOp : PTO_TOp<"tinsert", [ sOpt.value() == ::mlir::pto::AddressSpace::VEC && dOpt.value() == ::mlir::pto::AddressSpace::MAT) return ::mlir::pto::PIPE::PIPE_MTE3; + if (sOpt.has_value() && dOpt.has_value() && + sOpt.value() == ::mlir::pto::AddressSpace::ACC && + (dOpt.value() == ::mlir::pto::AddressSpace::MAT || + dOpt.value() == ::mlir::pto::AddressSpace::VEC)) + return ::mlir::pto::PIPE::PIPE_FIX; - // Default and Acc->Mat path. return ::mlir::pto::PIPE::PIPE_FIX; } ::mlir::MutableOperandRange getDpsInitsMutable() { return getDstMutable(); } @@ -4093,7 +4106,9 @@ def TInsertFPOp : PTO_TOp<"tinsert_fp", [ PTODpsType:$fp, Index:$indexRow, Index:$indexCol, - PTODpsType:$dst + PTODpsType:$dst, + OptionalAttr:$accToVecMode, + DefaultValuedAttr:$reluPreMode ); let results = (outs); diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 6952a7e288..8516e63be9 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -5196,6 +5196,7 @@ mlir::LogicalResult mlir::pto::TExtractOp::verify() { }; return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } +static bool isA5VectorPreQuantTypePair(Type srcElem, Type dstElem); mlir::LogicalResult mlir::pto::TInsertOp::verify() { auto isColMajorRowMajorNZ = [&](pto::TileBufType ty) -> bool { return ty.getBLayoutValueI32() != static_cast(pto::BLayout::RowMajor) && @@ -5215,6 +5216,9 @@ mlir::LogicalResult mlir::pto::TInsertOp::verify() { auto isA2A3VecInsertElemType = [&](Type ty) -> bool { return ty.isInteger(8) || ty.isF16() || ty.isBF16() || ty.isF32(); }; + auto getSpace = [](Type ty) -> std::optional { + return getPTOMemorySpaceEnum(ty); + }; auto verifyCommon = [&]() -> FailureOr, @@ -5241,12 +5245,79 @@ mlir::LogicalResult mlir::pto::TInsertOp::verify() { return std::make_tuple(srcTy, dstTy, srcTb, dstTb, srcElem, dstElem, srcSpace, dstSpace); }; + // Shared validation for optional operands and attributes. + auto verifyOptionalArgs = [&](const std::optional &srcSpace, + const std::optional &dstSpace, + bool isA5) -> LogicalResult { + const bool hasFp = static_cast(getFp()); + const bool hasPreQuantScalar = static_cast(getPreQuantScalar()); + const bool hasAccToVecMode = static_cast(getAccToVecModeAttr()); + const bool hasInsertMode = static_cast(getTinsertModeAttr()); + const bool reluNonDefault = getReluPreMode() != pto::ReluPreMode::NoRelu; + + if (hasFp && hasPreQuantScalar) + return emitOpError("fp and preQuantScalar are mutually exclusive"); + + // fp tile is only valid with Acc source. + if (hasFp) { + if (!srcSpace || *srcSpace != pto::AddressSpace::ACC) + return emitOpError("fp is only valid with src loc=acc"); + auto fpTy = getFp().getType(); + auto fpTb = dyn_cast(fpTy); + if (!fpTb) return emitOpError("expects fp to be !pto.tile_buf"); + auto fpSpace = getSpace(fpTy); + if (!fpSpace || *fpSpace != pto::AddressSpace::SCALING) + return emitOpError("expects fp to be loc=scaling"); + } + + // preQuantScalar is only valid with Acc source. + if (hasPreQuantScalar) { + if (!srcSpace || *srcSpace != pto::AddressSpace::ACC) + return emitOpError("preQuantScalar is only valid with src loc=acc"); + } + + // reluPreMode is only valid with Acc source. + if (reluNonDefault) { + if (!srcSpace || *srcSpace != pto::AddressSpace::ACC) + return emitOpError("reluPreMode is only valid with src loc=acc"); + } + + // accToVecMode is only valid with Acc->Vec (A5 only). + if (hasAccToVecMode) { + if (!isA5) + return emitOpError("accToVecMode is only supported on A5"); + if (!srcSpace || !dstSpace || + *srcSpace != pto::AddressSpace::ACC || + *dstSpace != pto::AddressSpace::VEC) + return emitOpError("accToVecMode is only valid with src=acc, dst=vec"); + } + + // tinsertMode (SPLIT2/SPLIT4) is only valid with Vec(NZ)->Mat on A5. + if (hasInsertMode) { + if (!isA5) + return emitOpError("tinsertMode is only supported on A5"); + if (!srcSpace || !dstSpace || + *srcSpace != pto::AddressSpace::VEC || + *dstSpace != pto::AddressSpace::MAT) + return emitOpError( + "tinsertMode (SPLIT2/SPLIT4) is only valid with src=vec, dst=mat"); + auto srcTb = dyn_cast(getSrc().getType()); + if (!srcTb || !isColMajorRowMajorNZ(srcTb)) + return emitOpError( + "tinsertMode (SPLIT2/SPLIT4) requires src NZ layout " + "(blayout=col_major, slayout=row_major)"); + } + + return success(); + }; auto verifyA2A3 = [&]() -> LogicalResult { auto common = verifyCommon(); if (failed(common)) return failure(); auto [srcTy, dstTy, srcTb, dstTb, srcElem, dstElem, srcSpace, dstSpace] = *common; + if (failed(verifyOptionalArgs(srcSpace, dstSpace, /*isA5=*/false))) + return failure(); if (srcSpace && dstSpace && *srcSpace == pto::AddressSpace::VEC && *dstSpace == pto::AddressSpace::VEC) { if (srcElem != dstElem || !isA2A3VecInsertElemType(srcElem)) @@ -5278,20 +5349,58 @@ mlir::LogicalResult mlir::pto::TInsertOp::verify() { *common; if (!srcSpace || !dstSpace) return emitOpError("expects A5 tinsert src/dst to have explicit loc"); + if (failed(verifyOptionalArgs(srcSpace, dstSpace, /*isA5=*/true))) + return failure(); - // A5 regular acc->mat path. + // A5 acc->mat path. if (*srcSpace == pto::AddressSpace::ACC && *dstSpace == pto::AddressSpace::MAT) { if (!isColMajorRowMajorNZ(srcTb)) return emitOpError("expects A5 acc->mat tinsert src to use blayout=col_major and slayout=row_major"); if (!isColMajorRowMajorNZ(dstTb)) return emitOpError("expects A5 acc->mat tinsert dst to use blayout=col_major and slayout=row_major"); - bool okTypes = (srcElem.isF32() && - (dstElem.isF16() || dstElem.isBF16() || dstElem.isF32())) || - (srcElem.isInteger(32) && dstElem.isInteger(32)); + const bool hasQuant = static_cast(getFp()) || + static_cast(getPreQuantScalar()); + bool okTypes; + if (hasQuant) { + // With fp/scalar quantization, allow wider dst types (i8/fp8/f16/bf16/f32). + okTypes = isA5VectorPreQuantTypePair(srcElem, dstElem); + } else { + okTypes = (srcElem.isF32() && + (dstElem.isF16() || dstElem.isBF16() || dstElem.isF32())) || + (srcElem.isInteger(32) && dstElem.isInteger(32)); + } if (!okTypes) return emitOpError( "expects A5 acc->mat tinsert element types to be " - "(src=f32,dst=f16/bf16/f32) or (src=i32,dst=i32)"); + "(src=f32,dst=f16/bf16/f32) or (src=i32,dst=i32)" + + (hasQuant ? std::string("; with fp/scalar: (src=f32,dst=i8/fp8/f16/bf16/f32) or (src=i32,dst=i8/f16/bf16)") : std::string())); + return success(); + } + + // A5 acc->vec path. + if (*srcSpace == pto::AddressSpace::ACC && *dstSpace == pto::AddressSpace::VEC) { + if (!isColMajorRowMajorNZ(srcTb)) + return emitOpError("expects A5 acc->vec tinsert src to use blayout=col_major and slayout=row_major"); + bool srcIsND = isRowMajorNoneBoxND(dstTb); + bool srcIsNZ = isColMajorRowMajorNZ(dstTb); + if (!srcIsND && !srcIsNZ) + return emitOpError( + "expects A5 acc->vec tinsert dst to use ND(row_major/none_box) or NZ(col_major/row_major) layout"); + const bool hasQuant = static_cast(getFp()) || + static_cast(getPreQuantScalar()); + bool okTypes; + if (hasQuant) { + okTypes = isA5VectorPreQuantTypePair(srcElem, dstElem); + } else { + okTypes = (srcElem.isF32() && + (dstElem.isF16() || dstElem.isBF16() || dstElem.isF32())) || + (srcElem.isInteger(32) && dstElem.isInteger(32)); + } + if (!okTypes) + return emitOpError( + "expects A5 acc->vec tinsert element types to be " + "(src=f32,dst=f16/bf16/f32) or (src=i32,dst=i32)" + + (hasQuant ? std::string("; with fp/scalar: (src=f32,dst=i8/fp8/f16/bf16/f32) or (src=i32,dst=i8/f16/bf16)") : std::string())); return success(); } @@ -5311,12 +5420,21 @@ mlir::LogicalResult mlir::pto::TInsertOp::verify() { return success(); } - // A5 vec->vec path (PR561 ND_VEC). + // A5 vec->vec path: supports ND->ND and NZ->NZ. if (*srcSpace == pto::AddressSpace::VEC && *dstSpace == pto::AddressSpace::VEC) { - if (!isRowMajorNoneBoxND(srcTb) || !isRowMajorNoneBoxND(dstTb)) + bool srcIsND = isRowMajorNoneBoxND(srcTb); + bool dstIsND = isRowMajorNoneBoxND(dstTb); + bool srcIsNZ = isColMajorRowMajorNZ(srcTb); + bool dstIsNZ = isColMajorRowMajorNZ(dstTb); + if (srcIsND && dstIsND) { + // ND->ND path + } else if (srcIsNZ && dstIsNZ) { + // NZ->NZ path + } else { return emitOpError( - "expects A5 vec->vec tinsert src/dst to use ND layout " - "(blayout=row_major, slayout=none_box)"); + "expects A5 vec->vec tinsert src/dst layouts to match: " + "both ND(row_major/none_box) or both NZ(col_major/row_major)"); + } if (srcElem != dstElem || !isA5SupportedVecElemType(srcElem)) return emitOpError( "expects A5 vec->vec tinsert src/dst to have same supported dtype " @@ -5326,7 +5444,7 @@ mlir::LogicalResult mlir::pto::TInsertOp::verify() { return emitOpError( "expects A5 tinsert to use a supported src/dst loc pair: " - "acc->mat, vec->mat, or vec->vec"); + "acc->mat, acc->vec, vec->mat, or vec->vec"); }; return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } @@ -5468,7 +5586,7 @@ mlir::LogicalResult mlir::pto::TExtractFPOp::verify() { } mlir::LogicalResult mlir::pto::TInsertFPOp::verify() { - auto verifyCommon = [&]() -> FailureOr FailureOr> { @@ -5499,17 +5617,24 @@ mlir::LogicalResult mlir::pto::TInsertFPOp::verify() { return emitOpError("expects src to use loc=acc"); if (*fpSpace != pto::AddressSpace::SCALING) return emitOpError("expects fp to use loc=scaling"); - if (*dstSpace != pto::AddressSpace::MAT) - return emitOpError("expects dst to use loc=mat"); + // A2/A3: only acc->mat; A5: acc->mat or acc->vec. + if (*dstSpace != pto::AddressSpace::MAT && + !(isA5 && *dstSpace == pto::AddressSpace::VEC)) + return emitOpError("expects dst to use loc=mat" + + (isA5 ? StringRef(" or loc=vec (A5)") : StringRef(""))); if (!isColMajorRowMajorNZTileBuf(srcTb)) return emitOpError("expects src to use blayout=col_major and slayout=row_major"); - if (!isColMajorRowMajorNZTileBuf(dstTb)) - return emitOpError("expects dst to use blayout=col_major and slayout=row_major"); + if (*dstSpace == pto::AddressSpace::MAT && !isColMajorRowMajorNZTileBuf(dstTb)) + return emitOpError("expects dst (mat) to use blayout=col_major and slayout=row_major"); + // accToVecMode is only valid when dst=vec. + if (static_cast(getAccToVecModeAttr()) && + *dstSpace != pto::AddressSpace::VEC) + return emitOpError("accToVecMode is only valid with dst=vec"); return std::make_tuple(srcTy, fpTy, dstTy, srcTb, fpTb, dstTb, *srcSpace, *fpSpace, *dstSpace); }; auto verifyA2A3 = [&]() -> LogicalResult { - auto common = verifyCommon(); + auto common = verifyCommon(/*isA5=*/false); if (failed(common)) return failure(); auto [srcTy, fpTy, dstTy, srcTb, fpTb, dstTb, srcSpace, fpSpace, dstSpace] = @@ -5531,7 +5656,7 @@ mlir::LogicalResult mlir::pto::TInsertFPOp::verify() { return success(); }; auto verifyA5 = [&]() -> LogicalResult { - auto common = verifyCommon(); + auto common = verifyCommon(/*isA5=*/true); if (failed(common)) return failure(); auto [srcTy, fpTy, dstTy, srcTb, fpTb, dstTb, srcSpace, fpSpace, dstSpace] = @@ -10964,8 +11089,11 @@ void TExtractOp::getEffects( // TINSERT: Read(src) -> Write(dst) void TInsertOp::getEffects( SmallVectorImpl> &effects) { - PTO_ADD_READ(getSrcMutable()); - PTO_ADD_WRITE(getDstMutable()); + addEffect(effects, &getSrcMutable(), MemoryEffects::Read::get()); + auto fpRange = getFpMutable(); + if (!fpRange.empty()) + addEffect(effects, &*fpRange.begin(), MemoryEffects::Read::get()); + addEffect(effects, &getDstMutable(), MemoryEffects::Write::get()); } // TEXTRACT_FP: Read(src), Read(fp) -> Write(dst) diff --git a/lib/PTO/IR/VPTO.cpp b/lib/PTO/IR/VPTO.cpp index afe98d76bb..0e526031bb 100644 --- a/lib/PTO/IR/VPTO.cpp +++ b/lib/PTO/IR/VPTO.cpp @@ -2627,10 +2627,10 @@ static ParseResult parseStructuredAccStoreClauses( return success(); } StringRef keyword; - if (parser.parseKeyword(&keyword)) { + if (parser.parseOptionalKeyword(&keyword)) { if (!seenClause) return success(); - return failure(); + return parser.emitError(parser.getCurrentLocation(), "expected valid keyword"); } seenClause = true; diff --git a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp index 398b13df8d..02bfad31a0 100644 --- a/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp +++ b/lib/PTO/Transforms/FoldTileBufIntrinsics.cpp @@ -76,6 +76,12 @@ struct TileHandleInfo { static std::optional resolveTileHandle(Value tileBuf, Operation *user) { + while (auto cast = tileBuf.getDefiningOp()) { + if (cast.getNumOperands() != 1 || cast.getNumResults() != 1) + break; + tileBuf = cast.getOperand(0); + } + if (auto alloc = tileBuf.getDefiningOp()) { auto tileTy = dyn_cast(alloc.getResult().getType()); if (!tileTy) { @@ -580,6 +586,20 @@ struct FoldTileBufIntrinsicsPass for (auto castOp : llvm::reverse(deadCasts)) castOp.erase(); + // Clean up dead unrealized_conversion_cast ops that bridged between + // two tile_buf types with different configs (e.g. after + // PTOMaterializeTileHandles creates a new alloc_tile with default + // config that is bridged to the daemon-generated template type). + SmallVector deadTileBufCasts; + func.walk([&](UnrealizedConversionCastOp castOp) { + if (castOp.use_empty() && castOp.getNumOperands() == 1 && + isa(castOp.getOperand(0).getType()) && + isa(castOp.getResult(0).getType())) + deadTileBufCasts.push_back(castOp); + }); + for (auto castOp : llvm::reverse(deadTileBufCasts)) + castOp.erase(); + while (true) { SmallVector deadMemrefOps; func.walk([&](Operation *op) { diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index c664712446..3e058f3be0 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -7790,8 +7790,9 @@ struct PTOExtractFPToEmitC : public OpConversionPattern { } }; //===----------------------------------------------------------------------===// -// pto.tinsert lowering -> TINSERT(dst, src, indexRow, indexCol) -// Keep lowering arch-agnostic and let PTO-ISA infer proper A5 path. +// pto.tinsert lowering -> TINSERT(dst, src, [fp|scalar,] indexRow, indexCol) +// Template args are emitted when optional parameters are present to match +// the corresponding pto-isa TINSERT_IMPL overloads. //===----------------------------------------------------------------------===// struct PTOInsertToEmitC : public OpConversionPattern { @@ -7800,16 +7801,141 @@ struct PTOInsertToEmitC : public OpConversionPattern { LogicalResult matchAndRewrite(pto::TInsertOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { auto loc = op.getLoc(); + auto *ctx = rewriter.getContext(); Value src = peelUnrealized(adaptor.getSrc()); Value dst = peelUnrealized(adaptor.getDst()); Value r0 = peelUnrealized(adaptor.getIndexRow()); Value c0 = peelUnrealized(adaptor.getIndexCol()); + Value fp; + if (op.getFp()) + fp = peelUnrealized(adaptor.getFp()); + Value preQuantScalar; + if (op.getPreQuantScalar()) + preQuantScalar = peelUnrealized(adaptor.getPreQuantScalar()); + + auto dstOT = mlir::dyn_cast(dst.getType()); + auto srcOT = mlir::dyn_cast(src.getType()); + + auto reluTok = [&](pto::ReluPreMode mode) -> StringRef { + switch (mode) { + case pto::ReluPreMode::NoRelu: return "ReluPreMode::NoRelu"; + case pto::ReluPreMode::NormalRelu: return "ReluPreMode::NormalRelu"; + case pto::ReluPreMode::ScalarRelu: return "ReluPreMode::ScalarRelu"; + case pto::ReluPreMode::VectorRelu: return "ReluPreMode::VectorRelu"; + case pto::ReluPreMode::Pwl: return "ReluPreMode::Pwl"; + } + llvm_unreachable("unknown ReluPreMode"); + }; + + auto modeTok = [&](pto::AccToVecMode mode) -> StringRef { + switch (mode) { + case pto::AccToVecMode::SingleModeVec0: return "pto::AccToVecMode::SingleModeVec0"; + case pto::AccToVecMode::SingleModeVec1: return "pto::AccToVecMode::SingleModeVec1"; + case pto::AccToVecMode::DualModeSplitM: return "pto::AccToVecMode::DualModeSplitM"; + case pto::AccToVecMode::DualModeSplitN: return "pto::AccToVecMode::DualModeSplitN"; + } + llvm_unreachable("unknown AccToVecMode"); + }; + + auto splitTok = [&](pto::TInsertMode mode) -> StringRef { + switch (mode) { + case pto::TInsertMode::SPLIT2: return "pto::TInsertMode::SPLIT2"; + case pto::TInsertMode::SPLIT4: return "pto::TInsertMode::SPLIT4"; + } + llvm_unreachable("unknown TInsertMode"); + }; + + const bool hasFp = static_cast(fp); + const bool hasPreQuantScalar = static_cast(preQuantScalar); + const bool hasMode = static_cast(op.getAccToVecModeAttr()); + const bool reluNonDefault = op.getReluPreMode() != pto::ReluPreMode::NoRelu; + const bool hasSplitMode = static_cast(op.getTinsertModeAttr()); + + SmallVector operands{dst, src}; + + // Split mode: TINSERT(dst, src, r0, c0) + if (hasSplitMode) { + SmallVector templateArgVec; + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, splitTok(op.getTinsertModeAttr().getValue()))); + if (dstOT && srcOT) { + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, dstOT.getValue().str())); + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, srcOT.getValue().str())); + } + operands.push_back(r0); + operands.push_back(c0); + rewriter.create( + loc, TypeRange{}, "TINSERT", + /*args=*/ArrayAttr{}, + /*templateArgs=*/rewriter.getArrayAttr(templateArgVec), + /*operands=*/operands); + rewriter.eraseOp(op); + return success(); + } + + // Build template args vector (type-templated path). + SmallVector templateArgVec; + if (dstOT && srcOT) { + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, dstOT.getValue().str())); + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, srcOT.getValue().str())); + } + + if (hasFp) { + // Vector quantization (fp tile). + auto fpOT = mlir::dyn_cast(fp.getType()); + if (fpOT) + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, fpOT.getValue().str())); + if (hasMode) + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, modeTok(op.getAccToVecModeAttr().getValue()))); + if (hasMode || reluNonDefault) + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, reluTok(op.getReluPreMode()))); + operands.push_back(fp); + } else if (hasPreQuantScalar) { + // Scalar quantization. + if (hasMode) + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, modeTok(op.getAccToVecModeAttr().getValue()))); + if (hasMode || reluNonDefault) + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, reluTok(op.getReluPreMode()))); + operands.push_back(preQuantScalar); + } else if (hasMode) { + // AccToVecMode without fp/preQuantScalar. + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, modeTok(op.getAccToVecModeAttr().getValue()))); + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, reluTok(op.getReluPreMode()))); + } else if (reluNonDefault) { + // Only ReLU mode. + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, reluTok(op.getReluPreMode()))); + } + + operands.push_back(r0); + operands.push_back(c0); + + // Determine callee: use TINSERT_FP for the dedicated fp-tile op (kept + // separately), otherwise always TINSERT. + StringRef callee = "TINSERT"; + + ArrayAttr templateArgs = + templateArgVec.size() == 2 && !hasFp && !hasPreQuantScalar && + !hasMode && !reluNonDefault + ? ArrayAttr{} + : rewriter.getArrayAttr(templateArgVec); rewriter.create( - loc, TypeRange{}, "TINSERT", - /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, - /*operands=*/ValueRange{dst, src, r0, c0}); + loc, TypeRange{}, callee, + /*args=*/ArrayAttr{}, /*templateArgs=*/templateArgs, + /*operands=*/operands); rewriter.eraseOp(op); return success(); @@ -7817,6 +7943,7 @@ struct PTOInsertToEmitC : public OpConversionPattern { }; //===----------------------------------------------------------------------===// // pto.tinsert_fp lowering -> TINSERT_FP(dst, src, fp, indexRow, indexCol) +// Emits template args when accToVecMode / reluPreMode are present. //===----------------------------------------------------------------------===// struct PTOInsertFPToEmitC : public OpConversionPattern { @@ -7825,17 +7952,66 @@ struct PTOInsertFPToEmitC : public OpConversionPattern { LogicalResult matchAndRewrite(pto::TInsertFPOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { auto loc = op.getLoc(); + auto *ctx = rewriter.getContext(); Value src = peelUnrealized(adaptor.getSrc()); - Value fp = peelUnrealized(adaptor.getFp()); + Value fp = peelUnrealized(adaptor.getFp()); Value dst = peelUnrealized(adaptor.getDst()); - Value r0 = peelUnrealized(adaptor.getIndexRow()); - Value c0 = peelUnrealized(adaptor.getIndexCol()); + Value r0 = peelUnrealized(adaptor.getIndexRow()); + Value c0 = peelUnrealized(adaptor.getIndexCol()); - rewriter.create( - loc, TypeRange{}, "TINSERT_FP", - /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, - /*operands=*/ValueRange{dst, src, fp, r0, c0}); + auto reluTok = [&](pto::ReluPreMode mode) -> StringRef { + switch (mode) { + case pto::ReluPreMode::NoRelu: return "ReluPreMode::NoRelu"; + case pto::ReluPreMode::NormalRelu: return "ReluPreMode::NormalRelu"; + case pto::ReluPreMode::ScalarRelu: return "ReluPreMode::ScalarRelu"; + case pto::ReluPreMode::VectorRelu: return "ReluPreMode::VectorRelu"; + case pto::ReluPreMode::Pwl: return "ReluPreMode::Pwl"; + } + llvm_unreachable("unknown ReluPreMode"); + }; + + auto modeTok = [&](pto::AccToVecMode mode) -> StringRef { + switch (mode) { + case pto::AccToVecMode::SingleModeVec0: return "pto::AccToVecMode::SingleModeVec0"; + case pto::AccToVecMode::SingleModeVec1: return "pto::AccToVecMode::SingleModeVec1"; + case pto::AccToVecMode::DualModeSplitM: return "pto::AccToVecMode::DualModeSplitM"; + case pto::AccToVecMode::DualModeSplitN: return "pto::AccToVecMode::DualModeSplitN"; + } + llvm_unreachable("unknown AccToVecMode"); + }; + + const bool hasMode = static_cast(op.getAccToVecModeAttr()); + const bool reluNonDefault = op.getReluPreMode() != pto::ReluPreMode::NoRelu; + + SmallVector operands{dst, src, fp, r0, c0}; + + auto dstOT = mlir::dyn_cast(dst.getType()); + auto srcOT = mlir::dyn_cast(src.getType()); + auto fpOT = mlir::dyn_cast(fp.getType()); + + if (!hasMode && !reluNonDefault) { + rewriter.create( + loc, TypeRange{}, "TINSERT_FP", + /*args=*/ArrayAttr{}, /*templateArgs=*/ArrayAttr{}, + operands); + } else { + // Path with accToVecMode / reluPreMode: emit TINSERT with full template args. + SmallVector templateArgVec; + if (dstOT) templateArgVec.push_back(emitc::OpaqueAttr::get(ctx, dstOT.getValue().str())); + if (srcOT) templateArgVec.push_back(emitc::OpaqueAttr::get(ctx, srcOT.getValue().str())); + if (fpOT) templateArgVec.push_back(emitc::OpaqueAttr::get(ctx, fpOT.getValue().str())); + if (hasMode) + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, modeTok(op.getAccToVecModeAttr().getValue()))); + if (hasMode || reluNonDefault) + templateArgVec.push_back( + emitc::OpaqueAttr::get(ctx, reluTok(op.getReluPreMode()))); + rewriter.create( + loc, TypeRange{}, "TINSERT", + /*args=*/ArrayAttr{}, rewriter.getArrayAttr(templateArgVec), + operands); + } rewriter.eraseOp(op); return success(); diff --git a/lib/PTO/Transforms/PTOViewToMemref.cpp b/lib/PTO/Transforms/PTOViewToMemref.cpp index c62d3340e6..e3d19e7aa8 100644 --- a/lib/PTO/Transforms/PTOViewToMemref.cpp +++ b/lib/PTO/Transforms/PTOViewToMemref.cpp @@ -1969,6 +1969,19 @@ struct PTOViewToMemrefPass op.getReluPreModeAttr()); } + // --- TInsertOp --- + SmallVector tinsertOps; + func.walk([&](mlir::pto::TInsertOp op) { tinsertOps.push_back(op); }); + for (auto op : tinsertOps) { + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(op); + rewriter.replaceOpWithNewOp( + op, TypeRange{}, op.getSrc(), op.getIndexRow(), op.getIndexCol(), + op.getDst(), op.getFp(), op.getPreQuantScalar(), + op.getAccToVecModeAttr(), op.getReluPreModeAttr(), + op.getTinsertModeAttr()); + } + SmallVector abseops; func.walk([&](mlir::pto::TAbsOp op) { abseops.push_back(op); }); @@ -3125,6 +3138,18 @@ struct PTOViewToMemrefPass dst); } + // --- TInsertFPOp --- + SmallVector tinsertFpOps; + func.walk([&](mlir::pto::TInsertFPOp op) { tinsertFpOps.push_back(op); }); + for (auto op : tinsertFpOps) { + IRRewriter rewriter(ctx); + rewriter.setInsertionPoint(op); + rewriter.replaceOpWithNewOp( + op, TypeRange{}, op.getSrc(), op.getFp(), + op.getIndexRow(), op.getIndexCol(), op.getDst(), + op.getAccToVecModeAttr(), op.getReluPreModeAttr()); + } + SmallVector quantops; func.walk([&](mlir::pto::TQuantOp op) { quantops.push_back(op); }); diff --git a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp index 8291521580..f67bd5fbb1 100644 --- a/lib/PTO/Transforms/VPTOLLVMEmitter.cpp +++ b/lib/PTO/Transforms/VPTOLLVMEmitter.cpp @@ -25,6 +25,7 @@ #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Func/Transforms/FuncConversions.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/SCF/Transforms/Patterns.h" #include "mlir/IR/Builders.h" #include "mlir/IR/BuiltinOps.h" @@ -8708,6 +8709,39 @@ class ConvertPtoCastPtrOp final : public OpConversionPattern { Value input = adaptor.getInput(); Type inputType = input.getType(); + + if (isa(op.getInput().getType())) { + Value replacement = input; + Type replacementType = inputType; + if (auto cast = + op.getInput().getDefiningOp()) { + if (cast->getNumOperands() == 1 && cast->getNumResults() == 1) { + Value original = cast.getOperand(0); + Type convertedOriginal = + getTypeConverter()->convertType(original.getType()); + if (convertedOriginal) { + replacement = original; + replacementType = convertedOriginal; + } + } + } + if (replacementType == convertedResultType) { + rewriter.replaceOp(op, replacement); + return success(); + } + auto replacementPtr = dyn_cast(replacementType); + auto targetPtr = dyn_cast(convertedResultType); + if (replacementPtr && targetPtr) { + if (replacementPtr.getAddressSpace() == targetPtr.getAddressSpace()) + rewriter.replaceOpWithNewOp(op, targetPtr, + replacement); + else + rewriter.replaceOpWithNewOp(op, targetPtr, + replacement); + return success(); + } + } + if (inputType == convertedResultType) { rewriter.replaceOp(op, input); return success(); @@ -8718,16 +8752,17 @@ class ConvertPtoCastPtrOp final : public OpConversionPattern { rewriter.replaceOpWithNewOp(op, llvmPtrType, input); return success(); } - auto sourcePtrType = dyn_cast(inputType); - if (!sourcePtrType) - return rewriter.notifyMatchFailure(op, - "expected integer or LLVM pointer input"); - if (sourcePtrType.getAddressSpace() == llvmPtrType.getAddressSpace()) { - rewriter.replaceOpWithNewOp(op, llvmPtrType, input); + if (auto sourcePtrType = dyn_cast(inputType)) { + if (sourcePtrType.getAddressSpace() == llvmPtrType.getAddressSpace()) { + rewriter.replaceOpWithNewOp(op, llvmPtrType, input); + return success(); + } + rewriter.replaceOpWithNewOp(op, llvmPtrType, + input); return success(); } - return rewriter.notifyMatchFailure( - op, "cross-address-space ptr casts are unsupported"); + return rewriter.notifyMatchFailure(op, + "expected integer or LLVM pointer input"); } if (auto resultIntType = dyn_cast(convertedResultType)) { diff --git a/lib/PTO/Transforms/VPTOPtrNormalize.cpp b/lib/PTO/Transforms/VPTOPtrNormalize.cpp index e61b4e8bb9..081e0ff03a 100644 --- a/lib/PTO/Transforms/VPTOPtrNormalize.cpp +++ b/lib/PTO/Transforms/VPTOPtrNormalize.cpp @@ -226,10 +226,17 @@ static Value materializeScalarAccessPtr(Value source, PatternRewriter &rewriter, return materializeScalarAccessPtr(addr, rewriter, loc); } - // Restrict normalization to memref views that already sit on top of a ptr-like - // boundary bridge. Materializing fresh memref -> ptr casts here would leave - // illegal pto.castptr(memref) behind in this pass. - return {}; + auto memrefType = dyn_cast(source.getType()); + if (!memrefType) + return {}; + auto memorySpace = + getPointerMemorySpace(memrefType.getMemorySpace(), rewriter.getContext()); + if (!memorySpace) + return {}; + auto ptrType = + pto::PtrType::get(rewriter.getContext(), memrefType.getElementType(), + memorySpace); + return rewriter.create(loc, ptrType, source); } static Value materializeBoundaryOperandPtr(Value source, @@ -330,9 +337,22 @@ struct ConvertCastPtrPattern : public OpConversionPattern { Value input = adaptor.getInput(); Type inputType = input.getType(); - if (isMemRefType(inputType) || isMemRefType(convertedResultType)) + if (isMemRefType(convertedResultType)) return rewriter.notifyMatchFailure(op, - "memref castptr must be eliminated"); + "memref castptr result must be eliminated"); + + if (isMemRefType(inputType)) { + auto ptrType = dyn_cast(convertedResultType); + if (!ptrType) + return rewriter.notifyMatchFailure(op, + "expected pto.ptr result for memref input"); + Value basePtrIdx = rewriter.create( + op.getLoc(), input); + Value basePtrI64 = rewriter.create( + op.getLoc(), rewriter.getI64Type(), basePtrIdx); + rewriter.replaceOpWithNewOp(op, ptrType, basePtrI64); + return success(); + } if (!isa(inputType) || !isa(convertedResultType)) diff --git a/lib/TileOps/tinsert_template.py b/lib/TileOps/tinsert_template.py new file mode 100644 index 0000000000..6ce4f714b0 --- /dev/null +++ b/lib/TileOps/tinsert_template.py @@ -0,0 +1,1459 @@ +# 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. + +"""TileLang DSL template for pto.tinsert - tile data insertion + +Implements six data movement paths using @pto.ckernel: + - Acc→Mat NZ: pto.mte_l0c_l1 (copy_matrix_cc_to_cbuf) + - Acc→Vec ND/DN/NZ: pto.mte_l0c_ub (copy_matrix_cc_to_ub) + - Vec→Vec ND/NZ: pto.copy_ubuf_to_ubuf + - Vec→Mat NZ/ND: pto.mte_ub_l1 (copy_ubuf_to_cbuf) +""" + +import tilelang_dsl as pto + +BLOCK_BYTE_SIZE = 32 +BLOCK_BYTE_BITS = 256 +FRACTAL_NZ_ROW = 16 + + +# --------------------------------------------------------------------------- +# Constraint functions +# --------------------------------------------------------------------------- + + +def _acc_to_mat_constraint(src, dst) -> bool: + return ( + src.memory_space == "acc" + and dst.memory_space == "mat" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.ROW_MAJOR + ) + + +def _acc_to_vec_nd_constraint(src, dst) -> bool: + return ( + src.memory_space == "acc" + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.ROW_MAJOR + and dst.config.s_layout == pto.SLayout.NONE_BOX + ) + + +def _acc_to_vec_dn_constraint(src, dst) -> bool: + return ( + src.memory_space == "acc" + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.NONE_BOX + ) + + +def _acc_to_vec_nz_constraint(src, dst) -> bool: + return ( + src.memory_space == "acc" + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.ROW_MAJOR + ) + + +def _vec_to_vec_nd_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.NONE_BOX + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.ROW_MAJOR + and dst.config.s_layout == pto.SLayout.NONE_BOX + and not (src.valid_shape[0] == 1 and src.valid_shape[1] == 1) + ) + + +def _vec_to_vec_nd_scalar_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.NONE_BOX + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.ROW_MAJOR + and dst.config.s_layout == pto.SLayout.NONE_BOX + and src.valid_shape[0] == 1 + and src.valid_shape[1] == 1 + ) + + +def _vec_to_vec_nz_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.COL_MAJOR + and src.config.s_layout == pto.SLayout.ROW_MAJOR + and dst.memory_space == "ub" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.ROW_MAJOR + ) + + +def _vec_to_mat_nz_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.COL_MAJOR + and src.config.s_layout == pto.SLayout.ROW_MAJOR + and dst.memory_space == "mat" + and dst.config.b_layout == pto.BLayout.COL_MAJOR + and dst.config.s_layout == pto.SLayout.ROW_MAJOR + ) + + +def _vec_to_mat_nd_constraint(src, dst) -> bool: + return ( + src.memory_space == "ub" + and src.config.b_layout == pto.BLayout.ROW_MAJOR + and src.config.s_layout == pto.SLayout.NONE_BOX + and dst.memory_space == "mat" + and dst.config.s_layout == pto.SLayout.NONE_BOX + and src.shape[1] * pto.bytewidth(pto.ScalarType(src.dtype)) >= BLOCK_BYTE_SIZE + ) + + +# --------------------------------------------------------------------------- +# Acc -> Mat +# --------------------------------------------------------------------------- + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.i8, pto.f32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), + ], + constraints=[_acc_to_mat_constraint], +) +def template_tinsert_acc_to_mat( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, + fp: pto.Tile, pre_quant_scalar: pto.i64, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + dst_elem = dst.element_type + elem_bytes = pto.bytewidth(dst_elem) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + n_size = (valid_cols + c0_size - 1) // c0_size * c0_size + + dst_rows = dst.shape[0] + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * c0_size * elem_bytes + src_stride = src.shape[0] * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + has_fp = fp is not None + has_scalar = pre_quant_scalar is not None + + if pto.constexpr(has_fp): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_relu=("normal_relu", None, None), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + else: + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_quant=(fp, "qf322f16_pre_vec"), + ) + elif pto.constexpr(has_scalar): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_relu=("normal_relu", None, None), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + else: + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + elif pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + ) + return None + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i32, pto.i32, pto.f16), + (pto.f32, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.f32, pto.i32, pto.i32, pto.i8), + (pto.i32, pto.i32, pto.i32, pto.f16), + (pto.i32, pto.i32, pto.i32, pto.bf16), + (pto.i32, pto.i32, pto.i32, pto.i8), + ], + constraints=[_acc_to_mat_constraint], +) +def template_tinsert_acc_to_mat_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + dst_elem = dst.element_type + elem_bytes = pto.bytewidth(dst_elem) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + n_size = (valid_cols + c0_size - 1) // c0_size * c0_size + + dst_rows = dst.shape[0] + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * c0_size * elem_bytes + src_stride = src.shape[0] * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_l1( + src_ptr, dst_ptr, + valid_rows, n_size, src_stride, dst_stride, + ) + return None + + +# --------------------------------------------------------------------------- +# Acc -> Vec (ND, NONE_BOX) +# --------------------------------------------------------------------------- + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i64, pto.i64, pto.f32, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i32, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), + ], + priority=10, + constraints=[_acc_to_vec_nd_constraint], +) +def template_tinsert_acc_to_vec_nd( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, + fp: pto.Tile, pre_quant_scalar: pto.i64, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + elem_bytes = pto.bytewidth(dst.element_type) + c0_size = 32 // elem_bytes + + valid_rows = src.shape[0] + valid_cols_raw = src.shape[1] + valid_cols = (valid_cols_raw + c0_size - 1) // c0_size * c0_size + + dst_cols = dst.shape[1] + dst_offset = index_row * dst_cols + index_col + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_cols * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + has_fp = fp is not None + has_scalar = pre_quant_scalar is not None + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(has_fp): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_relu=("normal_relu", None, None), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_quant=(fp, "qf322f16_pre_vec"), + ) + elif pto.constexpr(has_scalar): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_relu=("normal_relu", None, None), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + elif pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + ) + return None + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.f32, pto.i32, pto.i32, pto.f16), + (pto.f32, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.i8), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i32, pto.i32, pto.i32, pto.f16), + (pto.i32, pto.i32, pto.i32, pto.bf16), + (pto.i32, pto.i32, pto.i32, pto.i8), + ], + priority=10, + constraints=[_acc_to_vec_nd_constraint], +) +def template_tinsert_acc_to_vec_nd_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + elem_bytes = pto.bytewidth(dst.element_type) + c0_size = 32 // elem_bytes + + valid_rows = src.shape[0] + valid_cols_raw = src.shape[1] + valid_cols = (valid_cols_raw + c0_size - 1) // c0_size * c0_size + + dst_cols = dst.shape[1] + dst_offset = index_row * dst_cols + index_col + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_cols * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout="nz2nd", + ) + return None + + +# --------------------------------------------------------------------------- +# Acc -> Vec (DN, NONE_BOX) +# --------------------------------------------------------------------------- + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i64, pto.i64, pto.f32, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i32, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), + ], + priority=10, + constraints=[_acc_to_vec_dn_constraint], +) +def template_tinsert_acc_to_vec_dn( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, + fp: pto.Tile, pre_quant_scalar: pto.i64, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + elem_bytes = pto.bytewidth(dst.element_type) + c0_size = 32 // elem_bytes + + valid_rows_raw = src.shape[0] + valid_rows = (valid_rows_raw + c0_size - 1) // c0_size * c0_size + valid_cols = src.shape[1] + + dst_rows = dst.shape[0] + dst_offset = index_col * dst_rows + index_row + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + has_fp = fp is not None + has_scalar = pre_quant_scalar is not None + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(has_fp): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_relu=("normal_relu", None, None), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + elif pto.constexpr(has_scalar): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_relu=("normal_relu", None, None), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + elif pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + ) + return None + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.f32, pto.i32, pto.i32, pto.f16), + (pto.f32, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.i32), + (pto.f32, pto.i32, pto.i32, pto.i8), + (pto.i32, pto.i32, pto.i32, pto.f32), + (pto.i32, pto.i32, pto.i32, pto.f16), + (pto.i32, pto.i32, pto.i32, pto.bf16), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i32, pto.i32, pto.i32, pto.i8), + ], + priority=10, + constraints=[_acc_to_vec_dn_constraint], +) +def template_tinsert_acc_to_vec_dn_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + elem_bytes = pto.bytewidth(dst.element_type) + c0_size = 32 // elem_bytes + + valid_rows_raw = src.shape[0] + valid_rows = (valid_rows_raw + c0_size - 1) // c0_size * c0_size + valid_cols = src.shape[1] + + dst_rows = dst.shape[0] + dst_offset = index_col * dst_rows + index_row + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2dn", pto.i64(0)), + ) + return None + + +# --------------------------------------------------------------------------- +# Acc -> Vec (NZ, ROW_MAJOR) +# --------------------------------------------------------------------------- + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i64, pto.i64, pto.f32, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.f16, pto.f32, pto.i64), + (pto.f32, pto.i64, pto.i64, pto.bf16, pto.f32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i32, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.f16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.bf16, pto.i32, pto.i64), + (pto.i32, pto.i64, pto.i64, pto.i8, pto.i32, pto.i64), + ], + priority=10, + constraints=[_acc_to_vec_nz_constraint], +) +def template_tinsert_acc_to_vec_nz( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, + fp: pto.Tile, pre_quant_scalar: pto.i64, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + dst_elem = dst.element_type + elem_bytes = pto.bytewidth(dst_elem) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols_raw = src.shape[1] + + if pto.constexpr(dst_elem == pto.f32): + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == BLOCK_BYTE_BITS): + valid_cols_align = c0_size + else: + valid_cols_align = FRACTAL_NZ_ROW + else: + valid_cols_align = c0_size + valid_cols = (valid_cols_raw + valid_cols_align - 1) // valid_cols_align * valid_cols_align + + dst_rows = dst.shape[0] + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * c0_size * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + has_fp = fp is not None + has_scalar = pre_quant_scalar is not None + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(has_fp): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_relu=("normal_relu", None, None), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_quant=(fp, "qf322f16_pre_vec"), + ) + elif pto.constexpr(has_scalar): + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_relu=("normal_relu", None, None), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_quant=(pre_quant_scalar, "qf322f16_pre_scalar"), + ) + elif pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + ) + return None + + +@pto.ckernel( + target="a5", + op="pto.tinsert", + dtypes=[ + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.f32, pto.i32, pto.i32, pto.f16), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i32, pto.i32, pto.i32, pto.f16), + ], + priority=10, + constraints=[_acc_to_vec_nz_constraint], +) +def template_tinsert_acc_to_vec_nz_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + + dst_elem = dst.element_type + elem_bytes = pto.bytewidth(dst_elem) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols_raw = src.shape[1] + + if pto.constexpr(dst_elem == pto.f32): + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == BLOCK_BYTE_BITS): + valid_cols_align = c0_size + else: + valid_cols_align = FRACTAL_NZ_ROW + else: + valid_cols_align = c0_size + valid_cols = (valid_cols_raw + valid_cols_align - 1) // valid_cols_align * valid_cols_align + + dst_rows = dst.shape[0] + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst_ptr, dst_offset) + + dst_stride = dst_rows * c0_size * elem_bytes + src_stride = (valid_rows + 15) // 16 * 16 * elem_bytes + + relu_mode_name = pto.get_op_attr("relu_pre_mode", "no_relu") + acc_mode_name = pto.get_op_attr("acc_to_vec_mode", "single_mode_vec0") + + dst_mode = 0 + if pto.constexpr(acc_mode_name == "single_mode_vec1"): + dst_mode = 1 + elif pto.constexpr(acc_mode_name == "dual_mode_split_m"): + dst_mode = "split_m" + elif pto.constexpr(acc_mode_name == "dual_mode_split_n"): + dst_mode = "split_n" + + if pto.constexpr(relu_mode_name == "normal_relu"): + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + pre_relu=("normal_relu", None, None), + ) + else: + pto.mte_l0c_ub( + src_ptr, dst_ptr, + valid_rows, valid_cols, src_stride, dst_stride, dst_mode, + layout=("nz2nz", pto.i64(0)), + ) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Vec (ND, ROW_MAJOR + NONE_BOX) - O1 +# --------------------------------------------------------------------------- + + +_VEC_TO_VEC_DTYPES = [ + (pto.f16, pto.i64, pto.i64, pto.f16), + (pto.bf16, pto.i64, pto.i64, pto.bf16), + (pto.f32, pto.i64, pto.i64, pto.f32), + (pto.i32, pto.i64, pto.i64, pto.i32), + (pto.i8, pto.i64, pto.i64, pto.i8), +] + +_VEC_TO_VEC_BASIC_DTYPES = [ + (pto.f16, pto.i32, pto.i32, pto.f16), + (pto.bf16, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i8, pto.i32, pto.i32, pto.i8), +] + +_VEC_TO_MAT_SPLIT_DTYPES = [ + (pto.f16, pto.i32, pto.i32, pto.f16), + (pto.bf16, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i8, pto.i32, pto.i32, pto.i8), +] + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_DTYPES, + constraints=[_vec_to_vec_nd_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nd( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + lanes = pto.get_lanes(dtype) + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + src_stride = src.shape[1] + dst_stride = dst.shape[1] + + src_ptr = src.as_ptr() + + dst_offset = index_row * dst_stride + index_col + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + src_stride_bytes = src_stride * elem_bytes + dst_stride_bytes = dst_stride * elem_bytes + strides_aligned = src_stride_bytes % BLOCK_BYTE_SIZE == 0 and dst_stride_bytes % BLOCK_BYTE_SIZE == 0 + + if pto.constexpr(strides_aligned): + if index_col * elem_bytes % BLOCK_BYTE_SIZE == 0: + if pto.constexpr(valid_cols * elem_bytes % BLOCK_BYTE_SIZE == 0): + row_bytes = valid_cols * elem_bytes + total_bytes = valid_rows * row_bytes + row_burst_len = row_bytes // BLOCK_BYTE_SIZE + if pto.constexpr(valid_cols == src_stride and valid_cols == dst_stride and total_bytes >= BLOCK_BYTE_SIZE): + burst_len = total_bytes // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) + elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): + src_gap = (src_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + dst_gap = (dst_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, valid_rows, row_burst_len, src_gap, dst_gap) + else: + burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) + else: + repeat_times = (valid_cols + lanes - 1) // lanes + for i in range(0, valid_rows, 1): + remained = valid_cols + for j in range(0, repeat_times, 1): + pred, remained = pto.make_mask(dtype, remained) + src_off = i * src_stride + j * lanes + dst_off = i * dst_stride + j * lanes + data = pto.vlds(pto.addptr(src_ptr, src_off), 0) + pto.vsts(data, pto.addptr(dst_ptr, dst_off), 0, pred) + else: + full_repeats = valid_cols // lanes + remainder = valid_cols % lanes + for i in range(0, valid_rows, 1): + ureg = pto.init_align() + src_row_off = i * src_stride + dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) + for j in range(0, full_repeats, 1): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) + ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) + dst_row_ptr = pto.addptr(dst_row_ptr, lanes) + if pto.constexpr(remainder > 0): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) + ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) + pto.vstas(ureg, dst_row_ptr, 0) + else: + full_repeats = valid_cols // lanes + remainder = valid_cols % lanes + for i in range(0, valid_rows, 1): + ureg = pto.init_align() + src_row_off = i * src_stride + dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) + for j in range(0, full_repeats, 1): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) + ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) + dst_row_ptr = pto.addptr(dst_row_ptr, lanes) + if pto.constexpr(remainder > 0): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) + ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) + pto.vstas(ureg, dst_row_ptr, 0) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_BASIC_DTYPES, + constraints=[_vec_to_vec_nd_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nd_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + lanes = pto.get_lanes(dtype) + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + src_stride = src.shape[1] + dst_stride = dst.shape[1] + + src_ptr = src.as_ptr() + + dst_offset = index_row * dst_stride + index_col + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + src_stride_bytes = src_stride * elem_bytes + dst_stride_bytes = dst_stride * elem_bytes + strides_aligned = src_stride_bytes % BLOCK_BYTE_SIZE == 0 and dst_stride_bytes % BLOCK_BYTE_SIZE == 0 + + if pto.constexpr(strides_aligned): + if index_col * elem_bytes % BLOCK_BYTE_SIZE == 0: + if pto.constexpr(valid_cols * elem_bytes % BLOCK_BYTE_SIZE == 0): + row_bytes = valid_cols * elem_bytes + total_bytes = valid_rows * row_bytes + row_burst_len = row_bytes // BLOCK_BYTE_SIZE + if pto.constexpr(valid_cols == src_stride and valid_cols == dst_stride and total_bytes >= BLOCK_BYTE_SIZE): + burst_len = total_bytes // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) + elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): + src_gap = (src_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + dst_gap = (dst_stride - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, valid_rows, row_burst_len, src_gap, dst_gap) + else: + burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, 1, burst_len, 0, 0) + else: + repeat_times = (valid_cols + lanes - 1) // lanes + for i in range(0, valid_rows, 1): + remained = valid_cols + for j in range(0, repeat_times, 1): + pred, remained = pto.make_mask(dtype, remained) + src_off = i * src_stride + j * lanes + dst_off = i * dst_stride + j * lanes + data = pto.vlds(pto.addptr(src_ptr, src_off), 0) + pto.vsts(data, pto.addptr(dst_ptr, dst_off), 0, pred) + else: + full_repeats = valid_cols // lanes + remainder = valid_cols % lanes + for i in range(0, valid_rows, 1): + ureg = pto.init_align() + src_row_off = i * src_stride + dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) + for j in range(0, full_repeats, 1): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) + ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) + dst_row_ptr = pto.addptr(dst_row_ptr, lanes) + if pto.constexpr(remainder > 0): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) + ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) + pto.vstas(ureg, dst_row_ptr, 0) + else: + full_repeats = valid_cols // lanes + remainder = valid_cols % lanes + for i in range(0, valid_rows, 1): + ureg = pto.init_align() + src_row_off = i * src_stride + dst_row_ptr = pto.addptr(dst_ptr, i * dst_stride) + for j in range(0, full_repeats, 1): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + j * lanes), 0) + ureg = pto.vstus(ureg, lanes, data, dst_row_ptr) + dst_row_ptr = pto.addptr(dst_row_ptr, lanes) + if pto.constexpr(remainder > 0): + data = pto.vlds(pto.addptr(src_ptr, src_row_off + full_repeats * lanes), 0) + ureg = pto.vstus(ureg, remainder, data, dst_row_ptr) + pto.vstas(ureg, dst_row_ptr, 0) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_DTYPES, + constraints=[_vec_to_vec_nd_scalar_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nd_scalar( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dst_stride = dst.shape[1] + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + src_val = pto.load_scalar(src_ptr, 0) + dst_elem_offset = index_row * dst_stride + index_col + pto.store_scalar(src_val, dst_ptr, dst_elem_offset) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_BASIC_DTYPES, + constraints=[_vec_to_vec_nd_scalar_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nd_scalar_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dst_stride = dst.shape[1] + src_ptr = src.as_ptr() + dst_ptr = dst.as_ptr() + src_val = pto.load_scalar(src_ptr, 0) + dst_elem_offset = index_row * dst_stride + index_col + pto.store_scalar(src_val, dst_ptr, dst_elem_offset) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Vec (NZ, COL_MAJOR + ROW_MAJOR) - O2 +# --------------------------------------------------------------------------- + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_DTYPES, + constraints=[_vec_to_vec_nz_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nz( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 + else: + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + src_gap = src_stride_rows - valid_rows + dst_gap = dst_rows - valid_rows + + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, burst_num, burst_len, src_gap, dst_gap) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_VEC_BASIC_DTYPES, + constraints=[_vec_to_vec_nz_constraint], + advanced=True, +) +def template_tinsert_vec_to_vec_nz_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 + else: + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + src_gap = src_stride_rows - valid_rows + dst_gap = dst_rows - valid_rows + + pto.copy_ubuf_to_ubuf(src_ptr, dst_ptr, 0, burst_num, burst_len, src_gap, dst_gap) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Mat (NZ, COL_MAJOR + ROW_MAJOR) - O3 +# --------------------------------------------------------------------------- + + +_VEC_TO_MAT_DTYPES = [ + (pto.f16, pto.i64, pto.i64, pto.f16), + (pto.bf16, pto.i64, pto.i64, pto.bf16), + (pto.f32, pto.i64, pto.i64, pto.f32), + (pto.i32, pto.i64, pto.i64, pto.i32), + (pto.i8, pto.i64, pto.i64, pto.i8), +] + +_VEC_TO_MAT_BASIC_DTYPES = [ + (pto.f16, pto.i32, pto.i32, pto.f16), + (pto.bf16, pto.i32, pto.i32, pto.bf16), + (pto.f32, pto.i32, pto.i32, pto.f32), + (pto.i32, pto.i32, pto.i32, pto.i32), + (pto.i8, pto.i32, pto.i32, pto.i8), +] + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_DTYPES, + constraints=[_vec_to_mat_nz_constraint], + advanced=True, +) +def template_tinsert_vec_to_mat_nz( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 + else: + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + src_gap = src_stride_rows - valid_rows + dst_gap = dst_rows - valid_rows + + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(burst_num, src_gap, dst_gap)) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_BASIC_DTYPES, + constraints=[_vec_to_mat_nz_constraint], + advanced=True, +) +def template_tinsert_vec_to_mat_nz_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = valid_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + 1 + else: + src_stride_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + src_gap = src_stride_rows - valid_rows + dst_gap = dst_rows - valid_rows + + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(burst_num, src_gap, dst_gap)) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Mat (ND, ROW_MAJOR + NONE_BOX) - O4 +# --------------------------------------------------------------------------- + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_DTYPES, + constraints=[_vec_to_mat_nd_constraint], + advanced=True, +) +def template_tinsert_vec_to_mat_nd( + src: pto.Tile, + index_row: pto.i64, index_col: pto.i64, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + src_cols = src.shape[1] + dst_cols = dst.shape[1] + + src_ptr = src.as_ptr() + + dst_offset = index_row * dst_cols + index_col + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + row_bytes = valid_cols * elem_bytes + total_bytes = valid_rows * row_bytes + + if pto.constexpr(valid_cols == src_cols and valid_cols == dst_cols and total_bytes >= BLOCK_BYTE_SIZE): + burst_len = total_bytes // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) + elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): + row_burst_len = row_bytes // BLOCK_BYTE_SIZE + src_row_gap = (src_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + dst_row_gap = (dst_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, row_burst_len, nburst=(valid_rows, src_row_gap, dst_row_gap)) + else: + burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) + return None + + +@pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_BASIC_DTYPES, + constraints=[_vec_to_mat_nd_constraint], + advanced=True, +) +def template_tinsert_vec_to_mat_nd_basic( + src: pto.Tile, + index_row: pto.i32, index_col: pto.i32, + dst: pto.Tile, +): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + src_cols = src.shape[1] + dst_cols = dst.shape[1] + + src_ptr = src.as_ptr() + + dst_offset = index_row * dst_cols + index_col + dst_ptr = pto.addptr(dst.as_ptr(), dst_offset) + + row_bytes = valid_cols * elem_bytes + total_bytes = valid_rows * row_bytes + + if pto.constexpr(valid_cols == src_cols and valid_cols == dst_cols and total_bytes >= BLOCK_BYTE_SIZE): + burst_len = total_bytes // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) + elif pto.constexpr(row_bytes >= BLOCK_BYTE_SIZE): + row_burst_len = row_bytes // BLOCK_BYTE_SIZE + src_row_gap = (src_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + dst_row_gap = (dst_cols - valid_cols) * elem_bytes // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, row_burst_len, nburst=(valid_rows, src_row_gap, dst_row_gap)) + else: + burst_len = (total_bytes + BLOCK_BYTE_SIZE - 1) // BLOCK_BYTE_SIZE + pto.mte_ub_l1(src_ptr, dst_ptr, burst_len, nburst=(1, 0, 0)) + return None + + +# --------------------------------------------------------------------------- +# Vec -> Mat (NZ, Split2/Split4) - O5 +# Split large-tile NZ DMA into 2/4 independent segments to reduce L1 bank +# conflicts (mirrors pto-isa TInsertMode::SPLIT2 / SPLIT4). +# --------------------------------------------------------------------------- + + +def _make_split_template(split_count): + @pto.vkernel( + target="a5", + op="pto.tinsert", + dtypes=_VEC_TO_MAT_SPLIT_DTYPES, + name=f"template_tinsert_vec_to_mat_nz_split{split_count}", + constraints=[_vec_to_mat_nz_constraint], + advanced=True, + ) + def _split_fn(src: pto.Tile, index_row: pto.i32, index_col: pto.i32, dst: pto.Tile): + dtype = dst.element_type + elem_bytes = pto.bytewidth(dtype) + c0_standard = BLOCK_BYTE_SIZE // elem_bytes + s_frac_bits = dst.config.s_fractal_size + if pto.constexpr(s_frac_bits == 2 * BLOCK_BYTE_BITS): + c0_size = 2 * c0_standard + else: + c0_size = c0_standard + + valid_rows = src.shape[0] + valid_cols = src.shape[1] + dst_rows = dst.shape[0] + aligned_rows = (valid_rows + FRACTAL_NZ_ROW - 1) // FRACTAL_NZ_ROW * FRACTAL_NZ_ROW + + src_ptr = src.as_ptr() + + col_block = index_col // c0_size + col_mod = index_col - col_block * c0_size + dst_offset = dst_rows * c0_size * col_block + index_row * c0_size + col_mod + dst_base = pto.addptr(dst.as_ptr(), dst_offset) + + total_burst_num = (valid_cols + c0_size - 1) // c0_size + burst_len = aligned_rows * c0_size * elem_bytes // BLOCK_BYTE_SIZE + + compact = src.config.compact_mode + if pto.constexpr(compact == pto.CompactMode.NULL): + src_stride_rows = src.shape[0] + elif pto.constexpr(compact == pto.CompactMode.ROW_PLUS_ONE): + src_stride_rows = aligned_rows + 1 + else: + src_stride_rows = aligned_rows + src_gap = src_stride_rows - aligned_rows + dst_gap = dst_rows - aligned_rows + + part_num = total_burst_num // split_count + last_num = total_burst_num - part_num * (split_count - 1) + src_block_size = (burst_len + src_gap) * BLOCK_BYTE_SIZE // elem_bytes + dst_block_size = dst_rows * c0_size + + pto.mte_ub_l1(src_ptr, dst_base, burst_len, nburst=(part_num, src_gap, dst_gap)) + + src_ptr1 = pto.addptr(src_ptr, part_num * src_block_size) + dst_ptr1 = pto.addptr(dst_base, part_num * dst_block_size) + if pto.constexpr(split_count == 2): + pto.mte_ub_l1(src_ptr1, dst_ptr1, burst_len, nburst=(last_num, src_gap, dst_gap)) + else: + pto.mte_ub_l1(src_ptr1, dst_ptr1, burst_len, nburst=(part_num, src_gap, dst_gap)) + + if pto.constexpr(split_count == 4): + src_ptr2 = pto.addptr(src_ptr, 2 * part_num * src_block_size) + dst_ptr2 = pto.addptr(dst_base, 2 * part_num * dst_block_size) + pto.mte_ub_l1(src_ptr2, dst_ptr2, burst_len, nburst=(part_num, src_gap, dst_gap)) + + src_ptr3 = pto.addptr(src_ptr, 3 * part_num * src_block_size) + dst_ptr3 = pto.addptr(dst_base, 3 * part_num * dst_block_size) + pto.mte_ub_l1(src_ptr3, dst_ptr3, burst_len, nburst=(last_num, src_gap, dst_gap)) + return None + + return _split_fn + + +template_tinsert_vec_to_mat_nz_split2 = _make_split_template(2) +template_tinsert_vec_to_mat_nz_split4 = _make_split_template(4) diff --git a/test/lit/pto/tinsert_a5_acc_vec_lowering.pto b/test/lit/pto/tinsert_a5_acc_vec_lowering.pto new file mode 100644 index 0000000000..7e29892c6f --- /dev/null +++ b/test/lit/pto/tinsert_a5_acc_vec_lowering.pto @@ -0,0 +1,57 @@ +// RUN: ptoas --pto-arch=a5 %s | FileCheck %s + +module attributes {"pto.target_arch" = "a5"} { + // A5: acc->vec with default (SingleModeVec0, NoRelu) -- no template args. + // CHECK-LABEL: AICORE void tinsert_acc_vec_default( + // CHECK: TINSERT( + // CHECK-NOT: ReluPreMode:: + // CHECK-NOT: AccToVecMode:: + // CHECK-NOT: TInsertMode:: + func.func @tinsert_acc_vec_default() { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } + + // A5: acc->vec with ReLU. + // CHECK-LABEL: AICORE void tinsert_acc_vec_relu( + // CHECK: TINSERT<{{.*}}, ReluPreMode::NormalRelu>( + func.func @tinsert_acc_vec_relu() { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + {reluPreMode = #pto} + return + } + + // A5: acc->vec with AccToVecMode + ReLU. + // CHECK-LABEL: AICORE void tinsert_acc_vec_splitm( + // CHECK: TINSERT<{{.*}}, {{.*}}AccToVecMode::DualModeSplitM, ReluPreMode::NormalRelu>( + func.func @tinsert_acc_vec_splitm() { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + {accToVecMode = #pto.acc_to_vec_mode, reluPreMode = #pto} + return + } + + // A5: acc->mat with ReLU (cast quantization). + // CHECK-LABEL: AICORE void tinsert_acc_mat_relu( + // CHECK: TINSERT<{{.*}}, ReluPreMode::NormalRelu>( + func.func @tinsert_acc_mat_relu() { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + {reluPreMode = #pto} + return + } +} diff --git a/test/lit/pto/tinsert_a5_extended_modes.pto b/test/lit/pto/tinsert_a5_extended_modes.pto new file mode 100644 index 0000000000..8b91cf1871 --- /dev/null +++ b/test/lit/pto/tinsert_a5_extended_modes.pto @@ -0,0 +1,73 @@ +// RUN: ptoas --pto-arch=a5 %s | FileCheck %s + +module attributes {"pto.target_arch" = "a5"} { + // A5: vec->mat NZ with SPLIT2. + // CHECK-LABEL: AICORE void tinsert_vec_mat_split2( + // CHECK: TINSERT + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + {tinsertMode = #pto.tinsert_mode} + return + } + + // A5: vec->mat NZ with SPLIT4. + // CHECK-LABEL: AICORE void tinsert_vec_mat_split4( + // CHECK: TINSERT + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + {tinsertMode = #pto.tinsert_mode} + return + } + + // A5: acc->mat with scalar pre-quantization (generic MLIR form). + // CHECK-LABEL: AICORE void tinsert_acc_mat_scalar_quant( + // CHECK: TINSERT<{{.*}}TileType::Mat, int8_t{{.*}}, {{.*}}TileType::Acc, float{{.*}}>( + func.func @tinsert_acc_mat_scalar_quant() { + %c0 = arith.constant 0 : index + %scalar = arith.constant 123456789 : i64 + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + "pto.tinsert"(%src, %c0, %c0, %dst, %scalar) <{operandSegmentSizes = array}> : + (!pto.tile_buf, + index, index, + !pto.tile_buf, + i64) -> () + return + } + + // A5: acc->mat with fp (vector) quantization (generic MLIR form). + // CHECK-LABEL: AICORE void tinsert_acc_mat_fp_quant( + // CHECK: TINSERT<{{.*}}TileType::Mat, int8_t{{.*}}, {{.*}}TileType::Acc, float{{.*}}, {{.*}}TileType::Scaling{{.*}}>( + func.func @tinsert_acc_mat_fp_quant() { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %fp = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + "pto.tinsert"(%src, %c0, %c0, %dst, %fp) <{operandSegmentSizes = array}> : + (!pto.tile_buf, + index, index, + !pto.tile_buf, + !pto.tile_buf) -> () + return + } + + // A5: vec->vec NZ->NZ. + // CHECK-LABEL: AICORE void tinsert_vec_vec_nz( + // CHECK: TINSERT( + func.func @tinsert_vec_vec_nz() { + %c0 = arith.constant 0 : index + %src = pto.alloc_tile : !pto.tile_buf + %dst = pto.alloc_tile : !pto.tile_buf + pto.tinsert ins(%src, %c0, %c0 : !pto.tile_buf, index, index) + outs(%dst : !pto.tile_buf) + return + } +} diff --git a/test/lit/pto/tinsert_verify_a5_invalid_loc.pto b/test/lit/pto/tinsert_verify_a5_invalid_loc.pto index 7125132e5a..e01a1c18e9 100644 --- a/test/lit/pto/tinsert_verify_a5_invalid_loc.pto +++ b/test/lit/pto/tinsert_verify_a5_invalid_loc.pto @@ -13,4 +13,4 @@ module attributes {"pto.target_arch" = "a5"} { } } -// CHECK: error: 'pto.tinsert' op expects A5 tinsert to use a supported src/dst loc pair: acc->mat, vec->mat, or vec->vec +// CHECK: error: 'pto.tinsert' op expects A5 tinsert to use a supported src/dst loc pair: acc->mat, acc->vec, vec->mat, or vec->vec diff --git a/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt index 22fbc3f731..e2be53efbe 100644 --- a/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt +++ b/test/tilelang_st/npu/a5/src/st/testcase/CMakeLists.txt @@ -110,6 +110,8 @@ endfunction() # -------------------------------------------------------------------------- set(ALL_TESTCASES tadd + tinsert + tinsert_vec tsub tmul tdiv diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/CMakeLists.txt new file mode 100644 index 0000000000..6878082ab1 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/CMakeLists.txt @@ -0,0 +1,9 @@ +# 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. + +pto_tilelang_cube_st(tinsert PTO_LEVEL level3) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/cases.py new file mode 100644 index 0000000000..53da565802 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/cases.py @@ -0,0 +1,87 @@ +# 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. + +# coding=utf-8 + +"""Test cases for pto.tinsert ST (Acc->Mat NZ, Acc->Vec ND/NZ).""" + +import numpy as np + + +CASES = [ + { + "name": "acc2mat_f16_16x16", + "kernel": "TINSERT_acc2mat_f16_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float16, + "path": "acc2mat_nz", + "has_output": False, + "eps": 1e-2, + }, + { + "name": "acc2mat_f16_32x32", + "kernel": "TINSERT_acc2mat_f16_32x32", + "m": 32, "k": 32, "n": 32, + "dtype": np.float16, + "dtype_out": np.float16, + "path": "acc2mat_nz", + "has_output": False, + "eps": 1e-2, + }, + { + "name": "acc2mat_bf16_16x16", + "kernel": "TINSERT_acc2mat_bf16_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float16, + "path": "acc2mat_nz", + "has_output": False, + "eps": 1e-2, + }, + { + "name": "acc2mat_f32_16x16", + "kernel": "TINSERT_acc2mat_f32_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float32, + "path": "acc2mat_nz", + "has_output": False, + "eps": 1e-2, + }, + { + "name": "acc2vec_nd_f16_16x16", + "kernel": "TINSERT_acc2vec_nd_f16_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float32, + "path": "acc2vec_nd", + "has_output": True, + "eps": 1e-2, + }, + { + "name": "acc2vec_nd_f32_16x16", + "kernel": "TINSERT_acc2vec_nd_f32_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float32, + "path": "acc2vec_nd", + "has_output": True, + "eps": 1e-2, + }, + { + "name": "acc2vec_nz_f32_16x16", + "kernel": "TINSERT_acc2vec_nz_f32_16x16", + "m": 16, "k": 16, "n": 16, + "dtype": np.float16, + "dtype_out": np.float32, + "path": "acc2vec_nz", + "has_output": True, + "eps": 1e-2, + }, +] diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/compare.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/compare.py new file mode 100644 index 0000000000..90ae4a1667 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/compare.py @@ -0,0 +1,66 @@ +# 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. + +# coding=utf-8 + +import os +import sys +import numpy as np + +from cases import CASES +from st_common import result_cmp, style_fail, style_pass + + +def main(): + case_filter = sys.argv[1] if len(sys.argv) > 1 else None + + all_passed = True + for case in CASES: + if case_filter is not None and case["name"] != case_filter: + continue + + case_dir = case["name"] + m, n = case["m"], case["n"] + dtype_out = case["dtype_out"] + + if not case["has_output"]: + print(style_pass(f"[INFO] {case['name']}: compile-only (no output comparison)")) + continue + + golden_path = os.path.join(case_dir, "golden.bin") + output_path = os.path.join(case_dir, "output.bin") + + if not os.path.exists(output_path): + print(style_fail(f"[ERROR] {case['name']}: output.bin not found (kernel may not have written output)")) + all_passed = False + continue + + golden = np.fromfile(golden_path, dtype=dtype_out) + output = np.fromfile(output_path, dtype=dtype_out) + + if golden.shape != output.shape: + print(style_fail( + f"[ERROR] {case['name']}: shape mismatch golden={golden.shape} output={output.shape}" + )) + all_passed = False + continue + + ok = result_cmp(golden, output, case["eps"]) + if ok: + print(style_pass(f"[INFO] {case['name']}: compare passed")) + else: + print(style_fail(f"[ERROR] {case['name']}: compare failed")) + all_passed = False + + if not all_passed: + sys.exit(2) + print(style_pass("[INFO] all cases passed")) + + +if __name__ == "__main__": + main() diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/gen_data.py new file mode 100644 index 0000000000..ddceacec0e --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/gen_data.py @@ -0,0 +1,35 @@ +# 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. + +# coding=utf-8 + +import os +import numpy as np + +from cases import CASES +from st_common import setup_case_rng, save_case_data + + +for case in CASES: + setup_case_rng(case) + m, k, n = case["m"], case["k"], case["n"] + dtype = case["dtype"] + dtype_out = case["dtype_out"] + + A = np.random.uniform(-1.0, 1.0, size=(m, k)).astype(dtype) + B = np.random.uniform(-1.0, 1.0, size=(k, n)).astype(dtype) + golden_f32 = np.matmul(A.astype(np.float32), B.astype(np.float32)) + golden = golden_f32.astype(dtype_out) + + data = {"input1": A, "input2": B, "golden": golden} + + save_case_data(case["name"], data) + print( + f"[INFO] gen_data: {case['name']} A=({m},{k}) B=({k},{n}) " + f"dtype={dtype.__name__} dtype_out={dtype_out.__name__}" + ) diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/launch.cpp new file mode 100644 index 0000000000..9671721364 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/launch.cpp @@ -0,0 +1,56 @@ +// 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. + +#include + +#ifndef AICORE +#define AICORE [aicore] +#endif + +extern "C" __global__ AICORE void TINSERT_acc2mat_f16_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint16_t *c); +extern "C" __global__ AICORE void TINSERT_acc2mat_f16_32x32( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint16_t *c); +extern "C" __global__ AICORE void TINSERT_acc2mat_bf16_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint16_t *c); +extern "C" __global__ AICORE void TINSERT_acc2mat_f32_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint32_t *c); +extern "C" __global__ AICORE void TINSERT_acc2vec_nd_f16_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint32_t *c); +extern "C" __global__ AICORE void TINSERT_acc2vec_nd_f32_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint32_t *c); +extern "C" __global__ AICORE void TINSERT_acc2vec_nz_f32_16x16( + __gm__ uint16_t *a, __gm__ uint16_t *b, __gm__ uint32_t *c); + +void LaunchAcc2Mat_f16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2mat_f16_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint16_t *)c); +} + +void LaunchAcc2Mat_f16_32x32(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2mat_f16_32x32<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint16_t *)c); +} + +void LaunchAcc2Mat_bf16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2mat_bf16_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint16_t *)c); +} + +void LaunchAcc2Mat_f32_16x16(uint16_t *a, uint16_t *b, uint32_t *c, void *stream) { + TINSERT_acc2mat_f32_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint32_t *)c); +} + +void LaunchAcc2VecND_f16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2vec_nd_f16_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint32_t *)c); +} + +void LaunchAcc2VecND_f32_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2vec_nd_f32_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint32_t *)c); +} + +void LaunchAcc2VecNZ_f32_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream) { + TINSERT_acc2vec_nz_f32_16x16<<<1, nullptr, stream>>>((__gm__ uint16_t *)a, (__gm__ uint16_t *)b, (__gm__ uint32_t *)c); +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/main.cpp new file mode 100644 index 0000000000..7da0c77cb4 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/main.cpp @@ -0,0 +1,150 @@ +// 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. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include +#include +#include + +using namespace PtoTestCommon; + +using LaunchFn = void (*)(uint16_t *, uint16_t *, uint16_t *, void *); + +void LaunchAcc2Mat_f16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2Mat_f16_32x32(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2Mat_bf16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2Mat_f32_16x16(uint16_t *a, uint16_t *b, uint32_t *c, void *stream); +void LaunchAcc2VecND_f16_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2VecND_f32_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); +void LaunchAcc2VecNZ_f32_16x16(uint16_t *a, uint16_t *b, uint16_t *c, void *stream); + +struct TestCase { + const char *name; + LaunchFn launch; + size_t m, k, n; + bool has_output; + size_t out_elem_bytes; +}; + +static const TestCase kCases[] = { + {"acc2mat_f16_16x16", LaunchAcc2Mat_f16_16x16, 16, 16, 16, false, 2}, + {"acc2mat_f16_32x32", LaunchAcc2Mat_f16_32x32, 32, 32, 32, false, 2}, + {"acc2mat_bf16_16x16", LaunchAcc2Mat_bf16_16x16, 16, 16, 16, false, 2}, + {"acc2mat_f32_16x16", reinterpret_cast(LaunchAcc2Mat_f32_16x16), 16, 16, 16, false, 4}, + {"acc2vec_nd_f16_16x16", LaunchAcc2VecND_f16_16x16, 16, 16, 16, true, 4}, + {"acc2vec_nd_f32_16x16", LaunchAcc2VecND_f32_16x16, 16, 16, 16, true, 4}, + {"acc2vec_nz_f32_16x16", LaunchAcc2VecNZ_f32_16x16, 16, 16, 16, true, 4}, +}; +static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); + +static int RunCase(const TestCase &tc, int deviceId, aclrtStream stream) { + (void)deviceId; + int rc = 0; + const size_t aElems = tc.m * tc.k; + const size_t bElems = tc.k * tc.n; + const size_t outElems = tc.m * tc.n; + const size_t aBytes = aElems * sizeof(uint16_t); + const size_t bBytes = bElems * sizeof(uint16_t); + const size_t outBytes = outElems * tc.out_elem_bytes; + size_t aFileSize = aBytes; + size_t bFileSize = bBytes; + + std::printf("[INFO] === case: %s (m=%zu, k=%zu, n=%zu) ===\n", tc.name, tc.m, tc.k, tc.n); + + std::string caseDir = std::string("./") + tc.name; + + void *aHost = nullptr, *bHost = nullptr, *outHost = nullptr; + void *aDev = nullptr, *bDev = nullptr, *outDev = nullptr; + + aclrtMallocHost(&aHost, aBytes); + aclrtMallocHost(&bHost, bBytes); + aclrtMallocHost(&outHost, outBytes); + aclrtMalloc(&aDev, aBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&bDev, bBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&outDev, outBytes, ACL_MEM_MALLOC_HUGE_FIRST); + + if (!ReadFile((caseDir + "/input1.bin").c_str(), aFileSize, aHost, aBytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input1.bin\n", caseDir.c_str()); + rc = 1; + } + if (rc == 0 && !ReadFile((caseDir + "/input2.bin").c_str(), bFileSize, bHost, bBytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input2.bin\n", caseDir.c_str()); + rc = 1; + } + + if (rc == 0) { + aclrtMemcpy(aDev, aBytes, aHost, aBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemcpy(bDev, bBytes, bHost, bBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemset(outDev, outBytes, 0, outBytes); + + tc.launch( + static_cast(aDev), + static_cast(bDev), + static_cast(outDev), + stream + ); + + aclrtSynchronizeStream(stream); + + if (tc.has_output) { + aclrtMemcpy(outHost, outBytes, outDev, outBytes, ACL_MEMCPY_DEVICE_TO_HOST); + if (!WriteFile((caseDir + "/output.bin").c_str(), outHost, outBytes)) { + std::fprintf(stderr, "[ERROR] failed to write %s/output.bin\n", caseDir.c_str()); + rc = 1; + } + } + } + + if (aDev) aclrtFree(aDev); + if (bDev) aclrtFree(bDev); + if (outDev) aclrtFree(outDev); + if (aHost) aclrtFreeHost(aHost); + if (bHost) aclrtFreeHost(bHost); + if (outHost) aclrtFreeHost(outHost); + + if (rc == 0) + std::printf("[INFO] case %s done\n", tc.name); + return rc; +} + +int main(int argc, char *argv[]) { + const char *caseFilter = (argc > 1) ? argv[1] : nullptr; + + int rc = 0; + int deviceId = 0; + aclrtStream stream = nullptr; + + aclInit(nullptr); + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) { + deviceId = std::atoi(envDevice); + } + aclrtSetDevice(deviceId); + aclrtCreateStream(&stream); + + for (size_t i = 0; i < kNumCases; ++i) { + if (caseFilter != nullptr && std::strcmp(kCases[i].name, caseFilter) != 0) { + continue; + } + int ret = RunCase(kCases[i], deviceId, stream); + if (ret != 0) { + std::fprintf(stderr, "[ERROR] case %s failed\n", kCases[i].name); + rc = 1; + break; + } + } + + if (stream != nullptr) + aclrtDestroyStream(stream); + aclrtResetDevice(deviceId); + aclFinalize(); + + return rc; +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert/tinsert.pto b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/tinsert.pto new file mode 100644 index 0000000000..5385679f23 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert/tinsert.pto @@ -0,0 +1,595 @@ +// 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 and copyright holder. + +// ST kernels for pto.tinsert ckernel templates. +// Each kernel: load A,B from GM (MTE) -> matmul -> tinsert(Acc->dst) -> store result to GM. +// Uses low-level MTE ops (level3) for boundary data movement. + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + // ------------------------------------------------------------------------- + // Case 3: Acc → Vec ND (16×16×16, f16×f16→f32→f16) + // Pattern: mte(GM→L1→L0) → matmul → tinsert(acc→vec_ND) → mte(l0c→GM) + // Output: L0C accumulator (f32, ND layout via nz2nd). + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2vec_nd_f16_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_vec_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + %dst_vec = pto.tile_buf_addr %dst_vec_tile + : !pto.tile_buf + -> !pto.ptr + %l0c = pto.tile_buf_addr %l0c_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_vec_tile : !pto.tile_buf) + + pto.mte_l0c_gm %l0c, %out_gm, %c16_i64, %c16_i64, %c16_i64, %c16_i64, %c0_i64, %c0_i64, + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Case 5: Acc → Vec NZ (16×16×16, f16×f16→f32→f32) + // Pattern: mte(GM→L1→L0) → matmul → tinsert(acc→vec_NZ) → mte(l0c→GM) + // Output: L0C accumulator (f32, ND layout via nz2nd). + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2vec_nz_f32_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c64_i64 = arith.constant 64 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_vec_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + %dst_vec = pto.tile_buf_addr %dst_vec_tile + : !pto.tile_buf + -> !pto.ptr + %l0c = pto.tile_buf_addr %l0c_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_vec_tile : !pto.tile_buf) + + pto.mte_l0c_gm %l0c, %out_gm, %c16_i64, %c16_i64, %c16_i64, %c16_i64, %c0_i64, %c0_i64, + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Case 1: Acc → Mat NZ (16×16×16, f16×f16→f32→f16) + // Pattern: mte(GM→L1→L0) → matmul → tinsert(acc→mat_NZ) + // Note: compile+run only (no output comparison). + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2mat_f16_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_mat_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_mat_tile : !pto.tile_buf) + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Case 2: Acc → Mat NZ (32×32×32, f16×f16→f32→f16) — larger tile variant + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2mat_f16_32x32(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_mat_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c32_i64, %c32_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c32_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c32_i64, %c32_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c32_i64, %c32_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c32_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c32_i64, %c32_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_mat_tile : !pto.tile_buf) + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Acc → Mat NZ (16×16×16, f16×f16→f32→bf16) + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2mat_bf16_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_mat_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_mat_tile : !pto.tile_buf) + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Acc → Mat NZ (16×16×16, f16×f16→f32→f32) + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2mat_f32_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_mat_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_mat_tile : !pto.tile_buf) + + pto.barrier #pto.pipe + return + } + + // ------------------------------------------------------------------------- + // Acc → Vec ND (16×16×16, f16×f16→f32→f32) + // ------------------------------------------------------------------------- + func.func @TINSERT_acc2vec_nd_f32_16x16(%a_gm: !pto.ptr, + %b_gm: !pto.ptr, + %out_gm: !pto.ptr) attributes {pto.kernel} { + %c0_i64 = arith.constant 0 : i64 + %c1_i64 = arith.constant 1 : i64 + %c16_i64 = arith.constant 16 : i64 + %c32_i64 = arith.constant 32 : i64 + %c512_i64 = arith.constant 512 : i64 + %false = arith.constant false + + %l1_a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l1_b_tile = pto.alloc_tile addr = %c512_i64 + : !pto.tile_buf + %l0a_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0b_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %l0c_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + %dst_vec_tile = pto.alloc_tile addr = %c0_i64 + : !pto.tile_buf + + %l1_a = pto.tile_buf_addr %l1_a_tile + : !pto.tile_buf + -> !pto.ptr + %l1_b = pto.tile_buf_addr %l1_b_tile + : !pto.tile_buf + -> !pto.ptr + %l0a = pto.tile_buf_addr %l0a_tile + : !pto.tile_buf + -> !pto.ptr + %l0b = pto.tile_buf_addr %l0b_tile + : !pto.tile_buf + -> !pto.ptr + %l0c = pto.tile_buf_addr %l0c_tile + : !pto.tile_buf + -> !pto.ptr + + pto.mte_gm_l1_frac %a_gm, %l1_a, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID0"] + pto.mte_l1_l0a %l1_a, %l0a, %c16_i64, %c16_i64 + : !pto.ptr, !pto.ptr, i64, i64 + + pto.mte_gm_l1_frac %b_gm, %l1_b, nd2nz, + shape(%c16_i64, %c16_i64), src_layout(%c32_i64), + dst_group(%c1_i64, %c1_i64, %c16_i64, %c0_i64), ctrl(%c0_i64, %false) + : !pto.ptr, !pto.ptr, nd2nz, + shape i64, i64, src_layout(i64), dst_group i64, i64, i64, i64, ctrl i64, i1 + pto.set_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.wait_flag["PIPE_MTE2", "PIPE_MTE1", "EVENT_ID1"] + pto.mte_l1_l0b %l1_b, %l0b, %c16_i64, %c16_i64 {transpose = true} + : !pto.ptr, !pto.ptr, i64, i64 + + pto.set_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.wait_flag["PIPE_MTE1", "PIPE_M", "EVENT_ID0"] + pto.tmatmul ins(%l0a_tile, %l0b_tile : + !pto.tile_buf, + !pto.tile_buf) + outs(%l0c_tile : !pto.tile_buf) + + pto.set_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + pto.wait_flag["PIPE_M", "PIPE_FIX", "EVENT_ID1"] + %c0_idx = arith.constant 0 : index + pto.tinsert ins(%l0c_tile, %c0_idx, %c0_idx : + !pto.tile_buf, + index, index) + outs(%dst_vec_tile : !pto.tile_buf) + + pto.mte_l0c_gm %l0c, %out_gm, %c16_i64, %c16_i64, %c16_i64, %c16_i64, %c0_i64, %c0_i64, + nz2nd + : !pto.ptr, !pto.ptr, i64, i64, i64, i64, i64, i64 + + pto.barrier #pto.pipe + return + } + +} diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/CMakeLists.txt b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/CMakeLists.txt new file mode 100644 index 0000000000..be1519c4ae --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/CMakeLists.txt @@ -0,0 +1,9 @@ +# 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. + +pto_tilelang_vec_st(tinsert_vec) \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/cases.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/cases.py new file mode 100644 index 0000000000..4f2c5d1b66 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/cases.py @@ -0,0 +1,73 @@ +# 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. + +# coding=utf-8 + +"""Test cases for pto.tinsert Vec->Vec ND vkernel ST.""" + +import numpy as np +from ml_dtypes import bfloat16 + + +CASES = [ + { + "name": "vec2vec_nd_f16_16x16_into_32x32_idx00", + "kernel": "TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx00", + "dtype": np.float16, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 0, + "index_col": 0, + "has_output": True, + "eps": 1e-2, + }, + { + "name": "vec2vec_nd_f16_16x16_into_32x32_idx816", + "kernel": "TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx816", + "dtype": np.float16, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 8, + "index_col": 16, + "has_output": True, + "eps": 1e-2, + }, + { + "name": "vec2vec_nd_f32_16x16_into_32x32_idx00", + "kernel": "TINSERT_vec2vec_nd_f32_16x16_into_32x32_idx00", + "dtype": np.float32, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 0, + "index_col": 0, + "has_output": True, + "eps": 1e-6, + }, + { + "name": "vec2vec_nd_bf16_16x16_into_32x32_idx00", + "kernel": "TINSERT_vec2vec_nd_bf16_16x16_into_32x32_idx00", + "dtype": bfloat16, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 0, + "index_col": 0, + "has_output": True, + "eps": 1e-2, + }, + { + "name": "vec2vec_nd_i32_16x16_into_32x32_idx00", + "kernel": "TINSERT_vec2vec_nd_i32_16x16_into_32x32_idx00", + "dtype": np.int32, + "src_shape": (16, 16), + "dst_shape": (32, 32), + "index_row": 0, + "index_col": 0, + "has_output": True, + "eps": 0, + }, +] \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/compare.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/compare.py new file mode 100644 index 0000000000..0272300355 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/compare.py @@ -0,0 +1,66 @@ +# 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. + +# coding=utf-8 + +import os +import sys +import numpy as np + +from cases import CASES +from st_common import result_cmp, style_fail, style_pass + + +def main(): + case_filter = sys.argv[1] if len(sys.argv) > 1 else None + + all_passed = True + for case in CASES: + if case_filter is not None and case["name"] != case_filter: + continue + + case_dir = case["name"] + dtype = case["dtype"] + dst_rows, dst_cols = case["dst_shape"] + + if not case["has_output"]: + print(style_pass(f"[INFO] {case['name']}: compile-only (no output comparison)")) + continue + + golden_path = os.path.join(case_dir, "golden.bin") + output_path = os.path.join(case_dir, "output.bin") + + if not os.path.exists(output_path): + print(style_fail(f"[ERROR] {case['name']}: output.bin not found")) + all_passed = False + continue + + golden = np.fromfile(golden_path, dtype=dtype).reshape(dst_rows, dst_cols) + output = np.fromfile(output_path, dtype=dtype).reshape(dst_rows, dst_cols) + + if golden.shape != output.shape: + print(style_fail( + f"[ERROR] {case['name']}: shape mismatch golden={golden.shape} output={output.shape}" + )) + all_passed = False + continue + + ok = result_cmp(golden, output, case["eps"]) + if ok: + print(style_pass(f"[INFO] {case['name']}: compare passed")) + else: + print(style_fail(f"[ERROR] {case['name']}: compare failed")) + all_passed = False + + if not all_passed: + sys.exit(2) + print(style_pass("[INFO] all cases passed")) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/gen_data.py b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/gen_data.py new file mode 100644 index 0000000000..e3faf1b586 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/gen_data.py @@ -0,0 +1,35 @@ +# 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. + +# coding=utf-8 + +import numpy as np + +from cases import CASES +from st_common import setup_case_rng, save_case_data + + +for case in CASES: + setup_case_rng(case) + dtype = case["dtype"] + src_rows, src_cols = case["src_shape"] + dst_rows, dst_cols = case["dst_shape"] + idx_row, idx_col = case["index_row"], case["index_col"] + + src = np.random.uniform(-1.0, 1.0, size=(src_rows, src_cols)).astype(dtype) + dst = np.random.uniform(-1.0, 1.0, size=(dst_rows, dst_cols)).astype(dtype) + + golden = dst.copy() + golden[idx_row:idx_row + src_rows, idx_col:idx_col + src_cols] = src + + data = {"input1": src, "input2": dst, "golden": golden} + save_case_data(case["name"], data) + print( + f"[INFO] gen_data: {case['name']} src=({src_rows},{src_cols}) dst=({dst_rows},{dst_cols}) " + f"idx=({idx_row},{idx_col}) dtype={dtype.__name__}" + ) \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/launch.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/launch.cpp new file mode 100644 index 0000000000..4e36cf8bf1 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/launch.cpp @@ -0,0 +1,44 @@ +// 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. + +#include + +#ifndef AICORE +#define AICORE [aicore] +#endif + +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx00( + __gm__ uint16_t *src, __gm__ uint16_t *dst, __gm__ uint16_t *out); +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx816( + __gm__ uint16_t *src, __gm__ uint16_t *dst, __gm__ uint16_t *out); +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_f32_16x16_into_32x32_idx00( + __gm__ float *src, __gm__ float *dst, __gm__ float *out); +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_bf16_16x16_into_32x32_idx00( + __gm__ uint16_t *src, __gm__ uint16_t *dst, __gm__ uint16_t *out); +extern "C" __global__ AICORE void TINSERT_vec2vec_nd_i32_16x16_into_32x32_idx00( + __gm__ int32_t *src, __gm__ int32_t *dst, __gm__ int32_t *out); + +void LaunchVec2VecND_f16_16x16_into_32x32_idx00(uint16_t *src, uint16_t *dst, uint16_t *out, void *stream) { + TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx00<<<1, nullptr, stream>>>((__gm__ uint16_t *)src, (__gm__ uint16_t *)dst, (__gm__ uint16_t *)out); +} + +void LaunchVec2VecND_f16_16x16_into_32x32_idx816(uint16_t *src, uint16_t *dst, uint16_t *out, void *stream) { + TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx816<<<1, nullptr, stream>>>((__gm__ uint16_t *)src, (__gm__ uint16_t *)dst, (__gm__ uint16_t *)out); +} + +void LaunchVec2VecND_f32_16x16_into_32x32_idx00(float *src, float *dst, float *out, void *stream) { + TINSERT_vec2vec_nd_f32_16x16_into_32x32_idx00<<<1, nullptr, stream>>>((__gm__ float *)src, (__gm__ float *)dst, (__gm__ float *)out); +} + +void LaunchVec2VecND_bf16_16x16_into_32x32_idx00(uint16_t *src, uint16_t *dst, uint16_t *out, void *stream) { + TINSERT_vec2vec_nd_bf16_16x16_into_32x32_idx00<<<1, nullptr, stream>>>((__gm__ uint16_t *)src, (__gm__ uint16_t *)dst, (__gm__ uint16_t *)out); +} + +void LaunchVec2VecND_i32_16x16_into_32x32_idx00(int32_t *src, int32_t *dst, int32_t *out, void *stream) { + TINSERT_vec2vec_nd_i32_16x16_into_32x32_idx00<<<1, nullptr, stream>>>((__gm__ int32_t *)src, (__gm__ int32_t *)dst, (__gm__ int32_t *)out); +} \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/main.cpp b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/main.cpp new file mode 100644 index 0000000000..0dd7b743b5 --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/main.cpp @@ -0,0 +1,158 @@ +// 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. + +#include "acl/acl.h" +#include "test_common.h" +#include +#include +#include +#include +#include + +using namespace PtoTestCommon; + +void LaunchVec2VecND_f16_16x16_into_32x32_idx00( + uint16_t *src, uint16_t *dst, uint16_t *out, void *stream); +void LaunchVec2VecND_f16_16x16_into_32x32_idx816( + uint16_t *src, uint16_t *dst, uint16_t *out, void *stream); +void LaunchVec2VecND_f32_16x16_into_32x32_idx00( + float *src, float *dst, float *out, void *stream); +void LaunchVec2VecND_bf16_16x16_into_32x32_idx00( + uint16_t *src, uint16_t *dst, uint16_t *out, void *stream); +void LaunchVec2VecND_i32_16x16_into_32x32_idx00( + int32_t *src, int32_t *dst, int32_t *out, void *stream); + +using LaunchFn = void (*)(void *, void *, void *, void *); + +struct TestCase { + const char *name; + LaunchFn launch; + size_t src_rows, src_cols; + size_t dst_rows, dst_cols; + size_t idx_row, idx_col; + bool has_output; + size_t elem_bytes; +}; + +static const TestCase kCases[] = { + {"vec2vec_nd_f16_16x16_into_32x32_idx00", + reinterpret_cast(LaunchVec2VecND_f16_16x16_into_32x32_idx00), + 16, 16, 32, 32, 0, 0, true, 2}, + {"vec2vec_nd_f16_16x16_into_32x32_idx816", + reinterpret_cast(LaunchVec2VecND_f16_16x16_into_32x32_idx816), + 16, 16, 32, 32, 8, 16, true, 2}, + {"vec2vec_nd_f32_16x16_into_32x32_idx00", + reinterpret_cast(LaunchVec2VecND_f32_16x16_into_32x32_idx00), + 16, 16, 32, 32, 0, 0, true, 4}, + {"vec2vec_nd_bf16_16x16_into_32x32_idx00", + reinterpret_cast(LaunchVec2VecND_bf16_16x16_into_32x32_idx00), + 16, 16, 32, 32, 0, 0, true, 2}, + {"vec2vec_nd_i32_16x16_into_32x32_idx00", + reinterpret_cast(LaunchVec2VecND_i32_16x16_into_32x32_idx00), + 16, 16, 32, 32, 0, 0, true, 4}, +}; +static constexpr size_t kNumCases = sizeof(kCases) / sizeof(kCases[0]); + +static int RunCase(const TestCase &tc, int deviceId, aclrtStream stream) { + (void)deviceId; + int rc = 0; + const size_t srcElems = tc.src_rows * tc.src_cols; + const size_t dstElems = tc.dst_rows * tc.dst_cols; + const size_t srcBytes = srcElems * tc.elem_bytes; + const size_t dstBytes = dstElems * tc.elem_bytes; + size_t srcFileSize = srcBytes; + size_t dstFileSize = dstBytes; + + std::printf("[INFO] === case: %s src=%zux%zu dst=%zux%zu idx=(%zu,%zu) ===\n", + tc.name, tc.src_rows, tc.src_cols, tc.dst_rows, tc.dst_cols, + tc.idx_row, tc.idx_col); + + std::string caseDir = std::string("./") + tc.name; + + void *srcHost = nullptr, *dstHost = nullptr, *outHost = nullptr; + void *srcDev = nullptr, *dstDev = nullptr, *outDev = nullptr; + + aclrtMallocHost(&srcHost, srcBytes); + aclrtMallocHost(&dstHost, dstBytes); + aclrtMallocHost(&outHost, dstBytes); + aclrtMalloc(&srcDev, srcBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&dstDev, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST); + aclrtMalloc(&outDev, dstBytes, ACL_MEM_MALLOC_HUGE_FIRST); + + if (!ReadFile((caseDir + "/input1.bin").c_str(), srcFileSize, srcHost, srcBytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input1.bin\n", caseDir.c_str()); + rc = 1; + } + if (rc == 0 && !ReadFile((caseDir + "/input2.bin").c_str(), dstFileSize, dstHost, dstBytes)) { + std::fprintf(stderr, "[ERROR] failed to read %s/input2.bin\n", caseDir.c_str()); + rc = 1; + } + + if (rc == 0) { + aclrtMemcpy(srcDev, srcBytes, srcHost, srcBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemcpy(dstDev, dstBytes, dstHost, dstBytes, ACL_MEMCPY_HOST_TO_DEVICE); + aclrtMemset(outDev, dstBytes, 0, dstBytes); + + tc.launch(srcDev, dstDev, outDev, stream); + + aclrtSynchronizeStream(stream); + + if (tc.has_output) { + aclrtMemcpy(outHost, dstBytes, outDev, dstBytes, ACL_MEMCPY_DEVICE_TO_HOST); + if (!WriteFile((caseDir + "/output.bin").c_str(), outHost, dstBytes)) { + std::fprintf(stderr, "[ERROR] failed to write %s/output.bin\n", caseDir.c_str()); + rc = 1; + } + } + } + + if (srcDev) aclrtFree(srcDev); + if (dstDev) aclrtFree(dstDev); + if (outDev) aclrtFree(outDev); + if (srcHost) aclrtFreeHost(srcHost); + if (dstHost) aclrtFreeHost(dstHost); + if (outHost) aclrtFreeHost(outHost); + + if (rc == 0) + std::printf("[INFO] case %s done\n", tc.name); + return rc; +} + +int main(int argc, char *argv[]) { + const char *caseFilter = (argc > 1) ? argv[1] : nullptr; + + int rc = 0; + int deviceId = 0; + aclrtStream stream = nullptr; + + aclInit(nullptr); + if (const char *envDevice = std::getenv("ACL_DEVICE_ID")) { + deviceId = std::atoi(envDevice); + } + aclrtSetDevice(deviceId); + aclrtCreateStream(&stream); + + for (size_t i = 0; i < kNumCases; ++i) { + if (caseFilter != nullptr && std::strcmp(kCases[i].name, caseFilter) != 0) { + continue; + } + int ret = RunCase(kCases[i], deviceId, stream); + if (ret != 0) { + std::fprintf(stderr, "[ERROR] case %s failed\n", kCases[i].name); + rc = 1; + break; + } + } + + if (stream != nullptr) + aclrtDestroyStream(stream); + aclrtResetDevice(deviceId); + aclFinalize(); + + return rc; +} \ No newline at end of file diff --git a/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/tinsert_vec.pto b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/tinsert_vec.pto new file mode 100644 index 0000000000..29d023c19f --- /dev/null +++ b/test/tilelang_st/npu/a5/src/st/testcase/tinsert_vec/tinsert_vec.pto @@ -0,0 +1,305 @@ +// 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 OR FITNESS FOR A PARTICULAR PURPOSE. +// See LICENSE in the root of the software and copyright holder. + +// ST kernels for pto.tinsert Vec->Vec ND vkernel templates. +// Pattern: tload(src->src_tile) + tload(dst->dst_tile) + tinsert(src_tile,idx->dst_tile) + tstore(dst_tile->out) + +module attributes {pto.target_arch = "a5", pto.kernel_kind = #pto.kernel_kind} { + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) f16 -> dst=(32,32) f16, index=(0,0) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx00( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx0 = arith.constant 0 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> -> !pto.partition_tensor_view<1x1x1x32x32xf16> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> -> !pto.partition_tensor_view<1x1x1x32x32xf16> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx0, %idx0 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + return + } + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) f16 -> dst=(32,32) f16, index=(8,16) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_f16_16x16_into_32x32_idx816( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx8 = arith.constant 8 : index + %idx16 = arith.constant 16 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf16> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf16> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf16> -> !pto.partition_tensor_view<1x1x1x16x16xf16> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> -> !pto.partition_tensor_view<1x1x1x32x32xf16> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf16> -> !pto.partition_tensor_view<1x1x1x32x32xf16> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xf16>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx8, %idx16 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xf16>) + return + } + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) f32 -> dst=(32,32) f32, index=(0,0) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_f32_16x16_into_32x32_idx00( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx0 = arith.constant 0 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xf32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf32> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xf32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xf32> -> !pto.partition_tensor_view<1x1x1x16x16xf32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf32> -> !pto.partition_tensor_view<1x1x1x32x32xf32> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xf32> -> !pto.partition_tensor_view<1x1x1x32x32xf32> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xf32>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xf32>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx0, %idx0 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xf32>) + return + } + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) bf16 -> dst=(32,32) bf16, index=(0,0) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_bf16_16x16_into_32x32_idx00( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx0 = arith.constant 0 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xbf16> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xbf16> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xbf16> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xbf16> -> !pto.partition_tensor_view<1x1x1x16x16xbf16> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xbf16> -> !pto.partition_tensor_view<1x1x1x32x32xbf16> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xbf16> -> !pto.partition_tensor_view<1x1x1x32x32xbf16> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xbf16>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xbf16>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx0, %idx0 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xbf16>) + return + } + + // ------------------------------------------------------------------------- + // Vec -> Vec ND: src=(16,16) i32 -> dst=(32,32) i32, index=(0,0) + // ------------------------------------------------------------------------- + func.func @TINSERT_vec2vec_nd_i32_16x16_into_32x32_idx00( + %src_ptr: !pto.ptr, %dst_ptr: !pto.ptr, %out_ptr: !pto.ptr + ) attributes {pto.kernel} { + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c16 = arith.constant 16 : index + %c32 = arith.constant 32 : index + %c256 = arith.constant 256 : index + %c1024 = arith.constant 1024 : index + %idx0 = arith.constant 0 : index + + %src_view = pto.make_tensor_view %src_ptr, + shape = [%c1, %c1, %c1, %c16, %c16], + strides = [%c256, %c256, %c256, %c16, %c1] + : !pto.tensor_view<1x1x1x16x16xi32> + %dst_view = pto.make_tensor_view %dst_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xi32> + %out_view = pto.make_tensor_view %out_ptr, + shape = [%c1, %c1, %c1, %c32, %c32], + strides = [%c1024, %c1024, %c1024, %c32, %c1] + : !pto.tensor_view<1x1x1x32x32xi32> + + %src_part = pto.partition_view %src_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c16, %c16] + : !pto.tensor_view<1x1x1x16x16xi32> -> !pto.partition_tensor_view<1x1x1x16x16xi32> + %dst_part = pto.partition_view %dst_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xi32> -> !pto.partition_tensor_view<1x1x1x32x32xi32> + %out_part = pto.partition_view %out_view, + offsets = [%c0, %c0, %c0, %c0, %c0], + sizes = [%c1, %c1, %c1, %c32, %c32] + : !pto.tensor_view<1x1x1x32x32xi32> -> !pto.partition_tensor_view<1x1x1x32x32xi32> + + %src_tile = pto.alloc_tile : !pto.tile_buf + %dst_tile = pto.alloc_tile : !pto.tile_buf + + pto.tload ins(%src_part : !pto.partition_tensor_view<1x1x1x16x16xi32>) + outs(%src_tile : !pto.tile_buf) + pto.tload ins(%dst_part : !pto.partition_tensor_view<1x1x1x32x32xi32>) + outs(%dst_tile : !pto.tile_buf) + + pto.tinsert ins(%src_tile, %idx0, %idx0 : + !pto.tile_buf, + index, index) + outs(%dst_tile : !pto.tile_buf) + + pto.tstore ins(%dst_tile : !pto.tile_buf) + outs(%out_part : !pto.partition_tensor_view<1x1x1x32x32xi32>) + return + } + +} diff --git a/tilelang-dsl/python/tilelang_dsl/lowering.py b/tilelang-dsl/python/tilelang_dsl/lowering.py index 65b8bd258d..71e00bf5f4 100644 --- a/tilelang-dsl/python/tilelang_dsl/lowering.py +++ b/tilelang-dsl/python/tilelang_dsl/lowering.py @@ -2786,6 +2786,8 @@ def _render_mte_l0c_store( op_text = f"pto.{expr.name} " + ", ".join(value.name for value in operands) if clause_parts: op_text += ", " + ", ".join(clause_parts) + else: + op_text += "," into.append(self._indent(indent) + op_text + " : " + ", ".join(type_parts)) def _lower_cube_loop_groups( diff --git a/tilelang-dsl/python/tilelang_dsl/types.py b/tilelang-dsl/python/tilelang_dsl/types.py index ea688cd048..2479f5ca60 100644 --- a/tilelang-dsl/python/tilelang_dsl/types.py +++ b/tilelang-dsl/python/tilelang_dsl/types.py @@ -667,6 +667,8 @@ def s_layout(self) -> SLayout: @property def s_fractal_size(self) -> int: value = dict(self.fields).get("s_fractal_size", 512) + if value == 0: + value = 512 if isinstance(value, bool) or not isinstance(value, int): raise TypeError("TileConfig.s_fractal_size must be an integer") return value