[AArch64] ORRWrs is copy instruction when there's no implicit def of the X register

llvm/include/llvm/CodeGen/TargetInstrInfo.h

@@ -1025,6 +1025,11 @@ class TargetInstrInfo : public MCInstrInfo {
     return std::nullopt;
   }
 
+  virtual std::optional<DestSourcePair>
+  isCopyLikeInstrImpl(const MachineInstr &MI) const {
+    return std::nullopt;
+  }
+
   /// Return true if the given terminator MI is not expected to spill. This
   /// sets the live interval as not spillable and adjusts phi node lowering to
   /// not introduce copies after the terminator. Use with care, these are
@@ -1050,6 +1055,14 @@ class TargetInstrInfo : public MCInstrInfo {
     return isCopyInstrImpl(MI);
   }
 
+  // Similar to `isCopyInstr`, but adds non-copy semantics on MIR, but
+  // ultimately generates a copy instruction.
+  std::optional<DestSourcePair> isCopyLikeInstr(const MachineInstr &MI) const {
+    if (auto IsCopyInstr = isCopyInstr(MI))
+      return IsCopyInstr;
+    return isCopyLikeInstrImpl(MI);
+  }
+
   bool isFullCopyInstr(const MachineInstr &MI) const {
     auto DestSrc = isCopyInstr(MI);
     if (!DestSrc)

llvm/lib/CodeGen/LiveDebugValues/InstrRefBasedImpl.cpp

@@ -2116,7 +2116,7 @@ bool InstrRefBasedLDV::transferSpillOrRestoreInst(MachineInstr &MI) {
 }
 
 bool InstrRefBasedLDV::transferRegisterCopy(MachineInstr &MI) {
-  auto DestSrc = TII->isCopyInstr(MI);
+  auto DestSrc = TII->isCopyLikeInstr(MI);
   if (!DestSrc)
     return false;
 

llvm/lib/CodeGen/LiveDebugValues/VarLocBasedImpl.cpp

@@ -1364,7 +1364,7 @@ void VarLocBasedLDV::removeEntryValue(const MachineInstr &MI,
     // TODO: Try to keep tracking of an entry value if we encounter a propagated
     // DBG_VALUE describing the copy of the entry value. (Propagated entry value
     // does not indicate the parameter modification.)
-    auto DestSrc = TII->isCopyInstr(*TransferInst);
+    auto DestSrc = TII->isCopyLikeInstr(*TransferInst);
     if (DestSrc) {
       const MachineOperand *SrcRegOp, *DestRegOp;
       SrcRegOp = DestSrc->Source;
@@ -1840,7 +1840,7 @@ void VarLocBasedLDV::transferRegisterCopy(MachineInstr &MI,
                                            OpenRangesSet &OpenRanges,
                                            VarLocMap &VarLocIDs,
                                            TransferMap &Transfers) {
-  auto DestSrc = TII->isCopyInstr(MI);
+  auto DestSrc = TII->isCopyLikeInstr(MI);
   if (!DestSrc)
     return;
 

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

@@ -9180,24 +9180,46 @@ void AArch64InstrInfo::buildClearRegister(Register Reg, MachineBasicBlock &MBB,
 
 std::optional<DestSourcePair>
 AArch64InstrInfo::isCopyInstrImpl(const MachineInstr &MI) const {
-
   // AArch64::ORRWrs and AArch64::ORRXrs with WZR/XZR reg
   // and zero immediate operands used as an alias for mov instruction.
-  if (MI.getOpcode() == AArch64::ORRWrs &&
-      MI.getOperand(1).getReg() == AArch64::WZR &&
-      MI.getOperand(3).getImm() == 0x0) {
+  bool OpIsORRWrs = MI.getOpcode() == AArch64::ORRWrs;
+  bool OpIsORRXrs = MI.getOpcode() == AArch64::ORRXrs;
+  if (!(OpIsORRWrs || OpIsORRXrs) || MI.getOperand(3).getImm() != 0x0)
+    return std::nullopt;
+  Register Reg1 = MI.getOperand(1).getReg();
+
+  if (OpIsORRWrs && Reg1 == AArch64::WZR) {
+    Register Reg0 = MI.getOperand(0).getReg();
+    // ORRWrs is copy instruction when there's no implicit def of the X
+    // register.
+    if (Reg0.isPhysical()) {
+      const MachineFunction *MF = MI.getMF();
+      const TargetRegisterInfo *TRI = MF->getSubtarget().getRegisterInfo();
+      for (const MachineOperand &MO : MI.implicit_operands())
+        if (MO.isDef() && MO.isImplicit() &&
+            TRI->isSubRegister(MO.getReg(), Reg0)) {
+          return std::nullopt;
+        }
+    }
     return DestSourcePair{MI.getOperand(0), MI.getOperand(2)};
   }
 
-  if (MI.getOpcode() == AArch64::ORRXrs &&
-      MI.getOperand(1).getReg() == AArch64::XZR &&
-      MI.getOperand(3).getImm() == 0x0) {
+  if (OpIsORRXrs && Reg1 == AArch64::XZR) {
     return DestSourcePair{MI.getOperand(0), MI.getOperand(2)};
   }
 
   return std::nullopt;
 }
 
+std::optional<DestSourcePair>
+AArch64InstrInfo::isCopyLikeInstrImpl(const MachineInstr &MI) const {
+  if (MI.getOpcode() == AArch64::ORRWrs &&
+      MI.getOperand(1).getReg() == AArch64::WZR &&
+      MI.getOperand(3).getImm() == 0x0)
+    return DestSourcePair{MI.getOperand(0), MI.getOperand(2)};
+  return std::nullopt;
+}
+
 std::optional<RegImmPair>
 AArch64InstrInfo::isAddImmediate(const MachineInstr &MI, Register Reg) const {
   int Sign = 1;
@@ -9241,7 +9263,7 @@ static std::optional<ParamLoadedValue>
 describeORRLoadedValue(const MachineInstr &MI, Register DescribedReg,
                        const TargetInstrInfo *TII,
                        const TargetRegisterInfo *TRI) {
-  auto DestSrc = TII->isCopyInstr(MI);
+  auto DestSrc = TII->isCopyLikeInstr(MI);
   if (!DestSrc)
     return std::nullopt;
 

llvm/lib/Target/AArch64/AArch64InstrInfo.h

@@ -401,6 +401,8 @@ class AArch64InstrInfo final : public AArch64GenInstrInfo {
   /// registers as machine operands.
   std::optional<DestSourcePair>
   isCopyInstrImpl(const MachineInstr &MI) const override;
+  std::optional<DestSourcePair>
+  isCopyLikeInstrImpl(const MachineInstr &MI) const override;
 
 private:
   unsigned getInstBundleLength(const MachineInstr &MI) const;

llvm/test/CodeGen/AArch64/machine-cp-sub-reg.mir

@@ -0,0 +1,33 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 4
+# RUN: llc -o - %s -O3 --run-pass=machine-cp -mcp-use-is-copy-instr -mtriple=arm64-apple-macos -mcpu=apple-m1 --verify-machineinstrs | FileCheck %s
+
+---
+name: test
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: test
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT:   liveins: $w0
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   $x8 = ORRXrs $xzr, undef $x0, 0, implicit $w0
+  ; CHECK-NEXT:   $w8 = ORRWrs $wzr, $w0, 0, implicit-def $x8
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   liveins: $x8
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   $x0 = ADDXri $x8, 1, 0
+  ; CHECK-NEXT:   RET undef $lr, implicit $x0
+  bb.0:
+    successors: %bb.1(0x80000000)
+    liveins: $w0
+
+    $x8 = ORRXrs $xzr, undef $x0, 0, implicit $w0
+    $w8 = ORRWrs $wzr, $w0, 0, implicit-def $x8
+
+  bb.1:
+    liveins: $x8
+    $x0 = ADDXri $x8, 1, 0
+
+    RET undef $lr, implicit $x0
+...