[mlir][tosa] Add tosa.max_pool2d lowering to linalg int max pooling additions

[rsuderman]

Summary

Lowerings tosa.max_pool2d to linalg equivalent operations. Includes
adding max pooling operations for linalg, with corresponding tests.

[rsuderman]

Fixed clang-tidy error

[mravishankar]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

1466

If pad[0] and pad[1] etc. are integers, you could just create an IntegerAttr. So

OpFoldResult ... = rewriter.createI64IntegerAttr(pad[1])

Also, I dont think createOrFold does anything for ConstantIndexOp . Could just use create (i have made the same mistake previously, so stating comments from then)

[rsuderman]

Remove createOrFold uses

[rsuderman]

Fixed build error due to OpFoldResult.

[mravishankar]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

1475

I am not sure I follow this (sorry if my previous comment was not clear). The RHS values are all constants right?

So you could just create an IntegerAttr for these statically known values. Basically there is no need for the ConstantIndexOp used to create any of these values AFAICS.

(Sorry I had commented earlier, but forgot to hit enter)

[rsuderman]

Updated to use constant attributes instead of values for padding behavior.

[mravishankar]

Looks fine to me except for use of C++-14 feature which might cause build issues. So LGTM apart from that issue. Please make sure no bots are broken

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

1490

This is a C++-14 feature. I am not sure what compilers/level of c++14 support the buildbots use. I try to stick to C++11?

[rriddle]

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

1490

LLVM is C++14.

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

def pooling_nhwc_sum
   -> (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<PoolingNHWCMaxI8Op>:
+def pooling_nhwc_i8_max
+    (I: i8(N, H, W, C), K: i8(KH, KW))
+  -> (O: i8(N, OH, OW, C))
+  attr(strides: 2xi64, dilations: 2xi64)
+{
+  O(n, oh, ow, c) =
+      std_select<kh, kw>(std_cmpi_sgt(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<PoolingNHWCMaxI16Op>:
+def pooling_nhwc_i16_max
+    (I: i16(N, H, W, C), K: i16(KH, KW))
+  -> (O: i16(N, OH, OW, C))
+  attr(strides: 2xi64, dilations: 2xi64)
+{
+  O(n, oh, ow, c) =
+      std_select<kh, kw>(std_cmpi_sgt(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<PoolingNHWCMaxI32Op>:
+def pooling_nhwc_i32_max
+    (I: i32(N, H, W, C), K: i32(KH, KW))
+  -> (O: i32(N, OH, OW, C))
+  attr(strides: 2xi64, dilations: 2xi64)
+{
+  O(n, oh, ow, c) =
+      std_select<kh, kw>(std_cmpi_sgt(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<PoolingNHWCMaxFOp>:
 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],

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

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>;
 
+template <CmpIPredicate Predicate>
+struct CmpIValueBuilder : public ValueBuilder<CmpIOp> {
+  using ValueBuilder<CmpIOp>::ValueBuilder;
+  template <typename... Args>
+  CmpIValueBuilder(Args... args) : ValueBuilder<CmpIOp>(Predicate, args...) {}
+};
+
+using std_cmpi_sgt = CmpIValueBuilder<CmpIPredicate::sgt>;
+
 /// 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);
 } // namespace intrinsics
 } // namespace edsc
 } // namespace mlir
 
 #endif // MLIR_DIALECT_STANDARDOPS_EDSC_INTRINSICS_H_

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

LogicalResult matchAndRewrite(tosa::PadOp padOp,
     int64_t rank = inputTy.getRank();
 
     if (!inputTy.hasStaticShape() || !paddingTy.hasStaticShape()) {
       return rewriter.notifyMatchFailure(
           padOp,
           "Pad converter requires static shaped input / padding values.");
     }
 
+    Attribute constantAttr;
+    if (elementTy.isa<FloatType>())
+      constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
+    else if (elementTy.isa<IntegerType>() && !padOp.quantization_info())
+      constantAttr = rewriter.getIntegerAttr(elementTy, 0);
+    else if (elementTy.isa<IntegerType>() && padOp.quantization_info()) {
+      auto value = padOp.quantization_info().getValue().input_zp().getValue();
+      constantAttr = rewriter.getIntegerAttr(elementTy, value.getZExtValue());
+    }
+
+    if (!constantAttr) {
+      return rewriter.notifyMatchFailure(
+          padOp,
+          "tosa.pad to linalg lowering encountered an unknown element type");
+    }
+
     Value lowIndex = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(0));
     Value highIndex =
         rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(1));
 
     SmallVector<OpFoldResult, 3> lowValues;
     SmallVector<OpFoldResult, 3> highValues;
 
     lowValues.reserve(rank);
for (int i = 0; i < rank; i++) {
                                                   lowVal);
       highVal = rewriter.createOrFold<IndexCastOp>(loc, rewriter.getIndexType(),
                                                    highVal);
 
       lowValues.push_back(lowVal);
       highValues.push_back(highVal);
     }
 
-    Attribute constantAttr;
-    if (elementTy.isa<FloatType>())
-      constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
-    else if (elementTy.isa<IntegerType>() && !padOp.quantization_info())
-      constantAttr = rewriter.getIntegerAttr(elementTy, 0);
-    else if (elementTy.isa<IntegerType>() && padOp.quantization_info()) {
-      auto value = padOp.quantization_info().getValue().input_zp().getValue();
-      constantAttr = rewriter.getIntegerAttr(elementTy, value.getZExtValue());
-    }
-
-    if (!constantAttr) {
-      return rewriter.notifyMatchFailure(
-          padOp,
-          "tosa.pad to linalg lowering encountered an unknown element type");
-    }
-
     Value constant = rewriter.create<ConstantOp>(loc, constantAttr);
 
     auto newPadOp = linalg::PadTensorOp::createPadScalarOp(
         padOp.getType(), input, constant, lowValues, highValues, loc, rewriter);
 
     rewriter.replaceOp(padOp, newPadOp.getResult());
     return success();
   }
rewriter.replaceOp(op, genericOp.getResult(0));
           resultElementTy.isInteger(8)) {
         Value index = rewriter.create<IndexCastOp>(loc, rewriter.getIndexType(),
                                                    inputValue);
         Value extract =
             rewriter.create<tensor::ExtractOp>(loc, table, ValueRange{index});
         rewriter.create<linalg::YieldOp>(loc, extract);
         return success();
       }
 
       if (inputElementTy.isInteger(16) && tableElementTy.isInteger(16) &&
           resultElementTy.isInteger(32)) {
         Value extend = rewriter.create<SignExtendIOp>(
             loc, rewriter.getI32Type(), inputValue);
 
         auto offset =
             rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(32768));
         auto seven =
             rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(7));
         auto one =
             rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(1));
         auto b1111111 =
             rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(127));
 
         // Compute the index and fractional part from the input value:
         // value = value + 32768
         // index = value >> 7;
         // fraction = 0x01111111 & value
         auto extendAdd = rewriter.create<AddIOp>(loc, extend, offset);
         Value index =
             rewriter.create<UnsignedShiftRightOp>(loc, extendAdd, seven);
         Value fraction = rewriter.create<mlir::AndOp>(loc, extendAdd, b1111111);
 
         // Extract the base and next values from the table.
         // base = (int32_t) table[index];
         // next = (int32_t) table[index + 1];
         Value indexPlusOne = rewriter.create<AddIOp>(loc, index, one);
 
         index =
             rewriter.create<IndexCastOp>(loc, rewriter.getIndexType(), index);
         indexPlusOne = rewriter.create<IndexCastOp>(
             loc, rewriter.getIndexType(), indexPlusOne);
rewriter.replaceOp(op, genericOp.getResult(0));
       }
     }
 
     return rewriter.notifyMatchFailure(
         op, "unable to create body for tosa.table op");
   }
 };
 
+class MaxPool2dConverter : public OpRewritePattern<tosa::MaxPool2dOp> {
+public:
+  using OpRewritePattern<tosa::MaxPool2dOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tosa::MaxPool2dOp op,
+                                PatternRewriter &rewriter) const final {
+    Location loc = op.getLoc();
+    Value input = op.input();
+    ShapedType inputTy = input.getType().cast<ShapedType>();
+    Type inElementTy = inputTy.getElementType();
+
+    ShapedType resultTy = op.getType().cast<ShapedType>();
+    Type outElementTy = inputTy.getElementType();
+    int64_t rank = inputTy.getRank();
+
+    if (!inputTy.hasStaticShape())
+      return failure();
+
+    // Determine what the initial value needs to be for the max pool op.
+    Attribute initialAttr;
+    if (outElementTy.isF32())
+      initialAttr = rewriter.getFloatAttr(
+          outElementTy,
+          APFloat::getLargest(
+              outElementTy.cast<FloatType>().getFloatSemantics(), true));
+
+    if (outElementTy.isa<IntegerType>())
+      initialAttr = rewriter.getIntegerAttr(
+          outElementTy,
+          APInt::getSignedMinValue(outElementTy.getIntOrFloatBitWidth()));
+
+    if (!initialAttr)
+      return rewriter.notifyMatchFailure(
+          op, "Unsupported initial value for tosa.maxpool_2d op");
+
+    Value initialValue = rewriter.create<ConstantOp>(loc, initialAttr);
+
+    SmallVector<int64_t> kernel, stride, pad;
+    getValuesFromIntArrayAttribute(op.kernel(), kernel);
+    getValuesFromIntArrayAttribute(op.stride(), stride);
+    getValuesFromIntArrayAttribute(op.pad(), pad);
+
+    Attribute strideAttr = rewriter.getI64VectorAttr(stride);
+    Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
+
+    // If non-zero padding we need to pad the input
+    if (llvm::any_of(pad, [](int64_t v) { return v != 0; })) {
+      SmallVector<int64_t, 4> paddedShape;
+      for (int64_t i = 0; i < rank; i++)
+        paddedShape.push_back(inputTy.getDimSize(i));
+
+      paddedShape[1] += pad[0] + pad[1];
+      paddedShape[2] += pad[2] + pad[3];
+
+      OpFoldResult zeroIndex = rewriter.getIndexAttr(0);
+      OpFoldResult heightLowPadIndex = rewriter.getIndexAttr(pad[0]);
+      OpFoldResult heightHighPadIndex = rewriter.getIndexAttr(pad[1]);
+      OpFoldResult widthLowPadIndex = rewriter.getIndexAttr(pad[2]);
+      OpFoldResult widthHighPadIndex = rewriter.getIndexAttr(pad[3]);
+
+      SmallVector<OpFoldResult, 4> lowIndices = {zeroIndex, heightLowPadIndex,
+                                                 widthLowPadIndex, zeroIndex};
+      SmallVector<OpFoldResult, 4> highIndices = {zeroIndex, heightHighPadIndex,
+                                                  widthHighPadIndex, zeroIndex};
+
+      input = linalg::PadTensorOp::createPadScalarOp(
+                  RankedTensorType::get(paddedShape, inElementTy), input,
+                  initialValue, lowIndices, highIndices, loc, rewriter)
+                  .result();
+    }
+
+    Value initTensor = rewriter.create<linalg::InitTensorOp>(
+        loc, resultTy.getShape(), resultTy.getElementType());
+
+    Value filledInitTensor =
+        rewriter.create<linalg::FillOp>(loc, initTensor, initialValue).result();
+
+    Value fakeWindowDims =
+        rewriter.create<linalg::InitTensorOp>(loc, kernel, outElementTy);
+
+    auto createOp = [&](auto *typePtr) -> linalg::LinalgOp {
+      return cast<linalg::LinalgOp>(
+          rewriter
+              .create<std::remove_pointer_t<decltype(typePtr)>>(
+                  loc, ArrayRef<Type>{resultTy},
+                  ValueRange{input, fakeWindowDims}, filledInitTensor,
+                  dilationAttr, strideAttr)
+              .getOperation());
+    };
+
+    if (inElementTy.isF32()) {
+      linalg::LinalgOp poolingOp =
+          createOp(static_cast<linalg::PoolingNHWCMaxFOp *>(nullptr));
+      rewriter.replaceOp(op, poolingOp->getResult(0));
+      return success();
+    }
+
+    if (inElementTy.isInteger(8)) {
+      linalg::LinalgOp poolingOp =
+          createOp(static_cast<linalg::PoolingNHWCMaxI8Op *>(nullptr));
+      rewriter.replaceOp(op, poolingOp->getResult(0));
+      return success();
+    }
+
+    if (inElementTy.isInteger(16)) {
+      linalg::LinalgOp poolingOp =
+          createOp(static_cast<linalg::PoolingNHWCMaxI16Op *>(nullptr));
+      rewriter.replaceOp(op, poolingOp->getResult(0));
+      return success();
+    }
+
+    if (inElementTy.isInteger(32)) {
+      linalg::LinalgOp poolingOp =
+          createOp(static_cast<linalg::PoolingNHWCMaxI32Op *>(nullptr));
+      rewriter.replaceOp(op, poolingOp->getResult(0));
+      return success();
+    }
+
+    return failure();
+  }
+};
+
 } // namespace
 
 void mlir::tosa::populateTosaToLinalgOnTensorsConversionPatterns(
     RewritePatternSet *patterns) {
   patterns->add<
       // clang-format off
       PointwiseConverter<tosa::AddOp>,
       PointwiseConverter<tosa::SubOp>,
patterns->add<
       PadConverter,
       ReshapeConverter,
       RescaleConverter,
       ReverseConverter,
       TableConverter,
       TileConverter,
       TransposeConverter,
       MatMulConverter,
+      MaxPool2dConverter,
       FullyConnectedConverter>(patterns->getContext());
       // clang-format on
 }

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir

func @table16(%arg0: tensor<6xi16>, %arg1: tensor<513xi16>) -> () {
   // CHECK: %[[BASE_MUL:.+]] = shift_left %[[BASE_EXT]], %[[C7]]
   // CHECK: %[[DIFF:.+]] = subi %[[NEXT_EXT]], %[[BASE_EXT]]
   // CHECK: %[[DIFF_MUL:.+]] = muli %[[DIFF]], %[[FRACTION]]
   // CHECK: %[[RESULT:.+]] = addi %[[BASE_MUL]], %[[DIFF_MUL]]
   // CHECK: linalg.yield %[[RESULT]]
   %0 = "tosa.table"(%arg0, %arg1)  : (tensor<6xi16>, tensor<513xi16>)  -> (tensor<6xi32>)
   return
 }
+
+// -----
+
+// CHECK-LABEL: @max_pool
+func @max_pool(%arg0: tensor<1x6x34x62xf32>) -> () {
+  // CHECK-DAG: [[CONST:%.+]] = constant -3.40282347E+38
+  // CHECK-DAG: [[INIT:%.+]] = linalg.init_tensor [1, 4, 32, 62]
+  // CHECK-DAG: [[FILL:%.+]] = linalg.fill([[INIT]], [[CONST]])
+  // CHECK-DAG: [[KERNEL:%.+]] = linalg.init_tensor [3, 3]
+  // CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%arg0, [[KERNEL]] : tensor<1x6x34x62xf32>, tensor<3x3xf32>) outs([[FILL]] : tensor<1x4x32x62xf32>)
+  %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xf32>)  -> (tensor<1x4x32x62xf32>)
+  return
+}
+
+// CHECK-LABEL: @max_pool_padded
+func @max_pool_padded(%arg0: tensor<1x6x34x62xf32>) -> () {
+  // CHECK-DAG: [[CONST:%.+]] = constant -3.40282347E+38 : f32
+  // CHECK-DAG: [[PAD:%.+]] = linalg.pad_tensor %arg0 low[0, 0, 0, 0] high[0, 0, 1, 0]
+  // CHECK-DAG:   linalg.yield [[CONST]]
+  // CHECK-DAG: [[INIT:%.+]] = linalg.init_tensor [1, 4, 33, 62]
+  // CHECK-DAG: [[FILL:%.+]] = linalg.fill([[INIT]], [[CONST]])
+  // CHECK-DAG: [[KERNEL:%.+]] = linalg.init_tensor [3, 3]
+  // CHECK: linalg.pooling_nhwc_max {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins([[PAD]], [[KERNEL]] : tensor<1x6x35x62xf32>, tensor<3x3xf32>) outs([[FILL]] : tensor<1x4x33x62xf32>)
+  %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 1], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xf32>)  -> (tensor<1x4x33x62xf32>)
+  return
+}
+
+// CHECK-LABEL: @max_pool_i8
+func @max_pool_i8(%arg0: tensor<1x6x34x62xi8>) -> () {
+  // CHECK: constant -128
+  // CHECK: linalg.pooling_nhwc_i8_max
+  %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xi8>)  -> (tensor<1x4x32x62xi8>)
+  return
+}
+
+// CHECK-LABEL: @max_pool_i16
+func @max_pool_i16(%arg0: tensor<1x6x34x62xi16>) -> () {
+  // CHECK: constant -32768
+  // CHECK: linalg.pooling_nhwc_i16_max
+  %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xi16>)  -> (tensor<1x4x32x62xi16>)
+  return
+}
+
+// CHECK-LABEL: @max_pool_i32
+func @max_pool_i32(%arg0: tensor<1x6x34x62xi32>) -> () {
+  // CHECK: constant -2147483648
+  // CHECK: linalg.pooling_nhwc_i32_max
+  %0 = "tosa.max_pool2d"(%arg0) {pad = [0, 0, 0, 0], kernel = [3, 3], stride = [1, 1]} : (tensor<1x6x34x62xi32>)  -> (tensor<1x4x32x62xi32>)
+  return
+}

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

 
 // 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_i8_max(%input: memref<?x?x?x?xi8>, %fake: memref<2x3xi8>, %init: memref<?x?x?x?xi8>) {
+  linalg.pooling_nhwc_i8_max {dilations = dense<1> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64>}
+    ins(%input, %fake: memref<?x?x?x?xi8>, memref<2x3xi8>)
+    outs(%init: memref<?x?x?x?xi8>)
+  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_i8_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?xi8>, memref<2x3xi8>)
+// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xi8>)
+
+// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: i8, %[[BBARG1:.+]]: i8, %[[BBARG2:.+]]: i8)
+// CHECK-NEXT:      %[[CMP:.+]] = cmpi sgt, %[[BBARG0]], %[[BBARG2]] : i8
+// CHECK-NEXT:      %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : i8
+// CHECK-NEXT:      linalg.yield %[[RES]] : i8
+
+// -----
+
+func @pooling_nhwc_i16_max(%input: memref<?x?x?x?xi16>, %fake: memref<2x3xi16>, %init: memref<?x?x?x?xi16>) {
+  linalg.pooling_nhwc_i16_max {dilations = dense<1> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64>}
+    ins(%input, %fake: memref<?x?x?x?xi16>, memref<2x3xi16>)
+    outs(%init: memref<?x?x?x?xi16>)
+  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_i16_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?xi16>, memref<2x3xi16>)
+// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xi16>)
+
+// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: i16, %[[BBARG1:.+]]: i16, %[[BBARG2:.+]]: i16)
+// CHECK-NEXT:      %[[CMP:.+]] = cmpi sgt, %[[BBARG0]], %[[BBARG2]] : i16
+// CHECK-NEXT:      %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : i16
+// CHECK-NEXT:      linalg.yield %[[RES]] : i16
+
+// -----
+
+func @pooling_nhwc_i32_max(%input: memref<?x?x?x?xi32>, %fake: memref<2x3xi32>, %init: memref<?x?x?x?xi32>) {
+  linalg.pooling_nhwc_i32_max {dilations = dense<1> : tensor<2xi64>, strides = dense<[2, 3]> : tensor<2xi64>}
+    ins(%input, %fake: memref<?x?x?x?xi32>, memref<2x3xi32>)
+    outs(%init: memref<?x?x?x?xi32>)
+  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_i32_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?xi32>, memref<2x3xi32>)
+// CHECK-SAME: outs(%{{.+}} : memref<?x?x?x?xi32>)
+
+// CHECK:         ^{{.+}}(%[[BBARG0:.+]]: i32, %[[BBARG1:.+]]: i32, %[[BBARG2:.+]]: i32)
+// CHECK-NEXT:      %[[CMP:.+]] = cmpi sgt, %[[BBARG0]], %[[BBARG2]] : i32
+// CHECK-NEXT:      %[[RES:.+]] = select %[[CMP]], %[[BBARG0]], %[[BBARG2]] : i32
+// CHECK-NEXT:      linalg.yield %[[RES]] : i32
+
+// -----
+
 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)>

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

func @pooling_nhwc_max(%input: memref<1x4x4x1xf32>, %fake: memref<3x3xf32>, %output: memref<1x2x2x1xf32>) {
   linalg.pooling_nhwc_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
     ins(%input, %fake: memref<1x4x4x1xf32>, memref<3x3xf32>)
     outs(%output: memref<1x2x2x1xf32>)
   return
 }
 
 // -----
 
+// CHECK-LABEL: func @pooling_nhwc_i8_max_tensor
+// CHECK:         %{{.+}} = linalg.pooling_nhwc_i8_max
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xi8>, tensor<3x3xi8>)
+// CHECK-SAME:      outs(%{{.+}} : tensor<1x2x2x1xi8>) -> tensor<1x2x2x1xi8>
+func @pooling_nhwc_i8_max_tensor(%input: tensor<1x4x4x1xi8>) -> tensor<1x2x2x1xi8> {
+  %fake = linalg.init_tensor [3, 3] : tensor<3x3xi8>
+  %init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xi8>
+  %cst = constant 0 : i8
+  %fill = linalg.fill(%init, %cst) : tensor<1x2x2x1xi8>, i8 -> tensor<1x2x2x1xi8>
+  %res = linalg.pooling_nhwc_i8_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: tensor<1x4x4x1xi8>, tensor<3x3xi8>)
+    outs(%fill: tensor<1x2x2x1xi8>) -> tensor<1x2x2x1xi8>
+  return %res : tensor<1x2x2x1xi8>
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_i8_max
+// CHECK:         linalg.pooling_nhwc_i8_max
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : memref<1x4x4x1xi8>, memref<3x3xi8>)
+// CHECK-SAME:      outs(%{{.+}} : memref<1x2x2x1xi8>)
+func @pooling_nhwc_i8_max(%input: memref<1x4x4x1xi8>, %fake: memref<3x3xi8>, %output: memref<1x2x2x1xi8>) {
+  linalg.pooling_nhwc_i8_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: memref<1x4x4x1xi8>, memref<3x3xi8>)
+    outs(%output: memref<1x2x2x1xi8>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_i16_max_tensor
+// CHECK:         %{{.+}} = linalg.pooling_nhwc_i16_max
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xi16>, tensor<3x3xi16>)
+// CHECK-SAME:      outs(%{{.+}} : tensor<1x2x2x1xi16>) -> tensor<1x2x2x1xi16>
+func @pooling_nhwc_i16_max_tensor(%input: tensor<1x4x4x1xi16>) -> tensor<1x2x2x1xi16> {
+  %fake = linalg.init_tensor [3, 3] : tensor<3x3xi16>
+  %init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xi16>
+  %cst = constant 0 : i16
+  %fill = linalg.fill(%init, %cst) : tensor<1x2x2x1xi16>, i16 -> tensor<1x2x2x1xi16>
+  %res = linalg.pooling_nhwc_i16_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: tensor<1x4x4x1xi16>, tensor<3x3xi16>)
+    outs(%fill: tensor<1x2x2x1xi16>) -> tensor<1x2x2x1xi16>
+  return %res : tensor<1x2x2x1xi16>
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_i16_max
+// CHECK:         linalg.pooling_nhwc_i16_max
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : memref<1x4x4x1xi16>, memref<3x3xi16>)
+// CHECK-SAME:      outs(%{{.+}} : memref<1x2x2x1xi16>)
+func @pooling_nhwc_i16_max(%input: memref<1x4x4x1xi16>, %fake: memref<3x3xi16>, %output: memref<1x2x2x1xi16>) {
+  linalg.pooling_nhwc_i16_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: memref<1x4x4x1xi16>, memref<3x3xi16>)
+    outs(%output: memref<1x2x2x1xi16>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_i32_max_tensor
+// CHECK:         %{{.+}} = linalg.pooling_nhwc_i32_max
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : tensor<1x4x4x1xi32>, tensor<3x3xi32>)
+// CHECK-SAME:      outs(%{{.+}} : tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32>
+func @pooling_nhwc_i32_max_tensor(%input: tensor<1x4x4x1xi32>) -> tensor<1x2x2x1xi32> {
+  %fake = linalg.init_tensor [3, 3] : tensor<3x3xi32>
+  %init = linalg.init_tensor [1, 2, 2, 1] : tensor<1x2x2x1xi32>
+  %cst = constant 0 : i32
+  %fill = linalg.fill(%init, %cst) : tensor<1x2x2x1xi32>, i32 -> tensor<1x2x2x1xi32>
+  %res = linalg.pooling_nhwc_i32_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: tensor<1x4x4x1xi32>, tensor<3x3xi32>)
+    outs(%fill: tensor<1x2x2x1xi32>) -> tensor<1x2x2x1xi32>
+  return %res : tensor<1x2x2x1xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @pooling_nhwc_i32_max
+// CHECK:         linalg.pooling_nhwc_i32_max
+// CHECK-SAME:      dilations = dense<1> : tensor<2xi64>
+// CHECK-SAME:      strides = dense<1> : tensor<2xi64>
+// CHECK-SAME:      ins(%{{.+}}, %{{.+}} : memref<1x4x4x1xi32>, memref<3x3xi32>)
+// CHECK-SAME:      outs(%{{.+}} : memref<1x2x2x1xi32>)
+func @pooling_nhwc_i32_max(%input: memref<1x4x4x1xi32>, %fake: memref<3x3xi32>, %output: memref<1x2x2x1xi32>) {
+  linalg.pooling_nhwc_i32_max {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
+    ins(%input, %fake: memref<1x4x4x1xi32>, memref<3x3xi32>)
+    outs(%output: memref<1x2x2x1xi32>)
+  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<1x4x4x1xf32>, tensor<3x3xf32>)
 // CHECK-SAME:      outs(%{{.+}} : tensor<1x2x2x1xf32>) -> tensor<1x2x2x1xf32>
 func @pooling_nhwc_min_tensor(%input: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
   %fake = linalg.init_tensor [3, 3] : tensor<3x3xf32>