onnxruntime/js/node/src/directml_load_helper.cc
Arthur Islamov c262879214
Added DML and CUDA provider support in onnxruntime-node (#16050)
### Description
I've added changes to support CUDA and DML (only on Windows, on other
platforms it will throw an error)



### Motivation and Context
It fixes this feature request
https://github.com/microsoft/onnxruntime/issues/14127 which is tracked
here https://github.com/microsoft/onnxruntime/issues/14529

I was working on StableDiffusion implementation for node.js and it is
very slow on CPU, so GPU support is essential.

Here is a working demo with a patched and precompiled version
https://github.com/dakenf/stable-diffusion-nodejs

---------
2023-08-25 16:57:06 -07:00

37 lines
1.3 KiB
C++

// 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;
std::wstring path(pathLen, L'\0');
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, const_cast<wchar_t *>(path.c_str()), pathLen);
while (getModuleFileNameResult == 0 || getModuleFileNameResult == pathLen) {
int ret = GetLastError();
if (ret == ERROR_INSUFFICIENT_BUFFER && pathLen < 32768) {
pathLen *= 2;
path.resize(pathLen);
getModuleFileNameResult = GetModuleFileNameW(moduleHandle, const_cast<wchar_t *>(path.c_str()), pathLen);
} else {
ORT_NAPI_THROW_ERROR(env, "Failed getting path to load DirectML.dll, error code: ", ret);
}
}
path.resize(path.rfind(L'\\') + 1);
path.append(L"DirectML.dll");
HMODULE libraryLoadResult = LoadLibraryW(path.c_str());
if (!libraryLoadResult) {
int ret = GetLastError();
ORT_NAPI_THROW_ERROR(env, "Failed loading bundled DirectML.dll, error code: ", ret);
}
}
#endif