Use 32bit jump table entries for AArch64

[joerg]

Summary

AArch64 currently uses 64bit jump table entries for non-PIC mode. With this patch, it will use the same encoding for small/medium code models as the current PIC mode does. For other (larger) code models, the function itself is used as relocation base.

[t.p.northover]

Thanks for working on this. I think the approach is good, so just have a couple of minor comments on the implementation:

lib/Target/AArch64/AArch64ISelLowering.cpp

2322

I think we should be explicit about what cases this covers. JITDefault, Kernel & Large? I didn't think we actually supported a Kernel model on AArch64, but I could sort of believe it wanted this.

2350–2351

It looks like getTargetMachine().getSymbol(MF->getFunction()) would be correct (going via getOrCreateSymbol) and save stashing the symbol away in the FunctionInfo.

[joerg]

lib/Target/AArch64/AArch64ISelLowering.cpp

2322

I'm not attached to using default here, enumerating all code models works fine for me, too. The primary concern for me is that we are pessimistic in our assumptions, e.g. whether the 2GB difference works or not.

2350–2351

TM::getSymbol needs access to the mangler. The implementation in AsmPrinter uses getObjFileLowering() for accessing that. It doesn't seem to be simpler that way.

I'm consider whether the symbol should be stashed into the generic MachineFunction data, but that looks like an optimisation for a separate patch.

[t.p.northover]

lib/Target/AArch64/AArch64ISelLowering.cpp

2350–2351

I still think the locality makes it simpler.

Actually, I think it makes little sense for getSymbol to require a mangler. The implementation already goes most of the way towards getting one of its own (it has a TLOF local variable). I think I'll fix that this morning.

lib/Target/AArch64/AArch64AsmPrinter.cpp

public:
 
   void EmitInstruction(const MachineInstr *MI) override;
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AsmPrinter::getAnalysisUsage(AU);
     AU.setPreservesAll();
   }
 
-  bool runOnMachineFunction(MachineFunction &F) override {
-    AArch64FI = F.getInfo<AArch64FunctionInfo>();
-    STI = static_cast<const AArch64Subtarget*>(&F.getSubtarget());
-    return AsmPrinter::runOnMachineFunction(F);
+  bool runOnMachineFunction(MachineFunction &MF) override {
+    AArch64FI = MF.getInfo<AArch64FunctionInfo>();
+    STI = static_cast<const AArch64Subtarget*>(&MF.getSubtarget());
+    SetupMachineFunction(MF);
+    AArch64FI->setCurrentFnSymbol(CurrentFnSym);
+    EmitFunctionBody();
+    return false;
   }
 
 private:
   void printOperand(const MachineInstr *MI, unsigned OpNum, raw_ostream &O);
   bool printAsmMRegister(const MachineOperand &MO, char Mode, raw_ostream &O);
   bool printAsmRegInClass(const MachineOperand &MO,
                           const TargetRegisterClass *RC, bool isVector,
                           raw_ostream &O);

lib/Target/AArch64/AArch64ISelLowering.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_AARCH64_AARCH64ISELLOWERING_H
 #define LLVM_LIB_TARGET_AARCH64_AARCH64ISELLOWERING_H
 
 #include "AArch64.h"
 #include "llvm/CodeGen/CallingConvLower.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/SelectionDAG.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/Target/TargetLowering.h"
 
 namespace llvm {
 
 namespace AArch64ISD {
private:
   SDValue addTokenForArgument(SDValue Chain, SelectionDAG &DAG,
                               MachineFrameInfo &MFI, int ClobberedFI) const;
 
   bool DoesCalleeRestoreStack(CallingConv::ID CallCC, bool TailCallOpt) const;
 
   void saveVarArgRegisters(CCState &CCInfo, SelectionDAG &DAG, const SDLoc &DL,
                            SDValue &Chain) const;
 
+  unsigned getJumpTableEncoding() const override {
+    return MachineJumpTableInfo::EK_LabelDifference32;
+  }
+  SDValue getPICJumpTableRelocBase(SDValue Table,
+                                   SelectionDAG &DAG) const override;
+  const MCExpr *getPICJumpTableRelocBaseExpr(const MachineFunction *MF,
+                                             unsigned JTI,
+                                             MCContext &Ctx) const override;
+
+  bool isJumpTableRelative() const override {
+    return true;
+  }
+
   bool CanLowerReturn(CallingConv::ID CallConv, MachineFunction &MF,
                       bool isVarArg,
                       const SmallVectorImpl<ISD::OutputArg> &Outs,
                       LLVMContext &Context) const override;
 
   SDValue LowerReturn(SDValue Chain, CallingConv::ID CallConv, bool isVarArg,
                       const SmallVectorImpl<ISD::OutputArg> &Outs,
                       const SmallVectorImpl<SDValue> &OutVals, const SDLoc &DL,

lib/Target/AArch64/AArch64ISelLowering.cpp

     return DAG.getNode(ISD::SMIN, dl, Op.getValueType(),
                        Op.getOperand(1), Op.getOperand(2));
   case Intrinsic::aarch64_neon_umin:
     return DAG.getNode(ISD::UMIN, dl, Op.getValueType(),
                        Op.getOperand(1), Op.getOperand(2));
   }
 }
 
+SDValue
+AArch64TargetLowering::getPICJumpTableRelocBase(SDValue Table,
+                                                SelectionDAG &DAG) const {
+  switch (getTargetMachine().getCodeModel()) {
+  case CodeModel::Default:
+  case CodeModel::Small:
+  case CodeModel::Medium:
+    return TargetLowering::getPICJumpTableRelocBase(Table, DAG);
+  default: {
+    MachineFunction &MF = DAG.getMachineFunction();
+    const Function *F = MF.getFunction();
+    SDLoc DL(Table);
+    // Since the target is the function itself, always resolve it locally.
+    unsigned char OpFlags = AArch64II::MO_NO_FLAG;
+    auto PtrVT = getPointerTy(DAG.getDataLayout());
+    SDValue Hi = DAG.getTargetGlobalAddress(F, DL, PtrVT, 0,
+                                            OpFlags | AArch64II::MO_PAGE);
+    unsigned char LoFlags = OpFlags | AArch64II::MO_PAGEOFF | AArch64II::MO_NC;
+    SDValue Lo = DAG.getTargetGlobalAddress(F, DL, PtrVT, 0, LoFlags);
+
+    SDValue ADRP = DAG.getNode(AArch64ISD::ADRP, DL, PtrVT, Hi);
+    return DAG.getNode(AArch64ISD::ADDlow, DL, PtrVT, ADRP, Lo);
+  }
+  }
+}
+
+const MCExpr *
+AArch64TargetLowering::getPICJumpTableRelocBaseExpr(const MachineFunction *MF,
+                                                    unsigned JTI,
+                                                    MCContext &Ctx) const {
+  switch (getTargetMachine().getCodeModel()) {
+  case CodeModel::Default:
+  case CodeModel::Small:
+  case CodeModel::Medium:
+    return TargetLowering::getPICJumpTableRelocBaseExpr(MF, JTI, Ctx);
+  default: {
+    auto *AArch64FI = MF->getInfo<AArch64FunctionInfo>();
+    return MCSymbolRefExpr::create(AArch64FI->getCurrentFnSymbol(), Ctx);
+  }
+  }
+}
+
 SDValue AArch64TargetLowering::LowerOperation(SDValue Op,
                                               SelectionDAG &DAG) const {
   switch (Op.getOpcode()) {
   default:
     llvm_unreachable("unimplemented operand");
     return SDValue();
   case ISD::BITCAST:
     return LowerBITCAST(Op, DAG);

lib/Target/AArch64/AArch64MachineFunctionInfo.h

class AArch64FunctionInfo final : public MachineFunctionInfo {
   /// True when the stack gets realigned dynamically because the size of stack
   /// frame is unknown at compile time. e.g., in case of VLAs.
   bool StackRealigned;
 
   /// True when the callee-save stack area has unused gaps that may be used for
   /// other stack allocations.
   bool CalleeSaveStackHasFreeSpace;
 
+  /// \brief The symbol of the current MachineFunction during emission.
+  MCSymbol *CurrentFnSym;
+
 public:
   AArch64FunctionInfo()
       : BytesInStackArgArea(0), ArgumentStackToRestore(0), HasStackFrame(false),
         NumLocalDynamicTLSAccesses(0), VarArgsStackIndex(0), VarArgsGPRIndex(0),
         VarArgsGPRSize(0), VarArgsFPRIndex(0), VarArgsFPRSize(0),
         IsSplitCSR(false), StackRealigned(false),
-        CalleeSaveStackHasFreeSpace(false) {}
+        CalleeSaveStackHasFreeSpace(false), CurrentFnSym(nullptr) {}
 
   explicit AArch64FunctionInfo(MachineFunction &MF)
       : BytesInStackArgArea(0), ArgumentStackToRestore(0), HasStackFrame(false),
         NumLocalDynamicTLSAccesses(0), VarArgsStackIndex(0), VarArgsGPRIndex(0),
         VarArgsGPRSize(0), VarArgsFPRIndex(0), VarArgsFPRSize(0),
         IsSplitCSR(false), StackRealigned(false),
-        CalleeSaveStackHasFreeSpace(false) {
+        CalleeSaveStackHasFreeSpace(false), CurrentFnSym(nullptr) {
     (void)MF;
   }
 
   unsigned getBytesInStackArgArea() const { return BytesInStackArgArea; }
   void setBytesInStackArgArea(unsigned bytes) { BytesInStackArgArea = bytes; }
 
   unsigned getArgumentStackToRestore() const { return ArgumentStackToRestore; }
   void setArgumentStackToRestore(unsigned bytes) {
public:
   const MILOHContainer &getLOHContainer() const { return LOHContainerSet; }
 
   /// Add a LOH directive of this @p Kind and this @p Args.
   void addLOHDirective(MCLOHType Kind, MILOHArgs Args) {
     LOHContainerSet.push_back(MILOHDirective(Kind, Args));
     LOHRelated.insert(Args.begin(), Args.end());
   }
 
+  MCSymbol *getCurrentFnSymbol() const { return CurrentFnSym; }
+  void setCurrentFnSymbol(MCSymbol *Value) { CurrentFnSym = Value; }
+
 private:
   // Hold the lists of LOHs.
   MILOHContainer LOHContainerSet;
   SetOfInstructions LOHRelated;
 };
 } // End llvm namespace
 
 #endif

test/CodeGen/AArch64/jump-table.ll

 ; RUN: llc -verify-machineinstrs -o - %s -mtriple=aarch64-none-linux-gnu -aarch64-enable-atomic-cfg-tidy=0 | FileCheck %s
 ; RUN: llc -code-model=large -verify-machineinstrs -o - %s -mtriple=aarch64-none-linux-gnu -aarch64-enable-atomic-cfg-tidy=0 | FileCheck --check-prefix=CHECK-LARGE %s
 ; RUN: llc -mtriple=aarch64-none-linux-gnu -verify-machineinstrs -relocation-model=pic -aarch64-enable-atomic-cfg-tidy=0 -o - %s | FileCheck --check-prefix=CHECK-PIC %s
+; XXX: pic+large
 
 define i32 @test_jumptable(i32 %in) {
 ; CHECK: test_jumptable
 
   switch i32 %in, label %def [
     i32 0, label %lbl1
     i32 1, label %lbl2
     i32 2, label %lbl3
     i32 4, label %lbl4
   ]
-; CHECK: adrp [[JTPAGE:x[0-9]+]], .LJTI0_0
-; CHECK: add x[[JT:[0-9]+]], [[JTPAGE]], {{#?}}:lo12:.LJTI0_0
-; CHECK: ldr [[DEST:x[0-9]+]], [x[[JT]], {{x[0-9]+}}, lsl #3]
-; CHECK: br [[DEST]]
-
-; CHECK-LARGE: movz x[[JTADDR:[0-9]+]], #:abs_g3:.LJTI0_0
-; CHECK-LARGE: movk x[[JTADDR]], #:abs_g2_nc:.LJTI0_0
-; CHECK-LARGE: movk x[[JTADDR]], #:abs_g1_nc:.LJTI0_0
-; CHECK-LARGE: movk x[[JTADDR]], #:abs_g0_nc:.LJTI0_0
-; CHECK-LARGE: ldr [[DEST:x[0-9]+]], [x[[JTADDR]], {{x[0-9]+}}, lsl #3]
-; CHECK-LARGE: br [[DEST]]
-
-; CHECK-PIC: adrp [[JTPAGE:x[0-9]+]], .LJTI0_0
-; CHECK-PIC: add x[[JT:[0-9]+]], [[JTPAGE]], {{#?}}:lo12:.LJTI0_0
-; CHECK-PIC: ldrsw [[DEST:x[0-9]+]], [x[[JT]], {{x[0-9]+}}, lsl #2]
-; CHECK-PIC: add [[TABLE:x[0-9]+]], [[DEST]], x[[JT]]
-; CHECK-PIC: br [[TABLE]]
+; CHECK:       adrp [[JTPAGE:x[0-9]+]], .LJTI0_0
+; CHECK:       add  [[JT:x[0-9]+]], [[JTPAGE]], {{#?}}:lo12:.LJTI0_0
+; CHECK-PIC:   adrp [[JTPAGE:x[0-9]+]], .LJTI0_0
+; CHECK-PIC:   add  [[JT:x[0-9]+]], [[JTPAGE]], {{#?}}:lo12:.LJTI0_0
+; CHECK-LARGE: movz [[JT:x[0-9]+]], #:abs_g3:.LJTI0_0
+; CHECK-LARGE: movk [[JT]], #:abs_g2_nc:.LJTI0_0
+; CHECK-LARGE: movk [[JT]], #:abs_g1_nc:.LJTI0_0
+; CHECK-LARGE: movk [[JT]], #:abs_g0_nc:.LJTI0_0
+
+; CHECK:       ldrsw [[JTE:x[0-9]+]], {{\[}}[[JT]], {{x[0-9]+}}, lsl #2]
+; CHECK-PIC:   ldrsw [[JTE:x[0-9]+]], {{\[}}[[JT]], {{x[0-9]+}}, lsl #2]
+; CHECK-LARGE: ldrsw [[JTE:x[0-9]+]], {{\[}}[[JT]], {{x[0-9]+}}, lsl #2]
+
+; CHECK:       add [[DEST:x[0-9]+]], [[JTE]], [[JT]]
+; CHECK-PIC:   add [[DEST:x[0-9]+]], [[JTE]], [[JT]]
+; CHECK-LARGE: adrp [[JTBASE_HI:x[0-9]+]], test_jumptable
+; CHECK-LARGE: add  [[JTBASE:x[0-9]+]], [[JTBASE_HI]], :lo12:test_jumptable
+; CHECK-LARGE: add  [[DEST:x[0-9]+]], [[JTE]], [[JTBASE]]
 
 def:
   ret i32 0
 
 lbl1:
   ret i32 1
 
 lbl2:
   ret i32 2
 
 lbl3:
   ret i32 4
 
 lbl4:
   ret i32 8
 
 }
 
-; CHECK: .rodata
+; CHECK: .section .rodata
+; CHECK-LABEL: .LJTI0_0:
+; CHECK-NEXT: .word .LBB{{.*}}-.LJTI0_0
+; CHECK-NEXT: .word .LBB{{.*}}-.LJTI0_0
+; CHECK-NEXT: .word .LBB{{.*}}-.LJTI0_0
+; CHECK-NEXT: .word .LBB{{.*}}-.LJTI0_0
+; CHECK-NEXT: .word .LBB{{.*}}-.LJTI0_0
+; CHECK-NOT: .end_data_region
 
-; CHECK: .LJTI0_0:
-; CHECK-NEXT: .xword
-; CHECK-NEXT: .xword
-; CHECK-NEXT: .xword
-; CHECK-NEXT: .xword
-; CHECK-NEXT: .xword
-
-; CHECK-PIC-NOT: .data_region
-; CHECK-PIC-NOT: .LJTI0_0
-; CHECK-PIC: .LJTI0_0:
+; CHECK-PIC: .section .rodata
+; CHECK-PIC-LABEL: .LJTI0_0:
 ; CHECK-PIC-NEXT: .word .LBB{{.*}}-.LJTI0_0
 ; CHECK-PIC-NEXT: .word .LBB{{.*}}-.LJTI0_0
 ; CHECK-PIC-NEXT: .word .LBB{{.*}}-.LJTI0_0
 ; CHECK-PIC-NEXT: .word .LBB{{.*}}-.LJTI0_0
 ; CHECK-PIC-NEXT: .word .LBB{{.*}}-.LJTI0_0
 ; CHECK-PIC-NOT: .end_data_region
+
+; CHECK-LARGE: .section .rodata
+; CHECK-LARGE-LABEL: .LJTI0_0:
+; CHECK-LARGE-NEXT: .word .LBB{{.*}}-test_jumptable
+; CHECK-LARGE-NEXT: .word .LBB{{.*}}-test_jumptable
+; CHECK-LARGE-NEXT: .word .LBB{{.*}}-test_jumptable
+; CHECK-LARGE-NEXT: .word .LBB{{.*}}-test_jumptable
+; CHECK-LARGE-NEXT: .word .LBB{{.*}}-test_jumptable
+; CHECK-LARGE-NOT: .end_data_region