[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

1682

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

5333–5334

Looks like FoldAsLoadDefReg will never be 0 here.

5334

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

1682

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

5333–5334

That's what I thought, glad you agree.

5334

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

1682

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

5334

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

1586

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

1586

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

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;
       }
 
       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()) {
+
+        // We visit each operand even after successfully folding a previous
+        // one.  This allows us to fold multiple loads into a single
+        // instruction.  We do assume that optimizeLoadInstr doesn't insert
+        // foldable uses earlier in the argument list.  Since we don't restart
+        // iteration, we'd miss such cases.
         const MCInstrDesc &MIDesc = MI->getDesc();
-        for (unsigned i = MIDesc.getNumDefs(); i != MIDesc.getNumOperands();
+        for (unsigned i = MIDesc.getNumDefs(); i != MI->getNumOperands();
              ++i) {
           const MachineOperand &MOp = MI->getOperand(i);
           if (!MOp.isReg())
             continue;
           unsigned FoldAsLoadDefReg = MOp.getReg();
           if (FoldAsLoadDefCandidates.count(FoldAsLoadDefReg)) {
             // We need to fold load after optimizeCmpInstr, since
             // optimizeCmpInstr can enable folding by converting SUB to CMP.
for (MachineBasicBlock::iterator MII = MBB.begin(), MIE = MBB.end();
               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.
+              
+              // MI is replaced with FoldMI so we can continue trying to fold
               Changed = true;
-              break;
-            }
-          }
-        }
+              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 {
-  if (FoldAsLoadDefReg == 0)
-    return nullptr;
-  // To be conservative, if there exists another load, clear the load candidate.
-  if (MI.mayLoad()) {
-    FoldAsLoadDefReg = 0;
-    return nullptr;
-  }
-
   // 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/anyregcc.ll

 ; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _property_access2
 ; CHECK-NEXT:   .quad 24
 ; CHECK-NEXT:   .quad _property_access3
 ; CHECK-NEXT:   .quad 24
 ; CHECK-NEXT:   .quad _anyreg_test1
 ; CHECK-NEXT:   .quad 56
 ; CHECK-NEXT:   .quad _anyreg_test2
-; CHECK-NEXT:   .quad 56
+; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _patchpoint_spilldef
 ; CHECK-NEXT:   .quad 56
 ; CHECK-NEXT:   .quad _patchpoint_spillargs
 ; CHECK-NEXT:   .quad 88
 
 ; No constants
 
 ; Callsites

 ; CHECK-NEXT:   .byte 8
 ; CHECK-NEXT:   .short {{[0-9]+}}
 ; CHECK-NEXT:   .long 0
 ; Loc 8: Register
 ; CHECK-NEXT:   .byte 1
 ; CHECK-NEXT:   .byte 8
 ; CHECK-NEXT:   .short {{[0-9]+}}
 ; CHECK-NEXT:   .long 0
-; Loc 9: Register
-; CHECK-NEXT:   .byte 1
-; CHECK-NEXT:   .byte 8
-; CHECK-NEXT:   .short {{[0-9]+}}
-; CHECK-NEXT:   .long 0
-; Loc 10: Register
-; CHECK-NEXT:   .byte 1
-; CHECK-NEXT:   .byte 8
-; CHECK-NEXT:   .short {{[0-9]+}}
-; CHECK-NEXT:   .long 0
-; Loc 11: Register
-; CHECK-NEXT:   .byte 1
-; CHECK-NEXT:   .byte 8
-; CHECK-NEXT:   .short {{[0-9]+}}
-; CHECK-NEXT:   .long 0
-; Loc 12: Register
-; CHECK-NEXT:   .byte 1
-; CHECK-NEXT:   .byte 8
-; CHECK-NEXT:   .short {{[0-9]+}}
-; CHECK-NEXT:   .long 0
-; Loc 13: Register
-; CHECK-NEXT:   .byte 1
-; CHECK-NEXT:   .byte 8
-; CHECK-NEXT:   .short {{[0-9]+}}
-; CHECK-NEXT:   .long 0
+; Loc 9: Argument, still on stack
+; CHECK-NEXT: .byte  3
+; CHECK-NEXT: .byte  8
+; CHECK-NEXT: .short 6
+; CHECK-NEXT: .long
+; Loc 10: Argument, still on stack
+; CHECK-NEXT: .byte  3
+; CHECK-NEXT: .byte  8
+; CHECK-NEXT: .short 6
+; CHECK-NEXT: .long
+; Loc 11: Argument, still on stack
+; CHECK-NEXT: .byte  3
+; CHECK-NEXT: .byte  8
+; CHECK-NEXT: .short 6
+; CHECK-NEXT: .long
+; Loc 12: Argument, still on stack
+; CHECK-NEXT: .byte  3
+; CHECK-NEXT: .byte  8
+; CHECK-NEXT: .short 6
+; CHECK-NEXT: .long
+; Loc 13: Argument, still on stack
+; CHECK-NEXT: .byte  3
+; CHECK-NEXT: .byte  8
+; CHECK-NEXT: .short 6
+; CHECK-NEXT: .long
 define i64 @anyreg_test2(i8* %a1, i8* %a2, i8* %a3, i8* %a4, i8* %a5, i8* %a6, i8* %a7, i8* %a8, i8* %a9, i8* %a10, i8* %a11, i8* %a12, i8* %a13) nounwind ssp uwtable {
 entry:
   %f = inttoptr i64 12297829382473034410 to i8*
   %ret = call anyregcc i64 (i64, i32, i8*, i32, ...) @llvm.experimental.patchpoint.i64(i64 5, i32 15, i8* %f, i32 8, i8* %a1, i8* %a2, i8* %a3, i8* %a4, i8* %a5, i8* %a6, i8* %a7, i8* %a8, i8* %a9, i8* %a10, i8* %a11, i8* %a12, i8* %a13)
   ret i64 %ret
 }
 
 ; Test spilling the return value of an anyregcc call.

test/CodeGen/X86/stackmap.ll

 ; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _propertyWrite
 ; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _jsVoidCall
 ; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _jsIntCall
 ; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _spilledValue
-; CHECK-NEXT:   .quad 56
+; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _spilledStackMapValue
-; CHECK-NEXT:   .quad 56
+; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _spillSubReg
 ; CHECK-NEXT:   .quad 56
 ; CHECK-NEXT:   .quad _subRegOffset
 ; CHECK-NEXT:   .quad 56
 ; CHECK-NEXT:   .quad _liveConstant
 ; CHECK-NEXT:   .quad 8
 ; CHECK-NEXT:   .quad _directFrameIdx
 ; CHECK-NEXT:   .quad 56

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