[llvm] [WebAssembly] Unstackify operands for removed terminators (PR #149432)

[Heejin Ahn via llvm-commits]

https://github.com/aheejin updated https://github.com/llvm/llvm-project/pull/149432

>From 0746e1101273c9633abdc7d213a8254634f08dba Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Thu, 17 Jul 2025 22:27:46 +0000
Subject: [PATCH 1/2] [WebAssembly] Unstackify operands for removed terminators

There are cases we end up removing a conditional branch with a stackified
register condition operand. For example:
```wasm
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.
---
 .../Target/WebAssembly/WebAssemblyCFGSort.cpp | 35 ++++++++++++++++++-
 .../CodeGen/WebAssembly/removed-terminator.ll | 15 ++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/WebAssembly/removed-terminator.ll

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
+}

>From 31ad3d880ae40e7cd3373b6bd97ce24563f05932 Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Fri, 18 Jul 2025 01:09:55 +0000
Subject: [PATCH 2/2] Fix comments

---
 llvm/test/CodeGen/WebAssembly/removed-terminator.ll | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/llvm/test/CodeGen/WebAssembly/removed-terminator.ll b/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
index 4a3c590ea13f7..8c9b6ff4bca97 100644
--- a/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
+++ b/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
@@ -4,10 +4,9 @@ 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.
+  ; This br_if's operand (%y) is stackified in RegStackify. 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: