[mlir][linalg] Add pooling ops to Linalg TC ops.

[hanchung]

Summary

• Add EDSC builders for std_cmpf_ogt and std_cmpf_olt.
• Add pooling_nhwc_min/max/sum ops

Depends On D97384

[hanchung]

remove -canonicalize pass from the test

[hanchung]

rebase

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOpsSpec.tc

 """
 {
   O(n, f, d, h, w) = std_addf<kd, kh, kw>(
       O(n, f, d, h, w), std_mulf(I(n, c, d * strides[0] + kd * dilations[0],
                                    h * strides[1] + kh * dilations[1],
                                    w * strides[2] + kw * dilations[2]),
                                  K(kd, kh, kw, c, f)));
 }
+
+ods_def<PoolingNHWCSum>:
+def pooling_nhwc_sum
+    (I: f32(N, H, W, C), K: f32(KH, KW))
+  -> (O: f32(N, OH, OW, C))
+  attr(strides: 2xi64, dilations: 2xi64)
+{
+  O(n, oh, ow, c) = std_addf<kh, kw>(O(n, oh, ow, c),
+                                     I(n, oh * strides[0] + kh * dilations[0],
+                                       ow * strides[1] + kw * dilations[1], c));
+}
+
+ods_def<PoolingNHWCMax>:
+def pooling_nhwc_max
+    (I: f32(N, H, W, C), K: f32(KH, KW))
+  -> (O: f32(N, OH, OW, C))
+  attr(strides: 2xi64, dilations: 2xi64)
+{
+  O(n, oh, ow, c) =
+      std_select<kh, kw>(std_cmpf_ogt(I(n, oh * strides[0] + kh * dilations[0],
+                                        ow * strides[1] + kw * dilations[1], c),
+                                      O(n, oh, ow, c)),
+                         I(n, oh * strides[0] + kh * dilations[0],
+                           ow * strides[1] + kw * dilations[1], c),
+                         O(n, oh, ow, c));
+}
+
+ods_def<PoolingNHWCMin>:
+def pooling_nhwc_min
+    (I: f32(N, H, W, C), K: f32(KH, KW))
+  -> (O: f32(N, OH, OW, C))
+  attr(strides: 2xi64, dilations: 2xi64)
+{
+  O(n, oh, ow, c) =
+      std_select<kh, kw>(std_cmpf_olt(I(n, oh * strides[0] + kh * dilations[0],
+                                        ow * strides[1] + kw * dilations[1], c),
+                                      O(n, oh, ow, c)),
+                         I(n, oh * strides[0] + kh * dilations[0],
+                           ow * strides[1] + kw * dilations[1], c),
+                         O(n, oh, ow, c));
 }

mlir/include/mlir/Dialect/StandardOps/EDSC/Intrinsics.h

struct SExtiValueBuilder : public ValueBuilder<SignExtendIOp> {
   template <typename... Args>
   SExtiValueBuilder(Args... args)
       : ValueBuilder<SignExtendIOp>(ScopedContext::getBuilderRef().getI32Type(),
                                     args...) {}
 };
 
 using std_sexti32 = SExtiValueBuilder<32>;
 
+template <CmpFPredicate Predicate>
+struct CmpFValueBuilder : public ValueBuilder<CmpFOp> {
+  using ValueBuilder<CmpFOp>::ValueBuilder;
+  template <typename... Args>
+  CmpFValueBuilder(Args... args) : ValueBuilder<CmpFOp>(Predicate, args...) {}
+};
+
+using std_cmpf_ogt = CmpFValueBuilder<CmpFPredicate::OGT>;
+using std_cmpf_olt = CmpFValueBuilder<CmpFPredicate::OLT>;
+
 /// Branches into `block` with `operands`.
 BranchOp std_br(Block *block, ValueRange operands);
 
 /// Branches into `trueBranch` with `trueOperands` if `cond` evaluates to `true`
 /// or to `falseBranch` and `falseOperand` if `cond` evaluates to `false`.
 CondBranchOp std_cond_br(Value cond, Block *trueBranch, ValueRange trueOperands,
                          Block *falseBranch, ValueRange falseOperands);
 
 /// Provide an index notation around sdt_load and std_store.
 using StdIndexedValue =
     TemplatedIndexedValue<intrinsics::std_load, intrinsics::std_store>;
 } // namespace intrinsics
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_STANDARDOPS_EDSC_INTRINSICS_H_

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

-// RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -linalg-generalize-named-ops -canonicalize | FileCheck %s
 
 func @generalize_conv(%input : memref<1x225x225x3xf32>, %filter: memref<3x3x3x32xf32>, %output: memref<1x112x112x32xf32>) {
   linalg.conv(%filter, %input, %output) {dilations = [2, 3], strides = [4, 5]} : memref<3x3x3x32xf32>, memref<1x225x225x3xf32>, memref<1x112x112x32xf32>
   return
 }
 
 // CHECK: #[[FILTER_MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d5, d6, d4, d3)>
 // CHECK:  #[[INPUT_MAP:.+]] = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 * 4 + d5 * 2, d2 * 5 + d6 * 3, d4)>

 // CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction", "parallel"]}
 // CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
 // CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?x?xf32>)
 
 // CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
 // CHECK-NEXT:      %[[MUL:.+]] = mulf %[[BBARG0]], %[[BBARG1]] : f32
 // CHECK-NEXT:      %[[ADD:.+]] = addf %[[BBARG2]], %[[MUL]] : f32
 // CHECK-NEXT:      linalg.yield %[[ADD]] : f32
+
+// -----
+
+func @pooling_nhwc_sum(%input: memref<?x?x?x?xf32>, %fake: memref<2x3xf32>, %init: memref<?x?x?x?xf32>) {
+  linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: memref<?x?x?x?xf32>, memref<2x3xf32>)
+    outs(%init: memref<?x?x?x?xf32>)
+  return
+}
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 + d4, d2 + d5, d3)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
+
+// CHECK: func @pooling_nhwc_sum
+
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
+// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
+
+// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
+// CHECK-NEXT:      %[[RES:.+]] = addf %[[BBARG2]], %[[BBARG0]] : f32
+// CHECK-NEXT:      linalg.yield %[[RES]] : f32
+
+// -----
+
+func @pooling_nhwc_max(%input: memref<?x?x?x?xf32>, %fake: memref<2x3xf32>, %init: memref<?x?x?x?xf32>) {
+  linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64>}
+    ins(%input, %fake: memref<?x?x?x?xf32>, memref<2x3xf32>)
+    outs(%init: memref<?x?x?x?xf32>)
+  return
+}
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4, d2 * 3 + d5, d3)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
+
+// CHECK: func @pooling_nhwc_max
+
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
+// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
+
+// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
+// CHECK-NEXT:      %[[CMP:.+]] = cmpf ogt, %[[BBARG0]], %[[BBARG2]] : f32
+// CHECK-NEXT:      %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : f32
+// CHECK-NEXT:      linalg.yield %[[RES]] : f32
+
+// -----
+
+func @pooling_nhwc_min(%input: memref<?x?x?x?xf32>, %fake: memref<2x3xf32>, %init: memref<?x?x?x?xf32>) {
+  linalg.pooling_nhwc_min {dilations = dense<3> : tensor<2xi64>, strides = dense<2> : tensor<2xi64>}
+    ins(%input, %fake: memref<?x?x?x?xf32>, memref<2x3xf32>)
+    outs(%init: memref<?x?x?x?xf32>)
+  return
+}
+
+// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1 * 2 + d4 * 3, d2 * 2 + d5 * 3, d3)>
+// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d4, d5)>
+// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3)>
+
+// CHECK: func @pooling_nhwc_min
+
+// CHECK: linalg.generic
+// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]}
+// CHECK-SAME: ins(%{{.+}}, %{{.+}} : memref<?x?x?x?xf32>, memref<2x3xf32>)
+// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xf32>)
+
+// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: f32, %[[BBARG1:.+]]: f32, %[[BBARG2:.+]]: f32)
+// CHECK-NEXT:      %[[CMP:.+]] = cmpf olt, %[[BBARG0]], %[[BBARG2]] : f32
+// CHECK-NEXT:      %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : f32
+// CHECK-NEXT:      linalg.yield %[[RES]] : f32

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

func @conv_3d_input_ncdhw_filter_dhwcf(%input: memref<?x?x?x?x?xf32>, %filter: memref<?x?x?x?x?xf32>, %output: memref<?x?x?x?x?xf32>) {
   // 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_3d_input_ncdhw_filter_dhwcf {dilations = dense<1> : tensor<3xi64>,
                                            strides = dense<1> : tensor<3xi64>}
      ins (%input, %filter: memref<?x?x?x?x?xf32>, memref<?x?x?x?x?xf32>)
     outs (%output: memref<?x?x?x?x?xf32>)
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_sum_tensor
+// CHECK:         %{{.+}} = linalg.pooling_nhwc_sum
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : tensor<1x6x6x1xf32>, tensor<3x3xf32>)
+// CHECK-SAME:      outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
+func @pooling_nhwc_sum_tensor(%input: tensor<1x6x6x1xf32>) -> tensor<1x2x2x1xf32> {
+  %fake = linalg.init_tensor [3, 3] : tensor<3x3xf32>
+  %init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xf32>
+  %cst = constant 0.000000e+00 : f32
+  %fill = linalg.fill(%init, %cst) : tensor<1x2x2x1xf32>, f32 -> tensor<1x2x2x1xf32>
+  %res = linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: tensor<1x6x6x1xf32>, tensor<3x3xf32>)
+    outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
+  return %res : tensor<1x2x2x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_sum
+// CHECK:         linalg.pooling_nhwc_sum
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : memref<1x6x6x1xf32>, memref<3x3xf32>)
+// CHECK-SAME:      outs(%{{.+}} : memref<1x2x2x1xf32>)
+func @pooling_nhwc_sum(%input: memref<1x6x6x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
+  linalg.pooling_nhwc_sum {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: memref<1x6x6x1xf32>, memref<3x3xf32>)
+    outs(%output: memref<1x2x2x1xf32>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_max_tensor
+// CHECK:         %{{.+}} = linalg.pooling_nhwc_max
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : tensor<1x6x6x1xf32>, tensor<3x3xf32>)
+// CHECK-SAME:      outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
+func @pooling_nhwc_max_tensor(%input: tensor<1x6x6x1xf32>) -> tensor<1x2x2x1xf32> {
+  %fake = linalg.init_tensor [3, 3] : tensor<3x3xf32>
+  %init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xf32>
+  %cst = constant 0.000000e+00 : f32
+  %fill = linalg.fill(%init, %cst) : tensor<1x2x2x1xf32>, f32 -> tensor<1x2x2x1xf32>
+  %res = linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: tensor<1x6x6x1xf32>, tensor<3x3xf32>)
+    outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
+  return %res : tensor<1x2x2x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_max
+// CHECK:         linalg.pooling_nhwc_max
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : memref<1x6x6x1xf32>, memref<3x3xf32>)
+// CHECK-SAME:      outs(%{{.+}} : memref<1x2x2x1xf32>)
+func @pooling_nhwc_max(%input: memref<1x6x6x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
+  linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: memref<1x6x6x1xf32>, memref<3x3xf32>)
+    outs(%output: memref<1x2x2x1xf32>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_min_tensor
+// CHECK:         %{{.+}} = linalg.pooling_nhwc_min
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : tensor<1x6x6x1xf32>, tensor<3x3xf32>)
+// CHECK-SAME:      outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
+func @pooling_nhwc_min_tensor(%input: tensor<1x6x6x1xf32>) -> tensor<1x2x2x1xf32> {
+  %fake = linalg.init_tensor [3, 3] : tensor<3x3xf32>
+  %init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xf32>
+  %cst = constant 0.000000e+00 : f32
+  %fill = linalg.fill(%init, %cst) : tensor<1x2x2x1xf32>, f32 -> tensor<1x2x2x1xf32>
+  %res = linalg.pooling_nhwc_min {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: tensor<1x6x6x1xf32>, tensor<3x3xf32>)
+    outs(%fill: tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
+  return %res : tensor<1x2x2x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_min
+// CHECK:         linalg.pooling_nhwc_min
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : memref<1x6x6x1xf32>, memref<3x3xf32>)
+// CHECK-SAME:      outs(%{{.+}} : memref<1x2x2x1xf32>)
+func @pooling_nhwc_min(%input: memref<1x6x6x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
+  linalg.pooling_nhwc_min {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: memref<1x6x6x1xf32>, memref<3x3xf32>)
+    outs(%output: memref<1x2x2x1xf32>)
+  return
+}