[llvm] [AMDGPU] Avoid wmma bank conflict for GFX11 and GFX12 (PR #205530)

[Matt Arsenault via llvm-commits]

================
@@ -262,6 +266,41 @@ void GCNPreRAOptimizationsImpl::hintTrue16Copy(const MachineInstr &MI) {
     MRI->setRegAllocationHint(Dst, AMDGPURI::Size16, Src);
 }
 
+void GCNPreRAOptimizationsImpl::recordWMMABankSiblings(
+    const MachineInstr &MI, SIMachineFunctionInfo &MFI) {
+  // The consecutive-allocation VGPR bank collision only happens when each
+  // matrix operand is a multiple of 4 VGPRs (128 bits) wide; with narrower
+  // operands the start registers already fall on different banks. Restrict the
+  // hint to WMMAs whose matrix source operands are a multiple of 128 bits.
+  auto Is128BitMultiple = [&](AMDGPU::OpName Name) {
+    const MachineOperand *MO = TII->getNamedOperand(MI, Name);
+    return MO && MO->isReg() && MO->getReg() &&
+           TRI->getRegSizeInBits(MO->getReg(), *MRI) % 128 == 0;
+  };
+  if (!Is128BitMultiple(AMDGPU::OpName::src0) ||
+      !Is128BitMultiple(AMDGPU::OpName::src1))
+    return;
+
+  // Record, for each virtual register operand, the other operands of this WMMA
+  // so SIRegisterInfo::getRegAllocationHints can steer them onto distinct
+  // VGPR banks. A register used by several WMMAs accumulates the union of their
+  // siblings.
+  SmallVector<Register, 4> Regs;
+  for (const MachineOperand &MO : MI.operands())
+    if (MO.isReg() && MO.getReg() && MO.getReg().isVirtual() &&
+        !is_contained(Regs, MO.getReg()))
+      Regs.push_back(MO.getReg());
+
+  DenseMap<Register, SmallVector<Register, 3>> &Siblings =
+      MFI.getWMMABankSiblings();
+  for (Register Reg : Regs) {
+    SmallVector<Register, 3> &S = Siblings[Reg];
+    for (Register Other : Regs)
+      if (Other != Reg && !is_contained(S, Other))
+        S.push_back(Other);
+  }
----------------
arsenm wrote:

This feels clumsy 

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