Add a preprocessor in Turbo

.gitignore

@@ -60,9 +60,9 @@ allure-*
 !build_android.sh
 !build_ios.sh
 
-# congfig
+# config
 doc/
 config/
 examples/python/
 examples/c_api/
-logs/
+logs/

src/include/zvec/turbo/turbo.h

@@ -13,6 +13,8 @@
 // limitations under the License.
 #pragma once
 
+#include <cstddef>
+#include <cstdint>
 #include <functional>
 #include <zvec/ailego/math_batch/utils.h>
 
@@ -33,6 +35,19 @@ using QueryPreprocessFunc =
 using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale,
                                      float bias, int8_t *out);
 
+// Generic rotate / unrotate function pointer types.
+// ctx is an opaque context (e.g. FhtCtx*) managed by the caller.
+using RotateFunc = void (*)(const float *in, float *out, size_t in_dim,
+                            size_t out_dim, void *ctx);
+using UnrotateFunc = void (*)(const float *in, float *out, size_t in_dim,
+                              size_t out_dim, void *ctx);
+
+// ISA-dispatched rotate/unrotate kernels.
+struct RotatorKernels {
+  RotateFunc rotate = nullptr;
+  UnrotateFunc unrotate = nullptr;
+};
+
 enum class MetricType {
   kSquaredEuclidean,
   kCosine,
@@ -59,6 +74,10 @@ enum class QuantizeType {
   kRabit
 };
 
+enum class RotateType : uint16_t {
+  kFht = 1,  //!< O(d log d) FHT-based Kac random rotation
+};
+
 enum class CpuArchType {
   kAuto,
   kScalar,
@@ -95,4 +114,8 @@ QueryPreprocessFunc get_query_preprocess_func(
 // interface can grow to cover other output types (e.g. fp16) in the future.
 UniformQuantizeFunc get_uniform_quantize_func(DataType data_type);
 
+// Returns rotator kernels dispatched for the current CPU.
+RotatorKernels get_rotator_kernels(
+    RotateType rotate_type, CpuArchType cpu_arch_type = CpuArchType::kAuto);
+
 }  // namespace zvec::turbo

src/turbo/CMakeLists.txt

@@ -5,23 +5,41 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH)
     if (HOST_ARCH MATCHES "^(x86|x64)$")
         setup_compiler_march_for_x86(TURBO_MARCH_FLAG_SSE TURBO_MARCH_FLAG_AVX2 TURBO_MARCH_FLAG_AVX512 TURBO_MARCH_FLAG_AVX512FP16)
     elseif (HOST_ARCH MATCHES "^(arm|arm64)$")
-        # ARM64 architecture - no special march flags needed for now
-        # NEON implementations can be added here if needed
-        message(STATUS "turbo: ARM64 detected, skipping x86-specific optimizations")
+        # ARM64 architecture - NEON is enabled by default on aarch64,
+        # no special march flags needed.
+        message(STATUS "turbo: ARM64 detected, NEON enabled by default")
     endif()
 endif()
 
 file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h)
 
-# Set per-file compile flags for AVX512-VNNI sources.
+# Set per-file compile flags for SIMD sources.
 # set_source_files_properties is directory-scoped, so it must be called in the
 # same directory that adds the sources to a target (i.e. here, not in a
 # subdirectory).
 if(NOT ANDROID AND AUTO_DETECT_ARCH)
     if (HOST_ARCH MATCHES "^(x86|x64)$")
-        file(GLOB_RECURSE AVX512_VNNI_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc)
+        # SSE
+        file(GLOB_RECURSE SSE_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/sse/*.cc)
         set_source_files_properties(
-            ${AVX512_VNNI_SRCS}
+            ${SSE_SRCS}
+            PROPERTIES
+            COMPILE_FLAGS "${TURBO_MARCH_FLAG_SSE}"
+        )
+
+        # AVX2
+        file(GLOB_RECURSE AVX2_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx2/*.cc)
+        set_source_files_properties(
+            ${AVX2_SRCS}
+            PROPERTIES
+            COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX2}"
+        )
+
+        # AVX512 (avx512_vnni, avx512/fht, and turbo.cc dispatch all use AVX512 march)
+        file(GLOB_RECURSE AVX512_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512/*.cc)
+        list(APPEND AVX512_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/turbo.cc)
+        set_source_files_properties(
+            ${AVX512_SRCS}
             PROPERTIES
             COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX512}"
         )

src/turbo/distance/avx2/rotate/fht/fht.cc

@@ -0,0 +1,154 @@
+// Copyright 2025-present the zvec project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#if defined(__AVX2__)
+
+#include "fht.h"
+#include <immintrin.h>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include "common/fht_common.h"
+
+namespace zvec::turbo::avx2 {
+
+void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim) {
+  size_t simd_end = dim & ~31u;
+  constexpr size_t kChunk = 32;
+  const __m256i bit_select =
+      _mm256_setr_epi32(0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80);
+  const __m256 sign_flip = _mm256_castsi256_ps(_mm256_set1_epi32(0x80000000));
+  for (size_t i = 0; i < simd_end; i += kChunk) {
+    uint32_t mask_bits;
+    std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits));
+    for (int b = 0; b < 4; ++b) {
+      __m256i mb = _mm256_set1_epi32((mask_bits >> (b * 8)) & 0xFF);
+      __m256i test = _mm256_and_si256(mb, bit_select);
+      __m256i cmp = _mm256_cmpeq_epi32(test, bit_select);
+      __m256 xor_mask = _mm256_and_ps(_mm256_castsi256_ps(cmp), sign_flip);
+      __m256 v = _mm256_loadu_ps(&data[i + b * 8]);
+      v = _mm256_xor_ps(v, xor_mask);
+      _mm256_storeu_ps(&data[i + b * 8], v);
+    }
+  }
+  // Scalar tail
+  for (size_t i = simd_end; i < dim; ++i) {
+    if (flip[i / 8] & (1u << (i % 8))) {
+      data[i] = -data[i];
+    }
+  }
+}
+
+void fht_kacs_walk_avx2(float *data, size_t len) {
+  size_t half = len / 2;
+  size_t base = len % 2;
+  size_t offset = base + half;
+  size_t half_end = half & ~7u;
+  for (size_t i = 0; i < half_end; i += 8) {
+    __m256 x = _mm256_loadu_ps(&data[i]);
+    __m256 y = _mm256_loadu_ps(&data[i + offset]);
+    _mm256_storeu_ps(&data[i], _mm256_add_ps(x, y));
+    _mm256_storeu_ps(&data[i + offset], _mm256_sub_ps(x, y));
+  }
+  // Scalar tail
+  for (size_t i = half_end; i < half; ++i) {
+    float x = data[i];
+    float y = data[i + offset];
+    data[i] = x + y;
+    data[i + offset] = x - y;
+  }
+  if (base != 0) {
+    data[half] *= std::sqrt(2.0f);
+  }
+}
+
+void fht_inv_kacs_walk_avx2(float *data, size_t len) {
+  size_t half = len / 2;
+  size_t base = len % 2;
+  size_t offset = base + half;
+  if (base != 0) {
+    data[half] *= std::sqrt(0.5f);
+  }
+  size_t half_end = half & ~7u;
+  const __m256 half_fac = _mm256_set1_ps(0.5f);
+  for (size_t i = 0; i < half_end; i += 8) {
+    __m256 a = _mm256_loadu_ps(&data[i]);
+    __m256 b = _mm256_loadu_ps(&data[i + offset]);
+    _mm256_storeu_ps(&data[i], _mm256_mul_ps(_mm256_add_ps(a, b), half_fac));
+    _mm256_storeu_ps(&data[i + offset],
+                     _mm256_mul_ps(_mm256_sub_ps(a, b), half_fac));
+  }
+  // Scalar tail
+  for (size_t i = half_end; i < half; ++i) {
+    float a = data[i];
+    float b = data[i + offset];
+    data[i] = (a + b) * 0.5f;
+    data[i + offset] = (a - b) * 0.5f;
+  }
+}
+
+void fht_inplace_avx2(float *data, size_t n) {
+  for (size_t len = 1; len < n; len <<= 1) {
+    size_t step = len << 1;
+    size_t simd_end = len & ~7u;
+    for (size_t i = 0; i < n; i += step) {
+      for (size_t j = 0; j < simd_end; j += 8) {
+        __m256 u = _mm256_loadu_ps(&data[i + j]);
+        __m256 v = _mm256_loadu_ps(&data[i + j + len]);
+        _mm256_storeu_ps(&data[i + j], _mm256_add_ps(u, v));
+        _mm256_storeu_ps(&data[i + j + len], _mm256_sub_ps(u, v));
+      }
+      for (size_t j = simd_end; j < len; ++j) {
+        float u = data[i + j];
+        float v = data[i + j + len];
+        data[i + j] = u + v;
+        data[i + j + len] = u - v;
+      }
+    }
+  }
+}
+
+void fht_vec_rescale_avx2(float *data, size_t n, float factor) {
+  const __m256 fac = _mm256_set1_ps(factor);
+  size_t simd_end = n & ~7u;
+  for (size_t i = 0; i < simd_end; i += 8) {
+    __m256 v = _mm256_loadu_ps(&data[i]);
+    _mm256_storeu_ps(&data[i], _mm256_mul_ps(v, fac));
+  }
+  // Scalar tail
+  for (size_t i = simd_end; i < n; ++i) {
+    data[i] *= factor;
+  }
+}
+
+void fht_rotate_avx2(const float * /*in*/, float *out, size_t in_dim,
+                     size_t /*out_dim*/, void *ctx) {
+  static constexpr FhtPrimitives kPrim = {
+      fht_flip_sign_avx2, fht_inplace_avx2, fht_kacs_walk_avx2,
+      fht_inv_kacs_walk_avx2, fht_vec_rescale_avx2};
+  fht_rotate_impl(out, in_dim, ctx, kPrim);
+}
+
+void fht_unrotate_avx2(const float * /*in*/, float *out, size_t in_dim,
+                       size_t /*out_dim*/, void *ctx) {
+  static constexpr FhtPrimitives kPrim = {
+      fht_flip_sign_avx2, fht_inplace_avx2, fht_kacs_walk_avx2,
+      fht_inv_kacs_walk_avx2, fht_vec_rescale_avx2};
+  fht_unrotate_impl(out, in_dim, ctx, kPrim);
+}
+
+}  // namespace zvec::turbo::avx2
+
+#endif  // __AVX2__

src/turbo/distance/avx2/rotate/fht/fht.h

@@ -0,0 +1,49 @@
+// Copyright 2025-present the zvec project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#if defined(__AVX2__)
+
+#include <cstddef>
+#include <cstdint>
+
+namespace zvec::turbo::avx2 {
+
+//! Apply bitwise sign-flip mask to a float vector (AVX2).
+void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim);
+
+//! Apply KacsWalk butterfly operation (AVX2).
+void fht_kacs_walk_avx2(float *data, size_t len);
+
+//! Inverse KacsWalk butterfly operation (AVX2).
+void fht_inv_kacs_walk_avx2(float *data, size_t len);
+
+//! In-place Fast Hadamard Transform (AVX2, n must be power-of-2).
+void fht_inplace_avx2(float *data, size_t n);
+
+//! Element-wise rescale: data[i] *= factor (AVX2).
+void fht_vec_rescale_avx2(float *data, size_t n, float factor);
+
+//! Forward FHT rotation (AVX2).
+void fht_rotate_avx2(const float *in, float *out, size_t in_dim, size_t out_dim,
+                     void *ctx);
+
+//! Inverse FHT rotation (AVX2).
+void fht_unrotate_avx2(const float *in, float *out, size_t in_dim,
+                       size_t out_dim, void *ctx);
+
+}  // namespace zvec::turbo::avx2
+
+#endif  // __AVX2__