enhance: memcpy_nontemporal with alignment check

Author: sundanc

src/simd.rs

@@ -911,6 +911,23 @@ unsafe fn memcpy_nontemporal_sse2(dst: *mut u8, src: *const u8, len: usize) {
     }
 }
 
+#[cfg(all(feature = "simd", target_arch = "x86_64"))]
+#[target_feature(enable = "avx2")]
+unsafe fn memcpy_nontemporal_avx2(dst: *mut u8, src: *const u8, len: usize) {
+    unsafe {
+        let mut i = 0;
+        while i + 32 <= len {
+            let v = _mm256_loadu_si256(src.add(i) as *const __m256i);
+            _mm256_stream_si256(dst.add(i) as *mut __m256i, v);
+            i += 32;
+        }
+        _mm_sfence();
+        if i < len {
+            std::ptr::copy_nonoverlapping(src.add(i), dst.add(i), len - i);
+        }
+    }
+}
+
 #[cfg(all(feature = "simd", target_arch = "x86_64"))]
 #[inline(always)]
 /// # Safety
@@ -921,14 +938,46 @@ unsafe fn memcpy_nontemporal_sse2(dst: *mut u8, src: *const u8, len: usize) {
 /// - the regions do not overlap
 pub unsafe fn memcpy_nontemporal(dst: *mut u8, src: *const u8, len: usize) {
     unsafe {
-        if is_simd_available() && len >= 64 {
+        if is_simd_available() && len >= 64 && (dst as usize).is_multiple_of(16) {
             memcpy_nontemporal_sse2(dst, src, len);
             return;
         }
         std::ptr::copy_nonoverlapping(src, dst, len);
     }
 }
 
+pub fn safe_memcpy_nontemporal(dst: &mut [u8], src: &[u8]) -> Result<(), &'static str> {
+    if dst.len() != src.len() {
+        return Err("length mismatch");
+    }
+    let len = dst.len();
+    let src_ptr = src.as_ptr() as usize;
+    let dst_ptr = dst.as_mut_ptr() as usize;
+    let src_end = src_ptr.checked_add(len).ok_or("overflow")?;
+    let dst_end = dst_ptr.checked_add(len).ok_or("overflow")?;
+    if src_ptr < dst_end && dst_ptr < src_end {
+        return Err("overlap");
+    }
+
+    #[cfg(all(feature = "simd", target_arch = "x86_64"))]
+    {
+        let info = simd_info();
+        if info.avx2 && is_simd_available() && len >= 64 && (dst_ptr).is_multiple_of(32) {
+            unsafe { memcpy_nontemporal_avx2(dst.as_mut_ptr(), src.as_ptr(), len) };
+            return Ok(());
+        }
+        if is_simd_available() && len >= 64 && (dst_ptr).is_multiple_of(16) {
+            unsafe { memcpy_nontemporal_sse2(dst.as_mut_ptr(), src.as_ptr(), len) };
+            return Ok(());
+        }
+    }
+
+    unsafe {
+        std::ptr::copy_nonoverlapping(src.as_ptr(), dst.as_mut_ptr(), len);
+    }
+    Ok(())
+}
+
 #[cfg(not(all(feature = "simd", target_arch = "x86_64")))]
 #[inline(always)]
 /// # Safety
@@ -1965,6 +2014,41 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_memcpy_nontemporal_unaligned_large() {
+        use std::alloc::{Layout, alloc, dealloc};
+        let layout = Layout::from_size_align(136, 8).unwrap();
+        let base_ptr = unsafe { alloc(layout) };
+        let dst_ptr = unsafe { base_ptr.add(8) };
+        let src = [99; 128];
+        unsafe {
+            memcpy_nontemporal(dst_ptr, src.as_ptr(), 128);
+            for i in 0..128 {
+                assert_eq!(*dst_ptr.add(i), 99);
+            }
+            dealloc(base_ptr, layout);
+        }
+    }
+
+    #[test]
+    fn test_safe_memcpy_nontemporal_slice() {
+        let src = [7u8; 128];
+        let mut dst = [0u8; 128];
+        assert!(safe_memcpy_nontemporal(&mut dst, &src).is_ok());
+        assert_eq!(dst, src);
+    }
+
+    #[test]
+    fn test_safe_memcpy_nontemporal_overlap() {
+        let mut buf = vec![0u8; 256];
+        for (i, item) in buf.iter_mut().enumerate().take(256) {
+            *item = i as u8;
+        }
+        let src = unsafe { std::slice::from_raw_parts(buf.as_ptr(), 128) };
+        let dst = unsafe { std::slice::from_raw_parts_mut(buf.as_mut_ptr().add(64), 128) };
+        assert!(safe_memcpy_nontemporal(dst, src).is_err());
+    }
+
     #[test]
     fn test_simd_level_fallback() {
         let info = SimdInfo {