[MOD-14955] Add DistanceCalculatorWithNorm

src/VecSim/spaces/computer/calculator.h

@@ -11,6 +11,8 @@
 
 #include "VecSim/memory/vecsim_base.h"
 #include "VecSim/spaces/spaces.h"
+#include "VecSim/types/sq8.h"
+#include "VecSim/utils/alignment.h"
 
 enum class DistanceMode {
     StoredToStored,
@@ -120,3 +122,121 @@ class DistanceCalculatorCommon
         return DistanceDispatch<DistType>::stateless(func);
     }
 };
+
+/**
+ * Distance calculator for mean-normalized SQ8 indices.
+ *
+ * Stored vectors are SQ8-quantized from x' = x - mean. This calculator applies analytical
+ * correction terms derived from the per-vector x_mean_ip / y_mean_ip metadata stored in
+ * the blobs by QuantPreprocessor WithNorm=True to recover correct distances between the
+ * original (un-shifted) vectors x and y.
+ *
+ * Correction formulas (expressed in terms of distance semantics):
+ *
+ * Asymmetric (query y as FP32/FP16, stored vector x as SQ8 of x'):
+ *   IP distance:  dist(x,y) = asym_func(x_blob, y_blob) - y_mean_ip
+ *   L2 distance:  dist(x,y) = asym_func(x_blob, y_blob) + 2*(x_mean_ip - y_mean_ip)
+ *                              - mean_sum_squares
+ *
+ * Symmetric (both x and y as SQ8 blobs):
+ *   IP distance:  dist(x,y) = sym_func(x_blob, y_blob) - x_mean_ip - y_mean_ip
+ *                              + mean_sum_squares
+ *   L2 distance:  dist(x,y) = sym_func(x_blob, y_blob)   (mean terms cancel exactly)
+ *
+ * Note: IP dist functions return 1 - IP(x',y'), not raw inner product. L2 dist functions
+ * return ||x'-y'||² directly.
+ *
+ * DataType is float or float16
+ * DistType is float
+ */
+template <typename DataType, typename DistType, VecSimMetric Metric>
+class DistanceCalculatorWithNorm : public IndexCalculatorInterface<DistType> {
+    static_assert(Metric == VecSimMetric_L2 || Metric == VecSimMetric_IP,
+                  "DistanceCalculatorWithNorm only supports L2 and IP metrics");
+
+private:
+    using sq8 = vecsim_types::sq8;
+
+    struct WithNormDistanceContext {
+        spaces::dist_func_t<DistType> stored_func;
+        spaces::dist_func_t<DistType> query_func;
+        float mean_sum_squares;
+    };
+
+    const WithNormDistanceContext context_;
+
+    static DistType calcStoredWithContext(const void *opaque_context, const void *v1,
+                                          const void *v2, size_t dim) {
+        const auto *context = static_cast<const WithNormDistanceContext *>(opaque_context);
+        DistType base = context->stored_func(v1, v2, dim);
+        if constexpr (Metric == VecSimMetric_IP) {
+            // base = 1 - IP(x', y'). We want 1 - IP(x, y).
+            // IP(x, y) = IP(x', y') + x_mean_ip + y_mean_ip - mean_sum_squares
+            // => 1 - IP(x, y) = base - x_mean_ip - y_mean_ip + mean_sum_squares
+            float x_mean_ip =
+                load_unaligned<float>(static_cast<const uint8_t *>(v1) + dim +
+                                      sq8::template mean_ip_index<Metric>() * sizeof(float));
+            float y_mean_ip =
+                load_unaligned<float>(static_cast<const uint8_t *>(v2) + dim +
+                                      sq8::template mean_ip_index<Metric>() * sizeof(float));
+            return base - x_mean_ip - y_mean_ip + context->mean_sum_squares;
+        } else { // L2: ||x - y||² = ||x' - y'||² (mean terms cancel exactly)
+            return base;
+        }
+    }
+
+    static DistType calcQueryWithContext(const void *opaque_context, const void *candidate,
+                                         const void *query, size_t dim) {
+        const auto *context = static_cast<const WithNormDistanceContext *>(opaque_context);
+        DistType base = context->query_func(candidate, query, dim);
+        float y_mean_ip =
+            load_unaligned<float>(static_cast<const uint8_t *>(query) + dim * sizeof(DataType) +
+                                  sq8::template query_mean_ip_index<Metric>() * sizeof(float));
+        if constexpr (Metric == VecSimMetric_IP) {
+            // base = 1 - IP(x', y). We want 1 - IP(x, y).
+            // IP(x, y) = IP(x', y) + y_mean_ip
+            // => 1 - IP(x, y) = base - y_mean_ip
+            return base - y_mean_ip;
+        } else { // L2
+            // base = ||x' - y||². We want ||x - y||².
+            // ||x - y||² = ||x' - y||² + 2*(x_mean_ip - y_mean_ip) - mean_sum_squares
+            float x_mean_ip =
+                load_unaligned<float>(static_cast<const uint8_t *>(candidate) + dim +
+                                      sq8::template mean_ip_index<Metric>() * sizeof(float));
+            return base + 2.0f * (x_mean_ip - y_mean_ip) - context->mean_sum_squares;
+        }
+    }
+
+public:
+    DistanceCalculatorWithNorm(std::shared_ptr<VecSimAllocator> allocator,
+                               spaces::dist_func_t<DistType> asym_func,
+                               spaces::dist_func_t<DistType> sym_func, float mean_sum_squares)
+        : IndexCalculatorInterface<DistType>(allocator), context_(WithNormDistanceContext{
+                                                             .stored_func = sym_func,
+                                                             .query_func = asym_func,
+                                                             .mean_sum_squares = mean_sum_squares,
+                                                         }) {}
+
+    // Symmetric: both v1 and v2 are stored SQ8-of-x' blobs.
+    DistType calcDistance(const void *v1, const void *v2, size_t dim) const override {
+        return calcStoredWithContext(&context_, v1, v2, dim);
+    }
+
+    // Asymmetric: query is raw FP32/FP16 blob, candidate is stored SQ8-of-x' blob.
+    // Note: query_dist_func expects (storage, query) argument order.
+    DistType calcDistanceForQuery(const void *candidate, const void *query,
+                                  size_t dim) const override {
+        return calcQueryWithContext(&context_, candidate, query, dim);
+    }
+
+    DistanceDispatch<DistType> getDistanceDispatch(DistanceMode mode) const override {
+        if (mode == DistanceMode::StoredToStored) {
+            if constexpr (Metric == VecSimMetric_L2) {
+                // The mean terms cancel for stored-to-stored L2, so retain the stateless fast path.
+                return DistanceDispatch<DistType>::stateless(context_.stored_func);
+            }
+            return DistanceDispatch<DistType>::stateful(&context_, calcStoredWithContext);
+        }
+        return DistanceDispatch<DistType>::stateful(&context_, calcQueryWithContext);
+    }
+};

src/VecSim/spaces/computer/preprocessors.h

@@ -17,12 +17,14 @@
 #include <cstring>
 #include <memory>
 #include <type_traits>
+#include <variant>
 
 #include "VecSim/memory/vecsim_base.h"
 #include "VecSim/spaces/spaces.h"
 #include "VecSim/memory/memory_utils.h"
 #include "VecSim/types/float16.h"
 #include "VecSim/types/sq8.h"
+#include "VecSim/utils/vecsim_stl.h"
 
 class PreprocessorInterface : public VecsimBaseObject {
 public:
@@ -235,7 +237,7 @@ static inline float to_fp32(T x) {
     }
 }
 
-template <QuantInput DataType, VecSimMetric Metric>
+template <QuantInput DataType, VecSimMetric Metric, bool WithNorm = false>
 class QuantPreprocessor : public PreprocessorInterface {
     using OUTPUT_TYPE = uint8_t;
     using MetadataType = float; // SQ8 metadata is always FP32 (see class doc).
@@ -244,13 +246,35 @@ class QuantPreprocessor : public PreprocessorInterface {
     static_assert(Metric == VecSimMetric_L2 || Metric == VecSimMetric_IP ||
                       Metric == VecSimMetric_Cosine,
                   "QuantPreprocessor only supports L2, IP and Cosine metrics");
+    static_assert(!WithNorm || Metric != VecSimMetric_Cosine,
+                  "WithNorm does not support Cosine metric.");
 
     // Helper function to perform quantization. This function is used by the storage preprocessing
     // methods.
     void quantize(const DataType *input, OUTPUT_TYPE *quantized) const {
         assert(input && quantized);
-        // Find min and max values (computed in MetadataType regardless of DataType).
-        auto [min_val, max_val] = find_min_max(input);
+
+        if constexpr (std::is_same_v<DataType, float> && !WithNorm) {
+            quantize_impl(input, quantized);
+        } else {
+            float x_mean_ip = 0.0f; // only used for WithNorm
+
+            // Get transformed values (convert to float, mean-subtracted for WithNorm)
+            vecsim_stl::vector<float> values(this->dim, this->allocator);
+            for (size_t i = 0; i < this->dim; ++i) {
+                values[i] = to_fp32<DataType>(input[i]);
+                if constexpr (WithNorm) {
+                    x_mean_ip += values[i] * mean[i];
+                    values[i] -= mean[i];
+                }
+            }
+            quantize_impl(values.data(), quantized, x_mean_ip);
+        }
+    }
+
+    void quantize_impl(const float *values, OUTPUT_TYPE *quantized, float x_mean_ip = 0.0f) const {
+        // Find min and max values.
+        auto [min_val, max_val] = find_min_max(values);
 
         // Calculate scaling factor (typed as MetadataType because they end up as metadata).
         const MetadataType diff = (max_val - min_val);
@@ -272,10 +296,10 @@ class QuantPreprocessor : public PreprocessorInterface {
         // Quantize the values
         for (; i < dim_round_down; i += 4) {
             // Load once (widened to FP32 if DataType is FP16).
-            const float x0 = to_fp32<DataType>(input[i + 0]);
-            const float x1 = to_fp32<DataType>(input[i + 1]);
-            const float x2 = to_fp32<DataType>(input[i + 2]);
-            const float x3 = to_fp32<DataType>(input[i + 3]);
+            const float x0 = values[i + 0];
+            const float x1 = values[i + 1];
+            const float x2 = values[i + 2];
+            const float x3 = values[i + 3];
             // We know (input - min) => 0
             // If min == max, all values are the same and should be quantized to 0.
             // reconstruction will yield the same original value for all vectors.
@@ -305,7 +329,7 @@ class QuantPreprocessor : public PreprocessorInterface {
         MetadataType sum_squares = (q0 + q1) + (q2 + q3);
 
         for (; i < this->dim; ++i) {
-            const float x = to_fp32<DataType>(input[i]);
+            const float x = values[i];
             quantized[i] = static_cast<OUTPUT_TYPE>(std::round((x - min_val) * inv_delta));
             sum += x;
             if constexpr (Metric == VecSimMetric_L2) {
@@ -316,26 +340,31 @@ class QuantPreprocessor : public PreprocessorInterface {
         // Metadata uses MetadataType. Use memcpy because the metadata offset
         // (dim * sizeof(uint8_t)) is not guaranteed to be sizeof(MetadataType)-aligned.
         void *meta_dst = quantized + this->dim;
-        if constexpr (Metric == VecSimMetric_L2) {
-            const MetadataType buf[4] = {min_val, delta, sum, sum_squares};
-            memcpy(meta_dst, buf, sizeof(buf));
-        } else {
-            const MetadataType buf[3] = {min_val, delta, sum};
-            memcpy(meta_dst, buf, sizeof(buf));
-        }
+        MetadataType buf[5] = {min_val, delta, sum};
+        size_t n = 3;
+        if constexpr (Metric == VecSimMetric_L2)
+            buf[n++] = sum_squares;
+        if constexpr (WithNorm)
+            buf[n++] = x_mean_ip;
+        memcpy(meta_dst, buf, n * sizeof(MetadataType));
     }
 
     // Computes and writes query metadata (FP32) in a single pass over the input vector.
-    // For IP/Cosine: writes y_sum = Σy_i
-    // For L2: writes y_sum = Σy_i and y_sum_squares = Σy_i²
+    // Without norm:
+    //   For IP/Cosine: writes y_sum = Σy_i
+    //   For L2: writes y_sum = Σy_i and y_sum_squares = Σy_i²
+    // With norm: additionally appends y_mean_ip = Σ(mean_i * y_i).
     // The output pointer addresses the metadata region after the query body and may not be
     // 4-byte aligned (e.g. FP16 query body with odd dim), so writes go through memcpy.
     void assign_query_metadata(const DataType *input, void *output_metadata) const {
+
         // Accumulators are FP32 to preserve precision for FP16 inputs.
         // 4 independent accumulators for sum
         float s0{}, s1{}, s2{}, s3{};
         // 4 independent accumulators for sum of squares (only used for L2)
         float q0{}, q1{}, q2{}, q3{};
+        // 4 independent accumulators for y_mean_ip (only used for WithNorm)
+        float m0{}, m1{}, m2{}, m3{};
 
         size_t i = 0;
         // round dim down to the nearest multiple of 4
@@ -358,11 +387,19 @@ class QuantPreprocessor : public PreprocessorInterface {
                 q2 += y2 * y2;
                 q3 += y3 * y3;
             }
+
+            if constexpr (WithNorm) {
+                m0 += mean[i + 0] * y0;
+                m1 += mean[i + 1] * y1;
+                m2 += mean[i + 2] * y2;
+                m3 += mean[i + 3] * y3;
+            }
         }
 
-        // Sum/sum_squares become metadata, so they are MetadataType.
+        // Sum/sum_squares/y_mean_ip become metadata, so they are MetadataType.
         MetadataType sum = (s0 + s1) + (s2 + s3);
         MetadataType sum_squares = (q0 + q1) + (q2 + q3);
+        MetadataType y_mean_ip = (m0 + m1) + (m2 + m3);
 
         // Tail: handle remaining elements
         for (; i < this->dim; ++i) {
@@ -371,29 +408,45 @@ class QuantPreprocessor : public PreprocessorInterface {
             if constexpr (Metric == VecSimMetric_L2) {
                 sum_squares += y * y;
             }
+            if constexpr (WithNorm) {
+                y_mean_ip += mean[i] * y;
+            }
         }
 
         // Metadata uses MetadataType. Use memcpy because the metadata offset (after the query
         // body of dim * sizeof(DataType)) is not guaranteed to be sizeof(MetadataType)-aligned
         // when DataType is float16 and dim is odd.
-        if constexpr (Metric == VecSimMetric_L2) {
-            const MetadataType buf[2] = {sum, sum_squares};
-            memcpy(output_metadata, buf, sizeof(buf));
-        } else {
-            const MetadataType buf[1] = {sum};
-            memcpy(output_metadata, buf, sizeof(buf));
-        }
+        MetadataType buf[3] = {sum};
+        size_t n = 1;
+        if constexpr (Metric == VecSimMetric_L2)
+            buf[n++] = sum_squares;
+        if constexpr (WithNorm)
+            buf[n++] = y_mean_ip;
+        memcpy(output_metadata, buf, n * sizeof(MetadataType));
     }
 
 public:
+    // Standard constructor (WithNorm == false): no mean vector.
     QuantPreprocessor(std::shared_ptr<VecSimAllocator> allocator, size_t dim)
+        requires(!WithNorm)
         : PreprocessorInterface(allocator), dim(dim),
           storage_bytes_count(dim * sizeof(OUTPUT_TYPE) +
-                              (vecsim_types::sq8::storage_metadata_count<Metric>()) *
+                              sq8::storage_metadata_count<Metric>() * sizeof(MetadataType)),
+          query_bytes_count(dim * sizeof(DataType) +
+                            sq8::query_metadata_count<Metric>() * sizeof(MetadataType)) {}
+
+    // WithNorm constructor: accepts a pre-computed mean vector (FP32, length == dim).
+    QuantPreprocessor(std::shared_ptr<VecSimAllocator> allocator, size_t dim,
+                      const vecsim_stl::vector<float> &mean_vec)
+        requires(WithNorm)
+        : PreprocessorInterface(allocator), mean(mean_vec), dim(dim),
+          storage_bytes_count(dim * sizeof(OUTPUT_TYPE) +
+                              sq8::storage_metadata_count<Metric, WithNorm>() *
                                   sizeof(MetadataType)),
           query_bytes_count(dim * sizeof(DataType) +
-                            (vecsim_types::sq8::query_metadata_count<Metric>()) *
-                                sizeof(MetadataType)) {}
+                            sq8::query_metadata_count<Metric, WithNorm>() * sizeof(MetadataType)) {
+        assert(this->mean.size() == dim && "mean vector size must equal dim");
+    }
 
     /**
      * Preprocesses the original blob into separate storage and query blobs.
@@ -495,15 +548,15 @@ class QuantPreprocessor : public PreprocessorInterface {
     }
 
 private:
-    // Returns (min, max) of the input vector evaluated in MetadataType. Both float and float16
-    // expose a usable operator< (float16's overload delegates to FP32 semantics), so a single
-    // std::minmax_element call covers both DataTypes. The returned values are written verbatim
-    // into the metadata region.
-    std::pair<MetadataType, MetadataType> find_min_max(const DataType *input) const {
+    std::pair<float, float> find_min_max(const float *input) const {
         auto [min_it, max_it] = std::minmax_element(input, input + dim);
-        return {to_fp32<DataType>(*min_it), to_fp32<DataType>(*max_it)};
+        return {*min_it, *max_it};
     }
 
+    // Mean vector for WithNorm=true; zero-size placeholder otherwise (no runtime overhead).
+    [[no_unique_address]] std::conditional_t<WithNorm, vecsim_stl::vector<float>, std::monostate>
+        mean;
+
     const size_t dim;
     const size_t storage_bytes_count;
     const size_t query_bytes_count;

src/VecSim/types/sq8.h

@@ -31,15 +31,27 @@ struct sq8 {
         SUM_SQUARES_QUERY = 1 // Only for L2
     };
 
-    // Template on metric — compile-time constant when metric is known
-    template <VecSimMetric Metric>
+    // Template on Metric and WithNorm — compile-time constants
+    // WithNorm: one extra metadata slot for x_mean_ip / y_mean_ip
+    template <VecSimMetric Metric, bool WithNorm = false>
     static constexpr size_t storage_metadata_count() {
-        return (Metric == VecSimMetric_L2) ? 4 : 3;
+        return ((Metric == VecSimMetric_L2) ? 4 : 3) + (WithNorm ? 1 : 0);
     }
 
-    template <VecSimMetric Metric>
+    template <VecSimMetric Metric, bool WithNorm = false>
     static constexpr size_t query_metadata_count() {
-        return (Metric == VecSimMetric_L2) ? 2 : 1;
+        return ((Metric == VecSimMetric_L2) ? 2 : 1) + (WithNorm ? 1 : 0);
+    }
+
+    // Index of x_mean_ip / y_mean_ip in the last slot in metadata array
+    template <VecSimMetric Metric>
+    static constexpr size_t mean_ip_index() {
+        return storage_metadata_count<Metric, true>() - 1;
+    }
+
+    template <VecSimMetric Metric>
+    static constexpr size_t query_mean_ip_index() {
+        return query_metadata_count<Metric, true>() - 1;
     }
 };