[MachineOutliner] Keep track of fns that use a redzone in AArch64FunctionInfo

This patch adds a hasRedZone() function to AArch64MachineFunctionInfo. It
returns true if the function is known to use a redzone, false if it is known
to not use a redzone, and no value otherwise.

This removes the requirement to pass -mno-red-zone when outlining for AArch64.

https://reviews.llvm.org/D45189

llvm-svn: 329120

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

@@ -531,9 +531,11 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF,
       return;
     // REDZONE: If the stack size is less than 128 bytes, we don't need
     // to actually allocate.
-    if (canUseRedZone(MF))
+    if (canUseRedZone(MF)) {
+      AFI->setHasRedZone(true);
       ++NumRedZoneFunctions;
-    else {
+    } else {
+      AFI->setHasRedZone(false);
       emitFrameOffset(MBB, MBBI, DL, AArch64::SP, AArch64::SP, -NumBytes, TII,
                       MachineInstr::FrameSetup);
 

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -4985,15 +4985,18 @@ bool AArch64InstrInfo::isFunctionSafeToOutlineFrom(
     MachineFunction &MF, bool OutlineFromLinkOnceODRs) const {
   const Function &F = MF.getFunction();
 
-  // If F uses a redzone, then don't outline from it because it might mess up
-  // the stack.
-  if (!F.hasFnAttribute(Attribute::NoRedZone))
-    return false;
-
   // Can F be deduplicated by the linker? If it can, don't outline from it.
   if (!OutlineFromLinkOnceODRs && F.hasLinkOnceODRLinkage())
     return false;
 
+  // Outlining from functions with redzones is unsafe since the outliner may
+  // modify the stack. Check if hasRedZone is true or unknown; if yes, don't
+  // outline from it.
+  AArch64FunctionInfo *AFI = MF.getInfo<AArch64FunctionInfo>();
+  if (!AFI || AFI->hasRedZone().getValueOr(true))
+    return false;
+
+  // It's safe to outline from MF.
   return true;
 }
 

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

@@ -15,6 +15,7 @@
 #define LLVM_LIB_TARGET_AARCH64_AARCH64MACHINEFUNCTIONINFO_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -90,11 +91,22 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
   /// other stack allocations.
   bool CalleeSaveStackHasFreeSpace = false;
 
+  /// \brief Has a value when it is known whether or not the function uses a
+  /// redzone, and no value otherwise.
+  /// Initialized during frame lowering, unless the function has the noredzone
+  /// attribute, in which case it is set to false at construction.
+  Optional<bool> HasRedZone;
+
 public:
   AArch64FunctionInfo() = default;
 
   explicit AArch64FunctionInfo(MachineFunction &MF) {
     (void)MF;
+
+    // If we already know that the function doesn't have a redzone, set
+    // HasRedZone here.
+    if (MF.getFunction().hasFnAttribute(Attribute::NoRedZone))
+      HasRedZone = false;
   }
 
   unsigned getBytesInStackArgArea() const { return BytesInStackArgArea; }
@@ -132,6 +144,9 @@ public:
     return NumLocalDynamicTLSAccesses;
   }
 
+  Optional<bool> hasRedZone() const { return HasRedZone; }
+  void setHasRedZone(bool s) { HasRedZone = s; }
+  
   int getVarArgsStackIndex() const { return VarArgsStackIndex; }
   void setVarArgsStackIndex(int Index) { VarArgsStackIndex = Index; }
 

llvm/test/CodeGen/AArch64/machine-outliner-noredzone.ll

@@ -0,0 +1,47 @@
+; RUN: llc -enable-machine-outliner %s -o - | FileCheck %s
+; CHECK: OUTLINED_FUNCTION
+; RUN: llc -enable-machine-outliner -aarch64-redzone %s -o - | FileCheck %s -check-prefix=REDZONE
+; REDZONE-NOT: OUTLINED_FUNCTION
+
+target triple = "arm64----"
+
+; Ensure that the MachineOutliner does not fire on functions which use a
+; redzone. foo() should have a redzone when compiled with -aarch64-redzone, and
+; no redzone otherwise.
+define void @foo() #0 {
+  %1 = alloca i32, align 4
+  %2 = alloca i32, align 4
+  %3 = alloca i32, align 4
+  %4 = alloca i32, align 4
+  store i32 0, i32* %1, align 4
+  store i32 0, i32* %2, align 4
+  store i32 0, i32* %3, align 4
+  store i32 0, i32* %4, align 4
+  %5 = load i32, i32* %1, align 4
+  %6 = add nsw i32 %5, 1
+  store i32 %6, i32* %1, align 4
+  %7 = load i32, i32* %3, align 4
+  %8 = add nsw i32 %7, 1
+  store i32 %8, i32* %3, align 4
+  %9 = load i32, i32* %4, align 4
+  %10 = add nsw i32 %9, 1
+  store i32 %10, i32* %4, align 4
+  %11 = load i32, i32* %2, align 4
+  %12 = add nsw i32 %11, 1
+  store i32 %12, i32* %2, align 4
+  %13 = load i32, i32* %1, align 4
+  %14 = add nsw i32 %13, 1
+  store i32 %14, i32* %1, align 4
+  %15 = load i32, i32* %3, align 4
+  %16 = add nsw i32 %15, 1
+  store i32 %16, i32* %3, align 4
+  %17 = load i32, i32* %4, align 4
+  %18 = add nsw i32 %17, 1
+  store i32 %18, i32* %4, align 4
+  %19 = load i32, i32* %2, align 4
+  %20 = add nsw i32 %19, -1
+  store i32 %20, i32* %2, align 4
+  ret void
+}
+
+attributes #0 = { noinline nounwind optnone }