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
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

use std::sync::{Arc, Weak};

use vulkanalia::{vk, vk::DeviceV1_0};

use crate::VKMLError;

use super::vk_gpu::Gpu;

pub struct GPUMemory {
    pub buffer: vk::Buffer,
    pub memory: vk::DeviceMemory,
    pub size: vk::DeviceSize,
    pub properties: vk::MemoryPropertyFlags,
    gpu: Weak<Gpu>,
}

impl GPUMemory {
    pub fn new(
        buffer: vk::Buffer,
        memory: vk::DeviceMemory,
        size: vk::DeviceSize,
        properties: vk::MemoryPropertyFlags,
        gpu: &Arc<Gpu>,
    ) -> Self {
        Self {
            buffer,
            memory,
            size,
            properties,
            gpu: Arc::downgrade(gpu),
        }
    }

    fn upgrade_gpu(&self) -> Result<Arc<Gpu>, VKMLError> {
        self.gpu.upgrade().ok_or_else(|| {
            VKMLError::Gpu("GPU allocation reference dropped before use".to_string())
        })
    }

    /// Copy raw bytes into GPU memory. Falls back to staging when not host-visible.
    pub fn copy_into(&self, data: &[u8]) -> Result<(), VKMLError> {
        let data_size = data.len() as vk::DeviceSize;

        if data_size > self.size {
            return Err(VKMLError::Gpu(format!(
                "Data size {} exceeds GPU buffer size {}",
                data_size, self.size
            )));
        }

        if self
            .properties
            .contains(vk::MemoryPropertyFlags::HOST_VISIBLE)
        {
            let gpu = self.upgrade_gpu()?;
            unsafe {
                let data_ptr = gpu.get_device().map_memory(
                    self.memory,
                    0,
                    data_size,
                    vk::MemoryMapFlags::empty(),
                )? as *mut u8;

                std::ptr::copy_nonoverlapping(data.as_ptr(), data_ptr, data.len());

                gpu.get_device().unmap_memory(self.memory);
            }
            Ok(())
        } else {
            let gpu = self.upgrade_gpu()?;
            gpu.write_through_staging(self, data)
        }
    }

    /// Read raw bytes from GPU memory. Falls back to staging when not host-visible.
    pub fn read_memory(&self) -> Result<Box<[u8]>, VKMLError> {
        if self
            .properties
            .contains(vk::MemoryPropertyFlags::HOST_VISIBLE)
        {
            let gpu = self.upgrade_gpu()?;
            let mut buffer = Box::new_uninit_slice(self.size as usize);

            unsafe {
                let data_ptr = gpu.get_device().map_memory(
                    self.memory,
                    0,
                    self.size,
                    vk::MemoryMapFlags::empty(),
                )? as *const u8;

                let buffer_ptr = buffer.as_mut_ptr() as *mut u8;
                std::ptr::copy_nonoverlapping(data_ptr, buffer_ptr, buffer.len());

                gpu.get_device().unmap_memory(self.memory);
            }

            let output = unsafe { buffer.assume_init() };

            Ok(output)
        } else {
            let gpu = self.upgrade_gpu()?;
            gpu.read_through_staging(self)
        }
    }
}