[WebAssembly] Fixed 64-bit indices in br_table

[aardappel]

Summary

LLVM selection dag assumes "switch" indices are pointer sized, which causes problems for our 32-bit br_table. The new function ensures 32-bit operands don't get unnecessarily extended, and 64-bit operands get truncated.

Note that the changes to the existing test test exactly that: the addition of -NEXT in 2 places ensures no extension is inserted (which the test previously ignored) and that the wrap is present (previously omitted in wasm64 mode).

[aardappel]

Replace fragile way of finding the extend instruction by just getting the def of the br_table operand.

[tlively]

LGTM!

llvm/lib/Target/WebAssembly/WebAssemblyFixBrTableDefaults.cpp

 
 public:
   static char ID; // Pass identification, replacement for typeid
   WebAssemblyFixBrTableDefaults() : MachineFunctionPass(ID) {}
 };
 
 char WebAssemblyFixBrTableDefaults::ID = 0;
 
+// Target indepedent selection dag assumes that it is ok to use PointerTy
+// as the index for a "switch", whereas Wasm so far only has a 32-bit br_table.
+// See e.g. SelectionDAGBuilder::visitJumpTableHeader
+// We have a 64-bit br_table in the tablegen defs as a result, which does get
+// selected, and thus we get incorrect truncates/extensions happening on
+// wasm64. Here we fix that.
+void fixBrTableIndex(MachineInstr &MI, MachineBasicBlock *MBB,
+                     MachineFunction &MF) {
+  // Only happens on wasm64.
+  auto &WST = MF.getSubtarget<WebAssemblySubtarget>();
+  if (!WST.hasAddr64())
+    return;
+
+  assert(MI.getDesc().getOpcode() == WebAssembly::BR_TABLE_I64 &&
+         "64-bit br_table pseudo instruction expected");
+
+  // Find extension op, if any. It sits in the previous BB before the branch.
+  auto ExtMI = MF.getRegInfo().getVRegDef(MI.getOperand(0).getReg());
+  if (ExtMI->getOpcode() == WebAssembly::I64_EXTEND_U_I32) {
+    // Unnecessarily extending a 32-bit value to 64, remove it.
+    assert(MI.getOperand(0).getReg() == ExtMI->getOperand(0).getReg());
+    MI.getOperand(0).setReg(ExtMI->getOperand(1).getReg());
+    ExtMI->eraseFromParent();
+  } else {
+    // Incoming 64-bit value that needs to be truncated.
+    Register Reg32 =
+        MF.getRegInfo().createVirtualRegister(&WebAssembly::I32RegClass);
+    BuildMI(*MBB, MI.getIterator(), MI.getDebugLoc(),
+            WST.getInstrInfo()->get(WebAssembly::I32_WRAP_I64), Reg32)
+        .addReg(MI.getOperand(0).getReg());
+    MI.getOperand(0).setReg(Reg32);
+  }
+
+  // We now have a 32-bit operand in all cases, so change the instruction
+  // accordingly.
+  MI.setDesc(WST.getInstrInfo()->get(WebAssembly::BR_TABLE_I32));
+}
+
 // `MI` is a br_table instruction with a dummy default target argument. This
 // function finds and adds the default target argument and removes any redundant
 // range check preceding the br_table. Returns the MBB that the br_table is
 // moved into so it can be removed from further consideration, or nullptr if the
 // br_table cannot be optimized.
-MachineBasicBlock *fixBrTable(MachineInstr &MI, MachineBasicBlock *MBB,
+MachineBasicBlock *fixBrTableDefault(MachineInstr &MI, MachineBasicBlock *MBB,
                                      MachineFunction &MF) {
   // Get the header block, which contains the redundant range check.
   assert(MBB->pred_size() == 1 && "Expected a single guard predecessor");
   auto *HeaderMBB = *MBB->pred_begin();
 
   // Find the conditional jump to the default target. If it doesn't exist, the
   // default target is unreachable anyway, so we can keep the existing dummy
   // target.
   MachineBasicBlock *TBB = nullptr, *FBB = nullptr;
bool WebAssemblyFixBrTableDefaults::runOnMachineFunction(MachineFunction &MF) {
   for (auto &MBB : MF)
     MBBSet.insert(&MBB);
 
   while (!MBBSet.empty()) {
     MachineBasicBlock *MBB = *MBBSet.begin();
     MBBSet.erase(MBB);
     for (auto &MI : *MBB) {
       if (WebAssembly::isBrTable(MI)) {
-        auto *Fixed = fixBrTable(MI, MBB, MF);
+        fixBrTableIndex(MI, MBB, MF);
+        auto *Fixed = fixBrTableDefault(MI, MBB, MF);
         if (Fixed != nullptr) {
           MBBSet.erase(Fixed);
           Changed = true;
         }
         break;
       }
     }
   }

llvm/test/CodeGen/WebAssembly/switch.ll

-; RUN: llc < %s -asm-verbose=false -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -disable-block-placement -verify-machineinstrs | FileCheck %s
+; RUN: llc < %s --mtriple=wasm32-unknown-unknown -asm-verbose=false -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -disable-block-placement -verify-machineinstrs | FileCheck %s
+; RUN: llc < %s --mtriple=wasm64-unknown-unknown -asm-verbose=false -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -disable-block-placement -verify-machineinstrs | FileCheck %s
 
 ; Test switch instructions. Block placement is disabled because it reorders
 ; the blocks in a way that isn't interesting here.
 
-target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
-target triple = "wasm32-unknown-unknown"
-
 declare void @foo0()
 declare void @foo1()
 declare void @foo2()
 declare void @foo3()
 declare void @foo4()
 declare void @foo5()
 
 ; CHECK-LABEL: bar32:
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: br_table {{[^,]+}}, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 4, 5, 6{{$}}
-; CHECK: .LBB{{[0-9]+}}_1:
-; CHECK:   call foo0{{$}}
-; CHECK: .LBB{{[0-9]+}}_2:
-; CHECK:   call foo1{{$}}
-; CHECK: .LBB{{[0-9]+}}_3:
-; CHECK:   call foo2{{$}}
-; CHECK: .LBB{{[0-9]+}}_4:
-; CHECK:   call foo3{{$}}
-; CHECK: .LBB{{[0-9]+}}_5:
-; CHECK:   call foo4{{$}}
-; CHECK: .LBB{{[0-9]+}}_6:
-; CHECK:   call foo5{{$}}
-; CHECK: .LBB{{[0-9]+}}_7:
-; CHECK:   return{{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK-NEXT: br_table {{[^,]+}}, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 4, 5, 6{{$}}
+; CHECK:      .LBB{{[0-9]+}}_1:
+; CHECK:        call foo0{{$}}
+; CHECK:      .LBB{{[0-9]+}}_2:
+; CHECK:        call foo1{{$}}
+; CHECK:      .LBB{{[0-9]+}}_3:
+; CHECK:        call foo2{{$}}
+; CHECK:      .LBB{{[0-9]+}}_4:
+; CHECK:        call foo3{{$}}
+; CHECK:      .LBB{{[0-9]+}}_5:
+; CHECK:        call foo4{{$}}
+; CHECK:      .LBB{{[0-9]+}}_6:
+; CHECK:        call foo5{{$}}
+; CHECK:      .LBB{{[0-9]+}}_7:
+; CHECK:        return{{$}}
 define void @bar32(i32 %n) {
 entry:
   switch i32 %n, label %sw.epilog [
     i32 0, label %sw.bb
     i32 1, label %sw.bb
     i32 2, label %sw.bb
     i32 3, label %sw.bb
     i32 4, label %sw.bb
sw.bb.5:                                          ; preds = %entry
   tail call void @foo5()
   br label %sw.epilog
 
 sw.epilog:                                        ; preds = %entry, %sw.bb.5, %sw.bb.4, %sw.bb.3, %sw.bb.2, %sw.bb.1, %sw.bb
   ret void
 }
 
 ; CHECK-LABEL: bar64:
-; CHECK: block   {{$}}
-; CHECK: i64.const
-; CHECK: i64.gt_u
-; CHECK: br_if 0
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: block   {{$}}
-; CHECK: i32.wrap_i64
-; CHECK: br_table {{[^,]+}}, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 4, 5, 0{{$}}
-; CHECK: .LBB{{[0-9]+}}_2:
-; CHECK:   call foo0{{$}}
-; CHECK: .LBB{{[0-9]+}}_3:
-; CHECK:   call foo1{{$}}
-; CHECK: .LBB{{[0-9]+}}_4:
-; CHECK:   call foo2{{$}}
-; CHECK: .LBB{{[0-9]+}}_5:
-; CHECK:   call foo3{{$}}
-; CHECK: .LBB{{[0-9]+}}_6:
-; CHECK:   call foo4{{$}}
-; CHECK: .LBB{{[0-9]+}}_7:
-; CHECK:   call foo5{{$}}
-; CHECK: .LBB{{[0-9]+}}_8:
-; CHECK:   return{{$}}
+; CHECK:      block   {{$}}
+; CHECK:      i64.const
+; CHECK:      i64.gt_u
+; CHECK:      br_if 0
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      block   {{$}}
+; CHECK:      i32.wrap_i64
+; CHECK-NEXT: br_table {{[^,]+}}, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 4, 5, 0{{$}}
+; CHECK:      .LBB{{[0-9]+}}_2:
+; CHECK:        call foo0{{$}}
+; CHECK:      .LBB{{[0-9]+}}_3:
+; CHECK:        call foo1{{$}}
+; CHECK:      .LBB{{[0-9]+}}_4:
+; CHECK:        call foo2{{$}}
+; CHECK:      .LBB{{[0-9]+}}_5:
+; CHECK:        call foo3{{$}}
+; CHECK:      .LBB{{[0-9]+}}_6:
+; CHECK:        call foo4{{$}}
+; CHECK:      .LBB{{[0-9]+}}_7:
+; CHECK:        call foo5{{$}}
+; CHECK:      .LBB{{[0-9]+}}_8:
+; CHECK:        return{{$}}
 define void @bar64(i64 %n) {
 entry:
   switch i64 %n, label %sw.epilog [
     i64 0, label %sw.bb
     i64 1, label %sw.bb
     i64 2, label %sw.bb
     i64 3, label %sw.bb
     i64 4, label %sw.bb
sw.bb.5:                                          ; preds = %entry
   tail call void @foo5()
   br label %sw.epilog
 
 sw.epilog:                                        ; preds = %entry, %sw.bb.5, %sw.bb.4, %sw.bb.3, %sw.bb.2, %sw.bb.1, %sw.bb
   ret void
 }
 
 ; CHECK-LABEL: truncated:
-; CHECK:   block
-; CHECK:   block
-; CHECK:   block
-; CHECK:   i32.wrap_i64
-; CHECK:   br_table {{[^,]+}}, 0, 1, 2{{$}}
-; CHECK: .LBB{{[0-9]+}}_1
-; CHECK:   end_block
-; CHECK:   call foo0{{$}}
-; CHECK:   return{{$}}
-; CHECK: .LBB{{[0-9]+}}_2
-; CHECK:   end_block
-; CHECK:   call foo1{{$}}
-; CHECK:   return{{$}}
-; CHECK: .LBB{{[0-9]+}}_3
-; CHECK:   end_block
-; CHECK:   call foo2{{$}}
-; CHECK:   return{{$}}
-; CHECK:   end_function
+; CHECK:        block
+; CHECK:        block
+; CHECK:        block
+; CHECK:        i32.wrap_i64
+; CHECK-NEXT:   br_table {{[^,]+}}, 0, 1, 2{{$}}
+; CHECK:      .LBB{{[0-9]+}}_1
+; CHECK:        end_block
+; CHECK:        call foo0{{$}}
+; CHECK:        return{{$}}
+; CHECK:      .LBB{{[0-9]+}}_2
+; CHECK:        end_block
+; CHECK:        call foo1{{$}}
+; CHECK:        return{{$}}
+; CHECK:      .LBB{{[0-9]+}}_3
+; CHECK:        end_block
+; CHECK:        call foo2{{$}}
+; CHECK:        return{{$}}
+; CHECK:        end_function
 define void @truncated(i64 %n) {
 entry:
   %m = trunc i64 %n to i32
   switch i32 %m, label %default [
     i32 0, label %bb1
     i32 1, label %bb2
   ]