[mlir][tosa] Always generated pad_const and remove input_zp attr for PadOp

Co-authored-by: Udaya Ranga <[email protected]>
Co-authored-by: Tai Ly <[email protected]>

Signed-off-by: Jerry Ge <[email protected]>
Change-Id: I2b7a0169b7ec1158d28779713ad125c061e04592

Author: Jerry-Ge

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

@@ -197,14 +197,6 @@ def Tosa_PadOpQuantInfoBuilder : OpBuilder<
                             input, paddings);
   }]>;
 
-def Tosa_ExplicitValuePadOpQuantInfoBuilder : OpBuilder<
-  (ins "Type":$outputType, "Value":$input, "Value":$paddings,
-       "Value":$pad_value),
-  [{
-    buildExplicitValuePadOpWithQuantInfo($_builder, $_state, outputType,
-                                         input, paddings, pad_value);
-  }]>;
-
 // Wrapper over base I32EnumAttr to set common fields.
 class Tosa_I32Enum<string name, string description, list<I32EnumAttrCase> cases>
      : I32EnumAttr<name, description, cases> {

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

@@ -168,6 +168,10 @@ namespace tosa {
 std::optional<Value> createZeroPointTensor(OpBuilder &builder, Location loc,
                                            Type srcElemType, int64_t zp = 0);
 
+// Create a pad-const const tensor with value of `val` of required data-type
+std::optional<Value> createPadConstTensor(OpBuilder &builder, Location loc,
+                                          Value src, int32_t val = 0);
+
 } // namespace tosa
 } // namespace mlir
 

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

@@ -1882,8 +1882,7 @@ def Tosa_PadOp : Tosa_InferShapedTypeOp<"pad"> {
   let arguments = (ins
     Tosa_RankedTensor:$input1,
     Tosa_Shape:$padding,
-    Optional<Tosa_ScalarTensor>:$pad_const,
-    OptionalAttr<I32Attr>:$input_zp
+    Tosa_ScalarTensor:$pad_const
   );
 
   let results = (outs
@@ -1895,10 +1894,8 @@ def Tosa_PadOp : Tosa_InferShapedTypeOp<"pad"> {
     Extension<[Tosa_EXT_FP8E4M3, Tosa_EXT_FP8E5M2, Tosa_EXT_BF16]>,
   ];
 
-  let builders = [Tosa_PadOpQuantInfoBuilder,
-                  Tosa_ExplicitValuePadOpQuantInfoBuilder];
+  let builders = [Tosa_PadOpQuantInfoBuilder];
 
-  let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
 }

mlir/lib/Conversion/TosaToTensor/TosaToTensor.cpp

@@ -36,10 +36,10 @@ TensorType inferReshapeInputType(TypedValue<TensorType> input,
     return input.getType();
 
   // The input type must be cast into a tensor with the same rank and all static
-  // dimensions set to 1. This prevents the generation of a tensor.collapse_shape
-  // op that converts a dynamically shaped tensor into a 0D tensor. While such
-  // construct is not incorrect on its own, bufferization cannot properly handle
-  // it at the moment, so we avoid it.
+  // dimensions set to 1. This prevents the generation of a
+  // tensor.collapse_shape op that converts a dynamically shaped tensor into a
+  // 0D tensor. While such construct is not incorrect on its own, bufferization
+  // cannot properly handle it at the moment, so we avoid it.
   SmallVector<int64_t> shape(input.getType().getRank(), 1);
   return input.getType().clone(shape);
 }
@@ -58,29 +58,31 @@ TensorType inferReshapeExpandedType(TensorType inputType,
   int64_t totalSize = inputIsStatic ? inputType.getNumElements() : -1;
 
   // Compute result shape
-  auto resultShape = llvm::map_to_vector(newShape, [&](int64_t size) -> int64_t {
-    // If this is not a placeholder, do not change it
-    if (size >= 0)
-      return size;
-
-    // If we do not know the total size of the tensor, keep this dimension
-    // dynamic in the result shape.
-    if (!inputIsStatic)
-      return ShapedType::kDynamic;
-
-    // Calculate the product of all elements in 'newShape' except for the -1
-    // placeholder, which we discard by negating the result.
-    int64_t totalSizeNoPlaceholder = -std::accumulate(
-        newShape.begin(), newShape.end(), 1, std::multiplies<int64_t>());
-
-    // If there is a 0 component in 'newShape', resolve the placeholder as 0.
-    if (totalSizeNoPlaceholder == 0)
-      return 0;
-
-    // Resolve the placeholder as the quotient between the total tensor size and
-    // the product of all other sizes.
-    return totalSize / totalSizeNoPlaceholder;
-  });
+  auto resultShape =
+      llvm::map_to_vector(newShape, [&](int64_t size) -> int64_t {
+        // If this is not a placeholder, do not change it
+        if (size >= 0)
+          return size;
+
+        // If we do not know the total size of the tensor, keep this dimension
+        // dynamic in the result shape.
+        if (!inputIsStatic)
+          return ShapedType::kDynamic;
+
+        // Calculate the product of all elements in 'newShape' except for the -1
+        // placeholder, which we discard by negating the result.
+        int64_t totalSizeNoPlaceholder = -std::accumulate(
+            newShape.begin(), newShape.end(), 1, std::multiplies<int64_t>());
+
+        // If there is a 0 component in 'newShape', resolve the placeholder as
+        // 0.
+        if (totalSizeNoPlaceholder == 0)
+          return 0;
+
+        // Resolve the placeholder as the quotient between the total tensor size
+        // and the product of all other sizes.
+        return totalSize / totalSizeNoPlaceholder;
+      });
 
   bool resultIsStatic = !ShapedType::isDynamicShape(resultShape);
 
@@ -108,7 +110,8 @@ TensorType inferReshapeCollapsedType(TensorType lhsType, TensorType rhsType) {
   if (lhsShape.empty() || rhsShape.empty())
     return lhsType.clone(ArrayRef<int64_t>{});
 
-  if (ShapedType::isDynamicShape(lhsShape) || ShapedType::isDynamicShape(rhsShape))
+  if (ShapedType::isDynamicShape(lhsShape) ||
+      ShapedType::isDynamicShape(rhsShape))
     return lhsType.clone({ShapedType::kDynamic});
 
   SmallVector<int64_t> intermediateShape;
@@ -150,14 +153,16 @@ TensorType inferReshapeCollapsedType(TensorType lhsType, TensorType rhsType) {
 }
 
 SmallVector<ReassociationExprs>
-createReassociationMapForCollapse(OpBuilder &builder, Type srcType, Type dstType) {
+createReassociationMapForCollapse(OpBuilder &builder, Type srcType,
+                                  Type dstType) {
   auto srcShape = cast<TensorType>(srcType).getShape();
   auto dstShape = cast<TensorType>(dstType).getShape();
 
   if (srcShape.empty() || dstShape.empty())
     return {};
 
-  if (ShapedType::isDynamicShape(srcShape) || ShapedType::isDynamicShape(dstShape)) {
+  if (ShapedType::isDynamicShape(srcShape) ||
+      ShapedType::isDynamicShape(dstShape)) {
     assert(dstShape.size() == 1);
     SmallVector<AffineExpr, 2> exprs;
     for (auto i : llvm::seq<int64_t>(srcShape.size()))
@@ -249,14 +254,16 @@ class ReshapeConverter : public OpConversionPattern<tosa::ReshapeOp> {
     auto collapsedType = inferReshapeCollapsedType(inputType, expandedType);
 
     // Cast input if needed
-    auto castInput = rewriter.createOrFold<tensor::CastOp>(loc, inputType, input);
+    auto castInput =
+        rewriter.createOrFold<tensor::CastOp>(loc, inputType, input);
 
     // Emit collaspe-expand pair
     auto collapsed = createCollapse(rewriter, loc, collapsedType, castInput);
     auto expanded = createExpand(rewriter, loc, expandedType, collapsed);
 
     // Cast to final result type if needed
-    auto result = rewriter.createOrFold<tensor::CastOp>(loc, resultType, expanded);
+    auto result =
+        rewriter.createOrFold<tensor::CastOp>(loc, resultType, expanded);
     rewriter.replaceOp(reshape, result);
     return success();
   }
@@ -350,29 +357,12 @@ class PadConverter : public OpConversionPattern<tosa::PadOp> {
     }
 
     ShapedType inputTy = cast<ShapedType>(input.getType());
-    Type elementTy = inputTy.getElementType();
     int64_t rank = inputTy.getRank();
 
     // Setup the default constantAttr.
 
-    Value padConstant;
-
-    if (padOp.getPadConst()) {
-      padConstant = rewriter.createOrFold<tensor::ExtractOp>(
-          loc, padOp.getPadConst(), ValueRange({}));
-    } else {
-      TypedAttr constantAttr;
-      if (isa<FloatType>(elementTy)) {
-        constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
-      } else if (isa<IntegerType>(elementTy) && !padOp.getInputZpAttr()) {
-        constantAttr = rewriter.getIntegerAttr(elementTy, 0);
-      } else if (isa<IntegerType>(elementTy) && padOp.getInputZpAttr()) {
-        int64_t value = padOp.getInputZpAttr().getInt();
-        constantAttr = rewriter.getIntegerAttr(elementTy, value);
-      }
-      if (constantAttr)
-        padConstant = rewriter.create<arith::ConstantOp>(loc, constantAttr);
-    }
+    Value padConstant = rewriter.createOrFold<tensor::ExtractOp>(
+        loc, padOp.getPadConst(), ValueRange({}));
 
     if (!padConstant) {
       return rewriter.notifyMatchFailure(

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

@@ -175,53 +175,6 @@ void TransposeOp::getCanonicalizationPatterns(RewritePatternSet &results,
   results.add<ConsolidateTransposeOptimization, TransposeIsReshape>(context);
 }
 
-struct MaterializePadValue : public OpRewritePattern<tosa::PadOp> {
-  using OpRewritePattern::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(tosa::PadOp op,
-                                PatternRewriter &rewriter) const override {
-    if (op.getPadConst())
-      return failure();
-
-    auto input = op.getInput1();
-    auto padding = op.getPadding();
-
-    ShapedType inputTy = llvm::cast<ShapedType>(input.getType());
-    Type elementTy = inputTy.getElementType();
-
-    Attribute constantAttr;
-    if (llvm::isa<FloatType>(elementTy)) {
-      constantAttr = rewriter.getFloatAttr(elementTy, 0.0);
-    } else if (llvm::isa<IntegerType>(elementTy) && !op.getInputZpAttr()) {
-      constantAttr = rewriter.getIntegerAttr(elementTy, 0);
-    } else if (llvm::isa<IntegerType>(elementTy) && op.getInputZpAttr()) {
-      int64_t value = op.getInputZpAttr().getInt();
-      constantAttr = rewriter.getIntegerAttr(elementTy, value);
-    }
-
-    if (!constantAttr) {
-      return rewriter.notifyMatchFailure(
-          op,
-          "tosa.pad to linalg lowering encountered an unknown element type");
-    }
-
-    auto denseAttr = DenseElementsAttr::get(
-        RankedTensorType::get({1}, elementTy), constantAttr);
-    auto constantVal = rewriter.create<tosa::ConstOp>(
-        op.getLoc(), denseAttr.getType(), denseAttr);
-
-    rewriter.replaceOpWithNewOp<tosa::PadOp>(
-        op, op.getType(), ValueRange{input, padding, constantVal},
-        op->getAttrs());
-    return success();
-  }
-};
-
-void PadOp::getCanonicalizationPatterns(RewritePatternSet &results,
-                                        MLIRContext *context) {
-  results.add<MaterializePadValue>(context);
-}
-
 struct MaxPool2dIsNoOp : public OpRewritePattern<tosa::MaxPool2dOp> {
   using OpRewritePattern::OpRewritePattern;
 

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

@@ -214,6 +214,23 @@ void mlir::tosa::printTypeOrAttr(OpAsmPrinter &p, Operation *op, TypeAttr type,
   }
 }
 
+// Create a pad-const const tensor with value of `val` of required data-type
+std::optional<Value> mlir::tosa::createPadConstTensor(OpBuilder &builder,
+                                                      Location loc, Value src,
+                                                      int32_t val) {
+  auto const srcType = getElementTypeOrSelf(src);
+  auto const srcElemType = getElementTypeOrSelf(src);
+  auto const padConstType = mlir::RankedTensorType::get({1}, srcType);
+  auto const padConstEType = mlir::RankedTensorType::get({1}, srcElemType);
+  auto const pad_const_attr{
+      llvm::isa<FloatType>(srcElemType)
+          ? DenseElementsAttr::get(padConstEType,
+                                   builder.getFloatAttr(srcElemType, val))
+          : DenseElementsAttr::get(padConstEType,
+                                   builder.getIntegerAttr(srcElemType, val))};
+  return builder.create<tosa::ConstOp>(loc, padConstType, pad_const_attr);
+}
+
 //===----------------------------------------------------------------------===//
 // Tosa utilities.
 //===----------------------------------------------------------------------===//
@@ -679,30 +696,14 @@ static void buildUnaryOpWithQuantInfo(OpBuilder &builder,
 static void buildPadOpWithQuantInfo(OpBuilder &builder, OperationState &result,
                                     Type outputType, Value input,
                                     Value paddings) {
-  result.addOperands({input, paddings});
-  auto quantAttr = buildPadOpQuantizationAttr(builder, input);
+  const Location loc{result.location};
+  int32_t zp{0};
+  auto const quantAttr = buildPadOpQuantizationAttr(builder, input);
   if (quantAttr) {
-    result.addAttribute("input_zp",
-                        builder.getI32IntegerAttr(
-                            static_cast<int32_t>(quantAttr.getInputZp())));
-  }
-  result.types.push_back(outputType);
-}
-
-/// This builder is called on TOSA pad operator when an explicit pad_const
-/// value is passed in. It also optionally constructs quantization_attr.
-static void buildExplicitValuePadOpWithQuantInfo(OpBuilder &builder,
-                                                 OperationState &result,
-                                                 Type outputType, Value input,
-                                                 Value paddings,
-                                                 Value padConst) {
-  result.addOperands({input, paddings, padConst});
-  auto quantAttr = buildPadOpQuantizationAttr(builder, input);
-  if (quantAttr) {
-    result.addAttribute("input_zp",
-                        builder.getI32IntegerAttr(
-                            static_cast<int32_t>(quantAttr.getInputZp())));
+    zp = static_cast<int32_t>(quantAttr.getInputZp());
   }
+  auto const pad_const_op{createPadConstTensor(builder, loc, input, zp)};
+  result.addOperands({input, paddings, pad_const_op.value()});
   result.types.push_back(outputType);
 }
 

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

@@ -148,16 +148,16 @@ class TransposeConvStridedConverter
       return rewriter.notifyMatchFailure(
           op, "zero point must be zero for non-int8 integer types");
 
-    if (weightZpVal != 0) {
-      weight = CreateOpAndInferShape<tosa::PadOp>(
-          rewriter, loc, UnrankedTensorType::get(weightETy), weight,
-          weightPaddingVal, nullptr, rewriter.getI32IntegerAttr(weightZpVal));
-
-    } else {
-      weight = CreateOpAndInferShape<tosa::PadOp>(
-          rewriter, loc, UnrankedTensorType::get(weightETy), weight,
-          weightPaddingVal);
-    }
+    // construct pad_const values from zp values
+    ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
+    Value const inputPadConst =
+        createPadConstTensor(builder, op->getLoc(), input, inputZpVal).value();
+    Value const weightPadConst =
+        createPadConstTensor(builder, op->getLoc(), input, weightZpVal).value();
+
+    weight = CreateOpAndInferShape<tosa::PadOp>(
+        rewriter, loc, UnrankedTensorType::get(weightETy), weight,
+        weightPaddingVal, weightPadConst);
 
     weightTy = cast<ShapedType>(weight.getType());
     weightHeight = weightTy.getDimSize(1);
@@ -169,7 +169,6 @@ class TransposeConvStridedConverter
         stride[0],      weightWidth / stride[1],
         stride[1],      inputChannels};
 
-    ImplicitLocOpBuilder builder(op->getLoc(), rewriter);
     weight = CreateOpAndInferShape<tosa::ReshapeOp>(
         builder, UnrankedTensorType::get(weightETy), weight,
         getTosaConstShape(rewriter, loc, weightReshapeDims0));
@@ -206,15 +205,9 @@ class TransposeConvStridedConverter
     Value inputPaddingVal =
         getTosaConstShape(rewriter, op->getLoc(), inputPadding);
 
-    if (inputZpVal != 0) {
-      input = CreateOpAndInferShape<tosa::PadOp>(
-          rewriter, loc, UnrankedTensorType::get(inputETy), input,
-          inputPaddingVal, nullptr, rewriter.getI32IntegerAttr(inputZpVal));
-    } else {
-      input = CreateOpAndInferShape<tosa::PadOp>(
-          rewriter, loc, UnrankedTensorType::get(inputETy), input,
-          inputPaddingVal);
-    }
+    input = CreateOpAndInferShape<tosa::PadOp>(
+        rewriter, loc, UnrankedTensorType::get(inputETy), input,
+        inputPaddingVal, inputPadConst);
 
     // We use a zero bias as we need to broadcast the bias.
     auto zeroBias = rewriter.create<tosa::ConstOp>(