tract_gpu/memory/
pool.rs

1use crate::device::get_context;
2use crate::memory::DeviceResolvedMemSchema;
3use crate::tensor::DeviceArenaView;
4use crate::tensor::DeviceTensor;
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 => {
49                Ok(DeviceArenaView {
50                    arena: Arc::clone(&self.storage),
51                    dt,
52                    len: shape.iter().product(),
53                    shape: shape.into(),
54                    strides: Tensor::natural_strides(shape),
55                    offset_bytes: offsets[slot][0],
56                    exotic_fact: None,
57                }
58                .into())
59            }
60            _ => DeviceTensor::uninitialized_dt(dt, shape),
61        }
62    }
63
64    pub fn scalar_exotic_tensor_for_node(
65        &self,
66        node_id: usize,
67        dt: DatumType,
68        exotic_fact: Box<dyn ExoticFact>,
69    ) -> TractResult<DeviceTensor> {
70        match self.resolved_schema.offsets_by_node[node_id].as_ref() {
71            Some(offsets) => {
72                ensure!(
73                    offsets.len() == 1 && offsets[0].len() == 2,
74                    "'scalar_exotic_tensor_for_node' is for mono-output nodes only"
75                );
76                Ok(DeviceArenaView {
77                    arena: Arc::clone(&self.storage),
78                    dt,
79                    len: 1,
80                    shape: tvec!(),
81                    strides: tvec!(),
82                    offset_bytes: offsets[0][1],
83                    exotic_fact: Some(exotic_fact.clone()),
84                }
85                .into())
86            }
87            None => DeviceTensor::uninitialized_exotic(exotic_fact),
88        }
89    }
90}
tract_gpu/memory/pool.rs

tract_gpu/memory/
pool.rs
