1 files changed, 20 insertions, 30 deletions

diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index 69c9020e060..cf93de6f759 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -587,25 +587,9 @@ private:
 
   /// An empty volatile inline asm that prevents callback merge.
   InlineAsm *EmptyAsm;
-};
 
-void insertModuleCtor(Module &M) {
-  getOrCreateSanitizerCtorAndInitFunctions(
-      M, kMsanModuleCtorName, kMsanInitName,
-      /*InitArgTypes=*/{},
-      /*InitArgs=*/{},
-      // This callback is invoked when the functions are created the first
-      // time. Hook them into the global ctors list in that case:
-      [&](Function *Ctor, FunctionCallee) {
-        if (!ClWithComdat) {
-          appendToGlobalCtors(M, Ctor, 0);
-          return;
-        }
-        Comdat *MsanCtorComdat = M.getOrInsertComdat(kMsanModuleCtorName);
-        Ctor->setComdat(MsanCtorComdat);
-        appendToGlobalCtors(M, Ctor, 0, Ctor);
-      });
-}
+  Function *MsanCtorFunction;
+};
 
 /// A legacy function pass for msan instrumentation.
 ///
@@ -651,14 +635,6 @@ PreservedAnalyses MemorySanitizerPass::run(Function &F,
   return PreservedAnalyses::all();
 }
 
-PreservedAnalyses MemorySanitizerPass::run(Module &M,
-                                           ModuleAnalysisManager &AM) {
-  if (Options.Kernel)
-    return PreservedAnalyses::all();
-  insertModuleCtor(M);
-  return PreservedAnalyses::none();
-}
-
 char MemorySanitizerLegacyPass::ID = 0;
 
 INITIALIZE_PASS_BEGIN(MemorySanitizerLegacyPass, "msan",
@@ -944,6 +920,23 @@ void MemorySanitizer::initializeModule(Module &M) {
   OriginStoreWeights = MDBuilder(*C).createBranchWeights(1, 1000);
 
   if (!CompileKernel) {
+    std::tie(MsanCtorFunction, std::ignore) =
+        getOrCreateSanitizerCtorAndInitFunctions(
+            M, kMsanModuleCtorName, kMsanInitName,
+            /*InitArgTypes=*/{},
+            /*InitArgs=*/{},
+            // This callback is invoked when the functions are created the first
+            // time. Hook them into the global ctors list in that case:
+            [&](Function *Ctor, FunctionCallee) {
+              if (!ClWithComdat) {
+                appendToGlobalCtors(M, Ctor, 0);
+                return;
+              }
+              Comdat *MsanCtorComdat = M.getOrInsertComdat(kMsanModuleCtorName);
+              Ctor->setComdat(MsanCtorComdat);
+              appendToGlobalCtors(M, Ctor, 0, Ctor);
+            });
+
     if (TrackOrigins)
       M.getOrInsertGlobal("__msan_track_origins", IRB.getInt32Ty(), [&] {
         return new GlobalVariable(
@@ -961,8 +954,6 @@ void MemorySanitizer::initializeModule(Module &M) {
 }
 
 bool MemorySanitizerLegacyPass::doInitialization(Module &M) {
-  if (!Options.Kernel)
-    insertModuleCtor(M);
   MSan.emplace(M, Options);
   return true;
 }
@@ -4587,9 +4578,8 @@ static VarArgHelper *CreateVarArgHelper(Function &Func, MemorySanitizer &Msan,
 }
 
 bool MemorySanitizer::sanitizeFunction(Function &F, TargetLibraryInfo &TLI) {
-  if (!CompileKernel && F.getName() == kMsanModuleCtorName)
+  if (!CompileKernel && (&F == MsanCtorFunction))
     return false;
-
   MemorySanitizerVisitor Visitor(F, *this, TLI);
 
   // Clear out readonly/readnone attributes.