[mlir][tosa] Fix tosa average_pool2d to linalg type issue

[rsuderman]

Summary

Average pool assumed the same input/output type. Result type for integers
is always an i32, should be updated appropriately.

[GMNGeoffrey]

If the output should always be i32, could you add verification to the op for this? That would have avoided this error.

[rsuderman]

Added tosa.avg_pool2d verifier.

[GMNGeoffrey]

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

72

So this is limited to f32, i8, or i32? It would be nice to express those constraints with type constraints instead of custom verification.

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

355–360

Please add failure messages with emitOpError(). Longer term, it would be better to do this with return type inference.

[rsuderman]

Added op error.

[rsuderman]

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

72

I agree we should update this restriction, but that should be done as a separate pass with fixing tests. It is surprisingly common to find tests that use slightly incorrect types.

[GMNGeoffrey]

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

362–364

This returns LogicalResult so you can combine them.

[rsuderman]

Changed return statement.

[GMNGeoffrey]

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

72

Yeah upsettingly common :-/ That's why I push hard for verifiers

[rsuderman]

Rebase on head.

mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td

def Tosa_AvgPool2dOp : Tosa_Op<"avg_pool2d", [
 
   let description = [{
     This performs an average pooling over the given input tensor. A sliding
     window of size given by <kernel size> is passed over the input tensor, with
     the mean value being placed in the output tensor.
   }];
 
   let arguments = (ins
     Tosa_Tensor4D:$input,
 
     Tosa_IntArrayAttr2:$kernel,
     Tosa_IntArrayAttr2:$stride,
     Tosa_IntArrayAttr4:$pad,
     OptionalAttr<Tosa_UnaryOpQuantizationAttr>:$quantization_info
   );
 
   let results = (outs
     Tosa_Tensor4D:$output
   );
 
   let builders = [Tosa_AvgPool2dOpQuantInfoBuilder];
+
+  let verifier = [{ return verifyAveragePoolOp(*this); }];
 }
 
 //===----------------------------------------------------------------------===//
 // Operator: conv2d
 //===----------------------------------------------------------------------===//
 def Tosa_Conv2DOp : Tosa_Op<"conv2d", [
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
                               ["inferReturnTypeComponents"]>,

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

public:
   LogicalResult matchAndRewrite(tosa::AvgPool2dOp 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().template cast<ShapedType>();
-    Type resultETy = inputTy.getElementType();
+    Type resultETy = op.getType().cast<ShapedType>().getElementType();
 
     Type accETy =
         inElementTy.isa<IntegerType>() ? rewriter.getI32Type() : inElementTy;
     ShapedType accTy = resultTy.clone(accETy);
 
     if (!inputTy.hasStaticShape())
       return failure();
 
     // Apply padding as necessary.
     llvm::SmallVector<int64_t> pad;
     pad.resize(2, 0);
     getValuesFromIntArrayAttribute(op.pad(), pad);
     pad.resize(pad.size() + 2, 0);
-    Attribute initialAttr = rewriter.getZeroAttr(accETy);
-    Value paddedInput = applyPad(loc, input, pad, initialAttr, rewriter);
+    Attribute padAttr = rewriter.getZeroAttr(inElementTy);
+    Value paddedInput = applyPad(loc, input, pad, padAttr, rewriter);
 
+    Attribute initialAttr = rewriter.getZeroAttr(accETy);
     Value initialValue = rewriter.create<ConstantOp>(loc, initialAttr);
 
     SmallVector<int64_t> kernel, stride;
     getValuesFromIntArrayAttribute(op.kernel(), kernel);
     getValuesFromIntArrayAttribute(op.stride(), stride);
 
     Attribute strideAttr = rewriter.getI64VectorAttr(stride);
     Attribute dilationAttr = rewriter.getI64VectorAttr({1, 1});
auto genericOp = rewriter.create<linalg::GenericOp>(
           auto countI = rewriter.create<mlir::IndexCastOp>(
               loc, rewriter.getI32Type(), count);
 
           // Divide by the number of summed values. For floats this is just
           // a div however for quantized values input normalization had
           // to be applied.
           Value poolVal = args[0];
           if (accETy.isa<FloatType>()) {
-            auto countF =
-                rewriter.create<mlir::SIToFPOp>(loc, inElementTy, countI);
+            auto countF = rewriter.create<mlir::SIToFPOp>(loc, accETy, countI);
             poolVal =
                 rewriter.create<DivFOp>(loc, poolVal, countF)->getResult(0);
           } else {
 
             // If we have quantization information we need to apply an offset
             // for the input zp value.
             if (op.quantization_info()) {
               auto quantizationInfo = op.quantization_info().getValue();
auto genericOp = rewriter.create<linalg::GenericOp>(
                          APInt::getSignedMinValue(outBitwidth).getSExtValue()));
             auto max = rewriter.create<ConstantOp>(
                 loc, rewriter.getIntegerAttr(
                          accETy,
                          APInt::getSignedMaxValue(outBitwidth).getSExtValue()));
             auto clamp = clampHelper<mlir::CmpIOp>(
                 loc, scaled, min, max, CmpIPredicate::slt, rewriter);
 
+            poolVal = clamp;
             // Convert type.
-            poolVal = rewriter.create<TruncateIOp>(loc, resultETy, clamp);
+            if (resultETy != clamp.getType()) {
+              poolVal = rewriter.create<TruncateIOp>(loc, resultETy, poolVal);
+            }
           }
 
           // Cast to output type.
 
           rewriter.create<linalg::YieldOp>(loc, poolVal);
         });
 
     rewriter.replaceOp(op, genericOp.getResult(0));

mlir/lib/Dialect/Tosa/IR/TosaOps.cpp

static LogicalResult verifyConvOp(T op) {
   // types should not have a quantizationattr.
   if ((inputIsQuant && !op.quantization_info()) ||
       (!inputIsQuant && op.quantization_info()))
     return failure();
 
   return success();
 }
 
+static LogicalResult verifyAveragePoolOp(tosa::AvgPool2dOp op) {
+  auto inputETy = op.input().getType().cast<ShapedType>().getElementType();
+  auto resultETy = op.getType().cast<ShapedType>().getElementType();
+
+  if (auto quantType = inputETy.dyn_cast<mlir::quant::UniformQuantizedType>())
+    inputETy = quantType.getStorageType();
+
+  if (auto quantType = resultETy.dyn_cast<mlir::quant::UniformQuantizedType>())
+    resultETy = quantType.getStorageType();
+
+  if (inputETy.isF32() && resultETy.isF32())
+    return success();
+  if (inputETy.isInteger(8) && resultETy.isInteger(32))
+    return success();
+  if (inputETy.isInteger(16) && resultETy.isInteger(32))
+    return success();
+
+  return op.emitOpError("input/output element types are incompatible.");
+}
+
 //===----------------------------------------------------------------------===//
 // TOSA Operator Quantization Builders.
 //===----------------------------------------------------------------------===//
 
 /// This builder is called on all convolution operators except TransposeConv,
 /// which has specialized output shape semantics. The builder also defines the
 /// bitwidth of the output given the bit width of the input & weight content.
 static void buildConvOpWithQuantInfo(OpBuilder &builder, OperationState &result,

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

func @avg_pool_i8(%arg0 : tensor<1x128x128x2xi8>) -> () {
   // CHECK: %[[INZP_OFF:.+]] = muli %{{.+}}, %[[INZP]]
   // CHECK: %[[OFFSETED:.+]] = subi %arg1, %[[INZP_OFF]]
   // CHECK: %[[NUMERATOR:.+]] = constant 1073741825
   // CHECK: %[[MULTIPLIER:.+]] = divi_unsigned %[[NUMERATOR]], %{{.+}}
   // CHECK: %[[SHIFT:.+]] = constant 30
   // CHECK: %[[SCALE:.+]] = "tosa.apply_scale"(%{{.+}}, %[[MULTIPLIER]], %[[SHIFT]]) {double_round = false}
   // CHECK: %[[OUTZP:.+]] = constant -128
   // CHECK: %[[OUT:.+]] = addi %[[SCALE]], %[[OUTZP]]
-  // CHECK: %[[MIN:.+]] = constant -128
-  // CHECK: %[[MAX:.+]] = constant 127
+  // CHECK: %[[MIN:.+]] = constant -2147483648
+  // CHECK: %[[MAX:.+]] = constant 2147483647
   // CHECK: %[[CMP_MIN:.+]] = cmpi slt, %[[OUT]], %[[MIN]]
   // CHECK: %[[CLMP_MIN:.+]] = select %[[CMP_MIN]], %[[MIN]], %[[OUT]]
   // CHECK: %[[CMP_MAX:.+]] = cmpi slt, %[[MAX]], %[[OUT]]
   // CHECK: %[[CLMP_MAX:.+]] = select %[[CMP_MAX]], %[[MAX]], %[[CLMP_MIN]]
-  // CHECK: %[[TRUNC:.+]] = trunci %[[CLMP_MAX]]
-  // CHECK: linalg.yield %[[TRUNC]]
-  %0 = "tosa.avg_pool2d"(%arg0) {kernel = [4, 4], pad = [0, 0, 0, 0], quantization_info = {input_zp = -128 : i32, output_zp = -128 : i32}, stride = [4, 4]} : (tensor<1x128x128x2xi8>) -> tensor<1x32x32x2xi8>
+  // CHECK: linalg.yield %[[CLMP_MAX]]
+  %0 = "tosa.avg_pool2d"(%arg0) {kernel = [4, 4], pad = [0, 0, 0, 0], quantization_info = {input_zp = -128 : i32, output_zp = -128 : i32}, stride = [4, 4]} : (tensor<1x128x128x2xi8>) -> tensor<1x32x32x2xi32>
   return
 }
 
 // -----
 
 // CHECK: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>
 // CHECK: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
 // CHECK: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>

mlir/test/Dialect/Tosa/ops.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s
 // RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
 
 
 // -----
 // CHECK-LABEL: argmax
 func @test_argmax(%arg0: tensor<14x19xf32>) -> tensor<14xi32> {
   %0 = "tosa.argmax"(%arg0) {axis = 1 : i64} : (tensor<14x19xf32>) -> tensor<14xi32>
   return %0 : tensor<14xi32>
 }
 
 // -----
-// CHECK-LABEL: avg_pool2d
-func @test_avg_pool2d(%arg0: tensor<1x7x7x9xf32>) -> tensor<1x7x7x9xf32> {
+// CHECK-LABEL: avg_pool2d_f32
+func @test_avg_pool2d_f32(%arg0: tensor<1x7x7x9xf32>) -> tensor<1x7x7x9xf32> {
     %0 = "tosa.avg_pool2d"(%arg0) {kernel = [2, 2], pad = [0, 1, 0, 1], stride = [1, 1]} : (tensor<1x7x7x9xf32>) -> tensor<1x7x7x9xf32>
     return %0 : tensor<1x7x7x9xf32>
 }
 
 // -----
+// CHECK-LABEL: avg_pool2d_i8
+func @test_avg_pool2d_i8(%arg0: tensor<1x7x7x9xi8>) -> tensor<1x7x7x9xi32> {
+    %0 = "tosa.avg_pool2d"(%arg0) {kernel = [2, 2], pad = [0, 1, 0, 1], stride = [1, 1]} : (tensor<1x7x7x9xi8>) -> tensor<1x7x7x9xi32>
+    return %0 : tensor<1x7x7x9xi32>
+}
+
+// -----
+// CHECK-LABEL: avg_pool2d_i16
+func @test_avg_pool2d_i16(%arg0: tensor<1x7x7x9xi16>) -> tensor<1x7x7x9xi32> {
+    %0 = "tosa.avg_pool2d"(%arg0) {kernel = [2, 2], pad = [0, 1, 0, 1], stride = [1, 1]} : (tensor<1x7x7x9xi16>) -> tensor<1x7x7x9xi32>
+    return %0 : tensor<1x7x7x9xi32>
+}
+
+// -----
+// CHECK-LABEL: avg_pool2d_q8
+func @test_avg_pool2d_q8(%arg0: tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>) -> tensor<1x7x7x9x!quant.uniform<i32:f32, 0.01>> {
+    %0 = "tosa.avg_pool2d"(%arg0) {kernel = [2, 2], pad = [0, 1, 0, 1], stride = [1, 1]} : (tensor<1x7x7x9x!quant.uniform<i8:f32, 0.01>>) -> tensor<1x7x7x9x!quant.uniform<i32:f32, 0.01>>
+    return %0 : tensor<1x7x7x9x!quant.uniform<i32:f32, 0.01>>
+}
+
+// -----
 // CHECK-LABEL: conv2d
 func @test_conv2d(%arg0: tensor<1x4x4x4xf32>, %arg1: tensor<8x1x1x4xf32>, %arg2: tensor<8xf32>) -> tensor<1x4x4x8xf32> {
     %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = [1, 1], pad = [0, 0, 0, 0], stride = [1, 1]} : (tensor<1x4x4x4xf32>, tensor<8x1x1x4xf32>, tensor<8xf32>) -> tensor<1x4x4x8xf32>
     return %0 : tensor<1x4x4x8xf32>
 }
 
 // -----
 // CHECK-LABEL: depthwise_conv2d