[mlir][tensor] Enhance SimplifyPackToExpandShape for unit dim cases. (#79247)

Progress on iree-org/iree#16181

Author: hanhanW

mlir/lib/Dialect/Tensor/Transforms/PackAndUnpackPatterns.cpp

@@ -22,6 +22,12 @@ static bool areAllConstantIntValue(ArrayRef<OpFoldResult> ofrs, int64_t value) {
       ofrs, [&](OpFoldResult ofr) { return isConstantIntValue(ofr, value); });
 }
 
+/// Returns the number of shape sizes that is either dynamic or greater than 1.
+static int64_t getNumGtOneDims(ArrayRef<int64_t> shape) {
+  return llvm::count_if(
+      shape, [](int64_t v) { return ShapedType::isDynamic(v) || v > 1; });
+}
+
 /// Packing one-dimensional tensor can be expressed as an expand shape op.
 struct SimplifyPackToExpandShape : public OpRewritePattern<PackOp> {
   using OpRewritePattern<PackOp>::OpRewritePattern;
@@ -34,26 +40,60 @@ struct SimplifyPackToExpandShape : public OpRewritePattern<PackOp> {
                                                   reassociation);
   }
 
-  LogicalResult matchAndRewrite(PackOp packOp,
-                                PatternRewriter &rewriter) const override {
-    if (packOp.getPaddingValue())
-      return rewriter.notifyMatchFailure(packOp, "expects no padding value");
-
+  /// Returns success() if it is only packing on the innermost dimension.
+  LogicalResult isPackOnInnerMostDim(RewriterBase &rewriter,
+                                     PackOp packOp) const {
     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();
+    int64_t srcRank = packOp.getSourceRank();
     ArrayRef<int64_t> dimsPos = packOp.getInnerDimsPos();
-    if (dimsPos.size() != 1 || (dimsPos[0] + 1 != sourceType.getRank())) {
+    if (dimsPos.size() != 1 || (dimsPos[0] + 1 != srcRank)) {
       return rewriter.notifyMatchFailure(
           packOp, "expects packing at the innermost dimension");
     }
+    return success();
+  }
+
+  /// Returns success() if there is only 1 dimension size in source being
+  /// greater than 1 and packing only happens on the dimension. It assumes that
+  /// the pack op does not have padding value.
+  LogicalResult isPack1DSrc(RewriterBase &rewriter, PackOp packOp) const {
+    assert(!packOp.getPaddingValue() &&
+           "expect the op does not have padding value.");
+    ArrayRef<int64_t> srcShape = packOp.getSourceType().getShape();
+    if (getNumGtOneDims(srcShape) > 1) {
+      return rewriter.notifyMatchFailure(
+          packOp, "expects source to have at most one non-unit dims");
+    }
 
+    // The pack op does not have padding value. Non-unit inner tile size must be
+    // be used by the non-unit dimension.
+    SmallVector<int64_t> innerTiles = packOp.getStaticTiles();
+    if (getNumGtOneDims(innerTiles) > 1) {
+      return rewriter.notifyMatchFailure(
+          packOp, "expects at most one non-unit inner tiles");
+    }
+
+    return success();
+  }
+
+  LogicalResult matchAndRewrite(PackOp packOp,
+                                PatternRewriter &rewriter) const override {
+    if (packOp.getPaddingValue())
+      return rewriter.notifyMatchFailure(packOp, "expects no padding value");
+
+    if (failed(isPackOnInnerMostDim(rewriter, packOp)) &&
+        failed(isPack1DSrc(rewriter, packOp))) {
+      return failure();
+    }
+
+    RankedTensorType sourceType = packOp.getSourceType();
+    RankedTensorType destType = packOp.getDestType();
     auto reassociation =
         getReassociationIndicesForReshape(sourceType, destType);
     if (!reassociation)

mlir/test/Dialect/Tensor/simplify-pack-unpack.mlir

@@ -83,6 +83,57 @@ func.func @single_first_inner_dim_packing(%arg0: tensor<256x5xf32>) -> tensor<8x
 
 // -----
 
+// CHECK-LABEL: func.func @pack_1x32_to_1x32x1x1
+// CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK:         %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
+// CHECK:         return %[[EXPANDED]]
+func.func @pack_1x32_to_1x32x1x1(%arg0 : tensor<1x32xf32>) -> tensor<1x32x1x1xf32> {
+  %empty = tensor.empty() : tensor<1x32x1x1xf32>
+  %pack = tensor.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [1, 1] into %empty
+    : tensor<1x32xf32> -> tensor<1x32x1x1xf32>
+  return %pack : tensor<1x32x1x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_1x32_to_1x16x1x2
+// CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK:         %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2, 3]]
+// CHECK:         return %[[EXPANDED]]
+func.func @pack_1x32_to_1x16x1x2(%arg0 : tensor<1x32xf32>) -> tensor<1x16x1x2xf32> {
+  %empty = tensor.empty() : tensor<1x16x1x2xf32>
+  %pack = tensor.pack %arg0 inner_dims_pos = [0, 1] inner_tiles = [1, 2] into %empty
+    : tensor<1x32xf32> -> tensor<1x16x1x2xf32>
+  return %pack : tensor<1x16x1x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_32x1_to_16x1x2x1
+// CHECK-SAME:    %[[ARG0:[0-9a-zA-Z]+]]
+// CHECK:         %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]
+// CHECK:         return %[[EXPANDED]]
+func.func @pack_32x1_to_16x1x2x1(%arg0 : tensor<32x1xf32>) -> tensor<1x16x2x1xf32> {
+  %empty = tensor.empty() : tensor<1x16x2x1xf32>
+  %pack = tensor.pack %arg0 outer_dims_perm = [1, 0] inner_dims_pos = [0, 1] inner_tiles = [2, 1] into %empty
+    : tensor<32x1xf32> -> tensor<1x16x2x1xf32>
+  return %pack : tensor<1x16x2x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func @pack_32x1_to_16x1x1x2
+// CHECK-NOT:     tensor.expand_shape
+// CHECK:         tensor.pack
+func.func @pack_32x1_to_16x1x1x2(%arg0 : tensor<32x1xf32>) -> tensor<16x1x1x2xf32> {
+  %empty = tensor.empty() : tensor<16x1x1x2xf32>
+  %pack = tensor.pack %arg0 inner_dims_pos = [1, 0] inner_tiles = [1, 2] into %empty
+    : tensor<32x1xf32> -> tensor<16x1x1x2xf32>
+  return %pack : tensor<16x1x1x2xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func.func @unpack_1d_to_collapse
 // CHECK-SAME:    %[[ARG0:.+]]: tensor<8x32xf32>)
 // CHECK:         %[[COLLAPSED:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1]] : tensor<8x32xf32> into tensor<256xf32>