[WebAssembly] Fix scheduling dependencies in register-stackified code

[sunfish]

Summary

Add physical register defs to instructions used from stackified instructions to prevent them from being scheduled into the middle of a stack sequence. This is a conservative measure which may be loosened in the future.

[jfb]

lgtm

ncbray and I had a little geek-out about coloring combined with tree formation. Fun problem.

lib/Target/WebAssembly/WebAssemblyRegStackify.cpp

75

"on"

lib/Target/WebAssembly/WebAssemblyRegStackify.cpp

 } // end anonymous namespace
 
 char WebAssemblyRegStackify::ID = 0;
 FunctionPass *llvm::createWebAssemblyRegStackify() {
   return new WebAssemblyRegStackify();
 }
 
 // Decorate the given instruction with implicit operands that enforce the
-// expression stack ordering constraints.
-static void ImposeStackOrdering(MachineInstr *MI) {
-  // Read and write the opaque EXPR_STACK register.
+// expression stack ordering constraints needed for an instruction which is
+// consumed by an instruction using the expression stack.
+static void ImposeStackInputOrdering(MachineInstr *MI) {
+  // Write the opaque EXPR_STACK register.
+  if (!MI->definesRegister(WebAssembly::EXPR_STACK))
     MI->addOperand(MachineOperand::CreateReg(WebAssembly::EXPR_STACK,
                                              /*isDef=*/true,
                                              /*isImp=*/true));
+}
+
+// Decorate the given instruction with implicit operands that enforce the
+// expression stack ordering constraints for an instruction which is one
+// the expression stack.
+static void ImposeStackOrdering(MachineInstr *MI, MachineRegisterInfo &MRI) {
+  ImposeStackInputOrdering(MI);
+
+  // Also read the opaque EXPR_STACK register.
   MI->addOperand(MachineOperand::CreateReg(WebAssembly::EXPR_STACK,
                                            /*isDef=*/false,
                                            /*isImp=*/true));
+
+  // Also, mark any inputs to this instruction as being consumed by an
+  // instruction on the expression stack.
+  // TODO: Find a lighter way to describe the appropriate constraints.
+  for (MachineOperand &MO : MI->uses()) {
+    if (!MO.isReg())
+      continue;
+    unsigned Reg = MO.getReg();
+    if (!TargetRegisterInfo::isVirtualRegister(Reg))
+      continue;
+    MachineInstr *Def = MRI.getVRegDef(Reg);
+    if (Def->getOpcode() == TargetOpcode::PHI)
+      continue;
+    ImposeStackInputOrdering(Def);
+  }
 }
 
 // Test whether it's safe to move Def to just before Insert. Note that this
 // doesn't account for physical register dependencies, because WebAssembly
 // doesn't have any (other than special ones like EXPR_STACK).
 // 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
for (MachineInstr &MI : reverse(MBB)) {
       // operands off the stack in LIFO order.
       bool AnyStackified = false;
       for (MachineOperand &Op : reverse(Insert->uses())) {
         // We're only interested in explicit virtual register operands.
         if (!Op.isReg() || Op.isImplicit() || !Op.isUse())
           continue;
 
         unsigned Reg = Op.getReg();
-        if (!TargetRegisterInfo::isVirtualRegister(Reg))
+        if (TargetRegisterInfo::isPhysicalRegister(Reg)) {
+          // An instruction with a physical register. Conservatively mark it as
+          // an expression stack input so that it isn't reordered with anything
+          // in an expression stack which might use it (physical registers
+          // aren't in SSA form so it's not trivial to determine this).
+          // TODO: Be less conservative.
+          ImposeStackInputOrdering(Insert);
           continue;
+        }
 
         // Only consider registers with a single definition.
         // TODO: Eventually we may relax this, to stackify phi transfers.
         MachineInstr *Def = MRI.getVRegDef(Reg);
         if (!Def)
           continue;
 
         // There's no use in nesting implicit defs inside anything.
for (MachineInstr &MI : reverse(MBB)) {
           continue;
 
         Changed = true;
         AnyStackified = true;
         // Move the def down and nest it in the current instruction.
         MBB.insert(MachineBasicBlock::instr_iterator(Insert),
                    Def->removeFromParent());
         MFI.stackifyVReg(Reg);
-        ImposeStackOrdering(Def);
+        ImposeStackOrdering(Def, MRI);
         Insert = Def;
       }
       if (AnyStackified)
-        ImposeStackOrdering(&MI);
+        ImposeStackOrdering(&MI, MRI);
     }
   }
 
   // If we used EXPR_STACK anywhere, add it to the live-in sets everywhere
   // so that it never looks like a use-before-def.
   if (Changed) {
     MF.getRegInfo().addLiveIn(WebAssembly::EXPR_STACK);
     for (MachineBasicBlock &MBB : MF)

test/CodeGen/WebAssembly/load-store-i1.ll

 ; CHECK-NEXT: return $pop[[NUM0]]{{$}}
 define i32 @load_u_i1_i32(i1* %p) {
   %v = load i1, i1* %p
   %e = zext i1 %v to i32
   ret i32 %e
 }
 
 ; CHECK-LABEL: load_s_i1_i32:
-; CHECK:      i32.load8_u $push[[NUM0:[0-9]+]]=, $0{{$}}
-; CHECK-NEXT: i32.const $[[NUM1:[0-9]+]]=, 31{{$}}
+; CHECK: i32.const $[[NUM1:[0-9]+]]=, 31{{$}}
+; CHECK-NEXT: i32.load8_u $push[[NUM0:[0-9]+]]=, $0{{$}}
 ; CHECK-NEXT: shl $push[[NUM2:[0-9]+]]=, $pop[[NUM0]], $[[NUM1]]{{$}}
 ; CHECK-NEXT: shr_s $push[[NUM3:[0-9]+]]=, $pop[[NUM2]], $[[NUM1]]{{$}}
 ; CHECK-NEXT: return $pop[[NUM3]]{{$}}
 define i32 @load_s_i1_i32(i1* %p) {
   %v = load i1, i1* %p
   %e = sext i1 %v to i32
   ret i32 %e
 }
 
 ; CHECK-LABEL: load_u_i1_i64:
 ; CHECK:      i64.load8_u $push[[NUM0:[0-9]+]]=, $0{{$}}
 ; CHECK-NEXT: return $pop[[NUM0]]{{$}}
 define i64 @load_u_i1_i64(i1* %p) {
   %v = load i1, i1* %p
   %e = zext i1 %v to i64
   ret i64 %e
 }
 
 ; CHECK-LABEL: load_s_i1_i64:
-; CHECK:      i64.load8_u $push[[NUM0:[0-9]+]]=, $0{{$}}
-; CHECK-NEXT: i64.const $[[NUM1:[0-9]+]]=, 63{{$}}
+; CHECK: i64.const $[[NUM1:[0-9]+]]=, 63{{$}}
+; CHECK-NEXT: i64.load8_u $push[[NUM0:[0-9]+]]=, $0{{$}}
 ; CHECK-NEXT: shl $push[[NUM2:[0-9]+]]=, $pop[[NUM0]], $[[NUM1]]{{$}}
 ; CHECK-NEXT: shr_s $push[[NUM3:[0-9]+]]=, $pop[[NUM2]], $[[NUM1]]{{$}}
 ; CHECK-NEXT: return $pop[[NUM3]]{{$}}
 define i64 @load_s_i1_i64(i1* %p) {
   %v = load i1, i1* %p
   %e = sext i1 %v to i64
   ret i64 %e
 }

test/CodeGen/WebAssembly/reg-stackify.ll

 
 ; CHECK-LABEL: yes1:
 ; CHECK: return $pop0{{$}}
 define i32 @yes1(i32* %q) {
   %t = load volatile i32, i32* %q
   ret i32 %t
 }
 
+; Don't schedule stack uses into the stack. To reduce register pressure, the
+; scheduler might be tempted to move the definition of $2 down. However, this
+; would risk getting incorrect liveness if the instructions are later
+; rearranged to make the stack contiguous.
+
+; CHECK-LABEL: stack_uses:
+; CHECK-NEXT: .param i32{{$}}
+; CHECK-NEXT: .result i32{{$}}
+; CHECK-NEXT: local i32, i32{{$}}
+; CHECK-NEXT: i32.const       $1=, 1{{$}}
+; CHECK-NEXT: i32.const       $2=, 0{{$}}
+; CHECK-NEXT: i32.and         $push0=, $0, $1{{$}}
+; CHECK-NEXT: i32.eq          $push1=, $pop0, $2{{$}}
+; CHECK-NEXT: block           BB4_2{{$}}
+; CHECK-NEXT: br_if           $pop1, BB4_2{{$}}
+; CHECK-NEXT: return          $2{{$}}
+; CHECK-NEXT: BB4_2:{{$}}
+; CHECK-NEXT: return          $1{{$}}
+define i32 @stack_uses(i32 %x) {
+entry:
+  %c = trunc i32 %x to i1
+  br i1 %c, label %true, label %false
+true:
+  ret i32 0
+false:
+  ret i32 1
+}
+
 !0 = !{}