Index: llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp =================================================================== --- llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp +++ llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp @@ -86,10 +86,8 @@ void emitStartOfAsmFile(Module &M) override; void emitJumpTableInfo() override; - void emitJumpTableEntry(const MachineJumpTableInfo *MJTI, - const MachineBasicBlock *MBB, unsigned JTI); - void LowerJumpTableDestSmall(MCStreamer &OutStreamer, const MachineInstr &MI); + void LowerJumpTableDest(MCStreamer &OutStreamer, const MachineInstr &MI); void LowerSTACKMAP(MCStreamer &OutStreamer, StackMaps &SM, const MachineInstr &MI); @@ -793,33 +791,24 @@ emitAlignment(Align(Size)); OutStreamer->emitLabel(GetJTISymbol(JTI)); - for (auto *JTBB : JTBBs) - emitJumpTableEntry(MJTI, JTBB, JTI); - } -} - -void AArch64AsmPrinter::emitJumpTableEntry(const MachineJumpTableInfo *MJTI, - const MachineBasicBlock *MBB, - unsigned JTI) { - const MCExpr *Value = MCSymbolRefExpr::create(MBB->getSymbol(), OutContext); - auto AFI = MF->getInfo(); - unsigned Size = AFI->getJumpTableEntrySize(JTI); - - if (Size == 4) { - // .word LBB - LJTI - const TargetLowering *TLI = MF->getSubtarget().getTargetLowering(); - const MCExpr *Base = TLI->getPICJumpTableRelocBaseExpr(MF, JTI, OutContext); - Value = MCBinaryExpr::createSub(Value, Base, OutContext); - } else { - // .byte (LBB - LBB) >> 2 (or .hword) - const MCSymbol *BaseSym = AFI->getJumpTableEntryPCRelSymbol(JTI); + const MCSymbol *BaseSym = AArch64FI->getJumpTableEntryPCRelSymbol(JTI); const MCExpr *Base = MCSymbolRefExpr::create(BaseSym, OutContext); - Value = MCBinaryExpr::createSub(Value, Base, OutContext); - Value = MCBinaryExpr::createLShr( - Value, MCConstantExpr::create(2, OutContext), OutContext); - } - OutStreamer->emitValue(Value, Size); + for (auto *JTBB : JTBBs) { + // Each entry is: + // .byte/.hword (LBB - Lbase)>>2 + // or plain: + // .word LBB - Lbase + const MCExpr *Value = + MCSymbolRefExpr::create(JTBB->getSymbol(), OutContext); + Value = MCBinaryExpr::createSub(Value, Base, OutContext); + if (Size != 4) + Value = MCBinaryExpr::createLShr( + Value, MCConstantExpr::create(2, OutContext), OutContext); + + OutStreamer->emitValue(Value, Size); + } + } } /// Small jump tables contain an unsigned byte or half, representing the offset @@ -830,9 +819,9 @@ /// /// adr xDest, .LBB0_0 /// ldrb wScratch, [xTable, xEntry] (with "lsl #1" for ldrh). -/// add xDest, xDest, xScratch, lsl #2 -void AArch64AsmPrinter::LowerJumpTableDestSmall(llvm::MCStreamer &OutStreamer, - const llvm::MachineInstr &MI) { +/// add xDest, xDest, xScratch (with "lsl #2" for smaller entries) +void AArch64AsmPrinter::LowerJumpTableDest(llvm::MCStreamer &OutStreamer, + const llvm::MachineInstr &MI) { Register DestReg = MI.getOperand(0).getReg(); Register ScratchReg = MI.getOperand(1).getReg(); Register ScratchRegW = @@ -840,33 +829,49 @@ Register TableReg = MI.getOperand(2).getReg(); Register EntryReg = MI.getOperand(3).getReg(); int JTIdx = MI.getOperand(4).getIndex(); - bool IsByteEntry = MI.getOpcode() == AArch64::JumpTableDest8; + int Size = AArch64FI->getJumpTableEntrySize(JTIdx); // This has to be first because the compression pass based its reachability // calculations on the start of the JumpTableDest instruction. - auto Label = - MF->getInfo()->getJumpTableEntryPCRelSymbol(JTIdx); + auto Label = AArch64FI->getJumpTableEntryPCRelSymbol(JTIdx); + + // If we don't already have a symbol to use as the base, use the ADR + // instruction itself. + if (!Label) { + Label = MF->getContext().createTempSymbol(); + AArch64FI->setJumpTableEntryInfo(JTIdx, Size, Label); + OutStreamer.emitLabel(Label); + } + + auto LabelExpr = MCSymbolRefExpr::create(Label, MF->getContext()); EmitToStreamer(OutStreamer, MCInstBuilder(AArch64::ADR) .addReg(DestReg) - .addExpr(MCSymbolRefExpr::create( - Label, MF->getContext()))); + .addExpr(LabelExpr)); // Load the number of instruction-steps to offset from the label. - unsigned LdrOpcode = IsByteEntry ? AArch64::LDRBBroX : AArch64::LDRHHroX; + unsigned LdrOpcode; + switch (Size) { + case 1: LdrOpcode = AArch64::LDRBBroX; break; + case 2: LdrOpcode = AArch64::LDRHHroX; break; + case 4: LdrOpcode = AArch64::LDRSWroX; break; + default: + llvm_unreachable("Unknown jump table size"); + } + EmitToStreamer(OutStreamer, MCInstBuilder(LdrOpcode) - .addReg(ScratchRegW) + .addReg(Size == 4 ? ScratchReg : ScratchRegW) .addReg(TableReg) .addReg(EntryReg) .addImm(0) - .addImm(IsByteEntry ? 0 : 1)); + .addImm(Log2_32(Size))); - // Multiply the steps by 4 and add to the already materialized base label - // address. + // Add to the already materialized base label address, multiplying by 4 if + // compressed. EmitToStreamer(OutStreamer, MCInstBuilder(AArch64::ADDXrs) .addReg(DestReg) .addReg(DestReg) .addReg(ScratchReg) - .addImm(2)); + .addImm(Size == 4 ? 0 : 2)); } void AArch64AsmPrinter::LowerSTACKMAP(MCStreamer &OutStreamer, StackMaps &SM, @@ -1187,30 +1192,10 @@ return; } - case AArch64::JumpTableDest32: { - // We want: - // ldrsw xScratch, [xTable, xEntry, lsl #2] - // add xDest, xTable, xScratch - unsigned DestReg = MI->getOperand(0).getReg(), - ScratchReg = MI->getOperand(1).getReg(), - TableReg = MI->getOperand(2).getReg(), - EntryReg = MI->getOperand(3).getReg(); - EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::LDRSWroX) - .addReg(ScratchReg) - .addReg(TableReg) - .addReg(EntryReg) - .addImm(0) - .addImm(1)); - EmitToStreamer(*OutStreamer, MCInstBuilder(AArch64::ADDXrs) - .addReg(DestReg) - .addReg(TableReg) - .addReg(ScratchReg) - .addImm(0)); - return; - } + case AArch64::JumpTableDest32: case AArch64::JumpTableDest16: case AArch64::JumpTableDest8: - LowerJumpTableDestSmall(*OutStreamer, *MI); + LowerJumpTableDest(*OutStreamer, *MI); return; case AArch64::FMOVH0: Index: llvm/lib/Target/AArch64/AArch64ISelLowering.cpp =================================================================== --- llvm/lib/Target/AArch64/AArch64ISelLowering.cpp +++ llvm/lib/Target/AArch64/AArch64ISelLowering.cpp @@ -6219,6 +6219,9 @@ SDValue Entry = Op.getOperand(2); int JTI = cast(JT.getNode())->getIndex(); + auto *AFI = DAG.getMachineFunction().getInfo(); + AFI->setJumpTableEntryInfo(JTI, 4, nullptr); + SDNode *Dest = DAG.getMachineNode(AArch64::JumpTableDest32, DL, MVT::i64, MVT::i64, JT, Entry, DAG.getTargetJumpTable(JTI, MVT::i32)); Index: llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h =================================================================== --- llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h +++ llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h @@ -285,15 +285,14 @@ void setSRetReturnReg(unsigned Reg) { SRetReturnReg = Reg; } unsigned getJumpTableEntrySize(int Idx) const { - auto It = JumpTableEntryInfo.find(Idx); - if (It != JumpTableEntryInfo.end()) - return It->second.first; - return 4; + return JumpTableEntryInfo[Idx].first; } MCSymbol *getJumpTableEntryPCRelSymbol(int Idx) const { - return JumpTableEntryInfo.find(Idx)->second.second; + return JumpTableEntryInfo[Idx].second; } void setJumpTableEntryInfo(int Idx, unsigned Size, MCSymbol *PCRelSym) { + if ((unsigned)Idx >= JumpTableEntryInfo.size()) + JumpTableEntryInfo.resize(Idx + 1); JumpTableEntryInfo[Idx] = std::make_pair(Size, PCRelSym); } @@ -354,7 +353,7 @@ MILOHContainer LOHContainerSet; SetOfInstructions LOHRelated; - DenseMap> JumpTableEntryInfo; + SmallVector, 2> JumpTableEntryInfo; }; namespace yaml { Index: llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp =================================================================== --- llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp +++ llvm/lib/Target/AArch64/GISel/AArch64InstructionSelector.cpp @@ -2934,6 +2934,8 @@ Register TargetReg = MRI.createVirtualRegister(&AArch64::GPR64RegClass); Register ScratchReg = MRI.createVirtualRegister(&AArch64::GPR64spRegClass); + + MF->getInfo()->setJumpTableEntryInfo(JTI, 4, nullptr); auto JumpTableInst = MIB.buildInstr(AArch64::JumpTableDest32, {TargetReg, ScratchReg}, {JTAddr, Index}) .addJumpTableIndex(JTI); Index: llvm/test/CodeGen/AArch64/jump-table-exynos.ll =================================================================== --- llvm/test/CodeGen/AArch64/jump-table-exynos.ll +++ llvm/test/CodeGen/AArch64/jump-table-exynos.ll @@ -11,7 +11,17 @@ i32 4, label %lbl4 ] ; CHECK-LABEL: test_jumptable: -; CHECK-NOT: ldrb +; CHECK: adrp [[JTPAGE:x[0-9]+]], .LJTI0_0 +; CHECK: add x[[JT:[0-9]+]], [[JTPAGE]], {{#?}}:lo12:.LJTI0_0 +; CHECK: [[PCREL_LBL:.Ltmp.*]]: +; CHECK-NEXT: adr [[PCBASE:x[0-9]+]], [[PCREL_LBL]] +; CHECK: ldrsw x[[OFFSET:[0-9]+]], [x[[JT]], {{x[0-9]+}}, lsl #2] +; CHECK: add [[DEST:x[0-9]+]], [[PCBASE]], x[[OFFSET]] +; CHECK: br [[DEST]] + + +; CHECK: .LJTI0_0: +; CHECK-NEXT: .word .LBB{{.*}}-[[PCREL_LBL]] def: ret i32 0 Index: llvm/test/CodeGen/AArch64/jump-table.ll =================================================================== --- llvm/test/CodeGen/AArch64/jump-table.ll +++ llvm/test/CodeGen/AArch64/jump-table.ll @@ -155,3 +155,56 @@ ; CHECK-IOS-NEXT: .byte (LBB{{.*}}-[[JTBASE]])>>2 ; CHECK-IOS-NEXT: .byte (LBB{{.*}}-[[JTBASE]])>>2 ; CHECK-IOS-NOT: .end_data_region + +; Compressing just the first table has the opportunity to truncate the vector of +; sizes. Make sure it doesn't. +define i32 @test_twotables(i32 %in1, i32 %in2) { +; CHECK-LABEL: test_twotables: +; CHECK: .LJTI2_0 +; CHECK: .LJTI2_1 + + switch i32 %in1, label %def [ + i32 0, label %lbl1 + i32 1, label %lbl2 + i32 2, label %lbl3 + i32 4, label %lbl4 + ] + +def: + ret i32 0 + +lbl1: + ret i32 1 + +lbl2: + ret i32 2 + +lbl3: + ret i32 4 + +lbl4: + switch i32 %in1, label %def [ + i32 0, label %lbl5 + i32 1, label %lbl6 + i32 2, label %lbl7 + i32 4, label %lbl8 + ] + +lbl5: + call i64 @llvm.aarch64.space(i32 262144, i64 undef) + ret i32 1 + +lbl6: + call i64 @llvm.aarch64.space(i32 262144, i64 undef) + ret i32 2 + +lbl7: + call i64 @llvm.aarch64.space(i32 262144, i64 undef) + ret i32 4 +lbl8: + call i64 @llvm.aarch64.space(i32 262144, i64 undef) + ret i32 8 + +} + +declare i64 @llvm.aarch64.space(i32, i64) Index: llvm/test/CodeGen/AArch64/win64-jumptable.ll =================================================================== --- llvm/test/CodeGen/AArch64/win64-jumptable.ll +++ llvm/test/CodeGen/AArch64/win64-jumptable.ll @@ -40,10 +40,10 @@ ; CHECK-NEXT: .seh_endfunclet ; CHECK-NEXT: .p2align 2 ; CHECK-NEXT: .LJTI0_0: -; CHECK: .word .LBB0_2-.LJTI0_0 -; CHECK: .word .LBB0_3-.LJTI0_0 -; CHECK: .word .LBB0_4-.LJTI0_0 -; CHECK: .word .LBB0_5-.LJTI0_0 +; CHECK: .word .LBB0_2-.Ltmp0 +; CHECK: .word .LBB0_3-.Ltmp0 +; CHECK: .word .LBB0_4-.Ltmp0 +; CHECK: .word .LBB0_5-.Ltmp0 ; CHECK: .section .xdata,"dr" ; CHECK: .seh_handlerdata ; CHECK: .text