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

[vitalybuka]

Summary

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

PR27453

[vitalybuka]

typo

[eugenis]

Please extend the test with different combinations of inline/non-inline shadow ranges.

lib/Transforms/Instrumentation/AddressSanitizer.cpp

2037 ↗

(On Diff #68664)

do you mean "break"?
Or even replace this with assert().

[eugenis]

LGTM

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

 // set it to 10000.
 static cl::opt<int> ClMaxInsnsToInstrumentPerBB(
     "asan-max-ins-per-bb", cl::init(10000),
     cl::desc("maximal number of instructions to instrument in any given BB"),
     cl::Hidden);
 // 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));
 static cl::opt<bool> ClUseAfterScope("asan-use-after-scope",
                                      cl::desc("Check stack-use-after-scope"),
                                      cl::Hidden, cl::init(false));
 static cl::opt<bool> ClExperimentalPoisoning(
     "asan-experimental-poisoning",
bool doesDominateAllExits(const Instruction *I) const {
     for (auto Ret : RetVec) {
       if (!ASan.getDominatorTree().dominates(I, Ret)) return false;
     }
     return true;
   }
 
   /// 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);
 
   Value *createAllocaForLayout(IRBuilder<> &IRB, const ASanStackFrameLayout &L,
                                bool Dynamic);
   PHINode *createPHI(IRBuilder<> &IRB, Value *Cond, Value *ValueIfTrue,
                      Instruction *ThenTerm, Value *ValueIfFalse);

 }
 
 // If DoPoison is true function copies ShadowBytes into shadow memory.
 // If DoPoison is false function sets 0s into shadow memory.
 // Function assumes that if ShadowBytes[i] is 0, then corresponding shadow
 // 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,
+  if (Begin >= End)
+    return;
 
   const size_t LargestStoreSizeInBytes =
       std::min<size_t>(sizeof(uint64_t), ASan.LongSize / 8);
 
   const bool IsLittleEndian = F.getParent()->getDataLayout().isLittleEndian();
 
   // Poison given range in shadow using larges store size with out leading and
   // 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;
     }
 
     size_t StoreSizeInBytes = LargestStoreSizeInBytes;
     // Fit store size into the range.
     while (StoreSizeInBytes > End - i)
for (size_t i = Begin; i < End;) {
     Value *Poison = IRB.getIntN(StoreSizeInBytes * 8, Val);
     IRB.CreateStore(Poison,
                     IRB.CreateIntToPtr(Ptr, Poison->getType()->getPointerTo()));
 
     i += StoreSizeInBytes;
   }
 }
 
+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) {
   assert(LocalStackSize <= kMaxStackMallocSize);
   uint64_t MaxSize = kMinStackMallocSize;
   for (int i = 0;; i++, MaxSize *= 2)
     if (LocalStackSize <= MaxSize) return i;
   llvm_unreachable("impossible LocalStackSize");

llvm/trunk/test/Instrumentation/AddressSanitizer/stack-poisoning-experimental.ll

-; 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"
 
 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_