[Mlir-commits] [mlir] [mlir] Avoid unset results in test.tile_using_forall (PR #220498)

[llvmlistbot at llvm.org]

https://github.com/AkunCube created https://github.com/llvm/llvm-project/pull/220498

`transform.test.tile_using_forall` can leave variadic loop results unset when the target handle is empty or when the number of generated loops does not match the number of declared loop results.

The transform interpreter expects every result of a successful transform to be initialized. Leaving a result unset causes `TransformResults::get()` to hit an assertion, as reported in #220460.

  This change:
  - reports an error when tiling produces more loops than declared loop results;
  - initializes missing loop results to empty handles;
  - adds regression coverage for empty targets;
  - adds diagnostic coverage for two payload matmuls with only one declared loop
    result.

  No changes are made to the public `transform.structured.tile_using_forall` operation or to the variadic result design.

Fixes: https://github.com/llvm/llvm-project/issues/220460

>From a313fe286344f114793fe08c2c16af3839a222ab Mon Sep 17 00:00:00 2001
From: cubehan3 <cubehan3 at gmail.com>
Date: Wed, 2 Sep 2026 07:10:19 +0000
Subject: [PATCH] [mlir] Avoid unset results in test.tile_using_forall

---
 .../tile-using-scfforall-failure.mlir         | 25 +++++++++++++++++++
 .../TilingInterface/tile-using-scfforall.mlir | 16 ++++++++++++
 .../TestTilingInterfaceTransformOps.cpp       |  8 ++++++
 3 files changed, 49 insertions(+)
 create mode 100644 mlir/test/Interfaces/TilingInterface/tile-using-scfforall-failure.mlir

diff --git a/mlir/test/Interfaces/TilingInterface/tile-using-scfforall-failure.mlir b/mlir/test/Interfaces/TilingInterface/tile-using-scfforall-failure.mlir
new file mode 100644
index 0000000000000..95068832a7b43
--- /dev/null
+++ b/mlir/test/Interfaces/TilingInterface/tile-using-scfforall-failure.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-opt -transform-interpreter -verify-diagnostics %s
+
+func.func @two_matmuls(
+    %arg0 : tensor<?x?xf32>, %arg1 : tensor<?x?xf32>,
+    %arg2 : tensor<?x?xf32>, %arg3 : tensor<?x?xf32>,
+    %arg4 : tensor<?x?xf32>, %arg5 : tensor<?x?xf32>) -> tensor<?x?xf32> {
+  %0 = linalg.matmul
+    ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  %1 = linalg.matmul
+    ins(%arg3, %arg4 : tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%arg5 : tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%arg0 : !transform.any_op {transform.readonly}) {
+    %matmul = transform.structured.match ops{["linalg.matmul"]} in %arg0
+      : (!transform.any_op) -> !transform.any_op
+    // expected-error @+1 {{generated 2 loops but only 1 loop results were declared}}
+    %tiled, %loop = transform.test.tile_using_forall %matmul [10, 20]
+      : (!transform.any_op) -> (!transform.any_op, !transform.any_op)
+    transform.yield
+  }
+}
diff --git a/mlir/test/Interfaces/TilingInterface/tile-using-scfforall.mlir b/mlir/test/Interfaces/TilingInterface/tile-using-scfforall.mlir
index 745a82fc0da75..00458f72f0552 100644
--- a/mlir/test/Interfaces/TilingInterface/tile-using-scfforall.mlir
+++ b/mlir/test/Interfaces/TilingInterface/tile-using-scfforall.mlir
@@ -349,3 +349,19 @@ module attributes {transform.with_named_sequence} {
 // CHECK-LABEL: func @check_scalar_memref_operation
 //   CHECK-NOT:   scf.for
 //       CHECK:   linalg.generic
+
+// -----
+
+// Empty handles must still initialize all declared results.
+module {
+  module attributes {transform.with_named_sequence} {
+    transform.named_sequence @__transform_main(%arg0 : !transform.any_op {transform.readonly}) {
+      %0 = transform.structured.match ops{["linalg.matmul"]} in %arg0
+        : (!transform.any_op) -> !transform.any_op
+      %tiled_op, %loops:2 = transform.test.tile_using_forall %0 [10, 20]
+        : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op)
+      transform.yield
+    }
+  }
+}
+// CHECK: transform.test.tile_using_forall
diff --git a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
index 9467c925e543c..b3d5bb43eb70f 100644
--- a/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
+++ b/mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp
@@ -356,9 +356,17 @@ applyTileToAll(RewriterBase &rewriter, Operation *transformOp,
       loopOps.push_back(loop);
   }
 
+  unsigned numLoopResults = transformOp->getNumResults() - 1;
+  if (loopOps.size() > numLoopResults)
+    return transformOp->emitError()
+           << "generated " << loopOps.size() << " loops but only "
+           << numLoopResults << " loop results were declared";
+
   transformResults.set(transformOp->getOpResult(0), tiledOps);
   for (auto [index, loop] : llvm::enumerate(loopOps))
     transformResults.set(transformOp->getOpResult(index + 1), {loop});
+  for (unsigned index : llvm::seq<unsigned>(loopOps.size(), numLoopResults))
+    transformResults.set(transformOp->getOpResult(index + 1), {});
 
   return success();
 }