[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

8391–8398

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

8391–8398

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

8400

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

8391–8398

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?

8400

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

[dhoekwater]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8400

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

8391–8398

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

8400

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 ↗

(On Diff #553295)

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

8400

Done!

[dhoekwater]

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

521 ↗

(On Diff #553295)

Good call; done.

[dhoekwater]

Rebase onto main to reduce diff

[snehasish]

lgtm

[aeubanks]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8391–8398

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

8391–8398

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

8391–8398

ah ok that sounds reasonable

[snehasish]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

8391–8398

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

8391–8398

+@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

8391–8398

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

llvm_unreachable("Target didn't implement "
                      "TargetInstrInfo::isFunctionSafeToOutlineFrom!");
   }
 
   /// Return true if the function should be outlined from by default.
   virtual bool shouldOutlineFromFunctionByDefault(MachineFunction &MF) const {
     return false;
   }
 
+  /// 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

 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/CodeGen/BasicBlockSectionUtils.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #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"
 #include <optional>
 
 using namespace llvm;
 
 // FIXME: This cutoff value is CPU dependent and should be moved to
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:
   static bool isFalkorShiftExtFast(const MachineInstr &MI);
   /// Return true if the instructions is a SEH instruciton used for unwinding
   /// on Windows.
   static bool isSEHInstruction(const MachineInstr &MI);
 
   std::optional<RegImmPair> isAddImmediate(const MachineInstr &MI,
                                            Register Reg) 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

describeORRLoadedValue(const MachineInstr &MI, Register DescribedReg,
   }
 
   assert(!TRI->isSuperOrSubRegisterEq(DestReg, DescribedReg) &&
          "Unhandled ORR[XW]rs copy case");
 
   return std::nullopt;
 }
 
+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 else relocation fixup handling will throw a
+  // fit.
+  auto isJumpTableLookup = [](const MachineInstr &MI) {
+    switch (MI.getOpcode()) {
+    case TargetOpcode::G_BRJT:
+    case AArch64::JumpTableDest32:
+    case AArch64::JumpTableDest16:
+    case AArch64::JumpTableDest8:
+      return true;
+    default:
+      return false;
+    }
+  };
+  auto containsJumpTableLookup = [&](const MachineBasicBlock &MBB) {
+    return llvm::any_of(MBB, isJumpTableLookup);
+  };
+
+  auto isInJumpTable = [&](const MachineJumpTableEntry &JTE) {
+    return llvm::find(JTE.MBBs, &MBB) != JTE.MBBs.end();
+  };
+  auto isJumpTableTarget = [&](const MachineBasicBlock &MBB) {
+    const MachineJumpTableInfo *MJTI = MBB.getParent()->getJumpTableInfo();
+    return MJTI != nullptr &&
+           llvm::any_of(MJTI->getJumpTables(), isInJumpTable);
+  };
+
+  return !containsJumpTableLookup(MBB) && !isJumpTableTarget(MBB);
+}
+
 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/AArch64/machine-function-splitter.mir

--- |
   4:                                                ; preds = %1
     %5 = call i32 @baz()
     br label %6
 
   6:                                                ; preds = %4, %2
     %7 = tail call i32 @qux()
     ret void
   }
-  
+  ; Function Attrs: nounwind
+  define i32 @nosplit_jumptable(i32 %in) #0 !prof !14 !section_prefix !15 {
+    switch i32 %in, label %common.ret [
+      i32 0, label %hot1
+      i32 1, label %hot2
+      i32 2, label %cold1
+      i32 3, label %cold2
+    ], !prof !17
+  
+  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()
   
--- |
   !9 = !{!"DetailedSummary", !10}
   !10 = !{!11, !12, !13}
   !11 = !{i32 10000, i64 100, i32 1}
   !12 = !{i32 999900, i64 100, i32 1}
   !13 = !{i32 999999, i64 1, i32 2}
   !14 = !{!"function_entry_count", i64 9000}
   !15 = !{!"function_section_prefix", !"hot"}
   !16 = !{!"branch_weights", i32 7000, i32 0}
+  !17 = !{!"branch_weights", i32 1000, i32 4000, i32 4000, i32 0, i32 0}
 
 ...
 ---
 name:            split_basic
 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: '' }
bb.1 (%ir-block.2):
     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
 
 ...
+---
+name:            nosplit_jumptable
+machineFunctionInfo:
+  hasRedZone:      false
+jumpTable:
+  kind:            block-address
+  entries:
+    - id:              0
+      blocks:          [ '%bb.3', '%bb.4', '%bb.5', '%bb.6' ]
+body:             |
+  ; CHECK-LABEL: name: nosplit_jumptable
+  ; COM: Check that a cold block targeted by a jump table is not split.
+  ; CHECK-NOT: bbsections Cold
+  ; CHECK:     TCRETURNdi @qux
+  bb.0 (%ir-block.0):
+    successors: %bb.2(0x0e38e38e), %bb.1(0x71c71c72)
+    liveins: $w0
+  
+    dead $wzr = SUBSWri renamable $w0, 3, 0, implicit-def $nzcv
+    Bcc 8, %bb.2, implicit killed $nzcv
+  
+  bb.1 (%ir-block.0):
+    successors: %bb.3(0x40000000), %bb.4(0x40000000), %bb.5(0x00000000), %bb.6(0x00000000)
+    liveins: $w0
+  
+    renamable $w8 = ORRWrs $wzr, killed renamable $w0, 0, implicit-def $x8
+    $x9 = ADRP target-flags(aarch64-page) %jump-table.0
+    renamable $x9 = ADDXri killed $x9, target-flags(aarch64-pageoff, aarch64-nc) %jump-table.0, 0
+    early-clobber renamable $x10, dead early-clobber renamable $x11 = JumpTableDest32 killed renamable $x9, killed renamable $x8, %jump-table.0
+    BR killed renamable $x10
+  
+  bb.3.hot1:
+    TCRETURNdi @bar, 0, csr_aarch64_aapcs, implicit $sp
+  
+  bb.4.hot2:
+    TCRETURNdi @baz, 0, csr_aarch64_aapcs, implicit $sp
+  
+  bb.2.common.ret:
+    $w0 = ORRWrs $wzr, $wzr, 0
+    RET undef $lr, implicit killed $w0
+  
+  bb.5.cold1:
+    TCRETURNdi @bam, 0, csr_aarch64_aapcs, implicit $sp
+  
+  bb.6.cold2:
+    TCRETURNdi @qux, 0, csr_aarch64_aapcs, implicit $sp
+
+...