Skip to content

Commit

Permalink
Add SQfp16 support to ARM Architecture
Browse files Browse the repository at this point in the history 
Signed-off-by: Naveen Tatikonda <[email protected]>
  • Loading branch information
EC2 Default User committed Dec 4, 2023
1 parent bb66831 commit f2d2944
Show file tree
Hide file tree
Showing 5 changed files with 157 additions and 73 deletions.
4 changes: 2 additions & 2 deletions benchmarks/perf-tool/release-configs/faiss-hnsw-sqfp16/index.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
"index": {
"knn": true,
"number_of_shards": 24,
"number_of_replicas": 1
"number_of_replicas": 0
}
},
"mappings": {
Expand All @@ -24,4 +24,4 @@
}
}
}
}
}
11 changes: 6 additions & 5 deletions benchmarks/perf-tool/release-configs/faiss-hnsw-sqfp16/test.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,20 @@ endpoint: "localhost"
test_name: "Faiss HNSW SQfp16 Test"
test_id: "Faiss HNSW SQfp16 Test"
num_runs: 10
show_runs: false
timeout: 600
show_runs: true
steps:
- name: delete_index
index_name: target_index
- name: create_index
index_name: target_index
index_spec: /home/ec2-user/[PATH]/index.json
index_spec: /home/ec2-user/k-NN/benchmarks/perf-tool/release-configs/faiss-hnsw-sqfp16/index.json
- name: ingest
index_name: target_index
field_name: target_field
bulk_size: 500
dataset_format: hdf5
dataset_path: [DATASET_PATH]/sift-128-euclidean.hdf5
dataset_path: /home/ec2-user/data/sift-128-euclidean.hdf5
- name: refresh_index
index_name: target_index
- name: force_merge
Expand All @@ -29,6 +30,6 @@ steps:
index_name: target_index
field_name: target_field
dataset_format: hdf5
dataset_path: [DATASET_PATH]/sift-128-euclidean.hdf5
dataset_path: /home/ec2-user/data/sift-128-euclidean.hdf5
neighbors_format: hdf5
neighbors_path: [DATASET_PATH]/sift-128-euclidean.hdf5
neighbors_path: /home/ec2-user/data/sift-128-euclidean.hdf5
2 changes: 1 addition & 1 deletion jni/src/faiss_wrapper.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ void knn_jni::faiss_wrapper::CreateIndex(knn_jni::JNIUtilInterface * jniUtil, JN
jniUtil->DeleteLocalRef(env, parametersJ);

// TODO: Add a condition to set it true for sq encoder based on index description
// indexWriter->is_trained = true;
indexWriter->is_trained = true;

// Check that the index does not need to be trained
if(!indexWriter->is_trained) {
Expand Down
84 changes: 84 additions & 0 deletions jni/tests/faiss_wrapper_test.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -247,6 +247,90 @@ TEST(FaissInitLibraryTest, BasicAssertions) {
knn_jni::faiss_wrapper::InitLibrary();
}

/*
#include "faiss/impl/ScalarQuantizer.h"
TEST(FaissSQComputeDistanceTest, BasicAssertions) {
faiss::idx_t numIds = 2;
std::vector<std::vector<float>> vectors;
int dim = 8;
for (int64_t i = 0; i < numIds; ++i) {
// ids.push_back(i);
std::cout << "\nVector " << i << ": " << std::flush;
std::vector<float> vect;
vect.reserve(dim);
for (int j = 0; j < dim; ++j) {
float f = test_util::RandomFloat(-500.0, 500.0);
std::cout << f << ", " << std::flush;
vect.push_back(f);
}
vectors.push_back(vect);
}
uint8_t codes;
faiss::ScalarQuantizer * sq = new faiss::ScalarQuantizer(dim, faiss::ScalarQuantizer::QT_fp16);
sq->compute_codes(&vectors[0][0], &codes, numIds);
}*/

#include "faiss/impl/ScalarQuantizer.h"
#include <cstring>
#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
TEST(FaissSQComputeDistanceTest, BasicAssertions) {
faiss::idx_t numIds = 20;
std::vector<std::vector<float>> vectors;
int dim = 8;
std::ifstream MyReadFile("/home/ec2-user/k-NN/test_vectors.txt");

while(!MyReadFile.eof()){
string text, T;
getline(MyReadFile, text);
stringstream X(text);
std::vector<float> vect;
vect.reserve(dim);
int j = 0;
while (std::getline(X, T, ',')) {
if (j < dim) {
float f = std::stof(T);
vect.push_back(f);
j++;
}
}
vectors.push_back(vect);
}

std::ifstream MyReadFile1("/home/ec2-user/k-NN/query_vectors.txt");
std::vector<std::vector<float>> query_vectors;
while(!MyReadFile1.eof()) {
string text, T;
getline(MyReadFile1, text);
stringstream X(text);
std::vector<float> vect;
vect.reserve(dim);
int j = 0;
while (std::getline(X, T, ',')) {
if (j < dim) {
float f = std::stof(T);
vect.push_back(f);
j++;
}
}
query_vectors.push_back(vect);
}
uint8_t codes;
faiss::ScalarQuantizer * sq = new faiss::ScalarQuantizer(dim, faiss::ScalarQuantizer::QT_fp16);
sq->compute_codes(&vectors[0][0], &codes, numIds);
faiss::ScalarQuantizer::SQDistanceComputer * sqdc = sq -> get_distance_computer(faiss::METRIC_L2);
for(int i=0; i<1; i++) {
sqdc->set_query(&query_vectors[i][0]);
sqdc->query_to_code(&(codes));
}
}


TEST(FaissTrainIndexTest, BasicAssertions) {
// Define the index configuration
int dim = 2;
Expand Down
129 changes: 64 additions & 65 deletions src/main/java/org/opensearch/knn/index/util/Faiss.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,6 @@
import static org.opensearch.knn.common.KNNConstants.FAISS_HNSW_DESCRIPTION;
import static org.opensearch.knn.common.KNNConstants.FAISS_IVF_DESCRIPTION;
import static org.opensearch.knn.common.KNNConstants.FAISS_PQ_DESCRIPTION;
import static org.opensearch.knn.common.KNNConstants.FAISS_SQFP16_DESCRIPTION;
import static org.opensearch.knn.common.KNNConstants.METHOD_ENCODER_PARAMETER;
import static org.opensearch.knn.common.KNNConstants.METHOD_HNSW;
import static org.opensearch.knn.common.KNNConstants.METHOD_IVF;
Expand Down Expand Up @@ -78,71 +77,71 @@ class Faiss extends NativeLibrary {
).build())
)
.build(),
KNNConstants.ENCODER_SQFP16,
KNNConstants.ENCODER_SQFP16,
MethodComponent.Builder.builder(KNNConstants.ENCODER_SQFP16)
.setRequiresTraining(false)
.setMapGenerator(
((methodComponent, methodComponentContext) -> MethodAsMapBuilder.builder(
KNNConstants.FAISS_SQFP16_DESCRIPTION,
methodComponent,
methodComponentContext
).build())
)
.build(),
// KNNConstants.ENCODER_SQFP16,
// MethodComponent.Builder.builder(KNNConstants.ENCODER_SQFP16)
// .addParameter(
// ENCODER_PARAMETER_PQ_M,
// new Parameter.IntegerParameter(
// ENCODER_PARAMETER_PQ_M,
// ENCODER_PARAMETER_PQ_CODE_COUNT_DEFAULT,
// v -> v > 0 && v < ENCODER_PARAMETER_PQ_CODE_COUNT_LIMIT
// )
// )
// .addParameter(
// ENCODER_PARAMETER_PQ_CODE_SIZE,
// new Parameter.IntegerParameter(
// ENCODER_PARAMETER_PQ_CODE_SIZE,
// ENCODER_PARAMETER_PQ_CODE_SIZE_DEFAULT,
// v -> v > 0 && v < ENCODER_PARAMETER_PQ_CODE_SIZE_LIMIT
// )
// )
// .setRequiresTraining(true)
// .setMapGenerator(
// ((methodComponent, methodComponentContext) -> MethodAsMapBuilder.builder(
// FAISS_SQFP16_DESCRIPTION,
// methodComponent,
// methodComponentContext
// )
// //.addParameter(ENCODER_PARAMETER_PQ_M, "", "").addParameter(ENCODER_PARAMETER_PQ_CODE_SIZE, "", "")
// .build())
// )
// .setOverheadInKBEstimator((methodComponent, methodComponentContext, dimension) -> {
// // Size estimate formula: (4 * d * 2^code_size) / 1024 + 1
//
// // Get value of code size passed in by user
// Object codeSizeObject = methodComponentContext.getParameters().get(ENCODER_PARAMETER_PQ_CODE_SIZE);
//
// // If not specified, get default value of code size
// if (codeSizeObject == null) {
// Parameter<?> codeSizeParameter = methodComponent.getParameters().get(ENCODER_PARAMETER_PQ_CODE_SIZE);
// if (codeSizeParameter == null) {
// throw new IllegalStateException(
// String.format("%s is not a valid parameter. This is a bug.", ENCODER_PARAMETER_PQ_CODE_SIZE)
// );
// }
//
// codeSizeObject = codeSizeParameter.getDefaultValue();
// }
//
// if (!(codeSizeObject instanceof Integer)) {
// throw new IllegalStateException(String.format("%s must be an integer.", ENCODER_PARAMETER_PQ_CODE_SIZE));
// }
//
// int codeSize = (Integer) codeSizeObject;
// return ((4L * (1L << codeSize) * dimension) / BYTES_PER_KILOBYTES) + 1;
// })
// .build(),
.setRequiresTraining(false)
.setMapGenerator(
((methodComponent, methodComponentContext) -> MethodAsMapBuilder.builder(
KNNConstants.FAISS_SQFP16_DESCRIPTION,
methodComponent,
methodComponentContext
).build())
)
.build(),
// KNNConstants.ENCODER_SQFP16,
// MethodComponent.Builder.builder(KNNConstants.ENCODER_SQFP16)
// .addParameter(
// ENCODER_PARAMETER_PQ_M,
// new Parameter.IntegerParameter(
// ENCODER_PARAMETER_PQ_M,
// ENCODER_PARAMETER_PQ_CODE_COUNT_DEFAULT,
// v -> v > 0 && v < ENCODER_PARAMETER_PQ_CODE_COUNT_LIMIT
// )
// )
// .addParameter(
// ENCODER_PARAMETER_PQ_CODE_SIZE,
// new Parameter.IntegerParameter(
// ENCODER_PARAMETER_PQ_CODE_SIZE,
// ENCODER_PARAMETER_PQ_CODE_SIZE_DEFAULT,
// v -> v > 0 && v < ENCODER_PARAMETER_PQ_CODE_SIZE_LIMIT
// )
// )
// .setRequiresTraining(true)
// .setMapGenerator(
// ((methodComponent, methodComponentContext) -> MethodAsMapBuilder.builder(
// FAISS_SQFP16_DESCRIPTION,
// methodComponent,
// methodComponentContext
// )
// //.addParameter(ENCODER_PARAMETER_PQ_M, "", "").addParameter(ENCODER_PARAMETER_PQ_CODE_SIZE, "", "")
// .build())
// )
// .setOverheadInKBEstimator((methodComponent, methodComponentContext, dimension) -> {
// // Size estimate formula: (4 * d * 2^code_size) / 1024 + 1
//
// // Get value of code size passed in by user
// Object codeSizeObject = methodComponentContext.getParameters().get(ENCODER_PARAMETER_PQ_CODE_SIZE);
//
// // If not specified, get default value of code size
// if (codeSizeObject == null) {
// Parameter<?> codeSizeParameter = methodComponent.getParameters().get(ENCODER_PARAMETER_PQ_CODE_SIZE);
// if (codeSizeParameter == null) {
// throw new IllegalStateException(
// String.format("%s is not a valid parameter. This is a bug.", ENCODER_PARAMETER_PQ_CODE_SIZE)
// );
// }
//
// codeSizeObject = codeSizeParameter.getDefaultValue();
// }
//
// if (!(codeSizeObject instanceof Integer)) {
// throw new IllegalStateException(String.format("%s must be an integer.", ENCODER_PARAMETER_PQ_CODE_SIZE));
// }
//
// int codeSize = (Integer) codeSizeObject;
// return ((4L * (1L << codeSize) * dimension) / BYTES_PER_KILOBYTES) + 1;
// })
// .build(),
KNNConstants.ENCODER_PQ,
MethodComponent.Builder.builder(KNNConstants.ENCODER_PQ)
.addParameter(
Expand Down

0 comments on commit f2d2944

Please sign in to comment.