[AArch64][SVE] Add instcombine for PTEST_ANY(X=OP(PG,...), X) -> PTEST_ANY(PG, X))

[c-rhodes]

Summary

Given this is an OR reduction the two are equivalent and later
optimizations (AArch64InstrInfo::optimizePTestInstr) may rewrite the
sequence to use the flag-setting variant of instruction X, to remove the
PTEST altogether.

[bsmith]

LGTM with minor nits

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1335 ↗

(On Diff #464198)

Extra parenthesis

1344–1346 ↗

(On Diff #464198)

This comment is out of place now

[c-rhodes]

In D134946#3830351, @bsmith wrote:

LGTM with minor nits

Thanks for review, I'll address those nits when landing unless there's any objections.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1335 ↗

(On Diff #464198)

Are you sure this is universally true? The s variants do an implicitly ptest using the original instruction's general predicate operand:

brkb p0.b, p1/z, p2.b; ptest p1, p0.b ==> brkbs p0.b, p1/z, p2.b

If you perform this transformation when the general predicates of both instructions differ the implicit ptest will calculate a result by considering a different set of lanes.

Whilst the logic might be safe for your ptest.any case because it's essentially an OR reduction (and thus safe when pg & data == data), the same is unlikely for things like cc.first and cc.none where it's critical to test the exact lanes requested.

Assuming you agree and I've not misunderstood you probably want to do something specific for ptest.any, which is best done during isel lowering.

[c-rhodes]

Restrict optimization to ptest.any and move to inst combine.

[c-rhodes]

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

1335 ↗

(On Diff #464198)

Are you sure this is universally true? The s variants do an implicitly ptest using the original instruction's general predicate operand:

brkb p0.b, p1/z, p2.b; ptest p1, p0.b ==> brkbs p0.b, p1/z, p2.b

If you perform this transformation when the general predicates of both instructions differ the implicit ptest will calculate a result by considering a different set of lanes.

Whilst the logic might be safe for your ptest.any case because it's essentially an OR reduction (and thus safe when pg & data == data), the same is unlikely for things like cc.first and cc.none where it's critical to test the exact lanes requested.

Assuming you agree and I've not misunderstood you probably want to do something specific for ptest.any, which is best done during isel lowering.

You're right this isn't safe for first/last, I missed that. I've updated the patch as suggested and moved the transform to inst combine. Cheers!

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

987

Perhaps PTEST_ANY(X=OP(PG,...), X) -> PTEST_ANY(PG, X)) to show where PG comes from.

993

Looking at the instruction pseudo code for brkns it works differently than the other instructions that set the flags in that its implicit predicate is an all active one. This likely means there's a bug in AArch64InstrInfo::optimizePTestInstr.

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

45

Is the zext necessary? Same goes for the other tests.

[c-rhodes]

• Don't transform brkn intrinsic.
• Update comment.

[c-rhodes]

Remove zext from tests.

[c-rhodes]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

993

Looking at the instruction pseudo code for brkns it works differently than the other instructions that set the flags in that its implicit predicate is an all active one. This likely means there's a bug in AArch64InstrInfo::optimizePTestInstr.

Good spot! I've removed brkn from this patch and will fix the bug in AArch64InstrInfo::optimizePTestInstr in a separate patch. I guess the optimization there can still be done as long as the ptest is called with all active mask rather than the mask from the BRKN instruction?

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

45

Is the zext necessary? Same goes for the other tests.

Hangover from previous revision, fixed.

[c-rhodes]

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

993

Looking at the instruction pseudo code for brkns it works differently than the other instructions that set the flags in that its implicit predicate is an all active one. This likely means there's a bug in AArch64InstrInfo::optimizePTestInstr.

Good spot! I've removed brkn from this patch and will fix the bug in AArch64InstrInfo::optimizePTestInstr in a separate patch. I guess the optimization there can still be done as long as the ptest is called with all active mask rather than the mask from the BRKN instruction?

Fixed in D135655.

llvm/lib/Target/AArch64/AArch64TargetTransformInfo.cpp

return MinNumElts && NumElts >= MinNumElts
              : None;
 }
 
 static Optional<Instruction *> instCombineSVEPTest(InstCombiner &IC,
                                                    IntrinsicInst &II) {
   IntrinsicInst *Op1 = dyn_cast<IntrinsicInst>(II.getArgOperand(0));
   IntrinsicInst *Op2 = dyn_cast<IntrinsicInst>(II.getArgOperand(1));
 
-  if (Op1 && Op2 &&
-      Op1->getIntrinsicID() == Intrinsic::aarch64_sve_convert_to_svbool &&
-      Op2->getIntrinsicID() == Intrinsic::aarch64_sve_convert_to_svbool &&
-      Op1->getArgOperand(0)->getType() == Op2->getArgOperand(0)->getType()) {
+  if (!Op1 || !Op2)
+    return None;
 
-    IRBuilder<> Builder(II.getContext());
-    Builder.SetInsertPoint(&II);
 
+  if (Op1->getIntrinsicID() == Intrinsic::aarch64_sve_convert_to_svbool &&
+      Op2->getIntrinsicID() == Intrinsic::aarch64_sve_convert_to_svbool &&
+      Op1->getArgOperand(0)->getType() == Op2->getArgOperand(0)->getType()) {
     Value *Ops[] = {Op1->getArgOperand(0), Op2->getArgOperand(0)};
     Type *Tys[] = {Op1->getArgOperand(0)->getType()};
 
     auto *PTest = Builder.CreateIntrinsic(II.getIntrinsicID(), Tys, Ops);
 
     PTest->takeName(&II);
     return IC.replaceInstUsesWith(II, PTest);
   }
 
+  // Transform PTEST_ANY(X=OP(PG,...), X) -> PTEST_ANY(PG, X)).
+  // Later optimizations may rewrite sequence to use the flag-setting variant
+  // of instruction X to remove PTEST.
+  if ((Op1 == Op2) &&
+      (II.getIntrinsicID() == Intrinsic::aarch64_sve_ptest_any) &&
+      ((Op1->getIntrinsicID() == Intrinsic::aarch64_sve_brkb_z) ||
+       (Op1->getIntrinsicID() == Intrinsic::aarch64_sve_rdffr_z))) {
+    Value *Ops[] = {Op1->getArgOperand(0), Op1};
+    Type *Tys[] = {Op1->getType()};
+
+    auto *PTest = Builder.CreateIntrinsic(II.getIntrinsicID(), Tys, Ops);
+    PTest->takeName(&II);
+
+    return IC.replaceInstUsesWith(II, PTest);
+  }
+
   return None;
 }
 
 static Optional<Instruction *> instCombineSVEVectorFMLA(InstCombiner &IC,
                                                         IntrinsicInst &II) {
   // fold (fadd p a (fmul p b c)) -> (fma p a b c)
   Value *P = II.getOperand(0);
   Value *A = II.getOperand(1);

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

 ;
   %mask = tail call <vscale x 2 x i1> @llvm.aarch64.sve.ptrue.nxv2i1(i32 31)
   %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv2i1(<vscale x 2 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.any.nxv16i1(<vscale x 16 x i1> %1, <vscale x 16 x i1> %2)
   ret i1 %out
 }
 
+; Rewrite PTEST_ANY(X=OP(PG,...), X) -> PTEST_ANY(PG, X)).
+define i1 @ptest_any_brkb_z(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a) {
+; CHECK-LABEL: @ptest_any_brkb_z(
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call <vscale x 16 x i1> @llvm.aarch64.sve.brkb.z.nxv16i1(<vscale x 16 x i1> [[PG:%.*]], <vscale x 16 x i1> [[A:%.*]])
+; CHECK-NEXT:    [[OUT:%.*]] = call i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1> [[PG]], <vscale x 16 x i1> [[TMP1]])
+; CHECK-NEXT:    ret i1 [[OUT]]
+;
+  %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.brkb.z.nxv16i1(<vscale x 16 x i1> %pg, <vscale x 16 x i1> %a)
+  %out = call i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1> %1, <vscale x 16 x i1> %1)
+  ret i1 %out
+}
+
+define i1 @ptest_any_rdffr_z(<vscale x 16 x i1> %pg) {
+; CHECK-LABEL: @ptest_any_rdffr_z(
+; CHECK-NEXT:    [[TMP1:%.*]] = tail call <vscale x 16 x i1> @llvm.aarch64.sve.rdffr.z(<vscale x 16 x i1> [[PG:%.*]])
+; CHECK-NEXT:    [[OUT:%.*]] = call i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1> [[PG]], <vscale x 16 x i1> [[TMP1]])
+; CHECK-NEXT:    ret i1 [[OUT]]
+;
+  %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.rdffr.z(<vscale x 16 x i1> %pg)
+  %out = call i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1> %1, <vscale x 16 x i1> %1)
+  ret i1 %out
+}
+
 define i1 @ptest_first(<vscale x 4 x i1> %a) #0 {
 ; CHECK-LABEL: @ptest_first(
 ; CHECK-NEXT:    [[MASK:%.*]] = tail call <vscale x 4 x i1> @llvm.aarch64.sve.ptrue.nxv4i1(i32 0)
 ; CHECK-NEXT:    [[OUT:%.*]] = call i1 @llvm.aarch64.sve.ptest.first.nxv4i1(<vscale x 4 x i1> [[MASK]], <vscale x 4 x i1> [[A:%.*]])
 ; CHECK-NEXT:    ret i1 [[OUT]]
 ;
   %mask = tail call <vscale x 4 x i1> @llvm.aarch64.sve.ptrue.nxv4i1(i32 0)
   %1 = tail call <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1> %mask)

 declare i1 @llvm.aarch64.sve.ptest.any.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>)
 declare i1 @llvm.aarch64.sve.ptest.first.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>)
 declare i1 @llvm.aarch64.sve.ptest.last.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>)
 
 declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv8i1(<vscale x 8 x i1>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv4i1(<vscale x 4 x i1>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.convert.to.svbool.nxv2i1(<vscale x 2 x i1>)
 
+declare <vscale x 16 x i1> @llvm.aarch64.sve.brkb.z.nxv16i1(<vscale x 16 x i1>, <vscale x 16 x i1>)
+declare <vscale x 16 x i1> @llvm.aarch64.sve.rdffr.z(<vscale x 16 x i1>)
+
 attributes #0 = { "target-features"="+sve" }