LLVM should recognize more patterns that map to pmovmsk or similar instructions

[jhorstmann]

This comes up regularly in functions that process or detect ascii inputs. My current example is in rust, int a const function, which prevents me from using explicit simd intrinsics, but the code generated by clang is nearly the same.

pub unsafe fn ascii_prefix(input: &[u8]) -> usize {
    let mut mask = 0_u16;
    for i in 0..16 {
        mask |= ((*input.get_unchecked(i) < 128) as u16) << i;
    }
    mask.trailing_ones() as usize
}

Ideally this would map to 4 instructions: compare, movmsk, not and tzcnt. But currently it gets compiled to about 50 instructions with a mix of scalar and vector instructions: https://godbolt.org/z/jaKb5TMrn

A sightly simpler example that only tries to detect fully ascii chunks leads to similar complex output.

pub unsafe fn is_ascii_mask(input: &[u8]) -> bool {
    let mut is_ascii = [false; 16];
    for i in 0..16 {
        is_ascii[i] = *input.get_unchecked(i) <= 127;
    }

    let mut mask = 0_u16;
    for i in 0..16 {
        mask |= (is_ascii[i] as u16) << i;
    }

    mask == 0xFFFF
}

However, for detection only I found one pattern that generates the wanted instruction sequence, so LLVM has some logic to detect this pattern:

pub unsafe fn is_ascii_sum(input: &[u8]) -> bool {
    let mut is_ascii = [false; 16];
    for i in 0..16 {
        is_ascii[i] = *input.get_unchecked(i) <= 127;
    }

    let mut count = 0_u8;
    for i in 0..16 {
        count += is_ascii[i] as u8;
    }

    count == 16
}

example::is_ascii_sum::h00bc3be7e76426eb:
        vmovdqu xmm0, xmmword ptr [rdi]
        vpmovmskb       eax, xmm0
        test    eax, eax
        sete    al
        ret

Unfortunately this pattern using sum cannot be used for the initial example which needs a mask as output.

#45398 seems related, but focuses a bit more on the bit counting aspect, where movmsk might only be beneficial if the popcnt instruction is also available.

[llvmbot]

@llvm/issue-subscribers-backend-x86

Author: Jörn Horstmann (jhorstmann)

This comes up regularly in functions that process or detect ascii inputs. My current example is in rust, int a `const` function, which prevents me from using explicit simd intrinsics, but the code generated by clang is nearly the same.

pub unsafe fn ascii_prefix(input: &[u8]) -> usize {
    let mut mask = 0_u16;
    for i in 0..16 {
        mask |= ((*input.get_unchecked(i) < 128) as u16) << i;
    }
    mask.trailing_ones() as usize
}

Ideally this would map to 4 instructions: compare, movmsk, not and tzcnt. But currently it gets compiled to about 50 instructions with a mix of scalar and vector instructions: https://godbolt.org/z/jaKb5TMrn

A sightly simpler example that only tries to detect fully ascii chunks leads to similar complex output.

pub unsafe fn is_ascii_mask(input: &[u8]) -> bool {
    let mut is_ascii = [false; 16];
    for i in 0..16 {
        is_ascii[i] = *input.get_unchecked(i) <= 127;
    }

    let mut mask = 0_u16;
    for i in 0..16 {
        mask |= (is_ascii[i] as u16) << i;
    }

    mask == 0xFFFF
}

However, for detection only I found one pattern that generates the wanted instruction sequence, so LLVM has some logic to detect this pattern:

pub unsafe fn is_ascii_sum(input: &[u8]) -> bool {
    let mut is_ascii = [false; 16];
    for i in 0..16 {
        is_ascii[i] = *input.get_unchecked(i) <= 127;
    }

    let mut count = 0_u8;
    for i in 0..16 {
        count += is_ascii[i] as u8;
    }

    count == 16
}

example::is_ascii_sum::h00bc3be7e76426eb:
        vmovdqu xmm0, xmmword ptr [rdi]
        vpmovmskb       eax, xmm0
        test    eax, eax
        sete    al
        ret

Unfortunately this pattern using sum cannot be used for the initial example which needs a mask as output.

#45398 seems related, but focuses a bit more on the bit counting aspect, where movmsk might only be beneficial if the popcnt instruction is also available.

[RKSimon]

A lot of this is down to the SLP not correctly recognising the unrolled loop as all shifts:

pub unsafe fn ascii_prefix(input: &[u8]) -> usize {
    let mut mask = 0_u16;
    {
        mask |= ((*input.get_unchecked( 0) < 128) as u16) <<  0; // !!!! no shift
        mask |= ((*input.get_unchecked( 1) < 128) as u16) <<  1;
        mask |= ((*input.get_unchecked( 2) < 128) as u16) <<  2;
        mask |= ((*input.get_unchecked( 3) < 128) as u16) <<  3;
        mask |= ((*input.get_unchecked( 4) < 128) as u16) <<  4;
        mask |= ((*input.get_unchecked( 5) < 128) as u16) <<  5;
        mask |= ((*input.get_unchecked( 6) < 128) as u16) <<  6;
        mask |= ((*input.get_unchecked( 7) < 128) as u16) <<  7;
        mask |= ((*input.get_unchecked( 8) < 128) as u16) <<  8;
        mask |= ((*input.get_unchecked( 9) < 128) as u16) <<  9;
        mask |= ((*input.get_unchecked(10) < 128) as u16) << 10;
        mask |= ((*input.get_unchecked(11) < 128) as u16) << 11;
        mask |= ((*input.get_unchecked(12) < 128) as u16) << 12;
        mask |= ((*input.get_unchecked(13) < 128) as u16) << 13;
        mask |= ((*input.get_unchecked(14) < 128) as u16) << 14;
        mask |= ((*input.get_unchecked(15) < 128) as u16) << 15;
    }
    mask.trailing_ones() as usize
}

[jhorstmann]

Seems the same was previously reported for aarch64 in #81376

[okaneco]

I ran into similar issues trying to validate bytes as ASCII, only needing to detect if any byte was invalid.

I was able to successfully get constant chunk sizes to work like in the opening comment, but the auto-vectorized code for while-loops seemed like it could be better.

I was hoping the vector code could boil down to or-ing the bytes together and checking the most significant bit at the end, but the naive method of doing it was the only way I found not to generate a pmovmskb. Instead, it does an or-reduction and so does the C code.

pub const fn is_ascii_or_1(bytes: &[u8]) -> bool {
    let mut i = 0;
    let mut mask = 0u8;
    while i < bytes.len() {
        mask |= bytes[i];
        i += 1;
    }
   mask <= 127
}

#include <stdint.h>
int is_ascii(uint8_t *input, int len) {
    uint8_t mask = 0;
    for (int i = 0; i < len; i++) {
        mask |= input[i];
    }
    return mask <= 127;
}

From looking at the IR of the Rust code, the is_ascii_or_1 has a @llvm.vector.reduce.or.v16i8 in its middle block after the vector body instead of something like this pattern which produces a pmovmskb in the other examples.

  %2 = icmp slt <32 x i8> %1, zeroinitializer
  %3 = bitcast <32 x i1> %2 to i32
  %4 = icmp eq i32 %3, 0

Rust is_ascii functions with 2 chunk validating functions and 3 while loop example functions that check an entire slice, along with the emitted LLVM IR
https://rust.godbolt.org/z/Te18sc1G6

[RKSimon]

pub unsafe fn ascii_prefix(input: &[u8]) -> usize {
    let mut mask = 0_u16;
    for i in 0..16 {
        mask |= ((*input.get_unchecked(i) < 128) as u16) << i;
    }
    mask.trailing_ones() as usize
}

Much of this is the same old SLP problem - the shift with the i == 0 gets folded away and we then fail to vectorize it as only 15 of the 16 elements have a SHL node.