[MachineOutliner][AArch64] Keep track of functions that use a red zone in AArch64MachineFunctionInfo and use that instead of checking for noredzone in the MachineOutliner

[paquette]

Summary

This patch adds the a hasRedZone() function to AArch64MachineFunctionInfo.

Right now, to use the MachineOutliner, you have to pass -mno-red-zone to ensure that we never outline anything from a function which uses a red zone. This is kind of clunky, especially considering you have to pass -aarch64-redzone to use a red zone in the first place in AArch64.

This patch removes the -mno-red-zone requirement and teaches the outliner about which functions actually use a red zone. This would also allow for outlining from functions without red zones in the case where -aarch64-redzone is passed.

[thegameg]

lib/Target/AArch64/AArch64MachineFunctionInfo.h

95

Since this is available only after PEI, I wonder if it wouldn't be safer to initialize UsesRedZone to !F.hasFnAttribute(Attribute::NoRedZone). That way if the function is called before PEI we use the attribute, and after PEI we overwrite this.

[javed.absar]

lib/Target/AArch64/AArch64MachineFunctionInfo.h

147

nitpick - should this be called 'hasRedZone' to be consistent with other namings e.g. hasStackFrame?

[gberry]

lib/Target/AArch64/AArch64MachineFunctionInfo.h

95

If the value is going to change over time, perhaps you could name the function mightUseRedZone instead?
Or you could use an Option<bool> and assert if it is queried before it is set?

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

8

This could use a one line description of what it is testing.

10

I assume this test needs to be this large in order to cross some threshold to trigger outlining?

[gberry]

Another thought, would it be worthwhile to add a serialized MIR function property for this usesRedZone flag so you could write MIR tests for this instead?

[thegameg]

In D45189#1055625, @gberry wrote:

Another thought, would it be worthwhile to add a serialized MIR function property for this usesRedZone flag so you could write MIR tests for this instead?

In that case this would have to go to MachineFrameInfo since we don't have MIR serialization for (Target)FunctionInfo.

[paquette]

Updated diff to address comments.

usesRedZone is now hasRedZone.

hasRedZone is now a function that returns an optional. The optional is set either during frame lowering, or on the construction of AArch64FunctionInfo in the case where the IR-level function has the noredzone attribute.

I tried to simplify the test a little, but it needs to be somewhat long to trigger the desired outliner behaviour, unfortunately.

[paquette]

Marked comments, etc.

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

10

Yes, unfortunately. :(

[thegameg]

LGTM, thanks!

[paquette]

Committed in r329120.

lib/Target/AArch64/AArch64FrameLowering.cpp

if (!AFI->hasStackFrame() && !windowsRequiresStackProbe(MF, NumBytes)) {
 
     // All of the stack allocation is for locals.
     AFI->setLocalStackSize(NumBytes);
 
     if (!NumBytes)
       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);
 
       // Label used to tie together the PROLOG_LABEL and the MachineMoves.
       MCSymbol *FrameLabel = MMI.getContext().createTempSymbol();
       // Encode the stack size of the leaf function.
       unsigned CFIIndex = MF.addFrameInst(
           MCCFIInstruction::createDefCfaOffset(FrameLabel, -NumBytes));

lib/Target/AArch64/AArch64InstrInfo.cpp

AArch64InstrInfo::getOutlininingCandidateInfo(
   return MachineOutlinerInfo(NumInstrsForCall, NumInstrsToCreateFrame, CallID,
                              FrameID);
 }
 
 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;
 }
 
 unsigned
 AArch64InstrInfo::getMachineOutlinerMBBFlags(MachineBasicBlock &MBB) const {
   unsigned Flags = 0x0;
   // Check if there's a call inside this MachineBasicBlock. If there is, then
   // set a flag.

lib/Target/AArch64/AArch64MachineFunctionInfo.h

 // This file declares AArch64-specific per-machine-function information.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_AARCH64_AARCH64MACHINEFUNCTIONINFO_H
 #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"
 #include "llvm/MC/MCLinkerOptimizationHint.h"
 #include <cassert>
 
 namespace llvm {
 
class AArch64FunctionInfo final : public MachineFunctionInfo {
   /// True when the stack gets realigned dynamically because the size of stack
   /// frame is unknown at compile time. e.g., in case of VLAs.
   bool StackRealigned = false;
 
   /// True when the callee-save stack area has unused gaps that may be used for
   /// 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; }
   void setBytesInStackArgArea(unsigned bytes) { BytesInStackArgArea = bytes; }
 
   unsigned getArgumentStackToRestore() const { return ArgumentStackToRestore; }
   void setArgumentStackToRestore(unsigned bytes) {
     ArgumentStackToRestore = bytes;
public:
   void setCalleeSavedStackSize(unsigned Size) { CalleeSavedStackSize = Size; }
   unsigned getCalleeSavedStackSize() const { return CalleeSavedStackSize; }
 
   void incNumLocalDynamicTLSAccesses() { ++NumLocalDynamicTLSAccesses; }
   unsigned getNumLocalDynamicTLSAccesses() const {
     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; }
 
   int getVarArgsGPRIndex() const { return VarArgsGPRIndex; }
   void setVarArgsGPRIndex(int Index) { VarArgsGPRIndex = Index; }
 
   unsigned getVarArgsGPRSize() const { return VarArgsGPRSize; }
   void setVarArgsGPRSize(unsigned Size) { VarArgsGPRSize = Size; }

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

+; 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 }