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 25 using namespace llvm; 26 27 #define DEBUG_TYPE "wasm-fix-br-table-defaults" 28 29 namespace { 30 31 class 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 38 public: 39 static char ID; // Pass identification, replacement for typeid 40 WebAssemblyFixBrTableDefaults() : MachineFunctionPass(ID) {} 41 }; 42 43 char WebAssemblyFixBrTableDefaults::ID = 0; 44 45 // Target indepedent 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. 51 void 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(MI.getOperand(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(0).getReg(); 66 assert(MI.getOperand(0).getReg() == ExtDefReg); 67 MI.getOperand(0).setReg(ExtMI->getOperand(1).getReg()); 68 if (MF.getRegInfo().use_nodbg_empty(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(&WebAssembly::I32RegClass); 76 BuildMI(*MBB, MI.getIterator(), MI.getDebugLoc(), 77 WST.getInstrInfo()->get(WebAssembly::I32_WRAP_I64), Reg32) 78 .addReg(MI.getOperand(0).getReg()); 79 MI.getOperand(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(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. 92 MachineBasicBlock *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(*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(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(MI.getNumExplicitOperands() - 1); 135 MI.addOperand(MF, MachineOperand::CreateMBB(TBB)); 136 } 137 138 // Remove any branches from the header and splice in the jump table instead 139 TII.removeBranch(*HeaderMBB, nullptr); 140 HeaderMBB->splice(HeaderMBB->end(), MBB, MBB->begin(), 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(MBB); 145 for (auto &Succ : MBB->successors()) 146 if (HeaderMBB->isSuccessor(Succ)) 147 HeaderMBB->removeSuccessor(Succ); 148 HeaderMBB->transferSuccessorsAndUpdatePHIs(MBB); 149 150 // Remove the old jump table block from the function 151 MF.erase(MBB); 152 153 return HeaderMBB; 154 } 155 156 bool 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 SmallPtrSet<MachineBasicBlock *, 16> MBBSet; 163 for (auto &MBB : MF) 164 MBBSet.insert(&MBB); 165 166 while (!MBBSet.empty()) { 167 MachineBasicBlock *MBB = *MBBSet.begin(); 168 MBBSet.erase(MBB); 169 for (auto &MI : *MBB) { 170 if (WebAssembly::isBrTable(MI.getOpcode())) { 171 fixBrTableIndex(MI, MBB, MF); 172 auto *Fixed = fixBrTableDefault(MI, MBB, MF); 173 if (Fixed != nullptr) { 174 MBBSet.erase(Fixed); 175 Changed = true; 176 } 177 break; 178 } 179 } 180 } 181 182 if (Changed) { 183 // We rewrote part of the function; recompute relevant things. 184 MF.RenumberBlocks(); 185 return true; 186 } 187 188 return false; 189 } 190 191 } // end anonymous namespace 192 193 INITIALIZE_PASS(WebAssemblyFixBrTableDefaults, DEBUG_TYPE, 194 "Removes range checks and sets br_table default targets", false, 195 false) 196 197 FunctionPass *llvm::createWebAssemblyFixBrTableDefaults() { 198 return new WebAssemblyFixBrTableDefaults(); 199 } 200