How to make std::simd generic over integer types

An example function that supports i32:

const N: usize = 16;

pub fn position_simd(arr: &[i32], needle: i32) -> Option<usize> {
    let (prefix, simd_data, suffix) = arr.as_simd::<N>();

    // handle prefix ... (not relevant for the example)

    // SIMD
    let simd_needle = i32x16::splat(needle);
    for (chunk_idx, chunk) in simd_data.iter().enumerate() {
        let mask = chunk.simd_eq(simd_needle).to_bitmask();
        if mask != 0 {
            return Some(prefix.len() + (chunk_idx * N) + (mask.trailing_zeros() as usize));
        }
    }

    // handle suffix ... (not relevant for the example)
}

(The exact algorithm isn't the interesting part, rather how to make similar functions handle multiple types)

The function does something similar to what iter().position() does. But how should I make it work for different types, for example u32? or maybe even all integer types?

My attempt:

const N: usize = 16;

pub fn position_simd<T>(arr: &[T], needle: T) -> Option<usize>
where
    T: SimdElement + SimdPartialEq,
{
    let (prefix, simd_data, suffix) = arr.as_simd::<N>();

    // handle prefix ... (not relevant for the example)

    // SIMD
    let simd_needle = Simd::splat(needle);
    for (chunk_idx, chunk) in simd_data.iter().enumerate() {
        // no method named `simd_eq` found for reference `&Simd<T, 16>`
        let mask = chunk.simd_eq(simd_needle).to_bitmask();
        if mask != 0 {
            return Some(prefix.len() + (chunk_idx * N) + (mask.trailing_zeros() as usize));
        }
    }
    // handle suffix ... (not relevant for the example)
}

If you had searched the documentation for SimdPartialEq, you'd have found out that it's implemented by SIMD vectors, not the individual elements.

This compiles.

1 Like

Note that your version accidentally prohibits use with unsigned integers, because you require the mask elements to be the same type as the SIMD elements. Changing the where clause to:

where
    T: SimdElement,
    Simd<T, N>: SimdPartialEq<Mask = Mask<T::Mask, N>>,

fixes this, by asking the compiler to choose the mask type selected by the SimdElement implementation.

1 Like