[mlir][linalg] Introduce new linalg.conv op

mlir/include/mlir/Dialect/Linalg/IR/LinalgBase.td

@@ -81,4 +81,43 @@ def IteratorTypeEnum : EnumAttr<Linalg_Dialect, IteratorType, "iterator_type"> {
 def IteratorTypeArrayAttr : TypedArrayAttrBase<IteratorTypeEnum,
   "Iterator type should be an enum.">;
 
+
+def ConvolutionDimArray : ArrayRefParameter<"ConvDimEnum"> {
+  let printer = [{
+    $_printer << '{';
+    llvm::interleaveComma($_self, $_printer, [&](ConvDimEnum en) {
+        $_printer.printStrippedAttrOrType(en);
+    });
+    $_printer << '}';
+  }];
+
+  let parser = [{
+    [&]() -> FailureOr<SmallVector<ConvDimEnum>> {
+        using Result = SmallVector<ConvDimEnum>;
+        if ($_parser.parseLBrace())
+            return failure();
+        FailureOr<Result> result = FieldParser<Result>::parse($_parser);
+        if (failed(result))
+            return failure();
+        if ($_parser.parseRBrace())
+            return failure();
+        return result;
+    }()
+  }];
+}
+
+/// Attribute that represents an ordered set of tensor dimensions involved in
+/// convolution.
+def ConvDimsAttr : AttrDef<Linalg_Dialect, "ConvDims", [], "::mlir::Attribute"> {
+  let mnemonic = "conv_dims";
+
+  let parameters = (ins
+    ConvolutionDimArray:$dims
+  );
+
+  let assemblyFormat = "$dims";
+
+  let returnType = "mlir::linalg::ConvDims";
+  let convertFromStorage = "mlir::linalg::ConvDims($_self.getDims())";
+}
 #endif // LINALG_BASE

mlir/include/mlir/Dialect/Linalg/IR/LinalgEnums.td

@@ -63,4 +63,30 @@ def TypeFn : I32EnumAttr<"TypeFn", "", [
   let cppNamespace = "::mlir::linalg";
 }
 
+
+class ConvDimEnumAttrCase<string sym, int val, string str = sym>
+    : IntEnumAttrCaseBase<I8, sym, str, val>;
+
+def ConvDimEnumAttr :
+    IntEnumAttr<I8, "ConvDimEnum", "summary", [
+      /// Batch is a dimension of input and output, indexed from a parallel loop.
+      ConvDimEnumAttrCase<"BATCH", 0, "N">,
+      /// Input channel is a dimension in all tensors, indexed from a reduction loop.
+      /// Depthwise convolutions perform no reduction across channels and therefore
+      /// do not use this.
+      ConvDimEnumAttrCase<"INPUT_CHANNEL", 1, "C">,
+      /// Output channel is a dimension in filter and output, index from a parallel loop.
+      ConvDimEnumAttrCase<"OUTPUT_CHANNEL", 2, "F">,
+      /// Group is a dimension in all tensors and indexed from a parallel loop.
+      ConvDimEnumAttrCase<"GROUP", 3, "G">,
+      /// Spatial dimensions occur in all tensors. Output is indexed from a parallel
+      /// loop, filter from a reduction loop and input from both.
+      ConvDimEnumAttrCase<"SPATIAL_0", 4, "0">,
+      ConvDimEnumAttrCase<"SPATIAL_1", 5, "1">,
+      ConvDimEnumAttrCase<"SPATIAL_2", 6, "2">,
+    ]> {
+  let underlyingType = "uint8_t";
+  let cppNamespace = "::mlir::linalg";
+}
+
 #endif // LINALG_ENUMS

mlir/include/mlir/Dialect/Linalg/IR/LinalgInterfaces.h

@@ -117,6 +117,33 @@ FailureOr<ConvolutionDimensions> inferConvolutionDims(LinalgOp linalgOp);
 bool isaConvolutionOpInterface(LinalgOp linalgOp,
                                bool allowEmptyConvolvedDims = false);
 
+enum class ConvDimEnum : uint8_t;
+class ConvDims {
+  ArrayRef<ConvDimEnum> storage;
+
+public:
+  ConvDims() = default;
+  ConvDims(ArrayRef<ConvDimEnum> dims) : storage(dims) {}
+  ConvDims(SmallVectorImpl<ConvDimEnum> &dims) : storage(dims) {}
+
+  bool contains(ConvDimEnum dim) const {
+    return llvm::is_contained(storage, dim);
+  }
+
+  int64_t getPos(ConvDimEnum dim) const {
+    auto it = llvm::find(storage, dim);
+    assert(it != storage.end() && "expected dimension to be present");
+
+    return std::distance(storage.begin(), it);
+  }
+
+  int64_t size() const { return storage.size(); }
+  operator ArrayRef<ConvDimEnum>() const { return storage; }
+
+  auto begin() const { return storage.begin(); }
+  auto end() const { return storage.end(); }
+};
+
 /// Checks whether `linalgOp` is semantically equivalent to a `linalg.copyOp`.
 bool isaCopyOpInterface(LinalgOp linalgOp);
 

mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td

@@ -683,6 +683,122 @@ def MatmulOp : LinalgStructuredBase_Op<"matmul", [
     }];
 }
 
+//===----------------------------------------------------------------------===//
+// Op definition for ConvOp
+//===----------------------------------------------------------------------===//
+
+def ConvOp : LinalgStructuredBase_Op<"conv", [AttrSizedOperandSegments]> {
+
+  let summary = [{
+    Configurable convolution operation with configurable tensor layouts.
+  }];
+  let description = [{
+    Numeric casting is performed on the operands to the inner multiply,
+    promoting them to the same data type as the accumulator/output.
+
+    The subtype of convolution is defined by the tensor layouts of `input`,
+    `filter`, and `output`. For example, a standard batched 2D convolution:
+
+    ```
+      %0 = linalg.conv {
+          input_dims = #linalg<conv_dims {N, C, "1", "0"}>,
+          filter_dims = #linalg<conv_dims {F, C, "1", "0"}>,
+          output_dims = #linalg<conv_dims {N, F, "1", "0"}>
+        }
+        ins(%input, %filter : tensor<8x4x16x16xf32>, tensor<16x4x3x3xf32>)
+        outs(%output : tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32>
+    ```
+
+    This op could be turned into a depthwise convolution as follows:
+    ```
+      %0 = linalg.conv {
+          input_dims = #linalg<conv_dims {N, G, "1", "0"}>,
+          filter_dims = #linalg<conv_dims {G, "1", "0"}>,
+          output_dims = #linalg<conv_dims {N, G, "1", "0"}>
+        }
+        ins(%input, %filter : tensor<8x4x16x16xf32>, tensor<4x3x3xf32>)
+        outs(%output : tensor<8x4x14x14xf32>) -> tensor<8x4x14x14xf32>
+    ```
+
+    For the detailed semantics of the available tensor dimensions, refer to
+    `mlir::linalg::ConvDimsEnum`.
+
+    Strides and dilations can be supplied as optional attributes, where
+    `strides[0]` is the stride for the `SPATIAL_0` dimension, etc.
+  }];
+
+  let arguments = (ins
+    Variadic<AnyType>:$inputs, Variadic<AnyShaped>:$outputs,
+    ConvDimsAttr:$input_dims, ConvDimsAttr:$filter_dims, ConvDimsAttr:$output_dims,
+    OptionalAttr<I64ElementsAttr>:$strides, OptionalAttr<I64ElementsAttr>:$dilations
+  );
+  let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
+  let regions = (region AnyRegion:$region);
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilder<
+      (ins "TypeRange":$resTys, "Value":$input, "Value":$filter, "Value":$output, "ConvDims":$input_dims,
+            "ConvDims":$filter_dims, "ConvDims":$output_dims, "ArrayRef<int64_t>":$strides,
+            "ArrayRef<int64_t>":$dilations, CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+      [{
+        buildConvOp($_builder, $_state, resTys, input, filter, output,
+            input_dims, filter_dims, output_dims, strides, dilations,
+            attributes, ConvOp::getRegionBuilder());
+      }]>,
+    OpBuilder<
+      (ins "ValueRange":$inputs, "ValueRange":$outputs, "ConvDimsAttr":$input_dims,
+            "ConvDimsAttr":$filter_dims, "ConvDimsAttr":$output_dims,
+            CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+      [{
+        buildConvOp($_builder, $_state, std::nullopt, inputs, outputs,
+            input_dims, filter_dims, output_dims, nullptr, nullptr,
+            attributes, ConvOp::getRegionBuilder());
+      }]>,
+    OpBuilder<
+      (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
+            "ValueRange":$outputs, "ConvDimsAttr":$input_dims,
+            "ConvDimsAttr":$filter_dims, "ConvDimsAttr":$output_dims,
+            CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes),
+      [{
+        buildConvOp($_builder, $_state, resultTensorTypes,
+            inputs, outputs, input_dims, filter_dims, output_dims, nullptr, nullptr,
+            attributes, ConvOp::getRegionBuilder());
+      }]>
+  ];
+  let hasCustomAssemblyFormat = 1;
+  let hasFolder = 1;
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = structuredOpsBaseDecls # [{
+    SmallVector<utils::IteratorType> getIteratorTypesArray();
+    ArrayAttr getIndexingMaps();
+
+    /// Implements the block region builder.
+    static void regionBuilder(ImplicitLocOpBuilder &b,
+                              Block &block, ArrayRef<NamedAttribute> attrs);
+
+    /// Returns a list of AffineMap with the typical matmul indexing charactristic.
+    static SmallVector<AffineMap> getDefaultIndexingMaps(MLIRContext *context);
+
+    static std::function<void(ImplicitLocOpBuilder &,
+                              Block &, ArrayRef<NamedAttribute>)>
+    getRegionBuilder() { return regionBuilder; }
+
+    ::mlir::MutableOperandRange getDpsInitsMutable() { return getOutputsMutable(); }
+
+    bool hasDynamicIndexingMaps() { return true; }
+
+    /// Returns the number of spatial dimensions, i.e. 1 for 1D convolution,
+    /// 2 for 2D convolution, etc.
+    int64_t getNumSpatialDims();
+
+    bool isDepthwise();
+    bool isGrouped();
+    bool isBatched();
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // Named Linalg ops, implemented as a declarative configurations of generic ops.
 //===----------------------------------------------------------------------===//