[Mlir-commits] [mlir] [mlir] Fix consumer fusion for producer with multiple results (PR #125915)

[llvmlistbot at llvm.org]

================
@@ -1979,6 +2033,26 @@ mlir::scf::tileAndFuseConsumerOfSlice(RewriterBase &rewriter,
         consumerOp, "consumer op's operand doesn't seem to be an OpResult");
   }
 
+  SmallVector<OpOperand *> potentialOperands = {*maybeConsumerOpOperand};
+  SmallVector<unsigned> potentialOperandResultNos = {
+      consumerOpOperand->getOperandNumber()};
+  SmallVector<Operation *> potentialSliceOps = {candidateSliceOp};
+
+  // 1b. Get all the other operands of the consumer op and their corresponding
+  // slice ops. In the case of the consumer using multiple results
+  // from the producer, we need to update every operand.
+  for (OpOperand &otherOperand : consumerOp->getOpOperands()) {
+    if (&otherOperand == *maybeConsumerOpOperand)
+      continue;
+    auto maybePotentialSlice = getSliceOpFromConsumerOperand(otherOperand);
+    if (failed(maybePotentialSlice)) {
+      continue;
+    }
+    potentialSliceOps.push_back(*maybePotentialSlice);
----------------
MaheshRavishankar wrote:

Do you want to check right here that the producer of `otherOperand` is the same as `*maybeConsumerOpOperand`?

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