[AArch64][SVE] Fold predicate into compare

[c-rhodes]

Summary

Codegen of added testcase before this patch:

ptrue   p0.s
cmpgt   p1.s, p0/z, z0.s, z1.s
cmpge   p2.s, p0/z, z2.s, z1.s
and     p0.b, p0/z, p1.b, p2.b
ret

Patterns originally authored by Will Lovett.

[david-arm]

Nice improvement @c-rhodes! Just a couple of suggestions ...

llvm/lib/Target/AArch64/SVEInstrFormats.td

4644

I do wonder if it's also worth creating patterns for ands matching the first operand too, i.e.

def : Pat<(predvt (and (AArch64setcc_z (predvt (AArch64ptrue 31)), intvt:$Op2, intvt:$Op3, cc))),
          (cmp $Pg, $Op2, $Op3), predvt:$Pg>;

I'm thinking about the case where only one input to the and instruction is a setcc, i.e. see my comment about writing other test cases below.

llvm/test/CodeGen/AArch64/sve-intrinsics-int-compares.ll

967

Hi @c-rhodes, given your patterns also work for any general predicate for the first operand of and can you also add a testcase that demonstrates this? For example, something like:

define <vscale x 4 x i1> @predicated_icmp_unknown_lhs(<vscale x 4 x i1> %a, <vscale x 4 x i32> %b, <vscale x 4 x i32> %c) {
  %icmp = icmp sle <vscale x 4 x i32> %b, %c
  %and = and <vscale x 4 x i1> %a, %icmp
  ret <vscale x 4 x i1> %and
}

I'd expect to see this code being produced:

; CHECK:       // %bb.0:
; CHECK-NEXT:    cmpge p0.s, p0/z, z2.s, z1.s
; CHECK-NEXT:    ret

[c-rhodes]

Add test Dave suggested

[c-rhodes]

llvm/lib/Target/AArch64/SVEInstrFormats.td

4644

I do wonder if it's also worth creating patterns for ands matching the first operand too, i.e.

def : Pat<(predvt (and (AArch64setcc_z (predvt (AArch64ptrue 31)), intvt:$Op2, intvt:$Op3, cc))),
          (cmp $Pg, $Op2, $Op3), predvt:$Pg>;

I'm thinking about the case where only one input to the and instruction is a setcc, i.e. see my comment about writing other test cases below.

I've added the testcase you suggested and the codegen is as expected without this pattern, so I think we're good unless I'm missing something?

[c-rhodes]

Add further test for RHS

[david-arm]

LGTM!. Thanks for adding the tests @c-rhodes. :)

llvm/lib/Target/AArch64/SVEInstrFormats.td

 }
 
 multiclass SVE_SETCC_Pat<CondCode cc, CondCode invcc, ValueType predvt,
                          ValueType intvt, Instruction cmp> {
   def : Pat<(predvt (AArch64setcc_z predvt:$Op1, intvt:$Op2, intvt:$Op3, cc)),
             (cmp $Op1, $Op2, $Op3)>;
   def : Pat<(predvt (AArch64setcc_z predvt:$Op1, intvt:$Op2, intvt:$Op3, invcc)),
             (cmp $Op1, $Op3, $Op2)>;
+  def : Pat<(predvt (and predvt:$Pg, (AArch64setcc_z (predvt (AArch64ptrue 31)), intvt:$Op2, intvt:$Op3, cc))),
+            (cmp $Pg, $Op2, $Op3)>;
+  def : Pat<(predvt (and predvt:$Pg, (AArch64setcc_z (predvt (AArch64ptrue 31)), intvt:$Op2, intvt:$Op3, invcc))),
+            (cmp $Pg, $Op3, $Op2)>;
 }
 
 multiclass SVE_SETCC_Pat_With_Zero<CondCode cc, CondCode invcc, ValueType predvt,
                                    ValueType intvt, Instruction cmp> {
   def : Pat<(predvt (AArch64setcc_z predvt:$Op1, intvt:$Op2, (SVEDup0), cc)),
             (cmp $Op1, $Op2)>;
   def : Pat<(predvt (AArch64setcc_z predvt:$Op1, (SVEDup0), intvt:$Op2, invcc)),
             (cmp $Op1, $Op2)>;

llvm/test/CodeGen/AArch64/sve-intrinsics-int-compares.ll

 ; CHECK-NEXT: ret
   %splat = call <vscale x 2 x i64> @llvm.aarch64.sve.dup.x.nxv2i64(i64 %b)
   %out = call <vscale x 4 x i1> @llvm.aarch64.sve.cmpls.wide.nxv4i32(<vscale x 4 x i1> %pg,
                                                                      <vscale x 4 x i32> %a,
                                                                      <vscale x 2 x i64> %splat)
   ret <vscale x 4 x i1> %out
 }
 
+; Verify general predicate is folded into the compare
+define <vscale x 4 x i1> @predicated_icmp(<vscale x 4 x i32> %a, <vscale x 4 x i32> %b, <vscale x 4 x i32> %c) {
+; CHECK-LABEL: predicated_icmp:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p0.s
+; CHECK-NEXT:    cmpgt p0.s, p0/z, z0.s, z1.s
+; CHECK-NEXT:    cmpge p0.s, p0/z, z2.s, z1.s
+; CHECK-NEXT:    ret
+  %icmp1 = icmp sgt <vscale x 4 x i32> %a, %b
+  %icmp2 = icmp sle <vscale x 4 x i32> %b, %c
+  %and = and <vscale x 4 x i1> %icmp1, %icmp2
+  ret <vscale x 4 x i1> %and
+}
+
+define <vscale x 4 x i1> @predicated_icmp_unknown_lhs(<vscale x 4 x i1> %a, <vscale x 4 x i32> %b, <vscale x 4 x i32> %c) {
+; CHECK-LABEL: predicated_icmp_unknown_lhs:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmpge p0.s, p0/z, z1.s, z0.s
+; CHECK-NEXT:    ret
+  %icmp = icmp sle <vscale x 4 x i32> %b, %c
+  %and = and <vscale x 4 x i1> %a, %icmp
+  ret <vscale x 4 x i1> %and
+}
+
+define <vscale x 4 x i1> @predicated_icmp_unknown_rhs(<vscale x 4 x i1> %a, <vscale x 4 x i32> %b, <vscale x 4 x i32> %c) {
+; CHECK-LABEL: predicated_icmp_unknown_rhs:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    cmpge p0.s, p0/z, z1.s, z0.s
+; CHECK-NEXT:    ret
+  %icmp = icmp sle <vscale x 4 x i32> %b, %c
+  %and = and <vscale x 4 x i1> %icmp, %a
+  ret <vscale x 4 x i1> %and
+}
 
 declare <vscale x 16 x i1> @llvm.aarch64.sve.cmpeq.nxv16i8(<vscale x 16 x i1>, <vscale x 16 x i8>, <vscale x 16 x i8>)
 declare <vscale x 8 x i1> @llvm.aarch64.sve.cmpeq.nxv8i16(<vscale x 8 x i1>, <vscale x 8 x i16>, <vscale x 8 x i16>)
 declare <vscale x 4 x i1> @llvm.aarch64.sve.cmpeq.nxv4i32(<vscale x 4 x i1>, <vscale x 4 x i32>, <vscale x 4 x i32>)
 declare <vscale x 2 x i1> @llvm.aarch64.sve.cmpeq.nxv2i64(<vscale x 2 x i1>, <vscale x 2 x i64>, <vscale x 2 x i64>)
 declare <vscale x 16 x i1> @llvm.aarch64.sve.cmpeq.wide.nxv16i8(<vscale x 16 x i1>, <vscale x 16 x i8>, <vscale x 2 x i64>)
 declare <vscale x 8 x i1> @llvm.aarch64.sve.cmpeq.wide.nxv8i16(<vscale x 8 x i1>, <vscale x 8 x i16>, <vscale x 2 x i64>)
 declare <vscale x 4 x i1> @llvm.aarch64.sve.cmpeq.wide.nxv4i32(<vscale x 4 x i1>, <vscale x 4 x i32>, <vscale x 2 x i64>)