mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 11:56:36 +00:00
refactor: register CuTe attention and fusion in sd.cpp
This commit is contained in:
1 parent
c6a4666289
commit
4b274821e9
5 files changed
+946
-10
No files matched your search
@@ -8,6 +8,10 @@
|
||||
#include <mutex>
|
||||
#include <unordered_set>
|
||||
|
||||
#ifdef SD_CUDA_EXTENSIONS
|
||||
void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
|
||||
#endif
|
||||
|
||||
static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) {
|
||||
if (count != 2) {
|
||||
return false;
|
||||
@@ -155,9 +159,11 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
|
||||
if (!backend) {
|
||||
return;
|
||||
}
|
||||
auto device = ggml_backend_get_device(backend);
|
||||
auto reg = ggml_backend_dev_backend_reg(device);
|
||||
if (!ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")) {
|
||||
auto device = ggml_backend_get_device(backend);
|
||||
auto reg = ggml_backend_dev_backend_reg(device);
|
||||
const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr;
|
||||
const bool cuda = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1") != nullptr;
|
||||
if (!cpu && !cuda) {
|
||||
return;
|
||||
}
|
||||
static std::mutex mutex;
|
||||
@@ -166,6 +172,12 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
|
||||
if (!registered.insert(device).second) {
|
||||
return;
|
||||
}
|
||||
if (cuda) {
|
||||
#ifdef SD_CUDA_EXTENSIONS
|
||||
sd_ggml_register_cuda_attention(device);
|
||||
#endif
|
||||
return;
|
||||
}
|
||||
ggml_ext_kernel_desc desc = {};
|
||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||
desc.name = "sd.rope_apply.cpu.v1";
|
||||
|
||||
Reference in new issue
Block a user