From 9ec52275acd6120db9a33d4f97d28848166cf839 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me@m-sp.org>
Date: Tue, 30 May 2023 15:12:54 +0200
Subject: [PATCH] [mlir][linalg] FuseIntoContainingOp: Always set
 newContainingOp

All result handles must be set in case of success.

Differential Revision: https://reviews.llvm.org/D151705
---
 mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index f18f24d..9233ce9 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -699,11 +699,6 @@ transform::FuseIntoContainingOp::apply(transform::TransformResults &results,
                                        transform::TransformState &state) {
   SmallVector<Operation *> fusedOps;
   auto producerOps = state.getPayloadOps(getProducerOp());
-  // If nothing to fuse, propagate success.
-  if (std::empty(producerOps)) {
-    results.set(cast<OpResult>(getFusedOp()), SmallVector<mlir::Operation *>{});
-    return DiagnosedSilenceableFailure::success();
-  }
   auto containingOps = state.getPayloadOps(getContainingOp());
   if (!llvm::hasSingleElement(containingOps)) {
     return emitDefiniteFailure()
@@ -712,6 +707,13 @@ transform::FuseIntoContainingOp::apply(transform::TransformResults &results,
   }
   Operation *containingOp = *containingOps.begin();
 
+  // If nothing to fuse, propagate success.
+  if (std::empty(producerOps)) {
+    results.set(cast<OpResult>(getFusedOp()), SmallVector<mlir::Operation *>{});
+    results.set(cast<OpResult>(getNewContainingOp()), {containingOp});
+    return DiagnosedSilenceableFailure::success();
+  }
+
   // Helper function to find the next producer that should be fused. Take any
   // producer that has a use inside the containing op.
   SetVector<Operation *> remainingProducers(producerOps.begin(),
-- 
2.7.4