Skip to content

Commit

Permalink
Using consistent naming for faiss algos
Browse files Browse the repository at this point in the history
  • Loading branch information
@cjnolet
cjnolet committed Sep 25, 2023
1 parent 3fcd1e9 commit 1ec75ba
Show file tree
Hide file tree
Showing 18 changed files with 625 additions and 649 deletions.
40 changes: 21 additions & 19 deletions cpp/bench/ann/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,9 @@
# ##################################################################################################
# * benchmark options ------------------------------------------------------------------------------

option(RAFT_ANN_BENCH_USE_FAISS_FLAT "Include faiss' brute-force knn algorithm in benchmark" ON)
option(RAFT_ANN_BENCH_USE_FAISS_IVF_FLAT "Include faiss' ivf flat algorithm in benchmark" ON)
option(RAFT_ANN_BENCH_USE_FAISS_IVF_PQ "Include faiss' ivf pq algorithm in benchmark" ON)
option(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT "Include faiss' brute-force knn algorithm in benchmark" ON)
option(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT "Include faiss' ivf flat algorithm in benchmark" ON)
option(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_PQ "Include faiss' ivf pq algorithm in benchmark" ON)
option(RAFT_ANN_BENCH_USE_FAISS_CPU_FLAT
"Include faiss' cpu brute-force knn algorithm in benchmark" ON
)
Expand Down Expand Up @@ -48,9 +48,9 @@ if(BUILD_CPU_ONLY)
include(cmake/thirdparty/get_spdlog.cmake)

set(RAFT_FAISS_ENABLE_GPU OFF)
set(RAFT_ANN_BENCH_USE_FAISS_FLAT OFF)
set(RAFT_ANN_BENCH_USE_FAISS_IVF_FLAT OFF)
set(RAFT_ANN_BENCH_USE_FAISS_IVF_PQ OFF)
set(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT OFF)
set(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT OFF)
set(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_PQ OFF)
set(RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT OFF)
set(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ OFF)
set(RAFT_ANN_BENCH_USE_RAFT_CAGRA OFF)
Expand All @@ -60,17 +60,17 @@ else()
# https://github.com/rapidsai/raft/issues/1627
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0.0)
set(RAFT_FAISS_ENABLE_GPU OFF)
set(RAFT_ANN_BENCH_USE_FAISS_FLAT OFF)
set(RAFT_ANN_BENCH_USE_FAISS_IVF_FLAT OFF)
set(RAFT_ANN_BENCH_USE_FAISS_IVF_PQ OFF)
set(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT OFF)
set(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT OFF)
set(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_PQ OFF)
else()
set(RAFT_FAISS_ENABLE_GPU ON)
endif()
endif()

set(RAFT_ANN_BENCH_USE_FAISS OFF)
if(RAFT_ANN_BENCH_USE_FAISS_FLAT
OR RAFT_ANN_BENCH_USE_FAISS_IVFPQ
if(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT
OR RAFT_ANN_BENCH_USE_FAISS_GPU_IVFPQ
OR RAFT_ANN_BENCH_USE_FAISS_IFFLAT
OR RAFT_ANN_BENCH_USE_FAISS_CPU_FLAT
OR RAFT_ANN_BENCH_USE_FAISS_CPU_IVFPQ
Expand Down Expand Up @@ -243,8 +243,8 @@ if(RAFT_ANN_BENCH_USE_RAFT_CAGRA)
endif()

set(RAFT_FAISS_TARGETS faiss::faiss)
if(CXX_AVX_FOUND)
list(APPEND RAFT_FAISS_TARGETS faiss::faiss_avx2)
if(TARGET faiss::faiss_avx2)
set(RAFT_FAISS_TARGETS faiss::faiss_avx2)
endif()

if(RAFT_ANN_BENCH_USE_FAISS_CPU_FLAT)
Expand All @@ -268,21 +268,23 @@ if(RAFT_ANN_BENCH_USE_FAISS_CPU_IVF_PQ)
)
endif()

if(RAFT_ANN_BENCH_USE_FAISS_IVF_FLAT)
if(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT)
ConfigureAnnBench(
NAME FAISS_IVF_FLAT PATH bench/ann/src/faiss/faiss_benchmark.cu LINKS ${RAFT_FAISS_TARGETS}
NAME FAISS_GPU_IVF_FLAT PATH bench/ann/src/faiss/faiss_gpu_benchmark.cu LINKS
${RAFT_FAISS_TARGETS}
)
endif()

if(RAFT_ANN_BENCH_USE_FAISS_IVF_PQ)
if(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_PQ)
ConfigureAnnBench(
NAME FAISS_IVF_PQ PATH bench/ann/src/faiss/faiss_benchmark.cu LINKS ${RAFT_FAISS_TARGETS}
NAME FAISS_GPU_IVF_PQ PATH bench/ann/src/faiss/faiss_gpu_benchmark.cu LINKS
${RAFT_FAISS_TARGETS}
)
endif()

if(RAFT_ANN_BENCH_USE_FAISS_FLAT)
if(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT)
ConfigureAnnBench(
NAME FAISS_FLAT PATH bench/ann/src/faiss/faiss_benchmark.cu LINKS ${RAFT_FAISS_TARGETS}
NAME FAISS_GPU_FLAT PATH bench/ann/src/faiss/faiss_gpu_benchmark.cu LINKS ${RAFT_FAISS_TARGETS}
)
endif()

Expand Down
32 changes: 16 additions & 16 deletions cpp/bench/ann/src/faiss/faiss_cpu_benchmark.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,22 +31,22 @@ namespace raft::bench::ann {

template <typename T>
void parse_base_build_param(const nlohmann::json& conf,
typename raft::bench::ann::Faiss<T>::BuildParam& param)
typename raft::bench::ann::FaissCpu<T>::BuildParam& param)
{
param.nlist = conf.at("nlist");
if (conf.contains("ratio")) { param.ratio = conf.at("ratio"); }
}

template <typename T>
void parse_build_param(const nlohmann::json& conf,
typename raft::bench::ann::FaissIVFFlat<T>::BuildParam& param)
typename raft::bench::ann::FaissCpuIVFFlat<T>::BuildParam& param)
{
parse_base_build_param<T>(conf, param);
}

template <typename T>
void parse_build_param(const nlohmann::json& conf,
typename raft::bench::ann::FaissIVFPQ<T>::BuildParam& param)
typename raft::bench::ann::FaissCpuIVFPQ<T>::BuildParam& param)
{
parse_base_build_param<T>(conf, param);
param.M = conf.at("M");
Expand All @@ -64,15 +64,15 @@ void parse_build_param(const nlohmann::json& conf,

template <typename T>
void parse_build_param(const nlohmann::json& conf,
typename raft::bench::ann::FaissIVFSQ<T>::BuildParam& param)
typename raft::bench::ann::FaissCpuIVFSQ<T>::BuildParam& param)
{
parse_base_build_param<T>(conf, param);
param.quantizer_type = conf.at("quantizer_type");
}

template <typename T>
void parse_search_param(const nlohmann::json& conf,
typename raft::bench::ann::Faiss<T>::SearchParam& param)
typename raft::bench::ann::FaissCpu<T>::SearchParam& param)
{
param.nprobe = conf.at("nprobe");
if (conf.contains("refine_ratio")) { param.refine_ratio = conf.at("refine_ratio"); }
Expand Down Expand Up @@ -115,14 +115,14 @@ std::unique_ptr<raft::bench::ann::ANN<T>> create_algo(const std::string& algo,

if constexpr (std::is_same_v<T, float>) {
raft::bench::ann::Metric metric = parse_metric(distance);
if (algo == "faiss_ivf_flat") {
ann = make_algo<T, raft::bench::ann::FaissIVFFlat>(metric, dim, conf, dev_list);
} else if (algo == "faiss_ivf_pq") {
ann = make_algo<T, raft::bench::ann::FaissIVFPQ>(metric, dim, conf);
} else if (algo == "faiss_ivf_sq") {
ann = make_algo<T, raft::bench::ann::FaissIVFSQ>(metric, dim, conf);
} else if (algo == "faiss_flat") {
ann = std::make_unique<raft::bench::ann::FaissFlat<T>>(metric, dim);
if (algo == "faiss_gpu_ivf_flat") {
ann = make_algo<T, raft::bench::ann::FaissCpuIVFFlat>(metric, dim, conf, dev_list);
} else if (algo == "faiss_gpu_ivf_pq") {
ann = make_algo<T, raft::bench::ann::FaissCpuIVFPQ>(metric, dim, conf);
} else if (algo == "faiss_gpu_ivf_sq") {
ann = make_algo<T, raft::bench::ann::FaissCpuIVFSQ>(metric, dim, conf);
} else if (algo == "faiss_gpu_flat") {
ann = std::make_unique<raft::bench::ann::FaissCpuFlat<T>>(metric, dim);
}
}

Expand All @@ -137,11 +137,11 @@ template <typename T>
std::unique_ptr<typename raft::bench::ann::ANN<T>::AnnSearchParam> create_search_param(
const std::string& algo, const nlohmann::json& conf)
{
if (algo == "faiss_ivf_flat" || algo == "faiss_ivf_pq" || algo == "faiss_ivf_sq") {
auto param = std::make_unique<typename raft::bench::ann::Faiss<T>::SearchParam>();
if (algo == "faiss_gpu_ivf_flat" || algo == "faiss_gpu_ivf_pq" || algo == "faiss_gpu_ivf_sq") {
auto param = std::make_unique<typename raft::bench::ann::FaissCpu<T>::SearchParam>();
parse_search_param<T>(conf, *param);
return param;
} else if (algo == "faiss_flat") {
} else if (algo == "faiss_gpu_flat") {
auto param = std::make_unique<typename raft::bench::ann::ANN<T>::AnnSearchParam>();
return param;
}
Expand Down
77 changes: 27 additions & 50 deletions cpp/bench/ann/src/faiss/faiss_cpu_wrapper.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@
#include <faiss/IndexRefine.h>
#include <faiss/IndexScalarQuantizer.h>
#include <faiss/index_io.h>
#include <omp.h>

#include <cassert>
#include <memory>
Expand All @@ -44,32 +43,12 @@ faiss::MetricType parse_metric_type(raft::bench::ann::Metric metric)
throw std::runtime_error("faiss supports only metric type of inner product and L2");
}
}

// note BLAS library can still use multi-threading, and
// setting environment variable like OPENBLAS_NUM_THREADS can control it
class OmpSingleThreadScope {
public:
OmpSingleThreadScope()
{
max_threads_ = omp_get_max_threads();
omp_set_num_threads(1);
}
~OmpSingleThreadScope()
{
// the best we can do
omp_set_num_threads(max_threads_);
}

private:
int max_threads_;
};

} // namespace

namespace raft::bench::ann {

template <typename T>
class Faiss : public ANN<T> {
class FaissCpu : public ANN<T> {
public:
using typename ANN<T>::AnnSearchParam;
struct SearchParam : public AnnSearchParam {
Expand All @@ -82,7 +61,7 @@ class Faiss : public ANN<T> {
int ratio = 2;
};

Faiss(Metric metric, int dim, const BuildParam& param)
FaissCpu(Metric metric, int dim, const BuildParam& param)
: ANN<T>(metric, dim),
metric_type_(parse_metric_type(metric)),
nlist_{param.nlist},
Expand All @@ -91,7 +70,7 @@ class Faiss : public ANN<T> {
static_assert(std::is_same_v<T, float>, "faiss support only float type");
}

virtual ~Faiss() noexcept {}
virtual ~FaissCpu() noexcept {}

void build(const T* dataset, size_t nrow, cudaStream_t stream = 0) final;

Expand Down Expand Up @@ -140,9 +119,8 @@ class Faiss : public ANN<T> {
};

template <typename T>
void Faiss<T>::build(const T* dataset, size_t nrow, cudaStream_t stream)
void FaissCpu<T>::build(const T* dataset, size_t nrow, cudaStream_t stream)
{
OmpSingleThreadScope omp_single_thread;
auto index_ivf = dynamic_cast<faiss::IndexIVF*>(index_.get());
if (index_ivf != nullptr) {
// set the min/max training size for clustering to use the whole provided training set.
Expand Down Expand Up @@ -171,7 +149,7 @@ void Faiss<T>::build(const T* dataset, size_t nrow, cudaStream_t stream)
}

template <typename T>
void Faiss<T>::set_search_param(const AnnSearchParam& param)
void FaissCpu<T>::set_search_param(const AnnSearchParam& param)
{
auto search_param = dynamic_cast<const SearchParam&>(param);
int nprobe = search_param.nprobe;
Expand All @@ -185,12 +163,12 @@ void Faiss<T>::set_search_param(const AnnSearchParam& param)
}

template <typename T>
void Faiss<T>::search(const T* queries,
int batch_size,
int k,
size_t* neighbors,
float* distances,
cudaStream_t stream) const
void FaissCpu<T>::search(const T* queries,
int batch_size,
int k,
size_t* neighbors,
float* distances,
cudaStream_t stream) const
{
static_assert(sizeof(size_t) == sizeof(faiss::idx_t),
"sizes of size_t and faiss::idx_t are different");
Expand All @@ -199,26 +177,24 @@ void Faiss<T>::search(const T* queries,

template <typename T>
template <typename Index>
void Faiss<T>::save_(const std::string& file) const
void FaissCpu<T>::save_(const std::string& file) const
{
OmpSingleThreadScope omp_single_thread;
faiss::write_index(index_.get(), file.c_str());
}

template <typename T>
template <typename Index>
void Faiss<T>::load_(const std::string& file)
void FaissCpu<T>::load_(const std::string& file)
{
OmpSingleThreadScope omp_single_thread;
index_ = std::unique_ptr<Index>(dynamic_cast<Index*>(faiss::read_index(file.c_str())));
}

template <typename T>
class FaissIVFFlat : public Faiss<T> {
class FaissCpuIVFFlat : public FaissCpu<T> {
public:
using typename Faiss<T>::BuildParam;
using typename FaissCpu<T>::BuildParam;

FaissIVFFlat(Metric metric, int dim, const BuildParam& param) : Faiss<T>(metric, dim, param)
FaissCpuIVFFlat(Metric metric, int dim, const BuildParam& param) : FaissCpu<T>(metric, dim, param)
{
this->init_quantizer(dim);
this->index_ = std::make_unique<faiss::IndexIVFFlat>(
Expand All @@ -233,15 +209,15 @@ class FaissIVFFlat : public Faiss<T> {
};

template <typename T>
class FaissIVFPQ : public Faiss<T> {
class FaissCpuIVFPQ : public FaissCpu<T> {
public:
struct BuildParam : public Faiss<T>::BuildParam {
struct BuildParam : public FaissCpu<T>::BuildParam {
int M;
int bitsPerCode;
bool usePrecomputed;
};

FaissIVFPQ(Metric metric, int dim, const BuildParam& param) : Faiss<T>(metric, dim, param)
FaissCpuIVFPQ(Metric metric, int dim, const BuildParam& param) : FaissCpu<T>(metric, dim, param)
{
this->init_quantizer(dim);
this->index_ = std::make_unique<faiss::IndexIVFPQ>(
Expand All @@ -256,21 +232,21 @@ class FaissIVFPQ : public Faiss<T> {
};

template <typename T>
class FaissIVFSQ : public Faiss<T> {
class FaissCpuIVFSQ : public FaissCpu<T> {
public:
struct BuildParam : public Faiss<T>::BuildParam {
struct BuildParam : public FaissCpu<T>::BuildParam {
std::string quantizer_type;
};

FaissIVFSQ(Metric metric, int dim, const BuildParam& param) : Faiss<T>(metric, dim, param)
FaissCpuIVFSQ(Metric metric, int dim, const BuildParam& param) : FaissCpu<T>(metric, dim, param)
{
faiss::ScalarQuantizer::QuantizerType qtype;
if (param.quantizer_type == "fp16") {
qtype = faiss::ScalarQuantizer::QT_fp16;
} else if (param.quantizer_type == "int8") {
qtype = faiss::ScalarQuantizer::QT_8bit;
} else {
throw std::runtime_error("FaissIVFSQ supports only fp16 and int8 but got " +
throw std::runtime_error("FaissCpuIVFSQ supports only fp16 and int8 but got " +
param.quantizer_type);
}

Expand All @@ -290,14 +266,15 @@ class FaissIVFSQ : public Faiss<T> {
};

template <typename T>
class FaissFlat : public Faiss<T> {
class FaissCpuFlat : public FaissCpu<T> {
public:
FaissFlat(Metric metric, int dim) : Faiss<T>(metric, dim, typename Faiss<T>::BuildParam{})
FaissCpuFlat(Metric metric, int dim)
: FaissCpu<T>(metric, dim, typename FaissCpu<T>::BuildParam{})
{
this->index_ = std::make_unique<faiss::IndexFlat>(dim, this->metric_type_);
}

// class Faiss is more like a IVF class, so need special treating here
// class FaissCpu is more like a IVF class, so need special treating here
void set_search_param(const typename ANN<T>::AnnSearchParam&) override{};

void save(const std::string& file) const override
Expand Down
2 changes: 1 addition & 1 deletion cpp/bench/ann/src/faiss/faiss_benchmark.cu → ...ench/ann/src/faiss/faiss_gpu_benchmark.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@

#include "../common/ann_types.hpp"
#undef WARP_SIZE
#include "faiss_wrapper.h"
#include "faiss_gpu_wrapper.h"
#define JSON_DIAGNOSTICS 1
#include <nlohmann/json.hpp>

Expand Down
Loading

0 comments on commit 1ec75ba

Please sign in to comment.