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 16ac3410af..e05fa405ce 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 48f16a3b28..11e75802da 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 f4532e7912..b466cb8f82 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); }); @@ -12492,7 +12551,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 53bb9a7097..097bbfe115 100644 --- a/lib/PTO/Transforms/Utils.cpp +++ b/lib/PTO/Transforms/Utils.cpp @@ -279,6 +279,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; @@ -431,6 +438,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; if (!value || !matchPattern(value, m_Constant(&attr)) || attr.getInt() < 0) @@ -484,6 +506,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) { if (op.getOffsets().size() != 2) @@ -547,9 +577,19 @@ static std::optional resolveSemanticRange(Value value) { return std::nullopt; if (auto alloc = value.getDefiningOp()) { auto tileType = dyn_cast(alloc.getResult().getType()); - auto bytes = tileType ? getStaticTileBytes(tileType) : std::nullopt; - if (!tileType || !bytes) + auto convType = dyn_cast(alloc.getResult().getType()); + auto bytes = tileType ? getStaticTileBytes(tileType) + : (convType ? getStaticConvTileBytes(convType) + : std::nullopt); + const bool invalidRange = (!tileType && !convType) || !bytes; + if (invalidRange) { return std::nullopt; + } + if (convType) { + return SemanticRange{ + alloc.getResult(), 0, *bytes, getConstantAddress(alloc.getAddr()), + getConvTileAddressSpace(convType), std::nullopt, std::nullopt, 0}; + } auto strides = getStaticTileStrides(tileType); return SemanticRange{ alloc.getResult(), 0, *bytes, getConstantAddress(alloc.getAddr()),