[WebAssembly] Fix writeback of stack pointer with dynamic alloca

Previously the stack pointer was only written back to memory in the
prolog. But this is wrong for dynamic allocas, for which
target-independent codegen handles SP updates after the prolog (and
possibly even in another BB). Instead update the SP global in
ADJCALLSTACKDOWN which is generated after the SP update sequence.
This will have further refinements when we add red zone support.

llvm-svn: 261579

1 files changed, 35 insertions, 32 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyFrameLowering.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyFrameLowering.cpp
index cbf59d1a3ce..c5f43d806ed 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyFrameLowering.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyFrameLowering.cpp
@@ -59,12 +59,39 @@ bool WebAssemblyFrameLowering::hasReservedCallFrame(
   return !MF.getFrameInfo()->hasVarSizedObjects();
 }
 
+static void writeSPToMemory(unsigned SrcReg, MachineFunction &MF,
+                            MachineBasicBlock &MBB,
+                            MachineBasicBlock::iterator &InsertPt,
+                            DebugLoc DL) {
+  auto *SPSymbol = MF.createExternalSymbolName("__stack_pointer");
+  unsigned SPAddr =
+      MF.getRegInfo().createVirtualRegister(&WebAssembly::I32RegClass);
+  const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
+
+  BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::CONST_I32), SPAddr)
+      .addExternalSymbol(SPSymbol);
+  auto *MMO = new MachineMemOperand(MachinePointerInfo(),
+                                    MachineMemOperand::MOStore, 4, 4);
+  BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::STORE_I32),
+          WebAssembly::SP32)
+      .addImm(0)
+      .addReg(SPAddr)
+      .addImm(2)  // p2align
+      .addReg(SrcReg)
+      .addMemOperand(MMO);
+  MF.getInfo<WebAssemblyFunctionInfo>()->stackifyVReg(SPAddr);
+}
+
 void WebAssemblyFrameLowering::eliminateCallFramePseudoInstr(
     MachineFunction &MF, MachineBasicBlock &MBB,
     MachineBasicBlock::iterator I) const {
-  // TODO: can we avoid using call frame pseudos altogether?
-  assert(!I->getOperand(0).getImm() &&
-         "Stack should not be adjusted around calls");
+  assert(!I->getOperand(0).getImm() && hasFP(MF) &&
+         "Call frame pseudos should only be used for dynamic stack adjustment");
+  const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
+  if (I->getOpcode() == TII->getCallFrameDestroyOpcode()) {
+    DebugLoc DL = I->getDebugLoc();
+    writeSPToMemory(WebAssembly::SP32, MF, MBB, I, DL);
+  }
   MBB.erase(I);
 }
 
@@ -125,21 +152,8 @@ void WebAssemblyFrameLowering::emitPrologue(MachineFunction &MF,
             WebAssembly::FP32)
         .addReg(WebAssembly::SP32);
   }
-  if (StackSize || hasFP(MF)) {
-    SPAddr = MRI.createVirtualRegister(&WebAssembly::I32RegClass);
-    // The SP32 register now has the new stacktop. Also write it back to memory.
-    BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::CONST_I32), SPAddr)
-        .addExternalSymbol(SPSymbol);
-    auto *MMO = new MachineMemOperand(MachinePointerInfo(),
-                                      MachineMemOperand::MOStore, 4, 4);
-    BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::STORE_I32),
-            WebAssembly::SP32)
-        .addImm(0)
-        .addReg(SPAddr)
-        .addImm(2)  // p2align
-        .addReg(WebAssembly::SP32)
-        .addMemOperand(MMO);
-    WFI->stackifyVReg(SPAddr);
+  if (StackSize) {
+    writeSPToMemory(WebAssembly::SP32, MF, MBB, InsertPt, DL);
   }
 }
 
@@ -171,18 +185,7 @@ void WebAssemblyFrameLowering::emitEpilogue(MachineFunction &MF,
     WFI->stackifyVReg(OffsetReg);
   }
 
-  auto *SPSymbol = MF.createExternalSymbolName("__stack_pointer");
-  unsigned SPAddr = MRI.createVirtualRegister(&WebAssembly::I32RegClass);
-  BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::CONST_I32), SPAddr)
-      .addExternalSymbol(SPSymbol);
-  auto *MMO = new MachineMemOperand(MachinePointerInfo(),
-                                    MachineMemOperand::MOStore, 4, 4);
-  BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::STORE_I32),
-          WebAssembly::SP32)
-      .addImm(0)
-      .addReg(SPAddr)
-      .addImm(2)  // p2align
-      .addReg((!StackSize && hasFP(MF)) ? WebAssembly::FP32 : WebAssembly::SP32)
-      .addMemOperand(MMO);
-  WFI->stackifyVReg(SPAddr);
+  writeSPToMemory(
+      (!StackSize && hasFP(MF)) ? WebAssembly::FP32 : WebAssembly::SP32, MF,
+      MBB, InsertPt, DL);
 }