Skip to main content

spatialrust_core/
transfer.rs

1/// Direction of an explicit data transfer between host and device memory.
2#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
3pub enum TransferDirection {
4    /// Copy from host memory into device memory.
5    HostToDevice,
6    /// Copy between buffers on the same device.
7    DeviceToDevice,
8    /// Copy from device memory back into host memory.
9    DeviceToHost,
10}
11
12/// Byte accounting for explicit execution transfers.
13#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
14pub struct TransferStats {
15    host_to_device_bytes: u64,
16    device_to_device_bytes: u64,
17    device_to_host_bytes: u64,
18}
19
20impl TransferStats {
21    /// Records bytes transferred in the given direction.
22    pub fn record(&mut self, direction: TransferDirection, bytes: u64) {
23        let counter = match direction {
24            TransferDirection::HostToDevice => &mut self.host_to_device_bytes,
25            TransferDirection::DeviceToDevice => &mut self.device_to_device_bytes,
26            TransferDirection::DeviceToHost => &mut self.device_to_host_bytes,
27        };
28        *counter = counter.saturating_add(bytes);
29    }
30
31    /// Returns bytes copied from host memory into device memory.
32    #[must_use]
33    pub const fn host_to_device_bytes(self) -> u64 {
34        self.host_to_device_bytes
35    }
36
37    /// Returns bytes copied between buffers on a device.
38    #[must_use]
39    pub const fn device_to_device_bytes(self) -> u64 {
40        self.device_to_device_bytes
41    }
42
43    /// Returns bytes copied from device memory into host memory.
44    #[must_use]
45    pub const fn device_to_host_bytes(self) -> u64 {
46        self.device_to_host_bytes
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::{TransferDirection, TransferStats};
53
54    #[test]
55    fn records_transfer_bytes_by_direction() {
56        let mut stats = TransferStats::default();
57        stats.record(TransferDirection::HostToDevice, 12);
58        stats.record(TransferDirection::HostToDevice, 8);
59        stats.record(TransferDirection::DeviceToDevice, 4);
60        stats.record(TransferDirection::DeviceToHost, 16);
61
62        assert_eq!(stats.host_to_device_bytes(), 20);
63        assert_eq!(stats.device_to_device_bytes(), 4);
64        assert_eq!(stats.device_to_host_bytes(), 16);
65    }
66}