[X86] Use 64-bit jump table entries for large code model PIC

[aeubanks]

Summary

With the large code model, the label difference may not fit into 32 bits.
Even if we assume that any individual function is no larger than 2^32
and use a difference from the function entry to the target destination,
things like BOLT can rearrange blocks (even if BOLT doesn't necessarily
work with the large code model right now).

set directives avoid static relocations in some 32-bit entry cases, but
don't worry about set directives for 64-bit jump table entries (we can
do that later if somebody really cares about it).

check-llvm in a bootstrapped clang with the large code model passes.

Fixes #62894

[rnk]

lgtm

[efriedma]

If offsets within a function are larger than 2GB, we have other issues (e.g. branch relaxation for branches with a +-2GB offset encoding).

But this seems okay as a minima fix just to get things working.

llvm/include/llvm/CodeGen/MIRYamlMapping.h

static void enumeration(yaml::IO &IO,
     IO.enumCase(EntryKind, "block-address",
                 MachineJumpTableInfo::EK_BlockAddress);
     IO.enumCase(EntryKind, "gp-rel64-block-address",
                 MachineJumpTableInfo::EK_GPRel64BlockAddress);
     IO.enumCase(EntryKind, "gp-rel32-block-address",
                 MachineJumpTableInfo::EK_GPRel32BlockAddress);
     IO.enumCase(EntryKind, "label-difference32",
                 MachineJumpTableInfo::EK_LabelDifference32);
+    IO.enumCase(EntryKind, "label-difference64",
+                MachineJumpTableInfo::EK_LabelDifference64);
     IO.enumCase(EntryKind, "inline", MachineJumpTableInfo::EK_Inline);
     IO.enumCase(EntryKind, "custom32", MachineJumpTableInfo::EK_Custom32);
   }
 };
 
 template <> struct ScalarTraits<MaybeAlign> {
   static void output(const MaybeAlign &Alignment, void *,
                      llvm::raw_ostream &out) {

llvm/include/llvm/CodeGen/MachineJumpTableInfo.h

enum JTEntryKind {
     /// the address of the jump table.  This is used for PIC jump tables where
     /// gprel32 is not supported.  e.g.:
     ///      .word LBB123 - LJTI1_2
     /// If the .set directive is supported, this is emitted as:
     ///      .set L4_5_set_123, LBB123 - LJTI1_2
     ///      .word L4_5_set_123
     EK_LabelDifference32,
 
+    /// EK_LabelDifference64 - Each entry is the address of the block minus
+    /// the address of the jump table.  This is used for PIC jump tables where
+    /// gprel64 is not supported.  e.g.:
+    ///      .quad LBB123 - LJTI1_2
+    EK_LabelDifference64,
+
     /// EK_Inline - Jump table entries are emitted inline at their point of
     /// use. It is the responsibility of the target to emit the entries.
     EK_Inline,
 
     /// EK_Custom32 - Each entry is a 32-bit value that is custom lowered by the
     /// TargetLowering::LowerCustomJumpTableEntry hook.
     EK_Custom32
   };
+
 private:
   JTEntryKind EntryKind;
   std::vector<MachineJumpTableEntry> JumpTables;
 public:
   explicit MachineJumpTableInfo(JTEntryKind Kind): EntryKind(Kind) {}
 
   JTEntryKind getEntryKind() const { return EntryKind; }
 

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

void AsmPrinter::emitJumpTableInfo() {
   const std::vector<MachineJumpTableEntry> &JT = MJTI->getJumpTables();
   if (JT.empty()) return;
 
   // Pick the directive to use to print the jump table entries, and switch to
   // the appropriate section.
   const Function &F = MF->getFunction();
   const TargetLoweringObjectFile &TLOF = getObjFileLowering();
   bool JTInDiffSection = !TLOF.shouldPutJumpTableInFunctionSection(
-      MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32,
+      MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 ||
+          MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference64,
       F);
   if (JTInDiffSection) {
     // Drop it in the readonly section.
     MCSection *ReadOnlySection = TLOF.getSectionForJumpTable(F, TM);
     OutStreamer->switchSection(ReadOnlySection);
   }
 
   emitAlignment(Align(MJTI->getEntryAlignment(DL)));
case MachineJumpTableInfo::EK_GPRel64BlockAddress: {
     // EK_GPRel64BlockAddress - Each entry is an address of block, encoded
     // with a relocation as gp-relative, e.g.:
     //     .gpdword LBB123
     MCSymbol *MBBSym = MBB->getSymbol();
     OutStreamer->emitGPRel64Value(MCSymbolRefExpr::create(MBBSym, OutContext));
     return;
   }
 
-  case MachineJumpTableInfo::EK_LabelDifference32: {
+  case MachineJumpTableInfo::EK_LabelDifference32:
+  case MachineJumpTableInfo::EK_LabelDifference64: {
     // Each entry is the address of the block minus the address of the jump
     // table. This is used for PIC jump tables where gprel32 is not supported.
     // e.g.:
     //      .word LBB123 - LJTI1_2
     // If the .set directive avoids relocations, this is emitted as:
     //      .set L4_5_set_123, LBB123 - LJTI1_2
     //      .word L4_5_set_123
-    if (MAI->doesSetDirectiveSuppressReloc()) {
+    if (MJTI->getEntryKind() == MachineJumpTableInfo::EK_LabelDifference32 &&
+        MAI->doesSetDirectiveSuppressReloc()) {
       Value = MCSymbolRefExpr::create(GetJTSetSymbol(UID, MBB->getNumber()),
                                       OutContext);
       break;
     }
     Value = MCSymbolRefExpr::create(MBB->getSymbol(), OutContext);
     const TargetLowering *TLI = MF->getSubtarget().getTargetLowering();
     const MCExpr *Base = TLI->getPICJumpTableRelocBaseExpr(MF, UID, OutContext);
     Value = MCBinaryExpr::createSub(Value, Base, OutContext);

llvm/lib/CodeGen/AsmPrinter/CodeViewDebug.cpp

forEachJumpTableBranch(
               "EK_GPRel64BlockAddress should never be emitted for COFF");
         case MachineJumpTableInfo::EK_BlockAddress:
           // Each entry is an absolute address.
           EntrySize = JumpTableEntrySize::Pointer;
           Base = nullptr;
           break;
         case MachineJumpTableInfo::EK_Inline:
         case MachineJumpTableInfo::EK_LabelDifference32:
+        case MachineJumpTableInfo::EK_LabelDifference64:
           // Ask the AsmPrinter.
           std::tie(Base, BaseOffset, Branch, EntrySize) =
               Asm->getCodeViewJumpTableInfo(JumpTableIndex, &BranchMI, Branch);
           break;
         }
 
         CurFn->JumpTables.push_back(
             {EntrySize, Base, BaseOffset, Branch,

llvm/lib/CodeGen/MachineFunction.cpp

 /// Return the size of each entry in the jump table.
 unsigned MachineJumpTableInfo::getEntrySize(const DataLayout &TD) const {
   // The size of a jump table entry is 4 bytes unless the entry is just the
   // address of a block, in which case it is the pointer size.
   switch (getEntryKind()) {
   case MachineJumpTableInfo::EK_BlockAddress:
     return TD.getPointerSize();
   case MachineJumpTableInfo::EK_GPRel64BlockAddress:
+  case MachineJumpTableInfo::EK_LabelDifference64:
     return 8;
   case MachineJumpTableInfo::EK_GPRel32BlockAddress:
   case MachineJumpTableInfo::EK_LabelDifference32:
   case MachineJumpTableInfo::EK_Custom32:
     return 4;
   case MachineJumpTableInfo::EK_Inline:
     return 0;
   }
   llvm_unreachable("Unknown jump table encoding!");
 }
 
 /// Return the alignment of each entry in the jump table.
 unsigned MachineJumpTableInfo::getEntryAlignment(const DataLayout &TD) const {
   // The alignment of a jump table entry is the alignment of int32 unless the
   // entry is just the address of a block, in which case it is the pointer
   // alignment.
   switch (getEntryKind()) {
   case MachineJumpTableInfo::EK_BlockAddress:
     return TD.getPointerABIAlignment(0).value();
   case MachineJumpTableInfo::EK_GPRel64BlockAddress:
+  case MachineJumpTableInfo::EK_LabelDifference64:
     return TD.getABIIntegerTypeAlignment(64).value();
   case MachineJumpTableInfo::EK_GPRel32BlockAddress:
   case MachineJumpTableInfo::EK_LabelDifference32:
   case MachineJumpTableInfo::EK_Custom32:
     return TD.getABIIntegerTypeAlignment(32).value();
   case MachineJumpTableInfo::EK_Inline:
     return 1;
   }

llvm/lib/Target/X86/X86ISelLoweringCall.cpp

 /// Return the entry encoding for a jump table in the
 /// current function.  The returned value is a member of the
 /// MachineJumpTableInfo::JTEntryKind enum.
 unsigned X86TargetLowering::getJumpTableEncoding() const {
   // In GOT pic mode, each entry in the jump table is emitted as a @GOTOFF
   // symbol.
   if (isPositionIndependent() && Subtarget.isPICStyleGOT())
     return MachineJumpTableInfo::EK_Custom32;
+  if (isPositionIndependent() &&
+      getTargetMachine().getCodeModel() == CodeModel::Large)
+    return MachineJumpTableInfo::EK_LabelDifference64;
 
   // Otherwise, use the normal jump table encoding heuristics.
   return TargetLowering::getJumpTableEncoding();
 }
 
 bool X86TargetLowering::splitValueIntoRegisterParts(
     SelectionDAG &DAG, const SDLoc &DL, SDValue Val, SDValue *Parts,
     unsigned NumParts, MVT PartVT, std::optional<CallingConv::ID> CC) const {

 }
 
 /// This returns the relocation base for the given PIC jumptable,
 /// the same as getPICJumpTableRelocBase, but as an MCExpr.
 const MCExpr *X86TargetLowering::
 getPICJumpTableRelocBaseExpr(const MachineFunction *MF, unsigned JTI,
                              MCContext &Ctx) const {
   // X86-64 uses RIP relative addressing based on the jump table label.
-  if (Subtarget.isPICStyleRIPRel())
+  if (Subtarget.isPICStyleRIPRel() ||
+      (Subtarget.is64Bit() &&
+       getTargetMachine().getCodeModel() == CodeModel::Large))
     return TargetLowering::getPICJumpTableRelocBaseExpr(MF, JTI, Ctx);
 
   // Otherwise, the reference is relative to the PIC base.
   return MCSymbolRefExpr::create(MF->getPICBaseSymbol(), Ctx);
 }
 
 std::pair<const TargetRegisterClass *, uint8_t>
 X86TargetLowering::findRepresentativeClass(const TargetRegisterInfo *TRI,

llvm/test/CodeGen/X86/large-pic-jump-table.ll

+; RUN: llc -code-model=large -relocation-model=pic %s -o - | FileCheck %s
+
+target triple = "x86_64-linux-gnu"
+
+define i32 @f(i32 %i) {
+bb:
+  switch i32 %i, label %bb4 [
+    i32 0, label %bb0
+    i32 1, label %bb1
+    i32 2, label %bb2
+    i32 3, label %bb3
+  ]
+
+bb0:
+  ret i32 1
+
+bb1:
+  ret i32 100
+
+bb2:
+  ret i32 200
+
+bb3:
+  ret i32 400
+
+bb4:
+  ret i32 300
+}
+
+; CHECK:      movabsq $.LJTI0_0@GOTOFF, [[R1:%r[a-z]{2}]]
+; CHECK-NEXT: addq    [[R1]], [[R2:%r[a-z]{2}]]
+; CHECK-NEXT: addq    ([[R2]],[[R3:%r[a-z]{2}]],8), [[R2]]
+; CHECK-NEXT: jmpq    *[[R2]]
+
+; CHECK: .LJTI0_0:
+; CHECK-NEXT: .quad   .LBB0_2-.LJTI0_0
+; CHECK-NEXT: .quad   .LBB0_3-.LJTI0_0
+; CHECK-NEXT: .quad   .LBB0_4-.LJTI0_0
+; CHECK-NEXT: .quad   .LBB0_5-.LJTI0_0