[Flang][OpenMP][MLIR] Fix memory leak caused by D149368 causing sanitizer error and fix iterator invalidation error

This patch fixes two issues introduced by the D149368 patch, one is
a memory leak from using the removeFromParent rather
than eraseFromParent (the erase also had to be moved to not create
use after deletes).

And the other is a possible iterator invalidation bug, better to be safe
than sorry.

Author: agozillon

flang/lib/Optimizer/Transforms/OMPEarlyOutlining.cpp

@@ -127,14 +127,18 @@ class OMPEarlyOutliningPass
     llvm::SetVector<mlir::Value> inputs;
     mlir::Region &targetRegion = targetOp.getRegion();
     mlir::getUsedValuesDefinedAbove(targetRegion, inputs);
-
+    
     // filter out declareTarget and map entries which are specially handled
     // at the moment, so we do not wish these to end up as function arguments
     // which would just be more noise in the IR.
-    for (mlir::Value value : inputs)
-      if (mlir::isa_and_nonnull<mlir::omp::MapInfoOp>(value.getDefiningOp()) ||
-          isAddressOfGlobalDeclareTarget(value))
-        inputs.remove(value);
+    for (llvm::SetVector<mlir::Value>::iterator iter = inputs.begin(); iter != inputs.end();) {
+      if (mlir::isa_and_nonnull<mlir::omp::MapInfoOp>(iter->getDefiningOp()) ||
+          isAddressOfGlobalDeclareTarget(*iter)) {
+        iter = inputs.erase(iter);
+      } else {
+        ++iter;
+      }
+    }
 
     // Create new function and initialize
     mlir::FunctionType funcType = builder.getFunctionType(

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

@@ -1989,6 +1989,15 @@ handleDeclareTargetMapVar(llvm::ArrayRef<Value> mapOperands,
           user->replaceUsesOfWith(mapOpValue, load);
         }
       }
+
+      // A global has already been generated by this stage for device
+      // that's now dead after the remapping, we can remove it now,
+      // as we have replaced all usages with the new ref pointer.
+      if (llvm::GlobalValue *unusedGlobalVal =
+              dyn_cast<llvm::GlobalValue>(mapOpValue)) {
+        unusedGlobalVal->dropAllReferences();
+        unusedGlobalVal->eraseFromParent();
+      }
     }
   }
 }
@@ -2182,15 +2191,6 @@ convertDeclareTargetAttr(Operation *op, mlir::omp::DeclareTargetAttr attribute,
             generatedRefs, /*OpenMPSimd*/ false, targetTriple, gVal->getType(),
             /*GlobalInitializer*/ nullptr,
             /*VariableLinkage*/ nullptr);
-        // A global has already been generated by this stage, unlike Clang, so
-        // this needs to be specially removed here for device when we're
-        // anything but a To clause specified variable with no unified shared
-        // memory.
-        if (llvm::GlobalValue *llvmVal =
-                llvmModule->getNamedValue(mangledName)) {
-          llvmVal->removeFromParent();
-          llvmVal->dropAllReferences();
-        }
       }
     }
   }