DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-12 11:56:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: separate CPU kernels from shared operator extensions

Browse Source
This commit is contained in:
leejet committed 2026-11-01 00:44:28 +08:00
1 parent d11410eee7
commit d07d48d466
13 files changed
+570 -520
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

CMakeLists.txt
+2
View File
@@ -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
+12 -5
View File
@@ -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
`sd.quantize_i8_convrot.v1`, with CPU implementations and overrides of the old
builtins. The CPU kernels use the configured thread pool and scheduler workspace;
the old CPU INT8 GEMM and ConvRot implementations have been removed from ggml.
`src/core/ggml_extensions.cpp` defines `sd.linear_i8.v1` and
`sd.quantize_i8_convrot.v1`. `src/core/cpu/int8.cpp` supplies their CPU
implementations and overrides of the old builtins. The CPU kernels use the
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
+189
View File
@@ -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
+12
View File
@@ -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
+14
View File
@@ -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
+143
View File
@@ -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
+1 -1
View File
@@ -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
+1 -1
View File
@@ -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 -1
View File
@@ -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
+173 -110
View File
@@ -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 {
const ggml_tensor* rope;
const ggml_tensor* norm;
const ggml_tensor* weight;
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 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));
static bool valid_group(int group) {
return group == 1 || group == 4 || group == 16 || group == 64 || group == 256;
}
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);
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;
}
const float mean = static_cast<float>(sum / x->ne[0]);
scale = 1.0f / std::sqrt(mean + eps);
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);
}
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]);
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];
}
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;
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 {
output[i0] = a * c - b * s;
output[i1] = a * s + b * c;
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++];
}
}
for (int64_t col = 2 * half_dim; col < x->ne[0]; ++col) {
output[col] = value(col);
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;
}
static int rope_fusion_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
auto* view = const_cast<ggml_cgraph*>(graph);
if (first + 2 >= ggml_graph_n_nodes(view)) {
return 0;
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;
}
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;
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];
}
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;
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));
}
return 3;
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;
}
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];
if (ws) {
attr.flags |= I8_WEIGHT_SCALE;
inputs[count++] = ws;
}
return run(context->backend, nodes[count - 1], rope_compute, &execution);
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);
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));
sd_cpu::register_extensions(device);
}
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
+22 -1
View File
@@ -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
-373
View File
@@ -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
-28
View File
@@ -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