From b67297576df5dba0734c768ca49dcc3ec572a637 Mon Sep 17 00:00:00 2001
From: dixyes <dixyes@gmail.com>
Date: Tue, 16 Apr 2024 18:48:10 +0800
Subject: [PATCH] Add MUSA support

---
 CMakeLists.txt | 25 +++++++++++++++++++++++++
 ggml           |  2 +-
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index da0395e..884aca3 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -42,6 +42,7 @@ option(RWKV_OPENBLAS               "rwkv: use OpenBLAS"
 option(RWKV_CUBLAS                 "rwkv: use cuBLAS"                                     OFF)
 option(RWKV_CLBLAST                "rwkv: use CLBlast"                                    OFF)
 option(RWKV_HIPBLAS                "rwkv: use hipBLAS"                                    OFF)
+option(RWKV_MUSA                   "rwkv: use MUSA"                                       OFF)
 
 # Build only shared library without building tests and extras
 option(RWKV_STANDALONE             "rwkv: build only RWKV library"                        OFF)
@@ -232,6 +233,30 @@ if (RWKV_HIPBLAS)
     endif()
 endif()
 
+if (RWKV_MUSA)
+    option(MUSA_ARCH "MUSA architecture" "21")
+
+    list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}/ggml/cmake")
+
+    find_package(MUSA     REQUIRED)
+
+    message(STATUS "MUSA found")
+
+    enable_language(MUSA)
+
+    add_compile_definitions(GGML_USE_MUSA GGML_USE_CUDA GGML_USE_CUBLAS)
+
+    add_library(ggml-musa OBJECT
+                ${CMAKE_SOURCE_DIR}/ggml/src/ggml-cuda.cu
+                ${CMAKE_SOURCE_DIR}/ggml/src/ggml-cuda.h)
+
+    target_include_directories(ggml-musa PUBLIC ${CMAKE_SOURCE_DIR}/ggml/include/ggml)
+    set_source_files_properties(${CMAKE_SOURCE_DIR}/ggml/src/ggml-cuda.cu PROPERTIES LANGUAGE MUSA)
+    target_link_libraries(ggml-musa PRIVATE MUSA::musa MUSA::mublas MUSA::musart)
+
+    set(RWKV_EXTRA_LIBS ${RWKV_EXTRA_LIBS} ggml-musa)
+endif()
+
 if (RWKV_ALL_WARNINGS)
     if (NOT MSVC)
         set(c_flags
diff --git a/ggml b/ggml
index 46f083d..355e45a 160000
--- a/ggml
+++ b/ggml
@@ -1 +1 @@
-Subproject commit 46f083d15bb31c62933300ffbfffa5aa6ae2ecae
+Subproject commit 355e45af9e301c58e6152f8b03235a1cf6fd3fad