diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt index a661e710a2..63e20970a6 100644 --- a/ggml/src/ggml-metal/CMakeLists.txt +++ b/ggml/src/ggml-metal/CMakeLists.txt @@ -10,6 +10,7 @@ ggml_add_backend_library(ggml-metal ggml-metal-device.cpp ggml-metal-common.cpp ggml-metal-context.m + ggml-metal-fuse.cpp ggml-metal-ops.cpp ggml-metal-tuning.cpp ) diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp index 6f1638a114..02983f3d83 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.cpp +++ b/ggml/src/ggml-metal/ggml-metal-common.cpp @@ -1,4 +1,5 @@ #include "ggml-metal-common.h" +#include "ggml-metal-fuse.h" #include "ggml.h" #include "ggml-impl.h" @@ -390,59 +391,31 @@ static std::vector ggml_metal_graph_optimize_reorder(const std::vectorn_nodes; - enum ggml_op ops[MAX_FUSE]; - std::vector nodes; nodes.reserve(gf->n_nodes); // fuse nodes: // we don't want to make reorders that break fusing, so we first pack all fusable tensors // and perform the reorder over the fused nodes. after the reorder is done, we unfuse + // + // the fusable sequences are declared in the fusion table (ggml-metal-fuse.cpp), so the + // packing here is driven by the same patterns that the op encoders will later use for (int i = 0; i < n; i++) { node_info node = { /*.node =*/ gf->nodes[i], /*.fused =*/ {}, }; - // fuse only ops that start with these operations - // can be expanded when needed - if (node.op() == GGML_OP_ADD || - node.op() == GGML_OP_NORM || - node.op() == GGML_OP_RMS_NORM) { - ops[0] = node.op(); + const int f = ggml_metal_fuse_max(gf, i); - int f = i + 1; - while (f < n && f < i + MAX_FUSE) { - // conservatively allow fusing only these ops - // can be expanded when needed - if (gf->nodes[f]->op != GGML_OP_ADD && - gf->nodes[f]->op != GGML_OP_MUL && - gf->nodes[f]->op != GGML_OP_NORM && - gf->nodes[f]->op != GGML_OP_RMS_NORM) { - break; - } - ops[f - i] = gf->nodes[f]->op; - f++; - } + // add the fused tensors into the node info so we can unfuse them later + for (int k = 1; k < f; k++) { + ++i; - f -= i; - for (; f > 1; f--) { - if (ggml_can_fuse(gf, i, ops, f)) { - break; - } - } - - // add the fused tensors into the node info so we can unfuse them later - for (int k = 1; k < f; k++) { - ++i; - - // the .dst() becomes the last fused tensor - node.add_fused(gf->nodes[i]); - } + // the .dst() becomes the last fused tensor + node.add_fused(gf->nodes[i]); } nodes.push_back(std::move(node)); diff --git a/ggml/src/ggml-metal/ggml-metal-fuse.cpp b/ggml/src/ggml-metal/ggml-metal-fuse.cpp new file mode 100644 index 0000000000..d4c065e49d --- /dev/null +++ b/ggml/src/ggml-metal/ggml-metal-fuse.cpp @@ -0,0 +1,275 @@ +#include "ggml-metal-fuse.h" + +#include "ggml-backend-impl.h" +#include "ggml-metal-device.h" + +#include + +// ---- helpers ------------------------------------------------------------- + +// the pattern outputs, defaulting to the last node +static const int * ggml_metal_fuse_outputs(const struct ggml_metal_fuse * fuse, int * buf) { + if (fuse->outputs) { + return fuse->outputs; + } + + buf[0] = fuse->n_ops - 1; + + return buf; +} + +// true if two tensors live in the same Metal buffer +static bool ggml_metal_fuse_same_buffer(const struct ggml_tensor * a, const struct ggml_tensor * b) { + if (!a || !b) { + return false; + } + + ggml_backend_buffer_t ba = a->view_src ? a->view_src->buffer : a->buffer; + ggml_backend_buffer_t bb = b->view_src ? b->view_src->buffer : b->buffer; + + ggml_metal_buffer_t ca = (ggml_metal_buffer_t) ba->context; + ggml_metal_buffer_t cb = (ggml_metal_buffer_t) bb->context; + + return ggml_metal_buffer_get_id(ca, a).metal == ggml_metal_buffer_get_id(cb, b).metal; +} + +// ---- pattern checks ------------------------------------------------------ + +// NORM/RMS_NORM + MUL + ADD: the weight/bias of each fused step must match the norm input +// width, be contiguous rows, and the fused outputs must stay F32 +static bool ggml_metal_fuse_check_norm(const struct ggml_tensor * const * nodes, + const struct ggml_metal_fuse * fuse, + enum ggml_metal_fuse_mode mode) { + GGML_ASSERT(fuse->n_ops >= 2); + + for (int j = 1; j < fuse->n_ops; j++) { + // the fused MUL/ADD must read the previous node as src0 + if (nodes[j]->src[0] != nodes[j - 1]) { + return false; + } + + // the weight/bias must have the same row width as the norm input + if (nodes[j]->src[1]->ne[0] != nodes[0]->ne[0]) { + return false; + } + + if (!ggml_is_contiguous_rows(nodes[j]->src[1])) { + return false; + } + + if (nodes[j]->type != GGML_TYPE_F32) { + return false; + } + } + + return true; +} + +// ADD x N: each ADD reads the previous ADD as src0, and all addends must share layout +// (and, in FULL mode, live in the same Metal buffer) +static bool ggml_metal_fuse_check_add_chain(const struct ggml_tensor * const * nodes, + const struct ggml_metal_fuse * fuse, + enum ggml_metal_fuse_mode mode) { + GGML_ASSERT(fuse->n_ops >= 2); + + for (int j = 1; j < fuse->n_ops; j++) { + if (nodes[j]->src[0] != nodes[j - 1]) { + return false; + } + + if (!ggml_are_same_layout(nodes[j]->src[1], nodes[j - 1]->src[1])) { + return false; + } + + if (mode == GGML_METAL_FUSE_FULL) { + if (!ggml_metal_fuse_same_buffer(nodes[j]->src[1], nodes[0]->src[1])) { + return false; + } + } + } + + return true; +} + +// MUL + SIN + SQR + MUL + ADD (snake activation) +static bool ggml_metal_fuse_check_snake(const struct ggml_tensor * const * nodes, + const struct ggml_metal_fuse * fuse, + enum ggml_metal_fuse_mode mode) { + const struct ggml_tensor * mul0 = nodes[0]; + const struct ggml_tensor * sin_node = nodes[1]; + const struct ggml_tensor * sqr = nodes[2]; + const struct ggml_tensor * mul1 = nodes[3]; + const struct ggml_tensor * add = nodes[4]; + + // x carries the full activation shape, a is the broadcast operand + const struct ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; + const struct ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; + + // mul1 reads sqr and inv_b in either operand order + const struct ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; + + // closure check: the trailing add reads the same x as the leading mul + const struct ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0]; + + // x is in the supported whitelist and every chain intermediate shares x's type. + // a and inv_b bind as device const float * in the kernel, so they stay F32. + const bool types_ok = + (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && + (a->type == GGML_TYPE_F32) && (inv_b->type == GGML_TYPE_F32) && + (mul0->type == x->type) && (sin_node->type == x->type) && + (sqr->type == x->type) && (mul1->type == x->type) && + (add->type == x->type); + // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D + const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1]; + const bool dim_ok = + (x->ne[2] == 1) && (x->ne[3] == 1) && + (add->ne[2] == 1) && (add->ne[3] == 1) && + (a->ne[2] == 1) && (a->ne[3] == 1) && + (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1); + // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous + const bool contig_ok = + ggml_is_contiguous(x) && ggml_is_contiguous(add) && + ggml_is_contiguous(a) && ggml_is_contiguous(inv_b); + + return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x; +} + +// ---- patterns ------------------------------------------------------------ + +static const enum ggml_op ops_norm_mul[] = { GGML_OP_NORM, GGML_OP_MUL }; +static const enum ggml_op ops_norm_mul_add[] = { GGML_OP_NORM, GGML_OP_MUL, GGML_OP_ADD }; +static const enum ggml_op ops_rms_norm_mul[] = { GGML_OP_RMS_NORM, GGML_OP_MUL }; +static const enum ggml_op ops_rms_norm_mul_add[] = { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD }; +static const enum ggml_op ops_add_2[] = { GGML_OP_ADD, GGML_OP_ADD }; +static const enum ggml_op ops_add_3[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const enum ggml_op ops_add_4[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const enum ggml_op ops_add_5[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const enum ggml_op ops_add_6[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const enum ggml_op ops_add_7[] = { GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_ADD }; +static const enum ggml_op ops_snake[] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }; + +static const struct ggml_metal_fuse ggml_metal_fuses[] = { + { GGML_METAL_FUSE_NORM_MUL, ops_norm_mul, 2, nullptr, 0, ggml_metal_fuse_check_norm }, + { GGML_METAL_FUSE_NORM_MUL_ADD, ops_norm_mul_add, 3, nullptr, 0, ggml_metal_fuse_check_norm }, + { GGML_METAL_FUSE_NORM_MUL, ops_rms_norm_mul, 2, nullptr, 0, ggml_metal_fuse_check_norm }, + { GGML_METAL_FUSE_NORM_MUL_ADD, ops_rms_norm_mul_add, 3, nullptr, 0, ggml_metal_fuse_check_norm }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_2, 2, nullptr, 0, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_3, 3, nullptr, 0, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_4, 4, nullptr, 0, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_5, 5, nullptr, 0, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_6, 6, nullptr, 0, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_ADD_CHAIN, ops_add_7, 7, nullptr, 0, ggml_metal_fuse_check_add_chain }, + { GGML_METAL_FUSE_SNAKE, ops_snake, 5, nullptr, 0, ggml_metal_fuse_check_snake }, +}; + +const struct ggml_metal_fuse * ggml_metal_fuse_all(int * n) { + *n = (int) sizeof(ggml_metal_fuses) / sizeof(ggml_metal_fuses[0]); + + return ggml_metal_fuses; +} + +// ---- queries ------------------------------------------------------------- + +// find the longest pattern matching the node sequence starting at idx +// (idx is a position in node_idxs, which maps to graph node indices) +const struct ggml_metal_fuse * ggml_metal_fuse_next( + const struct ggml_cgraph * gf, + const int * node_idxs, + int n_idxs, + int idx, + enum ggml_metal_fuse_mode mode, + int * n_out) { + int n = 0; + const struct ggml_metal_fuse * all = ggml_metal_fuse_all(&n); + + const struct ggml_metal_fuse * res = nullptr; + int best = 1; + + for (int i = 0; i < n; i++) { + const struct ggml_metal_fuse * fuse = &all[i]; + + // only look for a longer match than the current best + if (fuse->n_ops <= best) { + continue; + } + if (idx + fuse->n_ops > n_idxs) { + continue; + } + + // the first node must match the pattern start + if (gf->nodes[node_idxs[idx]]->op != fuse->ops[0]) { + continue; + } + + const struct ggml_tensor * nodes[GGML_METAL_FUSE_MAX]; + + // common element-wise chain constraints: each node reads the previous one, + // and all nodes have the same shape + bool ok = true; + for (int j = 0; j < fuse->n_ops; j++) { + nodes[j] = gf->nodes[node_idxs[idx + j]]; + + if (j > 0) { + if (nodes[j]->src[0] != nodes[j - 1] && nodes[j]->src[1] != nodes[j - 1]) { + ok = false; + break; + } + if (!ggml_are_same_shape(nodes[j], nodes[j - 1])) { + ok = false; + break; + } + } + } + if (!ok) { + continue; + } + + int outputs_buf[GGML_MAX_SRC]; + const int * outputs = ggml_metal_fuse_outputs(fuse, outputs_buf); + const int n_outputs = fuse->n_outputs ? fuse->n_outputs : 1; + + // structural subgraph checks (op sequence, elidable uses, view containment) + if (!ggml_can_fuse_subgraph_ext(gf, node_idxs + idx, fuse->n_ops, fuse->ops, outputs, n_outputs)) { + continue; + } + + // pattern-specific checks + if (fuse->check && !fuse->check(nodes, fuse, mode)) { + continue; + } + + best = fuse->n_ops; + res = fuse; + } + + *n_out = best; + + return res; +} + +// optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that +// could be fused, chaining patterns back-to-back +int ggml_metal_fuse_max(const struct ggml_cgraph * gf, int idx) { + int idxs[GGML_METAL_FUSE_MAX]; + + int total = 0; + int i = idx; + + while (i < gf->n_nodes && total < GGML_METAL_FUSE_MAX) { + const int n_idxs = std::min(GGML_METAL_FUSE_MAX, (int) gf->n_nodes - i); + for (int j = 0; j < n_idxs; j++) { + idxs[j] = i + j; + } + + int len = 1; + const struct ggml_metal_fuse * fuse = ggml_metal_fuse_next(gf, idxs, n_idxs, 0, GGML_METAL_FUSE_STRUCTURAL, &len); + if (!fuse || total + len > GGML_METAL_FUSE_MAX) { + break; + } + + total += len; + i += len; + } + + return std::max(total, 1); +} diff --git a/ggml/src/ggml-metal/ggml-metal-fuse.h b/ggml/src/ggml-metal/ggml-metal-fuse.h new file mode 100644 index 0000000000..63ab0847fd --- /dev/null +++ b/ggml/src/ggml-metal/ggml-metal-fuse.h @@ -0,0 +1,71 @@ +// single source of truth for the fusions supported by the Metal backend +// +// every fusable subgraph is declared exactly once as a ggml_metal_fuse entry in +// the table in ggml-metal-fuse.cpp. both the graph optimizer (ggml_metal_fuse_max) +// and the op encoders (ggml_metal_fuse_next) consult this same table, so the two +// phases can never disagree about what can be fused. + +#pragma once + +#include "ggml-impl.h" + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// the maximum number of nodes that can be fused in a single kernel +// (also the maximum length of a packed fusion group during graph optimization) +#define GGML_METAL_FUSE_MAX 16 + +enum ggml_metal_fuse_mode { + // structural checks only; used by the graph optimizer, at which point the graph + // tensors are not allocated yet, so buffer placement cannot be verified + GGML_METAL_FUSE_STRUCTURAL = 0, + // full checks, including buffer placement; used by the op encoders + GGML_METAL_FUSE_FULL, +}; + +// identifier of each fusion pattern so the op encoders know which kernel to use +enum ggml_metal_fuse_id { + GGML_METAL_FUSE_NONE = 0, + GGML_METAL_FUSE_NORM_MUL, // NORM/RMS_NORM + MUL + GGML_METAL_FUSE_NORM_MUL_ADD, // NORM/RMS_NORM + MUL + ADD + GGML_METAL_FUSE_ADD_CHAIN, // ADD x N (N in [2, 7]) + GGML_METAL_FUSE_SNAKE, // MUL + SIN + SQR + MUL + ADD +}; + +struct ggml_metal_fuse { + enum ggml_metal_fuse_id id; + const enum ggml_op * ops; // op sequence (fixed length) + int n_ops; // number of ops + const int * outputs; // output node indices relative to the subgraph start (nullptr => { n_ops-1 }) + int n_outputs;// number of outputs (0 => default last node) + // extra backend constraints on top of ggml_can_fuse_subgraph + // nodes[j] is the j-th node of the pattern + bool (*check)(const struct ggml_tensor * const * nodes, + const struct ggml_metal_fuse * fuse, + enum ggml_metal_fuse_mode mode); +}; + +// the single table of all fusions supported by the Metal backend +const struct ggml_metal_fuse * ggml_metal_fuse_all(int * n); + +// compute phase: longest fusion starting at idx (a position in node_idxs) that matches in `mode`. +// returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes consumed. +const struct ggml_metal_fuse * ggml_metal_fuse_next( + const struct ggml_cgraph * gf, + const int * node_idxs, + int n_idxs, + int idx, + enum ggml_metal_fuse_mode mode, + int * n_out); + +// optimize phase: maximum number of nodes starting at idx (a raw sequential graph index) that +// could be fused, chaining patterns back-to-back. returns at least 1. +int ggml_metal_fuse_max(const struct ggml_cgraph * gf, int idx); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 3db8bca437..e5cc36bc10 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -7,6 +7,7 @@ #include "ggml-metal-impl.h" #include "ggml-metal-common.h" #include "ggml-metal-device.h" +#include "ggml-metal-fuse.h" #include "ggml-metal-tuning.h" #include @@ -78,15 +79,13 @@ struct ggml_metal_op { return ggml_graph_node(gf, idxs[i]); } - bool can_fuse(int i0, const ggml_op * ops, int n_ops) const { + // consult the fusion table for the longest pattern starting at i0 + // returns the matching pattern (nullptr if no fusion) and sets *n_out to the number of nodes + const ggml_metal_fuse * can_fuse(int i0, enum ggml_metal_fuse_mode mode, int * n_out) const { assert(use_fusion); assert(i0 >= 0 && i0 < n_nodes()); - if (i0 + n_ops > n_nodes()) { - return false; - } - - return ggml_can_fuse_ext(gf, idxs.data() + i0, ops, n_ops); + return ggml_metal_fuse_next(gf, idxs.data(), (int) idxs.size(), i0, mode, n_out); } ggml_metal_device_t dev; @@ -3718,56 +3717,19 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { return 1; } -// Snake activation autofuse: mul -> sin -> sqr -> mul -> add -static bool ggml_metal_op_can_fuse_snake(ggml_metal_op_t ctx, int idx) { - static constexpr ggml_op snake_ops[5] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }; - - if (ctx->node(idx)->op != GGML_OP_MUL || !ctx->can_fuse(idx, snake_ops, 5)) { - return false; - } - - const ggml_tensor * mul0 = ctx->node(idx + 0); - const ggml_tensor * sin_node = ctx->node(idx + 1); - const ggml_tensor * sqr = ctx->node(idx + 2); - const ggml_tensor * mul1 = ctx->node(idx + 3); - const ggml_tensor * add = ctx->node(idx + 4); - - // x carries the full activation shape, a is the broadcast operand - const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; - const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; - - // mul1 reads sqr and inv_b in either operand order - const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; - - // closure check: the trailing add reads the same x as the leading mul - const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0]; - - // x is in the supported whitelist and every chain intermediate shares x's type. - // a and inv_b bind as device const float * in the kernel, so they stay F32. - const bool types_ok = - (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && - (a->type == GGML_TYPE_F32) && (inv_b->type == GGML_TYPE_F32) && - (mul0->type == x->type) && (sin_node->type == x->type) && - (sqr->type == x->type) && (mul1->type == x->type) && - (add->type == x->type); - // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D - const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1]; - const bool dim_ok = - (x->ne[2] == 1) && (x->ne[3] == 1) && - (add->ne[2] == 1) && (add->ne[3] == 1) && - (a->ne[2] == 1) && (a->ne[3] == 1) && - (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1); - // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous - const bool contig_ok = - ggml_is_contiguous(x) && ggml_is_contiguous(add) && - ggml_is_contiguous(a) && ggml_is_contiguous(inv_b); - - return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x; -} - int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { - if (ctx->use_fusion && ggml_metal_op_can_fuse_snake(ctx, idx)) { - return ggml_metal_op_snake_fused(ctx, idx); + int n_fuse = 1; + const ggml_metal_fuse * fuse = nullptr; + + if (ctx->use_fusion) { + int n = 1; + fuse = ctx->can_fuse(idx, GGML_METAL_FUSE_FULL, &n); + n_fuse = n; + + // snake activation autofuse: mul -> sin -> sqr -> mul -> add + if (fuse && fuse->id == GGML_METAL_FUSE_SNAKE) { + return ggml_metal_op_snake_fused(ctx, idx); + } } ggml_tensor * op = ctx->node(idx); @@ -3822,57 +3784,17 @@ int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { /*.o1 =*/ { bid_src1.offs }, }; - ggml_op fops[8]; - - int n_fuse = 1; - // c[0] = add(a, b[0]) // c[1] = add(c[0], b[1]) // c[2] = add(c[1], b[2]) // ... - if (use_fusion) { - fops[0] = GGML_OP_ADD; - fops[1] = GGML_OP_ADD; - fops[2] = GGML_OP_ADD; - fops[3] = GGML_OP_ADD; - fops[4] = GGML_OP_ADD; - fops[5] = GGML_OP_ADD; - fops[6] = GGML_OP_ADD; - fops[7] = GGML_OP_ADD; - - // note: in metal, we sometimes encode the graph in parallel so we have to avoid fusing ops - // across splits. idx_end indicates the last node in the current split - for (n_fuse = 0; n_fuse <= 6; ++n_fuse) { - if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) { - break; - } - - ggml_tensor * f0 = ctx->node(idx + n_fuse); - ggml_tensor * f1 = ctx->node(idx + n_fuse + 1); - - if (f0 != f1->src[0]) { - break; - } - - // b[0] === b[1] === ... - if (!ggml_are_same_layout(f0->src[1], f1->src[1])) { - break; - } - - // only fuse ops if src1 is in the same Metal buffer - ggml_metal_buffer_id bid_fuse = ggml_metal_get_buffer_id(f1->src[1]); - if (bid_fuse.metal != bid_src1.metal) { - break; - } - - //ctx->fuse_cnt[ops[n_fuse + 1]->op]++; - - args.o1[n_fuse + 1] = bid_fuse.offs; + if (use_fusion && fuse && fuse->id == GGML_METAL_FUSE_ADD_CHAIN) { + // the offsets of the fused addends are relative to the start of the src1 buffer + for (int i = 1; i < n_fuse; i++) { + args.o1[i] = ggml_metal_get_buffer_id(ctx->node(idx + i)->src[1]).offs; } - ++n_fuse; - - if (debug_fusion > 1 && n_fuse > 1) { + if (debug_fusion > 1) { GGML_LOG_DEBUG("%s: fuse: ADD x %d\n", __func__, n_fuse); } } @@ -4110,8 +4032,6 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) { /*.nbf3 =*/ { nb03 }, }; - ggml_op fops[8]; - int n_fuse = 1; ggml_metal_buffer_id bid_fuse[2] = { bid_src0, bid_src0 }; @@ -4120,55 +4040,33 @@ int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) { // d[1] = mul(d[0], b) // d[2] = add(d[1], c) if (use_fusion) { - fops[0] = op->op; - fops[1] = GGML_OP_MUL; - fops[2] = GGML_OP_ADD; + int n = 1; + const ggml_metal_fuse * fuse = ctx->can_fuse(idx, GGML_METAL_FUSE_FULL, &n); - for (n_fuse = 0; n_fuse <= 1; ++n_fuse) { - if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) { - break; + if (fuse && (fuse->id == GGML_METAL_FUSE_NORM_MUL || fuse->id == GGML_METAL_FUSE_NORM_MUL_ADD)) { + n_fuse = n; + + for (int i = 1; i < n_fuse; i++) { + const ggml_tensor * fn = ctx->node(idx + i); + + bid_fuse[i - 1] = ggml_metal_get_buffer_id(fn->src[1]); + + args.nef1[i] = fn->src[1]->ne[1]; + args.nef2[i] = fn->src[1]->ne[2]; + args.nef3[i] = fn->src[1]->ne[3]; + + args.nbf1[i] = fn->src[1]->nb[1]; + args.nbf2[i] = fn->src[1]->nb[2]; + args.nbf3[i] = fn->src[1]->nb[3]; } - ggml_tensor * f0 = ctx->node(idx + n_fuse); - ggml_tensor * f1 = ctx->node(idx + n_fuse + 1); - - if (f0 != f1->src[0]) { - break; - } - - if (f1->src[1]->ne[0] != op->ne[0]) { - break; - } - - if (!ggml_is_contiguous_rows(f1->src[1])) { - break; - } - - if (f1->type != GGML_TYPE_F32) { - break; - } - - //ctx->fuse_cnt[f1->op]++; - - bid_fuse[n_fuse] = ggml_metal_get_buffer_id(f1->src[1]); - - args.nef1[n_fuse + 1] = f1->src[1]->ne[1]; - args.nef2[n_fuse + 1] = f1->src[1]->ne[2]; - args.nef3[n_fuse + 1] = f1->src[1]->ne[3]; - - args.nbf1[n_fuse + 1] = f1->src[1]->nb[1]; - args.nbf2[n_fuse + 1] = f1->src[1]->nb[2]; - args.nbf3[n_fuse + 1] = f1->src[1]->nb[3]; - } - - ++n_fuse; - - if (debug_fusion > 1 && n_fuse > 1) { - if (n_fuse == 2) { - GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op)); - } - if (n_fuse == 3) { - GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op)); + if (debug_fusion > 1) { + if (n_fuse == 2) { + GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op)); + } + if (n_fuse == 3) { + GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op)); + } } } }