polars_utils/
mem.rs

1use std::sync::LazyLock;
2
3pub static PAGE_SIZE: LazyLock<usize> = LazyLock::new(|| {
4    #[cfg(target_family = "unix")]
5    unsafe {
6        libc::sysconf(libc::_SC_PAGESIZE) as usize
7    }
8    #[cfg(not(target_family = "unix"))]
9    {
10        4096
11    }
12});
13
14pub mod prefetch {
15    use super::PAGE_SIZE;
16
17    /// # Safety
18    ///
19    /// This should only be called with pointers to valid memory.
20    unsafe fn prefetch_l2_impl(ptr: *const u8) {
21        _ = ptr; // Silence unused - not always used on all platforms.
22
23        #[cfg(target_arch = "x86_64")]
24        {
25            use std::arch::x86_64::*;
26            unsafe { _mm_prefetch(ptr as *const _, _MM_HINT_T1) };
27        }
28
29        #[cfg(all(target_arch = "aarch64", feature = "nightly"))]
30        {
31            use std::arch::aarch64::*;
32            unsafe { _prefetch(ptr as *const _, _PREFETCH_READ, _PREFETCH_LOCALITY2) };
33        }
34    }
35
36    /// Attempt to prefetch the memory in the slice to the L2 cache.
37    pub fn prefetch_l2(slice: &[u8]) {
38        if slice.is_empty() {
39            return;
40        }
41
42        // @TODO: We can play a bit more with this prefetching. Maybe introduce a maximum number of
43        // prefetches as to not overwhelm the processor. The linear prefetcher should pick it up
44        // at a certain point.
45
46        for i in (0..slice.len()).step_by(*PAGE_SIZE) {
47            unsafe { prefetch_l2_impl(slice[i..].as_ptr()) };
48        }
49
50        unsafe { prefetch_l2_impl(slice[slice.len() - 1..].as_ptr()) }
51    }
52
53    /// `madvise()` with `MADV_SEQUENTIAL` on unix systems. This is a no-op on non-unix systems.
54    pub fn madvise_sequential(#[allow(unused)] slice: &[u8]) {
55        #[cfg(target_family = "unix")]
56        madvise(slice, libc::MADV_SEQUENTIAL);
57    }
58
59    /// `madvise()` with `MADV_WILLNEED` on unix systems. This is a no-op on non-unix systems.
60    pub fn madvise_willneed(#[allow(unused)] slice: &[u8]) {
61        #[cfg(target_family = "unix")]
62        madvise(slice, libc::MADV_WILLNEED);
63    }
64
65    /// `madvise()` with `MADV_POPULATE_READ` on linux systems. This a no-op on non-linux systems.
66    pub fn madvise_populate_read(#[allow(unused)] slice: &[u8]) {
67        #[cfg(target_os = "linux")]
68        madvise(slice, libc::MADV_POPULATE_READ);
69    }
70
71    /// Forcibly reads at least one byte each page.
72    pub fn force_populate_read(slice: &[u8]) {
73        for i in (0..slice.len()).step_by(*PAGE_SIZE) {
74            std::hint::black_box(slice[i]);
75        }
76
77        std::hint::black_box(slice.last().copied());
78    }
79
80    #[cfg(target_family = "unix")]
81    fn madvise(slice: &[u8], advice: libc::c_int) {
82        if slice.is_empty() {
83            return;
84        }
85        let ptr = slice.as_ptr();
86
87        let align = ptr as usize % *PAGE_SIZE;
88        let ptr = ptr.wrapping_sub(align);
89        let len = slice.len() + align;
90
91        if unsafe { libc::madvise(ptr as *mut libc::c_void, len, advice) } != 0 {
92            let err = std::io::Error::last_os_error();
93            if let std::io::ErrorKind::InvalidInput = err.kind() {
94                panic!("{}", err);
95            }
96        }
97    }
98
99    pub fn no_prefetch(_: &[u8]) {}
100
101    /// Get the configured memory prefetch function.
102    pub fn get_memory_prefetch_func(verbose: bool) -> fn(&[u8]) -> () {
103        let memory_prefetch_func = match std::env::var("POLARS_MEMORY_PREFETCH").ok().as_deref() {
104            None => {
105                // madvise_willneed performed the best on both MacOS on Apple Silicon and Ubuntu on x86-64,
106                // using PDS-H query 3 SF=10 after clearing file cache as a benchmark.
107                #[cfg(target_family = "unix")]
108                {
109                    madvise_willneed
110                }
111                #[cfg(not(target_family = "unix"))]
112                {
113                    no_prefetch
114                }
115            },
116            Some("no_prefetch") => no_prefetch,
117            Some("prefetch_l2") => prefetch_l2,
118            Some("madvise_sequential") => {
119                #[cfg(target_family = "unix")]
120                {
121                    madvise_sequential
122                }
123                #[cfg(not(target_family = "unix"))]
124                {
125                    panic!(
126                        "POLARS_MEMORY_PREFETCH=madvise_sequential is not supported by this system"
127                    );
128                }
129            },
130            Some("madvise_willneed") => {
131                #[cfg(target_family = "unix")]
132                {
133                    madvise_willneed
134                }
135                #[cfg(not(target_family = "unix"))]
136                {
137                    panic!(
138                        "POLARS_MEMORY_PREFETCH=madvise_willneed is not supported by this system"
139                    );
140                }
141            },
142            Some("madvise_populate_read") => {
143                #[cfg(target_os = "linux")]
144                {
145                    madvise_populate_read
146                }
147                #[cfg(not(target_os = "linux"))]
148                {
149                    panic!(
150                        "POLARS_MEMORY_PREFETCH=madvise_populate_read is not supported by this system"
151                    );
152                }
153            },
154            Some("force_populate_read") => force_populate_read,
155            Some(v) => panic!("invalid value for POLARS_MEMORY_PREFETCH: {v}"),
156        };
157
158        if verbose {
159            let func_name = match memory_prefetch_func as usize {
160                v if v == no_prefetch as *const () as usize => "no_prefetch",
161                v if v == prefetch_l2 as *const () as usize => "prefetch_l2",
162                v if v == madvise_sequential as *const () as usize => "madvise_sequential",
163                v if v == madvise_willneed as *const () as usize => "madvise_willneed",
164                v if v == madvise_populate_read as *const () as usize => "madvise_populate_read",
165                v if v == force_populate_read as *const () as usize => "force_populate_read",
166                _ => unreachable!(),
167            };
168
169            eprintln!("memory prefetch function: {func_name}");
170        }
171
172        memory_prefetch_func
173    }
174}
polars_utils/mem.rs

polars_utils/
mem.rs
