rustalign_simd/

sse_wrapper.rs

//! SSE wrapper types and operations
//!
//! This module provides wrappers around SIMD intrinsics that match
//! the C++ sse_wrap.h interface, with runtime dispatch for SSE2/SSE4.2/AVX2.

// CPU feature detection imports - only available on x86_64

/// SIMD register type wrapper
///
/// This abstracts over __m128i (SSE2) and __m256i (AVX2) types,
/// matching the C++ SSERegI typedef.
#[derive(Copy, Clone, Debug)]
pub struct SseReg {
    /// Internal storage - uses u8 array for portability
    data: [u8; 32],
}

impl SseReg {
    /// Create a new zero register
    pub fn zero() -> Self {
        Self { data: [0; 32] }
    }

    /// Set all 16-bit elements to the same value
    pub fn set1_epi16(value: i16) -> Self {
        let mut result = Self::zero();
        let bytes = value.to_le_bytes();
        for i in 0..16 {
            result.data[i * 2] = bytes[0];
            result.data[i * 2 + 1] = bytes[1];
        }
        result
    }

    /// Load aligned 128-bit data
    #[cfg(target_arch = "x86_64")]
    pub fn load_128(ptr: &[u8; 16]) -> Self {
        use std::arch::x86_64::*;
        unsafe {
            let reg = _mm_loadu_si128(ptr.as_ptr() as *const __m128i);
            let mut result = Self::zero();
            std::ptr::copy_nonoverlapping(
                &reg as *const _ as *const u8,
                result.data.as_mut_ptr(),
                16,
            );
            result
        }
    }

    /// Load unaligned data
    #[cfg(not(target_arch = "x86_64"))]
    pub fn load_128(ptr: &[u8; 16]) -> Self {
        let mut result = Self::zero();
        result.data[..16].copy_from_slice(ptr);
        result
    }

    /// Store aligned 128-bit data
    #[cfg(target_arch = "x86_64")]
    pub fn store_128(&self, ptr: &mut [u8; 16]) {
        use std::arch::x86_64::*;
        unsafe {
            let reg = _mm_loadu_si128(self.data.as_ptr() as *const __m128i);
            _mm_storeu_si128(ptr.as_mut_ptr() as *mut __m128i, reg);
        }
    }

    /// Store data (non-x86 fallback)
    #[cfg(not(target_arch = "x86_64"))]
    pub fn store_128(&self, ptr: &mut [u8; 16]) {
        ptr.copy_from_slice(&self.data[..16]);
    }

    /// Extract a 16-bit element by index
    pub fn extract_epi16(&self, index: i32) -> i16 {
        let idx = index as usize;
        let byte_idx = idx * 2;
        i16::from_le_bytes([self.data[byte_idx], self.data[byte_idx + 1]])
    }

    /// Compare 16-bit integers for equality
    #[cfg(target_arch = "x86_64")]
    pub fn cmpeq_epi16(&self, other: &Self) -> Self {
        use std::arch::x86_64::*;
        unsafe {
            let a = _mm_loadu_si128(self.data.as_ptr() as *const __m128i);
            let b = _mm_loadu_si128(other.data.as_ptr() as *const __m128i);
            let result = _mm_cmpeq_epi16(a, b);
            let mut out = Self::zero();
            std::ptr::copy_nonoverlapping(
                &result as *const _ as *const u8,
                out.data.as_mut_ptr(),
                16,
            );
            out
        }
    }

    /// Compare 16-bit integers for equality (scalar fallback)
    #[cfg(not(target_arch = "x86_64"))]
    pub fn cmpeq_epi16(&self, other: &Self) -> Self {
        let mut result = Self::zero();
        for i in 0..8 {
            let a = self.extract_epi16(i as i32);
            let b = other.extract_epi16(i as i32);
            let val = if a == b { 0xffff } else { 0 };
            result.data[i * 2] = (val & 0xff) as u8;
            result.data[i * 2 + 1] = ((val >> 8) & 0xff) as u8;
        }
        result
    }

    /// Add signed 16-bit integers with saturation
    #[cfg(target_arch = "x86_64")]
    pub fn adds_epi16(&self, other: &Self) -> Self {
        use std::arch::x86_64::*;
        unsafe {
            let a = _mm_loadu_si128(self.data.as_ptr() as *const __m128i);
            let b = _mm_loadu_si128(other.data.as_ptr() as *const __m128i);
            let result = _mm_adds_epi16(a, b);
            let mut out = Self::zero();
            std::ptr::copy_nonoverlapping(
                &result as *const _ as *const u8,
                out.data.as_mut_ptr(),
                16,
            );
            out
        }
    }

    /// Add signed 16-bit integers with saturation (scalar fallback)
    #[cfg(not(target_arch = "x86_64"))]
    pub fn adds_epi16(&self, other: &Self) -> Self {
        let mut result = Self::zero();
        for i in 0..8 {
            let a = self.extract_epi16(i as i32);
            let b = other.extract_epi16(i as i32);
            let val = a.saturating_add(b);
            result.data[i * 2] = (val & 0xff) as u8;
            result.data[i * 2 + 1] = ((val >> 8) & 0xff) as u8;
        }
        result
    }

    /// Maximum of signed 16-bit integers
    #[cfg(target_arch = "x86_64")]
    pub fn max_epi16(&self, other: &Self) -> Self {
        use std::arch::x86_64::*;
        unsafe {
            let a = _mm_loadu_si128(self.data.as_ptr() as *const __m128i);
            let b = _mm_loadu_si128(other.data.as_ptr() as *const __m128i);
            let result = _mm_max_epi16(a, b);
            let mut out = Self::zero();
            std::ptr::copy_nonoverlapping(
                &result as *const _ as *const u8,
                out.data.as_mut_ptr(),
                16,
            );
            out
        }
    }

    /// Maximum of signed 16-bit integers (scalar fallback)
    #[cfg(not(target_arch = "x86_64"))]
    pub fn max_epi16(&self, other: &Self) -> Self {
        let mut result = Self::zero();
        for i in 0..8 {
            let a = self.extract_epi16(i as i32);
            let b = other.extract_epi16(i as i32);
            let val = a.max(b);
            result.data[i * 2] = (val & 0xff) as u8;
            result.data[i * 2 + 1] = ((val >> 8) & 0xff) as u8;
        }
        result
    }

    /// Compute horizontal maximum of i16 elements
    ///
    /// This matches the C++ sse_max_score_i16 macro.
    pub fn hmax_epi16(&self) -> i16 {
        let mut result = self.extract_epi16(0);
        for i in 1..8 {
            result = result.max(self.extract_epi16(i));
        }
        result
    }

    /// Bitwise XOR
    #[cfg(target_arch = "x86_64")]
    pub fn xor(&self, other: &Self) -> Self {
        use std::arch::x86_64::*;
        unsafe {
            let a = _mm_loadu_si128(self.data.as_ptr() as *const __m128i);
            let b = _mm_loadu_si128(other.data.as_ptr() as *const __m128i);
            let result = _mm_xor_si128(a, b);
            let mut out = Self::zero();
            std::ptr::copy_nonoverlapping(
                &result as *const _ as *const u8,
                out.data.as_mut_ptr(),
                16,
            );
            out
        }
    }

    /// Bitwise XOR (scalar fallback)
    #[cfg(not(target_arch = "x86_64"))]
    pub fn xor(&self, other: &Self) -> Self {
        let mut result = Self::zero();
        for i in 0..16 {
            result.data[i] = self.data[i] ^ other.data[i];
        }
        result
    }
}

impl Default for SseReg {
    fn default() -> Self {
        Self::zero()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sse_reg_zero() {
        let reg = SseReg::zero();
        assert_eq!(reg.extract_epi16(0), 0);
        assert_eq!(reg.extract_epi16(7), 0);
    }

    #[test]
    fn test_sse_reg_set1() {
        let reg = SseReg::set1_epi16(-1);
        for i in 0..8 {
            assert_eq!(reg.extract_epi16(i), -1);
        }
    }

    #[test]
    fn test_sse_reg_set1_42() {
        let reg = SseReg::set1_epi16(42);
        for i in 0..8 {
            assert_eq!(reg.extract_epi16(i), 42);
        }
    }

    #[test]
    fn test_sse_cmpeq_epi16() {
        let a = SseReg::set1_epi16(42);
        let b = SseReg::set1_epi16(42);
        let result = a.cmpeq_epi16(&b);
        // All elements should be equal (0xffff)
        for i in 0..8 {
            assert_eq!(result.extract_epi16(i), -1);
        }
    }

    #[test]
    fn test_sse_cmpeq_epi16_not_equal() {
        let a = SseReg::set1_epi16(42);
        let b = SseReg::set1_epi16(43);
        let result = a.cmpeq_epi16(&b);
        // All elements should be not equal (0)
        for i in 0..8 {
            assert_eq!(result.extract_epi16(i), 0);
        }
    }

    #[test]
    fn test_sse_adds_epi16() {
        let a = SseReg::set1_epi16(100);
        let b = SseReg::set1_epi16(50);
        let result = a.adds_epi16(&b);
        for i in 0..8 {
            assert_eq!(result.extract_epi16(i), 150);
        }
    }

    #[test]
    fn test_sse_adds_epi16_saturate() {
        let a = SseReg::set1_epi16(30000);
        let b = SseReg::set1_epi16(10000);
        let result = a.adds_epi16(&b);
        // Should saturate at i16::MAX = 32767
        for i in 0..8 {
            assert_eq!(result.extract_epi16(i), 32767);
        }
    }

    #[test]
    fn test_sse_max_epi16() {
        let a = SseReg::set1_epi16(42);
        let b = SseReg::set1_epi16(100);
        let result = a.max_epi16(&b);
        for i in 0..8 {
            assert_eq!(result.extract_epi16(i), 100);
        }
    }

    #[test]
    fn test_sse_hmax_epi16() {
        let mut data = [0u8; 32];
        for i in 0..8 {
            data[i * 2] = ((i * 10) & 0xff) as u8;
            data[i * 2 + 1] = (((i * 10) >> 8) & 0xff) as u8;
        }
        let reg = SseReg { data };
        assert_eq!(reg.hmax_epi16(), 70); // max of 0, 10, 20, ..., 70
    }

    #[test]
    fn test_sse_xor() {
        let a = SseReg::set1_epi16(-1);
        let b = SseReg::set1_epi16(-1);
        let result = a.xor(&b);
        for i in 0..8 {
            assert_eq!(result.extract_epi16(i), 0);
        }
    }

    #[test]
    fn test_sse_load_store() {
        let input: [u8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];
        let reg = SseReg::load_128(&input);
        let mut output = [0u8; 16];
        reg.store_128(&mut output);
        assert_eq!(input, output);
    }
}