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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: crZhao (cr-zhao)

<details>
<summary>Changes</summary>

## Summary

- stop deriving the generated `scf.forall` rank from the unadjusted tile-size list in the test transform
- keep the structural assertion that forall tiling returns a single `scf.forall`
- add regression coverage for excess tile sizes on a one-dimensional `linalg.generic`

## Root cause

`tileUsingSCF` resizes the requested tile sizes to the target operation's iteration domain. `transform.test.fuse_using_forall` instead computed an expected rank from the original tile-size list. For a one-dimensional target with `[10, 20]`, tiling correctly generated a rank-1 `scf.forall`, but the test transform asserted that its rank must be 2 and crashed.

The transform returns one handle for the generated `scf.forall`, independent of that operation's rank, so the rank assertion was not a valid invariant.

Fixes #<!-- -->206619.

## Testing

- `build/bin/mlir-opt --transform-interpreter --cse --split-input-file mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-scfforall.mlir | build/bin/FileCheck mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-scfforall.mlir`
- original issue reproducer with `build/bin/mlir-opt -transform-interpreter` (exit code 0)

Assisted-by: Codex


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


2 Files Affected:

- (modified) mlir/test/Interfaces/TilingInterface/tile-and-fuse-using-scfforall.mlir (+32) 
- (modified) mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp (+4-5) 


``````````diff
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,

``````````

</details>


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