[llvm] [WebAssembly] Unstackify registers with no uses in ExplicitLocals (PR #149626)

[Heejin Ahn via llvm-commits]

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

>From 4ba45dda0a32d205ba266735e9914fc1c7199f6f Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Sat, 19 Jul 2025 01:13:19 +0000
Subject: [PATCH 1/3] [WebAssembly] Unstackify registers with no uses in
 ExplicitLocals

There are cases we end up removing some intructions that use stackified
registers after RegStackify. For example,

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

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

Rather than handling this in case-by-case basis, this PR just
unstackifies all stackifies register with no uses, so that they can be
correctly dropped.

Fixes #149097.
---
 .../WebAssembly/WebAssemblyExplicitLocals.cpp      | 12 ++++++++++--
 .../test/CodeGen/WebAssembly/removed-terminator.ll | 14 ++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/WebAssembly/removed-terminator.ll

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp
index 2662241ef8499..ad28e36f29112 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp
@@ -256,9 +256,17 @@ bool WebAssemblyExplicitLocals::runOnMachineFunction(MachineFunction &MF) {
 
   // Precompute the set of registers that are unused, so that we can insert
   // drops to their defs.
+  // And unstackify any stackified registers that don't have any uses, so that
+  // they can be dropped later. This can happen when transformations after
+  // RegStackify removes instructions that use stackified registers.
   BitVector UseEmpty(MRI.getNumVirtRegs());
-  for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I)
-    UseEmpty[I] = MRI.use_empty(Register::index2VirtReg(I));
+  for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I) {
+    Register Reg = Register::index2VirtReg(I);
+    if (MRI.use_empty(Reg)) {
+      UseEmpty[I] = true;
+      MFI.unstackifyVReg(Reg);
+    }
+  }
 
   // Visit each instruction in the function.
   for (MachineBasicBlock &MBB : MF) {
diff --git a/llvm/test/CodeGen/WebAssembly/removed-terminator.ll b/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
new file mode 100644
index 0000000000000..8c9b6ff4bca97
--- /dev/null
+++ b/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
@@ -0,0 +1,14 @@
+; RUN: llc -O0 < %s
+
+target triple = "wasm32-unknown-unknown"
+
+define void @test(i1 %x) {
+  %y = xor i1 %x, true
+  ; 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:
+  ret void
+}

>From d6a2019251c22c1980b703cb2ca43cb53dafca20 Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Sat, 19 Jul 2025 21:32:43 +0000
Subject: [PATCH 2/3] Fix comment

---
 llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp
index ad28e36f29112..e6486e247209b 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp
@@ -258,7 +258,7 @@ bool WebAssemblyExplicitLocals::runOnMachineFunction(MachineFunction &MF) {
   // drops to their defs.
   // And unstackify any stackified registers that don't have any uses, so that
   // they can be dropped later. This can happen when transformations after
-  // RegStackify removes instructions that use stackified registers.
+  // RegStackify remove instructions using stackified registers.
   BitVector UseEmpty(MRI.getNumVirtRegs());
   for (unsigned I = 0, E = MRI.getNumVirtRegs(); I < E; ++I) {
     Register Reg = Register::index2VirtReg(I);

>From 72afdd8e04f6eb44bba1d65416aa0795dcc2268a Mon Sep 17 00:00:00 2001
From: Heejin Ahn <aheejin at gmail.com>
Date: Sat, 19 Jul 2025 21:43:53 +0000
Subject: [PATCH 3/3] Add test expectation

---
 .../test/CodeGen/WebAssembly/removed-terminator.ll | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/llvm/test/CodeGen/WebAssembly/removed-terminator.ll b/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
index 8c9b6ff4bca97..188f6f67eee8b 100644
--- a/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
+++ b/llvm/test/CodeGen/WebAssembly/removed-terminator.ll
@@ -1,8 +1,20 @@
-; RUN: llc -O0 < %s
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc -O0 -verify-machineinstrs < %s | FileCheck %s
 
 target triple = "wasm32-unknown-unknown"
 
 define void @test(i1 %x) {
+; CHECK-LABEL: test:
+; CHECK:         .functype test (i32) -> ()
+; CHECK-NEXT:  # %bb.0:
+; CHECK-NEXT:    local.get 0
+; CHECK-NEXT:    i32.const -1
+; CHECK-NEXT:    i32.xor
+; CHECK-NEXT:    i32.const 1
+; CHECK-NEXT:    i32.and
+; CHECK-NEXT:    drop
+; CHECK-NEXT:  # %bb.1: # %exit
+; CHECK-NEXT:    return
   %y = xor i1 %x, true
   ; 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