[AArch64][SVE] Fix bad PTEST(X, X) optimization

[c-rhodes]

Summary

AArch64InstrInfo::optimizePTestInstr attempts to remove a PTEST of a
predicate generating operation that identically sets flags (implictly).

When the mask is the same as the input predicate the PTEST is currently
removed. This is incorrect since the mask for the implicit PTEST
performed by the flag-setting instruction differs from the mask
specified to the explicit PTEST and could set different flags.

For example, consider

PG=<1, 1, x, x>
Z0=<1, 2, x, x>
Z1=<2, 1, x, x>

X=CMPLE(PG, Z0, Z1)
 =<0, 1, x, x>       NZCV=0xxx
PTEST(X, X),         NZCV=1xxx

where the first active flag (bit 'N' in NZCV) is set by the explicit
PTEST, but not by the implicit PTEST as part of the compare. Given the
PTEST mask and source are the same however, first is equivalent to any,
so the PTEST could be removed if the condition is changed. The same
applies to last active. It is safe to remove the PTEST for any active,
but this information isn't available in the current optimization.

This patch fixes the bad optimization, a later patch will implement the
optimization proposed above and fix the any active case.

[bsmith]

LGTM, modulo future changes.

llvm/lib/Target/AArch64/AArch64InstrInfo.cpp

if (isPTrueOpcode(MaskOpcode) && (PredIsPTestLike || PredIsWhileLike)) {
     uint64_t PredElementSize = getElementSizeForOpcode(PredOpcode);
 
     // Must be an all active predicate of matching element size.
     if ((PredElementSize != MaskElementSize) ||
         (Mask->getOperand(1).getImm() != 31))
       return false;
 
     // Fallthough to simply remove the PTEST.
-  } else if ((Mask == Pred) && (PredIsPTestLike || PredIsWhileLike)) {
-    // For PTEST(PG, PG), PTEST is redundant when PG is the result of an
-    // instruction that sets the flags as PTEST would.
-
-    // Fallthough to simply remove the PTEST.
   } else if (PredIsPTestLike) {
     // For PTEST(PG, PTEST_LIKE(PG, ...)), the PTEST is redundant since the
     // flags are set based on the same mask 'PG', but PTEST_LIKE must operate
     // on 8-bit predicates like the PTEST.  Otherwise, for instructions like
     // compare that also support 16/32/64-bit predicates, the implicit PTEST
     // performed by the compare could consider fewer lanes for these element
     // sizes.
     //

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

 ; 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:    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:    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.
 ;
 define i1 @cmp8_ptest_any_xx(<vscale x 16 x i1> %pg, <vscale x 16 x i8> %a, <vscale x 16 x i8> %b) {
 ; CHECK-LABEL: cmp8_ptest_any_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, 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.any.nxv16i1(<vscale x 16 x i1> %1, <vscale x 16 x i1> %1)
   ret i1 %2
 }
 
 ;
 ; 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:    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:    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
 }
 
 ;
 ; PTEST_ANY(X=CMP32(PG, A, B), X). PTEST is redundant.
 ;
 define i1 @cmp32_ptest_any_xx(<vscale x 16 x i1> %pg, <vscale x 4 x i32> %a, <vscale x 4 x i32> %b) {
 ; CHECK-LABEL: cmp32_ptest_any_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, 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.any.nxv4i1(<vscale x 4 x i1> %2, <vscale x 4 x i1> %2)
   ret i1 %3
 }
 

llvm/test/CodeGen/AArch64/sve-ptest-removal-whilegt.mir

liveins:
   - { reg: '$w1', virtual-reg: '%1' }
 frameInfo:
   maxCallFrameSize: 0
 body:             |
   bb.0.entry:
     liveins: $w0, $w1
 
     ; CHECK-LABEL: name: whilegt_b8_s32_ptest_with_matching_operands
-    ; CHECK-NOT: PTEST
+    ; CHECK: PTEST
     %1:gpr32 = COPY $w1
     %0:gpr32 = COPY $w0
     %2:ppr = WHILEGT_PWW_B %0, %1, implicit-def dead $nzcv
     PTEST_PP %2, killed %2, implicit-def $nzcv
     %3:gpr32 = COPY $wzr
     %4:gpr32 = CSINCWr %3, $wzr, 0, implicit $nzcv
     $w0 = COPY %4
     RET_ReallyLR implicit $w0
 
 ...