From b5d8c6b416fe456f0e7c766223a06d34cbb36342 Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 26 Aug 2026 09:55:43 +0800 Subject: [PATCH 01/12] Adapt conv tile. --- include/PTO/IR/PTOAttrs.td | 41 +++++ include/PTO/IR/PTOOps.td | 12 +- include/PTO/IR/PTOTypeDefs.td | 30 ++++ lib/PTO/IR/PTO.cpp | 120 ++++++++++++++- lib/PTO/IR/PTOTypeDefs.cpp | 169 +++++++++++++++++++++ lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 23 +++ lib/PTO/Transforms/PTOToEmitC.cpp | 115 +++++++++++++- lib/PTO/Transforms/Utils.cpp | 54 ++++++- 8 files changed, 546 insertions(+), 18 deletions(-) diff --git a/include/PTO/IR/PTOAttrs.td b/include/PTO/IR/PTOAttrs.td index 283d9ab105..19f6388ae1 100644 --- a/include/PTO/IR/PTOAttrs.td +++ b/include/PTO/IR/PTOAttrs.td @@ -504,6 +504,47 @@ def PTO_LayoutAttr : PTO_Attr<"Layout", "layout"> { }]; } +//===----------------------------------------------------------------------===// +// ConvTile layout +//===----------------------------------------------------------------------===// + +def PTO_ConvLayout_NC1HWC0 : + I32EnumAttrCase<"NC1HWC0", 0, "nc1hwc0">; +def PTO_ConvLayout_NDC1HWC0 : + I32EnumAttrCase<"NDC1HWC0", 1, "ndc1hwc0">; +def PTO_ConvLayout_FRACTAL_Z : + I32EnumAttrCase<"FRACTAL_Z", 2, "fractal_z">; +def PTO_ConvLayout_FRACTAL_Z_3D : + I32EnumAttrCase<"FRACTAL_Z_3D", 3, "fractal_z_3d">; +def PTO_ConvLayout_NCHW : + I32EnumAttrCase<"NCHW", 4, "nchw">; +def PTO_ConvLayout_NHWC : + I32EnumAttrCase<"NHWC", 5, "nhwc">; +def PTO_ConvLayout_GNCHW : + I32EnumAttrCase<"GNCHW", 6, "gnchw">; +def PTO_ConvLayout_GNC1HWC0 : + I32EnumAttrCase<"GNC1HWC0", 7, "gnc1hwc0">; + +def PTO_ConvLayoutEnum : PTO_I32Enum< + "ConvLayout", "PTO ConvTile storage layout", [ + PTO_ConvLayout_NC1HWC0, + PTO_ConvLayout_NDC1HWC0, + PTO_ConvLayout_FRACTAL_Z, + PTO_ConvLayout_FRACTAL_Z_3D, + PTO_ConvLayout_NCHW, + PTO_ConvLayout_NHWC, + PTO_ConvLayout_GNCHW, + PTO_ConvLayout_GNC1HWC0 + ]>; + +def PTO_ConvLayoutAttr : PTO_Attr<"ConvLayout", "conv_layout"> { + let parameters = (ins EnumParameter:$value); + let assemblyFormat = "`<` params `>`"; + let description = [{ + Physical layout carried by a PTO ConvTile. + }]; +} + //===----------------------------------------------------------------------===// // Function and Module Core Type //===----------------------------------------------------------------------===// diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index c746ec6781..349360b0c8 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -35,15 +35,17 @@ include "mlir/Interfaces/ViewLikeInterface.td" //===----------------------------------------------------------------------===// def PTODpsType : - AnyTypeOf<[AnyRankedTensor, PartitionTensorViewType, TileBufType]>; + AnyTypeOf<[AnyRankedTensor, PartitionTensorViewType, TileBufType, + ConvTileType]>; def PTOPipeEntryType : - AnyTypeOf<[AnyRankedTensor, TensorViewType, TileBufType], - "TensorView, TileBuf, or Tensor">; + AnyTypeOf<[AnyRankedTensor, TensorViewType, TileBufType, ConvTileType], + "TensorView, TileBuf, ConvTile, or Tensor">; def PTOCommType : AnyTypeOf<[AnyRankedTensor, TensorViewType, PartitionTensorViewType, - TileBufType], "TensorView, PartitionTensorView, TileBuf, or Tensor">; + TileBufType, ConvTileType], + "TensorView, PartitionTensorView, TileBuf, ConvTile, or Tensor">; def PtrOrMemRef : AnyTypeOf<[PtrType, AnyMemRef], "Ptr or MemRef">; @@ -331,7 +333,7 @@ def AllocTileOp : PTO_Op<"alloc_tile", [AttrSizedOperandSegments]> { Optional:$valid_col ); - let results = (outs TileBufType:$result); + let results = (outs AnyTypeOf<[TileBufType, ConvTileType]>:$result); let assemblyFormat = [{ (`addr` `=` $addr^)? diff --git a/include/PTO/IR/PTOTypeDefs.td b/include/PTO/IR/PTOTypeDefs.td index 569530185b..75e9b40a69 100644 --- a/include/PTO/IR/PTOTypeDefs.td +++ b/include/PTO/IR/PTOTypeDefs.td @@ -228,6 +228,36 @@ def TileBufType : TypeDef { }]; } +def ConvTileType : TypeDef { + let mnemonic = "conv_tile"; + let summary = "A convolution tile buffer with an explicit physical capacity"; + let description = [{ + Represents the PTO-ISA ConvTile object. Shape describes the logical + convolution dimensions, while bufferSize is the number of storage elements + reserved by the hardware tile object. + }]; + + let parameters = (ins + ArrayRefParameter<"int64_t">:$shape, + "mlir::Type":$elementType, + "mlir::Attribute":$memorySpace, + "int64_t":$bufferSize, + "mlir::pto::ConvLayoutAttr":$layout + ); + + let hasCustomAssemblyFormat = 1; + + let extraClassDeclaration = [{ + int64_t getRank() const { return getShape().size(); } + int64_t getDimSize(unsigned idx) const { return getShape()[idx]; } + bool hasDynamicShape() const { + return llvm::any_of(getShape(), [](int64_t dim) { + return dim == mlir::ShapedType::kDynamic; + }); + } + }]; +} + // ============================================================================= // MultiTileBufType // ============================================================================= diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 616cd8dbba..734b6d26c8 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -293,6 +293,10 @@ static int64_t getPTOTypeRank(Type type) { return tileBufTy.getRank(); } + if (auto convTileTy = dyn_cast(type)) { + return convTileTy.getRank(); + } + // 3. 不支持的类型 return -1; } @@ -3585,14 +3589,50 @@ static LogicalResult verifyConstantLocalAddress(Operation *op, Value addr, } LogicalResult AllocTileOp::verify() { - auto ty = getResult().getType(); // TileBufType + auto ty = getResult().getType(); - if (failed(verifyTileBufLayoutConstraints(*this, ty, "result"))) { + if (auto convTy = dyn_cast(ty)) { + const bool invalidRank = convTy.getRank() == 0 || convTy.getRank() > 6; + if (invalidRank) { + return emitOpError("ConvTile result rank must be between 1 and 6"); + } + for (int64_t dim : convTy.getShape()) { + if (dim <= 0) { + return emitOpError("ConvTile result dimensions must be positive"); + } + } + const bool invalidBuffer = convTy.getBufferSize() <= 0; + if (invalidBuffer) { + return emitOpError("ConvTile buffer size must be positive"); + } + const bool invalidElementSize = + getElemByteSize(convTy.getElementType()) == 0; + if (invalidElementSize) { + return emitOpError("ConvTile element type must have a byte size"); + } + const bool hasValidOperands = getValidRow() || getValidCol(); + if (hasValidOperands) { + return emitOpError( + "ConvTile allocation does not accept valid_row or valid_col operands"); + } + if (failed(verifyConstantLocalAddress(getOperation(), getAddr(), + convTy.getMemorySpace()))) { + return failure(); + } + return success(); + } + + auto tileTy = dyn_cast(ty); + if (!tileTy) { + return emitOpError("result must be !pto.tile_buf or !pto.conv_tile"); + } + + if (failed(verifyTileBufLayoutConstraints(*this, tileTy, "result"))) { return failure(); } if (failed(verifyConstantLocalAddress(getOperation(), getAddr(), - ty.getMemorySpace()))) { + tileTy.getMemorySpace()))) { return failure(); } @@ -3601,7 +3641,7 @@ LogicalResult AllocTileOp::verify() { bool hasVC = getValidCol() != nullptr; // type 上的 validShape - auto vs = ty.getValidShape(); + auto vs = tileTy.getValidShape(); if (vs.size() != 2) { return emitOpError("result tile_buf must have rank-2 validShape"); } @@ -3789,9 +3829,21 @@ LogicalResult TAssignOp::verify() { return emitOpError("result type must match tile operand type"); } + if (auto convTy = dyn_cast(getTile().getType())) { + const bool invalidConvTile = + convTy.getBufferSize() <= 0 || convTy.getRank() == 0 || + convTy.getRank() > 6; + if (invalidConvTile) { + return emitOpError("expects a valid ConvTile type"); + } + return verifyConstantLocalAddress(getOperation(), getAddr(), + convTy.getMemorySpace()); + } + auto tileTy = dyn_cast(getTile().getType()); if (!tileTy) { - return emitOpError("expects tile operand and result to be !pto.tile_buf"); + return emitOpError( + "expects tile operand and result to be !pto.tile_buf or !pto.conv_tile"); } if (failed(verifyConstantLocalAddress(getOperation(), getAddr(), @@ -3803,6 +3855,42 @@ LogicalResult TAssignOp::verify() { } LogicalResult TLoadOp::verify() { + if (auto convDst = dyn_cast(getDst().getType())) { + auto srcPart = dyn_cast(getSrc().getType()); + if (!srcPart) { + return emitOpError( + "ConvTile tload expects src to be !pto.partition_tensor_view"); + } + const bool invalidRank = convDst.getRank() == 0 || convDst.getRank() > 6; + if (invalidRank) { + return emitOpError("ConvTile tload dst rank must be between 1 and 6"); + } + const bool invalidCapacity = + convDst.getBufferSize() <= 0 || + getElemByteSize(convDst.getElementType()) == 0; + if (invalidCapacity) { + return emitOpError("ConvTile tload dst must have a positive buffer and " + "a byte-sized element type"); + } + auto dstSpace = getPTOMemorySpaceEnum(convDst); + if (!dstSpace || *dstSpace != AddressSpace::MAT) { + return emitOpError("ConvTile tload dst must use loc=mat"); + } + for (int64_t dim : srcPart.getShape()) { + if (dim != ShapedType::kDynamic && dim <= 0) { + return emitOpError() << "expects src shape dimension to be positive"; + } + } + const bool mismatchedElementSize = + getElemByteSize(srcPart.getElementType()) != + getElemByteSize(convDst.getElementType()); + if (mismatchedElementSize) { + return emitOpError( + "ConvTile tload src and dst must have the same element size"); + } + return success(); + } + auto verifyCommon = [&](bool allowLowPrecision) -> FailureOr> { @@ -4871,6 +4959,21 @@ static LogicalResult verifyCommPingPongSameType(Operation *op, Value ping, } static std::optional getStaticByteSize(Type ty) { + if (auto conv = dyn_cast(ty)) { + uint64_t elemBytes = getElemByteSize(conv.getElementType()); + const bool invalidCapacity = elemBytes == 0 || conv.getBufferSize() <= 0; + if (invalidCapacity) { + return std::nullopt; + } + uint64_t bufferSize = static_cast(conv.getBufferSize()); + const bool overflows = + bufferSize > std::numeric_limits::max() / elemBytes; + if (overflows) { + return std::nullopt; + } + return bufferSize * elemBytes; + } + SmallVector shape = getShapeVec(ty); if (shape.empty()) { return std::nullopt; @@ -4920,6 +5023,13 @@ static std::optional getPTOMemorySpaceEnum(Type ty) { } return std::nullopt; } + if (auto conv = dyn_cast(ty)) { + if (auto as = + dyn_cast_or_null(conv.getMemorySpace())) { + return as.getAddressSpace(); + } + return std::nullopt; + } return std::nullopt; } diff --git a/lib/PTO/IR/PTOTypeDefs.cpp b/lib/PTO/IR/PTOTypeDefs.cpp index fe0c9b4daa..b9773f8f2b 100644 --- a/lib/PTO/IR/PTOTypeDefs.cpp +++ b/lib/PTO/IR/PTOTypeDefs.cpp @@ -190,6 +190,19 @@ static std::optional resolveTileBufMemorySpace(StringRef locStr) { .Default(::std::nullopt); } +static std::optional resolveConvLayout(StringRef layoutStr) { + return ::llvm::StringSwitch<::std::optional>(layoutStr) + .Case("nc1hwc0", ConvLayout::NC1HWC0) + .Case("ndc1hwc0", ConvLayout::NDC1HWC0) + .Case("fractal_z", ConvLayout::FRACTAL_Z) + .Case("fractal_z_3d", ConvLayout::FRACTAL_Z_3D) + .Case("nchw", ConvLayout::NCHW) + .Case("nhwc", ConvLayout::NHWC) + .Case("gnchw", ConvLayout::GNCHW) + .Case("gnc1hwc0", ConvLayout::GNC1HWC0) + .Default(::std::nullopt); +} + static BLayout resolveTileBufBLayout(MLIRContext *context, AddressSpace memorySpace, BLayout parsedLayout) { @@ -757,6 +770,162 @@ void mlir::pto::TileBufType::print(mlir::AsmPrinter &printer) const { printer << ">"; } +// ---- ConvTileType custom asm ---- +// !pto.conv_tile +Type ConvTileType::parse(AsmParser &parser) { + if (failed(parser.parseLess())) { + return Type(); + } + + std::string locStr; + std::string layoutStr; + SmallVector shape; + Type dtype; + int64_t bufferSize = 0; + + ParseResult parseResult = parser.parseKeywordOrString(&locStr); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseComma(); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseKeyword("buffer"); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseEqual(); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseInteger(bufferSize); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseComma(); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseKeyword("layout"); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseEqual(); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseKeywordOrString(&layoutStr); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseComma(); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseKeyword("shape"); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseEqual(); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseDimensionList(shape, /*allowDynamic=*/false); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseType(dtype); + if (!parseResult.succeeded()) { + return Type(); + } + parseResult = parser.parseGreater(); + if (!parseResult.succeeded()) { + return Type(); + } + + auto emitError = [&]() -> InFlightDiagnostic { + return parser.emitError(parser.getNameLoc()); + }; + auto memorySpace = resolveTileBufMemorySpace(locStr); + if (!memorySpace.has_value()) { + emitError() << "unknown ConvTile loc: " << locStr; + return Type(); + } + auto layout = resolveConvLayout(layoutStr); + if (!layout.has_value()) { + emitError() << "unknown ConvTile layout: " << layoutStr; + return Type(); + } + if (bufferSize <= 0) { + emitError() << "ConvTile buffer must be positive"; + return Type(); + } + const bool invalidRank = shape.empty() || shape.size() > 6; + if (invalidRank) { + emitError() << "ConvTile shape rank must be between 1 and 6"; + return Type(); + } + for (int64_t dim : shape) { + if (dim <= 0) { + emitError() << "ConvTile shape dimensions must be positive"; + return Type(); + } + } + + auto memorySpaceAttr = AddressSpaceAttr::get(parser.getContext(), + memorySpace.value()); + auto layoutAttr = ConvLayoutAttr::get(parser.getContext(), layout.value()); + return ConvTileType::get(parser.getContext(), shape, dtype, memorySpaceAttr, + bufferSize, layoutAttr); +} + +void mlir::pto::ConvTileType::print(mlir::AsmPrinter &printer) const { + auto memorySpace = + llvm::dyn_cast_or_null(getMemorySpace()); + auto layout = getLayout(); + if (!memorySpace || !layout) { + printer << ""; + return; + } + + auto layoutName = [&]() -> llvm::StringRef { + switch (layout.getValue()) { + case ConvLayout::NC1HWC0: + return "nc1hwc0"; + case ConvLayout::NDC1HWC0: + return "ndc1hwc0"; + case ConvLayout::FRACTAL_Z: + return "fractal_z"; + case ConvLayout::FRACTAL_Z_3D: + return "fractal_z_3d"; + case ConvLayout::NCHW: + return "nchw"; + case ConvLayout::NHWC: + return "nhwc"; + case ConvLayout::GNCHW: + return "gnchw"; + case ConvLayout::GNC1HWC0: + return "gnc1hwc0"; + } + return "unknown"; + }; + + printer << "<" << stringifyLocFromMemorySpace(memorySpace) + << ", buffer=" << getBufferSize() + << ", layout=" << layoutName() + << ", shape="; + for (auto [index, dim] : llvm::enumerate(getShape())) { + if (index != 0) { + printer << "x"; + } + printTileBufDim(printer, dim); + } + printer << "x"; + printer.printType(getElementType()); + printer << ">"; +} + // ---- MultiTileBufType custom asm ---- LogicalResult MultiTileBufType::verify( function_ref emitError, diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 9b00833b9d..3880371d47 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -107,6 +107,14 @@ static std::optional getBufferAddressSpace(Type type) { return std::nullopt; } + if (auto convType = dyn_cast(type)) { + if (auto attr = + dyn_cast_or_null(convType.getMemorySpace())) { + return attr.getAddressSpace(); + } + return std::nullopt; + } + if (auto multiType = dyn_cast(type)) { return getBufferAddressSpace(multiType.getSlotType()); } @@ -248,6 +256,21 @@ static FailureOr computeStaticBufferBytes(Value value) { if (auto tileType = dyn_cast(value.getType())) { return computeTileBytes(tileType); } + if (auto convType = dyn_cast(value.getType())) { + uint64_t elemBytes = getPTOStorageElemByteSize(convType.getElementType()); + const bool invalidCapacity = + elemBytes == 0 || convType.getBufferSize() <= 0; + if (invalidCapacity) { + return failure(); + } + uint64_t bufferSize = static_cast(convType.getBufferSize()); + const bool overflows = + bufferSize > std::numeric_limits::max() / elemBytes; + if (overflows) { + return failure(); + } + return bufferSize * elemBytes; + } if (auto multiType = dyn_cast(value.getType())) { return computeTileBytes(multiType.getSlotType()); } diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 045c286f82..22b2bfe05a 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -843,6 +843,56 @@ static std::optional getEmitCTileTypeString(pto::TileBufType type) tileBufCompactToken(configAttr) + ">"; } +static StringRef convLayoutToken(pto::ConvLayout layout) { + switch (layout) { + case pto::ConvLayout::NC1HWC0: + return "Layout::NC1HWC0"; + case pto::ConvLayout::NDC1HWC0: + return "Layout::NDC1HWC0"; + case pto::ConvLayout::FRACTAL_Z: + return "Layout::FRACTAL_Z"; + case pto::ConvLayout::FRACTAL_Z_3D: + return "Layout::FRACTAL_Z_3D"; + case pto::ConvLayout::NCHW: + return "Layout::NCHW"; + case pto::ConvLayout::NHWC: + return "Layout::NHWC"; + case pto::ConvLayout::GNCHW: + return "Layout::GNCHW"; + case pto::ConvLayout::GNC1HWC0: + return "Layout::GNC1HWC0"; + } + return "Layout::NC1HWC0"; +} + +static std::optional +getEmitCConvTileTypeString(pto::ConvTileType type) { + auto memorySpace = + dyn_cast_or_null(type.getMemorySpace()); + auto layout = type.getLayout(); + const bool invalidType = + !memorySpace || !layout || type.getRank() == 0 || type.getRank() > 6 || + type.getBufferSize() <= 0; + if (invalidType) { + return std::nullopt; + } + + std::string shape = "ConvTileShape<"; + for (auto [index, dim] : llvm::enumerate(type.getShape())) { + if (index != 0) { + shape += ", "; + } + shape += std::to_string(dim); + } + shape += ">"; + + return std::string("ConvTile<") + + tileRoleToken(type.getMemorySpace(), type.getElementType(), nullptr) + + ", " + getEmitCScalarTypeToken(type.getElementType()) + ", " + + std::to_string(type.getBufferSize()) + ", " + + convLayoutToken(layout.getValue()).str() + ", " + shape + ">"; +} + //===----------------------------------------------------------------------===// // Type Converter //===----------------------------------------------------------------------===// @@ -1017,10 +1067,19 @@ class PTOToEmitCTypeConverter : public TypeConverter { type.getShape()); }); - addConversion([Ctx](pto::TileBufType type) -> std::optional { + addConversion([Ctx](pto::TileBufType type) -> std::optional { auto typeString = getEmitCTileTypeString(type); - if (!typeString) + if (!typeString) { return std::nullopt; + } + return emitc::OpaqueType::get(Ctx, *typeString); + }); + + addConversion([Ctx](pto::ConvTileType type) -> std::optional { + auto typeString = getEmitCConvTileTypeString(type); + if (!typeString) { + return std::nullopt; + } return emitc::OpaqueType::get(Ctx, *typeString); }); @@ -12476,7 +12535,57 @@ struct PTOAllocTileToEmitC ConversionPatternRewriter &rewriter) const override { Location loc = op.getLoc(); MLIRContext *ctx = rewriter.getContext(); - auto tileTy = cast(op.getResult().getType()); + Type resultTy = op.getResult().getType(); + if (auto convTy = dyn_cast(resultTy)) { + auto convTypeString = getEmitCConvTileTypeString(convTy); + if (!convTypeString) { + return rewriter.notifyMatchFailure( + op, "invalid ConvTile type for EmitC conversion"); + } + Type convertedTy = getTypeConverter()->convertType(convTy); + if (!convertedTy) { + convertedTy = emitc::OpaqueType::get(ctx, *convTypeString); + } + Value tile = + rewriter + .create( + loc, getEmitCVariableResultType(convertedTy), + emitc::OpaqueAttr::get(ctx, "")) + .getResult(); + tile = loadEmitCVariableIfNeeded(rewriter, loc, tile); + + Value addr = adaptor.getAddr(); + if (addr) { + addr = peelUnrealized(addr); + auto u64Ty = emitc::OpaqueType::get(ctx, "uint64_t"); + const bool isPointer = + isa(addr.getType()) || + (isa(addr.getType()) && + cast(addr.getType()).getValue().ends_with("*")); + if (isPointer) { + auto rcU64 = + rewriter.getArrayAttr({emitc::OpaqueAttr::get(ctx, "uint64_t")}); + addr = rewriter + .create( + loc, u64Ty, "reinterpret_cast", ArrayAttr{}, rcU64, + ValueRange{addr}) + .getResult(0); + } else if (addr.getType() != u64Ty) { + addr = rewriter.create(loc, u64Ty, addr).getResult(); + } + rewriter.create( + loc, TypeRange{}, "TASSIGN", ArrayAttr{}, ArrayAttr{}, + ValueRange{tile, addr}); + } + rewriter.replaceOp(op, tile); + return success(); + } + + auto tileTy = dyn_cast(resultTy); + if (!tileTy) { + return rewriter.notifyMatchFailure( + op, "expected tile_buf or conv_tile result"); + } auto tileTypeString = getEmitCTileTypeString(tileTy); if (!tileTypeString) return rewriter.notifyMatchFailure( diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index 63fa540a6f..c62f6efe9a 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -278,6 +278,13 @@ std::optional GetBufferSpaceAttr(Value operand) { } return std::nullopt; } + if (auto convTy = dyn_cast(operand.getType())) { + if (auto memorySpaceAttr = dyn_cast_or_null( + convTy.getMemorySpace())) { + return memorySpaceAttr; + } + return std::nullopt; + } if (!llvm::isa(operand.getType())) { return std::nullopt; @@ -500,6 +507,21 @@ static std::optional getStaticTileBytes(TileBufType type) { return *elements * elemBytes; } +static std::optional getStaticConvTileBytes(ConvTileType type) { + unsigned elemBytes = getPTOStorageElemByteSize(type.getElementType()); + const bool invalidCapacity = elemBytes == 0 || type.getBufferSize() <= 0; + if (invalidCapacity) { + return std::nullopt; + } + uint64_t bufferSize = static_cast(type.getBufferSize()); + const bool overflows = + bufferSize > std::numeric_limits::max() / elemBytes; + if (overflows) { + return std::nullopt; + } + return bufferSize * elemBytes; +} + static std::optional getConstantAddress(Value value) { IntegerAttr attr; bool isInvalid = @@ -568,6 +590,14 @@ static std::optional getTileAddressSpace(TileBufType type) { return attr.getAddressSpace(); } +static std::optional getConvTileAddressSpace(ConvTileType type) { + auto attr = dyn_cast_or_null(type.getMemorySpace()); + if (!attr) { + return std::nullopt; + } + return attr.getAddressSpace(); +} + static std::optional getSubviewByteOffset(SubViewOp op, const SemanticRange &source) { bool hasInvalidRank = op.getOffsets().size() != kValue2; @@ -667,13 +697,27 @@ static std::optional makeTileRange(Value root, TileBufType type, static std::optional resolveAllocTileRange(AllocTileOp alloc) { auto tileType = dyn_cast(alloc.getResult().getType()); - std::optional bytes = - tileType ? getStaticTileBytes(tileType) : std::nullopt; - if (!tileType || !bytes) { + if (tileType) { + std::optional bytes = getStaticTileBytes(tileType); + if (!bytes) { + return std::nullopt; + } + return makeTileRange(alloc.getResult(), tileType, *bytes, + getConstantAddress(alloc.getAddr())); + } + + auto convType = dyn_cast(alloc.getResult().getType()); + if (!convType) { + return std::nullopt; + } + + std::optional bytes = getStaticConvTileBytes(convType); + if (!bytes) { return std::nullopt; } - return makeTileRange(alloc.getResult(), tileType, *bytes, - getConstantAddress(alloc.getAddr())); + return SemanticRange{ + alloc.getResult(), 0, *bytes, getConstantAddress(alloc.getAddr()), + getConvTileAddressSpace(convType), std::nullopt, std::nullopt, 0}; } static FailureOr> getMultiTileSlotBase( From ee511d84fe784793667e4ad3fc4ca88ec5b84c1c Mon Sep 17 00:00:00 2001 From: andodo Date: Sat, 29 Aug 2026 18:34:11 +0800 Subject: [PATCH 02/12] Add conv ops. --- include/PTO/IR/PTOAttrs.td | 72 ++++ include/PTO/IR/PTOOps.td | 63 ++- include/PTO/IR/PTOTypeDefs.td | 27 +- lib/PTO/IR/PTO.cpp | 627 ++++++++++++++++++++++++------ lib/PTO/IR/PTOAttrs.cpp | 459 ++++++++++++++++++++++ lib/PTO/IR/PTOTypeDefs.cpp | 344 ++++++++++------ lib/PTO/Transforms/PTOToEmitC.cpp | 188 +++++---- 7 files changed, 1424 insertions(+), 356 deletions(-) diff --git a/include/PTO/IR/PTOAttrs.td b/include/PTO/IR/PTOAttrs.td index 19f6388ae1..71c2f9b117 100644 --- a/include/PTO/IR/PTOAttrs.td +++ b/include/PTO/IR/PTOAttrs.td @@ -1278,6 +1278,78 @@ def TileBufConfigAttr : AttrDef { }]; } +// ---------- conv_tile_config ---------- +def ConvTileConfigAttr : AttrDef { + let mnemonic = "conv_tile_config"; + let parameters = (ins + "mlir::IntegerAttr":$fmapH, + "mlir::IntegerAttr":$fmapW, + ArrayRefParameter<"int64_t">:$padList, + "mlir::IntegerAttr":$filterH, + "mlir::IntegerAttr":$filterW, + "mlir::IntegerAttr":$dilationH, + "mlir::IntegerAttr":$dilationW, + "mlir::IntegerAttr":$strideH, + "mlir::IntegerAttr":$strideW, + "mlir::Attribute":$padValue, + "mlir::IntegerAttr":$channelSize, + "mlir::IntegerAttr":$repeatStride, + "mlir::IntegerAttr":$repeatTime, + "mlir::IntegerAttr":$repeatMode, + "mlir::IntegerAttr":$dstStride, + "mlir::IntegerAttr":$dstMposition, + "mlir::BoolAttr":$transpose + ); + + let hasCustomAssemblyFormat = 1; + + let builders = [ + AttrBuilder<(ins + "mlir::IntegerAttr":$fmapH, + "mlir::IntegerAttr":$fmapW, + ArrayRefParameter<"int64_t">:$padList, + "mlir::IntegerAttr":$filterH, + "mlir::IntegerAttr":$filterW, + "mlir::IntegerAttr":$dilationH, + "mlir::IntegerAttr":$dilationW, + "mlir::IntegerAttr":$strideH, + "mlir::IntegerAttr":$strideW, + "mlir::Attribute":$padValue, + "mlir::IntegerAttr":$channelSize, + "mlir::IntegerAttr":$repeatStride, + "mlir::IntegerAttr":$repeatTime, + "mlir::IntegerAttr":$repeatMode, + "mlir::IntegerAttr":$dstStride, + "mlir::IntegerAttr":$dstMposition, + "mlir::BoolAttr":$transpose + )> + ]; + + let extraClassDeclaration = [{ + static ConvTileConfigAttr getDefault(MLIRContext *ctx); + bool isDefault() const; + + static LogicalResult verify(function_ref emitError, + mlir::IntegerAttr fmapH, + mlir::IntegerAttr fmapW, + ArrayRef padList, + mlir::IntegerAttr filterH, + mlir::IntegerAttr filterW, + mlir::IntegerAttr dilationH, + mlir::IntegerAttr dilationW, + mlir::IntegerAttr strideH, + mlir::IntegerAttr strideW, + mlir::Attribute padValue, + mlir::IntegerAttr channelSize, + mlir::IntegerAttr repeatStride, + mlir::IntegerAttr repeatTime, + mlir::IntegerAttr repeatMode, + mlir::IntegerAttr dstStride, + mlir::IntegerAttr dstMposition, + mlir::BoolAttr transpose); + }]; +} + //===----------------------------------------------------------------------===// // QuantType //===----------------------------------------------------------------------===// diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index 349360b0c8..ff62dbf6a1 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -333,7 +333,10 @@ def AllocTileOp : PTO_Op<"alloc_tile", [AttrSizedOperandSegments]> { Optional:$valid_col ); - let results = (outs AnyTypeOf<[TileBufType, ConvTileType]>:$result); + // ConvTile values reuse the same allocation surface but carry a different + // type/metadata payload. The lowering selects the emitted C++ type from the + // result type. + let results = (outs AnyTypeOf<[TileBufType, ConvTileType], "TileBuf or ConvTile">:$result); let assemblyFormat = [{ (`addr` `=` $addr^)? @@ -1362,6 +1365,12 @@ def TMovOp : PTO_TOp<"tmov", [ return as.getAddressSpace(); return std::nullopt; } + if (auto ct = llvm::dyn_cast<::mlir::pto::ConvTileType>(ty)) { + if (auto as = llvm::dyn_cast_or_null<::mlir::pto::AddressSpaceAttr>( + ct.getMemorySpace())) + return as.getAddressSpace(); + return std::nullopt; + } return std::nullopt; }; @@ -1884,6 +1893,57 @@ def SetQuantVectorOp : PTO_Op<"set_quant_vector", [ }]; } +//===----------------------------------------------------------------------===// +// ConvTile / IMG2COL config ops +//===----------------------------------------------------------------------===// + +def SetFmatrixOp : PTO_Op<"set_fmatrix", [MemoryEffects<[MemWrite]>]> { + let summary = "Set FMATRIX registers from a ConvTile config"; + let arguments = (ins ConvTileType:$src); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = [{ + $src attr-dict `:` qualified(type($src)) + }]; +} + +def SetImg2colRptOp : PTO_Op<"set_img2col_rpt", [MemoryEffects<[MemWrite]>]> { + let summary = "Set IMG2COL repeat control from a ConvTile config"; + let arguments = (ins ConvTileType:$src); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = [{ + $src attr-dict `:` qualified(type($src)) + }]; +} + +def SetImg2colPaddingOp : PTO_Op<"set_img2col_padding", [MemoryEffects<[MemWrite]>]> { + let summary = "Set IMG2COL padding control from a ConvTile config"; + let arguments = (ins ConvTileType:$src); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = [{ + $src attr-dict `:` qualified(type($src)) + }]; +} + +def TImg2colOp : PTO_Op<"timg2col", [ + DeclareOpInterfaceMethods +]> { + let summary = "Image-to-column transform from ConvTile to TileBuf"; + let arguments = (ins + TileBufType:$dst, + ConvTileType:$src, + DefaultValuedOptionalAttr:$posM, + DefaultValuedOptionalAttr:$posK + ); + let results = (outs); + let hasVerifier = 1; + let assemblyFormat = [{ + $dst `,` $src attr-dict `:` qualified(type($dst)) `,` qualified(type($src)) + }]; +} + def ReserveBufferOp : PTO_Op<"reserve_buffer"> { let summary = "Reserve a local consumer slot buffer"; @@ -6241,7 +6301,6 @@ def TDeInterleaveOp: PTO_TOp<"tdeinterleave", [ ::mlir::Value getDst1() { return getDsts()[1]; } }]; } - def TRowProdOp: PTO_TOp<"trowprod", [ PTO_DpsInitOpInterface, OpPipeInterface, diff --git a/include/PTO/IR/PTOTypeDefs.td b/include/PTO/IR/PTOTypeDefs.td index 75e9b40a69..def534e6d7 100644 --- a/include/PTO/IR/PTOTypeDefs.td +++ b/include/PTO/IR/PTOTypeDefs.td @@ -230,19 +230,13 @@ def TileBufType : TypeDef { def ConvTileType : TypeDef { let mnemonic = "conv_tile"; - let summary = "A convolution tile buffer with an explicit physical capacity"; - let description = [{ - Represents the PTO-ISA ConvTile object. Shape describes the logical - convolution dimensions, while bufferSize is the number of storage elements - reserved by the hardware tile object. - }]; - let parameters = (ins ArrayRefParameter<"int64_t">:$shape, "mlir::Type":$elementType, - "mlir::Attribute":$memorySpace, - "int64_t":$bufferSize, - "mlir::pto::ConvLayoutAttr":$layout + "mlir::IntegerAttr":$bufferSize, + "mlir::pto::AddressSpaceAttr":$memorySpace, + "mlir::pto::LayoutAttr":$layout, + "mlir::pto::ConvTileConfigAttr":$config ); let hasCustomAssemblyFormat = 1; @@ -250,11 +244,16 @@ def ConvTileType : TypeDef { let extraClassDeclaration = [{ int64_t getRank() const { return getShape().size(); } int64_t getDimSize(unsigned idx) const { return getShape()[idx]; } - bool hasDynamicShape() const { - return llvm::any_of(getShape(), [](int64_t dim) { - return dim == mlir::ShapedType::kDynamic; - }); + int64_t getNumElements() const { + int64_t num = 1; + for (int64_t dim : getShape()) num *= dim; + return num; } + + int64_t getBufferSizeValue() const { return getBufferSize().getInt(); } + + mlir::pto::ConvTileConfigAttr getConfigAttr() const; + bool hasNonDefaultConfig() const; }]; } diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 734b6d26c8..fbf9d6bbf9 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -293,10 +293,6 @@ static int64_t getPTOTypeRank(Type type) { return tileBufTy.getRank(); } - if (auto convTileTy = dyn_cast(type)) { - return convTileTy.getRank(); - } - // 3. 不支持的类型 return -1; } @@ -3589,50 +3585,14 @@ static LogicalResult verifyConstantLocalAddress(Operation *op, Value addr, } LogicalResult AllocTileOp::verify() { - auto ty = getResult().getType(); - - if (auto convTy = dyn_cast(ty)) { - const bool invalidRank = convTy.getRank() == 0 || convTy.getRank() > 6; - if (invalidRank) { - return emitOpError("ConvTile result rank must be between 1 and 6"); - } - for (int64_t dim : convTy.getShape()) { - if (dim <= 0) { - return emitOpError("ConvTile result dimensions must be positive"); - } - } - const bool invalidBuffer = convTy.getBufferSize() <= 0; - if (invalidBuffer) { - return emitOpError("ConvTile buffer size must be positive"); - } - const bool invalidElementSize = - getElemByteSize(convTy.getElementType()) == 0; - if (invalidElementSize) { - return emitOpError("ConvTile element type must have a byte size"); - } - const bool hasValidOperands = getValidRow() || getValidCol(); - if (hasValidOperands) { - return emitOpError( - "ConvTile allocation does not accept valid_row or valid_col operands"); - } - if (failed(verifyConstantLocalAddress(getOperation(), getAddr(), - convTy.getMemorySpace()))) { - return failure(); - } - return success(); - } - - auto tileTy = dyn_cast(ty); - if (!tileTy) { - return emitOpError("result must be !pto.tile_buf or !pto.conv_tile"); - } + auto ty = getResult().getType(); // TileBufType - if (failed(verifyTileBufLayoutConstraints(*this, tileTy, "result"))) { + if (failed(verifyTileBufLayoutConstraints(*this, ty, "result"))) { return failure(); } if (failed(verifyConstantLocalAddress(getOperation(), getAddr(), - tileTy.getMemorySpace()))) { + ty.getMemorySpace()))) { return failure(); } @@ -3641,7 +3601,7 @@ LogicalResult AllocTileOp::verify() { bool hasVC = getValidCol() != nullptr; // type 上的 validShape - auto vs = tileTy.getValidShape(); + auto vs = ty.getValidShape(); if (vs.size() != 2) { return emitOpError("result tile_buf must have rank-2 validShape"); } @@ -3829,25 +3789,20 @@ LogicalResult TAssignOp::verify() { return emitOpError("result type must match tile operand type"); } - if (auto convTy = dyn_cast(getTile().getType())) { - const bool invalidConvTile = - convTy.getBufferSize() <= 0 || convTy.getRank() == 0 || - convTy.getRank() > 6; - if (invalidConvTile) { - return emitOpError("expects a valid ConvTile type"); - } - return verifyConstantLocalAddress(getOperation(), getAddr(), - convTy.getMemorySpace()); + auto tileTy = getTile().getType(); + if (!isa(tileTy)) { + return emitOpError("expects tile operand and result to be !pto.tile_buf or !pto.conv_tile"); } - auto tileTy = dyn_cast(getTile().getType()); - if (!tileTy) { - return emitOpError( - "expects tile operand and result to be !pto.tile_buf or !pto.conv_tile"); + Attribute memorySpace; + if (auto tb = dyn_cast(tileTy)) { + memorySpace = tb.getMemorySpace(); + } else if (auto ct = dyn_cast(tileTy)) { + memorySpace = ct.getMemorySpace(); } if (failed(verifyConstantLocalAddress(getOperation(), getAddr(), - tileTy.getMemorySpace()))) { + memorySpace))) { return failure(); } @@ -3855,52 +3810,16 @@ LogicalResult TAssignOp::verify() { } LogicalResult TLoadOp::verify() { - if (auto convDst = dyn_cast(getDst().getType())) { - auto srcPart = dyn_cast(getSrc().getType()); - if (!srcPart) { - return emitOpError( - "ConvTile tload expects src to be !pto.partition_tensor_view"); - } - const bool invalidRank = convDst.getRank() == 0 || convDst.getRank() > 6; - if (invalidRank) { - return emitOpError("ConvTile tload dst rank must be between 1 and 6"); - } - const bool invalidCapacity = - convDst.getBufferSize() <= 0 || - getElemByteSize(convDst.getElementType()) == 0; - if (invalidCapacity) { - return emitOpError("ConvTile tload dst must have a positive buffer and " - "a byte-sized element type"); - } - auto dstSpace = getPTOMemorySpaceEnum(convDst); - if (!dstSpace || *dstSpace != AddressSpace::MAT) { - return emitOpError("ConvTile tload dst must use loc=mat"); - } - for (int64_t dim : srcPart.getShape()) { - if (dim != ShapedType::kDynamic && dim <= 0) { - return emitOpError() << "expects src shape dimension to be positive"; - } - } - const bool mismatchedElementSize = - getElemByteSize(srcPart.getElementType()) != - getElemByteSize(convDst.getElementType()); - if (mismatchedElementSize) { - return emitOpError( - "ConvTile tload src and dst must have the same element size"); - } - return success(); - } - auto verifyCommon = [&](bool allowLowPrecision) - -> FailureOr> { + -> FailureOr> { auto srcPart = dyn_cast(getSrc().getType()); - auto dstTile = dyn_cast(getDst().getType()); - if (!srcPart || !dstTile) { - emitOpError("expects src to be !pto.partition_tensor_view and dst to be !pto.tile_buf"); + Type dstTy = getDst().getType(); + if (!srcPart || !isa(dstTy)) { + emitOpError("expects src to be !pto.partition_tensor_view and dst to be !pto.tile_buf or !pto.conv_tile"); return failure(); } - if (failed(verifyTileBufCommon(*this, dstTile, "dst", allowLowPrecision))) { + if (failed(verifyTileLikeCommon(*this, dstTy, "dst", allowLowPrecision))) { return failure(); } @@ -3911,14 +3830,14 @@ LogicalResult TLoadOp::verify() { return failure(); } } - auto dstValid = dstTile.getValidShape(); + auto dstValid = getValidShapeVec(dstTy); for (unsigned i = 0; i < dstValid.size(); ++i) { if (dstValid[i] != ShapedType::kDynamic && dstValid[i] < 0) { emitOpError() << "expects dst valid_shape[" << i << "] to be non-negative"; return failure(); } } - return std::make_pair(srcPart, dstTile); + return std::make_pair(srcPart, dstTy); }; auto verifyA2A3 = [&]() -> LogicalResult { @@ -3929,7 +3848,7 @@ LogicalResult TLoadOp::verify() { auto [srcPart, dstTile] = *common; Type srcElem = srcPart.getElementType(); - Type dstElem = dstTile.getElementType(); + Type dstElem = getElemTy(dstTile); if (isPTOLowPrecisionType(srcElem) || isPTOLowPrecisionType(dstElem)) { return emitOpError("expects A2/A3 tload low-precision element types to be unsupported"); } @@ -3958,7 +3877,7 @@ LogicalResult TLoadOp::verify() { auto [srcPart, dstTile] = *common; Type srcElem = srcPart.getElementType(); - Type dstElem = dstTile.getElementType(); + Type dstElem = getElemTy(dstTile); unsigned srcBytes = getElemByteSize(srcElem); unsigned dstBytes = getElemByteSize(dstElem); if (srcBytes != dstBytes) { @@ -3984,8 +3903,12 @@ LogicalResult TLoadOp::verify() { auto dstSpace = getPTOMemorySpaceEnum(dstTile); if (dstSpace && *dstSpace == pto::AddressSpace::VEC) { - int32_t bl = dstTile.getBLayoutValueI32(); - int32_t sl = dstTile.getSLayoutValueI32(); + auto dstTB = dyn_cast(dstTile); + if (!dstTB) { + return emitOpError("expects A5 tload vec dst to be a tile_buf"); + } + int32_t bl = dstTB.getBLayoutValueI32(); + int32_t sl = dstTB.getSLayoutValueI32(); bool isND = (bl == static_cast(pto::BLayout::RowMajor) && sl == static_cast(pto::SLayout::NoneBox)); bool isDN = (bl == static_cast(pto::BLayout::ColMajor) && @@ -4003,6 +3926,52 @@ LogicalResult TLoadOp::verify() { return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } +LogicalResult mlir::pto::SetFmatrixOp::verify() { + return verifyConvTileCommon(*this, getSrc().getType(), "src", + /*allowLowPrecision=*/true); +} + +LogicalResult mlir::pto::SetImg2colRptOp::verify() { + return verifyConvTileCommon(*this, getSrc().getType(), "src", + /*allowLowPrecision=*/true); +} + +LogicalResult mlir::pto::SetImg2colPaddingOp::verify() { + return verifyConvTileCommon(*this, getSrc().getType(), "src", + /*allowLowPrecision=*/true); +} + +LogicalResult mlir::pto::TImg2colOp::verify() { + if (failed(verifyTileBufCommon(*this, getDst().getType(), "dst", + /*allowLowPrecision=*/true))) { + return failure(); + } + if (failed(verifyConvTileCommon(*this, getSrc().getType(), "src", + /*allowLowPrecision=*/true))) { + return failure(); + } + if (getElemTy(getDst().getType()) != getElemTy(getSrc().getType())) { + return emitOpError() << "expects src and dst to have the same element type"; + } + + auto checkPos = [&](IntegerAttr posAttr, StringRef name) -> LogicalResult { + if (!posAttr || !posAttr.getType().isSignlessInteger(32)) { + return emitOpError() << "expects " << name << " to be an i32 attr"; + } + int64_t value = posAttr.getInt(); + if (value < 0 || value > std::numeric_limits::max()) { + return emitOpError() << "expects " << name + << " to fit in an unsigned 16-bit integer"; + } + return success(); + }; + if (failed(checkPos(getPosM(), "posM")) || failed(checkPos(getPosK(), "posK"))) { + return failure(); + } + + return success(); +} + LogicalResult TPrefetchOp::verify() { auto verifyImpl = [&](bool allowLowPrecision) -> LogicalResult { Type srcTy = getSrc().getType(); @@ -4701,7 +4670,7 @@ static bool isPTOShapedLike(Type ty) { } static bool isTileLikeType(Type ty) { - return isa(ty); + return isa(ty); } static Type getElemTy(Type ty) { @@ -4714,6 +4683,9 @@ static Type getElemTy(Type ty) { if (auto tb = mlir::dyn_cast(ty)) { return tb.getElementType(); } + if (auto ct = mlir::dyn_cast(ty)) { + return ct.getElementType(); + } if (auto tv = mlir::dyn_cast(ty)) { return tv.getElementType(); } @@ -4731,6 +4703,9 @@ static SmallVector getShapeVec(Type ty) { if (auto tb = mlir::dyn_cast(ty)) { return SmallVector(tb.getShape().begin(), tb.getShape().end()); } + if (auto ct = mlir::dyn_cast(ty)) { + return SmallVector(ct.getShape().begin(), ct.getShape().end()); + } if (auto tv = mlir::dyn_cast(ty)) { return SmallVector(tv.getShape().begin(), tv.getShape().end()); } @@ -4959,21 +4934,6 @@ static LogicalResult verifyCommPingPongSameType(Operation *op, Value ping, } static std::optional getStaticByteSize(Type ty) { - if (auto conv = dyn_cast(ty)) { - uint64_t elemBytes = getElemByteSize(conv.getElementType()); - const bool invalidCapacity = elemBytes == 0 || conv.getBufferSize() <= 0; - if (invalidCapacity) { - return std::nullopt; - } - uint64_t bufferSize = static_cast(conv.getBufferSize()); - const bool overflows = - bufferSize > std::numeric_limits::max() / elemBytes; - if (overflows) { - return std::nullopt; - } - return bufferSize * elemBytes; - } - SmallVector shape = getShapeVec(ty); if (shape.empty()) { return std::nullopt; @@ -5023,9 +4983,8 @@ static std::optional getPTOMemorySpaceEnum(Type ty) { } return std::nullopt; } - if (auto conv = dyn_cast(ty)) { - if (auto as = - dyn_cast_or_null(conv.getMemorySpace())) { + if (auto ct = dyn_cast(ty)) { + if (auto as = dyn_cast_or_null(ct.getMemorySpace())) { return as.getAddressSpace(); } return std::nullopt; @@ -5394,6 +5353,75 @@ static bool isA5SupportedTCvtPair(Type srcElem, Type dstElem) { return true; } +static LogicalResult verifyConvTileCommon(Operation *op, Type ty, StringRef name, + bool allowLowPrecision = false) { + auto ct = dyn_cast(ty); + if (!ct) { + return op->emitOpError() << "expects " << name << " to be a !pto.conv_tile"; + } + + auto as = getPTOMemorySpaceEnum(ty); + if (!as || *as != pto::AddressSpace::MAT) { + return op->emitOpError() << "expects " << name + << " to be in the mat address space"; + } + + if (!allowLowPrecision && isPTOLowPrecisionType(ct.getElementType())) { + return op->emitOpError() << name << ": dtype " << ct.getElementType() + << " is not supported by this op yet"; + } + + auto shape = getShapeVec(ty); + if (shape.empty() || shape.size() > 6) { + return op->emitOpError() << "expects " << name + << " to have a rank in [1, 6]"; + } + for (unsigned i = 0; i < shape.size(); ++i) { + if (shape[i] <= 0) { + return op->emitOpError() << "expects " << name << " shape[" << i + << "] to be positive"; + } + } + + auto bufferSize = dyn_cast(ct.getBufferSize()); + if (!bufferSize) { + return op->emitOpError() << "expects " << name + << " to have a signless integer buffer_size"; + } + if (bufferSize.getInt() <= 0) { + return op->emitOpError() << "expects " << name + << " buffer_size to be positive"; + } + + int64_t logicalSize = 1; + for (int64_t dim : shape) { + if (logicalSize > std::numeric_limits::max() / dim) { + return op->emitOpError() << "cannot compute logical size of " << name + << " without overflow"; + } + logicalSize *= dim; + } + if (bufferSize.getInt() < logicalSize) { + return op->emitOpError() << "expects " << name + << " buffer_size to cover at least " + << logicalSize << " elements"; + } + + return success(); +} + +static LogicalResult verifyTileLikeCommon(Operation *op, Type ty, StringRef name, + bool allowLowPrecision = false) { + if (isa(ty)) { + return verifyTileBufCommon(op, ty, name, allowLowPrecision); + } + if (isa(ty)) { + return verifyConvTileCommon(op, ty, name, allowLowPrecision); + } + return op->emitOpError() << "expects " << name + << " to be a !pto.tile_buf or !pto.conv_tile"; +} + static LogicalResult verifyTileBufCommon(Operation *op, Type ty, StringRef name, bool allowLowPrecision) { auto tb = dyn_cast(ty); @@ -8042,6 +8070,17 @@ static bool isRowMajorNoneBoxND(pto::TileBufType ty) { ty.getSLayoutValueI32() == static_cast(pto::SLayout::NoneBox); } +static bool isConvTileExtractElem(Type ty) { + if (ty.isF16() || ty.isBF16() || ty.isF32()) { + return true; + } + if (auto intTy = dyn_cast(ty)) { + unsigned width = intTy.getWidth(); + return width == 8 || width == 16 || width == 32; + } + return false; +} + struct TExtractCommon { Type srcTy; Type dstTy; @@ -8185,6 +8224,54 @@ static LogicalResult verifyTExtractA2A3(TExtractOp op) { const bool hasFp = static_cast(op.getFp()); const bool hasPreQuantScalar = static_cast(op.getPreQuantScalar()); const bool hasRelu = op.getReluPreMode() != pto::ReluPreMode::NoRelu; + if (isa(getSrc().getType())) { + Type srcTy = getSrc().getType(); + Type dstTy = getDst().getType(); + if (failed(verifyTileLikeCommon(*this, srcTy, "src", + /*allowLowPrecision=*/false)) || + failed(verifyTileBufCommon(*this, dstTy, "dst", + /*allowLowPrecision=*/false)) || + failed(verifyNonNegativeIndexRowCol( + *getOperation(), getIndexRow(), getIndexCol(), + /*includeIndexAndIntOpsInConstFold=*/hasFp))) { + return failure(); + } + if (hasFp || hasPreQuantScalar || hasRelu || op.getAccToVecModeAttr()) { + return emitOpError("expects convtile textract to use the base form"); + } + auto srcCT = dyn_cast(srcTy); + auto srcSpace = getPTOMemorySpaceEnum(srcTy); + auto dstSpace = getPTOMemorySpaceEnum(dstTy); + if (!srcSpace || *srcSpace != pto::AddressSpace::MAT) { + return emitOpError("expects convtile textract src to use loc=mat"); + } + if (!dstSpace || *dstSpace != pto::AddressSpace::RIGHT) { + return emitOpError("expects convtile textract dst to use loc=right"); + } + auto srcLayout = dyn_cast_or_null(srcCT.getLayout()); + if (!srcLayout || + (srcLayout.getLayout() != pto::Layout::FRACTAL_Z && + srcLayout.getLayout() != pto::Layout::FRACTAL_Z_3D)) { + return emitOpError( + "expects convtile textract src to use layout=FRACTAL_Z or FRACTAL_Z_3D"); + } + auto dstTb = dyn_cast(dstTy); + if (!dstTb || + dstTb.getBLayoutValueI32() != static_cast(pto::BLayout::RowMajor) || + dstTb.getSLayoutValueI32() != static_cast(pto::SLayout::ColMajor)) { + return emitOpError( + "expects convtile textract dst to use blayout=row_major and slayout=col_major"); + } + Type srcElem = getElemTy(srcTy); + Type dstElem = getElemTy(dstTy); + if (!srcElem || !dstElem || srcElem != dstElem) { + return emitOpError("expects convtile textract src and dst to have the same element type"); + } + if (!isConvTileExtractElem(srcElem)) { + return emitOpError("expects convtile textract element type to be i8/i16/i32/f16/bf16/f32"); + } + return success(); + } if (!isA2A3ExtractElemType(c.dstElem) && !(hasFp && c.dstElem.isInteger(16))) { return op.emitOpError("expects A2/A3 textract element type to be i8/f16/bf16/f32"); } @@ -8288,6 +8375,54 @@ static LogicalResult verifyTExtractA5(TExtractOp op) { const TExtractCommon &c = *common; const bool hasPreQuantScalar = static_cast(op.getPreQuantScalar()); const bool hasRelu = op.getReluPreMode() != pto::ReluPreMode::NoRelu; + if (isa(getSrc().getType())) { + Type srcTy = getSrc().getType(); + Type dstTy = getDst().getType(); + if (failed(verifyTileLikeCommon(*this, srcTy, "src", + /*allowLowPrecision=*/true)) || + failed(verifyTileBufCommon(*this, dstTy, "dst", + /*allowLowPrecision=*/true)) || + failed(verifyNonNegativeIndexRowCol( + *getOperation(), getIndexRow(), getIndexCol(), + /*includeIndexAndIntOpsInConstFold=*/hasFp))) { + return failure(); + } + if (hasFp || hasPreQuantScalar || hasRelu || op.getAccToVecModeAttr()) { + return emitOpError("expects convtile textract to use the base form"); + } + auto srcCT = dyn_cast(srcTy); + auto srcSpace = getPTOMemorySpaceEnum(srcTy); + auto dstSpace = getPTOMemorySpaceEnum(dstTy); + if (!srcSpace || *srcSpace != pto::AddressSpace::MAT) { + return emitOpError("expects convtile textract src to use loc=mat"); + } + if (!dstSpace || *dstSpace != pto::AddressSpace::RIGHT) { + return emitOpError("expects convtile textract dst to use loc=right"); + } + auto srcLayout = dyn_cast_or_null(srcCT.getLayout()); + if (!srcLayout || + (srcLayout.getLayout() != pto::Layout::FRACTAL_Z && + srcLayout.getLayout() != pto::Layout::FRACTAL_Z_3D)) { + return emitOpError( + "expects convtile textract src to use layout=FRACTAL_Z or FRACTAL_Z_3D"); + } + auto dstTb = dyn_cast(dstTy); + if (!dstTb || + dstTb.getBLayoutValueI32() != static_cast(pto::BLayout::RowMajor) || + dstTb.getSLayoutValueI32() != static_cast(pto::SLayout::ColMajor)) { + return emitOpError( + "expects convtile textract dst to use blayout=row_major and slayout=col_major"); + } + Type srcElem = getElemTy(srcTy); + Type dstElem = getElemTy(dstTy); + if (!srcElem || !dstElem || srcElem != dstElem) { + return emitOpError("expects convtile textract src and dst to have the same element type"); + } + if (!isConvTileExtractElem(srcElem)) { + return emitOpError("expects convtile textract element type to be i8/i16/i32/f16/bf16/f32"); + } + return success(); + } if (!isA5ExtractElemType(c.dstElem)) { return op.emitOpError("expects A5 textract element type to be an fp8/f16/bf16/f32 or int8 family type"); } @@ -9766,6 +9901,48 @@ static LogicalResult verifyTMovImpl(TMovOp op, bool isA5) { if (classifyTMovForm(fp) == TMovForm::XToZz) { return verifyTMovXToZz(op, isA5); } + if (auto srcCT = dyn_cast(op.getSrc().getType())) { + Type srcTy = op.getSrc().getType(); + Type dstTy = op.getDst().getType(); + Value preQuantScalar = op.getPreQuantScalar(); + auto accToVecModeAttr = op.getAccToVecModeAttr(); + auto reluMode = op.getReluPreMode(); + const bool hasFp = static_cast(fp); + const bool hasPreQuantScalar = static_cast(preQuantScalar); + auto srcTb = dyn_cast(srcTy); + auto dstTb = dyn_cast(dstTy); + Type srcElem = getElemTy(srcTy); + Type dstElem = getElemTy(dstTy); + auto srcSpace = getPTOMemorySpaceEnum(srcTy); + auto dstSpace = getPTOMemorySpaceEnum(dstTy); + + if (!srcTb || !dstTb || !srcSpace || !dstSpace || !srcElem || !dstElem) { + return op.emitOpError("expects convtile tmov operands to be valid tile types"); + } + if (hasFp || hasPreQuantScalar || accToVecModeAttr || + reluMode != pto::ReluPreMode::NoRelu) { + return op.emitOpError("expects convtile tmov to use the base form"); + } + if (*srcSpace != pto::AddressSpace::MAT) { + return op.emitOpError("expects convtile tmov src to use loc=mat"); + } + if (*dstSpace != pto::AddressSpace::RIGHT) { + return op.emitOpError("expects convtile tmov dst to use loc=right"); + } + auto srcLayout = dyn_cast_or_null(srcCT.getLayout()); + if (!srcLayout || srcLayout.getLayout() != pto::Layout::FRACTAL_Z) { + return op.emitOpError("expects convtile tmov src to use layout=FRACTAL_Z"); + } + if (dstTb.getBLayoutValueI32() != static_cast(pto::BLayout::RowMajor) || + dstTb.getSLayoutValueI32() != static_cast(pto::SLayout::ColMajor)) { + return op.emitOpError( + "expects convtile tmov dst to use blayout=row_major and slayout=col_major"); + } + if (srcElem != dstElem) { + return op.emitOpError("expects convtile tmov src and dst to have the same element type"); + } + return success(); + } return verifyTMovGeneric(op, isA5); } @@ -16542,8 +16719,210 @@ static LogicalResult verifyTTransA5(TTransOp op) { } mlir::LogicalResult mlir::pto::TTransOp::verify() { + auto verifyConvTile = [&]() -> LogicalResult { + Type srcTy = getSrc().getType(); + Type tmpTy = getTmp() ? getTmp().getType() : Type{}; + Type dstTy = getDst().getType(); + auto srcCT = dyn_cast(srcTy); + auto dstCT = dyn_cast(dstTy); + auto tmpCT = dyn_cast(tmpTy); + + if (!srcCT || !dstCT) { + return emitOpError("expects convtile ttrans src and dst to be !pto.conv_tile"); + } + if (!tmpTy) { + return emitOpError("expects convtile ttrans to provide a tmp operand"); + } + if (!tmpCT) { + return emitOpError("expects convtile ttrans tmp to be !pto.conv_tile"); + } + + if (failed(verifyConvTileCommon(*this, srcTy, "src", + /*allowLowPrecision=*/false)) || + failed(verifyConvTileCommon(*this, dstTy, "dst", + /*allowLowPrecision=*/false)) || + failed(verifyConvTileCommon(*this, tmpTy, "tmp", + /*allowLowPrecision=*/false))) { + return failure(); + } + + Type srcElem = getElemTy(srcTy); + Type dstElem = getElemTy(dstTy); + Type tmpElem = getElemTy(tmpTy); + if (!srcElem || !dstElem || !tmpElem || srcElem != dstElem || srcElem != tmpElem) { + return emitOpError() << "expects convtile src, dst, and tmp to have the same element type"; + } + + auto srcLayout = dyn_cast_or_null(srcCT.getLayout()); + auto dstLayout = dyn_cast_or_null(dstCT.getLayout()); + if (!srcLayout || !dstLayout) { + return emitOpError("expects convtile ttrans src and dst to have a layout attr"); + } + + const int64_t elemBytes = getPTOStorageElemByteSize(srcElem); + if (elemBytes != 1 && elemBytes != 2 && elemBytes != 4) { + return emitOpError("expects convtile ttrans element size to be 1, 2, or 4 bytes"); + } + const int64_t c0 = 32 / elemBytes; + + auto checkedAdd = [](int64_t lhs, int64_t rhs) -> std::optional { + if (lhs < 0 || rhs < 0 || lhs > std::numeric_limits::max() - rhs) { + return std::nullopt; + } + return lhs + rhs; + }; + auto checkedMul = [](int64_t lhs, int64_t rhs) -> std::optional { + if (lhs < 0 || rhs < 0 || (rhs != 0 && lhs > std::numeric_limits::max() / rhs)) { + return std::nullopt; + } + return lhs * rhs; + }; + auto checkedCeilDiv = [&](int64_t lhs, int64_t rhs) -> std::optional { + if (lhs < 0 || rhs <= 0) { + return std::nullopt; + } + auto biased = checkedAdd(lhs, rhs - 1); + if (!biased) { + return std::nullopt; + } + return *biased / rhs; + }; + + auto srcShape = getShapeVec(srcTy); + auto dstShape = getShapeVec(dstTy); + auto expectShape = [&](ArrayRef actual, ArrayRef expected, + StringRef name) -> LogicalResult { + if (actual.size() != expected.size()) { + return emitOpError() << "expects " << name << " rank to be " << expected.size(); + } + for (size_t i = 0; i < expected.size(); ++i) { + if (actual[i] != expected[i]) { + return emitOpError() << "expects " << name << " shape[" << i << "] to be " + << expected[i] << " (got " << actual[i] << ")"; + } + } + return success(); + }; + + auto verifyPair = [&](pto::Layout srcLayoutKind, + pto::Layout dstLayoutKind) -> LogicalResult { + switch (srcLayoutKind) { + case pto::Layout::NCHW: { + if (dstLayoutKind != pto::Layout::NC1HWC0) { + return emitOpError("expects convtile ttrans NCHW src to lower to NC1HWC0 dst"); + } + if (srcShape.size() != 4 || dstShape.size() != 5) { + return emitOpError("expects convtile ttrans NCHW->NC1HWC0 to use 4D src and 5D dst"); + } + auto expectedC1 = checkedCeilDiv(srcShape[1], c0); + if (!expectedC1) { + return emitOpError("cannot compute convtile ttrans NCHW->NC1HWC0 channel split"); + } + SmallVector expected{srcShape[0], *expectedC1, srcShape[2], srcShape[3], c0}; + return expectShape(dstShape, expected, "dst"); + } + case pto::Layout::NC1HWC0: { + if (dstLayoutKind != pto::Layout::FRACTAL_Z) { + return emitOpError("expects convtile ttrans NC1HWC0 src to lower to FRACTAL_Z dst"); + } + if (srcShape.size() != 5 || dstShape.size() != 4) { + return emitOpError("expects convtile ttrans NC1HWC0->FRACTAL_Z to use 5D src and 4D dst"); + } + auto expectedN1 = checkedCeilDiv(srcShape[0], 16); + if (!expectedN1) { + return emitOpError("cannot compute convtile ttrans NC1HWC0->FRACTAL_Z batch split"); + } + auto expectedC1HW = checkedMul(srcShape[1], srcShape[2]); + if (!expectedC1HW) { + return emitOpError("cannot compute convtile ttrans NC1HWC0->FRACTAL_Z output size"); + } + expectedC1HW = checkedMul(*expectedC1HW, srcShape[3]); + if (!expectedC1HW) { + return emitOpError("cannot compute convtile ttrans NC1HWC0->FRACTAL_Z output size"); + } + SmallVector expected{*expectedC1HW, *expectedN1, 16, srcShape[4]}; + return expectShape(dstShape, expected, "dst"); + } + case pto::Layout::GNCHW: { + if (dstLayoutKind != pto::Layout::GNC1HWC0) { + return emitOpError("expects convtile ttrans GNCHW src to lower to GNC1HWC0 dst"); + } + if (srcShape.size() != 5 || dstShape.size() != 6) { + return emitOpError("expects convtile ttrans GNCHW->GNC1HWC0 to use 5D src and 6D dst"); + } + auto expectedC1 = checkedCeilDiv(srcShape[2], c0); + if (!expectedC1) { + return emitOpError("cannot compute convtile ttrans GNCHW->GNC1HWC0 channel split"); + } + SmallVector expected{srcShape[0], srcShape[1], *expectedC1, srcShape[3], + srcShape[4], c0}; + return expectShape(dstShape, expected, "dst"); + } + case pto::Layout::GNC1HWC0: { + if (dstLayoutKind != pto::Layout::FRACTAL_Z) { + return emitOpError("expects convtile ttrans GNC1HWC0 src to lower to FRACTAL_Z dst"); + } + if (srcShape.size() != 6 || dstShape.size() != 4) { + return emitOpError("expects convtile ttrans GNC1HWC0->FRACTAL_Z to use 6D src and 4D dst"); + } + auto expectedN1 = checkedCeilDiv(srcShape[1], 16); + if (!expectedN1) { + return emitOpError("cannot compute convtile ttrans GNC1HWC0->FRACTAL_Z batch split"); + } + auto expectedGC1HW = checkedMul(srcShape[0], srcShape[2]); + if (!expectedGC1HW) { + return emitOpError("cannot compute convtile ttrans GNC1HWC0->FRACTAL_Z output size"); + } + expectedGC1HW = checkedMul(*expectedGC1HW, srcShape[3]); + if (!expectedGC1HW) { + return emitOpError("cannot compute convtile ttrans GNC1HWC0->FRACTAL_Z output size"); + } + expectedGC1HW = checkedMul(*expectedGC1HW, srcShape[4]); + if (!expectedGC1HW) { + return emitOpError("cannot compute convtile ttrans GNC1HWC0->FRACTAL_Z output size"); + } + SmallVector expected{*expectedGC1HW, *expectedN1, 16, srcShape[5]}; + return expectShape(dstShape, expected, "dst"); + } + case pto::Layout::NCDHW: { + if (dstLayoutKind != pto::Layout::FRACTAL_Z_3D) { + return emitOpError("expects convtile ttrans NCDHW src to lower to FRACTAL_Z_3D dst"); + } + if (srcShape.size() != 5 || dstShape.size() != 4) { + return emitOpError("expects convtile ttrans NCDHW->FRACTAL_Z_3D to use 5D src and 4D dst"); + } + auto expectedC1 = checkedCeilDiv(srcShape[1], c0); + auto expectedN1 = checkedCeilDiv(srcShape[0], 16); + if (!expectedC1 || !expectedN1) { + return emitOpError("cannot compute convtile ttrans NCDHW->FRACTAL_Z_3D split"); + } + auto expectedDst0 = checkedMul(srcShape[2], *expectedC1); + if (!expectedDst0) { + return emitOpError("cannot compute convtile ttrans NCDHW->FRACTAL_Z_3D output size"); + } + expectedDst0 = checkedMul(*expectedDst0, srcShape[3]); + if (!expectedDst0) { + return emitOpError("cannot compute convtile ttrans NCDHW->FRACTAL_Z_3D output size"); + } + expectedDst0 = checkedMul(*expectedDst0, srcShape[4]); + if (!expectedDst0) { + return emitOpError("cannot compute convtile ttrans NCDHW->FRACTAL_Z_3D output size"); + } + SmallVector expected{*expectedDst0, *expectedN1, 16, c0}; + return expectShape(dstShape, expected, "dst"); + } + default: + return emitOpError("expects convtile ttrans src to use NCHW, NC1HWC0, GNCHW, GNC1HWC0, or NCDHW layout"); + } + }; + + return verifyPair(srcLayout.getLayout(), dstLayout.getLayout()); + }; auto verifyA2A3 = [&]() -> LogicalResult { return verifyTTransA2A3(*this); }; auto verifyA5 = [&]() -> LogicalResult { return verifyTTransA5(*this); }; + if (isa(getSrc().getType())) { + return verifyConvTile(); + } return dispatchVerifierByArch(getOperation(), verifyA2A3, verifyA5); } diff --git a/lib/PTO/IR/PTOAttrs.cpp b/lib/PTO/IR/PTOAttrs.cpp index a6b4a60a7f..5855423a6b 100644 --- a/lib/PTO/IR/PTOAttrs.cpp +++ b/lib/PTO/IR/PTOAttrs.cpp @@ -264,3 +264,462 @@ void TileBufConfigAttr::print(AsmPrinter &p) const { p << ", compact=" << getCompactMode(); p << ">"; } + +namespace { + +constexpr unsigned kConvTilePadListSize = 4; +constexpr int32_t kConvTileDefaultZero = 0; +constexpr int32_t kConvTileDefaultOne = 1; + +static LogicalResult parseConvTileIntField(AsmParser &parser, StringRef key, + IntegerAttr &value) { + if (failed(parseTileBufKeyEq(parser, key))) { + return failure(); + } + int64_t parsed = 0; + if (failed(parser.parseInteger(parsed))) { + return failure(); + } + value = IntegerAttr::get(IntegerType::get(parser.getContext(), kI32BitWidth), + parsed); + return success(); +} + +static LogicalResult parseConvTileBoolField(AsmParser &parser, StringRef key, + BoolAttr &value) { + if (failed(parseTileBufKeyEq(parser, key))) { + return failure(); + } + bool parsed = false; + if (succeeded(parser.parseOptionalKeyword("true"))) { + parsed = true; + } else if (succeeded(parser.parseOptionalKeyword("false"))) { + parsed = false; + } else { + parser.emitError(parser.getCurrentLocation()) + << key << " must be true or false"; + return failure(); + } + value = BoolAttr::get(parser.getContext(), parsed); + return success(); +} + +static LogicalResult parseConvTileAttrField(AsmParser &parser, StringRef key, + Attribute &value) { + if (failed(parseTileBufKeyEq(parser, key))) { + return failure(); + } + if (failed(parser.parseAttribute(value))) { + return failure(); + } + return success(); +} + +static LogicalResult parseConvTilePadListField(AsmParser &parser, + SmallVectorImpl &padList) { + if (failed(parseTileBufKeyEq(parser, "pad_list"))) { + return failure(); + } + SmallVector parsed; + if (failed(parser.parseDimensionList(parsed, /*allowDynamic=*/false, + /*withTrailingX=*/false))) { + return failure(); + } + if (parsed.size() != kConvTilePadListSize) { + parser.emitError(parser.getCurrentLocation()) + << "pad_list must have exactly " << kConvTilePadListSize + << " dimensions"; + return failure(); + } + padList.assign(parsed.begin(), parsed.end()); + return success(); +} + +static void printConvTileIntField(AsmPrinter &p, StringRef key, IntegerAttr value) { + p << key << "=" << value.getInt(); +} + +static void printConvTilePadList(AsmPrinter &p, ArrayRef padList) { + for (auto it = padList.begin(); it != padList.end(); ++it) { + if (it != padList.begin()) { + p << "x"; + } + p << *it; + } +} + +static StringRef printBoolAttr(BoolAttr value) { + return value.getValue() ? "true" : "false"; +} + +} // namespace + +ConvTileConfigAttr ConvTileConfigAttr::getDefault(MLIRContext *ctx) { + Builder b(ctx); + auto zero = b.getI32IntegerAttr(kConvTileDefaultZero); + auto one = b.getI32IntegerAttr(kConvTileDefaultOne); + SmallVector padList(kConvTilePadListSize, + kConvTileDefaultZero); + auto padValue = b.getI32IntegerAttr(kConvTileDefaultZero); + auto transpose = BoolAttr::get(ctx, false); + return ConvTileConfigAttr::get(ctx, zero, zero, padList, one, one, one, one, + one, one, padValue, zero, zero, one, zero, + zero, transpose); +} + +bool ConvTileConfigAttr::isDefault() const { + auto d = getDefault(getContext()); + return getFmapH() == d.getFmapH() && getFmapW() == d.getFmapW() && + getPadList() == d.getPadList() && getFilterH() == d.getFilterH() && + getFilterW() == d.getFilterW() && + getDilationH() == d.getDilationH() && + getDilationW() == d.getDilationW() && + getStrideH() == d.getStrideH() && getStrideW() == d.getStrideW() && + getPadValue() == d.getPadValue() && + getChannelSize() == d.getChannelSize() && + getRepeatStride() == d.getRepeatStride() && + getRepeatTime() == d.getRepeatTime() && + getRepeatMode() == d.getRepeatMode() && + getDstStride() == d.getDstStride() && + getDstMposition() == d.getDstMposition() && + getTranspose() == d.getTranspose(); +} + +LogicalResult ConvTileConfigAttr::verify(function_ref emitError, + IntegerAttr fmapH, + IntegerAttr fmapW, + ArrayRef padList, + IntegerAttr filterH, + IntegerAttr filterW, + IntegerAttr dilationH, + IntegerAttr dilationW, + IntegerAttr strideH, + IntegerAttr strideW, + Attribute padValue, + IntegerAttr channelSize, + IntegerAttr repeatStride, + IntegerAttr repeatTime, + IntegerAttr repeatMode, + IntegerAttr dstStride, + IntegerAttr dstMposition, + BoolAttr transpose) { + auto checkI32 = [&](IntegerAttr attr, StringRef name) -> LogicalResult { + if (!attr || !attr.getType().isSignlessInteger(kI32BitWidth)) { + return emitError() << name << " must be an i32 integer attr", failure(); + } + return success(); + }; + + if (failed(checkI32(fmapH, "fmap_h")) || failed(checkI32(fmapW, "fmap_w")) || + failed(checkI32(filterH, "filter_h")) || + failed(checkI32(filterW, "filter_w")) || + failed(checkI32(dilationH, "dilation_h")) || + failed(checkI32(dilationW, "dilation_w")) || + failed(checkI32(strideH, "stride_h")) || + failed(checkI32(strideW, "stride_w")) || + failed(checkI32(channelSize, "channel_size")) || + failed(checkI32(repeatStride, "repeat_stride")) || + failed(checkI32(repeatTime, "repeat_time")) || + failed(checkI32(repeatMode, "repeat_mode")) || + failed(checkI32(dstStride, "dst_stride")) || + failed(checkI32(dstMposition, "dst_mposition"))) { + return failure(); + } + + if (!transpose) { + return emitError() << "transpose must be a bool attr", failure(); + } + + if (padList.size() != kConvTilePadListSize) { + return emitError() << "pad_list must have exactly 4 elements", failure(); + } + for (int64_t dim : padList) { + if (dim < 0 || dim > 255) { + return emitError() << "pad_list entries must be in [0, 255]", failure(); + } + } + + auto checkNonNegative = [&](IntegerAttr attr, StringRef name) -> LogicalResult { + if (attr.getInt() < 0) { + return emitError() << name << " must be non-negative", failure(); + } + return success(); + }; + + if (failed(checkNonNegative(fmapH, "fmap_h")) || + failed(checkNonNegative(fmapW, "fmap_w")) || + failed(checkNonNegative(filterH, "filter_h")) || + failed(checkNonNegative(filterW, "filter_w")) || + failed(checkNonNegative(dilationH, "dilation_h")) || + failed(checkNonNegative(dilationW, "dilation_w")) || + failed(checkNonNegative(strideH, "stride_h")) || + failed(checkNonNegative(strideW, "stride_w")) || + failed(checkNonNegative(channelSize, "channel_size")) || + failed(checkNonNegative(repeatStride, "repeat_stride")) || + failed(checkNonNegative(repeatTime, "repeat_time")) || + failed(checkNonNegative(repeatMode, "repeat_mode")) || + failed(checkNonNegative(dstStride, "dst_stride")) || + failed(checkNonNegative(dstMposition, "dst_mposition"))) { + return failure(); + } + + if (!padValue || (!isa(padValue) && !isa(padValue))) { + return emitError() << "pad_value must be an integer or float attr", failure(); + } + + return success(); +} + +Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { + MLIRContext *ctx = p.getContext(); + auto def = ConvTileConfigAttr::getDefault(ctx); + IntegerAttr fmapH = def.getFmapH(); + IntegerAttr fmapW = def.getFmapW(); + SmallVector padList(def.getPadList().begin(), + def.getPadList().end()); + IntegerAttr filterH = def.getFilterH(); + IntegerAttr filterW = def.getFilterW(); + IntegerAttr dilationH = def.getDilationH(); + IntegerAttr dilationW = def.getDilationW(); + IntegerAttr strideH = def.getStrideH(); + IntegerAttr strideW = def.getStrideW(); + Attribute padValue = def.getPadValue(); + IntegerAttr channelSize = def.getChannelSize(); + IntegerAttr repeatStride = def.getRepeatStride(); + IntegerAttr repeatTime = def.getRepeatTime(); + IntegerAttr repeatMode = def.getRepeatMode(); + IntegerAttr dstStride = def.getDstStride(); + IntegerAttr dstMposition = def.getDstMposition(); + BoolAttr transpose = def.getTranspose(); + bool parsedGreater = false; + + auto consumeFieldTerminator = [&]() -> LogicalResult { + if (succeeded(p.parseOptionalGreater())) { + parsedGreater = true; + return success(); + } + return p.parseComma(); + }; + + if (p.parseLess()) { + return {}; + } + + if (succeeded(p.parseOptionalGreater())) { + return ConvTileConfigAttr::get(ctx, fmapH, fmapW, padList, filterH, filterW, + dilationH, dilationW, strideH, strideW, + padValue, channelSize, repeatStride, + repeatTime, repeatMode, dstStride, + dstMposition, transpose); + } + + while (!parsedGreater) { + StringRef key; + if (p.parseKeyword(&key)) { + return {}; + } + if (p.parseEqual()) { + return {}; + } + + if (key == "fmap_h") { + if (failed(parseConvTileIntField(p, key, fmapH))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "fmap_w") { + if (failed(parseConvTileIntField(p, key, fmapW))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "pad_list") { + if (failed(parseConvTilePadListField(p, padList))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "filter_h") { + if (failed(parseConvTileIntField(p, key, filterH))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "filter_w") { + if (failed(parseConvTileIntField(p, key, filterW))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "dilation_h") { + if (failed(parseConvTileIntField(p, key, dilationH))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "dilation_w") { + if (failed(parseConvTileIntField(p, key, dilationW))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "stride_h") { + if (failed(parseConvTileIntField(p, key, strideH))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "stride_w") { + if (failed(parseConvTileIntField(p, key, strideW))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "pad_value") { + if (failed(parseConvTileAttrField(p, key, padValue))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "channel_size") { + if (failed(parseConvTileIntField(p, key, channelSize))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "repeat_stride") { + if (failed(parseConvTileIntField(p, key, repeatStride))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "repeat_time") { + if (failed(parseConvTileIntField(p, key, repeatTime))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "repeat_mode") { + if (failed(parseConvTileIntField(p, key, repeatMode))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "dst_stride") { + if (failed(parseConvTileIntField(p, key, dstStride))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "dst_mposition") { + if (failed(parseConvTileIntField(p, key, dstMposition))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + if (key == "transpose") { + if (failed(parseConvTileBoolField(p, key, transpose))) { + return {}; + } + if (failed(consumeFieldTerminator())) { + return {}; + } + continue; + } + + p.emitError(p.getCurrentLocation(), "unknown key in conv_tile_config: ") + << key; + return {}; + } + + return ConvTileConfigAttr::get(ctx, fmapH, fmapW, padList, filterH, filterW, + dilationH, dilationW, strideH, strideW, + padValue, channelSize, repeatStride, + repeatTime, repeatMode, dstStride, dstMposition, + transpose); +} + +void ConvTileConfigAttr::print(AsmPrinter &p) const { + p << "<"; + printConvTileIntField(p, "fmap_h", getFmapH()); + p << ", "; + printConvTileIntField(p, "fmap_w", getFmapW()); + p << ", pad_list="; + printConvTilePadList(p, getPadList()); + p << ", "; + printConvTileIntField(p, "filter_h", getFilterH()); + p << ", "; + printConvTileIntField(p, "filter_w", getFilterW()); + p << ", "; + printConvTileIntField(p, "dilation_h", getDilationH()); + p << ", "; + printConvTileIntField(p, "dilation_w", getDilationW()); + p << ", "; + printConvTileIntField(p, "stride_h", getStrideH()); + p << ", "; + printConvTileIntField(p, "stride_w", getStrideW()); + p << ", pad_value=" << getPadValue(); + p << ", "; + printConvTileIntField(p, "channel_size", getChannelSize()); + p << ", "; + printConvTileIntField(p, "repeat_stride", getRepeatStride()); + p << ", "; + printConvTileIntField(p, "repeat_time", getRepeatTime()); + p << ", "; + printConvTileIntField(p, "repeat_mode", getRepeatMode()); + p << ", "; + printConvTileIntField(p, "dst_stride", getDstStride()); + p << ", "; + printConvTileIntField(p, "dst_mposition", getDstMposition()); + p << ", transpose=" << printBoolAttr(getTranspose()); + p << ">"; +} diff --git a/lib/PTO/IR/PTOTypeDefs.cpp b/lib/PTO/IR/PTOTypeDefs.cpp index b9773f8f2b..001041b991 100644 --- a/lib/PTO/IR/PTOTypeDefs.cpp +++ b/lib/PTO/IR/PTOTypeDefs.cpp @@ -190,19 +190,6 @@ static std::optional resolveTileBufMemorySpace(StringRef locStr) { .Default(::std::nullopt); } -static std::optional resolveConvLayout(StringRef layoutStr) { - return ::llvm::StringSwitch<::std::optional>(layoutStr) - .Case("nc1hwc0", ConvLayout::NC1HWC0) - .Case("ndc1hwc0", ConvLayout::NDC1HWC0) - .Case("fractal_z", ConvLayout::FRACTAL_Z) - .Case("fractal_z_3d", ConvLayout::FRACTAL_Z_3D) - .Case("nchw", ConvLayout::NCHW) - .Case("nhwc", ConvLayout::NHWC) - .Case("gnchw", ConvLayout::GNCHW) - .Case("gnc1hwc0", ConvLayout::GNC1HWC0) - .Default(::std::nullopt); -} - static BLayout resolveTileBufBLayout(MLIRContext *context, AddressSpace memorySpace, BLayout parsedLayout) { @@ -282,6 +269,26 @@ int32_t TileBufType::getCompactModeI32() const { return 0; } +ConvTileConfigAttr ConvTileType::getConfigAttr() const { + if constexpr (std::is_same_v) { + auto cfg = getConfig(); + if (!cfg) { + cfg = ConvTileConfigAttr::getDefault(getContext()); + } + return cfg; + } else { + auto cfg = llvm::dyn_cast_or_null(getConfig()); + if (!cfg) { + cfg = ConvTileConfigAttr::getDefault(getContext()); + } + return cfg; + } +} + +bool ConvTileType::hasNonDefaultConfig() const { + return !getConfigAttr().isDefault(); +} + namespace { struct ParsedTileBufFields { @@ -770,159 +777,266 @@ void mlir::pto::TileBufType::print(mlir::AsmPrinter &printer) const { printer << ">"; } -// ---- ConvTileType custom asm ---- -// !pto.conv_tile -Type ConvTileType::parse(AsmParser &parser) { - if (failed(parser.parseLess())) { - return Type(); - } +namespace { +struct ParsedConvTileFields { std::string locStr; - std::string layoutStr; SmallVector shape; Type dtype; - int64_t bufferSize = 0; + IntegerAttr bufferSize; + Attribute layoutAttr; + ConvTileConfigAttr config; +}; - ParseResult parseResult = parser.parseKeywordOrString(&locStr); - if (!parseResult.succeeded()) { - return Type(); +static std::optional computeConvTileBufferSize(ArrayRef shape) { + if (shape.empty()) { + return std::nullopt; } - parseResult = parser.parseComma(); - if (!parseResult.succeeded()) { - return Type(); + + int64_t capacity = 1; + for (int64_t dim : shape) { + if (dim <= 0) { + return std::nullopt; + } + if (capacity > std::numeric_limits::max() / dim) { + return std::nullopt; + } + capacity *= dim; } - parseResult = parser.parseKeyword("buffer"); - if (!parseResult.succeeded()) { - return Type(); + return capacity; +} + +static LogicalResult parseConvTileLayoutField(AsmParser &parser, + Attribute &layoutAttr) { + if (failed(parseTileBufKeyEq(parser, "layout"))) { + return failure(); } - parseResult = parser.parseEqual(); - if (!parseResult.succeeded()) { - return Type(); + if (failed(parser.parseAttribute(layoutAttr))) { + return failure(); } - parseResult = parser.parseInteger(bufferSize); - if (!parseResult.succeeded()) { - return Type(); + return success(); +} + +static LogicalResult parseConvTileConfigField(AsmParser &parser, + Attribute &configAttr) { + if (failed(parseTileBufKeyEq(parser, "config"))) { + return failure(); } - parseResult = parser.parseComma(); - if (!parseResult.succeeded()) { - return Type(); + if (failed(parser.parseAttribute(configAttr))) { + return failure(); } - parseResult = parser.parseKeyword("layout"); - if (!parseResult.succeeded()) { + return success(); +} + +static void printConvTileLayoutField(AsmPrinter &printer, Attribute layoutAttr) { + printer << "layout="; + printer.printAttribute(layoutAttr); +} + +static void printConvTileConfigField(AsmPrinter &printer, + ConvTileConfigAttr configAttr) { + printer << "config="; + printer.printAttribute(configAttr); +} + +static Type buildConvTileType(AsmParser &parser, + const ParsedConvTileFields &fields) { + MLIRContext *ctx = parser.getContext(); + auto emitError = [&]() -> InFlightDiagnostic { + return parser.emitError(parser.getNameLoc()); + }; + + if (fields.shape.empty()) { + emitError() << "conv_tile shape must be non-empty"; return Type(); } - parseResult = parser.parseEqual(); - if (!parseResult.succeeded()) { + if (fields.shape.size() > 6) { + emitError() << "conv_tile shape must have rank in [1, 6]"; return Type(); } - parseResult = parser.parseKeywordOrString(&layoutStr); - if (!parseResult.succeeded()) { + if (llvm::is_contained(fields.shape, ShapedType::kDynamic)) { + emitError() << "conv_tile shape must be static"; return Type(); } - parseResult = parser.parseComma(); - if (!parseResult.succeeded()) { + + auto defaultBufferSize = computeConvTileBufferSize(fields.shape); + if (!defaultBufferSize) { + emitError() << "conv_tile shape must have a positive, overflow-free element count"; return Type(); } - parseResult = parser.parseKeyword("shape"); - if (!parseResult.succeeded()) { + + IntegerAttr bufferSize = fields.bufferSize; + if (!bufferSize) { + bufferSize = IntegerAttr::get(parser.getBuilder().getI64Type(), + *defaultBufferSize); + } else if (bufferSize.getInt() < *defaultBufferSize) { + emitError() << "conv_tile buffer_size must be at least the logical element count"; return Type(); } - parseResult = parser.parseEqual(); - if (!parseResult.succeeded()) { + + auto memorySpace = resolveTileBufMemorySpace(fields.locStr); + if (!memorySpace.has_value()) { + emitError() << "unknown loc: " << fields.locStr; return Type(); } - parseResult = parser.parseDimensionList(shape, /*allowDynamic=*/false); - if (!parseResult.succeeded()) { + if (*memorySpace != AddressSpace::MAT) { + emitError() << "conv_tile only supports loc=mat"; return Type(); } - parseResult = parser.parseType(dtype); - if (!parseResult.succeeded()) { + + auto layoutAttr = dyn_cast_or_null(fields.layoutAttr); + if (!layoutAttr) { + emitError() << "layout must be a pto.layout attr"; return Type(); } - parseResult = parser.parseGreater(); - if (!parseResult.succeeded()) { + auto cfg = fields.config ? fields.config : ConvTileConfigAttr::getDefault(ctx); + + return ConvTileType::get(ctx, fields.shape, fields.dtype, bufferSize, + AddressSpaceAttr::get(ctx, *memorySpace), + layoutAttr, cfg); +} + +} // namespace + +Type ConvTileType::parse(AsmParser &parser) { + if (failed(parser.parseLess())) { return Type(); } - auto emitError = [&]() -> InFlightDiagnostic { - return parser.emitError(parser.getNameLoc()); - }; - auto memorySpace = resolveTileBufMemorySpace(locStr); - if (!memorySpace.has_value()) { - emitError() << "unknown ConvTile loc: " << locStr; + std::string firstToken; + if (failed(parser.parseKeywordOrString(&firstToken))) { return Type(); } - auto layout = resolveConvLayout(layoutStr); - if (!layout.has_value()) { - emitError() << "unknown ConvTile layout: " << layoutStr; + + ParsedConvTileFields fields; + fields.locStr = firstToken; + + if (failed(parser.parseComma())) { return Type(); } - if (bufferSize <= 0) { - emitError() << "ConvTile buffer must be positive"; + + if (failed(parser.parseDimensionList(fields.shape, /*allowDynamic=*/false, + /*withTrailingX=*/false))) { return Type(); } - const bool invalidRank = shape.empty() || shape.size() > 6; - if (invalidRank) { - emitError() << "ConvTile shape rank must be between 1 and 6"; + + if (failed(parser.parseType(fields.dtype))) { return Type(); } - for (int64_t dim : shape) { - if (dim <= 0) { - emitError() << "ConvTile shape dimensions must be positive"; + + bool parsedGreater = false; + bool seenLayout = false; + bool seenConfig = false; + bool seenBufferSize = false; + LayoutAttr layoutAttr; + IntegerAttr bufferSizeAttr; + ConvTileConfigAttr configAttr; + configAttr = ConvTileConfigAttr::getDefault(parser.getContext()); + while (!parsedGreater) { + if (succeeded(parser.parseOptionalGreater())) { + parsedGreater = true; + break; + } + if (failed(parser.parseComma())) { return Type(); } - } - auto memorySpaceAttr = AddressSpaceAttr::get(parser.getContext(), - memorySpace.value()); - auto layoutAttr = ConvLayoutAttr::get(parser.getContext(), layout.value()); - return ConvTileType::get(parser.getContext(), shape, dtype, memorySpaceAttr, - bufferSize, layoutAttr); -} + StringRef key; + if (failed(parser.parseKeyword(&key)) || failed(parser.parseEqual())) { + return Type(); + } -void mlir::pto::ConvTileType::print(mlir::AsmPrinter &printer) const { - auto memorySpace = - llvm::dyn_cast_or_null(getMemorySpace()); - auto layout = getLayout(); - if (!memorySpace || !layout) { - printer << ""; - return; + if (key == "layout") { + Attribute attr; + if (failed(parser.parseAttribute(attr))) { + return Type(); + } + layoutAttr = dyn_cast(attr); + if (!layoutAttr || seenLayout) { + return Type(); + } + seenLayout = true; + continue; + } + + if (key == "buffer_size") { + Attribute attr; + if (failed(parser.parseAttribute(attr))) { + return Type(); + } + bufferSizeAttr = dyn_cast_or_null(attr); + if (!bufferSizeAttr || seenBufferSize) { + return Type(); + } + seenBufferSize = true; + continue; + } + + if (key == "config") { + Attribute attr; + if (failed(parser.parseAttribute(attr))) { + return Type(); + } + configAttr = dyn_cast_or_null(attr); + if (!configAttr || seenConfig) { + return Type(); + } + seenConfig = true; + continue; + } + + parser.emitError(parser.getCurrentLocation(), + "unknown key in conv_tile syntax: ") + << key; + return Type(); + } + + if (!layoutAttr) { + layoutAttr = LayoutAttr::get(parser.getContext(), Layout::NC1HWC0); } - auto layoutName = [&]() -> llvm::StringRef { - switch (layout.getValue()) { - case ConvLayout::NC1HWC0: - return "nc1hwc0"; - case ConvLayout::NDC1HWC0: - return "ndc1hwc0"; - case ConvLayout::FRACTAL_Z: - return "fractal_z"; - case ConvLayout::FRACTAL_Z_3D: - return "fractal_z_3d"; - case ConvLayout::NCHW: - return "nchw"; - case ConvLayout::NHWC: - return "nhwc"; - case ConvLayout::GNCHW: - return "gnchw"; - case ConvLayout::GNC1HWC0: - return "gnc1hwc0"; + if (!bufferSizeAttr) { + auto defaultBufferSize = computeConvTileBufferSize(fields.shape); + if (!defaultBufferSize) { + parser.emitError(parser.getCurrentLocation(), + "conv_tile shape must have a positive, overflow-free element count"); + return Type(); } - return "unknown"; - }; + bufferSizeAttr = + IntegerAttr::get(parser.getBuilder().getI64Type(), *defaultBufferSize); + } - printer << "<" << stringifyLocFromMemorySpace(memorySpace) - << ", buffer=" << getBufferSize() - << ", layout=" << layoutName() - << ", shape="; - for (auto [index, dim] : llvm::enumerate(getShape())) { - if (index != 0) { + if (!seenConfig) { + configAttr = ConvTileConfigAttr::getDefault(parser.getContext()); + } + + return buildConvTileType(parser, {fields.locStr, fields.shape, fields.dtype, + bufferSizeAttr, layoutAttr, configAttr}); +} + +void ConvTileType::print(AsmPrinter &printer) const { + printer << "<"; + printer << stringifyLocFromMemorySpace(getMemorySpace()); + printer << ", "; + auto shape = getShape(); + for (auto it = shape.begin(); it != shape.end(); ++it) { + if (it != shape.begin()) { printer << "x"; } - printTileBufDim(printer, dim); + if (*it == ShapedType::kDynamic) { + printer << "?"; + } else { + printer << *it; + } } printer << "x"; printer.printType(getElementType()); + printer << ", buffer_size="; + printer.printAttribute(getBufferSize()); + printer << ", "; + printConvTileLayoutField(printer, getLayout()); + printer << ", "; + printConvTileConfigField(printer, getConfigAttr()); printer << ">"; } diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 22b2bfe05a..39f4af2817 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -217,7 +217,7 @@ static int64_t getIntegerAttrSignedValue(IntegerAttr attr) { static SmallVector collectTileOperandNumbers(Operation *op) { SmallVector tileOperandNumbers; for (OpOperand &operand : op->getOpOperands()) { - if (isa(operand.get().getType())) + if (isa(operand.get().getType())) tileOperandNumbers.push_back(operand.getOperandNumber()); } return tileOperandNumbers; @@ -843,54 +843,23 @@ static std::optional getEmitCTileTypeString(pto::TileBufType type) tileBufCompactToken(configAttr) + ">"; } -static StringRef convLayoutToken(pto::ConvLayout layout) { - switch (layout) { - case pto::ConvLayout::NC1HWC0: - return "Layout::NC1HWC0"; - case pto::ConvLayout::NDC1HWC0: - return "Layout::NDC1HWC0"; - case pto::ConvLayout::FRACTAL_Z: - return "Layout::FRACTAL_Z"; - case pto::ConvLayout::FRACTAL_Z_3D: - return "Layout::FRACTAL_Z_3D"; - case pto::ConvLayout::NCHW: - return "Layout::NCHW"; - case pto::ConvLayout::NHWC: - return "Layout::NHWC"; - case pto::ConvLayout::GNCHW: - return "Layout::GNCHW"; - case pto::ConvLayout::GNC1HWC0: - return "Layout::GNC1HWC0"; - } - return "Layout::NC1HWC0"; -} - -static std::optional -getEmitCConvTileTypeString(pto::ConvTileType type) { - auto memorySpace = - dyn_cast_or_null(type.getMemorySpace()); - auto layout = type.getLayout(); - const bool invalidType = - !memorySpace || !layout || type.getRank() == 0 || type.getRank() > 6 || - type.getBufferSize() <= 0; - if (invalidType) { +static std::optional getEmitCConvTileTypeString(pto::ConvTileType type) { + auto shape = type.getShape(); + if (shape.empty() || shape.size() > 6) return std::nullopt; - } - std::string shape = "ConvTileShape<"; - for (auto [index, dim] : llvm::enumerate(type.getShape())) { - if (index != 0) { - shape += ", "; - } - shape += std::to_string(dim); - } - shape += ">"; + Type elemTy = type.getElementType(); + auto layoutAttr = dyn_cast_or_null(type.getLayout()); + if (!layoutAttr) + layoutAttr = pto::LayoutAttr::get(type.getContext(), pto::Layout::NC1HWC0); + std::string shapeType = "pto::ConvTileShape<" + joinIntTemplateParams(shape) + ">"; return std::string("ConvTile<") + - tileRoleToken(type.getMemorySpace(), type.getElementType(), nullptr) + - ", " + getEmitCScalarTypeToken(type.getElementType()) + ", " + - std::to_string(type.getBufferSize()) + ", " + - convLayoutToken(layout.getValue()).str() + ", " + shape + ">"; + tileRoleToken(type.getMemorySpace(), elemTy) + ", " + + getEmitCScalarTypeToken(elemTy) + ", " + + std::to_string(type.getBufferSizeValue()) + ", " + + layoutToEmitCString(layoutAttr.getLayout()) + ", " + + shapeType + ">"; } //===----------------------------------------------------------------------===// @@ -1067,19 +1036,17 @@ class PTOToEmitCTypeConverter : public TypeConverter { type.getShape()); }); - addConversion([Ctx](pto::TileBufType type) -> std::optional { + addConversion([Ctx](pto::TileBufType type) -> std::optional { auto typeString = getEmitCTileTypeString(type); - if (!typeString) { + if (!typeString) return std::nullopt; - } return emitc::OpaqueType::get(Ctx, *typeString); }); addConversion([Ctx](pto::ConvTileType type) -> std::optional { auto typeString = getEmitCConvTileTypeString(type); - if (!typeString) { + if (!typeString) return std::nullopt; - } return emitc::OpaqueType::get(Ctx, *typeString); }); @@ -12411,6 +12378,72 @@ struct PTOXORSToEmitC : public OpConversionPattern { return success(); } }; + +struct PTOSetFmatrixToEmitC : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(pto::SetFmatrixOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, + "SETFMATRIX", + ValueRange{peelUnrealized(adaptor.getSrc())}); + rewriter.eraseOp(op); + return success(); + } +}; + +struct PTOSetImg2colRptToEmitC + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(pto::SetImg2colRptOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, + "SET_IMG2COL_RPT", + ValueRange{peelUnrealized(adaptor.getSrc())}); + rewriter.eraseOp(op); + return success(); + } +}; + +struct PTOSetImg2colPaddingToEmitC + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(pto::SetImg2colPaddingOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, + "SET_IMG2COL_PADDING", + ValueRange{peelUnrealized(adaptor.getSrc())}); + rewriter.eraseOp(op); + return success(); + } +}; + +struct PTOTImg2colToEmitC : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(pto::TImg2colOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto loc = op.getLoc(); + auto *ctx = rewriter.getContext(); + + Value dst = peelUnrealized(adaptor.getDst()); + Value src = peelUnrealized(adaptor.getSrc()); + Type u16Ty = emitc::OpaqueType::get(ctx, "uint16_t"); + Value posM = makeEmitCIntConstant(rewriter, loc, u16Ty, + static_cast(op.getPosM().getInt())); + Value posK = makeEmitCIntConstant(rewriter, loc, u16Ty, + static_cast(op.getPosK().getInt())); + + createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, + "TIMG2COL", + ValueRange{dst, src, posM, posK}); + rewriter.eraseOp(op); + return success(); + } +}; + struct PTOPrintToTPRINT : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; @@ -12535,57 +12568,7 @@ struct PTOAllocTileToEmitC ConversionPatternRewriter &rewriter) const override { Location loc = op.getLoc(); MLIRContext *ctx = rewriter.getContext(); - Type resultTy = op.getResult().getType(); - if (auto convTy = dyn_cast(resultTy)) { - auto convTypeString = getEmitCConvTileTypeString(convTy); - if (!convTypeString) { - return rewriter.notifyMatchFailure( - op, "invalid ConvTile type for EmitC conversion"); - } - Type convertedTy = getTypeConverter()->convertType(convTy); - if (!convertedTy) { - convertedTy = emitc::OpaqueType::get(ctx, *convTypeString); - } - Value tile = - rewriter - .create( - loc, getEmitCVariableResultType(convertedTy), - emitc::OpaqueAttr::get(ctx, "")) - .getResult(); - tile = loadEmitCVariableIfNeeded(rewriter, loc, tile); - - Value addr = adaptor.getAddr(); - if (addr) { - addr = peelUnrealized(addr); - auto u64Ty = emitc::OpaqueType::get(ctx, "uint64_t"); - const bool isPointer = - isa(addr.getType()) || - (isa(addr.getType()) && - cast(addr.getType()).getValue().ends_with("*")); - if (isPointer) { - auto rcU64 = - rewriter.getArrayAttr({emitc::OpaqueAttr::get(ctx, "uint64_t")}); - addr = rewriter - .create( - loc, u64Ty, "reinterpret_cast", ArrayAttr{}, rcU64, - ValueRange{addr}) - .getResult(0); - } else if (addr.getType() != u64Ty) { - addr = rewriter.create(loc, u64Ty, addr).getResult(); - } - rewriter.create( - loc, TypeRange{}, "TASSIGN", ArrayAttr{}, ArrayAttr{}, - ValueRange{tile, addr}); - } - rewriter.replaceOp(op, tile); - return success(); - } - - auto tileTy = dyn_cast(resultTy); - if (!tileTy) { - return rewriter.notifyMatchFailure( - op, "expected tile_buf or conv_tile result"); - } + auto tileTy = cast(op.getResult().getType()); auto tileTypeString = getEmitCTileTypeString(tileTy); if (!tileTypeString) return rewriter.notifyMatchFailure( @@ -13821,6 +13804,9 @@ static void populatePTOToEmitCPatterns(RewritePatternSet &patterns, patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); patterns.add(typeConverter, ctx); + patterns.add( + typeConverter, ctx); patterns.add< PTOTMatmulBiasToTMATMUL_BIAS, PTOTMatmulMXToTMATMUL_MX, From 226d5983799b537f9f7e1bcb369b329b4b081eaf Mon Sep 17 00:00:00 2001 From: andodo Date: Mon, 31 Aug 2026 01:33:58 +0800 Subject: [PATCH 03/12] =?UTF-8?q?Add=204=20op:=20SETFMATRIX=E3=80=81SET=5F?= =?UTF-8?q?IMG2COL=5FRPT=E3=80=81SET=5FIMG2COL=5FPADDING=E3=80=81TIMG2COL.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- include/PTO/IR/PTOAttrs.td | 23 ++++++++++ include/PTO/IR/PTOOps.td | 18 ++++++-- lib/PTO/IR/PTO.cpp | 73 ++++++++++++++++++++++++++++--- lib/PTO/Transforms/PTOToEmitC.cpp | 40 +++++++++++++++-- 4 files changed, 140 insertions(+), 14 deletions(-) diff --git a/include/PTO/IR/PTOAttrs.td b/include/PTO/IR/PTOAttrs.td index 71c2f9b117..327cf72670 100644 --- a/include/PTO/IR/PTOAttrs.td +++ b/include/PTO/IR/PTOAttrs.td @@ -94,6 +94,29 @@ def PTO_SignednessAttr : PTO_Attr<"Signedness", "signedness"> { let summary = "Integer signedness control for semantic ops"; } +//===----------------------------------------------------------------------===// +// FMATRIX mode +//===----------------------------------------------------------------------===// + +def PTO_FmatrixMode_AutoA : I32EnumAttrCase<"FMATRIX_A_AUTO", 0, "a_auto">; +def PTO_FmatrixMode_AutoB : I32EnumAttrCase<"FMATRIX_B_AUTO", 1, "b_auto">; +def PTO_FmatrixMode_ManualA : I32EnumAttrCase<"FMATRIX_A_MANUAL", 2, "a_manual">; +def PTO_FmatrixMode_ManualB : I32EnumAttrCase<"FMATRIX_B_MANUAL", 3, "b_manual">; + +def PTO_FmatrixModeEnum : PTO_I32Enum< + "FmatrixMode", "PTO FMATRIX selector", [ + PTO_FmatrixMode_AutoA, + PTO_FmatrixMode_AutoB, + PTO_FmatrixMode_ManualA, + PTO_FmatrixMode_ManualB + ]>; + +def PTO_FmatrixModeAttr : PTO_Attr<"FmatrixMode", "fmatrix_mode"> { + let parameters = (ins EnumParameter:$value); + let assemblyFormat = "`<` params `>`"; + let summary = "A/B selector for FMATRIX-related ops"; +} + //===----------------------------------------------------------------------===// // SIMT GM load/store L1 and L2 cache controls //===----------------------------------------------------------------------===// diff --git a/include/PTO/IR/PTOOps.td b/include/PTO/IR/PTOOps.td index ff62dbf6a1..c2bd746c10 100644 --- a/include/PTO/IR/PTOOps.td +++ b/include/PTO/IR/PTOOps.td @@ -1899,7 +1899,10 @@ def SetQuantVectorOp : PTO_Op<"set_quant_vector", [ def SetFmatrixOp : PTO_Op<"set_fmatrix", [MemoryEffects<[MemWrite]>]> { let summary = "Set FMATRIX registers from a ConvTile config"; - let arguments = (ins ConvTileType:$src); + let arguments = (ins + ConvTileType:$src, + DefaultValuedAttr:$fmatrixMode + ); let results = (outs); let hasVerifier = 1; let assemblyFormat = [{ @@ -1909,7 +1912,10 @@ def SetFmatrixOp : PTO_Op<"set_fmatrix", [MemoryEffects<[MemWrite]>]> { def SetImg2colRptOp : PTO_Op<"set_img2col_rpt", [MemoryEffects<[MemWrite]>]> { let summary = "Set IMG2COL repeat control from a ConvTile config"; - let arguments = (ins ConvTileType:$src); + let arguments = (ins + ConvTileType:$src, + DefaultValuedAttr:$fmatrixMode + ); let results = (outs); let hasVerifier = 1; let assemblyFormat = [{ @@ -1919,7 +1925,10 @@ def SetImg2colRptOp : PTO_Op<"set_img2col_rpt", [MemoryEffects<[MemWrite]>]> { def SetImg2colPaddingOp : PTO_Op<"set_img2col_padding", [MemoryEffects<[MemWrite]>]> { let summary = "Set IMG2COL padding control from a ConvTile config"; - let arguments = (ins ConvTileType:$src); + let arguments = (ins + ConvTileType:$src, + DefaultValuedAttr:$fmatrixMode + ); let results = (outs); let hasVerifier = 1; let assemblyFormat = [{ @@ -1935,7 +1944,8 @@ def TImg2colOp : PTO_Op<"timg2col", [ TileBufType:$dst, ConvTileType:$src, DefaultValuedOptionalAttr:$posM, - DefaultValuedOptionalAttr:$posK + DefaultValuedOptionalAttr:$posK, + DefaultValuedAttr:$fmatrixMode ); let results = (outs); let hasVerifier = 1; diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index fbf9d6bbf9..963ebf2fab 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3927,18 +3927,42 @@ LogicalResult TLoadOp::verify() { } LogicalResult mlir::pto::SetFmatrixOp::verify() { - return verifyConvTileCommon(*this, getSrc().getType(), "src", - /*allowLowPrecision=*/true); + if (failed(verifyConvTileCommon(*this, getSrc().getType(), "src", + /*allowLowPrecision=*/true))) { + return failure(); + } + auto mode = getFmatrixMode(); + if (mode != pto::FmatrixMode::FMATRIX_A_MANUAL && + mode != pto::FmatrixMode::FMATRIX_B_MANUAL) { + return emitOpError("expects fmatrix_mode to be a_manual or b_manual"); + } + return success(); } LogicalResult mlir::pto::SetImg2colRptOp::verify() { - return verifyConvTileCommon(*this, getSrc().getType(), "src", - /*allowLowPrecision=*/true); + if (failed(verifyConvTileCommon(*this, getSrc().getType(), "src", + /*allowLowPrecision=*/true))) { + return failure(); + } + auto mode = getFmatrixMode(); + if (mode != pto::FmatrixMode::FMATRIX_A_MANUAL && + mode != pto::FmatrixMode::FMATRIX_B_MANUAL) { + return emitOpError("expects fmatrix_mode to be a_manual or b_manual"); + } + return success(); } LogicalResult mlir::pto::SetImg2colPaddingOp::verify() { - return verifyConvTileCommon(*this, getSrc().getType(), "src", - /*allowLowPrecision=*/true); + if (failed(verifyConvTileCommon(*this, getSrc().getType(), "src", + /*allowLowPrecision=*/true))) { + return failure(); + } + auto mode = getFmatrixMode(); + if (mode != pto::FmatrixMode::FMATRIX_A_MANUAL && + mode != pto::FmatrixMode::FMATRIX_B_MANUAL) { + return emitOpError("expects fmatrix_mode to be a_manual or b_manual"); + } + return success(); } LogicalResult mlir::pto::TImg2colOp::verify() { @@ -3954,6 +3978,35 @@ LogicalResult mlir::pto::TImg2colOp::verify() { return emitOpError() << "expects src and dst to have the same element type"; } + auto srcCT = dyn_cast(getSrc().getType()); + if (!srcCT) { + return emitOpError("expects src to be a !pto.conv_tile"); + } + auto srcLayout = dyn_cast_or_null(srcCT.getLayout()); + if (!srcLayout || + (srcLayout.getLayout() != pto::Layout::NC1HWC0 && + srcLayout.getLayout() != pto::Layout::NDC1HWC0)) { + return emitOpError( + "expects src layout to be NC1HWC0 or NDC1HWC0"); + } + + auto dstTile = dyn_cast(getDst().getType()); + if (!dstTile) { + return emitOpError("expects dst to be a !pto.tile_buf"); + } + if (failed(verifyTileBufLayoutConstraints(*this, dstTile, "dst"))) { + return failure(); + } + if (dstTile.getBLayoutValueI32() != static_cast(pto::BLayout::ColMajor) || + dstTile.getSLayoutValueI32() != static_cast(pto::SLayout::RowMajor)) { + return emitOpError( + "expects dst layout to be BLayout=col_major and SLayout=row_major"); + } + auto dstSpace = getPTOMemorySpaceEnum(dstTile); + if (!dstSpace || *dstSpace != pto::AddressSpace::LEFT) { + return emitOpError("expects dst to use loc=left"); + } + auto checkPos = [&](IntegerAttr posAttr, StringRef name) -> LogicalResult { if (!posAttr || !posAttr.getType().isSignlessInteger(32)) { return emitOpError() << "expects " << name << " to be an i32 attr"; @@ -3969,6 +4022,14 @@ LogicalResult mlir::pto::TImg2colOp::verify() { return failure(); } + auto mode = getFmatrixMode(); + if (mode != pto::FmatrixMode::FMATRIX_A_AUTO && + mode != pto::FmatrixMode::FMATRIX_B_AUTO && + mode != pto::FmatrixMode::FMATRIX_A_MANUAL && + mode != pto::FmatrixMode::FMATRIX_B_MANUAL) { + return emitOpError("expects fmatrix_mode to be one of a_auto, b_auto, a_manual, or b_manual"); + } + return success(); } diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 39f4af2817..cd811fdfd7 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -336,6 +336,30 @@ static void createLastUseAwareOpaqueCall( operands); } +static StringRef getFmatrixModeToken(pto::FmatrixMode mode) { + switch (mode) { + case pto::FmatrixMode::FMATRIX_A_AUTO: + return "pto::SetFmatrixMode::FMATRIX_A_AUTO"; + case pto::FmatrixMode::FMATRIX_B_AUTO: + return "pto::SetFmatrixMode::FMATRIX_B_AUTO"; + case pto::FmatrixMode::FMATRIX_A_MANUAL: + return "pto::SetFmatrixMode::FMATRIX_A_MANUAL"; + case pto::FmatrixMode::FMATRIX_B_MANUAL: + return "pto::SetFmatrixMode::FMATRIX_B_MANUAL"; + } + llvm_unreachable("unknown FmatrixMode"); +} + +static ArrayAttr getFmatrixModeTemplateArgs(ConversionPatternRewriter &rewriter, + pto::FmatrixMode mode) { + if (mode == pto::FmatrixMode::FMATRIX_A_MANUAL) { + return ArrayAttr{}; + } + auto *ctx = rewriter.getContext(); + return rewriter.getArrayAttr( + {emitc::OpaqueAttr::get(ctx, getFmatrixModeToken(mode))}); +} + static Value buildGlobalTensorFromMemref(ConversionPatternRewriter &rewriter, Location loc, Value basePtr, MemRefType mrTy, Operation *anchor, @@ -12384,9 +12408,11 @@ struct PTOSetFmatrixToEmitC : public OpConversionPattern { LogicalResult matchAndRewrite(pto::SetFmatrixOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { + auto templateArgs = getFmatrixModeTemplateArgs(rewriter, op.getFmatrixMode()); createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, "SETFMATRIX", - ValueRange{peelUnrealized(adaptor.getSrc())}); + ValueRange{peelUnrealized(adaptor.getSrc())}, + ArrayAttr{}, templateArgs); rewriter.eraseOp(op); return success(); } @@ -12398,9 +12424,11 @@ struct PTOSetImg2colRptToEmitC LogicalResult matchAndRewrite(pto::SetImg2colRptOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { + auto templateArgs = getFmatrixModeTemplateArgs(rewriter, op.getFmatrixMode()); createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, "SET_IMG2COL_RPT", - ValueRange{peelUnrealized(adaptor.getSrc())}); + ValueRange{peelUnrealized(adaptor.getSrc())}, + ArrayAttr{}, templateArgs); rewriter.eraseOp(op); return success(); } @@ -12412,9 +12440,11 @@ struct PTOSetImg2colPaddingToEmitC LogicalResult matchAndRewrite(pto::SetImg2colPaddingOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { + auto templateArgs = getFmatrixModeTemplateArgs(rewriter, op.getFmatrixMode()); createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, "SET_IMG2COL_PADDING", - ValueRange{peelUnrealized(adaptor.getSrc())}); + ValueRange{peelUnrealized(adaptor.getSrc())}, + ArrayAttr{}, templateArgs); rewriter.eraseOp(op); return success(); } @@ -12435,10 +12465,12 @@ struct PTOTImg2colToEmitC : public OpConversionPattern { static_cast(op.getPosM().getInt())); Value posK = makeEmitCIntConstant(rewriter, loc, u16Ty, static_cast(op.getPosK().getInt())); + auto templateArgs = getFmatrixModeTemplateArgs(rewriter, op.getFmatrixMode()); createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, "TIMG2COL", - ValueRange{dst, src, posM, posK}); + ValueRange{dst, src, posM, posK}, ArrayAttr{}, + templateArgs); rewriter.eraseOp(op); return success(); } From baa914a41aa28bd66a2347db604c18c8cd936efa Mon Sep 17 00:00:00 2001 From: andodo Date: Mon, 31 Aug 2026 16:34:41 +0800 Subject: [PATCH 04/12] Rebase fix. --- include/PTO/IR/PTOAttrs.td | 51 ++++---- lib/Bindings/Python/PTOModule.cpp | 9 +- lib/PTO/IR/PTO.cpp | 138 ++++++++++++--------- lib/PTO/IR/PTOAttrs.cpp | 10 +- lib/PTO/IR/PTOTypeDefs.cpp | 8 +- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 4 +- lib/PTO/Transforms/PTOToEmitC.cpp | 22 +++- lib/PTO/Transforms/Utils.cpp | 4 +- 8 files changed, 146 insertions(+), 100 deletions(-) diff --git a/include/PTO/IR/PTOAttrs.td b/include/PTO/IR/PTOAttrs.td index 327cf72670..25a3e94015 100644 --- a/include/PTO/IR/PTOAttrs.td +++ b/include/PTO/IR/PTOAttrs.td @@ -111,8 +111,7 @@ def PTO_FmatrixModeEnum : PTO_I32Enum< PTO_FmatrixMode_ManualB ]>; -def PTO_FmatrixModeAttr : PTO_Attr<"FmatrixMode", "fmatrix_mode"> { - let parameters = (ins EnumParameter:$value); +def PTO_FmatrixModeAttr : EnumAttr { let assemblyFormat = "`<` params `>`"; let summary = "A/B selector for FMATRIX-related ops"; } @@ -507,6 +506,16 @@ def PTO_Layout_DN : I32EnumAttrCase<"DN", 1, "dn">; def PTO_Layout_NZ : I32EnumAttrCase<"NZ", 2, "nz">; def PTO_Layout_MX_A_ZZ : I32EnumAttrCase<"MX_A_ZZ", 3, "mx_a_zz">; def PTO_Layout_MX_B_NN : I32EnumAttrCase<"MX_B_NN", 4, "mx_b_nn">; +def PTO_Layout_NCHW : I32EnumAttrCase<"NCHW", 5, "nchw">; +def PTO_Layout_NC1HWC0 : I32EnumAttrCase<"NC1HWC0", 6, "nc1hwc0">; +def PTO_Layout_NCDHW : I32EnumAttrCase<"NCDHW", 7, "ncdhw">; +def PTO_Layout_NDC1HWC0 : I32EnumAttrCase<"NDC1HWC0", 8, "ndc1hwc0">; +def PTO_Layout_GNCHW : I32EnumAttrCase<"GNCHW", 9, "gnchw">; +def PTO_Layout_GNC1HWC0 : I32EnumAttrCase<"GNC1HWC0", 10, "gnc1hwc0">; +def PTO_Layout_NHWC : I32EnumAttrCase<"NHWC", 11, "nhwc">; +def PTO_Layout_FRACTAL_Z : I32EnumAttrCase<"FRACTAL_Z", 12, "fractal_z">; +def PTO_Layout_FRACTAL_Z_3D : + I32EnumAttrCase<"FRACTAL_Z_3D", 13, "fractal_z_3d">; def PTO_LayoutEnum : PTO_I32Enum< "Layout", "Global tensor layout (row/col/fractal)", [ @@ -514,16 +523,26 @@ def PTO_LayoutEnum : PTO_I32Enum< PTO_Layout_DN, PTO_Layout_NZ, PTO_Layout_MX_A_ZZ, - PTO_Layout_MX_B_NN + PTO_Layout_MX_B_NN, + PTO_Layout_NCHW, + PTO_Layout_NC1HWC0, + PTO_Layout_NCDHW, + PTO_Layout_NDC1HWC0, + PTO_Layout_GNCHW, + PTO_Layout_GNC1HWC0, + PTO_Layout_NHWC, + PTO_Layout_FRACTAL_Z, + PTO_Layout_FRACTAL_Z_3D ]>; def PTO_LayoutAttr : PTO_Attr<"Layout", "layout"> { let parameters = (ins EnumParameter:$layout); let assemblyFormat = "`<` params `>`"; let description = [{ - Layout inferred from shape/stride for GlobalTensor: + Layout inferred from shape/stride for GlobalTensor and ConvTile: ND (row-major), DN (col-major), NZ (fractal), - MX_A_ZZ (A-side MX scale), MX_B_NN (B-side MX scale). + MX_A_ZZ (A-side MX scale), MX_B_NN (B-side MX scale), + NCHW/NC1HWC0/NCDHW/NDC1HWC0/GNCHW/GNC1HWC0/NHWC/FRACTAL_Z/FRACTAL_Z_3D for conv layouts. }]; } @@ -1326,28 +1345,6 @@ def ConvTileConfigAttr : AttrDef { let hasCustomAssemblyFormat = 1; - let builders = [ - AttrBuilder<(ins - "mlir::IntegerAttr":$fmapH, - "mlir::IntegerAttr":$fmapW, - ArrayRefParameter<"int64_t">:$padList, - "mlir::IntegerAttr":$filterH, - "mlir::IntegerAttr":$filterW, - "mlir::IntegerAttr":$dilationH, - "mlir::IntegerAttr":$dilationW, - "mlir::IntegerAttr":$strideH, - "mlir::IntegerAttr":$strideW, - "mlir::Attribute":$padValue, - "mlir::IntegerAttr":$channelSize, - "mlir::IntegerAttr":$repeatStride, - "mlir::IntegerAttr":$repeatTime, - "mlir::IntegerAttr":$repeatMode, - "mlir::IntegerAttr":$dstStride, - "mlir::IntegerAttr":$dstMposition, - "mlir::BoolAttr":$transpose - )> - ]; - let extraClassDeclaration = [{ static ConvTileConfigAttr getDefault(MLIRContext *ctx); bool isDefault() const; diff --git a/lib/Bindings/Python/PTOModule.cpp b/lib/Bindings/Python/PTOModule.cpp index caef126588..c11cd3e300 100644 --- a/lib/Bindings/Python/PTOModule.cpp +++ b/lib/Bindings/Python/PTOModule.cpp @@ -258,7 +258,14 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { .value("DN", mlir::pto::Layout::DN) .value("NZ", mlir::pto::Layout::NZ) .value("MX_A_ZZ", mlir::pto::Layout::MX_A_ZZ) - .value("MX_B_NN", mlir::pto::Layout::MX_B_NN); + .value("MX_B_NN", mlir::pto::Layout::MX_B_NN) + .value("NCHW", mlir::pto::Layout::NCHW) + .value("NC1HWC0", mlir::pto::Layout::NC1HWC0) + .value("NCDHW", mlir::pto::Layout::NCDHW) + .value("NDC1HWC0", mlir::pto::Layout::NDC1HWC0) + .value("GNCHW", mlir::pto::Layout::GNCHW) + .value("GNC1HWC0", mlir::pto::Layout::GNC1HWC0) + .value("NHWC", mlir::pto::Layout::NHWC); py::enum_(m, "AccToVecMode") .value("SingleModeVec0", mlir::pto::AccToVecMode::SingleModeVec0) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 963ebf2fab..c9e0608100 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -126,6 +126,10 @@ static bool isKnownZeroOrUnitExtent(int64_t value); static bool isByteIntegerType(Type ty); static LogicalResult verifyTileBufCommon(Operation *op, Type ty, StringRef name, bool allowLowPrecision = false); +static LogicalResult verifyConvTileCommon(Operation *op, Type ty, StringRef name, + bool allowLowPrecision = false); +static LogicalResult verifyTileLikeCommon(Operation *op, Type ty, StringRef name, + bool allowLowPrecision = false); static LogicalResult verifyTmpCapacityAtLeast(Operation *op, Type tmpTy, uint64_t requiredBytes, StringRef tmpName = "tmp"); @@ -1956,9 +1960,12 @@ static unsigned getElemByteSize(Type ty) { return getPTOStorageElemByteSize(ty); } -static LogicalResult verifyTileBufLayoutConstraints(Operation *op, - pto::TileBufType tb, +static LogicalResult verifyTileBufLayoutConstraints(Operation *op, Type ty, StringRef name) { + auto tb = dyn_cast(ty); + if (!tb) { + return op->emitOpError() << "expects " << name << " to be a !pto.tile_buf"; + } auto shape = tb.getShape(); if (shape.size() != 2) { return op->emitOpError() << "expects " << name << " to be rank-2"; @@ -3585,7 +3592,10 @@ static LogicalResult verifyConstantLocalAddress(Operation *op, Value addr, } LogicalResult AllocTileOp::verify() { - auto ty = getResult().getType(); // TileBufType + auto ty = dyn_cast(getResult().getType()); + if (!ty) { + return emitOpError("result must be `!pto.tile_buf`"); + } if (failed(verifyTileBufLayoutConstraints(*this, ty, "result"))) { return failure(); @@ -3894,10 +3904,12 @@ LogicalResult TLoadOp::verify() { } if (dstElem.isInteger(64)) { - auto pad = dstTile.getPadValueI32(); - if (pad != static_cast(pto::PadValue::Null) && - pad != static_cast(pto::PadValue::Zero)) { - return emitOpError("expects A5 i64/u64 tload dst pad to be null or zero"); + if (auto dstTB = dyn_cast(dstTile)) { + auto pad = dstTB.getPadValueI32(); + if (pad != static_cast(pto::PadValue::Null) && + pad != static_cast(pto::PadValue::Zero)) { + return emitOpError("expects A5 i64/u64 tload dst pad to be null or zero"); + } } } @@ -4007,12 +4019,8 @@ LogicalResult mlir::pto::TImg2colOp::verify() { return emitOpError("expects dst to use loc=left"); } - auto checkPos = [&](IntegerAttr posAttr, StringRef name) -> LogicalResult { - if (!posAttr || !posAttr.getType().isSignlessInteger(32)) { - return emitOpError() << "expects " << name << " to be an i32 attr"; - } - int64_t value = posAttr.getInt(); - if (value < 0 || value > std::numeric_limits::max()) { + auto checkPos = [&](uint32_t value, StringRef name) -> LogicalResult { + if (value > std::numeric_limits::max()) { return emitOpError() << "expects " << name << " to fit in an unsigned 16-bit integer"; } @@ -5415,7 +5423,7 @@ static bool isA5SupportedTCvtPair(Type srcElem, Type dstElem) { } static LogicalResult verifyConvTileCommon(Operation *op, Type ty, StringRef name, - bool allowLowPrecision = false) { + bool allowLowPrecision) { auto ct = dyn_cast(ty); if (!ct) { return op->emitOpError() << "expects " << name << " to be a !pto.conv_tile"; @@ -5472,7 +5480,7 @@ static LogicalResult verifyConvTileCommon(Operation *op, Type ty, StringRef name } static LogicalResult verifyTileLikeCommon(Operation *op, Type ty, StringRef name, - bool allowLowPrecision = false) { + bool allowLowPrecision) { if (isa(ty)) { return verifyTileBufCommon(op, ty, name, allowLowPrecision); } @@ -8277,62 +8285,68 @@ static LogicalResult verifyTExtractA2A3Mat(TExtractOp op, } static LogicalResult verifyTExtractA2A3(TExtractOp op) { - auto common = verifyTExtractCommon(op, /*allowLowPrecision=*/false); - if (failed(common)) { - return failure(); - } - const TExtractCommon &c = *common; - const bool hasFp = static_cast(op.getFp()); - const bool hasPreQuantScalar = static_cast(op.getPreQuantScalar()); - const bool hasRelu = op.getReluPreMode() != pto::ReluPreMode::NoRelu; - if (isa(getSrc().getType())) { - Type srcTy = getSrc().getType(); - Type dstTy = getDst().getType(); - if (failed(verifyTileLikeCommon(*this, srcTy, "src", + if (isa(op.getSrc().getType())) { + Type srcTy = op.getSrc().getType(); + Type dstTy = op.getDst().getType(); + const bool hasFp = static_cast(op.getFp()); + if (failed(verifyTileLikeCommon(op, srcTy, "src", /*allowLowPrecision=*/false)) || - failed(verifyTileBufCommon(*this, dstTy, "dst", + failed(verifyTileBufCommon(op, dstTy, "dst", /*allowLowPrecision=*/false)) || failed(verifyNonNegativeIndexRowCol( - *getOperation(), getIndexRow(), getIndexCol(), - /*includeIndexAndIntOpsInConstFold=*/hasFp))) { + *op.getOperation(), op.getIndexRow(), op.getIndexCol(), + /*includeIndexAndIntOpsInConstFold=*/hasFp)) || + failed(verifyExtractStaticBoundsCommon( + *op.getOperation(), op.getIndexRow(), op.getIndexCol(), srcTy, + dstTy, /*includeIndexAndIntOpsInConstFold=*/hasFp))) { return failure(); } - if (hasFp || hasPreQuantScalar || hasRelu || op.getAccToVecModeAttr()) { - return emitOpError("expects convtile textract to use the base form"); + if (hasFp || op.getPreQuantScalar() || + op.getReluPreMode() != pto::ReluPreMode::NoRelu || + op.getAccToVecModeAttr()) { + return op.emitOpError("expects convtile textract to use the base form"); } auto srcCT = dyn_cast(srcTy); auto srcSpace = getPTOMemorySpaceEnum(srcTy); auto dstSpace = getPTOMemorySpaceEnum(dstTy); if (!srcSpace || *srcSpace != pto::AddressSpace::MAT) { - return emitOpError("expects convtile textract src to use loc=mat"); + return op.emitOpError("expects convtile textract src to use loc=mat"); } if (!dstSpace || *dstSpace != pto::AddressSpace::RIGHT) { - return emitOpError("expects convtile textract dst to use loc=right"); + return op.emitOpError("expects convtile textract dst to use loc=right"); } auto srcLayout = dyn_cast_or_null(srcCT.getLayout()); if (!srcLayout || (srcLayout.getLayout() != pto::Layout::FRACTAL_Z && srcLayout.getLayout() != pto::Layout::FRACTAL_Z_3D)) { - return emitOpError( + return op.emitOpError( "expects convtile textract src to use layout=FRACTAL_Z or FRACTAL_Z_3D"); } auto dstTb = dyn_cast(dstTy); if (!dstTb || dstTb.getBLayoutValueI32() != static_cast(pto::BLayout::RowMajor) || dstTb.getSLayoutValueI32() != static_cast(pto::SLayout::ColMajor)) { - return emitOpError( + return op.emitOpError( "expects convtile textract dst to use blayout=row_major and slayout=col_major"); } Type srcElem = getElemTy(srcTy); Type dstElem = getElemTy(dstTy); if (!srcElem || !dstElem || srcElem != dstElem) { - return emitOpError("expects convtile textract src and dst to have the same element type"); + return op.emitOpError("expects convtile textract src and dst to have the same element type"); } if (!isConvTileExtractElem(srcElem)) { - return emitOpError("expects convtile textract element type to be i8/i16/i32/f16/bf16/f32"); + return op.emitOpError("expects convtile textract element type to be i8/i16/i32/f16/bf16/f32"); } return success(); } + auto common = verifyTExtractCommon(op, /*allowLowPrecision=*/false); + if (failed(common)) { + return failure(); + } + const TExtractCommon &c = *common; + const bool hasFp = static_cast(op.getFp()); + const bool hasPreQuantScalar = static_cast(op.getPreQuantScalar()); + const bool hasRelu = op.getReluPreMode() != pto::ReluPreMode::NoRelu; if (!isA2A3ExtractElemType(c.dstElem) && !(hasFp && c.dstElem.isInteger(16))) { return op.emitOpError("expects A2/A3 textract element type to be i8/f16/bf16/f32"); } @@ -8429,61 +8443,67 @@ static LogicalResult verifyTExtractA5Acc(TExtractOp op, } static LogicalResult verifyTExtractA5(TExtractOp op) { - auto common = verifyTExtractCommon(op, /*allowLowPrecision=*/true); - if (failed(common)) { - return failure(); - } - const TExtractCommon &c = *common; - const bool hasPreQuantScalar = static_cast(op.getPreQuantScalar()); - const bool hasRelu = op.getReluPreMode() != pto::ReluPreMode::NoRelu; - if (isa(getSrc().getType())) { - Type srcTy = getSrc().getType(); - Type dstTy = getDst().getType(); - if (failed(verifyTileLikeCommon(*this, srcTy, "src", + if (isa(op.getSrc().getType())) { + Type srcTy = op.getSrc().getType(); + Type dstTy = op.getDst().getType(); + const bool hasFp = static_cast(op.getFp()); + const bool hasPreQuantScalar = static_cast(op.getPreQuantScalar()); + const bool hasRelu = op.getReluPreMode() != pto::ReluPreMode::NoRelu; + if (failed(verifyTileLikeCommon(op, srcTy, "src", /*allowLowPrecision=*/true)) || - failed(verifyTileBufCommon(*this, dstTy, "dst", + failed(verifyTileBufCommon(op, dstTy, "dst", /*allowLowPrecision=*/true)) || failed(verifyNonNegativeIndexRowCol( - *getOperation(), getIndexRow(), getIndexCol(), - /*includeIndexAndIntOpsInConstFold=*/hasFp))) { + *op.getOperation(), op.getIndexRow(), op.getIndexCol(), + /*includeIndexAndIntOpsInConstFold=*/hasFp)) || + failed(verifyExtractStaticBoundsCommon( + *op.getOperation(), op.getIndexRow(), op.getIndexCol(), srcTy, + dstTy, /*includeIndexAndIntOpsInConstFold=*/hasFp))) { return failure(); } if (hasFp || hasPreQuantScalar || hasRelu || op.getAccToVecModeAttr()) { - return emitOpError("expects convtile textract to use the base form"); + return op.emitOpError("expects convtile textract to use the base form"); } auto srcCT = dyn_cast(srcTy); auto srcSpace = getPTOMemorySpaceEnum(srcTy); auto dstSpace = getPTOMemorySpaceEnum(dstTy); if (!srcSpace || *srcSpace != pto::AddressSpace::MAT) { - return emitOpError("expects convtile textract src to use loc=mat"); + return op.emitOpError("expects convtile textract src to use loc=mat"); } if (!dstSpace || *dstSpace != pto::AddressSpace::RIGHT) { - return emitOpError("expects convtile textract dst to use loc=right"); + return op.emitOpError("expects convtile textract dst to use loc=right"); } auto srcLayout = dyn_cast_or_null(srcCT.getLayout()); if (!srcLayout || (srcLayout.getLayout() != pto::Layout::FRACTAL_Z && srcLayout.getLayout() != pto::Layout::FRACTAL_Z_3D)) { - return emitOpError( + return op.emitOpError( "expects convtile textract src to use layout=FRACTAL_Z or FRACTAL_Z_3D"); } auto dstTb = dyn_cast(dstTy); if (!dstTb || dstTb.getBLayoutValueI32() != static_cast(pto::BLayout::RowMajor) || dstTb.getSLayoutValueI32() != static_cast(pto::SLayout::ColMajor)) { - return emitOpError( + return op.emitOpError( "expects convtile textract dst to use blayout=row_major and slayout=col_major"); } Type srcElem = getElemTy(srcTy); Type dstElem = getElemTy(dstTy); if (!srcElem || !dstElem || srcElem != dstElem) { - return emitOpError("expects convtile textract src and dst to have the same element type"); + return op.emitOpError("expects convtile textract src and dst to have the same element type"); } if (!isConvTileExtractElem(srcElem)) { - return emitOpError("expects convtile textract element type to be i8/i16/i32/f16/bf16/f32"); + return op.emitOpError("expects convtile textract element type to be i8/i16/i32/f16/bf16/f32"); } return success(); } + auto common = verifyTExtractCommon(op, /*allowLowPrecision=*/true); + if (failed(common)) { + return failure(); + } + const TExtractCommon &c = *common; + const bool hasPreQuantScalar = static_cast(op.getPreQuantScalar()); + const bool hasRelu = op.getReluPreMode() != pto::ReluPreMode::NoRelu; if (!isA5ExtractElemType(c.dstElem)) { return op.emitOpError("expects A5 textract element type to be an fp8/f16/bf16/f32 or int8 family type"); } diff --git a/lib/PTO/IR/PTOAttrs.cpp b/lib/PTO/IR/PTOAttrs.cpp index 5855423a6b..011d65c42f 100644 --- a/lib/PTO/IR/PTOAttrs.cpp +++ b/lib/PTO/IR/PTOAttrs.cpp @@ -271,6 +271,8 @@ constexpr unsigned kConvTilePadListSize = 4; constexpr int32_t kConvTileDefaultZero = 0; constexpr int32_t kConvTileDefaultOne = 1; +static LogicalResult parseTileBufKeyEq(AsmParser &parser, StringRef expectedKey); + static LogicalResult parseConvTileIntField(AsmParser &parser, StringRef key, IntegerAttr &value) { if (failed(parseTileBufKeyEq(parser, key))) { @@ -364,7 +366,7 @@ ConvTileConfigAttr ConvTileConfigAttr::getDefault(MLIRContext *ctx) { auto transpose = BoolAttr::get(ctx, false); return ConvTileConfigAttr::get(ctx, zero, zero, padList, one, one, one, one, one, one, padValue, zero, zero, one, zero, - zero, transpose); + zero, zero, transpose); } bool ConvTileConfigAttr::isDefault() const { @@ -395,7 +397,7 @@ LogicalResult ConvTileConfigAttr::verify(function_ref emit IntegerAttr dilationW, IntegerAttr strideH, IntegerAttr strideW, - Attribute padValue, + Attribute padValueAttr, IntegerAttr channelSize, IntegerAttr repeatStride, IntegerAttr repeatTime, @@ -463,7 +465,9 @@ LogicalResult ConvTileConfigAttr::verify(function_ref emit return failure(); } - if (!padValue || (!isa(padValue) && !isa(padValue))) { + if (!padValueAttr || + (!llvm::isa(padValueAttr) && + !llvm::isa(padValueAttr))) { return emitError() << "pad_value must be an integer or float attr", failure(); } diff --git a/lib/PTO/IR/PTOTypeDefs.cpp b/lib/PTO/IR/PTOTypeDefs.cpp index 001041b991..75e60dc651 100644 --- a/lib/PTO/IR/PTOTypeDefs.cpp +++ b/lib/PTO/IR/PTOTypeDefs.cpp @@ -884,7 +884,7 @@ static Type buildConvTileType(AsmParser &parser, return Type(); } - auto layoutAttr = dyn_cast_or_null(fields.layoutAttr); + auto layoutAttr = llvm::dyn_cast_or_null(fields.layoutAttr); if (!layoutAttr) { emitError() << "layout must be a pto.layout attr"; return Type(); @@ -951,7 +951,7 @@ Type ConvTileType::parse(AsmParser &parser) { if (failed(parser.parseAttribute(attr))) { return Type(); } - layoutAttr = dyn_cast(attr); + layoutAttr = llvm::dyn_cast(attr); if (!layoutAttr || seenLayout) { return Type(); } @@ -964,7 +964,7 @@ Type ConvTileType::parse(AsmParser &parser) { if (failed(parser.parseAttribute(attr))) { return Type(); } - bufferSizeAttr = dyn_cast_or_null(attr); + bufferSizeAttr = llvm::dyn_cast_or_null(attr); if (!bufferSizeAttr || seenBufferSize) { return Type(); } @@ -977,7 +977,7 @@ Type ConvTileType::parse(AsmParser &parser) { if (failed(parser.parseAttribute(attr))) { return Type(); } - configAttr = dyn_cast_or_null(attr); + configAttr = llvm::dyn_cast_or_null(attr); if (!configAttr || seenConfig) { return Type(); } diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 3880371d47..2a1ce2e68b 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -259,11 +259,11 @@ static FailureOr computeStaticBufferBytes(Value value) { if (auto convType = dyn_cast(value.getType())) { uint64_t elemBytes = getPTOStorageElemByteSize(convType.getElementType()); const bool invalidCapacity = - elemBytes == 0 || convType.getBufferSize() <= 0; + elemBytes == 0 || convType.getBufferSizeValue() <= 0; if (invalidCapacity) { return failure(); } - uint64_t bufferSize = static_cast(convType.getBufferSize()); + uint64_t bufferSize = static_cast(convType.getBufferSizeValue()); const bool overflows = bufferSize > std::numeric_limits::max() / elemBytes; if (overflows) { diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index cd811fdfd7..393c17988a 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -424,6 +424,24 @@ static std::string layoutToEmitCString(mlir::pto::Layout layout) { return "pto::Layout::MX_A_ZZ"; case mlir::pto::Layout::MX_B_NN: return "pto::Layout::MX_B_NN"; + case mlir::pto::Layout::NCHW: + return "pto::Layout::NCHW"; + case mlir::pto::Layout::NC1HWC0: + return "pto::Layout::NC1HWC0"; + case mlir::pto::Layout::NCDHW: + return "pto::Layout::NCDHW"; + case mlir::pto::Layout::NDC1HWC0: + return "pto::Layout::NDC1HWC0"; + case mlir::pto::Layout::GNCHW: + return "pto::Layout::GNCHW"; + case mlir::pto::Layout::GNC1HWC0: + return "pto::Layout::GNC1HWC0"; + case mlir::pto::Layout::NHWC: + return "pto::Layout::NHWC"; + case mlir::pto::Layout::FRACTAL_Z: + return "pto::Layout::FRACTAL_Z"; + case mlir::pto::Layout::FRACTAL_Z_3D: + return "pto::Layout::FRACTAL_Z_3D"; } return "pto::Layout::ND"; } @@ -12462,9 +12480,9 @@ struct PTOTImg2colToEmitC : public OpConversionPattern { Value src = peelUnrealized(adaptor.getSrc()); Type u16Ty = emitc::OpaqueType::get(ctx, "uint16_t"); Value posM = makeEmitCIntConstant(rewriter, loc, u16Ty, - static_cast(op.getPosM().getInt())); + static_cast(op.getPosM())); Value posK = makeEmitCIntConstant(rewriter, loc, u16Ty, - static_cast(op.getPosK().getInt())); + static_cast(op.getPosK())); auto templateArgs = getFmatrixModeTemplateArgs(rewriter, op.getFmatrixMode()); createLastUseAwareOpaqueCall(rewriter, op.getOperation(), TypeRange{}, diff --git a/lib/PTO/Transforms/Utils.cpp b/lib/PTO/Transforms/Utils.cpp index c62f6efe9a..1ceaab019e 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -509,11 +509,11 @@ static std::optional getStaticTileBytes(TileBufType type) { static std::optional getStaticConvTileBytes(ConvTileType type) { unsigned elemBytes = getPTOStorageElemByteSize(type.getElementType()); - const bool invalidCapacity = elemBytes == 0 || type.getBufferSize() <= 0; + const bool invalidCapacity = elemBytes == 0 || type.getBufferSizeValue() <= 0; if (invalidCapacity) { return std::nullopt; } - uint64_t bufferSize = static_cast(type.getBufferSize()); + uint64_t bufferSize = static_cast(type.getBufferSizeValue()); const bool overflows = bufferSize > std::numeric_limits::max() / elemBytes; if (overflows) { From 10fbd7a134ff3a184263e071ed1b3f9b0b622ae5 Mon Sep 17 00:00:00 2001 From: andodo Date: Mon, 31 Aug 2026 17:21:44 +0800 Subject: [PATCH 05/12] Fix compile error. --- lib/PTO/IR/PTOAttrs.cpp | 8 ++++++++ lib/PTO/IR/PTOTypeDefs.cpp | 10 ++-------- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/lib/PTO/IR/PTOAttrs.cpp b/lib/PTO/IR/PTOAttrs.cpp index 011d65c42f..48862a9bdf 100644 --- a/lib/PTO/IR/PTOAttrs.cpp +++ b/lib/PTO/IR/PTOAttrs.cpp @@ -33,6 +33,14 @@ constexpr int32_t kCompactModeNull = static_cast(CompactMode::Null); constexpr int32_t kCompactModeRowPlusOne = static_cast(CompactMode::RowPlusOne); +static LogicalResult parseTileBufKeyEq(AsmParser &parser, + StringRef expectedKey) { + if (failed(parser.parseKeyword(expectedKey))) { + return failure(); + } + return parser.parseEqual(); +} + } // namespace TileBufConfigAttr TileBufConfigAttr::getDefault(MLIRContext *ctx) { diff --git a/lib/PTO/IR/PTOTypeDefs.cpp b/lib/PTO/IR/PTOTypeDefs.cpp index 75e60dc651..8ba8a1a486 100644 --- a/lib/PTO/IR/PTOTypeDefs.cpp +++ b/lib/PTO/IR/PTOTypeDefs.cpp @@ -808,9 +808,6 @@ static std::optional computeConvTileBufferSize(ArrayRef shape) static LogicalResult parseConvTileLayoutField(AsmParser &parser, Attribute &layoutAttr) { - if (failed(parseTileBufKeyEq(parser, "layout"))) { - return failure(); - } if (failed(parser.parseAttribute(layoutAttr))) { return failure(); } @@ -819,9 +816,6 @@ static LogicalResult parseConvTileLayoutField(AsmParser &parser, static LogicalResult parseConvTileConfigField(AsmParser &parser, Attribute &configAttr) { - if (failed(parseTileBufKeyEq(parser, "config"))) { - return failure(); - } if (failed(parser.parseAttribute(configAttr))) { return failure(); } @@ -948,7 +942,7 @@ Type ConvTileType::parse(AsmParser &parser) { if (key == "layout") { Attribute attr; - if (failed(parser.parseAttribute(attr))) { + if (failed(parseConvTileLayoutField(parser, attr))) { return Type(); } layoutAttr = llvm::dyn_cast(attr); @@ -974,7 +968,7 @@ Type ConvTileType::parse(AsmParser &parser) { if (key == "config") { Attribute attr; - if (failed(parser.parseAttribute(attr))) { + if (failed(parseConvTileConfigField(parser, attr))) { return Type(); } configAttr = llvm::dyn_cast_or_null(attr); From 1de04141e0014f10fabaf4ab455841c735e8543b Mon Sep 17 00:00:00 2001 From: andodo Date: Tue, 1 Sep 2026 14:35:25 +0800 Subject: [PATCH 06/12] Fix CI error. --- lib/PTO/IR/PTO.cpp | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index c9e0608100..6e15ea4817 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -18621,6 +18621,12 @@ void TPrefetchOp::getEffects( addEffect(effects, &getDstMutable(), MemoryEffects::Write::get()); } +void TImg2colOp::getEffects( + SmallVectorImpl> &effects) { + addEffect(effects, &getSrcMutable(), MemoryEffects::Read::get()); + addEffect(effects, &getDstMutable(), MemoryEffects::Write::get()); +} + // === TAbsOp === // Read: src, Write: dst void TAbsOp::getEffects( From 5a489fa0e5ffc638855f526d04cff7fa1f4f5257 Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 2 Sep 2026 15:29:51 +0800 Subject: [PATCH 07/12] Error fix. --- lib/PTO/IR/PTOAttrs.cpp | 56 +++++----------- lib/PTO/IR/PTOTypeDefs.cpp | 2 +- lib/PTO/Transforms/PTOToEmitC.cpp | 105 ++++++++++++++++++++++++++++-- 3 files changed, 118 insertions(+), 45 deletions(-) diff --git a/lib/PTO/IR/PTOAttrs.cpp b/lib/PTO/IR/PTOAttrs.cpp index 48862a9bdf..94dda7a176 100644 --- a/lib/PTO/IR/PTOAttrs.cpp +++ b/lib/PTO/IR/PTOAttrs.cpp @@ -33,14 +33,6 @@ constexpr int32_t kCompactModeNull = static_cast(CompactMode::Null); constexpr int32_t kCompactModeRowPlusOne = static_cast(CompactMode::RowPlusOne); -static LogicalResult parseTileBufKeyEq(AsmParser &parser, - StringRef expectedKey) { - if (failed(parser.parseKeyword(expectedKey))) { - return failure(); - } - return parser.parseEqual(); -} - } // namespace TileBufConfigAttr TileBufConfigAttr::getDefault(MLIRContext *ctx) { @@ -279,13 +271,8 @@ constexpr unsigned kConvTilePadListSize = 4; constexpr int32_t kConvTileDefaultZero = 0; constexpr int32_t kConvTileDefaultOne = 1; -static LogicalResult parseTileBufKeyEq(AsmParser &parser, StringRef expectedKey); - -static LogicalResult parseConvTileIntField(AsmParser &parser, StringRef key, +static LogicalResult parseConvTileIntField(AsmParser &parser, IntegerAttr &value) { - if (failed(parseTileBufKeyEq(parser, key))) { - return failure(); - } int64_t parsed = 0; if (failed(parser.parseInteger(parsed))) { return failure(); @@ -297,9 +284,6 @@ static LogicalResult parseConvTileIntField(AsmParser &parser, StringRef key, static LogicalResult parseConvTileBoolField(AsmParser &parser, StringRef key, BoolAttr &value) { - if (failed(parseTileBufKeyEq(parser, key))) { - return failure(); - } bool parsed = false; if (succeeded(parser.parseOptionalKeyword("true"))) { parsed = true; @@ -314,11 +298,8 @@ static LogicalResult parseConvTileBoolField(AsmParser &parser, StringRef key, return success(); } -static LogicalResult parseConvTileAttrField(AsmParser &parser, StringRef key, +static LogicalResult parseConvTileAttrField(AsmParser &parser, Attribute &value) { - if (failed(parseTileBufKeyEq(parser, key))) { - return failure(); - } if (failed(parser.parseAttribute(value))) { return failure(); } @@ -327,9 +308,6 @@ static LogicalResult parseConvTileAttrField(AsmParser &parser, StringRef key, static LogicalResult parseConvTilePadListField(AsmParser &parser, SmallVectorImpl &padList) { - if (failed(parseTileBufKeyEq(parser, "pad_list"))) { - return failure(); - } SmallVector parsed; if (failed(parser.parseDimensionList(parsed, /*allowDynamic=*/false, /*withTrailingX=*/false))) { @@ -535,7 +513,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { } if (key == "fmap_h") { - if (failed(parseConvTileIntField(p, key, fmapH))) { + if (failed(parseConvTileIntField(p, fmapH))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -544,7 +522,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "fmap_w") { - if (failed(parseConvTileIntField(p, key, fmapW))) { + if (failed(parseConvTileIntField(p, fmapW))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -562,7 +540,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "filter_h") { - if (failed(parseConvTileIntField(p, key, filterH))) { + if (failed(parseConvTileIntField(p, filterH))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -571,7 +549,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "filter_w") { - if (failed(parseConvTileIntField(p, key, filterW))) { + if (failed(parseConvTileIntField(p, filterW))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -580,7 +558,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "dilation_h") { - if (failed(parseConvTileIntField(p, key, dilationH))) { + if (failed(parseConvTileIntField(p, dilationH))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -589,7 +567,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "dilation_w") { - if (failed(parseConvTileIntField(p, key, dilationW))) { + if (failed(parseConvTileIntField(p, dilationW))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -598,7 +576,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "stride_h") { - if (failed(parseConvTileIntField(p, key, strideH))) { + if (failed(parseConvTileIntField(p, strideH))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -607,7 +585,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "stride_w") { - if (failed(parseConvTileIntField(p, key, strideW))) { + if (failed(parseConvTileIntField(p, strideW))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -616,7 +594,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "pad_value") { - if (failed(parseConvTileAttrField(p, key, padValue))) { + if (failed(parseConvTileAttrField(p, padValue))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -625,7 +603,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "channel_size") { - if (failed(parseConvTileIntField(p, key, channelSize))) { + if (failed(parseConvTileIntField(p, channelSize))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -634,7 +612,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "repeat_stride") { - if (failed(parseConvTileIntField(p, key, repeatStride))) { + if (failed(parseConvTileIntField(p, repeatStride))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -643,7 +621,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "repeat_time") { - if (failed(parseConvTileIntField(p, key, repeatTime))) { + if (failed(parseConvTileIntField(p, repeatTime))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -652,7 +630,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "repeat_mode") { - if (failed(parseConvTileIntField(p, key, repeatMode))) { + if (failed(parseConvTileIntField(p, repeatMode))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -661,7 +639,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "dst_stride") { - if (failed(parseConvTileIntField(p, key, dstStride))) { + if (failed(parseConvTileIntField(p, dstStride))) { return {}; } if (failed(consumeFieldTerminator())) { @@ -670,7 +648,7 @@ Attribute ConvTileConfigAttr::parse(AsmParser &p, Type) { continue; } if (key == "dst_mposition") { - if (failed(parseConvTileIntField(p, key, dstMposition))) { + if (failed(parseConvTileIntField(p, dstMposition))) { return {}; } if (failed(consumeFieldTerminator())) { diff --git a/lib/PTO/IR/PTOTypeDefs.cpp b/lib/PTO/IR/PTOTypeDefs.cpp index 8ba8a1a486..3e8a5fcee5 100644 --- a/lib/PTO/IR/PTOTypeDefs.cpp +++ b/lib/PTO/IR/PTOTypeDefs.cpp @@ -910,7 +910,7 @@ Type ConvTileType::parse(AsmParser &parser) { } if (failed(parser.parseDimensionList(fields.shape, /*allowDynamic=*/false, - /*withTrailingX=*/false))) { + /*withTrailingX=*/true))) { return Type(); } diff --git a/lib/PTO/Transforms/PTOToEmitC.cpp b/lib/PTO/Transforms/PTOToEmitC.cpp index 393c17988a..b43dd1be7b 100644 --- a/lib/PTO/Transforms/PTOToEmitC.cpp +++ b/lib/PTO/Transforms/PTOToEmitC.cpp @@ -336,6 +336,47 @@ static void createLastUseAwareOpaqueCall( operands); } +static void emitConvTileConfigInitCall(ConversionPatternRewriter &rewriter, + Location loc, Value tile, + pto::ConvTileType convTy) { + auto *ctx = rewriter.getContext(); + auto cfg = convTy.getConfigAttr(); + + SmallVector args; + args.push_back(IntegerAttr::get(IndexType::get(ctx), 0)); + + auto appendI32 = [&](IntegerAttr attr) { + args.push_back(IntegerAttr::get(rewriter.getI32Type(), attr.getInt())); + }; + auto appendPad = [&](int64_t value) { + args.push_back(IntegerAttr::get(rewriter.getI32Type(), value)); + }; + + appendI32(cfg.getFmapH()); + appendI32(cfg.getFmapW()); + for (int64_t pad : cfg.getPadList()) + appendPad(pad); + appendI32(cfg.getFilterH()); + appendI32(cfg.getFilterW()); + appendI32(cfg.getDilationH()); + appendI32(cfg.getDilationW()); + appendI32(cfg.getStrideH()); + appendI32(cfg.getStrideW()); + args.push_back(cfg.getPadValue()); + appendI32(cfg.getChannelSize()); + appendI32(cfg.getRepeatStride()); + appendI32(cfg.getRepeatTime()); + appendI32(cfg.getRepeatMode()); + appendI32(cfg.getDstStride()); + appendI32(cfg.getDstMposition()); + args.push_back(cfg.getTranspose()); + + rewriter.create( + loc, TypeRange{}, "PTOAS__INIT_CONVTILE_CONFIG", + /*args=*/rewriter.getArrayAttr(args), + /*templateArgs=*/ArrayAttr{}, /*operands=*/ValueRange{tile}); +} + static StringRef getFmatrixModeToken(pto::FmatrixMode mode) { switch (mode) { case pto::FmatrixMode::FMATRIX_A_AUTO: @@ -12618,19 +12659,30 @@ struct PTOAllocTileToEmitC ConversionPatternRewriter &rewriter) const override { Location loc = op.getLoc(); MLIRContext *ctx = rewriter.getContext(); - auto tileTy = cast(op.getResult().getType()); - auto tileTypeString = getEmitCTileTypeString(tileTy); + Type resultTy = op.getResult().getType(); + auto tileTy = dyn_cast(resultTy); + auto convTy = dyn_cast(resultTy); + if (!tileTy && !convTy) + return rewriter.notifyMatchFailure( + op, "expected alloc_tile to produce a tile_buf or conv_tile"); + + std::optional tileTypeString; + if (tileTy) + tileTypeString = getEmitCTileTypeString(tileTy); + else + tileTypeString = getEmitCConvTileTypeString(convTy); if (!tileTypeString) return rewriter.notifyMatchFailure( op, "only rank-2 alloc_tile handles can be converted to EmitC"); - Type convertedTy = getTypeConverter()->convertType(tileTy); + Type convertedTy = getTypeConverter()->convertType(resultTy); if (!convertedTy) convertedTy = emitc::OpaqueType::get(ctx, *tileTypeString); - auto validShape = tileTy.getValidShape(); + ArrayRef validShape = tileTy ? tileTy.getValidShape() + : ArrayRef{}; bool hasDynamicValidDim = - llvm::any_of(validShape, [](int64_t dim) { return dim < 0; }); + tileTy && llvm::any_of(validShape, [](int64_t dim) { return dim < 0; }); bool useConstructor = hasDynamicValidDim; SmallVector constructorArgs; @@ -12686,6 +12738,9 @@ struct PTOAllocTileToEmitC tile = loadEmitCVariableIfNeeded(rewriter, loc, tile); } + if (convTy) + emitConvTileConfigInitCall(rewriter, loc, tile, convTy); + Value addr = adaptor.getAddr(); if (addr) { addr = peelUnrealized(addr); @@ -13935,6 +13990,7 @@ struct EmitPTOManualPass bool needsEventIdArrayHelper = false; bool needsTRandomHelper = false; bool needsGlobalTensorDataHelper = false; + bool needsConvTileInitHelper = false; mop.walk([&](Operation *op) { if (isa(op)) needsEventIdArrayHelper = true; @@ -13950,6 +14006,10 @@ struct EmitPTOManualPass } if (isa(op)) needsGlobalTensorDataHelper = true; + if (auto alloc = dyn_cast(op)) { + if (isa(alloc.getResult().getType())) + needsConvTileInitHelper = true; + } }); // 1. 插入头文件 @@ -14018,6 +14078,41 @@ static AICORE inline void PTOAS__TRANDOM( TRandomCounter counter = {counter0, counter1, counter2, counter3}; TRANDOM(dst, key, counter); } +)cpp")); + } + if (needsConvTileInitHelper) { + builder.create( + loc, builder.getStringAttr(R"cpp( +template +static AICORE inline void PTOAS__INIT_CONVTILE_CONFIG( + Tile &tile, uint16_t fmapH, uint16_t fmapW, uint8_t padLeft, + uint8_t padRight, uint8_t padTop, uint8_t padBottom, uint16_t filterH, + uint16_t filterW, uint16_t dilationH, uint16_t dilationW, + uint16_t strideH, uint16_t strideW, PadValueT padValue, + uint16_t channelSize, uint16_t repeatStride, uint8_t repeatTime, + uint8_t repeatMode, uint16_t dstStride, uint16_t dstMposition, + bool transpose) { + const uint8_t padList[4] = {padLeft, padRight, padTop, padBottom}; + tile.SetFmapH(fmapH); + tile.SetFmapW(fmapW); + tile.SetPadListArray(padList); + tile.SetFilterH(filterH); + tile.SetFilterW(filterW); + tile.SetDilationH(dilationH); + tile.SetDilationW(dilationW); + tile.SetStrideH(strideH); + tile.SetStrideW(strideW); + tile.SetPadValue(padValue); + tile.SetChannelSize(channelSize); + tile.SetRepeatStride(repeatStride); + tile.SetRepeatTime(repeatTime); + tile.SetRepeatMode(repeatMode); +#ifndef PTO_NPU_ARCH_A2A3 + tile.SetDstStride(dstStride); + tile.SetDstMposition(dstMposition); +#endif + tile.SetTranspose(transpose); +} )cpp")); } builder.create( From a2cd45f2bb03265958824335d76d426c4792e26e Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 2 Sep 2026 15:33:54 +0800 Subject: [PATCH 08/12] Add test cases. --- test/samples/Img2col/timg2col-pto.pto | 64 ++++++++ test/samples/Img2col/timg2col_runtime.py | 146 ++++++++++++++++++ .../Img2col/timg2col_runtime_compare.py | 24 +++ .../Img2col/timg2col_runtime_golden.py | 95 ++++++++++++ test/samples/runop.sh | 6 +- 5 files changed, 333 insertions(+), 2 deletions(-) create mode 100644 test/samples/Img2col/timg2col-pto.pto create mode 100644 test/samples/Img2col/timg2col_runtime.py create mode 100644 test/samples/Img2col/timg2col_runtime_compare.py create mode 100644 test/samples/Img2col/timg2col_runtime_golden.py diff --git a/test/samples/Img2col/timg2col-pto.pto b/test/samples/Img2col/timg2col-pto.pto new file mode 100644 index 0000000000..d57c305b57 --- /dev/null +++ b/test/samples/Img2col/timg2col-pto.pto @@ -0,0 +1,64 @@ +// 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 OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +// See LICENSE in the root of the software repository for the full text of the License. + +module attributes {pto.target_arch = "a5"} { + func.func @timg2col_sample( + %src: !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>>, + %dst: !pto.tile_buf) { + pto.set_fmatrix %src {fmatrix_mode = #pto.fmatrix_mode} + : !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>> + pto.set_img2col_rpt %src {fmatrix_mode = #pto.fmatrix_mode} + : !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>> + pto.set_img2col_padding %src {fmatrix_mode = #pto.fmatrix_mode} + : !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>> + pto.timg2col %dst, %src {posM = 1 : i32, posK = 8 : i32, + fmatrix_mode = #pto.fmatrix_mode} + : !pto.tile_buf, + !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>> + return + } +} diff --git a/test/samples/Img2col/timg2col_runtime.py b/test/samples/Img2col/timg2col_runtime.py new file mode 100644 index 0000000000..af1ba76ab2 --- /dev/null +++ b/test/samples/Img2col/timg2col_runtime.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. Please ensure you do 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. + +import os + +from ptoas.mlir.dialects import arith, func, pto +from ptoas.mlir.ir import ( + BoolAttr, + Context, + F32Type, + InsertionPoint, + IndexType, + IntegerAttr, + IntegerType, + Location, + Module, + StringAttr, + UnitAttr, +) + + +def build(): + with Context() as ctx: + pto.register_dialect(ctx, load=True) + with Location.unknown(ctx): + module = Module.create() + arch = os.environ.get("PTOAS_SAMPLE_ARCH", "a5") + module.operation.attributes["pto.target_arch"] = StringAttr.get(arch) + + f32 = F32Type.get(ctx) + i32 = IntegerType.get_signless(32, ctx) + i64 = IntegerType.get_signless(64, ctx) + + ptr_f32 = pto.PtrType.get(f32, ctx) + tv5_f32 = pto.TensorViewType.get(5, f32, ctx) + src_ptv = pto.PartitionTensorViewType.get([1, 1, 3, 4, 8], f32, ctx) + dst_tv = pto.TensorViewType.get(2, f32, ctx) + dst_ptv = pto.PartitionTensorViewType.get([16, 32], f32, ctx) + + vec = pto.AddressSpaceAttr.get(pto.AddressSpace.VEC, ctx) + left = pto.AddressSpaceAttr.get(pto.AddressSpace.LEFT, ctx) + mat = pto.AddressSpaceAttr.get(pto.AddressSpace.MAT, ctx) + layout = pto.LayoutAttr.get(pto.Layout.NC1HWC0, ctx) + col_major = pto.BLayoutAttr.get(pto.BLayout.ColMajor, ctx) + row_box = pto.SLayoutAttr.get(pto.SLayout.RowMajor, ctx) + null_pad = pto.PadValueAttr.get(pto.PadValue.Null, ctx) + tile_cfg = pto.TileBufConfigAttr.get( + col_major, + row_box, + pto.TileConfig.fractalABSize, + null_pad, + ctx, + ) + dst_tile_ty = pto.TileBufType.get([16, 32], f32, left, [9, 32], tile_cfg, ctx) + + conv_cfg = pto.ConvTileConfigAttr.get( + IntegerAttr.get(i32, 3), + IntegerAttr.get(i32, 4), + [1, 0, 1, 0], + IntegerAttr.get(i32, 2), + IntegerAttr.get(i32, 2), + IntegerAttr.get(i32, 1), + IntegerAttr.get(i32, 1), + IntegerAttr.get(i32, 1), + IntegerAttr.get(i32, 1), + IntegerAttr.get(i32, 0), + IntegerAttr.get(i32, 8), + IntegerAttr.get(i32, 0), + IntegerAttr.get(i32, 1), + IntegerAttr.get(i32, 0), + IntegerAttr.get(i32, 1), + IntegerAttr.get(i32, 0), + BoolAttr.get(False), + ) + src_tile_ty = pto.ConvTileType.get( + [1, 1, 3, 4, 8], + f32, + IntegerAttr.get(i64, 1 * 1 * 3 * 4 * 8), + mat, + layout, + conv_cfg, + ctx, + ) + + fn_ty = func.FunctionType.get([ptr_f32, ptr_f32], []) + with InsertionPoint(module.body): + fn = func.FuncOp("timg2col_runtime_kernel", fn_ty) + fn.operation.attributes["pto.entry"] = UnitAttr.get(ctx) + entry = fn.add_entry_block() + + with InsertionPoint(entry): + c0 = arith.ConstantOp(IndexType.get(ctx), 0).result + c1 = arith.ConstantOp(IndexType.get(ctx), 1).result + c3 = arith.ConstantOp(IndexType.get(ctx), 3).result + c4 = arith.ConstantOp(IndexType.get(ctx), 4).result + c8 = arith.ConstantOp(IndexType.get(ctx), 8).result + c16 = arith.ConstantOp(IndexType.get(ctx), 16).result + c32 = arith.ConstantOp(IndexType.get(ctx), 32).result + c96 = arith.ConstantOp(IndexType.get(ctx), 96).result + src_ptr, dst_ptr = entry.arguments + + src_view = pto.MakeTensorViewOp( + tv5_f32, src_ptr, [c1, c1, c3, c4, c8], [c96, c96, c32, c8, c1] + ).result + dst_view = pto.MakeTensorViewOp( + dst_tv, dst_ptr, [c16, c32], [c32, c1] + ).result + + src_part = pto.PartitionViewOp( + src_ptv, + src_view, + offsets=[c0, c0, c0, c0, c0], + sizes=[c1, c1, c3, c4, c8], + ).result + dst_part = pto.PartitionViewOp( + dst_ptv, dst_view, offsets=[c0, c0], sizes=[c16, c32] + ).result + + src_tile = pto.AllocTileOp(src_tile_ty).result + dst_tile = pto.AllocTileOp(dst_tile_ty).result + pto.TLoadOp(None, src_part, src_tile) + pto.SetFmatrixOp(src_tile, pto.FmatrixModeAttr.get(pto.FmatrixMode.FMATRIX_B_MANUAL, ctx)) + pto.SetImg2colRptOp(src_tile, pto.FmatrixModeAttr.get(pto.FmatrixMode.FMATRIX_B_MANUAL, ctx)) + pto.SetImg2colPaddingOp(src_tile, pto.FmatrixModeAttr.get(pto.FmatrixMode.FMATRIX_B_MANUAL, ctx)) + pto.TImg2colOp( + dst_tile, + src_tile, + posM=1, + posK=8, + fmatrixMode=pto.FmatrixModeAttr.get(pto.FmatrixMode.FMATRIX_B_MANUAL, ctx), + ) + pto.TStoreOp(None, dst_tile, dst_part) + func.ReturnOp([]) + + module.operation.verify() + return module + + +if __name__ == "__main__": + print(build()) diff --git a/test/samples/Img2col/timg2col_runtime_compare.py b/test/samples/Img2col/timg2col_runtime_compare.py new file mode 100644 index 0000000000..b73696a12b --- /dev/null +++ b/test/samples/Img2col/timg2col_runtime_compare.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. Please ensure you do not use this file except in compliance with the +# License. THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS 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. + +from pathlib import Path +import sys + +import numpy as np + +for search_root in (Path(__file__).resolve().parent, Path(__file__).resolve().parents[1]): + if (search_root / "validation_runtime.py").is_file(): + sys.path.insert(0, str(search_root)) + break + +from validation_runtime import compare_outputs + + +if __name__ == "__main__": + compare_outputs(np.float32, atol=0.001) diff --git a/test/samples/Img2col/timg2col_runtime_golden.py b/test/samples/Img2col/timg2col_runtime_golden.py new file mode 100644 index 0000000000..92b6386e16 --- /dev/null +++ b/test/samples/Img2col/timg2col_runtime_golden.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python3 +# Copyright (c) 2026 Huawei Technologies Co., Ltd. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. Please ensure you do 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. + +from pathlib import Path +import sys + +import numpy as np + +for search_root in (Path(__file__).resolve().parent, Path(__file__).resolve().parents[1]): + if (search_root / "validation_runtime.py").is_file(): + sys.path.insert(0, str(search_root)) + break + +from validation_runtime import default_buffers, float_values, load_case_meta, rng, write_buffers, write_golden + + +def _src_offset(n: int, c1: int, h: int, w: int, c0: int, *, c1_size: int, h_size: int, w_size: int, c0_size: int) -> int: + return (((n * c1_size + c1) * h_size + h) * w_size + w) * c0_size + c0 + + +def main(): + meta = load_case_meta() + src_name, dst_name = meta.inputs + generator = rng() + src = np.asarray(float_values(generator, meta.elem_counts[src_name], style="signed"), dtype=np.float32) + src = src.reshape(1, 1, 3, 4, 8) + + fmap_h = 3 + fmap_w = 4 + c0_size = 8 + filter_h = 2 + filter_w = 2 + stride_h = 1 + stride_w = 1 + dilation_h = 1 + dilation_w = 1 + pad_left = 1 + pad_right = 0 + pad_top = 1 + pad_bottom = 0 + pos_m = 1 + pos_k = 8 + valid_rows = 9 + valid_cols = 32 + pad_value = np.float32(0.0) + + out_h = (fmap_h + pad_top + pad_bottom - dilation_h * (filter_h - 1) - 1) // stride_h + 1 + out_w = (fmap_w + pad_left + pad_right - dilation_w * (filter_w - 1) - 1) // stride_w + 1 + assert pos_m + valid_rows <= out_h * out_w + golden = np.zeros((16, 32), dtype=np.float32) + + for row in range(valid_rows): + m = pos_m + row + out_row = m // out_w + out_col = m % out_w + for col in range(valid_cols): + k = pos_k + col + c1 = k // (c0_size * filter_h * filter_w) + kernel_offset = (k % (c0_size * filter_h * filter_w)) // c0_size + c0 = k % c0_size + kernel_h = kernel_offset // filter_w + kernel_w = kernel_offset % filter_w + input_h = out_row * stride_h + kernel_h * dilation_h - pad_top + input_w = out_col * stride_w + kernel_w * dilation_w - pad_left + value = pad_value + if 0 <= input_h < fmap_h and 0 <= input_w < fmap_w and c1 < 1: + value = src[ + _src_offset( + 0, + c1, + input_h, + input_w, + c0, + c1_size=1, + h_size=fmap_h, + w_size=fmap_w, + c0_size=c0_size, + ) + ] + golden[row, col] = value + + buffers = default_buffers(meta) + buffers[src_name] = src.reshape(-1) + write_buffers(meta, buffers) + write_golden(meta, {dst_name: golden.reshape(-1)}) + + +if __name__ == "__main__": + main() diff --git a/test/samples/runop.sh b/test/samples/runop.sh index 2dd5282328..fd8f2720d4 100755 --- a/test/samples/runop.sh +++ b/test/samples/runop.sh @@ -35,7 +35,9 @@ for model_path in "${BASE_DIR}"/Qwen* "${BASE_DIR}"/Deepseek*; do ;; esac done -PTO_PTO_DIRS="${PTO_PTO_DIRS:-Sync${MODEL_PTO_DIRS} CommSync Prelu Rem Rems Gemvmx MatmulMxLowPrecision TquantMx TquantMxDn Movfp Interleave DeInterleave PairReduceSum}" +PTO_PTO_DIRS="${PTO_PTO_DIRS:-Sync${MODEL_PTO_DIRS} CommSync Img2col \ +Prelu Rem Rems Gemvmx MatmulMxLowPrecision TquantMx TquantMxDn Movfp \ +Interleave DeInterleave PairReduceSum}" ENABLE_BC=0 usage() { @@ -82,7 +84,7 @@ sample_dir_arch() { case "$1" in TPipe|TAxpy|TColArgMax|TColArgMin|TConcatIdx|\ TRowArgMax|TRowArgMin|Qwen*A3|Deepseek*A3) printf 'a3\n' ;; - Qwen*A5|Deepseek*A5|TquantMx|TquantMxDn|Interleave|DeInterleave|PairReduceSum) printf 'a5\n' ;; + Qwen*A5|Deepseek*A5|Img2col|TquantMx|TquantMxDn|Interleave|DeInterleave|PairReduceSum) printf 'a5\n' ;; esac } From f6d6302e5864d6ea1e5dfeb2c6a53bd293e4ec6e Mon Sep 17 00:00:00 2001 From: andodo Date: Wed, 2 Sep 2026 17:44:57 +0800 Subject: [PATCH 09/12] Add license header. --- test/samples/Img2col/timg2col-pto.pto | 1 + test/samples/Img2col/timg2col_runtime.py | 6 +++--- test/samples/Img2col/timg2col_runtime_compare.py | 6 +++--- test/samples/Img2col/timg2col_runtime_golden.py | 6 +++--- 4 files changed, 10 insertions(+), 9 deletions(-) diff --git a/test/samples/Img2col/timg2col-pto.pto b/test/samples/Img2col/timg2col-pto.pto index d57c305b57..4e72eee957 100644 --- a/test/samples/Img2col/timg2col-pto.pto +++ b/test/samples/Img2col/timg2col-pto.pto @@ -3,6 +3,7 @@ // 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 OR CONDITIONS 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. module attributes {pto.target_arch = "a5"} { diff --git a/test/samples/Img2col/timg2col_runtime.py b/test/samples/Img2col/timg2col_runtime.py index af1ba76ab2..adc70d8a24 100644 --- a/test/samples/Img2col/timg2col_runtime.py +++ b/test/samples/Img2col/timg2col_runtime.py @@ -2,9 +2,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. Please ensure you do 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. +# 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. import os diff --git a/test/samples/Img2col/timg2col_runtime_compare.py b/test/samples/Img2col/timg2col_runtime_compare.py index b73696a12b..ed5d566cb0 100644 --- a/test/samples/Img2col/timg2col_runtime_compare.py +++ b/test/samples/Img2col/timg2col_runtime_compare.py @@ -2,9 +2,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. Please ensure you do not use this file except in compliance with the -# License. THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, -# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR PURPOSE. +# 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. from pathlib import Path diff --git a/test/samples/Img2col/timg2col_runtime_golden.py b/test/samples/Img2col/timg2col_runtime_golden.py index 92b6386e16..1b5036b07e 100644 --- a/test/samples/Img2col/timg2col_runtime_golden.py +++ b/test/samples/Img2col/timg2col_runtime_golden.py @@ -2,9 +2,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. Please ensure you do 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. +# 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. from pathlib import Path From c5e26aad971c36cf85d343b234facc96ca3f10c0 Mon Sep 17 00:00:00 2001 From: andodo Date: Fri, 4 Sep 2026 11:16:23 +0800 Subject: [PATCH 10/12] Add ConvTile Python bindings --- include/pto-c/Dialect/PTO.h | 32 +++++++ lib/Bindings/Python/PTOModule.cpp | 149 ++++++++++++++++++++++++++++++ lib/CAPI/Dialect/PTO.cpp | 110 ++++++++++++++++++++++ python/pto/dialects/pto.py | 8 ++ 4 files changed, 299 insertions(+) diff --git a/include/pto-c/Dialect/PTO.h b/include/pto-c/Dialect/PTO.h index c506384a28..5e93273e39 100644 --- a/include/pto-c/Dialect/PTO.h +++ b/include/pto-c/Dialect/PTO.h @@ -152,6 +152,19 @@ MLIR_CAPI_EXPORTED MlirType mlirPTOTileBufTypeGet( MLIR_CAPI_EXPORTED MlirType mlirPTOTileBufTypeGetWithConfig( MlirContext ctx, intptr_t rank, const int64_t *shape, MlirType elementType, MlirAttribute memorySpace, MlirAttribute config); + +// ---- ConvTileType ---- +MLIR_CAPI_EXPORTED bool mlirPTOTypeIsAConvTileType(MlirType type); + +MLIR_CAPI_EXPORTED MlirType mlirPTOConvTileTypeGet( + MlirContext ctx, intptr_t rank, const int64_t *shape, + MlirType elementType, MlirAttribute bufferSize, + MlirAttribute memorySpace, MlirAttribute layout, MlirAttribute config); + +MLIR_CAPI_EXPORTED intptr_t mlirPTOConvTileTypeGetRank(MlirType type); +MLIR_CAPI_EXPORTED MlirType mlirPTOConvTileTypeGetElementType(MlirType type); +MLIR_CAPI_EXPORTED const int64_t *mlirPTOConvTileTypeGetShape(MlirType type, + intptr_t *numDimsOut); // ---- Enum attrs helpers (BLayout/SLayout/PadValue in mlir::pto) ---- MLIR_CAPI_EXPORTED bool mlirPTOAttrIsABLayoutAttr(MlirAttribute attr); MLIR_CAPI_EXPORTED MlirAttribute mlirPTOBLayoutAttrGet(MlirContext ctx, int32_t value); @@ -219,6 +232,9 @@ MLIR_CAPI_EXPORTED int32_t mlirPTOFmodPrecisionAttrGetValue(MlirAttribute attr); MLIR_CAPI_EXPORTED MlirAttribute mlirPTOSaturationModeAttrGet(MlirContext ctx, int32_t value); MLIR_CAPI_EXPORTED bool mlirPTOAttrIsASaturationModeAttr(MlirAttribute attr); MLIR_CAPI_EXPORTED int32_t mlirPTOSaturationModeAttrGetValue(MlirAttribute attr); +MLIR_CAPI_EXPORTED MlirAttribute mlirPTOFmatrixModeAttrGet(MlirContext ctx, int32_t value); +MLIR_CAPI_EXPORTED bool mlirPTOAttrIsAFmatrixModeAttr(MlirAttribute attr); +MLIR_CAPI_EXPORTED int32_t mlirPTOFmatrixModeAttrGetValue(MlirAttribute attr); // ---- Pipe attr ---- MLIR_CAPI_EXPORTED MlirAttribute mlirPTOPipeAttrGet(MlirContext ctx, int32_t value); MLIR_CAPI_EXPORTED bool mlirPTOAttrIsAPipeAttr(MlirAttribute attr); @@ -293,6 +309,22 @@ MLIR_CAPI_EXPORTED MlirAttribute mlirPTOTileBufConfigAttrGetWithCompactMode( MlirAttribute bLayout, MlirAttribute sLayout, MlirAttribute sFractalSize, MlirAttribute pad, MlirAttribute compactMode); + +// ---- ConvTileConfigAttr ---- +MLIR_CAPI_EXPORTED bool mlirPTOAttrIsAConvTileConfigAttr(MlirAttribute attr); + +MLIR_CAPI_EXPORTED MlirAttribute mlirPTOConvTileConfigAttrGetDefault(MlirContext ctx); + +MLIR_CAPI_EXPORTED MlirAttribute mlirPTOConvTileConfigAttrGet( + MlirContext ctx, MlirAttribute fmapH, MlirAttribute fmapW, + intptr_t padListSize, const int64_t *padList, + MlirAttribute filterH, MlirAttribute filterW, + MlirAttribute dilationH, MlirAttribute dilationW, + MlirAttribute strideH, MlirAttribute strideW, + MlirAttribute padValue, MlirAttribute channelSize, + MlirAttribute repeatStride, MlirAttribute repeatTime, + MlirAttribute repeatMode, MlirAttribute dstStride, + MlirAttribute dstMposition, MlirAttribute transpose); MLIR_CAPI_EXPORTED MlirType mlirPTOTileBufTypeGetWithValidShape( MlirContext ctx, intptr_t rank, const int64_t *shape, MlirType elementType, MlirAttribute memorySpace, intptr_t validRank, const int64_t *validShape); diff --git a/lib/Bindings/Python/PTOModule.cpp b/lib/Bindings/Python/PTOModule.cpp index c11cd3e300..2f75a59d50 100644 --- a/lib/Bindings/Python/PTOModule.cpp +++ b/lib/Bindings/Python/PTOModule.cpp @@ -66,6 +66,17 @@ static MlirContext inferContextFromElementType(MlirContext context, return mlirTypeGetContext(elementType); } +static MlirContext inferContextFromAttribute(py::object contextObj, + MlirAttribute attr) { + if (!contextObj.is_none()) { + return contextObj.cast(); + } + if (mlirAttributeIsNull(attr)) { + throw py::value_error("context is required when attribute is null"); + } + return mlirAttributeGetContext(attr); +} + static int32_t enumValueFromPy(py::object value, const char *attrName, const char *enumName) { if (py::isinstance(value)) { @@ -227,6 +238,12 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { .value("ON", mlir::pto::SaturationMode::ON) .value("OFF", mlir::pto::SaturationMode::OFF); + py::enum_(m, "FmatrixMode") + .value("FMATRIX_A_AUTO", mlir::pto::FmatrixMode::FMATRIX_A_AUTO) + .value("FMATRIX_B_AUTO", mlir::pto::FmatrixMode::FMATRIX_B_AUTO) + .value("FMATRIX_A_MANUAL", mlir::pto::FmatrixMode::FMATRIX_A_MANUAL) + .value("FMATRIX_B_MANUAL", mlir::pto::FmatrixMode::FMATRIX_B_MANUAL); + py::enum_(m, "CmpMode") .value("EQ", MlirPTOCmpMode_EQ) .value("NE", MlirPTOCmpMode_NE) @@ -621,6 +638,10 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { mlirPTOAttrIsAFmodPrecisionAttr, mlirPTOFmodPrecisionAttrGet, mlirPTOFmodPrecisionAttrGetValue); + bindPTOEnumAttr(m, "FmatrixModeAttr", "FmatrixMode", + mlirPTOAttrIsAFmatrixModeAttr, + mlirPTOFmatrixModeAttrGet, + mlirPTOFmatrixModeAttrGetValue); mlir_attribute_subclass( m, "SaturationModeAttr", @@ -1444,6 +1465,74 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { py::arg("context") = py::none(), py::arg("compact_mode") = py::none()); + // ---- ConvTileConfigAttr ---- + mlir_attribute_subclass(m, "ConvTileConfigAttr", + [](MlirAttribute a) -> bool { + return mlirPTOAttrIsAConvTileConfigAttr(a); + }) + .def_classmethod( + "get_default", + [](py::object cls, MlirContext ctx) -> py::object { + MlirAttribute a = mlirPTOConvTileConfigAttrGetDefault(ctx); + if (mlirAttributeIsNull(a)) { + return py::none(); + } + return cls(a); + }, + py::arg("cls"), py::arg("context")) + .def_classmethod( + "get", + [](py::object cls, + MlirAttribute fmapH, + MlirAttribute fmapW, + std::vector padList, + MlirAttribute filterH, + MlirAttribute filterW, + MlirAttribute dilationH, + MlirAttribute dilationW, + MlirAttribute strideH, + MlirAttribute strideW, + MlirAttribute padValue, + MlirAttribute channelSize, + MlirAttribute repeatStride, + MlirAttribute repeatTime, + MlirAttribute repeatMode, + MlirAttribute dstStride, + MlirAttribute dstMposition, + MlirAttribute transpose, + py::object contextObj) -> py::object { + MlirContext ctx = inferContextFromAttribute(contextObj, fmapH); + MlirAttribute a = mlirPTOConvTileConfigAttrGet( + ctx, fmapH, fmapW, static_cast(padList.size()), + padList.data(), filterH, filterW, dilationH, dilationW, + strideH, strideW, padValue, channelSize, repeatStride, + repeatTime, repeatMode, dstStride, dstMposition, + transpose); + if (mlirAttributeIsNull(a)) { + return py::none(); + } + return cls(a); + }, + py::arg("cls"), + py::arg("fmap_h"), + py::arg("fmap_w"), + py::arg("pad_list"), + py::arg("filter_h"), + py::arg("filter_w"), + py::arg("dilation_h"), + py::arg("dilation_w"), + py::arg("stride_h"), + py::arg("stride_w"), + py::arg("pad_value"), + py::arg("channel_size"), + py::arg("repeat_stride"), + py::arg("repeat_time"), + py::arg("repeat_mode"), + py::arg("dst_stride"), + py::arg("dst_mposition"), + py::arg("transpose"), + py::arg("context") = py::none()); + // ---- TileBufType ---- mlir_type_subclass(m, "TileBufType", [](MlirType t) -> bool { return mlirPTOTypeIsATileBufType(t); }) .def_classmethod( @@ -1549,5 +1638,65 @@ void mlir::pto::python::populatePTODialectBindings(pybind11::module_ &m) { return mlirPTOTileBufTypeGetSFractalSize(self); }); + // ---- ConvTileType ---- + mlir_type_subclass(m, "ConvTileType", + [](MlirType t) -> bool { + return mlirPTOTypeIsAConvTileType(t); + }) + .def_classmethod( + "get", + [](py::object cls, + std::vector shape, + MlirType elementType, + MlirAttribute bufferSize, + MlirAttribute memorySpace, + MlirAttribute layout, + py::object configObj, + MlirContext ctx) -> py::object { + ctx = inferContextFromElementType(ctx, elementType); + MlirAttribute config = optionalAttributeFromPy(configObj); + MlirType ty = mlirPTOConvTileTypeGet( + ctx, static_cast(shape.size()), shape.data(), + elementType, bufferSize, memorySpace, layout, config); + if (mlirTypeIsNull(ty)) { + return py::none(); + } + return cls(ty); + }, + py::arg("cls"), + py::arg("shape"), + py::arg("element_type"), + py::arg("buffer_size"), + py::arg("memory_space"), + py::arg("layout"), + py::arg("config") = py::none(), + py::arg("context") = py::none()) + .def_classmethod( + "upcast_type", + [](py::object cls, MlirType t) -> py::object { + if (mlirPTOTypeIsAConvTileType(t)) { + return cls(t); + } + return py::none(); + }, + py::arg("cls"), py::arg("type")) + .def_property_readonly( + "rank", + [](MlirType self) -> intptr_t { + return mlirPTOConvTileTypeGetRank(self); + }) + .def_property_readonly( + "element_type", + [](MlirType self) -> MlirType { + return mlirPTOConvTileTypeGetElementType(self); + }) + .def_property_readonly( + "shape", + [](MlirType self) -> py::list { + intptr_t n = 0; + const int64_t *d = mlirPTOConvTileTypeGetShape(self, &n); + return shapeToPyList(d, n); + }); + populatePTODialectSubmodule(m); } diff --git a/lib/CAPI/Dialect/PTO.cpp b/lib/CAPI/Dialect/PTO.cpp index 4966e2f3c1..4e5e8fce25 100644 --- a/lib/CAPI/Dialect/PTO.cpp +++ b/lib/CAPI/Dialect/PTO.cpp @@ -352,6 +352,57 @@ MlirType mlirPTOTileBufTypeGetWithValidShapeAndConfig(MlirContext ctx, return wrap(ty); } +bool mlirPTOTypeIsAConvTileType(MlirType type) { + return mlir::isa(unwrap(type)); +} + +MlirType mlirPTOConvTileTypeGet(MlirContext ctx, intptr_t rank, + const int64_t *shape, MlirType elementType, + MlirAttribute bufferSize, + MlirAttribute memorySpace, + MlirAttribute layout, + MlirAttribute config) { + MLIRContext *c = unwrap(ctx); + auto shp = llvm::ArrayRef(shape, rank); + auto bufferSizeAttr = mlir::dyn_cast(unwrap(bufferSize)); + auto memorySpaceAttr = + mlir::dyn_cast(unwrap(memorySpace)); + auto layoutAttr = mlir::dyn_cast(unwrap(layout)); + mlir::pto::ConvTileConfigAttr configAttr; + if (!mlirAttributeIsNull(config)) { + configAttr = + mlir::dyn_cast_or_null(unwrap(config)); + } + if (!configAttr) { + configAttr = mlir::pto::ConvTileConfigAttr::getDefault(c); + } + if (!bufferSizeAttr || !memorySpaceAttr || !layoutAttr) { + return MlirType{nullptr}; + } + + auto ty = mlir::pto::ConvTileType::get(c, shp, unwrap(elementType), + bufferSizeAttr, memorySpaceAttr, + layoutAttr, configAttr); + return wrap(ty); +} + +intptr_t mlirPTOConvTileTypeGetRank(MlirType type) { + return mlir::cast(unwrap(type)).getRank(); +} + +MlirType mlirPTOConvTileTypeGetElementType(MlirType type) { + return wrap(mlir::cast(unwrap(type)).getElementType()); +} + +const int64_t *mlirPTOConvTileTypeGetShape(MlirType type, + intptr_t *numDimsOut) { + static thread_local llvm::SmallVector shape; + auto dims = mlir::cast(unwrap(type)).getShape(); + shape.assign(dims.begin(), dims.end()); + *numDimsOut = static_cast(shape.size()); + return shape.data(); +} + bool mlirPTOAttrIsABLayoutAttr(MlirAttribute attr) { return mlir::isa(unwrap(attr)); } @@ -436,6 +487,7 @@ DEFINE_PTO_ENUM_ATTR_CAPI(RemPrecision, RemPrecisionAttr, RemPrecision) DEFINE_PTO_ENUM_ATTR_CAPI(RsqrtPrecision, RsqrtPrecisionAttr, RsqrtPrecision) DEFINE_PTO_ENUM_ATTR_CAPI(SqrtPrecision, SqrtPrecisionAttr, SqrtPrecision) DEFINE_PTO_ENUM_ATTR_CAPI(FmodPrecision, FmodPrecisionAttr, FmodPrecision) +DEFINE_PTO_ENUM_ATTR_CAPI(FmatrixMode, FmatrixModeAttr, FmatrixMode) #undef DEFINE_PTO_ENUM_ATTR_CAPI @@ -760,6 +812,14 @@ static mlir::pto::CompactModeAttr toCompactModeAttr(mlir::MLIRContext *c, return {}; } +static mlir::IntegerAttr toI32IntegerAttr(mlir::Attribute a) { + auto intAttr = mlir::dyn_cast(a); + if (!intAttr || !intAttr.getType().isSignlessInteger(kI32BitWidth)) { + return {}; + } + return intAttr; +} + bool mlirPTOAttrIsACompactModeAttr(MlirAttribute attr) { return mlir::isa(unwrap(attr)); } @@ -912,6 +972,56 @@ MlirAttribute mlirPTOTileBufConfigAttrGetWithCompactMode( return wrap(mlir::pto::TileBufConfigAttr::get(c, blA, slA, sz, pvA, cmA)); } +bool mlirPTOAttrIsAConvTileConfigAttr(MlirAttribute attr) { + return mlir::isa(unwrap(attr)); +} + +MlirAttribute mlirPTOConvTileConfigAttrGetDefault(MlirContext ctx) { + auto *c = unwrap(ctx); + return wrap(mlir::pto::ConvTileConfigAttr::getDefault(c)); +} + +MlirAttribute mlirPTOConvTileConfigAttrGet( + MlirContext ctx, MlirAttribute fmapH, MlirAttribute fmapW, + intptr_t padListSize, const int64_t *padList, MlirAttribute filterH, + MlirAttribute filterW, MlirAttribute dilationH, MlirAttribute dilationW, + MlirAttribute strideH, MlirAttribute strideW, MlirAttribute padValue, + MlirAttribute channelSize, MlirAttribute repeatStride, + MlirAttribute repeatTime, MlirAttribute repeatMode, MlirAttribute dstStride, + MlirAttribute dstMposition, MlirAttribute transpose) { + auto *c = unwrap(ctx); + auto fmapHAttr = toI32IntegerAttr(unwrap(fmapH)); + auto fmapWAttr = toI32IntegerAttr(unwrap(fmapW)); + auto filterHAttr = toI32IntegerAttr(unwrap(filterH)); + auto filterWAttr = toI32IntegerAttr(unwrap(filterW)); + auto dilationHAttr = toI32IntegerAttr(unwrap(dilationH)); + auto dilationWAttr = toI32IntegerAttr(unwrap(dilationW)); + auto strideHAttr = toI32IntegerAttr(unwrap(strideH)); + auto strideWAttr = toI32IntegerAttr(unwrap(strideW)); + auto channelSizeAttr = toI32IntegerAttr(unwrap(channelSize)); + auto repeatStrideAttr = toI32IntegerAttr(unwrap(repeatStride)); + auto repeatTimeAttr = toI32IntegerAttr(unwrap(repeatTime)); + auto repeatModeAttr = toI32IntegerAttr(unwrap(repeatMode)); + auto dstStrideAttr = toI32IntegerAttr(unwrap(dstStride)); + auto dstMpositionAttr = toI32IntegerAttr(unwrap(dstMposition)); + auto transposeAttr = mlir::dyn_cast(unwrap(transpose)); + if (!fmapHAttr || !fmapWAttr || !filterHAttr || !filterWAttr || + !dilationHAttr || !dilationWAttr || !strideHAttr || !strideWAttr || + !channelSizeAttr || !repeatStrideAttr || !repeatTimeAttr || + !repeatModeAttr || !dstStrideAttr || !dstMpositionAttr || + !transposeAttr) { + return MlirAttribute{nullptr}; + } + + auto pads = + llvm::ArrayRef(padList, static_cast(padListSize)); + return wrap(mlir::pto::ConvTileConfigAttr::get( + c, fmapHAttr, fmapWAttr, pads, filterHAttr, filterWAttr, dilationHAttr, + dilationWAttr, strideHAttr, strideWAttr, unwrap(padValue), + channelSizeAttr, repeatStrideAttr, repeatTimeAttr, repeatModeAttr, + dstStrideAttr, dstMpositionAttr, transposeAttr)); +} + MlirType mlirPTOGMTypeGet(MlirContext ctx, intptr_t rank, const int64_t *shape, MlirType elementType) { auto *c = unwrap(ctx); diff --git a/python/pto/dialects/pto.py b/python/pto/dialects/pto.py index c3fdac1700..86298b86bc 100644 --- a/python/pto/dialects/pto.py +++ b/python/pto/dialects/pto.py @@ -61,11 +61,13 @@ def _export_optional_cext_symbol(name): PartitionTensorViewType = _pto_mod.PartitionTensorViewType TileType = _pto_mod.TileType TileBufType = _pto_mod.TileBufType +ConvTileType = _pto_mod.ConvTileType AddressSpace = _pto_mod.AddressSpace AddressSpaceAttr = _pto_mod.AddressSpaceAttr FenceScope = _pto_mod.FenceScope FenceScopeAttr = _pto_mod.FenceScopeAttr TileBufConfigAttr = _pto_mod.TileBufConfigAttr +ConvTileConfigAttr = _pto_mod.ConvTileConfigAttr BLayout = _pto_mod.BLayout BLayoutAttr = _pto_mod.BLayoutAttr SLayout = _pto_mod.SLayout @@ -108,6 +110,8 @@ def _export_optional_cext_symbol(name): FmodPrecisionAttr = _pto_mod.FmodPrecisionAttr SaturationMode = _pto_mod.SaturationMode SaturationModeAttr = _pto_mod.SaturationModeAttr +FmatrixMode = _pto_mod.FmatrixMode +FmatrixModeAttr = _pto_mod.FmatrixModeAttr CmpMode = _pto_mod.CmpMode CmpModeAttr = _pto_mod.CmpModeAttr PIPE = _pto_mod.PIPE @@ -290,6 +294,8 @@ def fence_scope_attr_builder(value, context=None): "FmodPrecisionAttr", "SaturationMode", "SaturationModeAttr", + "FmatrixMode", + "FmatrixModeAttr", "CmpMode", "CmpModeAttr", "PIPE", @@ -313,6 +319,8 @@ def fence_scope_attr_builder(value, context=None): "VecStoreMode", "VecStoreModeAttr", "TileBufConfigAttr", + "ConvTileConfigAttr", + "ConvTileType", "TileConfig", # High-level sync helpers "record_event", From 70cdb85cf598a63320e68f34239fa9721240dcac Mon Sep 17 00:00:00 2001 From: andodo Date: Fri, 4 Sep 2026 11:20:50 +0800 Subject: [PATCH 11/12] Support ConvTile memory planning --- lib/PTO/IR/PTO.cpp | 15 ++++++++++++-- lib/PTO/Transforms/PTOPlanMemory.cpp | 23 ++++++++++++++++++++-- lib/PTO/Transforms/PTOPlanMemoryModern.cpp | 4 ++-- 3 files changed, 36 insertions(+), 6 deletions(-) diff --git a/lib/PTO/IR/PTO.cpp b/lib/PTO/IR/PTO.cpp index 6e15ea4817..656fb3d6fd 100644 --- a/lib/PTO/IR/PTO.cpp +++ b/lib/PTO/IR/PTO.cpp @@ -3592,9 +3592,20 @@ static LogicalResult verifyConstantLocalAddress(Operation *op, Value addr, } LogicalResult AllocTileOp::verify() { - auto ty = dyn_cast(getResult().getType()); + auto resultType = getResult().getType(); + if (auto convTy = dyn_cast(resultType)) { + bool hasValidShapeOperands = getValidRow() || getValidCol(); + if (hasValidShapeOperands) { + return emitOpError("valid_row and valid_col operands require result to be `!pto.tile_buf`"); + } + + return verifyConstantLocalAddress(getOperation(), getAddr(), + convTy.getMemorySpace()); + } + + auto ty = dyn_cast(resultType); if (!ty) { - return emitOpError("result must be `!pto.tile_buf`"); + return emitOpError("result must be `!pto.tile_buf` or `!pto.conv_tile`"); } if (failed(verifyTileBufLayoutConstraints(*this, ty, "result"))) { diff --git a/lib/PTO/Transforms/PTOPlanMemory.cpp b/lib/PTO/Transforms/PTOPlanMemory.cpp index 84881e422d..d70963a1e3 100644 --- a/lib/PTO/Transforms/PTOPlanMemory.cpp +++ b/lib/PTO/Transforms/PTOPlanMemory.cpp @@ -94,6 +94,22 @@ static std::optional getTileBufferFootprintBytes(TileBufType type) { static_cast(elemBytes); } +static std::optional getConvTileFootprintBytes(ConvTileType type) { + unsigned elemBytes = getPTOStorageElemByteSize(type.getElementType()); + int64_t bufferSize = type.getBufferSizeValue(); + bool invalidCapacity = elemBytes == 0 || bufferSize <= 0; + if (invalidCapacity) { + return std::nullopt; + } + + int64_t elemBytesI64 = static_cast(elemBytes); + bool overflows = bufferSize > std::numeric_limits::max() / elemBytesI64; + if (overflows) { + return std::nullopt; + } + return bufferSize * elemBytesI64; +} + static int64_t ceilDivBitsToBytes(int64_t bits) { return (bits + kBitsPerByte - 1) / kBitsPerByte; } @@ -1115,6 +1131,9 @@ BufferInfo MemLivenessAnalysis::GetBufferInfo(Operation *op, Value operand, if (auto tileType = dyn_cast(operand.getType())) { elementType = tileType.getElementType(); footprintBytes = getTileBufferFootprintBytes(tileType); + } else if (auto convType = dyn_cast(operand.getType())) { + elementType = convType.getElementType(); + footprintBytes = getConvTileFootprintBytes(convType); } else if (auto multiType = dyn_cast(operand.getType())) { TileBufType slotType = multiType.getSlotType(); elementType = slotType.getElementType(); @@ -2663,8 +2682,8 @@ class LegacyAllocTileOpAddPlannedAddressPattern return failure(); } - auto tileType = dyn_cast(op.getResult().getType()); - if (!tileType) { + Type tileType = op.getResult().getType(); + if (!isa(tileType)) { return failure(); } diff --git a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp index 2a1ce2e68b..e1793eb87b 100644 --- a/lib/PTO/Transforms/PTOPlanMemoryModern.cpp +++ b/lib/PTO/Transforms/PTOPlanMemoryModern.cpp @@ -1591,8 +1591,8 @@ class AllocTileOpAddPlannedAddressPattern return failure(); } - auto tileType = dyn_cast(op.getResult().getType()); - if (!tileType) { + Type tileType = op.getResult().getType(); + if (!isa(tileType)) { return failure(); } From 5d835ffacf5c7c89c6d7c35ff513e5019b657436 Mon Sep 17 00:00:00 2001 From: andodo Date: Fri, 4 Sep 2026 11:32:56 +0800 Subject: [PATCH 12/12] Add ConvTile Img2col tests --- test/lit/pto/conv_tile_config_parse.pto | 37 ++++++++++ test/lit/pto/fmatrix_mode_emitc.pto | 62 ++++++++++++++++ .../pto/fmatrix_mode_invalid_dst_layout.pto | 21 ++++++ test/lit/pto/fmatrix_mode_invalid_set.pto | 19 +++++ .../pto/fmatrix_mode_invalid_src_layout.pto | 21 ++++++ .../pto/timg2col_alloc_conv_tile_emitc.pto | 74 +++++++++++++++++++ test/lit/pto/timg2col_invalid_pos.pto | 21 ++++++ test/lit/pto/timg2col_memory_effects.pto | 23 ++++++ test/lit/pto/timg2col_pos_offsets.pto | 25 +++++++ test/samples/Img2col/timg2col-pto.pto | 8 +- test/samples/Img2col/timg2col_runtime.py | 41 +++++----- 11 files changed, 325 insertions(+), 27 deletions(-) create mode 100644 test/lit/pto/conv_tile_config_parse.pto create mode 100644 test/lit/pto/fmatrix_mode_emitc.pto create mode 100644 test/lit/pto/fmatrix_mode_invalid_dst_layout.pto create mode 100644 test/lit/pto/fmatrix_mode_invalid_set.pto create mode 100644 test/lit/pto/fmatrix_mode_invalid_src_layout.pto create mode 100644 test/lit/pto/timg2col_alloc_conv_tile_emitc.pto create mode 100644 test/lit/pto/timg2col_invalid_pos.pto create mode 100644 test/lit/pto/timg2col_memory_effects.pto create mode 100644 test/lit/pto/timg2col_pos_offsets.pto diff --git a/test/lit/pto/conv_tile_config_parse.pto b/test/lit/pto/conv_tile_config_parse.pto new file mode 100644 index 0000000000..9454df8ea3 --- /dev/null +++ b/test/lit/pto/conv_tile_config_parse.pto @@ -0,0 +1,37 @@ +// 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. + +// RUN: ptoas --pto-arch=a5 --emit-pto-ir %s -o - | FileCheck %s + +module { + func.func @conv_tile_config_parse( + %src: !pto.conv_tile< + mat, 1x1x3x4x8xf16, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x2x3x4, + filter_h=1, filter_w=1, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=16, repeat_time=2, + repeat_mode=0, dst_stride=32, dst_mposition=0, + transpose=false>>) { + pto.set_fmatrix %src : !pto.conv_tile< + mat, 1x1x3x4x8xf16, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x2x3x4, + filter_h=1, filter_w=1, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=16, repeat_time=2, + repeat_mode=0, dst_stride=32, dst_mposition=0, + transpose=false>> + return + } +} + +// CHECK-LABEL: func.func @conv_tile_config_parse +// CHECK: !pto.conv_tile, + %dst: !pto.tile_buf) { + pto.set_fmatrix %cfg : !pto.conv_tile + pto.set_img2col_rpt %cfg : !pto.conv_tile + pto.set_img2col_padding %cfg : !pto.conv_tile + pto.timg2col %dst, %cfg + : !pto.tile_buf, + !pto.conv_tile + return + } + + func.func @fmatrix_mode_b( + %cfg: !pto.conv_tile, + %dst: !pto.tile_buf) { + pto.set_fmatrix %cfg {fmatrixMode = #pto.fmatrix_mode} + : !pto.conv_tile + pto.set_img2col_rpt %cfg {fmatrixMode = #pto.fmatrix_mode} + : !pto.conv_tile + pto.set_img2col_padding %cfg {fmatrixMode = #pto.fmatrix_mode} + : !pto.conv_tile + pto.timg2col %dst, %cfg {fmatrixMode = #pto.fmatrix_mode} + : !pto.tile_buf, + !pto.conv_tile + return + } + + func.func @fmatrix_mode_a_auto( + %cfg: !pto.conv_tile, + %dst: !pto.tile_buf) { + pto.timg2col %dst, %cfg {fmatrixMode = #pto.fmatrix_mode} + : !pto.tile_buf, + !pto.conv_tile + return + } +} + +// CHECK-LABEL: fmatrix_mode_default +// CHECK: SETFMATRIX( +// CHECK: SET_IMG2COL_RPT( +// CHECK: SET_IMG2COL_PADDING( +// CHECK: TIMG2COL( + +// CHECK-LABEL: fmatrix_mode_b +// CHECK: SETFMATRIX +// CHECK: SET_IMG2COL_RPT +// CHECK: SET_IMG2COL_PADDING +// CHECK: TIMG2COL + +// CHECK-LABEL: fmatrix_mode_a_auto +// CHECK: TIMG2COL diff --git a/test/lit/pto/fmatrix_mode_invalid_dst_layout.pto b/test/lit/pto/fmatrix_mode_invalid_dst_layout.pto new file mode 100644 index 0000000000..2db5a2bdd1 --- /dev/null +++ b/test/lit/pto/fmatrix_mode_invalid_dst_layout.pto @@ -0,0 +1,21 @@ +// 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. + +// RUN: not ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s + +module { + func.func @invalid_timg2col_dst_layout( + %src: !pto.conv_tile, + %dst: !pto.tile_buf) { + // CHECK: expects dst layout to be BLayout=col_major and SLayout=row_major + pto.timg2col %dst, %src + : !pto.tile_buf, + !pto.conv_tile + return + } +} diff --git a/test/lit/pto/fmatrix_mode_invalid_set.pto b/test/lit/pto/fmatrix_mode_invalid_set.pto new file mode 100644 index 0000000000..2a872520f4 --- /dev/null +++ b/test/lit/pto/fmatrix_mode_invalid_set.pto @@ -0,0 +1,19 @@ +// 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. + +// RUN: not ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s + +module { + func.func @invalid_set_mode( + %cfg: !pto.conv_tile) { + // CHECK: expects fmatrix_mode to be a_manual or b_manual + pto.set_fmatrix %cfg {fmatrixMode = #pto.fmatrix_mode} + : !pto.conv_tile + return + } +} diff --git a/test/lit/pto/fmatrix_mode_invalid_src_layout.pto b/test/lit/pto/fmatrix_mode_invalid_src_layout.pto new file mode 100644 index 0000000000..ae11df3211 --- /dev/null +++ b/test/lit/pto/fmatrix_mode_invalid_src_layout.pto @@ -0,0 +1,21 @@ +// 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. + +// RUN: not ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s + +module { + func.func @invalid_timg2col_src_layout( + %src: !pto.conv_tile>, + %dst: !pto.tile_buf) { + // CHECK: expects src layout to be NC1HWC0 or NDC1HWC0 + pto.timg2col %dst, %src + : !pto.tile_buf, + !pto.conv_tile> + return + } +} diff --git a/test/lit/pto/timg2col_alloc_conv_tile_emitc.pto b/test/lit/pto/timg2col_alloc_conv_tile_emitc.pto new file mode 100644 index 0000000000..aa882828b9 --- /dev/null +++ b/test/lit/pto/timg2col_alloc_conv_tile_emitc.pto @@ -0,0 +1,74 @@ +// 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 OR CONDITIONS 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. + +// RUN: ptoas --pto-arch=a5 %s -o - | FileCheck %s + +module { + func.func @timg2col_alloc_conv_tile_emitc( + %dst: !pto.tile_buf) { + %src = pto.alloc_tile : !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>> + pto.set_fmatrix %src {fmatrixMode = #pto.fmatrix_mode} + : !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>> + pto.set_img2col_rpt %src {fmatrixMode = #pto.fmatrix_mode} + : !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>> + pto.set_img2col_padding %src {fmatrixMode = #pto.fmatrix_mode} + : !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>> + pto.timg2col %dst, %src {posM = 1 : i32, posK = 8 : i32, + fmatrixMode = #pto.fmatrix_mode} + : !pto.tile_buf, + !pto.conv_tile< + mat, 1x1x3x4x8xf32, + config=#pto.conv_tile_config< + fmap_h=3, fmap_w=4, pad_list=1x0x1x0, + filter_h=2, filter_w=2, dilation_h=1, dilation_w=1, + stride_h=1, stride_w=1, pad_value=0 : i32, + channel_size=8, repeat_stride=0, repeat_time=1, + repeat_mode=0, dst_stride=1, dst_mposition=0, + transpose=false>> + return + } +} + +// CHECK-LABEL: timg2col_alloc_conv_tile_emitc( +// CHECK: PTOAS__INIT_CONVTILE_CONFIG( +// CHECK: SETFMATRIX +// CHECK: SET_IMG2COL_RPT +// CHECK: SET_IMG2COL_PADDING +// CHECK: TIMG2COL diff --git a/test/lit/pto/timg2col_invalid_pos.pto b/test/lit/pto/timg2col_invalid_pos.pto new file mode 100644 index 0000000000..53107d1e88 --- /dev/null +++ b/test/lit/pto/timg2col_invalid_pos.pto @@ -0,0 +1,21 @@ +// 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 OR CONDITIONS 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. + +// RUN: not ptoas --pto-arch=a5 %s -o - 2>&1 | FileCheck %s + +module { + func.func @timg2col_invalid_pos( + %src: !pto.conv_tile, + %dst: !pto.tile_buf) { + // CHECK: expects posM to fit in an unsigned 16-bit integer + pto.timg2col %dst, %src {posM = 65536 : i32} + : !pto.tile_buf, + !pto.conv_tile + return + } +} diff --git a/test/lit/pto/timg2col_memory_effects.pto b/test/lit/pto/timg2col_memory_effects.pto new file mode 100644 index 0000000000..6341874a66 --- /dev/null +++ b/test/lit/pto/timg2col_memory_effects.pto @@ -0,0 +1,23 @@ +// 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 OR CONDITIONS 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. + +// RUN: ptoas --pto-arch=a5 --enable-insert-sync --emit-pto-ir %s -o - | FileCheck %s + +module attributes {pto.target_arch = "a5"} { + func.func @timg2col_memory_effects( + %src: !pto.conv_tile, + %dst: !pto.tile_buf) { + pto.timg2col %dst, %src + : !pto.tile_buf, + !pto.conv_tile + return + } +} + +// CHECK-LABEL: func.func @timg2col_memory_effects +// CHECK: pto.timg2col diff --git a/test/lit/pto/timg2col_pos_offsets.pto b/test/lit/pto/timg2col_pos_offsets.pto new file mode 100644 index 0000000000..103cd9c8cd --- /dev/null +++ b/test/lit/pto/timg2col_pos_offsets.pto @@ -0,0 +1,25 @@ +// 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. + +// RUN: ptoas --pto-arch=a5 %s -o - | FileCheck %s + +module { + func.func @timg2col_pos_offsets( + %src: !pto.conv_tile, + %dst: !pto.tile_buf) { + pto.timg2col %dst, %src {posM = 1 : i32, posK = 8 : i32} + : !pto.tile_buf, + !pto.conv_tile + return + } +} + +// CHECK-LABEL: timg2col_pos_offsets( +// CHECK: const uint16_t {{.*}} = 1; +// CHECK: const uint16_t {{.*}} = 8; +// CHECK: TIMG2COL( diff --git a/test/samples/Img2col/timg2col-pto.pto b/test/samples/Img2col/timg2col-pto.pto index 4e72eee957..f4b192030a 100644 --- a/test/samples/Img2col/timg2col-pto.pto +++ b/test/samples/Img2col/timg2col-pto.pto @@ -18,7 +18,7 @@ module attributes {pto.target_arch = "a5"} { repeat_mode=0, dst_stride=1, dst_mposition=0, transpose=false>>, %dst: !pto.tile_buf) { - pto.set_fmatrix %src {fmatrix_mode = #pto.fmatrix_mode} + pto.set_fmatrix %src {fmatrixMode = #pto.fmatrix_mode} : !pto.conv_tile< mat, 1x1x3x4x8xf32, config=#pto.conv_tile_config< @@ -28,7 +28,7 @@ module attributes {pto.target_arch = "a5"} { channel_size=8, repeat_stride=0, repeat_time=1, repeat_mode=0, dst_stride=1, dst_mposition=0, transpose=false>> - pto.set_img2col_rpt %src {fmatrix_mode = #pto.fmatrix_mode} + pto.set_img2col_rpt %src {fmatrixMode = #pto.fmatrix_mode} : !pto.conv_tile< mat, 1x1x3x4x8xf32, config=#pto.conv_tile_config< @@ -38,7 +38,7 @@ module attributes {pto.target_arch = "a5"} { channel_size=8, repeat_stride=0, repeat_time=1, repeat_mode=0, dst_stride=1, dst_mposition=0, transpose=false>> - pto.set_img2col_padding %src {fmatrix_mode = #pto.fmatrix_mode} + pto.set_img2col_padding %src {fmatrixMode = #pto.fmatrix_mode} : !pto.conv_tile< mat, 1x1x3x4x8xf32, config=#pto.conv_tile_config< @@ -49,7 +49,7 @@ module attributes {pto.target_arch = "a5"} { repeat_mode=0, dst_stride=1, dst_mposition=0, transpose=false>> pto.timg2col %dst, %src {posM = 1 : i32, posK = 8 : i32, - fmatrix_mode = #pto.fmatrix_mode} + fmatrixMode = #pto.fmatrix_mode} : !pto.tile_buf, !pto.conv_tile< mat, 1x1x3x4x8xf32, diff --git a/test/samples/Img2col/timg2col_runtime.py b/test/samples/Img2col/timg2col_runtime.py index adc70d8a24..f5fef7ae4a 100644 --- a/test/samples/Img2col/timg2col_runtime.py +++ b/test/samples/Img2col/timg2col_runtime.py @@ -40,8 +40,6 @@ def build(): ptr_f32 = pto.PtrType.get(f32, ctx) tv5_f32 = pto.TensorViewType.get(5, f32, ctx) src_ptv = pto.PartitionTensorViewType.get([1, 1, 3, 4, 8], f32, ctx) - dst_tv = pto.TensorViewType.get(2, f32, ctx) - dst_ptv = pto.PartitionTensorViewType.get([16, 32], f32, ctx) vec = pto.AddressSpaceAttr.get(pto.AddressSpace.VEC, ctx) left = pto.AddressSpaceAttr.get(pto.AddressSpace.LEFT, ctx) @@ -88,7 +86,7 @@ def build(): ctx, ) - fn_ty = func.FunctionType.get([ptr_f32, ptr_f32], []) + fn_ty = func.FunctionType.get([ptr_f32], []) with InsertionPoint(module.body): fn = func.FuncOp("timg2col_runtime_kernel", fn_ty) fn.operation.attributes["pto.entry"] = UnitAttr.get(ctx) @@ -100,17 +98,13 @@ def build(): c3 = arith.ConstantOp(IndexType.get(ctx), 3).result c4 = arith.ConstantOp(IndexType.get(ctx), 4).result c8 = arith.ConstantOp(IndexType.get(ctx), 8).result - c16 = arith.ConstantOp(IndexType.get(ctx), 16).result c32 = arith.ConstantOp(IndexType.get(ctx), 32).result c96 = arith.ConstantOp(IndexType.get(ctx), 96).result - src_ptr, dst_ptr = entry.arguments + src_ptr = entry.arguments[0] src_view = pto.MakeTensorViewOp( tv5_f32, src_ptr, [c1, c1, c3, c4, c8], [c96, c96, c32, c8, c1] ).result - dst_view = pto.MakeTensorViewOp( - dst_tv, dst_ptr, [c16, c32], [c32, c1] - ).result src_part = pto.PartitionViewOp( src_ptv, @@ -118,24 +112,25 @@ def build(): offsets=[c0, c0, c0, c0, c0], sizes=[c1, c1, c3, c4, c8], ).result - dst_part = pto.PartitionViewOp( - dst_ptv, dst_view, offsets=[c0, c0], sizes=[c16, c32] - ).result src_tile = pto.AllocTileOp(src_tile_ty).result dst_tile = pto.AllocTileOp(dst_tile_ty).result - pto.TLoadOp(None, src_part, src_tile) - pto.SetFmatrixOp(src_tile, pto.FmatrixModeAttr.get(pto.FmatrixMode.FMATRIX_B_MANUAL, ctx)) - pto.SetImg2colRptOp(src_tile, pto.FmatrixModeAttr.get(pto.FmatrixMode.FMATRIX_B_MANUAL, ctx)) - pto.SetImg2colPaddingOp(src_tile, pto.FmatrixModeAttr.get(pto.FmatrixMode.FMATRIX_B_MANUAL, ctx)) - pto.TImg2colOp( - dst_tile, - src_tile, - posM=1, - posK=8, - fmatrixMode=pto.FmatrixModeAttr.get(pto.FmatrixMode.FMATRIX_B_MANUAL, ctx), - ) - pto.TStoreOp(None, dst_tile, dst_part) + cube_section = pto.SectionCubeOp() + with InsertionPoint(cube_section.body.blocks.append()): + fmatrix_mode = pto.FmatrixModeAttr.get( + pto.FmatrixMode.FMATRIX_B_MANUAL, ctx + ) + pto.TLoadOp(None, src_part, src_tile) + pto.SetFmatrixOp(src_tile, fmatrixMode=fmatrix_mode) + pto.SetImg2colRptOp(src_tile, fmatrixMode=fmatrix_mode) + pto.SetImg2colPaddingOp(src_tile, fmatrixMode=fmatrix_mode) + pto.TImg2colOp( + dst_tile, + src_tile, + posM=1, + posK=8, + fmatrixMode=fmatrix_mode, + ) func.ReturnOp([]) module.operation.verify()