[AArch64] Avoid using JumpTableEntryInfo when lowering jump tables when it is not initialized

[sinan]

Summary

JumpTableEntryInfo is initialized during the instruction selection phase via AArch64FunctionInfo::setJumpTableEntryInfo. However, if we try to run a mir file containing jump tables with llc, then it will crash since no elements in JumpTableEntryInfo. With this patch, we don't need to look up JumpTableEntryInfo for element size and a local label will be created if JumpTableEntryInfo is not initialized.

e.g.
asm.ll (https://reviews.llvm.org/F26307823)

llc -mtriple=aarch64 -stop-after=unpack-mi-bundles asm.ll -o asm.mir
llc -mtriple=aarch64 -start-after=unpack-mi-bundles asm.mir -o asm.s

llc: /root/sinan/llvm-project/llvm/include/llvm/ADT/SmallVector.h:298: const T& llvm::SmallVectorTemplateCommon<T, <template-parameter-1-2> >::operator[](llvm::SmallVectorTemplateCommon<T, <template-parameter-1-2> >::size_type) const [with T = std::pair<unsigned int, llvm::MCSymbol*>; <template-parameter-1-2> = void; llvm::SmallVectorTemplateCommon<T, <template-parameter-1-2> >::const_reference = const std::pair<unsigned int, llvm::MCSymbol*>&; llvm::SmallVectorTemplateCommon<T, <template-parameter-1-2> >::size_type = long unsigned int]: Assertion `idx < size()' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.      Program arguments: ./../../../../llvm-dev-install/bin/llc -mtriple=aarch64 -start-after=unpack-mi-bundles asm.mir -o asm.s
1.      Running pass 'Function Pass Manager' on module 'asm.mir'.
2.      Running pass 'AArch64 Assembly Printer' on function '@sum'
 #0 0x00000000019fecfc llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) (./../../../../llvm-dev-install/bin/llc+0x19fecfc)
 #1 0x00000000019fd420 llvm::sys::RunSignalHandlers() (./../../../../llvm-dev-install/bin/llc+0x19fd420)
 #2 0x00000000019fd5e4 SignalHandler(int) Signals.cpp:0:0
 #3 0x0000ffff84f377b0 (linux-vdso.so.1+0x7b0)
 #4 0x0000ffff849e9e98 raise (/lib64/libc.so.6+0x37e98)
 #5 0x0000ffff849d650c abort (/lib64/libc.so.6+0x2450c)
 #6 0x0000ffff849e31e8 __assert_fail_base (/lib64/libc.so.6+0x311e8)
 #7 0x0000ffff849e324c .annobin___GI___assert_fail.end assert.c:0:0

This change has no impact on the normal routine and enables writing mir test with a jump table.

[efriedma]

I'd prefer to make the MIR-parsed codepath as similar as possible to the normal codepath. There are two ways you could do that:

1. Teach yaml::AArch64FunctionInfo to dump/parse JumpTableEntryInfo.
2. Drop the calls to setJumpTableEntryInfo from isel... which I guess implies we can't model the effect of AArch64CompressJumpTables, but otherwise does the right thing.

llvm/test/CodeGen/AArch64/arm64-jumptable.ll

6

I don't think you need a dedicated testcase for this, if you're planning to use this functionality in the near future.

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

                                            const llvm::MachineInstr &MI) {
   Register DestReg = MI.getOperand(0).getReg();
   Register ScratchReg = MI.getOperand(1).getReg();
   Register ScratchRegW =
       STI->getRegisterInfo()->getSubReg(ScratchReg, AArch64::sub_32);
   Register TableReg = MI.getOperand(2).getReg();
   Register EntryReg = MI.getOperand(3).getReg();
   int JTIdx = MI.getOperand(4).getIndex();
-  int Size = AArch64FI->getJumpTableEntrySize(JTIdx);
+
+  int Size;
+  switch (MI.getOpcode()) {
+  case AArch64::JumpTableDest32: Size = 4; break;
+  case AArch64::JumpTableDest16: Size = 2; break;
+  case AArch64::JumpTableDest8:  Size = 1; break;
+  default:
+    llvm_unreachable("Unknown jump table size.");
+  }
 
   // This has to be first because the compression pass based its reachability
   // calculations on the start of the JumpTableDest instruction.
   auto Label =
       MF->getInfo<AArch64FunctionInfo>()->getJumpTableEntryPCRelSymbol(JTIdx);
 
   // If we don't already have a symbol to use as the base, use the ADR
   // instruction itself.

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

#endif
 
   unsigned getSRetReturnReg() const { return SRetReturnReg; }
   void setSRetReturnReg(unsigned Reg) { SRetReturnReg = Reg; }
 
   unsigned getJumpTableEntrySize(int Idx) const {
     return JumpTableEntryInfo[Idx].first;
   }
   MCSymbol *getJumpTableEntryPCRelSymbol(int Idx) const {
-    return JumpTableEntryInfo[Idx].second;
+    return (unsigned)Idx < JumpTableEntryInfo.size()
+               ? JumpTableEntryInfo[Idx].second
+               : nullptr;
   }
   void setJumpTableEntryInfo(int Idx, unsigned Size, MCSymbol *PCRelSym) {
     if ((unsigned)Idx >= JumpTableEntryInfo.size())
       JumpTableEntryInfo.resize(Idx+1);
     JumpTableEntryInfo[Idx] = std::make_pair(Size, PCRelSym);
   }
 
   using SetOfInstructions = SmallPtrSet<const MachineInstr *, 16>;

llvm/test/CodeGen/AArch64/arm64-jumptable.ll

 ; RUN: llc -mtriple=arm64-apple-ios < %s | FileCheck %s
 ; RUN: llc -mtriple=arm64-linux-gnu < %s | FileCheck %s --check-prefix=CHECK-LINUX
 ; <rdar://11417675>
 
+; RUN: llc -mtriple=aarch64 -stop-after=unpack-mi-bundles -o %t.mir %s
+; RUN: llc %t.mir -mtriple=aarch64 -start-after=unpack-mi-bundles -o - | FileCheck %s --check-prefix=CHECK-MIR
+
 define void @sum(i32 %a, ptr %to, i32 %c) {
 entry:
   switch i32 %a, label %exit [
     i32 1, label %bb1
     i32 2, label %exit.sink.split
     i32 3, label %bb3
     i32 4, label %bb4
   ]

 
 ; CHECK-LABEL: sum:
 ; CHECK: adrp    {{x[0-9]+}}, LJTI0_0@PAGE
 ; CHECK:  add    {{x[0-9]+}}, {{x[0-9]+}}, LJTI0_0@PAGEOFF
 
 ; CHECK-LINUX-LABEL: sum:
 ; CHECK-LINUX: adrp    {{x[0-9]+}}, .LJTI0_0
 ; CHECK-LINUX:  add    {{x[0-9]+}}, {{x[0-9]+}}, :lo12:.LJTI0_0
+
+; CHECK-MIR: sum:
+; CHECK-MIR:  adrp  {{x[0-9]+}}, .LJTI0_0
+; CHECK-MIR:  add   {{x[0-9]+}}, {{x[0-9]+}}, :lo12:.LJTI0_0
+; CHECK-MIR: .Ltmp0:
+; CHECK-MIR:  adr   {{x[0-9]+}}, .Ltmp0