polars_utils/
mem.rs

1use once_cell::sync::Lazy;
2static PAGE_SIZE: Lazy<usize> = Lazy::new(|| {
3    #[cfg(target_family = "unix")]
4    unsafe {
5        libc::sysconf(libc::_SC_PAGESIZE) as usize
6    }
7    #[cfg(not(target_family = "unix"))]
8    {
9        4096
10    }
11});
12
13/// # Safety
14/// This may break aliasing rules, make sure you are the only owner.
15#[allow(clippy::mut_from_ref)]
16pub unsafe fn to_mutable_slice<T: Copy>(s: &[T]) -> &mut [T] {
17    let ptr = s.as_ptr() as *mut T;
18    let len = s.len();
19    std::slice::from_raw_parts_mut(ptr, len)
20}
21
22/// # Safety
23///
24/// This should only be called with pointers to valid memory.
25unsafe fn prefetch_l2_impl(ptr: *const u8) {
26    #[cfg(target_arch = "x86_64")]
27    {
28        use std::arch::x86_64::*;
29        unsafe { _mm_prefetch(ptr as *const _, _MM_HINT_T1) };
30    }
31
32    #[cfg(all(target_arch = "aarch64", feature = "nightly"))]
33    {
34        use std::arch::aarch64::*;
35        unsafe { _prefetch(ptr as *const _, _PREFETCH_READ, _PREFETCH_LOCALITY2) };
36    }
37}
38
39/// Attempt to prefetch the memory in the slice to the L2 cache.
40pub fn prefetch_l2(slice: &[u8]) {
41    if slice.is_empty() {
42        return;
43    }
44
45    // @TODO: We can play a bit more with this prefetching. Maybe introduce a maximum number of
46    // prefetches as to not overwhelm the processor. The linear prefetcher should pick it up
47    // at a certain point.
48
49    for i in (0..slice.len()).step_by(*PAGE_SIZE) {
50        unsafe { prefetch_l2_impl(slice[i..].as_ptr()) };
51    }
52
53    unsafe { prefetch_l2_impl(slice[slice.len() - 1..].as_ptr()) }
54}
55
56/// `madvise()` with `MADV_SEQUENTIAL` on unix systems. This is a no-op on non-unix systems.
57pub fn madvise_sequential(#[allow(unused)] slice: &[u8]) {
58    #[cfg(target_family = "unix")]
59    madvise(slice, libc::MADV_SEQUENTIAL);
60}
61
62/// `madvise()` with `MADV_WILLNEED` on unix systems. This is a no-op on non-unix systems.
63pub fn madvise_willneed(#[allow(unused)] slice: &[u8]) {
64    #[cfg(target_family = "unix")]
65    madvise(slice, libc::MADV_WILLNEED);
66}
67
68/// `madvise()` with `MADV_POPULATE_READ` on linux systems. This a no-op on non-linux systems.
69pub fn madvise_populate_read(#[allow(unused)] slice: &[u8]) {
70    #[cfg(target_os = "linux")]
71    madvise(slice, libc::MADV_POPULATE_READ);
72}
73
74#[cfg(target_family = "unix")]
75fn madvise(slice: &[u8], advice: libc::c_int) {
76    let ptr = slice.as_ptr();
77
78    let align = ptr as usize % *PAGE_SIZE;
79    let ptr = ptr.wrapping_sub(align);
80    let len = slice.len() + align;
81
82    if unsafe { libc::madvise(ptr as *mut libc::c_void, len, advice) } != 0 {
83        let err = std::io::Error::last_os_error();
84        if let std::io::ErrorKind::InvalidInput = err.kind() {
85            panic!("{}", err);
86        }
87    }
88}
polars_utils/mem.rs

polars_utils/
mem.rs
