[CodeGen][AArch64] Don't split functions with a red zone on AArch64

[leonardchan]

Summary

Because unconditional branch relaxation on AArch64 grows the stack to
spill a register, splitting a function would cause the red zone to be
overwritten. Explicitly disable MFS for such functions.

Depends on D157063

llvm/include/llvm/CodeGen/TargetInstrInfo.h

     return false;
   }
 
+  /// Return true if the function is a viable candidate for machine function
+  /// splitting. The criteria for if a function can be split may vary by target.
+  virtual bool isFunctionSafeToSplit(const MachineFunction &MF) const;
+
   /// Produce the expression describing the \p MI loading a value into
   /// the physical register \p Reg. This hook should only be used with
   /// \p MIs belonging to VReg-less functions.

llvm/lib/CodeGen/MachineFunctionSplitter.cpp

 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/Passes.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Support/CommandLine.h"

   if (!UseProfileData && !SplitAllEHCode)
     return false;
 
-  // TODO: We don't split functions where a section attribute has been set
-  // since the split part may not be placed in a contiguous region. It may also
-  // be more beneficial to augment the linker to ensure contiguous layout of
-  // split functions within the same section as specified by the attribute.
-  if (MF.getFunction().hasSection() ||
-      MF.getFunction().hasFnAttribute("implicit-section-name"))
+  const TargetInstrInfo &TII = *MF.getSubtarget().getInstrInfo();
+  if (!TII.isFunctionSafeToSplit(MF))
     return false;
 
-  // We don't want to proceed further for cold functions
-  // or functions of unknown hotness. Lukewarm functions have no prefix.
-  std::optional<StringRef> SectionPrefix = MF.getFunction().getSectionPrefix();
-  if (SectionPrefix &&
-      (*SectionPrefix == "unlikely" || *SectionPrefix == "unknown")) {
-    return false;
-  }
-
   // Renumbering blocks here preserves the order of the blocks as
   // sortBasicBlocksAndUpdateBranches uses the numeric identifier to sort
   // blocks. Preserving the order of blocks is essential to retaining decisions

llvm/lib/CodeGen/TargetInstrInfo.cpp

   return (DefCycle != -1 && DefCycle <= 1);
 }
 
+bool TargetInstrInfo::isFunctionSafeToSplit(const MachineFunction &MF) const {
+  // TODO: We don't split functions where a section attribute has been set
+  // since the split part may not be placed in a contiguous region. It may also
+  // be more beneficial to augment the linker to ensure contiguous layout of
+  // split functions within the same section as specified by the attribute.
+  if (MF.getFunction().hasSection() ||
+      MF.getFunction().hasFnAttribute("implicit-section-name"))
+    return false;
+
+  // We don't want to proceed further for cold functions
+  // or functions of unknown hotness. Lukewarm functions have no prefix.
+  std::optional<StringRef> SectionPrefix = MF.getFunction().getSectionPrefix();
+  if (SectionPrefix &&
+      (*SectionPrefix == "unlikely" || *SectionPrefix == "unknown")) {
+    return false;
+  }
+
+  return true;
+}
+
 std::optional<ParamLoadedValue>
 TargetInstrInfo::describeLoadedValue(const MachineInstr &MI,
                                      Register Reg) const {

llvm/lib/Target/AArch64/AArch64InstrInfo.h

   std::optional<RegImmPair> isAddImmediate(const MachineInstr &MI,
                                            Register Reg) const override;
 
+  bool isFunctionSafeToSplit(const MachineFunction &MF) const override;
+
   std::optional<ParamLoadedValue>
   describeLoadedValue(const MachineInstr &MI, Register Reg) const override;
 

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

   return std::nullopt;
 }
 
+bool AArch64InstrInfo::isFunctionSafeToSplit(const MachineFunction &MF) const {
+  // Functions cannot be split to different sections on AArch64 if they have
+  // a red zone. This is because relaxing a cross-section branch may require
+  // incrementing the stack pointer to spill a register, which would overwrite
+  // the red zone.
+  if (MF.getInfo<AArch64FunctionInfo>()->hasRedZone().value_or(true))
+    return false;
+
+  return TargetInstrInfo::isFunctionSafeToSplit(MF);
+}
+
 std::optional<ParamLoadedValue>
 AArch64InstrInfo::describeLoadedValue(const MachineInstr &MI,
                                       Register Reg) const {

llvm/test/CodeGen/AArch64/machine-function-splitter.mir

     %7 = tail call i32 @qux()
     ret void
   }
-  
+  ; Function Attrs: nounwind
+  define void @nosplit_redzone(i1 zeroext %0) #0 !prof !14 !section_prefix !15 {
+    br i1 %0, label %2, label %4, !prof !16
+
+  2:                                                ; preds = %1
+    %3 = call i32 @bar()
+    br label %6
+
+  4:                                                ; preds = %1
+    %5 = call i32 @baz()
+    br label %6
+
+  6:                                                ; preds = %4, %2
+    %7 = tail call i32 @qux()
+    ret void
+  }
+
   declare i32 @bar()
   
   declare i32 @baz()

     TCRETURNdi @qux, 0, csr_aarch64_aapcs, implicit $sp
 
 ...
+---
+name:            nosplit_redzone
+stack:
+  - { id: 0, name: '', type: spill-slot, offset: -16, size: 8, alignment: 16,
+      stack-id: default, callee-saved-register: '$lr', callee-saved-restored: true,
+      debug-info-variable: '', debug-info-expression: '', debug-info-location: '' }
+machineFunctionInfo:
+  hasRedZone:      true
+body:             |
+  ; CHECK-LABEL: name: nosplit_redzone
+  ; COM: Check that cold blocks in functions with red zones aren't split.
+  ; CHECK: bb.0 (%ir-block.1)
+  ; CHECK-NOT: bb.1 (%ir-block.2, bbsections Cold)
+  ; CHECK-NOT: bb.2 (%ir-block.4, bbsections Cold)
+
+  bb.0 (%ir-block.1):
+    successors: %bb.1(0x80000000), %bb.2(0x00000000)
+    liveins: $w0, $lr
+
+    early-clobber $sp = frame-setup STRXpre killed $lr, $sp, -16 :: (store (s64) into %stack.0)
+    CBZW killed renamable $w0, %bb.2
+
+  bb.1 (%ir-block.2):
+    BL @bar, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $w0
+    early-clobber $sp, $lr = frame-destroy LDRXpost $sp, 16 :: (load (s64) from %stack.0)
+    TCRETURNdi @qux, 0, csr_aarch64_aapcs, implicit $sp
+
+  bb.2 (%ir-block.4):
+    BL @baz, csr_aarch64_aapcs, implicit-def dead $lr, implicit $sp, implicit-def $sp, implicit-def dead $w0
+    early-clobber $sp, $lr = frame-destroy LDRXpost $sp, 16 :: (load (s64) from %stack.0)
+    TCRETURNdi @qux, 0, csr_aarch64_aapcs, implicit $sp
+
+...