[FEA] Scalar Quantized Dataset

c/include/cuvs/preprocessing/quantize/scalar.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -19,15 +19,24 @@ extern "C" {
  * @defgroup preprocessing_c_scalar C API for Scalar Quantizer
  * @{
  */
+/** Select global or per-dimension scalar quantization bounds. */
+enum cuvsScalarQuantizationMode {
+  CUVS_SCALAR_GLOBAL = 0,
+  CUVS_SCALAR_PER_DIMENSION = 1
+};
+
 /**
- * @brief Scalar quantizer parameters.
+ * @brief Scalar quantizer parameters. Allocate with cuvsScalarQuantizerParamsCreate.
  */
 struct cuvsScalarQuantizerParams {
-  /*
+  /**
    * specifies how many outliers at top & bottom will be ignored
    * needs to be within range of (0, 1]
    */
   float quantile;
+
+  /** Quantization mode. Defaults to CUVS_SCALAR_GLOBAL. */
+  enum cuvsScalarQuantizationMode mode;
 };
 
 typedef struct cuvsScalarQuantizerParams* cuvsScalarQuantizerParams_t;
@@ -53,6 +62,11 @@ CUVS_EXPORT cuvsError_t cuvsScalarQuantizerParamsDestroy(cuvsScalarQuantizerPara
  *
  * The quantization is performed by a linear mapping of an interval in the
  * float data type to the full range of the quantized int type.
+ * Allocate with cuvsScalarQuantizerCreate and release with cuvsScalarQuantizerDestroy;
+ * do not copy the handle's contents. The implementation also owns any per-dimension bounds.
+ * In global mode, min_ and max_ are the trained bounds. In per-dimension mode, they are
+ * numeric_limits<T>::min() and numeric_limits<T>::max() for the training data type T,
+ * and are not used for transformation.
  */
 typedef struct {
   double min_;
@@ -70,7 +84,7 @@ typedef cuvsScalarQuantizer* cuvsScalarQuantizer_t;
 CUVS_EXPORT cuvsError_t cuvsScalarQuantizerCreate(cuvsScalarQuantizer_t* quantizer);
 
 /**
- * @brief De-allocate Scalar Quantizer
+ * @brief De-allocate Scalar Quantizer after outstanding stream work completes.
  *
  * @param[in] quantizer
  * @return cuvsError_t
@@ -80,10 +94,15 @@ CUVS_EXPORT cuvsError_t cuvsScalarQuantizerDestroy(cuvsScalarQuantizer_t quantiz
 /**
  * @brief Trains a scalar quantizer to be used later for quantizing the dataset.
  *
+ * Supports float16, float32, and float64 in global or per-dimension mode.
+ * Per-dimension mode applies the quantile independently to each dimension using every row.
+ * Successful retraining replaces the previous bounds; failed training preserves them.
+ * Complete outstanding stream work using the quantizer before retraining.
+ *
  * @param[in] res raft resource
  * @param[in] params configure scalar quantizer, e.g. quantile
  * @param[in] dataset a row-major host or device matrix
- * @param[out] quantizer trained scalar quantizer
+ * @param[in,out] quantizer scalar quantizer allocated by cuvsScalarQuantizerCreate
  */
 CUVS_EXPORT cuvsError_t cuvsScalarQuantizerTrain(cuvsResources_t res,
                                      cuvsScalarQuantizerParams_t params,
@@ -93,6 +112,9 @@ CUVS_EXPORT cuvsError_t cuvsScalarQuantizerTrain(cuvsResources_t res,
 /**
  * @brief Applies quantization transform to given dataset
  *
+ * Per-dimension mode requires device matrices with the same input type and width as training.
+ * Its output codes are int8 in [-127, 127]. The quantizer must outlive stream work using its bounds.
+ *
  * @param[in] res raft resource
  * @param[in] quantizer a scalar quantizer
  * @param[in] dataset a row-major host or device matrix to transform
@@ -107,7 +129,8 @@ CUVS_EXPORT cuvsError_t cuvsScalarQuantizerTransform(cuvsResources_t res,
  * @brief Perform inverse quantization step on previously quantized dataset
  *
  * Note that depending on the chosen data types train dataset the conversion is
- * not lossless.
+ * not lossless. Per-dimension mode requires device matrices with the same output type and
+ * width as training. The quantizer must outlive stream work using its bounds.
  *
  * @param[in] res raft resource
  * @param[in] quantizer a scalar quantizer

c/src/preprocessing/quantize/scalar.cpp

@@ -1,9 +1,11 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include <cstdint>
+#include <cuda/std/limits>
+#include <memory>
 #include <dlpack/dlpack.h>
 
 #include <cuvs/core/c_api.h>
@@ -15,49 +17,40 @@
 
 namespace {
 
+namespace scalar = cuvs::preprocessing::quantize::scalar;
+
+// Keep the public min/max layout while owning per-dimension bounds behind the factory handle.
+struct scalar_quantizer : cuvsScalarQuantizer {
+  uintptr_t addr = 0;
+  DLDataType dtype{};
+};
+
 template <typename T>
-void _train(cuvsResources_t res,
-            cuvsScalarQuantizerParams params,
-            DLManagedTensor* dataset_tensor,
-            cuvsScalarQuantizer_t quantizer)
+scalar::quantizer<T> _train(cuvsResources_t res,
+                            scalar::params params,
+                            DLManagedTensor* dataset_tensor)
 {
   auto dataset = dataset_tensor->dl_tensor;
-
   auto res_ptr = reinterpret_cast<raft::resources*>(res);
-
-  auto quantizer_params     = cuvs::preprocessing::quantize::scalar::params();
-  quantizer_params.quantile = params.quantile;
-
-  cuvs::preprocessing::quantize::scalar::quantizer<T> ret;
-
   if (cuvs::core::is_dlpack_device_compatible(dataset)) {
     using mdspan_type = raft::device_matrix_view<T const, int64_t, raft::row_major>;
-    auto mds          = cuvs::core::from_dlpack<mdspan_type>(dataset_tensor);
-    ret = cuvs::preprocessing::quantize::scalar::train(*res_ptr, quantizer_params, mds);
+    return scalar::train(*res_ptr, params, cuvs::core::from_dlpack<mdspan_type>(dataset_tensor));
   } else if (cuvs::core::is_dlpack_host_compatible(dataset)) {
     using mdspan_type = raft::host_matrix_view<T const, int64_t, raft::row_major>;
-    auto mds          = cuvs::core::from_dlpack<mdspan_type>(dataset_tensor);
-    ret = cuvs::preprocessing::quantize::scalar::train(*res_ptr, quantizer_params, mds);
+    return scalar::train(*res_ptr, params, cuvs::core::from_dlpack<mdspan_type>(dataset_tensor));
   } else {
     RAFT_FAIL("dataset must be accessible on host or device memory");
   }
-
-  quantizer->min_ = ret.min_;
-  quantizer->max_ = ret.max_;
 }
 
 template <typename T, typename OutputT = int8_t>
 void _transform(cuvsResources_t res,
-                cuvsScalarQuantizer_t quantizer_,
+                const scalar::quantizer<T>& quantizer,
                 DLManagedTensor* dataset_tensor,
                 DLManagedTensor* out_tensor)
 {
   auto res_ptr = reinterpret_cast<raft::resources*>(res);
 
-  cuvs::preprocessing::quantize::scalar::quantizer<T> quantizer;
-  quantizer.min_ = quantizer_->min_;
-  quantizer.max_ = quantizer_->max_;
-
   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>;
@@ -85,16 +78,12 @@ void _transform(cuvsResources_t res,
 
 template <typename OutputT, typename InputT = int8_t>
 void _inverse_transform(cuvsResources_t res,
-                        cuvsScalarQuantizer_t quantizer_,
+                        const scalar::quantizer<OutputT>& quantizer,
                         DLManagedTensor* dataset_tensor,
                         DLManagedTensor* out_tensor)
 {
   auto res_ptr = reinterpret_cast<raft::resources*>(res);
 
-  cuvs::preprocessing::quantize::scalar::quantizer<OutputT> quantizer;
-  quantizer.min_ = quantizer_->min_;
-  quantizer.max_ = quantizer_->max_;
-
   auto dataset = dataset_tensor->dl_tensor;
   if (cuvs::core::is_dlpack_device_compatible(dataset)) {
     using mdspan_type     = raft::device_matrix_view<InputT const, int64_t, raft::row_major>;
@@ -119,12 +108,59 @@ void _inverse_transform(cuvsResources_t res,
     RAFT_FAIL("dataset must be accessible on host or device memory");
   }
 }
+
+template <typename T>
+void _transform(cuvsResources_t res,
+                cuvsScalarQuantizer_t quantizer,
+                DLManagedTensor* dataset,
+                DLManagedTensor* out)
+{
+  _transform(res,
+             scalar::quantizer<T>{static_cast<T>(quantizer->min_), static_cast<T>(quantizer->max_)},
+             dataset,
+             out);
+}
+
+template <typename T>
+void _inverse_transform(cuvsResources_t res,
+                        cuvsScalarQuantizer_t quantizer,
+                        DLManagedTensor* dataset,
+                        DLManagedTensor* out)
+{
+  _inverse_transform(res,
+                     scalar::quantizer<T>{static_cast<T>(quantizer->min_), static_cast<T>(quantizer->max_)},
+                     dataset,
+                     out);
+}
+
+template <typename Fn>
+void dispatch_float(DLDataType dtype, Fn&& fn)
+{
+  RAFT_EXPECTS(dtype.code == kDLFloat && dtype.lanes == 1, "Expected a scalar floating-point dtype");
+  switch (dtype.bits) {
+    case 16: fn.template operator()<half>(); break;
+    case 32: fn.template operator()<float>(); break;
+    case 64: fn.template operator()<double>(); break;
+    default: RAFT_FAIL("Unsupported floating-point dtype bits: %d", dtype.bits);
+  }
+}
+
+void destroy_per_dim(scalar_quantizer* quantizer)
+{
+  if (quantizer->addr) {
+    dispatch_float(quantizer->dtype, [&]<typename T>() {
+      delete reinterpret_cast<scalar::quantizer<T>*>(quantizer->addr);
+    });
+    quantizer->addr = 0;
+  }
+}
+
 }  // namespace
 
 extern "C" cuvsError_t cuvsScalarQuantizerParamsCreate(cuvsScalarQuantizerParams_t* params)
 {
   return cuvs::core::translate_exceptions(
-    [=] { *params = new cuvsScalarQuantizerParams{.quantile = 0.99}; });
+    [=] { *params = new cuvsScalarQuantizerParams{.quantile = 0.99, .mode = CUVS_SCALAR_GLOBAL}; });
 }
 
 extern "C" cuvsError_t cuvsScalarQuantizerParamsDestroy(cuvsScalarQuantizerParams_t params)
@@ -134,12 +170,21 @@ extern "C" cuvsError_t cuvsScalarQuantizerParamsDestroy(cuvsScalarQuantizerParam
 
 extern "C" cuvsError_t cuvsScalarQuantizerCreate(cuvsScalarQuantizer_t* quantizer)
 {
-  return cuvs::core::translate_exceptions([=] { *quantizer = new cuvsScalarQuantizer{}; });
+  return cuvs::core::translate_exceptions([=] {
+    RAFT_EXPECTS(quantizer != nullptr, "quantizer must not be null");
+    *quantizer = new scalar_quantizer{};
+  });
 }
 
 extern "C" cuvsError_t cuvsScalarQuantizerDestroy(cuvsScalarQuantizer_t quantizer)
 {
-  return cuvs::core::translate_exceptions([=] { delete quantizer; });
+  return cuvs::core::translate_exceptions([=] {
+    if (quantizer) {
+      auto q = static_cast<scalar_quantizer*>(quantizer);
+      destroy_per_dim(q);
+      delete q;
+    }
+  });
 }
 
 extern "C" cuvsError_t cuvsScalarQuantizerTrain(cuvsResources_t res,
@@ -148,18 +193,34 @@ extern "C" cuvsError_t cuvsScalarQuantizerTrain(cuvsResources_t res,
                                                 cuvsScalarQuantizer_t quantizer)
 {
   return cuvs::core::translate_exceptions([=] {
-    auto dataset = dataset_tensor->dl_tensor;
-    if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) {
-      _train<float>(res, *params, dataset_tensor, quantizer);
-    } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) {
-      _train<half>(res, *params, dataset_tensor, quantizer);
-    } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) {
-      _train<double>(res, *params, dataset_tensor, quantizer);
-    } else {
-      RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
-                dataset.dtype.code,
-                dataset.dtype.bits);
+    RAFT_EXPECTS(params && dataset_tensor && quantizer,
+                 "params, dataset, and quantizer must not be null");
+    scalar::params cpp_params;
+    cpp_params.quantile = params->quantile;
+    switch (params->mode) {
+      case CUVS_SCALAR_GLOBAL: cpp_params.range = scalar::quantization_mode::GLOBAL; break;
+      case CUVS_SCALAR_PER_DIMENSION:
+        cpp_params.range = scalar::quantization_mode::PER_DIMENSION;
+        break;
+      default: RAFT_FAIL("Unsupported scalar quantization mode");
     }
+    auto q = static_cast<scalar_quantizer*>(quantizer);
+    dispatch_float(dataset_tensor->dl_tensor.dtype, [&]<typename T>() {
+      auto trained = _train<T>(res, cpp_params, dataset_tensor);
+      if (auto const* bounds =
+            std::get_if<typename scalar::quantizer<T>::global_type>(&trained.value)) {
+        destroy_per_dim(q);
+        q->min_ = bounds->min;
+        q->max_ = bounds->max;
+      } else {
+        auto owned = std::make_unique<scalar::quantizer<T>>(std::move(trained));
+        destroy_per_dim(q);
+        q->addr = reinterpret_cast<uintptr_t>(owned.release());
+        q->dtype = dataset_tensor->dl_tensor.dtype;
+        q->min_ = static_cast<double>(cuda::std::numeric_limits<T>::min());
+        q->max_ = static_cast<double>(cuda::std::numeric_limits<T>::max());
+      }
+    });
   });
 }
 
@@ -169,18 +230,16 @@ extern "C" cuvsError_t cuvsScalarQuantizerTransform(cuvsResources_t res,
                                                     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, quantizer, dataset_tensor, out_tensor);
-    } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) {
-      _transform<half>(res, quantizer, dataset_tensor, out_tensor);
-    } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) {
-      _transform<double>(res, quantizer, dataset_tensor, out_tensor);
-    } else {
-      RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
-                dataset.dtype.code,
-                dataset.dtype.bits);
-    }
+    RAFT_EXPECTS(quantizer && dataset_tensor && out_tensor,
+                 "quantizer, dataset, and out must not be null");
+    auto q = static_cast<scalar_quantizer*>(quantizer);
+    dispatch_float(q->addr ? q->dtype : dataset_tensor->dl_tensor.dtype, [&]<typename T>() {
+      if (q->addr) {
+        _transform(res, *reinterpret_cast<scalar::quantizer<T>*>(q->addr), dataset_tensor, out_tensor);
+      } else {
+        _transform<T>(res, quantizer, dataset_tensor, out_tensor);
+      }
+    });
   });
 }
 
@@ -190,16 +249,14 @@ cuvsError_t cuvsScalarQuantizerInverseTransform(cuvsResources_t res,
                                                 DLManagedTensor* out)
 {
   return cuvs::core::translate_exceptions([=] {
-    auto dtype = out->dl_tensor.dtype;
-    if (dtype.code == kDLFloat && dtype.bits == 32) {
-      _inverse_transform<float>(res, quantizer, dataset, out);
-    } else if (dtype.code == kDLFloat && dtype.bits == 16) {
-      _inverse_transform<half>(res, quantizer, dataset, out);
-    } else if (dtype.code == kDLFloat && dtype.bits == 64) {
-      _inverse_transform<double>(res, quantizer, dataset, out);
-    } else {
-      RAFT_FAIL(
-        "Unsupported output dataset DLtensor dtype: %d and bits: %d", dtype.code, dtype.bits);
-    }
+    RAFT_EXPECTS(quantizer && dataset && out, "quantizer, dataset, and out must not be null");
+    auto q = static_cast<scalar_quantizer*>(quantizer);
+    dispatch_float(q->addr ? q->dtype : out->dl_tensor.dtype, [&]<typename T>() {
+      if (q->addr) {
+        _inverse_transform(res, *reinterpret_cast<scalar::quantizer<T>*>(q->addr), dataset, out);
+      } else {
+        _inverse_transform<T>(res, quantizer, dataset, out);
+      }
+    });
   });
 }

c/tests/CMakeLists.txt

@@ -95,6 +95,7 @@ ConfigureTest(
 )
 
 ConfigureTest(NAME PCA_C_TEST PATH preprocessing/pca_c.cu)
+ConfigureTest(NAME SCALAR_QUANTIZATION_C_TEST PATH preprocessing/scalar_c.cu)
 
 if(BUILD_CAGRA_HNSWLIB)
   ConfigureTest(NAME HNSW_C_TEST PATH neighbors/ann_hnsw_c.cu)