Skip to content

Add leanvec_primary_only build option to C++ runtime #323

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ibhati wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Add leanvec_primary_only build option to C++ runtime #323

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions bindings/cpp/include/svs/runtime/dynamic_vamana_index.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,7 +105,8 @@ struct SVS_RUNTIME_API DynamicVamanaIndexLeanVec : public DynamicVamanaIndex {
size_t leanvec_dims,
const VamanaIndex::BuildParams& params,
const VamanaIndex::SearchParams& default_search_params,
const VamanaIndex::DynamicIndexParams& dynamic_index_params
const VamanaIndex::DynamicIndexParams& dynamic_index_params,
bool leanvec_primary_only = false
) noexcept;

// Specialization to build LeanVec-based Vamana index with provided training data
Expand All @@ -128,7 +129,8 @@ struct SVS_RUNTIME_API DynamicVamanaIndexLeanVec : public DynamicVamanaIndex {
const LeanVecTrainingData* training_data,
const VamanaIndex::BuildParams& params,
const VamanaIndex::SearchParams& default_search_params,
const VamanaIndex::DynamicIndexParams& dynamic_index_params
const VamanaIndex::DynamicIndexParams& dynamic_index_params,
bool leanvec_primary_only = false
) noexcept;
};

Expand Down
6 changes: 4 additions & 2 deletions bindings/cpp/include/svs/runtime/vamana_index.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,7 +114,8 @@ struct SVS_RUNTIME_API VamanaIndexLeanVec : public VamanaIndex {
StorageKind storage_kind,
size_t leanvec_dims,
const VamanaIndex::BuildParams& params = {},
const VamanaIndex::SearchParams& default_search_params = {}
const VamanaIndex::SearchParams& default_search_params = {},
bool leanvec_primary_only = false
) noexcept;

// Specialization to build LeanVec-based Vamana index with provided training data
Expand All @@ -125,7 +126,8 @@ struct SVS_RUNTIME_API VamanaIndexLeanVec : public VamanaIndex {
StorageKind storage_kind,
const LeanVecTrainingData* training_data,
const VamanaIndex::BuildParams& params = {},
const VamanaIndex::SearchParams& default_search_params = {}
const VamanaIndex::SearchParams& default_search_params = {},
bool leanvec_primary_only = false
) noexcept;
};

Expand Down
38 changes: 30 additions & 8 deletions bindings/cpp/src/dynamic_vamana_index.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -271,7 +271,8 @@ Status DynamicVamanaIndexLeanVec::build(
size_t leanvec_dims,
const DynamicVamanaIndex::BuildParams& params,
const DynamicVamanaIndex::SearchParams& default_search_params,
const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params
const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params,
bool leanvec_primary_only
) noexcept {
using Impl = DynamicVamanaIndexLeanVecImpl;
*index = nullptr;
Expand All @@ -288,7 +289,8 @@ Status DynamicVamanaIndexLeanVec::build(
leanvec_dims,
params,
default_search_params,
dynamic_index_params
dynamic_index_params,
leanvec_primary_only
);
*index = new DynamicVamanaIndexManagerBase<Impl>{std::move(impl)};
});
Expand All @@ -303,7 +305,8 @@ Status DynamicVamanaIndexLeanVec::build(
const LeanVecTrainingData* training_data,
const DynamicVamanaIndex::BuildParams& params,
const DynamicVamanaIndex::SearchParams& default_search_params,
const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params
const DynamicVamanaIndex::DynamicIndexParams& dynamic_index_params,
bool leanvec_primary_only
) noexcept {
using Impl = DynamicVamanaIndexLeanVecImpl;
*index = nullptr;
Expand All @@ -322,24 +325,43 @@ Status DynamicVamanaIndexLeanVec::build(
training_data_impl,
params,
default_search_params,
dynamic_index_params
dynamic_index_params,
leanvec_primary_only
);
*index = new DynamicVamanaIndexManagerBase<Impl>{std::move(impl)};
});
}

#else // SVS_RUNTIME_HAVE_LVQ_LEANVEC
// LeanVec storage kind is not supported in this build configuration
Status DynamicVamanaIndexLeanVec::
build(DynamicVamanaIndex**, size_t, MetricType, StorageKind, size_t, const DynamicVamanaIndex::BuildParams&, const DynamicVamanaIndex::SearchParams&, const DynamicVamanaIndex::DynamicIndexParams&) noexcept {
Status DynamicVamanaIndexLeanVec::build(
DynamicVamanaIndex**,
size_t,
MetricType,
StorageKind,
size_t,
const DynamicVamanaIndex::BuildParams&,
const DynamicVamanaIndex::SearchParams&,
const DynamicVamanaIndex::DynamicIndexParams&,
bool
) noexcept {
return Status(
ErrorCode::NOT_IMPLEMENTED,
"DynamicVamanaIndexLeanVec is not supported in this build configuration."
);
}

Status DynamicVamanaIndexLeanVec::
build(DynamicVamanaIndex**, size_t, MetricType, StorageKind, const LeanVecTrainingData*, const DynamicVamanaIndex::BuildParams&, const DynamicVamanaIndex::SearchParams&, const DynamicVamanaIndex::DynamicIndexParams&) noexcept {
Status DynamicVamanaIndexLeanVec::build(
DynamicVamanaIndex**,
size_t,
MetricType,
StorageKind,
const LeanVecTrainingData*,
const DynamicVamanaIndex::BuildParams&,
const DynamicVamanaIndex::SearchParams&,
const DynamicVamanaIndex::DynamicIndexParams&,
bool
) noexcept {
return Status(
ErrorCode::NOT_IMPLEMENTED,
"DynamicVamanaIndexLeanVec is not supported in this build configuration."
Expand Down
16 changes: 11 additions & 5 deletions bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,11 +54,13 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl {
const LeanVecTrainingDataImpl& training_data,
const VamanaIndex::BuildParams& params,
const VamanaIndex::SearchParams& default_search_params,
const VamanaIndex::DynamicIndexParams& dynamic_index_params
const VamanaIndex::DynamicIndexParams& dynamic_index_params,
bool leanvec_primary_only = false
)
: DynamicVamanaIndexImpl{dim, metric, storage_kind, params, default_search_params, dynamic_index_params}
, leanvec_dims_{training_data.get_leanvec_dims()}
, leanvec_matrices_{training_data.get_leanvec_matrices()} {
, leanvec_matrices_{training_data.get_leanvec_matrices()}
, leanvec_primary_only_{leanvec_primary_only} {
check_storage_kind(storage_kind);
}

Expand All @@ -69,11 +71,13 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl {
size_t leanvec_dims,
const VamanaIndex::BuildParams& params,
const VamanaIndex::SearchParams& default_search_params,
const VamanaIndex::DynamicIndexParams& dynamic_index_params
const VamanaIndex::DynamicIndexParams& dynamic_index_params,
bool leanvec_primary_only = false
)
: DynamicVamanaIndexImpl{dim, metric, storage_kind, params, default_search_params, dynamic_index_params}
, leanvec_dims_{leanvec_dims}
, leanvec_matrices_{std::nullopt} {
, leanvec_matrices_{std::nullopt}
, leanvec_primary_only_{leanvec_primary_only} {
check_storage_kind(storage_kind);
}

Expand Down Expand Up @@ -124,7 +128,8 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl {
labels,
blocksize_bytes,
this->leanvec_dims_,
this->leanvec_matrices_
this->leanvec_matrices_,
this->leanvec_primary_only_
);
},
data,
Expand All @@ -136,6 +141,7 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl {
protected:
size_t leanvec_dims_;
std::optional<LeanVecMatricesType> leanvec_matrices_;
bool leanvec_primary_only_ = false;

StorageKind check_storage_kind(StorageKind kind) {
if (!storage::is_leanvec_storage(kind)) {
Expand Down
11 changes: 9 additions & 2 deletions bindings/cpp/src/svs_runtime_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -355,13 +355,20 @@ struct StorageFactory<LeanVecStorageType> {
Pool& pool,
const Alloc& alloc = {},
size_t leanvec_d = 0,
std::optional<svs::leanvec::LeanVecMatrices<svs::Dynamic>> matrices = std::nullopt
std::optional<svs::leanvec::LeanVecMatrices<svs::Dynamic>> matrices = std::nullopt,
bool leanvec_primary_only = false
) {
if (leanvec_d == 0) {
leanvec_d = (data.dimensions() + 1) / 2;
}
return LeanVecStorageType::reduce(
data, std::move(matrices), pool, 0, svs::lib::MaybeStatic{leanvec_d}, alloc
data,
std::move(matrices),
pool,
0,
svs::lib::MaybeStatic{leanvec_d},
alloc,
leanvec_primary_only
);
}
};
Expand Down
46 changes: 38 additions & 8 deletions bindings/cpp/src/vamana_index.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -153,14 +153,21 @@ Status VamanaIndexLeanVec::build(
StorageKind storage_kind,
size_t leanvec_dims,
const VamanaIndex::BuildParams& params,
const VamanaIndex::SearchParams& default_search_params
const VamanaIndex::SearchParams& default_search_params,
bool leanvec_primary_only
) noexcept {
using Impl = VamanaIndexLeanVecImpl;
*index = nullptr;

return runtime_error_wrapper([&] {
auto impl = std::make_unique<Impl>(
dim, metric, storage_kind, leanvec_dims, params, default_search_params
dim,
metric,
storage_kind,
leanvec_dims,
params,
default_search_params,
leanvec_primary_only
);
*index = new VamanaIndexManagerBase<Impl>{std::move(impl)};
});
Expand All @@ -174,7 +181,8 @@ Status VamanaIndexLeanVec::build(
StorageKind storage_kind,
const LeanVecTrainingData* training_data,
const VamanaIndex::BuildParams& params,
const VamanaIndex::SearchParams& default_search_params
const VamanaIndex::SearchParams& default_search_params,
bool leanvec_primary_only
) noexcept {
using Impl = VamanaIndexLeanVecImpl;
*index = nullptr;
Expand All @@ -187,24 +195,46 @@ Status VamanaIndexLeanVec::build(
auto training_data_impl =
static_cast<const LeanVecTrainingDataManager*>(training_data)->impl_;
auto impl = std::make_unique<Impl>(
dim, metric, storage_kind, training_data_impl, params, default_search_params
dim,
metric,
storage_kind,
training_data_impl,
params,
default_search_params,
leanvec_primary_only
);
*index = new VamanaIndexManagerBase<Impl>{std::move(impl)};
});
}

#else // SVS_RUNTIME_HAVE_LVQ_LEANVEC
// LeanVec storage kind is not supported in this build configuration
Status VamanaIndexLeanVec::
build(VamanaIndex**, size_t, MetricType, StorageKind, size_t, const VamanaIndex::BuildParams&, const VamanaIndex::SearchParams&) noexcept {
Status VamanaIndexLeanVec::build(
VamanaIndex**,
size_t,
MetricType,
StorageKind,
size_t,
const VamanaIndex::BuildParams&,
const VamanaIndex::SearchParams&,
bool
) noexcept {
return Status(
ErrorCode::NOT_IMPLEMENTED,
"VamanaIndexLeanVec is not supported in this build configuration."
);
}

Status VamanaIndexLeanVec::
build(VamanaIndex**, size_t, MetricType, StorageKind, const LeanVecTrainingData*, const VamanaIndex::BuildParams&, const VamanaIndex::SearchParams&) noexcept {
Status VamanaIndexLeanVec::build(
VamanaIndex**,
size_t,
MetricType,
StorageKind,
const LeanVecTrainingData*,
const VamanaIndex::BuildParams&,
const VamanaIndex::SearchParams&,
bool
) noexcept {
return Status(
ErrorCode::NOT_IMPLEMENTED,
"VamanaIndexLeanVec is not supported in this build configuration."
Expand Down
16 changes: 11 additions & 5 deletions bindings/cpp/src/vamana_index_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -515,11 +515,13 @@ struct VamanaIndexLeanVecImpl : public VamanaIndexImpl {
StorageKind storage_kind,
const LeanVecTrainingDataImpl& training_data,
const VamanaIndex::BuildParams& params,
const VamanaIndex::SearchParams& default_search_params
const VamanaIndex::SearchParams& default_search_params,
bool leanvec_primary_only = false
)
: VamanaIndexImpl{dim, metric, storage_kind, params, default_search_params}
, leanvec_dims_{training_data.get_leanvec_dims()}
, leanvec_matrices_{training_data.get_leanvec_matrices()} {
, leanvec_matrices_{training_data.get_leanvec_matrices()}
, leanvec_primary_only_{leanvec_primary_only} {
check_storage_kind(storage_kind);
}

Expand All @@ -529,11 +531,13 @@ struct VamanaIndexLeanVecImpl : public VamanaIndexImpl {
StorageKind storage_kind,
size_t leanvec_dims,
const VamanaIndex::BuildParams& params,
const VamanaIndex::SearchParams& default_search_params
const VamanaIndex::SearchParams& default_search_params,
bool leanvec_primary_only = false
)
: VamanaIndexImpl{dim, metric, storage_kind, params, default_search_params}
, leanvec_dims_{leanvec_dims}
, leanvec_matrices_{std::nullopt} {
, leanvec_matrices_{std::nullopt}
, leanvec_primary_only_{leanvec_primary_only} {
check_storage_kind(storage_kind);
}

Expand Down Expand Up @@ -573,7 +577,8 @@ struct VamanaIndexLeanVecImpl : public VamanaIndexImpl {
this->vamana_build_parameters(),
data,
leanvec_dims_,
leanvec_matrices_
leanvec_matrices_,
leanvec_primary_only_
);
},
data
Expand All @@ -584,6 +589,7 @@ struct VamanaIndexLeanVecImpl : public VamanaIndexImpl {
protected:
size_t leanvec_dims_;
std::optional<LeanVecMatricesType> leanvec_matrices_;
bool leanvec_primary_only_ = false;

StorageKind check_storage_kind(StorageKind kind) {
if (!storage::is_leanvec_storage(kind)) {
Expand Down
Loading