refactor: align ggml backend implementation

thxCode · thxCode · commit 5489ad2995e7 · 2024-12-01T09:58:53.000+08:00
Signed-off-by: thxCode &lt;thxcode0824@gmail.com&gt;
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
@@ -165,7 +165,7 @@ jobs:
           - build: "cuda12"
             defines: "-DSD_CUBLAS=ON -DSD_BUILD_SHARED_LIBS=ON"
           # - build: "rocm5.5"
-          #   defines: '-G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS="gfx1100;gfx1102;gfx1030" -DSD_BUILD_SHARED_LIBS=ON'
+          #   defines: '-G Ninja -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIP=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS="gfx1100;gfx1102;gfx1030" -DSD_BUILD_SHARED_LIBS=ON'
           - build: 'vulkan'
             defines: "-DSD_VULKAN=ON -DSD_BUILD_SHARED_LIBS=ON"
     steps:
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -1,6 +1,10 @@
 cmake_minimum_required(VERSION 3.12)
 project("stable-diffusion")
 
+if (NOT TARGET ggml)
+    cmake_policy(SET CMP0077 NEW)
+endif()
+
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
 if (NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE)
@@ -24,19 +28,20 @@ endif()
 # general
 #option(SD_BUILD_TESTS                "sd: build tests"    ${SD_STANDALONE})
 option(SD_BUILD_EXAMPLES             "sd: build examples" ${SD_STANDALONE})
-option(SD_CUBLAS                     "sd: cuda backend" OFF)
-option(SD_HIPBLAS                    "sd: rocm backend" OFF)
+option(SD_CUDA                       "sd: cuda backend" OFF)
+option(SD_HIP                        "sd: rocm backend" OFF)
 option(SD_METAL                      "sd: metal backend" OFF)
 option(SD_VULKAN                     "sd: vulkan backend" OFF)
 option(SD_SYCL                       "sd: sycl backend" OFF)
-option(SD_FAST_SOFTMAX               "sd: x1.5 faster softmax, indeterministic (sometimes, same seed don't generate same image), cuda only" OFF)
+option(SD_CANN                       "sd: cann backend" OFF)
+option(SD_MUSA                       "sd: musa backend" OFF)
 option(SD_BUILD_SHARED_LIBS          "sd: build shared libs" OFF)
 #option(SD_BUILD_SERVER               "sd: build server example"                           ON)
 
-if(SD_CUBLAS)
-    message("-- Use CUBLAS as backend stable-diffusion")
+if(SD_CUDA)
+    message("-- Use CUDA as backend stable-diffusion")
     set(GGML_CUDA ON)
-    add_definitions(-DSD_USE_CUBLAS)
+    add_definitions(-DSD_USE_CUDA)
 endif()
 
 if(SD_METAL)
@@ -51,13 +56,22 @@ if (SD_VULKAN)
     add_definitions(-DSD_USE_VULKAN)
 endif ()
 
-if (SD_HIPBLAS)
-    message("-- Use HIPBLAS as backend stable-diffusion")
-    set(GGML_HIPBLAS ON)
-    add_definitions(-DSD_USE_CUBLAS)
-    if(SD_FAST_SOFTMAX)
-        set(GGML_CUDA_FAST_SOFTMAX ON)
-    endif()
+if (SD_HIP)
+    message("-- Use HIP as backend stable-diffusion")
+    set(GGML_HIP ON)
+    add_definitions(-DSD_USE_CUDA)
+endif ()
+
+if (SD_CANN)
+    message("-- Use CANN as backend stable-diffusion")
+    set(GGML_CANN ON)
+    add_definitions(-DSD_USE_CANN)
+endif ()
+
+if (SD_MUSA)
+    message("-- Use MUSA as backend stable-diffusion")
+    set(GGML_MUSA ON)
+    add_definitions(-DSD_USE_CUDA)
 endif ()
 
 set(SD_LIB stable-diffusion)
@@ -98,8 +112,6 @@ if(SD_SYCL)
     target_compile_options(${SD_LIB} PRIVATE ${SYCL_COMPILE_OPTIONS})
 endif()
 
-set(CMAKE_POLICY_DEFAULT_CMP0077 NEW)
-
 # see https://github.com/ggerganov/ggml/pull/682
 add_definitions(-DGGML_MAX_NAME=128)
 
diff --git a/README.md b/README.md
@@ -118,7 +118,7 @@ cmake --build . --config Release
 This provides BLAS acceleration using the CUDA cores of your Nvidia GPU. Make sure to have the CUDA toolkit installed. You can download it from your Linux distro's package manager (e.g. `apt install nvidia-cuda-toolkit`) or from here: [CUDA Toolkit](https://developer.nvidia.com/cuda-downloads). Recommended to have at least 4 GB of VRAM.
 
 ```
-cmake .. -DSD_CUBLAS=ON
+cmake .. -DSD_CUDA=ON
 cmake --build . --config Release
 ```
 
@@ -128,7 +128,7 @@ This provides BLAS acceleration using the ROCm cores of your AMD GPU. Make sure
 Windows User Refer to [docs/hipBLAS_on_Windows.md](docs%2FhipBLAS_on_Windows.md) for a comprehensive guide.
 
 ```
-cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS=gfx1100
+cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIP=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS=gfx1100
 cmake --build . --config Release
 ```
 
diff --git a/docs/hipBLAS_on_Windows.md b/docs/hipBLAS_on_Windows.md
@@ -45,7 +45,7 @@ set ninja=C:\Program Files\ninja\ninja.exe
 ```
 ## Building stable-diffusion.cpp
 
-The thing different from the regular CPU build is `-DSD_HIPBLAS=ON` ,
+The thing different from the regular CPU build is `-DSD_HIP=ON` ,
 `-G "Ninja"`, `-DCMAKE_C_COMPILER=clang`, `-DCMAKE_CXX_COMPILER=clang++`, `-DAMDGPU_TARGETS=gfx1100`
 
 >**Notice**: check the `clang` and `clang++` information:
@@ -78,7 +78,7 @@ option:
 ```commandline
 mkdir build
 cd build
-cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIPBLAS=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS=gfx1100
+cmake .. -G "Ninja" -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ -DSD_HIP=ON -DCMAKE_BUILD_TYPE=Release -DAMDGPU_TARGETS=gfx1100
 cmake --build . --config Release
 ```
 
diff --git a/ggml_extend.hpp b/ggml_extend.hpp
@@ -27,7 +27,7 @@
 
 #include "model.h"
 
-#ifdef SD_USE_CUBLAS
+#ifdef SD_USE_CUDA
 #include "ggml-cuda.h"
 #endif
 
@@ -43,6 +43,14 @@
 #include "ggml-sycl.h"
 #endif
 
+#ifdef SD_USE_CANN
+#include "ggml-cann.h"
+#endif
+
+#ifdef SD_USE_MUSA
+#include "ggml-musa.h"
+#endif
+
 #include "rng.hpp"
 #include "util.h"
 
@@ -672,7 +680,7 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_attention(struct ggml_context* ctx
                                                         struct ggml_tensor* k,
                                                         struct ggml_tensor* v,
                                                         bool mask = false) {
-#if defined(SD_USE_FLASH_ATTENTION) && !defined(SD_USE_CUBLAS) && !defined(SD_USE_METAL) && !defined(SD_USE_VULKAN) && !defined(SD_USE_SYCL)
+#if defined(SD_USE_FLASH_ATTENTION) && !defined(SD_USE_CUDA) && !defined(SD_USE_METAL) && !defined(SD_USE_VULKAN) && !defined(SD_USE_SYCL) && !defined(SD_USE_CANN) && !defined(SD_USE_MUSA)
     struct ggml_tensor* kqv = ggml_flash_attn(ctx, q, k, v, false);  // [N * n_head, n_token, d_head]
 #else
     float d_head           = (float)q->ne[0];
@@ -828,7 +836,7 @@ __STATIC_INLINE__ struct ggml_tensor* ggml_nn_group_norm(struct ggml_context* ct
 }
 
 __STATIC_INLINE__ void ggml_backend_tensor_get_and_sync(ggml_backend_t backend, const struct ggml_tensor* tensor, void* data, size_t offset, size_t size) {
-#if defined(SD_USE_CUBLAS) || defined(SD_USE_SYCL)
+#if defined(SD_USE_CUDA) || defined(SD_USE_SYCL)
     if (!ggml_backend_is_cpu(backend)) {
         ggml_backend_tensor_get_async(backend, tensor, data, offset, size);
         ggml_backend_synchronize(backend);
@@ -1138,11 +1146,6 @@ struct GGMLRunner {
             ggml_backend_cpu_set_n_threads(backend, n_threads);
         }
 
-#ifdef SD_USE_METAL
-        if (ggml_backend_is_metal(backend)) {
-            ggml_backend_metal_set_n_cb(backend, n_threads);
-        }
-#endif
         ggml_backend_graph_compute(backend, gf);
 
 #ifdef GGML_PERF
diff --git a/model.cpp b/model.cpp
@@ -26,6 +26,14 @@
 #include "ggml-vulkan.h"
 #endif
 
+#ifdef SD_USE_CANN
+#include "ggml-cann.h"
+#endif
+
+#ifdef SD_USE_MUSA
+#include "ggml-musa.h"
+#endif
+
 #define ST_HEADER_SIZE_LEN 8
 
 uint64_t read_u64(uint8_t* buffer) {
diff --git a/stable-diffusion.cpp b/stable-diffusion.cpp
@@ -155,27 +155,48 @@ class StableDiffusionGGML {
                         bool vae_on_cpu,
                         bool diffusion_flash_attn) {
         use_tiny_autoencoder = taesd_path.size() > 0;
-#ifdef SD_USE_CUBLAS
+#ifdef SD_USE_CUDA
+#ifdef SD_USE_HIP
+        LOG_DEBUG("Using HIP backend");
+#else
+#ifdef SD_USE_MUSA
+        LOG_DEBUG("Using MUSA backend");
+#else
         LOG_DEBUG("Using CUDA backend");
+#endif
+#endif
         backend = ggml_backend_cuda_init(0);
+        if (!backend) {
+            LOG_ERROR("CUDA backend init failed");
+        }
 #endif
 #ifdef SD_USE_METAL
         LOG_DEBUG("Using Metal backend");
-        ggml_backend_metal_log_set_callback(ggml_log_callback_default, nullptr);
         backend = ggml_backend_metal_init();
+        if (!backend) {
+            LOG_ERROR("Metal backend init failed");
+        }
 #endif
 #ifdef SD_USE_VULKAN
         LOG_DEBUG("Using Vulkan backend");
-        for (int device = 0; device < ggml_backend_vk_get_device_count(); ++device) {
-            backend = ggml_backend_vk_init(device);
-        }
+        backend = ggml_backend_vk_init(0);
         if (!backend) {
-            LOG_WARN("Failed to initialize Vulkan backend");
+            LOG_ERROR("Vulkan backend init failed");
         }
 #endif
 #ifdef SD_USE_SYCL
         LOG_DEBUG("Using SYCL backend");
         backend = ggml_backend_sycl_init(0);
+        if (!backend) {
+            LOG_ERROR("SYCL backend init failed");
+        }
+#endif
+#ifdef SD_USE_CANN
+        LOG_DEBUG("Using CANN backend");
+        backend = ggml_backend_cann_init(0);
+        if (!backend) {
+            LOG_ERROR("CANN backend init failed");
+        }
 #endif
 
         if (!backend) {
diff --git a/upscaler.cpp b/upscaler.cpp
@@ -15,22 +15,48 @@ struct UpscalerGGML {
     }
 
     bool load_from_file(const std::string& esrgan_path) {
-#ifdef SD_USE_CUBLAS
+#ifdef SD_USE_CUDA
+#ifdef SD_USE_HIP
+        LOG_DEBUG("Using HIP backend");
+#else
+#ifdef SD_USE_MUSA
+        LOG_DEBUG("Using MUSA backend");
+#else
         LOG_DEBUG("Using CUDA backend");
+#endif
+#endif
         backend = ggml_backend_cuda_init(0);
+        if (!backend) {
+            LOG_ERROR("CUDA backend init failed");
+        }
 #endif
 #ifdef SD_USE_METAL
         LOG_DEBUG("Using Metal backend");
-        ggml_backend_metal_log_set_callback(ggml_log_callback_default, nullptr);
         backend = ggml_backend_metal_init();
+        if (!backend) {
+            LOG_ERROR("Metal backend init failed");
+        }
 #endif
 #ifdef SD_USE_VULKAN
         LOG_DEBUG("Using Vulkan backend");
         backend = ggml_backend_vk_init(0);
+        if (!backend) {
+            LOG_ERROR("Vulkan backend init failed");
+        }
 #endif
 #ifdef SD_USE_SYCL
         LOG_DEBUG("Using SYCL backend");
         backend = ggml_backend_sycl_init(0);
+        if (!backend) {
+            LOG_ERROR("SYCL backend init failed");
+        }
+#endif
+#ifdef SD_USE_CANN
+        LOG_DEBUG("Using CANN backend");
+        backend = ggml_backend_cann_init(0);
+        if (!backend) {
+            LOG_ERROR("CANN backend init failed");
+        }
 #endif
         ModelLoader model_loader;
         if (!model_loader.init_from_file(esrgan_path)) {
@@ -41,6 +67,7 @@ struct UpscalerGGML {
             LOG_DEBUG("Using CPU backend");
             backend = ggml_backend_cpu_init();
         }
+
         LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type));
         esrgan_upscaler = std::make_shared<ESRGAN>(backend, model_loader.tensor_storages_types);
         if (!esrgan_upscaler->load_from_file(esrgan_path)) {
