Add KMSAN instrumentation to MSan pass

[glider]

Summary

Introduce the -msan-kernel flag, which enables the kernel instrumentation.

The main differences between KMSAN and MSan instrumentations are:

• KMSAN implies msan-track-origins=2, msan-keep-going=true;
• there're no explicit accesses to shadow and origin memory. Shadow and origin values for a particular X-byte memory location are read and written via pointers returned by __msan_metadata_ptr_for_load_X(u8 *addr) and __msan_store_shadow_origin_X(u8 *addr, uptr shadow, uptr origin);
• TLS variables are stored in a single struct in per-task storage. A call to a function returning that struct is inserted into every instrumented function before the entry block;
• __msan_warning() takes a 32-bit origin parameter;
• local variables are poisoned with __msan_poison_alloca() upon function entry and unpoisoned with __msan_unpoison_alloca() before leaving the function;
• the pass doesn't declare any global variables or add global constructors to the translation unit.

[dvyukov]

I know, not the feedback you was looking for :)

lib/Transforms/Instrumentation/MemorySanitizer.cpp

455

Looks like a bug in clang-format. Why is it 3 spaces here? All other members start with 2 and 2 is the proper style.

[dvyukov]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

234

I see that for interface part (flags, pass ctor) you copied KASAN approach, so this looks good to me.

638

ShadowOriginStruct seems to be used only in the loop below. So drop it and use the actual types in the loop.

679

The interface looks bulky because of distinction between arg/param/retval/va_arg. It feels that we put too much details into the interface, such interfaces are usually unstable. Is it possible to combine all that data into a single struct including shadow and origins (for user-space as well) and then have just single function that requests data at offsets of that struct? Or maybe just returns a pointer to the thread-local struct. I should also reduce amount of differences between kernel and userspace.

We could go further and reuse __kmsan_load/store_shadow_origin functions by giving then some magic address that mean TLS. Not sure if it will be good or not.

[dvyukov]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

679

Note that we can't version compiler and kernel atomically. This interface will have to stay basically forever. If we will want to improve it later, we will only make situation worse because we will have to support 2 interfaces.

[dvyukov]

Evgenii, please take a look at actual pass changes.

[glider]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

689

@eugenis: Dima suggested we merge all msan_XXX_tls globals into a single struct, this way I can replace all the kmsan_get_XXX_tls() callbacks with a single one.
WDYT?

[glider]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

455

I think this is being caused by me formatting the diff, not the whole file.
clang-format handles this case correctly in the wild.

638

This is a legacy, will remove.

[eugenis]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

679

That would be an ABI break on Linux. And hiding the arg/param/retval/etc details in a struct does not make the interface any more stable.
On the other hand, it makes the interface simpler, so if you think you can do it just for the kernel w/o complicating the runtime library too much, then go for it.

679

Are you sure there is no way to version compiler and the runtime atomically? As I imagine, the part that must be in sync with the kernel is basically shadow mapping and the bool are_we_initialized_already() function. Is it possible to make that the forever-stable interface, and tie everything on top of it to the compiler?

[glider]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

679

Are you sure there is no way to version compiler and the runtime atomically?

If we're talking about the runtime residing in the upstream kernel, then yes.
Clang versions live long, so if we ever introduce a new API version, then the kernel will have to support all the existing versions.

As I imagine, the part that must be in sync with the kernel is basically shadow mapping and the bool are_we_initialized_already() function.

Neither is currently part of the compiler API. The kernel allocates the shadow on per-page basis (i.e. shadow is non-contiguous), and these details better be abstracted from the instrumentation.

Is it possible to make that the forever-stable interface, and tie everything on top of it to the compiler?

If we treat the userspace shadow mapping as an example of such a stable interface, we see that it's not enough (e.g. handling va_args require additional globals or functions).
The question is basically whether right now MSan interface is stable enough, or will there be changes to it in the future.

679

On the other hand, it makes the interface simpler, so if you think you can do it just for the kernel w/o complicating the runtime library too much, then go for it.

Ok, I'll make this change just for the kernel.

[glider]

Updated the patch.
Now all KMSAN-specific functions bear the "__msan_" prefix so that they can be used in other contexts (e.g. MSan instrumentation for Android).
I've also dropped aliasing functions that were working with argument shadow, except for __msan_load_arg_shadow() and __msan_store_arg_shadow().

Please take another look.

[eugenis]

In general, way too many small pieces of kernel-specific code. Please try to come up with a set of common utility functions that would capture the difference between kernel and userspace.

lib/Transforms/Instrumentation/MemorySanitizer.cpp

1139

There is some code duplication between this function and materializeStoresUserspace. It could be better to implement a kernel-specific storeOrigin() and keep this one common.

1154

When can origin be less than 32 bit?

1702

I think this is the wrong level of abstraction. You don't want kernel-specific handling of LoadInst, you want kernel-specific primitives for writing and reading shadow. You can use them in handleVectorLoadIntrinsic as well as other similar places. Ideally, there should be no checks for MS.CompileKernel outside of such common utility functgions.

3104

How can it be? The call to this function is inserted in the Kmsan prologue, and we specifically step over the prologue when running the visitor in runOnFunction.

[glider]

Updated the diff.
This is still work in progress, and I'm not sure what's the best way to proceed with it.
The diff of the instrumentation pass patch itself is around 1000 LoC, plus I've moved around some IR tests.
(In particular the tests that behave differently with and without -msan-kernel=1 were put into separate files, while those agnostic to the memory layout remained in msan_common_basic.ll)

We can try moving forward with this patch as a whole, or I can start landing small chunks of it just for the userspace (e.g. replace F.getEntryBlock() with MSV.ActualFnStart everywhere, factor some code into helper functions) so that it'll be easier to add a kernel-specific implementation to them.

[glider]

Please disregard. Looks like I've missed some opportunities to move the CompileKernel() checks to the lower level.

[glider]

Updated the diff.
This is more or less the final version of the KMSAN API. I'm still working on the tests and fixing minor nits.
I'll be OOO till the end of April, so take your time.

[eugenis]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

444

Lets make -mllvm flags take precedence, as in KASan (see ClRecover.getNumOccurrences() > 0).

675

What are these, why do they have "_arg_" in the name but not in the description, and why are they not used?

[glider]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

675

Removed these.

[glider]

Major update after landing the kernel-inspecific parts.

[eugenis]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

754–758

Please invert this check.

919–920

Please rebase on ToT.

1096–1102

Invert this check, and a few others below.

1337

I think you want DL.getTypeStoreSize() here.

1520

std::ignore instead of CpOriginPtr to avoid unused variable warning? Or is the compiler not smart enough to see that it is unused?

3294

Could you explain why this function is passed return address (i.e. the caller pc)?

Why is this api so different from userspace?

3497

I think you can detect this by looking at function attributes.

3580

This is not about kernel at all, is it?
Could you fix it for userspace as well, as a separate change?

test/Instrumentation/MemorySanitizer/msan_kernel_basic.ll

17

AFAIK, in no-asserts build IR temporary names are lost, so this "%param_shadow" is not going to match.

45

Please add checks for actual arguments etc. For example, make sure that the right member of context_state is called for param_shadow. No need to repeat this in every test case, just once would be enough.

[glider]

Rebased to r339138.

[glider]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

1520

Looks like it isn't.
But I've replaced the variable with std::ignore anyway.

3294

It's too expensive to unwind the stack for each alloca, and __builtin_return_address(1) is unreliable, therefore we take the caller's return address to get two stack frames instead of one.

3497

I've only found mentions of SSE in the "target-features" string, where they look like: "-sse,-sse2,-sse3,...", not sure it's easier to parse those.

3580

This is gonna change the userspace ABI, as we'll need to introduce the origin tls var.
I can sure fix it for the userspace, we'll just need to be accurate deploying the new version.

test/Instrumentation/MemorySanitizer/msan_kernel_basic.ll

17

Hm, I don't see any difference between files built with -f[no-]discard-value-names using Clang built with/without assertions.
I've replaced the value names with seven getelementptr instructions.

[glider]

Re: value names, there's code in Clang that explicitly overrides discarding them for ASan/MSan builds.

[eugenis]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

1520

Hmm, now it is clear that this chunk does not do anything at all.
Revert it (but keep the TODO)?

3294

Is that because you can not use frame pointers in the kernel for some reason?

Looking at the code for __msan_poison_alloca in kmsan, with the loops and the hash tables and origin setting (with even more loops), I find it hard to believe that taking one or even several steps through FP chain would slow it down.

Note that when a function has more than one local variable, you are doing some redundant work.

Having said that, I don't really mind this change.

It would be nice to change the function name though, because it has very different semantics from the userspace function.

3445

Remove this.

3497

It's not easier, but it is correct.

3580

It won't break old instrumented code running with new runtime library, right?

AFAIK that's all that we care about.

test/Instrumentation/MemorySanitizer/msan_kernel_basic.ll

17

Right. As you mention in the other changelist, we preserve value names in asan/msan builds.

Anyway, I like this new test more.

[glider]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

3294

Is that because you can not use frame pointers in the kernel for some reason?

There are configurations that disable frame pointers and use a different unwinding scheme.

Looking at the code for __msan_poison_alloca in kmsan, with the loops and the hash tables and origin setting (with even more loops), I find it hard to believe that taking one or even several steps through FP chain would slow it down.

Fair point.

Note that when a function has more than one local variable, you are doing some redundant work.

It might be a good idea to cache the unwound stack somehow, e.g. call the unwinder at the function prologue.

Having said that, I don't really mind this change.

Looks like it really doesn't make any difference whether we pass the pc or unwind it.
I'd better drop it then.

It would be nice to change the function name though, because it has very different semantics from the userspace function.

It's already different :)
There's no __msan_poison_alloca in the userspace, we're defining it in createKernelApi()

3497

Ok, I'll send a separate patch along the lines of:

AMD64FpEndOffset = AMD64FpEndOffsetSSE;
for (const auto &Attr : F.getAttributes().getFnAttributes()) {
  if (Attr.isStringAttribute() && (Attr.getKindAsString() == "target-features")) {
    if (Attr.getValueAsString().contains("-sse"))
      AMD64FpEndOffset = AMD64FpEndOffsetNoSSE;
    break;
  }
}

I'm not sure though that we're guaranteed to have those attributes (e.g. we currently don't have them in the tests, so we default to 176 bytes)

3580

Guess it shouldn't.
Ok, will do.

[eugenis]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

3294

It's already different :)
There's no __msan_poison_alloca in the userspace, we're defining it in createKernelApi()

Indeed.

3497

Then we should do whatever the target does when it does not find the attribute.

[glider]

Picked r339414, updated the test, PTAL

[eugenis]

Looks like all that's left is to split va_arg origin copy into a separate change and enable it in userspace.

[glider]

Incorporate https://reviews.llvm.org/D51412 (va_arg origin support in userspace)

[glider]

Now for realz (pulled the missing bit of va_arg support patch)

[eugenis]

LGTM, but you probably want to rebase it on top of D51412, instead of including that change in this one.

[glider]

Sure, will do.

[glider]

test/Instrumentation/MemorySanitizer/store-origin.ll

64

TODO: replace "01-9a-z" with "_0-9a-z"

[glider]

Rebased the CL on top of recent vararg changes

[eugenis]

LGTM

[glider]

Landed r341637, hope it sticks!

include/llvm/Transforms/Instrumentation.h

FunctionPass *createAddressSanitizerFunctionPass(bool CompileKernel = false,
                                                  bool Recover = false,
                                                  bool UseAfterScope = false);
 ModulePass *createAddressSanitizerModulePass(bool CompileKernel = false,
                                              bool Recover = false,
                                              bool UseGlobalsGC = true);
 
 // Insert MemorySanitizer instrumentation (detection of uninitialized reads)
 FunctionPass *createMemorySanitizerPass(int TrackOrigins = 0,
-                                        bool Recover = false);
+                                        bool Recover = false,
+                                        bool EnableKmsan = false);
 
 FunctionPass *createHWAddressSanitizerPass(bool CompileKernel = false,
                                            bool Recover = false);
 
 // Insert ThreadSanitizer (race detection) instrumentation
 FunctionPass *createThreadSanitizerPass();
 
 // Insert DataFlowSanitizer (dynamic data flow analysis) instrumentation

lib/Transforms/Instrumentation/MemorySanitizer.cpp

 ///
 /// This does not work very well with Compare-And-Swap (CAS) and
 /// Read-Modify-Write (RMW) operations. To follow the above logic, CAS and RMW
 /// must store the new shadow before the app operation, and load the shadow
 /// after the app operation. Computers don't work this way. Current
 /// implementation ignores the load aspect of CAS/RMW, always returning a clean
 /// value. It implements the store part as a simple atomic store by storing a
 /// clean shadow.
-//
+///
+///                  KernelMemorySanitizer (KMSAN) implementation.
+///
+/// The major differences between KMSAN and MSan instrumentation are:
+///  - KMSAN always tracks the origins and implies msan-keep-going=true;
+///  - KMSAN allocates shadow and origin memory for each page separately, so
+///    there are no explicit accesses to shadow and origin in the
+///    instrumentation.
+///    Shadow and origin values for a particular X-byte memory location
+///    (X=1,2,4,8) are accessed through pointers obtained via the
+///      __msan_metadata_ptr_for_load_X(ptr)
+///      __msan_metadata_ptr_for_store_X(ptr)
+///    functions. The corresponding functions check that the X-byte accesses
+///    are possible and returns the pointers to shadow and origin memory.
+///    Arbitrary sized accesses are handled with:
+///      __msan_metadata_ptr_for_load_n(ptr, size)
+///      __msan_metadata_ptr_for_store_n(ptr, size);
+///  - TLS variables are stored in a single per-task struct. A call to a
+///    function __msan_get_context_state() returning a pointer to that struct
+///    is inserted into every instrumented function before the entry block;
+///  - __msan_warning() takes a 32-bit origin parameter;
+///  - local variables are poisoned with __msan_poison_alloca() upon function
+///    entry and unpoisoned with __msan_unpoison_alloca() before leaving the
+///    function;
+///  - the pass doesn't declare any global variables or add global constructors
+///    to the translation unit.
+///
+/// Also, KMSAN currently ignores uninitialized memory passed into inline asm
+/// calls, making sure we're on the safe side wrt. possible false positives.
+///
+///  KernelMemorySanitizer only supports X86_64 at the moment.
+///
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringExtras.h"
static cl::opt<bool> ClHandleICmp("msan-handle-icmp",
        cl::desc("propagate shadow through ICmpEQ and ICmpNE"),
        cl::Hidden, cl::init(true));
 
 static cl::opt<bool> ClHandleICmpExact("msan-handle-icmp-exact",
        cl::desc("exact handling of relational integer ICmp"),
        cl::Hidden, cl::init(false));
 
 // When compiling the Linux kernel, we sometimes see false positives related to
 // MSan being unable to understand that inline assembly calls may initialize
 // local variables.
 // This flag makes the compiler conservatively unpoison every memory location
 // passed into an assembly call. Note that this may cause false positives.
 // Because it's impossible to figure out the array sizes, we can only unpoison
 // the first sizeof(type) bytes for each type* pointer.
 static cl::opt<bool> ClHandleAsmConservative(
     "msan-handle-asm-conservative",
     cl::desc("conservative handling of inline assembly"), cl::Hidden,

 static cl::opt<int> ClInstrumentationWithCallThreshold(
     "msan-instrumentation-with-call-threshold",
     cl::desc(
         "If the function being instrumented requires more than "
         "this number of checks and origin stores, use callbacks instead of "
         "inline checks (-1 means never use callbacks)."),
     cl::Hidden, cl::init(3500));
 
+static cl::opt<bool>
+    ClEnableKmsan("msan-kernel",
+                  cl::desc("Enable KernelMemorySanitizer instrumentation"),
+                  cl::Hidden, cl::init(false));
+
 // This is an experiment to enable handling of cases where shadow is a non-zero
 // compile-time constant. For some unexplainable reason they were silently
 // ignored in the instrumentation.
 static cl::opt<bool> ClCheckConstantShadow("msan-check-constant-shadow",
        cl::desc("Insert checks for constant shadow values"),
        cl::Hidden, cl::init(false));
 
 // This is off by default because of a bug in gold:

 ///
 /// MemorySanitizer: instrument the code in module to find
 /// uninitialized reads.
 class MemorySanitizer : public FunctionPass {
 public:
   // Pass identification, replacement for typeid.
   static char ID;
 
-  MemorySanitizer(int TrackOrigins = 0, bool Recover = false)
-      : FunctionPass(ID),
-        TrackOrigins(std::max(TrackOrigins, (int)ClTrackOrigins)),
-        Recover(Recover || ClKeepGoing) {}
-
+  MemorySanitizer(int TrackOrigins = 0, bool Recover = false,
+                  bool EnableKmsan = false)
+      : FunctionPass(ID) {
+    this->CompileKernel =
+        ClEnableKmsan.getNumOccurrences() > 0 ? ClEnableKmsan : EnableKmsan;
+    if (ClTrackOrigins.getNumOccurrences() > 0)
+      this->TrackOrigins = ClTrackOrigins;
+    else
+      this->TrackOrigins = this->CompileKernel ? 2 : TrackOrigins;
+    this->Recover = ClKeepGoing.getNumOccurrences() > 0
+                        ? ClKeepGoing
+                        : (this->CompileKernel | Recover);
+  }
   StringRef getPassName() const override { return "MemorySanitizer"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<TargetLibraryInfoWrapperPass>();
   }
 
   bool runOnFunction(Function &F) override;
   bool doInitialization(Module &M) override;
 
 private:
   friend struct MemorySanitizerVisitor;
   friend struct VarArgAMD64Helper;
   friend struct VarArgMIPS64Helper;
   friend struct VarArgAArch64Helper;
   friend struct VarArgPowerPC64Helper;
 
   void initializeCallbacks(Module &M);
+  void createKernelApi(Module &M);
   void createUserspaceApi(Module &M);
 
+  /// True if we're compiling the Linux kernel.
+  bool CompileKernel;
+
   /// Track origins (allocation points) of uninitialized values.
   int TrackOrigins;
   bool Recover;
 
   LLVMContext *C;
   Type *IntptrTy;
   Type *OriginTy;
 
+  // XxxTLS variables represent the per-thread state in MSan and per-task state
+  // in KMSAN.
+  // For the userspace these point to thread-local globals. In the kernel land
+  // they point to the members of a per-task struct obtained via a call to
+  // __msan_get_context_state().
+
   /// Thread-local shadow storage for function parameters.
-  GlobalVariable *ParamTLS;
+  Value *ParamTLS;
 
   /// Thread-local origin storage for function parameters.
-  GlobalVariable *ParamOriginTLS;
+  Value *ParamOriginTLS;
 
   /// Thread-local shadow storage for function return value.
-  GlobalVariable *RetvalTLS;
+  Value *RetvalTLS;
 
   /// Thread-local origin storage for function return value.
-  GlobalVariable *RetvalOriginTLS;
+  Value *RetvalOriginTLS;
 
   /// Thread-local shadow storage for in-register va_arg function
   /// parameters (x86_64-specific).
-  GlobalVariable *VAArgTLS;
+  Value *VAArgTLS;
+
+  /// Thread-local shadow storage for in-register va_arg function
+  /// parameters (x86_64-specific, KMSAN only).
+  Value *VAArgOriginTLS;
 
   /// Thread-local shadow storage for va_arg overflow area
   /// (x86_64-specific).
-  GlobalVariable *VAArgOverflowSizeTLS;
+  Value *VAArgOverflowSizeTLS;
 
   /// Thread-local space used to pass origin value to the UMR reporting
   /// function.
-  GlobalVariable *OriginTLS;
+  Value *OriginTLS;
 
   /// Are the instrumentation callbacks set up?
   bool CallbacksInitialized = false;
 
   /// The run-time callback to print a warning.
   Value *WarningFn;
 
   // These arrays are indexed by log2(AccessSize).
private:
 
   /// Run-time helper that records a store (or any event) of an
   /// uninitialized value and returns an updated origin id encoding this info.
   Value *MsanChainOriginFn;
 
   /// MSan runtime replacements for memmove, memcpy and memset.
   Value *MemmoveFn, *MemcpyFn, *MemsetFn;
 
+  /// KMSAN callback for task-local function argument shadow.
+  Value *MsanGetContextStateFn;
+
+  /// Functions for poisoning/unpoisoning local variables
+  Value *MsanPoisonAllocaFn, *MsanUnpoisonAllocaFn;
+
+  /// Each of the MsanMetadataPtrXxx functions returns a pair of shadow/origin
+  /// pointers.
+  Value *MsanMetadataPtrForLoadN, *MsanMetadataPtrForStoreN;
+  Value *MsanMetadataPtrForLoad_1_8[4];
+  Value *MsanMetadataPtrForStore_1_8[4];
+
+  /// Helper to choose between different MsanMetadataPtrXxx().
+  Value *getKmsanShadowOriginAccessFn(bool isStore, int size);
+
   /// Memory map parameters used in application-to-shadow calculation.
   const MemoryMapParams *MapParams;
 
   /// Custom memory map parameters used when -msan-shadow-base or
   // -msan-origin-base is provided.
   MemoryMapParams CustomMapParams;
 
   MDNode *ColdCallWeights;

 INITIALIZE_PASS_BEGIN(
     MemorySanitizer, "msan",
     "MemorySanitizer: detects uninitialized reads.", false, false)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_END(
     MemorySanitizer, "msan",
     "MemorySanitizer: detects uninitialized reads.", false, false)
 
-FunctionPass *llvm::createMemorySanitizerPass(int TrackOrigins, bool Recover) {
-  return new MemorySanitizer(TrackOrigins, Recover);
+FunctionPass *llvm::createMemorySanitizerPass(int TrackOrigins, bool Recover,
+                                              bool CompileKernel) {
+  return new MemorySanitizer(TrackOrigins, Recover, CompileKernel);
 }
 
 /// Create a non-const global initialized with the given string.
 ///
 /// Creates a writable global for Str so that we can pass it to the
 /// run-time lib. Runtime uses first 4 bytes of the string to store the
 /// frame ID, so the string needs to be mutable.
 static GlobalVariable *createPrivateNonConstGlobalForString(Module &M,
                                                             StringRef Str) {
   Constant *StrConst = ConstantDataArray::getString(M.getContext(), Str);
   return new GlobalVariable(M, StrConst->getType(), /*isConstant=*/false,
                             GlobalValue::PrivateLinkage, StrConst, "");
 }
 
+/// Create KMSAN API callbacks.
+void MemorySanitizer::createKernelApi(Module &M) {
+  IRBuilder<> IRB(*C);
+
+  // These will be initialized in insertKmsanPrologue().
+  RetvalTLS = nullptr;
+  RetvalOriginTLS = nullptr;
+  ParamTLS = nullptr;
+  ParamOriginTLS = nullptr;
+  VAArgTLS = nullptr;
+  VAArgOriginTLS = nullptr;
+  VAArgOverflowSizeTLS = nullptr;
+  // OriginTLS is unused in the kernel.
+  OriginTLS = nullptr;
+
+  // __msan_warning() in the kernel takes an origin.
+  WarningFn = M.getOrInsertFunction("__msan_warning", IRB.getVoidTy(),
+                                    IRB.getInt32Ty());
+  // Requests the per-task context state (kmsan_context_state*) from the
+  // runtime library.
+  MsanGetContextStateFn = M.getOrInsertFunction(
+      "__msan_get_context_state",
+      PointerType::get(
+          StructType::get(ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8),
+                          ArrayType::get(IRB.getInt64Ty(), kRetvalTLSSize / 8),
+                          ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8),
+                          ArrayType::get(IRB.getInt64Ty(),
+                                         kParamTLSSize / 8), /* va_arg_origin */
+                          IRB.getInt64Ty(),
+                          ArrayType::get(OriginTy, kParamTLSSize / 4), OriginTy,
+                          OriginTy),
+          0));
+
+  Type *RetTy = StructType::get(PointerType::get(IRB.getInt8Ty(), 0),
+                                PointerType::get(IRB.getInt32Ty(), 0));
+
+  for (int ind = 0, size = 1; ind < 4; ind++, size <<= 1) {
+    std::string name_load =
+        "__msan_metadata_ptr_for_load_" + std::to_string(size);
+    std::string name_store =
+        "__msan_metadata_ptr_for_store_" + std::to_string(size);
+    MsanMetadataPtrForLoad_1_8[ind] = M.getOrInsertFunction(
+        name_load, RetTy, PointerType::get(IRB.getInt8Ty(), 0));
+    MsanMetadataPtrForStore_1_8[ind] = M.getOrInsertFunction(
+        name_store, RetTy, PointerType::get(IRB.getInt8Ty(), 0));
+  }
+
+  MsanMetadataPtrForLoadN = M.getOrInsertFunction(
+      "__msan_metadata_ptr_for_load_n", RetTy,
+      PointerType::get(IRB.getInt8Ty(), 0), IRB.getInt64Ty());
+  MsanMetadataPtrForStoreN = M.getOrInsertFunction(
+      "__msan_metadata_ptr_for_store_n", RetTy,
+      PointerType::get(IRB.getInt8Ty(), 0), IRB.getInt64Ty());
+
+  // Functions for poisoning and unpoisoning memory.
+  MsanPoisonAllocaFn =
+      M.getOrInsertFunction("__msan_poison_alloca", IRB.getVoidTy(),
+                            IRB.getInt8PtrTy(), IntptrTy, IRB.getInt8PtrTy());
+  MsanUnpoisonAllocaFn = M.getOrInsertFunction(
+      "__msan_unpoison_alloca", IRB.getVoidTy(), IRB.getInt8PtrTy(), IntptrTy);
+}
+
 /// Insert declarations for userspace-specific functions and globals.
 void MemorySanitizer::createUserspaceApi(Module &M) {
   IRBuilder<> IRB(*C);
   // Create the callback.
   // FIXME: this function should have "Cold" calling conv,
   // which is not yet implemented.
   StringRef WarningFnName = Recover ? "__msan_warning"
                                     : "__msan_warning_noreturn";
   WarningFn = M.getOrInsertFunction(WarningFnName, IRB.getVoidTy());
 
   // Create the global TLS variables.
   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(
       M, ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8), false,
       GlobalVariable::ExternalLinkage, nullptr, "__msan_va_arg_tls", nullptr,
void MemorySanitizer::initializeCallbacks(Module &M) {
   MemsetFn = M.getOrInsertFunction(
     "__msan_memset", IRB.getInt8PtrTy(), IRB.getInt8PtrTy(), IRB.getInt32Ty(),
     IntptrTy);
   // We insert an empty inline asm after __msan_report* to avoid callback merge.
   EmptyAsm = InlineAsm::get(FunctionType::get(IRB.getVoidTy(), false),
                             StringRef(""), StringRef(""),
                             /*hasSideEffects=*/true);
 
+  if (CompileKernel) {
+    createKernelApi(M);
+  } else {
     createUserspaceApi(M);
+  }
   CallbacksInitialized = true;
 }
 
+Value *MemorySanitizer::getKmsanShadowOriginAccessFn(bool isStore, int size) {
+  Value **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.
 ///
 /// inserts a call to __msan_init to the module's constructor list.
 bool MemorySanitizer::doInitialization(Module &M) {
   auto &DL = M.getDataLayout();
 
   bool ShadowPassed = ClShadowBase.getNumOccurrences() > 0;
   bool OriginPassed = ClOriginBase.getNumOccurrences() > 0;
bool MemorySanitizer::doInitialization(Module &M) {
   C = &(M.getContext());
   IRBuilder<> IRB(*C);
   IntptrTy = IRB.getIntPtrTy(DL);
   OriginTy = IRB.getInt32Ty();
 
   ColdCallWeights = MDBuilder(*C).createBranchWeights(1, 1000);
   OriginStoreWeights = MDBuilder(*C).createBranchWeights(1, 1000);
 
+  if (!CompileKernel) {
     std::tie(MsanCtorFunction, std::ignore) =
-      createSanitizerCtorAndInitFunctions(M, kMsanModuleCtorName, kMsanInitName,
+        createSanitizerCtorAndInitFunctions(M, kMsanModuleCtorName,
+                                            kMsanInitName,
                                             /*InitArgTypes=*/{},
                                             /*InitArgs=*/{});
     if (ClWithComdat) {
       Comdat *MsanCtorComdat = M.getOrInsertComdat(kMsanModuleCtorName);
       MsanCtorFunction->setComdat(MsanCtorComdat);
       appendToGlobalCtors(M, MsanCtorFunction, 0, MsanCtorFunction);
     } else {
       appendToGlobalCtors(M, MsanCtorFunction, 0);
     }
 
-
     if (TrackOrigins)
       new GlobalVariable(M, IRB.getInt32Ty(), true, GlobalValue::WeakODRLinkage,
                          IRB.getInt32(TrackOrigins), "__msan_track_origins");
 
     if (Recover)
       new GlobalVariable(M, IRB.getInt32Ty(), true, GlobalValue::WeakODRLinkage,
                          IRB.getInt32(Recover), "__msan_keep_going");
-
+  }
   return true;
 }
 
 namespace {
 
 /// A helper class that handles instrumentation of VarArg
 /// functions on a particular platform.
 ///
MemorySanitizerVisitor(Function &F, MemorySanitizer &MS)
     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");
     TLI = &MS.getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
 
     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 '"
                << F.getName() << "'\n");
   }
 
   Value *updateOrigin(Value *V, IRBuilder<> &IRB) {
     if (MS.TrackOrigins <= 1) return V;
if (Shadow->getType()->isAggregateType()) {
           paintOrigin(IRB, updateOrigin(Origin, IRB), OriginPtr, StoreSize,
                       OriginAlignment);
         return;
       }
 
       unsigned TypeSizeInBits =
           DL.getTypeSizeInBits(ConvertedShadow->getType());
       unsigned SizeIndex = TypeSizeToSizeIndex(TypeSizeInBits);
-      if (AsCall && SizeIndex < kNumberOfAccessSizes) {
+      if (AsCall && SizeIndex < kNumberOfAccessSizes && !MS.CompileKernel) {
         Value *Fn = MS.MaybeStoreOriginFn[SizeIndex];
         Value *ConvertedShadow2 = IRB.CreateZExt(
             ConvertedShadow, IRB.getIntNTy(8 * (1 << SizeIndex)));
         IRB.CreateCall(Fn, {ConvertedShadow2,
                             IRB.CreatePointerCast(Addr, IRB.getInt8PtrTy()),
                             Origin});
       } else {
         Value *Cmp = IRB.CreateICmpNE(
for (StoreInst *SI : StoreList) {
       Value *ShadowPtr, *OriginPtr;
       Type *ShadowTy = Shadow->getType();
       unsigned Alignment = SI->getAlignment();
       unsigned OriginAlignment = std::max(kMinOriginAlignment, Alignment);
       std::tie(ShadowPtr, OriginPtr) =
           getShadowOriginPtr(Addr, IRB, ShadowTy, Alignment, /*isStore*/ true);
 
       StoreInst *NewSI = IRB.CreateAlignedStore(Shadow, ShadowPtr, Alignment);
       LLVM_DEBUG(dbgs() << "  STORE: " << *NewSI << "\n");
       (void)NewSI;
 
       if (SI->isAtomic())
         SI->setOrdering(addReleaseOrdering(SI->getOrdering()));
 
       if (MS.TrackOrigins && !SI->isAtomic())
         storeOrigin(IRB, Addr, Shadow, getOrigin(Val), OriginPtr,
                     OriginAlignment, InstrumentWithCalls);
     }
   }
 
   /// Helper function to insert a warning at IRB's current insert point.
   void insertWarningFn(IRBuilder<> &IRB, Value *Origin) {
     if (!Origin)
       Origin = (Value *)IRB.getInt32(0);
+    if (MS.CompileKernel) {
+      IRB.CreateCall(MS.WarningFn, Origin);
+    } else {
       if (MS.TrackOrigins) {
         IRB.CreateStore(Origin, MS.OriginTLS);
       }
       IRB.CreateCall(MS.WarningFn, {});
+    }
     IRB.CreateCall(MS.EmptyAsm, {});
     // 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) {
if (ConstantShadow) {
       }
       return;
     }
 
     const DataLayout &DL = OrigIns->getModule()->getDataLayout();
 
     unsigned TypeSizeInBits = DL.getTypeSizeInBits(ConvertedShadow->getType());
     unsigned SizeIndex = TypeSizeToSizeIndex(TypeSizeInBits);
-    if (AsCall && SizeIndex < kNumberOfAccessSizes) {
+    if (AsCall && SizeIndex < kNumberOfAccessSizes && !MS.CompileKernel) {
       Value *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)});
     } else {
       Value *Cmp = IRB.CreateICmpNE(ConvertedShadow,
                                     getCleanShadow(ConvertedShadow), "_mscmp");
       Instruction *CheckTerm = SplitBlockAndInsertIfThen(
           Cmp, OrigIns,
           /* Unreachable */ !MS.Recover, MS.ColdCallWeights);
 
       IRB.SetInsertPoint(CheckTerm);
       insertWarningFn(IRB, Origin);
       LLVM_DEBUG(dbgs() << "  CHECK: " << *Cmp << "\n");
     }
   }
 
   void materializeChecks(bool InstrumentWithCalls) {
     for (const auto &ShadowData : InstrumentationList) {
       Instruction *OrigIns = ShadowData.OrigIns;
       Value *Shadow = ShadowData.Shadow;
       Value *Origin = ShadowData.Origin;
       materializeOneCheck(OrigIns, Shadow, Origin, InstrumentWithCalls);
     }
     LLVM_DEBUG(dbgs() << "DONE:\n" << F);
   }
 
+  BasicBlock *insertKmsanPrologue(Function &F) {
+    BasicBlock *ret =
+        SplitBlock(&F.getEntryBlock(), F.getEntryBlock().getFirstNonPHI());
+    IRBuilder<> IRB(F.getEntryBlock().getFirstNonPHI());
+    Value *ContextState = IRB.CreateCall(MS.MsanGetContextStateFn, {});
+    Constant *Zero = IRB.getInt32(0);
+    MS.ParamTLS =
+        IRB.CreateGEP(ContextState, {Zero, IRB.getInt32(0)}, "param_shadow");
+    MS.RetvalTLS =
+        IRB.CreateGEP(ContextState, {Zero, IRB.getInt32(1)}, "retval_shadow");
+    MS.VAArgTLS =
+        IRB.CreateGEP(ContextState, {Zero, IRB.getInt32(2)}, "va_arg_shadow");
+    MS.VAArgOriginTLS =
+        IRB.CreateGEP(ContextState, {Zero, IRB.getInt32(3)}, "va_arg_origin");
+    MS.VAArgOverflowSizeTLS = IRB.CreateGEP(
+        ContextState, {Zero, IRB.getInt32(4)}, "va_arg_overflow_size");
+    MS.ParamOriginTLS =
+        IRB.CreateGEP(ContextState, {Zero, IRB.getInt32(5)}, "param_origin");
+    MS.RetvalOriginTLS =
+        IRB.CreateGEP(ContextState, {Zero, IRB.getInt32(6)}, "retval_origin");
+    return ret;
+  }
+
   /// Add MemorySanitizer instrumentation to a function.
   bool runOnFunction() {
     // In the presence of unreachable blocks, we may see Phi nodes with
     // incoming nodes from such blocks. Since InstVisitor skips unreachable
     // blocks, such nodes will not have any shadow value associated with them.
     // It's easier to remove unreachable blocks than deal with missing shadow.
     removeUnreachableBlocks(F);
 
if (MS.TrackOrigins) {
             IRB.CreateAnd(OriginLong, ConstantInt::get(MS.IntptrTy, ~Mask));
       }
       OriginPtr =
           IRB.CreateIntToPtr(OriginLong, PointerType::get(IRB.getInt32Ty(), 0));
     }
     return std::make_pair(ShadowPtr, OriginPtr);
   }
 
+  std::pair<Value *, Value *>
+  getShadowOriginPtrKernel(Value *Addr, IRBuilder<> &IRB, Type *ShadowTy,
+                           unsigned Alignment, bool isStore) {
+    Value *ShadowOriginPtrs;
+    const DataLayout &DL = F.getParent()->getDataLayout();
+    int Size = DL.getTypeStoreSize(ShadowTy);
+
+    Value *Getter = MS.getKmsanShadowOriginAccessFn(isStore, Size);
+    Value *AddrCast =
+        IRB.CreatePointerCast(Addr, PointerType::get(IRB.getInt8Ty(), 0));
+    if (Getter) {
+      ShadowOriginPtrs = IRB.CreateCall(Getter, AddrCast);
+    } else {
+      Value *SizeVal = ConstantInt::get(MS.IntptrTy, Size);
+      ShadowOriginPtrs = IRB.CreateCall(isStore ? MS.MsanMetadataPtrForStoreN
+                                                : MS.MsanMetadataPtrForLoadN,
+                                        {AddrCast, SizeVal});
+    }
+    Value *ShadowPtr = IRB.CreateExtractValue(ShadowOriginPtrs, 0);
+    ShadowPtr = IRB.CreatePointerCast(ShadowPtr, PointerType::get(ShadowTy, 0));
+    Value *OriginPtr = IRB.CreateExtractValue(ShadowOriginPtrs, 1);
+
+    return std::make_pair(ShadowPtr, OriginPtr);
+  }
+
   std::pair<Value *, Value *> getShadowOriginPtr(Value *Addr, IRBuilder<> &IRB,
                                                  Type *ShadowTy,
                                                  unsigned Alignment,
                                                  bool isStore) {
-    std::pair<Value *, Value *> ret =
-        getShadowOriginPtrUserspace(Addr, IRB, ShadowTy, Alignment);
+    std::pair<Value *, Value *> ret;
+    if (MS.CompileKernel)
+      ret = getShadowOriginPtrKernel(Addr, IRB, ShadowTy, Alignment, isStore);
+    else
+      ret = getShadowOriginPtrUserspace(Addr, IRB, ShadowTy, Alignment);
     return ret;
   }
 
   /// Compute the shadow address for a given function argument.
   ///
   /// Shadow = ParamTLS+ArgOffset.
   Value *getShadowPtrForArgument(Value *A, IRBuilder<> &IRB,
                                  int ArgOffset) {
     Value *Base = IRB.CreatePointerCast(MS.ParamTLS, MS.IntptrTy);
     if (ArgOffset)
       Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
     return IRB.CreateIntToPtr(Base, PointerType::get(getShadowTy(A), 0),
                               "_msarg");
   }
 
   /// Compute the origin address for a given function argument.
   Value *getOriginPtrForArgument(Value *A, IRBuilder<> &IRB,
                                  int ArgOffset) {
-    if (!MS.TrackOrigins) return nullptr;
+    if (!MS.TrackOrigins)
+      return nullptr;
     Value *Base = IRB.CreatePointerCast(MS.ParamOriginTLS, MS.IntptrTy);
     if (ArgOffset)
       Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
     return IRB.CreateIntToPtr(Base, PointerType::get(MS.OriginTy, 0),
                               "_msarg_o");
   }
 
   /// Compute the shadow address for a retval.
if (Argument *A = dyn_cast<Argument>(V)) {
             // argument shadow to the underlying memory.
             // Figure out maximal valid memcpy alignment.
             unsigned ArgAlign = FArg.getParamAlignment();
             if (ArgAlign == 0) {
               Type *EltType = A->getType()->getPointerElementType();
               ArgAlign = DL.getABITypeAlignment(EltType);
             }
             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 {
               unsigned CopyAlign = std::min(ArgAlign, kShadowTLSAlignment);
               Value *Cpy = EntryIRB.CreateMemCpy(CpShadowPtr, CopyAlign, Base,
if (MS.TrackOrigins) {
       if (PropagateShadow) {
         unsigned OriginAlignment = std::max(kMinOriginAlignment, Alignment);
         setOrigin(&I, IRB.CreateAlignedLoad(OriginPtr, OriginAlignment));
       } else {
         setOrigin(&I, getCleanOrigin());
       }
     }
   }
 
   /// Instrument StoreInst
   ///
   /// Stores the corresponding shadow and (optionally) origin.
   /// Optionally, checks that the store address is fully defined.
   void visitStoreInst(StoreInst &I) {
     StoreList.push_back(&I);
     if (ClCheckAccessAddress)
       insertShadowCheck(I.getPointerOperand(), &I);
if (CS.isCall()) {
       assert(!isa<IntrinsicInst>(&I) && "intrinsics are handled elsewhere");
 
       // We are going to insert code that relies on the fact that the callee
       // will become a non-readonly function after it is instrumented by us. To
       // prevent this code from being optimized out, mark that function
       // non-readonly in advance.
       if (Function *Func = Call->getCalledFunction()) {
         // Clear out readonly/readnone attributes.
         AttrBuilder B;
         B.addAttribute(Attribute::ReadOnly)
           .addAttribute(Attribute::ReadNone);
         Func->removeAttributes(AttributeList::FunctionIndex, B);
       }
 
       maybeMarkSanitizerLibraryCallNoBuiltin(Call, TLI);
     }
     IRBuilder<> IRB(&I);
for (CallSite::arg_iterator ArgIt = CS.arg_begin(), End = CS.arg_end();
       const DataLayout &DL = F.getParent()->getDataLayout();
       if (CS.paramHasAttr(i, Attribute::ByVal)) {
         assert(A->getType()->isPointerTy() &&
                "ByVal argument is not a pointer!");
         Size = DL.getTypeAllocSize(A->getType()->getPointerElementType());
         if (ArgOffset + Size > kParamTLSSize) break;
         unsigned ParamAlignment = CS.getParamAlignment(i);
         unsigned Alignment = std::min(ParamAlignment, kShadowTLSAlignment);
-        Value *AShadowPtr = getShadowOriginPtr(A, IRB, IRB.getInt8Ty(),
-                                               Alignment, /*isStore*/ false)
-                                .first;
+        Value *AShadowPtr =
+            getShadowOriginPtr(A, IRB, IRB.getInt8Ty(), Alignment,
+                               /*isStore*/ false)
 
         Store = IRB.CreateMemCpy(ArgShadowBase, Alignment, AShadowPtr,
                                  Alignment, Size);
+        // TODO(glider): need to copy origins.
       } else {
         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;
       }
void visitPHINode(PHINode &I) {
     ShadowPHINodes.push_back(&I);
     setShadow(&I, IRB.CreatePHI(getShadowTy(&I), I.getNumIncomingValues(),
                                 "_msphi_s"));
     if (MS.TrackOrigins)
       setOrigin(&I, IRB.CreatePHI(MS.OriginTy, I.getNumIncomingValues(),
                                   "_msphi_o"));
   }
 
-  void visitAllocaInst(AllocaInst &I) {
-    setShadow(&I, getCleanShadow(&I));
-    setOrigin(&I, getCleanOrigin());
-    IRBuilder<> IRB(I.getNextNode());
-    const DataLayout &DL = F.getParent()->getDataLayout();
-    uint64_t TypeSize = DL.getTypeAllocSize(I.getAllocatedType());
-    Value *Len = ConstantInt::get(MS.IntptrTy, TypeSize);
-    if (I.isArrayAllocation())
-      Len = IRB.CreateMul(Len, I.getArraySize());
-    if (PoisonStack && ClPoisonStackWithCall) {
-      IRB.CreateCall(MS.MsanPoisonStackFn,
-                     {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len});
-    } else {
-      Value *ShadowBase = getShadowOriginPtr(&I, IRB, IRB.getInt8Ty(),
-                                             I.getAlignment(), /*isStore*/ true)
-                              .first;
-
-      Value *PoisonValue = IRB.getInt8(PoisonStack ? ClPoisonStackPattern : 0);
-      IRB.CreateMemSet(ShadowBase, PoisonValue, Len, I.getAlignment());
-    }
-
-    if (PoisonStack && MS.TrackOrigins) {
-      SmallString<2048> StackDescriptionStorage;
-      raw_svector_ostream StackDescription(StackDescriptionStorage);
-      // We create a string with a description of the stack allocation and
-      // pass it into __msan_set_alloca_origin.
-      // It will be printed by the run-time if stack-originated UMR is found.
-      // The first 4 bytes of the string are set to '----' and will be replaced
-      // by __msan_va_arg_overflow_size_tls at the first call.
-      StackDescription << "----" << I.getName() << "@" << F.getName();
-      Value *Descr =
-          createPrivateNonConstGlobalForString(*F.getParent(),
+  Value *getLocalVarDescription(AllocaInst &I) {
+    return createPrivateNonConstGlobalForString(*F.getParent(),
                                                 StackDescription.str());
+  }
+
+  void instrumentAllocaUserspace(AllocaInst &I, IRBuilder<> &IRB, Value *Len) {
+    if (PoisonStack && ClPoisonStackWithCall) {
+      IRB.CreateCall(MS.MsanPoisonStackFn,
+                     {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len});
+    } else {
+      Value *ShadowBase, *OriginBase;
+      std::tie(ShadowBase, OriginBase) =
+          getShadowOriginPtr(&I, IRB, IRB.getInt8Ty(), 1, /*isStore*/ true);
 
+      Value *PoisonValue = IRB.getInt8(PoisonStack ? ClPoisonStackPattern : 0);
+      IRB.CreateMemSet(ShadowBase, PoisonValue, Len, I.getAlignment());
+    }
+
+    if (PoisonStack && MS.TrackOrigins) {
+      Value *Descr = getLocalVarDescription(I);
       IRB.CreateCall(MS.MsanSetAllocaOrigin4Fn,
                      {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len,
                       IRB.CreatePointerCast(Descr, IRB.getInt8PtrTy()),
                       IRB.CreatePointerCast(&F, MS.IntptrTy)});
     }
   }
 
+  void instrumentAllocaKmsan(AllocaInst &I, IRBuilder<> &IRB, Value *Len) {
+    Value *Descr = getLocalVarDescription(I);
+    if (PoisonStack) {
+      IRB.CreateCall(MS.MsanPoisonAllocaFn,
+                     {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len,
+                      IRB.CreatePointerCast(Descr, IRB.getInt8PtrTy())});
+    } else {
+      IRB.CreateCall(MS.MsanUnpoisonAllocaFn,
+                     {IRB.CreatePointerCast(&I, IRB.getInt8PtrTy()), Len});
+    }
+  }
+
+  void visitAllocaInst(AllocaInst &I) {
+    setShadow(&I, getCleanShadow(&I));
+    setOrigin(&I, getCleanOrigin());
+    IRBuilder<> IRB(I.getNextNode());
+    const DataLayout &DL = F.getParent()->getDataLayout();
+    uint64_t TypeSize = DL.getTypeAllocSize(I.getAllocatedType());
+    Value *Len = ConstantInt::get(MS.IntptrTy, TypeSize);
+    if (I.isArrayAllocation())
+      Len = IRB.CreateMul(Len, I.getArraySize());
+
+    if (MS.CompileKernel)
+      instrumentAllocaKmsan(I, IRB, Len);
+    else
+      instrumentAllocaUserspace(I, IRB, Len);
+  }
+
   void visitSelectInst(SelectInst& I) {
     IRBuilder<> IRB(&I);
     // a = select b, c, d
     Value *B = I.getCondition();
     Value *C = I.getTrueValue();
     Value *D = I.getFalseValue();
     Value *Sb = getShadow(B);
     Value *Sc = getShadow(C);
void visitCatchReturnInst(CatchReturnInst &CRI) {
     LLVM_DEBUG(dbgs() << "CatchReturn: " << CRI << "\n");
     // Nothing to do here.
   }
 
   void visitAsmInstruction(Instruction &I) {
     // Conservative inline assembly handling: check for poisoned shadow of
     // asm() arguments, then unpoison the result and all the memory locations
     // pointed to by those arguments.
     CallInst *CI = dyn_cast<CallInst>(&I);
 
     for (size_t i = 0, n = CI->getNumOperands(); i < n; i++) {
       Value *Operand = CI->getOperand(i);
       if (Operand->getType()->isSized())
         insertShadowCheck(Operand, &I);
     }
     setShadow(&I, getCleanShadow(&I));
     setOrigin(&I, getCleanOrigin());

 /// AMD64-specific implementation of VarArgHelper.
 struct VarArgAMD64Helper : public VarArgHelper {
   // An unfortunate workaround for asymmetric lowering of va_arg stuff.
   // See a comment in visitCallSite for more details.
   static const unsigned AMD64GpEndOffset = 48;  // AMD64 ABI Draft 0.99.6 p3.5.7
   static const unsigned AMD64FpEndOffsetSSE = 176;
   // If SSE is disabled, fp_offset in va_list is zero.
   static const unsigned AMD64FpEndOffsetNoSSE = AMD64GpEndOffset;
 
   unsigned AMD64FpEndOffset;
   Function &F;
   MemorySanitizer &MS;
   MemorySanitizerVisitor &MSV;
   Value *VAArgTLSCopy = nullptr;
+  Value *VAArgTLSOriginCopy = nullptr;
   Value *VAArgOverflowSize = nullptr;
 
   SmallVector<CallInst*, 16> VAStartInstrumentationList;
 
   enum ArgKind { AK_GeneralPurpose, AK_FloatingPoint, AK_Memory };
 
   VarArgAMD64Helper(Function &F, MemorySanitizer &MS,
                     MemorySanitizerVisitor &MSV)
for (CallSite::arg_iterator ArgIt = CS.arg_begin(), End = CS.arg_end();
         // over by va_start, so don't count them towards the offset.
         if (IsFixed)
           continue;
         assert(A->getType()->isPointerTy());
         Type *RealTy = A->getType()->getPointerElementType();
         uint64_t ArgSize = DL.getTypeAllocSize(RealTy);
         Value *ShadowBase =
             getShadowPtrForVAArgument(RealTy, IRB, OverflowOffset);
+        Value *OriginBase = nullptr;
+        if (MS.TrackOrigins)
+          OriginBase = getOriginPtrForVAArgument(RealTy, IRB, OverflowOffset);
         OverflowOffset += alignTo(ArgSize, 8);
         Value *ShadowPtr, *OriginPtr;
         std::tie(ShadowPtr, OriginPtr) =
             MSV.getShadowOriginPtr(A, IRB, IRB.getInt8Ty(), kShadowTLSAlignment,
                                    /*isStore*/ false);
 
         IRB.CreateMemCpy(ShadowBase, kShadowTLSAlignment, ShadowPtr,
                          kShadowTLSAlignment, ArgSize);
+        if (MS.TrackOrigins)
+          IRB.CreateMemCpy(OriginBase, kShadowTLSAlignment, OriginPtr,
+                           kShadowTLSAlignment, ArgSize);
       } else {
         ArgKind AK = classifyArgument(A);
         if (AK == AK_GeneralPurpose && GpOffset >= AMD64GpEndOffset)
           AK = AK_Memory;
         if (AK == AK_FloatingPoint && FpOffset >= AMD64FpEndOffset)
           AK = AK_Memory;
-        Value *ShadowBase;
+        Value *ShadowBase, *OriginBase = nullptr;
         switch (AK) {
           case AK_GeneralPurpose:
             ShadowBase = getShadowPtrForVAArgument(A->getType(), IRB, GpOffset);
+            if (MS.TrackOrigins)
+              OriginBase =
+                  getOriginPtrForVAArgument(A->getType(), IRB, GpOffset);
             GpOffset += 8;
             break;
           case AK_FloatingPoint:
             ShadowBase = getShadowPtrForVAArgument(A->getType(), IRB, FpOffset);
+            if (MS.TrackOrigins)
+              OriginBase =
+                  getOriginPtrForVAArgument(A->getType(), IRB, FpOffset);
             FpOffset += 16;
             break;
           case AK_Memory:
             if (IsFixed)
               continue;
             uint64_t ArgSize = DL.getTypeAllocSize(A->getType());
             ShadowBase =
                 getShadowPtrForVAArgument(A->getType(), IRB, OverflowOffset);
+            if (MS.TrackOrigins)
+              OriginBase =
+                  getOriginPtrForVAArgument(A->getType(), IRB, OverflowOffset);
             OverflowOffset += alignTo(ArgSize, 8);
         }
         // Take fixed arguments into account for GpOffset and FpOffset,
         // but don't actually store shadows for them.
+        // TODO(glider): don't call get*PtrForVAArgument() for them.
         if (IsFixed)
           continue;
-        IRB.CreateAlignedStore(MSV.getShadow(A), ShadowBase,
-                               kShadowTLSAlignment);
+        Value *Shadow = MSV.getShadow(A);
+        IRB.CreateAlignedStore(Shadow, ShadowBase, kShadowTLSAlignment);
+        if (MS.TrackOrigins) {
+          unsigned StoreSize = DL.getTypeStoreSize(Shadow->getType());
+          MSV.paintOrigin(IRB, MSV.getOrigin(A), OriginBase, StoreSize,
+                          std::max(kShadowTLSAlignment, kMinOriginAlignment));
+        }
       }
     }
     Constant *OverflowSize =
       ConstantInt::get(IRB.getInt64Ty(), OverflowOffset - AMD64FpEndOffset);
     IRB.CreateStore(OverflowSize, MS.VAArgOverflowSizeTLS);
   }
 
   /// Compute the shadow address for a given va_arg.
   Value *getShadowPtrForVAArgument(Type *Ty, IRBuilder<> &IRB,
                                    int ArgOffset) {
     Value *Base = IRB.CreatePointerCast(MS.VAArgTLS, MS.IntptrTy);
     Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
     return IRB.CreateIntToPtr(Base, PointerType::get(MSV.getShadowTy(Ty), 0),
-                              "_msarg");
+                              "_msarg_va_s");
+  }
+
+  /// Compute the origin address for a given va_arg.
+  Value *getOriginPtrForVAArgument(Type *Ty, IRBuilder<> &IRB, int ArgOffset) {
+    Value *Base = IRB.CreatePointerCast(MS.VAArgOriginTLS, MS.IntptrTy);
+    Base = IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
+    return IRB.CreateIntToPtr(Base, PointerType::get(MSV.getShadowTy(Ty), 0),
+                              "_msarg_va_o");
   }
 
   void unpoisonVAListTagForInst(IntrinsicInst &I) {
     IRBuilder<> IRB(&I);
     Value *VAListTag = I.getArgOperand(0);
     Value *ShadowPtr, *OriginPtr;
     unsigned Alignment = 8;
     std::tie(ShadowPtr, OriginPtr) =
if (!VAStartInstrumentationList.empty()) {
       // va_arg_tls somewhere in the function entry block.
       IRBuilder<> IRB(MSV.ActualFnStart->getFirstNonPHI());
       VAArgOverflowSize = IRB.CreateLoad(MS.VAArgOverflowSizeTLS);
       Value *CopySize =
         IRB.CreateAdd(ConstantInt::get(MS.IntptrTy, AMD64FpEndOffset),
                       VAArgOverflowSize);
       VAArgTLSCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
       IRB.CreateMemCpy(VAArgTLSCopy, 8, MS.VAArgTLS, 8, CopySize);
+      if (MS.TrackOrigins) {
+        VAArgTLSOriginCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
+        IRB.CreateMemCpy(VAArgTLSOriginCopy, 8, MS.VAArgOriginTLS, 8, CopySize);
+      }
     }
 
     // Instrument va_start.
     // Copy va_list shadow from the backup copy of the TLS contents.
     for (size_t i = 0, n = VAStartInstrumentationList.size(); i < n; i++) {
       CallInst *OrigInst = VAStartInstrumentationList[i];
       IRBuilder<> IRB(OrigInst->getNextNode());
       Value *VAListTag = OrigInst->getArgOperand(0);
 
       Value *RegSaveAreaPtrPtr = IRB.CreateIntToPtr(
           IRB.CreateAdd(IRB.CreatePtrToInt(VAListTag, MS.IntptrTy),
                         ConstantInt::get(MS.IntptrTy, 16)),
           PointerType::get(Type::getInt64PtrTy(*MS.C), 0));
       Value *RegSaveAreaPtr = IRB.CreateLoad(RegSaveAreaPtrPtr);
       Value *RegSaveAreaShadowPtr, *RegSaveAreaOriginPtr;
       unsigned Alignment = 16;
       std::tie(RegSaveAreaShadowPtr, RegSaveAreaOriginPtr) =
           MSV.getShadowOriginPtr(RegSaveAreaPtr, IRB, IRB.getInt8Ty(),
                                  Alignment, /*isStore*/ true);
       IRB.CreateMemCpy(RegSaveAreaShadowPtr, Alignment, VAArgTLSCopy, Alignment,
                        AMD64FpEndOffset);
+      if (MS.TrackOrigins)
+        IRB.CreateMemCpy(RegSaveAreaOriginPtr, Alignment, VAArgTLSOriginCopy,
+                         Alignment, AMD64FpEndOffset);
       Value *OverflowArgAreaPtrPtr = IRB.CreateIntToPtr(
           IRB.CreateAdd(IRB.CreatePtrToInt(VAListTag, MS.IntptrTy),
                         ConstantInt::get(MS.IntptrTy, 8)),
           PointerType::get(Type::getInt64PtrTy(*MS.C), 0));
       Value *OverflowArgAreaPtr = IRB.CreateLoad(OverflowArgAreaPtrPtr);
       Value *OverflowArgAreaShadowPtr, *OverflowArgAreaOriginPtr;
       std::tie(OverflowArgAreaShadowPtr, OverflowArgAreaOriginPtr) =
           MSV.getShadowOriginPtr(OverflowArgAreaPtr, IRB, IRB.getInt8Ty(),
                                  Alignment, /*isStore*/ true);
       Value *SrcPtr = IRB.CreateConstGEP1_32(IRB.getInt8Ty(), VAArgTLSCopy,
                                              AMD64FpEndOffset);
       IRB.CreateMemCpy(OverflowArgAreaShadowPtr, Alignment, SrcPtr, Alignment,
                        VAArgOverflowSize);
+      if (MS.TrackOrigins) {
+        SrcPtr = IRB.CreateConstGEP1_32(IRB.getInt8Ty(), VAArgTLSOriginCopy,
+                                        AMD64FpEndOffset);
+        IRB.CreateMemCpy(OverflowArgAreaOriginPtr, Alignment, SrcPtr, Alignment,
+                         VAArgOverflowSize);
+      }
     }
   }
 };
 
 /// MIPS64-specific implementation of VarArgHelper.
 struct VarArgMIPS64Helper : public VarArgHelper {
   Function &F;
   MemorySanitizer &MS;

test/Instrumentation/MemorySanitizer/alloca.ll

 ; RUN: opt < %s -msan -msan-check-access-address=0 -S | FileCheck %s --check-prefixes=CHECK,INLINE
 ; RUN: opt < %s -msan -msan-check-access-address=0 -msan-poison-stack-with-call=1 -S | FileCheck %s --check-prefixes=CHECK,CALL
 ; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=1 -S | FileCheck %s --check-prefixes=CHECK,ORIGIN
 ; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=2 -S | FileCheck %s --check-prefixes=CHECK,ORIGIN
+; RUN: opt < %s -msan -msan-kernel=1 -S | FileCheck %s --check-prefixes=CHECK,KMSAN
 
 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 void @static() sanitize_memory {
 entry:
   %x = alloca i32, align 4
   ret void
 }
 
 ; CHECK-LABEL: define void @static(
 ; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false)
 ; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
 ; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 4,
 ; CHECK: ret void
 
 
 define void @dynamic() sanitize_memory {
 entry:
   br label %l
 l:
   %x = alloca i32, align 4
   ret void
 }
 
 ; CHECK-LABEL: define void @dynamic(
 ; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 4, i1 false)
 ; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 4)
 ; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 4,
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 4,
 ; CHECK: ret void
 
 define void @array() sanitize_memory {
 entry:
   %x = alloca i32, i64 5, align 4
   ret void
 }
 
 ; CHECK-LABEL: define void @array(
 ; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 20, i1 false)
 ; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 20)
 ; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 20,
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 20,
 ; CHECK: ret void
 
 define void @array_non_const(i64 %cnt) sanitize_memory {
 entry:
   %x = alloca i32, i64 %cnt, align 4
   ret void
 }
 
 ; CHECK-LABEL: define void @array_non_const(
 ; CHECK: %[[A:.*]] = mul i64 4, %cnt
 ; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 -1, i64 %[[A]], i1 false)
 ; CALL: call void @__msan_poison_stack(i8* {{.*}}, i64 %[[A]])
 ; ORIGIN: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 %[[A]],
+; KMSAN: call void @__msan_poison_alloca(i8* {{.*}}, i64 %[[A]],
 ; CHECK: ret void
+
+; Check that the local is unpoisoned in the absence of sanitize_memory
+define void @unpoison_local() {
+entry:
+  %x = alloca i32, i64 5, align 4
+  ret void
+}
+
+; CHECK-LABEL: define void @unpoison_local(
+; INLINE: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 0, i64 20, i1 false)
+; CALL: call void @llvm.memset.p0i8.i64(i8* align 4 {{.*}}, i8 0, i64 20, i1 false)
+; ORIGIN-NOT: call void @__msan_set_alloca_origin4(i8* {{.*}}, i64 20,
+; KMSAN: call void @__msan_unpoison_alloca(i8* {{.*}}, i64 20)
+; CHECK: ret void
+

test/Instrumentation/MemorySanitizer/msan_kernel_basic.ll

+; KMSAN instrumentation tests
+; RUN: opt < %s -msan -msan-kernel=1 -S | FileCheck %s -check-prefixes=CHECK
+
+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 the instrumentation prologue.
+define void @Empty() nounwind uwtable sanitize_memory {
+entry:
+  ret void
+}
+
+; CHECK-LABEL: @Empty
+; CHECK: entry:
+; CHECK: @__msan_get_context_state()
+; %param_shadow:
+; CHECK: getelementptr {{.*}} i32 0, i32 0
+; %retval_shadow:
+; CHECK: getelementptr {{.*}} i32 0, i32 1
+; %va_arg_shadow:
+; CHECK: getelementptr {{.*}} i32 0, i32 2
+; %va_arg_origin:
+; CHECK: getelementptr {{.*}} i32 0, i32 3
+; %va_arg_overflow_size:
+; CHECK: getelementptr {{.*}} i32 0, i32 4
+; %param_origin:
+; CHECK: getelementptr {{.*}} i32 0, i32 5
+; %retval_origin:
+; CHECK: getelementptr {{.*}} i32 0, i32 6
+; CHECK: entry.split:
+
+; Check instrumentation of stores
+
+define void @Store1(i8* nocapture %p, i8 %x) nounwind uwtable sanitize_memory {
+entry:
+  store i8 %x, i8* %p
+  ret void
+}
+
+; CHECK-LABEL: @Store1
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: [[BASE2:%[0-9]+]] = ptrtoint {{.*}} [[PARAM_SHADOW]]
+; CHECK: [[BASE:%[0-9]+]] = ptrtoint {{.*}} [[PARAM_SHADOW]]
+; CHECK: [[SHADOW:%[a-z0-9_]+]] = inttoptr {{.*}} [[BASE]]
+; Load the shadow of %p and check it
+; CHECK: load i64, i64* [[SHADOW]]
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; CHECK: @__msan_metadata_ptr_for_store_1(i8* %p)
+; CHECK: store i8
+; If the new shadow is non-zero, jump to __msan_chain_origin()
+; CHECK: icmp
+; CHECK: br i1
+; CHECK: <label>
+; CHECK: @__msan_chain_origin
+; Storing origin here:
+; CHECK: store i32
+; CHECK: br label
+; CHECK: <label>
+; CHECK: store i8
+; CHECK: ret void
+
+define void @Store2(i16* nocapture %p, i16 %x) nounwind uwtable sanitize_memory {
+entry:
+  store i16 %x, i16* %p
+  ret void
+}
+
+; CHECK-LABEL: @Store2
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: ptrtoint {{.*}} [[PARAM_SHADOW]]
+; Load the shadow of %p and check it
+; CHECK: load i64
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; CHECK: [[REG:%[0-9]+]] = bitcast i16* %p to i8*
+; CHECK: @__msan_metadata_ptr_for_store_2(i8* [[REG]])
+; CHECK: store i16
+; If the new shadow is non-zero, jump to __msan_chain_origin()
+; CHECK: icmp
+; CHECK: br i1
+; CHECK: <label>
+; CHECK: @__msan_chain_origin
+; Storing origin here:
+; CHECK: store i32
+; CHECK: br label
+; CHECK: <label>
+; CHECK: store i16
+; CHECK: ret void
+
+
+define void @Store4(i32* nocapture %p, i32 %x) nounwind uwtable sanitize_memory {
+entry:
+  store i32 %x, i32* %p
+  ret void
+}
+
+; CHECK-LABEL: @Store4
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: ptrtoint {{.*}} [[PARAM_SHADOW]]
+; Load the shadow of %p and check it
+; CHECK: load i32
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; CHECK: [[REG:%[0-9]+]] = bitcast i32* %p to i8*
+; CHECK: @__msan_metadata_ptr_for_store_4(i8* [[REG]])
+; CHECK: store i32
+; If the new shadow is non-zero, jump to __msan_chain_origin()
+; CHECK: icmp
+; CHECK: br i1
+; CHECK: <label>
+; CHECK: @__msan_chain_origin
+; Storing origin here:
+; CHECK: store i32
+; CHECK: br label
+; CHECK: <label>
+; CHECK: store i32
+; CHECK: ret void
+
+define void @Store8(i64* nocapture %p, i64 %x) nounwind uwtable sanitize_memory {
+entry:
+  store i64 %x, i64* %p
+  ret void
+}
+
+; CHECK-LABEL: @Store8
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: ptrtoint {{.*}} [[PARAM_SHADOW]]
+; Load the shadow of %p and check it
+; CHECK: load i64
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; CHECK: [[REG:%[0-9]+]] = bitcast i64* %p to i8*
+; CHECK: @__msan_metadata_ptr_for_store_8(i8* [[REG]])
+; CHECK: store i64
+; If the new shadow is non-zero, jump to __msan_chain_origin()
+; CHECK: icmp
+; CHECK: br i1
+; CHECK: <label>
+; CHECK: @__msan_chain_origin
+; Storing origin here:
+; CHECK: store i32
+; CHECK: br label
+; CHECK: <label>
+; CHECK: store i64
+; CHECK: ret void
+
+define void @Store16(i128* nocapture %p, i128 %x) nounwind uwtable sanitize_memory {
+entry:
+  store i128 %x, i128* %p
+  ret void
+}
+
+; CHECK-LABEL: @Store16
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: ptrtoint {{.*}} [[PARAM_SHADOW]]
+; Load the shadow of %p and check it
+; CHECK: load i64
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; CHECK: [[REG:%[0-9]+]] = bitcast i128* %p to i8*
+; CHECK: @__msan_metadata_ptr_for_store_n(i8* [[REG]], i64 16)
+; CHECK: store i128
+; If the new shadow is non-zero, jump to __msan_chain_origin()
+; CHECK: icmp
+; CHECK: br i1
+; CHECK: <label>
+; CHECK: @__msan_chain_origin
+; Storing origin here:
+; CHECK: store i32
+; CHECK: br label
+; CHECK: <label>
+; CHECK: store i128
+; CHECK: ret void
+
+
+; Check instrumentation of loads
+
+define i8 @Load1(i8* nocapture %p) nounwind uwtable sanitize_memory {
+entry:
+  %0 = load i8, i8* %p
+  ret i8 %0
+}
+
+; CHECK-LABEL: @Load1
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: ptrtoint {{.*}} [[PARAM_SHADOW]]
+; Load the shadow of %p and check it
+; CHECK: load i64
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; Load the value from %p. This is done before accessing the shadow
+; to ease atomic handling.
+; CHECK: load i8
+; CHECK: @__msan_metadata_ptr_for_load_1(i8* %p)
+; Load the shadow and origin.
+; CHECK: load i8
+; CHECK: load i32
+
+
+define i16 @Load2(i16* nocapture %p) nounwind uwtable sanitize_memory {
+entry:
+  %0 = load i16, i16* %p
+  ret i16 %0
+}
+
+; CHECK-LABEL: @Load2
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: ptrtoint {{.*}} [[PARAM_SHADOW]]
+; Load the shadow of %p and check it
+; CHECK: load i64
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; Load the value from %p. This is done before accessing the shadow
+; to ease atomic handling.
+; CHECK: load i16
+; CHECK: [[REG:%[0-9]+]] = bitcast i16* %p to i8*
+; CHECK: @__msan_metadata_ptr_for_load_2(i8* [[REG]])
+; Load the shadow and origin.
+; CHECK: load i16
+; CHECK: load i32
+
+
+define i32 @Load4(i32* nocapture %p) nounwind uwtable sanitize_memory {
+entry:
+  %0 = load i32, i32* %p
+  ret i32 %0
+}
+
+; CHECK-LABEL: @Load4
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: ptrtoint {{.*}} [[PARAM_SHADOW]]
+; Load the shadow of %p and check it
+; CHECK: load i64
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; Load the value from %p. This is done before accessing the shadow
+; to ease atomic handling.
+; CHECK: load i32
+; CHECK: [[REG:%[0-9]+]] = bitcast i32* %p to i8*
+; CHECK: @__msan_metadata_ptr_for_load_4(i8* [[REG]])
+; Load the shadow and origin.
+; CHECK: load i32
+; CHECK: load i32
+
+define i64 @Load8(i64* nocapture %p) nounwind uwtable sanitize_memory {
+entry:
+  %0 = load i64, i64* %p
+  ret i64 %0
+}
+
+; CHECK-LABEL: @Load8
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: ptrtoint {{.*}} [[PARAM_SHADOW]]
+; Load the shadow of %p and check it
+; CHECK: load i64
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; Load the value from %p. This is done before accessing the shadow
+; to ease atomic handling.
+; CHECK: load i64
+; CHECK: [[REG:%[0-9]+]] = bitcast i64* %p to i8*
+; CHECK: @__msan_metadata_ptr_for_load_8(i8* [[REG]])
+; Load the shadow and origin.
+; CHECK: load i64
+; CHECK: load i32
+
+define i128 @Load16(i128* nocapture %p) nounwind uwtable sanitize_memory {
+entry:
+  %0 = load i128, i128* %p
+  ret i128 %0
+}
+
+; CHECK-LABEL: @Load16
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK-LABEL: entry.split:
+; CHECK: ptrtoint {{.*}} [[PARAM_SHADOW]]
+; Load the shadow of %p and check it
+; CHECK: load i64
+; CHECK: icmp
+; CHECK: br i1
+; CHECK-LABEL: <label>
+; Load the value from %p. This is done before accessing the shadow
+; to ease atomic handling.
+; CHECK: load i128
+; CHECK: [[REG:%[0-9]+]] = bitcast i128* %p to i8*
+; CHECK: @__msan_metadata_ptr_for_load_n(i8* [[REG]], i64 16)
+; Load the shadow and origin.
+; CHECK: load i128
+; CHECK: load i32
+
+
+; Test kernel-specific va_list instrumentation
+
+%struct.__va_list_tag = type { i32, i32, i8*, i8* }
+declare void @llvm.va_start(i8*) nounwind
+declare void @llvm.va_end(i8*)
+@.str = private unnamed_addr constant [4 x i8] c"%d\0A\00", align 1
+declare dso_local i32 @VAListFn(i8*, %struct.__va_list_tag*) local_unnamed_addr
+
+; Function Attrs: nounwind uwtable
+define dso_local i32 @VarArgFn(i8* %fmt, ...) local_unnamed_addr sanitize_memory #0 {
+entry:
+  %args = alloca [1 x %struct.__va_list_tag], align 16
+  %0 = bitcast [1 x %struct.__va_list_tag]* %args to i8*
+  %arraydecay = getelementptr inbounds [1 x %struct.__va_list_tag], [1 x %struct.__va_list_tag]* %args, i64 0, i64 0
+  call void @llvm.va_start(i8* nonnull %0)
+  %call = call i32 @VAListFn(i8* %fmt, %struct.__va_list_tag* nonnull %arraydecay)
+  call void @llvm.va_end(i8* nonnull %0)
+  ret i32 %call
+}
+
+; Kernel is built without SSE support.
+attributes #0 = { "target-features"="+fxsr,+x87,-sse" }
+
+; CHECK-LABEL: @VarArgFn
+; CHECK: @__msan_get_context_state()
+; CHECK: [[VA_ARG_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 2
+; CHECK: [[VA_ARG_ORIGIN:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 3
+; CHECK: [[VA_ARG_OVERFLOW_SIZE:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 4
+
+; CHECK-LABEL: entry.split:
+; CHECK: [[OSIZE:%[0-9]+]] = load i64, i64* [[VA_ARG_OVERFLOW_SIZE]]
+; Register save area is 48 bytes for non-SSE builds.
+; CHECK: [[SIZE:%[0-9]+]] = add i64 48, [[OSIZE]]
+; CHECK: [[SHADOWS:%[0-9]+]] = alloca i8, i64 [[SIZE]]
+; CHECK: [[VA_ARG_SHADOW]]
+; CHECK: call void @llvm.memcpy{{.*}}(i8* align 8 [[SHADOWS]], {{.*}}, i64 [[SIZE]]
+; CHECK: [[ORIGINS:%[0-9]+]] = alloca i8, i64 [[SIZE]]
+; CHECK: [[VA_ARG_ORIGIN]]
+; CHECK: call void @llvm.memcpy{{.*}}(i8* align 8 [[ORIGINS]], {{.*}}, i64 [[SIZE]]
+; CHECK: call i32 @VAListFn
+
+; Function Attrs: nounwind uwtable
+define dso_local void @VarArgCaller() local_unnamed_addr sanitize_memory {
+entry:
+  %call = tail call i32 (i8*, ...) @VarArgFn(i8* getelementptr inbounds ([4 x i8], [4 x i8]* @.str, i64 0, i64 0), i32 123)
+  ret void
+}
+
+; CHECK-LABEL: @VarArgCaller
+
+; CHECK-LABEL: entry:
+; CHECK: @__msan_get_context_state()
+; CHECK: [[PARAM_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 0
+; CHECK: [[VA_ARG_SHADOW:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 2
+; CHECK: [[VA_ARG_OVERFLOW_SIZE:%[a-z0-9_]+]] = getelementptr {{.*}} i32 0, i32 4
+
+; CHECK-LABEL: entry.split:
+; CHECK: [[PARAM_SI:%[_a-z0-9]+]] = ptrtoint {{.*}} [[PARAM_SHADOW]]
+; CHECK: [[ARG1_S:%[_a-z0-9]+]] = inttoptr i64 [[PARAM_SI]] to i64*
+; First argument is initialized
+; CHECK: store i64 0, i64* [[ARG1_S]]
+
+; Dangling cast of va_arg_shadow[0], unused because the first argument is fixed.
+; CHECK: [[VA_CAST0:%[_a-z0-9]+]] = ptrtoint {{.*}} [[VA_ARG_SHADOW]] to i64
+
+; CHECK: [[VA_CAST1:%[_a-z0-9]+]] = ptrtoint {{.*}} [[VA_ARG_SHADOW]] to i64
+; CHECK: [[ARG1_SI:%[_a-z0-9]+]] = add i64 [[VA_CAST1]], 8
+; CHECK: [[PARG1_S:%[_a-z0-9]+]] = inttoptr i64 [[ARG1_SI]] to i32*
+
+; Shadow for 123 is 0.
+; CHECK: store i32 0, i32* [[ARG1_S]]
+
+; CHECK: store i64 0, i64* [[VA_ARG_OVERFLOW_SIZE]]
+; CHECK: call i32 (i8*, ...) @VarArgFn({{.*}} @.str{{.*}} i32 123)

test/Instrumentation/MemorySanitizer/store-origin.ll

-; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=1 -S | FileCheck -check-prefix=CHECK -check-prefix=CHECK-ORIGINS1 %s
-; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=2 -S | FileCheck -check-prefix=CHECK -check-prefix=CHECK-ORIGINS2 %s
+; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=1 -S | FileCheck -check-prefixes=CHECK,CHECK-MSAN,CHECK-ORIGINS1 %s
+; RUN: opt < %s -msan -msan-check-access-address=0 -msan-track-origins=2 -S | FileCheck -check-prefixes=CHECK,CHECK-MSAN,CHECK-ORIGINS2 %s
+; RUN: opt < %s -msan -msan-kernel=1 -msan-check-access-address=0 -S | FileCheck -check-prefixes=CHECK,CHECK-KMSAN,CHECK-ORIGINS2 %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 origin instrumentation of stores.
 ; Check that debug info for origin propagation code is set correctly.
 

 !18 = !{!19, !19, i64 0}
 !19 = !{!"int", !20, i64 0}
 !20 = !{!"omnipotent char", !21, i64 0}
 !21 = !{!"Simple C/C++ TBAA"}
 !22 = !DILocation(line: 3, scope: !4)
 
 
 ; CHECK-LABEL: @Store
-; CHECK: load {{.*}} @__msan_param_tls
-; CHECK: [[ORIGIN:%[01-9a-z]+]] = load {{.*}} @__msan_param_origin_tls
+
+; CHECK-MSAN: load {{.*}} @__msan_param_tls
+; CHECK-MSAN: [[ORIGIN:%[01-9a-z]+]] = load {{.*}} @__msan_param_origin_tls
+
+; CHECK-KMSAN-LABEL: entry.split:
+; CHECK-KMSAN: %param_shadow
+; CHECK-KMSAN: load i32, i32*
+; CHECK-KMSAN: %param_origin
+; CHECK-KMSAN: [[ORIGIN:%[01-9a-z]+]] = load i32, i32*
+
 ; CHECK: store {{.*}}!dbg ![[DBG:[01-9]+]]
 ; CHECK: icmp
 ; CHECK: br i1
 ; CHECK: <label>
 
 ; Origin tracking level 1: simply store the origin value
 ; CHECK-ORIGINS1: store i32 {{.*}}[[ORIGIN]],{{.*}}!dbg !{{.*}}[[DBG]]
 
 ; Origin tracking level 2: pass origin value through __msan_chain_origin and store the result.
 ; CHECK-ORIGINS2: [[ORIGIN2:%[01-9a-z]+]] = call i32 @__msan_chain_origin(i32 {{.*}}[[ORIGIN]])
 ; CHECK-ORIGINS2: store i32 {{.*}}[[ORIGIN2]],{{.*}}!dbg !{{.*}}[[DBG]]
 
 ; CHECK: br label{{.*}}!dbg !{{.*}}[[DBG]]
 ; CHECK: <label>
 ; CHECK: store{{.*}}!dbg !{{.*}}[[DBG]]
 ; CHECK: ret void