[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 ↗

(On Diff #379144)

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 ↗

(On Diff #379144)

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 ↗

(On Diff #379144)

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 ↗

(On Diff #379151)

This returns LogicalResult so you can combine them.

[rsuderman]

Changed return statement.

[GMNGeoffrey]

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

72 ↗

(On Diff #379144)

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

[rsuderman]

Rebase on head.

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/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)>