https://github.com/ramosian-glider updated 
https://github.com/llvm/llvm-project/pull/202603

>From 1357dab4bdeb9e32992c870ca07d0a300e93fc33 Mon Sep 17 00:00:00 2001
From: Alexander Potapenko <[email protected]>
Date: Thu, 28 May 2026 16:11:22 +0200
Subject: [PATCH] [msan] Introduce KMSAN intrinsics for region instrumentation

This patch introduces `llvm.kmsan.instrumentation.begin`, `end`, and 
`update.context`
intrinsics to allow KernelMemorySanitizer to selectively instrument a specific 
region
of an otherwise uninstrumented function (marked with 
`disable_sanitizer_instrumentation`).

When the function is marked with `disable_sanitizer_instrumentation`, the 
standard
MemorySanitizer function-level ABI prologue is skipped. Instead, the intrinsic
`llvm.kmsan.instrumentation.begin` indicates the start of the code region that
should be instrumented, and `llvm.kmsan.instrumentation.end` marks the end. 
State
propagation and check insertion are only performed on instructions within this 
region.
`llvm.kmsan.instrumentation.update.context` re-initializes the MSan context 
state.

To support complex CFGs while avoiding dominance tree violations, the context 
state
pointer is stored in an entry block `alloca`, and region intrinsics update this
alloca with fresh TLS pointers. The CFG dataflow pass enforcing these regions 
relies
on a monotonic transfer function to ensure O(N) linear time convergence and 
prevent
infinite loops in cyclic graphs.

If these intrinsics are used in a function that is already fully instrumented 
or has
`no_sanitize("memory")`, they are ignored and a diagnostic warning is emitted.
If the MSan pass does not run on a module (e.g. if the module itself has the
`nosanitize_memory` flag), the backend (SelectionDAG, FastISel) trivially 
ignores
these intrinsics, preventing a backend crash.

Documentation is updated in LangRef.rst and MemorySanitizer.rst.
---
 clang/docs/MemorySanitizer.rst                |   5 +
 llvm/docs/LangRef.rst                         |  30 ++
 llvm/include/llvm/IR/Intrinsics.td            |   5 +
 llvm/lib/CodeGen/SelectionDAG/FastISel.cpp    |   3 +
 .../SelectionDAG/SelectionDAGBuilder.cpp      |   6 +
 .../Instrumentation/MemorySanitizer.cpp       | 277 +++++++++++++++---
 .../Generic/kmsan-region-intrinsics.ll        |  21 ++
 .../MemorySanitizer/kmsan-intrinsics.ll       | 143 +++++++++
 .../kmsan-region-domination.ll                |  39 +++
 9 files changed, 488 insertions(+), 41 deletions(-)
 create mode 100644 llvm/test/CodeGen/Generic/kmsan-region-intrinsics.ll
 create mode 100644 
llvm/test/Instrumentation/MemorySanitizer/kmsan-intrinsics.ll
 create mode 100644 
llvm/test/Instrumentation/MemorySanitizer/kmsan-region-domination.ll

diff --git a/clang/docs/MemorySanitizer.rst b/clang/docs/MemorySanitizer.rst
index b40846a75073c..c342a9ce06e1d 100644
--- a/clang/docs/MemorySanitizer.rst
+++ b/clang/docs/MemorySanitizer.rst
@@ -101,6 +101,11 @@ positives and therefore should be used with care, and only 
if absolutely
 required; for example for certain code that cannot tolerate any instrumentation
 and resulting side-effects. This attribute overrides ``no_sanitize("memory")``.
 
+In the context of KernelMemorySanitizer (KMSAN), this attribute skips the 
standard
+function-level ABI prologue and MSan instrumentation. However, you can use the
+``llvm.kmsan.instrumentation.*`` intrinsics within a function marked with this 
attribute
+to selectively enable instrumentation for a specific region of code.
+
 Interaction of Inlining with Disabling Sanitizer Instrumentation
 -----------------------------------------------------------------
 
diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index 2cbc927ba5e93..d558ace99b801 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -32854,3 +32854,33 @@ taken. A late optimization pass will convert this 
intrinsic to either
 ``llvm.cond.loop(true)`` or ``llvm.cond.loop(pred)``, where ``pred``
 is a predicate for a conditional branch leading to the intrinsic call,
 if possible.
+
+.. _llvm.kmsan.instrumentation:
+
+'``llvm.kmsan.instrumentation.*``' Intrinsics
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Syntax:
+"""""""
+
+::
+
+      declare void @llvm.kmsan.instrumentation.begin()
+      declare void @llvm.kmsan.instrumentation.update.context()
+      declare void @llvm.kmsan.instrumentation.end()
+
+Overview:
+"""""""""
+
+The '``llvm.kmsan.instrumentation.*``' intrinsics allow selectively enabling 
MemorySanitizer instrumentation within a specific region of an otherwise 
uninstrumented function (e.g., one marked with the 
``disable_sanitizer_instrumentation`` attribute).
+
+Semantics:
+""""""""""
+
+These intrinsics are used by KernelMemorySanitizer to mark a region of code 
where instrumentation is enabled within an otherwise uninstrumented function.
+
+- ``llvm.kmsan.instrumentation.begin`` marks the start of the instrumented 
region.
+- ``llvm.kmsan.instrumentation.update.context`` updates the shadow context for 
parameter passing/returning.
+- ``llvm.kmsan.instrumentation.end`` marks the end of the instrumented region.
+
+When the function is marked with ``disable_sanitizer_instrumentation``, the 
standard MemorySanitizer function-level ABI prologue is skipped. Instead, the 
intrinsic ``llvm.kmsan.instrumentation.begin`` indicates the start of the code 
region that should be instrumented, and ``llvm.kmsan.instrumentation.end`` 
marks the end of the region. State propagation is only performed on 
instructions within this region. ``llvm.kmsan.instrumentation.update.context`` 
can be used to re-initialize the MSan context state within the region if needed.
diff --git a/llvm/include/llvm/IR/Intrinsics.td 
b/llvm/include/llvm/IR/Intrinsics.td
index b400374ed1649..cece0cb76dd36 100644
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/llvm/include/llvm/IR/Intrinsics.td
@@ -2681,6 +2681,11 @@ def int_load_relative: 
DefaultAttrsIntrinsic<[llvm_ptr_ty], [llvm_ptr_ty, llvm_a
 def int_asan_check_memaccess :
   Intrinsic<[],[llvm_ptr_ty, llvm_i32_ty], [ImmArg<ArgIndex<1>>]>;
 
+// KMSAN intrinsics to delimit regions of interest and update the context.
+def int_kmsan_instrumentation_begin : Intrinsic<[], [], []>;
+def int_kmsan_instrumentation_end : Intrinsic<[], [], []>;
+def int_kmsan_instrumentation_update_context : Intrinsic<[], [], []>;
+
 // Spin in an infinite loop (using instructions specified by the target) iff 
the
 // argument is true. Used to implement efficient conditional traps.
 def int_cond_loop : Intrinsic<[], [llvm_i1_ty], [IntrNoMem, 
IntrHasSideEffects]>;
diff --git a/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp 
b/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
index 1b6d7b57c2c58..39922d6287340 100644
--- a/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/FastISel.cpp
@@ -1385,6 +1385,9 @@ bool FastISel::selectIntrinsicCall(const IntrinsicInst 
*II) {
   case Intrinsic::lifetime_end:
   // The donothing intrinsic does, well, nothing.
   case Intrinsic::donothing:
+  case Intrinsic::kmsan_instrumentation_begin:
+  case Intrinsic::kmsan_instrumentation_update_context:
+  case Intrinsic::kmsan_instrumentation_end:
   // Neither does the sideeffect intrinsic.
   case Intrinsic::sideeffect:
   // Neither does the assume intrinsic; it's also OK not to codegen its 
operand.
diff --git a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp 
b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
index 05c75c484889e..99874b36504a6 100644
--- a/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
+++ b/llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp
@@ -3420,6 +3420,9 @@ void SelectionDAGBuilder::visitInvoke(const InvokeInst 
&I) {
       llvm_unreachable("Cannot invoke this intrinsic");
     case Intrinsic::donothing:
       // Ignore invokes to @llvm.donothing: jump directly to the next BB.
+    case Intrinsic::kmsan_instrumentation_begin:
+    case Intrinsic::kmsan_instrumentation_update_context:
+    case Intrinsic::kmsan_instrumentation_end:
     case Intrinsic::seh_try_begin:
     case Intrinsic::seh_scope_begin:
     case Intrinsic::seh_try_end:
@@ -7810,6 +7813,9 @@ void SelectionDAGBuilder::visitIntrinsicCall(const 
CallInst &I,
     return;
   }
   case Intrinsic::donothing:
+  case Intrinsic::kmsan_instrumentation_begin:
+  case Intrinsic::kmsan_instrumentation_update_context:
+  case Intrinsic::kmsan_instrumentation_end:
   case Intrinsic::seh_try_begin:
   case Intrinsic::seh_scope_begin:
   case Intrinsic::seh_try_end:
diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp 
b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
index c45ec68f3cd07..0caa9860b8e47 100644
--- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -170,11 +170,13 @@
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/GlobalValue.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InlineAsm.h"
+#include "llvm/IR/InstIterator.h"
 #include "llvm/IR/InstVisitor.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
@@ -1247,10 +1249,55 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
   SmallVector<std::pair<IntrinsicInst *, AllocaInst *>, 16> LifetimeStartList;
   SmallVector<StoreInst *, 16> StoreList;
   int64_t SplittableBlocksCount = 0;
+  bool IsRegionInstrumentedOnly = false;
+
+  Value *ContextStateAlloca = nullptr;
+
+  Value *getKmsanTLS(IRBuilderBase &IRB, int Index, const Twine &Name) {
+    Value *Ctx = IRB.CreateLoad(MS.PtrTy, ContextStateAlloca);
+    Constant *Zero = IRB.getInt32(0);
+    return IRB.CreateGEP(MS.MsanContextStateTy, Ctx,
+                         {Zero, IRB.getInt32(Index)}, Name);
+  }
+
+  Value *getParamTLS(IRBuilderBase &IRB) {
+    return IsRegionInstrumentedOnly ? getKmsanTLS(IRB, 0, "param_shadow")
+                                    : MS.ParamTLS;
+  }
+  Value *getRetvalTLS(IRBuilderBase &IRB) {
+    return IsRegionInstrumentedOnly ? getKmsanTLS(IRB, 1, "retval_shadow")
+                                    : MS.RetvalTLS;
+  }
+  Value *getVAArgTLS(IRBuilderBase &IRB) {
+    return IsRegionInstrumentedOnly ? getKmsanTLS(IRB, 2, "va_arg_shadow")
+                                    : MS.VAArgTLS;
+  }
+  Value *getVAArgOriginTLS(IRBuilderBase &IRB) {
+    return IsRegionInstrumentedOnly ? getKmsanTLS(IRB, 3, "va_arg_origin")
+                                    : MS.VAArgOriginTLS;
+  }
+  Value *getVAArgOverflowSizeTLS(IRBuilderBase &IRB) {
+    return IsRegionInstrumentedOnly
+               ? getKmsanTLS(IRB, 4, "va_arg_overflow_size")
+               : MS.VAArgOverflowSizeTLS;
+  }
+  Value *getParamOriginTLS(IRBuilderBase &IRB) {
+    return IsRegionInstrumentedOnly ? getKmsanTLS(IRB, 5, "param_origin")
+                                    : MS.ParamOriginTLS;
+  }
+  Value *getRetvalOriginTLS(IRBuilderBase &IRB) {
+    return IsRegionInstrumentedOnly ? getKmsanTLS(IRB, 6, "retval_origin")
+                                    : MS.RetvalOriginTLS;
+  }
+
+  DenseMap<Instruction *, bool> InRegion;
 
   MemorySanitizerVisitor(Function &F, MemorySanitizer &MS,
                          const TargetLibraryInfo &TLI)
       : F(F), MS(MS), VAHelper(CreateVarArgHelper(F, MS, *this)), TLI(&TLI) {
+    IsRegionInstrumentedOnly =
+        F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation);
+
     bool SanitizeFunction =
         F.hasFnAttribute(Attribute::SanitizeMemory) && !ClDisableChecks;
     InsertChecks = SanitizeFunction;
@@ -1270,7 +1317,7 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
         IRBuilder<>(&F.getEntryBlock(), F.getEntryBlock().getFirstNonPHIIt())
             .CreateIntrinsic(Intrinsic::donothing, {});
 
-    if (MS.CompileKernel) {
+    if (MS.CompileKernel && !IsRegionInstrumentedOnly) {
       IRBuilder<> IRB(FnPrologueEnd);
       insertKmsanPrologue(IRB);
     }
@@ -1615,6 +1662,10 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
   // Returns the last instruction in the new prologue
   void insertKmsanPrologue(IRBuilder<> &IRB) {
     Value *ContextState = IRB.CreateCall(MS.MsanGetContextStateFn, {});
+    if (IsRegionInstrumentedOnly) {
+      IRB.CreateStore(ContextState, ContextStateAlloca);
+      return;
+    }
     Constant *Zero = IRB.getInt32(0);
     MS.ParamTLS = IRB.CreateGEP(MS.MsanContextStateTy, ContextState,
                                 {Zero, IRB.getInt32(0)}, "param_shadow");
@@ -1638,6 +1689,48 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
 
   /// Add MemorySanitizer instrumentation to a function.
   bool runOnFunction() {
+    if (IsRegionInstrumentedOnly) {
+      IRBuilder<> AllocaIRB(&F.getEntryBlock(),
+                            F.getEntryBlock().getFirstNonPHIIt());
+      ContextStateAlloca =
+          AllocaIRB.CreateAlloca(MS.PtrTy, nullptr, "msan_context_state");
+
+      DenseMap<BasicBlock *, bool> BlockInRegionIn;
+      DenseMap<BasicBlock *, bool> BlockInRegionOut;
+      bool Changed = true;
+      BlockInRegionIn[&F.getEntryBlock()] = false;
+      while (Changed) {
+        Changed = false;
+        for (BasicBlock &BB : F) {
+          bool State = BlockInRegionIn[&BB];
+          for (Instruction &I : BB) {
+            if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
+              if (II->getIntrinsicID() ==
+                  Intrinsic::kmsan_instrumentation_begin)
+                State = true;
+              else if (II->getIntrinsicID() ==
+                       Intrinsic::kmsan_instrumentation_end)
+                State = false;
+            }
+            InRegion[&I] = State;
+          }
+          if (BlockInRegionOut[&BB] != State) {
+            BlockInRegionOut[&BB] = State;
+            Changed = true;
+          }
+          for (BasicBlock *Succ : successors(&BB)) {
+            if (!BlockInRegionIn.count(Succ)) {
+              BlockInRegionIn[Succ] = State;
+              Changed = true;
+            } else if (State && !BlockInRegionIn[Succ]) {
+              BlockInRegionIn[Succ] = true;
+              Changed = true;
+            }
+          }
+        }
+      }
+    }
+
     // Iterate all BBs in depth-first order and create shadow instructions
     // for all instructions (where applicable).
     // For PHI nodes we create dummy shadow PHIs which will be finalized later.
@@ -1646,11 +1739,22 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
 
     // `visit` above only collects instructions. Process them after iterating
     // CFG to avoid requirement on CFG transformations.
-    for (Instruction *I : Instructions)
+    for (Instruction *I : Instructions) {
+      if (IsRegionInstrumentedOnly) {
+        bool State = InRegion[I];
+        PropagateShadow = State;
+        InsertChecks = State;
+      }
       InstVisitor<MemorySanitizerVisitor>::visit(*I);
+    }
 
     // Finalize PHI nodes.
     for (PHINode *PN : ShadowPHINodes) {
+      if (IsRegionInstrumentedOnly) {
+        bool State = InRegion[PN];
+        PropagateShadow = State;
+        InsertChecks = State;
+      }
       PHINode *PNS = cast<PHINode>(getShadow(PN));
       PHINode *PNO = MS.TrackOrigins ? cast<PHINode>(getOrigin(PN)) : nullptr;
       size_t NumValues = PN->getNumIncomingValues();
@@ -1667,14 +1771,17 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
     // instrumenting only allocas.
     if (ClHandleLifetimeIntrinsics) {
       for (auto Item : LifetimeStartList) {
-        instrumentAlloca(*Item.second, Item.first);
+        if (!IsRegionInstrumentedOnly || InRegion[Item.first])
+          instrumentAlloca(*Item.second, Item.first);
         AllocaSet.remove(Item.second);
       }
     }
     // Poison the allocas for which we didn't instrument the corresponding
     // lifetime intrinsics.
-    for (AllocaInst *AI : AllocaSet)
-      instrumentAlloca(*AI);
+    for (AllocaInst *AI : AllocaSet) {
+      if (!IsRegionInstrumentedOnly || InRegion[AI])
+        instrumentAlloca(*AI);
+    }
 
     // Insert shadow value checks.
     materializeChecks();
@@ -1976,7 +2083,7 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
   ///
   /// Shadow = ParamTLS+ArgOffset.
   Value *getShadowPtrForArgument(IRBuilder<> &IRB, int ArgOffset) {
-    return IRB.CreatePtrAdd(MS.ParamTLS,
+    return IRB.CreatePtrAdd(getParamTLS(IRB),
                             ConstantInt::get(MS.IntptrTy, ArgOffset), 
"_msarg");
   }
 
@@ -1984,20 +2091,20 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
   Value *getOriginPtrForArgument(IRBuilder<> &IRB, int ArgOffset) {
     if (!MS.TrackOrigins)
       return nullptr;
-    return IRB.CreatePtrAdd(MS.ParamOriginTLS,
+    return IRB.CreatePtrAdd(getParamOriginTLS(IRB),
                             ConstantInt::get(MS.IntptrTy, ArgOffset),
                             "_msarg_o");
   }
 
   /// Compute the shadow address for a retval.
   Value *getShadowPtrForRetval(IRBuilder<> &IRB) {
-    return IRB.CreatePointerCast(MS.RetvalTLS, IRB.getPtrTy(0), "_msret");
+    return IRB.CreatePointerCast(getRetvalTLS(IRB), IRB.getPtrTy(0), "_msret");
   }
 
   /// Compute the origin address for a retval.
-  Value *getOriginPtrForRetval() {
+  Value *getOriginPtrForRetval(IRBuilderBase &IRB) {
     // We keep a single origin for the entire retval. Might be too optimistic.
-    return MS.RetvalOriginTLS;
+    return getRetvalOriginTLS(IRB);
   }
 
   /// Set SV to be the shadow value for V.
@@ -2083,6 +2190,11 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
       return AllOnes;
     }
     if (Argument *A = dyn_cast<Argument>(V)) {
+      if (IsRegionInstrumentedOnly) {
+        if (MS.TrackOrigins && !OriginMap.count(A))
+          setOrigin(A, getCleanOrigin());
+        return getCleanShadow(V);
+      }
       // For arguments we compute the shadow on demand and store it in the map.
       Value *&ShadowPtr = ShadowMap[V];
       if (ShadowPtr)
@@ -2353,6 +2465,7 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
     // Don't want to visit if we're in the prologue
     if (isInPrologue(I))
       return;
+
     if (!DebugCounter::shouldExecute(DebugInstrumentInstruction)) {
       LLVM_DEBUG(dbgs() << "Skipping instruction: " << I << "\n");
       // We still need to set the shadow and origin to clean values.
@@ -2407,12 +2520,23 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
   /// Stores the corresponding shadow and (optionally) origin.
   /// Optionally, checks that the store address is fully defined.
   void visitStoreInst(StoreInst &I) {
+    if (IsRegionInstrumentedOnly && !PropagateShadow) {
+      setShadow(&I, getCleanShadow(&I));
+      setOrigin(&I, getCleanOrigin());
+      return;
+    }
     StoreList.push_back(&I);
     if (ClCheckAccessAddress)
       insertCheckShadowOf(I.getPointerOperand(), &I);
   }
 
   void handleCASOrRMW(Instruction &I) {
+    if (IsRegionInstrumentedOnly && !PropagateShadow) {
+      setShadow(&I, getCleanShadow(&I));
+      setOrigin(&I, getCleanOrigin());
+      return;
+    }
+
     assert(isa<AtomicRMWInst>(I) || isa<AtomicCmpXchgInst>(I));
 
     IRBuilder<> IRB(&I);
@@ -7354,6 +7478,34 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
   }
 
   void visitIntrinsicInst(IntrinsicInst &I) {
+    switch (I.getIntrinsicID()) {
+    case Intrinsic::kmsan_instrumentation_begin:
+    case Intrinsic::kmsan_instrumentation_update_context: {
+      if (IsRegionInstrumentedOnly) {
+        NextNodeIRBuilder IRB(&I);
+        insertKmsanPrologue(IRB);
+      } else {
+        F.getContext().diagnose(DiagnosticInfoOptimizationFailure(
+            F, F.getSubprogram(),
+            "llvm.kmsan.instrumentation intrinsics are ignored in fully "
+            "instrumented or no_sanitize(\"memory\") functions"));
+      }
+      I.eraseFromParent();
+      return;
+    }
+    case Intrinsic::kmsan_instrumentation_end:
+      if (!IsRegionInstrumentedOnly) {
+        F.getContext().diagnose(DiagnosticInfoOptimizationFailure(
+            F, F.getSubprogram(),
+            "llvm.kmsan.instrumentation intrinsics are ignored in fully "
+            "instrumented or no_sanitize(\"memory\") functions"));
+      }
+      I.eraseFromParent();
+      return;
+    default:
+      break;
+    }
+
     if (maybeHandleCrossPlatformIntrinsic(I))
       return;
 
@@ -7425,6 +7577,12 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
   }
 
   void visitCallBase(CallBase &CB) {
+    if (IsRegionInstrumentedOnly && !PropagateShadow) {
+      setShadow(&CB, getCleanShadow(&CB));
+      setOrigin(&CB, getCleanOrigin());
+      return;
+    }
+
     assert(!CB.getMetadata(LLVMContext::MD_nosanitize));
     if (CB.isInlineAsm()) {
       // For inline asm (either a call to asm function, or callbr instruction),
@@ -7625,7 +7783,8 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
         "_msret");
     setShadow(&CB, RetvalShadow);
     if (MS.TrackOrigins)
-      setOrigin(&CB, IRBAfter.CreateLoad(MS.OriginTy, 
getOriginPtrForRetval()));
+      setOrigin(&CB,
+                IRBAfter.CreateLoad(MS.OriginTy, getOriginPtrForRetval(IRB)));
   }
 
   bool isAMustTailRetVal(Value *RetVal) {
@@ -7639,6 +7798,9 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
   }
 
   void visitReturnInst(ReturnInst &I) {
+    if (IsRegionInstrumentedOnly && !PropagateShadow)
+      return;
+
     IRBuilder<> IRB(&I);
     Value *RetVal = I.getReturnValue();
     if (!RetVal)
@@ -7666,7 +7828,7 @@ struct MemorySanitizerVisitor : public 
InstVisitor<MemorySanitizerVisitor> {
     if (StoreShadow) {
       IRB.CreateAlignedStore(Shadow, ShadowPtr, kShadowTLSAlignment);
       if (MS.TrackOrigins && StoreOrigin)
-        IRB.CreateStore(getOrigin(RetVal), getOriginPtrForRetval());
+        IRB.CreateStore(getOrigin(RetVal), getOriginPtrForRetval(IRB));
     }
   }
 
@@ -8064,14 +8226,15 @@ struct VarArgHelperBase : public VarArgHelper {
       : F(F), MS(MS), MSV(MSV), VAListTagSize(VAListTagSize) {}
 
   Value *getShadowAddrForVAArgument(IRBuilder<> &IRB, unsigned ArgOffset) {
-    Value *Base = IRB.CreatePointerCast(MS.VAArgTLS, MS.IntptrTy);
+    Value *Base = IRB.CreatePointerCast(MSV.getVAArgTLS(IRB), MS.IntptrTy);
     return IRB.CreateAdd(Base, ConstantInt::get(MS.IntptrTy, ArgOffset));
   }
 
   /// Compute the shadow address for a given va_arg.
   Value *getShadowPtrForVAArgument(IRBuilder<> &IRB, unsigned ArgOffset) {
-    return IRB.CreatePtrAdd(
-        MS.VAArgTLS, ConstantInt::get(MS.IntptrTy, ArgOffset), "_msarg_va_s");
+    return IRB.CreatePtrAdd(MSV.getVAArgTLS(IRB),
+                            ConstantInt::get(MS.IntptrTy, ArgOffset),
+                            "_msarg_va_s");
   }
 
   /// Compute the shadow address for a given va_arg.
@@ -8088,7 +8251,7 @@ struct VarArgHelperBase : public VarArgHelper {
     // getOriginPtrForVAArgument() is always called after
     // getShadowPtrForVAArgument(), so __msan_va_arg_origin_tls can never
     // overflow.
-    return IRB.CreatePtrAdd(MS.VAArgOriginTLS,
+    return IRB.CreatePtrAdd(MSV.getVAArgOriginTLS(IRB),
                             ConstantInt::get(MS.IntptrTy, ArgOffset),
                             "_msarg_va_o");
   }
@@ -8279,7 +8442,7 @@ struct VarArgAMD64Helper : public VarArgHelperBase {
     }
     Constant *OverflowSize =
         ConstantInt::get(IRB.getInt64Ty(), OverflowOffset - AMD64FpEndOffset);
-    IRB.CreateStore(OverflowSize, MS.VAArgOverflowSizeTLS);
+    IRB.CreateStore(OverflowSize, MSV.getVAArgOverflowSizeTLS(IRB));
   }
 
   void finalizeInstrumentation() override {
@@ -8290,7 +8453,7 @@ struct VarArgAMD64Helper : public VarArgHelperBase {
       // va_arg_tls somewhere in the function entry block.
       IRBuilder<> IRB(MSV.FnPrologueEnd);
       VAArgOverflowSize =
-          IRB.CreateLoad(IRB.getInt64Ty(), MS.VAArgOverflowSizeTLS);
+          IRB.CreateLoad(IRB.getInt64Ty(), MSV.getVAArgOverflowSizeTLS(IRB));
       Value *CopySize = IRB.CreateAdd(
           ConstantInt::get(MS.IntptrTy, AMD64FpEndOffset), VAArgOverflowSize);
       VAArgTLSCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
@@ -8301,13 +8464,14 @@ struct VarArgAMD64Helper : public VarArgHelperBase {
       Value *SrcSize = IRB.CreateBinaryIntrinsic(
           Intrinsic::umin, CopySize,
           ConstantInt::get(MS.IntptrTy, kParamTLSSize));
-      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
+      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MSV.getVAArgTLS(IRB),
                        kShadowTLSAlignment, SrcSize);
       if (MS.TrackOrigins) {
         VAArgTLSOriginCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), 
CopySize);
         VAArgTLSOriginCopy->setAlignment(kShadowTLSAlignment);
         IRB.CreateMemCpy(VAArgTLSOriginCopy, kShadowTLSAlignment,
-                         MS.VAArgOriginTLS, kShadowTLSAlignment, SrcSize);
+                         MSV.getVAArgOriginTLS(IRB), kShadowTLSAlignment,
+                         SrcSize);
       }
     }
 
@@ -8456,7 +8620,7 @@ struct VarArgAArch64Helper : public VarArgHelperBase {
     }
     Constant *OverflowSize =
         ConstantInt::get(IRB.getInt64Ty(), OverflowOffset - 
AArch64VAEndOffset);
-    IRB.CreateStore(OverflowSize, MS.VAArgOverflowSizeTLS);
+    IRB.CreateStore(OverflowSize, MSV.getVAArgOverflowSizeTLS(IRB));
   }
 
   // Retrieve a va_list field of 'void*' size.
@@ -8482,7 +8646,7 @@ struct VarArgAArch64Helper : public VarArgHelperBase {
       // va_arg_tls somewhere in the function entry block.
       IRBuilder<> IRB(MSV.FnPrologueEnd);
       VAArgOverflowSize =
-          IRB.CreateLoad(IRB.getInt64Ty(), MS.VAArgOverflowSizeTLS);
+          IRB.CreateLoad(IRB.getInt64Ty(), MSV.getVAArgOverflowSizeTLS(IRB));
       Value *CopySize = IRB.CreateAdd(
           ConstantInt::get(MS.IntptrTy, AArch64VAEndOffset), 
VAArgOverflowSize);
       VAArgTLSCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), CopySize);
@@ -8493,7 +8657,7 @@ struct VarArgAArch64Helper : public VarArgHelperBase {
       Value *SrcSize = IRB.CreateBinaryIntrinsic(
           Intrinsic::umin, CopySize,
           ConstantInt::get(MS.IntptrTy, kParamTLSSize));
-      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
+      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MSV.getVAArgTLS(IRB),
                        kShadowTLSAlignment, SrcSize);
     }
 
@@ -8683,14 +8847,15 @@ struct VarArgPowerPC64Helper : public VarArgHelperBase {
         ConstantInt::get(MS.IntptrTy, VAArgOffset - VAArgBase);
     // Here using VAArgOverflowSizeTLS as VAArgSizeTLS to avoid creation of
     // a new class member i.e. it is the total size of all VarArgs.
-    IRB.CreateStore(TotalVAArgSize, MS.VAArgOverflowSizeTLS);
+    IRB.CreateStore(TotalVAArgSize, MSV.getVAArgOverflowSizeTLS(IRB));
   }
 
   void finalizeInstrumentation() override {
     assert(!VAArgSize && !VAArgTLSCopy &&
            "finalizeInstrumentation called twice");
     IRBuilder<> IRB(MSV.FnPrologueEnd);
-    VAArgSize = IRB.CreateLoad(IRB.getInt64Ty(), MS.VAArgOverflowSizeTLS);
+    VAArgSize =
+        IRB.CreateLoad(IRB.getInt64Ty(), MSV.getVAArgOverflowSizeTLS(IRB));
     Value *CopySize = VAArgSize;
 
     if (!VAStartInstrumentationList.empty()) {
@@ -8705,7 +8870,7 @@ struct VarArgPowerPC64Helper : public VarArgHelperBase {
       Value *SrcSize = IRB.CreateBinaryIntrinsic(
           Intrinsic::umin, CopySize,
           ConstantInt::get(IRB.getInt64Ty(), kParamTLSSize));
-      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
+      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MSV.getVAArgTLS(IRB),
                        kShadowTLSAlignment, SrcSize);
     }
 
@@ -8819,14 +8984,14 @@ struct VarArgPowerPC32Helper : public VarArgHelperBase {
         ConstantInt::get(MS.IntptrTy, VAArgOffset - VAArgBase);
     // Here using VAArgOverflowSizeTLS as VAArgSizeTLS to avoid creation of
     // a new class member i.e. it is the total size of all VarArgs.
-    IRB.CreateStore(TotalVAArgSize, MS.VAArgOverflowSizeTLS);
+    IRB.CreateStore(TotalVAArgSize, MSV.getVAArgOverflowSizeTLS(IRB));
   }
 
   void finalizeInstrumentation() override {
     assert(!VAArgSize && !VAArgTLSCopy &&
            "finalizeInstrumentation called twice");
     IRBuilder<> IRB(MSV.FnPrologueEnd);
-    VAArgSize = IRB.CreateLoad(MS.IntptrTy, MS.VAArgOverflowSizeTLS);
+    VAArgSize = IRB.CreateLoad(MS.IntptrTy, MSV.getVAArgOverflowSizeTLS(IRB));
     Value *CopySize = VAArgSize;
 
     if (!VAStartInstrumentationList.empty()) {
@@ -8841,7 +9006,7 @@ struct VarArgPowerPC32Helper : public VarArgHelperBase {
       Value *SrcSize = IRB.CreateBinaryIntrinsic(
           Intrinsic::umin, CopySize,
           ConstantInt::get(MS.IntptrTy, kParamTLSSize));
-      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
+      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MSV.getVAArgTLS(IRB),
                        kShadowTLSAlignment, SrcSize);
     }
 
@@ -9104,7 +9269,7 @@ struct VarArgSystemZHelper : public VarArgHelperBase {
     }
     Constant *OverflowSize = ConstantInt::get(
         IRB.getInt64Ty(), OverflowOffset - SystemZOverflowOffset);
-    IRB.CreateStore(OverflowSize, MS.VAArgOverflowSizeTLS);
+    IRB.CreateStore(OverflowSize, MSV.getVAArgOverflowSizeTLS(IRB));
   }
 
   void copyRegSaveArea(IRBuilder<> &IRB, Value *VAListTag) {
@@ -9165,7 +9330,7 @@ struct VarArgSystemZHelper : public VarArgHelperBase {
       // va_arg_tls somewhere in the function entry block.
       IRBuilder<> IRB(MSV.FnPrologueEnd);
       VAArgOverflowSize =
-          IRB.CreateLoad(IRB.getInt64Ty(), MS.VAArgOverflowSizeTLS);
+          IRB.CreateLoad(IRB.getInt64Ty(), MSV.getVAArgOverflowSizeTLS(IRB));
       Value *CopySize =
           IRB.CreateAdd(ConstantInt::get(MS.IntptrTy, SystemZOverflowOffset),
                         VAArgOverflowSize);
@@ -9177,13 +9342,14 @@ struct VarArgSystemZHelper : public VarArgHelperBase {
       Value *SrcSize = IRB.CreateBinaryIntrinsic(
           Intrinsic::umin, CopySize,
           ConstantInt::get(MS.IntptrTy, kParamTLSSize));
-      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
+      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MSV.getVAArgTLS(IRB),
                        kShadowTLSAlignment, SrcSize);
       if (MS.TrackOrigins) {
         VAArgTLSOriginCopy = IRB.CreateAlloca(Type::getInt8Ty(*MS.C), 
CopySize);
         VAArgTLSOriginCopy->setAlignment(kShadowTLSAlignment);
         IRB.CreateMemCpy(VAArgTLSOriginCopy, kShadowTLSAlignment,
-                         MS.VAArgOriginTLS, kShadowTLSAlignment, SrcSize);
+                         MSV.getVAArgOriginTLS(IRB), kShadowTLSAlignment,
+                         SrcSize);
       }
     }
 
@@ -9259,14 +9425,14 @@ struct VarArgI386Helper : public VarArgHelperBase {
     Constant *TotalVAArgSize = ConstantInt::get(MS.IntptrTy, VAArgOffset);
     // Here using VAArgOverflowSizeTLS as VAArgSizeTLS to avoid creation of
     // a new class member i.e. it is the total size of all VarArgs.
-    IRB.CreateStore(TotalVAArgSize, MS.VAArgOverflowSizeTLS);
+    IRB.CreateStore(TotalVAArgSize, MSV.getVAArgOverflowSizeTLS(IRB));
   }
 
   void finalizeInstrumentation() override {
     assert(!VAArgSize && !VAArgTLSCopy &&
            "finalizeInstrumentation called twice");
     IRBuilder<> IRB(MSV.FnPrologueEnd);
-    VAArgSize = IRB.CreateLoad(MS.IntptrTy, MS.VAArgOverflowSizeTLS);
+    VAArgSize = IRB.CreateLoad(MS.IntptrTy, MSV.getVAArgOverflowSizeTLS(IRB));
     Value *CopySize = VAArgSize;
 
     if (!VAStartInstrumentationList.empty()) {
@@ -9280,7 +9446,7 @@ struct VarArgI386Helper : public VarArgHelperBase {
       Value *SrcSize = IRB.CreateBinaryIntrinsic(
           Intrinsic::umin, CopySize,
           ConstantInt::get(MS.IntptrTy, kParamTLSSize));
-      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
+      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MSV.getVAArgTLS(IRB),
                        kShadowTLSAlignment, SrcSize);
     }
 
@@ -9344,14 +9510,14 @@ struct VarArgGenericHelper : public VarArgHelperBase {
     Constant *TotalVAArgSize = ConstantInt::get(MS.IntptrTy, VAArgOffset);
     // Here using VAArgOverflowSizeTLS as VAArgSizeTLS to avoid creation of
     // a new class member i.e. it is the total size of all VarArgs.
-    IRB.CreateStore(TotalVAArgSize, MS.VAArgOverflowSizeTLS);
+    IRB.CreateStore(TotalVAArgSize, MSV.getVAArgOverflowSizeTLS(IRB));
   }
 
   void finalizeInstrumentation() override {
     assert(!VAArgSize && !VAArgTLSCopy &&
            "finalizeInstrumentation called twice");
     IRBuilder<> IRB(MSV.FnPrologueEnd);
-    VAArgSize = IRB.CreateLoad(MS.IntptrTy, MS.VAArgOverflowSizeTLS);
+    VAArgSize = IRB.CreateLoad(MS.IntptrTy, MSV.getVAArgOverflowSizeTLS(IRB));
     Value *CopySize = VAArgSize;
 
     if (!VAStartInstrumentationList.empty()) {
@@ -9365,7 +9531,7 @@ struct VarArgGenericHelper : public VarArgHelperBase {
       Value *SrcSize = IRB.CreateBinaryIntrinsic(
           Intrinsic::umin, CopySize,
           ConstantInt::get(MS.IntptrTy, kParamTLSSize));
-      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MS.VAArgTLS,
+      IRB.CreateMemCpy(VAArgTLSCopy, kShadowTLSAlignment, MSV.getVAArgTLS(IRB),
                        kShadowTLSAlignment, SrcSize);
     }
 
@@ -9472,8 +9638,37 @@ bool MemorySanitizer::sanitizeFunction(Function &F, 
TargetLibraryInfo &TLI) {
   if (!CompileKernel && F.getName() == kMsanModuleCtorName)
     return false;
 
-  if (F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation))
-    return false;
+  bool HasKmsanRegionIntrinsics = false;
+  for (Instruction &I : instructions(F)) {
+    if (auto *II = dyn_cast<IntrinsicInst>(&I)) {
+      if (II->getIntrinsicID() == Intrinsic::kmsan_instrumentation_begin ||
+          II->getIntrinsicID() == Intrinsic::kmsan_instrumentation_end ||
+          II->getIntrinsicID() ==
+              Intrinsic::kmsan_instrumentation_update_context) {
+        HasKmsanRegionIntrinsics = true;
+        break;
+      }
+    }
+  }
+
+  bool IsFullyInstrumented =
+      F.hasFnAttribute(Attribute::SanitizeMemory) && !ClDisableChecks;
+  bool HasNoSanitizeMemory =
+      !IsFullyInstrumented &&
+      !F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation);
+
+  if (HasKmsanRegionIntrinsics &&
+      (IsFullyInstrumented || HasNoSanitizeMemory)) {
+    F.getContext().diagnose(DiagnosticInfoOptimizationFailure(
+        F, F.getSubprogram(),
+        "llvm.kmsan.instrumentation intrinsics are ignored in fully "
+        "instrumented or no_sanitize(\"memory\") functions"));
+  }
+
+  if (F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation)) {
+    if (!HasKmsanRegionIntrinsics)
+      return false;
+  }
 
   MemorySanitizerVisitor Visitor(F, *this, TLI);
 
diff --git a/llvm/test/CodeGen/Generic/kmsan-region-intrinsics.ll 
b/llvm/test/CodeGen/Generic/kmsan-region-intrinsics.ll
new file mode 100644
index 0000000000000..3d62828e87059
--- /dev/null
+++ b/llvm/test/CodeGen/Generic/kmsan-region-intrinsics.ll
@@ -0,0 +1,21 @@
+; RUN: llc -O0 < %s | FileCheck %s
+; RUN: llc -O2 < %s | FileCheck %s
+; RUN: llc -O0 -fast-isel < %s | FileCheck %s
+
+; Verify that the backend successfully lowers KMSAN region intrinsics to 
nothing
+; and does not crash when compiling a module where MSan did not run.
+
+declare void @llvm.kmsan.instrumentation.begin()
+declare void @llvm.kmsan.instrumentation.update.context()
+declare void @llvm.kmsan.instrumentation.end()
+
+define void @test_kmsan_region_intrinsics() {
+; CHECK-LABEL: test_kmsan_region_intrinsics:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    ret{{[ql]?}}
+entry:
+  call void @llvm.kmsan.instrumentation.begin()
+  call void @llvm.kmsan.instrumentation.update.context()
+  call void @llvm.kmsan.instrumentation.end()
+  ret void
+}
\ No newline at end of file
diff --git a/llvm/test/Instrumentation/MemorySanitizer/kmsan-intrinsics.ll 
b/llvm/test/Instrumentation/MemorySanitizer/kmsan-intrinsics.ll
new file mode 100644
index 0000000000000..5060ec52b730b
--- /dev/null
+++ b/llvm/test/Instrumentation/MemorySanitizer/kmsan-intrinsics.ll
@@ -0,0 +1,143 @@
+; RUN: opt < %s -msan-check-access-address=0 -msan-kernel=1 -S -passes=msan 
2>&1 | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define void @instrumented_function() sanitize_memory {
+entry:
+  call void @llvm.kmsan.instrumentation.begin()
+  call void @llvm.kmsan.instrumentation.end()
+  ret void
+}
+; CHECK: warning: {{.*}}llvm.kmsan.instrumentation intrinsics are ignored in 
fully instrumented or no_sanitize("memory") functions
+
+define void @no_sanitize_function() {
+entry:
+  call void @llvm.kmsan.instrumentation.begin()
+  call void @llvm.kmsan.instrumentation.end()
+  ret void
+}
+; CHECK: warning: {{.*}}llvm.kmsan.instrumentation intrinsics are ignored in 
fully instrumented or no_sanitize("memory") functions
+
+define void @noinstr_function(ptr %p, i64 %val) 
disable_sanitizer_instrumentation {
+entry:
+  %x = load i64, ptr %p
+  %add = add i64 %x, %val
+  call void @llvm.kmsan.instrumentation.begin()
+  %y = load i64, ptr %p
+  %add2 = add i64 %y, %add
+  store i64 %add2, ptr %p
+  call void @llvm.kmsan.instrumentation.update.context()
+  %z = load i64, ptr %p
+  store i64 %z, ptr %p
+  call void @llvm.kmsan.instrumentation.end()
+  %w = load i64, ptr %p
+  ret void
+}
+
+; CHECK-LABEL: @noinstr_function(
+; CHECK-NOT: __msan_get_context_state
+; CHECK: %x = load i64, ptr %p
+; CHECK: %add = add i64 %x, %val
+; CHECK: [[CTX1:%.*]] = call ptr @__msan_get_context_state()
+; CHECK: store ptr [[CTX1]], ptr %msan_context_state
+; CHECK: %y = load i64, ptr %p
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_load_8(ptr %p)
+; CHECK: %add2 = add i64 %y, %add
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_store_8(ptr %p)
+; CHECK: store i64 %add2, ptr %p
+; CHECK: [[CTX2:%.*]] = call ptr @__msan_get_context_state()
+; CHECK: store ptr [[CTX2]], ptr %msan_context_state
+; CHECK: %z = load i64, ptr %p
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_load_8(ptr %p)
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_store_8(ptr %p)
+; CHECK: store i64 %z, ptr %p
+; CHECK-NOT: call ptr @__msan_get_context_state()
+; CHECK-NOT: call { ptr, ptr } @__msan_metadata_ptr_for_load_8
+; CHECK: %w = load i64, ptr %p
+; CHECK: ret void
+
+declare void @llvm.kmsan.instrumentation.begin()
+declare void @llvm.kmsan.instrumentation.end()
+declare void @llvm.kmsan.instrumentation.update.context()
+
+define void @noinstr_with_branch(ptr %p, i1 %cond) 
disable_sanitizer_instrumentation {
+entry:
+  call void @llvm.kmsan.instrumentation.begin()
+  br i1 %cond, label %if.then, label %if.else
+
+if.then:
+  %a = load i64, ptr %p
+  store i64 %a, ptr %p
+  br label %if.end
+
+if.else:
+  %b = load i64, ptr %p
+  store i64 %b, ptr %p
+  br label %if.end
+
+if.end:
+  call void @llvm.kmsan.instrumentation.end()
+  ret void
+}
+
+; CHECK-LABEL: @noinstr_with_branch(
+; CHECK: [[CTX3:%.*]] = call ptr @__msan_get_context_state()
+; CHECK: store ptr [[CTX3]], ptr %msan_context_state
+; CHECK: if.then:
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_load_8(ptr %p)
+; CHECK: if.else:
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_load_8(ptr %p)
+
+define void @noinstr_with_loop(ptr %p, i64 %n) 
disable_sanitizer_instrumentation {
+entry:
+  call void @llvm.kmsan.instrumentation.begin()
+  br label %loop
+
+loop:
+  %i = phi i64 [ 0, %entry ], [ %inc, %loop ]
+  %c = load i64, ptr %p
+  store i64 %c, ptr %p
+  %inc = add i64 %i, 1
+  %cmp = icmp ult i64 %inc, %n
+  br i1 %cmp, label %loop, label %exit
+
+exit:
+  call void @llvm.kmsan.instrumentation.end()
+  ret void
+}
+
+; CHECK-LABEL: @noinstr_with_loop(
+; CHECK: [[CTX4:%.*]] = call ptr @__msan_get_context_state()
+; CHECK: store ptr [[CTX4]], ptr %msan_context_state
+; CHECK: loop:
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_load_8(ptr %p)
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_store_8(ptr %p)
+; CHECK: exit:
+
+declare i64 @dummy_call(i64)
+
+define i64 @noinstr_with_various_insts(ptr %p, i64 %val) 
disable_sanitizer_instrumentation {
+entry:
+  %a = alloca i64
+  %rmw1 = atomicrmw add ptr %p, i64 %val seq_cst
+  %call1 = call i64 @dummy_call(i64 %rmw1)
+  store i64 %call1, ptr %a
+  
+  call void @llvm.kmsan.instrumentation.begin()
+  %rmw2 = atomicrmw add ptr %p, i64 %val seq_cst
+  %call2 = call i64 @dummy_call(i64 %rmw2)
+  store i64 %call2, ptr %a
+  call void @llvm.kmsan.instrumentation.end()
+  
+  %rmw3 = atomicrmw add ptr %p, i64 %val seq_cst
+  %call3 = call i64 @dummy_call(i64 %rmw3)
+  store i64 %call3, ptr %a
+  ret i64 %call3
+}
+
+; CHECK-LABEL: @noinstr_with_various_insts(
+; CHECK-NOT: __msan_metadata
+; CHECK: call ptr @__msan_get_context_state()
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_store_8(ptr %p)
+; CHECK: ret i64 %call3
\ No newline at end of file
diff --git 
a/llvm/test/Instrumentation/MemorySanitizer/kmsan-region-domination.ll 
b/llvm/test/Instrumentation/MemorySanitizer/kmsan-region-domination.ll
new file mode 100644
index 0000000000000..e2df182ebac8d
--- /dev/null
+++ b/llvm/test/Instrumentation/MemorySanitizer/kmsan-region-domination.ll
@@ -0,0 +1,39 @@
+; RUN: opt < %s -msan-check-access-address=0 -msan-kernel=1 -S -passes=msan 
2>&1 | FileCheck %s
+
+; This test verifies that the MemorySanitizer region instrumentation does not
+; generate broken IR (violating dominance constraints) when `update.context` 
+; is called after `begin` in complex CFGs.
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @llvm.kmsan.instrumentation.begin()
+declare void @llvm.kmsan.instrumentation.update.context()
+declare void @llvm.kmsan.instrumentation.end()
+
+define i1 @do_syscall_64(ptr %regs, i32 %nr) disable_sanitizer_instrumentation 
{
+entry:
+  call void @llvm.kmsan.instrumentation.begin()
+  br label %loop
+
+loop:
+  %x = load i64, ptr %regs
+  call void @llvm.kmsan.instrumentation.update.context()
+  %y = load i64, ptr %regs
+  br i1 false, label %loop, label %exit
+
+exit:
+  call void @llvm.kmsan.instrumentation.end()
+  ret i1 true
+}
+
+; CHECK-LABEL: @do_syscall_64(
+; CHECK: entry:
+; CHECK: [[CTX_ALLOCA:%.*]] = alloca ptr
+; CHECK: [[CTX1:%.*]] = call ptr @__msan_get_context_state()
+; CHECK: store ptr [[CTX1]], ptr [[CTX_ALLOCA]]
+; CHECK: loop:
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_load_8
+; CHECK: [[CTX2:%.*]] = call ptr @__msan_get_context_state()
+; CHECK: store ptr [[CTX2]], ptr [[CTX_ALLOCA]]
+; CHECK: call { ptr, ptr } @__msan_metadata_ptr_for_load_8
\ No newline at end of file

_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to