hrxdb 0.2.0

//! Owned, coordinated device queries and bounded result slots.
use crate::{Corpus, DeviceNeighbors, DeviceQueries, Neighbor, Result, invalid};
use hrx::{
    Access,
    execution::GpuAccess,
    inference::{Inference, ModelContext, PreparedModel},
    tensor::{DType, DeviceTensor, Layout, TensorDesc},
};

/// A fixed query shape with shared corpus storage and a bounded worker count.
pub struct PreparedSearch {
    plan: PreparedModel,
}
impl Corpus {
    /// Prepare shared-context query normalization, exhaustive scan and top-k.
    /// Each slot owns private search workspace. `workers=1` is the conservative
    /// bandwidth-bound default; increasing it never duplicates the corpus.
    pub fn prepare_search(
        &self,
        context: &ModelContext,
        batch: usize,
        k: usize,
        workers: usize,
    ) -> Result<PreparedSearch> {
        if batch == 0 || batch > crate::MAX_BATCH || !(1..=crate::MAX_K).contains(&k) {
            return Err(invalid("invalid coordinated search shape"));
        }
        if self.inner.device_id
            != hrx::Device::open(context.runtime().gpu()?.index())?
                .stream()?
                .device_id()
        {
            return Err(invalid("corpus belongs to another device"));
        }
        let dimensions = self.dimensions();
        let input =
            TensorDesc::new(DType::F32, vec![batch, dimensions])?.with_layout(Layout::Rows)?;
        let output = [
            TensorDesc::new(DType::U32, vec![batch])?,
            TensorDesc::new(DType::U32, vec![batch])?,
            TensorDesc::new(DType::F32, vec![batch, k])?.with_layout(Layout::Rows)?,
            TensorDesc::new(DType::U32, vec![batch, k])?.with_layout(Layout::Rows)?,
        ];
        let plan = PreparedModel::prepare(
            context,
            &[input],
            &output,
            workers,
            |context, inputs, outputs| {
                let mut corpus = self.clone();
                if let Some(budget) = context.runtime().memory_budget() {
                    corpus = corpus.with_workspace_budget(budget.clone());
                }
                let mut searcher = corpus.searcher()?;
                searcher.reserve_device(batch, k)?;
                let mut results = DeviceNeighbors::new(&searcher.stream, batch, k)?;
                searcher.stream.synchronize()?;
                let bindings = std::iter::once(GpuAccess {
                    view: inputs[0].binding().unwrap(),
                    access: Access::Read,
                })
                .chain(outputs.iter().map(|tensor| GpuAccess {
                    view: tensor.binding().unwrap(),
                    access: Access::Write,
                }))
                .collect::<Vec<_>>();
                let mut graph = context.runtime().graph();
                // SAFETY: The closure owns its private searcher, result allocation
                // and corpus clone. Only input is read; all four declared outputs
                // are fully copied. Every stream operation completes before return.
                unsafe {
                    graph.gpu_scoped(&bindings, move |views| {
                        let queries = DeviceQueries::new(views[0], batch, dimensions, dimensions)?;
                        searcher.search_device(queries, None, &mut results)?;
                        for (destination, source) in views[1..].iter().zip([
                            results.status(),
                            results.counts(),
                            results.scores(),
                            results.ids(),
                        ]) {
                            searcher.stream.copy(*destination, source)?;
                        }
                        searcher.stream.synchronize()
                    })?;
                }
                graph.prepare()
            },
        )?;
        Ok(PreparedSearch { plan })
    }
}
impl PreparedSearch {
    /// Submit resident contiguous f32 rows after their producer, without host waits.
    pub fn submit(&self, queries: &DeviceTensor) -> Result<SearchInference> {
        Ok(SearchInference {
            inference: self.plan.submit(std::slice::from_ref(queries))?,
        })
    }
    /// Whether every worker and exported output is idle, for budgeted eviction.
    pub fn is_idle(&self) -> bool {
        self.plan.is_idle()
    }
}
/// Resident search results; only each row's `counts` entries are valid.
pub struct SearchInference {
    inference: Inference,
}
impl SearchInference {
    /// Row status: zero is success; one means a non-finite or zero-norm query.
    pub fn status(&self) -> &DeviceTensor {
        &self.inference.outputs()[0]
    }
    /// Valid neighbor counts, one u32 per query.
    pub fn counts(&self) -> &DeviceTensor {
        &self.inference.outputs()[1]
    }
    /// Resident row-major f32 cosine scores.
    pub fn scores(&self) -> &DeviceTensor {
        &self.inference.outputs()[2]
    }
    /// Resident row-major u32 insertion IDs.
    pub fn ids(&self) -> &DeviceTensor {
        &self.inference.outputs()[3]
    }
    /// Read results at an explicit host boundary. Invalid rows are errors.
    pub fn read(self) -> Result<Vec<Vec<Neighbor>>> {
        let batch = self.status().desc().elements();
        let k = self.ids().desc().shape()[1];
        let values = self.inference.download()?.wait()?;
        let word = |output: usize, i: usize| -> [u8; 4] {
            values[output][i * 4..i * 4 + 4].try_into().unwrap()
        };
        let mut rows = Vec::with_capacity(batch);
        for row in 0..batch {
            if u32::from_le_bytes(word(0, row)) != 0 {
                return Err(invalid("device query is nonfinite or has zero norm"));
            }
            let count = u32::from_le_bytes(word(1, row)) as usize;
            if count > k {
                return Err(invalid("invalid device result count"));
            }
            rows.push(
                (0..count)
                    .map(|column| Neighbor {
                        id: u32::from_le_bytes(word(3, row * k + column)),
                        similarity: f32::from_le_bytes(word(2, row * k + column)),
                    })
                    .collect(),
            );
        }
        Ok(rows)
    }
}