Expand test coverage Part 2

achirkin · achirkin · commit fee4b622edda · 2026-03-24T06:33:39.000-07:00
diff --git a/cpp/tests/sparse/csr_transpose.cu b/cpp/tests/sparse/csr_transpose.cu
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -94,17 +94,23 @@ class CSRTransposeTest : public ::testing::TestWithParam<CSRTransposeInputs<valu
 
     make_data();
 
-    raft::sparse::linalg::csr_transpose(handle,
-                                        indptr.data(),
-                                        indices.data(),
-                                        data.data(),
-                                        out_indptr.data(),
-                                        out_indices.data(),
-                                        out_data.data(),
-                                        params.nrows,
-                                        params.ncols,
-                                        params.nnz,
-                                        stream);
+    raft::execute_with_dry_run_check(
+      handle,
+      [&](raft::resources const& h) {
+        raft::sparse::linalg::csr_transpose(h,
+                                            indptr.data(),
+                                            indices.data(),
+                                            data.data(),
+                                            out_indptr.data(),
+                                            out_indices.data(),
+                                            out_data.data(),
+                                            params.nrows,
+                                            params.ncols,
+                                            params.nnz,
+                                            resource::get_cuda_stream(h));
+      },
+      raft::alloc_behavior::ARGUMENT_DRIVEN,
+      1);
 
     resource::sync_stream(handle, stream);
   }
diff --git a/cpp/tests/sparse/masked_matmul.cu b/cpp/tests/sparse/masked_matmul.cu
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -338,12 +338,18 @@ class MaskedMatmulTest
 
     if constexpr (bits_layout == BitsLayout::Bitmap) {
       auto mask = raft::core::bitmap_view<const bits_t, index_t>(bits_d.data(), params.m, params.n);
-      raft::sparse::linalg::masked_matmul(handle, A, B, mask, C);
+      raft::execute_with_dry_run_check(
+        handle,
+        [&](raft::resources const& h) { raft::sparse::linalg::masked_matmul(h, A, B, mask, C); },
+        raft::alloc_behavior::ARGUMENT_DRIVEN,
+        c_data_d.size() * sizeof(output_t));
     } else if constexpr (bits_layout == BitsLayout::Bitset) {
       auto mask = raft::core::bitset_view<const bits_t, index_t>(bits_d.data(), params.n);
-      raft::sparse::linalg::masked_matmul(handle, A, B, mask, C);
-    } else {
-      GTEST_SKIP() << "Unsupported BitsLayout!";
+      raft::execute_with_dry_run_check(
+        handle,
+        [&](raft::resources const& h) { raft::sparse::linalg::masked_matmul(h, A, B, mask, C); },
+        raft::alloc_behavior::ARGUMENT_DRIVEN,
+        c_data_d.size() * sizeof(output_t));
     }
 
     resource::sync_stream(handle);
diff --git a/cpp/tests/sparse/norm.cu b/cpp/tests/sparse/norm.cu
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -52,7 +52,12 @@ class CSRRowNormTest : public ::testing::TestWithParam<CSRRowNormInputs<Type_f,
     raft::update_device(data.data(), params.data.data(), nnz, stream);
     raft::update_device(verify.data(), params.verify.data(), n_rows, stream);
 
-    linalg::rowNormCsr(handle, indptr.data(), data.data(), nnz, n_rows, result.data(), params.norm);
+    raft::execute_with_dry_run_check(
+      handle,
+      [&](raft::resources const& h) {
+        linalg::rowNormCsr(h, indptr.data(), data.data(), nnz, n_rows, result.data(), params.norm);
+      },
+      raft::alloc_behavior::NO_ALLOCATIONS);
     RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
 
     ASSERT_TRUE(
diff --git a/cpp/tests/sparse/preprocess.cu b/cpp/tests/sparse/preprocess.cu
@@ -137,9 +137,21 @@ class SparsePreprocessCSR
       auto bm25_vals = raft::make_device_vector<Type_f, int64_t>(handle, int(coo_a.nnz));
       raft::util::calc_tfidf_bm25<Index_, Type_f>(handle, csr_matrix.view(), bm25_vals.view());
       if (coo_on) {
-        raft::sparse::matrix::encode_bm25<float, int>(handle, coo_a_matrix, result.view());
+        raft::execute_with_dry_run_check(
+          handle,
+          [&](raft::resources const& h) {
+            raft::sparse::matrix::encode_bm25<float, int>(h, coo_a_matrix, result.view());
+          },
+          raft::alloc_behavior::DATA_DRIVEN,
+          sizeof(float) * coo_a.nnz);
       } else {
-        raft::sparse::matrix::encode_bm25<float, int>(handle, csr_matrix, result.view());
+        raft::execute_with_dry_run_check(
+          handle,
+          [&](raft::resources const& h) {
+            raft::sparse::matrix::encode_bm25<float, int>(h, csr_matrix, result.view());
+          },
+          raft::alloc_behavior::DATA_DRIVEN,
+          sizeof(float) * coo_a.nnz);
       }
       ASSERT_TRUE(raft::devArrMatch<Type_f>(bm25_vals.data_handle(),
                                             result.data_handle(),
@@ -151,9 +163,21 @@ class SparsePreprocessCSR
       raft::util::calc_tfidf_bm25<Index_, Type_f>(
         handle, csr_matrix.view(), tfidf_vals.view(), true);
       if (coo_on) {
-        raft::sparse::matrix::encode_tfidf<float, int>(handle, coo_a_matrix, result.view());
+        raft::execute_with_dry_run_check(
+          handle,
+          [&](raft::resources const& h) {
+            raft::sparse::matrix::encode_tfidf<float, int>(h, coo_a_matrix, result.view());
+          },
+          raft::alloc_behavior::ARGUMENT_DRIVEN,
+          1);
       } else {
-        raft::sparse::matrix::encode_tfidf<float, int>(handle, csr_matrix, result.view());
+        raft::execute_with_dry_run_check(
+          handle,
+          [&](raft::resources const& h) {
+            raft::sparse::matrix::encode_tfidf<float, int>(h, csr_matrix, result.view());
+          },
+          raft::alloc_behavior::ARGUMENT_DRIVEN,
+          1);
       }
       ASSERT_TRUE(raft::devArrMatch<Type_f>(tfidf_vals.data_handle(),
                                             result.data_handle(),
diff --git a/cpp/tests/sparse/reduce.cu b/cpp/tests/sparse/reduce.cu
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -63,14 +63,21 @@ class SparseReduceTest : public ::testing::TestWithParam<SparseReduceInputs<valu
     raft::update_device(out_vals.data(), params.out_vals.data(), params.out_vals.size(), stream);
 
     raft::sparse::COO<value_t, value_idx, value_idx> out(stream);
-    raft::sparse::op::max_duplicates(handle,
-                                     out,
-                                     in_rows.data(),
-                                     in_cols.data(),
-                                     in_vals.data(),
-                                     (value_idx)params.in_rows.size(),
-                                     (value_idx)params.m,
-                                     (value_idx)params.n);
+    auto min_alloc = (2 * sizeof(value_idx) + sizeof(value_t)) * params.out_vals.size();
+    raft::execute_with_dry_run_check(
+      handle,
+      [&](raft::resources const& h) {
+        raft::sparse::op::max_duplicates(h,
+                                         out,
+                                         in_rows.data(),
+                                         in_cols.data(),
+                                         in_vals.data(),
+                                         (value_idx)params.in_rows.size(),
+                                         (value_idx)params.m,
+                                         (value_idx)params.n);
+      },
+      raft::alloc_behavior::DATA_DRIVEN,
+      min_alloc);
     RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
     ASSERT_TRUE(raft::devArrMatch<value_idx>(
       out_rows.data(), out.rows(), out.nnz, raft::Compare<value_idx>()));
diff --git a/cpp/tests/sparse/sddmm.cu b/cpp/tests/sparse/sddmm.cu
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -313,14 +313,20 @@ class SDDMMTest : public ::testing::TestWithParam<SDDMMInputs<ValueType, IndexTy
     auto op_b = params.transpose_b ? raft::linalg::Operation::TRANSPOSE
                                    : raft::linalg::Operation::NON_TRANSPOSE;
 
-    raft::sparse::linalg::sddmm(handle,
-                                a,
-                                b,
-                                c,
-                                op_a,
-                                op_b,
-                                raft::make_host_scalar_view<OutputType>(&params.alpha),
-                                raft::make_host_scalar_view<OutputType>(&params.beta));
+    raft::execute_with_dry_run_check(
+      handle,
+      [&](raft::resources const& h) {
+        raft::sparse::linalg::sddmm(h,
+                                    a,
+                                    b,
+                                    c,
+                                    op_a,
+                                    op_b,
+                                    raft::make_host_scalar_view<OutputType>(&params.alpha),
+                                    raft::make_host_scalar_view<OutputType>(&params.beta));
+      },
+      raft::alloc_behavior::ARGUMENT_DRIVEN,
+      1);
 
     resource::sync_stream(handle);
 
diff --git a/cpp/tests/sparse/select_k_csr.cu b/cpp/tests/sparse/select_k_csr.cu
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -287,8 +287,14 @@ class SelectKCsrTest : public ::testing::TestWithParam<SelectKCsrInputs<index_t>
     auto out_idx = raft::make_device_matrix_view<index_t, index_t, raft::row_major>(
       dst_indices_d.data(), params.n_rows, params.top_k);
 
-    raft::sparse::matrix::select_k(
-      handle, in_val, in_idx, out_val, out_idx, params.select_min, true);
+    raft::execute_with_dry_run_check(
+      handle,
+      [&](raft::resources const& h) {
+        raft::sparse::matrix::select_k(
+          h, in_val, in_idx, out_val, out_idx, params.select_min, true);
+      },
+      raft::alloc_behavior::ARGUMENT_DRIVEN,
+      1);
 
     ASSERT_TRUE(raft::devArrMatch<index_t>(dst_indices_expected_d.data(),
                                            out_idx.data_handle(),
diff --git a/cpp/tests/sparse/solver/lanczos.cu b/cpp/tests/sparse/solver/lanczos.cu
@@ -192,13 +192,19 @@ class rmat_lanczos_tests
     auto csr_matrix = raft::make_device_csr_matrix_view<ValueType, IndexType, IndexType, IndexType>(
       const_cast<ValueType*>(symmetric_coo.vals()), csr_structure);
 
-    std::get<0>(stats) = raft::sparse::solver::lanczos_compute_eigenpairs<IndexType, ValueType>(
+    raft::execute_with_dry_run_check(
       handle,
-      config,
-      csr_matrix,
-      std::make_optional(v0.view()),
-      eigenvalues.view(),
-      eigenvectors.view());
+      [&](raft::resources const& h) {
+        std::get<0>(stats) = raft::sparse::solver::lanczos_compute_eigenpairs<IndexType, ValueType>(
+          h,
+          config,
+          csr_matrix,
+          std::make_optional(v0.view()),
+          eigenvalues.view(),
+          eigenvectors.view());
+      },
+      raft::alloc_behavior::ARGUMENT_DRIVEN,
+      sizeof(ValueType) * symmetric_coo.n_rows * config.ncv);
 
     ASSERT_TRUE(raft::devArrMatch<ValueType>(eigenvalues.data_handle(),
                                              expected_eigenvalues.data_handle(),
@@ -340,13 +346,19 @@ class lanczos_tests : public ::testing::TestWithParam<lanczos_inputs<IndexType,
     auto csr_matrix = raft::make_device_csr_matrix_view<ValueType, IndexType, IndexType, IndexType>(
       const_cast<ValueType*>(vals.data_handle()), csr_structure);
 
-    std::get<0>(stats) = raft::sparse::solver::lanczos_compute_eigenpairs<IndexType, ValueType>(
+    raft::execute_with_dry_run_check(
       handle,
-      config,
-      csr_matrix,
-      std::make_optional(v0.view()),
-      eigenvalues.view(),
-      eigenvectors.view());
+      [&](raft::resources const& h) {
+        std::get<0>(stats) = raft::sparse::solver::lanczos_compute_eigenpairs<IndexType, ValueType>(
+          h,
+          config,
+          csr_matrix,
+          std::make_optional(v0.view()),
+          eigenvalues.view(),
+          eigenvectors.view());
+      },
+      raft::alloc_behavior::ARGUMENT_DRIVEN,
+      sizeof(ValueType) * n * config.ncv);
 
     ASSERT_TRUE(raft::devArrMatch<ValueType>(eigenvalues.data_handle(),
                                              expected_eigenvalues.data_handle(),
diff --git a/cpp/tests/sparse/spmm.cu b/cpp/tests/sparse/spmm.cu
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -189,8 +189,13 @@ class SpmmTest : public ::testing::TestWithParam<SpmmInputs<T>> {
                                               ldz,
                                               params.row_major);
 
-    spmm(
-      handle, params.trans_x, params.trans_y, &alpha, X_csr, y_stride_view, &beta, z_stride_view);
+    raft::execute_with_dry_run_check(
+      handle,
+      [&](raft::resources const& h) {
+        spmm(h, params.trans_x, params.trans_y, &alpha, X_csr, y_stride_view, &beta, z_stride_view);
+      },
+      raft::alloc_behavior::ARGUMENT_DRIVEN,
+      z_size * sizeof(T));
 
     resource::sync_stream(handle, stream);
 
diff --git a/cpp/tests/sparse/symmetrize.cu b/cpp/tests/sparse/symmetrize.cu
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -94,8 +94,14 @@ class SparseSymmetrizeTest
 
     raft::sparse::COO<value_t, value_idx, nnz_t> out(stream);
 
-    raft::sparse::linalg::symmetrize(
-      handle, coo_rows.data(), indices.data(), data.data(), m, n, coo_rows.size(), out);
+    raft::execute_with_dry_run_check(
+      handle,
+      [&](raft::resources const& h) {
+        raft::sparse::linalg::symmetrize(
+          h, coo_rows.data(), indices.data(), data.data(), m, n, coo_rows.size(), out);
+      },
+      raft::alloc_behavior::DATA_DRIVEN,
+      nnz * 2 * (2 * sizeof(value_idx) + sizeof(value_t)));
 
     rmm::device_scalar<value_idx> sum(stream);
     sum.set_value_to_zero_async(stream);
diff --git a/cpp/tests/stats/cov.cu b/cpp/tests/stats/cov.cu
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -64,21 +64,31 @@ class CovTest : public ::testing::TestWithParam<CovInputs<T>> {
     if (params.rowMajor) {
       using layout = raft::row_major;
       raft::stats::mean<true>(mean_act.data(), data.data(), cols, rows, stream);
-      cov(handle,
-          raft::make_device_matrix_view<T, std::uint32_t, layout>(data.data(), rows, cols),
-          raft::make_device_vector_view<const T, std::uint32_t>(mean_act.data(), cols),
-          raft::make_device_matrix_view<T, std::uint32_t, layout>(cov_act.data(), cols, cols),
-          params.sample,
-          params.stable);
+      raft::execute_with_dry_run_check(
+        handle,
+        [&](raft::resources const& h) {
+          cov(h,
+              raft::make_device_matrix_view<T, std::uint32_t, layout>(data.data(), rows, cols),
+              raft::make_device_vector_view<const T, std::uint32_t>(mean_act.data(), cols),
+              raft::make_device_matrix_view<T, std::uint32_t, layout>(cov_act.data(), cols, cols),
+              params.sample,
+              params.stable);
+        },
+        raft::alloc_behavior::NO_ALLOCATIONS);
     } else {
       using layout = raft::col_major;
       raft::stats::mean<false>(mean_act.data(), data.data(), cols, rows, stream);
-      cov(handle,
-          raft::make_device_matrix_view<T, std::uint32_t, layout>(data.data(), rows, cols),
-          raft::make_device_vector_view<const T, std::uint32_t>(mean_act.data(), cols),
-          raft::make_device_matrix_view<T, std::uint32_t, layout>(cov_act.data(), cols, cols),
-          params.sample,
-          params.stable);
+      raft::execute_with_dry_run_check(
+        handle,
+        [&](raft::resources const& h) {
+          cov(h,
+              raft::make_device_matrix_view<T, std::uint32_t, layout>(data.data(), rows, cols),
+              raft::make_device_vector_view<const T, std::uint32_t>(mean_act.data(), cols),
+              raft::make_device_matrix_view<T, std::uint32_t, layout>(cov_act.data(), cols, cols),
+              params.sample,
+              params.stable);
+        },
+        raft::alloc_behavior::NO_ALLOCATIONS);
     }
 
     T data_h[6]       = {1.0, 2.0, 5.0, 4.0, 2.0, 1.0};
diff --git a/cpp/tests/stats/mean.cu b/cpp/tests/stats/mean.cu
