candela-tensor 0.2.0

A lazy, graph-based tensor engine in Rust
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

use std::sync::Arc;

use crate::Layout;
use crate::tensor::backend::Backend;
use crate::tensor::graph::{TensorGraphBaked, TensorGraphCacheNode, TensorGraphEdge};
use crate::tensor::ops::def_op::OpKind;
use crate::tensor::planner::plan::CorePlan;
use crate::tensor::planner::{ComputeKind, OutputKind};

/// The owned, graph-detached payload of an [`OwnedComputeKind::Op`] step: the node's
/// id, op, and output layout cloned out of the graph so the plan owns them. Boxed in
/// `Op` to keep the enum small. The owned mirror of the borrowed `Op`'s `&node`.
pub(crate) struct OwnedOp<T> {
    pub(crate) id: usize,
    pub(crate) op: OpKind<T>,
    pub(crate) layout: Layout,
}

pub(crate) enum OwnedComputeKind<T, B: Backend> {
    Leaf {
        edge: Arc<TensorGraphEdge<T, B>>,
    },
    Op {
        node: Box<OwnedOp<T>>,
        output: OutputKind,
        resolved_inputs: Vec<usize>,
        dealloc_after: Vec<usize>,
    },
    CachedOp {
        cache: Arc<TensorGraphCacheNode<T, B>>,
        output: OutputKind,
        resolved_inputs: Vec<usize>,
        dealloc_after: Vec<usize>,
    },
    Baked {
        baked: Arc<TensorGraphBaked<T, B>>,
        resolved_inputs: Vec<usize>,
        dealloc_after: Vec<usize>,
    },
}

pub(crate) struct OwnedCorePlan<T, B: Backend> {
    pub(crate) plan: Vec<OwnedComputeKind<T, B>>,
    pub(crate) root_id: usize,
}

#[inline]
fn from_borrowed_compute_kind_to_owned<T: Clone, B: Backend>(
    borrowed: Vec<ComputeKind<'_, T, B>>,
) -> Vec<OwnedComputeKind<T, B>> {
    borrowed
        .into_iter()
        .map(|borrowed| match borrowed {
            ComputeKind::Leaf { edge } => OwnedComputeKind::Leaf { edge: edge.clone() },
            ComputeKind::CachedOp {
                cache,
                output,
                resolved_inputs,
                dealloc_after,
            } => OwnedComputeKind::CachedOp {
                cache: cache.clone(),
                output,
                resolved_inputs,
                dealloc_after,
            },
            ComputeKind::Op {
                node,
                output,
                resolved_inputs,
                dealloc_after,
            } => OwnedComputeKind::Op {
                node: Box::new(OwnedOp {
                    id: node.id,
                    op: node.op.clone(),
                    layout: node.layout().clone(),
                }),
                output,
                resolved_inputs,
                dealloc_after,
            },
            ComputeKind::Baked {
                baked,
                resolved_inputs,
                dealloc_after,
            } => OwnedComputeKind::Baked {
                baked: baked.clone(),
                resolved_inputs,
                dealloc_after,
            },
        })
        .collect()
}

#[inline]
pub(crate) fn from_borrowed_core_to_owned<T: Clone, B: Backend>(
    core: CorePlan<'_, T, B>,
) -> OwnedCorePlan<T, B> {
    OwnedCorePlan {
        plan: from_borrowed_compute_kind_to_owned(core.plan),
        root_id: core.root_id,
    }
}