From 685816bb0a648e4467b8fa710a301ef3bd3d5774 Mon Sep 17 00:00:00 2001 From: Jhen-Jie Hong Date: Fri, 16 Jun 2023 17:38:41 +0800 Subject: [PATCH] [js/rn] Add executionProviders support (#16233) ### Description This PR adds support for `executionProviders` option for react-native package, support: - Android: cpu / xnnpack / nnapi - iOS: cpu / xnnpack / coreml ### Motivation and Context In my case I want to enable Core ML / NNAPI EP for react-native project. --- js/common/lib/inference-session.ts | 16 ++++++++ .../reactnative/OnnxruntimeModule.java | 41 +++++++++++++++++++ js/react_native/ios/OnnxruntimeModule.mm | 40 ++++++++++++++++++ js/react_native/lib/index.ts | 7 ++++ 4 files changed, 104 insertions(+) diff --git a/js/common/lib/inference-session.ts b/js/common/lib/inference-session.ts index a79bcd7ebc..846b3ea37a 100644 --- a/js/common/lib/inference-session.ts +++ b/js/common/lib/inference-session.ts @@ -167,6 +167,7 @@ export declare namespace InferenceSession { // Backend Node.js binding: supports 'cpu' and 'cuda'. // Backend WebAssembly: supports 'cpu', 'wasm', 'xnnpack' and 'webnn'. // Backend ONNX.js: supports 'webgl'. + // Backend React Native: supports 'cpu', 'xnnpack', 'coreml' (iOS), 'nnapi' (Android). interface ExecutionProviderOptionMap { cpu: CpuExecutionProviderOption; cuda: CudaExecutionProviderOption; @@ -174,6 +175,8 @@ export declare namespace InferenceSession { webgl: WebGLExecutionProviderOption; xnnpack: XnnpackExecutionProviderOption; webnn: WebNNExecutionProviderOption; + coreml: CoreMLExecutionProviderOption; + nnapi: NnapiExecutionProviderOption; } type ExecutionProviderName = keyof ExecutionProviderOptionMap; @@ -206,6 +209,19 @@ export declare namespace InferenceSession { deviceType?: 'cpu'|'gpu'; powerPreference?: 'default'|'low-power'|'high-performance'; } + export interface CoreMLExecutionProviderOption extends ExecutionProviderOption { + readonly name: 'coreml'; + useCPUOnly?: boolean; + enableOnSubgraph?: boolean; + onlyEnableDeviceWithANE?: boolean; + } + export interface NnapiExecutionProviderOption extends ExecutionProviderOption { + readonly name: 'nnapi'; + useFP16?: boolean; + useNCHW?: boolean; + cpuDisabled?: boolean; + cpuOnly?: boolean; + } // #endregion // #endregion diff --git a/js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java b/js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java index 6ef550cde2..fd085f9533 100644 --- a/js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java +++ b/js/react_native/android/src/main/java/ai/onnxruntime/reactnative/OnnxruntimeModule.java @@ -11,6 +11,7 @@ import ai.onnxruntime.OrtSession; import ai.onnxruntime.OrtSession.Result; import ai.onnxruntime.OrtSession.RunOptions; import ai.onnxruntime.OrtSession.SessionOptions; +import ai.onnxruntime.providers.NNAPIFlags; import android.net.Uri; import android.os.Build; import android.util.Log; @@ -34,6 +35,8 @@ import java.io.InputStream; import java.io.InputStreamReader; import java.io.Reader; import java.math.BigInteger; +import java.util.Collections; +import java.util.EnumSet; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; @@ -367,6 +370,44 @@ public class OnnxruntimeModule extends ReactContextBaseJavaModule implements Lif } } + if (options.hasKey("executionProviders")) { + ReadableArray executionProviders = options.getArray("executionProviders"); + for (int i = 0; i < executionProviders.size(); ++i) { + String epName = null; + ReadableMap epOptions = null; + if (executionProviders.getType(i) == ReadableType.String) { + epName = executionProviders.getString(i); + } else { + epOptions = executionProviders.getMap(i); + epName = epOptions.getString("name"); + } + if (epName.equals("nnapi")) { + EnumSet flags = EnumSet.noneOf(NNAPIFlags.class); + if (epOptions != null) { + if (epOptions.hasKey("useFP16") && epOptions.getBoolean("useFP16")) { + flags.add(NNAPIFlags.USE_FP16); + } + if (epOptions.hasKey("useNCHW") && epOptions.getBoolean("useNCHW")) { + flags.add(NNAPIFlags.USE_NCHW); + } + if (epOptions.hasKey("cpuDisabled") && epOptions.getBoolean("cpuDisabled")) { + flags.add(NNAPIFlags.CPU_DISABLED); + } + if (epOptions.hasKey("cpuOnly") && epOptions.getBoolean("cpuOnly")) { + flags.add(NNAPIFlags.CPU_ONLY); + } + } + sessionOptions.addNnapi(flags); + } else if (epName.equals("xnnpack")) { + sessionOptions.addXnnpack(Collections.emptyMap()); + } else if (epName.equals("cpu")) { + continue; + } else { + throw new OrtException("Unsupported execution provider: " + epName); + } + } + } + if (options.hasKey("logId")) { String logId = options.getString("logId"); sessionOptions.setLoggerId(logId); diff --git a/js/react_native/ios/OnnxruntimeModule.mm b/js/react_native/ios/OnnxruntimeModule.mm index 117b737974..078821a467 100644 --- a/js/react_native/ios/OnnxruntimeModule.mm +++ b/js/react_native/ios/OnnxruntimeModule.mm @@ -17,9 +17,11 @@ // https://google.github.io/styleguide/objcguide.html#import-and-include // https://microsoft.github.io/objc-guide/Headers/ImportAndInclude.html #ifdef ORT_ENABLE_EXTENSIONS +#include "coreml_provider_factory.h" #include "onnxruntime_cxx_api.h" #include "onnxruntime_extensions.h" #else +#include "onnxruntime/coreml_provider_factory.h" #include "onnxruntime/onnxruntime_cxx_api.h" #endif @@ -371,6 +373,44 @@ static NSDictionary *executionModeTable = @{@"sequential" : @(ORT_SEQUENTIAL), @ } } + if ([options objectForKey:@"executionProviders"]) { + NSArray *executionProviders = [options objectForKey:@"executionProviders"]; + for (auto *executionProvider in executionProviders) { + NSString *epName = nil; + bool useOptions = false; + if ([executionProvider isKindOfClass:[NSString class]]) { + epName = (NSString *)executionProvider; + } else { + epName = [executionProvider objectForKey:@"name"]; + useOptions = true; + } + if ([epName isEqualToString:@"coreml"]) { + uint32_t coreml_flags = 0; + if (useOptions) { + if ([[executionProvider objectForKey:@"useCPUOnly"] boolValue]) { + coreml_flags |= COREML_FLAG_USE_CPU_ONLY; + } + if ([[executionProvider objectForKey:@"enableOnSubgraph"] boolValue]) { + coreml_flags |= COREML_FLAG_ENABLE_ON_SUBGRAPH; + } + if ([[executionProvider objectForKey:@"onlyEnableDeviceWithANE"] boolValue]) { + coreml_flags |= COREML_FLAG_ONLY_ENABLE_DEVICE_WITH_ANE; + } + } + Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CoreML(sessionOptions, coreml_flags)); + } else if ([epName isEqualToString:@"xnnpack"]) { + sessionOptions.AppendExecutionProvider("XNNPACK", {}); + } else if ([epName isEqualToString:@"cpu"]) { + continue; + } else { + NSException *exception = [NSException exceptionWithName:@"onnxruntime" + reason:@"unsupported execution provider" + userInfo:nil]; + @throw exception; + } + } + } + if ([options objectForKey:@"logId"]) { NSString *logId = [[options objectForKey:@"logId"] stringValue]; sessionOptions.SetLogId([logId UTF8String]); diff --git a/js/react_native/lib/index.ts b/js/react_native/lib/index.ts index a20c5c6a46..b6b559ceb3 100644 --- a/js/react_native/lib/index.ts +++ b/js/react_native/lib/index.ts @@ -3,9 +3,16 @@ export * from 'onnxruntime-common'; import {registerBackend, env} from 'onnxruntime-common'; +import {Platform} from 'react-native'; import {onnxruntimeBackend} from './backend'; import {version} from './version'; registerBackend('cpu', onnxruntimeBackend, 1); +registerBackend('xnnpack', onnxruntimeBackend, 1); +if (Platform.OS === 'android') { + registerBackend('nnapi', onnxruntimeBackend, 1); +} else if (Platform.OS === 'ios') { + registerBackend('coreml', onnxruntimeBackend, 1); +} env.versions['react-native'] = version;