Skip to main content

hrx/
access_graph.rs

1//! Raw GPU graphs with dependencies inferred from byte-range access.
2
3use crate::{
4    Constants, Graph, GraphExec, Kernel, Node, Result, Stream, View,
5    dependency_frontier::{Frontier, Use},
6    execution::Access,
7};
8
9/// One graph binding and the complete access performed through it.
10#[derive(Clone, Copy, Debug)]
11pub struct AccessView<'a> {
12    /// Bound allocation range.
13    pub view: View<'a>,
14    /// Reads and writes performed through the binding.
15    pub access: Access,
16}
17
18/// A raw GPU graph whose edges are inferred from declared buffer access.
19pub struct AccessGraph<'a> {
20    graph: Graph<'a>,
21    frontier: Frontier,
22    nodes: Vec<Node>,
23}
24
25impl Stream {
26    /// Begin an access-aware graph recording on this stream.
27    pub fn access_graph(&self) -> Result<AccessGraph<'_>> {
28        Ok(AccessGraph {
29            graph: self.graph()?,
30            frontier: Frontier::default(),
31            nodes: Vec::new(),
32        })
33    }
34}
35
36impl<'a> View<'a> {
37    /// Declare how a graph operation accesses this view.
38    #[must_use]
39    pub fn access(self, access: Access) -> AccessView<'a> {
40        AccessView { view: self, access }
41    }
42
43    /// Declare read-only graph access.
44    #[must_use]
45    pub fn read(self) -> AccessView<'a> {
46        self.access(Access::Read)
47    }
48
49    /// Declare write-only graph access.
50    #[must_use]
51    pub fn write(self) -> AccessView<'a> {
52        self.access(Access::Write)
53    }
54
55    /// Declare graph access that reads and writes existing bytes.
56    #[must_use]
57    pub fn read_write(self) -> AccessView<'a> {
58        self.access(Access::ReadWrite)
59    }
60}
61
62impl<'a> AccessGraph<'a> {
63    fn dependencies(&mut self, uses: impl IntoIterator<Item = AccessView<'a>>) -> Vec<Node> {
64        let index = self.nodes.len();
65        self.frontier
66            .dependencies(
67                index,
68                uses.into_iter().map(|binding| Use {
69                    allocation: binding.view.owner() as *const crate::Buffer as usize,
70                    range: binding.view.offset()
71                        ..binding.view.offset().saturating_add(binding.view.len()),
72                    access: binding.access,
73                }),
74            )
75            .into_iter()
76            .map(|dependency| self.nodes[dependency])
77            .collect()
78    }
79
80    fn record(&mut self, node: Node) -> Node {
81        self.nodes.push(node);
82        node
83    }
84
85    /// Record a byte-pattern fill.
86    pub fn fill(&mut self, destination: View<'a>, pattern: u8) -> Result<Node> {
87        let after = self.dependencies([destination.write()]);
88        let node = self.graph.fill(&after, destination, pattern)?;
89        Ok(self.record(node))
90    }
91
92    /// Record a copy between equal spans.
93    pub fn copy(&mut self, destination: View<'a>, source: View<'a>) -> Result<Node> {
94        let after = self.dependencies([source.read(), destination.write()]);
95        let node = self.graph.copy(&after, destination, source)?;
96        Ok(self.record(node))
97    }
98
99    /// Record a kernel dispatch.
100    ///
101    /// # Safety
102    ///
103    /// In addition to [`Graph::dispatch`]'s contract, every binding must declare
104    /// all bytes and the complete access performed by the kernel.
105    pub unsafe fn dispatch(
106        &mut self,
107        kernel: &'a Kernel,
108        grid: [u32; 3],
109        block: [u32; 3],
110        constants: &Constants,
111        bindings: &[AccessView<'a>],
112    ) -> Result<Node> {
113        let after = self.dependencies(bindings.iter().copied());
114        let views = bindings
115            .iter()
116            .map(|binding| binding.view)
117            .collect::<Vec<_>>();
118        let node = unsafe {
119            self.graph
120                .dispatch(&after, kernel, grid, block, constants, &views)?
121        };
122        Ok(self.record(node))
123    }
124
125    /// Instantiate the recorded graph.
126    pub fn finish(self) -> Result<GraphExec> {
127        self.graph.finish()
128    }
129
130    /// Number of recorded operations.
131    #[must_use]
132    pub fn len(&self) -> usize {
133        self.nodes.len()
134    }
135
136    /// Whether no operations have been recorded.
137    #[must_use]
138    pub fn is_empty(&self) -> bool {
139        self.nodes.is_empty()
140    }
141}