llvm/lib/Target/WebAssembly/WebAssemblyFixBrTableDefaults.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-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 }