Add Wasm peephole optimisation to remove dead code around unreachables

lld/test/wasm/init-fini.ll

@@ -77,7 +77,7 @@ entry:
 ; CHECK-NEXT:         Body:            10031004100A100F1012100F10141003100C100F10161001100E0B
 ; CHECK:            - Index:           22
 ; CHECK-NEXT:         Locals:
-; CHECK-NEXT:         Body:            02404186808080004100418088808000108780808000450D0000000B0B
+; CHECK-NEXT:         Body:            02404186808080004100418088808000108780808000450D00000B0B
 ; CHECK-NEXT:   - Type:            CUSTOM
 ; CHECK-NEXT:     Name:            name
 ; CHECK-NEXT:     FunctionNames:

llvm/lib/Target/WebAssembly/WebAssemblyDebugFixup.cpp

@@ -122,9 +122,25 @@ bool WebAssemblyDebugFixup::runOnMachineFunction(MachineFunction &MF) {
           // it will be culled later.
         }
       } else {
+
+        // There is a WebAssembly peephole optimisation can remove drop
+        // instructions before a wasm unreachable. This is a valid
+        // transformation because unreachable is "stack polymorphic", but stack
+        // polymorphism is not modeled in the llvm wasm backend.
+        // In current codegen, virtual registers can only be stackified within
+        // a basic block, and cannot be stackified across an UNREACHABLE
+        // because it is marked as having unmodeled side effects.
+        // So we can assume the operand stack is empty after an UNREACHABLE.
+        if (MI.getOpcode() == WebAssembly::UNREACHABLE) {
+          Stack.clear();
+          break;
+        }
+
         // Track stack depth.
         for (MachineOperand &MO : reverse(MI.explicit_uses())) {
           if (MO.isReg() && MFI.isVRegStackified(MO.getReg())) {
+            assert(!Stack.empty() &&
+                   "WebAssemblyDebugFixup: Pop: Operand stack empty!");
             auto Prev = Stack.back();
             Stack.pop_back();
             assert(Prev.Reg == MO.getReg() &&

llvm/lib/Target/WebAssembly/WebAssemblyPeephole.cpp

@@ -20,6 +20,7 @@
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include <iterator>
 using namespace llvm;
 
 #define DEBUG_TYPE "wasm-peephole"
@@ -109,6 +110,53 @@ static bool maybeRewriteToFallthrough(MachineInstr &MI, MachineBasicBlock &MBB,
   return true;
 }
 
+static bool eraseDeadCodeAroundUnreachable(MachineInstr &UnreachbleMI,
+                                           MachineBasicBlock &MBB) {
+  SmallVector<MachineInstr *, 16> ToDelete;
+
+  // Because wasm unreachable is stack polymorphic and unconditionally ends
+  // control, all instructions after it can be removed until the end of this
+  // block. We remove the common case of double unreachable.
+  auto ForwardsIterator = UnreachbleMI.getIterator();
+  for (ForwardsIterator++; !ForwardsIterator.isEnd(); ForwardsIterator++) {
+    MachineInstr &MI = *ForwardsIterator;
+    if (MI.getOpcode() == WebAssembly::UNREACHABLE) {
+      ToDelete.push_back(&MI);
+    } else {
+      break;
+    }
+  }
+
+  // For the same reasons as above, previous instructions that only affect
+  // local function state can be removed (e.g. local.set, drop, various reads).
+  // We remove the common case of "drop unreachable".
+  auto BackwardsIterator = UnreachbleMI.getReverseIterator();
+  for (BackwardsIterator++; !BackwardsIterator.isEnd(); BackwardsIterator++) {
+    MachineInstr &MI = *BackwardsIterator;
+    switch (MI.getOpcode()) {
+    case WebAssembly::DROP_I32:
+    case WebAssembly::DROP_I64:
+    case WebAssembly::DROP_F32:
+    case WebAssembly::DROP_F64:
+    case WebAssembly::DROP_EXTERNREF:
+    case WebAssembly::DROP_FUNCREF:
+    case WebAssembly::DROP_V128:
+      ToDelete.push_back(&MI);
+      continue;
+    default:
+      goto exit_loop;
+    }
+  }
+exit_loop:;
+
+  bool Changed = false;
+  for (MachineInstr *MI : ToDelete) {
+    MI->eraseFromParent();
+    Changed = true;
+  }
+  return Changed;
+}
+
 bool WebAssemblyPeephole::runOnMachineFunction(MachineFunction &MF) {
   LLVM_DEBUG({
     dbgs() << "********** Peephole **********\n"
@@ -159,6 +207,9 @@ bool WebAssemblyPeephole::runOnMachineFunction(MachineFunction &MF) {
       case WebAssembly::RETURN:
         Changed |= maybeRewriteToFallthrough(MI, MBB, MF, MFI, MRI, TII);
         break;
+      case WebAssembly::UNREACHABLE:
+        Changed |= eraseDeadCodeAroundUnreachable(MI, MBB);
+        break;
       }
 
   return Changed;

llvm/test/CodeGen/WebAssembly/unreachable.ll

@@ -54,7 +54,6 @@ define void @trap_unreacheable() {
 ; CHECK:         .functype trap_unreacheable () -> ()
 ; CHECK-NEXT:  # %bb.0:
 ; CHECK-NEXT:    unreachable
-; CHECK-NEXT:    unreachable
 ; CHECK-NEXT:    end_function
   call void @llvm.trap()
   unreachable
@@ -80,11 +79,10 @@ define i32 @missing_ret_unreachable() {
   unreachable
 }
 
-; This is similar to the above test, but ensures wasm unreachable is emitted
-; This is similar to the above test, but the callee has a 'noreturn' attribute.    
-; There is an optimization that removes an 'unreachable' after a noreturn call,  
-; but Wasm backend doesn't use it and ignore `--no-trap-after-noreturn`, if      
-; given, to generate valid code. 
+; This is similar to the above test, but the callee has a 'noreturn' attribute.
+; There is an optimization that removes an 'unreachable' after a noreturn call,
+; but Wasm backend doesn't use it and ignore `--no-trap-after-noreturn`, if
+; given, to generate valid code.
 define i32 @missing_ret_noreturn_unreachable() {
 ; CHECK-LABEL: missing_ret_noreturn_unreachable:
 ; CHECK:         .functype missing_ret_noreturn_unreachable () -> (i32)
@@ -95,3 +93,18 @@ define i32 @missing_ret_noreturn_unreachable() {
   call void @ext_never_return()
   unreachable
 }
+
+; This is a test for the wasm peephole optimization that erases unnecessary
+; code before and after a wasm unreachable instruction.
+; It creates an unused stack variable by calling ext_func_i32() that would
+; otherwise require a drop instruction before the wasm unreachable instruction.
+define i64 @drop_unreachable() {
+; CHECK-LABEL: drop_unreachable:
+; CHECK:         .functype drop_unreachable () -> (i64)
+; CHECK-NEXT:  # %bb.0:
+; CHECK-NEXT:    call ext_func_i32
+; CHECK-NEXT:    unreachable
+; CHECK-NEXT:    end_function
+  call i32 @ext_func_i32()
+  unreachable
+}

llvm/test/MC/WebAssembly/global-ctor-dtor.ll

@@ -80,29 +80,29 @@ declare void @func3()
 ; CHECK-NEXT:         Offset:          0x1D
 ; CHECK-NEXT:       - Type:            R_WASM_FUNCTION_INDEX_LEB
 ; CHECK-NEXT:         Index:           6
-; CHECK-NEXT:         Offset:          0x2C
+; CHECK-NEXT:         Offset:          0x2B
 ; CHECK-NEXT:       - Type:            R_WASM_TABLE_INDEX_SLEB
 ; CHECK-NEXT:         Index:           5
-; CHECK-NEXT:         Offset:          0x37
+; CHECK-NEXT:         Offset:          0x36
 ; CHECK-NEXT:       - Type:            R_WASM_MEMORY_ADDR_SLEB
 ; CHECK-NEXT:         Index:           3
-; CHECK-NEXT:         Offset:          0x3F
+; CHECK-NEXT:         Offset:          0x3E
 ; CHECK-NEXT:       - Type:            R_WASM_FUNCTION_INDEX_LEB
 ; CHECK-NEXT:         Index:           4
-; CHECK-NEXT:         Offset:          0x45
+; CHECK-NEXT:         Offset:          0x44
 ; CHECK-NEXT:     Functions:
 ; CHECK-NEXT:       - Index:           5
 ; CHECK-NEXT:         Locals:
 ; CHECK-NEXT:         Body:            1080808080000B
 ; CHECK-NEXT:       - Index:           6
 ; CHECK-NEXT:         Locals:
-; CHECK-NEXT:         Body:            02404181808080004100418080808000108180808000450D0000000B0B
+; CHECK-NEXT:         Body:            02404181808080004100418080808000108180808000450D00000B0B
 ; CHECK-NEXT:       - Index:           7
 ; CHECK-NEXT:         Locals:
 ; CHECK-NEXT:         Body:            1082808080000B
 ; CHECK-NEXT:       - Index:           8
 ; CHECK-NEXT:         Locals:
-; CHECK-NEXT:         Body:            02404182808080004100418080808000108180808000450D0000000B0B
+; CHECK-NEXT:         Body:            02404182808080004100418080808000108180808000450D00000B0B
 ; CHECK-NEXT:   - Type:            DATA
 ; CHECK-NEXT:     Segments:
 ; CHECK-NEXT:       - SectionOffset:   6