[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.

[SingleAccretion]

Did you consider the fix I was thinking about (dropping all unused defs not just those not-stackified) in explicit-locals?

[aheejin]

Did you consider the fix I was thinking about (dropping all unused defs not just those not-stackified) in explicit-locals?

I think that is a better idea. See #149626.

[aheejin]

Closing in favor of #149626.

[aheejin]

@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?

The only case I can think of is a really contrived case like this:
(I've removed all implicit arguments for readability, so this may not run with llc out of the box)

bb.0:
  successors: %bb.1; %bb.1
  %0:i32 = CALL @foo
  %1:i32 = CONST_I32 -1
  %2:i32 = XOR_I32 %0:i32, %1:i32
  BR_IF %bb.1, %2:i32
  BR_IF %bb.1, %2:i32
  BR_IF %bb.1, %2:i32

bb.1:
  ; predecessors: %bb.0
  RETURN

Because %2 is used in both CALL @use and BR_IF, we will get a TEE:

bb.0:                                                                        
  successors: %bb.1; %bb.1                                                            
  %0:i32 = CALL @foo
  %1:i32 = CONST_I32 -1
  %4:i32 = XOR_I32 %0:i32, %1:i32
  %3:i32, %2:i32 = TEE_I32 %4:i32
  BR_IF %bb.1, %3:i32
  BR_IF %bb.1, %2:i32
  BR_IF %bb.1, %2:i32

bb.1:
; predecessors: %bb.0
  RETURN

with %4 and %3 stackified. (See this for how it's done)

If all these BR_IFs get removed in CFGSort I think the situation you are worried about can occur. If I run this from RegStackify, this fails in here:

llvm-project/llvm/lib/CodeGen/MachineBasicBlock.cpp

Line 706 in cfddb40

assert(!B && "UpdateTerminators requires analyzable predecessors!");

which is called from here:

llvm-project/llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp

Line 101 in cfddb40

MBB->updateTerminator(OriginalSuccessor);

So I think it's OK to assume that this kind of code is not generated without manually modifying mir files.

In valid code, there will be only one BR_IF, so the code could be something like

bb.0:
  successors: %bb.1; %bb.1
  %0:i32 = CALL @foo
  %1:i32 = CONST_I32 -1
  %2:i32 = XOR_I32 %0:i32, %1:i32
  CALL @use, %2:i32
  CALL @use, %2:i32
  BR_IF %bb.1, %2:i32

bb.1:
  ; predecessors: %bb.0
  RETURN

bb.0:                                                                        
  successors: %bb.1; %bb.1                                                            
  %0:i32 = CALL @foo
  %1:i32 = CONST_I32 -1
  %4:i32 = XOR_I32 %0:i32, %1:i32
  %3:i32, %2:i32 = TEE_I32 %4:i32
  CALL @use, %3:i32
  CALL @use, %2:i32
  BR_IF %bb.1, %2:i32

bb.1:
; predecessors: %bb.0
  RETURN

with %4 and %3 stackified.

In this case, even if the BR_IF is removed, its operand (%2) is not the stackified register.