[llvm] [AMDGPU] Overwrite shouldCoalesce to stop merging imm mov with virtual reg with large valnos (PR #209608)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Guo Chen (broxigarchen)

<details>
<summary>Changes</summary>

Isel pattern putting imm inside reg_sequence create side effects when register coalescer join these `copy` from imm
```
%1 = v_mov_b16_t16_e64 ...
%2.hi16 = copy %1
%2.lo16 = ....
...
%3.hi16 = copy %1
%3.lo16 = ....
....
%4.hi16 = copy %1
%4.lo16 = ....
```
to
```
%1 = v_mov_b16_t16_e64 ...
%2.hi16 = copy %1
%2.lo16 = ....
...
%2.lo16 = ... (reuse %2 and repeat)
....
```
When the number of copy increase this inserts a large number of WAR hazzards on the reused virutal reg, and the machine scheduler bail out getting higher reg pressure after sorting.

Overwrite the `shouldCoalsce` to add additional checks on merging copy with mov imms, and stop merging these copies when the size of valnos is too large.

---
Full diff: https://github.com/llvm/llvm-project/pull/209608.diff


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp (+62) 
- (modified) llvm/lib/Target/AMDGPU/SIRegisterInfo.h (+5) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
index 358aa054b3a3f..cda81d0143a07 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
@@ -4339,3 +4339,65 @@ SIRegisterInfo::getVRegFlagsOfReg(Register Reg,
     RegFlags.push_back("WWM_REG");
   return RegFlags;
 }
+
+// Check if a VReg Def is a Mov Imm using LIS
+static bool IsVRegDefImm(Register VReg, LaneBitmask LaneMask,
+                         const LiveIntervals &LIS, SlotIndex UseIdx) {
+  if (!LIS.hasInterval(VReg))
+    return false;
+
+  const LiveInterval &LI = LIS.getInterval(VReg);
+  if (!LI.hasSubRanges()) {
+    const VNInfo *VNI = LI.getVNInfoBefore(UseIdx);
+    if (!VNI || VNI->isUnused() || VNI->isPHIDef())
+      return false;
+
+    MachineInstr *DefMI = LIS.getInstructionFromIndex(VNI->def);
+    if (DefMI && DefMI->isMoveImmediate())
+      return true;
+
+    return false;
+  }
+
+  for (const LiveInterval::SubRange &SR : LI.subranges()) {
+    const VNInfo *VNI = SR.getVNInfoBefore(UseIdx);
+    if (!VNI || VNI->isUnused() || VNI->isPHIDef())
+      return false;
+
+    MachineInstr *DefMI = LIS.getInstructionFromIndex(VNI->def);
+    if (DefMI && DefMI->isMoveImmediate())
+      return true;
+  }
+  return false;
+}
+
+bool SIRegisterInfo::shouldCoalesce(
+    MachineInstr *MI, const TargetRegisterClass *SrcRC, unsigned SubReg,
+    const TargetRegisterClass *DstRC, unsigned DstSubReg,
+    const TargetRegisterClass *NewRC, LiveIntervals &LIS) const {
+  assert(MI->isCopy() && "Only expecting COPY instructions");
+
+  // Do not coalesce if src is a constant while size of LI subrange sum is
+  // more than 100. Let RA decide if need to rematerialize
+  MachineFunction *MF = MI->getParent()->getParent();
+  MachineRegisterInfo &MRI = MF->getRegInfo();
+  auto *TRI = static_cast<const SIRegisterInfo *>(MRI.getTargetRegisterInfo());
+  Register DstReg = MI->getOperand(0).getReg();
+  Register SrcReg = MI->getOperand(1).getReg();
+
+  if (SrcReg.isVirtual() && DstReg.isVirtual()) {
+    LaneBitmask LaneMask =
+        TRI->getSubRegIndexLaneMask(MI->getOperand(1).getSubReg());
+    SlotIndex MIIdx = LIS.getInstructionIndex(*MI);
+    if (IsVRegDefImm(SrcReg, LaneMask, LIS, MIIdx)) {
+      LiveInterval &LIDst = LIS.getInterval(SrcReg);
+      LiveInterval &LISrc = LIS.getInterval(DstReg);
+      if (LISrc.valnos.size() + LIDst.valnos.size() > 100) {
+        return false;
+      }
+    }
+  }
+
+  return TargetRegisterInfo::shouldCoalesce(MI, SrcRC, SubReg, DstRC, DstSubReg,
+                                            NewRC, LIS);
+}
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
index e464c9334ffea..01a94f56382d7 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
@@ -506,6 +506,11 @@ class SIRegisterInfo final : public AMDGPUGenRegisterInfo {
                 ? 2.0
                 : 1.0);
   }
+
+  bool shouldCoalesce(MachineInstr *MI, const TargetRegisterClass *SrcRC,
+                      unsigned SubReg, const TargetRegisterClass *DstRC,
+                      unsigned DstSubReg, const TargetRegisterClass *NewRC,
+                      LiveIntervals &LIS) const override;
 };
 
 } // End namespace llvm

``````````

</details>


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