[mlir][Tensor] Add a helper build method for pad operations with constant padding.

[hanchung]

I think we already have a similar method in tensor/Utils.h. Does it fit your needs?

// Return a PadOp that pads `source to `type` size with `pad` value.
// I.e., a block will be created and the `pad` value will be yielded
// directly. If the type passed is nullptr, it is inferred.
PadOp createPadScalarOp(Type type, Value source, Value pad,
                        ArrayRef<OpFoldResult> low, ArrayRef<OpFoldResult> high,
                        bool nofold, Location loc, OpBuilder &builder);

[mravishankar]

In D136493#3875849, @hanchung wrote:

I think we already have a similar method in tensor/Utils.h. Does it fit your needs?

// Return a PadOp that pads `source to `type` size with `pad` value.
// I.e., a block will be created and the `pad` value will be yielded
// directly. If the type passed is nullptr, it is inferred.
PadOp createPadScalarOp(Type type, Value source, Value pad,
                        ArrayRef<OpFoldResult> low, ArrayRef<OpFoldResult> high,
                        bool nofold, Location loc, OpBuilder &builder);

Maybe we should just drop that method....

[mravishankar]

Drop duplicate method in Tensor/Utils.h.

mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td

OpBuilder<(ins "Value":$source, "ValueRange":$low, "ValueRange":$high,
       CArg<"bool", "false">:$nofold,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
     // Build a PadOp with mixed static and dynamic entries and custom
     // result type. If the type passed is nullptr, it is inferred.
     OpBuilder<(ins "Type":$resultType, "Value":$source,
       "ArrayRef<OpFoldResult>":$low, "ArrayRef<OpFoldResult>":$high,
       CArg<"bool", "false">:$nofold,
       CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>,
+    // Build a PadOp with constant padding,  mixed static and dynamic entries
+    // and custom result type. If the type passed is nullptr, it is inferred.
+    OpBuilder<(ins "Type":$resultType, "Value":$source,
+      "ArrayRef<OpFoldResult>":$low, "ArrayRef<OpFoldResult>":$high,
+      "Value":$constantPadValue, CArg<"bool", "false">:$nofold,
+      CArg<"ArrayRef<NamedAttribute>", "{}">:$attrs)>
   ];
 
   let hasCanonicalizer = 1;
   let hasFolder = 1;
   let hasVerifier = 1;
   let hasRegionVerifier = 1;
 }
 

mlir/include/mlir/Dialect/Tensor/Utils/Utils.h

 
 // Return a PadOp that pads `source` to `type` size where the static
 // sizes are assumed to be greater than the dynamic sizes. If `type` has dynamic
 // dimensions the padding width is set to zero. The op performs "high" padding
 // (i.e. it adds trailing padding values until the desired size is met).
 PadOp createPadHighOp(RankedTensorType type, Value source, Value pad,
                       bool nofold, Location loc, OpBuilder &builder);
 
-// Return a PadOp that pads `source to `type` size with `pad` value.
-// I.e., a block will be created and the `pad` value will be yielded
-// directly. If the type passed is nullptr, it is inferred.
-PadOp createPadScalarOp(Type type, Value source, Value pad,
-                        ArrayRef<OpFoldResult> low, ArrayRef<OpFoldResult> high,
-                        bool nofold, Location loc, OpBuilder &builder);
-
 // Creates dim ops for each dynamic dimension of the ranked tensor argument and
 // returns these as values.
 SmallVector<Value> createDynamicDimValues(OpBuilder &b, Location loc,
                                           Value rankedTensor);
 
 // Creates dim ops or constant ops for each dimension of the ranked tensor
 // argument and returns these as values.
 SmallVector<Value> createDimValues(OpBuilder &b, Location loc,
                                    Value rankedTensor);
 
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_UTILS_UTILS_H_

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp

for (int i = 0; i < rank; i++) {
           loc, rewriter.getIndexType(), lowVal);
       highVal = rewriter.createOrFold<arith::IndexCastOp>(
           loc, rewriter.getIndexType(), highVal);
 
       lowValues.push_back(lowVal);
       highValues.push_back(highVal);
     }
 
-    auto newPadOp = tensor::createPadScalarOp(
-        padOp.getType(), input, padConstant, lowValues, highValues,
-        /*nofold=*/false, loc, rewriter);
+    auto newPadOp = rewriter.create<tensor::PadOp>(
+        loc, padOp.getType(), input, lowValues, highValues, padConstant);
 
     rewriter.replaceOp(padOp, newPadOp.getResult());
     return success();
   }
 };
 
 // Tosa argmax lowering represents the ArgMax op as an linalg.indexed_generic
 // op, producing two output buffers.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp

for (int i = 0, s = inputShape.size(); i < s; i++) {
     else
       paddedShape.push_back(inputShape[i] + highPad + lowPad);
     lowIndices.push_back(rewriter.getIndexAttr(lowPad));
     highIndices.push_back(rewriter.getIndexAttr(highPad));
   }
 
   Value padValue = rewriter.create<arith::ConstantOp>(loc, padAttr);
 
-  return tensor::createPadScalarOp(RankedTensorType::get(paddedShape, inputETy),
-                                   input, padValue, lowIndices, highIndices,
-                                   /*nofold=*/false, loc, rewriter)
-      .getResult();
+  return rewriter.create<tensor::PadOp>(
+      loc, RankedTensorType::get(paddedShape, inputETy), input, lowIndices,
+      highIndices, padValue);
 }
 
 static mlir::Value reifyConstantDim(Attribute attr,
                                     ImplicitLocOpBuilder &builder) {
   return builder.createOrFold<arith::IndexCastOp>(
       builder.getIndexType(), builder.create<arith::ConstantOp>(attr));
 }
 

mlir/lib/Dialect/Tensor/IR/TensorOps.cpp

void PadOp::build(OpBuilder &b, OperationState &result, Type resultType,
   // value as well.
   dispatchIndexOpFoldResults(low, dynamicLow, staticLow,
                              ShapedType::kDynamicSize);
   dispatchIndexOpFoldResults(high, dynamicHigh, staticHigh,
                              ShapedType::kDynamicSize);
   if (!resultType) {
     resultType = PadOp::inferResultType(sourceType, staticLow, staticHigh);
   }
+  assert(resultType.isa<RankedTensorType>());
   build(b, result, resultType, source, dynamicLow, dynamicHigh,
         b.getI64ArrayAttr(staticLow), b.getI64ArrayAttr(staticHigh),
         nofold ? b.getUnitAttr() : UnitAttr());
   result.addAttributes(attrs);
 }
 
+void PadOp::build(OpBuilder &b, OperationState &result, Type resultType,
+                  Value source, ArrayRef<OpFoldResult> low,
+                  ArrayRef<OpFoldResult> high, Value constantPadValue,
+                  bool nofold, ArrayRef<NamedAttribute> attrs) {
+  build(b, result, resultType, source, low, high, nofold, attrs);
+
+  // Add a region and a block to yield the pad value.
+  Region *region = result.regions[0].get();
+  int sourceRank = source.getType().cast<RankedTensorType>().getRank();
+  SmallVector<Type> blockArgTypes(sourceRank, b.getIndexType());
+  SmallVector<Location> blockArgLocs(sourceRank, result.location);
+
+  // `builder.createBlock` changes the insertion point within the block. Create
+  // a guard to reset the insertion point of the builder after it is destroyed.
+  OpBuilder::InsertionGuard guard(b);
+  b.createBlock(region, region->end(), blockArgTypes, blockArgLocs);
+  b.create<tensor::YieldOp>(result.location, constantPadValue);
+}
+
 llvm::SmallBitVector PadOp::getPaddedDims() {
   llvm::SmallBitVector paddedDims(getSourceType().getRank());
   auto extractPaddedDims = [&](ArrayRef<OpFoldResult> paddingWidths) {
     for (const auto &en : enumerate(paddingWidths))
       if (getConstantIntValue(en.value()) != static_cast<int64_t>(0))
         paddedDims.set(en.index());
   };
   extractPaddedDims(getMixedLowPad());

mlir/lib/Dialect/Tensor/Utils/Utils.cpp

 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
-PadOp mlir::tensor::createPadScalarOp(Type type, Value source, Value pad,
-                                      ArrayRef<OpFoldResult> low,
-                                      ArrayRef<OpFoldResult> high, bool nofold,
-                                      Location loc, OpBuilder &builder) {
-  auto padTensorOp =
-      builder.create<PadOp>(loc, type, source, low, high, nofold);
-  int rank = padTensorOp.getResultType().getRank();
-  SmallVector<Type> blockArgTypes(rank, builder.getIndexType());
-  SmallVector<Location> blockArgLocs(rank, loc);
-  auto &region = padTensorOp.getRegion();
-  // `builder.createBlock` changes the insertion point within the block. Create
-  // a guard to reset the insertion point of the builder after it is destroyed.
-  OpBuilder::InsertionGuard guard(builder);
-  builder.createBlock(&region, region.end(), blockArgTypes, blockArgLocs);
-  builder.create<YieldOp>(loc, pad);
-  return padTensorOp;
-}
-
 PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
                                     Value pad, bool nofold, Location loc,
                                     OpBuilder &b) {
   auto zero = b.createOrFold<arith::ConstantIndexOp>(loc, 0);
   SmallVector<OpFoldResult> low(type.getRank(), zero);
   SmallVector<OpFoldResult> high(type.getRank(), zero);
   for (const auto &en : enumerate(type.getShape())) {
     // Pad only the static dimensions of the result tensor type.
     if (ShapedType::isDynamic(en.value()))
       continue;
     // Compute the padding width.
     AffineExpr d0;
     bindDims(b.getContext(), d0);
     auto dimOp = b.createOrFold<tensor::DimOp>(loc, source, en.index());
     high[en.index()] =
         makeComposedAffineApply(b, loc, en.value() - d0, {dimOp}).getResult();
   }
-  return createPadScalarOp(type, source, pad, low, high, nofold, loc, b);
+  return b.create<PadOp>(loc, type, source, low, high, pad, nofold);
 }
 
 SmallVector<Value> mlir::tensor::createDynamicDimValues(OpBuilder &b,
                                                         Location loc,
                                                         Value rankedTensor) {
   auto tensorTy = rankedTensor.getType().cast<RankedTensorType>();
   SmallVector<Value> dynamicDims;
   for (const auto &en : llvm::enumerate(tensorTy.getShape())) {