[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.

Author: davemgreen

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)

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
+
+...