forked from mindspore-Ecosystem/mindspore
!7721 fix java msconfig bug
Merge pull request !7721 from yeyunpeng2020/java
This commit is contained in:
commit
c9d6a7880c
|
@ -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");
|
||||
|
|
Loading…
Reference in New Issue