[AArch64] Support all extend/shift op for pattern: (ExtendOrShfitNode - Y) + Z --> (Z - Y) + ExtendOrShfitNode

[bcl5980]

Summary

Followup rG325a3083b5b24
Add SRL/SRA/xxx_EXTEND/AND support

[dmgreen]

Sounds good. Should we be handling uxtw and sxtw type operands too? As in getExtendTypeForNode.

[bcl5980]

minor format update.

[bcl5980]

Ping

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16910

Perhaps change canbeExtend to isExtendOrShiftOperand, or something like it?

16915

Do you have tests for the other cases?

16928

What overflow are you thinking of? Shifting past the bitwidth?

[bcl5980]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16915

Not sure how to generate ISD::ANY_EXTEND/ISD::SIGN_EXTEND_INREG from llvm-ir. Can you give me some suggestions for that?

16928

Yeah, I'm not sure if we can do this when the shift amount is larger than bitwidth.

[dmgreen]

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

16915

SIGN_EXTEND_INREG is usually lsl 48; ashr 48. Any extend can be difficult (especially if we expect the value to be used). Adding some extra 64bit tests would be useful too, for the different And cases for example.

[bcl5980]

rebase

[dmgreen]

Thanks. From what I can tell this looks OK, and the performance should be better. LGTM

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (Elt0->getOpcode() == ISD::EXTRACT_VECTOR_ELT &&
     SDValue Ext = DAG.getNode(ISD::ANY_EXTEND, DL, ExtVT, VecToExtend);
     return DAG.getNode(ISD::EXTRACT_SUBVECTOR, DL, MVT::v2i32, Ext,
                        SubvectorIdx);
   }
 
   return SDValue();
 }
 
-// ((X >> C) - Y) + Z --> (Z - Y) + (X >> C)
-static SDValue performAddCombineSubShift(SDNode *N, SDValue SUB, SDValue Z,
-                                         SelectionDAG &DAG) {
-  auto IsOneUseShiftC = [&](SDValue Shift) {
-    if (!Shift.hasOneUse())
-      return false;
+// Check an node is an extend or shift operand
+static bool isExtendOrShiftOperand(SDValue N) {
+  unsigned Opcode = N.getOpcode();
+  if (Opcode == ISD::SIGN_EXTEND || Opcode == ISD::SIGN_EXTEND_INREG ||
+      Opcode == ISD::ZERO_EXTEND || Opcode == ISD::ANY_EXTEND) {
+    EVT SrcVT;
+    if (Opcode == ISD::SIGN_EXTEND_INREG)
+      SrcVT = cast<VTSDNode>(N.getOperand(1))->getVT();
+    else
+      SrcVT = N.getOperand(0).getValueType();
 
-    // TODO: support SRL and SRA also
-    if (Shift.getOpcode() != ISD::SHL)
+    return SrcVT == MVT::i32 || SrcVT == MVT::i16 || SrcVT == MVT::i8;
+  } else if (Opcode == ISD::AND) {
+    ConstantSDNode *CSD = dyn_cast<ConstantSDNode>(N.getOperand(1));
+    if (!CSD)
       return false;
+    uint64_t AndMask = CSD->getZExtValue();
+    return AndMask == 0xff || AndMask == 0xffff || AndMask == 0xffffffff;
+  } else if (Opcode == ISD::SHL || Opcode == ISD::SRL || Opcode == ISD::SRA) {
+    return isa<ConstantSDNode>(N.getOperand(1));
+  }
 
-    if (!isa<ConstantSDNode>(Shift.getOperand(1)))
-      return false;
-    return true;
+}
+
+// (N - Y) + Z --> (Z - Y) + N
+// when N is an extend or shift operand
+static SDValue performAddCombineSubShift(SDNode *N, SDValue SUB, SDValue Z,
+                                         SelectionDAG &DAG) {
+  auto IsOneUseExtend = [](SDValue N) {
+    return N.hasOneUse() && isExtendOrShiftOperand(N);
   };
 
   // DAGCombiner will revert the combination when Z is constant cause
   // dead loop. So don't enable the combination when Z is constant.
   // If Z is one use shift C, we also can't do the optimization.
   // It will falling to self infinite loop.
-  if (isa<ConstantSDNode>(Z) || IsOneUseShiftC(Z))
+  if (isa<ConstantSDNode>(Z) || IsOneUseExtend(Z))
     return SDValue();
 
   if (SUB.getOpcode() != ISD::SUB || !SUB.hasOneUse())
     return SDValue();
 
   SDValue Shift = SUB.getOperand(0);
-  if (!IsOneUseShiftC(Shift))
+  if (!IsOneUseExtend(Shift))
     return SDValue();
 
   SDLoc DL(N);
   EVT VT = N->getValueType(0);
 
   SDValue Y = SUB.getOperand(1);
   SDValue NewSub = DAG.getNode(ISD::SUB, DL, VT, Z, Y);
   return DAG.getNode(ISD::ADD, DL, VT, NewSub, Shift);

llvm/test/CodeGen/AArch64/addsub.ll

; CHECK-NEXT:    ret
   %add = add i32 %sub, %z
   ret i32 %add
 }
 
 ; ((X >> C) - Y) + Z --> (Z - Y) + (X >> C)
 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
   %add = add i32 %sub, %z
   ret i32 %add
 }
 
 ; ((X >> C) - Y) + Z --> (Z - Y) + (X >> C)
 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
   %add = add i32 %sub, %z
   ret i32 %add
 }
 
 ; ((sext X) - Y) + Z --> (Z - Y) + (sext X)
 define i64 @commute_subop0_sext(i32 %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: commute_subop0_sext:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    // kill: def $w0 killed $w0 def $x0
-; CHECK-NEXT:    sxtw x8, w0
-; CHECK-NEXT:    sub x8, x8, x1
-; CHECK-NEXT:    add x0, x8, x2
+; CHECK-NEXT:    sub x8, x2, x1
+; CHECK-NEXT:    add x0, x8, w0, sxtw
 ; CHECK-NEXT:    ret
   %sext = sext i32 %x to i64
   %sub = sub i64 %sext, %y
   %add = add i64 %sub, %z
   ret i64 %add
 }
 
 ; ((sext_inreg X) - Y) + Z --> (Z - Y) + (sext_inreg X)
 define i64 @commute_subop0_sext_inreg(i64 %x, i64 %y, i64 %z) {
 ; CHECK-LABEL: commute_subop0_sext_inreg:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    sxth x8, w0
-; CHECK-NEXT:    sub x8, x8, x1
-; CHECK-NEXT:    add x0, x8, x2
+; CHECK-NEXT:    sub x8, x2, x1
+; CHECK-NEXT:    add x0, x8, w0, sxth
 ; CHECK-NEXT:    ret
   %shl = shl i64 %x, 48
   %ashr = ashr i64 %shl, 48
   %sub = sub i64 %ashr, %y
   %add = add i64 %sub, %z
   ret i64 %add
 }
 
 ; ((zext X) - Y) + Z --> (Z - Y) + (zext X)
 define i32 @commute_subop0_zext(i16 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: commute_subop0_zext:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    and w8, w0, #0xffff
-; CHECK-NEXT:    sub w8, w8, w1
-; CHECK-NEXT:    add w0, w8, w2
+; CHECK-NEXT:    sub w8, w2, w1
+; CHECK-NEXT:    add w0, w8, w0, uxth
 ; CHECK-NEXT:    ret
   %zext = zext i16 %x to i32
   %sub = sub i32 %zext, %y
   %add = add i32 %sub, %z
   ret i32 %add
 }
 
 
 ; ((anyext X) - Y) + Z --> (Z - Y) + (anyext X)
 define i8 @commute_subop0_anyext(i16 %a, i16 %b, i32 %c) {
 ; CHECK-LABEL: commute_subop0_anyext:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    mov w8, #111
-; CHECK-NEXT:    neg w9, w1
+; CHECK-NEXT:    sub w9, w2, w1
 ; CHECK-NEXT:    madd w8, w0, w8, w9
-; CHECK-NEXT:    add w8, w8, w2
 ; CHECK-NEXT:    lsl w8, w8, #3
 ; CHECK-NEXT:    sub w0, w8, #1776
 ; CHECK-NEXT:    ret
   %aa = mul i16 %a, 111
   %bb = add i16 %b, 222
   %a_32 = zext i16 %aa to i32
   %b_32 = zext i16 %bb to i32
   %sub = sub i32 %a_32, %b_32
   %add = add i32 %sub, %c
   %trunc = trunc i32 %add to i8
   %r = shl i8 %trunc, 3
   ret i8 %r
 }
 
 ; ((X and C) - Y) + Z --> (Z - Y) + (X and C)
 define i32 @commute_subop0_and(i32 %x, i32 %y, i32 %z) {
 ; CHECK-LABEL: commute_subop0_and:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    and w8, w0, #0xff
-; CHECK-NEXT:    sub w8, w8, w1
-; CHECK-NEXT:    add w0, w8, w2
+; CHECK-NEXT:    sub w8, w2, w1
+; CHECK-NEXT:    add w0, w8, w0, uxtb
 ; CHECK-NEXT:    ret
   %and = and i32 %x, 255
   %sub = sub i32 %and, %y
   %add = add i32 %sub, %z
   ret i32 %add
 }
 
 ; Z + ((X << C) - Y) --> (Z - Y) + (X << C)