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

[Sander de Smalen via llvm-commits]

https://github.com/sdesmalen-arm updated https://github.com/llvm/llvm-project/pull/70682

>From 51d58121a57d46e83f140bbb5fe75a53271c872a Mon Sep 17 00:00:00 2001
From: Sander de Smalen <sander.desmalen at arm.com>
Date: Mon, 30 Oct 2023 16:02:12 +0000
Subject: [PATCH 1/2] [AArch64] Also implement PNR -> PNR copies.

Previously we only implemented PNR -> PPR and PPR -> PNR copies.
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp  | 47 ++++++++-----------
 .../{PNRtoPPRCopy.mir => sve2p1_copy_pnr.mir} | 23 ++++++++-
 2 files changed, 41 insertions(+), 29 deletions(-)
 rename llvm/test/CodeGen/AArch64/{PNRtoPPRCopy.mir => sve2p1_copy_pnr.mir} (70%)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index e9b9f8013abeaf6..c54b1d9f340dec7 100644
--- 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 b/llvm/test/CodeGen/AArch64/sve2p1_copy_pnr.mir
similarity index 70%
rename from llvm/test/CodeGen/AArch64/PNRtoPPRCopy.mir
rename to llvm/test/CodeGen/AArch64/sve2p1_copy_pnr.mir
index 5b1e24ea732f25d..d6a87a42a79e000 100644
--- a/llvm/test/CodeGen/AArch64/PNRtoPPRCopy.mir
+++ b/llvm/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
+
+...

>From bb6ef33a868396d18d6e8d03484c01f10a1f9726 Mon Sep 17 00:00:00 2001
From: Sander de Smalen <sander.desmalen at arm.com>
Date: Mon, 30 Oct 2023 16:38:22 +0000
Subject: [PATCH 2/2] clang-format

---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index c54b1d9f340dec7..c9f7fef66c9f5bc 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -4441,9 +4441,9 @@ void AArch64InstrInfo::copyPhysReg(MachineBasicBlock &MBB,
 
     if (PPRSrcReg != PPRDestReg) {
       auto NewMI = BuildMI(MBB, I, DL, get(AArch64::ORR_PPzPP), PPRDestReg)
-          .addReg(PPRSrcReg) // Pg
-          .addReg(PPRSrcReg)
-          .addReg(PPRSrcReg, getKillRegState(KillSrc));
+                       .addReg(PPRSrcReg) // Pg
+                       .addReg(PPRSrcReg)
+                       .addReg(PPRSrcReg, getKillRegState(KillSrc));
       if (DestIsPNR)
         NewMI.addDef(DestReg, RegState::Implicit);
     }