[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.
This commit is contained in:
Adam Pocock 2021-02-24 13:25:05 -05:00 committed by GitHub
parent 47c8e9ad28
commit 5a473216b7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
7 changed files with 214 additions and 4 deletions

View file

@ -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)

View file

@ -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<String, OrtProvider> valueMap = new HashMap<>(values().length);

View file

@ -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<NNAPIFlags> 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<CoreMLFlags> 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}. */

View file

@ -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;
}
}

View file

@ -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;
}
}

View file

@ -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 <E> The enum type to aggregate.
* @return The aggregated values
*/
public static <E extends Enum<E> & OrtFlags> int aggregateToInt(EnumSet<E> set) {
int value = 0;
for (OrtFlags flag : set) {
value |= flag.getValue();
}
return value;
}
}

View file

@ -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
}