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