[Mlir-commits] [mlir] [mlir][linalg] add more pattern to fold pack op padding_value. (PR #198468)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-linalg

Author: Jerry Shih (JerryShih)

<details>
<summary>Changes</summary>

No padding is needed for unit tile size.

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


2 Files Affected:

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


``````````diff
diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 6ebf147394356..485dafffd6de1 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5849,11 +5849,14 @@ static bool haveSameTiles(PackOp packOp, UnPackOp unPackOp) {
 /// Returns true if the pack op does not need a padding value.
 static bool paddingIsNotNeeded(PackOp op) {
   auto srcType = op.getSourceType();
-  if (llvm::any_of(op.getInnerDimsPos(),
-                   [&](int64_t pos) { return srcType.isDynamicDim(pos); }))
-    return false;
-  if (ShapedType::isDynamicShape(op.getStaticInnerTiles()))
+  auto innerDimsPos = op.getInnerDimsPos();
+  auto innerTiles = op.getStaticInnerTiles();
+  if (ShapedType::isDynamicShape(innerTiles))
     return false;
+  for (auto [pos, tileSize] : llvm::zip_equal(innerDimsPos, innerTiles)) {
+    if (srcType.isDynamicDim(pos) && tileSize != 1)
+      return false;
+  }
   return !PackOp::requirePaddingValue(
       srcType.getShape(), op.getInnerDimsPos(), op.getDestType().getShape(),
       op.getOuterDimsPerm(), op.getMixedTiles());
diff --git a/mlir/test/Dialect/Linalg/canonicalize.mlir b/mlir/test/Dialect/Linalg/canonicalize.mlir
index 8896cf1ab7e2d..3318c13da1776 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1569,6 +1569,24 @@ func.func @fold_padding_value_pack(%arg0: tensor<1200x500000xf32>) -> tensor<312
 
 // -----
 
+func.func @fold_padding_value_pack_dynamic_with_unit_tile_size(%arg0: tensor<?x500000xf32>) -> tensor<31250x?x16x1xf32> {
+  %cst = arith.constant 0.000000e+00 : f32
+  %c0 = arith.constant 0 : index
+  %dim0 = tensor.dim %arg0, %c0 : tensor<?x500000xf32>
+  %0 = tensor.empty(%dim0) : tensor<31250x?x16x1xf32>
+  %pack = linalg.pack %arg0
+    padding_value(%cst : f32)
+    outer_dims_perm = [1, 0]
+    inner_dims_pos = [1, 0]
+    inner_tiles = [16, 1]
+    into %0 : tensor<?x500000xf32> -> tensor<31250x?x16x1xf32>
+  return %pack : tensor<31250x?x16x1xf32>
+}
+// CHECK-LABEL: func @fold_padding_value_pack_dynamic_with_unit_tile_size
+// CHECK-NOT:     padding_value
+
+// -----
+
 func.func @infer_src_shape_pack(%src: tensor<?x?x?x?xf32>, %dest: tensor<10x20x30x40x16xf32>) -> tensor<10x20x30x40x16xf32> {
   %cst = arith.constant 0.000000e+00 : f32
    %pack = linalg.pack %src

``````````

</details>


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