[Mlir-commits] [mlir] [mlir][tensor] Fix `tensor.pad` to remove newly static values (PR #79938)
Rob Suderman
llvmlistbot at llvm.org
Mon Jan 29 19:39:30 PST 2024
https://github.com/rsuderman updated https://github.com/llvm/llvm-project/pull/79938
>From bc9d42a83049237af43c5dcf5132d60c441b94b9 Mon Sep 17 00:00:00 2001
From: Rob Suderman <rob.suderman at gmail.com>
Date: Mon, 29 Jan 2024 18:30:09 -0800
Subject: [PATCH 1/4] [mlir] Fix `tensor.pad` to remote newly static values
The canonicalization incrementally converts foldable dynamic hi/lo
padding to static hi/lo values. During this canonicalization the
static-fied valued should be removed from the dynamic values.
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 6 +++++-
mlir/test/Dialect/Tensor/canonicalize.mlir | 2 +-
2 files changed, 6 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index b2fe58099b2fb..ff549c697c7fa 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3158,19 +3158,23 @@ struct FoldStaticPadding : public OpRewritePattern<PadOp> {
// Extract the static info from the high and low operands.
SmallVector<int64_t> constOperandsLow;
+ llvm::SmallVector<Value> newLows;
for (auto operand : padTensorOp.getLow()) {
APSInt intOp;
if (!matchPattern(operand, m_ConstantInt(&intOp))) {
constOperandsLow.push_back(ShapedType::kDynamic);
+ newLows.push_back(operand);
continue;
}
constOperandsLow.push_back(intOp.getExtValue());
}
SmallVector<int64_t> constOperandsHigh;
+ llvm::SmallVector<Value> newHighs;
for (auto operand : padTensorOp.getHigh()) {
APSInt intOp;
if (!matchPattern(operand, m_ConstantInt(&intOp))) {
constOperandsHigh.push_back(ShapedType::kDynamic);
+ newHighs.push_back(operand);
continue;
}
constOperandsHigh.push_back(intOp.getExtValue());
@@ -3222,7 +3226,7 @@ struct FoldStaticPadding : public OpRewritePattern<PadOp> {
newOutDims, padTensorOp.getType().getElementType());
auto newOp = rewriter.create<PadOp>(
padTensorOp->getLoc(), newResultType, input, staticLow, staticHigh,
- padTensorOp.getLow(), padTensorOp.getHigh(), padTensorOp.getNofold(),
+ newLows, newLows, padTensorOp.getNofold(),
getPrunedAttributeList(padTensorOp, PadOp::getAttributeNames()));
IRMapping mapper;
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index ed964071358ac..0ee38410400f2 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1361,7 +1361,7 @@ func.func @pad_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
// CHECK-LABEL: func @pad_fold_static(
// CHECK-SAME: %[[INPUT:.*]]: tensor<?x64x?x?xf32>) -> tensor<?x?x?x?xf32> {
// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[PADDING:.*]] = arith.constant 4 : index
+// CHECK-NOT: %[[PADDING:.*]] = arith.constant 4 : index
// CHECK: %[[PADDED:.*]] = tensor.pad %[[INPUT]]
// CHECK-SAME: low[0, 4, 1, 1] high[0, 4, 1, 1] {
// CHECK: ^bb0(%[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index):
>From 98ece5159bf4766e97cca950c659574b1a53bb48 Mon Sep 17 00:00:00 2001
From: Rob Suderman <rob.suderman at gmail.com>
Date: Mon, 29 Jan 2024 19:26:51 -0800
Subject: [PATCH 2/4] fix typo
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index ff549c697c7fa..7b032caa666c3 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3226,7 +3226,7 @@ struct FoldStaticPadding : public OpRewritePattern<PadOp> {
newOutDims, padTensorOp.getType().getElementType());
auto newOp = rewriter.create<PadOp>(
padTensorOp->getLoc(), newResultType, input, staticLow, staticHigh,
- newLows, newLows, padTensorOp.getNofold(),
+ newLows, newHighs, padTensorOp.getNofold(),
getPrunedAttributeList(padTensorOp, PadOp::getAttributeNames()));
IRMapping mapper;
>From 273b72896128fa8c67d94bce89e5b8a976b756e5 Mon Sep 17 00:00:00 2001
From: Rob Suderman <rob.suderman at gmail.com>
Date: Mon, 29 Jan 2024 19:31:52 -0800
Subject: [PATCH 3/4] fix for jakub comments
---
mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 4 ++--
mlir/test/Dialect/Tensor/canonicalize.mlir | 2 +-
2 files changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index 7b032caa666c3..b21e89ae3a571 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -3158,7 +3158,7 @@ struct FoldStaticPadding : public OpRewritePattern<PadOp> {
// Extract the static info from the high and low operands.
SmallVector<int64_t> constOperandsLow;
- llvm::SmallVector<Value> newLows;
+ SmallVector<Value> newLows;
for (auto operand : padTensorOp.getLow()) {
APSInt intOp;
if (!matchPattern(operand, m_ConstantInt(&intOp))) {
@@ -3169,7 +3169,7 @@ struct FoldStaticPadding : public OpRewritePattern<PadOp> {
constOperandsLow.push_back(intOp.getExtValue());
}
SmallVector<int64_t> constOperandsHigh;
- llvm::SmallVector<Value> newHighs;
+ SmallVector<Value> newHighs;
for (auto operand : padTensorOp.getHigh()) {
APSInt intOp;
if (!matchPattern(operand, m_ConstantInt(&intOp))) {
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index 0ee38410400f2..ed964071358ac 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1361,7 +1361,7 @@ func.func @pad_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
// CHECK-LABEL: func @pad_fold_static(
// CHECK-SAME: %[[INPUT:.*]]: tensor<?x64x?x?xf32>) -> tensor<?x?x?x?xf32> {
// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK-NOT: %[[PADDING:.*]] = arith.constant 4 : index
+// CHECK: %[[PADDING:.*]] = arith.constant 4 : index
// CHECK: %[[PADDED:.*]] = tensor.pad %[[INPUT]]
// CHECK-SAME: low[0, 4, 1, 1] high[0, 4, 1, 1] {
// CHECK: ^bb0(%[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index):
>From f7f692eeff17833bb5cf12b23fd693c081bbe3e0 Mon Sep 17 00:00:00 2001
From: Rob Suderman <rob.suderman at gmail.com>
Date: Mon, 29 Jan 2024 19:39:18 -0800
Subject: [PATCH 4/4] refix
---
mlir/test/Dialect/Tensor/canonicalize.mlir | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/test/Dialect/Tensor/canonicalize.mlir b/mlir/test/Dialect/Tensor/canonicalize.mlir
index ed964071358ac..7192a719ceb13 100644
--- a/mlir/test/Dialect/Tensor/canonicalize.mlir
+++ b/mlir/test/Dialect/Tensor/canonicalize.mlir
@@ -1361,7 +1361,7 @@ func.func @pad_same_static_shape(%arg0: tensor<5x6xf32>, %a: index)
// CHECK-LABEL: func @pad_fold_static(
// CHECK-SAME: %[[INPUT:.*]]: tensor<?x64x?x?xf32>) -> tensor<?x?x?x?xf32> {
// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
-// CHECK: %[[PADDING:.*]] = arith.constant 4 : index
+// CHECK-NOT: arith.constant 4 : index
// CHECK: %[[PADDED:.*]] = tensor.pad %[[INPUT]]
// CHECK-SAME: low[0, 4, 1, 1] high[0, 4, 1, 1] {
// CHECK: ^bb0(%[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index):
More information about the Mlir-commits
mailing list