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

[Jerry Shih]

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

>From 8745ddb01c13c459a5ff11b9465daa923f904056 Mon Sep 17 00:00:00 2001
From: Jerry Shih <jerry.shih at sifive.com>
Date: Thu, 30 Apr 2026 02:15:01 -0700
Subject: [PATCH] [mlir][linalg] add more pattern to fold pack op
 padding_value.

No padding is needed for unit tile size.
---
 mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp   | 11 +++---
 mlir/test/Dialect/Linalg/canonicalize.mlir | 41 +++++++++++++++++++++-
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index de7f4d1610bd0..a7eef0c50c2f3 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5879,11 +5879,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 12bdeb84e47e0..bb11ce0d4dfb8 100644
--- a/mlir/test/Dialect/Linalg/canonicalize.mlir
+++ b/mlir/test/Dialect/Linalg/canonicalize.mlir
@@ -1610,7 +1610,46 @@ func.func @fold_padding_value_pack(%arg0: tensor<1200x500000xf32>) -> tensor<312
   return %pack : tensor<31250x1200x16x1xf32>
 }
 // CHECK-LABEL: func @fold_padding_value_pack
-// CHECK-NOT:     padding_value
+//       CHECK:   linalg.pack
+//   CHECK-NOT:   padding_value
+
+// -----
+
+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:   linalg.pack
+//   CHECK-NOT:   padding_value
+
+// -----
+
+func.func @nofold_padding_value_pack_dynamic_with_non_unit_tile_size(%arg0: tensor<?x500000xf32>) -> tensor<31250x?x16x2xf32> {
+  %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?x16x2xf32>
+  %pack = linalg.pack %arg0
+    padding_value(%cst : f32)
+    outer_dims_perm = [1, 0]
+    inner_dims_pos = [1, 0]
+    inner_tiles = [16, 2]
+    into %0 : tensor<?x500000xf32> -> tensor<31250x?x16x2xf32>
+  return %pack : tensor<31250x?x16x2xf32>
+}
+// CHECK-LABEL: func @nofold_padding_value_pack_dynamic_with_non_unit_tile_size
+//       CHECK:   linalg.pack
+//       CHECK:   padding_value
 
 // -----