Skip to content
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.

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

Use FAISS with RMM #363

Merged
merged 14 commits into from
Jan 19, 2022
2 changes: 1 addition & 1 deletion cpp/include/raft/spatial/knn/ann.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
#include "detail/ann_quantized_faiss.cuh"

#include <faiss/gpu/GpuIndex.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <raft/spatial/knn/faiss_mr.hpp>

#include <raft/mr/device/buffer.hpp>

Expand Down
4 changes: 2 additions & 2 deletions cpp/include/raft/spatial/knn/ann_common.h
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
#include <raft/linalg/distance_type.h>

#include <faiss/gpu/GpuIndex.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <raft/spatial/knn/faiss_mr.hpp>

namespace raft {
namespace spatial {
Expand All @@ -30,7 +30,7 @@ struct knnIndex {
raft::distance::DistanceType metric;
float metricArg;

faiss::gpu::StandardGpuResources *gpu_res;
raft::spatial::knn::RmmGpuResources *gpu_res;
int device;
~knnIndex() {
delete index;
Expand Down
6 changes: 3 additions & 3 deletions cpp/include/raft/spatial/knn/detail/ann_quantized_faiss.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -28,14 +28,14 @@

#include <label/classlabels.cuh>
#include <raft/distance/distance.hpp>
#include <raft/spatial/knn/faiss_mr.hpp>

#include <faiss/gpu/GpuDistance.h>
#include <faiss/gpu/GpuIndexFlat.h>
#include <faiss/gpu/GpuIndexIVFFlat.h>
#include <faiss/gpu/GpuIndexIVFPQ.h>
#include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
#include <faiss/gpu/GpuResources.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/utils/Heap.h>
#include <faiss/gpu/utils/Limits.cuh>
#include <faiss/gpu/utils/Select.cuh>
Expand Down Expand Up @@ -131,8 +131,8 @@ void approx_knn_build_index(raft::handle_t &handle,
int device;
CUDA_CHECK(cudaGetDevice(&device));

faiss::gpu::StandardGpuResources *gpu_res =
new faiss::gpu::StandardGpuResources();
raft::spatial::knn::RmmGpuResources *gpu_res =
new raft::spatial::knn::RmmGpuResources();
gpu_res->noTempMemory();
gpu_res->setDefaultStream(device, handle.get_stream());
index->gpu_res = gpu_res;
Expand Down
2 changes: 1 addition & 1 deletion cpp/include/raft/spatial/knn/detail/haversine_distance.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -21,13 +21,13 @@

#include <faiss/gpu/GpuDistance.h>
#include <faiss/gpu/GpuResources.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/utils/Heap.h>
#include <faiss/gpu/utils/Limits.cuh>
#include <faiss/gpu/utils/Select.cuh>

#include <raft/linalg/distance_type.h>
#include <raft/handle.hpp>
#include <raft/spatial/knn/faiss_mr.hpp>

namespace raft {
namespace spatial {
Expand Down
4 changes: 2 additions & 2 deletions cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,6 @@

#include <faiss/gpu/GpuDistance.h>
#include <faiss/gpu/GpuResources.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/utils/Heap.h>
#include <faiss/gpu/utils/Limits.cuh>
#include <faiss/gpu/utils/Select.cuh>
Expand All @@ -33,6 +32,7 @@
#include <cstdint>
#include <iostream>
#include <raft/handle.hpp>
#include <raft/spatial/knn/faiss_mr.hpp>
#include <set>

#include "fused_l2_knn.cuh"
Expand Down Expand Up @@ -284,7 +284,7 @@ void brute_force_knn_impl(std::vector<float *> &input,
default:
faiss::MetricType m = build_faiss_metric(metric);

faiss::gpu::StandardGpuResources gpu_res;
raft::spatial::knn::RmmGpuResources gpu_res;

gpu_res.noTempMemory();
gpu_res.setDefaultStream(device, stream);
Expand Down
2 changes: 1 addition & 1 deletion cpp/include/raft/spatial/knn/detail/selection_faiss.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,11 @@
#pragma once

#include <raft/cudart_utils.h>
#include <raft/spatial/knn/faiss_mr.hpp>

#include <faiss/gpu/GpuDistance.h>
#include <faiss/gpu/GpuIndexFlat.h>
#include <faiss/gpu/GpuResources.h>
#include <faiss/gpu/StandardGpuResources.h>
#include <faiss/utils/Heap.h>
#include <faiss/gpu/utils/Limits.cuh>
#include <faiss/gpu/utils/Select.cuh>
Expand Down
Loading