Partial Optimization for IVFPQ

Author: littleniuer

faiss/CMakeLists.txt

@@ -335,6 +335,45 @@ if(FAISS_ENABLE_SVS)
   )
 endif()
 
+# -------------------- ARM aarch64 sra_krl sources --------------------
+# On aarch64, add sra_krl sources. Source code uses #ifdef __aarch64__.
+#
+# These .c files contain C++ features (template, constexpr, #include <cstring>)
+# so they MUST be compiled as C++. The public API headers (krl.h) already wrap
+# function declarations with extern "C" { }, ensuring C linkage for symbols.
+set(KRL_SRC)
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
+  message(STATUS "ARM aarch64 detected: adding sra_krl sources")
+
+  set(KRL_SRC
+    sra_krl/src/IPdistance_simd_f16f32.c
+    sra_krl/src/IPdistance_simd.c
+    sra_krl/src/IPdistance_simd_s8.c
+    sra_krl/src/krl_handles.c
+    sra_krl/src/L2distance_simd_f16f32.c
+    sra_krl/src/L2distance_simd.c
+    sra_krl/src/L2distance_simd_u8.c
+    sra_krl/src/matrix_block_transpose.c
+    sra_krl/src/MinMax_quant.c
+    sra_krl/src/pq_search_with_table_8bit.c
+  )
+
+
+  # Force C++ compilation: these .c files use template, constexpr, <cstring>
+  set_source_files_properties(${KRL_SRC} PROPERTIES LANGUAGE CXX)
+
+  set(KRL_COMPILE_FLAGS "-march=armv8.2-a+fp16+fp16fml+dotprod+sve+rcpc")
+
+  if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang|AppleClang")
+    set_source_files_properties(${KRL_SRC} PROPERTIES
+      COMPILE_FLAGS "${KRL_COMPILE_FLAGS}"
+    )
+  endif()
+
+  # Add KRL sources to main FAISS_SRC
+  list(APPEND FAISS_SRC ${KRL_SRC})
+endif()
+
 if(NOT WIN32)
   list(APPEND FAISS_SRC invlists/OnDiskInvertedLists.cpp)
   list(APPEND FAISS_HEADERS invlists/OnDiskInvertedLists.h)
@@ -345,6 +384,14 @@ set(FAISS_HEADERS ${FAISS_HEADERS} PARENT_SCOPE)
 
 add_library(faiss ${FAISS_SRC})
 
+# -------------------- aarch64 sra_krl include dirs for faiss --------------------
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
+  target_include_directories(faiss PUBLIC
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/sra_krl/include>
+    $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}/faiss/sra_krl>
+  )
+endif()
+
 add_library(faiss_avx2 ${FAISS_SRC})
 if(NOT FAISS_OPT_LEVEL STREQUAL "avx2" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512" AND NOT FAISS_OPT_LEVEL STREQUAL "avx512_spr")
   set_target_properties(faiss_avx2 PROPERTIES EXCLUDE_FROM_ALL TRUE)
@@ -422,6 +469,14 @@ endif()
 target_sources(faiss_sve PRIVATE ${FAISS_SIMD_NEON_SRC} ${FAISS_SIMD_SVE_SRC})
 target_compile_definitions(faiss_sve PRIVATE COMPILE_SIMD_ARM_NEON COMPILE_SIMD_ARM_SVE)
 
+# -------------------- aarch64 sra_krl include dirs for faiss_sve --------------------
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
+  target_include_directories(faiss_sve PUBLIC
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/sra_krl/include>
+    $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}/faiss/sra_krl>
+  )
+endif()
+
 # =============================================================================
 # Dynamic Dispatch Mode
 # When FAISS_OPT_LEVEL=dd, the main faiss library is built with runtime SIMD
@@ -616,6 +671,13 @@ endif()
 # Note: FAISS_OPT_LEVEL=dd builds DD support into main faiss target,
 # so no separate faiss_dd install is needed
 
+# -------------------- Install KRL headers (aarch64 only) --------------------
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "(aarch64|arm64|ARM64)")
+  install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/sra_krl/include/
+          DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/faiss/sra_krl/include
+          FILES_MATCHING PATTERN "*.h")
+endif()
+
 foreach(header ${FAISS_HEADERS})
   get_filename_component(dir ${header} DIRECTORY )
   install(FILES ${header}

faiss/IndexIVF.cpp

@@ -28,6 +28,12 @@
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/expanded_scanners.h>
 
+#ifdef __aarch64__
+extern "C" {
+#include <faiss/sra_krl/include/krl.h>
+}
+#endif
+
 namespace faiss {
 
 using ScopedIds = InvertedLists::ScopedIds;
@@ -464,6 +470,9 @@ void IndexIVF::search_preassigned(
         std::unique_ptr<InvertedListScanner> scanner(
                 get_InvertedListScanner(store_pairs, sel, params));
 
+#ifdef __aarch64__
+        krl_create_LUT8b_handle(&(scanner->klh),(int)(sel != nullptr), tmp_buffer_size);
+#endif
         /*****************************************************
          * Depending on parallel_mode, there are two possible ways
          * to organize the search. Here we define local functions

faiss/IndexIVF.h

@@ -20,6 +20,14 @@
 #include <faiss/invlists/InvertedLists.h>
 #include <faiss/utils/Heap.h>
 
+#ifdef __aarch64__
+extern "C" {
+    typedef struct KRLLookupTable8bitHandle KRLLUT8bHandle;
+    void krl_clean_LUT8b_handle(KRLLUT8bHandle** klh);
+}
+#endif
+
+
 namespace faiss {
 
 /** Encapsulates a quantizer object for the IndexIVF
@@ -198,6 +206,10 @@ struct IndexIVF : Index, IndexIVFInterface {
     /// centroids?
     bool by_residual = true;
 
+#ifdef __aarch64__
+    size_t tmp_buffer_size = 0;
+#endif
+
     /** The Inverted file takes a quantizer (an Index) on input,
      * which implements the function mapping a vector to a list
      * identifier.
@@ -553,7 +565,17 @@ struct InvertedListScanner {
             const idx_t* ids,
             ResultHandler& handler) const;
 
+#ifdef __aarch64__
+    KRLLUT8bHandle* klh = nullptr;
+    virtual ~InvertedListScanner() {
+        if(klh){
+            krl_clean_LUT8b_handle(&klh);
+            klh = nullptr;
+        }
+    }
+#else
     virtual ~InvertedListScanner() {}
+#endif
 };
 
 // whether to check that coarse quantizers are the same

faiss/IndexIVFPQ.cpp

@@ -33,6 +33,12 @@
 #include <faiss/impl/pq_code_distance/pq_code_distance-inl.h>
 #include <faiss/impl/simd_dispatch.h>
 
+#ifdef __aarch64__
+extern "C" {
+#include <faiss/sra_krl/include/krl.h>
+}
+#endif
+
 namespace faiss {
 
 /*****************************************
@@ -1229,7 +1235,46 @@ struct IVFPQScanner : IVFPQScannerT<idx_t, METRIC_TYPE, PQCodeDist>,
             assert(precompute_mode == 2);
             this->scan_list_polysemous(ncode, codes, res);
         } else if (precompute_mode == 2) {
+#ifdef __aarch64__
+            if (this->pq.nbits == 8 && this->klh) {
+				if constexpr (use_sel) {
+                    size_t j = 0;
+                    size_t* idx_tmp_buffer = krl_get_idx_pointer(klh);
+                    float* distance_tmp_buffer = krl_get_dist_pointer(klh);
+                    for(size_t i = 0; i < ncode; ++i) {
+                        if(this->sel->is_member(ids[i])){
+                            idx_tmp_buffer[j++] = i;
+                        }
+                    }
+                    ncode = j;
+                    krl_table_lookup_8b_f32_by_idx(
+                        this->pq.M, ncode, codes, this->sim_table, distance_tmp_buffer, this->dis0, idx_tmp_buffer, 
+						this->pq.M * ncode, this->pq.M * 256, ncode);
+                    for(size_t i = 0; i < ncode; ++i) {
+                        res.add(idx_tmp_buffer[i],distance_tmp_buffer[i]);
+                    }
+                } else {
+					float* distance_tmp_buffer = krl_get_dist_pointer(klh);
+					krl_table_lookup_8b_f32(
+						this->pq.M, ncode, codes, this->sim_table, distance_tmp_buffer, this->dis0, this->pq.M * ncode, 
+						this->pq.M * 256, ncode);
+					size_t i = 0;
+					for(; i + 4 <= ncode; i+=4) {
+						res.add(i, distance_tmp_buffer[i]);
+						res.add(i + 1, distance_tmp_buffer[i + 1]);
+						res.add(i + 2, distance_tmp_buffer[i + 2]);
+						res.add(i + 3, distance_tmp_buffer[i + 3]);
+					}
+					for(; i < ncode; ++i) {
+						res.add(i, distance_tmp_buffer[i]);
+					}
+                }
+            } else {
+				this->scan_list_with_table(ncode, codes, res);
+            }
+#else 
             this->scan_list_with_table(ncode, codes, res);
+#endif
         } else if (precompute_mode == 1) {
             this->scan_list_with_pointer(ncode, codes, res);
         } else if (precompute_mode == 0) {

faiss/impl/ProductQuantizer.cpp

@@ -24,6 +24,10 @@
 
 extern "C" {
 
+#ifdef __aarch64__
+#include <faiss/sra_krl/include/krl.h>
+#endif
+
 /* declare BLAS functions, see http://www.netlib.org/clapack/cblas/ */
 
 int sgemm_(
@@ -436,6 +440,12 @@ 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 {
+#ifdef __aarch64__
+    if(use_transpose) {
+        krl_L2sqr_ny_with_handle(kdh, dis_table, x, M * ksub, dsub * M);
+        return;
+    }
+#endif
     with_simd_level([&]<SIMDLevel SL>() {
         if (transposed_centroids.empty()) {
             // use regular version
@@ -466,6 +476,12 @@ void ProductQuantizer::compute_distance_table(const float* x, float* dis_table)
 void ProductQuantizer::compute_inner_prod_table(
         const float* x,
         float* dis_table) const {
+#ifdef __aarch64__
+    if(use_transpose) {
+        krl_inner_product_ny_with_handle(kdh, dis_table, x, M * ksub, dsub * M);
+        return;
+    }
+#endif
     with_simd_level([&]<SIMDLevel SL>() {
         for (size_t m = 0; m < M; m++) {
             fvec_inner_products_ny<SL>(
@@ -482,6 +498,15 @@ void ProductQuantizer::compute_distance_tables(
         size_t nx,
         const float* x,
         float* dis_tables) const {
+#ifdef __aarch64__
+    if(use_transpose) {
+    	#pragma omp parallel for if (nx > 1)
+        for (int64_t i = 0; i < nx; i++) {
+            krl_L2sqr_ny_with_handle(kdh, dis_tables + i * ksub * M ,x + i * d, M * ksub, dsub * M);
+        }
+        return;
+    }
+#endif
 #if defined(__AVX2__) || defined(__aarch64__)
     if (dsub == 2 && nbits < 8) { // interesting for a narrow range of settings
         compute_PQ_dis_tables_dsub2(
@@ -516,6 +541,15 @@ void ProductQuantizer::compute_inner_prod_tables(
         size_t nx,
         const float* x,
         float* dis_tables) const {
+#ifdef __aarch64__
+    if(use_transpose) {
+		#pragma omp parallel for if (nx > 1)
+        for (int64_t i = 0; i < nx; i++) {
+			krl_inner_product_ny_with_handle(kdh, dis_tables + i * ksub * M ,x + i * d, M * ksub, dsub * M);
+        }
+        return;
+    }
+#endif
 #if defined(__AVX2__) || defined(__aarch64__)
     if (dsub == 2 && nbits < 8) {
         compute_PQ_dis_tables_dsub2(

faiss/impl/ProductQuantizer.h

@@ -18,6 +18,13 @@
 #include <faiss/impl/Quantizer.h>
 #include <faiss/impl/platform_macros.h>
 #include <faiss/utils/Heap.h>
+#include <iostream>
+
+#ifdef __aarch64__
+extern "C" {
+#include <faiss/sra_krl/include/krl.h>
+}
+#endif
 
 namespace faiss {
 
@@ -51,6 +58,67 @@ struct ProductQuantizer : Quantizer {
     /// d / M)
     Index* assign_index;
 
+#ifdef __aarch64__
+    bool use_transpose = false;
+    KRLDistanceHandle* kdh = nullptr; 
+
+    ProductQuantizer(const ProductQuantizer& other)
+            : Quantizer(other),
+              M(other.M),
+              nbits(other.nbits),
+              dsub(other.dsub),
+              ksub(other.ksub),
+              verbose(other.verbose),
+              train_type(other.train_type),
+              cp(other.cp),
+              assign_index(other.assign_index),
+              use_transpose(false),
+              kdh(nullptr),
+              centroids(other.centroids),
+              transposed_centroids(other.transposed_centroids),
+              centroids_sq_lengths(other.centroids_sq_lengths),
+              sdc_table(other.sdc_table) {}
+
+    ProductQuantizer& operator=(const ProductQuantizer& other) {
+        if (this != &other) {
+            Quantizer::operator=(other);
+            M = other.M;
+            nbits = other.nbits;
+            dsub = other.dsub;
+            ksub = other.ksub;
+            verbose = other.verbose;
+            train_type = other.train_type;
+            cp = other.cp;
+            assign_index = other.assign_index;
+            centroids = other.centroids;
+            transposed_centroids = other.transposed_centroids;
+            centroids_sq_lengths = other.centroids_sq_lengths;
+            sdc_table = other.sdc_table;
+            if (kdh) {
+                krl_clean_distance_handle(&kdh);
+                kdh = nullptr;
+            }
+            use_transpose = false;
+        }
+        return *this;
+    }
+
+    void initialize_krl_transpose_centroids(size_t batchsize, int metric_type) {
+        if(!kdh && !centroids.empty()) {
+            krl_create_distance_handle(
+                &kdh, 3, batchsize, ksub, dsub, M, metric_type, (const uint8_t *)centroids.data(), M * ksub * dsub * sizeof(float));
+			use_transpose = (kdh != nullptr);
+        }
+    }
+    ~ProductQuantizer() {
+        if(kdh) {
+            krl_clean_distance_handle(&kdh);
+            kdh = nullptr;
+			use_transpose = false;			
+        }
+    }
+#endif
+
     /// Centroid table, size M * ksub * dsub.
     /// Layout: (M, ksub, dsub)
     std::vector<float> centroids;