Skip to content

Commit

Permalink
Support GPU index types (milvus-io#708)
Browse files Browse the repository at this point in the history
Signed-off-by: Yudong Cai <[email protected]>
  • Loading branch information
cydrain authored Feb 28, 2023
1 parent 8025c3b commit 6d5a17a
Show file tree
Hide file tree
Showing 15 changed files with 520 additions and 179 deletions.
6 changes: 4 additions & 2 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ knowhere_option(WITH_DISKANN "Build with diskann index" OFF)
knowhere_option(WITH_BENCHMARK "Build with benchmark" OFF)
knowhere_option(WITH_COVERAGE "Build with coverage" OFF)
knowhere_option(WITH_CCACHE "Build with ccache" ON)
knowhere_option(WITH_PROFILER "Build with profiler" OFF)

if(KNOWHERE_VERSION)
message(STATUS "Building KNOWHERE version: ${KNOWHERE_VERSION}")
Expand All @@ -46,6 +47,7 @@ if(WITH_CCACHE)
endif()

if(USE_CUDA)
add_definitions(-DUSE_CUDA)
set(CMAKE_CUDA_ARCHITECTURES 75;70;61;60)
enable_language(CUDA)
find_package(CUDAToolkit REQUIRED)
Expand Down Expand Up @@ -90,8 +92,8 @@ else()
endif()

if(NOT USE_CUDA)
knowhere_file_glob(GLOB_RECURSE KNOWHERE_GPU_SRCS src/index/ivf_gpu/*.cc
src/index/flat_gpu/*.cc)
knowhere_file_glob(GLOB_RECURSE KNOWHERE_GPU_SRCS src/index/flat_gpu/*.cc
src/index/ivf_gpu/*.cc)
list(REMOVE_ITEM KNOWHERE_SRCS ${KNOWHERE_GPU_SRCS})
endif()

Expand Down
11 changes: 6 additions & 5 deletions benchmark/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
# or implied. See the License for the specific language governing permissions and limitations under the License

include_directories(${CMAKE_SOURCE_DIR})
include_directories(${CMAKE_SOURCE_DIR}/include)

include_directories(/usr/local/hdf5/include)
Expand All @@ -24,12 +25,12 @@ set(depend_libs
${LAPACK_LIBRARIES}
)

if ( LINUX AND ENABLE_PROFILING )
set( depend_libs
if(WITH_PROFILER)
set(depend_libs
${depend_libs}
gperftools
)
endif ()
tcmalloc_and_profiler
)
endif()

#==============================================================================
macro(benchmark_test target file)
Expand Down
25 changes: 25 additions & 0 deletions benchmark/hdf5/benchmark_knowhere_float.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
#include "knowhere/comp/knowhere_config.h"
#include "knowhere/dataset.h"

const int32_t GPU_DEVICE_ID = 0;

class Benchmark_knowhere_float : public Benchmark_knowhere, public ::testing::Test {
public:
void
Expand Down Expand Up @@ -108,11 +110,18 @@ class Benchmark_knowhere_float : public Benchmark_knowhere, public ::testing::Te
cfg_[knowhere::meta::METRIC_TYPE] = metric_type_;
knowhere::KnowhereConfig::SetSimdType(knowhere::KnowhereConfig::SimdType::AVX2);
printf("faiss::distance_compute_blas_threshold: %ld\n", knowhere::KnowhereConfig::GetBlasThreshold());
#ifdef USE_CUDA
knowhere::KnowhereConfig::InitGPUResource(GPU_DEVICE_ID);
cfg_[knowhere::meta::DEVICE_ID] = GPU_DEVICE_ID;
#endif
}

void
TearDown() override {
free_all();
#ifdef USE_CUDA
knowhere::KnowhereConfig::FreeGPUResource();
#endif
}

protected:
Expand All @@ -138,7 +147,11 @@ class Benchmark_knowhere_float : public Benchmark_knowhere, public ::testing::Te
};

TEST_F(Benchmark_knowhere_float, TEST_IDMAP) {
#ifdef USE_CUDA
index_type_ = knowhere::IndexEnum::INDEX_FAISS_GPU_IDMAP;
#else
index_type_ = knowhere::IndexEnum::INDEX_FAISS_IDMAP;
#endif

knowhere::Json conf = cfg_;
std::string index_file_name = get_index_name({});
Expand All @@ -149,7 +162,11 @@ TEST_F(Benchmark_knowhere_float, TEST_IDMAP) {
}

TEST_F(Benchmark_knowhere_float, TEST_IVF_FLAT_NM) {
#ifdef USE_CUDA
index_type_ = knowhere::IndexEnum::INDEX_FAISS_GPU_IVFFLAT;
#else
index_type_ = knowhere::IndexEnum::INDEX_FAISS_IVFFLAT;
#endif

knowhere::Json conf = cfg_;
for (auto nlist : NLISTs_) {
Expand All @@ -171,7 +188,11 @@ TEST_F(Benchmark_knowhere_float, TEST_IVF_FLAT_NM) {
}

TEST_F(Benchmark_knowhere_float, TEST_IVF_SQ8) {
#ifdef USE_CUDA
index_type_ = knowhere::IndexEnum::INDEX_FAISS_GPU_IVFSQ8;
#else
index_type_ = knowhere::IndexEnum::INDEX_FAISS_IVFSQ8;
#endif

knowhere::Json conf = cfg_;
for (auto nlist : NLISTs_) {
Expand All @@ -186,7 +207,11 @@ TEST_F(Benchmark_knowhere_float, TEST_IVF_SQ8) {
}

TEST_F(Benchmark_knowhere_float, TEST_IVF_PQ) {
#ifdef USE_CUDA
index_type_ = knowhere::IndexEnum::INDEX_FAISS_GPU_IVFPQ;
#else
index_type_ = knowhere::IndexEnum::INDEX_FAISS_IVFPQ;
#endif

knowhere::Json conf = cfg_;
conf[knowhere::indexparam::NBITS] = NBITS_;
Expand Down
96 changes: 96 additions & 0 deletions include/knowhere/comp/blocking_queue.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software distributed under the License
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
// or implied. See the License for the specific language governing permissions and limitations under the License.

#pragma once

#include <assert.h>

#include <condition_variable>
#include <iostream>
#include <queue>
#include <vector>

namespace knowhere {

template <typename T>
class BlockingQueue {
public:
BlockingQueue() : mtx(), full_(), empty_() {
}

virtual ~BlockingQueue() {
}

BlockingQueue(const BlockingQueue& rhs) = delete;

BlockingQueue&
operator=(const BlockingQueue& rhs) = delete;

void
Put(const T& task) {
std::unique_lock<std::mutex> lock(mtx);
full_.wait(lock, [this] { return (queue_.size() < capacity_); });
queue_.push(task);
empty_.notify_all();
}

T
Take() {
std::unique_lock<std::mutex> lock(mtx);
empty_.wait(lock, [this] { return !queue_.empty(); });
T front(queue_.front());
queue_.pop();
full_.notify_all();
return front;
}

T
Front() {
std::unique_lock<std::mutex> lock(mtx);
empty_.wait(lock, [this] { return !queue_.empty(); });
T front(queue_.front());
return front;
}

T
Back() {
std::unique_lock<std::mutex> lock(mtx);
empty_.wait(lock, [this] { return !queue_.empty(); });
T back(queue_.back());
return back;
}

size_t
Size() const {
std::lock_guard<std::mutex> lock(mtx);
return queue_.size();
}

bool
Empty() const {
std::unique_lock<std::mutex> lock(mtx);
return queue_.empty();
}

void
SetCapacity(const size_t capacity) {
capacity_ = (capacity > 0 ? capacity : capacity_);
}

protected:
mutable std::mutex mtx;
std::condition_variable full_;
std::condition_variable empty_;
std::queue<T> queue_;
size_t capacity_ = 32;
};

} // namespace knowhere
5 changes: 5 additions & 0 deletions include/knowhere/comp/index_param.h
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,11 @@ constexpr const char* INDEX_FAISS_IVFFLAT = "IVF_FLAT";
constexpr const char* INDEX_FAISS_IVFPQ = "IVF_PQ";
constexpr const char* INDEX_FAISS_IVFSQ8 = "IVF_SQ8";

constexpr const char* INDEX_FAISS_GPU_IDMAP = "GPU_FLAT";
constexpr const char* INDEX_FAISS_GPU_IVFFLAT = "GPU_IVF_FLAT";
constexpr const char* INDEX_FAISS_GPU_IVFPQ = "GPU_IVF_PQ";
constexpr const char* INDEX_FAISS_GPU_IVFSQ8 = "GPU_IVF_SQ8";

constexpr const char* INDEX_ANNOY = "ANNOY";
constexpr const char* INDEX_HNSW = "HNSW";

Expand Down
12 changes: 12 additions & 0 deletions include/knowhere/comp/knowhere_config.h
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,18 @@ class KnowhereConfig {
*/
static void
SetAioContextPool(size_t num_ctx, size_t max_events);

/**
* init GPU Resource
*/
static void
InitGPUResource(int64_t gpu_id, int64_t res_num = 2);

/**
* free GPU Resource
*/
static void
FreeGPUResource();
};

} // namespace knowhere
Expand Down
Loading

0 comments on commit 6d5a17a

Please sign in to comment.