[Mlir-commits] [mlir] 171d8c4 - [Flang][OpenMP][MLIR] Fix memory leak caused by D149368 causing sanitizer error and fix iterator invalidation error

[Andrew Gozillon]

Author: Andrew Gozillon
Date: 2023-09-20T22:28:11-05:00
New Revision: 171d8c40288a995f7c34b0479f31241132af67a5

URL: https://github.com/llvm/llvm-project/commit/171d8c40288a995f7c34b0479f31241132af67a5
DIFF: https://github.com/llvm/llvm-project/commit/171d8c40288a995f7c34b0479f31241132af67a5.diff

LOG: [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.

Added: 
    

Modified: 
    flang/lib/Optimizer/Transforms/OMPEarlyOutlining.cpp
    mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

Removed: 
    


################################################################################
diff  --git a/flang/lib/Optimizer/Transforms/OMPEarlyOutlining.cpp b/flang/lib/Optimizer/Transforms/OMPEarlyOutlining.cpp
index 4eba921636e8abb..355a687d5f88588 100644
--- a/flang/lib/Optimizer/Transforms/OMPEarlyOutlining.cpp
+++ b/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(

diff  --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index e0c139be0f77f61..8f7f1963b3e5a4f 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/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();
-        }
       }
     }
   }