trueno_graph/gpu/
paged_bfs.rs

1//! Paged GPU BFS for graphs larger than VRAM
2//!
3//! Processes graph tiles incrementally using LRU caching.
4//! Based on morsel-driven parallelism (Umbra, Neumann & Freitag 2020).
5
6use super::cache::LruTileCache;
7use super::paging::PagingCoordinator;
8use super::{GpuBfsResult, GpuDevice};
9use crate::{CsrGraph, NodeId};
10use anyhow::{Context, Result};
11use std::collections::VecDeque;
12
13/// Run paged BFS on a graph that may exceed VRAM capacity
14///
15/// Automatically tiles the graph and processes tiles with LRU caching.
16///
17/// # Arguments
18///
19/// * `device` - GPU device
20/// * `graph` - Graph to process
21/// * `source` - Source node for BFS
22///
23/// # Errors
24///
25/// Returns error if:
26/// - GPU operations fail
27/// - Graph is empty
28///
29/// # Example
30///
31/// ```ignore
32/// # use trueno_graph::gpu::{GpuDevice, gpu_bfs_paged};
33/// # use trueno_graph::{CsrGraph, NodeId};
34/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
35/// let device = GpuDevice::new().await?;
36/// let graph = CsrGraph::new();
37/// // ... add many edges to create large graph ...
38///
39/// let result = gpu_bfs_paged(&device, &graph, NodeId(0)).await?;
40/// println!("Visited {} nodes", result.visited_count);
41/// # Ok(())
42/// # }
43/// ```
44/// Process a single node's neighbors within its tile, updating distances
45#[allow(clippy::cast_possible_truncation)]
46fn process_node_neighbors(
47    coordinator: &PagingCoordinator,
48    node: NodeId,
49    distances: &mut [u32],
50    current_level: u32,
51    next_frontier: &mut Vec<NodeId>,
52) -> Result<()> {
53    let tile = coordinator.get_tile_for_node(node).context("Node not in any tile")?;
54    let node_idx_in_tile = node.0 as usize - tile.start_node;
55
56    if tile.row_offsets.len() < 2 || node_idx_in_tile >= tile.row_offsets.len() - 1 {
57        return Ok(());
58    }
59
60    let start = tile.row_offsets[node_idx_in_tile] as usize;
61    let end = tile.row_offsets[node_idx_in_tile + 1] as usize;
62
63    for &neighbor in &tile.col_indices[start..end] {
64        let neighbor_idx = neighbor as usize;
65        if distances[neighbor_idx] == u32::MAX {
66            distances[neighbor_idx] = current_level + 1;
67            next_frontier.push(NodeId(neighbor));
68        }
69    }
70    Ok(())
71}
72
73/// Run GPU breadth-first search over a graph that may exceed VRAM.
74///
75/// When the graph fits in GPU memory this delegates to the in-VRAM
76/// [`super::gpu_bfs`]; otherwise it tiles the graph via [`PagingCoordinator`] and
77/// processes each tile in turn.
78///
79/// # Errors
80///
81/// Returns an error if GPU paging setup or any tile traversal fails.
82#[allow(clippy::cast_possible_truncation)]
83pub async fn gpu_bfs_paged(
84    device: &GpuDevice,
85    graph: &CsrGraph,
86    source: NodeId,
87) -> Result<GpuBfsResult> {
88    let coordinator = PagingCoordinator::new(device, graph)?;
89
90    if coordinator.fits_in_vram() {
91        return super::gpu_bfs(
92            device,
93            &super::GpuCsrBuffers::from_csr_graph(device, graph)?,
94            source,
95        )
96        .await;
97    }
98
99    let num_nodes = graph.num_nodes();
100    let mut distances = vec![u32::MAX; num_nodes];
101    distances[source.0 as usize] = 0;
102
103    let mut frontier = VecDeque::new();
104    frontier.push_back(source);
105    let mut current_level = 0_u32;
106
107    let cache_capacity = coordinator.limits().max_morsels.min(coordinator.num_tiles());
108    let mut _tile_cache = LruTileCache::new(cache_capacity);
109
110    while !frontier.is_empty() {
111        let mut next_frontier = Vec::new();
112
113        for &node in &frontier {
114            process_node_neighbors(
115                &coordinator,
116                node,
117                &mut distances,
118                current_level,
119                &mut next_frontier,
120            )?;
121        }
122
123        frontier = VecDeque::from(next_frontier);
124        current_level += 1;
125
126        if current_level > num_nodes as u32 {
127            break;
128        }
129    }
130
131    let visited_count = distances.iter().filter(|&&d| d != u32::MAX).count();
132    Ok(GpuBfsResult { distances, visited_count })
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[tokio::test]
140    async fn test_paged_bfs_small_graph() {
141        if !GpuDevice::is_gpu_available().await {
142            eprintln!("⚠️  Skipping test_paged_bfs_small_graph: GPU not available");
143            return;
144        }
145
146        let device = GpuDevice::new().await.unwrap();
147
148        // Create small graph: 0 -> 1 -> 2
149        let mut graph = CsrGraph::new();
150        graph.add_edge(NodeId(0), NodeId(1), 1.0).unwrap();
151        graph.add_edge(NodeId(1), NodeId(2), 1.0).unwrap();
152
153        let result = gpu_bfs_paged(&device, &graph, NodeId(0)).await.unwrap();
154
155        assert_eq!(result.distance(NodeId(0)), Some(0));
156        assert_eq!(result.distance(NodeId(1)), Some(1));
157        assert_eq!(result.distance(NodeId(2)), Some(2));
158        assert_eq!(result.visited_count, 3);
159    }
160
161    #[tokio::test]
162    async fn test_paged_bfs_disconnected() {
163        if !GpuDevice::is_gpu_available().await {
164            eprintln!("⚠️  Skipping test_paged_bfs_disconnected: GPU not available");
165            return;
166        }
167
168        let device = GpuDevice::new().await.unwrap();
169
170        // Create disconnected: 0 -> 1, 2 (isolated)
171        let mut graph = CsrGraph::new();
172        graph.add_edge(NodeId(0), NodeId(1), 1.0).unwrap();
173        graph.add_edge(NodeId(2), NodeId(2), 1.0).unwrap(); // Self-loop
174
175        let result = gpu_bfs_paged(&device, &graph, NodeId(0)).await.unwrap();
176
177        assert_eq!(result.distance(NodeId(0)), Some(0));
178        assert_eq!(result.distance(NodeId(1)), Some(1));
179        assert_eq!(result.distance(NodeId(2)), None); // Unreachable
180        assert_eq!(result.visited_count, 2);
181    }
182
183    #[tokio::test]
184    async fn test_paged_bfs_larger_graph() {
185        if !GpuDevice::is_gpu_available().await {
186            eprintln!("⚠️  Skipping test_paged_bfs_larger_graph: GPU not available");
187            return;
188        }
189
190        let device = GpuDevice::new().await.unwrap();
191
192        // Create larger ring graph
193        let mut graph = CsrGraph::new();
194        for i in 0..100 {
195            graph.add_edge(NodeId(i), NodeId((i + 1) % 100), 1.0).unwrap();
196        }
197
198        let result = gpu_bfs_paged(&device, &graph, NodeId(0)).await.unwrap();
199
200        // All nodes should be reachable in a ring
201        assert_eq!(result.visited_count, 100);
202        assert_eq!(result.distance(NodeId(0)), Some(0));
203        assert_eq!(result.distance(NodeId(1)), Some(1));
204        assert_eq!(result.distance(NodeId(50)), Some(50));
205    }
206
207    #[tokio::test]
208    async fn test_paged_bfs_star_graph() {
209        if !GpuDevice::is_gpu_available().await {
210            eprintln!("⚠️  Skipping test_paged_bfs_star_graph: GPU not available");
211            return;
212        }
213
214        let device = GpuDevice::new().await.unwrap();
215
216        // Create star graph: center node 0 connected to all others
217        let mut graph = CsrGraph::new();
218        for i in 1..20 {
219            graph.add_edge(NodeId(0), NodeId(i), 1.0).unwrap();
220        }
221
222        let result = gpu_bfs_paged(&device, &graph, NodeId(0)).await.unwrap();
223
224        assert_eq!(result.distance(NodeId(0)), Some(0));
225        // All other nodes at distance 1
226        for i in 1..20 {
227            assert_eq!(result.distance(NodeId(i)), Some(1));
228        }
229        assert_eq!(result.visited_count, 20);
230    }
231
232    #[tokio::test]
233    async fn test_paged_bfs_multiple_levels() {
234        if !GpuDevice::is_gpu_available().await {
235            eprintln!("⚠️  Skipping test_paged_bfs_multiple_levels: GPU not available");
236            return;
237        }
238
239        let device = GpuDevice::new().await.unwrap();
240
241        // Create multi-level tree
242        let mut graph = CsrGraph::new();
243        // Level 0: node 0
244        // Level 1: nodes 1, 2
245        // Level 2: nodes 3, 4, 5, 6
246        graph.add_edge(NodeId(0), NodeId(1), 1.0).unwrap();
247        graph.add_edge(NodeId(0), NodeId(2), 1.0).unwrap();
248        graph.add_edge(NodeId(1), NodeId(3), 1.0).unwrap();
249        graph.add_edge(NodeId(1), NodeId(4), 1.0).unwrap();
250        graph.add_edge(NodeId(2), NodeId(5), 1.0).unwrap();
251        graph.add_edge(NodeId(2), NodeId(6), 1.0).unwrap();
252
253        let result = gpu_bfs_paged(&device, &graph, NodeId(0)).await.unwrap();
254
255        assert_eq!(result.distance(NodeId(0)), Some(0));
256        assert_eq!(result.distance(NodeId(1)), Some(1));
257        assert_eq!(result.distance(NodeId(2)), Some(1));
258        assert_eq!(result.distance(NodeId(3)), Some(2));
259        assert_eq!(result.distance(NodeId(4)), Some(2));
260        assert_eq!(result.distance(NodeId(5)), Some(2));
261        assert_eq!(result.distance(NodeId(6)), Some(2));
262        assert_eq!(result.visited_count, 7);
263    }
264
265    #[tokio::test]
266    async fn test_paged_bfs_with_duplicate_edges() {
267        if !GpuDevice::is_gpu_available().await {
268            eprintln!("⚠️  Skipping test_paged_bfs_with_duplicate_edges: GPU not available");
269            return;
270        }
271
272        let device = GpuDevice::new().await.unwrap();
273
274        // Graph with multiple paths between nodes
275        let mut graph = CsrGraph::new();
276        graph.add_edge(NodeId(0), NodeId(1), 1.0).unwrap();
277        graph.add_edge(NodeId(0), NodeId(2), 1.0).unwrap();
278        graph.add_edge(NodeId(1), NodeId(3), 1.0).unwrap();
279        graph.add_edge(NodeId(2), NodeId(3), 1.0).unwrap(); // Two paths to 3
280
281        let result = gpu_bfs_paged(&device, &graph, NodeId(0)).await.unwrap();
282
283        assert_eq!(result.distance(NodeId(0)), Some(0));
284        assert_eq!(result.distance(NodeId(1)), Some(1));
285        assert_eq!(result.distance(NodeId(2)), Some(1));
286        assert_eq!(result.distance(NodeId(3)), Some(2)); // Shortest path
287        assert_eq!(result.visited_count, 4);
288    }
289
290    #[tokio::test]
291    async fn test_paged_bfs_empty_graph() {
292        if !GpuDevice::is_gpu_available().await {
293            eprintln!("⚠️  Skipping test_paged_bfs_empty_graph: GPU not available");
294            return;
295        }
296
297        let device = GpuDevice::new().await.unwrap();
298
299        // Single node, no edges
300        let mut graph = CsrGraph::new();
301        graph.add_edge(NodeId(0), NodeId(0), 1.0).unwrap(); // Self-loop to create node
302
303        let result = gpu_bfs_paged(&device, &graph, NodeId(0)).await.unwrap();
304
305        assert_eq!(result.distance(NodeId(0)), Some(0));
306        assert_eq!(result.visited_count, 1);
307    }
308}
trueno_graph/gpu/paged_bfs.rs

trueno_graph/gpu/
paged_bfs.rs
