mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-04 02:37:27 +02:00
metal : rework fusion patterns into a single table
All fusable op patterns for the Metal backend are now declared once in a fusion table (ggml-metal-fuse.cpp) and consumed by both the graph optimizer (ggml_metal_fuse_max, packing) and the op encoders (ggml_metal_fuse_next, compute). The two phases share the same pattern table plus ggml_can_fuse_subgraph_ext for the structural checks, and differ only in the mode used for the pattern check (STRUCTURAL at optimize time, since tensors are not allocated yet, and FULL at compute time, including Metal buffer placement). This also protects the snake activation (MUL + SIN + SQR + MUL + ADD) from being reordered during graph optimization, which was previously unprotected. Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-0731
This commit is contained in:
Georgi Gerganov
@@ -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
|
||||
)
|
||||
|
||||
@@ -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<int> ggml_metal_graph_optimize_reorder(const std::vector<node
|
||||
}
|
||||
|
||||
void ggml_graph_optimize(ggml_cgraph * gf) {
|
||||
constexpr int MAX_FUSE = 16;
|
||||
|
||||
const int n = gf->n_nodes;
|
||||
|
||||
enum ggml_op ops[MAX_FUSE];
|
||||
|
||||
std::vector<node_info> 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));
|
||||
|
||||
@@ -0,0 +1,275 @@
|
||||
#include "ggml-metal-fuse.h"
|
||||
|
||||
#include "ggml-backend-impl.h"
|
||||
#include "ggml-metal-device.h"
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
// ---- 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);
|
||||
}
|
||||
@@ -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 <stdbool.h>
|
||||
|
||||
#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
|
||||
@@ -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 <cassert>
|
||||
@@ -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));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user