Add a preprocessor in Turbo (#548)
Co-authored-by: ray <rui.xing@alibaba-inc.com>
This commit is contained in:
Zhuanglin Zheng
GitHub
parent
764754427c
commit
a4347e31f1
|
|
@ -60,9 +60,9 @@ allure-*
|
|||
!build_android.sh
|
||||
!build_ios.sh
|
||||
|
||||
# congfig
|
||||
# config
|
||||
doc/
|
||||
config/
|
||||
examples/python/
|
||||
examples/c_api/
|
||||
logs/
|
||||
logs/
|
||||
|
|
|
|||
|
|
@ -23,6 +23,13 @@
|
|||
#include <cpuid.h>
|
||||
#endif
|
||||
|
||||
#if defined(__aarch64__) && defined(__linux__)
|
||||
#include <sys/auxv.h>
|
||||
#ifndef HWCAP_ASIMD
|
||||
#define HWCAP_ASIMD (1 << 1)
|
||||
#endif
|
||||
#endif
|
||||
|
||||
namespace zvec {
|
||||
namespace ailego {
|
||||
namespace internal {
|
||||
|
|
@ -336,6 +343,17 @@ bool CpuFeatures::HYPERVISOR(void) {
|
|||
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) {
|
||||
return ""
|
||||
#if defined(__ARM_NEON)
|
||||
|
|
|
|||
|
|
@ -182,6 +182,9 @@ class CpuFeatures {
|
|||
// ！Running on a hypervisor
|
||||
static bool HYPERVISOR(void);
|
||||
|
||||
//! ARM NEON (ASIMD) support
|
||||
static bool NEON(void);
|
||||
|
||||
//! Intrinsics of compiling
|
||||
static const char *Intrinsics(void);
|
||||
|
||||
|
|
@ -361,6 +364,9 @@ class CpuFeatures {
|
|||
|
||||
// ！Running on a hypervisor
|
||||
bool HYPERVISOR = CpuFeatures::HYPERVISOR();
|
||||
|
||||
//! ARM NEON (ASIMD) support
|
||||
bool NEON = CpuFeatures::NEON();
|
||||
};
|
||||
static StaticFlags static_flags_;
|
||||
};
|
||||
|
|
|
|||
|
|
@ -13,11 +13,35 @@
|
|||
// limitations under the License.
|
||||
#pragma once
|
||||
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <functional>
|
||||
#include <zvec/ailego/math_batch/utils.h>
|
||||
|
||||
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 =
|
||||
std::function<void(const void *m, const void *q, size_t dim, float *out)>;
|
||||
using BatchDistanceFunc = std::function<void(
|
||||
|
|
@ -33,6 +57,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 +96,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 +136,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
|
||||
|
|
|
|||
|
|
@ -5,23 +5,40 @@ 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 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
|
||||
COMPILE_FLAGS "${TURBO_MARCH_FLAG_AVX512}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -28,26 +28,6 @@ namespace turbo {
|
|||
|
||||
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.
|
||||
//! The type-specific payload (scalar params, codebook, rotation matrix, ...)
|
||||
//! follows immediately after this header.
|
||||
|
|
|
|||
|
|
@ -12,18 +12,29 @@
|
|||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <cassert>
|
||||
#include <ailego/internal/cpu_features.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/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,35 @@ 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 (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
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@ TEST(QueryParamsTest, QueryParamsBaseClass) {
|
|||
// Test constructor
|
||||
QueryParams params(IndexType::HNSW);
|
||||
EXPECT_EQ(params.type(), IndexType::HNSW);
|
||||
|
||||
}
|
||||
|
||||
TEST(QueryParamsTest, HnswQueryParams) {
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue