[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,118 @@ class DistanceCalculatorCommon
         return DistanceDispatch<DistType>::stateless(func);
     }
 };
+
+/**
+ * Distance calculator for mean-normalized SQ8 indices.
+ *
+ * Stored vectors are SQ8-quantized from x' = x - mean. IP queries remain unshifted and use
+ * analytical correction terms derived from the per-vector x_mean_ip / y_mean_ip metadata.
+ * L2 queries are centered once during preprocessing, allowing the base kernel to compute the
+ * original distance directly without cancellation-prone correction terms.
+ *
+ * Correction formulas (expressed in terms of distance semantics):
+ *
+ * Asymmetric (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, centered_y_blob)
+ *                              where centered_y = y - mean
+ *
+ * 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);
+        if constexpr (Metric == VecSimMetric_IP) {
+            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));
+            // 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;
+        }
+        // L2 query preprocessing centers y, so base = ||(x - mean) - (y - mean)||².
+        return base;
+    }
+
+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: IP queries are raw; L2 queries are centered during preprocessing. The candidate
+    // is a 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 constexpr (Metric == VecSimMetric_L2) {
+            auto func =
+                mode == DistanceMode::StoredToStored ? context_.stored_func : context_.query_func;
+            return DistanceDispatch<DistType>::stateless(func);
+        }
+        if (mode == DistanceMode::StoredToStored) {
+            return DistanceDispatch<DistType>::stateful(&context_, calcStoredWithContext);
+        }
+        return DistanceDispatch<DistType>::stateful(&context_, calcQueryWithContext);
+    }
+};

src/VecSim/spaces/computer/preprocessors.h

@@ -167,8 +167,9 @@ class CosinePreprocessor : public PreprocessorInterface {
  *
  * Query processing:
  * The query vector is not quantized. It remains in DataType width (FP32 stays FP32, FP16 stays
- * FP16), but we precompute and store metric-specific FP32 values to accelerate asymmetric
- * distance computation:
+ * FP16). Normalized L2 stores y' = y - mean; all other modes store y unchanged. We precompute and
+ * store metric-specific FP32 values from the query body to accelerate asymmetric distance
+ * computation:
  * - For IP/Cosine: y_sum = Σy_i (sum of query values)
  * - For L2: y_sum = Σy_i (sum of query values), y_sum_squares = Σy_i² (sum of squared query values)
  *
@@ -199,6 +200,8 @@ class CosinePreprocessor : public PreprocessorInterface {
  *     - x_sum_squares = Σx_i² is precomputed and stored in the storage blob
  *     - IP(x, y) is computed using the formula above
  *     - y_sum_squares = Σy_i² is precomputed and stored in the query blob
+ *   For normalized L2, x and y in this formula are the centered values x' and y'; their distance
+ *   is identical to the distance between the original vectors.
  *
  * === Symmetric distance (both x and y are quantized) ===
  *
@@ -213,13 +216,13 @@ class CosinePreprocessor : public PreprocessorInterface {
  *            = min_x * sum_y + min_y * sum_x - dim * min_x * min_y
  *              + delta_x * delta_y * Σ(qx_i * qy_i)
  *   where:
- *     - sum_x, sum_y are precomputed sums of original values
+ *     - sum_x, sum_y are precomputed sums of the values represented by each blob
  *     - Σqx_i = (sum_x - dim * min_x) / delta_x  (sum of quantized values, derived from stored sum)
  *     - Σqy_i = (sum_y - dim * min_y) / delta_y
  *
  * For L2:
  *   ||x - y||² = sum_sq_x + sum_sq_y - 2 * IP(x, y)
- *   where sum_sq_x, sum_sq_y are precomputed sums of squared original values.
+ *   where sum_sq_x, sum_sq_y are precomputed sums of the squared represented values.
  */
 // Input types accepted by QuantPreprocessor. Opt-in via std::same_as so unrelated types
 // (e.g. integers, double, bfloat16) are rejected at the template head with a named constraint.
@@ -237,6 +240,15 @@ static inline float to_fp32(T x) {
     }
 }
 
+template <QuantInput T>
+static inline T from_fp32(float x) {
+    if constexpr (std::is_same_v<T, vecsim_types::float16>) {
+        return vecsim_types::FP32_to_FP16(x);
+    } else {
+        return x;
+    }
+}
+
 template <QuantInput DataType, VecSimMetric Metric, bool WithNorm = false>
 class QuantPreprocessor : public PreprocessorInterface {
     using OUTPUT_TYPE = uint8_t;
@@ -277,14 +289,14 @@ 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 = transformed_value(input, i + 0);
+            const float x0 = transformed_value(input, i);
             const float x1 = transformed_value(input, i + 1);
             const float x2 = transformed_value(input, i + 2);
             const float x3 = transformed_value(input, 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.
-            quantized[i + 0] = static_cast<OUTPUT_TYPE>(std::round((x0 - min_val) * inv_delta));
+            quantized[i] = static_cast<OUTPUT_TYPE>(std::round((x0 - min_val) * inv_delta));
             quantized[i + 1] = static_cast<OUTPUT_TYPE>(std::round((x1 - min_val) * inv_delta));
             quantized[i + 2] = static_cast<OUTPUT_TYPE>(std::round((x2 - min_val) * inv_delta));
             quantized[i + 3] = static_cast<OUTPUT_TYPE>(std::round((x3 - min_val) * inv_delta));
@@ -325,19 +337,22 @@ class QuantPreprocessor : public PreprocessorInterface {
         size_t n = 3;
         if constexpr (Metric == VecSimMetric_L2)
             buf[n++] = sum_squares;
-        if constexpr (WithNorm)
+        if constexpr (WithNorm && Metric == VecSimMetric_IP)
             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.
+    // Computes and writes query metadata (FP32) in a single pass over the query values.
     // 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).
+    // With norm and Metric == IP: additionally appends y_mean_ip = Σ(mean_i * original_y_i).
+    // Normalized L2 needs no extra metadata; the mean terms cancel exactly during centering.
+    // For normalized L2, values contains y - mean while original_input contains y.
     // 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 {
+    void assign_query_metadata(const DataType *values, const DataType *original_input,
+                               void *output_metadata) const {
 
         // Accumulators are FP32 to preserve precision for FP16 inputs.
         // 4 independent accumulators for sum
@@ -352,10 +367,10 @@ class QuantPreprocessor : public PreprocessorInterface {
         size_t dim_round_down = this->dim & ~size_t(3);
 
         for (; i < dim_round_down; i += 4) {
-            const float y0 = to_fp32<DataType>(input[i + 0]);
-            const float y1 = to_fp32<DataType>(input[i + 1]);
-            const float y2 = to_fp32<DataType>(input[i + 2]);
-            const float y3 = to_fp32<DataType>(input[i + 3]);
+            const float y0 = to_fp32<DataType>(values[i]);
+            const float y1 = to_fp32<DataType>(values[i + 1]);
+            const float y2 = to_fp32<DataType>(values[i + 2]);
+            const float y3 = to_fp32<DataType>(values[i + 3]);
 
             s0 += y0;
             s1 += y1;
@@ -369,11 +384,11 @@ class QuantPreprocessor : public PreprocessorInterface {
                 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;
+            if constexpr (WithNorm && Metric == VecSimMetric_IP) {
+                m0 += mean[i] * to_fp32<DataType>(original_input[i]);
+                m1 += mean[i + 1] * to_fp32<DataType>(original_input[i + 1]);
+                m2 += mean[i + 2] * to_fp32<DataType>(original_input[i + 2]);
+                m3 += mean[i + 3] * to_fp32<DataType>(original_input[i + 3]);
             }
         }
 
@@ -384,13 +399,13 @@ class QuantPreprocessor : public PreprocessorInterface {
 
         // Tail: handle remaining elements
         for (; i < this->dim; ++i) {
-            const float y = to_fp32<DataType>(input[i]);
+            const float y = to_fp32<DataType>(values[i]);
             sum += y;
             if constexpr (Metric == VecSimMetric_L2) {
                 sum_squares += y * y;
             }
-            if constexpr (WithNorm) {
-                y_mean_ip += mean[i] * y;
+            if constexpr (WithNorm && Metric == VecSimMetric_IP) {
+                y_mean_ip += mean[i] * to_fp32<DataType>(original_input[i]);
             }
         }
 
@@ -401,7 +416,7 @@ class QuantPreprocessor : public PreprocessorInterface {
         size_t n = 1;
         if constexpr (Metric == VecSimMetric_L2)
             buf[n++] = sum_squares;
-        if constexpr (WithNorm)
+        if constexpr (WithNorm && Metric == VecSimMetric_IP)
             buf[n++] = y_mean_ip;
         memcpy(output_metadata, buf, n * sizeof(MetadataType));
     }
@@ -488,7 +503,11 @@ class QuantPreprocessor : public PreprocessorInterface {
     /**
      * Preprocesses the query vector for asymmetric distance computation.
      *
-     * The query blob contains the original DataType values followed by FP32 precomputed values:
+     * The query blob contains DataType values followed by FP32 precomputed values. Normalized L2
+     * stores centered query values (y - mean), so its distance kernel directly computes
+     * ||(x - mean) - (y - mean)||² without cancellation-prone correction terms. Other modes keep
+     * the original query values.
+     *
      * - For IP/Cosine: y_sum = Σy_i (sum of query values)
      * - For L2: y_sum = Σy_i (sum of query values), y_sum_squares = Σy_i² (sum of squared query
      *                                                                      values)
@@ -508,13 +527,20 @@ class QuantPreprocessor : public PreprocessorInterface {
         // Allocate aligned memory for the query blob
         blob = this->allocator->allocate_aligned(this->query_bytes_count, alignment);
         const size_t body_bytes = this->dim * sizeof(DataType);
-        memcpy(blob, original_blob, body_bytes);
         const DataType *input = static_cast<const DataType *>(original_blob);
+        DataType *query_values = static_cast<DataType *>(blob);
+        if constexpr (WithNorm && Metric == VecSimMetric_L2) {
+            for (size_t i = 0; i < this->dim; ++i) {
+                query_values[i] = from_fp32<DataType>(to_fp32<DataType>(input[i]) - this->mean[i]);
+            }
+        } else {
+            memcpy(query_values, original_blob, body_bytes);
+        }
 
         // Compute and write FP32 query metadata after the query body. The metadata offset is
         // body_bytes, which is not guaranteed to be 4-byte aligned for FP16 query bodies.
         void *metadata_dst = static_cast<uint8_t *>(blob) + body_bytes;
-        assign_query_metadata(input, metadata_dst);
+        assign_query_metadata(query_values, input, metadata_dst);
 
         query_blob_size = this->query_bytes_count;
     }
@@ -546,14 +572,18 @@ class QuantPreprocessor : public PreprocessorInterface {
             float value = first_input_value - mean[0];
             float min_val = value;
             float max_val = value;
-            x_mean_ip = first_input_value * mean[0];
+            // x_mean_ip only feeds the IP correction term (see DistanceCalculatorWithNorm); the
+            // normalized L2 correction cancels the mean terms exactly, so skip the extra work.
+            if constexpr (Metric == VecSimMetric_IP)
+                x_mean_ip = first_input_value * mean[0];
 
             for (size_t i = 1; i < dim; ++i) {
                 const float input_value = to_fp32<DataType>(input[i]);
                 value = input_value - mean[i];
                 min_val = std::min(min_val, value);
                 max_val = std::max(max_val, value);
-                x_mean_ip += input_value * mean[i];
+                if constexpr (Metric == VecSimMetric_IP)
+                    x_mean_ip += input_value * mean[i];
             }
             return {min_val, max_val};
         }

src/VecSim/types/sq8.h

@@ -32,15 +32,19 @@ struct sq8 {
     };
 
     // Template on Metric and WithNorm — compile-time constants
-    // WithNorm: one extra metadata slot for x_mean_ip / y_mean_ip
+    // WithNorm: one extra metadata slot for x_mean_ip / y_mean_ip. Only needed for IP: the
+    // normalized L2 correction cancels the mean terms exactly during centering and never reads
+    // mean_ip (see DistanceCalculatorWithNorm), so no extra slot is allocated for L2.
     template <VecSimMetric Metric, bool WithNorm = false>
     static constexpr size_t storage_metadata_count() {
-        return ((Metric == VecSimMetric_L2) ? 4 : 3) + (WithNorm ? 1 : 0);
+        return ((Metric == VecSimMetric_L2) ? 4 : 3) +
+               ((WithNorm && Metric == VecSimMetric_IP) ? 1 : 0);
     }
 
     template <VecSimMetric Metric, bool WithNorm = false>
     static constexpr size_t query_metadata_count() {
-        return ((Metric == VecSimMetric_L2) ? 2 : 1) + (WithNorm ? 1 : 0);
+        return ((Metric == VecSimMetric_L2) ? 2 : 1) +
+               ((WithNorm && Metric == VecSimMetric_IP) ? 1 : 0);
     }
 
     // Index of x_mean_ip / y_mean_ip in the last slot in metadata array