From 365a5232873e72fd0fbebeb734d41711bb8919d9 Mon Sep 17 00:00:00 2001 From: Victoriya Fedotova Date: Thu, 15 Oct 2026 06:18:40 -0700 Subject: [PATCH] Initial attempt to split fp16->fp32 argument conversion and the distance compuation in heuristic prune Signed-off-by: Victoriya Fedotova --- include/svs/core/distance.h | 1 + include/svs/core/distance/conversion.h | 242 +++++++++++++++++++++ include/svs/core/distance/cosine.h | 9 + include/svs/core/distance/distance_core.h | 25 +++ include/svs/core/distance/inner_product.h | 16 ++ include/svs/index/vamana/prune.h | 114 +++++++++- include/svs/index/vamana/vamana_build.h | 28 ++- include/svs/multi-arch/x86/avx2.cpp | 3 + include/svs/multi-arch/x86/avx512.cpp | 3 + include/svs/multi-arch/x86/preprocessor.h | 9 + tests/CMakeLists.txt | 1 + tests/svs/core/distances/conversion.cpp | 234 ++++++++++++++++++++ tests/svs/core/distances/cosine.cpp | 9 + tests/svs/core/distances/inner_product.cpp | 8 + tests/svs/index/vamana/prune.cpp | 234 ++++++++++++++++++++ tests/utils/preconversion.h | 164 ++++++++++++++ tests/utils/utils.h | 35 +++ 17 files changed, 1130 insertions(+), 5 deletions(-) create mode 100644 include/svs/core/distance/conversion.h create mode 100644 tests/svs/core/distances/conversion.cpp create mode 100644 tests/utils/preconversion.h diff --git a/include/svs/core/distance.h b/include/svs/core/distance.h index 121827643..46736c8ec 100644 --- a/include/svs/core/distance.h +++ b/include/svs/core/distance.h @@ -16,6 +16,7 @@ #pragma once +#include "svs/core/distance/conversion.h" #include "svs/core/distance/cosine.h" #include "svs/core/distance/euclidean.h" #include "svs/core/distance/inner_product.h" diff --git a/include/svs/core/distance/conversion.h b/include/svs/core/distance/conversion.h new file mode 100644 index 000000000..1e1db4167 --- /dev/null +++ b/include/svs/core/distance/conversion.h @@ -0,0 +1,242 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// svs +#include "svs/core/distance/distance_core.h" +#include "svs/core/distance/simd_utils.h" +#include "svs/lib/avx_detection.h" +#include "svs/lib/float16.h" +#include "svs/lib/preprocessor.h" +#include "svs/lib/static.h" +#include "svs/lib/threads/threadlocal.h" + +// stl +#include +#include +#include +#include +#include +#include +#include + +namespace svs::distance { + +/// Alignment, in bytes, that ``Convert`` requires of its destination. +/// It equals one Intel(R) AVX-512 register, so every full-width store is aligned. +inline constexpr size_t conversion_alignment = 64; + +// `PreconvertedArgument` obtains its alignment from `CacheAlignedAllocator`. If the two +// diverge, its buffers stop satisfying the precondition of `Convert`. +static_assert(conversion_alignment == threads::CACHE_LINE_BYTES); + +// Forward declare implementation to allow entry point to be near the top. +template struct ConvertImpl; + +/// +/// @brief Convert ``length`` elements of ``src`` into ``dst``. +/// +/// Separates the element conversion that distance kernels otherwise repeat on every call +/// from the distance computation itself: a vector converted once can be passed to +/// ``distance::compute`` any number of times without being converted again. +/// +class Convert { + public: + /// + /// @param src Source elements. Need not be aligned. + /// @param dst Destination for ``length`` elements. Must be aligned to + /// ``conversion_alignment`` bytes and must not overlap ``src``. + /// @param length The number of elements to convert. + /// + template + static void apply(const From* src, To* dst, size_t length) { + assert(reinterpret_cast(dst) % conversion_alignment == 0); + if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx512f_supported(), 1)) { + return ConvertImpl::apply( + src, dst, lib::MaybeStatic(length) + ); + } + if (__builtin_expect(svs::detail::avx_runtime_flags.is_avx2_supported(), 1)) { + return ConvertImpl::apply( + src, dst, lib::MaybeStatic(length) + ); + } + return ConvertImpl::apply( + src, dst, lib::MaybeStatic(length) + ); + } +}; + +///// +///// Generic Implementation +///// + +template struct ConvertImpl { + SVS_NOINLINE static void apply(const From* src, To* dst, lib::MaybeStatic length) { + std::copy(src, src + length.size(), dst); + } +}; + +///// +///// Intel(R) AVX-512 Implementations +///// + +SVS_VALIDATE_BOOL_ENV(SVS_AVX512_F) +#if SVS_AVX512_F + +template struct ConvertImpl { + SVS_NOINLINE static void + apply(const Float16* src, float* dst, lib::MaybeStatic length) { + using op = simd::ConvertToFloat<16>; + size_t i = 0; + for (; i + op::simd_width <= length.size(); i += op::simd_width) { + _mm512_store_ps(dst + i, op::load(src + i)); + } + if (i < length.size()) { + auto mask = create_mask(length); + _mm512_mask_store_ps(dst + i, mask, op::load(mask, src + i)); + } + } +}; + +#endif + +///// +///// Intel(R) AVX2 Implementations +///// + +SVS_VALIDATE_BOOL_ENV(SVS_AVX512_F) +SVS_VALIDATE_BOOL_ENV(SVS_AVX2) +#if !SVS_AVX512_F && SVS_AVX2 + +template struct ConvertImpl { + SVS_NOINLINE static void + apply(const Float16* src, float* dst, lib::MaybeStatic length) { + using op = simd::ConvertToFloat<8>; + size_t i = 0; + for (; i + op::simd_width <= length.size(); i += op::simd_width) { + _mm256_store_ps(dst + i, op::load(src + i)); + } + if (i < length.size()) { + // A full-width store here would write past the last converted element. + auto mask = create_mask(length); + _mm256_maskstore_ps( + dst + i, + _mm256_castps_si256(simd::create_blend_mask_avx2(mask)), + op::load(mask, src + i) + ); + } + } +}; + +#endif + +/// +/// @brief Pre-allocated storage for one distance argument in the working type of ``Dist``. +/// +/// @tparam Dist The distance functor the converted argument is passed to. +/// @tparam Argument The argument type as stored. ``distance_argument_t`` +/// names the type it is converted to. +/// +/// Construct once with the dimensionality of the vectors it will hold, then call +/// ``operator()`` for each vector. Converting a vector once lets every later +/// ``distance::compute`` call that uses it, in either argument position, skip the +/// conversion. +/// +/// The span returned by ``operator()`` stays valid until the next ``operator()`` call on +/// the same object. When a conversion takes place, its ``data()`` is aligned to +/// ``conversion_alignment`` bytes. An object is not safe to share between threads: use one +/// per thread. Copies own separate storage. +/// +template class PreconvertedArgument; + +/// +/// @brief No conversion: holds no storage and forwards each argument unchanged. +/// +template + requires std::same_as, Argument> +class PreconvertedArgument { + public: + /// + /// @brief Construct the storage. + /// + /// @param dimensions Ignored: nothing is converted, so nothing is stored. + /// + explicit PreconvertedArgument(size_t SVS_UNUSED(dimensions)) {} + + /// + /// @brief Return ``argument`` unchanged. + /// + /// A prvalue argument is returned by value and an lvalue by reference, so binding the + /// result to ``const auto&`` behaves exactly as binding ``argument`` itself would. + /// + template A operator()(A&& argument) const { + return std::forward(argument); + } +}; + +/// +/// @brief Converts each arithmetic span into one buffer, allocated at construction. +/// +template + requires(!std::same_as< + distance_argument_t>, + std::span>) +class PreconvertedArgument> { + public: + /// The type ``operator()`` returns. + using converted_type = distance_argument_t>; + + private: + using value_type = std::remove_const_t; + + public: + /// + /// @brief Allocate the storage. + /// + /// @param dimensions The number of elements of every vector later passed to + /// ``operator()``. + /// + explicit PreconvertedArgument(size_t dimensions) + : buffer_(dimensions) {} + + /// + /// @brief Convert ``argument`` into the storage and return a view of the result. + /// + /// @param argument The vector to convert. Its size must equal the ``dimensions`` + /// given at construction. + /// + converted_type operator()(std::span argument) { + assert(argument.size() == buffer_.size()); + Convert::apply(argument.data(), buffer_.data(), argument.size()); + return converted_type(buffer_.data(), buffer_.size()); + } + + private: + std::vector> buffer_; +}; + +#if defined(__x86_64__) + +#include "svs/multi-arch/x86/preprocessor.h" + +DISTANCE_CONVERT_EXTERN_TEMPLATE(AVX_AVAILABILITY::AVX512); +DISTANCE_CONVERT_EXTERN_TEMPLATE(AVX_AVAILABILITY::AVX2); + +#endif + +} // namespace svs::distance diff --git a/include/svs/core/distance/cosine.h b/include/svs/core/distance/cosine.h index 9f4924997..858d37c8d 100644 --- a/include/svs/core/distance/cosine.h +++ b/include/svs/core/distance/cosine.h @@ -178,6 +178,15 @@ float compute(DistanceCosineSimilarity distance, std::span a, std::span< } } +// The cosine kernels convert `Float16` operands to `float` on every SIMD step. An operand +// converted once ahead of time selects the mixed `float` kernel, which loads it directly. +// Removing this sends the iterative prune back to reconverting its query for every +// candidate. +template +struct distance_argument> { + using type = std::span; +}; + ///// ///// Generic Implementation ///// diff --git a/include/svs/core/distance/distance_core.h b/include/svs/core/distance/distance_core.h index 4f59f9de0..6d74f3534 100644 --- a/include/svs/core/distance/distance_core.h +++ b/include/svs/core/distance/distance_core.h @@ -42,6 +42,31 @@ template constexpr bool is_dim_supported() { using default_accum_type = float; +/// +/// @brief Type that an argument of type ``Argument`` should be converted to before it is +/// compared against many other vectors by the distance functor ``Distance``. +/// +/// @tparam Distance The distance functor the converted argument is passed to. +/// @tparam Argument The argument type as stored, for example the +/// ``std::span`` a dataset accessor returns. +/// +/// The mapping applies equally to the left- and right-hand arguments of +/// ``distance::compute``. The primary template maps ``Argument`` to itself, meaning +/// converting ahead of time brings no benefit: ``Argument`` is already in the kernel's +/// working type, the kernel for it would be slower or numerically different with a +/// converted operand, or ``Argument`` is not a plain arithmetic span (for example a +/// compressed vector). +/// +/// A specialization belongs next to the ``distance::compute`` overload of the functor +/// whose kernels it describes. +/// +template struct distance_argument { + using type = Argument; +}; + +template +using distance_argument_t = typename distance_argument::type; + template Accum norm_square(lib::Type /*unused*/, std::span data) { Accum accum{0}; diff --git a/include/svs/core/distance/inner_product.h b/include/svs/core/distance/inner_product.h index 0f7837a53..b36cda0bf 100644 --- a/include/svs/core/distance/inner_product.h +++ b/include/svs/core/distance/inner_product.h @@ -160,6 +160,15 @@ float compute(DistanceIP /*unused*/, std::span a, std::span b) { } } +// The IP kernels convert `Float16` operands to `float` on every SIMD step. An operand +// converted once ahead of time selects the mixed `float` kernel, which loads it directly. +// Removing this sends the iterative prune back to reconverting its query for every +// candidate. +template +struct distance_argument> { + using type = std::span; +}; + ///// ///// Generic Implementation ///// @@ -362,6 +371,13 @@ template struct IPImpl { } }; +template struct IPImpl { + SVS_NOINLINE static float + compute(const Float16* a, const float* b, lib::MaybeStatic length) { + return svs::simd::generic_simd_op(IPFloatOp<8>{}, a, b, length); + } +}; + template struct IPImpl { SVS_NOINLINE static float compute(const float* a, const int8_t* b, lib::MaybeStatic length) { diff --git a/include/svs/index/vamana/prune.h b/include/svs/index/vamana/prune.h index aeab27ac7..be6eb3847 100644 --- a/include/svs/index/vamana/prune.h +++ b/include/svs/index/vamana/prune.h @@ -19,6 +19,7 @@ #include "svs/concepts/data.h" #include "svs/concepts/distance.h" #include "svs/core/distance.h" +#include "svs/core/distance/conversion.h" #include "svs/lib/neighbor.h" #include "svs/lib/type_traits.h" @@ -88,6 +89,21 @@ inline PruneState reenable(PruneState state) { inline bool excluded(PruneState state) { return state != PruneState::Available; } +/// +/// @brief Pre-allocated storage for the query of the ``IterativePruneStrategy`` overload of +/// ``heuristic_prune_neighbors``. +/// +/// @tparam Data The full type of the dataset passed to the prune. +/// @tparam Accessor The accessor passed to the prune. +/// @tparam Dist The distance functor passed to the prune. +/// +/// Construct with ``dataset.dimensions()`` and reuse across prune calls on one thread. +/// +template +using prune_query_storage_t = distance::PreconvertedArgument< + std::remove_cvref_t, + std::remove_cvref_t>>; + /// /// @brief Function to prune neighbors using MRNG rule (extended with alpha as in Vamana). /// @@ -97,6 +113,10 @@ inline bool excluded(PruneState state) { return state != PruneState::Available; /// @tparam I The type of the resulting index for each neighbor. /// @tparam Alloc Allocator for the result vector. /// +/// @param query_storage Storage the query of each pruning step is converted into, so the +/// comparisons against the remaining candidates do not reconvert it. Must have been +/// constructed with ``dataset.dimensions()``. Its contents on return are unspecified. +/// template < data::ImmutableMemoryDataset Data, data::AccessorFor Accessor, @@ -113,7 +133,8 @@ void heuristic_prune_neighbors( Dist& distance_function, size_t current_node_id, const std::span& pool, - std::vector& result + std::vector& result, + prune_query_storage_t& query_storage ) { auto cmp = distance::comparator(distance_function); assert(std::is_sorted(pool.begin(), pool.end(), cmp)); @@ -142,7 +163,15 @@ void heuristic_prune_neighbors( // Only once we know this item needs to be processed to we retrieve // the corresponding data and perform preprocessing. - const auto& query = accessor(dataset, id); + const auto& query = query_storage(accessor(dataset, id)); + static_assert( + distance::Distance< + Dist, + std::remove_cvref_t, + std::remove_cvref_t>, + "The distance functor has no `compute` overload for the converted query " + "type named by `distance::distance_argument`." + ); distance::maybe_fix_argument(distance_function, query); result.push_back(detail::construct_as(lib::Type(), pool[start])); for (size_t t = start + 1; t < poolsize; ++t) { @@ -173,6 +202,45 @@ void heuristic_prune_neighbors( } } +/// +/// @brief Function to prune neighbors using MRNG rule (extended with alpha as in Vamana). +/// +/// Allocates the query storage for this call only. Callers that prune repeatedly should +/// hold a ``prune_query_storage_t`` and pass it to the overload that takes one. +/// +template < + data::ImmutableMemoryDataset Data, + data::AccessorFor Accessor, + distance::Distance, data::const_value_type_t> Dist, + NeighborLike Neighbors, + detail::IntegerOrNeighbor I, + typename Alloc> +void heuristic_prune_neighbors( + IterativePruneStrategy dispatch, + size_t max_result_size, + float alpha, + const Data& dataset, + const Accessor& accessor, + Dist& distance_function, + size_t current_node_id, + const std::span& pool, + std::vector& result +) { + auto query_storage = prune_query_storage_t(dataset.dimensions()); + heuristic_prune_neighbors( + dispatch, + max_result_size, + alpha, + dataset, + accessor, + distance_function, + current_node_id, + pool, + result, + query_storage + ); +} + template < data::ImmutableMemoryDataset Data, data::AccessorFor Accessor, @@ -306,4 +374,46 @@ void heuristic_prune_neighbors( ++start; } } + +/// +/// @brief Prune with a strategy that does not use pre-allocated query storage. +/// +/// Lets a caller that holds a ``prune_query_storage_t`` pass it whichever strategy its +/// distance functor selects. ``ProgressivePruneStrategy`` and ``LegacyPruneStrategy`` do +/// not convert their query, so ``query_storage`` is ignored. +/// +template < + typename Strategy, + data::ImmutableMemoryDataset Data, + data::AccessorFor Accessor, + distance::Distance, data::const_value_type_t> Dist, + NeighborLike Neighbors, + detail::IntegerOrNeighbor I, + typename Alloc> + requires std::same_as || + std::same_as +void heuristic_prune_neighbors( + Strategy dispatch, + size_t max_result_size, + float alpha, + const Data& dataset, + const Accessor& accessor, + Dist& distance_function, + size_t current_node_id, + const std::span& pool, + std::vector& result, + prune_query_storage_t& SVS_UNUSED(query_storage) +) { + heuristic_prune_neighbors( + dispatch, + max_result_size, + alpha, + dataset, + accessor, + distance_function, + current_node_id, + pool, + result + ); +} } // namespace svs::index::vamana diff --git a/include/svs/index/vamana/vamana_build.h b/include/svs/index/vamana/vamana_build.h index f658a974f..52cd774ed 100644 --- a/include/svs/index/vamana/vamana_build.h +++ b/include/svs/index/vamana/vamana_build.h @@ -47,6 +47,7 @@ #include #include #include +#include #include namespace svs::index::vamana { @@ -178,6 +179,16 @@ class VamanaBuilder { using update_type = threads::SequentialTLS>>>; + // Storage for the query of each prune, typed after the accessor and distance the build + // adaptor hands to `heuristic_prune_neighbors`. + using build_adaptor_type = decltype(extensions::build_adaptor( + std::declval(), std::declval() + )); + using prune_query_storage_type = prune_query_storage_t< + Data, + decltype(std::declval().general_accessor()), + decltype(std::declval().general_distance())>; + /// Constructor VamanaBuilder( Graph& graph, @@ -371,6 +382,10 @@ class VamanaBuilder { auto&& general_distance = build_adaptor.general_distance(); auto general_accessor = build_adaptor.general_accessor(); + // `StaticPartition` runs this body once per thread, so this storage is + // allocated once per thread and reused for every node it prunes. + auto prune_query_storage = prune_query_storage_type(data_.dimensions()); + for (auto node_id : local_indices) { pool.clear(); search_buffer.clear(); @@ -472,7 +487,8 @@ class VamanaBuilder { general_distance, node_id, lib::as_const_span(pool), - pruned_results + pruned_results, + prune_query_storage ); } } @@ -525,10 +541,15 @@ class VamanaBuilder { // // Take care to avoid duplicate entries. auto prune_timer = timer.push_back("pruning backedges"); + // `DynamicPartition` runs the body below once per bucket, so storage declared in it + // would be reallocated for every bucket. Allocate it once per thread instead. + auto prune_query_storage = threads::SequentialTLS( + prune_query_storage_type(data_.dimensions()), threadpool_.size() + ); threads::parallel_for( threadpool_, threads::DynamicPartition{backedge_buffer_.buckets(), 1}, - [&](auto& buckets, uint64_t SVS_UNUSED(tid)) { + [&](auto& buckets, uint64_t tid) { // Thread local auxiliary data structures. std::vector> candidates{}; std::vector pruned_results{}; @@ -582,7 +603,8 @@ class VamanaBuilder { general_distance, src, lib::as_const_span(candidates), - pruned_results + pruned_results, + prune_query_storage.at(tid) ); graph_.replace_node(src, pruned_results); } diff --git a/include/svs/multi-arch/x86/avx2.cpp b/include/svs/multi-arch/x86/avx2.cpp index bff53ae10..f7e446349 100644 --- a/include/svs/multi-arch/x86/avx2.cpp +++ b/include/svs/multi-arch/x86/avx2.cpp @@ -15,6 +15,7 @@ */ #if defined(__x86_64__) +#include "svs/core/distance/conversion.h" #include "svs/core/distance/cosine.h" #include "svs/core/distance/euclidean.h" #include "svs/core/distance/inner_product.h" @@ -52,6 +53,8 @@ DISTANCE_CS_INSTANTIATE_TEMPLATE(512, AVX_AVAILABILITY::AVX2); DISTANCE_CS_INSTANTIATE_TEMPLATE(768, AVX_AVAILABILITY::AVX2); DISTANCE_CS_INSTANTIATE_TEMPLATE(Dynamic, AVX_AVAILABILITY::AVX2); +DISTANCE_CONVERT_INSTANTIATE_TEMPLATE(AVX_AVAILABILITY::AVX2); + } // namespace svs::distance #endif diff --git a/include/svs/multi-arch/x86/avx512.cpp b/include/svs/multi-arch/x86/avx512.cpp index bee150d75..e6c4c5351 100644 --- a/include/svs/multi-arch/x86/avx512.cpp +++ b/include/svs/multi-arch/x86/avx512.cpp @@ -15,6 +15,7 @@ */ #if defined(__x86_64__) +#include "svs/core/distance/conversion.h" #include "svs/core/distance/cosine.h" #include "svs/core/distance/euclidean.h" #include "svs/core/distance/inner_product.h" @@ -52,6 +53,8 @@ DISTANCE_CS_INSTANTIATE_TEMPLATE(512, AVX_AVAILABILITY::AVX512); DISTANCE_CS_INSTANTIATE_TEMPLATE(768, AVX_AVAILABILITY::AVX512); DISTANCE_CS_INSTANTIATE_TEMPLATE(Dynamic, AVX_AVAILABILITY::AVX512); +DISTANCE_CONVERT_INSTANTIATE_TEMPLATE(AVX_AVAILABILITY::AVX512); + } // namespace svs::distance #endif diff --git a/include/svs/multi-arch/x86/preprocessor.h b/include/svs/multi-arch/x86/preprocessor.h index 4e0cb941d..4149aa288 100644 --- a/include/svs/multi-arch/x86/preprocessor.h +++ b/include/svs/multi-arch/x86/preprocessor.h @@ -87,3 +87,12 @@ #define DISTANCE_CS_EXTERN_TEMPLATE(N, AVX) \ DISTANCE_CS_TEMPLATE_HELPER(extern template, N, AVX); + +#define DISTANCE_CONVERT_TEMPLATE_HELPER(SPEC, AVX) \ + SPEC struct ConvertImpl; + +#define DISTANCE_CONVERT_INSTANTIATE_TEMPLATE(AVX) \ + DISTANCE_CONVERT_TEMPLATE_HELPER(template, AVX); + +#define DISTANCE_CONVERT_EXTERN_TEMPLATE(AVX) \ + DISTANCE_CONVERT_TEMPLATE_HELPER(extern template, AVX); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 98a1cb29d..b0e4950b3 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -111,6 +111,7 @@ set(TEST_SOURCES ${TEST_DIR}/svs/core/data/view.cpp ${TEST_DIR}/svs/core/distance.cpp ${TEST_DIR}/svs/core/distances/simd_utils.cpp + ${TEST_DIR}/svs/core/distances/conversion.cpp ${TEST_DIR}/svs/core/distances/distance_euclidean.cpp ${TEST_DIR}/svs/core/distances/inner_product.cpp ${TEST_DIR}/svs/core/distances/cosine.cpp diff --git a/tests/svs/core/distances/conversion.cpp b/tests/svs/core/distances/conversion.cpp new file mode 100644 index 000000000..0fe9ef179 --- /dev/null +++ b/tests/svs/core/distances/conversion.cpp @@ -0,0 +1,234 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// header under test +#include "svs/core/distance/conversion.h" + +// svs +#include "svs/core/distance/cosine.h" +#include "svs/core/distance/euclidean.h" +#include "svs/core/distance/inner_product.h" +#include "svs/lib/float16.h" +#include "svs/lib/threads/threadlocal.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" + +// tests +#include "tests/utils/generators.h" +#include "tests/utils/preconversion.h" +#include "tests/utils/utils.h" + +// stl +#include +#include +#include +#include +#include +#include + +namespace { + +using svs::Float16; +namespace distance = svs::distance; + +using svs_test::aligned_vector; + +// Cover the full-width loops and ragged tails of both the 8-wide (AVX2) and 16-wide +// (AVX-512) converters, the empty input, and dimensionalities into the thousands. +constexpr auto test_lengths = std::to_array( + {0, 1, 7, 8, 9, 15, 16, 17, 31, 33, 100, 128, 768, 1000, 4096, 4099} +); + +// Elements after the converted range that every converter must leave untouched. +constexpr size_t num_sentinels = 32; +constexpr float sentinel = -12345.0F; + +bool is_aligned(const void* ptr) { + return reinterpret_cast(ptr) % distance::conversion_alignment == 0; +} + +std::vector random_float16(size_t length) { + auto values = std::vector{}; + auto generator = svs_test::make_generator(-10, 10); + svs_test::populate(values, generator, length); + svs_test::insert_subnormals(values); + return values; +} + +// IEEE 754 binary16 to binary32, exact for every finite input including subnormals. This +// is what the hardware conversion in the Intel(R) AVX kernels computes. +float exact_float16_to_float(Float16 x) { + const uint16_t raw = x.raw(); + const int exponent = (raw >> 10) & 0x1F; + const int mantissa = raw & 0x3FF; + const float sign = (raw & 0x8000) != 0 ? -1.0F : 1.0F; + if (exponent == 0) { + return sign * std::ldexp(static_cast(mantissa), -24); + } + return sign * std::ldexp(static_cast(mantissa + 0x400), exponent - 25); +} + +// The value `Convert` must produce at the dispatched level: the same conversion the +// distance kernels at that level apply, so that pre-converted and unconverted operands +// agree. The generic kernels use `Float16`'s software conversion, which flushes +// subnormals to zero. +float expected_conversion(std::string_view isa, Float16 x) { + return isa == "generic" ? static_cast(x) : exact_float16_to_float(x); +} + +} // namespace + +CATCH_TEST_CASE("Element Conversion", "[distance][conversion]") { + CATCH_SECTION("Float16 to float") { + svs_test::for_each_supported_isa([](std::string_view isa) { + for (size_t length : test_lengths) { + CATCH_INFO("ISA: " << isa << ", length: " << length); + // One extra leading element lets the source start one element past an + // allocation boundary: `Convert` must not require an aligned source. + auto storage = random_float16(length + 1); + const Float16* src = storage.data() + 1; + + auto dst = aligned_vector(length + num_sentinels, sentinel); + CATCH_REQUIRE(is_aligned(dst.data())); + distance::Convert::apply(src, dst.data(), length); + + for (size_t i = 0; i < length; ++i) { + CATCH_REQUIRE(dst[i] == expected_conversion(isa, src[i])); + } + for (size_t i = length; i < dst.size(); ++i) { + CATCH_REQUIRE(dst[i] == sentinel); + } + } + }); + } +} + +CATCH_TEST_CASE("Distance Argument", "[distance][conversion]") { + // `Float16` is converted for the distances whose kernels convert it on every step, + // keeping the extent. + static_assert(std::is_same_v< + distance::distance_argument_t< + distance::DistanceIP, + std::span>, + std::span>); + static_assert(std::is_same_v< + distance:: + distance_argument_t>, + std::span>); + static_assert(std::is_same_v< + distance::distance_argument_t< + distance::DistanceCosineSimilarity, + std::span>, + std::span>); + + // Not converted: `float` already is the working type, and `int8_t` would lose the + // VNNI kernel. L2 is not covered by this change. + static_assert(std::is_same_v< + distance:: + distance_argument_t>, + std::span>); + static_assert(std::is_same_v< + distance:: + distance_argument_t>, + std::span>); + static_assert(std::is_same_v< + distance:: + distance_argument_t>, + std::span>); +} + +CATCH_TEST_CASE("Preconverted Argument", "[distance][conversion]") { + CATCH_SECTION("Identity") { + auto values = std::vector(100, 1.0F); + auto span = std::span(values); + auto storage = + distance::PreconvertedArgument>( + values.size() + ); + + // An lvalue comes back as a reference to the same object, a prvalue by value. + static_assert(std::is_same_v&>); + static_assert(std::is_same_v< + decltype(storage(std::span(values))), + std::span>); + CATCH_REQUIRE(storage(span).data() == values.data()); + + auto int8_values = std::vector(100, 1); + auto int8_storage = + distance::PreconvertedArgument>( + int8_values.size() + ); + auto int8_span = std::span(int8_values); + CATCH_REQUIRE(int8_storage(int8_span).data() == int8_values.data()); + } + + CATCH_SECTION("Converting, dynamic extent") { + constexpr size_t dims = 100; + using storage_type = + distance::PreconvertedArgument>; + static_assert(std::is_same_v>); + + auto storage = storage_type(dims); + auto first = random_float16(dims); + auto second = random_float16(dims); + auto expected_first = svs_test::convert_like_kernels(first); + auto expected_second = svs_test::convert_like_kernels(second); + + auto converted_first = storage(std::span(first)); + CATCH_REQUIRE(converted_first.size() == dims); + CATCH_REQUIRE(is_aligned(converted_first.data())); + for (size_t i = 0; i < dims; ++i) { + CATCH_REQUIRE(converted_first[i] == expected_first[i]); + } + + // The storage is reused, not reallocated, for the next vector. + auto converted_second = storage(std::span(second)); + CATCH_REQUIRE(converted_second.data() == converted_first.data()); + for (size_t i = 0; i < dims; ++i) { + CATCH_REQUIRE(converted_second[i] == expected_second[i]); + } + + // Copies, as made by `SequentialTLS` for each thread, own separate storage. + auto copy = storage; + auto converted_copy = copy(std::span(first)); + CATCH_REQUIRE(converted_copy.data() != converted_second.data()); + CATCH_REQUIRE(is_aligned(converted_copy.data())); + for (size_t i = 0; i < dims; ++i) { + CATCH_REQUIRE(converted_copy[i] == expected_first[i]); + CATCH_REQUIRE(converted_second[i] == expected_second[i]); + } + } + + CATCH_SECTION("Converting, static extent") { + constexpr size_t dims = 128; + using storage_type = distance::PreconvertedArgument< + distance::DistanceCosineSimilarity, + std::span>; + static_assert(std::is_same_v< + storage_type::converted_type, + std::span>); + + auto storage = storage_type(dims); + auto values = random_float16(dims); + auto expected = svs_test::convert_like_kernels(values); + auto converted = storage(std::span(values.data(), dims)); + CATCH_REQUIRE(is_aligned(converted.data())); + for (size_t i = 0; i < dims; ++i) { + CATCH_REQUIRE(converted[i] == expected[i]); + } + } +} diff --git a/tests/svs/core/distances/cosine.cpp b/tests/svs/core/distances/cosine.cpp index 8c68ebc1a..01758edd5 100644 --- a/tests/svs/core/distances/cosine.cpp +++ b/tests/svs/core/distances/cosine.cpp @@ -38,6 +38,7 @@ // Testing Utilities #include "tests/utils/benchmarks.h" #include "tests/utils/generators.h" +#include "tests/utils/preconversion.h" #include "tests/utils/utils.h" // The values for the floating point tolerance are chose somewhat heuristically based @@ -135,6 +136,14 @@ CATCH_TEST_CASE("Testing CosineSimilarity", "[distance][cosinesimilarity_distanc } } +CATCH_TEST_CASE( + "CosineSimilarity with pre-converted operands", + "[distance][cosinesimilarity_distance][conversion]" +) { + svs_test::require_preconverted_operands_match( + ); +} + // This function is capable of running both the static and dynamically sized versions // of the distance computation kernels. // diff --git a/tests/svs/core/distances/inner_product.cpp b/tests/svs/core/distances/inner_product.cpp index a074a058e..b1c5f095d 100644 --- a/tests/svs/core/distances/inner_product.cpp +++ b/tests/svs/core/distances/inner_product.cpp @@ -37,6 +37,7 @@ // Testing Utilities #include "tests/utils/benchmarks.h" #include "tests/utils/generators.h" +#include "tests/utils/preconversion.h" #include "tests/utils/utils.h" // The values for the floating point tolerance are chose somewhat heuristically based @@ -134,6 +135,13 @@ CATCH_TEST_CASE("Testing InnerProduct Distance", "[distance][innerproduct_distan } } +CATCH_TEST_CASE( + "InnerProduct with pre-converted operands", + "[distance][innerproduct_distance][conversion]" +) { + svs_test::require_preconverted_operands_match(); +} + // This function is capable of running both the static and dynamically sized versions // of the distance computation kernels. // diff --git a/tests/svs/index/vamana/prune.cpp b/tests/svs/index/vamana/prune.cpp index 1c67a5afa..3f3db62b1 100644 --- a/tests/svs/index/vamana/prune.cpp +++ b/tests/svs/index/vamana/prune.cpp @@ -17,9 +17,165 @@ // header under test #include "svs/index/vamana/prune.h" +// svs +#include "svs/core/data/simple.h" +#include "svs/core/distance.h" +#include "svs/lib/float16.h" +#include "svs/lib/neighbor.h" + // catch2 +#include "catch2/catch_template_test_macros.hpp" #include "catch2/catch_test_macros.hpp" +// tests +#include "tests/utils/generators.h" +#include "tests/utils/preconversion.h" +#include "tests/utils/utils.h" + +// stl +#include +#include +#include +#include +#include + +namespace { + +namespace v = svs::index::vamana; + +template +svs::data::SimpleData random_dataset(size_t size, size_t dims) { + auto data = svs::data::SimpleData(size, dims); + auto generator = svs_test::make_generator(-1, 1); + auto datum = std::vector{}; + for (size_t i = 0; i < size; ++i) { + svs_test::populate(datum, generator, dims); + svs_test::insert_subnormals(datum); + data.set_datum(i, datum); + } + return data; +} + +// A `float` copy holding exactly what the kernels at the dispatched instruction-set level +// compute with when they convert `data`. +template +svs::data::SimpleData +to_float(const svs::data::SimpleData& data) { + auto copy = svs::data::SimpleData(data.size(), data.dimensions()); + for (size_t i = 0; i < data.size(); ++i) { + copy.set_datum(i, svs_test::convert_like_kernels(data.get_datum(i))); + } + return copy; +} + +// Every vector of `data`, ordered by its distance to `node` as a build orders its +// candidate pool. The pool includes `node` itself, which the prune must skip. +template +std::vector> +sorted_pool(Distance distance, const Data& data, size_t node) { + auto query = data.get_datum(node); + svs::distance::maybe_fix_argument(distance, query); + auto pool = std::vector>{}; + for (size_t i = 0; i < data.size(); ++i) { + pool.emplace_back(i, svs::distance::compute(distance, query, data.get_datum(i))); + } + std::sort( + pool.begin(), pool.end(), svs::TotalOrder(svs::distance::comparator(distance)) + ); + return pool; +} + +template +std::vector prune( + const Data& data, + Distance distance, + float alpha, + size_t max_result_size, + size_t node, + const std::vector>& pool +) { + auto result = std::vector{}; + v::heuristic_prune_neighbors( + v::prune_strategy(distance), + max_result_size, + alpha, + data, + svs::data::GetDatumAccessor{}, + distance, + node, + svs::lib::as_const_span(pool), + result + ); + return result; +} + +// The query of every prune is converted to `float` once and compared against `Float16` +// candidates. Pruning a `float` copy converted the way the kernels convert runs the +// `float x float` kernels, which compute bit-identical distances at the same +// instruction-set level, so the pruned neighbors must be identical too. +template +void require_float16_prune_matches_float(size_t dims) { + constexpr size_t num_vectors = 300; + // `alpha < 1` reruns the pruning rounds with re-enabled candidates. + constexpr auto alphas = std::to_array({1.0F, 0.95F}); + // Fewer slots than survivors stops a round early; more lets every round run out. + constexpr auto result_sizes = std::to_array({8, 64}); + constexpr auto nodes = std::to_array({0, 17, 299}); + + auto half_data = random_dataset(num_vectors, dims); + using accessor_type = svs::data::GetDatumAccessor; + + svs_test::for_each_supported_isa([&](std::string_view isa) { + CATCH_INFO("ISA: " << isa); + auto float_data = to_float(half_data); + auto storage = + v::prune_query_storage_t( + half_data.dimensions() + ); + + for (float alpha : alphas) { + for (size_t max_result_size : result_sizes) { + for (size_t node : nodes) { + CATCH_INFO( + "alpha: " << alpha << ", max_result_size: " << max_result_size + << ", node: " << node + ); + auto pool = sorted_pool(Distance(), float_data, node); + auto expected = + prune(float_data, Distance(), alpha, max_result_size, node, pool); + CATCH_REQUIRE(!expected.empty()); + + // Through the signature that allocates its own storage. + CATCH_REQUIRE( + prune(half_data, Distance(), alpha, max_result_size, node, pool) == + expected + ); + + // Through caller-held storage, reused across every prune of this + // loop: nothing a previous query left in it may change the result. + auto distance = Distance(); + auto result = std::vector{}; + v::heuristic_prune_neighbors( + v::IterativePruneStrategy{}, + max_result_size, + alpha, + half_data, + accessor_type{}, + distance, + node, + svs::lib::as_const_span(pool), + result, + storage + ); + CATCH_REQUIRE(result == expected); + } + } + } + }); +} + +} // namespace + CATCH_TEST_CASE("Pruning", "[index][vamana]") { namespace v = svs::index::vamana; // Protect against changes to the default strategies getting merged. @@ -47,3 +203,81 @@ CATCH_TEST_CASE("Pruning", "[index][vamana]") { } } } + +CATCH_TEMPLATE_TEST_CASE( + "Iterative pruning with a pre-converted query", + "[index][vamana][conversion]", + svs::distance::DistanceIP, + svs::distance::DistanceCosineSimilarity +) { + using Distance = TestType; + static_assert(std::is_same_v, v::IterativePruneStrategy>); + + // Ragged tail, the specialized static extent, and a high dimensionality. + CATCH_SECTION("Dynamic extent, 100") { + require_float16_prune_matches_float(100); + } + CATCH_SECTION("Static extent, 128") { + require_float16_prune_matches_float(128); + } + CATCH_SECTION("Dynamic extent, 1000") { + require_float16_prune_matches_float(1000); + } +} + +CATCH_TEST_CASE("Pruning with query storage it does not use", "[index][vamana]") { + // Callers pass storage regardless of the strategy their distance selects. The + // strategies that do not convert their query must ignore it and prune as before. + using Distance = svs::distance::DistanceL2; + static_assert(std::is_same_v< + v::prune_strategy_t, + v::ProgressivePruneStrategy>); + + auto data = random_dataset(300, 100); + using accessor_type = svs::data::GetDatumAccessor; + using storage_type = v::prune_query_storage_t; + // L2 does not convert `Float16`, so the storage is the forwarding kind. + static_assert(std::is_same_v< + svs::distance:: + distance_argument_t>, + std::span>); + auto storage = storage_type(data.dimensions()); + + constexpr size_t node = 0; + constexpr size_t max_result_size = 16; + constexpr float alpha = 1.2F; + auto pool = sorted_pool(Distance(), data, node); + auto require_storage_ignored = [&](auto dispatch) { + auto distance = Distance(); + auto expected = std::vector{}; + v::heuristic_prune_neighbors( + dispatch, + max_result_size, + alpha, + data, + accessor_type{}, + distance, + node, + svs::lib::as_const_span(pool), + expected + ); + auto result = std::vector{}; + v::heuristic_prune_neighbors( + dispatch, + max_result_size, + alpha, + data, + accessor_type{}, + distance, + node, + svs::lib::as_const_span(pool), + result, + storage + ); + CATCH_REQUIRE(!expected.empty()); + CATCH_REQUIRE(result == expected); + }; + + CATCH_SECTION("Progressive") { require_storage_ignored(v::ProgressivePruneStrategy{}); } + CATCH_SECTION("Legacy") { require_storage_ignored(v::LegacyPruneStrategy{}); } +} diff --git a/tests/utils/preconversion.h b/tests/utils/preconversion.h new file mode 100644 index 000000000..cc2a64f88 --- /dev/null +++ b/tests/utils/preconversion.h @@ -0,0 +1,164 @@ +/* + * Copyright 2026 Intel Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +// svs +#include "svs/concepts/distance.h" +#include "svs/core/distance/conversion.h" +#include "svs/core/distance/distance_core.h" +#include "svs/lib/float16.h" +#include "svs/lib/threads/threadlocal.h" + +// tests +#include "tests/utils/generators.h" +#include "tests/utils/utils.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" + +// stl +#include +#include +#include +#include +#include +#include +#include + +namespace svs_test { + +template +using aligned_vector = std::vector>; + +/// +/// @brief Convert ``values`` to ``float`` the way the distance kernels at the currently +/// dispatched instruction-set level convert ``Float16`` operands. +/// +/// The Intel(R) AVX kernels convert in hardware, which keeps subnormal halves exact. The +/// generic kernels use ``svs::Float16``'s software conversion, which flushes them to zero. +/// A ``float`` copy built any other way does not match the kernels it is compared against. +/// +inline aligned_vector convert_like_kernels(std::span values) { + auto converted = aligned_vector(values.size()); + svs::distance::Convert::apply(values.data(), converted.data(), values.size()); + return converted; +} + +/// +/// @brief Overwrite a few elements of ``values`` with subnormal halves of both signs. +/// +/// Uniformly drawn values are subnormal too rarely to exercise the one input class on +/// which the hardware and software conversions disagree. +/// +inline void insert_subnormals(std::vector& values) { + // Smallest, a middle, and the largest subnormal; positive and negative. + constexpr auto subnormals = + std::to_array({0x0001, 0x8001, 0x0200, 0x8200, 0x03FF, 0x83FF}); + for (size_t i = 0; i < subnormals.size() && i < values.size(); ++i) { + // Spread them out so they land in different SIMD lanes and loop iterations. + values[(i * 37) % values.size()] = svs::Float16::from_raw(subnormals[i]); + } +} + +namespace detail { + +template +void require_preconverted_operands_match(size_t dims, size_t num_trials) { + CATCH_INFO("Extent: " << Extent << ", dimensions: " << dims); + auto generator = make_generator(-1, 1); + auto a = std::vector{}; + auto b = std::vector{}; + + for (size_t trial = 0; trial < num_trials; ++trial) { + populate(a, generator, dims); + populate(b, generator, dims); + insert_subnormals(a); + insert_subnormals(b); + auto a_float = convert_like_kernels(a); + auto b_float = convert_like_kernels(b); + + auto a_half = std::span(a.data(), dims); + auto b_half = std::span(b.data(), dims); + auto a_full = std::span(a_float.data(), dims); + auto b_full = std::span(b_float.data(), dims); + + // At one instruction-set level, every kernel below shares one SIMD operation and + // `Convert` yields the register values the kernel's own conversion would, so the + // results must be bit-identical. + // + // One functor, fixed once, serves all four calls. `fix_argument` is not a kernel: + // for cosine, given `Float16` it computes the left-hand norm through `Float16`'s + // software conversion, which flushes subnormals to zero; given converted `float`, + // it uses the converted values, which keep subnormals exact at the Intel(R) AVX + // levels. Fixing it per call would compare those two norms, not the kernels. + auto f = Distance{}; + svs::distance::maybe_fix_argument(f, a_full); + // + // Compared as bit patterns: a vector whose every element flushes to zero has an + // undefined cosine, and NaN does not compare equal to itself. + auto bits = [&](auto left, auto right) { + return std::bit_cast(svs::distance::compute(f, left, right)); + }; + auto expected = bits(a_half, b_half); + CATCH_REQUIRE(bits(a_full, b_half) == expected); + CATCH_REQUIRE(bits(a_half, b_full) == expected); + CATCH_REQUIRE(bits(a_full, b_full) == expected); + } +} + +// Each static extent in `supported_dim_list`, the list the kernels are specialized for. +template +void require_static_extents_match(size_t num_trials, std::index_sequence) { + ( + [&]() { + constexpr size_t extent = svs::distance::supported_dim_list[Is]; + if constexpr (extent != svs::Dynamic) { + require_preconverted_operands_match(extent, num_trials); + } + }(), + ... + ); +} + +} // namespace detail + +/// +/// @brief Require that converting either or both ``Float16`` arguments of +/// ``distance::compute`` to ``float`` ahead of time leaves the result bit-identical. +/// +/// @tparam Distance The distance functor under test. +/// +/// Runs at every supported instruction-set level, for each statically sized dimension the +/// kernels specialize and for dynamically sized arguments of several lengths, including +/// ragged tails and dimensionalities in the thousands. +/// +template void require_preconverted_operands_match() { + constexpr size_t num_trials = 20; + for_each_supported_isa([&](std::string_view isa) { + CATCH_INFO("ISA: " << isa); + detail::require_static_extents_match( + num_trials, std::make_index_sequence{} + ); + for (size_t dims : {1, 17, 100, 128, 1000, 4096, 4099}) { + detail::require_preconverted_operands_match( + dims, num_trials + ); + } + }); +} + +} // namespace svs_test diff --git a/tests/utils/utils.h b/tests/utils/utils.h index 6044cbe57..9b36b0219 100644 --- a/tests/utils/utils.h +++ b/tests/utils/utils.h @@ -24,13 +24,16 @@ #include #include #include +#include #include #include "svs/core/data/simple.h" #include "svs/core/query_result.h" +#include "svs/lib/avx_detection.h" #include "svs/lib/exception.h" #include "svs/lib/float16.h" #include "svs/lib/misc.h" +#include "svs/lib/scopeguard.h" #include "svs/third-party/fmt.h" #include "svs/third-party/toml.h" @@ -317,4 +320,36 @@ struct GetDistanceTester { } } }; + +///// +///// Instruction-set dispatch +///// + +/// +/// @brief Invoke ``f(isa_name)`` once per instruction-set level the host supports. +/// +/// Covers each level the ``svs::distance`` dispatchers select between, from the highest +/// down to the generic path. Lower levels are reached by temporarily clearing the runtime +/// flags the dispatchers read, so every level's kernels run on a single machine. The flags +/// are restored before returning, including when ``f`` throws, because every later test +/// in the process dispatches on them. +/// +template void for_each_supported_isa(F&& f) { +#if defined(__x86_64__) + auto& flags = const_cast(svs::detail::avx_runtime_flags); + const auto original = flags; + auto restore = svs::lib::make_scope_guard([&]() noexcept { flags = original; }); + + if (original.avx512f) { + f(std::string_view{"AVX512"}); + } + flags.avx512vnni = false; + flags.avx512f = false; + if (original.avx2) { + f(std::string_view{"AVX2"}); + } + flags.avx2 = false; +#endif + f(std::string_view{"generic"}); +} } // namespace svs_test