From 5a473216b7e0143dd2b59a471aec8384afcf5c05 Mon Sep 17 00:00:00 2001 From: Adam Pocock Date: Wed, 24 Feb 2021 13:25:05 -0500 Subject: [PATCH] [Java] Adds extra providers (#6770) Add providers for CoreML, ROCM, NNAPI, ArmNN Adding the structs for OrtCUDAProviderOptions and OrtOpenVINOProviderOptions Updating NNAPI flags. Adding the new CoreML flag. Adding hooks to the build system to tell Java about the new providers. --- cmake/onnxruntime_java.cmake | 9 +++ .../main/java/ai/onnxruntime/OrtProvider.java | 6 +- .../main/java/ai/onnxruntime/OrtSession.java | 70 ++++++++++++++++++- .../ai/onnxruntime/providers/CoreMLFlags.java | 23 ++++++ .../ai/onnxruntime/providers/NNAPIFlags.java | 23 ++++++ .../ai/onnxruntime/providers/OrtFlags.java | 35 ++++++++++ ...ai_onnxruntime_OrtSession_SessionOptions.c | 52 ++++++++++++++ 7 files changed, 214 insertions(+), 4 deletions(-) create mode 100644 java/src/main/java/ai/onnxruntime/providers/CoreMLFlags.java create mode 100644 java/src/main/java/ai/onnxruntime/providers/NNAPIFlags.java create mode 100644 java/src/main/java/ai/onnxruntime/providers/OrtFlags.java 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 +} +