diff --git a/cpp/src/preprocessing/quantize/detail/scalar.cuh b/cpp/src/preprocessing/quantize/detail/scalar.cuh index 63a55aaf54..8a4aa64526 100644 --- a/cpp/src/preprocessing/quantize/detail/scalar.cuh +++ b/cpp/src/preprocessing/quantize/detail/scalar.cuh @@ -1,5 +1,5 @@ /* - * 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 */ @@ -7,14 +7,17 @@ #include #include +#include #include #include #include +#include #include #include #include -#include -#include +#include + +#include namespace cuvs::preprocessing::quantize::detail { @@ -30,6 +33,11 @@ _RAFT_HOST_DEVICE bool fp_lt(const half& a, const half& b) return static_cast(a) < static_cast(b); } +template +struct fp_lt_op { + _RAFT_HOST_DEVICE bool operator()(const T& a, const T& b) const { return fp_lt(a, b); } +}; + template struct quantize_op { const T min_; @@ -85,16 +93,31 @@ std::tuple 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(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{}, stream)); + auto sort_ws = raft::make_device_vector(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{}, + 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(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(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]); diff --git a/cpp/tests/preprocessing/scalar_quantization.cu b/cpp/tests/preprocessing/scalar_quantization.cu index 055e80d83a..7aea707818 100644 --- a/cpp/tests/preprocessing/scalar_quantization.cu +++ b/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> 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(max_num_samples / static_cast(cols_), static_cast(rows_)); + size_t min_train_alloc = n_sample_rows * static_cast(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 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(rows_, cols_); auto quantized_input_d = raft::make_device_matrix(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> auto quantized_input_h_const_view = raft::make_host_matrix_view( quantized_input_h.data_handle(), rows_, cols_); auto re_transformed_input_h = raft::make_host_matrix(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( quantized_input_d.data_handle(), rows_, cols_); auto re_transformed_input_d = raft::make_device_matrix(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> } // train quantizer_2 on host - auto quantizer_2 = - cuvs::preprocessing::quantize::scalar::train(handle, params_.quantization_params, dataset_h); + cuvs::preprocessing::quantize::scalar::quantizer 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; diff --git a/cpp/tests/test_utils.cuh b/cpp/tests/test_utils.cuh index 7554596eee..6bf770546f 100644 --- a/cpp/tests/test_utils.cuh +++ b/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 #include #include +#include +#include +#include +#include #include #include #include @@ -323,4 +327,96 @@ inline std::vector 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 +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)(static_cast(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