[CodeGen] Make the MachineFunctionSplitter pass testable via MIR

[dhoekwater]

Summary

Because the MachineFunctionSplitter pass isn't initialized in
CodeGen.cpp, running LLC with --run-pass machine-function-splitter
doesn't work. Adding it to the initializeCodeGen makes the function
splitting directly testable with .mir tests.

[dhoekwater]

Remove an unnecessary line.

[dhoekwater]

Rebase onto main to hopefully make unrelated test failures go away

[snehasish]

Is there a specific reason to test MFS (on ARM) with MIR instead of .ll inputs? Only ~11% of the tests in test/CodeGen/AArch64 are in MIR, the rest are bitcode files. IMO bitcode is much easier to read, maintain and we can reuse most (all?) of the existing tests we have in machine-function-splitter.ll. So my preference would be to use bitcode unless there is a fundamental reason which blocks us from testing with bitcode.

[dhoekwater]

In D157063#4573702, @snehasish wrote:

Is there a specific reason to test MFS (on ARM) with MIR instead of .ll inputs?

In a different review, I was asked to make MIR tests so that they only targeted a single pass. That doesn't seem necessary here, so I'll use bitcode tests instead. Dropping this revision and updating its dependents.

[dhoekwater]

Try to replace with main to get this to work

[dhoekwater]

Try to replace with a different branch to get this to work

[dhoekwater]

Try to replace with yet another

llvm/include/llvm/CodeGen/TargetInstrInfo.h

public:
   virtual bool shouldOutlineFromFunctionByDefault(MachineFunction &MF) const {
     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;
 
+  /// Return true if the MachineBasicBlock can safely be split to the cold
+  /// section. On AArch64, certain instructions may cause a block to be unsafe
+  /// to split to the cold section.
+  virtual bool isMBBSafeToSplitToCold(const MachineBasicBlock &MBB) const {
+    return true;
+  }
+
   /// 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.
   virtual std::optional<ParamLoadedValue>
   describeLoadedValue(const MachineInstr &MI, Register Reg) const;
 
   /// Given the generic extension instruction \p ExtMI, returns true if this
   /// extension is a likely candidate for being folded into an another

llvm/lib/CodeGen/MachineFunctionSplitter.cpp

static void finishAdjustingBasicBlocksAndLandingPads(MachineFunction &MF) {
   llvm::sortBasicBlocksAndUpdateBranches(MF, Comparator);
   llvm::avoidZeroOffsetLandingPad(MF);
 }
 
 static bool isColdBlock(const MachineBasicBlock &MBB,
                         const MachineBlockFrequencyInfo *MBFI,
                         ProfileSummaryInfo *PSI) {
   std::optional<uint64_t> Count = MBFI->getBlockProfileCount(&MBB);
+
+  // Temporary hack to cope with AArch64's jump table encoding
+  const TargetInstrInfo &TII = *MBB.getParent()->getSubtarget().getInstrInfo();
+  if (!TII.isMBBSafeToSplitToCold(MBB))
+    return false;
+
   // For instrumentation profiles and sample profiles, we use different ways
   // to judge whether a block is cold and should be split.
   if (PSI->hasInstrumentationProfile() || PSI->hasCSInstrumentationProfile()) {
     // If using instrument profile, which is deemed "accurate", no count means
     // cold.
     if (!Count)
       return true;
     if (PercentileCutoff > 0)

llvm/lib/Target/AArch64/AArch64InstrInfo.h

public:
   /// on Windows.
   static bool isSEHInstruction(const MachineInstr &MI);
 
   std::optional<RegImmPair> isAddImmediate(const MachineInstr &MI,
                                            Register Reg) const override;
 
   bool isFunctionSafeToSplit(const MachineFunction &MF) const override;
 
+  bool isMBBSafeToSplitToCold(const MachineBasicBlock &MBB) const override;
+
   std::optional<ParamLoadedValue>
   describeLoadedValue(const MachineInstr &MI, Register Reg) const override;
 
   unsigned int getTailDuplicateSize(CodeGenOpt::Level OptLevel) const override;
 
   bool isExtendLikelyToBeFolded(MachineInstr &ExtMI,
                                 MachineRegisterInfo &MRI) const override;
 

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

bool AArch64InstrInfo::isFunctionSafeToSplit(const MachineFunction &MF) const {
   // 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);
 }
 
+bool AArch64InstrInfo::isMBBSafeToSplitToCold(
+    const MachineBasicBlock &MBB) const {
+  // Because jump tables are label-relative instead of table-relative, they all
+  // must be in the same section or relocation fixup handling will fail.
+
+  // Check if MBB is a jump table target
+  const MachineJumpTableInfo *MJTI = MBB.getParent()->getJumpTableInfo();
+  auto containsMBB = [&MBB](const MachineJumpTableEntry &JTE) {
+    return llvm::is_contained(JTE.MBBs, &MBB);
+  };
+  if (MJTI != nullptr && llvm::any_of(MJTI->getJumpTables(), containsMBB))
+    return false;
+
+  // Check if MBB contains a jump table lookup
+  for (const MachineInstr &MI : MBB) {
+    switch (MI.getOpcode()) {
+    case TargetOpcode::G_BRJT:
+    case AArch64::JumpTableDest32:
+    case AArch64::JumpTableDest16:
+    case AArch64::JumpTableDest8:
+      return false;
+    default:
+      continue;
+    }
+  }
+
+  // MBB isn't a special case, so it's safe to be split to the cold section.
+  return true;
+}
+
 std::optional<ParamLoadedValue>
 AArch64InstrInfo::describeLoadedValue(const MachineInstr &MI,
                                       Register Reg) const {
   const MachineFunction *MF = MI.getMF();
   const TargetRegisterInfo *TRI = MF->getSubtarget().getRegisterInfo();
   switch (MI.getOpcode()) {
   case AArch64::MOVZWi:
   case AArch64::MOVZXi: {

llvm/test/CodeGen/Generic/machine-function-splitter.ll

4:                                                ; preds = %3, %2
   %tmp = load i32, ptr %a.addr, align 4
   %tmp1 = load i32, ptr %b.addr, align 4
   %add = add nsw i32 %tmp, %tmp1
   store i32 %add, ptr %x, align 4
   %tmp2 = load i32, ptr %x, align 4
   ret i32 %tmp2
 }
 
+define i32 @foo18(i32 %in) !prof !14 !section_prefix !15 {
+;; Check that a cold block targeted by a jump table is not split
+;; on AArch64.
+; MFS-DEFAULTS-LABEL:        foo18
+; MFS-DEFAULTS:              .section        .text.split.foo18
+; MFS-DEFAULTS-NEXT:         foo18.cold:
+; MFS-DEFAULTS-SAME:           %common.ret
+; MFS-DEFAULTS-X86-DAG:        jmp     qux
+; MFS-DEFAULTS-X86-DAG:        jmp     bam
+; MFS-DEFAULTS-AARCH64-NOT:    b       bar
+; MFS-DEFAULTS-AARCH64-NOT:    b       baz
+; MFS-DEFAULTS-AARCH64-NOT:    b       qux
+; MFS-DEFAULTS-AARCH64-NOT:    b       bam
+
+  switch i32 %in, label %common.ret [
+    i32 0, label %hot1
+    i32 1, label %hot2
+    i32 2, label %cold1
+    i32 3, label %cold2
+  ], !prof !28
+
+common.ret:                                       ; preds = %0
+  ret i32 0
+
+hot1:                                             ; preds = %0
+  %1 = tail call i32 @bar()
+  ret i32 %1
+
+hot2:                                             ; preds = %0
+  %2 = tail call i32 @baz()
+  ret i32 %2
+
+cold1:                                            ; preds = %0
+  %3 = tail call i32 @bam()
+  ret i32 %3
+
+cold2:                                            ; preds = %0
+  %4 = tail call i32 @qux()
+  ret i32 %4
+}
+
+define i32 @foo19(i32 %in) !prof !14 !section_prefix !15 {
+;; Check that a cold block that contains a jump table dispatch is
+;; not split on AArch64.
+; MFS-DEFAULTS-LABEL:        foo19
+; MFS-DEFAULTS:              .section        .text.split.foo19
+; MFS-DEFAULTS-NEXT:         foo19.cold:
+; MFS-DEFAULTS-X86:            .LJTI18_0
+; MFS-DEFAULTS-AARCH64-NOT:    .LJTI18_0
+; MFS-DEFAULTS:              .section        .rodata
+; MFS-DEFAULTS:                .LJTI18_0
+  %cmp = icmp sgt i32 %in, 3
+  br i1 %cmp, label %hot, label %cold_switch, !prof !17
+
+hot:                                              ; preds = %0
+ret i32 1
+
+cold_switch:                                      ; preds = %0
+  switch i32 %in, label %common.ret [
+    i32 0, label %hot1
+    i32 1, label %hot2
+    i32 2, label %cold1
+    i32 3, label %cold2
+  ], !prof !28
+
+common.ret:                                       ; preds = %0
+  ret i32 0
+
+hot1:                                             ; preds = %0
+  %1 = tail call i32 @bar()
+  ret i32 %1
+
+hot2:                                             ; preds = %0
+  %2 = tail call i32 @baz()
+  ret i32 %2
+
+cold1:                                            ; preds = %0
+  %3 = tail call i32 @bam()
+  ret i32 %3
+
+cold2:                                            ; preds = %0
+  %4 = tail call i32 @qux()
+  ret i32 %4
+}
+
 declare i32 @bar()
 declare i32 @baz()
 declare i32 @bam()
 declare i32 @qux()
 declare void @_Z1fv()
 declare i32 @__gxx_personality_v0(...)
 
 @_ZTIi = external constant ptr

 !20 = !{!"function_entry_count", i64 10000}
 !21 = !{!"branch_weights", i32 6000, i32 4000}
 !22 = !{!"branch_weights", i32 80, i32 9920}
 !23 = !{!"function_entry_count", i64 7}
 !24 = !{!"function_entry_count", i64 10000}
 !25 = !{!"branch_weights", i32 0, i32 7000}
 !26 = !{!"branch_weights", i32 1000, i32 6000}
 !27 = !{!"function_entry_count", i64 10000}
+!28 = !{!"branch_weights", i32 0, i32 4000, i32 4000, i32 0, i32 0}