[AArch64] Also implement PNR -> PNR copies. (#70682)

Previously we only implemented PNR -> PPR and PPR -> PNR copies.
llvm · Oct 31, 2023 · 73498d2 · 73498d2
1 parent f8742b8
commit 73498d2
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 29 deletions.
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -4426,36 +4426,27 @@ void AArch64InstrInfo::copyPhysReg(MachineBasicBlock &MBB,
     return;
   }
 
-  if (AArch64::PNRRegClass.contains(DestReg) &&
-      AArch64::PPRRegClass.contains(SrcReg)) {
+  // Copy a predicate-as-counter register by ORRing with itself as if it
+  // were a regular predicate (mask) register.
+  bool DestIsPNR = AArch64::PNRRegClass.contains(DestReg);
+  bool SrcIsPNR = AArch64::PNRRegClass.contains(SrcReg);
+  if (DestIsPNR || SrcIsPNR) {
     assert((Subtarget.hasSVE2p1() || Subtarget.hasSME2()) &&
            "Unexpected predicate-as-counter register.");
-    // Copy from pX to pnX is a no-op
-    if ((DestReg.id() - AArch64::PN0) == (SrcReg.id() - AArch64::P0))
-      return;
-    MCRegister PPRDestReg = (DestReg - AArch64::PN0) + AArch64::P0;
-    BuildMI(MBB, I, DL, get(AArch64::ORR_PPzPP), PPRDestReg)
-        .addReg(SrcReg)
-        .addReg(SrcReg)
-        .addReg(SrcReg, getKillRegState(KillSrc))
-        .addDef(DestReg, RegState::Implicit);
-    return;
-  }
-
-  if (AArch64::PPRRegClass.contains(DestReg) &&
-      AArch64::PNRRegClass.contains(SrcReg)) {
-    assert((Subtarget.hasSVE2p1() || Subtarget.hasSME2()) &&
-           "Unexpected predicate-as-counter register.");
-    // Copy from pnX to pX is a no-op
-    if ((DestReg.id() - AArch64::P0) == (SrcReg.id() - AArch64::PN0))
-      return;
-    MCRegister PNRDestReg = (DestReg - AArch64::P0) + AArch64::PN0;
-    MCRegister PPRSrcReg = (SrcReg - AArch64::PN0) + AArch64::P0;
-    BuildMI(MBB, I, DL, get(AArch64::ORR_PPzPP), DestReg)
-        .addReg(PPRSrcReg)
-        .addReg(PPRSrcReg)
-        .addReg(PPRSrcReg, getKillRegState(KillSrc))
-        .addDef(PNRDestReg, RegState::Implicit);
+    auto ToPPR = [](MCRegister R) -> MCRegister {
+      return (R - AArch64::PN0) + AArch64::P0;
+    };
+    MCRegister PPRSrcReg = SrcIsPNR ? ToPPR(SrcReg) : SrcReg;
+    MCRegister PPRDestReg = DestIsPNR ? ToPPR(DestReg) : DestReg;
+
+    if (PPRSrcReg != PPRDestReg) {
+      auto NewMI = BuildMI(MBB, I, DL, get(AArch64::ORR_PPzPP), PPRDestReg)
+                       .addReg(PPRSrcReg) // Pg
+                       .addReg(PPRSrcReg)
+                       .addReg(PPRSrcReg, getKillRegState(KillSrc));
+      if (DestIsPNR)
+        NewMI.addDef(DestReg, RegState::Implicit);
+    }
     return;
   }
 

diff --git a/llvm/test/CodeGen/AArch64/PNRtoPPRCopy.mir → .../test/CodeGen/AArch64/sve2p1_copy_pnr.mir b/llvm/test/CodeGen/AArch64/PNRtoPPRCopy.mir → .../test/CodeGen/AArch64/sve2p1_copy_pnr.mir
@@ -14,7 +14,7 @@ body:             |
   bb.0:
     ; CHECK-LABEL: name: pnr_to_ppr
     ; CHECK: renamable $pn8 = PTRUE_C_D
-    ; CHECK-NEXT: $p0 = ORR_PPzPP $p8, $p8, killed $p8, implicit-def $pn0
+    ; CHECK-NEXT: $p0 = ORR_PPzPP $p8, $p8, killed $p8
     ; CHECK-NEXT: RET_ReallyLR implicit killed $p0
     renamable $pn8 = PTRUE_C_D
     $p0 = COPY killed renamable $pn8
@@ -42,3 +42,24 @@ body:             |
     RET_ReallyLR implicit killed $pn0
 
 ...
+---
+name:            pnr_to_pnr
+alignment:       4
+tracksRegLiveness: true
+tracksDebugUserValues: true
+frameInfo:
+  maxAlignment:    1
+  maxCallFrameSize: 0
+machineFunctionInfo:
+  hasRedZone:      false
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: pnr_to_pnr
+    ; CHECK: renamable $pn8 = PTRUE_C_H
+    ; CHECK-NEXT: $p0 = ORR_PPzPP $p8, $p8, killed $p8, implicit-def $pn0
+    ; CHECK-NEXT: RET_ReallyLR implicit killed $pn0
+    renamable $pn8 = PTRUE_C_H
+    $pn0 = COPY killed renamable $pn8
+    RET_ReallyLR implicit killed $pn0
+
+...