diff --git a/llvm/include/llvm/CodeGen/TargetInstrInfo.h b/llvm/include/llvm/CodeGen/TargetInstrInfo.h --- a/llvm/include/llvm/CodeGen/TargetInstrInfo.h +++ b/llvm/include/llvm/CodeGen/TargetInstrInfo.h @@ -2066,6 +2066,10 @@ 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. diff --git a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp b/llvm/lib/CodeGen/MachineFunctionSplitter.cpp --- a/llvm/lib/CodeGen/MachineFunctionSplitter.cpp +++ b/llvm/lib/CodeGen/MachineFunctionSplitter.cpp @@ -35,6 +35,7 @@ #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" @@ -135,22 +136,10 @@ 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 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 diff --git a/llvm/lib/CodeGen/TargetInstrInfo.cpp b/llvm/lib/CodeGen/TargetInstrInfo.cpp --- a/llvm/lib/CodeGen/TargetInstrInfo.cpp +++ b/llvm/lib/CodeGen/TargetInstrInfo.cpp @@ -1373,6 +1373,26 @@ 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 SectionPrefix = MF.getFunction().getSectionPrefix(); + if (SectionPrefix && + (*SectionPrefix == "unlikely" || *SectionPrefix == "unknown")) { + return false; + } + + return true; +} + std::optional TargetInstrInfo::describeLoadedValue(const MachineInstr &MI, Register Reg) const { diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.h b/llvm/lib/Target/AArch64/AArch64InstrInfo.h --- a/llvm/lib/Target/AArch64/AArch64InstrInfo.h +++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.h @@ -325,6 +325,8 @@ std::optional isAddImmediate(const MachineInstr &MI, Register Reg) const override; + bool isFunctionSafeToSplit(const MachineFunction &MF) const override; + std::optional describeLoadedValue(const MachineInstr &MI, Register Reg) const override; diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp --- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp +++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp @@ -8368,6 +8368,17 @@ 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()->hasRedZone().value_or(true)) + return false; + + return TargetInstrInfo::isFunctionSafeToSplit(MF); +} + std::optional AArch64InstrInfo::describeLoadedValue(const MachineInstr &MI, Register Reg) const { diff --git a/llvm/test/CodeGen/AArch64/machine-function-splitter.mir b/llvm/test/CodeGen/AArch64/machine-function-splitter.mir --- a/llvm/test/CodeGen/AArch64/machine-function-splitter.mir +++ b/llvm/test/CodeGen/AArch64/machine-function-splitter.mir @@ -17,7 +17,24 @@ %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() @@ -82,3 +99,36 @@ 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 + +...