[AArch64] Add patterns to fold zext(cmpeq(x, splat(0)))

[c-rhodes]

Add patterns for all legal FP types.

[c-rhodes]

Remove FP patterns.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

451–452

Are you sure these produce the same result? It looks like the input pattern sets the result of inactive lanes to zero but CNOT_ZPmZ_B will set them to whatever is in $Op2. Perhaps the pattern should be restricted to all active predicates.

Is it possible to add these patterns to the PatFrags used to define AArch64cnot_mt?

[c-rhodes]

Restrict patterns to all active predicate.

[c-rhodes]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

451–452

Are you sure these produce the same result? It looks like the input pattern sets the result of inactive lanes to zero but CNOT_ZPmZ_B will set them to whatever is in $Op2. Perhaps the pattern should be restricted to all active predicates.

Good spot! I've updated the patterns to only match all active predicate.

Is it possible to add these patterns to the PatFrags used to define AArch64cnot_mt?

I can't figure out how to express the all active predicate restriction in a PatFrag, happy to make the change if you know how.

[paulwalker-arm]

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

451–452

Fair enough. With the all active restriction they're not really the same thing anymore so my PatFrag idea goes out of the window.

459

Here you should be able to use (SVEAllActive):$Pg that will catch more cases whilst also enabling you to use $Pg in place of (PTRUE_# 31).

[c-rhodes]

Use SVEAllActive

llvm/lib/Target/AArch64/AArch64SVEInstrInfo.td

let Predicates = [HasSVEorSME] in {
   defm SXTH_ZPmZ : sve_int_un_pred_arit_0_w<0b010, "sxth", AArch64sxt_mt>;
   defm UXTH_ZPmZ : sve_int_un_pred_arit_0_w<0b011, "uxth", AArch64uxt_mt>;
   defm SXTW_ZPmZ : sve_int_un_pred_arit_0_d<0b100, "sxtw", AArch64sxt_mt>;
   defm UXTW_ZPmZ : sve_int_un_pred_arit_0_d<0b101, "uxtw", AArch64uxt_mt>;
   defm ABS_ZPmZ  : sve_int_un_pred_arit_0<  0b110, "abs",  AArch64abs_mt>;
   defm NEG_ZPmZ  : sve_int_un_pred_arit_0<  0b111, "neg",  AArch64neg_mt>;
 
   defm CLS_ZPmZ  : sve_int_un_pred_arit_1<   0b000, "cls",  AArch64cls_mt>;
   defm CLZ_ZPmZ  : sve_int_un_pred_arit_1<   0b001, "clz",  AArch64clz_mt>;
   defm CNT_ZPmZ  : sve_int_un_pred_arit_1<   0b010, "cnt",  AArch64cnt_mt>;
   defm CNOT_ZPmZ : sve_int_un_pred_arit_1<   0b011, "cnot", AArch64cnot_mt>;
   defm NOT_ZPmZ  : sve_int_un_pred_arit_1<   0b110, "not",  AArch64not_mt>;
   defm FABS_ZPmZ : sve_int_un_pred_arit_1_fp<0b100, "fabs", AArch64fabs_mt>;
   defm FNEG_ZPmZ : sve_int_un_pred_arit_1_fp<0b101, "fneg", AArch64fneg_mt>;
 
+  // zext(cmpeq(x, splat(0))) -> cnot(x)
+  def : Pat<(nxv16i8 (zext (nxv16i1 (AArch64setcc_z (nxv16i1 (SVEAllActive):$Pg), nxv16i8:$Op2, (SVEDup0), SETEQ)))),
+            (CNOT_ZPmZ_B $Op2, $Pg, $Op2)>;
+  def : Pat<(nxv8i16 (zext (nxv8i1 (AArch64setcc_z (nxv8i1 (SVEAllActive):$Pg), nxv8i16:$Op2, (SVEDup0), SETEQ)))),
+            (CNOT_ZPmZ_H $Op2, $Pg, $Op2)>;
+  def : Pat<(nxv4i32 (zext (nxv4i1 (AArch64setcc_z (nxv4i1 (SVEAllActive):$Pg), nxv4i32:$Op2, (SVEDup0), SETEQ)))),
+            (CNOT_ZPmZ_S $Op2, $Pg, $Op2)>;
+  def : Pat<(nxv2i64 (zext (nxv2i1 (AArch64setcc_z (nxv2i1 (SVEAllActive):$Pg), nxv2i64:$Op2, (SVEDup0), SETEQ)))),
+            (CNOT_ZPmZ_D $Op2, $Pg, $Op2)>;
+
   defm SMAX_ZPmZ : sve_int_bin_pred_arit_1<0b000, "smax", "SMAX_ZPZZ", int_aarch64_sve_smax, DestructiveBinaryComm>;
   defm UMAX_ZPmZ : sve_int_bin_pred_arit_1<0b001, "umax", "UMAX_ZPZZ", int_aarch64_sve_umax, DestructiveBinaryComm>;
   defm SMIN_ZPmZ : sve_int_bin_pred_arit_1<0b010, "smin", "SMIN_ZPZZ", int_aarch64_sve_smin, DestructiveBinaryComm>;
   defm UMIN_ZPmZ : sve_int_bin_pred_arit_1<0b011, "umin", "UMIN_ZPZZ", int_aarch64_sve_umin, DestructiveBinaryComm>;
   defm SABD_ZPmZ : sve_int_bin_pred_arit_1<0b100, "sabd", "SABD_ZPZZ", int_aarch64_sve_sabd, DestructiveBinaryComm>;
   defm UABD_ZPmZ : sve_int_bin_pred_arit_1<0b101, "uabd", "UABD_ZPZZ", int_aarch64_sve_uabd, DestructiveBinaryComm>;
 
   defm SMAX_ZPZZ : sve_int_bin_pred_bhsd<AArch64smax_p>;

llvm/test/CodeGen/AArch64/sve-cmp-folds.ll

; CHECK-NEXT:    ret
   %not = xor <vscale x 4 x i1> %icmp, %ones
   ret <vscale x 4 x i1> %not
 }
 
 define <vscale x 16 x i8> @icmp_cnot_nxv16i8(<vscale x 16 x i8> %a) {
 ; CHECK-LABEL: icmp_cnot_nxv16i8:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.b
-; CHECK-NEXT:    cmpeq p0.b, p0/z, z0.b, #0
-; CHECK-NEXT:    mov z0.b, p0/z, #1 // =0x1
+; CHECK-NEXT:    cnot z0.b, p0/m, z0.b
 ; CHECK-NEXT:    ret
   %mask = icmp eq <vscale x 16 x i8> %a, zeroinitializer
   %zext = zext <vscale x 16 x i1> %mask to <vscale x 16 x i8>
   ret <vscale x 16 x i8> %zext
 }
 
 define <vscale x 8 x i16> @icmp_cnot_nxv8i16(<vscale x 8 x i16> %a) {
 ; CHECK-LABEL: icmp_cnot_nxv8i16:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.h
-; CHECK-NEXT:    cmpeq p0.h, p0/z, z0.h, #0
-; CHECK-NEXT:    mov z0.h, p0/z, #1 // =0x1
+; CHECK-NEXT:    cnot z0.h, p0/m, z0.h
 ; CHECK-NEXT:    ret
   %mask = icmp eq <vscale x 8 x i16> %a, zeroinitializer
   %zext = zext <vscale x 8 x i1> %mask to <vscale x 8 x i16>
   ret <vscale x 8 x i16> %zext
 }
 
 define <vscale x 4 x i32> @icmp_cnot_nxv4i32(<vscale x 4 x i32> %a) {
 ; CHECK-LABEL: icmp_cnot_nxv4i32:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.s
-; CHECK-NEXT:    cmpeq p0.s, p0/z, z0.s, #0
-; CHECK-NEXT:    mov z0.s, p0/z, #1 // =0x1
+; CHECK-NEXT:    cnot z0.s, p0/m, z0.s
 ; CHECK-NEXT:    ret
   %mask = icmp eq <vscale x 4 x i32> %a, zeroinitializer
   %zext = zext <vscale x 4 x i1> %mask to <vscale x 4 x i32>
   ret <vscale x 4 x i32> %zext
 }
 
 define <vscale x 2 x i64> @icmp_cnot_nxv2i64(<vscale x 2 x i64> %a) {
 ; CHECK-LABEL: icmp_cnot_nxv2i64:
 ; CHECK:       // %bb.0:
 ; CHECK-NEXT:    ptrue p0.d
-; CHECK-NEXT:    cmpeq p0.d, p0/z, z0.d, #0
-; CHECK-NEXT:    mov z0.d, p0/z, #1 // =0x1
+; CHECK-NEXT:    cnot z0.d, p0/m, z0.d
 ; CHECK-NEXT:    ret
   %mask = icmp eq <vscale x 2 x i64> %a, zeroinitializer
   %zext = zext <vscale x 2 x i1> %mask to <vscale x 2 x i64>
   ret <vscale x 2 x i64> %zext
 }
 
 define i1 @foo_first(<vscale x 4 x float> %a, <vscale x 4 x float> %b) {
 ; CHECK-LABEL: foo_first: