[WIP][shrink-wrap]Sink COPYs to CSR from entry to successors

[junbuml]

Summary

In AArch64, ISel create Copy instructions to move function parameters to virtual registers in the entry block. If the parameters are live across a function call, then RA would allocate CSRs for such Copys in the entry.

We should be able to sink such Copy instructions if there is no use in the entry block, but MachineSink (before RA) cannot sink such Copys because SreReg is an allocatable PhyReg. After RA, there is no pass which try to sink such Copy instruction, so shrink-wrapping pass often miss opportunity because the first def of CSR is in the entry block.

This change try to encourage more shrink-wrapping by sinking Copys of which DestReg is CSR from the entry to successors if there there is no use in the entry.

In the IR below, one of the parameter (%paramAcrossCall) live across a function call, but it's only used in %BB0. Currently, in AArch64, we create a "COPY %X19, %W0 " in the entry for this parameter, which prevent this function from being shrink-wrapped. Since this COPY in the entry block is not used in the entry, we can sink it to %BB0, then this function will be shrink-wrapped.

define i32 @shrinkwrapme(i32 %paramAcrossCall, i32 %paramNotAcrossCall) {
entry:
  %cmp5 = icmp sgt i32 %paramNotAcrossCall, 0
  br i1 %cmp5, label %BB0, label %Exit
BB0:
  %call = call i32 @fun()
  %add = add i32 %call, %paramAcrossCall
  ret i32 %add
Exit:
  ret i32 0
}

With this change I observed 10% more shrink-wrapping happened in spec2000/2006/2017.

For now, I perform this only in the entry block as a way to encourage shrink-wrapping, but we might be able to make this more general as a separate pass like Post-RA machine sink. Please let me know any thought.

[thegameg]

Thanks for working on this!

A few things I am wondering about:

• What if we can sink through but in the end it's not useful for shrink-wrapping (something else is blocking) ? We should be able to revert, or anticipate this in the first place.
• I liked the idea that shrink-wrapping is only analyzing and not modifying the code until PEI, is there any other pass where this would fit better?
• Would a post-RA MachineSink pass make sense here ? (in a way that's what shrink-wrapping really is)
• I wonder how much extra work the register allocator would do if we sink (during pre-RA MachineSink) the %vreg = COPY %arg into the successor and add %arg as a live-in.

test/CodeGen/AArch64/sink-copy-for-shrink-wrap.ll

2

I think a MIR test would show the use case of this patch much better here.

[aemerson]

Do you have benchmark performance results for this change?

[junbuml]

What if we can sink through but in the end it's not useful for shrink-wrapping (something else is blocking) ? We should be able to revert, or anticipate this in the first place.

• This is certainly possible, but I'm not sure about the negative cases by sinking such Copys in post-RA.
• By moving Copys into a successor, we can avoid executing instructions in case the result is not used.
• If there is any undesirable impact by sinking with no use of shrink-wrappping, I may want to check if doing so can help shrink wrapping first instead of reverting.

I liked the idea that shrink-wrapping is only analyzing and not modifying the code until PEI, is there any other pass where this would fit better?

• Actaully, I'm also not fully happy with doing this in shrink-wrapping pass. That's why I said this is WIP.
• It should be done between RA and shrink-wrapping, but I cannot see any good existing pass where this can be placed.

Would a post-RA MachineSink pass make sense here ? (in a way that's what shrink-wrapping really is)

• I don't sink this is a bad option if there is any other case which is not handled in pre-RA MachineSink pass.

I wonder how much extra work the register allocator would do if we sink (during pre-RA MachineSink) the %vreg = COPY %arg into the successor and add %arg as a live-in.

• It sounds like another option. However, for me, based on the comment in FindSuccToSinkTo() in MachineSink, it seems require pretty fundamental changes in RA. Quentin may have opinion about this.

[junbuml]

Do you have benchmark performance results for this change?

As I mentioned in summary, I observed 10% more shrink-wrapping happened in spec2000/2006/2017. However, I didn't see significant improvement in scores on AArch64.
One observation I want to share is that if this is performed after the copy source forwaring (like rL312328, which was reverted though at this moment), it will improve spec2006/astar score about 10%.

[junbuml]

Added MIR tests which should show use cases of this patch better.
Again, this patch is preliminary, but I posted this to hear any early high level feedback. Please let me know any thought.

[junbuml]

Submitted https://reviews.llvm.org/D41463 to cover this.

lib/CodeGen/ShrinkWrap.cpp

 // Property #3 is ensured via the MachineBlockFrequencyInfo.
 //
 // If this pass found points matching all these properties, then
 // MachineFrameInfo is updated with this information.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"

 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Target/TargetMachine.h"
 #include <cassert>
 #include <cstdint>
 #include <memory>
 
+#include "llvm/CodeGen/MachineBasicBlock.h"
+#include "llvm/CodeGen/MachineFunction.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineOperand.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/CodeGen/TargetRegisterInfo.h"
+#include "llvm/CodeGen/TargetSubtargetInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
+
 using namespace llvm;
 
 #define DEBUG_TYPE "shrink-wrap"
 
 STATISTIC(NumFunc, "Number of functions");
 STATISTIC(NumCandidates, "Number of shrink-wrapping candidates");
 STATISTIC(NumCandidatesDropped,
           "Number of shrink-wrapping candidates dropped because of frequency");
class ShrinkWrap : public MachineFunctionPass {
   /// Current MachineFunction.
   MachineFunction *MachineFunc;
 
   /// \brief Check if \p MI uses or defines a callee-saved register or
   /// a frame index. If this is the case, this means \p MI must happen
   /// after Save and before Restore.
   bool useOrDefCSROrFI(const MachineInstr &MI, RegScavenger *RS) const;
 
+  /// \brief Sink Copys to a CSR in the entry block into successors if we can
+  /// safely move them.
+  bool sinkCSRFromEntry(MachineFunction &MF, RegScavenger *RS,
+                        const TargetRegisterInfo *TRI);
+
+  /// Return true if \p MI define a CSR.
+  bool defCSR(const MachineInstr &MI, RegScavenger *RS) const;
+
   const SetOfRegs &getCurrentCSRs(RegScavenger *RS) const {
     if (CurrentCSRs.empty()) {
       BitVector SavedRegs;
       const TargetFrameLowering *TFI =
           MachineFunc->getSubtarget().getFrameLowering();
 
       TFI->determineCalleeSaves(*MachineFunc, SavedRegs, RS);
 
if (UseOrDefCSR || MO.isFI()) {
       DEBUG(dbgs() << "Use or define CSR(" << UseOrDefCSR << ") or FI("
                    << MO.isFI() << "): " << MI << '\n');
       return true;
     }
   }
   return false;
 }
 
+bool ShrinkWrap::defCSR(const MachineInstr &MI, RegScavenger *RS) const {
+  for (const MachineOperand &MO : MI.operands()) {
+    if (MO.isReg()) {
+      unsigned PhysReg = MO.getReg();
+      if (!PhysReg)
+        continue;
+      assert(TargetRegisterInfo::isPhysicalRegister(PhysReg) &&
+             "Unallocated register?!");
+      if (MO.isDef() && RCI.getLastCalleeSavedAlias(PhysReg))
+        return true;
+    } else if (MO.isRegMask()) {
+      // Check if this regmask clobbers any of the CSRs.
+      for (unsigned Reg : getCurrentCSRs(RS))
+        if (MO.clobbersPhysReg(Reg))
+          return true;
+    }
+  }
+  return false;
+}
+
 /// \brief Helper function to find the immediate (post) dominator.
 template <typename ListOfBBs, typename DominanceAnalysis>
 static MachineBasicBlock *FindIDom(MachineBasicBlock &Block, ListOfBBs BBs,
                                    DominanceAnalysis &Dom) {
   MachineBasicBlock *IDom = &Block;
   for (MachineBasicBlock *BB : BBs) {
     IDom = Dom.findNearestCommonDominator(IDom, BB);
     if (!IDom)
for (const MachineBasicBlock *SuccBB : MBB->successors()) {
       // Otherwise, we have an irreducible graph.
       if (!isProperBackedge(MLI, MBB, SuccBB))
         return true;
     }
   }
   return false;
 }
 
+static bool isLiveIn(MachineBasicBlock *MBB, SmallSet<unsigned, 8> &RegAliasSet,
+                     const TargetRegisterInfo *TRI) {
+  for (const auto LI : MBB->liveins())
+    if (RegAliasSet.count(LI.PhysReg))
+      return true;
+  return false;
+}
+
+static bool isUsed(MachineInstr &MI, unsigned Reg,
+                   const TargetRegisterInfo *TRI) {
+  for (unsigned i = 0, e = MI.getNumOperands(); i != e; ++i) {
+    const MachineOperand &MO = MI.getOperand(i);
+    if (!MO.isReg() || !MO.isUse())
+      continue;
+    unsigned MOReg = MO.getReg();
+    if (!MOReg)
+      continue;
+    if (TRI->isSuperOrSubRegisterEq(MOReg, Reg)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+static bool canSinkThrough(MachineInstr &CI, const TargetRegisterInfo *TRI) {
+  assert(CI.isCopy() && "Expect only copy");
+  unsigned DefReg = CI.getOperand(0).getReg();
+  unsigned SrcReg = CI.getOperand(1).getReg();
+  MachineBasicBlock::iterator I = CI, E = CI.getParent()->end();
+  ++I;
+  while (I != E) {
+    MachineInstr &MI = *I;
+    if (MI.isCall())
+      return false;
+    if (isUsed(MI, DefReg, TRI) || MI.modifiesRegister(DefReg, TRI) ||
+        MI.modifiesRegister(SrcReg, TRI))
+      return false;
+    ++I;
+  }
+  return true;
+}
+
+/// Sink Copy instructions which define a CSR in the entry to successors.
+bool ShrinkWrap::sinkCSRFromEntry(MachineFunction &MF, RegScavenger *RS,
+                                  const TargetRegisterInfo *TRI) {
+  MachineBasicBlock *Entry = &MF.front();
+  SmallVector<MachineBasicBlock *, 2> SuccBBs(Entry->succ_begin(),
+                                              Entry->succ_end());
+  DenseMap<MachineBasicBlock *, MachineBasicBlock *> NewSuccBBs;
+  bool Changed = false;
+
+  if (SuccBBs.size() != 2 || SuccBBs[0] == SuccBBs[1])
+    return Changed;
+
+  for (auto I = Entry->rbegin(), E = Entry->rend(); I != E;) {
+    MachineInstr *MI = &*I;
+    ++I;
+
+    // FIXME: For now, only handle Copys which define CSRs.
+    if (!MI->isCopy() || !defCSR(*MI, RS))
+      continue;
+
+    unsigned CSRDefReg = MI->getOperand(0).getReg();
+    unsigned SrcReg = MI->getOperand(1).getReg();
+
+    SmallSet<unsigned, 8> RegAliasSet;
+    for (MCRegAliasIterator AI(CSRDefReg, TRI, true); AI.isValid(); ++AI)
+      RegAliasSet.insert(*AI);
+
+    bool isLiveOutBB0 = isLiveIn(SuccBBs[0], RegAliasSet, TRI);
+    bool isLiveOutBB1 = isLiveIn(SuccBBs[1], RegAliasSet, TRI);
+
+    if (!((!isLiveOutBB0 && isLiveOutBB1) || (isLiveOutBB0 && !isLiveOutBB1)))
+      continue;
+
+    // Sink if Def of CSR is not used/defed and SrcReg is defed in
+    if (!canSinkThrough(*MI, TRI))
+      continue;
+
+    // Create a new successor for the entry, and place the Copy which def a CSR
+    // to the new successor.
+    MachineBasicBlock *SuccBB = isLiveOutBB0 ? SuccBBs[0] : SuccBBs[1];
+    if (!NewSuccBBs.count(SuccBB)) {
+      MachineBasicBlock *NMBB = MF.CreateMachineBasicBlock();
+      MF.insert(MachineFunction::iterator(SuccBB), NMBB);
+      Entry->ReplaceUsesOfBlockWith(SuccBB, NMBB);
+      NMBB->addSuccessor(SuccBB);
+
+      for (const auto &LI : SuccBB->liveins())
+        NMBB->addLiveIn(LI);
+
+      for (auto *PI : SuccBB->predecessors())
+        if (PI != Entry)
+          PI->updateTerminator();
+      NewSuccBBs[SuccBB] = NMBB;
+    }
+
+    MachineBasicBlock *NewSuccBB = NewSuccBBs[SuccBB];
+
+    // Move the instr to the new successor.
+    MachineBasicBlock::iterator InsertPos = NewSuccBB->getFirstNonPHI();
+    NewSuccBB->splice(InsertPos, Entry, MI);
+
+    NewSuccBB->removeLiveIn(CSRDefReg);
+    if (!NewSuccBB->isLiveIn(SrcReg))
+      NewSuccBB->addLiveIn(SrcReg);
+
+    Changed = true;
+  }
+  return Changed;
+}
+
 bool ShrinkWrap::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(*MF.getFunction()) || MF.empty() || !isShrinkWrapEnabled(MF))
     return false;
 
   DEBUG(dbgs() << "**** Analysing " << MF.getName() << '\n');
 
   init(MF);
 
   if (isIrreducibleCFG(MF, *MLI)) {
     // If MF is irreducible, a block may be in a loop without
     // MachineLoopInfo reporting it. I.e., we may use the
     // post-dominance property in loops, which lead to incorrect
     // results. Moreover, we may miss that the prologue and
     // epilogue are not in the same loop, leading to unbalanced
     // construction/deconstruction of the stack frame.
     DEBUG(dbgs() << "Irreducible CFGs are not supported yet\n");
     return false;
   }
 
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   std::unique_ptr<RegScavenger> RS(
       TRI->requiresRegisterScavenging(MF) ? new RegScavenger() : nullptr);
 
+  if (sinkCSRFromEntry(MF, RS.get(), TRI)) {
+    MDT->runOnMachineFunction(MF);
+    MPDT->runOnMachineFunction(MF);
+  }
+
   for (MachineBasicBlock &MBB : MF) {
     DEBUG(dbgs() << "Look into: " << MBB.getNumber() << ' ' << MBB.getName()
                  << '\n');
 
     if (MBB.isEHFuncletEntry()) {
       DEBUG(dbgs() << "EH Funclets are not supported yet.\n");
       return false;
     }

test/CodeGen/AArch64/post-rs-copy-sink.mir

+# RUN: llc -mtriple=aarch64-none-linux-gnu -run-pass shrink-wrap -verify-machineinstrs  -o - %s | FileCheck %s
+---
+# Sink w19 to the new predecessor of bb.2.
+# CHECK-LABEL: name: sinkcopy1
+# CHECK-LABEL: bb.0:
+# CHECK-NOT: %w19 = COPY %w0
+# CHECK: Bcc 11, %bb.3
+# CHECK-LABEL: bb.3:
+# CHECK: %w19 = COPY %w0
+
+name: sinkcopy1
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: %w0, %w1
+    %w1 = SUBSWri %w1, 1, 0, implicit-def %nzcv
+    %w19 = COPY %w0
+    Bcc 11, %bb.1, implicit %nzcv
+    B %bb.2
+
+  bb.1:
+    liveins: %w0, %w19
+    %w0 = ADDWrr %w0, %w19
+    RET %x0
+
+  bb.2:
+    %w0 = COPY %wzr
+    RET   %x0
+...
+
+---
+# Sink w19 to the new predecessor of bb.2 and then it should branch from bb.1 to
+# bb.2 as bb.3 is placed between bb.1 and bb.2.
+# CHECK-LABEL: name: sinkcopy2
+# CHECK-LABEL: bb.0:
+# CHECK-NOT: %w19 = COPY %w0
+# CHECK: Bcc 11, %bb.3
+# CHECK-LABEL: bb.1:
+# CHECK: B %bb.2
+# CHECK-LABEL: bb.3:
+# CHECK: %w19 = COPY %w0
+# CHECK-LABEL: bb.2:
+name: sinkcopy2
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: %w0, %w1
+    %w1 = SUBSWri %w1, 1, 0, implicit-def %nzcv
+    %w19 = COPY %w0
+    Bcc 11, %bb.2, implicit %nzcv
+    B %bb.1
+
+  bb.1:
+    %w19 = COPY %wzr
+
+  bb.2:
+    liveins: %w0, %w19
+    %w0 = ADDWrr %w0, %w19
+    RET %x0
+...
+
+---
+# Don't sink w19 as w0 is defined in bb.0
+# CHECK-LABEL: name: sinkcopy3
+# CHECK-LABEL: bb.0:
+# CHECK: %w19 = COPY %w0
+name: sinkcopy3
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: %w0, %w1
+    %w1 = SUBSWri %w1, 1, 0, implicit-def %nzcv
+    %w19 = COPY %w0
+    %w0 = LDRWui %sp, 0 :: (load 4)
+    Bcc 11, %bb.2, implicit %nzcv
+    B %bb.1
+
+  bb.1:
+    %w19 = COPY %wzr
+
+  bb.2:
+    liveins: %w0, %w19
+    %w0 = ADDWrr %w0, %w19
+    RET %x0
+...
+
+---
+# Don't sink w19 as w19 is used in bb.0
+# CHECK-LABEL: name: sinkcopy4
+# CHECK-LABEL: bb.0:
+# CHECK: %w19 = COPY %w0
+name: sinkcopy4
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: %w0, %w1
+    %w1 = SUBSWri %w1, 1, 0, implicit-def %nzcv
+    %w19 = COPY %w0
+    STRWui %w1, %x19, 0 :: (store 4)
+    Bcc 11, %bb.2, implicit %nzcv
+    B %bb.1
+
+  bb.1:
+    %w19 = COPY %wzr
+
+  bb.2:
+    liveins: %w0, %w19
+    %w0 = ADDWrr %w0, %w19
+    RET %x0
+...

test/CodeGen/AArch64/sink-copy-for-shrink-wrap.ll

+; RUN: llc -mtriple=aarch64-linux-gnu -o - %s | FileCheck %s
+
+; CHECK-LABEL: BB#0
+; CHECK-NOT: stp
+; CHECK-NOT: mov w{{[0-9]+}}, w0
+; CHECK-LABEL: BB#1
+; CHECK: stp x19
+; CHECK: mov w{{[0-9]+}}, w0
+
+define i32 @shrinkwrapme(i32 %paramAcrossCall, i32 %paramNotAcrossCall) {
+entry:
+  %cmp5 = icmp sgt i32 %paramNotAcrossCall, 0
+  br i1 %cmp5, label %CallBB, label %Exit
+CallBB:
+  %call = call i32 @fun()
+  %add = add i32 %call, %paramAcrossCall
+  ret i32 %add
+Exit:
+  ret i32 0
+}
+
+declare i32 @fun()