Dry run: preprocessing::quantize::scalar

cpp/src/preprocessing/quantize/detail/scalar.cuh

@@ -1,20 +1,23 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
 #include <cuvs/preprocessing/quantize/scalar.hpp>
 #include <raft/core/copy.cuh>
+#include <raft/core/device_mdarray.hpp>
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/host_mdspan.hpp>
 #include <raft/core/operators.hpp>
+#include <raft/core/resource/dry_run_flag.hpp>
 #include <raft/linalg/map.cuh>
 #include <raft/matrix/sample_rows.cuh>
 #include <raft/random/rng.cuh>
-#include <thrust/execution_policy.h>
-#include <thrust/sort.h>
+#include <raft/util/cudart_utils.hpp>
+
+#include <cub/device/device_merge_sort.cuh>
 
 namespace cuvs::preprocessing::quantize::detail {
 
@@ -30,6 +33,11 @@ _RAFT_HOST_DEVICE bool fp_lt(const half& a, const half& b)
   return static_cast<float>(a) < static_cast<float>(b);
 }
 
+template <class T>
+struct fp_lt_op {
+  _RAFT_HOST_DEVICE bool operator()(const T& a, const T& b) const { return fp_lt(a, b); }
+};
+
 template <typename T, typename QuantI, typename TempT = double>
 struct quantize_op {
   const T min_;
@@ -85,16 +93,31 @@ std::tuple<T, T> quantile_min_max(
   auto subset = raft::matrix::sample_rows(res, rng, dataset, (IdxT)n_sample_rows);
 
   // quantile / sort element-wise and pick for now
-  size_t subset_size = n_sample_rows * dim;
-  thrust::sort(raft::resource::get_thrust_policy(res),
-               subset.data_handle(),
-               subset.data_handle() + subset_size);
+  auto subset_size = static_cast<int64_t>(n_sample_rows * dim);
+
+  // Sort through CUB rather than thrust so that the sort workspace is queried and allocated
+  // explicitly: thrust::sort takes its temporary storage from the thrust policy, which the
+  // dry-run tracker would not see once the sort itself is skipped.
+  size_t sort_ws_bytes = 0;
+  RAFT_CUDA_TRY(cub::DeviceMergeSort::SortKeys(
+    nullptr, sort_ws_bytes, subset.data_handle(), subset_size, fp_lt_op<T>{}, stream));
+  auto sort_ws = raft::make_device_vector<char, int64_t>(res, sort_ws_bytes);
+  if (!raft::resource::get_dry_run_flag(res)) {
+    RAFT_CUDA_TRY(cub::DeviceMergeSort::SortKeys(sort_ws.data_handle(),
+                                                 sort_ws_bytes,
+                                                 subset.data_handle(),
+                                                 subset_size,
+                                                 fp_lt_op<T>{},
+                                                 stream));
+  }
 
   double half_quantile_pos = (0.5 + 0.5 * quantile) * subset_size;
   int pos_max              = std::ceil(half_quantile_pos) - 1;
   int pos_min              = subset_size - pos_max - 1;
 
-  T minmax_h[2];
+  // The copies below do not run in dry-run mode, so these placeholders are what the caller
+  // receives; they keep quantize_op's scale and offset finite.
+  T minmax_h[2] = {T(0), T(1)};
   raft::copy(res,
              raft::make_host_scalar_view(&minmax_h[0]),
              raft::make_device_scalar_view(subset.data_handle() + pos_min));
@@ -162,6 +185,9 @@ void transform(raft::resources const& res,
   auto main_op      = quantize_op<T, QuantI>(quantizer.min_, quantizer.max_);
   size_t n_elements = dataset.extent(0) * dataset.extent(1);
 
+  // In dry-run mode the views may be unbacked or alias the shared probe buffer.
+  if (raft::resource::get_dry_run_flag(res)) { return; }
+
 #pragma omp parallel for
   for (size_t i = 0; i < n_elements; ++i) {
     out.data_handle()[i] = main_op(dataset.data_handle()[i]);
@@ -191,6 +217,9 @@ void inverse_transform(raft::resources const& res,
   auto main_op      = quantize_op<T, QuantI>(quantizer.min_, quantizer.max_);
   size_t n_elements = dataset.extent(0) * dataset.extent(1);
 
+  // In dry-run mode the views may be unbacked or alias the shared probe buffer.
+  if (raft::resource::get_dry_run_flag(res)) { return; }
+
 #pragma omp parallel for
   for (size_t i = 0; i < n_elements; ++i) {
     out.data_handle()[i] = main_op(dataset.data_handle()[i]);

cpp/tests/preprocessing/scalar_quantization.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -92,19 +92,49 @@ class QuantizationTest : public ::testing::TestWithParam<QuantizationInputs<T>>
 
     size_t print_size = std::min(input_.size(), 20ul);
 
+    // `train` samples this many rows into a device buffer that stays alive for the whole call, so
+    // whatever else it allocates, the peak cannot be below the subsample. Mirrors the sizing in
+    // quantile_min_max, integer division included.
+    constexpr size_t max_num_samples = 1000000;
+    size_t n_sample_rows =
+      std::min<size_t>(max_num_samples / static_cast<size_t>(cols_), static_cast<size_t>(rows_));
+    size_t min_train_alloc = n_sample_rows * static_cast<size_t>(cols_) * sizeof(T);
+
     // train quantizer_1 on device
-    auto quantizer_1 =
-      cuvs::preprocessing::quantize::scalar::train(handle, params_.quantization_params, dataset);
+    // `train` sizes its allocations from the dataset extents and the queried sort workspace, so
+    // the prediction is exact once the sort workspace dominates the peak. On very small inputs it
+    // is not: the peak then falls inside raft::matrix::sample_rows, where skipping the gather
+    // keeps the index buffer alive across the peak, and the dry run reports 8 bytes per sampled
+    // row more than the real run. Hence the upper-bound check rather than an exact match.
+    cuvs::preprocessing::quantize::scalar::quantizer<T> quantizer_1;
+    cuvs::execute_with_dry_run_check(
+      handle,
+      [&](raft::resources const& h) {
+        quantizer_1 =
+          cuvs::preprocessing::quantize::scalar::train(h, params_.quantization_params, dataset);
+      },
+      cuvs::alloc_behavior::DATA_DRIVEN,
+      min_train_alloc);
     std::cerr << "Q1: min = " << (double)quantizer_1.min_ << ", max = " << (double)quantizer_1.max_
               << std::endl;
 
     {
       auto quantized_input_h = raft::make_host_matrix<QuantI, int64_t>(rows_, cols_);
       auto quantized_input_d = raft::make_device_matrix<QuantI, int64_t>(handle, rows_, cols_);
-      cuvs::preprocessing::quantize::scalar::transform(
-        handle, quantizer_1, dataset, quantized_input_d.view());
-      cuvs::preprocessing::quantize::scalar::transform(
-        handle, quantizer_1, dataset_h, quantized_input_h.view());
+      cuvs::execute_with_dry_run_check(
+        handle,
+        [&](raft::resources const& h) {
+          cuvs::preprocessing::quantize::scalar::transform(
+            h, quantizer_1, dataset, quantized_input_d.view());
+        },
+        cuvs::alloc_behavior::NO_ALLOCATIONS);
+      cuvs::execute_with_dry_run_check(
+        handle,
+        [&](raft::resources const& h) {
+          cuvs::preprocessing::quantize::scalar::transform(
+            h, quantizer_1, dataset_h, quantized_input_h.view());
+        },
+        cuvs::alloc_behavior::NO_ALLOCATIONS);
 
       {
         raft::print_device_vector("Input array: ", input_.data(), print_size, std::cerr);
@@ -130,14 +160,24 @@ class QuantizationTest : public ::testing::TestWithParam<QuantizationInputs<T>>
       auto quantized_input_h_const_view = raft::make_host_matrix_view<const QuantI, int64_t>(
         quantized_input_h.data_handle(), rows_, cols_);
       auto re_transformed_input_h = raft::make_host_matrix<T, int64_t>(rows_, cols_);
-      cuvs::preprocessing::quantize::scalar::inverse_transform(
-        handle, quantizer_1, quantized_input_h_const_view, re_transformed_input_h.view());
+      cuvs::execute_with_dry_run_check(
+        handle,
+        [&](raft::resources const& h) {
+          cuvs::preprocessing::quantize::scalar::inverse_transform(
+            h, quantizer_1, quantized_input_h_const_view, re_transformed_input_h.view());
+        },
+        cuvs::alloc_behavior::NO_ALLOCATIONS);
 
       auto quantized_input_d_const_view = raft::make_device_matrix_view<const QuantI, int64_t>(
         quantized_input_d.data_handle(), rows_, cols_);
       auto re_transformed_input_d = raft::make_device_matrix<T, int64_t>(handle, rows_, cols_);
-      cuvs::preprocessing::quantize::scalar::inverse_transform(
-        handle, quantizer_1, quantized_input_d_const_view, re_transformed_input_d.view());
+      cuvs::execute_with_dry_run_check(
+        handle,
+        [&](raft::resources const& h) {
+          cuvs::preprocessing::quantize::scalar::inverse_transform(
+            h, quantizer_1, quantized_input_d_const_view, re_transformed_input_d.view());
+        },
+        cuvs::alloc_behavior::NO_ALLOCATIONS);
       raft::print_device_vector(
         "re-transformed array: ", re_transformed_input_d.data_handle(), print_size, std::cerr);
 
@@ -154,8 +194,15 @@ class QuantizationTest : public ::testing::TestWithParam<QuantizationInputs<T>>
     }
 
     // train quantizer_2 on host
-    auto quantizer_2 =
-      cuvs::preprocessing::quantize::scalar::train(handle, params_.quantization_params, dataset_h);
+    cuvs::preprocessing::quantize::scalar::quantizer<T> quantizer_2;
+    cuvs::execute_with_dry_run_check(
+      handle,
+      [&](raft::resources const& h) {
+        quantizer_2 =
+          cuvs::preprocessing::quantize::scalar::train(h, params_.quantization_params, dataset_h);
+      },
+      cuvs::alloc_behavior::DATA_DRIVEN,
+      min_train_alloc);
     std::cerr << "Q2: min = " << (double)quantizer_2.min_ << ", max = " << (double)quantizer_2.max_
               << std::endl;
 

cpp/tests/test_utils.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -11,6 +11,10 @@
 #include <gtest/gtest.h>
 #include <iostream>
 #include <memory>
+#include <raft/core/dry_run_resources.hpp>
+#include <raft/core/memory_stats_resources.hpp>
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resources.hpp>
 #include <raft/random/rng.cuh>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
@@ -323,4 +327,96 @@ inline std::vector<float> read_csv(std::string filename, bool skip_first_n_colum
   return result;
 }
 
+enum class alloc_behavior {
+  NO_ALLOCATIONS,
+  ARGUMENT_DRIVEN,
+  DATA_DRIVEN,
+};
+
+/**
+ * @brief Execute an action and check dry-run protocol compliance.
+ *
+ * Runs @p action once in dry-run mode (via dry_run_execute) to record predicted
+ * allocations, then runs it for real with all six memory resources wrapped in
+ * statistics adaptors to record actual peak usage. Compares the predicted and
+ * actual peaks according to the specified @p behavior.
+ *
+ * @tparam Action callable with signature void(raft::resources const&)
+ */
+template <typename Action>
+void execute_with_dry_run_check(raft::resources const& res,
+                                Action&& action,
+                                alloc_behavior behavior,
+                                std::size_t min_alloc = 0)
+{
+  auto dry = raft::util::dry_run_execute(res, action);
+
+  raft::memory_stats_resources stat_res(res);
+  std::forward<Action>(action)(static_cast<const raft::resources&>(stat_res));
+  raft::resource::sync_stream(stat_res);
+  auto actual = stat_res.get_bytes_peak();
+
+  auto total_dry    = dry.total();
+  auto total_actual = actual.total();
+
+  if (dry.device_workspace != actual.device_workspace ||
+      dry.device_large_workspace != actual.device_large_workspace ||
+      dry.device_global != actual.device_global || dry.device_managed != actual.device_managed ||
+      dry.host != actual.host || dry.host_pinned != actual.host_pinned) {
+    printf(
+      "  dry-run: ws=%zu large_ws=%zu global=%zu managed=%zu host=%zu pinned=%zu (total=%zu)\n"
+      "  actual:  ws=%zu large_ws=%zu global=%zu managed=%zu host=%zu pinned=%zu (total=%zu)\n",
+      dry.device_workspace,
+      dry.device_large_workspace,
+      dry.device_global,
+      dry.device_managed,
+      dry.host,
+      dry.host_pinned,
+      total_dry,
+      actual.device_workspace,
+      actual.device_large_workspace,
+      actual.device_global,
+      actual.device_managed,
+      actual.host,
+      actual.host_pinned,
+      total_actual);
+  }
+
+  EXPECT_GE(total_actual, min_alloc);
+  EXPECT_GE(total_dry, min_alloc);
+
+  switch (behavior) {
+    case alloc_behavior::NO_ALLOCATIONS:
+      EXPECT_EQ(dry.device_workspace, std::size_t{0});
+      EXPECT_EQ(dry.device_large_workspace, std::size_t{0});
+      EXPECT_EQ(dry.device_global, std::size_t{0});
+      EXPECT_EQ(dry.device_managed, std::size_t{0});
+      EXPECT_EQ(dry.host, std::size_t{0});
+      EXPECT_EQ(dry.host_pinned, std::size_t{0});
+      EXPECT_EQ(actual.device_workspace, std::size_t{0});
+      EXPECT_EQ(actual.device_large_workspace, std::size_t{0});
+      EXPECT_EQ(actual.device_global, std::size_t{0});
+      EXPECT_EQ(actual.device_managed, std::size_t{0});
+      EXPECT_EQ(actual.host, std::size_t{0});
+      EXPECT_EQ(actual.host_pinned, std::size_t{0});
+      break;
+    case alloc_behavior::ARGUMENT_DRIVEN:
+      EXPECT_EQ(dry.device_workspace, actual.device_workspace);
+      EXPECT_EQ(dry.device_large_workspace, actual.device_large_workspace);
+      EXPECT_EQ(dry.device_global, actual.device_global);
+      EXPECT_EQ(dry.device_managed, actual.device_managed);
+      EXPECT_EQ(dry.host, actual.host);
+      EXPECT_EQ(dry.host_pinned, actual.host_pinned);
+      break;
+    case alloc_behavior::DATA_DRIVEN:
+      EXPECT_GE(dry.device_workspace, actual.device_workspace);
+      EXPECT_GE(dry.device_large_workspace, actual.device_large_workspace);
+      EXPECT_GE(dry.device_global, actual.device_global);
+      EXPECT_GE(dry.device_managed, actual.device_managed);
+      EXPECT_GE(dry.host, actual.host);
+      EXPECT_GE(dry.host_pinned, actual.host_pinned);
+      break;
+  }
+}
+
 };  // end namespace cuvs