[CodeGen][AArch64] Don't split jump table basic blocks

[dhoekwater]

Summary

Jump tables on AArch64 are label-relative rather than table-relative, so
having jump table destinations that are in different sections causes
problems with relocation. Jump table lookups have a max range of 1MB, so
all destinations must be in the same section as the lookup code. Both of
these restrictions can be mitigated with some careful and complex logic,
but doing so doesn't gain a huge performance benefit.

Efficiently ensuring jump tables are correct and can be compressed on
AArch64 is a TODO item. In the meantime, don't split blocks that can
cause problems.

[dhoekwater]

Use an LLVM IR test instead of an MIR test since those are easier to read. Removes dependency on https://reviews.llvm.org/D157063.

[dhoekwater]

Use the correct test this time.

[mingmingl]

thanks! Mostly LG with one comment from compile-time perspective

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

Mostly from compile-time perspective, could this be pre-computed per MachineFunction in machine-function-splitter pass for AArch64?

Something like:

// precompute 'JumpTableTargets' for each machine function 
const auto& MJTI = F->getJumpTableInfo();

SmallSet<MachineBasicBlock*> JumpTableTargetsSet;

for (auto JT : MJTI->getJumpTables()) {
  for (auto* MBB : JT.MBBs) { 
    JumpTableTargetsSet.insert(MBB);
  }
}

As long as 'JumpTableTargetsSet' remains valid (i.e., not invalidated by any code transformations, like branch folding, etc)
For each 'AArch64InstrInfo::isMBBSafeToSplitToCold' method call

auto isJumpTableTarget = [&](const MachineBasicBlock& MBB) {
  return JumpTableTargets.count(MBB);
}

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

32–34 ↗

(On Diff #548776)

For my understanding, is cold1 split out to nosplit_jumptable.cold before this patch?

63 ↗

(On Diff #548776)

nit: if this is not used, maybe remove this?

[dhoekwater]

Move tests to Generic

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

Yeah, that would definitely be more efficient (O(n) worst case complexity in the number of blocks instead of O(n^2)).

However, I don't think this has a huge performance impact in real builds, and this logic only exists temporarily before I push changes that allow splitting jump table dispatch and targets. I'll collect a breakdown of how it impacts compile time.

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

32–34 ↗

(On Diff #548776)

Yeah, cold1 and cold2 are split out to the cold section before this patch.

63 ↗

(On Diff #548776)

Resolved in updated test.

[snehasish]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8427

IMO structuring the code with functors makes it hard to follow. Additionally, there are several cases where we could return false early instead which would lower the compile time in practice.

I would prefer the code to be structured as follows:

bool AArch64InstrInfo::isMBBSafeToSplitToCold(
    const MachineBasicBlock &MBB) const {
  // Check if MBB is JumpTableTarget
  const MachineJumpTableInfo *MJTI = MBB.getParent()->getJumpTableInfo();
 bool isJumpTableTarget = MJTI != nullptr &&
           llvm::any_of(MJTI->getJumpTables(), 
      [](const MachineJumpTableEntry &JTE){ return  llvm::find(JTE.MBBs, &MBB) != JTE.MBBs.end()});

  if (isJumpTableTarget) return false;
  
  // Check if contains lookup
  for(const auto &MI : MBB) {
        switch (MI.getOpcode()) {
    case TargetOpcode::G_BRJT:
    case AArch64::JumpTableDest32:
    case AArch64::JumpTableDest16:
    case AArch64::JumpTableDest8:
      return false;
    default:
      continue;
    }
  }
  // All checks passed.
  return true;
}

What do you think?

[dhoekwater]

Refactor isMBBSafeToSplitToCold to be more readable.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

I just timed the backend compile. Although I don't have per-pass numbers, enabling function splitting only increases the build time of a large binary by 0.15%.

@mingmingl @snehasish Thoughts on this?

8427

That sounds pretty reasonable to me; it's definitely an easier read.

[dhoekwater]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8427

I refactored the function per your feedback, but I did rework and keep one lambda. It follows the common pattern of defining a lambda then using it in an STL library on the next line as its only use, which I think is pretty readable.

[snehasish]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

Seems fine in the interim while we wait for jump table specific fixes to land.

8427

Agree, looks good. Maybe limit the capture in the lambda to just &MBB to make it clearer?

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

521

Do we need another test case when the MBB contains a jump table lookup?

[dhoekwater]

Reduce the scope of the lambda capture and add a test to make sure jump table dispatch isn't split on AArch64

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8427

Done!

[dhoekwater]

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

521

Good call; done.

[dhoekwater]

Rebase onto main to reduce diff

[snehasish]

lgtm

[aeubanks]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

0.15% is a fairly large regression. do you mean that enabling function splitting on its own regresses compile times compared to not enabling it, or that this patch regresses function splitting builds by 0.15%?

[dhoekwater]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

Enabling function splitting for an AArch64 target makes the build 0.15% slower than leaving it disabled. This patch has no impact on compile times of builds that don't use machine function splitting for AArch64.

I'll collect the impact of just this patch to see its individual impact.

[aeubanks]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

ah ok that sounds reasonable

[snehasish]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

I believe that builds which enable function splitting on AArch64 can break without this fix. So it seemed reasonable to me to submit a fix while we wait for a more optimized approach to land.

Also @aeubanks do we have some guidelines on compile time increases across configurations that we can refer to in the future? Thanks!

[aeubanks]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

+@nikic
I don't think we have anything official, but we do yell at people who noticeably regress compile times, especially if it shows up on https://llvm-compile-time-tracker.com/ as red. (although I believe that only tracks x86-64 for now)
One 0.15% regression might be fine, but they add up over time so we'd really like to not introduce noticeable regressions when possible. We probably should have a document somewhere that talks about acceptable compile time tradeoffs.

[dhoekwater]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8418–8425

Builds which enable function splitting on AArch64 without this fix can break unless they disable jump tables. For AArch64 MFS builds, compiling with this fix takes virtually the same time (technically 0.07% faster) as compiling with jump tables disabled.

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}