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 {
31   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
39   WebAssemblyFixBrTableDefaults() : MachineFunctionPass(ID) {}
40 };
41 
42 char WebAssemblyFixBrTableDefaults::ID = 0;
43 
44 // `MI` is a br_table instruction with a dummy default target argument. This
45 // function finds and adds the default target argument and removes any redundant
46 // range check preceding the br_table. Returns the MBB that the br_table is
47 // moved into so it can be removed from further consideration, or nullptr if the
48 // br_table cannot be optimized.
49 MachineBasicBlock *fixBrTable(MachineInstr &MI, MachineBasicBlock *MBB,
50                               MachineFunction &MF) {
51   // Get the header block, which contains the redundant range check.
52   assert(MBB->pred_size() == 1 && "Expected a single guard predecessor");
53   auto *HeaderMBB = *MBB->pred_begin();
54 
55   // Find the conditional jump to the default target. If it doesn't exist, the
56   // default target is unreachable anyway, so we can keep the existing dummy
57   // target.
58   MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
59   SmallVector<MachineOperand, 2> Cond;
60   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
61   bool Analyzed = !TII.analyzeBranch(*HeaderMBB, TBB, FBB, Cond);
62   assert(Analyzed && "Could not analyze jump header branches");
63 
64   // Here are the possible outcomes. '_' is nullptr, `J` is the jump table block
65   // aka MBB, 'D' is the default block.
66   //
67   // TBB | FBB | Meaning
68   //  _  |  _  | No default block, header falls through to jump table
69   //  J  |  _  | No default block, header jumps to the jump table
70   //  D  |  _  | Header jumps to the default and falls through to the jump table
71   //  D  |  J  | Header jumps to the default and also to the jump table
72   if (TBB && TBB != MBB) {
73     assert((FBB == nullptr || FBB == MBB) &&
74            "Expected jump or fallthrough to br_table block");
75     assert(Cond.size() == 2 && Cond[1].isReg() && "Unexpected condition info");
76 
77     // If the range check checks an i64 value, we cannot optimize it out because
78     // the i64 index is truncated to an i32, making values over 2^32
79     // indistinguishable from small numbers. There are also other strange edge
80     // cases that can arise in practice that we don't want to reason about, so
81     // conservatively only perform the optimization if the range check is the
82     // normal case of an i32.gt_u.
83     MachineRegisterInfo &MRI = MF.getRegInfo();
84     auto *RangeCheck = MRI.getVRegDef(Cond[1].getReg());
85     assert(RangeCheck != nullptr);
86     if (RangeCheck->getOpcode() != WebAssembly::GT_U_I32)
87       return nullptr;
88 
89     // Remove the dummy default target and install the real one.
90     MI.RemoveOperand(MI.getNumExplicitOperands() - 1);
91     MI.addOperand(MF, MachineOperand::CreateMBB(TBB));
92   }
93 
94   // Remove any branches from the header and splice in the jump table instead
95   TII.removeBranch(*HeaderMBB, nullptr);
96   HeaderMBB->splice(HeaderMBB->end(), MBB, MBB->begin(), MBB->end());
97 
98   // Update CFG to skip the old jump table block. Remove shared successors
99   // before transferring to avoid duplicated successors.
100   HeaderMBB->removeSuccessor(MBB);
101   for (auto &Succ : MBB->successors())
102     if (HeaderMBB->isSuccessor(Succ))
103       HeaderMBB->removeSuccessor(Succ);
104   HeaderMBB->transferSuccessorsAndUpdatePHIs(MBB);
105 
106   // Remove the old jump table block from the function
107   MF.erase(MBB);
108 
109   return HeaderMBB;
110 }
111 
112 bool WebAssemblyFixBrTableDefaults::runOnMachineFunction(MachineFunction &MF) {
113   LLVM_DEBUG(dbgs() << "********** Fixing br_table Default Targets **********\n"
114                        "********** Function: "
115                     << MF.getName() << '\n');
116 
117   bool Changed = false;
118   SmallPtrSet<MachineBasicBlock *, 16> MBBSet;
119   for (auto &MBB : MF)
120     MBBSet.insert(&MBB);
121 
122   while (!MBBSet.empty()) {
123     MachineBasicBlock *MBB = *MBBSet.begin();
124     MBBSet.erase(MBB);
125     for (auto &MI : *MBB) {
126       if (WebAssembly::isBrTable(MI)) {
127         auto *Fixed = fixBrTable(MI, MBB, MF);
128         if (Fixed != nullptr) {
129           MBBSet.erase(Fixed);
130           Changed = true;
131         }
132         break;
133       }
134     }
135   }
136 
137   if (Changed) {
138     // We rewrote part of the function; recompute relevant things.
139     MF.RenumberBlocks();
140     return true;
141   }
142 
143   return false;
144 }
145 
146 } // end anonymous namespace
147 
148 INITIALIZE_PASS(WebAssemblyFixBrTableDefaults, DEBUG_TYPE,
149                 "Removes range checks and sets br_table default targets", false,
150                 false)
151 
152 FunctionPass *llvm::createWebAssemblyFixBrTableDefaults() {
153   return new WebAssemblyFixBrTableDefaults();
154 }
155