[AArch64][SVE] Add instcombine to convert ptest.last/first to ptest.any

[bsmith]

Summary

This allow for better optimization later in the backend.

This fixes the remaining missed optimizations in D137717.

Depends on D137930

[c-rhodes]

just one suggestion but otherwise LGTM, cheers

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

983

perhaps worth adding a comment describing the combine?

PTEST_<FIRST|LAST>(X, X) is equivalent to PTEST_ANY(X, X).

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

instCombineSVECntElts(InstCombiner &IC, IntrinsicInst &II, unsigned NumElts) {
   return MinNumElts && NumElts >= MinNumElts
              ? Optional<Instruction *>(IC.replaceInstUsesWith(
                    II, ConstantInt::get(II.getType(), MinNumElts)))
              : None;
 }
 
 static Optional<Instruction *> instCombineSVEPTest(InstCombiner &IC,
                                                    IntrinsicInst &II) {
-  IntrinsicInst *Pg = dyn_cast<IntrinsicInst>(II.getArgOperand(0));
-  IntrinsicInst *Op = dyn_cast<IntrinsicInst>(II.getArgOperand(1));
+  Value *PgVal = II.getArgOperand(0);
+  Value *OpVal = II.getArgOperand(1);
+
+  IRBuilder<> Builder(II.getContext());
+  Builder.SetInsertPoint(&II);
+
+  // PTEST_<FIRST|LAST>(X, X) is equivalent to PTEST_ANY(X, X).
+  // Later optimizations prefer this form.
+  if (PgVal == OpVal &&
+      (II.getIntrinsicID() == Intrinsic::aarch64_sve_ptest_first ||
+       II.getIntrinsicID() == Intrinsic::aarch64_sve_ptest_last)) {
+    Value *Ops[] = {PgVal, OpVal};
+    Type *Tys[] = {PgVal->getType()};
+
+    auto *PTest =
+        Builder.CreateIntrinsic(Intrinsic::aarch64_sve_ptest_any, Tys, Ops);
+    PTest->takeName(&II);
+
+    return IC.replaceInstUsesWith(II, PTest);
+  }
+
+  IntrinsicInst *Pg = dyn_cast<IntrinsicInst>(PgVal);
+  IntrinsicInst *Op = dyn_cast<IntrinsicInst>(OpVal);
 
   if (!Pg || !Op)
     return None;
 
   Intrinsic::ID OpIID = Op->getIntrinsicID();
 
-  IRBuilder<> Builder(II.getContext());
-  Builder.SetInsertPoint(&II);
-
   if (Pg->getIntrinsicID() == Intrinsic::aarch64_sve_convert_to_svbool &&
       OpIID == Intrinsic::aarch64_sve_convert_to_svbool &&
       Pg->getArgOperand(0)->getType() == Op->getArgOperand(0)->getType()) {
     Value *Ops[] = {Pg->getArgOperand(0), Op->getArgOperand(0)};
     Type *Tys[] = {Pg->getArgOperand(0)->getType()};
 
     auto *PTest = Builder.CreateIntrinsic(II.getIntrinsicID(), Tys, Ops);
 

llvm/test/CodeGen/AArch64/sve-ptest-removal-cmple.ll

 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc -mtriple=aarch64--linux-gnu -mattr=+sve %s -o - | FileCheck %s
+; RUN: opt -instcombine -S %s | llc -mtriple=aarch64--linux-gnu -mattr=+sve -o - | FileCheck %s
+
+target triple = "aarch64-unknown-linux-gnu"
 
 ;
 ; Immediate Compares
 ;
 
 define i32 @cmple_imm_nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %a) {
 ; CHECK-LABEL: cmple_imm_nxv16i8:
 ; CHECK:       // %bb.0:

 ; but not by the implicit PTEST as part of the compare. However, given the
 ; PTEST mask and source are the same, first is equivalent to any. The same
 ; applies to last active.
 ;
 define i1 @cmp8_ptest_first_xx(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
 ; CHECK-LABEL: cmp8_ptest_first_xx:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cmpge p0.b, p0/z, z0.b, z1.b
-; CHECK-NEXT:    ptest p0, p0.b
-; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.cmpge.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
   %2 = tail call i1 @llvm.aarch64.sve.ptest.first.nxv16i1(<vscale x 16 x i1> %1, <vscale x 16 x i1> %1)
   ret i1 %2
 }
 
 ;
 ; PTEST_LAST(X=CMP8(PG, A, B), X). PTEST is redundant if condition is changed
 ; to any.
 ;
 define i1 @cmp8_ptest_last_xx(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
 ; CHECK-LABEL: cmp8_ptest_last_xx:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cmpge p0.b, p0/z, z0.b, z1.b
-; CHECK-NEXT:    ptest p0, p0.b
-; CHECK-NEXT:    cset w0, lo
+; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.cmpge.nxv16i8(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b)
   %2 = tail call i1 @llvm.aarch64.sve.ptest.last.nxv16i1(<vscale x 16 x i1> %1, <vscale x 16 x i1> %1)
   ret i1 %2
 }
 
 ;
 ; PTEST_ANY(X=CMP8(PG, A, B), X). PTEST is redundant.

 ;
 ; PTEST_FIRST(X=CMP32(PG, A, B), X). PTEST is redundant if condition is changed
 ; to any.
 ;
 define i1 @cmp32_ptest_first_xx(<vscale x 16 x i1> %pg, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
 ; CHECK-LABEL: cmp32_ptest_first_xx:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cmpge p0.s, p0/z, z0.s, z1.s
-; CHECK-NEXT:    ptest p0, p0.b
-; CHECK-NEXT:    cset w0, mi
+; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %1 = tail call <vscale x 4 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv4i1(<vscale x 16 x i1> %pg)
   %2 = tail call <vscale x 4 x i1> @llvm.aarch64.sve.cmpge.nxv4i32(<vscale x 4 x i1> %1, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b)
   %3 = tail call i1 @llvm.aarch64.sve.ptest.first.nxv4i1(<vscale x 4 x i1> %2, <vscale x 4 x i1> %2)
   ret i1 %3
 }
 
 ;
 ; PTEST_LAST(X=CMP32(PG, A, B), X). PTEST is redundant if condition is changed
 ; to any.
 ;
 define i1 @cmp32_ptest_last_xx(<vscale x 16 x i1> %pg, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
 ; CHECK-LABEL: cmp32_ptest_last_xx:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    cmpge p0.s, p0/z, z0.s, z1.s
-; CHECK-NEXT:    ptest p0, p0.b
-; CHECK-NEXT:    cset w0, lo
+; CHECK-NEXT:    cset w0, ne
 ; CHECK-NEXT:    ret
   %1 = tail call <vscale x 4 x i1> @llvm.aarch64.sve.convert.from.svbool.nxv4i1(<vscale x 16 x i1> %pg)
   %2 = tail call <vscale x 4 x i1> @llvm.aarch64.sve.cmpge.nxv4i32(<vscale x 4 x i1> %1, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b)
   %3 = tail call i1 @llvm.aarch64.sve.ptest.last.nxv4i1(<vscale x 4 x i1> %2, <vscale x 4 x i1> %2)
   ret i1 %3
 }
 
 ;

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsics-ptest.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt -S -passes=instcombine < %s | FileCheck %s
 
 target triple = "aarch64-unknown-linux-gnu"
 
+; PTEST first can be changed to any if the mask and operand are the same
+define i1 @ptest_first_to_any(<vscale x 16 x i1> %a) #0 {
+; CHECK-LABEL: @ptest_first_to_any(
+; CHECK-NEXT:    [[OUT:%.*]] = call i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1> [[A:%.*]], <vscale x 16 x i1> [[A]])
+; CHECK-NEXT:    ret i1 [[OUT]]
+;
+  %out = call i1 @llvm.aarch64.sve.ptest.first.nxv16i1(<vscale x 16 x i1> %a, <vscale x 16 x i1> %a)
+  ret i1 %out
+}
+
+; PTEST last can be changed to any if the mask and operand are the same
+define i1 @ptest_last_to_any(<vscale x 16 x i1> %a) #0 {
+; CHECK-LABEL: @ptest_last_to_any(
+; CHECK-NEXT:    [[OUT:%.*]] = call i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1> [[A:%.*]], <vscale x 16 x i1> [[A]])
+; CHECK-NEXT:    ret i1 [[OUT]]
+;
+  %out = call i1 @llvm.aarch64.sve.ptest.last.nxv16i1(<vscale x 16 x i1> %a, <vscale x 16 x i1> %a)
+  ret i1 %out
+}
+
 define i1 @ptest_any1(<vscale x 2 x i1> %a) #0 {
 ; CHECK-LABEL: @ptest_any1(
 ; CHECK-NEXT:    [[MASK:%.*]] = tail call <vscale x 2 x i1> @llvm.aarch64.sve.ptrue.nxv2i1(i32 0)
 ; CHECK-NEXT:    [[OUT:%.*]] = call i1 @llvm.aarch64.sve.ptest.any.nxv2i1(<vscale x 2 x i1> [[MASK]], <vscale x 2 x i1> [[A:%.*]])
 ; CHECK-NEXT:    ret i1 [[OUT]]
 ;
   %mask = tail call <vscale x 2 x i1> @llvm.aarch64.sve.ptrue.nxv2i1(i32 0)
   %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv2i1(<vscale x 2 x i1> %mask)
;
   %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1> %mask)
   %2 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1> %a)
   %out = call i1 @llvm.aarch64.sve.ptest.first.nxv16i1(<vscale x 16 x i1> %1, <vscale x 16 x i1> %2)
   ret i1 %out
 }
 
 define i1 @ptest_first_same_ops(<vscale x 2 x i1> %a) #0 {
 ; CHECK-LABEL: @ptest_first_same_ops(
-; CHECK-NEXT:    [[TMP1:%.*]] = call i1 @llvm.aarch64.sve.ptest.first.nxv2i1(<vscale x 2 x i1> [[A:%.*]], <vscale x 2 x i1> [[A]])
+; CHECK-NEXT:    [[TMP1:%.*]] = call i1 @llvm.aarch64.sve.ptest.any.nxv2i1(<vscale x 2 x i1> [[A:%.*]], <vscale x 2 x i1> [[A]])
 ; CHECK-NEXT:    ret i1 [[TMP1]]
 ;
   %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv2i1(<vscale x 2 x i1> %a)
   %2 = tail call i1 @llvm.aarch64.sve.ptest.first.nxv16i1(<vscale x 16 x i1> %1, <vscale x 16 x i1> %1)
   ret i1 %2
 }
 
 define i1 @ptest_last(<vscale x 8 x i1> %a) #0 {