From 2025d2d15b8571643241d145085d7cc6cd1d331b Mon Sep 17 00:00:00 2001
From: Jared Van Bortel <jared@nomic.ai>
Date: Thu, 16 May 2024 17:39:49 -0400
Subject: [PATCH] llmodel: add CUDA to the DLL search path if CUDA_PATH is set
 (#2357)

Signed-off-by: Jared Van Bortel <jared@nomic.ai>
---
 gpt4all-backend/dlhandle.h  |  5 +++--
 gpt4all-backend/llmodel.cpp | 26 +++++++++++++++++++++++++-
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/gpt4all-backend/dlhandle.h b/gpt4all-backend/dlhandle.h
index e76162d7..f24a2a25 100644
--- a/gpt4all-backend/dlhandle.h
+++ b/gpt4all-backend/dlhandle.h
@@ -58,14 +58,15 @@ public:
 #include <string>
 #include <exception>
 #include <stdexcept>
+
+#define WIN32_LEAN_AND_MEAN
 #ifndef NOMINMAX
-    #define NOMINMAX
+#   define NOMINMAX
 #endif
 #include <windows.h>
 #include <libloaderapi.h>
 
 
-
 class Dlhandle {
     HMODULE chandle;
 
diff --git a/gpt4all-backend/llmodel.cpp b/gpt4all-backend/llmodel.cpp
index 3f359c9e..fa756056 100644
--- a/gpt4all-backend/llmodel.cpp
+++ b/gpt4all-backend/llmodel.cpp
@@ -15,8 +15,16 @@
 #include <unordered_map>
 #include <vector>
 
+#ifdef _WIN32
+#   define WIN32_LEAN_AND_MEAN
+#   ifndef NOMINMAX
+#       define NOMINMAX
+#   endif
+#   include <windows.h>
+#endif
+
 #ifdef _MSC_VER
-#include <intrin.h>
+#   include <intrin.h>
 #endif
 
 #ifndef __APPLE__
@@ -85,6 +93,20 @@ static bool isImplementation(const Dlhandle &dl) {
     return dl.get<bool(uint32_t)>("is_g4a_backend_model_implementation");
 }
 
+// Add the CUDA Toolkit to the DLL search path on Windows.
+// This is necessary for chat.exe to find CUDA when started from Qt Creator.
+static void addCudaSearchPath() {
+#ifdef _WIN32
+    if (const auto *cudaPath = _wgetenv(L"CUDA_PATH")) {
+        auto libDir = std::wstring(cudaPath) + L"\\bin";
+        if (!AddDllDirectory(libDir.c_str())) {
+            auto err = GetLastError();
+            std::wcerr << L"AddDllDirectory(\"" << libDir << L"\") failed with error 0x" << std::hex << err << L"\n";
+        }
+    }
+#endif
+}
+
 const std::vector<LLModel::Implementation> &LLModel::Implementation::implementationList() {
     if (cpu_supports_avx() == 0) {
         throw std::runtime_error("CPU does not support AVX");
@@ -95,6 +117,8 @@ const std::vector<LLModel::Implementation> &LLModel::Implementation::implementat
     static auto* libs = new std::vector<Implementation>([] () {
         std::vector<Implementation> fres;
 
+        addCudaSearchPath();
+
         std::string impl_name_re = "(gptj|llamamodel-mainline)-(cpu|metal|kompute|vulkan|cuda)";
         if (cpu_supports_avx2() == 0) {
             impl_name_re += "-avxonly";