1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181
// 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)
}