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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
//! Asynchronous, validated device-input search.
use crate::*;
use hrx::{Event, View};
/// Borrowed FP32 queries in row-major device storage. Values are normalized on
/// the GPU; nonfinite or zero-norm rows produce status=1 and zero matches.
/// Stride is in FP32 elements. All accessed rows must be ready on the searcher's
/// stream (use an event for a producer on another stream).
#[derive(Clone, Copy, Debug)]
pub struct DeviceQueries<'a> {
values: View<'a>,
batch: usize,
dimensions: usize,
stride: usize,
}
impl<'a> DeviceQueries<'a> {
/// Validate a 0..=64 query batch, dimensions 1..=16384, and row stride between
/// dimensions and 16384. The binding must cover batch*stride FP32 elements.
pub fn new(values: View<'a>, batch: usize, dimensions: usize, stride: usize) -> Result<Self> {
if batch > MAX_BATCH
|| !(1..=MAX_DIMENSIONS).contains(&dimensions)
|| stride < dimensions
|| stride > MAX_DIMENSIONS
{
return Err(invalid("invalid device query shape or stride"));
}
Ok(Self {
values: values.slice(0, batch * stride * 4)?,
batch,
dimensions,
stride,
})
}
}
pub(crate) struct QueryWorkspace {
query: Buffer,
width: usize,
plans: std::collections::HashMap<usize, Kernel>,
finish: Kernel,
}
impl Searcher {
/// Prepare device normalization, scan, and selection for repeated submissions.
/// First use can allocate/compile and wait while replacing host-imported
/// batch workspace. Fully reserved device submissions perform no host wait.
pub fn reserve_device(&mut self, batch: usize, k: usize) -> Result<()> {
if batch > MAX_BATCH || !(1..=MAX_K).contains(&k) {
return Err(invalid("invalid device search shape"));
}
if batch == 0 {
return Ok(());
}
if batch == 1 {
self.reserve_search(k)?;
} else {
self.reserve_batch(batch, k)?;
}
let width = if batch == 1 {
1
} else {
batch.next_power_of_two().max(8)
};
if self.device_queries.is_none() {
let (finish, _) = compile(
&self.compiler,
&self.stream,
include_str!("../kernels/finish.loom"),
kernels::named_spec("finish"),
)?;
self.device_queries = Some(QueryWorkspace {
query: self.stream.allocate(width * self.padded * 4)?,
width,
plans: Default::default(),
finish,
});
}
let workspace = self.device_queries.as_mut().unwrap();
if workspace.width < width {
workspace.query = self.stream.allocate(width * self.padded * 4)?;
workspace.width = width;
}
if let std::collections::hash_map::Entry::Vacant(entry) = workspace.plans.entry(width) {
let mut spec = kernels::named_spec("prepare_queries");
for (key, value) in [
("query.dimensions", self.dimensions),
("query.padded", self.padded),
("query.width", width),
] {
spec.config.insert(key.into(), value.to_string());
}
let (plan, report) = compile(
&self.compiler,
&self.stream,
include_str!("../kernels/prepare_queries.loom"),
spec,
)?;
entry.insert(plan);
self.reports.push(report);
}
Ok(())
}
/// Queue GPU normalization, exhaustive cosine search, and top-k into owned
/// device outputs. Return a completion event without reading results back.
///
/// `excluded` is an optional shared bitmap (u32 words, least-significant bit
/// first) over insertion IDs, at least ceil(corpus.len()/32)*4 bytes. It is
/// read in place, allowing device producers to update it incrementally.
/// Output batch count must match input. `output.k()` selects k. Invalid
/// queries set output status=1 and count=0; `DeviceNeighbors::read` reports
/// these as errors. Inputs/bitmap must be ready on this searcher's stream.
/// Keep buffers read-only until completion. Slots beyond counts are invalid.
///
/// Scaled FP32 norm accumulation handles extreme finite magnitudes without
/// overflow. Its rounding can differ from host FP64 normalization; very tiny
/// relative components may underflow. Scores use FP32 accumulation.
pub fn search_device(
&mut self,
queries: DeviceQueries<'_>,
excluded: Option<View<'_>>,
output: &mut DeviceNeighbors,
) -> Result<Event> {
if queries.dimensions != self.dimensions {
return Err(invalid("device query dimension mismatch"));
}
output.validate(&self.stream, queries.batch)?;
let excluded = excluded
.map(|v| v.slice(0, self.count.div_ceil(32) * 4))
.transpose()?;
self.reserve_device(queries.batch, output.k)?;
if queries.batch == 0 {
return self.stream.record_event();
}
let width = if queries.batch == 1 {
1
} else {
queries.batch.next_power_of_two().max(8)
};
let workspace = self.device_queries.as_ref().unwrap();
self.stream.fill(workspace.query.binding(), 0)?;
let mut c = Constants::new();
c.push(queries.batch as u32)?;
c.push(queries.stride as u32)?;
// SAFETY: query shape/stride and output extents are validated. GPU
// normalization checks every logical row before producing packed values.
unsafe {
self.stream.dispatch(
&workspace.plans[&width],
[queries.batch as u32, 1, 1],
[128, 1, 1],
&c,
&[queries.values, workspace.query.binding(), output.status()],
)?;
}
let k = output.k.min(self.count);
if k == 0 {
device::finish(&self.stream, &workspace.finish, output, None)?;
} else if queries.batch == 1 {
self.dispatch_scan_from(false, workspace.query.binding())?;
if let Some(bitmap) = excluded {
self.apply_exclusions_from(bitmap)?;
}
let selected = self.select(k)?;
device::finish(
&self.stream,
&workspace.finish,
output,
Some((
self.candidates[selected].0.binding(),
self.candidates[selected].1.binding(),
k,
)),
)?;
} else {
let scratch = self.batch.as_ref().unwrap();
scratch.scan(
&self.stream,
&self.corpus,
workspace.query.binding(),
excluded,
width,
k,
)?;
device::finish(
&self.stream,
&workspace.finish,
output,
Some((scratch.running.0.binding(), scratch.running.1.binding(), k)),
)?;
}
self.stream.record_event()
}
pub(crate) fn device_query_bytes(&self) -> usize {
self.device_queries.as_ref().map_or(0, |w| w.query.bytes())
}
}