From f084d66314b7ff88bf66d94ffe438f0e22773713 Mon Sep 17 00:00:00 2001
From: Jhen-Jie Hong <developer@jhen.me>
Date: Sun, 19 Jan 2025 18:25:25 +0800
Subject: [PATCH] feat: build cuda bin for linux (x86/arm) (#77)

* ci: use new linux arm runner for build build-linux-arm64

* ci: install yarn by self

* ci: enable linux cuda build

* feat: add CMAKE_CUDA_ARCHITECTURES

* feat: update scripts/prepare-linux.sh

* feat: set GGML_NATIVE=OFF
---
 .github/workflows/build-release.yml |  6 ++++++
 lib/binding.ts                      |  2 +-
 scripts/build-linux.sh              | 21 +++++++++++++++++++++
 scripts/prepare-linux.sh            |  5 ++++-
 4 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/build-release.yml b/.github/workflows/build-release.yml
index 2d73546..44eecd6 100644
--- a/.github/workflows/build-release.yml
+++ b/.github/workflows/build-release.yml
@@ -20,7 +20,10 @@ on:
 jobs:
   build-linux-x86_64:
     runs-on: ubuntu-22.04
+    container: nvidia/cuda:12.6.2-devel-ubuntu22.04
     steps:
+      - name: Upgrade git
+        run: apt-get update && apt-get install -y git
       - uses: actions/checkout@v4
         with:
           submodules: "true"
@@ -55,7 +58,10 @@ jobs:
 
   build-linux-arm64:
     runs-on: ubuntu-22.04-arm
+    container: nvidia/cuda:12.6.2-devel-ubuntu22.04
     steps:
+      - name: Upgrade git
+        run: apt-get update && apt-get install -y git
       - uses: actions/checkout@v4
         with:
           submodules: "true"
diff --git a/lib/binding.ts b/lib/binding.ts
index 83e60fc..def14bf 100644
--- a/lib/binding.ts
+++ b/lib/binding.ts
@@ -117,7 +117,7 @@ export interface Module {
   LlamaContext: LlamaContext
 }
 
-export type LibVariant = 'default' | 'vulkan'
+export type LibVariant = 'default' | 'vulkan' | 'cuda'
 
 const setupEnv = (variant?: string) => {
   const postfix = variant ? `-${variant}` : ''
diff --git a/scripts/build-linux.sh b/scripts/build-linux.sh
index ccf4818..ec8c3c3 100755
--- a/scripts/build-linux.sh
+++ b/scripts/build-linux.sh
@@ -9,7 +9,28 @@ ARCH=${ARCH:-${1:-$(uname -m)}}
 if [ $ARCH == "x86_64" ]; then
   yarn clean && yarn build-native
   yarn clean && yarn build-native --CDLLAMA_VULKAN=1 --CDVARIANT=vulkan
+
+  # Check CUDA is available
+  if [ -f /usr/local/cuda/bin/nvcc ]; then
+    yarn clean && yarn build-native \
+      --CDLLAMA_CUDA=1 \
+      --CDVARIANT=cuda \
+      --CDCMAKE_CUDA_ARCHITECTURES=89 # > GeForce RTX 40 series
+  else
+    echo "CUDA is not available, skipping CUDA build"
+  fi
 else
   yarn clean && yarn build-native --CDGGML_NATIVE=OFF
   yarn clean && yarn build-native --CDGGML_NATIVE=OFF --CDLLAMA_VULKAN=1 --CDVULKAN_SDK="$(realpath 'externals/arm64-Vulkan-SDK')" --CDVARIANT=vulkan
+
+  # Check CUDA is available
+  if [ -f /usr/local/cuda/bin/nvcc ]; then
+    yarn clean && yarn build-native \
+      --CDLLAMA_CUDA=1 \
+      --CDVARIANT=cuda \
+      --CDGGML_NATIVE=OFF \
+      --CDCMAKE_CUDA_ARCHITECTURES=87 # > Jetson Orin series
+  else
+    echo "CUDA is not available, skipping CUDA build"
+  fi
 fi
diff --git a/scripts/prepare-linux.sh b/scripts/prepare-linux.sh
index 4702306..24f8e7d 100755
--- a/scripts/prepare-linux.sh
+++ b/scripts/prepare-linux.sh
@@ -14,12 +14,15 @@ export DEBIAN_FRONTEND=noninteractive
 
 ARCH=${ARCH:-${1:-$(uname -m)}}
 
+run_as_root apt-get update
+run_as_root apt-get install -qy lsb-release wget
+
 if [ $ARCH == "x86_64" ]; then
   DISTRO=$(lsb_release -c -s)
   wget -qO- https://packages.lunarg.com/lunarg-signing-key-pub.asc | run_as_root tee /etc/apt/trusted.gpg.d/lunarg.asc
   run_as_root wget -qO /etc/apt/sources.list.d/lunarg-vulkan-1.3.280-$DISTRO.list https://packages.lunarg.com/vulkan/1.3.280/lunarg-vulkan-1.3.280-$DISTRO.list
   run_as_root apt-get update
-  run_as_root apt-get install -qy vulkan-sdk
+  run_as_root apt-get install -qy vulkan-sdk cmake pkg-config build-essential  libx11-xcb-dev libxkbcommon-dev libwayland-dev libxrandr-dev
 else
   run_as_root apt-get update
   run_as_root apt-get install -qy curl gnupg2