[llvm] [RegAlloc] Consider rematerialization over CSR use (PR #206756)

[Lukas Sommer via llvm-commits]

================
@@ -2372,6 +2372,87 @@ BlockFrequency RAGreedy::calcSpillCost(const LiveInterval &LI) {
   return BlockFrequency(SpillCost);
 }
 
+bool RAGreedy::shouldAvoidCSRForRemat(const LiveInterval &VirtReg,
+                                      AllocationOrder &Order) const {
+  if (!VirtReg.isSpillable())
+    return false;
+
+  // This logic is intentionally narrow: handle a single concrete value
+  // whose def can be cheaply rematerialized at every use.
+  const VNInfo *OnlyVNI = nullptr;
+  for (const VNInfo *VNI : VirtReg.vnis()) {
+    if (!VNI || VNI->isUnused())
+      continue;
+    if (VNI->isPHIDef())
+      return false;
+    if (OnlyVNI)
+      return false;
+    OnlyVNI = VNI;
+  }
+  if (!OnlyVNI)
+    return false;
+
+  MachineInstr *DefMI = LIS->getInstructionFromIndex(OnlyVNI->def);
+  if (!DefMI || DefMI->isImplicitDef() || !TII->isReMaterializable(*DefMI) ||
+      !TII->isAsCheapAsAMove(*DefMI))
+    return false;
+
+  // This logic aims to address a specific problem: The first use of CSR for
+  // cheap-to-rematerialize live ranges because they cross calls, not because
+  // pressure requires it. With this check and the check for an available for a
+  // register that is not a first-use CSR below, we can scope this.
+  if (!Matrix->checkRegMaskInterference(VirtReg))
+    return false;
+
+  SmallPtrSet<MachineInstr *, 8> VisitedUses;
+  unsigned NumUses = 0;
+  for (MachineOperand &MO : MRI->use_nodbg_operands(VirtReg.reg())) {
+    if (MO.isUndef())
+      continue;
+
+    MachineInstr *UseMI = MO.getParent();
+    if (!VisitedUses.insert(UseMI).second)
+      continue;
+
+    SlotIndex UseIdx = LIS->getInstructionIndex(*UseMI).getRegSlot(true);
+    if (!VirtRegAuxInfo::allUsesAvailableAt(DefMI, UseIdx, *LIS, *MRI, *TII))
+      return false;
+
+    // Check if any register that is not a first-use CSR is available as
+    // destination for the rematerialization.
+    SlotIndex PrevIdx = UseIdx.getPrevSlot();
+    bool HasNoFirstCSRReg = false;
+    for (MCRegister RematPhysReg : Order) {
+      if (EvictAdvisor->isUnusedCalleeSavedReg(RematPhysReg))
+        continue;
+      if (!Matrix->checkInterference(PrevIdx, UseIdx, RematPhysReg)) {
+        HasNoFirstCSRReg = true;
+        break;
+      }
+    }
+    if (!HasNoFirstCSRReg)
+      return false;
+
+    ++NumUses;
+  }
+
+  if (!NumUses)
+    return false;
+
+  // Ideally, this would be handled with the CSR cost model, but the scales
+  // differ. The rationale for allowing a fan-out of 3 is simple: Saving and
+  // restoring the CSR will require at least two moves, so we can allow three
+  // uses, as the first is free and the two materializations will at most have
+  // the same cost as CSR save/restore.
----------------
sommerlukas wrote:

I agree that the current logic doesn't take the block frequency into account. 

The issue that led me to use this approximation is that AMD GPU currently still uses the legacy scale for CSR cost. This means that `initializeCSRCost` scales the actual block frequency with a fixed value `2^14` [here](https://github.com/llvm/llvm-project/blob/d0cc9123eb14cae514e6f4e77cf5d1e3c32b6ebe/llvm/lib/CodeGen/RegAllocGreedy.cpp#L2421-L2449).

Even if there were only a single block (the entry block) and a single use, comparing the block frequency of the use with the CSR cost would lead to rematerialization being considered non-profitable, because the scales are not comparable. We would effectively compare:
```
Rematerialization cost < CSR cost
= EntryFreq < ((100 * EntryFreq) / 2^14)
= 1 < (100 / 16384)
``` 

PR https://github.com/llvm/llvm-project/pull/202007 refactors this, also for AMD GPU. After that refactor, we could compare rematerialization cost (as sum of the block frequency over all uses) against the CSR cost.

If we want to take block frequency into account, I would pause this PR until #202007 or a similar fix for AMDGPU has landed.

https://github.com/llvm/llvm-project/pull/206756