Separate query and storage alignment requirements [MOD-13837]

src/VecSim/index_factories/components/preprocessors_factory.h

@@ -14,24 +14,21 @@
 struct PreprocessorsContainerParams {
     VecSimMetric metric;
     size_t dim;
-    unsigned char alignment;
+    unsigned char query_alignment;
+    unsigned char storage_alignment;
     size_t processed_bytes_count;
 };
 
 /**
  * @brief Creates parameters for a preprocessors container based on the given metric, dimension,
- *        normalization flag, and alignment.
+ *        normalization flag, and alignments.
  *
  * @tparam DataType The data type of the vector elements (e.g., float, int).
  * @param metric The similarity metric to be used (e.g., Cosine, Inner Product).
  * @param dim The dimensionality of the vectors.
  * @param is_normalized A flag indicating whether the vectors are already normalized.
- * @param alignment The alignment requirement for the data.
- * @return A PreprocessorsContainerParams object containing the processed parameters:
- *         - metric: The adjusted metric based on the input and normalization flag.
- *         - dim: The dimensionality of the vectors.
- *         - alignment: The alignment requirement for the data.
- *         - processed_bytes_count: The size of the processed data blob in bytes.
+ * @param query_alignment The alignment requirement for query blobs.
+ * @param storage_alignment The alignment requirement for storage blobs.
  *
  * @details
  * If the metric is Cosine and the data type is integral, the processed bytes count may include
@@ -40,9 +37,9 @@ struct PreprocessorsContainerParams {
  * redundant normalization during preprocessing.
  */
 template <typename DataType>
-PreprocessorsContainerParams CreatePreprocessorsContainerParams(VecSimMetric metric, size_t dim,
-                                                                bool is_normalized,
-                                                                unsigned char alignment) {
+PreprocessorsContainerParams
+CreatePreprocessorsContainerParams(VecSimMetric metric, size_t dim, bool is_normalized,
+                                   unsigned char query_alignment, unsigned char storage_alignment) {
     // By default the processed blob size is the same as the original blob size.
     size_t processed_bytes_count = dim * sizeof(DataType);
 
@@ -61,18 +58,29 @@ PreprocessorsContainerParams CreatePreprocessorsContainerParams(VecSimMetric met
     }
     return {.metric = pp_metric,
             .dim = dim,
-            .alignment = alignment,
+            .query_alignment = query_alignment,
+            .storage_alignment = storage_alignment,
             .processed_bytes_count = processed_bytes_count};
 }
 
+// Single-alignment overload: applies the same alignment to both query and storage (homogeneous
+// case). Most existing callers use this form.
+template <typename DataType>
+PreprocessorsContainerParams CreatePreprocessorsContainerParams(VecSimMetric metric, size_t dim,
+                                                                bool is_normalized,
+                                                                unsigned char alignment) {
+    return CreatePreprocessorsContainerParams<DataType>(metric, dim, is_normalized, alignment,
+                                                        alignment);
+}
+
 template <typename DataType>
 PreprocessorsContainerAbstract *
 CreatePreprocessorsContainer(std::shared_ptr<VecSimAllocator> allocator,
                              PreprocessorsContainerParams params) {
 
     if (params.metric == VecSimMetric_Cosine) {
-        auto multiPPContainer =
-            new (allocator) MultiPreprocessorsContainer<DataType, 1>(allocator, params.alignment);
+        auto multiPPContainer = new (allocator) MultiPreprocessorsContainer<DataType, 1>(
+            allocator, params.query_alignment, params.storage_alignment);
         auto cosine_preprocessor = new (allocator)
             CosinePreprocessor<DataType>(allocator, params.dim, params.processed_bytes_count);
         int next_valid_pp_index = multiPPContainer->addPreprocessor(cosine_preprocessor);
@@ -81,19 +89,31 @@ CreatePreprocessorsContainer(std::shared_ptr<VecSimAllocator> allocator,
         return multiPPContainer;
     }
 
-    return new (allocator) PreprocessorsContainerAbstract(allocator, params.alignment);
+    return new (allocator)
+        PreprocessorsContainerAbstract(allocator, params.query_alignment, params.storage_alignment);
 }
 
 template <typename DataType>
 PreprocessorsContainerAbstract *
 CreatePreprocessorsContainer(std::shared_ptr<VecSimAllocator> allocator, VecSimMetric metric,
-                             size_t dim, bool is_normalized, unsigned char alignment) {
+                             size_t dim, bool is_normalized, unsigned char query_alignment,
+                             unsigned char storage_alignment) {
 
-    PreprocessorsContainerParams ppParams =
-        CreatePreprocessorsContainerParams<DataType>(metric, dim, is_normalized, alignment);
+    PreprocessorsContainerParams ppParams = CreatePreprocessorsContainerParams<DataType>(
+        metric, dim, is_normalized, query_alignment, storage_alignment);
     return CreatePreprocessorsContainer<DataType>(allocator, ppParams);
 }
 
+// Single-alignment overload: applies the same alignment to both query and storage (homogeneous
+// case). Most existing callers use this form.
+template <typename DataType>
+PreprocessorsContainerAbstract *
+CreatePreprocessorsContainer(std::shared_ptr<VecSimAllocator> allocator, VecSimMetric metric,
+                             size_t dim, bool is_normalized, unsigned char alignment) {
+    return CreatePreprocessorsContainer<DataType>(allocator, metric, dim, is_normalized, alignment,
+                                                  alignment);
+}
+
 template <typename DataType>
 size_t EstimatePreprocessorsContainerMemory(VecSimMetric metric, bool is_normalized = false) {
     size_t allocations_overhead = VecSimAllocator::getAllocationOverheadSize();

src/VecSim/spaces/IP_space.cpp

@@ -70,23 +70,32 @@ dist_func_t<float> IP_SQ8_FP32_GetDistFunc(size_t dim, unsigned char *alignment,
     if (dim < 16) {
         return ret_dist_func;
     }
+    // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract.
 #ifdef OPT_AVX512_F_BW_VL_VNNI
     if (features.avx512f && features.avx512bw && features.avx512vnni) {
+        if (dim % 16 == 0) // SQ8 chunk = 16 bytes
+            *alignment = 16 * sizeof(uint8_t);
         return Choose_SQ8_FP32_IP_implementation_AVX512F_BW_VL_VNNI(dim);
     }
 #endif
 #ifdef OPT_AVX2_FMA
     if (features.avx2 && features.fma3) {
+        if (dim % 8 == 0) // SQ8 chunk = 8 bytes
+            *alignment = 8 * sizeof(uint8_t);
         return Choose_SQ8_FP32_IP_implementation_AVX2_FMA(dim);
     }
 #endif
 #ifdef OPT_AVX2
     if (features.avx2) {
+        if (dim % 8 == 0) // SQ8 chunk = 8 bytes
+            *alignment = 8 * sizeof(uint8_t);
         return Choose_SQ8_FP32_IP_implementation_AVX2(dim);
     }
 #endif
 #ifdef OPT_SSE4
     if (features.sse4_1) {
+        if (dim % 4 == 0) // SQ8 chunk = 4 bytes
+            *alignment = 4 * sizeof(uint8_t);
         return Choose_SQ8_FP32_IP_implementation_SSE4(dim);
     }
 #endif
@@ -129,23 +138,32 @@ dist_func_t<float> Cosine_SQ8_FP32_GetDistFunc(size_t dim, unsigned char *alignm
     if (dim < 16) {
         return ret_dist_func;
     }
+    // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract.
 #ifdef OPT_AVX512_F_BW_VL_VNNI
     if (features.avx512f && features.avx512bw && features.avx512vnni) {
+        if (dim % 16 == 0) // SQ8 chunk = 16 bytes
+            *alignment = 16 * sizeof(uint8_t);
         return Choose_SQ8_FP32_Cosine_implementation_AVX512F_BW_VL_VNNI(dim);
     }
 #endif
 #ifdef OPT_AVX2_FMA
     if (features.avx2 && features.fma3) {
+        if (dim % 8 == 0) // SQ8 chunk = 8 bytes
+            *alignment = 8 * sizeof(uint8_t);
         return Choose_SQ8_FP32_Cosine_implementation_AVX2_FMA(dim);
     }
 #endif
 #ifdef OPT_AVX2
     if (features.avx2) {
+        if (dim % 8 == 0) // SQ8 chunk = 8 bytes
+            *alignment = 8 * sizeof(uint8_t);
         return Choose_SQ8_FP32_Cosine_implementation_AVX2(dim);
     }
 #endif
 #ifdef OPT_SSE4
     if (features.sse4_1) {
+        if (dim % 4 == 0) // SQ8 chunk = 4 bytes
+            *alignment = 4 * sizeof(uint8_t);
         return Choose_SQ8_FP32_Cosine_implementation_SSE4(dim);
     }
 #endif
@@ -218,7 +236,10 @@ dist_func_t<float> IP_SQ8_SQ8_GetDistFunc(size_t dim, unsigned char *alignment,
 
 #ifdef CPU_FEATURES_ARCH_X86_64
 #ifdef OPT_AVX512_F_BW_VL_VNNI
+    // AVX512 VNNI SQ8_SQ8 uses 64-element chunks; residual handling is in 32-byte sub-chunks.
     if (dim >= 64 && features.avx512f && features.avx512bw && features.avx512vnni) {
+        if (dim % 32 == 0) // align to 256 bits when there is no offsetting residual
+            *alignment = 32 * sizeof(uint8_t);
         return Choose_SQ8_SQ8_IP_implementation_AVX512F_BW_VL_VNNI(dim);
     }
 #endif
@@ -262,7 +283,10 @@ dist_func_t<float> Cosine_SQ8_SQ8_GetDistFunc(size_t dim, unsigned char *alignme
 
 #ifdef CPU_FEATURES_ARCH_X86_64
 #ifdef OPT_AVX512_F_BW_VL_VNNI
+    // AVX512 VNNI SQ8_SQ8 uses 64-element chunks; residual handling is in 32-byte sub-chunks.
     if (dim >= 64 && features.avx512f && features.avx512bw && features.avx512vnni) {
+        if (dim % 32 == 0) // align to 256 bits when there is no offsetting residual
+            *alignment = 32 * sizeof(uint8_t);
         return Choose_SQ8_SQ8_Cosine_implementation_AVX512F_BW_VL_VNNI(dim);
     }
 #endif

src/VecSim/spaces/L2_space.cpp

@@ -70,23 +70,32 @@ dist_func_t<float> L2_SQ8_FP32_GetDistFunc(size_t dim, unsigned char *alignment,
     if (dim < 16) {
         return ret_dist_func;
     }
+    // Alignment hints below refer to the SQ8 (first) operand per the GetDistFunc contract.
 #ifdef OPT_AVX512_F_BW_VL_VNNI
     if (features.avx512f && features.avx512bw && features.avx512vnni) {
+        if (dim % 16 == 0) // SQ8 chunk = 16 bytes; no point in aligning if there's a residual
+            *alignment = 16 * sizeof(uint8_t);
         return Choose_SQ8_FP32_L2_implementation_AVX512F_BW_VL_VNNI(dim);
     }
 #endif
 #ifdef OPT_AVX2_FMA
     if (features.avx2 && features.fma3) {
+        if (dim % 8 == 0) // SQ8 chunk = 8 bytes
+            *alignment = 8 * sizeof(uint8_t);
         return Choose_SQ8_FP32_L2_implementation_AVX2_FMA(dim);
     }
 #endif
 #ifdef OPT_AVX2
     if (features.avx2) {
+        if (dim % 8 == 0) // SQ8 chunk = 8 bytes
+            *alignment = 8 * sizeof(uint8_t);
         return Choose_SQ8_FP32_L2_implementation_AVX2(dim);
     }
 #endif
 #ifdef OPT_SSE4
     if (features.sse4_1) {
+        if (dim % 4 == 0) // SQ8 chunk = 4 bytes
+            *alignment = 4 * sizeof(uint8_t);
         return Choose_SQ8_FP32_L2_implementation_SSE4(dim);
     }
 #endif
@@ -470,8 +479,10 @@ dist_func_t<float> L2_SQ8_SQ8_GetDistFunc(size_t dim, unsigned char *alignment,
 
 #ifdef CPU_FEATURES_ARCH_X86_64
 #ifdef OPT_AVX512_F_BW_VL_VNNI
-    // AVX512 VNNI SQ8_SQ8 uses 64-element chunks
+    // AVX512 VNNI SQ8_SQ8 uses 64-element chunks; residual handling is in 32-byte sub-chunks.
     if (dim >= 64 && features.avx512f && features.avx512bw && features.avx512vnni) {
+        if (dim % 32 == 0) // align to 256 bits when there is no offsetting residual
+            *alignment = 32 * sizeof(uint8_t);
         return Choose_SQ8_SQ8_L2_implementation_AVX512F_BW_VL_VNNI(dim);
     }
 #endif

src/VecSim/spaces/computer/preprocessor_container.cpp

@@ -31,11 +31,13 @@ void PreprocessorsContainerAbstract::preprocessStorageInPlace(void *blob,
 
 MemoryUtils::unique_blob PreprocessorsContainerAbstract::maybeCopyToAlignedMem(
     const void *original_blob, size_t input_blob_size, bool force_copy) const {
-    bool needs_copy =
-        force_copy || (this->alignment && ((uintptr_t)original_blob % this->alignment != 0));
+    // This helper aligns query buffers; storage allocation paths use storage_alignment elsewhere.
+    bool needs_copy = force_copy || (this->query_alignment &&
+                                     ((uintptr_t)original_blob % this->query_alignment != 0));
 
     if (needs_copy) {
-        auto aligned_mem = this->allocator->allocate_aligned(input_blob_size, this->alignment);
+        auto aligned_mem =
+            this->allocator->allocate_aligned(input_blob_size, this->query_alignment);
         memcpy(aligned_mem, original_blob, input_blob_size);
         return this->wrapAllocated(aligned_mem);
     }

src/VecSim/spaces/computer/preprocessor_container.h

@@ -19,9 +19,16 @@ struct ProcessedBlobs;
 
 class PreprocessorsContainerAbstract : public VecsimBaseObject {
 public:
+    // Homogeneous ctor: same value applies to both query and storage alignment.
     PreprocessorsContainerAbstract(std::shared_ptr<VecSimAllocator> allocator,
                                    unsigned char alignment)
-        : VecsimBaseObject(allocator), alignment(alignment) {}
+        : PreprocessorsContainerAbstract(allocator, alignment, alignment) {}
+
+    PreprocessorsContainerAbstract(std::shared_ptr<VecSimAllocator> allocator,
+                                   unsigned char query_alignment, unsigned char storage_alignment)
+        : VecsimBaseObject(allocator), query_alignment(query_alignment),
+          storage_alignment(storage_alignment) {}
+
     // It is assumed that the resulted query blob is aligned.
     virtual ProcessedBlobs preprocess(const void *original_blob, size_t input_blob_size) const;
 
@@ -35,10 +42,12 @@ class PreprocessorsContainerAbstract : public VecsimBaseObject {
 
     virtual void preprocessStorageInPlace(void *blob, size_t input_blob_size) const;
 
-    unsigned char getAlignment() const { return alignment; }
+    unsigned char getQueryAlignment() const { return query_alignment; }
+    unsigned char getStorageAlignment() const { return storage_alignment; }
 
 protected:
-    const unsigned char alignment;
+    const unsigned char query_alignment;
+    const unsigned char storage_alignment;
 
     // Allocate and copy the blob only if the original blob is not aligned.
     MemoryUtils::unique_blob maybeCopyToAlignedMem(const void *original_blob,
@@ -61,8 +70,13 @@ class MultiPreprocessorsContainer : public PreprocessorsContainerAbstract {
     std::array<PreprocessorInterface *, n_preprocessors> preprocessors;
 
 public:
+    // Homogeneous ctor: same value applies to both query and storage alignment.
     MultiPreprocessorsContainer(std::shared_ptr<VecSimAllocator> allocator, unsigned char alignment)
-        : PreprocessorsContainerAbstract(allocator, alignment) {
+        : MultiPreprocessorsContainer(allocator, alignment, alignment) {}
+
+    MultiPreprocessorsContainer(std::shared_ptr<VecSimAllocator> allocator,
+                                unsigned char query_alignment, unsigned char storage_alignment)
+        : PreprocessorsContainerAbstract(allocator, query_alignment, storage_alignment) {
         assert(n_preprocessors);
         std::fill_n(preprocessors.begin(), n_preprocessors, nullptr);
     }
@@ -178,7 +192,7 @@ MultiPreprocessorsContainer<DataType, n_preprocessors>::preprocess(const void *o
         if (!pp)
             break;
         pp->preprocess(original_blob, storage_blob, query_blob, storage_blob_size, query_blob_size,
-                       this->alignment);
+                       this->storage_alignment, this->query_alignment);
     }
     // At least one blob was allocated.
 
@@ -214,7 +228,8 @@ MultiPreprocessorsContainer<DataType, n_preprocessors>::preprocessForStorage(
     for (auto pp : preprocessors) {
         if (!pp)
             break;
-        pp->preprocessForStorage(original_blob, storage_blob, input_blob_size);
+        pp->preprocessForStorage(original_blob, storage_blob, input_blob_size,
+                                 this->storage_alignment);
     }
 
     return storage_blob ? std::move(this->wrapAllocated(storage_blob))
@@ -230,7 +245,7 @@ MemoryUtils::unique_blob MultiPreprocessorsContainer<DataType, n_preprocessors>:
         if (!pp)
             break;
         // modifies the memory in place
-        pp->preprocessQuery(original_blob, query_blob, input_blob_size, this->alignment);
+        pp->preprocessQuery(original_blob, query_blob, input_blob_size, this->query_alignment);
     }
     return query_blob
                ? std::move(this->wrapAllocated(query_blob))