Hoist SIMD dispatch outside loops in 4 call sites (#5074)

Summary:
Pull Request resolved: #5074

Move `with_simd_level` / `with_simd_level_256bit` calls outside the
enclosing loops so the SIMD level is resolved once rather than on every
iteration.

Sites fixed:
- distances.cpp: knn_inner_products_by_idx, knn_L2sqr_by_idx
- NeuralNet.cpp: ZnLUTCodec::encode
- ClusteringInitialization.cpp: init_kmpp_plus_plus

Reviewed By: mdouze

Differential Revision: D100144174

fbshipit-source-id: bd2369ed4fd9c3b5b54e435c7ee66a03f0e152df

Author: algoriddle

faiss/impl/ClusteringInitialization.cpp

@@ -221,27 +221,27 @@ void ClusteringInitialization::init_kmeans_plus_plus(
     std::vector<double> cumsum(n);
 
     // Select remaining centroids using D² sampling
-    for (size_t c = result.first_new_centroid_idx; c < k; c++) {
-        // Compute cumulative sum
-        cumsum[0] = min_distances[0];
-        for (size_t i = 1; i < n; i++) {
-            cumsum[i] = cumsum[i - 1] + min_distances[i];
-        }
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (size_t c = result.first_new_centroid_idx; c < k; c++) {
+            // Compute cumulative sum
+            cumsum[0] = min_distances[0];
+            for (size_t i = 1; i < n; i++) {
+                cumsum[i] = cumsum[i - 1] + min_distances[i];
+            }
 
-        // Sample using precomputed cumsum
-        size_t next_idx = sample_from_cumsum(cumsum, rng);
+            // Sample using precomputed cumsum
+            size_t next_idx = sample_from_cumsum(cumsum, rng);
 
-        float* new_centroid = centroids + c * d;
-        std::memcpy(new_centroid, x + next_idx * d, d * sizeof(float));
+            float* new_centroid = centroids + c * d;
+            std::memcpy(new_centroid, x + next_idx * d, d * sizeof(float));
 
-        // Update min distances incrementally
-        with_simd_level([&]<SIMDLevel SL>() {
+            // Update min distances incrementally
             for (size_t i = 0; i < n; i++) {
                 double dist = fvec_L2sqr<SL>(x + i * d, new_centroid, d);
                 min_distances[i] = std::min(min_distances[i], dist);
             }
-        });
-    }
+        }
+    });
 }
 
 void ClusteringInitialization::init_afkmc2(

faiss/utils/NeuralNet.cpp

@@ -268,12 +268,12 @@ nn::Int32Tensor2D QINCoStep::encode(
         res = residuals->data();
     }
 
-    for (size_t i = 0; i < n; i++) {
-        const float* q = x.data() + i * d;
-        const float* db = zqs_r.data() + i * K * d;
-        float dis_min = HUGE_VALF;
-        int64_t idx = -1;
-        with_simd_level([&]<SIMDLevel SL>() {
+    with_simd_level([&]<SIMDLevel SL>() {
+        for (size_t i = 0; i < n; i++) {
+            const float* q = x.data() + i * d;
+            const float* db = zqs_r.data() + i * K * d;
+            float dis_min = HUGE_VALF;
+            int64_t idx = -1;
             for (size_t j = 0; j < static_cast<size_t>(K); j++) {
                 float dis = fvec_L2sqr<SL>(q, db, d);
                 if (dis < dis_min) {
@@ -282,17 +282,17 @@ nn::Int32Tensor2D QINCoStep::encode(
                 }
                 db += d;
             }
-        });
-        codes.v[i] = idx;
-        if (res) {
-            const float* xhat_row = xhat.data() + i * d;
-            const float* xhat_next_row = zqs_r.data() + (i * K + idx) * d;
-            for (size_t j = 0; j < static_cast<size_t>(d); j++) {
-                res[j] = xhat_next_row[j] - xhat_row[j];
+            codes.v[i] = idx;
+            if (res) {
+                const float* xhat_row = xhat.data() + i * d;
+                const float* xhat_next_row = zqs_r.data() + (i * K + idx) * d;
+                for (size_t j = 0; j < static_cast<size_t>(d); j++) {
+                    res[j] = xhat_next_row[j] - xhat_row[j];
+                }
+                res += d;
             }
-            res += d;
         }
-    }
+    });
     return codes;
 }
 

faiss/utils/distances.cpp

@@ -838,16 +838,16 @@ void knn_inner_products_by_idx(
         ld_ids = ny;
     }
 
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for if (nx > 100)
-    for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
-        const float* x_ = x + i * d;
-        const int64_t* idsi = ids + i * ld_ids;
-        size_t j;
-        float* __restrict simi = res_vals + i * k;
-        int64_t* __restrict idxi = res_ids + i * k;
-        minheap_heapify(k, simi, idxi);
+        for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
+            const float* x_ = x + i * d;
+            const int64_t* idsi = ids + i * ld_ids;
+            size_t j;
+            float* __restrict simi = res_vals + i * k;
+            int64_t* __restrict idxi = res_ids + i * k;
+            minheap_heapify(k, simi, idxi);
 
-        with_simd_level([&]<SIMDLevel SL>() {
             for (j = 0; j < nsubset; j++) {
                 if (idsi[j] < 0 || static_cast<size_t>(idsi[j]) >= ny) {
                     break;
@@ -858,9 +858,9 @@ void knn_inner_products_by_idx(
                     minheap_replace_top(k, simi, idxi, ip, idsi[j]);
                 }
             }
-        });
-        minheap_reorder(k, simi, idxi);
-    }
+            minheap_reorder(k, simi, idxi);
+        }
+    });
 }
 
 void knn_L2sqr_by_idx(
@@ -878,14 +878,14 @@ void knn_L2sqr_by_idx(
     if (ld_ids < 0) {
         ld_ids = ny;
     }
+    with_simd_level([&]<SIMDLevel SL>() {
 #pragma omp parallel for if (nx > 100)
-    for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
-        const float* x_ = x + i * d;
-        const int64_t* __restrict idsi = ids + i * ld_ids;
-        float* __restrict simi = res_vals + i * k;
-        int64_t* __restrict idxi = res_ids + i * k;
-        maxheap_heapify(k, simi, idxi);
-        with_simd_level([&]<SIMDLevel SL>() {
+        for (int64_t i = 0; i < static_cast<int64_t>(nx); i++) {
+            const float* x_ = x + i * d;
+            const int64_t* __restrict idsi = ids + i * ld_ids;
+            float* __restrict simi = res_vals + i * k;
+            int64_t* __restrict idxi = res_ids + i * k;
+            maxheap_heapify(k, simi, idxi);
             for (size_t j = 0; j < nsubset; j++) {
                 if (idsi[j] < 0 || static_cast<size_t>(idsi[j]) >= ny) {
                     break;
@@ -896,9 +896,9 @@ void knn_L2sqr_by_idx(
                     maxheap_replace_top(k, simi, idxi, disij, idsi[j]);
                 }
             }
-        });
-        maxheap_reorder(k, simi, idxi);
-    }
+            maxheap_reorder(k, simi, idxi);
+        }
+    });
 }
 
 void pairwise_L2sqr(