Update ggml_extend.hpp

Author: idostyle

ggml_extend.hpp

@@ -1045,9 +1045,25 @@ struct GGMLRunner {
         reset_compute_ctx();
         struct ggml_cgraph* gf = get_graph();
         backend_tensor_data_map.clear();
-        compute_allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
+        
+        ggml_backend_buffer_type_t bufts[2] = {
+            ggml_backend_get_default_buffer_type(backend),
+            ggml_backend_get_default_buffer_type(backend)
+        };
 
-        if (!ggml_gallocr_reserve(compute_allocr, gf)) {
+        compute_allocr = ggml_gallocr_new_n_multi(bufts, 2);
+
+        int n_nodes = ggml_graph_n_nodes(gf);
+
+        int total_n_half = n_nodes / 2;
+
+        int * node_buffer_ids = (int *)calloc(n_nodes, sizeof(int));
+
+        for (int i = 0; i < n_nodes; i++) {
+            node_buffer_ids[i] = i < total_n_half ? 0 : 1;
+        }
+
+        if (!ggml_gallocr_reserve_n(compute_allocr, gf, node_buffer_ids, NULL)) {
             // failed to allocate the compute buffer
             LOG_ERROR("%s: failed to allocate the compute buffer\n", get_desc().c_str());
             free_compute_buffer();