refactor: separate CPU kernels from shared operator extensions

CMakeLists.txt

@@ -225,6 +225,8 @@ file(GLOB SD_LIB_SOURCES CONFIGURE_DEPENDS
     "src/core/*.h"
     "src/core/*.cpp"
     "src/core/*.hpp"
+    "src/core/cpu/*.h"
+    "src/core/cpu/*.cpp"
     "src/core/probe/*.h"
     "src/core/probe/*.cpp"
     "src/extensions/*.h"

docs/ggml_extensions.md

@@ -129,7 +129,13 @@ the command recording services described below.
 
 ## sd.cpp integration
 
-`src/core/ggml_extensions.cpp` registers the CPU implementation of
+`src/core/ggml_extensions.cpp` and `.h` contain the shared operator schemas,
+validation, graph construction helpers, and backend registration entry point.
+CPU kernels and fusion rules live in `src/core/cpu/`, with registration collected
+in `cpu/registration.cpp`. CUDA-style and Vulkan implementations live in their
+respective backend directories.
+
+`src/core/cpu/rope.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
@@ -180,10 +186,11 @@ activation layout, and separate weight scales. External consumers prevent fusion
 aliasing terminal copies use scheduler workspace. The CUDA projection kernels
 and matchers now live in sd.cpp. `GGML_CUDA_FORCE_CUBLAS` disables these epilogues.
 
-`src/core/int8_ops.cpp` defines `sd.linear_i8.v1` and
+`src/core/ggml_extensions.cpp` defines `sd.linear_i8.v1` and
-`sd.quantize_i8_convrot.v1`, with CPU implementations and overrides of the old
+`sd.quantize_i8_convrot.v1`. `src/core/cpu/int8.cpp` supplies their CPU
-builtins. The CPU kernels use the configured thread pool and scheduler workspace;
+implementations and overrides of the old builtins. The CPU kernels use the
-the old CPU INT8 GEMM and ConvRot implementations have been removed from ggml.
+configured thread pool and scheduler workspace; the old CPU INT8 GEMM and ConvRot
+implementations have been removed from ggml.
 Custom nodes execute on CPU or CUDA when supported. Model code retains builtin
 nodes while the existing CUDA and Vulkan fusion rules are being migrated.
 

src/core/cpu/int8.cpp

@@ -0,0 +1,189 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+
+namespace sd_cpu {
+
+    static void transform(float* values, int group) {
+        const float scale = 1.f / std::sqrt(float(group));
+        for (int i = 0; i < group; ++i)
+            values[i] *= scale;
+        for (int stride = 1; stride < group; stride *= 4)
+            for (int base = 0; base < group; base += 4 * stride)
+                for (int j = 0; j < stride; ++j) {
+                    const int i   = base + j;
+                    const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride];
+                    values[i]              = a + b + c - d;
+                    values[i + stride]     = a + b - c + d;
+                    values[i + 2 * stride] = a - b + c + d;
+                    values[i + 3 * stride] = -a + b + c + d;
+                }
+    }
+    struct I8Execution {
+        SDI8Operation op;
+        int64_t rows;
+        int8_t* quant;
+        float* scales;
+        float* output;
+    };
+    static void quantize(ggml_tensor*, int ith, int nth, void* data) {
+        auto& e           = *static_cast<I8Execution*>(data);
+        const int64_t k   = e.op.input->ne[0];
+        const int group   = e.op.group_size;
+        const auto* input = static_cast<const float*>(e.op.input->data);
+        for (int64_t row = ith; row < e.rows; row += nth) {
+            const float* src = input + row * k;
+            int8_t* dst      = e.quant + row * k;
+            float amax       = 0.f, values[256];
+            auto visit       = [&](auto fn) { if (!group) { for (int64_t i = 0; i < k; ++i) fn(i, src[i]); }
+            else for (int64_t base = 0; base < k; base += group) { std::memcpy(values, src + base, group * sizeof(float)); transform(values, group); for (int j = 0; j < group; ++j) fn(base + j, values[j]); } };
+            visit([&](int64_t, float value) { amax = std::max(amax, std::fabs(value)); });
+            const float scale = amax / 127.f, inverse = scale == 0.f ? 0.f : 1.f / scale;
+            e.scales[row] = scale;
+            if (scale == 0.f)
+                std::memset(dst, 0, k);
+            else
+                visit([&](int64_t i, float value) { dst[i] = static_cast<int8_t>(std::max(-127L, std::min(127L, std::lrintf(value * inverse)))); });
+        }
+    }
+    static void multiply(ggml_tensor*, int ith, int nth, void* data) {
+        auto& e         = *static_cast<I8Execution*>(data);
+        const auto& op  = e.op;
+        const int64_t k = op.weight->ne[0], n = op.weight->ne[1];
+        const auto* w    = static_cast<const int8_t*>(op.weight->data);
+        const auto* ws   = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
+        const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+        for (int64_t index = ith; index < e.rows * n; index += nth) {
+            const int64_t row = index / n, column = index % n;
+            uint32_t sum = 0;
+            for (int64_t i = 0; i < k; ++i)
+                sum += uint32_t(int32_t(w[column * k + i]) * int32_t(e.quant[row * k + i]));
+            int32_t signed_sum;
+            std::memcpy(&signed_sum, &sum, sizeof(sum));
+            float value = float(signed_sum) * e.scales[row];
+            if (ws)
+                value *= ws[column];
+            if (bias)
+                value += bias[column];
+            e.output[index] = value;
+        }
+    }
+    static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
+        if (!src)
+            return false;
+        const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(src->data);
+        return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src);
+    }
+    static size_t quant_workspace(const ggml_tensor* node, const SDI8Operation& op) {
+        return op.input->type == GGML_TYPE_F32 ? (size_t(ggml_nelements(op.input)) + 3) / 4 * 4 + size_t(ggml_nrows(node)) * 4 : 0;
+    }
+    static bool i8_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        SDI8Operation op;
+        return sd_ggml_i8_describe(node, &op);
+    }
+    static bool convrot_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        int group;
+        return sd_ggml_convrot_describe(node, &group);
+    }
+    static ggml_status i8_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        SDI8Operation op;
+        if (!sd_ggml_i8_describe(nodes[0], &op))
+            return GGML_STATUS_FAILED;
+        *bytes = quant_workspace(nodes[0], op) + ggml_nbytes(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static ggml_status convrot_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        *bytes = ggml_nbytes(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static void destroy(void*, void*) {}
+    static ggml_ext_cpu_parallel_fn parallel(ggml_backend_t backend) {
+        return reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)), "ggml_ext_cpu_parallel_v1"));
+    }
+    static ggml_status i8_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+        auto* dst     = nodes[0];
+        I8Execution e = {};
+        if (!sd_ggml_i8_describe(dst, &e.op))
+            return GGML_STATUS_FAILED;
+        e.rows   = ggml_nrows(dst);
+        auto run = parallel(context->backend);
+        if (!run)
+            return GGML_STATUS_FAILED;
+        char* scratch           = static_cast<char*>(ggml_backend_buffer_get_base(context->workspace));
+        const size_t quant_size = quant_workspace(dst, e.op);
+        if (e.op.input->type == GGML_TYPE_F32) {
+            e.quant     = reinterpret_cast<int8_t*>(scratch);
+            e.scales    = reinterpret_cast<float*>(scratch + (size_t(ggml_nelements(e.op.input)) + 3) / 4 * 4);
+            auto status = run(context->backend, dst, quantize, &e);
+            if (status != GGML_STATUS_SUCCESS)
+                return status;
+        } else {
+            e.quant  = static_cast<int8_t*>(e.op.input->data);
+            e.scales = e.op.packed ? reinterpret_cast<float*>(e.quant + e.op.input->ne[0] * ((e.rows + 3) / 4 * 4)) : static_cast<float*>(e.op.input_scale->data);
+        }
+        bool alias = false;
+        for (auto* input : dst->src)
+            alias |= overlaps(dst, input);
+        e.output    = alias ? reinterpret_cast<float*>(scratch + quant_size) : static_cast<float*>(dst->data);
+        auto status = run(context->backend, dst, multiply, &e);
+        if (status == GGML_STATUS_SUCCESS && alias)
+            std::memcpy(dst->data, e.output, ggml_nbytes(dst));
+        return status;
+    }
+    static ggml_status convrot_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+        auto* dst     = nodes[0];
+        I8Execution e = {};
+        if (!sd_ggml_convrot_describe(dst, &e.op.group_size))
+            return GGML_STATUS_FAILED;
+        auto run = parallel(context->backend);
+        if (!run)
+            return GGML_STATUS_FAILED;
+        e.op.input      = dst->src[0];
+        e.rows          = ggml_nrows(e.op.input);
+        bool alias      = overlaps(dst, e.op.input);
+        e.quant         = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(context->workspace) : dst->data);
+        const int64_t k = dst->ne[0], padded = (e.rows + 3) / 4 * 4;
+        e.scales    = reinterpret_cast<float*>(e.quant + k * padded);
+        auto status = run(context->backend, dst, quantize, &e);
+        if (status == GGML_STATUS_SUCCESS) {
+            std::memset(e.quant + k * e.rows, 0, k * (padded - e.rows));
+            if (alias)
+                std::memcpy(dst->data, e.quant, ggml_nbytes(dst));
+        }
+        return status;
+    }
+    void register_i8(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc d = {};
+        d.abi_version          = GGML_EXT_ABI_VERSION;
+        d.priority             = 100;
+        d.destroy              = destroy;
+        d.name                 = "sd.linear_i8.cpu.v1";
+        d.op                   = GGML_OP_CUSTOM;
+        d.custom_op            = sd_ggml_i8_op();
+        d.supports             = i8_supports;
+        d.prepare              = i8_prepare;
+        d.execute              = i8_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+        d.name      = "sd.linear_i8.cpu.override.v1";
+        d.op        = GGML_OP_MUL_MAT;
+        d.custom_op = nullptr;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+        d.name      = "sd.quantize_i8_convrot.cpu.v1";
+        d.op        = GGML_OP_CUSTOM;
+        d.custom_op = sd_ggml_convrot_op();
+        d.supports  = convrot_supports;
+        d.prepare   = convrot_prepare;
+        d.execute   = convrot_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+        d.name      = "sd.quantize_i8_convrot.cpu.override.v1";
+        d.op        = GGML_OP_QUANTIZE_I8_CONVROT;
+        d.custom_op = nullptr;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    }
+
+}
+#endif

src/core/cpu/registration.cpp

@@ -0,0 +1,12 @@
+#include "core/cpu/registration.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+namespace sd_cpu {
+
+    void register_extensions(ggml_backend_dev_t device) {
+        register_i8(device);
+        register_rope(device);
+    }
+
+}
+#endif

src/core/cpu/registration.h

@@ -0,0 +1,14 @@
+#ifndef __SD_CORE_CPU_REGISTRATION_H__
+#define __SD_CORE_CPU_REGISTRATION_H__
+
+#include "ggml-backend.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+namespace sd_cpu {
+    void register_extensions(ggml_backend_dev_t device);
+    void register_rope(ggml_backend_dev_t device);
+    void register_i8(ggml_backend_dev_t device);
+}
+#endif
+
+#endif

src/core/cpu/rope.cpp

@@ -0,0 +1,143 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <cmath>
+#include <cstring>
+
+namespace sd_cpu {
+
+    static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        return sd_ggml_rope_describe(node);
+    }
+
+    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<const RopeExecution*>(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<const char*>(x->data) + head * x->nb[1] + token * x->nb[2] + batch * x->nb[3];
+            const char* pos     = static_cast<const char*>(pe->data) + token * pe->nb[3];
+            float* output       = static_cast<float*>(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<const float*>(input + col * x->nb[0]);
+                    sum += static_cast<double>(v * v);
+                }
+                const float mean = static_cast<float>(sum / x->ne[0]);
+                scale            = 1.0f / std::sqrt(mean + eps);
+            }
+            auto value = [&](int64_t col) {
+                float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
+                if (execution.norm) {
+                    const auto* weight = static_cast<const char*>(execution.weight->data);
+                    v                  = (v * scale) * *reinterpret_cast<const float*>(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<const float*>(m);
+                const float s    = *reinterpret_cast<const float*>(m + pe->nb[1]);
+                if (interleaved) {
+                    const float neg_s = *reinterpret_cast<const float*>(m + pe->nb[0]);
+                    const float c1    = *reinterpret_cast<const float*>(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, const ggml_cgraph* graph, int first, void*) {
+        auto* view = const_cast<ggml_cgraph*>(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) != sd_ggml_rope_op()) ||
+            !sd_ggml_rope_describe(rope) || !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_ext_cpu_parallel_fn>(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 register_rope(ggml_backend_dev_t device) {
+        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            = sd_ggml_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));
+    }
+
+}
+#endif

src/core/cuda/convrot.cu

@@ -4,7 +4,7 @@
 #include <cstring>
 #include <string>
 #include "convrot_kernel.h"
-#include "core/int8_ops.h"
+#include "core/ggml_extensions.h"
 #include "runtime.h"
 
 namespace SD_CUDA_NAMESPACE {

src/core/cuda/i8_gated_residual.cu

@@ -9,7 +9,7 @@
 #include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
 #include <mutex>
 #include <unordered_map>
-#include "core/int8_ops.h"
+#include "core/ggml_extensions.h"
 
 namespace sd_i8_gated {
     struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {

src/core/cuda/int8.cu

@@ -1,6 +1,6 @@
 #include <cstdlib>
 #include "convrot_kernel.h"
-#include "core/int8_ops.h"
+#include "core/ggml_extensions.h"
 #ifdef SD_CUDA_I8_CUTLASS
 #include "core/cuda/i8_gemm.h"
 #else

src/core/ggml_extensions.cpp

@@ -1,13 +1,14 @@
 #include "core/ggml_extensions.h"
 
 #ifndef SD_USE_UPSTREAM_GGML
-#include "core/int8_ops.h"
+#include "core/cpu/registration.h"
 #include "ggml-ext.h"
 
-#include <cmath>
+#include <algorithm>
 #include <cstring>
 #include <mutex>
 #include <unordered_set>
+#include <utility>
 
 #ifdef SD_BACKEND_EXTENSIONS
 #include "sd_backend_extensions.h"
@@ -42,7 +43,7 @@ ggml_ext_op_t sd_ggml_rope_op() {
     return op;
 }
 
-static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+bool sd_ggml_rope_describe(const ggml_tensor* node) {
     if (node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node) || !node->src[0] || !node->src[1]) {
         return false;
     }
@@ -52,108 +53,191 @@ static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
            std::memcmp(output.ne, node->ne, sizeof(output.ne)) == 0;
 }
 
-struct RopeExecution {
+struct I8Attributes {
-    const ggml_tensor* rope;
+    int64_t shape[3];
-    const ggml_tensor* norm;
+    int32_t group_size;
-    const ggml_tensor* weight;
+    int32_t flags;
 };
+enum I8Flags { I8_PACKED       = 1,
+               I8_INPUT_SCALE  = 2,
+               I8_WEIGHT_SCALE = 4,
+               I8_BIAS         = 8 };
 
-static void rope_compute(ggml_tensor* dst, int ith, int nth, void* user_data) {
+static bool valid_group(int group) {
-    const auto& execution = *static_cast<const RopeExecution*>(user_data);
+    return group == 1 || group == 4 || group == 16 || group == 64 || group == 256;
-    const auto* rope      = execution.rope;
+}
-    const auto* x         = execution.norm ? execution.norm->src[0] : rope->src[0];
+static constexpr int64_t i8_size_limit = int64_t(std::min(uint64_t(INT64_MAX), uint64_t(SIZE_MAX)));
-    const auto* pe        = rope->src[1];
+static bool packed_height(int64_t k, int64_t rows, int64_t* height) {
-    int32_t interleaved;
+    if (k <= 0 || rows <= 0 || rows > i8_size_limit / 4)
-    std::memcpy(&interleaved, ggml_ext_op_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, sizeof(interleaved));
+        return false;
-    float eps = 0.0f;
+    const int64_t padded = (rows + 3) / 4 * 4;
-    if (execution.norm) {
+    const int64_t scales = (rows * 4 - 1) / k + 1;
-        std::memcpy(&eps, execution.norm->op_params, sizeof(eps));
+    if (padded > i8_size_limit / k || scales > i8_size_limit / k - padded)
+        return false;
+    *height = padded + scales;
+    return true;
+}
+static bool scale_matches(const ggml_tensor* scale, int64_t count) {
+    return !scale || (scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale) && ggml_nelements(scale) == count);
+}
+static bool i8_shape(const SDI8Operation& op, const int64_t* shape) {
+    const auto* w = op.weight;
+    const auto* x = op.input;
+    if (!w || !x || w->type != GGML_TYPE_I8 || w->ne[2] != 1 || w->ne[3] != 1 ||
+        !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || x->ne[0] != w->ne[0] ||
+        (op.group_size != 0 && (!valid_group(op.group_size) || w->ne[0] % op.group_size)))
+        return false;
+    int64_t rows = 1;
+    for (int i = 0; i < 3; ++i) {
+        if (shape[i] <= 0 || shape[i] > (INT64_MAX - 3) / rows)
+            return false;
+        rows *= shape[i];
     }
-    const int64_t half_dim = pe->ne[2];
+    const int64_t k = w->ne[0], n = w->ne[1];
-    const int64_t rows     = x->ne[1] * x->ne[2] * x->ne[3];
+    int64_t height;
-    for (int64_t row = ith; row < rows; row += nth) {
+    if (n <= 0 || n > (i8_size_limit - 4) / 4 || k <= 0 || k > i8_size_limit - 4 - n * 4 ||
-        const int64_t token = row % x->ne[2];
+        rows > (i8_size_limit - 3) / (k + 4 + n * 4) || !packed_height(k, rows, &height))
-        const int64_t head  = (row / x->ne[2]) % x->ne[1];
+        return false;
-        const int64_t batch = row / (x->ne[2] * x->ne[1]);
+    if (op.packed) {
-        const char* input   = static_cast<const char*>(x->data) + head * x->nb[1] + token * x->nb[2] + batch * x->nb[3];
+        if (x->type != GGML_TYPE_I8 || !valid_group(op.group_size) || op.input_scale ||
-        const char* pos     = static_cast<const char*>(pe->data) + token * pe->nb[3];
+            x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1)
-        float* output       = static_cast<float*>(dst->data) + row * x->ne[0];
+            return false;
-        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<const float*>(input + col * x->nb[0]);
-                sum += static_cast<double>(v * v);
-            }
-            const float mean = static_cast<float>(sum / x->ne[0]);
-            scale            = 1.0f / std::sqrt(mean + eps);
-        }
-        auto value = [&](int64_t col) {
-            float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
-            if (execution.norm) {
-                const auto* weight = static_cast<const char*>(execution.weight->data);
-                v                  = (v * scale) * *reinterpret_cast<const float*>(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<const float*>(m);
-            const float s    = *reinterpret_cast<const float*>(m + pe->nb[1]);
-            if (interleaved) {
-                const float neg_s = *reinterpret_cast<const float*>(m + pe->nb[0]);
-                const float c1    = *reinterpret_cast<const float*>(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;
+        if (std::memcmp(shape, x->ne + 1, 3 * sizeof(int64_t)))
-                output[i1] = a * s + b * c;
+            return false;
-            }
+        if (x->type == GGML_TYPE_I8) {
-        }
+            if (!op.input_scale || op.group_size != 0)
-        for (int64_t col = 2 * half_dim; col < x->ne[0]; ++col) {
+                return false;
-            output[col] = value(col);
+        } else if (x->type != GGML_TYPE_F32 || op.input_scale)
-        }
+            return false;
     }
+    return scale_matches(op.input_scale, rows) && scale_matches(op.weight_scale, w->ne[1]) && scale_matches(op.bias, w->ne[1]);
 }
-
+static bool i8_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
-static int rope_fusion_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+    I8Attributes attr;
-    auto* view = const_cast<ggml_cgraph*>(graph);
+    std::memcpy(&attr, attributes, sizeof(attr));
-    if (first + 2 >= ggml_graph_n_nodes(view)) {
+    if (count < 2 || (attr.flags & ~15))
-        return 0;
+        return false;
+    SDI8Operation op = {inputs[0], inputs[1], nullptr, nullptr, nullptr, attr.group_size, (attr.flags & I8_PACKED) != 0};
+    int index        = 2;
+    for (auto entry : {std::make_pair(I8_INPUT_SCALE, &op.input_scale), std::make_pair(I8_WEIGHT_SCALE, &op.weight_scale), std::make_pair(I8_BIAS, &op.bias)}) {
+        if (attr.flags & entry.first) {
+            if (index >= count)
+                return false;
+            *entry.second = inputs[index++];
         }
-    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) != sd_ggml_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 (index != count || !i8_shape(op, attr.shape))
-    if (!weight || weight->type != GGML_TYPE_F32 || weight->ne[0] != norm->ne[0] ||
+        return false;
-        weight->ne[1] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1) {
+    output->type  = GGML_TYPE_F32;
-        return 0;
+    output->ne[0] = op.weight->ne[1];
-    }
+    std::memcpy(output->ne + 1, attr.shape, sizeof(attr.shape));
-    return 3;
+    return true;
 }
-
+static bool convrot_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
-static ggml_status rope_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int count, void*, void*) {
+    int32_t group;
-    auto device = ggml_backend_get_device(context->backend);
+    std::memcpy(&group, attributes, sizeof(group));
-    auto reg    = ggml_backend_dev_backend_reg(device);
+    if (count != 1 || inputs[0]->type != GGML_TYPE_F32 || !ggml_is_contiguous(inputs[0]) || !valid_group(group) || inputs[0]->ne[0] <= 0 || inputs[0]->ne[0] % group)
-    auto run    = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+        return false;
-    GGML_ASSERT(run);
+    const int64_t k = inputs[0]->ne[0], rows = ggml_nrows(inputs[0]);
-    RopeExecution execution = {nodes[count - 1], nullptr, nullptr};
+    int64_t height;
-    if (count == 3) {
+    if (!packed_height(k, rows, &height))
-        execution.norm   = nodes[0];
+        return false;
-        execution.weight = nodes[1]->src[0] == nodes[0] ? nodes[1]->src[1] : nodes[1]->src[0];
+    output->type  = GGML_TYPE_I8;
+    output->ne[0] = k;
+    output->ne[1] = height;
+    output->ne[2] = output->ne[3] = 1;
+    return true;
+}
+ggml_ext_op_t sd_ggml_i8_op() {
+    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.linear_i8.v1", sizeof(I8Attributes), i8_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
+    return op;
+}
+ggml_ext_op_t sd_ggml_convrot_op() {
+    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.quantize_i8_convrot.v1", sizeof(int32_t), convrot_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
+    return op;
+}
+bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* op) {
+    if (!node || node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node))
+        return false;
+    *op = {node->src[0], node->src[1], nullptr, nullptr, nullptr, 0, false};
+    if (node->op == GGML_OP_MUL_MAT) {
+        op->weight_scale = node->src[2];
+        op->bias         = node->src[3];
+        op->group_size   = node->op_params[2];
+        op->packed       = node->op_params[4] > 0;
+        if (op->packed)
+            op->group_size = node->op_params[4];
+    } else if (ggml_ext_op_get(node) == sd_ggml_i8_op()) {
+        I8Attributes attr;
+        std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr));
+        if (attr.flags & ~15 || std::memcmp(attr.shape, node->ne + 1, sizeof(attr.shape)))
+            return false;
+        op->group_size = attr.group_size;
+        op->packed     = (attr.flags & I8_PACKED) != 0;
+        int index      = 2;
+        if (attr.flags & I8_INPUT_SCALE)
+            op->input_scale = node->src[index++];
+        if (attr.flags & I8_WEIGHT_SCALE)
+            op->weight_scale = node->src[index++];
+        if (attr.flags & I8_BIAS)
+            op->bias = node->src[index++];
+        if (((attr.flags & I8_INPUT_SCALE) && !op->input_scale) || ((attr.flags & I8_WEIGHT_SCALE) && !op->weight_scale) || ((attr.flags & I8_BIAS) && !op->bias))
+            return false;
+    } else
+        return false;
+    return i8_shape(*op, node->ne + 1) && node->ne[0] == op->weight->ne[1];
+}
+bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) {
+    if (!node || !node->src[0] || !ggml_is_contiguous(node))
+        return false;
+    if (node->op == GGML_OP_QUANTIZE_I8_CONVROT)
+        *group = node->op_params[0];
+    else if (ggml_ext_op_get(node) == sd_ggml_convrot_op())
+        std::memcpy(group, ggml_ext_op_attributes(node), sizeof(int32_t));
+    else
+        return false;
+    ggml_ext_tensor_desc output;
+    return convrot_infer(node->src, 1, group, &output) && output.type == node->type && !std::memcmp(output.ne, node->ne, sizeof(output.ne));
+}
+static ggml_tensor* i8_create(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* xs, ggml_tensor* ws, ggml_tensor* bias, int group, const ggml_tensor* logical) {
+    I8Attributes attr      = {{logical->ne[1], logical->ne[2], logical->ne[3]}, group, logical != x ? I8_PACKED : 0};
+    ggml_tensor* inputs[5] = {w, x};
+    int count              = 2;
+    if (xs) {
+        attr.flags |= I8_INPUT_SCALE;
+        inputs[count++] = xs;
     }
-    return run(context->backend, nodes[count - 1], rope_compute, &execution);
+    if (ws) {
+        attr.flags |= I8_WEIGHT_SCALE;
+        inputs[count++] = ws;
+    }
+    if (bias) {
+        attr.flags |= I8_BIAS;
+        inputs[count++] = bias;
+    }
+    return ggml_ext_op_create(ctx, sd_ggml_i8_op(), inputs, count, &attr);
+}
+ggml_tensor* sd_ggml_mul_mat_i8(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* xs, ggml_tensor* ws, ggml_tensor* bias) {
+    if (!w || !x || !xs || x->type != GGML_TYPE_I8)
+        return nullptr;
+    return i8_create(ctx, w, x, xs, ws, bias, 0, x);
+}
+ggml_tensor* sd_ggml_linear_i8(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* bias, int group) {
+    if (!w || !x)
+        return nullptr;
+    const auto* logical = x;
+    if (x->type == GGML_TYPE_I8) {
+        int packed_group;
+        if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != group)
+            return nullptr;
+        logical = x->src[0];
+    }
+    return i8_create(ctx, w, x, nullptr, ws, bias, group, logical);
+}
+ggml_tensor* sd_ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int group) {
+    int32_t attr = group;
+    return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr);
 }
 
 void sd_ggml_register_extensions(ggml_backend_t backend) {
@@ -176,28 +260,7 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
 #endif
         return;
     }
-    sd_ggml_register_cpu_i8(device);
+    sd_cpu::register_extensions(device);
-    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            = sd_ggml_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) {

src/core/ggml_extensions.h

@@ -6,8 +6,29 @@
 void sd_ggml_register_extensions(ggml_backend_t backend);
 
 #ifndef SD_USE_UPSTREAM_GGML
-const struct ggml_ext_op* sd_ggml_rope_op();
+#include "ggml-ext.h"
+
+ggml_ext_op_t sd_ggml_rope_op();
+bool sd_ggml_rope_describe(const ggml_tensor* node);
 ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved);
+
+struct SDI8Operation {
+    const ggml_tensor* weight;
+    const ggml_tensor* input;
+    const ggml_tensor* input_scale;
+    const ggml_tensor* weight_scale;
+    const ggml_tensor* bias;
+    int group_size;
+    bool packed;
+};
+
+ggml_ext_op_t sd_ggml_i8_op();
+ggml_ext_op_t sd_ggml_convrot_op();
+bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* operation);
+bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group_size);
+ggml_tensor* sd_ggml_mul_mat_i8(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* input_scale, ggml_tensor* weight_scale = nullptr, ggml_tensor* bias = nullptr);
+ggml_tensor* sd_ggml_linear_i8(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* bias, int group_size);
+ggml_tensor* sd_ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size);
 #endif
 
 #endif

src/core/int8_ops.cpp

@@ -1,373 +0,0 @@
-#include "core/int8_ops.h"
-
-#ifndef SD_USE_UPSTREAM_GGML
-#include <algorithm>
-#include <cmath>
-#include <cstring>
-#include <limits>
-#include <utility>
-
-struct I8Attributes {
-    int64_t shape[3];
-    int32_t group_size;
-    int32_t flags;
-};
-enum I8Flags { I8_PACKED       = 1,
-               I8_INPUT_SCALE  = 2,
-               I8_WEIGHT_SCALE = 4,
-               I8_BIAS         = 8 };
-
-static bool valid_group(int group) {
-    return group == 1 || group == 4 || group == 16 || group == 64 || group == 256;
-}
-static constexpr int64_t i8_size_limit = int64_t(std::min(uint64_t(INT64_MAX), uint64_t(SIZE_MAX)));
-static bool packed_height(int64_t k, int64_t rows, int64_t* height) {
-    if (k <= 0 || rows <= 0 || rows > i8_size_limit / 4)
-        return false;
-    const int64_t padded = (rows + 3) / 4 * 4;
-    const int64_t scales = (rows * 4 - 1) / k + 1;
-    if (padded > i8_size_limit / k || scales > i8_size_limit / k - padded)
-        return false;
-    *height = padded + scales;
-    return true;
-}
-static bool scale_matches(const ggml_tensor* scale, int64_t count) {
-    return !scale || (scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale) && ggml_nelements(scale) == count);
-}
-static bool i8_shape(const SDI8Operation& op, const int64_t* shape) {
-    const auto* w = op.weight;
-    const auto* x = op.input;
-    if (!w || !x || w->type != GGML_TYPE_I8 || w->ne[2] != 1 || w->ne[3] != 1 ||
-        !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || x->ne[0] != w->ne[0] ||
-        (op.group_size != 0 && (!valid_group(op.group_size) || w->ne[0] % op.group_size)))
-        return false;
-    int64_t rows = 1;
-    for (int i = 0; i < 3; ++i) {
-        if (shape[i] <= 0 || shape[i] > (INT64_MAX - 3) / rows)
-            return false;
-        rows *= shape[i];
-    }
-    const int64_t k = w->ne[0], n = w->ne[1];
-    int64_t height;
-    if (n <= 0 || n > (i8_size_limit - 4) / 4 || k <= 0 || k > i8_size_limit - 4 - n * 4 ||
-        rows > (i8_size_limit - 3) / (k + 4 + n * 4) || !packed_height(k, rows, &height))
-        return false;
-    if (op.packed) {
-        if (x->type != GGML_TYPE_I8 || !valid_group(op.group_size) || op.input_scale ||
-            x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1)
-            return false;
-    } else {
-        if (std::memcmp(shape, x->ne + 1, 3 * sizeof(int64_t)))
-            return false;
-        if (x->type == GGML_TYPE_I8) {
-            if (!op.input_scale || op.group_size != 0)
-                return false;
-        } else if (x->type != GGML_TYPE_F32 || op.input_scale)
-            return false;
-    }
-    return scale_matches(op.input_scale, rows) && scale_matches(op.weight_scale, w->ne[1]) && scale_matches(op.bias, w->ne[1]);
-}
-static bool i8_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
-    I8Attributes attr;
-    std::memcpy(&attr, attributes, sizeof(attr));
-    if (count < 2 || (attr.flags & ~15))
-        return false;
-    SDI8Operation op = {inputs[0], inputs[1], nullptr, nullptr, nullptr, attr.group_size, (attr.flags & I8_PACKED) != 0};
-    int index        = 2;
-    for (auto entry : {std::make_pair(I8_INPUT_SCALE, &op.input_scale), std::make_pair(I8_WEIGHT_SCALE, &op.weight_scale), std::make_pair(I8_BIAS, &op.bias)}) {
-        if (attr.flags & entry.first) {
-            if (index >= count)
-                return false;
-            *entry.second = inputs[index++];
-        }
-    }
-    if (index != count || !i8_shape(op, attr.shape))
-        return false;
-    output->type  = GGML_TYPE_F32;
-    output->ne[0] = op.weight->ne[1];
-    std::memcpy(output->ne + 1, attr.shape, sizeof(attr.shape));
-    return true;
-}
-static bool convrot_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
-    int32_t group;
-    std::memcpy(&group, attributes, sizeof(group));
-    if (count != 1 || inputs[0]->type != GGML_TYPE_F32 || !ggml_is_contiguous(inputs[0]) || !valid_group(group) || inputs[0]->ne[0] <= 0 || inputs[0]->ne[0] % group)
-        return false;
-    const int64_t k = inputs[0]->ne[0], rows = ggml_nrows(inputs[0]);
-    int64_t height;
-    if (!packed_height(k, rows, &height))
-        return false;
-    output->type  = GGML_TYPE_I8;
-    output->ne[0] = k;
-    output->ne[1] = height;
-    output->ne[2] = output->ne[3] = 1;
-    return true;
-}
-ggml_ext_op_t sd_ggml_i8_op() {
-    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.linear_i8.v1", sizeof(I8Attributes), i8_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
-    return op;
-}
-ggml_ext_op_t sd_ggml_convrot_op() {
-    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.quantize_i8_convrot.v1", sizeof(int32_t), convrot_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
-    return op;
-}
-bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* op) {
-    if (!node || node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node))
-        return false;
-    *op = {node->src[0], node->src[1], nullptr, nullptr, nullptr, 0, false};
-    if (node->op == GGML_OP_MUL_MAT) {
-        op->weight_scale = node->src[2];
-        op->bias         = node->src[3];
-        op->group_size   = node->op_params[2];
-        op->packed       = node->op_params[4] > 0;
-        if (op->packed)
-            op->group_size = node->op_params[4];
-    } else if (ggml_ext_op_get(node) == sd_ggml_i8_op()) {
-        I8Attributes attr;
-        std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr));
-        if (attr.flags & ~15 || std::memcmp(attr.shape, node->ne + 1, sizeof(attr.shape)))
-            return false;
-        op->group_size = attr.group_size;
-        op->packed     = (attr.flags & I8_PACKED) != 0;
-        int index      = 2;
-        if (attr.flags & I8_INPUT_SCALE)
-            op->input_scale = node->src[index++];
-        if (attr.flags & I8_WEIGHT_SCALE)
-            op->weight_scale = node->src[index++];
-        if (attr.flags & I8_BIAS)
-            op->bias = node->src[index++];
-        if (((attr.flags & I8_INPUT_SCALE) && !op->input_scale) || ((attr.flags & I8_WEIGHT_SCALE) && !op->weight_scale) || ((attr.flags & I8_BIAS) && !op->bias))
-            return false;
-    } else
-        return false;
-    return i8_shape(*op, node->ne + 1) && node->ne[0] == op->weight->ne[1];
-}
-bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) {
-    if (!node || !node->src[0] || !ggml_is_contiguous(node))
-        return false;
-    if (node->op == GGML_OP_QUANTIZE_I8_CONVROT)
-        *group = node->op_params[0];
-    else if (ggml_ext_op_get(node) == sd_ggml_convrot_op())
-        std::memcpy(group, ggml_ext_op_attributes(node), sizeof(int32_t));
-    else
-        return false;
-    ggml_ext_tensor_desc output;
-    return convrot_infer(node->src, 1, group, &output) && output.type == node->type && !std::memcmp(output.ne, node->ne, sizeof(output.ne));
-}
-static ggml_tensor* i8_create(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* xs, ggml_tensor* ws, ggml_tensor* bias, int group, const ggml_tensor* logical) {
-    I8Attributes attr      = {{logical->ne[1], logical->ne[2], logical->ne[3]}, group, logical != x ? I8_PACKED : 0};
-    ggml_tensor* inputs[5] = {w, x};
-    int count              = 2;
-    if (xs) {
-        attr.flags |= I8_INPUT_SCALE;
-        inputs[count++] = xs;
-    }
-    if (ws) {
-        attr.flags |= I8_WEIGHT_SCALE;
-        inputs[count++] = ws;
-    }
-    if (bias) {
-        attr.flags |= I8_BIAS;
-        inputs[count++] = bias;
-    }
-    return ggml_ext_op_create(ctx, sd_ggml_i8_op(), inputs, count, &attr);
-}
-ggml_tensor* sd_ggml_mul_mat_i8(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* xs, ggml_tensor* ws, ggml_tensor* bias) {
-    if (!w || !x || !xs || x->type != GGML_TYPE_I8)
-        return nullptr;
-    return i8_create(ctx, w, x, xs, ws, bias, 0, x);
-}
-ggml_tensor* sd_ggml_linear_i8(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* bias, int group) {
-    if (!w || !x)
-        return nullptr;
-    const auto* logical = x;
-    if (x->type == GGML_TYPE_I8) {
-        int packed_group;
-        if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != group)
-            return nullptr;
-        logical = x->src[0];
-    }
-    return i8_create(ctx, w, x, nullptr, ws, bias, group, logical);
-}
-ggml_tensor* sd_ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int group) {
-    int32_t attr = group;
-    return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr);
-}
-
-static void transform(float* values, int group) {
-    const float scale = 1.f / std::sqrt(float(group));
-    for (int i = 0; i < group; ++i)
-        values[i] *= scale;
-    for (int stride = 1; stride < group; stride *= 4)
-        for (int base = 0; base < group; base += 4 * stride)
-            for (int j = 0; j < stride; ++j) {
-                const int i   = base + j;
-                const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride];
-                values[i]              = a + b + c - d;
-                values[i + stride]     = a + b - c + d;
-                values[i + 2 * stride] = a - b + c + d;
-                values[i + 3 * stride] = -a + b + c + d;
-            }
-}
-struct I8Execution {
-    SDI8Operation op;
-    int64_t rows;
-    int8_t* quant;
-    float* scales;
-    float* output;
-};
-static void quantize(ggml_tensor*, int ith, int nth, void* data) {
-    auto& e           = *static_cast<I8Execution*>(data);
-    const int64_t k   = e.op.input->ne[0];
-    const int group   = e.op.group_size;
-    const auto* input = static_cast<const float*>(e.op.input->data);
-    for (int64_t row = ith; row < e.rows; row += nth) {
-        const float* src = input + row * k;
-        int8_t* dst      = e.quant + row * k;
-        float amax       = 0.f, values[256];
-        auto visit       = [&](auto fn) { if (!group) { for (int64_t i = 0; i < k; ++i) fn(i, src[i]); }
-            else for (int64_t base = 0; base < k; base += group) { std::memcpy(values, src + base, group * sizeof(float)); transform(values, group); for (int j = 0; j < group; ++j) fn(base + j, values[j]); } };
-        visit([&](int64_t, float value) { amax = std::max(amax, std::fabs(value)); });
-        const float scale = amax / 127.f, inverse = scale == 0.f ? 0.f : 1.f / scale;
-        e.scales[row] = scale;
-        if (scale == 0.f)
-            std::memset(dst, 0, k);
-        else
-            visit([&](int64_t i, float value) { dst[i] = static_cast<int8_t>(std::max(-127L, std::min(127L, std::lrintf(value * inverse)))); });
-    }
-}
-static void multiply(ggml_tensor*, int ith, int nth, void* data) {
-    auto& e         = *static_cast<I8Execution*>(data);
-    const auto& op  = e.op;
-    const int64_t k = op.weight->ne[0], n = op.weight->ne[1];
-    const auto* w    = static_cast<const int8_t*>(op.weight->data);
-    const auto* ws   = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
-    const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
-    for (int64_t index = ith; index < e.rows * n; index += nth) {
-        const int64_t row = index / n, column = index % n;
-        uint32_t sum = 0;
-        for (int64_t i = 0; i < k; ++i)
-            sum += uint32_t(int32_t(w[column * k + i]) * int32_t(e.quant[row * k + i]));
-        int32_t signed_sum;
-        std::memcpy(&signed_sum, &sum, sizeof(sum));
-        float value = float(signed_sum) * e.scales[row];
-        if (ws)
-            value *= ws[column];
-        if (bias)
-            value += bias[column];
-        e.output[index] = value;
-    }
-}
-static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
-    if (!src)
-        return false;
-    const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(src->data);
-    return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src);
-}
-static size_t quant_workspace(const ggml_tensor* node, const SDI8Operation& op) {
-    return op.input->type == GGML_TYPE_F32 ? (size_t(ggml_nelements(op.input)) + 3) / 4 * 4 + size_t(ggml_nrows(node)) * 4 : 0;
-}
-static bool i8_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
-    SDI8Operation op;
-    return sd_ggml_i8_describe(node, &op);
-}
-static bool convrot_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
-    int group;
-    return sd_ggml_convrot_describe(node, &group);
-}
-static ggml_status i8_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
-    SDI8Operation op;
-    if (!sd_ggml_i8_describe(nodes[0], &op))
-        return GGML_STATUS_FAILED;
-    *bytes = quant_workspace(nodes[0], op) + ggml_nbytes(nodes[0]);
-    return GGML_STATUS_SUCCESS;
-}
-static ggml_status convrot_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
-    *bytes = ggml_nbytes(nodes[0]);
-    return GGML_STATUS_SUCCESS;
-}
-static void destroy(void*, void*) {}
-static ggml_ext_cpu_parallel_fn parallel(ggml_backend_t backend) {
-    return reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)), "ggml_ext_cpu_parallel_v1"));
-}
-static ggml_status i8_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
-    auto* dst     = nodes[0];
-    I8Execution e = {};
-    if (!sd_ggml_i8_describe(dst, &e.op))
-        return GGML_STATUS_FAILED;
-    e.rows   = ggml_nrows(dst);
-    auto run = parallel(context->backend);
-    if (!run)
-        return GGML_STATUS_FAILED;
-    char* scratch           = static_cast<char*>(ggml_backend_buffer_get_base(context->workspace));
-    const size_t quant_size = quant_workspace(dst, e.op);
-    if (e.op.input->type == GGML_TYPE_F32) {
-        e.quant     = reinterpret_cast<int8_t*>(scratch);
-        e.scales    = reinterpret_cast<float*>(scratch + (size_t(ggml_nelements(e.op.input)) + 3) / 4 * 4);
-        auto status = run(context->backend, dst, quantize, &e);
-        if (status != GGML_STATUS_SUCCESS)
-            return status;
-    } else {
-        e.quant  = static_cast<int8_t*>(e.op.input->data);
-        e.scales = e.op.packed ? reinterpret_cast<float*>(e.quant + e.op.input->ne[0] * ((e.rows + 3) / 4 * 4)) : static_cast<float*>(e.op.input_scale->data);
-    }
-    bool alias = false;
-    for (auto* input : dst->src)
-        alias |= overlaps(dst, input);
-    e.output    = alias ? reinterpret_cast<float*>(scratch + quant_size) : static_cast<float*>(dst->data);
-    auto status = run(context->backend, dst, multiply, &e);
-    if (status == GGML_STATUS_SUCCESS && alias)
-        std::memcpy(dst->data, e.output, ggml_nbytes(dst));
-    return status;
-}
-static ggml_status convrot_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
-    auto* dst     = nodes[0];
-    I8Execution e = {};
-    if (!sd_ggml_convrot_describe(dst, &e.op.group_size))
-        return GGML_STATUS_FAILED;
-    auto run = parallel(context->backend);
-    if (!run)
-        return GGML_STATUS_FAILED;
-    e.op.input      = dst->src[0];
-    e.rows          = ggml_nrows(e.op.input);
-    bool alias      = overlaps(dst, e.op.input);
-    e.quant         = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(context->workspace) : dst->data);
-    const int64_t k = dst->ne[0], padded = (e.rows + 3) / 4 * 4;
-    e.scales    = reinterpret_cast<float*>(e.quant + k * padded);
-    auto status = run(context->backend, dst, quantize, &e);
-    if (status == GGML_STATUS_SUCCESS) {
-        std::memset(e.quant + k * e.rows, 0, k * (padded - e.rows));
-        if (alias)
-            std::memcpy(dst->data, e.quant, ggml_nbytes(dst));
-    }
-    return status;
-}
-void sd_ggml_register_cpu_i8(ggml_backend_dev_t device) {
-    ggml_ext_kernel_desc d = {};
-    d.abi_version          = GGML_EXT_ABI_VERSION;
-    d.priority             = 100;
-    d.destroy              = destroy;
-    d.name                 = "sd.linear_i8.cpu.v1";
-    d.op                   = GGML_OP_CUSTOM;
-    d.custom_op            = sd_ggml_i8_op();
-    d.supports             = i8_supports;
-    d.prepare              = i8_prepare;
-    d.execute              = i8_execute;
-    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
-    d.name      = "sd.linear_i8.cpu.override.v1";
-    d.op        = GGML_OP_MUL_MAT;
-    d.custom_op = nullptr;
-    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
-    d.name      = "sd.quantize_i8_convrot.cpu.v1";
-    d.op        = GGML_OP_CUSTOM;
-    d.custom_op = sd_ggml_convrot_op();
-    d.supports  = convrot_supports;
-    d.prepare   = convrot_prepare;
-    d.execute   = convrot_execute;
-    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
-    d.name      = "sd.quantize_i8_convrot.cpu.override.v1";
-    d.op        = GGML_OP_QUANTIZE_I8_CONVROT;
-    d.custom_op = nullptr;
-    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
-}
-#endif

src/core/int8_ops.h

@@ -1,28 +0,0 @@
-#ifndef __SD_CORE_INT8_OPS_H__
-#define __SD_CORE_INT8_OPS_H__
-
-#include "ggml-backend.h"
-
-#ifndef SD_USE_UPSTREAM_GGML
-#include "ggml-ext.h"
-struct SDI8Operation {
-    const ggml_tensor* weight;
-    const ggml_tensor* input;
-    const ggml_tensor* input_scale;
-    const ggml_tensor* weight_scale;
-    const ggml_tensor* bias;
-    int group_size;
-    bool packed;
-};
-
-ggml_ext_op_t sd_ggml_i8_op();
-ggml_ext_op_t sd_ggml_convrot_op();
-bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* operation);
-bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group_size);
-ggml_tensor* sd_ggml_mul_mat_i8(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* input_scale, ggml_tensor* weight_scale = nullptr, ggml_tensor* bias = nullptr);
-ggml_tensor* sd_ggml_linear_i8(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* bias, int group_size);
-ggml_tensor* sd_ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size);
-void sd_ggml_register_cpu_i8(ggml_backend_dev_t device);
-#endif
-
-#endif