CodeGen: Fix ARM cmpxchg64 register fragmentation in fast-regalloc

[strager]

Summary

ARM's CMP_SWAP_64 pseudo-instruction (introduced in r266679,
used with -O0) uses three GPRPair-class registers and two
GPR-class registers. The fast register allocator (also used
with -O0) allocates GPR-class registers before allocating
GPRPair-class registers.

GPRPair includes is r0+r1, r2+r3, r4+r5, r6+r7, r8+r9,
r10+r11, and r12+sp. With Clang's -ffixed-r9 option, with
the frame pointer enabled, and with sp reserved, the
register allocator can only use r0+r1, r2+r3, r4+r5, and
r6+r7 from GPRPair.

If the fast register allocator allocates CMP_SWAP_64's GPR
operands first, it may decide to allocate r1 and r3. Later,
when the allocator allocates CMP_SWAP_64's GPRPair oprands,
it realizes only two GPRPair-class registers are available
(r4+r5 and r6+r7), and it can't spill the already-allocated
registers to make room. LLVM then fails with the message:

error: ran out of registers during register allocation

In short, the fast register allocator fragments the
registers and LLVM can't compile 64-bit compare-exchanges.

As a workaround, reduce the risk of fragmentation by
allocating registers for bigger register classes before
smaller ones. For CMP_SWAP_64 on ARM, this means
registers are allocated for GPRPair operands before GPR
operands.

For consistency, change all architectures, not just ARM.

This fixes PR30228.

Test Plan:
See the included test case (test/CodeGen/ARM/cmpxchg-O0.ll
test_cmpxchg_64_register_pressure). It is based on the
following C program compiled with Clang with -O0:

void f(unsigned long long *addr, unsigned long long desired, unsigned long long new) {
  while (!__sync_bool_compare_and_swap(addr, desired, new)) {
  }
}

[strager]

Let me know if there's a better solution than sorting by size! (Aside from fixing the bug which necessitates CMP_SWAP_64 in the first place...)

lib/CodeGen/RegAllocFast.cpp

997–1000

(I know I should delete this before merging. Locally I need to work with LLVM 4.0.0.)

[strager]

Ping! This is ready for design and code review.

[hans]

+Matthias for register allocation

[MatzeB]

+ Quentin for register allocation

I'd prefer not to add additional logic into the main path of the fast register allocator. How about just reordering the operands of CMP_SWAP_64 (or for all of the CMP_SWAP operations for consistency) instead?

[qcolombet]

Hi,

The approach sounds sensible to me however, I'd like we use the AllocationPriority instead of the size of the register class.
Moreover, please refactor the code (I am thinking helper function) so that the definitions are processed in the same way.
Finally, keep the order of the operand as tie breaker.

Cheers,
-Quentin
PS: If AllocationPriority does not fix the problem, then you'll have to fix the description of the related register class, because size alone is too magic.

[qcolombet]

lib/CodeGen/RegAllocFast.cpp

984

Reject Reg == 0

991

Why can't you just get the the reg class here?

[MatzeB]

In D33814#834086, @qcolombet wrote:

Hi,

The approach sounds sensible to me however, I'd like we use the AllocationPriority instead of the size of the register class.
Moreover, please refactor the code (I am thinking helper function) so that the definitions are processed in the same way.
Finally, keep the order of the operand as tie breaker.

Cheers,
-Quentin
PS: If AllocationPriority does not fix the problem, then you'll have to fix the description of the related register class, because size alone is too magic.

The ARM target does not use AllocationPriority yet. Though it would be great to start using it for the GPRPair and tuple classes (in a separate commit though).

[strager]

I don't plan to work on this anymore.

lib/CodeGen/RegAllocFast.cpp

if (MI->isCopy()) {
       CopySrcSub = MI->getOperand(1).getSubReg();
     }
 
     // Track registers used by instruction.
     UsedInInstr.clear();
 
     // First scan.
     // Mark physreg uses and early clobbers as used.
-    // Find the end of the virtreg operands
-    unsigned VirtOpEnd = 0;
     bool hasTiedOps = false;
     bool hasEarlyClobbers = false;
     bool hasPartialRedefs = false;
     bool hasPhysDefs = false;
     for (unsigned i = 0, e = MI->getNumOperands(); i != e; ++i) {
       MachineOperand &MO = MI->getOperand(i);
       // Make sure MRI knows about registers clobbered by regmasks.
       if (MO.isRegMask()) {
         MRI->addPhysRegsUsedFromRegMask(MO.getRegMask());
         continue;
       }
       if (!MO.isReg()) continue;
       unsigned Reg = MO.getReg();
       if (!Reg) continue;
       if (TargetRegisterInfo::isVirtualRegister(Reg)) {
-        VirtOpEnd = i+1;
         if (MO.isUse()) {
           hasTiedOps = hasTiedOps ||
                               MCID.getOperandConstraint(i, MCOI::TIED_TO) != -1;
         } else {
           if (MO.isEarlyClobber())
             hasEarlyClobbers = true;
           if (MO.getSubReg() && MI->readsVirtualRegister(Reg))
             hasPartialRedefs = true;
if (MI->isInlineAsm() || hasEarlyClobbers || hasPartialRedefs ||
       CopyDst = 0;
       // Pretend we have early clobbers so the use operands get marked below.
       // This is not necessary for the common case of a single tied use.
       hasEarlyClobbers = true;
     }
 
     // Second scan.
     // Allocate virtreg uses.
-    for (unsigned i = 0; i != VirtOpEnd; ++i) {
-      MachineOperand &MO = MI->getOperand(i);
-      if (!MO.isReg()) continue;
+    // HACK(strager): Allocate larger registers first.
+    SmallVector<unsigned, 8> VirtOps;
+    for (unsigned i = 0, e = MI->getNumOperands(); i != e; ++i) {
+      const MachineOperand &MO = MI->getOperand(i);
+      if (!MO.isReg())
+        continue;
+      unsigned Reg = MO.getReg();
+      if (!TargetRegisterInfo::isVirtualRegister(Reg))
+        continue;
+      if (!MO.isUse())
+        continue;
+      VirtOps.emplace_back(i);
+    }
+    std::sort(VirtOps.begin(), VirtOps.end(),
+              [this, &MI](unsigned i, unsigned j) {
+                auto OpSize = [this](const MachineOperand &MO) -> unsigned {
+                  unsigned Reg = MO.getReg();
+                  if (auto RC = MRI->getRegClassOrRegBank(Reg)) {
+                    if (RC.is<const TargetRegisterClass *>()) {
+#if LLVM_VERSION_MAJOR >= 5
+                      // LLVM 5.0.0.
+                      return TRI->getRegSizeInBits(
+                          *RC.get<const TargetRegisterClass *>());
+#else
+                      // LLVM 4.0.0.
+                      return RC.get<const TargetRegisterClass *>()->getSize();
+#endif
+                    } else {
+                      return 0;
+                    }
+                  } else {
+                    return 0;
+                  }
+                };
+                return OpSize(MI->getOperand(i)) >= OpSize(MI->getOperand(j));
+              });
+    for (unsigned i : VirtOps) {
+      const MachineOperand &MO = MI->getOperand(i);
       unsigned Reg = MO.getReg();
-      if (!TargetRegisterInfo::isVirtualRegister(Reg)) continue;
-      if (MO.isUse()) {
-        LiveRegMap::iterator LRI = reloadVirtReg(*MI, i, Reg, CopyDst);
-        unsigned PhysReg = LRI->PhysReg;
-        CopySrc = (CopySrc == Reg || CopySrc == PhysReg) ? PhysReg : 0;
-        if (setPhysReg(MI, i, PhysReg))
-          killVirtReg(LRI);
-      }
-    }
 
     // Track registers defined by instruction - early clobbers and tied uses at
     // this point.
     UsedInInstr.clear();
     if (hasEarlyClobbers) {
       for (unsigned i = 0, e = MI->getNumOperands(); i != e; ++i) {
         MachineOperand &MO = MI->getOperand(i);
         if (!MO.isReg()) continue;

test/CodeGen/ARM/cmpxchg-O0.ll

 ; CHECK:     cmp{{(\.w)?}} [[STATUS]], #0
 ; CHECK:     bne [[RETRY]]
 ; CHECK: [[DONE]]:
 ; CHECK:     dmb ish
   %res = cmpxchg i64* %addr, i64 %desired, i64 %new seq_cst monotonic
   ret { i64, i1 } %res
 }
 
+; If r9 and fp are reserved, cmpxchg can only use r0/r1, r2/r3, r4/r5, or r6/r7
+; for the 64-bit inputs to ldrexd and strexd. Ensure fast-regalloc can find
+; enough registers without spilling.
+define i64 @test_cmpxchg_64_register_pressure(i64* %addr, i64 %desired, i64 %new) nounwind "no-frame-pointer-elim"="true" "target-features"="+reserve-r9" {
+; CHECK-LABEL: test_cmpxchg_64_register_pressure:
+  %addr.addr = alloca i64*, align 4
+  %desired.addr = alloca i64, align 8
+  %new.addr = alloca i64, align 8
+  store i64* %addr, i64** %addr.addr, align 4
+  store i64 %desired, i64* %desired.addr, align 8
+  store i64 %new, i64* %new.addr, align 8
+  br label %while.cond
+
+while.cond:
+  %addr.tmp = load i64*, i64** %addr.addr, align 4
+  %desired.tmp = load i64, i64* %desired.addr, align 8
+  %new.tmp = load i64, i64* %new.addr, align 8
+
+; CHECK-DAG: mov [[NEWLO:r[0-9]+]], r3
+; CHECK-NEXT:mov [[NEWHI:r[0-9]+]], {{r[0-9]+}}
+; CHECK-DAG: ldrd [[DESIREDLO:r[0-9]+]], [[DESIREDHI:r[0-9]+]], [sp, #{{[0-9]+}}] @ 8-byte Reload
+; CHECK-DAG: dmb ish
+; CHECK: [[INNERRETRY:.LBB[0-9]+_[0-9]+]]:
+; CHECK-NOT: {{ldr[^e]|str}}
+; CHECK:     ldrexd [[OLDLO:r[0-9]+]], [[OLDHI:r[0-9]+]], {{\[}}[[ADDR:r[0-9]+]]{{\]}}
+; CHECK-NOT: {{ldr|str}}
+; CHECK:     cmp [[OLDLO]], [[DESIREDLO]]
+; CHECK-NOT: {{ldr|str}}
+; CHECK:     cmpeq [[OLDHI]], [[DESIREDHI]]
+; CHECK-NOT: {{ldr|str}}
+; CHECK:     bne [[INNERDONE:.LBB[0-9]+_[0-9]+]]
+; CHECK-NOT: {{ldr|str[^e]}}
+; CHECK:     strexd [[STATUS:r[0-9]+]], [[NEWLO]], [[NEWHI]], {{\[}}[[ADDR]]{{\]}}
+; CHECK-NOT: {{ldr|str}}
+; CHECK:     cmp{{(\.w)?}} [[STATUS]], #0
+; CHECK-NOT: {{ldr|str}}
+; CHECK:     bne [[INNERRETRY]]
+; CHECK: [[INNERDONE]]:
+; CHECK:     dmb ish
+  %tmp = cmpxchg i64* %addr.tmp, i64 %desired.tmp, i64 %new.tmp seq_cst seq_cst
+
+  %status = extractvalue { i64, i1 } %tmp, 1
+  br i1 %status, label %done, label %while.cond
+
+done:
+  ret i64 0
+}
+
 define { i64, i1 } @test_nontrivial_args(i64* %addr, i64 %desired, i64 %new) {
 ; CHECK-LABEL: test_nontrivial_args:
 ; CHECK:     dmb ish
 ; CHECK-NOT: uxt
 ; CHECK: [[RETRY:.LBB[0-9]+_[0-9]+]]:
 ; CHECK:     ldrexd [[OLDLO:r[0-9]+]], [[OLDHI:r[0-9]+]], [r0]
 ; CHECK:     cmp [[OLDLO]], {{r[0-9]+}}
 ; CHECK:     cmpeq [[OLDHI]], {{r[0-9]+}}