src/packed/vector.rs

// This file contains a set of fairly generic utility functions when working
// with SIMD vectors.
//
// SAFETY: All of the routines below are unsafe to call because they assume
// the necessary CPU target features in order to use particular vendor
// intrinsics. Calling these routines when the underlying CPU does not support
// the appropriate target features is NOT safe. Callers must ensure this
// themselves.
//
// Note that it may not look like this safety invariant is being upheld when
// these routines are called. Namely, the CPU feature check is typically pretty
// far away from when these routines are used. Instead, we rely on the fact
// that certain types serve as a guaranteed receipt that pertinent target
// features are enabled. For example, the only way TeddySlim3Mask256 can be
// constructed is if the AVX2 CPU feature is available. Thus, any code running
// inside of TeddySlim3Mask256 can use any of the functions below without any
// additional checks: its very existence *is* the check.

use std::arch::x86_64::*;

/// Shift `a` to the left by two bytes (removing its two most significant
/// bytes), and concatenate it with the the two most significant bytes of `b`.
#[target_feature(enable = "avx2")]
pub unsafe fn alignr256_14(a: __m256i, b: __m256i) -> __m256i {
    // Credit goes to jneem for figuring this out:
    // https://github.com/jneem/teddy/blob/9ab5e899ad6ef6911aecd3cf1033f1abe6e1f66c/src/x86/teddy_simd.rs#L145-L184
    //
    // TL;DR avx2's PALIGNR instruction is actually just two 128-bit PALIGNR
    // instructions, which is not what we want, so we need to do some extra
    // shuffling.

    // This permute gives us the low 16 bytes of a concatenated with the high
    // 16 bytes of b, in order of most significant to least significant. So
    // `v = a[15:0] b[31:16]`.
    let v = _mm256_permute2x128_si256(b, a, 0x21);
    // This effectively does this (where we deal in terms of byte-indexing
    // and byte-shifting, and use inclusive ranges):
    //
    //   ret[15:0]  := ((a[15:0] << 16) | v[15:0]) >> 14
    //               = ((a[15:0] << 16) | b[31:16]) >> 14
    //   ret[31:16] := ((a[31:16] << 16) | v[31:16]) >> 14
    //               = ((a[31:16] << 16) | a[15:0]) >> 14
    //
    // Which therefore results in:
    //
    //   ret[31:0]  := a[29:16] a[15:14] a[13:0] b[31:30]
    //
    // The end result is that we've effectively done this:
    //
    //   (a << 2) | (b >> 30)
    //
    // When `A` and `B` are strings---where the beginning of the string is in
    // the least significant bits---we effectively result in the following
    // semantic operation:
    //
    //   (A >> 2) | (B << 30)
    //
    // The reversal being attributed to the fact that we are in little-endian.
    _mm256_alignr_epi8(a, v, 14)
}

/// Shift `a` to the left by one byte (removing its most significant byte), and
/// concatenate it with the the most significant byte of `b`.
#[target_feature(enable = "avx2")]
pub unsafe fn alignr256_15(a: __m256i, b: __m256i) -> __m256i {
    // For explanation, see alignr256_14.
    let v = _mm256_permute2x128_si256(b, a, 0x21);
    _mm256_alignr_epi8(a, v, 15)
}

/// Unpack the given 128-bit vector into its 64-bit components. The first
/// element of the array returned corresponds to the least significant 64-bit
/// lane in `a`.
#[target_feature(enable = "ssse3")]
pub unsafe fn unpack64x128(a: __m128i) -> [u64; 2] {
    [
        _mm_cvtsi128_si64(a) as u64,
        _mm_cvtsi128_si64(_mm_srli_si128(a, 8)) as u64,
    ]
}

/// Unpack the given 256-bit vector into its 64-bit components. The first
/// element of the array returned corresponds to the least significant 64-bit
/// lane in `a`.
#[target_feature(enable = "avx2")]
pub unsafe fn unpack64x256(a: __m256i) -> [u64; 4] {
    // Using transmute here is precisely equivalent, but actually slower. It's
    // not quite clear why.
    let lo = _mm256_extracti128_si256(a, 0);
    let hi = _mm256_extracti128_si256(a, 1);
    [
        _mm_cvtsi128_si64(lo) as u64,
        _mm_cvtsi128_si64(_mm_srli_si128(lo, 8)) as u64,
        _mm_cvtsi128_si64(hi) as u64,
        _mm_cvtsi128_si64(_mm_srli_si128(hi, 8)) as u64,
    ]
}

/// Unpack the low 128-bits of `a` and `b`, and return them as 4 64-bit
/// integers.
///
/// More precisely, if a = a4 a3 a2 a1 and b = b4 b3 b2 b1, where each element
/// is a 64-bit integer and a1/b1 correspond to the least significant 64 bits,
/// then the return value is `b2 b1 a2 a1`.
#[target_feature(enable = "avx2")]
pub unsafe fn unpacklo64x256(a: __m256i, b: __m256i) -> [u64; 4] {
    let lo = _mm256_castsi256_si128(a);
    let hi = _mm256_castsi256_si128(b);
    [
        _mm_cvtsi128_si64(lo) as u64,
        _mm_cvtsi128_si64(_mm_srli_si128(lo, 8)) as u64,
        _mm_cvtsi128_si64(hi) as u64,
        _mm_cvtsi128_si64(_mm_srli_si128(hi, 8)) as u64,
    ]
}

/// Returns true if and only if all bits in the given 128-bit vector are 0.
#[target_feature(enable = "ssse3")]
pub unsafe fn is_all_zeroes128(a: __m128i) -> bool {
    let cmp = _mm_cmpeq_epi8(a, zeroes128());
    _mm_movemask_epi8(cmp) as u32 == 0xFFFF
}

/// Returns true if and only if all bits in the given 256-bit vector are 0.
#[target_feature(enable = "avx2")]
pub unsafe fn is_all_zeroes256(a: __m256i) -> bool {
    let cmp = _mm256_cmpeq_epi8(a, zeroes256());
    _mm256_movemask_epi8(cmp) as u32 == 0xFFFFFFFF
}

/// Load a 128-bit vector from slice at the given position. The slice does
/// not need to be unaligned.
///
/// Since this code assumes little-endian (there is no big-endian x86), the
/// bytes starting in `slice[at..]` will be at the least significant bits of
/// the returned vector. This is important for the surrounding code, since for
/// example, shifting the resulting vector right is equivalent to logically
/// shifting the bytes in `slice` left.
#[target_feature(enable = "sse2")]
pub unsafe fn loadu128(slice: &[u8], at: usize) -> __m128i {
    let ptr = slice.get_unchecked(at..).as_ptr();
    _mm_loadu_si128(ptr as *const u8 as *const __m128i)
}

/// Load a 256-bit vector from slice at the given position. The slice does
/// not need to be unaligned.
///
/// Since this code assumes little-endian (there is no big-endian x86), the
/// bytes starting in `slice[at..]` will be at the least significant bits of
/// the returned vector. This is important for the surrounding code, since for
/// example, shifting the resulting vector right is equivalent to logically
/// shifting the bytes in `slice` left.
#[target_feature(enable = "avx2")]
pub unsafe fn loadu256(slice: &[u8], at: usize) -> __m256i {
    let ptr = slice.get_unchecked(at..).as_ptr();
    _mm256_loadu_si256(ptr as *const u8 as *const __m256i)
}

/// Returns a 128-bit vector with all bits set to 0.
#[target_feature(enable = "sse2")]
pub unsafe fn zeroes128() -> __m128i {
    _mm_set1_epi8(0)
}

/// Returns a 256-bit vector with all bits set to 0.
#[target_feature(enable = "avx2")]
pub unsafe fn zeroes256() -> __m256i {
    _mm256_set1_epi8(0)
}

/// Returns a 128-bit vector with all bits set to 1.
#[target_feature(enable = "sse2")]
pub unsafe fn ones128() -> __m128i {
    _mm_set1_epi8(0xFF as u8 as i8)
}

/// Returns a 256-bit vector with all bits set to 1.
#[target_feature(enable = "avx2")]
pub unsafe fn ones256() -> __m256i {
    _mm256_set1_epi8(0xFF as u8 as i8)
}