[MSAN] Cache stack traces and chained origins

[guiand]

Summary

Prevent (excessive) expensive stack unwinding and depot lookups by introducing smaller caches external to the depots.

The stack trace depot is cached by "stack trace hash" -- a TLS value constantly updated by instrumentation in the MSAN LLVM pass. The origin depot is cached by a cheap 32-bit hash over the two IDs comprising the edge in the graph.

Each depot is accessed only after missing in a L1 and L2 cache. The L1 is small and per-thread, and the L2 is a few times larger and is global to all threads. Both are direct-mapped using the hash. On collision, the old entry is evicted.

In order to prevent stack unwinding in the case of a cache hit, MSan no longer passes StackTrace objects among its internal functions. Instead, I use a new StackUnwindCtx object which stores the parameters we need for unwinding on demand.

[guiand]

compiler-rt/lib/msan/msan_origin.h

130

I'm only not really sure what to do about these use counts. Updating them in the depot somewhat defeats the purpose of trying to prevent depot accesses through caching.

[eugenis]

compiler-rt/lib/msan/msan_origin.h

130

The purpose of this is to limit the number of new chained origins that one access trace can create. We could cache this value in l1 and l2, and when we are about to create a chained origin, double-check against the global stack depot and refresh the caches.

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1304

KMSan uses this SplitBlock thing to avoid visiting the prologue setup instructions, but it has a cost: it turns all allocas into dynamic allocas, as they are no longer in the entry basic block.

Try to avoid this by turning ActualFnStart into an instruction pointer, that should help with the code size and maybe performance a little.

This could be what is causing the local shadow addresses to be pre-calculated.

3774

The same should be done on ResumeInst and probably more, otherwise this would keep generating unique hash values in an exception-enabled code.

[guiand]

Addressed comments. For handling number of uses per stack trace, this uses a bit of a heuristic:

For some unique origin count O, some unique stack trace count T, and some origins/trace limit L:

check_limit = (O > T * L * 7/8)
if (check_limit) {
  hnd = FindHandle()
  if (hnd.uses() > L * 1/8)
    return early
}
inserted = PutChainedOrigin()
if (check_limit && inserted)
  hnd.uses()++

[eugenis]

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

1114

You need to move ActualFnStart below the new prologue, otherwise ex. the TLS store store will get instrumented with an address check, and a shadow store.

That's why I suggested to make it Instruction* and tweak the visitor loop to skip earlier instruction in the starting BB.

[guiand]

Updated to depend on https://reviews.llvm.org/D85985 (ActualFnStart becomes a Instruction *)

[guiand]

Added a test for hash codegen.

[vitalybuka]

What is the plan for this patch?

compiler-rt/lib/msan/msan.h

 
 bool ProtectRange(uptr beg, uptr end);
 bool InitShadow(bool init_origins);
 char *GetProcSelfMaps();
 void InitializeInterceptors();
 
 void MsanAllocatorInit();
 void MsanAllocatorThreadFinish();
-void MsanDeallocate(StackTrace *stack, void *ptr);
+void MsanDeallocate(StackUnwindCtx *stack, void *ptr);
 
-void *msan_malloc(uptr size, StackTrace *stack);
-void *msan_calloc(uptr nmemb, uptr size, StackTrace *stack);
-void *msan_realloc(void *ptr, uptr size, StackTrace *stack);
-void *msan_reallocarray(void *ptr, uptr nmemb, uptr size, StackTrace *stack);
-void *msan_valloc(uptr size, StackTrace *stack);
-void *msan_pvalloc(uptr size, StackTrace *stack);
-void *msan_aligned_alloc(uptr alignment, uptr size, StackTrace *stack);
-void *msan_memalign(uptr alignment, uptr size, StackTrace *stack);
+void *msan_malloc(uptr size, StackUnwindCtx *stack);
+void *msan_calloc(uptr nmemb, uptr size, StackUnwindCtx *stack);
+void *msan_realloc(void *ptr, uptr size, StackUnwindCtx *stack);
+void *msan_reallocarray(void *ptr, uptr nmemb, uptr size,
+                        StackUnwindCtx *stack);
+void *msan_valloc(uptr size, StackUnwindCtx *stack);
+void *msan_pvalloc(uptr size, StackUnwindCtx *stack);
+void *msan_aligned_alloc(uptr alignment, uptr size, StackUnwindCtx *stack);
+void *msan_memalign(uptr alignment, uptr size, StackUnwindCtx *stack);
 int msan_posix_memalign(void **memptr, uptr alignment, uptr size,
-                        StackTrace *stack);
+                        StackUnwindCtx *stack);
 
 void InstallTrapHandler();
 void InstallAtExitHandler();
 
 const char *GetStackOriginDescr(u32 id, uptr *pc);
 
 void EnterSymbolizer();
 void ExitSymbolizer();
 bool IsInSymbolizer();
 
 struct SymbolizerScope {
   SymbolizerScope() { EnterSymbolizer(); }
   ~SymbolizerScope() { ExitSymbolizer(); }
 };
 
 void PrintWarning(uptr pc, uptr bp);
 void PrintWarningWithOrigin(uptr pc, uptr bp, u32 origin);
 
 // Unpoison first n function arguments.
 void UnpoisonParam(uptr n);
 void UnpoisonThreadLocalState();
 
 // Returns a "chained" origin id, pointing to the given stack trace followed by
 // the previous origin id.
-u32 ChainOrigin(u32 id, StackTrace *stack);
+u32 ChainOrigin(u32 id, StackUnwindCtx *stack);
+
+TraceHash GetTraceHash(uptr pc);
 
 const int STACK_TRACE_TAG_POISON = StackTrace::TAG_CUSTOM + 1;
 
-#define GET_MALLOC_STACK_TRACE                                            \
-  BufferedStackTrace stack;                                               \
-  if (__msan_get_track_origins() && msan_inited)                          \
-    stack.Unwind(StackTrace::GetCurrentPc(), GET_CURRENT_FRAME(),         \
-                 nullptr, common_flags()->fast_unwind_on_malloc,          \
-                 common_flags()->malloc_context_size)
+#define GET_MALLOC_STACK_TRACE                                          \
+  StackUnwindCtx stack(StackTrace::GetCurrentPc(), GET_CURRENT_FRAME(), \
+                       nullptr);                                        \
+  if (__msan_get_track_origins() && msan_inited)                        \
+    stack.RequestFast(common_flags()->fast_unwind_on_malloc)            \
+        .WithMaxDepth(common_flags()->malloc_context_size)              \
+        .WithTraceHash(GetTraceHash(GET_CALLER_PC()));
 
 // For platforms which support slow unwinder only, we restrict the store context
 // size to 1, basically only storing the current pc. We do this because the slow
 // unwinder which is based on libunwind is not async signal safe and causes
 // random freezes in forking applications as well as in signal handlers.
-#define GET_STORE_STACK_TRACE_PC_BP(pc, bp)                                    \
-  BufferedStackTrace stack;                                                    \
-  if (__msan_get_track_origins() > 1 && msan_inited) {                         \
-    int size = flags()->store_context_size;                                    \
-    if (!SANITIZER_CAN_FAST_UNWIND)                                            \
-      size = Min(size, 1);                                                     \
-    stack.Unwind(pc, bp, nullptr, common_flags()->fast_unwind_on_malloc, size);\
+#define GET_STORE_STACK_TRACE_PC_BP(pc, bp)                  \
+  StackUnwindCtx stack(pc, bp, nullptr);                     \
+  if (__msan_get_track_origins() > 1 && msan_inited) {       \
+    int size = flags()->store_context_size;                  \
+    if (!SANITIZER_CAN_FAST_UNWIND)                          \
+      size = Min(size, 1);                                   \
+    stack.RequestFast(common_flags()->fast_unwind_on_malloc) \
+        .WithMaxDepth(size)                                  \
+        .WithTraceHash(GetTraceHash(GET_CALLER_PC()));       \
   }
 
 #define GET_STORE_STACK_TRACE \
   GET_STORE_STACK_TRACE_PC_BP(StackTrace::GetCurrentPc(), GET_CURRENT_FRAME())
 
 #define GET_FATAL_STACK_TRACE_PC_BP(pc, bp)                              \
   BufferedStackTrace stack;                                              \
   if (msan_inited) {                                                     \

compiler-rt/lib/msan/msan.cpp

 THREADLOCAL u32 __msan_va_arg_origin_tls[kMsanParamTlsSize / sizeof(u32)];
 
 SANITIZER_INTERFACE_ATTRIBUTE
 THREADLOCAL u64 __msan_va_arg_overflow_size_tls;
 
 SANITIZER_INTERFACE_ATTRIBUTE
 THREADLOCAL u32 __msan_origin_tls;
 
+SANITIZER_INTERFACE_ATTRIBUTE
+THREADLOCAL TraceHash __msan_ctrace_tls;
+
 static THREADLOCAL int is_in_symbolizer;
 
 extern "C" SANITIZER_WEAK_ATTRIBUTE const int __msan_track_origins;
 
 int __msan_get_track_origins() {
   return &__msan_track_origins ? __msan_track_origins : 0;
 }
 

 }
 
 const char *GetStackOriginDescr(u32 id, uptr *pc) {
   CHECK_LT(id, kNumStackOriginDescrs);
   if (pc) *pc = StackOriginPC[id];
   return StackOriginDescr[id];
 }
 
-u32 ChainOrigin(u32 id, StackTrace *stack) {
+u32 ChainOrigin(u32 id, StackUnwindCtx *stack) {
   MsanThread *t = GetCurrentThread();
   if (t && t->InSignalHandler())
     return id;
 
   Origin o = Origin::FromRawId(id);
-  stack->tag = StackTrace::TAG_UNKNOWN;
+  stack->WithTag(StackTrace::TAG_UNKNOWN);
   Origin chained = Origin::CreateChainedOrigin(o, stack);
   return chained.raw_id();
 }
 
+TraceHash GetTraceHash(uptr pc) {
+  const static auto rotl5 = [](u32 x) { return (x << 5) | (x >> 27); };
+  if (__msan_ctrace_tls)
+    return rotl5(__msan_ctrace_tls) ^ pc;
+  return 0;
+}
+
 } // namespace __msan
 
 void __sanitizer::BufferedStackTrace::UnwindImpl(
     uptr pc, uptr bp, void *context, bool request_fast, u32 max_depth) {
   using namespace __msan;
   MsanThread *t = GetCurrentThread();
   if (!t || !StackTrace::WillUseFastUnwind(request_fast)) {
     // Block reports from our interceptors during _Unwind_Backtrace.

compiler-rt/lib/msan/msan_allocator.cpp

AllocatorCache *GetAllocatorCache(MsanThreadLocalMallocStorage *ms) {
   CHECK_LE(sizeof(AllocatorCache), sizeof(ms->allocator_cache));
   return reinterpret_cast<AllocatorCache *>(ms->allocator_cache);
 }
 
 void MsanThreadLocalMallocStorage::CommitBack() {
   allocator.SwallowCache(GetAllocatorCache(this));
 }
 
-static void *MsanAllocate(StackTrace *stack, uptr size, uptr alignment,
+static void *MsanAllocate(StackUnwindCtx *stack, uptr size, uptr alignment,
                           bool zeroise) {
   if (size > max_malloc_size) {
     if (AllocatorMayReturnNull()) {
       Report("WARNING: MemorySanitizer failed to allocate 0x%zx bytes\n", size);
       return nullptr;
     }
-    ReportAllocationSizeTooBig(size, max_malloc_size, stack);
+    ReportAllocationSizeTooBig(size, max_malloc_size,
+                               &BufferedStackTrace::TLSUnwind(*stack));
   }
   MsanThread *t = GetCurrentThread();
   void *allocated;
   if (t) {
     AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
     allocated = allocator.Allocate(cache, size, alignment);
   } else {
     SpinMutexLock l(&fallback_mutex);
     AllocatorCache *cache = &fallback_allocator_cache;
     allocated = allocator.Allocate(cache, size, alignment);
   }
   if (UNLIKELY(!allocated)) {
     SetAllocatorOutOfMemory();
     if (AllocatorMayReturnNull())
       return nullptr;
-    ReportOutOfMemory(size, stack);
+    ReportOutOfMemory(size, &BufferedStackTrace::TLSUnwind(*stack));
   }
   Metadata *meta =
       reinterpret_cast<Metadata *>(allocator.GetMetaData(allocated));
   meta->requested_size = size;
   if (zeroise) {
     __msan_clear_and_unpoison(allocated, size);
   } else if (flags()->poison_in_malloc) {
     __msan_poison(allocated, size);
     if (__msan_get_track_origins()) {
-      stack->tag = StackTrace::TAG_ALLOC;
+      stack->WithTag(StackTrace::TAG_ALLOC);
       Origin o = Origin::CreateHeapOrigin(stack);
       __msan_set_origin(allocated, size, o.raw_id());
     }
   }
   MSAN_MALLOC_HOOK(allocated, size);
   return allocated;
 }
 
-void MsanDeallocate(StackTrace *stack, void *p) {
+void MsanDeallocate(StackUnwindCtx *stack, void *p) {
   CHECK(p);
   MSAN_FREE_HOOK(p);
   Metadata *meta = reinterpret_cast<Metadata *>(allocator.GetMetaData(p));
   uptr size = meta->requested_size;
   meta->requested_size = 0;
   // This memory will not be reused by anyone else, so we are free to keep it
   // poisoned.
   if (flags()->poison_in_free) {
     __msan_poison(p, size);
     if (__msan_get_track_origins()) {
-      stack->tag = StackTrace::TAG_DEALLOC;
+      stack->WithTag(StackTrace::TAG_DEALLOC);
       Origin o = Origin::CreateHeapOrigin(stack);
       __msan_set_origin(p, size, o.raw_id());
     }
   }
   MsanThread *t = GetCurrentThread();
   if (t) {
     AllocatorCache *cache = GetAllocatorCache(&t->malloc_storage());
     allocator.Deallocate(cache, p);
   } else {
     SpinMutexLock l(&fallback_mutex);
     AllocatorCache *cache = &fallback_allocator_cache;
     allocator.Deallocate(cache, p);
   }
 }
 
-void *MsanReallocate(StackTrace *stack, void *old_p, uptr new_size,
+void *MsanReallocate(StackUnwindCtx *stack, void *old_p, uptr new_size,
                      uptr alignment) {
   Metadata *meta = reinterpret_cast<Metadata*>(allocator.GetMetaData(old_p));
   uptr old_size = meta->requested_size;
   uptr actually_allocated_size = allocator.GetActuallyAllocatedSize(old_p);
   if (new_size <= actually_allocated_size) {
     // We are not reallocating here.
     meta->requested_size = new_size;
     if (new_size > old_size) {
       if (flags()->poison_in_malloc) {
-        stack->tag = StackTrace::TAG_ALLOC;
+        stack->WithTag(StackTrace::TAG_ALLOC);
         PoisonMemory((char *)old_p + old_size, new_size - old_size, stack);
       }
     }
     return old_p;
   }
   uptr memcpy_size = Min(new_size, old_size);
   void *new_p = MsanAllocate(stack, new_size, alignment, false /*zeroise*/);
   if (new_p) {
     CopyMemory(new_p, old_p, memcpy_size, stack);
     MsanDeallocate(stack, old_p);
   }
   return new_p;
 }
 
-void *MsanCalloc(StackTrace *stack, uptr nmemb, uptr size) {
+void *MsanCalloc(StackUnwindCtx *stack, uptr nmemb, uptr size) {
   if (UNLIKELY(CheckForCallocOverflow(size, nmemb))) {
     if (AllocatorMayReturnNull())
       return nullptr;
-    ReportCallocOverflow(nmemb, size, stack);
+    ReportCallocOverflow(nmemb, size, &BufferedStackTrace::TLSUnwind(*stack));
   }
   return MsanAllocate(stack, nmemb * size, sizeof(u64), true);
 }
 
 static uptr AllocationSize(const void *p) {
   if (!p) return 0;
   const void *beg = allocator.GetBlockBegin(p);
   if (beg != p) return 0;
   Metadata *b = (Metadata *)allocator.GetMetaData(p);
   return b->requested_size;
 }
 
-void *msan_malloc(uptr size, StackTrace *stack) {
+void *msan_malloc(uptr size, StackUnwindCtx *stack) {
   return SetErrnoOnNull(MsanAllocate(stack, size, sizeof(u64), false));
 }
 
-void *msan_calloc(uptr nmemb, uptr size, StackTrace *stack) {
+void *msan_calloc(uptr nmemb, uptr size, StackUnwindCtx *stack) {
   return SetErrnoOnNull(MsanCalloc(stack, nmemb, size));
 }
 
-void *msan_realloc(void *ptr, uptr size, StackTrace *stack) {
+void *msan_realloc(void *ptr, uptr size, StackUnwindCtx *stack) {
   if (!ptr)
     return SetErrnoOnNull(MsanAllocate(stack, size, sizeof(u64), false));
   if (size == 0) {
     MsanDeallocate(stack, ptr);
     return nullptr;
   }
   return SetErrnoOnNull(MsanReallocate(stack, ptr, size, sizeof(u64)));
 }
 
-void *msan_reallocarray(void *ptr, uptr nmemb, uptr size, StackTrace *stack) {
+void *msan_reallocarray(void *ptr, uptr nmemb, uptr size,
+                        StackUnwindCtx *stack) {
   if (UNLIKELY(CheckForCallocOverflow(size, nmemb))) {
     errno = errno_ENOMEM;
     if (AllocatorMayReturnNull())
       return nullptr;
-    ReportReallocArrayOverflow(nmemb, size, stack);
+    ReportReallocArrayOverflow(nmemb, size,
+                               &BufferedStackTrace::TLSUnwind(*stack));
   }
   return msan_realloc(ptr, nmemb * size, stack);
 }
 
-void *msan_valloc(uptr size, StackTrace *stack) {
+void *msan_valloc(uptr size, StackUnwindCtx *stack) {
   return SetErrnoOnNull(MsanAllocate(stack, size, GetPageSizeCached(), false));
 }
 
-void *msan_pvalloc(uptr size, StackTrace *stack) {
+void *msan_pvalloc(uptr size, StackUnwindCtx *stack) {
   uptr PageSize = GetPageSizeCached();
   if (UNLIKELY(CheckForPvallocOverflow(size, PageSize))) {
     errno = errno_ENOMEM;
     if (AllocatorMayReturnNull())
       return nullptr;
-    ReportPvallocOverflow(size, stack);
+    ReportPvallocOverflow(size, &BufferedStackTrace::TLSUnwind(*stack));
   }
   // pvalloc(0) should allocate one page.
   size = size ? RoundUpTo(size, PageSize) : PageSize;
   return SetErrnoOnNull(MsanAllocate(stack, size, PageSize, false));
 }
 
-void *msan_aligned_alloc(uptr alignment, uptr size, StackTrace *stack) {
+void *msan_aligned_alloc(uptr alignment, uptr size, StackUnwindCtx *stack) {
   if (UNLIKELY(!CheckAlignedAllocAlignmentAndSize(alignment, size))) {
     errno = errno_EINVAL;
     if (AllocatorMayReturnNull())
       return nullptr;
-    ReportInvalidAlignedAllocAlignment(size, alignment, stack);
+    ReportInvalidAlignedAllocAlignment(size, alignment,
+                                       &BufferedStackTrace::TLSUnwind(*stack));
   }
   return SetErrnoOnNull(MsanAllocate(stack, size, alignment, false));
 }
 
-void *msan_memalign(uptr alignment, uptr size, StackTrace *stack) {
+void *msan_memalign(uptr alignment, uptr size, StackUnwindCtx *stack) {
   if (UNLIKELY(!IsPowerOfTwo(alignment))) {
     errno = errno_EINVAL;
     if (AllocatorMayReturnNull())
       return nullptr;
-    ReportInvalidAllocationAlignment(alignment, stack);
+    ReportInvalidAllocationAlignment(alignment,
+                                     &BufferedStackTrace::TLSUnwind(*stack));
   }
   return SetErrnoOnNull(MsanAllocate(stack, size, alignment, false));
 }
 
 int msan_posix_memalign(void **memptr, uptr alignment, uptr size,
-                        StackTrace *stack) {
+                        StackUnwindCtx *stack) {
   if (UNLIKELY(!CheckPosixMemalignAlignment(alignment))) {
     if (AllocatorMayReturnNull())
       return errno_EINVAL;
-    ReportInvalidPosixMemalignAlignment(alignment, stack);
+    ReportInvalidPosixMemalignAlignment(alignment,
+                                        &BufferedStackTrace::TLSUnwind(*stack));
   }
   void *ptr = MsanAllocate(stack, size, alignment, false);
   if (UNLIKELY(!ptr))
     // OOM error is already taken care of by MsanAllocate.
     return errno_ENOMEM;
   CHECK(IsAligned((uptr)ptr, alignment));
   *memptr = ptr;
   return 0;

compiler-rt/lib/msan/msan_interceptors.cpp

if (UNLIKELY(!msan_inited))
     // Hack: dlsym calls malloc before REAL(malloc) is retrieved from dlsym.
     return AllocateFromLocalPool(size);
   return msan_malloc(size, &stack);
 }
 
 void __msan_allocated_memory(const void *data, uptr size) {
   GET_MALLOC_STACK_TRACE;
   if (flags()->poison_in_malloc) {
-    stack.tag = STACK_TRACE_TAG_POISON;
+    stack.WithTag(STACK_TRACE_TAG_POISON);
     PoisonMemory(data, size, &stack);
   }
 }
 
 void __msan_copy_shadow(void *dest, const void *src, uptr n) {
   GET_STORE_STACK_TRACE;
   MoveShadowAndOrigin(dest, src, n, &stack);
 }
 
 void __sanitizer_dtor_callback(const void *data, uptr size) {
   GET_MALLOC_STACK_TRACE;
   if (flags()->poison_in_dtor) {
-    stack.tag = STACK_TRACE_TAG_POISON;
+    stack.WithTag(STACK_TRACE_TAG_POISON);
     PoisonMemory(data, size, &stack);
   }
 }
 
 template <class Mmap>
 static void *mmap_interceptor(Mmap real_mmap, void *addr, SIZE_T length,
                               int prot, int flags, int fd, OFF64_T offset) {
   SIZE_T rounded_length = RoundUpTo(length, GetPageSize());

compiler-rt/lib/msan/msan_new_delete.cpp

 // Fake std::nothrow_t and std::align_val_t to avoid including <new>.
 namespace std {
   struct nothrow_t {};
   enum class align_val_t: size_t {};
 }  // namespace std
 
 
 // TODO(alekseys): throw std::bad_alloc instead of dying on OOM.
 #define OPERATOR_NEW_BODY(nothrow) \
   GET_MALLOC_STACK_TRACE; \
   void *res = msan_malloc(size, &stack);\
-  if (!nothrow && UNLIKELY(!res)) ReportOutOfMemory(size, &stack);\
+  if (!nothrow && UNLIKELY(!res))     \
+    ReportOutOfMemory(size, &BufferedStackTrace::TLSUnwind(stack)); \
   return res
 #define OPERATOR_NEW_BODY_ALIGN(nothrow) \
   GET_MALLOC_STACK_TRACE;\
   void *res = msan_memalign((uptr)align, size, &stack);\
-  if (!nothrow && UNLIKELY(!res)) ReportOutOfMemory(size, &stack);\
+  if (!nothrow && UNLIKELY(!res))     \
+    ReportOutOfMemory(size, &BufferedStackTrace::TLSUnwind(stack)); \
   return res;
 
 INTERCEPTOR_ATTRIBUTE
 void *operator new(size_t size) { OPERATOR_NEW_BODY(false /*nothrow*/); }
 INTERCEPTOR_ATTRIBUTE
 void *operator new[](size_t size) { OPERATOR_NEW_BODY(false /*nothrow*/); }
 INTERCEPTOR_ATTRIBUTE
 void *operator new(size_t size, std::nothrow_t const&) {

compiler-rt/lib/msan/msan_origin.h

 //===-- msan_origin.h ----------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Origin id utils.
 //===----------------------------------------------------------------------===//
 #ifndef MSAN_ORIGIN_H
 #define MSAN_ORIGIN_H
 
-#include "sanitizer_common/sanitizer_stackdepot.h"
 #include "msan_chained_origin_depot.h"
+#include "sanitizer_common/sanitizer_common.h"
+#include "sanitizer_common/sanitizer_hash.h"
+#include "sanitizer_common/sanitizer_hashcache.h"
+#include "sanitizer_common/sanitizer_stackdepot.h"
 
 namespace __msan {
 
 // Origin handling.
 //
 // Origin is a 32-bit identifier that is attached to any uninitialized value in
 // the program and describes, more or less exactly, how this memory came to be
 // uninitialized.
StackTrace getStackTraceForHeapOrigin() const {
     return StackDepotGet(getHeapId());
   }
 
   static Origin CreateStackOrigin(u32 id) {
     CHECK((id & kStackIdMask) == id);
     return Origin((1 << kHeapShift) | id);
   }
 
-  static Origin CreateHeapOrigin(StackTrace *stack) {
-    u32 stack_id = StackDepotPut(*stack);
+  static Origin CreateHeapOrigin(StackUnwindCtx *stack) {
+    u32 stack_id = StackDepotPutUnwound(*stack);
     CHECK(stack_id);
     CHECK((stack_id & kHeapIdMask) == stack_id);
     return Origin(stack_id);
   }
 
-  static Origin CreateChainedOrigin(Origin prev, StackTrace *stack) {
+  static Origin CreateChainedOrigin(Origin prev, StackUnwindCtx *stack) {
     int depth = prev.isChainedOrigin() ? prev.depth() : 0;
     // depth is the length of the chain minus 1.
     // origin_history_size of 0 means unlimited depth.
     if (flags()->origin_history_size > 0) {
       if (depth + 1 >= flags()->origin_history_size) {
         return prev;
       } else {
         ++depth;
         CHECK(depth < (1 << kDepthBits));
       }
     }
 
-    StackDepotHandle h = StackDepotPut_WithHandle(*stack);
-    if (!h.valid()) return prev;
+    u32 id = StackDepotPutUnwound(*stack);
+    if (!id)
+      return prev;
 
     if (flags()->origin_history_per_stack_limit > 0) {
-      int use_count = h.use_count();
-      if (use_count > flags()->origin_history_per_stack_limit) return prev;
+      // int use_count = h.use_count();
+      // if (use_count > flags()->origin_history_per_stack_limit) return prev;
     }
 
+    static thread_local HashCache<128> chain_cache;
+    static ConcurrentHashCache<1024> chain_cache_l2;
+
     u32 chained_id;
-    bool inserted = ChainedOriginDepotPut(h.id(), prev.raw_id(), &chained_id);
+    u32 chain_hash = id * 4132367431 + prev.raw_id() * 3462451541;
+
+    if (!chain_cache.get(chain_hash, &chained_id)) {
+      // L1 cache miss
+      if (chain_cache_l2.get(chain_hash, &chained_id)) {
+        // L2 cache hit; update L1
+        chain_cache.insert(chain_hash, chained_id);
+      } else {
+        // L1+L2 cache miss; update both caches
+        ChainedOriginDepotPut(id, prev.raw_id(), &chained_id);
         CHECK((chained_id & kChainedIdMask) == chained_id);
+        chain_cache.insert(chain_hash, chained_id);
+        chain_cache_l2.insert(chain_hash, chained_id);
+      }
+    }
 
-    if (inserted && flags()->origin_history_per_stack_limit > 0)
-      h.inc_use_count_unsafe();
+    // if (inserted && flags()->origin_history_per_stack_limit > 0)
+      // h.inc_use_count_unsafe();
 
     return Origin((1 << kHeapShift) | (depth << kDepthShift) | chained_id);
   }
 
   static Origin FromRawId(u32 id) {
     return Origin(id);
   }
 

compiler-rt/lib/msan/msan_poisoning.h

 // Walk [addr, addr+size) app memory region, copying origin tags from the
 // corresponding positions in [src_origin, src_origin+size) where the
 // corresponding shadow in [src_shadow, src_shadow+size) is non-zero.
 void SetOriginIfPoisoned(uptr addr, uptr src_shadow, uptr size, u32 src_origin);
 
 // Copy origin from src (app address) to dst (app address), creating chained
 // origin ids as necessary, without overriding origin for fully initialized
 // quads.
-void CopyOrigin(const void *dst, const void *src, uptr size, StackTrace *stack);
+void CopyOrigin(const void *dst, const void *src, uptr size,
+                StackUnwindCtx *stack);
 
 // memmove() shadow and origin. Dst and src are application addresses.
 // See CopyOrigin() for the origin copying logic.
 void MoveShadowAndOrigin(const void *dst, const void *src, uptr size,
-                         StackTrace *stack);
+                         StackUnwindCtx *stack);
 
 // memcpy() shadow and origin. Dst and src are application addresses.
 // See CopyOrigin() for the origin copying logic.
 void CopyShadowAndOrigin(const void *dst, const void *src, uptr size,
-                         StackTrace *stack);
+                         StackUnwindCtx *stack);
 
 // memcpy() app memory, and do "the right thing" to the corresponding shadow and
 // origin regions.
-void CopyMemory(void *dst, const void *src, uptr size, StackTrace *stack);
+void CopyMemory(void *dst, const void *src, uptr size, StackUnwindCtx *stack);
 
 // Fill shadow will value. Ptr is an application address.
 void SetShadow(const void *ptr, uptr size, u8 value);
 
 // Set origin for the memory region.
 void SetOrigin(const void *dst, uptr size, u32 origin);
 
 // Mark memory region uninitialized, with origins.
-void PoisonMemory(const void *dst, uptr size, StackTrace *stack);
+void PoisonMemory(const void *dst, uptr size, StackUnwindCtx *stack);
 
 }  // namespace __msan
 
 #endif  // MSAN_POISONING_H

compiler-rt/lib/msan/msan_poisoning.cpp

void SetOriginIfPoisoned(uptr addr, uptr src_shadow, uptr size,
   uptr src_s = src_shadow;
   uptr src_s_end = src_s + size;
 
   for (; src_s < src_s_end; ++dst_s, ++src_s)
     if (*(u8 *)src_s) *(u32 *)SHADOW_TO_ORIGIN(dst_s & ~3UL) = src_origin;
 }
 
 void CopyOrigin(const void *dst, const void *src, uptr size,
-                StackTrace *stack) {
+                StackUnwindCtx *stack) {
   if (!MEM_IS_APP(dst) || !MEM_IS_APP(src)) return;
 
   uptr d = (uptr)dst;
   uptr beg = d & ~3UL;
   // Copy left unaligned origin if that memory is poisoned.
   if (beg < d) {
     u32 o = GetOriginIfPoisoned((uptr)src, d - beg);
     if (o) {
if (beg < end) {
     } else {
       REAL(memcpy)((void *)MEM_TO_ORIGIN(beg), (void *)MEM_TO_ORIGIN(s),
                    end - beg);
     }
   }
 }
 
 void MoveShadowAndOrigin(const void *dst, const void *src, uptr size,
-                         StackTrace *stack) {
+                         StackUnwindCtx *stack) {
   if (!MEM_IS_APP(dst)) return;
   if (!MEM_IS_APP(src)) return;
   if (src == dst) return;
   REAL(memmove)((void *)MEM_TO_SHADOW((uptr)dst),
                 (void *)MEM_TO_SHADOW((uptr)src), size);
   if (__msan_get_track_origins()) CopyOrigin(dst, src, size, stack);
 }
 
 void CopyShadowAndOrigin(const void *dst, const void *src, uptr size,
-                         StackTrace *stack) {
+                         StackUnwindCtx *stack) {
   if (!MEM_IS_APP(dst)) return;
   if (!MEM_IS_APP(src)) return;
   REAL(memcpy)((void *)MEM_TO_SHADOW((uptr)dst),
                (void *)MEM_TO_SHADOW((uptr)src), size);
   if (__msan_get_track_origins()) CopyOrigin(dst, src, size, stack);
 }
 
-void CopyMemory(void *dst, const void *src, uptr size, StackTrace *stack) {
+void CopyMemory(void *dst, const void *src, uptr size, StackUnwindCtx *stack) {
   REAL(memcpy)(dst, src, size);
   CopyShadowAndOrigin(dst, src, size, stack);
 }
 
 void SetShadow(const void *ptr, uptr size, u8 value) {
   uptr PageSize = GetPageSizeCached();
   uptr shadow_beg = MEM_TO_SHADOW(ptr);
   uptr shadow_end = shadow_beg + size;
void SetOrigin(const void *dst, uptr size, u32 origin) {
   if (beg & 7ULL) {
     *(u32 *)beg = origin;
     beg += 4;
   }
   for (uptr addr = beg; addr < (end & ~7UL); addr += 8) *(u64 *)addr = origin64;
   if (end & 7ULL) *(u32 *)(end - 4) = origin;
 }
 
-void PoisonMemory(const void *dst, uptr size, StackTrace *stack) {
+void PoisonMemory(const void *dst, uptr size, StackUnwindCtx *stack) {
   SetShadow(dst, size, (u8)-1);
 
   if (__msan_get_track_origins()) {
     Origin o = Origin::CreateHeapOrigin(stack);
     SetOrigin(dst, size, o.raw_id());
   }
 }
 
 }  // namespace __msan

compiler-rt/lib/sanitizer_common/sanitizer_hashcache.h

+#pragma once
+
+#include <utility>
+
+#include "sanitizer_allocator_internal.h"
+#include "sanitizer_common.h"
+#include "sanitizer_internal_defs.h"
+#include "sanitizer_platform.h"
+
+namespace __sanitizer {
+
+template <uptr capacity>
+struct HashCache {
+  u32 keys[capacity] = {};
+  u32 vals[capacity] = {};
+
+  constexpr HashCache() {}
+
+  bool get(u32 key, u32 *val) {
+    uptr idx = key & (capacity - 1);
+    if (keys[idx] != key)
+      return false;
+    if (val)
+      *val = vals[idx];
+    return true;
+  }
+
+  bool insert(u32 key, u32 val) {
+    uptr idx = key & (capacity - 1);
+    if (keys[idx] == key)
+      return false;
+
+    u32 *keyslot = &keys[idx];
+    u32 *valslot = &vals[idx];
+
+    *keyslot = key;
+    *valslot = val;
+    return true;
+  }
+};
+
+template <uptr N>
+struct ConcurrentHashCache {
+  atomic_uint64_t data[N];
+
+  bool get(u32 key, u32 *val) {
+    u64 data = atomic_load(&this->data[key % N], memory_order_relaxed);
+    u32 found_key = data;
+    u32 found_val = data >> 32;
+    *val = found_val;
+    return key == found_key;
+  }
+
+  void insert(u32 key, u32 val) {
+    u64 data = (u64)val << 32 | key;
+    atomic_store(&this->data[key % N], data, memory_order_relaxed);
+  }
+};
+}  // namespace __sanitizer

compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.h

struct StackDepotHandle {
   int use_count();
   void inc_use_count_unsafe();
 };
 
 const int kStackDepotMaxUseCount = 1U << (SANITIZER_ANDROID ? 16 : 20);
 
 StackDepotStats *StackDepotGetStats();
 u32 StackDepotPut(StackTrace stack);
-StackDepotHandle StackDepotPut_WithHandle(StackTrace stack);
+u32 StackDepotPutUnwound(const StackUnwindCtx &ctx);
+
 // Retrieves a stored stack trace by the id.
 StackTrace StackDepotGet(u32 id);
 
 void StackDepotLockAll();
 void StackDepotUnlockAll();
 
 // Instantiating this class creates a snapshot of StackDepot which can be
 // efficiently queried with StackDepotGet(). You can use it concurrently with

compiler-rt/lib/sanitizer_common/sanitizer_stackdepot.cpp

 //===-- sanitizer_stackdepot.cpp ------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file is shared between AddressSanitizer and ThreadSanitizer
 // run-time libraries.
 //===----------------------------------------------------------------------===//
 
 #include "sanitizer_stackdepot.h"
 
 #include "sanitizer_common.h"
 #include "sanitizer_hash.h"
+#include "sanitizer_hashcache.h"
 #include "sanitizer_stackdepotbase.h"
 
 namespace __sanitizer {
 
 struct StackDepotNode {
   StackDepotNode *link;
   u32 id;
   atomic_uint32_t hash_and_use_count; // hash_bits : 12; use_count : 20;
struct StackDepotNode {
   static const u32 kHashMask = ~kUseCountMask;
 
   typedef StackTrace args_type;
   bool eq(u32 hash, const args_type &args) const {
     u32 hash_bits =
         atomic_load(&hash_and_use_count, memory_order_relaxed) & kHashMask;
     if ((hash & kHashMask) != hash_bits || args.size != size || args.tag != tag)
       return false;
-    uptr i = 0;
-    for (; i < size; i++) {
-      if (stack[i] != args.trace[i]) return false;
-    }
+    // uptr i = 0;
+    // for (; i < size; i++) {
+    //   if (stack[i] != args.trace[i]) return false;
+    // }
     return true;
   }
   static uptr storage_size(const args_type &args) {
     return sizeof(StackDepotNode) + (args.size - 1) * sizeof(uptr);
   }
   static u32 hash(const args_type &args) {
     MurMur2HashBuilder H(args.size * sizeof(uptr));
     for (uptr i = 0; i < args.size; i++) H.add(args.trace[i]);
StackDepotStats *StackDepotGetStats() {
   return theDepot.GetStats();
 }
 
 u32 StackDepotPut(StackTrace stack) {
   StackDepotHandle h = theDepot.Put(stack);
   return h.valid() ? h.id() : 0;
 }
 
-StackDepotHandle StackDepotPut_WithHandle(StackTrace stack) {
-  return theDepot.Put(stack);
+u32 StackDepotPutUnwound(const StackUnwindCtx &ctx) {
+  static thread_local HashCache<1024> trace_hash_map;
+  static ConcurrentHashCache<16384> trace_map_l2;
+
+  const static auto PutTrace = [](const StackUnwindCtx &ctx) {
+    // Keep the temp trace in TLS because it's too big to be comfortable
+    // in the stack.
+    static thread_local BufferedStackTrace tmp_trace;
+    tmp_trace.Reset();
+    tmp_trace.Unwind(ctx);
+    return StackDepotPut(tmp_trace);
+  };
+
+  if (!ctx.trace_hash) {
+    return PutTrace(ctx);
+  }
+  u32 used_hash = ctx.trace_hash ^ ctx.tag;
+
+  u32 depot_id;
+  if (trace_hash_map.get(used_hash, &depot_id))
+    // L1 cache hit
+    return depot_id;
+  if (trace_map_l2.get(used_hash, &depot_id)) {
+    // L2 cache hit, so update L1
+    trace_hash_map.insert(used_hash, depot_id);
+    return depot_id;
+  }
+  // L1+L2 cache misses; update both caches
+  depot_id = PutTrace(ctx);
+  trace_hash_map.insert(used_hash, depot_id);
+  trace_map_l2.insert(used_hash, depot_id);
+  return depot_id;
 }
 
 StackTrace StackDepotGet(u32 id) {
   return theDepot.Get(id);
 }
 
 void StackDepotLockAll() {
   theDepot.LockAll();

compiler-rt/lib/sanitizer_common/sanitizer_stacktrace.h

#elif defined(__powerpc__) || defined(__powerpc64__) || defined(__aarch64__)
   return pc - 4;
 #elif defined(__sparc__) || defined(__mips__)
   return pc - 8;
 #else
   return pc - 1;
 #endif
 }
 
+typedef u32 TraceHash;
+
+struct StackUnwindCtx {
+  uptr pc;
+  uptr bp;
+  void *context;
+  uptr stack_top, stack_bottom;
+  TraceHash trace_hash = 0;
+  u32 max_depth = kStackTraceMax;
+  u32 tag = 0;
+  bool use_stack_bounds = false;
+  bool request_fast = false;
+
+  StackUnwindCtx(uptr pc, uptr bp, void *ctx) : pc(pc), bp(bp), context(ctx) {}
+  StackUnwindCtx &WithTraceHash(TraceHash hash) {
+    this->trace_hash = hash;
+    return *this;
+  }
+  StackUnwindCtx &WithStackBounds(uptr stack_top, uptr stack_bottom) {
+    this->stack_top = stack_top;
+    this->stack_bottom = stack_bottom;
+    this->use_stack_bounds = true;
+    return *this;
+  };
+  StackUnwindCtx &RequestFast(bool fast) {
+    this->request_fast = fast;
+    return *this;
+  }
+  StackUnwindCtx &WithMaxDepth(u32 max) {
+    this->max_depth = max;
+    return *this;
+  }
+  StackUnwindCtx &WithTag(u32 tag) {
+    this->tag = tag;
+    return *this;
+  }
+};
+
 // StackTrace that owns the buffer used to store the addresses.
 struct BufferedStackTrace : public StackTrace {
   uptr trace_buffer[kStackTraceMax];
   uptr top_frame_bp;  // Optional bp of a top frame.
 
   BufferedStackTrace() : StackTrace(trace_buffer, 0), top_frame_bp(0) {}
 
   void Init(const uptr *pcs, uptr cnt, uptr extra_top_pc = 0);
if (max_depth <= 1) {
       return;
     }
     UnwindImpl(pc, bp, context, request_fast, max_depth);
   }
 
   void Unwind(u32 max_depth, uptr pc, uptr bp, void *context, uptr stack_top,
               uptr stack_bottom, bool request_fast_unwind);
 
+  void Unwind(const StackUnwindCtx &ctx) {
+    if (ctx.use_stack_bounds) {
+      Unwind(ctx.max_depth, ctx.pc, ctx.bp, ctx.context, ctx.stack_top,
+             ctx.stack_bottom, ctx.request_fast);
+    } else {
+      Unwind(ctx.pc, ctx.bp, ctx.context, ctx.request_fast, ctx.max_depth);
+    }
+    tag = ctx.tag;
+  }
+
+  static const BufferedStackTrace &TLSUnwind(const StackUnwindCtx &ctx) {
+    thread_local BufferedStackTrace trace;
+    trace.Reset();
+    trace.Unwind(ctx);
+    return trace;
+  }
+
   void Reset() {
     *static_cast<StackTrace *>(this) = StackTrace(trace_buffer, 0);
     top_frame_bp = 0;
   }
 
  private:
   // Every runtime defines its own implementation of this method
-  void UnwindImpl(uptr pc, uptr bp, void *context, bool request_fast,
-                  u32 max_depth);
+  __attribute__((weak)) void UnwindImpl(uptr pc, uptr bp, void *context,
+                                        bool request_fast, u32 max_depth);
 
   // UnwindFast/Slow have platform-specific implementations
   void UnwindFast(uptr pc, uptr bp, uptr stack_top, uptr stack_bottom,
                   u32 max_depth);
   void UnwindSlow(uptr pc, u32 max_depth);
   void UnwindSlow(uptr pc, void *context, u32 max_depth);
 
   void PopStackFrames(uptr count);

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

static cl::opt<bool> ClCheckAccessAddress("msan-check-access-address",
        cl::desc("report accesses through a pointer which has poisoned shadow"),
        cl::Hidden, cl::init(true));
 
 static cl::opt<bool> ClEagerChecks(
     "msan-eager-checks",
     cl::desc("check arguments and return values at function call boundaries"),
     cl::Hidden, cl::init(false));
 
+static cl::opt<bool>
+    ClCachedUnwinding("msan-cached-unwinding",
+                      cl::desc("hash active stack traces for fast unwinding"),
+                      cl::Hidden, cl::init(false));
+
 static cl::opt<bool> ClDumpStrictInstructions("msan-dump-strict-instructions",
        cl::desc("print out instructions with default strict semantics"),
        cl::Hidden, cl::init(false));
 
 static cl::opt<int> ClInstrumentationWithCallThreshold(
     "msan-instrumentation-with-call-threshold",
     cl::desc(
         "If the function being instrumented requires more than "
private:
   /// Thread-local shadow storage for in-register va_arg function
   /// parameters (x86_64-specific).
   Value *VAArgOriginTLS;
 
   /// Thread-local shadow storage for va_arg overflow area
   /// (x86_64-specific).
   Value *VAArgOverflowSizeTLS;
 
+  /// Thread-local shadow storage for current trace hash
+  Value *CurrTraceTLS;
+
   /// Are the instrumentation callbacks set up?
   bool CallbacksInitialized = false;
 
   /// The run-time callback to print a warning.
   FunctionCallee WarningFn;
 
   // These arrays are indexed by log2(AccessSize).
   FunctionCallee MaybeWarningFn[kNumberOfAccessSizes];
void MemorySanitizer::createKernelApi(Module &M) {
   // These will be initialized in insertKmsanPrologue().
   RetvalTLS = nullptr;
   RetvalOriginTLS = nullptr;
   ParamTLS = nullptr;
   ParamOriginTLS = nullptr;
   VAArgTLS = nullptr;
   VAArgOriginTLS = nullptr;
   VAArgOverflowSizeTLS = nullptr;
+  CurrTraceTLS = nullptr;
 
   WarningFn = M.getOrInsertFunction("__msan_warning", IRB.getVoidTy(),
                                     IRB.getInt32Ty());
   // Requests the per-task context state (kmsan_context_state*) from the
   // runtime library.
   MsanContextStateTy = StructType::get(
       ArrayType::get(IRB.getInt64Ty(), kParamTLSSize / 8),
       ArrayType::get(IRB.getInt64Ty(), kRetvalTLSSize / 8),
void MemorySanitizer::createUserspaceApi(Module &M) {
 
   VAArgOriginTLS =
       getOrInsertGlobal(M, "__msan_va_arg_origin_tls",
                         ArrayType::get(OriginTy, kParamTLSSize / 4));
 
   VAArgOverflowSizeTLS =
       getOrInsertGlobal(M, "__msan_va_arg_overflow_size_tls", IRB.getInt64Ty());
 
+  CurrTraceTLS = getOrInsertGlobal(M, "__msan_ctrace_tls", IRB.getInt32Ty());
+
   for (size_t AccessSizeIndex = 0; AccessSizeIndex < kNumberOfAccessSizes;
        AccessSizeIndex++) {
     unsigned AccessSize = 1 << AccessSizeIndex;
     std::string FunctionName = "__msan_maybe_warning_" + itostr(AccessSize);
     SmallVector<std::pair<unsigned, Attribute>, 2> MaybeWarningFnAttrs;
     MaybeWarningFnAttrs.push_back(std::make_pair(
         AttributeList::FirstArgIndex, Attribute::get(*C, Attribute::ZExt)));
     MaybeWarningFnAttrs.push_back(std::make_pair(
MemorySanitizerVisitor(Function &F, MemorySanitizer &MS,
     if (MS.CompileKernel)
       ActualFnStart = insertKmsanPrologue(F);
     else
       ActualFnStart = &F.getEntryBlock();
 
     LLVM_DEBUG(if (!InsertChecks) dbgs()
                << "MemorySanitizer is not inserting checks into '"
                << F.getName() << "'\n");
+
+    if (ClCachedUnwinding)
+      ActualFnStart = insertTracePrologue(ActualFnStart);
   }
 
   Value *updateOrigin(Value *V, IRBuilder<> &IRB) {
     if (MS.TrackOrigins <= 1) return V;
     return IRB.CreateCall(MS.MsanChainOriginFn, V);
   }
 
   Value *originToIntptr(IRBuilder<> &IRB, Value *Origin) {
     const DataLayout &DL = F.getParent()->getDataLayout();
BasicBlock *insertKmsanPrologue(Function &F) {
     MS.ParamOriginTLS = IRB.CreateGEP(MS.MsanContextStateTy, ContextState,
                                       {Zero, IRB.getInt32(5)}, "param_origin");
     MS.RetvalOriginTLS =
         IRB.CreateGEP(MS.MsanContextStateTy, ContextState,
                       {Zero, IRB.getInt32(6)}, "retval_origin");
     return ret;
   }
 
+  BasicBlock *insertTracePrologue(BasicBlock *B) {
+    BasicBlock *ret = SplitBlock(B, B->getFirstNonPHI());
+    IRBuilder<> IRB(B->getFirstNonPHI());
+
+    // CurrTrace allows us to approximately hash the entire call trace leading
+    // to this function. We XOR in the prologue to mark the beginning of this
+    // frame, and then XOR in the epilogue (ret instructions) to mark the end.
+    Value *CurrTraceHash = IRB.CreateLoad(IRB.getInt32Ty(), MS.CurrTraceTLS);
+    Value *ReturnAddress = IRB.CreatePtrToInt(
+        IRB.CreateIntrinsic(Intrinsic::returnaddress, {}, {IRB.getInt32(0)}),
+        IRB.getInt32Ty());
+    CurrTraceHash =
+        IRB.CreateIntrinsic(Intrinsic::fshl, {IRB.getInt32Ty()},
+                            {CurrTraceHash, CurrTraceHash, IRB.getInt32(5)});
+    CurrTraceHash = IRB.CreateXor(CurrTraceHash, ReturnAddress);
+    IRB.CreateStore(CurrTraceHash, MS.CurrTraceTLS);
+    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);
 
bool isAMustTailRetVal(Value *RetVal) {
     if (auto *I = dyn_cast<CallInst>(RetVal)) {
       return I->isMustTailCall();
     }
     return false;
   }
 
   void visitReturnInst(ReturnInst &I) {
     IRBuilder<> IRB(&I);
+    if (ClCachedUnwinding) {
+      // XOR out the current stack frame's hash to indicate returning up a
+      // frame.
+      Value *CurrTraceHash = IRB.CreateLoad(IRB.getInt32Ty(), MS.CurrTraceTLS);
+      Value *ReturnAddress = IRB.CreatePtrToInt(
+          IRB.CreateIntrinsic(Intrinsic::returnaddress, {}, {IRB.getInt32(0)}),
+          IRB.getInt32Ty());
+      CurrTraceHash = IRB.CreateXor(CurrTraceHash, ReturnAddress);
+      CurrTraceHash =
+          IRB.CreateIntrinsic(Intrinsic::fshr, {IRB.getInt32Ty()},
+                              {CurrTraceHash, CurrTraceHash, IRB.getInt32(5)});
+      IRB.CreateStore(CurrTraceHash, MS.CurrTraceTLS);
+    }
+
     Value *RetVal = I.getReturnValue();
     if (!RetVal) return;
     // Don't emit the epilogue for musttail call returns.
     if (isAMustTailRetVal(RetVal)) return;
     Value *ShadowPtr = getShadowPtrForRetval(RetVal, IRB);
     bool HasNoUndef =
         F.hasAttribute(AttributeList::ReturnIndex, Attribute::NoUndef);
     bool StoreShadow = !(ClEagerChecks && HasNoUndef);