mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-19 21:32:23 +00:00
### 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 ---------
37 lines
1.3 KiB
C++
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
|