RegisterCoalescer: Fix joinReservedPhysReg()

[MatzeB]

Summary

• Some code cleanup
• Fix incorrect handling of early clobber definitions
• Allow multiple vreg uses

[kparzysz]

lib/CodeGen/RegisterCoalescer.cpp

1825

I think starting from the next index is somewhat dangerous. Consider this scenario: Hexagon has an instruction "J4_jumpsetr Rd, Rs, #bb", which is equivalent to "Rd = Rs; jump #bb". This one will not cause a problem here, but it inspired me to invent a new hypothetical instruction where the branch is indirect: "J4_jumprsetr Rd, Rs, Rt", defined as "Rd = Rs; jumpr Rt".

I converted the testcase to Hexagon (use R30 instead of FP), with the change that it now uses the invented instruction to define virtual register: "J4_jumprsetr %0, R0, R30". Before coalescing, this would be equivalent to "%0 = R0; jumpr R30", after coalescing it becomes "R30 = R0; jumpr R30". While with the Hexagon semantics the R30 in the target of the branch would still be the old value, I'm not sure that this can be said about any potential architecture (which would support this kind of instructions). Reserved physical registers may not obey the same use/def ordering, and there is no guarantee that the branch to R30 above would not be redirected to the newly defined value (or that is still a legal instruction to begin with).

Here's the modified testcase:

# RUN: llc -mtriple=hexagon -run-pass=simple-register-coalescing %s -o - | FileCheck %s
--- |
  define void @func1() { ret void }
...
---
name: func1
body: |
  bb.0:
    successors: %bb.3, %bb.4
    J2_jumpt undef %p0, %bb.3, implicit def %pc
    J2_jump %bb.4, implicit-def %pc

  bb.1:
    successors: %bb.2, %bb.5
    %r30 = COPY %0
    J2_jumpt undef %p0, %bb.2, implicit-def %pc
    J2_jump %bb.5, implicit-def %pc

  bb.2:
    successors: %bb.6
    %r30 = COPY %r0; outside the lifetime of %0, so shouldn't matter
    J2_jump %bb.6, implicit-def %pc

  bb.3:
    %r0 = COPY %r30 ; outside the lifetime of %0, so shouldn't matter
    J2_jumpr %r31, implicit-def %pc

  bb.5:
    S2_storeri_io %r30, 0, %r30
    J2_jumpr %r31, implicit-def %pc

  bb.6:
    J2_jumpr %r31, implicit-def %pc

  bb.4:
    successors: %bb.1
    %0 : intregs = J4_jumprsetr %r0, %r30, implicit-def %pc
...

You'd need to add a definition of it to the Hexagon target for the testcase to work, but you can change the testcase to use J4_jumpsetr (no "r" in the middle) to still get an instruction that uses R30.

Here's the output:

********** SIMPLE REGISTER COALESCING **********
********** Function: func1
********** JOINING INTERVALS ***********
(null):
64B     %R30<def> = COPY %vreg0; IntRegs:%vreg0
        Considering merging %vreg0 with %R30
                RHS = %vreg0 [48B,64r:0)[304r,320B:0)  0@304r
                Removing phys reg def of %R30 at 64B
                updated: 304B   %R30<def> = J4_jumprsetr %R0, %R30, %PC<imp-def>
        Success: %vreg0 -> %R30
        Result = %R30
(null):
(null):
128B    %R30<def> = COPY %R0
        Not coalescable.
(null):
(null):
176B    %R0<def> = COPY %R30
        Not coalescable.
(null):
(null):
Trying to inflate 0 regs.

and we end up with

288B    BB#6:
            Predecessors according to CFG: BB#0
304B            %R30<def> = J4_jumprsetr %R0, %R30, %PC<imp-def>
            Successors according to CFG: BB#1(0x80000000 / 0x80000000 = 100.00%)

[MatzeB]

lib/CodeGen/RegisterCoalescer.cpp

1825

Sorry, I am getting confused here. For my understand: Your concern is about the defining instruction of the virtual register. You think of "R30 = R0; jumpr R30" as execution in parallel (VLIW style). IN terms of register allocation/semantics this would be: read R0 and R30, the execute the two instructions; then write R30? That would match the semantics we model in our liveranges. If this is not the case we would need to set the definition to the earlyclobber register slot. You are right that this case is not handled correctly in the patch right now. I'll update.

[MatzeB]

Addressed Krzysztofs comments. I will try to extend the tests for those sneaky corner cases tomorrow.

[MatzeB]

I will try to update the tests tomorrow. BTW: For testing regalloc infrastructure I found that you do not need to invent new instructions, just take an existing one (like NOP) and add implicit def and use operands in your .mir test case :)

[MatzeB]

Add tests for early-clobber corner cases.

[kparzysz]

lib/CodeGen/RegisterCoalescer.cpp

1825

Yes, although I missed the earlyclobber solution (I was thinking about the concept of "late use" and the earlyclobber flag eluded me here).

[wmi]

Sorry to chime in late.

test/CodeGen/AArch64/regcoal-physreg.mir

125–132

If bb.4 has another successor to bb.5, after %0 = ADRP 0 is changed to %fp = ADRP 0, it may change the value of %fp in bb.5. So the original check that the def of physical register has no overlap with the interval of SrcReg is not enough in the case of complex control flow?

bb.4:
  successors: %bb.1(0x40000000), %bb.5(0x40000000)

  %fp = COPY %x3
  %0 : gpr64 = ADRP 0
  CBZX undef %x0, %bb.5
  B %bb.1

bb.5:
  STRXui %fp, %fp, 0
  RET_ReallyLR

[MatzeB]

Good catch Wei! Looks like the only sensible way to fix this today is to require the liverange to be whithin a single basic block.

I put this simple fix into a different review: https://reviews.llvm.org/D29436 that we can hopefully pick into the llvm 4.0 branch.
I will update this review as a followup patch to D29436 as it has grown some interesting features like proper support for earlyclobber defs and multiple uses of the vreg and some code cleanups (they are however not critical enough that they must go into the 4.0 branch).

[MatzeB]

Addressed review comments. Note that I put the fix of the original problem into D29611 now.

lib/CodeGen/RegisterCoalescer.cpp

 bool RegisterCoalescer::joinReservedPhysReg(CoalescerPair &CP) {
   unsigned DstReg = CP.getDstReg();
   unsigned SrcReg = CP.getSrcReg();
   assert(CP.isPhys() && "Must be a physreg copy");
   assert(MRI->isReserved(DstReg) && "Not a reserved register");
   LiveInterval &RHS = LIS->getInterval(SrcReg);
   DEBUG(dbgs() << "\t\tRHS = " << RHS << '\n');
 
+  if (!MRI->hasOneDef(SrcReg))
+    return false;
+
   assert(RHS.containsOneValue() && "Invalid join with reserved register");
 
   // Optimization for reserved registers like ESP. We can only merge with a
   // reserved physreg if RHS has a single value that is a copy of DstReg.
   // The live range of the reserved register will look like a set of dead defs
   // - we don't properly track the live range of reserved registers.
 
   // Deny any overlapping intervals.  This depends on all the reserved
bool RegisterCoalescer::joinReservedPhysReg(CoalescerPair &CP) {
   } else {
     // VReg is copied into physreg:
     //   %vregX = def
     //   ... //< no other def or use of %Y here
     //   %Y = COPY %vregX
     // =>
     //   %Y = def
     //   ...
-    if (!MRI->hasOneNonDBGUse(SrcReg)) {
-      DEBUG(dbgs() << "\t\tMultiple vreg uses!\n");
-      return false;
-    }
+    const MachineOperand &DefMO = *MRI->def_begin(SrcReg);
 
-    MachineInstr &DestMI = *MRI->getVRegDef(SrcReg);
     CopyMI = &*MRI->use_instr_nodbg_begin(SrcReg);
-    SlotIndex CopyRegIdx = LIS->getInstructionIndex(*CopyMI).getRegSlot();
-    SlotIndex DestRegIdx = LIS->getInstructionIndex(DestMI).getRegSlot();
+    SlotIndex CopyIdx = LIS->getInstructionIndex(*CopyMI);
+    const MachineInstr &DestMI = *DefMO.getParent();
+    SlotIndex DestRegIdx =
+      LIS->getInstructionIndex(DestMI).getRegSlot(DefMO.isEarlyClobber());
 
     if (!MRI->isConstantPhysReg(DstReg)) {
       // We checked above that there are no interfering defs of the physical
       // register. However, for this case, where we intent to move up the def of
-      // the physical register, we also need to check for interfering uses.
-      SlotIndexes *Indexes = LIS->getSlotIndexes();
-      for (SlotIndex SI = Indexes->getNextNonNullIndex(DestRegIdx);
-           SI != CopyRegIdx; SI = Indexes->getNextNonNullIndex(SI)) {
-        MachineInstr *MI = LIS->getInstructionFromIndex(SI);
-        if (MI->readsRegister(DstReg, TRI)) {
-          DEBUG(dbgs() << "\t\tInterference (read): " << *MI);
+      // the physical register, we also need to check that there are no uses
+      // of DstReg during the lifetime of SrcReg.
+      SlotIndexes &Indexes = *LIS->getSlotIndexes();
+      for (const LiveRange::Segment &S : RHS) {
+        // Adjust Start to point to an instruction. The segment either starts
+        // at block begin or at the defining instruction.
+        // Segments start at the block begin which has no
+        // instruction associated or at the definition instruction of SrcReg.
+        // We can ignore the defining instructions uses unless the def is in the
+        // earlyclobber slot.
+        SlotIndex Start = S.start;
+        if (Start.isBlock() || !DestRegIdx.isEarlyClobber()) {
+          assert(Start.isBlock() ||
+                 (SlotIndex::isSameInstr(Start, DestRegIdx) &&
+                  !DestRegIdx.isEarlyClobber()) && "block begin or unique def");
+          Start = Indexes.getNextNonNullIndex(Start);
+        }
+
+        // Check instructions in range for uses of DstReg.
+        for (SlotIndex SI = Start; SI < S.end;
+             SI = Indexes.getNextNonNullIndex(SI)) {
+          MachineInstr &MI = *LIS->getInstructionFromIndex(SI);
+          if (MI.readsRegister(DstReg, TRI)) {
+            DEBUG(dbgs() << "\t\tInterference (read): " << MI);
             return false;
           }
         }
       }
+    }
 
     // We're going to remove the copy which defines a physical reserved
     // register, so remove its valno, etc.
     DEBUG(dbgs() << "\t\tRemoving phys reg def of " << PrintReg(DstReg, TRI)
-          << " at " << CopyRegIdx << "\n");
+          << " at " << CopyIdx << "\n");
 
-    LIS->removePhysRegDefAt(DstReg, CopyRegIdx);
+    LIS->removePhysRegDefAt(DstReg, CopyIdx.getRegSlot());
     // Create a new dead def at the new def location.
     for (MCRegUnitIterator UI(DstReg, TRI); UI.isValid(); ++UI) {
       LiveRange &LR = LIS->getRegUnit(*UI);
       LR.createDeadDef(DestRegIdx, LIS->getVNInfoAllocator());
     }
   }
 
   LIS->RemoveMachineInstrFromMaps(*CopyMI);

test/CodeGen/AArch64/regcoal-physreg.mir

 # RUN: llc -mtriple=aarch64-apple-ios -run-pass=simple-register-coalescing %s -o - | FileCheck %s
 --- |
   declare void @f2()
 
-  define void @func() { ret void }
+  define void @func0() { ret void }
+  define void @func1() { ret void }
+  define void @func2() { ret void }
 ...
 ---
 # Check coalescing of COPYs from reserved physregs.
-# CHECK-LABEL: name: func
-name: func
+# CHECK-LABEL: name: func0
+name: func0
 registers:
   - { id: 0, class: gpr32 }
   - { id: 1, class: gpr64 }
   - { id: 2, class: gpr64 }
   - { id: 3, class: gpr32 }
   - { id: 4, class: gpr64 }
   - { id: 5, class: gpr32 }
   - { id: 6, class: xseqpairsclass }
bb.0:
     ; CHECK: STRXui %7, %x1, 0
 
     ; Need a def of x18 so that it's not deduced as "constant".
     %x18 = COPY %xzr
     %7 = COPY %x18
     BL @f2, csr_aarch64_aapcs, implicit-def dead %lr, implicit %sp, implicit-def %sp
     STRXui %7, %x1, 0
 ...
+---
+# It should be possible to coalesce "%0 = COPY %fp" even with the contrived CFG
+# CHECK-LABEL: name: func1
+# CHECK-NOT: %fp = COPY %0
+# CHECK: %fp = ADRP
+name: func1
+body: |
+  bb.0:
+    successors: %bb.3, %bb.4
+    CBZX undef %x0, %bb.3
+    B %bb.4
+
+  bb.1:
+    successors: %bb.2, %bb.5
+    %fp = COPY %0
+    CBZX undef %x0, %bb.2
+    B %bb.5
+
+  bb.2:
+    successors: %bb.6
+    %fp = COPY %xzr; outside the lifetime of %0, so shouldn't matter
+    B %bb.6
+
+  bb.3:
+    %xzr = COPY %fp ; outside the lifetime of %0, so shouldn't matter
+    RET_ReallyLR
+
+  bb.4:
+    successors: %bb.1
+    %0 : gpr64 = ADRP 0
+    B %bb.1
+
+  bb.5:
+    STRXui %fp, %fp, 0
+    RET_ReallyLR
+
+  bb.6:
+    RET_ReallyLR
+...
+---
+# Variant of func1 were we should not be able to coalesce because of extra %fp
+# reads during the lifetime of %0
+# CHECK-LABEL: name: func2
+# CHECK: %fp = COPY %0
+# CHECK: %0 = ADRP
+name: func2
+body: |
+  bb.0:
+    successors: %bb.4
+    B %bb.4
+
+  bb.1:
+    successors: %bb.2, %bb.5
+    %fp = COPY %0
+    CBZX undef %x0, %bb.2
+    B %bb.5
+
+  bb.2:
+    successors: %bb.6
+    B %bb.6
+
+  bb.3:
+    successors: %bb.1
+    %xzr = COPY %fp ; %0 is live here so this should block coalescing
+    B %bb.1
+
+  bb.4:
+    successors: %bb.3
+    %0 : gpr64 = ADRP 0
+    B %bb.3
+
+  bb.5:
+    STRXui %fp, %fp, 0
+    RET_ReallyLR
+
+  bb.6:
+    RET_ReallyLR
+...