From f4bf1d514f537af9166f72fa00feda04556fc3d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alexander=20F=2E=20R=C3=B8dseth?= <rodseth@gmail.com>
Date: Sun, 14 Jan 2024 13:40:36 +0100
Subject: [PATCH] Let gpu.go and gen_linux.sh also find CUDA on Arch Linux

---
 gpu/gpu.go                |  1 +
 llm/generate/gen_linux.sh | 50 ++++++++++++++++++++-------------------
 2 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/gpu/gpu.go b/gpu/gpu.go
index 9b3d51b6..57cf7d5d 100644
--- a/gpu/gpu.go
+++ b/gpu/gpu.go
@@ -39,6 +39,7 @@ var CudaLinuxGlobs = []string{
 	"/usr/lib/x86_64-linux-gnu/libnvidia-ml.so*",
 	"/usr/lib/wsl/lib/libnvidia-ml.so*",
 	"/opt/cuda/lib64/libnvidia-ml.so*",
+	"/opt/cuda/targets/x86_64-linux/lib/stubs/libnvidia-ml.so*",
 	"/usr/lib*/libnvidia-ml.so*",
 	"/usr/local/lib*/libnvidia-ml.so*",
 	"/usr/lib/aarch64-linux-gnu/nvidia/current/libnvidia-ml.so*",
diff --git a/llm/generate/gen_linux.sh b/llm/generate/gen_linux.sh
index 0c940ba5..a6e665c3 100755
--- a/llm/generate/gen_linux.sh
+++ b/llm/generate/gen_linux.sh
@@ -109,31 +109,33 @@ else
     echo "Skipping CPU generation step as requested"
 fi
 
-if [ -d /usr/local/cuda/lib64/ ]; then
-    echo "CUDA libraries detected - building dynamic CUDA library"
-    init_vars
-    CUDA_MAJOR=$(ls /usr/local/cuda/lib64/libcudart.so.* | head -1 | cut -f3 -d. || true)
-    if [ -n "${CUDA_MAJOR}" ]; then
-        CUDA_VARIANT=_v${CUDA_MAJOR}
+for cudalibpath in "/usr/local/cuda/lib64" "/opt/cuda/targets/x86_64-linux/lib"; do
+    if [ -d "$cudalibpath" ]; then
+        echo "CUDA libraries detected - building dynamic CUDA library"
+        init_vars
+        CUDA_MAJOR=$(find "$cudalibpath" -name 'libcudart.so.*' -print | head -1 | cut -f3 -d. || true)
+        if [ -n "${CUDA_MAJOR}" ]; then
+            CUDA_VARIANT="_v${CUDA_MAJOR}"
+        fi
+        CMAKE_DEFS="-DLLAMA_CUBLAS=on ${COMMON_CMAKE_DEFS} ${CMAKE_DEFS}"
+        BUILD_DIR="${LLAMACPP_DIR}/build/linux/cuda${CUDA_VARIANT}"
+        CUDA_LIB_DIR="$cudalibpath"
+        build
+        install
+        gcc -fPIC -g -shared -o "${BUILD_DIR}/lib/libext_server.so" \
+            -Wl,--whole-archive \
+            "${BUILD_DIR}/lib/libext_server.a" \
+            "${BUILD_DIR}/lib/libcommon.a" \
+            "${BUILD_DIR}/lib/libllama.a" \
+            -Wl,--no-whole-archive \
+            "${CUDA_LIB_DIR}/libcudart_static.a" \
+            "${CUDA_LIB_DIR}/libcublas_static.a" \
+            "${CUDA_LIB_DIR}/libcublasLt_static.a" \
+            "${CUDA_LIB_DIR}/libcudadevrt.a" \
+            "${CUDA_LIB_DIR}/libculibos.a" \
+            -lrt -lpthread -ldl -lstdc++ -lm
     fi
-    CMAKE_DEFS="-DLLAMA_CUBLAS=on ${COMMON_CMAKE_DEFS} ${CMAKE_DEFS}"
-    BUILD_DIR="${LLAMACPP_DIR}/build/linux/cuda${CUDA_VARIANT}"
-    CUDA_LIB_DIR=/usr/local/cuda/lib64
-    build
-    install
-    gcc -fPIC -g -shared -o ${BUILD_DIR}/lib/libext_server.so \
-        -Wl,--whole-archive \
-        ${BUILD_DIR}/lib/libext_server.a \
-        ${BUILD_DIR}/lib/libcommon.a \
-        ${BUILD_DIR}/lib/libllama.a \
-        -Wl,--no-whole-archive \
-        ${CUDA_LIB_DIR}/libcudart_static.a \
-        ${CUDA_LIB_DIR}/libcublas_static.a \
-        ${CUDA_LIB_DIR}/libcublasLt_static.a \
-        ${CUDA_LIB_DIR}/libcudadevrt.a \
-        ${CUDA_LIB_DIR}/libculibos.a \
-        -lrt -lpthread -ldl -lstdc++ -lm
-fi
+done
 
 if [ -z "${ROCM_PATH}" ]; then
     # Try the default location in case it exists