[flang-commits] [llvm] [lldb] [clang-tools-extra] [clang] [libunwind] [libc] [compiler-rt] [flang] [lld] [libclc] [mlir] [libcxx] [mlir][tensor] Enhance pack/unpack simplification for identity outer_dims_perm cases. (PR #77409)
Han-Chung Wang via flang-commits
flang-commits at lists.llvm.org
Wed Jan 10 07:00:31 PST 2024
https://github.com/hanhanW updated https://github.com/llvm/llvm-project/pull/77409
>From e74b859897cdf1b1effbbd48a4e5376a231f7132 Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Mon, 8 Jan 2024 20:17:30 -0800
Subject: [PATCH 1/2] [mlir][tensor] Enhance pack/unpack simplication patterns
for identity outer_dims_perm cases.
They can be simplified to reshape ops if outer_dims_perm is an identity
permutation. The revision adds a `isIdentityPermutation` method to
IndexingUtils.
---
.../mlir/Dialect/Utils/IndexingUtils.h | 3 +++
.../Dialect/Tensor/Transforms/CMakeLists.txt | 1 +
.../Transforms/PackAndUnpackPatterns.cpp | 17 +++++++++----
mlir/lib/Dialect/Utils/IndexingUtils.cpp | 8 +++++++
.../Dialect/Tensor/simplify-pack-unpack.mlir | 24 +++++++++++++++++++
5 files changed, 48 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Utils/IndexingUtils.h b/mlir/include/mlir/Dialect/Utils/IndexingUtils.h
index f51a8b28b7548e..2453d841f633e4 100644
--- a/mlir/include/mlir/Dialect/Utils/IndexingUtils.h
+++ b/mlir/include/mlir/Dialect/Utils/IndexingUtils.h
@@ -228,6 +228,9 @@ void applyPermutationToVector(SmallVector<T, N> &inVec,
/// Helper method to apply to inverse a permutation.
SmallVector<int64_t> invertPermutationVector(ArrayRef<int64_t> permutation);
+/// Returns true if `permutation` is an identity permutation.
+bool isIdentityPermutation(ArrayRef<int64_t> permutation);
+
/// Method to check if an interchange vector is a permutation.
bool isPermutationVector(ArrayRef<int64_t> interchange);
diff --git a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
index cbc0d499d9d52c..c6ef6ed86e0d9d 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
@@ -27,6 +27,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
MLIRArithUtils
MLIRBufferizationDialect
MLIRBufferizationTransforms
+ MLIRDialectUtils
MLIRIR
MLIRLinalgDialect
MLIRMemRefDialect
diff --git a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
index cfd838e85c1b80..2853cb8fe77a3b 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
@@ -9,6 +9,7 @@
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/IR/PatternMatch.h"
#include "llvm/Support/Debug.h"
@@ -38,8 +39,12 @@ struct SimplifyPackToExpandShape : public OpRewritePattern<PackOp> {
if (packOp.getPaddingValue())
return rewriter.notifyMatchFailure(packOp, "expects no padding value");
- if (!packOp.getOuterDimsPerm().empty())
- return rewriter.notifyMatchFailure(packOp, "expects no outer_dims_perm");
+ auto outerDimsPerm = packOp.getOuterDimsPerm();
+ if (!outerDimsPerm.empty() && !isIdentityPermutation(outerDimsPerm)) {
+ return rewriter.notifyMatchFailure(
+ packOp,
+ "expects outer_dims_perm is empty or an identity permutation");
+ }
RankedTensorType sourceType = packOp.getSourceType();
RankedTensorType destType = packOp.getDestType();
@@ -74,9 +79,11 @@ struct SimplifyUnPackToCollapseShape : public OpRewritePattern<UnPackOp> {
LogicalResult matchAndRewrite(UnPackOp unpackOp,
PatternRewriter &rewriter) const override {
- if (!unpackOp.getOuterDimsPerm().empty()) {
- return rewriter.notifyMatchFailure(unpackOp,
- "expects no outer_dims_perm");
+ auto outerDimsPerm = unpackOp.getOuterDimsPerm();
+ if (!outerDimsPerm.empty() && !isIdentityPermutation(outerDimsPerm)) {
+ return rewriter.notifyMatchFailure(
+ unpackOp,
+ "expects outer_dims_perm is empty or an identity permutation");
}
RankedTensorType sourceType = unpackOp.getSourceType();
diff --git a/mlir/lib/Dialect/Utils/IndexingUtils.cpp b/mlir/lib/Dialect/Utils/IndexingUtils.cpp
index bb8a0d5912d7c1..f3de454dc4b81a 100644
--- a/mlir/lib/Dialect/Utils/IndexingUtils.cpp
+++ b/mlir/lib/Dialect/Utils/IndexingUtils.cpp
@@ -213,6 +213,14 @@ mlir::invertPermutationVector(ArrayRef<int64_t> permutation) {
return inversion;
}
+bool mlir::isIdentityPermutation(ArrayRef<int64_t> permutation) {
+ int n = permutation.size();
+ for (int i = 0; i < n; ++i)
+ if (permutation[i] != i)
+ return false;
+ return true;
+}
+
bool mlir::isPermutationVector(ArrayRef<int64_t> interchange) {
assert(llvm::all_of(interchange, [](int64_t s) { return s >= 0; }) &&
"permutation must be non-negative");
diff --git a/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir b/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
index b78ab9bb3fd87e..ffbb2278a2e327 100644
--- a/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
+++ b/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
@@ -37,6 +37,18 @@ func.func @single_last_inner_dim_packing(%arg0: tensor<5x256xf32>) -> tensor<5x8
// -----
+// CHECK-LABEL: func.func @single_last_inner_dim_packing_with_identity_outer_dims_perm(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<5x256xf32>)
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2]] : tensor<5x256xf32> into tensor<5x8x32xf32>
+// CHECK: return %[[EXPANDED]] : tensor<5x8x32xf32>
+func.func @single_last_inner_dim_packing_with_identity_outer_dims_perm(%arg0: tensor<5x256xf32>) -> tensor<5x8x32xf32> {
+ %empty = tensor.empty() : tensor<5x8x32xf32>
+ %0 = tensor.pack %arg0 outer_dims_perm = [0, 1] inner_dims_pos = [1] inner_tiles = [32] into %empty : tensor<5x256xf32> -> tensor<5x8x32xf32>
+ return %0 : tensor<5x8x32xf32>
+}
+
+// -----
+
// CHECK-LABEL: func.func @packing_with_outer_dims_perm(
// CHECK-NOT: tensor.expand_shape
// CHECK: tensor.pack
@@ -109,6 +121,18 @@ func.func @single_last_inner_dim_unpacking(%arg0: tensor<5x8x32xf32>) -> tensor<
// -----
+// CHECK-LABEL: func.func @single_last_inner_dim_unpacking_with_identity_outer_dims_perm(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<5x8x32xf32>)
+// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2]] : tensor<5x8x32xf32> into tensor<5x256xf32>
+// CHECK: return %[[COLLAPSED]] : tensor<5x256xf32>
+func.func @single_last_inner_dim_unpacking_with_identity_outer_dims_perm(%arg0: tensor<5x8x32xf32>) -> tensor<5x256xf32> {
+ %empty = tensor.empty() : tensor<5x256xf32>
+ %0 = tensor.unpack %arg0 outer_dims_perm = [0, 1] inner_dims_pos = [1] inner_tiles = [32] into %empty : tensor<5x8x32xf32> -> tensor<5x256xf32>
+ return %0 : tensor<5x256xf32>
+}
+
+// -----
+
// CHECK-LABEL: func.func @unpacking_with_outer_dims_perm(
// CHECK-NOT: tensor.collpase_shape
// CHECK: tensor.unpack
>From ab99dc5d6cc979a57c6f93f4c7c00c4311b5cd59 Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Wed, 10 Jan 2024 23:00:02 +0800
Subject: [PATCH 2/2] address comments from chelini
---
mlir/lib/Dialect/Utils/IndexingUtils.cpp | 3 +--
mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir | 12 ++++++++++++
2 files changed, 13 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Utils/IndexingUtils.cpp b/mlir/lib/Dialect/Utils/IndexingUtils.cpp
index f3de454dc4b81a..2765d1eb1000da 100644
--- a/mlir/lib/Dialect/Utils/IndexingUtils.cpp
+++ b/mlir/lib/Dialect/Utils/IndexingUtils.cpp
@@ -214,8 +214,7 @@ mlir::invertPermutationVector(ArrayRef<int64_t> permutation) {
}
bool mlir::isIdentityPermutation(ArrayRef<int64_t> permutation) {
- int n = permutation.size();
- for (int i = 0; i < n; ++i)
+ for (auto i : llvm::seq<int64_t>(0, permutation.size()))
if (permutation[i] != i)
return false;
return true;
diff --git a/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir b/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
index ffbb2278a2e327..82bfe6fe8689ab 100644
--- a/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
+++ b/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
@@ -37,6 +37,18 @@ func.func @single_last_inner_dim_packing(%arg0: tensor<5x256xf32>) -> tensor<5x8
// -----
+// CHECK-LABEL: func.func @pack_1d_with_outer_dims_perm(
+// CHECK-SAME: %[[ARG0:.+]]: tensor<64xf32>)
+// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1]] : tensor<64xf32> into tensor<2x32xf32>
+// CHECK: return %[[EXPANDED]] : tensor<2x32xf32>
+func.func @pack_1d_with_outer_dims_perm(%arg0: tensor<64xf32>) -> tensor<2x32xf32> {
+ %empty = tensor.empty() : tensor<2x32xf32>
+ %pack = tensor.pack %arg0 outer_dims_perm = [0] inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<64xf32> -> tensor<2x32xf32>
+ return %pack : tensor<2x32xf32>
+}
+
+// -----
+
// CHECK-LABEL: func.func @single_last_inner_dim_packing_with_identity_outer_dims_perm(
// CHECK-SAME: %[[ARG0:.+]]: tensor<5x256xf32>)
// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2]] : tensor<5x256xf32> into tensor<5x8x32xf32>
More information about the flang-commits
mailing list