diff --git a/c/include/cuvs/core/dataset.h b/c/include/cuvs/core/dataset.h index 78d3547495..6d87b407ce 100644 --- a/c/include/cuvs/core/dataset.h +++ b/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 diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 82dba5424d..6b9abc3bce 100644 --- a/c/src/neighbors/cagra.cpp +++ b/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; std::unique_ptr 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; auto mds = cuvs::core::from_dlpack(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; auto mds = cuvs::core::from_dlpack(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; std::unique_ptr 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; auto mds = cuvs::core::from_dlpack(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; auto mds = cuvs::core::from_dlpack(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"); } diff --git a/c/tests/core/dataset_c.cu b/c/tests/core/dataset_c.cu index d2d2ddacb3..3dfc515c0d 100644 --- a/c/tests/core/dataset_c.cu +++ b/c/tests/core/dataset_c.cu @@ -5,6 +5,7 @@ #include #include +#include #include #include @@ -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 +auto device_padded_owner(cuvsDataset_t dataset) +{ + return reinterpret_cast*>(dataset->addr); +} + +template +auto device_padded_view(cuvsDataset_t dataset) +{ + return reinterpret_cast*>(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(padded); + uint32_t const required_stride = + cuvs::neighbors::cagra_required_row_width(static_cast(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 host(n_rows * n_cols, 3.0f); + rmm::device_uvector 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(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); } diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 3be6cf1fa1..d1589a567a 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/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 { 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 { 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(logical_dim), static_cast(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>(src, static_cast(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 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> { 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>(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 @@ -1269,10 +1283,20 @@ auto make_host_padded_dataset_view(SrcT const& src, uint32_t align_bytes = 16) host_padded_dataset_view>(src, static_cast(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 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> { 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>(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 diff --git a/python/cuvs/cuvs/common/dataset.pyx b/python/cuvs/cuvs/common/dataset.pyx index 0c83633d13..098af9fb91 100644 --- a/python/cuvs/cuvs/common/dataset.pyx +++ b/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 diff --git a/python/cuvs/cuvs/tests/test_dataset.py b/python/cuvs/cuvs/tests/test_dataset.py index fc0b82b943..f49de9100a 100644 --- a/python/cuvs/cuvs/tests/test_dataset.py +++ b/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():