runmat-accelerate 0.4.5

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
27
28
29
30

use std::collections::HashMap;
use std::sync::{Arc, Mutex};

pub struct PipelineRegistry {
    inner: Mutex<HashMap<u64, Arc<wgpu::ComputePipeline>>>,
}

impl Default for PipelineRegistry {
    fn default() -> Self {
        Self::new()
    }
}

impl PipelineRegistry {
    pub fn new() -> Self {
        Self {
            inner: Mutex::new(HashMap::new()),
        }
    }

    pub fn get(&self, key: &u64) -> Option<Arc<wgpu::ComputePipeline>> {
        self.inner.lock().ok().and_then(|g| g.get(key).cloned())
    }

    pub fn insert(&self, key: u64, pipeline: Arc<wgpu::ComputePipeline>) {
        if let Ok(mut g) = self.inner.lock() {
            g.insert(key, pipeline);
        }
    }
}