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