[WebAssembly] Emit a BasePointer when we have overly-aligned stack objects

[jgravelle-google]

Summary

Because we shift the stack pointer by an unknown amount, we need an
additional pointer. In the case where we have variable-size objects
as well, we can't reuse the frame pointer, thus three pointers.

[jgravelle-google]

• Move HasBP closer to its use

[dschuff]

lib/Target/WebAssembly/WebAssemblyFrameLowering.cpp

40 ↗

(On Diff #76791)

Maybe clarify that "overaligned" means more aligned than the stack pointer (rather than more than natural alignment for the slot itself).

test/CodeGen/WebAssembly/alignment.ll

1 ↗

(On Diff #76791)

alignment.ll seems a bit generic. It could probably be stack-aligment or else these tests could just be added to userstack.ll

29 ↗

(On Diff #76791)

Might as well check for the value of the const too, so we know the alignment is 16.

34 ↗

(On Diff #76791)

Seems like we should be storing directly from BP. Where does this copy_local come from?
The epilog codegen looks ok; I guess this copy is the one generated in the prolog and gets moved down by the stackifier?

[jgravelle-google]

test/CodeGen/WebAssembly/alignment.ll

34 ↗

(On Diff #76791)

Huh. This is caused by us emitting a frame pointer, and register stackification doing magic to it.
tl;dr: Yes.

Prolog/epilogue:

1. load __stack_pointer into SP
2. copy SP into BP
3. decrement SP by fixed amount
4. store new SP as FP
5. [rest of code]
6. store BP into __stack_pointer

Reg stackify:

1. Moves the single copy into BP to be next to its use
2. Tees the two uses of the original load of __stack_pointer, stackifies one for use in computing FP, and saves one in BP

The good news is this just gets translated to get_local in s2wasm. And adding -wasm-explicit-locals models it that was as well, so there's no actual cost here, it's just kinda quirky.
... is what I've been able to piece together.

[jgravelle-google]

• Simplify SPReg, clarify comment, rename test, check stack size in test

llvm/trunk/lib/Target/WebAssembly/WebAssemblyFrameLowering.h

public:
   /// These methods insert prolog and epilog code into the function.
   void emitPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
   void emitEpilogue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
 
   bool hasFP(const MachineFunction &MF) const override;
   bool hasReservedCallFrame(const MachineFunction &MF) const override;
 
  private:
+  bool hasBP(const MachineFunction &MF) const;
   bool needsSP(const MachineFunction &MF, const MachineFrameInfo &MFI) const;
   bool needsSPWriteback(const MachineFunction &MF,
                         const MachineFrameInfo &MFI) const;
 };
 
 }  // end namespace llvm
 
 #endif

llvm/trunk/lib/Target/WebAssembly/WebAssemblyFrameLowering.cpp

 #include "llvm/Support/Debug.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "wasm-frame-info"
 
 // TODO: wasm64
 // TODO: Emit TargetOpcode::CFI_INSTRUCTION instructions
 
+/// We need a base pointer in the case of having items on the stack that
+/// require stricter alignment than the stack pointer itself.  Because we need
+/// to shift the stack pointer by some unknown amount to force the alignment,
+/// we need to record the value of the stack pointer on entry to the function.
+bool WebAssemblyFrameLowering::hasBP(
+    const MachineFunction &MF) const {
+  const auto *RegInfo =
+      MF.getSubtarget<WebAssemblySubtarget>().getRegisterInfo();
+  return RegInfo->needsStackRealignment(MF);
+}
+
 /// Return true if the specified function should have a dedicated frame pointer
 /// register.
 bool WebAssemblyFrameLowering::hasFP(const MachineFunction &MF) const {
   const MachineFrameInfo &MFI = MF.getFrameInfo();
-  const auto *RegInfo =
-      MF.getSubtarget<WebAssemblySubtarget>().getRegisterInfo();
-  return MFI.isFrameAddressTaken() || MFI.hasVarSizedObjects() ||
-         MFI.hasStackMap() || MFI.hasPatchPoint() ||
-         RegInfo->needsStackRealignment(MF);
+
+  // When we have var-sized objects, we move the stack pointer by an unknown
+  // amount, and need to emit a frame pointer to restore the stack to where we
+  // were on function entry.
+  // If we already need a base pointer, we use that to fix up the stack pointer.
+  // If there are no fixed-size objects, we would have no use of a frame
+  // pointer, and thus should not emit one.
+  bool HasFixedSizedObjects = MFI.getStackSize() > 0;
+  bool NeedsFixedReference = !hasBP(MF) || HasFixedSizedObjects;
+
+  return MFI.isFrameAddressTaken() ||
+         (MFI.hasVarSizedObjects() && NeedsFixedReference) ||
+         MFI.hasStackMap() || MFI.hasPatchPoint();
 }
 
 /// Under normal circumstances, when a frame pointer is not required, we reserve
 /// argument space for call sites in the function immediately on entry to the
 /// current function. This eliminates the need for add/sub sp brackets around
 /// call sites. Returns true if the call frame is included as part of the stack
 /// frame.
 bool WebAssemblyFrameLowering::hasReservedCallFrame(
BuildMI(MBB, InsertStore, DL, TII->get(WebAssembly::STORE_I32))
       .addReg(SrcReg)
       .addMemOperand(MMO);
 }
 
 MachineBasicBlock::iterator
 WebAssemblyFrameLowering::eliminateCallFramePseudoInstr(
     MachineFunction &MF, MachineBasicBlock &MBB,
     MachineBasicBlock::iterator I) const {
-  assert(!I->getOperand(0).getImm() && hasFP(MF) &&
+  assert(!I->getOperand(0).getImm() && (hasFP(MF) || hasBP(MF)) &&
          "Call frame pseudos should only be used for dynamic stack adjustment");
   const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
   if (I->getOpcode() == TII->getCallFrameDestroyOpcode() &&
       needsSPWriteback(MF, MF.getFrameInfo())) {
     DebugLoc DL = I->getDebugLoc();
     writeSPToMemory(WebAssembly::SP32, MF, MBB, I, I, DL);
   }
   return MBB.erase(I);
void WebAssemblyFrameLowering::emitPrologue(MachineFunction &MF,
   auto &MRI = MF.getRegInfo();
 
   auto InsertPt = MBB.begin();
   DebugLoc DL;
 
   const TargetRegisterClass *PtrRC =
       MRI.getTargetRegisterInfo()->getPointerRegClass(MF);
   unsigned Zero = MRI.createVirtualRegister(PtrRC);
-  unsigned SPReg = MRI.createVirtualRegister(PtrRC);
+  unsigned SPReg = WebAssembly::SP32;
+  if (StackSize)
+    SPReg = MRI.createVirtualRegister(PtrRC);
   const char *ES = "__stack_pointer";
   auto *SPSymbol = MF.createExternalSymbolName(ES);
   BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::CONST_I32), Zero)
       .addImm(0);
   MachineMemOperand *LoadMMO = MF.getMachineMemOperand(
       MachinePointerInfo(MF.getPSVManager().getExternalSymbolCallEntry(ES)),
       MachineMemOperand::MOLoad, 4, 4);
   // Load the SP value.
-  BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::LOAD_I32),
-          StackSize ? SPReg : (unsigned)WebAssembly::SP32)
+  BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::LOAD_I32), SPReg)
       .addImm(2)       // p2align
       .addExternalSymbol(SPSymbol)
       .addReg(Zero)    // addr
       .addMemOperand(LoadMMO);
 
+  bool HasBP = hasBP(MF);
+  if (HasBP) {
+    auto FI = MF.getInfo<WebAssemblyFunctionInfo>();
+    unsigned BasePtr = MRI.createVirtualRegister(PtrRC);
+    FI->setBasePointerVreg(BasePtr);
+    BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::COPY), BasePtr)
+        .addReg(SPReg);
+  }
   if (StackSize) {
     // Subtract the frame size
     unsigned OffsetReg = MRI.createVirtualRegister(PtrRC);
     BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::CONST_I32), OffsetReg)
         .addImm(StackSize);
     BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::SUB_I32),
             WebAssembly::SP32)
         .addReg(SPReg)
         .addReg(OffsetReg);
   }
+  if (HasBP) {
+    unsigned BitmaskReg = MRI.createVirtualRegister(PtrRC);
+    unsigned Alignment = MFI.getMaxAlignment();
+    assert((1 << countTrailingZeros(Alignment)) == Alignment &&
+      "Alignment must be a power of 2");
+    BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::CONST_I32), BitmaskReg)
+        .addImm((int)~(Alignment - 1));
+    BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::AND_I32),
+            WebAssembly::SP32)
+        .addReg(WebAssembly::SP32)
+        .addReg(BitmaskReg);
+  }
   if (hasFP(MF)) {
     // Unlike most conventional targets (where FP points to the saved FP),
     // FP points to the bottom of the fixed-size locals, so we can use positive
     // offsets in load/store instructions.
     BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::COPY),
             WebAssembly::FP32)
         .addReg(WebAssembly::SP32);
   }
void WebAssemblyFrameLowering::emitEpilogue(MachineFunction &MF,
 
   if (InsertPt != MBB.end())
     DL = InsertPt->getDebugLoc();
 
   // Restore the stack pointer. If we had fixed-size locals, add the offset
   // subtracted in the prolog.
   unsigned SPReg = 0;
   MachineBasicBlock::iterator InsertAddr = InsertPt;
-  if (StackSize) {
+  if (hasBP(MF)) {
+    auto FI = MF.getInfo<WebAssemblyFunctionInfo>();
+    SPReg = FI->getBasePointerVreg();
+  } else if (StackSize) {
     const TargetRegisterClass *PtrRC =
         MRI.getTargetRegisterInfo()->getPointerRegClass(MF);
     unsigned OffsetReg = MRI.createVirtualRegister(PtrRC);
     InsertAddr =
         BuildMI(MBB, InsertPt, DL, TII->get(WebAssembly::CONST_I32), OffsetReg)
             .addImm(StackSize);
     // In the epilog we don't need to write the result back to the SP32 physreg
     // because it won't be used again. We can use a stackified register instead.

llvm/trunk/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h

class WebAssemblyFunctionInfo final : public MachineFunctionInfo {
   ///   - defined and used in LIFO order with other stack registers
   BitVector VRegStackified;
 
   // A virtual register holding the pointer to the vararg buffer for vararg
   // functions. It is created and set in TLI::LowerFormalArguments and read by
   // TLI::LowerVASTART
   unsigned VarargVreg = -1U;
 
+  // A virtual register holding the base pointer for functions that have
+  // overaligned values on the user stack.
+  unsigned BasePtrVreg = -1U;
+
  public:
   explicit WebAssemblyFunctionInfo(MachineFunction &MF) : MF(MF) {}
   ~WebAssemblyFunctionInfo() override;
 
   void addParam(MVT VT) { Params.push_back(VT); }
   const std::vector<MVT> &getParams() const { return Params; }
 
   void addResult(MVT VT) { Results.push_back(VT); }
   const std::vector<MVT> &getResults() const { return Results; }
 
   void addLocal(MVT VT) { Locals.push_back(VT); }
   const std::vector<MVT> &getLocals() const { return Locals; }
 
   unsigned getVarargBufferVreg() const {
     assert(VarargVreg != -1U && "Vararg vreg hasn't been set");
     return VarargVreg;
   }
   void setVarargBufferVreg(unsigned Reg) { VarargVreg = Reg; }
 
+  unsigned getBasePointerVreg() const {
+    assert(BasePtrVreg != -1U && "Base ptr vreg hasn't been set");
+    return BasePtrVreg;
+  }
+  void setBasePointerVreg(unsigned Reg) { BasePtrVreg = Reg; }
+
   static const unsigned UnusedReg = -1u;
 
   void stackifyVReg(unsigned VReg) {
     assert(MF.getRegInfo().getUniqueVRegDef(VReg));
     if (TargetRegisterInfo::virtReg2Index(VReg) >= VRegStackified.size())
       VRegStackified.resize(TargetRegisterInfo::virtReg2Index(VReg) + 1);
     VRegStackified.set(TargetRegisterInfo::virtReg2Index(VReg));
   }

llvm/trunk/test/CodeGen/WebAssembly/stack-alignment.ll

+; RUN: llc < %s -asm-verbose=false -disable-wasm-fallthrough-return-opt | FileCheck %s
+
+target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
+target triple = "wasm32-unknown-unknown"
+
+declare void @somefunc(i32*)
+
+; CHECK-LABEL: underalign:
+; CHECK:      i32.load  $push[[L1:.+]]=, __stack_pointer{{.+}}
+; CHECK-NEXT: i32.const $push[[L2:.+]]=, 16
+; CHECK-NEXT: i32.sub   $push[[L10:.+]]=, $pop[[L1]], $pop[[L2]]
+; CHECK-NEXT: tee_local $push{{.+}}=, $[[SP:.+]]=, $pop[[L10]]
+
+; CHECK:      i32.add   $push[[underaligned:.+]]=, $[[SP]], $pop{{.+}}
+; CHECK-NEXT: call      somefunc@FUNCTION, $pop[[underaligned]]
+
+; CHECK:      i32.add   $push[[L5:.+]]=, $[[SP]], $pop{{.+}}
+; CHECK-NEXT: i32.store __stack_pointer($pop{{.+}}), $pop[[L5]]
+define void @underalign() {
+entry:
+  %underaligned = alloca i32, align 8
+  call void @somefunc(i32* %underaligned)
+  ret void
+}
+
+; CHECK-LABEL: overalign:
+; CHECK:      i32.load   $push[[L10:.+]]=, __stack_pointer
+; CHECK-NEXT: tee_local  $push[[L9:.+]]=, $[[BP:.+]]=, $pop[[L10]]
+; CHECK-NEXT: i32.const  $push[[L2:.+]]=, 32
+; CHECK-NEXT: i32.sub    $push[[L8:.+]]=, $pop[[L9]], $pop[[L2]]
+; CHECK-NEXT: i32.const  $push[[L3:.+]]=, -32
+; CHECK-NEXT: i32.and    $push[[L7:.+]]=, $pop[[L8]], $pop[[L3]]
+; CHECK-NEXT: tee_local  $push{{.+}}=, $[[SP:.+]]=, $pop[[L7]]
+
+; CHECK:      call       somefunc@FUNCTION, $[[SP]]
+
+; CHECK:      copy_local $push[[L5:.+]]=, $[[BP]]
+; CHECK-NEXT: i32.store  __stack_pointer($pop{{.+}}), $pop[[L5]]
+define void @overalign() {
+entry:
+  %overaligned = alloca i32, align 32
+  call void @somefunc(i32* %overaligned)
+  ret void
+}
+
+; CHECK-LABEL: over_and_normal_align:
+; CHECK:      i32.load   $push[[L14:.+]]=, __stack_pointer
+; CHECK-NEXT: tee_local  $push[[L13:.+]]=, $[[BP:.+]]=, $pop[[L14]]
+; CHECK:      i32.sub    $push[[L12:.+]]=, $pop[[L13]], $pop{{.+}}
+; CHECK:      i32.and    $push[[L11:.+]]=, $pop[[L12]], $pop{{.+}}
+; CHECK-NEXT: tee_local  $push{{.+}}=, $[[SP]]=, $pop[[L11]]
+
+; CHECK:      i32.add    $push[[L6:.+]]=, $[[SP]], $pop{{.+}}
+; CHECK-NEXT: call       somefunc@FUNCTION, $pop[[L6]]
+; CHECK:      i32.add    $push[[L8:.+]]=, $[[SP]], $pop{{.+}}
+; CHECK-NEXT: call       somefunc@FUNCTION, $pop[[L8]]
+
+; CHECK:      copy_local $push[[L9:.+]]=, $[[BP]]
+; CHECK-NEXT: i32.store  __stack_pointer({{.+}}), $pop[[L9]]
+define void @over_and_normal_align() {
+entry:
+  %over = alloca i32, align 32
+  %normal = alloca i32
+  call void @somefunc(i32* %over)
+  call void @somefunc(i32* %normal)
+  ret void
+}
+
+; CHECK-LABEL: dynamic_overalign:
+; CHECK:      i32.load   $push[[L18:.+]]=, __stack_pointer
+; CHECK-NEXT: tee_local  $push[[L17:.+]]=, $[[SP:.+]]=, $pop[[L18]]
+; CHECK-NEXT: copy_local $[[BP:.+]]=, $pop[[L17]]
+; CHECK:      tee_local  $push{{.+}}=, $[[SP_2:.+]]=, $pop{{.+}}
+
+; CHECK:      call       somefunc@FUNCTION, $[[SP_2]]
+
+; CHECK: i32.store __stack_pointer($pop{{.+}}), $[[BP]]
+define void @dynamic_overalign(i32 %num) {
+entry:
+  %dynamic = alloca i32, i32 %num, align 32
+  call void @somefunc(i32* %dynamic)
+  ret void
+}
+
+; CHECK-LABEL: overalign_and_dynamic:
+; CHECK:      i32.load   $push[[L21:.+]]=, __stack_pointer
+; CHECK-NEXT: tee_local  $push[[L20:.+]]=, $[[BP:.+]]=, $pop[[L21]]
+; CHECK:      i32.sub    $push[[L19:.+]]=, $pop[[L20]], $pop{{.+}}
+; CHECK:      i32.and    $push[[L18:.+]]=, $pop[[L19]], $pop{{.+}}
+; CHECK:      tee_local  $push{{.+}}=, $[[FP:.+]]=, $pop[[L18]]
+; CHECK:      i32.sub    $push[[L16:.+]]=, $[[FP]], $pop{{.+}}
+; CHECK-NEXT: tee_local  $push{{.+}}=, $[[SP:.+]]=, $pop[[L16]]
+
+; CHECK:      copy_local $push[[over:.+]]=, $[[FP]]
+; CHECK-NEXT: call       somefunc@FUNCTION, $pop[[over]]
+; CHECK-NEXT: call       somefunc@FUNCTION, $[[SP]]
+
+; CHECK:      copy_local $push[[L12:.+]]=, $[[BP]]
+; CHECK-NEXT: i32.store  __stack_pointer($pop{{.+}}), $pop[[L12]]
+define void @overalign_and_dynamic(i32 %num) {
+entry:
+  %over = alloca i32, align 32
+  %dynamic = alloca i32, i32 %num
+  call void @somefunc(i32* %over)
+  call void @somefunc(i32* %dynamic)
+  ret void
+}
+
+; CHECK-LABEL: overalign_static_and_dynamic:
+; CHECK:      i32.load   $push[[L26:.+]]=, __stack_pointer
+; CHECK-NEXT: tee_local  $push[[L25:.+]]=, $[[BP:.+]]=, $pop[[L26]]
+; CHECK:      i32.sub    $push[[L24:.+]]=, $pop[[L25]], $pop{{.+}}
+; CHECK:      i32.and    $push[[L23:.+]]=, $pop[[L24]], $pop{{.+}}
+; CHECK:      tee_local  $push{{.+}}=, $[[FP:.+]]=, $pop[[L23]]
+; CHECK:      i32.sub    $push[[L21:.+]]=, $[[FP]], $pop{{.+}}
+; CHECK-NEXT: tee_local  $push{{.+}}=, $[[SP:.+]]=, $pop[[L21]]
+
+; CHECK:      copy_local $push[[L19:.+]]=, $[[FP]]
+; CHECK:      tee_local  $push[[L18:.+]]=, $[[FP_2:.+]]=, $pop[[L19]]
+; CHECK:      i32.add    $push[[over:.+]]=, $pop[[L18]], $pop{{.+}}
+; CHECK-NEXT: call       somefunc@FUNCTION, $pop[[over]]
+; CHECK:      call       somefunc@FUNCTION, $[[SP]]
+; CHECK:      i32.add    $push[[static:.+]]=, $[[FP_2]], $pop{{.+}}
+; CHECK-NEXT: call       somefunc@FUNCTION, $pop[[static]]
+
+; CHECK:      copy_local $push[[L16:.+]]=, $[[BP]]
+; CHECK-NEXT: i32.store  __stack_pointer({{.+}}), $pop[[L16]]
+define void @overalign_static_and_dynamic(i32 %num) {
+entry:
+  %over = alloca i32, align 32
+  %dynamic = alloca i32, i32 %num
+  %static = alloca i32
+  call void @somefunc(i32* %over)
+  call void @somefunc(i32* %dynamic)
+  call void @somefunc(i32* %static)
+  ret void
+}