[WebAssembly] Make register stackification more conservative

[jgravelle-google]

Summary

Register stackification currently checks VNInfo for changes. Make that
more accurate by testing each intervening instruction for any other defs
to the same virtual register.

[sunfish]

lgtm.

lib/Target/WebAssembly/WebAssemblyRegStackify.cpp

281

Instead of holding MachineOperand*, this could just hold unsigned register numbers, which would simplify some of the code below. Also, this is a good candidate for a SmallVector, with a size of something like 4.

307

MRI.hasOneDef(Reg) would make this a little clearer than using getUniqueVRegDef.

316–317

!MutableOperands.empty() would be a little more concise than .size() > 0.

317

Ah, the bug is the case where InsVNI is null. The register may not be live at the insert point, however it may still be clobbered between Def and the insert point.

It's unfortunate then that LiveIntervals doesn't seem to have a convenient way to ask "is this value clobbered anywhere between here and there" if it isn't already live there. But given that, this was the only use of LIS in this function, so you can remove the LIS argument now too.

[jgravelle-google]

lib/Target/WebAssembly/WebAssemblyRegStackify.cpp

281

Indeed it does.

And I didn't think about SmallVector. Thanks

307

That it would.

316–317

Right, forgot that was a thing. Thanks

317

Exactly.

And, nice catch on the LIS.

340

Looking at this again, I could reorder this to only loop over the mutable operands if the instruction's operands are defs, which would be slightly more performant. I'll stick with this because it's more readable.

[jgravelle-google]

• Code review feedback

lib/Target/WebAssembly/WebAssemblyRegStackify.cpp

 }
 
 // Test whether it's safe to move Def to just before Insert.
 // TODO: Compute memory dependencies in a way that doesn't require always
 // walking the block.
 // TODO: Compute memory dependencies in a way that uses AliasAnalysis to be
 // more precise.
 static bool IsSafeToMove(const MachineInstr *Def, const MachineInstr *Insert,
-                         AliasAnalysis &AA, const LiveIntervals &LIS,
-                         const MachineRegisterInfo &MRI) {
+                         AliasAnalysis &AA, const MachineRegisterInfo &MRI) {
   assert(Def->getParent() == Insert->getParent());
 
   // Check for register dependencies.
+  SmallVector<unsigned, 4> MutableRegisters;
   for (const MachineOperand &MO : Def->operands()) {
     if (!MO.isReg() || MO.isUndef())
       continue;
     unsigned Reg = MO.getReg();
 
     // If the register is dead here and at Insert, ignore it.
     if (MO.isDead() && Insert->definesRegister(Reg) &&
         !Insert->readsRegister(Reg))
       continue;
 
     if (TargetRegisterInfo::isPhysicalRegister(Reg)) {
       // Ignore ARGUMENTS; it's just used to keep the ARGUMENT_* instructions
       // from moving down, and we've already checked for that.
       if (Reg == WebAssembly::ARGUMENTS)
         continue;
       // If the physical register is never modified, ignore it.
       if (!MRI.isPhysRegModified(Reg))
         continue;
       // Otherwise, it's a physical register with unknown liveness.
       return false;
     }
 
-    // Ask LiveIntervals whether moving this virtual register use or def to
-    // Insert will change which value numbers are seen.
-    //
-    // If the operand is a use of a register that is also defined in the same
-    // instruction, test that the newly defined value reaches the insert point,
-    // since the operand will be moving along with the def.
-    const LiveInterval &LI = LIS.getInterval(Reg);
-    VNInfo *DefVNI =
-        (MO.isDef() || Def->definesRegister(Reg)) ?
-        LI.getVNInfoAt(LIS.getInstructionIndex(*Def).getRegSlot()) :
-        LI.getVNInfoBefore(LIS.getInstructionIndex(*Def));
-    assert(DefVNI && "Instruction input missing value number");
-    VNInfo *InsVNI = LI.getVNInfoBefore(LIS.getInstructionIndex(*Insert));
-    if (InsVNI && DefVNI != InsVNI)
-      return false;
+    // If one of the operands isn't in SSA form, it has different values at
+    // different times, and we need to make sure we don't move our use across
+    // a different def.
+    if (!MO.isDef() && !MRI.hasOneDef(Reg))
+      MutableRegisters.push_back(Reg);
   }
 
   bool Read = false, Write = false, Effects = false, StackPointer = false;
   Query(*Def, AA, Read, Write, Effects, StackPointer);
 
   // If the instruction does not access memory and has no side effects, it has
   // no additional dependencies.
-  if (!Read && !Write && !Effects && !StackPointer)
+  bool HasMutableRegisters = !MutableRegisters.empty();
+  if (!Read && !Write && !Effects && !StackPointer && !HasMutableRegisters)
     return true;
 
   // Scan through the intervening instructions between Def and Insert.
   MachineBasicBlock::const_iterator D(Def), I(Insert);
   for (--I; I != D; --I) {
     bool InterveningRead = false;
     bool InterveningWrite = false;
     bool InterveningEffects = false;
     bool InterveningStackPointer = false;
     Query(*I, AA, InterveningRead, InterveningWrite, InterveningEffects,
           InterveningStackPointer);
     if (Effects && InterveningEffects)
       return false;
     if (Read && InterveningWrite)
       return false;
     if (Write && (InterveningRead || InterveningWrite))
       return false;
     if (StackPointer && InterveningStackPointer)
       return false;
+
+    for (unsigned Reg : MutableRegisters)
+      for (const MachineOperand &MO : I->operands())
+        if (MO.isReg() && MO.isDef() && MO.getReg() == Reg)
+          return false;
   }
 
   return true;
 }
 
 /// Test whether OneUse, a use of Reg, dominates all of Reg's other uses.
 static bool OneUseDominatesOtherUses(unsigned Reg, const MachineOperand &OneUse,
                                      const MachineBasicBlock &MBB,
for (auto MII = MBB.rbegin(); MII != MBB.rend(); ++MII) {
 
         // Decide which strategy to take. Prefer to move a single-use value
         // over cloning it, and prefer cloning over introducing a tee_local.
         // For moving, we require the def to be in the same block as the use;
         // this makes things simpler (LiveIntervals' handleMove function only
         // supports intra-block moves) and it's MachineSink's job to catch all
         // the sinking opportunities anyway.
         bool SameBlock = Def->getParent() == &MBB;
-        bool CanMove = SameBlock && IsSafeToMove(Def, Insert, AA, LIS, MRI) &&
+        bool CanMove = SameBlock && IsSafeToMove(Def, Insert, AA, MRI) &&
                        !TreeWalker.IsOnStack(Reg);
         if (CanMove && HasOneUse(Reg, Def, MRI, MDT, LIS)) {
           Insert = MoveForSingleUse(Reg, Op, Def, MBB, Insert, LIS, MFI, MRI);
         } else if (ShouldRematerialize(*Def, AA, TII)) {
           Insert =
               RematerializeCheapDef(Reg, Op, *Def, MBB, Insert->getIterator(),
                                     LIS, MFI, MRI, TII, TRI);
         } else if (CanMove &&

test/CodeGen/WebAssembly/userstack.ll

 
  ; Writeback to memory.
  ; CHECK: i32.const $push[[L17:.+]]=, 16
  ; CHECK-NEXT: i32.add $push[[L18:.+]]=, $[[FP]], $pop[[L17]]
  ; CHECK-NEXT: i32.store $drop=, __stack_pointer($pop{{.+}}), $pop[[L18]]
  ret void
 }
 
+declare i8* @llvm.stacksave()
+declare void @llvm.stackrestore(i8*)
+
+; CHECK-LABEL: llvm_stack_builtins:
+define void @llvm_stack_builtins(i32 %alloc) noredzone {
+ ; CHECK: i32.load $push[[L11:.+]]=, __stack_pointer($pop{{.+}})
+ ; CHECK-NEXT: tee_local $push[[L10:.+]]=, ${{.+}}=, $pop[[L11]]
+ ; CHECK-NEXT: copy_local $[[STACK:.+]]=, $pop[[L10]]
+ %stack = call i8* @llvm.stacksave()
+
+ ; Ensure we don't reassign the stacksave local
+ ; CHECK-NOT: $[[STACK]]=
+ %dynamic = alloca i32, i32 %alloc
+
+ ; CHECK: i32.store $drop=, __stack_pointer($pop{{.+}}), $[[STACK]]
+ call void @llvm.stackrestore(i8* %stack)
+
+ ret void
+}
+
+; Not actually using the alloca'd variables exposed an issue with register
+; stackification, where copying the stack pointer into the frame pointer was
+; moved after the stack pointer was updated for the dynamic alloca.
+; CHECK-LABEL: dynamic_alloca_nouse:
+define void @dynamic_alloca_nouse(i32 %alloc) noredzone {
+ ; CHECK: i32.load $push[[L11:.+]]=, __stack_pointer($pop{{.+}})
+ ; CHECK-NEXT: tee_local $push[[L10:.+]]=, ${{.+}}=, $pop[[L11]]
+ ; CHECK-NEXT: copy_local $[[FP:.+]]=, $pop[[L10]]
+ %dynamic = alloca i32, i32 %alloc
+
+ ; CHECK-NOT: $[[FP]]=,
+
+ ; CHECK: i32.store $drop=, __stack_pointer($pop{{.+}}), $[[FP]]
+ ret void
+}
+
 ; The use of the alloca in a phi causes a CopyToReg DAG node to be generated,
 ; which has to have special handling because CopyToReg can't have a FI operand
 ; CHECK-LABEL: copytoreg_fi:
 define void @copytoreg_fi(i1 %cond, i32* %b) {
 entry:
  ; CHECK: i32.const $push[[L1:.+]]=, 16
  ; CHECK-NEXT: i32.sub $push[[L3:.+]]=, {{.+}}, $pop[[L1]]
  %addr = alloca i32