1//=- WebAssemblyFixBrTableDefaults.cpp - Fix br_table default branch targets -//
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//===----------------------------------------------------------------------===//
8///
9/// \file This file implements a pass that eliminates redundant range checks
10/// guarding br_table instructions. Since jump tables on most targets cannot
11/// handle out of range indices, LLVM emits these checks before most jump
12/// tables. But br_table takes a default branch target as an argument, so it
13/// does not need the range checks.
14///
15//===----------------------------------------------------------------------===//
16
17#include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
18#include "WebAssembly.h"
19#include "WebAssemblySubtarget.h"
20#include "llvm/CodeGen/MachineFunction.h"
21#include "llvm/CodeGen/MachineFunctionAnalysisManager.h"
22#include "llvm/CodeGen/MachineFunctionPass.h"
23#include "llvm/CodeGen/MachinePassManager.h"
24#include "llvm/CodeGen/MachineRegisterInfo.h"
25#include "llvm/IR/Analysis.h"
26#include "llvm/Pass.h"
27
28using namespace llvm;
29
30#define DEBUG_TYPE "wasm-fix-br-table-defaults"
31
32namespace {
33
34class WebAssemblyFixBrTableDefaultsLegacy final : public MachineFunctionPass {
35 StringRef getPassName() const override {
36 return "WebAssembly Fix br_table Defaults";
37 }
38
39 bool runOnMachineFunction(MachineFunction &MF) override;
40
41public:
42 static char ID; // Pass identification, replacement for typeid
43 WebAssemblyFixBrTableDefaultsLegacy() : MachineFunctionPass(ID) {}
44};
45
46char WebAssemblyFixBrTableDefaultsLegacy::ID = 0;
47
48// Target independent selection dag assumes that it is ok to use PointerTy
49// as the index for a "switch", whereas Wasm so far only has a 32-bit br_table.
50// See e.g. SelectionDAGBuilder::visitJumpTableHeader
51// We have a 64-bit br_table in the tablegen defs as a result, which does get
52// selected, and thus we get incorrect truncates/extensions happening on
53// wasm64. Here we fix that.
54void fixBrTableIndex(MachineInstr &MI, MachineBasicBlock *MBB,
55 MachineFunction &MF) {
56 // Only happens on wasm64.
57 auto &WST = MF.getSubtarget<WebAssemblySubtarget>();
58 if (!WST.hasAddr64())
59 return;
60
61 assert(MI.getDesc().getOpcode() == WebAssembly::BR_TABLE_I64 &&
62 "64-bit br_table pseudo instruction expected");
63
64 // Find extension op, if any. It sits in the previous BB before the branch.
65 auto ExtMI = MF.getRegInfo().getVRegDef(Reg: MI.getOperand(i: 0).getReg());
66 if (ExtMI->getOpcode() == WebAssembly::I64_EXTEND_U_I32) {
67 // Unnecessarily extending a 32-bit value to 64, remove it.
68 auto ExtDefReg = ExtMI->getOperand(i: 0).getReg();
69 assert(MI.getOperand(0).getReg() == ExtDefReg);
70 MI.getOperand(i: 0).setReg(ExtMI->getOperand(i: 1).getReg());
71 if (MF.getRegInfo().use_nodbg_empty(RegNo: ExtDefReg)) {
72 // No more users of extend, delete it.
73 ExtMI->eraseFromParent();
74 }
75 } else {
76 // Incoming 64-bit value that needs to be truncated.
77 Register Reg32 =
78 MF.getRegInfo().createVirtualRegister(RegClass: &WebAssembly::I32RegClass);
79 BuildMI(BB&: *MBB, I: MI.getIterator(), MIMD: MI.getDebugLoc(),
80 MCID: WST.getInstrInfo()->get(Opcode: WebAssembly::I32_WRAP_I64), DestReg: Reg32)
81 .addReg(RegNo: MI.getOperand(i: 0).getReg());
82 MI.getOperand(i: 0).setReg(Reg32);
83 }
84
85 // We now have a 32-bit operand in all cases, so change the instruction
86 // accordingly.
87 MI.setDesc(WST.getInstrInfo()->get(Opcode: WebAssembly::BR_TABLE_I32));
88}
89
90// `MI` is a br_table instruction with a dummy default target argument. This
91// function finds and adds the default target argument and removes any redundant
92// range check preceding the br_table. Returns the MBB that the br_table is
93// moved into so it can be removed from further consideration, or nullptr if the
94// br_table cannot be optimized.
95MachineBasicBlock *fixBrTableDefault(MachineInstr &MI, MachineBasicBlock *MBB,
96 MachineFunction &MF) {
97 // Get the header block, which contains the redundant range check.
98 assert(MBB->pred_size() == 1 && "Expected a single guard predecessor");
99 auto *HeaderMBB = *MBB->pred_begin();
100
101 // Find the conditional jump to the default target. If it doesn't exist, the
102 // default target is unreachable anyway, so we can keep the existing dummy
103 // target.
104 MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
105 SmallVector<MachineOperand, 2> Cond;
106 const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
107 bool Analyzed = !TII.analyzeBranch(MBB&: *HeaderMBB, TBB, FBB, Cond);
108 assert(Analyzed && "Could not analyze jump header branches");
109 (void)Analyzed;
110
111 // Here are the possible outcomes. '_' is nullptr, `J` is the jump table block
112 // aka MBB, 'D' is the default block.
113 //
114 // TBB | FBB | Meaning
115 // _ | _ | No default block, header falls through to jump table
116 // J | _ | No default block, header jumps to the jump table
117 // D | _ | Header jumps to the default and falls through to the jump table
118 // D | J | Header jumps to the default and also to the jump table
119 if (TBB && TBB != MBB) {
120 assert((FBB == nullptr || FBB == MBB) &&
121 "Expected jump or fallthrough to br_table block");
122 assert(Cond.size() == 2 && Cond[1].isReg() && "Unexpected condition info");
123
124 // If the range check checks an i64 value, we cannot optimize it out because
125 // the i64 index is truncated to an i32, making values over 2^32
126 // indistinguishable from small numbers. There are also other strange edge
127 // cases that can arise in practice that we don't want to reason about, so
128 // conservatively only perform the optimization if the range check is the
129 // normal case of an i32.gt_u.
130 MachineRegisterInfo &MRI = MF.getRegInfo();
131 auto *RangeCheck = MRI.getVRegDef(Reg: Cond[1].getReg());
132 assert(RangeCheck != nullptr);
133 if (RangeCheck->getOpcode() != WebAssembly::GT_U_I32)
134 return nullptr;
135
136 // Remove the dummy default target and install the real one.
137 MI.removeOperand(OpNo: MI.getNumExplicitOperands() - 1);
138 MI.addOperand(MF, Op: MachineOperand::CreateMBB(MBB: TBB));
139 }
140
141 // Remove any branches from the header and splice in the jump table instead
142 TII.removeBranch(MBB&: *HeaderMBB, BytesRemoved: nullptr);
143 HeaderMBB->splice(Where: HeaderMBB->end(), Other: MBB, From: MBB->begin(), To: MBB->end());
144
145 // Update CFG to skip the old jump table block. Remove shared successors
146 // before transferring to avoid duplicated successors.
147 HeaderMBB->removeSuccessor(Succ: MBB);
148 for (auto &Succ : MBB->successors())
149 if (HeaderMBB->isSuccessor(MBB: Succ))
150 HeaderMBB->removeSuccessor(Succ);
151 HeaderMBB->transferSuccessorsAndUpdatePHIs(FromMBB: MBB);
152
153 // Remove the old jump table block from the function
154 MF.erase(MBBI: MBB);
155
156 return HeaderMBB;
157}
158
159bool fixBrTableDefaults(MachineFunction &MF) {
160 LLVM_DEBUG(dbgs() << "********** Fixing br_table Default Targets **********\n"
161 "********** Function: "
162 << MF.getName() << '\n');
163
164 bool Changed = false;
165 SetVector<MachineBasicBlock *, SmallVector<MachineBasicBlock *, 16>,
166 DenseSet<MachineBasicBlock *>, 16>
167 MBBSet;
168 for (auto &MBB : MF)
169 MBBSet.insert(X: &MBB);
170
171 while (!MBBSet.empty()) {
172 MachineBasicBlock *MBB = *MBBSet.begin();
173 MBBSet.remove(X: MBB);
174 for (auto &MI : *MBB) {
175 if (WebAssembly::isBrTable(Opc: MI.getOpcode())) {
176 fixBrTableIndex(MI, MBB, MF);
177 auto *Fixed = fixBrTableDefault(MI, MBB, MF);
178 if (Fixed != nullptr) {
179 MBBSet.remove(X: Fixed);
180 Changed = true;
181 }
182 break;
183 }
184 }
185 }
186
187 if (Changed) {
188 // We rewrote part of the function; recompute relevant things.
189 MF.RenumberBlocks();
190 return true;
191 }
192
193 return false;
194}
195
196} // end anonymous namespace
197
198INITIALIZE_PASS(WebAssemblyFixBrTableDefaultsLegacy, DEBUG_TYPE,
199 "Removes range checks and sets br_table default targets", false,
200 false)
201
202FunctionPass *llvm::createWebAssemblyFixBrTableDefaultsLegacyPass() {
203 return new WebAssemblyFixBrTableDefaultsLegacy();
204}
205
206bool WebAssemblyFixBrTableDefaultsLegacy::runOnMachineFunction(
207 MachineFunction &MF) {
208 return fixBrTableDefaults(MF);
209}
210
211PreservedAnalyses
212WebAssemblyFixBrTableDefaultsPass::run(MachineFunction &MF,
213 MachineFunctionAnalysisManager &MFAM) {
214 return fixBrTableDefaults(MF) ? getMachineFunctionPassPreservedAnalyses()
215 : PreservedAnalyses::all();
216}
217