[REVIEW] Fix padded dataset creation for aligned device tensors

c/include/cuvs/core/dataset.h

@@ -59,8 +59,11 @@ CUVS_EXPORT cuvsError_t cuvsDatasetCreate(cuvsDataset_t* dataset);
 /**
  * @brief Create an owning padded dataset in the requested memory space.
  *
- * The source tensor may reside in host- or device-accessible memory. Its contents are copied into
- * newly allocated padded storage in `target_mem_type`.
+ * The source tensor may reside in host- or device-accessible memory. Its contents are always
+ * copied into newly allocated padded storage in `target_mem_type` and the returned handle always
+ * owns that storage (`cuvsDatasetGetIsOwning()` reports `true`), even if the source tensor's row
+ * stride already satisfies the padding requirement. Use `cuvsDatasetMakePaddedView()` instead to
+ * avoid the copy when a non-owning view of already-correctly-strided storage will do.
  *
  * @param[in] res cuVS resources
  * @param[in] dataset source tensor

c/src/neighbors/cagra.cpp

@@ -382,14 +382,20 @@ static void make_device_padded_dataset(raft::resources* res_ptr,
   auto dataset = dataset_tensor->dl_tensor;
   using owner_type = cuvs::neighbors::device_padded_dataset<T, int64_t>;
   std::unique_ptr<owner_type> owner;
+  // cuvsDatasetMakePadded() documents an owning copy unconditionally, so force the copy even if
+  // `src` is device-resident and already padded to the required stride (see
+  // make_device_padded_dataset()'s docs: skipping the copy there is only safe for callers able to
+  // track a borrow back to the source, which the C API's owning cuvsDataset handle cannot do).
   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);
-    owner             = cuvs::neighbors::make_device_padded_dataset(*res_ptr, mds);
+    owner             = cuvs::neighbors::make_device_padded_dataset(
+      *res_ptr, mds, /*align_bytes=*/16, /*force_copy=*/true);
   } 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);
-    owner             = cuvs::neighbors::make_device_padded_dataset(*res_ptr, mds);
+    owner             = cuvs::neighbors::make_device_padded_dataset(
+      *res_ptr, mds, /*align_bytes=*/16, /*force_copy=*/true);
   } else {
     RAFT_FAIL("cuvsDatasetMakePadded: unsupported source tensor memory type");
   }
@@ -411,14 +417,18 @@ static void make_host_padded_dataset(raft::resources* res_ptr,
   auto dataset = dataset_tensor->dl_tensor;
   using owner_type = cuvs::neighbors::host_padded_dataset<T, int64_t>;
   std::unique_ptr<owner_type> owner;
+  // See the matching comment in make_device_padded_dataset() above: cuvsDatasetMakePadded() always
+  // copies, even when `src` is host-resident and already padded to the required stride.
   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);
-    owner             = cuvs::neighbors::make_host_padded_dataset(*res_ptr, mds);
+    owner             = cuvs::neighbors::make_host_padded_dataset(
+      *res_ptr, mds, /*align_bytes=*/16, /*force_copy=*/true);
   } else 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);
-    owner             = cuvs::neighbors::make_host_padded_dataset(*res_ptr, mds);
+    owner             = cuvs::neighbors::make_host_padded_dataset(
+      *res_ptr, mds, /*align_bytes=*/16, /*force_copy=*/true);
   } else {
     RAFT_FAIL("cuvsDatasetMakePadded: unsupported source tensor memory type");
   }

c/tests/core/dataset_c.cu

@@ -5,6 +5,7 @@
 
 #include <cuvs/core/c_api.h>
 #include <cuvs/core/dataset.h>
+#include <cuvs/neighbors/common.hpp>
 #include <dlpack/dlpack.h>
 
 #include <cuda_runtime.h>
@@ -36,6 +37,25 @@ struct MatrixTensor {
   }
 };
 
+/**
+ * The C API's `cuvsDataset::addr` (a public struct field, see `cuvs/core/dataset.h`) points at
+ * the concrete C++ dataset object backing the handle. These helpers reach into it the same way
+ * `c/src/neighbors/cagra.cpp` itself does (see `with_dataset_view()`), so tests can verify the
+ * stride/ownership/aliasing behavior that isn't observable through the narrower
+ * `cuvsDatasetGet*()` accessors alone.
+ */
+template <typename T>
+auto device_padded_owner(cuvsDataset_t dataset)
+{
+  return reinterpret_cast<cuvs::neighbors::device_padded_dataset<T, int64_t>*>(dataset->addr);
+}
+
+template <typename T>
+auto device_padded_view(cuvsDataset_t dataset)
+{
+  return reinterpret_cast<cuvs::neighbors::device_padded_dataset_view<T, int64_t>*>(dataset->addr);
+}
+
 }  // namespace
 
 TEST(DatasetC, CreateDestroy)
@@ -90,7 +110,11 @@ TEST(DatasetC, MakePaddedFromDeviceUnalignedOwnsCopy)
   ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS);
 }
 
-TEST(DatasetC, MakePaddedFromDeviceAlignedFailsUseView)
+// Regression test for https://github.com/NVIDIA/cuvs/issues/2482: a device tensor whose row
+// stride already satisfies CAGRA's padding requirement (dim=32 floats = 128 bytes, already
+// 16-byte aligned) must still succeed and return an owning, independent copy -- matching
+// cuvsDatasetMakePadded()'s documented contract -- instead of being rejected.
+TEST(DatasetC, MakePaddedFromDeviceAlignedOwnsCopy)
 {
   cuvsResources_t res;
   ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS);
@@ -104,16 +128,55 @@ TEST(DatasetC, MakePaddedFromDeviceAlignedFailsUseView)
   raft::copy(device.data(), host.data(), host.size(), stream);
 
   MatrixTensor matrix(device.data(), n_rows, n_cols, kDLCUDA, 32);
-  cuvsDataset_t padded;
-  EXPECT_EQ(
+  cuvsDataset_t padded = nullptr;
+  ASSERT_EQ(
     cuvsDatasetMakePadded(res, &matrix.tensor, CUVS_DATASET_MEM_TYPE_DEVICE, &padded),
-    CUVS_ERROR);
-  EXPECT_EQ(padded, nullptr);
+    CUVS_SUCCESS);
+  ASSERT_NE(padded, nullptr);
+
+  bool is_owning = false;
+  ASSERT_EQ(cuvsDatasetGetIsOwning(padded, &is_owning), CUVS_SUCCESS);
+  EXPECT_TRUE(is_owning) << "cuvsDatasetMakePadded() must always return an owning dataset";
+
+  auto* owner = device_padded_owner<float>(padded);
+  uint32_t const required_stride =
+    cuvs::neighbors::cagra_required_row_width<float>(static_cast<uint32_t>(n_cols));
+  EXPECT_EQ(owner->stride(), required_stride);
+  EXPECT_NE(owner->data_handle(), device.data())
+    << "expected an independent copy, not an alias of the source buffer";
+
+  ASSERT_EQ(cuvsDatasetDestroy(padded), CUVS_SUCCESS);
+  ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS);
+}
 
-  cuvsDataset_t view;
+// No-regression check: cuvsDatasetMakePaddedView() on the same already-aligned device tensor
+// must still behave exactly as before (non-owning, zero-copy alias of the source buffer). The
+// fix for #2482 only changes cuvsDatasetMakePadded()'s behavior.
+TEST(DatasetC, MakePaddedViewFromDeviceAlignedIsNonOwningAlias)
+{
+  cuvsResources_t res;
+  ASSERT_EQ(cuvsResourcesCreate(&res), CUVS_SUCCESS);
+  cudaStream_t stream;
+  ASSERT_EQ(cuvsStreamGet(res, &stream), CUVS_SUCCESS);
+
+  constexpr int64_t n_rows = 64;
+  constexpr int64_t n_cols = 32;
+  std::vector<float> host(n_rows * n_cols, 3.0f);
+  rmm::device_uvector<float> device(host.size(), stream);
+  raft::copy(device.data(), host.data(), host.size(), stream);
+
+  MatrixTensor matrix(device.data(), n_rows, n_cols, kDLCUDA, 32);
+  cuvsDataset_t view = nullptr;
   ASSERT_EQ(cuvsDatasetMakePaddedView(res, &matrix.tensor, &view), CUVS_SUCCESS);
   ASSERT_NE(view, nullptr);
 
+  bool is_owning = true;
+  ASSERT_EQ(cuvsDatasetGetIsOwning(view, &is_owning), CUVS_SUCCESS);
+  EXPECT_FALSE(is_owning);
+
+  auto* view_obj = device_padded_view<float>(view);
+  EXPECT_EQ(view_obj->view().data_handle(), device.data()) << "expected a true alias, not a copy";
+
   ASSERT_EQ(cuvsDatasetDestroy(view), CUVS_SUCCESS);
   ASSERT_EQ(cuvsResourcesDestroy(res), CUVS_SUCCESS);
 }

cpp/include/cuvs/neighbors/common.hpp

@@ -1144,7 +1144,8 @@ auto make_device_dense_row_major_dataset_from_src(raft::resources const& res,
                                                   SrcT const& src,
                                                   uint32_t logical_dim,
                                                   uint32_t target_stride,
-                                                  char const* view_factory_name)
+                                                  char const* view_factory_name,
+                                                  bool force_copy = false)
   -> std::unique_ptr<DatasetT>
 {
   uint32_t const src_stride = mdspan_row_stride_elements(src);
@@ -1158,7 +1159,7 @@ auto make_device_dense_row_major_dataset_from_src(raft::resources const& res,
   RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle()));
   bool const device_src =
     (ptr_attrs.type == cudaMemoryTypeDevice) || (ptr_attrs.type == cudaMemoryTypeManaged);
-  if (device_src && src_stride == target_stride) {
+  if (!force_copy && device_src && src_stride == target_stride) {
     RAFT_EXPECTS(false,
                  "source is device and stride is already correct. "
                  "Use %s() to get a view instead.",
@@ -1184,7 +1185,8 @@ auto make_host_dense_row_major_dataset_from_src(raft::resources const& res,
                                                 SrcT const& src,
                                                 uint32_t logical_dim,
                                                 uint32_t target_stride,
-                                                char const* view_factory_name)
+                                                char const* view_factory_name,
+                                                bool force_copy = false)
   -> std::unique_ptr<DatasetT>
 {
   uint32_t const src_stride = mdspan_row_stride_elements(src);
@@ -1193,7 +1195,7 @@ auto make_host_dense_row_major_dataset_from_src(raft::resources const& res,
                "logical dim (%u) must not exceed row stride (%u).",
                static_cast<unsigned>(logical_dim),
                static_cast<unsigned>(target_stride));
-  if (!device_src && src_stride == target_stride) {
+  if (!force_copy && !device_src && src_stride == target_stride) {
     RAFT_EXPECTS(false,
                  "source stride is already correct. Use %s() to get a view instead.",
                  view_factory_name);
@@ -1235,10 +1237,21 @@ auto make_device_padded_dataset_view(const raft::resources& res,
     device_padded_dataset_view<value_type, index_type>>(src, static_cast<uint32_t>(src.extent(1)));
 }
 
+/**
+ * @brief Create an owning device-padded copy of `src`.
+ *
+ * By default (`force_copy = false`), this rejects a device-resident `src` whose row stride
+ * already matches the required padded width: `make_device_padded_dataset_view()` should be used
+ * instead in that case to avoid a redundant device-to-device copy. Callers that must always get
+ * back independent, owning storage regardless of `src`'s existing layout (for example, the
+ * `cuvsDatasetMakePadded()` C API, which documents an unconditional copy) should pass
+ * `force_copy = true`.
+ */
 template <typename SrcT>
 auto make_device_padded_dataset(const raft::resources& res,
                                 SrcT const& src,
-                                uint32_t align_bytes = 16)
+                                uint32_t align_bytes = 16,
+                                bool force_copy      = false)
   -> std::unique_ptr<device_padded_dataset<typename SrcT::value_type, typename SrcT::index_type>>
 {
   using value_type               = typename SrcT::value_type;
@@ -1248,7 +1261,8 @@ auto make_device_padded_dataset(const raft::resources& res,
   return detail::make_device_dense_row_major_dataset_from_src<
     device_padded_dataset<value_type, index_type>,
     value_type,
-    index_type>(res, src, logical_dim, required_stride, "make_device_padded_dataset_view");
+    index_type>(
+    res, src, logical_dim, required_stride, "make_device_padded_dataset_view", force_copy);
 }
 
 template <typename SrcT>
@@ -1269,10 +1283,20 @@ auto make_host_padded_dataset_view(SrcT const& src, uint32_t align_bytes = 16)
     host_padded_dataset_view<value_type, index_type>>(src, static_cast<uint32_t>(src.extent(1)));
 }
 
+/**
+ * @brief Create an owning host-padded copy of `src`.
+ *
+ * By default (`force_copy = false`), this rejects a host-resident `src` whose row stride already
+ * matches the required padded width: `make_host_padded_dataset_view()` should be used instead in
+ * that case to avoid a redundant copy. Callers that must always get back independent, owning
+ * storage regardless of `src`'s existing layout (for example, the `cuvsDatasetMakePadded()` C
+ * API, which documents an unconditional copy) should pass `force_copy = true`.
+ */
 template <typename SrcT>
 auto make_host_padded_dataset(const raft::resources& res,
                               SrcT const& src,
-                              uint32_t align_bytes = 16)
+                              uint32_t align_bytes = 16,
+                              bool force_copy      = false)
   -> std::unique_ptr<host_padded_dataset<typename SrcT::value_type, typename SrcT::index_type>>
 {
   using value_type               = typename SrcT::value_type;
@@ -1282,7 +1306,8 @@ auto make_host_padded_dataset(const raft::resources& res,
   return detail::make_host_dense_row_major_dataset_from_src<
     host_padded_dataset<value_type, index_type>,
     value_type,
-    index_type>(res, src, logical_dim, required_stride, "make_host_padded_dataset_view");
+    index_type>(
+    res, src, logical_dim, required_stride, "make_host_padded_dataset_view", force_copy);
 }
 
 template <typename SrcT>

python/cuvs/cuvs/common/dataset.pyx

@@ -10,12 +10,12 @@ from libcpp cimport bool as cbool
 
 cimport cuvs.common.cydlpack
 from cuvs.common cimport cydlpack
-from cuvs.common.c_api cimport cuvsError_t, cuvsResources_t
+from cuvs.common.c_api cimport cuvsResources_t
 
 from pylibraft.common.cai_wrapper import wrap_array
 from pylibraft.common.interruptible import cuda_interruptible
 
-from cuvs.common.exceptions import check_cuvs, get_last_error_text
+from cuvs.common.exceptions import check_cuvs
 from cuvs.common.resources import auto_sync_resources
 
 
@@ -70,21 +70,11 @@ cdef Dataset make_device_padded_dataset_handle(
         cuvsResources_t res,
         cydlpack.DLManagedTensor* dataset_dlpack):
     cdef Dataset padded = Dataset()
-    cdef cuvsError_t status = cuvsDatasetMakePadded(
+    check_cuvs(cuvsDatasetMakePadded(
         res,
         dataset_dlpack,
         CUVS_DATASET_MEM_TYPE_DEVICE,
         &padded.dataset
-    )
-    if status == cuvsError_t.CUVS_SUCCESS:
-        return padded
-    err = get_last_error_text() or ""
-    if "stride is already correct" not in err:
-        check_cuvs(status)
-    check_cuvs(cuvsDatasetMakePaddedView(
-        res,
-        dataset_dlpack,
-        &padded.dataset
     ))
     return padded
 
@@ -104,12 +94,12 @@ def make_device_padded_dataset(dataset, resources=None):
     """
     Create a device-padded ``Dataset`` from a host or device array.
 
-    The input must be a row-major 2-D matrix. Host arrays are always copied into
-    newly allocated device-padded storage (``is_owning`` is ``True``). Device
-    arrays are copied when their row stride does not already match the
-    required padded width; if the stride is already correct, a non-owning
-    padded view of the input is returned and the caller must keep ``dataset``
-    alive for as long as the ``Dataset`` is used.
+    The input must be a row-major 2-D matrix. Its contents are always copied
+    into newly allocated, owning device-padded storage (``is_owning`` is
+    always ``True``), regardless of whether the input's row stride already
+    satisfies the padding requirement. The returned ``Dataset`` does not
+    keep a reference to ``dataset``; it is independent of it once this
+    function returns.
 
     Parameters
     ----------
@@ -121,8 +111,7 @@ def make_device_padded_dataset(dataset, resources=None):
     Returns
     -------
     dataset : Dataset
-        A device-resident padded dataset handle. Check ``is_owning`` to see
-        whether the handle owns its storage or is a view of ``dataset``.
+        An owning, device-resident padded dataset handle.
 
     Examples
     --------
@@ -143,6 +132,4 @@ def make_device_padded_dataset(dataset, resources=None):
     cdef Dataset padded
     with cuda_interruptible():
         padded = make_device_padded_dataset_handle(res, dataset_dlpack)
-    if not padded.is_owning:
-        padded._source = dataset
     return padded

python/cuvs/cuvs/tests/test_dataset.py

@@ -31,19 +31,23 @@ def test_make_device_padded_dataset(dtype, n_cols, from_host):
     assert ds.layout == "padded"
     assert ds.memory_type == "device"
     assert ds.dtype is not None
-    if from_host:
-        assert ds.is_owning is True
-    else:
-        assert ds.is_owning in (True, False)
+    # cuvsDatasetMakePadded() always copies into newly allocated, owning
+    # storage, regardless of source residency or whether its row stride
+    # already satisfies the padding requirement. See issue #2482.
+    assert ds.is_owning is True
 
 
-def test_make_device_padded_dataset_device_aligned_is_view():
+def test_make_device_padded_dataset_device_aligned_is_owning():
+    # dim=32 float32 rows are already 16-byte aligned (no padding needed), but
+    # cuvsDatasetMakePadded() must still return an independent, owning copy
+    # rather than rejecting the input or silently returning a view. See
+    # issue #2482.
     data = generate_data((64, 32), np.float32)
     source = device_ndarray(data)
     ds = make_device_padded_dataset(source)
     assert ds.layout == "padded"
     assert ds.memory_type == "device"
-    assert ds.is_owning is False
+    assert ds.is_owning is True
 
 
 def test_make_device_padded_dataset_device_unaligned_is_owning():