Loading cpp/src/core/CMakeLists.txt +1 −1 Original line number Diff line number Diff line Loading @@ -93,7 +93,7 @@ endif() set(CORE_INCLUDE_DIRS ${CORE_INCLUDE_DIRS} PARENT_SCOPE) if(BUILD_UNIT_TEST STREQUAL "ON") # add_subdirectory(test) add_subdirectory(test) endif() config_summary() cpp/src/core/knowhere/CMakeLists.txt +1 −1 Original line number Diff line number Diff line Loading @@ -8,7 +8,6 @@ link_directories(${CUDA_TOOLKIT_ROOT_DIR}/lib64) include_directories(${CORE_SOURCE_DIR}/knowhere) include_directories(${CORE_SOURCE_DIR}/thirdparty) include_directories(${CORE_SOURCE_DIR}/thirdparty/SPTAG/AnnService) include_directories(${CORE_SOURCE_DIR}/thirdparty/jsoncons-0.126.0/include) set(SPTAG_SOURCE_DIR ${CORE_SOURCE_DIR}/thirdparty/SPTAG) file(GLOB HDR_FILES Loading Loading @@ -55,6 +54,7 @@ set(index_srcs knowhere/index/vector_index/IndexGPUIVFPQ.cpp knowhere/index/vector_index/FaissBaseIndex.cpp knowhere/index/vector_index/helpers/FaissIO.cpp knowhere/index/vector_index/helpers/IndexParameter.cpp ) set(depend_libs Loading cpp/src/core/knowhere/knowhere/adapter/SptagAdapter.cpp +1 −2 Original line number Diff line number Diff line Loading @@ -69,8 +69,7 @@ ConvertToQueryResult(const DatasetPtr &dataset, const Config &config) { auto dimension = tensor->shape()[1]; auto rows = tensor->shape()[0]; auto k = config[META_K].as<int64_t>(); std::vector<SPTAG::QueryResult> query_results(rows, SPTAG::QueryResult(nullptr, k, true)); std::vector<SPTAG::QueryResult> query_results(rows, SPTAG::QueryResult(nullptr, config->k, true)); for (auto i = 0; i < rows; ++i) { query_results[i].SetTarget(&p_data[i * dimension]); } Loading cpp/src/core/knowhere/knowhere/common/Config.h +32 −5 Original line number Diff line number Diff line Loading @@ -18,15 +18,42 @@ #pragma once #include <jsoncons/json.hpp> #include <memory> namespace zilliz { namespace knowhere { using Config = jsoncons::json; enum class METRICTYPE { INVALID = 0, L2 = 1, IP = 2, }; // General Config constexpr int64_t INVALID_VALUE = -1; constexpr int64_t DEFAULT_K = INVALID_VALUE; constexpr int64_t DEFAULT_DIM = INVALID_VALUE; constexpr int64_t DEFAULT_GPUID = INVALID_VALUE; constexpr METRICTYPE DEFAULT_TYPE = METRICTYPE::INVALID; struct Cfg { METRICTYPE metric_type = DEFAULT_TYPE; int64_t k = DEFAULT_K; int64_t gpu_id = DEFAULT_GPUID; int64_t d = DEFAULT_DIM; Cfg(const int64_t &dim, const int64_t &k, const int64_t &gpu_id, METRICTYPE type) : d(dim), k(k), gpu_id(gpu_id), metric_type(type) {} Cfg() = default; virtual bool CheckValid(){}; }; using Config = std::shared_ptr<Cfg>; } // namespace knowhere } // namespace zilliz cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp +11 −7 Original line number Diff line number Diff line Loading @@ -36,10 +36,11 @@ namespace zilliz { namespace knowhere { IndexModelPtr GPUIVF::Train(const DatasetPtr &dataset, const Config &config) { auto nlist = config["nlist"].as<size_t>(); gpu_id_ = config.get_with_default("gpu_id", gpu_id_); auto metric_type = config["metric_type"].as_string() == "L2" ? faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; auto build_cfg = std::dynamic_pointer_cast<IVFCfg>(config); if (build_cfg != nullptr) { build_cfg->CheckValid(); // throw exception } gpu_id_ = build_cfg->gpu_id; GETTENSOR(dataset) Loading @@ -48,7 +49,9 @@ IndexModelPtr GPUIVF::Train(const DatasetPtr &dataset, const Config &config) { ResScope rs(temp_resource, gpu_id_, true); faiss::gpu::GpuIndexIVFFlatConfig idx_config; idx_config.device = gpu_id_; faiss::gpu::GpuIndexIVFFlat device_index(temp_resource->faiss_res.get(), dim, nlist, metric_type, idx_config); faiss::gpu::GpuIndexIVFFlat device_index(temp_resource->faiss_res.get(), dim, build_cfg->nlist, GetMetricType(build_cfg->metric_type), idx_config); device_index.train(rows, (float *) p_data); std::shared_ptr<faiss::Index> host_index = nullptr; Loading Loading @@ -143,9 +146,10 @@ void GPUIVF::search_impl(int64_t n, const Config &cfg) { std::lock_guard<std::mutex> lk(mutex_); // TODO(linxj): gpu index support GenParams if (auto device_index = std::static_pointer_cast<faiss::gpu::GpuIndexIVF>(index_)) { auto nprobe = cfg.get_with_default("nprobe", size_t(1)); device_index->setNumProbes(nprobe); auto search_cfg = std::dynamic_pointer_cast<IVFCfg>(cfg); device_index->setNumProbes(search_cfg->nprobe); { // TODO(linxj): allocate mem Loading Loading
cpp/src/core/CMakeLists.txt +1 −1 Original line number Diff line number Diff line Loading @@ -93,7 +93,7 @@ endif() set(CORE_INCLUDE_DIRS ${CORE_INCLUDE_DIRS} PARENT_SCOPE) if(BUILD_UNIT_TEST STREQUAL "ON") # add_subdirectory(test) add_subdirectory(test) endif() config_summary()
cpp/src/core/knowhere/CMakeLists.txt +1 −1 Original line number Diff line number Diff line Loading @@ -8,7 +8,6 @@ link_directories(${CUDA_TOOLKIT_ROOT_DIR}/lib64) include_directories(${CORE_SOURCE_DIR}/knowhere) include_directories(${CORE_SOURCE_DIR}/thirdparty) include_directories(${CORE_SOURCE_DIR}/thirdparty/SPTAG/AnnService) include_directories(${CORE_SOURCE_DIR}/thirdparty/jsoncons-0.126.0/include) set(SPTAG_SOURCE_DIR ${CORE_SOURCE_DIR}/thirdparty/SPTAG) file(GLOB HDR_FILES Loading Loading @@ -55,6 +54,7 @@ set(index_srcs knowhere/index/vector_index/IndexGPUIVFPQ.cpp knowhere/index/vector_index/FaissBaseIndex.cpp knowhere/index/vector_index/helpers/FaissIO.cpp knowhere/index/vector_index/helpers/IndexParameter.cpp ) set(depend_libs Loading
cpp/src/core/knowhere/knowhere/adapter/SptagAdapter.cpp +1 −2 Original line number Diff line number Diff line Loading @@ -69,8 +69,7 @@ ConvertToQueryResult(const DatasetPtr &dataset, const Config &config) { auto dimension = tensor->shape()[1]; auto rows = tensor->shape()[0]; auto k = config[META_K].as<int64_t>(); std::vector<SPTAG::QueryResult> query_results(rows, SPTAG::QueryResult(nullptr, k, true)); std::vector<SPTAG::QueryResult> query_results(rows, SPTAG::QueryResult(nullptr, config->k, true)); for (auto i = 0; i < rows; ++i) { query_results[i].SetTarget(&p_data[i * dimension]); } Loading
cpp/src/core/knowhere/knowhere/common/Config.h +32 −5 Original line number Diff line number Diff line Loading @@ -18,15 +18,42 @@ #pragma once #include <jsoncons/json.hpp> #include <memory> namespace zilliz { namespace knowhere { using Config = jsoncons::json; enum class METRICTYPE { INVALID = 0, L2 = 1, IP = 2, }; // General Config constexpr int64_t INVALID_VALUE = -1; constexpr int64_t DEFAULT_K = INVALID_VALUE; constexpr int64_t DEFAULT_DIM = INVALID_VALUE; constexpr int64_t DEFAULT_GPUID = INVALID_VALUE; constexpr METRICTYPE DEFAULT_TYPE = METRICTYPE::INVALID; struct Cfg { METRICTYPE metric_type = DEFAULT_TYPE; int64_t k = DEFAULT_K; int64_t gpu_id = DEFAULT_GPUID; int64_t d = DEFAULT_DIM; Cfg(const int64_t &dim, const int64_t &k, const int64_t &gpu_id, METRICTYPE type) : d(dim), k(k), gpu_id(gpu_id), metric_type(type) {} Cfg() = default; virtual bool CheckValid(){}; }; using Config = std::shared_ptr<Cfg>; } // namespace knowhere } // namespace zilliz
cpp/src/core/knowhere/knowhere/index/vector_index/IndexGPUIVF.cpp +11 −7 Original line number Diff line number Diff line Loading @@ -36,10 +36,11 @@ namespace zilliz { namespace knowhere { IndexModelPtr GPUIVF::Train(const DatasetPtr &dataset, const Config &config) { auto nlist = config["nlist"].as<size_t>(); gpu_id_ = config.get_with_default("gpu_id", gpu_id_); auto metric_type = config["metric_type"].as_string() == "L2" ? faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT; auto build_cfg = std::dynamic_pointer_cast<IVFCfg>(config); if (build_cfg != nullptr) { build_cfg->CheckValid(); // throw exception } gpu_id_ = build_cfg->gpu_id; GETTENSOR(dataset) Loading @@ -48,7 +49,9 @@ IndexModelPtr GPUIVF::Train(const DatasetPtr &dataset, const Config &config) { ResScope rs(temp_resource, gpu_id_, true); faiss::gpu::GpuIndexIVFFlatConfig idx_config; idx_config.device = gpu_id_; faiss::gpu::GpuIndexIVFFlat device_index(temp_resource->faiss_res.get(), dim, nlist, metric_type, idx_config); faiss::gpu::GpuIndexIVFFlat device_index(temp_resource->faiss_res.get(), dim, build_cfg->nlist, GetMetricType(build_cfg->metric_type), idx_config); device_index.train(rows, (float *) p_data); std::shared_ptr<faiss::Index> host_index = nullptr; Loading Loading @@ -143,9 +146,10 @@ void GPUIVF::search_impl(int64_t n, const Config &cfg) { std::lock_guard<std::mutex> lk(mutex_); // TODO(linxj): gpu index support GenParams if (auto device_index = std::static_pointer_cast<faiss::gpu::GpuIndexIVF>(index_)) { auto nprobe = cfg.get_with_default("nprobe", size_t(1)); device_index->setNumProbes(nprobe); auto search_cfg = std::dynamic_pointer_cast<IVFCfg>(cfg); device_index->setNumProbes(search_cfg->nprobe); { // TODO(linxj): allocate mem Loading