[Mlir-commits] [mlir] [mlir][test] Allow excess tile sizes in fuse_using_forall (PR #216370)

[llvmlistbot at llvm.org]

https://github.com/cr-zhao updated https://github.com/llvm/llvm-project/pull/216370

>From bb10500f4b3a1ce13d0ac0e44b261710b8ea380a Mon Sep 17 00:00:00 2001
From: real-cpu <zhaochenrui757 at gmail.com>
Date: Fri, 14 Aug 2026 10:51:41 -0700
Subject: [PATCH] [mlir][test] Allow excess tile sizes in fuse_using_forall

---
 .../tile-and-fuse-using-scfforall.mlir        | 32 +++++++++++++++++++
 .../TestTilingInterfaceTransformOps.cpp       |  9 +++---
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-scfforall.mlir b/mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-scfforall.mlir
index 0bd2546e082b5..6682c0de3e133 100644
--- a/mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-scfforall.mlir
+++ b/mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-scfforall.mlir
@@ -174,3 +174,35 @@ module attributes {transform.with_named_sequence} {
 //       CHECK:       tensor.parallel_insert_slice %[[GENERIC2]] into %[[ITERARG0]][%[[IV]], 0]
 //       CHECK:     }
 //       CHECK:   return %[[RESULT]]
+
+// -----
+
+// Tile sizes beyond the iteration domain are ignored.
+func.func @excess_tile_sizes(%arg0: tensor<?xf32>,
+                             %arg1: tensor<?xf32>) -> tensor<?xf32> {
+  %0 = linalg.generic {
+      indexing_maps = [affine_map<(d0) -> (d0)>,
+                       affine_map<(d0) -> (d0)>],
+      iterator_types = ["parallel"]}
+      ins(%arg0 : tensor<?xf32>) outs(%arg1 : tensor<?xf32>) {
+    ^bb0(%in: f32, %out: f32):
+      %sum = arith.addf %in, %out : f32
+      linalg.yield %sum : f32
+  } -> tensor<?xf32>
+  return %0 : tensor<?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(
+      %arg0: !transform.any_op {transform.readonly}) {
+    %generic = transform.structured.match ops{["linalg.generic"]} in %arg0
+      : (!transform.any_op) -> !transform.any_op
+    %tiled, %loop = transform.test.fuse_using_forall %generic [10, 20]
+      : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
+
+// CHECK-LABEL: func.func @excess_tile_sizes
+// CHECK:         scf.forall (%[[IV:.+]]) = (0) to
+// CHECK:           linalg.generic
diff --git a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
index 9467c925e543c..5902001094f5c 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
@@ -397,7 +397,7 @@ void transform::TestTileUsingForallOp::getEffects(
 template <typename Range>
 static LogicalResult applyTilingToAll(
     RewriterBase &rewriter, Operation *transformOp, Range &&payloadOps,
-    unsigned numLoops, TransformResults &transformResults,
+    TransformResults &transformResults,
     function_ref<FailureOr<scf::SCFTileAndFuseResult>(TilingInterface)>
         applyFn) {
   SmallVector<Operation *> tiledLinalgOps;
@@ -428,9 +428,8 @@ static LogicalResult applyTilingToAll(
     // Report back the relevant handles to the transform op.
     tiledLinalgOps.push_back(tiledResults->tiledAndFusedOps.front());
     assert(tiledResults->loops.size() == 1 &&
-           cast<scf::ForallOp>(tiledResults->loops[0]).getRank() == numLoops &&
-           "Mismatched number of loops, tile and fuse transform should have "
-           "failed");
+           isa<scf::ForallOp>(tiledResults->loops[0]) &&
+           "expected a single scf.forall op");
     loopOps[0] = {tiledResults->loops[0]};
   }
 
@@ -460,7 +459,7 @@ transform::TestFuseUsingForallOp::apply(TransformRewriter &rewriter,
   tileAndFuseOptions.tilingOptions = tilingOptions;
   LogicalResult result = applyTilingToAll(
       rewriter, getOperation(), state.getPayloadOps(getRootOp()),
-      tileSizes.size() - llvm::count(tileSizes, 0), transformResults,
+      transformResults,
       [&](TilingInterface tilingInterfaceOp)
           -> FailureOr<scf::SCFTileAndFuseResult> {
         return tileConsumerAndFuseProducersUsingSCF(rewriter, tilingInterfaceOp,