[asan] Speed up compilation of large C++ stringmaps (tons of allocas) with ASan

[kubamracek]

Summary

This addresses PR30746 https://llvm.org/bugs/show_bug.cgi?id=30746. The ASan pass iterates over entry-block instructions and checks each alloca whether it's in NonInstrumentedStaticAllocaVec, which is apparently slow. This patch gathers the instructions to move during visitAllocaInst.

[aprantl]

lib/Transforms/Instrumentation/AddressSanitizer.cpp

772 ↗

(On Diff #77124)

Nit: clang-format would put the return on the next line.

[eugenis]

LGTM w/ nit

lib/Transforms/Instrumentation/AddressSanitizer.cpp

641 ↗

(On Diff #77124)

This can be just SmallVector.

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

struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   AddressSanitizer &ASan;
   DIBuilder DIB;
   LLVMContext *C;
   Type *IntptrTy;
   Type *IntptrPtrTy;
   ShadowMapping Mapping;
 
   SmallVector<AllocaInst *, 16> AllocaVec;
-  SmallSetVector<AllocaInst *, 16> NonInstrumentedStaticAllocaVec;
+  SmallVector<AllocaInst *, 16> StaticAllocasToMoveUp;
   SmallVector<Instruction *, 8> RetVec;
   unsigned StackAlignment;
 
   Function *AsanStackMallocFunc[kMaxAsanStackMallocSizeClass + 1],
       *AsanStackFreeFunc[kMaxAsanStackMallocSizeClass + 1];
   Function *AsanSetShadowFunc[0x100] = {};
   Function *AsanPoisonStackMemoryFunc, *AsanUnpoisonStackMemoryFunc;
   Function *AsanAllocaPoisonFunc, *AsanAllocasUnpoisonFunc;
struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
   //   addr = tmp + 32 (first 32 bytes are for the left redzone).
   // Additional_size is added to make new memory allocation contain not only
   // requested memory, but also left, partial and right redzones.
   void handleDynamicAllocaCall(AllocaInst *AI);
 
   /// \brief Collect Alloca instructions we want (and can) handle.
   void visitAllocaInst(AllocaInst &AI) {
     if (!ASan.isInterestingAlloca(AI)) {
-      if (AI.isStaticAlloca()) NonInstrumentedStaticAllocaVec.insert(&AI);
+      if (AI.isStaticAlloca()) {
+        // Skip over allocas that are present *before* the first instrumented
+        // alloca, we don't want to move those around.
+        if (AllocaVec.empty())
+          return;
+
+        StaticAllocasToMoveUp.push_back(&AI);
+      }
       return;
     }
 
     StackAlignment = std::max(StackAlignment, AI.getAlignment());
     if (!AI.isStaticAlloca())
       DynamicAllocaVec.push_back(&AI);
     else
       AllocaVec.push_back(&AI);
void FunctionStackPoisoner::processStaticAllocas() {
   IRBuilder<> IRB(InsBefore);
   IRB.SetCurrentDebugLocation(EntryDebugLocation);
 
   // Make sure non-instrumented allocas stay in the entry block. Otherwise,
   // debug info is broken, because only entry-block allocas are treated as
   // regular stack slots.
   auto InsBeforeB = InsBefore->getParent();
   assert(InsBeforeB == &F.getEntryBlock());
-  for (BasicBlock::iterator I(InsBefore); I != InsBeforeB->end(); ++I)
-    if (auto *AI = dyn_cast<AllocaInst>(I))
-      if (NonInstrumentedStaticAllocaVec.count(AI) > 0)
-        AI->moveBefore(InsBefore);
+  for (auto *AI : StaticAllocasToMoveUp)
+    if (AI->getParent() == InsBeforeB)
 
   // If we have a call to llvm.localescape, keep it in the entry block.
   if (LocalEscapeCall) LocalEscapeCall->moveBefore(InsBefore);
 
   SmallVector<ASanStackVariableDescription, 16> SVD;
   SVD.reserve(AllocaVec.size());
   for (AllocaInst *AI : AllocaVec) {
     ASanStackVariableDescription D = {AI->getName().data(),

llvm/trunk/test/Instrumentation/AddressSanitizer/debug_info_noninstrumented_alloca2.ll

+; Make sure we don't break the IR when moving non-instrumented allocas
+
+; RUN: opt < %s -asan -asan-module -S | FileCheck %s
+; RUN: opt < %s -asan -asan-module -asan-instrument-dynamic-allocas -S | FileCheck %s
+
+target datalayout = "e-m:o-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-apple-macosx10.10.0"
+
+define i32 @foo() sanitize_address {
+entry:
+  %non_instrumented1 = alloca i32, align 4
+  %t = load i32, i32* %non_instrumented1, align 4
+  %instrumented = alloca i32, align 4
+  %ptr = ptrtoint i32* %instrumented to i32
+  ret i32 %t
+}
+
+; CHECK: entry:
+; CHECK: %non_instrumented1 = alloca i32, align 4
+; CHECK: load i32, i32* %non_instrumented1
+; CHECK: load i32, i32* @__asan_option_detect_stack_use_after_return