|
|
|
@@ -14,16 +14,12 @@ |
|
|
|
* limitations under the License. |
|
|
|
*/ |
|
|
|
|
|
|
|
|
|
|
|
#include <jni.h> |
|
|
|
#include "common/ms_log.h" |
|
|
|
#include "include/context.h" |
|
|
|
|
|
|
|
extern "C" JNIEXPORT jlong JNICALL Java_com_mindspore_lite_config_MSConfig_createMSConfig(JNIEnv *env, jobject thiz, |
|
|
|
jint device_type, |
|
|
|
jint thread_num, |
|
|
|
jint cpu_bind_mode, |
|
|
|
jboolean enable_float16) { |
|
|
|
extern "C" JNIEXPORT jlong JNICALL Java_com_mindspore_lite_config_MSConfig_createMSConfig( |
|
|
|
JNIEnv *env, jobject thiz, jint device_type, jint thread_num, jint cpu_bind_mode, jboolean enable_float16) { |
|
|
|
auto *context = new (std::nothrow) mindspore::lite::Context(); |
|
|
|
if (context == nullptr) { |
|
|
|
MS_LOGE("new Context fail!"); |
|
|
|
@@ -35,11 +31,13 @@ extern "C" JNIEXPORT jlong JNICALL Java_com_mindspore_lite_config_MSConfig_creat |
|
|
|
case 0: |
|
|
|
context->device_list_[0].device_type_ = mindspore::lite::DT_CPU; |
|
|
|
break; |
|
|
|
case 1: |
|
|
|
context->device_list_[0].device_type_ = mindspore::lite::DT_GPU; |
|
|
|
case 1: // DT_GPU |
|
|
|
MS_LOGE("We only support CPU now."); |
|
|
|
return (jlong)context; |
|
|
|
break; |
|
|
|
case 2: |
|
|
|
context->device_list_[0].device_type_ = mindspore::lite::DT_NPU; |
|
|
|
case 2: // DT_NPU |
|
|
|
MS_LOGE("We only support CPU now."); |
|
|
|
return (jlong)context; |
|
|
|
break; |
|
|
|
default: |
|
|
|
MS_LOGE("Invalid device_type : %d", device_type); |
|
|
|
@@ -63,17 +61,12 @@ extern "C" JNIEXPORT jlong JNICALL Java_com_mindspore_lite_config_MSConfig_creat |
|
|
|
|
|
|
|
cpu_device_ctx.device_info_.cpu_device_info_.enable_float16_ = enable_float16; |
|
|
|
|
|
|
|
if (device_type == mindspore::lite::DT_GPU) { |
|
|
|
mindspore::lite::DeviceContext gpu_device_ctx{mindspore::lite::DT_GPU, {false}}; |
|
|
|
gpu_device_ctx.device_info_.gpu_device_info_.enable_float16_ = enable_float16; |
|
|
|
context->device_list_.push_back(gpu_device_ctx); |
|
|
|
} |
|
|
|
context->thread_num_ = thread_num; |
|
|
|
return (jlong)context; |
|
|
|
} |
|
|
|
|
|
|
|
extern "C" JNIEXPORT void JNICALL Java_com_mindspore_lite_config_MSConfig_free(JNIEnv *env, jobject thiz, |
|
|
|
jlong context_ptr) { |
|
|
|
jlong context_ptr) { |
|
|
|
auto *pointer = reinterpret_cast<void *>(context_ptr); |
|
|
|
if (pointer == nullptr) { |
|
|
|
MS_LOGE("Context pointer from java is nullptr"); |
|
|
|
|