qualia_core_db/tensor/
volume_gpu.rs

1//! GPU tensor volume search (B8) — mirrors `visit_tensor_search_into` on U1 queue lane.
2
3use bytemuck::{Pod, Zeroable};
4
5#[cfg(not(target_arch = "wasm32"))]
6use super::resident_substrate::{global_resident_substrate, MAX_KNN_HITS, MAX_RESIDENT_NODES};
7use super::Tensor10D;
8
9pub const TENSOR_VOLUME_STRIDE_FLOATS: u32 = 10;
10
11#[repr(C)]
12#[derive(Clone, Copy, Pod, Zeroable)]
13struct VolumeGpuParams {
14    node_count: u32,
15    max_distance: f32,
16    stride_floats: u32,
17    max_hits: u32,
18}
19
20#[cfg(not(target_arch = "wasm32"))]
21pub struct TensorVolumeGpu {
22    pipeline: wgpu::ComputePipeline,
23    query_buf: wgpu::Buffer,
24    nodes_buf: wgpu::Buffer,
25    params_buf: wgpu::Buffer,
26    hits_buf: wgpu::Buffer,
27    count_buf: wgpu::Buffer,
28    staging_hits: wgpu::Buffer,
29    staging_count: wgpu::Buffer,
30    max_nodes: u32,
31}
32
33#[cfg(not(target_arch = "wasm32"))]
34impl TensorVolumeGpu {
35    pub fn try_new(device: &wgpu::Device) -> Option<Self> {
36        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
37            label: Some("TensorVolumeShader"),
38            source: wgpu::ShaderSource::Wgsl(include_str!("../shaders/tensor_volume.wgsl").into()),
39        });
40        let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
41            label: Some("TensorVolumePipeline"),
42            layout: None,
43            module: &shader,
44            entry_point: Some("main"),
45            compilation_options: Default::default(),
46            cache: None,
47        });
48        let max_nodes = MAX_RESIDENT_NODES as u32;
49        let node_floats = max_nodes * TENSOR_VOLUME_STRIDE_FLOATS;
50        Some(Self {
51            pipeline,
52            query_buf: device.create_buffer(&wgpu::BufferDescriptor {
53                label: Some("TensorVolumeQuery"),
54                size: std::mem::size_of::<Tensor10D>() as wgpu::BufferAddress,
55                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
56                mapped_at_creation: false,
57            }),
58            nodes_buf: device.create_buffer(&wgpu::BufferDescriptor {
59                label: Some("TensorVolumeNodes"),
60                size: (node_floats as usize * 4) as wgpu::BufferAddress,
61                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
62                mapped_at_creation: false,
63            }),
64            params_buf: device.create_buffer(&wgpu::BufferDescriptor {
65                label: Some("TensorVolumeParams"),
66                size: std::mem::size_of::<VolumeGpuParams>() as wgpu::BufferAddress,
67                usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
68                mapped_at_creation: false,
69            }),
70            hits_buf: device.create_buffer(&wgpu::BufferDescriptor {
71                label: Some("TensorVolumeHits"),
72                size: (MAX_KNN_HITS * 4) as wgpu::BufferAddress,
73                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
74                mapped_at_creation: false,
75            }),
76            count_buf: device.create_buffer(&wgpu::BufferDescriptor {
77                label: Some("TensorVolumeHitCount"),
78                size: 4,
79                // COPY_DST: reset to 0 each cycle via queue.write_buffer (line ~145);
80                // COPY_SRC: copied to staging_count for readback; STORAGE: shader binding.
81                usage: wgpu::BufferUsages::STORAGE
82                    | wgpu::BufferUsages::COPY_SRC
83                    | wgpu::BufferUsages::COPY_DST,
84                mapped_at_creation: false,
85            }),
86            staging_hits: device.create_buffer(&wgpu::BufferDescriptor {
87                label: Some("TensorVolumeHitsStaging"),
88                size: (MAX_KNN_HITS * 4) as wgpu::BufferAddress,
89                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
90                mapped_at_creation: false,
91            }),
92            staging_count: device.create_buffer(&wgpu::BufferDescriptor {
93                label: Some("TensorVolumeCountStaging"),
94                size: 4,
95                usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
96                mapped_at_creation: false,
97            }),
98            max_nodes,
99        })
100    }
101
102    fn upload_nodes(&self, queue: &wgpu::Queue, node_count: u32) -> bool {
103        let substrate = global_resident_substrate();
104        let count = node_count.min(substrate.node_count()).min(self.max_nodes);
105        if count == 0 {
106            return false;
107        }
108        let mut flat = [0f32; MAX_RESIDENT_NODES * TENSOR_VOLUME_STRIDE_FLOATS as usize];
109        for i in 0..count as usize {
110            if let Some(t) = substrate.tensor_at(i as u32) {
111                let base = i * TENSOR_VOLUME_STRIDE_FLOATS as usize;
112                flat[base] = t.q;
113                flat[base + 1] = t.v;
114                flat[base + 2] = t.w;
115                flat[base + 3] = t.x;
116                flat[base + 4] = t.y;
117                flat[base + 5] = t.z;
118                flat[base + 6] = t.t;
119                flat[base + 7] = t.alpha;
120                flat[base + 8] = t.mu;
121                flat[base + 9] = t.sigma;
122            }
123        }
124        let bytes =
125            (count as usize * TENSOR_VOLUME_STRIDE_FLOATS as usize * 4) as wgpu::BufferAddress;
126        queue.write_buffer(
127            &self.nodes_buf,
128            0,
129            bytemuck::cast_slice(&flat[..count as usize * 10]),
130        );
131        let _ = bytes;
132        true
133    }
134
135    /// GPU kNN filter; returns hit count written into `out` (caller stack buffer).
136    pub fn tensor_search_into(
137        &self,
138        device: &wgpu::Device,
139        queue: &wgpu::Queue,
140        query: &Tensor10D,
141        max_distance: f32,
142        out: &mut [usize],
143    ) -> usize {
144        let node_count = global_resident_substrate().node_count().min(self.max_nodes);
145        if node_count == 0 || out.is_empty() || !self.upload_nodes(queue, node_count) {
146            return 0;
147        }
148
149        queue.write_buffer(&self.query_buf, 0, bytemuck::bytes_of(query));
150        let params = VolumeGpuParams {
151            node_count,
152            max_distance,
153            stride_floats: TENSOR_VOLUME_STRIDE_FLOATS,
154            max_hits: MAX_KNN_HITS as u32,
155        };
156        queue.write_buffer(&self.params_buf, 0, bytemuck::bytes_of(&params));
157        queue.write_buffer(&self.count_buf, 0, &[0u8; 4]);
158
159        let bind = device.create_bind_group(&wgpu::BindGroupDescriptor {
160            label: Some("TensorVolumeBindGroup"),
161            layout: &self.pipeline.get_bind_group_layout(0),
162            entries: &[
163                wgpu::BindGroupEntry {
164                    binding: 0,
165                    resource: self.query_buf.as_entire_binding(),
166                },
167                wgpu::BindGroupEntry {
168                    binding: 1,
169                    resource: self.nodes_buf.as_entire_binding(),
170                },
171                wgpu::BindGroupEntry {
172                    binding: 2,
173                    resource: self.params_buf.as_entire_binding(),
174                },
175                wgpu::BindGroupEntry {
176                    binding: 3,
177                    resource: self.hits_buf.as_entire_binding(),
178                },
179                wgpu::BindGroupEntry {
180                    binding: 4,
181                    resource: self.count_buf.as_entire_binding(),
182                },
183            ],
184        });
185
186        let wg = (node_count + 63) / 64;
187        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
188            label: Some("TensorVolumeEncoder"),
189        });
190        {
191            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
192                label: Some("TensorVolumePass"),
193                timestamp_writes: None,
194            });
195            pass.set_pipeline(&self.pipeline);
196            pass.set_bind_group(0, &bind, &[]);
197            pass.dispatch_workgroups(wg, 1, 1);
198        }
199        encoder.copy_buffer_to_buffer(
200            &self.hits_buf,
201            0,
202            &self.staging_hits,
203            0,
204            (MAX_KNN_HITS * 4) as u64,
205        );
206        encoder.copy_buffer_to_buffer(&self.count_buf, 0, &self.staging_count, 0, 4);
207        queue.submit(Some(encoder.finish()));
208
209        let _ = device.poll(wgpu::PollType::wait_indefinitely());
210        let count_slice = self.staging_count.slice(..4);
211        let (tx, rx) = std::sync::mpsc::channel();
212        count_slice.map_async(wgpu::MapMode::Read, move |r| {
213            let _ = tx.send(r);
214        });
215        let _ = device.poll(wgpu::PollType::wait_indefinitely());
216        if rx.recv().ok().and_then(|r| r.ok()).is_none() {
217            return 0;
218        }
219        let count_data = count_slice
220            .get_mapped_range()
221            .expect("wgpu buffer map_range failed");
222        let total = u32::from_le_bytes(count_data[..4].try_into().unwrap_or([0; 4])) as usize;
223        drop(count_data);
224        self.staging_count.unmap();
225
226        let hits_slice = self.staging_hits.slice(..(MAX_KNN_HITS * 4) as u64);
227        let (tx2, rx2) = std::sync::mpsc::channel();
228        hits_slice.map_async(wgpu::MapMode::Read, move |r| {
229            let _ = tx2.send(r);
230        });
231        let _ = device.poll(wgpu::PollType::wait_indefinitely());
232        if rx2.recv().ok().and_then(|r| r.ok()).is_none() {
233            return 0;
234        }
235        let hits_data = hits_slice
236            .get_mapped_range()
237            .expect("wgpu buffer map_range failed");
238        let indices: &[u32] = bytemuck::cast_slice(&hits_data);
239        let n = total.min(out.len()).min(indices.len());
240        for i in 0..n {
241            out[i] = indices[i] as usize;
242        }
243        drop(hits_data);
244        self.staging_hits.unmap();
245        n
246    }
247}
248
249#[cfg(not(target_arch = "wasm32"))]
250static VOLUME_GPU: std::sync::OnceLock<Option<TensorVolumeGpu>> = std::sync::OnceLock::new();
251
252/// Try GPU tensor search; returns None when GPU path unavailable (caller uses SIMD).
253#[cfg(not(target_arch = "wasm32"))]
254pub fn try_gpu_tensor_search_into(
255    query: &Tensor10D,
256    max_distance: f32,
257    out: &mut [usize],
258) -> Option<usize> {
259    let gpu_ctx = crate::gpu_context::shared_gpu();
260    let vol = VOLUME_GPU.get_or_init(|| TensorVolumeGpu::try_new(&gpu_ctx.device));
261    let vol = vol.as_ref()?;
262    let n = vol.tensor_search_into(
263        &gpu_ctx.device,
264        crate::gpu_context::shared_gpu()
265            .queue_for_universe(crate::gpu_context::ComputeUniverse::Tensor10D),
266        query,
267        max_distance,
268        out,
269    );
270    if n > 0 {
271        Some(n)
272    } else {
273        None
274    }
275}
276
277#[cfg(target_arch = "wasm32")]
278pub fn try_gpu_tensor_search_into(
279    _query: &Tensor10D,
280    _max_distance: f32,
281    _out: &mut [usize],
282) -> Option<usize> {
283    None
284}
285
286/// CPU reference for the GPU tensor-volume search (ALGEBRA_MANIFOLD_PLAN.md Phase 4.2).
287///
288/// A plain linear scan using the CANONICAL `Tensor10D::full_distance` metric — the same
289/// metric `shaders/tensor_volume.wgsl` now ports (chosen by the QUERY's `v` topology
290/// class: euclidean / cyclic / hyperbolic / boundary). All three paths — GPU, this
291/// reference, and the substrate CPU fallback — therefore agree for ALL `v`. Makes the
292/// 10D manifold search testable without a GPU.
293///
294/// Semantics match the shader: a node is a hit iff `distance <= max_distance`; indices
295/// are written into `out` up to its capacity, and the FULL match count is returned
296/// (which may exceed `out.len()`, exactly like the shader's `hit_count`).
297pub fn cpu_tensor_search_into(
298    query: &Tensor10D,
299    nodes: &[Tensor10D],
300    max_distance: f32,
301    out: &mut [usize],
302) -> usize {
303    let mut matches = 0usize;
304    for (idx, node) in nodes.iter().enumerate() {
305        if query.full_distance(node) <= max_distance {
306            if matches < out.len() {
307                out[matches] = idx;
308            }
309            matches += 1;
310        }
311    }
312    matches
313}
314
315#[cfg(test)]
316mod cpu_reference_tests {
317    use super::*;
318
319    #[test]
320    fn cpu_tensor_search_matches_metric() {
321        // Nodes vary along single dimensions; query is the origin tensor. The metric
322        // uses x,y,z,t,alpha,mu,sigma and ignores q,v,w.
323        let zeros = Tensor10D::new(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);
324        let nodes = [
325            zeros,                                                            // 0: dist 0
326            Tensor10D::new(0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), // 1: x=1 → dist 1
327            Tensor10D::new(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.0), // 2: sigma=3 → dist 3
328            // q,v,w set but METRIC-IRRELEVANT → still dist 0, must be a hit.
329            Tensor10D::new(9.0, 9.0, 9.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0), // 3: dist 0
330        ];
331        let mut out = [usize::MAX; 8];
332
333        // radius 1.5 → nodes 0,1,3 (not 2 at dist 3).
334        let n = cpu_tensor_search_into(&zeros, &nodes, 1.5, &mut out);
335        assert_eq!(n, 3);
336        let hits: std::collections::BTreeSet<usize> = out[..n].iter().copied().collect();
337        assert_eq!(hits, [0, 1, 3].into_iter().collect());
338
339        // radius 5 → all four.
340        assert_eq!(cpu_tensor_search_into(&zeros, &nodes, 5.0, &mut out), 4);
341        // radius 0 → only the exact matches (nodes 0 and 3).
342        assert_eq!(cpu_tensor_search_into(&zeros, &nodes, 0.0, &mut out), 2);
343    }
344
345    #[test]
346    fn cpu_tensor_search_honors_topology_class() {
347        // Proves the v-switched metric unification (§4.1). A CYCLIC query (v == 1) wraps
348        // x modulo 1, so x=0 and x=0.9 are 0.1 apart — NOT 0.9 as under euclidean.
349        // new(q, v, w, x, y, z, t, alpha, mu, sigma)
350        let query = Tensor10D::new(0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);
351        let near_wrap = Tensor10D::new(0.0, 1.0, 0.0, 0.9, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);
352        let mut out = [usize::MAX; 4];
353        // radius 0.2: cyclic distance 0.1 → HIT (euclidean 0.9 would miss).
354        assert_eq!(
355            cpu_tensor_search_into(&query, &[near_wrap], 0.2, &mut out),
356            1,
357            "cyclic metric must wrap x; node should be within 0.2"
358        );
359        // Same geometry under a euclidean query (v == 0) → MISS at radius 0.2.
360        let euclid_query = Tensor10D::new(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);
361        let euclid_node = Tensor10D::new(0.0, 0.0, 0.0, 0.9, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0);
362        assert_eq!(
363            cpu_tensor_search_into(&euclid_query, &[euclid_node], 0.2, &mut out),
364            0
365        );
366    }
367}
qualia_core_db/tensor/volume_gpu.rs

qualia_core_db/tensor/
volume_gpu.rs
