[AArch64] Guard against non-vector abd long nodes. (#102026)

This fixes a problem if abd nodes are generated more readily (#92576).
The folding of abd nodes into abdl needs to check that the inputs are
the correct form of vector. The added test requires vector legalization
to occur in order to hit the combine at the wrong time.

3 files changed, 49 insertions, 0 deletions

diff --git a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
index 7704321..f0c3afc 100644
--- a/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
+++ b/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp
@@ -21769,6 +21769,7 @@ static SDValue performExtendCombine(SDNode *N,
   // helps the backend to decide that an sabdl2 would be useful, saving a real
   // extract_high operation.
   if (!DCI.isBeforeLegalizeOps() && N->getOpcode() == ISD::ZERO_EXTEND &&
+      N->getOperand(0).getValueType().is64BitVector() &&
       (N->getOperand(0).getOpcode() == ISD::ABDU ||
        N->getOperand(0).getOpcode() == ISD::ABDS)) {
     SDNode *ABDNode = N->getOperand(0).getNode();
diff --git a/llvm/test/CodeGen/AArch64/abds.ll b/llvm/test/CodeGen/AArch64/abds.ll
index d4ad33f..215907c 100644
--- a/llvm/test/CodeGen/AArch64/abds.ll
+++ b/llvm/test/CodeGen/AArch64/abds.ll
@@ -571,6 +571,28 @@ define i32 @abd_sub_i32(i32 %a, i32 %b) nounwind {
   ret i32 %abs
 }
 
+define i64 @vector_legalized(i16 %a, i16 %b) {
+; CHECK-LABEL: vector_legalized:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    sxth w8, w0
+; CHECK-NEXT:    sub w8, w8, w1, sxth
+; CHECK-NEXT:    addp d0, v0.2d
+; CHECK-NEXT:    cmp w8, #0
+; CHECK-NEXT:    cneg w8, w8, mi
+; CHECK-NEXT:    fmov x9, d0
+; CHECK-NEXT:    add x0, x9, x8
+; CHECK-NEXT:    ret
+  %ea = sext i16 %a to i32
+  %eb = sext i16 %b to i32
+  %s = sub i32 %ea, %eb
+  %ab = call i32 @llvm.abs.i32(i32 %s, i1 false)
+  %e = zext i32 %ab to i64
+  %red = call i64 @llvm.vector.reduce.add.v32i64(<32 x i64> zeroinitializer)
+  %z = add i64 %red, %e
+  ret i64 %z
+}
+
 
 declare i8 @llvm.abs.i8(i8, i1)
 declare i16 @llvm.abs.i16(i16, i1)
diff --git a/llvm/test/CodeGen/AArch64/abdu.ll b/llvm/test/CodeGen/AArch64/abdu.ll
index 983db62..f70f095 100644
--- a/llvm/test/CodeGen/AArch64/abdu.ll
+++ b/llvm/test/CodeGen/AArch64/abdu.ll
@@ -409,6 +409,32 @@ define i128 @abd_cmp_i128(i128 %a, i128 %b) nounwind {
   ret i128 %sel
 }
 
+;
+; negative tests
+;
+
+define i64 @vector_legalized(i16 %a, i16 %b) {
+; CHECK-LABEL: vector_legalized:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    and w8, w0, #0xffff
+; CHECK-NEXT:    sub w8, w8, w1, uxth
+; CHECK-NEXT:    cmp w8, #0
+; CHECK-NEXT:    addp d0, v0.2d
+; CHECK-NEXT:    cneg w8, w8, mi
+; CHECK-NEXT:    fmov x9, d0
+; CHECK-NEXT:    add x0, x9, x8
+; CHECK-NEXT:    ret
+  %ea = zext i16 %a to i32
+  %eb = zext i16 %b to i32
+  %s = sub i32 %ea, %eb
+  %ab = call i32 @llvm.abs.i32(i32 %s, i1 false)
+  %e = zext i32 %ab to i64
+  %red = call i64 @llvm.vector.reduce.add.v32i64(<32 x i64> zeroinitializer)
+  %z = add i64 %red, %e
+  ret i64 %z
+}
+
 declare i8 @llvm.abs.i8(i8, i1)
 declare i16 @llvm.abs.i16(i16, i1)
 declare i32 @llvm.abs.i32(i32, i1)