Skip to content

[MOD-17102] Add cosine-family support for COSINE_SIMILARITY #997

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rbs333 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[MOD-17102] Add cosine-family support for COSINE_SIMILARITY #997

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 0 additions & 5 deletions src/VecSim/algorithms/svs/svs_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,11 +66,6 @@ inline double toVecSimDistance<svs::distance::DistanceIP>(float v) {
return 1.0 - static_cast<double>(v);
}

template <>
inline double toVecSimDistance<svs::distance::DistanceCosineSimilarity>(float v) {
return 1.0 - static_cast<double>(v);
}

// VecSim allocator wrapper for SVS containers
template <typename T>
using SVSAllocator = VecsimSTLAllocator<T>;
Expand Down
8 changes: 4 additions & 4 deletions src/VecSim/index_factories/components/preprocessors_factory.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,7 +44,7 @@ CreatePreprocessorsContainerParams(VecSimMetric metric, size_t dim, bool is_norm
size_t processed_bytes_count = dim * sizeof(DataType);

VecSimMetric pp_metric = metric;
if (metric == VecSimMetric_Cosine) {
if (VecSimMetric_IsCosineFamily(metric)) {
// if metric is cosine and DataType is integral, the processed_bytes_count includes the
// norm appended to the vector.
if (std::is_integral<DataType>::value) {
Expand Down Expand Up @@ -78,7 +78,7 @@ PreprocessorsContainerAbstract *
CreatePreprocessorsContainer(std::shared_ptr<VecSimAllocator> allocator,
PreprocessorsContainerParams params) {

if (params.metric == VecSimMetric_Cosine) {
if (VecSimMetric_IsCosineFamily(params.metric)) {
auto multiPPContainer = new (allocator) MultiPreprocessorsContainer<DataType, 1>(
allocator, params.query_alignment, params.storage_alignment);
auto cosine_preprocessor = new (allocator)
Expand Down Expand Up @@ -118,13 +118,13 @@ template <typename DataType>
size_t EstimatePreprocessorsContainerMemory(VecSimMetric metric, bool is_normalized = false) {
size_t allocations_overhead = VecSimAllocator::getAllocationOverheadSize();
VecSimMetric pp_metric;
if (is_normalized && metric == VecSimMetric_Cosine) {
if (is_normalized && VecSimMetric_IsCosineFamily(metric)) {
pp_metric = VecSimMetric_IP;
} else {
pp_metric = metric;
}

if (pp_metric == VecSimMetric_Cosine) {
if (VecSimMetric_IsCosineFamily(pp_metric)) {
constexpr size_t n_preprocessors = 1;
// One entry in preprocessors array
size_t est =
Expand Down
3 changes: 2 additions & 1 deletion src/VecSim/index_factories/svs_factory.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ VecSimIndex *NewIndexImpl(const VecSimParams *params, bool is_normalized) {
abstractInitParams.allocator, svsParams.metric, svsParams.dim, is_normalized, 0);
IndexComponents<svs_details::vecsim_dt<DataType>, float> components = {
nullptr, preprocessors}; // calculator is not in use in svs.
bool forcePreprocessing = !is_normalized && svsParams.metric == VecSimMetric_Cosine;
bool forcePreprocessing = !is_normalized && VecSimMetric_IsCosineFamily(svsParams.metric);
if (svsParams.multi) {
return new (abstractInitParams.allocator)
SVSIndex<MetricType, DataType, true, QuantBits, ResidualBits, IsLeanVec>(
Expand Down Expand Up @@ -113,6 +113,7 @@ VecSimIndex *NewIndexImpl(const VecSimParams *params, bool is_normalized) {
return NewIndexImpl<svs::distance::DistanceL2>(params, is_normalized);
case VecSimMetric_IP:
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
return NewIndexImpl<svs::distance::DistanceIP>(params, is_normalized);
default:
// If we got here something is wrong.
Expand Down
5 changes: 3 additions & 2 deletions src/VecSim/spaces/computer/preprocessors.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,8 +242,9 @@ class QuantPreprocessor : public PreprocessorInterface {
using sq8 = vecsim_types::sq8;

static_assert(Metric == VecSimMetric_L2 || Metric == VecSimMetric_IP ||
Metric == VecSimMetric_Cosine,
"QuantPreprocessor only supports L2, IP and Cosine metrics");
Metric == VecSimMetric_Cosine ||
Metric == VecSimMetric_CosineSimilarity,
"QuantPreprocessor only supports L2, IP and cosine-based metrics");

// Helper function to perform quantization. This function is used by the storage preprocessing
// methods.
Expand Down
9 changes: 9 additions & 0 deletions src/VecSim/spaces/spaces.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ dist_func_t<float> GetDistFunc<vecsim_types::bfloat16, float>(VecSimMetric metri
unsigned char *alignment) {
switch (metric) {
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
case VecSimMetric_IP:
return IP_BF16_GetDistFunc(dim, alignment);
case VecSimMetric_L2:
Expand All @@ -39,6 +40,7 @@ dist_func_t<float> GetDistFunc<vecsim_types::float16, float>(VecSimMetric metric
unsigned char *alignment) {
switch (metric) {
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
case VecSimMetric_IP:
return IP_FP16_GetDistFunc(dim, alignment);
case VecSimMetric_L2:
Expand All @@ -52,6 +54,7 @@ dist_func_t<float> GetDistFunc<float, float>(VecSimMetric metric, size_t dim,
unsigned char *alignment) {
switch (metric) {
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
case VecSimMetric_IP:
return IP_FP32_GetDistFunc(dim, alignment);
case VecSimMetric_L2:
Expand All @@ -65,6 +68,7 @@ dist_func_t<double> GetDistFunc<double, double>(VecSimMetric metric, size_t dim,
unsigned char *alignment) {
switch (metric) {
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
case VecSimMetric_IP:
return IP_FP64_GetDistFunc(dim, alignment);
case VecSimMetric_L2:
Expand All @@ -78,6 +82,7 @@ dist_func_t<float> GetDistFunc<int8_t, float>(VecSimMetric metric, size_t dim,
unsigned char *alignment) {
switch (metric) {
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
return Cosine_INT8_GetDistFunc(dim, alignment);
case VecSimMetric_IP:
return IP_INT8_GetDistFunc(dim, alignment);
Expand All @@ -92,6 +97,7 @@ dist_func_t<float> GetDistFunc<uint8_t, float>(VecSimMetric metric, size_t dim,
unsigned char *alignment) {
switch (metric) {
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
return Cosine_UINT8_GetDistFunc(dim, alignment);
case VecSimMetric_IP:
return IP_UINT8_GetDistFunc(dim, alignment);
Expand All @@ -106,6 +112,7 @@ dist_func_t<float> GetDistFunc<vecsim_types::sq8, float>(VecSimMetric metric, si
unsigned char *alignment) {
switch (metric) {
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
return Cosine_SQ8_SQ8_GetDistFunc(dim, alignment);
case VecSimMetric_IP:
return IP_SQ8_SQ8_GetDistFunc(dim, alignment);
Expand All @@ -120,6 +127,7 @@ dist_func_t<float> GetDistFunc<vecsim_types::sq8, float, float>(VecSimMetric met
unsigned char *alignment) {
switch (metric) {
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
return Cosine_SQ8_FP32_GetDistFunc(dim, alignment);
case VecSimMetric_IP:
return IP_SQ8_FP32_GetDistFunc(dim, alignment);
Expand All @@ -135,6 +143,7 @@ GetDistFunc<vecsim_types::sq8, float, vecsim_types::float16>(VecSimMetric metric
unsigned char *alignment) {
switch (metric) {
case VecSimMetric_Cosine:
case VecSimMetric_CosineSimilarity:
return Cosine_SQ8_FP16_GetDistFunc(dim, alignment);
case VecSimMetric_IP:
return IP_SQ8_FP16_GetDistFunc(dim, alignment);
Expand Down
5 changes: 4 additions & 1 deletion src/VecSim/utils/vec_utils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,6 +223,8 @@ const char *VecSimMetric_ToString(VecSimMetric vecsimMetric) {
switch (vecsimMetric) {
case VecSimMetric_Cosine:
return "COSINE";
case VecSimMetric_CosineSimilarity:
return "COSINE_SIMILARITY";
case VecSimMetric_IP:
return "IP";
case VecSimMetric_L2:
Expand Down Expand Up @@ -295,7 +297,8 @@ size_t VecSimType_sizeof(VecSimType type) {

size_t VecSimParams_GetStoredDataSize(VecSimType type, size_t dim, VecSimMetric metric) {
size_t storedDataSize = VecSimType_sizeof(type) * dim;
if (metric == VecSimMetric_Cosine && (type == VecSimType_INT8 || type == VecSimType_UINT8)) {
if (VecSimMetric_IsCosineFamily(metric) &&
(type == VecSimType_INT8 || type == VecSimType_UINT8)) {
storedDataSize += sizeof(float); // For the norm
}
return storedDataSize;
Expand Down
2 changes: 1 addition & 1 deletion src/VecSim/vec_sim.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,7 +259,7 @@ extern "C" size_t VecSimParams_GetQueryBlobSize(VecSimType type, size_t dim, Vec
type == VecSimType_BFLOAT16 || type == VecSimType_FLOAT16 || type == VecSimType_INT8 ||
type == VecSimType_UINT8);
size_t blobSize = VecSimType_sizeof(type) * dim;
if (metric == VecSimMetric_Cosine && (type == VecSimType_INT8 || type == VecSimType_UINT8)) {
if (VecSimMetric_IsCosineFamily(metric) && (type == VecSimType_INT8 || type == VecSimType_UINT8)) {
blobSize += sizeof(float); // For the norm
}
return blobSize;
Expand Down
13 changes: 12 additions & 1 deletion src/VecSim/vec_sim_common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,18 @@ typedef enum {
} VecSimBool;

// Distance metric
typedef enum { VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine } VecSimMetric;
typedef enum {
VecSimMetric_L2,
VecSimMetric_IP,
VecSimMetric_Cosine,
VecSimMetric_CosineSimilarity
} VecSimMetric;

// Returns true for any metric whose internal execution path uses cosine distance,
// regardless of how scores are presented at the API boundary.
static inline bool VecSimMetric_IsCosineFamily(VecSimMetric metric) {
return metric == VecSimMetric_Cosine || metric == VecSimMetric_CosineSimilarity;
}

typedef size_t labelType;
typedef unsigned int idType;
Expand Down
1 change: 1 addition & 0 deletions src/python_bindings/bindings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -679,6 +679,7 @@ PYBIND11_MODULE(VecSim, m) {
.value("VecSimMetric_L2", VecSimMetric_L2)
.value("VecSimMetric_IP", VecSimMetric_IP)
.value("VecSimMetric_Cosine", VecSimMetric_Cosine)
.value("VecSimMetric_CosineSimilarity", VecSimMetric_CosineSimilarity)
.export_values();

py::enum_<VecSimOptionMode>(m, "VecSimOptionMode")
Expand Down
12 changes: 10 additions & 2 deletions tests/unit/test_components.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1313,11 +1313,15 @@ TEST_P(QuantPreprocessorMetricTest, QuantizationBlobSizeAndMetadata) {
case VecSimMetric_Cosine:
runQuantizationTest<VecSimMetric_Cosine>();
break;
case VecSimMetric_CosineSimilarity:
runQuantizationTest<VecSimMetric_CosineSimilarity>();
break;
}
}

INSTANTIATE_TEST_SUITE_P(QuantPreprocessorTests, QuantPreprocessorMetricTest,
testing::Values(VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine),
testing::Values(VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine,
VecSimMetric_CosineSimilarity),
[](const testing::TestParamInfo<VecSimMetric> &info) {
return VecSimMetric_ToString(info.param);
});
Expand Down Expand Up @@ -1477,11 +1481,15 @@ TEST_P(QuantPreprocessorFP16MetricTest, QuantizationBlobSizeAndMetadata) {
case VecSimMetric_Cosine:
runQuantizationTest<VecSimMetric_Cosine>();
break;
case VecSimMetric_CosineSimilarity:
runQuantizationTest<VecSimMetric_CosineSimilarity>();
break;
}
}

INSTANTIATE_TEST_SUITE_P(QuantPreprocessorFP16Tests, QuantPreprocessorFP16MetricTest,
testing::Values(VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine),
testing::Values(VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine,
VecSimMetric_CosineSimilarity),
[](const testing::TestParamInfo<VecSimMetric> &info) {
return VecSimMetric_ToString(info.param);
});
Expand Down
25 changes: 16 additions & 9 deletions tests/unit/test_spaces.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -492,49 +492,49 @@ TEST_F(SpacesTest, SQ8_FP16_l2sqr_odd_dim_unaligned_metadata_test) {

TEST_F(SpacesTest, GetDistFuncInvalidMetricFP32) {
EXPECT_THROW(
(spaces::GetDistFunc<float, float>((VecSimMetric)(VecSimMetric_Cosine + 1), 10, nullptr)),
(spaces::GetDistFunc<float, float>((VecSimMetric)(VecSimMetric_CosineSimilarity + 1), 10, nullptr)),
std::invalid_argument);
}
TEST_F(SpacesTest, GetDistFuncInvalidMetricFP64) {
EXPECT_THROW(
(spaces::GetDistFunc<double, double>((VecSimMetric)(VecSimMetric_Cosine + 1), 10, nullptr)),
(spaces::GetDistFunc<double, double>((VecSimMetric)(VecSimMetric_CosineSimilarity + 1), 10, nullptr)),
std::invalid_argument);
}
TEST_F(SpacesTest, GetDistFuncInvalidMetricBF16) {
EXPECT_THROW((spaces::GetDistFunc<bfloat16, float>((VecSimMetric)(VecSimMetric_Cosine + 1), 10,
EXPECT_THROW((spaces::GetDistFunc<bfloat16, float>((VecSimMetric)(VecSimMetric_CosineSimilarity + 1), 10,
nullptr)),
std::invalid_argument);
}
TEST_F(SpacesTest, GetDistFuncInvalidMetricFP16) {
EXPECT_THROW(
(spaces::GetDistFunc<float16, float>((VecSimMetric)(VecSimMetric_Cosine + 1), 10, nullptr)),
(spaces::GetDistFunc<float16, float>((VecSimMetric)(VecSimMetric_CosineSimilarity + 1), 10, nullptr)),
std::invalid_argument);
}
TEST_F(SpacesTest, GetDistFuncInvalidMetricINT8) {
EXPECT_THROW(
(spaces::GetDistFunc<int8_t, float>((VecSimMetric)(VecSimMetric_Cosine + 1), 10, nullptr)),
(spaces::GetDistFunc<int8_t, float>((VecSimMetric)(VecSimMetric_CosineSimilarity + 1), 10, nullptr)),
std::invalid_argument);
}
TEST_F(SpacesTest, GetDistFuncInvalidMetricUINT8) {
EXPECT_THROW(
(spaces::GetDistFunc<uint8_t, float>((VecSimMetric)(VecSimMetric_Cosine + 1), 10, nullptr)),
(spaces::GetDistFunc<uint8_t, float>((VecSimMetric)(VecSimMetric_CosineSimilarity + 1), 10, nullptr)),
std::invalid_argument);
}
TEST_F(SpacesTest, GetDistFuncInvalidMetricSQ8) {
// SQ8 to SQ8 (symmetric)
EXPECT_THROW(
(spaces::GetDistFunc<sq8, float>((VecSimMetric)(VecSimMetric_Cosine + 1), 10, nullptr)),
(spaces::GetDistFunc<sq8, float>((VecSimMetric)(VecSimMetric_CosineSimilarity + 1), 10, nullptr)),
std::invalid_argument);
}
TEST_F(SpacesTest, GetDistFuncInvalidMetricSQ8ToFloat) {
// SQ8 to float (asymmetric)
EXPECT_THROW((spaces::GetDistFunc<sq8, float, float>((VecSimMetric)(VecSimMetric_Cosine + 1),
EXPECT_THROW((spaces::GetDistFunc<sq8, float, float>((VecSimMetric)(VecSimMetric_CosineSimilarity + 1),
10, nullptr)),
std::invalid_argument);
}
TEST_F(SpacesTest, GetDistFuncInvalidMetricSQ8ToFP16) {
// SQ8 storage with FP16 query (asymmetric)
EXPECT_THROW((spaces::GetDistFunc<sq8, float, float16>((VecSimMetric)(VecSimMetric_Cosine + 1),
EXPECT_THROW((spaces::GetDistFunc<sq8, float, float16>((VecSimMetric)(VecSimMetric_CosineSimilarity + 1),
10, nullptr)),
std::invalid_argument);
}
Expand All @@ -546,9 +546,11 @@ TEST_F(SpacesTest, GetDistFuncSQ8Symmetric) {
auto l2_func = spaces::GetDistFunc<sq8, float>(VecSimMetric_L2, dim, nullptr);
auto ip_func = spaces::GetDistFunc<sq8, float>(VecSimMetric_IP, dim, nullptr);
auto cosine_func = spaces::GetDistFunc<sq8, float>(VecSimMetric_Cosine, dim, nullptr);
auto cosine_similarity_func = spaces::GetDistFunc<sq8, float>(VecSimMetric_CosineSimilarity, 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));
ASSERT_EQ(cosine_similarity_func, Cosine_SQ8_SQ8_GetDistFunc(dim, nullptr));
}

TEST_F(SpacesTest, GetDistFuncSQ8Asymmetric) {
Expand All @@ -557,9 +559,11 @@ TEST_F(SpacesTest, GetDistFuncSQ8Asymmetric) {
auto l2_func = spaces::GetDistFunc<sq8, float, float>(VecSimMetric_L2, dim, nullptr);
auto ip_func = spaces::GetDistFunc<sq8, float, float>(VecSimMetric_IP, dim, nullptr);
auto cosine_func = spaces::GetDistFunc<sq8, float, float>(VecSimMetric_Cosine, dim, nullptr);
auto cosine_similarity_func = spaces::GetDistFunc<sq8, float, float>(VecSimMetric_CosineSimilarity, dim, nullptr);
ASSERT_EQ(l2_func, L2_SQ8_FP32_GetDistFunc(dim, nullptr));
ASSERT_EQ(ip_func, IP_SQ8_FP32_GetDistFunc(dim, nullptr));
ASSERT_EQ(cosine_func, Cosine_SQ8_FP32_GetDistFunc(dim, nullptr));
ASSERT_EQ(cosine_similarity_func, Cosine_SQ8_FP32_GetDistFunc(dim, nullptr));
}

TEST_F(SpacesTest, GetDistFuncSQ8FP16Asymmetric) {
Expand All @@ -569,9 +573,12 @@ TEST_F(SpacesTest, GetDistFuncSQ8FP16Asymmetric) {
auto l2_func = spaces::GetDistFunc<sq8, float, float16>(VecSimMetric_L2, dim, nullptr);
auto ip_func = spaces::GetDistFunc<sq8, float, float16>(VecSimMetric_IP, dim, nullptr);
auto cosine_func = spaces::GetDistFunc<sq8, float, float16>(VecSimMetric_Cosine, dim, nullptr);
auto cosine_similarity_func =
spaces::GetDistFunc<sq8, float, float16>(VecSimMetric_CosineSimilarity, dim, nullptr);
ASSERT_EQ(l2_func, L2_SQ8_FP16_GetDistFunc(dim, nullptr));
ASSERT_EQ(ip_func, IP_SQ8_FP16_GetDistFunc(dim, nullptr));
ASSERT_EQ(cosine_func, Cosine_SQ8_FP16_GetDistFunc(dim, nullptr));
ASSERT_EQ(cosine_similarity_func, Cosine_SQ8_FP16_GetDistFunc(dim, nullptr));

// dim < 16 takes the scalar early-return in every SQ8_FP16 dispatcher (no SIMD tier).
size_t small_dim = 8;
Expand Down
Loading