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

[Snehasish Kumar via Phabricator via llvm-commits]

snehasish added inline comments.


================
Comment at: llvm/lib/Target/AArch64/AArch64InstrInfo.cpp:8391-8398
+  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);
----------------
dhoekwater wrote:
> dhoekwater wrote:
> > mingmingl wrote:
> > > 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);
> > > }
> > > ```
> > 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.
> 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?
Seems fine in the interim while we wait for jump table specific fixes to land. 


================
Comment at: llvm/lib/Target/AArch64/AArch64InstrInfo.cpp:8421
+
+  return !containsJumpTableLookup(MBB) && !isJumpTableTarget(MBB);
+}
----------------
dhoekwater wrote:
> dhoekwater wrote:
> > snehasish wrote:
> > > 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?
> > That sounds pretty reasonable to me; it's definitely an easier read.
> 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.
Agree, looks good. Maybe limit the capture in the lambda to just &MBB to make it clearer?


================
Comment at: llvm/test/CodeGen/Generic/machine-function-splitter.ll:521
+define i32 @foo18(i32 %in) !prof !14 !section_prefix !15 {
+;; Check that a cold block targeted by a jump table is not split.
+; MFS-DEFAULTS-LABEL:        foo18
----------------
Do we need another test case when the MBB contains a jump table lookup?


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D157124/new/

https://reviews.llvm.org/D157124