[SVE][CodeGen] CTLZ, CTTZ & CTPOP operations (predicates)

[kmclaughlin]

Summary

Canonicalise the following operations in getNode() for predicate types:

• CTLZ(Pred) -> bitwise_NOT(Pred)
• CTTZ(Pred) -> bitwise_NOT(Pred)
• CTPOP(Pred) -> Pred

[david-arm]

LGTM! The canonicalisation looks sensible to me.

llvm/test/CodeGen/AArch64/sve-bit-counting-pred.ll

9

nit: Perhaps worth changing this comment to CTPOP to match the intrinsic name, similar to the CTLZ and CTTZ comments below?

llvm/lib/CodeGen/SelectionDAG/SelectionDAG.cpp

case ISD::FNEG:
     break;
   case ISD::FABS:
     if (OpOpcode == ISD::FNEG)  // abs(-X) -> abs(X)
       return getNode(ISD::FABS, DL, VT, Operand.getOperand(0));
     break;
   case ISD::VSCALE:
     assert(VT == Operand.getValueType() && "Unexpected VT!");
     break;
+  case ISD::CTPOP:
+    if (Operand.getValueType().getScalarType() == MVT::i1)
+      return Operand;
+    break;
+  case ISD::CTLZ:
+  case ISD::CTTZ:
+    if (Operand.getValueType().getScalarType() == MVT::i1)
+      return getNOT(DL, Operand, Operand.getValueType());
+    break;
   case ISD::VECREDUCE_SMIN:
   case ISD::VECREDUCE_UMAX:
     if (Operand.getValueType().getScalarType() == MVT::i1)
       return getNode(ISD::VECREDUCE_OR, DL, VT, Operand);
     break;
   case ISD::VECREDUCE_SMAX:
   case ISD::VECREDUCE_UMIN:
     if (Operand.getValueType().getScalarType() == MVT::i1)

llvm/test/CodeGen/AArch64/sve-bit-counting-pred.ll

+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=aarch64-linux-gnu -mattr=+sve < %s 2>%t | FileCheck %s
+; RUN: FileCheck --check-prefix=WARN --allow-empty %s <%t
+
+; If this check fails please read test/CodeGen/AArch64/README for instructions on how to resolve it.
+; WARN-NOT: warning
+
+;
+; CNT
+;
+
+define <vscale x 16 x i1> @ctpop_nxv16i1(<vscale x 16 x i1> %a) {
+; CHECK-LABEL: ctpop_nxv16i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ret
+  %res = call <vscale x 16 x i1> @llvm.ctpop.nxv16i1(<vscale x 16 x i1> %a)
+  ret <vscale x 16 x i1> %res
+}
+
+define <vscale x 8 x i1> @ctpop_nxv8i1(<vscale x 8 x i1> %a) {
+; CHECK-LABEL: ctpop_nxv8i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ret
+  %res = call <vscale x 8 x i1> @llvm.ctpop.nxv8i1(<vscale x 8 x i1> %a)
+  ret <vscale x 8 x i1> %res
+}
+
+define <vscale x 4 x i1> @ctpop_nxv4i1(<vscale x 4 x i1> %a) {
+; CHECK-LABEL: ctpop_nxv4i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ret
+  %res = call <vscale x 4 x i1> @llvm.ctpop.nxv4i1(<vscale x 4 x i1> %a)
+  ret <vscale x 4 x i1> %res
+}
+
+define <vscale x 2 x i1> @ctpop_nxv2i1(<vscale x 2 x i1> %a) {
+; CHECK-LABEL: ctpop_nxv2i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i1> @llvm.ctpop.nxv2i1(<vscale x 2 x i1> %a)
+  ret <vscale x 2 x i1> %res
+}
+
+; CTLZ
+
+define <vscale x 16 x i1> @ctlz_nxv16i1(<vscale x 16 x i1> %a) {
+; CHECK-LABEL: ctlz_nxv16i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p1.b
+; CHECK-NEXT:    not p0.b, p1/z, p0.b
+; CHECK-NEXT:    ret
+  %res = call <vscale x 16 x i1> @llvm.ctlz.nxv16i1(<vscale x 16 x i1> %a)
+  ret <vscale x 16 x i1> %res
+}
+
+define <vscale x 8 x i1> @ctlz_nxv8i1(<vscale x 8 x i1> %a) {
+; CHECK-LABEL: ctlz_nxv8i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p1.h
+; CHECK-NEXT:    not p0.b, p1/z, p0.b
+; CHECK-NEXT:    ret
+  %res = call <vscale x 8 x i1> @llvm.ctlz.nxv8i1(<vscale x 8 x i1> %a)
+  ret <vscale x 8 x i1> %res
+}
+
+define <vscale x 4 x i1> @ctlz_nxv4i1(<vscale x 4 x i1> %a) {
+; CHECK-LABEL: ctlz_nxv4i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    not p0.b, p1/z, p0.b
+; CHECK-NEXT:    ret
+  %res = call <vscale x 4 x i1> @llvm.ctlz.nxv4i1(<vscale x 4 x i1> %a)
+  ret <vscale x 4 x i1> %res
+}
+
+define <vscale x 2 x i1> @ctlz_nxv2i1(<vscale x 2 x i1> %a) {
+; CHECK-LABEL: ctlz_nxv2i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    not p0.b, p1/z, p0.b
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i1> @llvm.ctlz.nxv2i1(<vscale x 2 x i1> %a)
+  ret <vscale x 2 x i1> %res
+}
+
+; CTTZ
+
+define <vscale x 16 x i1> @cttz_nxv16i1(<vscale x 16 x i1> %a) {
+; CHECK-LABEL: cttz_nxv16i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p1.b
+; CHECK-NEXT:    not p0.b, p1/z, p0.b
+; CHECK-NEXT:    ret
+  %res = call <vscale x 16 x i1> @llvm.cttz.nxv16i1(<vscale x 16 x i1> %a)
+  ret <vscale x 16 x i1> %res
+}
+
+define <vscale x 8 x i1> @cttz_nxv8i1(<vscale x 8 x i1> %a) {
+; CHECK-LABEL: cttz_nxv8i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p1.h
+; CHECK-NEXT:    not p0.b, p1/z, p0.b
+; CHECK-NEXT:    ret
+  %res = call <vscale x 8 x i1> @llvm.cttz.nxv8i1(<vscale x 8 x i1> %a)
+  ret <vscale x 8 x i1> %res
+}
+
+define <vscale x 4 x i1> @cttz_nxv4i1(<vscale x 4 x i1> %a) {
+; CHECK-LABEL: cttz_nxv4i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p1.s
+; CHECK-NEXT:    not p0.b, p1/z, p0.b
+; CHECK-NEXT:    ret
+  %res = call <vscale x 4 x i1> @llvm.cttz.nxv4i1(<vscale x 4 x i1> %a)
+  ret <vscale x 4 x i1> %res
+}
+
+define <vscale x 2 x i1> @cttz_nxv2i1(<vscale x 2 x i1> %a) {
+; CHECK-LABEL: cttz_nxv2i1:
+; CHECK:       // %bb.0:
+; CHECK-NEXT:    ptrue p1.d
+; CHECK-NEXT:    not p0.b, p1/z, p0.b
+; CHECK-NEXT:    ret
+  %res = call <vscale x 2 x i1> @llvm.cttz.nxv2i1(<vscale x 2 x i1> %a)
+  ret <vscale x 2 x i1> %res
+}
+
+declare <vscale x 16 x i1> @llvm.ctpop.nxv16i1(<vscale x 16 x i1>)
+declare <vscale x 8 x i1> @llvm.ctpop.nxv8i1(<vscale x 8 x i1>)
+declare <vscale x 4 x i1> @llvm.ctpop.nxv4i1(<vscale x 4 x i1>)
+declare <vscale x 2 x i1> @llvm.ctpop.nxv2i1(<vscale x 2 x i1>)
+
+declare <vscale x 16 x i1> @llvm.ctlz.nxv16i1(<vscale x 16 x i1>)
+declare <vscale x 8 x i1> @llvm.ctlz.nxv8i1(<vscale x 8 x i1>)
+declare <vscale x 4 x i1> @llvm.ctlz.nxv4i1(<vscale x 4 x i1>)
+declare <vscale x 2 x i1> @llvm.ctlz.nxv2i1(<vscale x 2 x i1>)
+
+declare <vscale x 16 x i1> @llvm.cttz.nxv16i1(<vscale x 16 x i1>)
+declare <vscale x 8 x i1> @llvm.cttz.nxv8i1(<vscale x 8 x i1>)
+declare <vscale x 4 x i1> @llvm.cttz.nxv4i1(<vscale x 4 x i1>)
+declare <vscale x 2 x i1> @llvm.cttz.nxv2i1(<vscale x 2 x i1>)