Adding a named op for grouped convolutions

[gpetters94]

Summary

This is very much a WIP but I'd like to get feedback on if this is the correct way to go about implementing this. As per Sean Silva's advice, this uses the group size as a second batch dimension in both input and weight.

[silvas]

I suspect the g dimension should be contiguous with the C dimension, and I'm not sure if it should be major or minor to it.

[gpetters94]

I think it makes more sense for it to be major, if only for legibility's sake.

[mravishankar]

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

364

mods and divs in indexing map is going to be problematic..... I'd expect input shape to be [NG][N][CG][...][..] where NG is assumed to be C/G and CG is assumed to be c%G. We cannot enforce those conditions, they are expected to be true....

[silvas]

In D127225#3565062, @gpetters94 wrote:

I think it makes more sense for it to be major, if only for legibility's sake.

It's not really something we get to choose -- the op definition in frontends determines whether the groups are major or minor to the "real channels" -- I think it is major but please verify that. If we don't get this right then we have to transpose before calling this op which is wasteful.

[gpetters94]

In D127225#3567554, @silvas wrote:

In D127225#3565062, @gpetters94 wrote:

I think it makes more sense for it to be major, if only for legibility's sake.

It's not really something we get to choose -- the op definition in frontends determines whether the groups are major or minor to the "real channels" -- I think it is major but please verify that. If we don't get this right then we have to transpose before calling this op which is wasteful.

It looks like PyTorch in the slow CPU backend implements this with batches of convolutions rather than a batch dimension. After some digging through mkldnn, it looks like group is major there, although that's opaque to the user. So for PyTorch at least it looks like [N, G, C, ...] is the way to go.

[gpetters94]

Addressed comments.

[gpetters94]

Updated indexing maps.

[nicolasvasilache]

Thanks @gpetters94 !

[gpetters94]

Finally finished the bugfixing, with a lot of help from Mahesh. Should be ready to be merged unless anyone has any more requests.

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

                 - !ScalarExpression
                   scalar_arg: K
 --- !LinalgOpConfig
+metadata: !LinalgOpMetadata
+  name: conv_2d_ngchw_fgchw
+  cpp_class_name: Conv2DNgchwFgchwOp
+  doc: |-
+    Performs 2-D convolution.
+
+    Layout:
+      * Input: NGCHW.
+      * Kernel: FGCHW.
+
+    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)>
+  - !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] -> (s10,
+      s1, s11, s4, s8)>
+  - !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,
+      s1, s10, s2, s6)>
+  - !LinalgOperandDefConfig
+    name: strides
+    kind: index_attr
+    index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
+      (s3, s7)>
+    default_indices:
+    - 1
+    - 1
+  - !LinalgOperandDefConfig
+    name: dilations
+    kind: index_attr
+    index_attr_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11] ->
+      (s5, s9)>
+    default_indices:
+    - 1
+    - 1
+  indexing_maps: !LinalgIndexingMapsConfig
+    static_indexing_maps:
+    - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
+      s9, s10, s11] -> (d0, d1, d5, d3 * s3 + d6 * s5, d4 * s7 + d7 * s9)>
+    - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
+      s9, s10, s11] -> (d2, d1, d5, d6, d7)>
+    - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7, s8,
+      s9, s10, s11] -> (d0, d1, d2, d3, d4)>
+  iterator_types:
+  - parallel
+  - parallel
+  - parallel
+  - parallel
+  - parallel
+  - reduction
+  - reduction
+  - reduction
+  assignments:
+  - !ScalarAssign
+    arg: O
+    value: !ScalarExpression
+      scalar_fn:
+        kind: binary
+        fn_name: add
+        operands:
+        - !ScalarExpression
+          scalar_arg: O
+        - !ScalarExpression
+          scalar_fn:
+            kind: binary
+            fn_name: mul
+            operands:
+            - !ScalarExpression
+              scalar_fn:
+                kind: type
+                fn_name: cast_signed
+                type_var: U
+                operands:
+                - !ScalarExpression
+                  scalar_arg: I
+            - !ScalarExpression
+              scalar_fn:
+                kind: type
+                fn_name: cast_signed
+                type_var: U
+                operands:
+                - !ScalarExpression
+                  scalar_arg: K
+--- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: conv_3d_ndhwc_dhwcf
   cpp_class_name: Conv3DNdhwcDhwcfOp

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

       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, 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.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.
+
+  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 +
+          D.kw * S.DW]) * TypeFn.cast_signed(U, K[D.fg, D.g, D.c, D.kh, D.kw])
 
 @linalg_structured_op
 def conv_3d_ndhwc_dhwcf(I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD,