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