vortex_array/
buffer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Debug;
5use std::hash::Hash;
6use std::sync::Arc;
7
8use vortex_buffer::ByteBuffer;
9use vortex_error::VortexResult;
10use vortex_error::vortex_bail;
11use vortex_utils::dyn_traits::DynEq;
12use vortex_utils::dyn_traits::DynHash;
13
14/// A buffer can be either on the CPU or on an attached device (e.g. GPU).
15/// The Device implementation will come later.
16#[derive(Debug, Clone)]
17pub enum BufferHandle {
18    /// On the host/cpu.
19    Host(ByteBuffer),
20    /// On the device.
21    Device(Arc<dyn DeviceBuffer>),
22}
23
24/// A buffer that is stored on a device (e.g. GPU).
25pub trait DeviceBuffer: 'static + Send + Sync + Debug + DynEq + DynHash {
26    /// Returns the length of the buffer in bytes.
27    fn len(&self) -> usize;
28
29    /// Returns true if the buffer is empty.
30    fn is_empty(&self) -> bool {
31        self.len() == 0
32    }
33
34    /// Attempts to copy the device buffer to a host ByteBuffer.
35    fn to_host(self: Arc<Self>) -> VortexResult<ByteBuffer>;
36}
37
38impl Hash for dyn DeviceBuffer {
39    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
40        self.dyn_hash(state);
41    }
42}
43
44impl PartialEq for dyn DeviceBuffer {
45    fn eq(&self, other: &Self) -> bool {
46        self.dyn_eq(other)
47    }
48}
49impl Eq for dyn DeviceBuffer {}
50
51impl BufferHandle {
52    /// Fetches the cpu buffer and fails otherwise.
53    pub fn bytes(&self) -> &ByteBuffer {
54        match self {
55            BufferHandle::Host(b) => b,
56            BufferHandle::Device(_) => todo!(),
57        }
58    }
59
60    /// Fetches the cpu buffer and fails otherwise.
61    pub fn into_bytes(self) -> ByteBuffer {
62        match self {
63            BufferHandle::Host(b) => b,
64            BufferHandle::Device(_) => todo!(),
65        }
66    }
67
68    /// Attempts to convert this handle into a CPU ByteBuffer.
69    /// Returns an error if the buffer is on a device.
70    pub fn try_to_bytes(self) -> VortexResult<ByteBuffer> {
71        match self {
72            BufferHandle::Host(b) => Ok(b),
73            BufferHandle::Device(_) => vortex_bail!("cannot move device_buffer to buffer"),
74        }
75    }
76}
vortex_array/buffer.rs

vortex_array/
buffer.rs
