[peephole] Enhance folding logic to work for STATEPOINTs

[reames]

Summary

The general idea here is to get enough of the existing restrictions out of the way that the already existing folding logic in foldMemoryOperand can kick in for STATEPOINTs and fold references to immutable stack slots. The key changes are:

• Support for folding multiple operands at once which reference the same load
• Support for folding multiple loads into a single instruction
• Walk all the operands of the instruction for varidic instructions (this is a bug fix!)

Once this lands, I'll post another patch which refactors the TII interface here. There's nothing actually x86 specific about the x86 code used here.

[wmi]

lib/CodeGen/PeepholeOptimizer.cpp

1683

Why recursive folding is needed here? I guess it is trying to fold the case that the address of a load is another load? If that is true, can we put the use operands of newly folded load into a set and only visit the set in the next iteration?

lib/Target/X86/X86InstrInfo.cpp

5336

Looks like FoldAsLoadDefReg will never be 0 here.

5348

If MI is load, foldMemoryOperand will return nullptr in the end after it goes through a lot of effort. I agree with you it is better to remove the check here and add an early exit in foldMemoryOperand.

[reames]

lib/CodeGen/PeepholeOptimizer.cpp

1683

The case I'm trying to handle is a single instruction which can fold multiple distinct loads into it. The problem is that when I replace MI, I need to restart the operand iteration.

lib/Target/X86/X86InstrInfo.cpp

5336

That's what I thought, glad you agree.

5348

It will? I don't see why a instruction which can load definitely can't fold another load into itself. An instruction set with an "reg = add mem, mem" instruction would be a great example. Just because an "reg = add mem, reg" form loads doesn't mean you're done folding.

[wmi]

lib/CodeGen/PeepholeOptimizer.cpp

1683

I see. An idea is at the second iteration, don't have to restart the index i from the first use operand. we can start from "the last value of i in the previous iteration" + 1, i.e., keep i monotonically increasing from one iteration to the next, so the complexity is O(n).

lib/Target/X86/X86InstrInfo.cpp

5348

In X86InstrInfo::foldMemoryOperandImpl, if the opcode of current instruction is not in the table pointed by OpcodeTablePtr, it means it is impossible to get a valid instruction after folding, and foldMemoryOperandImpl will return nullptr. Like MemoryFoldTable1 specifies the set of instructions which can have their operand1 replaced by a mem operand.

[reames]

A real patch now that we've settled the design questions.

[reames]

ping

[wmi]

lib/CodeGen/PeepholeOptimizer.cpp

1589

can we keep the continue and just let statepoint builtin to slip through? It is a little worrysome to let inlineasm or other instructions with sideeffect to try folding.

[reames]

lib/CodeGen/PeepholeOptimizer.cpp

1589

We won't fold across anything considered a isLoadFoldBarrier(). This is handled below. Is your concern that the folding logic might crash when folding a use into an instruction with side effects? If so, that would seem like a separate bug we should fix.

I am strongly resistant to special casing statepoint where generic code should work instead.

[wmi]

LGTM.

lib/CodeGen/PeepholeOptimizer.cpp

bool PeepholeOptimizer::isLoadFoldable(
 
   unsigned Reg = MI->getOperand(0).getReg();
   // To reduce compilation time, we check MRI->hasOneNonDBGUse when inserting
   // loads. It should be checked when processing uses of the load, since
   // uses can be removed during peephole.
   if (!MI->getOperand(0).getSubReg() &&
       TargetRegisterInfo::isVirtualRegister(Reg) &&
       MRI->hasOneNonDBGUse(Reg)) {
+    //dbgs() << Reg << "might be forldable\n";
     FoldAsLoadDefCandidates.insert(Reg);
     return true;
   }
   return false;
 }
 
 bool PeepholeOptimizer::isMoveImmediate(
     MachineInstr *MI, SmallSet<unsigned, 4> &ImmDefRegs,
for (MachineBasicBlock::iterator MII = MBB.begin(), MIE = MBB.end();
       // We may be erasing MI below, increment MII now.
       ++MII;
       LocalMIs.insert(MI);
 
       // Skip debug values. They should not affect this peephole optimization.
       if (MI->isDebugValue())
           continue;
 
-      // If we run into an instruction we can't fold across, discard
-      // the load candidates.
-      if (MI->isLoadFoldBarrier())
-        FoldAsLoadDefCandidates.clear();
-
       if (MI->isPosition() || MI->isPHI())
         continue;
 
       if (!MI->isCopy()) {
         for (const auto &Op : MI->operands()) {
           // Visit all operands: definitions can be implicit or explicit.
           if (Op.isReg()) {
             unsigned Reg = Op.getReg();
for (MachineBasicBlock::iterator MII = MBB.begin(), MIE = MBB.end();
       if (MI->isInlineAsm() || MI->hasUnmodeledSideEffects()) {
         // Blow away all non-allocatable physical registers knowledge since we
         // don't know what's correct anymore.
         //
         // FIXME: handle explicit asm clobbers.
         DEBUG(dbgs() << "NAPhysCopy: blowing away all info due to " << *MI
                      << '\n');
         NAPhysToVirtMIs.clear();
-        continue;
+        // FOR DISCUSSION - Is it legal to delete this continue?  
+        //        continue;
       }
 
       if ((isUncoalescableCopy(*MI) &&
            optimizeUncoalescableCopy(MI, LocalMIs)) ||
           (MI->isCompare() && optimizeCmpInstr(MI, &MBB)) ||
           (MI->isSelect() && optimizeSelect(MI, LocalMIs))) {
         // MI is deleted.
         LocalMIs.erase(MI);
         Changed = true;
for (MachineBasicBlock::iterator MII = MBB.begin(), MIE = MBB.end();
           Changed |= foldImmediate(MI, &MBB, ImmDefRegs, ImmDefMIs);
       }
 
       // Check whether MI is a load candidate for folding into a later
       // instruction. If MI is not a candidate, check whether we can fold an
       // earlier load into MI.
       if (!isLoadFoldable(MI, FoldAsLoadDefCandidates) &&
           !FoldAsLoadDefCandidates.empty()) {
+
+        auto tryOptimizeLoad = [&]() {
           const MCInstrDesc &MIDesc = MI->getDesc();
-        for (unsigned i = MIDesc.getNumDefs(); i != MIDesc.getNumOperands();
+          for (unsigned i = MIDesc.getNumDefs(); i != MI->getNumOperands();
                ++i) {
+            //dbgs() << i << "\n";
             const MachineOperand &MOp = MI->getOperand(i);
+            //dbgs() << MOp << "\n";
             if (!MOp.isReg())
               continue;
             unsigned FoldAsLoadDefReg = MOp.getReg();
+            //dbgs() << "Candidate: " << FoldAsLoadDefReg << "\n";
             if (FoldAsLoadDefCandidates.count(FoldAsLoadDefReg)) {
+              //dbgs() << "Found candidate to fold\n";
               // We need to fold load after optimizeCmpInstr, since
               // optimizeCmpInstr can enable folding by converting SUB to CMP.
               // Save FoldAsLoadDefReg because optimizeLoadInstr() resets it and
               // we need it for markUsesInDebugValueAsUndef().
               unsigned FoldedReg = FoldAsLoadDefReg;
               MachineInstr *DefMI = nullptr;
               if (MachineInstr *FoldMI =
-                    TII->optimizeLoadInstr(*MI, MRI, FoldAsLoadDefReg, DefMI)) {
                 // Update LocalMIs since we replaced MI with FoldMI and deleted
                 // DefMI.
                 DEBUG(dbgs() << "Replacing: " << *MI);
                 DEBUG(dbgs() << "     With: " << *FoldMI);
                 LocalMIs.erase(MI);
                 LocalMIs.erase(DefMI);
                 LocalMIs.insert(FoldMI);
                 MI->eraseFromParent();
                 DefMI->eraseFromParent();
                 MRI->markUsesInDebugValueAsUndef(FoldedReg);
                 FoldAsLoadDefCandidates.erase(FoldedReg);
                 ++NumLoadFold;
-              // MI is replaced with FoldMI.
-              Changed = true;
-              break;
-            }
-          }
-        }
+                return FoldMI;
+              }
+            }
+          }
+          return (MachineInstr*)nullptr;
+        };
+
+        // Keep trying to optimize the load until we can't fold any further.
+        // FOR DISCUSSION
+        // Note: This is O(n^2) in the number of operands if
+        // TII->optimizeLoadInstr were to fail for every other operand.
+        while(true) {
+          auto *FoldMI = tryOptimizeLoad();
+          if (!FoldMI)
+            break; // No folding occurred
+          
+          MI = FoldMI;
+      // If we run into an instruction we can't fold across, discard
+      // the load candidates.  Note: We might be able to fold *into* this
+      // instruction, so this needs to be after the folding logic.
+      if (MI->isLoadFoldBarrier()) {
+        DEBUG(dbgs() << "Encountered load fold barrier on " << *MI << "\n");
+        FoldAsLoadDefCandidates.clear();
       }
+
     }
   }
 
   return Changed;
 }
 
 ValueTrackerResult ValueTracker::getNextSourceFromCopy() {
   assert(Def->isCopy() && "Invalid definition");

lib/Target/X86/X86InstrInfo.cpp

 /// Try to remove the load by folding it to a register
 /// operand at the use. We fold the load instructions if load defines a virtual
 /// register, the virtual register is used once in the same BB, and the
 /// instructions in-between do not load or store, and have no side effects.
 MachineInstr *X86InstrInfo::optimizeLoadInstr(MachineInstr &MI,
                                               const MachineRegisterInfo *MRI,
                                               unsigned &FoldAsLoadDefReg,
                                               MachineInstr *&DefMI) const {
+  // FOR DISCUSSION - Is this left over from when we had a single variable as
+  // opposed to a set of FoldAsLoadDefs?  Can we remove?
   if (FoldAsLoadDefReg == 0)
     return nullptr;
+
+  // FOR DISCUSSION - Do we need this code at all?  If so, why?  Shouldn't
+  // foldMemoryOperand be robust against this case?
+#if 0
   // To be conservative, if there exists another load, clear the load candidate.
-  if (MI.mayLoad()) {
+  if (MI.mayLoad() &&
+      MI.getOpcode() != TargetOpcode::STATEPOINT) {
     FoldAsLoadDefReg = 0;
     return nullptr;
   }
+#endif
 
   // Check whether we can move DefMI here.
   DefMI = MRI->getVRegDef(FoldAsLoadDefReg);
   assert(DefMI);
   bool SawStore = false;
   if (!DefMI->isSafeToMove(nullptr, SawStore))
     return nullptr;
 
   // Collect information about virtual register operands of MI.
-  unsigned SrcOperandId = 0;
-  bool FoundSrcOperand = false;
-  for (unsigned i = 0, e = MI.getDesc().getNumOperands(); i != e; ++i) {
+  SmallVector<unsigned, 1> SrcOperandIds;
+  for (unsigned i = 0, e = MI.getNumOperands(); i != e; ++i) {
     MachineOperand &MO = MI.getOperand(i);
     if (!MO.isReg())
       continue;
     unsigned Reg = MO.getReg();
     if (Reg != FoldAsLoadDefReg)
       continue;
-    // Do not fold if we have a subreg use or a def or multiple uses.
-    if (MO.getSubReg() || MO.isDef() || FoundSrcOperand)
+    // Do not fold if we have a subreg use or a def.
+    if (MO.getSubReg() || MO.isDef())
       return nullptr;
-
-    SrcOperandId = i;
-    FoundSrcOperand = true;
+    SrcOperandIds.push_back(i);
   }
-  if (!FoundSrcOperand)
+  if (SrcOperandIds.empty())
     return nullptr;
 
   // Check whether we can fold the def into SrcOperandId.
-  if (MachineInstr *FoldMI = foldMemoryOperand(MI, SrcOperandId, *DefMI)) {
+  if (MachineInstr *FoldMI = foldMemoryOperand(MI, SrcOperandIds, *DefMI)) {
     FoldAsLoadDefReg = 0;
     return FoldMI;
   }
 
   return nullptr;
 }
 
 /// Expand a single-def pseudo instruction to a two-addr

test/CodeGen/X86/statepoint-live-in.ll

; CHECK-NEXT: callq _bar
   call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 2882400000, i32 0, void ()* @bar, i32 0, i32 2, i32 0, i32 2, i32 %b, i32 %a)
   ret void
 }
 
 define void @test3(i32 %a, i32 %b, i32 %c, i32 %d, i32 %e, i32 %f, i32 %g, i32 %h, i32 %i) gc "statepoint-example" {
 entry:
 ; TODO: We should have folded the reload into the statepoint.
 ; CHECK-LABEL: @test3
-; CHECK:        movl  32(%rsp), %r10d
-; CHECK-NEXT:   movl  24(%rsp), %r11d
-; CHECK-NEXT:   movl  16(%rsp), %eax
+; CHECK:        pushq %rax
+; CHECK-NEXT:   Ltmp9:
+; CHECK-NEXT:   .cfi_def_cfa_offset 16
 ; CHECK-NEXT:   callq _bar
   %statepoint_token1 = call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 2882400000, i32 0, void ()* @bar, i32 0, i32 2, i32 0, i32 9, i32 %a, i32 %b, i32 %c, i32 %d, i32 %e, i32 %f, i32 %g, i32 %h, i32 %i)
   ret void
 }
 
 ; This case just confirms that we don't crash when given more live values
-; than registers.  This is a case where we *have* to use a stack slot.
+; than registers.  This is a case where we *have* to use a stack slot.  This
+; also ends up being a good test of whether we can fold loads from immutable
+; stack slots into the statepoint.
 define void @test4(i32 %a, i32 %b, i32 %c, i32 %d, i32 %e, i32 %f, i32 %g, i32 %h, i32 %i, i32 %j, i32 %k, i32 %l, i32 %m, i32 %n, i32 %o, i32 %p, i32 %q, i32 %r, i32 %s, i32 %t, i32 %u, i32 %v, i32 %w, i32 %x, i32 %y, i32 %z) gc "statepoint-example" {
 entry:
-; TODO: We should have folded the reload into the statepoint.  
 ; CHECK-LABEL: test4
-; CHECK: pushq %r15
-; CHECK: pushq %r14
-; CHECK: pushq %r13
-; CHECK: pushq %r12
-; CHECK: pushq %rbx
-; CHECK: pushq %rax
-; CHECK:        movl  128(%rsp), %r13d
-; CHECK-NEXT:   movl  120(%rsp), %r12d
-; CHECK-NEXT:   movl  112(%rsp), %r15d
-; CHECK-NEXT:   movl  104(%rsp), %r14d
-; CHECK-NEXT:   movl  96(%rsp), %ebp
-; CHECK-NEXT:   movl  88(%rsp), %ebx
-; CHECK-NEXT:   movl  80(%rsp), %r11d
-; CHECK-NEXT:   movl  72(%rsp), %r10d
-; CHECK-NEXT:   movl  64(%rsp), %eax
+; CHECK:        pushq %rax
+; CHECK-NEXT:   Ltmp11:
+; CHECK-NEXT:   .cfi_def_cfa_offset 16
 ; CHECK-NEXT:   callq _bar
   %statepoint_token1 = call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 2882400000, i32 0, void ()* @bar, i32 0, i32 2, i32 0, i32 26, i32 %a, i32 %b, i32 %c, i32 %d, i32 %e, i32 %f, i32 %g, i32 %h, i32 %i, i32 %j, i32 %k, i32 %l, i32 %m, i32 %n, i32 %o, i32 %p, i32 %q, i32 %r, i32 %s, i32 %t, i32 %u, i32 %v, i32 %w, i32 %x, i32 %y, i32 %z)
   ret void
 }
 
 ; A live-through gc-value must be spilled even if it is also a live-in deopt
 ; value.  For live-in, we could technically report the register copy, but from
 ; a code quality perspective it's better to reuse the required stack slot so 

 ; Show the interaction of live-through spilling followed by live-in.
 define void @test6(i32 %a) gc "statepoint-example" {
 entry:
 ; TODO: We could have reused the previous spill slot at zero additional cost.
 ; CHECK-LABEL: test6
 ; CHECK:        movl %edi, %ebx
 ; CHECK:        movl %ebx, 12(%rsp)
 ; CHECK-NEXT:   callq _baz
-; CHECK-NEXT:  Ltmp30:
+; CHECK-NEXT:  Ltmp
 ; CHECK-NEXT:   callq _bar
   call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 2882400000, i32 0, void ()* @baz, i32 0, i32 0, i32 0, i32 1, i32 %a)
   call token (i64, i32, void ()*, i32, i32, ...) @llvm.experimental.gc.statepoint.p0f_isVoidf(i64 2882400000, i32 0, void ()* @bar, i32 0, i32 2, i32 0, i32 1, i32 %a)
   ret void
 }
 
 
 ; CHECK: Ltmp1-_test1