Index: llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
===================================================================
--- llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -16854,15 +16854,15 @@
   return SDValue();
 }
 
-// ((X >> C) - Y) + Z --> (Z - Y) + (X >> C)
+// ((X shift C) - Y) + Z --> (Z - Y) + (X shift C)
 static SDValue performAddCombineSubShift(SDNode *N, SDValue SUB, SDValue Z,
                                          SelectionDAG &DAG) {
   auto IsOneUseShiftC = [&](SDValue Shift) {
     if (!Shift.hasOneUse())
       return false;
 
-    // TODO: support SRL and SRA also
-    if (Shift.getOpcode() != ISD::SHL)
+    if (Shift.getOpcode() != ISD::SHL && Shift.getOpcode() != ISD::SRL &&
+        Shift.getOpcode() != ISD::SRA)
       return false;
 
     if (!isa<ConstantSDNode>(Shift.getOperand(1)))
Index: llvm/test/CodeGen/AArch64/addsub.ll
===================================================================
--- llvm/test/CodeGen/AArch64/addsub.ll
+++ llvm/test/CodeGen/AArch64/addsub.ll
@@ -711,9 +711,8 @@
 define i32 @commute_subop0_lshr(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: commute_subop0_lshr:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    lsr w8, w0, #3
-; CHECK-NEXT:    sub w8, w8, w1
-; CHECK-NEXT:    add w0, w8, w2
+; CHECK-NEXT:    sub w8, w2, w1
+; CHECK-NEXT:    add w0, w8, w0, lsr #3
 ; CHECK-NEXT:    ret
   %lshr = lshr i32 %x, 3
   %sub = sub i32 %lshr, %y
@@ -725,9 +724,8 @@
 define i32 @commute_subop0_ashr(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: commute_subop0_ashr:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    asr w8, w0, #3
-; CHECK-NEXT:    sub w8, w8, w1
-; CHECK-NEXT:    add w0, w8, w2
+; CHECK-NEXT:    sub w8, w2, w1
+; CHECK-NEXT:    add w0, w8, w0, asr #3
 ; CHECK-NEXT:    ret
   %ashr = ashr i32 %x, 3
   %sub = sub i32 %ashr, %y