big refactor

src/PIM/Dialect/Spatial/Spatial.td

@@ -5,6 +5,7 @@ include "mlir/IR/OpBase.td"
 include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/BuiltinTypes.td"
 include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/EnumAttr.td"
 include "mlir/IR/RegionKindInterface.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/ParallelCombiningOpInterface.td"
@@ -14,11 +15,84 @@ def SpatialDialect : Dialect {
   let name = "spat";
   let summary = "Dialect designed for deep learning computation in a spatial architecture";
   let cppNamespace = "::onnx_mlir::spatial";
+  let useDefaultAttributePrinterParser = 0;
+  let extraClassDeclaration = [{
+    ::mlir::Attribute parseAttribute(::mlir::DialectAsmParser &parser,
+                                     ::mlir::Type type) const override;
+    void printAttribute(::mlir::Attribute attr,
+                        ::mlir::DialectAsmPrinter &printer) const override;
+  }];
+}
+
+def SpatialLayoutCapabilityInterface : OpInterface<"SpatialLayoutCapabilityInterface"> {
+  let description = [{
+    Contract implemented by logical Spatial planning operations that expose
+    their legal physical layout alternatives to the Spatial planner.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      "Return legal physical layout alternatives for this operation and its current operand layouts.",
+      "::llvm::SmallVector<::onnx_mlir::spatial::LayoutAlternative>",
+      "getLayoutAlternatives",
+      (ins "const ::onnx_mlir::spatial::SpatialTargetResources &":$target,
+           "::llvm::ArrayRef<::onnx_mlir::spatial::PhysicalLayout>":$operandLayouts)>
+  ];
+
+  let cppNamespace = "::onnx_mlir::spatial";
+}
+
+def SpatLogicalLayoutNCHW : I32EnumAttrCase<"NCHW", 0, "nchw">;
+def SpatLogicalLayout : I32EnumAttr<"LogicalLayout", "Logical tensor layout", [
+  SpatLogicalLayoutNCHW
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::onnx_mlir::spatial";
+}
+
+def SpatLogicalLayoutAttr : EnumAttr<SpatialDialect, SpatLogicalLayout, "logical_layout"> {
+  let assemblyFormat = "$value";
+}
+
+def SpatPhysicalLayoutDenseNCHW : I32EnumAttrCase<"DenseNCHW", 0, "dense_nchw">;
+def SpatPhysicalLayoutNCHWRowStrip : I32EnumAttrCase<"NCHWRowStrip", 1, "nchw_row_strip">;
+def SpatPhysicalLayoutNHWCRowStrip : I32EnumAttrCase<"NHWCRowStrip", 2, "nhwc_row_strip">;
+def SpatPhysicalLayoutFragmented : I32EnumAttrCase<"Fragmented", 3, "fragmented">;
+def SpatPhysicalLayout : I32EnumAttr<"PhysicalLayout", "Physical tensor layout", [
+  SpatPhysicalLayoutDenseNCHW,
+  SpatPhysicalLayoutNCHWRowStrip,
+  SpatPhysicalLayoutNHWCRowStrip,
+  SpatPhysicalLayoutFragmented
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::onnx_mlir::spatial";
+}
+
+def SpatPhysicalLayoutAttr : EnumAttr<SpatialDialect, SpatPhysicalLayout, "physical_layout"> {
+  let assemblyFormat = "$value";
+}
+
+def SpatBlueprintModePhysicalView : I32EnumAttrCase<"PhysicalView", 0, "physical_view">;
+def SpatBlueprintModeFragmentAssembly : I32EnumAttrCase<"FragmentAssembly", 1, "fragment_assembly">;
+def SpatBlueprintMode : I32EnumAttr<"BlueprintMode", "Blueprint reconstruction mode", [
+  SpatBlueprintModePhysicalView,
+  SpatBlueprintModeFragmentAssembly
+]> {
+  let genSpecializedAttr = 0;
+  let cppNamespace = "::onnx_mlir::spatial";
+}
+
+def SpatBlueprintModeAttr : EnumAttr<SpatialDialect, SpatBlueprintMode, "blueprint_mode"> {
+  let assemblyFormat = "$value";
 }
 
 class SpatOp<string mnemonic, list<Trait> traits = []> :
     Op<SpatialDialect, mnemonic, traits>;
 
+class SpatLayoutPlanOp<string mnemonic> : SpatOp<mnemonic,
+    [SpatialLayoutCapabilityInterface,
+     DeclareOpInterfaceMethods<SpatialLayoutCapabilityInterface>]>;
+
 // TODO maybe remove and use AnyRankedTensor directly
 def SpatTensor :
   AnyTypeOf<[AnyMemRef, AnyRankedTensor], "", "::mlir::ShapedType">;
@@ -252,7 +326,7 @@ def SpatConcatOp : SpatOp<"concat", []> {
 // Planning
 //===----------------------------------------------------------------------===//
 
-def SpatConv2DPlanOp : SpatOp<"conv2d_plan", []> {
+def SpatConv2DPlanOp : SpatLayoutPlanOp<"conv2d_plan"> {
   let summary = "Structured Conv2D planning op that preserves logical ONNX geometry";
 
   let arguments = (ins
@@ -263,7 +337,7 @@ def SpatConv2DPlanOp : SpatOp<"conv2d_plan", []> {
     DenseI64ArrayAttr:$strides,
     DenseI64ArrayAttr:$dilations,
     I64Attr:$group,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -273,12 +347,28 @@ def SpatConv2DPlanOp : SpatOp<"conv2d_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatReluPlanOp : SpatOp<"relu_plan", []> {
+def SpatFlattenPlanOp : SpatLayoutPlanOp<"flatten_plan"> {
+  let summary = "Layout-aware static Flatten planning op";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    I64Attr:$axis,
+    SpatLogicalLayoutAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
+def SpatReluPlanOp : SpatLayoutPlanOp<"relu_plan"> {
   let summary = "Layout-aware ReLU planning op";
 
   let arguments = (ins
     SpatTensor:$input,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -288,12 +378,12 @@ def SpatReluPlanOp : SpatOp<"relu_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatSiluPlanOp : SpatOp<"silu_plan", []> {
+def SpatSiluPlanOp : SpatLayoutPlanOp<"silu_plan"> {
   let summary = "Layout-aware SiLU planning op";
 
   let arguments = (ins
     SpatTensor:$input,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -303,7 +393,22 @@ def SpatSiluPlanOp : SpatOp<"silu_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> {
+def SpatResizeNearestPlanOp : SpatLayoutPlanOp<"resize_nearest_plan"> {
+  let summary = "Layout-aware nearest asymmetric Resize planning op";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    SpatLogicalLayoutAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
+def SpatMaxPool2DPlanOp : SpatLayoutPlanOp<"max_pool2d_plan"> {
   let summary = "Layout-aware 2D NCHW MaxPool planning op";
 
   let arguments = (ins
@@ -312,7 +417,7 @@ def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> {
     DenseI64ArrayAttr:$pads,
     DenseI64ArrayAttr:$strides,
     DenseI64ArrayAttr:$dilations,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -322,12 +427,12 @@ def SpatMaxPool2DPlanOp : SpatOp<"max_pool2d_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatGlobalAveragePoolPlanOp : SpatOp<"global_average_pool_plan", []> {
+def SpatGlobalAveragePoolPlanOp : SpatLayoutPlanOp<"global_average_pool_plan"> {
   let summary = "Layout-aware NCHW global average-pool planning op";
 
   let arguments = (ins
     SpatTensor:$input,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -337,13 +442,13 @@ def SpatGlobalAveragePoolPlanOp : SpatOp<"global_average_pool_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> {
+def SpatBiasAddPlanOp : SpatLayoutPlanOp<"bias_add_plan"> {
   let summary = "Layout-aware Conv-style bias add planning op";
 
   let arguments = (ins
     SpatTensor:$input,
     SpatTensor:$bias,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -353,13 +458,13 @@ def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatAddPlanOp : SpatOp<"add_plan", []> {
+def SpatAddPlanOp : SpatLayoutPlanOp<"add_plan"> {
   let summary = "Layout-aware elementwise add planning op";
 
   let arguments = (ins
     SpatTensor:$lhs,
     SpatTensor:$rhs,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -369,13 +474,13 @@ def SpatAddPlanOp : SpatOp<"add_plan", []> {
   let hasVerifier = 1;
 }
 
-def SpatConcatPlanOp : SpatOp<"concat_plan", []> {
+def SpatConcatPlanOp : SpatLayoutPlanOp<"concat_plan"> {
   let summary = "Layout-aware tensor concatenation planning op";
 
   let arguments = (ins
     Variadic<SpatTensor>:$inputs,
     I64Attr:$axis,
-    StrAttr:$logicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout
   );
 
   let results = (outs
@@ -391,12 +496,12 @@ def SpatBlueprintOp : SpatOp<"blueprint", []> {
   let arguments = (ins
     SpatTensor:$input,
     Variadic<SpatTensor>:$fragments,
-    StrAttr:$logicalLayout,
-    StrAttr:$physicalLayout,
+    SpatLogicalLayoutAttr:$logicalLayout,
+    SpatPhysicalLayoutAttr:$physicalLayout,
     DenseI64ArrayAttr:$fragmentOffsets,
     DenseI64ArrayAttr:$fragmentSizes,
     StrAttr:$indexMap,
-    OptionalAttr<StrAttr>:$mode,
+    OptionalAttr<SpatBlueprintModeAttr>:$mode,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentOperandIndices,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceSlots,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceOffsets,
@@ -418,9 +523,9 @@ def SpatMaterializeLayoutOp : SpatOp<"materialize_layout", []> {
 
   let arguments = (ins
     SpatTensor:$input,
-    StrAttr:$logicalLayout,
-    StrAttr:$sourcePhysicalLayout,
-    StrAttr:$targetPhysicalLayout
+    SpatLogicalLayoutAttr:$logicalLayout,
+    SpatPhysicalLayoutAttr:$sourcePhysicalLayout,
+    SpatPhysicalLayoutAttr:$targetPhysicalLayout
   );
 
   let results = (outs