[asan] Minimize code size by using __asan_set_shadow_* for large blocks

Summary:
We can insert function call instead of multiple store operation.
Current default is blocks larger than 64 bytes.
Changes are hidden behind -asan-experimental-poisoning flag.

PR27453

Differential Revision: https://reviews.llvm.org/D23711

llvm-svn: 279383

2 files changed, 93 insertions, 15 deletions

diff --git a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
index dde544d1e89..7bf958eb82c 100644
--- a/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp
@@ -167,6 +167,11 @@ static cl::opt<int> ClMaxInsnsToInstrumentPerBB(
 // This flag may need to be replaced with -f[no]asan-stack.
 static cl::opt<bool> ClStack("asan-stack", cl::desc("Handle stack memory"),
                              cl::Hidden, cl::init(true));
+static cl::opt<uint32_t> ClMaxInlinePoisoningSize(
+    "asan-max-inline-poisoning-size",
+    cl::desc(
+        "Inline shadow poisoning for blocks up to the given size in bytes."),
+    cl::Hidden, cl::init(64));
 static cl::opt<bool> ClUseAfterReturn("asan-use-after-return",
                                       cl::desc("Check stack-use-after-return"),
                                       cl::Hidden, cl::init(true));
@@ -806,6 +811,9 @@ struct FunctionStackPoisoner : public InstVisitor<FunctionStackPoisoner> {
 
   /// Finds alloca where the value comes from.
   AllocaInst *findAllocaForValue(Value *V);
+  void poisonStackFrameInline(ArrayRef<uint8_t> ShadowBytes, size_t Begin,
+                              size_t End, IRBuilder<> &IRB, Value *ShadowBase,
+                              bool DoPoison);
   void poisonStackFrame(ArrayRef<uint8_t> ShadowBytes, IRBuilder<> &IRB,
                         Value *ShadowBase, bool DoPoison);
   void poisonAlloca(Value *V, uint64_t Size, IRBuilder<> &IRB, bool DoPoison);
@@ -1956,10 +1964,11 @@ void FunctionStackPoisoner::initializeCallbacks(Module &M) {
 // memory is constant for duration of the function and it contains 0s. So we
 // will try to minimize writes into corresponding addresses of the real shadow
 // memory.
-void FunctionStackPoisoner::poisonStackFrame(ArrayRef<uint8_t> ShadowBytes,
-                                             IRBuilder<> &IRB,
-                                             Value *ShadowBase, bool DoPoison) {
-  const size_t End = ShadowBytes.size();
+void FunctionStackPoisoner::poisonStackFrameInline(
+    ArrayRef<uint8_t> ShadowBytes, size_t Begin, size_t End, IRBuilder<> &IRB,
+    Value *ShadowBase, bool DoPoison) {
+  if (Begin >= End)
+    return;
 
   const size_t LargestStoreSizeInBytes =
       std::min<size_t>(sizeof(uint64_t), ASan.LongSize / 8);
@@ -1970,7 +1979,7 @@ void FunctionStackPoisoner::poisonStackFrame(ArrayRef<uint8_t> ShadowBytes,
   // trailing zeros. Zeros never change, so they need neither poisoning nor
   // up-poisoning, but we don't mind if some of them get into a middle of a
   // store.
-  for (size_t i = 0; i < End;) {
+  for (size_t i = Begin; i < End;) {
     if (!ShadowBytes[i]) {
       ++i;
       continue;
@@ -2009,6 +2018,40 @@ void FunctionStackPoisoner::poisonStackFrame(ArrayRef<uint8_t> ShadowBytes,
   }
 }
 
+void FunctionStackPoisoner::poisonStackFrame(ArrayRef<uint8_t> ShadowBytes,
+                                             IRBuilder<> &IRB,
+                                             Value *ShadowBase, bool DoPoison) {
+  auto ValueToWrite = [&](size_t i) {
+    if (DoPoison)
+      return ShadowBytes[i];
+    return static_cast<uint8_t>(0);
+  };
+
+  const size_t End = ShadowBytes.size();
+  size_t Done = 0;
+  for (size_t i = 0, j = 1; i < End; i = j++) {
+    if (!ShadowBytes[i])
+      continue;
+    uint8_t Val = ValueToWrite(i);
+    if (!AsanSetShadowFunc[Val])
+      continue;
+
+    // Skip same values.
+    for (; j < End && ShadowBytes[j] && Val == ValueToWrite(j); ++j) {
+    }
+
+    if (j - i >= ClMaxInlinePoisoningSize) {
+      poisonStackFrameInline(ShadowBytes, Done, i, IRB, ShadowBase, DoPoison);
+      IRB.CreateCall(AsanSetShadowFunc[Val],
+                     {IRB.CreateAdd(ShadowBase, ConstantInt::get(IntptrTy, i)),
+                      ConstantInt::get(IntptrTy, j - i)});
+      Done = j;
+    }
+  }
+
+  poisonStackFrameInline(ShadowBytes, Done, End, IRB, ShadowBase, DoPoison);
+}
+
 // Fake stack allocator (asan_fake_stack.h) has 11 size classes
 // for every power of 2 from kMinStackMallocSize to kMaxAsanStackMallocSizeClass
 static int StackMallocSizeClass(uint64_t LocalStackSize) {
diff --git a/llvm/test/Instrumentation/AddressSanitizer/stack-poisoning-experimental.ll b/llvm/test/Instrumentation/AddressSanitizer/stack-poisoning-experimental.ll
index bf3cf9cecb6..a9dd83a2183 100644
--- a/llvm/test/Instrumentation/AddressSanitizer/stack-poisoning-experimental.ll
+++ b/llvm/test/Instrumentation/AddressSanitizer/stack-poisoning-experimental.ll
@@ -1,5 +1,5 @@
-; RUN: opt < %s -asan -asan-module -asan-experimental-poisoning -S | FileCheck %s
-; RUN: opt < %s -asan -asan-module -S | FileCheck --check-prefix=CHECK-OFF %s
+; RUN: opt < %s -asan -asan-module -asan-experimental-poisoning -S | FileCheck --check-prefixes=CHECK-ON,CHECK %s
+; RUN: opt < %s -asan -asan-module -S | FileCheck --check-prefixes=CHECK-OFF,CHECK %s
 
 target datalayout = "e-i64:64-f80:128-s:64-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
@@ -8,17 +8,52 @@ declare void @Foo(i8*)
 
 define void @Bar() uwtable sanitize_address {
 entry:
-  %x = alloca [20 x i8], align 16
-  %arraydecay = getelementptr inbounds [20 x i8], [20 x i8]* %x, i64 0, i64 0
+  ; CHECK: store i32 -235802127
+  ; CHECK: store i64 -868082074056920318
+  ; CHECK: store i64 -868082074056920077
+  ; CHECK: store i16 -3085
+  ; CHECK: store i8 -13
+  ; CHECK-LABEL: call void @Foo
+  
+  ; CHECK-LABEL: <label>
+  ; CHECK-ON-NOT: store i64
+  ; CHECK-ON: call void @__asan_set_shadow_f5(i64 %{{[0-9]+}}, i64 128)
+ 
+  ; CHECK-OFF-NOT: call void @__asan_set_shadow_f5
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; CHECK-OFF: store i64 -723401728380766731
+  ; And more...
+
+  ; CHECK-LABEL: <label>
+  ; CHECK-NOT: call void @__asan_set_shadow_00
+  ; CHECK: store i32 0
+  ; CHECK: store i64 0
+  ; CHECK: store i64 0
+  ; CHECK: store i16 0
+  ; CHECK: store i8 0
+
+  ; CHECK-LABEL: <label>
+  ; CHECK: ret void
+
+  %x = alloca [650 x i8], align 16
+  %arraydecay = getelementptr inbounds [650 x i8], [650 x i8]* %x, i64 0, i64 0
   call void @Foo(i8* %arraydecay)
   ret void
 }
 
-; CHECK: declare void @__asan_set_shadow_00(i64, i64)
-; CHECK: declare void @__asan_set_shadow_f1(i64, i64)
-; CHECK: declare void @__asan_set_shadow_f2(i64, i64)
-; CHECK: declare void @__asan_set_shadow_f3(i64, i64)
-; CHECK: declare void @__asan_set_shadow_f5(i64, i64)
-; CHECK: declare void @__asan_set_shadow_f8(i64, i64)
+; CHECK-ON: declare void @__asan_set_shadow_00(i64, i64)
+; CHECK-ON: declare void @__asan_set_shadow_f1(i64, i64)
+; CHECK-ON: declare void @__asan_set_shadow_f2(i64, i64)
+; CHECK-ON: declare void @__asan_set_shadow_f3(i64, i64)
+; CHECK-ON: declare void @__asan_set_shadow_f5(i64, i64)
+; CHECK-ON: declare void @__asan_set_shadow_f8(i64, i64)
 
 ; CHECK-OFF-NOT: declare void @__asan_set_shadow_