LLVM generates branch soup for array partition point

[ds84182]

I tried this code:

pub fn partition_point_short_array_n<const N: usize>(array: &[usize; N]) -> usize {
    // array.partition_point(|x| *x != 0)
    let mut i = 0;
    while i < N {
        if array[i] == 0 { break }
        i += 1;
    }
    i
}

pub fn partition_point_short_array(array: &[usize; 24]) -> usize {
    partition_point_short_array_n(array)
}

I expected to see this happen: Assembly composed of cmov and/or adox instruction. Or at least mov + je to a single exit branch.

Instead, this happened:

playground::partition_point_short_array:
	cmpq	$0, (%rdi)
	je	.LBB0_1
	cmpq	$0, 8(%rdi)
	je	.LBB0_3
	cmpq	$0, 16(%rdi)
	je	.LBB0_5
	cmpq	$0, 24(%rdi)
	je	.LBB0_7
	cmpq	$0, 32(%rdi)
	je	.LBB0_9
	cmpq	$0, 40(%rdi)
	je	.LBB0_11
	cmpq	$0, 48(%rdi)
	je	.LBB0_13
	cmpq	$0, 56(%rdi)
	je	.LBB0_15
	cmpq	$0, 64(%rdi)
	je	.LBB0_17
	cmpq	$0, 72(%rdi)
	je	.LBB0_19
	cmpq	$0, 80(%rdi)
	je	.LBB0_21
	cmpq	$0, 88(%rdi)
	je	.LBB0_23
	cmpq	$0, 96(%rdi)
	je	.LBB0_25
	cmpq	$0, 104(%rdi)
	je	.LBB0_27
	cmpq	$0, 112(%rdi)
	je	.LBB0_29
	cmpq	$0, 120(%rdi)
	je	.LBB0_31
	cmpq	$0, 128(%rdi)
	je	.LBB0_33
	cmpq	$0, 136(%rdi)
	je	.LBB0_35
	cmpq	$0, 144(%rdi)
	je	.LBB0_37
	cmpq	$0, 152(%rdi)
	je	.LBB0_39
	cmpq	$0, 160(%rdi)
	je	.LBB0_41
	cmpq	$0, 168(%rdi)
	je	.LBB0_43
	cmpq	$0, 176(%rdi)
	je	.LBB0_45
	cmpq	$1, 184(%rdi)
	movl	$24, %eax
	sbbq	$0, %rax
	retq

.LBB0_1:
	xorl	%eax, %eax
	retq

.LBB0_3:
	movl	$1, %eax
	retq

.LBB0_5:
	movl	$2, %eax
	retq

.LBB0_7:
	movl	$3, %eax
	retq

.LBB0_9:
	movl	$4, %eax
	retq

.LBB0_11:
	movl	$5, %eax
	retq

.LBB0_13:
	movl	$6, %eax
	retq

.LBB0_15:
	movl	$7, %eax
	retq

.LBB0_17:
	movl	$8, %eax
	retq

.LBB0_19:
	movl	$9, %eax
	retq

.LBB0_21:
	movl	$10, %eax
	retq

.LBB0_23:
	movl	$11, %eax
	retq

.LBB0_25:
	movl	$12, %eax
	retq

.LBB0_27:
	movl	$13, %eax
	retq

.LBB0_29:
	movl	$14, %eax
	retq

.LBB0_31:
	movl	$15, %eax
	retq

.LBB0_33:
	movl	$16, %eax
	retq

.LBB0_35:
	movl	$17, %eax
	retq

.LBB0_37:
	movl	$18, %eax
	retq

.LBB0_39:
	movl	$19, %eax
	retq

.LBB0_41:
	movl	$20, %eax
	retq

.LBB0_43:
	movl	$21, %eax
	retq

.LBB0_45:
	movl	$22, %eax
	retq

First occurs in 1.19.0 with the alternative code snippet:

#[no_mangle]
pub fn partition_point_short_array(array: &[usize; 24]) -> usize {
    let mut i = 0;
    while i < 24 {
        if array[i] == 0 { break }
        i += 1;
    }
    i
}

[saethlin]

This is loop unrolling. So...

I expected to see this happen: Assembly composed of cmov and/or adox instruction. Or at least mov + je to a single exit branch.

Why do you think this would be better? If you think the unrolled version is slower, do you have a benchmark? If you think the code size is problematic, does -Copt-level=s do what you want? Is that optimization setting a better fit for your codebase?

[the8472]

std's partition_point implements a binary search, your simplified version is a linear scan. LLVM has no optimizations that would turn a linear search into a binary one since it wouldn't be able to know that the array is sorted/partitioned.

[workingjubilee]

Is there a reason that std's partition_point does not serve well on your code, @ds84182?

Is this code not faster for the small array size you are concerned about?

[theemathas]

The "a bunch of jumps" approach will need fewer comparisons than the "a bunch of cmovs" approach if the zero is usually near the start of the array. Since the compiler doesn't know if this is often the case in your workload, it assumes that you know what you're doing and therefore preserves what your code does.

[cvijdea-bd]

FWIW if want a branchless implementation, then SIMD will likely be 2-5x faster than the scalar version, depending on what target features you can afford to use.
Reducing the array element size from usize would also help.

#![feature(portable_simd)]
use std::simd::cmp::SimdPartialEq;
use std::simd::Simd;

pub fn partition_point_simd(array: &[usize; 24]) -> usize {
    let mut array_zext = [0; 32];
    array_zext[..24].copy_from_slice(array);

    let array = Simd::from_array(array_zext);
    let mask = array.simd_eq(Simd::splat(0));
    mask.to_bitmask().trailing_zeros() as usize
}

However, in the latest nightly, std::binary_search is also branchless (#128254), and it turns out it's the fastest implementation on inputs where the partition point is randomly distributed in 0..=24:

pub fn partition_point_std(array: &[usize; 24]) -> usize {
    array.partition_point(|x| *x != 0)
}

It's also remarkably small:

playground::partition_point_std:
	mov rax, qword ptr [rcx + 96]
	test rax, rax
	mov edx, 12
	cmove rdx, rax
	lea rax, [rdx + 6]
	mov r8d, eax
	cmp qword ptr [rcx + 8*r8], 0
	cmovne rdx, rax
	lea rax, [rdx + 3]
	mov r8d, eax
	cmp qword ptr [rcx + 8*r8], 0
	cmovne rdx, rax
	lea r8, [rdx + 1]
	cmp qword ptr [rcx + 8*rdx + 8], 0
	cmove r8, rdx
	lea rax, [r8 + 1]
	cmp qword ptr [rcx + 8*r8 + 8], 0
	cmove rax, r8
	cmp qword ptr [rcx + 8*rax], 1
	sbb rax, -1
	ret

[usize; 24] arg with -Ctarget-cpu=x86-64

test bench_partition_point_branchless ... bench:   1,209,903.75 ns/iter (+/- 683,113.25) = 15869 MB/s
test bench_partition_point_linear     ... bench:   1,889,390.00 ns/iter (+/- 371,780.75) = 10162 MB/s
test bench_partition_point_simd       ... bench:     792,851.25 ns/iter (+/- 637,814.88) = 24216 MB/s
test bench_partition_point_std        ... bench:     558,115.62 ns/iter (+/- 306,352.19) = 34401 MB/s

[u32; 24] arg with -Ctarget-cpu=x86-64-v3

test bench_partition_point_branchless ... bench:     312,764.69 ns/iter (+/- 104,581.22) = 30694 MB/s
test bench_partition_point_linear     ... bench:   1,153,777.50 ns/iter (+/- 239,229.50) = 8320 MB/s
test bench_partition_point_simd       ... bench:     198,122.81 ns/iter (+/- 39,717.12) = 48454 MB/s
test bench_partition_point_std        ... bench:     235,001.56 ns/iter (+/- 98,562.62) = 40850 MB/s

partition_point_branchless is the scalar equivalent of partition_point_simd - it's faster in the second case because it gets auto-vectorized, but badly.

fn partition_point_branchless<const N: usize>(array: &[u32; N]) -> usize {
    let mut mask: u32 = 0;
    for i in 0..N {
        if array[i] == 0 {
            mask |= 1 << i;
        }
    }
    return mask.trailing_zeros() as usize;
}

[tgross35]

I think this can be closed - LLVM is doing something reasonable here with the context it has and there are some alternatives provided.