[WebAssembly] Unstackify operands for removed terminators

[aheejin]

There are cases we end up removing a conditional branch with a stackified register condition operand. For example:

bb.0:
  %0 = ...    ;; %0 is stackified
  br_if %bb.1, %0
bb.1:

In this code, br_if will be removed, so we should unstackify %0 so that it can be correctly dropped in ExplicitLocals.

There seems no good method to determine which registers to unstackify without analyzing branches thoroughly, which is supposed to be done within updateTerminator(). So we just unstackify all of them and restackify the still-remaining registers after updateTerminator().

Fixes #149097.

[llvmbot]

@llvm/pr-subscribers-backend-webassembly

Author: Heejin Ahn (aheejin)

Changes

There are cases we end up removing a conditional branch with a stackified register condition operand. For example:

bb.0:
  %0 = ...    ;; %0 is stackified
  br_if %bb.1, %0
bb.1:
```wasm

In this code, br_if will be removed, so we should unstackify %0 so that it can be correctly dropped in ExplicitLocals.

There seems no good method to determine which registers to unstackify without analyzing branches thoroughly, which is supposed to be done within `updateTerminator()`. So we just unstackify all of them and restackify the still-remaining registers after `updateTerminator()`.

Fixes #<!-- -->149097.

---
Full diff: https://github.com/llvm/llvm-project/pull/149432.diff


2 Files Affected:

- (modified) llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp (+34-1) 
- (added) llvm/test/CodeGen/WebAssembly/removed-terminator.ll (+15) 


``````````diff
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp
index 6525c6d93bee8..378579d8776bb 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp
@@ -18,6 +18,7 @@
 
 #include "WebAssembly.h"
 #include "WebAssemblyExceptionInfo.h"
+#include "WebAssemblyMachineFunctionInfo.h"
 #include "WebAssemblySortRegion.h"
 #include "WebAssemblyUtilities.h"
 #include "llvm/ADT/PriorityQueue.h"
@@ -97,8 +98,40 @@ static void maybeUpdateTerminator(MachineBasicBlock *MBB) {
           ? MF->getBlockNumbered(MBB->getNumber() + 1)
           : nullptr;
 
-  if (AllAnalyzable)
+  if (AllAnalyzable) {
+    // There are cases we end up removing a conditional branch with a stackified
+    // register condition operand. For example:
+    // bb.0:
+    //   %0 = ...    ;; %0 is stackified
+    //   br_if %bb.1, %0
+    // bb.1:
+    //
+    // In this code, br_if will be removed, so we should unstackify %0 so that
+    // it can be correctly dropped in ExplicitLocals.
+    //
+    // There seems no good method to determine which registers to unstackify
+    // without analyzing branches thoroughly, which is supposed to be done
+    // within updateTerminator(). So we just unstackify all of them and
+    // restackify the still-remaining registers after updateTerminator().
+    SmallSet<Register, 2> StackifiedRegs;
+    auto *MF = MBB->getParent();
+    WebAssemblyFunctionInfo *MFI = MF->getInfo<WebAssemblyFunctionInfo>();
+    for (const MachineInstr &Term : MBB->terminators()) {
+      for (auto &MO : Term.explicit_uses()) {
+        if (MO.isReg() && MFI->isVRegStackified(MO.getReg())) {
+          StackifiedRegs.insert(MO.getReg());
+          MFI->unstackifyVReg(MO.getReg());
+        }
+      }
+    }
+
     MBB->updateTerminator(OriginalSuccessor);
+
+    for (const MachineInstr &Term : MBB->terminators())
+      for (auto &MO : Term.explicit_uses())
+        if (MO.isReg() && StackifiedRegs.contains(MO.getReg()))
+          MFI->stackifyVReg(MF->getRegInfo(), MO.getReg());
+  }
 }
 
 namespace {
diff --git a/llvm/test/CodeGen/WebAssembly/removed-terminator.ll b/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
new file mode 100644
index 0000000000000..4a3c590ea13f7
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
@@ -0,0 +1,15 @@
+; RUN: llc -O0 < %s
+
+target triple = "wasm32-unknown-unknown"
+
+define void @test(i1 %x) {
+  %y = xor i1 %x, true
+  ; We now do a limited amount of register stackification in RegStackify even in
+  ; -O0, so its operand (%y) is stackified. But this terminator will be removed
+  ; in CFGSort after that. We need to make sure we unstackify %y so that it can
+  ; be dropped in ExplicitLocals.
+  br i1 %y, label %exit, label %exit
+
+exit:
+  ret void
+}

[aheejin]

@SingleAccretion (you don't show up in the reviewers list)

[aheejin]

Some background: #149097 (comment)

[nikic]

Hm, couldn't this still potentially cause issues with tee instructions? That is, if a register does stay unstackified and it happens to be the first result of tee?

Otherwise this looks reasonable to me.