whisper.rn 0.4.0-rc.8 → 0.4.0-rc.9

package/cpp/ggml-alloc.h

@@ -6,88 +6,70 @@
 extern "C" {
 #endif
 
-struct wsp_ggml_backend;
-struct wsp_ggml_backend_buffer;
-struct wsp_ggml_backend_buffer_type;
+typedef struct wsp_ggml_backend_buffer_type * wsp_ggml_backend_buffer_type_t;
+typedef struct      wsp_ggml_backend_buffer * wsp_ggml_backend_buffer_t;
+typedef struct             wsp_ggml_backend * wsp_ggml_backend_t;
 
-//
-// Legacy API
-//
-
-typedef struct wsp_ggml_allocr * wsp_ggml_allocr_t;
-
-// initialize allocator for use with CPU backend only
-WSP_GGML_API wsp_ggml_allocr_t wsp_ggml_allocr_new(void * data, size_t size, size_t alignment);
-WSP_GGML_API wsp_ggml_allocr_t wsp_ggml_allocr_new_measure(size_t alignment);
-
-// initialize allocator for use with ggml-backend
-WSP_GGML_API wsp_ggml_allocr_t wsp_ggml_allocr_new_from_buffer(struct wsp_ggml_backend_buffer * buffer);
-WSP_GGML_API wsp_ggml_allocr_t wsp_ggml_allocr_new_from_backend(struct wsp_ggml_backend * backend, size_t size); // allocates an owned buffer
-WSP_GGML_API wsp_ggml_allocr_t wsp_ggml_allocr_new_measure_from_backend(struct wsp_ggml_backend * backend);
-
-WSP_GGML_API struct wsp_ggml_backend_buffer * wsp_ggml_allocr_get_buffer(wsp_ggml_allocr_t alloc);
-
-// tell the allocator to parse nodes following the order described in the list
-// you should call this if your graph are optimized to execute out-of-order
-WSP_GGML_API void   wsp_ggml_allocr_set_parse_seq(wsp_ggml_allocr_t alloc, const int * list, int n);
-
-WSP_GGML_API void   wsp_ggml_allocr_free       (wsp_ggml_allocr_t alloc);
-WSP_GGML_API bool   wsp_ggml_allocr_is_measure (wsp_ggml_allocr_t alloc);
-WSP_GGML_API void   wsp_ggml_allocr_reset      (wsp_ggml_allocr_t alloc);
-WSP_GGML_API void   wsp_ggml_allocr_alloc      (wsp_ggml_allocr_t alloc, struct wsp_ggml_tensor * tensor);
-WSP_GGML_API size_t wsp_ggml_allocr_max_size   (wsp_ggml_allocr_t alloc);
+// Tensor allocator
+struct wsp_ggml_tallocr {
+    wsp_ggml_backend_buffer_t buffer;
+    void * base;
+    size_t alignment;
+    size_t offset;
+};
 
-WSP_GGML_API size_t wsp_ggml_allocr_alloc_graph(wsp_ggml_allocr_t alloc, struct wsp_ggml_cgraph * graph);
+WSP_GGML_API struct wsp_ggml_tallocr wsp_ggml_tallocr_new(wsp_ggml_backend_buffer_t buffer);
+WSP_GGML_API void                wsp_ggml_tallocr_alloc(struct wsp_ggml_tallocr * talloc, struct wsp_ggml_tensor * tensor);
 
-//
-// ggml-backend v2 API
-//
+// Graph allocator
+/*
+  Example usage:
+    wsp_ggml_gallocr_t galloc = wsp_ggml_gallocr_new(wsp_ggml_backend_cpu_buffer_type());
 
-// Separate 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
+    // optional: create a worst-case graph and reserve the buffers to avoid reallocations
+    wsp_ggml_gallocr_reserve(galloc, build_graph(max_batch));
 
-// Tensor allocator
-typedef struct wsp_ggml_tallocr * wsp_ggml_tallocr_t;
+    // allocate the graph
+    struct wsp_ggml_cgraph * graph = build_graph(batch);
+    wsp_ggml_gallocr_alloc_graph(galloc, graph);
 
-WSP_GGML_API wsp_ggml_tallocr_t wsp_ggml_tallocr_new(void * data, size_t size, size_t alignment);
-WSP_GGML_API wsp_ggml_tallocr_t wsp_ggml_tallocr_new_measure(size_t alignment);
-WSP_GGML_API wsp_ggml_tallocr_t wsp_ggml_tallocr_new_from_buft(struct wsp_ggml_backend_buffer_type * buft, size_t size);
-WSP_GGML_API wsp_ggml_tallocr_t wsp_ggml_tallocr_new_from_backend(struct wsp_ggml_backend * backend, size_t size); // allocates an owned buffer
-WSP_GGML_API wsp_ggml_tallocr_t wsp_ggml_tallocr_new_from_buffer(struct wsp_ggml_backend_buffer * buffer);
-WSP_GGML_API wsp_ggml_tallocr_t wsp_ggml_tallocr_new_measure_from_buft(struct wsp_ggml_backend_buffer_type * buft);
-WSP_GGML_API wsp_ggml_tallocr_t wsp_ggml_tallocr_new_measure_from_backend(struct wsp_ggml_backend * backend);
+    printf("compute buffer size: %zu bytes\n", wsp_ggml_gallocr_get_buffer_size(galloc, 0));
 
-WSP_GGML_API struct wsp_ggml_backend_buffer * wsp_ggml_tallocr_get_buffer(wsp_ggml_tallocr_t talloc);
+    // evaluate the graph
+    wsp_ggml_backend_graph_compute(backend, graph);
+*/
 
-WSP_GGML_API void   wsp_ggml_tallocr_free       (wsp_ggml_tallocr_t talloc);
-WSP_GGML_API bool   wsp_ggml_tallocr_is_measure (wsp_ggml_tallocr_t talloc);
-WSP_GGML_API void   wsp_ggml_tallocr_reset      (wsp_ggml_tallocr_t talloc);
-WSP_GGML_API void   wsp_ggml_tallocr_alloc      (wsp_ggml_tallocr_t talloc, struct wsp_ggml_tensor * tensor);
-WSP_GGML_API size_t wsp_ggml_tallocr_max_size   (wsp_ggml_tallocr_t talloc);
+// special tensor flags for use with the graph allocator:
+//   wsp_ggml_set_input(): all input tensors are allocated at the beginning of the graph in non-overlapping addresses
+//   wsp_ggml_set_output(): output tensors are never freed and never overwritten
 
-
-// Graph allocator
 typedef struct wsp_ggml_gallocr * wsp_ggml_gallocr_t;
 
-WSP_GGML_API wsp_ggml_gallocr_t wsp_ggml_gallocr_new(void);
-WSP_GGML_API void   wsp_ggml_gallocr_free(wsp_ggml_gallocr_t galloc);
+WSP_GGML_API wsp_ggml_gallocr_t wsp_ggml_gallocr_new(wsp_ggml_backend_buffer_type_t buft);
+WSP_GGML_API wsp_ggml_gallocr_t wsp_ggml_gallocr_new_n(wsp_ggml_backend_buffer_type_t * bufts, int n_bufs);
+WSP_GGML_API void           wsp_ggml_gallocr_free(wsp_ggml_gallocr_t galloc);
 
-WSP_GGML_API void   wsp_ggml_gallocr_set_parse_seq(wsp_ggml_gallocr_t galloc, const int * list, int n);
-WSP_GGML_API size_t wsp_ggml_gallocr_alloc_graph(wsp_ggml_gallocr_t galloc, wsp_ggml_tallocr_t talloc, struct wsp_ggml_cgraph * graph);
+// pre-allocate buffers from a measure graph - does not allocate or modify the graph
+// call with a worst-case graph to avoid buffer reallocations
+// not strictly required for single buffer usage: wsp_ggml_gallocr_alloc_graph will reallocate the buffers automatically if needed
+// returns false if the buffer allocation failed
+WSP_GGML_API bool wsp_ggml_gallocr_reserve(wsp_ggml_gallocr_t galloc, struct wsp_ggml_cgraph * graph);
+WSP_GGML_API bool wsp_ggml_gallocr_reserve_n(
+    wsp_ggml_gallocr_t galloc,
+    struct wsp_ggml_cgraph * graph,
+    const int * node_buffer_ids,
+    const int * leaf_buffer_ids);
 
-// Allocate tensors from the allocators given by the hash table
-WSP_GGML_API void   wsp_ggml_gallocr_alloc_graph_n(
-                    wsp_ggml_gallocr_t galloc,
-                    struct wsp_ggml_cgraph * graph,
-                    struct wsp_ggml_hash_set hash_set,
-                    wsp_ggml_tallocr_t * hash_node_talloc);
+// automatic reallocation if the topology changes when using a single buffer
+// returns false if using multiple buffers and a re-allocation is needed (call wsp_ggml_gallocr_reserve_n first to set the node buffers)
+WSP_GGML_API bool wsp_ggml_gallocr_alloc_graph(wsp_ggml_gallocr_t galloc, struct wsp_ggml_cgraph * graph);
 
+WSP_GGML_API size_t wsp_ggml_gallocr_get_buffer_size(wsp_ggml_gallocr_t galloc, int buffer_id);
 
 // Utils
 // Create a buffer and allocate all the tensors in a wsp_ggml_context
-WSP_GGML_API struct wsp_ggml_backend_buffer * wsp_ggml_backend_alloc_ctx_tensors_from_buft(struct wsp_ggml_context * ctx, struct wsp_ggml_backend_buffer_type * buft);
-WSP_GGML_API struct wsp_ggml_backend_buffer * wsp_ggml_backend_alloc_ctx_tensors(struct wsp_ggml_context * ctx, struct wsp_ggml_backend * backend);
+WSP_GGML_API struct wsp_ggml_backend_buffer * wsp_ggml_backend_alloc_ctx_tensors_from_buft(struct wsp_ggml_context * ctx, wsp_ggml_backend_buffer_type_t buft);
+WSP_GGML_API struct wsp_ggml_backend_buffer * wsp_ggml_backend_alloc_ctx_tensors(struct wsp_ggml_context * ctx, wsp_ggml_backend_t backend);
 
 #ifdef  __cplusplus
 }

package/cpp/ggml-backend-impl.h

@@ -9,107 +9,218 @@ extern "C" {
 #endif
 
     //
-    // Backend buffer
+    // Backend buffer type
     //
 
-    // buffer type
-    typedef void * wsp_ggml_backend_buffer_type_context_t;
-
     struct wsp_ggml_backend_buffer_type_i {
-        const char *          (*WSP_GGML_CALL get_name)        (wsp_ggml_backend_buffer_type_t buft);
-        wsp_ggml_backend_buffer_t (*WSP_GGML_CALL alloc_buffer)    (wsp_ggml_backend_buffer_type_t buft, size_t size);
-        size_t                (*WSP_GGML_CALL get_alignment)   (wsp_ggml_backend_buffer_type_t buft); // tensor alignment
-        size_t                (*WSP_GGML_CALL get_alloc_size)  (wsp_ggml_backend_buffer_type_t buft, const struct wsp_ggml_tensor * tensor); // data size needed to allocate the tensor, including padding
-        bool                  (*WSP_GGML_CALL supports_backend)(wsp_ggml_backend_buffer_type_t buft, wsp_ggml_backend_t backend); // check if the buffer type is usable by the backend
-        // check if tensor data is in host memory
-        // should be equivalent to supports_backend(buft, wsp_ggml_backend_cpu_init())
-        bool                  (*WSP_GGML_CALL is_host)         (wsp_ggml_backend_buffer_type_t buft);
+        const char *          (*get_name)      (wsp_ggml_backend_buffer_type_t buft);
+        // allocate a buffer of this type
+        wsp_ggml_backend_buffer_t (*alloc_buffer)  (wsp_ggml_backend_buffer_type_t buft, size_t size);
+        // tensor alignment
+        size_t                (*get_alignment) (wsp_ggml_backend_buffer_type_t buft);
+        // (optional) max buffer size that can be allocated (defaults to SIZE_MAX)
+        size_t                (*get_max_size)  (wsp_ggml_backend_buffer_type_t buft);
+        // (optional) data size needed to allocate the tensor, including padding (defaults to wsp_ggml_nbytes)
+        size_t                (*get_alloc_size)(wsp_ggml_backend_buffer_type_t buft, const struct wsp_ggml_tensor * tensor);
+        // (optional) check if tensor data is in host memory (defaults to false)
+        bool                  (*is_host)       (wsp_ggml_backend_buffer_type_t buft);
     };
 
     struct wsp_ggml_backend_buffer_type {
         struct wsp_ggml_backend_buffer_type_i  iface;
-        wsp_ggml_backend_buffer_type_context_t context;
+        wsp_ggml_backend_dev_t device;
+        void * context;
     };
 
-    // buffer
-    typedef void * wsp_ggml_backend_buffer_context_t;
+    //
+    // Backend buffer
+    //
 
     struct wsp_ggml_backend_buffer_i {
-        const char * (*WSP_GGML_CALL get_name)   (wsp_ggml_backend_buffer_t buffer);
-        void         (*WSP_GGML_CALL free_buffer)(wsp_ggml_backend_buffer_t buffer);
-        void *       (*WSP_GGML_CALL get_base)   (wsp_ggml_backend_buffer_t buffer);
-        void         (*WSP_GGML_CALL init_tensor)(wsp_ggml_backend_buffer_t buffer, struct wsp_ggml_tensor * tensor);
-        void         (*WSP_GGML_CALL set_tensor) (wsp_ggml_backend_buffer_t buffer,       struct wsp_ggml_tensor * tensor, const void * data, size_t offset, size_t size);
-        void         (*WSP_GGML_CALL get_tensor) (wsp_ggml_backend_buffer_t buffer, const struct wsp_ggml_tensor * tensor,       void * data, size_t offset, size_t size);
-        bool         (*WSP_GGML_CALL cpy_tensor) (wsp_ggml_backend_buffer_t buffer, const struct wsp_ggml_tensor * src, struct wsp_ggml_tensor * dst); // dst is in the buffer, src may be in any buffer
-        void         (*WSP_GGML_CALL clear)      (wsp_ggml_backend_buffer_t buffer, uint8_t value);
-        void         (*WSP_GGML_CALL reset)      (wsp_ggml_backend_buffer_t buffer); // reset any internal state due to tensor initialization, such as tensor extras
+        const char * (*get_name)     (wsp_ggml_backend_buffer_t buffer);
+        // (optional) free the buffer
+        void         (*free_buffer)  (wsp_ggml_backend_buffer_t buffer);
+        // base address of the buffer
+        void *       (*get_base)     (wsp_ggml_backend_buffer_t buffer);
+        // (optional) initialize a tensor in the buffer (eg. add tensor extras)
+        void         (*init_tensor)  (wsp_ggml_backend_buffer_t buffer, struct wsp_ggml_tensor * tensor);
+        // tensor data access
+        void         (*memset_tensor)(wsp_ggml_backend_buffer_t buffer,       struct wsp_ggml_tensor * tensor,     uint8_t value, size_t offset, size_t size);
+        void         (*set_tensor)   (wsp_ggml_backend_buffer_t buffer,       struct wsp_ggml_tensor * tensor, const void * data, size_t offset, size_t size);
+        void         (*get_tensor)   (wsp_ggml_backend_buffer_t buffer, const struct wsp_ggml_tensor * tensor,       void * data, size_t offset, size_t size);
+        // (optional) tensor copy: dst is in the buffer, src may be in any buffer, including buffers from a different backend (return false if not supported)
+        bool         (*cpy_tensor)   (wsp_ggml_backend_buffer_t buffer, const struct wsp_ggml_tensor * src, struct wsp_ggml_tensor * dst);
+        // clear the entire buffer
+        void         (*clear)        (wsp_ggml_backend_buffer_t buffer, uint8_t value);
+        // (optional) reset any internal state due to tensor initialization, such as tensor extras
+        void         (*reset)        (wsp_ggml_backend_buffer_t buffer);
     };
 
     struct wsp_ggml_backend_buffer {
         struct wsp_ggml_backend_buffer_i  iface;
         wsp_ggml_backend_buffer_type_t    buft;
-        wsp_ggml_backend_buffer_context_t context;
+        void * context;
         size_t size;
         enum wsp_ggml_backend_buffer_usage usage;
     };
 
-    WSP_GGML_CALL wsp_ggml_backend_buffer_t wsp_ggml_backend_buffer_init(
-                   wsp_ggml_backend_buffer_type_t      buft,
-            struct wsp_ggml_backend_buffer_i           iface,
-                   wsp_ggml_backend_buffer_context_t   context,
-                   size_t                          size);
+    wsp_ggml_backend_buffer_t wsp_ggml_backend_buffer_init(
+                   wsp_ggml_backend_buffer_type_t buft,
+            struct wsp_ggml_backend_buffer_i      iface,
+                   void *                     context,
+                   size_t                     size);
 
     // do not use directly, use wsp_ggml_backend_tensor_copy instead
     bool wsp_ggml_backend_buffer_copy_tensor(const struct wsp_ggml_tensor * src, struct wsp_ggml_tensor * dst);
 
+    // multi-buffer
+    // buffer that contains a collection of buffers
+    wsp_ggml_backend_buffer_t wsp_ggml_backend_multi_buffer_alloc_buffer(wsp_ggml_backend_buffer_t * buffers, size_t n_buffers);
+    bool                  wsp_ggml_backend_buffer_is_multi_buffer(wsp_ggml_backend_buffer_t buffer);
+    void                  wsp_ggml_backend_multi_buffer_set_usage(wsp_ggml_backend_buffer_t buffer, enum wsp_ggml_backend_buffer_usage usage);
+
     //
-    // Backend
+    // Backend (stream)
     //
 
-    typedef void * wsp_ggml_backend_context_t;
-
     struct wsp_ggml_backend_i {
-        const char * (*WSP_GGML_CALL get_name)(wsp_ggml_backend_t backend);
+        const char * (*get_name)(wsp_ggml_backend_t backend);
 
-        void (*WSP_GGML_CALL free)(wsp_ggml_backend_t backend);
+        void (*free)(wsp_ggml_backend_t backend);
 
+        // Will be moved to the device interface
         // buffer allocation
-        wsp_ggml_backend_buffer_type_t (*WSP_GGML_CALL get_default_buffer_type)(wsp_ggml_backend_t backend);
+        wsp_ggml_backend_buffer_type_t (*get_default_buffer_type)(wsp_ggml_backend_t backend);
 
         // (optional) asynchronous tensor data access
-        void (*WSP_GGML_CALL set_tensor_async)(wsp_ggml_backend_t backend,       struct wsp_ggml_tensor * tensor, const void * data, size_t offset, size_t size);
-        void (*WSP_GGML_CALL get_tensor_async)(wsp_ggml_backend_t backend, const struct wsp_ggml_tensor * tensor,       void * data, size_t offset, size_t size);
-        bool (*WSP_GGML_CALL cpy_tensor_async)(wsp_ggml_backend_t backend, const struct wsp_ggml_tensor * src, struct wsp_ggml_tensor * dst);
+        void (*set_tensor_async)(wsp_ggml_backend_t backend,       struct wsp_ggml_tensor * tensor, const void * data, size_t offset, size_t size);
+        void (*get_tensor_async)(wsp_ggml_backend_t backend, const struct wsp_ggml_tensor * tensor,       void * data, size_t offset, size_t size);
+        bool (*cpy_tensor_async)(wsp_ggml_backend_t backend_src, wsp_ggml_backend_t backend_dst, const struct wsp_ggml_tensor * src, struct wsp_ggml_tensor * dst);
 
         // (optional) complete all pending operations
-        void (*WSP_GGML_CALL synchronize)(wsp_ggml_backend_t backend);
-
-        // compute graph with a plan
-        wsp_ggml_backend_graph_plan_t (*WSP_GGML_CALL graph_plan_create) (wsp_ggml_backend_t backend, const struct wsp_ggml_cgraph * cgraph);
-        void                      (*WSP_GGML_CALL graph_plan_free)   (wsp_ggml_backend_t backend, wsp_ggml_backend_graph_plan_t plan);
-        void                      (*WSP_GGML_CALL graph_plan_compute)(wsp_ggml_backend_t backend, wsp_ggml_backend_graph_plan_t plan);
-
-        // compute graph without a plan (async)
-        bool (*WSP_GGML_CALL graph_compute)(wsp_ggml_backend_t backend, struct wsp_ggml_cgraph * cgraph);
-
-        // check if the backend supports an operation
-        bool (*WSP_GGML_CALL supports_op)(wsp_ggml_backend_t backend, const struct wsp_ggml_tensor * op);
+        void (*synchronize)(wsp_ggml_backend_t backend);
+
+        // (optional) compute graph with a plan (not used currently)
+        wsp_ggml_backend_graph_plan_t (*graph_plan_create) (wsp_ggml_backend_t backend, const struct wsp_ggml_cgraph * cgraph);
+        void                      (*graph_plan_free)   (wsp_ggml_backend_t backend, wsp_ggml_backend_graph_plan_t plan);
+        // update the plan with a new graph - this should be faster than creating a new plan when the graph has the same topology
+        void                      (*graph_plan_update) (wsp_ggml_backend_t backend, wsp_ggml_backend_graph_plan_t plan, const struct wsp_ggml_cgraph * cgraph);
+        // compute the graph with the plan
+        enum wsp_ggml_status          (*graph_plan_compute)(wsp_ggml_backend_t backend, wsp_ggml_backend_graph_plan_t plan);
+
+        // compute graph (always async if supported by the backend)
+        enum wsp_ggml_status          (*graph_compute)     (wsp_ggml_backend_t backend, struct wsp_ggml_cgraph * cgraph);
+
+        // IMPORTANT: these functions have been moved to the device interface and will be removed from the backend interface
+        //            new backends should implement the device interface instead
+        // These functions are being moved to the device interface
+        bool (*supports_op)  (wsp_ggml_backend_t backend, const struct wsp_ggml_tensor * op);
+        bool (*supports_buft)(wsp_ggml_backend_t backend, wsp_ggml_backend_buffer_type_t buft);
+        bool (*offload_op)   (wsp_ggml_backend_t backend, const struct wsp_ggml_tensor * op);
+
+        // (optional) event synchronization
+        // record an event on this stream
+        void (*event_record)(wsp_ggml_backend_t backend, wsp_ggml_backend_event_t event);
+        // wait for an event on on a different stream
+        void (*event_wait)  (wsp_ggml_backend_t backend, wsp_ggml_backend_event_t event);
     };
 
     struct wsp_ggml_backend {
+        wsp_ggml_guid_t guid;
         struct wsp_ggml_backend_i iface;
+        wsp_ggml_backend_dev_t device;
+        void * context;
+    };
+
+    struct wsp_ggml_backend_event {
+        struct wsp_ggml_backend_device * device;
+        void * context;
+    };
+
+    //
+    // Backend device
+    //
+
+    // Note: if additional properties are needed, we should add a struct with all of them
+    //       the current functions to obtain the properties can remain, since they are more convenient for often used properties
+    struct wsp_ggml_backend_device_i {
+        // device name: short identifier for this device, such as "CPU" or "CUDA0"
+        const char * (*get_name)(wsp_ggml_backend_dev_t dev);
+
+        // device description: short informative description of the device, could be the model name
+        const char * (*get_description)(wsp_ggml_backend_dev_t dev);
 
-        wsp_ggml_backend_context_t context;
+        // device memory in bytes
+        void         (*get_memory)(wsp_ggml_backend_dev_t dev, size_t * free, size_t * total);
+
+        // device type
+        enum wsp_ggml_backend_dev_type (*get_type)(wsp_ggml_backend_dev_t dev);
+
+        // device properties
+        void (*get_props)(wsp_ggml_backend_dev_t dev, struct wsp_ggml_backend_dev_props * props);
+
+        // backend (stream) initialization
+        wsp_ggml_backend_t (*init_backend)(wsp_ggml_backend_dev_t dev, const char * params);
+
+        // preferred buffer type
+        wsp_ggml_backend_buffer_type_t (*get_buffer_type)(wsp_ggml_backend_dev_t dev);
+
+        // (optional) host buffer type (in system memory, typically this is a pinned memory buffer for faster transfers between host and device)
+        wsp_ggml_backend_buffer_type_t (*get_host_buffer_type)(wsp_ggml_backend_dev_t dev);
+
+        // (optional) buffer from pointer: create a buffer from a host pointer (useful for memory mapped models and importing data from other libraries)
+        wsp_ggml_backend_buffer_t (*buffer_from_host_ptr)(wsp_ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size);
+
+        // check if the backend can compute an operation
+        bool (*supports_op)(wsp_ggml_backend_dev_t dev, const struct wsp_ggml_tensor * op);
+
+        // check if the backend can use tensors allocated in a buffer type
+        bool (*supports_buft)(wsp_ggml_backend_dev_t dev, wsp_ggml_backend_buffer_type_t buft);
+
+        // (optional) check if the backend wants to run an operation, even if the weights are allocated in an incompatible buffer
+        // these should be expensive operations that may benefit from running on this backend instead of the CPU backend
+        bool (*offload_op)(wsp_ggml_backend_dev_t dev, const struct wsp_ggml_tensor * op);
+
+        // (optional) event synchronization
+        wsp_ggml_backend_event_t (*event_new)         (wsp_ggml_backend_dev_t dev);
+        void                 (*event_free)        (wsp_ggml_backend_dev_t dev, wsp_ggml_backend_event_t event);
+        void                 (*event_synchronize) (wsp_ggml_backend_dev_t dev, wsp_ggml_backend_event_t event);
+    };
+
+    struct wsp_ggml_backend_device {
+        struct wsp_ggml_backend_device_i iface;
+        wsp_ggml_backend_reg_t reg;
+        void * context;
     };
 
     //
-    // Backend registry
+    // Backend (reg)
     //
 
-    typedef wsp_ggml_backend_t (*WSP_GGML_CALL wsp_ggml_backend_init_fn)(const char * params, void * user_data);
+    struct wsp_ggml_backend_reg_i {
+        const char * (*get_name)(wsp_ggml_backend_reg_t reg);
+
+        // enumerate available devices
+        size_t             (*get_device_count)(wsp_ggml_backend_reg_t reg);
+        wsp_ggml_backend_dev_t (*get_device)(wsp_ggml_backend_reg_t reg, size_t index);
+
+        // (optional) get a pointer to a function in the backend
+        // backends can add custom functions that are not part of the standard ggml-backend interface
+        void * (*get_proc_address)(wsp_ggml_backend_reg_t reg, const char * name);
+    };
+
+    struct wsp_ggml_backend_reg {
+        // int api_version; // TODO: for dynamic loading
+        struct wsp_ggml_backend_reg_i iface;
+        void * context;
+    };
+
 
-    WSP_GGML_CALL void wsp_ggml_backend_register(const char * name, wsp_ggml_backend_init_fn init_fn, wsp_ggml_backend_buffer_type_t default_buffer_type, void * user_data);
+    // Internal backend registry API
+    void wsp_ggml_backend_register(wsp_ggml_backend_reg_t reg);
+    void wsp_ggml_backend_device_register(wsp_ggml_backend_dev_t device);
+    // TODO: backends can be loaded as a dynamic library, in which case it needs to export this function
+    // typedef wsp_ggml_backend_register_t * (*wsp_ggml_backend_init)(void);
 
 #ifdef  __cplusplus
 }