From b4fdc7d3af0f108403f0b4102ecdba77ae593376 Mon Sep 17 00:00:00 2001
From: Konstantin S Bobrovsky <konstantin.s.bobrovsky@intel.com>
Date: Tue, 27 Sep 2022 12:20:27 -0700
Subject: [PATCH] [ESIMD] Fix intermittent memory corruption in
 ESIMDOptimizeVecArgCallConv.cpp

Functions were erased from Module too early, while iterating over them.
Also fix signed/unsigned comparison warnings.

Signed-off-by: Konstantin S Bobrovsky <konstantin.s.bobrovsky@intel.com>
---
 .../ESIMD/ESIMDOptimizeVecArgCallConv.cpp     | 29 +++++++++----------
 1 file changed, 14 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/SYCLLowerIR/ESIMD/ESIMDOptimizeVecArgCallConv.cpp b/llvm/lib/SYCLLowerIR/ESIMD/ESIMDOptimizeVecArgCallConv.cpp
index d9add86a00080..4c88e794b7195 100644
--- a/llvm/lib/SYCLLowerIR/ESIMD/ESIMDOptimizeVecArgCallConv.cpp
+++ b/llvm/lib/SYCLLowerIR/ESIMD/ESIMDOptimizeVecArgCallConv.cpp
@@ -374,7 +374,7 @@ optimizeFunction(Function *OldF,
   }
   Instruction *At = &*(NewF->getEntryBlock().begin());
 
-  for (int I = 0; I < NewInsts.size(); ++I) {
+  for (unsigned I = 0; I < NewInsts.size(); ++I) {
     NewInsts[I]->insertBefore(At);
   }
   return NewF;
@@ -394,7 +394,7 @@ void optimizeCall(CallInst *CI, Function *OptF,
   int SretInd = -1;
   IRBuilder<> Bld(CI); // insert before CI
 
-  for (int I = 0; I < OptimizeableParams.size(); ++I) {
+  for (unsigned I = 0; I < OptimizeableParams.size(); ++I) {
     const auto &PI = OptimizeableParams[I];
     auto ArgNo = PI.getFormalParam().getArgNo();
 
@@ -475,20 +475,14 @@ static bool processFunction(Function *F) {
   }
   // Optimize the function.
   Function *NewF = optimizeFunction(F, OptimizeableParams, NewParamTs);
-  // Copy users to a separate container, to enable safe eraseFromParent.
-  SmallVector<User *> FUsers;
-  std::copy(F->users().begin(), F->users().end(), std::back_inserter(FUsers));
 
   // Optimize calls to the function.
-  // Iterate over FUsers, to enable safe eraseFromParent in optimizeCall.
-  for (auto *U : FUsers) {
+  for (auto *U : F->users()) {
     auto *Call = cast<CallInst>(U);
     assert(Call->getCalledFunction() == F);
     optimizeCall(Call, NewF, OptimizeableParams);
   }
-  std::string Name = F->getName().str();
-  F->eraseFromParent();
-  NewF->setName(Name);
+  NewF->takeName(F);
   return true;
 }
 
@@ -506,13 +500,18 @@ ESIMDOptimizeVecArgCallConvPass::run(Module &M, ModuleAnalysisManager &MAM) {
   }
 #endif // DEBUG_OPT_VEC_ARG_CALL_CONV
 
-  SmallVector<Function *, 16> Funcs;
-  std::for_each(M.begin(), M.end(), [&](Function &F) { Funcs.push_back(&F); });
+  SmallVector<Function *, 16> ToErase;
 
-  // Iterate over Funcs, to enable safe eraseFromParent in processFunction.
-  for (Function *F : Funcs) {
-    Modified &= processFunction(F);
+  for (Function &F : M) {
+    const bool FReplaced = processFunction(&F);
+    Modified &= FReplaced;
+
+    if (FReplaced) {
+      ToErase.push_back(&F);
+    }
   }
+  std::for_each(ToErase.begin(), ToErase.end(),
+                [](Function *F) { F->eraseFromParent(); });
 #ifdef DEBUG_OPT_VEC_ARG_CALL_CONV
   {
     std::error_code EC;