[msan] Allow using a MemorySanitizer instance for functions across different modules

[loladiro]

Summary

We use MemorySanitizer in a JIT setting where we run all of our functions through a function pass manager to apply some of our own lowering passes before passing everything off to MCJIT. I tried adding MemorySanitizer to this setup, but the current implementation assumes that all functions MemorySanitizer will run on belong to the same module. Address this by reinitializing MemorySanitizer if the passed function is in a different module and looking or creating the correct globals in that module.

[eugenis]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

370

Do we need to reinit C as well? What about things like DataLayout, TargetTriple (I'm looking at MemorySanitizer::doInitialization) etc? Where do we stop?

Perhaps it would be cleaner to detect the change of ModuleID and reinit everything?

[eugenis]

Is it possible to test this?

[loladiro]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

370

Hmm, I'm not sure. This was the only changes I needed to make it work for my use case, which is a single FunctionPassManager running over all my Functions independent of which module they're in. That works fine for every pass but msan, but perhaps somebody more familiar with the pass manager can chime in here what the actual expectation for a function pass is.

lib/Transforms/Instrumentation/MemorySanitizer.cpp

 /// \brief Insert extern declaration of runtime-provided functions and globals.
 void MemorySanitizer::initializeCallbacks(Module &M) {
   // Only do this once.
-  if (WarningFn)
+  if (WarningFn && cast<Function>(WarningFn)->getParent() == &M)
     return;
 
   IRBuilder<> IRB(*C);

     IntptrTy, nullptr);
 
   // Create globals.
-  RetvalTLS = new GlobalVariable(
-    M, ArrayType::get(IRB.getInt64Ty(), kRetvalTLSSize / 8), false,
-    GlobalVariable::ExternalLinkage, nullptr, "__msan_retval_tls", nullptr,
-    GlobalVariable::InitialExecTLSModel);
-  RetvalOriginTLS = new GlobalVariable(
-    M, OriginTy, false, GlobalVariable::ExternalLinkage, nullptr,
-    "__msan_retval_origin_tls", nullptr, GlobalVariable::InitialExecTLSModel);
-
-  ParamTLS = new GlobalVariable(
-    M, ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8), false,
-    GlobalVariable::ExternalLinkage, nullptr, "__msan_param_tls", nullptr,
-    GlobalVariable::InitialExecTLSModel);
-  ParamOriginTLS = new GlobalVariable(
-    M, ArrayType::get(OriginTy, kParamTLSSize / 4), false,
-    GlobalVariable::ExternalLinkage, nullptr, "__msan_param_origin_tls",
-    nullptr, GlobalVariable::InitialExecTLSModel);
-
-  VAArgTLS = new GlobalVariable(
+  RetvalTLS = M.getGlobalVariable("__msan_retval_tls");
+  if (!RetvalTLS) {
+    RetvalTLS = new GlobalVariable(
+      M, ArrayType::get(IRB.getInt64Ty(), kRetvalTLSSize / 8), false,
+      GlobalVariable::ExternalLinkage, nullptr, "__msan_retval_tls", nullptr,
+      GlobalVariable::InitialExecTLSModel);
+  }
+
+  RetvalOriginTLS = M.getGlobalVariable("__msan_retval_origin_tls");
+  if (!RetvalOriginTLS) {
+    RetvalOriginTLS = new GlobalVariable(
+      M, OriginTy, false, GlobalVariable::ExternalLinkage, nullptr,
+      "__msan_retval_origin_tls", nullptr, GlobalVariable::InitialExecTLSModel);
+  }
+
+  ParamTLS = M.getGlobalVariable("__msan_param_tls");
+  if (!ParamTLS) {
+    ParamTLS = new GlobalVariable(
       M, ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8), false,
-    GlobalVariable::ExternalLinkage, nullptr, "__msan_va_arg_tls", nullptr,
+      GlobalVariable::ExternalLinkage, nullptr, "__msan_param_tls", nullptr,
       GlobalVariable::InitialExecTLSModel);
-  VAArgOverflowSizeTLS = new GlobalVariable(
-    M, IRB.getInt64Ty(), false, GlobalVariable::ExternalLinkage, nullptr,
-    "__msan_va_arg_overflow_size_tls", nullptr,
-    GlobalVariable::InitialExecTLSModel);
-  OriginTLS = new GlobalVariable(
-    M, IRB.getInt32Ty(), false, GlobalVariable::ExternalLinkage, nullptr,
-    "__msan_origin_tls", nullptr, GlobalVariable::InitialExecTLSModel);
+  }
+
+  ParamOriginTLS = M.getGlobalVariable("__msan_param_origin_tls");
+  if (!ParamOriginTLS) {
+    ParamOriginTLS = new GlobalVariable(
+      M, ArrayType::get(OriginTy, kParamTLSSize / 4), false,
+      GlobalVariable::ExternalLinkage, nullptr, "__msan_param_origin_tls",
+      nullptr, GlobalVariable::InitialExecTLSModel);
+  }
+
+  VAArgTLS = M.getGlobalVariable("__msan_va_arg_tls");
+  if (!VAArgTLS) {
+    VAArgTLS = new GlobalVariable(
+      M, ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8), false,
+      GlobalVariable::ExternalLinkage, nullptr, "__msan_va_arg_tls", nullptr,
+      GlobalVariable::InitialExecTLSModel);
+  }
+
+  VAArgOverflowSizeTLS = M.getGlobalVariable("__msan_va_arg_overflow_size_tls");
+  if (!VAArgOverflowSizeTLS) {
+    VAArgOverflowSizeTLS = new GlobalVariable(
+      M, IRB.getInt64Ty(), false, GlobalVariable::ExternalLinkage, nullptr,
+      "__msan_va_arg_overflow_size_tls", nullptr,
+      GlobalVariable::InitialExecTLSModel);
+  }
+
+  OriginTLS = M.getGlobalVariable("__msan_origin_tls");
+  if (!OriginTLS) {
+    OriginTLS = new GlobalVariable(
+      M, IRB.getInt32Ty(), false, GlobalVariable::ExternalLinkage, nullptr,
+      "__msan_origin_tls", nullptr, GlobalVariable::InitialExecTLSModel);
+  }
 
   // We insert an empty inline asm after __msan_report* to avoid callback merge.
   EmptyAsm = InlineAsm::get(FunctionType::get(IRB.getVoidTy(), false),