diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md new file mode 100644 index 00000000..7048f48b --- /dev/null +++ b/docs/ggml_extensions.md @@ -0,0 +1,107 @@ +# ggml operator extensions + +The patched ggml exposes `ggml-ext.h` for application-owned operators, kernel +overrides, and graph fusion rules. Registrations are in-process and scoped to a +backend device. Register them before capability queries or graph allocation; +callback code and user data must remain alive until all dependent graphs and +backends have been destroyed. Registration names are unique, and duplicate +registrations fail rather than replacing live implementations. + +## Operator definitions + +`ggml_ext_register_op` registers a versioned name, attribute size, and an output +metadata inference callback. `ggml_ext_op_create` creates an ordinary ggml tensor +with explicit input dependencies. The initial interface supports one contiguous, +out-of-place output with an existing `ggml_type`. + +Attributes are copied into the unused portion of `op_params`. They are small, +plain values; use `memcpy` when reading attributes that need stricter alignment. +Auxiliary tensors must appear in `src`, not in an attribute pointer or user data. +Registered operators use `GGML_OP_CUSTOM` internally, with an identifiable +callback marker. Unsupported devices do not mistake them for CPU callbacks. +These process-local descriptors are not a graph serialization format. + +## Kernel selection + +`ggml_ext_register_kernel(device, &desc)` associates a kernel with a built-in +`ggml_op`, or with a registered custom operator. `supports` describes correctness +constraints; `preferred` optionally decides whether to override an available +native implementation. Higher priorities are considered first. When native +execution is unsupported, a supported extension can run regardless of its +preference callback. + +Capability queries include registered single-node implementations, so the +scheduler can select a device even when its native implementation rejects the +operator. Callbacks inspect metadata only and must not recursively call the +public capability query for the same device and operator. A fusion-only +registration does not claim single-node support. + +`prepare` creates optional host-side plan state and reports workspace bytes. +It can run during memory measurement, before tensor addresses are assigned, and +must not launch work or allocate device buffers. Workspace is allocated by ggml +before execution and included in scheduler memory estimates and buffer-size +reports. `execute` receives the backend, workspace buffer, and current tensor +addresses. `destroy` releases prepared state. A preparation or execution error +is propagated; execution never retries a native implementation after launch. + +Overriding a built-in operator preserves its numerical and aliasing contract, +including allocator reuse of input storage. Explicit views/in-place nodes are +not intercepted by this first interface. Rebuild the scheduler allocation after +changing graph metadata, selection policy, or registrations. + +## Fusion + +A `match` callback returns the number of consecutive nodes it can replace. +The scheduler freezes this selection before allocating tensors. ggml checks +compute flags, external uses, output flags, and views; the initial implementation +accepts a conservative set of pure operators with only the final node exposed +as the fused output. The rule remains responsible for mathematical equivalence, +dtype, shape, layout, and hardware constraints. + +External inputs stay allocated through the final node. Native graph optimization +and execution operate on the remaining contiguous ranges, so native fusion cannot +consume nodes selected by an extension. Intermediate allocations are retained +in this initial implementation; only their computation is removed. + +Automatic extension fusion is enabled in scheduler execution with one pipeline +copy and without evaluation callbacks. Multiple pipeline copies and evaluation +callbacks retain individual operations. Direct backend graph computation supports +single-node extensions; it does not fuse an already allocated graph. Direct calls +use temporary plans and synchronize before releasing them; scheduler plans and +workspaces are reused across executions. + +Use the backend graph APIs or the scheduler. The legacy CPU `ggml_graph_compute` +and backend `graph_plan_*` interfaces do not dispatch these extensions. + +## Backend services + +Services are discovered through `ggml_backend_reg_get_proc_address`: + +| Service | Purpose | +| --- | --- | +| `ggml_ext_cpu_parallel_v1` | Run an application callback on the CPU backend's configured thread pool, with its abort callback. | +| `ggml_ext_cuda_context_v1` | Borrow the current CUDA device and stream, ordered with native work. | +| `ggml_ext_cuda_device_v1` | Get the physical CUDA ordinal for capability checks before backend creation, including virtual ggml devices. | + +Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; callers +must use the representation appropriate to their backend. A Vulkan command +recording service is not implemented yet. + +## sd.cpp integration + +`src/core/ggml_extensions.cpp` registers the CPU implementation of +`sd.rope_apply.v1`, an override for existing `GGML_OP_ROPE_APPLY` nodes, and a +`RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight +vector. Both interleaved and split-half RoPE, partial rotation, and strided input +rows are supported. Model construction uses the custom operator on CPU and +retains existing backend implementations elsewhere. + +`SD_USE_UPSTREAM_GGML` excludes these registrations. The stable-diffusion public C +API is unchanged. INT8 ConvRot retains its existing packed activation layout and +separate weight scale tensor. + +CuTe Flash Attention still resides in ggml, including its existing fusion paths. +Migrating it requires moving its pool/conversion dependencies onto the extension +workspace and CUDA service, then registering its selection and fusion rules in +sd.cpp. This initial implementation establishes and validates the mechanism; +it does not yet externalize all model-specific kernels or add dynamic dtypes. diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 5747d625..fe6b48f1 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -6,6 +6,7 @@ #include "ggml-backend.h" #include "ggml.h" +#include "core/ggml_extensions.h" #define EPS 1e-05f diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index 8586d876..1ea86a8c 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -1,4 +1,5 @@ #include "core/ggml_extend_backend.h" +#include "core/ggml_extensions.h" #include #include @@ -596,7 +597,9 @@ bool sd_backend_supports_conv_implicit_gemm(ggml_backend_t backend, bool conv3d) ggml_backend_t sd_backend_cpu_init() { ggml_backend_load_all_once(); - return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + auto backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + sd_ggml_register_extensions(backend); + return backend; } bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) { @@ -1207,6 +1210,7 @@ ggml_backend_t SDBackendManager::init_cached_backend(const std::string& name) { } SDBackendHandle handle(backend); + sd_ggml_register_extensions(backend); backends_.emplace(actual_key, std::move(handle)); return backend; } diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp new file mode 100644 index 00000000..fa3d406d --- /dev/null +++ b/src/core/ggml_extensions.cpp @@ -0,0 +1,207 @@ +#include "core/ggml_extensions.h" + +#ifndef SD_USE_UPSTREAM_GGML +#include "ggml-ext.h" + +#include +#include +#include +#include + +static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) { + if (count != 2) { + return false; + } + const auto* x = inputs[0]; + const auto* pe = inputs[1]; + if (x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || x->ne[0] % 2 != 0 || + pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 || + pe->ne[3] != x->ne[2]) { + return false; + } + out->type = GGML_TYPE_F32; + out->ne[0] = x->ne[0]; + out->ne[1] = x->ne[2]; + out->ne[2] = x->ne[1] * x->ne[3]; + out->ne[3] = 1; + return true; +} + +static ggml_ext_op_t rope_op() { + static const ggml_ext_op_t op = [] { + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.rope_apply.v1", sizeof(int32_t), rope_infer}; + auto result = ggml_ext_register_op(&desc); + GGML_ASSERT(result); + return result; + }(); + return op; +} + +static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) { + if (node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node) || !node->src[0] || !node->src[1]) { + return false; + } + ggml_tensor* inputs[] = {node->src[0], node->src[1]}; + ggml_ext_tensor_desc output; + return rope_infer(inputs, 2, nullptr, &output) && + std::memcmp(output.ne, node->ne, sizeof(output.ne)) == 0; +} + +struct RopeExecution { + const ggml_tensor* rope; + const ggml_tensor* norm; + const ggml_tensor* weight; +}; + +static void rope_compute(ggml_tensor* dst, int ith, int nth, void* user_data) { + const auto& execution = *static_cast(user_data); + const auto* rope = execution.rope; + const auto* x = execution.norm ? execution.norm->src[0] : rope->src[0]; + const auto* pe = rope->src[1]; + int32_t interleaved; + std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved)); + float eps = 0.0f; + if (execution.norm) { + std::memcpy(&eps, execution.norm->op_params, sizeof(eps)); + } + const int64_t half_dim = pe->ne[2]; + const int64_t rows = x->ne[1] * x->ne[2] * x->ne[3]; + for (int64_t row = ith; row < rows; row += nth) { + const int64_t token = row % x->ne[2]; + const int64_t head = (row / x->ne[2]) % x->ne[1]; + const int64_t batch = row / (x->ne[2] * x->ne[1]); + const char* input = static_cast(x->data) + head * x->nb[1] + token * x->nb[2] + batch * x->nb[3]; + const char* pos = static_cast(pe->data) + token * pe->nb[3]; + float* output = static_cast(dst->data) + row * x->ne[0]; + float scale = 1.0f; + if (execution.norm) { + double sum = 0.0; + for (int64_t col = 0; col < x->ne[0]; ++col) { + const float v = *reinterpret_cast(input + col * x->nb[0]); + sum += static_cast(v * v); + } + const float mean = static_cast(sum / x->ne[0]); + scale = 1.0f / std::sqrt(mean + eps); + } + auto value = [&](int64_t col) { + float v = *reinterpret_cast(input + col * x->nb[0]); + if (execution.norm) { + const auto* weight = static_cast(execution.weight->data); + v = (v * scale) * *reinterpret_cast(weight + col * execution.weight->nb[0]); + } + return v; + }; + for (int64_t j = 0; j < half_dim; ++j) { + const int64_t i0 = interleaved ? 2 * j : j; + const int64_t i1 = interleaved ? 2 * j + 1 : j + half_dim; + const float a = value(i0); + const float b = value(i1); + const char* m = pos + j * pe->nb[2]; + const float c = *reinterpret_cast(m); + const float s = *reinterpret_cast(m + pe->nb[1]); + if (interleaved) { + const float neg_s = *reinterpret_cast(m + pe->nb[0]); + const float c1 = *reinterpret_cast(m + pe->nb[0] + pe->nb[1]); + output[i0] = a * c + b * neg_s; + output[i1] = a * s + b * c1; + } else { + output[i0] = a * c - b * s; + output[i1] = a * s + b * c; + } + } + for (int64_t col = 2 * half_dim; col < x->ne[0]; ++col) { + output[col] = value(col); + } + } +} + +static int rope_fusion_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + auto* view = const_cast(graph); + if (first + 2 >= ggml_graph_n_nodes(view)) { + return 0; + } + const auto* norm = ggml_graph_node(view, first); + const auto* mul = ggml_graph_node(view, first + 1); + const auto* rope = ggml_graph_node(view, first + 2); + if (norm->op != GGML_OP_RMS_NORM || norm->type != GGML_TYPE_F32 || !norm->src[0] || + norm->src[0]->type != GGML_TYPE_F32 || norm->src[0]->nb[0] != sizeof(float) || + mul->op != GGML_OP_MUL || mul->type != GGML_TYPE_F32 || rope->src[0] != mul || + (rope->op != GGML_OP_ROPE_APPLY && ggml_ext_op_get(rope) != rope_op()) || + !rope_supports(device, rope, nullptr) || !ggml_are_same_shape(norm, mul)) { + return 0; + } + const auto* weight = mul->src[0] == norm ? mul->src[1] : (mul->src[1] == norm ? mul->src[0] : nullptr); + if (!weight || weight->type != GGML_TYPE_F32 || weight->ne[0] != norm->ne[0] || + weight->ne[1] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1) { + return 0; + } + return 3; +} + +static ggml_status rope_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int count, void*, void*) { + auto device = ggml_backend_get_device(context->backend); + auto reg = ggml_backend_dev_backend_reg(device); + auto run = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")); + GGML_ASSERT(run); + RopeExecution execution = {nodes[count - 1], nullptr, nullptr}; + if (count == 3) { + execution.norm = nodes[0]; + execution.weight = nodes[1]->src[0] == nodes[0] ? nodes[1]->src[1] : nodes[1]->src[0]; + } + return run(context->backend, nodes[count - 1], rope_compute, &execution); +} + +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")) { + return; + } + static std::mutex mutex; + static std::unordered_set registered; + std::lock_guard lock(mutex); + if (!registered.insert(device).second) { + return; + } + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.rope_apply.cpu.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = rope_op(); + desc.priority = 100; + desc.supports = rope_supports; + desc.execute = rope_execute; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + + desc.name = "sd.rope_apply.cpu.override.v1"; + desc.op = GGML_OP_ROPE_APPLY; + desc.custom_op = nullptr; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + + desc.name = "sd.rms_norm_mul_rope.cpu.v1"; + desc.op = GGML_OP_RMS_NORM; + desc.priority = 200; + desc.supports = nullptr; + desc.match = rope_fusion_match; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); +} + +ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved) { + sd_ggml_register_extensions(backend); + if (backend) { + auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); + if (ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")) { + ggml_tensor* inputs[] = {x, pe}; + int32_t attributes = interleaved ? 1 : 0; + return ggml_ext_op_create(ctx, rope_op(), inputs, 2, &attributes); + } + } + return ggml_rope_apply(ctx, x, pe, interleaved); +} + +#else +void sd_ggml_register_extensions(ggml_backend_t) {} +#endif diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h new file mode 100644 index 00000000..acee6385 --- /dev/null +++ b/src/core/ggml_extensions.h @@ -0,0 +1,12 @@ +#ifndef __SD_CORE_GGML_EXTENSIONS_H__ +#define __SD_CORE_GGML_EXTENSIONS_H__ + +#include "ggml-backend.h" + +void sd_ggml_register_extensions(ggml_backend_t backend); + +#ifndef SD_USE_UPSTREAM_GGML +ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved); +#endif + +#endif diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index 63584bc5..b6fdd328 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -1077,8 +1077,8 @@ namespace Rope { if (backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 && x->ne[0] % 2 == 0 && pe->ne[0] == 2 && pe->ne[1] == 2 && pe->ne[2] == x->ne[0] / 2 && pe->ne[3] == x->ne[2]) { - auto out = ggml_rope_apply(ctx, x, pe, rope_interleaved); - if (ggml_backend_supports_op(backend, out)) { + auto out = sd_ggml_rope_apply(ctx, backend, x, pe, rope_interleaved); + if (out && ggml_backend_supports_op(backend, out)) { return out; } } diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 739ef77f..3dd720ca 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -180,8 +180,8 @@ namespace MiniMaxH3 { GGML_ASSERT(rot_dim <= x->ne[0]); #ifndef SD_USE_UPSTREAM_GGML if (runner_ctx->backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32) { - auto out = ggml_rope_apply(ctx, x, pe, false); - if (ggml_backend_supports_op(runner_ctx->backend, out)) { + auto out = sd_ggml_rope_apply(ctx, runner_ctx->backend, x, pe, false); + if (out && ggml_backend_supports_op(runner_ctx->backend, out)) { return out; } }