Dry Run Protocol

cpp/include/raft/core/coo_matrix.hpp

@@ -1,6 +1,6 @@
 
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -179,8 +179,8 @@ class coordinate_structure : public coordinate_structure_t<RowType, ColType, NZT
   void initialize_sparsity(nnz_type nnz)
   {
     sparse_structure_type::initialize_sparsity(nnz);
-    c_rows_.resize(nnz);
-    c_cols_.resize(nnz);
+    c_rows_.reallocate(nnz);
+    c_cols_.reallocate(nnz);
   }
 
  protected:

cpp/include/raft/core/csr_matrix.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -188,8 +188,8 @@ class compressed_structure
   void initialize_sparsity(NZType nnz) override
   {
     sparse_structure_type::initialize_sparsity(nnz);
-    c_indptr_.resize(this->get_n_rows() + 1);
-    c_indices_.resize(nnz);
+    c_indptr_.reallocate(this->get_n_rows() + 1);
+    c_indices_.reallocate(nnz);
   }
 
  protected:

cpp/include/raft/core/detail/copy.hpp

@@ -10,6 +10,7 @@
 #include <raft/core/host_mdspan.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/core/mdspan.hpp>
+#include <raft/core/resource/dry_run_flag.hpp>
 #include <raft/core/resource/stream_view.hpp>
 #include <raft/core/resources.hpp>
 
@@ -398,6 +399,10 @@ mdspan_copyable_t<DstType, SrcType> copy(resources const& res, DstType&& dst, Sr
     RAFT_EXPECTS(src.extent(i) == dst.extent(i), "Must copy between mdspans of the same shape");
   }
 
+  // Dry-run guard: raft::copy is a pure data-movement utility with no
+  // allocations that callers would need tracked.
+  if (resource::get_dry_run_flag(res)) { return; }
+
   if constexpr (config::use_intermediate_src) {
 #ifndef RAFT_DISABLE_CUDA
     // Copy to intermediate source on device, then perform necessary

cpp/include/raft/core/device_container_policy.hpp

@@ -126,6 +126,29 @@ class device_uvector {
 
   void resize(size_type size) { data_.resize(size, data_.stream()); }
 
+  /**
+   * @brief Resize the internal buffer without copying old data.
+   *
+   * Unlike resize(), this never copies old data.
+   * Thus, unlike in resize(), there's no point in time where the old and the new buffers are both
+   * alive, and the peak memory usage is lower.
+   *
+   * Unlike resize(), this deallocates the old buffer even if the new size is smaller.
+   * This ensures the memory is released promptly.
+   */
+  void reallocate(size_type size)
+  {
+    if (size != data_.size()) {
+      auto stream = data_.stream();
+      auto mr     = data_.memory_resource();
+      // Resize and shrink rmm::device_uvector: force deallocation without copying old data
+      data_.resize(0, data_.stream());
+      data_.shrink_to_fit(data_.stream());
+      // Assign a new value after the old one is deallocated
+      data_ = rmm::device_uvector<T>(size, stream, mr);
+    }
+  }
+
   [[nodiscard]] auto data() noexcept -> pointer { return data_.data(); }
   [[nodiscard]] auto data() const noexcept -> const_pointer { return data_.data(); }
 };

cpp/include/raft/core/device_mdarray.hpp

@@ -8,6 +8,7 @@
 #include <raft/core/device_container_policy.hpp>
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/mdarray.hpp>
+#include <raft/core/resource/dry_run_flag.hpp>
 #include <raft/core/resources.hpp>
 
 #include <rmm/resource_ref.hpp>
@@ -163,7 +164,7 @@ auto make_device_scalar(raft::resources const& handle, ElementType const& v)
   using policy_t = typename device_scalar<ElementType, IndexType>::container_policy_type;
   policy_t policy{};
   auto scalar = device_scalar<ElementType, IndexType>{handle, extents, policy};
-  scalar(0)   = v;
+  if (!resource::get_dry_run_flag(handle)) { scalar(0) = v; }
   return scalar;
 }
 

cpp/include/raft/core/host_container_policy.hpp

@@ -104,6 +104,27 @@ requires cuda::mr::synchronous_resource_with<MR, cuda::mr::host_accessible>
     *this = std::move(new_container);
   }
 
+  /**
+   * @brief Resize the internal buffer without copying old data.
+   *
+   * Unlike resize(), this never copies old data.
+   * Thus, unlike in resize(), there's no point in time where the old and the new buffers are both
+   * alive, and the peak memory usage is lower.
+   *
+   * Unlike resize(), this deallocates the old buffer even if the new size is smaller.
+   * This ensures the memory is released promptly.
+   */
+  void reallocate(size_type count)
+  {
+    if (bytesize_ == sizeof(value_type) * count) { return; }
+    if (data_ != nullptr) {
+      mr_.deallocate_sync(data_, bytesize_);
+      data_ = nullptr;
+    }
+    auto tmp = host_container{count, mr_};
+    std::swap(tmp, *this);
+  }
+
   [[nodiscard]] auto data() noexcept -> pointer { return data_; }
   [[nodiscard]] auto data() const noexcept -> const_pointer { return data_; }
 };

cpp/include/raft/core/host_mdarray.hpp

@@ -8,6 +8,7 @@
 #include <raft/core/host_container_policy.hpp>
 #include <raft/core/host_mdspan.hpp>
 #include <raft/core/mdarray.hpp>
+#include <raft/core/resource/dry_run_flag.hpp>
 #include <raft/core/resources.hpp>
 
 #include <cstdint>
@@ -223,7 +224,7 @@ auto make_host_scalar(raft::resources const& res, ElementType const& v)
   using policy_t = typename host_scalar<ElementType, IndexType>::container_policy_type;
   policy_t policy;
   auto scalar = host_scalar<ElementType, IndexType>{res, extents, policy};
-  scalar(0)   = v;
+  if (!resource::get_dry_run_flag(res)) { scalar(0) = v; }
   return scalar;
 }
 

cpp/include/raft/core/managed_mdarray.hpp

@@ -8,6 +8,7 @@
 #include <raft/core/managed_container_policy.hpp>
 #include <raft/core/managed_mdspan.hpp>
 #include <raft/core/mdarray.hpp>
+#include <raft/core/resource/dry_run_flag.hpp>
 #include <raft/core/resources.hpp>
 
 #include <cstdint>
@@ -117,7 +118,7 @@ auto make_managed_scalar(raft::resources const& handle, ElementType const& v)
   using policy_t = typename managed_scalar<ElementType>::container_policy_type;
   policy_t policy{};
   auto scalar = managed_scalar<ElementType>{handle, extents, policy};
-  scalar(0)   = v;
+  if (!resource::get_dry_run_flag(handle)) { scalar(0) = v; }
   return scalar;
 }
 

cpp/include/raft/core/pinned_mdarray.hpp

@@ -8,6 +8,7 @@
 #include <raft/core/mdarray.hpp>
 #include <raft/core/pinned_container_policy.hpp>
 #include <raft/core/pinned_mdspan.hpp>
+#include <raft/core/resource/dry_run_flag.hpp>
 #include <raft/core/resources.hpp>
 
 #include <cstdint>
@@ -117,7 +118,7 @@ auto make_pinned_scalar(raft::resources const& handle, ElementType const& v)
   using policy_t = typename pinned_scalar<ElementType>::container_policy_type;
   policy_t policy{};
   auto scalar = pinned_scalar<ElementType>{handle, extents, policy};
-  scalar(0)   = v;
+  if (!resource::get_dry_run_flag(handle)) { scalar(0) = v; }
   return scalar;
 }
 

cpp/include/raft/core/resource/cuda_stream.hpp

@@ -1,10 +1,11 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 
 #include <raft/core/interruptible.hpp>
+#include <raft/core/resource/dry_run_flag.hpp>
 #include <raft/core/resource/resource_types.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/cudart_utils.hpp>
@@ -82,13 +83,18 @@ inline void set_cuda_stream(resources const& res, rmm::cuda_stream_view stream_v
  */
 inline void sync_stream(const resources& res, rmm::cuda_stream_view stream)
 {
+  if (raft::resource::get_dry_run_flag(res)) { return; }
   interruptible::synchronize(stream);
 }
 
 /**
  * @brief synchronize main stream on the resources instance
  */
-inline void sync_stream(const resources& res) { sync_stream(res, get_cuda_stream(res)); }
+inline void sync_stream(const resources& res)
+{
+  if (raft::resource::get_dry_run_flag(res)) { return; }
+  sync_stream(res, get_cuda_stream(res));
+}
 
 /**
  * @}

cpp/include/raft/core/resource/dry_run_flag.hpp

@@ -0,0 +1,89 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <raft/core/resource/resource_types.hpp>
+#include <raft/core/resources.hpp>
+
+#include <memory>
+
+namespace raft::resource {
+
+/**
+ * @defgroup dry_run_flag Dry-run flag resource
+ * @{
+ */
+
+/**
+ * @brief Resource that holds a boolean dry-run flag.
+ *
+ * When the dry-run flag is set, algorithms should skip kernel execution
+ * and only perform allocations to measure memory usage.
+ */
+class dry_run_flag_resource : public resource {
+ public:
+  dry_run_flag_resource() = default;
+  explicit dry_run_flag_resource(bool value) : flag_(value) {}
+  ~dry_run_flag_resource() override = default;
+
+  auto get_resource() -> void* override { return &flag_; }
+
+  void set(bool value) { flag_ = value; }
+  [[nodiscard]] auto get() const -> bool { return flag_; }
+
+ private:
+  bool flag_{false};
+};
+
+/**
+ * @brief Factory that creates a dry_run_flag_resource.
+ */
+class dry_run_flag_resource_factory : public resource_factory {
+ public:
+  explicit dry_run_flag_resource_factory(bool initial_value = false) : initial_value_(initial_value)
+  {
+  }
+
+  auto get_resource_type() -> resource_type override { return resource_type::DRY_RUN_FLAG; }
+  auto make_resource() -> resource* override { return new dry_run_flag_resource(initial_value_); }
+
+ private:
+  bool initial_value_;
+};
+
+/**
+ * @brief Get the dry-run flag from a resources handle.
+ *
+ * @param res raft resources object
+ * @return true if dry-run mode is active
+ */
+inline auto get_dry_run_flag(resources const& res) -> bool
+{
+  if (!res.has_resource_factory(resource_type::DRY_RUN_FLAG)) {
+    res.add_resource_factory(std::make_shared<dry_run_flag_resource_factory>());
+  }
+  return *res.get_resource<bool>(resource_type::DRY_RUN_FLAG);
+}
+
+/**
+ * @brief Set the dry-run flag on a resources handle.
+ *
+ * @param res raft resources object
+ * @param value true to enable dry-run mode, false to disable
+ */
+inline void set_dry_run_flag(resources const& res, bool value)
+{
+  if (!res.has_resource_factory(resource_type::DRY_RUN_FLAG)) {
+    res.add_resource_factory(std::make_shared<dry_run_flag_resource_factory>(value));
+  } else {
+    // The resource may already be instantiated; update it directly
+    auto* flag = res.get_resource<bool>(resource_type::DRY_RUN_FLAG);
+    *flag      = value;
+  }
+}
+
+/** @} */
+
+}  // namespace raft::resource

cpp/include/raft/core/resource/resource_types.hpp

@@ -40,6 +40,7 @@ enum resource_type {
   MULTI_GPU,                 // resource that tracks resource of each device in multi-gpu world
   PINNED_MEMORY_RESOURCE,    // memory resource for pinned (page-locked) host allocations
   MANAGED_MEMORY_RESOURCE,   // resource for managed (unified) allocations
+  DRY_RUN_FLAG,              // dry-run mode flag for allocation profiling
 
   LAST_KEY  // reserved for the last key
 };

cpp/include/raft/core/sparse_types.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -177,7 +177,7 @@ class sparse_matrix {
 
   ~sparse_matrix() noexcept(std::is_nothrow_destructible<container_type>::value) = default;
 
-  void initialize_sparsity(nnz_type nnz) { c_elements_.resize(nnz); };
+  void initialize_sparsity(nnz_type nnz) { c_elements_.reallocate(nnz); };
 
   raft::span<ElementType, is_device> get_elements()
   {

cpp/include/raft/label/classlabels.cuh

@@ -1,17 +1,43 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2022, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #ifndef __CLASS_LABELS_H
 #define __CLASS_LABELS_H
 
 #pragma once
 
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/dry_run_flag.hpp>
+#include <raft/core/resources.hpp>
 #include <raft/label/detail/classlabels.cuh>
 
 namespace raft {
 namespace label {
 
+/**
+ * Get unique class labels.
+ *
+ * The y array is assumed to store class labels. The unique values are selected
+ * from this array.
+ *
+ * @tparam value_t numeric type of the arrays with class labels
+ * @param [in] handle raft resources handle (dry-run aware)
+ * @param [inout] unique output unique labels
+ * @param [in] y device array of labels, size [n]
+ * @param [in] n number of labels
+ * @returns number of unique labels (upper bound in dry-run mode)
+ */
+template <typename value_t>
+int getUniquelabels(raft::resources const& handle,
+                    rmm::device_uvector<value_t>& unique,
+                    value_t* y,
+                    size_t n)
+{
+  return detail::getUniquelabels<value_t>(
+    resource::get_dry_run_flag(handle), unique, y, n, resource::get_cuda_stream(handle));
+}
+
 /**
  * Get unique class labels.
  *