hrx-rs 0.8.4

Owned native AMD GPU and NPU execution with Loom compilation
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
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

//! Allocation interval frontiers shared by execution schedulers and model graphs.

use crate::execution::Access;
use std::{
    collections::{BTreeMap, BTreeSet},
    ops::Range,
};

#[derive(Clone)]
pub(crate) struct Use {
    pub allocation: usize,
    pub range: Range<usize>,
    pub access: Access,
}

#[derive(Clone, Default)]
struct Segment {
    range: Range<usize>,
    writer: Option<usize>,
    readers: BTreeSet<usize>,
}

/// Per-allocation interval frontiers retain only dependencies relevant to
/// future overlapping accesses instead of producing a dense graph.
#[derive(Default)]
pub(crate) struct Frontier(BTreeMap<usize, Vec<Segment>>);

impl Frontier {
    pub(crate) fn dependencies(
        &mut self,
        node: usize,
        uses: impl IntoIterator<Item = Use>,
    ) -> Vec<usize> {
        let mut dependencies = BTreeSet::new();
        for usage in uses {
            if usage.range.is_empty() {
                continue;
            }
            let segments = self.0.entry(usage.allocation).or_default();
            for point in [usage.range.start, usage.range.end] {
                if let Some(index) = segments
                    .iter()
                    .position(|segment| segment.range.start < point && point < segment.range.end)
                {
                    let mut right = segments[index].clone();
                    right.range.start = point;
                    segments[index].range.end = point;
                    segments.insert(index + 1, right);
                }
            }
            let mut position = usage.range.start;
            let mut index = segments.partition_point(|segment| segment.range.end <= position);
            while position < usage.range.end {
                if index == segments.len() || segments[index].range.start > position {
                    let end = segments.get(index).map_or(usage.range.end, |segment| {
                        segment.range.start.min(usage.range.end)
                    });
                    segments.insert(
                        index,
                        Segment {
                            range: position..end,
                            ..Segment::default()
                        },
                    );
                }
                let segment = &mut segments[index];
                dependencies.extend(segment.writer);
                if usage.access != Access::Read {
                    dependencies.append(&mut segment.readers);
                    segment.writer = Some(node);
                } else {
                    segment.readers.insert(node);
                }
                position = segment.range.end;
                index += 1;
            }
        }
        // Multiple arguments in one operation may overlap. Those uses update
        // the frontier in declaration order but cannot make the node depend on itself.
        dependencies.remove(&node);
        dependencies.into_iter().collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn usage(allocation: usize, range: Range<usize>, access: Access) -> Use {
        Use {
            allocation,
            range,
            access,
        }
    }

    #[test]
    fn preserves_only_overlapping_hazards() {
        let mut frontier = Frontier::default();
        assert!(
            frontier
                .dependencies(0, [usage(0, 0..64, Access::Write)])
                .is_empty()
        );
        assert_eq!(
            frontier.dependencies(1, [usage(0, 0..32, Access::Read)]),
            [0]
        );
        assert!(
            frontier
                .dependencies(2, [usage(0, 64..96, Access::Read)])
                .is_empty()
        );
        assert_eq!(
            frontier.dependencies(3, [usage(0, 16..80, Access::Write)]),
            [0, 1, 2]
        );
        assert_eq!(
            frontier.dependencies(4, [usage(0, 24..32, Access::Read)]),
            [3]
        );
        assert!(
            frontier
                .dependencies(5, [usage(1, 0..32, Access::Write)])
                .is_empty()
        );
    }

    #[test]
    fn one_node_can_read_and_write_the_same_region() {
        let mut frontier = Frontier::default();
        assert!(
            frontier
                .dependencies(
                    0,
                    [
                        usage(0, 0..32, Access::Read),
                        usage(0, 0..32, Access::Write),
                    ],
                )
                .is_empty()
        );
        assert_eq!(
            frontier.dependencies(1, [usage(0, 0..32, Access::Read)]),
            [0]
        );
    }
}