device_search/
device_search.rs

1//! Queue searches and update visited IDs on the GPU, then read only the final result.
2use hrxdb::{Corpus, Device, DeviceExclusions, DeviceNeighbors, DeviceQueries};
3
4fn main() -> hrxdb::Result<()> {
5    let device = Device::open(0)?;
6    let corpus = Corpus::build(
7        &device,
8        3,
9        [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]],
10    )?;
11    let mut worker = corpus.searcher()?;
12    worker.reserve_device(1, 1)?;
13    let query = worker.stream().allocate(12)?;
14    worker.stream().upload(
15        query.binding(),
16        &[1.0f32, 0.0, 0.0]
17            .into_iter()
18            .flat_map(f32::to_le_bytes)
19            .collect::<Vec<_>>(),
20    )?;
21    let mut visited = DeviceExclusions::new(worker.stream(), corpus.len())?;
22    let mut result = DeviceNeighbors::new(worker.stream(), 1, 1)?;
23    for _ in 0..2 {
24        let _done = worker.search_device(
25            DeviceQueries::new(query.binding(), 1, 3, 3)?,
26            Some(visited.binding()),
27            &mut result,
28        )?;
29        visited.insert_device(worker.stream(), result.ids(), 1)?;
30    }
31    let done = worker.search_device(
32        DeviceQueries::new(query.binding(), 1, 3, 3)?,
33        Some(visited.binding()),
34        &mut result,
35    )?;
36    let mut consumer = corpus.stream()?;
37    consumer.wait_event(&done)?;
38    let matches = result.read(&mut consumer)?;
39    assert_eq!(matches[0][0].id, 2);
40    println!("third unvisited match: {:?}", matches[0][0]);
41    Ok(())
42}
device_search/device_search.rs

device_search/
device_search.rs
