[Mlir-commits] [mlir] [mlir][linalg] Add PackOp canonicalization pattern (PR #215785)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Hendrik Klug (HendrikKlug-synthara)

<details>
<summary>Changes</summary>

Fold pack(empty) to the destination tensor if no padding value is provided.

---
Full diff: https://github.com/llvm/llvm-project/pull/215785.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp (+7) 
- (modified) mlir/test/Dialect/Linalg/canonicalize.mlir (+28) 


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 170e1edf8a55d..aacaf7d3c3e91 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -6008,6 +6008,13 @@ LogicalResult PackOp::canonicalize(PackOp packOp, PatternRewriter &rewriter) {
     }
   }
 
+  // Fold pack(empty) to the destination tensor if no padding value is provided.
+  if (packOp.getSource().getDefiningOp<tensor::EmptyOp>() &&
+      !packOp.getPaddingValue()) {
+    rewriter.replaceOp(packOp, packOp.getDest());
+    return success();
+  }
+
   // Fold optional PaddingValue operand away if padding is not needed.
   if (packOp.getPaddingValue() && paddingIsNotNeeded(packOp)) {
     rewriter.startOpModification(packOp);
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index bb11ce0d4dfb8..2a37a843aef41 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1535,6 +1535,34 @@ func.func @recursive_effect(%arg : tensor<1xf32>) {
 // linalg.pack
 //===----------------------------------------------------------------------===//
 
+// CHECK-LABEL: func @fold_pack_empty
+//   CHECK-NOT: linalg.pack
+//       CHECK: %[[DEST:.+]] = tensor.empty() : tensor<4x8x8x32xf32>
+//       CHECK: return %[[DEST]]
+func.func @fold_pack_empty() -> tensor<4x8x8x32xf32> {
+  %src = tensor.empty() : tensor<64x128xf32>
+  %dest = tensor.empty() : tensor<4x8x8x32xf32>
+  %0 = linalg.pack %src outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
+    inner_tiles = [8, 32] into %dest : tensor<64x128xf32> -> tensor<4x8x8x32xf32>
+  return %0 : tensor<4x8x8x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @negative_fold_pack_empty_padded
+//       CHECK: linalg.pack
+func.func @negative_fold_pack_empty_padded() -> tensor<4x8x8x32xf32> {
+  %cst = arith.constant 0.0 : f32
+  %src = tensor.empty() : tensor<63x127xf32>
+  %dest = tensor.empty() : tensor<4x8x8x32xf32>
+  %0 = linalg.pack %src padding_value(%cst : f32)
+    outer_dims_perm = [1, 0] inner_dims_pos = [0, 1]
+    inner_tiles = [8, 32] into %dest : tensor<63x127xf32> -> tensor<4x8x8x32xf32>
+  return %0 : tensor<4x8x8x32xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @fold_pack_constant_splat
 //   CHECK-NOT: linalg.pack
 //       CHECK: arith.constant dense<1.000000e-01> : tensor<4x8x8x32xf32>

``````````

</details>


https://github.com/llvm/llvm-project/pull/215785