diff --git a/.gitignore b/.gitignore index 65859ca54..99623a413 100644 --- a/.gitignore +++ b/.gitignore @@ -60,9 +60,9 @@ allure-* !build_android.sh !build_ios.sh -# congfig +# config doc/ config/ examples/python/ examples/c_api/ -logs/ \ No newline at end of file +logs/ diff --git a/src/include/zvec/turbo/turbo.h b/src/include/zvec/turbo/turbo.h index e6f1587ac..8b48e36a6 100644 --- a/src/include/zvec/turbo/turbo.h +++ b/src/include/zvec/turbo/turbo.h @@ -13,6 +13,8 @@ // limitations under the License. #pragma once +#include +#include #include #include @@ -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 diff --git a/src/turbo/CMakeLists.txt b/src/turbo/CMakeLists.txt index 7ec9c0e5d..b5a116d1b 100644 --- a/src/turbo/CMakeLists.txt +++ b/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}" ) diff --git a/src/turbo/distance/avx2/rotate/fht/fht.cc b/src/turbo/distance/avx2/rotate/fht/fht.cc new file mode 100644 index 000000000..a77dc087b --- /dev/null +++ b/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 +#include +#include +#include +#include +#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__ diff --git a/src/turbo/distance/avx2/rotate/fht/fht.h b/src/turbo/distance/avx2/rotate/fht/fht.h new file mode 100644 index 000000000..2fb011236 --- /dev/null +++ b/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 +#include + +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__ diff --git a/src/turbo/distance/avx512/rotate/fht/fht.cc b/src/turbo/distance/avx512/rotate/fht/fht.cc new file mode 100644 index 000000000..9a237214d --- /dev/null +++ b/src/turbo/distance/avx512/rotate/fht/fht.cc @@ -0,0 +1,159 @@ +// 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(__AVX512F__) + +#include "fht.h" +#include +#include +#include +#include +#include +#include "common/fht_common.h" + +namespace zvec::turbo::avx512 { + +void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) { + size_t simd_end = dim & ~63u; + constexpr size_t kChunk = 64; + const __m512 sign_flip = _mm512_castsi512_ps(_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); + __m512 v0 = _mm512_loadu_ps(&data[i]); + v0 = _mm512_mask_xor_ps(v0, m0, v0, sign_flip); + _mm512_storeu_ps(&data[i], v0); + __m512 v1 = _mm512_loadu_ps(&data[i + 16]); + v1 = _mm512_mask_xor_ps(v1, m1, v1, sign_flip); + _mm512_storeu_ps(&data[i + 16], v1); + __m512 v2 = _mm512_loadu_ps(&data[i + 32]); + v2 = _mm512_mask_xor_ps(v2, m2, v2, sign_flip); + _mm512_storeu_ps(&data[i + 32], v2); + __m512 v3 = _mm512_loadu_ps(&data[i + 48]); + v3 = _mm512_mask_xor_ps(v3, m3, v3, sign_flip); + _mm512_storeu_ps(&data[i + 48], v3); + } + // 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_avx512(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 & ~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); + } +} + +void fht_inv_kacs_walk_avx512(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 & ~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; + } +} + +void fht_inplace_avx512(float *data, size_t n) { + 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; + } + } + } +} + +void fht_vec_rescale_avx512(float *data, size_t n, float factor) { + 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; + } +} + +void fht_rotate_avx512(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + 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(out, in_dim, ctx, kPrim); +} + +void fht_unrotate_avx512(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + 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(out, in_dim, ctx, kPrim); +} + +} // namespace zvec::turbo::avx512 + +#endif // __AVX512F__ diff --git a/src/turbo/distance/avx512/rotate/fht/fht.h b/src/turbo/distance/avx512/rotate/fht/fht.h new file mode 100644 index 000000000..bcb3a9fe9 --- /dev/null +++ b/src/turbo/distance/avx512/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(__AVX512F__) + +#include +#include + +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 + +#endif // __AVX512F__ diff --git a/src/turbo/distance/common/fht_common.h b/src/turbo/distance/common/fht_common.h new file mode 100644 index 000000000..556338863 --- /dev/null +++ b/src/turbo/distance/common/fht_common.h @@ -0,0 +1,124 @@ +// 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 +#include + +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(float *data, size_t dim, void *ctx, + const FhtPrimitives &p) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(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(float *data, size_t dim, void *ctx, + const FhtPrimitives &p) { + auto *base = reinterpret_cast(ctx); + const size_t flip_offset = *reinterpret_cast(base); + const size_t trunc_dim = *reinterpret_cast(base + 8); + const float fac = *reinterpret_cast(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(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 diff --git a/src/turbo/distance/neon/rotate/fht/fht.cc b/src/turbo/distance/neon/rotate/fht/fht.cc new file mode 100644 index 000000000..cc736a0a2 --- /dev/null +++ b/src/turbo/distance/neon/rotate/fht/fht.cc @@ -0,0 +1,137 @@ +// 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(__ARM_NEON) && defined(__aarch64__) + +#include "fht.h" +#include +#include +#include +#include +#include +#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) { + 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]; + } + } +} + +void fht_kacs_walk_neon(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 & ~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); + } +} + +void fht_inv_kacs_walk_neon(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 & ~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; + } +} + +void fht_vec_rescale_neon(float *data, size_t n, float factor) { + 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; + } +} + +void fht_rotate_neon(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + 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(out, in_dim, ctx, kPrim); +} + +void fht_unrotate_neon(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + 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(out, in_dim, ctx, kPrim); +} + +} // namespace zvec::turbo::neon + +#endif // __ARM_NEON && __aarch64__ diff --git a/src/turbo/distance/neon/rotate/fht/fht.h b/src/turbo/distance/neon/rotate/fht/fht.h new file mode 100644 index 000000000..ee6bcacdc --- /dev/null +++ b/src/turbo/distance/neon/rotate/fht/fht.h @@ -0,0 +1,46 @@ +// 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(__ARM_NEON) && defined(__aarch64__) + +#include +#include + +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 + +#endif // __ARM_NEON && __aarch64__ diff --git a/src/turbo/distance/scalar/rotate/fht/fht.cc b/src/turbo/distance/scalar/rotate/fht/fht.cc new file mode 100644 index 000000000..1d4e6e779 --- /dev/null +++ b/src/turbo/distance/scalar/rotate/fht/fht.cc @@ -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 +#include +#include +#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(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(out, in_dim, ctx, kPrim); +} + +} // namespace zvec::turbo::scalar diff --git a/src/turbo/distance/scalar/rotate/fht/fht.h b/src/turbo/distance/scalar/rotate/fht/fht.h new file mode 100644 index 000000000..c31b82543 --- /dev/null +++ b/src/turbo/distance/scalar/rotate/fht/fht.h @@ -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 +#include + +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 diff --git a/src/turbo/distance/sse/rotate/fht/fht.cc b/src/turbo/distance/sse/rotate/fht/fht.cc new file mode 100644 index 000000000..279bb3e96 --- /dev/null +++ b/src/turbo/distance/sse/rotate/fht/fht.cc @@ -0,0 +1,136 @@ +// 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(__SSE2__) + +#include "fht.h" +#include +#include +#include +#include +#include +#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) { + 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]; + } + } +} + +void fht_kacs_walk_sse(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 & ~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); + } +} + +void fht_inv_kacs_walk_sse(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 & ~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; + } +} + +void fht_vec_rescale_sse(float *data, size_t n, float factor) { + 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; + } +} + +void fht_rotate_sse(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + 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(out, in_dim, ctx, kPrim); +} + +void fht_unrotate_sse(const float * /*in*/, float *out, size_t in_dim, + size_t /*out_dim*/, void *ctx) { + 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(out, in_dim, ctx, kPrim); +} + +} // namespace zvec::turbo::sse + +#endif // __SSE2__ diff --git a/src/turbo/distance/sse/rotate/fht/fht.h b/src/turbo/distance/sse/rotate/fht/fht.h new file mode 100644 index 000000000..283ced493 --- /dev/null +++ b/src/turbo/distance/sse/rotate/fht/fht.h @@ -0,0 +1,46 @@ +// 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(__SSE2__) + +#include +#include + +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 + +#endif // __SSE2__ diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.cc b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc new file mode 100644 index 000000000..119b4a7ac --- /dev/null +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.cc @@ -0,0 +1,176 @@ +// 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 +#include +#include +#include +#include "quantizer/quantizer.h" + +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(dim) + kByteLen - 1) / kByteLen; + r->kernels_ = get_rotator_kernels(RotateType::kFht); + + const size_t trunc_dim = floor_pow2(static_cast(dim)); + const float fac = 1.0f / std::sqrt(static_cast(trunc_dim)); + const size_t flip_size = 4 * r->flip_offset_; + + // Single allocation: FhtCtx header + trailing flip data. + r->fht_ctx_ = static_cast(std::malloc(sizeof(FhtCtx) + flip_size)); + 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 dist(0, 255); + for (size_t i = 0; i < flip_size; ++i) + r->fht_ctx_->flip[i] = static_cast(dist(gen)); + + return r; +} + +FhtRotator::Pointer FhtRotator::from_blob(const void *data, size_t len) { + if (!data || len < sizeof(RotatorSerHeader)) return nullptr; + + const auto *hdr = reinterpret_cast(data); + if (hdr->magic != kRotatorMagic) return nullptr; + if (hdr->version != kRotatorSerVersion) return nullptr; + if (static_cast(hdr->rotator_type) != RotateType::kFht) { + return nullptr; + } + + Pointer r(new FhtRotator()); + const size_t expected_total = + sizeof(RotatorSerHeader) + static_cast(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 { + std::memcpy(out, in, sizeof(float) * static_cast(in_dim_)); + kernels_.rotate(in, out, static_cast(in_dim_), + static_cast(out_dim_), static_cast(fht_ctx_)); +} + +// --------------------------------------------------------------------------- +// apply_inverse (inverse rotation) +// --------------------------------------------------------------------------- + +void FhtRotator::apply_inverse(const float *in, float *out) const { + std::memcpy(out, in, sizeof(float) * static_cast(in_dim_)); + kernels_.unrotate(in, out, static_cast(in_dim_), + static_cast(out_dim_), + static_cast(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(RotateType::kFht); + hdr.in_dim = static_cast(in_dim_); + hdr.out_dim = static_cast(out_dim_); + hdr.payload_size = static_cast(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; + + const auto *hdr = reinterpret_cast(data); + if (hdr->magic != kRotatorMagic) return kErrUnsupported; + if (hdr->version != kRotatorSerVersion) return kErrUnsupported; + if (static_cast(hdr->rotator_type) != RotateType::kFht) { + return kErrUnsupported; + } + + const size_t total = sizeof(RotatorSerHeader) + hdr->payload_size; + if (len < total) return kErrInvalidArgument; + + in_dim_ = static_cast(hdr->in_dim); + out_dim_ = static_cast(hdr->out_dim); + flip_offset_ = (static_cast(in_dim_) + kByteLen - 1) / kByteLen; + kernels_ = get_rotator_kernels(RotateType::kFht); + + const size_t trunc_dim = floor_pow2(static_cast(in_dim_)); + const float fac = 1.0f / std::sqrt(static_cast(trunc_dim)); + + // Free old ctx, allocate new one with trailing flip data. + std::free(fht_ctx_); + fht_ctx_ = + static_cast(std::malloc(sizeof(FhtCtx) + hdr->payload_size)); + fht_ctx_->flip_offset = flip_offset_; + fht_ctx_->trunc_dim = trunc_dim; + fht_ctx_->fac = fac; + std::memcpy(fht_ctx_->flip, + reinterpret_cast(data) + sizeof(RotatorSerHeader), + hdr->payload_size); + + return 0; +} + +} // namespace turbo +} // namespace zvec diff --git a/src/turbo/preprocessor/fht_rotator/fht_rotator.h b/src/turbo/preprocessor/fht_rotator/fht_rotator.h new file mode 100644 index 000000000..0ff6879d4 --- /dev/null +++ b/src/turbo/preprocessor/fht_rotator/fht_rotator.h @@ -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 +#include +#include +#include +#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; + + //! 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 diff --git a/src/turbo/preprocessor/preprocessor.h b/src/turbo/preprocessor/preprocessor.h new file mode 100644 index 000000000..c51f19b14 --- /dev/null +++ b/src/turbo/preprocessor/preprocessor.h @@ -0,0 +1,87 @@ +// 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 +#include +#include + +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; + + virtual ~Preprocessor() = default; + + //! Input dimensionality accepted by apply(). + virtual int in_dim() const = 0; + + //! Output dimensionality produced by apply(). May differ from in_dim() + //! (e.g. FHT pads to the next power of two), but for FhtRotator they are + //! equal. + 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 diff --git a/src/turbo/quantizer/quantizer.h b/src/turbo/quantizer/quantizer.h index 3b7a7d944..2a9d54cb5 100644 --- a/src/turbo/quantizer/quantizer.h +++ b/src/turbo/quantizer/quantizer.h @@ -39,8 +39,10 @@ using namespace zvec::core; //! //! 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; diff --git a/src/turbo/turbo.cc b/src/turbo/turbo.cc index 4fe8309d0..493b720d7 100644 --- a/src/turbo/turbo.cc +++ b/src/turbo/turbo.cc @@ -12,18 +12,29 @@ // See the License for the specific language governing permissions and // limitations under the License. +#include #include #include +#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/squared_euclidean.h" #include "avx512_vnni/uniform_int8/quantize.h" #include "avx512_vnni/uniform_int8/squared_euclidean.h" +#include "neon/rotate/fht/fht.h" #include "scalar/fp32/cosine.h" #include "scalar/fp32/inner_product.h" #include "scalar/fp32/squared_euclidean.h" +#include "scalar/rotate/fht/fht.h" +#include "sse/rotate/fht/fht.h" 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, QuantizeType quantize_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 (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { 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 (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { 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 (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { 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 (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { 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 (quantize_type == QuantizeType::kDefault) { if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512_VNNI && - (cpu_arch_type == CpuArchType::kAuto || - cpu_arch_type == CpuArchType::kAVX512VNNI)) { + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512VNNI)) { if (metric_type == MetricType::kSquaredEuclidean) { return avx512_vnni::squared_euclidean_int8_query_preprocess; } @@ -147,4 +153,41 @@ UniformQuantizeFunc get_uniform_quantize_func(DataType data_type) { return nullptr; } +RotatorKernels get_rotator_kernels(RotateType rotate_type, + CpuArchType cpu_arch_type) { + switch (rotate_type) { + case RotateType::kFht: { +#if defined(__AVX512F__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX512F && + zvec::ailego::internal::CpuFeatures::static_flags_.AVX512DQ && + IsArchMatch(cpu_arch_type, CpuArchType::kAVX512)) { + return {avx512::fht_rotate_avx512, avx512::fht_unrotate_avx512}; + } +#endif +#if defined(__AVX2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.AVX2 && + IsArchMatch(cpu_arch_type, CpuArchType::kAVX2)) { + return {avx2::fht_rotate_avx2, avx2::fht_unrotate_avx2}; + } +#endif +#if defined(__SSE2__) + if (zvec::ailego::internal::CpuFeatures::static_flags_.SSE2 && + IsArchMatch(cpu_arch_type, CpuArchType::kSSE)) { + return {sse::fht_rotate_sse, sse::fht_unrotate_sse}; + } +#endif +#if defined(__ARM_NEON) && defined(__aarch64__) + if (IsArchMatch(cpu_arch_type, CpuArchType::kNEON)) { + return {neon::fht_rotate_neon, neon::fht_unrotate_neon}; + } +#endif + return {scalar::fht_rotate, scalar::fht_unrotate}; + } + } + + // Fallback (unreachable for valid RotateType values). + assert(false && "unsupported RotateType"); + return {}; +} + } // namespace zvec::turbo diff --git a/tests/db/index/common/query_params_test.cc b/tests/db/index/common/query_params_test.cc index ceb510746..13e8bd00b 100644 --- a/tests/db/index/common/query_params_test.cc +++ b/tests/db/index/common/query_params_test.cc @@ -21,7 +21,6 @@ TEST(QueryParamsTest, QueryParamsBaseClass) { // Test constructor QueryParams params(IndexType::HNSW); EXPECT_EQ(params.type(), IndexType::HNSW); - } TEST(QueryParamsTest, HnswQueryParams) { diff --git a/tests/turbo/turbo_fht_rotator_test.cc b/tests/turbo/turbo_fht_rotator_test.cc new file mode 100644 index 000000000..1e454ed9b --- /dev/null +++ b/tests/turbo/turbo_fht_rotator_test.cc @@ -0,0 +1,447 @@ +// 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 +#include +#include +#include +#include +#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 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 &input, + float tol = 1e-3f) { + const int dim = rot.in_dim(); + std::vector rotated(dim); + std::vector 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 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 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 input(dim); + fill_random(input.data(), dim, gen); + + std::vector 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 payload is not all zeros (extremely unlikely for random bits). + const auto *hdr = reinterpret_cast(blob.data()); + EXPECT_EQ(hdr->magic, kRotatorMagic); + EXPECT_EQ(hdr->version, kRotatorSerVersion); + EXPECT_EQ(static_cast(hdr->rotator_type), RotateType::kFht); + EXPECT_EQ(static_cast(hdr->in_dim), dim); + EXPECT_EQ(static_cast(hdr->out_dim), dim); + EXPECT_GT(hdr->payload_size, 0u); + + // Check that the flip payload is not all-zero. + const uint8_t *payload = reinterpret_cast(blob.data()) + + sizeof(RotatorSerHeader); + bool any_nonzero = false; + for (uint32_t i = 0; i < hdr->payload_size; ++i) { + if (payload[i] != 0) { + any_nonzero = true; + break; + } + } + EXPECT_TRUE(any_nonzero) << "flip payload is all-zero for dim=" << dim; + } +} + +// --------------------------------------------------------------------------- +// 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(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); + std::uniform_real_distribution dist(-1.0f, 1.0f); + + for (int dim : {32, 64, 97, 128}) { + auto rot = FhtRotator::create(dim); + ASSERT_TRUE(rot); + + const int N = 50; + std::vector> raw(N, std::vector(dim)); + std::vector> rotated(N, std::vector(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); + std::uniform_real_distribution dist(0.1f, 1.0f); + + 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> raw(N, std::vector(dim)); + std::vector> rotated(N, std::vector(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 input(dim); + fill_random(input.data(), dim, gen); + + std::vector 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 input(dim); + fill_random(input.data(), dim, gen); + + std::vector 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 input(dim); + fill_random(input.data(), dim, gen); + + std::vector 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. + auto *hdr = reinterpret_cast(blob.data()); + 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 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 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 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 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; + } +}