[Mlir-commits] [llvm] [mlir] [MLIR][OpenMP] Add Taskloop Collapse Support (PR #175924)

[Jack Styles]

================
@@ -2367,10 +2367,20 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createTaskloop(
       Value *IVPlusTaskLB = Builder.CreateAdd(
           CLI->getIndVar(),
           Builder.CreateSub(CastedTaskLB, ConstantInt::get(IVTy, 1)));
-      for (User *IVUser : CLI->getIndVar()->users()) {
-        if (IVUser == IVPlusTaskLB)
-          continue;
-        IVUser->replaceUsesOfWith(CLI->getIndVar(), IVPlusTaskLB);
+      for (auto IVUse = CLI->getIndVar()->uses().begin();
+           IVUse != CLI->getIndVar()->uses().end();) {
+        User *IVUser = IVUse->getUser();
+        // To ensure every Use is correctly captured, we want to iterate before
+        // replacing the uses of the loop index. If this is done after replacing
+        // the uses, then it is possible for uses to be missed, and values are
+        // not calculated correctly
+        IVUse++;
----------------
Stylie777 wrote:

I've updated this.

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