Repository: ggerganov/llama.cpp
File: ggml/src/ggml.c (C)

Code before:

static int ggml_find_tensor_node_list(const struct ggml_cgraph * cgraph,
                                      const int *                idxs,
                                      int                        count,
                                      const struct ggml_tensor * tensor) {
    GGML_ASSERT(cgraph && idxs);
    for (int i = 0; i < count; ++i) {
        const int node_idx = idxs[i];

        if (node_idx >= cgraph->n_nodes) {
            return -1;
        }
        if (cgraph->nodes[node_idx] == tensor) {
            return i;
        }
    }
    return -1;
}

bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph,
                                const int *                node_idxs,
                                int                        count,
                                const enum ggml_op *       ops,
                                const int *                outputs,
                                int                        num_outputs) {
    GGML_ASSERT(count < 32 && outputs && num_outputs > 0);

    for (int i = 0; i < count; ++i) {
        if (node_idxs[i] >= cgraph->n_nodes || cgraph->nodes[node_idxs[i]]->op != ops[i]) {
            return false;
        }

        const struct ggml_tensor * node = cgraph->nodes[node_idxs[i]];

        if (ggml_find_tensor_node_list(cgraph, outputs, num_outputs, node) != -1) {
            continue;
        }

        if (node->flags & GGML_TENSOR_FLAG_OUTPUT) {
            return false;
        }

        int subgraph_uses = 0;
        for (int j = i + 1; j < count; ++j) {
            const struct ggml_tensor * other_node = cgraph->nodes[node_idxs[j]];
            for (int src_idx = 0; src_idx < GGML_MAX_SRC; src_idx++) {
                if (other_node->src[src_idx] == node) {
                    subgraph_uses++;
                }
            }
        }

Review comment (suggestion):
The check for `count < 32` seems more suitable for `ggml_can_fuse_subgraph` - here it's not clear why it is necessary.

Code after:

static int ggml_node_list_find_tensor(const struct ggml_cgraph * cgraph,
                                      const int *                idxs,
                                      int                        count,
                                      const struct ggml_tensor * tensor) {
    GGML_ASSERT(cgraph && idxs);
    for (int i = 0; i < count; ++i) {
        const int node_idx = idxs[i];

        if (node_idx >= cgraph->n_nodes) {
            return -1;
        }
        if (cgraph->nodes[node_idx] == tensor) {
            return i;
        }
    }
    return -1;
}

bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph,
                                const int *                node_idxs,
                                int                        count,
                                const enum ggml_op *       ops,
                                const int *                outputs,
                                int                        num_outputs) {
    GGML_ASSERT(outputs && num_outputs > 0);

    for (int i = 0; i < count; ++i) {
        if (node_idxs[i] >= cgraph->n_nodes) {
            return false;
        }

        const struct ggml_tensor * node = cgraph->nodes[node_idxs[i]];

        if (node->op != ops[i]) {
            return false;
        }

        if (ggml_node_list_find_tensor(cgraph, outputs, num_outputs, node) != -1) {
            continue;
        }

        if (node->flags & GGML_TENSOR_FLAG_OUTPUT) {
            return false;
        }

        int subgraph_uses = 0;
        for (int j = i + 1; j < count; ++j) {
            const struct ggml_tensor * other_node = cgraph->nodes[node_idxs[j]];
            for (int src_idx = 0; src_idx < GGML_MAX_SRC; src_idx++) {
                if (other_node->src[src_idx] == node) {