faiss 0.5.3 → 0.6.1

data/vendor/faiss/faiss/impl/ProductQuantizer.cpp

@@ -19,6 +19,7 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/VectorTransform.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
 
 extern "C" {
@@ -47,8 +48,8 @@ namespace faiss {
  * PQ implementation
  *********************************************/
 
-ProductQuantizer::ProductQuantizer(size_t d, size_t M, size_t nbits)
-        : Quantizer(d, 0), M(M), nbits(nbits), assign_index(nullptr) {
+ProductQuantizer::ProductQuantizer(size_t d_in, size_t M_in, size_t nbits_in)
+        : Quantizer(d_in, 0), M(M_in), nbits(nbits_in), assign_index(nullptr) {
     set_derived_values();
 }
 
@@ -56,14 +57,15 @@ ProductQuantizer::ProductQuantizer() : ProductQuantizer(0, 1, 0) {}
 
 void ProductQuantizer::set_derived_values() {
     // quite a few derived values
+    FAISS_THROW_IF_NOT_MSG(M > 0, "M must be > 0");
     FAISS_THROW_IF_NOT_MSG(
             d % M == 0,
             "The dimension of the vector (d) should be a multiple of the number of subquantizers (M)");
     dsub = d / M;
-    code_size = (nbits * M + 7) / 8;
     FAISS_THROW_IF_MSG(nbits > 24, "nbits larger than 24 is not practical.");
+    code_size = (nbits * M + 7) / 8;
     ksub = 1 << nbits;
-    centroids.resize(d * ksub);
+    centroids.resize(mul_no_overflow(d, (size_t)ksub, "PQ centroids"));
     verbose = false;
     train_type = Train_default;
 }
@@ -137,8 +139,8 @@ void ProductQuantizer::train(size_t n, const float* x) {
         }
 
         std::unique_ptr<float[]> xslice(new float[n * dsub]);
-        for (int m = 0; m < M; m++) {
-            for (int j = 0; j < n; j++)
+        for (size_t m = 0; m < M; m++) {
+            for (size_t j = 0; j < n; j++)
                 memcpy(xslice.get() + j * dsub,
                        x + j * d + m * dsub,
                        dsub * sizeof(float));
@@ -177,7 +179,7 @@ void ProductQuantizer::train(size_t n, const float* x) {
 
             if (verbose) {
                 clus.verbose = true;
-                printf("Training PQ slice %d/%zd\n", m, M);
+                printf("Training PQ slice %zd/%zd\n", m, M);
             }
             IndexFlatL2 index(dsub);
             clus.train(n, xslice.get(), assign_index ? *assign_index : index);
@@ -195,14 +197,16 @@ void ProductQuantizer::train(size_t n, const float* x) {
         IndexFlatL2 index(dsub);
 
         clus.train(n * M, x, assign_index ? *assign_index : index);
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             set_params(clus.centroids.data(), m);
         }
     }
 }
 
-template <class PQEncoder>
-void compute_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
+namespace {
+
+template <class PQEncoder, SIMDLevel SL>
+void compute_1_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
     std::vector<float> distances(pq.ksub);
 
     // It seems to be meaningless to allocate std::vector<float> distances.
@@ -248,7 +252,7 @@ void compute_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
         uint64_t idxm = 0;
         if (pq.transposed_centroids.empty()) {
             // the regular version
-            idxm = fvec_L2sqr_ny_nearest(
+            idxm = fvec_L2sqr_ny_nearest<SL>(
                     distances.data(),
                     xsub,
                     pq.get_centroids(m, 0),
@@ -256,7 +260,7 @@ void compute_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
                     pq.ksub);
         } else {
             // transposed centroids are available, use'em
-            idxm = fvec_L2sqr_ny_nearest_y_transposed(
+            idxm = fvec_L2sqr_ny_nearest_y_transposed<SL>(
                     distances.data(),
                     xsub,
                     pq.transposed_centroids.data() + m * pq.ksub,
@@ -270,20 +274,24 @@ void compute_code(const ProductQuantizer& pq, const float* x, uint8_t* code) {
     }
 }
 
+} // namespace
+
 void ProductQuantizer::compute_code(const float* x, uint8_t* code) const {
-    switch (nbits) {
-        case 8:
-            faiss::compute_code<PQEncoder8>(*this, x, code);
-            break;
+    with_simd_level([&]<SIMDLevel SL>() {
+        switch (nbits) {
+            case 8:
+                compute_1_code<PQEncoder8, SL>(*this, x, code);
+                break;
 
-        case 16:
-            faiss::compute_code<PQEncoder16>(*this, x, code);
-            break;
+            case 16:
+                compute_1_code<PQEncoder16, SL>(*this, x, code);
+                break;
 
-        default:
-            faiss::compute_code<PQEncoderGeneric>(*this, x, code);
-            break;
-    }
+            default:
+                compute_1_code<PQEncoderGeneric, SL>(*this, x, code);
+                break;
+        }
+    }); // with_simd_level
 }
 
 template <class PQDecoder>
@@ -314,8 +322,9 @@ void ProductQuantizer::decode(const uint8_t* code, float* x) const {
 }
 
 void ProductQuantizer::decode(const uint8_t* code, float* x, size_t n) const {
+    int64_t n_signed = n;
 #pragma omp parallel for if (n > 100)
-    for (int64_t i = 0; i < n; i++) {
+    for (int64_t i = 0; i < n_signed; i++) {
         this->decode(code + code_size * i, x + d * i);
     }
 }
@@ -345,7 +354,8 @@ void ProductQuantizer::compute_codes_with_assign_index(
         const float* x,
         uint8_t* codes,
         size_t n) {
-    FAISS_THROW_IF_NOT(assign_index && assign_index->d == dsub);
+    FAISS_THROW_IF_NOT(
+            assign_index && static_cast<size_t>(assign_index->d) == dsub);
 
     for (size_t m = 0; m < M; m++) {
         assign_index->reset();
@@ -407,10 +417,11 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
         return;
     }
 
+    int64_t n_signed = n;
     if (dsub < 16) { // simple direct computation
 
 #pragma omp parallel for
-        for (int64_t i = 0; i < n; i++)
+        for (int64_t i = 0; i < n_signed; i++)
             compute_code(x + i * d, codes + i * code_size);
 
     } else { // worthwhile to use BLAS
@@ -418,7 +429,7 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
         compute_distance_tables(n, x, dis_tables.get());
 
 #pragma omp parallel for
-        for (int64_t i = 0; i < n; i++) {
+        for (int64_t i = 0; i < n_signed; i++) {
             uint8_t* code = codes + i * code_size;
             const float* tab = dis_tables.get() + i * ksub * M;
             compute_code_from_distance_table(tab, code);
@@ -428,51 +439,54 @@ void ProductQuantizer::compute_codes(const float* x, uint8_t* codes, size_t n)
 
 void ProductQuantizer::compute_distance_table(const float* x, float* dis_table)
         const {
-    if (transposed_centroids.empty()) {
-        // use regular version
-        for (size_t m = 0; m < M; m++) {
-            fvec_L2sqr_ny(
-                    dis_table + m * ksub,
-                    x + m * dsub,
-                    get_centroids(m, 0),
-                    dsub,
-                    ksub);
+    with_simd_level([&]<SIMDLevel SL>() {
+        if (transposed_centroids.empty()) {
+            // use regular version
+            for (size_t m = 0; m < M; m++) {
+                fvec_L2sqr_ny<SL>(
+                        dis_table + m * ksub,
+                        x + m * dsub,
+                        get_centroids(m, 0),
+                        dsub,
+                        ksub);
+            }
+        } else {
+            // transposed centroids are available, use'em
+            for (size_t m = 0; m < M; m++) {
+                fvec_L2sqr_ny_transposed<SL>(
+                        dis_table + m * ksub,
+                        x + m * dsub,
+                        transposed_centroids.data() + m * ksub,
+                        centroids_sq_lengths.data() + m * ksub,
+                        dsub,
+                        M * ksub,
+                        ksub);
+            }
         }
-    } else {
-        // transposed centroids are available, use'em
+    });
+}
+
+void ProductQuantizer::compute_inner_prod_table(
+        const float* x,
+        float* dis_table) const {
+    with_simd_level([&]<SIMDLevel SL>() {
         for (size_t m = 0; m < M; m++) {
-            fvec_L2sqr_ny_transposed(
+            fvec_inner_products_ny<SL>(
                     dis_table + m * ksub,
                     x + m * dsub,
-                    transposed_centroids.data() + m * ksub,
-                    centroids_sq_lengths.data() + m * ksub,
+                    get_centroids(m, 0),
                     dsub,
-                    M * ksub,
                     ksub);
         }
-    }
-}
-
-void ProductQuantizer::compute_inner_prod_table(
-        const float* x,
-        float* dis_table) const {
-    size_t m;
-
-    for (m = 0; m < M; m++) {
-        fvec_inner_products_ny(
-                dis_table + m * ksub,
-                x + m * dsub,
-                get_centroids(m, 0),
-                dsub,
-                ksub);
-    }
+    });
 }
 
 void ProductQuantizer::compute_distance_tables(
         size_t nx,
         const float* x,
         float* dis_tables) const {
-#if defined(__AVX2__) || defined(__aarch64__)
+    int64_t nx_signed = nx;
+#if defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_ARM_NEON)
     if (dsub == 2 && nbits < 8) { // interesting for a narrow range of settings
         compute_PQ_dis_tables_dsub2(
                 d, ksub, centroids.data(), nx, x, false, dis_tables);
@@ -481,13 +495,13 @@ void ProductQuantizer::compute_distance_tables(
             if (dsub < 16) {
 
 #pragma omp parallel for if (nx > 1)
-        for (int64_t i = 0; i < nx; i++) {
+        for (int64_t i = 0; i < nx_signed; i++) {
             compute_distance_table(x + i * d, dis_tables + i * ksub * M);
         }
 
     } else { // use BLAS
 
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             pairwise_L2sqr(
                     dsub,
                     nx,
@@ -506,7 +520,8 @@ void ProductQuantizer::compute_inner_prod_tables(
         size_t nx,
         const float* x,
         float* dis_tables) const {
-#if defined(__AVX2__) || defined(__aarch64__)
+    int64_t nx_signed = nx;
+#if defined(COMPILE_SIMD_AVX2) || defined(COMPILE_SIMD_ARM_NEON)
     if (dsub == 2 && nbits < 8) {
         compute_PQ_dis_tables_dsub2(
                 d, ksub, centroids.data(), nx, x, true, dis_tables);
@@ -515,14 +530,14 @@ void ProductQuantizer::compute_inner_prod_tables(
             if (dsub < 16) {
 
 #pragma omp parallel for if (nx > 1)
-        for (int64_t i = 0; i < nx; i++) {
+        for (int64_t i = 0; i < nx_signed; i++) {
             compute_inner_prod_table(x + i * d, dis_tables + i * ksub * M);
         }
 
     } else { // use BLAS
 
         // compute distance tables
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             FINTEGER ldc = ksub * M, nxi = nx, ksubi = ksub, dsubi = dsub,
                      di = d;
             float one = 1.0, zero = 0;
@@ -566,7 +581,7 @@ void pq_estimators_from_tables_Mmul4(
         float dis = 0;
         const float* dt = dis_table;
 
-        for (size_t m = 0; m < M; m += 4) {
+        for (int m = 0; m < M; m += 4) {
             float dism = 0;
             dism = dt[*codes++];
             dt += ksub;
@@ -638,7 +653,7 @@ void pq_estimators_from_tables(
     for (size_t j = 0; j < ncodes; j++) {
         float dis = 0;
         const float* __restrict dt = dis_table;
-        for (int m = 0; m < M; m++) {
+        for (size_t m = 0; m < M; m++) {
             dis += dt[*codes++];
             dt += ksub;
         }
@@ -686,10 +701,11 @@ void pq_knn_search_with_tables(
         HeapArray<C>* res,
         bool init_finalize_heap) {
     size_t k = res->k, nx = res->nh;
+    int64_t nx_signed = nx;
     size_t ksub = pq.ksub, M = pq.M;
 
 #pragma omp parallel for if (nx > 1)
-    for (int64_t i = 0; i < nx; i++) {
+    for (int64_t i = 0; i < nx_signed; i++) {
         /* query preparation for asymmetric search: compute look-up tables */
         const float* dis_table = dis_tables + i * ksub * M;
 
@@ -785,22 +801,24 @@ void ProductQuantizer::compute_sdc_table() {
     sdc_table.resize(M * ksub * ksub);
 
     if (dsub < 4) {
+        with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for
-        for (int mk = 0; mk < M * ksub; mk++) {
-            // allow omp to schedule in a more fine-grained way
-            // `collapse` is not supported in OpenMP 2.x
-            int m = mk / ksub;
-            int k = mk % ksub;
-            const float* cents = centroids.data() + m * ksub * dsub;
-            const float* centi = cents + k * dsub;
-            float* dis_tab = sdc_table.data() + m * ksub * ksub;
-            fvec_L2sqr_ny(dis_tab + k * ksub, centi, cents, dsub, ksub);
-        }
+            for (int64_t mk = 0; mk < static_cast<int64_t>(M * ksub); mk++) {
+                // allow omp to schedule in a more fine-grained way
+                // `collapse` is not supported in OpenMP 2.x
+                int m = mk / ksub;
+                int k = mk % ksub;
+                const float* cents = centroids.data() + m * ksub * dsub;
+                const float* centi = cents + k * dsub;
+                float* dis_tab = sdc_table.data() + m * ksub * ksub;
+                fvec_L2sqr_ny<SL>(dis_tab + k * ksub, centi, cents, dsub, ksub);
+            }
+        });
     } else {
         // NOTE: it would disable the omp loop in pairwise_L2sqr
         // but still accelerate especially when M >= 4
 #pragma omp parallel for
-        for (int m = 0; m < M; m++) {
+        for (int64_t m = 0; m < static_cast<int64_t>(M); m++) {
             const float* cents = centroids.data() + m * ksub * dsub;
             float* dis_tab = sdc_table.data() + m * ksub * ksub;
             pairwise_L2sqr(
@@ -819,9 +837,10 @@ void ProductQuantizer::search_sdc(
     FAISS_THROW_IF_NOT(sdc_table.size() == M * ksub * ksub);
     FAISS_THROW_IF_NOT(nbits == 8);
     size_t k = res->k;
+    int64_t nq_signed = nq;
 
 #pragma omp parallel for
-    for (int64_t i = 0; i < nq; i++) {
+    for (int64_t i = 0; i < nq_signed; i++) {
         /* Compute distances and keep smallest values */
         idx_t* heap_ids = res->ids + i * k;
         float* heap_dis = res->val + i * k;
@@ -834,7 +853,7 @@ void ProductQuantizer::search_sdc(
         for (size_t j = 0; j < nb; j++) {
             float dis = 0;
             const float* tab = sdc_table.data();
-            for (int m = 0; m < M; m++) {
+            for (size_t m = 0; m < M; m++) {
                 dis += tab[bcode[m] + qcode[m] * ksub];
                 tab += ksub * ksub;
             }

data/vendor/faiss/faiss/impl/Quantizer.h

@@ -16,8 +16,8 @@ struct Quantizer {
     size_t d;         ///< size of the input vectors
     size_t code_size; ///< bytes per indexed vector
 
-    explicit Quantizer(size_t d = 0, size_t code_size = 0)
-            : d(d), code_size(code_size) {}
+    explicit Quantizer(size_t d_in = 0, size_t code_size_in = 0)
+            : d(d_in), code_size(code_size_in) {}
 
     /** Train the quantizer
      *

data/vendor/faiss/faiss/impl/RaBitQUtils.cpp

@@ -8,9 +8,12 @@
 #include <faiss/impl/RaBitQUtils.h>
 
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/simd_dispatch.h>
 #include <faiss/utils/distances.h>
+#include <faiss/utils/rabitq_simd.h>
 #include <algorithm>
 #include <cmath>
+#include <cstring>
 #include <limits>
 
 namespace faiss {
@@ -151,6 +154,7 @@ QueryFactorsData compute_query_factors(
         std::vector<uint8_t>& rotated_qq) {
     FAISS_THROW_IF_NOT(qb <= 8);
     FAISS_THROW_IF_NOT(qb > 0);
+    FAISS_THROW_IF_NOT(d > 0);
 
     QueryFactorsData query_factors;
 
@@ -163,38 +167,42 @@ QueryFactorsData compute_query_factors(
     query_factors.g_error = std::sqrt(query_factors.qr_to_c_L2sqr);
 
     // Rotate the query (subtract centroid)
+    // Save aliasing state before resize(), which may reallocate the buffer.
+    const bool query_aliased = (query == rotated_q.data());
+    FAISS_THROW_IF_NOT_MSG(
+            !query_aliased || centroid == nullptr,
+            "query aliasing is only supported in the IVF residual path "
+            "(centroid == nullptr)");
     rotated_q.resize(d);
-    for (size_t i = 0; i < d; i++) {
-        if (i < rotated_q.size()) {
-            rotated_q[i] =
-                    query[i] - ((centroid == nullptr) ? 0.0f : centroid[i]);
+    if (centroid == nullptr) {
+        // Caller may pass query == rotated_q.data() (IVF residual path);
+        // memcpy with overlapping src/dst is UB, so skip the copy in that case.
+        if (!query_aliased) {
+            memcpy(rotated_q.data(), query, d * sizeof(float));
+        }
+    } else {
+        for (size_t i = 0; i < d; i++) {
+            rotated_q[i] = query[i] - centroid[i];
         }
     }
 
-    const float inv_d_sqrt =
-            (d == 0) ? 1.0f : (1.0f / std::sqrt(static_cast<float>(d)));
+    const float inv_d_sqrt = 1.0f / std::sqrt(static_cast<float>(d));
 
     // Compute quantization range
     float v_min = std::numeric_limits<float>::max();
     float v_max = std::numeric_limits<float>::lowest();
 
+    const float* rq = rotated_q.data();
     if (centered) {
         float z_max = Z_MAX_BY_QB[qb - 1];
         float v_radius = z_max * std::sqrt(query_factors.qr_to_c_L2sqr / d);
         v_min = -v_radius;
         v_max = v_radius;
     } else {
-        // Only compute min/max if we have dimensions to process
-        if (d > 0 && !rotated_q.empty()) {
-            for (size_t i = 0; i < d; i++) {
-                const float v_q = rotated_q[i];
-                v_min = std::min(v_min, v_q);
-                v_max = std::max(v_max, v_q);
-            }
-        } else {
-            // For empty dimensions, use default range
-            v_min = 0.0f;
-            v_max = 1.0f;
+        for (size_t i = 0; i < d; i++) {
+            const float v_q = rq[i];
+            v_min = std::min(v_min, v_q);
+            v_max = std::max(v_max, v_q);
         }
     }
 
@@ -207,25 +215,18 @@ QueryFactorsData compute_query_factors(
     size_t sum_qq = 0;
     int64_t sum2_signed_odd_int = 0;
 
-    // Process arrays - throw error if they are unexpectedly empty
-    if (d > 0 && !rotated_q.empty() && !rotated_qq.empty()) {
-        for (size_t i = 0; i < d; i++) {
-            const float v_q = rotated_q[i];
-            // Non-randomized scalar quantization
-            const uint8_t v_qq = std::clamp<float>(
-                    std::round((v_q - v_min) * inv_delta), 0, max_code);
-            rotated_qq[i] = v_qq;
-            sum_qq += v_qq;
-
-            if (centered) {
-                int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
-                sum2_signed_odd_int += signed_odd_int * signed_odd_int;
-            }
+    uint8_t* rqq = rotated_qq.data();
+    for (size_t i = 0; i < d; i++) {
+        const float v_q = rq[i];
+        const uint8_t v_qq = std::clamp<float>(
+                std::round((v_q - v_min) * inv_delta), 0, max_code);
+        rqq[i] = v_qq;
+        sum_qq += v_qq;
+
+        if (centered) {
+            int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
+            sum2_signed_odd_int += signed_odd_int * signed_odd_int;
         }
-    } else {
-        FAISS_THROW_MSG(
-                "Arrays unexpectedly empty when d=" + std::to_string(d) +
-                "or d is incorrectly set");
     }
 
     // Compute query factors
@@ -240,10 +241,18 @@ QueryFactorsData compute_query_factors(
         query_factors.int_dot_scale = 1.0f;
     }
 
-    // Compute query norm for inner product metric
+    // Compute query norm for inner product metric.
+    // When centroid is nullptr (IVF residual path), qr_to_c_L2sqr already
+    // holds fvec_norm_L2sqr(query, d) from line 164, so reuse it.
     query_factors.qr_norm_L2sqr = 0.0f;
+    query_factors.q_dot_c = 0.0f;
     if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
-        query_factors.qr_norm_L2sqr = fvec_norm_L2sqr(query, d);
+        query_factors.qr_norm_L2sqr = (centroid == nullptr)
+                ? query_factors.qr_to_c_L2sqr
+                : fvec_norm_L2sqr(query, d);
+        if (centroid != nullptr) {
+            query_factors.q_dot_c = fvec_inner_product(query, centroid, d);
+        }
     }
 
     return query_factors;
@@ -290,5 +299,94 @@ void set_bit_fastscan(uint8_t* code, size_t bit_index) {
     }
 }
 
+size_t compute_per_vector_storage_size(size_t nb_bits, size_t d) {
+    const size_t ex_bits = nb_bits - 1;
+    if (ex_bits == 0) {
+        return sizeof(SignBitFactors);
+    } else {
+        return sizeof(SignBitFactorsWithError) + sizeof(ExtraBitsFactors) +
+                (d * ex_bits + 7) / 8;
+    }
+}
+
+// Non-template wrapper with dynamic dispatch (one dispatch per call).
+// The hot path in RaBitQuantizer dispatches once at distance computer
+// construction, so per-vector dispatch only affects this utility path.
+float compute_full_multibit_distance(
+        const uint8_t* sign_bits,
+        const uint8_t* ex_code,
+        const ExtraBitsFactors& ex_fac,
+        const float* rotated_q,
+        float qr_base,
+        size_t d,
+        size_t ex_bits,
+        MetricType metric_type) {
+    return with_selected_simd_levels<AVAILABLE_SIMD_LEVELS_A0>(
+            [&]<SIMDLevel SL>() {
+                return compute_full_multibit_distance<SL>(
+                        sign_bits,
+                        ex_code,
+                        ex_fac,
+                        rotated_q,
+                        qr_base,
+                        d,
+                        ex_bits,
+                        metric_type);
+            });
+}
+
+void populate_block_aux_from_flat_storage(
+        const std::vector<uint8_t>& flat_storage,
+        AlignedTable<uint8_t>& codes,
+        size_t num_vectors,
+        size_t bbs,
+        size_t M2,
+        size_t old_block_stride,
+        size_t new_block_stride,
+        size_t storage_size,
+        const int64_t* id_map) {
+    if (flat_storage.empty() || num_vectors == 0) {
+        return;
+    }
+
+    const size_t packed_block_size = ((M2 + 1) / 2) * bbs;
+    const size_t n_blocks = (num_vectors + bbs - 1) / bbs;
+
+    if (old_block_stride < new_block_stride) {
+        AlignedTable<uint8_t> old_data;
+        old_data.resize(codes.size());
+        memcpy(old_data.data(), codes.data(), codes.size());
+
+        codes.resize(n_blocks * new_block_stride);
+        memset(codes.data(), 0, n_blocks * new_block_stride);
+        for (size_t b = 0; b < n_blocks; b++) {
+            memcpy(codes.data() + b * new_block_stride,
+                   old_data.data() + b * old_block_stride,
+                   packed_block_size);
+        }
+    }
+
+    for (size_t offset = 0; offset < num_vectors; offset++) {
+        const int64_t global_id =
+                id_map ? id_map[offset] : static_cast<int64_t>(offset);
+        FAISS_THROW_IF_NOT_MSG(
+                global_id >= 0 &&
+                        static_cast<size_t>(global_id) * storage_size +
+                                        storage_size <=
+                                flat_storage.size(),
+                "global_id out of bounds for flat_storage during migration");
+
+        const uint8_t* src = flat_storage.data() + global_id * storage_size;
+        uint8_t* dst = get_block_aux_ptr(
+                codes.data(),
+                offset,
+                bbs,
+                packed_block_size,
+                new_block_stride,
+                storage_size);
+        memcpy(dst, src, storage_size);
+    }
+}
+
 } // namespace rabitq_utils
 } // namespace faiss