orthotope 0.2.0

Arena-based memory allocator with per-thread caches and fixed size classes, designed for allocation-heavy workloads like ML inference and tensor pipelines
Documentation 
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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

use orthotope::{Allocator, AllocatorConfig, SizeClass, ThreadCache};

const fn draining_config() -> AllocatorConfig {
    AllocatorConfig {
        arena_size: 1 << 20,
        alignment: 64,
        refill_target_bytes: 512,
        local_cache_target_bytes: 256,
    }
}

const fn large_stats_config() -> AllocatorConfig {
    AllocatorConfig {
        arena_size: 64 << 20,
        alignment: 64,
        refill_target_bytes: 512,
        local_cache_target_bytes: 256,
    }
}

#[test]
fn freeing_into_a_slab_backed_cache_drains_a_full_batch_when_over_limit() {
    let config = draining_config();
    let allocator = Allocator::new(config)
        .unwrap_or_else(|error| panic!("expected allocator to initialize: {error}"));
    let mut cache = ThreadCache::new(config);
    let class = SizeClass::B64;
    let ptr = allocator
        .allocate_with_cache(&mut cache, 32)
        .unwrap_or_else(|error| panic!("expected allocation to succeed: {error}"));

    assert_eq!(config.refill_count(class), 4);
    assert_eq!(config.local_limit(class), 2);
    assert_eq!(config.drain_count(class), 2);

    // SAFETY: `ptr` is still live and is freed exactly once here.
    unsafe {
        allocator
            .deallocate_with_cache(&mut cache, ptr)
            .unwrap_or_else(|error| panic!("expected free to succeed: {error}"));
    }

    let local = cache.stats();
    let shared = allocator.stats();

    assert_eq!(
        local.local[class.index()].blocks,
        config.local_limit(class),
        "freeing into an over-limit cache should leave at most the configured local limit"
    );
    assert_eq!(
        shared.small_central[class.index()].blocks,
        config.drain_count(class),
        "the central pool should receive a full configured drain batch"
    );
}

#[test]
fn large_reuse_stats_track_bucketed_free_blocks_correctly() {
    let config = large_stats_config();
    let allocator = Allocator::new(config)
        .unwrap_or_else(|error| panic!("expected allocator to initialize: {error}"));
    let mut cache = ThreadCache::new(config);
    let first_request = SizeClass::max_small_request() + 1;
    let second_request = first_request + 1_024;

    let first = allocator
        .allocate_with_cache(&mut cache, first_request)
        .unwrap_or_else(|error| panic!("expected first large allocation to succeed: {error}"));
    let second = allocator
        .allocate_with_cache(&mut cache, second_request)
        .unwrap_or_else(|error| panic!("expected second large allocation to succeed: {error}"));

    // SAFETY: both pointers are still live large allocations returned by this allocator.
    unsafe {
        allocator
            .deallocate_with_size_checked(&mut cache, first, first_request)
            .unwrap_or_else(|error| panic!("expected first large free to succeed: {error}"));
        allocator
            .deallocate_with_size_checked(&mut cache, second, second_request)
            .unwrap_or_else(|error| panic!("expected second large free to succeed: {error}"));
    }

    let freed = allocator.stats();
    assert_eq!(freed.large_live_allocations, 0);
    assert_eq!(freed.large_free_blocks, 2);
    assert!(freed.large_free_bytes > 0);

    let reused = allocator
        .allocate_with_cache(&mut cache, first_request)
        .unwrap_or_else(|error| panic!("expected large reuse allocation to succeed: {error}"));

    let after_reuse = allocator.stats();
    assert_eq!(after_reuse.large_live_allocations, 1);
    assert_eq!(after_reuse.large_free_blocks, 1);
    assert!(after_reuse.large_free_bytes < freed.large_free_bytes);

    // SAFETY: `reused` is the currently live large allocation returned above.
    unsafe {
        allocator
            .deallocate_with_size_checked(&mut cache, reused, first_request)
            .unwrap_or_else(|error| panic!("expected reused large free to succeed: {error}"));
    }
}

#[test]
fn embedding_batch_sized_bursts_stay_local_with_default_medium_class() {
    let config = AllocatorConfig::default();
    let allocator = Allocator::new(config)
        .unwrap_or_else(|error| panic!("expected allocator to initialize: {error}"));
    let mut cache = ThreadCache::new(config);
    let class = SizeClass::from_request(6_144)
        .unwrap_or_else(|| panic!("expected embedding-sized request to use the small path"));
    let mut pointers = Vec::with_capacity(8);

    assert_eq!(class, SizeClass::B6K);
    assert_eq!(config.refill_count(class), 5);
    assert_eq!(config.local_limit(class), 10);
    assert_eq!(config.drain_count(class), 2);

    for _ in 0..8 {
        let ptr = allocator
            .allocate_with_cache(&mut cache, 6_144)
            .unwrap_or_else(|error| {
                panic!("expected embedding-sized allocation to succeed: {error}")
            });
        pointers.push(ptr);
    }

    while let Some(ptr) = pointers.pop() {
        // SAFETY: every pointer in `pointers` is live and is freed exactly once here.
        unsafe {
            allocator
                .deallocate_with_size_checked(&mut cache, ptr, 6_144)
                .unwrap_or_else(|error| {
                    panic!("expected embedding-sized free to succeed: {error}")
                });
        }
    }

    let local = cache.stats();
    let shared = allocator.stats();

    assert!(local.local[class.index()].blocks >= 8);
    assert_eq!(shared.small_central[class.index()].blocks, 0);
}