enable SIMD, memory pool

Signed-off-by: eric-epsilla <eric@epsilla.com>

Author: eric-epsilla

VERSION

@@ -1 +1 @@
-0.3.22
+0.3.23

engine/CMakeLists.txt

@@ -68,6 +68,8 @@ list(APPEND LIB_FILES ${UTILS_FILES})
 file(GLOB UTILS_FILES "db/execution/*.cpp")
 list(APPEND LIB_FILES ${UTILS_FILES})
 file(GLOB UTILS_FILES "db/index/*.cpp")
+# 排除有问题的 distances.cpp，使用 distance_simd.cpp
+list(REMOVE_ITEM UTILS_FILES "${CMAKE_CURRENT_SOURCE_DIR}/db/index/distances.cpp")
 list(APPEND LIB_FILES ${UTILS_FILES})
 file(GLOB UTILS_FILES "db/index/nsg/*.cpp")
 list(APPEND LIB_FILES ${UTILS_FILES})
@@ -311,3 +313,27 @@ if(CLANG_TIDY_PATH)
     DEPENDS vector_db_test
   )
 endif()
+
+# ============================================================
+# Phase 1 优化: SIMD 距离计算
+# ============================================================
+include(CheckCXXCompilerFlag)
+
+# 检测 AVX2 支持
+check_cxx_compiler_flag("-mavx2" COMPILER_SUPPORTS_AVX2)
+if(COMPILER_SUPPORTS_AVX2)
+    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mavx2 -mfma")
+    add_definitions(-DUSE_AVX2)
+    message(STATUS "✓ AVX2 SIMD support enabled")
+endif()
+
+# 检测 ARM NEON 支持
+if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64|arm64|ARM64")
+    add_definitions(-DUSE_NEON)
+    message(STATUS "✓ ARM NEON SIMD support enabled")
+endif()
+
+# 使用 SIMD 优化的距离计算
+add_definitions(-DUSE_SIMD_OPTIMIZED)
+message(STATUS "✓ SIMD optimized distance calculation enabled")
+

engine/db/execution/vec_search_executor.cpp

@@ -1,5 +1,6 @@
 #include "db/execution/vec_search_executor.hpp"
 #include "utils/safe_memory_ops.hpp"
+#include "utils/memory_pool.hpp"
 
 #include <omp.h>
 
@@ -68,6 +69,31 @@ VecSearchExecutor::VecSearchExecutor(
       prefilter_enabled_(prefilter_enabled) {
   ann_index_ = ann_index;
 
+  // Initialize memory pool (singleton, only initialized once)
+  static bool memory_pool_initialized = false;
+  if (!memory_pool_initialized) {
+    utils::MemoryPoolConfig config;
+    config.initial_size = 128 * 1024 * 1024;  // 128MB
+    config.max_size = 1024 * 1024 * 1024;     // 1GB
+    config.enable_thread_cache = true;
+    config.enable_stats = true;
+    utils::MemoryPool::GetInstance().Initialize(config);
+    memory_pool_initialized = true;
+  }
+  
+  // Pre-allocate vectors to reduce runtime allocations
+  // These are already sized in the initializer list, but ensure capacity
+  search_result_.reserve(L_master_);
+  distance_.reserve(L_master_);
+  init_ids_.reserve(L_master_);
+  is_visited_.reserve(ann_index->record_number_);
+  set_L_.reserve((num_threads - 1) * L_local + L_master);
+  local_queues_sizes_.reserve(num_threads);
+  local_queues_starts_.reserve(num_threads);
+  if (brute_force_search_) {
+    brute_force_queue_.reserve(ann_index->record_number_);
+  }
+  
   // Log thread configuration for debugging - only in debug builds or when explicitly enabled
 #ifdef VECTORDB_DEBUG_BUILD
   static std::atomic<int> executor_count(0);

engine/db/index/distances.cpp

@@ -5,6 +5,7 @@
 #include <cmath>
 #include <cstdio>
 #include <cstring>
+#include <memory>
 
 
 #ifdef __AVX2__

engine/db/index/simd_config.hpp

@@ -0,0 +1,25 @@
+#pragma once
+
+// SIMD 配置
+#if defined(USE_AVX2)
+    #define SIMD_WIDTH 8
+    #define SIMD_TYPE "AVX2"
+#elif defined(USE_NEON)
+    #define SIMD_WIDTH 4
+    #define SIMD_TYPE "NEON"
+#else
+    #define SIMD_WIDTH 1
+    #define SIMD_TYPE "Scalar"
+#endif
+
+namespace vectordb {
+namespace simd {
+    inline const char* GetSIMDType() {
+        return SIMD_TYPE;
+    }
+    
+    inline int GetSIMDWidth() {
+        return SIMD_WIDTH;
+    }
+}
+}

engine/db/vector_normalizer.cpp

@@ -0,0 +1,59 @@
+#include "db/batch_insertion_optimizer.hpp"
+
+#ifdef __AVX2__
+#include <immintrin.h>
+
+namespace vectordb {
+namespace engine {
+namespace db {
+
+void VectorNormalizer::NormalizeBatchAVX2(float* vectors, size_t count, size_t dimension) {
+    for (size_t i = 0; i < count; ++i) {
+        float* vec = vectors + i * dimension;
+        
+        // Compute norm using AVX2
+        __m256 sum_vec = _mm256_setzero_ps();
+        size_t j = 0;
+        
+        for (; j + 7 < dimension; j += 8) {
+            __m256 v = _mm256_loadu_ps(vec + j);
+            sum_vec = _mm256_fmadd_ps(v, v, sum_vec);
+        }
+        
+        // Horizontal sum
+        __m128 sum_high = _mm256_extractf128_ps(sum_vec, 1);
+        __m128 sum_low = _mm256_castps256_ps128(sum_vec);
+        __m128 sum128 = _mm_add_ps(sum_high, sum_low);
+        sum128 = _mm_hadd_ps(sum128, sum128);
+        sum128 = _mm_hadd_ps(sum128, sum128);
+        
+        float sum = _mm_cvtss_f32(sum128);
+        
+        // Handle remaining elements
+        for (; j < dimension; ++j) {
+            sum += vec[j] * vec[j];
+        }
+        
+        // Normalize using AVX2
+        float inv_norm = 1.0f / std::sqrt(sum);
+        __m256 inv_norm_vec = _mm256_set1_ps(inv_norm);
+        
+        j = 0;
+        for (; j + 7 < dimension; j += 8) {
+            __m256 v = _mm256_loadu_ps(vec + j);
+            v = _mm256_mul_ps(v, inv_norm_vec);
+            _mm256_storeu_ps(vec + j, v);
+        }
+        
+        // Handle remaining elements
+        for (; j < dimension; ++j) {
+            vec[j] *= inv_norm;
+        }
+    }
+}
+
+} // namespace db
+} // namespace engine
+} // namespace vectordb
+
+#endif // __AVX2__

engine/utils/memory_pool.hpp

@@ -69,6 +69,11 @@ struct MemoryPoolStats {
     std::atomic<size_t> cache_hits{0};
     std::atomic<size_t> cache_misses{0};
 
+    // Delete copy constructor and assignment operator
+    MemoryPoolStats() = default;
+    MemoryPoolStats(const MemoryPoolStats&) = delete;
+    MemoryPoolStats& operator=(const MemoryPoolStats&) = delete;
+    
     double GetCacheHitRate() const {
         size_t total = cache_hits + cache_misses;
         return total > 0 ? static_cast<double>(cache_hits) / total : 0.0;
@@ -416,10 +421,17 @@ class MemoryPool {
     }
 
     /**
-     * @brief Get pool statistics
+     * @brief Get pool statistics (returns a snapshot copy)
      */
-    MemoryPoolStats GetStats() const {
-        return stats_;
+    void GetStats(MemoryPoolStats& out_stats) const {
+        out_stats.total_allocated.store(stats_.total_allocated.load());
+        out_stats.total_freed.store(stats_.total_freed.load());
+        out_stats.current_usage.store(stats_.current_usage.load());
+        out_stats.peak_usage.store(stats_.peak_usage.load());
+        out_stats.allocation_count.store(stats_.allocation_count.load());
+        out_stats.free_count.store(stats_.free_count.load());
+        out_stats.cache_hits.store(stats_.cache_hits.load());
+        out_stats.cache_misses.store(stats_.cache_misses.load());
     }
 
     /**
@@ -437,8 +449,15 @@ class MemoryPool {
         // Clear large allocations
         large_allocations_.clear();
 
-        // Reset stats
-        stats_ = MemoryPoolStats();
+        // Reset stats (manually reset each atomic)
+        stats_.total_allocated.store(0);
+        stats_.total_freed.store(0);
+        stats_.current_usage.store(0);
+        stats_.peak_usage.store(0);
+        stats_.allocation_count.store(0);
+        stats_.free_count.store(0);
+        stats_.cache_hits.store(0);
+        stats_.cache_misses.store(0);
 
         // Reinitialize if needed
         if (initialized_) {
@@ -601,7 +620,7 @@ class MemoryPool {
     std::unordered_map<std::thread::id, std::unique_ptr<ThreadCache>> thread_caches_;
 
     MemoryPoolStats stats_;
-    Logger logger_;
+    engine::Logger logger_;
 };
 
 /**