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

[llvmlistbot at llvm.org]

Author: Jerry Shih
Date: 2026-06-10T18:01:08+02:00
New Revision: 3443243ded167229ddc37e64b7e754854ae1ba2c

URL: https://github.com/llvm/llvm-project/commit/3443243ded167229ddc37e64b7e754854ae1ba2c
DIFF: https://github.com/llvm/llvm-project/commit/3443243ded167229ddc37e64b7e754854ae1ba2c.diff

LOG: [mlir][linalg] add more pattern to fold pack op padding_value. (#198468)

No padding is needed for unit tile size.

Added: 
    

Modified: 
    mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
    mlir/test/Dialect/Linalg/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
index 6b61564eb5fcc..a921b8d69ee45 100644
--- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
+++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp
@@ -5885,11 +5885,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
 
 // -----