tests : copy ggml graph copy functions for backend ops tests

This commit copies a few ggml graph copy functions into
test-backend-ops.cpp to allow for some more control over testing extra
backend buffers.

Author: danbev

tests/CMakeLists.txt

@@ -198,6 +198,7 @@ if (NOT LLAMA_SANITIZE_ADDRESS)
 endif()
 llama_build_and_test(test-gguf.cpp)
 llama_build_and_test(test-backend-ops.cpp)
+target_include_directories(test-backend-ops PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src)
 
 llama_build_and_test(test-model-load-cancel.cpp  LABEL "model")
 llama_build_and_test(test-autorelease.cpp        LABEL "model")

tests/test-backend-ops.cpp

@@ -20,6 +20,8 @@
 #include <ggml-backend.h>
 #include <ggml-cpp.h>
 
+#include "ggml-impl.h"
+
 #include <algorithm>
 #include <array>
 #include <cfloat>
@@ -1085,7 +1087,232 @@ struct test_case {
         }
     }
 
-    bool eval(ggml_backend_t backend1, ggml_backend_t backend2, const char * op_names_filter, printer * output_printer) {
+    struct ggml_tensor * ggml_dup_tensor_layout(struct ggml_context * ctx, const struct ggml_tensor * tensor) {
+        struct ggml_tensor * dup = ggml_dup_tensor(ctx, tensor);
+        for (int i = 0; i < GGML_MAX_DIMS; i++) {
+            dup->nb[i] = tensor->nb[i];
+        }
+        return dup;
+    }
+
+    struct ggml_tensor * graph_copy_dup_tensor(struct ggml_hash_set hash_set, struct ggml_tensor ** node_copies,
+        struct ggml_context * ctx_allocated, struct ggml_context * ctx_unallocated, struct ggml_tensor * src) {
+
+        GGML_ASSERT(src != NULL);
+        GGML_ASSERT(src->data && "graph must be allocated");
+
+        size_t id = ggml_hash_insert(&hash_set, src);
+        if (id == GGML_HASHSET_ALREADY_EXISTS) {
+            return node_copies[ggml_hash_find(&hash_set, src)];
+        }
+
+        struct ggml_tensor * dst = ggml_dup_tensor_layout(src->data && !src->view_src ? ctx_allocated : ctx_unallocated, src);
+        if (src->view_src != NULL) {
+            dst->view_src = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, src->view_src);
+            dst->view_offs = src->view_offs;
+        }
+        dst->op = src->op;
+        memcpy(dst->op_params, src->op_params, sizeof(dst->op_params));
+        ggml_set_name(dst, src->name);
+
+        // copy src
+        for (int i = 0; i < GGML_MAX_SRC; i++) {
+            struct ggml_tensor * s = src->src[i];
+            if (s == NULL) {
+                continue;
+            }
+            dst->src[i] = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, s);
+        }
+
+        node_copies[id] = dst;
+        return dst;
+    }
+
+    void graph_copy_init_tensor(struct ggml_hash_set * hash_set, struct ggml_tensor ** node_copies, bool * node_init, struct ggml_tensor * src) {
+        size_t id = ggml_hash_find(hash_set, src);
+        if (node_init[id]) {
+            return;
+        }
+        node_init[id] = true;
+
+        struct ggml_tensor * dst = node_copies[id];
+        if (dst->view_src != NULL) {
+            graph_copy_init_tensor(hash_set, node_copies, node_init, src->view_src);
+            enum ggml_status status = ggml_backend_view_init(dst);
+            GGML_ASSERT(status == GGML_STATUS_SUCCESS);
+        }
+        else {
+            ggml_backend_tensor_copy(src, dst);
+        }
+
+        // init src
+        for (int i = 0; i < GGML_MAX_SRC; i++) {
+            struct ggml_tensor * s = src->src[i];
+            if (s == NULL) {
+                continue;
+            }
+            graph_copy_init_tensor(hash_set, node_copies, node_init, s);
+        }
+    }
+
+    struct ggml_backend_graph_copy ggml_backend_graph_copy(ggml_backend_t backend, struct ggml_cgraph * graph,
+            std::unordered_map<ggml_backend_buffer_type_t, ggml_backend_buffer_t> extra_buf_map) {
+        GGML_ASSERT(graph);
+        struct ggml_hash_set hash_set = ggml_hash_set_new(graph->visited_hash_set.size);
+        struct ggml_tensor ** node_copies = (ggml_tensor **) calloc(hash_set.size, sizeof(node_copies[0])); // NOLINT
+        bool * node_init = (bool *) calloc(hash_set.size, sizeof(node_init[0]));
+
+        struct ggml_init_params params = {
+            /* .mem_size   = */ ggml_tensor_overhead()*hash_set.size + ggml_graph_overhead_custom(graph->size, false),
+            /* .mem_buffer = */ NULL,
+            /* .no_alloc   = */ true
+        };
+
+        struct ggml_context * ctx_allocated = ggml_init(params);
+        struct ggml_context * ctx_unallocated = ggml_init(params);
+
+        if (ctx_allocated == NULL || ctx_unallocated == NULL) {
+            GGML_LOG_ERROR("%s: failed to allocate context for graph copy\n", __func__);
+            ggml_hash_set_free(&hash_set);
+            free(node_copies);
+            free(node_init);
+            ggml_free(ctx_allocated);
+            ggml_free(ctx_unallocated);
+            return {
+                /* .buffer           = */ NULL,
+                /* .ctx_allocated    = */ NULL,
+                /* .ctx_unallocated  = */ NULL,
+                /* .graph            = */ NULL,
+            };
+        }
+
+        // dup nodes
+        for (int i = 0; i < graph->n_nodes; i++) {
+            struct ggml_tensor * node = graph->nodes[i];
+            graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, node);
+        }
+
+        // allocate nodes
+        ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx_allocated, backend);
+
+        if (buffer == NULL) {
+            GGML_LOG_ERROR("%s: failed to allocate buffer for graph copy\n", __func__);
+            ggml_hash_set_free(&hash_set);
+            free(node_copies);
+            free(node_init);
+            ggml_free(ctx_allocated);
+            ggml_free(ctx_unallocated);
+            for (auto buft : extra_buf_map) {
+                ggml_backend_buffer_free(buft.second);
+            }
+            return {
+                /* .buffer           = */ NULL,
+                /* .ctx_allocated    = */ NULL,
+                /* .ctx_unallocated  = */ NULL,
+                /* .graph            = */ NULL,
+            };
+        }
+
+        //printf("copy buffer size: %zu MB\n", ggml_backend_buffer_get_size(buffer) / 1024 / 1024);
+
+        // copy data and init views
+        for (int i = 0; i < graph->n_nodes; i++) {
+            struct ggml_tensor * node = graph->nodes[i];
+
+            if (node->op != GGML_OP_NONE && node->src[0]) {
+                for (const auto& [buft, buf] : extra_buf_map) {
+                    size_t id = ggml_hash_find(&hash_set, node);
+                    ggml_status status = ggml_backend_buffer_init_tensor(buf, node_copies[id]);
+                    if (status != GGML_STATUS_SUCCESS) {
+                        GGML_LOG_ERROR("%s: failed to initialize tensor in extra buffer type '%s' for graph copy\n", __func__, ggml_backend_buft_name(buft));
+                    }
+                }
+            }
+
+            graph_copy_init_tensor(&hash_set, node_copies, node_init, node);
+        }
+
+        // build graph copy
+        struct ggml_cgraph * graph_copy = ggml_new_graph_custom(ctx_allocated, graph->size, false);
+        for (int i = 0; i < graph->n_nodes; i++) {
+            struct ggml_tensor * node = graph->nodes[i];
+            struct ggml_tensor * node_copy = node_copies[ggml_hash_find(&hash_set, node)];
+            graph_copy->nodes[i] = node_copy;
+        }
+        graph_copy->n_nodes = graph->n_nodes;
+
+        ggml_hash_set_free(&hash_set);
+        free(node_copies);
+        free(node_init);
+
+        return {
+            /* .buffer           = */ buffer,
+            /* .ctx_allocated    = */ ctx_allocated,
+            /* .ctx_unallocated  = */ ctx_unallocated,
+            /* .graph            = */ graph_copy,
+        };
+    }
+
+    bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t backend2,
+            struct ggml_cgraph * graph, ggml_backend_eval_callback callback, void * user_data,
+            struct ggml_tensor * test_node,
+            std::unordered_map<ggml_backend_buffer_type_t, ggml_backend_buffer_t> extra_buf_map) {
+        struct ggml_backend_graph_copy copy = ggml_backend_graph_copy(backend2, graph, extra_buf_map);
+        if (copy.buffer == NULL) {
+            return false;
+        }
+
+        struct ggml_cgraph * g1 = graph;
+        struct ggml_cgraph * g2 = copy.graph;
+
+        assert(g1->n_nodes == g2->n_nodes);
+
+        if (test_node != nullptr) {
+            // Compute the whole graph and only test the output for a specific tensor
+            ggml_backend_graph_compute(backend1, g1);
+            ggml_backend_graph_compute(backend2, g2);
+
+            int test_node_idx = -1;
+            for (int i = 0; i < g1->n_nodes; i++) {
+                struct ggml_tensor * t1 = g1->nodes[i];
+                if (t1 == test_node) {
+                    test_node_idx = i;
+                    break;
+                }
+            }
+            GGML_ASSERT(test_node_idx != -1);
+
+            callback(test_node_idx, g1->nodes[test_node_idx], g2->nodes[test_node_idx], user_data);
+        } else {
+            for (int i = 0; i < g1->n_nodes; i++) {
+                struct ggml_tensor * t1 = g1->nodes[i];
+                struct ggml_tensor * t2 = g2->nodes[i];
+
+                assert(t1->op == t2->op && ggml_are_same_layout(t1, t2));
+
+                struct ggml_cgraph g1v = ggml_graph_view(g1, i, i + 1);
+                struct ggml_cgraph g2v = ggml_graph_view(g2, i, i + 1);
+
+                ggml_backend_graph_compute(backend1, &g1v);
+                ggml_backend_graph_compute(backend2, &g2v);
+
+                if (ggml_is_view_op(t1->op)) {
+                    continue;
+                }
+
+                // compare results, calculate rms etc
+                if (!callback(i, t1, t2, user_data)) {
+                    break;
+                }
+            }
+        }
+        ggml_backend_graph_copy_free(copy);
+
+        return true;
+    }
+
+    bool eval(ggml_backend_t backend1, ggml_backend_t backend2, const char * op_names_filter, printer * output_printer,
+            std::unordered_map<ggml_backend_buffer_type_t, ggml_backend_buffer_t> extra_buf_map) {
         mode = MODE_TEST;
 
         ggml_init_params params = {
@@ -1225,7 +1452,8 @@ struct test_case {
             GGML_UNUSED(index);
         };
 
-        const bool cmp_ok = ggml_backend_compare_graph_backend(backend1, backend2, gf, callback, &ud, run_whole_graph() ? out : nullptr);
+        const bool cmp_ok = ggml_backend_compare_graph_backend(backend1, backend2, gf, callback,
+                &ud, run_whole_graph() ? out : nullptr, extra_buf_map);
 
         ggml_backend_buffer_free(buf);
 
@@ -6764,7 +6992,7 @@ static bool test_backend(ggml_backend_t backend, test_mode mode, const char * op
 
         size_t n_ok = 0;
         for (auto & test : test_cases) {
-            if (test->eval(backend, backend_cpu, op_names_filter, output_printer)) {
+            if (test->eval(backend, backend_cpu, op_names_filter, output_printer, {})) {
                 n_ok++;
             }
         }
@@ -6923,7 +7151,34 @@ static void print_backend_features(ggml_backend_t backend) {
 static bool test_cpu_variant(const char * variant_name, const char * op_names_filter,
         const char * params_filter, printer * output_printer) {
     std::string backend_ref_name = "CPU-ref";
-    ggml_backend_load_all_variants("cpu");
+    std::string variant = std::string(variant_name).substr(4);
+    ggml_backend_load_variant("cpu", variant.c_str());
+
+    std::unordered_map<ggml_backend_buffer_type_t, ggml_backend_buffer_t> extra_buf_map;
+    {
+        auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+        auto * cpu_reg = ggml_backend_dev_backend_reg(cpu_dev);
+
+        auto ggml_backend_dev_get_extra_bufts_fn = (ggml_backend_dev_get_extra_bufts_t)
+            ggml_backend_reg_get_proc_address(cpu_reg, "ggml_backend_dev_get_extra_bufts");
+        if (ggml_backend_dev_get_extra_bufts_fn) {
+            ggml_backend_buffer_type_t * extra_bufts = ggml_backend_dev_get_extra_bufts_fn(cpu_dev);
+            while (extra_bufts && *extra_bufts) {
+                // TODO: What should the size be here? Do extra buffer types need a size even?
+                // We need to have a value larger than 0 to avoid the default buffer to be used.
+                extra_buf_map[*extra_bufts] = ggml_backend_buft_alloc_buffer(*extra_bufts, 1);
+                ++extra_bufts;
+            }
+        }
+    }
+
+    printf("\n");
+    for (auto buft : extra_buf_map) {
+        printf("Using extra buffer type: %s\n", ggml_backend_buft_name(buft.first));
+    }
+    printf("\n");
+
+    ggml_backend_load_variant("cpu", "ref");
 
     ggml_backend_t backend_ref = ggml_backend_init_by_name(backend_ref_name.c_str(), nullptr);
     if (backend_ref == nullptr) {
@@ -6941,6 +7196,8 @@ static bool test_cpu_variant(const char * variant_name, const char * op_names_fi
     }
     print_backend_features(backend_variant);
 
+
+
     printf("Testing CPU variant '%s' against '%s' backend...\n\n", variant_name, backend_ref_name.c_str());
 
     auto test_cases = make_test_cases_eval();
@@ -6962,7 +7219,7 @@ static bool test_cpu_variant(const char * variant_name, const char * op_names_fi
     for (auto & test : test_cases) {
         // Switch the order so that we copy from the reference backend to the
         // variant backend.
-        if (test->eval(backend_ref, backend_variant, op_names_filter, output_printer)) {
+        if (test->eval(backend_ref, backend_variant, op_names_filter, output_printer, extra_buf_map)) {
             n_ok++;
         }
     }
@@ -6972,6 +7229,10 @@ static bool test_cpu_variant(const char * variant_name, const char * op_names_fi
     ggml_backend_free(backend_variant);
     ggml_backend_free(backend_ref);
 
+    for (auto buft : extra_buf_map) {
+        ggml_backend_buffer_free(buft.second);
+    }
+
     return n_ok == test_cases.size();
 }