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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>

`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

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


3 Files Affected:

- (added) mlir/test/Interfaces/TilingInterface/tile-using-scfforall-failure.mlir (+25) 
- (modified) mlir/test/Interfaces/TilingInterface/tile-using-scfforall.mlir (+16) 
- (modified) mlir/test/lib/Interfaces/TilingInterface/TestTilingInterfaceTransformOps.cpp (+8) 


``````````diff
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();
 }

``````````

</details>


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