Felipe97/llama-cpp-compiled

1#include "arg.h"2#include "common.h"3#include "log.h"4#include "llama-cpp.h"5#include "../src/llama-ext.h"6#include "ggml.h"7#include "gguf-model-data.h"8#include "gguf.h"9#include "ggml-backend.h"10#include "download.h"11 12#include <array>13#include <vector>14#include <set>15#include <fstream>16#include <iostream>17#include <random>18 19// Noop because weights are not needed20static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) {21    GGML_UNUSED(tensor);22    GGML_UNUSED(userdata);23}24 25struct input_tensor {26    ggml_type type;27    std::array<int64_t, 4> ne;28    std::array<size_t, 4> nb;29 30    input_tensor(ggml_type type, int64_t * ne, size_t * nb): type(type) {31        memcpy(this->ne.data(), ne, 4 * sizeof(int64_t));32        memcpy(this->nb.data(), nb, 4 * sizeof(size_t));33    }34 35    bool operator<(const input_tensor &b) const {36        return std::tie(type, ne, nb) <37               std::tie(b.type, b.ne, b.nb);38    }39 40    void serialize(std::ostream& out) const {41        out << type << ' ';42        for (size_t i = 0; i < 4; i++) {43            out << ne[i] << ' ';44        }45        for (size_t i = 0; i < 4; i++) {46            out << nb[i] << ' ';47        }48    }49};50 51struct test_object {52    ggml_op op;53    ggml_type type;54    std::array<int64_t, 4> ne;55    std::vector<int32_t> op_params;56    std::vector<input_tensor> sources;57    std::string name;58 59    void serialize(std::ostream& out) const {60        out << op << ' ' << type << ' ';61        for (size_t i = 0; i < 4; i++) {62            out << ne[i] << ' ';63        }64 65        out << op_params.size() << ' ';66        for (size_t i = 0; i < op_params.size(); i++) {67            out << op_params[i] << ' ';68        }69 70        out << sources.size() << ' ';71        for (size_t s = 0; s < sources.size(); s++) {72            sources[s].serialize(out);73        }74 75        if (!name.empty()) {76            out << name;77        } else {78            out << '-';79        }80 81        out << '\n';82    }83 84    bool operator<(const test_object &b) const {85        return std::tie(op, type, ne, op_params, sources) <86               std::tie(b.op, b.type, b.ne, b.op_params, b.sources);87    }88};89 90static void extract_graph_ops(ggml_cgraph * cgraph, const char * label, std::set<test_object> & tests) {91    int n_nodes = ggml_graph_n_nodes(cgraph);92    int n_skipped = 0;93    int n_before = (int) tests.size();94    for (int i = 0; i < n_nodes; i++) {95        ggml_tensor * node = ggml_graph_node(cgraph, i);96 97        if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) {98            n_skipped++;99            continue;100        }101 102        test_object test;103 104        test.op = node->op;105        test.type = node->type;106        memcpy(&test.ne, node->ne, 4 * sizeof(int64_t));107 108        test.op_params.resize(GGML_MAX_OP_PARAMS / sizeof(int32_t));109        memcpy(test.op_params.data(), node->op_params, GGML_MAX_OP_PARAMS);110 111        for (size_t s = 0; s < GGML_MAX_SRC; s++) {112            if (node->src[s] == nullptr) {113                break;114            }115 116            test.sources.emplace_back(node->src[s]->type, node->src[s]->ne, node->src[s]->nb);117        }118 119        test.name = node->name;120        tests.insert(test);121    }122 123    int n_new = (int) tests.size() - n_before;124    LOG_INF("%s: %d unique ops, %d total nodes, %d skipped (view ops)\n",125            label, n_new, n_nodes, n_skipped);126}127 128int main(int argc, char ** argv) {129    common_params params;130    params.out_file = "tests.txt";131 132    common_init();133 134    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_EXPORT_GRAPH_OPS)) {135        return 1;136    }137 138    // Load CPU-only139    ggml_backend_dev_t cpu_device = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);140    params.devices = { cpu_device, nullptr };141    params.fit_params = false;142    params.n_gpu_layers = 0;143 144    params.warmup = false;145 146    llama_context * ctx;147    common_init_result_ptr init_result;148    llama_context_ptr ctx2;149    llama_model_ptr model;150 151    if (params.model.hf_repo.empty()) {152        init_result = common_init_from_params(params);153 154        ctx = init_result->context();155        if (!ctx) {156            LOG_ERR("failed to initialize params\n");157            return 1;158        }159    } else {160#ifdef LLAMA_HF_FETCH161        auto [hf_repo, hf_quant] = common_download_split_repo_tag(params.model.hf_repo);162        if (hf_quant.empty() || hf_quant == "latest") {163            hf_quant = "Q4_K_M";164        }165 166        gguf_context_ptr gguf_ctx = gguf_fetch_gguf_ctx(hf_repo, hf_quant);167        if (!gguf_ctx) {168            LOG_ERR("failed to fetch GGUF metadata from %s\n", hf_repo.c_str());169            return 1;170        }171 172        llama_model_params model_params = llama_model_default_params();173        model_params.devices = params.devices.data();174        model_params.no_alloc = true;175 176        model.reset(llama_model_init_from_user(gguf_ctx.get(), set_tensor_data, nullptr, model_params));177 178        if (!model) {179            LOG_ERR("failed to create llama_model from %s\n", hf_repo.c_str());180            return 1;181        }182 183        llama_context_params ctx_params = llama_context_default_params();184        ctx2.reset(llama_init_from_model(model.get(), ctx_params));185        ctx = ctx2.get();186 187        if (!ctx) {188            LOG_ERR("failed to create llama_context\n");189            return 1;190        }191#else192        LOG_ERR("test-export-graph-ops compiled without HF fetch support\n");193        return 1;194#endif195    }196 197    const uint32_t n_seqs  = llama_n_seq_max(ctx);198    const uint32_t n_tokens = std::min(llama_n_ctx(ctx), llama_n_ubatch(ctx));199 200    std::set<test_object> tests;201 202    auto * gf_pp = llama_graph_reserve(ctx, n_tokens, n_seqs, n_tokens);203    if (!gf_pp) {204        LOG_ERR("failed to reserve prompt processing graph\n");205        return 1;206    }207    extract_graph_ops(gf_pp, "pp", tests);208 209    auto * gf_tg = llama_graph_reserve(ctx, n_seqs, n_seqs, n_seqs);210    if (!gf_tg) {211        LOG_ERR("failed to reserve token generation graph\n");212        return 1;213    }214    extract_graph_ops(gf_tg, "tg", tests);215 216    LOG_INF("%d unique ops total\n", (int) tests.size());217 218    std::ofstream f(params.out_file);219 220    if (!f.is_open()) {221        LOG_ERR("unable to open output file: %s\n", params.out_file.c_str());222        return 1;223    }224 225    for (const auto& test : tests) {226        test.serialize(f);227    }228 229    return 0;230}231