[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 cfe-commits cfe-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 cfe-commits mailing list