-
Notifications
You must be signed in to change notification settings - Fork 126
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Introduce a loading layer in NMSLIB. (#2185)
* Introduce a loading layer in NMSLIB. Signed-off-by: Dooyong Kim <[email protected]> * Added NMSLIB istream implementation. Signed-off-by: Dooyong Kim <[email protected]> * Fix integer overflow issue when passing read size for loading NMSLIB vector index. Signed-off-by: Dooyong Kim <[email protected]> * Added unit test for NMSLIB loading layer. Signed-off-by: Dooyong Kim <[email protected]> * Made a patch in NMSLIB to avoid frequently calling JNI for better loading index performance. Signed-off-by: Dooyong Kim <[email protected]> * Compliance constexpr function in C++11 having nullstatement. Signed-off-by: Dooyong Kim <[email protected]> --------- Signed-off-by: Dooyong Kim <[email protected]> Co-authored-by: Dooyong Kim <[email protected]> (cherry picked from commit 7cf45c8)
- Loading branch information
1 parent
cbc7dac
commit 19ad84e
Showing
24 changed files
with
1,103 additions
and
562 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,125 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| * | ||
| * Modifications Copyright OpenSearch Contributors. See | ||
| * GitHub history for details. | ||
| */ | ||
|
|
||
| #ifndef OPENSEARCH_KNN_JNI_STREAM_SUPPORT_H | ||
| #define OPENSEARCH_KNN_JNI_STREAM_SUPPORT_H | ||
|
|
||
| #include "jni_util.h" | ||
|
|
||
| #include <jni.h> | ||
| #include <stdexcept> | ||
| #include <iostream> | ||
| #include <cstring> | ||
|
|
||
| namespace knn_jni { | ||
| namespace stream { | ||
|
|
||
|
|
||
|
|
||
| /** | ||
| * This class contains Java IndexInputWithBuffer reference and calls its API to copy required bytes into a read buffer. | ||
| */ | ||
| class NativeEngineIndexInputMediator { | ||
| public: | ||
| // Expect IndexInputWithBuffer is given as `_indexInput`. | ||
| NativeEngineIndexInputMediator(JNIUtilInterface *_jni_interface, | ||
| JNIEnv *_env, | ||
| jobject _indexInput) | ||
| : jni_interface(_jni_interface), | ||
| env(_env), | ||
| indexInput(_indexInput), | ||
| bufferArray((jbyteArray) (_jni_interface->GetObjectField(_env, | ||
| _indexInput, | ||
| getBufferFieldId(_jni_interface, _env)))), | ||
| copyBytesMethod(getCopyBytesMethod(_jni_interface, _env)), | ||
| remainingBytesMethod(getRemainingBytesMethod(_jni_interface, _env)) { | ||
| } | ||
|
|
||
| void copyBytes(int64_t nbytes, uint8_t *destination) { | ||
| auto jclazz = getIndexInputWithBufferClass(jni_interface, env); | ||
|
|
||
| while (nbytes > 0) { | ||
| // Call `copyBytes` to read bytes as many as possible. | ||
| jvalue args; | ||
| args.j = nbytes; | ||
| const auto readBytes = | ||
| jni_interface->CallNonvirtualIntMethodA(env, indexInput, jclazz, copyBytesMethod, &args); | ||
|
|
||
| // === Critical Section Start === | ||
|
|
||
| // Get primitive array pointer, no copy is happening in OpenJDK. | ||
| auto primitiveArray = | ||
| (jbyte *) jni_interface->GetPrimitiveArrayCritical(env, bufferArray, nullptr); | ||
|
|
||
| // Copy Java bytes to C++ destination address. | ||
| std::memcpy(destination, primitiveArray, readBytes); | ||
|
|
||
| // Release the acquired primitive array pointer. | ||
| // JNI_ABORT tells JVM to directly free memory without copying back to Java byte[]. | ||
| // Since we're merely copying data, we don't need to copying back. | ||
| jni_interface->ReleasePrimitiveArrayCritical(env, bufferArray, primitiveArray, JNI_ABORT); | ||
|
|
||
| // === Critical Section End === | ||
|
|
||
| destination += readBytes; | ||
| nbytes -= readBytes; | ||
| } // End while | ||
| } | ||
|
|
||
| int64_t remainingBytes() { | ||
| return jni_interface->CallNonvirtualLongMethodA(env, | ||
| indexInput, | ||
| getIndexInputWithBufferClass(jni_interface, env), | ||
| remainingBytesMethod, | ||
| nullptr); | ||
| } | ||
|
|
||
| private: | ||
| static jclass getIndexInputWithBufferClass(JNIUtilInterface *jni_interface, JNIEnv *env) { | ||
| static jclass INDEX_INPUT_WITH_BUFFER_CLASS = | ||
| jni_interface->FindClassFromJNIEnv(env, "org/opensearch/knn/index/store/IndexInputWithBuffer"); | ||
| return INDEX_INPUT_WITH_BUFFER_CLASS; | ||
| } | ||
|
|
||
| static jmethodID getCopyBytesMethod(JNIUtilInterface *jni_interface, JNIEnv *env) { | ||
| static jmethodID COPY_METHOD_ID = | ||
| jni_interface->GetMethodID(env, getIndexInputWithBufferClass(jni_interface, env), "copyBytes", "(J)I"); | ||
| return COPY_METHOD_ID; | ||
| } | ||
|
|
||
| static jmethodID getRemainingBytesMethod(JNIUtilInterface *jni_interface, JNIEnv *env) { | ||
| static jmethodID COPY_METHOD_ID = | ||
| jni_interface->GetMethodID(env, getIndexInputWithBufferClass(jni_interface, env), "remainingBytes", "()J"); | ||
| return COPY_METHOD_ID; | ||
| } | ||
|
|
||
| static jfieldID getBufferFieldId(JNIUtilInterface *jni_interface, JNIEnv *env) { | ||
| static jfieldID BUFFER_FIELD_ID = | ||
| jni_interface->GetFieldID(env, getIndexInputWithBufferClass(jni_interface, env), "buffer", "[B"); | ||
| return BUFFER_FIELD_ID; | ||
| } | ||
|
|
||
| JNIUtilInterface *jni_interface; | ||
| JNIEnv *env; | ||
|
|
||
| // `IndexInputWithBuffer` instance having `IndexInput` instance obtained from `Directory` for reading. | ||
| jobject indexInput; | ||
| jbyteArray bufferArray; | ||
| jmethodID copyBytesMethod; | ||
| jmethodID remainingBytesMethod; | ||
| }; // class NativeEngineIndexInputMediator | ||
|
|
||
|
|
||
|
|
||
| } | ||
| } | ||
|
|
||
| #endif //OPENSEARCH_KNN_JNI_STREAM_SUPPORT_H |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,51 @@ | ||
| /* | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| * | ||
| * The OpenSearch Contributors require contributions made to | ||
| * this file be licensed under the Apache-2.0 license or a | ||
| * compatible open source license. | ||
| * | ||
| * Modifications Copyright OpenSearch Contributors. See | ||
| * GitHub history for details. | ||
| */ | ||
|
|
||
| #ifndef OPENSEARCH_KNN_JNI_NMSLIB_STREAM_SUPPORT_H | ||
| #define OPENSEARCH_KNN_JNI_NMSLIB_STREAM_SUPPORT_H | ||
|
|
||
| #include "native_engines_stream_support.h" | ||
|
|
||
| namespace knn_jni { | ||
| namespace stream { | ||
|
|
||
|
|
||
|
|
||
| /** | ||
| * NmslibIOReader implementation delegating NativeEngineIndexInputMediator to read bytes. | ||
| */ | ||
| class NmslibOpenSearchIOReader final : public similarity::NmslibIOReader { | ||
| public: | ||
| explicit NmslibOpenSearchIOReader(NativeEngineIndexInputMediator *_mediator) | ||
| : mediator(_mediator) { | ||
| } | ||
|
|
||
| void read(char *bytes, size_t len) final { | ||
| if (len > 0) { | ||
| // Mediator calls IndexInput, then copy read bytes to `ptr`. | ||
| mediator->copyBytes(len, (uint8_t *) bytes); | ||
| } | ||
| } | ||
|
|
||
| size_t remainingBytes() final { | ||
| return mediator->remainingBytes(); | ||
| } | ||
|
|
||
| private: | ||
| NativeEngineIndexInputMediator *mediator; | ||
| }; // class NmslibOpenSearchIOReader | ||
|
|
||
|
|
||
|
|
||
| } | ||
| } | ||
|
|
||
| #endif //OPENSEARCH_KNN_JNI_NMSLIB_STREAM_SUPPORT_H |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.