topos 0.13.0

An autodiff compiler stack 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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

use cow_vec::CowVec;
use smallvec::SmallVec;

use crate::Shape;
use crate::op::Op;

use super::opcode::Node;
use super::{Origin, Symbol, ValueId};

use super::Operands;

/// The node columns of a recorded graph.
///
/// Three equal-length columns describe every recorded node: what it
/// computes, which earlier nodes it reads, and the shape inferred when
/// it was recorded. Runs replay ops and operands; shapes are the
/// cold column used at record time and by structure consumers (plans,
/// zero placeholders). Parameter and input payloads live outside this
/// type: initials and defaults in the spec's stores, live state in the
/// caller's [`Parameters`](crate::Parameters).
///
/// Cloning shares the append-only column arena in O(1), which is how
/// plans and `differentiate` freeze the structure they read.
#[derive(Debug, Clone)]
pub(crate) struct Structure<Data> {
    pub(crate) ops: CowVec<Op<Data>>,
    pub(crate) operands: CowVec<Operands>,
    pub(crate) shapes: CowVec<Shape>,
}

impl<Data> Structure<Data> {
    /// Creates empty columns.
    pub(crate) fn new() -> Self {
        Self {
            ops: CowVec::new(),
            operands: CowVec::new(),
            shapes: CowVec::new(),
        }
    }

    /// Returns the number of recorded nodes.
    pub(crate) fn len(&self) -> usize {
        self.ops.len()
    }

    /// Appends one node and returns its handle.
    ///
    /// The three columns stay equal length; callers supply a shape that
    /// has already been inferred and validated against the operands.
    pub(crate) fn push(&mut self, op: Op<Data>, operands: Operands, shape: Shape) -> ValueId {
        self.ops.push(op);
        self.operands.push(operands);
        self.shapes.push(shape);
        debug_assert_eq!(self.ops.len(), self.operands.len());
        debug_assert_eq!(self.ops.len(), self.shapes.len());
        ValueId(self.ops.len() - 1)
    }
}

impl<Data> Structure<Data> {
    /// Marks the ancestor closure of `seeds` over the operand links:
    /// one descending sweep, correct because operands live below
    /// their consumers. The shared reachability of the interpreter's
    /// entry closure and the plan's wanted set; the reverse scans
    /// keep their own integrated marking, welded to this pair by the
    /// closure suite.
    pub(crate) fn ancestors(&self, seeds: impl IntoIterator<Item = ValueId>) -> Vec<bool> {
        let mut wanted = vec![false; self.len()];
        for seed in seeds {
            wanted[seed.index()] = true;
        }
        for index in (0..wanted.len()).rev() {
            if !wanted[index] {
                continue;
            }
            let links = self
                .operands
                .get(index)
                .expect("operand links cover the columns");
            for link in links.as_slice() {
                wanted[link.index()] = true;
            }
        }
        wanted
    }

    /// Builds the public snapshot of the node at `index`, minting
    /// symbols of the `origin` family for the node and its operands.
    ///
    /// # Panics
    /// Panics if `index` is out of bounds.
    pub(crate) fn node_at(&self, origin: Origin, index: usize) -> Node {
        let op = self
            .ops
            .get(index)
            .expect("`node_at` index is in bounds for its columns");
        let operands: SmallVec<[Symbol; 2]> = self
            .operands
            .get(index)
            .expect("operand links cover the columns")
            .as_slice()
            .iter()
            .map(|link| Symbol { origin, id: *link })
            .collect();
        Node {
            symbol: Symbol {
                origin,
                id: ValueId(index),
            },
            opcode: op.opcode(),
            shape: self
                .shapes
                .get(index)
                .expect("shapes cover the columns")
                .clone(),
            operands,
        }
    }
}