feat: add custom modulate operator and ConvRot fusion (#2258)

src/core/cpu/modulate.cpp

@@ -0,0 +1,82 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <cstring>
+
+namespace sd_cpu {
+    static bool modulate_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        return sd_ggml_modulate_describe(node);
+    }
+
+    static float modulate_value(const ggml_tensor* t, const int64_t* index) {
+        size_t offset = 0;
+        for (int d = 0; d < GGML_MAX_DIMS; ++d) {
+            offset += (index[d] % t->ne[d]) * t->nb[d];
+        }
+        return *reinterpret_cast<const float*>(static_cast<const char*>(t->data) + offset);
+    }
+
+    static void modulate_compute(ggml_tensor* dst, int ith, int nth, void*) {
+        auto* output       = static_cast<float*>(dst->data);
+        const int64_t rows = ggml_nrows(dst);
+        for (int64_t row = ith; row < rows; row += nth) {
+            int64_t index[] = {0, row % dst->ne[1], (row / dst->ne[1]) % dst->ne[2], row / (dst->ne[1] * dst->ne[2])};
+            for (int64_t col = 0; col < dst->ne[0]; ++col) {
+                index[0]      = col;
+                const float x = modulate_value(dst->src[0], index);
+                // Match separate ggml MUL/ADD rounding even with contraction enabled.
+                volatile float product         = x * modulate_value(dst->src[1], index);
+                volatile float scaled          = x + product;
+                output[row * dst->ne[0] + col] = dst->src[2] ? scaled + modulate_value(dst->src[2], index) : scaled;
+            }
+        }
+    }
+
+    static ggml_status modulate_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        *bytes = sd_ggml_modulate_workspace_size(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+
+    static void modulate_destroy(void*, void*) {}
+
+    static ggml_status modulate_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+        auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(context->backend));
+        auto run = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+        GGML_ASSERT(run);
+        auto* dst          = nodes[0];
+        const auto address = reinterpret_cast<uintptr_t>(dst->data);
+        bool alias         = false;
+        for (const auto* input : dst->src) {
+            if (input) {
+                const auto src = reinterpret_cast<uintptr_t>(input->data);
+                alias |= address < src + ggml_nbytes(input) && src < address + ggml_nbytes(dst);
+            }
+        }
+        if (!alias) {
+            return run(context->backend, dst, modulate_compute, nullptr);
+        }
+        ggml_tensor temporary = *dst;
+        temporary.data        = ggml_backend_buffer_get_base(context->workspace);
+        const auto status     = run(context->backend, &temporary, modulate_compute, nullptr);
+        if (status == GGML_STATUS_SUCCESS) {
+            std::memcpy(dst->data, temporary.data, ggml_nbytes(dst));
+        }
+        return status;
+    }
+
+    void register_modulate(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.modulate.cpu.v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_modulate_op();
+        desc.priority             = 100;
+        desc.supports             = modulate_supports;
+        desc.prepare              = modulate_prepare;
+        desc.destroy              = modulate_destroy;
+        desc.execute              = modulate_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+    }
+}
+#endif

src/core/cpu/registration.cpp

@@ -6,6 +6,7 @@ namespace sd_cpu {
     void register_extensions(ggml_backend_dev_t device) {
         register_i8(device);
         register_rope(device);
+        register_modulate(device);
     }
 
 }

src/core/cpu/registration.h

@@ -7,6 +7,7 @@
 namespace sd_cpu {
     void register_extensions(ggml_backend_dev_t device);
     void register_rope(ggml_backend_dev_t device);
+    void register_modulate(ggml_backend_dev_t device);
     void register_i8(ggml_backend_dev_t device);
 }
 #endif