diff --git a/src/VecSim/spaces/spaces.cpp b/src/VecSim/spaces/spaces.cpp index dbb18943a..42569ba80 100644 --- a/src/VecSim/spaces/spaces.cpp +++ b/src/VecSim/spaces/spaces.cpp @@ -8,6 +8,7 @@ */ #include "VecSim/types/bfloat16.h" #include "VecSim/types/float16.h" +#include "VecSim/types/sq8.h" #include "VecSim/spaces/space_includes.h" #include "VecSim/spaces/spaces.h" #include "VecSim/spaces/IP_space.h" @@ -100,6 +101,34 @@ dist_func_t GetDistFunc(VecSimMetric metric, size_t dim, throw std::invalid_argument("Invalid metric"); } +template <> +dist_func_t GetDistFunc(VecSimMetric metric, size_t dim, + unsigned char *alignment) { + switch (metric) { + case VecSimMetric_Cosine: + return Cosine_SQ8_SQ8_GetDistFunc(dim, alignment); + case VecSimMetric_IP: + return IP_SQ8_SQ8_GetDistFunc(dim, alignment); + case VecSimMetric_L2: + return L2_SQ8_SQ8_GetDistFunc(dim, alignment); + } + throw std::invalid_argument("Invalid metric"); +} + +template <> +dist_func_t GetDistFunc(VecSimMetric metric, size_t dim, + unsigned char *alignment) { + switch (metric) { + case VecSimMetric_Cosine: + return Cosine_SQ8_GetDistFunc(dim, alignment); + case VecSimMetric_IP: + return IP_SQ8_GetDistFunc(dim, alignment); + case VecSimMetric_L2: + return L2_SQ8_GetDistFunc(dim, alignment); + } + throw std::invalid_argument("Invalid metric"); +} + template <> normalizeVector_f GetNormalizeFunc(void) { return normalizeVector_imp; diff --git a/src/VecSim/spaces/spaces.h b/src/VecSim/spaces/spaces.h index 7d3c76282..982d3f749 100644 --- a/src/VecSim/spaces/spaces.h +++ b/src/VecSim/spaces/spaces.h @@ -16,10 +16,11 @@ namespace spaces { template using dist_func_t = RET_TYPE (*)(const void *, const void *, size_t); -// Set the distance function for a given data type, metric and dimension. The alignment hint is -// determined according to the chosen implementation and available optimizations. - -template +// Get the distance function for comparing vectors of type VecType1 and VecType2, for a given metric +// and dimension. The returned function has the signature: dist(VecType1*, VecType2*, size_t) -> +// DistType. VecType2 defaults to VecType1 when both vectors are of the same type. The alignment +// hint is set based on the chosen implementation and available optimizations. +template dist_func_t GetDistFunc(VecSimMetric metric, size_t dim, unsigned char *alignment); template diff --git a/src/VecSim/types/sq8.h b/src/VecSim/types/sq8.h new file mode 100644 index 000000000..a2bf738e0 --- /dev/null +++ b/src/VecSim/types/sq8.h @@ -0,0 +1,22 @@ +/* + * Copyright (c) 2006-Present, Redis Ltd. + * All rights reserved. + * + * Licensed under your choice of the Redis Source Available License 2.0 + * (RSALv2); or (b) the Server Side Public License v1 (SSPLv1); or (c) the + * GNU Affero General Public License v3 (AGPLv3). + */ +#pragma once + +#include +#include +#include "VecSim/vec_sim_common.h" + +namespace vecsim_types { + +// Represents a scalar-quantized 8-bit blob with reconstruction metadata +struct sq8 { + using value_type = uint8_t; +}; + +} // namespace vecsim_types diff --git a/tests/unit/test_spaces.cpp b/tests/unit/test_spaces.cpp index 6e6e92d69..dcb2171d6 100644 --- a/tests/unit/test_spaces.cpp +++ b/tests/unit/test_spaces.cpp @@ -20,6 +20,7 @@ #include "VecSim/spaces/IP_space.h" #include "VecSim/spaces/L2_space.h" #include "VecSim/types/float16.h" +#include "VecSim/types/sq8.h" #include "VecSim/spaces/functions/AVX512F.h" #include "VecSim/spaces/functions/AVX.h" #include "VecSim/spaces/functions/SSE.h" @@ -43,6 +44,7 @@ using bfloat16 = vecsim_types::bfloat16; using float16 = vecsim_types::float16; +using sq8 = vecsim_types::sq8; using namespace spaces; class SpacesTest : public ::testing::Test { @@ -509,6 +511,42 @@ TEST_F(SpacesTest, GetDistFuncInvalidMetricUINT8) { (spaces::GetDistFunc((VecSimMetric)(VecSimMetric_Cosine + 1), 10, nullptr)), std::invalid_argument); } +TEST_F(SpacesTest, GetDistFuncInvalidMetricSQ8) { + // SQ8 to SQ8 (symmetric) + EXPECT_THROW( + (spaces::GetDistFunc((VecSimMetric)(VecSimMetric_Cosine + 1), 10, nullptr)), + std::invalid_argument); +} +TEST_F(SpacesTest, GetDistFuncInvalidMetricSQ8ToFloat) { + // SQ8 to float (asymmetric) + EXPECT_THROW((spaces::GetDistFunc((VecSimMetric)(VecSimMetric_Cosine + 1), + 10, nullptr)), + std::invalid_argument); +} + +// Positive tests for GetDistFunc - verify correct function is returned +TEST_F(SpacesTest, GetDistFuncSQ8Symmetric) { + // SQ8 to SQ8 (symmetric) - should return SQ8_SQ8 functions + size_t dim = 128; + auto l2_func = spaces::GetDistFunc(VecSimMetric_L2, dim, nullptr); + auto ip_func = spaces::GetDistFunc(VecSimMetric_IP, dim, nullptr); + auto cosine_func = spaces::GetDistFunc(VecSimMetric_Cosine, dim, nullptr); + ASSERT_EQ(l2_func, L2_SQ8_SQ8_GetDistFunc(dim, nullptr)); + ASSERT_EQ(ip_func, IP_SQ8_SQ8_GetDistFunc(dim, nullptr)); + ASSERT_EQ(cosine_func, Cosine_SQ8_SQ8_GetDistFunc(dim, nullptr)); +} + +TEST_F(SpacesTest, GetDistFuncSQ8Asymmetric) { + // SQ8 to float (asymmetric) - should return SQ8 functions + size_t dim = 128; + auto l2_func = spaces::GetDistFunc(VecSimMetric_L2, dim, nullptr); + auto ip_func = spaces::GetDistFunc(VecSimMetric_IP, dim, nullptr); + auto cosine_func = spaces::GetDistFunc(VecSimMetric_Cosine, dim, nullptr); + ASSERT_EQ(l2_func, L2_SQ8_GetDistFunc(dim, nullptr)); + ASSERT_EQ(ip_func, IP_SQ8_GetDistFunc(dim, nullptr)); + ASSERT_EQ(cosine_func, Cosine_SQ8_GetDistFunc(dim, nullptr)); +} + #ifdef CPU_FEATURES_ARCH_X86_64 TEST_F(SpacesTest, smallDimChooser) { // Verify that small dimensions gets the no optimization function.