[mlir][linalg] Rename conv_2d_ngchw_fgchw to conv_2d_ngchw_gfchw

[kon72]

Summary

The name conv_2d_ngchw_fgchw suggests it takes a kernel in FGCHW format, but it actually expected the kernel to be in GFCHW format, and Torch-MLIR is using this op on the assumption that it takes GFCHW kernel.
https://github.com/llvm/torch-mlir/blob/0cf9ee340bc3264922e03b0df002e41c69184702/lib/Conversion/TorchToLinalg/Linear.cpp#L826

This change renames conv_2d_ngchw_fgchw to conv_2d_ngchw_gfchw and fixes the shape map.
Note that the behavior of this op remains the same as before.

Also adds a test against static tensors to make sure the data format is correct.

[kon72]

Afterthought, it seems this change was meant to be a part of D128880

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

value: !ScalarExpression
                 kind: type
                 fn_name: cast_signed
                 type_var: U
                 operands:
                 - !ScalarExpression
                   scalar_arg: K
 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
-  name: conv_2d_ngchw_fgchw
-  cpp_class_name: Conv2DNgchwFgchwOp
+  name: conv_2d_ngchw_gfchw
+  cpp_class_name: Conv2DNgchwGfchwOp
   doc: |-
     Performs 2-D grouped convolution.
 
     Layout:
       * Input: NGCHW.
-      * Kernel: FGCHW.
+      * Kernel: GFCHW.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output.
   implements:
   - LinalgConvolutionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
   - !LinalgOperandDefConfig
     name: I
     kind: input_tensor
     type_var: T1
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
       (s0, s1, s2, s3 * s4 + s5 * s6, s7 * s8 + s9 * s10)>
   - !LinalgOperandDefConfig
     name: K
     kind: input_tensor
     type_var: T2
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
-      (s11, s1, s2, s5, s9)>
+      (s1, s11, s2, s5, s9)>
   - !LinalgOperandDefConfig
     name: O
     kind: output_tensor
     type_var: U
     shape_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
-      (s0, s11, s1, s3, s7)>
+      (s0, s1, s11, s3, s7)>
   - !LinalgOperandDefConfig
     name: strides
     kind: index_attr
     index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11]
       -> (s4, s8)>
     default_indices:
     - 1
     - 1

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

def conv_2d_nchw_fchw(I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH,
   implements(ConvolutionOpInterface)
   domain(D.n, D.f, D.oh, D.ow, D.c, D.kh, D.kw)
   O[D.n, D.f, D.oh, D.ow] += TypeFn.cast_signed(
       U, I[D.n, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW +
            D.kw * S.DW]) * TypeFn.cast_signed(U, K[D.f, D.c, D.kh, D.kw])
 
 
 @linalg_structured_op
-def conv_2d_ngchw_fgchw(I=TensorDef(T1, S.N, S.G, S.C,
+def conv_2d_ngchw_gfchw(I=TensorDef(T1, S.N, S.G, S.C,
                                     S.OH * S.SH + S.KH * S.DH,
                                     S.OW * S.SW + S.KW * S.DW),
-                        K=TensorDef(T2, S.FG, S.G, S.C, S.KH, S.KW),
-                        O=TensorDef(U, S.N, S.FG, S.G, S.OH, S.OW, output=True),
+                        K=TensorDef(T2, S.G, S.FG, S.C, S.KH, S.KW),
+                        O=TensorDef(U, S.N, S.G, S.FG, S.OH, S.OW, output=True),
                         strides=IndexAttrDef(S.SH, S.SW, default=[1, 1]),
                         dilations=IndexAttrDef(S.DH, S.DW, default=[1, 1])):
   """Performs 2-D grouped convolution.
 
   Layout:
     * Input: NGCHW.
-    * Kernel: FGCHW.
+    * Kernel: GFCHW.
 
   Numeric casting is performed on the operands to the inner multiply, promoting
   them to the same data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.g, D.fg, D.oh, D.ow, D.c, D.kh, D.kw)
   O[D.n, D.g, D.fg, D.oh, D.ow] += TypeFn.cast_signed(
       U, I[D.n, D.g, D.c, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW +

mlir/test/Dialect/Linalg/named-ops.mlir

%0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
                                               strides = dense<1> : tensor<2xi64>}
      ins (%input, %filter: tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
     outs (%init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
   return %0 : tensor<?x?x?x?xf32>
 }
 
 // -----
 
-// CHECK-LABEL: func @conv_2d_ngchw_fgchw
-func.func @conv_2d_ngchw_fgchw(%input: tensor<?x?x?x?x?xf32>, %filter: tensor<?x?x?x?x?xf32>, %init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> {
-  // CHECK:      %{{.+}} = linalg.conv_2d_ngchw_fgchw
+// CHECK-LABEL: func @conv_2d_ngchw_gfchw
+func.func @conv_2d_ngchw_gfchw(%input: tensor<?x?x?x?x?xf32>, %filter: tensor<?x?x?x?x?xf32>, %init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32> {
+  // CHECK:      %{{.+}} = linalg.conv_2d_ngchw_gfchw
   // CHECK-SAME:   dilations = dense<1> : tensor<2xi64>
   // CHECK-SAME:   strides = dense<1> : tensor<2xi64>
   // CHECK-SAME:   ins(%{{.+}}, %{{.+}} : tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>)
   // CHECK-SAME:   outs(%{{.+}} : tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
-  %0 = linalg.conv_2d_ngchw_fgchw {dilations = dense<1> : tensor<2xi64>,
+  %0 = linalg.conv_2d_ngchw_gfchw {dilations = dense<1> : tensor<2xi64>,
                                               strides = dense<1> : tensor<2xi64>}
      ins (%input, %filter: tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>)
     outs (%init: tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
   return %0 : tensor<?x?x?x?x?xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @conv_2d_ngchw_gfchw_static
+func.func @conv_2d_ngchw_gfchw_static(%input: tensor<1x2x32x128x256xf32>, %filter: tensor<2x64x32x3x5xf32>, %init: tensor<1x2x64x126x252xf32>) -> tensor<1x2x64x126x252xf32> {
+  // CHECK:      %{{.+}} = linalg.conv_2d_ngchw_gfchw
+  // CHECK-SAME:   dilations = dense<1> : tensor<2xi64>
+  // CHECK-SAME:   strides = dense<1> : tensor<2xi64>
+  // CHECK-SAME:   ins(%{{.+}}, %{{.+}} : tensor<1x2x32x128x256xf32>, tensor<2x64x32x3x5xf32>)
+  // CHECK-SAME:   outs(%{{.+}} : tensor<1x2x64x126x252xf32>) -> tensor<1x2x64x126x252xf32>
+  %0 = linalg.conv_2d_ngchw_gfchw {dilations = dense<1> : tensor<2xi64>,
+                                              strides = dense<1> : tensor<2xi64>}
+     ins (%input, %filter: tensor<1x2x32x128x256xf32>, tensor<2x64x32x3x5xf32>)
+    outs (%init: tensor<1x2x64x126x252xf32>) -> tensor<1x2x64x126x252xf32>
+  return %0 : tensor<1x2x64x126x252xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @conv_2d_nhwc_fhwc
 func.func @conv_2d_nhwc_fhwc(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?x?x?xf32>, %init: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {
   // CHECK:      %{{.+}} = linalg.conv_2d_nhwc_fhwc
   // CHECK-SAME:   dilations = dense<1> : tensor<2xi64>
   // CHECK-SAME:   strides = dense<1> : tensor<2xi64>
   // CHECK-SAME:   ins(%{{.+}}, %{{.+}} : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>)
   // CHECK-SAME:   outs(%{{.+}} : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
   %0 = linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>,
linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
                                          strides = dense<1> : tensor<2xi64>}
      ins (%input, %filter: memref<?x?x?x?xf32>, memref<?x?x?x?xf32>)
     outs (%output: memref<?x?x?x?xf32>)
   return
 }
 
 // -----
 
-// CHECK-LABEL: func @conv_2d_ngchw_fgchw
-func.func @conv_2d_ngchw_fgchw(%input: memref<?x?x?x?x?xf32>, %filter: memref<?x?x?x?x?xf32>, %output: memref<?x?x?x?x?xf32>) {
-  // CHECK:      linalg.conv_2d_ngchw_fgchw
+// CHECK-LABEL: func @conv_2d_ngchw_gfchw
+func.func @conv_2d_ngchw_gfchw(%input: memref<?x?x?x?x?xf32>, %filter: memref<?x?x?x?x?xf32>, %output: memref<?x?x?x?x?xf32>) {
+  // CHECK:      linalg.conv_2d_ngchw_gfchw
   // CHECK-SAME:   dilations = dense<1> : tensor<2xi64>
   // CHECK-SAME:   strides = dense<1> : tensor<2xi64>
   // CHECK-SAME:   ins(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
   // CHECK-SAME:   outs(%{{.+}} : memref<?x?x?x?x?xf32>)
-  linalg.conv_2d_ngchw_fgchw {dilations = dense<1> : tensor<2xi64>,
+  linalg.conv_2d_ngchw_gfchw {dilations = dense<1> : tensor<2xi64>,
                                          strides = dense<1> : tensor<2xi64>}
      ins (%input, %filter: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
     outs (%output: memref<?x?x?x?x?xf32>)
   return
 }
 
 // -----