/* * 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). */ #include #include "gtest/gtest.h" #include "VecSim/vec_sim.h" #include "VecSim/spaces/computer/preprocessor_container.h" #include "VecSim/spaces/computer/preprocessors.h" #include "VecSim/spaces/computer/calculator.h" #include "unit_test_utils.h" #include "tests_utils.h" class IndexCalculatorTest : public ::testing::Test {}; namespace dummyCalcultor { using DummyType = int; using dummy_dist_func_t = DummyType (*)(int); int dummyDistFunc(int value) { return value; } template class DistanceCalculatorDummy : public DistanceCalculatorInterface { public: DistanceCalculatorDummy(std::shared_ptr allocator, dummy_dist_func_t dist_func) : DistanceCalculatorInterface(allocator, dist_func) {} virtual DistType calcDistance(const void *v1, const void *v2, size_t dim) const { return this->dist_func(7); } // Dummy uses a non-standard dist func signature, so the standard slot is unavailable. spaces::dist_func_t getDistFunc() const override { return nullptr; } }; } // namespace dummyCalcultor TEST(IndexCalculatorTest, TestIndexCalculator) { std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); // Test computer with a distance function signature different from dim(v1, v2, dim()). using namespace dummyCalcultor; auto distance_calculator = DistanceCalculatorDummy(allocator, dummyDistFunc); ASSERT_EQ(distance_calculator.calcDistance(nullptr, nullptr, 0), 7); } class PreprocessorsTest : public ::testing::Test {}; namespace dummyPreprocessors { using DummyType = int; enum pp_mode { STORAGE_ONLY, QUERY_ONLY, BOTH, EMPTY }; // Dummy storage preprocessor template class DummyStoragePreprocessor : public PreprocessorInterface { public: DummyStoragePreprocessor(std::shared_ptr allocator, DataType value_to_add_storage, DataType value_to_add_query = 0) : PreprocessorInterface(allocator), value_to_add_storage(value_to_add_storage), value_to_add_query(value_to_add_query) { if (!value_to_add_query) value_to_add_query = value_to_add_storage; } void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, size_t &storage_blob_size, size_t &query_blob_size, unsigned char storage_alignment, unsigned char query_alignment) const override { assert(storage_blob_size == query_blob_size); this->preprocessForStorage(original_blob, storage_blob, storage_blob_size, storage_alignment); } void preprocessForStorage(const void *original_blob, void *&blob, size_t &input_blob_size, unsigned char storage_alignment) const override { // If the blob was not allocated yet, allocate it. if (blob == nullptr) { blob = this->allocator->allocate(input_blob_size); memcpy(blob, original_blob, input_blob_size); } static_cast(blob)[0] += value_to_add_storage; } void preprocessStorageInPlace(void *blob, size_t input_blob_size) const override { assert(blob); static_cast(blob)[0] += value_to_add_storage; } void preprocessQuery(const void *original_blob, void *&blob, size_t &input_blob_size, unsigned char alignment) const override { /* do nothing*/ } private: DataType value_to_add_storage; DataType value_to_add_query; }; // Dummy query preprocessor template class DummyQueryPreprocessor : public PreprocessorInterface { public: DummyQueryPreprocessor(std::shared_ptr allocator, DataType value_to_add_storage, DataType _value_to_add_query = 0) : PreprocessorInterface(allocator), value_to_add_storage(value_to_add_storage), value_to_add_query(_value_to_add_query) { if (!_value_to_add_query) value_to_add_query = value_to_add_storage; } void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, size_t &storage_blob_size, size_t &query_blob_size, unsigned char storage_alignment, unsigned char query_alignment) const override { assert(storage_blob_size == query_blob_size); this->preprocessQuery(original_blob, query_blob, query_blob_size, query_alignment); } void preprocessForStorage(const void *original_blob, void *&blob, size_t &input_blob_size, unsigned char storage_alignment) const override { /* do nothing*/ } void preprocessStorageInPlace(void *blob, size_t input_blob_size) const override {} void preprocessQuery(const void *original_blob, void *&blob, size_t &input_blob_size, unsigned char alignment) const override { // If the blob was not allocated yet, allocate it. if (blob == nullptr) { blob = this->allocator->allocate_aligned(input_blob_size, alignment); memcpy(blob, original_blob, input_blob_size); } static_cast(blob)[0] += value_to_add_query; } private: DataType value_to_add_storage; DataType value_to_add_query; }; // Dummy mixed preprocessor (precesses the blobs differently) template class DummyMixedPreprocessor : public PreprocessorInterface { public: DummyMixedPreprocessor(std::shared_ptr allocator, DataType value_to_add_storage, DataType value_to_add_query) : PreprocessorInterface(allocator), value_to_add_storage(value_to_add_storage), value_to_add_query(value_to_add_query) {} void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, size_t &storage_blob_size, size_t &query_blob_size, unsigned char storage_alignment, unsigned char query_alignment) const override { assert(storage_blob_size == query_blob_size); // One blob was already allocated by a previous preprocessor(s) that process both blobs the // same. The blobs are pointing to the same memory, we need to allocate another memory slot // to split them. if ((storage_blob == query_blob) && (query_blob != nullptr)) { storage_blob = this->allocator->allocate(storage_blob_size); memcpy(storage_blob, query_blob, storage_blob_size); } // Either both are nullptr or they are pointing to different memory slots. Both cases are // handled by the designated functions. this->preprocessForStorage(original_blob, storage_blob, storage_blob_size, storage_alignment); this->preprocessQuery(original_blob, query_blob, query_blob_size, query_alignment); } void preprocessForStorage(const void *original_blob, void *&blob, size_t &input_blob_size, unsigned char storage_alignment) const override { // If the blob was not allocated yet, allocate it. if (blob == nullptr) { blob = this->allocator->allocate(input_blob_size); memcpy(blob, original_blob, input_blob_size); } static_cast(blob)[0] += value_to_add_storage; } void preprocessStorageInPlace(void *blob, size_t input_blob_size) const override {} void preprocessQuery(const void *original_blob, void *&blob, size_t &input_blob_size, unsigned char alignment) const override { // If the blob was not allocated yet, allocate it. if (blob == nullptr) { blob = this->allocator->allocate_aligned(input_blob_size, alignment); memcpy(blob, original_blob, input_blob_size); } static_cast(blob)[0] += value_to_add_query; } private: DataType value_to_add_storage; DataType value_to_add_query; }; // A preprocessor that changes the allocation size of the blobs in the same manner. // set excess bytes to (char)2 template class DummyChangeAllocSizePreprocessor : public PreprocessorInterface { private: size_t processed_bytes_count; static constexpr unsigned char excess_value = 2; public: DummyChangeAllocSizePreprocessor(std::shared_ptr allocator, size_t processed_bytes_count) : PreprocessorInterface(allocator), processed_bytes_count(processed_bytes_count) {} static constexpr unsigned char getExcessValue() { return excess_value; } void preprocess(const void *original_blob, void *&storage_blob, void *&query_blob, size_t &storage_blob_size, size_t &query_blob_size, unsigned char storage_alignment, unsigned char query_alignment) const override { // if the blobs are equal, allocate a single shared buffer aligned to satisfy both hints. if (storage_blob == query_blob) { assert(storage_blob_size == query_blob_size); const unsigned char shared_alignment = spaces::combineAlignments(storage_alignment, query_alignment); preprocessGeneral(original_blob, storage_blob, storage_blob_size, shared_alignment); query_blob = storage_blob; query_blob_size = storage_blob_size; return; } // The blobs are not equal auto alloc_and_process = [&](void *&blob, size_t &input_blob_size, unsigned char alignment) { // If the input blob size is not enough if (input_blob_size < processed_bytes_count) { auto new_blob = this->allocator->allocate_aligned(processed_bytes_count, alignment); if (blob == nullptr) { memcpy(new_blob, original_blob, input_blob_size); } else { // copy the blob to the new blob, and free the old one memcpy(new_blob, blob, input_blob_size); this->allocator->free_allocation(blob); } blob = new_blob; memset((char *)blob + input_blob_size, excess_value, processed_bytes_count - input_blob_size); } else { if (blob == nullptr) { blob = this->allocator->allocate_aligned(processed_bytes_count, alignment); memcpy(blob, original_blob, processed_bytes_count); } else { memset((char *)blob + processed_bytes_count, excess_value, input_blob_size - processed_bytes_count); } } input_blob_size = processed_bytes_count; }; alloc_and_process(storage_blob, storage_blob_size, storage_alignment); alloc_and_process(query_blob, query_blob_size, query_alignment); } void preprocessForStorage(const void *original_blob, void *&blob, size_t &input_blob_size, unsigned char storage_alignment) const override { this->preprocessGeneral(original_blob, blob, input_blob_size, storage_alignment); } void preprocessStorageInPlace(void *blob, size_t input_blob_size) const override { // only supported if the blob in already large enough assert(input_blob_size >= processed_bytes_count); // set excess bytes to 0 memset((char *)blob + processed_bytes_count, excess_value, input_blob_size - processed_bytes_count); } void preprocessQuery(const void *original_blob, void *&blob, size_t &input_blob_size, unsigned char alignment) const override { this->preprocessGeneral(original_blob, blob, input_blob_size, alignment); } private: void preprocessGeneral(const void *original_blob, void *&blob, size_t &input_blob_size, unsigned char alignment = 0) const { // if the size of the input is not enough. if (input_blob_size < processed_bytes_count) { // calloc doesn't have an alignment version, so we need to allocate aligned memory and // cset the excess bytes to 0. auto new_blob = this->allocator->allocate_aligned(processed_bytes_count, alignment); if (blob == nullptr) { // copy thr original blob memcpy(new_blob, original_blob, input_blob_size); } else { // copy the blob to the new blob memcpy(new_blob, blob, input_blob_size); this->allocator->free_allocation(blob); } blob = new_blob; // set excess bytes to 0 memset((char *)blob + input_blob_size, excess_value, processed_bytes_count - input_blob_size); } else { // input size is larger than output if (blob == nullptr) { blob = this->allocator->allocate_aligned(processed_bytes_count, alignment); memcpy(blob, original_blob, processed_bytes_count); } else { // set excess bytes to 0 memset((char *)blob + processed_bytes_count, excess_value, input_blob_size - processed_bytes_count); } } // update the input blob size input_blob_size = processed_bytes_count; } }; } // namespace dummyPreprocessors TEST(PreprocessorsTest, PreprocessorsTestBasicAlignmentTest) { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); unsigned char alignment = 8; auto preprocessor = PreprocessorsContainerAbstract(allocator, alignment); const int original_blob[4] = {1, 1, 1, 1}; size_t processed_bytes_count = sizeof(original_blob); { auto aligned_query = preprocessor.preprocessQuery(original_blob, processed_bytes_count); unsigned char address_alignment = (uintptr_t)(aligned_query.get()) % alignment; ASSERT_EQ(address_alignment, 0); } } template void MultiPPContainerEmpty() { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t dim = 4; const int original_blob[dim] = {1, 2, 3, 4}; const int original_blob_cpy[dim] = {1, 2, 3, 4}; constexpr size_t n_preprocessors = 3; auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, sizeof(original_blob)); // Original blob should not be changed CompareVectors(original_blob, original_blob_cpy, dim); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // Storage blob should not be reallocated or changed ASSERT_EQ(storage_blob, (const int *)original_blob); CompareVectors(original_blob, (const int *)storage_blob, dim); // query blob *values* should not be changed CompareVectors(original_blob, (const int *)query_blob, dim); // If alignment is set the query blob address should be aligned to the specified alignment. if constexpr (alignment) { unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; ASSERT_EQ(address_alignment, 0); } } } TEST(PreprocessorsTest, MultiPPContainerEmptyNoAlignment) { using namespace dummyPreprocessors; MultiPPContainerEmpty<0>(); } TEST(PreprocessorsTest, MultiPPContainerEmptyAlignment) { using namespace dummyPreprocessors; MultiPPContainerEmpty<5>(); } template void MultiPreprocessorsContainerNoAlignment(dummyPreprocessors::pp_mode MODE) { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t n_preprocessors = 2; unsigned char alignment = 0; int initial_value = 1; int value_to_add = 7; const int original_blob[4] = {initial_value, initial_value, initial_value, initial_value}; size_t original_blob_size = sizeof(original_blob); // Test computer with multiple preprocessors of the same type. auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); auto verify_preprocess = [&](int expected_processed_value) { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); // Original blob should not be changed ASSERT_EQ(original_blob[0], initial_value); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); if (MODE == STORAGE_ONLY) { // New storage blob should be allocated ASSERT_NE(storage_blob, original_blob); // query blob should be unprocessed ASSERT_EQ(query_blob, original_blob); ASSERT_EQ(((const int *)storage_blob)[0], expected_processed_value); } else if (MODE == QUERY_ONLY) { // New query blob should be allocated ASSERT_NE(query_blob, original_blob); // Storage blob should be unprocessed ASSERT_EQ(storage_blob, original_blob); ASSERT_EQ(((const int *)query_blob)[0], expected_processed_value); } }; /* ==== Add the first preprocessor ==== */ auto preprocessor0 = new (allocator) PreprocessorType(allocator, value_to_add); // add preprocessor returns next free spot in its preprocessors array. ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor0), 1); verify_preprocess(initial_value + value_to_add); /* ==== Add the second preprocessor ==== */ auto preprocessor1 = new (allocator) PreprocessorType(allocator, value_to_add); // add preprocessor returns 0 when adding the last preprocessor. ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor1), 0); ASSERT_NO_FATAL_FAILURE(verify_preprocess(initial_value + 2 * value_to_add)); } TEST(PreprocessorsTest, MultiPreprocessorsContainerStorageNoAlignment) { using namespace dummyPreprocessors; MultiPreprocessorsContainerNoAlignment>( pp_mode::STORAGE_ONLY); } TEST(PreprocessorsTest, MultiPreprocessorsContainerQueryNoAlignment) { using namespace dummyPreprocessors; MultiPreprocessorsContainerNoAlignment>(pp_mode::QUERY_ONLY); } template void multiPPContainerMixedPreprocessorNoAlignment() { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t n_preprocessors = 3; unsigned char alignment = 0; int initial_value = 1; int value_to_add_storage = 7; int value_to_add_query = 2; const int original_blob[4] = {initial_value, initial_value, initial_value, initial_value}; size_t original_blob_size = sizeof(original_blob); // Test multiple preprocessors of the same type. auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); /* ==== Add one preprocessor of each type ==== */ auto preprocessor0 = new (allocator) FirstPreprocessorType(allocator, value_to_add_storage, value_to_add_query); ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor0), 1); auto preprocessor1 = new (allocator) SecondPreprocessorType(allocator, value_to_add_storage, value_to_add_query); ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor1), 2); // scope this section so the blobs are released before the allocator. { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); // Original blob should not be changed ASSERT_EQ(original_blob[0], initial_value); // Both blobs should be allocated const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // Ensure the computer process returns a new allocation of the expected processed blob with // the new value. ASSERT_NE(storage_blob, original_blob); ASSERT_NE(query_blob, original_blob); ASSERT_NE(query_blob, storage_blob); ASSERT_EQ(((const int *)storage_blob)[0], initial_value + value_to_add_storage); ASSERT_EQ(((const int *)query_blob)[0], initial_value + value_to_add_query); } /* ==== Add a preprocessor that processes both storage and query ==== */ auto preprocessor2 = new (allocator) DummyMixedPreprocessor(allocator, value_to_add_storage, value_to_add_query); // add preprocessor returns 0 when adding the last preprocessor. ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor2), 0); { ProcessedBlobs mixed_processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); const void *mixed_pp_storage_blob = mixed_processed_blobs.getStorageBlob(); const void *mixed_pp_query_blob = mixed_processed_blobs.getQueryBlob(); // Ensure the computer process both blobs. ASSERT_EQ(((const int *)mixed_pp_storage_blob)[0], initial_value + 2 * value_to_add_storage); ASSERT_EQ(((const int *)mixed_pp_query_blob)[0], initial_value + 2 * value_to_add_query); } // try adding another preprocessor and fail. ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor2), -1); } TEST(PreprocessorsTest, multiPPContainerMixedPreprocessorQueryFirst) { using namespace dummyPreprocessors; multiPPContainerMixedPreprocessorNoAlignment, DummyStoragePreprocessor>(); } TEST(PreprocessorsTest, multiPPContainerMixedPreprocessorStorageFirst) { using namespace dummyPreprocessors; multiPPContainerMixedPreprocessorNoAlignment, DummyQueryPreprocessor>(); } template void multiPPContainerAlignment(dummyPreprocessors::pp_mode MODE) { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); unsigned char alignment = 8; constexpr size_t n_preprocessors = 1; int initial_value = 1; int value_to_add = 7; const int original_blob[4] = {initial_value, initial_value, initial_value, initial_value}; size_t original_blob_size = sizeof(original_blob); auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); auto verify_preprocess = [&](int expected_processed_value) { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); if (MODE == STORAGE_ONLY) { // New storage blob should be allocated and processed ASSERT_NE(storage_blob, original_blob); ASSERT_EQ(((const int *)storage_blob)[0], expected_processed_value); // query blob *values* should be unprocessed, however, it might be allocated if the // original blob is not aligned. ASSERT_EQ(((const int *)query_blob)[0], original_blob[0]); } else if (MODE == QUERY_ONLY) { // New query blob should be allocated ASSERT_NE(query_blob, original_blob); // Storage blob should be unprocessed and not allocated. ASSERT_EQ(storage_blob, original_blob); ASSERT_EQ(((const int *)query_blob)[0], expected_processed_value); } // anyway the query blob should be aligned unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; ASSERT_EQ(address_alignment, 0); }; auto preprocessor0 = new (allocator) PreprocessorType(allocator, value_to_add); // add preprocessor returns next free spot in its preprocessors array. ASSERT_EQ(multiPPContainer.addPreprocessor(preprocessor0), 0); verify_preprocess(initial_value + value_to_add); } TEST(PreprocessorsTest, StoragePreprocessorWithAlignment) { using namespace dummyPreprocessors; multiPPContainerAlignment>(pp_mode::STORAGE_ONLY); } TEST(PreprocessorsTest, QueryPreprocessorWithAlignment) { using namespace dummyPreprocessors; multiPPContainerAlignment>(pp_mode::QUERY_ONLY); } TEST(PreprocessorsTest, multiPPContainerCosineThenMixedPreprocess) { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t n_preprocessors = 2; constexpr size_t dim = 4; unsigned char alignment = 8; float initial_value = 1.0f; float normalized_value = 0.5f; float value_to_add_storage = 7.0f; float value_to_add_query = 2.0f; const float original_blob[dim] = {initial_value, initial_value, initial_value, initial_value}; constexpr size_t original_blob_size = sizeof(original_blob); auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); // adding cosine preprocessor auto cosine_preprocessor = new (allocator) CosinePreprocessor(allocator, dim, original_blob_size); multiPPContainer.addPreprocessor(cosine_preprocessor); { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // blobs should point to the same memory slot ASSERT_EQ(storage_blob, query_blob); // memory should be aligned unsigned char address_alignment = (uintptr_t)(storage_blob) % alignment; ASSERT_EQ(address_alignment, 0); // They need to be allocated and processed ASSERT_NE(storage_blob, nullptr); ASSERT_EQ(((const float *)storage_blob)[0], normalized_value); // the original blob should not change ASSERT_NE(storage_blob, original_blob); } // adding mixed preprocessor auto mixed_preprocessor = new (allocator) DummyMixedPreprocessor(allocator, value_to_add_storage, value_to_add_query); multiPPContainer.addPreprocessor(mixed_preprocessor); { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // blobs should point to a different memory slot ASSERT_NE(storage_blob, query_blob); ASSERT_NE(storage_blob, nullptr); ASSERT_NE(query_blob, nullptr); // query blob should be aligned unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; ASSERT_EQ(address_alignment, 0); // They need to be processed by both processors. ASSERT_EQ(((const float *)storage_blob)[0], normalized_value + value_to_add_storage); ASSERT_EQ(((const float *)query_blob)[0], normalized_value + value_to_add_query); // the original blob should not change ASSERT_NE(storage_blob, original_blob); ASSERT_NE(query_blob, original_blob); } // The preprocessors should be released by the preprocessors container. } // Cosine pp receives different allocation for the storage and query blobs. TEST(PreprocessorsTest, multiPPContainerMixedThenCosinePreprocess) { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t n_preprocessors = 2; constexpr size_t dim = 4; unsigned char alignment = 8; // In this test the first preprocessor allocates the memory for both blobs, according to the // size passed by the pp container. The second preprocessor expects that if the blobs are // already allocated, their size matches its processed_bytes_count. Hence, the blob size should // be sufficient for the cosine preprocessing. constexpr int8_t normalized_blob_bytes_count = dim * sizeof(int8_t) + sizeof(float); auto blob_alloc = allocator->allocate_unique(normalized_blob_bytes_count); int8_t *original_blob = static_cast(blob_alloc.get()); test_utils::populate_int8_vec(original_blob, dim); // Processing params int8_t value_to_add_storage = 7; int8_t value_to_add_query = 4; // Creating multi preprocessors container auto mixed_preprocessor = new (allocator) DummyMixedPreprocessor(allocator, value_to_add_storage, value_to_add_query); auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); multiPPContainer.addPreprocessor(mixed_preprocessor); { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, normalized_blob_bytes_count); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // blobs should point to a different memory slot ASSERT_NE(storage_blob, query_blob); ASSERT_NE(storage_blob, nullptr); ASSERT_NE(query_blob, nullptr); // query blob should be aligned unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; ASSERT_EQ(address_alignment, 0); // Both blobs were processed. ASSERT_EQ((static_cast(storage_blob))[0], original_blob[0] + value_to_add_storage); ASSERT_EQ((static_cast(query_blob))[0], original_blob[0] + value_to_add_query); // the original blob should not change ASSERT_NE(storage_blob, original_blob); ASSERT_NE(query_blob, original_blob); } // Processed blob expected output auto expected_processed_blob = [&](int8_t *blob, int8_t value_to_add) { memcpy(blob, original_blob, normalized_blob_bytes_count); blob[0] += value_to_add; VecSim_Normalize(blob, dim, VecSimType_INT8); }; int8_t expected_processed_storage[normalized_blob_bytes_count] = {0}; expected_processed_blob(expected_processed_storage, value_to_add_storage); int8_t expected_processed_query[normalized_blob_bytes_count] = {0}; expected_processed_blob(expected_processed_query, value_to_add_query); // normalize the original blob auto cosine_preprocessor = new (allocator) CosinePreprocessor(allocator, dim, normalized_blob_bytes_count); multiPPContainer.addPreprocessor(cosine_preprocessor); { // An assertion should be raised by the cosine preprocessor for unmatching blob sizes. // in valgrind the test continues, but the assertion appears in its log looking like an // error, so to avoid confusion we skip this line in valgrind. #if !defined(RUNNING_ON_VALGRIND) && !defined(NDEBUG) EXPECT_EXIT( { ProcessedBlobs processed_blobs = multiPPContainer.preprocess( original_blob, normalized_blob_bytes_count - sizeof(float)); }, testing::KilledBySignal(SIGABRT), "blob_size == processed_bytes_count"); #endif // Use the correct size ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, normalized_blob_bytes_count); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // blobs should point to a different memory slot ASSERT_NE(storage_blob, query_blob); // query memory should be aligned unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; ASSERT_EQ(address_alignment, 0); // They need to be allocated and processed ASSERT_NE(storage_blob, nullptr); ASSERT_NE(query_blob, nullptr); // the original blob should not change ASSERT_NE(storage_blob, original_blob); ASSERT_NE(query_blob, original_blob); EXPECT_NO_FATAL_FAILURE(CompareVectors(static_cast(storage_blob), expected_processed_storage, normalized_blob_bytes_count)); EXPECT_NO_FATAL_FAILURE(CompareVectors(static_cast(query_blob), expected_processed_query, normalized_blob_bytes_count)); } // The preprocessors should be released by the preprocessors container. } template void AsymmetricPPThenCosine(dummyPreprocessors::pp_mode MODE) { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t n_preprocessors = 2; constexpr size_t dim = 4; unsigned char alignment = 8; float original_blob[dim] = {0}; constexpr size_t original_blob_size = dim * sizeof(float); test_utils::populate_float_vec(original_blob, dim); // Processing params float value_to_add_storage = 0; float value_to_add_query = 0; if (MODE == STORAGE_ONLY) { value_to_add_storage = 7; } else if (MODE == QUERY_ONLY) { value_to_add_query = 4; } // Processed blob expected output auto expected_processed_blob = [&](float *blob, float value_to_add) { memcpy(blob, original_blob, original_blob_size); blob[0] += value_to_add; VecSim_Normalize(blob, dim, VecSimType_FLOAT32); }; float expected_processed_storage[dim] = {0}; expected_processed_blob(expected_processed_storage, value_to_add_storage); float expected_processed_query[dim] = {0}; expected_processed_blob(expected_processed_query, value_to_add_query); auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); // will allocate either the storage or the query blob, depending on the preprocessor type. // TODO : can we pass just value_to_add? auto asymmetric_preprocessor = new (allocator) PreprocessorType(allocator, value_to_add_storage, value_to_add_query); auto cosine_preprocessor = new (allocator) CosinePreprocessor(allocator, dim, original_blob_size); multiPPContainer.addPreprocessor(asymmetric_preprocessor); multiPPContainer.addPreprocessor(cosine_preprocessor); { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // blobs should point to a different memory slot ASSERT_NE(storage_blob, query_blob); ASSERT_NE(storage_blob, nullptr); ASSERT_NE(query_blob, nullptr); // query blob should be aligned unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; ASSERT_EQ(address_alignment, 0); EXPECT_NO_FATAL_FAILURE(CompareVectors(static_cast(storage_blob), expected_processed_storage, dim)); EXPECT_NO_FATAL_FAILURE(CompareVectors(static_cast(query_blob), expected_processed_query, dim)); } } // The first preprocessor in the chain allocates only the storage blob, // and the responsibility of allocating the query blob is delegated to the next preprocessor // in the chain (cosine preprocessor in this case). TEST(PreprocessorsTest, STORAGE_then_cosine_no_size_change) { using namespace dummyPreprocessors; EXPECT_NO_FATAL_FAILURE( AsymmetricPPThenCosine>(pp_mode::STORAGE_ONLY)); } TEST(PreprocessorsTest, QUERY_then_cosine_no_size_change) { using namespace dummyPreprocessors; EXPECT_NO_FATAL_FAILURE( AsymmetricPPThenCosine>(pp_mode::QUERY_ONLY)); } // A test where the value of input_blob_size is modified by the first pp. // The cosine pp processed_bytes_count should be set at initialization with the *modified* value, // otherwise, an assertion will be raised by it for an allocated blob that is smaller than the // expected size. TEST(PreprocessorsTest, DecreaseSizeThenFloatNormalize) { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t n_preprocessors = 2; constexpr size_t alignment = 8; constexpr size_t elements = 8; constexpr size_t decrease_amount = 2; constexpr size_t new_elem_amount = elements - decrease_amount; // valgrind detects out of bound reads only if the considered memory is allocated on the heap, // rather than on the stack. constexpr size_t original_blob_size = elements * sizeof(float); auto original_blob_alloc = allocator->allocate_unique(original_blob_size); float *original_blob = static_cast(original_blob_alloc.get()); test_utils::populate_float_vec(original_blob, elements); constexpr size_t new_processed_bytes_count = original_blob_size - decrease_amount * sizeof(float); // Processed blob expected output float expected_processed_blob[elements] = {0}; memcpy(expected_processed_blob, original_blob, new_processed_bytes_count); // Only use half of the blob for normalization VecSim_Normalize(expected_processed_blob, new_elem_amount, VecSimType_FLOAT32); // A pp that decreases the original blob size auto decrease_size_preprocessor = new (allocator) DummyChangeAllocSizePreprocessor(allocator, new_processed_bytes_count); // A normalize pp auto cosine_preprocessor = new (allocator) CosinePreprocessor(allocator, new_elem_amount, new_processed_bytes_count); // Creating multi preprocessors container auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); multiPPContainer.addPreprocessor(decrease_size_preprocessor); multiPPContainer.addPreprocessor(cosine_preprocessor); { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // blobs should point to the same memory slot ASSERT_EQ(storage_blob, query_blob); ASSERT_NE(storage_blob, nullptr); // memory should be aligned unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; ASSERT_EQ(address_alignment, 0); // They need to be allocated and processed EXPECT_NO_FATAL_FAILURE(CompareVectors(static_cast(storage_blob), expected_processed_blob, new_elem_amount)); } } // In this test the original blob size is smaller than the processed_bytes_count of the // cosine preprocessor. Before the bug fix, the cosine pp would try to copy processed_bytes_count // bytes of the original blob to the allocated memory, which would cause an out of bound read, // that should be detected by valgrind. TEST(PreprocessorsTest, Int8NormalizeThenIncreaseSize) { using namespace dummyPreprocessors; std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t n_preprocessors = 2; constexpr size_t alignment = 8; constexpr size_t elements = 7; // valgrind detects out of bound reads only if the considered memory is allocated on the heap, // rather than on the stack. constexpr size_t original_blob_size = elements * sizeof(int8_t); auto original_blob_alloc = allocator->allocate_unique(original_blob_size); int8_t *original_blob = static_cast(original_blob_alloc.get()); test_utils::populate_int8_vec(original_blob, elements); // size after normalization constexpr size_t normalized_blob_bytes_count = original_blob_size + sizeof(float); // size after increasing pp constexpr size_t elements_addition = 3; constexpr size_t final_blob_bytes_count = normalized_blob_bytes_count + elements_addition * sizeof(unsigned char); // Processed blob expected output int8_t expected_processed_blob[elements + sizeof(float) + elements_addition] = {0}; memcpy(expected_processed_blob, original_blob, original_blob_size); // normalize the original blob VecSim_Normalize(expected_processed_blob, elements, VecSimType_INT8); // add the values of the pp that increases the blob size unsigned char added_value = DummyChangeAllocSizePreprocessor::getExcessValue(); for (size_t i = 0; i < elements_addition; i++) { expected_processed_blob[elements + sizeof(float) + i] = static_cast(added_value); } // A normalize pp - will allocate the blob + sizeof(float) auto cosine_preprocessor = new (allocator) CosinePreprocessor(allocator, elements, normalized_blob_bytes_count); // A pp that will increase the *normalized* blob size auto increase_size_preprocessor = new (allocator) DummyChangeAllocSizePreprocessor(allocator, final_blob_bytes_count); // Creating multi preprocessors container auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); multiPPContainer.addPreprocessor(cosine_preprocessor); multiPPContainer.addPreprocessor(increase_size_preprocessor); { // Note: we pass the original blob size to detect out of bound reads in the cosine pp. ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // blobs should point to the same memory slot ASSERT_EQ(storage_blob, query_blob); ASSERT_NE(storage_blob, nullptr); // memory should be aligned unsigned char address_alignment = (uintptr_t)(query_blob) % alignment; ASSERT_EQ(address_alignment, 0); // They need to be allocated and processed EXPECT_NO_FATAL_FAILURE(CompareVectors(static_cast(storage_blob), expected_processed_blob, final_blob_bytes_count)); } } // Tests the quantization preprocessor with a single preprocessor in the chain. // The QuantPreprocessor allocates and processes both the storage blob (quantized) and the query // blob (original values + precomputed sum_squares for L2). TEST(PreprocessorsTest, QuantizationTest) { std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t n_preprocessors = 1; constexpr size_t alignment = 8; constexpr size_t dim = 6; constexpr size_t original_blob_size = dim * sizeof(float); float original_blob[dim] = {1, 2, 3, 4, 5, 6}; // === Storage blob expected values === // For L2 metric: quantized values + min + delta + sum + sum_squares = dim bytes + 4 floats constexpr size_t quantized_blob_bytes_count = dim * sizeof(uint8_t) + sq8::storage_metadata_count() * sizeof(float); uint8_t expected_storage_blob[quantized_blob_bytes_count] = {0}; ComputeSQ8Quantization(original_blob, dim, expected_storage_blob); // === Query blob expected values === // Query layout: | query_values[dim] | y_sum_squares (for L2) | constexpr size_t query_blob_bytes_count = (dim + sq8::query_metadata_count()) * sizeof(float); // Compute expected sum and sum of squares for L2: float expected_query_sum = 0; float expected_query_sum_squares = 0; for (size_t i = 0; i < dim; ++i) { expected_query_sum += original_blob[i]; expected_query_sum_squares += original_blob[i] * original_blob[i]; } auto quant_preprocessor = new (allocator) QuantPreprocessor(allocator, dim); auto multiPPContainer = MultiPreprocessorsContainer(allocator, alignment); multiPPContainer.addPreprocessor(quant_preprocessor); // Test preprocess (both storage and query) { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); // Verify storage and query blobs are separate ASSERT_NE(storage_blob, nullptr); ASSERT_NE(query_blob, nullptr); ASSERT_NE(storage_blob, query_blob); // Verify storage blob content EXPECT_NO_FATAL_FAILURE(CompareVectors(static_cast(storage_blob), expected_storage_blob, quantized_blob_bytes_count)); // Verify query blob content const float *query_floats = static_cast(query_blob); EXPECT_NO_FATAL_FAILURE(CompareVectors(query_floats, original_blob, dim)); ASSERT_FLOAT_EQ(query_floats[dim + sq8::SUM_QUERY], expected_query_sum); ASSERT_FLOAT_EQ(query_floats[dim + sq8::SUM_SQUARES_QUERY], expected_query_sum_squares); } // Test preprocessForStorage { auto storage_blob = multiPPContainer.preprocessForStorage(original_blob, original_blob_size); ASSERT_NE(storage_blob.get(), nullptr); EXPECT_NO_FATAL_FAILURE( CompareVectors(static_cast(storage_blob.get()), expected_storage_blob, quantized_blob_bytes_count)); } // Test preprocessQuery (query values + precomputed sum_squares for L2) { auto query_blob = multiPPContainer.preprocessQuery(original_blob, original_blob_size); ASSERT_NE(query_blob.get(), nullptr); // Verify query blob content: original floats followed by sum_squares const float *query_floats = static_cast(query_blob.get()); EXPECT_NO_FATAL_FAILURE(CompareVectors(query_floats, original_blob, dim)); ASSERT_FLOAT_EQ(query_floats[dim + sq8::SUM_QUERY], expected_query_sum); ASSERT_FLOAT_EQ(query_floats[dim + sq8::SUM_SQUARES_QUERY], expected_query_sum_squares); // Check address is aligned unsigned char address_alignment = (uintptr_t)(query_blob.get()) % alignment; ASSERT_EQ(address_alignment, 0) << "expected alignment " << alignment; } // Test preprocessStorageInPlace { // Allocate buffer large enough for in-place quantization auto buffer_alloc = allocator->allocate_unique(original_blob_size); float *buffer = static_cast(buffer_alloc.get()); memcpy(buffer, original_blob, original_blob_size); multiPPContainer.preprocessStorageInPlace(buffer, original_blob_size); EXPECT_NO_FATAL_FAILURE(CompareVectors(reinterpret_cast(buffer), expected_storage_blob, quantized_blob_bytes_count)); #if !defined(NDEBUG) EXPECT_EXIT( { multiPPContainer.preprocessStorageInPlace(buffer, sizeof(uint8_t)); }, testing::KilledBySignal(SIGABRT), "Input buffer too small for in-place quantization"); #endif } } // Verifies that the QuantPreprocessor honors distinct query_alignment and storage_alignment hints // independently. This guards the MOD-13837 contract: storage and query buffers can have different // SIMD alignment requirements (e.g. SQ8 storage vs FP32 query). TEST(PreprocessorsTest, QuantizationAsymmetricAlignment) { std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t n_preprocessors = 1; constexpr unsigned char query_alignment = 32; constexpr unsigned char storage_alignment = 16; constexpr size_t dim = 6; constexpr size_t original_blob_size = dim * sizeof(float); float original_blob[dim] = {1, 2, 3, 4, 5, 6}; auto quant_preprocessor = new (allocator) QuantPreprocessor(allocator, dim); auto multiPPContainer = MultiPreprocessorsContainer( allocator, query_alignment, storage_alignment); multiPPContainer.addPreprocessor(quant_preprocessor); // preprocess() exercises the joint storage+query allocation path. { ProcessedBlobs processed_blobs = multiPPContainer.preprocess(original_blob, original_blob_size); const void *storage_blob = processed_blobs.getStorageBlob(); const void *query_blob = processed_blobs.getQueryBlob(); ASSERT_NE(storage_blob, nullptr); ASSERT_NE(query_blob, nullptr); ASSERT_NE(storage_blob, query_blob); ASSERT_EQ(reinterpret_cast(storage_blob) % storage_alignment, 0u) << "storage blob not aligned to " << static_cast(storage_alignment); ASSERT_EQ(reinterpret_cast(query_blob) % query_alignment, 0u) << "query blob not aligned to " << static_cast(query_alignment); } // preprocessForStorage() is the storage-only path; must honor storage_alignment. { auto storage_blob = multiPPContainer.preprocessForStorage(original_blob, original_blob_size); ASSERT_NE(storage_blob.get(), nullptr); ASSERT_EQ(reinterpret_cast(storage_blob.get()) % storage_alignment, 0u) << "storage blob not aligned to " << static_cast(storage_alignment); } // preprocessQuery() is the query-only path; must honor query_alignment. { auto query_blob = multiPPContainer.preprocessQuery(original_blob, original_blob_size); ASSERT_NE(query_blob.get(), nullptr); ASSERT_EQ(reinterpret_cast(query_blob.get()) % query_alignment, 0u) << "query blob not aligned to " << static_cast(query_alignment); } } // Test edge case where all entries are equal TEST(PreprocessorsTest, QuantizationTestAllEntriesEqual) { std::shared_ptr allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t dim = 5; constexpr unsigned char alignment = 0; float original_blob[dim] = {3.5f, 3.5f, 3.5f, 3.5f, 3.5f}; auto quant_preprocessor = new (allocator) QuantPreprocessor(allocator, dim); void *storage_blob = nullptr; void *query_blob = nullptr; size_t storage_blob_size = dim * sizeof(float); size_t query_blob_size = dim * sizeof(float); quant_preprocessor->preprocess(original_blob, storage_blob, query_blob, storage_blob_size, query_blob_size, alignment, alignment); ASSERT_NE(storage_blob, nullptr); // When all values are equal: min == max, delta = 1, all quantized values should be 0 const uint8_t *quantized = static_cast(storage_blob); for (size_t i = 0; i < dim; ++i) { ASSERT_EQ(quantized[i], 0) << "All equal values should quantize to 0"; } // Verify metadata: min_val = 3.5f, delta = 1.0f (fallback when diff == 0) const float *metadata = reinterpret_cast(quantized + dim); ASSERT_FLOAT_EQ(metadata[sq8::MIN_VAL], 3.5f); // min_val ASSERT_FLOAT_EQ(metadata[sq8::DELTA], 1.0f); // delta (fallback) // Verify sum and sum_squares for L2 metric float expected_sum = 3.5f * dim; float expected_sum_squares = 3.5f * 3.5f * dim; ASSERT_FLOAT_EQ(metadata[sq8::SUM], expected_sum); // sum ASSERT_FLOAT_EQ(metadata[sq8::SUM_SQUARES], expected_sum_squares); // sum_squares // Reconstruct and verify: min + quantized * delta = 3.5 + 0 * 1 = 3.5 for (size_t i = 0; i < dim; ++i) { float reconstructed = metadata[sq8::MIN_VAL] + quantized[i] * metadata[sq8::DELTA]; ASSERT_FLOAT_EQ(reconstructed, original_blob[i]); } allocator->free_allocation(storage_blob); allocator->free_allocation(query_blob); delete quant_preprocessor; } // Parameterized test class for QuantPreprocessor with different metrics class QuantPreprocessorMetricTest : public testing::TestWithParam { protected: static constexpr size_t dim = 5; static constexpr unsigned char alignment = 0; static constexpr size_t original_blob_size = dim * sizeof(float); std::shared_ptr allocator; float original_blob[dim] = {1, 2, 3, 4, 5}; void SetUp() override { allocator = VecSimAllocator::newVecsimAllocator(); } // === Storage blob helpers === // Storage layout: | quantized_values[dim] | min | delta | sum | (sum_squares for L2) | // L2: dim bytes + 4 floats (min, delta, sum, sum_squares) // IP/Cosine: dim bytes + 3 floats (min, delta, sum) template static size_t getExpectedStorageSize() { constexpr size_t extra_floats = sq8::storage_metadata_count(); return dim * sizeof(uint8_t) + extra_floats * sizeof(float); } // === Query blob helpers === // Query layout: | query_values[dim] | y_sum (IP/Cosine) OR y_sum_squares (L2) | // All metrics: (dim + 1) floats template static constexpr size_t getExpectedQuerySize() { return (dim + sq8::query_metadata_count()) * sizeof(float); } // Helper to run quantization test for a specific metric template void runQuantizationTest() { size_t expected_storage_size = getExpectedStorageSize(); size_t expected_query_size = getExpectedQuerySize(); float expected_query_sum = 0; float expected_query_sum_squares = 0; for (size_t i = 0; i < dim; ++i) { expected_query_sum += original_blob[i]; expected_query_sum_squares += original_blob[i] * original_blob[i]; } auto quant_preprocessor = new (allocator) QuantPreprocessor(allocator, dim); // Test preprocess (both storage and query) { void *storage_blob = nullptr; void *query_blob = nullptr; size_t storage_blob_size = original_blob_size; size_t query_blob_size = original_blob_size; quant_preprocessor->preprocess(original_blob, storage_blob, query_blob, storage_blob_size, query_blob_size, alignment, alignment); // Verify storage blob ASSERT_NE(storage_blob, nullptr); ASSERT_EQ(storage_blob_size, expected_storage_size); // Verify query blob ASSERT_NE(query_blob, nullptr); ASSERT_EQ(query_blob_size, expected_query_size); // === Verify storage blob content === constexpr size_t max_storage_size = dim * sizeof(uint8_t) + 4 * sizeof(float); uint8_t expected_storage_blob[max_storage_size]; ComputeSQ8Quantization(original_blob, dim, expected_storage_blob); // For non-L2 metrics, compare only dim + 3 floats (excluding sum_squares) size_t compare_size = (Metric == VecSimMetric_L2) ? expected_storage_size : dim * sizeof(uint8_t) + 3 * sizeof(float); EXPECT_NO_FATAL_FAILURE(CompareVectors( static_cast(storage_blob), expected_storage_blob, compare_size)); // === Verify query blob content === const float *query_floats = static_cast(query_blob); EXPECT_NO_FATAL_FAILURE(CompareVectors(query_floats, original_blob, dim)); // Verify precomputed value (sum for IP/Cosine, sum and sum_squares for L2) ASSERT_FLOAT_EQ(query_floats[dim + sq8::SUM_QUERY], expected_query_sum); if constexpr (Metric == VecSimMetric_L2) { ASSERT_FLOAT_EQ(query_floats[dim + sq8::SUM_SQUARES_QUERY], expected_query_sum_squares); } allocator->free_allocation(storage_blob); allocator->free_allocation(query_blob); } // Test preprocessForStorage { void *blob = nullptr; size_t blob_size = original_blob_size; quant_preprocessor->preprocessForStorage(original_blob, blob, blob_size, alignment); ASSERT_EQ(blob_size, expected_storage_size); allocator->free_allocation(blob); } // Test preprocessQuery { void *blob = nullptr; size_t blob_size = original_blob_size; quant_preprocessor->preprocessQuery(original_blob, blob, blob_size, alignment); ASSERT_NE(blob, nullptr); ASSERT_EQ(blob_size, expected_query_size); // Verify query blob content: original values followed by precomputed value const float *query_floats = static_cast(blob); EXPECT_NO_FATAL_FAILURE(CompareVectors(query_floats, original_blob, dim)); // Verify precomputed value (sum for IP/Cosine, sum and sum_squares for L2) ASSERT_FLOAT_EQ(query_floats[dim + sq8::SUM_QUERY], expected_query_sum); if constexpr (Metric == VecSimMetric_L2) { ASSERT_FLOAT_EQ(query_floats[dim + sq8::SUM_SQUARES_QUERY], expected_query_sum_squares); } allocator->free_allocation(blob); } delete quant_preprocessor; } }; TEST_P(QuantPreprocessorMetricTest, QuantizationBlobSizeAndMetadata) { VecSimMetric metric = GetParam(); switch (metric) { case VecSimMetric_L2: runQuantizationTest(); break; case VecSimMetric_IP: runQuantizationTest(); break; case VecSimMetric_Cosine: runQuantizationTest(); break; } } INSTANTIATE_TEST_SUITE_P(QuantPreprocessorTests, QuantPreprocessorMetricTest, testing::Values(VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine), [](const testing::TestParamInfo &info) { return VecSimMetric_ToString(info.param); }); // Parameterized test class for QuantPreprocessor. Verifies the hybrid layout: // storage = [uint8 * dim][float * N], query = [float16 * dim][float * M], with FP32 metadata // matching the FP32-quantized baseline of the same input widened to FP32. class QuantPreprocessorFP16MetricTest : public testing::TestWithParam { protected: using float16 = vecsim_types::float16; static constexpr size_t dim = 5; static constexpr unsigned char alignment = 0; static constexpr size_t original_blob_size = dim * sizeof(float16); std::shared_ptr allocator; float16 original_blob[dim]; // FP16 input float widened_blob[dim]; // FP32 view of the same input (round-trip through FP16) void SetUp() override { allocator = VecSimAllocator::newVecsimAllocator(); const float src[dim] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f}; for (size_t i = 0; i < dim; ++i) { original_blob[i] = vecsim_types::FP32_to_FP16(src[i]); widened_blob[i] = vecsim_types::FP16_to_FP32(original_blob[i]); } } template static size_t getExpectedStorageSize() { return dim * sizeof(uint8_t) + sq8::storage_metadata_count() * sizeof(float); } template static size_t getExpectedQuerySize() { return dim * sizeof(float16) + sq8::query_metadata_count() * sizeof(float); } // Reads an FP32 metadata scalar at the given byte offset from `base` via memcpy (the // metadata region is not guaranteed to be 4-byte aligned for FP16 query bodies). static float load_meta(const void *base, size_t byte_offset) { float v; std::memcpy(&v, static_cast(base) + byte_offset, sizeof(float)); return v; } template void runQuantizationTest() { const size_t expected_storage_size = getExpectedStorageSize(); const size_t expected_query_size = getExpectedQuerySize(); const size_t storage_meta_offset = dim * sizeof(uint8_t); const size_t query_meta_offset = dim * sizeof(float16); // FP32 baseline: quantize the widened (FP16->FP32) input through the same algorithm. // Read metadata via load_meta() because baseline_storage is a uint8_t buffer and the // metadata region (offset = dim) is not guaranteed to be 4-byte aligned. constexpr size_t max_storage_size = dim * sizeof(uint8_t) + 4 * sizeof(float); uint8_t baseline_storage[max_storage_size]; ComputeSQ8Quantization(widened_blob, dim, baseline_storage); const float baseline_min = load_meta(baseline_storage, dim + sq8::MIN_VAL * sizeof(float)); const float baseline_delta = load_meta(baseline_storage, dim + sq8::DELTA * sizeof(float)); const float baseline_sum = load_meta(baseline_storage, dim + sq8::SUM * sizeof(float)); const float baseline_sum_sq = load_meta(baseline_storage, dim + sq8::SUM_SQUARES * sizeof(float)); auto quant_preprocessor = new (allocator) QuantPreprocessor(allocator, dim); // Test preprocess (both storage and query) { void *storage_blob = nullptr; void *query_blob = nullptr; size_t storage_blob_size = original_blob_size; size_t query_blob_size = original_blob_size; quant_preprocessor->preprocess(original_blob, storage_blob, query_blob, storage_blob_size, query_blob_size, alignment, alignment); // Verify storage blob layout/size ASSERT_NE(storage_blob, nullptr); ASSERT_EQ(storage_blob_size, expected_storage_size); // Verify query blob layout/size ASSERT_NE(query_blob, nullptr); ASSERT_EQ(query_blob_size, expected_query_size); // Storage quantized values must match the FP32 baseline. EXPECT_NO_FATAL_FAILURE(CompareVectors( static_cast(storage_blob), baseline_storage, dim)); // Storage FP32 metadata must match the baseline values. ASSERT_FLOAT_EQ( load_meta(storage_blob, storage_meta_offset + sq8::MIN_VAL * sizeof(float)), baseline_min); ASSERT_FLOAT_EQ( load_meta(storage_blob, storage_meta_offset + sq8::DELTA * sizeof(float)), baseline_delta); ASSERT_FLOAT_EQ(load_meta(storage_blob, storage_meta_offset + sq8::SUM * sizeof(float)), baseline_sum); if constexpr (Metric == VecSimMetric_L2) { ASSERT_FLOAT_EQ( load_meta(storage_blob, storage_meta_offset + sq8::SUM_SQUARES * sizeof(float)), baseline_sum_sq); } // Query body must be a bit-equal copy of the FP16 input. EXPECT_NO_FATAL_FAILURE(CompareVectors( static_cast(query_blob), original_blob, dim)); // Query FP32 metadata: y_sum (and y_sum_squares for L2) match the FP32 baseline. ASSERT_FLOAT_EQ( load_meta(query_blob, query_meta_offset + sq8::SUM_QUERY * sizeof(float)), baseline_sum); if constexpr (Metric == VecSimMetric_L2) { ASSERT_FLOAT_EQ(load_meta(query_blob, query_meta_offset + sq8::SUM_SQUARES_QUERY * sizeof(float)), baseline_sum_sq); } allocator->free_allocation(storage_blob); allocator->free_allocation(query_blob); } // Test preprocessQuery alone. { void *blob = nullptr; size_t blob_size = original_blob_size; quant_preprocessor->preprocessQuery(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)); ASSERT_FLOAT_EQ(load_meta(blob, query_meta_offset + sq8::SUM_QUERY * sizeof(float)), baseline_sum); if constexpr (Metric == VecSimMetric_L2) { ASSERT_FLOAT_EQ( load_meta(blob, query_meta_offset + sq8::SUM_SQUARES_QUERY * sizeof(float)), baseline_sum_sq); } allocator->free_allocation(blob); } delete quant_preprocessor; } }; TEST_P(QuantPreprocessorFP16MetricTest, QuantizationBlobSizeAndMetadata) { VecSimMetric metric = GetParam(); switch (metric) { case VecSimMetric_L2: runQuantizationTest(); break; case VecSimMetric_IP: runQuantizationTest(); break; case VecSimMetric_Cosine: runQuantizationTest(); break; } } INSTANTIATE_TEST_SUITE_P(QuantPreprocessorFP16Tests, QuantPreprocessorFP16MetricTest, testing::Values(VecSimMetric_L2, VecSimMetric_IP, VecSimMetric_Cosine), [](const testing::TestParamInfo &info) { return VecSimMetric_ToString(info.param); }); // Quantize -> reconstruct round-trip for FP16 input. Verifies that for each quantized value // q_i, reconstructed = min + delta * q_i is within one quantization step of the original // FP16 value (widened to FP32). Also covers the in-place quantization path. TEST(QuantPreprocessorFP16Test, QuantizeReconstructRoundTripL2) { using float16 = vecsim_types::float16; auto allocator = VecSimAllocator::newVecsimAllocator(); constexpr size_t dim = 17; // odd, exercises the tail loop and unaligned metadata writes const float src[dim] = {-3.5f, -2.0f, -1.25f, -0.5f, -0.125f, 0.0f, 0.125f, 0.5f, 1.0f, 1.5f, 2.0f, 2.5f, 3.0f, 3.25f, 3.4f, 3.45f, 3.5f}; float16 input[dim]; float widened[dim]; for (size_t i = 0; i < dim; ++i) { input[i] = vecsim_types::FP32_to_FP16(src[i]); widened[i] = vecsim_types::FP16_to_FP32(input[i]); } auto preprocessor = new (allocator) QuantPreprocessor(allocator, dim); void *storage_blob = nullptr; size_t storage_blob_size = 0; preprocessor->preprocessForStorage(input, storage_blob, storage_blob_size, 0); ASSERT_NE(storage_blob, nullptr); ASSERT_EQ(storage_blob_size, dim * sizeof(uint8_t) + 4 * sizeof(float)); const uint8_t *quantized = static_cast(storage_blob); float min_val, delta; std::memcpy(&min_val, quantized + dim + sq8::MIN_VAL * sizeof(float), sizeof(float)); std::memcpy(&delta, quantized + dim + sq8::DELTA * sizeof(float), sizeof(float)); // Reconstruction error should be bounded by the quantization step (delta). for (size_t i = 0; i < dim; ++i) { const float reconstructed = min_val + delta * static_cast(quantized[i]); EXPECT_NEAR(reconstructed, widened[i], delta); } // In-place path: seed a buffer large enough to hold both the FP16 input and the SQ8 // storage layout, copy the FP16 input in, and quantize in place. The resulting SQ8 blob // must match the one produced by preprocessForStorage. constexpr size_t input_size = dim * sizeof(float16); constexpr size_t storage_size = dim * sizeof(uint8_t) + 4 * sizeof(float); constexpr size_t buf_size = (input_size > storage_size) ? input_size : storage_size; alignas(float) uint8_t in_place_buf[buf_size]{}; std::memcpy(in_place_buf, input, input_size); preprocessor->preprocessStorageInPlace(in_place_buf, buf_size); EXPECT_NO_FATAL_FAILURE(CompareVectors( in_place_buf, static_cast(storage_blob), storage_size)); allocator->free_allocation(storage_blob); delete preprocessor; }