[WebAssembly] Fixed FrameBaseLocal not being set.

Summary:
Fixes: https://bugs.llvm.org/show_bug.cgi?id=44920

WebAssemblyRegColoring may merge the vreg that currently represents
the FrameBase with one representing an argument.
WebAssemblyExplicitLocals picks up the corresponding local when
a vreg is first added to the Reg2Local mapping, except when it is
an argument instruction which are handled separately.

Note that this does not change that vregs representing the FrameBase
may get merged, it is not clear to me that this may have other
effects we may want to avoid?

Reviewers: dschuff

Reviewed By: dschuff

Subscribers: azakai, sbc100, hiraditya, aheejin, sunfish, llvm-commits, jgravelle-google

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D75718

1 files changed, 16 insertions, 9 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp
index 4da97b6..56da3ec 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyExplicitLocals.cpp
@@ -69,6 +69,18 @@ FunctionPass *llvm::createWebAssemblyExplicitLocals() {
   return new WebAssemblyExplicitLocals();
 }
 
+static void checkFrameBase(WebAssemblyFunctionInfo &MFI, unsigned Local,
+                           unsigned Reg) {
+  // Mark a local for the frame base vreg.
+  if (MFI.isFrameBaseVirtual() && Reg == MFI.getFrameBaseVreg()) {
+    LLVM_DEBUG({
+      dbgs() << "Allocating local " << Local << "for VReg "
+             << Register::virtReg2Index(Reg) << '\n';
+    });
+    MFI.setFrameBaseLocal(Local);
+  }
+}
+
 /// Return a local id number for the given register, assigning it a new one
 /// if it doesn't yet have one.
 static unsigned getLocalId(DenseMap<unsigned, unsigned> &Reg2Local,
@@ -76,14 +88,7 @@ static unsigned getLocalId(DenseMap<unsigned, unsigned> &Reg2Local,
                            unsigned Reg) {
   auto P = Reg2Local.insert(std::make_pair(Reg, CurLocal));
   if (P.second) {
-    // Mark the local allocated for the frame base vreg.
-    if (MFI.isFrameBaseVirtual() && Reg == MFI.getFrameBaseVreg()) {
-      LLVM_DEBUG({
-        dbgs() << "Allocating local " << CurLocal << "for VReg "
-               << Register::virtReg2Index(Reg) << '\n';
-      });
-      MFI.setFrameBaseLocal(CurLocal);
-    }
+    checkFrameBase(MFI, CurLocal, Reg);
     ++CurLocal;
   }
   return P.first->second;
@@ -227,7 +232,9 @@ bool WebAssemblyExplicitLocals::runOnMachineFunction(MachineFunction &MF) {
       break;
     Register Reg = MI.getOperand(0).getReg();
     assert(!MFI.isVRegStackified(Reg));
-    Reg2Local[Reg] = static_cast<unsigned>(MI.getOperand(1).getImm());
+    auto Local = static_cast<unsigned>(MI.getOperand(1).getImm());
+    Reg2Local[Reg] = Local;
+    checkFrameBase(MFI, Local, Reg);
     MI.eraseFromParent();
     Changed = true;
   }