AArch64: Allow ZEXT+COPY -> FMOV peephole for ZPR registers as well (#135436)

Author: MatzeB

llvm/lib/Target/AArch64/AArch64MIPeepholeOpt.cpp

@@ -261,15 +261,18 @@ bool AArch64MIPeepholeOpt::visitORR(MachineInstr &MI) {
     // A COPY from an FPR will become a FMOVSWr, so do so now so that we know
     // that the upper bits are zero.
     if (RC != &AArch64::FPR32RegClass &&
-        ((RC != &AArch64::FPR64RegClass && RC != &AArch64::FPR128RegClass) ||
+        ((RC != &AArch64::FPR64RegClass && RC != &AArch64::FPR128RegClass &&
+          RC != &AArch64::ZPRRegClass) ||
          SrcMI->getOperand(1).getSubReg() != AArch64::ssub))
       return false;
-    Register CpySrc = SrcMI->getOperand(1).getReg();
+    Register CpySrc;
     if (SrcMI->getOperand(1).getSubReg() == AArch64::ssub) {
       CpySrc = MRI->createVirtualRegister(&AArch64::FPR32RegClass);
       BuildMI(*SrcMI->getParent(), SrcMI, SrcMI->getDebugLoc(),
               TII->get(TargetOpcode::COPY), CpySrc)
           .add(SrcMI->getOperand(1));
+    } else {
+      CpySrc = SrcMI->getOperand(1).getReg();
     }
     BuildMI(*SrcMI->getParent(), SrcMI, SrcMI->getDebugLoc(),
             TII->get(AArch64::FMOVSWr), SrcMI->getOperand(0).getReg())

llvm/test/CodeGen/AArch64/peephole-orr.mir

@@ -1,6 +1,49 @@
 # NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
 # RUN: llc -run-pass=aarch64-mi-peephole-opt -o - -mtriple=aarch64-unknown-linux -verify-machineinstrs %s | FileCheck %s
-
+---
+name: copy_fpr128_gpr32
+body: |
+  bb.0:
+    liveins: $q0
+    ; CHECK-LABEL: name: copy_fpr128_gpr32
+    ; CHECK: liveins: $q0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:fpr128 = COPY $q0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:fpr32 = COPY [[COPY]].ssub
+    ; CHECK-NEXT: [[FMOVSWr:%[0-9]+]]:gpr32 = FMOVSWr [[COPY1]]
+    %0:fpr128 = COPY $q0
+    %1:gpr32 = COPY %0.ssub:fpr128
+    %2:gpr32 = ORRWrs $wzr, killed %1:gpr32, 0
+...
+---
+name: copy_fpr32_gpr32
+body: |
+  bb.0:
+    liveins: $s0
+    ; CHECK-LABEL: name: copy_fpr32_gpr32
+    ; CHECK: liveins: $s0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:fpr32 = COPY $s0
+    ; CHECK-NEXT: [[FMOVSWr:%[0-9]+]]:gpr32 = FMOVSWr [[COPY]]
+    %0:fpr32 = COPY $s0
+    %1:gpr32 = COPY %0:fpr32
+    %2:gpr32 = ORRWrs $wzr, killed %1:gpr32, 0
+...
+---
+name: copy_zpr_gpr32
+body: |
+  bb.0:
+    liveins: $z0
+    ; CHECK-LABEL: name: copy_zpr_gpr32
+    ; CHECK: liveins: $z0
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[COPY:%[0-9]+]]:zpr = COPY $z0
+    ; CHECK-NEXT: [[COPY1:%[0-9]+]]:fpr32 = COPY [[COPY]].ssub
+    ; CHECK-NEXT: [[FMOVSWr:%[0-9]+]]:gpr32 = FMOVSWr [[COPY1]]
+    %0:zpr = COPY $z0
+    %1:gpr32 = COPY %0.ssub:zpr
+    %2:gpr32 = ORRWrs $wzr, killed %1:gpr32, 0
+...
 ---
 name: copy_multiple_uses
 tracksRegLiveness: true