edit/simd/

memset.rs

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

//! `memchr` for arbitrary sizes (1/2/4/8 bytes).
//!
//! Clang calls the C `memset` function only for byte-sized types (or 0 fills).
//! We however need to fill other types as well. For that, clang generates
//! SIMD loops under higher optimization levels. With `-Os` however, it only
//! generates a trivial loop which is too slow for our needs.
//!
//! This implementation uses SWAR to only have a single implementation for all
//! 4 sizes: By duplicating smaller types into a larger `u64` register we can
//! treat all sizes as if they were `u64`. The only thing we need to take care
//! of is the tail end of the array, which needs to write 0-7 additional bytes.

use std::mem;

/// A marker trait for types that are safe to `memset`.
///
/// # Safety
///
/// Just like with C's `memset`, bad things happen
/// if you use this with non-trivial types.
pub unsafe trait MemsetSafe: Copy {}

unsafe impl MemsetSafe for u8 {}
unsafe impl MemsetSafe for u16 {}
unsafe impl MemsetSafe for u32 {}
unsafe impl MemsetSafe for u64 {}
unsafe impl MemsetSafe for usize {}

unsafe impl MemsetSafe for i8 {}
unsafe impl MemsetSafe for i16 {}
unsafe impl MemsetSafe for i32 {}
unsafe impl MemsetSafe for i64 {}
unsafe impl MemsetSafe for isize {}

/// Fills a slice with the given value.
#[inline]
pub fn memset<T: MemsetSafe>(dst: &mut [T], val: T) {
    unsafe {
        match mem::size_of::<T>() {
            1 => {
                // LLVM will compile this to a call to `memset`,
                // which hopefully should be better optimized than my code.
                let beg = dst.as_mut_ptr();
                let val = mem::transmute_copy::<_, u8>(&val);
                beg.write_bytes(val, dst.len());
            }
            2 => {
                let beg = dst.as_mut_ptr();
                let end = beg.add(dst.len());
                let val = mem::transmute_copy::<_, u16>(&val);
                memset_raw(beg as *mut u8, end as *mut u8, val as u64 * 0x0001000100010001);
            }
            4 => {
                let beg = dst.as_mut_ptr();
                let end = beg.add(dst.len());
                let val = mem::transmute_copy::<_, u32>(&val);
                memset_raw(beg as *mut u8, end as *mut u8, val as u64 * 0x0000000100000001);
            }
            8 => {
                let beg = dst.as_mut_ptr();
                let end = beg.add(dst.len());
                let val = mem::transmute_copy::<_, u64>(&val);
                memset_raw(beg as *mut u8, end as *mut u8, val);
            }
            _ => unreachable!(),
        }
    }
}

#[inline]
fn memset_raw(beg: *mut u8, end: *mut u8, val: u64) {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    return unsafe { MEMSET_DISPATCH(beg, end, val) };

    #[cfg(target_arch = "aarch64")]
    return unsafe { memset_neon(beg, end, val) };

    #[allow(unreachable_code)]
    return unsafe { memset_fallback(beg, end, val) };
}

#[inline(never)]
unsafe fn memset_fallback(mut beg: *mut u8, end: *mut u8, val: u64) {
    unsafe {
        let mut remaining = end.offset_from_unsigned(beg);

        while remaining >= 8 {
            (beg as *mut u64).write_unaligned(val);
            beg = beg.add(8);
            remaining -= 8;
        }

        if remaining >= 4 {
            // 4-7 bytes remaining
            (beg as *mut u32).write_unaligned(val as u32);
            (end.sub(4) as *mut u32).write_unaligned(val as u32);
        } else if remaining >= 2 {
            // 2-3 bytes remaining
            (beg as *mut u16).write_unaligned(val as u16);
            (end.sub(2) as *mut u16).write_unaligned(val as u16);
        } else if remaining >= 1 {
            // 1 byte remaining
            beg.write(val as u8);
        }
    }
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
static mut MEMSET_DISPATCH: unsafe fn(beg: *mut u8, end: *mut u8, val: u64) = memset_dispatch;

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
fn memset_dispatch(beg: *mut u8, end: *mut u8, val: u64) {
    let func = if is_x86_feature_detected!("avx2") { memset_avx2 } else { memset_sse2 };
    unsafe { MEMSET_DISPATCH = func };
    unsafe { func(beg, end, val) }
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "sse2")]
unsafe fn memset_sse2(mut beg: *mut u8, end: *mut u8, val: u64) {
    unsafe {
        #[cfg(target_arch = "x86")]
        use std::arch::x86::*;
        #[cfg(target_arch = "x86_64")]
        use std::arch::x86_64::*;

        let mut remaining = end.offset_from_unsigned(beg);

        if remaining >= 16 {
            let fill = _mm_set1_epi64x(val as i64);

            while remaining >= 32 {
                _mm_storeu_si128(beg as *mut _, fill);
                _mm_storeu_si128(beg.add(16) as *mut _, fill);

                beg = beg.add(32);
                remaining -= 32;
            }

            if remaining >= 16 {
                // 16-31 bytes remaining
                _mm_storeu_si128(beg as *mut _, fill);
                _mm_storeu_si128(end.sub(16) as *mut _, fill);
                return;
            }
        }

        if remaining >= 8 {
            // 8-15 bytes remaining
            (beg as *mut u64).write_unaligned(val);
            (end.sub(8) as *mut u64).write_unaligned(val);
        } else if remaining >= 4 {
            // 4-7 bytes remaining
            (beg as *mut u32).write_unaligned(val as u32);
            (end.sub(4) as *mut u32).write_unaligned(val as u32);
        } else if remaining >= 2 {
            // 2-3 bytes remaining
            (beg as *mut u16).write_unaligned(val as u16);
            (end.sub(2) as *mut u16).write_unaligned(val as u16);
        } else if remaining >= 1 {
            // 1 byte remaining
            beg.write(val as u8);
        }
    }
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "avx2")]
fn memset_avx2(mut beg: *mut u8, end: *mut u8, val: u64) {
    unsafe {
        #[cfg(target_arch = "x86")]
        use std::arch::x86::*;
        #[cfg(target_arch = "x86_64")]
        use std::arch::x86_64::*;
        use std::hint::black_box;

        let mut remaining = end.offset_from_unsigned(beg);

        if remaining >= 128 {
            let fill = _mm256_set1_epi64x(val as i64);

            loop {
                _mm256_storeu_si256(beg as *mut _, fill);
                _mm256_storeu_si256(beg.add(32) as *mut _, fill);
                _mm256_storeu_si256(beg.add(64) as *mut _, fill);
                _mm256_storeu_si256(beg.add(96) as *mut _, fill);

                beg = beg.add(128);
                remaining -= 128;
                if remaining < 128 {
                    break;
                }
            }
        }

        if remaining >= 16 {
            let fill = _mm_set1_epi64x(val as i64);

            loop {
                // LLVM is _very_ eager to unroll loops. In the absence of an unroll attribute, black_box does the job.
                // Note that this must not be applied to the intrinsic parameters, as they're otherwise misoptimized.
                #[allow(clippy::unit_arg)]
                black_box(_mm_storeu_si128(beg as *mut _, fill));

                beg = beg.add(16);
                remaining -= 16;
                if remaining < 16 {
                    break;
                }
            }
        }

        // `remaining` is between 0 and 15 at this point.
        // By overlapping the stores we can write all of them in at most 2 stores. This approach
        // can be seen in various libraries, such as wyhash which uses it for loading data in `wyr3`.
        if remaining >= 8 {
            // 8-15 bytes
            (beg as *mut u64).write_unaligned(val);
            (end.sub(8) as *mut u64).write_unaligned(val);
        } else if remaining >= 4 {
            // 4-7 bytes
            (beg as *mut u32).write_unaligned(val as u32);
            (end.sub(4) as *mut u32).write_unaligned(val as u32);
        } else if remaining >= 2 {
            // 2-3 bytes
            (beg as *mut u16).write_unaligned(val as u16);
            (end.sub(2) as *mut u16).write_unaligned(val as u16);
        } else if remaining >= 1 {
            // 1 byte
            beg.write(val as u8);
        }
    }
}

#[cfg(target_arch = "aarch64")]
unsafe fn memset_neon(mut beg: *mut u8, end: *mut u8, val: u64) {
    unsafe {
        use std::arch::aarch64::*;
        let mut remaining = end.offset_from_unsigned(beg);

        if remaining >= 32 {
            let fill = vdupq_n_u64(val);

            loop {
                // Compiles to a single `stp` instruction.
                vst1q_u64(beg as *mut _, fill);
                vst1q_u64(beg.add(16) as *mut _, fill);

                beg = beg.add(32);
                remaining -= 32;
                if remaining < 32 {
                    break;
                }
            }
        }

        if remaining >= 16 {
            // 16-31 bytes remaining
            let fill = vdupq_n_u64(val);
            vst1q_u64(beg as *mut _, fill);
            vst1q_u64(end.sub(16) as *mut _, fill);
        } else if remaining >= 8 {
            // 8-15 bytes remaining
            (beg as *mut u64).write_unaligned(val);
            (end.sub(8) as *mut u64).write_unaligned(val);
        } else if remaining >= 4 {
            // 4-7 bytes remaining
            (beg as *mut u32).write_unaligned(val as u32);
            (end.sub(4) as *mut u32).write_unaligned(val as u32);
        } else if remaining >= 2 {
            // 2-3 bytes remaining
            (beg as *mut u16).write_unaligned(val as u16);
            (end.sub(2) as *mut u16).write_unaligned(val as u16);
        } else if remaining >= 1 {
            // 1 byte remaining
            beg.write(val as u8);
        }
    }
}

#[cfg(test)]
mod tests {
    use std::fmt;
    use std::ops::Not;

    use super::*;

    fn check_memset<T>(val: T, len: usize)
    where
        T: MemsetSafe + Not<Output = T> + PartialEq + fmt::Debug,
    {
        let mut buf = vec![!val; len];
        memset(&mut buf, val);
        assert!(buf.iter().all(|&x| x == val));
    }

    #[test]
    fn test_memset_empty() {
        check_memset(0u8, 0);
        check_memset(0u16, 0);
        check_memset(0u32, 0);
        check_memset(0u64, 0);
    }

    #[test]
    fn test_memset_single() {
        check_memset(0u8, 1);
        check_memset(0xFFu8, 1);
        check_memset(0xABu16, 1);
        check_memset(0x12345678u32, 1);
        check_memset(0xDEADBEEFu64, 1);
    }

    #[test]
    fn test_memset_small() {
        for &len in &[2, 3, 4, 5, 7, 8, 9] {
            check_memset(0xAAu8, len);
            check_memset(0xBEEFu16, len);
            check_memset(0xCAFEBABEu32, len);
            check_memset(0x1234567890ABCDEFu64, len);
        }
    }

    #[test]
    fn test_memset_large() {
        check_memset(0u8, 1000);
        check_memset(0xFFu8, 1024);
        check_memset(0xBEEFu16, 512);
        check_memset(0xCAFEBABEu32, 256);
        check_memset(0x1234567890ABCDEFu64, 128);
    }

    #[test]
    fn test_memset_various_values() {
        check_memset(0u8, 17);
        check_memset(0x7Fu8, 17);
        check_memset(0x8001u16, 17);
        check_memset(0xFFFFFFFFu32, 17);
        check_memset(0x8000000000000001u64, 17);
    }

    #[test]
    fn test_memset_signed_types() {
        check_memset(-1i8, 8);
        check_memset(-2i16, 8);
        check_memset(-3i32, 8);
        check_memset(-4i64, 8);
        check_memset(-5isize, 8);
    }

    #[test]
    fn test_memset_usize_isize() {
        check_memset(0usize, 4);
        check_memset(usize::MAX, 4);
        check_memset(0isize, 4);
        check_memset(isize::MIN, 4);
    }

    #[test]
    fn test_memset_alignment() {
        // Check that memset works for slices not aligned to 8 bytes
        let mut buf = [0u8; 15];
        for offset in 0..8 {
            let slice = &mut buf[offset..(offset + 7)];
            memset(slice, 0x5A);
            assert!(slice.iter().all(|&x| x == 0x5A));
        }
    }
}