Faster stack-protector for Android/AArch64.

[eugenis]

Summary

Bionic has a defined thread-local location for the stack protector
cookie. Emit a direct load instead of going through __stack_chk_guard.

The current signature of getStackCookieLocation (returning addressspace + offset) is very x86 specific; let the function generate arbitrary IR for the cookie access instead. This is the same way getSafeStackPointerLocation works.

[echristo]

Adding Tim since he's in here as well.

[timshen]

LGTM, glad that AArch64 has Intrinsic::aarch64_thread_pointer to make your job easier. :)

I'm not the owner of SSP, but I'm recently working on it. Let's wait for others comments for a while, then I'm happy to stamp the patch.

For a bit background: I'm currently working on a cleanup on SSP (D17736), which also involves a similar change to yours. No worries, I will rebase to include your change.

[eugenis]

Cool. Would it be better to avoid the use of IRBuilder, and switch the targetlowering hooks to the same signature as in D17736?
Value *..(Module &);

[timshen]

I slightly prefer taking a Module&, since by taking a Module&, getStackCookieLocation doesn't have a chance to insert a concrete instruction to a potentially specified place. That's taking Module& makes the interface narrower. But I don't have strong opinion on this.

[eugenis]

Good point. AArch64 implementation actually needs to insert instructions for the intrinsic call. I guess we keep IRBuilder then, or replace it with Instruction *InsertBefore argument.

[eugenis]

Friendly ping.

[eugenis]

Thanks for the review.
r265481

include/llvm/Target/TargetLowering.h

unsigned getPrefFunctionAlignment() const {
     return PrefFunctionAlignment;
   }
 
   /// Return the preferred loop alignment.
   virtual unsigned getPrefLoopAlignment(MachineLoop *ML = nullptr) const {
     return PrefLoopAlignment;
   }
 
-  /// Return true if the target stores stack protector cookies at a fixed offset
-  /// in some non-standard address space, and populates the address space and
-  /// offset as appropriate.
-  virtual bool getStackCookieLocation(unsigned &/*AddressSpace*/,
-                                      unsigned &/*Offset*/) const {
-    return false;
+  /// If the target has a standard location for the stack protector cookie,
+  /// returns the address of that location. Otherwise, returns nullptr.
+  virtual Value *getStackCookieLocation(IRBuilder<> &IRB) const {
+    return nullptr;
   }
 
   /// If the target has a standard location for the unsafe stack pointer,
   /// returns the address of that location. Otherwise, returns nullptr.
   virtual Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const;
 
   /// Returns true if a cast between SrcAS and DestAS is a noop.
   virtual bool isNoopAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const {

lib/CodeGen/StackProtector.cpp

 ///
 /// Returns true if the platform/triple supports the stackprotectorcreate pseudo
 /// node.
 static bool CreatePrologue(Function *F, Module *M, ReturnInst *RI,
                            const TargetLoweringBase *TLI, const Triple &TT,
                            AllocaInst *&AI, Value *&StackGuardVar) {
   bool SupportsSelectionDAGSP = false;
   PointerType *PtrTy = Type::getInt8PtrTy(RI->getContext());
-  unsigned AddressSpace, Offset;
-  if (TLI->getStackCookieLocation(AddressSpace, Offset)) {
-    Constant *OffsetVal =
-        ConstantInt::get(Type::getInt32Ty(RI->getContext()), Offset);
-
-    StackGuardVar =
-        ConstantExpr::getIntToPtr(OffsetVal, PointerType::get(PtrTy,
-                                                              AddressSpace));
-  } else if (TT.isOSOpenBSD()) {
+  IRBuilder<> B(&F->getEntryBlock().front());
+
+  StackGuardVar = TLI->getStackCookieLocation(B);
+  if (!StackGuardVar) {
+    if (TT.isOSOpenBSD()) {
       StackGuardVar = M->getOrInsertGlobal("__guard_local", PtrTy);
       cast<GlobalValue>(StackGuardVar)
           ->setVisibility(GlobalValue::HiddenVisibility);
     } else {
       SupportsSelectionDAGSP = true;
       StackGuardVar = M->getOrInsertGlobal("__stack_chk_guard", PtrTy);
     }
+  }
 
-  IRBuilder<> B(&F->getEntryBlock().front());
   AI = B.CreateAlloca(PtrTy, nullptr, "StackGuardSlot");
   LoadInst *LI = B.CreateLoad(StackGuardVar, "StackGuard");
   B.CreateCall(Intrinsic::getDeclaration(M, Intrinsic::stackprotector),
                {LI, AI});
 
   return SupportsSelectionDAGSP;
 }
 

lib/Target/AArch64/AArch64ISelLowering.h

public:
   shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const override;
 
   bool shouldExpandAtomicCmpXchgInIR(AtomicCmpXchgInst *AI) const override;
 
   bool useLoadStackGuardNode() const override;
   TargetLoweringBase::LegalizeTypeAction
   getPreferredVectorAction(EVT VT) const override;
 
+  /// If the target has a standard location for the stack protector cookie,
+  /// returns the address of that location. Otherwise, returns nullptr.
+  Value *getStackCookieLocation(IRBuilder<> &IRB) const override;
+
   /// If the target has a standard location for the unsafe stack pointer,
   /// returns the address of that location. Otherwise, returns nullptr.
   Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const override;
 
   /// If a physical register, this returns the register that receives the
   /// exception address on entry to an EH pad.
   unsigned
   getExceptionPointerRegister(const Constant *PersonalityFn) const override {

lib/Target/AArch64/AArch64ISelLowering.cpp

bool AArch64TargetLowering::functionArgumentNeedsConsecutiveRegisters(
   return Ty->isArrayTy();
 }
 
 bool AArch64TargetLowering::shouldNormalizeToSelectSequence(LLVMContext &,
                                                             EVT) const {
   return false;
 }
 
+Value *AArch64TargetLowering::getStackCookieLocation(IRBuilder<> &IRB) const {
+  if (!Subtarget->isTargetAndroid())
+    return TargetLowering::getStackCookieLocation(IRB);
+
+  // Android provides a fixed TLS slot for the stack cookie. See the definition
+  // of TLS_SLOT_STACK_GUARD in
+  // https://android.googlesource.com/platform/bionic/+/master/libc/private/bionic_tls.h
+  const unsigned TlsOffset = 0x28;
+  Module *M = IRB.GetInsertBlock()->getParent()->getParent();
+  Function *ThreadPointerFunc =
+      Intrinsic::getDeclaration(M, Intrinsic::aarch64_thread_pointer);
+  return IRB.CreatePointerCast(
+      IRB.CreateConstGEP1_32(IRB.CreateCall(ThreadPointerFunc), TlsOffset),
+      Type::getInt8PtrTy(IRB.getContext())->getPointerTo(0));
+}
+
 Value *AArch64TargetLowering::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
   if (!Subtarget->isTargetAndroid())
     return TargetLowering::getSafeStackPointerLocation(IRB);
 
   // Android provides a fixed TLS slot for the SafeStack pointer. See the
   // definition of TLS_SLOT_SAFESTACK in
   // https://android.googlesource.com/platform/bionic/+/master/libc/private/bionic_tls.h
   const unsigned TlsOffset = 0x48;

lib/Target/X86/X86ISelLowering.h

public:
 
     virtual bool needsFixedCatchObjects() const override;
 
     /// This method returns a target specific FastISel object,
     /// or null if the target does not support "fast" ISel.
     FastISel *createFastISel(FunctionLoweringInfo &funcInfo,
                              const TargetLibraryInfo *libInfo) const override;
 
-    /// Return true if the target stores stack protector cookies at a fixed
-    /// offset in some non-standard address space, and populates the address
-    /// space and offset as appropriate.
-    bool getStackCookieLocation(unsigned &AddressSpace,
-                                unsigned &Offset) const override;
+    /// If the target has a standard location for the stack protector cookie,
+    /// returns the address of that location. Otherwise, returns nullptr.
+    Value *getStackCookieLocation(IRBuilder<> &IRB) const override;
 
     /// Return true if the target stores SafeStack pointer at a fixed offset in
     /// some non-standard address space, and populates the address space and
     /// offset as appropriate.
     Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const override;
 
     SDValue BuildFILD(SDValue Op, EVT SrcVT, SDValue Chain, SDValue StackSlot,
                       SelectionDAG &DAG) const;

lib/Target/X86/X86ISelLowering.cpp

 }
 
 unsigned X86TargetLowering::getAddressSpace() const {
   if (Subtarget.is64Bit())
     return (getTargetMachine().getCodeModel() == CodeModel::Kernel) ? 256 : 257;
   return 256;
 }
 
-bool X86TargetLowering::getStackCookieLocation(unsigned &AddressSpace,
-                                               unsigned &Offset) const {
+Value *X86TargetLowering::getStackCookieLocation(IRBuilder<> &IRB) const {
   if (!Subtarget.isTargetLinux())
-    return false;
+    return TargetLowering::getStackCookieLocation(IRB);
 
   // %fs:0x28, unless we're using a Kernel code model, in which case it's %gs:
   // %gs:0x14 on i386
-  Offset = (Subtarget.is64Bit()) ? 0x28 : 0x14;
-  AddressSpace = getAddressSpace();
-  return true;
+  unsigned Offset = (Subtarget.is64Bit()) ? 0x28 : 0x14;
+  unsigned AddressSpace = getAddressSpace();
+  return ConstantExpr::getIntToPtr(
+      ConstantInt::get(Type::getInt32Ty(IRB.getContext()), Offset),
+      Type::getInt8PtrTy(IRB.getContext())->getPointerTo(AddressSpace));
 }
 
 Value *X86TargetLowering::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
   if (!Subtarget.isTargetAndroid())
     return TargetLowering::getSafeStackPointerLocation(IRB);
 
   // Android provides a fixed TLS slot for the SafeStack pointer. See the
   // definition of TLS_SLOT_SAFESTACK in

test/CodeGen/AArch64/stack-protector-target.ll

+; Test target-specific stack cookie location.
+; RUN: llc -mtriple=aarch64-linux-android < %s -o - | FileCheck --check-prefix=ANDROID-AARCH64 %s
+
+define void @_Z1fv() sspreq {
+entry:
+  %x = alloca i32, align 4
+  %0 = bitcast i32* %x to i8*
+  call void @_Z7CapturePi(i32* nonnull %x)
+  ret void
+}
+
+declare void @_Z7CapturePi(i32*)
+
+; ANDROID-AARCH64: mrs [[A:.*]], TPIDR_EL0
+; ANDROID-AARCH64: ldr [[B:.*]], {{\[}}[[A]], #40]
+; ANDROID-AARCH64: str [[B]], [sp,
+; ANDROID-AARCH64: ldr [[C:.*]], {{\[}}[[A]], #40]
+; ANDROID-AARCH64: ldr [[D:.*]], [sp,
+; ANDROID-AARCH64: cmp [[C]], [[D]]

test/CodeGen/X86/stack-protector-target.ll

+; Test target-specific stack cookie location.
+; RUN: llc -mtriple=i386-linux < %s -o - | FileCheck --check-prefix=LINUX-I386 %s
+; RUN: llc -mtriple=x86_64-linux < %s -o - | FileCheck --check-prefix=LINUX-X64 %s
+; RUN: llc -mtriple=i386-linux-android < %s -o - | FileCheck --check-prefix=LINUX-I386 %s
+; RUN: llc -mtriple=x86_64-linux-android < %s -o - | FileCheck --check-prefix=LINUX-X64 %s
+
+define void @_Z1fv() sspreq {
+entry:
+  %x = alloca i32, align 4
+  %0 = bitcast i32* %x to i8*
+  call void @_Z7CapturePi(i32* nonnull %x)
+  ret void
+}
+
+declare void @_Z7CapturePi(i32*)
+
+; LINUX-X64: movq %fs:40, %[[B:.*]]
+; LINUX-X64: movq %[[B]], 16(%rsp)
+; LINUX-X64: movq %fs:40, %[[C:.*]]
+; LINUX-X64: cmpq 16(%rsp), %[[C]]
+
+; LINUX-I386: movl %gs:20, %[[B:.*]]
+; LINUX-I386: movl %[[B]], 8(%esp)
+; LINUX-I386: movl %gs:20, %[[C:.*]]
+; LINUX-I386: cmpl 8(%esp), %[[C]]