[Linalg] Add matchers to infer which Convolution Op a given linalg.generic is

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

@@ -110,6 +110,15 @@ GenericOp makeMemRefCopyOp(OpBuilder &b, Location loc, Value from, Value to);
 std::optional<SmallVector<ReassociationIndices>>
 getReassociationMapForFoldingUnitDims(ArrayRef<OpFoldResult> mixedSizes);
 
+//===----------------------------------------------------------------------===//
+// Convolution matcher utility
+//===----------------------------------------------------------------------===//
+
+template <typename ConvOpTy>
+bool isaConvolutionOpOfType(LinalgOp op,
+                            SmallVector<int64_t> *dilations = nullptr,
+                            SmallVector<int64_t> *strides = nullptr);
+
 //===----------------------------------------------------------------------===//
 // Fusion / Tiling utilities
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Linalg/Transforms/Specialize.cpp

@@ -237,6 +237,286 @@ static FailureOr<LinalgOp> specializeLinalgContractions(RewriterBase &rewriter,
   return replaceWithMatmulVariant<MatmulOp>(rewriter, genericOp);
 }
 
+/// Utility to create a `genericOp` with a convolution op of type `ConvOpTy`
+/// with `dilations` and `strides`.
+template <typename ConvOpTy>
+static FailureOr<LinalgOp>
+specializeToConvOp(RewriterBase &rewriter, GenericOp genericOp,
+                   ArrayRef<int64_t> dilations, ArrayRef<int64_t> strides) {
+  SmallVector<Value> inputs = genericOp.getDpsInputs();
+  ValueRange outputs = genericOp.getDpsInits();
+  SmallVector<AffineMap> indexingMaps = genericOp.getIndexingMapsArray();
+  SmallVector<Type> resultTypes = genericOp.hasPureTensorSemantics()
+                                      ? TypeRange(ValueRange(outputs))
+                                      : TypeRange{};
+  LinalgOp namedOp;
+  if constexpr (std::is_same_v<ConvOpTy, linalg::Conv1DOp> ||
+                std::is_same_v<ConvOpTy, linalg::Conv2DOp> ||
+                std::is_same_v<ConvOpTy, linalg::Conv3DOp>) {
+    namedOp = rewriter.replaceOpWithNewOp<ConvOpTy>(genericOp, resultTypes,
+                                                    inputs, outputs);
+  } else {
+    Attribute stridesAttr = rewriter.getI64TensorAttr(strides);
+    Attribute dilationsAttr = rewriter.getI64TensorAttr(dilations);
+    namedOp = rewriter.replaceOpWithNewOp<ConvOpTy>(
+        genericOp, resultTypes, inputs, outputs, stridesAttr, dilationsAttr);
+  }
+  return namedOp;
+}
+
+static FailureOr<LinalgOp>
+inferAndSpecializeBasedOnRank2ConvIteratorTypes(RewriterBase &rewriter,
+                                                GenericOp genericOp) {
+  SmallVector<int64_t> dilations, strides;
+  if (isaConvolutionOpOfType<linalg::Conv1DOp>(genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::Conv1DOp>(rewriter, genericOp, dilations,
+                                                strides);
+  return failure();
+}
+
+static FailureOr<LinalgOp>
+inferAndSpecializeBasedOnRank4ConvIteratorTypes(RewriterBase &rewriter,
+                                                GenericOp genericOp) {
+  SmallVector<int64_t> dilations, strides;
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv1DNcwCwOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv1DNcwCwOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv1DNwcWcOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv1DNwcWcOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DOp>(genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::Conv2DOp>(rewriter, genericOp, dilations,
+                                                strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNcwMaxOp>(genericOp, &dilations,
+                                                      &strides))
+    return specializeToConvOp<linalg::PoolingNcwMaxOp>(rewriter, genericOp,
+                                                       dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNcwSumOp>(genericOp, &dilations,
+                                                      &strides))
+    return specializeToConvOp<linalg::PoolingNcwSumOp>(rewriter, genericOp,
+                                                       dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNwcMaxOp>(genericOp, &dilations,
+                                                      &strides))
+    return specializeToConvOp<linalg::PoolingNwcMaxOp>(rewriter, genericOp,
+                                                       dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNwcMinOp>(genericOp, &dilations,
+                                                      &strides))
+    return specializeToConvOp<linalg::PoolingNwcMinOp>(rewriter, genericOp,
+                                                       dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNwcSumOp>(genericOp, &dilations,
+                                                      &strides))
+    return specializeToConvOp<linalg::PoolingNwcSumOp>(rewriter, genericOp,
+                                                       dilations, strides);
+  return failure();
+}
+
+static FailureOr<LinalgOp>
+inferAndSpecializeBasedOnRank5ConvIteratorTypes(RewriterBase &rewriter,
+                                                GenericOp genericOp) {
+  SmallVector<int64_t> dilations, strides;
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv1DNwcWcmOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv1DNwcWcmOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv1DNwcWcfOp>(genericOp, &dilations,
+                                                     &strides))
+    return specializeToConvOp<linalg::Conv1DNwcWcfOp>(rewriter, genericOp,
+                                                      dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv1DNcwFcwOp>(genericOp, &dilations,
+                                                     &strides))
+    return specializeToConvOp<linalg::Conv1DNcwFcwOp>(rewriter, genericOp,
+                                                      dilations, strides);
+  return failure();
+}
+
+static FailureOr<LinalgOp>
+inferAndSpecializeBasedOnRank6ConvIteratorTypes(RewriterBase &rewriter,
+                                                GenericOp genericOp) {
+  SmallVector<int64_t> dilations, strides;
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv2DNchwChwOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv2DNchwChwOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv2DNhwcHwcOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv2DNhwcHwcOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv2DNhwcHwcQOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv2DNhwcHwcQOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv3DOp>(genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::Conv3DOp>(rewriter, genericOp, dilations,
+                                                strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNchwMaxOp>(genericOp, &dilations,
+                                                       &strides))
+    return specializeToConvOp<linalg::PoolingNchwMaxOp>(rewriter, genericOp,
+                                                        dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNchwSumOp>(genericOp, &dilations,
+                                                       &strides))
+    return specializeToConvOp<linalg::PoolingNchwSumOp>(rewriter, genericOp,
+                                                        dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNhwcMaxOp>(genericOp, &dilations,
+                                                       &strides))
+    return specializeToConvOp<linalg::PoolingNhwcMaxOp>(rewriter, genericOp,
+                                                        dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNhwcMinOp>(genericOp, &dilations,
+                                                       &strides))
+    return specializeToConvOp<linalg::PoolingNhwcMinOp>(rewriter, genericOp,
+                                                        dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNhwcSumOp>(genericOp, &dilations,
+                                                       &strides))
+    return specializeToConvOp<linalg::PoolingNhwcSumOp>(rewriter, genericOp,
+                                                        dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNhwcMaxUnsignedOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::PoolingNhwcMaxUnsignedOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNhwcMinUnsignedOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::PoolingNhwcMinUnsignedOp>(
+        rewriter, genericOp, dilations, strides);
+  return failure();
+}
+
+static FailureOr<LinalgOp>
+inferAndSpecializeBasedOnRank7ConvIteratorTypes(RewriterBase &rewriter,
+                                                GenericOp genericOp) {
+  SmallVector<int64_t> dilations, strides;
+  if (isaConvolutionOpOfType<linalg::Conv2DNhwcFhwcOp>(genericOp, &dilations,
+                                                       &strides))
+    return specializeToConvOp<linalg::Conv2DNhwcFhwcOp>(rewriter, genericOp,
+                                                        dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DNhwcHwcfOp>(genericOp, &dilations,
+                                                       &strides))
+    return specializeToConvOp<linalg::Conv2DNhwcHwcfOp>(rewriter, genericOp,
+                                                        dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DNchwFchwOp>(genericOp, &dilations,
+                                                       &strides))
+    return specializeToConvOp<linalg::Conv2DNchwFchwOp>(rewriter, genericOp,
+                                                        dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DNhwcFhwcQOp>(genericOp, &dilations,
+                                                        &strides))
+    return specializeToConvOp<linalg::Conv2DNhwcFhwcQOp>(rewriter, genericOp,
+                                                         dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DNchwFchwQOp>(genericOp, &dilations,
+                                                        &strides))
+    return specializeToConvOp<linalg::Conv2DNchwFchwQOp>(rewriter, genericOp,
+                                                         dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DNhwcHwcfQOp>(genericOp, &dilations,
+                                                        &strides))
+    return specializeToConvOp<linalg::Conv2DNhwcHwcfQOp>(rewriter, genericOp,
+                                                         dilations, strides);
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv2DNhwcHwcmOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv2DNhwcHwcmOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv2DNhwcHwcmQOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv2DNhwcHwcmQOp>(
+        rewriter, genericOp, dilations, strides);
+  return failure();
+}
+
+static FailureOr<LinalgOp>
+inferAndSpecializeBasedOnRank8ConvIteratorTypes(RewriterBase &rewriter,
+                                                GenericOp genericOp) {
+  SmallVector<int64_t> dilations, strides;
+  if (isaConvolutionOpOfType<linalg::Conv2DNgchwFgchwOp>(genericOp, &dilations,
+                                                         &strides))
+    return specializeToConvOp<linalg::Conv2DNgchwFgchwOp>(rewriter, genericOp,
+                                                          dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DNgchwGfchwOp>(genericOp, &dilations,
+                                                         &strides))
+    return specializeToConvOp<linalg::Conv2DNgchwGfchwOp>(rewriter, genericOp,
+                                                          dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DNgchwGfchwQOp>(genericOp, &dilations,
+                                                          &strides))
+    return specializeToConvOp<linalg::Conv2DNgchwGfchwQOp>(rewriter, genericOp,
+                                                           dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DNhwgcGfhwcOp>(genericOp, &dilations,
+                                                         &strides))
+    return specializeToConvOp<linalg::Conv2DNhwgcGfhwcOp>(rewriter, genericOp,
+                                                          dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv2DNhwgcGfhwcQOp>(genericOp, &dilations,
+                                                          &strides))
+    return specializeToConvOp<linalg::Conv2DNhwgcGfhwcQOp>(rewriter, genericOp,
+                                                           dilations, strides);
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv3DNcdhwCdhwOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv3DNcdhwCdhwOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv3DNdhwcDhwcOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv3DNdhwcDhwcOp>(
+        rewriter, genericOp, dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNdhwcMaxOp>(genericOp, &dilations,
+                                                        &strides))
+    return specializeToConvOp<linalg::PoolingNdhwcMaxOp>(rewriter, genericOp,
+                                                         dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNdhwcMinOp>(genericOp, &dilations,
+                                                        &strides))
+    return specializeToConvOp<linalg::PoolingNdhwcMinOp>(rewriter, genericOp,
+                                                         dilations, strides);
+  if (isaConvolutionOpOfType<linalg::PoolingNdhwcSumOp>(genericOp, &dilations,
+                                                        &strides))
+    return specializeToConvOp<linalg::PoolingNdhwcSumOp>(rewriter, genericOp,
+                                                         dilations, strides);
+  return failure();
+}
+
+static FailureOr<LinalgOp>
+inferAndSpecializeBasedOnRank9ConvIteratorTypes(RewriterBase &rewriter,
+                                                GenericOp genericOp) {
+  SmallVector<int64_t> dilations, strides;
+  if (isaConvolutionOpOfType<linalg::Conv3DNcdhwFcdhwOp>(genericOp, &dilations,
+                                                         &strides))
+    return specializeToConvOp<linalg::Conv3DNcdhwFcdhwOp>(rewriter, genericOp,
+                                                          dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv3DNdhwcDhwcfOp>(genericOp, &dilations,
+                                                         &strides))
+    return specializeToConvOp<linalg::Conv3DNdhwcDhwcfOp>(rewriter, genericOp,
+                                                          dilations, strides);
+  if (isaConvolutionOpOfType<linalg::Conv3DNdhwcDhwcfQOp>(genericOp, &dilations,
+                                                          &strides))
+    return specializeToConvOp<linalg::Conv3DNdhwcDhwcfQOp>(rewriter, genericOp,
+                                                           dilations, strides);
+  if (isaConvolutionOpOfType<linalg::DepthwiseConv3DNdhwcDhwcmOp>(
+          genericOp, &dilations, &strides))
+    return specializeToConvOp<linalg::DepthwiseConv3DNdhwcDhwcmOp>(
+        rewriter, genericOp, dilations, strides);
+  return failure();
+}
+
+// Converts linalg.generic to named linalg.*conv/pooling* where possible. To
+// improve the search speed, the convolution ops have been segregated based on
+// the rank of iterator types array.
+static FailureOr<LinalgOp>
+inferAndSpecializeToConvolutionOp(RewriterBase &rewriter, GenericOp genericOp) {
+  SmallVector<utils::IteratorType> iteratorTypes =
+      genericOp.getIteratorTypesArray();
+  unsigned totalIterators = iteratorTypes.size();
+  switch (totalIterators) {
+  case 2:
+    return inferAndSpecializeBasedOnRank2ConvIteratorTypes(rewriter, genericOp);
+  case 4:
+    return inferAndSpecializeBasedOnRank4ConvIteratorTypes(rewriter, genericOp);
+  case 5:
+    return inferAndSpecializeBasedOnRank5ConvIteratorTypes(rewriter, genericOp);
+  case 6:
+    return inferAndSpecializeBasedOnRank6ConvIteratorTypes(rewriter, genericOp);
+  case 7:
+    return inferAndSpecializeBasedOnRank7ConvIteratorTypes(rewriter, genericOp);
+  case 8:
+    return inferAndSpecializeBasedOnRank8ConvIteratorTypes(rewriter, genericOp);
+  case 9:
+    return inferAndSpecializeBasedOnRank9ConvIteratorTypes(rewriter, genericOp);
+  }
+  return failure();
+}
+
 } // namespace
 
 //===----------------------------------------------------------------------===//
@@ -316,6 +596,11 @@ FailureOr<LinalgOp> mlir::linalg::specializeGenericOp(RewriterBase &rewriter,
   if (isaContractionOpInterface(genericOp)) {
     return specializeLinalgContractions(rewriter, genericOp);
   }
+
+  // Convolution - e.g. *conv/pooling*
+  if (isaConvolutionOpInterface(genericOp)) {
+    return inferAndSpecializeToConvolutionOp(rewriter, genericOp);
+  }
   return failure();
 }