From d2075c78517535ff764b3138b2b5457d6657b8ff Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=97=BA=E6=97=BA=E7=A2=8E=E5=86=B0=E5=86=B0?=
 <38837039+Cyberhan123@users.noreply.github.com>
Date: Sat, 2 Dec 2023 15:21:01 +0800
Subject: [PATCH] Support build cuda rocm (#155)

* add cuda and hip support

* fix ci fail

* add ninja

* add ninja

* peek ninja

* peek ninja

* peek ninja

* peek ninja

* peek ninja

* Update build.yml

* test build with rocm

* change set ninja action

* fix rocm build fail build.yml
---
 .github/workflows/build.yml | 28 +++++++++++++++++++++++++++-
 1 file changed, 27 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 04208ee..19dab1e 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -183,6 +183,10 @@ jobs:
            defines: '-DRWKV_AVX2=OFF'
          - build: 'avx512'
            defines: '-DRWKV_AVX512=ON'
+         - build: 'cuda12'
+           defines: '-DRWKV_CUBLAS=ON'
+         - build: 'rocm5.5'
+           defines: '-G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DRWKV_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS="gfx1100;gfx1102;gfx1030"'
 
     steps:
       - name: Clone
@@ -191,6 +195,28 @@ jobs:
         with:
           submodules: 'recursive'
 
+      - name: Install cuda-toolkit
+        id: cuda_toolkit
+        if: ${{ matrix.build == 'cuda12' }}
+        uses: Jimver/cuda-toolkit@v0.2.11
+        with:
+          cuda: '12.2.0'
+          method: 'local'
+
+      - name: Install rocm-toolkit
+        id: rocm-toolkit
+        if: ${{ matrix.build == 'rocm5.5' }}
+        uses: Cyberhan123/rocm-toolkit@v0.1.0
+        with:
+          rocm: '5.5.0'
+
+      - name: Install Ninja
+        id: install-ninja
+        if: ${{ matrix.build == 'rocm5.5' }}
+        uses: urkle/action-get-ninja@v1
+        with:
+          version: 1.11.1
+
       - name: Build
         id: cmake_build
         run: |
@@ -215,7 +241,7 @@ jobs:
       - name: Test
         id: cmake_test
         # Test AVX-512 only when possible
-        if: ${{ matrix.build != 'avx512' || env.HAS_AVX512F == '1' }}
+        if: ${{ (matrix.build != 'avx512' || env.HAS_AVX512F == '1') && matrix.build != 'cuda12' && matrix.build != 'rocm5.5'}}
         run: |
           cd build
           ctest -C Release --verbose