faiss 0.6.0 → 0.6.2

data/vendor/faiss/faiss/impl/pq_code_distance/pq_scan_impl.h

@@ -0,0 +1,105 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#ifndef THE_SIMD_LEVEL
+#error "THE_SIMD_LEVEL must be defined before including pq_scan_impl.h"
+#endif
+
+#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
+#include <faiss/utils/Heap.h>
+
+namespace faiss {
+namespace pq_code_distance {
+
+namespace {
+
+template <class C>
+void pq_scan_8bit_at_level(
+        size_t M,
+        const float* dis_table,
+        const uint8_t* codes,
+        size_t ncodes,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids) {
+    const size_t code_size = M;
+
+    size_t j = 0;
+    for (; j + 3 < ncodes; j += 4) {
+        float d0, d1, d2, d3;
+        pq_code_distance_8bit_four_impl<THE_SIMD_LEVEL>( // NOLINT(facebook-modularize-issue-check)
+                M,
+                dis_table,
+                codes + j * code_size,
+                codes + (j + 1) * code_size,
+                codes + (j + 2) * code_size,
+                codes + (j + 3) * code_size,
+                d0,
+                d1,
+                d2,
+                d3);
+
+        if (C::cmp(heap_dis[0], d0)) {
+            heap_replace_top<C>(k, heap_dis, heap_ids, d0, j);
+        }
+        if (C::cmp(heap_dis[0], d1)) {
+            heap_replace_top<C>(k, heap_dis, heap_ids, d1, j + 1);
+        }
+        if (C::cmp(heap_dis[0], d2)) {
+            heap_replace_top<C>(k, heap_dis, heap_ids, d2, j + 2);
+        }
+        if (C::cmp(heap_dis[0], d3)) {
+            heap_replace_top<C>(k, heap_dis, heap_ids, d3, j + 3);
+        }
+    }
+
+    for (; j < ncodes; j++) {
+        float dis = pq_code_distance_8bit_single_impl<THE_SIMD_LEVEL>( // NOLINT(facebook-modularize-issue-check)
+                M, dis_table, codes + j * code_size);
+        if (C::cmp(heap_dis[0], dis)) {
+            heap_replace_top<C>(k, heap_dis, heap_ids, dis, j);
+        }
+    }
+}
+
+} // anonymous namespace
+
+template <SIMDLevel SL>
+void pq_scan_8bit_impl(
+        size_t M,
+        const float* dis_table,
+        const uint8_t* codes,
+        size_t ncodes,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids,
+        bool max_heap);
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization,facebook-modularize-issue-check)
+template <>
+void pq_scan_8bit_impl<THE_SIMD_LEVEL>(
+        size_t M,
+        const float* dis_table,
+        const uint8_t* codes,
+        size_t ncodes,
+        size_t k,
+        float* heap_dis,
+        int64_t* heap_ids,
+        bool max_heap) {
+    if (max_heap) {
+        pq_scan_8bit_at_level<CMax<float, int64_t>>(
+                M, dis_table, codes, ncodes, k, heap_dis, heap_ids);
+    } else {
+        pq_scan_8bit_at_level<CMin<float, int64_t>>(
+                M, dis_table, codes, ncodes, k, heap_dis, heap_ids);
+    }
+}
+
+} // namespace pq_code_distance
+} // namespace faiss

data/vendor/faiss/faiss/impl/pq_code_distance/rvv.cpp

@@ -0,0 +1,70 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#ifdef COMPILE_SIMD_RISCV_RVV
+
+#include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
+
+namespace faiss {
+namespace pq_code_distance {
+
+// RISCV_RVV: no RVV-optimized PQ code distance exists yet. Use scalar.
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+float pq_code_distance_8bit_single_impl<SIMDLevel::RISCV_RVV>(
+        size_t M,
+        const float* sim_table,
+        const uint8_t* code) {
+    return PQCodeDistanceScalar<PQDecoder8>::distance_single_code(
+            M, 8, sim_table, code);
+}
+
+// NOLINTNEXTLINE(facebook-hte-MisplacedTemplateSpecialization)
+template <>
+void pq_code_distance_8bit_four_impl<SIMDLevel::RISCV_RVV>(
+        size_t M,
+        const float* sim_table,
+        const uint8_t* __restrict code0,
+        const uint8_t* __restrict code1,
+        const uint8_t* __restrict code2,
+        const uint8_t* __restrict code3,
+        float& result0,
+        float& result1,
+        float& result2,
+        float& result3) {
+    PQCodeDistanceScalar<PQDecoder8>::distance_four_codes(
+            M,
+            8,
+            sim_table,
+            code0,
+            code1,
+            code2,
+            code3,
+            result0,
+            result1,
+            result2,
+            result3);
+}
+
+} // namespace pq_code_distance
+} // namespace faiss
+
+#define THE_SIMD_LEVEL SIMDLevel::RISCV_RVV
+
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/pq_code_distance/pq_scan_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/utils/hamming_distance/hamming_computer-rvv.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/pq_code_distance/PQDistanceComputer_impl.h>
+// NOLINTNEXTLINE(facebook-hte-InlineHeader)
+#include <faiss/impl/pq_code_distance/IVFPQScanner_impl.h>
+
+#undef THE_SIMD_LEVEL
+
+#endif // COMPILE_SIMD_RISCV_RVV