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

polars_utils/
mem.rs
