vkml 0.0.3

High-level Vulkan-based machine learning library
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

use std::sync::atomic::{AtomicU64, Ordering};

pub struct MemoryTracker {
    maximum: u64,
    in_use: AtomicU64,
}

// This implementation doesn't require a mutable reference to update
// The trade off of checking after the change is that there's only one operation, so no race conditions

impl MemoryTracker {
    pub fn new(maximum: u64) -> Self {
        Self {
            maximum,
            in_use: AtomicU64::new(0),
        }
    }

    pub fn allocate(&self, size: u64) {
        let prev = self.in_use.fetch_add(size, Ordering::Release);

        if prev + size > self.maximum {
            panic!(
                "Memory limit exceeded: tried to allocate {} bytes when {} of {} bytes are used",
                size, prev, self.maximum
            );
        }
    }

    pub fn get_current(&self) -> u64 {
        self.in_use.load(Ordering::Acquire)
    }

    pub fn get_available(&self) -> u64 {
        self.maximum - self.get_current()
    }

    pub fn get_maximum(&self) -> u64 {
        self.maximum
    }
}