[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.

Author: hanhanW

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);
 

mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt

@@ -27,6 +27,7 @@ add_mlir_dialect_library(MLIRTensorTransforms
   MLIRArithUtils
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
+  MLIRDialectUtils
   MLIRIR
   MLIRLinalgDialect
   MLIRMemRefDialect

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();

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");

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