[MSan] Apply the ctor creation scheme of TSan

[philip.pfaffe]

Summary

To avoid adding an extern function to the global ctors list, apply the changes of D56538 also to MSan.

[chandlerc]

LGTM with some minor cleanups!

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

925 ↗

(On Diff #181823)

Add a comment at the start of the lambda to remind the reader that this is used when we are the first to try and get the sanitizer ctor/init function to set it up?

931 ↗

(On Diff #181823)

Strange vertical whitespace here...

932–934 ↗

(On Diff #181823)

Instead of this pattern, I'd suggest early exit:

if (!ClWithComdat) {
  appendToGlobalCtors(M, cTor, 0);
  return;
}
...

[philip.pfaffe]

Address comments and update two missing testcases.

llvm/trunk/lib/Transforms/Instrumentation/MemorySanitizer.cpp

 static cl::opt<unsigned long long> ClShadowBase("msan-shadow-base",
        cl::desc("Define custom MSan ShadowBase"),
        cl::Hidden, cl::init(0));
 
 static cl::opt<unsigned long long> ClOriginBase("msan-origin-base",
        cl::desc("Define custom MSan OriginBase"),
        cl::Hidden, cl::init(0));
 
+static const char *const kMsanModuleCtorName = "msan.module_ctor";
 static const char *const kMsanInitName = "__msan_init";
 
 namespace {
 
 // Memory map parameters used in application-to-shadow address calculation.
 // Offset = (Addr & ~AndMask) ^ XorMask
 // Shadow = ShadowBase + Offset
 // Origin = OriginBase + Offset
private:
 
   MDNode *ColdCallWeights;
 
   /// Branch weights for origin store.
   MDNode *OriginStoreWeights;
 
   /// An empty volatile inline asm that prevents callback merge.
   InlineAsm *EmptyAsm;
+
+  Function *MsanCtorFunction;
 };
 
 /// A legacy function pass for msan instrumentation.
 ///
 /// Instruments functions to detect unitialized reads.
 struct MemorySanitizerLegacyPass : public FunctionPass {
   // Pass identification, replacement for typeid.
   static char ID;
Value *MemorySanitizer::getKmsanShadowOriginAccessFn(bool isStore, int size) {
   case 8:
     return Fns[3];
   default:
     return nullptr;
   }
 }
 
 /// Module-level initialization.
+///
+/// inserts a call to __msan_init to the module's constructor list.
 void MemorySanitizer::initializeModule(Module &M) {
   auto &DL = M.getDataLayout();
 
   bool ShadowPassed = ClShadowBase.getNumOccurrences() > 0;
   bool OriginPassed = ClOriginBase.getNumOccurrences() > 0;
   // Check the overrides first
   if (ShadowPassed || OriginPassed) {
     CustomMapParams.AndMask = ClAndMask;
void MemorySanitizer::initializeModule(Module &M) {
   IRBuilder<> IRB(*C);
   IntptrTy = IRB.getIntPtrTy(DL);
   OriginTy = IRB.getInt32Ty();
 
   ColdCallWeights = MDBuilder(*C).createBranchWeights(1, 1000);
   OriginStoreWeights = MDBuilder(*C).createBranchWeights(1, 1000);
 
   if (!CompileKernel) {
-    getOrCreateInitFunction(M, kMsanInitName);
+    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, Function *) {
+              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(
             M, IRB.getInt32Ty(), true, GlobalValue::WeakODRLinkage,
             IRB.getInt32(TrackOrigins), "__msan_track_origins");
       });
 
static VarArgHelper *CreateVarArgHelper(Function &Func, MemorySanitizer &Msan,
   else if (TargetTriple.getArch() == Triple::ppc64 ||
            TargetTriple.getArch() == Triple::ppc64le)
     return new VarArgPowerPC64Helper(Func, Msan, Visitor);
   else
     return new VarArgNoOpHelper(Func, Msan, Visitor);
 }
 
 bool MemorySanitizer::sanitizeFunction(Function &F, TargetLibraryInfo &TLI) {
+  if (!CompileKernel && (&F == MsanCtorFunction))
+    return false;
   MemorySanitizerVisitor Visitor(F, *this, TLI);
 
   // Clear out readonly/readnone attributes.
   AttrBuilder B;
   B.addAttribute(Attribute::ReadOnly)
     .addAttribute(Attribute::ReadNone);
   F.removeAttributes(AttributeList::FunctionIndex, B);
 
   return Visitor.runOnFunction();
 }

llvm/trunk/test/Instrumentation/MemorySanitizer/global_ctors_2to3.ll

+; MSan converts 2-element global_ctors to 3-element when adding the new entry.
+; RUN: opt < %s -msan-with-comdat -S -passes=msan 2>&1 | FileCheck %s
+; RUN: opt < %s -msan -msan-with-comdat -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+; CHECK: $msan.module_ctor = comdat any
+; CHECK: @llvm.global_ctors = appending global [2 x { i32, void ()*, i8* }] [{ i32, void ()*, i8* } { i32 65535, void ()* @f, i8* null }, { i32, void ()*, i8* } { i32 0, void ()* @msan.module_ctor, i8* bitcast (void ()* @msan.module_ctor to i8*) }]
+
[e-mail]_ctors = appending global [1 x { i32, void ()* }] [{ i32, void ()* } { i32 65535, void ()* @f }]
+
+define internal void @f() {
+entry:
+  ret void
+}
+
+; CHECK: define internal void @msan.module_ctor() comdat {

llvm/trunk/test/Instrumentation/MemorySanitizer/msan_basic.ll

 ; RUN: opt < %s -msan-check-access-address=0 -S -passes=msan 2>&1 | FileCheck  \
 ; RUN: -allow-deprecated-dag-overlap %s
 ; RUN: opt < %s -msan -msan-check-access-address=0 -S | FileCheck -allow-deprecated-dag-overlap %s
 ; RUN: opt < %s -msan-check-access-address=0 -msan-track-origins=1 -S          \
 ; RUN: -passes=msan 2>&1 | FileCheck -allow-deprecated-dag-overlap             \
 ; RUN: -check-prefix=CHECK -check-prefix=CHECK-ORIGINS %s
 ; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=1 -S | FileCheck -allow-deprecated-dag-overlap -check-prefix=CHECK -check-prefix=CHECK-ORIGINS %s
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
-; CHECK: @llvm.global_ctors {{.*}} { i32 0, void ()* @__msan_init, i8* null }
+; CHECK: @llvm.global_ctors {{.*}} { i32 0, void ()* @msan.module_ctor, i8* null }
 
 ; Check the presence and the linkage type of __msan_track_origins and
 ; other interface symbols.
 ; CHECK-NOT: @__msan_track_origins
 ; CHECK-ORIGINS: @__msan_track_origins = weak_odr constant i32 1
 ; CHECK-NOT: @__msan_keep_going = weak_odr constant i32 0
 ; CHECK: @__msan_retval_tls = external thread_local(initialexec) global [{{.*}}]
 ; CHECK: @__msan_retval_origin_tls = external thread_local(initialexec) global i32

 
 ; CHECK-LABEL: define i8* @MismatchingCallMustTailCall
 ; CHECK: musttail call i32* @MismatchingMustTailCall
 ; No instrumentation between call and ret.
 ; CHECK-NEXT: bitcast i32* {{.*}} to i8*
 ; CHECK-NEXT: ret i8*
 
 
-; CHECK: declare void @__msan_init()
+; CHECK-LABEL: define internal void @msan.module_ctor() {
+; CHECK: call void @__msan_init()

llvm/trunk/test/Instrumentation/MemorySanitizer/msan_llvm_is_constant.ll

 ; Make sure MSan doesn't insert shadow checks for @llvm.is.constant.* arguments.
 
+; RUN: opt < %s -msan-kernel=1 -S -passes=msan 2>&1 | FileCheck                \
+; RUN: -check-prefixes=CHECK %s
 ; RUN: opt < %s -msan -msan-kernel=1 -S | FileCheck -check-prefixes=CHECK %s
+; RUN: opt < %s -S -passes=msan 2>&1 | FileCheck -check-prefixes=CHECK %s
 ; RUN: opt < %s -msan -S | FileCheck -check-prefixes=CHECK %s
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
 ; Function Attrs: nounwind readnone uwtable
 define dso_local i32 @bar(i32 %v) local_unnamed_addr sanitize_memory {
 entry: