From a4f04d3d26988d28ecfe3049644470b55a010723 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Wed, 21 Oct 2026 11:38:10 +0800 Subject: [PATCH] Align NumPy inputs before native processing --- docs/docs/index/usage.md | 3 ++ python_bindings/bindings_common.hpp | 14 ++++++ tests/python/test_array_alignment.py | 65 ++++++++++++++++++++++++++++ 3 files changed, 82 insertions(+) create mode 100644 tests/python/test_array_alignment.py diff --git a/docs/docs/index/usage.md b/docs/docs/index/usage.md index c334ef64..e4ec2c0e 100644 --- a/docs/docs/index/usage.md +++ b/docs/docs/index/usage.md @@ -11,6 +11,9 @@ SymphonyQG. Start with the [Python quick start](../quick_start.md) for a complet to a multiple of 32. Queries and centroids must have the same original dimension as the indexed data. +Arrays whose data pointers are not aligned for float32 are copied into aligned +storage before native processing; inputs remain unchanged. + These limits apply to the Python indexes. The C++ [stable core API](../rabitq/core_api.md) accepts working dimensions from 2, with multiples of 32 required for `SplitPacked`. C++ [MatrixRotator](../rabitq/rotator.md) has separate dimension and storage requirements; diff --git a/python_bindings/bindings_common.hpp b/python_bindings/bindings_common.hpp index b0525565..de30e833 100644 --- a/python_bindings/bindings_common.hpp +++ b/python_bindings/bindings_common.hpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -98,6 +99,19 @@ inline py::array_t ensure_2d_array if (array.ndim() != 2) { throw std::invalid_argument(std::string(name) + " must be a 2D NumPy array"); } + // C-contiguous NumPy views may start at an unaligned byte offset. Native + // scalar accesses require alignof(T), independently of SIMD load alignment. + if (reinterpret_cast(array.data()) % alignof(T) != 0) { + py::array_t aligned( + {array.shape(0), array.shape(1)} + ); + if (array.size() != 0) { + std::memcpy( + aligned.mutable_data(), array.data(), static_cast(array.nbytes()) + ); + } + return aligned; + } return array; } diff --git a/tests/python/test_array_alignment.py b/tests/python/test_array_alignment.py new file mode 100644 index 00000000..7b31caf5 --- /dev/null +++ b/tests/python/test_array_alignment.py @@ -0,0 +1,65 @@ +"""Contiguous arrays may still require an aligned copy at the native boundary.""" + +import numpy as np +import pytest +from rabitqlib import HnswIndex, IvfIndex, QGKMeans, RaBitQKMeans, SymqgIndex + + +def unaligned_copy(values, offset): + result = np.ndarray( + values.shape, + dtype=values.dtype, + buffer=bytearray(values.nbytes + offset), + offset=offset, + ) + result[:] = values + result.flags.writeable = False + assert result.flags.c_contiguous + assert not result.flags.aligned + return result + + +@pytest.mark.parametrize("model_type", [RaBitQKMeans, QGKMeans]) +@pytest.mark.parametrize("offset", [1, 2, 3]) +def test_unaligned_training_matches_aligned_input(model_type, offset): + data = np.random.default_rng(42).standard_normal((97, 65), dtype=np.float32) + unaligned = unaligned_copy(data, offset) + expected = model_type(65, 33, niter=2, num_threads=1, seed=42) + actual = model_type(65, 33, niter=2, num_threads=1, seed=42) + assert actual.train(unaligned) == expected.train(data) + for name in ("centroids", "assignments", "distances", "obj"): + np.testing.assert_array_equal(getattr(actual, name), getattr(expected, name)) + np.testing.assert_array_equal(unaligned, data) + + +@pytest.mark.parametrize("kind", ["ivf", "hnsw", "qg"]) +@pytest.mark.parametrize("offset", [1, 2, 3]) +def test_indexes_accept_unaligned_build_and_query_arrays(kind, offset): + rng = np.random.default_rng(43) + data = rng.standard_normal((65, 65), dtype=np.float32) + centroids = data[:2].copy() + labels = (np.arange(len(data)) % 2).astype(np.uint32) + queries = rng.standard_normal((3, 65), dtype=np.float32) + unaligned_data = unaligned_copy(data, offset) + unaligned_centroids = unaligned_copy(centroids, offset) + unaligned_queries = unaligned_copy(queries, offset) + if kind == "ivf": + index = IvfIndex(65, len(data), 2, 4) + parameters = {"nprobe": 2} + elif kind == "hnsw": + index = HnswIndex(65, len(data), M=8, ef_construction=32, nbits=4) + parameters = {"ef": 65} + else: + index = SymqgIndex(65, 32) + parameters = {"ef": 65} + if kind == "qg": + index.build(unaligned_data, ef_construction=32, num_threads=1) + else: + index.build(unaligned_data, unaligned_centroids, labels, num_threads=1) + expected = index.search(queries, k=5, num_threads=1, **parameters) + actual = index.search(unaligned_queries, k=5, num_threads=1, **parameters) + for result, reference in zip(actual, expected, strict=True): + np.testing.assert_array_equal(result, reference) + np.testing.assert_array_equal(unaligned_data, data) + np.testing.assert_array_equal(unaligned_centroids, centroids) + np.testing.assert_array_equal(unaligned_queries, queries)