[Mlir-commits] [mlir] [mlir][linalg] Propagate filter tensor encoding in im2col (PR #160099)

[Andrzej Warzyński]

================
@@ -327,6 +327,34 @@ module attributes {transform.with_named_sequence} {
 
 // -----
 
+// CHECK: func.func @conv2d_decompose_im2col_with_filter_encoding
+// CHECK-SAME: %[[INPUT:.+]]: tensor<1x16x16x4xf32>,
+// CHECK-SAME: %[[FILTER:.*]]: tensor<16x3x3x4xf32, 42 : i32>,
+// CHECK-SAME: %[[OUTPUT:.*]]: tensor<1x14x14x16xf32>
+//  CHECK-DAG: %[[COLLAPSED_FILTER:.+]] = tensor.collapse_shape %[[FILTER]]
+  // CHECK-SAME{LITERAL}: [[0], [1, 2, 3]] : tensor<16x3x3x4xf32, 42 : i32> into tensor<16x36xf32, 42 : i32>
+//  CHECK-DAG: %[[COLLAPSED_OUT:.+]] = tensor.collapse_shape %[[OUTPUT]] {{\[}}[0], [1, 2], [3]] : tensor<1x14x14x16xf32> into tensor<1x196x16xf32>
+//  CHECK: %[[INIT_COL_TENSOR:.+]] = tensor.empty() : tensor<1x196x36xf32>
+//  CHECK: %[[COL_TENSOR:.+]] = linalg.generic {{.*}} ins(%[[INPUT]] : tensor<1x16x16x4xf32>) outs(%[[INIT_COL_TENSOR]] : tensor<1x196x36xf32>)
+//  CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic
+func.func @conv2d_decompose_im2col_with_filter_encoding(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<16x3x3x4xf32, 42 : i32>, %arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> {
+    %0 = linalg.conv_2d_nhwc_fhwc
+      { dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }
+      ins(%arg0, %arg1: tensor<1x16x16x4xf32>, tensor<16x3x3x4xf32, 42 : i32>)
+      outs(%arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32>
+    return %0 : tensor<1x14x14x16xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
+    %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op
+    %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
----------------
banach-space wrote:

IIUC, this is simply `@conv_2d_nhwc_fhwc` with encoding? If yes, then I suggest renaming 
* `@conv2d_decompose_im2col_with_filter_encoding` to `@conv_2d_nhwc_fhwc_with_encoding`

This way you will make the difference between the two very clear ("self-documenting"). I would also use more descriptive names, e.g. `%filter` instead of `%arg1` - it took me a while to find where that _filter_ with encoding was 😅 

Finally, if this is identical to `@conv_2d_nhwc_fhwc` then you only need to check the following:
* `COLLAPSED_FILTER` includes encoding (please add comments, it's very hard to see)
* There are no convolutions (checking for `linalg.generic` should be fine)
* `COLLAPSED_FILTER` is used by in matmul (not checked ATM)

Everything else can probably be stripped :) 

https://github.com/llvm/llvm-project/pull/160099