DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-09 18:36:33 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
45a63ac460b9403df7a739bbc8662cb7d867eace
stable-diffusion.cpp/src/core/vulkan/int8.cpp
T

807 lines
39 KiB
C++
Raw Blame History

#include <algorithm>
#include <array>
#include <climits>
#include <cstring>
#include <vector>
#include "core/ggml_extensions.h"
#include "ggml-ext-vulkan.h"
#include "modulate_convrot.h"
#include "mul_mat_i8_residual_cm2.h"
#include "mul_mat_i8_swiglu_split_cm2.h"
#include "mul_mat_i8_tensorwise.h"
#include "mul_mat_i8_tensorwise_cm1.h"
#include "mul_mat_i8_tensorwise_cm2.h"
#include "quantize_i8_convrot.h"
#include "quantize_i8_convrot_register.h"
#include "quantize_i8_convrot_subgroup.h"
#include "rms_modulate_convrot.h"
#include "sigmoid_quantize_i8_convrot_6144.h"
#include "swiglu_quantize_i8_convrot.h"
#include "swiglu_quantize_i8_convrot_16384.h"
namespace sd_vulkan_i8 {
struct Device {
const ggml_ext_vulkan_api* api;
ggml_ext_vulkan_device_info info;
ggml_ext_vulkan_compute_info compute;
bool subgroup, cm2, cm2_small;
uint32_t cm_m, cm_n, cm_k;
};
static bool device_info(ggml_backend_dev_t dev, Device& d) {
auto reg = ggml_backend_dev_backend_reg(dev);
auto query = reinterpret_cast<ggml_ext_vulkan_api_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1"));
auto compute = reinterpret_cast<ggml_ext_vulkan_compute_api_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_compute_v1"));
d = {};
if (!query || !compute)
return false;
d.api = query();
auto service = compute();
if (!d.api->device_info(dev, &d.info) || !service->properties(dev, &d.compute))
return false;
d.subgroup = d.info.subgroup_size == 32 && d.info.full_subgroups && d.info.subgroup_shuffle && d.compute.subgroup_arithmetic && d.info.max_group_invocations >= 256;
ggml_ext_vulkan_matrix_properties prop;
for (uint32_t i = 0; service->matrix_properties(dev, i, &prop); ++i) {
if (prop.a != GGML_TYPE_I8 || prop.b != GGML_TYPE_I8 || prop.accumulator != GGML_TYPE_I32 || prop.result != GGML_TYPE_I32 || prop.saturating || !prop.m || !prop.n || !prop.k)
continue;
if (!prop.flexible && !prop.workgroup_scope && !d.cm_m && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise_cm1) > 4) {
d.cm_m = prop.m;
d.cm_n = prop.n;
d.cm_k = prop.k;
}
if (prop.flexible && prop.workgroup_scope && prop.invocations == 128 && 128 % prop.m == 0 && 64 % prop.n == 0 && 64 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise_cm2) > 4) {
d.cm2 = true;
d.cm2_small |= 64 % prop.m == 0 && 32 % prop.n == 0;
}
}
return d.info.integer_dot_product && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise) > 4;
}
static const ggml_tensor* node(const ggml_cgraph* graph, int i) {
return ggml_graph_node(const_cast<ggml_cgraph*>(graph), i);
}
static int count_nodes(const ggml_cgraph* graph) {
return ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph));
}
static int param_i32(const ggml_tensor* node, int i) {
return node->op_params[i];
}
static float param_f32(const ggml_tensor* node, int i) {
float value;
std::memcpy(&value, node->op_params + i, 4);
return value;
}
static bool is_convrot(const ggml_tensor* node) {
int group;
return sd_ggml_convrot_describe(node, &group);
}
static bool is_mat(const ggml_tensor* node) {
SDI8Operation op;
return sd_ggml_i8_describe(node, &op);
}
static bool aligned(const ggml_tensor* t, size_t alignment) {
size_t offset = 0;
while (!t->data && t->view_src) {
offset += t->view_offs;
t = t->view_src;
}
return (reinterpret_cast<uintptr_t>(t->data) + offset) % alignment == 0;
}
static bool bounded(const Device& d, const ggml_tensor* t, bool descriptor_aligned = true) {
return t && ggml_nbytes(t) <= d.info.max_buffer_range - (d.info.buffer_alignment - 1) &&
aligned(t, descriptor_aligned ? d.info.buffer_alignment : 4);
}
static bool can_modulate(const Device& d, int64_t k, bool affine) {
return (k == 4096 || k == 6144 || (affine && k == 5376)) && d.info.max_group_invocations >= 512 && d.info.max_group_size[0] >= 512 && d.compute.shared_memory >= (k + 1024) * 4;
}
static bool can_quantize(const Device& d, const ggml_tensor* dst) {
int group;
if (!sd_ggml_convrot_describe(dst, &group) || group != 256 || !bounded(d, dst) || !bounded(d, dst->src[0]))
return false;
const auto* x = dst->src[0];
return x->ne[0] <= UINT32_MAX && ggml_nrows(x) <= UINT32_MAX && d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 256;
}
static bool can_matmul(const Device& d, const ggml_tensor* dst) {
SDI8Operation op;
if (!sd_ggml_i8_describe(dst, &op) || !op.packed || op.group_size != 256 || !op.weight_scale || !bounded(d, dst))
return false;
for (auto* t : {op.weight, op.input, op.weight_scale, op.bias})
if (t && !bounded(d, t))
return false;
const uint64_t cols = op.weight->ne[1], rows = ggml_nrows(dst);
return op.weight->ne[0] % 256 == 0 && op.weight->ne[0] <= UINT32_MAX && cols <= UINT32_MAX && rows <= UINT32_MAX &&
(cols + 15) / 16 <= d.info.max_group_count[0] && (rows + 15) / 16 <= d.info.max_group_count[1] &&
d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16;
}
static bool can_fuse_sigmoid_convrot(const Device& device, const ggml_cgraph* graph, int i) {
if (i + 2 >= count_nodes(graph) || !device.subgroup) {
return false;
}
const ggml_tensor* sigmoid = node(graph, i);
const ggml_tensor* mul = node(graph, i + 1);
const ggml_tensor* quant = node(graph, i + 2);
if (sigmoid->op != GGML_OP_UNARY || ggml_get_unary_op(sigmoid) != GGML_UNARY_OP_SIGMOID ||
mul->op != GGML_OP_MUL || mul->src[1] != sigmoid ||
!is_convrot(quant) || quant->src[0] != mul ||
sigmoid->ne[0] != 6144 || !ggml_are_same_shape(sigmoid, mul->src[0]) ||
!can_quantize(device, quant)) {
return false;
}
for (const ggml_tensor* tensor : {sigmoid->src[0], mul->src[0]}) {
if (tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor) ||
tensor->view_offs % device.info.buffer_alignment != 0) {
return false;
}
}
return true;
}
static bool can_fuse_swiglu_convrot(const Device& device, const ggml_cgraph* graph, int i) {
if (i + 1 >= count_nodes(graph) || !device.subgroup) {
return false;
}
const ggml_tensor* glu = node(graph, i);
const ggml_tensor* quant = node(graph, i + 1);
if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU ||
glu->src[1] == nullptr || param_i32(glu, 1) != 0 ||
!is_convrot(quant) || quant->src[0] != glu ||
!can_quantize(device, quant) || !ggml_is_contiguous(glu->src[0]) || !ggml_is_contiguous(glu->src[1])) {
return false;
}
for (const ggml_tensor* tensor : {glu->src[0], glu->src[1]}) {
if (tensor->type != GGML_TYPE_F32 || tensor->view_offs % device.info.buffer_alignment != 0) {
return false;
}
}
return true;
}
static int match_modulate_convrot(const Device& device, const ggml_cgraph* graph, int quant_idx, int* indices) {
if (quant_idx < 1 || quant_idx >= count_nodes(graph)) {
return 0;
}
const ggml_tensor* quant = node(graph, quant_idx);
if (!is_convrot(quant) || !can_quantize(device, quant)) {
return 0;
}
const ggml_tensor* nodes[4]{};
const ggml_tensor* mod = nullptr;
const bool affine = quant->src[0]->op == GGML_OP_ADD;
const int count = affine ? 4 : 3;
if (affine) {
const ggml_tensor* shift = quant->src[0];
const ggml_tensor* add = shift->src[0];
if (add->op != GGML_OP_ADD || add->src[1]->op != GGML_OP_MUL || add->src[1]->src[0] != add->src[0]) {
return 0;
}
nodes[0] = add->src[1];
nodes[1] = add;
nodes[2] = shift;
nodes[3] = quant;
mod = nodes[0]->src[1];
const ggml_tensor* vector = shift->src[1];
if (vector->type != GGML_TYPE_F32 || !ggml_is_contiguous(vector) ||
vector->ne[0] != add->ne[0] || ggml_nelements(vector) != add->ne[0]) {
return 0;
}
} else {
const ggml_tensor* mul = quant->src[0];
if (mul->op != GGML_OP_MUL || mul->src[0]->op != GGML_OP_NORM) {
return 0;
}
nodes[0] = mul->src[0];
nodes[1] = mul;
nodes[2] = quant;
mod = mul->src[1];
}
const ggml_tensor* x = nodes[0]->src[0];
if ((x->ne[0] != 4096 && x->ne[0] != 6144 && x->ne[0] != 5376) ||
!can_modulate(device, x->ne[0], affine) ||
x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || mod->type != GGML_TYPE_F32 ||
!ggml_is_contiguous(mod) || mod->ne[0] != x->ne[0] || ggml_nelements(mod) != x->ne[0]) {
return 0;
}
for (int j = 0; j < count; ++j) {
if (j < count - 1 && (nodes[j]->type != GGML_TYPE_F32 || !ggml_is_contiguous(nodes[j]) || !ggml_are_same_shape(x, nodes[j]))) {
return 0;
}
}
int next = count - 1;
for (int j = quant_idx; j >= std::max(0, quant_idx - 64) && next >= 0; --j) {
if (node(graph, j) == nodes[next]) {
indices[next--] = j;
}
}
return next < 0 ? count : 0;
}
static int match_rms_modulate_convrot(const Device& device, const ggml_cgraph* graph, int q, int* indices) {
if (!can_modulate(device, 5376, true) || match_modulate_convrot(device, graph, q, indices + 2) != 4) {
return 0;
}
const auto weighted = node(graph, indices[2])->src[0];
if (weighted->op != GGML_OP_MUL || weighted->src[0]->op != GGML_OP_RMS_NORM || weighted->ne[0] != 5376) {
return 0;
}
const auto norm = weighted->src[0];
const auto x = norm->src[0];
const auto weight = weighted->src[1];
if (x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || weight->type != GGML_TYPE_F32 ||
!ggml_is_contiguous(weight) || ggml_nelements(weight) != 5376 || weight->ne[0] != 5376) {
return 0;
}
int next = 1;
for (int j = indices[2] - 1; j >= std::max(0, indices[2] - 64) && next >= 0; --j) {
if (node(graph, j) == (next == 1 ? weighted : norm)) {
indices[next--] = j;
}
}
return next < 0 ? 6 : 0;
}
static int can_fuse_modulate_convrot(const Device& device, const ggml_cgraph* graph, int i) {
const ggml_op op = node(graph, i)->op;
if (op == GGML_OP_RMS_NORM) {
int indices[6];
if (match_rms_modulate_convrot(device, graph, i + 5, indices) != 6) {
return 0;
}
for (int j = 0; j < 6; ++j) {
if (indices[j] != i + j) {
return 0;
}
}
return 6;
}
if (op != GGML_OP_NORM && op != GGML_OP_MUL) {
return 0;
}
const int count = op == GGML_OP_NORM ? 3 : 4;
int indices[4];
if (match_modulate_convrot(device, graph, i + count - 1, indices) != count) {
return 0;
}
for (int j = 0; j < count; ++j) {
if (indices[j] != i + j) {
return 0;
}
}
return count;
}
static bool match_i8_swiglu_split(const Device& device, const ggml_cgraph* graph, int i, int* indices) {
if (i < 2 || i >= count_nodes(graph) || !device.cm2_small) {
return false;
}
const ggml_tensor* glu = node(graph, i);
if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU ||
param_i32(glu, 1) != 0 || glu->src[1] == nullptr || glu->ne[0] % 128 != 0 ||
glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu) ||
glu->view_offs % device.info.buffer_alignment != 0) {
return false;
}
const ggml_tensor* gate = glu->src[0];
const ggml_tensor* up = glu->src[1];
if (gate == up || gate->src[1] != up->src[1] || !ggml_are_same_shape(gate, up)) {
return false;
}
for (auto mat : {gate, up}) {
if (!is_mat(mat) || mat->view_src != nullptr || mat->src[3] != nullptr ||
!can_matmul(device, mat)) {
return false;
}
}
int found = 0;
for (int j = std::max(0, i - 32); j < i; ++j) {
if (node(graph, j) == gate || node(graph, j) == up) {
indices[found++] = j;
}
}
if (found != 2) {
return false;
}
indices[2] = i;
return true;
}
static bool can_fuse_i8_swiglu_split(const Device& device, const ggml_cgraph* graph, int i) {
int indices[3];
return match_i8_swiglu_split(device, graph, i + 2, indices) && indices[0] == i && indices[1] == i + 1;
}
static bool can_fuse_i8_swiglu(const Device& device, const ggml_cgraph* graph, int i) {
if (i + 1 >= count_nodes(graph) || !device.cm2_small) {
return false;
}
const ggml_tensor* mat = node(graph, i);
const ggml_tensor* glu = node(graph, i + 1);
return is_mat(mat) && can_matmul(device, mat) &&
mat->view_src == nullptr && mat->src[3] == nullptr && mat->ne[0] % 256 == 0 &&
glu->op == GGML_OP_GLU && ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU &&
param_i32(glu, 1) == 0 && glu->src[0] == mat && glu->src[1] == nullptr &&
glu->type == GGML_TYPE_F32 && ggml_is_contiguous(glu) &&
glu->view_offs % device.info.buffer_alignment == 0;
}
static int match_i8_residual(const Device& device, const ggml_cgraph* graph, int add_idx, int* indices) {
if (add_idx < 2 || add_idx >= count_nodes(graph) || !device.cm2) {
return 0;
}
const ggml_tensor* dst = node(graph, add_idx);
if (dst->op != GGML_OP_ADD) {
return 0;
}
const ggml_tensor* mul = dst->src[0]->op == GGML_OP_MUL ? dst->src[0] : dst->src[1];
if (mul->op != GGML_OP_MUL) {
return 0;
}
const ggml_tensor* scaled_mat = mul->src[0];
const bool scaled = scaled_mat->op == GGML_OP_SCALE;
const ggml_tensor* mat = scaled ? scaled_mat->src[0] : scaled_mat;
if (!is_mat(mat) || !can_matmul(device, mat) || mat->src[3] != nullptr ||
(scaled && param_f32(scaled_mat, 1) != 0.0f)) {
return 0;
}
const int count = scaled ? 4 : 3;
const ggml_tensor* nodes[4] = {mat, scaled ? scaled_mat : mul, scaled ? mul : dst, dst};
const ggml_tensor* gate = mul->src[1];
const ggml_tensor* residual = dst->src[0] == mul ? dst->src[1] : dst->src[0];
if (!ggml_are_same_shape(mat, dst) || !ggml_are_same_shape(residual, dst) ||
gate->ne[0] != mat->ne[0] || ggml_nelements(gate) != mat->ne[0]) {
return 0;
}
for (const ggml_tensor* tensor : {mat, scaled_mat, mul, gate, residual, dst}) {
if (tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor) ||
tensor->view_offs % device.info.buffer_alignment != 0) {
return 0;
}
}
int next = count - 1;
for (int j = add_idx; j >= std::max(0, add_idx - 32) && next >= 0; --j) {
if (node(graph, j) == nodes[next]) {
indices[next--] = j;
}
}
return next < 0 ? count : 0;
}
static int can_fuse_i8_residual(const Device& device, const ggml_cgraph* graph, int i) {
if (i + 2 >= count_nodes(graph) || !is_mat(node(graph, i))) {
return 0;
}
const int count = node(graph, i + 1)->op == GGML_OP_SCALE ? 4 : 3;
int indices[4];
if (match_i8_residual(device, graph, i + count - 1, indices) != count) {
return 0;
}
for (int j = 0; j < count; ++j) {
if (indices[j] != i + j) {
return 0;
}
}
return count;
}
static int match_i8_row_concat(const Device& device, const ggml_cgraph* graph, int end, int* indices) {
const ggml_tensor* concat = node(graph, end);
if (concat->op != GGML_OP_CONCAT || param_i32(concat, 0) != 1 ||
concat->type != GGML_TYPE_F32 || concat->ne[2] != 1 || concat->ne[3] != 1 ||
!ggml_is_contiguous(concat) || !ggml_is_contiguous(concat->src[0]) ||
concat->view_offs % device.info.buffer_alignment != 0 ||
ggml_nbytes(concat->src[0]) % device.info.buffer_alignment != 0) {
return 0;
}
const ggml_tensor* tail = concat->src[1];
int count = 0;
for (int j = end - 1; j >= std::max(0, end - 32); --j) {
if (node(graph, j) != tail) {
continue;
}
if (is_mat(tail) && can_matmul(device, tail)) {
indices[0] = j;
count = 1;
} else {
count = match_i8_residual(device, graph, j, indices);
}
break;
}
if (!count) {
return 0;
}
indices[count++] = end;
return count;
}
static int can_fuse_i8_row_concat(const Device& device, const ggml_cgraph* graph, int i) {
if (!is_mat(node(graph, i))) {
return 0;
}
for (int end = i + 1; end < std::min(i + 5, count_nodes(graph)); ++end) {
int indices[5];
const int count = match_i8_row_concat(device, graph, end, indices);
if (count != end - i + 1) {
continue;
}
bool consecutive = true;
for (int j = 0; j < count; ++j) {
consecutive &= indices[j] == i + j;
}
if (consecutive) {
return count;
}
}
return 0;
}
struct QuantParameters {
uint32_t k, rows;
};
struct ModulateParameters {
uint32_t rows, x_offset, m_offset, h_offset;
float eps;
uint32_t k, w_offset;
};
struct MatParameters {
uint32_t n, rows, k, has_bias, row_offset;
float output_scale;
uint32_t group_cols;
};
static_assert(sizeof(QuantParameters) == 8 && sizeof(ModulateParameters) == 28 && sizeof(MatParameters) == 28, "INT8 push constant layout mismatch");
template <size_t N>
static ggml_ext_vulkan_pipeline shader(const char* name, const uint32_t (&spirv)[N], uint32_t bindings, uint32_t constants) {
return {name, spirv, sizeof(spirv), bindings, constants, nullptr, 0, 0, false};
}
#define SD_VK_PIPELINE(NAME, BINDINGS, PARAMS) shader("sd." #NAME ".vulkan", sd_vulkan_shaders::NAME, BINDINGS, sizeof(PARAMS))
static ggml_ext_vulkan_binding range(ggml_ext_vulkan_binding binding, size_t offset, size_t size) {
binding.offset += offset;
binding.size = size;
return binding;
}
static bool overlap(const ggml_ext_vulkan_buffer_info& a, const ggml_ext_vulkan_buffer_info& b) {
return a.storage == b.storage && a.offset < b.offset + b.size && b.offset < a.offset + a.size;
}
static bool output_binding(const ggml_ext_execution_context* e, const Device& d, const ggml_tensor* output, const std::vector<ggml_ext_vulkan_binding>& inputs, ggml_ext_vulkan_binding& binding, bool& alias, const ggml_tensor* residual = nullptr) {
binding = {output};
alias = false;
ggml_ext_vulkan_buffer_info dst, src;
if (!d.api->buffer_info(e->backend, &binding, &dst))
return false;
for (const auto& input : inputs) {
if (!d.api->buffer_info(e->backend, &input, &src))
return false;
// The residual epilogue reads each element before writing that same element.
if (residual && input.tensor == residual && dst.storage == src.storage && dst.offset == src.offset && dst.size == src.size)
continue;
alias |= overlap(dst, src);
}
if (alias)
binding = {nullptr, e->workspace, 0, ggml_nbytes(output)};
return true;
}
static ggml_status finish(const ggml_ext_execution_context* e, const Device& d, const ggml_tensor* output, const ggml_ext_vulkan_binding& binding, bool alias, ggml_status status) {
if (status == GGML_STATUS_SUCCESS && alias) {
ggml_ext_vulkan_binding dst = {output};
return d.api->copy(e->backend, &binding, &dst);
}
return status;
}
static ggml_status quantize(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* const* nodes, int count) {
const auto* output = nodes[count - 1];
const auto* first = nodes[0];
const bool modulate = count > 1 && (first->op == GGML_OP_NORM || first->op == GGML_OP_MUL || first->op == GGML_OP_RMS_NORM);
const bool rms = modulate && first->op == GGML_OP_RMS_NORM;
const bool affine = modulate && (rms || first->op == GGML_OP_MUL);
const bool sigmoid = count > 1 && first->op == GGML_OP_UNARY;
const auto* x = count == 1 ? output->src[0] : first->src[0];
const auto* up = count == 1 || modulate ? nullptr : sigmoid ? nodes[1]->src[0]
: first->src[1];
const int extra = rms ? 2 : 0;
const auto* mod = modulate ? nodes[extra + (affine ? 0 : 1)]->src[1] : nullptr;
const auto* shift = modulate ? (affine ? nodes[extra + 2]->src[1] : mod) : nullptr;
const auto* weight = rms ? nodes[1]->src[1] : x;
std::vector<ggml_ext_vulkan_binding> inputs = {{x}};
if (up)
inputs.push_back({up});
if (modulate) {
inputs.push_back({mod});
inputs.push_back({shift});
if (rms)
inputs.push_back({weight});
}
ggml_ext_vulkan_binding dst;
bool alias;
if (!output_binding(e, d, output, inputs, dst, alias))
return GGML_STATUS_FAILED;
const uint32_t k = uint32_t(x->ne[0]), rows = uint32_t(ggml_nrows(x));
const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
auto packed = range(dst, 0, scale_offset), scales = range(dst, scale_offset, size_t(rows) * 4);
const uint32_t groups[] = {std::min(rows, d.info.max_group_count[0]), 1, 1};
ggml_status status;
if (modulate) {
std::vector<ggml_ext_vulkan_binding> bindings = {{x}, {mod}, {shift}, packed, scales};
if (rms)
bindings.push_back({weight});
ModulateParameters pc = {rows, 0, 0, 0, affine && !rms ? 0.f : param_f32(first, 0), k, 0};
uint32_t* offsets[] = {&pc.x_offset, &pc.m_offset, &pc.h_offset, nullptr, nullptr, &pc.w_offset};
for (size_t i = 0; i < bindings.size(); ++i)
if (offsets[i]) {
ggml_ext_vulkan_buffer_info info;
if (!d.api->buffer_info(e->backend, &bindings[i], &info))
return GGML_STATUS_FAILED;
*offsets[i] = info.shader_offset / 4;
}
auto pipeline = rms ? SD_VK_PIPELINE(rms_modulate_convrot, 6, ModulateParameters) : SD_VK_PIPELINE(modulate_convrot, 5, ModulateParameters);
const uint32_t constants[] = {k, uint32_t(affine)};
pipeline.constants = constants;
pipeline.constant_count = 2;
status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
} else {
ggml_ext_vulkan_binding bindings[] = {{x}, packed, scales, {up}};
auto pipeline = SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters);
bool subgroup = d.subgroup && k >= 1024;
uint32_t width = k;
if (subgroup) {
pipeline = SD_VK_PIPELINE(quantize_i8_convrot_subgroup, 3, QuantParameters);
if (!up && ((uint64_t(k) + 2047) / 2048) * 8 <= 64) {
pipeline = SD_VK_PIPELINE(quantize_i8_convrot_register, 3, QuantParameters);
pipeline.constants = &width;
pipeline.constant_count = 1;
}
}
if (up) {
subgroup = true;
pipeline = k == 16384 ? SD_VK_PIPELINE(swiglu_quantize_i8_convrot_16384, 4, QuantParameters) : SD_VK_PIPELINE(swiglu_quantize_i8_convrot, 4, QuantParameters);
if (sigmoid)
pipeline = SD_VK_PIPELINE(sigmoid_quantize_i8_convrot_6144, 4, QuantParameters);
}
if (subgroup) {
pipeline.full_subgroups = d.info.full_subgroups;
pipeline.required_subgroup_size = d.info.subgroup_size_control ? 32 : 0;
}
QuantParameters pc = {k, rows};
status = d.api->dispatch(e->backend, &pipeline, bindings, &pc, groups);
}
return finish(e, d, output, dst, alias, status);
}
static uint32_t group_columns(uint32_t cols, uint32_t k, uint32_t tile_cols, bool swiglu) {
if (!cols || !k || !tile_cols)
return 1;
const uint64_t col_tiles = (uint64_t(cols) + tile_cols - 1) / tile_cols;
const uint64_t weight_tile_bytes = uint64_t(tile_cols) * k * (swiglu ? 2 : 1);
const uint64_t growth = std::clamp(uint64_t(k) / cols, uint64_t(1), uint64_t(2));
return uint32_t(std::clamp(uint64_t(8) * 1024 * 1024 * growth / weight_tile_bytes, uint64_t(1), std::min(col_tiles, uint64_t(32))));
}
static ggml_status multiply(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* const* nodes, int count) {
const auto* output = nodes[count - 1];
const auto* concat = output->op == GGML_OP_CONCAT ? output : nullptr;
const auto* tail = concat ? concat->src[1] : output;
const auto* glu = tail->op == GGML_OP_GLU ? tail : nullptr;
const auto* mat = glu ? glu->src[0] : tail;
const ggml_tensor *gate = nullptr, *residual = nullptr;
float output_scale = 1.f;
if (tail->op == GGML_OP_ADD) {
const auto* mul = tail->src[0]->op == GGML_OP_MUL ? tail->src[0] : tail->src[1];
gate = mul->src[1];
residual = tail->src[0] == mul ? tail->src[1] : tail->src[0];
mat = mul->src[0];
if (mat->op == GGML_OP_SCALE) {
output_scale = param_f32(mat, 0);
mat = mat->src[0];
}
}
const auto* up_mat = glu ? glu->src[1] : nullptr;
SDI8Operation op, up;
if (!sd_ggml_i8_describe(mat, &op) || (up_mat && !sd_ggml_i8_describe(up_mat, &up)))
return GGML_STATUS_FAILED;
const uint32_t k = uint32_t(op.weight->ne[0]), n = uint32_t(op.weight->ne[1]), rows = uint32_t(ggml_nrows(mat));
const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
std::vector<ggml_ext_vulkan_binding> inputs = {{op.weight}, {op.input}, {op.weight_scale}, {op.bias ? op.bias : op.weight_scale}};
if (gate) {
inputs.push_back({gate});
inputs.push_back({residual});
}
if (up_mat) {
inputs.push_back({up.weight});
inputs.push_back({up.weight_scale});
}
if (concat && ggml_nbytes(concat->src[0]))
inputs.push_back({concat->src[0]});
ggml_ext_vulkan_binding dst;
bool alias;
if (!output_binding(e, d, output, inputs, dst, alias, concat ? nullptr : residual))
return GGML_STATUS_FAILED;
auto destination = concat ? range(dst, ggml_nbytes(concat->src[0]), ggml_nbytes(tail)) : dst;
ggml_ext_vulkan_binding bindings[] = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}, {op.input, nullptr, scale_offset, size_t(rows) * 4}, {op.bias ? op.bias : op.weight_scale}, destination, {up_mat ? up.weight : gate}, {up_mat ? up.weight_scale : residual}};
auto dispatch = [&](ggml_ext_vulkan_pipeline pipeline, uint32_t tile_cols, uint32_t tile_rows,
std::initializer_list<uint32_t> constants, uint32_t row_offset, uint32_t dispatch_rows) {
if (!dispatch_rows)
return GGML_STATUS_SUCCESS;
const uint32_t cols = glu ? uint32_t(glu->ne[0]) : n;
MatParameters pc = {n * (up_mat ? 2u : 1u), rows, k, op.bias ? 1u : 0u, row_offset, output_scale, group_columns(cols, k, tile_cols, glu != nullptr)};
pipeline.constants = constants.begin();
pipeline.constant_count = uint32_t(constants.size());
const uint32_t groups[] = {(cols + tile_cols - 1) / tile_cols, (dispatch_rows + tile_rows - 1) / tile_rows, 1};
return d.api->dispatch(e->backend, &pipeline, bindings, &pc, groups);
};
ggml_status status = GGML_STATUS_SUCCESS;
if (d.cm2) {
const bool large_k = k % 4096 == 0 || (rows >= n && k >= n);
const uint32_t full_rows = (glu || n % 128 == 0) ? rows / 128 * 128 : 0;
auto pipeline = residual ? SD_VK_PIPELINE(mul_mat_i8_residual_cm2, 8, MatParameters) : up_mat ? SD_VK_PIPELINE(mul_mat_i8_swiglu_split_cm2, 8, MatParameters)
: SD_VK_PIPELINE(mul_mat_i8_tensorwise_cm2, 6, MatParameters);
const uint32_t tile_cols = glu ? 64 : 128;
status = dispatch(pipeline, tile_cols, 128, {128, 1, glu ? 1u : 0u, tile_cols, glu || large_k ? 128u : 64u}, 0, full_rows);
if (status == GGML_STATUS_SUCCESS) {
const bool small = d.cm2_small && (glu || n % 128 == 0) && rows - full_rows <= 32;
const uint32_t tail_cols = small || glu ? 64 : 128, tail_rows = small ? 32 : glu ? 128
: 64;
status = dispatch(pipeline, tail_cols, tail_rows, {tail_rows, small ? 2u : 0u, glu ? 1u : 0u, tail_cols, small ? 256u : 64u}, full_rows, rows - full_rows);
}
} else {
uint32_t coop_rows = 0;
if (d.cm_m && k % d.cm_k == 0 && n % (4 * d.cm_n) == 0) {
const uint32_t tile_rows = 4 * d.cm_m;
coop_rows = rows / tile_rows * tile_rows;
auto pipeline = SD_VK_PIPELINE(mul_mat_i8_tensorwise_cm1, 6, MatParameters);
pipeline.full_subgroups = d.info.full_subgroups;
pipeline.required_subgroup_size = d.info.subgroup_size_control ? d.info.subgroup_size : 0;
status = dispatch(pipeline, 4 * d.cm_n, tile_rows, {d.info.subgroup_size, d.cm_m, d.cm_n, d.cm_k}, 0, coop_rows);
}
if (status == GGML_STATUS_SUCCESS)
status = dispatch(SD_VK_PIPELINE(mul_mat_i8_tensorwise, 6, MatParameters), 16, 16, {}, coop_rows, rows - coop_rows);
}
if (status == GGML_STATUS_SUCCESS && concat && ggml_nbytes(concat->src[0])) {
ggml_ext_vulkan_binding src = {concat->src[0]}, prefix = range(dst, 0, ggml_nbytes(concat->src[0]));
status = d.api->copy(e->backend, &src, &prefix);
}
return finish(e, d, output, dst, alias, status);
}
static bool quant_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
Device d;
return device_info(dev, d) && can_quantize(d, n);
}
static bool mat_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
Device d;
return device_info(dev, d) && can_matmul(d, n);
}
static int match(ggml_backend_dev_t dev, const ggml_cgraph* graph, int first, void*) {
Device d;
if (!device_info(dev, d) || d.info.fusion_disabled)
return 0;
if (int count = can_fuse_modulate_convrot(d, graph, first))
return count;
if (int count = can_fuse_i8_row_concat(d, graph, first))
return count;
if (int count = can_fuse_i8_residual(d, graph, first))
return count;
if (can_fuse_i8_swiglu_split(d, graph, first))
return 3;
if (can_fuse_i8_swiglu(d, graph, first))
return 2;
if (can_fuse_sigmoid_convrot(d, graph, first))
return 3;
if (can_fuse_swiglu_convrot(d, graph, first))
return 2;
return 0;
}
static void optimize(ggml_backend_dev_t dev, ggml_cgraph* graph, void*) {
Device d;
if (!device_info(dev, d) || d.info.fusion_disabled || d.compute.graph_optimization_disabled)
return;
for (int last = 1; last < count_nodes(graph); ++last) {
int indices[6];
int count = match_rms_modulate_convrot(d, graph, last, indices);
if (!count)
count = match_modulate_convrot(d, graph, last, indices);
if (!count)
count = match_i8_row_concat(d, graph, last, indices);
if (!count)
count = match_i8_residual(d, graph, last, indices);
if (!count && match_i8_swiglu_split(d, graph, last, indices))
count = 3;
if (!count)
continue;
int first = indices[0];
for (int j = 1; j < count; ++j) {
if (!ggml_ext_graph_delay(graph, first, j, indices[j]))
break;
first = indices[j] - j;
}
}
}
static bool may_alias(const ggml_tensor* output, const ggml_tensor* input) {
size_t offset = 0;
const size_t size = ggml_nbytes(input);
while (!input->data && input->view_src) {
offset += input->view_offs;
input = input->view_src;
}
// Unallocated inputs get allocator storage, separate from a preallocated output.
if (!input->data)
return false;
if (!input->buffer)
return true;
if (output->buffer != input->buffer)
return false;
const auto dst = reinterpret_cast<uintptr_t>(output->data);
const auto src = reinterpret_cast<uintptr_t>(input->data) + offset;
return dst < src + size && src < dst + ggml_nbytes(output);
}
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) {
const auto* output = nodes[count - 1];
// Fresh outputs cannot alias external inputs except the residual's elementwise reuse.
const bool fresh = !output->data && !output->view_src &&
(output->op == GGML_OP_CUSTOM || output->op == GGML_OP_GLU || output->op == GGML_OP_CONCAT ||
(output->op == GGML_OP_ADD && count > 1));
*bytes = fresh ? 0 : ggml_nbytes(output);
if (output->data && output->buffer && !output->view_src) {
bool alias = false;
for (int i = 0; i < count; ++i)
for (const auto* src : nodes[i]->src)
if (src && std::find(nodes, nodes + count, src) == nodes + count)
alias |= may_alias(output, src);
if (!alias)
*bytes = 0;
}
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}
static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int count, void*, void*) {
Device d;
if (!device_info(ggml_backend_get_device(e->backend), d))
return GGML_STATUS_FAILED;
return is_convrot(nodes[count - 1]) ? quantize(e, d, nodes, count) : multiply(e, d, nodes, count);
}
#undef SD_VK_PIPELINE
}
void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) {
using namespace sd_vulkan_i8;
Device info;
if (!device_info(device, info))
return;
ggml_ext_kernel_desc d = {};
d.abi_version = GGML_EXT_ABI_VERSION;
d.priority = 100;
d.prepare = prepare;
d.destroy = destroy;
d.execute = execute;
d.name = "sd.quantize_i8_convrot.vulkan.v1";
d.op = GGML_OP_CUSTOM;
d.custom_op = sd_ggml_convrot_op();
d.supports = quant_supports;
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.name = "sd.linear_i8.vulkan.v1";
d.op = GGML_OP_CUSTOM;
d.custom_op = sd_ggml_i8_op();
d.supports = mat_supports;
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.priority = 200;
d.supports = nullptr;
d.match = match;
d.optimize = optimize;
d.name = "sd.linear_i8.fused.vulkan.v1";
d.op = GGML_OP_CUSTOM;
d.custom_op = sd_ggml_i8_op();
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.optimize = nullptr;
d.custom_op = nullptr;
const struct {
ggml_op op;
const char* name;
} fusions[] = {
{GGML_OP_NORM, "sd.norm_convrot.vulkan.v1"},
{GGML_OP_RMS_NORM, "sd.rms_modulate_convrot.vulkan.v1"},
{GGML_OP_MUL, "sd.modulate_convrot.vulkan.v1"},
{GGML_OP_GLU, "sd.swiglu_convrot.vulkan.v1"},
{GGML_OP_UNARY, "sd.sigmoid_convrot.vulkan.v1"}};
for (const auto& fusion : fusions) {
d.op = fusion.op;
d.name = fusion.name;
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
}
}
Reference in new issue View Git Blame Copy Permalink