[mlir][linalg] Improve pooling op iterator order consistency

[antiagainst]

Summary

All named ops list iterators for accessing output first except
pooling ops. This commit made the pooling ops consistent with
the rest.

[antiagainst]

Rebase

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

structured_op: !LinalgStructuredOpConfig
   - !LinalgOperandDefConfig
     name: dilations
     usage: IndexAttribute
     type_var: I64
     attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4, s8)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+      -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d3, d4)>
+      -> (d4, d5)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1, d2, d5)>
+      -> (d0, d1, d2, d3)>
   iterator_types:
   - parallel
   - parallel
   - parallel
+  - parallel
   - reduction
   - reduction
-  - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         fn_name: add
         operands:
         - !ScalarExpression
structured_op: !LinalgStructuredOpConfig
   - !LinalgOperandDefConfig
     name: dilations
     usage: IndexAttribute
     type_var: I64
     attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4, s8)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+      -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d3, d4)>
+      -> (d4, d5)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1, d2, d5)>
+      -> (d0, d1, d2, d3)>
   iterator_types:
   - parallel
   - parallel
   - parallel
+  - parallel
   - reduction
   - reduction
-  - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         fn_name: max
         operands:
         - !ScalarExpression
structured_op: !LinalgStructuredOpConfig
   - !LinalgOperandDefConfig
     name: dilations
     usage: IndexAttribute
     type_var: I64
     attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4, s8)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+      -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d3, d4)>
+      -> (d4, d5)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1, d2, d5)>
+      -> (d0, d1, d2, d3)>
   iterator_types:
   - parallel
   - parallel
   - parallel
+  - parallel
   - reduction
   - reduction
-  - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         fn_name: max_unsigned
         operands:
         - !ScalarExpression
structured_op: !LinalgStructuredOpConfig
   - !LinalgOperandDefConfig
     name: dilations
     usage: IndexAttribute
     type_var: I64
     attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4, s8)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+      -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d3, d4)>
+      -> (d4, d5)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1, d2, d5)>
+      -> (d0, d1, d2, d3)>
   iterator_types:
   - parallel
   - parallel
   - parallel
+  - parallel
   - reduction
   - reduction
-  - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         fn_name: min
         operands:
         - !ScalarExpression
structured_op: !LinalgStructuredOpConfig
   - !LinalgOperandDefConfig
     name: dilations
     usage: IndexAttribute
     type_var: I64
     attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9] -> (s4, s8)>
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1 * s2 + d3 * s4, d2 * s6 + d4 * s8, d5)>
+      -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d3, d4)>
+      -> (d4, d5)>
     - affine_map<(d0, d1, d2, d3, d4, d5)[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9]
-      -> (d0, d1, d2, d5)>
+      -> (d0, d1, d2, d3)>
   iterator_types:
   - parallel
   - parallel
   - parallel
+  - parallel
   - reduction
   - reduction
-  - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         fn_name: min_unsigned
         operands:
         - !ScalarExpression
- !LinalgOperandDefConfig
     name: dilations
     usage: IndexAttribute
     type_var: I64
     attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13] -> (s4, s8, s12)>
   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, s12, s13] -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3
-      * s10 + d6 * s12, d7)>
+      s8, s9, s10, s11, s12, s13] -> (d0, d1 * s2 + d5 * s4, d2 * s6 + d6 * s8, d3
+      * s10 + d7 * s12, d4)>
     - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7,
-      s8, s9, s10, s11, s12, s13] -> (d4, d5, d6)>
+      s8, s9, s10, s11, s12, s13] -> (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, s12, s13] -> (d0, d1, d2, d3, d7)>
+      s8, s9, s10, s11, s12, s13] -> (d0, d1, d2, d3, d4)>
   iterator_types:
   - parallel
   - parallel
   - parallel
   - parallel
+  - parallel
   - reduction
   - reduction
   - reduction
-  - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         fn_name: add
         operands:
         - !ScalarExpression
- !LinalgOperandDefConfig
     name: dilations
     usage: IndexAttribute
     type_var: I64
     attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13] -> (s4, s8, s12)>
   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, s12, s13] -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3
-      * s10 + d6 * s12, d7)>
+      s8, s9, s10, s11, s12, s13] -> (d0, d1 * s2 + d5 * s4, d2 * s6 + d6 * s8, d3
+      * s10 + d7 * s12, d4)>
     - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7,
-      s8, s9, s10, s11, s12, s13] -> (d4, d5, d6)>
+      s8, s9, s10, s11, s12, s13] -> (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, s12, s13] -> (d0, d1, d2, d3, d7)>
+      s8, s9, s10, s11, s12, s13] -> (d0, d1, d2, d3, d4)>
   iterator_types:
   - parallel
   - parallel
   - parallel
   - parallel
+  - parallel
   - reduction
   - reduction
   - reduction
-  - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         fn_name: max
         operands:
         - !ScalarExpression
- !LinalgOperandDefConfig
     name: dilations
     usage: IndexAttribute
     type_var: I64
     attribute_map: affine_map<()[s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11,
       s12, s13] -> (s4, s8, s12)>
   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, s12, s13] -> (d0, d1 * s2 + d4 * s4, d2 * s6 + d5 * s8, d3
-      * s10 + d6 * s12, d7)>
+      s8, s9, s10, s11, s12, s13] -> (d0, d1 * s2 + d5 * s4, d2 * s6 + d6 * s8, d3
+      * s10 + d7 * s12, d4)>
     - affine_map<(d0, d1, d2, d3, d4, d5, d6, d7)[s0, s1, s2, s3, s4, s5, s6, s7,
-      s8, s9, s10, s11, s12, s13] -> (d4, d5, d6)>
+      s8, s9, s10, s11, s12, s13] -> (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, s12, s13] -> (d0, d1, d2, d3, d7)>
+      s8, s9, s10, s11, s12, s13] -> (d0, d1, d2, d3, d4)>
   iterator_types:
   - parallel
   - parallel
   - parallel
   - parallel
+  - parallel
   - reduction
   - reduction
   - reduction
-  - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
       arith_fn:
         fn_name: min
         operands:
         - !ScalarExpression

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

def pooling_nhwc_sum(
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs sum pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
-  domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+  domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
   O[D.n, D.oh, D.ow, D.c] += TypeFn.cast(
       U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c])
 
 
 @linalg_structured_op
 def pooling_nhwc_max(
     I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
                 S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
-  domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+  domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.max[D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_nhwc_max_unsigned(
     I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
                 S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs unsigned max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
-  domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+  domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.max_unsigned[D.kh, D.kw](
       TypeFn.cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_nchw_max(
     I=TensorDef(T1, S.N, S.C, S.OH * S.SH + S.KH * S.DH,
def pooling_nhwc_min(
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
-  domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+  domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.min[D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_nhwc_min_unsigned(
     I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH, S.OW * S.SW + S.KW * S.DW,
                 S.C),
     K=TensorDef(T2, S.KH, S.KW, index_dims=[D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SH, S.SW),
     dilations=IndexAttrDef(S.DH, S.DW)):
   """Performs unsigned min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
-  domain(D.n, D.oh, D.ow, D.kh, D.kw, D.c)
+  domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
   O[D.n, D.oh, D.ow, D.c] = ReduceFn.min_unsigned[D.kh, D.kw](
       TypeFn.cast_unsigned(
           U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_ndhwc_sum(
     I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
                 S.OW * S.SW + S.KW * S.DW, S.C),
     K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SD, S.SH, S.SW),
     dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
   """Performs 3D sum pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
-  domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
+  domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw)
   O[D.n, D.od, D.oh, D.ow, D.c] += TypeFn.cast(
       U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
            D.ow * S.SW + D.kw * S.DW, D.c])
 
 
 @linalg_structured_op
 def pooling_ndhwc_max(
     I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
                 S.OW * S.SW + S.KW * S.DW, S.C),
     K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SD, S.SH, S.SW),
     dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
   """Performs 3D max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
-  domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
+  domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw)
   O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.max[D.kd, D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
                D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def pooling_ndhwc_min(
     I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD, S.OH * S.SH + S.KH * S.DH,
                 S.OW * S.SW + S.KW * S.DW, S.C),
     K=TensorDef(T2, S.KD, S.KH, S.KW, index_dims=[D.kd, D.kh, D.kw]),
     O=TensorDef(U, S.N, S.OD, S.OH, S.OW, S.C, output=True),
     strides=IndexAttrDef(S.SD, S.SH, S.SW),
     dilations=IndexAttrDef(S.DD, S.DH, S.DW)):
   """Performs 3D min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
-  domain(D.n, D.od, D.oh, D.ow, D.kd, D.kh, D.kw, D.c)
+  domain(D.n, D.od, D.oh, D.ow, D.c, D.kd, D.kh, D.kw)
   O[D.n, D.od, D.oh, D.ow, D.c] = ReduceFn.min[D.kd, D.kh, D.kw](
       TypeFn.cast(
           U, I[D.n, D.od * S.SD + D.kd * S.DD, D.oh * S.SH + D.kh * S.DH,
                D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
 def fill_rng_2d(