1 //==-- AArch64CompressJumpTables.cpp - Compress jump tables for AArch64 --====//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 // This pass looks at the basic blocks each jump-table refers to and works out
8 // whether they can be emitted in a compressed form (with 8 or 16-bit
9 // entries). If so, it changes the opcode and flags them in the associated
10 // AArch64FunctionInfo.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #include "AArch64.h"
15 #include "AArch64MachineFunctionInfo.h"
16 #include "AArch64Subtarget.h"
17 #include "llvm/ADT/Statistic.h"
18 #include "llvm/CodeGen/MachineFunctionPass.h"
19 #include "llvm/CodeGen/MachineJumpTableInfo.h"
20 #include "llvm/CodeGen/TargetInstrInfo.h"
21 #include "llvm/CodeGen/TargetSubtargetInfo.h"
22 #include "llvm/MC/MCContext.h"
23 #include "llvm/Support/Alignment.h"
24 #include "llvm/Support/Debug.h"
25
26 using namespace llvm;
27
28 #define DEBUG_TYPE "aarch64-jump-tables"
29
30 STATISTIC(NumJT8, "Number of jump-tables with 1-byte entries");
31 STATISTIC(NumJT16, "Number of jump-tables with 2-byte entries");
32 STATISTIC(NumJT32, "Number of jump-tables with 4-byte entries");
33
34 namespace {
35 class AArch64CompressJumpTables : public MachineFunctionPass {
36 const TargetInstrInfo *TII;
37 MachineFunction *MF;
38 SmallVector<int, 8> BlockInfo;
39
40 /// Returns the size of instructions in the block \p MBB, or std::nullopt if
41 /// we couldn't get a safe upper bound.
42 std::optional<int> computeBlockSize(MachineBasicBlock &MBB);
43
44 /// Gather information about the function, returns false if we can't perform
45 /// this optimization for some reason.
46 bool scanFunction();
47
48 bool compressJumpTable(MachineInstr &MI, int Offset);
49
50 public:
51 static char ID;
AArch64CompressJumpTables()52 AArch64CompressJumpTables() : MachineFunctionPass(ID) {
53 initializeAArch64CompressJumpTablesPass(*PassRegistry::getPassRegistry());
54 }
55
56 bool runOnMachineFunction(MachineFunction &MF) override;
57
getRequiredProperties() const58 MachineFunctionProperties getRequiredProperties() const override {
59 return MachineFunctionProperties().set(
60 MachineFunctionProperties::Property::NoVRegs);
61 }
getPassName() const62 StringRef getPassName() const override {
63 return "AArch64 Compress Jump Tables";
64 }
65 };
66 char AArch64CompressJumpTables::ID = 0;
67 } // namespace
68
69 INITIALIZE_PASS(AArch64CompressJumpTables, DEBUG_TYPE,
70 "AArch64 compress jump tables pass", false, false)
71
72 std::optional<int>
computeBlockSize(MachineBasicBlock & MBB)73 AArch64CompressJumpTables::computeBlockSize(MachineBasicBlock &MBB) {
74 int Size = 0;
75 for (const MachineInstr &MI : MBB) {
76 // Inline asm may contain some directives like .bytes which we don't
77 // currently have the ability to parse accurately. To be safe, just avoid
78 // computing a size and bail out.
79 if (MI.getOpcode() == AArch64::INLINEASM ||
80 MI.getOpcode() == AArch64::INLINEASM_BR)
81 return std::nullopt;
82 Size += TII->getInstSizeInBytes(MI);
83 }
84 return Size;
85 }
86
scanFunction()87 bool AArch64CompressJumpTables::scanFunction() {
88 BlockInfo.clear();
89 BlockInfo.resize(MF->getNumBlockIDs());
90
91 // NOTE: BlockSize, Offset, OffsetAfterAlignment are all upper bounds.
92
93 unsigned Offset = 0;
94 for (MachineBasicBlock &MBB : *MF) {
95 const Align Alignment = MBB.getAlignment();
96 unsigned OffsetAfterAlignment = Offset;
97 // We don't know the exact size of MBB so assume worse case padding.
98 if (Alignment > Align(4))
99 OffsetAfterAlignment += Alignment.value() - 4;
100 BlockInfo[MBB.getNumber()] = OffsetAfterAlignment;
101 auto BlockSize = computeBlockSize(MBB);
102 if (!BlockSize)
103 return false;
104 Offset = OffsetAfterAlignment + *BlockSize;
105 }
106 return true;
107 }
108
compressJumpTable(MachineInstr & MI,int Offset)109 bool AArch64CompressJumpTables::compressJumpTable(MachineInstr &MI,
110 int Offset) {
111 if (MI.getOpcode() != AArch64::JumpTableDest32)
112 return false;
113
114 int JTIdx = MI.getOperand(4).getIndex();
115 auto &JTInfo = *MF->getJumpTableInfo();
116 const MachineJumpTableEntry &JT = JTInfo.getJumpTables()[JTIdx];
117
118 // The jump-table might have been optimized away.
119 if (JT.MBBs.empty())
120 return false;
121
122 int MaxOffset = std::numeric_limits<int>::min(),
123 MinOffset = std::numeric_limits<int>::max();
124 MachineBasicBlock *MinBlock = nullptr;
125 for (auto *Block : JT.MBBs) {
126 int BlockOffset = BlockInfo[Block->getNumber()];
127 assert(BlockOffset % 4 == 0 && "misaligned basic block");
128
129 MaxOffset = std::max(MaxOffset, BlockOffset);
130 if (BlockOffset <= MinOffset) {
131 MinOffset = BlockOffset;
132 MinBlock = Block;
133 }
134 }
135 assert(MinBlock && "Failed to find minimum offset block");
136
137 // The ADR instruction needed to calculate the address of the first reachable
138 // basic block can address +/-1MB.
139 if (!isInt<21>(MinOffset - Offset)) {
140 ++NumJT32;
141 return false;
142 }
143
144 int Span = MaxOffset - MinOffset;
145 auto *AFI = MF->getInfo<AArch64FunctionInfo>();
146 if (isUInt<8>(Span / 4)) {
147 AFI->setJumpTableEntryInfo(JTIdx, 1, MinBlock->getSymbol());
148 MI.setDesc(TII->get(AArch64::JumpTableDest8));
149 ++NumJT8;
150 return true;
151 }
152 if (isUInt<16>(Span / 4)) {
153 AFI->setJumpTableEntryInfo(JTIdx, 2, MinBlock->getSymbol());
154 MI.setDesc(TII->get(AArch64::JumpTableDest16));
155 ++NumJT16;
156 return true;
157 }
158
159 ++NumJT32;
160 return false;
161 }
162
runOnMachineFunction(MachineFunction & MFIn)163 bool AArch64CompressJumpTables::runOnMachineFunction(MachineFunction &MFIn) {
164 bool Changed = false;
165 MF = &MFIn;
166
167 const auto &ST = MF->getSubtarget<AArch64Subtarget>();
168 TII = ST.getInstrInfo();
169
170 if (ST.force32BitJumpTables() && !MF->getFunction().hasMinSize())
171 return false;
172
173 if (!scanFunction())
174 return false;
175
176 for (MachineBasicBlock &MBB : *MF) {
177 int Offset = BlockInfo[MBB.getNumber()];
178 for (MachineInstr &MI : MBB) {
179 Changed |= compressJumpTable(MI, Offset);
180 Offset += TII->getInstSizeInBytes(MI);
181 }
182 }
183
184 return Changed;
185 }
186
createAArch64CompressJumpTablesPass()187 FunctionPass *llvm::createAArch64CompressJumpTablesPass() {
188 return new AArch64CompressJumpTables();
189 }
190