From b71ef73ca48e55563bc049cd6425605522cff43e Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Tue, 31 Oct 2023 20:31:22 -0400 Subject: [PATCH] Turning off faiss refinement for the time being. --- cpp/bench/ann/src/common/ann_types.hpp | 4 +- cpp/bench/ann/src/common/benchmark.cpp | 1 + cpp/bench/ann/src/common/benchmark.hpp | 20 +++---- cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h | 60 ++++++++++--------- .../run/conf/algos/faiss_gpu_ivf_flat.yaml | 10 ++-- .../run/conf/algos/faiss_gpu_ivf_pq.yaml | 2 +- 6 files changed, 48 insertions(+), 49 deletions(-) diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index f89743ac81..e964a81efa 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -120,11 +120,11 @@ class ANN : public AnnBase { // The advantage of this way is that index has smaller size // and many indices can share one dataset. // - // AlgoProperty::need_dataset_when_search of such algorithm should be true, + // SearchParam::needs_dataset() of such algorithm should be true, // and set_search_dataset() should save the passed-in pointer somewhere. // The client code should call set_search_dataset() before searching, // and should not release dataset before searching is finished. - virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/) { printf("Setting \n"); }; + virtual void set_search_dataset(const T* /*dataset*/, size_t /*nrow*/){}; }; } // namespace raft::bench::ann diff --git a/cpp/bench/ann/src/common/benchmark.cpp b/cpp/bench/ann/src/common/benchmark.cpp index 6424a36471..62c91b0791 100644 --- a/cpp/bench/ann/src/common/benchmark.cpp +++ b/cpp/bench/ann/src/common/benchmark.cpp @@ -88,6 +88,7 @@ template std::unique_ptr::AnnSearchParam> create_search_param( const std::string& algo, const nlohmann::json& conf) { + printf("INside create_search_param\n"); static auto fname = get_fun_name(reinterpret_cast(&create_search_param)); auto handle = load_lib(algo); auto fun_addr = dlsym(handle, fname.c_str()); diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index ae36779c86..3a930d288e 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -172,7 +172,6 @@ void bench_search(::benchmark::State& state, std::ptrdiff_t batch_offset = 0; std::size_t queries_processed = 0; - printf("Starting benchmark search\n"); double total_time = 0; const auto& sp_json = index.search_params[search_param_ix]; @@ -202,7 +201,6 @@ void bench_search(::benchmark::State& state, index_file = index.file; } - printf("Loading index from file\n"); std::unique_ptr::AnnSearchParam> search_param; ANN* algo; try { @@ -217,25 +215,15 @@ void bench_search(::benchmark::State& state, search_param->metric_objective = metric_objective; } catch (const std::exception& e) { state.SkipWithError("Failed to create an algo: " + std::string(e.what())); - } - - printf("Set search params\n"); - try { - algo->set_search_param(*search_param); - } catch (const std::exception& ex) { - state.SkipWithError("An error occurred setting search parameters: " + std::string(ex.what())); return; } - printf("Setting search dataset\n"); auto algo_property = parse_algo_property(algo->get_preference(), sp_json); current_algo_props = std::make_shared(algo_property.dataset_memory_type, algo_property.query_memory_type); - printf("AFTER!\n"); if (search_param->needs_dataset()) { try { - printf("About to set search datast\n"); algo->set_search_dataset(dataset->base_set(current_algo_props->dataset_memory_type), dataset->base_set_size()); } catch (const std::exception& ex) { @@ -245,6 +233,14 @@ void bench_search(::benchmark::State& state, return; } } + + try { + algo->set_search_param(*search_param); + + } catch (const std::exception& ex) { + state.SkipWithError("An error occurred setting search parameters: " + std::string(ex.what())); + return; + } } const auto algo_property = *current_algo_props; diff --git a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h index c94450d480..4f13ff8a49 100644 --- a/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h +++ b/cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h @@ -35,6 +35,9 @@ #include #include +#include +#include + #include #include #include @@ -102,6 +105,7 @@ class FaissGpu : public ANN { RAFT_CUDA_TRY(cudaGetDevice(&device_)); RAFT_CUDA_TRY(cudaEventCreate(&sync_, cudaEventDisableTiming)); faiss_default_stream_ = gpu_resource_.getDefaultStream(device_); + raft::resource::set_cuda_stream(handle_, faiss_default_stream_); } virtual ~FaissGpu() noexcept { RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(sync_)); } @@ -110,7 +114,7 @@ class FaissGpu : public ANN { virtual void set_search_param(const FaissGpu::AnnSearchParam& param) {} - virtual void set_search_dataset(const T* dataset, size_t nrow) {} + void set_search_dataset(const T* dataset, size_t nrow) override { dataset_ = dataset; } // TODO: if the number of results is less than k, the remaining elements of 'neighbors' // will be filled with (size_t)-1 @@ -126,7 +130,7 @@ class FaissGpu : public ANN { AlgoProperty property; // to enable building big dataset which is larger than GPU memory property.dataset_memory_type = MemoryType::Host; - property.query_memory_type = MemoryType::Device; + property.query_memory_type = MemoryType::Host; return property; } @@ -145,7 +149,7 @@ class FaissGpu : public ANN { mutable faiss::gpu::StandardGpuResources gpu_resource_; std::unique_ptr index_; - std::unique_ptr index_refine_; + std::unique_ptr index_refine_{nullptr}; faiss::MetricType metric_type_; int nlist_; int device_; @@ -153,6 +157,9 @@ class FaissGpu : public ANN { cudaStream_t faiss_default_stream_{nullptr}; double training_sample_fraction_; std::unique_ptr search_params_; + const T* dataset_; + raft::device_resources handle_; + float refine_ratio_ = 1.0; }; template @@ -198,12 +205,23 @@ void FaissGpu::search(const T* queries, static_assert(sizeof(size_t) == sizeof(faiss::idx_t), "sizes of size_t and faiss::idx_t are different"); - if (index_refine_->k_factor > 1) { - printf("Using refine!\n"); - index_refine_->search( - batch_size, queries, k, distances, reinterpret_cast(neighbors)); + if (this->refine_ratio_ > 1.0) { + // TODO: FAISS changed their search APIs to accept the search parameters as a struct object + // but their refine API doesn't allow the struct to be passed in. Once this is fixed, we + // need to re-enable refinement below + // index_refine_->search(batch_size, queries, k, distances, + // reinterpret_cast(neighbors), this->search_params_.get()); Related FAISS issue: + // https://github.com/facebookresearch/faiss/issues/3118 + throw std::runtime_error( + "FAISS doesn't support refinement in their new APIs so this feature is disabled in the " + "benchmarks for the time being."); } else { - index_->search(batch_size, queries, k, distances, reinterpret_cast(neighbors)); + index_->search(batch_size, + queries, + k, + distances, + reinterpret_cast(neighbors), + this->search_params_.get()); } stream_wait(stream); } @@ -258,6 +276,7 @@ class FaissGpuIVFFlat : public FaissGpu { faiss::IVFSearchParameters faiss_search_params; faiss_search_params.nprobe = nprobe; this->search_params_ = std::make_unique(faiss_search_params); + this->refine_ratio_ = search_param.refine_ratio; } void save(const std::string& file) const override @@ -296,19 +315,12 @@ class FaissGpuIVFPQ : public FaissGpu { config); } - void set_search_dataset(const T* dataset, size_t nrow) override - { - printf("Setting search ataset for refine\n"); - dataset_ = dataset; - } - void set_search_param(const typename FaissGpu::AnnSearchParam& param) override { - printf("Setting ivfpq search params\n"); auto search_param = dynamic_cast::SearchParam&>(param); int nprobe = search_param.nprobe; assert(nprobe <= nlist_); - + this->refine_ratio_ = search_param.refine_ratio; faiss::IVFPQSearchParameters faiss_search_params; faiss_search_params.nprobe = nprobe; @@ -329,8 +341,6 @@ class FaissGpuIVFPQ : public FaissGpu { { this->template load_(file); } - - const T* dataset_; }; // TODO: Enable this in cmake @@ -342,12 +352,6 @@ class FaissGpuIVFSQ : public FaissGpu { std::string quantizer_type; }; - struct SearchParam : public FaissGpu::SearchParam { - int nprobe; - float refine_ratio = 1.0; - auto needs_dataset() const -> bool override { return true; } - }; - FaissGpuIVFSQ(Metric metric, int dim, const BuildParam& param) : FaissGpu(metric, dim, param) { faiss::ScalarQuantizer::QuantizerType qtype; @@ -366,7 +370,6 @@ class FaissGpuIVFSQ : public FaissGpu { &(this->gpu_resource_), dim, param.nlist, qtype, this->metric_type_, true, config); } - void set_search_dataset(const T* dataset, size_t nrow) override { this->dataset_ = dataset; } void set_search_param(const typename FaissGpu::AnnSearchParam& param) override { auto search_param = dynamic_cast::SearchParam&>(param); @@ -377,9 +380,10 @@ class FaissGpuIVFSQ : public FaissGpu { faiss_search_params.nprobe = nprobe; this->search_params_ = std::make_unique(faiss_search_params); - + this->refine_ratio_ = search_param.refine_ratio; if (search_param.refine_ratio > 1.0) { - this->index_refine_ = std::make_unique(this->index_.get(), dataset_); + this->index_refine_ = + std::make_unique(this->index_.get(), this->dataset_); this->index_refine_.get()->k_factor = search_param.refine_ratio; } } @@ -394,8 +398,6 @@ class FaissGpuIVFSQ : public FaissGpu { this->template load_( file); } - - const T* dataset_; }; template diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml index 6542bbab4c..ed237becb3 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_flat.yaml @@ -2,9 +2,9 @@ name: faiss_gpu_ivf_flat groups: base: build: - nlists: [1024, 2048, 4096, 8192, 16000, 32000] - ratio: [1, 10, 25] - useFloat16: [True, False] + nlist: [2048] + ratio: [1, 4, 10] + useFloat16: [False] search: - numProbes: [1, 5, 10, 50, 100, 200, 500, 1000, 2000] - refine_ratio: [1, 2, 4, 10] \ No newline at end of file + nprobe: [2048] + refine_ratio: [1] diff --git a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml index 3a17eacc2d..9c8b99c21f 100644 --- a/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml +++ b/python/raft-ann-bench/src/raft-ann-bench/run/conf/algos/faiss_gpu_ivf_pq.yaml @@ -9,7 +9,7 @@ groups: useFloat16: [False] search: nprobe: [1, 5, 10, 50, 100, 200] - refine_ratio: [1, 2, 4] + refine_ratio: [1] test: build: nlist: [1024]