[asan] Prevent rematerialization of &__asan_shadow.

[eugenis]

Summary

In the mode when ASan shadow base is computed as the address of an
external global (__asan_shadow, currently on android/arm32 only),
regalloc prefers to rematerialize this value to save register spills.
Even in -Os. On arm32 it is rather expensive (2 loads + 1 constant
pool entry).

This changes adds an inline asm in the function prologue to suppress
this behavior. It reduces AsanTest binary size by 7%.

[pcc]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

216

Does this need to be a flag? Could we instead do this only on arm32?

[eugenis]

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

216

Flag is useful for benchmarking and also as a way to disable this feature if something goes wrong. We already do this only on android/arm32.

[pcc]

LGTM

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

216

Please at least remember to remove this flag once you know that it works well.

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

static cl::opt<bool> ClForceDynamicShadow(
     "asan-force-dynamic-shadow",
     cl::desc("Load shadow address into a local variable for each function"),
     cl::Hidden, cl::init(false));
 
 static cl::opt<bool>
     ClWithIfunc("asan-with-ifunc",
                 cl::desc("Access dynamic shadow through an ifunc global on "
                          "platforms that support this"),
-                cl::Hidden, cl::init(false));
+                cl::Hidden, cl::init(true));
+
+static cl::opt<bool> ClWithIfuncSuppressRemat(
+    "asan-with-ifunc-suppress-remat",
+    cl::desc("Suppress rematerialization of dynamic shadow address by passing "
+             "it through inline asm in prologue."),
+    cl::Hidden, cl::init(true));
 
 // This flag limits the number of instructions to be instrumented
 // in any given BB. Normally, this should be set to unlimited (INT_MAX),
 // but due to http://llvm.org/bugs/show_bug.cgi?id=12652 we temporary
 // 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"),
if (AI->isStaticAlloca())
       StaticAllocaPoisonCallVec.push_back(APC);
     else if (ClInstrumentDynamicAllocas)
       DynamicAllocaPoisonCallVec.push_back(APC);
   }
 
   void visitCallSite(CallSite CS) {
     Instruction *I = CS.getInstruction();
     if (CallInst *CI = dyn_cast<CallInst>(I)) {
-      HasNonEmptyInlineAsm |=
-          CI->isInlineAsm() && !CI->isIdenticalTo(EmptyInlineAsm.get());
+      HasNonEmptyInlineAsm |= CI->isInlineAsm() &&
+                              !CI->isIdenticalTo(EmptyInlineAsm.get()) &&
+                              I != ASan.LocalDynamicShadow;
       HasReturnsTwiceCall |= CI->canReturnTwice();
     }
   }
 
   // ---------------------- Helpers.
   void initializeCallbacks(Module &M);
 
   bool doesDominateAllExits(const Instruction *I) const {

 }
 
 Value *AddressSanitizer::memToShadow(Value *Shadow, IRBuilder<> &IRB) {
   // Shadow >> scale
   Shadow = IRB.CreateLShr(Shadow, Mapping.Scale);
   if (Mapping.Offset == 0) return Shadow;
   // (Shadow >> scale) | offset
   Value *ShadowBase;
-  if (Mapping.InGlobal)
-    return IRB.CreatePtrToInt(
-        IRB.CreateGEP(AsanShadowGlobal,
-                      {ConstantInt::get(IntptrTy, 0), Shadow}),
-        IntptrTy);
   if (LocalDynamicShadow)
     ShadowBase = LocalDynamicShadow;
   else
     ShadowBase = ConstantInt::get(IntptrTy, Mapping.Offset);
   if (Mapping.OrShadowOffset)
     return IRB.CreateOr(Shadow, ShadowBase);
   else
     return IRB.CreateAdd(Shadow, ShadowBase);
if (F.getName().find(" load]") != std::string::npos) {
     IRB.CreateCall(AsanInitFunction, {});
     return true;
   }
   return false;
 }
 
 void AddressSanitizer::maybeInsertDynamicShadowAtFunctionEntry(Function &F) {
   // Generate code only when dynamic addressing is needed.
-  if (Mapping.Offset != kDynamicShadowSentinel || Mapping.InGlobal)
+  if (Mapping.Offset != kDynamicShadowSentinel)
     return;
 
   IRBuilder<> IRB(&F.front().front());
+  if (Mapping.InGlobal) {
+    if (ClWithIfuncSuppressRemat) {
+      // An empty inline asm with input reg == output reg.
+      // An opaque pointer-to-int cast, basically.
+      InlineAsm *Asm = InlineAsm::get(
+          FunctionType::get(IntptrTy, {AsanShadowGlobal->getType()}, false),
+          StringRef(""), StringRef("=r,0"),
+          /*hasSideEffects=*/false);
+      LocalDynamicShadow =
+          IRB.CreateCall(Asm, {AsanShadowGlobal}, ".asan.shadow");
+    } else {
+      LocalDynamicShadow =
+          IRB.CreatePointerCast(AsanShadowGlobal, IntptrTy, ".asan.shadow");
+    }
+  } else {
     Value *GlobalDynamicAddress = F.getParent()->getOrInsertGlobal(
         kAsanShadowMemoryDynamicAddress, IntptrTy);
     LocalDynamicShadow = IRB.CreateLoad(GlobalDynamicAddress);
   }
+}
 
 void AddressSanitizer::markEscapedLocalAllocas(Function &F) {
   // Find the one possible call to llvm.localescape and pre-mark allocas passed
   // to it as uninteresting. This assumes we haven't started processing allocas
   // yet. This check is done up front because iterating the use list in
   // isInterestingAlloca would be algorithmically slower.
   assert(ProcessedAllocas.empty() && "must process localescape before allocas");
 

llvm/test/Instrumentation/AddressSanitizer/with-ifunc.ll

-; Test -asan-force-dynamic-shadow flag.
+; Test -asan-with-ifunc flag.
 ;
-; RUN: opt -asan -asan-module -S -asan-with-ifunc=1 < %s | FileCheck %s --check-prefixes=CHECK,CHECK-IFUNC
-; RUN: opt -asan -asan-module -S -asan-with-ifunc=0 < %s | FileCheck %s --check-prefixes=CHECK,CHECK-NOIFUNC
+; RUN: opt -asan -asan-module -S -asan-with-ifunc=0 < %s | \
+; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-NOIFUNC
+; RUN: opt -asan -asan-module -S -asan-with-ifunc=1 -asan-with-ifunc-suppress-remat=0 < %s | \
+; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-IFUNC
+; RUN: opt -asan -asan-module -S -asan-with-ifunc=1 -asan-with-ifunc-suppress-remat=1 < %s | \
+; RUN:     FileCheck %s --check-prefixes=CHECK,CHECK-IFUNC-NOREMAT
 
 target datalayout = "e-m:e-p:32:32-i64:64-v128:64:128-a:0:32-n32-S64"
 target triple = "armv7--linux-android"
 
 ; CHECK-IFUNC: @__asan_shadow = external global [0 x i8]
 ; CHECK-NOIFUNC: @__asan_shadow_memory_dynamic_address = external global i32
 
 define i32 @test_load(i32* %a) sanitize_address {
 ; First instrumentation in the function must be to load the dynamic shadow
 ; address into a local variable.
 ; CHECK-LABEL: @test_load
 ; CHECK: entry:
 
 ; CHECK-IFUNC-NEXT: %[[A:[^ ]*]] = ptrtoint i32* %a to i32
 ; CHECK-IFUNC-NEXT: %[[B:[^ ]*]] = lshr i32 %[[A]], 3
-; CHECK-IFUNC-NEXT: %[[C:[^ ]*]] = getelementptr [0 x i8], [0 x i8]* @__asan_shadow, i32 0, i32 %[[B]]
+; CHECK-IFUNC-NEXT: %[[C:[^ ]*]] = add i32 %[[B]], ptrtoint ([0 x i8]* @__asan_shadow to i32)
+
+; CHECK-IFUNC-NOREMAT-NEXT: %[[S:[^ ]*]] = call i32 asm "", "=r,0"([0 x i8]* @__asan_shadow)
+; CHECK-IFUNC-NOREMAT-NEXT: %[[A:[^ ]*]] = ptrtoint i32* %a to i32
+; CHECK-IFUNC-NOREMAT-NEXT: %[[B:[^ ]*]] = lshr i32 %[[A]], 3
+; CHECK-IFUNC-NOREMAT-NEXT: %[[C:[^ ]*]] = add i32 %[[B]], %[[S]]
 
 ; CHECK-NOIFUNC-NEXT: %[[SHADOW:[^ ]*]] = load i32, i32* @__asan_shadow_memory_dynamic_address
 ; CHECK-NOIFUNC-NEXT: %[[A:[^ ]*]] = ptrtoint i32* %a to i32
 ; CHECK-NOIFUNC-NEXT: %[[B:[^ ]*]] = lshr i32 %[[A]], 3
 ; CHECK-NOIFUNC-NEXT: %[[C:[^ ]*]] = add i32 %[[B]], %[[SHADOW]]
 
 entry:
   %x = load i32, i32* %a, align 4
   ret i32 %x
 }