From 946431d5b073ff2f620faabb182336895cd174dc Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Fri, 22 Dec 2023 12:17:37 -0800
Subject: [PATCH] build cuda and rocm

---
 .github/workflows/test.yaml | 82 +++++++++++++++++++++++++++++++++++--
 1 file changed, 79 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 36691b263..5ba4c6e86 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -28,8 +28,66 @@ jobs:
       - uses: actions/upload-artifact@v4
         with:
           name: ${{ matrix.os }}-${{ matrix.arch }}-libraries
-          path: |
-            llm/llama.cpp/build/**/lib/*
+          path: llm/llama.cpp/build/**/lib/*
+  generate-cuda:
+    strategy:
+      matrix:
+        cuda-version:
+          - '11.8.0'
+    runs-on: ubuntu-latest
+    container: nvidia/cuda:${{ matrix.cuda-version }}-devel-ubuntu20.04
+    steps:
+      - run: |
+          apt-get update && apt-get install -y git build-essential curl
+          curl -fsSL https://github.com/Kitware/CMake/releases/download/v3.28.1/cmake-3.28.1-linux-x86_64.tar.gz \
+            | tar -zx -C /usr --strip-components 1
+        env:
+          DEBIAN_FRONTEND: noninteractive
+      - uses: actions/checkout@v4
+      - uses: actions/setup-go@v4
+        with:
+          go-version: '1.21'
+          cache: true
+      - run: go get ./...
+      - run: |
+          git config --global --add safe.directory /__w/ollama/ollama
+          go generate -x ./...
+        env:
+          OLLAMA_SKIP_CPU_GENERATE: '1'
+      - uses: actions/upload-artifact@v4
+        with:
+          name: cuda-${{ matrix.cuda-version }}-libraries
+          path: llm/llama.cpp/build/**/lib/*
+  generate-rocm:
+    strategy:
+      matrix:
+        rocm-version:
+          - '5.7.1'
+          - '6.0'
+    runs-on: ubuntu-latest
+    container: rocm/dev-ubuntu-20.04:${{ matrix.rocm-version }}
+    steps:
+      - run: |
+          apt-get update && apt-get install -y git build-essential curl rocm-libs
+          curl -fsSL https://github.com/Kitware/CMake/releases/download/v3.28.1/cmake-3.28.1-linux-x86_64.tar.gz \
+            | tar -zx -C /usr --strip-components 1
+        env:
+          DEBIAN_FRONTEND: noninteractive
+      - uses: actions/checkout@v4
+      - uses: actions/setup-go@v4
+        with:
+          go-version: '1.21'
+          cache: true
+      - run: go get ./...
+      - run: |
+          git config --global --add safe.directory /__w/ollama/ollama
+          go generate -x ./...
+        env:
+          OLLAMA_SKIP_CPU_GENERATE: '1'
+      - uses: actions/upload-artifact@v4
+        with:
+          name: rocm-${{ matrix.rocm-version }}-libraries
+          path: llm/llama.cpp/build/**/lib/*
   lint:
     strategy:
       matrix:
@@ -69,7 +127,10 @@ jobs:
         if: ${{ startsWith(matrix.os, 'windows-') }}
       - uses: golangci/golangci-lint-action@v3
   test:
-    needs: generate
+    needs:
+      - generate
+      - generate-cuda
+      - generate-rocm
     strategy:
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
@@ -96,5 +157,20 @@ jobs:
         with:
           name: ${{ matrix.os }}-${{ matrix.arch }}-libraries
           path: llm/llama.cpp/build
+      - if: ${{ matrix.os == 'ubuntu-latest' }}
+        uses: actions/download-artifact@v4
+        with:
+          name: cuda-11.8.0-libraries
+          path: llm/llama.cpp/build
+      - if: ${{ matrix.os == 'ubuntu-latest' }}
+        uses: actions/download-artifact@v4
+        with:
+          name: rocm-5.7.1-libraries
+          path: llm/llama.cpp/build
+      - if: ${{ matrix.os == 'ubuntu-latest' }}
+        uses: actions/download-artifact@v4
+        with:
+          name: rocm-6.0-libraries
+          path: llm/llama.cpp/build
       - run: go build
       - run: go test -v ./...