Adapt conv tile.

include/PTO/IR/PTOAttrs.td

@@ -94,6 +94,28 @@ 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 : EnumAttr<PTO_Dialect, PTO_FmatrixModeEnum, "fmatrix_mode"> {
+  let assemblyFormat = "`<` params `>`";
+  let summary = "A/B selector for FMATRIX-related ops";
+}
+
 //===----------------------------------------------------------------------===//
 // SIMT GM load/store L1 and L2 cache controls
 //===----------------------------------------------------------------------===//
@@ -484,23 +506,84 @@ 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)", [
     PTO_Layout_ND,
     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<PTO_LayoutEnum>:$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.
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// 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.
   }];
 }
 
@@ -1237,6 +1320,56 @@ def TileBufConfigAttr : AttrDef<PTO_Dialect, "TileBufConfig"> {
   }];
 }
 
+// ---------- conv_tile_config ----------
+def ConvTileConfigAttr : AttrDef<PTO_Dialect, "ConvTileConfig"> {
+  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 extraClassDeclaration = [{
+    static ConvTileConfigAttr getDefault(MLIRContext *ctx);
+    bool isDefault() const;
+
+    static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
+                                mlir::IntegerAttr fmapH,
+                                mlir::IntegerAttr fmapW,
+                                ArrayRef<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);
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // QuantType
 //===----------------------------------------------------------------------===//

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,10 @@ def AllocTileOp : PTO_Op<"alloc_tile", [AttrSizedOperandSegments]> {
     Optional<Index>:$valid_col
   );
 
-  let results = (outs TileBufType:$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^)?
@@ -1360,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;
       };
 
@@ -1882,6 +1893,67 @@ 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,
+    DefaultValuedAttr<PTO_FmatrixModeAttr, "::mlir::pto::FmatrixMode::FMATRIX_A_MANUAL">:$fmatrixMode
+  );
+  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,
+    DefaultValuedAttr<PTO_FmatrixModeAttr, "::mlir::pto::FmatrixMode::FMATRIX_A_MANUAL">:$fmatrixMode
+  );
+  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,
+    DefaultValuedAttr<PTO_FmatrixModeAttr, "::mlir::pto::FmatrixMode::FMATRIX_A_MANUAL">:$fmatrixMode
+  );
+  let results = (outs);
+  let hasVerifier = 1;
+  let assemblyFormat = [{
+    $src attr-dict `:` qualified(type($src))
+  }];
+}
+
+def TImg2colOp : PTO_Op<"timg2col", [
+    DeclareOpInterfaceMethods<MemoryEffectsOpInterface>
+]> {
+  let summary = "Image-to-column transform from ConvTile to TileBuf";
+  let arguments = (ins
+    TileBufType:$dst,
+    ConvTileType:$src,
+    DefaultValuedOptionalAttr<I32Attr, "0">:$posM,
+    DefaultValuedOptionalAttr<I32Attr, "0">:$posK,
+    DefaultValuedAttr<PTO_FmatrixModeAttr, "::mlir::pto::FmatrixMode::FMATRIX_A_MANUAL">:$fmatrixMode
+  );
+  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";
 
@@ -6239,7 +6311,6 @@ def TDeInterleaveOp: PTO_TOp<"tdeinterleave", [
     ::mlir::Value getDst1() { return getDsts()[1]; }
   }];
 }
-
 def TRowProdOp: PTO_TOp<"trowprod", [
   PTO_DpsInitOpInterface,
   OpPipeInterface,

include/PTO/IR/PTOTypeDefs.td

@@ -228,6 +228,35 @@ def TileBufType : TypeDef<PTO_Dialect, "TileBuf"> {
   }];
 }
 
+def ConvTileType : TypeDef<PTO_Dialect, "ConvTile"> {
+  let mnemonic = "conv_tile";
+  let parameters = (ins
+    ArrayRefParameter<"int64_t">:$shape,
+    "mlir::Type":$elementType,
+    "mlir::IntegerAttr":$bufferSize,
+    "mlir::pto::AddressSpaceAttr":$memorySpace,
+    "mlir::pto::LayoutAttr":$layout,
+    "mlir::pto::ConvTileConfigAttr":$config
+  );
+
+  let hasCustomAssemblyFormat = 1;
+
+  let extraClassDeclaration = [{
+    int64_t getRank() const { return getShape().size(); }
+    int64_t getDimSize(unsigned idx) const { return getShape()[idx]; }
+    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;
+  }];
+}
+
 // =============================================================================
 // MultiTileBufType
 // =============================================================================

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);