[AArch64] Follow-up from llvm#166926

Ensures the hints are only added once, and ensures that hints
inserted by the register allocator take priority over hints to
reduce movprfx.

Author: sdesmalen-arm

llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp

@@ -1157,6 +1157,9 @@ bool AArch64RegisterInfo::getRegAllocationHints(
   // a movprfx.
   const TargetRegisterClass *RegRC = MRI.getRegClass(VirtReg);
   if (AArch64::ZPRRegClass.hasSubClassEq(RegRC)) {
+    bool ConsiderOnlyHints = TargetRegisterInfo::getRegAllocationHints(
+        VirtReg, Order, Hints, MF, VRM);
+
     for (const MachineOperand &DefOp : MRI.def_operands(VirtReg)) {
       const MachineInstr &Def = *DefOp.getParent();
       if (DefOp.isImplicit() ||
@@ -1168,26 +1171,33 @@ bool AArch64RegisterInfo::getRegAllocationHints(
           TII->get(AArch64::getSVEPseudoMap(Def.getOpcode())).TSFlags;
 
       for (MCPhysReg R : Order) {
-        auto AddHintIfSuitable = [&](MCPhysReg R, const MachineOperand &MO) {
-          // R is a suitable register hint if there exists an operand for the
-          // instruction that is not yet allocated a register or if R matches
-          // one of the other source operands.
-          if (!VRM->hasPhys(MO.getReg()) || VRM->getPhys(MO.getReg()) == R)
+        auto AddHintIfSuitable = [&](MCPhysReg R,
+                                     const MachineOperand &MO) -> bool {
+          // R is a suitable register hint if:
+          // * R is one of the source operands.
+          // * The register allocator has not suggested any hints and one of the
+          //   instruction's source operands does not yet have a register
+          //   allocated for it.
+          if (VRM->getPhys(MO.getReg()) == R ||
+              (!VRM->hasPhys(MO.getReg()) && Hints.empty())) {
             Hints.push_back(R);
+            return true;
+          }
+          return false;
         };
 
         switch (InstFlags & AArch64::DestructiveInstTypeMask) {
         default:
           break;
         case AArch64::DestructiveTernaryCommWithRev:
-          AddHintIfSuitable(R, Def.getOperand(2));
-          AddHintIfSuitable(R, Def.getOperand(3));
-          AddHintIfSuitable(R, Def.getOperand(4));
+          AddHintIfSuitable(R, Def.getOperand(2)) ||
+              AddHintIfSuitable(R, Def.getOperand(3)) ||
+              AddHintIfSuitable(R, Def.getOperand(4));
           break;
         case AArch64::DestructiveBinaryComm:
         case AArch64::DestructiveBinaryCommWithRev:
-          AddHintIfSuitable(R, Def.getOperand(2));
-          AddHintIfSuitable(R, Def.getOperand(3));
+          AddHintIfSuitable(R, Def.getOperand(2)) ||
+              AddHintIfSuitable(R, Def.getOperand(3));
           break;
         case AArch64::DestructiveBinary:
         case AArch64::DestructiveBinaryImm:
@@ -1198,8 +1208,7 @@ bool AArch64RegisterInfo::getRegAllocationHints(
     }
 
     if (Hints.size())
-      return TargetRegisterInfo::getRegAllocationHints(VirtReg, Order, Hints,
-                                                       MF, VRM);
+      return ConsiderOnlyHints;
   }
 
   if (!ST.hasSME() || !ST.isStreaming())

llvm/test/CodeGen/AArch64/regalloc-hint-movprfx.mir

@@ -0,0 +1,66 @@
+# RUN: llc -mtriple=aarch64 -mattr=+sve -start-before=greedy -stop-after=virtregrewriter -debug-only=regalloc %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=DBG
+
+# Check that the register allocator gets hints to reuse registers of one of it's operands.
+---
+name:            prioritize_movprfx_hints
+tracksRegLiveness: true
+isSSA:           false
+noVRegs:         false
+body:             |
+  bb.0.entry:
+    liveins: $z0, $z1, $z2, $z3, $p0
+
+    ; DBG: Machine code for function prioritize_movprfx_hints
+    ;
+    ; DBG: selectOrSplit ZPR:%4
+    ; DBG-NEXT: hints: $z0 $z1{{$}}
+    ;
+    ; DBG: selectOrSplit ZPR:%5
+    ; DBG-NEXT: hints: $z2 $z3{{$}}
+    ;
+    ; DBG: [%0 -> $z3] ZPR
+    ; DBG: [%1 -> $z2] ZPR
+    ; DBG: [%2 -> $z1] ZPR
+    ; DBG: [%3 -> $z0] ZPR
+    ; DBG: [%4 -> $z0] ZPR
+    ; DBG: [%5 -> $z2] ZPR
+    ; DBG: [%6 -> $z0] ZPR
+    %0:zpr = COPY $z3
+    %1:zpr = COPY $z2
+    %2:zpr = COPY $z1
+    %3:zpr = COPY $z0
+    %4:zpr = SDIV_ZPZZ_D_UNDEF $p0, %3:zpr, %2:zpr
+    %5:zpr = MUL_ZPZZ_D_UNDEF $p0, %1:zpr, %0:zpr
+    %6:zpr = MUL_ZPZZ_D_UNDEF $p0, %5:zpr, %4:zpr
+    $z0 = COPY %6:zpr
+    RET_ReallyLR implicit $z0
+...
+
+# Check that the register allocator prioritises hints that are set by the register
+# allocator itself (i.e. to use z4 for the result register).
+---
+name:            prioritize_regalloc_hints
+isSSA:           false
+noVRegs:         false
+body:             |
+  bb.0.entry:
+    %0:zpr = FDUP_ZI_S 0, implicit $vg
+    %1:zpr = FDUP_ZI_S 16, implicit $vg
+    %2:zpr = FDUP_ZI_S 32, implicit $vg
+    %3:ppr_3b = PTRUE_S 31, implicit $vg
+
+    ; DBG: Machine code for function prioritize_regalloc_hints
+    ;
+    ; DBG: selectOrSplit ZPR:%4
+    ; DBG-NEXT: hints: $z4{{$}}
+    ;
+    ; DBG: [%0 -> $z0] ZPR
+    ; DBG: [%1 -> $z1] ZPR
+    ; DBG: [%2 -> $z2] ZPR
+    ; DBG: [%3 -> $p0] PPR_3b
+    ; DBG: [%4 -> $z4] ZPR
+
+    %4:zpr = FMLA_ZPZZZ_S_UNDEF %3, %0, %1, %2
+    $z4 = COPY %4
+    RET_ReallyLR implicit $z4
+...