[Mlir-commits] [mlir] [mlir][tensor] Enhance SimplifyUnPackToCollapseShape for unit dim cases. (PR #79262)

[Han-Chung Wang]

https://github.com/hanhanW updated https://github.com/llvm/llvm-project/pull/79262

>From 2090a36b4dff50c365ba1d9faf33adac2c427189 Mon Sep 17 00:00:00 2001
From: hanhanW <hanhan0912 at gmail.com>
Date: Wed, 24 Jan 2024 16:02:45 +0800
Subject: [PATCH] [mlir][tensor] Enhance SimplifyUnPackToCollapseShape for unit
 dim cases.

---
 .../Transforms/PackAndUnpackPatterns.cpp      | 34 +++++++++++--
 .../Dialect/Tensor/simplify-pack-unpack.mlir  | 51 +++++++++++++++++++
 2 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
index b404543ddef867e..bcb00f02d98b55a 100644
--- a/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
+++ b/mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp
@@ -117,8 +117,9 @@ struct SimplifyUnPackToCollapseShape : public OpRewritePattern<UnPackOp> {
                                                     operand, reassociation);
   }
 
-  LogicalResult matchAndRewrite(UnPackOp unpackOp,
-                                PatternRewriter &rewriter) const override {
+  /// Returns success() if it is unpacking on the innermost dimension.
+  LogicalResult isUnpackOnInnerMostDim(RewriterBase &rewriter,
+                                       UnPackOp unpackOp) const {
     auto outerDimsPerm = unpackOp.getOuterDimsPerm();
     if (!outerDimsPerm.empty() && !isIdentityPermutation(outerDimsPerm)) {
       return rewriter.notifyMatchFailure(
@@ -134,9 +135,36 @@ struct SimplifyUnPackToCollapseShape : public OpRewritePattern<UnPackOp> {
     ArrayRef<int64_t> dimsPos = unpackOp.getInnerDimsPos();
     if (dimsPos.size() != 1 || (dimsPos[0] + 1 != destType.getRank())) {
       return rewriter.notifyMatchFailure(
-          unpackOp, "expects unpacking at the innermost dimension");
+          unpackOp, "expects unpacking on the innermost dimension");
     }
 
+    return success();
+  }
+
+  /// Returns success() if it unpacks a 1D or 2D source operand to a 1D
+  /// destination.
+  LogicalResult isUnpackTo1DDest(RewriterBase &rewriter,
+                                 UnPackOp unpackOp) const {
+    ArrayRef<int64_t> destShape = unpackOp.getDestType().getShape();
+    if (getNumGtOneDims(destShape) > 1)
+      return failure();
+
+    SmallVector<int64_t> innerTiles = unpackOp.getStaticTiles();
+    if (getNumGtOneDims(innerTiles) > 1)
+      return failure();
+
+    return success();
+  }
+
+  LogicalResult matchAndRewrite(UnPackOp unpackOp,
+                                PatternRewriter &rewriter) const override {
+    if (failed(isUnpackOnInnerMostDim(rewriter, unpackOp)) &&
+        failed(isUnpackTo1DDest(rewriter, unpackOp))) {
+      return failure();
+    }
+
+    RankedTensorType sourceType = unpackOp.getSourceType();
+    RankedTensorType destType = unpackOp.getDestType();
     auto reassociation =
         getReassociationIndicesForReshape(sourceType, destType);
     if (!reassociation)
diff --git a/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir b/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
index 859eb5ee2a7061d..9948c0246e6ed69 100644
--- a/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
+++ b/mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir
@@ -215,3 +215,54 @@ func.func @single_first_inner_dim_unpacking(%arg0: tensor<8x5x32xf32>) -> tensor
   %0 = tensor.unpack %arg0 inner_dims_pos = [0] inner_tiles = [32] into %empty : tensor<8x5x32xf32> -> tensor<256x5xf32>
   return %0 : tensor<256x5xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_1x32x1x1_to_1x32
+// CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK:         %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
+// CHECK:         return %[[COLLAPSED]]
+func.func @unpack_1x32x1x1_to_1x32(%arg0 : tensor<1x32x1x1xf32>) -> tensor<1x32xf32> {
+  %empty = tensor.empty() : tensor<1x32xf32>
+  %unpack = tensor.unpack %arg0 inner_dims_pos = [0, 1] inner_tiles = [1, 1] into %empty
+    : tensor<1x32x1x1xf32> -> tensor<1x32xf32>
+  return %unpack : tensor<1x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_1x2x1x16_to_1x32
+// CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK:         %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
+// CHECK:         return %[[COLLAPSED]]
+func.func @unpack_1x2x1x16_to_1x32(%arg0 : tensor<1x2x1x16xf32>) -> tensor<1x32xf32> {
+  %empty = tensor.empty() : tensor<1x32xf32>
+  %unpack = tensor.unpack %arg0 outer_dims_perm = [0, 1] inner_dims_pos = [0, 1] inner_tiles = [1, 16] into %empty
+    : tensor<1x2x1x16xf32> -> tensor<1x32xf32>
+  return %unpack : tensor<1x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_16x1x2x1_to_32x1
+// CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK:         %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]
+// CHECK:         return %[[COLLAPSED]]
+func.func @unpack_16x1x2x1_to_32x1(%arg0 : tensor<1x16x2x1xf32>) -> tensor<32x1xf32> {
+  %empty = tensor.empty() : tensor<32x1xf32>
+  %unpack = tensor.unpack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 1] into %empty
+    : tensor<1x16x2x1xf32> -> tensor<32x1xf32>
+  return %unpack : tensor<32x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @unpack_16x1x1x2_to_32x1
+// CHECK-NOT:     tensor.collapse_shape
+// CHECK:         tensor.unpack
+func.func @unpack_16x1x1x2_to_32x1(%arg0 : tensor<16x1x1x2xf32>) -> tensor<32x1xf32> {
+  %empty = tensor.empty() : tensor<32x1xf32>
+  %unpack = tensor.unpack %arg0 inner_dims_pos = [1, 0] inner_tiles = [1, 2] into %empty
+    : tensor<16x1x1x2xf32> -> tensor<32x1xf32>
+  return %unpack : tensor<32x1xf32>
+}