185 lines · cpp
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-exception6//7// This pass looks at the basic blocks each jump-table refers to and works out8// whether they can be emitted in a compressed form (with 8 or 16-bit9// entries). If so, it changes the opcode and flags them in the associated10// 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 if39 /// 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 perform43 /// 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} // namespace63 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't72 // currently have the ability to parse accurately. To be safe, just avoid73 // 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(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(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(MaxOffset, 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 reachable133 // basic block can address +/-1MB.134 if (!isInt<21>(MinOffset - Offset)) {135 ++NumJT32;136 return false;137 }138 139 int Span = MaxOffset - MinOffset;140 auto *AFI = MF->getInfo<AArch64FunctionInfo>();141 if (isUInt<8>(Span / 4)) {142 AFI->setJumpTableEntryInfo(JTIdx, 1, MinBlock->getSymbol());143 MI.setDesc(TII->get(AArch64::JumpTableDest8));144 ++NumJT8;145 return true;146 }147 if (isUInt<16>(Span / 4)) {148 AFI->setJumpTableEntryInfo(JTIdx, 2, MinBlock->getSymbol());149 MI.setDesc(TII->get(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