diff --git a/src/VecSim/spaces/computer/calculator.h b/src/VecSim/spaces/computer/calculator.h index c12699946..aec1015d6 100644 --- a/src/VecSim/spaces/computer/calculator.h +++ b/src/VecSim/spaces/computer/calculator.h @@ -11,6 +11,8 @@ #include "VecSim/memory/vecsim_base.h" #include "VecSim/spaces/spaces.h" +#include "VecSim/types/sq8.h" +#include "VecSim/utils/alignment.h" enum class DistanceMode { StoredToStored, @@ -120,3 +122,118 @@ class DistanceCalculatorCommon return DistanceDispatch::stateless(func); } }; + +/** + * Distance calculator for mean-normalized SQ8 indices. + * + * Stored vectors are SQ8-quantized from x' = x - mean. IP queries remain unshifted and use + * analytical correction terms derived from the per-vector x_mean_ip / y_mean_ip metadata. + * L2 queries are centered once during preprocessing, allowing the base kernel to compute the + * original distance directly without cancellation-prone correction terms. + * + * Correction formulas (expressed in terms of distance semantics): + * + * Asymmetric (stored vector x as SQ8 of x'): + * IP distance: dist(x,y) = asym_func(x_blob, y_blob) - y_mean_ip + * L2 distance: dist(x,y) = asym_func(x_blob, centered_y_blob) + * where centered_y = y - mean + * + * Symmetric (both x and y as SQ8 blobs): + * IP distance: dist(x,y) = sym_func(x_blob, y_blob) - x_mean_ip - y_mean_ip + * + mean_sum_squares + * L2 distance: dist(x,y) = sym_func(x_blob, y_blob) (mean terms cancel exactly) + * + * Note: IP dist functions return 1 - IP(x',y'), not raw inner product. L2 dist functions + * return ||x'-y'||² directly. + * + * DataType is float or float16 + * DistType is float + */ +template +class DistanceCalculatorWithNorm : public IndexCalculatorInterface { + static_assert(Metric == VecSimMetric_L2 || Metric == VecSimMetric_IP, + "DistanceCalculatorWithNorm only supports L2 and IP metrics"); + +private: + using sq8 = vecsim_types::sq8; + + struct WithNormDistanceContext { + spaces::dist_func_t stored_func; + spaces::dist_func_t query_func; + float mean_sum_squares; + }; + + const WithNormDistanceContext context_; + + static DistType calcStoredWithContext(const void *opaque_context, const void *v1, + const void *v2, size_t dim) { + const auto *context = static_cast(opaque_context); + DistType base = context->stored_func(v1, v2, dim); + if constexpr (Metric == VecSimMetric_IP) { + // base = 1 - IP(x', y'). We want 1 - IP(x, y). + // IP(x, y) = IP(x', y') + x_mean_ip + y_mean_ip - mean_sum_squares + // => 1 - IP(x, y) = base - x_mean_ip - y_mean_ip + mean_sum_squares + float x_mean_ip = + load_unaligned(static_cast(v1) + dim + + sq8::template mean_ip_index() * sizeof(float)); + float y_mean_ip = + load_unaligned(static_cast(v2) + dim + + sq8::template mean_ip_index() * sizeof(float)); + return base - x_mean_ip - y_mean_ip + context->mean_sum_squares; + } else { // L2: ||x - y||² = ||x' - y'||² (mean terms cancel exactly) + return base; + } + } + + static DistType calcQueryWithContext(const void *opaque_context, const void *candidate, + const void *query, size_t dim) { + const auto *context = static_cast(opaque_context); + DistType base = context->query_func(candidate, query, dim); + if constexpr (Metric == VecSimMetric_IP) { + float y_mean_ip = + load_unaligned(static_cast(query) + dim * sizeof(DataType) + + sq8::template query_mean_ip_index() * sizeof(float)); + // base = 1 - IP(x', y). We want 1 - IP(x, y). + // IP(x, y) = IP(x', y) + y_mean_ip + // => 1 - IP(x, y) = base - y_mean_ip + return base - y_mean_ip; + } + // L2 query preprocessing centers y, so base = ||(x - mean) - (y - mean)||². + return base; + } + +public: + DistanceCalculatorWithNorm(std::shared_ptr allocator, + spaces::dist_func_t asym_func, + spaces::dist_func_t sym_func, float mean_sum_squares) + : IndexCalculatorInterface(allocator), context_(WithNormDistanceContext{ + .stored_func = sym_func, + .query_func = asym_func, + .mean_sum_squares = mean_sum_squares, + }) {} + + // Symmetric: both v1 and v2 are stored SQ8-of-x' blobs. + DistType calcDistance(const void *v1, const void *v2, size_t dim) const override { + return calcStoredWithContext(&context_, v1, v2, dim); + } + + // Asymmetric: IP queries are raw; L2 queries are centered during preprocessing. The candidate + // is a stored SQ8-of-x' blob. + // Note: query_dist_func expects (storage, query) argument order. + DistType calcDistanceForQuery(const void *candidate, const void *query, + size_t dim) const override { + return calcQueryWithContext(&context_, candidate, query, dim); + } + + DistanceDispatch getDistanceDispatch(DistanceMode mode) const override { + if constexpr (Metric == VecSimMetric_L2) { + auto func = + mode == DistanceMode::StoredToStored ? context_.stored_func : context_.query_func; + return DistanceDispatch::stateless(func); + } + if (mode == DistanceMode::StoredToStored) { + return DistanceDispatch::stateful(&context_, calcStoredWithContext); + } + return DistanceDispatch::stateful(&context_, calcQueryWithContext); + } +}; diff --git a/src/VecSim/spaces/computer/preprocessors.h b/src/VecSim/spaces/computer/preprocessors.h index c7db23c57..e2f0d387c 100644 --- a/src/VecSim/spaces/computer/preprocessors.h +++ b/src/VecSim/spaces/computer/preprocessors.h @@ -167,8 +167,9 @@ class CosinePreprocessor : public PreprocessorInterface { * * Query processing: * The query vector is not quantized. It remains in DataType width (FP32 stays FP32, FP16 stays - * FP16), but we precompute and store metric-specific FP32 values to accelerate asymmetric - * distance computation: + * FP16). Normalized L2 stores y' = y - mean; all other modes store y unchanged. We precompute and + * store metric-specific FP32 values from the query body to accelerate asymmetric distance + * computation: * - For IP/Cosine: y_sum = Σy_i (sum of query values) * - For L2: y_sum = Σy_i (sum of query values), y_sum_squares = Σy_i² (sum of squared query values) * @@ -199,6 +200,8 @@ class CosinePreprocessor : public PreprocessorInterface { * - x_sum_squares = Σx_i² is precomputed and stored in the storage blob * - IP(x, y) is computed using the formula above * - y_sum_squares = Σy_i² is precomputed and stored in the query blob + * For normalized L2, x and y in this formula are the centered values x' and y'; their distance + * is identical to the distance between the original vectors. * * === Symmetric distance (both x and y are quantized) === * @@ -213,13 +216,13 @@ class CosinePreprocessor : public PreprocessorInterface { * = min_x * sum_y + min_y * sum_x - dim * min_x * min_y * + delta_x * delta_y * Σ(qx_i * qy_i) * where: - * - sum_x, sum_y are precomputed sums of original values + * - sum_x, sum_y are precomputed sums of the values represented by each blob * - Σqx_i = (sum_x - dim * min_x) / delta_x (sum of quantized values, derived from stored sum) * - Σqy_i = (sum_y - dim * min_y) / delta_y * * For L2: * ||x - y||² = sum_sq_x + sum_sq_y - 2 * IP(x, y) - * where sum_sq_x, sum_sq_y are precomputed sums of squared original values. + * where sum_sq_x, sum_sq_y are precomputed sums of the squared represented values. */ // Input types accepted by QuantPreprocessor. Opt-in via std::same_as so unrelated types // (e.g. integers, double, bfloat16) are rejected at the template head with a named constraint. @@ -237,6 +240,15 @@ static inline float to_fp32(T x) { } } +template +static inline T from_fp32(float x) { + if constexpr (std::is_same_v) { + return vecsim_types::FP32_to_FP16(x); + } else { + return x; + } +} + template class QuantPreprocessor : public PreprocessorInterface { using OUTPUT_TYPE = uint8_t; @@ -330,14 +342,16 @@ class QuantPreprocessor : public PreprocessorInterface { memcpy(meta_dst, buf, n * sizeof(MetadataType)); } - // Computes and writes query metadata (FP32) in a single pass over the input vector. + // Computes and writes query metadata (FP32) in a single pass over the query values. // Without norm: // For IP/Cosine: writes y_sum = Σy_i // For L2: writes y_sum = Σy_i and y_sum_squares = Σy_i² - // With norm: additionally appends y_mean_ip = Σ(mean_i * y_i). + // With norm: additionally appends y_mean_ip = Σ(mean_i * original_y_i). + // For normalized L2, values contains y - mean while original_input contains y. // The output pointer addresses the metadata region after the query body and may not be // 4-byte aligned (e.g. FP16 query body with odd dim), so writes go through memcpy. - void assign_query_metadata(const DataType *input, void *output_metadata) const { + void assign_query_metadata(const DataType *values, const DataType *original_input, + void *output_metadata) const { // Accumulators are FP32 to preserve precision for FP16 inputs. // 4 independent accumulators for sum @@ -352,10 +366,10 @@ class QuantPreprocessor : public PreprocessorInterface { size_t dim_round_down = this->dim & ~size_t(3); for (; i < dim_round_down; i += 4) { - const float y0 = to_fp32(input[i + 0]); - const float y1 = to_fp32(input[i + 1]); - const float y2 = to_fp32(input[i + 2]); - const float y3 = to_fp32(input[i + 3]); + const float y0 = to_fp32(values[i + 0]); + const float y1 = to_fp32(values[i + 1]); + const float y2 = to_fp32(values[i + 2]); + const float y3 = to_fp32(values[i + 3]); s0 += y0; s1 += y1; @@ -370,10 +384,10 @@ class QuantPreprocessor : public PreprocessorInterface { } if constexpr (WithNorm) { - m0 += mean[i + 0] * y0; - m1 += mean[i + 1] * y1; - m2 += mean[i + 2] * y2; - m3 += mean[i + 3] * y3; + m0 += mean[i + 0] * to_fp32(original_input[i + 0]); + m1 += mean[i + 1] * to_fp32(original_input[i + 1]); + m2 += mean[i + 2] * to_fp32(original_input[i + 2]); + m3 += mean[i + 3] * to_fp32(original_input[i + 3]); } } @@ -384,13 +398,13 @@ class QuantPreprocessor : public PreprocessorInterface { // Tail: handle remaining elements for (; i < this->dim; ++i) { - const float y = to_fp32(input[i]); + const float y = to_fp32(values[i]); sum += y; if constexpr (Metric == VecSimMetric_L2) { sum_squares += y * y; } if constexpr (WithNorm) { - y_mean_ip += mean[i] * y; + y_mean_ip += mean[i] * to_fp32(original_input[i]); } } @@ -488,7 +502,11 @@ class QuantPreprocessor : public PreprocessorInterface { /** * Preprocesses the query vector for asymmetric distance computation. * - * The query blob contains the original DataType values followed by FP32 precomputed values: + * The query blob contains DataType values followed by FP32 precomputed values. Normalized L2 + * stores centered query values (y - mean), so its distance kernel directly computes + * ||(x - mean) - (y - mean)||² without cancellation-prone correction terms. Other modes keep + * the original query values. + * * - For IP/Cosine: y_sum = Σy_i (sum of query values) * - For L2: y_sum = Σy_i (sum of query values), y_sum_squares = Σy_i² (sum of squared query * values) @@ -508,13 +526,20 @@ class QuantPreprocessor : public PreprocessorInterface { // Allocate aligned memory for the query blob blob = this->allocator->allocate_aligned(this->query_bytes_count, alignment); const size_t body_bytes = this->dim * sizeof(DataType); - memcpy(blob, original_blob, body_bytes); const DataType *input = static_cast(original_blob); + DataType *query_values = static_cast(blob); + if constexpr (WithNorm && Metric == VecSimMetric_L2) { + for (size_t i = 0; i < this->dim; ++i) { + query_values[i] = from_fp32(to_fp32(input[i]) - this->mean[i]); + } + } else { + memcpy(query_values, original_blob, body_bytes); + } // Compute and write FP32 query metadata after the query body. The metadata offset is // body_bytes, which is not guaranteed to be 4-byte aligned for FP16 query bodies. void *metadata_dst = static_cast(blob) + body_bytes; - assign_query_metadata(input, metadata_dst); + assign_query_metadata(query_values, input, metadata_dst); query_blob_size = this->query_bytes_count; } diff --git a/tests/unit/test_components.cpp b/tests/unit/test_components.cpp index 224388952..2dad33f8b 100644 --- a/tests/unit/test_components.cpp +++ b/tests/unit/test_components.cpp @@ -13,6 +13,8 @@ #include "VecSim/spaces/computer/preprocessor_container.h" #include "VecSim/spaces/computer/preprocessors.h" #include "VecSim/spaces/computer/calculator.h" +#include "VecSim/spaces/IP_space.h" +#include "VecSim/spaces/L2_space.h" #include "unit_test_utils.h" #include "tests_utils.h" @@ -1663,15 +1665,22 @@ class QuantPreprocessorWithNormMetricTestBase : public testing::TestWithParam(centered[i]); + } else { + expected_query_body[i] = original_blob[i]; + } + const float query_value = to_fp32(expected_query_body[i]); expected_x_mean_ip += widened_blob[i] * mean_vec[i]; - expected_y_sum += widened_blob[i]; - expected_y_sum_squares += widened_blob[i] * widened_blob[i]; + expected_y_sum += query_value; + expected_y_sum_squares += query_value * query_value; expected_y_mean_ip += widened_blob[i] * mean_vec[i]; } @@ -1710,7 +1719,7 @@ class QuantPreprocessorWithNormMetricTestBase : public testing::TestWithParam() * sizeof(float)), expected_x_mean_ip); EXPECT_NO_FATAL_FAILURE(CompareVectors( - static_cast(query_blob), original_blob, dim)); + static_cast(query_blob), expected_query_body, dim)); ASSERT_FLOAT_EQ( load_meta(query_blob, query_meta_offset + sq8::SUM_QUERY * sizeof(float)), expected_y_sum); @@ -1742,8 +1751,8 @@ class QuantPreprocessorWithNormMetricTestBase : public testing::TestWithParampreprocessQuery(original_blob, blob, blob_size, alignment); ASSERT_NE(blob, nullptr); ASSERT_EQ(blob_size, expected_query_size); - EXPECT_NO_FATAL_FAILURE( - CompareVectors(static_cast(blob), original_blob, dim)); + EXPECT_NO_FATAL_FAILURE(CompareVectors(static_cast(blob), + expected_query_body, dim)); ASSERT_FLOAT_EQ(load_meta(blob, query_meta_offset + sq8::SUM_QUERY * sizeof(float)), expected_y_sum); if constexpr (Metric == VecSimMetric_L2) { @@ -1802,3 +1811,529 @@ INSTANTIATE_TEST_SUITE_P(QuantPreprocessorFP16WithNormTests, [](const testing::TestParamInfo &info) { return VecSimMetric_ToString(info.param); }); + +// Helper: build storage blob from original vector x and mean. +static void buildStorageBlob(const std::shared_ptr &allocator, const float *x, + const float *mean, size_t dim, VecSimMetric metric, + void *&storage_blob) { + vecsim_stl::vector mean_vec(allocator); + for (size_t i = 0; i < dim; ++i) + mean_vec.push_back(mean[i]); + + storage_blob = nullptr; + size_t sz = dim * sizeof(float); + if (metric == VecSimMetric_IP) { + auto *pp = new (allocator) + QuantPreprocessor(allocator, dim, mean_vec); + pp->preprocessForStorage(x, storage_blob, sz, 0); + delete pp; + } else { + auto *pp = new (allocator) + QuantPreprocessor(allocator, dim, mean_vec); + pp->preprocessForStorage(x, storage_blob, sz, 0); + delete pp; + } +} + +// Helper: build query blob from original vector y and mean. +static void buildQueryBlob(const std::shared_ptr &allocator, const float *y, + const float *mean, size_t dim, VecSimMetric metric, void *&query_blob) { + vecsim_stl::vector mean_vec(allocator); + for (size_t i = 0; i < dim; ++i) + mean_vec.push_back(mean[i]); + + query_blob = nullptr; + size_t sz = dim * sizeof(float); + if (metric == VecSimMetric_IP) { + auto *pp = new (allocator) + QuantPreprocessor(allocator, dim, mean_vec); + pp->preprocessQuery(y, query_blob, sz, 0); + delete pp; + } else { + auto *pp = new (allocator) + QuantPreprocessor(allocator, dim, mean_vec); + pp->preprocessQuery(y, query_blob, sz, 0); + delete pp; + } +} + +// Brute-force IP distance on original (unshifted) float vectors: 1 - dot(x, y). +static float bruteForceIPDist(const float *x, const float *y, size_t dim) { + float dot = 0.0f; + for (size_t i = 0; i < dim; ++i) + dot += x[i] * y[i]; + return 1.0f - dot; +} + +// Brute-force L2 squared distance on original float vectors: sum((x_i - y_i)^2). +static float bruteForceL2Dist(const float *x, const float *y, size_t dim) { + float sum = 0.0f; + for (size_t i = 0; i < dim; ++i) { + float d = x[i] - y[i]; + sum += d * d; + } + return sum; +} + +// Compute mean_sum_squares = sum(mean_i^2). +static float computeMeanSumSquares(const float *mean, size_t dim) { + float s = 0.0f; + for (size_t i = 0; i < dim; ++i) + s += mean[i] * mean[i]; + return s; +} + +TEST(DistanceCalculatorWithNormTest, CalcDistanceForQuery_IP_FP32) { + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 8; + float x[dim] = {1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f}; + float y[dim] = {0.5f, 1.5f, 2.5f, 3.5f, 0.5f, 1.5f, 2.5f, 3.5f}; + float mean[dim] = {0.5f, 1.0f, 1.5f, 2.0f, 0.5f, 1.0f, 1.5f, 2.0f}; + + float mean_sum_sq = computeMeanSumSquares(mean, dim); + + void *storage_blob = nullptr; + void *query_blob = nullptr; + buildStorageBlob(allocator, x, mean, dim, VecSimMetric_IP, storage_blob); + buildQueryBlob(allocator, y, mean, dim, VecSimMetric_IP, query_blob); + + auto asym_func = spaces::IP_SQ8_FP32_GetDistFunc(dim); + auto sym_func = spaces::IP_SQ8_SQ8_GetDistFunc(dim); + auto *calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, mean_sum_sq); + + float got = calc->calcDistanceForQuery(storage_blob, query_blob, dim); + float expected = bruteForceIPDist(x, y, dim); + auto dispatch = calc->getDistanceDispatch(DistanceMode::StoredToQuery); + + // Allow quantization error + EXPECT_NEAR(got, expected, 0.05f) << "Asymmetric IP distance mismatch"; + ASSERT_TRUE(dispatch.isValid()); + EXPECT_EQ(dispatch.stateless_func, nullptr); + EXPECT_NE(dispatch.stateful_func, nullptr); + EXPECT_NEAR(dispatch(storage_blob, query_blob, dim), expected, 0.05f); + + allocator->free_allocation(storage_blob); + allocator->free_allocation(query_blob); + delete calc; +} + +TEST(DistanceCalculatorWithNormTest, CalcDistanceForQuery_L2_FP32) { + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 8; + float x[dim] = {1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f}; + float y[dim] = {0.5f, 1.5f, 2.5f, 3.5f, 0.5f, 1.5f, 2.5f, 3.5f}; + float mean[dim] = {0.5f, 1.0f, 1.5f, 2.0f, 0.5f, 1.0f, 1.5f, 2.0f}; + + float mean_sum_sq = computeMeanSumSquares(mean, dim); + + void *storage_blob = nullptr; + void *query_blob = nullptr; + buildStorageBlob(allocator, x, mean, dim, VecSimMetric_L2, storage_blob); + buildQueryBlob(allocator, y, mean, dim, VecSimMetric_L2, query_blob); + + auto asym_func = spaces::L2_SQ8_FP32_GetDistFunc(dim); + auto sym_func = spaces::L2_SQ8_SQ8_GetDistFunc(dim); + auto *calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, mean_sum_sq); + + float got = calc->calcDistanceForQuery(storage_blob, query_blob, dim); + float expected = bruteForceL2Dist(x, y, dim); + auto dispatch = calc->getDistanceDispatch(DistanceMode::StoredToQuery); + + EXPECT_NEAR(got, expected, 0.05f) << "Asymmetric L2 distance mismatch"; + ASSERT_TRUE(dispatch.isValid()); + EXPECT_EQ(dispatch.stateless_func, asym_func); + EXPECT_EQ(dispatch.stateful_func, nullptr); + EXPECT_NEAR(dispatch(storage_blob, query_blob, dim), expected, 0.05f); + + allocator->free_allocation(storage_blob); + allocator->free_allocation(query_blob); + delete calc; +} + +TEST(DistanceCalculatorWithNormTest, L2LargeOffsetSmallDistance) { + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 128; + float x[dim]; + float y[dim]; + float mean[dim]; + for (size_t i = 0; i < dim; ++i) { + mean[i] = 1000.0f; + x[i] = 1001.0f; + y[i] = 1001.1f; + } + + void *storage_blob = nullptr; + void *query_blob = nullptr; + buildStorageBlob(allocator, x, mean, dim, VecSimMetric_L2, storage_blob); + buildQueryBlob(allocator, y, mean, dim, VecSimMetric_L2, query_blob); + + auto asym_func = spaces::L2_SQ8_FP32_GetDistFunc(dim); + auto sym_func = spaces::L2_SQ8_SQ8_GetDistFunc(dim); + auto *calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, computeMeanSumSquares(mean, dim)); + + const float expected = bruteForceL2Dist(x, y, dim); + const float direct = calc->calcDistanceForQuery(storage_blob, query_blob, dim); + const auto dispatch = calc->getDistanceDispatch(DistanceMode::StoredToQuery); + + ASSERT_TRUE(dispatch.isValid()); + EXPECT_EQ(dispatch.stateless_func, asym_func); + EXPECT_EQ(dispatch.stateful_func, nullptr); + EXPECT_NEAR(direct, expected, 0.001f); + EXPECT_NEAR(dispatch(storage_blob, query_blob, dim), expected, 0.001f); + EXPECT_GE(direct, 0.0f); + + allocator->free_allocation(storage_blob); + allocator->free_allocation(query_blob); + delete calc; +} + +TEST(DistanceCalculatorWithNormTest, CalcDistance_IP_Symmetric) { + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 8; + float x[dim] = {1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f}; + float y[dim] = {0.5f, 1.5f, 2.5f, 3.5f, 0.5f, 1.5f, 2.5f, 3.5f}; + float mean[dim] = {0.5f, 1.0f, 1.5f, 2.0f, 0.5f, 1.0f, 1.5f, 2.0f}; + + float mean_sum_sq = computeMeanSumSquares(mean, dim); + + void *x_blob = nullptr; + void *y_blob = nullptr; + buildStorageBlob(allocator, x, mean, dim, VecSimMetric_IP, x_blob); + buildStorageBlob(allocator, y, mean, dim, VecSimMetric_IP, y_blob); + + auto asym_func = spaces::IP_SQ8_FP32_GetDistFunc(dim); + auto sym_func = spaces::IP_SQ8_SQ8_GetDistFunc(dim); + auto *calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, mean_sum_sq); + + float got = calc->calcDistance(x_blob, y_blob, dim); + float expected = bruteForceIPDist(x, y, dim); + auto dispatch = calc->getDistanceDispatch(DistanceMode::StoredToStored); + + EXPECT_NEAR(got, expected, 0.05f) << "Symmetric IP distance mismatch"; + ASSERT_TRUE(dispatch.isValid()); + EXPECT_EQ(dispatch.stateless_func, nullptr); + EXPECT_NE(dispatch.stateful_func, nullptr); + EXPECT_NEAR(dispatch(x_blob, y_blob, dim), expected, 0.05f); + + allocator->free_allocation(x_blob); + allocator->free_allocation(y_blob); + delete calc; +} + +TEST(DistanceCalculatorWithNormTest, CalcDistance_L2_Symmetric) { + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 8; + float x[dim] = {1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f}; + float y[dim] = {0.5f, 1.5f, 2.5f, 3.5f, 0.5f, 1.5f, 2.5f, 3.5f}; + float mean[dim] = {0.5f, 1.0f, 1.5f, 2.0f, 0.5f, 1.0f, 1.5f, 2.0f}; + + float mean_sum_sq = computeMeanSumSquares(mean, dim); + + void *x_blob = nullptr; + void *y_blob = nullptr; + buildStorageBlob(allocator, x, mean, dim, VecSimMetric_L2, x_blob); + buildStorageBlob(allocator, y, mean, dim, VecSimMetric_L2, y_blob); + + auto asym_func = spaces::L2_SQ8_FP32_GetDistFunc(dim); + auto sym_func = spaces::L2_SQ8_SQ8_GetDistFunc(dim); + auto *calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, mean_sum_sq); + + float got = calc->calcDistance(x_blob, y_blob, dim); + float expected = bruteForceL2Dist(x, y, dim); + auto dispatch = calc->getDistanceDispatch(DistanceMode::StoredToStored); + + EXPECT_NEAR(got, expected, 0.05f) << "Symmetric L2 distance mismatch"; + ASSERT_TRUE(dispatch.isValid()); + EXPECT_EQ(dispatch.stateless_func, sym_func); + EXPECT_EQ(dispatch.stateful_func, nullptr); + EXPECT_NEAR(dispatch(x_blob, y_blob, dim), expected, 0.05f); + + allocator->free_allocation(x_blob); + allocator->free_allocation(y_blob); + delete calc; +} + +TEST(DistanceCalculatorWithNormTest, CalcDistanceForQuery_FP16) { + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 8; + float x_fp32[dim] = {1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f}; + float y_fp32[dim] = {0.5f, 1.5f, 2.5f, 3.5f, 0.5f, 1.5f, 2.5f, 3.5f}; + float mean[dim] = {0.5f, 1.0f, 1.5f, 2.0f, 0.5f, 1.0f, 1.5f, 2.0f}; + + using DataType = vecsim_types::float16; + + DataType x[dim], y[dim]; + for (size_t i = 0; i < dim; ++i) { + x[i] = vecsim_types::FP32_to_FP16(x_fp32[i]); + y[i] = vecsim_types::FP32_to_FP16(y_fp32[i]); + } + + vecsim_stl::vector mean_vec(allocator); + for (size_t i = 0; i < dim; ++i) + mean_vec.push_back(mean[i]); + float mean_sum_sq = computeMeanSumSquares(mean, dim); + + // Build storage blob from FP16 x + void *storage_blob = nullptr; + size_t sz = dim * sizeof(DataType); + auto *pp_ip = new (allocator) + QuantPreprocessor(allocator, dim, mean_vec); + pp_ip->preprocessForStorage(x, storage_blob, sz, 0); + delete pp_ip; + + // Build query blob from FP16 y + void *query_blob = nullptr; + sz = dim * sizeof(DataType); + auto *pp_qr = new (allocator) + QuantPreprocessor(allocator, dim, mean_vec); + pp_qr->preprocessQuery(y, query_blob, sz, 0); + delete pp_qr; + + auto asym_func = spaces::IP_SQ8_FP16_GetDistFunc(dim); + auto sym_func = spaces::IP_SQ8_SQ8_GetDistFunc(dim); + auto *calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, mean_sum_sq); + + float got = calc->calcDistanceForQuery(storage_blob, query_blob, dim); + float expected = bruteForceIPDist(x_fp32, y_fp32, dim); + + EXPECT_NEAR(got, expected, 0.05f) << "Asymmetric IP FP16 distance mismatch"; + + allocator->free_allocation(storage_blob); + allocator->free_allocation(query_blob); + delete calc; +} + +TEST(DistanceCalculatorWithNormTest, CalcDistanceForQuery_L2_FP16) { + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 8; + const float x_fp32[dim] = {1.0f, 2.0f, 3.0f, 4.0f, 1.0f, 2.0f, 3.0f, 4.0f}; + const float y_fp32[dim] = {0.5f, 1.5f, 2.5f, 3.5f, 0.5f, 1.5f, 2.5f, 3.5f}; + const float mean[dim] = {0.5f, 1.0f, 1.5f, 2.0f, 0.5f, 1.0f, 1.5f, 2.0f}; + using DataType = vecsim_types::float16; + + DataType x[dim], y[dim]; + vecsim_stl::vector mean_vec(allocator); + for (size_t i = 0; i < dim; ++i) { + x[i] = vecsim_types::FP32_to_FP16(x_fp32[i]); + y[i] = vecsim_types::FP32_to_FP16(y_fp32[i]); + mean_vec.push_back(mean[i]); + } + + void *storage_blob = nullptr; + size_t storage_size = dim * sizeof(DataType); + auto *storage_preprocessor = new (allocator) + QuantPreprocessor(allocator, dim, mean_vec); + storage_preprocessor->preprocessForStorage(x, storage_blob, storage_size, 0); + delete storage_preprocessor; + + void *query_blob = nullptr; + size_t query_size = dim * sizeof(DataType); + auto *query_preprocessor = new (allocator) + QuantPreprocessor(allocator, dim, mean_vec); + query_preprocessor->preprocessQuery(y, query_blob, query_size, 0); + delete query_preprocessor; + + auto asym_func = spaces::L2_SQ8_FP16_GetDistFunc(dim); + auto sym_func = spaces::L2_SQ8_SQ8_GetDistFunc(dim); + auto *calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, computeMeanSumSquares(mean, dim)); + const float expected = bruteForceL2Dist(x_fp32, y_fp32, dim); + const auto dispatch = calc->getDistanceDispatch(DistanceMode::StoredToQuery); + + EXPECT_NEAR(calc->calcDistanceForQuery(storage_blob, query_blob, dim), expected, 0.05f); + ASSERT_TRUE(dispatch.isValid()); + EXPECT_EQ(dispatch.stateless_func, asym_func); + EXPECT_EQ(dispatch.stateful_func, nullptr); + EXPECT_NEAR(dispatch(storage_blob, query_blob, dim), expected, 0.05f); + + allocator->free_allocation(storage_blob); + allocator->free_allocation(query_blob); + delete calc; +} + +TEST(DistanceCalculatorWithNormTest, ZeroMean_MatchesBaseSQ8) { + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 8; + float x[dim] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}; + float y[dim] = {0.5f, 1.5f, 2.5f, 3.5f, 4.5f, 5.5f, 6.5f, 7.5f}; + vecsim_stl::vector zero_mean(dim, 0.0f, allocator); + + // Build WithNorm storage and query blobs using zero mean + auto *pp_ip = + new (allocator) QuantPreprocessor(allocator, dim, zero_mean); + void *x_norm_blob = nullptr, *y_norm_query = nullptr; + size_t sx = dim * sizeof(float), sq = dim * sizeof(float); + pp_ip->preprocessForStorage(x, x_norm_blob, sx, 0); + pp_ip->preprocessQuery(y, y_norm_query, sq, 0); + delete pp_ip; + + auto asym_func = spaces::IP_SQ8_FP32_GetDistFunc(dim); + auto sym_func = spaces::IP_SQ8_SQ8_GetDistFunc(dim); + + auto *norm_calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, 0.0f); + + // With zero mean, correction is 0: result equals raw base function call. + // asym_func expects (storage, query) order. + float norm_asym = norm_calc->calcDistanceForQuery(x_norm_blob, y_norm_query, dim); + float base_asym = asym_func(x_norm_blob, y_norm_query, dim); + EXPECT_FLOAT_EQ(norm_asym, base_asym) + << "WithNorm(zero mean) asymmetric IP should match raw base dist function"; + + // Build a second storage blob for y to test symmetric distance + void *y_norm_blob = nullptr; + sx = dim * sizeof(float); + auto *pp_ip2 = + new (allocator) QuantPreprocessor(allocator, dim, zero_mean); + pp_ip2->preprocessForStorage(y, y_norm_blob, sx, 0); + delete pp_ip2; + + float norm_sym = norm_calc->calcDistance(x_norm_blob, y_norm_blob, dim); + float base_sym = sym_func(x_norm_blob, y_norm_blob, dim); + EXPECT_FLOAT_EQ(norm_sym, base_sym) + << "WithNorm(zero mean) symmetric IP should match raw base dist function"; + + allocator->free_allocation(x_norm_blob); + allocator->free_allocation(y_norm_query); + allocator->free_allocation(y_norm_blob); + delete norm_calc; +} + +TEST(DistanceCalculatorWithNormTest, SymmetricVsAsymmetric_Sanity) { + // For the same two stored vectors, calcDistance (symmetric) and calcDistanceForQuery + // (asymmetric) must agree to within quantization error. + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 8; + float x[dim] = {3.0f, 1.0f, 4.0f, 1.0f, 5.0f, 3.0f, 2.0f, 6.0f}; + float y[dim] = {2.0f, 7.0f, 1.0f, 4.0f, 2.0f, 5.0f, 1.0f, 2.0f}; + float mean[dim] = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f}; + float mean_sum_sq = computeMeanSumSquares(mean, dim); + + // IP + { + void *x_blob = nullptr, *y_blob = nullptr, *y_query = nullptr; + buildStorageBlob(allocator, x, mean, dim, VecSimMetric_IP, x_blob); + buildStorageBlob(allocator, y, mean, dim, VecSimMetric_IP, y_blob); + buildQueryBlob(allocator, y, mean, dim, VecSimMetric_IP, y_query); + + auto asym_func = spaces::IP_SQ8_FP32_GetDistFunc(dim); + auto sym_func = spaces::IP_SQ8_SQ8_GetDistFunc(dim); + auto *calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, mean_sum_sq); + + float sym_dist = calc->calcDistance(x_blob, y_blob, dim); + float asym_dist = calc->calcDistanceForQuery(x_blob, y_query, dim); + // Both should be close to the brute-force answer; use relative tolerance + // since the IP magnitude (~157) amplifies absolute quantization error. + float bf = bruteForceIPDist(x, y, dim); + EXPECT_NEAR(sym_dist, bf, 0.05f) << "Symmetric IP vs brute-force"; + EXPECT_NEAR(asym_dist, bf, 0.05f) << "Asymmetric IP vs brute-force"; + + allocator->free_allocation(x_blob); + allocator->free_allocation(y_blob); + allocator->free_allocation(y_query); + delete calc; + } + + // L2 + { + void *x_blob = nullptr, *y_blob = nullptr, *y_query = nullptr; + buildStorageBlob(allocator, x, mean, dim, VecSimMetric_L2, x_blob); + buildStorageBlob(allocator, y, mean, dim, VecSimMetric_L2, y_blob); + buildQueryBlob(allocator, y, mean, dim, VecSimMetric_L2, y_query); + + auto asym_func = spaces::L2_SQ8_FP32_GetDistFunc(dim); + auto sym_func = spaces::L2_SQ8_SQ8_GetDistFunc(dim); + auto *calc = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_func, sym_func, mean_sum_sq); + + float sym_dist = calc->calcDistance(x_blob, y_blob, dim); + float asym_dist = calc->calcDistanceForQuery(x_blob, y_query, dim); + float bf = bruteForceL2Dist(x, y, dim); + EXPECT_NEAR(sym_dist, bf, 0.05f) << "Symmetric L2 vs brute-force"; + EXPECT_NEAR(asym_dist, bf, 0.05f) << "Asymmetric L2 vs brute-force"; + + allocator->free_allocation(x_blob); + allocator->free_allocation(y_blob); + allocator->free_allocation(y_query); + delete calc; + } +} + +TEST(DistanceCalculatorWithNormTest, RandomVectors) { + // Generate random vector pairs, compute WithNorm distances and verify against brute-force. + std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); + constexpr size_t dim = 16; + std::mt19937 rng(12345); + std::uniform_real_distribution dist(-1.0f, 1.0f); + + float mean[dim]; + for (size_t i = 0; i < dim; ++i) + mean[i] = dist(rng) * 0.5f; + float mean_sum_sq = computeMeanSumSquares(mean, dim); + + auto asym_ip = spaces::IP_SQ8_FP32_GetDistFunc(dim); + auto sym_ip = spaces::IP_SQ8_SQ8_GetDistFunc(dim); + auto asym_l2 = spaces::L2_SQ8_FP32_GetDistFunc(dim); + auto sym_l2 = spaces::L2_SQ8_SQ8_GetDistFunc(dim); + + auto *calc_ip = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_ip, sym_ip, mean_sum_sq); + auto *calc_l2 = new (allocator) DistanceCalculatorWithNorm( + allocator, asym_l2, sym_l2, mean_sum_sq); + + int failures = 0; + for (int trial = 0; trial < 100; ++trial) { + float x[dim], y[dim]; + for (size_t i = 0; i < dim; ++i) { + x[i] = dist(rng); + y[i] = dist(rng); + } + + void *x_blob = nullptr, *y_blob = nullptr; + void *y_query_ip = nullptr, *y_query_l2 = nullptr; + buildStorageBlob(allocator, x, mean, dim, VecSimMetric_IP, x_blob); + buildStorageBlob(allocator, y, mean, dim, VecSimMetric_IP, y_blob); + buildQueryBlob(allocator, y, mean, dim, VecSimMetric_IP, y_query_ip); + + void *x_blob_l2 = nullptr, *y_blob_l2 = nullptr; + buildStorageBlob(allocator, x, mean, dim, VecSimMetric_L2, x_blob_l2); + buildStorageBlob(allocator, y, mean, dim, VecSimMetric_L2, y_blob_l2); + buildQueryBlob(allocator, y, mean, dim, VecSimMetric_L2, y_query_l2); + + float bf_ip = bruteForceIPDist(x, y, dim); + float bf_l2 = bruteForceL2Dist(x, y, dim); + + float got_ip_asym = calc_ip->calcDistanceForQuery(x_blob, y_query_ip, dim); + float got_ip_sym = calc_ip->calcDistance(x_blob, y_blob, dim); + float got_l2_asym = calc_l2->calcDistanceForQuery(x_blob_l2, y_query_l2, dim); + float got_l2_sym = calc_l2->calcDistance(x_blob_l2, y_blob_l2, dim); + + if (std::abs(got_ip_asym - bf_ip) > 0.1f) + ++failures; + if (std::abs(got_ip_sym - bf_ip) > 0.1f) + ++failures; + if (std::abs(got_l2_asym - bf_l2) > 0.1f) + ++failures; + if (std::abs(got_l2_sym - bf_l2) > 0.1f) + ++failures; + + allocator->free_allocation(x_blob); + allocator->free_allocation(y_blob); + allocator->free_allocation(y_query_ip); + allocator->free_allocation(x_blob_l2); + allocator->free_allocation(y_blob_l2); + allocator->free_allocation(y_query_l2); + } + + EXPECT_EQ(failures, 0) << failures << " distance computations exceeded tolerance"; + + delete calc_ip; + delete calc_l2; +}