Adapt conv tile.

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<PTO_ConvLayoutEnum>:$value);
+  let assemblyFormat = "`<` params `>`";
+  let description = [{
+    Physical layout carried by a PTO ConvTile.
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Function and Module Core Type
 //===----------------------------------------------------------------------===//

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<Index>:$valid_col
   );
 
-  let results = (outs TileBufType:$result);
+  let results = (outs AnyTypeOf<[TileBufType, ConvTileType]>:$result);
 
   let assemblyFormat = [{
     (`addr` `=` $addr^)?

include/PTO/IR/PTOTypeDefs.td

@@ -228,6 +228,36 @@ def TileBufType : TypeDef<PTO_Dialect, "TileBuf"> {
   }];
 }
 
+def ConvTileType : TypeDef<PTO_Dialect, "ConvTile"> {
+  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
 // =============================================================================

lib/PTO/IR/PTO.cpp

@@ -293,6 +293,10 @@ static int64_t getPTOTypeRank(Type type) {
     return tileBufTy.getRank();
   }
 
+  if (auto convTileTy = dyn_cast<pto::ConvTileType>(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<ConvTileType>(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<TileBufType>(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<ConvTileType>(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<TileBufType>(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<ConvTileType>(getDst().getType())) {
+    auto srcPart = dyn_cast<PartitionTensorViewType>(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<std::pair<pto::PartitionTensorViewType, pto::TileBufType>> {
@@ -4871,6 +4959,21 @@ static LogicalResult verifyCommPingPongSameType(Operation *op, Value ping,
 }
 
 static std::optional<uint64_t> getStaticByteSize(Type ty) {
+  if (auto conv = dyn_cast<pto::ConvTileType>(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<uint64_t>(conv.getBufferSize());
+    const bool overflows =
+        bufferSize > std::numeric_limits<uint64_t>::max() / elemBytes;
+    if (overflows) {
+      return std::nullopt;
+    }
+    return bufferSize * elemBytes;
+  }
+
   SmallVector<int64_t, 4> shape = getShapeVec(ty);
   if (shape.empty()) {
     return std::nullopt;
@@ -4920,6 +5023,13 @@ static std::optional<pto::AddressSpace> getPTOMemorySpaceEnum(Type ty) {
     }
     return std::nullopt;
   }
+  if (auto conv = dyn_cast<pto::ConvTileType>(ty)) {
+    if (auto as =
+            dyn_cast_or_null<pto::AddressSpaceAttr>(conv.getMemorySpace())) {
+      return as.getAddressSpace();
+    }
+    return std::nullopt;
+  }
   return std::nullopt;
 }