use BitsetView instead of ConcurrentBitsetptr (#4428)

* use bitsetview instead of concurrentbitsetptr

Signed-off-by: fluorinedog <fluorinedog@gmail.com>

* fix cpplint check

Signed-off-by: shengjun.li <shengjun.li@zilliz.com>

* fix gpu

Signed-off-by: fluorinedog <fluorinedog@gmail.com>

Co-authored-by: shengjun.li <shengjun.li@zilliz.com>
pull/4438/head
FluorineDog 2020-12-10 12:57:08 +08:00 committed by GitHub
parent b12309f742
commit 33dba268d5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
119 changed files with 427 additions and 311 deletions

View File

@ -109,7 +109,7 @@ IndexAnnoy::BuildAll(const DatasetPtr& dataset_ptr, const Config& config) {
}
DatasetPtr
IndexAnnoy::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
IndexAnnoy::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_) {
KNOWHERE_THROW_MSG("index not initialize or trained");
}

View File

@ -54,7 +54,7 @@ class IndexAnnoy : public VecIndex {
}
DatasetPtr
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) override;
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) override;
int64_t
Count() override;

View File

@ -46,7 +46,7 @@ BinaryIDMAP::Load(const BinarySet& index_binary) {
}
DatasetPtr
BinaryIDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
BinaryIDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_) {
KNOWHERE_THROW_MSG("index not initialize");
}
@ -148,7 +148,7 @@ BinaryIDMAP::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config)
void
BinaryIDMAP::QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels,
const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
const Config& config, const faiss::BitsetView& bitset) {
// assign the metric type
auto bin_flat_index = dynamic_cast<faiss::IndexBinaryIDMap*>(index_.get())->index;
bin_flat_index->metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());

View File

@ -48,7 +48,7 @@ class BinaryIDMAP : public VecIndex, public FaissBaseBinaryIndex {
AddWithoutIds(const DatasetPtr&, const Config&) override;
DatasetPtr
Query(const DatasetPtr&, const Config&, const faiss::ConcurrentBitsetPtr& bitset) override;
Query(const DatasetPtr&, const Config&, const faiss::BitsetView& bitset) override;
int64_t
Count() override;
@ -70,7 +70,7 @@ class BinaryIDMAP : public VecIndex, public FaissBaseBinaryIndex {
protected:
virtual void
QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset);
const faiss::BitsetView& bitset);
protected:
std::mutex mutex_;

View File

@ -48,7 +48,7 @@ BinaryIVF::Load(const BinarySet& index_binary) {
}
DatasetPtr
BinaryIVF::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
BinaryIVF::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_ || !index_->is_trained) {
KNOWHERE_THROW_MSG("index not initialize or trained");
}
@ -133,7 +133,7 @@ BinaryIVF::GenParams(const Config& config) {
void
BinaryIVF::QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset) {
const faiss::BitsetView& bitset) {
auto params = GenParams(config);
auto ivf_index = dynamic_cast<faiss::IndexBinaryIVF*>(index_.get());
ivf_index->nprobe = params->nprobe;

View File

@ -60,7 +60,7 @@ class BinaryIVF : public VecIndex, public FaissBaseBinaryIndex {
}
DatasetPtr
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) override;
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) override;
int64_t
Count() override;
@ -77,7 +77,7 @@ class BinaryIVF : public VecIndex, public FaissBaseBinaryIndex {
virtual void
QueryImpl(int64_t n, const uint8_t* data, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset);
const faiss::BitsetView& bitset);
protected:
std::mutex mutex_;

View File

@ -140,7 +140,7 @@ IndexHNSW::Add(const DatasetPtr& dataset_ptr, const Config& config) {
}
DatasetPtr
IndexHNSW::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
IndexHNSW::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_) {
KNOWHERE_THROW_MSG("index not initialize or trained");
}

View File

@ -46,7 +46,7 @@ class IndexHNSW : public VecIndex {
}
DatasetPtr
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) override;
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) override;
int64_t
Count() override;

View File

@ -100,7 +100,7 @@ IDMAP::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
}
DatasetPtr
IDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
IDMAP::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_) {
KNOWHERE_THROW_MSG("index not initialize");
}
@ -229,7 +229,7 @@ IDMAP::GetVectorById(const DatasetPtr& dataset_ptr, const Config& config) {
void
IDMAP::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset) {
const faiss::BitsetView& bitset) {
// assign the metric type
auto flat_index = dynamic_cast<faiss::IndexIDMap*>(index_.get())->index;
flat_index->metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());

View File

@ -46,7 +46,7 @@ class IDMAP : public VecIndex, public FaissBaseIndex {
AddWithoutIds(const DatasetPtr&, const Config&) override;
DatasetPtr
Query(const DatasetPtr&, const Config&, const faiss::ConcurrentBitsetPtr&) override;
Query(const DatasetPtr&, const Config&, const faiss::BitsetView&) override;
#if 0
DatasetPtr
@ -80,7 +80,7 @@ class IDMAP : public VecIndex, public FaissBaseIndex {
protected:
virtual void
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::ConcurrentBitsetPtr&);
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::BitsetView&);
protected:
std::mutex mutex_;

View File

@ -104,7 +104,7 @@ IVF::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
}
DatasetPtr
IVF::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
IVF::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_ || !index_->is_trained) {
KNOWHERE_THROW_MSG("index not initialize or trained");
}
@ -326,7 +326,7 @@ IVF::GenParams(const Config& config) {
void
IVF::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset) {
const faiss::BitsetView& bitset) {
auto params = GenParams(config);
auto ivf_index = dynamic_cast<faiss::IndexIVF*>(index_.get());
ivf_index->nprobe = std::min(params->nprobe, ivf_index->invlists->nlist);

View File

@ -51,7 +51,7 @@ class IVF : public VecIndex, public FaissBaseIndex {
AddWithoutIds(const DatasetPtr&, const Config&) override;
DatasetPtr
Query(const DatasetPtr&, const Config&, const faiss::ConcurrentBitsetPtr&) override;
Query(const DatasetPtr&, const Config&, const faiss::BitsetView&) override;
#if 0
DatasetPtr
@ -86,7 +86,7 @@ class IVF : public VecIndex, public FaissBaseIndex {
GenParams(const Config&);
virtual void
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::ConcurrentBitsetPtr&);
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::BitsetView&);
void
SealImpl() override;

View File

@ -122,7 +122,7 @@ IndexNGT::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
#endif
DatasetPtr
IndexNGT::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
IndexNGT::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_) {
KNOWHERE_THROW_MSG("index not initialize");
}

View File

@ -54,7 +54,7 @@ class IndexNGT : public VecIndex {
}
DatasetPtr
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) override;
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) override;
int64_t
Count() override;

View File

@ -77,7 +77,7 @@ NSG::Load(const BinarySet& index_binary) {
}
DatasetPtr
NSG::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
NSG::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_ || !index_->is_trained) {
KNOWHERE_THROW_MSG("index not initialize or trained");
}

View File

@ -59,7 +59,7 @@ class NSG : public VecIndex {
}
DatasetPtr
Query(const DatasetPtr&, const Config&, const faiss::ConcurrentBitsetPtr&) override;
Query(const DatasetPtr&, const Config&, const faiss::BitsetView&) override;
int64_t
Count() override;

View File

@ -80,7 +80,7 @@ IndexRHNSW::Add(const DatasetPtr& dataset_ptr, const Config& config) {
}
DatasetPtr
IndexRHNSW::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
IndexRHNSW::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_) {
KNOWHERE_THROW_MSG("index not initialize or trained");
}

View File

@ -52,7 +52,7 @@ class IndexRHNSW : public VecIndex, public FaissBaseIndex {
AddWithoutIds(const DatasetPtr&, const Config&) override;
DatasetPtr
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) override;
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) override;
int64_t
Count() override;

View File

@ -180,7 +180,7 @@ CPUSPTAGRNG::SetParameters(const Config& config) {
}
DatasetPtr
CPUSPTAGRNG::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
CPUSPTAGRNG::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
SetParameters(config);
float* p_data = (float*)dataset_ptr->Get<const void*>(meta::TENSOR);

View File

@ -52,7 +52,7 @@ class CPUSPTAGRNG : public VecIndex {
}
DatasetPtr
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) override;
Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) override;
int64_t
Count() override;

View File

@ -11,6 +11,7 @@
#pragma once
#include <faiss/utils/BitsetView.h>
#include <faiss/utils/ConcurrentBitset.h>
#include <memory>
#include <utility>
@ -47,7 +48,7 @@ class VecIndex : public Index {
AddWithoutIds(const DatasetPtr& dataset, const Config& config) = 0;
virtual DatasetPtr
Query(const DatasetPtr& dataset, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) = 0;
Query(const DatasetPtr& dataset, const Config& config, const faiss::BitsetView& bitset) = 0;
#if 0
virtual DatasetPtr

View File

@ -105,7 +105,7 @@ GPUIDMAP::GetRawIds() {
void
GPUIDMAP::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset) {
const faiss::BitsetView& bitset) {
ResScope rs(res_, gpu_id_);
// assign the metric type

View File

@ -56,7 +56,7 @@ class GPUIDMAP : public IDMAP, public GPUIndex {
void
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&,
const faiss::ConcurrentBitsetPtr& bitset) override;
const faiss::BitsetView& bitset) override;
};
using GPUIDMAPPtr = std::shared_ptr<GPUIDMAP>;

View File

@ -138,7 +138,7 @@ GPUIVF::LoadImpl(const BinarySet& binary_set, const IndexType& type) {
void
GPUIVF::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset) {
const faiss::BitsetView& bitset) {
std::lock_guard<std::mutex> lk(mutex_);
auto device_index = std::dynamic_pointer_cast<faiss::gpu::GpuIndexIVF>(index_);

View File

@ -52,7 +52,7 @@ class GPUIVF : public IVF, public GPUIndex {
void
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&,
const faiss::ConcurrentBitsetPtr& bitset) override;
const faiss::BitsetView& bitset) override;
};
using GPUIVFPtr = std::shared_ptr<GPUIVF>;

View File

@ -242,7 +242,7 @@ IVFSQHybrid::LoadImpl(const BinarySet& binary_set, const IndexType& type) {
void
IVFSQHybrid::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset) {
const faiss::BitsetView& bitset) {
if (gpu_mode_ == 2) {
GPUIVF::QueryImpl(n, data, k, distances, labels, config, bitset);
// index_->search(n, (float*)data, k, distances, labels);

View File

@ -89,7 +89,7 @@ class IVFSQHybrid : public GPUIVFSQ {
void
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&,
const faiss::ConcurrentBitsetPtr& bitset) override;
const faiss::BitsetView& bitset) override;
protected:
int64_t gpu_mode_ = 0; // 0: CPU, 1: Hybrid, 2: GPU

View File

@ -851,7 +851,7 @@ NsgIndex::FindUnconnectedNode(float* data, boost::dynamic_bitset<>& has_linked,
void
NsgIndex::Search(const float* query, float* data, const unsigned& nq, const unsigned& dim, const unsigned& k,
float* dist, int64_t* ids, SearchParams& params, faiss::ConcurrentBitsetPtr bitset) {
float* dist, int64_t* ids, SearchParams& params, const faiss::BitsetView& bitset) {
std::vector<std::vector<Neighbor>> resset(nq);
TimeRecorder rc("NsgIndex::search", 1);
@ -873,7 +873,7 @@ NsgIndex::Search(const float* query, float* data, const unsigned& nq, const unsi
if (pos >= k) {
break; // already top k
}
if (!bitset || !bitset->test(node.id)) {
if (!bitset || !bitset.test(node.id)) {
ids[i * k + pos] = ids_[node.id];
dist[i * k + pos] = is_ip ? -node.distance : node.distance;
++pos;

View File

@ -84,7 +84,7 @@ class NsgIndex {
void
Search(const float* query, float* data, const unsigned& nq, const unsigned& dim, const unsigned& k, float* dist,
int64_t* ids, SearchParams& params, faiss::ConcurrentBitsetPtr bitset = nullptr);
int64_t* ids, SearchParams& params, const faiss::BitsetView& bitset = nullptr);
int64_t
GetSize();

View File

@ -145,7 +145,7 @@ IVF_NM::AddWithoutIds(const DatasetPtr& dataset_ptr, const Config& config) {
}
DatasetPtr
IVF_NM::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
IVF_NM::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_ || !index_->is_trained) {
KNOWHERE_THROW_MSG("index not initialize or trained");
}
@ -313,7 +313,7 @@ IVF_NM::GenParams(const Config& config) {
void
IVF_NM::QueryImpl(int64_t n, const float* query, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset) {
const faiss::BitsetView& bitset) {
auto params = GenParams(config);
auto ivf_index = dynamic_cast<faiss::IndexIVF*>(index_.get());
ivf_index->nprobe = params->nprobe;

View File

@ -51,7 +51,7 @@ class IVF_NM : public VecIndex, public OffsetBaseIndex {
AddWithoutIds(const DatasetPtr&, const Config&) override;
DatasetPtr
Query(const DatasetPtr&, const Config&, const faiss::ConcurrentBitsetPtr& bitset) override;
Query(const DatasetPtr&, const Config&, const faiss::BitsetView& bitset) override;
#if 0
DatasetPtr
@ -86,8 +86,7 @@ class IVF_NM : public VecIndex, public OffsetBaseIndex {
GenParams(const Config&);
virtual void
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&,
const faiss::ConcurrentBitsetPtr& bitset);
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&, const faiss::BitsetView& bitset);
void
SealImpl() override;

View File

@ -78,7 +78,7 @@ NSG_NM::Load(const BinarySet& index_binary) {
}
DatasetPtr
NSG_NM::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::ConcurrentBitsetPtr& bitset) {
NSG_NM::Query(const DatasetPtr& dataset_ptr, const Config& config, const faiss::BitsetView& bitset) {
if (!index_ || !index_->is_trained) {
KNOWHERE_THROW_MSG("index not initialize or trained");
}

View File

@ -59,7 +59,7 @@ class NSG_NM : public VecIndex {
}
DatasetPtr
Query(const DatasetPtr&, const Config&, const faiss::ConcurrentBitsetPtr& bitset) override;
Query(const DatasetPtr&, const Config&, const faiss::BitsetView& bitset) override;
int64_t
Count() override;

View File

@ -119,7 +119,7 @@ GPUIVF_NM::SerializeImpl(const IndexType& type) {
void
GPUIVF_NM::QueryImpl(int64_t n, const float* data, int64_t k, float* distances, int64_t* labels, const Config& config,
const faiss::ConcurrentBitsetPtr& bitset) {
const faiss::BitsetView& bitset) {
std::lock_guard<std::mutex> lk(mutex_);
auto device_index = std::dynamic_pointer_cast<faiss::gpu::GpuIndexIVF>(index_);

View File

@ -52,7 +52,7 @@ class GPUIVF_NM : public IVF, public GPUIndex {
void
QueryImpl(int64_t, const float*, int64_t, float*, int64_t*, const Config&,
const faiss::ConcurrentBitsetPtr& bitset) override;
const faiss::BitsetView& bitset) override;
protected:
uint8_t* arranged_data;

View File

@ -599,7 +599,7 @@ NeighborhoodGraph::setupSeeds(NGT::SearchContainer &sc, ObjectDistances &seeds,
}
// for milvus
void NeighborhoodGraph::search(NGT::SearchContainer & sc, ObjectDistances & seeds, const faiss::ConcurrentBitsetPtr & bitset)
void NeighborhoodGraph::search(NGT::SearchContainer & sc, ObjectDistances & seeds, const faiss::BitsetView& bitset)
{
if (sc.explorationCoefficient == 0.0)
{
@ -710,7 +710,7 @@ NeighborhoodGraph::setupSeeds(NGT::SearchContainer &sc, ObjectDistances &seeds,
distanceChecked.insert(neighbor.id);
// judge if id in blacklist
if (bitset != nullptr && bitset->test((faiss::ConcurrentBitset::id_type_t)neighbor.id - 1)) {
if (!bitset.empty() && bitset.test((faiss::ConcurrentBitset::id_type_t)neighbor.id - 1)) {
continue;
}

View File

@ -24,6 +24,7 @@
#include "NGT/ObjectSpaceRepository.h"
#include "faiss/utils/ConcurrentBitset.h"
#include "faiss/utils/BitsetView.h"
#include "NGT/HashBasedBooleanSet.h"
@ -697,7 +698,7 @@ namespace NGT {
void search(NGT::SearchContainer &sc, ObjectDistances &seeds);
// for milvus
void search(NGT::SearchContainer & sc, ObjectDistances & seeds, const faiss::ConcurrentBitsetPtr & bitset);
void search(NGT::SearchContainer & sc, ObjectDistances & seeds, const faiss::BitsetView&bitset);
#ifdef NGT_GRAPH_READ_ONLY_GRAPH
template <typename COMPARATOR, typename CHECK_LIST> void searchReadOnlyGraph(NGT::SearchContainer &sc, ObjectDistances &seeds);

View File

@ -655,7 +655,7 @@ public:
virtual void linearSearch(NGT::SearchContainer & sc) { getIndex().linearSearch(sc); }
virtual void linearSearch(NGT::SearchQuery & sc) { getIndex().linearSearch(sc); }
// for milvus
virtual void search(NGT::SearchContainer & sc, const faiss::ConcurrentBitsetPtr & bitset) { getIndex().search(sc, bitset); }
virtual void search(NGT::SearchContainer & sc, const faiss::BitsetView&bitset) { getIndex().search(sc, bitset); }
virtual void search(NGT::SearchContainer & sc) { getIndex().search(sc); }
virtual void search(NGT::SearchQuery & sc) { getIndex().search(sc); }
virtual void search(NGT::SearchContainer & sc, ObjectDistances & seeds) { getIndex().search(sc, seeds); }
@ -1058,7 +1058,7 @@ public:
}
// for milvus
virtual void search(NGT::SearchContainer & sc, const faiss::ConcurrentBitsetPtr & bitset)
virtual void search(NGT::SearchContainer & sc, const faiss::BitsetView&bitset)
{
sc.distanceComputationCount = 0;
sc.visitCount = 0;
@ -1586,7 +1586,7 @@ protected:
}
// for milvus
virtual void search(NGT::SearchContainer & sc, ObjectDistances & seeds, const faiss::ConcurrentBitsetPtr & bitset)
virtual void search(NGT::SearchContainer & sc, ObjectDistances & seeds, const faiss::BitsetView&bitset)
{
if (sc.size == 0)
{
@ -2147,7 +2147,7 @@ public:
// for milvus
void
getSeedsFromTree(NGT::SearchContainer& sc, ObjectDistances& seeds, const faiss::ConcurrentBitsetPtr& bitset) {
getSeedsFromTree(NGT::SearchContainer& sc, ObjectDistances& seeds, const faiss::BitsetView& bitset) {
DVPTree::SearchContainer tso(sc.object);
tso.mode = DVPTree::SearchContainer::SearchLeaf;
tso.radius = 0.0;
@ -2204,7 +2204,7 @@ public:
}
// for milvus
void search(NGT::SearchContainer & sc, const faiss::ConcurrentBitsetPtr & bitset)
void search(NGT::SearchContainer & sc, const faiss::BitsetView&bitset)
{
sc.distanceComputationCount = 0;
sc.visitCount = 0;

View File

@ -20,6 +20,7 @@
#include "NGT/Node.h"
#include "NGT/defines.h"
#include "faiss/utils/ConcurrentBitset.h"
#include "faiss/utils/BitsetView.h"
#include <sstream>
#include <string>
@ -262,7 +263,7 @@ namespace NGT {
// for milvus
void
getObjectIDsFromLeaf(Node::ID nid, ObjectDistances& rl, const faiss::ConcurrentBitsetPtr& bitset) {
getObjectIDsFromLeaf(Node::ID nid, ObjectDistances& rl, const faiss::BitsetView& bitset) {
LeafNode& ln = *(LeafNode*)getNode(nid);
rl.clear();
ObjectDistance r;
@ -274,7 +275,7 @@ namespace NGT {
r.id = ln.getObjectIDs()[i].id;
r.distance = ln.getObjectIDs()[i].distance;
#endif
if (bitset != nullptr && bitset->test(r.id - 1)) {
if (!bitset.empty() && bitset.test(r.id - 1)) {
continue;
}
rl.push_back(r);

View File

@ -121,6 +121,7 @@ inline void set_error_from_string(char **error, const char* msg) {
#elif defined(__GNUC__)
#include <x86intrin.h>
#include <src/index/thirdparty/faiss/utils/ConcurrentBitset.h>
#include <src/index/thirdparty/faiss/utils/BitsetView.h>
#endif
#endif
@ -838,9 +839,9 @@ class AnnoyIndexInterface {
virtual bool load_index(void* index_data, const int64_t& index_size, char** error = nullptr) = 0;
virtual T get_distance(S i, S j) const = 0;
virtual void get_nns_by_item(S item, size_t n, int64_t search_k, vector<S>* result, vector<T>* distances,
const faiss::ConcurrentBitsetPtr& bitset = nullptr) const = 0;
const faiss::BitsetView& bitset = nullptr) const = 0;
virtual void get_nns_by_vector(const T* w, size_t n, int64_t search_k, vector<S>* result, vector<T>* distances,
const faiss::ConcurrentBitsetPtr& bitset = nullptr) const = 0;
const faiss::BitsetView& bitset = nullptr) const = 0;
virtual S get_n_items() const = 0;
virtual S get_dim() const = 0;
virtual S get_n_trees() const = 0;
@ -1177,14 +1178,14 @@ public:
}
void get_nns_by_item(S item, size_t n, int64_t search_k, vector<S>* result, vector<T>* distances,
const faiss::ConcurrentBitsetPtr& bitset) const {
const faiss::BitsetView& bitset) const {
// TODO: handle OOB
const Node* m = _get(item);
_get_all_nns(m->v, n, search_k, result, distances, bitset);
}
void get_nns_by_vector(const T* w, size_t n, int64_t search_k, vector<S>* result, vector<T>* distances,
const faiss::ConcurrentBitsetPtr& bitset) const {
const faiss::BitsetView& bitset) const {
_get_all_nns(w, n, search_k, result, distances, bitset);
}
@ -1334,7 +1335,7 @@ protected:
}
void _get_all_nns(const T* v, size_t n, int64_t search_k, vector<S>* result, vector<T>* distances,
const faiss::ConcurrentBitsetPtr& bitset) const {
const faiss::BitsetView& bitset) const {
Node* v_node = (Node *)alloca(_s);
D::template zero_value<Node>(v_node);
memcpy(v_node->v, v, sizeof(T) * _f);
@ -1358,12 +1359,12 @@ protected:
Node* nd = _get(i);
q.pop();
if (nd->n_descendants == 1 && i < _n_items) { // raw data
if (bitset == nullptr || !bitset->test((faiss::ConcurrentBitset::id_type_t)i))
if (bitset.empty() || !bitset.test((faiss::ConcurrentBitset::id_type_t)i))
nns.push_back(i);
} else if (nd->n_descendants <= _K) {
const S* dst = nd->children;
for (auto ii = 0; ii < nd->n_descendants; ++ ii) {
if (bitset == nullptr || !bitset->test((faiss::ConcurrentBitset::id_type_t)dst[ii]))
if (bitset.empty() || !bitset.test((faiss::ConcurrentBitset::id_type_t)dst[ii]))
nns.push_back(dst[ii]);
// nns.insert(nns.end(), dst, &dst[nd->n_descendants]);
}

View File

@ -31,7 +31,7 @@ void Index::train(idx_t /*n*/, const float* /*x*/) {
void Index::range_search (idx_t , const float *, float,
RangeSearchResult *,
ConcurrentBitsetPtr) const
const BitsetView&) const
{
FAISS_THROW_MSG ("range search not implemented");
}
@ -62,12 +62,12 @@ void Index::add_with_ids_without_codes(idx_t n, const float* x, const idx_t* xid
}
#if 0
void Index::get_vector_by_id (idx_t n, const idx_t *xid, float *x, ConcurrentBitsetPtr bitset) {
void Index::get_vector_by_id (idx_t n, const idx_t *xid, float *x, const BitsetView& bitset) {
FAISS_THROW_MSG ("get_vector_by_id not implemented for this type of index");
}
void Index::search_by_id (idx_t n, const idx_t *xid, idx_t k, float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) {
const BitsetView& bitset) {
FAISS_THROW_MSG ("search_by_id not implemented for this type of index");
}
#endif

View File

@ -12,6 +12,7 @@
#include <faiss/MetricType.h>
#include <faiss/utils/ConcurrentBitset.h>
#include <faiss/utils/BitsetView.h>
#include <cstdio>
#include <typeinfo>
#include <string>
@ -128,7 +129,7 @@ struct Index {
*/
virtual void search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const = 0;
const BitsetView& bitset = nullptr) const = 0;
#if 0
/** query n raw vectors from the index by ids.
@ -140,7 +141,7 @@ struct Index {
* @param x output raw vectors, size n * d
* @param bitset flags to check the validity of vectors
*/
virtual void get_vector_by_id (idx_t n, const idx_t *xid, float *x, ConcurrentBitsetPtr bitset = nullptr);
virtual void get_vector_by_id (idx_t n, const idx_t *xid, float *x, const BitsetView& bitset = nullptr);
/** query n vectors of dimension d to the index by ids.
*
@ -153,7 +154,7 @@ struct Index {
* @param bitset flags to check the validity of vectors
*/
virtual void search_by_id (idx_t n, const idx_t *xid, idx_t k, float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr);
const BitsetView& bitset = nullptr);
#endif
/** query n vectors of dimension d to the index.
@ -168,7 +169,7 @@ struct Index {
*/
virtual void range_search (idx_t n, const float *x, float radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset = nullptr) const;
const BitsetView& bitset = nullptr) const;
/** return the indexes of the k vectors closest to the query x.
*

View File

@ -166,7 +166,7 @@ void Index2Layer::search(
idx_t /*k*/,
float* /*distances*/,
idx_t* /*labels*/,
ConcurrentBitsetPtr) const {
const BitsetView&) const {
FAISS_THROW_MSG("not implemented");
}

View File

@ -61,7 +61,7 @@ struct Index2Layer: Index {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reconstruct_n(idx_t i0, idx_t ni, float* recons) const override;

View File

@ -22,7 +22,7 @@ void IndexBinary::train(idx_t, const uint8_t *) {
void IndexBinary::range_search(idx_t, const uint8_t *, int,
RangeSearchResult *,
ConcurrentBitsetPtr) const {
const BitsetView&) const {
FAISS_THROW_MSG("range search not implemented");
}
@ -37,12 +37,12 @@ void IndexBinary::add_with_ids(idx_t, const uint8_t *, const idx_t *) {
}
#if 0
void IndexBinary::get_vector_by_id (idx_t n, const idx_t *xid, uint8_t *x, ConcurrentBitsetPtr bitset) {
void IndexBinary::get_vector_by_id (idx_t n, const idx_t *xid, uint8_t *x, const BitsetView& bitset) {
FAISS_THROW_MSG("get_vector_by_id not implemented for this type of index");
}
void IndexBinary::search_by_id (idx_t n, const idx_t *xid, idx_t k, int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) {
const BitsetView& bitset) {
FAISS_THROW_MSG("search_by_id not implemented for this type of index");
}
#endif

View File

@ -97,7 +97,7 @@ struct IndexBinary {
*/
virtual void search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const = 0;
const BitsetView& bitset = nullptr) const = 0;
#if 0
/** Query n raw vectors from the index by ids.
@ -109,7 +109,7 @@ struct IndexBinary {
* @param x output raw vectors, size n * d
* @param bitset flags to check the validity of vectors
*/
virtual void get_vector_by_id (idx_t n, const idx_t *xid, uint8_t *x, ConcurrentBitsetPtr bitset = nullptr);
virtual void get_vector_by_id (idx_t n, const idx_t *xid, uint8_t *x, const BitsetView& bitset = nullptr);
/** query n vectors of dimension d to the index by ids.
*
@ -122,7 +122,7 @@ struct IndexBinary {
* @param bitset flags to check the validity of vectors
*/
virtual void search_by_id (idx_t n, const idx_t *xid, idx_t k, int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr);
const BitsetView& bitset = nullptr);
#endif
/** Query n vectors of dimension d to the index.
@ -141,7 +141,7 @@ struct IndexBinary {
*/
virtual void range_search(idx_t n, const uint8_t *x, int radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset = nullptr) const;
const BitsetView& bitset = nullptr) const;
/** Return the indexes of the k vectors closest to the query x.
*

View File

@ -40,7 +40,7 @@ void IndexBinaryFlat::reset() {
void IndexBinaryFlat::search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
const idx_t block_size = query_batch_size;
if (metric_type == METRIC_Jaccard || metric_type == METRIC_Tanimoto) {
float *D = reinterpret_cast<float*>(distances);
@ -124,7 +124,7 @@ void IndexBinaryFlat::reconstruct(idx_t key, uint8_t *recons) const {
void IndexBinaryFlat::range_search(idx_t n, const uint8_t *x, int radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
hamming_range_search (x, xb.data(), n, ntotal, radius, code_size, result);
}

View File

@ -39,11 +39,11 @@ struct IndexBinaryFlat : IndexBinary {
void search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void range_search(idx_t n, const uint8_t *x, int radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reconstruct(idx_t key, uint8_t *recons) const override;

View File

@ -51,7 +51,7 @@ void IndexBinaryFromFloat::reset() {
void IndexBinaryFromFloat::search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
constexpr idx_t bs = 32768;
std::unique_ptr<float[]> xf(new float[bs * d]);
std::unique_ptr<float[]> df(new float[bs * k]);

View File

@ -42,7 +42,7 @@ struct IndexBinaryFromFloat : IndexBinary {
void search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void train(idx_t n, const uint8_t *x) override;
};

View File

@ -197,7 +197,7 @@ void IndexBinaryHNSW::train(idx_t n, const uint8_t *x)
void IndexBinaryHNSW::search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
#pragma omp parallel
{

View File

@ -46,7 +46,7 @@ struct IndexBinaryHNSW : IndexBinary {
/// entry point for search
void search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reconstruct(idx_t key, uint8_t* recons) const override;

View File

@ -213,7 +213,7 @@ search_single_query(const IndexBinaryHash & index, const uint8_t *q,
void IndexBinaryHash::range_search(idx_t n, const uint8_t *x, int radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
size_t nlist = 0, ndis = 0, n0 = 0;
@ -241,7 +241,7 @@ void IndexBinaryHash::range_search(idx_t n, const uint8_t *x, int radius,
void IndexBinaryHash::search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
using HeapForL2 = CMax<int32_t, idx_t>;
@ -430,7 +430,7 @@ search_1_query_multihash(const IndexBinaryMultiHash & index, const uint8_t *xi,
void IndexBinaryMultiHash::range_search(idx_t n, const uint8_t *x, int radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
size_t nlist = 0, ndis = 0, n0 = 0;
@ -458,7 +458,7 @@ void IndexBinaryMultiHash::range_search(idx_t n, const uint8_t *x, int radius,
void IndexBinaryMultiHash::search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
using HeapForL2 = CMax<int32_t, idx_t>;

View File

@ -52,11 +52,11 @@ struct IndexBinaryHash : IndexBinary {
void range_search(idx_t n, const uint8_t *x, int radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void display() const;
size_t hashtable_size() const;
@ -105,11 +105,11 @@ struct IndexBinaryMultiHash: IndexBinary {
void range_search(idx_t n, const uint8_t *x, int radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
size_t hashtable_size() const;

View File

@ -141,7 +141,7 @@ void IndexBinaryIVF::set_direct_map_type (DirectMap::Type type)
void IndexBinaryIVF::search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);
@ -158,12 +158,12 @@ void IndexBinaryIVF::search(idx_t n, const uint8_t *x, idx_t k,
}
#if 0
void IndexBinaryIVF::get_vector_by_id(idx_t n, const idx_t *xid, uint8_t *x, ConcurrentBitsetPtr bitset) {
void IndexBinaryIVF::get_vector_by_id(idx_t n, const idx_t *xid, uint8_t *x, const BitsetView& bitset) {
make_direct_map(true);
/* only get vector by 1 id */
FAISS_ASSERT(n == 1);
if (!bitset || !bitset->test(xid[0])) {
if (!bitset || !bitset.test(xid[0])) {
reconstruct(xid[0], x + 0 * d);
} else {
memset(x, UINT8_MAX, d * sizeof(uint8_t));
@ -171,7 +171,7 @@ void IndexBinaryIVF::get_vector_by_id(idx_t n, const idx_t *xid, uint8_t *x, Con
}
void IndexBinaryIVF::search_by_id (idx_t n, const idx_t *xid, idx_t k, int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) {
const BitsetView& bitset) {
make_direct_map(true);
auto x = new uint8_t[n * d];
@ -370,13 +370,13 @@ struct IVFBinaryScannerL2: BinaryInvertedListScanner {
const idx_t *ids,
int32_t *simi, idx_t *idxi,
size_t k,
ConcurrentBitsetPtr bitset) const override
const BitsetView& bitset) const override
{
using C = CMax<int32_t, idx_t>;
size_t nup = 0;
for (size_t j = 0; j < n; j++) {
if (!bitset || !bitset->test(ids[j])) {
if (!bitset || !bitset.test(ids[j])) {
uint32_t dis = hc.hamming (codes);
if (dis < simi[0]) {
idx_t id = store_pairs ? (list_no << 32 | j) : ids[j];
@ -433,13 +433,13 @@ struct IVFBinaryScannerJaccard: BinaryInvertedListScanner {
const idx_t *ids,
int32_t *simi, idx_t *idxi,
size_t k,
ConcurrentBitsetPtr bitset = nullptr) const override
const BitsetView& bitset = nullptr) const override
{
using C = CMax<float, idx_t>;
float* psimi = (float*)simi;
size_t nup = 0;
for (size_t j = 0; j < n; j++) {
if(!bitset || !bitset->test(ids[j])){
if(!bitset || !bitset.test(ids[j])){
float dis = hc.compute (codes);
if (dis < psimi[0]) {
@ -512,7 +512,7 @@ void search_knn_hamming_heap(const IndexBinaryIVF& ivf,
int32_t *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params,
ConcurrentBitsetPtr bitset = nullptr)
const BitsetView& bitset = nullptr)
{
long nprobe = params ? params->nprobe : ivf.nprobe;
long max_codes = params ? params->max_codes : ivf.max_codes;
@ -606,7 +606,7 @@ void search_knn_binary_dis_heap(const IndexBinaryIVF& ivf,
idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params,
ConcurrentBitsetPtr bitset = nullptr)
const BitsetView& bitset = nullptr)
{
long nprobe = params ? params->nprobe : ivf.nprobe;
long max_codes = params ? params->max_codes : ivf.max_codes;
@ -689,7 +689,7 @@ void search_knn_hamming_count(const IndexBinaryIVF& ivf,
int32_t *distances,
idx_t *labels,
const IVFSearchParameters *params,
ConcurrentBitsetPtr bitset = nullptr) {
const BitsetView& bitset = nullptr) {
const int nBuckets = ivf.d + 1;
std::vector<int> all_counters(nx * nBuckets, 0);
std::unique_ptr<idx_t[]> all_ids_per_dis(new idx_t[nx * nBuckets * k]);
@ -737,7 +737,7 @@ void search_knn_hamming_count(const IndexBinaryIVF& ivf,
: ivf.invlists->get_ids(key);
for (size_t j = 0; j < list_size; j++) {
if (!bitset || !bitset->test(ids[j])) {
if (!bitset || !bitset.test(ids[j])) {
const uint8_t *yj = list_vecs + ivf.code_size * j;
idx_t id = store_pairs ? (key << 32 | j) : ids[j];
csi.update_counter(yj, id);
@ -784,7 +784,7 @@ void search_knn_hamming_count_1 (
int32_t *distances,
idx_t *labels,
const IVFSearchParameters *params,
ConcurrentBitsetPtr bitset = nullptr) {
const BitsetView& bitset = nullptr) {
switch (ivf.code_size) {
#define HANDLE_CS(cs) \
case cs: \
@ -845,7 +845,7 @@ void IndexBinaryIVF::search_preassigned(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params,
ConcurrentBitsetPtr bitset
const BitsetView& bitset
) const {
if (metric_type == METRIC_Jaccard || metric_type == METRIC_Tanimoto) {
if (use_heap) {
@ -888,7 +888,7 @@ void IndexBinaryIVF::search_preassigned(idx_t n, const uint8_t *x, idx_t k,
void IndexBinaryIVF::range_search(
idx_t n, const uint8_t *x, int radius,
RangeSearchResult *res,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
std::unique_ptr<int32_t[]> coarse_dis(new int32_t[n * nprobe]);

View File

@ -105,7 +105,7 @@ struct IndexBinaryIVF : IndexBinary {
int32_t *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params=nullptr,
ConcurrentBitsetPtr bitset = nullptr
const BitsetView& bitset = nullptr
) const;
virtual BinaryInvertedListScanner *get_InvertedListScanner (
@ -113,19 +113,19 @@ struct IndexBinaryIVF : IndexBinary {
/** assign the vectors, then call search_preassign */
void search(idx_t n, const uint8_t *x, idx_t k,
int32_t *distances, idx_t *labels, ConcurrentBitsetPtr bitset = nullptr) const override;
int32_t *distances, idx_t *labels, const BitsetView& bitset = nullptr) const override;
#if 0
/** get raw vectors by ids */
void get_vector_by_id(idx_t n, const idx_t *xid, uint8_t *x, ConcurrentBitsetPtr bitset = nullptr) override;
void get_vector_by_id(idx_t n, const idx_t *xid, uint8_t *x, const BitsetView& bitset = nullptr) override;
void search_by_id (idx_t n, const idx_t *xid, idx_t k, int32_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) override;
const BitsetView& bitset = nullptr) override;
#endif
void range_search(idx_t n, const uint8_t *x, int radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reconstruct(idx_t key, uint8_t *recons) const override;
@ -216,7 +216,7 @@ struct BinaryInvertedListScanner {
const idx_t *ids,
int32_t *distances, idx_t *labels,
size_t k,
ConcurrentBitsetPtr bitset = nullptr) const = 0;
const BitsetView& bitset = nullptr) const = 0;
virtual void scan_codes_range (size_t n,
const uint8_t *codes,

View File

@ -41,7 +41,7 @@ void IndexFlat::reset() {
void IndexFlat::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
// we see the distances and labels as heaps
@ -91,7 +91,7 @@ void IndexFlat::assign(idx_t n, const float * x, idx_t * labels, float* distance
void IndexFlat::range_search (idx_t n, const float *x, float radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
switch (metric_type) {
case METRIC_INNER_PRODUCT:
@ -271,7 +271,7 @@ void IndexFlatL2BaseShift::search (
idx_t k,
float *distances,
idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT (shift.size() == ntotal);
@ -355,7 +355,7 @@ static void reorder_2_heaps (
void IndexRefineFlat::search (
idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT (is_trained);
idx_t k_base = idx_t (k * k_factor);
@ -449,7 +449,7 @@ void IndexFlat1D::search (
idx_t k,
float *distances,
idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT_MSG (perm.size() == ntotal,
"Call update_permutation before search");

View File

@ -35,7 +35,7 @@ struct IndexFlat: Index {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void assign (
idx_t n,
@ -48,7 +48,7 @@ struct IndexFlat: Index {
const float* x,
float radius,
RangeSearchResult* result,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reconstruct(idx_t key, float* recons) const override;
@ -115,7 +115,7 @@ struct IndexFlatL2BaseShift: IndexFlatL2 {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
};
@ -151,7 +151,7 @@ struct IndexRefineFlat: Index {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
~IndexRefineFlat() override;
};
@ -180,7 +180,7 @@ struct IndexFlat1D:IndexFlatL2 {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
};

View File

@ -284,7 +284,7 @@ void IndexHNSW::train(idx_t n, const float* x)
}
void IndexHNSW::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels, ConcurrentBitsetPtr bitset) const
float *distances, idx_t *labels, const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT_MSG(storage,
@ -1013,7 +1013,7 @@ int search_from_candidates_2(const HNSW & hnsw,
} // namespace
void IndexHNSW2Level::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels, ConcurrentBitsetPtr bitset) const
float *distances, idx_t *labels, const BitsetView& bitset) const
{
if (dynamic_cast<const Index2Layer*>(storage)) {
IndexHNSW::search (n, x, k, distances, labels);

View File

@ -92,7 +92,7 @@ struct IndexHNSW : Index {
/// entry point for search
void search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reconstruct(idx_t key, float* recons) const override;
@ -164,7 +164,7 @@ struct IndexHNSW2Level : IndexHNSW {
/// entry point for search
void search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
};

View File

@ -316,7 +316,7 @@ void IndexIVF::set_direct_map_type (DirectMap::Type type)
void IndexIVF::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
@ -352,7 +352,7 @@ void IndexIVF::search (idx_t n, const float *x, idx_t k,
void IndexIVF::search_without_codes (idx_t n, const float *x,
const uint8_t *arranged_codes, std::vector<size_t> prefix_sum,
bool is_sq8, idx_t k, float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset)
const BitsetView& bitset)
{
std::unique_ptr<idx_t[]> idx(new idx_t[n * nprobe]);
std::unique_ptr<float[]> coarse_dis(new float[n * nprobe]);
@ -386,12 +386,12 @@ void IndexIVF::search_without_codes (idx_t n, const float *x,
}
#if 0
void IndexIVF::get_vector_by_id (idx_t n, const idx_t *xid, float *x, ConcurrentBitsetPtr bitset) {
void IndexIVF::get_vector_by_id (idx_t n, const idx_t *xid, float *x, const BitsetView& bitset) {
make_direct_map(true);
/* only get vector by 1 id */
FAISS_ASSERT(n == 1);
if (!bitset || !bitset->test(xid[0])) {
if (!bitset || !bitset.test(xid[0])) {
reconstruct(xid[0], x + 0 * d);
} else {
memset(x, UINT8_MAX, d * sizeof(float));
@ -399,7 +399,7 @@ void IndexIVF::get_vector_by_id (idx_t n, const idx_t *xid, float *x, Concurrent
}
void IndexIVF::search_by_id (idx_t n, const idx_t *xid, idx_t k, float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) {
const BitsetView& bitset) {
make_direct_map(true);
auto x = new float[n * d];
@ -418,7 +418,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
long nprobe = params ? params->nprobe : this->nprobe;
long max_codes = params ? params->max_codes : this->max_codes;
@ -474,7 +474,7 @@ void IndexIVF::search_preassigned (idx_t n, const float *x, idx_t k,
// set porperly) and storing results in simi and idxi
auto scan_one_list = [&] (idx_t key, float coarse_dis_i,
float *simi, idx_t *idxi,
ConcurrentBitsetPtr bitset) {
const BitsetView& bitset) {
if (key < 0) {
// not enough centroids for multiprobe
@ -623,7 +623,7 @@ void IndexIVF::search_preassigned_without_codes (idx_t n, const float *x,
float *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params,
ConcurrentBitsetPtr bitset)
const BitsetView& bitset)
{
long nprobe = params ? params->nprobe : this->nprobe;
long max_codes = params ? params->max_codes : this->max_codes;
@ -678,7 +678,7 @@ void IndexIVF::search_preassigned_without_codes (idx_t n, const float *x,
// single list scan using the current scanner (with query
// set porperly) and storing results in simi and idxi
auto scan_one_list = [&] (idx_t key, float coarse_dis_i, const uint8_t *arranged_codes,
float *simi, idx_t *idxi, ConcurrentBitsetPtr bitset) {
float *simi, idx_t *idxi, const BitsetView& bitset) {
if (key < 0) {
// not enough centroids for multiprobe
@ -823,7 +823,7 @@ void IndexIVF::search_preassigned_without_codes (idx_t n, const float *x,
void IndexIVF::range_search (idx_t nx, const float *x, float radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
std::unique_ptr<idx_t[]> keys (new idx_t[nx * nprobe]);
std::unique_ptr<float []> coarse_dis (new float[nx * nprobe]);
@ -845,7 +845,7 @@ void IndexIVF::range_search_preassigned (
idx_t nx, const float *x, float radius,
const idx_t *keys, const float *coarse_dis,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
size_t nlistv = 0, ndis = 0;
@ -1260,7 +1260,7 @@ void InvertedListScanner::scan_codes_range (size_t ,
const idx_t *,
float ,
RangeQueryResult &,
ConcurrentBitsetPtr) const
const BitsetView&) const
{
FAISS_THROW_MSG ("scan_codes_range not implemented");
}

View File

@ -190,7 +190,7 @@ struct IndexIVF: Index, Level1Quantizer {
float *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params=nullptr,
ConcurrentBitsetPtr bitset = nullptr
const BitsetView& bitset = nullptr
) const;
/** Similar to search_preassigned, but does not store codes **/
@ -203,35 +203,35 @@ struct IndexIVF: Index, Level1Quantizer {
float *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params = nullptr,
ConcurrentBitsetPtr bitset = nullptr);
const BitsetView& bitset = nullptr);
/** assign the vectors, then call search_preassign */
void search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
/** Similar to search, but does not store codes **/
void search_without_codes (idx_t n, const float *x,
const uint8_t *arranged_codes, std::vector<size_t> prefix_sum,
bool is_sq8, idx_t k, float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr);
const BitsetView& bitset = nullptr);
#if 0
/** get raw vectors by ids */
void get_vector_by_id (idx_t n, const idx_t *xid, float *x, ConcurrentBitsetPtr bitset = nullptr) override;
void get_vector_by_id (idx_t n, const idx_t *xid, float *x, const BitsetView& bitset = nullptr) override;
void search_by_id (idx_t n, const idx_t *xid, idx_t k, float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) override;
const BitsetView& bitset = nullptr) override;
#endif
void range_search (idx_t n, const float* x, float radius,
RangeSearchResult* result,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void range_search_preassigned(idx_t nx, const float *x, float radius,
const idx_t *keys, const float *coarse_dis,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset = nullptr) const;
const BitsetView& bitset = nullptr) const;
/// get a scanner for this index (store_pairs means ignore labels)
virtual InvertedListScanner *get_InvertedListScanner (
@ -381,7 +381,7 @@ struct InvertedListScanner {
const idx_t *ids,
float *distances, idx_t *labels,
size_t k,
ConcurrentBitsetPtr bitset = nullptr) const = 0;
const BitsetView& bitset = nullptr) const = 0;
/** scan a set of codes, compute distances to current query and
* update results if distances are below radius
@ -392,7 +392,7 @@ struct InvertedListScanner {
const idx_t *ids,
float radius,
RangeQueryResult &result,
ConcurrentBitsetPtr bitset = nullptr) const;
const BitsetView& bitset = nullptr) const;
virtual ~InvertedListScanner () {}

View File

@ -184,12 +184,12 @@ struct IVFFlatScanner: InvertedListScanner {
const idx_t *ids,
float *simi, idx_t *idxi,
size_t k,
ConcurrentBitsetPtr bitset) const override
const BitsetView& bitset) const override
{
const float *list_vecs = (const float*)codes;
size_t nup = 0;
for (size_t j = 0; j < list_size; j++) {
if (!bitset || !bitset->test(ids[j])) {
if (!bitset || !bitset.test(ids[j])) {
const float * yj = list_vecs + d * j;
float dis = metric == METRIC_INNER_PRODUCT ?
fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
@ -208,7 +208,7 @@ struct IVFFlatScanner: InvertedListScanner {
const idx_t *ids,
float radius,
RangeQueryResult & res,
ConcurrentBitsetPtr bitset = nullptr) const override
const BitsetView& bitset = nullptr) const override
{
const float *list_vecs = (const float*)codes;
for (size_t j = 0; j < list_size; j++) {
@ -354,7 +354,7 @@ void IndexIVFFlatDedup::search_preassigned (
float *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT_MSG (
!store_pairs, "store_pairs not supported in IVFDedup");
@ -483,7 +483,7 @@ void IndexIVFFlatDedup::range_search(
const float* ,
float ,
RangeSearchResult* ,
ConcurrentBitsetPtr) const
const BitsetView&) const
{
FAISS_THROW_MSG ("not implemented");
}

View File

@ -81,7 +81,7 @@ struct IndexIVFFlatDedup: IndexIVFFlat {
float *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params=nullptr,
ConcurrentBitsetPtr bitset = nullptr
const BitsetView& bitset = nullptr
) const override;
size_t remove_ids(const IDSelector& sel) override;
@ -92,7 +92,7 @@ struct IndexIVFFlatDedup: IndexIVFFlat {
const float* x,
float radius,
RangeSearchResult* result,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
/// not implemented
void update_vectors (int nv, const idx_t *idx, const float *v) override;

View File

@ -801,10 +801,10 @@ struct KnnSearchResults {
size_t nup;
inline void add (idx_t j, float dis, ConcurrentBitsetPtr bitset = nullptr) {
inline void add (idx_t j, float dis, const BitsetView& bitset = nullptr) {
if (C::cmp (heap_sim[0], dis)) {
idx_t id = ids ? ids[j] : lo_build (key, j);
if (bitset != nullptr && bitset->test((faiss::ConcurrentBitset::id_type_t)id))
if (!bitset.empty() && bitset.test((faiss::ConcurrentBitset::id_type_t)id))
return;
heap_swap_top<C> (k, heap_sim, heap_ids, dis, id);
nup++;
@ -822,7 +822,7 @@ struct RangeSearchResults {
float radius;
RangeQueryResult & rres;
inline void add (idx_t j, float dis, faiss::ConcurrentBitsetPtr bitset = nullptr) {
inline void add (idx_t j, float dis, const faiss::BitsetView& bitset = nullptr) {
if (C::cmp (radius, dis)) {
idx_t id = ids ? ids[j] : lo_build (key, j);
rres.add (dis, id);
@ -872,7 +872,7 @@ struct IVFPQScannerT: QueryTables {
template<class SearchResultType>
void scan_list_with_table (size_t ncode, const uint8_t *codes,
SearchResultType & res,
ConcurrentBitsetPtr bitset = nullptr) const
const BitsetView& bitset = nullptr) const
{
for (size_t j = 0; j < ncode; j++) {
PQDecoder decoder(codes, pq.nbits);
@ -895,7 +895,7 @@ struct IVFPQScannerT: QueryTables {
template<class SearchResultType>
void scan_list_with_pointer (size_t ncode, const uint8_t *codes,
SearchResultType & res,
faiss::ConcurrentBitsetPtr bitset = nullptr) const
const faiss::BitsetView& bitset = nullptr) const
{
for (size_t j = 0; j < ncode; j++) {
PQDecoder decoder(codes, pq.nbits);
@ -918,7 +918,7 @@ struct IVFPQScannerT: QueryTables {
template<class SearchResultType>
void scan_on_the_fly_dist (size_t ncode, const uint8_t *codes,
SearchResultType &res,
faiss::ConcurrentBitsetPtr bitset = nullptr) const
const faiss::BitsetView& bitset = nullptr) const
{
const float *dvec;
float dis0 = 0;
@ -958,7 +958,7 @@ struct IVFPQScannerT: QueryTables {
void scan_list_polysemous_hc (
size_t ncode, const uint8_t *codes,
SearchResultType & res,
faiss::ConcurrentBitsetPtr bitset = nullptr) const
const faiss::BitsetView& bitset = nullptr) const
{
int ht = ivfpq.polysemous_ht;
size_t n_hamming_pass = 0, nup = 0;
@ -996,7 +996,7 @@ struct IVFPQScannerT: QueryTables {
void scan_list_polysemous (
size_t ncode, const uint8_t *codes,
SearchResultType &res,
faiss::ConcurrentBitsetPtr bitset = nullptr) const
const faiss::BitsetView& bitset = nullptr) const
{
switch (pq.code_size) {
#define HANDLE_CODE_SIZE(cs) \
@ -1075,7 +1075,7 @@ struct IVFPQScanner:
const idx_t *ids,
float *heap_sim, idx_t *heap_ids,
size_t k,
faiss::ConcurrentBitsetPtr bitset) const override
const faiss::BitsetView& bitset) const override
{
KnnSearchResults<C> res = {
/* key */ this->key,
@ -1106,7 +1106,7 @@ struct IVFPQScanner:
const idx_t *ids,
float radius,
RangeQueryResult & rres,
faiss::ConcurrentBitsetPtr bitset = nullptr) const override
const faiss::BitsetView& bitset = nullptr) const override
{
RangeSearchResults<C> res = {
/* key */ this->key,

View File

@ -101,7 +101,7 @@ void IndexIVFPQR::search_preassigned (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params,
ConcurrentBitsetPtr bitset
const BitsetView& bitset
) const
{
uint64_t t0;

View File

@ -56,7 +56,7 @@ struct IndexIVFPQR: IndexIVFPQ {
float *distances, idx_t *labels,
bool store_pairs,
const IVFSearchParameters *params=nullptr,
ConcurrentBitsetPtr bitset = nullptr
const BitsetView& bitset = nullptr
) const override;
IndexIVFPQR();

View File

@ -262,11 +262,11 @@ struct IVFScanner: InvertedListScanner {
const idx_t *ids,
float *simi, idx_t *idxi,
size_t k,
ConcurrentBitsetPtr bitset) const override
const BitsetView& bitset) const override
{
size_t nup = 0;
for (size_t j = 0; j < list_size; j++) {
if (!bitset || !bitset->test(ids[j])) {
if (!bitset || !bitset.test(ids[j])) {
float dis = hc.hamming (codes);
if (dis < simi [0]) {
@ -285,7 +285,7 @@ struct IVFScanner: InvertedListScanner {
const idx_t *ids,
float radius,
RangeQueryResult & res,
ConcurrentBitsetPtr bitset = nullptr) const override
const BitsetView& bitset = nullptr) const override
{
for (size_t j = 0; j < list_size; j++) {
float dis = hc.hamming (codes);

View File

@ -131,7 +131,7 @@ void IndexLSH::search (
idx_t k,
float *distances,
idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT (is_trained);
const float *xt = apply_preprocess (n, x);

View File

@ -58,7 +58,7 @@ struct IndexLSH:Index {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reset() override;

View File

@ -128,7 +128,7 @@ void IndexLattice::add(idx_t , const float* )
void IndexLattice::search(idx_t , const float* , idx_t ,
float* , idx_t* , ConcurrentBitsetPtr ) const
float* , idx_t* , const BitsetView&) const
{
FAISS_THROW_MSG("not implemented");
}

View File

@ -59,7 +59,7 @@ struct IndexLattice: Index {
void add(idx_t n, const float* x) override;
void search(idx_t n, const float* x, idx_t k,
float* distances, idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reset() override;
};

View File

@ -205,7 +205,7 @@ DistanceComputer * IndexPQ::get_distance_computer() const {
void IndexPQ::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT (is_trained);
if (search_type == ST_PQ) { // Simple PQ search
@ -947,7 +947,7 @@ void MultiIndexQuantizer::train(idx_t n, const float *x)
void MultiIndexQuantizer::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
if (n == 0) return;
// the allocation just below can be severe...
@ -1101,7 +1101,7 @@ void MultiIndexQuantizer2::train(idx_t n, const float* x)
void MultiIndexQuantizer2::search(
idx_t n, const float* x, idx_t K,
float* distances, idx_t* labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
if (n == 0) return;

View File

@ -54,7 +54,7 @@ struct IndexPQ: Index {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reset() override;
@ -159,7 +159,7 @@ struct MultiIndexQuantizer: Index {
void search(
idx_t n, const float* x, idx_t k,
float* distances, idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
/// add and reset will crash at runtime
void add(idx_t n, const float* x) override;
@ -193,7 +193,7 @@ struct MultiIndexQuantizer2: MultiIndexQuantizer {
void search(
idx_t n, const float* x, idx_t k,
float* distances, idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
};

View File

@ -181,7 +181,7 @@ void IndexPreTransform::add_with_ids (idx_t n, const float * x,
void IndexPreTransform::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT (is_trained);
const float *xt = apply_chain (n, x);
@ -191,7 +191,7 @@ void IndexPreTransform::search (idx_t n, const float *x, idx_t k,
void IndexPreTransform::range_search (idx_t n, const float* x, float radius,
RangeSearchResult* result,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT (is_trained);
const float *xt = apply_chain (n, x);

View File

@ -52,13 +52,13 @@ struct IndexPreTransform: Index {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
/* range search, no attempt is done to change the radius */
void range_search (idx_t n, const float* x, float radius,
RangeSearchResult* result,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reconstruct (idx_t key, float * recons) const override;

View File

@ -203,7 +203,7 @@ void IndexRHNSW::train(idx_t n, const float* x)
}
void IndexRHNSW::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels, ConcurrentBitsetPtr bitset) const
float *distances, idx_t *labels, const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT_MSG(storage,
@ -682,7 +682,7 @@ int search_from_candidates_2(const RHNSW & hnsw,
} // namespace
void IndexRHNSW2Level::search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels, ConcurrentBitsetPtr bitset) const
float *distances, idx_t *labels, const BitsetView& bitset) const
{
if (dynamic_cast<const Index2Layer*>(storage)) {
IndexRHNSW::search (n, x, k, distances, labels);

View File

@ -92,7 +92,7 @@ struct IndexRHNSW : Index {
/// entry point for search
void search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reconstruct(idx_t key, float* recons) const override;
@ -144,7 +144,7 @@ struct IndexRHNSW2Level : IndexRHNSW {
/// entry point for search
void search (idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
size_t cal_size();
};

View File

@ -81,7 +81,7 @@ IndexReplicasTemplate<IndexT>::search(idx_t n,
idx_t k,
distance_t* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
if (n == 0) {

View File

@ -61,7 +61,7 @@ class IndexReplicasTemplate : public ThreadedIndex<IndexT> {
idx_t k,
distance_t* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
/// reconstructs from the first index
void reconstruct(idx_t, component_t *v) const override;

View File

@ -66,7 +66,7 @@ void IndexScalarQuantizer::search(
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT (is_trained);
FAISS_THROW_IF_NOT (metric_type == METRIC_L2 ||

View File

@ -59,7 +59,7 @@ struct IndexScalarQuantizer: Index {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reset() override;

View File

@ -265,7 +265,7 @@ IndexShardsTemplate<IndexT>::search(idx_t n,
idx_t k,
distance_t *distances,
idx_t *labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
long nshard = this->count();
std::vector<distance_t> all_distances(nshard * k * n);

View File

@ -76,7 +76,7 @@ struct IndexShardsTemplate : public ThreadedIndex<IndexT> {
void search(idx_t n, const component_t* x, idx_t k,
distance_t* distances, idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void train(idx_t n, const component_t* x) override;

View File

@ -83,7 +83,7 @@ template <typename IndexT>
void IndexIDMapTemplate<IndexT>::search
(idx_t n, const typename IndexT::component_t *x, idx_t k,
typename IndexT::distance_t *distances, typename IndexT::idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
index->search (n, x, k, distances, labels, bitset);
idx_t *li = labels;
@ -96,11 +96,11 @@ void IndexIDMapTemplate<IndexT>::search
#if 0
template <typename IndexT>
void IndexIDMapTemplate<IndexT>::get_vector_by_id(idx_t n, const idx_t *xid, component_t *x,
ConcurrentBitsetPtr bitset)
const BitsetView& bitset)
{
/* only get vector by 1 id */
FAISS_ASSERT(n == 1);
if (!bitset || !bitset->test(xid[0])) {
if (!bitset || !bitset.test(xid[0])) {
index->reconstruct(xid[0], x + 0 * IndexT::d);
} else {
memset(x, UINT8_MAX, IndexT::d * sizeof(component_t));
@ -109,7 +109,7 @@ void IndexIDMapTemplate<IndexT>::get_vector_by_id(idx_t n, const idx_t *xid, com
template <typename IndexT>
void IndexIDMapTemplate<IndexT>::search_by_id (idx_t n, const idx_t *xid, idx_t k,
typename IndexT::distance_t *distances, idx_t *labels, ConcurrentBitsetPtr bitset)
typename IndexT::distance_t *distances, idx_t *labels, const BitsetView& bitset)
{
auto x = new typename IndexT::component_t[n * IndexT::d];
for (idx_t i = 0; i < n; i++) {
@ -124,7 +124,7 @@ template <typename IndexT>
void IndexIDMapTemplate<IndexT>::range_search
(typename IndexT::idx_t n, const typename IndexT::component_t *x,
typename IndexT::distance_t radius, RangeSearchResult *result,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
index->range_search(n, x, radius, result, bitset);
#pragma omp parallel for
@ -284,7 +284,7 @@ void IndexSplitVectors::add(idx_t /*n*/, const float* /*x*/) {
void IndexSplitVectors::search (
idx_t n, const float *x, idx_t k,
float *distances, idx_t *labels,
ConcurrentBitsetPtr bitset) const
const BitsetView& bitset) const
{
FAISS_THROW_IF_NOT_MSG (k == 1,
"search implemented only for k=1");

View File

@ -40,13 +40,13 @@ struct IndexIDMapTemplate : IndexT {
void search(
idx_t n, const component_t* x, idx_t k,
distance_t* distances, idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
#if 0
void get_vector_by_id(idx_t n, const idx_t *xid, component_t *x, ConcurrentBitsetPtr bitset = nullptr) override;
void get_vector_by_id(idx_t n, const idx_t *xid, component_t *x, const BitsetView& bitset = nullptr) override;
void search_by_id (idx_t n, const idx_t *xid, idx_t k, distance_t *distances, idx_t *labels,
ConcurrentBitsetPtr bitset = nullptr) override;
const BitsetView& bitset = nullptr) override;
#endif
void train(idx_t n, const component_t* x) override;
@ -58,7 +58,7 @@ struct IndexIDMapTemplate : IndexT {
void range_search (idx_t n, const component_t *x, distance_t radius,
RangeSearchResult *result,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
~IndexIDMapTemplate () override;
IndexIDMapTemplate () {own_fields=false; index=nullptr; }
@ -119,7 +119,7 @@ struct IndexSplitVectors: Index {
idx_t k,
float* distances,
idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void train(idx_t n, const float* x) override;

View File

@ -199,7 +199,7 @@ GpuIndex::search(Index::idx_t n,
Index::idx_t k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
FAISS_THROW_IF_NOT_MSG(this->is_trained, "Index not trained");
// For now, only support <= max int results
@ -276,7 +276,7 @@ GpuIndex::searchNonPaged_(int n,
int k,
float* outDistancesData,
Index::idx_t* outIndicesData,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
auto stream = resources_->getDefaultStream(device_);
// Make sure arguments are on the device we desire; use temporary
@ -296,7 +296,7 @@ GpuIndex::searchFromCpuPaged_(int n,
int k,
float* outDistancesData,
Index::idx_t* outIndicesData,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
Tensor<float, 2, true> outDistances(outDistancesData, {n, k});
Tensor<Index::idx_t, 2, true> outIndices(outIndicesData, {n, k});

View File

@ -73,7 +73,7 @@ class GpuIndex : public faiss::Index {
Index::idx_t k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
/// Overridden to force GPU indices to provide their own GPU-friendly
/// implementation
@ -112,7 +112,7 @@ class GpuIndex : public faiss::Index {
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const = 0;
const BitsetView& bitset = nullptr) const = 0;
private:
/// Handles paged adds if the add set is too large, passes to
@ -132,7 +132,7 @@ private:
int k,
float* outDistancesData,
Index::idx_t* outIndicesData,
ConcurrentBitsetPtr bitset = nullptr) const;
const BitsetView& bitset = nullptr) const;
/// Calls searchImpl_ for a single page of GPU-resident data,
/// handling paging of the data and copies from the CPU
@ -141,7 +141,7 @@ private:
int k,
float* outDistancesData,
Index::idx_t* outIndicesData,
ConcurrentBitsetPtr bitset = nullptr) const;
const BitsetView& bitset = nullptr) const;
protected:
/// Manages streams, cuBLAS handles and scratch memory for devices

View File

@ -146,7 +146,7 @@ GpuIndexBinaryFlat::search(faiss::IndexBinary::idx_t n,
faiss::IndexBinary::idx_t k,
int32_t* distances,
faiss::IndexBinary::idx_t* labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
if (n == 0) {
return;
}

View File

@ -55,7 +55,7 @@ class GpuIndexBinaryFlat : public IndexBinary {
faiss::IndexBinary::idx_t k,
int32_t* distances,
faiss::IndexBinary::idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
void reconstruct(faiss::IndexBinary::idx_t key,
uint8_t* recons) const override;

View File

@ -196,7 +196,7 @@ GpuIndexFlat::searchImpl_(int n,
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
auto stream = resources_->getDefaultStream(device_);
// Input and output data are already resident on the GPU
@ -214,8 +214,8 @@ GpuIndexFlat::searchImpl_(int n,
data_->query(queries, bitsetDevice, k, metric_type, metric_arg, outDistances, outIntLabels, true);
} else {
auto bitsetDevice = toDevice<uint8_t, 1>(resources_, device_,
const_cast<uint8_t*>(bitset->data()), stream,
{(int) bitset->size()});
const_cast<uint8_t*>(bitset.data()), stream,
{(int) bitset.u8size()});
data_->query(queries, bitsetDevice, k, metric_type, metric_arg, outDistances, outIntLabels, true);
}

View File

@ -122,7 +122,7 @@ class GpuIndexFlat : public GpuIndex {
int k,
float* distances,
faiss::Index::idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
protected:
/// Our config object

View File

@ -300,7 +300,7 @@
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
// Device is already set in GpuIndex::search
FAISS_ASSERT(index_);
FAISS_ASSERT(n > 0);
@ -319,8 +319,8 @@
index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels);
} else {
auto bitsetDevice = toDevice<uint8_t, 1>(resources_, device_,
const_cast<uint8_t*>(bitset->data()), stream,
{(int) bitset->size()});
const_cast<uint8_t*>(bitset.data()), stream,
{(int) bitset.u8size()});
index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels);
}
}

View File

@ -76,7 +76,7 @@ class GpuIndexIVFFlat : public GpuIndexIVF {
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
private:
GpuIndexIVFFlatConfig ivfFlatConfig_;

View File

@ -359,7 +359,7 @@ GpuIndexIVFPQ::searchImpl_(int n,
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
// Device is already set in GpuIndex::search
FAISS_ASSERT(index_);
FAISS_ASSERT(n > 0);
@ -378,8 +378,8 @@ GpuIndexIVFPQ::searchImpl_(int n,
index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels);
} else {
auto bitsetDevice = toDevice<uint8_t, 1>(resources_, device_,
const_cast<uint8_t*>(bitset->data()), stream,
{(int) bitset->size()});
const_cast<uint8_t*>(bitset.data()), stream,
{(int) bitset.u8size()});
index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels);
}
}

View File

@ -118,7 +118,7 @@ class GpuIndexIVFPQ : public GpuIndexIVF {
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
private:
void verifySettings_() const;

View File

@ -330,7 +330,7 @@ GpuIndexIVFSQHybrid::searchImpl_(int n,
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
// Device is already set in GpuIndex::search
FAISS_ASSERT(index_);
FAISS_ASSERT(n > 0);
@ -349,7 +349,7 @@ GpuIndexIVFSQHybrid::searchImpl_(int n,
index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels);
} else {
auto bitsetDevice = toDevice<uint8_t, 1>(resources_, device_,
const_cast<uint8_t*>(bitset->data()), stream,{(int) bitset->size()});
const_cast<uint8_t*>(bitset.data()), stream,{(int) bitset.u8size()});
index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels);
}
}

View File

@ -80,7 +80,7 @@ class GpuIndexIVFSQHybrid : public GpuIndexIVF {
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
/// Called from train to handle SQ residual training
void trainResiduals_(Index::idx_t n, const float* x);

View File

@ -342,7 +342,7 @@ GpuIndexIVFScalarQuantizer::searchImpl_(int n,
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset) const {
const BitsetView& bitset) const {
// Device is already set in GpuIndex::search
FAISS_ASSERT(index_);
FAISS_ASSERT(n > 0);
@ -361,8 +361,8 @@ GpuIndexIVFScalarQuantizer::searchImpl_(int n,
index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels);
} else {
auto bitsetDevice = toDevice<uint8_t, 1>(resources_, device_,
const_cast<uint8_t*>(bitset->data()), stream,
{(int) bitset->size()});
const_cast<uint8_t*>(bitset.data()), stream,
{(int) bitset.u8size()});
index_->query(queries, bitsetDevice, nprobe, k, outDistances, outLabels);
}
}

View File

@ -80,7 +80,7 @@ class GpuIndexIVFScalarQuantizer : public GpuIndexIVF {
int k,
float* distances,
Index::idx_t* labels,
ConcurrentBitsetPtr bitset = nullptr) const override;
const BitsetView& bitset = nullptr) const override;
/// Called from train to handle SQ residual training
void trainResiduals_(Index::idx_t n, const float* x);

Some files were not shown because too many files have changed in this diff Show More