Expose binary quantizer to C and Python

cpp/CMakeLists.txt

@@ -693,6 +693,7 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$<BOOL:${CUVS_NVTX}>:NVTX_ENAB
       $<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:src/neighbors/hnsw_c.cpp>
       src/neighbors/nn_descent_c.cpp
       src/neighbors/refine/refine_c.cpp
+      src/preprocessing/quantize/binary_c.cpp
       src/preprocessing/quantize/scalar_c.cpp
       src/distance/pairwise_distance_c.cpp
     )

cpp/include/cuvs/preprocessing/quantize/binary.h

@@ -0,0 +1,40 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cuvs/core/c_api.h>
+#include <dlpack/dlpack.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @brief Applies binary quantization transform to given dataset
+ *
+ * @param[in] res raft resource
+ * @param[in] dataset a row-major host or device matrix to transform
+ * @param[out] out a row-major host or device matrix to store transformed data
+ */
+cuvsError_t cuvsBinaryQuantizerTransform(cuvsResources_t res,
+                                         DLManagedTensor* dataset,
+                                         DLManagedTensor* out);
+
+#ifdef __cplusplus
+}
+#endif

cpp/src/preprocessing/quantize/binary_c.cpp

@@ -0,0 +1,76 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cstdint>
+#include <dlpack/dlpack.h>
+
+#include <cuvs/core/c_api.h>
+#include <cuvs/core/exceptions.hpp>
+#include <cuvs/core/interop.hpp>
+#include <cuvs/preprocessing/quantize/binary.h>
+#include <cuvs/preprocessing/quantize/binary.hpp>
+
+namespace {
+
+template <typename T, typename OutputT = uint8_t>
+void _transform(cuvsResources_t res, DLManagedTensor* dataset_tensor, DLManagedTensor* out_tensor)
+{
+  auto res_ptr = reinterpret_cast<raft::resources*>(res);
+
+  auto dataset = dataset_tensor->dl_tensor;
+  if (cuvs::core::is_dlpack_device_compatible(dataset)) {
+    using mdspan_type     = raft::device_matrix_view<T const, int64_t, raft::row_major>;
+    using out_mdspan_type = raft::device_matrix_view<OutputT, int64_t, raft::row_major>;
+
+    cuvs::preprocessing::quantize::binary::transform(
+      *res_ptr,
+      cuvs::core::from_dlpack<mdspan_type>(dataset_tensor),
+      cuvs::core::from_dlpack<out_mdspan_type>(out_tensor));
+
+  } else if (cuvs::core::is_dlpack_host_compatible(dataset)) {
+    using mdspan_type     = raft::host_matrix_view<T const, int64_t, raft::row_major>;
+    using out_mdspan_type = raft::host_matrix_view<OutputT, int64_t, raft::row_major>;
+
+    cuvs::preprocessing::quantize::binary::transform(
+      *res_ptr,
+      cuvs::core::from_dlpack<mdspan_type>(dataset_tensor),
+      cuvs::core::from_dlpack<out_mdspan_type>(out_tensor));
+  } else {
+    RAFT_FAIL("dataset must be accessible on host or device memory");
+  }
+}
+
+}  // namespace
+
+extern "C" cuvsError_t cuvsBinaryQuantizerTransform(cuvsResources_t res,
+                                                    DLManagedTensor* dataset_tensor,
+                                                    DLManagedTensor* out_tensor)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto dataset = dataset_tensor->dl_tensor;
+    if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) {
+      _transform<float>(res, dataset_tensor, out_tensor);
+    } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) {
+      _transform<half>(res, dataset_tensor, out_tensor);
+    } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) {
+      _transform<double>(res, dataset_tensor, out_tensor);
+    } else {
+      RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
+                dataset.dtype.code,
+                dataset.dtype.bits);
+    }
+  });
+}

python/cuvs/cuvs/preprocessing/CMakeLists.txt

@@ -12,4 +12,5 @@
 # the License.
 # =============================================================================
 
+add_subdirectory(quantize/binary)
 add_subdirectory(quantize/scalar)

python/cuvs/cuvs/preprocessing/quantize/binary/CMakeLists.txt

@@ -0,0 +1,28 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+# in compliance with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+# or implied. See the License for the specific language governing permissions and limitations under
+# the License.
+# =============================================================================
+
+# Set the list of Cython files to build
+set(cython_sources binary.pyx)
+set(linked_libraries cuvs::cuvs cuvs::c_api)
+
+# Build all of the Cython targets
+rapids_cython_create_modules(
+  CXX
+  SOURCE_FILES "${cython_sources}"
+  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS cuvs MODULE_PREFIX
+                   preprocessing_quantize_scalar_
+)
+
+foreach(tgt IN LISTS RAPIDS_CYTHON_CREATED_TARGETS)
+  target_link_libraries(${tgt} PRIVATE cuvs_rmm_logger)
+endforeach()

python/cuvs/cuvs/preprocessing/quantize/binary/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .binary import transform

python/cuvs/cuvs/preprocessing/quantize/binary/binary.pxd

@@ -0,0 +1,25 @@
+#
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# cython: language_level=3
+
+from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t
+from cuvs.common.cydlpack cimport DLDataType, DLManagedTensor
+
+
+cdef extern from "cuvs/preprocessing/quantize/binary.h" nogil:
+    cuvsError_t cuvsBinaryQuantizerTransform(cuvsResources_t res,
+                                             DLManagedTensor* dataset,
+                                             DLManagedTensor* out)

python/cuvs/cuvs/preprocessing/quantize/binary/binary.pyx

@@ -0,0 +1,84 @@
+#
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# cython: language_level=3
+
+import numpy as np
+
+from cuvs.common cimport cydlpack
+
+from pylibraft.common import auto_convert_output, device_ndarray
+from pylibraft.common.cai_wrapper import wrap_array
+
+from cuvs.common.exceptions import check_cuvs
+from cuvs.common.resources import auto_sync_resources
+from cuvs.neighbors.common import _check_input_array
+
+
+@auto_sync_resources
+@auto_convert_output
+def transform(dataset, output=None, resources=None):
+    """
+    Applies binary quantization transform to given dataset
+
+    Parameters
+    ----------
+    dataset : row major host or device dataset to transform
+    output : optional preallocated output memory, on host or device memory
+    {resources_docstring}
+
+    Returns
+    -------
+    output : transformed dataset quantized into a uint8
+
+    Examples
+    --------
+    >>> import cupy as cp
+    >>> from cuvs.preprocessing.quantize import binary
+    >>> n_samples = 50000
+    >>> n_features = 50
+    >>> dataset = cp.random.random_sample((n_samples, n_features),
+    ...                                   dtype=cp.float32)
+    >>> transformed = binary.transform(dataset)
+    """
+
+    dataset_ai = wrap_array(dataset)
+
+    _check_input_array(dataset_ai,
+                       [np.dtype("float32"),
+                        np.dtype("float64"),
+                        np.dtype("float16")])
+
+    if output is None:
+        on_device = hasattr(dataset, "__cuda_array_interface__")
+        ndarray = device_ndarray if on_device else np
+        cols = int(np.ceil(dataset_ai.shape[1] / 8))
+        output = ndarray.empty((dataset_ai.shape[0], cols), dtype="uint8")
+
+    output_ai = wrap_array(output)
+    _check_input_array(output_ai, [np.dtype("uint8")])
+
+    cdef cuvsResources_t res = <cuvsResources_t>resources.get_c_obj()
+
+    cdef cydlpack.DLManagedTensor* dataset_dlpack = \
+        cydlpack.dlpack_c(dataset_ai)
+    cdef cydlpack.DLManagedTensor* output_dlpack = \
+        cydlpack.dlpack_c(output_ai)
+
+    check_cuvs(cuvsBinaryQuantizerTransform(res,
+                                            dataset_dlpack,
+                                            output_dlpack))
+
+    return output

python/cuvs/cuvs/tests/test_binary_quantizer.py

@@ -0,0 +1,52 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import numpy as np
+import pytest
+from pylibraft.common import device_ndarray
+
+from cuvs.preprocessing.quantize import binary
+
+
+@pytest.mark.parametrize("n_rows", [50, 100])
+@pytest.mark.parametrize("n_cols", [10, 50])
+@pytest.mark.parametrize("inplace", [True, False])
+@pytest.mark.parametrize("device_memory", [True, False])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.float16])
+def test_binary_quantizer(n_rows, n_cols, inplace, device_memory, dtype):
+    input1 = np.random.random_sample((n_rows, n_cols)).astype(dtype)
+
+    output_cols = int(np.ceil(n_cols / 8))
+    output = (
+        np.zeros((n_rows, output_cols), dtype="uint8") if inplace else None
+    )
+
+    input1_device = device_ndarray(input1)
+    output_device = device_ndarray(output) if inplace else None
+
+    transformed = binary.transform(
+        input1_device if device_memory else input1,
+        output=(output_device if device_memory else output)
+        if inplace
+        else None,
+    )
+    if device_memory:
+        actual = transformed if not inplace else output_device
+        actual = actual.copy_to_host()
+    else:
+        actual = transformed if not inplace else output
+
+    expected = np.packbits(input1 > 0, axis=-1, bitorder="little")
+    assert np.all(actual == expected)