diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 97cc9d8bac..9610800016 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -47,7 +47,7 @@ namespace { * Heap-allocated bundle for the C API: owns only `cagra::index`. * Lives behind `cuvsCagraIndex::addr` via `sg_cagra_c_api_index_box`. */ -template +template struct cuvs_cagra_c_api_index_lifetime_holder { cuvs::neighbors::cagra::index idx; }; @@ -66,18 +66,18 @@ struct sg_cagra_c_api_index_box { cuvs::neighbors::c_api::detail::owner_record owner_rec; }; -template +template constexpr auto sg_cagra_index_layout_from_view() { - if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { + if constexpr (cuvs::core::is_device_standard_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::device_standard; - } else if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + } else if constexpr (cuvs::core::is_device_padded_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::device_padded; - } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { + } else if constexpr (cuvs::preprocessing::quantize::pq::is_device_vpq_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::device_vpq; - } else if constexpr (cuvs::neighbors::is_device_bbq_dataset_view_v) { + } else if constexpr (cuvs::preprocessing::quantize::bbq::is_device_bbq_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::device_bbq; - } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { + } else if constexpr (cuvs::core::is_host_standard_dataset_view_v) { return sg_cagra_c_api_index_box::dataset_layout::host_standard; } else { return sg_cagra_c_api_index_box::dataset_layout::host_padded; @@ -120,7 +120,7 @@ static void with_index_by_layout(sg_cagra_c_api_index_box* box, } case sg_cagra_c_api_index_box::dataset_layout::device_vpq: { using index_t = cuvs::neighbors::cagra:: - index>; + index>; auto* idx = reinterpret_cast(box->index_ptr); fn(*idx); break; @@ -151,7 +151,7 @@ static void with_index_by_layout(sg_cagra_c_api_index_box* box, template static void destroy_typed_addr(void* ptr); -template +template static void merge_indices_for_layout( raft::resources* res_ptr, cuvs::neighbors::cagra::index_params const& params_cpp, @@ -166,7 +166,7 @@ static void merge_indices_for_layout( "cuvsCagraMerge: merged dataset handle must be empty"); constexpr auto output_layout = - cuvs::neighbors::is_padded_dataset_view_v ? CUVS_DATASET_LAYOUT_PADDED + cuvs::core::is_padded_dataset_view_v ? CUVS_DATASET_LAYOUT_PADDED : CUVS_DATASET_LAYOUT_STANDARD; int64_t merged_row_count = 0; @@ -179,11 +179,11 @@ static void merge_indices_for_layout( cuvs::neighbors::cagra::detail::merged_dataset_size( *res_ptr, index_ptrs, row_filter); auto const dim = static_cast(index_ptrs.front()->dim()); - auto const stride = static_cast(index_ptrs.front()->dataset().stride()); + auto const stride = static_cast(index_ptrs.front()->dataset().as_matrix_view().stride()); try { auto matrix = raft::make_device_matrix(*res_ptr, final_row_count, stride); - using owner_t = cuvs::neighbors::owning_dataset_for_view_t; + using owner_t = cuvs::core::owning_dataset_for_view_t; auto owner = std::make_unique(std::move(matrix), dim); auto view = owner->as_dataset_view(); auto merged_idx = @@ -213,10 +213,10 @@ static void merge_indices_for_layout( } using host_view_t = std::conditional_t< - cuvs::neighbors::is_padded_dataset_view_v, - cuvs::neighbors::host_padded_dataset_view, - cuvs::neighbors::host_standard_dataset_view>; - using host_owner_t = cuvs::neighbors::owning_dataset_for_view_t; + cuvs::core::is_padded_dataset_view_v, + cuvs::core::host_padded_dataset_view, + cuvs::core::host_standard_dataset_view>; + using host_owner_t = cuvs::core::owning_dataset_for_view_t; auto matrix = raft::make_host_matrix(final_row_count, stride); std::fill_n(matrix.data_handle(), static_cast(matrix.size()), T{}); @@ -227,8 +227,8 @@ static void merge_indices_for_layout( auto const& input = index->dataset(); raft::copy_matrix(matrix.data_handle() + row_offset * static_cast(stride), static_cast(stride), - input.view().data_handle(), - static_cast(input.stride()), + input.as_matrix_view().data_handle(), + static_cast(input.as_matrix_view().stride()), static_cast(dim), static_cast(input.n_rows()), stream); @@ -268,7 +268,7 @@ static void merge_indices_for_layout( } } -template +template static auto convert_opaque_indices_to_concrete_types(cuvsCagraIndex_t* indices, size_t num_indices) -> std::vector*> { @@ -296,14 +296,14 @@ static void with_dataset_view_for_layout(raft::resources* res_ptr, if (cuvs::core::is_dlpack_device_compatible(dataset)) { using mdspan_type = raft::device_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - auto ds_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds); + auto ds_view = cuvs::core::make_device_padded_dataset_view(*res_ptr, mds); fn(ds_view); return; } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { if constexpr (!AllowHost) { RAFT_FAIL("%s", host_not_allowed_err); } using mdspan_type = raft::host_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - auto ds_view = cuvs::neighbors::make_host_padded_dataset_view(mds); + auto ds_view = cuvs::core::make_host_padded_dataset_view(mds); fn(ds_view); return; } @@ -311,14 +311,14 @@ static void with_dataset_view_for_layout(raft::resources* res_ptr, if (cuvs::core::is_dlpack_device_compatible(dataset)) { using mdspan_type = raft::device_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - auto ds_view = cuvs::neighbors::make_device_standard_dataset_view(mds); + auto ds_view = cuvs::core::make_device_standard_dataset_view(mds); fn(ds_view); return; } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { if constexpr (!AllowHost) { RAFT_FAIL("%s", host_not_allowed_err); } using mdspan_type = raft::host_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - auto ds_view = cuvs::neighbors::make_host_standard_dataset_view(mds); + auto ds_view = cuvs::core::make_host_standard_dataset_view(mds); fn(ds_view); return; } @@ -328,7 +328,7 @@ static void with_dataset_view_for_layout(raft::resources* res_ptr, RAFT_FAIL("%s: dataset must have host- or device-compatible memory", err_prefix); } -template +template static void compute_ivfpq_shape_from_indices(cuvsCagraIndex_t* indices, size_t num_indices, int64_t* total_size, @@ -348,7 +348,7 @@ static void compute_ivfpq_shape_from_indices(cuvsCagraIndex_t* indices, } } -template +template static auto make_sg_cagra_c_api_index_box( cuvs_cagra_c_api_index_lifetime_holder* holder) -> std::unique_ptr @@ -359,7 +359,7 @@ static auto make_sg_cagra_c_api_index_box( cuvs::neighbors::c_api::detail::make_owner_record(holder)}); } -template +template static void bind_index_lifetime_holder_to_C_index( cuvsCagraIndex_t out, DLDataType dtype, @@ -370,7 +370,7 @@ static void bind_index_lifetime_holder_to_C_index( out->dtype = dtype; } -template +template static void wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index( cuvsCagraIndex_t out, DLDataType dtype, @@ -410,11 +410,11 @@ static void with_dataset_view(cuvsDataset_t dataset, Fn&& fn) } template -void validate_bbq_layouts(cuvs::neighbors::device_bbq_dataset_view const& dataset) +void validate_bbq_layouts(cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) { using layout_t = cuvs::preprocessing::quantize::bbq::bbq_code_layout; - if (dataset.quantizers.size() == 1) { - auto layout = dataset.quantizers.front().layout; + if (dataset.data().quantizers.size() == 1) { + auto layout = dataset.data().quantizers.front().layout; RAFT_EXPECTS(layout == layout_t::packed_1b || layout == layout_t::transposed_2b || layout == layout_t::packed_4b || layout == layout_t::packed_7b || layout == layout_t::packed_8b, @@ -422,10 +422,10 @@ void validate_bbq_layouts(cuvs::neighbors::device_bbq_dataset_view c return; } - const bool has_1b = dataset.has_layout(layout_t::packed_1b); - const bool has_2bt = dataset.has_layout(layout_t::transposed_2b); - const bool has_4b = dataset.has_layout(layout_t::packed_4b); - const bool has_4bt = dataset.has_layout(layout_t::transposed_4b); + const bool has_1b = dataset.data().has_layout(layout_t::packed_1b); + const bool has_2bt = dataset.data().has_layout(layout_t::transposed_2b); + const bool has_4b = dataset.data().has_layout(layout_t::packed_4b); + const bool has_4bt = dataset.data().has_layout(layout_t::transposed_4b); RAFT_EXPECTS((has_1b && (has_4b || has_2bt || has_4bt)) || (has_2bt && has_4bt), "cuvsDatasetMakeBbqView: unsupported asymmetric BBQ layout pair"); } @@ -446,9 +446,9 @@ auto get_cpp_bbq_quantizer_view(cuvsBbqQuantizer_t quantizer) template auto make_bbq_dataset_view(cuvsBbqQuantizer_t* quantizers, std::size_t num_quantizers) - -> std::unique_ptr> + -> std::unique_ptr> { - using dataset_view_t = cuvs::neighbors::device_bbq_dataset_view; + using dataset_view_t = cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view; auto dataset = std::make_unique(); int64_t expected_rows{-1}; uint32_t expected_dim{}; @@ -469,7 +469,7 @@ auto make_bbq_dataset_view(cuvsBbqQuantizer_t* quantizers, std::size_t num_quant RAFT_EXPECTS(quantizer.metric == expected_metric, "cuvsDatasetMakeBbqView: all quantizers must use the same metric"); } - dataset->add_quantizer(quantizer); + dataset->data().add_quantizer(quantizer); } return dataset; @@ -481,7 +481,7 @@ void make_and_bind_bbq_dataset(cuvsBbqQuantizer_t* quantizers, DLDataType dtype, cuvsDataset_t* output) { - using dataset_view_t = cuvs::neighbors::device_bbq_dataset_view; + using dataset_view_t = cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view; auto view = make_bbq_dataset_view(quantizers, num_quantizers); validate_bbq_layouts(*view); @@ -495,8 +495,8 @@ void make_and_bind_bbq_dataset(cuvsBbqQuantizer_t* quantizers, *output = handle.release(); } -using device_vpq_owner_t = cuvs::neighbors::device_vpq_dataset; -using device_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; +using device_vpq_owner_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; +using device_vpq_view_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset_view; static void bind_vpq_owner_to_dataset(std::unique_ptr owner, cuvsDataset_t* output) @@ -515,9 +515,9 @@ static void bind_vpq_owner_to_dataset(std::unique_ptr owner, } static auto make_cpp_vpq_params(cuvsPqParams const& params) - -> cuvs::neighbors::vpq_params + -> cuvs::preprocessing::quantize::pq::vpq_params { - auto out = cuvs::neighbors::vpq_params{}; + auto out = cuvs::preprocessing::quantize::pq::vpq_params{}; out.pq_bits = params.pq_bits; out.pq_dim = params.pq_dim; out.vq_n_centers = params.vq_n_centers; @@ -546,19 +546,19 @@ static auto make_device_pq_dataset(raft::resources* res_ptr, const bool padded = dataset->layout == CUVS_DATASET_LAYOUT_PADDED; if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE) { if (padded) { - with_dataset_view, - cuvs::neighbors::device_padded_dataset_view>(dataset, make); + with_dataset_view, + cuvs::core::device_padded_dataset_view>(dataset, make); } else { - with_dataset_view, - cuvs::neighbors::device_standard_dataset_view>(dataset, make); + with_dataset_view, + cuvs::core::device_standard_dataset_view>(dataset, make); } } else if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_HOST) { if (padded) { - with_dataset_view, - cuvs::neighbors::host_padded_dataset_view>(dataset, make); + with_dataset_view, + cuvs::core::host_padded_dataset_view>(dataset, make); } else { - with_dataset_view, - cuvs::neighbors::host_standard_dataset_view>(dataset, make); + with_dataset_view, + cuvs::core::host_standard_dataset_view>(dataset, make); } } else { RAFT_FAIL("cuvsDatasetMakePQ: invalid source dataset memory type"); @@ -572,16 +572,16 @@ static void make_device_padded_dataset(raft::resources* res_ptr, cuvsDataset_t* output_padded_dataset) { auto dataset = dataset_tensor->dl_tensor; - using owner_type = cuvs::neighbors::device_padded_dataset; + using owner_type = cuvs::core::device_padded_dataset; std::unique_ptr owner; 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::core::make_device_padded_dataset(*res_ptr, mds); } 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::core::make_device_padded_dataset(*res_ptr, mds); } else { RAFT_FAIL("cuvsDatasetMakePadded: unsupported source tensor memory type"); } @@ -601,16 +601,16 @@ static void make_host_padded_dataset(raft::resources* res_ptr, cuvsDataset_t* output_padded_dataset) { auto dataset = dataset_tensor->dl_tensor; - using owner_type = cuvs::neighbors::host_padded_dataset; + using owner_type = cuvs::core::host_padded_dataset; std::unique_ptr owner; 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::core::make_host_padded_dataset(*res_ptr, mds); } 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::core::make_host_padded_dataset(*res_ptr, mds); } else { RAFT_FAIL("cuvsDatasetMakePadded: unsupported source tensor memory type"); } @@ -637,7 +637,7 @@ static void make_device_padded_dataset_view(raft::resources* res_ptr, } using mdspan_type = raft::device_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - auto ds_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds); + auto ds_view = cuvs::core::make_device_padded_dataset_view(*res_ptr, mds); auto* owned_view = new decltype(ds_view){ds_view}; out->addr = reinterpret_cast(owned_view); out->destroy_addr = &destroy_typed_addr; @@ -661,7 +661,7 @@ static void make_host_padded_dataset_view(raft::resources*, } using mdspan_type = raft::host_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - auto ds_view = cuvs::neighbors::make_host_padded_dataset_view(mds); + auto ds_view = cuvs::core::make_host_padded_dataset_view(mds); auto* owned_view = new decltype(ds_view){ds_view}; out->addr = reinterpret_cast(owned_view); out->destroy_addr = &destroy_typed_addr; @@ -685,7 +685,7 @@ static void make_device_standard_dataset_view(raft::resources*, } using mdspan_type = raft::device_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - auto ds_view = cuvs::neighbors::make_device_standard_dataset_view(mds); + auto ds_view = cuvs::core::make_device_standard_dataset_view(mds); auto* owned_view = new decltype(ds_view){ds_view}; out->addr = reinterpret_cast(owned_view); out->destroy_addr = &destroy_typed_addr; @@ -709,7 +709,7 @@ static void make_host_standard_dataset_view(raft::resources*, } using mdspan_type = raft::host_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - auto ds_view = cuvs::neighbors::make_host_standard_dataset_view(mds); + auto ds_view = cuvs::core::make_host_standard_dataset_view(mds); auto* owned_view = new decltype(ds_view){ds_view}; out->addr = reinterpret_cast(owned_view); out->destroy_addr = &destroy_typed_addr; @@ -755,8 +755,8 @@ static void update_dataset(raft::resources* res_ptr, if (dataset->layout == CUVS_DATASET_LAYOUT_PQ) { with_dataset_view(dataset, rebind); } else { - using owner_t = cuvs::neighbors::device_padded_dataset; - using view_t = cuvs::neighbors::device_padded_dataset_view; + using owner_t = cuvs::core::device_padded_dataset; + using view_t = cuvs::core::device_padded_dataset_view; with_dataset_view(dataset, rebind); } } @@ -864,8 +864,8 @@ void _from_args(cuvsResources_t res, if (cuvs::core::is_dlpack_device_compatible(dataset)) { using mdspan_type = raft::device_matrix_view; auto mds = cuvs::core::from_dlpack(dataset_tensor); - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)) { - auto dataset_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds); + if (cuvs::core::matrix_has_padded_row_width(mds)) { + auto dataset_view = cuvs::core::make_device_padded_dataset_view(*res_ptr, mds); auto* raw = new cuvs::neighbors::cagra::device_padded_index( *res_ptr, metric); *raw = @@ -873,10 +873,10 @@ void _from_args(cuvsResources_t res, update_graph_from_dlpack(raw); wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index< T, - cuvs::neighbors::device_padded_dataset_view>( + cuvs::core::device_padded_dataset_view>( output_index, output_index->dtype, raw); } else { - auto dataset_view = cuvs::neighbors::make_device_standard_dataset_view(mds); + auto dataset_view = cuvs::core::make_device_standard_dataset_view(mds); auto* raw = new cuvs::neighbors::cagra::device_standard_index( *res_ptr, metric); *raw = @@ -884,7 +884,7 @@ void _from_args(cuvsResources_t res, update_graph_from_dlpack(raw); wrap_CPP_index_in_lifetime_holder_and_bind_to_C_index< T, - cuvs::neighbors::device_standard_dataset_view>( + cuvs::core::device_standard_dataset_view>( output_index, output_index->dtype, raw); } } else if (cuvs::core::is_dlpack_host_compatible(dataset)) { @@ -929,8 +929,8 @@ void _extend(cuvsResources_t res, if constexpr (!idx_is_padded) { RAFT_FAIL("cuvsCagraExtend: only device_padded indices are extendable"); } else { - using out_owner_t = cuvs::neighbors::device_padded_dataset; - using out_view_t = cuvs::neighbors::device_padded_dataset_view; + using out_owner_t = cuvs::core::device_padded_dataset; + using out_view_t = cuvs::core::device_padded_dataset_view; with_dataset_view(extended_dataset, [&](auto& out_dataset) { cuvs::neighbors::cagra::extend( *res_ptr, extend_params, out_dataset, new_start_row, idx); @@ -1119,11 +1119,11 @@ void _serialize(cuvsResources_t res, const char *filename, with_index_by_layout(box, null_handle_err, "", [&](auto &idx) { using index_dataset_view_t = std::remove_cvref_t; - if constexpr (cuvs::neighbors::is_bbq_dataset_view_v) { + if constexpr (cuvs::preprocessing::quantize::bbq::is_bbq_dataset_view_v) { RAFT_EXPECTS(!include_dataset, "cuvsCagraSerializeGraphAndDataset is not supported for BBQ indices"); cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), idx); - } else if constexpr (cuvs::neighbors::is_vpq_dataset_view_v) { + } else if constexpr (cuvs::preprocessing::quantize::pq::is_vpq_dataset_view_v) { RAFT_EXPECTS( !include_dataset, "cuvsCagraSerializeGraphAndDataset is not supported for PQ indices; serialize the PQ " @@ -1212,34 +1212,34 @@ void dispatch_serialized_dataset_kind( switch (kind) { case serialized_kind::device_padded: fn.template operator()< - cuvs::neighbors::device_padded_dataset_view>(); + cuvs::core::device_padded_dataset_view>(); break; case serialized_kind::device_standard: fn.template operator()< - cuvs::neighbors::device_standard_dataset_view>(); + cuvs::core::device_standard_dataset_view>(); break; case serialized_kind::host_padded: fn.template operator()< - cuvs::neighbors::host_padded_dataset_view>(); + cuvs::core::host_padded_dataset_view>(); break; case serialized_kind::host_standard: fn.template operator()< - cuvs::neighbors::host_standard_dataset_view>(); + cuvs::core::host_standard_dataset_view>(); break; case serialized_kind::none: fn.template operator()< - cuvs::neighbors::device_padded_dataset_view>(); + cuvs::core::device_padded_dataset_view>(); break; } } -template +template void _deserialize(cuvsResources_t res, const char *filename, cuvsCagraIndex_t output_index, DLDataType dtype, bool include_dataset, cuvsDataset_t *out_dataset) { auto res_ptr = reinterpret_cast(res); using view_t = ViewT; - using owner_dataset_t = cuvs::neighbors::owning_dataset_for_view_t; + using owner_dataset_t = cuvs::core::owning_dataset_for_view_t; using holder_t = cuvs_cagra_c_api_index_lifetime_holder; auto holder = std::make_unique( @@ -1262,11 +1262,11 @@ void _deserialize(cuvsResources_t res, const char *filename, dataset_handle->destroy_addr = &destroy_typed_addr; dataset_handle->dtype = dtype; dataset_handle->mem_type = - cuvs::neighbors::is_device_dataset_view_v + cuvs::core::is_device_dataset_view_v ? CUVS_DATASET_MEM_TYPE_DEVICE : CUVS_DATASET_MEM_TYPE_HOST; dataset_handle->layout = - cuvs::neighbors::is_padded_dataset_view_v + cuvs::core::is_padded_dataset_view_v ? CUVS_DATASET_LAYOUT_PADDED : CUVS_DATASET_LAYOUT_STANDARD; dataset_handle->is_owning = true; @@ -1295,8 +1295,8 @@ void _serialize_to_hnswlib(cuvsResources_t res, const char *filename, "cuvsCagraSerializeToHnswlib: host indices are allowed", [&](auto &idx) { using index_dataset_view_t = std::remove_cvref_t; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v || - cuvs::neighbors::is_bbq_dataset_view_v) { + if constexpr (cuvs::preprocessing::quantize::pq::is_vpq_dataset_view_v || + cuvs::preprocessing::quantize::bbq::is_bbq_dataset_view_v) { RAFT_FAIL("cuvsCagraSerializeToHnswlib is not supported for quantized dataset layouts"); } else { cuvs::neighbors::cagra::serialize_to_hnswlib( @@ -1340,10 +1340,10 @@ void _merge(cuvsResources_t res, } if (params.build_algo == cuvsCagraGraphBuildAlgo::IVF_PQ) { if (layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) { - compute_ivfpq_shape_from_indices>( + compute_ivfpq_shape_from_indices>( indices, num_indices, &total_size, &dim); } else { - compute_ivfpq_shape_from_indices>( + compute_ivfpq_shape_from_indices>( indices, num_indices, &total_size, &dim); } } @@ -1355,15 +1355,15 @@ void _merge(cuvsResources_t res, dim); if (layout == sg_cagra_c_api_index_box::dataset_layout::device_padded) { auto index_ptrs = - convert_opaque_indices_to_concrete_types>( + convert_opaque_indices_to_concrete_types>( indices, num_indices); - merge_indices_for_layout>( + merge_indices_for_layout>( res_ptr, params_cpp, index_ptrs, filter, merge_params, merged_dataset, output_index); } else { auto index_ptrs = - convert_opaque_indices_to_concrete_types>( + convert_opaque_indices_to_concrete_types>( indices, num_indices); - merge_indices_for_layout>( + merge_indices_for_layout>( res_ptr, params_cpp, index_ptrs, filter, merge_params, merged_dataset, output_index); } } @@ -1378,12 +1378,12 @@ void get_dataset_view(cuvsCagraIndex_t index, DLManagedTensor* dataset) "cuvsCagraIndexGetDataset: host indices are allowed", [&](auto& idx) { using index_dataset_view_t = std::remove_cvref_t; - if constexpr (cuvs::neighbors::is_vpq_dataset_view_v || - cuvs::neighbors::is_bbq_dataset_view_v) { + if constexpr (cuvs::preprocessing::quantize::pq::is_vpq_dataset_view_v || + cuvs::preprocessing::quantize::bbq::is_bbq_dataset_view_v) { RAFT_FAIL( "cuvsCagraIndexGetDataset does not expose quantized datasets as dense DLPack tensors"); } else { - cuvs::core::to_dlpack(idx.dataset().view(), dataset); + cuvs::core::to_dlpack(idx.dataset().as_matrix_view(), dataset); } }); } @@ -1892,7 +1892,7 @@ extern "C" cuvsError_t cuvsCagraUpdateDataset(cuvsResources_t res, * Build from an already-constructed C++ dataset view. `DatasetViewT` selects the * `cuvs::neighbors::cagra::build` overload, and therefore the resulting index type. */ -template +template static void build_index_from_dataset_view(raft::resources* res_ptr, cuvsCagraIndexParams_t params, DatasetViewT const& ds_view, @@ -1921,8 +1921,8 @@ static void build_dispatch_on_mem_type_and_layout(raft::resources* res_ptr, if (dataset->layout == CUVS_DATASET_LAYOUT_BBQ) { RAFT_EXPECTS(dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE, "cuvsCagraBuild: BBQ dataset must be device-resident"); - using owner_t = cuvs::neighbors::device_bbq_dataset; - using view_t = cuvs::neighbors::device_bbq_dataset_view; + using owner_t = cuvs::preprocessing::quantize::bbq::device_bbq_dataset; + using view_t = cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view; with_dataset_view(dataset, [&](auto const& view) { build_index_from_dataset_view(res_ptr, params, view, index); }); @@ -1933,28 +1933,28 @@ static void build_dispatch_on_mem_type_and_layout(raft::resources* res_ptr, if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_DEVICE) { if (is_padded) { - using owner_t = cuvs::neighbors::device_padded_dataset; - using view_t = cuvs::neighbors::device_padded_dataset_view; + using owner_t = cuvs::core::device_padded_dataset; + using view_t = cuvs::core::device_padded_dataset_view; with_dataset_view(dataset, [&](auto const& view) { build_index_from_dataset_view(res_ptr, params, view, index); }); } else { - using owner_t = cuvs::neighbors::device_standard_dataset; - using view_t = cuvs::neighbors::device_standard_dataset_view; + using owner_t = cuvs::core::device_standard_dataset; + using view_t = cuvs::core::device_standard_dataset_view; with_dataset_view(dataset, [&](auto const& view) { build_index_from_dataset_view(res_ptr, params, view, index); }); } } else if (dataset->mem_type == CUVS_DATASET_MEM_TYPE_HOST) { if (is_padded) { - using owner_t = cuvs::neighbors::host_padded_dataset; - using view_t = cuvs::neighbors::host_padded_dataset_view; + using owner_t = cuvs::core::host_padded_dataset; + using view_t = cuvs::core::host_padded_dataset_view; with_dataset_view(dataset, [&](auto const& view) { build_index_from_dataset_view(res_ptr, params, view, index); }); } else { - using owner_t = cuvs::neighbors::host_standard_dataset; - using view_t = cuvs::neighbors::host_standard_dataset_view; + using owner_t = cuvs::core::host_standard_dataset; + using view_t = cuvs::core::host_standard_dataset_view; with_dataset_view(dataset, [&](auto const& view) { build_index_from_dataset_view(res_ptr, params, view, index); }); @@ -2302,7 +2302,7 @@ extern "C" cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params extern "C" cuvsError_t cuvsCagraCompressionParamsCreate(cuvsCagraCompressionParams_t* params) { return cuvs::core::translate_exceptions([=] { - auto ps = cuvs::neighbors::vpq_params(); + auto ps = cuvs::preprocessing::quantize::pq::vpq_params(); *params = new cuvsCagraCompressionParams{.pq_bits = ps.pq_bits, .pq_dim = ps.pq_dim, @@ -2433,7 +2433,7 @@ extern "C" cuvsError_t cuvsCagraDeserializeGraph(cuvsResources_t res, "cuvsCagraDeserializeGraph: null index handle"); auto const header = read_serialized_header(res, filename); dispatch_serialized_dtype(header.dtype, [&]() { - using view_t = cuvs::neighbors::device_padded_dataset_view; + using view_t = cuvs::core::device_padded_dataset_view; _deserialize( res, filename, index, header.dtype, false, nullptr); }); diff --git a/c/src/neighbors/cagra.hpp b/c/src/neighbors/cagra.hpp index 562f78f124..db579eebff 100644 --- a/c/src/neighbors/cagra.hpp +++ b/c/src/neighbors/cagra.hpp @@ -20,7 +20,7 @@ void convert_c_search_params(cuvsCagraSearchParams params, void* cagra_c_api_index_ptr(cuvsCagraIndex const* idx); namespace detail { -template +template int64_t merged_dataset_size( raft::resources const& res, std::vector*> const& indices, diff --git a/c/src/neighbors/mg_cagra.cpp b/c/src/neighbors/mg_cagra.cpp index 664883cd6f..496a93b807 100644 --- a/c/src/neighbors/mg_cagra.cpp +++ b/c/src/neighbors/mg_cagra.cpp @@ -75,8 +75,8 @@ static void with_mg_index_by_layout(mg_cagra_c_api_index_box* box, template static void with_device_padded_dataset_view(cuvsDataset_t dataset, Fn&& fn) { - using owner_t = cuvs::neighbors::device_padded_dataset; - using view_t = cuvs::neighbors::device_padded_dataset_view; + using owner_t = cuvs::core::device_padded_dataset; + using view_t = cuvs::core::device_padded_dataset_view; if (dataset->is_owning) { auto* owner = reinterpret_cast(dataset->addr); auto view = owner->as_dataset_view(); @@ -204,13 +204,13 @@ void* _mg_build(cuvsResources_t res, if (layout == mg_cagra_dataset_layout::device_padded) { using padded_ann_t = cuvs::neighbors::cagra::device_padded_index; - auto padded_mds = cuvs::neighbors::make_host_padded_dataset_view(mds); + auto padded_mds = cuvs::core::make_host_padded_dataset_view(mds); auto* mg_index = new mg_cagra_index_t( cuvs::neighbors::cagra::build(*res_ptr, mg_params, padded_mds)); return make_mg_cagra_box(mg_index, mg_cagra_dataset_layout::device_padded); } using standard_ann_t = cuvs::neighbors::cagra::device_standard_index; - auto standard_mds = cuvs::neighbors::make_host_standard_dataset_view(mds); + auto standard_mds = cuvs::core::make_host_standard_dataset_view(mds); auto* mg_index = new mg_cagra_index_t( cuvs::neighbors::cagra::build(*res_ptr, mg_params, standard_mds)); return make_mg_cagra_box(mg_index, mg_cagra_dataset_layout::device_standard); @@ -296,13 +296,13 @@ void _mg_extend(cuvsResources_t res, using padded_ann_t = cuvs::neighbors::cagra::device_padded_index; auto* mg_index_ptr = reinterpret_cast*>(box->index_ptr); - auto new_vectors = cuvs::neighbors::make_host_padded_dataset_view(new_vectors_mds); + auto new_vectors = cuvs::core::make_host_padded_dataset_view(new_vectors_mds); cuvs::neighbors::cagra::extend(*res_ptr, *mg_index_ptr, new_vectors, new_indices_mds); } else { using standard_ann_t = cuvs::neighbors::cagra::device_standard_index; auto* mg_index_ptr = reinterpret_cast*>(box->index_ptr); - auto new_vectors = cuvs::neighbors::make_host_standard_dataset_view(new_vectors_mds); + auto new_vectors = cuvs::core::make_host_standard_dataset_view(new_vectors_mds); cuvs::neighbors::cagra::extend(*res_ptr, *mg_index_ptr, new_vectors, new_indices_mds); } } @@ -372,28 +372,28 @@ extern "C" cuvsError_t cuvsMultiGpuCagraBuild(cuvsResources_t res, if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { auto mds = cuvs::core::from_dlpack>( dataset_tensor); - auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds) + auto layout = cuvs::core::matrix_has_padded_row_width(mds) ? mg_cagra_dataset_layout::device_padded : mg_cagra_dataset_layout::device_standard; index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor, layout)); } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) { auto mds = cuvs::core::from_dlpack>( dataset_tensor); - auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds) + auto layout = cuvs::core::matrix_has_padded_row_width(mds) ? mg_cagra_dataset_layout::device_padded : mg_cagra_dataset_layout::device_standard; index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor, layout)); } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) { auto mds = cuvs::core::from_dlpack>( dataset_tensor); - auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds) + auto layout = cuvs::core::matrix_has_padded_row_width(mds) ? mg_cagra_dataset_layout::device_padded : mg_cagra_dataset_layout::device_standard; index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor, layout)); } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) { auto mds = cuvs::core::from_dlpack>( dataset_tensor); - auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds) + auto layout = cuvs::core::matrix_has_padded_row_width(mds) ? mg_cagra_dataset_layout::device_padded : mg_cagra_dataset_layout::device_standard; index->addr = reinterpret_cast(_mg_build(res, *params, dataset_tensor, layout)); diff --git a/c/src/neighbors/tiered_index.cpp b/c/src/neighbors/tiered_index.cpp index 5ef16833d5..98c74e90c8 100644 --- a/c/src/neighbors/tiered_index.cpp +++ b/c/src/neighbors/tiered_index.cpp @@ -140,8 +140,8 @@ void* _build(cuvsResources_t res, cuvsTieredIndexParams params, DLManagedTensor* case CUVS_TIERED_INDEX_ALGO_CAGRA: { auto build_params = tiered_index::index_params(); convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params); - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)) { - auto padded_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds); + if (cuvs::core::matrix_has_padded_row_width(mds)) { + auto padded_view = cuvs::core::make_device_padded_dataset_view(*res_ptr, mds); auto* ptr = new tiered_index::index>( tiered_index::build(*res_ptr, build_params, padded_view)); return make_tiered_index_box( diff --git a/c/src/preprocessing/quantize/pq.cpp b/c/src/preprocessing/quantize/pq.cpp index 1e3a48694a..394061484f 100644 --- a/c/src/preprocessing/quantize/pq.cpp +++ b/c/src/preprocessing/quantize/pq.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -242,9 +242,10 @@ extern "C" cuvsError_t cuvsProductQuantizerGetPqCodebook(cuvsProductQuantizer_t if (quantizer != nullptr) { auto quant_addr = quantizer->addr; if (quantizer->dtype.code == kDLFloat && quantizer->dtype.bits == 32) { - auto pq_mdspan = + auto const& vpq = (reinterpret_cast*>(quant_addr)) - ->vpq_codebooks.pq_code_book.view(); + ->vpq_codebooks.data(); + auto pq_mdspan = vpq.pq_code_book.view(); cuvs::core::to_dlpack(pq_mdspan, pq_codebook); } else { RAFT_FAIL("Unsupported quantizer dtype: %d and bits: %d", @@ -264,9 +265,10 @@ extern "C" cuvsError_t cuvsProductQuantizerGetVqCodebook(cuvsProductQuantizer_t if (quantizer != nullptr) { auto quant_addr = quantizer->addr; if (quantizer->dtype.code == kDLFloat && quantizer->dtype.bits == 32) { - auto pq_mdspan = + auto const& vpq = (reinterpret_cast*>(quant_addr)) - ->vpq_codebooks.vq_code_book.view(); + ->vpq_codebooks.data(); + auto pq_mdspan = vpq.vq_code_book.view(); cuvs::core::to_dlpack(pq_mdspan, vq_codebook); } else { RAFT_FAIL("Unsupported quantizer dtype: %d and bits: %d", diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 3b5dfd16ab..083def411a 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -271,7 +271,8 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::nn_descent:: } } -inline void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::vpq_params& param) +inline void parse_build_param(const nlohmann::json& conf, + cuvs::preprocessing::quantize::pq::vpq_params& param) { if (conf.contains("pq_bits")) { param.pq_bits = conf.at("pq_bits"); } if (conf.contains("pq_dim")) { param.pq_dim = conf.at("pq_dim"); } @@ -445,7 +446,7 @@ void parse_build_param(const nlohmann::json& conf, nlohmann::json comp_search_conf = collect_conf_with_prefix(conf, "compression_"); if (!comp_search_conf.empty()) { - auto vpq_pams = param.compression.value_or(cuvs::neighbors::vpq_params{}); + auto vpq_pams = param.compression.value_or(cuvs::preprocessing::quantize::pq::vpq_params{}); parse_build_param(comp_search_conf, vpq_pams); param.compression.emplace(vpq_pams); } diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index da5e5baa9d..e6a0755c89 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -85,11 +85,11 @@ template auto make_padded_view(const raft::resources& res, SrcT src, raft::device_matrix& buffer) - -> cuvs::neighbors::device_padded_dataset_view + -> cuvs::core::device_padded_dataset_view { if constexpr (SrcT::accessor_type::is_device_accessible) { - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(src)) { - return cuvs::neighbors::make_device_padded_dataset_view(res, src); + if (cuvs::core::matrix_has_padded_row_width(src)) { + return cuvs::core::make_device_padded_dataset_view(res, src); } } cuvs::neighbors::cagra::detail::copy_with_padding(res, buffer, src); @@ -165,9 +165,9 @@ class cuvs_cagra : public algo, public algo_gpu { using dataset_dependent_params = std::function, cuvs::distance::DistanceType)>; dataset_dependent_params cagra_params; - std::optional compression = std::nullopt; - size_t num_dataset_splits = 1; - CagraMergeType merge_type = CagraMergeType::kPhysical; + std::optional compression = std::nullopt; + size_t num_dataset_splits = 1; + CagraMergeType merge_type = CagraMergeType::kPhysical; cuvs::neighbors::cagra::merge_params merge_params; }; @@ -267,7 +267,8 @@ class cuvs_cagra : public algo, public algo_gpu { std::shared_ptr>> sub_dataset_buffers_ = std::make_shared>>(); - std::shared_ptr> vpq_dataset_; + std::shared_ptr> + vpq_dataset_; std::shared_ptr> vpq_index_; inline rmm::device_async_resource_ref get_mr(AllocatorType mem_type) @@ -302,7 +303,7 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) // (and dispatches to ACE when it is configured), so the dataset is not uploaded here at all. // The single device copy needed for search is made later, by set_search_param. host_index_ = std::make_shared(cuvs::neighbors::cagra::build( - handle_, host_params, cuvs::neighbors::make_host_standard_dataset_view(dataset_view_host))); + handle_, host_params, cuvs::core::make_host_standard_dataset_view(dataset_view_host))); index_ = std::make_shared(detail::to_graph_only_index(handle_, *host_index_)); // The graph moved into the index along with the file descriptors; nothing views the host @@ -348,7 +349,7 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) // As in the single-split case: graph only, the rows are uploaded by set_search_dataset. sub_host_indices_.push_back( std::make_shared(cuvs::neighbors::cagra::build( - handle_, host_params, cuvs::neighbors::make_host_standard_dataset_view(sub_host)))); + handle_, host_params, cuvs::core::make_host_standard_dataset_view(sub_host)))); sub_index = detail::to_graph_only_index(handle_, *sub_host_indices_.back()); if (sub_index.graph_fd().has_value()) { sub_host_indices_.pop_back(); } } else { @@ -370,10 +371,10 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) for (auto* index : indices) { merged_rows += static_cast(index->size()); } - auto const stride = static_cast( - cuvs::neighbors::cagra_required_row_width(static_cast(dim_))); + auto const stride = + static_cast(cuvs::core::padded_row_width(static_cast(dim_))); *dataset_ = raft::make_device_matrix(handle_, merged_rows, stride); - auto merged_dataset_view = cuvs::neighbors::device_padded_dataset_view( + auto merged_dataset_view = cuvs::core::device_padded_dataset_view( raft::make_const_mdspan(dataset_->view()), static_cast(dim_)); index_ = std::make_shared(cuvs::neighbors::cagra::merge(handle_, @@ -410,13 +411,15 @@ void cuvs_cagra::compress_dataset(const T* dataset, size_t nrow) // make_vpq_dataset() reads the rows wherever they are: host-resident ones are subsampled and // encoded in bounded batches instead of being staged on the device. auto src = raft::make_device_matrix_view(dataset, rows, dim_); - vpq_dataset_ = std::make_shared>( - cuvs::preprocessing::quantize::pq::make_vpq_dataset(handle_, *index_params_.compression, src)); + vpq_dataset_ = + std::make_shared>( + cuvs::preprocessing::quantize::pq::make_vpq_dataset( + handle_, *index_params_.compression, src)); vpq_index_ = std::make_shared>( handle_, parse_metric_type(metric_), vpq_dataset_->as_dataset_view(), index_->graph()); // Search runs on the compressed rows and the graph, so release the dense copy of the dataset. - cuvs::neighbors::device_padded_dataset_view empty_dv( + cuvs::core::device_padded_dataset_view empty_dv( raft::make_device_matrix_view(static_cast(nullptr), 0, this->dim_), this->dim_); *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv); *dataset_ = raft::make_device_matrix(handle_, 0, 0); @@ -483,7 +486,7 @@ void cuvs_cagra::set_search_param(const search_param_base& param, // First free up existing memory *dataset_ = raft::make_device_matrix(handle_, 0, 0); - cuvs::neighbors::device_padded_dataset_view empty_dv( + cuvs::core::device_padded_dataset_view empty_dv( raft::make_device_matrix_view(static_cast(nullptr), 0, this->dim_), this->dim_); *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv); @@ -494,7 +497,7 @@ void cuvs_cagra::set_search_param(const search_param_base& param, auto mr = get_mr(dataset_mem_); cuvs::neighbors::cagra::detail::copy_with_padding(handle_, *dataset_, *input_dataset_v_, mr); - cuvs::neighbors::device_padded_dataset_view dv( + cuvs::core::device_padded_dataset_view dv( raft::make_device_matrix_view( dataset_->data_handle(), dataset_->extent(0), dataset_->extent(1)), this->dim_); diff --git a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h index 7a89b68394..439697a19f 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h @@ -95,8 +95,8 @@ void cuvs_mg_cagra::build(const T* dataset, size_t nrow) raft::make_host_matrix_view(dataset, nrow, dim_); // The row alignment of the host view is irrelevant: every per-rank device shard is padded // individually during the multi-GPU build. - cuvs::neighbors::host_padded_dataset_view dataset_view(dataset_mds, - static_cast(dim_)); + cuvs::core::host_padded_dataset_view dataset_view(dataset_mds, + static_cast(dim_)); auto idx = cuvs::neighbors::cagra::build(clique_, build_params, dataset_view); index_ = std::make_shared< cuvs::neighbors::mg_index, T, IdxT>>( diff --git a/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff b/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff index 738ae8fa3b..fcb61a8b8b 100644 --- a/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff +++ b/cpp/cmake/patches/faiss-1.14-cuvs-26.08.diff @@ -23,7 +23,7 @@ index 4be7cebb8..42315d027 100644 raft::make_device_matrix_view( train_dataset, n, dim / 8); + auto dataset_view = -+ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset_mds); ++ cuvs::core::make_device_padded_dataset_view(raft_handle, dataset_mds); cuvs_index = std::make_shared< cuvs::neighbors::cagra::index>( @@ -39,7 +39,7 @@ index 4be7cebb8..42315d027 100644 auto dataset_mds = raft::make_host_matrix_view( train_dataset, n, dim / 8); + host_to_device_dataset_ = -+ cuvs::neighbors::make_device_padded_dataset(raft_handle, dataset_mds); ++ cuvs::core::make_device_padded_dataset(raft_handle, dataset_mds); cuvs_index = std::make_shared< cuvs::neighbors::cagra::index>( @@ -55,7 +55,7 @@ index 4be7cebb8..42315d027 100644 auto dataset = raft::make_device_matrix_view( x, n, dim_ / 8); + auto dataset_view = -+ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset); ++ cuvs::core::make_device_padded_dataset_view(raft_handle, dataset); cuvs_index = std::make_shared< cuvs::neighbors::cagra::index>( cuvs::neighbors::cagra::build( @@ -66,7 +66,7 @@ index 4be7cebb8..42315d027 100644 auto dataset = raft::make_host_matrix_view( x, n, dim_ / 8); + host_to_device_dataset_ = -+ cuvs::neighbors::make_device_padded_dataset(raft_handle, dataset); ++ cuvs::core::make_device_padded_dataset(raft_handle, dataset); cuvs_index = std::make_shared< cuvs::neighbors::cagra::index>( cuvs::neighbors::cagra::build( @@ -86,7 +86,7 @@ index 4be7cebb8..42315d027 100644 storage_, n_, dim_ / 8); - cuvs_index->update_dataset(raft_handle, dataset); + auto dataset_view = -+ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset); ++ cuvs::core::make_device_padded_dataset_view(raft_handle, dataset); + *cuvs_index = cuvs::neighbors::cagra::update_dataset( + raft_handle, std::move(*cuvs_index), dataset_view); } else { @@ -95,7 +95,7 @@ index 4be7cebb8..42315d027 100644 storage_, n_, dim_ / 8); - cuvs_index->update_dataset(raft_handle, dataset); + host_to_device_dataset_ = -+ cuvs::neighbors::make_device_padded_dataset(raft_handle, host_dataset); ++ cuvs::core::make_device_padded_dataset(raft_handle, host_dataset); + *cuvs_index = cuvs::neighbors::cagra::update_dataset( + raft_handle, + std::move(*cuvs_index), @@ -135,7 +135,7 @@ index 3c7a9eb68..1038380a2 100644 size_t nn_descent_niter_ = 20; + /// Device padded copy when `storage_` is host memory (KNN-graph ctor path). -+ std::unique_ptr> ++ std::unique_ptr> + host_to_device_dataset_; + /// Instance of trained cuVS CAGRA index @@ -166,7 +166,7 @@ index 6ff3419b7..f03a5e5f0 100644 auto dataset_mds = raft::make_device_matrix_view( dataset, n, dim); + auto dataset_view = -+ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset_mds); ++ cuvs::core::make_device_padded_dataset_view(raft_handle, dataset_mds); cuvs_index = std::make_shared< cuvs::neighbors::cagra::index>( @@ -182,7 +182,7 @@ index 6ff3419b7..f03a5e5f0 100644 auto dataset_mds = raft::make_host_matrix_view( dataset, n, dim); + host_to_device_dataset_ = -+ cuvs::neighbors::make_device_padded_dataset(raft_handle, dataset_mds); ++ cuvs::core::make_device_padded_dataset(raft_handle, dataset_mds); cuvs_index = std::make_shared< cuvs::neighbors::cagra::index>( @@ -198,7 +198,7 @@ index 6ff3419b7..f03a5e5f0 100644 auto dataset = raft::make_device_matrix_view( x, n, dim_); + auto dataset_view = -+ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset); ++ cuvs::core::make_device_padded_dataset_view(raft_handle, dataset); cuvs_index = std::make_shared< cuvs::neighbors::cagra::index>( cuvs::neighbors::cagra::build( @@ -209,7 +209,7 @@ index 6ff3419b7..f03a5e5f0 100644 auto dataset = raft::make_host_matrix_view(x, n, dim_); + host_to_device_dataset_ = -+ cuvs::neighbors::make_device_padded_dataset(raft_handle, dataset); ++ cuvs::core::make_device_padded_dataset(raft_handle, dataset); cuvs_index = std::make_shared< cuvs::neighbors::cagra::index>( cuvs::neighbors::cagra::build( @@ -228,7 +228,7 @@ index 6ff3419b7..f03a5e5f0 100644 storage_, n_, dim_); - cuvs_index->update_dataset(raft_handle, dataset); + auto dataset_view = -+ cuvs::neighbors::make_device_padded_dataset_view(raft_handle, dataset); ++ cuvs::core::make_device_padded_dataset_view(raft_handle, dataset); + *cuvs_index = cuvs::neighbors::cagra::update_dataset( + raft_handle, std::move(*cuvs_index), dataset_view); } else { @@ -237,7 +237,7 @@ index 6ff3419b7..f03a5e5f0 100644 storage_, n_, dim_); - cuvs_index->update_dataset(raft_handle, dataset); + host_to_device_dataset_ = -+ cuvs::neighbors::make_device_padded_dataset(raft_handle, host_dataset); ++ cuvs::core::make_device_padded_dataset(raft_handle, host_dataset); + *cuvs_index = cuvs::neighbors::cagra::update_dataset( + raft_handle, + std::move(*cuvs_index), @@ -278,7 +278,7 @@ index b4c1989c6..3ad960b92 100644 bool guarantee_connectivity_ = false; + /// Device padded copy when `storage_` is host memory (KNN-graph ctor path). -+ std::unique_ptr> ++ std::unique_ptr> + host_to_device_dataset_; + /// Instance of trained cuVS CAGRA index diff --git a/cpp/include/cuvs/core/dataset.hpp b/cpp/include/cuvs/core/dataset.hpp new file mode 100644 index 0000000000..2786d92fe0 --- /dev/null +++ b/cpp/include/cuvs/core/dataset.hpp @@ -0,0 +1,928 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include // get_device_for_address, copy_matrix +#include // rounding up + +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace CUVS_EXPORT cuvs { +namespace core { + +/** + * @brief Spec-based `dataset` / `dataset_view`. + * + * `dataset` and `dataset_view` are single generic templates that know + * nothing about any particular kind of dataset. They hold exactly one payload (`data_type` / + * `view_type`, chosen by the spec) and expose only what every dataset has: `n_rows()`, `dim()`, + * `as_matrix_view()`, `as_dataset_view()` and `data()`. Each is a one-line forward to one of the + * three spec functions `get_data_view()`, `get_n_rows()` and `get_dim()`. Anything else a kind + * needs (e.g. codebooks or quantizers of a compressed dataset) is state and methods of that + * kind's payload type, reached through `data()`; `dataset`/`dataset_view` never name or branch on + * it. Compressed kinds define their payloads and specs in their own headers (quantize/pq.hpp, + * quantize/bbq.hpp) as children of these two structs. `dataset` and + * `dataset_view` are deliberately two independent, non-inheriting types (no shared_ptr, no + * "sometimes owning" object): `dataset` holds the owning payload, `dataset_view` the + * corresponding non-owning payload. + */ + +namespace detail { + +// Default owning/view accessors for public dataset aliases. +template +using device_owning_accessor = raft::device_accessor>; + +template +using host_owning_accessor = raft::host_accessor>; + +template +using device_view_accessor = raft::device_accessor>; + +template +using host_view_accessor = raft::host_accessor>; + +/** View accessor paired with an owning dataset accessor (same residency). */ +template +using dataset_view_accessor_for_owning = std::conditional_t, + host_view_accessor>; + +/** Owning accessor paired with a view accessor (same residency). */ +template +using dataset_owning_accessor_for_view = std::conditional_t, + host_owning_accessor>; + +// Accessor here is already device_owning_accessor / host_owning_accessor at every +// call site -- exactly the container policy raft::device_mdarray/host_mdarray default to for +// element type DataT -- so pass it straight through instead of re-deriving a +// raft::device_matrix/host_matrix from scratch. +template +using dense_owning_matrix = + raft::mdarray, raft::row_major, Accessor>; + +template +using dense_view_matrix = raft::mdspan, + raft::row_major, + dataset_view_accessor_for_owning>; + +// ----------------------------------------------------------------------------- +// empty +// ----------------------------------------------------------------------------- + +template +struct empty_dataset_storage { + uint32_t suggested_dim{}; + empty_dataset_storage() noexcept = default; + explicit empty_dataset_storage(uint32_t dim) noexcept : suggested_dim(dim) {} + [[nodiscard]] auto n_rows() const noexcept -> IdxT { return 0; } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return suggested_dim; } +}; + +// ----------------------------------------------------------------------------- +// dense row-major (logical dim may differ from row pitch; shared by padded & standard) +// ----------------------------------------------------------------------------- + +/** + * Dense row-major payload shared by padded and standard dataset specs, owning and non-owning + * alike: the owning payload passes its `raft::mdarray` as `BaseT`, the view payload passes the + * `raft::mdspan`. Publicly inherits from `BaseT` so `view()`/`data_handle()`/`extent()` etc. are + * reused as-is rather than hand-forwarded; `logical_dim_` is the only state this struct adds. + * + * Template parameters: + * - BaseT: the owning matrix (`raft::mdarray`) or the non-owning row-major view (`raft::mdspan`). + */ +template +struct dense_row_major_storage : public BaseT { + using index_type = typename BaseT::index_type; + + uint32_t logical_dim_{}; + + // BaseT (mdarray/mdspan) also has its own stride(size_t); pull it back into scope since + // declaring our own no-arg stride() below would otherwise hide it entirely (C++ name hiding), + // and the body of that stride() itself needs to call the inherited one. + using BaseT::stride; + + dense_row_major_storage() noexcept = default; + + // Takes BaseT by value so an owning matrix is moved in and a view is simply copied. + explicit dense_row_major_storage(BaseT base) noexcept + : BaseT(std::move(base)), logical_dim_(static_cast(this->extent(1))) + { + } + + dense_row_major_storage(BaseT base, uint32_t logical_dim) noexcept + : BaseT(std::move(base)), logical_dim_(logical_dim) + { + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type { return this->extent(0); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } + [[nodiscard]] auto stride() const noexcept -> uint32_t + { + return static_cast(BaseT::stride(0) > 0 ? BaseT::stride(0) : this->extent(1)); + } +}; + +/** Spec-side implementation shared by `padded_dataset_spec`/`standard_dataset_spec`; those two + * stay distinct top-level types (identical bodies) purely so classification traits can tell them + * apart -- exactly mirroring today's `padded_dataset_container`/`standard_dataset_container`, + * which are likewise two differently-named tags over one shared storage implementation. */ +template +struct dense_dataset_spec_impl { + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using MatrixT = dense_owning_matrix; + using ViewT = dense_view_matrix; + using data_type = dense_row_major_storage; + using view_type = dense_row_major_storage; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return view_type(data.view(), data.dim()); + } + template + [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type + { + return data.n_rows(); + } + template + [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t + { + return data.dim(); + } + }; +}; + +} // namespace detail + +// ----------------------------------------------------------------------------- +// Public specs -- the only place per-kind logic lives. +// ----------------------------------------------------------------------------- + +template +struct empty_dataset_spec { + using accessor_type = Accessor; + template + using rebind_accessor = empty_dataset_spec; + + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using data_type = detail::empty_dataset_storage; + using view_type = detail::empty_dataset_storage; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return data; + } + [[nodiscard]] static auto get_n_rows(data_type const& data) noexcept -> index_type + { + return static_cast(data.n_rows()); + } + [[nodiscard]] static auto get_dim(data_type const& data) noexcept -> uint32_t + { + return data.dim(); + } + }; +}; + +template +struct padded_dataset_spec { + using accessor_type = ContainerPolicy; + template + using rebind_accessor = padded_dataset_spec; + template + struct apply : detail::dense_dataset_spec_impl::template apply {}; +}; + +template +struct standard_dataset_spec { + using accessor_type = ContainerPolicy; + template + using rebind_accessor = standard_dataset_spec; + template + struct apply : detail::dense_dataset_spec_impl::template apply {}; +}; + +// ----------------------------------------------------------------------------- +// dataset / dataset_view +// ----------------------------------------------------------------------------- + +/** Non-owning dataset view: holds only the view-shaped payload. Deliberately not derived from + * `dataset` -- a view type holds "all view state" with no inheritance and no shared ownership tying + * it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as `dataset`, fed + * the view payload instead of the owning one. */ +template +struct dataset_view { + using spec_type = typename SpecT::template apply; + using value_type = typename spec_type::value_type; + using index_type = typename spec_type::index_type; + using view_type = typename spec_type::view_type; + + dataset_view() noexcept = default; + + // Already-constructed view payload -- the shape `as_dataset_view()` always constructs with, for + // every kind. Not a template, so it's preferred over the forwarding constructor below whenever + // both could apply. + explicit dataset_view(view_type data_view) noexcept : data_view_{data_view} {} + + // Forward raw constructor args straight to view_type's own constructor (e.g. (ViewT, uint32_t + // logical_dim) for dense, (uint32_t dim) for empty) -- preserves today's direct-construction + // call sites (e.g. `device_padded_dataset_view(raw_mdspan, dim)`) unchanged. Explicit + // only for a single argument, so that a raw matrix view never converts to a dataset view + // implicitly (which would let it silently bind to overloads meant for dataset views), while + // multi-argument brace initialization such as `return {view, dim};` keeps working. + template + explicit(sizeof...(Args) == 1) dataset_view(Args&&... args) + requires(std::is_constructible_v) + : data_view_(std::forward(args)...) + { + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type + { + return spec_type::get_n_rows(data_view_); + } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return spec_type::get_dim(data_view_); } + [[nodiscard]] auto as_matrix_view() const noexcept -> view_type { return data_view_; } + + /** The view payload; kind-specific state and methods are reached through it. */ + [[nodiscard]] auto data() const noexcept -> view_type const& { return data_view_; } + [[nodiscard]] auto data() noexcept -> view_type& { return data_view_; } + + private: + view_type data_view_{}; +}; + +/** Owning dataset: value-held payload (no shared_ptr -- exclusive ownership). Every member is a + * one-line forward to `spec_type::get_*` or to the payload; all per-kind state and logic lives in + * the spec's `data_type`, never inside this struct. */ +template +struct dataset { + using spec_type = typename SpecT::template apply; + using value_type = typename spec_type::value_type; + using index_type = typename spec_type::index_type; + using data_type = typename spec_type::data_type; + + // Forward constructor args straight to data_type's own constructor (e.g. (MatrixT&&, uint32_t + // logical_dim) for dense, (uint32_t dim) for empty, or whatever a compressed kind's payload + // takes). Explicit only for a single argument, so that a lone value never converts to a dataset + // implicitly, while multi-argument brace initialization such as `return {matrix, dim};` keeps + // working as it did before datasets were generic. + template + explicit(sizeof...(Args) == 1) dataset(Args&&... args) + requires(std::is_constructible_v) + : data_(std::forward(args)...) + { + } + + [[nodiscard]] auto n_rows() const noexcept -> index_type { return spec_type::get_n_rows(data_); } + [[nodiscard]] auto dim() const noexcept -> uint32_t { return spec_type::get_dim(data_); } + /** The spec-defined non-owning view of the payload (an mdspan derivative for dense kinds). */ + [[nodiscard]] auto as_matrix_view() const noexcept { return spec_type::get_data_view(data_); } + [[nodiscard]] auto as_dataset_view() const noexcept -> dataset_view + { + return dataset_view(as_matrix_view()); + } + + /** The owning payload; kind-specific state and methods are reached through it. */ + [[nodiscard]] auto data() const noexcept -> data_type const& { return data_; } + [[nodiscard]] auto data() noexcept -> data_type& { return data_; } + + private: + data_type data_; +}; + +/** + * @brief Aliases for concrete `dataset` / `dataset_view` layouts. + */ +template +using device_empty_dataset = + dataset>>; + +template +using device_empty_dataset_view = + dataset_view>>; + +template +using host_empty_dataset = + dataset>>; + +template +using host_empty_dataset_view = + dataset_view>>; + +template +using device_padded_dataset = + dataset>>; + +template +using device_padded_dataset_view = + dataset_view>>; + +template +using host_padded_dataset = + dataset>>; + +template +using host_padded_dataset_view = + dataset_view>>; + +template +using device_standard_dataset = + dataset>>; + +template +using device_standard_dataset_view = + dataset_view>>; + +template +using host_standard_dataset = + dataset>>; + +template +using host_standard_dataset_view = + dataset_view>>; + +// Maps a dataset view type to its owning (allocating) dataset counterpart. Trivial and total under +// the Spec design: the owning type for `dataset_view` is always +// `dataset` +// -- no per-kind specialization table needed (unlike the old Container-tagged design). +template +struct owning_dataset_for_view; + +template +struct owning_dataset_for_view> { + using type = dataset; +}; + +template +using owning_dataset_for_view_t = typename owning_dataset_for_view::type; + +// ----------------------------------------------------------------------------- +// Spec-kind classification. Only the kinds that live in this header are named here; every other +// kind (e.g. the compressed kinds in quantize/pq.hpp and quantize/bbq.hpp) defines its own spec +// predicate next to its own spec and reuses `dataset_view_has_spec_v` below. +// ----------------------------------------------------------------------------- + +template +struct is_empty_spec : std::false_type {}; +template +struct is_empty_spec> : std::true_type {}; +template +inline constexpr bool is_empty_spec_v = is_empty_spec::value; + +template +struct is_padded_spec : std::false_type {}; +template +struct is_padded_spec> : std::true_type {}; +template +inline constexpr bool is_padded_spec_v = is_padded_spec::value; + +template +struct is_standard_spec : std::false_type {}; +template +struct is_standard_spec> : std::true_type {}; +template +inline constexpr bool is_standard_spec_v = is_standard_spec::value; + +/** Owning-side kind traits (true for both `dataset<...>` and `dataset_view<...>` of that kind). */ +template +struct is_padded_dataset : std::false_type {}; +template +struct is_padded_dataset> : std::bool_constant> {}; +template +struct is_padded_dataset> + : std::bool_constant> {}; +template +inline constexpr bool is_padded_dataset_v = is_padded_dataset::value; + +template +struct is_standard_dataset : std::false_type {}; +template +struct is_standard_dataset> + : std::bool_constant> {}; +template +struct is_standard_dataset> + : std::bool_constant> {}; +template +inline constexpr bool is_standard_dataset_v = is_standard_dataset::value; + +// ----------------------------------------------------------------------------- +// Dataset view compile-time classification (replaces runtime std::variant dispatch). +// ----------------------------------------------------------------------------- + +/** Any type that behaves like a dataset: it exposes a row count (`n_rows()`) and a logical + * dimension (`dim()`). This is a structural check, so owning datasets and dataset views both + * satisfy it. To ask whether a type is literally a `dataset_view<...>`, use `is_dataset_view_v`. */ +template +concept dataset_like = requires(V const& v) { + { v.n_rows() } -> std::convertible_to; + { v.dim() } -> std::convertible_to; +}; + +template +using dataset_view_type_t = std::remove_cvref_t; + +/** True for any `dataset_view<...>` specialization. Evaluates to `false` (never a hard error) for + * everything else, e.g. a plain mdspan passed to a deprecated `build(matrix_view)` overload. */ +template +struct is_dataset_view : std::false_type {}; +template +struct is_dataset_view> : std::true_type {}; +template +inline constexpr bool is_dataset_view_v = is_dataset_view>::value; + +/** True when `V` is a `dataset_view` whose spec satisfies the predicate `SpecPred::value`. + * This is how a kind that lives outside this header classifies its own views. */ +template typename SpecPred> +struct dataset_view_has_spec : std::false_type {}; +template typename SpecPred> +struct dataset_view_has_spec, SpecPred> + : std::bool_constant::value> {}; +template typename SpecPred> +inline constexpr bool dataset_view_has_spec_v = + dataset_view_has_spec, SpecPred>::value; + +/** True when the dataset view accessor is device-accessible. */ +template +struct dataset_view_is_device_accessible : std::false_type {}; + +template +struct dataset_view_is_device_accessible> + : std::bool_constant {}; + +template +inline constexpr bool dataset_view_is_device_accessible_v = + dataset_view_is_device_accessible>::value; + +template +inline constexpr bool is_device_empty_dataset_view_v = + dataset_view_has_spec_v && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_empty_dataset_view_v = + dataset_view_has_spec_v && !dataset_view_is_device_accessible_v; + +/** True for any empty dataset view (device or host). */ +template +inline constexpr bool is_empty_dataset_view_v = + is_device_empty_dataset_view_v || is_host_empty_dataset_view_v; + +template +inline constexpr bool is_device_padded_dataset_view_v = + dataset_view_has_spec_v && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_padded_dataset_view_v = + dataset_view_has_spec_v && !dataset_view_is_device_accessible_v; + +/** True for either `device_padded_dataset_view` or `host_padded_dataset_view`. */ +template +inline constexpr bool is_padded_dataset_view_v = + is_device_padded_dataset_view_v || is_host_padded_dataset_view_v; + +template +inline constexpr bool is_device_standard_dataset_view_v = + dataset_view_has_spec_v && dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_standard_dataset_view_v = + dataset_view_has_spec_v && !dataset_view_is_device_accessible_v; + +/** True for either `device_standard_dataset_view` or `host_standard_dataset_view`. */ +template +inline constexpr bool is_standard_dataset_view_v = + is_device_standard_dataset_view_v || is_host_standard_dataset_view_v; + +/** True for any device-resident dataset view. */ +template +inline constexpr bool is_device_dataset_view_v = + is_dataset_view_v && dataset_view_is_device_accessible_v; + +/** True for any host-resident dataset view. */ +template +inline constexpr bool is_host_dataset_view_v = + is_dataset_view_v && !dataset_view_is_device_accessible_v; + +/** + * Generic accessor retargeting while preserving the spec kind and value/index types: + * `dataset> -> dataset>` + * `dataset_view> -> dataset_view>` + * Every spec provides `rebind_accessor` for this, so this header does not need to + * know about any particular kind. + */ +template +struct with_accessor; + +template +struct with_accessor, NewAccessor> { + using type = dataset>; +}; + +template +struct with_accessor, NewAccessor> { + using type = dataset_view>; +}; + +template +using with_accessor_t = + typename with_accessor, NewAccessor>::type; + +/** Map any host accessor to its device counterpart (same payload policy). */ +template +struct to_device_accessor { + using type = Accessor; +}; + +template +struct to_device_accessor> { + using type = detail::device_view_accessor; +}; + +template +struct to_device_accessor> { + using type = detail::device_owning_accessor; +}; + +template +using to_device_accessor_t = typename to_device_accessor::type; + +/** Maps a host dataset view type to its device-resident counterpart. */ +template +struct device_counterpart; + +template +struct device_counterpart> { + using type = with_accessor_t, + to_device_accessor_t>; +}; + +template +using device_counterpart_t = typename device_counterpart>::type; + +/** + * True when a host view `H` and device view `D` represent the same storage kind and differ + * only in residency (host vs. device). Used by host/device conversion helpers. + */ +template +inline constexpr bool compatible_host_device_dataset_views_v = + is_host_dataset_view_v && is_device_dataset_view_v && + std::is_same_v, dataset_view_type_t>; + +/** True for device padded or standard (dense row-major) dataset views. */ +template +inline constexpr bool is_dense_row_major_device_dataset_view_v = + is_device_padded_dataset_view_v || is_device_standard_dataset_view_v; + +/** True for host or device padded or standard (dense row-major) dataset views. */ +template +inline constexpr bool is_dense_row_major_dataset_view_v = + is_padded_dataset_view_v || is_standard_dataset_view_v; + +/** Element type `T` of a dataset view, deduced from the view. Trivial under the Spec design: every + * `dataset_view` already carries `T` directly. */ +template +using dataset_view_value_t = typename dataset_view_type_t::value_type; + +// ----------------------------------------------------------------------------- +// Padded row width in elements (shared by the make_*_padded_dataset* factories and row-width +// checks). +// ----------------------------------------------------------------------------- + +/** + * @brief Minimum row width in elements (the leading dimension) for `logical_columns` feature + * columns, such that each row occupies a whole multiple of `align_bytes` bytes (default 16, + * combined with `sizeof` of the element type). + */ +[[nodiscard]] inline uint32_t padded_row_width(uint32_t logical_columns, + std::size_t sizeof_value, + uint32_t align_bytes = 16) +{ + return static_cast( + raft::round_up_safe(static_cast(logical_columns) * sizeof_value, + std::lcm(align_bytes, static_cast(sizeof_value))) / + sizeof_value); +} + +template +[[nodiscard]] inline uint32_t padded_row_width(uint32_t logical_columns, uint32_t align_bytes = 16) +{ + return padded_row_width(logical_columns, sizeof(ValueT), align_bytes); +} + +/** Actual row width in elements (leading dimension) of a 2D row-major matrix view. */ +template +[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::device_matrix_view m) +{ + return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); +} + +template +[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::host_matrix_view m) +{ + return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); +} + +/** + * @brief True if the matrix's row width in elements equals `padded_row_width` for `m.extent(1)` + * and element type `T`, i.e. its rows are already padded. + */ +template +[[nodiscard]] inline bool matrix_has_padded_row_width(raft::device_matrix_view m, + uint32_t align_bytes = 16) +{ + using value_type = std::remove_const_t; + const uint32_t need = + padded_row_width(static_cast(m.extent(1)), align_bytes); + return matrix_actual_row_width(m) == need; +} + +template +[[nodiscard]] inline bool matrix_has_padded_row_width(raft::host_matrix_view m, + uint32_t align_bytes = 16) +{ + using value_type = std::remove_const_t; + const uint32_t need = + padded_row_width(static_cast(m.extent(1)), align_bytes); + return matrix_actual_row_width(m) == need; +} + +namespace detail { + +template +[[nodiscard]] inline uint32_t mdspan_row_stride_elements(SrcT const& src) +{ + return src.stride(0) > 0 ? static_cast(src.stride(0)) + : static_cast(src.extent(1)); +} + +template +[[nodiscard]] inline ValueT* expect_device_accessible_data_handle(SrcT const& src, + char const* error_msg) +{ + cudaPointerAttributes ptr_attrs; + RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); + // `devicePointer` is relative to the *current* device: it is null for an allocation owned by + // another device without peer access, even though that allocation is perfectly usable once the + // caller switches to the owning device (as the multi-GPU paths do). Accept device and managed + // allocations on their own merit and only consult `devicePointer` for host memory, which needs a + // mapping to be reachable at all. + if (ptr_attrs.type == cudaMemoryTypeDevice || ptr_attrs.type == cudaMemoryTypeManaged) { + return const_cast(src.data_handle()); + } + auto* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); + RAFT_EXPECTS(device_ptr != nullptr, "%s", error_msg); + return device_ptr; +} + +template +[[nodiscard]] inline ViewT make_device_dense_row_major_view_from_src(SrcT const& src, + uint32_t logical_dim) +{ + auto* device_ptr = expect_device_accessible_data_handle( + src, "make_device_*_dataset_view: source must be device-accessible."); + auto v = raft::make_device_matrix_view( + device_ptr, src.extent(0), static_cast(mdspan_row_stride_elements(src))); + return ViewT(v, logical_dim); +} + +template +[[nodiscard]] inline ViewT make_host_dense_row_major_view_from_src(SrcT const& src, + uint32_t logical_dim) +{ + RAFT_EXPECTS(raft::get_device_for_address(src.data_handle()) == -1, + "make_host_*_dataset_view: source must be host-accessible."); + auto v = raft::make_host_matrix_view(const_cast(src.data_handle()), + src.extent(0), + static_cast(mdspan_row_stride_elements(src))); + return ViewT(v, logical_dim); +} + +template +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) + -> std::unique_ptr +{ + uint32_t const src_stride = mdspan_row_stride_elements(src); + RAFT_EXPECTS(logical_dim <= target_stride, + "logical dim (%u) must not exceed row stride (%u).", + static_cast(logical_dim), + static_cast(target_stride)); + RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, + "Source row length must not exceed required stride."); + cudaPointerAttributes ptr_attrs; + 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) { + RAFT_EXPECTS(false, + "source is device and stride is already correct. " + "Use %s() to get a view instead.", + view_factory_name); + } + auto out_array = raft::make_device_matrix(res, src.extent(0), target_stride); + RAFT_CUDA_TRY(cudaMemsetAsync(out_array.data_handle(), + 0, + out_array.size() * sizeof(ValueT), + raft::resource::get_cuda_stream(res).get())); + raft::copy_matrix(out_array.data_handle(), + target_stride, + src.data_handle(), + src_stride, + logical_dim, + src.extent(0), + raft::resource::get_cuda_stream(res)); + return std::make_unique(std::move(out_array), logical_dim); +} + +template +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) + -> std::unique_ptr +{ + uint32_t const src_stride = mdspan_row_stride_elements(src); + constexpr bool device_src = SrcT::accessor_type::is_device_accessible; + RAFT_EXPECTS(logical_dim <= target_stride, + "logical dim (%u) must not exceed row stride (%u).", + static_cast(logical_dim), + static_cast(target_stride)); + if (!device_src && src_stride == target_stride) { + RAFT_EXPECTS(false, + "source stride is already correct. Use %s() to get a view instead.", + view_factory_name); + } + RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, + "Source row length must not exceed required stride."); + auto out_array = raft::make_host_matrix(src.extent(0), target_stride); + std::memset(out_array.data_handle(), 0, out_array.size() * sizeof(ValueT)); + raft::copy_matrix(out_array.data_handle(), + target_stride, + src.data_handle(), + src_stride, + logical_dim, + src.extent(0), + raft::resource::get_cuda_stream(res)); + if (device_src) { raft::resource::sync_stream(res); } + return std::make_unique(std::move(out_array), logical_dim); +} + +} // namespace detail + +template +auto make_device_padded_dataset_view(const raft::resources& res, + SrcT const& src, + uint32_t align_bytes = 16) + -> device_padded_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t required_stride = + padded_row_width(static_cast(src.extent(1)), align_bytes); + RAFT_EXPECTS( + detail::mdspan_row_stride_elements(src) == required_stride, + "make_device_padded_dataset_view: stride is incorrect (required stride for alignment). " + "Use make_device_padded_dataset() to get an owning padded copy."); + return detail::make_device_dense_row_major_view_from_src< + value_type, + index_type, + device_padded_dataset_view>(src, static_cast(src.extent(1))); +} + +template +auto make_device_padded_dataset(const raft::resources& res, + SrcT const& src, + uint32_t align_bytes = 16) + -> std::unique_ptr> +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t const logical_dim = static_cast(src.extent(1)); + uint32_t const required_stride = padded_row_width(logical_dim, align_bytes); + 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"); +} + +template +auto make_host_padded_dataset_view(SrcT const& src, uint32_t align_bytes = 16) + -> host_padded_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t required_stride = + padded_row_width(static_cast(src.extent(1)), align_bytes); + RAFT_EXPECTS( + detail::mdspan_row_stride_elements(src) == required_stride, + "make_host_padded_dataset_view: stride is incorrect (required stride for alignment). " + "Use make_host_padded_dataset() to get an owning padded copy."); + return detail::make_host_dense_row_major_view_from_src< + value_type, + index_type, + host_padded_dataset_view>(src, static_cast(src.extent(1))); +} + +template +auto make_host_padded_dataset(const raft::resources& res, + SrcT const& src, + uint32_t align_bytes = 16) + -> std::unique_ptr> +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + uint32_t const logical_dim = static_cast(src.extent(1)); + uint32_t const required_stride = padded_row_width(logical_dim, align_bytes); + 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"); +} + +template +auto make_device_standard_dataset_view(SrcT const& src) + -> device_standard_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + return detail::make_device_dense_row_major_view_from_src< + value_type, + index_type, + device_standard_dataset_view>(src, + static_cast(src.extent(1))); +} + +/** + * @brief Create an owning device standard dataset with explicit row layout. + * + * Internal use only: the sole caller today deserializes a dataset from disk and must pass + * wire-format `(logical_dim, stride)` because the deserialized host buffer is tight `[n_rows x + * dim]` while the on-disk stride may be larger. Do not call from user code; prefer + * `make_device_standard_dataset_view()` when wrapping existing correctly-strided storage. + */ +template +auto make_device_standard_dataset(const raft::resources& res, + SrcT const& src, + uint32_t logical_dim, + uint32_t target_stride) + -> std::unique_ptr> +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + return detail::make_device_dense_row_major_dataset_from_src< + device_standard_dataset, + value_type, + index_type>(res, src, logical_dim, target_stride, "make_device_standard_dataset_view"); +} + +template +auto make_host_standard_dataset_view(SrcT const& src) + -> host_standard_dataset_view +{ + using value_type = typename SrcT::value_type; + using index_type = typename SrcT::index_type; + return detail::make_host_dense_row_major_view_from_src< + value_type, + index_type, + host_standard_dataset_view>(src, static_cast(src.extent(1))); +} + +} // namespace core +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 418db9d88f..c9621109e9 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include @@ -344,7 +345,7 @@ struct index_params : cuvs::neighbors::index_params { * Disk-based ACE builds manage file-backed dataset state separately and ignore this flag. * * @code{.cpp} - * auto dataset = cuvs::neighbors::make_device_padded_dataset(res, host_matrix.view()); + * auto dataset = cuvs::core::make_device_padded_dataset(res, host_matrix.view()); * cagra::index_params index_params; * // Build graph only — caller attaches dataset later. * index_params.attach_dataset_on_build = false; @@ -482,7 +483,8 @@ static_assert(std::is_aggregate_v); */ template > + cuvs::core::dataset_like DatasetViewT = + cuvs::core::device_padded_dataset_view> struct CUVS_EXPORT index : cuvs::neighbors::index { using index_params_type = cagra::index_params; using search_params_type = cagra::search_params; @@ -588,7 +590,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { /** Construct a graph-only index with a zero-row dataset view placeholder. */ explicit index(raft::resources const& res, cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded) - requires(cuvs::neighbors::ann_dataset_view) + requires(cuvs::core::dataset_like) : cuvs::neighbors::index(), metric_(metric), graph_(raft::make_device_matrix(res, 0, 0)), @@ -606,7 +608,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { * that matrix must outlive the index): * @code{.cpp} * raft::device_matrix_view dataset = ...; - * auto view = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + * auto view = cuvs::core::make_device_padded_dataset_view(res, dataset); * auto graph = raft::make_device_matrix_view(...); * cuvs::neighbors::cagra::device_padded_index idx(res, metric, view, * raft::make_const_mdspan(graph)); @@ -617,7 +619,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { * **keep that object alive** (e.g. hold the `unique_ptr` in a variable or member) for as long as * the index uses the dataset; the index does not take ownership of the buffer. * @code{.cpp} - * auto padded_owner = cuvs::neighbors::make_device_padded_dataset(res, dataset_mdspan); + * auto padded_owner = cuvs::core::make_device_padded_dataset(res, dataset_mdspan); * auto view = padded_owner->as_dataset_view(); * cuvs::neighbors::cagra::device_padded_index idx(res, metric, view, * raft::make_const_mdspan(graph)); @@ -642,7 +644,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { "Dataset and knn_graph must have equal number of rows"); update_graph(res, knn_graph); - if constexpr (cuvs::neighbors::is_device_dataset_view_v) { + if constexpr (cuvs::core::is_device_dataset_view_v) { if (metric_ == cuvs::distance::DistanceType::CosineExpanded && dataset.n_rows() > 0) { compute_dataset_norms_(res); } @@ -653,7 +655,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { /* Construct an index with a new dataset type by moving the old index and passing in a new * dataset*/ - template + template index(raft::resources const& res, index&& other, DatasetViewT dataset) : metric_(other.metric_), graph_(std::move(other.graph_)), @@ -667,7 +669,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { dim_(other.dim_), graph_degree_(other.graph_degree_) { - if constexpr (is_device_dataset_view_v) { + if constexpr (cuvs::core::is_device_dataset_view_v) { if (metric() == cuvs::distance::DistanceType::CosineExpanded) { if (dataset_.n_rows() > 0) { compute_dataset_norms_(res); } } @@ -793,23 +795,23 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { // Re-open the file descriptor in read-only mode for subsequent operations dataset_fd_.emplace(std::move(fd)); - if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + if constexpr (cuvs::core::is_device_padded_dataset_view_v) { auto v = raft::make_device_matrix_view( static_cast(nullptr), int64_t{0}, dim_); dataset_ = DatasetViewT(v, dim_); - } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { + } else if constexpr (cuvs::core::is_device_standard_dataset_view_v) { auto v = raft::make_device_matrix_view( static_cast(nullptr), int64_t{0}, dim_); dataset_ = DatasetViewT(v); - } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v) { + } else if constexpr (cuvs::core::is_host_padded_dataset_view_v) { auto v = raft::make_host_matrix_view( static_cast(nullptr), int64_t{0}, dim_); dataset_ = DatasetViewT(v, dim_); - } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { + } else if constexpr (cuvs::core::is_host_standard_dataset_view_v) { auto v = raft::make_host_matrix_view( static_cast(nullptr), int64_t{0}, dim_); dataset_ = DatasetViewT(v); - } else if constexpr (cuvs::neighbors::is_empty_dataset_view_v) { + } else if constexpr (cuvs::core::is_empty_dataset_view_v) { dataset_ = DatasetViewT{dim_}; } else { RAFT_FAIL("update_dataset(fd): unsupported DatasetViewT for disk-backed dataset"); @@ -894,7 +896,7 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { } private: - template + template friend struct index; friend struct detail::fd_transfer; @@ -937,38 +939,38 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { /** CAGRA index with the usual padded device dataset view (graph build output type). */ template -using device_padded_index = index>; +using device_padded_index = index>; /** CAGRA index with a host-resident padded dataset view (returned by host build path). */ template -using host_padded_index = index>; +using host_padded_index = index>; /** CAGRA index with a device-resident standard (arbitrary stride) dataset view. */ template -using device_standard_index = - index>; +using device_standard_index = index>; /** CAGRA index with a host-resident standard dataset view. */ template -using host_standard_index = index>; +using host_standard_index = index>; /** CAGRA index with a device-resident VPQ dataset. */ template using device_pq_index = - index>; + index>; /** CAGRA index with a device-resident BBQ-quantized dataset. */ template -using device_bbq_index = index>; +using device_bbq_index = + index>; /** Index type returned by `cagra::build(res, params, dataset_view)`. */ template -using cagra_index_t = - std::conditional_t, - device_pq_index, - index, - uint32_t, - cuvs::neighbors::dataset_view_type_t>>; +using cagra_index_t = std::conditional_t< + cuvs::preprocessing::quantize::pq::is_device_vpq_f16_dataset_view_v, + device_pq_index, + index, + uint32_t, + cuvs::core::dataset_view_type_t>>; /** * @} @@ -1020,11 +1022,12 @@ using cagra_index_t = * @param[in] res raft resources * @param[in] params CAGRA index build parameters * @param[in] dataset device VPQ dataset view - * @return built `index>` + * @return built `index>` */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_vpq_dataset_view const& dataset) + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_pq_index; /** @@ -1036,7 +1039,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_padded_dataset_view const& dataset) + cuvs::core::device_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_padded_index; /** @@ -1048,7 +1051,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_standard_dataset_view const& dataset) + cuvs::core::device_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_standard_index; /** @@ -1060,7 +1063,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::host_padded_dataset_view const& dataset) + cuvs::core::host_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_padded_index; /** @@ -1072,7 +1075,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::host_standard_dataset_view const& dataset) + cuvs::core::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; /** @@ -1084,7 +1087,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_padded_dataset_view const& dataset) + cuvs::core::device_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_padded_index; /** @@ -1096,7 +1099,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_standard_dataset_view const& dataset) + cuvs::core::device_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_standard_index; /** @@ -1108,7 +1111,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::host_padded_dataset_view const& dataset) + cuvs::core::host_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_padded_index; /** @@ -1120,7 +1123,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::host_standard_dataset_view const& dataset) + cuvs::core::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; /** @@ -1132,7 +1135,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_padded_dataset_view const& dataset) + cuvs::core::device_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_padded_index; /** @@ -1144,7 +1147,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_standard_dataset_view const& dataset) + cuvs::core::device_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_standard_index; /** @@ -1156,7 +1159,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::host_padded_dataset_view const& dataset) + cuvs::core::host_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_padded_index; /** @@ -1168,7 +1171,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::host_standard_dataset_view const& dataset) + cuvs::core::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; /** @@ -1180,7 +1183,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_padded_dataset_view const& dataset) + cuvs::core::device_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_padded_index; /** @@ -1192,7 +1195,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_standard_dataset_view const& dataset) + cuvs::core::device_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_standard_index; /** @@ -1204,7 +1207,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::host_padded_dataset_view const& dataset) + cuvs::core::host_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_padded_index; /** @@ -1216,7 +1219,7 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::host_standard_dataset_view const& dataset) + cuvs::core::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; /** @@ -1236,30 +1239,34 @@ auto build(raft::resources const& res, * @param[in] dataset device BBQ dataset view [n_rows, dim] * @return built `device_bbq_index` */ -auto build(raft::resources const& res, - const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_bbq_dataset_view const& dataset) +auto build( + raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_bbq_index; /** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - * cuvs::neighbors::device_bbq_dataset_view const& dataset) */ -auto build(raft::resources const& res, - const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_bbq_dataset_view const& dataset) + * cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) */ +auto build( + raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_bbq_index; /** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - * cuvs::neighbors::device_bbq_dataset_view const& dataset) */ -auto build(raft::resources const& res, - const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_bbq_dataset_view const& dataset) + * cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) */ +auto build( + raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_bbq_index; /** @copydoc build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, - * cuvs::neighbors::device_bbq_dataset_view const& dataset) */ -auto build(raft::resources const& res, - const cuvs::neighbors::cagra::index_params& params, - cuvs::neighbors::device_bbq_dataset_view const& dataset) + * cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) */ +auto build( + raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_bbq_index; /** @@ -1292,7 +1299,7 @@ auto build(raft::resources const& res, * @code{.cpp} * using namespace cuvs::neighbors; * // Build `extended` = old || new on device, padded for CAGRA. - * auto extended = make_device_padded_dataset(res, concatenated_view); + * auto extended = cuvs::core::make_device_padded_dataset(res, concatenated_view); * auto extended_view = extended->as_dataset_view(); * * cagra::extend_params params; @@ -1308,28 +1315,28 @@ auto build(raft::resources const& res, */ void extend(raft::resources const& handle, const cagra::extend_params& params, - cuvs::neighbors::device_padded_dataset_view extended_dataset, + cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row, cuvs::neighbors::cagra::device_padded_index& idx); /** @brief Add new vectors to a CAGRA index. See the float overload for the full contract. */ void extend(raft::resources const& handle, const cagra::extend_params& params, - cuvs::neighbors::device_padded_dataset_view extended_dataset, + cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row, cuvs::neighbors::cagra::device_padded_index& idx); /** @brief Add new vectors to a CAGRA index. See the float overload for the full contract. */ void extend(raft::resources const& handle, const cagra::extend_params& params, - cuvs::neighbors::device_padded_dataset_view extended_dataset, + cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row, cuvs::neighbors::cagra::device_padded_index& idx); /** @brief Add new vectors to a CAGRA index. See the float overload for the full contract. */ void extend(raft::resources const& handle, const cagra::extend_params& params, - cuvs::neighbors::device_padded_dataset_view extended_dataset, + cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row, cuvs::neighbors::cagra::device_padded_index& idx); @@ -2226,7 +2233,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** * Write the index to an output stream @@ -2283,7 +2290,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** * Save the index to file. * @@ -2341,7 +2348,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** * Write the index to an output stream @@ -2398,7 +2405,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** * Save the index to file. @@ -2456,7 +2463,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** * Write the index to an output stream @@ -2513,7 +2520,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** * Save the index to file. @@ -2571,7 +2578,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** * Write the index to an output stream @@ -2628,7 +2635,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, const std::string& filename, @@ -2639,7 +2646,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_standard_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, std::ostream& os, @@ -2650,7 +2657,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_standard_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, const std::string& filename, @@ -2661,7 +2668,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_standard_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, std::ostream& os, @@ -2672,51 +2679,51 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_standard_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, const std::string& filename, const cuvs::neighbors::cagra::device_standard_index& index, bool include_dataset = true); -void deserialize(raft::resources const& handle, - const std::string& filename, - cuvs::neighbors::cagra::device_standard_index* index, - std::unique_ptr>* - out_dataset = nullptr); +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, std::ostream& os, const cuvs::neighbors::cagra::device_standard_index& index, bool include_dataset = true); -void deserialize(raft::resources const& handle, - std::istream& is, - cuvs::neighbors::cagra::device_standard_index* index, - std::unique_ptr>* - out_dataset = nullptr); +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, const std::string& filename, const cuvs::neighbors::cagra::device_standard_index& index, bool include_dataset = true); -void deserialize(raft::resources const& handle, - const std::string& filename, - cuvs::neighbors::cagra::device_standard_index* index, - std::unique_ptr>* - out_dataset = nullptr); +void deserialize( + raft::resources const& handle, + const std::string& filename, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* out_dataset = nullptr); void serialize(raft::resources const& handle, std::ostream& os, const cuvs::neighbors::cagra::device_standard_index& index, bool include_dataset = true); -void deserialize(raft::resources const& handle, - std::istream& is, - cuvs::neighbors::cagra::device_standard_index* index, - std::unique_ptr>* - out_dataset = nullptr); +void deserialize( + raft::resources const& handle, + std::istream& is, + cuvs::neighbors::cagra::device_standard_index* index, + std::unique_ptr>* out_dataset = nullptr); /* FP16-codebook device_pq_index graph-only overloads (CAGRA-Q). * @@ -2957,56 +2964,56 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** @copydoc deserialize */ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_standard_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** @copydoc deserialize */ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** @copydoc deserialize */ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_standard_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** @copydoc deserialize */ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** @copydoc deserialize */ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_standard_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** @copydoc deserialize */ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_padded_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** @copydoc deserialize */ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_standard_index* index, - std::unique_ptr>* out_dataset = nullptr); + std::unique_ptr>* out_dataset = nullptr); /** * Write the CAGRA built index as a base layer HNSW index to an output stream @@ -3513,7 +3520,7 @@ struct merge_params { * explicit FASTENER. * @return The merged physical CAGRA index. */ -template +template auto merge(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, std::vector*>& indices, @@ -3525,7 +3532,7 @@ auto merge(raft::resources const& res, /** @copydoc merge * @param[in] merge_params Parameters for the merge, including the algorithm selection. */ -template +template auto merge(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, std::vector*>& indices, @@ -3560,7 +3567,7 @@ auto merge(raft::resources const& res, */ auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - cuvs::neighbors::host_standard_dataset_view const& index_dataset) + cuvs::core::host_standard_dataset_view const& index_dataset) -> cuvs::neighbors::mg_index, float, uint32_t>; /// \ingroup mg_cpp_index_build @@ -3582,7 +3589,7 @@ auto build(const raft::resources& clique, */ auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - cuvs::neighbors::host_standard_dataset_view const& index_dataset) + cuvs::core::host_standard_dataset_view const& index_dataset) -> cuvs::neighbors::mg_index, half, uint32_t>; /// \ingroup mg_cpp_index_build @@ -3604,7 +3611,7 @@ auto build(const raft::resources& clique, */ auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - cuvs::neighbors::host_standard_dataset_view const& index_dataset) + cuvs::core::host_standard_dataset_view const& index_dataset) -> cuvs::neighbors::mg_index, int8_t, uint32_t>; /// \ingroup mg_cpp_index_build @@ -3626,27 +3633,27 @@ auto build(const raft::resources& clique, */ auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - cuvs::neighbors::host_standard_dataset_view const& index_dataset) + cuvs::core::host_standard_dataset_view const& index_dataset) -> cuvs::neighbors::mg_index, uint8_t, uint32_t>; auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - cuvs::neighbors::host_padded_dataset_view const& index_dataset) + cuvs::core::host_padded_dataset_view const& index_dataset) -> cuvs::neighbors::mg_index, float, uint32_t>; auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - cuvs::neighbors::host_padded_dataset_view const& index_dataset) + cuvs::core::host_padded_dataset_view const& index_dataset) -> cuvs::neighbors::mg_index, half, uint32_t>; auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - cuvs::neighbors::host_padded_dataset_view const& index_dataset) + cuvs::core::host_padded_dataset_view const& index_dataset) -> cuvs::neighbors::mg_index, int8_t, uint32_t>; auto build(const raft::resources& clique, const cuvs::neighbors::mg_index_params& index_params, - cuvs::neighbors::host_padded_dataset_view const& index_dataset) + cuvs::core::host_padded_dataset_view const& index_dataset) -> cuvs::neighbors::mg_index, uint8_t, uint32_t>; /** @@ -3657,26 +3664,26 @@ auto build(const raft::resources& clique, auto update_dataset( const raft::resources& clique, cuvs::neighbors::mg_index, float, uint32_t>&& idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + cuvs::core::device_padded_dataset_view const& padded_dataset) -> cuvs::neighbors::mg_index, float, uint32_t>; auto update_dataset( const raft::resources& clique, cuvs::neighbors::mg_index, half, uint32_t>&& idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + cuvs::core::device_padded_dataset_view const& padded_dataset) -> cuvs::neighbors::mg_index, half, uint32_t>; auto update_dataset( const raft::resources& clique, cuvs::neighbors::mg_index, int8_t, uint32_t>&& idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + cuvs::core::device_padded_dataset_view const& padded_dataset) -> cuvs::neighbors::mg_index, int8_t, uint32_t>; auto update_dataset( const raft::resources& clique, cuvs::neighbors::mg_index, uint8_t, uint32_t>&& idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + cuvs::core::device_padded_dataset_view const& padded_dataset) -> cuvs::neighbors::mg_index, uint8_t, uint32_t>; /** @@ -3685,22 +3692,22 @@ auto update_dataset( void update_dataset( const raft::resources& clique, cuvs::neighbors::mg_index, float, uint32_t>& idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset); + cuvs::core::device_padded_dataset_view const& padded_dataset); void update_dataset( const raft::resources& clique, cuvs::neighbors::mg_index, half, uint32_t>& idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset); + cuvs::core::device_padded_dataset_view const& padded_dataset); void update_dataset( const raft::resources& clique, cuvs::neighbors::mg_index, int8_t, uint32_t>& idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset); + cuvs::core::device_padded_dataset_view const& padded_dataset); void update_dataset( const raft::resources& clique, cuvs::neighbors::mg_index, uint8_t, uint32_t>& idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset); + cuvs::core::device_padded_dataset_view const& padded_dataset); /// \defgroup mg_cpp_index_extend ANN MG index extend @@ -3726,7 +3733,7 @@ void update_dataset( void extend( const raft::resources& clique, cuvs::neighbors::mg_index, float, uint32_t>& index, - cuvs::neighbors::host_padded_dataset_view new_vectors, + cuvs::core::host_padded_dataset_view new_vectors, std::optional> new_indices); /// \ingroup mg_cpp_index_extend @@ -3751,7 +3758,7 @@ void extend( void extend( const raft::resources& clique, cuvs::neighbors::mg_index, float, uint32_t>& index, - cuvs::neighbors::host_standard_dataset_view new_vectors, + cuvs::core::host_standard_dataset_view new_vectors, std::optional> new_indices); /// \ingroup mg_cpp_index_extend @@ -3776,14 +3783,14 @@ void extend( void extend( const raft::resources& clique, cuvs::neighbors::mg_index, half, uint32_t>& index, - cuvs::neighbors::host_padded_dataset_view new_vectors, + cuvs::core::host_padded_dataset_view new_vectors, std::optional> new_indices); /** @copydoc extend */ void extend( const raft::resources& clique, cuvs::neighbors::mg_index, half, uint32_t>& index, - cuvs::neighbors::host_standard_dataset_view new_vectors, + cuvs::core::host_standard_dataset_view new_vectors, std::optional> new_indices); /// \ingroup mg_cpp_index_extend @@ -3808,7 +3815,7 @@ void extend( void extend( const raft::resources& clique, cuvs::neighbors::mg_index, int8_t, uint32_t>& index, - cuvs::neighbors::host_padded_dataset_view new_vectors, + cuvs::core::host_padded_dataset_view new_vectors, std::optional> new_indices); /** @copydoc extend */ @@ -3816,7 +3823,7 @@ void extend( const raft::resources& clique, cuvs::neighbors::mg_index, int8_t, uint32_t>& index, - cuvs::neighbors::host_standard_dataset_view new_vectors, + cuvs::core::host_standard_dataset_view new_vectors, std::optional> new_indices); /// \ingroup mg_cpp_index_extend @@ -3842,7 +3849,7 @@ void extend( const raft::resources& clique, cuvs::neighbors::mg_index, uint8_t, uint32_t>& index, - cuvs::neighbors::host_padded_dataset_view new_vectors, + cuvs::core::host_padded_dataset_view new_vectors, std::optional> new_indices); /** @copydoc extend */ @@ -3850,7 +3857,7 @@ void extend( const raft::resources& clique, cuvs::neighbors::mg_index, uint8_t, uint32_t>& index, - cuvs::neighbors::host_standard_dataset_view new_vectors, + cuvs::core::host_standard_dataset_view new_vectors, std::optional> new_indices); /// \defgroup mg_cpp_index_search ANN MG index search @@ -4400,7 +4407,7 @@ void distribute(const raft::resources& clique, * auto optimized_graph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::helpers::optimize(res, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto dataset_view = make_host_standard_dataset_view(dataset); + * auto dataset_view = cuvs::core::make_host_standard_dataset_view(dataset); * auto index = cagra::host_standard_index( * res, metric, dataset_view, raft::make_const_mdspan(optimized_graph.view())); * @endcode @@ -4441,7 +4448,7 @@ void build_knn_graph(raft::resources const& res, * auto optimized_graph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::helpers::optimize(res, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto dataset_view = make_host_standard_dataset_view(dataset); + * auto dataset_view = cuvs::core::make_host_standard_dataset_view(dataset); * auto index = cagra::host_standard_index( * res, metric, dataset_view, raft::make_const_mdspan(optimized_graph.view())); * @endcode @@ -4482,7 +4489,7 @@ void build_knn_graph(raft::resources const& res, * auto optimized_graph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::helpers::optimize(res, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto dataset_view = make_host_standard_dataset_view(dataset); + * auto dataset_view = cuvs::core::make_host_standard_dataset_view(dataset); * auto index = cagra::host_standard_index( * res, metric, dataset_view, raft::make_const_mdspan(optimized_graph.view())); * @endcode @@ -4523,7 +4530,7 @@ void build_knn_graph(raft::resources const& res, * auto optimized_graph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::helpers::optimize(res, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto dataset_view = make_host_standard_dataset_view(dataset); + * auto dataset_view = cuvs::core::make_host_standard_dataset_view(dataset); * auto index = cagra::host_standard_index( * res, metric, dataset_view, raft::make_const_mdspan(optimized_graph.view())); * @endcode @@ -4548,8 +4555,8 @@ namespace detail { struct fd_transfer { template + cuvs::core::dataset_like SrcDatasetViewT, + cuvs::core::dataset_like DstDatasetViewT> static inline void steal_disk_fds_to(raft::resources const& res, index& src, index& dst) @@ -4580,16 +4587,15 @@ struct fd_transfer { template auto convert_standard_to_padded_index( raft::resources const& res, - index> const& standard_idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset) + index> const& standard_idx, + cuvs::core::device_padded_dataset_view const& padded_dataset) -> device_padded_index { RAFT_EXPECTS(padded_dataset.n_rows() == standard_idx.size(), "Padded dataset row count must match the index size"); using GraphIndexType = - typename index>:: - graph_index_type; + typename index>::graph_index_type; auto graph_host = raft::make_host_matrix(standard_idx.graph().extent(0), standard_idx.graph().extent(1)); if (standard_idx.graph().size() > 0) { @@ -4609,277 +4615,383 @@ auto convert_standard_to_padded_index( auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_standard_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_standard_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_standard_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_standard_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& + cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; + +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_standard_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_standard_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_standard_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_standard_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_standard_dataset_view dataset) - -> index>; + index>&& + cagra_index, + cuvs::core::device_standard_dataset_view dataset) + -> index>; + auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_standard_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_standard_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; + auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_standard_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_standard_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_standard_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_standard_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& + cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; + auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; - -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& + cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& + cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; + auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; + auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; - -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; - -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; - -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; - -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; - -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_padded_dataset_view dataset) - -> index>; + index>&& + cagra_index, + cuvs::core::device_padded_dataset_view dataset) + -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; -auto update_dataset(raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; +auto update_dataset( + raft::resources const& res, + index>&& cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; auto update_dataset( raft::resources const& res, - index>&& cagra_index, - device_vpq_dataset_view dataset) - -> index>; + index>&& + cagra_index, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view dataset) + -> index>; } // namespace cagra } // namespace neighbors diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index baa5b44da2..e17d67a7c6 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -24,6 +24,7 @@ #include #include +#include #include #include @@ -37,6 +38,7 @@ #include #include #include +#include #ifdef __cpp_lib_bitops #include #endif @@ -56,64 +58,6 @@ namespace neighbors { /* Graph build algo used in cagra and all_neighbors */ enum GRAPH_BUILD_ALGO { BRUTE_FORCE = 0, IVF_PQ = 1, NN_DESCENT = 2, ACE = 3 }; -/** Parameters for VPQ compression. */ -struct vpq_params { - /** - * The bit length of the vector element after compression by PQ. - * - * Possible values: [4, 5, 6, 7, 8]. - * - * Hint: the smaller the 'pq_bits', the smaller the index size and the better the search - * performance, but the lower the recall. - */ - uint32_t pq_bits = 8; - /** - * The dimensionality of the vector after compression by PQ. - * When zero, an optimal value is selected using a heuristic. - * - * TODO: at the moment `dim` must be a multiple `pq_dim`. - */ - uint32_t pq_dim = 0; - /** - * Vector Quantization (VQ) codebook size - number of "coarse cluster centers". - * When zero, an optimal value is selected using a heuristic. - */ - uint32_t vq_n_centers = 0; - /** The number of iterations searching for kmeans centers (both VQ & PQ phases). */ - uint32_t kmeans_n_iters = 25; - /** - * The fraction of data to use during iterative kmeans building (VQ phase). - * When zero, an optimal value is selected using a heuristic. - * @deprecated Prefer using `max_train_points_per_vq_cluster` instead. - */ - double vq_kmeans_trainset_fraction = 0; - /** - * The fraction of data to use during iterative kmeans building (PQ phase). - * When zero, an optimal value is selected using a heuristic. - * @deprecated Prefer using `max_train_points_per_pq_code` instead. - */ - double pq_kmeans_trainset_fraction = 0; - /** - * Type of k-means algorithm for PQ training. - * Balanced k-means tends to be faster than regular k-means for PQ training, for - * problem sets where the number of points per cluster are approximately equal. - * Regular k-means may be better for skewed cluster distributions. - */ - cuvs::cluster::kmeans::kmeans_type pq_kmeans_type = - cuvs::cluster::kmeans::kmeans_type::KMeansBalanced; - /** - * The max number of data points to use per PQ code during PQ codebook training. Using more data - * points per PQ code may increase the quality of PQ codebook but may also increase the build - * time. We will use `pq_n_centers * max_train_points_per_pq_code` training - * points to train each PQ codebook. - */ - uint32_t max_train_points_per_pq_code = 256; - /** - * The max number of data points to use per VQ cluster during training. - */ - uint32_t max_train_points_per_vq_cluster = 1024; -}; - /** @} */ // end group cagra_cpp_index_params /** @@ -149,1200 +93,6 @@ enum class MergeStrategy { /** @} */ // end group neighbors_index -/** - * @brief Tags selecting dataset representation for `dataset` / `dataset_view`. - * - * Each container defines nested `owning_storage` then `view_storage` (aliases into `detail::*` - * storage types shared by device/host). Accessibility (device vs host) is selected by the - * `Accessor` template parameter on `dataset` / `dataset_view`, not by duplicating containers. - * Layout kinds: empty, padded, standard, VPQ. `dataset` / `dataset_view` only express ownership - * vs view. - */ - -template -struct dataset; - -template -struct dataset_view; - -namespace detail { - -// Default owning/view accessors for public dataset aliases. -template -using device_owning_accessor = raft::device_accessor>; - -template -using host_owning_accessor = raft::host_accessor>; - -template -using device_view_accessor = raft::device_accessor>; - -template -using host_view_accessor = raft::host_accessor>; - -/** View accessor paired with an owning dataset accessor (same residency). */ -template -using dataset_view_accessor_for_owning = std::conditional_t, - host_view_accessor>; - -/** Owning accessor paired with a view accessor (same residency). */ -template -using dataset_owning_accessor_for_view = std::conditional_t, - host_owning_accessor>; - -template -using dense_owning_matrix = std::conditional_t, - raft::host_matrix>; - -template -using dense_view_matrix = - std::conditional_t, - raft::host_matrix_view>; - -template -using vpq_vq_book_matrix = std::conditional_t, - raft::host_matrix>; - -template -using vpq_data_matrix = std::conditional_t, - raft::host_matrix>; - -// ----------------------------------------------------------------------------- -// empty -// ----------------------------------------------------------------------------- - -template -struct empty_dataset_storage { - uint32_t suggested_dim{}; - empty_dataset_storage() noexcept = default; - explicit empty_dataset_storage(uint32_t dim) noexcept : suggested_dim(dim) {} - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return 0; } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return suggested_dim; } -}; - -template -using empty_dataset_owning_storage = empty_dataset_storage; - -template -using empty_dataset_view_storage = empty_dataset_storage; - -// ----------------------------------------------------------------------------- -// dense row-major (logical dim may differ from row pitch; shared by padded & standard) -// ----------------------------------------------------------------------------- - -/** - * Dense row-major owning storage shared by padded and standard dataset containers. - * - * Template parameters: - * - MatrixT: owning matrix type that stores the payload (host/device matrix). - * - ViewT: non-owning row-major view type returned by `view()`. - * - DataT: scalar element type of the dataset payload. - * - IdxT: index type used for row counts (`n_rows()` return type). - */ -template -struct dense_row_major_dataset_owning_storage { - MatrixT data_; - uint32_t logical_dim_; - - dense_row_major_dataset_owning_storage(MatrixT&& data, uint32_t logical_dim) noexcept - : data_{std::move(data)}, logical_dim_{logical_dim} - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } - [[nodiscard]] auto stride() const noexcept -> uint32_t - { - return static_cast(data_.extent(1)); - } - [[nodiscard]] auto view() const noexcept -> ViewT { return data_.view(); } - [[nodiscard]] auto data_handle() noexcept -> DataT* { return data_.data_handle(); } - [[nodiscard]] auto data_handle() const noexcept -> const DataT* { return data_.data_handle(); } -}; - -template -struct dense_row_major_dataset_view_storage { - ViewT data_; - uint32_t logical_dim_; - - dense_row_major_dataset_view_storage() noexcept = default; - - explicit dense_row_major_dataset_view_storage(ViewT v) noexcept - : data_(v), logical_dim_(static_cast(v.extent(1))) - { - } - - dense_row_major_dataset_view_storage(ViewT v, uint32_t logical_dim) noexcept - : data_(v), logical_dim_(logical_dim) - { - } - - dense_row_major_dataset_view_storage(dense_row_major_dataset_view_storage const& other) noexcept - : data_(other.data_), logical_dim_(other.logical_dim_) - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data_.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return logical_dim_; } - [[nodiscard]] auto stride() const noexcept -> uint32_t - { - return static_cast(data_.stride(0) > 0 ? data_.stride(0) : data_.extent(1)); - } - [[nodiscard]] auto view() const noexcept -> ViewT { return data_; } -}; - -template -using padded_dataset_owning_storage = - dense_row_major_dataset_owning_storage; - -template -using padded_dataset_view_storage = dense_row_major_dataset_view_storage; - -template -using standard_dataset_owning_storage = - dense_row_major_dataset_owning_storage; - -template -using standard_dataset_view_storage = dense_row_major_dataset_view_storage; - -// ----------------------------------------------------------------------------- -// VPQ compressed -// ----------------------------------------------------------------------------- - -/** - * Owning storage for VPQ-compressed datasets. - * - * Template parameters: - * - VqBookMatrixT: owning matrix type for the VQ codebook. - * - PqBookMatrixT: owning matrix type for the PQ codebook. - * - DataMatrixT: owning matrix type for encoded row data (uint8 codes). - * - MathT: floating-point type used by VQ/PQ codebooks. - * - IdxT: index type used for row counts (`n_rows()` return type). - */ -template -struct vpq_dataset_owning_storage { - /** Floating-point type used for VQ/PQ codebooks (rows are still uint8 codes). */ - using math_type = MathT; - - VqBookMatrixT vq_code_book; - PqBookMatrixT pq_code_book; - DataMatrixT data; - - vpq_dataset_owning_storage(VqBookMatrixT&& vq_code_book, - PqBookMatrixT&& pq_code_book, - DataMatrixT&& data) noexcept - : vq_code_book{std::move(vq_code_book)}, - pq_code_book{std::move(pq_code_book)}, - data{std::move(data)} - { - } - - [[nodiscard]] auto n_rows() const noexcept -> IdxT { return data.extent(0); } - [[nodiscard]] auto dim() const noexcept -> uint32_t { return vq_code_book.extent(1); } - - [[nodiscard]] constexpr inline auto encoded_row_length() const noexcept -> uint32_t - { - return data.extent(1); - } - [[nodiscard]] constexpr inline auto vq_n_centers() const noexcept -> uint32_t - { - return vq_code_book.extent(0); - } - [[nodiscard]] constexpr inline auto pq_bits() const noexcept -> uint32_t - { - auto pq_width = pq_n_centers(); -#ifdef __cpp_lib_bitops - return std::countr_zero(pq_width); -#else - uint32_t pq_bits = 0; - while (pq_width > 1) { - pq_bits++; - pq_width >>= 1; - } - return pq_bits; -#endif - } - [[nodiscard]] constexpr inline auto pq_dim() const noexcept -> uint32_t - { - return raft::div_rounding_up_unsafe(dim(), pq_len()); - } - [[nodiscard]] constexpr inline auto pq_len() const noexcept -> uint32_t - { - return pq_code_book.extent(1); - } - [[nodiscard]] constexpr inline auto pq_n_centers() const noexcept -> uint32_t - { - return pq_code_book.extent(0); - } -}; - -template -struct vpq_dataset_view_storage { - using owning_dataset_type = - dataset>; - - owning_dataset_type const* dataset_{nullptr}; - - vpq_dataset_view_storage() = default; - - explicit vpq_dataset_view_storage(owning_dataset_type const* ptr) : dataset_(ptr) - { - RAFT_EXPECTS(ptr != nullptr, "vpq_dataset_view: null dataset pointer"); - } - - [[nodiscard]] auto n_rows() const noexcept - { - using idx_type = decltype(std::declval().n_rows()); - return dataset_ != nullptr ? dataset_->n_rows() : idx_type{0}; - } - [[nodiscard]] auto dim() const noexcept -> uint32_t - { - return dataset_ != nullptr ? dataset_->dim() : uint32_t{0}; - } - [[nodiscard]] owning_dataset_type const& dset() const noexcept { return *dataset_; } -}; - -} // namespace detail - -// ----------------------------------------------------------------------------- -// empty -// ----------------------------------------------------------------------------- - -struct empty_dataset_container { - template - using owning_storage = detail::empty_dataset_owning_storage; - template - using view_storage = detail::empty_dataset_view_storage; -}; - -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- - -struct padded_dataset_container { - template - using owning_storage = - detail::padded_dataset_owning_storage, - detail::dense_view_matrix, - DataT, - IdxT>; - template - using view_storage = detail:: - padded_dataset_view_storage, DataT, IdxT>; -}; - -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride; no CAGRA alignment requirement) -// ----------------------------------------------------------------------------- - -struct standard_dataset_container { - template - using owning_storage = - detail::standard_dataset_owning_storage, - detail::dense_view_matrix, - DataT, - IdxT>; - template - using view_storage = detail:: - standard_dataset_view_storage, DataT, IdxT>; -}; - -// ----------------------------------------------------------------------------- -// VPQ compressed -// ----------------------------------------------------------------------------- - -struct vpq_dataset_container { - template - using owning_storage = - detail::vpq_dataset_owning_storage, - detail::vpq_vq_book_matrix, - detail::vpq_data_matrix, - MathT, - IdxT>; - template - using view_storage = - detail::vpq_dataset_view_storage; -}; - -template -struct dataset { - static_assert(!std::is_same_v, - "dataset: unsupported ContainerType / type-parameter combination"); -}; - -template -struct dataset_view { - static_assert(!std::is_same_v, - "dataset_view: unsupported ContainerType / type-parameter combination"); -}; - -// ----------------------------------------------------------------------------- -// empty -// ----------------------------------------------------------------------------- - -template -struct dataset - : empty_dataset_container::template owning_storage { - using container_type = empty_dataset_container; - using owning_storage_type = typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> - { - return dataset_view>{this->dim()}; - } -}; - -template -struct dataset_view - : empty_dataset_container::template view_storage { - using container_type = empty_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -// ----------------------------------------------------------------------------- -// standard (row-major with arbitrary stride) -// ----------------------------------------------------------------------------- - -template -struct dataset - : standard_dataset_container::template owning_storage { - using container_type = standard_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> - { - return dataset_view>(this->view(), - this->dim()); - } -}; - -template -struct dataset_view - : standard_dataset_container::template view_storage { - using container_type = standard_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -// ----------------------------------------------------------------------------- -// padded (row-major with logical dim vs stride) -// ----------------------------------------------------------------------------- - -template -struct dataset - : padded_dataset_container::template owning_storage { - using container_type = padded_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> - { - return dataset_view>(this->view(), - this->dim()); - } -}; - -template -struct dataset_view - : padded_dataset_container::template view_storage { - using container_type = padded_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -// ----------------------------------------------------------------------------- -// VPQ compressed (view holds non-owning pointer to owning dataset) -// ----------------------------------------------------------------------------- - -template -struct dataset - : vpq_dataset_container::template owning_storage { - using container_type = vpq_dataset_container; - using owning_storage_type = - typename container_type::template owning_storage; - using owning_storage_type::owning_storage_type; - - [[nodiscard]] auto as_dataset_view() const - -> dataset_view> - { - return dataset_view>{this}; - } -}; - -template -struct dataset_view - : vpq_dataset_container::template view_storage { - using container_type = vpq_dataset_container; - using view_storage_type = typename container_type::template view_storage; - using view_storage_type::view_storage_type; -}; - -/** - * @brief Aliases for concrete `dataset` / `dataset_view` layouts. - */ -template -using device_empty_dataset = - dataset>; - -template -using device_empty_dataset_view = - dataset_view>; - -template -using host_empty_dataset = - dataset>; - -template -using host_empty_dataset_view = - dataset_view>; - -template -using device_padded_dataset = - dataset>; - -template -using device_padded_dataset_view = - dataset_view>; - -template -using host_padded_dataset = - dataset>; - -template -using host_padded_dataset_view = - dataset_view>; - -template -using device_standard_dataset = - dataset>; - -template -using device_standard_dataset_view = - dataset_view>; - -template -using host_standard_dataset = - dataset>; - -template -using host_standard_dataset_view = - dataset_view>; - -template -using device_vpq_dataset = - dataset>; - -template -using device_vpq_dataset_view = - dataset_view>; - -template -using host_vpq_dataset = - dataset>; - -template -using host_vpq_dataset_view = - dataset_view>; - -// Maps a dataset view type to its owning (allocating) dataset counterpart. -// Used by serialize/deserialize to type the out_dataset output parameter; -// adding a new dataset type only requires adding a new specialization here. -template -struct owning_dataset_for_view; - -template -struct owning_dataset_for_view> { - using type = device_padded_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = device_standard_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = host_padded_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = host_standard_dataset; -}; - -template -struct owning_dataset_for_view> { - using type = device_vpq_dataset; -}; - -template -using owning_dataset_for_view_t = typename owning_dataset_for_view::type; - -template -struct is_padded_dataset : std::false_type {}; - -template -struct is_padded_dataset> - : std::true_type {}; - -template -struct is_padded_dataset> - : std::true_type {}; - -template -inline constexpr bool is_padded_dataset_v = is_padded_dataset::value; - -template -struct is_standard_dataset : std::false_type {}; - -template -struct is_standard_dataset> - : std::true_type {}; - -template -struct is_standard_dataset> - : std::true_type {}; - -template -inline constexpr bool is_standard_dataset_v = is_standard_dataset::value; - -template -struct is_vpq_dataset : std::false_type {}; - -template -struct is_vpq_dataset> : std::true_type {}; - -template -inline constexpr bool is_vpq_dataset_v = is_vpq_dataset::value; - -// ----------------------------------------------------------------------------- -// Dataset view compile-time classification (replaces runtime std::variant dispatch). -// ----------------------------------------------------------------------------- - -/** Any non-owning dataset view exposing row count and logical dimension. */ -template -concept ann_dataset_view = requires(V const& v) { - { v.n_rows() } -> std::convertible_to; - { v.dim() } -> std::convertible_to; -}; - -enum class dataset_view_kind { - // TODO(removal): Remove `unknown` once all deprecated host_matrix_view / device_matrix_view / - // mdspan overloads are deleted. It exists solely so that overload resolution on the deprecated - // build(host_matrix_view) / build(device_matrix_view) shims does not cause a hard error when - // the compiler evaluates is_host/device_dataset_view_v for a plain mdspan type. - unknown, - empty, - padded, - standard, - vpq_f16, - vpq_f32, - bbq, -}; - -/** Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. - */ -template -struct dataset_view_kind_of { - static constexpr dataset_view_kind value = dataset_view_kind::unknown; -}; - -template -struct dataset_view_kind_of> { - static constexpr dataset_view_kind value = dataset_view_kind::empty; -}; - -template -struct dataset_view_kind_of> { - static constexpr dataset_view_kind value = dataset_view_kind::padded; -}; - -template -struct dataset_view_kind_of> { - static constexpr dataset_view_kind value = dataset_view_kind::standard; -}; - -template -struct dataset_view_kind_of> { - static_assert(std::is_same_v || std::is_same_v, - "VPQ dataset_view_kind_of expects MathT to be half or float"); - static constexpr dataset_view_kind value = - std::is_same_v ? dataset_view_kind::vpq_f16 : dataset_view_kind::vpq_f32; -}; - -template -using dataset_view_type_t = std::remove_cvref_t; - -/** True when the dataset view accessor is device-accessible. */ -template -struct dataset_view_is_device_accessible : std::false_type {}; - -template -struct dataset_view_is_device_accessible> - : std::bool_constant {}; - -template -inline constexpr bool dataset_view_is_device_accessible_v = - dataset_view_is_device_accessible>::value; - -template -inline constexpr dataset_view_kind dataset_view_kind_v = - dataset_view_kind_of>::value; - -template -inline constexpr bool is_device_empty_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::empty && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_empty_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::empty && !dataset_view_is_device_accessible_v; - -/** True for any empty dataset view (device or host). */ -template -inline constexpr bool is_empty_dataset_view_v = - is_device_empty_dataset_view_v || is_host_empty_dataset_view_v; - -template -inline constexpr bool is_device_padded_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::padded && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_padded_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::padded && !dataset_view_is_device_accessible_v; - -/** True for either `device_padded_dataset_view` or `host_padded_dataset_view`. */ -template -inline constexpr bool is_padded_dataset_view_v = - is_device_padded_dataset_view_v || is_host_padded_dataset_view_v; - -template -inline constexpr bool is_device_standard_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::standard && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_standard_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::standard && !dataset_view_is_device_accessible_v; - -/** True for either `device_standard_dataset_view` or `host_standard_dataset_view`. */ -template -inline constexpr bool is_standard_dataset_view_v = - is_device_standard_dataset_view_v || is_host_standard_dataset_view_v; - -template -inline constexpr bool is_device_vpq_f16_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::vpq_f16 && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_vpq_f16_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::vpq_f16 && !dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_vpq_f16_dataset_view_v = - is_device_vpq_f16_dataset_view_v || is_host_vpq_f16_dataset_view_v; - -template -inline constexpr bool is_device_vpq_f32_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::vpq_f32 && dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_host_vpq_f32_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::vpq_f32 && !dataset_view_is_device_accessible_v; - -template -inline constexpr bool is_vpq_f32_dataset_view_v = - is_device_vpq_f32_dataset_view_v || is_host_vpq_f32_dataset_view_v; - -template -inline constexpr bool is_device_vpq_dataset_view_v = - is_device_vpq_f16_dataset_view_v || is_device_vpq_f32_dataset_view_v; - -template -inline constexpr bool is_host_vpq_dataset_view_v = - is_host_vpq_f16_dataset_view_v || is_host_vpq_f32_dataset_view_v; - -template -inline constexpr bool is_vpq_dataset_view_v = - is_device_vpq_dataset_view_v || is_host_vpq_dataset_view_v; - -/** True for any device-resident dataset view. */ -template -inline constexpr bool is_device_dataset_view_v = - dataset_view_kind_v != dataset_view_kind::unknown && dataset_view_is_device_accessible_v; - -/** True for any host-resident dataset view. */ -template -inline constexpr bool is_host_dataset_view_v = - dataset_view_kind_v != dataset_view_kind::unknown && !dataset_view_is_device_accessible_v; - -/** - * True when a host view `H` and device view `D` represent the same storage kind and differ - * only in residency (host vs. device). Used by host/device conversion helpers. - */ -template -inline constexpr bool compatible_host_device_dataset_views_v = - is_host_dataset_view_v && is_device_dataset_view_v && - (dataset_view_kind_v == dataset_view_kind_v); - -/** - * Generic accessor retargeting while preserving the dataset tag/layout and value/index types: - * `dataset -> dataset` - * `dataset_view -> dataset_view` - */ -template -struct with_accessor; - -template -struct with_accessor, NewAccessor> { - using type = dataset; -}; - -template -struct with_accessor, NewAccessor> { - using type = dataset_view; -}; - -template -using with_accessor_t = - typename with_accessor, NewAccessor>::type; - -/** Map any host accessor to its device counterpart (same payload policy). */ -template -struct to_device_accessor { - using type = Accessor; -}; - -template -struct to_device_accessor> { - using type = detail::device_view_accessor; -}; - -template -struct to_device_accessor> { - using type = detail::device_owning_accessor; -}; - -template -using to_device_accessor_t = typename to_device_accessor::type; - -/** Maps a host dataset view type to its device-resident counterpart. */ -template -struct device_counterpart; - -template -struct device_counterpart> { - using type = with_accessor_t, - to_device_accessor_t>; -}; - -template -using device_counterpart_t = typename device_counterpart>::type; - -/** True for device padded or standard views accepted by dense graph build (VPQ excluded). */ -template -inline constexpr bool is_dense_row_major_device_dataset_view_v = - is_device_padded_dataset_view_v || is_device_standard_dataset_view_v; - -/** True for host or device padded/standard views (dense graph build; VPQ excluded). */ -template -inline constexpr bool is_dense_row_major_dataset_view_v = - is_padded_dataset_view_v || is_standard_dataset_view_v; - -/** Element type `T` for `cagra::build(res, params, dataset_view)` (deduced, not a template arg). */ -template -struct cagra_view_element_type; - -template -struct cagra_view_element_type> { - using type = DataT; -}; - -template -struct cagra_view_element_type> { - using type = DataT; -}; - -template -struct cagra_view_element_type> { - using type = DataT; -}; - -template -struct cagra_view_element_type> { - using type = DataT; -}; - -template -struct cagra_view_element_type> { - using type = MathT; -}; - -template -using cagra_view_element_type_t = typename cagra_view_element_type>::type; - -// ----------------------------------------------------------------------------- -// CAGRA row width in elements (same for make_device_padded_dataset* and index layout checks). -// ----------------------------------------------------------------------------- - -/** - * @brief Required row width in elements for CAGRA: minimum leading dimension (LDA) per row for the - * default per-row byte alignment (16 bytes, combined with `sizeof` element type), given - * `logical_columns` feature columns. - */ -[[nodiscard]] inline uint32_t cagra_required_row_width(uint32_t logical_columns, - std::size_t sizeof_value, - uint32_t align_bytes = 16) -{ - return static_cast( - raft::round_up_safe(static_cast(logical_columns) * sizeof_value, - std::lcm(align_bytes, static_cast(sizeof_value))) / - sizeof_value); -} - -template -[[nodiscard]] inline uint32_t cagra_required_row_width(uint32_t logical_columns, - uint32_t align_bytes = 16) -{ - return cagra_required_row_width(logical_columns, sizeof(ValueT), align_bytes); -} - -/** Actual row width in elements (leading dimension) of a 2D row-major matrix view. */ -template -[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::device_matrix_view m) -{ - return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); -} - -template -[[nodiscard]] inline uint32_t matrix_actual_row_width(raft::host_matrix_view m) -{ - return m.stride(0) > 0 ? static_cast(m.stride(0)) : static_cast(m.extent(1)); -} - -/** - * @brief True if the matrix's row width in elements matches `cagra_required_row_width` for - * `m.extent(1)` and element type `T` (CAGRA row layout is satisfied for this view). - */ -template -[[nodiscard]] inline bool matrix_row_width_matches_cagra_required( - raft::device_matrix_view m, uint32_t align_bytes = 16) -{ - using value_type = std::remove_const_t; - const uint32_t need = - cagra_required_row_width(static_cast(m.extent(1)), align_bytes); - return matrix_actual_row_width(m) == need; -} - -template -[[nodiscard]] inline bool matrix_row_width_matches_cagra_required(raft::host_matrix_view m, - uint32_t align_bytes = 16) -{ - using value_type = std::remove_const_t; - const uint32_t need = - cagra_required_row_width(static_cast(m.extent(1)), align_bytes); - return matrix_actual_row_width(m) == need; -} - -namespace detail { - -template -[[nodiscard]] inline uint32_t mdspan_row_stride_elements(SrcT const& src) -{ - return src.stride(0) > 0 ? static_cast(src.stride(0)) - : static_cast(src.extent(1)); -} - -template -[[nodiscard]] inline ValueT* expect_device_accessible_data_handle(SrcT const& src, - char const* error_msg) -{ - cudaPointerAttributes ptr_attrs; - RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src.data_handle())); - // `devicePointer` is relative to the *current* device: it is null for an allocation owned by - // another device without peer access, even though that allocation is perfectly usable once the - // caller switches to the owning device (as the multi-GPU paths do). Accept device and managed - // allocations on their own merit and only consult `devicePointer` for host memory, which needs a - // mapping to be reachable at all. - if (ptr_attrs.type == cudaMemoryTypeDevice || ptr_attrs.type == cudaMemoryTypeManaged) { - return const_cast(src.data_handle()); - } - auto* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); - RAFT_EXPECTS(device_ptr != nullptr, "%s", error_msg); - return device_ptr; -} - -template -[[nodiscard]] inline ViewT make_device_dense_row_major_view_from_src(SrcT const& src, - uint32_t logical_dim) -{ - auto* device_ptr = expect_device_accessible_data_handle( - src, "make_device_*_dataset_view: source must be device-accessible."); - auto v = raft::make_device_matrix_view( - device_ptr, src.extent(0), static_cast(mdspan_row_stride_elements(src))); - return ViewT(v, logical_dim); -} - -template -[[nodiscard]] inline ViewT make_host_dense_row_major_view_from_src(SrcT const& src, - uint32_t logical_dim) -{ - RAFT_EXPECTS(raft::get_device_for_address(src.data_handle()) == -1, - "make_host_*_dataset_view: source must be host-accessible."); - auto v = raft::make_host_matrix_view(const_cast(src.data_handle()), - src.extent(0), - static_cast(mdspan_row_stride_elements(src))); - return ViewT(v, logical_dim); -} - -template -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) - -> std::unique_ptr -{ - uint32_t const src_stride = mdspan_row_stride_elements(src); - RAFT_EXPECTS(logical_dim <= target_stride, - "logical dim (%u) must not exceed row stride (%u).", - static_cast(logical_dim), - static_cast(target_stride)); - RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, - "Source row length must not exceed required stride."); - cudaPointerAttributes ptr_attrs; - 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) { - RAFT_EXPECTS(false, - "source is device and stride is already correct. " - "Use %s() to get a view instead.", - view_factory_name); - } - auto out_array = raft::make_device_matrix(res, src.extent(0), target_stride); - RAFT_CUDA_TRY(cudaMemsetAsync(out_array.data_handle(), - 0, - out_array.size() * sizeof(ValueT), - raft::resource::get_cuda_stream(res).get())); - raft::copy_matrix(out_array.data_handle(), - target_stride, - src.data_handle(), - src_stride, - logical_dim, - src.extent(0), - raft::resource::get_cuda_stream(res)); - return std::make_unique(std::move(out_array), logical_dim); -} - -template -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) - -> std::unique_ptr -{ - uint32_t const src_stride = mdspan_row_stride_elements(src); - constexpr bool device_src = SrcT::accessor_type::is_device_accessible; - RAFT_EXPECTS(logical_dim <= target_stride, - "logical dim (%u) must not exceed row stride (%u).", - static_cast(logical_dim), - static_cast(target_stride)); - if (!device_src && src_stride == target_stride) { - RAFT_EXPECTS(false, - "source stride is already correct. Use %s() to get a view instead.", - view_factory_name); - } - RAFT_EXPECTS(static_cast(src.extent(1)) <= target_stride, - "Source row length must not exceed required stride."); - auto out_array = raft::make_host_matrix(src.extent(0), target_stride); - std::memset(out_array.data_handle(), 0, out_array.size() * sizeof(ValueT)); - raft::copy_matrix(out_array.data_handle(), - target_stride, - src.data_handle(), - src_stride, - logical_dim, - src.extent(0), - raft::resource::get_cuda_stream(res)); - if (device_src) { raft::resource::sync_stream(res); } - return std::make_unique(std::move(out_array), logical_dim); -} - -} // namespace detail - -template -auto make_device_padded_dataset_view(const raft::resources& res, - SrcT const& src, - uint32_t align_bytes = 16) - -> device_padded_dataset_view -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - uint32_t required_stride = - cagra_required_row_width(static_cast(src.extent(1)), align_bytes); - RAFT_EXPECTS( - detail::mdspan_row_stride_elements(src) == required_stride, - "make_device_padded_dataset_view: stride is incorrect (required stride for alignment). " - "Use make_device_padded_dataset() to get an owning padded copy."); - return detail::make_device_dense_row_major_view_from_src< - value_type, - index_type, - device_padded_dataset_view>(src, static_cast(src.extent(1))); -} - -template -auto make_device_padded_dataset(const raft::resources& res, - SrcT const& src, - uint32_t align_bytes = 16) - -> std::unique_ptr> -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - uint32_t const logical_dim = static_cast(src.extent(1)); - uint32_t const required_stride = cagra_required_row_width(logical_dim, align_bytes); - 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"); -} - -template -auto make_host_padded_dataset_view(SrcT const& src, uint32_t align_bytes = 16) - -> host_padded_dataset_view -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - uint32_t required_stride = - cagra_required_row_width(static_cast(src.extent(1)), align_bytes); - RAFT_EXPECTS( - detail::mdspan_row_stride_elements(src) == required_stride, - "make_host_padded_dataset_view: stride is incorrect (required stride for alignment). " - "Use make_host_padded_dataset() to get an owning padded copy."); - return detail::make_host_dense_row_major_view_from_src< - value_type, - index_type, - host_padded_dataset_view>(src, static_cast(src.extent(1))); -} - -template -auto make_host_padded_dataset(const raft::resources& res, - SrcT const& src, - uint32_t align_bytes = 16) - -> std::unique_ptr> -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - uint32_t const logical_dim = static_cast(src.extent(1)); - uint32_t const required_stride = cagra_required_row_width(logical_dim, align_bytes); - 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"); -} - -template -auto make_device_standard_dataset_view(SrcT const& src) - -> device_standard_dataset_view -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - return detail::make_device_dense_row_major_view_from_src< - value_type, - index_type, - device_standard_dataset_view>(src, - static_cast(src.extent(1))); -} - -/** - * @brief Create an owning device standard dataset with explicit row layout. - * - * Internal use only: the sole call site today is - * `cuvs::neighbors::detail::deserialize_standard()` in `dataset_serialize.hpp`, which must pass - * wire-format `(logical_dim, stride)` because the deserialized host buffer is tight `[n_rows x - * dim]` while the on-disk stride may be larger. Do not call from user code; prefer - * `make_device_standard_dataset_view()` when wrapping existing correctly-strided storage. - * - * Potential future call sites if an owning copy with explicit stride is needed: - * - C API dataset upload (mirroring `make_device_padded_dataset` in `c/src/neighbors/cagra.cpp`) - * - `tiered_index` / composite index paths that materialize standard-layout device storage - * - Multigpu (MG) index build or merge when rehydrating a strided dataset from host fragments - */ -template -auto make_device_standard_dataset(const raft::resources& res, - SrcT const& src, - uint32_t logical_dim, - uint32_t target_stride) - -> std::unique_ptr> -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - return detail::make_device_dense_row_major_dataset_from_src< - device_standard_dataset, - value_type, - index_type>(res, src, logical_dim, target_stride, "make_device_standard_dataset_view"); -} - -template -auto make_host_standard_dataset_view(SrcT const& src) - -> host_standard_dataset_view -{ - using value_type = typename SrcT::value_type; - using index_type = typename SrcT::index_type; - return detail::make_host_dense_row_major_view_from_src< - value_type, - index_type, - host_standard_dataset_view>(src, static_cast(src.extent(1))); -} - namespace filtering { /** @@ -1892,10 +642,9 @@ struct iface { std::optional index_; /** Used by CAGRA when deserializing an index that contains a dataset; keeps it alive for the * view. */ - std::unique_ptr> cagra_owned_padded_dataset_; + std::unique_ptr> cagra_owned_padded_dataset_; /** Used by CAGRA standard-layout paths to keep deserialized/attached dataset views alive. */ - std::unique_ptr> - cagra_owned_standard_dataset_; + std::unique_ptr> cagra_owned_standard_dataset_; std::shared_ptr mutex_; }; diff --git a/cpp/include/cuvs/neighbors/nn_descent.hpp b/cpp/include/cuvs/neighbors/nn_descent.hpp index 1e1546c651..907fc2f089 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.hpp +++ b/cpp/include/cuvs/neighbors/nn_descent.hpp @@ -539,40 +539,40 @@ auto build(raft::resources const& res, */ auto build(raft::resources const& res, index_params const& params, - cuvs::neighbors::device_bbq_dataset_view dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, std::optional> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index; /** * @copydoc build(raft::resources const& res, index_params const& params, - * cuvs::neighbors::device_bbq_dataset_view dataset, + * cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, * std::optional> graph = std::nullopt) * -> cuvs::neighbors::nn_descent::index */ auto build(raft::resources const& res, index_params const& params, - cuvs::neighbors::device_bbq_dataset_view dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, std::optional> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index; /** * @copydoc build(raft::resources const& res, index_params const& params, - * cuvs::neighbors::device_bbq_dataset_view dataset, + * cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, * std::optional> graph = std::nullopt) * -> cuvs::neighbors::nn_descent::index */ auto build(raft::resources const& res, index_params const& params, - cuvs::neighbors::device_bbq_dataset_view dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, std::optional> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index; /** * @copydoc build(raft::resources const& res, index_params const& params, - * cuvs::neighbors::device_bbq_dataset_view dataset, + * cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, * std::optional> graph = std::nullopt) * -> cuvs::neighbors::nn_descent::index */ auto build(raft::resources const& res, index_params const& params, - cuvs::neighbors::device_bbq_dataset_view dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, std::optional> graph = std::nullopt) -> cuvs::neighbors::nn_descent::index; diff --git a/cpp/include/cuvs/neighbors/tiered_index.hpp b/cpp/include/cuvs/neighbors/tiered_index.hpp index 015e48c7e2..695e37d6e3 100644 --- a/cpp/include/cuvs/neighbors/tiered_index.hpp +++ b/cpp/include/cuvs/neighbors/tiered_index.hpp @@ -98,12 +98,12 @@ auto build(raft::resources const& res, auto convert_standard_to_padded_index( raft::resources const& res, const index>& idx, - cuvs::neighbors::device_padded_dataset_view padded_dataset) + cuvs::core::device_padded_dataset_view padded_dataset) -> index>; auto build(raft::resources const& res, const index_params& index_params, - cuvs::neighbors::device_padded_dataset_view dataset) + cuvs::core::device_padded_dataset_view dataset) -> tiered_index::index>; /** @copydoc build */ diff --git a/cpp/include/cuvs/neighbors/vamana.hpp b/cpp/include/cuvs/neighbors/vamana.hpp index 517e1f76e1..efec7bd953 100644 --- a/cpp/include/cuvs/neighbors/vamana.hpp +++ b/cpp/include/cuvs/neighbors/vamana.hpp @@ -133,7 +133,7 @@ struct index : cuvs::neighbors::index { /** Non-owning dataset view stored by the index (full-precision vectors may live in * `full_precision_storage_`). */ [[nodiscard]] inline auto data() const noexcept - -> const cuvs::neighbors::device_padded_dataset_view& + -> const cuvs::core::device_padded_dataset_view& { return dataset_.value(); } @@ -204,13 +204,13 @@ struct index : cuvs::neighbors::index { dataset.stride(0) > 0 ? static_cast(dataset.stride(0)) : dataset.extent(1); auto d_m = raft::make_device_matrix_view( dataset.data_handle(), dataset.extent(0), row_stride); - use_padded_view = cuvs::neighbors::matrix_row_width_matches_cagra_required(d_m); + use_padded_view = cuvs::core::matrix_has_padded_row_width(d_m); } if (use_padded_view) { - dataset_ = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + dataset_ = cuvs::core::make_device_padded_dataset_view(res, dataset); } else { - full_precision_storage_ = cuvs::neighbors::make_device_padded_dataset(res, dataset); + full_precision_storage_ = cuvs::core::make_device_padded_dataset(res, dataset); dataset_ = full_precision_storage_->as_dataset_view(); } update_graph(res, vamana_graph); @@ -288,8 +288,8 @@ struct index : cuvs::neighbors::index { raft::device_matrix graph_; raft::device_matrix_view graph_view_; /** Owns CAGRA-padded full-precision device storage for the index dataset view. */ - std::unique_ptr> full_precision_storage_; - std::optional> dataset_; + std::unique_ptr> full_precision_storage_; + std::optional> dataset_; raft::device_matrix quantized_dataset_; IdxT medoid_id_; }; diff --git a/cpp/include/cuvs/preprocessing/quantize/bbq.hpp b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp index 21366d47d0..45b6d50710 100644 --- a/cpp/include/cuvs/preprocessing/quantize/bbq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/bbq.hpp @@ -5,9 +5,9 @@ #pragma once +#include #include #include -#include #include #include @@ -242,30 +242,32 @@ void resolve_dequant_factors( } // namespace preprocessing::quantize::bbq -namespace neighbors { -struct bbq_dataset_container { - template - using owning_storage = cuvs::preprocessing::quantize::bbq::quantizer; - template - using view_storage = cuvs::preprocessing::quantize::bbq::quantizer_view; -}; +namespace preprocessing::quantize::bbq { + +// ----------------------------------------------------------------------------- +// BBQ dataset: a child of `cuvs::core::dataset` / `dataset_view`. The quantizers and the +// methods that manage them live in the payload types below; `dataset` itself knows nothing of them. +// ----------------------------------------------------------------------------- + +namespace detail { + +/** + * BBQ payloads: a BBQ dataset is a small bag of alternate encodings of the *same* rows, one per + * `bbq_code_layout`, selected at query time. The quantizers and the methods that manage them live + * here, in the BBQ payload, not in `dataset`/`dataset_view`; they are reached through `data()`. + */ +template +struct bbq_view_storage; -template -struct dataset { - using owning_storage_type = bbq_dataset_container::owning_storage; +template +struct bbq_owning_storage { + using value_type = DataT; + using owning_storage_type = cuvs::preprocessing::quantize::bbq::quantizer; std::vector quantizers; - dataset(owning_storage_type&& quantizer) noexcept { add_quantizer(std::move(quantizer)); } - [[nodiscard]] auto as_dataset_view() const noexcept - -> dataset_view> + explicit bbq_owning_storage(owning_storage_type&& quantizer) noexcept { - return dataset_view>{quantizers}; + add_quantizer(std::move(quantizer)); } [[nodiscard]] constexpr auto n_rows() const noexcept -> IdxT { @@ -290,15 +292,16 @@ struct dataset { } }; -template -struct dataset_view { - using owning_storage_type = bbq_dataset_container::owning_storage; - using view_storage_type = bbq_dataset_container::view_storage; +template +struct bbq_view_storage { + using value_type = DataT; + using owning_storage_type = cuvs::preprocessing::quantize::bbq::quantizer; + using view_storage_type = cuvs::preprocessing::quantize::bbq::quantizer_view; std::vector quantizers; - dataset_view() noexcept = default; + bbq_view_storage() noexcept = default; - dataset_view(const std::vector& quantizers) noexcept + bbq_view_storage(const std::vector& quantizers) noexcept { for (const auto& quantizer : quantizers) { add_quantizer(quantizer); @@ -339,50 +342,82 @@ struct dataset_view { } }; -template -using device_bbq_dataset = - dataset>; +} // namespace detail -template -using device_bbq_dataset_view = - dataset_view>; +/** BBQ is just another dataset type: it plugs its payloads into the shared `dataset`/ + * `dataset_view` through a spec, like every other kind does. */ +template +struct bbq_dataset_spec { + using accessor_type = Accessor; + template + using rebind_accessor = bbq_dataset_spec; -template -struct owning_dataset_for_view> { - using type = device_bbq_dataset; + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + + using data_type = detail::bbq_owning_storage; + using view_type = detail::bbq_view_storage; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return view_type{data.quantizers}; + } + template + [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type + { + return data.n_rows(); + } + template + [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t + { + return data.dim(); + } + }; }; +template +using device_bbq_dataset = cuvs::core:: + dataset>>; + +template +using device_bbq_dataset_view = cuvs::core:: + dataset_view>>; + +/** Spec predicate for `cuvs::core::dataset_view_has_spec_v`. */ +template +struct is_bbq_spec : std::false_type {}; +template +struct is_bbq_spec> : std::true_type {}; +template +inline constexpr bool is_bbq_spec_v = is_bbq_spec::value; + +/** True for an owning `dataset<...>` of the BBQ kind. */ template struct is_bbq_dataset : std::false_type {}; -template -struct is_bbq_dataset> : std::true_type {}; +template +struct is_bbq_dataset> + : std::bool_constant> {}; template inline constexpr bool is_bbq_dataset_v = is_bbq_dataset::value; -template -struct dataset_view_kind_of> { - static constexpr dataset_view_kind value = dataset_view_kind::bbq; -}; - -template -struct cagra_view_element_type> { - using type = DataT; -}; - template inline constexpr bool is_device_bbq_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::bbq && dataset_view_is_device_accessible_v; + cuvs::core::dataset_view_has_spec_v && + cuvs::core::dataset_view_is_device_accessible_v; template inline constexpr bool is_host_bbq_dataset_view_v = - dataset_view_kind_v == dataset_view_kind::bbq && !dataset_view_is_device_accessible_v; + cuvs::core::dataset_view_has_spec_v && + !cuvs::core::dataset_view_is_device_accessible_v; template inline constexpr bool is_bbq_dataset_view_v = is_device_bbq_dataset_view_v || is_host_bbq_dataset_view_v; -} // namespace neighbors +} // namespace preprocessing::quantize::bbq } // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/preprocessing/quantize/pq.hpp b/cpp/include/cuvs/preprocessing/quantize/pq.hpp index e77405ad7e..24a2cff52e 100644 --- a/cpp/include/cuvs/preprocessing/quantize/pq.hpp +++ b/cpp/include/cuvs/preprocessing/quantize/pq.hpp @@ -6,7 +6,7 @@ #pragma once #include -#include +#include #include #include #include @@ -16,6 +16,9 @@ #include #include #include +#ifdef __cpp_lib_bitops +#include +#endif namespace CUVS_EXPORT cuvs { namespace preprocessing { @@ -139,6 +142,308 @@ struct params { uint32_t max_train_points_per_vq_cluster = 1024; }; +/** Parameters for VPQ compression. */ +struct vpq_params { + /** + * The bit length of the vector element after compression by PQ. + * + * Possible values: [4, 5, 6, 7, 8]. + * + * Hint: the smaller the 'pq_bits', the smaller the index size and the better the search + * performance, but the lower the recall. + */ + uint32_t pq_bits = 8; + /** + * The dimensionality of the vector after compression by PQ. + * When zero, an optimal value is selected using a heuristic. + * + * TODO: at the moment `dim` must be a multiple `pq_dim`. + */ + uint32_t pq_dim = 0; + /** + * Vector Quantization (VQ) codebook size - number of "coarse cluster centers". + * When zero, an optimal value is selected using a heuristic. + */ + uint32_t vq_n_centers = 0; + /** The number of iterations searching for kmeans centers (both VQ & PQ phases). */ + uint32_t kmeans_n_iters = 25; + /** + * The fraction of data to use during iterative kmeans building (VQ phase). + * When zero, an optimal value is selected using a heuristic. + * @deprecated Prefer using `max_train_points_per_vq_cluster` instead. + */ + double vq_kmeans_trainset_fraction = 0; + /** + * The fraction of data to use during iterative kmeans building (PQ phase). + * When zero, an optimal value is selected using a heuristic. + * @deprecated Prefer using `max_train_points_per_pq_code` instead. + */ + double pq_kmeans_trainset_fraction = 0; + /** + * Type of k-means algorithm for PQ training. + * Balanced k-means tends to be faster than regular k-means for PQ training, for + * problem sets where the number of points per cluster are approximately equal. + * Regular k-means may be better for skewed cluster distributions. + */ + cuvs::cluster::kmeans::kmeans_type pq_kmeans_type = + cuvs::cluster::kmeans::kmeans_type::KMeansBalanced; + /** + * The max number of data points to use per PQ code during PQ codebook training. Using more data + * points per PQ code may increase the quality of PQ codebook but may also increase the build + * time. We will use `pq_n_centers * max_train_points_per_pq_code` training + * points to train each PQ codebook. + */ + uint32_t max_train_points_per_pq_code = 256; + /** + * The max number of data points to use per VQ cluster during training. + */ + uint32_t max_train_points_per_vq_cluster = 1024; +}; + +// ----------------------------------------------------------------------------- +// VPQ dataset: a child of `cuvs::core::dataset` / `dataset_view`. All VPQ-specific state +// (codebooks) and methods live in the payload types below; `dataset` itself knows nothing of them. +// ----------------------------------------------------------------------------- + +namespace detail { + +// The accessor aliases are shared with every dataset kind and live next to `dataset`. +using cuvs::core::detail::device_owning_accessor; +using cuvs::core::detail::host_owning_accessor; + +// VPQ codes are always uint8_t regardless of MathT, so retarget the owning accessor's element +// type instead of re-deriving a device/host matrix; residency is still driven by Accessor. +template +using owning_accessor_with_value_type = std::conditional_t, + host_owning_accessor>; + +// The two matrix types of a VPQ dataset. Both are owning (`mdarray`); the non-owning form of each +// is its `const_view_type`, which the view payload below uses directly. + +/** Encoded rows (`uint8_t`): each row holds the VQ code followed by the PQ codes. Owning. */ +template +using vpq_data_matrix = raft::mdarray, + raft::row_major, + owning_accessor_with_value_type>; + +/** A codebook (used for both the VQ and the PQ codebook). Owning. */ +template +using vpq_codebook_matrix = + raft::mdarray, raft::row_major, Accessor>; + +/** Read-only helpers derived from the codebook shapes; shared by the owning and view payloads. + * `Derived` provides `vq_code_book`, `pq_code_book` and the codes' `extent(r)`. */ +template +struct vpq_codebook_helpers { + /** Logical dimension: it comes from the VQ codebook, not from the encoded rows (row padding + * makes the encoded-row width ambiguous as a dimension). */ + [[nodiscard]] auto dim() const noexcept -> uint32_t + { + return static_cast(self().vq_code_book.extent(1)); + } + [[nodiscard]] auto vq_n_centers() const noexcept -> uint32_t + { + return static_cast(self().vq_code_book.extent(0)); + } + [[nodiscard]] auto pq_n_centers() const noexcept -> uint32_t + { + return static_cast(self().pq_code_book.extent(0)); + } + [[nodiscard]] auto pq_len() const noexcept -> uint32_t + { + return static_cast(self().pq_code_book.extent(1)); + } + [[nodiscard]] auto pq_bits() const noexcept -> uint32_t + { + auto pq_width = pq_n_centers(); +#ifdef __cpp_lib_bitops + return std::countr_zero(pq_width); +#else + uint32_t bits = 0; + while (pq_width > 1) { + bits++; + pq_width >>= 1; + } + return bits; +#endif + } + [[nodiscard]] auto pq_dim() const noexcept -> uint32_t + { + return raft::div_rounding_up_unsafe(dim(), pq_len()); + } + [[nodiscard]] auto encoded_row_length() const noexcept -> uint32_t + { + return static_cast(self().extent(1)); + } + + private: + [[nodiscard]] auto self() const noexcept -> Derived const& + { + return static_cast(*this); + } +}; + +/** The VPQ payload, defined once. `CodesT` is the encoded-rows matrix and `BookT` the type of each + * codebook: raft `mdarray`s for an owning dataset, `mdspan`s for a non-owning view. The payload + * *is* the `uint8_t` codes matrix (it derives from `CodesT`) and additionally holds the VQ and PQ + * codebooks, so `vq_code_book`, `pq_code_book` and the helpers below mean the same thing in both + * forms and only the ownership of the arrays differs. Use the aliases below rather than naming + * this template directly. */ +template +struct vpq_storage : public CodesT, public vpq_codebook_helpers> { + BookT vq_code_book; + BookT pq_code_book; + + // Only usable when every member is default-constructible, i.e. for the view form. + vpq_storage() noexcept = default; + + vpq_storage(CodesT&& codes, BookT&& vq_codes, BookT&& pq_codes) noexcept + : CodesT(std::move(codes)), vq_code_book(std::move(vq_codes)), pq_code_book(std::move(pq_codes)) + { + } +}; + +/** Owning VPQ payload: `mdarray` codes and codebooks. `Accessor` drives both codebook and code + * residency. */ +template +using vpq_owning_storage = + vpq_storage, vpq_codebook_matrix>; + +/** Non-owning VPQ payload: `mdspan` views of the codes and of both codebooks. */ +template +using vpq_view_storage = + vpq_storage::const_view_type, + typename vpq_codebook_matrix::const_view_type>; + +} // namespace detail + +/** `Accessor` drives both codebook and code residency, mirroring today's + * single-`Accessor`-per-VPQ-dataset design. The payload (`detail::vpq_owning_storage` / + * `detail::vpq_view_storage`) holds the encoded rows and the VQ/PQ codebooks. */ +template +struct vpq_dataset_spec { + using accessor_type = Accessor; + template + using rebind_accessor = vpq_dataset_spec; + + template + struct apply { + using value_type = std::remove_cv_t; + using index_type = std::remove_cv_t; + using math_type = MathT; + + using data_type = detail::vpq_owning_storage; + using view_type = detail::vpq_view_storage; + + [[nodiscard]] static auto get_data_view(data_type const& data) noexcept -> view_type + { + return view_type(data.view(), data.vq_code_book.view(), data.pq_code_book.view()); + } + template + [[nodiscard]] static auto get_n_rows(AnyStorage const& data) noexcept -> index_type + { + return static_cast(data.extent(0)); + } + template + [[nodiscard]] static auto get_dim(AnyStorage const& data) noexcept -> uint32_t + { + return data.dim(); + } + }; +}; + +template +using device_vpq_dataset = + cuvs::core::dataset>>; + +template +using device_vpq_dataset_view = cuvs::core:: + dataset_view>>; + +template +using host_vpq_dataset = + cuvs::core::dataset>>; + +template +using host_vpq_dataset_view = cuvs::core:: + dataset_view>>; + +/** Spec predicate for `cuvs::core::dataset_view_has_spec_v`. */ +template +struct is_vpq_spec : std::false_type {}; +template +struct is_vpq_spec> : std::true_type {}; +template +inline constexpr bool is_vpq_spec_v = is_vpq_spec::value; + +template +struct vpq_spec_math_type {}; +template +struct vpq_spec_math_type> { + using type = MathT; +}; +template +using vpq_spec_math_type_t = typename vpq_spec_math_type::type; + +/** True for an owning `dataset<...>` of the VPQ kind. */ +template +struct is_vpq_dataset : std::false_type {}; +template +struct is_vpq_dataset> + : std::bool_constant> {}; +template +inline constexpr bool is_vpq_dataset_v = is_vpq_dataset::value; + +/** True when `V` is a VPQ `dataset_view` whose codebooks have element type `MathT`. */ +template +struct is_vpq_dataset_view_with_math : std::false_type {}; +template +struct is_vpq_dataset_view_with_math< + cuvs::core::dataset_view>, + MathT> : std::is_same {}; +template +inline constexpr bool is_vpq_dataset_view_with_math_v = + is_vpq_dataset_view_with_math, MathT>::value; + +template +inline constexpr bool is_device_vpq_f16_dataset_view_v = + is_vpq_dataset_view_with_math_v && cuvs::core::dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_vpq_f16_dataset_view_v = + is_vpq_dataset_view_with_math_v && !cuvs::core::dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_vpq_f16_dataset_view_v = + is_device_vpq_f16_dataset_view_v || is_host_vpq_f16_dataset_view_v; + +template +inline constexpr bool is_device_vpq_f32_dataset_view_v = + is_vpq_dataset_view_with_math_v && cuvs::core::dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_host_vpq_f32_dataset_view_v = + is_vpq_dataset_view_with_math_v && !cuvs::core::dataset_view_is_device_accessible_v; + +template +inline constexpr bool is_vpq_f32_dataset_view_v = + is_device_vpq_f32_dataset_view_v || is_host_vpq_f32_dataset_view_v; + +template +inline constexpr bool is_device_vpq_dataset_view_v = + is_device_vpq_f16_dataset_view_v || is_device_vpq_f32_dataset_view_v; + +template +inline constexpr bool is_host_vpq_dataset_view_v = + is_host_vpq_f16_dataset_view_v || is_host_vpq_f32_dataset_view_v; + +/** True for any VPQ dataset view (host or device, f16 or f32 codebooks). */ +template +inline constexpr bool is_vpq_dataset_view_v = + is_device_vpq_dataset_view_v || is_host_vpq_dataset_view_v; + /** * @brief Defines and stores VPQ codebooks upon training * @@ -150,7 +455,7 @@ struct quantizer { /** Parameters used to build this quantizer. */ params params_quantizer; /** VPQ codebooks produced during training. */ - cuvs::neighbors::device_vpq_dataset vpq_codebooks; + device_vpq_dataset vpq_codebooks; }; /** @@ -255,9 +560,9 @@ namespace detail { // default visibility, an instantiation cannot be exported from the shared library when one of its // template arguments (`half`, or any mdspan type) is itself hidden, because the visibility of an // instantiation is capped by that of its template arguments. -[[nodiscard]] CUVS_EXPORT cuvs::neighbors::device_vpq_dataset vpq_train_from_rows( +[[nodiscard]] CUVS_EXPORT device_vpq_dataset vpq_train_from_rows( raft::resources const& res, - cuvs::neighbors::vpq_params const& params, + vpq_params const& params, void const* src_ptr, cudaDataType_t dtype, int64_t n_rows, @@ -270,7 +575,7 @@ namespace detail { * @brief Train VPQ storage (codebooks + encoded rows) from a row-major mdspan/mdarray/dataset. * * Accepts either a row-major mdspan with `value_type`, `extent`, `stride`, and `data_handle` (same - * pattern as `cuvs::neighbors::make_device_padded_dataset`), or any cuVS dense dataset / dataset + * pattern as `cuvs::core::make_device_padded_dataset`), or any cuVS dense dataset / dataset * view exposing `view`, `dim` and `stride`, in which case the logical `dim()` is quantized and the * row padding is skipped. The rows may be device-accessible or host-resident. Device-accessible * rows (device, managed or pinned) with tight row-major storage (logical stride equals dimension) @@ -288,24 +593,23 @@ namespace detail { * #include * * // `padded` is a `device_padded_dataset_view` over the source rows. - * cuvs::neighbors::vpq_params vpq_params{}; + * cuvs::preprocessing::quantize::pq::vpq_params vpq_params{}; * auto vpq = cuvs::preprocessing::quantize::pq::make_vpq_dataset(res, vpq_params, padded); * auto idx = cuvs::neighbors::cagra::build(res, cagra_params, vpq.as_dataset_view()); * @endcode */ template [[nodiscard]] auto make_vpq_dataset(raft::resources const& res, - cuvs::neighbors::vpq_params const& params, - SrcT const& src) - -> cuvs::neighbors::device_vpq_dataset + vpq_params const& params, + SrcT const& src) -> device_vpq_dataset { - // A cuVS dataset keeps its logical width in `dim()` while `view()` spans the full row pitch. + // A cuVS dataset keeps its logical width in `dim()` while `as_matrix_view()` spans the full row + // pitch. if constexpr (requires { - src.view(); + src.as_matrix_view(); src.dim(); - src.stride(); }) { - auto const rows = src.view(); + auto const rows = src.as_matrix_view(); using value_type = typename decltype(rows)::value_type; using extents_type = raft::matrix_extent; return make_vpq_dataset( @@ -314,7 +618,7 @@ template raft::mdspan{ rows.data_handle(), raft::make_strided_layout(extents_type{rows.extent(0), int64_t{src.dim()}}, - cuda::std::array{int64_t{src.stride()}, 1})}); + cuda::std::array{int64_t{rows.stride()}, 1})}); } else { using value_type = typename SrcT::value_type; static_assert(std::is_same_v || std::is_same_v || diff --git a/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp b/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp index 603517e518..6b17492e6d 100644 --- a/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp +++ b/cpp/internal/cuvs_internal/preprocessing/bbq_cpu_quantize.hpp @@ -331,10 +331,11 @@ inline host_quantizer_storage quantize(const float* data, template auto copy_bbq_owning_storage_host_to_device(raft::resources const& res, host_quantizer_storage const& host_storage) -> - typename cuvs::neighbors::device_bbq_dataset::owning_storage_type + typename cuvs::preprocessing::quantize::bbq::device_bbq_dataset::data_type:: + owning_storage_type { - using device_storage = - typename cuvs::neighbors::device_bbq_dataset::owning_storage_type; + using device_storage = typename cuvs::preprocessing::quantize::bbq:: + device_bbq_dataset::data_type::owning_storage_type; auto stream = raft::resource::get_cuda_stream(res); device_storage device{res, static_cast(host_storage.codes.extent(0)), @@ -383,13 +384,13 @@ auto copy_bbq_owning_storage_host_to_device(raft::resources const& res, template auto make_device_bbq_dataset(raft::resources const& res, std::vector const& host) - -> cuvs::neighbors::device_bbq_dataset + -> cuvs::preprocessing::quantize::bbq::device_bbq_dataset { RAFT_EXPECTS(host.size() != 0, "host BBQ dataset has no storage"); - cuvs::neighbors::device_bbq_dataset device{ + cuvs::preprocessing::quantize::bbq::device_bbq_dataset device{ copy_bbq_owning_storage_host_to_device(res, host[0])}; for (std::size_t i = 1; i < host.size(); ++i) { - device.add_quantizer(copy_bbq_owning_storage_host_to_device(res, host[i])); + device.data().add_quantizer(copy_bbq_owning_storage_host_to_device(res, host[i])); } return device; } @@ -595,7 +596,7 @@ inline auto quantize_to_device(raft::resources const& res, cuvs::distance::DistanceType metric, bbq_code_layout query_layout, bbq_code_layout doc_layout) - -> cuvs::neighbors::device_bbq_dataset + -> cuvs::preprocessing::quantize::bbq::device_bbq_dataset { validate_layout_pair(query_layout, doc_layout); std::vector host; @@ -616,7 +617,7 @@ inline auto quantize_to_device(raft::resources const& res, cuvs::distance::DistanceType metric, std::string_view query_token, std::string_view doc_token) - -> cuvs::neighbors::device_bbq_dataset + -> cuvs::preprocessing::quantize::bbq::device_bbq_dataset { return quantize_to_device(res, rows, diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 0e1454d987..b294b50ad5 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -32,21 +32,20 @@ namespace cuvs::neighbors::cagra { -template +template CUVS_EXPORT void index::compute_dataset_norms_(raft::resources const& res) { // raft::linalg::reduce wants row-major with leading dim = row pitch in elements. // Skip norm precomputation for VPQ/empty/non-dense views; CosineExpanded with VPQ is handled // (or rejected) on the search path. - namespace nb = cuvs::neighbors; bool skip_norms = false; std::optional> rm_dataset; - if constexpr (nb::is_padded_dataset_view_v || - nb::is_standard_dataset_view_v) { - rm_dataset = dataset_.view(); - } else if constexpr (nb::is_vpq_dataset_view_v || - nb::is_bbq_dataset_view_v) { + if constexpr (cuvs::core::is_padded_dataset_view_v || + cuvs::core::is_standard_dataset_view_v) { + rm_dataset = dataset_.as_matrix_view(); + } else if constexpr (cuvs::preprocessing::quantize::pq::is_vpq_dataset_view_v || + cuvs::preprocessing::quantize::bbq::is_bbq_dataset_view_v) { skip_norms = true; } @@ -106,7 +105,7 @@ CUVS_EXPORT void index::compute_dataset_norms_(raft::reso * auto optimized_gaph = raft::make_host_matrix(dataset.extent(0), 64); * cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto index = cagra::index>( + * auto index = cagra::index>( * res, build_params.metric(), dataset, optimized_graph.view()); * @endcode * @@ -210,7 +209,7 @@ void build_knn_graph( * // optimize graph * cagra::optimize(res, dataset, knn_graph.view(), optimized_graph.view()); * // Construct an index from dataset and optimized knn_graph - * auto index = cagra::index>( + * auto index = cagra::index>( * res, build_params.metric(), dataset, optimized_graph.view()); * @endcode * @@ -293,10 +292,10 @@ void optimize( * device-padded dataset to get a search-ready device index. */ template - requires(!cuvs::neighbors::is_empty_dataset_view_v && - (cuvs::neighbors::is_device_dataset_view_v || - cuvs::neighbors::is_host_dataset_view_v)) -auto build(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) + requires(!cuvs::core::is_empty_dataset_view_v && + (cuvs::core::is_device_dataset_view_v || + cuvs::core::is_host_dataset_view_v)) +auto build(raft::resources const& res, const index_params& params, DatasetViewT const& dataset_view) -> cuvs::neighbors::cagra::cagra_index_t { using index_type = cuvs::neighbors::cagra::cagra_index_t; @@ -305,10 +304,11 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT // Dense paths build the graph and optionally attach the input dataset view. Host indexes remain // non-searchable until the type-changing update_dataset(...) supplies a device-padded dataset. - if constexpr (cuvs::neighbors::is_device_bbq_dataset_view_v) { + if constexpr (cuvs::preprocessing::quantize::bbq::is_device_bbq_dataset_view_v) { return cuvs::neighbors::cagra::detail::build_from_bbq_dataset( - res, params, dataset); - } else if constexpr (cuvs::neighbors::is_device_vpq_dataset_view_v) { + res, params, dataset_view); + } else if constexpr (cuvs::preprocessing::quantize::pq::is_device_vpq_dataset_view_v< + DatasetViewT>) { auto effective_params = params; if (std::holds_alternative(effective_params.graph_build_params)) { effective_params.graph_build_params = graph_build_params::iterative_search_params{}; @@ -319,39 +319,41 @@ auto build(raft::resources const& res, const index_params& params, DatasetViewT "cagra::build: a VPQ dataset requires iterative_search_params graph construction"); RAFT_EXPECTS(effective_params.metric == cuvs::distance::DistanceType::L2Expanded, "cagra::build: a VPQ dataset supports only L2Expanded distance"); - RAFT_EXPECTS(dataset.n_rows() > 0, "cagra::build: VPQ dataset must not be empty"); - RAFT_EXPECTS(dataset.dset().pq_bits() == 8, + RAFT_EXPECTS(dataset_view.n_rows() > 0, "cagra::build: VPQ dataset must not be empty"); + auto const& vpq_view = dataset_view.data(); + RAFT_EXPECTS(vpq_view.pq_bits() == 8, "cagra::build: VPQ dataset requires pq_bits == 8, got %u", - dataset.dset().pq_bits()); - auto const pq_len = dataset.dset().pq_len(); + vpq_view.pq_bits()); + auto const pq_len = vpq_view.pq_len(); RAFT_EXPECTS(pq_len == 2 || pq_len == 4 || pq_len == 8, "cagra::build: VPQ dataset requires pq_len in {2, 4, 8}, got %u", pq_len); detail::check_graph_degree(effective_params.intermediate_graph_degree, effective_params.graph_degree, - static_cast(dataset.n_rows())); - auto cagra_graph = detail::iterative_build_graph(res, effective_params, dataset); + static_cast(dataset_view.n_rows())); + auto cagra_graph = detail::iterative_build_graph(res, effective_params, dataset_view); index_type idx(res, effective_params.metric); idx.update_graph(res, std::move(cagra_graph)); if (effective_params.attach_dataset_on_build) { - idx = cuvs::neighbors::cagra::update_dataset(res, std::move(idx), dataset); + idx = cuvs::neighbors::cagra::update_dataset(res, std::move(idx), dataset_view); } return idx; - } else if constexpr (cuvs::neighbors::is_dense_row_major_device_dataset_view_v) { + } else if constexpr (cuvs::core::is_dense_row_major_device_dataset_view_v) { auto idx = cuvs::neighbors::cagra::detail::build_from_device_matrix( - res, params, dataset); + res, params, dataset_view); if (params.attach_dataset_on_build) { - idx = cuvs::neighbors::cagra::update_dataset(res, std::move(idx), dataset); + idx = cuvs::neighbors::cagra::update_dataset(res, std::move(idx), dataset_view); } return idx; } else { if (std::holds_alternative(params.graph_build_params)) { - return cuvs::neighbors::cagra::detail::build_ace(res, params, dataset); + return cuvs::neighbors::cagra::detail::build_ace( + res, params, dataset_view); } return cuvs::neighbors::cagra::detail::build_from_host_matrix( - res, params, dataset); + res, params, dataset_view); } } @@ -398,7 +400,7 @@ template + cuvs::core::dataset_like DatasetViewT> void search_with_filtering(raft::resources const& res, const search_params& params, const index& idx, @@ -420,10 +422,7 @@ void search_with_filtering(raft::resources const& res, res, params, idx, queries, neighbors, distances, sample_filter); } -template +template void search(raft::resources const& res, const search_params& params, const index& idx, @@ -519,19 +518,19 @@ void search(raft::resources const& res, } } -template +template void extend(raft::resources const& handle, const cagra::extend_params& params, - cuvs::neighbors::device_padded_dataset_view extended_dataset, + cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row, cuvs::neighbors::cagra::index& index) { - static_assert(cuvs::neighbors::is_padded_dataset_view_v, + static_assert(cuvs::core::is_padded_dataset_view_v, "cagra::extend requires a padded index dataset type"); extend_core(handle, index, params, extended_dataset, new_start_row); } -template +template cuvs::neighbors::cagra::index merge( raft::resources const& handle, const cagra::index_params& params, @@ -543,7 +542,7 @@ cuvs::neighbors::cagra::index merge( handle, params, indices, merged_dataset, row_filter); } -template +template cuvs::neighbors::cagra::index merge( raft::resources const& handle, const cagra::index_params& params, @@ -622,8 +621,8 @@ void search( template + cuvs::core::dataset_like SrcDatasetViewT, + cuvs::core::dataset_like DstDatasetViewT> auto update_dataset(raft::resources const& res, index&& cagra_index, DstDatasetViewT dataset) -> index diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index 2a3e2b43c1..8145c62def 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -15,12 +15,13 @@ namespace { using data_t = @data_type@; using index_t = @index_type@; -using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; -using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; -using inst_host_padded_view_t = cuvs::neighbors::host_padded_dataset_view; -using inst_host_standard_view_t = cuvs::neighbors::host_standard_dataset_view; -using inst_vpq_view_t = cuvs::neighbors::device_vpq_dataset_view; -using inst_bbq_view_t = cuvs::neighbors::device_bbq_dataset_view; +using inst_device_padded_view_t = cuvs::core::device_padded_dataset_view; +using inst_device_standard_view_t = cuvs::core::device_standard_dataset_view; +using inst_host_padded_view_t = cuvs::core::host_padded_dataset_view; +using inst_host_standard_view_t = cuvs::core::host_standard_dataset_view; +using inst_vpq_view_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset_view; +using inst_bbq_view_t = + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view; } // namespace namespace cuvs::neighbors::cagra { diff --git a/cpp/src/neighbors/cagra_extend_inst.cu.in b/cpp/src/neighbors/cagra_extend_inst.cu.in index 6f45fe1fd3..b9573a0f91 100644 --- a/cpp/src/neighbors/cagra_extend_inst.cu.in +++ b/cpp/src/neighbors/cagra_extend_inst.cu.in @@ -10,7 +10,7 @@ namespace { using data_t = @data_type@; using index_t = @index_type@; -using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; +using inst_device_padded_view_t = cuvs::core::device_padded_dataset_view; } // namespace @@ -19,15 +19,15 @@ namespace cuvs::neighbors::cagra { extern template void index::compute_dataset_norms_( raft::resources const&); -#define CUVS_INST_CAGRA_EXTEND(T, IdxT, DatasetViewT) \ - void extend(raft::resources const& handle, \ - const cuvs::neighbors::cagra::extend_params& params, \ - cuvs::neighbors::device_padded_dataset_view extended_dataset, \ - int64_t new_start_row, \ - cuvs::neighbors::cagra::index& index) \ - { \ - cuvs::neighbors::cagra::extend( \ - handle, params, extended_dataset, new_start_row, index); \ +#define CUVS_INST_CAGRA_EXTEND(T, IdxT, DatasetViewT) \ + void extend(raft::resources const& handle, \ + const cuvs::neighbors::cagra::extend_params& params, \ + cuvs::core::device_padded_dataset_view extended_dataset, \ + int64_t new_start_row, \ + cuvs::neighbors::cagra::index& index) \ + { \ + cuvs::neighbors::cagra::extend( \ + handle, params, extended_dataset, new_start_row, index); \ } CUVS_INST_CAGRA_EXTEND(data_t, index_t, inst_device_padded_view_t); diff --git a/cpp/src/neighbors/cagra_merge_inst.cu.in b/cpp/src/neighbors/cagra_merge_inst.cu.in index da95e279b6..4ed7fe078f 100644 --- a/cpp/src/neighbors/cagra_merge_inst.cu.in +++ b/cpp/src/neighbors/cagra_merge_inst.cu.in @@ -12,13 +12,13 @@ namespace { using data_t = @data_type@; using index_t = @index_type@; -using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; +using inst_device_padded_view_t = cuvs::core::device_padded_dataset_view; } // namespace namespace cuvs::neighbors::cagra { -using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; +using inst_device_standard_view_t = cuvs::core::device_standard_dataset_view; extern template void index::compute_dataset_norms_( raft::resources const&); diff --git a/cpp/src/neighbors/cagra_search_inst.cu.in b/cpp/src/neighbors/cagra_search_inst.cu.in index 9e42ac476f..4e91a50a8d 100644 --- a/cpp/src/neighbors/cagra_search_inst.cu.in +++ b/cpp/src/neighbors/cagra_search_inst.cu.in @@ -10,11 +10,13 @@ namespace { using data_t = @data_type@; -using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; -using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; -using inst_vpq_f16_view_t = cuvs::neighbors::device_vpq_dataset_view; -using inst_bbq_view_t = cuvs::neighbors::device_bbq_dataset_view; -using inst_empty_view_t = cuvs::neighbors::device_empty_dataset_view; +using inst_device_padded_view_t = cuvs::core::device_padded_dataset_view; +using inst_device_standard_view_t = cuvs::core::device_standard_dataset_view; +using inst_vpq_f16_view_t = + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view; +using inst_bbq_view_t = + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view; +using inst_empty_view_t = cuvs::core::device_empty_dataset_view; } // namespace diff --git a/cpp/src/neighbors/cagra_serialize.cuh b/cpp/src/neighbors/cagra_serialize.cuh index 18e807b90f..284cc2bb5f 100644 --- a/cpp/src/neighbors/cagra_serialize.cuh +++ b/cpp/src/neighbors/cagra_serialize.cuh @@ -23,7 +23,7 @@ namespace cuvs::neighbors::cagra { raft::resources const& handle, \ const std::string& filename, \ cuvs::neighbors::cagra::device_padded_index* index, \ - std::unique_ptr>* out_dataset) \ + std::unique_ptr>* out_dataset) \ { \ cuvs::neighbors::cagra::detail::deserialize( \ handle, filename, index, out_dataset); \ @@ -41,7 +41,7 @@ namespace cuvs::neighbors::cagra { raft::resources const& handle, \ std::istream& is, \ cuvs::neighbors::cagra::device_padded_index* index, \ - std::unique_ptr>* out_dataset) \ + std::unique_ptr>* out_dataset) \ { \ cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ } \ @@ -59,7 +59,7 @@ namespace cuvs::neighbors::cagra { raft::resources const& handle, \ const std::string& filename, \ cuvs::neighbors::cagra::device_standard_index* index, \ - std::unique_ptr>* out_dataset) \ + std::unique_ptr>* out_dataset) \ { \ cuvs::neighbors::cagra::detail::deserialize( \ handle, filename, index, out_dataset); \ @@ -77,7 +77,7 @@ namespace cuvs::neighbors::cagra { raft::resources const& handle, \ std::istream& is, \ cuvs::neighbors::cagra::device_standard_index* index, \ - std::unique_ptr>* out_dataset) \ + std::unique_ptr>* out_dataset) \ { \ cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ } \ @@ -90,11 +90,10 @@ namespace cuvs::neighbors::cagra { handle, filename, index, include_dataset); \ } \ \ - void deserialize( \ - raft::resources const& handle, \ - const std::string& filename, \ - cuvs::neighbors::cagra::host_padded_index* index, \ - std::unique_ptr>* out_dataset) \ + void deserialize(raft::resources const& handle, \ + const std::string& filename, \ + cuvs::neighbors::cagra::host_padded_index* index, \ + std::unique_ptr>* out_dataset) \ { \ cuvs::neighbors::cagra::detail::deserialize( \ handle, filename, index, out_dataset); \ @@ -109,11 +108,10 @@ namespace cuvs::neighbors::cagra { handle, os, index, include_dataset); \ } \ \ - void deserialize( \ - raft::resources const& handle, \ - std::istream& is, \ - cuvs::neighbors::cagra::host_padded_index* index, \ - std::unique_ptr>* out_dataset) \ + void deserialize(raft::resources const& handle, \ + std::istream& is, \ + cuvs::neighbors::cagra::host_padded_index* index, \ + std::unique_ptr>* out_dataset) \ { \ cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ } \ @@ -131,7 +129,7 @@ namespace cuvs::neighbors::cagra { raft::resources const& handle, \ const std::string& filename, \ cuvs::neighbors::cagra::host_standard_index* index, \ - std::unique_ptr>* out_dataset) \ + std::unique_ptr>* out_dataset) \ { \ cuvs::neighbors::cagra::detail::deserialize( \ handle, filename, index, out_dataset); \ @@ -150,7 +148,7 @@ namespace cuvs::neighbors::cagra { raft::resources const& handle, \ std::istream& is, \ cuvs::neighbors::cagra::host_standard_index* index, \ - std::unique_ptr>* out_dataset) \ + std::unique_ptr>* out_dataset) \ { \ cuvs::neighbors::cagra::detail::deserialize(handle, is, index, out_dataset); \ } \ diff --git a/cpp/src/neighbors/cagra_serialize_inst.cu.in b/cpp/src/neighbors/cagra_serialize_inst.cu.in index 58e555d17e..0d8d399c69 100644 --- a/cpp/src/neighbors/cagra_serialize_inst.cu.in +++ b/cpp/src/neighbors/cagra_serialize_inst.cu.in @@ -10,9 +10,10 @@ namespace { using data_t = @data_type@; -using inst_device_padded_view_t = cuvs::neighbors::device_padded_dataset_view; -using inst_device_standard_view_t = cuvs::neighbors::device_standard_dataset_view; -using inst_vpq_f16_view_t = cuvs::neighbors::device_vpq_dataset_view; +using inst_device_padded_view_t = cuvs::core::device_padded_dataset_view; +using inst_device_standard_view_t = cuvs::core::device_standard_dataset_view; +using inst_vpq_f16_view_t = + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view; } // namespace diff --git a/cpp/src/neighbors/detail/cagra/add_nodes.cuh b/cpp/src/neighbors/detail/cagra/add_nodes.cuh index 7f2b372592..e8ff3d95ac 100644 --- a/cpp/src/neighbors/detail/cagra/add_nodes.cuh +++ b/cpp/src/neighbors/detail/cagra/add_nodes.cuh @@ -20,7 +20,7 @@ namespace cuvs::neighbors::cagra { -template +template void add_node_core( raft::resources const& handle, const cuvs::neighbors::cagra::index& idx, @@ -276,7 +276,7 @@ void add_node_core( } } -template +template void add_graph_nodes( raft::resources const& handle, raft::device_matrix_view input_updated_dataset_view, @@ -301,7 +301,7 @@ void add_graph_nodes( updated_graph_view.data_handle(), initial_dataset_size, degree); raft::copy(handle, updated_graph_prefix, raft::make_const_mdspan(index.graph())); - using padded_view_t = cuvs::neighbors::device_padded_dataset_view; + using padded_view_t = cuvs::core::device_padded_dataset_view; auto zero_row = raft::make_device_matrix_view( static_cast(nullptr), int64_t{0}, static_cast(dim)); padded_view_t device_empty_dataset_view(zero_row, static_cast(dim)); @@ -324,7 +324,7 @@ void add_graph_nodes( // add_node_core() uses CAGRA search internally, which requires a padded device dataset. // Keep this path allocation-free by requiring pre-padded chunk views. - auto pdv = cuvs::neighbors::make_device_padded_dataset_view(handle, dataset_view); + auto pdv = cuvs::core::make_device_padded_dataset_view(handle, dataset_view); internal_index = cuvs::neighbors::cagra::update_dataset(handle, std::move(internal_index), pdv); // Note: The graph is copied to the device memory. @@ -348,14 +348,14 @@ void add_graph_nodes( } } -template +template void extend_core(raft::resources const& handle, cuvs::neighbors::cagra::index& index, const cagra::extend_params& params, - cuvs::neighbors::device_padded_dataset_view extended_dataset, + cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row) { - static_assert(cuvs::neighbors::is_padded_dataset_view_v, + static_assert(cuvs::core::is_padded_dataset_view_v, "cagra::extend requires a padded dataset view index type"); RAFT_EXPECTS(!index.dataset_fd().has_value(), "Cannot extend a disk-backed CAGRA index. Convert it with " @@ -363,7 +363,7 @@ void extend_core(raft::resources const& handle, "cuvs::neighbors::hnsw::deserialize() before calling extend()."); const std::size_t initial_dataset_size = index.size(); - const auto extended_view = extended_dataset.view(); + const auto extended_view = extended_dataset.as_matrix_view(); const std::size_t new_dataset_size = static_cast(extended_view.extent(0)); const std::size_t degree = index.graph_degree(); const std::size_t dim = index.dim(); @@ -385,12 +385,12 @@ void extend_core(raft::resources const& handle, dim); auto const& leaf = index.dataset(); - if constexpr (cuvs::neighbors::is_empty_dataset_view_v>) { + if constexpr (cuvs::core::is_empty_dataset_view_v>) { RAFT_FAIL( "cagra::extend only supports an index to which the dataset is attached. Please check if the " "index has an empty dataset; attach one with update_dataset before " "extend."); - } else if constexpr (!cuvs::neighbors::is_padded_dataset_view_v>) { + } else if constexpr (!cuvs::core::is_padded_dataset_view_v>) { RAFT_FAIL("cagra::extend only supports an uncompressed padded dataset index"); } else { // Caller owns dataset concatenation. Extend only grows the graph and rebinds the view. diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 9a22676334..472e16656e 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -1303,14 +1303,14 @@ void ace_validate_disk_mode_partitions(raft::resources const& res, } template - requires cuvs::neighbors::is_dense_row_major_device_dataset_view_v + requires cuvs::core::is_dense_row_major_device_dataset_view_v auto build_from_device_matrix(raft::resources const& res, const index_params& params, DatasetViewT const& device_dataset) -> cuvs::neighbors::cagra::index; template - requires cuvs::neighbors::is_host_dataset_view_v + requires cuvs::core::is_host_dataset_view_v auto build_from_host_matrix(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) @@ -1326,7 +1326,7 @@ auto build_from_host_matrix(raft::resources const& res, // In disk mode, the graph is stored in build_dir and dataset is reordered on disk. // The returned index is not usable for search. Use the created files for search instead. template - requires cuvs::neighbors::is_host_dataset_view_v + requires cuvs::core::is_host_dataset_view_v auto build_ace(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) -> cuvs::neighbors::cagra::index { @@ -1347,7 +1347,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi params.graph_degree, npartitions); - auto dataset_view = dataset.view(); + auto dataset_view = dataset.as_matrix_view(); size_t dataset_size = dataset.n_rows(); size_t dataset_dim = dataset.dim(); @@ -1614,14 +1614,12 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi sub_dataset_tight.view()); read_end = std::chrono::high_resolution_clock::now(); auto sub_dataset_view = raft::make_const_mdspan(sub_dataset_tight.view()); - std::unique_ptr> - sub_dataset_padded; + std::unique_ptr> sub_dataset_padded; auto sub_dataset_dev = [&]() { - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(sub_dataset_view)) { - return cuvs::neighbors::make_device_padded_dataset_view(res, sub_dataset_view); + if (cuvs::core::matrix_has_padded_row_width(sub_dataset_view)) { + return cuvs::core::make_device_padded_dataset_view(res, sub_dataset_view); } - sub_dataset_padded = - cuvs::neighbors::make_device_padded_dataset(res, sub_dataset_view); + sub_dataset_padded = cuvs::core::make_device_padded_dataset(res, sub_dataset_view); return sub_dataset_padded->as_dataset_view(); }(); auto direct_index = @@ -1650,7 +1648,7 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi augmented_header_size, sub_dataset.view()); read_end = std::chrono::high_resolution_clock::now(); - auto sub_dataset_view = cuvs::neighbors::make_host_standard_dataset_view( + auto sub_dataset_view = cuvs::core::make_host_standard_dataset_view( raft::make_const_mdspan(sub_dataset.view())); auto host_index = ::cuvs::neighbors::cagra::build(res, sub_index_params, sub_dataset_view); @@ -1669,9 +1667,9 @@ auto build_ace(raft::resources const& res, const index_params& params, DatasetVi core_partition_offsets.view(), augmented_partition_offsets.view(), sub_dataset.view()); - read_end = std::chrono::high_resolution_clock::now(); - auto sub_dataset_view = cuvs::neighbors::make_host_standard_dataset_view( - raft::make_const_mdspan(sub_dataset.view())); + read_end = std::chrono::high_resolution_clock::now(); + auto sub_dataset_view = + cuvs::core::make_host_standard_dataset_view(raft::make_const_mdspan(sub_dataset.view())); auto host_index = ::cuvs::neighbors::cagra::build(res, sub_index_params, sub_dataset_view); static_assert(std::is_same_v); return sub_index_t{std::in_place_type, std::move(host_index)}; @@ -2213,10 +2211,11 @@ void build_knn_graph( } template -void build_knn_graph(raft::resources const& res, - cuvs::neighbors::device_bbq_dataset_view dataset, - raft::host_matrix_view knn_graph, - cuvs::neighbors::nn_descent::index_params build_params) +void build_knn_graph( + raft::resources const& res, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, + raft::host_matrix_view knn_graph, + cuvs::neighbors::nn_descent::index_params build_params) { raft::common::nvtx::range fun_scope( "cagra::build_knn_graph(%zu, %zu, %u)", @@ -2306,15 +2305,16 @@ __global__ void kern_reconstruct_vpq_queries(const uint8_t* encoded_data, } template -void reconstruct_vpq_queries(raft::resources const& res, - cuvs::neighbors::device_vpq_dataset_view const& vpq_view, - uint64_t offset, - uint32_t batch_size, - raft::device_matrix_view output) +void reconstruct_vpq_queries( + raft::resources const& res, + cuvs::preprocessing::quantize::pq::device_vpq_dataset_view const& dataset_view, + uint64_t offset, + uint32_t batch_size, + raft::device_matrix_view output) { - auto const& vpq_dset = vpq_view.dset(); - const uint32_t dim = vpq_dset.dim(); - const uint32_t pq_len = vpq_dset.pq_len(); + auto const& vpq_view = dataset_view.data(); + const uint32_t dim = dataset_view.dim(); + const uint32_t pq_len = vpq_view.pq_len(); const uint32_t output_ld = static_cast(output.extent(1)); const uint32_t threads = std::min(dim, 256u); RAFT_EXPECTS(output_ld >= dim, @@ -2324,10 +2324,10 @@ void reconstruct_vpq_queries(raft::resources const& res, kern_reconstruct_vpq_queries <<>>( - vpq_dset.data.data_handle(), - vpq_dset.encoded_row_length(), - vpq_dset.vq_code_book.data_handle(), - vpq_dset.pq_code_book.data_handle(), + vpq_view.data_handle(), + vpq_view.encoded_row_length(), + vpq_view.vq_code_book.data_handle(), + vpq_view.pq_code_book.data_handle(), dim, pq_len, offset, @@ -2365,9 +2365,9 @@ auto search_and_optimize( size_t next_graph_degree, uint64_t max_chunk_size, bool guarantee_connectivity, - std::optional> reconstructed_batch_queries = std::nullopt, - std::optional> vpq_queries = std::nullopt) - -> raft::device_matrix + std::optional> reconstructed_batch_queries = std::nullopt, + std::optional> + vpq_queries = std::nullopt) -> raft::device_matrix { auto stream = raft::resource::get_cuda_stream(res); auto const curr_query_size = knn_graph.extent(0); @@ -2449,8 +2449,8 @@ auto search_and_optimize( } template - requires(cuvs::neighbors::is_dense_row_major_device_dataset_view_v || - cuvs::neighbors::is_device_vpq_f16_dataset_view_v) + requires(cuvs::core::is_dense_row_major_device_dataset_view_v || + cuvs::preprocessing::quantize::pq::is_device_vpq_f16_dataset_view_v) auto iterative_build_graph(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) -> raft::device_matrix @@ -2473,21 +2473,22 @@ auto iterative_build_graph(raft::resources const& res, raft::make_device_matrix_view(static_cast(nullptr), 0, 0); uint32_t logical_dim = dataset.dim(); uint64_t final_graph_size; - auto vpq_dataset = cuvs::neighbors::device_vpq_dataset_view{}; + auto vpq_dataset = cuvs::preprocessing::quantize::pq::device_vpq_dataset_view{}; - if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + if constexpr (cuvs::preprocessing::quantize::pq::is_device_vpq_f16_dataset_view_v) { final_graph_size = static_cast(dataset.n_rows()); vpq_dataset = dataset; } else { - auto const required_stride = cuvs::neighbors::cagra_required_row_width(dataset.dim()); - RAFT_EXPECTS(dataset.stride() == required_stride, + auto const required_stride = cuvs::core::padded_row_width(dataset.dim()); + auto const data_view = dataset.as_matrix_view(); + RAFT_EXPECTS(data_view.stride() == required_stride, "iterative CAGRA build requires a CAGRA-aligned device dataset " "(stride %u, required %u). Pass a device_padded_dataset_view, or a " "device_standard_dataset_view whose row width already matches " - "cagra_required_row_width.", - dataset.stride(), + "padded_row_width.", + data_view.stride(), required_stride); - dev_dataset = dataset.view(); + dev_dataset = data_view; logical_dim = dataset.dim(); final_graph_size = static_cast(dataset.n_rows()); } @@ -2519,8 +2520,8 @@ auto iterative_build_graph(raft::resources const& res, std::optional> reconstructed_batch_queries; if (vpq_dataset.n_rows() > 0) { - auto const query_stride_i64 = static_cast( - cuvs::neighbors::cagra_required_row_width(static_cast(logical_dim))); + auto const query_stride_i64 = + static_cast(cuvs::core::padded_row_width(static_cast(logical_dim))); reconstructed_batch_queries.emplace( raft::make_device_matrix(res, chunk_i64, query_stride_i64)); // Padding columns must be zero: search_main cosine post-process reduces over the full row @@ -2651,8 +2652,7 @@ auto iterative_build_graph(raft::resources const& res, } else { auto dev_dataset_view = raft::make_device_matrix_view( dev_dataset.data_handle(), static_cast(curr_graph_size), dev_dataset.extent(1)); - cuvs::neighbors::device_padded_dataset_view sub_padded(dev_dataset_view, - logical_dim); + cuvs::core::device_padded_dataset_view sub_padded(dev_dataset_view, logical_dim); auto idx = cuvs::neighbors::cagra::update_dataset( res, cuvs::neighbors::cagra::device_padded_index(res, params.metric), sub_padded); idx.update_graph(res, raft::make_const_mdspan(dev_graph.view())); @@ -2804,7 +2804,7 @@ auto build_cagra_host_graph_from_knn_params(raft::resources const& res, * it still requires a device dataset before search. */ template - requires cuvs::neighbors::is_host_dataset_view_v + requires cuvs::core::is_host_dataset_view_v auto build_from_host_matrix(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) @@ -2841,7 +2841,7 @@ auto build_from_host_matrix(raft::resources const& res, static_cast(n_rows), intermediate_degree, graph_degree, - dataset.view()); + dataset.as_matrix_view()); }(); RAFT_LOG_TRACE("Graph optimized, creating index"); @@ -2862,7 +2862,7 @@ auto build_from_host_matrix(raft::resources const& res, * `cagra::update_dataset` before search. */ template - requires cuvs::neighbors::is_dense_row_major_device_dataset_view_v + requires cuvs::core::is_dense_row_major_device_dataset_view_v auto build_from_device_matrix(raft::resources const& res, const index_params& params, DatasetViewT const& device_dataset) @@ -2888,13 +2888,14 @@ auto build_from_device_matrix(raft::resources const& res, auto cagra_graph = iterative_build_graph(res, params, device_dataset); idx.update_graph(res, std::move(cagra_graph)); } else { - auto cagra_graph = build_cagra_host_graph_from_knn_params(res, - params, - knn_build_params, - device_dataset.n_rows(), - intermediate_degree, - graph_degree, - device_dataset.view()); + auto cagra_graph = + build_cagra_host_graph_from_knn_params(res, + params, + knn_build_params, + device_dataset.n_rows(), + intermediate_degree, + graph_degree, + device_dataset.as_matrix_view()); idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view())); } @@ -2941,13 +2942,13 @@ auto build_from_device_matrix(raft::resources const& res, * the same graph. */ template - requires cuvs::neighbors::is_device_bbq_dataset_view_v + requires cuvs::preprocessing::quantize::bbq::is_device_bbq_dataset_view_v auto build_from_bbq_dataset(raft::resources const& res, const index_params& params, DatasetViewT const& dataset) -> cuvs::neighbors::cagra::index { - RAFT_EXPECTS(!dataset.quantizers.empty(), "cagra::build: the BBQ dataset is empty."); + RAFT_EXPECTS(!dataset.data().quantizers.empty(), "cagra::build: the BBQ dataset is empty."); RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded || params.metric == cuvs::distance::DistanceType::L2SqrtExpanded || params.metric == cuvs::distance::DistanceType::CosineExpanded || diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index cf15999842..7b8feea1e4 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -37,7 +37,7 @@ namespace cuvs::neighbors::cagra::detail { -template +template int64_t merged_dataset_size( raft::resources const& handle, std::vector*> const& indices, @@ -59,7 +59,7 @@ int64_t merged_dataset_size( return merged_rows; } -template +template cuvs::neighbors::cagra::index merge_rebuild( raft::resources const& handle, const cagra::index_params& params, @@ -84,8 +84,7 @@ cuvs::neighbors::cagra::index merge_rebuild( RAFT_EXPECTS(index != nullptr, "Null pointer detected in 'indices'. Ensure all elements are valid before usage."); auto const& dataset = index->dataset(); - if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v< - std::decay_t>) { + if constexpr (cuvs::core::is_dense_row_major_dataset_view_v>) { RAFT_EXPECTS( dataset.n_rows() != 0, "cagra::merge only supports an index to which the dataset is attached. Please check if " @@ -93,10 +92,10 @@ cuvs::neighbors::cagra::index merge_rebuild( "before merge."); if (dim == 0) { dim = index->dim(); - stride = static_cast(dataset.stride()); + stride = static_cast(dataset.as_matrix_view().stride()); } else { RAFT_EXPECTS(dim == index->dim(), "Dimension of datasets in indices must be equal."); - RAFT_EXPECTS(stride == static_cast(dataset.stride()), + RAFT_EXPECTS(stride == static_cast(dataset.as_matrix_view().stride()), "Row stride of datasets in indices must be equal."); } merged_rows += static_cast(index->size()); @@ -118,12 +117,12 @@ cuvs::neighbors::cagra::index merge_rebuild( "merged_dataset dimension (%u) must equal the input dimension (%u)", unsigned(merged_dataset.dim()), unsigned(dim)); - RAFT_EXPECTS(merged_dataset.stride() == stride, + RAFT_EXPECTS(merged_dataset.as_matrix_view().stride() == stride, "merged_dataset stride (%u) must equal the input stride (%ld)", - unsigned(merged_dataset.stride()), + unsigned(merged_dataset.as_matrix_view().stride()), long(stride)); - auto output_const_view = merged_dataset.view(); + auto output_const_view = merged_dataset.as_matrix_view(); auto output_view = raft::make_device_matrix_view( const_cast(output_const_view.data_handle()), final_rows, stride); @@ -133,8 +132,8 @@ cuvs::neighbors::cagra::index merge_rebuild( const T* src_ptr = nullptr; std::size_t n_rows = 0; auto const& v = index->dataset(); - if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v>) { - src_ptr = v.view().data_handle(); + if constexpr (cuvs::core::is_dense_row_major_dataset_view_v>) { + src_ptr = v.as_matrix_view().data_handle(); n_rows = static_cast(v.n_rows()); } else { RAFT_FAIL("cagra::merge: unexpected dataset type while copying rows"); @@ -215,7 +214,7 @@ struct fastener_preflight_result { }; /** Validate every input and option without mutating anything. */ -template +template auto preflight_fastener( raft::resources const& handle, cagra::index_params const& params, @@ -236,7 +235,7 @@ auto preflight_fastener( } // Fastener reads the dataset densely per row with an explicit stride, so it needs a dense // device view; VPQ and host views are rejected here rather than deep inside a kernel. - if constexpr (!cuvs::neighbors::is_dense_row_major_device_dataset_view_v) { + if constexpr (!cuvs::core::is_dense_row_major_device_dataset_view_v) { return reject("only dense row-major device datasets are supported"); } if (indices.size() < 2) { return reject("at least two input indices are required"); } @@ -302,14 +301,14 @@ auto preflight_fastener( } if (result.offsets.size() == 1) { result.dim = static_cast(index->dim()); - result.stride = static_cast(dataset.stride()); + result.stride = static_cast(dataset.as_matrix_view().stride()); } else { if (result.dim != static_cast(index->dim())) { return reject("all input dimensions must match"); } // The merged dataset has a single row pitch, so mixed input strides cannot be consolidated // without re-padding each input separately. - if (result.stride != static_cast(dataset.stride())) { + if (result.stride != static_cast(dataset.as_matrix_view().stride())) { return reject("all input row strides must match"); } } @@ -371,7 +370,7 @@ auto preflight_fastener( /** Copy every input dataset into its row range of the caller-supplied merged dataset. Both sides * carry a row pitch: the inputs share one stride (enforced by preflight) and the destination uses * the merged dataset's own stride. */ -template +template void copy_input_datasets( raft::resources const& handle, std::vector*> const& indices, @@ -384,15 +383,15 @@ void copy_input_datasets( auto const& source = indices[i]->dataset(); raft::copy_matrix(destination + offsets[i] * destination_stride, static_cast(destination_stride), - source.view().data_handle(), - static_cast(source.stride()), + source.as_matrix_view().data_handle(), + static_cast(source.as_matrix_view().stride()), static_cast(dim), static_cast(source.n_rows()), raft::resource::get_cuda_stream(handle)); } } -template +template auto merge_fastener(raft::resources const& handle, cagra::index_params const& params, cagra::merge_params const& merge_params, @@ -401,7 +400,7 @@ auto merge_fastener(raft::resources const& handle, fastener_preflight_result const& preflight) -> cuvs::neighbors::cagra::index { - auto const stride = static_cast(merged_dataset.stride()); + auto const stride = static_cast(merged_dataset.as_matrix_view().stride()); RAFT_EXPECTS(merged_dataset.n_rows() == preflight.rows, "merged_dataset rows (%ld) must equal the merged row count (%ld)", long(merged_dataset.n_rows()), @@ -411,7 +410,7 @@ auto merge_fastener(raft::resources const& handle, unsigned(merged_dataset.dim()), long(preflight.dim)); - auto const output_const_view = merged_dataset.view(); + auto const output_const_view = merged_dataset.as_matrix_view(); auto* destination = const_cast(output_const_view.data_handle()); { raft::common::nvtx::range scope("cagra::merge/consolidate"); @@ -487,7 +486,7 @@ auto merge_fastener(raft::resources const& handle, return merged_index; } -template +template auto merge(raft::resources const& handle, cagra::index_params const& params, std::vector*>& indices, @@ -540,7 +539,7 @@ auto merge(raft::resources const& handle, } /** AUTO-algorithm convenience overload matching the base `merge` signature. */ -template +template auto merge(raft::resources const& handle, cagra::index_params const& params, std::vector*>& indices, diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh index 897991d2d5..9c22cacf1e 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh @@ -1142,7 +1142,7 @@ auto build(raft::resources const& res, * * The maximum input graph degree defines the base width, allowing partitions with mixed degrees. */ -template +template void append_to_input_graphs( raft::resources const& res, std::vector*> const& indices, diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index e9902d049e..f52280233c 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -100,16 +100,16 @@ void search_main_core( // batched path when stride==dim). const DataT* queries_buf{}; uint32_t query_row_stride{}; - std::unique_ptr> queries_padded_own; - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(queries)) { - auto v = cuvs::neighbors::make_device_padded_dataset_view(res, queries); - queries_buf = v.view().data_handle(); - query_row_stride = v.stride(); + std::unique_ptr> queries_padded_own; + if (cuvs::core::matrix_has_padded_row_width(queries)) { + auto v = cuvs::core::make_device_padded_dataset_view(res, queries); + queries_buf = v.as_matrix_view().data_handle(); + query_row_stride = v.as_matrix_view().stride(); } else { - queries_padded_own = cuvs::neighbors::make_device_padded_dataset(res, queries); + queries_padded_own = cuvs::core::make_device_padded_dataset(res, queries); auto v = queries_padded_own->as_dataset_view(); - queries_buf = v.view().data_handle(); - query_row_stride = v.stride(); + queries_buf = v.as_matrix_view().data_handle(); + query_row_stride = v.as_matrix_view().stride(); } const bool can_batch_n_queries = (query_row_stride == query_logical_dim); @@ -188,7 +188,7 @@ template + cuvs::core::dataset_like DatasetViewT> void search_main(raft::resources const& res, search_params params, const index& index, @@ -206,7 +206,7 @@ void search_main(raft::resources const& res, const uint32_t query_logical_dim = index.dim(); const uint32_t query_row_width = static_cast(queries.extent(1)); - const uint32_t required_stride = cuvs::neighbors::cagra_required_row_width(query_logical_dim); + const uint32_t required_stride = cuvs::core::padded_row_width(query_logical_dim); RAFT_EXPECTS(query_row_width == query_logical_dim || query_row_width == required_stride, "CAGRA search queries must have %u logical dimensions or CAGRA-padded row width %u " "(got %u).", @@ -243,11 +243,12 @@ void search_main(raft::resources const& res, sample_filter); }; - if constexpr (cuvs::neighbors::is_empty_dataset_view_v) { + if constexpr (cuvs::core::is_empty_dataset_view_v) { RAFT_FAIL( "Attempted to search without a dataset. Please call " "cagra::update_dataset(res, std::move(index), dataset) first."); - } else if constexpr (cuvs::neighbors::is_device_vpq_f16_dataset_view_v) { + } else if constexpr (cuvs::preprocessing::quantize::pq::is_device_vpq_f16_dataset_view_v< + DatasetViewT>) { auto const& vv = index.dataset(); if (params.smem_dtype == cuvs::neighbors::cagra::internal_dtype::E5M2 && raft::getComputeCapability().first < 9) { @@ -256,7 +257,7 @@ void search_main(raft::resources const& res, params.smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; } auto desc = dataset_descriptor_init_with_cache( - res, params, vv.dset(), index.metric(), nullptr); + res, params, vv, index.metric(), nullptr); search_main_core( res, params, @@ -268,15 +269,16 @@ void search_main(raft::resources const& res, distances, query_logical_dim, sample_filter); - } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { + } else if constexpr (cuvs::core::is_device_standard_dataset_view_v) { RAFT_EXPECTS( - cuvs::neighbors::matrix_row_width_matches_cagra_required(index.dataset().view()), + cuvs::core::matrix_has_padded_row_width(index.dataset().as_matrix_view()), "CAGRA search requires each dataset row to have the CAGRA-aligned stride. Create a padded " - "dataset with make_device_padded_dataset() and attach it with cagra::update_dataset()."); + "dataset with cuvs::core::make_device_padded_dataset() and attach it with " + "cagra::update_dataset()."); run_strided_like(index.dataset()); - } else if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + } else if constexpr (cuvs::core::is_device_padded_dataset_view_v) { run_strided_like(index.dataset()); - } else if constexpr (cuvs::neighbors::is_host_dataset_view_v) { + } else if constexpr (cuvs::core::is_host_dataset_view_v) { static_assert(sizeof(DatasetViewT) == 0, "search requires a device-resident dataset. " "Call cagra::update_dataset(res, std::move(index), padded_view) " @@ -370,7 +372,7 @@ void search_multi_partition( // The index type in this signature pins the dataset view to the default, so every partition is // statically known to hold a padded (non-compressed) device dataset. using partition_dataset_view_t = std::remove_cvref_tdataset())>; - static_assert(cuvs::neighbors::is_device_padded_dataset_view_v, + static_assert(cuvs::core::is_device_padded_dataset_view_v, "Multi-partition search requires padded device datasets"); const uint32_t num_partitions = static_cast(indices.size()); diff --git a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh index 94ea20c1ad..d09c9cc2e6 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_serialize.cuh @@ -53,17 +53,17 @@ inline constexpr bool is_host_cagra_hnsw_serialize_index_v = constexpr int serialization_version = cuvs::neighbors::cagra::cagra_serialization_version; -template +template constexpr auto serialized_dataset_kind_for_view() -> cuvs::neighbors::cagra::serialized_dataset_kind { using kind = cuvs::neighbors::cagra::serialized_dataset_kind; - if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + if constexpr (cuvs::core::is_device_padded_dataset_view_v) { return kind::device_padded; - } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { + } else if constexpr (cuvs::core::is_device_standard_dataset_view_v) { return kind::device_standard; - } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v) { + } else if constexpr (cuvs::core::is_host_padded_dataset_view_v) { return kind::host_padded; - } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { + } else if constexpr (cuvs::core::is_host_standard_dataset_view_v) { return kind::host_standard; } else { static_assert(sizeof(DatasetViewT) == 0, @@ -83,8 +83,8 @@ constexpr bool is_valid_serialized_dataset_kind(std::uint32_t raw) */ template inline constexpr bool is_graph_only_dataset_view_v = - cuvs::neighbors::is_vpq_dataset_view_v || - cuvs::neighbors::is_bbq_dataset_view_v; + cuvs::preprocessing::quantize::pq::is_vpq_dataset_view_v || + cuvs::preprocessing::quantize::bbq::is_bbq_dataset_view_v; inline constexpr char const* kGraphOnlyDatasetMessage = "CAGRA indexes with a quantized dataset store only the graph; serialize the quantized dataset " @@ -106,7 +106,7 @@ void serialize_index_mdspan(raft::resources const& res, std::ostream& os, Mdspan * @param[in] index_ CAGRA index * */ -template +template void serialize(raft::resources const& res, std::ostream& os, const cuvs::neighbors::cagra::index& index_, @@ -149,7 +149,7 @@ void serialize(raft::resources const& res, raft::serialize_scalar(res, os, content_map); if (include_dataset) { RAFT_LOG_DEBUG("Saving CAGRA index with dataset"); - if constexpr (cuvs::neighbors::is_dense_row_major_dataset_view_v) { + if constexpr (cuvs::core::is_dense_row_major_dataset_view_v) { neighbors::detail::serialize_cagra_dense_dataset(res, os, index_.dataset()); } else if constexpr (is_graph_only_dataset_view_v) { RAFT_FAIL(kGraphOnlyDatasetMessage); @@ -167,7 +167,7 @@ void serialize(raft::resources const& res, if (has_source_indices) { serialize_index_mdspan(res, os, index_.source_indices().value()); } } -template +template void serialize(raft::resources const& res, const std::string& filename, const cuvs::neighbors::cagra::index& index_, @@ -312,8 +312,8 @@ void write_hnswlib_rows_host( "CAGRA dataset rows (%zu) do not match index size (%zu)", static_cast(dataset_view.n_rows()), n_rows); - dataset_data = dataset_view.view().data_handle(); - dataset_stride = dataset_view.stride(); + dataset_data = dataset_view.as_matrix_view().data_handle(); + dataset_stride = dataset_view.as_matrix_view().stride(); dataset_is_device = is_device_cagra_hnsw_serialize_index_v; } @@ -450,12 +450,12 @@ void write_hnswlib_rows_device(raft::resources const& res, output.data_handle(), row_size, graph.data_handle(), - dataset.view().data_handle(), + dataset.as_matrix_view().data_handle(), first_row, rows, graph_degree, dim, - dataset.stride()); + dataset.as_matrix_view().stride()); RAFT_CUDA_TRY(cudaPeekAtLastError()); raft::resource::sync_stream(res); @@ -540,12 +540,12 @@ void serialize_to_hnswlib( * @param[in] index_ CAGRA index * */ -template +template void deserialize_impl( raft::resources const& res, Input& input, cuvs::neighbors::cagra::index* index_, - std::unique_ptr>* out_dataset = nullptr) + std::unique_ptr>* out_dataset = nullptr) { raft::common::nvtx::range fun_scope("cagra::deserialize"); auto& is = cuvs::util::detail::input_stream(input); @@ -595,7 +595,7 @@ void deserialize_impl( RAFT_EXPECTS(has_dataset == (dataset_kind != kind::none), "cagra::deserialize: dataset kind and content map disagree"); - using owner_t = cuvs::neighbors::owning_dataset_for_view_t; + using owner_t = cuvs::core::owning_dataset_for_view_t; std::unique_ptr dataset_owner{}; if (has_dataset) { if (out_dataset == nullptr) { @@ -614,16 +614,16 @@ void deserialize_impl( "cagra::deserialize: serialized dataset kind %u does not match requested kind %u", dataset_kind_raw, static_cast(expected_kind)); - if constexpr (cuvs::neighbors::is_device_padded_dataset_view_v) { + if constexpr (cuvs::core::is_device_padded_dataset_view_v) { dataset_owner = cuvs::neighbors::detail::deserialize_padded_dataset(res, input); - } else if constexpr (cuvs::neighbors::is_device_standard_dataset_view_v) { + } else if constexpr (cuvs::core::is_device_standard_dataset_view_v) { dataset_owner = cuvs::neighbors::detail::deserialize_standard_dataset(res, input); - } else if constexpr (cuvs::neighbors::is_host_padded_dataset_view_v) { + } else if constexpr (cuvs::core::is_host_padded_dataset_view_v) { dataset_owner = cuvs::neighbors::detail::deserialize_host_padded_dataset(res, input); - } else if constexpr (cuvs::neighbors::is_host_standard_dataset_view_v) { + } else if constexpr (cuvs::core::is_host_standard_dataset_view_v) { dataset_owner = cuvs::neighbors::detail::deserialize_host_standard_dataset(res, input); } else { @@ -678,22 +678,22 @@ void deserialize_impl( } } -template +template void deserialize( raft::resources const& res, std::istream& is, cuvs::neighbors::cagra::index* index_, - std::unique_ptr>* out_dataset = nullptr) + std::unique_ptr>* out_dataset = nullptr) { deserialize_impl(res, is, index_, out_dataset); } -template +template void deserialize( raft::resources const& res, const std::string& filename, cuvs::neighbors::cagra::index* index_, - std::unique_ptr>* out_dataset = nullptr) + std::unique_ptr>* out_dataset = nullptr) { cuvs::util::kvikio_file_reader reader(filename); deserialize_impl(res, reader, index_, out_dataset); diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp b/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp index 5f7c8efb10..fba16934fc 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance_standard.hpp @@ -28,7 +28,7 @@ struct standard_descriptor_spec : public instance_spec template constexpr static inline bool accepts_dataset() { - return is_padded_dataset_v; + return cuvs::core::is_padded_dataset_view_v; } template @@ -37,11 +37,12 @@ struct standard_descriptor_spec : public instance_spec cuvs::distance::DistanceType metric, const DistanceT* dataset_norms = nullptr) -> host_type { + auto const data_view = dataset.as_matrix_view(); return init_(params, - dataset.view().data_handle(), + data_view.data_handle(), IndexT(dataset.n_rows()), dataset.dim(), - dataset.stride(), + data_view.stride(), dataset_norms); } diff --git a/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp b/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp index 7ae9dc87ba..569fd7b24e 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance_vpq.hpp @@ -41,46 +41,52 @@ struct vpq_descriptor_spec : public instance_spec { using typename base_type::host_type; using typename base_type::index_type; + // `DatasetT` here is the non-owning dataset_view (not the owning dataset) -- callers pass the + // view directly rather than reaching back through a `.dset()`-style owner pointer, so + // classification and the codebook element type are read off the view via its own VPQ payload + // (`is_vpq_dataset_view_v`, `DatasetT::value_type`), not an owning-only trait/typedef. template constexpr static inline auto accepts_dataset() - -> std::enable_if_t, bool> + -> std::enable_if_t, bool> { - return std::is_same_v; + return std::is_same_v; } template constexpr static inline auto accepts_dataset() - -> std::enable_if_t, bool> + -> std::enable_if_t, bool> { return false; } template static auto init(const cagra::search_params& params, - const DatasetT& dataset, + const DatasetT& dataset_view, cuvs::distance::DistanceType metric, const DistanceT* dataset_norms = nullptr) -> host_type { + auto const& vpq_view = dataset_view.data(); return init_(params, - dataset.data.data_handle(), - dataset.encoded_row_length(), - dataset.vq_code_book.data_handle(), - dataset.pq_code_book.data_handle(), - IndexT(dataset.n_rows()), - dataset.dim()); + vpq_view.data_handle(), + static_cast(vpq_view.extent(1)), + vpq_view.vq_code_book.data_handle(), + vpq_view.pq_code_book.data_handle(), + IndexT(dataset_view.n_rows()), + dataset_view.dim()); } template static auto priority(const cagra::search_params& params, - const DatasetT& dataset, + const DatasetT& dataset_view, cuvs::distance::DistanceType metric) -> double { // If explicit team_size is specified and doesn't match the instance, discard it if (params.team_size != 0 && TeamSize != params.team_size) { return -1.0; } if (cuvs::distance::DistanceType::L2Expanded != metric) { return -1.0; } // Match codebook params - if (dataset.pq_bits() != PqBits) { return -1.0; } - if (dataset.pq_len() != PqLen) { return -1.0; } + auto const& vpq_view = dataset_view.data(); + if (vpq_view.pq_bits() != PqBits) { return -1.0; } + if (vpq_view.pq_len() != PqLen) { return -1.0; } if (select_supported_vpq_smem_dtype(params) != SmemDType) { return -1.0; } // Keep auto-selection on the tuned VPQ diagonal while allowing explicit team_size requests to // use the expanded team_size / dataset_block_dim grid. @@ -91,8 +97,8 @@ struct vpq_descriptor_spec : public instance_spec { // Otherwise, favor the closest dataset dimensionality. constexpr std::uint32_t preferred_load_elmes_per_thread = 16; /*magic number that is good based on experiments.*/ - return 1.0 / (0.1 + std::abs(double(dataset.dim()) - double(DatasetBlockDim))) * TeamSize + - 1.0 / (0.1 + std::abs(double(dataset.dim()) / TeamSize / PqLen - + return 1.0 / (0.1 + std::abs(double(dataset_view.dim()) - double(DatasetBlockDim))) * TeamSize + + 1.0 / (0.1 + std::abs(double(dataset_view.dim()) / TeamSize / PqLen - preferred_load_elmes_per_thread)); } diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index ace28b7842..444efcd88a 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -91,18 +91,23 @@ struct key { uint32_t smem_dtype; }; +// `DatasetT` here is the non-owning dataset_view passed in by the search path, so all state comes +// off the view's own `as_matrix_view()`/`data()`, not owning-only members. template auto make_key(const cagra::search_params& params, - const DatasetT& dataset, + const DatasetT& dataset_view, cuvs::distance::DistanceType metric, const void* dataset_norms) - -> std::enable_if_t || is_standard_dataset_v, key> + -> std::enable_if_t || + cuvs::core::is_standard_dataset_v, + key> { - return key{reinterpret_cast(dataset.view().data_handle()), + auto const data_view = dataset_view.as_matrix_view(); + return key{reinterpret_cast(data_view.data_handle()), reinterpret_cast(dataset_norms), - uint64_t(dataset.n_rows()), - dataset.dim(), - dataset.stride(), + uint64_t(dataset_view.n_rows()), + dataset_view.dim(), + data_view.stride(), uint32_t(params.team_size), uint32_t(metric), uint32_t(params.smem_dtype)}; @@ -110,15 +115,17 @@ auto make_key(const cagra::search_params& params, template auto make_key(const cagra::search_params& params, - const DatasetT& dataset, + const DatasetT& dataset_view, cuvs::distance::DistanceType metric, - const void* dataset_norms) -> std::enable_if_t, key> + const void* dataset_norms) + -> std::enable_if_t, key> { - return key{reinterpret_cast(dataset.data.data_handle()), + auto const& vpq_view = dataset_view.data(); + return key{reinterpret_cast(vpq_view.data_handle()), reinterpret_cast(dataset_norms), - uint64_t(dataset.n_rows()), - dataset.dim(), - uint32_t(reinterpret_cast(dataset.pq_code_book.data_handle()) >> 6), + uint64_t(dataset_view.n_rows()), + dataset_view.dim(), + uint32_t(reinterpret_cast(vpq_view.pq_code_book.data_handle()) >> 6), uint32_t(params.team_size), uint32_t(metric), uint32_t(params.smem_dtype)}; diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cu b/cpp/src/neighbors/detail/cagra/graph_shared.cu index 8dcb4d8b6f..8965966aa1 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cu +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cu @@ -252,17 +252,20 @@ auto select_sort_bbq_kernel(uint32_t degree) -> sort_bbq_kernel_type } template -auto select_sort_quantizers(cuvs::neighbors::device_bbq_dataset_view const& dataset) +auto select_sort_quantizers( + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) -> std::pair, bbq_quantizer_view_t> { using bbq_code_layout = cuvs::preprocessing::quantize::bbq::bbq_code_layout; - if (dataset.quantizers.size() == 1) { return {dataset.quantizers[0], dataset.quantizers[0]}; } + if (dataset.data().quantizers.size() == 1) { + return {dataset.data().quantizers[0], dataset.data().quantizers[0]}; + } - const bool has_1b = dataset.has_layout(bbq_code_layout::packed_1b); - const bool has_4b = dataset.has_layout(bbq_code_layout::packed_4b); - const bool has_2bt = dataset.has_layout(bbq_code_layout::transposed_2b); - const bool has_4bt = dataset.has_layout(bbq_code_layout::transposed_4b); + const bool has_1b = dataset.data().has_layout(bbq_code_layout::packed_1b); + const bool has_4b = dataset.data().has_layout(bbq_code_layout::packed_4b); + const bool has_2bt = dataset.data().has_layout(bbq_code_layout::transposed_2b); + const bool has_4bt = dataset.data().has_layout(bbq_code_layout::transposed_4b); const bool tc_pair = has_1b && has_4b; const bool simt_pair = (has_4bt && (has_1b || has_2bt)) || (has_2bt && has_1b); @@ -270,22 +273,23 @@ auto select_sort_quantizers(cuvs::neighbors::device_bbq_dataset_view -void sort_knn_graph_bbq_impl(raft::resources const& res, - cuvs::distance::DistanceType metric, - cuvs::neighbors::device_bbq_dataset_view dataset, - raft::host_matrix_view knn_graph) +void sort_knn_graph_bbq_impl( + raft::resources const& res, + cuvs::distance::DistanceType metric, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, + raft::host_matrix_view knn_graph) { namespace bbq = cuvs::preprocessing::quantize::bbq; - RAFT_EXPECTS(!dataset.quantizers.empty(), "the BBQ dataset holds no quantizer"); + RAFT_EXPECTS(!dataset.data().quantizers.empty(), "the BBQ dataset holds no quantizer"); RAFT_EXPECTS(dataset.n_rows() == knn_graph.extent(0), "dataset size is expected to have the same number of graph index size"); RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Expanded || @@ -367,13 +371,14 @@ CUVS_DEFINE_CAGRA_GRAPH_SORT(uint8_t) #undef CUVS_DEFINE_CAGRA_GRAPH_SORT -#define CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(DataT) \ - void sort_knn_graph_bbq(raft::resources const& res, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::device_bbq_dataset_view dataset, \ - raft::host_matrix_view knn_graph) \ - { \ - sort_knn_graph_bbq_impl(res, metric, dataset, knn_graph); \ +#define CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(DataT) \ + void sort_knn_graph_bbq( \ + raft::resources const& res, \ + cuvs::distance::DistanceType metric, \ + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, \ + raft::host_matrix_view knn_graph) \ + { \ + sort_knn_graph_bbq_impl(res, metric, dataset, knn_graph); \ } CUVS_DEFINE_CAGRA_GRAPH_SORT_BBQ(float) diff --git a/cpp/src/neighbors/detail/cagra/graph_shared.cuh b/cpp/src/neighbors/detail/cagra/graph_shared.cuh index 33ec7f92ed..3cb367872a 100644 --- a/cpp/src/neighbors/detail/cagra/graph_shared.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_shared.cuh @@ -37,11 +37,11 @@ CUVS_DECL_CAGRA_GRAPH_SORT(uint8_t); #undef CUVS_DECL_CAGRA_GRAPH_SORT -#define CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(DataT) \ - CUVS_EXPORT void sort_knn_graph_bbq( \ - raft::resources const& res, \ - cuvs::distance::DistanceType metric, \ - cuvs::neighbors::device_bbq_dataset_view dataset, \ +#define CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(DataT) \ + CUVS_EXPORT void sort_knn_graph_bbq( \ + raft::resources const& res, \ + cuvs::distance::DistanceType metric, \ + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, \ raft::host_matrix_view knn_graph) CUVS_DECL_CAGRA_GRAPH_SORT_BBQ(float); diff --git a/cpp/src/neighbors/detail/dataset_serialize.hpp b/cpp/src/neighbors/detail/dataset_serialize.hpp index 7d518171fc..0c811dd234 100644 --- a/cpp/src/neighbors/detail/dataset_serialize.hpp +++ b/cpp/src/neighbors/detail/dataset_serialize.hpp @@ -5,6 +5,7 @@ #pragma once #include +#include #include #include #include @@ -82,16 +83,17 @@ auto dense_matrix_elements(IdxT n_rows, uint32_t dim, char const* context) -> st } template - requires cuvs::neighbors::is_dense_row_major_dataset_view_v + requires cuvs::core::is_dense_row_major_dataset_view_v void serialize(const raft::resources& res, std::ostream& os, ViewT const& dataset) { - auto n_rows = dataset.n_rows(); - auto dim = dataset.dim(); - auto stride = dataset.stride(); + auto n_rows = dataset.n_rows(); + auto dim = dataset.dim(); + auto data_view = dataset.as_matrix_view(); + auto stride = data_view.stride(); raft::serialize_scalar(res, os, n_rows); raft::serialize_scalar(res, os, dim); raft::serialize_scalar(res, os, stride); - auto src = dataset.view(); + auto src = data_view; auto const elements = dense_matrix_elements(n_rows, dim, "serialize_dense_dataset"); raft::numpy_serializer::write_header(os, {raft::numpy_serializer::get_numpy_dtype(), @@ -100,7 +102,7 @@ void serialize(const raft::resources& res, std::ostream& os, ViewT const& datase static_cast(dim)}}); if (elements == 0) { return; } - if constexpr (cuvs::neighbors::is_device_dataset_view_v) { + if constexpr (cuvs::core::is_device_dataset_view_v) { if (auto* kvikio_stream = dynamic_cast(&os); kvikio_stream != nullptr) { auto const row_bytes = static_cast(dim) * sizeof(DataT); @@ -154,7 +156,7 @@ void serialize(const raft::resources& res, std::ostream& os, ViewT const& datase /** Write CAGRA index dataset blob (tag + element dtype + strided payload). */ template - requires cuvs::neighbors::is_dense_row_major_dataset_view_v + requires cuvs::core::is_dense_row_major_dataset_view_v void serialize_cagra_dense_dataset(const raft::resources& res, std::ostream& os, ViewT const& dataset) @@ -176,10 +178,10 @@ void serialize_cagra_dense_dataset(const raft::resources& res, template auto deserialize_empty(raft::resources const& res, std::istream& is) - -> std::unique_ptr> + -> std::unique_ptr> { auto suggested_dim = raft::deserialize_scalar(res, is); - return std::make_unique>(suggested_dim); + return std::make_unique>(suggested_dim); } /** Read and validate shared dense wire metadata and the tight `[n_rows x dim]` NumPy header. */ @@ -414,7 +416,7 @@ auto deserialize_host_dense(raft::resources const& res, std::istream& is) template auto deserialize_vpq(raft::resources const& res, std::istream& is) - -> std::unique_ptr> + -> std::unique_ptr> { auto n_rows = raft::deserialize_scalar(res, is); auto dim = raft::deserialize_scalar(res, is); @@ -434,8 +436,9 @@ auto deserialize_vpq(raft::resources const& res, std::istream& is) raft::deserialize_mdspan(res, is, pq_code_book.view()); raft::deserialize_mdspan(res, is, data.view()); - return std::make_unique>( - std::move(vq_code_book), std::move(pq_code_book), std::move(data)); + using owning_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; + return std::make_unique( + std::move(data), std::move(vq_code_book), std::move(pq_code_book)); } template @@ -456,13 +459,16 @@ auto deserialize_dense_dataset(raft::resources const& res, Input& input) "deserialize_dataset: serialized dtype (%d) does not match expected (%d)", static_cast(dtype), static_cast(expected_dtype)); - if constexpr (std::is_same_v>) { + if constexpr (std::is_same_v>) { return deserialize_device_dense(res, input); - } else if constexpr (std::is_same_v>) { + } else if constexpr (std::is_same_v>) { return deserialize_device_dense(res, input); - } else if constexpr (std::is_same_v>) { + } else if constexpr (std::is_same_v>) { return deserialize_host_dense(res, is); - } else if constexpr (std::is_same_v>) { + } else if constexpr (std::is_same_v>) { return deserialize_host_dense(res, is); } else { static_assert(!std::is_same_v, @@ -495,30 +501,34 @@ void skip_dense_dataset(raft::resources const& res, std::istream& is) // type-erased variant routing. template auto deserialize_padded_dataset(raft::resources const& res, Input& input) - -> std::unique_ptr> + -> std::unique_ptr> { - return deserialize_dense_dataset>(res, input); + return deserialize_dense_dataset>( + res, input); } template auto deserialize_standard_dataset(raft::resources const& res, Input& input) - -> std::unique_ptr> + -> std::unique_ptr> { - return deserialize_dense_dataset>(res, input); + return deserialize_dense_dataset>( + res, input); } template auto deserialize_host_padded_dataset(raft::resources const& res, Input& input) - -> std::unique_ptr> + -> std::unique_ptr> { - return deserialize_dense_dataset>(res, input); + return deserialize_dense_dataset>( + res, input); } template auto deserialize_host_standard_dataset(raft::resources const& res, Input& input) - -> std::unique_ptr> + -> std::unique_ptr> { - return deserialize_dense_dataset>(res, input); + return deserialize_dense_dataset>( + res, input); } } // namespace cuvs::neighbors::detail diff --git a/cpp/src/neighbors/detail/hnsw.hpp b/cpp/src/neighbors/detail/hnsw.hpp index 19c0afd820..4af370a15f 100644 --- a/cpp/src/neighbors/detail/hnsw.hpp +++ b/cpp/src/neighbors/detail/hnsw.hpp @@ -394,8 +394,8 @@ from_cagra(raft::resources const& res, host_dataset = raft::make_host_matrix(dataset_view.n_rows(), dataset_view.dim()); raft::copy_matrix(host_dataset.data_handle(), host_dataset.extent(1), - dataset_view.view().data_handle(), - dataset_view.stride(), + dataset_view.as_matrix_view().data_handle(), + dataset_view.as_matrix_view().stride(), host_dataset.extent(1), dataset_view.n_rows(), raft::resource::get_cuda_stream(res)); @@ -1839,12 +1839,13 @@ void serialize_to_hnswlib_from_inmem( source_stride = dim; } else if constexpr (is_host_cagra_hnsw_export_index_v) { RAFT_FAIL("serialize_to_hnswlib_from_inmem requires dataset for host CAGRA index"); - } else if (auto dataset_view = index_.dataset(); dataset_view.view().data_handle() != nullptr) { + } else if (auto dataset_view = index_.dataset(); + dataset_view.as_matrix_view().data_handle() != nullptr) { n_rows = dataset_view.n_rows(); dim = dataset_view.dim(); device_dataset = true; - source_dataset = dataset_view.view().data_handle(); - source_stride = dataset_view.stride(); + source_dataset = dataset_view.as_matrix_view().data_handle(); + source_stride = dataset_view.as_matrix_view().stride(); } else { RAFT_FAIL("serialize_to_hnswlib_from_inmem: No dataset provided"); } @@ -1964,12 +1965,12 @@ from_cagra(raft::resources const& res, } else if constexpr (is_host_cagra_hnsw_export_index_v) { RAFT_FAIL("hnsw::from_cagra requires dataset for host CAGRA index"); } else if (auto dataset_view = cagra_index.dataset(); - dataset_view.view().data_handle() != nullptr) { + dataset_view.as_matrix_view().data_handle() != nullptr) { n_rows = dataset_view.n_rows(); dim = dataset_view.dim(); device_copy = true; - source_dataset = dataset_view.view().data_handle(); - source_stride = dataset_view.stride(); + source_dataset = dataset_view.as_matrix_view().data_handle(); + source_stride = dataset_view.as_matrix_view().stride(); } else { RAFT_FAIL("hnsw::from_cagra: No dataset provided"); } @@ -3003,7 +3004,7 @@ std::unique_ptr> build(raft::resources const& res, // Public HNSW API uses host_matrix_view; CAGRA build expects a padded dataset view. // Host build stores only the graph; vectors are passed separately to from_cagra below. - cuvs::neighbors::host_padded_dataset_view host_padded_view( + cuvs::core::host_padded_dataset_view host_padded_view( dataset, static_cast(dataset.extent(1))); auto ace_host_index = cuvs::neighbors::cagra::build(res, cagra_params, host_padded_view); diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index ad9ea2e573..38926557d2 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -2477,7 +2477,8 @@ template template void GNND::local_join( cudaStream_t stream, - cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view, int64_t> + dataset, DistEpilogue_t dist_epilogue) { raft::matrix::fill(res, dists_buffer_.view(), std::numeric_limits::max()); @@ -2485,11 +2486,11 @@ void GNND::local_join( // Both kernels take the same (document, query) pair, so there is no symmetric/asymmetric split // here: a single quantizer just means the same one on both operands, which is exactly what // SelfJoin encodes. Picking the two quantizers is all that differs. - const bool self_join = dataset.quantizers.size() == 1; - const bool has_1b = dataset.has_layout(bbq_code_layout::packed_1b); - const bool has_4b = dataset.has_layout(bbq_code_layout::packed_4b); - const bool has_2bt = dataset.has_layout(bbq_code_layout::transposed_2b); - const bool has_4bt = dataset.has_layout(bbq_code_layout::transposed_4b); + const bool self_join = dataset.data().quantizers.size() == 1; + const bool has_1b = dataset.data().has_layout(bbq_code_layout::packed_1b); + const bool has_4b = dataset.data().has_layout(bbq_code_layout::packed_4b); + const bool has_2bt = dataset.data().has_layout(bbq_code_layout::transposed_2b); + const bool has_4bt = dataset.data().has_layout(bbq_code_layout::transposed_4b); // Asymmetric: a packed_4b query selects the tensor-core path, a transposed query the SIMT one. // Only packed_1b promotes to the tensor-core path; transposed_2b is SIMT-only (it would need @@ -2501,13 +2502,14 @@ void GNND::local_join( "packed_1b x packed_4b (tensor core); packed_1b x transposed_2b, " "packed_1b x transposed_4b, transposed_2b x transposed_4b (SIMT)."); auto quantizer_query = - self_join ? dataset.quantizers[0] - : (tc_pair ? dataset.get_quantizer(bbq_code_layout::packed_4b) - : (has_4bt ? dataset.get_quantizer(bbq_code_layout::transposed_4b) - : dataset.get_quantizer(bbq_code_layout::transposed_2b))); - auto quantizer_document = self_join ? dataset.quantizers[0] - : has_1b ? dataset.get_quantizer(bbq_code_layout::packed_1b) - : dataset.get_quantizer(bbq_code_layout::transposed_2b); + self_join + ? dataset.data().quantizers[0] + : (tc_pair ? dataset.data().get_quantizer(bbq_code_layout::packed_4b) + : (has_4bt ? dataset.data().get_quantizer(bbq_code_layout::transposed_4b) + : dataset.data().get_quantizer(bbq_code_layout::transposed_2b))); + auto quantizer_document = self_join ? dataset.data().quantizers[0] + : has_1b ? dataset.data().get_quantizer(bbq_code_layout::packed_1b) + : dataset.data().get_quantizer(bbq_code_layout::transposed_2b); // stage_tile_simt / stage_promoted_tile cast code buffers to uint32_t*, so every plane stride // must be 4-byte aligned. @@ -2859,7 +2861,8 @@ void GNND::build(Data_t* data, template template void GNND::build( - cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view, int64_t> + dataset, Index_t* output_graph, bool return_distances, DistData_t* output_distances, @@ -2979,11 +2982,11 @@ void GNND::build( template void build(raft::resources const& res, const index_params& params, - cuvs::neighbors::device_bbq_dataset_view dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, index& idx) { - RAFT_EXPECTS(dataset.quantizers.size() > 0, "BBQ dataset must not be empty."); - auto front_quantizer = dataset.quantizers[0]; + RAFT_EXPECTS(dataset.data().quantizers.size() > 0, "BBQ dataset must not be empty."); + auto front_quantizer = dataset.data().quantizers[0]; cuvs::common::nvtx::range fun_scope( "neighbors::nn_descent::detail::build-bbq(%zu, %zu, %zu, %zu, %zu)", size_t(dataset.n_rows()), @@ -3001,7 +3004,7 @@ void build(raft::resources const& res, "BBQ dataset metric does not match the NN-Descent metric."); // packed_4b is the only layout dispatched to local_join_kernel_bbq_wmma, and the int4 MMA that // kernel is built around only exists from sm_75 on. - if (dataset.has_layout(bbq_code_layout::packed_4b)) { + if (dataset.data().has_layout(bbq_code_layout::packed_4b)) { auto kernel = local_join_kernel_bbq_wmma -index build(raft::resources const& res, - const index_params& params, - cuvs::neighbors::device_bbq_dataset_view dataset) +index build( + raft::resources const& res, + const index_params& params, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset) { size_t graph_degree = params.graph_degree; if (params.intermediate_graph_degree < graph_degree) { diff --git a/cpp/src/neighbors/detail/nn_descent_gnnd.hpp b/cpp/src/neighbors/detail/nn_descent_gnnd.hpp index 07db8b7064..0179f8864f 100644 --- a/cpp/src/neighbors/detail/nn_descent_gnnd.hpp +++ b/cpp/src/neighbors/detail/nn_descent_gnnd.hpp @@ -207,11 +207,13 @@ class CUVS_EXPORT GNND { DistEpilogue_t dist_epilogue = DistEpilogue_t{}); template - void build(cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, - Index_t* output_graph, - bool return_distances, - DistData_t* output_distances, - DistEpilogue_t dist_epilogue = DistEpilogue_t{}); + void build( + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view, + int64_t> dataset, + Index_t* output_graph, + bool return_distances, + DistData_t* output_distances, + DistEpilogue_t dist_epilogue = DistEpilogue_t{}); ~GNND() = default; using ID_t = InternalID_t; void reset(raft::resources const& res); @@ -229,7 +231,8 @@ class CUVS_EXPORT GNND { template void local_join( cudaStream_t stream, - cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view, + int64_t> dataset, DistEpilogue_t dist_epilogue = DistEpilogue_t{}); raft::resources const& res; diff --git a/cpp/src/neighbors/detail/tiered_index.cuh b/cpp/src/neighbors/detail/tiered_index.cuh index 5e81958aa6..feaf06b75e 100644 --- a/cpp/src/neighbors/detail/tiered_index.cuh +++ b/cpp/src/neighbors/detail/tiered_index.cuh @@ -127,7 +127,7 @@ struct index_state { auto index = std::forward(build_fn)(res, tiered_params, dataset); if constexpr (std::is_same_v>) { index = cuvs::neighbors::cagra::update_dataset( - res, std::move(index), cuvs::neighbors::make_device_standard_dataset_view(dataset)); + res, std::move(index), cuvs::core::make_device_standard_dataset_view(dataset)); } return std::make_shared(std::move(index)); } @@ -303,7 +303,7 @@ inline void update_cagra_ann_dataset_for_stride( raft::device_matrix_view dataset) { ann_index = cuvs::neighbors::cagra::update_dataset( - res, std::move(ann_index), cuvs::neighbors::make_device_standard_dataset_view(dataset)); + res, std::move(ann_index), cuvs::core::make_device_standard_dataset_view(dataset)); } /** diff --git a/cpp/src/neighbors/detail/vamana/vamana_build.cuh b/cpp/src/neighbors/detail/vamana/vamana_build.cuh index 66cec4f630..15fe796e37 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_build.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_build.cuh @@ -677,12 +677,12 @@ index build( // process in batches const uint32_t n_rows = dataset.extent(0); - auto quantizer = cuvs::preprocessing::quantize::pq::quantizer( + using vpq_owning_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; + auto quantizer = cuvs::preprocessing::quantize::pq::quantizer( pq_params, - cuvs::neighbors::device_vpq_dataset{ - raft::make_device_matrix(res, 0, 0), - std::move(pq_codebook), - raft::make_device_matrix(res, 0, 0)}); + vpq_owning_t(raft::make_device_matrix(res, 0, 0), + raft::make_device_matrix(res, 0, 0), + std::move(pq_codebook))); const int64_t codes_rowlen = cuvs::preprocessing::quantize::pq::get_quantized_dim(pq_params); quantized_vectors = raft::make_device_matrix(res, n_rows, codes_rowlen); diff --git a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh index 8a21ed3f1b..b6576f602d 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh @@ -90,12 +90,12 @@ void serialize_dataset_view(raft::resources const& res, */ template void serialize_dataset(raft::resources const& res, - const cuvs::neighbors::device_padded_dataset_view* dataset, + const cuvs::core::device_padded_dataset_view* dataset, const std::string& dataset_base_file) { if (dataset == nullptr) { return; } try { - serialize_dataset_view(res, dataset->view(), dataset_base_file); + serialize_dataset_view(res, dataset->as_matrix_view(), dataset_base_file); } catch (std::bad_alloc& e) { RAFT_LOG_INFO("Failed to serialize dataset"); } catch (raft::logic_error& e) { @@ -129,12 +129,11 @@ void serialize_dataset(raft::resources const& res, * */ template -void serialize_sector_aligned( - raft::resources const& res, - const HostMatT& h_graph, - const cuvs::neighbors::device_padded_dataset_view& dataset, - const uint64_t medoid, - std::ostream& output_writer) +void serialize_sector_aligned(raft::resources const& res, + const HostMatT& h_graph, + const cuvs::core::device_padded_dataset_view& dataset, + const uint64_t medoid, + std::ostream& output_writer) { if constexpr (!std::is_same_v) { RAFT_FAIL("serialization is only implemented for uint32_t graph"); @@ -172,8 +171,8 @@ void serialize_sector_aligned( auto h_data = raft::make_host_matrix(npts, ndims); raft::copy_matrix(h_data.data_handle(), ndims, - dataset.view().data_handle(), - dataset.stride(), + dataset.as_matrix_view().data_handle(), + dataset.as_matrix_view().stride(), ndims, npts, raft::resource::get_cuda_stream(res)); diff --git a/cpp/src/neighbors/detail/vpq_dataset.cuh b/cpp/src/neighbors/detail/vpq_dataset.cuh index 4d27c5424a..ee17e5b1f8 100644 --- a/cpp/src/neighbors/detail/vpq_dataset.cuh +++ b/cpp/src/neighbors/detail/vpq_dataset.cuh @@ -122,11 +122,13 @@ void train_pq_centers( } template -auto fill_missing_params_heuristics(const vpq_params& params, const DatasetT& dataset) -> vpq_params +auto fill_missing_params_heuristics(const cuvs::preprocessing::quantize::pq::vpq_params& params, + const DatasetT& dataset) + -> cuvs::preprocessing::quantize::pq::vpq_params { - vpq_params r = params; - double n_rows = dataset.extent(0); - size_t dim = dataset.extent(1); + cuvs::preprocessing::quantize::pq::vpq_params r = params; + double n_rows = dataset.extent(0); + size_t dim = dataset.extent(1); if (r.pq_dim == 0) { r.pq_dim = raft::div_rounding_up_safe(dim, size_t{4}); } if (r.pq_bits == 0) { r.pq_bits = 8; } if (r.vq_n_centers == 0) { r.vq_n_centers = raft::round_up_safe(std::sqrt(n_rows), 8); } @@ -168,8 +170,9 @@ auto transform_data(const raft::resources& res, DatasetT dataset) using ix_t = int64_t; template -auto train_vq(const raft::resources& res, const vpq_params& params, const DatasetT& dataset) - -> raft::device_matrix +auto train_vq(const raft::resources& res, + const cuvs::preprocessing::quantize::pq::vpq_params& params, + const DatasetT& dataset) -> raft::device_matrix { using kmeans_in_type = typename DatasetT::value_type; const ix_t n_rows = dataset.extent(0); @@ -413,7 +416,7 @@ __launch_bounds__(BlockSize) RAFT_KERNEL process_and_fill_codes_kernel( template void process_and_fill_codes( const raft::resources& res, - const vpq_params& params, + const cuvs::preprocessing::quantize::pq::vpq_params& params, const DatasetT& dataset, raft::device_matrix_view pq_centers, raft::device_matrix_view vq_centers, @@ -422,7 +425,7 @@ void process_and_fill_codes( bool inline_vq_labels = false) { using data_t = typename DatasetT::value_type; - using cdataset_t = device_vpq_dataset; + using cdataset_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; using label_t = uint32_t; const ix_t n_rows = dataset.extent(0); @@ -807,7 +810,7 @@ __launch_bounds__(BlockSize) RAFT_KERNEL process_and_fill_codes_subspaces_kernel template void process_and_fill_codes_subspaces( const raft::resources& res, - const vpq_params& params, + const cuvs::preprocessing::quantize::pq::vpq_params& params, const DatasetT& dataset, raft::device_matrix_view pq_centers, raft::device_matrix_view vq_centers, @@ -815,7 +818,7 @@ void process_and_fill_codes_subspaces( raft::device_matrix_view codes) { using data_t = typename DatasetT::value_type; - using cdataset_t = device_vpq_dataset; + using cdataset_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; using label_t = uint32_t; const ix_t n_rows = dataset.extent(0); diff --git a/cpp/src/neighbors/iface/iface.hpp b/cpp/src/neighbors/iface/iface.hpp index 574859ae7a..dc950df185 100644 --- a/cpp/src/neighbors/iface/iface.hpp +++ b/cpp/src/neighbors/iface/iface.hpp @@ -33,7 +33,7 @@ bool dataset_mdspan_uses_padded_device_view( { using value_type = T; uint32_t const required_stride = - cagra_required_row_width(static_cast(mds.extent(1))); + cuvs::core::padded_row_width(static_cast(mds.extent(1))); uint32_t const src_stride = mds.stride(0) > 0 ? static_cast(mds.stride(0)) : static_cast(mds.extent(1)); cudaPointerAttributes a{}; @@ -55,12 +55,12 @@ void cagra_build_from_device_dataset( auto dview = raft::make_device_strided_matrix_view( m.data_handle(), m.extent(0), m.extent(1), stride); if constexpr (std::is_same_v>) { - auto padded = cuvs::neighbors::make_device_padded_dataset_view(h, dview); + auto padded = cuvs::core::make_device_padded_dataset_view(h, dview); auto index = cuvs::neighbors::cagra::build(h, cagra_params, padded); index = cuvs::neighbors::cagra::update_dataset(h, std::move(index), padded); interface.index_.emplace(std::move(index)); } else { - auto standard = cuvs::neighbors::make_device_standard_dataset_view(dview); + auto standard = cuvs::core::make_device_standard_dataset_view(dview); auto index = cuvs::neighbors::cagra::build(h, cagra_params, standard); index = cuvs::neighbors::cagra::update_dataset(h, std::move(index), standard); interface.index_.emplace(std::move(index)); @@ -104,23 +104,23 @@ void build(const raft::resources& handle, auto host_view = raft::make_host_matrix_view( index_dataset.data_handle(), index_dataset.extent(0), index_dataset.extent(1)); if constexpr (std::is_same>::value) { - cuvs::neighbors::host_padded_dataset_view host_padded( + cuvs::core::host_padded_dataset_view host_padded( host_view, static_cast(host_view.extent(1))); auto host_idx = cuvs::neighbors::cagra::build(handle, cagra_params, host_padded); - auto padded_r = cuvs::neighbors::make_device_padded_dataset(handle, index_dataset); + auto padded_r = cuvs::core::make_device_padded_dataset(handle, index_dataset); auto device_idx = cuvs::neighbors::cagra::update_dataset( handle, std::move(host_idx), padded_r->as_dataset_view()); interface.cagra_owned_padded_dataset_ = std::move(padded_r); interface.cagra_owned_standard_dataset_.reset(); interface.index_.emplace(std::move(device_idx)); } else { - auto host_standard = cuvs::neighbors::make_host_standard_dataset_view(host_view); + auto host_standard = cuvs::core::make_host_standard_dataset_view(host_view); auto host_idx = cuvs::neighbors::cagra::build(handle, cagra_params, host_standard); - auto standard_r = cuvs::neighbors::make_device_standard_dataset( - handle, - index_dataset, - static_cast(index_dataset.extent(1)), - static_cast(index_dataset.stride(0))); + auto standard_r = + cuvs::core::make_device_standard_dataset(handle, + index_dataset, + static_cast(index_dataset.extent(1)), + static_cast(index_dataset.stride(0))); auto device_idx = cuvs::neighbors::cagra::update_dataset( handle, std::move(host_idx), standard_r->as_dataset_view()); interface.cagra_owned_standard_dataset_ = std::move(standard_r); @@ -251,7 +251,7 @@ void deserialize(const raft::resources& handle, interface.index_.emplace(std::move(idx)); } else if constexpr (std::is_same>::value) { cagra::device_padded_index idx(handle); - std::unique_ptr> out_dataset; + std::unique_ptr> out_dataset; cagra::deserialize(handle, is, &idx, &out_dataset); interface.cagra_owned_padded_dataset_.reset(); interface.cagra_owned_standard_dataset_.reset(); @@ -260,7 +260,7 @@ void deserialize(const raft::resources& handle, interface.index_.emplace(std::move(idx)); } else if constexpr (std::is_same>::value) { cagra::device_standard_index idx(handle); - std::unique_ptr> out_dataset; + std::unique_ptr> out_dataset; cagra::deserialize(handle, is, &idx, &out_dataset); interface.cagra_owned_padded_dataset_.reset(); interface.cagra_owned_standard_dataset_.reset(); @@ -289,7 +289,7 @@ void deserialize(const raft::resources& handle, interface.index_.emplace(std::move(idx)); } else if constexpr (std::is_same>::value) { cagra::device_padded_index idx(handle); - std::unique_ptr> out_dataset; + std::unique_ptr> out_dataset; cagra::deserialize(handle, filename, &idx, &out_dataset); interface.cagra_owned_padded_dataset_.reset(); interface.cagra_owned_standard_dataset_.reset(); @@ -298,7 +298,7 @@ void deserialize(const raft::resources& handle, interface.index_.emplace(std::move(idx)); } else if constexpr (std::is_same>::value) { cagra::device_standard_index idx(handle); - std::unique_ptr> out_dataset; + std::unique_ptr> out_dataset; cagra::deserialize(handle, filename, &idx, &out_dataset); interface.cagra_owned_padded_dataset_.reset(); interface.cagra_owned_standard_dataset_.reset(); diff --git a/cpp/src/neighbors/mg/mg_cagra_inst.cu.in b/cpp/src/neighbors/mg/mg_cagra_inst.cu.in index 077a065e58..d155f20461 100644 --- a/cpp/src/neighbors/mg/mg_cagra_inst.cu.in +++ b/cpp/src/neighbors/mg/mg_cagra_inst.cu.in @@ -11,10 +11,10 @@ template void distribute_padded_dataset( const raft::resources& res, const cuvs::neighbors::mg_index& idx, - cuvs::neighbors::device_padded_dataset_view const& padded_dataset, + cuvs::core::device_padded_dataset_view const& padded_dataset, RankUpdate&& rank_update) { - auto padded_mds = padded_dataset.view(); + auto padded_mds = padded_dataset.as_matrix_view(); auto stride = padded_mds.extent(1); const raft::resources& root_res = raft::resource::set_current_device_to_root_rank(res); auto padded_host = @@ -38,7 +38,7 @@ void distribute_padded_dataset( auto rank_device_matrix = raft::make_device_matrix(dev_res, rank_rows, stride); raft::copy(dev_res, rank_device_matrix.view(), rank_host_mds); - auto rank_device_padded = std::make_unique>( + auto rank_device_padded = std::make_unique>( std::move(rank_device_matrix), padded_dataset.dim()); rank_update(dev_res, rank, std::move(rank_device_padded)); if (idx.mode_ == cuvs::neighbors::SHARDED) { offset += rank_rows; } @@ -53,7 +53,7 @@ void distribute_padded_dataset( cuvs::neighbors::mg_index, T, IdxT> build( \ const raft::resources& res, \ const mg_index_params& index_params, \ - cuvs::neighbors::host_standard_dataset_view const& index_dataset) \ + cuvs::core::host_standard_dataset_view const& index_dataset) \ { \ cuvs::neighbors::mg_index, T, IdxT> index( \ res, index_params.mode); \ @@ -61,14 +61,14 @@ void distribute_padded_dataset( res, \ index, \ static_cast(&index_params), \ - index_dataset.view()); \ + index_dataset.as_matrix_view()); \ return index; \ } \ \ cuvs::neighbors::mg_index, T, IdxT> build( \ const raft::resources& res, \ const mg_index_params& index_params, \ - cuvs::neighbors::host_padded_dataset_view const& index_dataset) \ + cuvs::core::host_padded_dataset_view const& index_dataset) \ { \ cuvs::neighbors::mg_index, T, IdxT> index( \ res, index_params.mode); \ @@ -76,14 +76,14 @@ void distribute_padded_dataset( res, \ index, \ static_cast(&index_params), \ - index_dataset.view()); \ + index_dataset.as_matrix_view()); \ return index; \ } \ \ auto update_dataset( \ const raft::resources& res, \ cuvs::neighbors::mg_index, T, IdxT>&& idx, \ - cuvs::neighbors::device_padded_dataset_view const& padded_dataset) \ + cuvs::core::device_padded_dataset_view const& padded_dataset) \ -> cuvs::neighbors::mg_index, T, IdxT> \ { \ cuvs::neighbors::mg_index, T, IdxT> out(res, idx.mode_); \ @@ -104,7 +104,7 @@ void distribute_padded_dataset( void update_dataset( \ const raft::resources& res, \ cuvs::neighbors::mg_index, T, IdxT>& idx, \ - cuvs::neighbors::device_padded_dataset_view const& padded_dataset) \ + cuvs::core::device_padded_dataset_view const& padded_dataset) \ { \ detail::distribute_padded_dataset( \ res, idx, padded_dataset, [&](const raft::resources& dev_res, int rank, auto dataset) { \ @@ -119,18 +119,18 @@ void distribute_padded_dataset( \ void extend(const raft::resources& res, \ cuvs::neighbors::mg_index, T, IdxT>& index, \ - cuvs::neighbors::host_padded_dataset_view new_vectors, \ + cuvs::core::host_padded_dataset_view new_vectors, \ std::optional> new_indices) \ { \ - cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.view(), new_indices); \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.as_matrix_view(), new_indices); \ } \ \ void extend(const raft::resources& res, \ cuvs::neighbors::mg_index, T, IdxT>& index, \ - cuvs::neighbors::host_standard_dataset_view new_vectors, \ + cuvs::core::host_standard_dataset_view new_vectors, \ std::optional> new_indices) \ { \ - cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.view(), new_indices); \ + cuvs::neighbors::snmg::detail::extend(res, index, new_vectors.as_matrix_view(), new_indices); \ } \ \ void search( \ diff --git a/cpp/src/neighbors/nn_descent.cuh b/cpp/src/neighbors/nn_descent.cuh index 9c9b3c5540..89f35dce92 100644 --- a/cpp/src/neighbors/nn_descent.cuh +++ b/cpp/src/neighbors/nn_descent.cuh @@ -187,7 +187,8 @@ void build(raft::resources const& res, template auto build(raft::resources const& res, index_params const& params, - cuvs::neighbors::device_bbq_dataset_view dataset) -> index + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset) + -> index { return detail::build(res, params, dataset); } @@ -195,7 +196,7 @@ auto build(raft::resources const& res, template void build(raft::resources const& res, index_params const& params, - cuvs::neighbors::device_bbq_dataset_view dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, index& idx) { detail::build(res, params, dataset, idx); diff --git a/cpp/src/neighbors/nn_descent_gnnd_inst.cu b/cpp/src/neighbors/nn_descent_gnnd_inst.cu index a5b2517b9c..ecb9c01cff 100644 --- a/cpp/src/neighbors/nn_descent_gnnd_inst.cu +++ b/cpp/src/neighbors/nn_descent_gnnd_inst.cu @@ -40,7 +40,8 @@ template void detail::GNND::build( float* output_distances, raft::identity_op dist_epilogue); template void detail::GNND::build( - cuvs::neighbors::device_bbq_dataset_view, int64_t> dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view, int64_t> + dataset, int* output_graph, bool return_distances, float* output_distances, diff --git a/cpp/src/neighbors/nn_descent_inst.cu.in b/cpp/src/neighbors/nn_descent_inst.cu.in index 62d749e5b6..141ccb6881 100644 --- a/cpp/src/neighbors/nn_descent_inst.cu.in +++ b/cpp/src/neighbors/nn_descent_inst.cu.in @@ -54,7 +54,7 @@ auto build(raft::resources const& handle, auto build(raft::resources const& handle, const cuvs::neighbors::nn_descent::index_params& params, - cuvs::neighbors::device_bbq_dataset_view dataset, + cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view dataset, std::optional> graph) -> cuvs::neighbors::nn_descent::index { diff --git a/cpp/src/neighbors/scann/detail/scann_build.cuh b/cpp/src/neighbors/scann/detail/scann_build.cuh index d65f50d36a..b0f93a7979 100644 --- a/cpp/src/neighbors/scann/detail/scann_build.cuh +++ b/cpp/src/neighbors/scann/detail/scann_build.cuh @@ -437,7 +437,8 @@ index build( // Codebooks from VPQ have the shape [subspace idx, subspace dim, code] // This converts the codebook into matrix format for easy interoperability // with open-source ScaNN search - auto full_codebook_view = pq_quantizer.vpq_codebooks.pq_code_book.view(); + auto const& vpq = pq_quantizer.vpq_codebooks.data(); + auto full_codebook_view = vpq.pq_code_book.view(); raft::linalg::map_offset( res, diff --git a/cpp/src/neighbors/tiered_index.cu b/cpp/src/neighbors/tiered_index.cu index fd7848454f..592888c171 100644 --- a/cpp/src/neighbors/tiered_index.cu +++ b/cpp/src/neighbors/tiered_index.cu @@ -37,7 +37,7 @@ cuvs::neighbors::cagra::device_standard_index cagra_build_for_t cuvs::neighbors::cagra::index_params const& params, raft::device_matrix_view dataset) { - auto view = cuvs::neighbors::make_device_standard_dataset_view(dataset); + auto view = cuvs::core::make_device_standard_dataset_view(dataset); return cuvs::neighbors::cagra::build(res, params, view); } @@ -46,7 +46,7 @@ cuvs::neighbors::cagra::device_padded_index cagra_build_for_tie cuvs::neighbors::cagra::index_params const& params, raft::device_matrix_view dataset) { - auto view = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + auto view = cuvs::core::make_device_padded_dataset_view(res, dataset); return cuvs::neighbors::cagra::build(res, params, view); } @@ -55,11 +55,11 @@ cuvs::neighbors::cagra::device_padded_index cagra_build_for_tie namespace cuvs::neighbors::tiered_index { auto build(raft::resources const& res, const index_params& params, - cuvs::neighbors::device_padded_dataset_view dataset) + cuvs::core::device_padded_dataset_view dataset) -> tiered_index::index> { auto state = detail::build>( - res, params, cagra_build_for_tiered_padded, dataset.view()); + res, params, cagra_build_for_tiered_padded, dataset.as_matrix_view()); return cuvs::neighbors::tiered_index::index>(state); } @@ -76,7 +76,7 @@ auto build(raft::resources const& res, auto convert_standard_to_padded_index( raft::resources const& res, const tiered_index::index>& idx, - cuvs::neighbors::device_padded_dataset_view padded_dataset) + cuvs::core::device_padded_dataset_view padded_dataset) -> tiered_index::index> { RAFT_EXPECTS(padded_dataset.n_rows() == idx.size(), @@ -92,12 +92,12 @@ auto convert_standard_to_padded_index( next_state->ann_index.reset(); if (idx.state->ann_index) { - auto padded_mds = padded_dataset.view(); + auto padded_mds = padded_dataset.as_matrix_view(); auto ann_rows = static_cast(idx.state->ann_rows()); auto ann_mds = raft::make_device_matrix_view( padded_mds.data_handle(), ann_rows, static_cast(padded_mds.extent(1))); auto ann_padded_view = - cuvs::neighbors::device_padded_dataset_view(ann_mds, padded_dataset.dim()); + cuvs::core::device_padded_dataset_view(ann_mds, padded_dataset.dim()); auto ann_padded_idx = cuvs::neighbors::cagra::convert_standard_to_padded_index( res, *idx.state->ann_index, ann_padded_view); next_state->ann_index = @@ -243,7 +243,7 @@ void search(raft::resources const& res, auto storage = index.state->storage; auto vectors = raft::make_device_matrix_view( storage->dataset.data(), index.size(), static_cast(storage->dim)); - auto padded_dataset = cuvs::neighbors::make_device_padded_dataset(res, vectors); + auto padded_dataset = cuvs::core::make_device_padded_dataset(res, vectors); auto padded_index = convert_standard_to_padded_index(res, index, padded_dataset->as_dataset_view()); padded_index.state->search( diff --git a/cpp/src/preprocessing/quantize/detail/pq.cuh b/cpp/src/preprocessing/quantize/detail/pq.cuh index 8a2a6c3cc9..5f7e8c2145 100644 --- a/cpp/src/preprocessing/quantize/detail/pq.cuh +++ b/cpp/src/preprocessing/quantize/detail/pq.cuh @@ -56,11 +56,11 @@ inline cuvs::distance::DistanceType get_kmeans_metric( } inline auto to_vpq_params(const cuvs::preprocessing::quantize::pq::params& params) - -> cuvs::neighbors::vpq_params + -> cuvs::preprocessing::quantize::pq::vpq_params { auto kmeans_type = is_balanced_kmeans(params) ? cuvs::cluster::kmeans::kmeans_type::KMeansBalanced : cuvs::cluster::kmeans::kmeans_type::KMeans; - return cuvs::neighbors::vpq_params{ + return cuvs::preprocessing::quantize::pq::vpq_params{ .pq_bits = params.pq_bits, .pq_dim = params.pq_dim, .vq_n_centers = params.vq_n_centers, @@ -192,9 +192,9 @@ quantizer build( pq_code_book = cuvs::neighbors::detail::train_pq( res, filled_params, dataset, raft::make_const_mdspan(vq_code_book.view())); } + using owning_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; return {filled_params, - cuvs::neighbors::device_vpq_dataset{ - std::move(vq_code_book), std::move(pq_code_book), std::move(empty_codes)}}; + owning_t(std::move(empty_codes), std::move(vq_code_book), std::move(pq_code_book))}; } template @@ -217,7 +217,8 @@ void transform( RAFT_EXPECTS(quantizer.params_quantizer.pq_bits >= 4 && quantizer.params_quantizer.pq_bits <= 16, "PQ bits must be within [4, 16]"); // Encode dataset - auto vq_centers = raft::make_const_mdspan(quantizer.vpq_codebooks.vq_code_book.view()); + auto const& vpq = quantizer.vpq_codebooks.data(); + auto vq_centers = vpq.vq_code_book.view(); auto vq_labels_view = raft::make_device_vector_view(nullptr, 0); if (vq_labels.has_value()) { vq_labels_view = vq_labels.value(); } @@ -226,7 +227,7 @@ void transform( res, to_vpq_params(quantizer.params_quantizer), dataset, - raft::make_const_mdspan(quantizer.vpq_codebooks.pq_code_book.view()), + vpq.pq_code_book.view(), vq_centers, vq_labels_view, pq_codes_out); @@ -235,7 +236,7 @@ void transform( res, to_vpq_params(quantizer.params_quantizer), dataset, - raft::make_const_mdspan(quantizer.vpq_codebooks.pq_code_book.view()), + vpq.pq_code_book.view(), vq_centers, vq_labels_view, pq_codes_out); @@ -356,33 +357,36 @@ void inverse_transform( "Codes matrix doesn't have the correct number of columns"); RAFT_EXPECTS(quant.params_quantizer.pq_bits >= 4 && quant.params_quantizer.pq_bits <= 16, "PQ bits must be within [4, 16]"); - reconstruct_vectors( - res, - quant.params_quantizer, - codes, - raft::make_const_mdspan(quant.vpq_codebooks.pq_code_book.view()), - raft::make_const_mdspan(quant.vpq_codebooks.vq_code_book.view()), - vq_labels, - out, - quant.params_quantizer.use_subspaces); + auto const& vpq = quant.vpq_codebooks.data(); + reconstruct_vectors(res, + quant.params_quantizer, + codes, + vpq.pq_code_book.view(), + vpq.vq_code_book.view(), + vq_labels, + out, + quant.params_quantizer.use_subspaces); } template -void vpq_convert_math_type(const raft::resources& res, - const cuvs::neighbors::device_vpq_dataset& src, - cuvs::neighbors::device_vpq_dataset& dst) +void vpq_convert_math_type( + const raft::resources& res, + const cuvs::preprocessing::quantize::pq::device_vpq_dataset& src, + raft::device_matrix_view dst_vq_code_book, + raft::device_matrix_view dst_pq_code_book) { + auto const& src_vpq = src.data(); raft::linalg::map(res, - dst.vq_code_book.view(), + dst_vq_code_book, cuvs::spatial::knn::detail::utils::mapping{}, - raft::make_const_mdspan(src.vq_code_book.view())); + src_vpq.vq_code_book.view()); raft::linalg::map(res, - dst.pq_code_book.view(), + dst_pq_code_book, cuvs::spatial::knn::detail::utils::mapping{}, - raft::make_const_mdspan(src.pq_code_book.view())); + src_vpq.pq_code_book.view()); } -inline auto make_pq_params_from_vpq(const cuvs::neighbors::vpq_params& in_params, +inline auto make_pq_params_from_vpq(const cuvs::preprocessing::quantize::pq::vpq_params& in_params, const uint64_t n_rows) -> cuvs::preprocessing::quantize::pq::params { @@ -408,8 +412,9 @@ inline auto make_pq_params_from_vpq(const cuvs::neighbors::vpq_params& in_params template auto vpq_build(const raft::resources& res, - const cuvs::neighbors::vpq_params& params, - const DatasetT& dataset) -> cuvs::neighbors::device_vpq_dataset + const cuvs::preprocessing::quantize::pq::vpq_params& params, + const DatasetT& dataset) + -> cuvs::preprocessing::quantize::pq::device_vpq_dataset { using label_t = uint32_t; // Use a heuristic to impute missing parameters. @@ -437,21 +442,25 @@ auto vpq_build(const raft::resources& res, codes.view(), true); - return cuvs::neighbors::device_vpq_dataset{ - std::move(vq_code_book), std::move(pq_code_book), std::move(codes)}; + using owning_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; + return owning_t(std::move(codes), std::move(vq_code_book), std::move(pq_code_book)); } template auto vpq_build_half(const raft::resources& res, - const cuvs::neighbors::vpq_params& params, - const DatasetT& dataset) -> cuvs::neighbors::device_vpq_dataset + const cuvs::preprocessing::quantize::pq::vpq_params& params, + const DatasetT& dataset) + -> cuvs::preprocessing::quantize::pq::device_vpq_dataset { - auto old_type = vpq_build(res, params, dataset); - auto new_type = cuvs::neighbors::device_vpq_dataset{ - raft::make_device_mdarray(res, old_type.vq_code_book.extents()), - raft::make_device_mdarray(res, old_type.pq_code_book.extents()), - std::move(old_type.data)}; - vpq_convert_math_type(res, old_type, new_type); - return new_type; + auto old_type = vpq_build(res, params, dataset); + using new_owning_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; + auto const& old_vpq = old_type.data(); + auto new_vq_code_book = raft::make_device_mdarray(res, old_vpq.vq_code_book.extents()); + auto new_pq_code_book = raft::make_device_mdarray(res, old_vpq.pq_code_book.extents()); + vpq_convert_math_type( + res, old_type, new_vq_code_book.view(), new_pq_code_book.view()); + // Reuse the already-encoded codes (the owning payload's mdarray base) with the new codebooks. + return new_owning_t( + std::move(old_type.data()), std::move(new_vq_code_book), std::move(new_pq_code_book)); } } // namespace cuvs::preprocessing::quantize::pq::detail diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 20b8f21d36..55276553c2 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -57,13 +57,13 @@ CUVS_INST_QUANTIZATION(float, uint8_t); #define CUVS_INST_VPQ_BUILD(T) \ auto vpq_build(const raft::resources& res, \ - const cuvs::neighbors::vpq_params& params, \ + const cuvs::preprocessing::quantize::pq::vpq_params& params, \ const raft::host_matrix_view& dataset) \ { \ return detail::vpq_build_half(res, params, dataset); \ } \ auto vpq_build(const raft::resources& res, \ - const cuvs::neighbors::vpq_params& params, \ + const cuvs::preprocessing::quantize::pq::vpq_params& params, \ const raft::device_matrix_view& dataset) \ { \ return detail::vpq_build_half(res, params, dataset); \ @@ -80,11 +80,12 @@ namespace detail { template auto train_from_rows(raft::resources const& res, - cuvs::neighbors::vpq_params const& params, + cuvs::preprocessing::quantize::pq::vpq_params const& params, T const* src_ptr, int64_t n_rows, int64_t dim, - int64_t stride) -> cuvs::neighbors::device_vpq_dataset + int64_t stride) + -> cuvs::preprocessing::quantize::pq::device_vpq_dataset { cudaPointerAttributes ptr_attrs; RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src_ptr)); @@ -114,12 +115,13 @@ auto train_from_rows(raft::resources const& res, } auto vpq_train_from_rows(raft::resources const& res, - cuvs::neighbors::vpq_params const& params, + cuvs::preprocessing::quantize::pq::vpq_params const& params, void const* src_ptr, cudaDataType_t dtype, int64_t n_rows, int64_t dim, - int64_t stride) -> cuvs::neighbors::device_vpq_dataset + int64_t stride) + -> cuvs::preprocessing::quantize::pq::device_vpq_dataset { switch (dtype) { case CUDA_R_32F: diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index a662c32f03..40866de655 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -65,11 +65,11 @@ void cagra_build_into_index( raft::resources const& res, cagra::index_params const& params, std::optional> ace_host_dataset, - cuvs::neighbors::device_padded_dataset_view const& padded, + cuvs::core::device_padded_dataset_view const& padded, cagra::device_padded_index& index) { if (ace_host_dataset.has_value()) { - cuvs::neighbors::host_padded_dataset_view host_view( + cuvs::core::host_padded_dataset_view host_view( *ace_host_dataset, static_cast(ace_host_dataset->extent(1))); auto host_idx = cagra::build(res, params, host_view); // In-memory ACE returns graph-only; attach device padded storage for search. @@ -488,7 +488,7 @@ class AnnCagraTest : public ::testing::TestWithParam { } cagra::device_padded_index index(handle_); - std::unique_ptr> loaded_dataset; + std::unique_ptr> loaded_dataset; cagra::deserialize(handle_, index_file.filename, &index, &loaded_dataset); if (!ps.include_serialized_dataset) { @@ -1329,8 +1329,8 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam( handle_, ps.n_rows - static_cast(test_cagra_sample_filter::offset), - static_cast(index0.dataset().stride())); - auto merged_dataset = cuvs::neighbors::device_padded_dataset( + static_cast(index0.dataset().as_matrix_view().stride())); + auto merged_dataset = cuvs::core::device_padded_dataset( std::move(merged_matrix), static_cast(ps.dim)); auto merge_idx = cuvs::neighbors::cagra::merge( handle_, index_params, indices, merged_dataset.as_dataset_view(), bitset_filter_obj); @@ -1575,8 +1575,8 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { auto const merged_rows = static_cast(index0.size()) + static_cast(index1.size()); auto merged_matrix = raft::make_device_matrix( - handle_, merged_rows, static_cast(index0.dataset().stride())); - auto merged_dataset = cuvs::neighbors::device_padded_dataset( + handle_, merged_rows, static_cast(index0.dataset().as_matrix_view().stride())); + auto merged_dataset = cuvs::core::device_padded_dataset( std::move(merged_matrix), static_cast(ps.dim)); auto merged_idx = ps.physical_merge_params.has_value() diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu index a6ad284a1a..a60215eab5 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -83,10 +83,10 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { // Step 2: Update to FULL dataset (1000 points) but keep small graph (500 nodes) // This creates the exact bug scenario: dataset.size=1000, graph.extent(0)=500 - small_index = cagra::update_dataset(res, - std::move(small_index), - cuvs::neighbors::make_device_padded_dataset_view( - res, raft::make_const_mdspan(dataset.view()))); + small_index = cagra::update_dataset( + res, + std::move(small_index), + cuvs::core::make_device_padded_dataset_view(res, raft::make_const_mdspan(dataset.view()))); // Verify the mismatch - THIS IS THE BUG SCENARIO! ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 500 nodes diff --git a/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu b/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu index 8a0b7c3d36..d418a548bc 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_iterative_cagra_build.cu @@ -59,7 +59,7 @@ class CagraIterativeBuildBugTest : public ::testing::Test { // so the searches driving the build run on compressed rows instead of dense ones. void run_compressed() { - cuvs::neighbors::vpq_params vpq_params; + cuvs::preprocessing::quantize::pq::vpq_params vpq_params; // pq_len = n_dim / pq_dim must be 2, 4 or 8 for CAGRA-Q. Codebook quality is irrelevant here, // since only graph construction is under test, so training stays short. vpq_params.pq_dim = static_cast(n_dim / 4); diff --git a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu index 68794c3151..ad9370e158 100644 --- a/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu +++ b/cpp/tests/neighbors/ann_cagra/test_filter_udf.cu @@ -126,10 +126,10 @@ class CagraUdfFilterTest : public ::testing::TestWithParam { index_params.graph_build_params = cagra::graph_build_params::nn_descent_params(index_params.intermediate_graph_degree); - index.emplace(cagra::build(res, - index_params, - cuvs::neighbors::make_device_padded_dataset_view( - res, raft::make_const_mdspan(dataset->view())))); + index.emplace(cagra::build( + res, + index_params, + cuvs::core::make_device_padded_dataset_view(res, raft::make_const_mdspan(dataset->view())))); raft::resource::sync_stream(res); } @@ -205,10 +205,10 @@ class CagraUdfFilterHalfTest : public ::testing::TestWithParamview())))); + index.emplace(cagra::build( + res, + index_params, + cuvs::core::make_device_padded_dataset_view(res, raft::make_const_mdspan(dataset->view())))); raft::resource::sync_stream(res); } diff --git a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu index abc99b2842..84e6014ab5 100644 --- a/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu +++ b/cpp/tests/neighbors/ann_cagra/test_iterative_cagra_q.cu @@ -37,7 +37,7 @@ namespace cuvs::neighbors::cagra { -using vpq_dataset_t = cuvs::neighbors::device_vpq_dataset; +using vpq_dataset_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; namespace { @@ -46,7 +46,7 @@ auto compress(const raft::resources& res, uint32_t pq_dim, uint32_t pq_bits = 8) -> vpq_dataset_t { - cuvs::neighbors::vpq_params params; + cuvs::preprocessing::quantize::pq::vpq_params params; params.pq_dim = pq_dim; params.pq_bits = pq_bits; params.vq_n_centers = 32; @@ -179,7 +179,7 @@ class CagraQBuildTest : public CagraQCompressedTestBase, TEST_P(CagraQBuildTest, BuildsAndSearchesAFreshlyCompressedDataset) { auto compressed = compress(res_, dataset(), params_.pq_dim); - ASSERT_EQ(compressed.pq_len(), static_cast(params_.dim / params_.pq_dim)); + ASSERT_EQ(compressed.data().pq_len(), static_cast(params_.dim / params_.pq_dim)); auto idx = cagra::build(res_, iterative_params(), compressed.as_dataset_view()); ASSERT_EQ(idx.size(), params_.n_rows); @@ -272,7 +272,7 @@ TEST_F(CagraQContractTest, RejectsMetricOtherThanL2Expanded) TEST_F(CagraQContractTest, RejectsPqBitsOtherThan8) { auto compressed = compress(res_, dataset(), 32, /* pq_bits */ 6); - ASSERT_EQ(compressed.pq_bits(), 6u); + ASSERT_EQ(compressed.data().pq_bits(), 6u); EXPECT_THROW(cagra::build(res_, iterative_params(), compressed.as_dataset_view()), raft::exception); } @@ -280,7 +280,7 @@ TEST_F(CagraQContractTest, RejectsPqBitsOtherThan8) TEST_F(CagraQContractTest, RejectsPqLenOutsideSupportedSet) { auto compressed = compress(res_, dataset(), /* pq_dim */ 4); // pq_len = 64 / 4 = 16 - ASSERT_EQ(compressed.pq_len(), 16u); + ASSERT_EQ(compressed.data().pq_len(), 16u); EXPECT_THROW(cagra::build(res_, iterative_params(), compressed.as_dataset_view()), raft::exception); } @@ -293,7 +293,7 @@ TEST_F(CagraQContractTest, RejectsEmptyDataset) auto vq_code_book = raft::make_device_matrix(res_, 1, width); auto pq_code_book = raft::make_device_matrix(res_, 256, 2); auto codes = raft::make_device_matrix(res_, 0, 4 + dim / 2); - vpq_dataset_t empty{std::move(vq_code_book), std::move(pq_code_book), std::move(codes)}; + vpq_dataset_t empty{std::move(codes), std::move(vq_code_book), std::move(pq_code_book)}; ASSERT_EQ(empty.n_rows(), 0); EXPECT_THROW(cagra::build(res_, iterative_params(), empty.as_dataset_view()), raft::exception); diff --git a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu index b7f0e44c1b..939078c909 100644 --- a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -57,7 +57,7 @@ inline auto make_ranges(raft::resources const& res, template struct padded_storage { raft::device_matrix matrix; - cuvs::neighbors::device_padded_dataset_view view; + cuvs::core::device_padded_dataset_view view; }; template @@ -66,7 +66,7 @@ auto make_padded(raft::resources const& res, raft::host_matrix_view(src.extent(1)); - auto const stride = cuvs::neighbors::cagra_required_row_width(dim, 16); + auto const stride = cuvs::core::padded_row_width(dim, 16); auto matrix = raft::make_device_matrix(res, src.extent(0), static_cast(stride)); RAFT_CUDA_TRY(cudaMemsetAsync( @@ -79,8 +79,8 @@ auto make_padded(raft::resources const& res, raft::host_matrix_view(src.extent(0)), stream); raft::resource::sync_stream(res); - cuvs::neighbors::device_padded_dataset_view view( - raft::make_const_mdspan(matrix.view()), dim); + cuvs::core::device_padded_dataset_view view(raft::make_const_mdspan(matrix.view()), + dim); return padded_storage{std::move(matrix), view}; } @@ -88,10 +88,10 @@ auto make_padded(raft::resources const& res, raft::host_matrix_view auto make_merged_storage(raft::resources const& res, int64_t rows, int64_t dim) -> padded_storage { - auto const stride = cuvs::neighbors::cagra_required_row_width(static_cast(dim), 16); + auto const stride = cuvs::core::padded_row_width(static_cast(dim), 16); auto matrix = raft::make_device_matrix(res, rows, static_cast(stride)); - cuvs::neighbors::device_padded_dataset_view view( - raft::make_const_mdspan(matrix.view()), static_cast(dim)); + cuvs::core::device_padded_dataset_view view(raft::make_const_mdspan(matrix.view()), + static_cast(dim)); return padded_storage{std::move(matrix), view}; } @@ -216,11 +216,13 @@ void expect_dataset_order(raft::resources const& res, // honouring that stride rather than as one contiguous block. auto host = raft::make_host_matrix(res, expected.extent(0), expected.extent(1)); auto stream = raft::resource::get_cuda_stream(res); - int64_t const row_stride = static_cast(view.stride()); + int64_t const row_stride = static_cast(view.as_matrix_view().stride()); int64_t const dim = static_cast(view.dim()); for (int64_t row = 0; row < view.n_rows(); ++row) { - raft::copy( - host.data_handle() + row * dim, view.view().data_handle() + row * row_stride, dim, stream); + raft::copy(host.data_handle() + row * dim, + view.as_matrix_view().data_handle() + row * row_stride, + dim, + stream); } raft::resource::sync_stream(res); for (int64_t row = 0; row < expected.extent(0); ++row) { @@ -237,10 +239,10 @@ void expect_zero_padding(raft::resources const& res, auto view = merged.dataset(); int64_t const rows = view.n_rows(); int64_t const dim = static_cast(view.dim()); - int64_t const stride = static_cast(view.stride()); + int64_t const stride = static_cast(view.as_matrix_view().stride()); std::vector host(static_cast(rows * stride)); auto stream = raft::resource::get_cuda_stream(res); - raft::copy(host.data(), view.view().data_handle(), host.size(), stream); + raft::copy(host.data(), view.as_matrix_view().data_handle(), host.size(), stream); raft::resource::sync_stream(res); for (int64_t row = 0; row < rows; ++row) { for (int64_t column = dim; column < stride; ++column) { diff --git a/cpp/tests/neighbors/ann_cagra_bbq.cuh b/cpp/tests/neighbors/ann_cagra_bbq.cuh index a044c8aaa3..0e67560fc2 100644 --- a/cpp/tests/neighbors/ann_cagra_bbq.cuh +++ b/cpp/tests/neighbors/ann_cagra_bbq.cuh @@ -67,7 +67,7 @@ class AnnCagraBbqTest : public ::testing::TestWithParam { protected: /** Quantize the float database on the host and upload the codes. */ - auto quantize_database() -> cuvs::neighbors::device_bbq_dataset + auto quantize_database() -> cuvs::preprocessing::quantize::bbq::device_bbq_dataset { std::vector host_data(static_cast(ps.n_rows) * ps.dim); raft::update_host(host_data.data(), database.data(), host_data.size(), stream_); @@ -263,7 +263,7 @@ class AnnCagraBbqTest : public ::testing::TestWithParam { ASSERT_EQ(restored.graph_degree(), graph_index.graph_degree()); EXPECT_EQ(restored.metric(), graph_index.metric()); EXPECT_EQ(restored.dataset().n_rows(), 0); - EXPECT_TRUE(restored.dataset().quantizers.empty()); + EXPECT_TRUE(restored.dataset().data().quantizers.empty()); auto original = cagra::update_dataset(handle_, std::move(graph_index), device_padded.view); auto reattached = cagra::update_dataset(handle_, std::move(restored), device_padded.view); @@ -286,7 +286,7 @@ class AnnCagraBbqTest : public ::testing::TestWithParam { ASSERT_EQ(index.graph_size(), static_cast(ps.n_rows)); ASSERT_EQ(index.graph_degree(), static_cast(ps.graph_degree)); EXPECT_EQ(index.dataset().n_rows(), 0); - EXPECT_TRUE(index.dataset().quantizers.empty()); + EXPECT_TRUE(index.dataset().data().quantizers.empty()); } /** Only NN-descent graph construction and the four BBQ metrics are accepted. */ diff --git a/cpp/tests/neighbors/ann_hnsw_ace.cuh b/cpp/tests/neighbors/ann_hnsw_ace.cuh index 90a3e46caf..2716e028ad 100644 --- a/cpp/tests/neighbors/ann_hnsw_ace.cuh +++ b/cpp/tests/neighbors/ann_hnsw_ace.cuh @@ -108,7 +108,7 @@ void build_ace_with_workspace(raft::resources const& resources, ace_params.use_disk = true; params.graph_build_params = ace_params; - auto dataset_view = cuvs::neighbors::make_host_standard_dataset_view(dataset); + auto dataset_view = cuvs::core::make_host_standard_dataset_view(dataset); [[maybe_unused]] auto index = cagra::build(resources, params, dataset_view); } diff --git a/cpp/tests/neighbors/ann_scann.cuh b/cpp/tests/neighbors/ann_scann.cuh index 8f6deb30de..824fa44e1a 100644 --- a/cpp/tests/neighbors/ann_scann.cuh +++ b/cpp/tests/neighbors/ann_scann.cuh @@ -185,10 +185,10 @@ class scann_test : public ::testing::TestWithParam { vq_codebook.data_handle(), idx.centers().data_handle(), idx.centers().size(), stream_); auto empty_data = raft::make_device_matrix(handle_, 0, 0); + using vpq_owning_t = cuvs::preprocessing::quantize::pq::device_vpq_dataset; cuvs::preprocessing::quantize::pq::quantizer quantizer{ pq_params, - cuvs::neighbors::device_vpq_dataset{ - std::move(vq_codebook), std::move(pq_codebook_copy), std::move(empty_data)}}; + vpq_owning_t(std::move(empty_data), std::move(vq_codebook), std::move(pq_codebook_copy))}; auto quantized_residuals_device = raft::make_device_matrix(handle_, ps.num_db_vecs, num_subspaces); diff --git a/cpp/tests/neighbors/cagra_padded_build_helpers.cuh b/cpp/tests/neighbors/cagra_padded_build_helpers.cuh index 215b845d9d..e0529bb4fd 100644 --- a/cpp/tests/neighbors/cagra_padded_build_helpers.cuh +++ b/cpp/tests/neighbors/cagra_padded_build_helpers.cuh @@ -20,8 +20,8 @@ namespace cuvs::neighbors::test { */ template struct padded_device_matrix_for_cagra { - std::unique_ptr> owned; - cuvs::neighbors::device_padded_dataset_view view; + std::unique_ptr> owned; + cuvs::core::device_padded_dataset_view view; padded_device_matrix_for_cagra( raft::resources const& res, raft::device_matrix_view src) @@ -31,8 +31,8 @@ struct padded_device_matrix_for_cagra { private: struct build_result { - std::unique_ptr> owned; - cuvs::neighbors::device_padded_dataset_view view; + std::unique_ptr> owned; + cuvs::core::device_padded_dataset_view view; }; // device_padded_dataset_view has no default constructor; fill both members from one build step. @@ -46,10 +46,10 @@ struct padded_device_matrix_for_cagra { -> build_result { using namespace cuvs::neighbors; - if (matrix_row_width_matches_cagra_required(src)) { - return build_result{nullptr, make_device_padded_dataset_view(res, src)}; + if (cuvs::core::matrix_has_padded_row_width(src)) { + return build_result{nullptr, cuvs::core::make_device_padded_dataset_view(res, src)}; } else { - auto own = make_device_padded_dataset(res, src); + auto own = cuvs::core::make_device_padded_dataset(res, src); auto vw = own->as_dataset_view(); return build_result{std::move(own), vw}; } diff --git a/cpp/tests/neighbors/dynamic_batching/test_cagra.cu b/cpp/tests/neighbors/dynamic_batching/test_cagra.cu index 752c43708c..5edacbccf7 100644 --- a/cpp/tests/neighbors/dynamic_batching/test_cagra.cu +++ b/cpp/tests/neighbors/dynamic_batching/test_cagra.cu @@ -20,7 +20,7 @@ auto build_cagra_with_dataset(raft::resources const& res, raft::device_matrix_view dataset) -> cagra::device_padded_index { - auto padded = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + auto padded = cuvs::core::make_device_padded_dataset_view(res, dataset); auto index = cagra::build(res, params, padded); index = cagra::update_dataset(res, std::move(index), padded); return index; diff --git a/cpp/tests/neighbors/mg.cuh b/cpp/tests/neighbors/mg.cuh index 336563d9c4..e1023f0452 100644 --- a/cpp/tests/neighbors/mg.cuh +++ b/cpp/tests/neighbors/mg.cuh @@ -226,7 +226,7 @@ class AnnMGTest : public ::testing::TestWithParam { auto index_dataset = raft::make_host_matrix_view( h_index_dataset.data(), ps.num_db_vecs, ps.dim); - auto index_dataset_view = cuvs::neighbors::make_host_standard_dataset_view(index_dataset); + auto index_dataset_view = cuvs::core::make_host_standard_dataset_view(index_dataset); auto queries = raft::make_host_matrix_view( h_queries.data(), ps.num_queries, ps.dim); auto neighbors = raft::make_host_matrix_view( @@ -248,7 +248,7 @@ class AnnMGTest : public ::testing::TestWithParam { auto index_dataset_device = raft::make_device_matrix_view( d_index_dataset.data(), ps.num_db_vecs, ps.dim); auto padded_index_dataset = - cuvs::neighbors::make_device_padded_dataset_view(clique_, index_dataset_device); + cuvs::core::make_device_padded_dataset_view(clique_, index_dataset_device); auto search_index = cuvs::neighbors::cagra::update_dataset(clique_, std::move(new_index), padded_index_dataset); @@ -387,7 +387,7 @@ class AnnMGTest : public ::testing::TestWithParam { { auto index_dataset = raft::make_device_matrix_view( d_index_dataset.data(), ps.num_db_vecs, ps.dim); - auto standard_view = cuvs::neighbors::make_device_standard_dataset_view(index_dataset); + auto standard_view = cuvs::core::make_device_standard_dataset_view(index_dataset); auto index = cuvs::neighbors::cagra::build(clique_, index_params, standard_view); cuvs::neighbors::cagra::serialize(clique_, index_file.filename, index); } @@ -408,7 +408,7 @@ class AnnMGTest : public ::testing::TestWithParam { auto index_dataset_device = raft::make_device_matrix_view( d_index_dataset.data(), ps.num_db_vecs, ps.dim); auto padded_index_dataset = - cuvs::neighbors::make_device_padded_dataset_view(clique_, index_dataset_device); + cuvs::core::make_device_padded_dataset_view(clique_, index_dataset_device); auto search_index = cuvs::neighbors::cagra::update_dataset( clique_, std::move(distributed_index), padded_index_dataset); @@ -585,7 +585,7 @@ class AnnMGTest : public ::testing::TestWithParam { auto index_dataset = raft::make_host_matrix_view( h_index_dataset.data(), ps.num_db_vecs, ps.dim); - auto index_dataset_view = cuvs::neighbors::make_host_standard_dataset_view(index_dataset); + auto index_dataset_view = cuvs::core::make_host_standard_dataset_view(index_dataset); auto small_batch_query = raft::make_host_matrix_view( h_queries.data(), ps.num_queries, ps.dim); @@ -593,7 +593,7 @@ class AnnMGTest : public ::testing::TestWithParam { auto index_dataset_device = raft::make_device_matrix_view( d_index_dataset.data(), ps.num_db_vecs, ps.dim); auto padded_index_dataset = - cuvs::neighbors::make_device_padded_dataset_view(clique_, index_dataset_device); + cuvs::core::make_device_padded_dataset_view(clique_, index_dataset_device); auto search_index = cuvs::neighbors::cagra::update_dataset(clique_, std::move(index), padded_index_dataset); diff --git a/cpp/tests/neighbors/tiered_index.cu b/cpp/tests/neighbors/tiered_index.cu index 5c71c31319..7382fb011d 100644 --- a/cpp/tests/neighbors/tiered_index.cu +++ b/cpp/tests/neighbors/tiered_index.cu @@ -163,16 +163,15 @@ class ANNTieredIndexTest : public ::testing::TestWithParam if constexpr (std::is_same_v>) { auto full_database_view = raft::make_device_matrix_view( (const value_type*)database.data(), ps.n_rows, ps.dim); - if (cuvs::neighbors::matrix_row_width_matches_cagra_required(full_database_view)) { + if (cuvs::core::matrix_has_padded_row_width(full_database_view)) { auto padded_view = - cuvs::neighbors::make_device_padded_dataset_view(handle_, full_database_view); + cuvs::core::make_device_padded_dataset_view(handle_, full_database_view); auto attached_index = cuvs::neighbors::tiered_index::convert_standard_to_padded_index( handle_, *final_index, padded_view); cuvs::neighbors::tiered_index::search( handle_, search_params, attached_index, queries_view, indices_view, distances_view); } else { - auto padded_dataset = - cuvs::neighbors::make_device_padded_dataset(handle_, full_database_view); + auto padded_dataset = cuvs::core::make_device_padded_dataset(handle_, full_database_view); auto attached_index = cuvs::neighbors::tiered_index::convert_standard_to_padded_index( handle_, *final_index, padded_dataset->as_dataset_view()); cuvs::neighbors::tiered_index::search( diff --git a/cpp/tests/neighbors/vpq_utils.cuh b/cpp/tests/neighbors/vpq_utils.cuh index 23dba3218c..d29f4428c5 100644 --- a/cpp/tests/neighbors/vpq_utils.cuh +++ b/cpp/tests/neighbors/vpq_utils.cuh @@ -5,6 +5,7 @@ #pragma once #include +#include #include #include @@ -45,15 +46,17 @@ __global__ void decode_vpq_dataset_kernel(data_t* const decoded_dataset_ptr, } template -void decode_vpq_dataset(raft::device_matrix_view decoded_dataset, - const cuvs::neighbors::device_vpq_dataset& vpq_dataset, - cudaStream_t cuda_stream) +void decode_vpq_dataset( + raft::device_matrix_view decoded_dataset, + const cuvs::preprocessing::quantize::pq::device_vpq_dataset& vpq_dataset, + cudaStream_t cuda_stream) { const auto dataset_size = decoded_dataset.extent(0); - RAFT_EXPECTS(vpq_dataset.data.extent(0) == dataset_size, "Dataset sizes mismatch"); - RAFT_EXPECTS(vpq_dataset.pq_bits() == 8, + auto const& vpq = vpq_dataset.data(); + RAFT_EXPECTS(vpq.extent(0) == dataset_size, "Dataset sizes mismatch"); + RAFT_EXPECTS(vpq.pq_bits() == 8, "decode_vpq_dataset currently only supports pq_bits == 8 (got %u)", - vpq_dataset.pq_bits()); + vpq.pq_bits()); constexpr uint32_t block_size = 256; constexpr uint32_t warp_size = 32; @@ -63,14 +66,14 @@ void decode_vpq_dataset(raft::device_matrix_view decoded_datase decode_vpq_dataset_kernel <<>>(decoded_dataset.data_handle(), decoded_dataset.stride(0), - vpq_dataset.vq_code_book.data_handle(), - vpq_dataset.vq_code_book.stride(0), - vpq_dataset.pq_code_book.data_handle(), - vpq_dataset.pq_len(), - 1u << vpq_dataset.pq_bits(), + vpq.vq_code_book.data_handle(), + vpq.vq_code_book.stride(0), + vpq.pq_code_book.data_handle(), + vpq.pq_len(), + 1u << vpq.pq_bits(), vpq_dataset.dim(), dataset_size, - vpq_dataset.data.data_handle(), - vpq_dataset.data.stride(0)); + vpq.data_handle(), + vpq.stride(0)); } } // namespace cuvs::neighbors diff --git a/cpp/tests/preprocessing/product_quantization.cu b/cpp/tests/preprocessing/product_quantization.cu index ff44edde20..42ca8c9a61 100644 --- a/cpp/tests/preprocessing/product_quantization.cu +++ b/cpp/tests/preprocessing/product_quantization.cu @@ -322,14 +322,14 @@ TEST(ProductQuantizationTestF, MakeVpqDatasetFromHost) dataset.data_handle()[i] = static_cast(i % 31) / 31.0f; } - cuvs::neighbors::vpq_params params{ + cuvs::preprocessing::quantize::pq::vpq_params params{ .pq_bits = 4, .pq_dim = 4, .vq_n_centers = 1, .kmeans_n_iters = 2}; auto vpq = make_vpq_dataset(handle, params, raft::make_const_mdspan(dataset.view())); raft::resource::sync_stream(handle); EXPECT_EQ(vpq.n_rows(), n_rows); EXPECT_EQ(vpq.dim(), dim); - EXPECT_NE(vpq.data.data_handle(), nullptr); + EXPECT_NE(vpq.data().data_handle(), nullptr); } TEST(ProductQuantizationTestF, MakeVpqDatasetFromPaddedView) @@ -351,18 +351,18 @@ TEST(ProductQuantizationTestF, MakeVpqDatasetFromPaddedView) host_rows.data_handle(), host_rows.size(), raft::resource::get_cuda_stream(handle)); - cuvs::neighbors::device_padded_dataset_view padded( + cuvs::core::device_padded_dataset_view padded( raft::make_device_matrix_view(device_rows.data_handle(), n_rows, stride), dim); - cuvs::neighbors::vpq_params params{ + cuvs::preprocessing::quantize::pq::vpq_params params{ .pq_bits = 4, .pq_dim = 4, .vq_n_centers = 1, .kmeans_n_iters = 2}; auto vpq = make_vpq_dataset(handle, params, padded); raft::resource::sync_stream(handle); EXPECT_EQ(vpq.n_rows(), n_rows); EXPECT_EQ(vpq.dim(), dim); - EXPECT_NE(vpq.data.data_handle(), nullptr); + EXPECT_NE(vpq.data().data_handle(), nullptr); } // Define test cases with different parameters diff --git a/examples/cpp/src/cagra_bloom_filter_example.cu b/examples/cpp/src/cagra_bloom_filter_example.cu index f84ab08fd9..444ae13f9e 100644 --- a/examples/cpp/src/cagra_bloom_filter_example.cu +++ b/examples/cpp/src/cagra_bloom_filter_example.cu @@ -71,7 +71,7 @@ int main() std::cout << "Building CAGRA index" << std::endl; auto padded = - cuvs::neighbors::make_device_padded_dataset_view(res, raft::make_const_mdspan(dataset.view())); + cuvs::core::make_device_padded_dataset_view(res, raft::make_const_mdspan(dataset.view())); auto index = cuvs::neighbors::cagra::build(res, index_params, padded); index = cuvs::neighbors::cagra::update_dataset(res, std::move(index), padded); diff --git a/examples/cpp/src/cagra_example.cu b/examples/cpp/src/cagra_example.cu index 5fa5c55f6e..4e52c07207 100644 --- a/examples/cpp/src/cagra_example.cu +++ b/examples/cpp/src/cagra_example.cu @@ -34,7 +34,7 @@ void cagra_build_search_simple(raft::device_resources const& dev_resources, cagra::index_params index_params; std::cout << "Building CAGRA index (search graph)" << std::endl; - auto padded = cuvs::neighbors::make_device_padded_dataset_view(dev_resources, dataset); + auto padded = cuvs::core::make_device_padded_dataset_view(dev_resources, dataset); auto index = cagra::build(dev_resources, index_params, padded); index = cagra::update_dataset(dev_resources, std::move(index), padded); diff --git a/examples/cpp/src/cagra_filter_udf_example.cu b/examples/cpp/src/cagra_filter_udf_example.cu index 23b3a316cc..a4ade31e3d 100644 --- a/examples/cpp/src/cagra_filter_udf_example.cu +++ b/examples/cpp/src/cagra_filter_udf_example.cu @@ -146,7 +146,7 @@ int main() index_params.intermediate_graph_degree); std::cout << "Building CAGRA index" << std::endl; - auto padded = cuvs::neighbors::make_device_padded_dataset_view(res, dataset.view()); + auto padded = cuvs::core::make_device_padded_dataset_view(res, dataset.view()); auto index = cuvs::neighbors::cagra::build(res, index_params, padded); index = cuvs::neighbors::cagra::update_dataset(res, std::move(index), padded); diff --git a/examples/cpp/src/cagra_hnsw_ace_example.cu b/examples/cpp/src/cagra_hnsw_ace_example.cu index e27d4dfa47..df57a9b5e7 100644 --- a/examples/cpp/src/cagra_hnsw_ace_example.cu +++ b/examples/cpp/src/cagra_hnsw_ace_example.cu @@ -68,7 +68,7 @@ void cagra_build_search_ace(raft::device_resources const& dev_resources, dataset_host.data_handle(), dataset_host.extent(0), dataset_host.extent(1)); // Wrap in a host_padded_dataset_view. ACE graph construction is host-side CPU work and does not // require CUDA row-alignment; construct the view directly to avoid the alignment check. - cuvs::neighbors::host_padded_dataset_view host_padded_view( + cuvs::core::host_padded_dataset_view host_padded_view( dataset_host_view, static_cast(dataset_host_view.extent(1))); std::cout << "Building CAGRA index (search graph)" << std::endl; @@ -87,7 +87,7 @@ void cagra_build_search_ace(raft::device_resources const& dev_resources, hnsw_params.hierarchy = hnsw::HnswHierarchy::GPU; // Offload hierarchy construction to GPU std::unique_ptr> hnsw_index; - std::unique_ptr> padded_owner; + std::unique_ptr> padded_owner; if (ace_host_index.dataset_fd().has_value()) { // Disk ACE path: ACE artifacts (dataset, graph, mapping) live on disk. Transfer file // descriptors to a device index so from_cagra can serialize to hnsw_index.bin on disk. @@ -98,7 +98,7 @@ void cagra_build_search_ace(raft::device_resources const& dev_resources, } else { // In-memory ACE path: graph is in host memory. Upload the original dataset to device and // attach it before from_cagra builds the HNSW hierarchy in memory. - padded_owner = cuvs::neighbors::make_device_padded_dataset(dev_resources, dataset_host_view); + padded_owner = cuvs::core::make_device_padded_dataset(dev_resources, dataset_host_view); auto device_index = cagra::update_dataset( dev_resources, std::move(ace_host_index), padded_owner->as_dataset_view()); hnsw_index = hnsw::from_cagra(dev_resources, hnsw_params, device_index, dataset_host_view); diff --git a/examples/cpp/src/cagra_persistent_example.cu b/examples/cpp/src/cagra_persistent_example.cu index ccc5930efd..24bbd6e13e 100644 --- a/examples/cpp/src/cagra_persistent_example.cu +++ b/examples/cpp/src/cagra_persistent_example.cu @@ -70,7 +70,7 @@ void cagra_build_search_variants(raft::device_resources const& res, cagra::index_params index_params; std::cout << "Building CAGRA index (search graph)" << std::endl; - auto padded = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + auto padded = cuvs::core::make_device_padded_dataset_view(res, dataset); auto index = cagra::build(res, index_params, padded); index = cagra::update_dataset(res, std::move(index), padded); diff --git a/examples/cpp/src/cagra_roaring_bitmap_filter_example.cu b/examples/cpp/src/cagra_roaring_bitmap_filter_example.cu index 929fcaa5a5..e5261aef72 100644 --- a/examples/cpp/src/cagra_roaring_bitmap_filter_example.cu +++ b/examples/cpp/src/cagra_roaring_bitmap_filter_example.cu @@ -45,7 +45,7 @@ int main() index_params.graph_degree = 32; index_params.intermediate_graph_degree = 64; auto padded = - cuvs::neighbors::make_device_padded_dataset_view(res, raft::make_const_mdspan(dataset.view())); + cuvs::core::make_device_padded_dataset_view(res, raft::make_const_mdspan(dataset.view())); auto index = cuvs::neighbors::cagra::build(res, index_params, padded); // Each owner is independently reusable. The filter supplies the query-to-allowlist mapping by diff --git a/examples/cpp/src/dynamic_batching_example.cu b/examples/cpp/src/dynamic_batching_example.cu index 2745b5e4a7..b8666be49b 100644 --- a/examples/cpp/src/dynamic_batching_example.cu +++ b/examples/cpp/src/dynamic_batching_example.cu @@ -116,7 +116,7 @@ void dynamic_batching_example(raft::resources const& res, cagra::index_params orig_index_params; std::cout << "Building CAGRA index (search graph)" << std::endl; - auto padded = cuvs::neighbors::make_device_padded_dataset_view(res, dataset); + auto padded = cuvs::core::make_device_padded_dataset_view(res, dataset); auto orig_index = cagra::build(res, orig_index_params, padded); orig_index = cagra::update_dataset(res, std::move(orig_index), padded); diff --git a/fern/docs.yml b/fern/docs.yml index 77ee0dae2a..b055d844f1 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -358,6 +358,8 @@ navigation: path: "./pages/cpp_api/cpp-api-common-types-copy-serialization-and-utility-apis.md" - page: "Errors and Logging" path: "./pages/cpp_api/cpp-api-common-types-errors-and-logging.md" + - page: "Core Dataset" + path: "./pages/cpp_api/cpp-api-core-dataset.md" - page: "Core Roaring Allowlist" path: "./pages/cpp_api/cpp-api-core-roaring-allowlist.md" - page: "Distance Distance" diff --git a/fern/pages/cpp_api/cpp-api-core-dataset.md b/fern/pages/cpp_api/cpp-api-core-dataset.md new file mode 100644 index 0000000000..8c82369cd2 --- /dev/null +++ b/fern/pages/cpp_api/cpp-api-core-dataset.md @@ -0,0 +1,109 @@ +--- +slug: api-reference/cpp-api-core-dataset +--- + +# Dataset + +_Source header: `cuvs/core/dataset.hpp`_ + +## Types + + +### core::dataset_view + +Non-owning dataset view: holds only the view-shaped payload. Deliberately not derived from + +`dataset` -- a view type holds "all view state" with no inheritance and no shared ownership tying it to the owning type. Reuses the same `get_n_rows`/`get_dim` spec functions as `dataset`, fed the view payload instead of the owning one. + +```cpp +template +struct dataset_view; +``` + + +### core::dataset + +Owning dataset: value-held payload (no shared_ptr -- exclusive ownership). Every member is a + +one-line forward to `spec_type::get_*` or to the payload; all per-kind state and logic lives in the spec's `data_type`, never inside this struct. + +```cpp +template +struct dataset; +``` + + +### core::is_padded_dataset + +Owning-side kind traits (true for both `dataset<...>` and `dataset_view<...>` of that kind). + +```cpp +template +struct is_padded_dataset; +``` + + +### core::is_dataset_view + +True for any `dataset_view<...>` specialization. Evaluates to `false` (never a hard error) for + +everything else, e.g. a plain mdspan passed to a deprecated `build(matrix_view)` overload. + +```cpp +template +struct is_dataset_view; +``` + + +### core::dataset_view_has_spec + +True when `V` is a `dataset_view` whose spec satisfies the predicate `SpecPred<SpecT>::value`. + +This is how a kind that lives outside this header classifies its own views. + +```cpp +template typename SpecPred> +struct dataset_view_has_spec; +``` + + +### core::dataset_view_is_device_accessible + +True when the dataset view accessor is device-accessible. + +```cpp +template +struct dataset_view_is_device_accessible; +``` + + +### core::with_accessor + +Generic accessor retargeting while preserving the spec kind and value/index types: + +`dataset<T, IdxT, SpecT<..., OldAccessor>> -> dataset<T, IdxT, SpecT<..., NewAccessor>>` `dataset_view<T, IdxT, SpecT<..., OldAccessor>> -> dataset_view<T, IdxT, SpecT<..., NewAccessor>>` Every spec provides `rebind_accessor<NewAccessor>` for this, so this header does not need to know about any particular kind. + +```cpp +template +struct with_accessor; +``` + + +### core::to_device_accessor + +Map any host accessor to its device counterpart (same payload policy). + +```cpp +template +struct to_device_accessor; +``` + + +### core::device_counterpart + +Maps a host dataset view type to its device-resident counterpart. + +```cpp +template +struct device_counterpart; +``` diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 2560f59ef1..b6464fe4b2 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -31,24 +31,32 @@ enum class search_algo { | `MULTI_KERNEL` | `2` | | `AUTO` | `100` | -## CAGRA index build parameters +## Types + + +### neighbors::graph_build_params::iterative_search_params + +Parameters for the iterative CAGRA graph build algorithm. + +Inherits from cagra::search_params so that all search tuning knobs (search_width, max_iterations, itopk_size, etc.) are available for controlling the search-and-optimize loop during graph construction. The defaults are tuned for the build loop (e.g. search_width=1, max_iterations=8) and may differ from the regular search defaults. + +```cpp +struct iterative_search_params; +``` - -### neighbors::vpq_params + +### neighbors::graph_build_params::ace_params -Parameters for VPQ compression. +Specialized parameters for ACE (Augmented Core Extraction) graph build ```cpp -struct vpq_params { - uint32_t pq_bits; - uint32_t pq_dim; - uint32_t vq_n_centers; - uint32_t kmeans_n_iters; - double vq_kmeans_trainset_fraction; - double pq_kmeans_trainset_fraction; - cuvs::cluster::kmeans::kmeans_type pq_kmeans_type; - uint32_t max_train_points_per_pq_code; - uint32_t max_train_points_per_vq_cluster; +struct ace_params { + size_t npartitions; + size_t ef_construction; + std::string build_dir; + bool use_disk; + double max_host_memory_gb; + double max_gpu_memory_gb; }; ``` @@ -56,15 +64,14 @@ struct vpq_params { | Name | Type | Description | | --- | --- | --- | -| `pq_bits` | `uint32_t` | The bit length of the vector element after compression by PQ.

Possible values: [4, 5, 6, 7, 8].

Hint: the smaller the 'pq_bits', the smaller the index size and the better the search performance, but the lower the recall. | -| `pq_dim` | `uint32_t` | The dimensionality of the vector after compression by PQ. When zero, an optimal value is selected using a heuristic.

TODO: at the moment `dim` must be a multiple `pq_dim`. | -| `vq_n_centers` | `uint32_t` | Vector Quantization (VQ) codebook size - number of "coarse cluster centers". When zero, an optimal value is selected using a heuristic. | -| `kmeans_n_iters` | `uint32_t` | The number of iterations searching for kmeans centers (both VQ & PQ phases). | -| `vq_kmeans_trainset_fraction` | `double` | The fraction of data to use during iterative kmeans building (VQ phase). When zero, an optimal value is selected using a heuristic. | -| `pq_kmeans_trainset_fraction` | `double` | The fraction of data to use during iterative kmeans building (PQ phase). When zero, an optimal value is selected using a heuristic. | -| `pq_kmeans_type` | [`cuvs::cluster::kmeans::kmeans_type`](/api-reference/cpp-api-cluster-kmeans#cluster-kmeans-kmeans-type) | Type of k-means algorithm for PQ training. Balanced k-means tends to be faster than regular k-means for PQ training, for problem sets where the number of points per cluster are approximately equal. Regular k-means may be better for skewed cluster distributions. | -| `max_train_points_per_pq_code` | `uint32_t` | The max number of data points to use per PQ code during PQ codebook training. Using more data points per PQ code may increase the quality of PQ codebook but may also increase the build time. We will use `pq_n_centers * max_train_points_per_pq_code` training points to train each PQ codebook. | -| `max_train_points_per_vq_cluster` | `uint32_t` | The max number of data points to use per VQ cluster during training. | +| `npartitions` | `size_t` | Number of partitions for ACE (Augmented Core Extraction) partitioned build.

When set to 0 (default), the number of partitions is automatically derived based on available host and GPU memory to maximize partition size while ensuring the build fits in memory.

Small values might improve recall but potentially degrade performance and increase memory usage. Partitions should not be too small to prevent issues in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the partition sizes (up to 3x in our tests).

If the specified number of partitions results in partitions that exceed available memory, the value will be automatically increased to fit memory constraints and a warning will be issued. | +| `ef_construction` | `size_t` | The index quality for the ACE build.

Bigger values increase the index quality. At some point, increasing this will no longer improve the quality. | +| `build_dir` | `std::string` | Directory to store ACE build artifacts (e.g., KNN graph, optimized graph).

Used when `use_disk` is true or when the graph does not fit in host and GPU memory. This should be the fastest disk in the system and hold enough space for twice the dataset, final graph, and label mapping. The directory may already exist, but ACE's named artifacts must not already exist. Simultaneous builds must use different directories. On failure, ACE removes only artifacts it created and never deletes unrelated directory contents. | +| `use_disk` | `bool` | Whether to use disk-based storage for ACE build.

When true, enables disk-based operations for memory-efficient graph construction. | +| `max_host_memory_gb` | `double` | Maximum host memory to use for ACE build in GiB.

When set to 0 (default), uses available host memory. When set to a positive value, limits host memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. | +| `max_gpu_memory_gb` | `double` | Maximum GPU memory to use for ACE build in GiB.

When set to 0 (default), uses available GPU memory. When set to a positive value, limits GPU memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. | + +## CAGRA index build parameters ### graph_build_params_t @@ -182,46 +189,6 @@ Usage example: `static cagra::index_params` -## Types - - -### neighbors::graph_build_params::iterative_search_params - -Parameters for the iterative CAGRA graph build algorithm. - -Inherits from cagra::search_params so that all search tuning knobs (search_width, max_iterations, itopk_size, etc.) are available for controlling the search-and-optimize loop during graph construction. The defaults are tuned for the build loop (e.g. search_width=1, max_iterations=8) and may differ from the regular search defaults. - -```cpp -struct iterative_search_params; -``` - - -### neighbors::graph_build_params::ace_params - -Specialized parameters for ACE (Augmented Core Extraction) graph build - -```cpp -struct ace_params { - size_t npartitions; - size_t ef_construction; - std::string build_dir; - bool use_disk; - double max_host_memory_gb; - double max_gpu_memory_gb; -}; -``` - -**Fields** - -| Name | Type | Description | -| --- | --- | --- | -| `npartitions` | `size_t` | Number of partitions for ACE (Augmented Core Extraction) partitioned build.

When set to 0 (default), the number of partitions is automatically derived based on available host and GPU memory to maximize partition size while ensuring the build fits in memory.

Small values might improve recall but potentially degrade performance and increase memory usage. Partitions should not be too small to prevent issues in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the partition sizes (up to 3x in our tests).

If the specified number of partitions results in partitions that exceed available memory, the value will be automatically increased to fit memory constraints and a warning will be issued. | -| `ef_construction` | `size_t` | The index quality for the ACE build.

Bigger values increase the index quality. At some point, increasing this will no longer improve the quality. | -| `build_dir` | `std::string` | Directory to store ACE build artifacts (e.g., KNN graph, optimized graph).

Used when `use_disk` is true or when the graph does not fit in host and GPU memory. This should be the fastest disk in the system and hold enough space for twice the dataset, final graph, and label mapping. The directory may already exist, but ACE's named artifacts must not already exist. Simultaneous builds must use different directories. On failure, ACE removes only artifacts it created and never deletes unrelated directory contents. | -| `use_disk` | `bool` | Whether to use disk-based storage for ACE build.

When true, enables disk-based operations for memory-efficient graph construction. | -| `max_host_memory_gb` | `double` | Maximum host memory to use for ACE build in GiB.

When set to 0 (default), uses available host memory. When set to a positive value, limits host memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. | -| `max_gpu_memory_gb` | `double` | Maximum GPU memory to use for ACE build in GiB.

When set to 0 (default), uses available GPU memory. When set to a positive value, limits GPU memory usage to the specified amount. Useful for testing or when running alongside other memory-intensive processes. | - ## CAGRA index extend parameters @@ -253,7 +220,8 @@ The index stores the dataset and a kNN graph in device memory. ```cpp template > +cuvs::core::dataset_like DatasetViewT = +cuvs::core::device_padded_dataset_view> struct index; ``` @@ -443,7 +411,7 @@ Construct a graph-only index with a zero-row dataset view placeholder. ```cpp explicit index(raft::resources const& res, cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded) -requires(cuvs::neighbors::ann_dataset_view) +requires(cuvs::core::dataset_like) : cuvs::neighbors::index(),; ``` @@ -676,7 +644,7 @@ Build directly from a device VPQ dataset view with FP16 codebooks. ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_vpq_dataset_view const& dataset) +cuvs::preprocessing::quantize::pq::device_vpq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_pq_index; ``` @@ -690,13 +658,13 @@ The returned index accepts float queries and stores a non-owning copy of `datase | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_vpq_dataset_view const&` | device VPQ dataset view | +| `dataset` | in | `cuvs::preprocessing::quantize::pq::device_vpq_dataset_view const&` | device VPQ dataset view | **Returns** `cuvs::neighbors::cagra::device_pq_index` -built `index<float, uint32_t, device_vpq_dataset_view<half, int64_t>>` +built `index<float, uint32_t, cuvs::preprocessing::quantize::pq::device_vpq_dataset_view<half, int64_t>>` **Additional overload:** `neighbors::cagra::build` @@ -705,7 +673,7 @@ Build from a device padded dataset view (`float`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_padded_dataset_view const& dataset) +cuvs::core::device_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_padded_index; ``` @@ -715,7 +683,7 @@ cuvs::neighbors::device_padded_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_padded_dataset_view const&` | device padded dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::device_padded_dataset_view const&` | device padded dataset view [n_rows, dim] | **Returns** @@ -730,7 +698,7 @@ Build from a device standard dataset view (`float`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_standard_dataset_view const& dataset) +cuvs::core::device_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_standard_index; ``` @@ -740,7 +708,7 @@ cuvs::neighbors::device_standard_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_standard_dataset_view const&` | device standard dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::device_standard_dataset_view const&` | device standard dataset view [n_rows, dim] | **Returns** @@ -755,7 +723,7 @@ Build from a host padded dataset view (`float`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::host_padded_dataset_view const& dataset) +cuvs::core::host_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_padded_index; ``` @@ -765,7 +733,7 @@ cuvs::neighbors::host_padded_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::host_padded_dataset_view const&` | host padded dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::host_padded_dataset_view const&` | host padded dataset view [n_rows, dim] | **Returns** @@ -780,7 +748,7 @@ Build from a host standard dataset view (`float`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::host_standard_dataset_view const& dataset) +cuvs::core::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; ``` @@ -790,7 +758,7 @@ cuvs::neighbors::host_standard_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::host_standard_dataset_view const&` | host standard dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::host_standard_dataset_view const&` | host standard dataset view [n_rows, dim] | **Returns** @@ -805,7 +773,7 @@ Build from a device padded dataset view (`half`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_padded_dataset_view const& dataset) +cuvs::core::device_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_padded_index; ``` @@ -815,7 +783,7 @@ cuvs::neighbors::device_padded_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_padded_dataset_view const&` | device padded dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::device_padded_dataset_view const&` | device padded dataset view [n_rows, dim] | **Returns** @@ -830,7 +798,7 @@ Build from a device standard dataset view (`half`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_standard_dataset_view const& dataset) +cuvs::core::device_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_standard_index; ``` @@ -840,7 +808,7 @@ cuvs::neighbors::device_standard_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_standard_dataset_view const&` | device standard dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::device_standard_dataset_view const&` | device standard dataset view [n_rows, dim] | **Returns** @@ -855,7 +823,7 @@ Build from a host padded dataset view (`half`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::host_padded_dataset_view const& dataset) +cuvs::core::host_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_padded_index; ``` @@ -865,7 +833,7 @@ cuvs::neighbors::host_padded_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::host_padded_dataset_view const&` | host padded dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::host_padded_dataset_view const&` | host padded dataset view [n_rows, dim] | **Returns** @@ -880,7 +848,7 @@ Build from a host standard dataset view (`half`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::host_standard_dataset_view const& dataset) +cuvs::core::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; ``` @@ -890,7 +858,7 @@ cuvs::neighbors::host_standard_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::host_standard_dataset_view const&` | host standard dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::host_standard_dataset_view const&` | host standard dataset view [n_rows, dim] | **Returns** @@ -905,7 +873,7 @@ Build from a device padded dataset view (`int8_t`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_padded_dataset_view const& dataset) +cuvs::core::device_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_padded_index; ``` @@ -915,7 +883,7 @@ cuvs::neighbors::device_padded_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_padded_dataset_view const&` | device padded dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::device_padded_dataset_view const&` | device padded dataset view [n_rows, dim] | **Returns** @@ -930,7 +898,7 @@ Build from a device standard dataset view (`int8_t`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_standard_dataset_view const& dataset) +cuvs::core::device_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_standard_index; ``` @@ -940,7 +908,7 @@ cuvs::neighbors::device_standard_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_standard_dataset_view const&` | device standard dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::device_standard_dataset_view const&` | device standard dataset view [n_rows, dim] | **Returns** @@ -955,7 +923,7 @@ Build from a host padded dataset view (`int8_t`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::host_padded_dataset_view const& dataset) +cuvs::core::host_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_padded_index; ``` @@ -965,7 +933,7 @@ cuvs::neighbors::host_padded_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::host_padded_dataset_view const&` | host padded dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::host_padded_dataset_view const&` | host padded dataset view [n_rows, dim] | **Returns** @@ -980,7 +948,7 @@ Build from a host standard dataset view (`int8_t`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::host_standard_dataset_view const& dataset) +cuvs::core::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; ``` @@ -990,7 +958,7 @@ cuvs::neighbors::host_standard_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::host_standard_dataset_view const&` | host standard dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::host_standard_dataset_view const&` | host standard dataset view [n_rows, dim] | **Returns** @@ -1005,7 +973,7 @@ Build from a device padded dataset view (`uint8_t`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_padded_dataset_view const& dataset) +cuvs::core::device_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_padded_index; ``` @@ -1015,7 +983,7 @@ cuvs::neighbors::device_padded_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_padded_dataset_view const&` | device padded dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::device_padded_dataset_view const&` | device padded dataset view [n_rows, dim] | **Returns** @@ -1030,7 +998,7 @@ Build from a device standard dataset view (`uint8_t`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_standard_dataset_view const& dataset) +cuvs::core::device_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_standard_index; ``` @@ -1040,7 +1008,7 @@ cuvs::neighbors::device_standard_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_standard_dataset_view const&` | device standard dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::device_standard_dataset_view const&` | device standard dataset view [n_rows, dim] | **Returns** @@ -1055,7 +1023,7 @@ Build from a host padded dataset view (`uint8_t`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::host_padded_dataset_view const& dataset) +cuvs::core::host_padded_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_padded_index; ``` @@ -1065,7 +1033,7 @@ cuvs::neighbors::host_padded_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::host_padded_dataset_view const&` | host padded dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::host_padded_dataset_view const&` | host padded dataset view [n_rows, dim] | **Returns** @@ -1080,7 +1048,7 @@ Build from a host standard dataset view (`uint8_t`). ```cpp auto build(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::host_standard_dataset_view const& dataset) +cuvs::core::host_standard_dataset_view const& dataset) -> cuvs::neighbors::cagra::host_standard_index; ``` @@ -1090,7 +1058,7 @@ cuvs::neighbors::host_standard_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::host_standard_dataset_view const&` | host standard dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::core::host_standard_dataset_view const&` | host standard dataset view [n_rows, dim] | **Returns** @@ -1103,9 +1071,10 @@ built `host_standard_index<uint8_t, uint32_t>` Build from a device BBQ-quantized dataset view. ```cpp -auto build(raft::resources const& res, +auto build( +raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_bbq_dataset_view const& dataset) +cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_bbq_index; ``` @@ -1119,7 +1088,7 @@ The returned index cannot be searched: CAGRA has no BBQ search kernels. Call the | --- | --- | --- | --- | | `res` | in | `raft::resources const&` | raft resources | | `params` | in | `const cuvs::neighbors::cagra::index_params&` | CAGRA index build parameters | -| `dataset` | in | `cuvs::neighbors::device_bbq_dataset_view const&` | device BBQ dataset view [n_rows, dim] | +| `dataset` | in | `cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const&` | device BBQ dataset view [n_rows, dim] | **Returns** @@ -1129,12 +1098,13 @@ built `device_bbq_index<float, uint32_t>` **Additional overload:** `neighbors::cagra::build` -cuvs::neighbors::device_bbq_dataset_view<float, int64_t> const& dataset) +cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view<float, int64_t> const& dataset) ```cpp -auto build(raft::resources const& res, +auto build( +raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_bbq_dataset_view const& dataset) +cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_bbq_index; ``` @@ -1144,7 +1114,7 @@ cuvs::neighbors::device_bbq_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | | `raft::resources const&` | | | `params` | | `const cuvs::neighbors::cagra::index_params&` | | -| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view const&` | | +| `dataset` | | `cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const&` | | **Returns** @@ -1152,12 +1122,13 @@ cuvs::neighbors::device_bbq_dataset_view const& dataset) **Additional overload:** `neighbors::cagra::build` -cuvs::neighbors::device_bbq_dataset_view<float, int64_t> const& dataset) +cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view<float, int64_t> const& dataset) ```cpp -auto build(raft::resources const& res, +auto build( +raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_bbq_dataset_view const& dataset) +cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_bbq_index; ``` @@ -1167,7 +1138,7 @@ cuvs::neighbors::device_bbq_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | | `raft::resources const&` | | | `params` | | `const cuvs::neighbors::cagra::index_params&` | | -| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view const&` | | +| `dataset` | | `cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const&` | | **Returns** @@ -1175,12 +1146,13 @@ cuvs::neighbors::device_bbq_dataset_view const& dataset) **Additional overload:** `neighbors::cagra::build` -cuvs::neighbors::device_bbq_dataset_view<float, int64_t> const& dataset) +cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view<float, int64_t> const& dataset) ```cpp -auto build(raft::resources const& res, +auto build( +raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, -cuvs::neighbors::device_bbq_dataset_view const& dataset) +cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const& dataset) -> cuvs::neighbors::cagra::device_bbq_index; ``` @@ -1190,7 +1162,7 @@ cuvs::neighbors::device_bbq_dataset_view const& dataset) | --- | --- | --- | --- | | `res` | | `raft::resources const&` | | | `params` | | `const cuvs::neighbors::cagra::index_params&` | | -| `dataset` | | `cuvs::neighbors::device_bbq_dataset_view const&` | | +| `dataset` | | `cuvs::preprocessing::quantize::bbq::device_bbq_dataset_view const&` | | **Returns** @@ -1206,7 +1178,7 @@ Add new vectors to a CAGRA index ```cpp void extend(raft::resources const& handle, const cagra::extend_params& params, -cuvs::neighbors::device_padded_dataset_view extended_dataset, +cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row, cuvs::neighbors::cagra::device_padded_index& idx); ``` @@ -1224,7 +1196,7 @@ Usage example: | --- | --- | --- | --- | | `handle` | in | `raft::resources const&` | raft resources | | `params` | in | [`const cagra::extend_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-extend-params) | extend params | -| `extended_dataset` | in | `cuvs::neighbors::device_padded_dataset_view` | caller-owned device-padded view already containing old \|\| new rows | +| `extended_dataset` | in | `cuvs::core::device_padded_dataset_view` | caller-owned device-padded view already containing old \|\| new rows | | `new_start_row` | in | `int64_t` | row index where the additional vectors begin (must equal `idx.size()`) | | `idx` | in,out | `cuvs::neighbors::cagra::device_padded_index&` | CAGRA index; graph is extended and dataset view is rebound | @@ -1239,7 +1211,7 @@ Add new vectors to a CAGRA index. See the float overload for the full contract. ```cpp void extend(raft::resources const& handle, const cagra::extend_params& params, -cuvs::neighbors::device_padded_dataset_view extended_dataset, +cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row, cuvs::neighbors::cagra::device_padded_index& idx); ``` @@ -1250,7 +1222,7 @@ cuvs::neighbors::cagra::device_padded_index& idx); | --- | --- | --- | --- | | `handle` | | `raft::resources const&` | | | `params` | | [`const cagra::extend_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-extend-params) | | -| `extended_dataset` | | `cuvs::neighbors::device_padded_dataset_view` | | +| `extended_dataset` | | `cuvs::core::device_padded_dataset_view` | | | `new_start_row` | | `int64_t` | | | `idx` | | `cuvs::neighbors::cagra::device_padded_index&` | | @@ -1265,7 +1237,7 @@ Add new vectors to a CAGRA index. See the float overload for the full contract. ```cpp void extend(raft::resources const& handle, const cagra::extend_params& params, -cuvs::neighbors::device_padded_dataset_view extended_dataset, +cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row, cuvs::neighbors::cagra::device_padded_index& idx); ``` @@ -1276,7 +1248,7 @@ cuvs::neighbors::cagra::device_padded_index& idx); | --- | --- | --- | --- | | `handle` | | `raft::resources const&` | | | `params` | | [`const cagra::extend_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-extend-params) | | -| `extended_dataset` | | `cuvs::neighbors::device_padded_dataset_view` | | +| `extended_dataset` | | `cuvs::core::device_padded_dataset_view` | | | `new_start_row` | | `int64_t` | | | `idx` | | `cuvs::neighbors::cagra::device_padded_index&` | | @@ -1291,7 +1263,7 @@ Add new vectors to a CAGRA index. See the float overload for the full contract. ```cpp void extend(raft::resources const& handle, const cagra::extend_params& params, -cuvs::neighbors::device_padded_dataset_view extended_dataset, +cuvs::core::device_padded_dataset_view extended_dataset, int64_t new_start_row, cuvs::neighbors::cagra::device_padded_index& idx); ``` @@ -1302,7 +1274,7 @@ cuvs::neighbors::cagra::device_padded_index& idx); | --- | --- | --- | --- | | `handle` | | `raft::resources const&` | | | `params` | | [`const cagra::extend_params&`](/api-reference/cpp-api-neighbors-cagra#neighbors-cagra-extend-params) | | -| `extended_dataset` | | `cuvs::neighbors::device_padded_dataset_view` | | +| `extended_dataset` | | `cuvs::core::device_padded_dataset_view` | | | `new_start_row` | | `int64_t` | | | `idx` | | `cuvs::neighbors::cagra::device_padded_index&` | | @@ -1374,7 +1346,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` Experimental, both the API and the serialization format are subject to change. @@ -1386,7 +1358,7 @@ Experimental, both the API and the serialization format are subject to change. | `handle` | in | `raft::resources const&` | the raft handle | | `filename` | in | `const std::string&` | the name of the file that stores the index | | `index` | out | `cuvs::neighbors::cagra::device_padded_index*` | the cagra index | -| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the file includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | +| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the file includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | **Returns** @@ -1427,7 +1399,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` Experimental, both the API and the serialization format are subject to change. @@ -1439,7 +1411,7 @@ Experimental, both the API and the serialization format are subject to change. | `handle` | in | `raft::resources const&` | the raft handle | | `is` | in | `std::istream&` | input stream | | `index` | out | `cuvs::neighbors::cagra::device_padded_index*` | the cagra index | -| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the stream includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | +| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the stream includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | **Returns** @@ -1480,7 +1452,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` Experimental, both the API and the serialization format are subject to change. @@ -1492,7 +1464,7 @@ Experimental, both the API and the serialization format are subject to change. | `handle` | in | `raft::resources const&` | the raft handle | | `filename` | in | `const std::string&` | the name of the file that stores the index | | `index` | out | `cuvs::neighbors::cagra::device_padded_index*` | the cagra index | -| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the file includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | +| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the file includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | **Returns** @@ -1533,7 +1505,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` Experimental, both the API and the serialization format are subject to change. @@ -1545,7 +1517,7 @@ Experimental, both the API and the serialization format are subject to change. | `handle` | in | `raft::resources const&` | the raft handle | | `is` | in | `std::istream&` | input stream | | `index` | out | `cuvs::neighbors::cagra::device_padded_index*` | the cagra index | -| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the stream includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | +| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the stream includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | **Returns** @@ -1586,7 +1558,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` Experimental, both the API and the serialization format are subject to change. @@ -1598,7 +1570,7 @@ Experimental, both the API and the serialization format are subject to change. | `handle` | in | `raft::resources const&` | the raft handle | | `filename` | in | `const std::string&` | the name of the file that stores the index | | `index` | out | `cuvs::neighbors::cagra::device_padded_index*` | the cagra index | -| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the file includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | +| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the file includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | **Returns** @@ -1639,7 +1611,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` Experimental, both the API and the serialization format are subject to change. @@ -1651,7 +1623,7 @@ Experimental, both the API and the serialization format are subject to change. | `handle` | in | `raft::resources const&` | the raft handle | | `is` | in | `std::istream&` | input stream | | `index` | out | `cuvs::neighbors::cagra::device_padded_index*` | the cagra index | -| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the stream includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | +| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the stream includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | **Returns** @@ -1692,7 +1664,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::device_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` Experimental, both the API and the serialization format are subject to change. @@ -1704,7 +1676,7 @@ Experimental, both the API and the serialization format are subject to change. | `handle` | in | `raft::resources const&` | the raft handle | | `filename` | in | `const std::string&` | the name of the file that stores the index | | `index` | out | `cuvs::neighbors::cagra::device_padded_index*` | the cagra index | -| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the file includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | +| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the file includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | **Returns** @@ -1745,7 +1717,7 @@ void deserialize( raft::resources const& handle, std::istream& is, cuvs::neighbors::cagra::device_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` Experimental, both the API and the serialization format are subject to change. @@ -1757,7 +1729,7 @@ Experimental, both the API and the serialization format are subject to change. | `handle` | in | `raft::resources const&` | the raft handle | | `is` | in | `std::istream&` | input stream | | `index` | out | `cuvs::neighbors::cagra::device_padded_index*` | the cagra index | -| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the stream includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | +| `out_dataset` | out | `std::unique_ptr>*` | if non-null, on success may be set to an owned deserialized dataset when the stream includes dataset data; may be left unchanged otherwise. Optional; pass nullptr to ignore.
Default: `nullptr`. | **Returns** @@ -2122,7 +2094,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` **Parameters** @@ -2132,7 +2104,7 @@ std::unique_ptr>* out_datas | `handle` | | `raft::resources const&` | | | `filename` | | `const std::string&` | | | `index` | | `cuvs::neighbors::cagra::host_padded_index*` | | -| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | +| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | **Returns** @@ -2145,7 +2117,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_standard_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` **Parameters** @@ -2155,7 +2127,7 @@ std::unique_ptr>* out_dat | `handle` | | `raft::resources const&` | | | `filename` | | `const std::string&` | | | `index` | | `cuvs::neighbors::cagra::host_standard_index*` | | -| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | +| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | **Returns** @@ -2168,7 +2140,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` **Parameters** @@ -2178,7 +2150,7 @@ std::unique_ptr>* out_datase | `handle` | | `raft::resources const&` | | | `filename` | | `const std::string&` | | | `index` | | `cuvs::neighbors::cagra::host_padded_index*` | | -| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | +| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | **Returns** @@ -2191,7 +2163,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_standard_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` **Parameters** @@ -2201,7 +2173,7 @@ std::unique_ptr>* out_data | `handle` | | `raft::resources const&` | | | `filename` | | `const std::string&` | | | `index` | | `cuvs::neighbors::cagra::host_standard_index*` | | -| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | +| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | **Returns** @@ -2214,7 +2186,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` **Parameters** @@ -2224,7 +2196,7 @@ std::unique_ptr>* out_data | `handle` | | `raft::resources const&` | | | `filename` | | `const std::string&` | | | `index` | | `cuvs::neighbors::cagra::host_padded_index*` | | -| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | +| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | **Returns** @@ -2237,7 +2209,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_standard_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` **Parameters** @@ -2247,7 +2219,7 @@ std::unique_ptr>* out_da | `handle` | | `raft::resources const&` | | | `filename` | | `const std::string&` | | | `index` | | `cuvs::neighbors::cagra::host_standard_index*` | | -| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | +| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | **Returns** @@ -2260,7 +2232,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_padded_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` **Parameters** @@ -2270,7 +2242,7 @@ std::unique_ptr>* out_dat | `handle` | | `raft::resources const&` | | | `filename` | | `const std::string&` | | | `index` | | `cuvs::neighbors::cagra::host_padded_index*` | | -| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | +| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | **Returns** @@ -2283,7 +2255,7 @@ void deserialize( raft::resources const& handle, const std::string& filename, cuvs::neighbors::cagra::host_standard_index* index, -std::unique_ptr>* out_dataset = nullptr); +std::unique_ptr>* out_dataset = nullptr); ``` **Parameters** @@ -2293,7 +2265,7 @@ std::unique_ptr>* out_d | `handle` | | `raft::resources const&` | | | `filename` | | `const std::string&` | | | `index` | | `cuvs::neighbors::cagra::host_standard_index*` | | -| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | +| `out_dataset` | | `std::unique_ptr>*` | Default: `nullptr`. | **Returns** diff --git a/fern/pages/cpp_api/cpp-api-neighbors-common.md b/fern/pages/cpp_api/cpp-api-neighbors-common.md index 1caee6f248..da72bab316 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-common.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-common.md @@ -57,139 +57,6 @@ enum class MergeStrategy { | `MERGE_STRATEGY_PHYSICAL` | `0` | | `MERGE_STRATEGY_LOGICAL` | `1` | -## Types - - -### neighbors::dataset - -Tags selecting dataset representation for `dataset` / `dataset_view`. - -Each container defines nested `owning_storage` then `view_storage` (aliases into `detail::*` storage types shared by device/host). Accessibility (device vs host) is selected by the `Accessor` template parameter on `dataset` / `dataset_view`, not by duplicating containers. Layout kinds: empty, padded, standard, VPQ. `dataset` / `dataset_view` only express ownership vs view. - -```cpp -template -struct dataset; -``` - - -### math_type - -Floating-point type used for VQ/PQ codebooks (rows are still uint8 codes). - -```cpp -using math_type = MathT; -``` - - -### neighbors::dataset_view_kind_of - -Primary template returns `unknown` so traits safely return `false` for non-dataset-view types. - -```cpp -template -struct dataset_view_kind_of { - static constexpr dataset_view_kind value; -}; -``` - -**Fields** - -| Name | Type | Description | -| --- | --- | --- | -| `value` | `static constexpr dataset_view_kind` | | - - -### neighbors::dataset_view_is_device_accessible - -True when the dataset view accessor is device-accessible. - -```cpp -template -struct dataset_view_is_device_accessible; -``` - - -### neighbors::with_accessor - -Generic accessor retargeting while preserving the dataset tag/layout and value/index types: - -`dataset<Tag, DataT, IdxT, OldAccessor> -> dataset<Tag, DataT, IdxT, NewAccessor>` `dataset_view<Tag, DataT, IdxT, OldAccessor> -> dataset_view<Tag, DataT, IdxT, NewAccessor>` - -```cpp -template -struct with_accessor; -``` - - -### neighbors::to_device_accessor - -Map any host accessor to its device counterpart (same payload policy). - -```cpp -template -struct to_device_accessor; -``` - - -### neighbors::device_counterpart - -Maps a host dataset view type to its device-resident counterpart. - -```cpp -template -struct device_counterpart; -``` - - -### neighbors::cagra_view_element_type - -Element type `T` for `cagra::build(res, params, dataset_view)` (deduced, not a template arg). - -```cpp -template -struct cagra_view_element_type; -``` - - -### neighbors::ivf::list_base - -Abstract base class for IVF list data. - -This allows polymorphic access to list data regardless of the underlying layout. - -TODO: Make this struct internal (tracking issue: https://github.com/nvidia/cuvs/issues/1726) - -```cpp -template -struct list_base; -``` - - -### neighbors::ivf::list - -The data for a single IVF list. - -```cpp -template