ggml-alloc : split into tensor and graph allocators

ggerganov · Oct 20, 2023 · 2257cf7 · 2257cf7
1 parent a186173
commit 2257cf7
Show file tree

Hide file tree

Showing 4 changed files with 326 additions and 212 deletions.
diff --git a/include/ggml/ggml-alloc.h b/include/ggml/ggml-alloc.h
@@ -8,6 +8,11 @@ extern "C" {
 
 struct ggml_backend;
 struct ggml_backend_buffer;
+
+//
+// Legacy API
+//
+
 typedef struct ggml_allocr * ggml_allocr_t;
 
 // initialize allocator for use with CPU backend only
@@ -33,12 +38,47 @@ GGML_API size_t ggml_allocr_max_size   (ggml_allocr_t alloc);
 
 GGML_API size_t ggml_allocr_alloc_graph(ggml_allocr_t alloc, struct ggml_cgraph * graph);
 
+//
+// ggml-backend v2 API
+//
+
+// Seperate tensor and graph allocator objects
+// This is necessary for multi-backend allocation because the graph allocator needs to use multiple tensor allocators
+// The original API is kept as a wrapper around the new API
+
+// Tensor allocator
+typedef struct ggml_tallocr * ggml_tallocr_t;
+
+GGML_API ggml_tallocr_t ggml_tallocr_new(void * data, size_t size, size_t alignment);
+GGML_API ggml_tallocr_t ggml_tallocr_new_measure(size_t alignment);
+GGML_API ggml_tallocr_t ggml_tallocr_new_from_buffer(struct ggml_backend_buffer * buffer);
+GGML_API ggml_tallocr_t ggml_tallocr_new_from_backend(struct ggml_backend * backend, size_t size); // allocates an owned buffer
+GGML_API ggml_tallocr_t ggml_tallocr_new_measure_from_backend(struct ggml_backend * backend);
+
+GGML_API struct ggml_backend_buffer * ggml_tallocr_get_buffer(ggml_tallocr_t talloc);
+
+GGML_API void   ggml_tallocr_free       (ggml_tallocr_t talloc);
+GGML_API bool   ggml_tallocr_is_measure (ggml_tallocr_t talloc);
+GGML_API void   ggml_tallocr_reset      (ggml_tallocr_t talloc);
+GGML_API void   ggml_tallocr_alloc      (ggml_tallocr_t talloc, struct ggml_tensor * tensor);
+GGML_API size_t ggml_tallocr_max_size   (ggml_tallocr_t talloc);
+
+
+// Graph allocator
+typedef struct ggml_gallocr * ggml_gallocr_t;
+
+GGML_API ggml_gallocr_t ggml_gallocr_new(void);
+GGML_API void   ggml_gallocr_free(ggml_gallocr_t galloc);
+
+GGML_API void   ggml_gallocr_set_parse_seq(ggml_gallocr_t galloc, const int * list, int n);
+GGML_API size_t ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, ggml_tallocr_t talloc, struct ggml_cgraph * graph);
+
 // Allocate tensors from the allocators given by the hash table
-GGML_API void ggml_allocr_alloc_graph_n(
+GGML_API void   ggml_gallocr_alloc_graph_n(
+                    ggml_gallocr_t galloc,
                     struct ggml_cgraph * graph,
                     struct ggml_hash_set hash_set,
-                    ggml_allocr_t hash_node_alloct[]);
-
+                    ggml_tallocr_t * hash_node_talloc);
 
 #ifdef  __cplusplus
 }

diff --git a/include/ggml/ggml-backend.h b/include/ggml/ggml-backend.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include "ggml.h"
+#include "ggml-alloc.h"
 
 #ifdef  __cplusplus
 extern "C" {
@@ -109,8 +110,6 @@ extern "C" {
         ggml_backend_sched_graph_compute(sched, graph);
     */
 
-    struct ggml_allocr;
-
     struct ggml_backend_sched;
     typedef struct ggml_backend_sched * ggml_backend_sched_t;
 
@@ -122,8 +121,8 @@ extern "C" {
     // Initialize backend buffers from a measure graph
     GGML_API void ggml_backend_sched_init_measure(ggml_backend_sched_t sched, struct ggml_cgraph * measure_graph);
 
-    GGML_API struct ggml_allocr * ggml_backend_sched_get_allocr(ggml_backend_sched_t sched, ggml_backend_t backend);
-    GGML_API ggml_backend_buffer_t ggml_backend_sched_get_buffer(ggml_backend_sched_t sched, ggml_backend_t backend);
+    GGML_API ggml_tallocr_t        ggml_backend_sched_get_tallocr(ggml_backend_sched_t sched, ggml_backend_t backend);
+    GGML_API ggml_backend_buffer_t ggml_backend_sched_get_buffer (ggml_backend_sched_t sched, ggml_backend_t backend);
 
     // Allocate a graph on the backend scheduler
     GGML_API void ggml_backend_sched_graph_compute(