hpt-allocator 0.1.2

An internal library for memory allocator for hpt
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

use std::{
    collections::{HashMap, HashSet},
    panic::Location,
};

use hpt_common::error::{base::TensorError, memory::MemoryError};
use lru::LruCache;

use crate::{
    ptr::SafePtr,
    storage::{CommonStorage, Storage},
};

fn allocate_mem(
    allocate_fn: impl Fn() -> *mut u8,
    deallocate_fn: &impl Fn(*mut u8, std::alloc::Layout),
    device_id: usize,
    layout: std::alloc::Layout,
    allocated: &mut HashSet<SafePtr>,
    cache: &mut LruCache<std::alloc::Layout, Vec<SafePtr>>,
) -> std::result::Result<*mut u8, TensorError> {
    let ptr = allocate_fn();
    if !ptr.is_null() {
        allocated.insert(SafePtr { ptr });
        Ok(ptr)
    } else {
        let needed_size = layout.size();
        let mut freed_size = 0;
        while freed_size < needed_size {
            if let Some((layout, ptrs)) = cache.pop_lru() {
                for safe_ptr in ptrs {
                    deallocate_fn(safe_ptr.ptr, layout);
                    freed_size += layout.size();
                    let ptr = allocate_fn();
                    if !ptr.is_null() {
                        allocated.insert(SafePtr { ptr });
                        return Ok(ptr);
                    }
                }
            } else {
                break;
            }
        }
        let ptr = allocate_fn();
        if !ptr.is_null() {
            allocated.insert(SafePtr { ptr });
            Ok(ptr)
        } else {
            Err(TensorError::Memory(MemoryError::AllocationFailed {
                device: "cpu".to_string(),
                id: device_id,
                size: layout.size() / 1024 / 1024,
                source: None,
                location: Location::caller(),
            }))
        }
    }
}

/// allocate memory based on layout provided, if the layout is not found in the cache, allocate, otherwise pop from the cache
///
/// this function internally checks if the cache is full, if it is full, it pops the least recently used layout and deallocates the memory
///
/// if the cache is not full, it inserts the allocated memory into the cache, and increments the reference count in the storage
///
/// # Safety
///
/// This function checks `null` ptr internally, any memory allocated through this method, downstream don't need to check for `null` ptr
#[track_caller]
pub(crate) fn allocate_helper(
    cache: &mut LruCache<std::alloc::Layout, Vec<SafePtr>>,
    allocated: &mut HashSet<SafePtr>,
    storage: &mut HashMap<usize, CommonStorage>,
    allocate_fn: impl Fn() -> *mut u8,
    zero_fn: impl Fn(*mut u8, std::alloc::Layout),
    deallocate_fn: impl Fn(*mut u8, std::alloc::Layout),
    layout: std::alloc::Layout,
    device_id: usize,
) -> std::result::Result<*mut u8, TensorError> {
    let ptr = if let Some(ptr) = cache.get_mut(&layout)
    /*check if we previously allocated same layout of memory */
    {
        // try pop the memory out, if it return None, there is no available cached memory, we need to allocate new memory
        if let Some(safe_ptr) = ptr.pop() {
            zero_fn(safe_ptr.ptr, layout);
            safe_ptr.ptr
        } else {
            allocate_mem(
                allocate_fn,
                &deallocate_fn,
                device_id,
                layout,
                allocated,
                cache,
            )?
        }
    } else {
        allocate_mem(
            allocate_fn,
            &deallocate_fn,
            device_id,
            layout,
            allocated,
            cache,
        )?
    };
    // check if the cache is full, if it is full, pop the least recently used layout and deallocate the memory
    if cache.cap().get() == cache.len() {
        if let Some((layout, ptrs)) = cache.pop_lru() {
            for safe_ptr in ptrs {
                deallocate_fn(safe_ptr.ptr, layout);
            }
        }
    }
    // increment the reference count in the storage of the ptr allocated
    if let Some(storage) = storage.get_mut(&device_id) {
        storage.increment_ref(SafePtr { ptr });
    } else {
        let mut new_storage = CommonStorage::new();
        new_storage.increment_ref(SafePtr { ptr });
        storage.insert(device_id, new_storage);
    }
    Ok(ptr)
}