From 6459377ae081021a0ec2b5b408c6ac4503de01b2 Mon Sep 17 00:00:00 2001
From: Daniel Hiltgen <dhiltgen@users.noreply.github.com>
Date: Thu, 14 Mar 2024 18:00:16 -0700
Subject: [PATCH] Add ROCm support to linux install script (#2966)

---
 docs/troubleshooting.md |  9 +++++----
 scripts/install.sh      | 45 ++++++++++++++++++++++++++++++++++-------
 2 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md
index 46f5bf51..983613fe 100644
--- a/docs/troubleshooting.md
+++ b/docs/troubleshooting.md
@@ -101,11 +101,12 @@ future release which should increase support for more GPUs.
 Reach out on [Discord](https://discord.gg/ollama) or file an
 [issue](https://github.com/ollama/ollama/issues) for additional help.
 
-## Installing older versions on Linux
+## Installing older or pre-release versions on Linux
 
-If you run into problems on Linux and want to install an older version you can tell the install script
-which version to install.
+If you run into problems on Linux and want to install an older version, or you'd
+like to try out a pre-release before it's officially released, you can tell the
+install script which version to install.
 
 ```sh
-curl -fsSL https://ollama.com/install.sh | OLLAMA_VERSION="0.1.27" sh
+curl -fsSL https://ollama.com/install.sh | OLLAMA_VERSION="0.1.29" sh
 ```
diff --git a/scripts/install.sh b/scripts/install.sh
index f8f1d6f4..eb3ff504 100644
--- a/scripts/install.sh
+++ b/scripts/install.sh
@@ -40,6 +40,7 @@ case "$KERN" in
     *) ;;
 esac
 
+VER_PARAM="${OLLAMA_VERSION:+?version=$OLLAMA_VERSION}"
 
 SUDO=
 if [ "$(id -u)" -ne 0 ]; then
@@ -61,7 +62,7 @@ if [ -n "$NEEDS" ]; then
 fi
 
 status "Downloading ollama..."
-curl --fail --show-error --location --progress-bar -o $TEMP_DIR/ollama "https://ollama.com/download/ollama-linux-$ARCH"
+curl --fail --show-error --location --progress-bar -o $TEMP_DIR/ollama "https://ollama.com/download/ollama-linux-${ARCH}${VER_PARAM}"
 
 for BINDIR in /usr/local/bin /usr/bin /bin; do
     echo $PATH | grep -q $BINDIR && break || continue
@@ -82,7 +83,7 @@ trap install_success EXIT
 configure_systemd() {
     if ! id ollama >/dev/null 2>&1; then
         status "Creating ollama user..."
-        $SUDO useradd -r -s /bin/false -m -d /usr/share/ollama ollama
+        $SUDO useradd -r -s /bin/false -U -m -d /usr/share/ollama ollama
     fi
     if getent group render >/dev/null 2>&1; then
         status "Adding ollama user to render group..."
@@ -131,14 +132,23 @@ if available systemctl; then
 fi
 
 if ! available lspci && ! available lshw; then
-    warning "Unable to detect NVIDIA GPU. Install lspci or lshw to automatically detect and install NVIDIA CUDA drivers."
+    warning "Unable to detect NVIDIA/AMD GPU. Install lspci or lshw to automatically detect and install GPU dependencies."
     exit 0
 fi
 
 check_gpu() {
+    # Look for devices based on vendor ID for NVIDIA and AMD
     case $1 in
-        lspci) available lspci && lspci -d '10de:' | grep -q 'NVIDIA' || return 1 ;;
-        lshw) available lshw && $SUDO lshw -c display -numeric | grep -q 'vendor: .* \[10DE\]' || return 1 ;;
+        lspci) 
+            case $2 in
+                nvidia) available lspci && lspci -d '10de:' | grep -q 'NVIDIA' || return 1 ;;
+                amdgpu) available lspci && lspci -d '1002:' | grep -q 'AMD' || return 1 ;;
+            esac ;;
+        lshw) 
+            case $2 in
+                nvidia) available lshw && $SUDO lshw -c display -numeric | grep -q 'vendor: .* \[10DE\]' || return 1 ;;
+                amdgpu) available lshw && $SUDO lshw -c display -numeric | grep -q 'vendor: .* \[1002\]' || return 1 ;;
+            esac ;;
         nvidia-smi) available nvidia-smi || return 1 ;;
     esac
 }
@@ -148,9 +158,30 @@ if check_gpu nvidia-smi; then
     exit 0
 fi
 
-if ! check_gpu lspci && ! check_gpu lshw; then
+if ! check_gpu lspci nvidia && ! check_gpu lshw nvidia && ! check_gpu lspci amdgpu && ! check_gpu lshw amdgpu; then
     install_success
-    warning "No NVIDIA GPU detected. Ollama will run in CPU-only mode."
+    warning "No NVIDIA/AMD GPU detected. Ollama will run in CPU-only mode."
+    exit 0
+fi
+
+if check_gpu lspci amdgpu || check_gpu lshw amdgpu; then
+    # Look for pre-existing ROCm v6 before downloading the dependencies
+    for search in "${HIP_PATH:-''}" "${ROCM_PATH:-''}" "/opt/rocm"; do
+        if [ -n "${search}" ] && [ -e "${search}/lib/libhipblas.so.2" ]; then
+            status "Compatible AMD GPU ROCm library detected at ${search}"
+            install_success
+            exit 0
+        fi
+    done
+
+    status "Downloading AMD GPU dependencies..."
+    $SUDO rm -rf /usr/share/ollama/lib
+    $SUDO chmod o+x /usr/share/ollama
+    $SUDO install -o ollama -g ollama -m 755 -d /usr/share/ollama/lib/rocm
+    curl --fail --show-error --location --progress-bar "https://ollama.com/download/ollama-linux-amd64-rocm.tgz${VER_PARAM}" \
+        | $SUDO tar zx --owner ollama --group ollama -C /usr/share/ollama/lib/rocm .
+    install_success
+    status "AMD GPU dependencies installed."
     exit 0
 fi