[safestack] Experimental mode where stack pointer is accessed with a function call.

[eugenis]

Summary

This adds an -mllvm flag that forces the use of a runtime function call to
get the unsafe stack pointer, the same that is currently used on non-x86, non-aarch64 android.

[eugenis]

The function should really have coldcc or preserve_allcc but those are either not implemented (i.e. crash the compiler) or just not do anything useful on arm and aarch64.
As is, generated code is atrocious.

[eugenis]

Attempt to use a custom calling convention for the safestack pointer accessor function.

Not ready for review.

[eugenis]

An attempt to inline the safestack accessor function.

To use, add -mllvm -safestack-pointer-address-inline=1.
Function must be called safestack_pointer_address_inline.
"attribute__((used)) static" is recommended so the function is not discarded early.

[eugenis]

Remove -customcc code. Added inlining logic.

Please review.

[eugenis]

removed one flag, fixed a condition, improved the test

[eugenis]

fix one test

[pcc]

LGTM

llvm/lib/CodeGen/SafeStack.cpp

 #include "SafeStackLayout.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/Argument.h"

 #include "llvm/Pass.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/Local.h"
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
 #include <string>
 #include <utility>
 
 using namespace llvm;

 STATISTIC(NumAllocas, "Total number of allocas");
 STATISTIC(NumUnsafeStaticAllocas, "Number of unsafe static allocas");
 STATISTIC(NumUnsafeDynamicAllocas, "Number of unsafe dynamic allocas");
 STATISTIC(NumUnsafeByValArguments, "Number of unsafe byval arguments");
 STATISTIC(NumUnsafeStackRestorePoints, "Number of setjmps and landingpads");
 
 } // namespace llvm
 
+/// Use __safestack_pointer_address even if the platform has a faster way of
+/// access safe stack pointer.
+static cl::opt<bool>
+    SafeStackUsePointerAddress("safestack-use-pointer-address",
+                                  cl::init(false), cl::Hidden);
+
+
 namespace {
 
 /// Rewrite an SCEV expression for a memory access address to an expression that
 /// represents offset from the given alloca.
 ///
 /// The implementation simply replaces all mentions of the alloca with zero.
 class AllocaOffsetRewriter : public SCEVRewriteVisitor<AllocaOffsetRewriter> {
   const Value *AllocaPtr;
class SafeStack {
 
   bool IsSafeStackAlloca(const Value *AllocaPtr, uint64_t AllocaSize);
 
   bool IsMemIntrinsicSafe(const MemIntrinsic *MI, const Use &U,
                           const Value *AllocaPtr, uint64_t AllocaSize);
   bool IsAccessSafe(Value *Addr, uint64_t Size, const Value *AllocaPtr,
                     uint64_t AllocaSize);
 
+  bool ShouldInlinePointerAddress(CallSite &CS);
+  void TryInlinePointerAddress();
+
 public:
   SafeStack(Function &F, const TargetLoweringBase &TL, const DataLayout &DL,
             ScalarEvolution &SE)
       : F(F), TL(TL), DL(DL), SE(SE),
         StackPtrTy(Type::getInt8PtrTy(F.getContext())),
         IntPtrTy(DL.getIntPtrType(F.getContext())),
         Int32Ty(Type::getInt32Ty(F.getContext())),
         Int8Ty(Type::getInt8Ty(F.getContext())) {}
for (inst_iterator It = inst_begin(&F), Ie = inst_end(&F); It != Ie;) {
         SI->takeName(II);
         assert(II->use_empty());
         II->eraseFromParent();
       }
     }
   }
 }
 
+bool SafeStack::ShouldInlinePointerAddress(CallSite &CS) {
+  Function *Callee = CS.getCalledFunction();
+  if (CS.hasFnAttr(Attribute::AlwaysInline) && isInlineViable(*Callee))
+    return true;
+  if (Callee->isInterposable() || Callee->hasFnAttribute(Attribute::NoInline) ||
+      CS.isNoInline())
+    return false;
+  return true;
+}
+
+void SafeStack::TryInlinePointerAddress() {
+  if (!isa<CallInst>(UnsafeStackPtr))
+    return;
+
+  if(F.hasFnAttribute(Attribute::OptimizeNone))
+    return;
+
+  CallSite CS(UnsafeStackPtr);
+  Function *Callee = CS.getCalledFunction();
+  if (!Callee || Callee->isDeclaration())
+    return;
+
+  if (!ShouldInlinePointerAddress(CS))
+    return;
+
+  InlineFunctionInfo IFI;
+  InlineFunction(CS, IFI);
+}
+
 bool SafeStack::run() {
   assert(F.hasFnAttribute(Attribute::SafeStack) &&
          "Can't run SafeStack on a function without the attribute");
   assert(!F.isDeclaration() && "Can't run SafeStack on a function declaration");
 
   ++NumFunctions;
 
   SmallVector<AllocaInst *, 16> StaticAllocas;
bool SafeStack::run() {
   if (!StaticAllocas.empty() || !DynamicAllocas.empty() ||
       !ByValArguments.empty())
     ++NumUnsafeStackFunctions; // This function has the unsafe stack.
 
   if (!StackRestorePoints.empty())
     ++NumUnsafeStackRestorePointsFunctions;
 
   IRBuilder<> IRB(&F.front(), F.begin()->getFirstInsertionPt());
+  if (SafeStackUsePointerAddress) {
+    Value *Fn = F.getParent()->getOrInsertFunction(
+        "__safestack_pointer_address", StackPtrTy->getPointerTo(0));
+    UnsafeStackPtr = IRB.CreateCall(Fn);
+  } else {
     UnsafeStackPtr = TL.getSafeStackPointerLocation(IRB);
+  }
 
   // Load the current stack pointer (we'll also use it as a base pointer).
   // FIXME: use a dedicated register for it ?
   Instruction *BasePointer =
       IRB.CreateLoad(UnsafeStackPtr, false, "unsafe_stack_ptr");
   assert(BasePointer->getType() == StackPtrTy);
 
   AllocaInst *StackGuardSlot = nullptr;
moveDynamicAllocasToUnsafeStack(F, UnsafeStackPtr, DynamicTop,
                                   DynamicAllocas);
 
   // Restore the unsafe stack pointer before each return.
   for (ReturnInst *RI : Returns) {
     IRB.SetInsertPoint(RI);
     IRB.CreateStore(BasePointer, UnsafeStackPtr);
   }
 
+  TryInlinePointerAddress();
+
   DEBUG(dbgs() << "[SafeStack]     safestack applied\n");
   return true;
 }
 
 class SafeStackLegacyPass : public FunctionPass {
   const TargetMachine *TM = nullptr;
 
 public:

llvm/test/CodeGen/X86/safestack.ll

 ; 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=ANDROID-I386 %s
 ; RUN: llc -mtriple=x86_64-linux-android < %s -o - | FileCheck --check-prefix=ANDROID-X64 %s
 ; RUN: llc -mtriple=x86_64-fuchsia < %s -o - | FileCheck --check-prefix=FUCHSIA-X64 %s
 
+; RUN: llc -mtriple=i386-linux -safestack-use-pointer-address < %s -o - | FileCheck --check-prefix=LINUX-I386-PA %s
+
 define void @_Z1fv() safestack {
 entry:
   %x = alloca i32, align 4
   %0 = bitcast i32* %x to i8*
   call void @_Z7CapturePi(i32* nonnull %x)
   ret void
 }
 

 
 ; ANDROID-X64: movq %fs:72, %[[A:.*]]
 ; ANDROID-X64: leaq -16(%[[A]]), %[[B:.*]]
 ; ANDROID-X64: movq %[[B]], %fs:72
 
 ; FUCHSIA-X64: movq %fs:24, %[[A:.*]]
 ; FUCHSIA-X64: leaq -16(%[[A]]), %[[B:.*]]
 ; FUCHSIA-X64: movq %[[B]], %fs:24
+
+; LINUX-I386-PA: calll __safestack_pointer_address
+; LINUX-I386-PA: movl %eax, %[[A:.*]]
+; LINUX-I386-PA: movl (%[[A]]), %[[B:.*]]
+; LINUX-I386-PA: leal -16(%[[B]]), %[[C:.*]]
+; LINUX-I386-PA: movl %[[C]], (%[[A]])

llvm/test/CodeGen/X86/safestack_inline.ll

+; RUN: sed -e "s/ATTR//" %s | llc -mtriple=x86_64-linux -safestack-use-pointer-address | FileCheck --check-prefix=INLINE %s
+; RUN: sed -e "s/ATTR/noinline/" %s | llc -mtriple=x86_64-linux -safestack-use-pointer-address | FileCheck --check-prefix=CALL %s
+
+@p = external thread_local global i8*, align 8
+
+define nonnull i8** @__safestack_pointer_address() local_unnamed_addr ATTR {
+entry:
+  ret i8** @p
+}
+
+define void @_Z1fv() safestack {
+entry:
+  %x = alloca i32, align 4
+  %0 = bitcast i32* %x to i8*
+  call void @_Z7CapturePi(i32* nonnull %x)
+  ret void
+}
+
+declare void @_Z7CapturePi(i32*)
+
+; INLINE: movq p@GOTTPOFF(%rip), %[[A:.*]]
+; INLINE: movq %fs:(%[[A]]), %[[B:.*]]
+; INLINE: leaq -16(%[[B]]), %[[C:.*]]
+; INLINE: movq %[[C]], %fs:(%[[A]])
+
+; CALL: callq __safestack_pointer_address
+; CALL: movq %rax, %[[A:.*]]
+; CALL: movq (%[[A]]), %[[B:.*]]
+; CALL: leaq -16(%[[B]]), %[[C:.*]]
+; CALL: movq %[[C]], (%[[A]])