MemorySanitizer: Add option to insert init checks at call site

[guiand]

Summary

This change depends on a previous patch adding the noundef attribute to LLVM calls.

With this change, specifying the option -msan-eager-checks will cause MemorySanitizer to insert initialization checks for function arguments before calling them, and for return values before returning.

This change depends on the noundef keyword to determining cases where it it sound to check these shadows, and falls back to passing shadows values by TLS.

Checking at call boundaries enforces undefined behavior rules with passing uninitialized arguments by value. It also provides substantial improvements to the optimized instrumented code, decreasing binary size by up to 17% and runtime up to 25%.

[vitalybuka]

May I ask you to move compiler-rt/ into a separate patch? With ClEagerChecks false by default we can do that and have smaller patches.

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

880

this line is to long. please clang-format the patch
I usually do that by:
git clang-format -f --style=file HEAD^

1082–1083

Please undo such irrelevant changes or move them into separate patches

1220

var names are inconsistent
same for i and it
llvm style is to start with uppercase

1222

so if we need integer index, then just iterate with getNumElements() getElementType(unsigned N) ?
range base loop and separate var is fine as well

BTW. the first "i" is 1, is it expected?

1223

declare and intialize variable at the same statement if possible.
I assume it's to destroy IRB before recursive call
maybe then extract lines in {} into a separate function?

1625

FArgCheck -> FArgEagerCheck?

1684

why do we need this?

3411

Check -> EagerCheck?

3467

ClEagerChecks && !PartialInit;

llvm/test/Instrumentation/MemorySanitizer/msan_eager.ll

2

would you like to try go generate test with llvm/utils/update_analyze_test_checks.py

10–21

llvm/utils/update_analyze_test_checks.py will likely drop this, but they are not important for this check

[glider]

Wrapping all variables in the existing tests in PartInit<> doesn't look good.
I believe instead of doing this you need to:

• change existing tests to also run with eager checks, adding expected MSan outputs as needed. Do not scatter PartInit<> across all tests.
• add more tests that use the partinit attribute to cover new functionality.

I also don't quite understand how this partinit thing works with programs written in C. We probably need a test for it as well.

compiler-rt/test/msan/in-struct-padding.cpp

4 ↗

(On Diff #270260)

It's unclear what this test is doing. Please consider adding comments.

compiler-rt/test/msan/param_tls_limit.cpp

39 ↗

(On Diff #270260)

Please don't introduce unnecessary diff.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3442

Please remove unrelated diff from the patch.

[aaron.ballman]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

867

Whitelist msan unaligned -> Allow msan unaligned

877

Whitelist msan unaligned -> Allow msan unaligned

[guiand]

It looks like folks are still deciding what to call the partialinit / frozen / nopoison attribute. Once that's settled I'll go through and update this patch and address the reviews.

One thing that's changed is that struct and union types in clang are never marked for checking anymore (at least, for now). So the tests I added to check that can likely go away.

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1222

Makes sense to use getElementType. But I think here the first i is still 0, as the i++ expression happens after each loop body.

1684

This change allows checked arguments to not take up TLS space since they don't use the TLS for anything anyway. That frees up more TLS for arguments which can't be checked.

[eugenis]

compiler-rt/test/msan/cxa_atexit.cpp

20 ↗

(On Diff #270260)

For tests that do not make sense with the eager checks, you could introduce a lit feature, and guard them with
// REQUIRES: !msan_eager_checks

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

288

... and return values ...

885

There is only one argument that can be uninit - the value to be stored. No need to set the attribute on all of them.

990

The module flag is gone, and in any case this should not be a fatal error - if the frontend can not promise correct attribute value, we always delay the check, that's all.

1223

It's also common to reuse the builder with IRB.SetInsertPoint

1233

This change can be done in a separate patch.

It's more efficient to OR everything together and have a single branch. I'm not sure which would generate better code:

• ShadowCombiner and materializeOneCheck on the result, this might do bad things if there are very large members (like vectors)
• icmp NE 0 each members' shadow and materializeOneCheck of that.

[guiand]

Addressed comments, updated for noundef

[guiand]

I'm splitting this patch into this new one, which depends on nothing but noundef correctly parsed by LLVM, and the rest which relies on clang emitting noundef everywhere.

[guiand]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3467

This one is for the return attribute, the other was for the arguments

[vitalybuka]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1684

clang-format this

llvm/test/Instrumentation/MemorySanitizer/msan_eager.ll

2

?

[guiand]

Fix test checking wrong __msan_warning variant, clang-format

[vitalybuka]

llvm/test/Instrumentation/MemorySanitizer/msan_eager.ll

2

; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
; RUN: opt < %s -msan-check-access-address=0 -msan-track-origins=1 -msan-eager-checks -S -passes='module(msan-module),function(msan)' 2>&1 | \
; RUN:   FileCheck -allow-deprecated-dag-overlap -check-prefixes=CHECK,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"

define noundef i32 @NormalRet() nounwind uwtable sanitize_memory {
; CHECK-LABEL: @NormalRet(
; CHECK-NEXT:    ret i32 123
;
  ret i32 123
}

define i32 @PartialRet() nounwind uwtable sanitize_memory {
; CHECK-LABEL: @PartialRet(
; CHECK-NEXT:    store i32 0, i32* bitcast ([100 x i64]* @__msan_retval_tls to i32*), align 8
; CHECK-NEXT:    store i32 0, i32* @__msan_retval_origin_tls, align 4
; CHECK-NEXT:    ret i32 123
;
  ret i32 123
}

define noundef i32 @LoadedRet() nounwind uwtable sanitize_memory {
; CHECK-LABEL: @LoadedRet(
; CHECK-NEXT:    [[P:%.*]] = inttoptr i64 0 to i32*
; CHECK-NEXT:    [[O:%.*]] = load i32, i32* [[P]], align 4
; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint i32* [[P]] to i64
; CHECK-NEXT:    [[TMP2:%.*]] = xor i64 [[TMP1]], 87960930222080
; CHECK-NEXT:    [[TMP3:%.*]] = inttoptr i64 [[TMP2]] to i32*
; CHECK-NEXT:    [[TMP4:%.*]] = add i64 [[TMP2]], 17592186044416
; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to i32*
; CHECK-NEXT:    [[_MSLD:%.*]] = load i32, i32* [[TMP3]], align 4
; CHECK-NEXT:    [[TMP6:%.*]] = load i32, i32* [[TMP5]], align 4
; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i32 [[_MSLD]], 0
; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP7:%.*]], label [[TMP8:%.*]], !prof !0
; CHECK:       7:
; CHECK-NEXT:    call void @__msan_warning_with_origin_noreturn(i32 [[TMP6]]) #1
; CHECK-NEXT:    unreachable
; CHECK:       8:
; CHECK-NEXT:    ret i32 [[O]]
;
  %p = inttoptr i64 0 to i32 *
  %o = load i32, i32 *%p
  ret i32 %o
}


define void @NormalArg(i32 noundef %a) nounwind uwtable sanitize_memory {
; CHECK-LABEL: @NormalArg(
; CHECK-NEXT:    [[P:%.*]] = inttoptr i64 0 to i32*
; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint i32* [[P]] to i64
; CHECK-NEXT:    [[TMP2:%.*]] = xor i64 [[TMP1]], 87960930222080
; CHECK-NEXT:    [[TMP3:%.*]] = inttoptr i64 [[TMP2]] to i32*
; CHECK-NEXT:    [[TMP4:%.*]] = add i64 [[TMP2]], 17592186044416
; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to i32*
; CHECK-NEXT:    store i32 0, i32* [[TMP3]], align 4
; CHECK-NEXT:    store i32 [[A:%.*]], i32* [[P]], align 4
; CHECK-NEXT:    ret void
;
  %p = inttoptr i64 0 to i32 *
  store i32 %a, i32 *%p
  ret void
}

define void @PartialArg(i32 %a) nounwind uwtable sanitize_memory {
; CHECK-LABEL: @PartialArg(
; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* bitcast ([100 x i64]* @__msan_param_tls to i32*), align 8
; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__msan_param_origin_tls, i32 0, i32 0), align 4
; CHECK-NEXT:    [[P:%.*]] = inttoptr i64 0 to i32*
; CHECK-NEXT:    [[TMP3:%.*]] = ptrtoint i32* [[P]] to i64
; CHECK-NEXT:    [[TMP4:%.*]] = xor i64 [[TMP3]], 87960930222080
; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to i32*
; CHECK-NEXT:    [[TMP6:%.*]] = add i64 [[TMP4]], 17592186044416
; CHECK-NEXT:    [[TMP7:%.*]] = inttoptr i64 [[TMP6]] to i32*
; CHECK-NEXT:    store i32 [[TMP1]], i32* [[TMP5]], align 4
; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i32 [[TMP1]], 0
; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP8:%.*]], label [[TMP9:%.*]], !prof !0
; CHECK:       8:
; CHECK-NEXT:    store i32 [[TMP2]], i32* [[TMP7]], align 4
; CHECK-NEXT:    br label [[TMP9]]
; CHECK:       9:
; CHECK-NEXT:    store i32 [[A:%.*]], i32* [[P]], align 4
; CHECK-NEXT:    ret void
;
  %p = inttoptr i64 0 to i32 *
  store i32 %a, i32 *%p
  ret void
}

define void @CallNormal() nounwind uwtable sanitize_memory {
; CHECK-LABEL: @CallNormal(
; CHECK-NEXT:    [[R:%.*]] = call i32 @NormalRet() #0
; CHECK-NEXT:    call void @NormalArg(i32 [[R]]) #0
; CHECK-NEXT:    ret void
;
  %r = call i32 @NormalRet() nounwind uwtable sanitize_memory
  call void @NormalArg(i32 %r) nounwind uwtable sanitize_memory
  ret void
}

define void @CallWithLoaded() nounwind uwtable sanitize_memory {
; CHECK-LABEL: @CallWithLoaded(
; CHECK-NEXT:    [[P:%.*]] = inttoptr i64 0 to i32*
; CHECK-NEXT:    [[O:%.*]] = load i32, i32* [[P]], align 4
; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint i32* [[P]] to i64
; CHECK-NEXT:    [[TMP2:%.*]] = xor i64 [[TMP1]], 87960930222080
; CHECK-NEXT:    [[TMP3:%.*]] = inttoptr i64 [[TMP2]] to i32*
; CHECK-NEXT:    [[TMP4:%.*]] = add i64 [[TMP2]], 17592186044416
; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to i32*
; CHECK-NEXT:    [[_MSLD:%.*]] = load i32, i32* [[TMP3]], align 4
; CHECK-NEXT:    [[TMP6:%.*]] = load i32, i32* [[TMP5]], align 4
; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i32 [[_MSLD]], 0
; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP7:%.*]], label [[TMP8:%.*]], !prof !0
; CHECK:       7:
; CHECK-NEXT:    call void @__msan_warning_with_origin_noreturn(i32 [[TMP6]]) #1
; CHECK-NEXT:    unreachable
; CHECK:       8:
; CHECK-NEXT:    call void @NormalArg(i32 [[O]]) #0
; CHECK-NEXT:    ret void
;
  %p = inttoptr i64 0 to i32 *
  %o = load i32, i32 *%p
  call void @NormalArg(i32 %o) nounwind uwtable sanitize_memory
  ret void
}

define void @CallPartial() nounwind uwtable sanitize_memory {
; CHECK-LABEL: @CallPartial(
; CHECK-NEXT:    store i32 0, i32* bitcast ([100 x i64]* @__msan_retval_tls to i32*), align 8
; CHECK-NEXT:    [[R:%.*]] = call i32 @PartialRet() #0
; CHECK-NEXT:    [[_MSRET:%.*]] = load i32, i32* bitcast ([100 x i64]* @__msan_retval_tls to i32*), align 8
; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* @__msan_retval_origin_tls, align 4
; CHECK-NEXT:    store i32 [[_MSRET]], i32* bitcast ([100 x i64]* @__msan_param_tls to i32*), align 8
; CHECK-NEXT:    store i32 [[TMP1]], i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__msan_param_origin_tls, i32 0, i32 0), align 4
; CHECK-NEXT:    call void @PartialArg(i32 [[R]]) #0
; CHECK-NEXT:    ret void
;
  %r = call i32 @PartialRet() nounwind uwtable sanitize_memory
  call void @PartialArg(i32 %r) nounwind uwtable sanitize_memory
  ret void
}

[guiand]

Use autogenerated test

[guiand]

llvm/test/Instrumentation/MemorySanitizer/msan_eager.ll

2

Thanks for helping me out with this, Vitaly!

[eugenis]

LGTM with a nit

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

3416

Remove "else".

3525–3534

This logic is a bit confusing.
Maybe move the "main" check into this function, and express it as

bool EagerCheck = ((ClEagerChecks && hasAttribute) || (name == "main"));
bool StoreShadow = !(ClEagerChecks && hasAttribute);

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

 // a garbage address typically results in SEGV, but still happen
 // (e.g. only lower bits of address are garbage, or the access happens
 // early at program startup where malloc-ed memory is more likely to
 // be zeroed. As of 2012-08-28 this flag adds 20% slowdown.
 static cl::opt<bool> ClCheckAccessAddress("msan-check-access-address",
        cl::desc("report accesses through a pointer which has poisoned shadow"),
        cl::Hidden, cl::init(true));
 
+static cl::opt<bool> ClEagerChecks(
+    "msan-eager-checks",
+    cl::desc("check arguments and return values at function call boundaries"),
+    cl::Hidden, cl::init(false));
+
 static cl::opt<bool> ClDumpStrictInstructions("msan-dump-strict-instructions",
        cl::desc("print out instructions with default strict semantics"),
        cl::Hidden, cl::init(false));
 
 static cl::opt<int> ClInstrumentationWithCallThreshold(
     "msan-instrumentation-with-call-threshold",
     cl::desc(
         "If the function being instrumented requires more than "
void MemorySanitizer::initializeCallbacks(Module &M) {
   MemsetFn = M.getOrInsertFunction(
     "__msan_memset", IRB.getInt8PtrTy(), IRB.getInt8PtrTy(), IRB.getInt32Ty(),
     IntptrTy);
 
   MsanInstrumentAsmStoreFn =
       M.getOrInsertFunction("__msan_instrument_asm_store", IRB.getVoidTy(),
                             PointerType::get(IRB.getInt8Ty(), 0), IntptrTy);
 
   if (CompileKernel) {
     createKernelApi(M);
   } else {
     createUserspaceApi(M);
   }
   CallbacksInitialized = true;
 }
 
 FunctionCallee MemorySanitizer::getKmsanShadowOriginAccessFn(bool isStore,
                                                              int size) {
   FunctionCallee *Fns =
       isStore ? MsanMetadataPtrForStore_1_8 : MsanMetadataPtrForLoad_1_8;
   switch (size) {
   case 1:
     return Fns[0];
   case 2:
     return Fns[1];
   case 4:
     return Fns[2];
   case 8:
     return Fns[3];
   default:
     return nullptr;
   }
 }
 
 /// Module-level initialization.
if (Recover)
         return new GlobalVariable(M, IRB.getInt32Ty(), true,
                                   GlobalValue::WeakODRLinkage,
                                   IRB.getInt32(Recover), "__msan_keep_going");
       });
 }
 }
 
 bool MemorySanitizerLegacyPass::doInitialization(Module &M) {
   if (!Options.Kernel)
     insertModuleCtor(M);
   MSan.emplace(M, Options);
   return true;
 }
 
 namespace {
 
 /// A helper class that handles instrumentation of VarArg
struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   BasicBlock *ActualFnStart;
 
   // The following flags disable parts of MSan instrumentation based on
   // exclusion list contents and command-line options.
   bool InsertChecks;
   bool PropagateShadow;
   bool PoisonStack;
   bool PoisonUndef;
-  bool CheckReturnValue;
 
   struct ShadowOriginAndInsertPoint {
     Value *Shadow;
     Value *Origin;
     Instruction *OrigIns;
 
     ShadowOriginAndInsertPoint(Value *S, Value *O, Instruction *I)
       : Shadow(S), Origin(O), OrigIns(I) {}
   };
   SmallVector<ShadowOriginAndInsertPoint, 16> InstrumentationList;
   bool InstrumentLifetimeStart = ClHandleLifetimeIntrinsics;
   SmallSet<AllocaInst *, 16> AllocaSet;
   SmallVector<std::pair<IntrinsicInst *, AllocaInst *>, 16> LifetimeStartList;
   SmallVector<StoreInst *, 16> StoreList;
 
   MemorySanitizerVisitor(Function &F, MemorySanitizer &MS,
                          const TargetLibraryInfo &TLI)
       : F(F), MS(MS), VAHelper(CreateVarArgHelper(F, MS, *this)), TLI(&TLI) {
     bool SanitizeFunction = F.hasFnAttribute(Attribute::SanitizeMemory);
     InsertChecks = SanitizeFunction;
     PropagateShadow = SanitizeFunction;
     PoisonStack = SanitizeFunction && ClPoisonStack;
     PoisonUndef = SanitizeFunction && ClPoisonUndef;
-    // FIXME: Consider using SpecialCaseList to specify a list of functions that
-    // must always return fully initialized values. For now, we hardcode "main".
-    CheckReturnValue = SanitizeFunction && (F.getName() == "main");
 
     MS.initializeCallbacks(*F.getParent());
     if (MS.CompileKernel)
       ActualFnStart = insertKmsanPrologue(F);
     else
       ActualFnStart = &F.getEntryBlock();
 
     LLVM_DEBUG(if (!InsertChecks) dbgs()
                << "MemorySanitizer is not inserting checks into '"
void insertWarningFn(IRBuilder<> &IRB, Value *Origin) {
     // FIXME: Insert UnreachableInst if !MS.Recover?
     // This may invalidate some of the following checks and needs to be done
     // at the very end.
   }
 
   void materializeOneCheck(Instruction *OrigIns, Value *Shadow, Value *Origin,
                            bool AsCall) {
     IRBuilder<> IRB(OrigIns);
     LLVM_DEBUG(dbgs() << "  SHAD0 : " << *Shadow << "\n");
     Value *ConvertedShadow = convertToShadowTyNoVec(Shadow, IRB);
     LLVM_DEBUG(dbgs() << "  SHAD1 : " << *ConvertedShadow << "\n");
 
     if (auto *ConstantShadow = dyn_cast<Constant>(ConvertedShadow)) {
       if (ClCheckConstantShadow && !ConstantShadow->isZeroValue()) {
         insertWarningFn(IRB, Origin);
       }
       return;
     }
 
     const DataLayout &DL = OrigIns->getModule()->getDataLayout();
 
     unsigned TypeSizeInBits = DL.getTypeSizeInBits(ConvertedShadow->getType());
     unsigned SizeIndex = TypeSizeToSizeIndex(TypeSizeInBits);
     if (AsCall && SizeIndex < kNumberOfAccessSizes && !MS.CompileKernel) {
       FunctionCallee Fn = MS.MaybeWarningFn[SizeIndex];
       Value *ConvertedShadow2 =
           IRB.CreateZExt(ConvertedShadow, IRB.getIntNTy(8 * (1 << SizeIndex)));
       IRB.CreateCall(Fn, {ConvertedShadow2, MS.TrackOrigins && Origin
                                                 ? Origin
                                                 : (Value *)IRB.getInt32(0)});
if (Argument *A = dyn_cast<Argument>(V)) {
       IRBuilder<> EntryIRB(ActualFnStart->getFirstNonPHI());
       unsigned ArgOffset = 0;
       const DataLayout &DL = F->getParent()->getDataLayout();
       for (auto &FArg : F->args()) {
         if (!FArg.getType()->isSized()) {
           LLVM_DEBUG(dbgs() << "Arg is not sized\n");
           continue;
         }
+
+        bool FArgByVal = FArg.hasByValAttr();
+        bool FArgNoUndef = FArg.hasAttribute(Attribute::NoUndef);
+        bool FArgEagerCheck = ClEagerChecks && !FArgByVal && FArgNoUndef;
         unsigned Size =
             FArg.hasByValAttr()
                 ? DL.getTypeAllocSize(FArg.getParamByValType())
                 : DL.getTypeAllocSize(FArg.getType());
+
         if (A == &FArg) {
           bool Overflow = ArgOffset + Size > kParamTLSSize;
+          if (FArgEagerCheck) {
+            *ShadowPtr = getCleanShadow(V);
+            setOrigin(A, getCleanOrigin());
+            continue;
+          } else if (FArgByVal) {
             Value *Base = getShadowPtrForArgument(&FArg, EntryIRB, ArgOffset);
-          if (FArg.hasByValAttr()) {
             // ByVal pointer itself has clean shadow. We copy the actual
             // argument shadow to the underlying memory.
             // Figure out maximal valid memcpy alignment.
             const Align ArgAlign = DL.getValueOrABITypeAlignment(
                 MaybeAlign(FArg.getParamAlignment()), FArg.getParamByValType());
             Value *CpShadowPtr =
                 getShadowOriginPtr(V, EntryIRB, EntryIRB.getInt8Ty(), ArgAlign,
                                    /*isStore*/ true)
                     .first;
             // TODO(glider): need to copy origins.
             if (Overflow) {
               // ParamTLS overflow.
               EntryIRB.CreateMemSet(
                   CpShadowPtr, Constant::getNullValue(EntryIRB.getInt8Ty()),
                   Size, ArgAlign);
             } else {
               const Align CopyAlign = std::min(ArgAlign, kShadowTLSAlignment);
               Value *Cpy = EntryIRB.CreateMemCpy(CpShadowPtr, CopyAlign, Base,
                                                  CopyAlign, Size);
               LLVM_DEBUG(dbgs() << "  ByValCpy: " << *Cpy << "\n");
               (void)Cpy;
             }
             *ShadowPtr = getCleanShadow(V);
           } else {
+            // Shadow over TLS
+            Value *Base = getShadowPtrForArgument(&FArg, EntryIRB, ArgOffset);
             if (Overflow) {
               // ParamTLS overflow.
               *ShadowPtr = getCleanShadow(V);
             } else {
               *ShadowPtr = EntryIRB.CreateAlignedLoad(getShadowTy(&FArg), Base,
                                                       kShadowTLSAlignment);
             }
           }
           LLVM_DEBUG(dbgs()
                      << "  ARG:    " << FArg << " ==> " << **ShadowPtr << "\n");
           if (MS.TrackOrigins && !Overflow) {
             Value *OriginPtr =
                 getOriginPtrForArgument(&FArg, EntryIRB, ArgOffset);
             setOrigin(A, EntryIRB.CreateLoad(MS.OriginTy, OriginPtr));
           } else {
             setOrigin(A, getCleanOrigin());
           }
         }
+
+        if (!FArgEagerCheck)
           ArgOffset += alignTo(Size, kShadowTLSAlignment);
       }
       assert(*ShadowPtr && "Could not find shadow for an argument");
       return *ShadowPtr;
     }
     // For everything else the shadow is zero.
     return getCleanShadow(V);
   }
 
for (auto ArgIt = CB.arg_begin(), End = CB.arg_end(); ArgIt != End;
       // in that case getShadow() will copy the actual arg shadow to
       // __msan_param_tls.
       Value *ArgShadow = getShadow(A);
       Value *ArgShadowBase = getShadowPtrForArgument(A, IRB, ArgOffset);
       LLVM_DEBUG(dbgs() << "  Arg#" << i << ": " << *A
                         << " Shadow: " << *ArgShadow << "\n");
       bool ArgIsInitialized = false;
       const DataLayout &DL = F.getParent()->getDataLayout();
-      if (CB.paramHasAttr(i, Attribute::ByVal)) {
+
+      bool ByVal = CB.paramHasAttr(i, Attribute::ByVal);
+      bool NoUndef = CB.paramHasAttr(i, Attribute::NoUndef);
+      bool EagerCheck = ClEagerChecks && !ByVal && NoUndef;
+
+      if (EagerCheck) {
+        insertShadowCheck(A, &CB);
+        continue;
+      }
+      if (ByVal) {
+        // ByVal requires some special handling as it's too big for a single
+        // load
         assert(A->getType()->isPointerTy() &&
                "ByVal argument is not a pointer!");
         Size = DL.getTypeAllocSize(CB.getParamByValType(i));
         if (ArgOffset + Size > kParamTLSSize) break;
         const MaybeAlign ParamAlignment(CB.getParamAlign(i));
         MaybeAlign Alignment = llvm::None;
         if (ParamAlignment)
           Alignment = std::min(*ParamAlignment, kShadowTLSAlignment);
         Value *AShadowPtr =
             getShadowOriginPtr(A, IRB, IRB.getInt8Ty(), Alignment,
                                /*isStore*/ false)
                 .first;
 
         Store = IRB.CreateMemCpy(ArgShadowBase, Alignment, AShadowPtr,
                                  Alignment, Size);
         // TODO(glider): need to copy origins.
       } else {
+        // Any other parameters mean we need bit-grained tracking of uninit data
         Size = DL.getTypeAllocSize(A->getType());
         if (ArgOffset + Size > kParamTLSSize) break;
         Store = IRB.CreateAlignedStore(ArgShadow, ArgShadowBase,
                                        kShadowTLSAlignment);
         Constant *Cst = dyn_cast<Constant>(ArgShadow);
         if (Cst && Cst->isNullValue()) ArgIsInitialized = true;
       }
       if (MS.TrackOrigins && !ArgIsInitialized)
         IRB.CreateStore(getOrigin(A),
                         getOriginPtrForArgument(A, IRB, ArgOffset));
       (void)Store;
       assert(Size != 0 && Store != nullptr);
       LLVM_DEBUG(dbgs() << "  Param:" << *Store << "\n");
       ArgOffset += alignTo(Size, 8);
     }
     LLVM_DEBUG(dbgs() << "  done with call args\n");
 
     FunctionType *FT = CB.getFunctionType();
     if (FT->isVarArg()) {
       VAHelper->visitCallBase(CB, IRB);
     }
 
     // Now, get the shadow for the RetVal.
     if (!CB.getType()->isSized())
       return;
     // Don't emit the epilogue for musttail call returns.
     if (isa<CallInst>(CB) && cast<CallInst>(CB).isMustTailCall())
       return;
+
+    if (ClEagerChecks && CB.hasRetAttr(Attribute::NoUndef)) {
+      setShadow(&CB, getCleanShadow(&CB));
+      setOrigin(&CB, getCleanOrigin());
+      return;
+    }
+
     IRBuilder<> IRBBefore(&CB);
     // Until we have full dynamic coverage, make sure the retval shadow is 0.
     Value *Base = getShadowPtrForRetval(&CB, IRBBefore);
     IRBBefore.CreateAlignedStore(getCleanShadow(&CB), Base,
                                  kShadowTLSAlignment);
     BasicBlock::iterator NextInsn;
     if (isa<CallInst>(CB)) {
       NextInsn = ++CB.getIterator();

 
   void visitReturnInst(ReturnInst &I) {
     IRBuilder<> IRB(&I);
     Value *RetVal = I.getReturnValue();
     if (!RetVal) return;
     // Don't emit the epilogue for musttail call returns.
     if (isAMustTailRetVal(RetVal)) return;
     Value *ShadowPtr = getShadowPtrForRetval(RetVal, IRB);
-    if (CheckReturnValue) {
-      insertShadowCheck(RetVal, &I);
-      Value *Shadow = getCleanShadow(RetVal);
-      IRB.CreateAlignedStore(Shadow, ShadowPtr, kShadowTLSAlignment);
-    } else {
-      Value *Shadow = getShadow(RetVal);
+    bool HasNoUndef =
+        F.hasAttribute(AttributeList::ReturnIndex, Attribute::NoUndef);
+    bool StoreShadow = !(ClEagerChecks && HasNoUndef);
+    // FIXME: Consider using SpecialCaseList to specify a list of functions that
+    // must always return fully initialized values. For now, we hardcode "main".
+    bool EagerCheck = (ClEagerChecks && HasNoUndef) || (F.getName() == "main");
+
+    bool StoreOrigin = true;
+    if (EagerCheck) {
+      insertShadowCheck(RetVal, &I);
+      Shadow = getCleanShadow(RetVal);
+      StoreOrigin = false;
+    }
+
+    // The caller may still expect information passed over TLS if we pass our
+    // check
+    if (StoreShadow) {
       IRB.CreateAlignedStore(Shadow, ShadowPtr, kShadowTLSAlignment);
-      if (MS.TrackOrigins)
+      if (MS.TrackOrigins && StoreOrigin)
         IRB.CreateStore(getOrigin(RetVal), getOriginPtrForRetval(IRB));
     }
   }
 
   void visitPHINode(PHINode &I) {
     IRBuilder<> IRB(&I);
     if (!PropagateShadow) {
       setShadow(&I, getCleanShadow(&I));

llvm/test/Instrumentation/MemorySanitizer/msan_eager.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -msan-check-access-address=0 -msan-track-origins=1 -msan-eager-checks -S -passes='module(msan-module),function(msan)' 2>&1 | \
+; RUN:   FileCheck -allow-deprecated-dag-overlap -check-prefixes=CHECK,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"
+
+define noundef i32 @NormalRet() nounwind uwtable sanitize_memory {
+; CHECK-LABEL: @NormalRet(
+; CHECK-NEXT:    ret i32 123
+;
+  ret i32 123
+}
+
+define i32 @PartialRet() nounwind uwtable sanitize_memory {
+; CHECK-LABEL: @PartialRet(
+; CHECK-NEXT:    store i32 0, i32* bitcast ([100 x i64]* @__msan_retval_tls to i32*), align 8
+; CHECK-NEXT:    store i32 0, i32* @__msan_retval_origin_tls, align 4
+; CHECK-NEXT:    ret i32 123
+;
+  ret i32 123
+}
+
+define noundef i32 @LoadedRet() nounwind uwtable sanitize_memory {
+; CHECK-LABEL: @LoadedRet(
+; CHECK-NEXT:    [[P:%.*]] = inttoptr i64 0 to i32*
+; CHECK-NEXT:    [[O:%.*]] = load i32, i32* [[P]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint i32* [[P]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = xor i64 [[TMP1]], 87960930222080
+; CHECK-NEXT:    [[TMP3:%.*]] = inttoptr i64 [[TMP2]] to i32*
+; CHECK-NEXT:    [[TMP4:%.*]] = add i64 [[TMP2]], 17592186044416
+; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to i32*
+; CHECK-NEXT:    [[_MSLD:%.*]] = load i32, i32* [[TMP3]], align 4
+; CHECK-NEXT:    [[TMP6:%.*]] = load i32, i32* [[TMP5]], align 4
+; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i32 [[_MSLD]], 0
+; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP7:%.*]], label [[TMP8:%.*]], !prof !0
+; CHECK:       7:
+; CHECK-NEXT:    call void @__msan_warning_with_origin_noreturn(i32 [[TMP6]]) #1
+; CHECK-NEXT:    unreachable
+; CHECK:       8:
+; CHECK-NEXT:    ret i32 [[O]]
+;
+  %p = inttoptr i64 0 to i32 *
+  %o = load i32, i32 *%p
+  ret i32 %o
+}
+
+
+define void @NormalArg(i32 noundef %a) nounwind uwtable sanitize_memory {
+; CHECK-LABEL: @NormalArg(
+; CHECK-NEXT:    [[P:%.*]] = inttoptr i64 0 to i32*
+; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint i32* [[P]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = xor i64 [[TMP1]], 87960930222080
+; CHECK-NEXT:    [[TMP3:%.*]] = inttoptr i64 [[TMP2]] to i32*
+; CHECK-NEXT:    [[TMP4:%.*]] = add i64 [[TMP2]], 17592186044416
+; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to i32*
+; CHECK-NEXT:    store i32 0, i32* [[TMP3]], align 4
+; CHECK-NEXT:    store i32 [[A:%.*]], i32* [[P]], align 4
+; CHECK-NEXT:    ret void
+;
+  %p = inttoptr i64 0 to i32 *
+  store i32 %a, i32 *%p
+  ret void
+}
+
+define void @PartialArg(i32 %a) nounwind uwtable sanitize_memory {
+; CHECK-LABEL: @PartialArg(
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* bitcast ([100 x i64]* @__msan_param_tls to i32*), align 8
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__msan_param_origin_tls, i32 0, i32 0), align 4
+; CHECK-NEXT:    [[P:%.*]] = inttoptr i64 0 to i32*
+; CHECK-NEXT:    [[TMP3:%.*]] = ptrtoint i32* [[P]] to i64
+; CHECK-NEXT:    [[TMP4:%.*]] = xor i64 [[TMP3]], 87960930222080
+; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to i32*
+; CHECK-NEXT:    [[TMP6:%.*]] = add i64 [[TMP4]], 17592186044416
+; CHECK-NEXT:    [[TMP7:%.*]] = inttoptr i64 [[TMP6]] to i32*
+; CHECK-NEXT:    store i32 [[TMP1]], i32* [[TMP5]], align 4
+; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i32 [[TMP1]], 0
+; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP8:%.*]], label [[TMP9:%.*]], !prof !0
+; CHECK:       8:
+; CHECK-NEXT:    store i32 [[TMP2]], i32* [[TMP7]], align 4
+; CHECK-NEXT:    br label [[TMP9]]
+; CHECK:       9:
+; CHECK-NEXT:    store i32 [[A:%.*]], i32* [[P]], align 4
+; CHECK-NEXT:    ret void
+;
+  %p = inttoptr i64 0 to i32 *
+  store i32 %a, i32 *%p
+  ret void
+}
+
+define void @CallNormal() nounwind uwtable sanitize_memory {
+; CHECK-LABEL: @CallNormal(
+; CHECK-NEXT:    [[R:%.*]] = call i32 @NormalRet() #0
+; CHECK-NEXT:    call void @NormalArg(i32 [[R]]) #0
+; CHECK-NEXT:    ret void
+;
+  %r = call i32 @NormalRet() nounwind uwtable sanitize_memory
+  call void @NormalArg(i32 %r) nounwind uwtable sanitize_memory
+  ret void
+}
+
+define void @CallWithLoaded() nounwind uwtable sanitize_memory {
+; CHECK-LABEL: @CallWithLoaded(
+; CHECK-NEXT:    [[P:%.*]] = inttoptr i64 0 to i32*
+; CHECK-NEXT:    [[O:%.*]] = load i32, i32* [[P]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = ptrtoint i32* [[P]] to i64
+; CHECK-NEXT:    [[TMP2:%.*]] = xor i64 [[TMP1]], 87960930222080
+; CHECK-NEXT:    [[TMP3:%.*]] = inttoptr i64 [[TMP2]] to i32*
+; CHECK-NEXT:    [[TMP4:%.*]] = add i64 [[TMP2]], 17592186044416
+; CHECK-NEXT:    [[TMP5:%.*]] = inttoptr i64 [[TMP4]] to i32*
+; CHECK-NEXT:    [[_MSLD:%.*]] = load i32, i32* [[TMP3]], align 4
+; CHECK-NEXT:    [[TMP6:%.*]] = load i32, i32* [[TMP5]], align 4
+; CHECK-NEXT:    [[_MSCMP:%.*]] = icmp ne i32 [[_MSLD]], 0
+; CHECK-NEXT:    br i1 [[_MSCMP]], label [[TMP7:%.*]], label [[TMP8:%.*]], !prof !0
+; CHECK:       7:
+; CHECK-NEXT:    call void @__msan_warning_with_origin_noreturn(i32 [[TMP6]]) #1
+; CHECK-NEXT:    unreachable
+; CHECK:       8:
+; CHECK-NEXT:    call void @NormalArg(i32 [[O]]) #0
+; CHECK-NEXT:    ret void
+;
+  %p = inttoptr i64 0 to i32 *
+  %o = load i32, i32 *%p
+  call void @NormalArg(i32 %o) nounwind uwtable sanitize_memory
+  ret void
+}
+
+define void @CallPartial() nounwind uwtable sanitize_memory {
+; CHECK-LABEL: @CallPartial(
+; CHECK-NEXT:    store i32 0, i32* bitcast ([100 x i64]* @__msan_retval_tls to i32*), align 8
+; CHECK-NEXT:    [[R:%.*]] = call i32 @PartialRet() #0
+; CHECK-NEXT:    [[_MSRET:%.*]] = load i32, i32* bitcast ([100 x i64]* @__msan_retval_tls to i32*), align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* @__msan_retval_origin_tls, align 4
+; CHECK-NEXT:    store i32 [[_MSRET]], i32* bitcast ([100 x i64]* @__msan_param_tls to i32*), align 8
+; CHECK-NEXT:    store i32 [[TMP1]], i32* getelementptr inbounds ([200 x i32], [200 x i32]* @__msan_param_origin_tls, i32 0, i32 0), align 4
+; CHECK-NEXT:    call void @PartialArg(i32 [[R]]) #0
+; CHECK-NEXT:    ret void
+;
+  %r = call i32 @PartialRet() nounwind uwtable sanitize_memory
+  call void @PartialArg(i32 %r) nounwind uwtable sanitize_memory
+  ret void
+}