memkit-gpu 0.2.0-beta.1

Backend-agnostic GPU memory management for memkit
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

//! GPU memory transfer operations.

use crate::backend::MkGpuBackend;
use crate::buffer::MkDeviceBuffer;
use crate::staging::MkStagingBuffer;
use crate::MkGpu;

/// A pending GPU memory transfer.
pub struct MkGpuTransfer {
    id: u64,
    complete: bool,
}

impl MkGpuTransfer {
    /// Create a new transfer handle.
    pub(crate) fn new(id: u64) -> Self {
        Self { id, complete: false }
    }

    /// Get the transfer ID.
    pub fn id(&self) -> u64 {
        self.id
    }

    /// Check if the transfer is complete.
    pub fn is_complete(&self) -> bool {
        self.complete
    }

    /// Mark as complete.
    pub(crate) fn mark_complete(&mut self) {
        self.complete = true;
    }
}

/// Builder for GPU transfer operations.
pub struct MkTransferBuilder<'a, B: MkGpuBackend> {
    gpu: &'a MkGpu<B>,
    staging: &'a MkStagingBuffer<B>,
    device: &'a MkDeviceBuffer<B>,
    offset: usize,
    size: usize,
}

impl<'a, B: MkGpuBackend> MkTransferBuilder<'a, B> {
    /// Create a new transfer builder.
    pub fn new(
        gpu: &'a MkGpu<B>,
        staging: &'a MkStagingBuffer<B>,
        device: &'a MkDeviceBuffer<B>,
    ) -> Self {
        Self {
            gpu,
            staging,
            device,
            offset: 0,
            size: staging.size().min(device.size()),
        }
    }

    /// Set the destination offset.
    pub fn offset(mut self, offset: usize) -> Self {
        self.offset = offset;
        self
    }

    /// Set the transfer size.
    pub fn size(mut self, size: usize) -> Self {
        self.size = size;
        self
    }

    /// Submit the transfer.
    pub fn submit(self) -> Result<MkGpuTransfer, B::Error> {
        self.gpu.transfer(self.staging, self.device)
    }
}

/// Batch transfer builder for multiple transfers.
pub struct MkBatchTransfer<'a, B: MkGpuBackend> {
    gpu: &'a MkGpu<B>,
    transfers: Vec<(B::BufferHandle, B::BufferHandle, usize)>,
}

impl<'a, B: MkGpuBackend> MkBatchTransfer<'a, B> {
    /// Create a new batch transfer.
    pub(crate) fn new(gpu: &'a MkGpu<B>) -> Self {
        Self {
            gpu,
            transfers: Vec::new(),
        }
    }

    /// Add a copy operation to the batch.
    pub fn copy(
        mut self,
        src: &MkStagingBuffer<B>,
        dst: &MkDeviceBuffer<B>,
    ) -> Self {
        let size = src.size().min(dst.size());
        self.transfers.push((src.handle().clone(), dst.handle().clone(), size));
        self
    }

    /// Submit all transfers and wait for completion.
    pub fn submit_and_wait(self) -> Result<(), B::Error> {
        for (src, dst, size) in self.transfers {
            self.gpu.backend().copy_buffer(&src, &dst, size)?;
        }
        self.gpu.wait_idle()
    }

    /// Submit all transfers without waiting.
    pub fn submit(self) -> Result<Vec<MkGpuTransfer>, B::Error> {
        let mut results = Vec::new();
        for (src, dst, size) in self.transfers {
            self.gpu.backend().copy_buffer(&src, &dst, size)?;
            results.push(MkGpuTransfer::new(results.len() as u64));
        }
        Ok(results)
    }
}