Named operand and stable partition applied

(cherry picked from commit ee6d876fcc3d84d6ea3a68b3eee1ce97e714b6e6)

Author: mssefat

llvm/include/llvm/CodeGen/MachineRegisterInfo.h

@@ -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/MachineRegisterInfo.cpp

@@ -682,22 +682,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

@@ -269,18 +269,29 @@ bool GCNPreRAOptimizationsImpl::run(MachineFunction &MF) {
         // 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) {
+            unsigned Opc = MI.getOpcode();
+            int OpIdx = AMDGPU::getNamedOperandIdx(Opc, OpName);
+            if (OpIdx == -1) {
+              LLVM_DEBUG(dbgs() << "    Named operand " << OpNameStr
+                                << " not found\n");
               return;
-
+            }
             const MachineOperand &MO = MI.getOperand(OpIdx);
-            if (MO.isReg() && MO.getReg().isVirtual())
+            if (MO.isReg() && MO.getReg().isVirtual()) {
               MFMARegisters.push_back(MO.getReg());
+              LLVM_DEBUG(dbgs()
+                         << "    Collected " << OpNameStr << " (Op" << OpIdx
+                         << "): " << printReg(MO.getReg(), 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));