[AArch64] Improve non-SVE popcount for 32bit and 64 bit using udot

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

@@ -9814,6 +9814,26 @@ SDValue AArch64TargetLowering::LowerCTPOP_PARITY(SDValue Op,
   Val = DAG.getBitcast(VT8Bit, Val);
   Val = DAG.getNode(ISD::CTPOP, DL, VT8Bit, Val);
 
+  if (Subtarget->hasDotProd() && VT.getScalarSizeInBits() != 16 &&
+      VT.getVectorNumElements() >= 2) {
+    EVT DT = VT == MVT::v2i64 ? MVT::v4i32 : VT;
+    SDValue Zeros = DAG.getConstant(0, DL, DT);
+    SDValue Ones = DAG.getConstant(1, DL, VT8Bit);
+
+    if (VT == MVT::v2i64) {
+      Val = DAG.getNode(AArch64ISD::UDOT, DL, DT, Zeros, Ones, Val);
+      Val = DAG.getNode(AArch64ISD::UADDLP, DL, VT, Val);
+    } else if (VT == MVT::v2i32) {
+      Val = DAG.getNode(AArch64ISD::UDOT, DL, DT, Zeros, Ones, Val);
+    } else if (VT == MVT::v4i32) {
+      Val = DAG.getNode(AArch64ISD::UDOT, DL, DT, Zeros, Ones, Val);
+    } else {
+      llvm_unreachable("Unexpected type for custom ctpop lowering");
+    }
+
+    return Val;
+  }
+
   // Widen v8i8/v16i8 CTPOP result to VT by repeatedly widening pairwise adds.
   unsigned EltSize = 8;
   unsigned NumElts = VT.is64BitVector() ? 8 : 16;

llvm/test/CodeGen/AArch64/popcount.ll

@@ -1,15 +1,28 @@
 ; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
-; RUN: llc < %s -O0 -mtriple=aarch64-unknown-unknown | FileCheck %s
+; RUN: llc < %s -O0 -mtriple=aarch64-unknown-unknown | FileCheck %s --check-prefix=CHECKO0
+; RUN: llc < %s -mtriple=aarch64-unknown-unknown -mattr=+neon | FileCheck %s --check-prefixes=CHECK,NEON
+; RUN: llc < %s -mtriple=aarch64-unknown-unknown -mattr=+neon,+dotprod | FileCheck %s --check-prefixes=CHECK,DOT
+; RUN: llc < %s -mtriple=aarch64-unknown-unknown -mattr=+sve | FileCheck %s --check-prefixes=CHECK,SVE
 
 ; Function Attrs: nobuiltin nounwind readonly
 define i8 @popcount128(ptr nocapture nonnull readonly %0) {
+; CHECKO0-LABEL: popcount128:
+; CHECKO0:       // %bb.0: // %Entry
+; CHECKO0-NEXT:    ldr q0, [x0]
+; CHECKO0-NEXT:    cnt v0.16b, v0.16b
+; CHECKO0-NEXT:    uaddlv h0, v0.16b
+; CHECKO0-NEXT:    // kill: def $q0 killed $h0
+; CHECKO0-NEXT:    // kill: def $s0 killed $s0 killed $q0
+; CHECKO0-NEXT:    fmov w0, s0
+; CHECKO0-NEXT:    ret
+;
 ; CHECK-LABEL: popcount128:
 ; CHECK:       // %bb.0: // %Entry
-; CHECK-NEXT:    ldr q0, [x0]
+; CHECK-NEXT:    ldr d0, [x0]
+; CHECK-NEXT:    add x8, x0, #8
+; CHECK-NEXT:    ld1 { v0.d }[1], [x8]
 ; CHECK-NEXT:    cnt v0.16b, v0.16b
 ; CHECK-NEXT:    uaddlv h0, v0.16b
-; CHECK-NEXT:    // kill: def $q0 killed $h0
-; CHECK-NEXT:    // kill: def $s0 killed $s0 killed $q0
 ; CHECK-NEXT:    fmov w0, s0
 ; CHECK-NEXT:    ret
 Entry:
@@ -24,37 +37,54 @@ declare i128 @llvm.ctpop.i128(i128)
 
 ; Function Attrs: nobuiltin nounwind readonly
 define i16 @popcount256(ptr nocapture nonnull readonly %0) {
+; CHECKO0-LABEL: popcount256:
+; CHECKO0:       // %bb.0: // %Entry
+; CHECKO0-NEXT:    ldr x11, [x0]
+; CHECKO0-NEXT:    ldr x10, [x0, #8]
+; CHECKO0-NEXT:    ldr x9, [x0, #16]
+; CHECKO0-NEXT:    ldr x8, [x0, #24]
+; CHECKO0-NEXT:    // implicit-def: $q1
+; CHECKO0-NEXT:    mov v1.d[0], x11
+; CHECKO0-NEXT:    mov v1.d[1], x10
+; CHECKO0-NEXT:    // implicit-def: $q0
+; CHECKO0-NEXT:    mov v0.d[0], x9
+; CHECKO0-NEXT:    mov v0.d[1], x8
+; CHECKO0-NEXT:    cnt v1.16b, v1.16b
+; CHECKO0-NEXT:    uaddlv h1, v1.16b
+; CHECKO0-NEXT:    // kill: def $q1 killed $h1
+; CHECKO0-NEXT:    // kill: def $s1 killed $s1 killed $q1
+; CHECKO0-NEXT:    fmov w0, s1
+; CHECKO0-NEXT:    mov w10, wzr
+; CHECKO0-NEXT:    mov w9, w0
+; CHECKO0-NEXT:    mov w8, w10
+; CHECKO0-NEXT:    bfi x9, x8, #32, #32
+; CHECKO0-NEXT:    cnt v0.16b, v0.16b
+; CHECKO0-NEXT:    uaddlv h0, v0.16b
+; CHECKO0-NEXT:    // kill: def $q0 killed $h0
+; CHECKO0-NEXT:    // kill: def $s0 killed $s0 killed $q0
+; CHECKO0-NEXT:    fmov w0, s0
+; CHECKO0-NEXT:    mov w8, w0
+; CHECKO0-NEXT:    // kill: def $x10 killed $w10
+; CHECKO0-NEXT:    bfi x8, x10, #32, #32
+; CHECKO0-NEXT:    adds x8, x8, x9
+; CHECKO0-NEXT:    mov w0, w8
+; CHECKO0-NEXT:    ret
+;
 ; CHECK-LABEL: popcount256:
 ; CHECK:       // %bb.0: // %Entry
-; CHECK-NEXT:    ldr x11, [x0]
-; CHECK-NEXT:    ldr x10, [x0, #8]
-; CHECK-NEXT:    ldr x9, [x0, #16]
-; CHECK-NEXT:    ldr x8, [x0, #24]
-; CHECK-NEXT:    // implicit-def: $q1
-; CHECK-NEXT:    mov v1.d[0], x11
-; CHECK-NEXT:    mov v1.d[1], x10
-; CHECK-NEXT:    // implicit-def: $q0
-; CHECK-NEXT:    mov v0.d[0], x9
-; CHECK-NEXT:    mov v0.d[1], x8
-; CHECK-NEXT:    cnt v1.16b, v1.16b
-; CHECK-NEXT:    uaddlv h1, v1.16b
-; CHECK-NEXT:    // kill: def $q1 killed $h1
-; CHECK-NEXT:    // kill: def $s1 killed $s1 killed $q1
-; CHECK-NEXT:    fmov w0, s1
-; CHECK-NEXT:    mov w10, wzr
-; CHECK-NEXT:    mov w9, w0
-; CHECK-NEXT:    mov w8, w10
-; CHECK-NEXT:    bfi x9, x8, #32, #32
+; CHECK-NEXT:    ldr d0, [x0, #16]
+; CHECK-NEXT:    ldr d1, [x0]
+; CHECK-NEXT:    add x8, x0, #8
+; CHECK-NEXT:    add x9, x0, #24
+; CHECK-NEXT:    ld1 { v0.d }[1], [x9]
+; CHECK-NEXT:    ld1 { v1.d }[1], [x8]
 ; CHECK-NEXT:    cnt v0.16b, v0.16b
+; CHECK-NEXT:    cnt v1.16b, v1.16b
 ; CHECK-NEXT:    uaddlv h0, v0.16b
-; CHECK-NEXT:    // kill: def $q0 killed $h0
-; CHECK-NEXT:    // kill: def $s0 killed $s0 killed $q0
-; CHECK-NEXT:    fmov w0, s0
-; CHECK-NEXT:    mov w8, w0
-; CHECK-NEXT:    // kill: def $x10 killed $w10
-; CHECK-NEXT:    bfi x8, x10, #32, #32
-; CHECK-NEXT:    adds x8, x8, x9
-; CHECK-NEXT:    mov w0, w8
+; CHECK-NEXT:    uaddlv h1, v1.16b
+; CHECK-NEXT:    fmov w8, s0
+; CHECK-NEXT:    fmov w9, s1
+; CHECK-NEXT:    add w0, w9, w8
 ; CHECK-NEXT:    ret
 Entry:
   %1 = load i256, ptr %0, align 16
@@ -67,25 +97,215 @@ Entry:
 declare i256 @llvm.ctpop.i256(i256)
 
 define <1 x i128> @popcount1x128(<1 x i128> %0) {
+; CHECKO0-LABEL: popcount1x128:
+; CHECKO0:       // %bb.0: // %Entry
+; CHECKO0-NEXT:    // implicit-def: $q0
+; CHECKO0-NEXT:    mov v0.d[0], x0
+; CHECKO0-NEXT:    mov v0.d[1], x1
+; CHECKO0-NEXT:    cnt v0.16b, v0.16b
+; CHECKO0-NEXT:    uaddlv h0, v0.16b
+; CHECKO0-NEXT:    // kill: def $q0 killed $h0
+; CHECKO0-NEXT:    mov x1, xzr
+; CHECKO0-NEXT:    // kill: def $s0 killed $s0 killed $q0
+; CHECKO0-NEXT:    fmov w0, s0
+; CHECKO0-NEXT:    mov w8, wzr
+; CHECKO0-NEXT:    // kill: def $x0 killed $w0
+; CHECKO0-NEXT:    // kill: def $x8 killed $w8
+; CHECKO0-NEXT:    bfi x0, x8, #32, #32
+; CHECKO0-NEXT:    ret
+;
 ; CHECK-LABEL: popcount1x128:
 ; CHECK:       // %bb.0: // %Entry
-; CHECK-NEXT:    // implicit-def: $q0
-; CHECK-NEXT:    mov v0.d[0], x0
-; CHECK-NEXT:    mov v0.d[1], x1
-; CHECK-NEXT:    cnt v0.16b, v0.16b
-; CHECK-NEXT:    uaddlv h0, v0.16b
-; CHECK-NEXT:    // kill: def $q0 killed $h0
-; CHECK-NEXT:    mov x1, xzr
-; CHECK-NEXT:    // kill: def $s0 killed $s0 killed $q0
-; CHECK-NEXT:    fmov w0, s0
-; CHECK-NEXT:    mov w8, wzr
-; CHECK-NEXT:    // kill: def $x0 killed $w0
-; CHECK-NEXT:    // kill: def $x8 killed $w8
-; CHECK-NEXT:    bfi x0, x8, #32, #32
+; CHECK-NEXT:    fmov d1, x0
+; CHECK-NEXT:    movi v0.2d, #0000000000000000
+; CHECK-NEXT:    mov v1.d[1], x1
+; CHECK-NEXT:    cnt v1.16b, v1.16b
+; CHECK-NEXT:    uaddlv h1, v1.16b
+; CHECK-NEXT:    mov v0.s[0], v1.s[0]
+; CHECK-NEXT:    mov x1, v0.d[1]
+; CHECK-NEXT:    fmov x0, d0
 ; CHECK-NEXT:    ret
 Entry:
-  %1 = tail call <1 x i128> @llvm.ctpop.v1.i128(<1 x i128> %0)
+  %1 = tail call <1 x i128> @llvm.ctpop.v1i128(<1 x i128> %0)
   ret <1 x i128> %1
 }
 
-declare <1 x i128> @llvm.ctpop.v1.i128(<1 x i128>)
+declare <1 x i128> @llvm.ctpop.v1i128(<1 x i128>)
+
+define <2 x i64> @popcount2x64(<2 x i64> %0) {
+; CHECKO0-LABEL: popcount2x64:
+; CHECKO0:       // %bb.0: // %Entry
+; CHECKO0-NEXT:    cnt v0.16b, v0.16b
+; CHECKO0-NEXT:    uaddlp v0.8h, v0.16b
+; CHECKO0-NEXT:    uaddlp v0.4s, v0.8h
+; CHECKO0-NEXT:    uaddlp v0.2d, v0.4s
+; CHECKO0-NEXT:    ret
+;
+; NEON-LABEL: popcount2x64:
+; NEON:       // %bb.0: // %Entry
+; NEON-NEXT:    cnt v0.16b, v0.16b
+; NEON-NEXT:    uaddlp v0.8h, v0.16b
+; NEON-NEXT:    uaddlp v0.4s, v0.8h
+; NEON-NEXT:    uaddlp v0.2d, v0.4s
+; NEON-NEXT:    ret
+;
+; DOT-LABEL: popcount2x64:
+; DOT:       // %bb.0: // %Entry
+; DOT-NEXT:    movi v1.16b, #1
+; DOT-NEXT:    cnt v0.16b, v0.16b
+; DOT-NEXT:    movi v2.2d, #0000000000000000
+; DOT-NEXT:    udot v2.4s, v1.16b, v0.16b
+; DOT-NEXT:    uaddlp v0.2d, v2.4s
+; DOT-NEXT:    ret
+;
+; SVE-LABEL: popcount2x64:
+; SVE:       // %bb.0: // %Entry
+; SVE-NEXT:    cnt v0.16b, v0.16b
+; SVE-NEXT:    uaddlp v0.8h, v0.16b
+; SVE-NEXT:    uaddlp v0.4s, v0.8h
+; SVE-NEXT:    uaddlp v0.2d, v0.4s
+; SVE-NEXT:    ret
+Entry:
+  %1 = tail call <2 x i64> @llvm.ctpop.v2i64(<2 x i64> %0)
+  ret <2 x i64> %1
+}
+
+declare <2 x i64> @llvm.ctpop.v2i64(<2 x i64>)
+
+define <1 x i64> @popcount1x64(<1 x i64> %0) {
+; CHECKO0-LABEL: popcount1x64:
+; CHECKO0:       // %bb.0: // %Entry
+; CHECKO0-NEXT:    fmov x0, d0
+; CHECKO0-NEXT:    fmov d0, x0
+; CHECKO0-NEXT:    cnt v0.8b, v0.8b
+; CHECKO0-NEXT:    uaddlv h0, v0.8b
+; CHECKO0-NEXT:    // kill: def $q0 killed $h0
+; CHECKO0-NEXT:    mov w8, v0.s[0]
+; CHECKO0-NEXT:    // kill: def $x8 killed $w8
+; CHECKO0-NEXT:    fmov d0, x8
+; CHECKO0-NEXT:    ret
+;
+; CHECK-LABEL: popcount1x64:
+; CHECK:       // %bb.0: // %Entry
+; CHECK-NEXT:    cnt v0.8b, v0.8b
+; CHECK-NEXT:    uaddlp v0.4h, v0.8b
+; CHECK-NEXT:    uaddlp v0.2s, v0.4h
+; CHECK-NEXT:    uaddlp v0.1d, v0.2s
+; CHECK-NEXT:    ret
+Entry:
+  %1 = tail call <1 x i64> @llvm.ctpop.v1i64(<1 x i64> %0)
+  ret <1 x i64> %1
+}
+
+declare <1 x i64> @llvm.ctpop.v1i64(<1 x i64>)
+
+define <4 x i32> @popcount4x32(<4 x i32> %0) {
+; CHECKO0-LABEL: popcount4x32:
+; CHECKO0:       // %bb.0: // %Entry
+; CHECKO0-NEXT:    cnt v0.16b, v0.16b
+; CHECKO0-NEXT:    uaddlp v0.8h, v0.16b
+; CHECKO0-NEXT:    uaddlp v0.4s, v0.8h
+; CHECKO0-NEXT:    ret
+;
+; NEON-LABEL: popcount4x32:
+; NEON:       // %bb.0: // %Entry
+; NEON-NEXT:    cnt v0.16b, v0.16b
+; NEON-NEXT:    uaddlp v0.8h, v0.16b
+; NEON-NEXT:    uaddlp v0.4s, v0.8h
+; NEON-NEXT:    ret
+;
+; DOT-LABEL: popcount4x32:
+; DOT:       // %bb.0: // %Entry
+; DOT-NEXT:    movi v1.16b, #1
+; DOT-NEXT:    cnt v2.16b, v0.16b
+; DOT-NEXT:    movi v0.2d, #0000000000000000
+; DOT-NEXT:    udot v0.4s, v1.16b, v2.16b
+; DOT-NEXT:    ret
+;
+; SVE-LABEL: popcount4x32:
+; SVE:       // %bb.0: // %Entry
+; SVE-NEXT:    cnt v0.16b, v0.16b
+; SVE-NEXT:    uaddlp v0.8h, v0.16b
+; SVE-NEXT:    uaddlp v0.4s, v0.8h
+; SVE-NEXT:    ret
+Entry:
+  %1 = tail call <4 x i32> @llvm.ctpop.v4i32(<4 x i32> %0)
+  ret <4 x i32> %1
+}
+
+declare <4 x i32> @llvm.ctpop.v4i32(<4 x i32>)
+
+define <2 x i32> @popcount2x32(<2 x i32> %0) {
+; CHECKO0-LABEL: popcount2x32:
+; CHECKO0:       // %bb.0: // %Entry
+; CHECKO0-NEXT:    cnt v0.8b, v0.8b
+; CHECKO0-NEXT:    uaddlp v0.4h, v0.8b
+; CHECKO0-NEXT:    uaddlp v0.2s, v0.4h
+; CHECKO0-NEXT:    ret
+;
+; NEON-LABEL: popcount2x32:
+; NEON:       // %bb.0: // %Entry
+; NEON-NEXT:    cnt v0.8b, v0.8b
+; NEON-NEXT:    uaddlp v0.4h, v0.8b
+; NEON-NEXT:    uaddlp v0.2s, v0.4h
+; NEON-NEXT:    ret
+;
+; DOT-LABEL: popcount2x32:
+; DOT:       // %bb.0: // %Entry
+; DOT-NEXT:    movi v1.2d, #0000000000000000
+; DOT-NEXT:    cnt v0.8b, v0.8b
+; DOT-NEXT:    movi v2.8b, #1
+; DOT-NEXT:    udot v1.2s, v2.8b, v0.8b
+; DOT-NEXT:    fmov d0, d1
+; DOT-NEXT:    ret
+;
+; SVE-LABEL: popcount2x32:
+; SVE:       // %bb.0: // %Entry
+; SVE-NEXT:    cnt v0.8b, v0.8b
+; SVE-NEXT:    uaddlp v0.4h, v0.8b
+; SVE-NEXT:    uaddlp v0.2s, v0.4h
+; SVE-NEXT:    ret
+Entry:
+  %1 = tail call <2 x i32> @llvm.ctpop.v2i32(<2 x i32> %0)
+  ret <2 x i32> %1
+}
+
+declare <2 x i32> @llvm.ctpop.v2i32(<2 x i32>)
+
+define <8 x i16> @popcount8x16(<8 x i16> %0) {
+; CHECKO0-LABEL: popcount8x16:
+; CHECKO0:       // %bb.0: // %Entry
+; CHECKO0-NEXT:    cnt v0.16b, v0.16b
+; CHECKO0-NEXT:    uaddlp v0.8h, v0.16b
+; CHECKO0-NEXT:    ret
+;
+; CHECK-LABEL: popcount8x16:
+; CHECK:       // %bb.0: // %Entry
+; CHECK-NEXT:    cnt v0.16b, v0.16b
+; CHECK-NEXT:    uaddlp v0.8h, v0.16b
+; CHECK-NEXT:    ret
+Entry:
+  %1 = tail call <8 x i16> @llvm.ctpop.v8i16(<8 x i16> %0)
+  ret <8 x i16> %1
+}
+
+declare <8 x i16> @llvm.ctpop.v8i16(<8 x i16>)
+
+define <4 x i16> @popcount4x16(<4 x i16> %0) {
+; CHECKO0-LABEL: popcount4x16:
+; CHECKO0:       // %bb.0: // %Entry
+; CHECKO0-NEXT:    cnt v0.8b, v0.8b
+; CHECKO0-NEXT:    uaddlp v0.4h, v0.8b
+; CHECKO0-NEXT:    ret
+;
+; CHECK-LABEL: popcount4x16:
+; CHECK:       // %bb.0: // %Entry
+; CHECK-NEXT:    cnt v0.8b, v0.8b
+; CHECK-NEXT:    uaddlp v0.4h, v0.8b
+; CHECK-NEXT:    ret
+Entry:
+  %1 = tail call <4 x i16> @llvm.ctpop.v4i16(<4 x i16> %0)
+  ret <4 x i16> %1
+}
+
+declare <4 x i16> @llvm.ctpop.v4i16(<4 x i16>)