Redesign the multi-vector kernels to decouple layout, micro-kernel and reduction #1333
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,52 +1,197 @@ | ||
| // Copyright (c) Microsoft Corporation. All rights reserved. | ||
| // Licensed under the MIT license. | ||
|
|
||
| //! f16 MaxSim. | ||
| //! | ||
| //! There is no f16 leaf: f16 widens to f32 and reuses the f32 pipeline. Both sides widen a | ||
| //! tile at a time into a buffer the walk reuses, which is what the lending [`TileWalk`] | ||
| //! exists for — the whole query never has to be staged at once, and the staged copy stays | ||
| //! inside the cache level its tile was sized for. | ||
|
|
||
| use diskann_vector::conversion::SliceCast; | ||
| use diskann_wide::Architecture; | ||
| #[cfg(target_arch = "x86_64")] | ||
| use diskann_wide::arch::x86_64::V3; | ||
| use diskann_wide::arch::{Scalar, Target2}; | ||
|
|
||
| use super::leaves::scalar::{A_PANEL as SC_A, B_PANEL as SC_B}; | ||
| #[cfg(target_arch = "x86_64")] | ||
| use super::leaves::v3::{A_PANEL as V3_A, B_PANEL as V3_B}; | ||
| use super::tiles::{Cursor, DocTile, QueryTile}; | ||
| use super::{Plan, TileAt, TileBudget, TileWalk, float}; | ||
| use crate::multi_vector::{BlockTransposedRef, MatRef, Standard}; | ||
|
|
||
| /// Stages one source tile at a time as f32. | ||
| struct Widen<'a, Arch> { | ||
| arch: Arch, | ||
| cursor: Cursor<'a, half::f16>, | ||
| buf: Vec<f32>, | ||
| k: usize, | ||
| } | ||
|
|
||
| impl<'a, Arch: Architecture> Widen<'a, Arch> | ||
| where | ||
| SliceCast<f32, half::f16>: for<'x> Target2<Arch, (), &'x mut [f32], &'x [half::f16]>, | ||
| { | ||
| /// # Panics | ||
| /// | ||
| /// Panics if `k` is zero — the entry guards that case before any walk is built. | ||
| fn new(arch: Arch, src: &'a [half::f16], k: usize, stride: usize) -> Self { | ||
| assert!(k > 0, "widening walk requires a non-empty contraction"); | ||
| let cursor = Cursor::new(src, stride); | ||
| let buf = vec![0.0f32; cursor.widest()]; | ||
| Self { | ||
| arch, | ||
| cursor, | ||
| buf, | ||
| k, | ||
| } | ||
| } | ||
|
|
||
| fn next(&mut self) -> Option<&[f32]> { | ||
| let arch = self.arch; | ||
| let src = self.cursor.next()?; | ||
| let len = src.len(); | ||
| arch.run2(SliceCast::new(), &mut self.buf[..len], src); | ||
| Some(&self.buf[..len]) | ||
| } | ||
| } | ||
|
|
||
| /// Widens the padded storage of a block-transposed f16 query. | ||
| /// | ||
| /// Widening is element-wise, so it preserves the block-transposed permutation. | ||
| struct QueryWiden<'a, Arch, const AR: usize>(Widen<'a, Arch>); | ||
|
|
||
| impl<'a, Arch: Architecture, const AR: usize> QueryWiden<'a, Arch, AR> | ||
| where | ||
| SliceCast<f32, half::f16>: for<'x> Target2<Arch, (), &'x mut [f32], &'x [half::f16]>, | ||
| { | ||
| fn new(arch: Arch, view: BlockTransposedRef<'a, half::f16, AR>, a_panels: usize) -> Self { | ||
| let k = view.padded_ncols(); | ||
| Self(Widen::new(arch, view.as_slice(), k, a_panels * AR * k)) | ||
| } | ||
| } | ||
|
|
||
| impl<'t, Arch, const AR: usize> TileAt<'t> for QueryWiden<'_, Arch, AR> { | ||
| type Tile = QueryTile<'t, f32, AR>; | ||
| } | ||
|
|
||
| impl<Arch: Architecture, const AR: usize> TileWalk for QueryWiden<'_, Arch, AR> | ||
| where | ||
| SliceCast<f32, half::f16>: for<'x> Target2<Arch, (), &'x mut [f32], &'x [half::f16]>, | ||
| { | ||
| fn next(&mut self) -> Option<QueryTile<'_, f32, AR>> { | ||
| let k = self.0.k; | ||
| self.0.next().map(|data| QueryTile::new(data, k)) | ||
| } | ||
|
|
||
| fn reset(&mut self) { | ||
| self.0.cursor.reset(); | ||
| } | ||
| } | ||
|
|
||
| /// Widens a row-major f16 doc matrix. | ||
| struct DocWiden<'a, Arch, const BR: usize>(Widen<'a, Arch>); | ||
|
|
||
| impl<'a, Arch: Architecture, const BR: usize> DocWiden<'a, Arch, BR> | ||
| where | ||
| SliceCast<f32, half::f16>: for<'x> Target2<Arch, (), &'x mut [f32], &'x [half::f16]>, | ||
| { | ||
| fn new(arch: Arch, docs: MatRef<'a, Standard<half::f16>>, b_panels: usize) -> Self { | ||
| let k = docs.vector_dim(); | ||
| Self(Widen::new(arch, docs.as_slice(), k, b_panels * BR * k)) | ||
| } | ||
| } | ||
|
|
||
| impl<'t, Arch, const BR: usize> TileAt<'t> for DocWiden<'_, Arch, BR> { | ||
| type Tile = DocTile<'t, f32, BR>; | ||
| } | ||
|
|
||
| impl<Arch: Architecture, const BR: usize> TileWalk for DocWiden<'_, Arch, BR> | ||
| where | ||
| SliceCast<f32, half::f16>: for<'x> Target2<Arch, (), &'x mut [f32], &'x [half::f16]>, | ||
| { | ||
| fn next(&mut self) -> Option<DocTile<'_, f32, BR>> { | ||
| let k = self.0.k; | ||
| self.0.next().map(|data| DocTile::new(data, k)) | ||
| } | ||
|
|
||
| fn reset(&mut self) { | ||
|
Contributor
There was a problem hiding this comment. The refactor removes the old asymmetric multi-tile f16 coverage, while the replacement tiny-budget tests only exercise the f32 walks. This leaves the new reusable widening buffer and |
||
| self.0.cursor.reset(); | ||
| } | ||
| } | ||
|
|
||
| // ── Entry ──────────────────────────────────────────────────────── | ||
|
|
||
| /// The f16 MaxSim entry — the f32 pipeline behind widening walks. | ||
| #[derive(Debug, Clone, Copy)] | ||
| pub(crate) struct MaxIpF16; | ||
|
|
||
| #[cfg(target_arch = "x86_64")] | ||
| impl | ||
| diskann_wide::arch::Target3< | ||
| V3, | ||
| (), | ||
| BlockTransposedRef<'_, half::f16, V3_A>, | ||
| MatRef<'_, Standard<half::f16>>, | ||
| &mut [f32], | ||
| > for MaxIpF16 | ||
| { | ||
| #[inline(always)] | ||
| fn run( | ||
| self, | ||
| arch: V3, | ||
| query: BlockTransposedRef<'_, half::f16, V3_A>, | ||
| docs: MatRef<'_, Standard<half::f16>>, | ||
| state: &mut [f32], | ||
| ) { | ||
| float::run( | ||
| arch, | ||
| query.padded_nrows(), | ||
| docs.num_vectors(), | ||
| query.padded_ncols(), | ||
| TileBudget::default(), | ||
| state, | ||
| |plan: Plan<V3_A, V3_B>| { | ||
| ( | ||
| QueryWiden::new(arch, query, plan.a_panels), | ||
| DocWiden::new(arch, docs, plan.b_panels), | ||
| ) | ||
| }, | ||
| ); | ||
| } | ||
| } | ||
|
|
||
| impl | ||
| diskann_wide::arch::Target3< | ||
| Scalar, | ||
| (), | ||
| BlockTransposedRef<'_, half::f16, SC_A>, | ||
| MatRef<'_, Standard<half::f16>>, | ||
| &mut [f32], | ||
| > for MaxIpF16 | ||
| { | ||
| #[inline(always)] | ||
| fn run( | ||
| self, | ||
| arch: Scalar, | ||
| query: BlockTransposedRef<'_, half::f16, SC_A>, | ||
| docs: MatRef<'_, Standard<half::f16>>, | ||
| state: &mut [f32], | ||
| ) { | ||
| float::run( | ||
| arch, | ||
| query.padded_nrows(), | ||
| docs.num_vectors(), | ||
| query.padded_ncols(), | ||
| TileBudget::default(), | ||
| state, | ||
| |plan: Plan<SC_A, SC_B>| { | ||
| ( | ||
| QueryWiden::new(arch, query, plan.a_panels), | ||
| DocWiden::new(arch, docs, plan.b_panels), | ||
| ) | ||
| }, | ||
| ); | ||
| } | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This call no longer performs the old
query.ncols() == doc.vector_dim()boundary check. For non-empty contractions, a mismatch reaches the opaque"B panel extent"assertion; for zero-dimensional queries or empty documents, the early returns accept the mismatch and return scores. Please validate dimensions before either early return for both f32 and f16, and report both values.