202 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8///9/// \file This file implements a pass that eliminates redundant range checks10/// guarding br_table instructions. Since jump tables on most targets cannot11/// handle out of range indices, LLVM emits these checks before most jump12/// tables. But br_table takes a default branch target as an argument, so it13/// 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 typeid40 WebAssemblyFixBrTableDefaults() : MachineFunctionPass(ID) {}41};42 43char WebAssemblyFixBrTableDefaults::ID = 0;44 45// Target independent selection dag assumes that it is ok to use PointerTy46// as the index for a "switch", whereas Wasm so far only has a 32-bit br_table.47// See e.g. SelectionDAGBuilder::visitJumpTableHeader48// We have a 64-bit br_table in the tablegen defs as a result, which does get49// selected, and thus we get incorrect truncates/extensions happening on50// 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(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 instruction83 // 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. This88// function finds and adds the default target argument and removes any redundant89// range check preceding the br_table. Returns the MBB that the br_table is90// moved into so it can be removed from further consideration, or nullptr if the91// 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, the99 // default target is unreachable anyway, so we can keep the existing dummy100 // 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 block109 // aka MBB, 'D' is the default block.110 //111 // TBB | FBB | Meaning112 // _ | _ | No default block, header falls through to jump table113 // J | _ | No default block, header jumps to the jump table114 // D | _ | Header jumps to the default and falls through to the jump table115 // D | J | Header jumps to the default and also to the jump table116 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 because122 // the i64 index is truncated to an i32, making values over 2^32123 // indistinguishable from small numbers. There are also other strange edge124 // cases that can arise in practice that we don't want to reason about, so125 // conservatively only perform the optimization if the range check is the126 // 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 instead139 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 successors143 // 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 function151 MF.erase(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(&MBB);167 168 while (!MBBSet.empty()) {169 MachineBasicBlock *MBB = *MBBSet.begin();170 MBBSet.remove(MBB);171 for (auto &MI : *MBB) {172 if (WebAssembly::isBrTable(MI.getOpcode())) {173 fixBrTableIndex(MI, MBB, MF);174 auto *Fixed = fixBrTableDefault(MI, MBB, MF);175 if (Fixed != nullptr) {176 MBBSet.remove(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 namespace194 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