prefetch-index 0.2.1

//! A small crate to prefetch an element of an array.
//!
//! Provides [`prefetch_index`] and [`prefetch_index_nta`] to prefetch the cache
//! line containing `slice[index]`.

#![cfg_attr(
    all(feature = "aarch64", target_arch = "aarch64"),
    feature(stdarch_aarch64_prefetch)
)]

/// Prefetches the cache line containing (the first byte of) `data[index]` into
/// _all_ levels of the cache.
///
/// On x86/x86_64, this uses _mm_prefetch which only requires the (commonly available) SSE feature.
///
/// On aarch64, this uses assembly intrinsics, or the nightly `aarch64::_prefetch` instruction when the `aarch64` feature is used.
///
/// On other architectures, this is a no-op.
///
/// ```
/// use prefetch_index::prefetch_index;
/// let data = vec![0u8; 1024];
/// prefetch_index(&data, 512);
/// ```
#[inline(always)]
pub fn prefetch_index<T>(data: impl AsRef<[T]>, index: usize) {
    let ptr = data.as_ref().as_ptr().wrapping_add(index) as *const i8;
    #[cfg(all(target_arch = "x86_64", target_feature = "sse"))]
    unsafe {
        std::arch::x86_64::_mm_prefetch(ptr, std::arch::x86_64::_MM_HINT_T0);
    }
    #[cfg(all(target_arch = "x86", target_feature = "sse"))]
    unsafe {
        std::arch::x86::_mm_prefetch(ptr, std::arch::x86::_MM_HINT_T0);
    }
    #[cfg(all(target_arch = "aarch64", feature = "aarch64"))]
    unsafe {
        std::arch::aarch64::_prefetch::<
            { std::arch::aarch64::_PREFETCH_READ },
            { std::arch::aarch64::_PREFETCH_LOCALITY3 },
        >(ptr);
    }
    #[cfg(all(target_arch = "aarch64", not(feature = "aarch64")))]
    unsafe {
        core::arch::asm!(
            "prfm pldl1keep, [{p}]",
            p = in(reg) ptr,
            options(nostack, preserves_flags, readonly),
        );
    }
    #[cfg(not(any(
        all(target_arch = "x86_64", target_feature = "sse"),
        all(target_arch = "x86", target_feature = "sse"),
        all(target_arch = "aarch64")
    )))]
    {
        let _ = ptr; // Silence unused variable warning.
    }
}

/// Prefetches the cache line containing (the first byte of) `data[index]` for non-temporal access.
///
/// On x86/x86_64, this uses _mm_prefetch which only requires the (commonly available) SSE feature.
///
/// On aarch64, this is gated behind the `aarch64` feature flag as `aarch64::_prefetch` is nightly.
///
/// On other architectures, this is a no-op.
#[inline(always)]
pub fn prefetch_index_nta<T>(data: impl AsRef<[T]>, index: usize) {
    let ptr = data.as_ref().as_ptr().wrapping_add(index) as *const i8;
    #[cfg(all(target_arch = "x86_64", target_feature = "sse"))]
    unsafe {
        std::arch::x86_64::_mm_prefetch(ptr, std::arch::x86_64::_MM_HINT_NTA);
    }
    #[cfg(all(target_arch = "x86", target_feature = "sse"))]
    unsafe {
        std::arch::x86::_mm_prefetch(ptr, std::arch::x86::_MM_HINT_NTA);
    }
    #[cfg(all(target_arch = "aarch64", feature = "aarch64"))]
    unsafe {
        std::arch::aarch64::_prefetch::<
            { std::arch::aarch64::_PREFETCH_READ },
            { std::arch::aarch64::_PREFETCH_LOCALITY0 },
        >(ptr);
    }
    #[cfg(all(target_arch = "aarch64", not(feature = "aarch64")))]
    unsafe {
        core::arch::asm!(
            "prfm pldl1strm, [{p}]",
            p = in(reg) ptr,
            options(nostack, preserves_flags, readonly),
        );
    }
    #[cfg(not(any(
        all(target_arch = "x86_64", target_feature = "sse"),
        all(target_arch = "x86", target_feature = "sse"),
        all(target_arch = "aarch64")
    )))]
    {
        let _ = ptr; // Silence unused variable warning.
    }
}

#[cfg(test)]
mod tests {
    use std::{array::from_fn, hint::black_box};

    use super::*;

    #[test]
    fn basic() {
        let data = vec![0u8; 1024];
        prefetch_index(&data, 512);
        prefetch_index_nta(&data, 512);
    }

    /// Returns a large vector with each index pointing to the next cache line.
    fn setup() -> Vec<u64> {
        let len = 1024 * 1024 * 1024 / 8;
        let mut data = vec![0u64; len];
        // triangular stride hits all position:
        // 0->1->3->6->...
        let mut idx = 0;
        for i in 0..len {
            let next = (idx + (i + 1)) % len;
            assert_eq!(data[idx], 0);
            data[idx] = next as u64;
            idx = next;
        }
        data
    }

    #[test]
    fn out_of_bounds_is_ok() {
        for j in 0..20 {
            let len = 1 << j;
            let data = vec![0u64; len];
            for i in 0..100usize {
                prefetch_index(&data, i.wrapping_neg());
                prefetch_index(&data, len + i);
            }
            prefetch_index(&data, usize::MAX);
            prefetch_index(&data, len + usize::MAX / 2);
        }
    }

    #[test]
    fn batched_pointer_chasing_prefetch() {
        // batch size
        const B: usize = 32;

        let data = setup();

        let time = |prefetch: bool| {
            let mut indices: [usize; B] = from_fn(|i| i * data.len() / B + i);
            let start = std::time::Instant::now();
            let mut sum = 0;
            for _ in 0..data.len() / B {
                for idx in &mut indices {
                    *idx = data[*idx] as usize;
                    if prefetch {
                        prefetch_index(&data, *idx);
                    }
                    // Simulate some work to fill the CPU reorder buffer.
                    let mut x: usize = 1;
                    for _ in 0..16 {
                        x = x.wrapping_mul(*idx);
                        sum += x;
                    }
                }
            }
            black_box(sum);
            let duration = start.elapsed();
            let ns_per_it = duration.as_nanos() as f64 / (data.len() as f64);
            eprintln!(
                "Prefetch={prefetch:<5}:    {:?}  {:5.2} ns/it",
                duration, ns_per_it
            );
            ns_per_it
        };
        let no_prefetch = time(false);
        let with_prefetch = time(true);
        assert!(with_prefetch < 0.8 * no_prefetch, "Time with prefetching {with_prefetch} is not sufficiently smaller than time without prefetching {no_prefetch}");
    }
}