logicaffeine_compile/codegen_cuda/lane/
branches.rs

1//! A form's launches as a graph over several streams. A launch runs after every earlier launch that writes memory it
2//! reads or writes, or reads memory it writes — memory by where it lives: a step-local array or a launch's scratch by
3//! its bytes in the arena (which several tensors share in turn), anything else by its buffer, an alias — a view — where
4//! its array lives (`arena::home`). Each launch takes a stream: the stream whose last launch it follows, or one whose work it already
5//! follows, or the least recently used; it waits on another stream's launch by an event only where nothing it already
6//! follows orders the two. The launches compute what they compute in order, every dependency kept.
7
8use super::arena::{self, Arena};
9use super::plan::{Schedule, Touch};
10
11/// The streams a form's launches spread over, the step's own among them.
12pub const STREAMS: usize = 4;
13
14/// A form's launches over its streams.
15#[derive(Clone, Debug, Default, PartialEq, Eq)]
16pub struct Branched {
17    /// Each launch's stream (0 the step's own).
18    pub stream: Vec<usize>,
19    /// The events each launch's stream waits on before it.
20    pub waits: Vec<Vec<usize>>,
21    /// The event recorded on each launch's stream after it, where a launch on another stream waits on it.
22    pub records: Vec<Option<usize>>,
23    /// The streams the launches take besides the step's own.
24    pub sides: Vec<usize>,
25    /// The events the step's own stream waits on at its end: each side stream's last launch's.
26    pub joins: Vec<usize>,
27    /// The events of the form.
28    pub events: usize,
29}
30
31/// Memory a launch touches: bytes of the arena, or a buffer of its own.
32#[derive(Clone, Debug, PartialEq, Eq)]
33pub enum Region {
34    Arena(usize, usize),
35    Named(String),
36}
37
38impl Region {
39    pub fn overlaps(&self, other: &Region) -> bool {
40        match (self, other) {
41            (Region::Arena(a0, a1), Region::Arena(b0, b1)) => a0 < b1 && b0 < a1,
42            (Region::Named(a), Region::Named(b)) => a == b,
43            _ => false,
44        }
45    }
46}
47
48/// Where tensor `t`'s array `touch` lives: in the arena, or a buffer of its own — an alias's in its home's.
49fn region(s: &Schedule, arena: &Arena, t: &str, touch: Touch) -> Region {
50    let name = arena::home(s, t, touch);
51    match arena.at(name, touch) {
52        Some(p) => Region::Arena(p.offset, p.offset + p.bytes),
53        None => Region::Named(format!("{name}/{touch:?}")),
54    }
55}
56
57/// The memory launch `i` reads and writes: its tensors' arrays and its scratch (written and read, as a launch's own).
58pub fn regions(s: &Schedule, arena: &Arena, i: usize) -> (Vec<Region>, Vec<Region>) {
59    let l = &s.launches[i];
60    let (r, w) = l.touches(s.numerics);
61    let reads: Vec<Region> = r.into_iter().map(|(t, touch)| region(s, arena, t, touch)).collect();
62    let mut writes: Vec<Region> = w.into_iter().map(|(t, touch)| region(s, arena, t, touch)).collect();
63    let owner = l.writes().first().map_or(String::new(), |t| t.to_string());
64    for sc in arena::scratch(l, s.numerics) {
65        writes.push(match arena.scratch_at(i, sc.tag) {
66            Some(p) if !sc.counter => Region::Arena(p.offset, p.offset + p.bytes),
67            _ => Region::Named(format!("{}#{owner}", sc.tag)),
68        });
69    }
70    (reads, writes)
71}
72
73/// Whether launch `j` must follow launch `i`: one writes memory the other reads or writes.
74fn conflict(a: &(Vec<Region>, Vec<Region>), b: &(Vec<Region>, Vec<Region>)) -> bool {
75    let hits = |x: &[Region], y: &[Region]| x.iter().any(|p| y.iter().any(|q| p.overlaps(q)));
76    hits(&a.1, &b.0) || hits(&a.1, &b.1) || hits(&a.0, &b.1)
77}
78
79/// The form's launches spread over `STREAMS` streams.
80pub fn branch(s: &Schedule, arena: &Arena) -> Branched {
81    let n = s.launches.len();
82    let mem: Vec<(Vec<Region>, Vec<Region>)> = (0..n).map(|i| regions(s, arena, i)).collect();
83    // `clock[i][t]`: the last launch on stream `t` that launch `i` follows (by order on a stream or by events), plus one.
84    let mut clock: Vec<[usize; STREAMS]> = Vec::with_capacity(n);
85    let mut last: [Option<usize>; STREAMS] = [None; STREAMS];
86    let mut stream_clock: [[usize; STREAMS]; STREAMS] = [[0; STREAMS]; STREAMS];
87    let mut out = Branched { stream: vec![0; n], waits: vec![Vec::new(); n], records: vec![None; n], ..Default::default() };
88    for j in 0..n {
89        let preds: Vec<usize> = (0..j).filter(|&i| conflict(&mem[i], &mem[j])).collect();
90        // Whether stream `t` has nothing running that `j` does not follow: unused, its last launch a predecessor, or
91        // ordered before one.
92        let follows = |t: usize| match last[t] {
93            None => true,
94            Some(l) => preds.iter().any(|&p| p == l || clock[p][t] > l),
95        };
96        let chained = preds.iter().copied().max().map(|p| out.stream[p]).filter(|&t| last[t] == preds.iter().copied().max());
97        let t = chained
98            .or_else(|| (0..STREAMS).filter(|&t| follows(t)).max_by_key(|&t| last[t].map_or(0, |l| l + 1)))
99            .unwrap_or_else(|| (0..STREAMS).min_by_key(|&t| last[t].map_or(0, |l| l + 1)).expect("a stream"));
100        let mut c = stream_clock[t];
101        // The latest first: waiting on a stream's later launch orders its earlier ones too.
102        for &p in preds.iter().rev() {
103            let ps = out.stream[p];
104            if ps != t && c[ps] <= p {
105                let e = *out.records[p].get_or_insert_with(|| {
106                    out.events += 1;
107                    out.events - 1
108                });
109                if !out.waits[j].contains(&e) {
110                    out.waits[j].push(e);
111                }
112                for (u, v) in c.iter_mut().zip(clock[p]) {
113                    *u = (*u).max(v);
114                }
115                c[ps] = c[ps].max(p + 1);
116            }
117        }
118        c[t] = j + 1;
119        out.stream[j] = t;
120        stream_clock[t] = c;
121        last[t] = Some(j);
122        clock.push(c);
123        if t != 0 && !out.sides.contains(&t) {
124            out.sides.push(t);
125        }
126    }
127    out.sides.sort_unstable();
128    for &t in &out.sides {
129        let l = last[t].expect("a side stream's launch");
130        let e = *out.records[l].get_or_insert_with(|| {
131            out.events += 1;
132            out.events - 1
133        });
134        out.joins.push(e);
135    }
136    out
137}
logicaffeine_compile/codegen_cuda/lane/branches.rs

logicaffeine_compile/codegen_cuda/lane/
branches.rs
