ruprim_host/simd/
aligned.rs

1//! SIMD-aligned memory allocation utilities.
2//!
3//! Provides aligned vectors for optimal SIMD performance. Uses 64-byte alignment
4//! to support all SIMD instruction sets (NEON: 16B, AVX2: 32B, AVX-512: 64B).
5
6use aligned_vec::{AVec, ConstAlign};
7use alloc::vec::Vec;
8
9/// Alignment for SIMD operations (64 bytes = AVX-512 width).
10/// Also works for NEON (16B) and AVX2 (32B) since 64 is a multiple of both.
11pub const SIMD_ALIGN: usize = 64;
12
13/// Type alias for 64-byte aligned vector.
14pub type AlignedVec<T> = AVec<T, ConstAlign<SIMD_ALIGN>>;
15
16/// Allocate a SIMD-aligned vector with the given capacity.
17#[inline]
18pub fn alloc_aligned<T>(capacity: usize) -> AlignedVec<T> {
19    AVec::with_capacity(SIMD_ALIGN, capacity)
20}
21
22/// Allocate a SIMD-aligned vector filled with zeros.
23#[inline]
24pub fn alloc_aligned_zeroed<T: bytemuck::Zeroable + Clone>(len: usize) -> AlignedVec<T> {
25    let mut vec = AVec::with_capacity(SIMD_ALIGN, len);
26    vec.resize(len, T::zeroed());
27    vec
28}
29
30/// Convert an aligned vector to a regular Vec.
31///
32/// This may copy the data if the alignment requirements differ.
33#[inline]
34pub fn to_vec<T: Clone>(aligned: AlignedVec<T>) -> Vec<T> {
35    aligned.to_vec()
36}
37
38/// Convert a slice to an aligned vector (copies data).
39#[inline]
40pub fn from_slice<T: Clone>(slice: &[T]) -> AlignedVec<T> {
41    let mut vec = AVec::with_capacity(SIMD_ALIGN, slice.len());
42    vec.extend_from_slice(slice);
43    vec
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn test_alignment() {
52        let vec: AlignedVec<f32> = alloc_aligned(1000);
53        assert_eq!(vec.as_ptr() as usize % SIMD_ALIGN, 0);
54    }
55
56    #[test]
57    fn test_zeroed() {
58        let vec: AlignedVec<f32> = alloc_aligned_zeroed(100);
59        assert_eq!(vec.len(), 100);
60        assert!(vec.iter().all(|&x| x == 0.0));
61    }
62
63    #[test]
64    fn test_from_slice() {
65        let data = [1.0f32, 2.0, 3.0, 4.0];
66        let aligned = from_slice(&data);
67        assert_eq!(aligned.as_ptr() as usize % SIMD_ALIGN, 0);
68        assert_eq!(&aligned[..], &data[..]);
69    }
70}
ruprim_host/simd/aligned.rs

ruprim_host/simd/
aligned.rs
