Skip to content

Commit 4926419

Browse files
am17anam17an
authored
ggml: add ggml_can_fuse_subgraph (ggml-org#16662)
* ggml: add ggml_can_fuse_subgraph * ggml-cuda: use ggml_can_fuse_subgraph for topk-moe * format * 1. remove inputs from signature as they are transient nodes 2. add check for views: view_src should be part of the subgraph * - combine check into one loop - check all view_src parents - other minor review comments * remove redudant if test * - rename and other minor review comments * add assert about count < 32
1 parent 6ea37f5 commit 4926419

File tree

3 files changed

+113
-19
lines changed

3 files changed

+113
-19
lines changed

‎ggml/src/ggml-cuda/ggml-cuda.cu‎

Lines changed: 4 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2821,15 +2821,8 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, int node_idx,
28212821
std::initializer_list<enum ggml_op> topk_moe_ops = ggml_cuda_topk_moe_ops(false);
28222822
std::initializer_list<enum ggml_op> topk_moe_ops_with_norm = ggml_cuda_topk_moe_ops(true);
28232823

2824-
if (ops.size() == topk_moe_ops_with_norm.size() && std::equal(ops.begin(), ops.end(), topk_moe_ops_with_norm.begin())) {
2825-
2826-
if (node_idx + topk_moe_ops_with_norm.size() > (size_t)cgraph->n_nodes) {
2827-
return false;
2828-
}
2829-
2830-
for (size_t i = 0; i < topk_moe_ops_with_norm.size(); i++) {
2831-
if (cgraph->nodes[node_idx + i]->op != topk_moe_ops_with_norm.begin()[i]) return false;
2832-
}
2824+
if (ops.size() == topk_moe_ops_with_norm.size() &&
2825+
ggml_can_fuse_subgraph(cgraph, node_idx, topk_moe_ops_with_norm, { node_idx + 3, node_idx + 8 })) {
28332826
ggml_tensor * softmax = cgraph->nodes[node_idx];
28342827
ggml_tensor * weights = cgraph->nodes[node_idx+8];
28352828

@@ -2838,16 +2831,8 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, int node_idx,
28382831
}
28392832
}
28402833

2841-
if (ops.size() == topk_moe_ops.size() && std::equal(ops.begin(), ops.end(), topk_moe_ops.begin())) {
2842-
2843-
if (node_idx + topk_moe_ops.size() > (size_t)cgraph->n_nodes) {
2844-
return false;
2845-
}
2846-
2847-
for (size_t i = 0; i < topk_moe_ops.size(); i++) {
2848-
if (cgraph->nodes[node_idx + i]->op != topk_moe_ops.begin()[i]) return false;
2849-
}
2850-
2834+
if (ops.size() == topk_moe_ops.size() &&
2835+
ggml_can_fuse_subgraph(cgraph, node_idx, topk_moe_ops, { node_idx + 3, node_idx + 4 })) {
28512836
ggml_tensor * softmax = cgraph->nodes[node_idx];
28522837
ggml_tensor * weights = cgraph->nodes[node_idx+4];
28532838
if (ggml_cuda_should_use_topk_moe(softmax, weights)) {

‎ggml/src/ggml-impl.h‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -647,6 +647,36 @@ static inline bool ggml_can_fuse(const struct ggml_cgraph * cgraph, int node_idx
647647
return ggml_can_fuse_ext(cgraph, idxs, ops, num_ops);
648648
}
649649

650+
GGML_API bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph,
651+
const int * node_idxs,
652+
int count,
653+
const enum ggml_op * ops,
654+
const int * outputs,
655+
int num_outputs);
656+
657+
// Returns true if the subgraph formed by {node_idxs} can be fused
658+
// checks whethers all nodes which are not part of outputs can be elided
659+
// by checking if their num_uses are confined to the subgraph
660+
static inline bool ggml_can_fuse_subgraph(const struct ggml_cgraph * cgraph,
661+
int node_idx,
662+
int count,
663+
const enum ggml_op * ops,
664+
const int * outputs,
665+
int num_outputs) {
666+
GGML_ASSERT(count < 32);
667+
if (node_idx + count > cgraph->n_nodes) {
668+
return false;
669+
}
670+
671+
int idxs[32];
672+
673+
for (int i = 0; i < count; ++i) {
674+
idxs[i] = node_idx + i;
675+
}
676+
677+
return ggml_can_fuse_subgraph_ext(cgraph, idxs, count, ops, outputs, num_outputs);
678+
}
679+
650680
#ifdef __cplusplus
651681
}
652682
#endif
@@ -660,6 +690,13 @@ inline bool ggml_can_fuse(const struct ggml_cgraph * cgraph, int node_idx, std::
660690
return ggml_can_fuse(cgraph, node_idx, ops.begin(), (int)ops.size());
661691
}
662692

693+
inline bool ggml_can_fuse_subgraph(const struct ggml_cgraph * cgraph,
694+
int start_idx,
695+
std::initializer_list<enum ggml_op> ops,
696+
std::initializer_list<int> outputs = {}) {
697+
return ggml_can_fuse_subgraph(cgraph, start_idx, ops.size(), ops.begin(), outputs.begin(), outputs.size());
698+
}
699+
663700
// expose GGUF internals for test code
664701
GGML_API size_t gguf_type_size(enum gguf_type type);
665702
GGML_API struct gguf_context * gguf_init_from_file_impl(FILE * file, struct gguf_init_params params);

‎ggml/src/ggml.c‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6964,6 +6964,78 @@ void ggml_graph_print(const struct ggml_cgraph * cgraph) {
69646964
GGML_LOG_INFO("========================================\n");
69656965
}
69666966

6967+
static int ggml_node_list_find_tensor(const struct ggml_cgraph * cgraph,
6968+
const int * idxs,
6969+
int count,
6970+
const struct ggml_tensor * tensor) {
6971+
GGML_ASSERT(cgraph && idxs);
6972+
for (int i = 0; i < count; ++i) {
6973+
const int node_idx = idxs[i];
6974+
6975+
if (node_idx >= cgraph->n_nodes) {
6976+
return -1;
6977+
}
6978+
if (cgraph->nodes[node_idx] == tensor) {
6979+
return i;
6980+
}
6981+
}
6982+
return -1;
6983+
}
6984+
6985+
bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph,
6986+
const int * node_idxs,
6987+
int count,
6988+
const enum ggml_op * ops,
6989+
const int * outputs,
6990+
int num_outputs) {
6991+
GGML_ASSERT(outputs && num_outputs > 0);
6992+
6993+
for (int i = 0; i < count; ++i) {
6994+
if (node_idxs[i] >= cgraph->n_nodes) {
6995+
return false;
6996+
}
6997+
6998+
const struct ggml_tensor * node = cgraph->nodes[node_idxs[i]];
6999+
7000+
if (node->op != ops[i]) {
7001+
return false;
7002+
}
7003+
7004+
if (ggml_node_list_find_tensor(cgraph, outputs, num_outputs, node) != -1) {
7005+
continue;
7006+
}
7007+
7008+
if (node->flags & GGML_TENSOR_FLAG_OUTPUT) {
7009+
return false;
7010+
}
7011+
7012+
int subgraph_uses = 0;
7013+
for (int j = i + 1; j < count; ++j) {
7014+
const struct ggml_tensor * other_node = cgraph->nodes[node_idxs[j]];
7015+
for (int src_idx = 0; src_idx < GGML_MAX_SRC; src_idx++) {
7016+
if (other_node->src[src_idx] == node) {
7017+
subgraph_uses++;
7018+
}
7019+
}
7020+
}
7021+
7022+
if (subgraph_uses != ggml_node_get_use_count(cgraph, node_idxs[i])) {
7023+
return false;
7024+
}
7025+
7026+
// if node is a view, check if the view_src and all it's parent view_srcs are within the subgraph
7027+
struct ggml_tensor * view_src = node->view_src;
7028+
while (view_src) {
7029+
if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1) {
7030+
return false;
7031+
}
7032+
view_src = view_src->view_src;
7033+
}
7034+
}
7035+
7036+
return true;
7037+
}
7038+
69677039
// check if node is part of the graph
69687040
static bool ggml_graph_find(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) {
69697041
if (cgraph == NULL) {

0 commit comments

Comments
 (0)