Skip to content

Commit

Permalink
pyknowhere support Int8Vector (#1067)
Browse files Browse the repository at this point in the history 
Signed-off-by: CaiYudong <[email protected]>
  • Loading branch information
@cydrain
cydrain authored Feb 12, 2025
1 parent a25575e commit 368e43b
Show file tree
Hide file tree
Showing 5 changed files with 102 additions and 16 deletions.
4 changes: 3 additions & 1 deletion benchmark/hdf5/benchmark_hdf5.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,9 @@ class Benchmark_hdf5 : public Benchmark_base {
set_metric_type(const std::string& str) {
if (str == METRIC_L2_STR || str == "l2") {
metric_type_ = "L2";
} else if (str == METRIC_COS_STR) {
} else if (str == "ip") {
metric_type_ = "IP";
} else if (str == METRIC_COS_STR || str == "cosine") {
metric_type_ = "COSINE";
} else if (str == METRIC_HAM_STR) {
metric_type_ = "HAMMING";
Expand Down
56 changes: 52 additions & 4 deletions benchmark/hdf5/gen_fbin_file.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
knowhere::DataSetPtr
GenDataSet(int rows, int dim) {
std::mt19937 rng(42);
std::uniform_real_distribution<> distrib(-1.0, 1.0);
std::uniform_int_distribution<> distrib(-128, 127);
float* ts = new float[rows * dim];
for (int i = 0; i < rows * dim; ++i) {
ts[i] = (float)distrib(rng);
Expand Down Expand Up @@ -137,13 +137,61 @@ class Create_FBIN : public Benchmark_hdf5, public ::testing::Test {
fbin_result_write(filename, nq, topk, gt_ids_int.data(), result.value()->GetDistance());
}
}

void
create_range_fbin_files(const int64_t nb, const int64_t nq, const int64_t dim,
const std::vector<std::pair<knowhere::MetricType, float>>& params) {
knowhere::DataSetPtr xb_ds, xq_ds;
xb_ds = GenDataSet(nb, dim);
xq_ds = GenDataSet(nq, dim);

std::string prefix = "rand-" + std::to_string(dim) + "-";
std::string postfix = ".fbin";
std::string filename;

filename = prefix + "base" + postfix;
fbin_write(filename, nb, dim, xb_ds->GetTensor());

filename = prefix + "query" + postfix;
fbin_write(filename, nq, dim, xq_ds->GetTensor());

for (auto [metric_type, radius] : params) {
std::string metric_str = metric_type;
transform(metric_str.begin(), metric_str.end(), metric_str.begin(), ::tolower);

knowhere::Json json;
json[knowhere::meta::DIM] = dim;
json[knowhere::meta::METRIC_TYPE] = metric_type;
json[knowhere::meta::RADIUS] = radius;

auto result = knowhere::BruteForce::RangeSearch<knowhere::fp32>(xb_ds, xq_ds, json, nullptr);
assert(result.has_value());

// convert golden_lims to int32
std::vector<uint32_t> gt_lims_int(nq + 1);
for (int32_t i = 0; i <= nq; i++) {
gt_lims_int[i] = result.value()->GetLims()[i];
}

// convert golden_ids to int32
auto elem_cnt = result.value()->GetLims()[nq];
std::vector<uint32_t> gt_ids_int(elem_cnt);
for (int32_t i = 0; i < elem_cnt; i++) {
gt_ids_int[i] = result.value()->GetIds()[i];
}

filename = prefix + metric_str + "-gt" + postfix;
fbin_range_result_write(filename, nq, radius, gt_lims_int.data(), gt_ids_int.data(),
result.value()->GetDistance());
}
}
};

TEST_F(Create_FBIN, CREATE_FLOAT) {
int64_t nb = 10000;
int64_t nq = 100;
int64_t dim = 128;
int64_t topk = 100;
int64_t topk = 5000;

create_fbin_files(nb, nq, dim, topk, {knowhere::metric::L2, knowhere::metric::IP, knowhere::metric::COSINE});
}
Expand Down Expand Up @@ -191,7 +239,7 @@ TEST_F(Create_FBIN, HDF5_RANGE_TO_FBIN) {
}

TEST_F(Create_FBIN, HDF5_BIN_TO_FBIN) {
set_ann_test_name("rand-1024-hamming");
set_ann_test_name("rand-4096-hamming");
parse_ann_test_name();
load_hdf5_data<knowhere::bin1>();

Expand All @@ -212,7 +260,7 @@ TEST_F(Create_FBIN, HDF5_BIN_TO_FBIN) {
}

TEST_F(Create_FBIN, HDF5_BIN_RANGE_TO_FBIN) {
set_ann_test_name("rand-1024-hamming-range");
set_ann_test_name("rand-4096-hamming-range");
parse_ann_test_name_with_range();
load_hdf5_data_range<knowhere::bin1>();

Expand Down
18 changes: 9 additions & 9 deletions benchmark/hdf5/gen_hdf5_file.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
knowhere::DataSetPtr
GenDataSet(int rows, int dim) {
std::mt19937 rng(42);
std::uniform_real_distribution<> distrib(-1.0, 1.0);
std::uniform_int_distribution<> distrib(-128, 127);
float* ts = new float[rows * dim];
for (int i = 0; i < rows * dim; ++i) {
ts[i] = (float)distrib(rng);
Expand Down Expand Up @@ -144,7 +144,7 @@ TEST_F(Create_HDF5, CREATE_FLOAT) {
int64_t nb = 10000;
int64_t nq = 100;
int64_t dim = 128;
int64_t topk = 100;
int64_t topk = 5000;

create_hdf5_file<knowhere::fp32>(knowhere::metric::L2, nb, nq, dim, topk);
create_hdf5_file<knowhere::fp32>(knowhere::metric::IP, nb, nq, dim, topk);
Expand All @@ -156,16 +156,16 @@ TEST_F(Create_HDF5, CREATE_FLOAT_RANGE) {
int64_t nq = 100;
int64_t dim = 128;

create_range_hdf5_file<knowhere::fp32>(knowhere::metric::L2, nb, nq, dim, 65.0);
create_range_hdf5_file<knowhere::fp32>(knowhere::metric::IP, nb, nq, dim, 8.7);
create_range_hdf5_file<knowhere::fp32>(knowhere::metric::L2, nb, nq, dim, 1070000.0);
create_range_hdf5_file<knowhere::fp32>(knowhere::metric::IP, nb, nq, dim, 142000.0);
create_range_hdf5_file<knowhere::fp32>(knowhere::metric::COSINE, nb, nq, dim, 0.2);
}

TEST_F(Create_HDF5, CREATE_BINARY) {
int64_t nb = 10000;
int64_t nq = 100;
int64_t dim = 1024;
int64_t topk = 100;
int64_t dim = 4096;
int64_t topk = 5000;

create_hdf5_file<knowhere::bin1>(knowhere::metric::HAMMING, nb, nq, dim, topk);
create_hdf5_file<knowhere::bin1>(knowhere::metric::JACCARD, nb, nq, dim, topk);
Expand All @@ -174,8 +174,8 @@ TEST_F(Create_HDF5, CREATE_BINARY) {
TEST_F(Create_HDF5, CREATE_BINARY_RANGE) {
int64_t nb = 10000;
int64_t nq = 100;
int64_t dim = 1024;
int64_t dim = 4096;

create_range_hdf5_file<knowhere::bin1>(knowhere::metric::HAMMING, nb, nq, dim, 476);
create_range_hdf5_file<knowhere::bin1>(knowhere::metric::JACCARD, nb, nq, dim, 0.63);
create_range_hdf5_file<knowhere::bin1>(knowhere::metric::HAMMING, nb, nq, dim, 1974);
create_range_hdf5_file<knowhere::bin1>(knowhere::metric::JACCARD, nb, nq, dim, 0.647);
}
11 changes: 11 additions & 0 deletions python/knowhere/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ def CreateIndex(name, version, type=np.float32):
return swigknowhere.IndexWrapBF16(name, version)
if type == np.uint8:
return swigknowhere.IndexWrapBin(name, version)
if type == np.int8:
return swigknowhere.IndexWrapInt8(name, version)

def BruteForceSearch(type=np.float32, *args):
if type == np.float32:
Expand Down Expand Up @@ -74,6 +76,8 @@ def ArrayToDataSet(arr):
return swigknowhere.Array2DataSetIds(arr)
if arr.ndim == 2:
if arr.dtype == np.uint8:
return swigknowhere.Array2DataSetU(arr)
if arr.dtype == np.int8:
return swigknowhere.Array2DataSetI(arr)
if arr.dtype == np.float32:
return swigknowhere.Array2DataSetF(arr)
Expand Down Expand Up @@ -176,6 +180,13 @@ def GetBinaryVectorDataSetToArray(ans):
swigknowhere.BinaryDataSetTensor2Array(ans, data)
return data

def GetInt8VectorDataSetToArray(ans):
dim = swigknowhere.DataSet_Dim(ans)
rows = swigknowhere.DataSet_Rows(ans)
data = np.zeros([rows, dim]).astype(np.int8)
swigknowhere.Int8DataSetTensor2Array(ans, data)
return data

def SetSimdType(type):
swigknowhere.SetSimdType(type)

Expand Down
29 changes: 27 additions & 2 deletions python/knowhere/knowhere.i
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,9 @@ import_array();
%apply (int* IN_ARRAY2, int DIM1, int DIM2) {(int* xb, int nb, int dim)}
%apply (uint8_t *IN_ARRAY1, int DIM1) {(uint8_t *block, int size)}
%apply (uint8_t* IN_ARRAY2, int DIM1, int DIM2) {(uint8_t* xb, int nb, int dim)}
%apply (uint8_t* INPLACE_ARRAY2, int DIM1, int DIM2) {(uint8_t *data,int rows,int dim)}
%apply (uint8_t* INPLACE_ARRAY2, int DIM1, int DIM2) {(uint8_t *data, int rows, int dim)}
%apply (int8_t* IN_ARRAY2, int DIM1, int DIM2) {(int8_t* xb, int nb, int dim)}
%apply (int8_t* INPLACE_ARRAY2, int DIM1, int DIM2) {(int8_t *data, int rows, int dim)}
%apply (int *IN_ARRAY1, int DIM1) {(int *lims, int len)}
%apply (int *IN_ARRAY1, int DIM1) {(int *ids, int len)}
%apply (float *IN_ARRAY1, int DIM1) {(float *dis, int len)}
Expand Down Expand Up @@ -336,6 +338,7 @@ Array2DataSetFP16(float* xb, int nb, int dim) {
ds->SetTensor(fp16_data);
return ds;
};

#pragma GCC push_options
#pragma GCC optimize("O0")
knowhere::DataSetPtr
Expand Down Expand Up @@ -394,7 +397,7 @@ Array2SparseDataSet(float* data, int nb1, int* ids, int nb2, int64_t* indptr, in
}

knowhere::DataSetPtr
Array2DataSetI(uint8_t* xb, int nb, int dim) {
Array2DataSetU(uint8_t* xb, int nb, int dim) {
auto ds = std::make_shared<DataSet>();
ds->SetIsOwner(false);
ds->SetRows(nb);
Expand All @@ -403,6 +406,16 @@ Array2DataSetI(uint8_t* xb, int nb, int dim) {
return ds;
};

knowhere::DataSetPtr
Array2DataSetI(int8_t* xb, int nb, int dim) {
auto ds = std::make_shared<DataSet>();
ds->SetIsOwner(false);
ds->SetRows(nb);
ds->SetDim(dim);
ds->SetTensor(xb);
return ds;
};

knowhere::DataSetPtr
Array2DataSetIds(int* ids, int len){
auto ds = std::make_shared<DataSet>();
Expand Down Expand Up @@ -519,6 +532,17 @@ BinaryDataSetTensor2Array(knowhere::DataSetPtr result, uint8_t* data, int rows,
}
}

void
Int8DataSetTensor2Array(knowhere::DataSetPtr result, int8_t* data, int rows, int dim) {
GILReleaser rel;
auto data_ = result->GetTensor();
for (int i = 0; i < rows; i++) {
for (int j = 0; j < dim; ++j) {
*(data + i * dim + j) = *((knowhere::int8*)(data_) + i * dim + j);
}
}
}

void
DumpRangeResultIds(knowhere::DataSetPtr result, int* ids, int len) {
GILReleaser rel;
Expand Down Expand Up @@ -683,6 +707,7 @@ SetSearchThreadPool(uint32_t num_threads) {
%template(IndexWrapFP16) IndexWrap<knowhere::fp16>;
%template(IndexWrapBF16) IndexWrap<knowhere::bf16>;
%template(IndexWrapBin) IndexWrap<knowhere::bin1>;
%template(IndexWrapInt8) IndexWrap<knowhere::int8>;

%template(BruteForceSearchFloat) BruteForceSearch<float>;
%template(BruteForceSearchFP16) BruteForceSearch<knowhere::fp16>;
Expand Down

0 comments on commit 368e43b

Please sign in to comment.