[Mlir-commits] [mlir] 2472c45 - [mlir][tensor] Enhance pack/unpack simplification for identity outer_dims_perm cases. (#77409)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jan 10 08:30:38 PST 2024
Author: Han-Chung Wang
Date: 2024-01-10T08:30:34-08:00
New Revision: 2472c45ba38828ee084360d52705955ff763e5b0
URL: https://github.com/llvm/llvm-project/commit/2472c45ba38828ee084360d52705955ff763e5b0
DIFF: https://github.com/llvm/llvm-project/commit/2472c45ba38828ee084360d52705955ff763e5b0.diff
LOG: [mlir][tensor] Enhance pack/unpack simplification for identity outer_dims_perm cases. (#77409)
They can be simplified to reshape ops if outer_dims_perm is an identity
permutation. The revision adds a `isIdentityPermutation` method to
IndexingUtils.
Added:
Modified:
mlir/include/mlir/Dialect/Utils/IndexingUtils.h
mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt
mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
mlir/lib/Dialect/Utils/IndexingUtils.cpp
mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
Removed:
################################################################################
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 8ab69d8c59b41f..06be017f24b823 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
@@ -39,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();
@@ -75,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..2765d1eb1000da 100644
--- a/mlir/lib/Dialect/Utils/IndexingUtils.cpp
+++ b/mlir/lib/Dialect/Utils/IndexingUtils.cpp
@@ -213,6 +213,13 @@ mlir::invertPermutationVector(ArrayRef<int64_t> permutation) {
return inversion;
}
+bool mlir::isIdentityPermutation(ArrayRef<int64_t> permutation) {
+ for (auto i : llvm::seq<int64_t>(0, permutation.size()))
+ 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..82bfe6fe8689ab 100644
--- a/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
+++ b/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
@@ -37,6 +37,30 @@ 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>
+// 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 +133,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
More information about the Mlir-commits
mailing list