[AArch64][SVE] Zero other lanes when doing OR reduction on unpacked predicate using ptest.

[sdesmalen]

Summary

When the predicate vector is unpacked, we cannot assume anything about the
values in the other lanes. We have to make sure we use the correct
predicate where we know that the other lanes have been zeroed.

[RosieSumpter]

LGTM

[sdesmalen]

Regenerated tests after rebase onto D128835.

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

if (!OpVT.isScalableVector() || OpVT.getVectorElementType() != MVT::i1)
     return SDValue();
 
   SDValue Pg = getPredicateForVector(DAG, DL, OpVT);
 
   switch (ReduceOp.getOpcode()) {
   default:
     return SDValue();
   case ISD::VECREDUCE_OR:
-    if (isAllActivePredicate(DAG, Pg))
+    if (isAllActivePredicate(DAG, Pg) && OpVT == MVT::nxv16i1)
       // The predicate can be 'Op' because
       // vecreduce_or(Op & <all true>) <=> vecreduce_or(Op).
       return getPTest(DAG, VT, Op, Op, AArch64CC::ANY_ACTIVE);
     else
       return getPTest(DAG, VT, Pg, Op, AArch64CC::ANY_ACTIVE);
   case ISD::VECREDUCE_AND: {
     Op = DAG.getNode(ISD::XOR, DL, OpVT, Op, Pg);
     return getPTest(DAG, VT, Pg, Op, AArch64CC::NONE_ACTIVE);

llvm/test/CodeGen/AArch64/sve-int-pred-reduce.ll

 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.or.i1.nxv16i1(<vscale x 16 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_or_nxv8i1(<vscale x 8 x i1> %vec) {
 ; CHECK-LABEL: reduce_or_nxv8i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.h
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.or.i1.nxv8i1(<vscale x 8 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_or_nxv4i1(<vscale x 4 x i1> %vec) {
 ; CHECK-LABEL: reduce_or_nxv4i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.or.i1.nxv4i1(<vscale x 4 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_or_nxv2i1(<vscale x 2 x i1> %vec) {
 ; CHECK-LABEL: reduce_or_nxv2i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.or.i1.nxv2i1(<vscale x 2 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_or_nxv1i1(<vscale x 1 x i1> %vec) {
 ; CHECK-LABEL: reduce_or_nxv1i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    punpklo p1.h, p1.b
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.or.i1.nxv1i1(<vscale x 1 x i1> %vec)
   ret i1 %res
 }
 
 ; XORV
 

 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.smin.i1.nxv16i1(<vscale x 16 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_smin_nxv8i1(<vscale x 8 x i1> %vec) {
 ; CHECK-LABEL: reduce_smin_nxv8i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.h
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.smin.i1.nxv8i1(<vscale x 8 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_smin_nxv4i1(<vscale x 4 x i1> %vec) {
 ; CHECK-LABEL: reduce_smin_nxv4i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.smin.i1.nxv4i1(<vscale x 4 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_smin_nxv2i1(<vscale x 2 x i1> %vec) {
 ; CHECK-LABEL: reduce_smin_nxv2i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.smin.i1.nxv2i1(<vscale x 2 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_smin_nxv1i1(<vscale x 1 x i1> %vec) {
 ; CHECK-LABEL: reduce_smin_nxv1i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    punpklo p1.h, p1.b
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.smin.i1.nxv1i1(<vscale x 1 x i1> %vec)
   ret i1 %res
 }
 
 ; UMAXV
 
 define i1 @reduce_umax_nxv16i1(<vscale x 16 x i1> %vec) {
 ; CHECK-LABEL: reduce_umax_nxv16i1:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptest p0, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.umax.i1.nxv16i1(<vscale x 16 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_umax_nxv8i1(<vscale x 8 x i1> %vec) {
 ; CHECK-LABEL: reduce_umax_nxv8i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.h
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.umax.i1.nxv8i1(<vscale x 8 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_umax_nxv4i1(<vscale x 4 x i1> %vec) {
 ; CHECK-LABEL: reduce_umax_nxv4i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.umax.i1.nxv4i1(<vscale x 4 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_umax_nxv2i1(<vscale x 2 x i1> %vec) {
 ; CHECK-LABEL: reduce_umax_nxv2i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.umax.i1.nxv2i1(<vscale x 2 x i1> %vec)
   ret i1 %res
 }
 
 define i1 @reduce_umax_nxv1i1(<vscale x 1 x i1> %vec) {
 ; CHECK-LABEL: reduce_umax_nxv1i1:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    punpklo p1.h, p1.b
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %res = call i1 @llvm.vector.reduce.umax.i1.nxv1i1(<vscale x 1 x i1> %vec)
   ret i1 %res
 }
 
 ; UMINV
 

llvm/test/CodeGen/AArch64/sve-vecreduce-fold.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
 ; RUN: llc -mtriple=aarch64 -mattr=+sve < %s | FileCheck %s
 
 ;
 ; OR reductions
 ;
 
 define i1 @reduce_or_insert_subvec_into_zero(<vscale x 4 x i1> %in) {
 ; CHECK-LABEL: reduce_or_insert_subvec_into_zero:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %t = call <vscale x 16 x i1> @llvm.vector.insert.nxv16i1.nxv4i1(<vscale x 16 x i1> zeroinitializer, <vscale x 4 x i1> %in, i64 0)
   %res = call i1 @llvm.vector.reduce.or.nxv16i1(<vscale x 16 x i1> %t)
   ret i1 %res
 }
 
 define i1 @reduce_or_insert_subvec_into_poison(<vscale x 4 x i1> %in) {
 ; CHECK-LABEL: reduce_or_insert_subvec_into_poison:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    ptest p0, p0.b
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    ptest p1, p0.b
 ; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %t = call <vscale x 16 x i1> @llvm.vector.insert.nxv16i1.nxv4i1(<vscale x 16 x i1> poison, <vscale x 4 x i1> %in, i64 0)
   %res = call i1 @llvm.vector.reduce.or.nxv16i1(<vscale x 16 x i1> %t)
   ret i1 %res
 }
 
 define i1 @reduce_or_insert_subvec_into_nonzero(<vscale x 4 x i1> %in, <vscale x 16 x i1> %vec) {