Added DML and CUDA provider support in onnxruntime-node

js/common/lib/inference-session.ts

@@ -170,6 +170,7 @@ export declare namespace InferenceSession {
   interface ExecutionProviderOptionMap {
     cpu: CpuExecutionProviderOption;
     cuda: CudaExecutionProviderOption;
+    dml: DmlExecutionProviderOption;
     wasm: WebAssemblyExecutionProviderOption;
     webgl: WebGLExecutionProviderOption;
     xnnpack: XnnpackExecutionProviderOption;
@@ -191,6 +192,10 @@ export declare namespace InferenceSession {
     readonly name: 'cuda';
     deviceId?: number;
   }
+  export interface DmlExecutionProviderOption extends ExecutionProviderOption {
+    readonly name: 'dml';
+    deviceId?: number;
+  }
   export interface WebAssemblyExecutionProviderOption extends ExecutionProviderOption {
     readonly name: 'wasm';
   }

js/node/CMakeLists.txt

@@ -28,6 +28,7 @@ endif()
 # include dirs
 include_directories(${CMAKE_JS_INC})
 include_directories(${CMAKE_SOURCE_DIR}/../../include/onnxruntime/core/session)
+include_directories(${CMAKE_SOURCE_DIR}/../../include/onnxruntime/core/providers/dml)
 include_directories(${CMAKE_SOURCE_DIR}/node_modules/node-addon-api)
 
 # source files
@@ -76,6 +77,9 @@ if (WIN32)
     COMMAND ${CMAKE_COMMAND} -E copy
       ${ONNXRUNTIME_BUILD_DIR}/${CMAKE_BUILD_TYPE}/onnxruntime.dll
       ${dist_folder}
+    COMMAND ${CMAKE_COMMAND} -E copy
+      ${ONNXRUNTIME_BUILD_DIR}/${CMAKE_BUILD_TYPE}/DirectML.dll
+      ${dist_folder}
   )
   if (CMAKE_BUILD_TYPE STREQUAL "Debug")
     add_custom_command(

js/node/lib/index.ts

@@ -6,3 +6,5 @@ import {registerBackend} from 'onnxruntime-common';
 import {onnxruntimeBackend} from './backend';
 
 registerBackend('cpu', onnxruntimeBackend, 100);
+registerBackend('cuda', onnxruntimeBackend, 100);
+registerBackend('dml', onnxruntimeBackend, 100);

js/node/src/directml_load_helper.cc

@@ -0,0 +1,44 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#ifdef _WIN32
+#include "common.h"
+#include "windows.h"
+
+void LoadDirectMLDll(Napi::Env env)
+{
+  DWORD pathLen = MAX_PATH;
+  wchar_t* path = new wchar_t[pathLen];
+  HMODULE moduleHandle = nullptr;
+
+  GetModuleHandleEx(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS |
+                                 GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+                             reinterpret_cast<LPCSTR>(&LoadDirectMLDll), &moduleHandle);
+
+  DWORD getModuleFileNameResult = GetModuleFileNameW(moduleHandle, path, pathLen);
+  while (getModuleFileNameResult == 0 || getModuleFileNameResult == pathLen) {
+    delete[] path;
+
+    pathLen = pathLen * 2;
+    int ret = GetLastError();
+    if (ret == ERROR_INSUFFICIENT_BUFFER && pathLen < 32768) {
+      path = new wchar_t[pathLen];
+      getModuleFileNameResult = GetModuleFileNameW(moduleHandle, path, pathLen);
+    } else {
+      ORT_NAPI_THROW_ERROR(env, "Failed getting path to load DirectML.dll, error code: ", ret);
+    }
+  }
+
+  // assume binding name will be always longer (onnxruntime_binding.node)
+  wchar_t* lastBackslash = wcsrchr(path, L'\\');
+  *lastBackslash = L'\0';
+  wcscat_s(path, pathLen, L"\\DirectML.dll");
+  auto libraryLoadResult = LoadLibraryW(path);
+  delete[] path;
+
+  if (!libraryLoadResult) {
+    int ret = GetLastError();
+    ORT_NAPI_THROW_ERROR(env, "Failed loading bundled DirectML.dll, error code: ", ret);
+  }
+}
+#endif

js/node/src/directml_load_helper.h

@@ -0,0 +1,6 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#ifdef _WIN32
+void LoadDirectMLDll(Napi::Env env);
+#endif

js/node/src/inference_session_wrap.cc

@@ -8,11 +8,15 @@
 #include "run_options_helper.h"
 #include "session_options_helper.h"
 #include "tensor_helper.h"
+#include "directml_load_helper.h"
 
 Napi::FunctionReference InferenceSessionWrap::constructor;
 Ort::Env *InferenceSessionWrap::ortEnv;
 
 Napi::Object InferenceSessionWrap::Init(Napi::Env env, Napi::Object exports) {
+#ifdef _WIN32
+  LoadDirectMLDll(env);
+#endif
   // create ONNX runtime env
   Ort::InitApi();
   ortEnv = new Ort::Env{ORT_LOGGING_LEVEL_WARNING, "onnxruntime-node"};
@@ -151,14 +155,15 @@ Napi::Value InferenceSessionWrap::Run(const Napi::CallbackInfo &info) {
   std::vector<bool> reuseOutput;
   size_t inputIndex = 0;
   size_t outputIndex = 0;
+  OrtMemoryInfo* memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault).release();
 
   try {
     for (auto &name : inputNames_) {
       if (feed.Has(name)) {
         inputIndex++;
         inputNames_cstr.push_back(name.c_str());
         auto value = feed.Get(name);
-        inputValues.push_back(NapiValueToOrtValue(env, value));
+        inputValues.push_back(NapiValueToOrtValue(env, value, memory_info));
       }
     }
     for (auto &name : outputNames_) {
@@ -167,7 +172,7 @@ Napi::Value InferenceSessionWrap::Run(const Napi::CallbackInfo &info) {
         outputNames_cstr.push_back(name.c_str());
         auto value = fetch.Get(name);
         reuseOutput.push_back(!value.IsNull());
-        outputValues.emplace_back(value.IsNull() ? Ort::Value{nullptr} : NapiValueToOrtValue(env, value));
+        outputValues.emplace_back(value.IsNull() ? Ort::Value{nullptr} : NapiValueToOrtValue(env, value, memory_info));
       }
     }
 

js/node/src/session_options_helper.cc

@@ -9,6 +9,9 @@
 
 #include "common.h"
 #include "session_options_helper.h"
+#ifdef _WIN32
+#include "dml_provider_factory.h"
+#endif
 
 const std::unordered_map<std::string, GraphOptimizationLevel> GRAPH_OPT_LEVEL_NAME_TO_ID_MAP = {
     {"disabled", ORT_DISABLE_ALL},
@@ -23,21 +26,32 @@ void ParseExecutionProviders(const Napi::Array epList, Ort::SessionOptions &sess
   for (uint32_t i = 0; i < epList.Length(); i++) {
     Napi::Value epValue = epList[i];
     std::string name;
+    int deviceId = 0;
     if (epValue.IsString()) {
       name = epValue.As<Napi::String>().Utf8Value();
     } else if (!epValue.IsObject() || epValue.IsNull() || !epValue.As<Napi::Object>().Has("name") ||
                !epValue.As<Napi::Object>().Get("name").IsString()) {
       ORT_NAPI_THROW_TYPEERROR(epList.Env(), "Invalid argument: sessionOptions.executionProviders[", i,
                                "] must be either a string or an object with property 'name'.");
     } else {
-      name = epValue.As<Napi::Object>().Get("name").As<Napi::String>().Utf8Value();
+      auto obj = epValue.As<Napi::Object>();
+      name = obj.Get("name").As<Napi::String>().Utf8Value();
+      if (obj.Has("deviceId")) {
+        deviceId = obj.Get("deviceId").As<Napi::Number>();
+      }
     }
 
     // CPU execution provider
     if (name == "cpu") {
       // TODO: handling CPU EP options
     } else if (name == "cuda") {
-      // TODO: handling Cuda EP options
+      Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_CUDA(sessionOptions, deviceId));
+    } else if (name == "dml") {
+#ifdef _WIN32
+      Ort::ThrowOnError(OrtSessionOptionsAppendExecutionProvider_DML(sessionOptions, deviceId));
+#else
+      ORT_NAPI_THROW_ERROR(epList.Env(), "DirectML EP is supported only on Windows");
+#endif
     } else {
       ORT_NAPI_THROW_ERROR(epList.Env(), "Invalid argument: sessionOptions.executionProviders[", i,
                            "] is unsupported: '", name, "'.");

js/node/src/tensor_helper.cc

@@ -106,7 +106,7 @@ const std::unordered_map<std::string, ONNXTensorElementDataType> DATA_TYPE_NAME_
     {"uint64", ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT64}};
 
 // currently only support tensor
-Ort::Value NapiValueToOrtValue(Napi::Env env, Napi::Value value) {
+Ort::Value NapiValueToOrtValue(Napi::Env env, Napi::Value value, OrtMemoryInfo* memory_info) {
   ORT_NAPI_THROW_TYPEERROR_IF(!value.IsObject(), env, "Tensor must be an object.");
 
   // check 'dims'
@@ -180,7 +180,6 @@ Ort::Value NapiValueToOrtValue(Napi::Env env, Napi::Value value) {
                                 "Tensor.data must be a typed array (", DATA_TYPE_TYPEDARRAY_MAP[elemType], ") for ",
                                 tensorTypeString, " tensors, but got typed array (", typedArrayType, ").");
 
-    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
     char *buffer = reinterpret_cast<char *>(tensorDataTypedArray.ArrayBuffer().Data());
     size_t bufferByteOffset = tensorDataTypedArray.ByteOffset();
     // there is a bug in TypedArray::ElementSize(): https://github.com/nodejs/node-addon-api/pull/705

js/node/src/tensor_helper.h

@@ -9,7 +9,7 @@
 #include "onnxruntime_cxx_api.h"
 
 // convert a Javascript OnnxValue object to an OrtValue object
-Ort::Value NapiValueToOrtValue(Napi::Env env, Napi::Value value);
+Ort::Value NapiValueToOrtValue(Napi::Env env, Napi::Value value, OrtMemoryInfo* memory_info);
 
 // convert an OrtValue object to a Javascript OnnxValue object
 Napi::Value OrtValueToNapiValue(Napi::Env env, Ort::Value &value);