diff --git a/cmake/onnxruntime_java.cmake b/cmake/onnxruntime_java.cmake index 54eb81f833..dfe81df048 100644 --- a/cmake/onnxruntime_java.cmake +++ b/cmake/onnxruntime_java.cmake @@ -85,6 +85,15 @@ endif() if (onnxruntime_USE_DML) target_compile_definitions(onnxruntime4j_jni PRIVATE USE_DIRECTML=1) endif() +if (onnxruntime_USE_ARMNN) + target_compile_definitions(onnxruntime4j_jni PRIVATE USE_ARMNN=1) +endif() +if (onnxruntime_USE_ROCM) + target_compile_definitions(onnxruntime4j_jni PRIVATE USE_ROCM=1) +endif() +if (onnxruntime_USE_COREML) + target_compile_definitions(onnxruntime4j_jni PRIVATE USE_COREML=1) +endif() # depend on java sources. if they change, the JNI should recompile add_dependencies(onnxruntime4j_jni onnxruntime4j) diff --git a/java/src/main/java/ai/onnxruntime/OrtProvider.java b/java/src/main/java/ai/onnxruntime/OrtProvider.java index abcb9594ee..cf5202c432 100644 --- a/java/src/main/java/ai/onnxruntime/OrtProvider.java +++ b/java/src/main/java/ai/onnxruntime/OrtProvider.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020, Oracle and/or its affiliates. All rights reserved. + * Copyright (c) 2020, 2021, Oracle and/or its affiliates. All rights reserved. * Licensed under the MIT License. */ package ai.onnxruntime; @@ -21,7 +21,9 @@ public enum OrtProvider { DIRECT_ML("DmlExecutionProvider"), MI_GRAPH_X("MIGraphXExecutionProvider"), ACL("ACLExecutionProvider"), - ARM_NN("ArmNNExecutionProvider"); + ARM_NN("ArmNNExecutionProvider"), + ROCM("ROCMExecutionProvider"), + CORE_ML("CoreMLExecutionProvider"); private static final Map valueMap = new HashMap<>(values().length); diff --git a/java/src/main/java/ai/onnxruntime/OrtSession.java b/java/src/main/java/ai/onnxruntime/OrtSession.java index dd72dd5dea..286159b207 100644 --- a/java/src/main/java/ai/onnxruntime/OrtSession.java +++ b/java/src/main/java/ai/onnxruntime/OrtSession.java @@ -4,10 +4,14 @@ */ package ai.onnxruntime; +import ai.onnxruntime.providers.CoreMLFlags; +import ai.onnxruntime.providers.NNAPIFlags; +import ai.onnxruntime.providers.OrtFlags; import java.io.IOException; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.EnumSet; import java.util.Iterator; import java.util.LinkedHashMap; import java.util.LinkedHashSet; @@ -809,13 +813,23 @@ public class OrtSession implements AutoCloseable { } /** - * Adds Android's NNAPI as an execution backend. + * Adds Android's NNAPI as an execution backend. Uses the default empty flag. * * @throws OrtException If there was an error in native code. */ public void addNnapi() throws OrtException { + addNnapi(EnumSet.noneOf(NNAPIFlags.class)); + } + + /** + * Adds Android's NNAPI as an execution backend. + * + * @param flags The flags which control the NNAPI configuration. + * @throws OrtException If there was an error in native code. + */ + public void addNnapi(EnumSet flags) throws OrtException { checkClosed(); - addNnapi(OnnxRuntime.ortApiHandle, nativeHandle, 0); + addNnapi(OnnxRuntime.ortApiHandle, nativeHandle, OrtFlags.aggregateToInt(flags)); } /** @@ -852,6 +866,49 @@ public class OrtSession implements AutoCloseable { addACL(OnnxRuntime.ortApiHandle, nativeHandle, useArena ? 1 : 0); } + /** + * Adds the ARM Neural Net library as an execution backend. + * + * @param useArena If true use the arena memory allocator. + * @throws OrtException If there was an error in native code. + */ + public void addArmNN(boolean useArena) throws OrtException { + checkClosed(); + addArmNN(OnnxRuntime.ortApiHandle, nativeHandle, useArena ? 1 : 0); + } + + /** + * Adds ROCM as an execution backend. + * + * @param deviceID The ROCM device ID. + * @param memLimit The maximum amount of memory available. + * @throws OrtException If there was an error in native code. + */ + public void addROCM(int deviceID, long memLimit) throws OrtException { + checkClosed(); + addROCM(OnnxRuntime.ortApiHandle, nativeHandle, deviceID, memLimit); + } + + /** + * Adds Apple's CoreML as an execution backend. Uses the default empty flag. + * + * @throws OrtException If there was an error in native code. + */ + public void addCoreML() throws OrtException { + addCoreML(EnumSet.noneOf(CoreMLFlags.class)); + } + + /** + * Adds Apple's CoreML as an execution backend. + * + * @param flags The flags which control the CoreML configuration. + * @throws OrtException If there was an error in native code. + */ + public void addCoreML(EnumSet flags) throws OrtException { + checkClosed(); + addCoreML(OnnxRuntime.ortApiHandle, nativeHandle, OrtFlags.aggregateToInt(flags)); + } + private native void setExecutionMode(long apiHandle, long nativeHandle, int mode) throws OrtException; @@ -944,6 +1001,15 @@ public class OrtSession implements AutoCloseable { throws OrtException; private native void addACL(long apiHandle, long nativeHandle, int useArena) throws OrtException; + + private native void addArmNN(long apiHandle, long nativeHandle, int useArena) + throws OrtException; + + private native void addROCM(long apiHandle, long nativeHandle, int deviceID, long memLimit) + throws OrtException; + + private native void addCoreML(long apiHandle, long nativeHandle, int coreMLFlags) + throws OrtException; } /** Used to control logging and termination of a call to {@link OrtSession#run}. */ diff --git a/java/src/main/java/ai/onnxruntime/providers/CoreMLFlags.java b/java/src/main/java/ai/onnxruntime/providers/CoreMLFlags.java new file mode 100644 index 0000000000..7a9db43081 --- /dev/null +++ b/java/src/main/java/ai/onnxruntime/providers/CoreMLFlags.java @@ -0,0 +1,23 @@ +/* + * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved. + * Licensed under the MIT License. + */ +package ai.onnxruntime.providers; + +/** Flags for the CoreML provider. */ +public enum CoreMLFlags implements OrtFlags { + CPU_ONLY(1), // COREML_FLAG_USE_CPU_ONLY(0x001) + ENABLE_ON_SUBGRAPH(2), // COREML_FLAG_ENABLE_ON_SUBGRAPH(0x002) + ONLY_ENABLE_DEVICE_WITH_ANE(4); // COREML_FLAG_ONLY_ENABLE_DEVICE_WITH_ANE(0x004), + + public final int value; + + CoreMLFlags(int value) { + this.value = value; + } + + @Override + public int getValue() { + return value; + } +} diff --git a/java/src/main/java/ai/onnxruntime/providers/NNAPIFlags.java b/java/src/main/java/ai/onnxruntime/providers/NNAPIFlags.java new file mode 100644 index 0000000000..560d15b081 --- /dev/null +++ b/java/src/main/java/ai/onnxruntime/providers/NNAPIFlags.java @@ -0,0 +1,23 @@ +/* + * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved. + * Licensed under the MIT License. + */ +package ai.onnxruntime.providers; + +/** Flags for the NNAPI provider. */ +public enum NNAPIFlags implements OrtFlags { + USE_FP16(1), // NNAPI_FLAG_USE_FP16(0x001) + USE_NCHW(2), // NNAPI_FLAG_USE_NCHW(0x002) + CPU_DISABLED(4); // NNAPI_FLAG_CPU_DISABLED(0x004) + + public final int value; + + NNAPIFlags(int value) { + this.value = value; + } + + @Override + public int getValue() { + return value; + } +} diff --git a/java/src/main/java/ai/onnxruntime/providers/OrtFlags.java b/java/src/main/java/ai/onnxruntime/providers/OrtFlags.java new file mode 100644 index 0000000000..73d3eeae64 --- /dev/null +++ b/java/src/main/java/ai/onnxruntime/providers/OrtFlags.java @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2021, Oracle and/or its affiliates. All rights reserved. + * Licensed under the MIT License. + */ +package ai.onnxruntime.providers; + +import java.util.EnumSet; + +/** An interface for bitset enums that should be aggregated into a single integer. */ +public interface OrtFlags { + + /** + * Gets the underlying flag value. + * + * @return The flag value. + */ + public int getValue(); + + /** + * Converts an EnumSet of flags into the value expected by the C API. + * + * @param set The enum set to aggregate the values from. + * @param The enum type to aggregate. + * @return The aggregated values + */ + public static & OrtFlags> int aggregateToInt(EnumSet set) { + int value = 0; + + for (OrtFlags flag : set) { + value |= flag.getValue(); + } + + return value; + } +} diff --git a/java/src/main/native/ai_onnxruntime_OrtSession_SessionOptions.c b/java/src/main/native/ai_onnxruntime_OrtSession_SessionOptions.c index 21f8af6930..fc5cab8c80 100644 --- a/java/src/main/native/ai_onnxruntime_OrtSession_SessionOptions.c +++ b/java/src/main/native/ai_onnxruntime_OrtSession_SessionOptions.c @@ -23,6 +23,9 @@ #include "onnxruntime/core/providers/tensorrt/tensorrt_provider_factory.h" #include "onnxruntime/core/providers/migraphx/migraphx_provider_factory.h" #include "onnxruntime/core/providers/acl/acl_provider_factory.h" +#include "onnxruntime/core/providers/armnn/armnn_provider_factory.h" +#include "onnxruntime/core/providers/coreml/coreml_provider_factory.h" +#include "onnxruntime/core/providers/rocm/rocm_provider_factory.h" #ifdef USE_DIRECTML #include "onnxruntime/core/providers/dml/dml_provider_factory.h" #endif @@ -502,3 +505,52 @@ JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtSession_00024SessionOptions_addACL throwOrtException(jniEnv,convertErrorCode(ORT_INVALID_ARGUMENT),"This binary was not compiled with ACL support."); #endif } + +/* + * Class: ai_onnxruntime_OrtSession_SessionOptions + * Method: addArmNN + * Signature: (JJI)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtSession_00024SessionOptions_addArmNN + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong handle, jint useArena) { + (void)jobj; + #ifdef USE_ARMNN + checkOrtStatus(jniEnv,(const OrtApi*)apiHandle,OrtSessionOptionsAppendExecutionProvider_ArmNN((OrtSessionOptions*) handle,useArena)); + #else + (void)apiHandle;(void)handle;(void)useArena; // Parameters used when ARMNN is defined. + throwOrtException(jniEnv,convertErrorCode(ORT_INVALID_ARGUMENT),"This binary was not compiled with ArmNN support."); + #endif +} + +/* + * Class: ai_onnxruntime_OrtSession_SessionOptions + * Method: addCoreML + * Signature: (JJI)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtSession_00024SessionOptions_addCoreML + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong handle, jint coreMLFlags) { + (void)jobj; + #ifdef USE_CORE_ML + checkOrtStatus(jniEnv,(const OrtApi*)apiHandle,OrtSessionOptionsAppendExecutionProvider_CoreML((OrtSessionOptions*) handle, (uint32_t) coreMLFlags)); + #else + (void)apiHandle;(void)handle;(void)coreMLFlags; // Parameters used when CoreML is defined. + throwOrtException(jniEnv,convertErrorCode(ORT_INVALID_ARGUMENT),"This binary was not compiled with CoreML support."); + #endif +} + +/* + * Class: ai_onnxruntime_OrtSession_SessionOptions + * Method: addROCM + * Signature: (JJI)V + */ +JNIEXPORT void JNICALL Java_ai_onnxruntime_OrtSession_00024SessionOptions_addROCM + (JNIEnv * jniEnv, jobject jobj, jlong apiHandle, jlong handle, jint deviceID, jlong memLimit) { + (void)jobj; + #ifdef USE_ROCM + checkOrtStatus(jniEnv,(const OrtApi*)apiHandle,OrtSessionOptionsAppendExecutionProvider_ROCM((OrtSessionOptions*) handle, deviceID, (size_t) memLimit)); + #else + (void)apiHandle;(void)handle;(void)deviceID;(void)memLimit; // Parameters used when ROCM is defined. + throwOrtException(jniEnv,convertErrorCode(ORT_INVALID_ARGUMENT),"This binary was not compiled with ROCM support."); + #endif +} +