rotate: add an optional random rotation feature in INT8/INT4 quantization method (#483)
Co-authored-by: rayx <rui.xing@alibaba-inc.com> Co-authored-by: Jalin Wang <wangjianning.wjn@alibaba-inc.com> Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
This commit is contained in:
parent
cdef8686c7
commit
c54f5e16e9
|
|
@ -60,3 +60,9 @@ allure-*
|
|||
!build_android.sh
|
||||
!build_ios.sh
|
||||
|
||||
# congfig
|
||||
doc/
|
||||
config/
|
||||
examples/python/
|
||||
examples/c_api/
|
||||
logs/
|
||||
|
|
@ -85,6 +85,12 @@ int main() {
|
|||
zvec_index_params_set_metric_type(hnsw_params_fast, ZVEC_METRIC_TYPE_L2);
|
||||
zvec_index_params_set_hnsw_params(hnsw_params_fast, 16, 100);
|
||||
|
||||
// Demonstrate INT8 quantization with random rotation preprocessing
|
||||
// (enable_rotate rotates vectors before INT8 quantization to reduce error)
|
||||
zvec_index_params_set_quantize_type(hnsw_params_fast,
|
||||
ZVEC_QUANTIZE_TYPE_INT8);
|
||||
zvec_index_params_set_quantizer_enable_rotate(hnsw_params_fast, true);
|
||||
|
||||
zvec_index_params_t *hnsw_params_balanced =
|
||||
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
|
||||
if (!hnsw_params_balanced) {
|
||||
|
|
|
|||
|
|
@ -36,6 +36,7 @@ from zvec import (
|
|||
IndexType,
|
||||
MetricType,
|
||||
QuantizeType,
|
||||
QuantizerParam,
|
||||
DataType,
|
||||
VectorSchema,
|
||||
)
|
||||
|
|
@ -452,3 +453,88 @@ class TestVectorQueryDeprecated:
|
|||
warnings.simplefilter("always")
|
||||
vq = VectorQuery(field_name="embedding", id="doc123")
|
||||
assert isinstance(vq, Query)
|
||||
|
||||
|
||||
# ----------------------------
|
||||
# QuantizerParam Test Case
|
||||
# ----------------------------
|
||||
|
||||
|
||||
class TestQuantizerParam:
|
||||
def test_default(self):
|
||||
qp = QuantizerParam()
|
||||
assert qp.enable_rotate is False
|
||||
|
||||
def test_enable_rotate_true(self):
|
||||
qp = QuantizerParam(enable_rotate=True)
|
||||
assert qp.enable_rotate is True
|
||||
|
||||
def test_enable_rotate_false(self):
|
||||
qp = QuantizerParam(enable_rotate=False)
|
||||
assert qp.enable_rotate is False
|
||||
|
||||
def test_equality(self):
|
||||
qp1 = QuantizerParam(enable_rotate=True)
|
||||
qp2 = QuantizerParam(enable_rotate=True)
|
||||
qp3 = QuantizerParam(enable_rotate=False)
|
||||
assert qp1 == qp2
|
||||
assert qp1 != qp3
|
||||
|
||||
def test_to_dict(self):
|
||||
qp = QuantizerParam(enable_rotate=True)
|
||||
d = qp.to_dict()
|
||||
assert isinstance(d, dict)
|
||||
assert d.get("enable_rotate") is True
|
||||
|
||||
def test_repr(self):
|
||||
qp = QuantizerParam(enable_rotate=True)
|
||||
r = repr(qp)
|
||||
assert "enable_rotate" in r or "QuantizerParam" in r
|
||||
|
||||
def test_pickle_roundtrip(self):
|
||||
import pickle
|
||||
|
||||
qp = QuantizerParam(enable_rotate=True)
|
||||
data = pickle.dumps(qp)
|
||||
qp2 = pickle.loads(data)
|
||||
assert qp2.enable_rotate is True
|
||||
assert qp == qp2
|
||||
|
||||
|
||||
# ----------------------------
|
||||
# HnswIndexParam with QuantizerParam
|
||||
# ----------------------------
|
||||
|
||||
|
||||
class TestHnswIndexParamQuantizer:
|
||||
def test_default_quantizer_param(self):
|
||||
param = HnswIndexParam()
|
||||
assert param.quantizer_param is not None
|
||||
assert param.quantizer_param.enable_rotate is False
|
||||
|
||||
def test_with_quantizer_param(self):
|
||||
qp = QuantizerParam(enable_rotate=True)
|
||||
param = HnswIndexParam(
|
||||
metric_type=MetricType.L2,
|
||||
quantize_type=QuantizeType.INT8,
|
||||
quantizer_param=qp,
|
||||
)
|
||||
assert param.quantizer_param.enable_rotate is True
|
||||
assert param.quantize_type == QuantizeType.INT8
|
||||
|
||||
|
||||
# ----------------------------
|
||||
# FlatIndexParam with QuantizerParam
|
||||
# ----------------------------
|
||||
|
||||
|
||||
class TestFlatIndexParamQuantizer:
|
||||
def test_with_quantizer_param(self):
|
||||
qp = QuantizerParam(enable_rotate=True)
|
||||
param = FlatIndexParam(
|
||||
metric_type=MetricType.L2,
|
||||
quantize_type=QuantizeType.INT8,
|
||||
quantizer_param=qp,
|
||||
)
|
||||
assert param.quantizer_param.enable_rotate is True
|
||||
assert param.quantize_type == QuantizeType.INT8
|
||||
|
|
|
|||
|
|
@ -108,6 +108,7 @@ from .model.param import (
|
|||
IVFIndexParam,
|
||||
IVFQueryParam,
|
||||
OptimizeOption,
|
||||
QuantizerParam,
|
||||
VamanaIndexParam,
|
||||
VamanaQueryParam,
|
||||
)
|
||||
|
|
@ -171,6 +172,7 @@ __all__ = [
|
|||
"HnswQueryParam",
|
||||
"HnswRabitqQueryParam",
|
||||
"IVFQueryParam",
|
||||
"QuantizerParam",
|
||||
"VamanaIndexParam",
|
||||
"VamanaQueryParam",
|
||||
# Extensions
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ from .model.param import (
|
|||
IVFIndexParam,
|
||||
IVFQueryParam,
|
||||
OptimizeOption,
|
||||
QuantizerParam,
|
||||
VamanaIndexParam,
|
||||
VamanaQueryParam,
|
||||
)
|
||||
|
|
@ -79,6 +80,7 @@ __all__: list = [
|
|||
"MetricType",
|
||||
"OptimizeOption",
|
||||
"QuantizeType",
|
||||
"QuantizerParam",
|
||||
"Query",
|
||||
"ReRanker",
|
||||
"RrfReRanker",
|
||||
|
|
|
|||
|
|
@ -31,6 +31,7 @@ from zvec._zvec.param import (
|
|||
IVFIndexParam,
|
||||
IVFQueryParam,
|
||||
OptimizeOption,
|
||||
QuantizerParam,
|
||||
VamanaIndexParam,
|
||||
VamanaQueryParam,
|
||||
)
|
||||
|
|
@ -53,6 +54,7 @@ __all__ = [
|
|||
"IndexOption",
|
||||
"InvertIndexParam",
|
||||
"OptimizeOption",
|
||||
"QuantizerParam",
|
||||
"VamanaIndexParam",
|
||||
"VamanaQueryParam",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ __all__: list[str] = [
|
|||
"IndexParam",
|
||||
"InvertIndexParam",
|
||||
"OptimizeOption",
|
||||
"QuantizerParam",
|
||||
"QueryParam",
|
||||
"SegmentOption",
|
||||
"VectorIndexParam",
|
||||
|
|
@ -147,6 +148,8 @@ class FlatIndexParam(VectorIndexParam):
|
|||
quantize_type (QuantizeType): Optional quantization type for vector
|
||||
compression (e.g., FP16, INT8). Use ``QuantizeType.UNDEFINED`` to
|
||||
disable quantization. Default is ``QuantizeType.UNDEFINED``.
|
||||
quantizer_param (QuantizerParam): Optional quantizer parameters. See
|
||||
``QuantizerParam`` for available options. Default is ``QuantizerParam()``.
|
||||
|
||||
Examples:
|
||||
>>> from zvec.typing import MetricType, QuantizeType
|
||||
|
|
@ -163,6 +166,7 @@ class FlatIndexParam(VectorIndexParam):
|
|||
self,
|
||||
metric_type: zvec._zvec.typing.MetricType = ...,
|
||||
quantize_type: zvec._zvec.typing.QuantizeType = ...,
|
||||
quantizer_param: QuantizerParam = ...,
|
||||
) -> None:
|
||||
"""
|
||||
Constructs a FlatIndexParam instance.
|
||||
|
|
@ -171,6 +175,8 @@ class FlatIndexParam(VectorIndexParam):
|
|||
metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.
|
||||
quantize_type (QuantizeType, optional): Vector quantization type.
|
||||
Defaults to QuantizeType.UNDEFINED (no quantization).
|
||||
quantizer_param (QuantizerParam, optional): Quantizer configuration.
|
||||
Defaults to QuantizerParam().
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
|
|
@ -226,6 +232,7 @@ class HnswIndexParam(VectorIndexParam):
|
|||
ef_construction: typing.SupportsInt = 500,
|
||||
quantize_type: zvec._zvec.typing.QuantizeType = ...,
|
||||
use_contiguous_memory: bool = False,
|
||||
quantizer_param: QuantizerParam = ...,
|
||||
) -> None: ...
|
||||
def __repr__(self) -> str: ...
|
||||
def __setstate__(self, arg0: tuple) -> None: ...
|
||||
|
|
@ -283,6 +290,7 @@ class HnswQueryParam(QueryParam):
|
|||
>>> print(params.to_dict() if hasattr(params, 'to_dict') else params)
|
||||
{"type":"HNSW", "ef":300}
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -307,6 +315,7 @@ class HnswQueryParam(QueryParam):
|
|||
- ``prefetch_lines`` (int): Number of 64B cache lines to prefetch
|
||||
per neighbour vector (PL). ``0`` (default) means auto-derive from vector size.
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
def __setstate__(self, arg0: tuple) -> None: ...
|
||||
@property
|
||||
|
|
@ -430,6 +439,7 @@ class HnswRabitqQueryParam(QueryParam):
|
|||
>>> print(params.ef)
|
||||
300
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -448,6 +458,7 @@ class HnswRabitqQueryParam(QueryParam):
|
|||
is_linear (bool, optional): Force linear search. Default is False.
|
||||
is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
def __setstate__(self, arg0: tuple) -> None: ...
|
||||
@property
|
||||
|
|
@ -498,6 +509,7 @@ class IVFIndexParam(VectorIndexParam):
|
|||
n_iters: typing.SupportsInt = 10,
|
||||
use_soar: bool = False,
|
||||
quantize_type: zvec._zvec.typing.QuantizeType = ...,
|
||||
quantizer_param: QuantizerParam = ...,
|
||||
) -> None:
|
||||
"""
|
||||
Constructs an IVFIndexParam instance.
|
||||
|
|
@ -511,6 +523,8 @@ class IVFIndexParam(VectorIndexParam):
|
|||
use_soar (bool, optional): Enable SOAR optimization. Defaults to False.
|
||||
quantize_type (QuantizeType, optional): Vector quantization type.
|
||||
Defaults to QuantizeType.UNDEFINED.
|
||||
quantizer_param (QuantizerParam, optional): Quantizer configuration.
|
||||
Defaults to QuantizerParam().
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
|
|
@ -558,6 +572,7 @@ class IVFQueryParam(QueryParam):
|
|||
>>> print(params.nprobe)
|
||||
20
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
def __init__(self, nprobe: typing.SupportsInt = 10) -> None:
|
||||
"""
|
||||
|
|
@ -567,6 +582,7 @@ class IVFQueryParam(QueryParam):
|
|||
nprobe (int, optional): Number of inverted lists to probe during search.
|
||||
Higher values improve accuracy. Defaults to 10.
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
def __setstate__(self, arg0: tuple) -> None: ...
|
||||
@property
|
||||
|
|
@ -592,6 +608,7 @@ class VamanaIndexParam(VectorIndexParam):
|
|||
Examples:
|
||||
>>> params = VamanaIndexParam(metric_type=MetricType.COSINE, max_degree=64)
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -610,18 +627,23 @@ class VamanaIndexParam(VectorIndexParam):
|
|||
@property
|
||||
def max_degree(self) -> int:
|
||||
"""int: Maximum out-degree (R) of every node in the Vamana graph."""
|
||||
|
||||
@property
|
||||
def search_list_size(self) -> int:
|
||||
"""int: Candidate list size during Vamana graph construction."""
|
||||
|
||||
@property
|
||||
def alpha(self) -> float:
|
||||
"""float: Vamana RobustPrune alpha factor."""
|
||||
|
||||
@property
|
||||
def saturate_graph(self) -> bool:
|
||||
"""bool: Whether to saturate every node to max_degree neighbors."""
|
||||
|
||||
@property
|
||||
def use_contiguous_memory(self) -> bool:
|
||||
"""bool: Whether to allocate a single contiguous memory arena."""
|
||||
|
||||
@property
|
||||
def use_id_map(self) -> bool:
|
||||
"""bool: Reserved flag for engine-level id remapping."""
|
||||
|
|
@ -644,6 +666,7 @@ class VamanaQueryParam(QueryParam):
|
|||
>>> print(params.ef_search)
|
||||
200
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -667,14 +690,17 @@ class VamanaQueryParam(QueryParam):
|
|||
- ``prefetch_lines`` (int): Cache lines to prefetch per vector (PL).
|
||||
``0`` (default) means auto-derive from vector size.
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
def __setstate__(self, arg0: tuple) -> None: ...
|
||||
@property
|
||||
def ef_search(self) -> int:
|
||||
"""int: Size of the dynamic candidate list during Vamana search."""
|
||||
|
||||
@property
|
||||
def prefetch_offset(self) -> int:
|
||||
"""int: Graph prefetch offset used by the Vamana fast path."""
|
||||
|
||||
@property
|
||||
def prefetch_lines(self) -> int:
|
||||
"""int: Override of prefetch cache lines per vector (0=auto)."""
|
||||
|
|
@ -760,6 +786,7 @@ class FtsQueryParam(QueryParam):
|
|||
>>> print(params.default_operator)
|
||||
AND
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -772,6 +799,7 @@ class FtsQueryParam(QueryParam):
|
|||
default_operator (str, optional): Default boolean operator for adjacent
|
||||
bare terms. Supported: "OR", "AND". Defaults to "" (uses engine default).
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
def __setstate__(self, arg0: tuple) -> None: ...
|
||||
@property
|
||||
|
|
@ -865,6 +893,7 @@ class InvertIndexParam(IndexParam):
|
|||
>>> print(config)
|
||||
{'enable_range_optimization': True, 'enable_extended_wildcard': False}
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -880,18 +909,21 @@ class InvertIndexParam(IndexParam):
|
|||
enable_extended_wildcard (bool, optional): If True, enables extended wildcard
|
||||
search including suffix and infix patterns. Defaults to False.
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
def __setstate__(self, arg0: tuple) -> None: ...
|
||||
def to_dict(self) -> dict:
|
||||
"""
|
||||
Convert to dictionary with all fields
|
||||
"""
|
||||
|
||||
@property
|
||||
def enable_extended_wildcard(self) -> bool:
|
||||
"""
|
||||
bool: Whether extended wildcard (suffix and infix) search is enabled.
|
||||
Note: Prefix search is always enabled regardless of this setting.
|
||||
"""
|
||||
|
||||
@property
|
||||
def enable_range_optimization(self) -> bool:
|
||||
"""
|
||||
|
|
@ -948,6 +980,7 @@ class QueryParam:
|
|||
using the index. Useful for debugging or small datasets. Default is False.
|
||||
is_using_refiner (bool, optional): Whether to use refiner for the query. Default is False.
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
def __setstate__(self, arg0: tuple) -> None: ...
|
||||
@property
|
||||
|
|
@ -955,16 +988,19 @@ class QueryParam:
|
|||
"""
|
||||
bool: Whether to bypass the index and use brute-force linear search.
|
||||
"""
|
||||
|
||||
@property
|
||||
def is_using_refiner(self) -> bool:
|
||||
"""
|
||||
bool: Whether to use refiner for the query.
|
||||
"""
|
||||
|
||||
@property
|
||||
def radius(self) -> float:
|
||||
"""
|
||||
IndexType: The type of index this query targets.
|
||||
"""
|
||||
|
||||
@property
|
||||
def type(self) -> zvec._zvec.typing.IndexType:
|
||||
"""
|
||||
|
|
@ -1015,6 +1051,49 @@ class SegmentOption:
|
|||
bool: Whether the segment is read-only.
|
||||
"""
|
||||
|
||||
class QuantizerParam:
|
||||
"""
|
||||
|
||||
Optional parameters for quantizer configuration.
|
||||
|
||||
This class is only needed when customizing quantizer behavior (e.g., enabling
|
||||
random rotation). It can be omitted for default quantization settings.
|
||||
|
||||
Attributes:
|
||||
enable_rotate (bool): Whether to apply random rotation before INT8/INT4
|
||||
quantization to reduce quantization error.
|
||||
Only effective with quantize_type=INT8 or INT4. Defaults to False.
|
||||
|
||||
Examples:
|
||||
>>> qp = QuantizerParam(enable_rotate=True)
|
||||
>>> print(qp.enable_rotate)
|
||||
True
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
def __init__(self, enable_rotate: bool = False) -> None:
|
||||
"""
|
||||
Constructs a QuantizerParam instance.
|
||||
|
||||
Args:
|
||||
enable_rotate (bool, optional): Whether to apply random rotation
|
||||
before INT8/INT4 quantization. Defaults to False.
|
||||
"""
|
||||
|
||||
def __repr__(self) -> str: ...
|
||||
def __setstate__(self, arg0: tuple) -> None: ...
|
||||
def __eq__(self, arg0: typing.Any) -> bool: ...
|
||||
def to_dict(self) -> dict:
|
||||
"""
|
||||
Convert to dictionary with all fields
|
||||
"""
|
||||
|
||||
@property
|
||||
def enable_rotate(self) -> bool:
|
||||
"""
|
||||
bool: Whether random rotation is enabled before INT8/INT4 quantization.
|
||||
"""
|
||||
|
||||
class VectorIndexParam(IndexParam):
|
||||
"""
|
||||
|
||||
|
|
@ -1026,6 +1105,7 @@ class VectorIndexParam(IndexParam):
|
|||
type (IndexType): The specific vector index type (e.g., HNSW, FLAT).
|
||||
metric_type (MetricType): Distance metric used for similarity search.
|
||||
quantize_type (QuantizeType): Optional vector quantization type.
|
||||
quantizer_param (QuantizerParam): Optional quantizer parameters.
|
||||
"""
|
||||
|
||||
def __getstate__(self) -> tuple: ...
|
||||
|
|
@ -1047,6 +1127,12 @@ class VectorIndexParam(IndexParam):
|
|||
QuantizeType: Vector quantization type (e.g., FP16, INT8).
|
||||
"""
|
||||
|
||||
@property
|
||||
def quantizer_param(self) -> QuantizerParam:
|
||||
"""
|
||||
QuantizerParam: Quantizer configuration including enable_rotate.
|
||||
"""
|
||||
|
||||
class _SearchQuery:
|
||||
field_name: str
|
||||
filter: str
|
||||
|
|
|
|||
|
|
@ -0,0 +1,38 @@
|
|||
// 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 {
|
||||
namespace ailego {
|
||||
|
||||
//! Flip the sign of elements based on a packed bit-array.
|
||||
void fht_flip_sign(const uint8_t *flip, float *data, size_t dim);
|
||||
|
||||
//! Kac random walk: butterfly add/sub between first and second halves.
|
||||
void fht_kacs_walk(float *data, size_t len);
|
||||
|
||||
//! Inverse Kac walk: undo butterfly add/sub with 0.5 factor.
|
||||
void fht_inv_kacs_walk(float *data, size_t len);
|
||||
|
||||
//! In-place Fast Hadamard Transform on a power-of-2 length array.
|
||||
void fht_inplace(float *data, size_t n);
|
||||
|
||||
//! Scale each element by a constant factor.
|
||||
void fht_vec_rescale(float *data, size_t n, float factor);
|
||||
|
||||
} // namespace ailego
|
||||
} // namespace zvec
|
||||
|
|
@ -0,0 +1,125 @@
|
|||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#if defined(__AVX2__)
|
||||
|
||||
#include <immintrin.h>
|
||||
#include <cmath>
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
namespace zvec {
|
||||
namespace ailego {
|
||||
|
||||
void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim) {
|
||||
size_t simd_end = dim & ~31u;
|
||||
constexpr size_t kChunk = 32;
|
||||
const __m256i bit_select =
|
||||
_mm256_setr_epi32(0x01, 0x02, 0x04, 0x08, 0x10, 0x20, 0x40, 0x80);
|
||||
const __m256 sign_flip = _mm256_castsi256_ps(_mm256_set1_epi32(0x80000000));
|
||||
for (size_t i = 0; i < simd_end; i += kChunk) {
|
||||
uint32_t mask_bits;
|
||||
std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits));
|
||||
for (int b = 0; b < 4; ++b) {
|
||||
__m256i mb = _mm256_set1_epi32((mask_bits >> (b * 8)) & 0xFF);
|
||||
__m256i test = _mm256_and_si256(mb, bit_select);
|
||||
__m256i cmp = _mm256_cmpeq_epi32(test, bit_select);
|
||||
__m256 xor_mask = _mm256_and_ps(_mm256_castsi256_ps(cmp), sign_flip);
|
||||
__m256 v = _mm256_loadu_ps(&data[i + b * 8]);
|
||||
v = _mm256_xor_ps(v, xor_mask);
|
||||
_mm256_storeu_ps(&data[i + b * 8], v);
|
||||
}
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = simd_end; i < dim; ++i) {
|
||||
if (flip[i / 8] & (1u << (i % 8))) {
|
||||
data[i] = -data[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void fht_kacs_walk_avx2(float *data, size_t len) {
|
||||
size_t half = len / 2;
|
||||
size_t base = len % 2;
|
||||
size_t offset = base + half;
|
||||
size_t half_end = half & ~7u;
|
||||
for (size_t i = 0; i < half_end; i += 8) {
|
||||
__m256 x = _mm256_loadu_ps(&data[i]);
|
||||
__m256 y = _mm256_loadu_ps(&data[i + offset]);
|
||||
_mm256_storeu_ps(&data[i], _mm256_add_ps(x, y));
|
||||
_mm256_storeu_ps(&data[i + offset], _mm256_sub_ps(x, y));
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = half_end; i < half; ++i) {
|
||||
float x = data[i];
|
||||
float y = data[i + offset];
|
||||
data[i] = x + y;
|
||||
data[i + offset] = x - y;
|
||||
}
|
||||
if (base != 0) {
|
||||
data[half] *= std::sqrt(2.0f);
|
||||
}
|
||||
}
|
||||
|
||||
void fht_inv_kacs_walk_avx2(float *data, size_t len) {
|
||||
size_t half = len / 2;
|
||||
size_t base = len % 2;
|
||||
size_t offset = base + half;
|
||||
if (base != 0) {
|
||||
data[half] *= std::sqrt(0.5f);
|
||||
}
|
||||
size_t half_end = half & ~7u;
|
||||
const __m256 half_fac = _mm256_set1_ps(0.5f);
|
||||
for (size_t i = 0; i < half_end; i += 8) {
|
||||
__m256 a = _mm256_loadu_ps(&data[i]);
|
||||
__m256 b = _mm256_loadu_ps(&data[i + offset]);
|
||||
_mm256_storeu_ps(&data[i], _mm256_mul_ps(_mm256_add_ps(a, b), half_fac));
|
||||
_mm256_storeu_ps(&data[i + offset],
|
||||
_mm256_mul_ps(_mm256_sub_ps(a, b), half_fac));
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = half_end; i < half; ++i) {
|
||||
float a = data[i];
|
||||
float b = data[i + offset];
|
||||
data[i] = (a + b) * 0.5f;
|
||||
data[i + offset] = (a - b) * 0.5f;
|
||||
}
|
||||
}
|
||||
|
||||
void fht_inplace_avx2(float *data, size_t n) {
|
||||
for (size_t len = 1; len < n; len <<= 1) {
|
||||
size_t step = len << 1;
|
||||
size_t simd_end = len & ~7u;
|
||||
for (size_t i = 0; i < n; i += step) {
|
||||
for (size_t j = 0; j < simd_end; j += 8) {
|
||||
__m256 u = _mm256_loadu_ps(&data[i + j]);
|
||||
__m256 v = _mm256_loadu_ps(&data[i + j + len]);
|
||||
_mm256_storeu_ps(&data[i + j], _mm256_add_ps(u, v));
|
||||
_mm256_storeu_ps(&data[i + j + len], _mm256_sub_ps(u, v));
|
||||
}
|
||||
for (size_t j = simd_end; j < len; ++j) {
|
||||
float u = data[i + j];
|
||||
float v = data[i + j + len];
|
||||
data[i + j] = u + v;
|
||||
data[i + j + len] = u - v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace ailego
|
||||
} // namespace zvec
|
||||
|
||||
#endif // __AVX2__
|
||||
|
|
@ -0,0 +1,130 @@
|
|||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#if defined(__AVX512F__)
|
||||
|
||||
#include <immintrin.h>
|
||||
#include <cmath>
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
namespace zvec {
|
||||
namespace ailego {
|
||||
|
||||
void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim) {
|
||||
size_t simd_end = dim & ~63u;
|
||||
constexpr size_t kChunk = 64;
|
||||
const __m512 sign_flip = _mm512_castsi512_ps(_mm512_set1_epi32(0x80000000));
|
||||
for (size_t i = 0; i < simd_end; i += kChunk) {
|
||||
uint64_t mask_bits;
|
||||
std::memcpy(&mask_bits, &flip[i / 8], sizeof(mask_bits));
|
||||
const __mmask16 m0 = _cvtu32_mask16(mask_bits & 0xFFFF);
|
||||
const __mmask16 m1 = _cvtu32_mask16((mask_bits >> 16) & 0xFFFF);
|
||||
const __mmask16 m2 = _cvtu32_mask16((mask_bits >> 32) & 0xFFFF);
|
||||
const __mmask16 m3 = _cvtu32_mask16((mask_bits >> 48) & 0xFFFF);
|
||||
__m512 v0 = _mm512_loadu_ps(&data[i]);
|
||||
v0 = _mm512_mask_xor_ps(v0, m0, v0, sign_flip);
|
||||
_mm512_storeu_ps(&data[i], v0);
|
||||
__m512 v1 = _mm512_loadu_ps(&data[i + 16]);
|
||||
v1 = _mm512_mask_xor_ps(v1, m1, v1, sign_flip);
|
||||
_mm512_storeu_ps(&data[i + 16], v1);
|
||||
__m512 v2 = _mm512_loadu_ps(&data[i + 32]);
|
||||
v2 = _mm512_mask_xor_ps(v2, m2, v2, sign_flip);
|
||||
_mm512_storeu_ps(&data[i + 32], v2);
|
||||
__m512 v3 = _mm512_loadu_ps(&data[i + 48]);
|
||||
v3 = _mm512_mask_xor_ps(v3, m3, v3, sign_flip);
|
||||
_mm512_storeu_ps(&data[i + 48], v3);
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = simd_end; i < dim; ++i) {
|
||||
if (flip[i / 8] & (1u << (i % 8))) {
|
||||
data[i] = -data[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void fht_kacs_walk_avx512(float *data, size_t len) {
|
||||
size_t half = len / 2;
|
||||
size_t base = len % 2;
|
||||
size_t offset = base + half;
|
||||
size_t half_end = half & ~15u;
|
||||
for (size_t i = 0; i < half_end; i += 16) {
|
||||
__m512 x = _mm512_loadu_ps(&data[i]);
|
||||
__m512 y = _mm512_loadu_ps(&data[i + offset]);
|
||||
_mm512_storeu_ps(&data[i], _mm512_add_ps(x, y));
|
||||
_mm512_storeu_ps(&data[i + offset], _mm512_sub_ps(x, y));
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = half_end; i < half; ++i) {
|
||||
float x = data[i];
|
||||
float y = data[i + offset];
|
||||
data[i] = x + y;
|
||||
data[i + offset] = x - y;
|
||||
}
|
||||
if (base != 0) {
|
||||
data[half] *= std::sqrt(2.0f);
|
||||
}
|
||||
}
|
||||
|
||||
void fht_inv_kacs_walk_avx512(float *data, size_t len) {
|
||||
size_t half = len / 2;
|
||||
size_t base = len % 2;
|
||||
size_t offset = base + half;
|
||||
if (base != 0) {
|
||||
data[half] *= std::sqrt(0.5f);
|
||||
}
|
||||
size_t half_end = half & ~15u;
|
||||
const __m512 half_fac = _mm512_set1_ps(0.5f);
|
||||
for (size_t i = 0; i < half_end; i += 16) {
|
||||
__m512 a = _mm512_loadu_ps(&data[i]);
|
||||
__m512 b = _mm512_loadu_ps(&data[i + offset]);
|
||||
_mm512_storeu_ps(&data[i], _mm512_mul_ps(_mm512_add_ps(a, b), half_fac));
|
||||
_mm512_storeu_ps(&data[i + offset],
|
||||
_mm512_mul_ps(_mm512_sub_ps(a, b), half_fac));
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = half_end; i < half; ++i) {
|
||||
float a = data[i];
|
||||
float b = data[i + offset];
|
||||
data[i] = (a + b) * 0.5f;
|
||||
data[i + offset] = (a - b) * 0.5f;
|
||||
}
|
||||
}
|
||||
|
||||
void fht_inplace_avx512(float *data, size_t n) {
|
||||
for (size_t len = 1; len < n; len <<= 1) {
|
||||
size_t step = len << 1;
|
||||
size_t simd_end = len & ~15u;
|
||||
for (size_t i = 0; i < n; i += step) {
|
||||
for (size_t j = 0; j < simd_end; j += 16) {
|
||||
__m512 u = _mm512_loadu_ps(&data[i + j]);
|
||||
__m512 v = _mm512_loadu_ps(&data[i + j + len]);
|
||||
_mm512_storeu_ps(&data[i + j], _mm512_add_ps(u, v));
|
||||
_mm512_storeu_ps(&data[i + j + len], _mm512_sub_ps(u, v));
|
||||
}
|
||||
for (size_t j = simd_end; j < len; ++j) {
|
||||
float u = data[i + j];
|
||||
float v = data[i + j + len];
|
||||
data[i + j] = u + v;
|
||||
data[i + j + len] = u - v;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace ailego
|
||||
} // namespace zvec
|
||||
|
||||
#endif // __AVX512F__
|
||||
|
|
@ -0,0 +1,150 @@
|
|||
// 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 <ailego/internal/cpu_features.h>
|
||||
#include "fht.h"
|
||||
|
||||
namespace zvec {
|
||||
namespace ailego {
|
||||
|
||||
// ISA-specific forward declarations (implementations in
|
||||
// fht_scalar/sse/avx2/avx512/neon)
|
||||
void fht_flip_sign_scalar(const uint8_t *flip, float *data, size_t dim);
|
||||
void fht_kacs_walk_scalar(float *data, size_t len);
|
||||
void fht_inv_kacs_walk_scalar(float *data, size_t len);
|
||||
void fht_inplace_scalar(float *data, size_t n);
|
||||
#if defined(__SSE2__)
|
||||
void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim);
|
||||
void fht_kacs_walk_sse(float *data, size_t len);
|
||||
void fht_inv_kacs_walk_sse(float *data, size_t len);
|
||||
#endif
|
||||
#if defined(__AVX2__)
|
||||
void fht_flip_sign_avx2(const uint8_t *flip, float *data, size_t dim);
|
||||
void fht_kacs_walk_avx2(float *data, size_t len);
|
||||
void fht_inv_kacs_walk_avx2(float *data, size_t len);
|
||||
void fht_inplace_avx2(float *data, size_t n);
|
||||
#endif
|
||||
#if defined(__AVX512F__)
|
||||
void fht_flip_sign_avx512(const uint8_t *flip, float *data, size_t dim);
|
||||
void fht_kacs_walk_avx512(float *data, size_t len);
|
||||
void fht_inv_kacs_walk_avx512(float *data, size_t len);
|
||||
void fht_inplace_avx512(float *data, size_t n);
|
||||
#endif
|
||||
#if defined(__ARM_NEON) && defined(__aarch64__)
|
||||
void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim);
|
||||
void fht_kacs_walk_neon(float *data, size_t len);
|
||||
void fht_inv_kacs_walk_neon(float *data, size_t len);
|
||||
#endif
|
||||
|
||||
// ============================================================================
|
||||
// Runtime dispatch entry points
|
||||
// ============================================================================
|
||||
|
||||
void fht_flip_sign(const uint8_t *flip, float *data, size_t dim) {
|
||||
#if defined(__ARM_NEON) && defined(__aarch64__)
|
||||
fht_flip_sign_neon(flip, data, dim);
|
||||
#else
|
||||
#if defined(__AVX512F__)
|
||||
if (internal::CpuFeatures::static_flags_.AVX512F &&
|
||||
internal::CpuFeatures::static_flags_.AVX512DQ) {
|
||||
fht_flip_sign_avx512(flip, data, dim);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
#if defined(__AVX2__)
|
||||
if (internal::CpuFeatures::static_flags_.AVX2) {
|
||||
fht_flip_sign_avx2(flip, data, dim);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
#if defined(__SSE2__)
|
||||
if (internal::CpuFeatures::static_flags_.SSE2) {
|
||||
fht_flip_sign_sse(flip, data, dim);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
fht_flip_sign_scalar(flip, data, dim);
|
||||
#endif // __ARM_NEON
|
||||
}
|
||||
|
||||
void fht_kacs_walk(float *data, size_t len) {
|
||||
#if defined(__ARM_NEON) && defined(__aarch64__)
|
||||
fht_kacs_walk_neon(data, len);
|
||||
#else
|
||||
#if defined(__AVX512F__)
|
||||
if (internal::CpuFeatures::static_flags_.AVX512F) {
|
||||
fht_kacs_walk_avx512(data, len);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
#if defined(__AVX2__)
|
||||
if (internal::CpuFeatures::static_flags_.AVX2) {
|
||||
fht_kacs_walk_avx2(data, len);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
#if defined(__SSE2__)
|
||||
if (internal::CpuFeatures::static_flags_.SSE2) {
|
||||
fht_kacs_walk_sse(data, len);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
fht_kacs_walk_scalar(data, len);
|
||||
#endif // __ARM_NEON
|
||||
}
|
||||
|
||||
void fht_inv_kacs_walk(float *data, size_t len) {
|
||||
#if defined(__ARM_NEON) && defined(__aarch64__)
|
||||
fht_inv_kacs_walk_neon(data, len);
|
||||
#else
|
||||
#if defined(__AVX512F__)
|
||||
if (internal::CpuFeatures::static_flags_.AVX512F) {
|
||||
fht_inv_kacs_walk_avx512(data, len);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
#if defined(__AVX2__)
|
||||
if (internal::CpuFeatures::static_flags_.AVX2) {
|
||||
fht_inv_kacs_walk_avx2(data, len);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
#if defined(__SSE2__)
|
||||
if (internal::CpuFeatures::static_flags_.SSE2) {
|
||||
fht_inv_kacs_walk_sse(data, len);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
fht_inv_kacs_walk_scalar(data, len);
|
||||
#endif // __ARM_NEON
|
||||
}
|
||||
|
||||
void fht_inplace(float *data, size_t n) {
|
||||
#if defined(__AVX512F__)
|
||||
if (internal::CpuFeatures::static_flags_.AVX512F) {
|
||||
fht_inplace_avx512(data, n);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
#if defined(__AVX2__)
|
||||
if (internal::CpuFeatures::static_flags_.AVX2) {
|
||||
fht_inplace_avx2(data, n);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
fht_inplace_scalar(data, n);
|
||||
}
|
||||
|
||||
} // namespace ailego
|
||||
} // namespace zvec
|
||||
|
|
@ -0,0 +1,107 @@
|
|||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#if defined(__ARM_NEON) && defined(__aarch64__)
|
||||
|
||||
#include <cmath>
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <arm_neon.h>
|
||||
|
||||
namespace zvec {
|
||||
namespace ailego {
|
||||
|
||||
void fht_flip_sign_neon(const uint8_t *flip, float *data, size_t dim) {
|
||||
const uint32x4_t sign_bit = vdupq_n_u32(0x80000000u);
|
||||
size_t simd_end = dim & ~3u;
|
||||
size_t flip_bytes = (dim + 7) / 8;
|
||||
for (size_t i = 0; i < simd_end; i += 4) {
|
||||
uint16_t bits16;
|
||||
size_t byte_pos = i / 8;
|
||||
if (byte_pos + 1 < flip_bytes) {
|
||||
std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16));
|
||||
} else {
|
||||
bits16 = flip[byte_pos];
|
||||
}
|
||||
bits16 >>= (i % 8);
|
||||
uint32_t b0 = bits16 & 1u;
|
||||
uint32_t b1 = (bits16 >> 1) & 1u;
|
||||
uint32_t b2 = (bits16 >> 2) & 1u;
|
||||
uint32_t b3 = (bits16 >> 3) & 1u;
|
||||
uint32x4_t bit_mask = {b0, b1, b2, b3};
|
||||
uint32x4_t sign_mask = vmulq_u32(bit_mask, sign_bit);
|
||||
float32x4_t v = vld1q_f32(&data[i]);
|
||||
v = vreinterpretq_f32_u32(veorq_u32(vreinterpretq_u32_f32(v), sign_mask));
|
||||
vst1q_f32(&data[i], v);
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = simd_end; i < dim; ++i) {
|
||||
if (flip[i / 8] & (1u << (i % 8))) {
|
||||
data[i] = -data[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void fht_kacs_walk_neon(float *data, size_t len) {
|
||||
size_t half = len / 2;
|
||||
size_t base = len % 2;
|
||||
size_t offset = base + half;
|
||||
size_t half_end = half & ~3u;
|
||||
for (size_t i = 0; i < half_end; i += 4) {
|
||||
float32x4_t x = vld1q_f32(&data[i]);
|
||||
float32x4_t y = vld1q_f32(&data[i + offset]);
|
||||
vst1q_f32(&data[i], vaddq_f32(x, y));
|
||||
vst1q_f32(&data[i + offset], vsubq_f32(x, y));
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = half_end; i < half; ++i) {
|
||||
float x = data[i];
|
||||
float y = data[i + offset];
|
||||
data[i] = x + y;
|
||||
data[i + offset] = x - y;
|
||||
}
|
||||
if (base != 0) {
|
||||
data[half] *= std::sqrt(2.0f);
|
||||
}
|
||||
}
|
||||
|
||||
void fht_inv_kacs_walk_neon(float *data, size_t len) {
|
||||
size_t half = len / 2;
|
||||
size_t base = len % 2;
|
||||
size_t offset = base + half;
|
||||
if (base != 0) {
|
||||
data[half] *= std::sqrt(0.5f);
|
||||
}
|
||||
size_t half_end = half & ~3u;
|
||||
const float32x4_t half_fac = vdupq_n_f32(0.5f);
|
||||
for (size_t i = 0; i < half_end; i += 4) {
|
||||
float32x4_t a = vld1q_f32(&data[i]);
|
||||
float32x4_t b = vld1q_f32(&data[i + offset]);
|
||||
vst1q_f32(&data[i], vmulq_f32(vaddq_f32(a, b), half_fac));
|
||||
vst1q_f32(&data[i + offset], vmulq_f32(vsubq_f32(a, b), half_fac));
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = half_end; i < half; ++i) {
|
||||
float a = data[i];
|
||||
float b = data[i + offset];
|
||||
data[i] = (a + b) * 0.5f;
|
||||
data[i + offset] = (a - b) * 0.5f;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace ailego
|
||||
} // namespace zvec
|
||||
|
||||
#endif // __ARM_NEON && __aarch64__
|
||||
|
|
@ -0,0 +1,80 @@
|
|||
// 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 <cstddef>
|
||||
#include <cstdint>
|
||||
|
||||
namespace zvec {
|
||||
namespace ailego {
|
||||
|
||||
void fht_flip_sign_scalar(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_scalar(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_scalar(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_scalar(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;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace ailego
|
||||
} // namespace zvec
|
||||
|
|
@ -0,0 +1,106 @@
|
|||
// Copyright 2025-present the zvec project
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#if defined(__SSE2__)
|
||||
|
||||
#include <emmintrin.h>
|
||||
#include <cmath>
|
||||
#include <cstddef>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
|
||||
namespace zvec {
|
||||
namespace ailego {
|
||||
|
||||
void fht_flip_sign_sse(const uint8_t *flip, float *data, size_t dim) {
|
||||
size_t simd_end = dim & ~3u;
|
||||
size_t flip_bytes = (dim + 7) / 8;
|
||||
for (size_t i = 0; i < simd_end; i += 4) {
|
||||
uint16_t bits16;
|
||||
size_t byte_pos = i / 8;
|
||||
if (byte_pos + 1 < flip_bytes) {
|
||||
std::memcpy(&bits16, &flip[byte_pos], sizeof(bits16));
|
||||
} else {
|
||||
bits16 = flip[byte_pos];
|
||||
}
|
||||
bits16 >>= (i % 8);
|
||||
uint32_t b0 = bits16 & 1u;
|
||||
uint32_t b1 = (bits16 >> 1) & 1u;
|
||||
uint32_t b2 = (bits16 >> 2) & 1u;
|
||||
uint32_t b3 = (bits16 >> 3) & 1u;
|
||||
__m128i bit_mask = _mm_set_epi32(b3, b2, b1, b0);
|
||||
__m128i sign_mask = _mm_slli_epi32(bit_mask, 31);
|
||||
__m128 v = _mm_loadu_ps(&data[i]);
|
||||
v = _mm_xor_ps(v, _mm_castsi128_ps(sign_mask));
|
||||
_mm_storeu_ps(&data[i], v);
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = simd_end; i < dim; ++i) {
|
||||
if (flip[i / 8] & (1u << (i % 8))) {
|
||||
data[i] = -data[i];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void fht_kacs_walk_sse(float *data, size_t len) {
|
||||
size_t half = len / 2;
|
||||
size_t base = len % 2;
|
||||
size_t offset = base + half;
|
||||
size_t half_end = half & ~3u;
|
||||
for (size_t i = 0; i < half_end; i += 4) {
|
||||
__m128 x = _mm_loadu_ps(&data[i]);
|
||||
__m128 y = _mm_loadu_ps(&data[i + offset]);
|
||||
_mm_storeu_ps(&data[i], _mm_add_ps(x, y));
|
||||
_mm_storeu_ps(&data[i + offset], _mm_sub_ps(x, y));
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = half_end; i < half; ++i) {
|
||||
float x = data[i];
|
||||
float y = data[i + offset];
|
||||
data[i] = x + y;
|
||||
data[i + offset] = x - y;
|
||||
}
|
||||
if (base != 0) {
|
||||
data[half] *= std::sqrt(2.0f);
|
||||
}
|
||||
}
|
||||
|
||||
void fht_inv_kacs_walk_sse(float *data, size_t len) {
|
||||
size_t half = len / 2;
|
||||
size_t base = len % 2;
|
||||
size_t offset = base + half;
|
||||
if (base != 0) {
|
||||
data[half] *= std::sqrt(0.5f);
|
||||
}
|
||||
size_t half_end = half & ~3u;
|
||||
const __m128 half_fac = _mm_set1_ps(0.5f);
|
||||
for (size_t i = 0; i < half_end; i += 4) {
|
||||
__m128 a = _mm_loadu_ps(&data[i]);
|
||||
__m128 b = _mm_loadu_ps(&data[i + offset]);
|
||||
_mm_storeu_ps(&data[i], _mm_mul_ps(_mm_add_ps(a, b), half_fac));
|
||||
_mm_storeu_ps(&data[i + offset], _mm_mul_ps(_mm_sub_ps(a, b), half_fac));
|
||||
}
|
||||
// Scalar tail
|
||||
for (size_t i = half_end; i < half; ++i) {
|
||||
float a = data[i];
|
||||
float b = data[i + offset];
|
||||
data[i] = (a + b) * 0.5f;
|
||||
data[i + offset] = (a - b) * 0.5f;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace ailego
|
||||
} // namespace zvec
|
||||
|
||||
#endif // __SSE2__
|
||||
|
|
@ -1476,6 +1476,60 @@ zvec_quantize_type_t zvec_index_params_get_quantize_type(
|
|||
return ZVEC_QUANTIZE_TYPE_UNDEFINED;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Set enable_rotate for quantizer parameters
|
||||
* @param params Index parameters (must be vector index type)
|
||||
* @param enable_rotate Whether to enable random rotation before quantization
|
||||
* @return ZVEC_OK on success, error code on failure
|
||||
*/
|
||||
zvec_error_code_t zvec_index_params_set_quantizer_enable_rotate(
|
||||
zvec_index_params_t *params, bool enable_rotate) {
|
||||
if (!params) {
|
||||
SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT,
|
||||
"Index params pointer cannot be null");
|
||||
return ZVEC_ERROR_INVALID_ARGUMENT;
|
||||
}
|
||||
auto *cpp_params = reinterpret_cast<zvec::IndexParams *>(params);
|
||||
|
||||
if (!cpp_params->is_vector_index_type()) {
|
||||
SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT,
|
||||
"Index params is not a vector index type");
|
||||
return ZVEC_ERROR_INVALID_ARGUMENT;
|
||||
}
|
||||
auto *vec_params = dynamic_cast<zvec::VectorIndexParams *>(cpp_params);
|
||||
if (!vec_params) {
|
||||
SET_LAST_ERROR(ZVEC_ERROR_INVALID_ARGUMENT,
|
||||
"Failed to cast to VectorIndexParams");
|
||||
return ZVEC_ERROR_INVALID_ARGUMENT;
|
||||
}
|
||||
zvec::QuantizerParam qp = vec_params->quantizer_param();
|
||||
qp.set_enable_rotate(enable_rotate);
|
||||
vec_params->set_quantizer_param(qp);
|
||||
return ZVEC_OK;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get enable_rotate setting from quantizer parameters
|
||||
* @param params Index parameters
|
||||
* @return true if rotation is enabled, false otherwise
|
||||
*/
|
||||
bool zvec_index_params_get_quantizer_enable_rotate(
|
||||
const zvec_index_params_t *params) {
|
||||
if (!params) {
|
||||
return false;
|
||||
}
|
||||
auto *cpp_params = reinterpret_cast<const zvec::IndexParams *>(params);
|
||||
|
||||
if (cpp_params->is_vector_index_type()) {
|
||||
auto *vec_params =
|
||||
dynamic_cast<const zvec::VectorIndexParams *>(cpp_params);
|
||||
if (vec_params) {
|
||||
return vec_params->quantizer_param().enable_rotate();
|
||||
}
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief Get index type from index parameters
|
||||
* @param params Index parameters
|
||||
|
|
|
|||
|
|
@ -338,6 +338,62 @@ Args:
|
|||
t[2].cast<std::string>());
|
||||
}));
|
||||
|
||||
// binding QuantizerParam
|
||||
py::class_<QuantizerParam, std::shared_ptr<QuantizerParam>> quantizer_param(
|
||||
m, "QuantizerParam", R"pbdoc(
|
||||
Parameters for quantizer configuration.
|
||||
|
||||
Encapsulates quantization-related settings such as enable_rotate.
|
||||
Designed for future extensibility.
|
||||
|
||||
Attributes:
|
||||
enable_rotate (bool): Whether to apply random rotation before INT8/INT4
|
||||
quantization to reduce quantization error.
|
||||
Only effective with quantize_type=INT8 or INT4. Defaults to False.
|
||||
|
||||
Examples:
|
||||
>>> qp = QuantizerParam(enable_rotate=True)
|
||||
>>> print(qp.enable_rotate)
|
||||
True
|
||||
)pbdoc");
|
||||
quantizer_param.def(py::init<bool>(), py::arg("enable_rotate") = false)
|
||||
.def_property_readonly(
|
||||
"enable_rotate",
|
||||
[](const QuantizerParam &self) -> bool {
|
||||
return self.enable_rotate();
|
||||
},
|
||||
"bool: Whether random rotation is enabled before INT8/INT4 "
|
||||
"quantization.")
|
||||
.def(
|
||||
"to_dict",
|
||||
[](const QuantizerParam &self) -> py::dict {
|
||||
py::dict dict;
|
||||
dict["enable_rotate"] = self.enable_rotate();
|
||||
return dict;
|
||||
},
|
||||
"Convert to dictionary with all fields")
|
||||
.def("__repr__",
|
||||
[](const QuantizerParam &self) -> std::string {
|
||||
return "{\"enable_rotate\":" +
|
||||
std::string(self.enable_rotate() ? "true" : "false") + "}";
|
||||
})
|
||||
.def(
|
||||
"__eq__",
|
||||
[](const QuantizerParam &self, const py::object &other) {
|
||||
if (!py::isinstance<QuantizerParam>(other)) return false;
|
||||
return self == other.cast<const QuantizerParam &>();
|
||||
},
|
||||
py::is_operator())
|
||||
.def(py::pickle(
|
||||
[](const QuantizerParam &self) {
|
||||
return py::make_tuple(self.enable_rotate());
|
||||
},
|
||||
[](py::tuple t) {
|
||||
if (t.size() != 1)
|
||||
throw std::runtime_error("Invalid state for QuantizerParam");
|
||||
return std::make_shared<QuantizerParam>(t[0].cast<bool>());
|
||||
}));
|
||||
|
||||
// binding base vector index params
|
||||
py::class_<VectorIndexParams, IndexParams, std::shared_ptr<VectorIndexParams>>
|
||||
vector_params(m, "VectorIndexParam", R"pbdoc(
|
||||
|
|
@ -349,6 +405,7 @@ Attributes:
|
|||
type (IndexType): The specific vector index type (e.g., HNSW, FLAT).
|
||||
metric_type (MetricType): Distance metric used for similarity search.
|
||||
quantize_type (QuantizeType): Optional vector quantization type.
|
||||
quantizer_param (QuantizerParam): Quantizer configuration (e.g., enable_rotate).
|
||||
)pbdoc");
|
||||
vector_params
|
||||
.def_property_readonly(
|
||||
|
|
@ -363,6 +420,12 @@ Attributes:
|
|||
return self.quantize_type();
|
||||
},
|
||||
"QuantizeType: Vector quantization type (e.g., FP16, INT8).")
|
||||
.def_property_readonly(
|
||||
"quantizer_param",
|
||||
[](const VectorIndexParams &self) -> QuantizerParam {
|
||||
return self.quantizer_param();
|
||||
},
|
||||
"QuantizerParam: Quantizer configuration including enable_rotate.")
|
||||
.def(
|
||||
"to_dict",
|
||||
[](const VectorIndexParams &self) -> py::dict {
|
||||
|
|
@ -371,6 +434,9 @@ Attributes:
|
|||
dict["metric_type"] = metric_type_to_string(self.metric_type());
|
||||
dict["quantize_type"] =
|
||||
quantize_type_to_string(self.quantize_type());
|
||||
py::dict qp_dict;
|
||||
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
|
||||
dict["quantizer_param"] = qp_dict;
|
||||
return dict;
|
||||
},
|
||||
"Convert to dictionary with all fields")
|
||||
|
|
@ -382,7 +448,7 @@ Attributes:
|
|||
[](py::tuple t) { // __setstate__
|
||||
if (t.size() != 3)
|
||||
throw std::runtime_error("Invalid state for VectorIndexParams");
|
||||
// 基类,不能直接实例化,用于子类
|
||||
// Base class, cannot instantiate directly, used by subclasses
|
||||
return std::shared_ptr<VectorIndexParams>();
|
||||
}));
|
||||
|
||||
|
|
@ -421,13 +487,20 @@ Examples:
|
|||
{'metric_type': 'IP', 'm': 16, 'ef_construction': 200, 'quantize_type': 'INT8', 'use_contiguous_memory': True}
|
||||
)pbdoc");
|
||||
hnsw_params
|
||||
.def(py::init<MetricType, int, int, QuantizeType, bool>(),
|
||||
.def(py::init([](MetricType metric_type, int m, int ef_construction,
|
||||
QuantizeType quantize_type, bool use_contiguous_memory,
|
||||
QuantizerParam quantizer_param) {
|
||||
return std::make_shared<HnswIndexParams>(
|
||||
metric_type, m, ef_construction, quantize_type,
|
||||
use_contiguous_memory, quantizer_param);
|
||||
}),
|
||||
py::arg("metric_type") = MetricType::IP,
|
||||
py::arg("m") = core_interface::kDefaultHnswNeighborCnt,
|
||||
py::arg("ef_construction") =
|
||||
core_interface::kDefaultHnswEfConstruction,
|
||||
py::arg("quantize_type") = QuantizeType::UNDEFINED,
|
||||
py::arg("use_contiguous_memory") = false)
|
||||
py::arg("use_contiguous_memory") = false,
|
||||
py::arg("quantizer_param") = QuantizerParam())
|
||||
.def_property_readonly(
|
||||
"m", &HnswIndexParams::m,
|
||||
"int: Maximum number of neighbors per node in upper layers.")
|
||||
|
|
@ -450,34 +523,43 @@ Examples:
|
|||
dict["quantize_type"] =
|
||||
quantize_type_to_string(self.quantize_type());
|
||||
dict["use_contiguous_memory"] = self.use_contiguous_memory();
|
||||
py::dict qp_dict;
|
||||
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
|
||||
dict["quantizer_param"] = qp_dict;
|
||||
return dict;
|
||||
},
|
||||
"Convert to dictionary with all fields")
|
||||
.def("__repr__",
|
||||
[](const HnswIndexParams &self) -> std::string {
|
||||
return "{"
|
||||
"\"metric_type\":" +
|
||||
metric_type_to_string(self.metric_type()) +
|
||||
", \"m\":" + std::to_string(self.m()) +
|
||||
", \"ef_construction\":" +
|
||||
std::to_string(self.ef_construction()) +
|
||||
", \"quantize_type\":" +
|
||||
quantize_type_to_string(self.quantize_type()) +
|
||||
", \"use_contiguous_memory\":" +
|
||||
(self.use_contiguous_memory() ? "true" : "false") + "}";
|
||||
})
|
||||
.def(
|
||||
"__repr__",
|
||||
[](const HnswIndexParams &self) -> std::string {
|
||||
return "{"
|
||||
"\"metric_type\":" +
|
||||
metric_type_to_string(self.metric_type()) +
|
||||
", \"m\":" + std::to_string(self.m()) +
|
||||
", \"ef_construction\":" +
|
||||
std::to_string(self.ef_construction()) +
|
||||
", \"quantize_type\":" +
|
||||
quantize_type_to_string(self.quantize_type()) +
|
||||
", \"use_contiguous_memory\":" +
|
||||
(self.use_contiguous_memory() ? "true" : "false") +
|
||||
", \"quantizer_param\":{" + "\"enable_rotate\":" +
|
||||
(self.quantizer_param().enable_rotate() ? "true" : "false") +
|
||||
"}}";
|
||||
})
|
||||
.def(py::pickle(
|
||||
[](const HnswIndexParams &self) {
|
||||
return py::make_tuple(self.metric_type(), self.m(),
|
||||
self.ef_construction(), self.quantize_type(),
|
||||
self.use_contiguous_memory());
|
||||
self.use_contiguous_memory(),
|
||||
self.quantizer_param().enable_rotate());
|
||||
},
|
||||
[](py::tuple t) {
|
||||
if (t.size() != 5)
|
||||
if (t.size() != 5 && t.size() != 6)
|
||||
throw std::runtime_error("Invalid state for HnswIndexParams");
|
||||
QuantizerParam qp(t.size() >= 6 ? t[5].cast<bool>() : false);
|
||||
return std::make_shared<HnswIndexParams>(
|
||||
t[0].cast<MetricType>(), t[1].cast<int>(), t[2].cast<int>(),
|
||||
t[3].cast<QuantizeType>(), t[4].cast<bool>());
|
||||
t[3].cast<QuantizeType>(), t[4].cast<bool>(), qp);
|
||||
}));
|
||||
|
||||
// binding hnsw rabitq index params
|
||||
|
|
@ -626,8 +708,16 @@ Examples:
|
|||
... )
|
||||
)pbdoc");
|
||||
vamana_params
|
||||
.def(py::init<MetricType, int, int, float, bool, bool, bool,
|
||||
QuantizeType>(),
|
||||
.def(py::init([](MetricType metric_type, int max_degree,
|
||||
int search_list_size, float alpha, bool saturate_graph,
|
||||
bool use_contiguous_memory, bool use_id_map,
|
||||
QuantizeType quantize_type,
|
||||
QuantizerParam quantizer_param) {
|
||||
return std::make_shared<VamanaIndexParams>(
|
||||
metric_type, max_degree, search_list_size, alpha,
|
||||
saturate_graph, use_contiguous_memory, use_id_map,
|
||||
quantize_type, quantizer_param);
|
||||
}),
|
||||
py::arg("metric_type") = MetricType::IP,
|
||||
py::arg("max_degree") = core_interface::kDefaultVamanaMaxDegree,
|
||||
py::arg("search_list_size") =
|
||||
|
|
@ -637,7 +727,8 @@ Examples:
|
|||
core_interface::kDefaultVamanaSaturateGraph,
|
||||
py::arg("use_contiguous_memory") = false,
|
||||
py::arg("use_id_map") = false,
|
||||
py::arg("quantize_type") = QuantizeType::UNDEFINED)
|
||||
py::arg("quantize_type") = QuantizeType::UNDEFINED,
|
||||
py::arg("quantizer_param") = QuantizerParam())
|
||||
.def_property_readonly(
|
||||
"max_degree", &VamanaIndexParams::max_degree,
|
||||
"int: Maximum out-degree (R) of every node in the Vamana graph.")
|
||||
|
|
@ -673,45 +764,53 @@ Examples:
|
|||
dict["use_id_map"] = self.use_id_map();
|
||||
dict["quantize_type"] =
|
||||
quantize_type_to_string(self.quantize_type());
|
||||
py::dict qp_dict;
|
||||
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
|
||||
dict["quantizer_param"] = qp_dict;
|
||||
return dict;
|
||||
},
|
||||
"Convert to dictionary with all fields")
|
||||
.def("__repr__",
|
||||
[](const VamanaIndexParams &self) -> std::string {
|
||||
return "{"
|
||||
"\"type\":\"" +
|
||||
index_type_to_string(self.type()) +
|
||||
"\", \"metric_type\":\"" +
|
||||
metric_type_to_string(self.metric_type()) +
|
||||
"\", \"max_degree\":" + std::to_string(self.max_degree()) +
|
||||
", \"search_list_size\":" +
|
||||
std::to_string(self.search_list_size()) +
|
||||
", \"alpha\":" + std::to_string(self.alpha()) +
|
||||
", \"saturate_graph\":" +
|
||||
std::string(self.saturate_graph() ? "true" : "false") +
|
||||
", \"use_contiguous_memory\":" +
|
||||
std::string(self.use_contiguous_memory() ? "true"
|
||||
: "false") +
|
||||
", \"use_id_map\":" +
|
||||
std::string(self.use_id_map() ? "true" : "false") +
|
||||
", \"quantize_type\":\"" +
|
||||
quantize_type_to_string(self.quantize_type()) + "\"}";
|
||||
})
|
||||
.def(
|
||||
"__repr__",
|
||||
[](const VamanaIndexParams &self) -> std::string {
|
||||
return "{"
|
||||
"\"type\":\"" +
|
||||
index_type_to_string(self.type()) +
|
||||
"\", \"metric_type\":\"" +
|
||||
metric_type_to_string(self.metric_type()) +
|
||||
"\", \"max_degree\":" + std::to_string(self.max_degree()) +
|
||||
", \"search_list_size\":" +
|
||||
std::to_string(self.search_list_size()) +
|
||||
", \"alpha\":" + std::to_string(self.alpha()) +
|
||||
", \"saturate_graph\":" +
|
||||
std::string(self.saturate_graph() ? "true" : "false") +
|
||||
", \"use_contiguous_memory\":" +
|
||||
std::string(self.use_contiguous_memory() ? "true"
|
||||
: "false") +
|
||||
", \"use_id_map\":" +
|
||||
std::string(self.use_id_map() ? "true" : "false") +
|
||||
", \"quantize_type\":\"" +
|
||||
quantize_type_to_string(self.quantize_type()) +
|
||||
"\", \"quantizer_param\":{" + "\"enable_rotate\":" +
|
||||
(self.quantizer_param().enable_rotate() ? "true" : "false") +
|
||||
"}}";
|
||||
})
|
||||
.def(py::pickle(
|
||||
[](const VamanaIndexParams &self) {
|
||||
return py::make_tuple(self.metric_type(), self.max_degree(),
|
||||
self.search_list_size(), self.alpha(),
|
||||
self.saturate_graph(),
|
||||
self.use_contiguous_memory(),
|
||||
self.use_id_map(), self.quantize_type());
|
||||
return py::make_tuple(
|
||||
self.metric_type(), self.max_degree(), self.search_list_size(),
|
||||
self.alpha(), self.saturate_graph(),
|
||||
self.use_contiguous_memory(), self.use_id_map(),
|
||||
self.quantize_type(), self.quantizer_param().enable_rotate());
|
||||
},
|
||||
[](py::tuple t) {
|
||||
if (t.size() != 8)
|
||||
if (t.size() != 8 && t.size() != 9)
|
||||
throw std::runtime_error("Invalid state for VamanaIndexParams");
|
||||
QuantizerParam qp(t.size() >= 9 ? t[8].cast<bool>() : false);
|
||||
return std::make_shared<VamanaIndexParams>(
|
||||
t[0].cast<MetricType>(), t[1].cast<int>(), t[2].cast<int>(),
|
||||
t[3].cast<float>(), t[4].cast<bool>(), t[5].cast<bool>(),
|
||||
t[6].cast<bool>(), t[7].cast<QuantizeType>());
|
||||
t[6].cast<bool>(), t[7].cast<QuantizeType>(), qp);
|
||||
}));
|
||||
|
||||
// FlatIndexParams
|
||||
|
|
@ -741,9 +840,14 @@ Examples:
|
|||
{'metric_type': 'L2', 'quantize_type': 'FP16'}
|
||||
)pbdoc");
|
||||
flat_params
|
||||
.def(py::init<MetricType, QuantizeType>(),
|
||||
.def(py::init([](MetricType metric_type, QuantizeType quantize_type,
|
||||
QuantizerParam quantizer_param) {
|
||||
return std::make_shared<FlatIndexParams>(
|
||||
metric_type, quantize_type, quantizer_param);
|
||||
}),
|
||||
py::arg("metric_type") = MetricType::IP,
|
||||
py::arg("quantize_type") = QuantizeType::UNDEFINED,
|
||||
py::arg("quantizer_param") = QuantizerParam(),
|
||||
R"pbdoc(
|
||||
Constructs a FlatIndexParam instance.
|
||||
|
||||
|
|
@ -751,6 +855,8 @@ Args:
|
|||
metric_type (MetricType, optional): Distance metric. Defaults to MetricType.IP.
|
||||
quantize_type (QuantizeType, optional): Vector quantization type.
|
||||
Defaults to QuantizeType.UNDEFINED (no quantization).
|
||||
quantizer_param (QuantizerParam, optional): Quantizer configuration.
|
||||
Defaults to QuantizerParam().
|
||||
)pbdoc")
|
||||
.def(
|
||||
"to_dict",
|
||||
|
|
@ -759,26 +865,35 @@ Args:
|
|||
dict["metric_type"] = metric_type_to_string(self.metric_type());
|
||||
dict["quantize_type"] =
|
||||
quantize_type_to_string(self.quantize_type());
|
||||
py::dict qp_dict;
|
||||
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
|
||||
dict["quantizer_param"] = qp_dict;
|
||||
return dict;
|
||||
},
|
||||
"Convert to dictionary with all fields")
|
||||
.def("__repr__",
|
||||
[](const FlatIndexParams &self) -> std::string {
|
||||
return "{"
|
||||
"\"metric_type\":" +
|
||||
metric_type_to_string(self.metric_type()) +
|
||||
", \"quantize_type\":" +
|
||||
quantize_type_to_string(self.quantize_type()) + "}";
|
||||
})
|
||||
.def(
|
||||
"__repr__",
|
||||
[](const FlatIndexParams &self) -> std::string {
|
||||
return "{"
|
||||
"\"metric_type\":" +
|
||||
metric_type_to_string(self.metric_type()) +
|
||||
", \"quantize_type\":" +
|
||||
quantize_type_to_string(self.quantize_type()) +
|
||||
", \"quantizer_param\":{" + "\"enable_rotate\":" +
|
||||
(self.quantizer_param().enable_rotate() ? "true" : "false") +
|
||||
"}}";
|
||||
})
|
||||
.def(py::pickle(
|
||||
[](const FlatIndexParams &self) {
|
||||
return py::make_tuple(self.metric_type(), self.quantize_type());
|
||||
return py::make_tuple(self.metric_type(), self.quantize_type(),
|
||||
self.quantizer_param().enable_rotate());
|
||||
},
|
||||
[](py::tuple t) {
|
||||
if (t.size() != 2)
|
||||
if (t.size() != 2 && t.size() != 3)
|
||||
throw std::runtime_error("Invalid state for FlatIndexParams");
|
||||
return std::make_shared<FlatIndexParams>(t[0].cast<MetricType>(),
|
||||
t[1].cast<QuantizeType>());
|
||||
QuantizerParam qp(t.size() >= 3 ? t[2].cast<bool>() : false);
|
||||
return std::make_shared<FlatIndexParams>(
|
||||
t[0].cast<MetricType>(), t[1].cast<QuantizeType>(), qp);
|
||||
}));
|
||||
|
||||
// IVFIndexParams
|
||||
|
|
@ -815,10 +930,17 @@ Examples:
|
|||
100
|
||||
)pbdoc");
|
||||
ivf_params
|
||||
.def(py::init<MetricType, int, int, bool, QuantizeType>(),
|
||||
.def(py::init([](MetricType metric_type, int n_list, int n_iters,
|
||||
bool use_soar, QuantizeType quantize_type,
|
||||
QuantizerParam quantizer_param) {
|
||||
return std::make_shared<IVFIndexParams>(
|
||||
metric_type, n_list, n_iters, use_soar, quantize_type,
|
||||
quantizer_param);
|
||||
}),
|
||||
py::arg("metric_type") = MetricType::IP, py::arg("n_list") = 10,
|
||||
py::arg("n_iters") = 10, py::arg("use_soar") = false,
|
||||
py::arg("quantize_type") = QuantizeType::UNDEFINED,
|
||||
py::arg("quantizer_param") = QuantizerParam(),
|
||||
R"pbdoc(
|
||||
Constructs an IVFIndexParam instance.
|
||||
|
||||
|
|
@ -831,6 +953,8 @@ Args:
|
|||
use_soar (bool, optional): Enable SOAR optimization. Defaults to False.
|
||||
quantize_type (QuantizeType, optional): Vector quantization type.
|
||||
Defaults to QuantizeType.UNDEFINED.
|
||||
quantizer_param (QuantizerParam, optional): Quantizer configuration.
|
||||
Defaults to QuantizerParam().
|
||||
)pbdoc")
|
||||
.def_property_readonly("n_list", &IVFIndexParams::n_list,
|
||||
"int: Number of inverted lists.")
|
||||
|
|
@ -850,32 +974,41 @@ Args:
|
|||
dict["use_soar"] = self.use_soar();
|
||||
dict["quantize_type"] =
|
||||
quantize_type_to_string(self.quantize_type());
|
||||
py::dict qp_dict;
|
||||
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
|
||||
dict["quantizer_param"] = qp_dict;
|
||||
return dict;
|
||||
},
|
||||
"Convert to dictionary with all fields")
|
||||
.def("__repr__",
|
||||
[](const IVFIndexParams &self) {
|
||||
return "{"
|
||||
"\"metric_type\":" +
|
||||
metric_type_to_string(self.metric_type()) +
|
||||
", \"n_list\":" + std::to_string(self.n_list()) +
|
||||
", \"n_iters\":" + std::to_string(self.n_iters()) +
|
||||
", \"use_soar\":" + std::to_string(self.use_soar()) +
|
||||
", \"quantize_type\":" +
|
||||
quantize_type_to_string(self.quantize_type()) + "}";
|
||||
})
|
||||
.def(
|
||||
"__repr__",
|
||||
[](const IVFIndexParams &self) {
|
||||
return "{"
|
||||
"\"metric_type\":" +
|
||||
metric_type_to_string(self.metric_type()) +
|
||||
", \"n_list\":" + std::to_string(self.n_list()) +
|
||||
", \"n_iters\":" + std::to_string(self.n_iters()) +
|
||||
", \"use_soar\":" + std::to_string(self.use_soar()) +
|
||||
", \"quantize_type\":" +
|
||||
quantize_type_to_string(self.quantize_type()) +
|
||||
", \"quantizer_param\":{" + "\"enable_rotate\":" +
|
||||
(self.quantizer_param().enable_rotate() ? "true" : "false") +
|
||||
"}}";
|
||||
})
|
||||
.def(py::pickle(
|
||||
[](const IVFIndexParams &self) {
|
||||
return py::make_tuple(self.metric_type(), self.n_list(),
|
||||
self.n_iters(), self.use_soar(),
|
||||
self.quantize_type());
|
||||
self.quantize_type(),
|
||||
self.quantizer_param().enable_rotate());
|
||||
},
|
||||
[](py::tuple t) {
|
||||
if (t.size() != 5)
|
||||
if (t.size() != 5 && t.size() != 6)
|
||||
throw std::runtime_error("Invalid state for IVFIndexParams");
|
||||
QuantizerParam qp(t.size() >= 6 ? t[5].cast<bool>() : false);
|
||||
return std::make_shared<IVFIndexParams>(
|
||||
t[0].cast<MetricType>(), t[1].cast<int>(), t[2].cast<int>(),
|
||||
t[3].cast<bool>(), t[4].cast<QuantizeType>());
|
||||
t[3].cast<bool>(), t[4].cast<QuantizeType>(), qp);
|
||||
}));
|
||||
|
||||
// DiskAnnIndexParams
|
||||
|
|
@ -915,10 +1048,17 @@ Examples:
|
|||
100
|
||||
)pbdoc");
|
||||
diskann_params
|
||||
.def(py::init<MetricType, int, int, int, QuantizeType>(),
|
||||
.def(py::init([](MetricType metric_type, int max_degree, int list_size,
|
||||
int pq_chunk_num, QuantizeType quantize_type,
|
||||
QuantizerParam quantizer_param) {
|
||||
return std::make_shared<DiskAnnIndexParams>(
|
||||
metric_type, max_degree, list_size, pq_chunk_num,
|
||||
quantize_type, quantizer_param);
|
||||
}),
|
||||
py::arg("metric_type") = MetricType::IP, py::arg("max_degree") = 100,
|
||||
py::arg("list_size") = 50, py::arg("pq_chunk_num") = 0,
|
||||
py::arg("quantize_type") = QuantizeType::UNDEFINED,
|
||||
py::arg("quantizer_param") = QuantizerParam(),
|
||||
R"pbdoc(
|
||||
Constructs an DiskAnnIndexParams instance.
|
||||
|
||||
|
|
@ -933,6 +1073,8 @@ Args:
|
|||
Clamped to [1, 1024]. Defaults to 0.
|
||||
quantize_type (QuantizeType, optional): Vector quantization type.
|
||||
Defaults to QuantizeType.UNDEFINED.
|
||||
quantizer_param (QuantizerParam, optional): Quantizer configuration.
|
||||
Defaults to QuantizerParam().
|
||||
)pbdoc")
|
||||
.def_property_readonly("max_degree", &DiskAnnIndexParams::max_degree,
|
||||
"int: max node degree.")
|
||||
|
|
@ -955,6 +1097,9 @@ Args:
|
|||
dict["pq_chunk_num"] = self.pq_chunk_num();
|
||||
dict["quantize_type"] =
|
||||
quantize_type_to_string(self.quantize_type());
|
||||
py::dict qp_dict;
|
||||
qp_dict["enable_rotate"] = self.quantizer_param().enable_rotate();
|
||||
dict["quantizer_param"] = qp_dict;
|
||||
return dict;
|
||||
},
|
||||
"Convert to dictionary with all fields")
|
||||
|
|
@ -968,20 +1113,25 @@ Args:
|
|||
", \"list_size\":" + std::to_string(self.list_size()) +
|
||||
", \"pq_chunk_num\":" + std::to_string(self.pq_chunk_num()) +
|
||||
", \"quantize_type\":" +
|
||||
quantize_type_to_string(self.quantize_type()) + "}";
|
||||
quantize_type_to_string(self.quantize_type()) +
|
||||
", \"quantizer_param\":{" + "\"enable_rotate\":" +
|
||||
(self.quantizer_param().enable_rotate() ? "true" : "false") +
|
||||
"}}";
|
||||
})
|
||||
.def(py::pickle(
|
||||
[](const DiskAnnIndexParams &self) {
|
||||
return py::make_tuple(self.metric_type(), self.max_degree(),
|
||||
self.list_size(), self.pq_chunk_num(),
|
||||
self.quantize_type());
|
||||
self.quantize_type(),
|
||||
self.quantizer_param().enable_rotate());
|
||||
},
|
||||
[](py::tuple t) {
|
||||
if (t.size() != 5)
|
||||
if (t.size() != 5 && t.size() != 6)
|
||||
throw std::runtime_error("Invalid state for DiskAnnIndexParams");
|
||||
QuantizerParam qp(t.size() >= 6 ? t[5].cast<bool>() : false);
|
||||
return std::make_shared<DiskAnnIndexParams>(
|
||||
t[0].cast<MetricType>(), t[1].cast<int>(), t[2].cast<int>(),
|
||||
t[3].cast<int>(), t[4].cast<QuantizeType>());
|
||||
t[3].cast<int>(), t[4].cast<QuantizeType>(), qp);
|
||||
}));
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -66,6 +66,7 @@ endif()
|
|||
# (real on Linux x86_64, stub on other platforms). Including them here causes
|
||||
# duplicate symbols and missing -laio when test binaries link both zvec_core
|
||||
# (via zvec) and core_knn_diskann.
|
||||
|
||||
list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/algorithm/diskann/.*")
|
||||
if(NOT DISKANN_SUPPORTED)
|
||||
list(FILTER ALL_CORE_SRCS EXCLUDE REGEX ".*/interface/indexes/diskann_index\\.cc")
|
||||
|
|
|
|||
|
|
@ -71,6 +71,16 @@ int IVFEntity::IVFReformerWrapper::init(const IndexMeta &imeta) {
|
|||
return 0;
|
||||
}
|
||||
|
||||
//! Load reformer state (e.g. rotation matrix) from storage
|
||||
int IVFEntity::IVFReformerWrapper::load(const IndexStorage::Pointer &storage) {
|
||||
if (!reformer_) {
|
||||
return 0;
|
||||
}
|
||||
int ret = reformer_->load(storage);
|
||||
ivf_check_with_msg(ret, "Failed to load reformer state");
|
||||
return 0;
|
||||
}
|
||||
|
||||
//! Update the params, Called by gpu searcher only
|
||||
int IVFEntity::IVFReformerWrapper::update(const IndexMeta &meta) {
|
||||
auto &name = meta.reformer_name();
|
||||
|
|
@ -503,6 +513,12 @@ int IVFEntity::load(const IndexStorage::Pointer &container) {
|
|||
|
||||
//! Load the remaining segments
|
||||
container_ = container;
|
||||
|
||||
//! Load reformer state (e.g. rotation matrix) from the main container,
|
||||
//! which holds the rotator segment dumped at build time.
|
||||
ret = reformer_.load(container);
|
||||
ivf_check_error_code(ret);
|
||||
|
||||
size_t expect_size = header_.inverted_body_size;
|
||||
inverted_ = load_segment(IVF_INVERTED_BODY_SEG_ID, expect_size);
|
||||
if (!inverted_) {
|
||||
|
|
|
|||
|
|
@ -267,6 +267,9 @@ class IVFEntity {
|
|||
//! Initialize
|
||||
int init(const IndexMeta &imeta);
|
||||
|
||||
//! Load reformer state (e.g. rotation matrix) from storage
|
||||
int load(const IndexStorage::Pointer &storage);
|
||||
|
||||
//! Update
|
||||
int update(const IndexMeta &meta);
|
||||
|
||||
|
|
|
|||
|
|
@ -86,6 +86,13 @@ int IVFSearcher::load(IndexStorage::Pointer container,
|
|||
}
|
||||
|
||||
auto reformer = centroid_index_->reformer();
|
||||
if (reformer) {
|
||||
//! The centroid index is loaded from the centroid sub-segment which does
|
||||
//! not contain the rotator segment. Load the reformer state (e.g. rotation
|
||||
//! matrix) from the main container instead.
|
||||
ret = reformer->load(container);
|
||||
ivf_check_error_code(ret);
|
||||
}
|
||||
params_.set(PARAM_IVF_SEARCHER_CONVERTER_REFORMER, reformer);
|
||||
|
||||
//! load iverted index
|
||||
|
|
|
|||
|
|
@ -86,6 +86,13 @@ int IVFStreamer::open(IndexStorage::Pointer storage) {
|
|||
}
|
||||
|
||||
auto reformer = centroid_index_->reformer();
|
||||
if (reformer) {
|
||||
//! The centroid index is loaded from the centroid sub-segment which does
|
||||
//! not contain the rotator segment. Load the reformer state (e.g. rotation
|
||||
//! matrix) from the main storage instead.
|
||||
ret = reformer->load(storage);
|
||||
ivf_check_error_code(ret);
|
||||
}
|
||||
params_.set(PARAM_IVF_SEARCHER_CONVERTER_REFORMER, reformer);
|
||||
|
||||
//! load iverted index
|
||||
|
|
|
|||
|
|
@ -182,6 +182,24 @@ int Index::CreateAndInitConverterReformer(const QuantizerParam ¶m,
|
|||
}
|
||||
}
|
||||
|
||||
// Pass enable_rotate to converter_params (effective for INT8 and INT4)
|
||||
if (param.enable_rotate) {
|
||||
if (param.type == QuantizerType::kInt8 ||
|
||||
param.type == QuantizerType::kInt4) {
|
||||
if (index_param.metric_type == MetricType::kCosine) {
|
||||
converter_params.set("cosine.converter.enable_rotate", true);
|
||||
} else {
|
||||
converter_params.set("integer_streaming.converter.enable_rotate", true);
|
||||
}
|
||||
} else {
|
||||
LOG_ERROR(
|
||||
"enable_rotate is only supported for INT8/INT4 quantizer, "
|
||||
"but got quantizer type: %d",
|
||||
static_cast<int>(param.type));
|
||||
return core::IndexError_Unsupported;
|
||||
}
|
||||
}
|
||||
|
||||
proxima_index_meta_.set_converter(converter_name, 0, converter_params);
|
||||
converter_ = core::IndexFactory::CreateConverter(converter_name);
|
||||
if (converter_ == nullptr ||
|
||||
|
|
@ -336,6 +354,25 @@ int Index::Open(const std::string &file_path, StorageOptions storage_options) {
|
|||
// converter/reformer/metric are created in IndexFactory::CreateIndex
|
||||
// TODO: init
|
||||
|
||||
// Load reformer data from storage (e.g., rotation matrix for
|
||||
// IntegerStreaming)
|
||||
if (reformer_ != nullptr) {
|
||||
// When building a new index, dump converter state (e.g., rotator) to
|
||||
// storage so the reformer can load it. This is needed for
|
||||
// enable_rotate with INT8 quantization.
|
||||
if (storage_options.create_new && converter_ != nullptr) {
|
||||
if (converter_->dump_to_storage(storage_) != 0) {
|
||||
LOG_ERROR("Failed to dump converter to storage, path: %s",
|
||||
file_path.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
}
|
||||
if (reformer_->load(storage_) != 0) {
|
||||
LOG_ERROR("Failed to load reformer, path: %s", file_path.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: context pool
|
||||
if (!init_context()) { // to validate if any error, will be overwritten
|
||||
LOG_ERROR("Failed to init context");
|
||||
|
|
|
|||
|
|
@ -251,12 +251,16 @@ ailego::JsonObject QuantizerParam::SerializeToJsonObject(
|
|||
json_obj.set("type",
|
||||
zvec::ailego::JsonValue(magic_enum::enum_name(type).data()));
|
||||
}
|
||||
if (!omit_empty_value || enable_rotate) {
|
||||
json_obj.set("enable_rotate", ailego::JsonValue(enable_rotate));
|
||||
}
|
||||
return json_obj;
|
||||
}
|
||||
|
||||
bool QuantizerParam::DeserializeFromJsonObject(
|
||||
const ailego::JsonObject &json_obj) {
|
||||
DESERIALIZE_ENUM_FIELD(json_obj, type, QuantizerType);
|
||||
DESERIALIZE_VALUE_FIELD(json_obj, enable_rotate);
|
||||
return true;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -121,6 +121,11 @@ int IVFIndex::Open(const std::string &file_path,
|
|||
LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
// Load reformer data from storage (e.g., rotation matrix for INT8+rotate)
|
||||
if (reformer_ != nullptr && reformer_->load(storage_) != 0) {
|
||||
LOG_ERROR("Failed to load reformer, path: %s", file_path_.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
is_trained_ = true;
|
||||
}
|
||||
is_open_ = true;
|
||||
|
|
@ -164,6 +169,11 @@ int IVFIndex::Train() {
|
|||
|
||||
dumper->create(file_path_);
|
||||
builder_->dump(dumper);
|
||||
// Dump converter state (e.g., rotator for INT8+rotate) to dumper
|
||||
if (converter_ && converter_->dump(dumper) != 0) {
|
||||
LOG_ERROR("Failed to dump converter, path: %s", file_path_.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
dumper->close();
|
||||
int ret = storage_->open(file_path_, false);
|
||||
if (ret != 0) {
|
||||
|
|
@ -175,6 +185,11 @@ int IVFIndex::Train() {
|
|||
LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
// Load reformer data from storage (e.g., rotation matrix)
|
||||
if (reformer_ != nullptr && reformer_->load(storage_) != 0) {
|
||||
LOG_ERROR("Failed to load reformer, path: %s", file_path_.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
is_trained_ = true;
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -209,7 +224,6 @@ int IVFIndex::_prepare_for_search(
|
|||
}
|
||||
|
||||
if (ivf_search_param->nprobe > 0) {
|
||||
// TODO: 1. sparse; 2. default ef
|
||||
ailego::Params params;
|
||||
params.set(core::PARAM_IVF_SEARCHER_NPROBE, ivf_search_param->nprobe);
|
||||
context->update(params);
|
||||
|
|
@ -227,6 +241,11 @@ int IVFIndex::Merge(const std::vector<Index::Pointer> &indexes,
|
|||
|
||||
dumper->create(file_path_);
|
||||
builder_->dump(dumper);
|
||||
// Dump converter state (e.g., rotator for INT8+rotate) to dumper
|
||||
if (converter_ && converter_->dump(dumper) != 0) {
|
||||
LOG_ERROR("Failed to dump converter, path: %s", file_path_.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
dumper->close();
|
||||
int ret = storage_->open(file_path_, false);
|
||||
if (ret != 0) {
|
||||
|
|
@ -238,6 +257,11 @@ int IVFIndex::Merge(const std::vector<Index::Pointer> &indexes,
|
|||
LOG_ERROR("Failed to open streamer, path: %s", file_path_.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
// Load reformer data from storage (e.g., rotation matrix)
|
||||
if (reformer_ != nullptr && reformer_->load(storage_) != 0) {
|
||||
LOG_ERROR("Failed to load reformer, path: %s", file_path_.c_str());
|
||||
return core::IndexError_Runtime;
|
||||
}
|
||||
is_trained_ = true;
|
||||
return 0;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ endif()
|
|||
cc_library(
|
||||
NAME core_quantizer
|
||||
STATIC SHARED STRICT ALWAYS_LINK
|
||||
SRCS *.cc
|
||||
SRCS *.cc rotator/*.cc
|
||||
LIBS zvec_ailego zvec_turbo core_framework
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core
|
||||
LDFLAGS "${CORE_QUANTIZER_LDFLAGS}"
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@
|
|||
#include <ailego/pattern/defer.h>
|
||||
#include <core/quantizer/quantizer_params.h>
|
||||
#include <zvec/core/framework/index_factory.h>
|
||||
#include "rotator/rotator.h"
|
||||
#include "record_quantizer.h"
|
||||
#include "../metric/metric_params.h"
|
||||
|
||||
|
|
@ -54,6 +55,10 @@ class CosineConverterHolder : public IndexHolder {
|
|||
type_ == IndexMeta::DataType::DT_INT8) {
|
||||
buffer_.resize(element_size, 0);
|
||||
}
|
||||
|
||||
if (owner_->rotator_) {
|
||||
rotate_buffer_.resize(owner_->rotator_->dimension());
|
||||
}
|
||||
}
|
||||
|
||||
this->convert_record();
|
||||
|
|
@ -116,17 +121,26 @@ class CosineConverterHolder : public IndexHolder {
|
|||
original_element_size);
|
||||
|
||||
float *buf = reinterpret_cast<float *>(&normalize_buffer_[0]);
|
||||
const float *vec = buf;
|
||||
|
||||
if (owner_->rotator_) {
|
||||
owner_->rotator_->rotate(vec, rotate_buffer_.data());
|
||||
vec = rotate_buffer_.data();
|
||||
}
|
||||
|
||||
float norm = 0.0f;
|
||||
ailego::Normalizer<float>::L2(buf, original_dimension_, &norm);
|
||||
ailego::Normalizer<float>::L2(const_cast<float *>(vec),
|
||||
original_dimension_, &norm);
|
||||
|
||||
if (type_ == IndexMeta::DataType::DT_FP32) {
|
||||
::memcpy(reinterpret_cast<float *>(&normalize_buffer_[0]), vec,
|
||||
original_dimension_ * sizeof(float));
|
||||
::memcpy(reinterpret_cast<float *>(&normalize_buffer_[0]) +
|
||||
original_dimension_,
|
||||
&norm, NORM_SIZE);
|
||||
} else if (type_ == IndexMeta::DataType::DT_FP16) {
|
||||
ailego::FloatHelper::ToFP16(
|
||||
buf, original_dimension_,
|
||||
const_cast<float *>(vec), original_dimension_,
|
||||
reinterpret_cast<uint16_t *>(&buffer_[0]));
|
||||
|
||||
::memcpy(
|
||||
|
|
@ -134,9 +148,8 @@ class CosineConverterHolder : public IndexHolder {
|
|||
&norm, NORM_SIZE);
|
||||
} else if (type_ == IndexMeta::DataType::DT_INT4 ||
|
||||
type_ == IndexMeta::DataType::DT_INT8) {
|
||||
RecordQuantizer::quantize_record(
|
||||
reinterpret_cast<const float *>(normalize_buffer_.data()),
|
||||
original_dimension_, type_, false, &buffer_[0]);
|
||||
RecordQuantizer::quantize_record(vec, original_dimension_, type_,
|
||||
false, &buffer_[0]);
|
||||
|
||||
::memcpy(reinterpret_cast<uint8_t *>(&buffer_[0]) + element_size -
|
||||
NORM_SIZE,
|
||||
|
|
@ -149,6 +162,7 @@ class CosineConverterHolder : public IndexHolder {
|
|||
const CosineConverterHolder *owner_{nullptr};
|
||||
std::string buffer_{};
|
||||
std::string normalize_buffer_{};
|
||||
std::vector<float> rotate_buffer_;
|
||||
IndexHolder::Iterator::Pointer front_iter_{};
|
||||
size_t dimension_{0u};
|
||||
size_t original_dimension_{0u};
|
||||
|
|
@ -159,11 +173,13 @@ class CosineConverterHolder : public IndexHolder {
|
|||
//! Constructor
|
||||
CosineConverterHolder(IndexHolder::Pointer front,
|
||||
IndexMeta::DataType original_type,
|
||||
IndexMeta::DataType type)
|
||||
IndexMeta::DataType type,
|
||||
std::shared_ptr<Rotator> rotator = nullptr)
|
||||
: front_(std::move(front)),
|
||||
original_type_(original_type),
|
||||
type_(type),
|
||||
dimension_(front_->dimension()) {}
|
||||
dimension_(front_->dimension()),
|
||||
rotator_(std::move(rotator)) {}
|
||||
|
||||
//! Retrieve count of elements in holder (-1 indicates unknown)
|
||||
size_t count(void) const override {
|
||||
|
|
@ -222,6 +238,7 @@ class CosineConverterHolder : public IndexHolder {
|
|||
IndexMeta::DataType original_type_{};
|
||||
IndexMeta::DataType type_{};
|
||||
uint32_t dimension_{0};
|
||||
std::shared_ptr<Rotator> rotator_{};
|
||||
};
|
||||
|
||||
/*! Converter of Cosine
|
||||
|
|
@ -264,8 +281,25 @@ class CosineConverter : public IndexConverter {
|
|||
return IndexError_Unsupported;
|
||||
}
|
||||
|
||||
params.get(COSINE_CONVERTER_ENABLE_ROTATE, &enable_rotate_);
|
||||
|
||||
ailego::Params reformer_params;
|
||||
|
||||
// Rotation only applies to integer quantization (INT8/INT4).
|
||||
if (enable_rotate_ && (dst_type_ == IndexMeta::DataType::DT_INT8 ||
|
||||
dst_type_ == IndexMeta::DataType::DT_INT4)) {
|
||||
size_t dim = index_meta.dimension();
|
||||
int ret = Rotator::create(&rotator_, dim);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("CosineConverter: create rotator failed, ret=%d, dim=%zu",
|
||||
ret, dim);
|
||||
return ret;
|
||||
}
|
||||
LOG_DEBUG("CosineConverter: rotation enabled, dim=%zu", dim);
|
||||
} else {
|
||||
enable_rotate_ = false;
|
||||
}
|
||||
|
||||
if (dst_type_ == IndexMeta::DataType::DT_INT8) {
|
||||
meta_.set_converter("CosineInt8Converter", 0, params);
|
||||
meta_.set_reformer("CosineInt8Reformer", 0, reformer_params);
|
||||
|
|
@ -333,12 +367,23 @@ class CosineConverter : public IndexConverter {
|
|||
*stats_.mutable_transformed_count() += holder->count();
|
||||
|
||||
holder_ = std::make_shared<CosineConverterHolder>(
|
||||
holder, holder->data_type(), dst_type_);
|
||||
holder, holder->data_type(), dst_type_, rotator_);
|
||||
return 0;
|
||||
}
|
||||
|
||||
//! Dump index into storage
|
||||
int dump(const IndexDumper::Pointer & /*dumper*/) override {
|
||||
int dump(const IndexDumper::Pointer &dumper) override {
|
||||
if (rotator_) {
|
||||
return rotator_->dump(dumper);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
//! Dump converter state to storage
|
||||
int dump_to_storage(const IndexStorage::Pointer &storage) override {
|
||||
if (rotator_) {
|
||||
return rotator_->dump(storage);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
|
@ -378,6 +423,8 @@ class CosineConverter : public IndexConverter {
|
|||
IndexHolder::Pointer holder_{};
|
||||
IndexMeta::DataType original_type_{IndexMeta::DataType::DT_UNDEFINED};
|
||||
IndexMeta::DataType dst_type_{IndexMeta::DataType::DT_UNDEFINED};
|
||||
bool enable_rotate_{false};
|
||||
std::shared_ptr<Rotator> rotator_{};
|
||||
};
|
||||
|
||||
INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(CosineNormalizeConverter,
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@
|
|||
#include <ailego/math/normalizer.h>
|
||||
#include <core/quantizer/quantizer_params.h>
|
||||
#include <zvec/core/framework/index_factory.h>
|
||||
#include "rotator/rotator.h"
|
||||
#include "record_quantizer.h"
|
||||
|
||||
namespace zvec {
|
||||
|
|
@ -53,7 +54,21 @@ class CosineReformer : public IndexReformer {
|
|||
}
|
||||
|
||||
//! Load index from container
|
||||
int load(IndexStorage::Pointer) override {
|
||||
//! Auto-detects rotation by checking for rotator segment in storage.
|
||||
int load(IndexStorage::Pointer storage) override {
|
||||
if (enable_rotate_ || storage->get(ROTATOR_SEG_ID)) {
|
||||
int ret = Rotator::open(&rotator_, storage);
|
||||
if (ret != 0) {
|
||||
if (enable_rotate_) {
|
||||
LOG_ERROR("CosineReformer: load rotator failed, ret=%d", ret);
|
||||
return ret;
|
||||
}
|
||||
} else {
|
||||
enable_rotate_ = true;
|
||||
LOG_DEBUG("CosineReformer: rotator auto-loaded, dim=%zu",
|
||||
rotator_->dimension());
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
|
@ -83,28 +98,35 @@ class CosineReformer : public IndexReformer {
|
|||
ometa->set_meta(dst_type_, qmeta.dimension() + ExtraDimension(dst_type_));
|
||||
out->resize(ometa->element_size());
|
||||
|
||||
float norm = 0.0f;
|
||||
size_t origin_dimension = qmeta.dimension();
|
||||
const float *vec = reinterpret_cast<const float *>(query);
|
||||
float norm = 0.0f;
|
||||
|
||||
// Fast path: no rotation — matches main branch behavior exactly
|
||||
std::string normalized_buffer(reinterpret_cast<const char *>(query),
|
||||
qmeta.element_size());
|
||||
|
||||
float *buf = reinterpret_cast<float *>(&normalized_buffer[0]);
|
||||
|
||||
if (enable_rotate_ && rotator_) {
|
||||
rotator_->rotate(vec, buf);
|
||||
}
|
||||
ailego::Normalizer<float>::L2(buf, origin_dimension, &norm);
|
||||
vec = buf;
|
||||
|
||||
::memcpy(reinterpret_cast<uint8_t *>(&(*out)[0]) + ometa->element_size() -
|
||||
NORM_SIZE,
|
||||
&norm, NORM_SIZE);
|
||||
|
||||
if (dst_type_ == IndexMeta::DataType::DT_FP32) {
|
||||
::memcpy(reinterpret_cast<uint8_t *>(&(*out)[0]), buf,
|
||||
::memcpy(reinterpret_cast<uint8_t *>(&(*out)[0]), vec,
|
||||
ometa->element_size() - NORM_SIZE);
|
||||
} else if (dst_type_ == IndexMeta::DataType::DT_FP16) {
|
||||
RecordQuantizer::quantize_record(buf, origin_dimension, dst_type_,
|
||||
false, &(*out)[0]);
|
||||
RecordQuantizer::quantize_record(const_cast<float *>(vec),
|
||||
qmeta.dimension(), dst_type_, false,
|
||||
&(*out)[0]);
|
||||
} else if (dst_type_ == IndexMeta::DataType::DT_INT4 ||
|
||||
dst_type_ == IndexMeta::DataType::DT_INT8) {
|
||||
RecordQuantizer::quantize_record(buf, qmeta.dimension(), dst_type_,
|
||||
RecordQuantizer::quantize_record(vec, qmeta.dimension(), dst_type_,
|
||||
false, &(*out)[0]);
|
||||
}
|
||||
} else if (type == IndexMeta::DataType::DT_FP16) {
|
||||
|
|
@ -186,6 +208,10 @@ class CosineReformer : public IndexReformer {
|
|||
NORM_SIZE,
|
||||
NORM_SIZE);
|
||||
|
||||
// Rotation only applies to INT8/INT4 targets (guarded at converter init).
|
||||
// For FP32/FP16 stored types, rotator_ is always null.
|
||||
const bool need_inv_rotate = (enable_rotate_ && rotator_);
|
||||
|
||||
if (type == IndexMeta::DataType::DT_FP32) {
|
||||
if (dst_type_ != IndexMeta::DataType::DT_FP32) {
|
||||
return IndexError_Unsupported;
|
||||
|
|
@ -195,6 +221,9 @@ class CosineReformer : public IndexReformer {
|
|||
const float *in_buf = reinterpret_cast<const float *>(in);
|
||||
|
||||
this->denormalize(in_buf, out_buf, qmeta, norm);
|
||||
if (need_inv_rotate) {
|
||||
rotator_->unrotate(out_buf, out_buf);
|
||||
}
|
||||
} else if (type == IndexMeta::DataType::DT_FP16) {
|
||||
if (dst_type_ != IndexMeta::DataType::DT_FP16) {
|
||||
return IndexError_Unsupported;
|
||||
|
|
@ -210,6 +239,7 @@ class CosineReformer : public IndexReformer {
|
|||
RecordQuantizer::unquantize_record(in, dimension, dst_type_, out_buf);
|
||||
|
||||
this->denormalize(out_buf, out_buf, qmeta, norm);
|
||||
// FP16 type path: no rotation was applied, skip inverse
|
||||
} else {
|
||||
ailego::Float16 *out_buf =
|
||||
reinterpret_cast<ailego::Float16 *>(&(*out)[0]);
|
||||
|
|
@ -228,6 +258,9 @@ class CosineReformer : public IndexReformer {
|
|||
RecordQuantizer::unquantize_record(in, dimension, dst_type_, out_buf);
|
||||
|
||||
this->denormalize(out_buf, out_buf, qmeta, norm);
|
||||
if (need_inv_rotate) {
|
||||
rotator_->unrotate(out_buf, out_buf);
|
||||
}
|
||||
}
|
||||
|
||||
return 0;
|
||||
|
|
@ -262,6 +295,8 @@ class CosineReformer : public IndexReformer {
|
|||
//! Members
|
||||
IndexMeta::DataType original_type_{IndexMeta::DataType::DT_UNDEFINED};
|
||||
IndexMeta::DataType dst_type_{IndexMeta::DataType::DT_UNDEFINED};
|
||||
bool enable_rotate_{false};
|
||||
std::shared_ptr<Rotator> rotator_{};
|
||||
};
|
||||
|
||||
INDEX_FACTORY_REGISTER_REFORMER_ALIAS(CosineNormalizeReformer, CosineReformer,
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@
|
|||
#include <ailego/pattern/defer.h>
|
||||
#include <core/quantizer/quantizer_params.h>
|
||||
#include <zvec/core/framework/index_factory.h>
|
||||
#include "rotator/rotator.h"
|
||||
#include "record_quantizer.h"
|
||||
#include "../metric/metric_params.h"
|
||||
|
||||
|
|
@ -378,6 +379,7 @@ class IntegerStreamingConverter : public IndexConverter {
|
|||
meta_ = index_meta;
|
||||
params.get(INTEGER_STREAMING_CONVERTER_ENABLE_NORMALIZE,
|
||||
&enable_normalize_);
|
||||
params.get(INTEGER_STREAMING_CONVERTER_ENABLE_ROTATE, &enable_rotate_);
|
||||
ailego::Params reformer_params;
|
||||
if (enable_normalize_) {
|
||||
reformer_params.set(INTEGER_STREAMING_REFORMER_ENABLE_NORMALIZE, true);
|
||||
|
|
@ -390,6 +392,17 @@ class IntegerStreamingConverter : public IndexConverter {
|
|||
reformer_params.set(INTEGER_STREAMING_REFORMER_IS_EUCLIDEAN, true);
|
||||
}
|
||||
|
||||
if (enable_rotate_) {
|
||||
int ret = Rotator::create(&rotator_, index_meta.dimension());
|
||||
if (ret != 0) {
|
||||
LOG_ERROR(
|
||||
"IntegerStreamingConverter: create rotator failed, ret=%d, dim=%u",
|
||||
ret, index_meta.dimension());
|
||||
return ret;
|
||||
}
|
||||
LOG_DEBUG("IntegerStreamingConverter: rotation enabled, dim=%zu",
|
||||
static_cast<size_t>(index_meta.dimension()));
|
||||
}
|
||||
|
||||
if (data_type_ == IndexMeta::DataType::DT_INT8) {
|
||||
meta_.set_converter("Int8StreamingConverter", 0, params);
|
||||
|
|
@ -433,12 +446,30 @@ class IntegerStreamingConverter : public IndexConverter {
|
|||
|
||||
*stats_.mutable_transformed_count() += holder->count();
|
||||
holder_ = std::make_shared<IntegerStreamingConverterHolder>(
|
||||
holder, data_type_, enable_normalize_, is_euclidean_);
|
||||
holder, data_type_, enable_normalize_, is_euclidean_, rotator_);
|
||||
return 0;
|
||||
}
|
||||
|
||||
//! Dump index into storage
|
||||
int dump(const IndexDumper::Pointer & /*dumper*/) override {
|
||||
//! Dump index into storage (writes rotator segment if rotate enabled)
|
||||
int dump(const IndexDumper::Pointer &dumper) override {
|
||||
if (enable_rotate_ && rotator_) {
|
||||
return rotator_->dump(dumper);
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
//! Dump converter state to IndexStorage for streaming build
|
||||
int dump_to_storage(const IndexStorage::Pointer &storage) override {
|
||||
if (enable_rotate_ && rotator_) {
|
||||
int ret = rotator_->dump(storage);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR(
|
||||
"IntegerStreamingConverter: dump rotator to storage failed, ret=%d",
|
||||
ret);
|
||||
return ret;
|
||||
}
|
||||
LOG_DEBUG("IntegerStreamingConverter: rotator dumped to storage");
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
|
@ -468,7 +499,8 @@ class IntegerStreamingConverter : public IndexConverter {
|
|||
IndexHolder::Iterator::Pointer &&iter)
|
||||
: owner_(owner),
|
||||
buffer_(owner->element_size(), 0),
|
||||
normalize_buffer_(owner->front_->element_size(), 0),
|
||||
normalize_buffer_(owner->dimension_ * sizeof(float), 0),
|
||||
rotate_buffer_(owner->dimension_ * sizeof(float), 0),
|
||||
front_iter_(std::move(iter)) {
|
||||
this->encode_record();
|
||||
}
|
||||
|
|
@ -503,18 +535,24 @@ class IntegerStreamingConverter : public IndexConverter {
|
|||
if (front_iter_->is_valid()) {
|
||||
const float *vec =
|
||||
reinterpret_cast<const float *>(front_iter_->data());
|
||||
size_t dim = owner_->dimension_;
|
||||
if (owner_->rotator_) {
|
||||
float *rotate_buf =
|
||||
reinterpret_cast<float *>(rotate_buffer_.data());
|
||||
owner_->rotator_->rotate(vec, rotate_buf);
|
||||
vec = rotate_buf;
|
||||
}
|
||||
if (owner_->enable_normalize_) {
|
||||
float norm = 0.0;
|
||||
memcpy((void *)normalize_buffer_.data(), vec,
|
||||
owner_->front_->element_size());
|
||||
memcpy((void *)normalize_buffer_.data(), vec, dim * sizeof(float));
|
||||
ailego::Normalizer<float>::L2((float *)normalize_buffer_.data(),
|
||||
owner_->dimension_, &norm);
|
||||
dim, &norm);
|
||||
vec = (float *)normalize_buffer_.data();
|
||||
}
|
||||
|
||||
RecordQuantizer::quantize_record(
|
||||
vec, owner_->dimension_, owner_->data_type(),
|
||||
owner_->is_euclidean_, buffer_.data());
|
||||
RecordQuantizer::quantize_record(vec, dim, owner_->data_type(),
|
||||
owner_->is_euclidean_,
|
||||
buffer_.data());
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -522,18 +560,21 @@ class IntegerStreamingConverter : public IndexConverter {
|
|||
const IntegerStreamingConverterHolder *owner_{nullptr};
|
||||
std::vector<uint8_t> buffer_{};
|
||||
std::string normalize_buffer_{};
|
||||
std::string rotate_buffer_{};
|
||||
IndexHolder::Iterator::Pointer front_iter_{};
|
||||
};
|
||||
|
||||
//! Constructor
|
||||
IntegerStreamingConverterHolder(IndexHolder::Pointer front,
|
||||
IndexMeta::DataType tp,
|
||||
bool enable_normalize, bool is_euclidean)
|
||||
bool enable_normalize, bool is_euclidean,
|
||||
std::shared_ptr<Rotator> rotator)
|
||||
: front_(std::move(front)),
|
||||
data_type_(tp),
|
||||
dimension_(front_->dimension()),
|
||||
enable_normalize_(enable_normalize),
|
||||
is_euclidean_(is_euclidean) {}
|
||||
is_euclidean_(is_euclidean),
|
||||
rotator_(std::move(rotator)) {}
|
||||
|
||||
//! Retrieve count of elements in holder (-1 indicates unknown)
|
||||
size_t count(void) const override {
|
||||
|
|
@ -576,6 +617,7 @@ class IntegerStreamingConverter : public IndexConverter {
|
|||
uint32_t dimension_{0};
|
||||
bool enable_normalize_{false};
|
||||
bool is_euclidean_{false};
|
||||
std::shared_ptr<Rotator> rotator_{};
|
||||
};
|
||||
|
||||
static size_t ExtraDimension(IndexMeta::DataType type) {
|
||||
|
|
@ -593,7 +635,9 @@ class IntegerStreamingConverter : public IndexConverter {
|
|||
IndexHolder::Pointer holder_{};
|
||||
IndexMeta::DataType data_type_{};
|
||||
bool enable_normalize_{false};
|
||||
bool enable_rotate_{false};
|
||||
bool is_euclidean_{false};
|
||||
std::shared_ptr<Rotator> rotator_{};
|
||||
};
|
||||
|
||||
INDEX_FACTORY_REGISTER_CONVERTER_ALIAS(
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@
|
|||
#include <ailego/pattern/defer.h>
|
||||
#include <core/quantizer/quantizer_params.h>
|
||||
#include <zvec/core/framework/index_factory.h>
|
||||
#include "rotator/rotator.h"
|
||||
#include "record_quantizer.h"
|
||||
|
||||
namespace zvec {
|
||||
|
|
@ -295,7 +296,27 @@ class IntegerStreamingReformer : public IndexReformer {
|
|||
}
|
||||
|
||||
//! Load index from container
|
||||
int load(IndexStorage::Pointer) override {
|
||||
//! Auto-detects rotation by checking for rotator segment in storage.
|
||||
//! No need for enable_rotate in search config.
|
||||
int load(IndexStorage::Pointer storage) override {
|
||||
// If config explicitly enables rotate but rotator not yet loaded, try
|
||||
// storage If config doesn't enable rotate, still try storage (auto-detect)
|
||||
if (enable_rotate_ || storage->get(ROTATOR_SEG_ID)) {
|
||||
int ret = Rotator::open(&rotator_, storage);
|
||||
if (ret != 0) {
|
||||
if (enable_rotate_) {
|
||||
// Config said enable_rotate but storage has no rotator — error
|
||||
LOG_ERROR("IntegerStreamingReformer: load rotator failed, ret=%d",
|
||||
ret);
|
||||
return ret;
|
||||
}
|
||||
// No rotator in storage, rotation not available
|
||||
} else {
|
||||
enable_rotate_ = true;
|
||||
LOG_DEBUG("IntegerStreamingReformer: rotator auto-loaded, dim=%zu",
|
||||
rotator_->dimension());
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
|
@ -319,10 +340,16 @@ class IntegerStreamingReformer : public IndexReformer {
|
|||
ometa->set_meta(data_type_, qmeta.dimension() + extra_dimension_);
|
||||
out->resize(ometa->element_size());
|
||||
const float *vec = reinterpret_cast<const float *>(query);
|
||||
std::unique_ptr<float[]> rotate_buffer;
|
||||
if (enable_rotate_ && rotator_) {
|
||||
rotate_buffer.reset(new float[rotator_->dimension()]);
|
||||
rotator_->rotate(vec, rotate_buffer.get());
|
||||
vec = rotate_buffer.get();
|
||||
}
|
||||
std::unique_ptr<float[]> normalized;
|
||||
if (enable_normalize_) {
|
||||
normalized.reset(new float[qmeta.dimension()]);
|
||||
vec = normalize(query, qmeta, normalized.get());
|
||||
vec = normalize(vec, qmeta, normalized.get());
|
||||
}
|
||||
|
||||
RecordQuantizer::quantize_record(vec, qmeta.dimension(), data_type_,
|
||||
|
|
@ -344,13 +371,21 @@ class IntegerStreamingReformer : public IndexReformer {
|
|||
*ometa = qmeta;
|
||||
ometa->set_meta(data_type_, qmeta.dimension() + extra_dimension_);
|
||||
out->resize(count * ometa->element_size());
|
||||
std::unique_ptr<float[]> rotate_buffer;
|
||||
std::unique_ptr<float[]> normalized;
|
||||
if (enable_rotate_ && rotator_) {
|
||||
rotate_buffer.reset(new float[rotator_->dimension()]);
|
||||
}
|
||||
if (enable_normalize_) {
|
||||
normalized.reset(new float[qmeta.dimension()]);
|
||||
}
|
||||
for (size_t i = 0; i < count; ++i) {
|
||||
const float *vec =
|
||||
reinterpret_cast<const float *>(query) + i * qmeta.dimension();
|
||||
if (enable_rotate_ && rotator_) {
|
||||
rotator_->rotate(vec, rotate_buffer.get());
|
||||
vec = rotate_buffer.get();
|
||||
}
|
||||
if (enable_normalize_) {
|
||||
vec = normalize(vec, qmeta, normalized.get());
|
||||
}
|
||||
|
|
@ -378,10 +413,16 @@ class IntegerStreamingReformer : public IndexReformer {
|
|||
ometa->set_meta(data_type_, rmeta.dimension() + extra_dimension_);
|
||||
out->resize(ometa->element_size());
|
||||
const float *vec = reinterpret_cast<const float *>(record);
|
||||
std::unique_ptr<float[]> rotate_buffer;
|
||||
if (enable_rotate_ && rotator_) {
|
||||
rotate_buffer.reset(new float[rotator_->dimension()]);
|
||||
rotator_->rotate(vec, rotate_buffer.get());
|
||||
vec = rotate_buffer.get();
|
||||
}
|
||||
std::unique_ptr<float[]> normalized;
|
||||
if (enable_normalize_) {
|
||||
normalized.reset(new float[rmeta.dimension()]);
|
||||
vec = normalize(record, rmeta, normalized.get());
|
||||
vec = normalize(vec, rmeta, normalized.get());
|
||||
}
|
||||
|
||||
RecordQuantizer::quantize_record(vec, rmeta.dimension(), data_type_,
|
||||
|
|
@ -404,13 +445,21 @@ class IntegerStreamingReformer : public IndexReformer {
|
|||
*ometa = rmeta;
|
||||
ometa->set_meta(data_type_, rmeta.dimension() + extra_dimension_);
|
||||
out->resize(count * ometa->element_size());
|
||||
std::unique_ptr<float[]> rotate_buffer;
|
||||
std::unique_ptr<float[]> normalized;
|
||||
if (enable_rotate_ && rotator_) {
|
||||
rotate_buffer.reset(new float[rotator_->dimension()]);
|
||||
}
|
||||
if (enable_normalize_) {
|
||||
normalized.reset(new float[rmeta.dimension()]);
|
||||
}
|
||||
for (size_t i = 0; i < count; ++i) {
|
||||
const float *vec =
|
||||
reinterpret_cast<const float *>(records) + i * rmeta.dimension();
|
||||
if (enable_rotate_ && rotator_) {
|
||||
rotator_->rotate(vec, rotate_buffer.get());
|
||||
vec = rotate_buffer.get();
|
||||
}
|
||||
if (enable_normalize_) {
|
||||
vec = normalize(vec, rmeta, normalized.get());
|
||||
}
|
||||
|
|
@ -447,15 +496,20 @@ class IntegerStreamingReformer : public IndexReformer {
|
|||
std::string *out) const override {
|
||||
if (enable_normalize_) {
|
||||
LOG_ERROR("Unsupported revert for normalized value");
|
||||
|
||||
return IndexError_Unsupported;
|
||||
}
|
||||
|
||||
out->resize((qmeta.dimension() - extra_dimension_) * sizeof(float));
|
||||
float *out_buf = reinterpret_cast<float *>(out->data());
|
||||
const size_t stored_dim = qmeta.dimension() - extra_dimension_;
|
||||
|
||||
RecordQuantizer::unquantize_record(in, qmeta.dimension() - extra_dimension_,
|
||||
data_type_, out_buf);
|
||||
// Step 1: Unquantize into out buffer (stored_dim floats)
|
||||
out->resize(stored_dim * sizeof(float));
|
||||
float *out_buf = reinterpret_cast<float *>(out->data());
|
||||
RecordQuantizer::unquantize_record(in, stored_dim, data_type_, out_buf);
|
||||
|
||||
// Step 2: Inverse rotate in-place if rotation was applied
|
||||
if (enable_rotate_ && rotator_) {
|
||||
rotator_->unrotate(out_buf, out_buf);
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
|
@ -465,6 +519,8 @@ class IntegerStreamingReformer : public IndexReformer {
|
|||
uint32_t extra_dimension_{0};
|
||||
bool enable_normalize_{false};
|
||||
bool is_euclidean_{false};
|
||||
bool enable_rotate_{false};
|
||||
std::shared_ptr<Rotator> rotator_{};
|
||||
};
|
||||
|
||||
INDEX_FACTORY_REGISTER_REFORMER_ALIAS(
|
||||
|
|
|
|||
|
|
@ -100,6 +100,8 @@ static const std::string INT4_QUANTIZER_REFORMER_METRIC =
|
|||
//! CosineConverter
|
||||
static const std::string COSINE_CONVERTER_FORCED_HALF_FLOAT =
|
||||
"cosine.converter.forced_half_float";
|
||||
static const std::string COSINE_CONVERTER_ENABLE_ROTATE =
|
||||
"cosine.converter.enable_rotate";
|
||||
|
||||
//! CosineReformer
|
||||
static const std::string COSINE_REFORMER_FORCED_HALF_FLOAT =
|
||||
|
|
@ -108,8 +110,10 @@ static const std::string COSINE_REFORMER_FORCED_HALF_FLOAT =
|
|||
//! IntegerStreamingConverter
|
||||
static const std::string INTEGER_STREAMING_CONVERTER_ENABLE_NORMALIZE =
|
||||
"integer_streaming.converter.enable_normalize";
|
||||
static const std::string INTEGER_STREAMING_CONVERTER_ENABLE_ROTATE =
|
||||
"integer_streaming.converter.enable_rotate";
|
||||
|
||||
//! IntegerStreamingConverter
|
||||
//! IntegerStreamingReformer
|
||||
static const std::string INTEGER_STREAMING_REFORMER_ENABLE_NORMALIZE =
|
||||
"integer_streaming.reformer.enable_normalize";
|
||||
static const std::string INTEGER_STREAMING_REFORMER_IS_EUCLIDEAN =
|
||||
|
|
|
|||
|
|
@ -0,0 +1,188 @@
|
|||
// 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 <cstring>
|
||||
#include <random>
|
||||
#include <zvec/core/framework/index_error.h>
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
namespace {
|
||||
|
||||
//! Largest power of 2 <= n (e.g. floor_pow2(97) = 64, floor_pow2(128) = 128).
|
||||
size_t floor_pow2(size_t n) {
|
||||
if (n == 0) return 0;
|
||||
size_t p = 1;
|
||||
while (p * 2 <= n) p *= 2;
|
||||
return p;
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// ============================================================================
|
||||
// FhtRotator method implementations
|
||||
// ============================================================================
|
||||
|
||||
int FhtRotator::init_impl(size_t dim) {
|
||||
if (dim == 0) {
|
||||
return IndexError_InvalidArgument;
|
||||
}
|
||||
trunc_dim = floor_pow2(dim);
|
||||
fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
|
||||
flip_offset_ = (dim + kByteLen - 1) / kByteLen;
|
||||
flip.resize(4 * flip_offset_);
|
||||
std::random_device rd;
|
||||
std::mt19937 gen(rd());
|
||||
std::uniform_int_distribution<int> dist(0, 255);
|
||||
for (auto &b : flip) b = static_cast<uint8_t>(dist(gen));
|
||||
return 0;
|
||||
}
|
||||
|
||||
void FhtRotator::rotate(const float *in, float *out) const {
|
||||
const size_t dim = dimension_;
|
||||
std::memcpy(out, in, sizeof(float) * dim);
|
||||
|
||||
if (trunc_dim == dim) {
|
||||
// Exact power-of-2: 4 rounds of (flip -> FHT -> rescale)
|
||||
ailego::fht_flip_sign(flip.data(), out, dim);
|
||||
ailego::fht_inplace(out, trunc_dim);
|
||||
ailego::fht_vec_rescale(out, trunc_dim, fac);
|
||||
|
||||
ailego::fht_flip_sign(flip.data() + flip_offset_, out, dim);
|
||||
ailego::fht_inplace(out, trunc_dim);
|
||||
ailego::fht_vec_rescale(out, trunc_dim, fac);
|
||||
|
||||
ailego::fht_flip_sign(flip.data() + 2 * flip_offset_, out, dim);
|
||||
ailego::fht_inplace(out, trunc_dim);
|
||||
ailego::fht_vec_rescale(out, trunc_dim, fac);
|
||||
|
||||
ailego::fht_flip_sign(flip.data() + 3 * flip_offset_, out, dim);
|
||||
ailego::fht_inplace(out, trunc_dim);
|
||||
ailego::fht_vec_rescale(out, trunc_dim, fac);
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
// Non-power-of-2 (e.g. 97, 100, 192, 320): 4 rounds with kacs_walk
|
||||
size_t start = dim - trunc_dim;
|
||||
float *trunc_ptr = out + start;
|
||||
|
||||
// Round 1: FHT on [0, trunc_dim)
|
||||
ailego::fht_flip_sign(flip.data(), out, dim);
|
||||
ailego::fht_inplace(out, trunc_dim);
|
||||
ailego::fht_vec_rescale(out, trunc_dim, fac);
|
||||
ailego::fht_kacs_walk(out, dim);
|
||||
|
||||
// Round 2: FHT on [start, start + trunc_dim)
|
||||
ailego::fht_flip_sign(flip.data() + flip_offset_, out, dim);
|
||||
ailego::fht_inplace(trunc_ptr, trunc_dim);
|
||||
ailego::fht_vec_rescale(trunc_ptr, trunc_dim, fac);
|
||||
ailego::fht_kacs_walk(out, dim);
|
||||
|
||||
// Round 3: FHT on [0, trunc_dim)
|
||||
ailego::fht_flip_sign(flip.data() + 2 * flip_offset_, out, dim);
|
||||
ailego::fht_inplace(out, trunc_dim);
|
||||
ailego::fht_vec_rescale(out, trunc_dim, fac);
|
||||
ailego::fht_kacs_walk(out, dim);
|
||||
|
||||
// Round 4: FHT on [start, start + trunc_dim)
|
||||
ailego::fht_flip_sign(flip.data() + 3 * flip_offset_, out, dim);
|
||||
ailego::fht_inplace(trunc_ptr, trunc_dim);
|
||||
ailego::fht_vec_rescale(trunc_ptr, trunc_dim, fac);
|
||||
ailego::fht_kacs_walk(out, dim);
|
||||
|
||||
// Final rescale: combine the 4 kacs_walk reductions
|
||||
ailego::fht_vec_rescale(out, dim, 0.25f);
|
||||
}
|
||||
|
||||
void FhtRotator::unrotate(const float *in, float *out) const {
|
||||
const size_t dim = dimension_;
|
||||
// Copy input into working buffer
|
||||
std::vector<float> data(in, in + dim);
|
||||
|
||||
if (trunc_dim == dim) {
|
||||
// Exact power-of-2: reverse 4 rounds in reverse order.
|
||||
const float inv_fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
|
||||
for (int round = 3; round >= 0; --round) {
|
||||
ailego::fht_inplace(data.data(), trunc_dim);
|
||||
ailego::fht_vec_rescale(data.data(), trunc_dim, inv_fac);
|
||||
ailego::fht_flip_sign(flip.data() + round * flip_offset_, data.data(),
|
||||
dim);
|
||||
}
|
||||
std::memcpy(out, data.data(), dim * sizeof(float));
|
||||
return;
|
||||
}
|
||||
|
||||
// Non-power-of-2: undo final rescale(0.25) first
|
||||
ailego::fht_vec_rescale(data.data(), dim, 4.0f);
|
||||
|
||||
const float inv_fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
|
||||
size_t start = dim - trunc_dim;
|
||||
float *trunc_ptr = data.data() + start;
|
||||
|
||||
// Undo Round 4 (FHT on [start, start+trunc_dim))
|
||||
ailego::fht_inv_kacs_walk(data.data(), dim);
|
||||
ailego::fht_inplace(trunc_ptr, trunc_dim);
|
||||
ailego::fht_vec_rescale(trunc_ptr, trunc_dim, inv_fac);
|
||||
ailego::fht_flip_sign(flip.data() + 3 * flip_offset_, data.data(), dim);
|
||||
|
||||
// Undo Round 3 (FHT on [0, trunc_dim))
|
||||
ailego::fht_inv_kacs_walk(data.data(), dim);
|
||||
ailego::fht_inplace(data.data(), trunc_dim);
|
||||
ailego::fht_vec_rescale(data.data(), trunc_dim, inv_fac);
|
||||
ailego::fht_flip_sign(flip.data() + 2 * flip_offset_, data.data(), dim);
|
||||
|
||||
// Undo Round 2 (FHT on [start, start+trunc_dim))
|
||||
ailego::fht_inv_kacs_walk(data.data(), dim);
|
||||
ailego::fht_inplace(trunc_ptr, trunc_dim);
|
||||
ailego::fht_vec_rescale(trunc_ptr, trunc_dim, inv_fac);
|
||||
ailego::fht_flip_sign(flip.data() + flip_offset_, data.data(), dim);
|
||||
|
||||
// Undo Round 1 (FHT on [0, trunc_dim))
|
||||
ailego::fht_inv_kacs_walk(data.data(), dim);
|
||||
ailego::fht_inplace(data.data(), trunc_dim);
|
||||
ailego::fht_vec_rescale(data.data(), trunc_dim, inv_fac);
|
||||
ailego::fht_flip_sign(flip.data(), data.data(), dim);
|
||||
|
||||
std::memcpy(out, data.data(), dim * sizeof(float));
|
||||
}
|
||||
|
||||
RotatorType FhtRotator::rotator_type() const {
|
||||
return RotatorType::FhtKac;
|
||||
}
|
||||
|
||||
void FhtRotator::save_blob(char *data) const {
|
||||
std::memcpy(data, flip.data(), flip.size());
|
||||
}
|
||||
|
||||
void FhtRotator::load_blob(const char *data) {
|
||||
// Recompute derived fields from dimension_ (init_impl is not called
|
||||
// during open, so trunc_dim/fac/flip_offset_ must be restored here)
|
||||
trunc_dim = floor_pow2(dimension_);
|
||||
fac = 1.0f / std::sqrt(static_cast<float>(trunc_dim));
|
||||
flip_offset_ = (dimension_ + kByteLen - 1) / kByteLen;
|
||||
// Load flip data
|
||||
flip.resize(4 * flip_offset_);
|
||||
std::memcpy(flip.data(), data, flip.size());
|
||||
}
|
||||
|
||||
size_t FhtRotator::blob_bytes() const {
|
||||
return flip.size();
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
@ -0,0 +1,60 @@
|
|||
// 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 <vector>
|
||||
#include <ailego/math/fht.h>
|
||||
#include "rotator.h"
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
// ============================================================================
|
||||
// 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 Rotator {
|
||||
public:
|
||||
FhtRotator() = default;
|
||||
~FhtRotator() override = default;
|
||||
|
||||
// Virtual interface
|
||||
void rotate(const float *in, float *out) const override;
|
||||
void unrotate(const float *in, float *out) const override;
|
||||
RotatorType rotator_type() const override;
|
||||
|
||||
protected:
|
||||
// Protected virtuals for base class factory/serialization
|
||||
int init_impl(size_t dim) override;
|
||||
size_t blob_bytes() const override;
|
||||
void save_blob(char *data) const override;
|
||||
void load_blob(const char *data) override;
|
||||
|
||||
private:
|
||||
std::vector<uint8_t> flip;
|
||||
size_t flip_offset_{0}; // bytes per round: ceil(dim / 8)
|
||||
size_t trunc_dim{0};
|
||||
float fac{0};
|
||||
|
||||
static constexpr size_t kByteLen = 8;
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
@ -0,0 +1,182 @@
|
|||
// 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 "matrix_rotator.h"
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstring>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
#include <zvec/core/framework/index_error.h>
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
namespace {
|
||||
|
||||
// Generate a dim x dim random Gaussian matrix (row-major) without Eigen.
|
||||
void random_gaussian_matrix(float *mat, size_t dim) {
|
||||
static std::random_device rd;
|
||||
static std::mt19937 gen(rd());
|
||||
std::normal_distribution<float> dist(0.0f, 1.0f);
|
||||
for (size_t i = 0; i < dim * dim; ++i) {
|
||||
mat[i] = dist(gen);
|
||||
}
|
||||
}
|
||||
|
||||
// Householder QR decomposition: A = Q * R.
|
||||
// Computes the orthogonal matrix Q from input matrix A (row-major, dim x dim).
|
||||
// Result is stored in q (row-major, dim x dim).
|
||||
//
|
||||
// Implemented manually to avoid rabitqlib/Eigen dependency whose ISA-sensitive
|
||||
// inline functions cause ODR violations (duplicate codegen with different
|
||||
// -march flags) leading to SEGFAULT on linux-x64-clang.
|
||||
void householder_qr(const float *A_in, float *q, size_t dim) {
|
||||
// R starts as a copy of A
|
||||
std::vector<float> R(A_in, A_in + dim * dim);
|
||||
|
||||
// Q starts as identity
|
||||
std::fill(q, q + dim * dim, 0.0f);
|
||||
for (size_t i = 0; i < dim; ++i) {
|
||||
q[i * dim + i] = 1.0f;
|
||||
}
|
||||
|
||||
std::vector<float> v(dim);
|
||||
|
||||
for (size_t k = 0; k < dim; ++k) {
|
||||
// x = R[k:dim, k] (sub-column below and including diagonal)
|
||||
float norm_x_sq = 0.0f;
|
||||
for (size_t i = k; i < dim; ++i) {
|
||||
norm_x_sq += R[i * dim + k] * R[i * dim + k];
|
||||
}
|
||||
if (norm_x_sq == 0.0f) continue;
|
||||
|
||||
float norm_x = std::sqrt(norm_x_sq);
|
||||
|
||||
// alpha = -sign(R[k][k]) * ||x|| (choose sign to avoid cancellation)
|
||||
float alpha = (R[k * dim + k] >= 0.0f) ? -norm_x : norm_x;
|
||||
|
||||
// v = x - alpha * e1 (only the sub-vector [k, dim) is non-zero)
|
||||
for (size_t i = k; i < dim; ++i) {
|
||||
v[i - k] = R[i * dim + k];
|
||||
}
|
||||
v[0] -= alpha;
|
||||
|
||||
// Normalize v
|
||||
float v_norm_sq = 0.0f;
|
||||
for (size_t i = 0; i < dim - k; ++i) {
|
||||
v_norm_sq += v[i] * v[i];
|
||||
}
|
||||
if (v_norm_sq == 0.0f) continue;
|
||||
float inv_v_norm = 1.0f / std::sqrt(v_norm_sq);
|
||||
for (size_t i = 0; i < dim - k; ++i) {
|
||||
v[i] *= inv_v_norm;
|
||||
}
|
||||
|
||||
// Apply Householder reflection to R: R[k:dim, k:dim] -= 2*v*(v^T * R)
|
||||
for (size_t j = k; j < dim; ++j) {
|
||||
float dot = 0.0f;
|
||||
for (size_t i = 0; i < dim - k; ++i) {
|
||||
dot += v[i] * R[(k + i) * dim + j];
|
||||
}
|
||||
dot *= 2.0f;
|
||||
for (size_t i = 0; i < dim - k; ++i) {
|
||||
R[(k + i) * dim + j] -= v[i] * dot;
|
||||
}
|
||||
}
|
||||
|
||||
// Accumulate Q: Q[:, k:dim] -= 2*(Q[:, k:dim] * v) * v^T
|
||||
for (size_t i = 0; i < dim; ++i) {
|
||||
float dot = 0.0f;
|
||||
for (size_t j = 0; j < dim - k; ++j) {
|
||||
dot += q[i * dim + k + j] * v[j];
|
||||
}
|
||||
dot *= 2.0f;
|
||||
for (size_t j = 0; j < dim - k; ++j) {
|
||||
q[i * dim + k + j] -= dot * v[j];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
int MatrixRotator::init_impl(size_t dim) {
|
||||
if (dim == 0) {
|
||||
return IndexError_InvalidArgument;
|
||||
}
|
||||
// Generate dim x dim random Gaussian matrix
|
||||
std::vector<float> rand_mat(dim * dim);
|
||||
random_gaussian_matrix(rand_mat.data(), dim);
|
||||
|
||||
// Householder QR: A = Q * R, use Q^T as the rotation matrix
|
||||
std::vector<float> Q(dim * dim);
|
||||
householder_qr(rand_mat.data(), Q.data(), dim);
|
||||
|
||||
// Store Q^T (transpose) as the rotation matrix
|
||||
matrix_.resize(dim * dim);
|
||||
for (size_t i = 0; i < dim; ++i) {
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
matrix_[j * dim + i] = Q[i * dim + j];
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
void MatrixRotator::rotate(const float *in, float *out) const {
|
||||
const size_t dim = dimension_;
|
||||
// out = in * matrix_ (1 x dim) * (dim x dim) -> (1 x dim)
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
float sum = 0.0f;
|
||||
for (size_t i = 0; i < dim; ++i) {
|
||||
sum += in[i] * matrix_[i * dim + j];
|
||||
}
|
||||
out[j] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
void MatrixRotator::unrotate(const float *in, float *out) const {
|
||||
const size_t dim = dimension_;
|
||||
// out = in * matrix_^T (1 x dim) * (dim x dim)^T -> (1 x dim)
|
||||
for (size_t j = 0; j < dim; ++j) {
|
||||
float sum = 0.0f;
|
||||
for (size_t i = 0; i < dim; ++i) {
|
||||
sum += in[i] * matrix_[j * dim + i];
|
||||
}
|
||||
out[j] = sum;
|
||||
}
|
||||
}
|
||||
|
||||
RotatorType MatrixRotator::rotator_type() const {
|
||||
return RotatorType::Matrix;
|
||||
}
|
||||
|
||||
void MatrixRotator::save_blob(char *data) const {
|
||||
std::memcpy(data, matrix_.data(), matrix_.size() * sizeof(float));
|
||||
}
|
||||
|
||||
void MatrixRotator::load_blob(const char *data) {
|
||||
// matrix_ must be pre-sized before loading
|
||||
if (matrix_.empty()) {
|
||||
matrix_.resize(dimension_ * dimension_);
|
||||
}
|
||||
std::memcpy(matrix_.data(), data, matrix_.size() * sizeof(float));
|
||||
}
|
||||
|
||||
size_t MatrixRotator::blob_bytes() const {
|
||||
return matrix_.size() * sizeof(float);
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
@ -0,0 +1,52 @@
|
|||
// 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 <vector>
|
||||
#include "rotator.h"
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
// ============================================================================
|
||||
// MatrixRotator - O(d^2) random orthogonal matrix rotation
|
||||
//
|
||||
// No alignment requirement on dimension. Uses a dim x dim square orthogonal
|
||||
// matrix generated via Householder QR on a random Gaussian matrix.
|
||||
// ============================================================================
|
||||
|
||||
class MatrixRotator : public Rotator {
|
||||
public:
|
||||
MatrixRotator() = default;
|
||||
~MatrixRotator() override = default;
|
||||
|
||||
// Virtual interface
|
||||
void rotate(const float *in, float *out) const override;
|
||||
void unrotate(const float *in, float *out) const override;
|
||||
RotatorType rotator_type() const override;
|
||||
|
||||
protected:
|
||||
// Protected virtuals for base class factory/serialization
|
||||
int init_impl(size_t dim) override;
|
||||
size_t blob_bytes() const override;
|
||||
void save_blob(char *data) const override;
|
||||
void load_blob(const char *data) override;
|
||||
|
||||
private:
|
||||
std::vector<float> matrix_; // dim x dim, row-major
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
@ -0,0 +1,335 @@
|
|||
// 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 "rotator.h"
|
||||
#include <cstring>
|
||||
#include <vector>
|
||||
#include <zvec/ailego/hash/crc32c.h>
|
||||
#include "zvec/core/framework/index_error.h"
|
||||
#include "zvec/core/framework/index_logger.h"
|
||||
#include "fht_rotator.h"
|
||||
#include "matrix_rotator.h"
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
namespace {
|
||||
|
||||
//! Read a little-endian uint32 from raw bytes.
|
||||
uint32_t read_u32_le(const char *p) {
|
||||
return static_cast<uint32_t>(static_cast<uint8_t>(p[0])) |
|
||||
(static_cast<uint32_t>(static_cast<uint8_t>(p[1])) << 8) |
|
||||
(static_cast<uint32_t>(static_cast<uint8_t>(p[2])) << 16) |
|
||||
(static_cast<uint32_t>(static_cast<uint8_t>(p[3])) << 24);
|
||||
}
|
||||
|
||||
//! Write a uint32 in little-endian to raw bytes.
|
||||
void write_u32_le(char *p, uint32_t v) {
|
||||
p[0] = static_cast<char>(v & 0xFF);
|
||||
p[1] = static_cast<char>((v >> 8) & 0xFF);
|
||||
p[2] = static_cast<char>((v >> 16) & 0xFF);
|
||||
p[3] = static_cast<char>((v >> 24) & 0xFF);
|
||||
}
|
||||
|
||||
//! Read a little-endian uint16 from raw bytes.
|
||||
uint16_t read_u16_le(const char *p) {
|
||||
return static_cast<uint16_t>(static_cast<uint8_t>(p[0])) |
|
||||
(static_cast<uint16_t>(static_cast<uint8_t>(p[1])) << 8);
|
||||
}
|
||||
|
||||
//! Write a uint16 in little-endian to raw bytes.
|
||||
void write_u16_le(char *p, uint16_t v) {
|
||||
p[0] = static_cast<char>(v & 0xFF);
|
||||
p[1] = static_cast<char>((v >> 8) & 0xFF);
|
||||
}
|
||||
|
||||
//! Serialization header (24 bytes, self-describing with magic).
|
||||
struct Header {
|
||||
uint32_t magic;
|
||||
uint16_t version;
|
||||
uint16_t rotator_type; // serialized: 0=Matrix, 1=Fht
|
||||
uint32_t in_dim;
|
||||
uint32_t out_dim;
|
||||
uint32_t payload_size;
|
||||
uint32_t reserved;
|
||||
|
||||
static uint16_t type_to_ser(RotatorType t) {
|
||||
return t == RotatorType::Matrix ? 0 : 1;
|
||||
}
|
||||
static RotatorType ser_to_type(uint16_t s) {
|
||||
return s == 0 ? RotatorType::Matrix : RotatorType::FhtKac;
|
||||
}
|
||||
|
||||
void write_to(char *buf) const {
|
||||
write_u32_le(buf + 0, magic);
|
||||
write_u16_le(buf + 4, version);
|
||||
write_u16_le(buf + 6, rotator_type);
|
||||
write_u32_le(buf + 8, in_dim);
|
||||
write_u32_le(buf + 12, out_dim);
|
||||
write_u32_le(buf + 16, payload_size);
|
||||
write_u32_le(buf + 20, reserved);
|
||||
}
|
||||
|
||||
void read_from(const char *buf) {
|
||||
magic = read_u32_le(buf + 0);
|
||||
version = read_u16_le(buf + 4);
|
||||
rotator_type = read_u16_le(buf + 6);
|
||||
in_dim = read_u32_le(buf + 8);
|
||||
out_dim = read_u32_le(buf + 12);
|
||||
payload_size = read_u32_le(buf + 16);
|
||||
reserved = read_u32_le(buf + 20);
|
||||
}
|
||||
};
|
||||
|
||||
} // anonymous namespace
|
||||
|
||||
// ============================================================================
|
||||
// Static factories
|
||||
// ============================================================================
|
||||
|
||||
int Rotator::create(std::shared_ptr<Rotator> *out, size_t dimension,
|
||||
RotatorType rotator_type) {
|
||||
*out = nullptr;
|
||||
std::unique_ptr<Rotator> rot;
|
||||
if (rotator_type == RotatorType::FhtKac) {
|
||||
rot = std::make_unique<FhtRotator>();
|
||||
} else {
|
||||
rot = std::make_unique<MatrixRotator>();
|
||||
}
|
||||
rot->dimension_ = dimension;
|
||||
int ret = rot->init_impl(dimension);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("Rotator::create: init_impl failed, ret=%d, dim=%zu", ret,
|
||||
dimension);
|
||||
return ret;
|
||||
}
|
||||
rot->initialized_ = true;
|
||||
*out = std::move(rot);
|
||||
return 0;
|
||||
}
|
||||
|
||||
int Rotator::open(std::shared_ptr<Rotator> *out, IndexStorage::Pointer storage,
|
||||
const std::string &seg_id) {
|
||||
*out = nullptr;
|
||||
if (!storage) {
|
||||
LOG_ERROR("Rotator::open: null storage");
|
||||
return IndexError_InvalidArgument;
|
||||
}
|
||||
|
||||
auto segment = storage->get(seg_id);
|
||||
if (!segment) {
|
||||
LOG_ERROR("Rotator::open: segment '%s' not found", seg_id.c_str());
|
||||
return IndexError_InvalidFormat;
|
||||
}
|
||||
|
||||
const size_t data_size = segment->data_size();
|
||||
if (data_size <= kHeaderSize) {
|
||||
LOG_ERROR("Rotator::open: data too small (%zu bytes)", data_size);
|
||||
return IndexError_InvalidFormat;
|
||||
}
|
||||
|
||||
IndexStorage::MemoryBlock block;
|
||||
size_t read_size = segment->read(0, block, data_size);
|
||||
if (read_size != data_size) {
|
||||
LOG_ERROR("Rotator::open: read failed, read=%zu, expected=%zu", read_size,
|
||||
data_size);
|
||||
return IndexError_InvalidFormat;
|
||||
}
|
||||
|
||||
// Verify CRC if available (covers header + blob)
|
||||
uint32_t expected_crc = segment->data_crc();
|
||||
if (expected_crc != 0) {
|
||||
uint32_t actual_crc = ailego::Crc32c::Hash(block.data(), data_size, 0);
|
||||
if (actual_crc != expected_crc) {
|
||||
LOG_ERROR("Rotator::open: CRC mismatch, expected=0x%08x, actual=0x%08x",
|
||||
expected_crc, actual_crc);
|
||||
return IndexError_InvalidFormat;
|
||||
}
|
||||
}
|
||||
|
||||
const char *raw = reinterpret_cast<const char *>(block.data());
|
||||
uint32_t maybe_magic = read_u32_le(raw);
|
||||
if (maybe_magic != kMagic) {
|
||||
LOG_ERROR("Rotator::open: invalid magic (expected 0x%08x, got 0x%08x)",
|
||||
kMagic, maybe_magic);
|
||||
return IndexError_InvalidFormat;
|
||||
}
|
||||
|
||||
Header header;
|
||||
header.read_from(raw);
|
||||
RotatorType type = Header::ser_to_type(header.rotator_type);
|
||||
size_t dim = static_cast<size_t>(header.in_dim);
|
||||
|
||||
// Reconstruct the rotator from header info and load blob
|
||||
std::unique_ptr<Rotator> rot;
|
||||
if (type == RotatorType::FhtKac) {
|
||||
rot = std::make_unique<FhtRotator>();
|
||||
} else {
|
||||
rot = std::make_unique<MatrixRotator>();
|
||||
}
|
||||
rot->dimension_ = dim;
|
||||
rot->load_blob(raw + kHeaderSize);
|
||||
rot->initialized_ = true;
|
||||
|
||||
LOG_DEBUG("Rotator::open done: seg=%s, dim=%zu, data_size=%zu",
|
||||
seg_id.c_str(), dim, data_size);
|
||||
|
||||
*out = std::move(rot);
|
||||
return 0;
|
||||
}
|
||||
|
||||
std::unique_ptr<Rotator> Rotator::load_matrix(const float *matrix,
|
||||
size_t dimension) {
|
||||
if (!matrix || dimension == 0) {
|
||||
LOG_ERROR("Rotator::load_matrix: invalid arguments");
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
std::unique_ptr<Rotator> rot = std::make_unique<MatrixRotator>();
|
||||
rot->dimension_ = dimension;
|
||||
rot->load_blob(reinterpret_cast<const char *>(matrix));
|
||||
rot->initialized_ = true;
|
||||
|
||||
LOG_DEBUG("Rotator::load_matrix done: dim=%zu", dimension);
|
||||
return rot;
|
||||
}
|
||||
|
||||
// ============================================================================
|
||||
// Non-virtual public methods
|
||||
// ============================================================================
|
||||
|
||||
std::vector<float> Rotator::rotate(const float *in) const {
|
||||
std::vector<float> out(dimension_);
|
||||
rotate(in, out.data());
|
||||
return out;
|
||||
}
|
||||
|
||||
std::vector<float> Rotator::unrotate(const float *in) const {
|
||||
std::vector<float> out(dimension_);
|
||||
unrotate(in, out.data());
|
||||
return out;
|
||||
}
|
||||
|
||||
size_t Rotator::dump_bytes() const {
|
||||
return kHeaderSize + blob_bytes();
|
||||
}
|
||||
|
||||
int Rotator::dump(const IndexStorage::Pointer &storage,
|
||||
const std::string &seg_id) const {
|
||||
if (!storage) {
|
||||
LOG_ERROR("Rotator::dump(storage): null storage");
|
||||
return IndexError_InvalidArgument;
|
||||
}
|
||||
if (!initialized_) {
|
||||
LOG_ERROR("Rotator::dump(storage): rotator not initialized");
|
||||
return IndexError_NoReady;
|
||||
}
|
||||
|
||||
auto align_size = [](size_t size) -> size_t {
|
||||
return (size + 0x1F) & (~0x1F);
|
||||
};
|
||||
|
||||
const size_t blob_size = blob_bytes();
|
||||
const size_t data_size = kHeaderSize + blob_size;
|
||||
const size_t total_size = align_size(data_size);
|
||||
std::vector<char> buffer(data_size);
|
||||
|
||||
Header header;
|
||||
header.magic = kMagic;
|
||||
header.version = kVersion;
|
||||
header.rotator_type = Header::type_to_ser(rotator_type());
|
||||
header.in_dim = static_cast<uint32_t>(dimension_);
|
||||
header.out_dim = static_cast<uint32_t>(dimension_);
|
||||
header.payload_size = static_cast<uint32_t>(blob_size);
|
||||
header.reserved = 0;
|
||||
header.write_to(buffer.data());
|
||||
save_blob(buffer.data() + kHeaderSize);
|
||||
|
||||
int ret = storage->append(seg_id, total_size);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("Rotator::dump(storage): append segment '%s' failed, ret=%d",
|
||||
seg_id.c_str(), ret);
|
||||
return ret;
|
||||
}
|
||||
|
||||
auto segment = storage->get(seg_id);
|
||||
if (!segment) {
|
||||
LOG_ERROR("Rotator::dump(storage): get segment '%s' failed",
|
||||
seg_id.c_str());
|
||||
return IndexError_WriteData;
|
||||
}
|
||||
|
||||
size_t written = segment->write(0, buffer.data(), data_size);
|
||||
if (written != data_size) {
|
||||
LOG_ERROR("Rotator::dump(storage): write failed, written=%zu, expected=%zu",
|
||||
written, data_size);
|
||||
return IndexError_WriteData;
|
||||
}
|
||||
segment->resize(data_size);
|
||||
segment->update_data_crc(ailego::Crc32c::Hash(buffer.data(), data_size, 0));
|
||||
|
||||
LOG_DEBUG("Rotator::dump(storage) done: seg=%s, data_size=%zu, total=%zu",
|
||||
seg_id.c_str(), data_size, total_size);
|
||||
return 0;
|
||||
}
|
||||
|
||||
int Rotator::dump(const IndexDumper::Pointer &dumper,
|
||||
const std::string &seg_id) const {
|
||||
if (!dumper) {
|
||||
LOG_ERROR("Rotator::dump(dumper): null dumper");
|
||||
return IndexError_InvalidArgument;
|
||||
}
|
||||
if (!initialized_) {
|
||||
LOG_ERROR("Rotator::dump(dumper): rotator not initialized");
|
||||
return IndexError_NoReady;
|
||||
}
|
||||
|
||||
const size_t blob_size = blob_bytes();
|
||||
const size_t data_size = kHeaderSize + blob_size;
|
||||
const size_t total_size = (data_size + 0x1F) & (~0x1F);
|
||||
|
||||
std::vector<char> buffer(total_size, 0);
|
||||
Header header;
|
||||
header.magic = kMagic;
|
||||
header.version = kVersion;
|
||||
header.rotator_type = Header::type_to_ser(rotator_type());
|
||||
header.in_dim = static_cast<uint32_t>(dimension_);
|
||||
header.out_dim = static_cast<uint32_t>(dimension_);
|
||||
header.payload_size = static_cast<uint32_t>(blob_size);
|
||||
header.reserved = 0;
|
||||
header.write_to(buffer.data());
|
||||
save_blob(buffer.data() + kHeaderSize);
|
||||
|
||||
const uint32_t crc = ailego::Crc32c::Hash(buffer.data(), data_size, 0);
|
||||
const size_t padding_size = total_size - data_size;
|
||||
|
||||
if (dumper->write(buffer.data(), total_size) != total_size) {
|
||||
LOG_ERROR("Rotator::dump(dumper): write failed, seg=%s", seg_id.c_str());
|
||||
return IndexError_WriteData;
|
||||
}
|
||||
|
||||
int ret = dumper->append(seg_id, data_size, padding_size, crc);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("Rotator::dump(dumper): append failed, seg=%s, ret=%d",
|
||||
seg_id.c_str(), ret);
|
||||
return ret;
|
||||
}
|
||||
|
||||
LOG_DEBUG("Rotator::dump(dumper) done: seg=%s, data_size=%zu, padding=%zu",
|
||||
seg_id.c_str(), data_size, padding_size);
|
||||
return 0;
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
@ -0,0 +1,153 @@
|
|||
// 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 <cstdint>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include "zvec/core/framework/index_dumper.h"
|
||||
#include "zvec/core/framework/index_storage.h"
|
||||
|
||||
namespace zvec {
|
||||
namespace core {
|
||||
|
||||
//! Segment ID used when dumping/loading the rotator data
|
||||
inline const std::string ROTATOR_SEG_ID{"enable_rotate"};
|
||||
|
||||
//! Rotator type
|
||||
enum class RotatorType : uint8_t {
|
||||
FhtKac = 0, //!< O(d log d) FHT-based Kac random rotation (default)
|
||||
Matrix = 1, //!< O(d^2) explicit random matrix rotation
|
||||
};
|
||||
|
||||
// Forward declarations for derived classes
|
||||
class FhtRotator;
|
||||
class MatrixRotator;
|
||||
|
||||
/*! Rotator provides per-vector rotation without external dependencies.
|
||||
*
|
||||
* Abstract base class for rotation algorithms. Use the static factory
|
||||
* methods to create instances:
|
||||
* - create() — build a new rotator from scratch
|
||||
* - open() — load from storage (auto-detects type from header)
|
||||
* - load_matrix() — load a user-specified rotation matrix
|
||||
*
|
||||
* Currently FhtRotator (O(d log d)) is the default and supports any
|
||||
* dimension. MatrixRotator (O(d^2)) is retained for future use.
|
||||
*
|
||||
* Rotation preserves dimension: output size == input size (no padding).
|
||||
*/
|
||||
class Rotator {
|
||||
public:
|
||||
virtual ~Rotator() = default;
|
||||
|
||||
// Static factories ---------------------------------------------------------
|
||||
|
||||
//! Create and initialize a new rotator.
|
||||
//! @param out on success, stores the rotator; on failure, nullptr
|
||||
//! @param dimension vector dimension (input and output size)
|
||||
//! @param rotator_type rotation algorithm (default: FhtKac)
|
||||
//! @return 0 on success, error code on failure
|
||||
static int create(std::shared_ptr<Rotator> *out, size_t dimension,
|
||||
RotatorType rotator_type = RotatorType::FhtKac);
|
||||
|
||||
//! Open a rotator from an IndexStorage segment (self-describing, no init
|
||||
//! needed). Parses header to get type/dimension, then reconstructs the
|
||||
//! appropriate derived class.
|
||||
//! @param out on success, stores the rotator; on failure, nullptr
|
||||
//! @param storage index storage
|
||||
//! @param seg_id segment identifier
|
||||
//! @return 0 on success, error code on failure
|
||||
static int open(std::shared_ptr<Rotator> *out, IndexStorage::Pointer storage,
|
||||
const std::string &seg_id = ROTATOR_SEG_ID);
|
||||
|
||||
//! Load a user-specified rotation matrix (always creates MatrixRotator).
|
||||
//! @param matrix row-major square matrix of shape dimension x dimension
|
||||
//! @param dimension vector dimension
|
||||
static std::unique_ptr<Rotator> load_matrix(const float *matrix,
|
||||
size_t dimension);
|
||||
|
||||
// Virtual interface --------------------------------------------------------
|
||||
|
||||
//! Rotate a single vector
|
||||
virtual void rotate(const float *in, float *out) const = 0;
|
||||
|
||||
//! Inverse-rotate a single vector (from rotated space back to original)
|
||||
virtual void unrotate(const float *in, float *out) const = 0;
|
||||
|
||||
//! Return the rotator type
|
||||
virtual RotatorType rotator_type() const = 0;
|
||||
|
||||
// Non-virtual public methods ----------------------------------------------
|
||||
|
||||
//! Rotate a single vector into a managed buffer
|
||||
std::vector<float> rotate(const float *in) const;
|
||||
|
||||
//! Inverse-rotate a single vector into a managed buffer
|
||||
std::vector<float> unrotate(const float *in) const;
|
||||
|
||||
//! Return the serialized size of the rotator in bytes (header + blob)
|
||||
size_t dump_bytes() const;
|
||||
|
||||
//! Dump the rotator to an IndexStorage as a named segment.
|
||||
int dump(const IndexStorage::Pointer &storage,
|
||||
const std::string &seg_id = ROTATOR_SEG_ID) const;
|
||||
|
||||
//! Dump the rotator to an IndexDumper as a named segment.
|
||||
//! Format: [Header (24B): magic|version|rotator_type|in_dim|
|
||||
//! out_dim|payload_size|reserved] [payload blob]
|
||||
//! Appends padding for 32-byte alignment.
|
||||
int dump(const IndexDumper::Pointer &dumper,
|
||||
const std::string &seg_id = ROTATOR_SEG_ID) const;
|
||||
|
||||
//! Return the vector dimension
|
||||
size_t dimension() const {
|
||||
return dimension_;
|
||||
}
|
||||
|
||||
//! Check if the rotator is initialized
|
||||
bool initialized() const {
|
||||
return initialized_;
|
||||
}
|
||||
|
||||
protected:
|
||||
// Protected virtuals — implemented by derived classes ---------------------
|
||||
|
||||
//! Initialize the rotator's internal state for the given dimension.
|
||||
//! @return 0 on success, error code on failure
|
||||
virtual int init_impl(size_t dim) = 0;
|
||||
|
||||
//! Return the serialized blob size (without header)
|
||||
virtual size_t blob_bytes() const = 0;
|
||||
|
||||
//! Write the payload blob to the given buffer
|
||||
virtual void save_blob(char *data) const = 0;
|
||||
|
||||
//! Read the payload blob from the given buffer
|
||||
virtual void load_blob(const char *data) = 0;
|
||||
|
||||
// Protected members --------------------------------------------------------
|
||||
|
||||
size_t dimension_{0};
|
||||
bool initialized_{false};
|
||||
|
||||
// Serialization constants (shared by dump/open)
|
||||
static constexpr size_t kHeaderSize = 24;
|
||||
static constexpr uint32_t kMagic = 0x52544F52; // "ROTR"
|
||||
static constexpr uint16_t kVersion = 1;
|
||||
};
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
|
@ -352,6 +352,7 @@ class ProximaEngineHelper {
|
|||
return tl::make_unexpected(
|
||||
Status::InvalidArgument("unsupported quantize type"));
|
||||
}
|
||||
index_param_builder->WithEnableRotate(db_index_params->enable_rotate());
|
||||
return index_param_builder;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -18,11 +18,12 @@ namespace zvec {
|
|||
|
||||
HnswIndexParams::OPtr ProtoConverter::FromPb(
|
||||
const proto::HnswIndexParams ¶ms_pb) {
|
||||
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
|
||||
auto params = std::make_shared<HnswIndexParams>(
|
||||
MetricTypeCodeBook::Get(params_pb.base().metric_type()), params_pb.m(),
|
||||
params_pb.ef_construction(),
|
||||
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
|
||||
params_pb.use_contiguous_memory());
|
||||
params_pb.use_contiguous_memory(), QuantizerParam(enable_rotate));
|
||||
|
||||
return params;
|
||||
}
|
||||
|
|
@ -33,6 +34,8 @@ proto::HnswIndexParams ProtoConverter::ToPb(const HnswIndexParams *params) {
|
|||
MetricTypeCodeBook::Get(params->metric_type()));
|
||||
params_pb.mutable_base()->set_quantize_type(
|
||||
QuantizeTypeCodeBook::Get(params->quantize_type()));
|
||||
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
|
||||
params->quantizer_param().enable_rotate());
|
||||
params_pb.set_ef_construction(params->ef_construction());
|
||||
params_pb.set_m(params->m());
|
||||
params_pb.set_use_contiguous_memory(params->use_contiguous_memory());
|
||||
|
|
@ -68,9 +71,11 @@ proto::HnswRabitqIndexParams ProtoConverter::ToPb(
|
|||
// FlatIndexParams
|
||||
FlatIndexParams::OPtr ProtoConverter::FromPb(
|
||||
const proto::FlatIndexParams ¶ms_pb) {
|
||||
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
|
||||
return std::make_shared<FlatIndexParams>(
|
||||
MetricTypeCodeBook::Get(params_pb.base().metric_type()),
|
||||
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()));
|
||||
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
|
||||
QuantizerParam(enable_rotate));
|
||||
}
|
||||
|
||||
proto::FlatIndexParams ProtoConverter::ToPb(const FlatIndexParams *params) {
|
||||
|
|
@ -79,16 +84,20 @@ proto::FlatIndexParams ProtoConverter::ToPb(const FlatIndexParams *params) {
|
|||
MetricTypeCodeBook::Get(params->metric_type()));
|
||||
params_pb.mutable_base()->set_quantize_type(
|
||||
QuantizeTypeCodeBook::Get(params->quantize_type()));
|
||||
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
|
||||
params->quantizer_param().enable_rotate());
|
||||
return params_pb;
|
||||
}
|
||||
|
||||
// IVFIndexParams
|
||||
IVFIndexParams::OPtr ProtoConverter::FromPb(
|
||||
const proto::IVFIndexParams ¶ms_pb) {
|
||||
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
|
||||
return std::make_shared<IVFIndexParams>(
|
||||
MetricTypeCodeBook::Get(params_pb.base().metric_type()),
|
||||
params_pb.n_list(), params_pb.n_iters(), params_pb.use_soar(),
|
||||
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()));
|
||||
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
|
||||
QuantizerParam(enable_rotate));
|
||||
}
|
||||
|
||||
proto::IVFIndexParams ProtoConverter::ToPb(const IVFIndexParams *params) {
|
||||
|
|
@ -97,6 +106,8 @@ proto::IVFIndexParams ProtoConverter::ToPb(const IVFIndexParams *params) {
|
|||
MetricTypeCodeBook::Get(params->metric_type()));
|
||||
params_pb.mutable_base()->set_quantize_type(
|
||||
QuantizeTypeCodeBook::Get(params->quantize_type()));
|
||||
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
|
||||
params->quantizer_param().enable_rotate());
|
||||
params_pb.set_n_list(params->n_list());
|
||||
params_pb.set_n_iters(params->n_iters());
|
||||
params_pb.set_use_soar(params->use_soar());
|
||||
|
|
@ -106,12 +117,14 @@ proto::IVFIndexParams ProtoConverter::ToPb(const IVFIndexParams *params) {
|
|||
// VamanaIndexParams
|
||||
VamanaIndexParams::OPtr ProtoConverter::FromPb(
|
||||
const proto::VamanaIndexParams ¶ms_pb) {
|
||||
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
|
||||
return std::make_shared<VamanaIndexParams>(
|
||||
MetricTypeCodeBook::Get(params_pb.base().metric_type()),
|
||||
params_pb.max_degree(), params_pb.search_list_size(), params_pb.alpha(),
|
||||
params_pb.saturate_graph(), params_pb.use_contiguous_memory(),
|
||||
params_pb.use_id_map(),
|
||||
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()));
|
||||
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
|
||||
QuantizerParam(enable_rotate));
|
||||
}
|
||||
|
||||
proto::VamanaIndexParams ProtoConverter::ToPb(const VamanaIndexParams *params) {
|
||||
|
|
@ -120,6 +133,8 @@ proto::VamanaIndexParams ProtoConverter::ToPb(const VamanaIndexParams *params) {
|
|||
MetricTypeCodeBook::Get(params->metric_type()));
|
||||
params_pb.mutable_base()->set_quantize_type(
|
||||
QuantizeTypeCodeBook::Get(params->quantize_type()));
|
||||
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
|
||||
params->quantizer_param().enable_rotate());
|
||||
params_pb.set_max_degree(params->max_degree());
|
||||
params_pb.set_search_list_size(params->search_list_size());
|
||||
params_pb.set_alpha(params->alpha());
|
||||
|
|
@ -147,10 +162,12 @@ proto::InvertIndexParams ProtoConverter::ToPb(const InvertIndexParams *params) {
|
|||
// DiskAnnIndexParams
|
||||
DiskAnnIndexParams::OPtr ProtoConverter::FromPb(
|
||||
const proto::DiskAnnIndexParams ¶ms_pb) {
|
||||
bool enable_rotate = params_pb.base().quantizer_param().enable_rotate();
|
||||
return std::make_shared<DiskAnnIndexParams>(
|
||||
MetricTypeCodeBook::Get(params_pb.base().metric_type()),
|
||||
params_pb.max_degree(), params_pb.list_size(), params_pb.pq_chunk_num(),
|
||||
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()));
|
||||
QuantizeTypeCodeBook::Get(params_pb.base().quantize_type()),
|
||||
QuantizerParam(enable_rotate));
|
||||
}
|
||||
|
||||
proto::DiskAnnIndexParams ProtoConverter::ToPb(
|
||||
|
|
@ -160,6 +177,8 @@ proto::DiskAnnIndexParams ProtoConverter::ToPb(
|
|||
MetricTypeCodeBook::Get(params->metric_type()));
|
||||
params_pb.mutable_base()->set_quantize_type(
|
||||
QuantizeTypeCodeBook::Get(params->quantize_type()));
|
||||
params_pb.mutable_base()->mutable_quantizer_param()->set_enable_rotate(
|
||||
params->quantizer_param().enable_rotate());
|
||||
params_pb.set_max_degree(params->max_degree());
|
||||
params_pb.set_list_size(params->list_size());
|
||||
params_pb.set_pq_chunk_num(params->pq_chunk_num());
|
||||
|
|
|
|||
|
|
@ -4036,7 +4036,8 @@ Status SegmentImpl::load_vector_index_blocks() {
|
|||
if (!segment_meta_->vector_indexed(column)) {
|
||||
new_field_params.set_index_params(MakeDefaultQuantVectorIndexParams(
|
||||
vector_index_params->metric_type(),
|
||||
vector_index_params->quantize_type()));
|
||||
vector_index_params->quantize_type(),
|
||||
vector_index_params->quantizer_param()));
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -4171,7 +4172,8 @@ Status SegmentImpl::init_memory_components() {
|
|||
block_id = allocate_block_id();
|
||||
FieldSchema normal_quant_field(*field);
|
||||
normal_quant_field.set_index_params(MakeDefaultQuantVectorIndexParams(
|
||||
index_params->metric_type(), index_params->quantize_type()));
|
||||
index_params->metric_type(), index_params->quantize_type(),
|
||||
index_params->quantizer_param()));
|
||||
auto quant_vector_indexer = create_vector_indexer(
|
||||
field->name(), normal_quant_field, block_id, true);
|
||||
|
||||
|
|
|
|||
|
|
@ -87,9 +87,19 @@ message InvertIndexParams {
|
|||
bool enable_range_optimization = 1;
|
||||
};
|
||||
|
||||
// Quantizer-related parameters for vector indexes.
|
||||
// Designed for future extensibility.
|
||||
message QuantizerParam {
|
||||
// When enabled, vectors are rotated before INT8 quantization to reduce
|
||||
// quantization error. Only effective with quantize_type=INT8.
|
||||
bool enable_rotate = 1;
|
||||
};
|
||||
|
||||
message BaseIndexParams {
|
||||
MetricType metric_type = 1;
|
||||
QuantizeType quantize_type = 2;
|
||||
// Quantizer parameters (enable_rotate, etc.)
|
||||
QuantizerParam quantizer_param = 4;
|
||||
};
|
||||
|
||||
message HnswIndexParams {
|
||||
|
|
|
|||
|
|
@ -961,6 +961,30 @@ ZVEC_EXPORT zvec_error_code_t ZVEC_CALL zvec_index_params_set_quantize_type(
|
|||
ZVEC_EXPORT zvec_quantize_type_t ZVEC_CALL
|
||||
zvec_index_params_get_quantize_type(const zvec_index_params_t *params);
|
||||
|
||||
/**
|
||||
* @brief Set enable_rotate for quantizer (only effective with INT8/INT4
|
||||
* quantize type)
|
||||
*
|
||||
* When enabled, vectors are randomly rotated before INT8/INT4 quantization to
|
||||
* reduce quantization error. The rotation matrix is stored with the index
|
||||
* and automatically applied to query vectors at search time.
|
||||
*
|
||||
* @param params Index parameters (must be vector index type)
|
||||
* @param enable_rotate Whether to enable random rotation before quantization
|
||||
* @return ZVEC_OK on success, error code on failure
|
||||
*/
|
||||
ZVEC_EXPORT zvec_error_code_t ZVEC_CALL
|
||||
zvec_index_params_set_quantizer_enable_rotate(zvec_index_params_t *params,
|
||||
bool enable_rotate);
|
||||
|
||||
/**
|
||||
* @brief Get enable_rotate setting from quantizer parameters
|
||||
* @param params Index parameters (must not be NULL)
|
||||
* @return true if rotation is enabled, false otherwise (default)
|
||||
*/
|
||||
ZVEC_EXPORT bool ZVEC_CALL zvec_index_params_get_quantizer_enable_rotate(
|
||||
const zvec_index_params_t *params);
|
||||
|
||||
/**
|
||||
* @brief Set HNSW specific parameters
|
||||
* @param params Index parameters (must be HNSW type)
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@
|
|||
#include <zvec/core/framework/index_holder.h>
|
||||
#include <zvec/core/framework/index_meta.h>
|
||||
#include <zvec/core/framework/index_stats.h>
|
||||
#include <zvec/core/framework/index_storage.h>
|
||||
#include "zvec/core/framework/index_reformer.h"
|
||||
|
||||
namespace zvec {
|
||||
|
|
@ -196,6 +197,13 @@ class IndexConverter : public IndexModule {
|
|||
//! Dump index into storage
|
||||
virtual int dump(const IndexDumper::Pointer &dumper) = 0;
|
||||
|
||||
//! Dump converter state (e.g. rotator) to IndexStorage for streaming build.
|
||||
//! Default is no-op; override in subclasses that need storage persistence.
|
||||
virtual int dump_to_storage(const IndexStorage::Pointer &storage) {
|
||||
(void)storage;
|
||||
return 0;
|
||||
}
|
||||
|
||||
//! Retrieve statistics
|
||||
virtual const Stats &stats(void) const = 0;
|
||||
|
||||
|
|
|
|||
|
|
@ -122,12 +122,17 @@ struct QuantizerParam : public SerializableBase {
|
|||
QuantizerType type = QuantizerType::kNone;
|
||||
int num_subquantizers = 8; // M
|
||||
int num_bits = 8; // bits per subquantizer
|
||||
bool enable_rotate =
|
||||
false; // rotate vectors before quantization to reduce error
|
||||
|
||||
// Constructors
|
||||
// QuantizerParam() = default;
|
||||
QuantizerParam(QuantizerType t = QuantizerType::kNone, int subquantizers = 8,
|
||||
int bits = 8)
|
||||
: type(t), num_subquantizers(subquantizers), num_bits(bits) {}
|
||||
int bits = 8, bool rotate = false)
|
||||
: type(t),
|
||||
num_subquantizers(subquantizers),
|
||||
num_bits(bits),
|
||||
enable_rotate(rotate) {}
|
||||
|
||||
|
||||
protected:
|
||||
|
|
|
|||
|
|
@ -87,6 +87,11 @@ class BaseIndexParamBuilder { // : public
|
|||
return static_cast<ActualIndexParamBuilderType &>(*this);
|
||||
}
|
||||
|
||||
ActualIndexParamBuilderType &WithEnableRotate(bool enable_rotate) {
|
||||
param->quantizer_param.enable_rotate = enable_rotate;
|
||||
return static_cast<ActualIndexParamBuilderType &>(*this);
|
||||
}
|
||||
|
||||
ActualIndexParamBuilderType &WithUseExternalVector(bool use_external_vector) {
|
||||
param->use_external_vector = use_external_vector;
|
||||
return static_cast<ActualIndexParamBuilderType &>(*this);
|
||||
|
|
|
|||
|
|
@ -118,16 +118,50 @@ class InvertIndexParams : public IndexParams {
|
|||
bool enable_extended_wildcard_{false};
|
||||
};
|
||||
|
||||
/*
|
||||
* Quantizer parameters for vector indexes.
|
||||
* Encapsulates quantization-related settings such as enable_rotate.
|
||||
* Designed for future extensibility (e.g., num_bits, calibration_size).
|
||||
*/
|
||||
class QuantizerParam {
|
||||
public:
|
||||
QuantizerParam() = default;
|
||||
explicit QuantizerParam(bool enable_rotate) : enable_rotate_(enable_rotate) {}
|
||||
|
||||
bool enable_rotate() const {
|
||||
return enable_rotate_;
|
||||
}
|
||||
|
||||
void set_enable_rotate(bool v) {
|
||||
enable_rotate_ = v;
|
||||
}
|
||||
|
||||
bool operator==(const QuantizerParam &other) const {
|
||||
return enable_rotate_ == other.enable_rotate_;
|
||||
}
|
||||
|
||||
bool operator!=(const QuantizerParam &other) const {
|
||||
return !(*this == other);
|
||||
}
|
||||
|
||||
private:
|
||||
// When enabled, vectors are rotated before INT8 quantization to reduce
|
||||
// quantization error. Only effective with quantize_type=INT8.
|
||||
bool enable_rotate_{false};
|
||||
};
|
||||
|
||||
/*
|
||||
* Column index params
|
||||
*/
|
||||
class VectorIndexParams : public IndexParams {
|
||||
public:
|
||||
VectorIndexParams(IndexType type, MetricType metric_type,
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED)
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED,
|
||||
QuantizerParam quantizer_param = {})
|
||||
: IndexParams(type),
|
||||
metric_type_(metric_type),
|
||||
quantize_type_(quantize_type) {}
|
||||
quantize_type_(quantize_type),
|
||||
quantizer_param_(quantizer_param) {}
|
||||
|
||||
~VectorIndexParams() override = default;
|
||||
|
||||
|
|
@ -151,9 +185,23 @@ class VectorIndexParams : public IndexParams {
|
|||
quantize_type_ = quantize_type;
|
||||
}
|
||||
|
||||
const QuantizerParam &quantizer_param() const {
|
||||
return quantizer_param_;
|
||||
}
|
||||
|
||||
void set_quantizer_param(const QuantizerParam &quantizer_param) {
|
||||
quantizer_param_ = quantizer_param;
|
||||
}
|
||||
|
||||
// Convenience getter for internal use (engine_helper, segment, etc.)
|
||||
bool enable_rotate() const {
|
||||
return quantizer_param_.enable_rotate();
|
||||
}
|
||||
|
||||
protected:
|
||||
MetricType metric_type_;
|
||||
QuantizeType quantize_type_;
|
||||
QuantizerParam quantizer_param_;
|
||||
};
|
||||
|
||||
/*
|
||||
|
|
@ -165,8 +213,9 @@ class HnswIndexParams : public VectorIndexParams {
|
|||
MetricType metric_type, int m = core_interface::kDefaultHnswNeighborCnt,
|
||||
int ef_construction = core_interface::kDefaultHnswEfConstruction,
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED,
|
||||
bool use_contiguous_memory = false)
|
||||
: VectorIndexParams(IndexType::HNSW, metric_type, quantize_type),
|
||||
bool use_contiguous_memory = false, QuantizerParam quantizer_param = {})
|
||||
: VectorIndexParams(IndexType::HNSW, metric_type, quantize_type,
|
||||
quantizer_param),
|
||||
m_(m),
|
||||
ef_construction_(ef_construction),
|
||||
use_contiguous_memory_(use_contiguous_memory) {}
|
||||
|
|
@ -175,9 +224,9 @@ class HnswIndexParams : public VectorIndexParams {
|
|||
|
||||
public:
|
||||
Ptr clone() const override {
|
||||
return std::make_shared<HnswIndexParams>(metric_type_, m_, ef_construction_,
|
||||
quantize_type_,
|
||||
use_contiguous_memory_);
|
||||
return std::make_shared<HnswIndexParams>(
|
||||
metric_type_, m_, ef_construction_, quantize_type_,
|
||||
use_contiguous_memory_, quantizer_param_);
|
||||
}
|
||||
|
||||
std::string to_string() const override {
|
||||
|
|
@ -186,7 +235,8 @@ class HnswIndexParams : public VectorIndexParams {
|
|||
std::ostringstream oss;
|
||||
oss << base_str << ",m:" << m_ << ",ef_construction:" << ef_construction_
|
||||
<< ",use_contiguous_memory:"
|
||||
<< (use_contiguous_memory_ ? "true" : "false") << "}";
|
||||
<< (use_contiguous_memory_ ? "true" : "false") << ",enable_rotate:"
|
||||
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
|
||||
return oss.str();
|
||||
}
|
||||
|
||||
|
|
@ -200,7 +250,9 @@ class HnswIndexParams : public VectorIndexParams {
|
|||
quantize_type() ==
|
||||
static_cast<const HnswIndexParams &>(other).quantize_type() &&
|
||||
use_contiguous_memory_ == static_cast<const HnswIndexParams &>(other)
|
||||
.use_contiguous_memory_;
|
||||
.use_contiguous_memory_ &&
|
||||
quantizer_param_ ==
|
||||
static_cast<const HnswIndexParams &>(other).quantizer_param_;
|
||||
}
|
||||
|
||||
void set_m(int m) {
|
||||
|
|
@ -348,21 +400,25 @@ class HnswRabitqIndexParams : public VectorIndexParams {
|
|||
class FlatIndexParams : public VectorIndexParams {
|
||||
public:
|
||||
FlatIndexParams(MetricType metric_type,
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED)
|
||||
: VectorIndexParams(IndexType::FLAT, metric_type, quantize_type) {}
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED,
|
||||
QuantizerParam quantizer_param = {})
|
||||
: VectorIndexParams(IndexType::FLAT, metric_type, quantize_type,
|
||||
quantizer_param) {}
|
||||
|
||||
using OPtr = std::shared_ptr<FlatIndexParams>;
|
||||
|
||||
public:
|
||||
Ptr clone() const override {
|
||||
return std::make_shared<FlatIndexParams>(metric_type_, quantize_type_);
|
||||
return std::make_shared<FlatIndexParams>(metric_type_, quantize_type_,
|
||||
quantizer_param_);
|
||||
}
|
||||
|
||||
std::string to_string() const override {
|
||||
auto base_str = vector_index_params_to_string("FlatIndexParams",
|
||||
metric_type_, quantize_type_);
|
||||
std::ostringstream oss;
|
||||
oss << base_str << "}";
|
||||
oss << base_str << ",enable_rotate:"
|
||||
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
|
||||
return oss.str();
|
||||
}
|
||||
|
||||
|
|
@ -371,7 +427,9 @@ class FlatIndexParams : public VectorIndexParams {
|
|||
metric_type() ==
|
||||
static_cast<const VectorIndexParams &>(other).metric_type() &&
|
||||
quantize_type() ==
|
||||
static_cast<const VectorIndexParams &>(other).quantize_type();
|
||||
static_cast<const VectorIndexParams &>(other).quantize_type() &&
|
||||
quantizer_param() ==
|
||||
static_cast<const VectorIndexParams &>(other).quantizer_param();
|
||||
}
|
||||
};
|
||||
|
||||
|
|
@ -383,16 +441,19 @@ inline FlatIndexParams MakeDefaultVectorIndexParams(MetricType metric_type) {
|
|||
}
|
||||
|
||||
inline FlatIndexParams MakeDefaultQuantVectorIndexParams(
|
||||
MetricType metric_type, QuantizeType quantize_type) {
|
||||
return FlatIndexParams(metric_type, quantize_type);
|
||||
MetricType metric_type, QuantizeType quantize_type,
|
||||
QuantizerParam quantizer_param = {}) {
|
||||
return FlatIndexParams(metric_type, quantize_type, quantizer_param);
|
||||
}
|
||||
|
||||
class IVFIndexParams : public VectorIndexParams {
|
||||
public:
|
||||
IVFIndexParams(MetricType metric_type, int n_list = 1024, int n_iters = 10,
|
||||
bool use_soar = false,
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED)
|
||||
: VectorIndexParams(IndexType::IVF, metric_type, quantize_type),
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED,
|
||||
QuantizerParam quantizer_param = {})
|
||||
: VectorIndexParams(IndexType::IVF, metric_type, quantize_type,
|
||||
quantizer_param),
|
||||
n_list_(n_list),
|
||||
n_iters_(n_iters),
|
||||
use_soar_(use_soar) {}
|
||||
|
|
@ -402,14 +463,17 @@ class IVFIndexParams : public VectorIndexParams {
|
|||
public:
|
||||
Ptr clone() const override {
|
||||
return std::make_shared<IVFIndexParams>(metric_type_, n_list_, n_iters_,
|
||||
use_soar_, quantize_type_);
|
||||
use_soar_, quantize_type_,
|
||||
quantizer_param_);
|
||||
}
|
||||
|
||||
std::string to_string() const override {
|
||||
auto base_str = vector_index_params_to_string("IVFIndexParams",
|
||||
metric_type_, quantize_type_);
|
||||
std::ostringstream oss;
|
||||
oss << base_str << ",n_list:" << n_list_ << ",n_iters:" << n_iters_ << "}";
|
||||
oss << base_str << ",n_list:" << n_list_ << ",n_iters:" << n_iters_
|
||||
<< ",enable_rotate:"
|
||||
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
|
||||
return oss.str();
|
||||
}
|
||||
|
||||
|
|
@ -445,7 +509,9 @@ class IVFIndexParams : public VectorIndexParams {
|
|||
n_iters_ == static_cast<const IVFIndexParams &>(other).n_iters_ &&
|
||||
use_soar_ == static_cast<const IVFIndexParams &>(other).use_soar_ &&
|
||||
quantize_type() ==
|
||||
static_cast<const IVFIndexParams &>(other).quantize_type();
|
||||
static_cast<const IVFIndexParams &>(other).quantize_type() &&
|
||||
quantizer_param_ ==
|
||||
static_cast<const IVFIndexParams &>(other).quantizer_param_;
|
||||
}
|
||||
|
||||
private:
|
||||
|
|
@ -458,8 +524,10 @@ class DiskAnnIndexParams : public VectorIndexParams {
|
|||
public:
|
||||
DiskAnnIndexParams(MetricType metric_type, int max_degree = 100,
|
||||
int list_size = 50, int pq_chunk_num = 0,
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED)
|
||||
: VectorIndexParams(IndexType::DISKANN, metric_type, quantize_type),
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED,
|
||||
QuantizerParam quantizer_param = {})
|
||||
: VectorIndexParams(IndexType::DISKANN, metric_type, quantize_type,
|
||||
quantizer_param),
|
||||
max_degree_{max_degree},
|
||||
list_size_{list_size},
|
||||
pq_chunk_num_{pq_chunk_num} {}
|
||||
|
|
@ -469,7 +537,8 @@ class DiskAnnIndexParams : public VectorIndexParams {
|
|||
public:
|
||||
Ptr clone() const override {
|
||||
return std::make_shared<DiskAnnIndexParams>(
|
||||
metric_type_, max_degree_, list_size_, pq_chunk_num_, quantize_type_);
|
||||
metric_type_, max_degree_, list_size_, pq_chunk_num_, quantize_type_,
|
||||
quantizer_param_);
|
||||
}
|
||||
|
||||
std::string to_string() const override {
|
||||
|
|
@ -478,7 +547,8 @@ class DiskAnnIndexParams : public VectorIndexParams {
|
|||
std::ostringstream oss;
|
||||
oss << base_str << ",max_degree:" << max_degree_
|
||||
<< ",list_size:" << list_size_ << ", pq_chunk_num:" << pq_chunk_num_
|
||||
<< "}";
|
||||
<< ",enable_rotate:"
|
||||
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
|
||||
return oss.str();
|
||||
}
|
||||
|
||||
|
|
@ -517,7 +587,9 @@ class DiskAnnIndexParams : public VectorIndexParams {
|
|||
pq_chunk_num_ ==
|
||||
static_cast<const DiskAnnIndexParams &>(other).pq_chunk_num_ &&
|
||||
quantize_type() ==
|
||||
static_cast<const DiskAnnIndexParams &>(other).quantize_type();
|
||||
static_cast<const DiskAnnIndexParams &>(other).quantize_type() &&
|
||||
quantizer_param_ ==
|
||||
static_cast<const DiskAnnIndexParams &>(other).quantizer_param_;
|
||||
}
|
||||
|
||||
private:
|
||||
|
|
@ -538,8 +610,10 @@ class VamanaIndexParams : public VectorIndexParams {
|
|||
float alpha = core_interface::kDefaultVamanaAlpha,
|
||||
bool saturate_graph = core_interface::kDefaultVamanaSaturateGraph,
|
||||
bool use_contiguous_memory = false, bool use_id_map = false,
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED)
|
||||
: VectorIndexParams(IndexType::VAMANA, metric_type, quantize_type),
|
||||
QuantizeType quantize_type = QuantizeType::UNDEFINED,
|
||||
QuantizerParam quantizer_param = {})
|
||||
: VectorIndexParams(IndexType::VAMANA, metric_type, quantize_type,
|
||||
quantizer_param),
|
||||
max_degree_(max_degree),
|
||||
search_list_size_(search_list_size),
|
||||
alpha_(alpha),
|
||||
|
|
@ -553,7 +627,7 @@ class VamanaIndexParams : public VectorIndexParams {
|
|||
Ptr clone() const override {
|
||||
return std::make_shared<VamanaIndexParams>(
|
||||
metric_type_, max_degree_, search_list_size_, alpha_, saturate_graph_,
|
||||
use_contiguous_memory_, use_id_map_, quantize_type_);
|
||||
use_contiguous_memory_, use_id_map_, quantize_type_, quantizer_param_);
|
||||
}
|
||||
|
||||
std::string to_string() const override {
|
||||
|
|
@ -565,7 +639,9 @@ class VamanaIndexParams : public VectorIndexParams {
|
|||
<< ",saturate_graph:" << (saturate_graph_ ? "true" : "false")
|
||||
<< ",use_contiguous_memory:"
|
||||
<< (use_contiguous_memory_ ? "true" : "false")
|
||||
<< ",use_id_map:" << (use_id_map_ ? "true" : "false") << "}";
|
||||
<< ",use_id_map:" << (use_id_map_ ? "true" : "false")
|
||||
<< ",enable_rotate:"
|
||||
<< (quantizer_param_.enable_rotate() ? "true" : "false") << "}";
|
||||
return oss.str();
|
||||
}
|
||||
|
||||
|
|
@ -580,7 +656,8 @@ class VamanaIndexParams : public VectorIndexParams {
|
|||
search_list_size_ == rhs.search_list_size_ && alpha_ == rhs.alpha_ &&
|
||||
saturate_graph_ == rhs.saturate_graph_ &&
|
||||
use_contiguous_memory_ == rhs.use_contiguous_memory_ &&
|
||||
use_id_map_ == rhs.use_id_map_;
|
||||
use_id_map_ == rhs.use_id_map_ &&
|
||||
quantizer_param_ == rhs.quantizer_param_;
|
||||
}
|
||||
|
||||
int max_degree() const {
|
||||
|
|
|
|||
|
|
@ -3491,6 +3491,179 @@ void test_index_params_functions(void) {
|
|||
TEST_END();
|
||||
}
|
||||
|
||||
void test_quantizer_enable_rotate(void) {
|
||||
TEST_START();
|
||||
|
||||
// Test 1: set enable_rotate=true on HNSW params and verify
|
||||
zvec_index_params_t *hnsw_params =
|
||||
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
|
||||
TEST_ASSERT(hnsw_params != NULL);
|
||||
|
||||
// Default should be false
|
||||
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(hnsw_params) ==
|
||||
false);
|
||||
|
||||
// Set to true and verify
|
||||
zvec_error_code_t err =
|
||||
zvec_index_params_set_quantizer_enable_rotate(hnsw_params, true);
|
||||
TEST_ASSERT(err == ZVEC_OK);
|
||||
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(hnsw_params) ==
|
||||
true);
|
||||
|
||||
// Set back to false and verify
|
||||
err = zvec_index_params_set_quantizer_enable_rotate(hnsw_params, false);
|
||||
TEST_ASSERT(err == ZVEC_OK);
|
||||
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(hnsw_params) ==
|
||||
false);
|
||||
|
||||
zvec_index_params_destroy(hnsw_params);
|
||||
|
||||
// Test 2: set enable_rotate on FLAT index params (also a vector index)
|
||||
zvec_index_params_t *flat_params =
|
||||
zvec_index_params_create(ZVEC_INDEX_TYPE_FLAT);
|
||||
TEST_ASSERT(flat_params != NULL);
|
||||
err = zvec_index_params_set_quantizer_enable_rotate(flat_params, true);
|
||||
TEST_ASSERT(err == ZVEC_OK);
|
||||
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(flat_params) ==
|
||||
true);
|
||||
zvec_index_params_destroy(flat_params);
|
||||
|
||||
// Test 3: set enable_rotate on non-vector index (INVERT) should fail
|
||||
zvec_index_params_t *invert_params =
|
||||
zvec_index_params_create(ZVEC_INDEX_TYPE_INVERT);
|
||||
TEST_ASSERT(invert_params != NULL);
|
||||
err = zvec_index_params_set_quantizer_enable_rotate(invert_params, true);
|
||||
TEST_ASSERT(err != ZVEC_OK);
|
||||
zvec_index_params_destroy(invert_params);
|
||||
|
||||
// Test 4: NULL params should return false for getter
|
||||
TEST_ASSERT(zvec_index_params_get_quantizer_enable_rotate(NULL) == false);
|
||||
|
||||
// Test 5: NULL params should return error for setter
|
||||
err = zvec_index_params_set_quantizer_enable_rotate(NULL, true);
|
||||
TEST_ASSERT(err != ZVEC_OK);
|
||||
|
||||
TEST_END();
|
||||
}
|
||||
|
||||
void test_int8_rotate_e2e(void) {
|
||||
TEST_START();
|
||||
|
||||
char temp_dir[] = "./zvec_test_int8_rotate_e2e";
|
||||
const size_t dim = 128;
|
||||
const size_t cnt = 2000;
|
||||
const size_t topk = 10;
|
||||
|
||||
// Create schema with HNSW + INT8 + enable_rotate
|
||||
zvec_collection_schema_t *schema =
|
||||
zvec_collection_schema_create("int8_rotate_test");
|
||||
TEST_ASSERT(schema != NULL);
|
||||
|
||||
// Add ID field
|
||||
zvec_field_schema_t *id_field =
|
||||
zvec_field_schema_create("id", ZVEC_DATA_TYPE_INT64, false, 0);
|
||||
zvec_collection_schema_add_field(schema, id_field);
|
||||
|
||||
// Add vector field with HNSW + INT8 + rotate
|
||||
zvec_index_params_t *hnsw_params =
|
||||
zvec_index_params_create(ZVEC_INDEX_TYPE_HNSW);
|
||||
TEST_ASSERT(hnsw_params != NULL);
|
||||
zvec_index_params_set_metric_type(hnsw_params, ZVEC_METRIC_TYPE_L2);
|
||||
zvec_index_params_set_hnsw_params(hnsw_params, 16, 100);
|
||||
zvec_index_params_set_quantize_type(hnsw_params, ZVEC_QUANTIZE_TYPE_INT8);
|
||||
zvec_index_params_set_quantizer_enable_rotate(hnsw_params, true);
|
||||
|
||||
zvec_field_schema_t *vec_field = zvec_field_schema_create(
|
||||
"embedding", ZVEC_DATA_TYPE_VECTOR_FP32, false, dim);
|
||||
zvec_field_schema_set_index_params(vec_field, hnsw_params);
|
||||
zvec_collection_schema_add_field(schema, vec_field);
|
||||
zvec_index_params_destroy(hnsw_params);
|
||||
|
||||
// Create and open collection
|
||||
zvec_collection_t *collection = NULL;
|
||||
zvec_error_code_t err =
|
||||
zvec_collection_create_and_open(temp_dir, schema, NULL, &collection);
|
||||
TEST_ASSERT(err == ZVEC_OK);
|
||||
TEST_ASSERT(collection != NULL);
|
||||
|
||||
// Insert 2000 random vectors
|
||||
srand(42);
|
||||
for (size_t i = 0; i < cnt; i++) {
|
||||
float *vec = (float *)malloc(dim * sizeof(float));
|
||||
TEST_ASSERT(vec != NULL);
|
||||
for (size_t j = 0; j < dim; j++) {
|
||||
vec[j] = (float)rand() / (float)RAND_MAX * 2.0f - 1.0f;
|
||||
}
|
||||
|
||||
zvec_doc_t *doc = zvec_doc_create();
|
||||
zvec_doc_set_pk(doc, zvec_test_make_pk(i + 1));
|
||||
zvec_doc_add_field_by_value(doc, "id", ZVEC_DATA_TYPE_INT64,
|
||||
&(int64_t){(int64_t)(i + 1)}, sizeof(int64_t));
|
||||
zvec_doc_add_field_by_value(doc, "embedding", ZVEC_DATA_TYPE_VECTOR_FP32,
|
||||
vec, dim * sizeof(float));
|
||||
|
||||
size_t success_count, error_count;
|
||||
const zvec_doc_t *docs[] = {doc};
|
||||
err = zvec_collection_insert(collection, docs, 1, &success_count,
|
||||
&error_count);
|
||||
TEST_ASSERT(err == ZVEC_OK);
|
||||
zvec_doc_destroy(doc);
|
||||
free(vec);
|
||||
}
|
||||
|
||||
// Flush to build index
|
||||
zvec_collection_flush(collection);
|
||||
|
||||
// Search
|
||||
float *query = (float *)malloc(dim * sizeof(float));
|
||||
TEST_ASSERT(query != NULL);
|
||||
for (size_t j = 0; j < dim; j++) {
|
||||
query[j] = (float)rand() / (float)RAND_MAX * 2.0f - 1.0f;
|
||||
}
|
||||
|
||||
zvec_vector_query_t *vq = zvec_vector_query_create();
|
||||
TEST_ASSERT(vq != NULL);
|
||||
zvec_vector_query_set_field_name(vq, "embedding");
|
||||
zvec_vector_query_set_query_vector(vq, query, dim * sizeof(float));
|
||||
zvec_vector_query_set_topk(vq, topk);
|
||||
|
||||
zvec_doc_t **results = NULL;
|
||||
size_t result_count = 0;
|
||||
err = zvec_collection_query(collection, vq, &results, &result_count);
|
||||
TEST_ASSERT(err == ZVEC_OK);
|
||||
TEST_ASSERT(result_count > 0);
|
||||
printf(" [int8_rotate_e2e] first search returned %zu results\n",
|
||||
result_count);
|
||||
zvec_docs_free(results, result_count);
|
||||
|
||||
// Close and reopen
|
||||
zvec_collection_close(collection);
|
||||
collection = NULL;
|
||||
|
||||
err = zvec_collection_open(temp_dir, NULL, &collection);
|
||||
TEST_ASSERT(err == ZVEC_OK);
|
||||
TEST_ASSERT(collection != NULL);
|
||||
|
||||
// Search again after reopen (rotator should auto-load from storage)
|
||||
results = NULL;
|
||||
result_count = 0;
|
||||
err = zvec_collection_query(collection, vq, &results, &result_count);
|
||||
TEST_ASSERT(err == ZVEC_OK);
|
||||
TEST_ASSERT(result_count > 0);
|
||||
printf(" [int8_rotate_e2e] reopen search returned %zu results\n",
|
||||
result_count);
|
||||
zvec_docs_free(results, result_count);
|
||||
|
||||
// Cleanup
|
||||
zvec_vector_query_destroy(vq);
|
||||
zvec_collection_destroy(collection);
|
||||
zvec_collection_schema_destroy(schema);
|
||||
free(query);
|
||||
cleanup_temp_directory(temp_dir);
|
||||
|
||||
TEST_END();
|
||||
}
|
||||
|
||||
void test_index_params_api_functions(void) {
|
||||
TEST_START();
|
||||
|
||||
|
|
@ -5992,6 +6165,8 @@ int main(void) {
|
|||
// Index tests
|
||||
test_index_params();
|
||||
test_index_params_functions();
|
||||
test_quantizer_enable_rotate();
|
||||
test_int8_rotate_e2e();
|
||||
test_index_params_api_functions();
|
||||
test_index_creation_and_management();
|
||||
|
||||
|
|
|
|||
|
|
@ -1209,6 +1209,100 @@ TEST_F(FlatStreamerTest, TestAddAndSearchWithID2) {
|
|||
EXPECT_GT(topk1Recall, 0.80f);
|
||||
}
|
||||
|
||||
// Test Flat + INT8 quantization + rotation end-to-end
|
||||
TEST_F(FlatStreamerTest, TestInt8WithRotate) {
|
||||
constexpr size_t kTestDim = 128;
|
||||
constexpr size_t kCnt = 2000U;
|
||||
constexpr size_t kTopk = 10;
|
||||
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
Params params;
|
||||
|
||||
IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, kTestDim);
|
||||
index_meta_raw.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
// Create INT8 converter with rotation enabled
|
||||
Params converter_params;
|
||||
converter_params.set("integer_streaming.converter.enable_rotate", true);
|
||||
auto converter = IndexFactory::CreateConverter("Int8StreamingConverter");
|
||||
ASSERT_NE(nullptr, converter);
|
||||
ASSERT_EQ(0, converter->init(index_meta_raw, converter_params));
|
||||
|
||||
IndexMeta index_meta = converter->meta();
|
||||
|
||||
auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name());
|
||||
ASSERT_NE(nullptr, reformer);
|
||||
ASSERT_EQ(0, reformer->init(index_meta.reformer_params()));
|
||||
|
||||
Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestInt8WithRotate.index", true));
|
||||
ASSERT_EQ(0, streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
// Add 2000 vectors
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kTestDim);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
|
||||
for (size_t i = 0; i < kCnt; i++) {
|
||||
NumericalVector<float> vec(kTestDim);
|
||||
for (size_t j = 0; j < kTestDim; ++j) vec[j] = dist(gen);
|
||||
|
||||
std::string new_vec;
|
||||
IndexQueryMeta new_meta;
|
||||
ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta));
|
||||
ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx));
|
||||
}
|
||||
|
||||
streamer->flush(0UL);
|
||||
streamer.reset();
|
||||
storage.reset();
|
||||
|
||||
// Reopen: reformer should auto-detect rotator from storage
|
||||
auto storage2 = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage2);
|
||||
ASSERT_EQ(0, storage2->init(stg_params));
|
||||
ASSERT_EQ(0, storage2->open(dir_ + "TestInt8WithRotate.index", false));
|
||||
|
||||
auto streamer2 = IndexFactory::CreateStreamer("FlatStreamer");
|
||||
ASSERT_NE(nullptr, streamer2);
|
||||
ASSERT_EQ(0, streamer2->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer2->open(storage2));
|
||||
|
||||
auto reformer2 = IndexFactory::CreateReformer(index_meta.reformer_name());
|
||||
ASSERT_NE(nullptr, reformer2);
|
||||
ASSERT_EQ(0, reformer2->init(index_meta.reformer_params()));
|
||||
ASSERT_EQ(0, reformer2->load(storage2));
|
||||
|
||||
// Search: verify results are non-empty
|
||||
auto knnCtx = streamer2->create_context();
|
||||
knnCtx->set_topk(kTopk);
|
||||
auto linearCtx = streamer2->create_context();
|
||||
linearCtx->set_topk(kTopk);
|
||||
|
||||
NumericalVector<float> query(kTestDim);
|
||||
for (size_t j = 0; j < kTestDim; ++j) query[j] = dist(gen);
|
||||
|
||||
std::string new_query;
|
||||
IndexQueryMeta new_qmeta;
|
||||
ASSERT_EQ(0,
|
||||
reformer2->transform(query.data(), qmeta, &new_query, &new_qmeta));
|
||||
ASSERT_EQ(0, streamer2->search_impl(new_query.data(), new_qmeta, knnCtx));
|
||||
ASSERT_EQ(0,
|
||||
streamer2->search_bf_impl(new_query.data(), new_qmeta, linearCtx));
|
||||
|
||||
EXPECT_EQ(kTopk, knnCtx->result().size());
|
||||
EXPECT_EQ(kTopk, linearCtx->result().size());
|
||||
}
|
||||
|
||||
#if defined(__GNUC__) || defined(__GNUG__)
|
||||
#pragma GCC diagnostic pop
|
||||
#endif
|
||||
|
|
@ -3779,6 +3779,105 @@ TEST_F(HnswStreamerTest, TestContiguousMultiThreadSearch) {
|
|||
s3.wait();
|
||||
}
|
||||
|
||||
// Test HNSW + INT8 quantization + rotation end-to-end
|
||||
TEST_F(HnswStreamerTest, TestInt8WithRotate) {
|
||||
constexpr size_t kTestDim = 128;
|
||||
constexpr size_t kCnt = 2000U;
|
||||
constexpr size_t kTopk = 10;
|
||||
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
ailego::Params params;
|
||||
params.set(PARAM_HNSW_STREAMER_MAX_NEIGHBOR_COUNT, 16U);
|
||||
params.set(PARAM_HNSW_STREAMER_SCALING_FACTOR, 5U);
|
||||
params.set(PARAM_HNSW_STREAMER_EFCONSTRUCTION, 100);
|
||||
params.set(PARAM_HNSW_STREAMER_EF, 100);
|
||||
params.set(PARAM_HNSW_STREAMER_BRUTE_FORCE_THRESHOLD, 1000U);
|
||||
|
||||
IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, kTestDim);
|
||||
index_meta_raw.set_metric("SquaredEuclidean", 0, ailego::Params());
|
||||
|
||||
// Create INT8 converter with rotation enabled
|
||||
ailego::Params converter_params;
|
||||
converter_params.set("integer_streaming.converter.enable_rotate", true);
|
||||
auto converter = IndexFactory::CreateConverter("Int8StreamingConverter");
|
||||
ASSERT_NE(nullptr, converter);
|
||||
ASSERT_EQ(0, converter->init(index_meta_raw, converter_params));
|
||||
|
||||
IndexMeta index_meta = converter->meta();
|
||||
|
||||
auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name());
|
||||
ASSERT_NE(nullptr, reformer);
|
||||
ASSERT_EQ(0, reformer->init(index_meta.reformer_params()));
|
||||
|
||||
ailego::Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestInt8WithRotate.index", true));
|
||||
ASSERT_EQ(0, streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
// Add 2000 vectors
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kTestDim);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
|
||||
for (size_t i = 0; i < kCnt; i++) {
|
||||
NumericalVector<float> vec(kTestDim);
|
||||
for (size_t j = 0; j < kTestDim; ++j) vec[j] = dist(gen);
|
||||
|
||||
std::string new_vec;
|
||||
IndexQueryMeta new_meta;
|
||||
ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta));
|
||||
ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx));
|
||||
}
|
||||
|
||||
streamer->flush(0UL);
|
||||
streamer.reset();
|
||||
storage.reset();
|
||||
|
||||
// Reopen: reformer should auto-detect rotator from storage
|
||||
auto storage2 = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage2);
|
||||
ASSERT_EQ(0, storage2->init(stg_params));
|
||||
ASSERT_EQ(0, storage2->open(dir_ + "TestInt8WithRotate.index", false));
|
||||
|
||||
auto streamer2 = IndexFactory::CreateStreamer("HnswStreamer");
|
||||
ASSERT_NE(nullptr, streamer2);
|
||||
ASSERT_EQ(0, streamer2->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer2->open(storage2));
|
||||
|
||||
auto reformer2 = IndexFactory::CreateReformer(index_meta.reformer_name());
|
||||
ASSERT_NE(nullptr, reformer2);
|
||||
ASSERT_EQ(0, reformer2->init(index_meta.reformer_params()));
|
||||
ASSERT_EQ(0, reformer2->load(storage2));
|
||||
|
||||
// Search: verify knn results are non-empty
|
||||
auto knnCtx = streamer2->create_context();
|
||||
knnCtx->set_topk(kTopk);
|
||||
auto linearCtx = streamer2->create_context();
|
||||
linearCtx->set_topk(kTopk);
|
||||
|
||||
NumericalVector<float> query(kTestDim);
|
||||
for (size_t j = 0; j < kTestDim; ++j) query[j] = dist(gen);
|
||||
|
||||
std::string new_query;
|
||||
IndexQueryMeta new_qmeta;
|
||||
ASSERT_EQ(0,
|
||||
reformer2->transform(query.data(), qmeta, &new_query, &new_qmeta));
|
||||
ASSERT_EQ(0, streamer2->search_impl(new_query.data(), new_qmeta, knnCtx));
|
||||
ASSERT_EQ(0,
|
||||
streamer2->search_bf_impl(new_query.data(), new_qmeta, linearCtx));
|
||||
|
||||
EXPECT_EQ(kTopk, knnCtx->result().size());
|
||||
EXPECT_EQ(kTopk, linearCtx->result().size());
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
||||
|
|
|
|||
|
|
@ -785,6 +785,105 @@ TEST_F(VamanaStreamerTest, TestConcurrentBuild) {
|
|||
ASSERT_GT(result.size(), 0UL);
|
||||
}
|
||||
|
||||
// Test Vamana + INT8 quantization + rotation end-to-end
|
||||
TEST_F(VamanaStreamerTest, TestInt8WithRotate) {
|
||||
constexpr size_t kTestDim = 128;
|
||||
constexpr size_t kCnt = 2000U;
|
||||
constexpr size_t kTopk = 10;
|
||||
|
||||
IndexStreamer::Pointer streamer =
|
||||
IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, streamer);
|
||||
|
||||
Params params;
|
||||
params.set(PARAM_VAMANA_STREAMER_MAX_DEGREE, 32U);
|
||||
params.set(PARAM_VAMANA_STREAMER_SEARCH_LIST_SIZE, 100U);
|
||||
params.set(PARAM_VAMANA_STREAMER_ALPHA, 1.2f);
|
||||
params.set(PARAM_VAMANA_STREAMER_EF, 64U);
|
||||
params.set(PARAM_VAMANA_STREAMER_BRUTE_FORCE_THRESHOLD, 500U);
|
||||
|
||||
IndexMeta index_meta_raw(IndexMeta::DataType::DT_FP32, kTestDim);
|
||||
index_meta_raw.set_metric("SquaredEuclidean", 0, Params());
|
||||
|
||||
// Create INT8 converter with rotation enabled
|
||||
Params converter_params;
|
||||
converter_params.set("integer_streaming.converter.enable_rotate", true);
|
||||
auto converter = IndexFactory::CreateConverter("Int8StreamingConverter");
|
||||
ASSERT_NE(nullptr, converter);
|
||||
ASSERT_EQ(0, converter->init(index_meta_raw, converter_params));
|
||||
|
||||
IndexMeta index_meta = converter->meta();
|
||||
|
||||
auto reformer = IndexFactory::CreateReformer(index_meta.reformer_name());
|
||||
ASSERT_NE(nullptr, reformer);
|
||||
ASSERT_EQ(0, reformer->init(index_meta.reformer_params()));
|
||||
|
||||
Params stg_params;
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage);
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(dir_ + "TestInt8WithRotate.index", true));
|
||||
ASSERT_EQ(0, streamer->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer->open(storage));
|
||||
|
||||
// Add 2000 vectors
|
||||
auto ctx = streamer->create_context();
|
||||
ASSERT_TRUE(!!ctx);
|
||||
IndexQueryMeta qmeta(IndexMeta::DataType::DT_FP32, kTestDim);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-1.0f, 1.0f);
|
||||
for (size_t i = 0; i < kCnt; i++) {
|
||||
NumericalVector<float> vec(kTestDim);
|
||||
for (size_t j = 0; j < kTestDim; ++j) vec[j] = dist(gen);
|
||||
|
||||
std::string new_vec;
|
||||
IndexQueryMeta new_meta;
|
||||
ASSERT_EQ(0, reformer->convert(vec.data(), qmeta, &new_vec, &new_meta));
|
||||
ASSERT_EQ(0, streamer->add_impl(i, new_vec.data(), new_meta, ctx));
|
||||
}
|
||||
|
||||
streamer->flush(0UL);
|
||||
streamer.reset();
|
||||
storage.reset();
|
||||
|
||||
// Reopen: reformer should auto-detect rotator from storage
|
||||
auto storage2 = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(nullptr, storage2);
|
||||
ASSERT_EQ(0, storage2->init(stg_params));
|
||||
ASSERT_EQ(0, storage2->open(dir_ + "TestInt8WithRotate.index", false));
|
||||
|
||||
auto streamer2 = IndexFactory::CreateStreamer("VamanaStreamer");
|
||||
ASSERT_NE(nullptr, streamer2);
|
||||
ASSERT_EQ(0, streamer2->init(index_meta, params));
|
||||
ASSERT_EQ(0, streamer2->open(storage2));
|
||||
|
||||
auto reformer2 = IndexFactory::CreateReformer(index_meta.reformer_name());
|
||||
ASSERT_NE(nullptr, reformer2);
|
||||
ASSERT_EQ(0, reformer2->init(index_meta.reformer_params()));
|
||||
ASSERT_EQ(0, reformer2->load(storage2));
|
||||
|
||||
// Search: verify knn results are non-empty
|
||||
auto knnCtx = streamer2->create_context();
|
||||
knnCtx->set_topk(kTopk);
|
||||
auto linearCtx = streamer2->create_context();
|
||||
linearCtx->set_topk(kTopk);
|
||||
|
||||
NumericalVector<float> query(kTestDim);
|
||||
for (size_t j = 0; j < kTestDim; ++j) query[j] = dist(gen);
|
||||
|
||||
std::string new_query;
|
||||
IndexQueryMeta new_qmeta;
|
||||
ASSERT_EQ(0,
|
||||
reformer2->transform(query.data(), qmeta, &new_query, &new_qmeta));
|
||||
ASSERT_EQ(0, streamer2->search_impl(new_query.data(), new_qmeta, knnCtx));
|
||||
ASSERT_EQ(0,
|
||||
streamer2->search_bf_impl(new_query.data(), new_qmeta, linearCtx));
|
||||
|
||||
EXPECT_EQ(kTopk, knnCtx->result().size());
|
||||
EXPECT_EQ(kTopk, linearCtx->result().size());
|
||||
}
|
||||
|
||||
} // namespace core
|
||||
} // namespace zvec
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ foreach(CC_SRCS ${ALL_TEST_SRCS})
|
|||
cc_gtest(
|
||||
NAME ${CC_TARGET}
|
||||
STRICT
|
||||
LIBS zvec_ailego core_framework core_quantizer
|
||||
LIBS zvec_ailego core_framework core_utility core_quantizer
|
||||
SRCS ${CC_SRCS}
|
||||
INCS . ${PROJECT_ROOT_DIR}/src/core/
|
||||
)
|
||||
|
|
|
|||
|
|
@ -12,10 +12,14 @@
|
|||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <cmath>
|
||||
#include <iostream>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
#include <gtest/gtest.h>
|
||||
#include <zvec/ailego/container/vector.h>
|
||||
#include "quantizer/rotator/rotator.h"
|
||||
#include "tests/test_util.h"
|
||||
#include "zvec/core/framework/index_factory.h"
|
||||
#include "zvec/core/framework/index_holder.h"
|
||||
|
||||
|
|
@ -821,3 +825,214 @@ TEST(IntegerReformer, Int4InitConverterWithTrainedParams) {
|
|||
EXPECT_EQ(buffer, buffer2);
|
||||
}
|
||||
}
|
||||
|
||||
// Test FhtKac rotator (dim=200, 4-aligned, non-power-of-2 kacs_walk path)
|
||||
TEST(RotatorTest, RotateUnrotateFhtKac_Dim200) {
|
||||
const size_t dim = 200;
|
||||
std::shared_ptr<Rotator> rotator;
|
||||
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
|
||||
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
|
||||
|
||||
std::vector<float> original(dim);
|
||||
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
|
||||
|
||||
std::vector<float> rotated(dim);
|
||||
rotator->rotate(original.data(), rotated.data());
|
||||
|
||||
std::vector<float> recovered(dim);
|
||||
rotator->unrotate(rotated.data(), recovered.data());
|
||||
|
||||
float max_err = 0.0f;
|
||||
for (size_t j = 0; j < dim; ++j)
|
||||
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
|
||||
std::cout << "FhtKac (dim=200) max error: " << max_err << std::endl;
|
||||
EXPECT_LT(max_err, 1e-3f);
|
||||
}
|
||||
|
||||
// Test FhtKac rotator (dim=96, 32-aligned but not 64-aligned, kacs_walk path)
|
||||
TEST(RotatorTest, RotateUnrotateFhtKac_Dim96) {
|
||||
const size_t dim = 96;
|
||||
std::shared_ptr<Rotator> rotator;
|
||||
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
|
||||
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
|
||||
|
||||
std::vector<float> original(dim);
|
||||
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
|
||||
|
||||
std::vector<float> rotated(dim);
|
||||
rotator->rotate(original.data(), rotated.data());
|
||||
|
||||
std::vector<float> recovered(dim);
|
||||
rotator->unrotate(rotated.data(), recovered.data());
|
||||
|
||||
float max_err = 0.0f;
|
||||
for (size_t j = 0; j < dim; ++j)
|
||||
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
|
||||
std::cout << "FhtKac (dim=96) max error: " << max_err << std::endl;
|
||||
EXPECT_LT(max_err, 1e-3f);
|
||||
}
|
||||
|
||||
// Test FhtKac rotator (dim=768, real-world embedding dimension, kacs_walk)
|
||||
TEST(RotatorTest, RotateUnrotateFhtKac_Dim768) {
|
||||
const size_t dim = 768;
|
||||
std::shared_ptr<Rotator> rotator;
|
||||
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
|
||||
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
|
||||
|
||||
std::vector<float> original(dim);
|
||||
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
|
||||
|
||||
std::vector<float> rotated(dim);
|
||||
rotator->rotate(original.data(), rotated.data());
|
||||
|
||||
std::vector<float> recovered(dim);
|
||||
rotator->unrotate(rotated.data(), recovered.data());
|
||||
|
||||
float max_err = 0.0f;
|
||||
for (size_t j = 0; j < dim; ++j)
|
||||
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
|
||||
std::cout << "FhtKac (dim=768) max error: " << max_err << std::endl;
|
||||
EXPECT_LT(max_err, 1e-3f);
|
||||
}
|
||||
|
||||
// Test FhtKac rotator (dim=128, power-of-2, pure FHT path)
|
||||
TEST(RotatorTest, RotateUnrotateFhtKac_Dim128) {
|
||||
const size_t dim = 128;
|
||||
std::shared_ptr<Rotator> rotator;
|
||||
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
|
||||
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
|
||||
|
||||
std::vector<float> original(dim);
|
||||
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
|
||||
|
||||
std::vector<float> rotated(dim);
|
||||
rotator->rotate(original.data(), rotated.data());
|
||||
|
||||
std::vector<float> recovered(dim);
|
||||
rotator->unrotate(rotated.data(), recovered.data());
|
||||
|
||||
float max_err = 0.0f;
|
||||
for (size_t j = 0; j < dim; ++j)
|
||||
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
|
||||
std::cout << "FhtKac (dim=128) max error: " << max_err << std::endl;
|
||||
EXPECT_LT(max_err, 1e-3f);
|
||||
}
|
||||
|
||||
// Test FhtKac rotator (dim=97, odd, non-4-aligned, non-power-of-2 kacs_walk)
|
||||
TEST(RotatorTest, RotateUnrotateFhtKac_Dim97) {
|
||||
const size_t dim = 97;
|
||||
std::shared_ptr<Rotator> rotator;
|
||||
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
|
||||
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
|
||||
|
||||
std::vector<float> original(dim);
|
||||
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
|
||||
|
||||
std::vector<float> rotated(dim);
|
||||
rotator->rotate(original.data(), rotated.data());
|
||||
|
||||
std::vector<float> recovered(dim);
|
||||
rotator->unrotate(rotated.data(), recovered.data());
|
||||
|
||||
float max_err = 0.0f;
|
||||
for (size_t j = 0; j < dim; ++j)
|
||||
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
|
||||
std::cout << "FhtKac (dim=97) max error: " << max_err << std::endl;
|
||||
EXPECT_LT(max_err, 1e-3f);
|
||||
}
|
||||
|
||||
// Test FhtKac rotator (dim=100, non-4-aligned, non-power-of-2 kacs_walk)
|
||||
TEST(RotatorTest, RotateUnrotateFhtKac_Dim100) {
|
||||
const size_t dim = 100;
|
||||
std::shared_ptr<Rotator> rotator;
|
||||
ASSERT_EQ(Rotator::create(&rotator, dim), 0);
|
||||
EXPECT_EQ(rotator->rotator_type(), RotatorType::FhtKac);
|
||||
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
|
||||
|
||||
std::vector<float> original(dim);
|
||||
for (size_t j = 0; j < dim; ++j) original[j] = dist(gen);
|
||||
|
||||
std::vector<float> rotated(dim);
|
||||
rotator->rotate(original.data(), rotated.data());
|
||||
|
||||
std::vector<float> recovered(dim);
|
||||
rotator->unrotate(rotated.data(), recovered.data());
|
||||
|
||||
float max_err = 0.0f;
|
||||
for (size_t j = 0; j < dim; ++j)
|
||||
max_err = std::max(max_err, std::abs(recovered[j] - original[j]));
|
||||
std::cout << "FhtKac (dim=100) max error: " << max_err << std::endl;
|
||||
EXPECT_LT(max_err, 1e-3f);
|
||||
}
|
||||
|
||||
// Test dump/open roundtrip: serialize then deserialize, verify rotate output
|
||||
// matches.
|
||||
TEST(RotatorTest, DumpOpenRoundtrip) {
|
||||
const std::string test_dir = "record_rotator_dump_test_dir/";
|
||||
zvec::test_util::RemoveTestPath(test_dir);
|
||||
|
||||
const size_t dim = 128;
|
||||
|
||||
// Build and dump original rotator
|
||||
std::shared_ptr<Rotator> original;
|
||||
ASSERT_EQ(Rotator::create(&original, dim), 0);
|
||||
EXPECT_EQ(original->rotator_type(), RotatorType::FhtKac);
|
||||
|
||||
auto storage = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(storage, nullptr);
|
||||
zvec::ailego::Params stg_params;
|
||||
ASSERT_EQ(0, storage->init(stg_params));
|
||||
ASSERT_EQ(0, storage->open(test_dir + "rotator.index", true));
|
||||
ASSERT_EQ(0, original->dump(storage));
|
||||
|
||||
// Close and reopen storage
|
||||
storage.reset();
|
||||
|
||||
auto storage2 = IndexFactory::CreateStorage("MMapFileStorage");
|
||||
ASSERT_NE(storage2, nullptr);
|
||||
ASSERT_EQ(0, storage2->init(stg_params));
|
||||
ASSERT_EQ(0, storage2->open(test_dir + "rotator.index", false));
|
||||
|
||||
// Load rotator from storage
|
||||
std::shared_ptr<Rotator> loaded;
|
||||
ASSERT_EQ(0, Rotator::open(&loaded, storage2));
|
||||
|
||||
// Verify metadata
|
||||
EXPECT_EQ(original->rotator_type(), loaded->rotator_type());
|
||||
EXPECT_EQ(original->dimension(), loaded->dimension());
|
||||
EXPECT_TRUE(loaded->initialized());
|
||||
|
||||
// Verify rotate output matches
|
||||
std::mt19937 gen(42);
|
||||
std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
|
||||
std::vector<float> vec(dim);
|
||||
for (size_t j = 0; j < dim; ++j) vec[j] = dist(gen);
|
||||
|
||||
auto rotated_orig = original->rotate(vec.data());
|
||||
auto rotated_loaded = loaded->rotate(vec.data());
|
||||
|
||||
float max_err = 0.0f;
|
||||
for (size_t j = 0; j < dim; ++j)
|
||||
max_err = std::max(max_err, std::abs(rotated_orig[j] - rotated_loaded[j]));
|
||||
std::cout << "DumpOpen roundtrip max error: " << max_err << std::endl;
|
||||
EXPECT_EQ(max_err, 0.0f);
|
||||
|
||||
zvec::test_util::RemoveTestPath(test_dir);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -470,4 +470,81 @@ TEST(ConverterTest, SegmentMetaWithEmptyFields) {
|
|||
EXPECT_EQ(pb_result.persisted_blocks_size(), 0);
|
||||
EXPECT_FALSE(pb_result.has_writing_forward_block());
|
||||
EXPECT_EQ(pb_result.indexed_vector_fields_size(), 0);
|
||||
}
|
||||
|
||||
// ==================== enable_rotate roundtrip tests ====================
|
||||
|
||||
TEST(ConverterTest, HnswIndexParamsWithEnableRotate) {
|
||||
// C++ -> PB -> C++ roundtrip with enable_rotate = true
|
||||
HnswIndexParams original(MetricType::COSINE, 16, 200, QuantizeType::INT8,
|
||||
false, QuantizerParam(true));
|
||||
EXPECT_TRUE(original.quantizer_param().enable_rotate());
|
||||
|
||||
auto pb = ProtoConverter::ToPb(&original);
|
||||
EXPECT_TRUE(pb.base().quantizer_param().enable_rotate());
|
||||
|
||||
auto restored = ProtoConverter::FromPb(pb);
|
||||
ASSERT_NE(restored, nullptr);
|
||||
EXPECT_TRUE(restored->quantizer_param().enable_rotate());
|
||||
EXPECT_TRUE(restored->enable_rotate()); // convenience getter
|
||||
EXPECT_EQ(restored->metric_type(), MetricType::COSINE);
|
||||
EXPECT_EQ(restored->m(), 16);
|
||||
EXPECT_EQ(restored->ef_construction(), 200);
|
||||
EXPECT_EQ(restored->quantize_type(), QuantizeType::INT8);
|
||||
|
||||
// C++ -> PB -> C++ roundtrip with enable_rotate = false
|
||||
HnswIndexParams original_no_rot(MetricType::L2, 32, 100, QuantizeType::FP16);
|
||||
auto pb2 = ProtoConverter::ToPb(&original_no_rot);
|
||||
EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate());
|
||||
auto restored2 = ProtoConverter::FromPb(pb2);
|
||||
ASSERT_NE(restored2, nullptr);
|
||||
EXPECT_FALSE(restored2->quantizer_param().enable_rotate());
|
||||
}
|
||||
|
||||
TEST(ConverterTest, FlatIndexParamsWithEnableRotate) {
|
||||
FlatIndexParams original(MetricType::IP, QuantizeType::INT8,
|
||||
QuantizerParam(true));
|
||||
EXPECT_TRUE(original.quantizer_param().enable_rotate());
|
||||
|
||||
auto pb = ProtoConverter::ToPb(&original);
|
||||
EXPECT_TRUE(pb.base().quantizer_param().enable_rotate());
|
||||
|
||||
auto restored = ProtoConverter::FromPb(pb);
|
||||
ASSERT_NE(restored, nullptr);
|
||||
EXPECT_TRUE(restored->quantizer_param().enable_rotate());
|
||||
EXPECT_EQ(restored->metric_type(), MetricType::IP);
|
||||
EXPECT_EQ(restored->quantize_type(), QuantizeType::INT8);
|
||||
|
||||
// enable_rotate = false
|
||||
FlatIndexParams original_no_rot(MetricType::L2, QuantizeType::FP16);
|
||||
auto pb2 = ProtoConverter::ToPb(&original_no_rot);
|
||||
EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate());
|
||||
auto restored2 = ProtoConverter::FromPb(pb2);
|
||||
EXPECT_FALSE(restored2->quantizer_param().enable_rotate());
|
||||
}
|
||||
|
||||
TEST(ConverterTest, IVFIndexParamsWithEnableRotate) {
|
||||
IVFIndexParams original(MetricType::COSINE, 256, 20, true, QuantizeType::INT8,
|
||||
QuantizerParam(true));
|
||||
EXPECT_TRUE(original.quantizer_param().enable_rotate());
|
||||
|
||||
auto pb = ProtoConverter::ToPb(&original);
|
||||
EXPECT_TRUE(pb.base().quantizer_param().enable_rotate());
|
||||
|
||||
auto restored = ProtoConverter::FromPb(pb);
|
||||
ASSERT_NE(restored, nullptr);
|
||||
EXPECT_TRUE(restored->quantizer_param().enable_rotate());
|
||||
EXPECT_EQ(restored->metric_type(), MetricType::COSINE);
|
||||
EXPECT_EQ(restored->n_list(), 256);
|
||||
EXPECT_EQ(restored->n_iters(), 20);
|
||||
EXPECT_TRUE(restored->use_soar());
|
||||
EXPECT_EQ(restored->quantize_type(), QuantizeType::INT8);
|
||||
|
||||
// enable_rotate = false
|
||||
IVFIndexParams original_no_rot(MetricType::L2, 128, 10, false,
|
||||
QuantizeType::FP16);
|
||||
auto pb2 = ProtoConverter::ToPb(&original_no_rot);
|
||||
EXPECT_FALSE(pb2.base().quantizer_param().enable_rotate());
|
||||
auto restored2 = ProtoConverter::FromPb(pb2);
|
||||
EXPECT_FALSE(restored2->quantizer_param().enable_rotate());
|
||||
}
|
||||
|
|
@ -186,4 +186,96 @@ TEST(IndexParamsTest, DynamicPointerCast) {
|
|||
IndexParams &base_ref = *base_ptr;
|
||||
auto &hnsw_ref = dynamic_cast<HnswIndexParams &>(base_ref);
|
||||
EXPECT_EQ(hnsw_ref.type(), IndexType::HNSW);
|
||||
}
|
||||
|
||||
// ==================== QuantizerParam tests ====================
|
||||
|
||||
TEST(IndexParamsTest, QuantizerParamBasic) {
|
||||
// Default constructor: enable_rotate should be false
|
||||
QuantizerParam qp_default;
|
||||
EXPECT_FALSE(qp_default.enable_rotate());
|
||||
|
||||
// Constructor with true
|
||||
QuantizerParam qp_true(true);
|
||||
EXPECT_TRUE(qp_true.enable_rotate());
|
||||
|
||||
// Constructor with false
|
||||
QuantizerParam qp_false(false);
|
||||
EXPECT_FALSE(qp_false.enable_rotate());
|
||||
|
||||
// Setter
|
||||
qp_default.set_enable_rotate(true);
|
||||
EXPECT_TRUE(qp_default.enable_rotate());
|
||||
qp_default.set_enable_rotate(false);
|
||||
EXPECT_FALSE(qp_default.enable_rotate());
|
||||
|
||||
// Equality
|
||||
EXPECT_TRUE(qp_true == QuantizerParam(true));
|
||||
EXPECT_TRUE(qp_false == QuantizerParam(false));
|
||||
EXPECT_FALSE(qp_true == qp_false);
|
||||
|
||||
// Inequality
|
||||
EXPECT_TRUE(qp_true != qp_false);
|
||||
EXPECT_FALSE(qp_true != QuantizerParam(true));
|
||||
}
|
||||
|
||||
TEST(IndexParamsTest, QuantizerParamWithVectorIndex) {
|
||||
// HnswIndexParams
|
||||
{
|
||||
HnswIndexParams params(MetricType::COSINE, 16, 100, QuantizeType::INT8);
|
||||
EXPECT_FALSE(params.quantizer_param().enable_rotate());
|
||||
EXPECT_FALSE(params.enable_rotate()); // convenience getter
|
||||
|
||||
params.set_quantizer_param(QuantizerParam(true));
|
||||
EXPECT_TRUE(params.quantizer_param().enable_rotate());
|
||||
EXPECT_TRUE(params.enable_rotate());
|
||||
|
||||
// Clone preserves quantizer_param
|
||||
auto cloned = params.clone();
|
||||
auto *cloned_hnsw = dynamic_cast<HnswIndexParams *>(cloned.get());
|
||||
ASSERT_NE(cloned_hnsw, nullptr);
|
||||
EXPECT_TRUE(cloned_hnsw->quantizer_param().enable_rotate());
|
||||
EXPECT_TRUE(*cloned == params);
|
||||
|
||||
// Equality: different enable_rotate -> not equal
|
||||
HnswIndexParams params2(MetricType::COSINE, 16, 100, QuantizeType::INT8);
|
||||
params2.set_quantizer_param(QuantizerParam(false));
|
||||
EXPECT_FALSE(params == params2);
|
||||
}
|
||||
|
||||
// FlatIndexParams
|
||||
{
|
||||
FlatIndexParams params(MetricType::L2, QuantizeType::INT8);
|
||||
EXPECT_FALSE(params.quantizer_param().enable_rotate());
|
||||
|
||||
params.set_quantizer_param(QuantizerParam(true));
|
||||
EXPECT_TRUE(params.quantizer_param().enable_rotate());
|
||||
EXPECT_TRUE(params.enable_rotate());
|
||||
|
||||
auto cloned = params.clone();
|
||||
auto *cloned_flat = dynamic_cast<FlatIndexParams *>(cloned.get());
|
||||
ASSERT_NE(cloned_flat, nullptr);
|
||||
EXPECT_TRUE(cloned_flat->quantizer_param().enable_rotate());
|
||||
|
||||
FlatIndexParams params2(MetricType::L2, QuantizeType::INT8);
|
||||
EXPECT_FALSE(params == params2);
|
||||
}
|
||||
|
||||
// IVFIndexParams
|
||||
{
|
||||
IVFIndexParams params(MetricType::IP, 128, 10, false, QuantizeType::INT8);
|
||||
EXPECT_FALSE(params.quantizer_param().enable_rotate());
|
||||
|
||||
params.set_quantizer_param(QuantizerParam(true));
|
||||
EXPECT_TRUE(params.quantizer_param().enable_rotate());
|
||||
EXPECT_TRUE(params.enable_rotate());
|
||||
|
||||
auto cloned = params.clone();
|
||||
auto *cloned_ivf = dynamic_cast<IVFIndexParams *>(cloned.get());
|
||||
ASSERT_NE(cloned_ivf, nullptr);
|
||||
EXPECT_TRUE(cloned_ivf->quantizer_param().enable_rotate());
|
||||
|
||||
IVFIndexParams params2(MetricType::IP, 128, 10, false, QuantizeType::INT8);
|
||||
EXPECT_FALSE(params == params2);
|
||||
}
|
||||
}
|
||||
|
|
@ -422,7 +422,8 @@ int do_build_sparse_by_streamer(IndexStreamer::Pointer &streamer,
|
|||
}
|
||||
|
||||
int build_sparse_by_streamer(IndexStreamer::Pointer &streamer,
|
||||
YAML::Node &config_common) {
|
||||
YAML::Node &config_common,
|
||||
const IndexConverter::Pointer &converter) {
|
||||
if (!config_common["IndexPath"]) {
|
||||
LOG_ERROR("Miss params IndexPath for Streamer");
|
||||
return IndexError_InvalidArgument;
|
||||
|
|
@ -451,6 +452,15 @@ int build_sparse_by_streamer(IndexStreamer::Pointer &streamer,
|
|||
return IndexError_Runtime;
|
||||
}
|
||||
|
||||
// Dump converter state (e.g. rotator) to storage for streaming build
|
||||
if (converter) {
|
||||
ret = converter->dump_to_storage(storage);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("Failed to dump converter to storage, ret=%d", ret);
|
||||
return ret;
|
||||
}
|
||||
}
|
||||
|
||||
size_t thread_count = config_common["ThreadCount"]
|
||||
? config_common["ThreadCount"].as<uint64_t>()
|
||||
: std::thread::hardware_concurrency();
|
||||
|
|
@ -464,7 +474,8 @@ int build_sparse_by_streamer(IndexStreamer::Pointer &streamer,
|
|||
}
|
||||
|
||||
int do_build_by_streamer(IndexStreamer::Pointer &streamer,
|
||||
uint32_t thread_count, RetrievalMode retrieval_mode) {
|
||||
uint32_t thread_count, RetrievalMode retrieval_mode,
|
||||
const IndexStorage::Pointer &storage = nullptr) {
|
||||
int ret;
|
||||
ailego::ThreadPool pool(thread_count, false);
|
||||
std::atomic<size_t> finished{0};
|
||||
|
|
@ -486,6 +497,14 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer,
|
|||
return IndexError_NoExist;
|
||||
}
|
||||
reformer->init(meta.reformer_params());
|
||||
// Load reformer state from storage (e.g. rotator for IntegerStreaming)
|
||||
if (storage) {
|
||||
ret = reformer->load(storage);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("Failed to load reformer from storage, ret=%d", ret);
|
||||
return ret;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -593,7 +612,8 @@ int do_build_by_streamer(IndexStreamer::Pointer &streamer,
|
|||
}
|
||||
|
||||
int build_by_streamer(IndexStreamer::Pointer &streamer,
|
||||
YAML::Node &config_common) {
|
||||
YAML::Node &config_common,
|
||||
const IndexConverter::Pointer &converter) {
|
||||
if (!config_common["IndexPath"]) {
|
||||
LOG_ERROR("Miss params IndexPath for Streamer");
|
||||
return IndexError_InvalidArgument;
|
||||
|
|
@ -624,6 +644,15 @@ int build_by_streamer(IndexStreamer::Pointer &streamer,
|
|||
return IndexError_Runtime;
|
||||
}
|
||||
|
||||
// Dump converter state (e.g. rotator) to storage for streaming build
|
||||
if (converter) {
|
||||
ret = converter->dump_to_storage(storage);
|
||||
if (ret != 0) {
|
||||
LOG_ERROR("Failed to dump converter to storage, ret=%d", ret);
|
||||
return ret;
|
||||
}
|
||||
}
|
||||
|
||||
size_t thread_count = config_common["ThreadCount"]
|
||||
? config_common["ThreadCount"].as<uint64_t>()
|
||||
: std::thread::hardware_concurrency();
|
||||
|
|
@ -639,14 +668,15 @@ int build_by_streamer(IndexStreamer::Pointer &streamer,
|
|||
|
||||
LOG_DEBUG("thread count: %zu, retrieval mode: %s", thread_count,
|
||||
retrieval_mode == 1 ? "Dense" : "Sparse");
|
||||
do_build_by_streamer(streamer, thread_count, retrieval_mode);
|
||||
do_build_by_streamer(streamer, thread_count, retrieval_mode, storage);
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
IndexSparseHolder::Pointer convert_sparse_holder(
|
||||
const std::string &name, const ailego::Params ¶ms,
|
||||
VecsIndexSparseHolder::Pointer &in_holder, IndexMeta &index_meta) {
|
||||
VecsIndexSparseHolder::Pointer &in_holder, IndexMeta &index_meta,
|
||||
IndexConverter::Pointer *out_converter) {
|
||||
IndexSparseHolder::Pointer cast_holder =
|
||||
std::dynamic_pointer_cast<IndexSparseHolder>(in_holder);
|
||||
if (name.empty()) {
|
||||
|
|
@ -679,13 +709,17 @@ IndexSparseHolder::Pointer convert_sparse_holder(
|
|||
|
||||
index_meta = converter->meta();
|
||||
|
||||
if (out_converter) {
|
||||
*out_converter = converter;
|
||||
}
|
||||
return converter->sparse_result();
|
||||
}
|
||||
|
||||
IndexHolder::Pointer convert_holder(const std::string &name,
|
||||
const ailego::Params ¶ms,
|
||||
VecsIndexHolder::Pointer &in_holder,
|
||||
IndexMeta &index_meta) {
|
||||
IndexMeta &index_meta,
|
||||
IndexConverter::Pointer *out_converter) {
|
||||
IndexHolder::Pointer cast_holder =
|
||||
std::dynamic_pointer_cast<IndexHolder>(in_holder);
|
||||
if (name.empty()) {
|
||||
|
|
@ -718,6 +752,9 @@ IndexHolder::Pointer convert_holder(const std::string &name,
|
|||
|
||||
index_meta = converter->meta();
|
||||
|
||||
if (out_converter) {
|
||||
*out_converter = converter;
|
||||
}
|
||||
return converter->result();
|
||||
}
|
||||
|
||||
|
|
@ -782,8 +819,9 @@ int do_build_sparse(YAML::Node &config_root, YAML::Node &config_common) {
|
|||
}
|
||||
cout << "Created builder " << builder_class << endl;
|
||||
|
||||
IndexConverter::Pointer build_converter;
|
||||
IndexSparseHolder::Pointer cv_build_holder = convert_sparse_holder(
|
||||
converter_name, converter_params, build_holder, meta);
|
||||
converter_name, converter_params, build_holder, meta, &build_converter);
|
||||
if (!cv_build_holder) {
|
||||
LOG_ERROR("Convert holder failed.");
|
||||
return -1;
|
||||
|
|
@ -819,7 +857,7 @@ int do_build_sparse(YAML::Node &config_root, YAML::Node &config_common) {
|
|||
}
|
||||
|
||||
IndexSparseHolder::Pointer cv_train_holder = convert_sparse_holder(
|
||||
converter_name, converter_params, train_holder, meta);
|
||||
converter_name, converter_params, train_holder, meta, nullptr);
|
||||
if (!cv_train_holder) {
|
||||
LOG_ERROR("Convert train holder failed.");
|
||||
return -1;
|
||||
|
|
@ -846,7 +884,7 @@ int do_build_sparse(YAML::Node &config_root, YAML::Node &config_common) {
|
|||
if (builder != nullptr) {
|
||||
ret = builder->build(std::move(cv_build_holder));
|
||||
} else {
|
||||
ret = build_sparse_by_streamer(streamer, config_common);
|
||||
ret = build_sparse_by_streamer(streamer, config_common, build_converter);
|
||||
}
|
||||
size_t build_time = timer.milli_seconds();
|
||||
if (ret < 0) {
|
||||
|
|
@ -987,8 +1025,9 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) {
|
|||
cout << "Created builder " << builder_class << endl;
|
||||
|
||||
|
||||
IndexHolder::Pointer cv_build_holder =
|
||||
convert_holder(converter_name, converter_params, build_holder, meta);
|
||||
IndexConverter::Pointer build_converter;
|
||||
IndexHolder::Pointer cv_build_holder = convert_holder(
|
||||
converter_name, converter_params, build_holder, meta, &build_converter);
|
||||
if (!cv_build_holder) {
|
||||
LOG_ERROR("Convert holder failed.");
|
||||
return -1;
|
||||
|
|
@ -1079,8 +1118,8 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) {
|
|||
|
||||
// support fp16 convert
|
||||
|
||||
IndexHolder::Pointer cv_train_holder =
|
||||
convert_holder(converter_name, converter_params, train_holder, meta);
|
||||
IndexHolder::Pointer cv_train_holder = convert_holder(
|
||||
converter_name, converter_params, train_holder, meta, nullptr);
|
||||
if (!cv_train_holder) {
|
||||
LOG_ERROR("Convert train holder failed.");
|
||||
return -1;
|
||||
|
|
@ -1136,8 +1175,8 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) {
|
|||
if (!metric_name.empty()) {
|
||||
train_holder->set_metric(metric_name, metric_params);
|
||||
}
|
||||
IndexHolder::Pointer cv_train_holder =
|
||||
convert_holder(converter_name, converter_params, train_holder, meta);
|
||||
IndexHolder::Pointer cv_train_holder = convert_holder(
|
||||
converter_name, converter_params, train_holder, meta, nullptr);
|
||||
if (!cv_train_holder) {
|
||||
LOG_ERROR("Convert train holder failed.");
|
||||
return -1;
|
||||
|
|
@ -1177,7 +1216,7 @@ int do_build(YAML::Node &config_root, YAML::Node &config_common) {
|
|||
retrieval_mode = "dense";
|
||||
}
|
||||
|
||||
ret = build_by_streamer(streamer, config_common);
|
||||
ret = build_by_streamer(streamer, config_common, build_converter);
|
||||
}
|
||||
size_t build_time = timer.milli_seconds();
|
||||
if (ret < 0) {
|
||||
|
|
|
|||
Loading…
Reference in New Issue