Improve im2col conversion to handle permuted filter and output maps (#20150)

The orginal generic im2col implementation assumed that the filter and
output affine map would not have permutations. While this is true in
most named ops, it is not necessarily required. See for example
llvm/llvm-project#129547 where it was not the
case for `conv_3d_ncdhw_fcdhw` . While in this case we just "normalized"
the map upstream, this PR adds support for such maps directly.
Fixes :   #20139

---------

Signed-off-by: Nirvedh Meshram <[email protected]>

Author: nirvedhmeshram

compiler/src/iree/compiler/Codegen/Common/ConvolutionToIGEMM.cpp

@@ -101,8 +101,8 @@ convertToIGEMMAndSetConfig(FunctionOpInterface funcOp,
   MLIRContext *context = funcOp->getContext();
   {
     RewritePatternSet patterns(context);
-    iree_compiler::IREE::LinalgExt::populateConv2DToIm2colOpPatterns(patterns,
-                                                                     controlFn);
+    iree_compiler::IREE::LinalgExt::populateConvToIm2colOpPatterns(patterns,
+                                                                   controlFn);
     if (configFn.has_value()) {
       patterns.add<SetIGEMMConfiguration>(context, configFn.value());
     }

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/BUILD.bazel

@@ -30,8 +30,8 @@ iree_gentbl_cc_library(
 iree_compiler_cc_library(
     name = "Transforms",
     srcs = [
-        "ConvertConv2DToIm2ColOp.cpp",
         "ConvertConv2DToWinograd.cpp",
+        "ConvertConvToIm2ColOp.cpp",
         "ConvertToLoops.cpp",
         "DecomposeAttention.cpp",
         "DecomposeIm2col.cpp",

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/CMakeLists.txt

@@ -27,8 +27,8 @@ iree_cc_library(
     "Passes.h.inc"
     "Transforms.h"
   SRCS
-    "ConvertConv2DToIm2ColOp.cpp"
     "ConvertConv2DToWinograd.cpp"
+    "ConvertConvToIm2ColOp.cpp"
     "ConvertToLoops.cpp"
     "DecomposeAttention.cpp"
     "DecomposeIm2col.cpp"

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ConvertConv2DToIm2ColOp.cpp renamed to compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/ConvertConvToIm2ColOp.cpp

@@ -14,7 +14,7 @@
 
 namespace mlir::iree_compiler::IREE::LinalgExt {
 
-#define GEN_PASS_DEF_CONVERTCONV2DTOIM2COLOPPASS
+#define GEN_PASS_DEF_CONVERTCONVTOIM2COLOPPASS
 #include "iree/compiler/Dialect/LinalgExt/Transforms/Passes.h.inc"
 
 static bool hasAllOneValues(ArrayRef<int64_t> attr) {
@@ -254,14 +254,14 @@ class ConvertConvGeneric final
   std::optional<ControlFnTy> controlFn;
 };
 
-struct ConvertConv2DToIm2ColOpPass final
-    : impl::ConvertConv2DToIm2ColOpPassBase<ConvertConv2DToIm2ColOpPass> {
+struct ConvertConvToIm2ColOpPass final
+    : impl::ConvertConvToIm2ColOpPassBase<ConvertConvToIm2ColOpPass> {
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<tensor::TensorDialect, IREELinalgExtDialect>();
   }
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
-    populateConv2DToIm2colOpPatterns(patterns);
+    populateConvToIm2colOpPatterns(patterns);
     if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
       return signalPassFailure();
     }
@@ -270,8 +270,8 @@ struct ConvertConv2DToIm2ColOpPass final
 
 } // namespace
 
-void populateConv2DToIm2colOpPatterns(RewritePatternSet &patterns,
-                                      std::optional<ControlFnTy> controlFn) {
+void populateConvToIm2colOpPatterns(RewritePatternSet &patterns,
+                                    std::optional<ControlFnTy> controlFn) {
   patterns.insert<ConvertConvGeneric>(patterns.getContext(),
                                       std::move(controlFn));
 }

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Passes.h

@@ -32,7 +32,7 @@ splitReduction(RewriterBase &rewriter, LinalgExt::TopkOp topkOp,
 /// op and reshapes on the inputs.
 /// TODO(Max191): Maybe move to transforms and use a funcOp walk instead of a
 ///               rewrite pattern for this.
-void populateConv2DToIm2colOpPatterns(
+void populateConvToIm2colOpPatterns(
     RewritePatternSet &patterns,
     std::optional<std::function<bool(Operation *)>> controlFn = std::nullopt);
 

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/Passes.td

@@ -66,8 +66,8 @@ def DecomposeWinogradTransformPass :
       "Decomposes winograd transform ops into linalg ops";
 }
 
-def ConvertConv2DToIm2ColOpPass :
-    InterfacePass<"iree-linalg-ext-convert-conv2d-to-im2col-op", "mlir::FunctionOpInterface"> {
+def ConvertConvToIm2ColOpPass :
+    InterfacePass<"iree-linalg-ext-convert-conv-to-im2col-op", "mlir::FunctionOpInterface"> {
   let summary = "Convert linalg convolution ops to im2col gemm based implementation.";
 }
 

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/BUILD.bazel

@@ -16,8 +16,8 @@ iree_lit_test_suite(
     name = "lit",
     srcs = enforce_glob(
         [
-            "conv2d_to_im2col.mlir",
             "conv2d_to_winograd.mlir",
+            "conv_to_im2col.mlir",
             "convert_to_loops.mlir",
             "convert_to_online_attention.mlir",
             "decompose_im2col.mlir",

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/CMakeLists.txt

@@ -14,8 +14,8 @@ iree_lit_test_suite(
   NAME
     lit
   SRCS
-    "conv2d_to_im2col.mlir"
     "conv2d_to_winograd.mlir"
+    "conv_to_im2col.mlir"
     "convert_to_loops.mlir"
     "convert_to_online_attention.mlir"
     "decompose_im2col.mlir"

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/conv2d_to_im2col.mlir renamed to compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/conv_to_im2col.mlir

@@ -1,4 +1,4 @@
-// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(util.func(iree-linalg-ext-convert-conv2d-to-im2col-op))" %s | FileCheck %s
+// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(util.func(iree-linalg-ext-convert-conv-to-im2col-op))" %s | FileCheck %s
 
 util.func public @conv_2d_nhwc_hwcf(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<3x3x4x16xf32>, %arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> {
   %0 = linalg.conv_2d_nhwc_hwcf
@@ -139,9 +139,9 @@ util.func public @conv_strided(%arg0: tensor<1x16x16x4xf16>, %arg1: tensor<3x3x4
 // CHECK:      util.return %[[MATMUL]] : tensor<1x7x7x16xf32>
 
 // -----
-#map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>
-#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d3, d6)>
-#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>
+#map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2 + d5, d3 + d6, d4)>
+#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d1, d4)>
+#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d2, d3, d1)>
 util.func public @conv_nhwc_hwfc(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<3x3x16x4xf32>, %arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> {
   %0 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%arg0, %arg1 : tensor<1x16x16x4xf32>, tensor<3x3x16x4xf32>) outs(%arg2 : tensor<1x14x14x16xf32>) {
   ^bb0(%in: f32, %in_0: f32, %out: f32):
@@ -195,3 +195,53 @@ util.func public @conv_2d_nhwc_fhwc(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<
 // CHECK-SAME:   iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction"]
 // CHECK-SAME:   ins(%[[IM2COL]], %[[COLLAPSED]] : tensor<1x14x14x36xf32>, tensor<16x36xf32>)
 // CHECK:      util.return %[[MATMUL]] : tensor<1x14x14x16xf32>
+
+// -----
+
+util.func public @conv_1d_ncw_fcw_transpose_maps(%arg0: tensor<1x8x130xf32>, %arg1: tensor<16x8x3xf32>) -> tensor<1x16x128xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %empty = tensor.empty() : tensor<1x16x128xf32>
+  %fill = linalg.fill ins(%cst : f32) outs(%empty : tensor<1x16x128xf32>) -> tensor<1x16x128xf32>
+  %0 = linalg.generic {
+    indexing_maps = [
+      affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d1 + d3)>,
+      affine_map<(d0, d1, d2, d3, d4) -> (d2, d4, d3)>,
+      affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d1)>],
+      iterator_types =
+      ["parallel", "parallel", "parallel", "reduction", "reduction"]}
+      ins(%arg0, %arg1 : tensor<1x8x130xf32>,tensor<16x8x3xf32>)
+      outs(%fill : tensor<1x16x128xf32>) {
+    ^bb0(%in: f32, %in_0: f32, %out: f32):
+      %8 = arith.mulf %in, %in_0 : f32
+      %9 = arith.addf %out, %8 : f32
+      linalg.yield %9 : f32
+    } -> tensor<1x16x128xf32>
+  util.return %0 : tensor<1x16x128xf32>
+}
+
+// CHECK-DAG:  #[[MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
+// CHECK-DAG:  #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+// CHECK-DAG:  #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>
+// CHECK:      util.func public @conv_1d_ncw_fcw_transpose_maps(
+// CHECK-SAME:   %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x8x130xf32>
+// CHECK-SAME:   %[[ARG1:[a-zA-Z0-9_]+]]: tensor<16x8x3xf32>
+// CHECK:      %[[CST:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK:      %[[EMPTY:.+]] = tensor.empty() : tensor<1x16x128xf32>
+// CHECK:      %[[FILL:.+]] = linalg.fill ins(%[[CST]] : f32) outs(%[[EMPTY]] : tensor<1x16x128xf32>) -> tensor<1x16x128xf32>
+// CHECK:      %[[EMPTY2:.+]] = tensor.empty() : tensor<1x128x24xf32>
+// CHECK:      %[[IM2COL:.+]] = iree_linalg_ext.im2col
+// CHECK-SAME:   strides = [1] dilations = [1] kernel_size = [3]
+// CHECK-SAME:   m_offset = [0] * [1] k_offset = [0] * [1]
+// CHECK-SAME:   batch_pos = [0] m_pos = [2] k_pos = [1]
+// CHECK-SAME:   ins(%[[ARG0]] : tensor<1x8x130xf32>)
+// CHECK-SAME:   outs(%[[EMPTY2]] : tensor<1x128x24xf32>) -> tensor<1x128x24xf32>
+// CHECK-DAG:  %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1, 2]] : tensor<16x8x3xf32> into tensor<16x24xf32>
+// CHECK:      %[[MATMUL:.+]] = linalg.generic
+// CHECK-SAME:   indexing_maps = [#[[MAP]], #[[MAP1]], #[[MAP2]]]
+// CHECK-SAME:   iterator_types = ["parallel", "parallel", "parallel", "reduction"]
+// CHECK-SAME:   ins(%[[COLLAPSED]], %[[IM2COL]] : tensor<16x24xf32>, tensor<1x128x24xf32>)
+// CHECK-SAME:   outs(%[[FILL]] : tensor<1x16x128xf32>) {
+// CHECK:          arith.mulf
+// CHECK:          arith.addf
+// CHECK:      } -> tensor<1x16x128xf32>
+// CHECK:      util.return %[[MATMUL]] : tensor<1x16x128xf32>

compiler/src/iree/compiler/Dialect/LinalgExt/Utils/Utils.cpp

@@ -445,20 +445,40 @@ getIGEMMGenericConvDetails(linalg::LinalgOp linalgOp) {
   auto outputShape = outputType.getShape();
   auto indexingMaps = linalgOp.getIndexingMapsArray();
   auto filterMap = indexingMaps[1];
+  auto outputMap = indexingMaps[2];
 
   SmallVector<int64_t> reductionDims;
   for (auto iter : llvm::enumerate(linalgOp.getIteratorTypesArray())) {
     if (linalg::isReductionIterator(iter.value())) {
       reductionDims.push_back(iter.index());
     }
   }
+
+  bool isOutputChannelFirst = false;
+  auto outputChannelPos = convDims.outputChannel;
+  auto outputImagePos = convDims.outputImage;
+
+  std::optional<int64_t> outputChannelLastDim = outputMap.getResultPosition(
+      getAffineDimExpr(outputChannelPos.back(), outputMap.getContext()));
+  std::optional<int64_t> outputImageFirstDim = outputMap.getResultPosition(
+      getAffineDimExpr(outputImagePos[0], outputMap.getContext()));
+  if (!outputImageFirstDim || !outputChannelLastDim) {
+    LDBG("output image or output channel dim not found in output.");
+    return failure();
+  }
+  if (outputChannelLastDim.value() < outputImageFirstDim.value())
+    isOutputChannelFirst = true;
+
   SmallVector<int64_t> filterkPos;
   for (auto reductionDim : reductionDims) {
     std::optional<int64_t> maybeDim = filterMap.getResultPosition(
         getAffineDimExpr(reductionDim, filterMap.getContext()));
     filterkPos.push_back(maybeDim.value());
   }
-  // group together adjacent reduction dimensions in the filter
+  // group together adjacent reduction dimensions in the filter.
+  // First we want to sort the dims as the look up from the filterMap
+  // can place the dims in arbitarty order.
+  std::sort(filterkPos.begin(), filterkPos.end());
   SmallVector<ReassociationIndices> collapsedFilterReductionDim;
   int64_t prevFilterIndex = filterkPos[0];
   int64_t currCollapsedIndex = 0;
@@ -526,12 +546,6 @@ getIGEMMGenericConvDetails(linalg::LinalgOp linalgOp) {
       SmallVector<AffineExpr>(dims.begin(), dims.begin() + numParallelDims),
       ctx);
 
-  bool isOutputChannelFirst = false;
-  auto outputChannelPos = convDims.outputChannel;
-  auto outputImagePos = convDims.outputImage;
-  if (outputChannelPos.back() < outputImagePos[0])
-    isOutputChannelFirst = true;
-
   // prepare the input map.
   SmallVector<AffineExpr> inputDims;
   // Add the batch dimensions.
@@ -579,7 +593,6 @@ getIGEMMGenericConvDetails(linalg::LinalgOp linalgOp) {
   SmallVector<int64_t> igemmLoopBounds;
   igemmLoopBounds.insert(igemmLoopBounds.end(), outputShape.begin(),
                          outputShape.begin() + numParallelDims);
-
   SmallVector<utils::IteratorType> igemmLoopIterators(outputShape.size(),
                                                       parallel);
 
@@ -594,7 +607,6 @@ getIGEMMGenericConvDetails(linalg::LinalgOp linalgOp) {
   igemmDetails.isOutputChannelFirst = isOutputChannelFirst;
   igemmDetails.convDims = convDims;
   igemmDetails.igemmLoopIterators = igemmLoopIterators;
-
   return igemmDetails;
 }