From 15016413de1cfeb9dbac8e8ddeec20759623bcde Mon Sep 17 00:00:00 2001
From: Jeremy <remy415@gmail.com>
Date: Thu, 18 Apr 2024 19:27:16 -0400
Subject: [PATCH] Update gen_windows.ps1

Added OLLAMA_CUSTOM_CUDA_DEFS and OLLAMA_CUSTOM_ROCM_DEFS to customize GPU builds on Windows
---
 llm/generate/gen_windows.ps1 | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/llm/generate/gen_windows.ps1 b/llm/generate/gen_windows.ps1
index 0d2ae57f..074516d6 100644
--- a/llm/generate/gen_windows.ps1
+++ b/llm/generate/gen_windows.ps1
@@ -243,6 +243,11 @@ if ($null -ne $script:CUDA_LIB_DIR) {
     init_vars
     $script:buildDir="../build/windows/${script:ARCH}/cuda$script:CUDA_VARIANT"
     $script:cmakeDefs += @("-A", "x64", "-DLLAMA_CUDA=ON", "-DLLAMA_AVX=on", "-DLLAMA_AVX2=off", "-DCUDAToolkit_INCLUDE_DIR=$script:CUDA_INCLUDE_DIR", "-DCMAKE_CUDA_ARCHITECTURES=${script:CMAKE_CUDA_ARCHITECTURES}")
+    if ($null -ne $script:OLLAMA_CUSTOM_CUDA_DEFS) {
+        write-host "OLLAMA_CUSTOM_CUDA_DEFS=`"${script:OLLAMA_CUSTOM_CUDA_DEFS}`"
+        $script:cmakeDefs +=@("${script:OLLAMA_CUSTOM_CUDA_DEFS}")
+        write-host "building custom CUDA GPU"
+    }
     build
     sign
     compress
@@ -274,7 +279,11 @@ if ($null -ne $env:HIP_PATH) {
 
     # We have to clobber the LIB var from the developer shell for clang to work properly
     $env:LIB=""
-
+    if ($null -ne $script:OLLAMA_CUSTOM_ROCM_DEFS) {
+        write-host "OLLAMA_CUSTOM_ROCM_DEFS=`"${script:OLLAMA_CUSTOM_ROCM_DEFS}`"
+        $script:cmakeDefs += @("${script:OLLAMA_CUSTOM_ROCM_DEFS}")
+        write-host "building custom ROCM GPU"
+    }
     write-host "Building ROCm"
     build
     # Ninja doesn't prefix with config name