[mlir][tosa] Fix padding for tosa.conv2d and tosa.depthwise_conv2d decomposition

[rsuderman]

Summary

Decomposition did not take padding into account when decomposing into fully
connected operation.

[rsuderman]

Missing newline

mlir/include/mlir/Dialect/Tosa/Utils/ConversionUtils.h

 Value clampFloatHelper(Location loc, Value arg, Value min, Value max,
                        OpBuilder &rewriter);
 
 // Takes the parameters for a clamp and turns it into a series of ops for
 // integer inputs.
 Value clampIntHelper(Location loc, Value arg, Value min, Value max,
                      OpBuilder &rewriter);
 
+// Determines whether the integer value falls witin the range of integer type.
+bool validIntegerRange(IntegerType ty, int64_t value);
+
 // Returns the values in an attribute as an array of values.
 template <typename T>
 void getValuesFromIntArrayAttribute(ArrayAttr attr,
                                     SmallVector<T> &arrayValues) {
   for (Attribute val : attr.getValue()) {
     arrayValues.push_back(val.cast<IntegerAttr>().getValue().getSExtValue());
   }
 }

mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeConv2D.cpp

 //===- TosaDecomposeConv2D.cpp --------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Decompose TOSA Conv2D operation to a series of TOSA Ops specifically
 // (1) Convert a 1x1 Convolution to a Reshape->FC->Reshape
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
 
 using namespace mlir;
 using namespace mlir::tosa;
 
 namespace {
 
 SmallVector<int64_t> convertFromMlirShape(ArrayRef<int64_t> shape) {
   return to_vector(llvm::map_range(shape, [](int64_t dim) {
for (APInt stride : op.getStride().getAsValueRange<IntegerAttr>()) {
         return failure();
     }
 
     // Only works for a 1x1 kernel.
     ArrayRef<int64_t> weightShape = weightType.getShape();
     if (weightShape[1] != 1 || weightShape[2] != 1)
       return failure();
 
+    auto padAttr = op.getPad();
+    llvm::SmallVector<int64_t> pad(8, 0);
+    for (auto it : llvm::enumerate(padAttr.getValue()))
+      pad[it.index() + 2] =
+          it.value().cast<IntegerAttr>().getValue().getSExtValue();
+
+    if (llvm::any_of(pad, [](int64_t p) { return p != 0; })) {
+      Type inputETy = inputType.getElementType();
+      Attribute zeroAttr = rewriter.getZeroAttr(inputETy);
+      if (op.getQuantizationInfo()) {
+        auto quantizationInfo = op.getQuantizationInfo();
+        int64_t iZp = quantizationInfo->getInputZp();
+
+        if (!validIntegerRange(inputETy.cast<IntegerType>(), iZp))
+          return rewriter.notifyMatchFailure(
+              op, "tosa.conv op quantization has zp outside of input range");
+
+        zeroAttr = rewriter.getIntegerAttr(inputETy, iZp);
+      }
+
+      llvm::SmallVector<int64_t> newShape(inputType.getShape());
+
+      for (int i = 0, s = newShape.size(); i < s; ++i) {
+        if (newShape[i] != ShapedType::kDynamic) {
+          newShape[i] += pad[i * 2] + pad[i * 2 + 1];
+        }
+      }
+
+      auto padSizeTy = RankedTensorType::get({4, 2}, rewriter.getI64Type());
+      auto padSize =
+          DenseIntElementsAttr::get(padSizeTy, ArrayRef<int64_t>(pad));
+      Value padSizeVal =
+          rewriter.create<tosa::ConstOp>(op->getLoc(), padSizeTy, padSize);
+
+      auto padTy = RankedTensorType::get({}, inputETy);
+      auto padAttr = DenseElementsAttr::get(padTy, zeroAttr);
+      Value padVal =
+          rewriter.create<tosa::ConstOp>(op->getLoc(), padTy, padAttr);
+      inputType = RankedTensorType::get(newShape, inputETy);
+      input = rewriter.create<tosa::PadOp>(op->getLoc(), inputType, input,
+                                           padSizeVal, padVal);
+    }
+
     // Reshape input to [N,IH,IW,IC] -> [N * IH * IW, IC].
     ArrayRef<int64_t> inputShape = inputType.getShape();
     int64_t combined = ShapedType::kDynamic;
     if (numDynamic == 0)
       combined = inputShape[0] * inputShape[1] * inputShape[2];
     llvm::SmallVector<int64_t, 2> revisedInputShape{combined, inputShape[3]};
     auto revisedInputShapeType =
         RankedTensorType::get(revisedInputShape, inputType.getElementType());

mlir/lib/Dialect/Tosa/Transforms/TosaDecomposeDepthwise.cpp

struct DepthwiseConv2DIsMul : public OpRewritePattern<tosa::DepthwiseConv2DOp> {
 
   LogicalResult matchAndRewrite(tosa::DepthwiseConv2DOp op,
                                 PatternRewriter &rewriter) const override {
     Value input = op.getInput();
     Value weight = op.getWeight();
     ShapedType inputType = input.getType().cast<ShapedType>();
     ShapedType weightType = weight.getType().cast<ShapedType>();
     ShapedType resultType = op.getOutput().getType().cast<ShapedType>();
-    Type inputEType = inputType.getElementType();
 
     if (!(inputType.hasStaticShape() && weightType.hasStaticShape() &&
           resultType.hasStaticShape())) {
       return failure();
     }
 
-    // Quantization information needs to still be performed.
-    if (op.getQuantizationInfo() || !inputEType.isa<FloatType>()) {
-      return failure();
-    }
-
     // Stride must be 1 for this optimization.
     for (Attribute stride : op.getStride().getValue()) {
       if (!stride.cast<IntegerAttr>().getValue().isOne()) {
         return failure();
       }
     }
 
     // Only works for a 1x1 kernel.
     ArrayRef<int64_t> weightShape = weightType.getShape();
     if (weightShape[0] != 1 || weightShape[1] != 1) {
       return failure();
     }
 
     // Reshape input to [N, H, W, C] -> [N, H, W, C, 1].
     ArrayRef<int64_t> inputShape = inputType.getShape();
     llvm::SmallVector<int64_t, 2> revisedInputShape{
         inputShape[0], inputShape[1], inputShape[2], inputShape[3], 1};
-    auto revisedInputShapeType = RankedTensorType::get(
+    inputType = RankedTensorType::get(
         revisedInputShape,
         input.getType().dyn_cast<RankedTensorType>().getElementType());
-    auto reshapedInput = rewriter
-                             .create<tosa::ReshapeOp>(
-                                 op.getLoc(), revisedInputShapeType, input,
-                                 rewriter.getI64ArrayAttr(revisedInputShape))
-                             .getResult();
+    input = rewriter
+                    op.getLoc(), inputType, input,
 
-    // Reshape kernel to [KH, KW, C, M] -> [1, 1, 1, C, M].
-    llvm::SmallVector<int64_t, 2> revisedWeightShape{1, 1, 1, weightShape[2],
-                                                     weightShape[3]};
-    auto revisedWeightShapeType = RankedTensorType::get(
-        revisedWeightShape,
-        weight.getType().dyn_cast<RankedTensorType>().getElementType());
-    auto reshapedWeight = rewriter
-                              .create<tosa::ReshapeOp>(
-                                  op.getLoc(), revisedWeightShapeType, weight,
-                                  rewriter.getI64ArrayAttr(revisedWeightShape))
-                              .getResult();
+    if (inputType.getElementType() != resultType.getElementType()) {
+      inputType = inputType.clone(resultType.getElementType());
+      input = rewriter.create<tosa::CastOp>(op.getLoc(), inputType, input);
+    }
+
+    if (weightType.getElementType() != resultType.getElementType()) {
+      weightType = weightType.clone(resultType.getElementType());
+      weight = rewriter.create<tosa::CastOp>(op.getLoc(), weightType, weight);
+    }
+
+    if (auto quantizationInfo = op.getQuantizationInfo()) {
+      auto iZp = quantizationInfo->getInputZp();
+      auto wZp = quantizationInfo->getWeightZp();
+
+      auto applyZp = [&](Value val, int64_t zp) -> Value {
+        if (zp == 0)
+          return val;
+        auto ety = val.getType().cast<ShapedType>().getElementType();
+        auto zpTy = RankedTensorType::get({}, ety);
+        auto zpAttr =
+            DenseElementsAttr::get(zpTy, rewriter.getIntegerAttr(ety, zp));
+        auto zpVal = rewriter.create<tosa::ConstOp>(op.getLoc(), zpTy, zpAttr);
+        return rewriter.create<tosa::SubOp>(op.getLoc(), val.getType(), val,
+                                            zpVal);
+      };
+
+      input = applyZp(input, iZp);
+      weight = applyZp(weight, wZp);
+    }
+
+    auto padAttr = op.getPad();
+    llvm::SmallVector<int64_t> pad(10, 0);
+    for (auto it : llvm::enumerate(padAttr.getValue()))
+      pad[it.index() + 2] =
+          it.value().cast<IntegerAttr>().getValue().getSExtValue();
+
+    if (llvm::any_of(pad, [](int64_t p) { return p != 0; })) {
+      Type inputETy = inputType.getElementType();
+      Attribute zeroAttr = rewriter.getZeroAttr(inputETy);
+
+      llvm::SmallVector<int64_t> newShape(inputType.getShape());
+      for (int i = 0, s = pad.size(); i < s; ++i) {
+        if (newShape[i / 2] != ShapedType::kDynamic) {
+          newShape[i / 2] += pad[i];
+        }
+      }
+
+      auto padSizeTy = RankedTensorType::get({5, 2}, rewriter.getI64Type());
+      auto padSize =
+          DenseIntElementsAttr::get(padSizeTy, ArrayRef<int64_t>(pad));
+      Value padSizeVal =
+          rewriter.create<tosa::ConstOp>(op->getLoc(), padSizeTy, padSize);
+
+      auto padTy = RankedTensorType::get({}, inputETy);
+      auto padAttr = DenseElementsAttr::get(padTy, zeroAttr);
+      Value padVal =
+          rewriter.create<tosa::ConstOp>(op->getLoc(), padTy, padAttr);
+      inputType = RankedTensorType::get(newShape, inputETy);
+      input = rewriter.create<tosa::PadOp>(op->getLoc(), inputType, input,
+                                           padSizeVal, padVal);
+    }
 
     // Perform an elementwise mul over the reshaped input and weight.
-    llvm::SmallVector<int64_t, 2> mulShape{inputShape[0], inputShape[1],
-                                           inputShape[2], inputShape[3],
-                                           weightShape[3]};
+    llvm::SmallVector<int64_t, 2> mulShape{
+        inputType.getDimSize(0), inputType.getDimSize(1),
+        inputType.getDimSize(2), inputType.getDimSize(3), weightShape[3]};
     auto mulShapeType = RankedTensorType::get(
         mulShape,
         weight.getType().dyn_cast<RankedTensorType>().getElementType());
-    Value mulValue =
-        rewriter
-            .create<tosa::MulOp>(op.getLoc(), mulShapeType, reshapedInput,
-                                 reshapedWeight, /*shift=*/0)
+    Value mulValue = rewriter
+                         .create<tosa::MulOp>(op.getLoc(), mulShapeType, input,
+                                              weight, /*shift=*/0)
                          .getResult();
 
     // Reshape output to [N, H, W, C * M].
     auto outputShape = op.getOutput().getType().cast<ShapedType>().getShape();
     auto outputShapeType = RankedTensorType::get(
         outputShape,
         input.getType().dyn_cast<RankedTensorType>().getElementType());
     auto outputValue =
         rewriter.create<tosa::ReshapeOp>(op.getLoc(), outputShapeType, mulValue,

mlir/lib/Dialect/Tosa/Utils/ConversionUtils.cpp

Value mlir::tosa::clampIntHelper(Location loc, Value arg, Value min, Value max,
   auto smallerThanMin =
       rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, arg, min);
   auto minOrArg =
       rewriter.create<arith::SelectOp>(loc, smallerThanMin, min, arg);
   auto largerThanMax =
       rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::slt, max, arg);
   return rewriter.create<arith::SelectOp>(loc, largerThanMax, max, minOrArg);
 }
+
+bool mlir::tosa::validIntegerRange(IntegerType ty, int64_t value) {
+  uint64_t bitwidth = ty.getIntOrFloatBitWidth();
+  if (ty.getSignedness() == IntegerType::Unsigned) {
+    uint64_t uvalue = value;
+    APInt intMin = APInt::getMinValue(bitwidth);
+    APInt intMax = APInt::getMaxValue(bitwidth);
+    return uvalue >= intMin.getZExtValue() && uvalue <= intMax.getZExtValue();
+  }
+
+  APInt intMin = APInt::getSignedMinValue(bitwidth);
+  APInt intMax = APInt::getSignedMaxValue(bitwidth);
+  return value >= intMin.getSExtValue() && value <= intMax.getSExtValue();
+}

mlir/test/Dialect/Tosa/tosa-decompose-conv2d.mlir

 // CHECK:           %[[VAL_5:.*]] = "tosa.fully_connected"(%[[VAL_3]], %[[VAL_4]], %[[VAL_2]]) {quantization_info = #tosa.conv_quant<input_zp = -6, weight_zp = 11>} : (tensor<?x64xi8>, tensor<384x64xi8>, tensor<384xi32>) -> tensor<?x384xi32>
 // CHECK:           %[[VAL_6:.*]] = "tosa.reshape"(%[[VAL_5]]) {new_shape = [-1, 14, 14, 384]} : (tensor<?x384xi32>) -> tensor<?x14x14x384xi32>
 // CHECK:           return %[[VAL_6]] : tensor<?x14x14x384xi32>
 // CHECK:         }
   %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {dilation = [1, 1], pad = [0, 0, 0, 0], quantization_info = #tosa.conv_quant<input_zp = -6, weight_zp = 11>, stride = [1, 1]} : (tensor<?x14x14x64xi8>, tensor<384x1x1x64xi8>, tensor<384xi32>) -> tensor<?x14x14x384xi32>
   return %0 : tensor<?x14x14x384xi32>
 }
 
+// -----
+
+// CHECK-LABEL: @conv2d_as_fully_connected_padded
+func.func @conv2d_as_fully_connected_padded(%arg0: tensor<4x10x10x2xi8>, %arg1: tensor<3x1x1x2xi8>, %arg2: tensor<3xi32>) -> tensor<4x12x12x3xi32> {
+  // CHECK-DAG: %[[PAD_SHAPE:.+]] = "tosa.const"() {value = dense<{{\[\[}}0, 0], [1, 1], [1, 1], [0, 0]]> : tensor<4x2xi64>}
+  // CHECK-DAG: %[[PAD_VAL:.+]] = "tosa.const"() {value = dense<42> : tensor<i8>}
+  // CHECK-DAG: %[[PAD:.+]] = "tosa.pad"(%arg0, %[[PAD_SHAPE]], %[[PAD_VAL]]) : (tensor<4x10x10x2xi8>, tensor<4x2xi64>, tensor<i8>) -> tensor<4x12x12x2xi8>
+  // CHECK-DAG: %[[RESHAPE_INPUT:.+]] = "tosa.reshape"(%[[PAD]]) {new_shape = [576, 2]}
+  // CHECK-DAG: %[[RESHAPE_FILTER:.+]] = "tosa.reshape"(%arg1) {new_shape = [3, 2]}
+  // CHECK-DAG: %[[FULLY:.+]] = "tosa.fully_connected"(%[[RESHAPE_INPUT]], %[[RESHAPE_FILTER]], %arg2) {quantization_info = #tosa.conv_quant<input_zp = 42, weight_zp = 24>}
+  // CHECK: %[[RESHAPE:.+]] = "tosa.reshape"(%[[FULLY]]) {new_shape = [4, 12, 12, 3]}
+  %0 = "tosa.conv2d"(%arg0, %arg1, %arg2) {pad = [1, 1, 1, 1], stride = [1, 1], dilation = [1, 1], quantization_info = #tosa.conv_quant<input_zp = 42, weight_zp = 24>} : (tensor<4x10x10x2xi8>, tensor<3x1x1x2xi8>, tensor<3xi32>) -> tensor<4x12x12x3xi32>
+  return %0 : tensor<4x12x12x3xi32>
+}
 No newline at end of file

mlir/test/Dialect/Tosa/tosa-decompose-depthwise.mlir

 // RUN: mlir-opt --split-input-file --tosa-optional-decompositions %s | FileCheck %s
 
 // -----
 
 // CHECK-LABEL: @depthwise_conv2d_as_mul
 func.func @depthwise_conv2d_as_mul(%arg0: tensor<4x10x10x2xf32>, %arg1: tensor<1x1x2x3xf32>, %arg2: tensor<6xf32>) -> tensor<4x10x10x6xf32> {
   // CHECK-NOT: "tosa.depthwise_conv2d"
   // CHECK: %[[VAR0:.*]] = "tosa.reshape"(%arg0) {new_shape = [4, 10, 10, 2, 1]}
   // CHECK-SAME: -> tensor<4x10x10x2x1xf32>
-  // CHECK: %[[VAR1:.*]] = "tosa.reshape"(%arg1) {new_shape = [1, 1, 1, 2, 3]}
-  // CHECK-SAME: -> tensor<1x1x1x2x3xf32>
-  // CHECK: %[[VAR2:.*]] = "tosa.mul"(%[[VAR0]], %[[VAR1]])
+  // CHECK: %[[VAR2:.*]] = "tosa.mul"(%[[VAR0]], %arg1)
   // CHECK-SAME: -> tensor<4x10x10x2x3xf32>
   // CHECK: %[[VAR3:.*]] = "tosa.reshape"(%[[VAR2]]) {new_shape = [4, 10, 10, 6]}
   // CHECK-SAME: -> tensor<4x10x10x6xf32>
   // CHECK: %[[VAR4:.*]] = "tosa.add"(%[[VAR3]], %arg2)
   // CHECK-SAME: -> tensor<4x10x10x6xf32>
   // CHECK: return %[[VAR4]]
   %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1]} : (tensor<4x10x10x2xf32>, tensor<1x1x2x3xf32>, tensor<6xf32>) -> tensor<4x10x10x6xf32>
   return %0 : tensor<4x10x10x6xf32>
 }
 
 // -----
 
 // CHECK-LABEL: @depthwise_conv2d_as_mul_q
 func.func @depthwise_conv2d_as_mul_q(%arg0: tensor<4x10x10x2xi8>, %arg1: tensor<1x1x2x3xi8>, %arg2: tensor<6xi32>) -> tensor<4x10x10x6xi32> {
-  // CHECK: "tosa.depthwise_conv2d"
-  %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1], quantization_info = #tosa.conv_quant<input_zp = 0, weight_zp = 0>} : (tensor<4x10x10x2xi8>, tensor<1x1x2x3xi8>, tensor<6xi32>) -> tensor<4x10x10x6xi32>
+  // CHECK: %[[iZp:.+]] = "tosa.const"() {value = dense<7> : tensor<i32>}
+  // CHECK: %[[wZp:.+]] = "tosa.const"() {value = dense<11> : tensor<i32>}
+  // CHECK: %[[rIn:.+]] = "tosa.reshape"(%arg0) {new_shape = [4, 10, 10, 2, 1]}
+  // CHECK: %[[cIn:.+]] = "tosa.cast"(%[[rIn]]) : (tensor<4x10x10x2x1xi8>) -> tensor<4x10x10x2x1xi32>
+  // CHECK: %[[cWe:.+]] = "tosa.cast"(%arg1) : (tensor<1x1x2x3xi8>) -> tensor<1x1x2x3xi32>
+  // CHECK: %[[sIn:.+]] = "tosa.sub"(%[[cIn]], %[[iZp]])
+  // CHECK: %[[sWe:.+]] = "tosa.sub"(%[[cWe]], %[[wZp]])
+  // CHECK: %[[mul:.+]] = "tosa.mul"(%[[sIn]], %[[sWe]]) {shift = 0 : i32}
+  // CHECK: %[[reO:.+]] = "tosa.reshape"(%[[mul]]) {new_shape = [4, 10, 10, 6]}
+  // CHECK: %[[add:.+]] = "tosa.add"(%[[reO]], %arg2)
+  %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1], quantization_info = #tosa.conv_quant<input_zp = 7, weight_zp = 11>} : (tensor<4x10x10x2xi8>, tensor<1x1x2x3xi8>, tensor<6xi32>) -> tensor<4x10x10x6xi32>
   return %0 : tensor<4x10x10x6xi32>
 }
 
 // -----
+
+// CHECK-LABEL: @depthwise_conv2d_as_mul_padded
+func.func @depthwise_conv2d_as_mul_padded(%arg0: tensor<4x10x10x2xf32>, %arg1: tensor<1x1x2x3xf32>, %arg2: tensor<6xf32>) -> tensor<4x12x12x6xf32> {
+  // CHECK: %[[pad:.+]] = "tosa.const"() {value = dense<{{\[\[}}0, 0], [1, 1], [1, 1], [0, 0], [0, 0]]> : tensor<5x2xi64>}
+  // CHECK: %[[zero:.+]] = "tosa.const"() {value = dense<0.000000e+00> : tensor<f32>}
+  // CHECK: %[[reIn:.+]] = "tosa.reshape"(%arg0) {new_shape = [4, 10, 10, 2, 1]}
+  // CHECK: %[[padded:.+]] = "tosa.pad"(%[[reIn]], %[[pad]], %[[zero]]) : (tensor<4x10x10x2x1xf32>, tensor<5x2xi64>, tensor<f32>) -> tensor<4x12x12x2x1xf32>
+  // CHECK: %[[mul:.+]] = "tosa.mul"(%3, %arg1) {shift = 0 : i32}
+  // CHECK: %[[reOut:.+]] = "tosa.reshape"(%[[mul]]) {new_shape = [4, 12, 12, 6]}
+  // CHECK: %[[add:.+]] = "tosa.add"(%[[reOut]], %arg2)
+  %0 = "tosa.depthwise_conv2d"(%arg0, %arg1, %arg2) {pad = [1, 1, 1, 1], stride = [1, 1], dilation = [1, 1]} : (tensor<4x10x10x2xf32>, tensor<1x1x2x3xf32>, tensor<6xf32>) -> tensor<4x12x12x6xf32>
+  return %0 : tensor<4x12x12x6xf32>
+}