[Mlir-commits] [mlir] 94b8469 - [mlir][Tensor] Add a helper build method for pad operations with constant padding.
Mahesh Ravishankar
llvmlistbot at llvm.org
Mon Oct 24 11:12:10 PDT 2022
Author: Mahesh Ravishankar
Date: 2022-10-24T18:11:53Z
New Revision: 94b8469a88bca49af8a71b6668eece1bbd48004b
URL: https://github.com/llvm/llvm-project/commit/94b8469a88bca49af8a71b6668eece1bbd48004b
DIFF: https://github.com/llvm/llvm-project/commit/94b8469a88bca49af8a71b6668eece1bbd48004b.diff
LOG: [mlir][Tensor] Add a helper build method for pad operations with constant padding.
Drop the `createPadScalarOp` from Utils.h since it is a duplicate of
the `build` method added here.
Differential Revision: https://reviews.llvm.org/D136493
Added:
Modified:
mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
mlir/lib/Dialect/Tensor/Utils/Utils.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
index 00887566812ef..f57852bc6c30b 100644
--- a/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
+++ b/mlir/include/mlir/Dialect/Tensor/IR/TensorOps.td
@@ -1298,6 +1298,12 @@ def Tensor_PadOp : Tensor_Op<"pad", [
"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;
diff --git a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
index c7547c050c90b..66035efa7f36e 100644
--- a/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
+++ b/mlir/include/mlir/Dialect/Tensor/Utils/Utils.h
@@ -21,13 +21,6 @@ namespace tensor {
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,
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
index 178b4b1f959f1..04cd00f61aa1d 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
@@ -2000,9 +2000,8 @@ class PadConverter : public OpRewritePattern<tosa::PadOp> {
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();
diff --git a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
index f08ac1908cf9a..b7c817bb60d7a 100644
--- a/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
+++ b/mlir/lib/Conversion/TosaToLinalg/TosaToLinalgNamed.cpp
@@ -58,10 +58,9 @@ static mlir::Value applyPad(Location loc, Value input, ArrayRef<int64_t> pad,
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,
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 6fa46c71c386c..282a342df1d8a 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2475,12 +2475,32 @@ void PadOp::build(OpBuilder &b, OperationState &result, Type resultType,
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) {
diff --git a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
index c7d04b2d0aa42..06c344be0a06e 100644
--- a/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
+++ b/mlir/lib/Dialect/Tensor/Utils/Utils.cpp
@@ -18,24 +18,6 @@
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 ®ion = 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(®ion, 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) {
@@ -53,7 +35,7 @@ PadOp mlir::tensor::createPadHighOp(RankedTensorType type, Value source,
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,
More information about the Mlir-commits
mailing list