tract_gpu/memory/
pool.rs

1use crate::device::get_context;
2use crate::memory::DeviceResolvedMemSchema;
3use crate::tensor::DeviceTensor;
4use crate::tensor::DeviceView;
5use crate::tensor::OwnedDeviceTensor;
6
7use tract_core::internal::*;
8
9#[derive(Debug)]
10pub struct DeviceMemoryPool {
11    storage: Arc<Box<dyn OwnedDeviceTensor>>,
12    resolved_schema: DeviceResolvedMemSchema,
13}
14
15impl DeviceMemoryPool {
16    pub fn from_schema(resolved_schema: DeviceResolvedMemSchema) -> TractResult<Self> {
17        Ok(Self {
18            storage: Arc::new(
19                get_context()?
20                    .uninitialized_device_tensor(&[resolved_schema.memory_size], DatumType::U8)?,
21            ),
22            resolved_schema,
23        })
24    }
25
26    pub fn tensor_for_node(
27        &self,
28        node_id: usize,
29        dt: DatumType,
30        shape: &[usize],
31    ) -> TractResult<DeviceTensor> {
32        if let Some(offsets) = self.resolved_schema.offsets_by_node[node_id].as_ref() {
33            ensure!(offsets.len() == 1, "'tensor_for_node' is for mono-output nodes only");
34        }
35        self.tensor_for_node_output(node_id, 0, dt, shape)
36    }
37
38    /// Per-output variant of [`Self::tensor_for_node`] for multi-output nodes:
39    /// each output slot has its own arena region in the schema.
40    pub fn tensor_for_node_output(
41        &self,
42        node_id: usize,
43        slot: usize,
44        dt: DatumType,
45        shape: &[usize],
46    ) -> TractResult<DeviceTensor> {
47        match self.resolved_schema.offsets_by_node[node_id].as_ref() {
48            Some(offsets) if slot < offsets.len() && offsets[slot].len() == 1 => Ok(DeviceView {
49                buffer: Arc::clone(&self.storage),
50                dt,
51                len: shape.iter().product(),
52                shape: shape.into(),
53                strides: Tensor::natural_strides(shape),
54                offset_bytes: offsets[slot][0],
55                exotic_fact: None,
56            }
57            .into()),
58            _ => DeviceTensor::uninitialized_dt(dt, shape),
59        }
60    }
61
62    pub fn scalar_exotic_tensor_for_node(
63        &self,
64        node_id: usize,
65        dt: DatumType,
66        exotic_fact: Box<dyn ExoticFact>,
67    ) -> TractResult<DeviceTensor> {
68        match self.resolved_schema.offsets_by_node[node_id].as_ref() {
69            Some(offsets) => {
70                ensure!(
71                    offsets.len() == 1 && offsets[0].len() == 2,
72                    "'scalar_exotic_tensor_for_node' is for mono-output nodes only"
73                );
74                Ok(DeviceView {
75                    buffer: Arc::clone(&self.storage),
76                    dt,
77                    len: 1,
78                    shape: tvec!(),
79                    strides: tvec!(),
80                    offset_bytes: offsets[0][1],
81                    exotic_fact: Some(exotic_fact.clone()),
82                }
83                .into())
84            }
85            None => DeviceTensor::uninitialized_exotic(exotic_fact),
86        }
87    }
88}
tract_gpu/memory/pool.rs

tract_gpu/memory/
pool.rs
