runmat-accelerate 0.4.4

Pluggable GPU acceleration layer for RunMat (CUDA, ROCm, Metal, Vulkan/Spir-V)
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

use std::collections::HashSet;
use std::sync::Mutex;

use once_cell::sync::Lazy;
use runmat_accelerate_api::GpuTensorHandle;

static RESIDENT_HANDLES: Lazy<Mutex<HashSet<u64>>> = Lazy::new(|| Mutex::new(HashSet::new()));

pub fn mark(handle: &GpuTensorHandle) {
    if let Ok(mut guard) = RESIDENT_HANDLES.lock() {
        guard.insert(handle.buffer_id);
    }
}

pub fn clear(handle: &GpuTensorHandle) {
    if let Ok(mut guard) = RESIDENT_HANDLES.lock() {
        guard.remove(&handle.buffer_id);
    }
}

pub fn is_resident(handle: &GpuTensorHandle) -> bool {
    RESIDENT_HANDLES
        .lock()
        .map(|guard| guard.contains(&handle.buffer_id))
        .unwrap_or(false)
}