Redesign the multi-vector kernels to decouple layout, micro-kernel and reduction

diskann-quantization/src/multi_vector/distance/factory.rs

@@ -1,5 +1,7 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT license.
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
 
 //! Factory + concrete `MaxSimKernel<T>` impls for the multi-vector distance
 //! API. BYOTE entry point — see [`build_max_sim`].
@@ -16,8 +18,7 @@ use diskann_wide::arch::x86_64::{V3, V4};
 
 use super::isa::{MaxSimIsa, NotSupported};
 use super::kernel::{Erase, MaxSimKernel};
-use super::kernels::f16::F16Entry;
-use super::kernels::f32::F32Kernel;
+use super::kernels::{MaxIp, MaxIpF16};
 use super::max_sim::{MaxSim, MaxSimError};
 use crate::multi_vector::distance::QueryMatRef;
 use crate::multi_vector::{BlockTransposed, BlockTransposedRef, Mat, MatRef, Standard};
@@ -35,7 +36,7 @@ struct Prepared<A, Q> {
 impl<A, const GROUP: usize> MaxSimKernel<f32> for Prepared<A, BlockTransposed<f32, GROUP>>
 where
     A: Architecture,
-    F32Kernel<GROUP>: for<'a> diskann_wide::arch::Target3<
+    MaxIp: for<'a> diskann_wide::arch::Target3<
             A,
             (),
             BlockTransposedRef<'a, f32, GROUP>,
@@ -59,14 +60,11 @@ where
             scores.fill(f32::MAX);
             return Ok(());
         }
-        let mut scratch = vec![f32::MIN; self.prepared.padded_nrows()];
-        self.arch.run3(
-            F32Kernel::<GROUP>,
-            self.prepared.reborrow(),
-            doc,
-            &mut scratch,
-        );
-        for (dst, &src) in scores.iter_mut().zip(&scratch[..self.prepared.nrows()]) {
+        // `run` seeds the max itself, so the fill value here is arbitrary.
+        let mut state = vec![0.0; self.prepared.padded_nrows()];
+        self.arch
+            .run3(MaxIp, self.prepared.reborrow(), doc, &mut state);
+        for (dst, &src) in scores.iter_mut().zip(&state[..self.prepared.nrows()]) {
             *dst = -src;
         }
         Ok(())
@@ -77,7 +75,7 @@ impl<A, const GROUP: usize> MaxSimKernel<half::f16>
     for Prepared<A, BlockTransposed<half::f16, GROUP>>
 where
     A: Architecture,
-    F16Entry<GROUP>: for<'a> diskann_wide::arch::Target3<
+    MaxIpF16: for<'a> diskann_wide::arch::Target3<
             A,
             (),
             BlockTransposedRef<'a, half::f16, GROUP>,
@@ -101,14 +99,11 @@ where
             scores.fill(f32::MAX);
             return Ok(());
         }
-        let mut scratch = vec![f32::MIN; self.prepared.padded_nrows()];
-        self.arch.run3(
-            F16Entry::<GROUP>,
-            self.prepared.reborrow(),
-            doc,
-            &mut scratch,
-        );
-        for (dst, &src) in scores.iter_mut().zip(&scratch[..self.prepared.nrows()]) {
+        // `run` seeds the max itself, so the fill value here is arbitrary.
+        let mut state = vec![0.0; self.prepared.padded_nrows()];
+        self.arch
+            .run3(MaxIpF16, self.prepared.reborrow(), doc, &mut state);
+        for (dst, &src) in scores.iter_mut().zip(&state[..self.prepared.nrows()]) {
             *dst = -src;
         }
         Ok(())

diskann-quantization/src/multi_vector/distance/fallback.rs

@@ -1,5 +1,7 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT license.
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
 
 //! Fallback kernel implementation of multi-vector distance computation.
 

diskann-quantization/src/multi_vector/distance/isa.rs

@@ -1,5 +1,7 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT license.
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
 
 //! Instruction Set Architecture (ISA) selector for the multi-vector MaxSim
 //! factory.

diskann-quantization/src/multi_vector/distance/kernel.rs

@@ -1,5 +1,7 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT license.
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
 
 //! Object-safe kernel boundary trait plus BYOTE visitor trait.
 

diskann-quantization/src/multi_vector/distance/kernels/f16.rs

@@ -1,52 +1,214 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
-// Licensed under the MIT license.
+/*
+ * Copyright (c) Microsoft Corporation.
+ * Licensed under the MIT license.
+ */
 
-//! f16 dispatch adapter for block-transposed multi-vector distance.
+//! f16 MaxSim.
 //!
-//! Reuses the f32 micro-kernel family with tile-level f16→f32 conversion
-//! via [`ConvertTo`](super::layouts::ConvertTo). No f16-specific micro-kernel
-//! code is needed — the [`F32Kernel`](super::f32::F32Kernel) does all the
-//! SIMD work after conversion.
-//!
-//! Conversion from f16 to f32 is performed at tile granularity via
-//! [`SliceCast`](diskann_vector::conversion::SliceCast), dispatched through
-//! the runtime architecture token — the same SIMD level used by the
-//! micro-kernel.
+//! 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 A side never has to be staged at once, and the staged copy stays
+//! inside the cache level its tile was sized for.
+
+use core::num::NonZeroUsize;
 
+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::Kernel;
-use super::TileBudget;
-use super::f32::{F32Kernel, max_ip_kernel};
-use super::layouts;
+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::{BlockTransposedTile, Cursor, RowMajorTile, contraction, tile_stride};
+use super::{Plan, TileAt, TileBudget, TileWalk, float};
 use crate::multi_vector::{BlockTransposedRef, MatRef, Standard};
 
-pub(crate) struct F16Entry<const GROUP: usize>;
+/// Stages one source tile at a time as f32.
+struct Widen<'a, Arch> {
+    arch: Arch,
+    cursor: Cursor<'a, half::f16>,
+    buf: Vec<f32>,
+    k: NonZeroUsize,
+}
+
+impl<'a, Arch: Architecture> Widen<'a, Arch>
+where
+    SliceCast<f32, half::f16>: for<'x> Target2<Arch, (), &'x mut [f32], &'x [half::f16]>,
+{
+    fn new(arch: Arch, src: &'a [half::f16], k: NonZeroUsize, stride: NonZeroUsize) -> Self {
+        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 an f16 [`BlockTransposedRef`].
+///
+/// Widening is element-wise, so it preserves the block-transposed permutation.
+struct BlockTransposedWiden<'a, Arch, const AR: usize>(Widen<'a, Arch>);
+
+impl<'a, Arch: Architecture, const AR: usize> BlockTransposedWiden<'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: NonZeroUsize,
+    ) -> Self {
+        let k = contraction(view.padded_ncols());
+        Self(Widen::new(
+            arch,
+            view.as_slice(),
+            k,
+            tile_stride(a_panels, AR, k),
+        ))
+    }
+}
+
+impl<'t, Arch, const AR: usize> TileAt<'t> for BlockTransposedWiden<'_, Arch, AR> {
+    type Tile = BlockTransposedTile<'t, f32, AR>;
+}
+
+impl<Arch: Architecture, const AR: usize> TileWalk for BlockTransposedWiden<'_, Arch, AR>
+where
+    SliceCast<f32, half::f16>: for<'x> Target2<Arch, (), &'x mut [f32], &'x [half::f16]>,
+{
+    fn next(&mut self) -> Option<BlockTransposedTile<'_, f32, AR>> {
+        let k = self.0.k;
+        self.0.next().map(|data| BlockTransposedTile::new(data, k))
+    }
+
+    fn reset(&mut self) {
+        self.0.cursor.reset();
+    }
+}
 
-impl<A, const GROUP: usize>
+/// Widens an f16 [`Standard`] matrix.
+struct RowMajorWiden<'a, Arch, const BR: usize>(Widen<'a, Arch>);
+
+impl<'a, Arch: Architecture, const BR: usize> RowMajorWiden<'a, Arch, BR>
+where
+    SliceCast<f32, half::f16>: for<'x> Target2<Arch, (), &'x mut [f32], &'x [half::f16]>,
+{
+    fn new(arch: Arch, mat: MatRef<'a, Standard<half::f16>>, b_panels: NonZeroUsize) -> Self {
+        let k = contraction(mat.vector_dim());
+        Self(Widen::new(
+            arch,
+            mat.as_slice(),
+            k,
+            tile_stride(b_panels, BR, k),
+        ))
+    }
+}
+
+impl<'t, Arch, const BR: usize> TileAt<'t> for RowMajorWiden<'_, Arch, BR> {
+    type Tile = RowMajorTile<'t, f32, BR>;
+}
+
+impl<Arch: Architecture, const BR: usize> TileWalk for RowMajorWiden<'_, Arch, BR>
+where
+    SliceCast<f32, half::f16>: for<'x> Target2<Arch, (), &'x mut [f32], &'x [half::f16]>,
+{
+    fn next(&mut self) -> Option<RowMajorTile<'_, f32, BR>> {
+        let k = self.0.k;
+        self.0.next().map(|data| RowMajorTile::new(data, k))
+    }
+
+    fn reset(&mut self) {
+        self.0.cursor.reset();
+    }
+}
+
+///////////
+// Entry //
+///////////
+
+/// The f16 MaxSim entry: the f32 pipeline behind widening walks.
+///
+/// Operand naming matches [`MaxIp`](super::MaxIp). The block-transposed A side is the
+/// query and the row-major B side the documents.
+#[derive(Debug, Clone, Copy)]
+pub(crate) struct MaxIpF16;
+
+#[cfg(target_arch = "x86_64")]
+impl
     diskann_wide::arch::Target3<
-        A,
+        V3,
         (),
-        BlockTransposedRef<'_, half::f16, GROUP>,
+        BlockTransposedRef<'_, half::f16, V3_A>,
         MatRef<'_, Standard<half::f16>>,
         &mut [f32],
-    > for F16Entry<GROUP>
-where
-    A: Architecture,
-    F32Kernel<GROUP>: Kernel<A>,
-    layouts::BlockTransposed<half::f16, GROUP>: layouts::ConvertTo<A, <F32Kernel<GROUP> as Kernel<A>>::Left>
-        + layouts::Layout<Element = half::f16>,
-    layouts::RowMajor<half::f16>: layouts::ConvertTo<A, <F32Kernel<GROUP> as Kernel<A>>::Right>
-        + layouts::Layout<Element = half::f16>,
+    > 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,
+            docs.num_vectors(),
+            query.padded_ncols(),
+            TileBudget::default(),
+            state,
+            |plan: Plan<V3_A, V3_B>| {
+                (
+                    BlockTransposedWiden::new(arch, query, plan.a_panels),
+                    RowMajorWiden::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: A,
-        lhs: BlockTransposedRef<'_, half::f16, GROUP>,
-        rhs: MatRef<'_, Standard<half::f16>>,
-        scratch: &mut [f32],
+        arch: Scalar,
+        query: BlockTransposedRef<'_, half::f16, SC_A>,
+        docs: MatRef<'_, Standard<half::f16>>,
+        state: &mut [f32],
     ) {
-        max_ip_kernel(arch, lhs, rhs, scratch, TileBudget::default());
+        float::run(
+            arch,
+            docs.num_vectors(),
+            query.padded_ncols(),
+            TileBudget::default(),
+            state,
+            |plan: Plan<SC_A, SC_B>| {
+                (
+                    BlockTransposedWiden::new(arch, query, plan.a_panels),
+                    RowMajorWiden::new(arch, docs, plan.b_panels),
+                )
+            },
+        );
     }
 }