diff --git a/c/CMakeLists.txt b/c/CMakeLists.txt index 1322c9ed25..1d9614225d 100644 --- a/c/CMakeLists.txt +++ b/c/CMakeLists.txt @@ -104,6 +104,7 @@ add_library( src/preprocessing/pca.cpp src/preprocessing/quantize/binary.cpp src/preprocessing/quantize/pq.cpp + src/preprocessing/quantize/rabitq.cpp src/preprocessing/quantize/scalar.cpp src/distance/pairwise_distance.cpp ) diff --git a/c/include/cuvs/core/all.h b/c/include/cuvs/core/all.h index 91958177c4..1c4ac48c40 100644 --- a/c/include/cuvs/core/all.h +++ b/c/include/cuvs/core/all.h @@ -44,4 +44,5 @@ #include #include #include +#include #include diff --git a/c/include/cuvs/preprocessing/quantize/rabitq.h b/c/include/cuvs/preprocessing/quantize/rabitq.h new file mode 100644 index 0000000000..9a5c4f8659 --- /dev/null +++ b/c/include/cuvs/preprocessing/quantize/rabitq.h @@ -0,0 +1,479 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @defgroup preprocessing_c_rabitq C API for RaBitQ Quantizer + * @{ + */ + +/** + * @brief Random-rotation implementation used before quantization. + * + * Both kinds produce an orthogonal transform of the padded working dimension; + * they differ only in cost and in the amount of state they carry. + */ +enum cuvsRabitqRotatorKind { + /** Dense `padded_dim x padded_dim` random orthogonal matrix (cuBLAS GEMM), O(n * dim^2). */ + CUVS_RABITQ_ROTATOR_KIND_MATMUL = 0, + /** Fast Hadamard Transform combined with Kac's walk, O(n * dim * log(dim)). */ + CUVS_RABITQ_ROTATOR_KIND_FHT_KAC = 1, +}; + +/** + * @brief Rule used to derive the per-vector reconstruction scale `delta`. + * + * Only relevant for cuvsRabitqTransform / cuvsRabitqTransformResiduals, which + * emit `(codes, delta, vl)`; the full-factor entry points emit + * `(f_add, f_rescale, f_error)` instead and ignore this setting. + */ +enum cuvsRabitqDeltaKind { + /** `delta = (|residual| / |code|) * cos(residual, code)` - minimizes reconstruction error. */ + CUVS_RABITQ_DELTA_KIND_RECONSTRUCTION = 0, + /** `delta = (|residual| / |code|) / cos(residual, code)` - unbiased inner-product estimate. */ + CUVS_RABITQ_DELTA_KIND_UNBIASED = 1, + /** `delta = |residual| / |code|` - plain norm ratio, no angular correction. */ + CUVS_RABITQ_DELTA_KIND_PLAIN = 2, +}; + +/** + * @brief RaBitQ quantizer parameters. + * + * These parameters fully describe both the rotator and the quantizer: RaBitQ is + * data-oblivious, so none of them are learned from a dataset. This is why the + * factories below take a `dim` instead of a training dataset. + */ +struct cuvsRabitqQuantizerParams { + /** + * Number of *extended* bits per dimension. The total code width is + * `ex_bits + 1` bits (one sign bit plus `ex_bits` magnitude bits), so + * `ex_bits = 0` is plain 1-bit RaBitQ and `ex_bits = 3` is the common + * 4-bit configuration. + * + * Possible value range: [0, 8]. + */ + uint32_t ex_bits; + /** Random-rotation implementation. */ + enum cuvsRabitqRotatorKind rotator; + /** + * Seed for the host RNG that draws the rotation state and the random probe + * vectors used to estimate the constant scaling factor. + */ + uint64_t seed; + /** Rule used to derive the per-vector `delta` / `vl` pair. */ + enum cuvsRabitqDeltaKind delta_mode; + /** + * When true, every vector is quantized with a single pre-estimated scaling + * factor (see cuvsRabitqQuantizerGetConstScalingFactor) instead of running + * the per-vector rescale search. This is substantially faster and is the + * recommended setting; disable it to trade build throughput for slightly + * tighter codes. + */ + bool use_fast; + /** Number of samples in the coarse stage of the per-vector rescale search. */ + uint32_t coarse_samples; + /** Number of samples in the refinement stage of the per-vector rescale search. */ + uint32_t fine_samples; +}; + +typedef struct cuvsRabitqQuantizerParams* cuvsRabitqQuantizerParams_t; + +/** + * @brief Allocate RaBitQ quantizer params, and populate with default values + * + * @param[in] params cuvsRabitqQuantizerParams_t to allocate + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqQuantizerParamsCreate(cuvsRabitqQuantizerParams_t* params); + +/** + * @brief De-allocate RaBitQ quantizer params + * + * @param[in] params + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqQuantizerParamsDestroy(cuvsRabitqQuantizerParams_t params); + +/** + * @brief Opaque handle to the random rotation applied before quantization. + * + * The rotator carries no learned information - its state is drawn from + * cuvsRabitqQuantizerParams::seed - but it *is* part of the code + * representation: codes produced with one rotator are only comparable to + * queries transformed by the very same rotator. There is deliberately no + * serialization API, so the caller owns this handle and is responsible for + * keeping it alongside the codes (either by storing the seed and re-running + * cuvsRabitqRotatorMake, or by copying the buffers returned by + * cuvsRabitqRotatorGetState into its own index format). + * + * A rotator is an independent peer of the quantizer: neither contains the + * other, and a single rotator may be shared by several quantizers. + */ +typedef struct { + uintptr_t addr; +} cuvsRabitqRotator; + +typedef cuvsRabitqRotator* cuvsRabitqRotator_t; + +/** + * @brief Allocate a RaBitQ rotator handle + * + * The handle starts out empty; populate it with cuvsRabitqRotatorMake or + * cuvsRabitqPipelineMake before use. + * + * @param[in] rotator cuvsRabitqRotator_t to allocate + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqRotatorCreate(cuvsRabitqRotator_t* rotator); + +/** + * @brief De-allocate a RaBitQ rotator, freeing its device buffers + * + * Invalidates every tensor previously filled in by cuvsRabitqRotatorGetState. + * + * @param[in] rotator + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqRotatorDestroy(cuvsRabitqRotator_t rotator); + +/** + * @brief Opaque handle to the (data-independent) configuration used to quantize + * rotated residuals. + * + * RaBitQ needs no training data, so this is plain metadata plus the single + * scalar estimated at init time. Like the rotator it is owned by the caller and + * has no serialization API - it is cheap to rebuild from the params and `dim`. + */ +typedef struct { + uintptr_t addr; +} cuvsRabitqQuantizer; + +typedef cuvsRabitqQuantizer* cuvsRabitqQuantizer_t; + +/** + * @brief Allocate a RaBitQ quantizer handle + * + * The handle starts out empty; populate it with cuvsRabitqQuantizerMake or + * cuvsRabitqPipelineMake before use. + * + * @param[in] quantizer cuvsRabitqQuantizer_t to allocate + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqQuantizerCreate(cuvsRabitqQuantizer_t* quantizer); + +/** + * @brief De-allocate a RaBitQ quantizer + * + * @param[in] quantizer + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqQuantizerDestroy(cuvsRabitqQuantizer_t quantizer); + +/** + * @brief Initializes a RaBitQ rotator. + * + * No training data is required: the rotation is drawn from + * cuvsRabitqQuantizerParams::seed alone, which is why this takes a `dim` + * rather than a dataset. The rotator owns its device buffers; keep it alive for + * as long as the codes it produced are in use. + * + * @param[in] res raft resource + * @param[in] params configure the rotator, e.g. kind and seed + * @param[in] dim dimensionality of the un-rotated input vectors + * @param[out] rotator an allocated rotator handle to initialize + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqRotatorMake(cuvsResources_t res, + cuvsRabitqQuantizerParams_t params, + int64_t dim, + cuvsRabitqRotator_t rotator); + +/** + * @brief Initializes a RaBitQ quantizer. + * + * No training data is required, which is why this takes a `dim` rather than a + * dataset. When cuvsRabitqQuantizerParams::use_fast is set, this estimates the + * constant scaling factor on the GPU, which is the only non-trivial work + * performed here. + * + * @param[in] res raft resource + * @param[in] params configure the quantizer, e.g. ex_bits + * @param[in] dim dimensionality of the un-rotated input vectors + * @param[out] quantizer an allocated quantizer handle to initialize + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqQuantizerMake(cuvsResources_t res, + cuvsRabitqQuantizerParams_t params, + int64_t dim, + cuvsRabitqQuantizer_t quantizer); + +/** + * @brief Initializes a rotator and a quantizer from the same parameters. + * + * Equivalent to calling cuvsRabitqRotatorMake and cuvsRabitqQuantizerMake with + * the same arguments. No training data is required. + * + * There is deliberately no opaque `cuvsRabitqPipeline` handle: the C++ + * `pipeline` struct is only a composition of the two independent peers - it + * adds no state and no behaviour of its own - so two out-parameters are the + * honest C spelling of it. + * + * @param[in] res raft resource + * @param[in] params configure both components + * @param[in] dim dimensionality of the un-rotated input vectors + * @param[out] rotator an allocated rotator handle to initialize + * @param[out] quantizer an allocated quantizer handle to initialize + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqPipelineMake(cuvsResources_t res, + cuvsRabitqQuantizerParams_t params, + int64_t dim, + cuvsRabitqRotator_t rotator, + cuvsRabitqQuantizer_t quantizer); + +/** + * @brief Applies the random rotation to already-padded rows. + * + * Both `in` and `out` must be `n_rows x padded_dim` row-major float32 device + * matrices (see cuvsRabitqRotatorGetPaddedDim); this entry point does not pad. + * In-place operation (`in` and `out` pointing at the same buffer) is supported + * by ::CUVS_RABITQ_ROTATOR_KIND_FHT_KAC only. + * + * Use this to rotate queries or externally computed centroids so that they live + * in the same rotated space as the codes. + * + * @param[in] res raft resource + * @param[in] rotator a RaBitQ rotator + * @param[in] in a row-major `n_rows x padded_dim` float32 device matrix + * @param[out] out a row-major `n_rows x padded_dim` float32 device matrix + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqRotate(cuvsResources_t res, + cuvsRabitqRotator_t rotator, + DLManagedTensor* in, + DLManagedTensor* out); + +/** + * @brief Quantizes a raw dataset, emitting per-vector `(delta, vl)` factors. + * + * Performs the whole chain: zero-pad `dim -> padded_dim`, rotate, optionally + * subtract the rotated centroid, then quantize the residuals. One code element + * is written per padded dimension; each holds an unsigned + * `cuvsRabitqQuantizerParams::ex_bits + 1` bit value, so `codes` must be + * `n_rows x padded_dim`. + * + * A temporary `n_rows x padded_dim` float buffer holding the rotated residuals + * is allocated from the workspace resource, so large datasets should be + * transformed in batches (a pool workspace resource is recommended). + * + * @param[in] res raft resource + * @param[in] quantizer a RaBitQ quantizer + * @param[in] rotator the rotator whose padded dim matches the quantizer + * @param[in] dataset a row-major `n_rows x dim` float32 device matrix + * @param[in] centroid optional un-rotated `dim` float32 device vector; when non-NULL, codes are + * computed on the residuals w.r.t. this centroid. Pass NULL to skip centering. + * @param[out] codes a row-major `n_rows x padded_dim` device matrix of uint8 or uint16 codes + * @param[out] delta per-vector scale, an `n_rows` float32 device vector + * @param[out] vl per-vector offset (`delta * -(2^ex_bits - 0.5)`), an `n_rows` float32 device + * vector + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqTransform(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + cuvsRabitqRotator_t rotator, + DLManagedTensor* dataset, + DLManagedTensor* centroid, + DLManagedTensor* codes, + DLManagedTensor* delta, + DLManagedTensor* vl); + +/** + * @brief Quantizes a raw dataset, emitting the full `(f_add, f_rescale, f_error)` factor triplet. + * + * Same chain as cuvsRabitqTransform, but the per-vector factors are the ones + * used for approximate-distance estimation during search: + * `factors[i] = { f_add, f_rescale, f_error }`. + * + * A temporary `n_rows x padded_dim` float buffer holding the rotated residuals + * is allocated from the workspace resource, so large datasets should be + * transformed in batches (a pool workspace resource is recommended). + * + * The rotated centroid is an internal temporary as well. A search + * implementation that needs it can reproduce it exactly by zero-padding + * `centroid` to `padded_dim` and calling cuvsRabitqRotate with the same + * rotator. + * + * @param[in] res raft resource + * @param[in] quantizer a RaBitQ quantizer + * @param[in] rotator the rotator whose padded dim matches the quantizer + * @param[in] dataset a row-major `n_rows x dim` float32 device matrix + * @param[in] centroid optional un-rotated `dim` float32 device vector; treated as zero when NULL + * @param[out] codes a row-major `n_rows x padded_dim` device matrix of uint8 or uint16 codes + * @param[out] factors a row-major `n_rows x 3` float32 device matrix + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqTransformFull(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + cuvsRabitqRotator_t rotator, + DLManagedTensor* dataset, + DLManagedTensor* centroid, + DLManagedTensor* codes, + DLManagedTensor* factors); + +/** + * @brief Quantizes pre-rotated residuals, emitting per-vector `(delta, vl)` factors. + * + * The rotator is not needed here: `residuals` must already be rotated, padded + * to the quantizer's padded dim and centered (centroid subtracted). Use this + * when the residuals are produced elsewhere, e.g. by an IVF build that has + * already rotated its data, or by cuvsRabitqRotate plus a subtraction of your + * own. + * + * @param[in] res raft resource + * @param[in] quantizer a RaBitQ quantizer + * @param[in] residuals a row-major `n_rows x padded_dim` float32 device matrix, rotated and + * centered + * @param[out] codes a row-major `n_rows x padded_dim` device matrix of uint8 or uint16 codes + * @param[out] delta per-vector scale, an `n_rows` float32 device vector + * @param[out] vl per-vector offset (`delta * -(2^ex_bits - 0.5)`), an `n_rows` float32 device + * vector + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqTransformResiduals(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + DLManagedTensor* residuals, + DLManagedTensor* codes, + DLManagedTensor* delta, + DLManagedTensor* vl); + +/** + * @brief Quantizes pre-rotated residuals, emitting the full factor triplet. + * + * As above, `residuals` must already be rotated, padded to the quantizer's + * padded dim and centered. The full-factor formula also reads the centroid, + * which must be supplied *in the rotated, padded space* (run cuvsRabitqRotate + * on it with the same rotator that produced the residuals); pass NULL to treat + * it as zero. + * + * @param[in] res raft resource + * @param[in] quantizer a RaBitQ quantizer + * @param[in] residuals a row-major `n_rows x padded_dim` float32 device matrix, rotated and + * centered + * @param[in] rotated_centroid optional rotated, padded centroid (a `padded_dim` float32 device + * vector); treated as zero when NULL + * @param[out] codes a row-major `n_rows x padded_dim` device matrix of uint8 or uint16 codes + * @param[out] factors a row-major `n_rows x 3` float32 device matrix + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqTransformResidualsFull(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + DLManagedTensor* residuals, + DLManagedTensor* rotated_centroid, + DLManagedTensor* codes, + DLManagedTensor* factors); + +/** + * @brief Get which random-rotation implementation a rotator uses. + * + * Together with cuvsRabitqRotatorGetPaddedDim and cuvsRabitqRotatorGetState + * this is the rotator's full identity. Because RaBitQ has no serialization API, + * reading these three back out is how a C caller persists the rotator next to + * the codes it produced. + * + * @param[in] rotator a RaBitQ rotator + * @param[out] kind the rotation implementation + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqRotatorGetKind(cuvsRabitqRotator_t rotator, + enum cuvsRabitqRotatorKind* kind); + +/** + * @brief Get the working dimension of a rotator: `dim` rounded up to a multiple of 32. + * + * Both transforms operate purely on the padded width, so the original unpadded + * `dim` is not part of the rotator. + * + * @param[in] rotator a RaBitQ rotator + * @param[out] padded_dim the working dimension + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqRotatorGetPaddedDim(cuvsRabitqRotator_t rotator, + int64_t* padded_dim); + +/** + * @brief Get the device-side state of a rotator. + * + * The shape and dtype depend on cuvsRabitqRotatorGetKind: + * - ::CUVS_RABITQ_ROTATOR_KIND_MATMUL fills in the `padded_dim x padded_dim` + * row-major float32 rotation matrix. + * - ::CUVS_RABITQ_ROTATOR_KIND_FHT_KAC fills in the packed sign bits of the four + * Kac's-walk passes, a `4 * padded_dim / 8` element uint8 vector. + * + * The returned DLManagedTensor is a **non-owning view**: the device buffer stays + * owned by `rotator` and becomes dangling once cuvsRabitqRotatorDestroy is + * called on it. Copy the contents out if you need them to outlive the rotator. + * Only the tensor's own shape/stride metadata is allocated here, and it is + * released by invoking the tensor's `deleter`. + * + * @param[in] rotator a RaBitQ rotator + * @param[out] state a view of the rotator's device state + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqRotatorGetState(cuvsRabitqRotator_t rotator, + DLManagedTensor* state); + +/** + * @brief Get the dimensionality of the un-rotated input vectors a quantizer expects. + * + * @param[in] quantizer a RaBitQ quantizer + * @param[out] dim dimensionality of the un-rotated input vectors + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqQuantizerGetDim(cuvsRabitqQuantizer_t quantizer, int64_t* dim); + +/** + * @brief Get the working dimension of a quantizer: `dim` rounded up to a multiple of 32. + * + * This is the number of code elements emitted per vector, i.e. the second + * extent that the `codes` tensor must have. + * + * @param[in] quantizer a RaBitQ quantizer + * @param[out] padded_dim the working dimension + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqQuantizerGetPaddedDim(cuvsRabitqQuantizer_t quantizer, + int64_t* padded_dim); + +/** + * @brief Get the scaling factor shared by all vectors on the fast path. + * + * Estimated from random Gaussian probe vectors when the quantizer was made. + * Left at 1.0 (and unused) when cuvsRabitqQuantizerParams::use_fast is false. + * + * @param[in] quantizer a RaBitQ quantizer + * @param[out] const_scaling_factor the shared scaling factor + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsRabitqQuantizerGetConstScalingFactor(cuvsRabitqQuantizer_t quantizer, + float* const_scaling_factor); + +/** + * @} + */ + +#ifdef __cplusplus +} +#endif diff --git a/c/src/preprocessing/quantize/rabitq.cpp b/c/src/preprocessing/quantize/rabitq.cpp new file mode 100644 index 0000000000..eca2e8db18 --- /dev/null +++ b/c/src/preprocessing/quantize/rabitq.cpp @@ -0,0 +1,489 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include +#include + +#include "../../core/exceptions.hpp" +#include "../../core/interop.hpp" + +namespace { + +namespace rabitq = cuvs::preprocessing::quantize::rabitq; + +/** Read-only `n_rows x dim` float input: `dataset`, `residuals`, `in`. */ +using in_matrix_mdspan_type = raft::device_matrix_view; +/** Read-only `dim` float vector: `centroid`, `rotated_centroid`. */ +using centroid_mdspan_type = raft::device_vector_view; +/** Writable float matrix output: `factors` (`n_rows x 3`) and `rotate`'s `out`. */ +using out_matrix_mdspan_type = raft::device_matrix_view; +/** Writable `n_rows` float vector output: `delta`, `vl`. */ +using out_vector_mdspan_type = raft::device_vector_view; + +template +using codes_mdspan_type = raft::device_matrix_view; + +/** Translate the C params struct into its C++ peer. */ +rabitq::params _convert_params(cuvsRabitqQuantizerParams_t params) +{ + if (params == nullptr) { RAFT_FAIL("params is not allocated"); } + + auto out = rabitq::params(); + out.ex_bits = params->ex_bits; + switch (params->rotator) { + case CUVS_RABITQ_ROTATOR_KIND_MATMUL: out.rotator = rabitq::rotator_kind::matmul; break; + case CUVS_RABITQ_ROTATOR_KIND_FHT_KAC: out.rotator = rabitq::rotator_kind::fht_kac; break; + default: RAFT_FAIL("Unsupported rotator kind: %d", static_cast(params->rotator)); + } + out.seed = params->seed; + switch (params->delta_mode) { + case CUVS_RABITQ_DELTA_KIND_RECONSTRUCTION: + out.delta_mode = rabitq::delta_kind::reconstruction; + break; + case CUVS_RABITQ_DELTA_KIND_UNBIASED: out.delta_mode = rabitq::delta_kind::unbiased; break; + case CUVS_RABITQ_DELTA_KIND_PLAIN: out.delta_mode = rabitq::delta_kind::plain; break; + default: RAFT_FAIL("Unsupported delta mode: %d", static_cast(params->delta_mode)); + } + out.use_fast = params->use_fast; + out.coarse_samples = params->coarse_samples; + out.fine_samples = params->fine_samples; + return out; +} + +rabitq::rotator* _rotator_ptr(cuvsRabitqRotator_t rotator) +{ + if (rotator == nullptr || rotator->addr == 0) { RAFT_FAIL("rotator is not initialized"); } + return reinterpret_cast(rotator->addr); +} + +rabitq::quantizer* _quantizer_ptr(cuvsRabitqQuantizer_t quantizer) +{ + if (quantizer == nullptr || quantizer->addr == 0) { RAFT_FAIL("quantizer is not initialized"); } + return reinterpret_cast(quantizer->addr); +} + +/** Every RaBitQ entry point is device-only; give a uniform diagnostic for host tensors. */ +void _expect_device(DLManagedTensor* tensor, char const* name) +{ + if (tensor == nullptr) { RAFT_FAIL("%s must not be NULL", name); } + if (!cuvs::core::is_dlpack_device_compatible(tensor->dl_tensor)) { + RAFT_FAIL("%s must be accessible on device memory", name); + } +} + +void _expect_float32(DLManagedTensor* tensor, char const* name) +{ + _expect_device(tensor, name); + auto dtype = tensor->dl_tensor.dtype; + if (dtype.code != kDLFloat || dtype.bits != 32) { + RAFT_FAIL("Unsupported %s DLtensor dtype: %d and bits: %d", name, dtype.code, dtype.bits); + } +} + +/** NULL maps to std::nullopt, which both transform() families treat as "no centroid". */ +std::optional _optional_centroid(DLManagedTensor* centroid_tensor, + char const* name) +{ + if (centroid_tensor == nullptr) { return std::nullopt; } + _expect_float32(centroid_tensor, name); + return cuvs::core::from_dlpack(centroid_tensor); +} + +template +void _transform(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + cuvsRabitqRotator_t rotator, + DLManagedTensor* dataset_tensor, + DLManagedTensor* centroid_tensor, + DLManagedTensor* codes_tensor, + DLManagedTensor* delta_tensor, + DLManagedTensor* vl_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto* q = _quantizer_ptr(quantizer); + auto* r = _rotator_ptr(rotator); + + _expect_float32(dataset_tensor, "dataset"); + _expect_float32(delta_tensor, "delta"); + _expect_float32(vl_tensor, "vl"); + _expect_device(codes_tensor, "codes"); + + rabitq::transform(*res_ptr, + *q, + *r, + cuvs::core::from_dlpack(dataset_tensor), + _optional_centroid(centroid_tensor, "centroid"), + cuvs::core::from_dlpack>(codes_tensor), + cuvs::core::from_dlpack(delta_tensor), + cuvs::core::from_dlpack(vl_tensor)); +} + +template +void _transform_full(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + cuvsRabitqRotator_t rotator, + DLManagedTensor* dataset_tensor, + DLManagedTensor* centroid_tensor, + DLManagedTensor* codes_tensor, + DLManagedTensor* factors_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto* q = _quantizer_ptr(quantizer); + auto* r = _rotator_ptr(rotator); + + _expect_float32(dataset_tensor, "dataset"); + _expect_float32(factors_tensor, "factors"); + _expect_device(codes_tensor, "codes"); + + rabitq::transform(*res_ptr, + *q, + *r, + cuvs::core::from_dlpack(dataset_tensor), + _optional_centroid(centroid_tensor, "centroid"), + cuvs::core::from_dlpack>(codes_tensor), + cuvs::core::from_dlpack(factors_tensor)); +} + +template +void _transform_residuals(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + DLManagedTensor* residuals_tensor, + DLManagedTensor* codes_tensor, + DLManagedTensor* delta_tensor, + DLManagedTensor* vl_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto* q = _quantizer_ptr(quantizer); + + _expect_float32(residuals_tensor, "residuals"); + _expect_float32(delta_tensor, "delta"); + _expect_float32(vl_tensor, "vl"); + _expect_device(codes_tensor, "codes"); + + rabitq::transform_residuals(*res_ptr, + *q, + cuvs::core::from_dlpack(residuals_tensor), + cuvs::core::from_dlpack>(codes_tensor), + cuvs::core::from_dlpack(delta_tensor), + cuvs::core::from_dlpack(vl_tensor)); +} + +template +void _transform_residuals_full(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + DLManagedTensor* residuals_tensor, + DLManagedTensor* rotated_centroid_tensor, + DLManagedTensor* codes_tensor, + DLManagedTensor* factors_tensor) +{ + auto res_ptr = reinterpret_cast(res); + auto* q = _quantizer_ptr(quantizer); + + _expect_float32(residuals_tensor, "residuals"); + _expect_float32(factors_tensor, "factors"); + _expect_device(codes_tensor, "codes"); + + rabitq::transform_residuals( + *res_ptr, + *q, + cuvs::core::from_dlpack(residuals_tensor), + _optional_centroid(rotated_centroid_tensor, "rotated_centroid"), + cuvs::core::from_dlpack>(codes_tensor), + cuvs::core::from_dlpack(factors_tensor)); +} + +} // namespace + +extern "C" cuvsError_t cuvsRabitqQuantizerParamsCreate(cuvsRabitqQuantizerParams_t* params) +{ + return cuvs::core::translate_exceptions([=] { + *params = new cuvsRabitqQuantizerParams{.ex_bits = 3, + .rotator = CUVS_RABITQ_ROTATOR_KIND_FHT_KAC, + .seed = 12345ULL, + .delta_mode = CUVS_RABITQ_DELTA_KIND_RECONSTRUCTION, + .use_fast = true, + .coarse_samples = 64, + .fine_samples = 64}; + }); +} + +extern "C" cuvsError_t cuvsRabitqQuantizerParamsDestroy(cuvsRabitqQuantizerParams_t params) +{ + return cuvs::core::translate_exceptions([=] { delete params; }); +} + +extern "C" cuvsError_t cuvsRabitqRotatorCreate(cuvsRabitqRotator_t* rotator) +{ + return cuvs::core::translate_exceptions([=] { *rotator = new cuvsRabitqRotator{}; }); +} + +extern "C" cuvsError_t cuvsRabitqRotatorDestroy(cuvsRabitqRotator_t rotator) +{ + return cuvs::core::translate_exceptions([=] { + if (rotator == nullptr) { return; } + delete reinterpret_cast(rotator->addr); + delete rotator; + }); +} + +extern "C" cuvsError_t cuvsRabitqQuantizerCreate(cuvsRabitqQuantizer_t* quantizer) +{ + return cuvs::core::translate_exceptions([=] { *quantizer = new cuvsRabitqQuantizer{}; }); +} + +extern "C" cuvsError_t cuvsRabitqQuantizerDestroy(cuvsRabitqQuantizer_t quantizer) +{ + return cuvs::core::translate_exceptions([=] { + if (quantizer == nullptr) { return; } + delete reinterpret_cast(quantizer->addr); + delete quantizer; + }); +} + +extern "C" cuvsError_t cuvsRabitqRotatorMake(cuvsResources_t res, + cuvsRabitqQuantizerParams_t params, + int64_t dim, + cuvsRabitqRotator_t rotator) +{ + return cuvs::core::translate_exceptions([=] { + if (rotator == nullptr) { RAFT_FAIL("rotator is not allocated"); } + auto res_ptr = reinterpret_cast(res); + auto quantizer_params = _convert_params(params); + auto* ret = new rabitq::rotator(rabitq::make_rotator(*res_ptr, quantizer_params, dim)); + delete reinterpret_cast(rotator->addr); + rotator->addr = reinterpret_cast(ret); + }); +} + +extern "C" cuvsError_t cuvsRabitqQuantizerMake(cuvsResources_t res, + cuvsRabitqQuantizerParams_t params, + int64_t dim, + cuvsRabitqQuantizer_t quantizer) +{ + return cuvs::core::translate_exceptions([=] { + if (quantizer == nullptr) { RAFT_FAIL("quantizer is not allocated"); } + auto res_ptr = reinterpret_cast(res); + auto quantizer_params = _convert_params(params); + auto* ret = new rabitq::quantizer(rabitq::make_quantizer(*res_ptr, quantizer_params, dim)); + delete reinterpret_cast(quantizer->addr); + quantizer->addr = reinterpret_cast(ret); + }); +} + +extern "C" cuvsError_t cuvsRabitqPipelineMake(cuvsResources_t res, + cuvsRabitqQuantizerParams_t params, + int64_t dim, + cuvsRabitqRotator_t rotator, + cuvsRabitqQuantizer_t quantizer) +{ + return cuvs::core::translate_exceptions([=] { + if (rotator == nullptr) { RAFT_FAIL("rotator is not allocated"); } + if (quantizer == nullptr) { RAFT_FAIL("quantizer is not allocated"); } + auto res_ptr = reinterpret_cast(res); + auto quantizer_params = _convert_params(params); + auto built = rabitq::make_pipeline(*res_ptr, quantizer_params, dim); + + auto* new_rotator = new rabitq::rotator(std::move(built.rotator)); + auto* new_quantizer = new rabitq::quantizer(built.quantizer); + + delete reinterpret_cast(rotator->addr); + delete reinterpret_cast(quantizer->addr); + rotator->addr = reinterpret_cast(new_rotator); + quantizer->addr = reinterpret_cast(new_quantizer); + }); +} + +extern "C" cuvsError_t cuvsRabitqRotate(cuvsResources_t res, + cuvsRabitqRotator_t rotator, + DLManagedTensor* in_tensor, + DLManagedTensor* out_tensor) +{ + return cuvs::core::translate_exceptions([=] { + auto res_ptr = reinterpret_cast(res); + auto* r = _rotator_ptr(rotator); + + _expect_float32(in_tensor, "in"); + _expect_float32(out_tensor, "out"); + + rabitq::rotate(*res_ptr, + *r, + cuvs::core::from_dlpack(in_tensor), + cuvs::core::from_dlpack(out_tensor)); + }); +} + +extern "C" cuvsError_t cuvsRabitqTransform(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + cuvsRabitqRotator_t rotator, + DLManagedTensor* dataset_tensor, + DLManagedTensor* centroid_tensor, + DLManagedTensor* codes_tensor, + DLManagedTensor* delta_tensor, + DLManagedTensor* vl_tensor) +{ + return cuvs::core::translate_exceptions([=] { + if (codes_tensor == nullptr) { RAFT_FAIL("codes must not be NULL"); } + auto codes = codes_tensor->dl_tensor; + if (codes.dtype.code == kDLUInt && codes.dtype.bits == 8) { + _transform(res, + quantizer, + rotator, + dataset_tensor, + centroid_tensor, + codes_tensor, + delta_tensor, + vl_tensor); + } else if (codes.dtype.code == kDLUInt && codes.dtype.bits == 16) { + _transform(res, + quantizer, + rotator, + dataset_tensor, + centroid_tensor, + codes_tensor, + delta_tensor, + vl_tensor); + } else { + RAFT_FAIL( + "Unsupported codes DLtensor dtype: %d and bits: %d", codes.dtype.code, codes.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsRabitqTransformFull(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + cuvsRabitqRotator_t rotator, + DLManagedTensor* dataset_tensor, + DLManagedTensor* centroid_tensor, + DLManagedTensor* codes_tensor, + DLManagedTensor* factors_tensor) +{ + return cuvs::core::translate_exceptions([=] { + if (codes_tensor == nullptr) { RAFT_FAIL("codes must not be NULL"); } + auto codes = codes_tensor->dl_tensor; + if (codes.dtype.code == kDLUInt && codes.dtype.bits == 8) { + _transform_full( + res, quantizer, rotator, dataset_tensor, centroid_tensor, codes_tensor, factors_tensor); + } else if (codes.dtype.code == kDLUInt && codes.dtype.bits == 16) { + _transform_full( + res, quantizer, rotator, dataset_tensor, centroid_tensor, codes_tensor, factors_tensor); + } else { + RAFT_FAIL( + "Unsupported codes DLtensor dtype: %d and bits: %d", codes.dtype.code, codes.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsRabitqTransformResiduals(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + DLManagedTensor* residuals_tensor, + DLManagedTensor* codes_tensor, + DLManagedTensor* delta_tensor, + DLManagedTensor* vl_tensor) +{ + return cuvs::core::translate_exceptions([=] { + if (codes_tensor == nullptr) { RAFT_FAIL("codes must not be NULL"); } + auto codes = codes_tensor->dl_tensor; + if (codes.dtype.code == kDLUInt && codes.dtype.bits == 8) { + _transform_residuals( + res, quantizer, residuals_tensor, codes_tensor, delta_tensor, vl_tensor); + } else if (codes.dtype.code == kDLUInt && codes.dtype.bits == 16) { + _transform_residuals( + res, quantizer, residuals_tensor, codes_tensor, delta_tensor, vl_tensor); + } else { + RAFT_FAIL( + "Unsupported codes DLtensor dtype: %d and bits: %d", codes.dtype.code, codes.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsRabitqTransformResidualsFull(cuvsResources_t res, + cuvsRabitqQuantizer_t quantizer, + DLManagedTensor* residuals_tensor, + DLManagedTensor* rotated_centroid_tensor, + DLManagedTensor* codes_tensor, + DLManagedTensor* factors_tensor) +{ + return cuvs::core::translate_exceptions([=] { + if (codes_tensor == nullptr) { RAFT_FAIL("codes must not be NULL"); } + auto codes = codes_tensor->dl_tensor; + if (codes.dtype.code == kDLUInt && codes.dtype.bits == 8) { + _transform_residuals_full(res, + quantizer, + residuals_tensor, + rotated_centroid_tensor, + codes_tensor, + factors_tensor); + } else if (codes.dtype.code == kDLUInt && codes.dtype.bits == 16) { + _transform_residuals_full(res, + quantizer, + residuals_tensor, + rotated_centroid_tensor, + codes_tensor, + factors_tensor); + } else { + RAFT_FAIL( + "Unsupported codes DLtensor dtype: %d and bits: %d", codes.dtype.code, codes.dtype.bits); + } + }); +} + +extern "C" cuvsError_t cuvsRabitqRotatorGetKind(cuvsRabitqRotator_t rotator, + enum cuvsRabitqRotatorKind* kind) +{ + return cuvs::core::translate_exceptions([=] { + auto* r = _rotator_ptr(rotator); + switch (r->kind) { + case rabitq::rotator_kind::matmul: *kind = CUVS_RABITQ_ROTATOR_KIND_MATMUL; break; + case rabitq::rotator_kind::fht_kac: *kind = CUVS_RABITQ_ROTATOR_KIND_FHT_KAC; break; + default: RAFT_FAIL("Unsupported rotator kind: %d", static_cast(r->kind)); + } + }); +} + +extern "C" cuvsError_t cuvsRabitqRotatorGetPaddedDim(cuvsRabitqRotator_t rotator, + int64_t* padded_dim) +{ + return cuvs::core::translate_exceptions([=] { *padded_dim = _rotator_ptr(rotator)->padded_dim; }); +} + +extern "C" cuvsError_t cuvsRabitqRotatorGetState(cuvsRabitqRotator_t rotator, + DLManagedTensor* state) +{ + return cuvs::core::translate_exceptions([=] { + auto* r = _rotator_ptr(rotator); + switch (r->kind) { + case rabitq::rotator_kind::matmul: + cuvs::core::to_dlpack(r->rotation_matrix.view(), state); + break; + case rabitq::rotator_kind::fht_kac: cuvs::core::to_dlpack(r->flip_bits.view(), state); break; + default: RAFT_FAIL("Unsupported rotator kind: %d", static_cast(r->kind)); + } + }); +} + +extern "C" cuvsError_t cuvsRabitqQuantizerGetDim(cuvsRabitqQuantizer_t quantizer, int64_t* dim) +{ + return cuvs::core::translate_exceptions([=] { *dim = _quantizer_ptr(quantizer)->dim; }); +} + +extern "C" cuvsError_t cuvsRabitqQuantizerGetPaddedDim(cuvsRabitqQuantizer_t quantizer, + int64_t* padded_dim) +{ + return cuvs::core::translate_exceptions( + [=] { *padded_dim = _quantizer_ptr(quantizer)->padded_dim; }); +} + +extern "C" cuvsError_t cuvsRabitqQuantizerGetConstScalingFactor(cuvsRabitqQuantizer_t quantizer, + float* const_scaling_factor) +{ + return cuvs::core::translate_exceptions( + [=] { *const_scaling_factor = _quantizer_ptr(quantizer)->const_scaling_factor; }); +} diff --git a/c/tests/CMakeLists.txt b/c/tests/CMakeLists.txt index 3ec7581b15..7fbb274524 100644 --- a/c/tests/CMakeLists.txt +++ b/c/tests/CMakeLists.txt @@ -93,6 +93,7 @@ ConfigureTest( ) ConfigureTest(NAME PCA_C_TEST PATH preprocessing/pca_c.cu) +ConfigureTest(NAME RABITQ_C_TEST PATH preprocessing/rabitq_c.cu) if(BUILD_CAGRA_HNSWLIB) ConfigureTest(NAME HNSW_C_TEST PATH neighbors/ann_hnsw_c.cu) diff --git a/c/tests/preprocessing/rabitq_c.cu b/c/tests/preprocessing/rabitq_c.cu new file mode 100644 index 0000000000..2d6a86eb5c --- /dev/null +++ b/c/tests/preprocessing/rabitq_c.cu @@ -0,0 +1,261 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include +#include +#include +#include + +#include "../../src/core/interop.hpp" +#include + +#include +#include +#include + +namespace { + +template +DLManagedTensor make_row_major_tensor(T* data, int64_t rows, int64_t cols) +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack( + raft::make_device_matrix_view(data, rows, cols), &tensor); + return tensor; +} + +template +DLManagedTensor make_vector_tensor(T* data, int64_t size) +{ + DLManagedTensor tensor{}; + cuvs::core::to_dlpack(raft::make_device_vector_view(data, size), &tensor); + return tensor; +} + +void free_tensor(DLManagedTensor& t) +{ + if (t.deleter) { t.deleter(&t); } +} + +template +std::vector to_host(raft::device_resources const& handle, T const* d_ptr, size_t len) +{ + std::vector h(len); + raft::copy(h.data(), d_ptr, len, raft::resource::get_cuda_stream(handle)); + handle.sync_stream(); + return h; +} + +/** Mirrors detail::padded_dim(): the working dimension is `dim` rounded up to 32. */ +constexpr int64_t expected_padded_dim(int64_t dim) { return ((dim + 31) / 32) * 32; } + +/** Mirrors detail::flip_bytes(): four sign-flip passes of `padded_dim` bits each. */ +constexpr int64_t expected_flip_bytes(int64_t padded_dim) { return 4 * padded_dim / 8; } + +/** + * Exercises the whole C surface for one rotator kind: params -> pipeline -> + * getters -> transform -> state readback -> teardown. + */ +void run_rabitq_c_smoke(enum cuvsRabitqRotatorKind kind) +{ + raft::device_resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + constexpr int64_t n_rows = 64; + // Deliberately not a multiple of 32, so the dim -> padded_dim padding is exercised. + constexpr int64_t dim = 40; + constexpr int64_t padded_dim = expected_padded_dim(dim); + constexpr uint32_t ex_bits = 3; + // Codes hold `ex_bits + 1` bit unsigned values, so this is the largest legal code. + constexpr uint32_t max_code = (1u << (ex_bits + 1)) - 1u; + // Pre-fill `codes` with a byte larger than max_code: any element the transform + // fails to write stays out of range and is caught by the range check below. + constexpr int unwritten_sentinel = 0xEE; + + static_assert(padded_dim == 64, "test expects dim=40 to pad up to 64"); + static_assert(unwritten_sentinel > static_cast(max_code), "sentinel must be detectable"); + + rmm::device_uvector dataset(n_rows * dim, stream); + rmm::device_uvector centroid(dim, stream); + raft::random::RngState rng(1234ULL); + raft::random::uniform(handle, rng, dataset.data(), n_rows * dim, -1.0f, 1.0f); + raft::random::uniform(handle, rng, centroid.data(), dim, -0.5f, 0.5f); + handle.sync_stream(); + + cuvsResources_t res; + ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS); + + // --- params: defaults must mirror the C++ member initializers --- + cuvsRabitqQuantizerParams_t params; + ASSERT_EQ(cuvsRabitqQuantizerParamsCreate(¶ms), CUVS_SUCCESS); + EXPECT_EQ(params->ex_bits, 3u); + EXPECT_EQ(params->rotator, CUVS_RABITQ_ROTATOR_KIND_FHT_KAC); + EXPECT_EQ(params->seed, 12345ULL); + EXPECT_EQ(params->delta_mode, CUVS_RABITQ_DELTA_KIND_RECONSTRUCTION); + EXPECT_TRUE(params->use_fast); + EXPECT_EQ(params->coarse_samples, 64u); + EXPECT_EQ(params->fine_samples, 64u); + + params->ex_bits = ex_bits; + params->rotator = kind; + params->seed = 42ULL; + + // --- factories: one call populates both peers --- + cuvsRabitqRotator_t rotator; + cuvsRabitqQuantizer_t quantizer; + ASSERT_EQ(cuvsRabitqRotatorCreate(&rotator), CUVS_SUCCESS); + ASSERT_EQ(cuvsRabitqQuantizerCreate(&quantizer), CUVS_SUCCESS); + ASSERT_EQ(cuvsRabitqPipelineMake(res, params, dim, rotator, quantizer), CUVS_SUCCESS); + + // --- metadata getters --- + enum cuvsRabitqRotatorKind read_kind = CUVS_RABITQ_ROTATOR_KIND_MATMUL; + ASSERT_EQ(cuvsRabitqRotatorGetKind(rotator, &read_kind), CUVS_SUCCESS); + EXPECT_EQ(read_kind, kind); + + int64_t rotator_padded_dim = 0; + int64_t quantizer_padded_dim = 0; + int64_t quantizer_dim = 0; + ASSERT_EQ(cuvsRabitqRotatorGetPaddedDim(rotator, &rotator_padded_dim), CUVS_SUCCESS); + ASSERT_EQ(cuvsRabitqQuantizerGetPaddedDim(quantizer, &quantizer_padded_dim), CUVS_SUCCESS); + ASSERT_EQ(cuvsRabitqQuantizerGetDim(quantizer, &quantizer_dim), CUVS_SUCCESS); + EXPECT_EQ(rotator_padded_dim, padded_dim); + EXPECT_EQ(quantizer_padded_dim, padded_dim); + EXPECT_EQ(quantizer_dim, dim); + + // use_fast is on, so the factor was actually estimated on the GPU. + float const_scaling_factor = 0.0f; + ASSERT_EQ(cuvsRabitqQuantizerGetConstScalingFactor(quantizer, &const_scaling_factor), + CUVS_SUCCESS); + EXPECT_TRUE(std::isfinite(const_scaling_factor)); + EXPECT_GT(const_scaling_factor, 0.0f); + + // --- transform --- + rmm::device_uvector codes(n_rows * padded_dim, stream); + rmm::device_uvector delta(n_rows, stream); + rmm::device_uvector vl(n_rows, stream); + RAFT_CUDA_TRY(cudaMemsetAsync(codes.data(), unwritten_sentinel, codes.size(), stream)); + handle.sync_stream(); + + auto dataset_t = make_row_major_tensor(dataset.data(), n_rows, dim); + auto centroid_t = make_vector_tensor(centroid.data(), dim); + auto codes_t = make_row_major_tensor(codes.data(), n_rows, padded_dim); + auto delta_t = make_vector_tensor(delta.data(), n_rows); + auto vl_t = make_vector_tensor(vl.data(), n_rows); + + // A NULL centroid must be accepted and mean "no centering". + ASSERT_EQ( + cuvsRabitqTransform(res, quantizer, rotator, &dataset_t, nullptr, &codes_t, &delta_t, &vl_t), + CUVS_SUCCESS); + ASSERT_EQ(cuvsStreamSync(res), CUVS_SUCCESS); + + { + auto h_codes = to_host(handle, codes.data(), codes.size()); + size_t n_out_of_range = 0; + bool all_zero = true; + for (auto c : h_codes) { + if (static_cast(c) > max_code) { ++n_out_of_range; } + if (c != 0) { all_zero = false; } + } + EXPECT_EQ(n_out_of_range, 0u) << n_out_of_range << " of " << h_codes.size() + << " codes are outside the " << (ex_bits + 1) + << "-bit range [0, " << max_code + << "] (an unwritten element keeps the 0xEE sentinel)"; + EXPECT_FALSE(all_zero) << "quantized codes are all zero"; + + auto h_delta = to_host(handle, delta.data(), delta.size()); + auto h_vl = to_host(handle, vl.data(), vl.size()); + // vl is documented as `delta * -(2^ex_bits - 0.5)`. + const float cb = -(static_cast(1u << ex_bits) - 0.5f); + for (int64_t i = 0; i < n_rows; ++i) { + EXPECT_TRUE(std::isfinite(h_delta[i])) << " row " << i; + EXPECT_GT(h_delta[i], 0.0f) << " row " << i; + EXPECT_NEAR(h_vl[i], h_delta[i] * cb, 1e-5f * std::fabs(h_delta[i] * cb) + 1e-20f) + << " row " << i; + } + } + + // A non-NULL centroid takes the residual path and must produce different codes. + auto h_codes_uncentered = to_host(handle, codes.data(), codes.size()); + RAFT_CUDA_TRY(cudaMemsetAsync(codes.data(), unwritten_sentinel, codes.size(), stream)); + handle.sync_stream(); + ASSERT_EQ(cuvsRabitqTransform( + res, quantizer, rotator, &dataset_t, ¢roid_t, &codes_t, &delta_t, &vl_t), + CUVS_SUCCESS); + ASSERT_EQ(cuvsStreamSync(res), CUVS_SUCCESS); + { + auto h_codes = to_host(handle, codes.data(), codes.size()); + size_t n_out_of_range = 0; + size_t n_diff = 0; + for (size_t i = 0; i < h_codes.size(); ++i) { + if (static_cast(h_codes[i]) > max_code) { ++n_out_of_range; } + if (h_codes[i] != h_codes_uncentered[i]) { ++n_diff; } + } + EXPECT_EQ(n_out_of_range, 0u) << "centered codes outside the " << (ex_bits + 1) << "-bit range"; + EXPECT_GT(n_diff, 0u) << "passing a centroid did not change the codes"; + } + + // --- rotator state readback --- + DLManagedTensor state{}; + ASSERT_EQ(cuvsRabitqRotatorGetState(rotator, &state), CUVS_SUCCESS); + EXPECT_NE(state.dl_tensor.data, nullptr); + EXPECT_TRUE(cuvs::core::is_dlpack_device_compatible(state.dl_tensor)); + if (kind == CUVS_RABITQ_ROTATOR_KIND_MATMUL) { + // padded_dim x padded_dim row-major float32 rotation matrix + ASSERT_EQ(state.dl_tensor.ndim, 2); + EXPECT_EQ(state.dl_tensor.dtype.code, kDLFloat); + EXPECT_EQ(state.dl_tensor.dtype.bits, 32); + EXPECT_EQ(state.dl_tensor.shape[0], padded_dim); + EXPECT_EQ(state.dl_tensor.shape[1], padded_dim); + // Non-empty in substance, not just in shape: a rotation matrix cannot be all zeros. + auto h_state = to_host( + handle, static_cast(state.dl_tensor.data), padded_dim * padded_dim); + bool all_zero = true; + for (auto v : h_state) { + if (v != 0.0f) { all_zero = false; } + } + EXPECT_FALSE(all_zero) << "rotation matrix read back as all zeros"; + } else { + // 4 * padded_dim / 8 packed uint8 sign bits + ASSERT_EQ(state.dl_tensor.ndim, 1); + EXPECT_EQ(state.dl_tensor.dtype.code, kDLUInt); + EXPECT_EQ(state.dl_tensor.dtype.bits, 8); + ASSERT_EQ(state.dl_tensor.shape[0], expected_flip_bytes(padded_dim)); + auto h_state = to_host(handle, + static_cast(state.dl_tensor.data), + static_cast(expected_flip_bytes(padded_dim))); + size_t n_set_bits = 0; + for (auto b : h_state) { + for (int i = 0; i < 8; ++i) { + n_set_bits += (b >> i) & 1u; + } + } + // Neither all-zero nor all-one: the flip bits are genuinely random. + EXPECT_GT(n_set_bits, 0u); + EXPECT_LT(n_set_bits, h_state.size() * 8); + } + free_tensor(state); + + // --- teardown --- + free_tensor(dataset_t); + free_tensor(centroid_t); + free_tensor(codes_t); + free_tensor(delta_t); + free_tensor(vl_t); + + EXPECT_EQ(cuvsRabitqRotatorDestroy(rotator), CUVS_SUCCESS); + EXPECT_EQ(cuvsRabitqQuantizerDestroy(quantizer), CUVS_SUCCESS); + EXPECT_EQ(cuvsRabitqQuantizerParamsDestroy(params), CUVS_SUCCESS); + EXPECT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS); +} + +} // namespace + +TEST(RabitqC, PipelineTransformFhtKac) { run_rabitq_c_smoke(CUVS_RABITQ_ROTATOR_KIND_FHT_KAC); } + +TEST(RabitqC, PipelineTransformMatmul) { run_rabitq_c_smoke(CUVS_RABITQ_ROTATOR_KIND_MATMUL); } diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 175669a620..dc6114fdc5 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1421,6 +1421,11 @@ if(NOT BUILD_CPU_ONLY) src/preprocessing/quantize/scalar.cu src/preprocessing/quantize/binary.cu src/preprocessing/quantize/pq.cu + src/preprocessing/quantize/rabitq.cu + src/preprocessing/quantize/detail/rabitq/rotator.cu + src/preprocessing/quantize/detail/rabitq/rescale_search.cu + src/preprocessing/quantize/detail/rabitq/quantize.cu + src/preprocessing/quantize/detail/rabitq/pipeline.cu src/preprocessing/spectral/spectral_embedding.cu src/preprocessing/pca/pca.cu ${select_k_inst_files} diff --git a/cpp/include/cuvs/preprocessing/quantize/rabitq.hpp b/cpp/include/cuvs/preprocessing/quantize/rabitq.hpp new file mode 100644 index 0000000000..0bbf3a3b7c --- /dev/null +++ b/cpp/include/cuvs/preprocessing/quantize/rabitq.hpp @@ -0,0 +1,468 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include + +#include +#include + +namespace CUVS_EXPORT cuvs { +namespace preprocessing { +namespace quantize { +namespace rabitq { + +/** + * @defgroup rabitq RaBitQ quantizer utilities + * @{ + */ + +/** + * @brief Random-rotation implementation used before quantization. + * + * Both kinds produce an orthogonal transform of the padded working dimension; + * they differ only in cost and in the amount of state they carry. + */ +enum class rotator_kind : uint8_t { + /** Dense `padded_dim x padded_dim` random orthogonal matrix (cuBLAS GEMM), O(n * dim^2). */ + matmul, + /** Fast Hadamard Transform combined with Kac's walk, O(n * dim * log(dim)). */ + fht_kac +}; + +/** + * @brief Rule used to derive the per-vector reconstruction scale `delta`. + * + * Only relevant for the `(codes, delta, vl)` output mode; the full-factor mode + * emits `(f_add, f_rescale, f_error)` instead and ignores this setting. + */ +enum class delta_kind : uint8_t { + /** `delta = (|residual| / |code|) * cos(residual, code)` - minimizes reconstruction error. */ + reconstruction = 0, + /** `delta = (|residual| / |code|) / cos(residual, code)` - unbiased inner-product estimate. */ + unbiased = 1, + /** `delta = |residual| / |code|` - plain norm ratio, no angular correction. */ + plain = 2 +}; + +/** + * @brief RaBitQ quantizer parameters. + * + * These parameters fully describe the quantizer: RaBitQ is data-oblivious, so + * none of them are learned from a dataset (see make_quantizer / make_rotator). + */ +struct params { + /** + * Number of *extended* bits per dimension. The total code width is + * `ex_bits + 1` bits (one sign bit plus `ex_bits` magnitude bits), so + * `ex_bits = 0` is plain 1-bit RaBitQ and `ex_bits = 3` is the common + * 4-bit configuration. + * + * Possible value range: [0, 8]. + */ + uint32_t ex_bits = 3; + /** Random-rotation implementation. */ + rotator_kind rotator = rotator_kind::fht_kac; + /** + * Seed for the host RNG that draws the rotation state and the random probe + * vectors used to estimate the constant scaling factor. + */ + uint64_t seed = 12345ULL; + /** Rule used to derive the per-vector `delta` / `vl` pair. */ + delta_kind delta_mode = delta_kind::reconstruction; + /** + * When true, every vector is quantized with a single pre-estimated scaling + * factor (`quantizer::const_scaling_factor`) instead of running the + * per-vector rescale search. This is substantially faster and is the + * recommended setting; disable it to trade build throughput for slightly + * tighter codes. + */ + bool use_fast = true; + /** Number of samples in the coarse stage of the per-vector rescale search. */ + uint32_t coarse_samples = 64; + /** Number of samples in the refinement stage of the per-vector rescale search. */ + uint32_t fine_samples = 64; +}; + +/** + * @brief Random rotation applied to the data before quantization. + * + * The rotator carries no learned information - its state is drawn from + * `params::seed` - but it *is* part of the code representation: codes produced + * with one rotator are only comparable to queries transformed by the very same + * rotator. There is deliberately no serialization API here, so the caller owns + * this object and is responsible for keeping it alongside the codes (e.g. by + * storing `params::seed` and re-running make_rotator, or by embedding the + * buffers below in its own index format). + * + * A rotator is an independent peer of the quantizer: neither contains the + * other, and a single rotator may be shared by several quantizers. + */ +struct rotator { + /** Which implementation the state below belongs to. */ + rotator_kind kind = rotator_kind::fht_kac; + /** + * Working dimension: `dim` rounded up to a multiple of 32. + * + * This is the rotator's full identity alongside `kind` and the state below โ€” + * both transforms operate purely on the padded width (`fht_kac` derives its + * segment size as `2^floor_log2(padded_dim)`), so the original unpadded `dim` + * is not part of the rotator. + */ + int64_t padded_dim = 0; + /** + * `padded_dim x padded_dim` row-major random orthogonal matrix. + * Empty unless `kind == rotator_kind::matmul`. + */ + raft::device_matrix rotation_matrix; + /** + * Packed sign bits of the four Kac's-walk passes (`4 * padded_dim / 8` bytes). + * Empty unless `kind == rotator_kind::fht_kac`. + */ + raft::device_vector flip_bits; + + /** @brief Construct an empty (uninitialized) rotator. */ + explicit rotator(raft::resources const& res) + : rotation_matrix(raft::make_device_matrix(res, 0, 0)), + flip_bits(raft::make_device_vector(res, 0)) + { + } +}; + +/** + * @brief Stores the (data-independent) configuration used to quantize rotated residuals. + * + * RaBitQ needs no training data, so this is plain metadata plus the single + * scalar estimated at init time. Like the rotator, it is owned by the caller + * and has no serialization API - it is cheap to rebuild from `params` and + * `dim`. + */ +struct quantizer { + /** Parameters used to build this quantizer. */ + params params_quantizer; + /** + * Scaling factor shared by all vectors on the `params::use_fast` path, + * estimated from random Gaussian probe vectors at init time. + * Left at 1.0 (and unused) when `params::use_fast` is false. + */ + float const_scaling_factor = 1.0f; + /** Dimensionality of the un-rotated input vectors. */ + int64_t dim = 0; + /** Working dimension: `dim` rounded up to a multiple of 32. */ + int64_t padded_dim = 0; +}; + +/** + * @brief Convenience composition of the two peers for the common end-to-end flow. + * + * Holds one rotator and one quantizer; it adds no state and no behaviour of its + * own. Pass `pipeline::quantizer` and `pipeline::rotator` to transform(). + */ +struct pipeline { + /** Rotation applied before quantization. */ + rabitq::rotator rotator; + /** Quantizer applied to the rotated residuals. */ + rabitq::quantizer quantizer; + + /** @brief Construct an empty (uninitialized) pipeline. */ + explicit pipeline(raft::resources const& res) : rotator(res) {} +}; + +/** + * @brief Initializes a RaBitQ rotator. + * + * No training data is required: the rotation is drawn from `params::seed` + * alone. The returned object owns its device buffers; keep it for as long as + * the codes it produced are in use (see rotator). + * + * Usage example: + * @code{.cpp} + * raft::resources res; + * cuvs::preprocessing::quantize::rabitq::params params; + * auto rotator = + * cuvs::preprocessing::quantize::rabitq::make_rotator(res, params, dataset.extent(1)); + * @endcode + * + * @param[in] res raft resource + * @param[in] params configure the rotator, e.g. kind and seed + * @param[in] dim dimensionality of the un-rotated input vectors + * + * @return rotator + */ +rotator make_rotator(raft::resources const& res, params const& params, int64_t dim); + +/** + * @brief Initializes a RaBitQ quantizer. + * + * No training data is required. When `params::use_fast` is set, this estimates + * the constant scaling factor on the GPU, which is the only non-trivial work + * performed here. + * + * Usage example: + * @code{.cpp} + * raft::resources res; + * cuvs::preprocessing::quantize::rabitq::params params; + * auto quantizer = + * cuvs::preprocessing::quantize::rabitq::make_quantizer(res, params, dataset.extent(1)); + * @endcode + * + * @param[in] res raft resource + * @param[in] params configure the quantizer, e.g. ex_bits + * @param[in] dim dimensionality of the un-rotated input vectors + * + * @return quantizer + */ +quantizer make_quantizer(raft::resources const& res, params const& params, int64_t dim); + +/** + * @brief Initializes a rotator and a quantizer from the same parameters. + * + * Equivalent to calling make_rotator() and make_quantizer() with the same + * arguments. No training data is required. + * + * Usage example: + * @code{.cpp} + * raft::resources res; + * cuvs::preprocessing::quantize::rabitq::params params; + * auto pipeline = + * cuvs::preprocessing::quantize::rabitq::make_pipeline(res, params, dataset.extent(1)); + * cuvs::preprocessing::quantize::rabitq::transform(res, pipeline.quantizer, pipeline.rotator, + * dataset, std::nullopt, codes.view(), delta.view(), vl.view()); + * @endcode + * + * @param[in] res raft resource + * @param[in] params configure both components + * @param[in] dim dimensionality of the un-rotated input vectors + * + * @return pipeline + */ +pipeline make_pipeline(raft::resources const& res, params const& params, int64_t dim); + +/** + * @brief Applies the random rotation to already-padded rows. + * + * Both `in` and `out` must be `n_rows x rotator::padded_dim`; this entry point + * does not pad. In-place operation (`in.data_handle() == out.data_handle()`) is + * supported by `rotator_kind::fht_kac` only. + * + * Use this to rotate queries or externally computed centroids so that they live + * in the same rotated space as the codes. + * + * Usage example: + * @code{.cpp} + * auto rotated = raft::make_device_matrix(res, n_queries, rotator.padded_dim); + * cuvs::preprocessing::quantize::rabitq::rotate(res, rotator, padded_queries, rotated.view()); + * @endcode + * + * @param[in] res raft resource + * @param[in] rotator a RaBitQ rotator + * @param[in] in a row-major `n_rows x padded_dim` matrix view on device + * @param[out] out a row-major `n_rows x padded_dim` matrix view on device + * + */ +void rotate(raft::resources const& res, + rotator const& rotator, + raft::device_matrix_view in, + raft::device_matrix_view out); + +/** + * @brief Quantizes a raw dataset, emitting per-vector `(delta, vl)` factors. + * + * Performs the whole chain: zero-pad `dim -> padded_dim`, rotate, optionally + * subtract the rotated centroid, then quantize the residuals. One code element + * is written per padded dimension; each holds an unsigned + * `params::ex_bits + 1` bit value, so `codes` must be `n_rows x padded_dim`. + * + * A temporary `n_rows x padded_dim` float buffer holding the rotated residuals + * is allocated from the workspace resource, so large datasets should be + * transformed in batches (a pool workspace resource is recommended). + * + * Usage example: + * @code{.cpp} + * raft::resources res; + * cuvs::preprocessing::quantize::rabitq::params params; + * auto quantizer = cuvs::preprocessing::quantize::rabitq::make_quantizer(res, params, dim); + * auto rotator = cuvs::preprocessing::quantize::rabitq::make_rotator(res, params, dim); + * auto codes = raft::make_device_matrix(res, n_rows, quantizer.padded_dim); + * auto delta = raft::make_device_vector(res, n_rows); + * auto vl = raft::make_device_vector(res, n_rows); + * cuvs::preprocessing::quantize::rabitq::transform(res, quantizer, rotator, dataset, + * std::nullopt, codes.view(), delta.view(), vl.view()); + * @endcode + * + * @param[in] res raft resource + * @param[in] quantizer a RaBitQ quantizer + * @param[in] rotator the rotator whose `padded_dim` matches the quantizer + * @param[in] dataset a row-major `n_rows x dim` matrix view on device + * @param[in] centroid optional un-rotated `dim` centroid on device; when set, codes are + * computed on the residuals w.r.t. this centroid + * @param[out] codes a row-major `n_rows x padded_dim` matrix view on device + * @param[out] delta per-vector scale, `n_rows` elements on device + * @param[out] vl per-vector offset (`delta * -(2^ex_bits - 0.5)`), `n_rows` elements on device + * + */ +void transform(raft::resources const& res, + quantizer const& quantizer, + rotator const& rotator, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl); + +/** @copydoc transform */ +void transform(raft::resources const& res, + quantizer const& quantizer, + rotator const& rotator, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl); + +/** + * @brief Quantizes a raw dataset, emitting the full `(f_add, f_rescale, f_error)` factor triplet. + * + * Same chain as the `(delta, vl)` overload, but the per-vector factors are the + * ones used for approximate-distance estimation during search: + * `factors[i] = { f_add, f_rescale, f_error }`. + * + * A temporary `n_rows x padded_dim` float buffer holding the rotated residuals + * is allocated from the workspace resource, so large datasets should be + * transformed in batches (a pool workspace resource is recommended). + * + * The rotated centroid is an internal temporary as well. A search implementation + * that needs it can reproduce it exactly by zero-padding `centroid` to + * `padded_dim` and calling rotate() with the same rotator. + * + * Usage example: + * @code{.cpp} + * auto codes = raft::make_device_matrix(res, n_rows, quantizer.padded_dim); + * auto factors = raft::make_device_matrix(res, n_rows, 3); + * cuvs::preprocessing::quantize::rabitq::transform(res, quantizer, rotator, dataset, + * std::make_optional(centroid.view()), codes.view(), factors.view()); + * @endcode + * + * @param[in] res raft resource + * @param[in] quantizer a RaBitQ quantizer + * @param[in] rotator the rotator whose `padded_dim` matches the quantizer + * @param[in] dataset a row-major `n_rows x dim` matrix view on device + * @param[in] centroid optional un-rotated `dim` centroid on device; treated as zero when unset + * @param[out] codes a row-major `n_rows x padded_dim` matrix view on device + * @param[out] factors a row-major `n_rows x 3` matrix view on device + * + */ +void transform(raft::resources const& res, + quantizer const& quantizer, + rotator const& rotator, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors); + +/** @copydoc transform */ +void transform(raft::resources const& res, + quantizer const& quantizer, + rotator const& rotator, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors); + +/** + * @brief Quantizes pre-rotated residuals, emitting per-vector `(delta, vl)` factors. + * + * The rotator is not needed here: `residuals` must already be rotated, padded + * to `quantizer::padded_dim` and centered (centroid subtracted). Use this when + * the residuals are produced elsewhere, e.g. by an IVF build that has already + * rotated its data, and rotate() plus a subtraction of your own. + * + * Usage example: + * @code{.cpp} + * auto codes = raft::make_device_matrix(res, n_rows, quantizer.padded_dim); + * auto delta = raft::make_device_vector(res, n_rows); + * auto vl = raft::make_device_vector(res, n_rows); + * cuvs::preprocessing::quantize::rabitq::transform_residuals(res, quantizer, residuals, + * codes.view(), delta.view(), vl.view()); + * @endcode + * + * @param[in] res raft resource + * @param[in] quantizer a RaBitQ quantizer + * @param[in] residuals a row-major `n_rows x padded_dim` matrix view on device, rotated and + * centered + * @param[out] codes a row-major `n_rows x padded_dim` matrix view on device + * @param[out] delta per-vector scale, `n_rows` elements on device + * @param[out] vl per-vector offset (`delta * -(2^ex_bits - 0.5)`), `n_rows` elements on device + * + */ +void transform_residuals(raft::resources const& res, + quantizer const& quantizer, + raft::device_matrix_view residuals, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl); + +/** @copydoc transform_residuals */ +void transform_residuals(raft::resources const& res, + quantizer const& quantizer, + raft::device_matrix_view residuals, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl); + +/** + * @brief Quantizes pre-rotated residuals, emitting the full factor triplet. + * + * As above, `residuals` must already be rotated, padded to + * `quantizer::padded_dim` and centered. The full-factor formula also reads the + * centroid, which must be supplied *in the rotated, padded space* (rotate() it + * with the same rotator that produced the residuals); pass `std::nullopt` to + * treat it as zero. + * + * Usage example: + * @code{.cpp} + * auto codes = raft::make_device_matrix(res, n_rows, quantizer.padded_dim); + * auto factors = raft::make_device_matrix(res, n_rows, 3); + * cuvs::preprocessing::quantize::rabitq::transform_residuals(res, quantizer, residuals, + * std::make_optional(rotated_centroid.view()), codes.view(), factors.view()); + * @endcode + * + * @param[in] res raft resource + * @param[in] quantizer a RaBitQ quantizer + * @param[in] residuals a row-major `n_rows x padded_dim` matrix view on device, rotated and + * centered + * @param[in] rotated_centroid optional rotated, padded centroid (`padded_dim` elements on device) + * @param[out] codes a row-major `n_rows x padded_dim` matrix view on device + * @param[out] factors a row-major `n_rows x 3` matrix view on device + * + */ +void transform_residuals( + raft::resources const& res, + quantizer const& quantizer, + raft::device_matrix_view residuals, + std::optional> rotated_centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors); + +/** @copydoc transform_residuals */ +void transform_residuals( + raft::resources const& res, + quantizer const& quantizer, + raft::device_matrix_view residuals, + std::optional> rotated_centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors); + +/** @} */ // end of group rabitq + +} // namespace rabitq +} // namespace quantize +} // namespace preprocessing +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/fht.cuh b/cpp/src/preprocessing/quantize/detail/rabitq/fht.cuh new file mode 100644 index 0000000000..6230c4592a --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/fht.cuh @@ -0,0 +1,695 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +/****************************************************************************** + * Fast Hadamard Transform + Kac's-walk rotation kernels (CUDA). + * + * Two execution strategies share one butterfly network. Both apply the + * transform levels in ascending element-bit order with identical arithmetic + * forms, so they produce bit-identical results: + * + * - Warp-per-vector kernels (large batches, padded dim <= 2048): one warp + * holds the whole vector in registers, lane-strided; levels 0-4 are warp + * shuffles, higher levels are in-register butterflies. No shared memory, + * no block barriers. + * + * - Cooperative kernels (small batches, and dims beyond the warp kernels' + * register budget): one block per vector; warp w owns the contiguous + * slice [w*S, (w+1)*S), lane-strided in registers. Levels below log2(S) + * stay inside the warp; each of the top log2(W) levels publishes the + * slices to a shared-memory mirror and combines with the partner + * slice's value at the same position. Threads never exchange element + * ownership. + ******************************************************************************/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +namespace fht { + +// ============================================================================ +// Compile-time helpers +// ============================================================================ + +constexpr __host__ __device__ int clog2(int val) { return val > 1 ? 1 + clog2(val >> 1) : 0; } + +/// Threads used by the cooperative kernels for a transform of 2^log_len +/// elements: length/4 clamped to [32, 256], except 512 for the largest size +/// (keeps registers per lane at <= 64 and avoids spills). +constexpr int coop_log_threads(int log_len) { + int t = log_len - 2; + if (t < 5) t = 5; + if (t > 8) t = 8; + if (log_len >= 15) t = 9; + return t; +} + +/// Generic compile-time dispatch: converts a runtime log_N (3..15) into a +/// template parameter by recursive if-constexpr. +/// Usage: dispatch_log_n(log_N, [&](auto K) { launch(...); }); +template +inline void dispatch_log_n(int log_n, Func&& func) { + if (log_n == kCur) { + func(std::integral_constant{}); + } else if constexpr (kCur < kMax) { + dispatch_log_n(log_n, std::forward(func)); + } else { + std::cerr << "fht::dispatch_log_n: unsupported log_N=" << log_n << std::endl; + exit(EXIT_FAILURE); + } +} + +template +inline void dispatch_p32(int p32, Func&& func) { + if (p32 == kCur) { + func(std::integral_constant{}); + } else if constexpr (kCur < kMax) { + dispatch_p32(p32, std::forward(func)); + } else { + std::cerr << "fht::dispatch_p32: unsupported p32=" << p32 << std::endl; + exit(EXIT_FAILURE); + } +} + +// ============================================================================ +// Device helpers +// ============================================================================ + +__device__ __forceinline__ float flip_sign_bit(float v, uint32_t bit) { + return __uint_as_float(__float_as_uint(v) ^ (bit << 31)); +} + +// ============================================================================ +// Warp-per-vector fused rotation kernels (large-batch path) +// +// Each warp holds one full (padded) vector in registers, so all butterfly +// levels are either in-thread ops or warp shuffles: no shared memory and no +// __syncthreads(). Used when N is large enough that one-warp-per-vector +// saturates the GPU (see kWarpPathMinN); small batches use the cooperative +// kernels below, which spread a single vector across more threads. +// ============================================================================ + +// ---------------------------------------------------------------------------- +// Power-of-2 path. Lane `l` of the warp owns float4 chunks at element offsets +// 128*c + 4*l, i.e. element-index bits [1:0] select within the float4, +// bits [6:2] the lane, bits [7+] the chunk register. +// ---------------------------------------------------------------------------- +template +__global__ void fht_kac_rotate_warp_kernel( + const float* input, // no __restrict__: in-place rotation allowed + float* output, + const uint8_t* __restrict__ flip, + int N, float total_scale) +{ + constexpr int kDim = 1 << kLogN; + constexpr int kC = kDim / 128; // float4 registers per lane + constexpr int kLogC = clog2(kC); + constexpr int kWordsPerRound = kDim / 32; + static_assert(kLogN >= 7, "warp kernel needs at least one float4 per lane"); + + const int lane = threadIdx.x & 31; + const int vec_id = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); + if (vec_id >= N) return; + + const uint32_t* flip32 = reinterpret_cast(flip); + + const float4* vin = reinterpret_cast(input + (size_t)vec_id * kDim); + float4 x[kC]; + #pragma unroll + for (int c = 0; c < kC; ++c) x[c] = vin[c * 32 + lane]; + + const int bit_shift = 4 * (lane & 7); // this lane's 4 flip bits inside the word + const int word_off = lane >> 3; + + #pragma unroll + for (int round = 0; round < 4; ++round) { + // Sign flip (element e uses bit e%32 of word e/32; all four of this + // lane's bits sit in one word). + #pragma unroll + for (int c = 0; c < kC; ++c) { + uint32_t w = __ldg(&flip32[round * kWordsPerRound + 4 * c + word_off]) >> bit_shift; + x[c].x = flip_sign_bit(x[c].x, w & 1u); + x[c].y = flip_sign_bit(x[c].y, (w >> 1) & 1u); + x[c].z = flip_sign_bit(x[c].z, (w >> 2) & 1u); + x[c].w = flip_sign_bit(x[c].w, (w >> 3) & 1u); + } + + // Levels 0-1: inside each float4. + #pragma unroll + for (int c = 0; c < kC; ++c) { + float a, b; + a = x[c].x; b = x[c].y; x[c].x = a + b; x[c].y = a - b; + a = x[c].z; b = x[c].w; x[c].z = a + b; x[c].w = a - b; + a = x[c].x; b = x[c].z; x[c].x = a + b; x[c].z = a - b; + a = x[c].y; b = x[c].w; x[c].y = a + b; x[c].w = a - b; + } + + // Levels 2-6: element bit (lv+2) == lane bit lv -> warp shuffles. + #pragma unroll + for (int lv = 0; lv < 5; ++lv) { + const int mask = 1 << lv; + const float sgn = (lane & mask) ? -1.f : 1.f; + #pragma unroll + for (int c = 0; c < kC; ++c) { + float o; + o = __shfl_xor_sync(0xffffffff, x[c].x, mask); x[c].x = sgn * x[c].x + o; + o = __shfl_xor_sync(0xffffffff, x[c].y, mask); x[c].y = sgn * x[c].y + o; + o = __shfl_xor_sync(0xffffffff, x[c].z, mask); x[c].z = sgn * x[c].z + o; + o = __shfl_xor_sync(0xffffffff, x[c].w, mask); x[c].w = sgn * x[c].w + o; + } + } + + // Levels 7+: across chunk registers, in-thread. Pair p of a level + // sits at offset (p % stride) inside butterfly group (p / stride). + #pragma unroll + for (int lv = 0; lv < kLogC; ++lv) { + const int stride = 1 << lv; + #pragma unroll + for (int p = 0; p < kC / 2; ++p) { + const int lo = (p / stride) * (2 * stride) + (p % stride); + const int hi = lo + stride; + float4 a = x[lo], b = x[hi]; + x[lo].x = a.x + b.x; x[lo].y = a.y + b.y; + x[lo].z = a.z + b.z; x[lo].w = a.w + b.w; + x[hi].x = a.x - b.x; x[hi].y = a.y - b.y; + x[hi].z = a.z - b.z; x[hi].w = a.w - b.w; + } + } + } + + float4* vout = reinterpret_cast(output + (size_t)vec_id * kDim); + #pragma unroll + for (int c = 0; c < kC; ++c) { + float4 v; + v.x = __fmul_rn(x[c].x, total_scale); v.y = __fmul_rn(x[c].y, total_scale); + v.z = __fmul_rn(x[c].z, total_scale); v.w = __fmul_rn(x[c].w, total_scale); + vout[c * 32 + lane] = v; + } +} + +// ---------------------------------------------------------------------------- +// Non-power-of-2 path. Lane-strided layout: lane `l` owns elements l + 32*j +// in register x[j]. Fully templated on (kLogT, kP32 = padded_dim/32) so all +// register indices are compile-time: +// - padded dim P = 32*kP32, registers per lane E = kP32 +// - FHT segment of T = 2^kLogT elements at offset 0 (even rounds) or +// P - T (odd rounds); both are multiples of 32, so a segment is a +// contiguous register window [base, base + T/32). +// - segment element bits [4:0] = lane -> levels 0-4 via shuffles, +// bits [5+] = register -> levels 5+ in-thread. +// - Kac's walk pairs (e, e + P/2). For even kP32, P/2 is a multiple of 32 +// and the partner is register j + kP32/2 in the same lane. For odd +// kP32, P/2 = 32*kHF + 16, so the partner lives in lane^16; the pairs +// form independent 2-cycles between register j of lanes < 16 and +// register m(j) of lanes >= 16, each resolved by one shfl_xor(16) with +// both register indices compile-time. +// ---------------------------------------------------------------------------- +template +__global__ void fht_kac_rotate_warp_nonpow2_kernel( + const float* input, // no __restrict__: in-place rotation allowed + float* output, + const uint8_t* __restrict__ flip, + int N, float fac, float final_scale) +{ + static_assert(kLogT >= 5 && kLogT <= 10, "warp nonpow2 kernel needs T in [32, 1024]"); + constexpr int kE = kP32; // registers = P/32 + constexpr int kP = 32 * kP32; // padded dim + constexpr int kT32 = 1 << (kLogT - 5); + constexpr int kB = kE - kT32; // (P - T)/32 + constexpr int kShflLevels = 5; + constexpr int kRegLevels = (kLogT > 5) ? (kLogT - 5) : 0; + static_assert(kB >= 0 && kB < kE, "padded dim must be in (T, 2T]"); + + const int lane = threadIdx.x & 31; + const int vec_id = blockIdx.x * (blockDim.x >> 5) + (threadIdx.x >> 5); + if (vec_id >= N) return; + + const uint32_t* flip32 = reinterpret_cast(flip); + + const float* vin = input + (size_t)vec_id * kP; + float x[kE]; + #pragma unroll + for (int j = 0; j < kE; ++j) x[j] = vin[lane + 32 * j]; + + #pragma unroll + for (int round = 0; round < 4; ++round) { + // Sign flip on all P elements (element l + 32j -> bit l of word j). + #pragma unroll + for (int j = 0; j < kE; ++j) { + uint32_t w = __ldg(&flip32[round * kE + j]); + x[j] = flip_sign_bit(x[j], (w >> lane) & 1u); + } + + const int base = (round & 1) ? kB : 0; // constant per unrolled round + + // FHT levels 0-4 on the segment: segment bit lv == lane bit lv. + #pragma unroll + for (int lv = 0; lv < kShflLevels; ++lv) { + const int mask = 1 << lv; + const float sgn = (lane & mask) ? -1.f : 1.f; + #pragma unroll + for (int j = 0; j < kT32; ++j) { + float o = __shfl_xor_sync(0xffffffff, x[base + j], mask); + x[base + j] = sgn * x[base + j] + o; + } + } + + // FHT levels 5+ across the segment's register window, in-thread. + #pragma unroll + for (int lv = 0; lv < kRegLevels; ++lv) { + const int stride = 1 << lv; + #pragma unroll + for (int p = 0; p < kT32 / 2; ++p) { + const int lo = base + (p / stride) * (2 * stride) + (p % stride); + const int hi = lo + stride; + float a = x[lo], b = x[hi]; + x[lo] = a + b; x[hi] = a - b; + } + } + + // Per-round fac on the FHT segment only (kacs walk mixes scales). + // __fmul_rn keeps this a standalone multiply (never FMA-contracted). + #pragma unroll + for (int j = 0; j < kT32; ++j) x[base + j] = __fmul_rn(x[base + j], fac); + + // Kac's walk. + if constexpr ((kP32 & 1) == 0) { + // P/2 = 32*(kP32/2): partner is the same lane, register j + half. + constexpr int kHalfR = kP32 / 2; + #pragma unroll + for (int j = 0; j < kHalfR; ++j) { + float a = x[j], b = x[j + kHalfR]; + x[j] = a + b; x[j + kHalfR] = a - b; + } + } else { + // P/2 = 32*kHF + 16: partner of element l + 32j is in lane^16. + // Register j of lanes < 16 pairs with register m(j) of lanes + // >= 16 (a mutual 2-cycle), where lanes-<16 slots j <= kHF and + // lanes->=16 slots m < kHF hold the first-half element of their + // pair. Exchange pre-values via one shuffle, then combine. + constexpr int kHF = (kP32 - 1) / 2; + #pragma unroll + for (int j = 0; j < kP32; ++j) { + const int m = (j < kHF) ? (j + kHF) + : (j == kHF) ? (kP32 - 1) : (j - kHF - 1); + float send = (lane < 16) ? x[j] : x[m]; + float o = __shfl_xor_sync(0xffffffff, send, 16); + if (lane < 16) x[j] = (j <= kHF) ? (x[j] + o) : (o - x[j]); + else x[m] = (m < kHF) ? (x[m] + o) : (o - x[m]); + } + } + } + + float* vout = output + (size_t)vec_id * kP; + #pragma unroll + for (int j = 0; j < kE; ++j) vout[lane + 32 * j] = __fmul_rn(x[j], final_scale); +} + +// ============================================================================ +// Cooperative fused rotation kernels (small-batch / large-dim path) +// +// One block per vector, W = threads/32 warps. Warp w owns the contiguous +// slice [w*S, (w+1)*S) of the transform, lane-strided within the slice: +// element(w, l, j) = w*S + l + 32*j. Levels below log2(S) are handled inside +// the warp exactly like the warp kernels above. Each of the top log2(W) +// levels publishes the slices to a shared-memory mirror and combines with +// the partner slice's value at the same (lane, register) position โ€” +// conflict-free stride-1 accesses, no ownership exchange. +// ============================================================================ + +template +__global__ void __launch_bounds__(1 << coop_log_threads(kLogN)) fht_kac_rotate_coop_kernel( + const float* input, // no __restrict__: in-place rotation allowed + float* output, + const uint8_t* __restrict__ flip, + int N, float total_scale) +{ + constexpr int kDim = 1 << kLogN; + constexpr int kLogT = coop_log_threads(kLogN); + constexpr int kLogW = kLogT - 5; // log2(warps per block) + constexpr int kLogS = kLogN - kLogW; // slice = 2^kLogS elements + constexpr int kS = 1 << kLogS; + constexpr int kE = (kS >= 32) ? (kS / 32) : 1; // registers per lane + constexpr int kShfl = (kLogS < 5) ? kLogS : 5; + constexpr int kRegLv = (kLogS > 5) ? (kLogS - 5) : 0; + static_assert(kLogS >= 3, "slice must hold at least 8 elements"); + + const int lane = threadIdx.x & 31; + const int w = threadIdx.x >> 5; + const int vec_id = blockIdx.x; + if (vec_id >= N) return; + + extern __shared__ float s_vec[]; // kDim floats when kW > 1 + + const uint32_t* flip32 = reinterpret_cast(flip); + const float* vin = input + (size_t)vec_id * kDim; + float* vout = output + (size_t)vec_id * kDim; + const int sbase = w * kS; + + float x[kE]; + #pragma unroll + for (int j = 0; j < kE; ++j) + x[j] = (kS >= 32 || lane < kS) ? vin[sbase + lane + 32 * j] : 0.f; + + for (int round = 0; round < 4; ++round) { + // Sign flip in registers; flip bit index of element e is + // round*kDim + e within the packed bit stream. + #pragma unroll + for (int j = 0; j < kE; ++j) { + if (kS >= 32 || lane < kS) { + const int b = round * kDim + sbase + 32 * j + lane; + const uint32_t word = __ldg(&flip32[b >> 5]); + x[j] = flip_sign_bit(x[j], (word >> (b & 31)) & 1u); + } + } + + // Levels 0..kShfl-1: lane shuffles. + #pragma unroll + for (int lv = 0; lv < kShfl; ++lv) { + const int mask = 1 << lv; + const float sgn = (lane & mask) ? -1.f : 1.f; + #pragma unroll + for (int j = 0; j < kE; ++j) { + float o = __shfl_xor_sync(0xffffffff, x[j], mask); + x[j] = sgn * x[j] + o; + } + } + + // Levels 5..kLogS-1: register butterflies inside the slice. + #pragma unroll + for (int lv = 0; lv < kRegLv; ++lv) { + const int stride = 1 << lv; + #pragma unroll + for (int g = 0; g < kE; g += 2 * stride) { + #pragma unroll + for (int m = 0; m < stride; ++m) { + const int lo = g + m, hi = lo + stride; + float a = x[lo], b = x[hi]; + x[lo] = a + b; x[hi] = a - b; + } + } + } + + // Levels kLogS..kLogN-1: cross-slice, one level at a time. Publish + // this slice to the shared mirror, then combine with the partner + // slice's value at the same (lane, register) position. + if constexpr (kLogW > 0) { + for (int m = 0; m < kLogW; ++m) { + #pragma unroll + for (int j = 0; j < kE; ++j) + s_vec[sbase + lane + 32 * j] = x[j]; + __syncthreads(); + const float sgn = (w & (1 << m)) ? -1.f : 1.f; + const int pbase = (w ^ (1 << m)) * kS; + #pragma unroll + for (int j = 0; j < kE; ++j) { + float o = s_vec[pbase + lane + 32 * j]; + x[j] = sgn * x[j] + o; + } + __syncthreads(); + } + } + } + + #pragma unroll + for (int j = 0; j < kE; ++j) { + if (kS >= 32 || lane < kS) + vout[sbase + lane + 32 * j] = __fmul_rn(x[j], total_scale); + } +} + +// ---------------------------------------------------------------------------- +// Cooperative non-power-of-2 kernel. The padded vector lives in shared +// memory; sign flips and the Kac walk operate on it directly with +// block-strided loops, while each round's FHT segment (2^kLogTrunc elements +// at offset 0 / P-T) is pulled into registers slice-per-warp and transformed +// with the same machinery as the pow2 cooperative kernel. +// ---------------------------------------------------------------------------- +template +__global__ void __launch_bounds__(1 << coop_log_threads(kLogTrunc)) fht_kac_rotate_coop_nonpow2_kernel( + const float* input, // no __restrict__: in-place rotation allowed + float* output, + const uint8_t* __restrict__ flip, + int N, int padded_dim, float fac, float final_scale) +{ + constexpr int kTrunc = 1 << kLogTrunc; + constexpr int kLogT = coop_log_threads(kLogTrunc); + constexpr int kNThreads = 1 << kLogT; + constexpr int kLogW = kLogT - 5; // log2(warps per block) + constexpr int kLogS = kLogTrunc - kLogW; + constexpr int kS = 1 << kLogS; + constexpr int kE = (kS >= 32) ? (kS / 32) : 1; + constexpr int kShfl = (kLogS < 5) ? kLogS : 5; + constexpr int kRegLv = (kLogS > 5) ? (kLogS - 5) : 0; + static_assert(kLogS >= 3, "slice must hold at least 8 elements"); + + const int lane = threadIdx.x & 31; + const int w = threadIdx.x >> 5; + const int tid = threadIdx.x; + const int vec_id = blockIdx.x; + if (vec_id >= N) return; + + // padded_dim floats for the vector, kTrunc scratch floats used as the + // alternate publish target of the cross-slice levels (one barrier per + // level instead of two), then the 4 rounds of flip words staged once so + // the per-round flip reads stay in shared memory. + extern __shared__ float s_vec[]; + + const uint32_t* flip32 = reinterpret_cast(flip); + const int P = padded_dim; + const int words = P / 32; // flip words per round + const int half = P / 2; + const int start = P - kTrunc; + + uint32_t* s_flip = reinterpret_cast(s_vec + P + kTrunc); + for (int i = tid; i < 4 * words; i += kNThreads) s_flip[i] = flip32[i]; + + const float* vin = input + (size_t)vec_id * P; + for (int i = tid; i < P; i += kNThreads) s_vec[i] = vin[i]; + __syncthreads(); + + for (int round = 0; round < 4; ++round) { + // Sign flip on the full padded vector. + for (int i = tid; i < P; i += kNThreads) { + const uint32_t word = s_flip[round * words + (i >> 5)]; + s_vec[i] = flip_sign_bit(s_vec[i], (word >> (i & 31)) & 1u); + } + __syncthreads(); + + // Pull this warp's slice of the segment into registers. + const int off = (round & 1) ? start : 0; + const int sbase = off + w * kS; + float x[kE]; + #pragma unroll + for (int j = 0; j < kE; ++j) + x[j] = (kS >= 32 || lane < kS) ? s_vec[sbase + lane + 32 * j] : 0.f; + + // Levels 0..kShfl-1: lane shuffles. + #pragma unroll + for (int lv = 0; lv < kShfl; ++lv) { + const int mask = 1 << lv; + const float sgn = (lane & mask) ? -1.f : 1.f; + #pragma unroll + for (int j = 0; j < kE; ++j) { + float o = __shfl_xor_sync(0xffffffff, x[j], mask); + x[j] = sgn * x[j] + o; + } + } + + // Levels 5..kLogS-1: register butterflies inside the slice. + #pragma unroll + for (int lv = 0; lv < kRegLv; ++lv) { + const int stride = 1 << lv; + #pragma unroll + for (int g = 0; g < kE; g += 2 * stride) { + #pragma unroll + for (int m = 0; m < stride; ++m) { + const int lo = g + m, hi = lo + stride; + float a = x[lo], b = x[hi]; + x[lo] = a + b; x[hi] = a - b; + } + } + } + + // Levels kLogS..kLogTrunc-1: cross-slice, one level at a time. + // Publishes alternate between the scratch area at s_vec[P..P+trunc) + // and the segment itself, so each level needs a single barrier (the + // next level writes the other buffer; per-warp program order makes + // its barrier also cover this level's reads). Parity is chosen so + // the LAST level publishes to scratch, keeping the fac writeback + // below race-free against partner reads. + if constexpr (kLogW > 0) { + for (int m = 0; m < kLogW; ++m) { + const bool to_scratch = ((kLogW - 1 - m) & 1) == 0; + const int pub = to_scratch ? (P + w * kS) : sbase; + #pragma unroll + for (int j = 0; j < kE; ++j) + s_vec[pub + lane + 32 * j] = x[j]; + __syncthreads(); + const float sgn = (w & (1 << m)) ? -1.f : 1.f; + const int pbase = (to_scratch ? P : off) + (w ^ (1 << m)) * kS; + #pragma unroll + for (int j = 0; j < kE; ++j) { + float o = s_vec[pbase + lane + 32 * j]; + x[j] = sgn * x[j] + o; + } + } + } + + // Write the segment back with the per-round fac (kacs walk mixes + // scales, so it cannot be deferred). Safe without a leading barrier: + // the segment's mirror was last read at cross-slice level kLogW-2, + // and the final level's publish barrier ordered those reads before + // this write (for kLogW == 0 the segment was only read by its own + // warp's slice load above). + #pragma unroll + for (int j = 0; j < kE; ++j) { + if (kS >= 32 || lane < kS) + s_vec[sbase + lane + 32 * j] = __fmul_rn(x[j], fac); + } + __syncthreads(); + + // Kac's walk. + for (int i = tid; i < half; i += kNThreads) { + float a = s_vec[i], b = s_vec[i + half]; + s_vec[i] = a + b; s_vec[i + half] = a - b; + } + __syncthreads(); + } + + float* vout = output + (size_t)vec_id * P; + for (int i = tid; i < P; i += kNThreads) vout[i] = __fmul_rn(s_vec[i], final_scale); +} + +// ============================================================================ +// Launch helpers + dispatch policy +// ============================================================================ + +/// Batch size at which the warp-per-vector kernels overtake the cooperative +/// kernels. Below this, one-block-per-vector wins by spreading a vector +/// across more threads (better latency at tiny batches); above it, warp +/// kernels win on removed barriers/smem traffic. Measured on RTX PRO 6000 +/// Blackwell: warp path is >=1.16x at N=1000 for every dim in [96, 2048]; +/// the cooperative path is faster for D >= 512 at N <= 100. Override with +/// env RABITQ_FHT_WARP_MIN_N (testing/tuning only). +constexpr int kWarpPathMinN = 1000; + +inline int warp_path_min_n() { + const char* e = std::getenv("RABITQ_FHT_WARP_MIN_N"); + return e ? std::atoi(e) : kWarpPathMinN; +} + +constexpr int kWarpKernelBlock = 256; // 8 vectors per block + +template +inline void launch_warp_rotate(const float* input, float* output, + const uint8_t* flip, int N, + float total_scale, cudaStream_t stream) { + const int vecs_per_block = kWarpKernelBlock / 32; + const int grid = (N + vecs_per_block - 1) / vecs_per_block; + fht_kac_rotate_warp_kernel<<>>( + input, output, flip, N, total_scale); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +template +inline void launch_coop_rotate(const float* input, float* output, + const uint8_t* flip, int N, + float total_scale, cudaStream_t stream) { + constexpr int kThreads = 1 << coop_log_threads(kLogN); + constexpr int kLogW = coop_log_threads(kLogN) - 5; + const int smem = (kLogW > 0) ? (1 << kLogN) * (int)sizeof(float) : 0; + auto kernel = &fht_kac_rotate_coop_kernel; + if (smem >= 48 * 1024) + RAFT_CUDA_TRY(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + kernel<<>>(input, output, flip, N, total_scale); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +template +inline void launch_coop_rotate_nonpow2(const float* input, float* output, + const uint8_t* flip, int N, + int padded_dim, float fac, + float final_scale, cudaStream_t stream) { + constexpr int kThreads = 1 << coop_log_threads(kLogTrunc); + const int smem = (padded_dim + (1 << kLogTrunc)) * (int)sizeof(float) + + 4 * (padded_dim / 8); // staged flip words + auto kernel = &fht_kac_rotate_coop_nonpow2_kernel; + if (smem >= 48 * 1024) + RAFT_CUDA_TRY(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + kernel<<>>(input, output, flip, N, padded_dim, fac, final_scale); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +inline bool warp_nonpow2_eligible(int log_trunc, int padded_dim) { + if (log_trunc < 5 || log_trunc > 10) return false; + const int T = 1 << log_trunc; + return padded_dim > T && padded_dim <= 2 * T && + padded_dim % 32 == 0 && padded_dim <= 2048; +} + +// ============================================================================ +// Runtime dispatch +// ============================================================================ + +inline void dispatch_fused_rotate(const float* input, float* output, + const uint8_t* flip, int N, int log_N, + float total_scale, cudaStream_t stream) { + if (N <= 0) return; + if (log_N >= 7 && log_N <= 11 && N >= warp_path_min_n()) { + dispatch_log_n<7, 11>(log_N, [&](auto K) { + launch_warp_rotate(input, output, flip, N, total_scale, stream); + }); + return; + } + dispatch_log_n(log_N, [&](auto K) { + launch_coop_rotate(input, output, flip, N, total_scale, stream); + }); +} + +inline void dispatch_fused_rotate_nonpow2(const float* input, float* output, + const uint8_t* flip, int N, + int log_trunc, int padded_dim, + float fac, float final_scale, + cudaStream_t stream) { + if (N <= 0) return; + if (warp_nonpow2_eligible(log_trunc, padded_dim) && N >= warp_path_min_n()) { + dispatch_log_n<5, 10>(log_trunc, [&](auto KT) { + constexpr int kLogT = KT.value; + constexpr int kMinP32 = (1 << kLogT) / 32 + 1; + constexpr int kMaxP32 = (1 << kLogT) / 16; + dispatch_p32(padded_dim / 32, [&](auto KP) { + const int vecs_per_block = kWarpKernelBlock / 32; + const int grid = (N + vecs_per_block - 1) / vecs_per_block; + fht_kac_rotate_warp_nonpow2_kernel + <<>>( + input, output, flip, N, fac, final_scale); + RAFT_CUDA_TRY(cudaGetLastError()); + }); + }); + return; + } + dispatch_log_n(log_trunc, [&](auto K) { + launch_coop_rotate_nonpow2(input, output, flip, N, + padded_dim, fac, final_scale, stream); + }); +} + +} // namespace fht + +} // namespace cuvs::preprocessing::quantize::rabitq::detail diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/pipeline.cu b/cpp/src/preprocessing/quantize/detail/rabitq/pipeline.cu new file mode 100644 index 0000000000..ba7893e538 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/pipeline.cu @@ -0,0 +1,196 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +// +// RaBitQ pipeline โ€” orchestration only. Composes the rotator launchers and the +// fused quantizer entry points; the only device code here is trivial glue +// (row padding, centroid subtraction). +// + +#include "pipeline.cuh" +#include "quantize.cuh" + +#include +#include +#include +#include + +#include + +#include + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +namespace { + +__global__ void pad_rows_kernel(const float* __restrict__ d_src, + float* __restrict__ d_dst, + int N, int dim, int D) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * D) return; + int row = idx / D; + int col = idx % D; + d_dst[idx] = (col < dim) ? d_src[row * dim + col] : 0.0f; +} + +__global__ void subtract_row_kernel(float* __restrict__ d_rows, + const float* __restrict__ d_row, + int N, int D) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= N * D) return; + d_rows[idx] -= d_row[idx % D]; +} + +/// Pad (if dim < rotator.padded_dim) and rotate d_src into d_dst. When padding +/// is needed, pads directly into d_dst and rotates in place if the rotator +/// supports it (fht_kac); otherwise stages through a temporary (matmul). +void pad_and_rotate(raft::resources const& res, + const float* d_src, size_t N, size_t dim, + rotator_ref const& rotator, float* d_dst) { + const size_t D = static_cast(rotator.padded_dim); + if (dim == D) { + rotate(res, rotator, d_src, d_dst, static_cast(N)); + return; + } + auto stream = raft::resource::get_cuda_stream(res); + constexpr int kBlock = 256; + const size_t total = N * D; + const int grid = static_cast((total + kBlock - 1) / kBlock); + if (supports_inplace_rotate(rotator)) { + pad_rows_kernel<<>>(d_src, d_dst, + static_cast(N), static_cast(dim), + static_cast(D)); + RAFT_CUDA_TRY(cudaGetLastError()); + rotate(res, rotator, d_dst, d_dst, static_cast(N)); + } else { + rmm::device_uvector d_tmp( + total, stream, raft::resource::get_workspace_resource_ref(res)); + pad_rows_kernel<<>>(d_src, d_tmp.data(), + static_cast(N), static_cast(dim), + static_cast(D)); + RAFT_CUDA_TRY(cudaGetLastError()); + rotate(res, rotator, d_tmp.data(), d_dst, static_cast(N)); + } +} + +/// d_residuals = rotate(pad(data)) [- rotate(pad(centroid))]. +/// require_rotated_c: full-factor mode always needs the rotated-centroid +/// buffer (zero-filled when there is no centroid). +void prepare_residuals(raft::resources const& res, + const float* d_data, size_t N, size_t dim, + rotator_ref const& rotator, + const float* d_centroid, float* d_rotated_centroid, + float* d_residuals, bool require_rotated_c) { + const size_t D = static_cast(rotator.padded_dim); + auto stream = raft::resource::get_cuda_stream(res); + pad_and_rotate(res, d_data, N, dim, rotator, d_residuals); + if (d_centroid != nullptr) { + pad_and_rotate(res, d_centroid, 1, dim, rotator, d_rotated_centroid); + constexpr int kBlock = 256; + const size_t total = N * D; + subtract_row_kernel<<((total + kBlock - 1) / kBlock), kBlock, 0, stream>>>( + d_residuals, d_rotated_centroid, static_cast(N), static_cast(D)); + RAFT_CUDA_TRY(cudaGetLastError()); + } else if (require_rotated_c) { + RAFT_CUDA_TRY(cudaMemsetAsync(d_rotated_centroid, 0, D * sizeof(float), stream)); + } +} + +template +void quantize_data_impl(raft::resources const& res, + const float* d_data, size_t N, size_t dim, + rotator_ref const& rotator, + const float* d_centroid, float* d_rotated_centroid, + size_t ex_bits, float const_scaling_factor, bool use_fast, + CodeT* d_total_code, float* d_delta, float* d_vl, + float* d_residuals, + int delta_mode, int coarse_samples, int fine_samples) { + prepare_residuals(res, d_data, N, dim, rotator, d_centroid, d_rotated_centroid, + d_residuals, /*require_rotated_c=*/false); + quantize_fused_on_residuals( + raft::resource::get_cuda_stream(res), + d_residuals, N, static_cast(rotator.padded_dim), ex_bits, + const_scaling_factor, use_fast, + d_total_code, d_delta, d_vl, delta_mode, coarse_samples, fine_samples); +} + +template +void quantize_data_full_impl(raft::resources const& res, + const float* d_data, size_t N, size_t dim, + rotator_ref const& rotator, + const float* d_centroid, float* d_rotated_centroid, + size_t ex_bits, float const_scaling_factor, bool use_fast, + CodeT* d_total_code, float* d_factors, + float* d_residuals) { + prepare_residuals(res, d_data, N, dim, rotator, d_centroid, d_rotated_centroid, + d_residuals, /*require_rotated_c=*/true); + quantize_full_on_residuals( + raft::resource::get_cuda_stream(res), + d_residuals, d_rotated_centroid, N, static_cast(rotator.padded_dim), ex_bits, + const_scaling_factor, use_fast, d_total_code, d_factors); +} + +} // namespace + +void quantize_data( + raft::resources const& res, + const float* d_data, size_t N, size_t dim, + rotator_ref const& rotator, + const float* d_centroid, float* d_rotated_centroid, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint16_t* d_total_code, float* d_delta, float* d_vl, + float* d_residuals, + int delta_mode, int coarse_samples, int fine_samples) { + quantize_data_impl(res, d_data, N, dim, rotator, d_centroid, d_rotated_centroid, + ex_bits, const_scaling_factor, use_fast, + d_total_code, d_delta, d_vl, d_residuals, + delta_mode, coarse_samples, fine_samples); +} + +void quantize_data( + raft::resources const& res, + const float* d_data, size_t N, size_t dim, + rotator_ref const& rotator, + const float* d_centroid, float* d_rotated_centroid, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint8_t* d_total_code, float* d_delta, float* d_vl, + float* d_residuals, + int delta_mode, int coarse_samples, int fine_samples) { + quantize_data_impl(res, d_data, N, dim, rotator, d_centroid, d_rotated_centroid, + ex_bits, const_scaling_factor, use_fast, + d_total_code, d_delta, d_vl, d_residuals, + delta_mode, coarse_samples, fine_samples); +} + +void quantize_data_full( + raft::resources const& res, + const float* d_data, size_t N, size_t dim, + rotator_ref const& rotator, + const float* d_centroid, float* d_rotated_centroid, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint16_t* d_total_code, float* d_factors, + float* d_residuals) { + quantize_data_full_impl(res, d_data, N, dim, rotator, d_centroid, d_rotated_centroid, + ex_bits, const_scaling_factor, use_fast, + d_total_code, d_factors, d_residuals); +} + +void quantize_data_full( + raft::resources const& res, + const float* d_data, size_t N, size_t dim, + rotator_ref const& rotator, + const float* d_centroid, float* d_rotated_centroid, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint8_t* d_total_code, float* d_factors, + float* d_residuals) { + quantize_data_full_impl(res, d_data, N, dim, rotator, d_centroid, d_rotated_centroid, + ex_bits, const_scaling_factor, use_fast, + d_total_code, d_factors, d_residuals); +} + +} // namespace cuvs::preprocessing::quantize::rabitq::detail diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/pipeline.cuh b/cpp/src/preprocessing/quantize/detail/rabitq/pipeline.cuh new file mode 100644 index 0000000000..3596661db1 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/pipeline.cuh @@ -0,0 +1,126 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +// +// One-call RaBitQ quantization pipeline. +// +// This layer only orchestrates the two underlying components โ€” the rotator +// (rotator.cuh) and the quantizer (quantize.cuh) โ€” plus two trivial glue +// kernels (row padding, centroid subtraction). It adds no quantization logic +// of its own; use the components directly for custom pipelines (e.g. +// externally computed residuals). +// +// Pipeline: pad(dim -> rotator.padded_dim) -> rotate -> [subtract rotated +// centroid] -> fused quantization on the rotated residuals. +// +// Design note: the centroid is rotated with a separate 1-row rotate() call +// rather than appended to the data batch (as cuVS does). Appending is only +// free when a gather/copy pass over the data exists anyway; for contiguous +// input it would cost a full extra N x D copy, while a 1-row rotate is a +// single tiny launch. +// +// The rotator state is NOT owned here: callers pass a non-owning rotator_ref +// that borrows the buffers held by the public `rabitq::rotator` struct. +// + +#ifndef RABITQ_GPU_PIPELINE_CUH +#define RABITQ_GPU_PIPELINE_CUH + +#include +#include + +#include + +#include "rotator.cuh" + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +/// Quantize raw (unrotated, unpadded) vectors end to end, emitting per-vector +/// scalar factors (delta, vl). +/// +/// d_data N x dim on device. +/// d_centroid optional (nullable) dim floats on device, UNrotated. +/// When non-null, residuals = rotate(pad(data)) - +/// rotate(pad(centroid)) and d_rotated_centroid must be +/// non-null (receives the rotated centroid, D floats). +/// When null, residuals = rotate(pad(data)) and +/// d_rotated_centroid may be null. +/// d_residuals N x D output (D = rotator.padded_dim): the rotated +/// residuals that the emitted codes refer to. +/// Remaining parameters mirror quantize_fused_on_residuals. +void quantize_data(raft::resources const& res, + const float* d_data, + size_t N, + size_t dim, + rotator_ref const& rotator, + const float* d_centroid, + float* d_rotated_centroid, + size_t ex_bits, + float const_scaling_factor, + bool use_fast, + uint16_t* d_total_code, + float* d_delta, + float* d_vl, + float* d_residuals, + int delta_mode = 0, + int coarse_samples = 64, + int fine_samples = 64); + +void quantize_data(raft::resources const& res, + const float* d_data, + size_t N, + size_t dim, + rotator_ref const& rotator, + const float* d_centroid, + float* d_rotated_centroid, + size_t ex_bits, + float const_scaling_factor, + bool use_fast, + uint8_t* d_total_code, + float* d_delta, + float* d_vl, + float* d_residuals, + int delta_mode = 0, + int coarse_samples = 64, + int fine_samples = 64); + +/// Full-factor variant: emits the (f_add, f_rescale, f_error) triplet used +/// for approximate-distance estimation. d_rotated_centroid (D floats) is +/// always required here โ€” it participates in the factor computation โ€” and is +/// zero-filled when d_centroid is null. +void quantize_data_full(raft::resources const& res, + const float* d_data, + size_t N, + size_t dim, + rotator_ref const& rotator, + const float* d_centroid, + float* d_rotated_centroid, + size_t ex_bits, + float const_scaling_factor, + bool use_fast, + uint16_t* d_total_code, + float* d_factors, + float* d_residuals); + +void quantize_data_full(raft::resources const& res, + const float* d_data, + size_t N, + size_t dim, + rotator_ref const& rotator, + const float* d_centroid, + float* d_rotated_centroid, + size_t ex_bits, + float const_scaling_factor, + bool use_fast, + uint8_t* d_total_code, + float* d_factors, + float* d_residuals); + +} // namespace cuvs::preprocessing::quantize::rabitq::detail + +#endif // RABITQ_GPU_PIPELINE_CUH diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/quantize.cu b/cpp/src/preprocessing/quantize/detail/rabitq/quantize.cu new file mode 100644 index 0000000000..6399e7de02 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/quantize.cu @@ -0,0 +1,366 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +// +// Fused RaBitQ scalar quantization kernels operating on pre-computed residuals. +// +// Extracted from IVF-RaBitQ-GPU-main/inc/gpu_index/quantizer_standalone.cu โ€” +// only the fused warp-cooperative path (sa_quantize_fused_kernel + +// sa_compute_delta_vl_kernel) is kept, plus the free-function entry points +// referenced by the benchmark. +// + +#include "quantize.cuh" +#include "rescale_search.cuh" +#include "tight_start_constants.cuh" +#include "reductions.cuh" + +#include +#include +#include +#include + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +static __device__ __forceinline__ float evaluate_rescale_sample( + const float* __restrict__ s_abs_norm, int D, int ex_bits, float t, int lane_id) +{ + constexpr float kEps = 1e-5f; + int max_code = (1 << ex_bits) - 1; + float numerator = 0.0f; + float sqr_denom = (lane_id == 0) ? static_cast(D) * 0.25f : 0.0f; + + for (int j = lane_id; j < D; j += 32) { + float val = s_abs_norm[j]; + int quantized = min(__float2int_rd(t * val + kEps), max_code); + numerator += (quantized + 0.5f) * val; + sqr_denom += quantized * quantized + quantized; + } + + numerator = blockred::warpReduceSum(numerator); + sqr_denom = blockred::warpReduceSum(sqr_denom); + + return numerator / sqrtf(sqr_denom); +} + +// --------------------------------------------------------------------------- +// Fused rescale search + quantize + factor kernel (fast and non-fast paths). +// The per-vector factors are accumulated during the final quantize loop while +// the codes are still in registers, so no second pass over the residual and +// codes is needed. kFullFactors selects the factor set: +// false: (delta, vl) via d_delta / d_vl / delta_mode +// true : (f_add, f_rescale, f_error) via d_centroid / d_factors +// --------------------------------------------------------------------------- +template +__global__ void sa_quantize_fused_kernel( + const float* __restrict__ d_residual, + CodeT* __restrict__ d_total_code, + int N, int padded_dim, int ex_bits, + float const_scaling_factor, + bool use_fast, + int n_coarse_samples, int n_fine_samples, + float* __restrict__ d_delta, + float* __restrict__ d_vl, + int delta_mode, + const float* __restrict__ d_centroid, + float* __restrict__ d_factors) +{ + constexpr int kNWarps = kBlockSize / 32; + constexpr float kEps = 1e-5f; + constexpr int kNEnum = 10; + int coarse_samples = n_coarse_samples > 1 ? n_coarse_samples : 1; + int fine_samples = n_fine_samples > 1 ? n_fine_samples : 1; + int coarse_denom = coarse_samples > 1 ? coarse_samples - 1 : 1; + int fine_denom = fine_samples > 1 ? fine_samples - 1 : 1; + + extern __shared__ char smem[]; + float* s_abs_norm = reinterpret_cast(smem); + float* s_warp_ip = s_abs_norm + padded_dim; + float* s_warp_t = s_warp_ip + kNWarps; + + int vec_id = blockIdx.x; + if (vec_id >= N) return; + + int warp_id = threadIdx.x / 32; + int lane_id = threadIdx.x % 32; + + const float* res = d_residual + (size_t)vec_id * padded_dim; + CodeT* code = d_total_code + (size_t)vec_id * padded_dim; + + // norm (fused block reduction; result in every thread) + float norm_sum[1] = {0.0f}; + for (int i = threadIdx.x; i < padded_dim; i += kBlockSize) + norm_sum[0] += res[i] * res[i]; + blockred::blockAllReduceSum(norm_sum); + float inv_norm = rsqrtf(norm_sum[0] + 1e-30f); + + float t; + + if (use_fast || ex_bits == 0) { + t = use_fast ? const_scaling_factor : 1.0f; + } else { + // cache abs(res)*inv_norm in smem + find block max. The barrier + // inside block_max_all also publishes the s_abs_norm writes to the + // whole block before the sampling loops below read them. + float local_max = 0.0f; + for (int i = threadIdx.x; i < padded_dim; i += kBlockSize) { + float val = fabsf(res[i]) * inv_norm; + s_abs_norm[i] = val; + local_max = fmaxf(local_max, val); + } + float max_o = blockred::blockAllReduceMax(local_max); + + if (max_o < kEps) { t = 1.0f; } + else { + float t_end = static_cast((1 << ex_bits) - 1 + kNEnum) / max_o; + float t_start = t_end * d_kTightStart_opt[ex_bits]; + + // coarse grid + float best_coarse_ip = 0.0f, best_coarse_t = t_start; + for (int base = 0; base < coarse_samples; base += kNWarps) { + int si = base + warp_id; + float tc = (si < coarse_samples) + ? t_start + (t_end - t_start) * si / coarse_denom : t_start; + float ip = (si < coarse_samples) + ? evaluate_rescale_sample(s_abs_norm, padded_dim, ex_bits, tc, lane_id) : 0.0f; + if (lane_id == 0 && ip > best_coarse_ip) { best_coarse_ip = ip; best_coarse_t = tc; } + } + + // Tournament winners are broadcast through a dedicated shared + // slot (never reused by the fine phase), so the pre-existing + // post-tournament barrier is the only synchronization needed โ€” + // no publish/overwrite race by construction. + __shared__ float s_best_t; + + if (lane_id == 0) { s_warp_ip[warp_id] = best_coarse_ip; s_warp_t[warp_id] = best_coarse_t; } + __syncthreads(); + if (warp_id == 0) { + float ip = (lane_id < kNWarps) ? s_warp_ip[lane_id] : -1.0f; + float tc = (lane_id < kNWarps) ? s_warp_t[lane_id] : 0.0f; + for (int s = kNWarps / 2; s > 0; s >>= 1) { + float oi = __shfl_down_sync(0xffffffff, ip, s); + float ot = __shfl_down_sync(0xffffffff, tc, s); + if (oi > ip) { ip = oi; tc = ot; } + } + if (lane_id == 0) s_best_t = tc; + } + __syncthreads(); + + float center_t = s_best_t; + float range = (t_end - t_start) / coarse_samples; + float fine_start = fmaxf(t_start, center_t - range); + float fine_end = fminf(t_end, center_t + range); + + // fine grid + float best_fine_ip = 0.0f, best_fine_t = center_t; + for (int base = 0; base < fine_samples; base += kNWarps) { + int si = base + warp_id; + float tf = (si < fine_samples) + ? fine_start + (fine_end - fine_start) * si / fine_denom : center_t; + float ip = (si < fine_samples) + ? evaluate_rescale_sample(s_abs_norm, padded_dim, ex_bits, tf, lane_id) : 0.0f; + if (lane_id == 0 && ip > best_fine_ip) { best_fine_ip = ip; best_fine_t = tf; } + } + + if (lane_id == 0) { s_warp_ip[warp_id] = best_fine_ip; s_warp_t[warp_id] = best_fine_t; } + __syncthreads(); + if (warp_id == 0) { + float ip = (lane_id < kNWarps) ? s_warp_ip[lane_id] : -1.0f; + float tf = (lane_id < kNWarps) ? s_warp_t[lane_id] : 0.0f; + for (int s = kNWarps / 2; s > 0; s >>= 1) { + float oi = __shfl_down_sync(0xffffffff, ip, s); + float ot = __shfl_down_sync(0xffffffff, tf, s); + if (oi > ip) { ip = oi; tf = ot; } + } + // Safe overwrite of s_best_t: the fine-publish barrier above + // ordered every thread's read of the coarse center before it. + if (lane_id == 0) s_best_t = tf; + } + __syncthreads(); + t = s_best_t; + } + } + + // quantize (abs + branch) + accumulate the factor sums in the same pass + int mask = (1 << ex_bits) - 1; + int offset = 1 << ex_bits; + float cb = -((float)(1 << ex_bits) - 0.5f); + + if constexpr (!kFullFactors) { + // sums = { |u+cb|^2, res . (u+cb) } + float sums[2] = {0.0f, 0.0f}; + for (int i = threadIdx.x; i < padded_dim; i += kBlockSize) { + float r = res[i]; + float abs_val = fabsf(r) * inv_norm; + int k = __float2int_rd(t * abs_val + kEps); + if (k > mask) k = mask; + int total = (r >= 0.0f) ? (k + offset) : (mask - k); + code[i] = static_cast(total); + float u = (float)total + cb; + sums[0] += u * u; + sums[1] += r * u; + } + blockred::blockReduceSum(sums); + if (threadIdx.x == 0) { + float norm_res = sqrtf(norm_sum[0]); + float norm_ucb = sqrtf(sums[0]); + float cos_sim = sums[1] / (norm_res * norm_ucb + 1e-30f); + + float ratio = norm_res / (norm_ucb + 1e-30f); + float delta; + if (delta_mode == 1) + delta = ratio / (cos_sim + 1e-30f); + else if (delta_mode == 2) + delta = ratio; + else + delta = ratio * cos_sim; + d_delta[vec_id] = delta; + d_vl[vec_id] = delta * cb; + } + } else { + // sums = { res . xu, cent . xu, |xu|^2 }; |res|^2 is norm_sum[0] + float sums[3] = {0.0f, 0.0f, 0.0f}; + for (int i = threadIdx.x; i < padded_dim; i += kBlockSize) { + float r = res[i]; + float abs_val = fabsf(r) * inv_norm; + int k = __float2int_rd(t * abs_val + kEps); + if (k > mask) k = mask; + int total = (r >= 0.0f) ? (k + offset) : (mask - k); + code[i] = static_cast(total); + float xu_cb = (float)total + cb; + sums[0] += r * xu_cb; + sums[1] += d_centroid[i] * xu_cb; + sums[2] += xu_cb * xu_cb; + } + blockred::blockReduceSum(sums); + if (threadIdx.x == 0) { + constexpr float kEpsilon = 1.9f; + float l2_sq = norm_sum[0]; + float ip_resi_xucb = sums[0]; + float ip_cent_xucb = sums[1]; + float xu_sq = sums[2]; + + float l2_norm = sqrtf(l2_sq); + float denom = ip_resi_xucb + 1e-30f; + + float f_add = l2_sq + 2.0f * l2_sq * (ip_cent_xucb / denom); + float f_rescale = -2.0f * l2_sq / denom; + + float ratio = (l2_sq * xu_sq) / (denom * denom); + float inner = fmaxf(0.0f, (ratio - 1.0f) / ((float)padded_dim - 1.0f)); + float f_error = 2.0f * l2_norm * kEpsilon * sqrtf(inner); + + d_factors[vec_id * 3 + 0] = f_add; + d_factors[vec_id * 3 + 1] = f_rescale; + d_factors[vec_id * 3 + 2] = f_error; + } + } +} + +// --------------------------------------------------------------------------- +// Launcher + free-function entry points. +// --------------------------------------------------------------------------- +template +static void launch_quantize_fused( + cudaStream_t stream, + const float* d_residual, size_t N, uint32_t padded_dim, size_t ex_bits, + float const_scaling_factor, bool use_fast, int delta_mode, + int coarse_samples, int fine_samples, + CodeT* d_total_code, float* d_delta, float* d_vl) +{ + constexpr int block = 256; + int iN = static_cast(N); + int iD = static_cast(padded_dim); + int iB = static_cast(ex_bits); + constexpr int nwarps = block / 32; + + size_t q_smem = (padded_dim + 2 * nwarps) * sizeof(float); + sa_quantize_fused_kernel<<>>( + d_residual, d_total_code, iN, iD, iB, const_scaling_factor, use_fast, + coarse_samples, fine_samples, + d_delta, d_vl, delta_mode, nullptr, nullptr); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +void quantize_fused_on_residuals( + cudaStream_t stream, + const float* d_residuals, size_t N, size_t padded_dim, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint16_t* d_total_code, float* d_delta, float* d_vl, int delta_mode, + int coarse_samples, int fine_samples) +{ + launch_quantize_fused(stream, + d_residuals, N, static_cast(padded_dim), ex_bits, + const_scaling_factor, use_fast, delta_mode, + coarse_samples, fine_samples, + d_total_code, d_delta, d_vl); +} + +void quantize_fused_on_residuals( + cudaStream_t stream, + const float* d_residuals, size_t N, size_t padded_dim, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint8_t* d_total_code, float* d_delta, float* d_vl, int delta_mode, + int coarse_samples, int fine_samples) +{ + launch_quantize_fused(stream, + d_residuals, N, static_cast(padded_dim), ex_bits, + const_scaling_factor, use_fast, delta_mode, + coarse_samples, fine_samples, + d_total_code, d_delta, d_vl); +} + +template +static void launch_quantize_full( + cudaStream_t stream, + const float* d_residual, const float* d_centroid, + size_t N, uint32_t padded_dim, size_t ex_bits, + float const_scaling_factor, bool use_fast, + CodeT* d_total_code, float* d_factors) +{ + constexpr int block = 256; + int iN = static_cast(N); + int iD = static_cast(padded_dim); + int iB = static_cast(ex_bits); + constexpr int nwarps = block / 32; + + size_t q_smem = (padded_dim + 2 * nwarps) * sizeof(float); + sa_quantize_fused_kernel<<>>( + d_residual, d_total_code, iN, iD, iB, const_scaling_factor, use_fast, + 64, 64, + nullptr, nullptr, 0, d_centroid, d_factors); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +void quantize_full_on_residuals( + cudaStream_t stream, + const float* d_residuals, const float* d_centroid, + size_t N, size_t padded_dim, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint16_t* d_total_code, float* d_factors) +{ + launch_quantize_full(stream, d_residuals, d_centroid, N, + static_cast(padded_dim), ex_bits, + const_scaling_factor, use_fast, + d_total_code, d_factors); +} + +void quantize_full_on_residuals( + cudaStream_t stream, + const float* d_residuals, const float* d_centroid, + size_t N, size_t padded_dim, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint8_t* d_total_code, float* d_factors) +{ + launch_quantize_full(stream, d_residuals, d_centroid, N, + static_cast(padded_dim), ex_bits, + const_scaling_factor, use_fast, + d_total_code, d_factors); +} + +} // namespace cuvs::preprocessing::quantize::rabitq::detail diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/quantize.cuh b/cpp/src/preprocessing/quantize/detail/rabitq/quantize.cuh new file mode 100644 index 0000000000..f454f2f6c0 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/quantize.cuh @@ -0,0 +1,78 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +// +// Standalone GPU RaBitQ scalar quantization on pre-computed residuals. +// Trimmed to just the fused-kernel free functions used by the benchmark. +// + +#ifndef RABITQ_GPU_QUANTIZER_STANDALONE_CUH +#define RABITQ_GPU_QUANTIZER_STANDALONE_CUH + +#include + +#include +#include + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +/// Quantize pre-computed, rotated residuals into RaBitQ scalar codes. +/// +/// stream : the caller owns the stream. Every one of these entry points +/// only enqueues kernel launches on it and is therefore +/// *asynchronous*: the outputs are not readable until the caller +/// synchronizes the stream (or orders subsequent work on it). +/// All input and output buffers must stay alive and unmodified +/// until that work completes. +/// d_residuals : must be N ร— padded_dim on device (rotated, zero-centroid). +/// +/// Writes N ร— padded_dim codes, plus a per-vector (delta, vl) pair. +/// +/// delta_mode: 0 = RECONSTRUCTION, 1 = UNBIASED, 2 = PLAIN. +void quantize_fused_on_residuals( + cudaStream_t stream, + const float* d_residuals, size_t N, size_t padded_dim, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint16_t* d_total_code, float* d_delta, float* d_vl, int delta_mode = 0, + int coarse_samples = 64, int fine_samples = 64); + +void quantize_fused_on_residuals( + cudaStream_t stream, + const float* d_residuals, size_t N, size_t padded_dim, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint8_t* d_total_code, float* d_delta, float* d_vl, int delta_mode = 0, + int coarse_samples = 64, int fine_samples = 64); + +/// Quantize pre-computed, rotated residuals and produce the full +/// (f_add, f_rescale, f_error) factor triplet used for approximate-distance +/// estimation during search. +/// +/// stream : the caller owns the stream; the call is asynchronous on it +/// (see quantize_fused_on_residuals above). +/// d_residuals : N ร— padded_dim on device (rotated, centroid already subtracted). +/// d_centroid : padded_dim floats on device (the rotated centroid) โ€” pass a +/// zero-filled buffer when there is no centroid. +/// d_factors : N ร— 3 floats on device โ€” stride 3, layout +/// [f_add, f_rescale, f_error] per vector. +void quantize_full_on_residuals( + cudaStream_t stream, + const float* d_residuals, const float* d_centroid, + size_t N, size_t padded_dim, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint16_t* d_total_code, float* d_factors); + +void quantize_full_on_residuals( + cudaStream_t stream, + const float* d_residuals, const float* d_centroid, + size_t N, size_t padded_dim, + size_t ex_bits, float const_scaling_factor, bool use_fast, + uint8_t* d_total_code, float* d_factors); + +} // namespace cuvs::preprocessing::quantize::rabitq::detail + +#endif // RABITQ_GPU_QUANTIZER_STANDALONE_CUH diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/reductions.cuh b/cpp/src/preprocessing/quantize/detail/rabitq/reductions.cuh new file mode 100644 index 0000000000..5980e23d49 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/reductions.cuh @@ -0,0 +1,130 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +// +// Warp/block reductions: warp shuffles + one shared-memory exchange and a +// single __syncthreads() per block-level call. Naming follows the +// cuVS/RAFT device-utility convention (camelCase; blockReduceSum leaves the +// result in thread 0 like raft::blockReduce; the AllReduce variants deliver +// it to every thread via butterfly shuffles, with no extra barrier). +// +// Requirements / contract: +// - blockDim.x is a multiple of 32, at most 1024 threads. +// - Each block-level template instantiation owns one static shared scratch +// per kernel. Call an instantiation at most once per kernel launch (or +// separate repeated calls with a __syncthreads()): a back-to-back second +// call would overwrite the scratch while stragglers of the first still +// read it. Distinct N (and Sum vs Max, thread-0 vs AllReduce) use +// distinct scratches. +// - The internal __syncthreads() also orders all shared/global writes +// issued before the call ahead of anything after it. +// + +#ifndef RABITQ_GPU_BLOCK_REDUCE_CUH +#define RABITQ_GPU_BLOCK_REDUCE_CUH + +#include + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +namespace blockred { + +/// Warp-level sum; the full warp sum lands in lane 0 (other lanes hold +/// partials), matching raft::warpReduce. +__device__ __forceinline__ float warpReduceSum(float val) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + val += __shfl_down_sync(0xffffffff, val, off); + return val; +} + +/// Sum-reduce N values across the block in one pass; on return vals[k] holds +/// the full block sum in THREAD 0 only (other threads hold partials) โ€” the +/// same result contract as raft::blockReduce / cuVS blockReduceSum. +template +__device__ __forceinline__ void blockReduceSum(float (&vals)[N]) { + __shared__ float sh[N][32]; + const int lane = threadIdx.x & 31; + const int wid = threadIdx.x >> 5; + const int nw = blockDim.x >> 5; + + #pragma unroll + for (int k = 0; k < N; ++k) { + float v = warpReduceSum(vals[k]); + if (lane == 0) sh[k][wid] = v; + } + __syncthreads(); + if (wid == 0) { + #pragma unroll + for (int k = 0; k < N; ++k) { + float v = (lane < nw) ? sh[k][lane] : 0.0f; + vals[k] = warpReduceSum(v); + } + } +} + +/// Single-value convenience routed through the N = 1 instantiation. +__device__ __forceinline__ float blockReduceSum(float val) { + float v[1] = {val}; + blockReduceSum(v); + return v[0]; +} + +/// Sum-reduce N values across the block; on return vals[k] holds the full +/// block sum in EVERY thread (butterfly all-reduce, no broadcast barrier). +template +__device__ __forceinline__ void blockAllReduceSum(float (&vals)[N]) { + __shared__ float sh[N][32]; + const int lane = threadIdx.x & 31; + const int wid = threadIdx.x >> 5; + const int nw = blockDim.x >> 5; + + #pragma unroll + for (int k = 0; k < N; ++k) { + float v = vals[k]; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_xor_sync(0xffffffff, v, off); + if (lane == 0) sh[k][wid] = v; + } + __syncthreads(); + #pragma unroll + for (int k = 0; k < N; ++k) { + float v = (lane < nw) ? sh[k][lane] : 0.0f; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_xor_sync(0xffffffff, v, off); + vals[k] = v; // every lane of every warp + } +} + +/// Max-reduce one value across the block; every thread receives the result. +__device__ __forceinline__ float blockAllReduceMax(float val) { + __shared__ float sh[32]; + const int lane = threadIdx.x & 31; + const int wid = threadIdx.x >> 5; + const int nw = blockDim.x >> 5; + const float kNegInf = __int_as_float(0xff800000); + + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, off)); + if (lane == 0) sh[wid] = val; + __syncthreads(); + float v = (lane < nw) ? sh[lane] : kNegInf; + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, off)); + return v; +} + +} // namespace blockred + +} // namespace cuvs::preprocessing::quantize::rabitq::detail + +#endif // RABITQ_GPU_BLOCK_REDUCE_CUH diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/rescale_search.cu b/cpp/src/preprocessing/quantize/detail/rabitq/rescale_search.cu new file mode 100644 index 0000000000..eecafc6df2 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/rescale_search.cu @@ -0,0 +1,301 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +// +// Implementation of the rescale-factor search helpers extracted from +// IVF-RaBitQ-GPU-main/inc/gpu_index/quantizer_gpu_fast.cu. +// + +#include "rescale_search.cuh" +#include "tight_start_constants.cuh" +#include "reductions.cuh" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +// --------------------------------------------------------------------------- +// Warp-cooperative single-sample evaluator for the rescale search. +// --------------------------------------------------------------------------- +static __device__ __forceinline__ float evaluate_rescale_sample_warp( + const float* __restrict__ s_xp_norm, int D, int EX_BITS, float t, int lane_id) +{ + constexpr float kEps = 1e-5f; + int max_code = (1 << EX_BITS) - 1; + float numerator = 0.0f; + float sqr_denom = (lane_id == 0) ? static_cast(D) * 0.25f : 0.0f; + + for (int j = lane_id; j < D; j += 32) { + float val = fabsf(s_xp_norm[j]); + int quantized = min(__float2int_rd(t * val + kEps), max_code); + numerator += (quantized + 0.5f) * val; + sqr_denom += quantized * quantized + quantized; + } + + numerator = blockred::warpReduceSum(numerator); + sqr_denom = blockred::warpReduceSum(sqr_denom); + return numerator / sqrtf(sqr_denom); +} + +// --------------------------------------------------------------------------- +// Warp-cooperative rescale search. See header for contract. +// --------------------------------------------------------------------------- +__device__ float compute_best_rescale_parallel( + float* s_xp_norm, + int D, + int EX_BITS, + float* reuse_space, + int BlockSize, + int n_coarse_samples, + int n_fine_samples) +{ + const int tid = threadIdx.x; + const int warp_id = tid >> 5; + const int lane_id = tid & 31; + const int nWarps = BlockSize / 32; + + constexpr float kEps = 1e-5f; + constexpr int kNEnum = 10; + int coarse_samples = n_coarse_samples > 1 ? n_coarse_samples : 1; + int fine_samples = n_fine_samples > 1 ? n_fine_samples : 1; + int coarse_denom = coarse_samples > 1 ? coarse_samples - 1 : 1; + int fine_denom = fine_samples > 1 ? fine_samples - 1 : 1; + + // block-wide max of |s_xp_norm| + float local_max = 0.0f; + for (int i = tid; i < D; i += BlockSize) { + local_max = fmaxf(local_max, fabsf(s_xp_norm[i])); + } + + float* s_reduce = reuse_space; + s_reduce[tid] = local_max; + __syncthreads(); + for (int stride = BlockSize / 2; stride > 0; stride >>= 1) { + if (tid < stride) s_reduce[tid] = fmaxf(s_reduce[tid], s_reduce[tid + stride]); + __syncthreads(); + } + __shared__ float max_o_shared; + if (tid == 0) max_o_shared = s_reduce[0]; + __syncthreads(); + float max_o = max_o_shared; + + if (max_o < kEps) return 1.0f; + + float t_end = static_cast((1 << EX_BITS) - 1 + kNEnum) / max_o; + float t_start = t_end * d_kTightStart_opt[EX_BITS]; + + float* s_warp_ip = reuse_space + BlockSize; + float* s_warp_t = s_warp_ip + nWarps; + + // Tournament winners are broadcast through a dedicated shared slot + // (never aliased by the workspace and never reused by the fine phase), + // so the pre-existing post-tournament barriers are the only + // synchronization needed โ€” no publish/overwrite race by construction. + __shared__ float s_best_t; + + // coarse grid search + float best_coarse_ip = 0.0f; + float best_coarse_t = t_start; + for (int base = 0; base < coarse_samples; base += nWarps) { + int si = base + warp_id; + float tc = (si < coarse_samples) + ? t_start + (t_end - t_start) * si / coarse_denom + : t_start; + float ip = (si < coarse_samples) + ? evaluate_rescale_sample_warp(s_xp_norm, D, EX_BITS, tc, lane_id) + : 0.0f; + if (lane_id == 0 && ip > best_coarse_ip) { + best_coarse_ip = ip; + best_coarse_t = tc; + } + } + + if (lane_id == 0) { + s_warp_ip[warp_id] = best_coarse_ip; + s_warp_t[warp_id] = best_coarse_t; + } + __syncthreads(); + + if (warp_id == 0) { + float ip = (lane_id < nWarps) ? s_warp_ip[lane_id] : -1.0f; + float tc = (lane_id < nWarps) ? s_warp_t[lane_id] : 0.0f; + for (int s = 16; s > 0; s >>= 1) { + float oi = __shfl_down_sync(0xffffffff, ip, s); + float ot = __shfl_down_sync(0xffffffff, tc, s); + if (oi > ip) { ip = oi; tc = ot; } + } + if (lane_id == 0) s_best_t = tc; + } + __syncthreads(); + + float center_t = s_best_t; + float range = (t_end - t_start) / coarse_samples; + float fine_start = fmaxf(t_start, center_t - range); + float fine_end = fminf(t_end, center_t + range); + + // fine grid search + float best_fine_ip = 0.0f; + float best_fine_t = center_t; + for (int base = 0; base < fine_samples; base += nWarps) { + int si = base + warp_id; + float tf = (si < fine_samples) + ? fine_start + (fine_end - fine_start) * si / fine_denom + : center_t; + float ip = (si < fine_samples) + ? evaluate_rescale_sample_warp(s_xp_norm, D, EX_BITS, tf, lane_id) + : 0.0f; + if (lane_id == 0 && ip > best_fine_ip) { + best_fine_ip = ip; + best_fine_t = tf; + } + } + + if (lane_id == 0) { + s_warp_ip[warp_id] = best_fine_ip; + s_warp_t[warp_id] = best_fine_t; + } + __syncthreads(); + + if (warp_id == 0) { + float ip = (lane_id < nWarps) ? s_warp_ip[lane_id] : -1.0f; + float tf = (lane_id < nWarps) ? s_warp_t[lane_id] : 0.0f; + for (int s = 16; s > 0; s >>= 1) { + float oi = __shfl_down_sync(0xffffffff, ip, s); + float ot = __shfl_down_sync(0xffffffff, tf, s); + if (oi > ip) { ip = oi; tf = ot; } + } + // Safe overwrite of s_best_t: the fine-publish barrier above ordered + // every thread's read of the coarse center before it. Returning the + // static slot also keeps the result immune to any workspace reuse + // in the caller. + if (lane_id == 0) s_best_t = tf; + } + __syncthreads(); + + return s_best_t; +} + +// --------------------------------------------------------------------------- +// Fully-fused kernel: generate a random Gaussian row, normalize, search for +// the optimal rescale factor. One block = one sample row. +// --------------------------------------------------------------------------- +__global__ void rabitq_rescale_sample_kernel( + float* __restrict__ output_factors, + int rows, + int cols, + int ex_bits, + unsigned long long seed, + int coarse_samples, + int fine_samples) +{ + const int row_id = blockIdx.x; + if (row_id >= rows) return; + + const int tid = threadIdx.x; + const int block_size = blockDim.x; + + extern __shared__ float shared_mem[]; + float* row_data = shared_mem; + float* reuse_space = &row_data[cols]; + + curandState rng_state; + curand_init(seed, row_id * block_size + tid, 0, &rng_state); + + for (int i = tid; i < cols; i += block_size) { + row_data[i] = curand_normal(&rng_state); + } + __syncthreads(); + + float local_sum = 0.0f; + for (int i = tid; i < cols; i += block_size) { + float val = row_data[i]; + local_sum += val * val; + } + + float norm_squared = blockred::blockReduceSum(local_sum); + + __shared__ float inv_norm; + if (tid == 0) inv_norm = rsqrtf(norm_squared); + __syncthreads(); + + for (int i = tid; i < cols; i += block_size) { + row_data[i] = fabsf(row_data[i] * inv_norm); + } + __syncthreads(); + + float rescale_factor = compute_best_rescale_parallel( + row_data, cols, ex_bits, reuse_space, block_size, + coarse_samples, fine_samples); + + if (tid == 0) output_factors[row_id] = rescale_factor; +} + +// --------------------------------------------------------------------------- +// Host entry point โ€” average the per-row factors on device. +// --------------------------------------------------------------------------- +float get_const_scaling_factor(raft::resources const& res, + size_t dim, size_t ex_bits, + uint64_t seed, + int coarse_samples, + int fine_samples) { + constexpr long kConstNum = 100; + + auto stream = raft::resource::get_cuda_stream(res); + auto workspace = raft::resource::get_workspace_resource_ref(res); + + rmm::device_uvector d_factors(static_cast(kConstNum), stream, workspace); + rmm::device_uvector d_sum(1, stream, workspace); + + int block_size = 256; + if (dim <= 512) block_size = 128; + if (dim >= 1536) block_size = 512; + + size_t shared_mem_size = (dim + 3 * block_size) * sizeof(float); + + // Device properties of the *resource's* device (cached on `res`), not device 0. + cudaDeviceProp const& prop = raft::resource::get_device_properties(res); + if (shared_mem_size > prop.sharedMemPerBlock) { + block_size = 128; + shared_mem_size = (dim + 3 * block_size) * sizeof(float); + } + + rabitq_rescale_sample_kernel<<>>( + d_factors.data(), kConstNum, static_cast(dim), static_cast(ex_bits), seed, + coarse_samples, fine_samples); + RAFT_CUDA_TRY(cudaGetLastError()); + + size_t temp_storage_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceReduce::Sum( + nullptr, temp_storage_bytes, d_factors.data(), d_sum.data(), kConstNum, stream)); + + rmm::device_uvector d_temp_storage(temp_storage_bytes, stream, workspace); + RAFT_CUDA_TRY(cub::DeviceReduce::Sum( + d_temp_storage.data(), temp_storage_bytes, d_factors.data(), d_sum.data(), kConstNum, + stream)); + + float sum; + RAFT_CUDA_TRY( + cudaMemcpyAsync(&sum, d_sum.data(), sizeof(float), cudaMemcpyDeviceToHost, stream)); + raft::resource::sync_stream(res); + + return sum / kConstNum; +} + +} // namespace cuvs::preprocessing::quantize::rabitq::detail diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/rescale_search.cuh b/cpp/src/preprocessing/quantize/detail/rabitq/rescale_search.cuh new file mode 100644 index 0000000000..578133ec3e --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/rescale_search.cuh @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +// +// Device-side helpers for the RaBitQ rescale-factor search, plus a host entry +// point that draws the constant scaling factor used in the fast-quantize path. +// +// Extracted from quantizer_gpu_fast.cu โ€” only the pieces needed by the +// standalone quantizer. +// + +#ifndef RABITQ_GPU_RESCALE_SEARCH_CUH +#define RABITQ_GPU_RESCALE_SEARCH_CUH + +#include + +#include +#include + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +/// Warp-cooperative rescale-factor search. Declared here so the standalone +/// quantizer kernel can call it via `extern __device__` linkage. +/// +/// s_xp_norm : shared-memory array of |x|/||x|| values, length D. +/// D : working dimension (padded). +/// EX_BITS : number of extended bits (total_bits - 1). +/// reuse_space: scratch array in shared memory, at least BlockSize + 2*nWarps +/// floats (BlockSize/32 == nWarps). +/// BlockSize : blockDim.x of the calling kernel. +/// +/// Returns the selected rescale factor `t` (broadcast via the return value +/// from thread 0's path; all threads read it from shared memory through the +/// kernel's own __syncthreads pattern). +__device__ float compute_best_rescale_parallel( + float* s_xp_norm, + int D, + int EX_BITS, + float* reuse_space, + int BlockSize, + int coarse_samples = 64, + int fine_samples = 64); + +/// Host: estimate the constant scaling factor used by the fast-quantize path. +/// Averages `kConstNum` rescale factors computed on random Gaussian vectors. +/// +/// All work is issued on `raft::resource::get_cuda_stream(res)` and all scratch +/// memory comes from `raft::resource::get_workspace_resource_ref(res)`; the +/// device properties used to pick the block size are read from `res` as well, +/// so the sizing follows the resource's device rather than device 0. +/// Because the estimate is returned by value, the function synchronizes that +/// stream once before returning. +float get_const_scaling_factor(raft::resources const& res, + size_t dim, size_t ex_bits, + uint64_t seed = 12345ULL, + int coarse_samples = 64, + int fine_samples = 64); + +} // namespace cuvs::preprocessing::quantize::rabitq::detail + +#endif // RABITQ_GPU_RESCALE_SEARCH_CUH diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/rotator.cu b/cpp/src/preprocessing/quantize/detail/rabitq/rotator.cu new file mode 100644 index 0000000000..4802299351 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/rotator.cu @@ -0,0 +1,216 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +// +// Rotator state generation and launchers โ€” matmul and fht_kac. +// The random orthogonal matrix is generated on the CPU via modified +// Gram-Schmidt so this file has no Eigen dependency. +// + +#include "fht.cuh" +#include "rotator.cuh" + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +namespace { + +inline int64_t rd_up(int64_t dim, int64_t mult) { return ((dim + mult - 1) / mult) * mult; } + +inline int64_t floor_log2(int64_t x) +{ + int64_t r = 0; + while (x >>= 1) { + ++r; + } + return r; +} + +// --------------------------------------------------------------------------- +// Modified Gram-Schmidt on a random Gaussian Dร—D matrix. Produces Q that is +// orthonormal in rows; equivalent to Q from a QR decomposition modulo sign +// conventions. Numerically stable enough for the few-thousand dimensions used +// by RaBitQ. Output is row-major. +// --------------------------------------------------------------------------- +std::vector random_orthogonal_matrix(int64_t D, uint64_t seed) +{ + const size_t d = static_cast(D); + std::vector M(d * d); + std::mt19937 gen(static_cast(seed)); + std::normal_distribution dist(0.0f, 1.0f); + for (size_t i = 0; i < d * d; ++i) + M[i] = dist(gen); + + // Modified Gram-Schmidt on rows (in double for stability). + std::vector row(d); + for (size_t i = 0; i < d; ++i) { + for (size_t k = 0; k < d; ++k) + row[k] = M[i * d + k]; + + for (size_t j = 0; j < i; ++j) { + double dot = 0.0; + for (size_t k = 0; k < d; ++k) + dot += row[k] * M[j * d + k]; + for (size_t k = 0; k < d; ++k) + row[k] -= dot * M[j * d + k]; + } + + double norm_sq = 0.0; + for (size_t k = 0; k < d; ++k) + norm_sq += row[k] * row[k]; + double inv_norm = 1.0 / std::sqrt(norm_sq); + for (size_t k = 0; k < d; ++k) + M[i * d + k] = static_cast(row[k] * inv_norm); + } + return M; +} + +/// Derived fht_kac parameters. Recomputed on every launch from the PADDED +/// dimension, matching the standalone rotator (see its commit "Derive the +/// FHT-Kac segment size from the padded dim": the segment size follows the +/// padded width, which is also the width the flip bits and the kernels work on). +/// A useful side effect: log_n then always lands in [5, 15], so small dims can +/// never fall below fht::dispatch_log_n's supported range. +struct fht_kac_derived { + int64_t trunc_dim; + int log_n; + float fac; +}; + +inline fht_kac_derived derive_fht_kac(int64_t D) +{ + const int64_t bottom_log = floor_log2(D); + const int64_t trunc_dim = int64_t{1} << bottom_log; + return fht_kac_derived{ + trunc_dim, static_cast(bottom_log), 1.0f / std::sqrt(static_cast(trunc_dim))}; +} + +} // namespace + +// --------------------------------------------------------------------------- +// Sizes +// --------------------------------------------------------------------------- + +int64_t padded_dim(int64_t dim) { return rd_up(dim, 32); } + +int64_t flip_bytes(int64_t D) { return 4 * D / 8; } + +// --------------------------------------------------------------------------- +// State generation +// --------------------------------------------------------------------------- + +void init_state_matmul(cudaStream_t stream, uint64_t seed, int64_t D, float* d_P_out) +{ + RAFT_EXPECTS(D > 0, "rotator: padded dimension must be positive"); + std::vector h_P = random_orthogonal_matrix(D, seed); + RAFT_CUDA_TRY(cudaMemcpyAsync( + d_P_out, h_P.data(), sizeof(float) * h_P.size(), cudaMemcpyHostToDevice, stream)); + // h_P is plain host memory and dies with this scope. + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); +} + +void init_state_fht_kac(cudaStream_t stream, uint64_t seed, int64_t D, uint8_t* d_flip_out) +{ + RAFT_EXPECTS(D > 0, "rotator: padded dimension must be positive"); + std::vector h_flip(static_cast(flip_bytes(D))); + + std::mt19937 gen(static_cast(seed)); + std::uniform_int_distribution dist(0, 255); + for (auto& b : h_flip) + b = static_cast(dist(gen)); + + RAFT_CUDA_TRY(cudaMemcpyAsync( + d_flip_out, h_flip.data(), h_flip.size(), cudaMemcpyHostToDevice, stream)); + // h_flip is plain host memory and dies with this scope. + RAFT_CUDA_TRY(cudaStreamSynchronize(stream)); +} + +// --------------------------------------------------------------------------- +// Launchers +// --------------------------------------------------------------------------- + +void rotate_matmul(raft::resources const& res, + const float* d_P, + const float* in, + float* out, + int64_t N, + int64_t D) +{ + if (N <= 0) { return; } + const float alpha = 1.0f; + const float beta = 0.0f; + auto cublas_h = raft::resource::get_cublas_handle(res); + RAFT_CUBLAS_TRY(cublasSetStream(cublas_h, raft::resource::get_cuda_stream(res))); + // Row-major row-of-A (size D) is reinterpreted as column-of-A in column-major; + // compute C = P ยท A by calling sgemm with leading dims = D in both inputs. + RAFT_CUBLAS_TRY(cublasSgemm(cublas_h, + CUBLAS_OP_N, + CUBLAS_OP_N, + static_cast(D), + static_cast(N), + static_cast(D), + &alpha, + d_P, + static_cast(D), + in, + static_cast(D), + &beta, + out, + static_cast(D))); +} + +void rotate_fht_kac(cudaStream_t stream, + const uint8_t* d_flip, + const float* in, + float* out, + int64_t N, + int64_t D) +{ + if (N <= 0) { return; } + const auto p = derive_fht_kac(D); + if (p.trunc_dim == D) { + const float total_scale = p.fac * p.fac * p.fac * p.fac; + fht::dispatch_fused_rotate( + in, out, d_flip, static_cast(N), p.log_n, total_scale, stream); + } else { + fht::dispatch_fused_rotate_nonpow2(in, + out, + d_flip, + static_cast(N), + p.log_n, + static_cast(D), + p.fac, + 0.25f, + stream); + } +} + +void rotate( + raft::resources const& res, rotator_ref const& rot, const float* in, float* out, int64_t N) +{ + if (rot.use_fht_kac) { + rotate_fht_kac( + raft::resource::get_cuda_stream(res), rot.flip_bits, in, out, N, rot.padded_dim); + } else { + rotate_matmul(res, rot.rotation_matrix, in, out, N, rot.padded_dim); + } +} + +} // namespace cuvs::preprocessing::quantize::rabitq::detail diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/rotator.cuh b/cpp/src/preprocessing/quantize/detail/rabitq/rotator.cuh new file mode 100644 index 0000000000..40e38b34a9 --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/rotator.cuh @@ -0,0 +1,116 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer. Behaviour unchanged; + * only include paths, error macros and the enclosing namespace differ. + */ + +// +// GPU rotator supporting two implementations: +// - matmul : full D x D random orthogonal matrix via cuBLAS sgemm, O(N*D^2) +// - fht_kac : Fast Hadamard Transform + Kac's walk, O(N*D*logD) +// +// This layer owns NO state. The rotation matrix / flip bits live in device +// buffers owned by the caller (the public `rabitq::rotator` struct holds them +// in raft containers); the functions below only +// (a) report the buffer sizes the caller must allocate, +// (b) generate the state into caller-provided buffers, and +// (c) launch the rotation kernels against that state. +// +// The derived fht_kac quantities (trunc_dim, fac, log_n) are recomputed from +// the padded dimension on every launch, so nothing has to be cached. +// + +#ifndef RABITQ_GPU_ROTATOR_GPU_CUH +#define RABITQ_GPU_ROTATOR_GPU_CUH + +#include + +#include +#include + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +// --------------------------------------------------------------------------- +// Sizes. The public layer allocates its raft containers from these. +// --------------------------------------------------------------------------- + +/// Padded working dimension: `dim` rounded up to a multiple of 32. Both +/// rotator kinds operate on padded rows of this width. +[[nodiscard]] int64_t padded_dim(int64_t dim); + +/// Bytes of flip-bit state the fht_kac rotator consumes at padded dimension +/// `D`: four sign passes of D bits each == 4 * D / 8 bytes. +[[nodiscard]] int64_t flip_bytes(int64_t D); + +// --------------------------------------------------------------------------- +// State generation. Writes into device buffers provided by the caller. The +// draw happens on the host (std::mt19937, identical to the standalone code), +// so both calls stage through pinned-free host memory and therefore +// synchronize `stream` before returning. +// --------------------------------------------------------------------------- + +/// Fill `d_P_out` (D * D floats, row-major) with a random orthogonal matrix. +void init_state_matmul(cudaStream_t stream, uint64_t seed, int64_t D, float* d_P_out); + +/// Fill `d_flip_out` (`flip_bytes(D)` bytes) with the Kac's-walk sign bits. +void init_state_fht_kac(cudaStream_t stream, uint64_t seed, int64_t D, uint8_t* d_flip_out); + +// --------------------------------------------------------------------------- +// Launchers. `in` and `out` are N x D row-major float matrices. +// In-place aliasing (`in == out`) is permitted for fht_kac only. +// --------------------------------------------------------------------------- + +/// out = P * in, computed with the cuBLAS handle carried by `res` (which is +/// bound to `res`'s stream). `d_P` is D * D floats, row-major. +void rotate_matmul(raft::resources const& res, + const float* d_P, + const float* in, + float* out, + int64_t N, + int64_t D); + +/// Fused FHT + Kac's-walk rotation. `d_flip` is `flip_bytes(D)` bytes. +/// +/// `D` is the padded row width, and the truncated-FHT segment size is derived +/// from it (trunc_dim = 2^floor_log2(D)) โ€” the padded width is what the flip bits +/// and the kernels operate on. The transform therefore depends only on `D`; the +/// original unpadded dimension is not needed here. +void rotate_fht_kac(cudaStream_t stream, + const uint8_t* d_flip, + const float* in, + float* out, + int64_t N, + int64_t D); + +// --------------------------------------------------------------------------- +// Non-owning shim over the two kinds, for internal callers (the pipeline) +// that are kind-agnostic. It only *borrows* the caller's state; the kind tag +// stays a plain bool here so that the detail layer does not have to duplicate +// the public `rabitq::rotator_kind` enum. +// --------------------------------------------------------------------------- +struct rotator_ref { + /// Padded row width; must equal padded_dim(dim) of the owning rotator. Both + /// kinds are fully determined by this width plus the state pointers below. + int64_t padded_dim = 0; + /// true -> fht_kac (uses `flip_bits`), false -> matmul (uses `rotation_matrix`). + bool use_fht_kac = true; + const float* rotation_matrix = nullptr; + const uint8_t* flip_bits = nullptr; +}; + +/// Only fht_kac can rotate a buffer onto itself. +[[nodiscard]] inline bool supports_inplace_rotate(rotator_ref const& rot) +{ + return rot.use_fht_kac; +} + +/// Dispatch to rotate_fht_kac / rotate_matmul according to `rot.use_fht_kac`. +/// Runs on `raft::resource::get_cuda_stream(res)` in both cases. +void rotate( + raft::resources const& res, rotator_ref const& rot, const float* in, float* out, int64_t N); + +} // namespace cuvs::preprocessing::quantize::rabitq::detail + +#endif // RABITQ_GPU_ROTATOR_GPU_CUH diff --git a/cpp/src/preprocessing/quantize/detail/rabitq/tight_start_constants.cuh b/cpp/src/preprocessing/quantize/detail/rabitq/tight_start_constants.cuh new file mode 100644 index 0000000000..b7818a6d0e --- /dev/null +++ b/cpp/src/preprocessing/quantize/detail/rabitq/tight_start_constants.cuh @@ -0,0 +1,35 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Imported from the standalone RaBitQ GPU quantizer (unchanged apart from the + * enclosing namespace). + */ + +// +// Per-TU __constant__ table used by the RaBitQ rescale-factor search. +// Defined `static` so each translation unit holding a kernel that reads it +// gets its own initialized copy โ€” avoids needing -rdc=true for cross-TU +// __constant__ linkage. +// + +#ifndef RABITQ_GPU_TIGHT_START_CONSTANTS_CUH +#define RABITQ_GPU_TIGHT_START_CONSTANTS_CUH + +namespace cuvs::preprocessing::quantize::rabitq::detail { + +static __device__ __constant__ float d_kTightStart_opt[9] = { + 0.0f, + 0.15f, + 0.20f, + 0.52f, + 0.59f, + 0.71f, + 0.75f, + 0.77f, + 0.81f, +}; + +} // namespace cuvs::preprocessing::quantize::rabitq::detail + +#endif // RABITQ_GPU_TIGHT_START_CONSTANTS_CUH diff --git a/cpp/src/preprocessing/quantize/rabitq.cu b/cpp/src/preprocessing/quantize/rabitq.cu new file mode 100644 index 0000000000..0000e9c11b --- /dev/null +++ b/cpp/src/preprocessing/quantize/rabitq.cu @@ -0,0 +1,467 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "./detail/rabitq/pipeline.cuh" +#include "./detail/rabitq/quantize.cuh" +#include "./detail/rabitq/rescale_search.cuh" +#include "./detail/rabitq/rotator.cuh" + +#include + +#include +#include +#include +#include +#include + +#include + +#include + +#include +#include +#include +#include + +namespace cuvs::preprocessing::quantize::rabitq { + +namespace { + +/** Largest `ex_bits` the detail layer's tight-start table covers. */ +constexpr uint32_t kMaxExBits = 8; + +/** Number of per-vector factors in the full-factor output mode. */ +constexpr int64_t kNumFactors = 3; + +/** Borrow the public rotator's buffers for the detail launchers. */ +[[nodiscard]] auto as_ref(rabitq::rotator const& rot) -> detail::rotator_ref +{ + detail::rotator_ref ref; + ref.padded_dim = rot.padded_dim; + ref.use_fht_kac = rot.kind == rotator_kind::fht_kac; + ref.rotation_matrix = ref.use_fht_kac ? nullptr : rot.rotation_matrix.data_handle(); + ref.flip_bits = ref.use_fht_kac ? rot.flip_bits.data_handle() : nullptr; + return ref; +} + +void check_params(params const& p) +{ + RAFT_EXPECTS(p.ex_bits <= kMaxExBits, "rabitq: params.ex_bits must be in [0, 8]"); + RAFT_EXPECTS(p.coarse_samples >= 1 && p.fine_samples >= 1, + "rabitq: params.coarse_samples and params.fine_samples must be >= 1"); +} + +void check_rotator(rabitq::rotator const& rot) +{ + RAFT_EXPECTS(rot.padded_dim > 0, + "rabitq: rotator is not initialized (padded_dim == 0); use make_rotator()"); + RAFT_EXPECTS(rot.padded_dim % 32 == 0, "rabitq: rotator.padded_dim must be a multiple of 32"); + if (rot.kind == rotator_kind::matmul) { + RAFT_EXPECTS(rot.rotation_matrix.extent(0) == rot.padded_dim && + rot.rotation_matrix.extent(1) == rot.padded_dim, + "rabitq: rotator.rotation_matrix must be padded_dim x padded_dim"); + } else { + RAFT_EXPECTS(rot.flip_bits.extent(0) == detail::flip_bytes(rot.padded_dim), + "rabitq: rotator.flip_bits must hold flip_bytes(padded_dim) bytes"); + } +} + +void check_quantizer(rabitq::quantizer const& q) +{ + RAFT_EXPECTS(q.dim > 0, "rabitq: quantizer is not initialized (dim == 0); use make_quantizer()"); + RAFT_EXPECTS(q.padded_dim == detail::padded_dim(q.dim), + "rabitq: quantizer.padded_dim is inconsistent with quantizer.dim"); + check_params(q.params_quantizer); +} + +/** The quantizer and the rotator must agree on the padded working dimension. */ +void check_peers(rabitq::quantizer const& q, rabitq::rotator const& rot) +{ + check_quantizer(q); + check_rotator(rot); + RAFT_EXPECTS(rot.padded_dim == q.padded_dim, + "rabitq: rotator.padded_dim must match quantizer.padded_dim"); +} + +/** One code element per padded dimension, each holding an `ex_bits + 1` bit value. */ +template +void check_codes(rabitq::quantizer const& q, + int64_t n_rows, + raft::device_matrix_view codes) +{ + constexpr uint32_t kCodeBits = sizeof(CodeT) * 8; + RAFT_EXPECTS(q.params_quantizer.ex_bits + 1 <= kCodeBits, + "rabitq: the code element type is too narrow for ex_bits + 1 bits"); + RAFT_EXPECTS(codes.extent(0) == n_rows, "rabitq: codes.extent(0) must match the row count"); + RAFT_EXPECTS(codes.extent(1) == q.padded_dim, + "rabitq: codes.extent(1) must equal quantizer.padded_dim"); +} + +void check_rows(int64_t n_rows) +{ + RAFT_EXPECTS(n_rows >= 0 && n_rows <= static_cast(std::numeric_limits::max()), + "rabitq: the number of rows passed in a single call must fit in a 32-bit int; " + "process larger datasets in batches"); +} + +void check_delta_vl(int64_t n_rows, + raft::device_vector_view delta, + raft::device_vector_view vl) +{ + RAFT_EXPECTS(delta.extent(0) == n_rows && vl.extent(0) == n_rows, + "rabitq: delta and vl must hold one element per row"); +} + +void check_factors(int64_t n_rows, + raft::device_matrix_view factors) +{ + RAFT_EXPECTS(factors.extent(0) == n_rows && factors.extent(1) == kNumFactors, + "rabitq: factors must be a n_rows x 3 matrix"); +} + +template +void transform_impl(raft::resources const& res, + rabitq::quantizer const& q, + rabitq::rotator const& rot, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl) +{ + const int64_t n_rows = dataset.extent(0); + check_peers(q, rot); + check_rows(n_rows); + RAFT_EXPECTS(dataset.extent(1) == q.dim, + "rabitq::transform: dataset.extent(1) must equal quantizer.dim"); + check_codes(q, n_rows, codes); + check_delta_vl(n_rows, delta, vl); + if (centroid.has_value()) { + RAFT_EXPECTS(centroid->extent(0) == q.dim, + "rabitq::transform: centroid must hold quantizer.dim elements"); + } + if (n_rows == 0) { return; } + + auto stream = raft::resource::get_cuda_stream(res); + auto workspace = raft::resource::get_workspace_resource_ref(res); + rmm::device_uvector residuals( + static_cast(n_rows) * static_cast(q.padded_dim), stream, workspace); + rmm::device_uvector rotated_centroid( + centroid.has_value() ? static_cast(q.padded_dim) : 0, stream, workspace); + + detail::quantize_data(res, + dataset.data_handle(), + static_cast(n_rows), + static_cast(q.dim), + as_ref(rot), + centroid.has_value() ? centroid->data_handle() : nullptr, + centroid.has_value() ? rotated_centroid.data() : nullptr, + static_cast(q.params_quantizer.ex_bits), + q.const_scaling_factor, + q.params_quantizer.use_fast, + codes.data_handle(), + delta.data_handle(), + vl.data_handle(), + residuals.data(), + static_cast(q.params_quantizer.delta_mode), + static_cast(q.params_quantizer.coarse_samples), + static_cast(q.params_quantizer.fine_samples)); +} + +template +void transform_full_impl(raft::resources const& res, + rabitq::quantizer const& q, + rabitq::rotator const& rot, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors) +{ + const int64_t n_rows = dataset.extent(0); + check_peers(q, rot); + check_rows(n_rows); + RAFT_EXPECTS(dataset.extent(1) == q.dim, + "rabitq::transform: dataset.extent(1) must equal quantizer.dim"); + check_codes(q, n_rows, codes); + check_factors(n_rows, factors); + if (centroid.has_value()) { + RAFT_EXPECTS(centroid->extent(0) == q.dim, + "rabitq::transform: centroid must hold quantizer.dim elements"); + } + if (n_rows == 0) { return; } + + auto stream = raft::resource::get_cuda_stream(res); + auto workspace = raft::resource::get_workspace_resource_ref(res); + rmm::device_uvector residuals( + static_cast(n_rows) * static_cast(q.padded_dim), stream, workspace); + // The full-factor formula always reads the rotated centroid; the detail layer + // zero-fills this buffer when no centroid is given. + rmm::device_uvector rotated_centroid( + static_cast(q.padded_dim), stream, workspace); + + detail::quantize_data_full(res, + dataset.data_handle(), + static_cast(n_rows), + static_cast(q.dim), + as_ref(rot), + centroid.has_value() ? centroid->data_handle() : nullptr, + rotated_centroid.data(), + static_cast(q.params_quantizer.ex_bits), + q.const_scaling_factor, + q.params_quantizer.use_fast, + codes.data_handle(), + factors.data_handle(), + residuals.data()); +} + +template +void transform_residuals_impl( + raft::resources const& res, + rabitq::quantizer const& q, + raft::device_matrix_view residuals, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl) +{ + const int64_t n_rows = residuals.extent(0); + check_quantizer(q); + check_rows(n_rows); + RAFT_EXPECTS(residuals.extent(1) == q.padded_dim, + "rabitq::transform_residuals: residuals.extent(1) must equal quantizer.padded_dim"); + check_codes(q, n_rows, codes); + check_delta_vl(n_rows, delta, vl); + if (n_rows == 0) { return; } + + detail::quantize_fused_on_residuals(raft::resource::get_cuda_stream(res), + residuals.data_handle(), + static_cast(n_rows), + static_cast(q.padded_dim), + static_cast(q.params_quantizer.ex_bits), + q.const_scaling_factor, + q.params_quantizer.use_fast, + codes.data_handle(), + delta.data_handle(), + vl.data_handle(), + static_cast(q.params_quantizer.delta_mode), + static_cast(q.params_quantizer.coarse_samples), + static_cast(q.params_quantizer.fine_samples)); +} + +template +void transform_residuals_full_impl( + raft::resources const& res, + rabitq::quantizer const& q, + raft::device_matrix_view residuals, + std::optional> rotated_centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors) +{ + const int64_t n_rows = residuals.extent(0); + check_quantizer(q); + check_rows(n_rows); + RAFT_EXPECTS(residuals.extent(1) == q.padded_dim, + "rabitq::transform_residuals: residuals.extent(1) must equal quantizer.padded_dim"); + check_codes(q, n_rows, codes); + check_factors(n_rows, factors); + if (rotated_centroid.has_value()) { + RAFT_EXPECTS( + rotated_centroid->extent(0) == q.padded_dim, + "rabitq::transform_residuals: rotated_centroid must hold quantizer.padded_dim elements"); + } + if (n_rows == 0) { return; } + + auto stream = raft::resource::get_cuda_stream(res); + auto workspace = raft::resource::get_workspace_resource_ref(res); + // The kernel always reads a centroid; stand in with zeros when there is none. + rmm::device_uvector zero_centroid( + rotated_centroid.has_value() ? 0 : static_cast(q.padded_dim), stream, workspace); + const float* d_centroid = nullptr; + if (rotated_centroid.has_value()) { + d_centroid = rotated_centroid->data_handle(); + } else { + RAFT_CUDA_TRY(cudaMemsetAsync( + zero_centroid.data(), 0, static_cast(q.padded_dim) * sizeof(float), stream)); + d_centroid = zero_centroid.data(); + } + + detail::quantize_full_on_residuals(stream, + residuals.data_handle(), + d_centroid, + static_cast(n_rows), + static_cast(q.padded_dim), + static_cast(q.params_quantizer.ex_bits), + q.const_scaling_factor, + q.params_quantizer.use_fast, + codes.data_handle(), + factors.data_handle()); +} + +} // namespace + +rotator make_rotator(raft::resources const& res, params const& params, int64_t dim) +{ + RAFT_EXPECTS(dim > 0, "rabitq::make_rotator: dim must be positive"); + check_params(params); + + rotator rot{res}; + rot.kind = params.rotator; + rot.padded_dim = detail::padded_dim(dim); + + auto stream = raft::resource::get_cuda_stream(res); + if (rot.kind == rotator_kind::matmul) { + rot.rotation_matrix = raft::make_device_matrix( + res, rot.padded_dim, rot.padded_dim); + detail::init_state_matmul( + stream, params.seed, rot.padded_dim, rot.rotation_matrix.data_handle()); + } else { + rot.flip_bits = + raft::make_device_vector(res, detail::flip_bytes(rot.padded_dim)); + detail::init_state_fht_kac(stream, params.seed, rot.padded_dim, rot.flip_bits.data_handle()); + } + return rot; +} + +// RaBitQ needs no training data: `res` only supplies the stream, the workspace +// memory resource and the device properties used by the scaling-factor +// estimator below (and is unused when `params.use_fast` is false). +quantizer make_quantizer(raft::resources const& res, params const& params, int64_t dim) +{ + RAFT_EXPECTS(dim > 0, "rabitq::make_quantizer: dim must be positive"); + check_params(params); + + quantizer q; + q.params_quantizer = params; + q.dim = dim; + q.padded_dim = detail::padded_dim(dim); + // RaBitQ is data-oblivious: the only thing to "learn" is the shared scaling + // factor of the fast path, and it is estimated from random probe vectors. + if (params.use_fast) { + q.const_scaling_factor = + detail::get_const_scaling_factor(res, + static_cast(q.padded_dim), + static_cast(params.ex_bits), + params.seed, + static_cast(params.coarse_samples), + static_cast(params.fine_samples)); + } + return q; +} + +pipeline make_pipeline(raft::resources const& res, params const& params, int64_t dim) +{ + pipeline p{res}; + p.rotator = make_rotator(res, params, dim); + p.quantizer = make_quantizer(res, params, dim); + return p; +} + +void rotate(raft::resources const& res, + rotator const& rotator, + raft::device_matrix_view in, + raft::device_matrix_view out) +{ + check_rotator(rotator); + const int64_t n_rows = in.extent(0); + check_rows(n_rows); + RAFT_EXPECTS(out.extent(0) == n_rows, "rabitq::rotate: in and out must have the same row count"); + RAFT_EXPECTS(in.extent(1) == rotator.padded_dim && out.extent(1) == rotator.padded_dim, + "rabitq::rotate: in and out must be n_rows x rotator.padded_dim (pad the rows " + "yourself; rotate() does not pad)"); + auto ref = as_ref(rotator); + RAFT_EXPECTS(in.data_handle() != out.data_handle() || detail::supports_inplace_rotate(ref), + "rabitq::rotate: in-place rotation is only supported by rotator_kind::fht_kac"); + if (n_rows == 0) { return; } + detail::rotate(res, ref, in.data_handle(), out.data_handle(), n_rows); +} + +void transform(raft::resources const& res, + quantizer const& quantizer, + rotator const& rotator, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl) +{ + transform_impl(res, quantizer, rotator, dataset, centroid, codes, delta, vl); +} + +void transform(raft::resources const& res, + quantizer const& quantizer, + rotator const& rotator, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl) +{ + transform_impl(res, quantizer, rotator, dataset, centroid, codes, delta, vl); +} + +void transform(raft::resources const& res, + quantizer const& quantizer, + rotator const& rotator, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors) +{ + transform_full_impl(res, quantizer, rotator, dataset, centroid, codes, factors); +} + +void transform(raft::resources const& res, + quantizer const& quantizer, + rotator const& rotator, + raft::device_matrix_view dataset, + std::optional> centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors) +{ + transform_full_impl(res, quantizer, rotator, dataset, centroid, codes, factors); +} + +void transform_residuals(raft::resources const& res, + quantizer const& quantizer, + raft::device_matrix_view residuals, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl) +{ + transform_residuals_impl(res, quantizer, residuals, codes, delta, vl); +} + +void transform_residuals(raft::resources const& res, + quantizer const& quantizer, + raft::device_matrix_view residuals, + raft::device_matrix_view codes, + raft::device_vector_view delta, + raft::device_vector_view vl) +{ + transform_residuals_impl(res, quantizer, residuals, codes, delta, vl); +} + +void transform_residuals( + raft::resources const& res, + quantizer const& quantizer, + raft::device_matrix_view residuals, + std::optional> rotated_centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors) +{ + transform_residuals_full_impl(res, quantizer, residuals, rotated_centroid, codes, factors); +} + +void transform_residuals( + raft::resources const& res, + quantizer const& quantizer, + raft::device_matrix_view residuals, + std::optional> rotated_centroid, + raft::device_matrix_view codes, + raft::device_matrix_view factors) +{ + transform_residuals_full_impl(res, quantizer, residuals, rotated_centroid, codes, factors); +} + +} // namespace cuvs::preprocessing::quantize::rabitq diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index c0decb2243..1ff92f6c08 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -409,6 +409,7 @@ ConfigureTest( preprocessing/binary_quantization.cu preprocessing/spectral_embedding.cu preprocessing/product_quantization.cu + preprocessing/rabitq_quantization.cu preprocessing/pca.cu GPUS 1 PERCENT 100 diff --git a/cpp/tests/preprocessing/rabitq_quantization.cu b/cpp/tests/preprocessing/rabitq_quantization.cu new file mode 100644 index 0000000000..be6b833a3d --- /dev/null +++ b/cpp/tests/preprocessing/rabitq_quantization.cu @@ -0,0 +1,710 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../test_utils.cuh" + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +/* + * RaBitQ is lossy and exposes no inverse_transform, so nothing below checks a + * round trip. What is checked instead: + * - the shape / consistency contract of the three factories, and that the + * rotator state is non-degenerate and reproducible from `params::seed`; + * - that rotate() is an orthogonal transform (norm preserving) - a cheap but + * strong numeric check that covers both rotator kinds and both FHT paths + * (power-of-two and non-power-of-two padded dims); + * - that transform() emits in-range, not-all-zero codes and finite factors; + * - that transform() is deterministic; + * - that transform() agrees with rotate() + transform_residuals() whenever no + * zero padding is involved (dim % 32 == 0). + */ + +namespace cuvs::preprocessing::quantize::rabitq { + +/** Mirrors detail::padded_dim(): the working dimension is `dim` rounded up to 32. */ +constexpr int64_t expected_padded_dim(int64_t dim) { return ((dim + 31) / 32) * 32; } + +/** Mirrors detail::flip_bytes(): four sign-flip passes of `padded_dim` bits each. */ +constexpr int64_t expected_flip_bytes(int64_t padded_dim) { return 4 * padded_dim / 8; } + +struct RabitqQuantizationInputs { + int rows; + int dim; + rotator_kind rotator_type; + uint32_t ex_bits; + bool use_fast; +}; + +std::ostream& operator<<(std::ostream& os, const RabitqQuantizationInputs& inputs) +{ + return os << "rows:" << inputs.rows << " dim:" << inputs.dim + << " padded_dim:" << expected_padded_dim(inputs.dim) << " rotator:" + << (inputs.rotator_type == rotator_kind::matmul ? "matmul" : "fht_kac") + << " ex_bits:" << inputs.ex_bits << " use_fast:" << inputs.use_fast; +} + +template +class RabitqQuantizationTest : public ::testing::TestWithParam { + public: + RabitqQuantizationTest() + : ps_(::testing::TestWithParam::GetParam()), + stream_(raft::resource::get_cuda_stream(handle_)), + dataset_(raft::make_device_matrix(handle_, ps_.rows, ps_.dim)), + centroid_(raft::make_device_vector(handle_, ps_.dim)) + { + } + + protected: + void SetUp() override + { + raft::random::RngState r(1234ULL); + raft::random::uniform(handle_, r, dataset_.data_handle(), dataset_.size(), -1.0f, 1.0f); + raft::random::uniform(handle_, r, centroid_.data_handle(), centroid_.size(), -0.5f, 0.5f); + raft::resource::sync_stream(handle_, stream_); + } + + // ------------------------------------------------------------------------- + // helpers + // ------------------------------------------------------------------------- + + params make_params() const + { + params p; + p.ex_bits = ps_.ex_bits; + p.rotator = ps_.rotator_type; + p.seed = 137ULL; + p.delta_mode = delta_kind::reconstruction; + p.use_fast = ps_.use_fast; + // The per-vector rescale search only runs when use_fast is false; keep it cheap. + p.coarse_samples = 16; + p.fine_samples = 16; + return p; + } + + template + std::vector to_host(const T* d_ptr, size_t len) + { + std::vector h(len); + raft::update_host(h.data(), d_ptr, len, stream_); + RAFT_CUDA_TRY(cudaStreamSynchronize(stream_)); + return h; + } + + static std::vector row_norms(const std::vector& h, int64_t rows, int64_t cols) + { + std::vector norms(rows); + for (int64_t i = 0; i < rows; ++i) { + double sum = 0.0; + for (int64_t j = 0; j < cols; ++j) { + const double v = h[static_cast(i) * static_cast(cols) + j]; + sum += v * v; + } + norms[i] = std::sqrt(sum); + } + return norms; + } + + /** The two peers must agree on the documented padding rule. */ + void verify_shape(rotator const& rot, quantizer const& quant) + { + const int64_t padded = expected_padded_dim(ps_.dim); + ASSERT_EQ(quant.dim, ps_.dim); + ASSERT_EQ(quant.padded_dim, padded); + ASSERT_EQ(rot.padded_dim, padded); + ASSERT_EQ(padded % 32, 0); + ASSERT_GE(padded, ps_.dim); + ASSERT_LT(padded - ps_.dim, 32); + ASSERT_EQ(static_cast(rot.kind), static_cast(ps_.rotator_type)); + ASSERT_EQ(quant.params_quantizer.ex_bits, ps_.ex_bits); + if (ps_.use_fast) { + ASSERT_TRUE(std::isfinite(quant.const_scaling_factor)); + ASSERT_GT(quant.const_scaling_factor, 0.0f); + } else { + // Unused on the slow path; the kernel substitutes 1.0 itself. + ASSERT_EQ(quant.const_scaling_factor, 1.0f); + } + } + + /** The dense rotation matrix must be padded_dim x padded_dim and orthonormal. */ + void verify_matmul_state(rotator const& rot) + { + const int64_t padded = rot.padded_dim; + ASSERT_EQ(rot.rotation_matrix.extent(0), padded); + ASSERT_EQ(rot.rotation_matrix.extent(1), padded); + ASSERT_EQ(rot.flip_bits.extent(0), 0); + + auto h = to_host(rot.rotation_matrix.data_handle(), + static_cast(padded) * static_cast(padded)); + // Spot-check the leading rows: unit norm and mutually orthogonal. + const int64_t n_check = std::min(padded, 8); + for (int64_t i = 0; i < n_check; ++i) { + for (int64_t j = i; j < n_check; ++j) { + double dot = 0.0; + for (int64_t k = 0; k < padded; ++k) { + dot += double(h[static_cast(i) * padded + k]) * + double(h[static_cast(j) * padded + k]); + } + ASSERT_NEAR(dot, (i == j) ? 1.0 : 0.0, 1e-4) + << "rotation_matrix rows " << i << " and " << j << " are not orthonormal"; + } + } + } + + /** The packed Kac's-walk sign bits must have the documented size and be mixed. */ + void verify_fht_kac_state(rotator const& rot) + { + const int64_t n_bytes = expected_flip_bytes(rot.padded_dim); + ASSERT_EQ(rot.flip_bits.extent(0), n_bytes); + ASSERT_EQ(rot.rotation_matrix.extent(0), 0); + ASSERT_EQ(rot.rotation_matrix.extent(1), 0); + + auto h = to_host(rot.flip_bits.data_handle(), static_cast(n_bytes)); + size_t n_set_bits = 0; + for (auto byte : h) { + for (int k = 0; k < 8; ++k) { + n_set_bits += (byte >> k) & 1u; + } + } + // An all-zero or all-one flip pattern would degenerate the random rotation + // into a plain Hadamard transform. + ASSERT_GT(n_set_bits, 0u); + ASSERT_LT(n_set_bits, static_cast(n_bytes) * 8); + } + + void verify_state(rotator const& rot) + { + if (ps_.rotator_type == rotator_kind::matmul) { + verify_matmul_state(rot); + } else { + verify_fht_kac_state(rot); + } + } + + /** Codes hold `ex_bits + 1` bit unsigned values and must not be all zero. */ + std::vector verify_codes(const CodeT* d_codes, size_t len) + { + auto h = to_host(d_codes, len); + const uint32_t max_code = (1u << (ps_.ex_bits + 1)) - 1u; + size_t n_out_of_range = 0; + bool all_zero = true; + for (size_t i = 0; i < len; ++i) { + if (static_cast(h[i]) > max_code) { ++n_out_of_range; } + if (h[i] != CodeT{0}) { all_zero = false; } + } + EXPECT_EQ(n_out_of_range, 0u) << n_out_of_range << " of " << len << " codes exceed the " + << (ps_.ex_bits + 1) << "-bit range [0, " << max_code << "]"; + EXPECT_FALSE(all_zero) << "quantized codes are all zero"; + return h; + } + + void verify_finite(const float* d_ptr, size_t len, const char* what, bool positive = false) + { + auto h = to_host(d_ptr, len); + size_t n_bad = 0; + for (size_t i = 0; i < len; ++i) { + if (!std::isfinite(h[i]) || (positive && !(h[i] > 0.0f))) { ++n_bad; } + } + EXPECT_EQ(n_bad, 0u) << n_bad << " of " << len << " " << what << " values are not finite" + << (positive ? " and strictly positive" : ""); + } + + /** vl is documented as `delta * -(2^ex_bits - 0.5)`. */ + void verify_vl(const float* d_delta, const float* d_vl, int64_t rows) + { + auto h_delta = to_host(d_delta, static_cast(rows)); + auto h_vl = to_host(d_vl, static_cast(rows)); + const float cb = -(static_cast(1u << ps_.ex_bits) - 0.5f); + for (int64_t i = 0; i < rows; ++i) { + const float expected = h_delta[i] * cb; + ASSERT_NEAR(h_vl[i], expected, 1e-5f * std::fabs(expected) + 1e-20f) << " row " << i; + } + } + + void expect_same_codes(std::vector const& a, + std::vector const& b, + const char* what) + { + ASSERT_EQ(a.size(), b.size()); + size_t n_diff = 0; + for (size_t i = 0; i < a.size(); ++i) { + if (a[i] != b[i]) { ++n_diff; } + } + EXPECT_EQ(n_diff, 0u) << what << ": " << n_diff << " of " << a.size() << " codes differ"; + } + + // ------------------------------------------------------------------------- + // 1. factories: consistent padded_dim, non-degenerate + reproducible state + // ------------------------------------------------------------------------- + void testInit() + { + const auto p = make_params(); + + auto pipe = make_pipeline(handle_, p, ps_.dim); + verify_shape(pipe.rotator, pipe.quantizer); + verify_state(pipe.rotator); + + // The peers built on their own must be interchangeable with the pipeline's: + // that is what makes the header's "store the seed and re-run make_rotator" + // recipe safe. + auto rot = make_rotator(handle_, p, ps_.dim); + auto quant = make_quantizer(handle_, p, ps_.dim); + verify_shape(rot, quant); + verify_state(rot); + ASSERT_EQ(rot.padded_dim, pipe.quantizer.padded_dim); + ASSERT_NEAR(quant.const_scaling_factor, + pipe.quantizer.const_scaling_factor, + 1e-6f * std::fabs(pipe.quantizer.const_scaling_factor) + 1e-20f); + if (ps_.rotator_type == rotator_kind::matmul) { + const size_t n = static_cast(rot.padded_dim) * static_cast(rot.padded_dim); + ASSERT_TRUE(devArrMatch(pipe.rotator.rotation_matrix.data_handle(), + rot.rotation_matrix.data_handle(), + n, + cuvs::Compare(), + stream_)); + } else { + ASSERT_TRUE(devArrMatch(pipe.rotator.flip_bits.data_handle(), + rot.flip_bits.data_handle(), + static_cast(expected_flip_bytes(rot.padded_dim)), + cuvs::Compare(), + stream_)); + } + + // dim must be positive, and ex_bits must stay inside the tight-start table. + EXPECT_THROW(make_rotator(handle_, p, 0), raft::logic_error); + EXPECT_THROW(make_quantizer(handle_, p, -1), raft::logic_error); + auto bad_params = p; + bad_params.ex_bits = 9; + EXPECT_THROW(make_quantizer(handle_, bad_params, ps_.dim), raft::logic_error); + } + + // ------------------------------------------------------------------------- + // 2. rotate() is an orthogonal transform + // ------------------------------------------------------------------------- + void testRotate() + { + const auto p = make_params(); + auto rot = make_rotator(handle_, p, ps_.dim); + const int64_t pd = rot.padded_dim; + const int64_t n = ps_.rows; + + auto in = raft::make_device_matrix(handle_, n, pd); + auto out = raft::make_device_matrix(handle_, n, pd); + raft::random::RngState r(4242ULL); + raft::random::uniform(handle_, r, in.data_handle(), in.size(), -1.0f, 1.0f); + RAFT_CUDA_TRY(cudaMemsetAsync(out.data_handle(), 0, out.size() * sizeof(float), stream_)); + raft::resource::sync_stream(handle_, stream_); + + rotate(handle_, rot, raft::make_const_mdspan(in.view()), out.view()); + raft::resource::sync_stream(handle_, stream_); + + auto h_in = to_host(in.data_handle(), in.size()); + auto h_out = to_host(out.data_handle(), out.size()); + auto norms_in = row_norms(h_in, n, pd); + auto norms_out = row_norms(h_out, n, pd); + for (int64_t i = 0; i < n; ++i) { + ASSERT_GT(norms_in[i], 0.0) << " row " << i; + // Orthogonal transform: ||Rx|| == ||x||. The tolerance is float32 slack + // over pd accumulations, not a modelling allowance. + ASSERT_NEAR(norms_out[i], norms_in[i], 1e-3 * norms_in[i]) << " row " << i; + } + + // ...and it must actually mix the coordinates rather than copy them. + size_t n_same = 0; + for (size_t i = 0; i < h_in.size(); ++i) { + if (h_in[i] == h_out[i]) { ++n_same; } + } + EXPECT_LT(n_same, h_in.size() / 2) << "rotate() left most coordinates untouched"; + + if (ps_.rotator_type == rotator_kind::fht_kac) { + // In-place rotation is documented as supported for fht_kac only. + auto inplace = raft::make_device_matrix(handle_, n, pd); + raft::copy(inplace.data_handle(), in.data_handle(), in.size(), stream_); + raft::resource::sync_stream(handle_, stream_); + rotate(handle_, rot, raft::make_const_mdspan(inplace.view()), inplace.view()); + raft::resource::sync_stream(handle_, stream_); + ASSERT_TRUE(devArrMatch(out.data_handle(), + inplace.data_handle(), + out.size(), + cuvs::CompareApprox(1e-5f), + stream_)); + } else { + EXPECT_THROW(rotate(handle_, rot, raft::make_const_mdspan(in.view()), in.view()), + raft::logic_error); + } + + // rotate() does not pad: an unpadded view must be rejected. + if (ps_.dim != pd) { + auto unpadded = raft::make_device_matrix_view( + (const float*)in.data_handle(), n, static_cast(ps_.dim)); + auto unpadded_out = raft::make_device_matrix_view( + out.data_handle(), n, static_cast(ps_.dim)); + EXPECT_THROW(rotate(handle_, rot, unpadded, unpadded_out), raft::logic_error); + } + } + + // ------------------------------------------------------------------------- + // 3 + 4. the codes / factors are sane, and transform() is deterministic + // ------------------------------------------------------------------------- + void testTransform() + { + const auto p = make_params(); + auto pipe = make_pipeline(handle_, p, ps_.dim); + const int64_t pd = pipe.quantizer.padded_dim; + const int64_t n = ps_.rows; + const size_t len = static_cast(n) * static_cast(pd); + auto data = raft::make_const_mdspan(dataset_.view()); + + auto codes_a = raft::make_device_matrix(handle_, n, pd); + auto codes_b = raft::make_device_matrix(handle_, n, pd); + auto delta_a = raft::make_device_vector(handle_, n); + auto delta_b = raft::make_device_vector(handle_, n); + auto vl_a = raft::make_device_vector(handle_, n); + auto vl_b = raft::make_device_vector(handle_, n); + + // --- (codes, delta, vl) output mode, no centroid --- + transform(handle_, + pipe.quantizer, + pipe.rotator, + data, + std::nullopt, + codes_a.view(), + delta_a.view(), + vl_a.view()); + transform(handle_, + pipe.quantizer, + pipe.rotator, + data, + std::nullopt, + codes_b.view(), + delta_b.view(), + vl_b.view()); + raft::resource::sync_stream(handle_, stream_); + + auto h_codes_a = verify_codes(codes_a.data_handle(), len); + auto h_codes_b = verify_codes(codes_b.data_handle(), len); + verify_finite(delta_a.data_handle(), static_cast(n), "delta", /* positive */ true); + verify_finite(vl_a.data_handle(), static_cast(n), "vl"); + verify_vl(delta_a.data_handle(), vl_a.data_handle(), n); + + // determinism: the same pipeline on the same input -> identical output + expect_same_codes(h_codes_a, h_codes_b, "transform() (delta, vl) is not deterministic"); + ASSERT_TRUE(devArrMatch(delta_a.data_handle(), + delta_b.data_handle(), + static_cast(n), + cuvs::Compare(), + stream_)); + ASSERT_TRUE(devArrMatch(vl_a.data_handle(), + vl_b.data_handle(), + static_cast(n), + cuvs::Compare(), + stream_)); + + // --- (codes, factors) output mode, with a centroid --- + auto ff_codes_a = raft::make_device_matrix(handle_, n, pd); + auto ff_codes_b = raft::make_device_matrix(handle_, n, pd); + auto factors_a = raft::make_device_matrix(handle_, n, 3); + auto factors_b = raft::make_device_matrix(handle_, n, 3); + + auto centroid_view = raft::make_const_mdspan(centroid_.view()); + std::optional> centroid_opt(centroid_view); + transform(handle_, + pipe.quantizer, + pipe.rotator, + data, + centroid_opt, + ff_codes_a.view(), + factors_a.view()); + transform(handle_, + pipe.quantizer, + pipe.rotator, + data, + centroid_opt, + ff_codes_b.view(), + factors_b.view()); + raft::resource::sync_stream(handle_, stream_); + + auto h_ff_a = verify_codes(ff_codes_a.data_handle(), len); + auto h_ff_b = verify_codes(ff_codes_b.data_handle(), len); + verify_finite(factors_a.data_handle(), static_cast(n) * 3, "factor"); + { + // f_error is an error bound, so it can never be negative. + auto h_factors = to_host(factors_a.data_handle(), static_cast(n) * 3); + size_t n_bad = 0; + for (int64_t i = 0; i < n; ++i) { + if (!(h_factors[static_cast(i) * 3 + 2] >= 0.0f)) { ++n_bad; } + } + EXPECT_EQ(n_bad, 0u) << n_bad << " of " << n << " f_error factors are negative"; + } + expect_same_codes(h_ff_a, h_ff_b, "transform() (factors) is not deterministic"); + ASSERT_TRUE(devArrMatch(factors_a.data_handle(), + factors_b.data_handle(), + static_cast(n) * 3, + cuvs::Compare(), + stream_)); + + // Codes are computed on residuals w.r.t. the centroid, so passing one must + // change the result. + { + size_t n_diff = 0; + for (size_t i = 0; i < len; ++i) { + if (h_codes_a[i] != h_ff_a[i]) { ++n_diff; } + } + EXPECT_GT(n_diff, 0u) << "subtracting a centroid did not change the codes"; + } + + if (ps_.dim % 32 == 0) { testResidualsMatchTransform(pipe, h_codes_a, delta_a, vl_a); } + } + + /** + * With no padding needed, transform() is exactly rotate() followed by + * transform_residuals() - same rotator, same launch shapes, same quantizer + * kernel - so the two paths must agree. + */ + void testResidualsMatchTransform(pipeline const& pipe, + std::vector const& h_codes_ref, + raft::device_vector const& delta_ref, + raft::device_vector const& vl_ref) + { + const int64_t pd = pipe.quantizer.padded_dim; + const int64_t n = ps_.rows; + const size_t len = static_cast(n) * static_cast(pd); + + auto residuals = raft::make_device_matrix(handle_, n, pd); + rotate(handle_, pipe.rotator, raft::make_const_mdspan(dataset_.view()), residuals.view()); + raft::resource::sync_stream(handle_, stream_); + auto residuals_view = raft::make_const_mdspan(residuals.view()); + + auto codes = raft::make_device_matrix(handle_, n, pd); + auto delta = raft::make_device_vector(handle_, n); + auto vl = raft::make_device_vector(handle_, n); + transform_residuals( + handle_, pipe.quantizer, residuals_view, codes.view(), delta.view(), vl.view()); + raft::resource::sync_stream(handle_, stream_); + + auto h_codes = verify_codes(codes.data_handle(), len); + expect_same_codes( + h_codes_ref, h_codes, "transform() and rotate() + transform_residuals() disagree"); + ASSERT_TRUE(devArrMatch(delta_ref.data_handle(), + delta.data_handle(), + static_cast(n), + cuvs::CompareApprox(1e-6f), + stream_)); + ASSERT_TRUE(devArrMatch(vl_ref.data_handle(), + vl.data_handle(), + static_cast(n), + cuvs::CompareApprox(1e-6f), + stream_)); + + // The full-factor residual overload treats std::nullopt as a zero centroid, + // which is what the (delta, vl) call above used too. + auto ff_codes = raft::make_device_matrix(handle_, n, pd); + auto factors = raft::make_device_matrix(handle_, n, 3); + transform_residuals( + handle_, pipe.quantizer, residuals_view, std::nullopt, ff_codes.view(), factors.view()); + raft::resource::sync_stream(handle_, stream_); + auto h_ff_codes = verify_codes(ff_codes.data_handle(), len); + if (ps_.use_fast) { + // Both modes then quantize with the same shared scaling factor, so the + // codes must match. (On the slow path they need not: the full-factor + // launcher hardcodes 64/64 rescale samples instead of forwarding + // params::coarse_samples / params::fine_samples.) + expect_same_codes(h_codes_ref, + h_ff_codes, + "the (delta, vl) and (factors) output modes emit different codes"); + } + verify_finite(factors.data_handle(), static_cast(n) * 3, "residual factor"); + + // Wrong extents must be rejected rather than silently reinterpreted. + if (n > 1) { + auto short_delta = raft::make_device_vector(handle_, n - 1); + EXPECT_THROW(transform_residuals(handle_, + pipe.quantizer, + residuals_view, + codes.view(), + short_delta.view(), + vl.view()), + raft::logic_error); + } + } + + private: + raft::resources handle_; + RabitqQuantizationInputs ps_; + cudaStream_t stream_; + raft::device_matrix dataset_; + raft::device_vector centroid_; +}; + +// --------------------------------------------------------------------------- +// Parameters +// --------------------------------------------------------------------------- + +/** + * Codes of `ex_bits + 1 <= 8` bits fit uint8_t. The dim sweep covers + * power-of-two padded dims (64, 128), non-power-of-two ones (96, 960 -> the + * truncated FHT plus Kac's walk) and one dim that actually gets padded + * (100 -> 128). The two row counts straddle the FHT warp/cooperative kernel + * switch (1000 vectors). + */ +const std::vector generate_inputs_narrow() +{ + auto inputs = raft::util::itertools::product( + {13, 1000}, + {64, 96, 100, 128}, + {rotator_kind::matmul, rotator_kind::fht_kac}, + {1u, 3u}, + {true}); + // dim = 960 exercises the non-power-of-two FHT with trunc_dim = 512. Only one + // matmul instance there: its rotation matrix costs O(dim^3) host work. + auto fht_960 = raft::util::itertools::product( + {13, 1000}, {960}, {rotator_kind::fht_kac}, {3u}, {true}); + auto matmul_960 = raft::util::itertools::product( + {1000}, {960}, {rotator_kind::matmul}, {3u}, {true}); + // use_fast = false runs the per-vector rescale search instead of the shared + // pre-estimated scaling factor. + auto slow = raft::util::itertools::product( + {64}, {96, 128}, {rotator_kind::fht_kac}, {1u, 3u}, {false}); + + inputs.insert(inputs.end(), fht_960.begin(), fht_960.end()); + inputs.insert(inputs.end(), matmul_960.begin(), matmul_960.end()); + inputs.insert(inputs.end(), slow.begin(), slow.end()); + return inputs; +} + +/** ex_bits = 8 needs 9 code bits, so those configurations require uint16_t. */ +const std::vector generate_inputs_wide() +{ + return raft::util::itertools::product( + {512}, {96, 128}, {rotator_kind::matmul, rotator_kind::fht_kac}, {5u, 8u}, {true}); +} + +typedef RabitqQuantizationTest RabitqQuantizationTest_uint8t; +TEST_P(RabitqQuantizationTest_uint8t, Init) { this->testInit(); } +TEST_P(RabitqQuantizationTest_uint8t, Rotate) { this->testRotate(); } +TEST_P(RabitqQuantizationTest_uint8t, Transform) { this->testTransform(); } + +typedef RabitqQuantizationTest RabitqQuantizationTest_uint16t; +TEST_P(RabitqQuantizationTest_uint16t, Init) { this->testInit(); } +TEST_P(RabitqQuantizationTest_uint16t, Rotate) { this->testRotate(); } +TEST_P(RabitqQuantizationTest_uint16t, Transform) { this->testTransform(); } + +INSTANTIATE_TEST_CASE_P(RabitqQuantizationTest, + RabitqQuantizationTest_uint8t, + ::testing::ValuesIn(generate_inputs_narrow())); +INSTANTIATE_TEST_CASE_P(RabitqQuantizationTest, + RabitqQuantizationTest_uint16t, + ::testing::ValuesIn(generate_inputs_wide())); + +// --------------------------------------------------------------------------- +// Checks that do not depend on the sweep +// --------------------------------------------------------------------------- + +/** A code element type too narrow for `ex_bits + 1` bits must be rejected. */ +TEST(RabitqQuantizationBasicTest, CodeWidthValidation) +{ + raft::resources handle; + constexpr int64_t kDim = 128; + constexpr int64_t kRows = 8; + + params p; + p.ex_bits = 8; // 9 bits per code element + auto pipe = make_pipeline(handle, p, kDim); + + const int64_t kPadded = pipe.quantizer.padded_dim; + + auto dataset = raft::make_device_matrix(handle, kRows, kDim); + raft::random::RngState r(1234ULL); + raft::random::uniform(handle, r, dataset.data_handle(), dataset.size(), -1.0f, 1.0f); + + auto data = raft::make_const_mdspan(dataset.view()); + auto delta = raft::make_device_vector(handle, kRows); + auto vl = raft::make_device_vector(handle, kRows); + + auto narrow = raft::make_device_matrix(handle, kRows, kPadded); + EXPECT_THROW(transform(handle, + pipe.quantizer, + pipe.rotator, + data, + std::nullopt, + narrow.view(), + delta.view(), + vl.view()), + raft::logic_error); + + auto wide = raft::make_device_matrix(handle, kRows, kPadded); + EXPECT_NO_THROW(transform(handle, + pipe.quantizer, + pipe.rotator, + data, + std::nullopt, + wide.view(), + delta.view(), + vl.view())); + raft::resource::sync_stream(handle); +} + +/** A rotator and a quantizer built for different dims are not peers. */ +TEST(RabitqQuantizationBasicTest, MismatchedPeers) +{ + raft::resources handle; + params p; + auto rot = make_rotator(handle, p, 128); + auto quant = make_quantizer(handle, p, 256); + ASSERT_NE(rot.padded_dim, quant.padded_dim); + + auto dataset = raft::make_device_matrix(handle, 4, 256); + raft::random::RngState r(1234ULL); + raft::random::uniform(handle, r, dataset.data_handle(), dataset.size(), -1.0f, 1.0f); + auto codes = raft::make_device_matrix(handle, 4, quant.padded_dim); + auto delta = raft::make_device_vector(handle, 4); + auto vl = raft::make_device_vector(handle, 4); + EXPECT_THROW(transform(handle, + quant, + rot, + raft::make_const_mdspan(dataset.view()), + std::nullopt, + codes.view(), + delta.view(), + vl.view()), + raft::logic_error); +} + +/** An empty batch is a no-op, not a zero-sized kernel launch. */ +TEST(RabitqQuantizationBasicTest, EmptyInput) +{ + raft::resources handle; + constexpr int64_t kDim = 64; + params p; + auto pipe = make_pipeline(handle, p, kDim); + + auto dataset = raft::make_device_matrix(handle, 0, kDim); + auto codes = raft::make_device_matrix(handle, 0, pipe.quantizer.padded_dim); + auto delta = raft::make_device_vector(handle, 0); + auto vl = raft::make_device_vector(handle, 0); + EXPECT_NO_THROW(transform(handle, + pipe.quantizer, + pipe.rotator, + raft::make_const_mdspan(dataset.view()), + std::nullopt, + codes.view(), + delta.view(), + vl.view())); + raft::resource::sync_stream(handle); +} + +} // namespace cuvs::preprocessing::quantize::rabitq