From d7263d6d6d120a833fb45a17924117aad7412a99 Mon Sep 17 00:00:00 2001
From: David Green <david.green@arm.com>
Date: Wed, 13 Nov 2024 09:22:08 +0000
Subject: [PATCH] [AArch64] Use second reg class in genSubAdd2SubSub machine
 combine.

In case the first operand is a physical register with no register class, use
the second operand of the sub as the register class for the new virtual
register in genSubAdd2SubSub machine combine.
---
 llvm/lib/Target/AArch64/AArch64InstrInfo.cpp  |  3 +-
 .../AArch64/machine-combiner-subadd2.mir      | 39 +++++++++++++++++--
 2 files changed, 38 insertions(+), 4 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index b6ac076f4a6d88..a470c03efd5eb4 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -7369,7 +7369,8 @@ genSubAdd2SubSub(MachineFunction &MF, MachineRegisterInfo &MRI,
   bool RegBIsKill = AddMI->getOperand(IdxOpd1).isKill();
   Register RegC = AddMI->getOperand(IdxOtherOpd).getReg();
   bool RegCIsKill = AddMI->getOperand(IdxOtherOpd).isKill();
-  Register NewVR = MRI.createVirtualRegister(MRI.getRegClass(RegA));
+  Register NewVR =
+      MRI.createVirtualRegister(MRI.getRegClass(Root.getOperand(2).getReg()));
 
   unsigned Opcode = Root.getOpcode();
   if (Opcode == AArch64::SUBSWrr)
diff --git a/llvm/test/CodeGen/AArch64/machine-combiner-subadd2.mir b/llvm/test/CodeGen/AArch64/machine-combiner-subadd2.mir
index 9da0808345a098..f0c9e18bc702b0 100644
--- a/llvm/test/CodeGen/AArch64/machine-combiner-subadd2.mir
+++ b/llvm/test/CodeGen/AArch64/machine-combiner-subadd2.mir
@@ -8,7 +8,7 @@
 # 32 bit.
 
 # CHECK-LABEL: name: test1
-# CHECK:       [[TMP:%[0-9]+]]:gpr32common = SUBWrr killed %3, %4
+# CHECK:       [[TMP:%[0-9]+]]:gpr32 = SUBWrr killed %3, %4
 # CHECK-NEXT:  %7:gpr32 = SUBWrr killed [[TMP]], %5
 
 name:            test1
@@ -41,7 +41,7 @@ body:              |
 # 64 bit.
 
 # CHECK-LABEL: name: test2
-# CHECK:       [[TMP:%[0-9]+]]:gpr64common = SUBXrr killed %3, %4
+# CHECK:       [[TMP:%[0-9]+]]:gpr64 = SUBXrr killed %3, %4
 # CHECK-NEXT:  %7:gpr64 = SUBXrr killed [[TMP]], %5
 
 name:            test2
@@ -107,7 +107,7 @@ body:              |
 # Dead define of flag registers should not block transformation.
 
 # CHECK-LABEL: name: test4
-# CHECK:       [[TMP:%[0-9]+]]:gpr64common = SUBXrr killed %3, %4
+# CHECK:       [[TMP:%[0-9]+]]:gpr64 = SUBXrr killed %3, %4
 # CHECK-NEXT:  %7:gpr64 = SUBXrr killed [[TMP]], %5
 
 name:            test4
@@ -264,3 +264,36 @@ body:              |
     RET_ReallyLR implicit $x0
 
 ...
+---
+# WZR use
+
+# CHECK-LABEL: name: wzr
+# CHECK:       [[TMP:%[0-9]+]]:gpr32 = SUBWrr killed $wzr, %4
+# CHECK-NEXT:  %7:gpr32 = SUBWrr killed [[TMP]], %5
+
+name:            wzr
+registers:
+  - { id: 0, class: gpr32common }
+  - { id: 1, class: gpr32 }
+  - { id: 2, class: gpr32 }
+  - { id: 3, class: gpr32common }
+  - { id: 4, class: gpr32common }
+  - { id: 5, class: gpr32 }
+  - { id: 6, class: gpr32 }
+  - { id: 7, class: gpr32 }
+  - { id: 8, class: gpr32 }
+body:              |
+  bb.0:
+    %2:gpr32 = COPY $w2
+    %1:gpr32 = COPY $w1
+    %0:gpr32common = COPY $w0
+    %3:gpr32common = ORRWri %2:gpr32, 1600
+    %4:gpr32common = ADDWri %0:gpr32common, 100, 0
+    %5:gpr32 = EORWrs %1:gpr32, %4:gpr32common, 8
+    %6:gpr32 = ADDWrr %5:gpr32, %4:gpr32common
+    %7:gpr32 = SUBWrr killed $wzr, killed %6:gpr32
+    %8:gpr32 = EORWrs killed %7:gpr32, %5:gpr32, 141
+    $w0 = COPY %8:gpr32
+    RET_ReallyLR implicit $w0
+
+...