1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
use mimalloc_rust_sys::{
aligned_allocation::{mi_malloc_aligned, mi_realloc_aligned, mi_zalloc_aligned},
basic_allocation::{mi_free, mi_malloc, mi_realloc, mi_zalloc},
extended_functions::{mi_process_info, mi_stats_reset},
};
use serde::Serialize;
use std::alloc::{GlobalAlloc, Layout};
/// `MI_MAX_ALIGN_SIZE` is 16 unless manually overridden:
/// <https://github.com/microsoft/mimalloc/blob/15220c68/include/mimalloc-types.h#L22>
const MI_MAX_ALIGN_SIZE: usize = 16;
/// Allocation and process statistics
#[derive(Debug, Clone, Copy, Default, Serialize)]
pub struct AllocStats {
pub elapsed_ms: usize,
pub user_ms: usize,
pub system_ms: usize,
pub current_rss: usize,
pub peak_rss: usize,
pub current_commit: usize,
pub peak_commit: usize,
pub page_faults: usize,
}
/// A [`GlobalAlloc`] implementation that uses `mimalloc`
#[derive(Debug, Clone, Copy, Default)]
pub struct MiMalloc;
// The dead code lint triggers if the function is unused within any benchmark,
// not if it's unused in all of them
#[allow(dead_code)]
impl MiMalloc {
/// Collect allocation stats
pub fn stats(&self) -> AllocStats {
let mut stats = AllocStats::default();
unsafe {
mi_process_info(
&mut stats.elapsed_ms,
&mut stats.user_ms,
&mut stats.system_ms,
&mut stats.current_rss,
&mut stats.peak_rss,
&mut stats.current_commit,
&mut stats.peak_commit,
&mut stats.page_faults,
);
}
stats
}
/// Reset allocation statistics
pub fn reset_stats(&self) {
unsafe { mi_stats_reset() };
}
}
// Using the normal (unaligned) mimalloc apis is apparently marginally faster
// and better supported than using the aligned api for everything, so we do a
// check to see if the allocation falls within the range where it's better to
// use the unaligned api
//
// See https://github.com/microsoft/mimalloc/issues/314
unsafe impl GlobalAlloc for MiMalloc {
#[inline]
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
unsafe {
if use_unaligned_api(layout.size(), layout.align()) {
mi_malloc(layout.size())
} else {
mi_malloc_aligned(layout.size(), layout.align())
}
.cast()
}
}
#[inline]
unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
unsafe {
if use_unaligned_api(layout.size(), layout.align()) {
mi_zalloc(layout.size())
} else {
mi_zalloc_aligned(layout.size(), layout.align())
}
.cast()
}
}
#[inline]
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
unsafe {
let ptr = ptr.cast();
if use_unaligned_api(layout.size(), layout.align()) {
mi_realloc(ptr, new_size)
} else {
mi_realloc_aligned(ptr, new_size, layout.align())
}
.cast()
}
}
#[inline]
unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
unsafe {
// Deallocation doesn't care about alignment which is nice
mi_free(ptr.cast());
}
}
}
#[inline(always)]
const fn use_unaligned_api(size: usize, alignment: usize) -> bool {
// This logic is based on the discussion [here]. We don't bother with the
// 3rd suggested test due to it being high cost (calling `mi_good_size`)
// compared to the other checks, and also feeling like it relies on too much
// implementation-specific behavior.
//
// [here]: https://github.com/microsoft/mimalloc/issues/314#issuecomment-708541845
(alignment <= MI_MAX_ALIGN_SIZE && size >= alignment)
|| (alignment == size && alignment <= 4096)
}