devanbenz commented on code in PR #11818:
URL: https://github.com/apache/arrow-rs/pull/11818#discussion_r4419786232
##########
arrow-buffer/src/util/bit_chunk_iterator.rs:
##########
@@ -168,10 +170,46 @@ impl<'a> UnalignedBitChunk<'a> {
/// Counts the number of ones
pub fn count_ones(&self) -> usize {
- self.iter().map(|x| x.count_ones() as usize).sum()
+ let edges = self.prefix.into_iter().chain(self.suffix);
+ edges.map(|x| x.count_ones() as usize).sum::<usize>() +
count_ones_words(self.chunks)
}
}
+#[cfg(not(target_arch = "aarch64"))]
+fn count_ones_words(words: &[u64]) -> usize {
+ words.iter().map(|x| x.count_ones() as usize).sum()
+}
+
+/// Summing four byte popcounts in `u8` lanes before widening lets LLVM emit
one
+/// widening per 64 bytes. A `u64::count_ones` sum widens every 16 bytes. On
x86
+/// the `u64` sum measured faster (#11776), so this is only used on aarch64.
+#[cfg(target_arch = "aarch64")]
+fn count_ones_words(words: &[u64]) -> usize {
Review Comment:
I was able to inspect this code on my M2 Macbook and didn't notice any
issues. Taking a peek at the llvm IR gen:
```llvm
vector.body: ; preds = %bb15.preheader,
%vector.body
%lane_counts.sroa.0.1 = phi <16 x i16> [ zeroinitializer, %bb15.preheader
], [ %21, %vector.body ]
%iter1.sroa.0.027 = phi ptr [ %iter.sroa.0.029, %bb15.preheader ], [ %_75,
%vector.body ]
%wide.load = load <16 x i8>, ptr %iter1.sroa.0.027, align 1
%10 = tail call range(i8 0, 9) <16 x i8> @llvm.ctpop.v16i8(<16 x i8>
%wide.load)
%11 = getelementptr inbounds nuw i8, ptr %iter1.sroa.0.027, i64 16
%wide.load48 = load <16 x i8>, ptr %11, align 1
%12 = tail call range(i8 0, 9) <16 x i8> @llvm.ctpop.v16i8(<16 x i8>
%wide.load48)
%13 = add nuw nsw <16 x i8> %12, %10
%14 = getelementptr inbounds nuw i8, ptr %iter1.sroa.0.027, i64 32
%wide.load49 = load <16 x i8>, ptr %14, align 1
%15 = tail call range(i8 0, 9) <16 x i8> @llvm.ctpop.v16i8(<16 x i8>
%wide.load49)
%16 = add nuw nsw <16 x i8> %13, %15
%17 = getelementptr inbounds nuw i8, ptr %iter1.sroa.0.027, i64 48
%wide.load50 = load <16 x i8>, ptr %17, align 1
%18 = tail call range(i8 0, 9) <16 x i8> @llvm.ctpop.v16i8(<16 x i8>
%wide.load50)
%19 = add nuw nsw <16 x i8> %16, %18
%20 = zext nneg <16 x i8> %19 to <16 x i16>
%21 = add <16 x i16> %lane_counts.sroa.0.1, %20
%_75 = getelementptr inbounds nuw i8, ptr %iter1.sroa.0.027, i64 64
%_69 = icmp eq ptr %_75, %_53
br i1 %_69, label %bb8.i12.preheader, label %vector.body
}
```
Looks like there is consistent vector lane usage as expected ` %wide.load48
= load <16 x i8>, ptr %11, align 1`. Doubt we can squeeze a ton more
performance out of it.
Checking in on the main loop assembly
```asm
LBB85_3:
ldp q2, q3, [x13]
cnt.16b v2, v2
cnt.16b v3, v3
add.16b v2, v3, v2
ldp q3, q4, [x13, #32]
cnt.16b v3, v3
cnt.16b v4, v4
add.16b v3, v3, v4
add.16b v2, v2, v3
uaddw2.8h v1, v1, v2
uaddw.8h v0, v0, v2
add x13, x13, #64
cmp x13, x12
b.ne LBB85_3
ushll.4s v2, v0, #0
ushll.4s v3, v1, #0
ushll2.4s v0, v0, #0
ushll2.4s v1, v1, #0
uaddl2.2d v4, v0, v1
uaddl2.2d v5, v2, v3
add.2d v4, v5, v4
uaddl.2d v0, v0, v1
uaddl.2d v1, v2, v3
add.2d v0, v1, v0
add.2d v0, v0, v4
addp.2d d0, v0
fmov x13, d0
add x8, x13, x8
cbnz x11, LBB85_2
```
This next portion you'll need to double check since I did this experiment
with the help of LLMs.
There is an operation within ARM called `UADALP`, it could take the `uaddw2`
and `uaddw` to combine them. Also, this shrinks the epilogue. Because the old
version ends with two accumulators, it needs the 13-instruction `ushll` /
`uaddl` / `add` / `addp` tree to fold 16 lanes down to one number. The new
version has 8 lanes in one register, so a single uaddlv.8h does it.
The code requires lots of intrinsics and is heavily AI generated:
```rust
fn count_ones_words(words: &[u64]) -> usize {
use core::arch::aarch64::*;
let (chunks, tail) = words.as_chunks::<8>();
// Per chunk each u8 lane is ≤ 32; UADALP adds pairs, so each u16 lane
gains ≤ 64.
const CHUNKS_PER_FLUSH: usize = (u16::MAX / 64) as usize; // 1023
let mut total = 0;
for group in chunks.chunks(CHUNKS_PER_FLUSH) {
let mut acc = unsafe { vdupq_n_u16(0) };
for chunk in group {
let p = chunk.as_ptr().cast::<u8>();
// SAFETY: a chunk is 8 u64s = 64 readable bytes.
let (a, b, c, d) = unsafe {
(
vld1q_u8(p),
vld1q_u8(p.add(16)),
vld1q_u8(p.add(32)),
vld1q_u8(p.add(48)),
)
};
let s = unsafe {
vaddq_u8(
vaddq_u8(vcntq_u8(a), vcntq_u8(b)),
vaddq_u8(vcntq_u8(c), vcntq_u8(d)),
)
};
acc = unsafe { vpadalq_u8(acc, s) }; // UADALP
}
total += unsafe { vaddlvq_u16(acc) as usize };
}
total + tail.iter().map(|w| w.count_ones() as usize).sum::<usize>()
}
```
With the resulting ASM:
```asm
LBB85_3:
ldp q1, q2, [x13]
ldp q3, q4, [x13, #32]
cnt.16b v1, v1
cnt.16b v2, v2
cnt.16b v3, v3
cnt.16b v4, v4
add.16b v1, v2, v1
add.16b v2, v3, v4
add.16b v1, v1, v2
uadalp.8h v0, v1
add x13, x13, #64
cmp x13, x12
b.ne LBB85_3
uaddlv.8h s0, v0
fmov w13, s0
add x8, x8, x13
cbnz x11, LBB85_2
```
That being said, this should decrease the operations by 1, but, the code is
horrendous. Also, there was not really a noticable performance improvement in
the boolean_array benchmarks, I only ran through a small amount but noticing
the perf didn't really change I exited early lol:
<details>
<summary>Benchmarks snippet this PR</summary>
```
true_count(all_true, 64)
time: [9.8319 ns 9.8490 ns 9.8649 ns]
Found 3 outliers among 100 measurements (3.00%)
3 (3.00%) high mild
has_true(all_true, 64) time: [8.0431 ns 8.0521 ns 8.0619 ns]
Found 6 outliers among 100 measurements (6.00%)
5 (5.00%) high mild
1 (1.00%) high severe
has_false(all_true, 64) time: [9.3496 ns 9.3617 ns 9.3747 ns]
Found 2 outliers among 100 measurements (2.00%)
2 (2.00%) high mild
true_count(all_false, 64)
time: [9.8090 ns 9.8213 ns 9.8349 ns]
Found 5 outliers among 100 measurements (5.00%)
5 (5.00%) high mild
has_true(all_false, 64) time: [8.9055 ns 8.9186 ns 8.9326 ns]
Found 1 outliers among 100 measurements (1.00%)
1 (1.00%) high mild
has_false(all_false, 64)
time: [8.8893 ns 8.9028 ns 8.9185 ns]
Found 3 outliers among 100 measurements (3.00%)
1 (1.00%) high mild
2 (2.00%) high severe
true_count(mixed_early, 64)
time: [9.8369 ns 9.8573 ns 9.8821 ns]
Found 6 outliers among 100 measurements (6.00%)
4 (4.00%) high mild
2 (2.00%) high severe
has_false(mixed_early, 64)
time: [9.0135 ns 9.0300 ns 9.0483 ns]
Found 5 outliers among 100 measurements (5.00%)
4 (4.00%) high mild
1 (1.00%) high severe
true_count(nulls_all_true, 64)
time: [9.2620 ns 9.2852 ns 9.3093 ns]
Found 1 outliers among 100 measurements (1.00%)
1 (1.00%) high mild
has_true(nulls_all_true, 64)
time: [7.9764 ns 8.0999 ns 8.2228 ns]
Found 2 outliers among 100 measurements (2.00%)
2 (2.00%) high mild
has_false(nulls_all_true, 64)
time: [10.102 ns 10.130 ns 10.159 ns]
Found 5 outliers among 100 measurements (5.00%)
4 (4.00%) high mild
1 (1.00%) high severe
take_n_true_unique(32) time: [87.153 ns 88.768 ns 90.064 ns]
```
</details>
<details>
<summary>Benchmark snippet with UADALP</summary>
```
true_count(all_true, 64)
time: [9.8282 ns 9.8448 ns 9.8605 ns]
change: [−0.3596% −0.1550% +0.0563%] (p = 0.13 >
0.05)
No change in performance detected.
Found 1 outliers among 100 measurements (1.00%)
1 (1.00%) high mild
has_true(all_true, 64) time: [8.1367 ns 8.1619 ns 8.1894 ns]
change: [+1.2133% +1.4913% +1.7940%] (p = 0.00 <
0.05)
Performance has regressed.
has_false(all_true, 64) time: [9.4249 ns 9.4461 ns 9.4662 ns]
change: [+0.5042% +0.6909% +0.8768%] (p = 0.00 <
0.05)
Change within noise threshold.
true_count(all_false, 64)
time: [9.8631 ns 9.8795 ns 9.8952 ns]
change: [−0.0373% +0.2251% +0.4831%] (p = 0.10 >
0.05)
No change in performance detected.
Found 2 outliers among 100 measurements (2.00%)
2 (2.00%) high mild
has_true(all_false, 64) time: [8.9118 ns 8.9331 ns 8.9598 ns]
change: [−0.2245% +0.0002% +0.2247%] (p = 1.00 >
0.05)
No change in performance detected.
Found 3 outliers among 100 measurements (3.00%)
3 (3.00%) high mild
has_false(all_false, 64)
time: [8.9489 ns 8.9603 ns 8.9718 ns]
change: [−1.1010% −0.1834% +0.3938%] (p = 0.75 >
0.05)
No change in performance detected.
Found 2 outliers among 100 measurements (2.00%)
2 (2.00%) high mild
true_count(mixed_early, 64)
time: [9.9050 ns 9.9222 ns 9.9405 ns]
change: [+0.2972% +0.6001% +0.8762%] (p = 0.00 <
0.05)
Change within noise threshold.
Found 1 outliers among 100 measurements (1.00%)
1 (1.00%) high mild
has_false(mixed_early, 64)
time: [8.9812 ns 9.0043 ns 9.0301 ns]
change: [−0.0181% +0.3545% +0.7718%] (p = 0.07 >
0.05)
No change in performance detected.
Found 4 outliers among 100 measurements (4.00%)
2 (2.00%) high mild
2 (2.00%) high severe
true_count(nulls_all_true, 64)
time: [9.2890 ns 9.3118 ns 9.3366 ns]
change: [+0.1582% +0.4465% +0.7484%] (p = 0.00 <
0.05)
Change within noise threshold.
has_true(nulls_all_true, 64)
time: [7.8813 ns 7.9714 ns 8.0606 ns]
change: [−3.3464% −1.5696% +0.1403%] (p = 0.09 >
0.05)
No change in performance detected.
has_false(nulls_all_true, 64)
time: [10.040 ns 10.052 ns 10.065 ns]
change: [−0.9417% −0.6809% −0.4366%] (p = 0.00 <
0.05)
Change within noise threshold.
Found 6 outliers among 100 measurements (6.00%)
2 (2.00%) high mild
4 (4.00%) high severe
take_n_true_unique(32) time: [87.706 ns 89.046 ns 90.072 ns]
change: [−1.9092% +1.5018% +4.6731%] (p = 0.38 >
0.05)
No change in performance detected.
```
</details>
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
