tract_gpu/
turn_handler.rs

1use crate::memory::DeviceMemSchema;
2use crate::memory::DeviceMemoryPool;
3use crate::tensor::DeviceTensor;
4use tract_core::internal::*;
5
6#[derive(Debug, Clone)]
7pub struct DeviceTurnHandler {
8    pub mem_schema: DeviceMemSchema,
9}
10
11impl DeviceTurnHandler {
12    pub fn from_plan(plan: &TypedSimplePlan, memory_hint: &SymbolValues) -> TractResult<Self> {
13        let mem_schema =
14            DeviceMemSchema::build(plan.model(), plan.order_without_consts(), memory_hint)?;
15        Ok(Self { mem_schema })
16    }
17
18    /// The arena a runtime installs on `plan`: sized by `hints` when given, and
19    /// without any when every shape in the model is concrete. A symbolic model
20    /// with no hints gets none, and its device tensors are allocated one by one.
21    /// wgpu does not call it: its dispatches cannot bind two regions of one
22    /// arena buffer as a read and a write.
23    pub fn for_plan(
24        plan: &TypedSimplePlan,
25        hints: Option<&SymbolValues>,
26    ) -> TractResult<Option<Self>> {
27        let no_hints = SymbolValues::default();
28        let hints = match hints {
29            Some(hints) => hints,
30            None if plan
31                .model()
32                .nodes
33                .iter()
34                .all(|n| n.outputs.iter().all(|o| o.fact.shape.is_concrete())) =>
35            {
36                &no_hints
37            }
38            None => return Ok(None),
39        };
40        Self::from_plan(plan, hints).context("While sizing memory arena. Missing hint ?").map(Some)
41    }
42}
43
44impl TurnStateHandler for DeviceTurnHandler {
45    fn before_plan_eval(&self, turn: &mut TurnState) -> TractResult<()> {
46        let resolved_mem_schema = self.mem_schema.resolve(&turn.resolved_symbols)?;
47        let memory_pool = DeviceMemoryPool::from_schema(resolved_mem_schema)?;
48
49        turn.shared.insert(memory_pool);
50        ensure!(turn.shared.get::<DeviceMemoryPool>().is_some());
51        Ok(())
52    }
53
54    fn after_plan_eval(&self, turn: &mut TurnState) -> TractResult<()> {
55        turn.shared.remove::<DeviceMemoryPool>();
56        Ok(())
57    }
58}
59
60pub fn make_tensor_for_node(
61    ctx: &EvalContext,
62    dt: DatumType,
63    shape: &[usize],
64) -> TractResult<DeviceTensor> {
65    ctx.shared
66        .and_then(|s| s.get::<DeviceMemoryPool>())
67        .map(|mem| mem.tensor_for_node(ctx.node_id, dt, shape))
68        .unwrap_or_else(|| DeviceTensor::uninitialized_dt(dt, shape))
69}
70
71/// Like [`make_tensor_for_node`] but for one output slot of a multi-output
72/// node: the memory schema reserves one arena region per device output.
73pub fn make_tensor_for_node_output(
74    ctx: &EvalContext,
75    slot: usize,
76    dt: DatumType,
77    shape: &[usize],
78) -> TractResult<DeviceTensor> {
79    ctx.shared
80        .and_then(|s| s.get::<DeviceMemoryPool>())
81        .map(|mem| mem.tensor_for_node_output(ctx.node_id, slot, dt, shape))
82        .unwrap_or_else(|| DeviceTensor::uninitialized_dt(dt, shape))
83}
84
85pub fn make_scalar_exotic_tensor_for_node(
86    ctx: &EvalContext,
87    dt: DatumType,
88    exotic_fact: Box<dyn ExoticFact>,
89) -> TractResult<DeviceTensor> {
90    match ctx.shared.and_then(|s| s.get::<DeviceMemoryPool>()) {
91        Some(mem) => mem.scalar_exotic_tensor_for_node(ctx.node_id, dt, exotic_fact),
92        None => DeviceTensor::uninitialized_exotic(exotic_fact),
93    }
94}
tract_gpu/turn_handler.rs

tract_gpu/
turn_handler.rs
