[GPU] Add support for conv padding when decomposing im2col (#20203)

Adds support for fused pad producer to linalg conv ops going down the
IGEMM path. This is done by swapping the pad with the extract slice of
the decomposed im2col op.
See this issue for details on flags needed to fuse the pad with the conv
in a single dispatch.
#20200

---------

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

Author: nirvedhmeshram

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/DecomposeIm2col.cpp

@@ -89,6 +89,10 @@ void DecomposeIm2colPass::runOnOperation() {
 
   RewritePatternSet patterns(context);
   memref::populateResolveRankedShapedTypeResultDimsPatterns(patterns);
+  // After im2col is decomposed, im2col extract slice can be swapped with input
+  // padding.
+  patterns.insert<linalg::ExtractSliceOfPadTensorSwapPattern>(
+      context, [](tensor::ExtractSliceOp) { return false; });
   if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
     return signalPassFailure();
   }

compiler/src/iree/compiler/Dialect/LinalgExt/Transforms/test/decompose_im2col.mlir

@@ -233,3 +233,27 @@ module {
 //       CHECK-UNROLL:   %[[INSERT3:.+]] = tensor.insert_slice %[[COPY]] into %[[INSERT2]][%[[C1]], %[[C1]], 0] [1, 1, 4] [1, 1, 1] : tensor<1x1x4xf32> into tensor<2x2x4xf32>
 
 //       CHECK-UNROLL:   return %[[INSERT3]] : tensor<2x2x4xf32>
+
+// -----
+
+module {
+  func.func @im2col_padding(%input: tensor<1x8x3x3xf32>) -> tensor<1x2x2x12xf32> {
+    %cst = arith.constant 0.000000e+00 : f32
+    %empty = tensor.empty() : tensor<1x2x2x12xf32>
+    %padded = tensor.pad %input low[0, 0, 3, 3] high[0, 0, 3, 3] {
+  ^bb0(%arg0: index, %arg1: index, %arg2: index, %arg3: index):
+    tensor.yield %cst : f32
+  } : tensor<1x8x3x3xf32> to tensor<1x8x9x9xf32>
+  %im2col = iree_linalg_ext.im2col strides = [1, 1] dilations = [1, 1] kernel_size = [3, 3]
+                              m_offset = [0, 0] * [2, 1] k_offset = [0] * [1]
+                              batch_pos = [0] m_pos = [2, 3] k_pos = [1]
+                              ins(%padded : tensor<1x8x9x9xf32>)
+                              outs(%empty : tensor<1x2x2x12xf32>) -> tensor<1x2x2x12xf32>
+  return %im2col : tensor<1x2x2x12xf32>
+  }
+}
+
+// CHECK-LABEL: func.func @im2col_padding
+//  CHECK-SAME:     %[[ARG0:[a-zA-Z0-9_]+]]
+//       CHECK: %[[T1:.+]] = tensor.extract_slice %[[ARG0]]
+//       CHECK: %[[T2:.+]] = tensor.pad %[[T1]]