[safestack] Fast access to the unsafe stack pointer on AArch64/Android.

[eugenis]

Summary

Android libc provides a fixed TLS slot for the unsafe stack pointer,
and this change implements direct access to that slot on AArch64 via
__builtin_thread_pointer() + offset.

This change also moves more code into TargetLowering and its
target-specific subclasses to get rid of target-specific codegen
in SafeStackPass.

This change does not touch the ARM backend because it lowers
builting_thread_pointer as aeabi_read_tp, which is not available
on Android.

[rengolin]

Hi Evgeniy,

This looks like a very Android-specific patch on the AArch64 lowering mechanism. Since no one else is using that, why don't you make that function specific to the pass?

cheers,
--renato

[eugenis]

Moving all of this into SafeStackPass would work, and would be much simpler, but it looks like llvm generally tries to avoid target-specific code in IR passes. I'm not really opposed to that, just want to see what people think.

Another option would be to implement x86.thread.pointer intrinsic (tied to __builtin_thread_pointer) and use that in safe stack pass. It would have to include an optimization that wraps x86.thread.pointer + offset into a single %gs:offset load (straightforward codegen would do *(%gs:0)+offset).

[rengolin]

Yeah, I'm not convinced either. But I don't like that android function name right in the middle of the TargetLowering.h. TBH, I don't like it either in the pass, but I assumed it was an android-specific pass.

If that's not true, than I'd think that an intrinsic would make more sense.

[echristo]

Hi Evgeniy,

A few inline comments here. I think if you fix those it'll also fix some of the complaints that Renato had - I'll leave it up to him to confirm. We've got a bit of sausage like this in the middle end, but it seems like we should avoid it if possible and take a slight code duplication hit to do so.

-eric

include/llvm/Target/TargetLowering.h

1951

Should probably be folded in.

lib/CodeGen/TargetLoweringBase.cpp

1665–1666

This comment should go with the Android bits below and a generic comment above the function. Might be better to just duplicate the code in every target rather than have a generic one only implemented for for android (i.e. if there were a generic function in compiler-rt it might make more sense).

lib/Target/AArch64/AArch64ISelLowering.cpp

9891–9893

Go ahead and put this over the android bits.

9898

Non-standard naming.

[eugenis]

It is Android-specific at the moment, but really it is intended for any platform that provides system-wide safestack support in libc. For example, there are some patches for FreeBSD (I don't know what their status is and if they are intended to be merged upstream or not): https://github.com/cpi-llvm/freebsd/commit/a542e1d0c55e71dce888736032deedb30c1a6c92#diff-f8caf191619c46a5ffd4f7f010fe97eaR53

[rengolin]

include/llvm/Target/TargetLowering.h

1951

I was thinking about this one, too. It's only ever used once.

lib/CodeGen/TargetLoweringBase.cpp

1665–1666

Especially when no one it using it...

[eugenis]

lib/CodeGen/TargetLoweringBase.cpp

1665–1666

Well there is a generic function in compiler-rt. Or, rather, a thread-local variable. See the remaining part of SafeStack::getOrCreateUnsafeStackPtr which could be moved here as well, and then this function would apply to all imaginable targets.

[echristo]

lib/CodeGen/TargetLoweringBase.cpp

1665–1666

SGTM. Let's see what that looks like.

It does give the assumption that anyone using safe stack is using compiler-rt, but...

[eugenis]

lib/CodeGen/TargetLoweringBase.cpp

1665–1666

If it's not a known platform (android) and not using compiler-rt, there is nothing we can do anyway.

lib/Target/AArch64/AArch64ISelLowering.cpp

9898

Should it be just "TlsOffset", without the "k" prefix?
CodingStandards does not say anything special about constants.

[echristo]

lib/Target/AArch64/AArch64ISelLowering.cpp

9898

Right. Constants aren't special in llvm naming here.

[eugenis]

lib/CodeGen/TargetLoweringBase.cpp

1665–1666

Done, PTAL

[echristo]

lib/CodeGen/TargetLoweringBase.cpp

1677

Naming, and you can just fold it in as well.

1686–1690

Pretty much all the arguments are fairly self explanatory, probably don't need the comments.

1693

No braces.

1697

Ditto.

lib/Target/AArch64/AArch64Subtarget.h

126–128 ↗

(On Diff #36813)

Given how pervasive this is now, it'll probably be worth splitting this out into a separate commit and doing a look for Triple.isAndroid() and updating all of the callers in llvm and clang.

[eugenis]

lib/CodeGen/TargetLoweringBase.cpp

1677

It's used in 4 places. Fixed the naming.

lib/Target/AArch64/AArch64Subtarget.h

126–128 ↗

(On Diff #36813)

Done and rebased.

[eugenis]

ping

[rengolin]

Right, I'm ok with this. Eric?

[echristo]

Sure. I'd still probably fold in the function name because I don't think local is getting us anything, but styles are different :)

-eric

[eugenis]

I just think that repeating a rather long string constant four times in the source code is typo-prone and less readable.

[rengolin]

In D13455#268231, @eugenis wrote:

I just think that repeating a rather long string constant four times in the source code is typo-prone and less readable.

I agree.

LGTM, too. Thanks!

[eugenis]

At least the name is checked by the compiler.
Thanks for the review! Landed as r250456.

[echristo]

Good point! Thanks for the work.

include/llvm/Target/TargetLowering.h

public:
   /// 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;
   }
 
-  /// 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.
-  virtual bool getSafeStackPointerLocation(unsigned & /*AddressSpace*/,
-                                           unsigned & /*Offset*/) const {
-    return false;
-  }
+  /// 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 {
     return false;
   }
 
   /// Return true if the pointer arguments to CI should be aligned by aligning
   /// the object whose address is being passed. If so then MinSize is set to the
private:
 
   /// The ISD::CondCode that should be used to test the result of each of the
   /// comparison libcall against zero.
   ISD::CondCode CmpLibcallCCs[RTLIB::UNKNOWN_LIBCALL];
 
   /// Stores the CallingConv that should be used for each libcall.
   CallingConv::ID LibcallCallingConvs[RTLIB::UNKNOWN_LIBCALL];
 
 protected:
   /// Return true if the extension represented by \p I is free.
   /// \pre \p I is a sign, zero, or fp extension and
   ///      is[Z|FP]ExtFree of the related types is not true.
   virtual bool isExtFreeImpl(const Instruction *I) const { return false; }
 
   /// \brief Specify maximum number of store instructions per memset call.
   ///
   /// When lowering \@llvm.memset this field specifies the maximum number of

lib/CodeGen/TargetLoweringBase.cpp

while (true) {
     if (LK.first == TypeSplitVector || LK.first == TypeExpandInteger)
       Cost *= 2;
 
     // Keep legalizing the type.
     MTy = LK.second;
   }
 }
 
+Value *TargetLoweringBase::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
+  Module *M = IRB.GetInsertBlock()->getParent()->getParent();
+  Type *StackPtrTy = Type::getInt8PtrTy(M->getContext());
+  if (TM.getTargetTriple().getEnvironment() == llvm::Triple::Android) {
+    // Android provides a libc function to retrieve the address of the current
+    // thread's unsafe stack pointer.
+    Value *Fn = M->getOrInsertFunction("__safestack_pointer_address",
+                                       StackPtrTy->getPointerTo(0), nullptr);
+    return IRB.CreateCall(Fn);
+  } else {
+    // Otherwise, assume the target links with compiler-rt, which provides a
+    // thread-local variable with a magic name.
+    const char *UnsafeStackPtrVar = "__safestack_unsafe_stack_ptr";
+    auto UnsafeStackPtr =
+        dyn_cast_or_null<GlobalVariable>(M->getNamedValue(UnsafeStackPtrVar));
+
+    if (!UnsafeStackPtr) {
+      // The global variable is not defined yet, define it ourselves.
+      // We use the initial-exec TLS model because we do not support the
+      // variable living anywhere other than in the main executable.
+      UnsafeStackPtr = new GlobalVariable(
+          *M, StackPtrTy, false, GlobalValue::ExternalLinkage, 0,
+          UnsafeStackPtrVar, nullptr, GlobalValue::InitialExecTLSModel);
+    } else {
+      // The variable exists, check its type and attributes.
+      if (UnsafeStackPtr->getValueType() != StackPtrTy)
+        report_fatal_error(Twine(UnsafeStackPtrVar) + " must have void* type");
+      if (!UnsafeStackPtr->isThreadLocal())
+        report_fatal_error(Twine(UnsafeStackPtrVar) + " must be thread-local");
+    }
+    return UnsafeStackPtr;
+  }
+}
+
 //===----------------------------------------------------------------------===//
 //  Loop Strength Reduction hooks
 //===----------------------------------------------------------------------===//
 
 /// isLegalAddressingMode - Return true if the addressing mode represented
 /// by AM is legal for this target, for a load/store of the specified type.
 bool TargetLoweringBase::isLegalAddressingMode(const DataLayout &DL,
                                                const AddrMode &AM, Type *Ty,

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 unsafe stack pointer,
+  /// returns the address of that location. Otherwise, returns nullptr.
+  Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const override;
+
 private:
   bool isExtFreeImpl(const Instruction *Ext) const override;
 
   /// Subtarget - Keep a pointer to the AArch64Subtarget around so that we can
   /// make the right decision when generating code for different targets.
   const AArch64Subtarget *Subtarget;
 
   void addTypeForNEON(EVT VT, EVT PromotedBitwiseVT);

lib/Target/AArch64/AArch64ISelLowering.cpp

bool AArch64TargetLowering::functionArgumentNeedsConsecutiveRegisters(
     Type *Ty, CallingConv::ID CallConv, bool isVarArg) const {
   return Ty->isArrayTy();
 }
 
 bool AArch64TargetLowering::shouldNormalizeToSelectSequence(LLVMContext &,
                                                             EVT) const {
   return false;
 }
+
+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;
+  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));
+}

lib/Target/X86/X86ISelLowering.h

public:
     /// offset in some non-standard address space, and populates the address
     /// space and offset as appropriate.
     bool getStackCookieLocation(unsigned &AddressSpace,
                                 unsigned &Offset) 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.
-    bool getSafeStackPointerLocation(unsigned &AddressSpace,
-                                     unsigned &Offset) const override;
+    Value *getSafeStackPointerLocation(IRBuilder<> &IRB) const override;
 
     SDValue BuildFILD(SDValue Op, EVT SrcVT, SDValue Chain, SDValue StackSlot,
                       SelectionDAG &DAG) const;
 
     bool isNoopAddrSpaceCast(unsigned SrcAS, unsigned DestAS) const override;
 
     bool useLoadStackGuardNode() const override;
     /// \brief Customize the preferred legalization strategy for certain types.

lib/Target/X86/X86ISelLowering.cpp

bool X86TargetLowering::getStackCookieLocation(unsigned &AddressSpace,
   } else {
     // %gs:0x14 on i386
     Offset = 0x14;
     AddressSpace = 256;
   }
   return true;
 }
 
-/// 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
-bool X86TargetLowering::getSafeStackPointerLocation(unsigned &AddressSpace,
-                                                    unsigned &Offset) const {
+Value *X86TargetLowering::getSafeStackPointerLocation(IRBuilder<> &IRB) const {
   if (!Subtarget->isTargetAndroid())
-    return false;
+    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
+  unsigned AddressSpace, Offset;
   if (Subtarget->is64Bit()) {
     // %fs:0x48, unless we're using a Kernel code model, in which case it's %gs:
     Offset = 0x48;
     if (getTargetMachine().getCodeModel() == CodeModel::Kernel)
       AddressSpace = 256;
     else
       AddressSpace = 257;
   } else {
     // %gs:0x24 on i386
     Offset = 0x24;
     AddressSpace = 256;
   }
-  return true;
+
+  return ConstantExpr::getIntToPtr(
+      ConstantInt::get(Type::getInt32Ty(IRB.getContext()), Offset),
+      Type::getInt8PtrTy(IRB.getContext())->getPointerTo(AddressSpace));
 }
 
 bool X86TargetLowering::isNoopAddrSpaceCast(unsigned SrcAS,
                                             unsigned DestAS) const {
   assert(SrcAS != DestAS && "Expected different address spaces!");
 
   return SrcAS < 256 && DestAS < 256;
 }

lib/Transforms/Instrumentation/SafeStack.cpp

 #include "llvm/Target/TargetSubtargetInfo.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include "llvm/Transforms/Utils/ModuleUtils.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "safestack"
 
-static const char *const kUnsafeStackPtrVar = "__safestack_unsafe_stack_ptr";
-static const char *const kUnsafeStackPtrAddrFn = "__safestack_pointer_address";
-
 namespace llvm {
 
 STATISTIC(NumFunctions, "Total number of functions");
 STATISTIC(NumUnsafeStackFunctions, "Number of functions with unsafe stack");
 STATISTIC(NumUnsafeStackRestorePointsFunctions,
           "Number of functions that use setjmp or exceptions");
 
 STATISTIC(NumAllocas, "Total number of allocas");
class SafeStack : public FunctionPass {
   /// Unsafe stack alignment. Each stack frame must ensure that the stack is
   /// aligned to this value. We need to re-align the unsafe stack if the
   /// alignment of any object on the stack exceeds this value.
   ///
   /// 16 seems like a reasonable upper bound on the alignment of objects that we
   /// might expect to appear on the stack on most common targets.
   enum { StackAlignment = 16 };
 
-  /// \brief Build a constant representing a pointer to the unsafe stack
-  /// pointer.
-  Value *getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F);
-
   /// \brief Find all static allocas, dynamic allocas, return instructions and
   /// stack restore points (exception unwind blocks and setjmp calls) in the
   /// given function and append them to the respective vectors.
   void findInsts(Function &F, SmallVectorImpl<AllocaInst *> &StaticAllocas,
                  SmallVectorImpl<AllocaInst *> &DynamicAllocas,
                  SmallVectorImpl<ReturnInst *> &Returns,
                  SmallVectorImpl<Instruction *> &StackRestorePoints);
 
bool doInitialization(Module &M) override {
     Int8Ty = Type::getInt8Ty(M.getContext());
 
     return false;
   }
 
   bool runOnFunction(Function &F) override;
 }; // class SafeStack
 
-Value *SafeStack::getOrCreateUnsafeStackPtr(IRBuilder<> &IRB, Function &F) {
-  Module &M = *F.getParent();
-  Triple TargetTriple(M.getTargetTriple());
-
-  unsigned Offset;
-  unsigned AddressSpace;
-  // Check if the target keeps the unsafe stack pointer at a fixed offset.
-  if (TLI && TLI->getSafeStackPointerLocation(AddressSpace, Offset)) {
-    Constant *OffsetVal =
-        ConstantInt::get(Type::getInt32Ty(F.getContext()), Offset);
-    return ConstantExpr::getIntToPtr(OffsetVal,
-                                     StackPtrTy->getPointerTo(AddressSpace));
-  }
-
-  // Android provides a libc function that returns the stack pointer address.
-  if (TargetTriple.isAndroid()) {
-    Value *Fn = M.getOrInsertFunction(kUnsafeStackPtrAddrFn,
-                                      StackPtrTy->getPointerTo(0), nullptr);
-    return IRB.CreateCall(Fn);
-  } else {
-    // Otherwise, declare a thread-local variable with a magic name.
-    auto UnsafeStackPtr =
-        dyn_cast_or_null<GlobalVariable>(M.getNamedValue(kUnsafeStackPtrVar));
-
-    if (!UnsafeStackPtr) {
-      // The global variable is not defined yet, define it ourselves.
-      // We use the initial-exec TLS model because we do not support the
-      // variable living anywhere other than in the main executable.
-      UnsafeStackPtr = new GlobalVariable(
-          /*Module=*/M, /*Type=*/StackPtrTy,
-          /*isConstant=*/false, /*Linkage=*/GlobalValue::ExternalLinkage,
-          /*Initializer=*/nullptr, /*Name=*/kUnsafeStackPtrVar,
-          /*InsertBefore=*/nullptr,
-          /*ThreadLocalMode=*/GlobalValue::InitialExecTLSModel);
-    } else {
-      // The variable exists, check its type and attributes.
-      if (UnsafeStackPtr->getValueType() != StackPtrTy) {
-        report_fatal_error(Twine(kUnsafeStackPtrVar) + " must have void* type");
-      }
-
-      if (!UnsafeStackPtr->isThreadLocal()) {
-        report_fatal_error(Twine(kUnsafeStackPtrVar) + " must be thread-local");
-      }
-    }
-    return UnsafeStackPtr;
-  }
-}
-
 void SafeStack::findInsts(Function &F,
                           SmallVectorImpl<AllocaInst *> &StaticAllocas,
                           SmallVectorImpl<AllocaInst *> &DynamicAllocas,
                           SmallVectorImpl<ReturnInst *> &Returns,
                           SmallVectorImpl<Instruction *> &StackRestorePoints) {
   for (Instruction &I : instructions(&F)) {
     if (auto AI = dyn_cast<AllocaInst>(&I)) {
       ++NumAllocas;
bool SafeStack::runOnFunction(Function &F) {
 
   if (!StaticAllocas.empty() || !DynamicAllocas.empty())
     ++NumUnsafeStackFunctions; // This function has the unsafe stack.
 
   if (!StackRestorePoints.empty())
     ++NumUnsafeStackRestorePointsFunctions;
 
   IRBuilder<> IRB(F.begin()->getFirstInsertionPt());
-  UnsafeStackPtr = getOrCreateUnsafeStackPtr(IRB, F);
+  UnsafeStackPtr = TLI->getSafeStackPointerLocation(IRB);
 
   // The top of the unsafe stack after all unsafe static allocas are allocated.
   Value *StaticTop = moveStaticAllocasToUnsafeStack(IRB, F, StaticAllocas, Returns);
 
   // Safe stack object that stores the current unsafe stack top. It is updated
   // as unsafe dynamic (non-constant-sized) allocas are allocated and freed.
   // This is only needed if we need to restore stack pointer after longjmp
   // or exceptions, and we have dynamic allocations.

test/Transforms/SafeStack/AArch64/abi.ll

 ; RUN: opt -safe-stack -S -mtriple=aarch64-linux-android < %s -o - | FileCheck %s
 
 
 define void @foo() nounwind uwtable safestack {
 entry:
-; CHECK: %[[SPA:.*]] = call i8** @__safestack_pointer_address()
+; CHECK: %[[TP:.*]] = call i8* @llvm.aarch64.thread.pointer()
+; CHECK: %[[SPA0:.*]] = getelementptr i8, i8* %[[TP]], i32 72
+; CHECK: %[[SPA:.*]] = bitcast i8* %[[SPA0]] to i8**
 ; CHECK: %[[USP:.*]] = load i8*, i8** %[[SPA]]
 ; CHECK: %[[USST:.*]] = getelementptr i8, i8* %[[USP]], i32 -16
 ; CHECK: store i8* %[[USST]], i8** %[[SPA]]
 
   %a = alloca i8, align 8
   call void @Capture(i8* %a)
 
 ; CHECK: store i8* %[[USP]], i8** %[[SPA]]
   ret void
 }
 
 declare void @Capture(i8*)