xiaoxue_admin
/
zvec
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add a preprocessor in Turbo (#548) 

Co-authored-by: ray <rui.xing@alibaba-inc.com>
This commit is contained in:
Zhuanglin Zheng 2026-07-24 14:01:52 +08:00 committed by GitHub
parent 764754427c
commit a4347e31f1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
23 changed files with 2210 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
.gitignore vendored
View File

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

18
src/ailego/internal/cpu_features.cc
View File

@ -23,6 +23,13 @@
#include <cpuid.h> #include <cpuid.h>
#endif #endif
#if defined(__aarch64__) && defined(__linux__)
#include <sys/auxv.h>
#ifndef HWCAP_ASIMD
#define HWCAP_ASIMD (1 << 1)
#endif
#endif
namespace zvec { namespace zvec {
namespace ailego { namespace ailego {
namespace internal { namespace internal {
@ -336,6 +343,17 @@ bool CpuFeatures::HYPERVISOR(void) {
return !!(flags_.L1_ECX & (1u << 31)); return !!(flags_.L1_ECX & (1u << 31));
} }
//! ARM NEON (ASIMD) support
bool CpuFeatures::NEON(void) {
#if defined(__aarch64__) && defined(__linux__)
return !!(getauxval(AT_HWCAP) & HWCAP_ASIMD);
#elif defined(__ARM_NEON)
return true;
#else
return false;
#endif
}
const char *CpuFeatures::Intrinsics(void) { const char *CpuFeatures::Intrinsics(void) {
return "" return ""
#if defined(__ARM_NEON) #if defined(__ARM_NEON)

6
src/ailego/internal/cpu_features.h
View File

@ -182,6 +182,9 @@ class CpuFeatures {
// ！Running on a hypervisor // ！Running on a hypervisor
static bool HYPERVISOR(void); static bool HYPERVISOR(void);
//! ARM NEON (ASIMD) support
static bool NEON(void);
//! Intrinsics of compiling //! Intrinsics of compiling
static const char *Intrinsics(void); static const char *Intrinsics(void);
@ -361,6 +364,9 @@ class CpuFeatures {
// ！Running on a hypervisor // ！Running on a hypervisor
bool HYPERVISOR = CpuFeatures::HYPERVISOR(); bool HYPERVISOR = CpuFeatures::HYPERVISOR();
//! ARM NEON (ASIMD) support
bool NEON = CpuFeatures::NEON();
}; };
static StaticFlags static_flags_; static StaticFlags static_flags_;
}; };

45
src/include/zvec/turbo/turbo.h
View File

@ -13,11 +13,35 @@
// limitations under the License. // limitations under the License.
#pragma once #pragma once
#include <cstddef>
#include <cstdint>
#include <functional> #include <functional>
#include <zvec/ailego/math_batch/utils.h> #include <zvec/ailego/math_batch/utils.h>
namespace zvec::turbo { namespace zvec::turbo {
//! Error code literals mirroring core::IndexError::Code integer values.
//!
//! Turbo quantizer sources use these directly instead of the
//! `IndexError_NotImplemented` / `IndexError_Unsupported` const objects
//! because MSVC's WINDOWS_EXPORT_ALL_SYMBOLS does not export const data
//! with constructors from zvec_shared.dll. zvec_turbo is a static library
//! linked with /WHOLEARCHIVE, so referencing those unexported symbols across
//! the DLL boundary triggers LNK2019 on Windows.
//!
//! IndexError::Code stores -val in its constructor, so NotImplemented(11)
//! yields -11 and Unsupported(12) yields -12.
constexpr int kErrRuntime = -1;
constexpr int kErrNotImplemented = -11;
constexpr int kErrUnsupported = -12;
constexpr int kErrInvalidArgument = -31;
//! Magic number ('QTZR') stamped at the start of a serialized quantizer blob.
constexpr uint32_t kQuantizerMagic = 0x52545A51u;
//! Current quantizer serialization format version.
constexpr uint16_t kQuantizerSerVersion = 1;
using DistanceFunc = using DistanceFunc =
std::function<void(const void *m, const void *q, size_t dim, float *out)>; std::function<void(const void *m, const void *q, size_t dim, float *out)>;
using BatchDistanceFunc = std::function<void( using BatchDistanceFunc = std::function<void(
@ -33,6 +57,19 @@ using QueryPreprocessFunc =
using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale, using UniformQuantizeFunc = void (*)(const float *in, size_t dim, float scale,
float bias, int8_t *out); 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 { enum class MetricType {
kSquaredEuclidean, kSquaredEuclidean,
kCosine, kCosine,
@ -59,6 +96,10 @@ enum class QuantizeType {
kRabit kRabit
}; };
enum class RotateType : uint16_t {
kFht = 1, //!< O(d log d) FHT-based Kac random rotation
};
enum class CpuArchType { enum class CpuArchType {
kAuto, kAuto,
kScalar, kScalar,
@ -95,4 +136,8 @@ QueryPreprocessFunc get_query_preprocess_func(
// interface can grow to cover other output types (e.g. fp16) in the future. // interface can grow to cover other output types (e.g. fp16) in the future.
UniformQuantizeFunc get_uniform_quantize_func(DataType data_type); 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 } // namespace zvec::turbo

29
src/turbo/CMakeLists.txt
View File

@ -5,23 +5,40 @@ if(NOT ANDROID AND AUTO_DETECT_ARCH)
if (HOST_ARCH MATCHES "^(x86|x64)$") 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) 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)$") elseif (HOST_ARCH MATCHES "^(arm|arm64)$")
# ARM64 architecture - no special march flags needed for now # ARM64 architecture - NEON is enabled by default on aarch64,
# NEON implementations can be added here if needed # no special march flags needed.
message(STATUS "turbo: ARM64 detected, skipping x86-specific optimizations") message(STATUS "turbo: ARM64 detected, NEON enabled by default")
endif() endif()
endif() endif()
file(GLOB_RECURSE ALL_SRCS *.cc *.c *.h) 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 # 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 # same directory that adds the sources to a target (i.e. here, not in a
# subdirectory). # subdirectory).
if(NOT ANDROID AND AUTO_DETECT_ARCH) if(NOT ANDROID AND AUTO_DETECT_ARCH)
if (HOST_ARCH MATCHES "^(x86|x64)$") 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( 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 and avx512/fht both use AVX512 march)
file(GLOB_RECURSE AVX512_SRCS ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512_vnni/*.cc ${CMAKE_CURRENT_SOURCE_DIR}/distance/avx512/*.cc)
set_source_files_properties(
${AVX512_SRCS}
PROPERTIES PROPERTIES
COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX512}" COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX512}"
) )

198
src/turbo/distance/avx2/rotate/fht/fht.cc Normal file
View File

@ -0,0 +1,198 @@
// 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.
// This file is compiled with per-file -march=core-avx2 (set in CMakeLists.txt)
// so that AVX2 intrinsics are available. When the build toolchain cannot emit
// AVX2 code, each function falls back to a no-op stub guarded by
// #if defined(__AVX2__).
#include "fht.h"
#if defined(__AVX2__)
#include <immintrin.h>
#endif
#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) {
#if defined(__AVX2__)
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];
}
}
#else
(void)flip;
(void)data;
(void)dim;
#endif
}
void fht_kacs_walk_avx2(float *data, size_t len) {
#if defined(__AVX2__)
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);
}
#else
(void)data;
(void)len;
#endif
}
void fht_inv_kacs_walk_avx2(float *data, size_t len) {
#if defined(__AVX2__)
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;
}
#else
(void)data;
(void)len;
#endif
}
void fht_inplace_avx2(float *data, size_t n) {
#if defined(__AVX2__)
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;
}
}
}
#else
(void)data;
(void)n;
#endif
}
void fht_vec_rescale_avx2(float *data, size_t n, float factor) {
#if defined(__AVX2__)
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;
}
#else
(void)data;
(void)n;
(void)factor;
#endif
}
void fht_rotate_avx2(const float *in, float *out, size_t in_dim,
size_t /*out_dim*/, void *ctx) {
#if defined(__AVX2__)
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(in, out, in_dim, ctx, kPrim);
#else
(void)in;
(void)out;
(void)in_dim;
(void)ctx;
#endif
}
void fht_unrotate_avx2(const float *in, float *out, size_t in_dim,
size_t /*out_dim*/, void *ctx) {
#if defined(__AVX2__)
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(in, out, in_dim, ctx, kPrim);
#else
(void)in;
(void)out;
(void)in_dim;
(void)ctx;
#endif
}
} // namespace zvec::turbo::avx2

45
src/turbo/distance/avx2/rotate/fht/fht.h Normal file
View File

@ -0,0 +1,45 @@
// 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
#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

206
src/turbo/distance/avx512/rotate/fht/fht.cc Normal file
View File

@ -0,0 +1,206 @@
// 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.
// This file is compiled with per-file -march=icelake-server (set in
// CMakeLists.txt) so that AVX512 intrinsics are available. When the build
// toolchain cannot emit AVX-512 code, each function falls back to a no-op
// stub guarded by #if defined(__AVX512F__).
#include "fht.h"
#if defined(__AVX512F__)
#include <immintrin.h>
#endif
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include "common/fht_common.h"
namespace zvec::turbo::avx512 {
void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) {
#if defined(__AVX512F__)
size_t simd_end = dim & ~63u;
constexpr size_t kChunk = 64;
// Sign-flip is a pure bitwise op (x ^ 0x80000000), done in the integer
// domain so only AVX512F is required. The float-domain VXORPS (zmm) would
// pull in an AVX512DQ dependency, which we deliberately avoid here.
const __m512i sign_bit = _mm512_set1_epi32(0x80000000);
for (size_t i = 0; i < simd_end; i += kChunk) {
uint64_t mask_bits;
std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits));
const __mmask16 m0 = _cvtu32_mask16(mask_bits & 0xFFFF);
const __mmask16 m1 = _cvtu32_mask16((mask_bits >> 16) & 0xFFFF);
const __mmask16 m2 = _cvtu32_mask16((mask_bits >> 32) & 0xFFFF);
const __mmask16 m3 = _cvtu32_mask16((mask_bits >> 48) & 0xFFFF);
__m512i v0 = _mm512_castps_si512(_mm512_loadu_ps(&data[i]));
v0 = _mm512_mask_xor_epi32(v0, m0, v0, sign_bit);
_mm512_storeu_ps(&data[i], _mm512_castsi512_ps(v0));
__m512i v1 = _mm512_castps_si512(_mm512_loadu_ps(&data[i + 16]));
v1 = _mm512_mask_xor_epi32(v1, m1, v1, sign_bit);
_mm512_storeu_ps(&data[i + 16], _mm512_castsi512_ps(v1));
__m512i v2 = _mm512_castps_si512(_mm512_loadu_ps(&data[i + 32]));
v2 = _mm512_mask_xor_epi32(v2, m2, v2, sign_bit);
_mm512_storeu_ps(&data[i + 32], _mm512_castsi512_ps(v2));
__m512i v3 = _mm512_castps_si512(_mm512_loadu_ps(&data[i + 48]));
v3 = _mm512_mask_xor_epi32(v3, m3, v3, sign_bit);
_mm512_storeu_ps(&data[i + 48], _mm512_castsi512_ps(v3));
}
// Scalar tail
for (size_t i = simd_end; i < dim; ++i) {
if (flip[i / 8] & (1u << (i % 8))) {
data[i] = -data[i];
}
}
#else
(void)flip;
(void)data;
(void)dim;
#endif
}
void fht_kacs_walk_avx512(float *data, size_t len) {
#if defined(__AVX512F__)
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
size_t half_end = half & ~15u;
for (size_t i = 0; i < half_end; i += 16) {
__m512 x = _mm512_loadu_ps(&data[i]);
__m512 y = _mm512_loadu_ps(&data[i + offset]);
_mm512_storeu_ps(&data[i], _mm512_add_ps(x, y));
_mm512_storeu_ps(&data[i + offset], _mm512_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);
}
#else
(void)data;
(void)len;
#endif
}
void fht_inv_kacs_walk_avx512(float *data, size_t len) {
#if defined(__AVX512F__)
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 & ~15u;
const __m512 half_fac = _mm512_set1_ps(0.5f);
for (size_t i = 0; i < half_end; i += 16) {
__m512 a = _mm512_loadu_ps(&data[i]);
__m512 b = _mm512_loadu_ps(&data[i + offset]);
_mm512_storeu_ps(&data[i], _mm512_mul_ps(_mm512_add_ps(a, b), half_fac));
_mm512_storeu_ps(&data[i + offset],
_mm512_mul_ps(_mm512_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;
}
#else
(void)data;
(void)len;
#endif
}
void fht_inplace_avx512(float *data, size_t n) {
#if defined(__AVX512F__)
for (size_t len = 1; len < n; len <<= 1) {
size_t step = len << 1;
size_t simd_end = len & ~15u;
for (size_t i = 0; i < n; i += step) {
for (size_t j = 0; j < simd_end; j += 16) {
__m512 u = _mm512_loadu_ps(&data[i + j]);
__m512 v = _mm512_loadu_ps(&data[i + j + len]);
_mm512_storeu_ps(&data[i + j], _mm512_add_ps(u, v));
_mm512_storeu_ps(&data[i + j + len], _mm512_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;
}
}
}
#else
(void)data;
(void)n;
#endif
}
void fht_vec_rescale_avx512(float *data, size_t n, float factor) {
#if defined(__AVX512F__)
const __m512 fac = _mm512_set1_ps(factor);
size_t simd_end = n & ~15u;
for (size_t i = 0; i < simd_end; i += 16) {
__m512 v = _mm512_loadu_ps(&data[i]);
_mm512_storeu_ps(&data[i], _mm512_mul_ps(v, fac));
}
// Scalar tail
for (size_t i = simd_end; i < n; ++i) {
data[i] *= factor;
}
#else
(void)data;
(void)n;
(void)factor;
#endif
}
void fht_rotate_avx512(const float *in, float *out, size_t in_dim,
size_t /*out_dim*/, void *ctx) {
#if defined(__AVX512F__)
static constexpr FhtPrimitives kPrim = {
fht_flip_sign_avx512, fht_inplace_avx512, fht_kacs_walk_avx512,
fht_inv_kacs_walk_avx512, fht_vec_rescale_avx512};
fht_rotate_impl(in, out, in_dim, ctx, kPrim);
#else
(void)in;
(void)out;
(void)in_dim;
(void)ctx;
#endif
}
void fht_unrotate_avx512(const float *in, float *out, size_t in_dim,
size_t /*out_dim*/, void *ctx) {
#if defined(__AVX512F__)
static constexpr FhtPrimitives kPrim = {
fht_flip_sign_avx512, fht_inplace_avx512, fht_kacs_walk_avx512,
fht_inv_kacs_walk_avx512, fht_vec_rescale_avx512};
fht_unrotate_impl(in, out, in_dim, ctx, kPrim);
#else
(void)in;
(void)out;
(void)in_dim;
(void)ctx;
#endif
}
} // namespace zvec::turbo::avx512

45
src/turbo/distance/avx512/rotate/fht/fht.h Normal file
View File

@ -0,0 +1,45 @@
// 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
#include <cstddef>
#include <cstdint>
namespace zvec::turbo::avx512 {
//! Apply bitwise sign-flip mask to a float vector (AVX512).
void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim);
//! Apply KacsWalk butterfly operation (AVX512).
void fht_kacs_walk_avx512(float *data, size_t len);
//! Inverse KacsWalk butterfly operation (AVX512).
void fht_inv_kacs_walk_avx512(float *data, size_t len);
//! In-place Fast Hadamard Transform (AVX512, n must be power-of-2).
void fht_inplace_avx512(float *data, size_t n);
//! Element-wise rescale: data[i] *= factor (AVX512).
void fht_vec_rescale_avx512(float *data, size_t n, float factor);
//! Forward FHT rotation (AVX512).
void fht_rotate_avx512(const float *in, float *out, size_t in_dim,
size_t out_dim, void *ctx);
//! Inverse FHT rotation (AVX512).
void fht_unrotate_avx512(const float *in, float *out, size_t in_dim,
size_t out_dim, void *ctx);
} // namespace zvec::turbo::avx512

133
src/turbo/distance/common/fht_common.h Normal file
View File

@ -0,0 +1,133 @@
// 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
#include <cstddef>
#include <cstdint>
#include <cstring>
namespace zvec::turbo {
/// ISA-level FHT primitive function pointers.
/// Each ISA fills in its own SIMD-optimized (or scalar-fallback) functions.
struct FhtPrimitives {
void (*flip_sign)(const uint8_t *flip, float *data, size_t dim);
void (*inplace)(float *data, size_t n);
void (*kacs_walk)(float *data, size_t len);
void (*inv_kacs_walk)(float *data, size_t len);
void (*rescale)(float *data, size_t n, float factor);
};
/// FhtCtx memory layout (accessed by address, NOT by type):
/// offset 0: size_t flip_offset
/// offset 8: size_t trunc_dim
/// offset 16: float fac
/// offset 20: uint8_t pad[4]
/// offset 24: uint8_t flip[]
inline void fht_rotate_impl(const float *in, float *out, size_t dim, void *ctx,
const FhtPrimitives &p) {
if (out != in) {
std::memcpy(out, in, sizeof(float) * dim);
}
float *data = out;
auto *base = reinterpret_cast<const uint8_t *>(ctx);
const size_t flip_offset = *reinterpret_cast<const size_t *>(base);
const size_t trunc_dim = *reinterpret_cast<const size_t *>(base + 8);
const float fac = *reinterpret_cast<const float *>(base + 16);
const uint8_t *flip = base + 24;
if (trunc_dim == dim) {
for (size_t r = 0; r < 4; ++r) {
p.flip_sign(flip + r * flip_offset, data, dim);
p.inplace(data, trunc_dim);
p.rescale(data, trunc_dim, fac);
}
return;
}
size_t start = dim - trunc_dim;
float *trunc_ptr = data + start;
p.flip_sign(flip, data, dim);
p.inplace(data, trunc_dim);
p.rescale(data, trunc_dim, fac);
p.kacs_walk(data, dim);
p.flip_sign(flip + flip_offset, data, dim);
p.inplace(trunc_ptr, trunc_dim);
p.rescale(trunc_ptr, trunc_dim, fac);
p.kacs_walk(data, dim);
p.flip_sign(flip + 2 * flip_offset, data, dim);
p.inplace(data, trunc_dim);
p.rescale(data, trunc_dim, fac);
p.kacs_walk(data, dim);
p.flip_sign(flip + 3 * flip_offset, data, dim);
p.inplace(trunc_ptr, trunc_dim);
p.rescale(trunc_ptr, trunc_dim, fac);
p.kacs_walk(data, dim);
p.rescale(data, dim, 0.25f);
}
inline void fht_unrotate_impl(const float *in, float *out, size_t dim,
void *ctx, const FhtPrimitives &p) {
if (out != in) {
std::memcpy(out, in, sizeof(float) * dim);
}
float *data = out;
auto *base = reinterpret_cast<const uint8_t *>(ctx);
const size_t flip_offset = *reinterpret_cast<const size_t *>(base);
const size_t trunc_dim = *reinterpret_cast<const size_t *>(base + 8);
const float fac = *reinterpret_cast<const float *>(base + 16);
const uint8_t *flip = base + 24;
if (trunc_dim == dim) {
for (int round = 3; round >= 0; --round) {
p.inplace(data, trunc_dim);
p.rescale(data, trunc_dim, fac);
p.flip_sign(flip + static_cast<size_t>(round) * flip_offset, data, dim);
}
return;
}
p.rescale(data, dim, 4.0f);
size_t start = dim - trunc_dim;
float *trunc_ptr = data + start;
p.inv_kacs_walk(data, dim);
p.inplace(trunc_ptr, trunc_dim);
p.rescale(trunc_ptr, trunc_dim, fac);
p.flip_sign(flip + 3 * flip_offset, data, dim);
p.inv_kacs_walk(data, dim);
p.inplace(data, trunc_dim);
p.rescale(data, trunc_dim, fac);
p.flip_sign(flip + 2 * flip_offset, data, dim);
p.inv_kacs_walk(data, dim);
p.inplace(trunc_ptr, trunc_dim);
p.rescale(trunc_ptr, trunc_dim, fac);
p.flip_sign(flip + flip_offset, data, dim);
p.inv_kacs_walk(data, dim);
p.inplace(data, trunc_dim);
p.rescale(data, trunc_dim, fac);
p.flip_sign(flip, data, dim);
}
} // namespace zvec::turbo

171
src/turbo/distance/neon/rotate/fht/fht.cc Normal file
View File

@ -0,0 +1,171 @@
// 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.
#include "fht.h"
#if defined(__ARM_NEON) && defined(__aarch64__)
#include <arm_neon.h>
#endif
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include "common/fht_common.h"
#include "scalar/rotate/fht/fht.h"
namespace zvec::turbo::neon {
void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim) {
#if defined(__ARM_NEON) && defined(__aarch64__)
const uint32x4_t sign_bit = vdupq_n_u32(0x80000000u);
size_t simd_end = dim & ~3u;
size_t flip_bytes = (dim + 7) / 8;
for (size_t i = 0; i < simd_end; i += 4) {
uint16_t bits16;
size_t byte_pos = i / 8;
if (byte_pos + 1 < flip_bytes) {
std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16));
} else {
bits16 = flip[byte_pos];
}
bits16 >>= (i % 8);
uint32_t b0 = bits16 & 1u;
uint32_t b1 = (bits16 >> 1) & 1u;
uint32_t b2 = (bits16 >> 2) & 1u;
uint32_t b3 = (bits16 >> 3) & 1u;
uint32x4_t bit_mask = {b0, b1, b2, b3};
uint32x4_t sign_mask = vmulq_u32(bit_mask, sign_bit);
float32x4_t v = vld1q_f32(&data[i]);
v = vreinterpretq_f32_u32(veorq_u32(vreinterpretq_u32_f32(v), sign_mask));
vst1q_f32(&data[i], v);
}
// Scalar tail
for (size_t i = simd_end; i < dim; ++i) {
if (flip[i / 8] & (1u << (i % 8))) {
data[i] = -data[i];
}
}
#else
(void)flip;
(void)data;
(void)dim;
#endif
}
void fht_kacs_walk_neon(float *data, size_t len) {
#if defined(__ARM_NEON) && defined(__aarch64__)
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
size_t half_end = half & ~3u;
for (size_t i = 0; i < half_end; i += 4) {
float32x4_t x = vld1q_f32(&data[i]);
float32x4_t y = vld1q_f32(&data[i + offset]);
vst1q_f32(&data[i], vaddq_f32(x, y));
vst1q_f32(&data[i + offset], vsubq_f32(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);
}
#else
(void)data;
(void)len;
#endif
}
void fht_inv_kacs_walk_neon(float *data, size_t len) {
#if defined(__ARM_NEON) && defined(__aarch64__)
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 & ~3u;
const float32x4_t half_fac = vdupq_n_f32(0.5f);
for (size_t i = 0; i < half_end; i += 4) {
float32x4_t a = vld1q_f32(&data[i]);
float32x4_t b = vld1q_f32(&data[i + offset]);
vst1q_f32(&data[i], vmulq_f32(vaddq_f32(a, b), half_fac));
vst1q_f32(&data[i + offset], vmulq_f32(vsubq_f32(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;
}
#else
(void)data;
(void)len;
#endif
}
void fht_vec_rescale_neon(float *data, size_t n, float factor) {
#if defined(__ARM_NEON) && defined(__aarch64__)
const float32x4_t fac = vdupq_n_f32(factor);
size_t simd_end = n & ~3u;
for (size_t i = 0; i < simd_end; i += 4) {
float32x4_t v = vld1q_f32(&data[i]);
vst1q_f32(&data[i], vmulq_f32(v, fac));
}
// Scalar tail
for (size_t i = simd_end; i < n; ++i) {
data[i] *= factor;
}
#else
(void)data;
(void)n;
(void)factor;
#endif
}
void fht_rotate_neon(const float *in, float *out, size_t in_dim,
size_t /*out_dim*/, void *ctx) {
#if defined(__ARM_NEON) && defined(__aarch64__)
static constexpr FhtPrimitives kPrim = {
fht_flip_sign_neon, scalar::fht_inplace, fht_kacs_walk_neon,
fht_inv_kacs_walk_neon, fht_vec_rescale_neon};
fht_rotate_impl(in, out, in_dim, ctx, kPrim);
#else
(void)in;
(void)out;
(void)in_dim;
(void)ctx;
#endif
}
void fht_unrotate_neon(const float *in, float *out, size_t in_dim,
size_t /*out_dim*/, void *ctx) {
#if defined(__ARM_NEON) && defined(__aarch64__)
static constexpr FhtPrimitives kPrim = {
fht_flip_sign_neon, scalar::fht_inplace, fht_kacs_walk_neon,
fht_inv_kacs_walk_neon, fht_vec_rescale_neon};
fht_unrotate_impl(in, out, in_dim, ctx, kPrim);
#else
(void)in;
(void)out;
(void)in_dim;
(void)ctx;
#endif
}
} // namespace zvec::turbo::neon

42
src/turbo/distance/neon/rotate/fht/fht.h Normal file
View File

@ -0,0 +1,42 @@
// 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
#include <cstddef>
#include <cstdint>
namespace zvec::turbo::neon {
//! Apply bitwise sign-flip mask to a float vector (NEON).
void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim);
//! Apply KacsWalk butterfly operation (NEON).
void fht_kacs_walk_neon(float *data, size_t len);
//! Inverse KacsWalk butterfly operation (NEON).
void fht_inv_kacs_walk_neon(float *data, size_t len);
//! Element-wise rescale: data[i] *= factor (NEON).
void fht_vec_rescale_neon(float *data, size_t n, float factor);
//! Forward FHT rotation (NEON). Inplace falls back to scalar.
void fht_rotate_neon(const float *in, float *out, size_t in_dim, size_t out_dim,
void *ctx);
//! Inverse FHT rotation (NEON). Inplace falls back to scalar.
void fht_unrotate_neon(const float *in, float *out, size_t in_dim,
size_t out_dim, void *ctx);
} // namespace zvec::turbo::neon

96
src/turbo/distance/scalar/rotate/fht/fht.cc Normal file
View File

@ -0,0 +1,96 @@
// 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.
#include "fht.h"
#include <cmath>
#include <cstddef>
#include <cstdint>
#include "common/fht_common.h"
namespace zvec::turbo::scalar {
void fht_flip_sign(const uint8_t *flip, float *data, size_t dim) {
for (size_t i = 0; i < dim; ++i) {
if (flip[i / 8] & (1u << (i % 8))) {
data[i] = -data[i];
}
}
}
void fht_kacs_walk(float *data, size_t len) {
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
for (size_t i = 0; 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(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);
}
for (size_t i = 0; 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(float *data, size_t n) {
for (size_t len = 1; len < n; len <<= 1) {
for (size_t i = 0; i < n; i += len << 1) {
for (size_t j = i; j < i + len; ++j) {
float u = data[j];
float v = data[j + len];
data[j] = u + v;
data[j + len] = u - v;
}
}
}
}
void fht_vec_rescale(float *data, size_t n, float factor) {
for (size_t i = 0; i < n; ++i) {
data[i] *= factor;
}
}
void fht_rotate(const float *in, float *out, size_t in_dim, size_t /*out_dim*/,
void *ctx) {
static constexpr FhtPrimitives kPrim = {fht_flip_sign, fht_inplace,
fht_kacs_walk, fht_inv_kacs_walk,
fht_vec_rescale};
fht_rotate_impl(in, out, in_dim, ctx, kPrim);
}
void fht_unrotate(const float *in, float *out, size_t in_dim,
size_t /*out_dim*/, void *ctx) {
static constexpr FhtPrimitives kPrim = {fht_flip_sign, fht_inplace,
fht_kacs_walk, fht_inv_kacs_walk,
fht_vec_rescale};
fht_unrotate_impl(in, out, in_dim, ctx, kPrim);
}
} // namespace zvec::turbo::scalar

47
src/turbo/distance/scalar/rotate/fht/fht.h Normal file
View File

@ -0,0 +1,47 @@
// 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
#include <cstddef>
#include <cstdint>
namespace zvec::turbo::scalar {
//! Apply bitwise sign-flip mask to a float vector.
//! Each bit in \p flip controls one element: bit set means negate.
void fht_flip_sign(const uint8_t *flip, float *data, size_t dim);
//! Apply KacsWalk butterfly operation to non-power-of-2 FHT.
void fht_kacs_walk(float *data, size_t len);
//! Inverse KacsWalk butterfly operation.
void fht_inv_kacs_walk(float *data, size_t len);
//! In-place Fast Hadamard Transform on \p n elements (must be power-of-2).
void fht_inplace(float *data, size_t n);
//! Element-wise rescale: data[i] *= factor.
void fht_vec_rescale(float *data, size_t n, float factor);
//! Forward FHT rotation (compose flip -> FHT -> rescale, 4 rounds).
//! ctx is a FhtCtx* defined in preprocessor/fht_rotator/fht_rotator.h.
void fht_rotate(const float *in, float *out, size_t in_dim, size_t out_dim,
void *ctx);
//! Inverse FHT rotation (undo 4 rounds in reverse order).
void fht_unrotate(const float *in, float *out, size_t in_dim, size_t out_dim,
void *ctx);
} // namespace zvec::turbo::scalar

175
src/turbo/distance/sse/rotate/fht/fht.cc Normal file
View File

@ -0,0 +1,175 @@
// 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.
// This file is compiled with per-file -march=corei7 (set in CMakeLists.txt)
// so that SSE2 intrinsics are available. When the build toolchain cannot emit
// SSE2 code, each function falls back to a no-op stub guarded by
// #if defined(__SSE2__).
#include "fht.h"
#if defined(__SSE2__)
#include <emmintrin.h>
#endif
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include "common/fht_common.h"
#include "scalar/rotate/fht/fht.h"
namespace zvec::turbo::sse {
void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim) {
#if defined(__SSE2__)
size_t simd_end = dim & ~3u;
size_t flip_bytes = (dim + 7) / 8;
for (size_t i = 0; i < simd_end; i += 4) {
uint16_t bits16;
size_t byte_pos = i / 8;
if (byte_pos + 1 < flip_bytes) {
std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16));
} else {
bits16 = flip[byte_pos];
}
bits16 >>= (i % 8);
uint32_t b0 = bits16 & 1u;
uint32_t b1 = (bits16 >> 1) & 1u;
uint32_t b2 = (bits16 >> 2) & 1u;
uint32_t b3 = (bits16 >> 3) & 1u;
__m128i bit_mask = _mm_set_epi32(b3, b2, b1, b0);
__m128i sign_mask = _mm_slli_epi32(bit_mask, 31);
__m128 v = _mm_loadu_ps(&data[i]);
v = _mm_xor_ps(v, _mm_castsi128_ps(sign_mask));
_mm_storeu_ps(&data[i], v);
}
// Scalar tail
for (size_t i = simd_end; i < dim; ++i) {
if (flip[i / 8] & (1u << (i % 8))) {
data[i] = -data[i];
}
}
#else
(void)flip;
(void)data;
(void)dim;
#endif
}
void fht_kacs_walk_sse(float *data, size_t len) {
#if defined(__SSE2__)
size_t half = len / 2;
size_t base = len % 2;
size_t offset = base + half;
size_t half_end = half & ~3u;
for (size_t i = 0; i < half_end; i += 4) {
__m128 x = _mm_loadu_ps(&data[i]);
__m128 y = _mm_loadu_ps(&data[i + offset]);
_mm_storeu_ps(&data[i], _mm_add_ps(x, y));
_mm_storeu_ps(&data[i + offset], _mm_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);
}
#else
(void)data;
(void)len;
#endif
}
void fht_inv_kacs_walk_sse(float *data, size_t len) {
#if defined(__SSE2__)
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 & ~3u;
const __m128 half_fac = _mm_set1_ps(0.5f);
for (size_t i = 0; i < half_end; i += 4) {
__m128 a = _mm_loadu_ps(&data[i]);
__m128 b = _mm_loadu_ps(&data[i + offset]);
_mm_storeu_ps(&data[i], _mm_mul_ps(_mm_add_ps(a, b), half_fac));
_mm_storeu_ps(&data[i + offset], _mm_mul_ps(_mm_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;
}
#else
(void)data;
(void)len;
#endif
}
void fht_vec_rescale_sse(float *data, size_t n, float factor) {
#if defined(__SSE2__)
const __m128 fac = _mm_set1_ps(factor);
size_t simd_end = n & ~3u;
for (size_t i = 0; i < simd_end; i += 4) {
__m128 v = _mm_loadu_ps(&data[i]);
_mm_storeu_ps(&data[i], _mm_mul_ps(v, fac));
}
// Scalar tail
for (size_t i = simd_end; i < n; ++i) {
data[i] *= factor;
}
#else
(void)data;
(void)n;
(void)factor;
#endif
}
void fht_rotate_sse(const float *in, float *out, size_t in_dim,
size_t /*out_dim*/, void *ctx) {
#if defined(__SSE2__)
static constexpr FhtPrimitives kPrim = {
fht_flip_sign_sse, scalar::fht_inplace, fht_kacs_walk_sse,
fht_inv_kacs_walk_sse, fht_vec_rescale_sse};
fht_rotate_impl(in, out, in_dim, ctx, kPrim);
#else
(void)in;
(void)out;
(void)in_dim;
(void)ctx;
#endif
}
void fht_unrotate_sse(const float *in, float *out, size_t in_dim,
size_t /*out_dim*/, void *ctx) {
#if defined(__SSE2__)
static constexpr FhtPrimitives kPrim = {
fht_flip_sign_sse, scalar::fht_inplace, fht_kacs_walk_sse,
fht_inv_kacs_walk_sse, fht_vec_rescale_sse};
fht_unrotate_impl(in, out, in_dim, ctx, kPrim);
#else
(void)in;
(void)out;
(void)in_dim;
(void)ctx;
#endif
}
} // namespace zvec::turbo::sse

42
src/turbo/distance/sse/rotate/fht/fht.h Normal file
View File

@ -0,0 +1,42 @@
// 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
#include <cstddef>
#include <cstdint>
namespace zvec::turbo::sse {
//! Apply bitwise sign-flip mask to a float vector (SSE).
void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim);
//! Apply KacsWalk butterfly operation (SSE).
void fht_kacs_walk_sse(float *data, size_t len);
//! Inverse KacsWalk butterfly operation (SSE).
void fht_inv_kacs_walk_sse(float *data, size_t len);
//! Element-wise rescale: data[i] *= factor (SSE).
void fht_vec_rescale_sse(float *data, size_t n, float factor);
//! Forward FHT rotation (SSE). Inplace falls back to scalar.
void fht_rotate_sse(const float *in, float *out, size_t in_dim, size_t out_dim,
void *ctx);
//! Inverse FHT rotation (SSE). Inplace falls back to scalar.
void fht_unrotate_sse(const float *in, float *out, size_t in_dim,
size_t out_dim, void *ctx);
} // namespace zvec::turbo::sse

217
src/turbo/preprocessor/fht_rotator/fht_rotator.cc Normal file
View File

@ -0,0 +1,217 @@
// 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.
#include "fht_rotator.h"
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <random>
namespace zvec {
namespace turbo {
// ============================================================================
// FhtRotator method implementations
// ============================================================================
size_t FhtRotator::floor_pow2(size_t n) {
if (n == 0) return 0;
size_t p = 1;
while (p * 2 <= n) p *= 2;
return p;
}
FhtRotator::~FhtRotator() {
std::free(fht_ctx_);
}
FhtRotator::Pointer FhtRotator::create(int dim) {
if (dim <= 0) return nullptr;
Pointer r(new FhtRotator());
r->in_dim_ = dim;
r->out_dim_ = dim;
r->flip_offset_ = (static_cast<size_t>(dim) + kByteLen - 1) / kByteLen;
r->kernels_ = get_rotator_kernels(RotateType::kFht);
const size_t trunc_dim = floor_pow2(static_cast<size_t>(dim));
const float fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
const size_t flip_size = 4 * r->flip_offset_;
// Single allocation: FhtCtx header + trailing flip data.
r->fht_ctx_ = static_cast<FhtCtx *>(std::malloc(sizeof(FhtCtx) + flip_size));
if (r->fht_ctx_ == nullptr) {
// Allocation failed: bail out before dereferencing. Returning nullptr lets
// the smart Pointer clean up the partially-built object (~FhtRotator frees
// the null fht_ctx_ safely).
return nullptr;
}
r->fht_ctx_->flip_offset = r->flip_offset_;
r->fht_ctx_->trunc_dim = trunc_dim;
r->fht_ctx_->fac = fac;
// Generate 4 rounds of random flip-sign arrays.
std::random_device rd;
std::mt19937 gen(rd());
std::uniform_int_distribution<int> dist(0, 255);
for (size_t i = 0; i < flip_size; ++i)
r->fht_ctx_->flip[i] = static_cast<uint8_t>(dist(gen));
return r;
}
FhtRotator::Pointer FhtRotator::from_blob(const void *data, size_t len) {
if (!data || len < sizeof(RotatorSerHeader)) return nullptr;
// Copy the header into a properly aligned local before reading any field.
// `data` may point to an unaligned byte buffer (e.g. std::string::data()),
// so dereferencing a reinterpret_cast<const RotatorSerHeader *> directly
// would be undefined behavior on architectures that require alignment.
RotatorSerHeader hdr;
std::memcpy(&hdr, data, sizeof(RotatorSerHeader));
if (hdr.magic != kRotatorMagic) return nullptr;
if (hdr.version != kRotatorSerVersion) return nullptr;
if (static_cast<RotateType>(hdr.rotator_type) != RotateType::kFht) {
return nullptr;
}
Pointer r(new FhtRotator());
const size_t expected_total =
sizeof(RotatorSerHeader) + static_cast<size_t>(hdr.payload_size);
if (len < expected_total) return nullptr;
int rc = r->deserialize(data, len);
if (rc != 0) return nullptr;
return r;
}
void FhtRotator::train(const void * /*data*/, size_t /*num*/,
size_t /*stride*/) {
// No-op: flip-sign arrays are generated in create().
}
// ---------------------------------------------------------------------------
// apply (forward rotation)
// ---------------------------------------------------------------------------
void FhtRotator::apply(const float *in, float *out) const {
kernels_.rotate(in, out, static_cast<size_t>(in_dim_),
static_cast<size_t>(out_dim_), static_cast<void *>(fht_ctx_));
}
// ---------------------------------------------------------------------------
// apply_inverse (inverse rotation)
// ---------------------------------------------------------------------------
void FhtRotator::apply_inverse(const float *in, float *out) const {
kernels_.unrotate(in, out, static_cast<size_t>(in_dim_),
static_cast<size_t>(out_dim_),
static_cast<void *>(fht_ctx_));
}
// ---------------------------------------------------------------------------
// serialize / deserialize
// ---------------------------------------------------------------------------
int FhtRotator::serialize(std::string *out) const {
if (!out) return kErrInvalidArgument;
if (!fht_ctx_) return kErrRuntime;
const size_t flip_size = 4 * flip_offset_;
RotatorSerHeader hdr{};
hdr.magic = kRotatorMagic;
hdr.version = kRotatorSerVersion;
hdr.rotator_type = static_cast<uint16_t>(RotateType::kFht);
hdr.in_dim = static_cast<uint32_t>(in_dim_);
hdr.out_dim = static_cast<uint32_t>(out_dim_);
hdr.payload_size = static_cast<uint32_t>(flip_size);
hdr.reserved = 0;
out->resize(sizeof(hdr) + flip_size);
std::memcpy(&(*out)[0], &hdr, sizeof(hdr));
std::memcpy(&(*out)[sizeof(hdr)], fht_ctx_->flip, flip_size);
return 0;
}
int FhtRotator::deserialize(const void *data, size_t len) {
if (!data || len < sizeof(RotatorSerHeader)) return kErrInvalidArgument;
// Copy the header into a properly aligned local before reading any field.
// `data` may point to an unaligned byte buffer (e.g. std::string::data()),
// so dereferencing a reinterpret_cast<const RotatorSerHeader *> directly
// would be undefined behavior on architectures that require alignment.
RotatorSerHeader hdr;
std::memcpy(&hdr, data, sizeof(RotatorSerHeader));
if (hdr.magic != kRotatorMagic) return kErrUnsupported;
if (hdr.version != kRotatorSerVersion) return kErrUnsupported;
if (static_cast<RotateType>(hdr.rotator_type) != RotateType::kFht) {
return kErrUnsupported;
}
// Validate dimensions before any cast to int: must be strictly positive and
// representable as int. FHT keeps dimensionality unchanged (in == out).
if (hdr.in_dim == 0 ||
hdr.in_dim > static_cast<uint32_t>(std::numeric_limits<int>::max())) {
return kErrInvalidArgument;
}
if (hdr.out_dim != hdr.in_dim) return kErrInvalidArgument;
// Length check via subtraction to avoid size_t overflow on 32-bit
// (len >= sizeof(header) was guaranteed above, so the subtraction is safe).
if (hdr.payload_size > len - sizeof(RotatorSerHeader)) {
return kErrInvalidArgument;
}
// Payload must hold exactly 4 rounds of ceil(in_dim/8) flip bytes. The
// rotation kernels read 4 * flip_offset bytes, so any smaller payload would
// read out of bounds; require an exact match.
const size_t new_flip_offset =
(static_cast<size_t>(hdr.in_dim) + kByteLen - 1) / kByteLen;
const size_t expected_flip_size = 4 * new_flip_offset;
if (hdr.payload_size != expected_flip_size) return kErrInvalidArgument;
// Build the new context into a temporary first, so a failed allocation
// leaves the existing object completely untouched (deserialize may target
// an already-initialized rotator). Only commit member state once the
// allocation and payload copy have both succeeded.
const int new_in_dim = static_cast<int>(hdr.in_dim);
const int new_out_dim = static_cast<int>(hdr.out_dim);
const size_t trunc_dim = floor_pow2(static_cast<size_t>(new_in_dim));
const float fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
FhtCtx *new_ctx =
static_cast<FhtCtx *>(std::malloc(sizeof(FhtCtx) + expected_flip_size));
if (new_ctx == nullptr) return kErrRuntime;
new_ctx->flip_offset = new_flip_offset;
new_ctx->trunc_dim = trunc_dim;
new_ctx->fac = fac;
std::memcpy(new_ctx->flip,
reinterpret_cast<const char *>(data) + sizeof(RotatorSerHeader),
expected_flip_size);
// All allocations succeeded: commit new state and free the old context.
std::free(fht_ctx_);
fht_ctx_ = new_ctx;
in_dim_ = new_in_dim;
out_dim_ = new_out_dim;
flip_offset_ = new_flip_offset;
kernels_ = get_rotator_kernels(RotateType::kFht);
return 0;
}
} // namespace turbo
} // namespace zvec

114
src/turbo/preprocessor/fht_rotator/fht_rotator.h Normal file
View File

@ -0,0 +1,114 @@
// 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
#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include "preprocessor/preprocessor.h"
namespace zvec {
namespace turbo {
// FHT context passed to ISA-level rotate/unrotate via void*.
// Layout (ISA code accesses by address, NOT by type):
// offset 0: size_t flip_offset (bytes per round)
// offset 8: size_t trunc_dim (largest power-of-2 <= in_dim)
// offset 16: float fac (1 / sqrt(trunc_dim))
// offset 20: uint8_t pad[4] (explicit padding)
// offset 24: uint8_t flip[] (4 * flip_offset bytes, trailing data)
//
// Allocated as a single block: malloc(sizeof(FhtCtx) + 4 * flip_offset).
struct FhtCtx {
size_t flip_offset;
size_t trunc_dim;
float fac;
uint8_t pad_[4];
uint8_t flip[]; // trailing flexible array (C++ extension)
};
static_assert(offsetof(FhtCtx, flip) == 24, "FhtCtx flip offset must be 24");
static_assert(sizeof(FhtCtx) == 24, "FhtCtx sizeof must be 24");
// ============================================================================
// FhtRotator - O(d log d) FHT-based Kac random rotation
//
// Works with any dimension (non-power-of-2 uses trunc_dim + KacsWalk).
// When dimension is a power of 2, uses 4 rounds of (flip -> FHT -> rescale).
// When dimension is NOT a power of 2, uses kacs_walk reduction.
// ============================================================================
class FhtRotator : public Preprocessor {
public:
using Pointer = std::shared_ptr<FhtRotator>;
//! Create a fully-initialized rotator for \p in_dim dimensions.
//! Random flip-sign arrays are generated during creation; the returned
//! object is immediately usable for apply() / apply_inverse().
static Pointer create(int in_dim);
//! Create and restore a rotator from a serialized blob (reads the type from
//! the embedded RotatorSerHeader). Returns nullptr on malformed input.
static Pointer from_blob(const void *data, size_t len);
// -- Preprocessor interface ------------------------------------------------
int in_dim() const override {
return in_dim_;
}
int out_dim() const override {
return out_dim_;
}
void apply(const float *in, float *out) const override;
void apply_inverse(const float *in, float *out) const override;
//! No-op for FhtRotator. Flip-sign arrays are generated in create().
//! Provided for interface compatibility with the Preprocessor contract.
void train(const void *data, size_t num, size_t stride) override;
int serialize(std::string *out) const override;
int deserialize(const void *data, size_t len) override;
//! Rotator type tag (kFht = 1).
RotateType rotate_type() const {
return RotateType::kFht;
}
~FhtRotator() override;
private:
FhtRotator() = default;
//! Largest power of 2 <= dim.
static size_t floor_pow2(size_t n);
int in_dim_{0};
int out_dim_{0};
//! Bytes per round: ceil(in_dim / 8). Kept for serialization.
size_t flip_offset_{0};
//! ISA-dispatched rotate/unrotate kernels.
RotatorKernels kernels_{};
//! FHT state (flip_offset, trunc_dim, fac, flip[]) -- single allocation.
FhtCtx *fht_ctx_{nullptr};
static constexpr size_t kByteLen = 8;
};
} // namespace turbo
} // namespace zvec

88
src/turbo/preprocessor/preprocessor.h Normal file
View File

@ -0,0 +1,88 @@
// 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
#include <memory>
#include <string>
#include <zvec/turbo/turbo.h>
namespace zvec {
namespace turbo {
//! Magic number ('ROTR') stamped at the start of a serialized rotator blob.
constexpr uint32_t kRotatorMagic = 0x52544F52u;
//! Current rotator serialization format version.
constexpr uint16_t kRotatorSerVersion = 1;
//! Self-describing, fixed-size header that prefixes every serialized rotator.
//! The type-specific payload (flip signs, rotation matrix, ...) follows
//! immediately after this header.
struct RotatorSerHeader {
uint32_t magic; // kRotatorMagic
uint16_t version; // kRotatorSerVersion
uint16_t rotator_type; // RotateType
uint32_t in_dim; // input dimensionality
uint32_t out_dim; // output dimensionality
uint32_t payload_size; // bytes following the header
uint32_t reserved; // 0, for future use / alignment
};
static_assert(sizeof(RotatorSerHeader) == 24,
"RotatorSerHeader must be 24 bytes");
//! Abstract preprocessor interface.
//!
//! A Preprocessor applies a deterministic, invertible transform to each
//! vector (e.g. random rotation). Concrete subclasses (FhtRotator, ...)
//! implement the actual algorithm.
class Preprocessor {
public:
using Pointer = std::shared_ptr<Preprocessor>;
virtual ~Preprocessor() = default;
//! Input dimensionality accepted by apply().
virtual int in_dim() const = 0;
//! Output dimensionality produced by apply(). A future preprocessor may
//! change dimensionality (out_dim() != in_dim()), but FhtRotator keeps it
//! unchanged: it operates on floor_pow2(in_dim) via the Hadamard transform
//! plus a Kac's walk over the remainder, so out_dim() == in_dim().
virtual int out_dim() const = 0;
//! Forward transform: map an input vector to the preprocessed space.
//! \p out must hold at least out_dim() elements.
virtual void apply(const float *in, float *out) const = 0;
//! Inverse transform: recover the original-space vector from a preprocessed
//! one. \p out must hold at least in_dim() elements.
virtual void apply_inverse(const float *in, float *out) const = 0;
//! Fit / initialize the preprocessor from a contiguous batch of training
//! data. For FhtRotator this generates the random flip-sign arrays.
//! \p data pointer to the first element of the batch.
//! \p num number of vectors in the batch.
//! \p stride byte offset between consecutive vectors (0 => packed).
virtual void train(const void *data, size_t num, size_t stride) = 0;
//! Serialize the preprocessor into a self-contained blob
//! (RotatorSerHeader + payload).
virtual int serialize(std::string *out) const = 0;
//! Deserialize the preprocessor from a raw, possibly mmap-backed buffer.
virtual int deserialize(const void *data, size_t len) = 0;
};
} // namespace turbo
} // namespace zvec

20
src/turbo/quantizer/quantizer.h
View File

@ -28,26 +28,6 @@ namespace turbo {
using namespace zvec::core; using namespace zvec::core;
//! Error code literals mirroring core::IndexError::Code integer values.
//!
//! Turbo quantizer sources use these directly instead of the
//! `IndexError_NotImplemented` / `IndexError_Unsupported` const objects
//! because MSVC's WINDOWS_EXPORT_ALL_SYMBOLS does not export const data
//! with constructors from zvec_shared.dll. zvec_turbo is a static library
//! linked with /WHOLEARCHIVE, so referencing those unexported symbols across
//! the DLL boundary triggers LNK2019 on Windows.
//!
//! IndexError::Code stores -val in its constructor, so NotImplemented(11)
//! yields -11 and Unsupported(12) yields -12.
constexpr int kErrNotImplemented = -11;
constexpr int kErrUnsupported = -12;
//! Magic number ('QTZR') stamped at the start of a serialized quantizer blob.
constexpr uint32_t kQuantizerMagic = 0x52545A51u;
//! Current quantizer serialization format version.
constexpr uint16_t kQuantizerSerVersion = 1;
//! Self-describing, fixed-size header that prefixes every serialized quantizer. //! Self-describing, fixed-size header that prefixes every serialized quantizer.
//! The type-specific payload (scalar params, codebook, rotation matrix, ...) //! The type-specific payload (scalar params, codebook, rotation matrix, ...)
//! follows immediately after this header. //! follows immediately after this header.

57
src/turbo/turbo.cc
View File

@ -12,18 +12,29 @@
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
#include <cassert>
#include <ailego/internal/cpu_features.h> #include <ailego/internal/cpu_features.h>
#include <zvec/turbo/turbo.h> #include <zvec/turbo/turbo.h>
#include "avx2/rotate/fht/fht.h"
#include "avx512/rotate/fht/fht.h"
#include "avx512_vnni/record_quantized_int8/cosine.h" #include "avx512_vnni/record_quantized_int8/cosine.h"
#include "avx512_vnni/record_quantized_int8/squared_euclidean.h" #include "avx512_vnni/record_quantized_int8/squared_euclidean.h"
#include "avx512_vnni/uniform_int8/quantize.h" #include "avx512_vnni/uniform_int8/quantize.h"
#include "avx512_vnni/uniform_int8/squared_euclidean.h" #include "avx512_vnni/uniform_int8/squared_euclidean.h"
#include "neon/rotate/fht/fht.h"
#include "scalar/fp32/cosine.h" #include "scalar/fp32/cosine.h"
#include "scalar/fp32/inner_product.h" #include "scalar/fp32/inner_product.h"
#include "scalar/fp32/squared_euclidean.h" #include "scalar/fp32/squared_euclidean.h"
#include "scalar/rotate/fht/fht.h"
#include "sse/rotate/fht/fht.h"
namespace zvec::turbo { namespace zvec::turbo {
// Helper: check if the requested arch matches the target or is auto-detect.
static bool IsArchMatch(CpuArchType actual, CpuArchType target) {
return actual == CpuArchType::kAuto || actual == target;
}
DistanceFunc get_distance_func(MetricType metric_type, DataType data_type, DistanceFunc get_distance_func(MetricType metric_type, DataType data_type,
QuantizeType quantize_type, QuantizeType quantize_type,
CpuArchType cpu_arch_type) { CpuArchType cpu_arch_type) {
@ -45,8 +56,7 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type,
if (data_type == DataType::kInt8) { if (data_type == DataType::kInt8) {
if (quantize_type == QuantizeType::kDefault) { if (quantize_type == QuantizeType::kDefault) {
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI &&
(cpu_arch_type == CpuArchType::kAuto || IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) {
cpu_arch_type == CpuArchType::kAVX512VNNI)) {
if (metric_type == MetricType::kSquaredEuclidean) { if (metric_type == MetricType::kSquaredEuclidean) {
return avx512_vnni::squared_euclidean_int8_distance; return avx512_vnni::squared_euclidean_int8_distance;
} }
@ -57,8 +67,7 @@ DistanceFunc get_distance_func(MetricType metric_type, DataType data_type,
} }
if (quantize_type == QuantizeType::kUniform) { if (quantize_type == QuantizeType::kUniform) {
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI &&
(cpu_arch_type == CpuArchType::kAuto || IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) {
cpu_arch_type == CpuArchType::kAVX512VNNI)) {
if (metric_type == MetricType::kSquaredEuclidean) { if (metric_type == MetricType::kSquaredEuclidean) {
return avx512_vnni::uniform_squared_euclidean_int8_distance; return avx512_vnni::uniform_squared_euclidean_int8_distance;
} }
@ -90,8 +99,7 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type,
if (data_type == DataType::kInt8) { if (data_type == DataType::kInt8) {
if (quantize_type == QuantizeType::kDefault) { if (quantize_type == QuantizeType::kDefault) {
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI &&
(cpu_arch_type == CpuArchType::kAuto || IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) {
cpu_arch_type == CpuArchType::kAVX512VNNI)) {
if (metric_type == MetricType::kSquaredEuclidean) { if (metric_type == MetricType::kSquaredEuclidean) {
return avx512_vnni::squared_euclidean_int8_batch_distance; return avx512_vnni::squared_euclidean_int8_batch_distance;
} }
@ -102,8 +110,7 @@ BatchDistanceFunc get_batch_distance_func(MetricType metric_type,
} }
if (quantize_type == QuantizeType::kUniform) { if (quantize_type == QuantizeType::kUniform) {
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI &&
(cpu_arch_type == CpuArchType::kAuto || IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) {
cpu_arch_type == CpuArchType::kAVX512VNNI)) {
if (metric_type == MetricType::kSquaredEuclidean) { if (metric_type == MetricType::kSquaredEuclidean) {
return avx512_vnni::uniform_squared_euclidean_int8_batch_distance; return avx512_vnni::uniform_squared_euclidean_int8_batch_distance;
} }
@ -121,8 +128,7 @@ QueryPreprocessFunc get_query_preprocess_func(MetricType metric_type,
if (data_type == DataType::kInt8) { if (data_type == DataType::kInt8) {
if (quantize_type == QuantizeType::kDefault) { if (quantize_type == QuantizeType::kDefault) {
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI &&
(cpu_arch_type == CpuArchType::kAuto || IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) {
cpu_arch_type == CpuArchType::kAVX512VNNI)) {
if (metric_type == MetricType::kSquaredEuclidean) { if (metric_type == MetricType::kSquaredEuclidean) {
return avx512_vnni::squared_euclidean_int8_query_preprocess; return avx512_vnni::squared_euclidean_int8_query_preprocess;
} }
@ -147,4 +153,35 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) {
return nullptr; return nullptr;
} }
RotatorKernels get_rotator_kernels(RotateType rotate_type,
CpuArchType cpu_arch_type) {
switch (rotate_type) {
case RotateType::kFht: {
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F &&
IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) {
return {avx512::fht_rotate_avx512, avx512::fht_unrotate_avx512};
}
if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 &&
IsArchMatch(cpu_arch_type, CpuArchType::kAVX2)) {
return {avx2::fht_rotate_avx2, avx2::fht_unrotate_avx2};
}
if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2 &&
IsArchMatch(cpu_arch_type, CpuArchType::kSSE)) {
return {sse::fht_rotate_sse, sse::fht_unrotate_sse};
}
if (zvec::ailego::internal::CpuFeatures::static_flags_.NEON &&
IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) {
return {neon::fht_rotate_neon, neon::fht_unrotate_neon};
}
return {scalar::fht_rotate, scalar::fht_unrotate};
}
}
// Unsupported RotateType: assert in debug for early detection, but always
// return the scalar kernels so release builds never hand back null function
// pointers (which would crash on the first call).
assert(false && "unsupported RotateType");
return {scalar::fht_rotate, scalar::fht_unrotate};
}
} // namespace zvec::turbo } // namespace zvec::turbo

1
tests/db/index/common/query_params_test.cc
View File

@ -21,7 +21,6 @@ TEST(QueryParamsTest, QueryParamsBaseClass) {
// Test constructor // Test constructor
QueryParams params(IndexType::HNSW); QueryParams params(IndexType::HNSW);
EXPECT_EQ(params.type(), IndexType::HNSW); EXPECT_EQ(params.type(), IndexType::HNSW);
} }
TEST(QueryParamsTest, HnswQueryParams) { TEST(QueryParamsTest, HnswQueryParams) {

450
tests/turbo/turbo_fht_rotator_test.cc Normal file
View File

@ -0,0 +1,450 @@
// 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.
#include <cmath>
#include <cstring>
#include <random>
#include <string>
#include <vector>
#include <gtest/gtest.h>
#include "preprocessor/fht_rotator/fht_rotator.h"
using namespace zvec::turbo;
namespace {
// Helper: fill a vector with random floats.
void fill_random(float *data, size_t dim, std::mt19937 &gen) {
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
for (size_t i = 0; i < dim; ++i) data[i] = dist(gen);
}
// Helper: check round-trip (apply_inverse(apply(x)) == x) within tolerance.
void check_round_trip(const FhtRotator &rot, const std::vector<float> &input,
float tol = 1e-3f) {
const int dim = rot.in_dim();
std::vector<float> rotated(dim);
std::vector<float> recovered(dim);
rot.apply(input.data(), rotated.data());
rot.apply_inverse(rotated.data(), recovered.data());
for (int i = 0; i < dim; ++i) {
EXPECT_NEAR(input[i], recovered[i], tol)
<< "mismatch at i=" << i << " dim=" << dim;
}
}
} // anonymous namespace
// ---------------------------------------------------------------------------
// Power-of-2 dimensions
// ---------------------------------------------------------------------------
TEST(FhtRotator, PowerOf2RoundTrip) {
std::mt19937 gen(42);
for (int dim : {1, 2, 4, 8, 16, 32, 64, 128, 256}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot) << "create failed for dim=" << dim;
std::vector<float> input(dim);
fill_random(input.data(), dim, gen);
check_round_trip(*rot, input);
}
}
// ---------------------------------------------------------------------------
// Non-power-of-2 dimensions
// ---------------------------------------------------------------------------
TEST(FhtRotator, NonPowerOf2RoundTrip) {
std::mt19937 gen(123);
for (int dim : {3, 5, 7, 10, 13, 31, 50, 97, 100, 127, 192, 320}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot) << "create failed for dim=" << dim;
std::vector<float> input(dim);
fill_random(input.data(), dim, gen);
check_round_trip(*rot, input);
}
}
// ---------------------------------------------------------------------------
// Serialize / Deserialize round-trip
// ---------------------------------------------------------------------------
TEST(FhtRotator, SerializeDeserialize) {
std::mt19937 gen(999);
for (int dim : {32, 97, 128}) {
// Build original rotator.
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot);
// Serialize.
std::string blob;
ASSERT_EQ(0, rot->serialize(&blob));
ASSERT_GT(blob.size(), sizeof(RotatorSerHeader));
// Restore from blob.
auto rot2 = FhtRotator::from_blob(blob.data(), blob.size());
ASSERT_TRUE(rot2) << "from_blob failed for dim=" << dim;
// Dimensions must match.
EXPECT_EQ(rot2->in_dim(), dim);
EXPECT_EQ(rot2->out_dim(), dim);
// Round-trip via the restored rotator must produce the same result
// as the original (same flip signs).
std::vector<float> input(dim);
fill_random(input.data(), dim, gen);
std::vector<float> r1(dim), r2(dim);
rot->apply(input.data(), r1.data());
rot2->apply(input.data(), r2.data());
for (int i = 0; i < dim; ++i) {
EXPECT_FLOAT_EQ(r1[i], r2[i]) << "apply mismatch at i=" << i;
}
// Inverse via restored rotator must recover the input.
check_round_trip(*rot2, input);
}
}
// ---------------------------------------------------------------------------
// Dimension preserved
// ---------------------------------------------------------------------------
TEST(FhtRotator, DimensionPreserved) {
for (int dim : {1, 7, 64, 97, 128}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot);
EXPECT_EQ(rot->in_dim(), dim);
EXPECT_EQ(rot->out_dim(), dim);
}
}
// ---------------------------------------------------------------------------
// Train generates non-zero flip signs
// ---------------------------------------------------------------------------
TEST(FhtRotator, CreateGeneratesFlip) {
for (int dim : {8, 64, 97}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot);
// After create, flip is already populated, so serialize must succeed.
std::string blob;
EXPECT_EQ(0, rot->serialize(&blob))
<< "serialize failed after create for dim=" << dim;
// Verify the serialized structure deterministically. Copy the header into
// an aligned local first: blob.data() may be unaligned and dereferencing a
// reinterpret_cast<const RotatorSerHeader *> would be undefined behavior on
// architectures that require alignment.
RotatorSerHeader hdr;
std::memcpy(&hdr, blob.data(), sizeof(RotatorSerHeader));
EXPECT_EQ(hdr.magic, kRotatorMagic);
EXPECT_EQ(hdr.version, kRotatorSerVersion);
EXPECT_EQ(static_cast<RotateType>(hdr.rotator_type), RotateType::kFht);
EXPECT_EQ(static_cast<int>(hdr.in_dim), dim);
EXPECT_EQ(static_cast<int>(hdr.out_dim), dim);
// Payload holds 4 rounds of ceil(dim/8) flip bytes. Checking the exact size
// (and total blob length) keeps the test deterministic: we avoid asserting
// on random bit values, which could theoretically be all-zero on a platform
// where std::random_device has low entropy.
const uint32_t expected_flip_size =
4u * ((static_cast<uint32_t>(dim) + 7u) / 8u);
EXPECT_EQ(hdr.payload_size, expected_flip_size);
EXPECT_EQ(blob.size(), sizeof(RotatorSerHeader) + expected_flip_size);
}
}
// ---------------------------------------------------------------------------
// Create with invalid dimension returns nullptr
// ---------------------------------------------------------------------------
TEST(FhtRotator, InvalidDimension) {
EXPECT_EQ(FhtRotator::create(0), nullptr);
EXPECT_EQ(FhtRotator::create(-1), nullptr);
}
// ---------------------------------------------------------------------------
// from_blob with malformed input returns nullptr
// ---------------------------------------------------------------------------
TEST(FhtRotator, FromBlobMalformed) {
EXPECT_EQ(FhtRotator::from_blob(nullptr, 0), nullptr);
// Too short.
char buf[4] = {};
EXPECT_EQ(FhtRotator::from_blob(buf, sizeof(buf)), nullptr);
// Wrong magic.
RotatorSerHeader hdr{};
hdr.magic = 0xDEADBEEF;
hdr.version = kRotatorSerVersion;
hdr.rotator_type = static_cast<uint16_t>(RotateType::kFht);
hdr.payload_size = 0;
EXPECT_EQ(FhtRotator::from_blob(&hdr, sizeof(hdr)), nullptr);
}
// ---------------------------------------------------------------------------
// L2 distance preserved (orthogonal transform)
// ---------------------------------------------------------------------------
TEST(FhtRotator, L2DistancePreserved) {
std::mt19937 gen(2024);
for (int dim : {32, 64, 97, 128}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot);
const int N = 50;
std::vector<std::vector<float>> raw(N, std::vector<float>(dim));
std::vector<std::vector<float>> rotated(N, std::vector<float>(dim));
for (int i = 0; i < N; ++i) {
fill_random(raw[i].data(), dim, gen);
rot->apply(raw[i].data(), rotated[i].data());
}
// Check that ||rotated[i] - rotated[j]|| ≈ ||raw[i] - raw[j]||.
for (int i = 1; i < N; ++i) {
float d_raw = 0.0f, d_rot = 0.0f;
for (int j = 0; j < dim; ++j) {
float dr = raw[i][j] - raw[0][j];
float dt = rotated[i][j] - rotated[0][j];
d_raw += dr * dr;
d_rot += dt * dt;
}
EXPECT_NEAR(d_raw, d_rot, 1e-2f)
<< "L2 mismatch for dim=" << dim << " i=" << i;
}
}
}
// ---------------------------------------------------------------------------
// Cosine distance preserved (orthogonal transform)
// ---------------------------------------------------------------------------
TEST(FhtRotator, CosineDistancePreserved) {
std::mt19937 gen(777);
auto cosine_dist = [](const float *a, const float *b, int dim) {
float dot = 0, na = 0, nb = 0;
for (int i = 0; i < dim; ++i) {
dot += a[i] * b[i];
na += a[i] * a[i];
nb += b[i] * b[i];
}
float denom = std::sqrt(na) * std::sqrt(nb);
return (denom < 1e-12f) ? 1.0f : 1.0f - dot / denom;
};
for (int dim : {32, 97, 128}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot);
const int N = 50;
std::vector<std::vector<float>> raw(N, std::vector<float>(dim));
std::vector<std::vector<float>> rotated(N, std::vector<float>(dim));
for (int i = 0; i < N; ++i) {
fill_random(raw[i].data(), dim, gen);
rot->apply(raw[i].data(), rotated[i].data());
}
for (int i = 1; i < N; ++i) {
float d_raw = cosine_dist(raw[i].data(), raw[0].data(), dim);
float d_rot = cosine_dist(rotated[i].data(), rotated[0].data(), dim);
EXPECT_NEAR(d_raw, d_rot, 1e-3f)
<< "Cosine mismatch for dim=" << dim << " i=" << i;
}
}
}
// ---------------------------------------------------------------------------
// Apply is non-trivial (not identity)
// ---------------------------------------------------------------------------
TEST(FhtRotator, ApplyIsNonTrivial) {
std::mt19937 gen(42);
for (int dim : {32, 97, 128}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot);
std::vector<float> input(dim);
fill_random(input.data(), dim, gen);
std::vector<float> output(dim);
rot->apply(input.data(), output.data());
// At least some elements should differ from the input.
bool any_diff = false;
for (int i = 0; i < dim; ++i) {
if (std::abs(input[i] - output[i]) > 1e-6f) {
any_diff = true;
break;
}
}
EXPECT_TRUE(any_diff) << "apply is identity for dim=" << dim;
}
}
// ---------------------------------------------------------------------------
// Apply is deterministic (same input → same output)
// ---------------------------------------------------------------------------
TEST(FhtRotator, ApplyDeterministic) {
std::mt19937 gen(55);
for (int dim : {32, 97, 128}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot);
std::vector<float> input(dim);
fill_random(input.data(), dim, gen);
std::vector<float> r1(dim), r2(dim);
rot->apply(input.data(), r1.data());
rot->apply(input.data(), r2.data());
for (int i = 0; i < dim; ++i) {
EXPECT_FLOAT_EQ(r1[i], r2[i])
<< "non-deterministic apply at i=" << i << " dim=" << dim;
}
}
}
// ---------------------------------------------------------------------------
// Deserialize on existing object (init → serialize → deserialize on new object)
// ---------------------------------------------------------------------------
TEST(FhtRotator, DeserializeOnExistingObject) {
std::mt19937 gen(314);
for (int dim : {32, 97, 128}) {
// Build original.
auto rot1 = FhtRotator::create(dim);
ASSERT_TRUE(rot1);
std::string blob;
ASSERT_EQ(0, rot1->serialize(&blob));
// Create a fresh rotator, then call deserialize() on it.
auto rot2 = FhtRotator::create(dim);
ASSERT_TRUE(rot2);
ASSERT_EQ(0, rot2->deserialize(blob.data(), blob.size()));
EXPECT_EQ(rot2->in_dim(), dim);
EXPECT_EQ(rot2->out_dim(), dim);
// Both rotators should produce identical results.
std::vector<float> input(dim);
fill_random(input.data(), dim, gen);
std::vector<float> r1(dim), r2(dim);
rot1->apply(input.data(), r1.data());
rot2->apply(input.data(), r2.data());
for (int i = 0; i < dim; ++i) {
EXPECT_FLOAT_EQ(r1[i], r2[i])
<< "apply mismatch at i=" << i << " dim=" << dim;
}
// Inverse via rot2 should recover input.
check_round_trip(*rot2, input);
}
}
// ---------------------------------------------------------------------------
// Deserialize with truncated payload fails
// ---------------------------------------------------------------------------
TEST(FhtRotator, DeserializeTruncatedPayload) {
std::mt19937 gen(42);
auto rot = FhtRotator::create(64);
ASSERT_TRUE(rot);
std::string blob;
ASSERT_EQ(0, rot->serialize(&blob));
// Truncate the blob: keep header but cut half the payload.
// Copy the header into an aligned local before reading fields (blob.data()
// may be unaligned; a direct reinterpret_cast dereference is UB on
// alignment-sensitive architectures).
RotatorSerHeader hdr;
std::memcpy(&hdr, blob.data(), sizeof(RotatorSerHeader));
size_t truncated_len = sizeof(RotatorSerHeader) + hdr.payload_size / 2;
auto rot2 = FhtRotator::create(64);
ASSERT_TRUE(rot2);
EXPECT_NE(0, rot2->deserialize(blob.data(), truncated_len));
// Also test from_blob with truncated data.
EXPECT_EQ(FhtRotator::from_blob(blob.data(), truncated_len), nullptr);
}
// ---------------------------------------------------------------------------
// Large dimension stress test
// ---------------------------------------------------------------------------
TEST(FhtRotator, LargeDimension) {
std::mt19937 gen(2025);
for (int dim : {1024, 2048, 4096}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot) << "create failed for dim=" << dim;
std::vector<float> input(dim);
fill_random(input.data(), dim, gen);
check_round_trip(*rot, input, 1e-2f);
// Verify serialize/deserialize round-trip.
std::string blob;
ASSERT_EQ(0, rot->serialize(&blob));
auto rot2 = FhtRotator::from_blob(blob.data(), blob.size());
ASSERT_TRUE(rot2);
std::vector<float> r1(dim), r2(dim);
rot->apply(input.data(), r1.data());
rot2->apply(input.data(), r2.data());
for (int i = 0; i < dim; ++i) {
EXPECT_FLOAT_EQ(r1[i], r2[i]) << "mismatch at i=" << i << " dim=" << dim;
}
}
}
// ---------------------------------------------------------------------------
// Norm preserved (orthogonal transform preserves vector norm)
// ---------------------------------------------------------------------------
TEST(FhtRotator, NormPreserved) {
std::mt19937 gen(99);
for (int dim : {32, 97, 128}) {
auto rot = FhtRotator::create(dim);
ASSERT_TRUE(rot);
std::vector<float> input(dim);
fill_random(input.data(), dim, gen);
float norm_in = 0.0f;
for (int i = 0; i < dim; ++i) norm_in += input[i] * input[i];
std::vector<float> output(dim);
rot->apply(input.data(), output.data());
float norm_out = 0.0f;
for (int i = 0; i < dim; ++i) norm_out += output[i] * output[i];
EXPECT_NEAR(norm_in, norm_out, 1e-2f)
<< "norm not preserved for dim=" << dim;
}
}