Named operand and stable partition applied

(cherry picked from commit ee6d876fcc3d84d6ea3a68b3eee1ce97e714b6e6)

Author: mssefat

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

@@ -867,13 +867,13 @@ class MachineRegisterInfo {
     return RegAllocHints.inBounds(VReg) ? &RegAllocHints[VReg] : nullptr;
   }
 
-  /// setRegAllocationAntiHint - Add a register allocation anti-hint for the
+  /// addRegAllocAntiHint - Add a register allocation anti-hint for the
   /// specified virtual register. This tells the allocator to avoid allocating
   /// VReg to the same physical register as AntiHintVReg (or overlapping ones).
-  void setRegAllocationAntiHint(Register VReg, Register AntiHintVReg) {
+  void addRegAllocAntiHint(Register VReg, Register AntiHintVReg) {
     assert(VReg.isVirtual() && "Anti-hints are only for virtual registers");
     assert(AntiHintVReg.isVirtual() && "Anti-hint target must be virtual");
-    AntiHintRegs.grow(Register::index2VirtReg(getNumVirtRegs()));
+    AntiHintRegs.grow(VReg);
     SmallVector<Register, 4> &AntiHints = AntiHintRegs[VReg];
     // Avoid duplicates
     if (llvm::find(AntiHints, AntiHintVReg) == AntiHints.end())
@@ -884,7 +884,7 @@ class MachineRegisterInfo {
   void addRegAllocationAntiHints(Register VReg,
                                  ArrayRef<Register> AntiHintVRegs) {
     for (Register AntiHint : AntiHintVRegs)
-      setRegAllocationAntiHint(VReg, AntiHint);
+      addRegAllocAntiHint(VReg, AntiHint);
   }
 
   /// clearRegAllocationAntiHints - Clear all anti-hints for a register.
@@ -915,7 +915,7 @@ class MachineRegisterInfo {
   /// VRM is the current virtual register map showing allocations made so far.
   void getPhysRegAntiHints(Register VReg,
                            SmallVectorImpl<MCPhysReg> &PhysAntiHints,
-                           const VirtRegMap *VRM) const;
+                           const VirtRegMap &VRM) const;
 
   /// markUsesInDebugValueAsUndef - Mark every DBG_VALUE referencing the
   /// specified register as undefined which causes the DBG_VALUE to be

llvm/lib/CodeGen/AllocationOrder.cpp

@@ -48,7 +48,7 @@ AllocationOrder AllocationOrder::create(Register VirtReg, const VirtRegMap &VRM,
 
   // Get anti-hints
   SmallVector<MCPhysReg, 16> AntiHintedPhysRegs;
-  MRI.getPhysRegAntiHints(VirtReg, AntiHintedPhysRegs, &VRM);
+  MRI.getPhysRegAntiHints(VirtReg, AntiHintedPhysRegs, VRM);
 
   LLVM_DEBUG({
     if (!AntiHintedPhysRegs.empty()) {
@@ -84,29 +84,34 @@ AllocationOrder AllocationOrder::create(Register VirtReg, const VirtRegMap &VRM,
 
 void AllocationOrder::applyAntiHints(ArrayRef<MCPhysReg> AntiHintedPhysRegs,
                                      const TargetRegisterInfo *TRI) {
+  // Helper to check if a register overlaps with any anti-hint
+  auto isAntiHinted = [&](MCPhysReg Reg) {
+    return std::any_of(
+        AntiHintedPhysRegs.begin(), AntiHintedPhysRegs.end(),
+        [&](MCPhysReg AntiHint) { return TRI->regsOverlap(Reg, AntiHint); });
+  };
+
   // Create filtered order
   FilteredOrderStorage.clear();
-  FilteredOrderStorage.reserve(Order.size());
-
-  // Add non-anti-hinted registers first
-  for (MCPhysReg PhysReg : Order) {
-    if (!is_contained(AntiHintedPhysRegs, PhysReg)) {
-      FilteredOrderStorage.push_back(PhysReg);
-    }
-  }
+  FilteredOrderStorage.assign(Order.begin(), Order.end());
 
-  // Add anti-hinted registers at the end as last resort
-  for (MCPhysReg PhysReg : Order) {
-    if (is_contained(AntiHintedPhysRegs, PhysReg)) {
-      FilteredOrderStorage.push_back(PhysReg);
-    }
-  }
+  // Partition: non-anti-hinted registers go first
+  auto PartitionPoint = std::stable_partition(
+      FilteredOrderStorage.begin(), FilteredOrderStorage.end(),
+      [&](MCPhysReg Reg) { return !isAntiHinted(Reg); });
 
   // Update Order
   Order = FilteredOrderStorage;
 
   LLVM_DEBUG({
-    dbgs() << "moved " << AntiHintedPhysRegs.size()
-           << " anti-hinted registers to end of allocation order\n";
+    size_t NonAntiHintedCount =
+        std::distance(FilteredOrderStorage.begin(), PartitionPoint);
+    size_t AntiHintedCount =
+        std::distance(PartitionPoint, FilteredOrderStorage.end());
+    dbgs() << "    Added " << NonAntiHintedCount
+           << " non-anti-hinted registers first\n"
+           << "    Added " << AntiHintedCount
+           << " anti-hinted registers at the end\n"
+           << "  Anti-hint filtering complete\n";
   });
 }

llvm/lib/CodeGen/AllocationOrder.h

@@ -30,6 +30,8 @@ class LiveRegMatrix;
 
 class LLVM_LIBRARY_VISIBILITY AllocationOrder {
   const SmallVector<MCPhysReg, 16> Hints;
+  // Used as storage if the Order received in the constructor needs to be
+  // altered.
   SmallVector<MCPhysReg, 16> FilteredOrderStorage;
   ArrayRef<MCPhysReg> Order;
   // How far into the Order we can iterate. This is 0 if the AllocationOrder is

llvm/lib/CodeGen/MachineRegisterInfo.cpp

@@ -680,22 +680,19 @@ bool MachineRegisterInfo::isReservedRegUnit(unsigned Unit) const {
 
 void MachineRegisterInfo::getPhysRegAntiHints(
     Register VReg, SmallVectorImpl<MCPhysReg> &PhysAntiHints,
-    const VirtRegMap *VRM) const {
+    const VirtRegMap &VRM) const {
   assert(VReg.isVirtual());
-  if (!AntiHintRegs.inBounds(VReg) || !VRM)
+  if (!AntiHintRegs.inBounds(VReg))
     return;
 
   const SmallVector<Register, 4> &AntiHints = AntiHintRegs[VReg];
-  const TargetRegisterInfo *TRI = getTargetRegisterInfo();
 
   for (Register AntiHintVReg : AntiHints) {
     // Check if the anti-hinted register has been allocated
-    if (VRM->hasPhys(AntiHintVReg)) {
-      MCPhysReg PhysReg = VRM->getPhys(AntiHintVReg);
-      // Add the physical register and all its aliases
-      for (MCRegAliasIterator AI(PhysReg, TRI, true); AI.isValid(); ++AI) {
-        PhysAntiHints.push_back(*AI);
-      }
+    if (VRM.hasPhys(AntiHintVReg)) {
+      MCPhysReg PhysReg = VRM.getPhys(AntiHintVReg);
+      // Add the physical register
+      PhysAntiHints.push_back(PhysReg);
     }
   }
 

llvm/lib/Target/AMDGPU/GCNPreRAOptimizations.cpp

@@ -37,6 +37,7 @@
 #include "SIRegisterInfo.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/Register.h"
 #include "llvm/InitializePasses.h"
 
 using namespace llvm;
@@ -253,37 +254,45 @@ bool GCNPreRAOptimizationsImpl::run(MachineFunction &MF) {
   TRI = ST.getRegisterInfo();
 
   bool Changed = false;
-
-  // Single pass implementation
+  // Add RA anti-hints to reduce MFMA hazard NOPs
   if (EnableAntiHintsForMFMARegs && ST.hasMAIInsts()) {
     // Max lookback window for RAW or WAW hazard
     constexpr unsigned MaxLookbackWindow = 19;
     for (const MachineBasicBlock &MBB : MF) {
-
-      SmallVector<std::pair<SlotIndex, SmallVector<Register, 4>>, 16>
-          RecentMFMAs;
+      SmallVector<SmallVector<Register, 4>, 16> RecentMFMAs;
       for (const MachineInstr &MI : MBB) {
         if (MI.isDebugInstr())
           continue;
-        const SlotIndex CurrentSlot = LIS->getInstructionIndex(MI).getRegSlot();
+
         // Handle MFMA instructions
         if (SIInstrInfo::isMFMA(MI)) {
           SmallVector<Register, 4> MFMARegisters;
-          auto collectMFMARegister = [&](unsigned OpIdx) {
-            if (OpIdx >= MI.getNumOperands())
+          // Helper to get named operand
+          auto collectNamedOperand = [&](AMDGPU::OpName OpName,
+                                         const char *OpNameStr) {
+            const MachineOperand *MO = TII->getNamedOperand(MI, OpName);
+            if (!MO) {
+              LLVM_DEBUG(dbgs() << "    Named operand " << OpNameStr
+                                << " not found\n");
               return;
-
-            const MachineOperand &MO = MI.getOperand(OpIdx);
-            if (MO.isReg() && MO.getReg().isVirtual())
-              MFMARegisters.push_back(MO.getReg());
+            }
+            if (MO->isReg() && MO->getReg().isVirtual()) {
+              Register Reg = MO->getReg();
+              const TargetRegisterClass *RC = MRI->getRegClass(Reg);
+              // Only consider VGPRs
+              if (TRI->hasVGPRs(RC))
+                MFMARegisters.push_back(Reg);
+              LLVM_DEBUG(dbgs() << "    Collected " << OpNameStr << " : "
+                                << printReg(Reg, TRI) << "\n");
+            }
           };
-          // Only collect Matrix C (operand 3) and destination (operand 0)
-          // registers
-          collectMFMARegister(0);
-          collectMFMARegister(3);
 
+          // Collect destination and source C registers
+          collectNamedOperand(AMDGPU::OpName::vdst, "vdst"); // Destination
+          collectNamedOperand(AMDGPU::OpName::src2,
+                              "src2"); // Matrix C (accumulator)
           if (!MFMARegisters.empty()) {
-            RecentMFMAs.emplace_back(CurrentSlot, std::move(MFMARegisters));
+            RecentMFMAs.emplace_back(std::move(MFMARegisters));
             // Maintain window
             if (RecentMFMAs.size() > MaxLookbackWindow)
               RecentMFMAs.erase(RecentMFMAs.begin());
@@ -309,17 +318,13 @@ bool GCNPreRAOptimizationsImpl::run(MachineFunction &MF) {
           // Only process VGPR registers
           if (!TRI->isVGPRClass(CandidateRC))
             continue;
-
           for (auto It = RecentMFMAs.rbegin(); It != RecentMFMAs.rend(); ++It) {
-            const SmallVector<Register, 4> &MFMARegs = It->second;
+            const SmallVector<Register, 4> &MFMARegs = *It;
             for (Register MFMAReg : MFMARegs) {
-              // Verify register class compatibility
-              const TargetRegisterClass *MFMARC = MRI->getRegClass(MFMAReg);
-              if (!TRI->hasVGPRs(MFMARC))
-                continue;
-
               // Check if MFMA register is dead at current instruction
               const LiveInterval &MFMAInterval = LIS->getInterval(MFMAReg);
+              const SlotIndex CurrentSlot =
+                  LIS->getInstructionIndex(MI).getRegSlot();
               if (!MFMAInterval.liveAt(CurrentSlot)) {
                 // Add bi-directional anti-hints
                 MRI->addRegAllocationAntiHints(CandidateReg, MFMAReg);