[Mlir-commits] [mlir] [mlir][scf] Fix `FoldTensorCastOfOutputIntoForallOp` for multi-result scf.forall (PR #173271)

[llvmlistbot at llvm.org]

================
@@ -1986,6 +1986,15 @@ struct FoldTensorCastOfOutputIntoForallOp
     if (tensorCastProducers.empty())
       return failure();
 
+    llvm::SmallMapVector<Operation *, int64_t, 2> yieldOpToIterArgsIndex;
+    for (auto [index, iterArg] :
+         llvm::enumerate(forallOp.getRegionIterArgs())) {
+      for (Operation *user : iterArg.getUsers()) {
+        if (isa<ParallelCombiningOpInterface>(user))
----------------
MaheshRavishankar wrote:

This maybe needs to error out if a mapping already exists. I dont know off the top of my head if the verifier for `scf::ForAllOp` already checks for a single `tensor.parallel_insert_slice` per init arg.

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