Implement basic CLBlast support

.github/workflows/build.yml

@@ -126,7 +126,6 @@ jobs:
       - name: Dependencies
         id: depends
         run: |
-          brew update
           brew install zip
 
       - name: Build
@@ -135,7 +134,7 @@ jobs:
         run: |
           mkdir build
           cd build
-          cmake -DRWKV_AVX2=OFF -DRWKV_FMA=OFF -DRWKV_SANITIZE_ADDRESS=ON ..
+          cmake -DRWKV_AVX2=OFF -DRWKV_FMA=OFF ..
           cmake --build . --config Release
 
       - name: Test

CMakeLists.txt

@@ -40,6 +40,7 @@ option(RWKV_FMA                    "rwkv: enable FMA"
 option(RWKV_ACCELERATE             "rwkv: enable Accelerate framework"                    ON)
 option(RWKV_OPENBLAS               "rwkv: use OpenBLAS"                                   OFF)
 option(RWKV_CUBLAS                 "rwkv: use cuBLAS"                                     OFF)
+option(RWKV_CLBLAST                "rwkv: use CLBlast"                                    OFF)
 
 # Build only shared library without building tests and extras
 option(RWKV_STANDALONE             "rwkv: build only RWKV library"                        OFF)
@@ -125,6 +126,52 @@ if (RWKV_CUBLAS)
     endif()
 endif()
 
+if (RWKV_CLBLAST)
+    cmake_minimum_required(VERSION 3.17)
+
+    file(GLOB OPENCL_INCLUDE_SEARCH_PATHS ${CMAKE_SOURCE_DIR}/OpenCL-*/)
+    file(GLOB CLBLAST_INCLUDE_SEARCH_PATHS ${CMAKE_SOURCE_DIR}/CLBlast-*/)
+
+    set(OPENCL_INCLUDE_SEARCH_PATHS
+        /usr/include
+        /usr/local/include
+        $ENV{OPENCL_HOME}
+        $ENV{OPENCL_HOME}/include
+        ${OPENCL_INCLUDE_SEARCH_PATHS}
+        )
+
+    set(CLBLAST_INCLUDE_SEARCH_PATHS
+        /usr/include
+        /usr/local/include
+        $ENV{CLBLAST_HOME}
+        $ENV{CLBLAST_HOME}/include
+        ${CLBLAST_INCLUDE_SEARCH_PATHS}
+        )
+
+    find_path(OPENCL_INC NAMES opencl.h PATHS ${OPENCL_INCLUDE_SEARCH_PATHS} PATH_SUFFIXES include/CL)
+    find_library(OPENCL_LIB NAMES OpenCL PATHS ${OPENCL_INCLUDE_SEARCH_PATHS} PATH_SUFFIXES lib)
+    find_path(CLBLAST_INC NAMES clblast.h PATHS ${CLBLAST_INCLUDE_SEARCH_PATHS} PATH_SUFFIXES include)
+    find_library(CLBLAST_LIB NAMES clblast PATHS ${CLBLAST_INCLUDE_SEARCH_PATHS} PATH_SUFFIXES lib)
+
+    if (OPENCL_LIB)
+        set(OPENCL_INC ${OPENCL_INC}/..) # disgusting
+        message(STATUS "OpenCL SDK found: ${OPENCL_INC}")
+
+        if (CLBLAST_LIB)
+            message(STATUS "CLBlast found: ${CLBLAST_INC}")
+            add_compile_definitions(GGML_USE_CLBLAST)
+            set(GGML_OPENCL_SOURCES ${CMAKE_SOURCE_DIR}/ggml/src/ggml-opencl.cpp ${CMAKE_SOURCE_DIR}/ggml/src/ggml-opencl.h)
+            set(GGML_OPENCL_DIRS ${GGML_OPENCL_DIRS} ${OPENCL_INC} ${CLBLAST_INC})
+            set(RWKV_EXTRA_LIBS ${RWKV_EXTRA_LIBS} ${OPENCL_LIB} ${CLBLAST_LIB})
+            link_libraries("-Wl,--copy-dt-needed-entries")
+        else()
+            message(WARNING "CLBlast not found")
+        endif()
+    else()
+       message(WARNING "OpenCL SDK not found, CLBlast cannot be enabled")
+    endif()
+endif()
+
 if (RWKV_ALL_WARNINGS)
     if (NOT MSVC)
         set(c_flags
@@ -250,9 +297,10 @@ endif()
 add_library(ggml OBJECT
             ${CMAKE_SOURCE_DIR}/ggml/src/ggml.c
             ${CMAKE_SOURCE_DIR}/ggml/include/ggml/ggml.h
-            ${GGML_CUDA_SOURCES})
+            ${GGML_CUDA_SOURCES}
+            ${GGML_OPENCL_SOURCES})
 
-target_include_directories(ggml PUBLIC ${CMAKE_SOURCE_DIR}/ggml/include/ggml)
+target_include_directories(ggml PUBLIC ${CMAKE_SOURCE_DIR}/ggml/include/ggml ${GGML_OPENCL_DIRS})
 target_compile_features(ggml PUBLIC c_std_11) # Don't bump
 
 if (MSVC)

rwkv.cpp

@@ -3,6 +3,8 @@
 
 #ifdef GGML_USE_CUBLAS
 #include "ggml/src/ggml-cuda.h"
+#elif defined(GGML_USE_CLBLAST)
+#include "ggml/src/ggml-opencl.h"
 #endif
 
 #include <string>
@@ -1544,11 +1546,15 @@ struct rwkv_context * rwkv_clone_context(struct rwkv_context * ctx, const uint32
 }
 
 bool rwkv_gpu_offload_layers(struct rwkv_context * ctx, const uint32_t n_layers) {
-#ifdef GGML_USE_CUBLAS
+#if defined(GGML_USE_CUBLAS) || defined(GGML_USE_CLBLAST)
     const auto offload = [&](struct ggml_tensor * tensor) {
         // TODO support multi-GPU
         tensor->backend = GGML_BACKEND_GPU;
+#ifdef GGML_USE_CUBLAS
         ggml_cuda_transform_tensor(tensor->data, tensor);
+#elif defined(GGML_USE_CLBLAST)
+        ggml_cl_transform_tensor(tensor->data, tensor);
+#endif
     };
 
     const size_t n_gpu = std::min(n_layers, ctx->instance->model.header.n_layer);