Repository navigation
perf(arrow-buffer): faster count_set_bits on aarch64 by summing byte popcounts before widening - #11818
perf(arrow-buffer): faster count_set_bits on aarch64 by summing byte popcounts before widening#11818mbutrovich wants to merge 7 commits into
Conversation
|
run benchmark boolean_array |
|
run benchmark boolean_array |
|
🤖 Benchmark starting (GKE) | trigger Target: boolean_array Sharding factor: 1 (1 workers). Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark starting (GKE) | trigger Target: boolean_array Sharding factor: 1 (1 workers). Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
resources:
arch: amd64
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: boolean_array Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
shards: 1
DetailsPer-runner informationboolean_array — shard 1/1Node: gk3-benchmark-cluster-nap-5fy40jly-b59e93e3-db75 Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench boolean_arrayCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: boolean_array Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
resources:
arch: amd64
shards: 1
DetailsPer-runner informationboolean_array — shard 1/1Node: gk3-benchmark-cluster-nap-vp60frle-472b4730-cln4 Instance: c4-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench boolean_arrayCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark boolean_array |
|
🤖 Benchmark starting (GKE) | trigger Target: boolean_array Sharding factor: 1 (1 workers). Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
resources:
arch: amd64
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: boolean_array Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
resources:
arch: amd64
shards: 1
DetailsPer-runner informationboolean_array — shard 1/1Node: gk3-benchmark-cluster-nap-1uz6vgaz-77581093-zsnw Instance: c4-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench boolean_arrayCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
run benchmark boolean_array |
|
🤖 Benchmark starting (GKE) | trigger Target: boolean_array Sharding factor: 1 (1 workers). Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
env:
BENCH_FILTER: true_count
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
run benchmark boolean_array |
|
🤖 Benchmark starting (GKE) | trigger Target: boolean_array Sharding factor: 1 (1 workers). Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
env:
BENCH_FILTER: true_count
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
resources:
arch: amd64
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: boolean_array Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
env:
BENCH_FILTER: true_count
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
shards: 1
DetailsPer-runner informationboolean_array — shard 1/1Node: gk3-benchmark-cluster-nap-rw8ujupk-523dc0ee-zbk6 Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench boolean_arrayCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: boolean_array Comparing count-set-bits-aarch64 (dcfd4da) to 01a9917 (merge-base) diff Run configurationrun benchmark boolean_array
env:
BENCH_FILTER: true_count
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
resources:
arch: amd64
shards: 1
DetailsPer-runner informationboolean_array — shard 1/1Node: gk3-benchmark-cluster-nap-1bw8wop1-2d2ce690-zt4s Instance: c4-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench boolean_arrayCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
Marking draft so I can open some bug-fix PRs. Reviews are still welcome, I just can't have more PRs open. |
|
cc @devanbenz @mightsleep since you gave such great feedback on #11467. |
|
I should have some time tomorrow or in the next few days to take a look at this. 🫡 |
|
I relied on AI to help with this review, and I ran the checks below myself. Please push back if I've got something wrong. Thanks @mbutrovich. The approach looks right to me. Correctness
Benchmark on Apple Silicon
Suggestion: count the tail in words, not bytes I tried this instead: let tail = &words[words.len() - remainder.len() / 8..];
total + tail.iter().map(|w| w.count_ones() as usize).sum::<usize>()
Nits
I couldn't test the non-aarch64 |
|
run benchmark boolean_array |
|
🤖 Benchmark starting (GKE) | trigger Target: boolean_array Sharding factor: 1 (1 workers). Comparing count-set-bits-aarch64 (110c2e8) to f6c3989 diff Run configurationrun benchmark boolean_array
baseline:
ref: f6c3989363
env:
BENCH_FILTER: true_count
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
run benchmark boolean_array |
|
🤖 Benchmark starting (GKE) | trigger Target: boolean_array Sharding factor: 1 (1 workers). Comparing count-set-bits-aarch64 (110c2e8) to f6c3989 diff Run configurationrun benchmark boolean_array
baseline:
ref: f6c3989363
env:
BENCH_FILTER: true_count
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
shards: 1
Results will be posted when all workers finish. File an issue against this benchmark runner |
|
Thanks @andygrove. I switched the tail to count words. I think it reads better, and it gets a bit of speedup.
A 64-bit input doesn't reach The tail runs when the aligned middle isn't a multiple of 8 words. I added 448-bit (7 tail words, no full chunk) and 960-bit (one chunk and 7 tail words) cases to the bench. Per word, the byte tail compiles to
The 64, 1024, and 65536-bit cases moved by less than 3.1%. With I split with
I added a seeded random input to |
|
🤖 Benchmark completed (GKE) | trigger Target: boolean_array Comparing count-set-bits-aarch64 (110c2e8) to f6c3989 diff Run configurationrun benchmark boolean_array
baseline:
ref: f6c3989363
env:
BENCH_FILTER: true_count
RUSTFLAGS: -C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6
shards: 1
DetailsPer-runner informationboolean_array — shard 1/1Node: gk3-benchmark-cluster-nap-1mvvlccl-0352b74c-j27r Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench boolean_arrayCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Target: boolean_array Comparing count-set-bits-aarch64 (110c2e8) to f6c3989 diff Run configurationrun benchmark boolean_array
baseline:
ref: f6c3989363
env:
BENCH_FILTER: true_count
shards: 1
DetailsPer-runner informationboolean_array — shard 1/1Node: gk3-benchmark-cluster-nap-1fboxlbk-f8e512a2-rbc7 Instance: c4a-highmem-16 (12 vCPU / 65 GiB) uname: BENCH_COMMAND: cargo bench --features=arrow,async,test_common,experimental,object_store --bench boolean_arrayCPU Details (lscpu)Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
@andygrove the bot ran the tail change on Axion against the bench commit. The 448-bit and 960-bit cases are 1.39x to 1.49x faster, and the 64, 1024, and 65536-bit cases moved by less than 1.04x (default, aligned). Looks like a good improvement! |
devanbenz
left a comment
There was a problem hiding this comment.
Just a thought experiment on exciting ARM intrinsics. I'm not really a huge ARM user so it's mostly research with an LLM helper. The resulting code I worked on didn't really improve much just the amount of instructions generated. The actual throughput seems to be the same. Overall this looks good to me!
| /// 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 { |
There was a problem hiding this 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:
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
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_2This 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:
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:
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_2That 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:
Benchmarks snippet this PR
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]
Benchmark snippet with UADALP
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.
Which issue does this PR close?
count_set_bitson aarch64 by summing byte popcounts before widening #11776.Rationale for this change
See #11776.
What changes are included in this PR?
On aarch64,
UnalignedBitChunk::count_onessums byte popcounts inu8lanes before widening them, as the issue proposes. The words after the last 64-byte chunk are counted withu64::count_ones. Other targets keep theu64::count_onessum for all words. On all targets,count_onesnow counts the prefix and suffix words separately from the aligned middle.Are these changes tested?
Adds
test_unaligned_bit_chunk_count_onesinarrow-buffer/src/util/bit_chunk_iterator.rs. It covers sizes around the 64-byte chunk and theu16lane flush, with uniform and seeded random bytes.The
true_countcases of theboolean_arraybench measure the change. This PR adds 448-bit and 960-bit cases, which end in 7 words after the last 64-byte chunk. On the bot's default runner (Google Axion, Neoverse V2),true_countwithout nulls is 1.87x faster at 65536 bits (results). A run with-C llvm-args=-align-all-functions=6 -C llvm-args=-align-all-nofallthru-blocks=6gave the same result, so code layout doesn't explain it (results).On amd64 (Xeon Platinum 8581C), the 65536-bit cases are unchanged within noise and the generated x86_64 inner loop is the same before and after (first, second).
The 64-bit and 1024-bit cases are faster on both architectures in the bot's build (thin LTO, 16 codegen units), where
maincopies theUnalignedBitChunkinto an out-of-lineChain::fold. Withcodegen-units = 1, as DataFusion and Comet build,maininlines that fold, so I don't expect consumers to see most of that gain.Are there any significant user-facing changes?
No.
AI Disclosure
I wrote this PR with help from Claude but understand and support the core approach, implementation, and test coverage.