whisper.rn 0.4.0-rc.3 → 0.4.0-rc.5

package/cpp/ggml-alloc.h

@@ -6,20 +6,86 @@
 extern "C" {
 #endif
 
+struct wsp_ggml_backend;
+struct wsp_ggml_backend_buffer;
+struct wsp_ggml_backend_buffer_type;
 
-WSP_GGML_API struct wsp_ggml_allocr * wsp_ggml_allocr_new(void * data, size_t size, size_t alignment);
-WSP_GGML_API struct wsp_ggml_allocr * wsp_ggml_allocr_new_measure(size_t alignment);
+//
+// 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(struct wsp_ggml_allocr * alloc, const int * list, int n);
+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);
+
+WSP_GGML_API size_t wsp_ggml_allocr_alloc_graph(wsp_ggml_allocr_t alloc, struct wsp_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 wsp_ggml_tallocr * wsp_ggml_tallocr_t;
+
+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_buffer(struct wsp_ggml_backend_buffer * buffer);
+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_measure_from_backend(struct wsp_ggml_backend * backend);
+
+WSP_GGML_API struct wsp_ggml_backend_buffer * wsp_ggml_tallocr_get_buffer(wsp_ggml_tallocr_t talloc);
+
+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);
+
+
+// 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 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);
+
+// 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);
 
-WSP_GGML_API void   wsp_ggml_allocr_free(struct wsp_ggml_allocr * alloc);
-WSP_GGML_API bool   wsp_ggml_allocr_is_measure(struct wsp_ggml_allocr * alloc);
-WSP_GGML_API void   wsp_ggml_allocr_reset(struct wsp_ggml_allocr * alloc);
-WSP_GGML_API void   wsp_ggml_allocr_alloc(struct wsp_ggml_allocr * alloc, struct wsp_ggml_tensor * tensor);
-WSP_GGML_API size_t wsp_ggml_allocr_alloc_graph(struct wsp_ggml_allocr * alloc, struct wsp_ggml_cgraph * graph);
 
+// 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);
 
 #ifdef  __cplusplus
 }

package/cpp/ggml-backend-impl.h

@@ -0,0 +1,112 @@
+#pragma once
+
+// ggml-backend internal header
+
+#include "ggml-backend.h"
+
+#ifdef  __cplusplus
+extern "C" {
+#endif
+
+    //
+    // Backend buffer
+    //
+
+    // buffer type
+    typedef void * wsp_ggml_backend_buffer_type_context_t;
+
+    struct wsp_ggml_backend_buffer_type_i {
+        wsp_ggml_backend_buffer_t (*alloc_buffer)    (wsp_ggml_backend_buffer_type_t buft, size_t size);
+        size_t                (*get_alignment)   (wsp_ggml_backend_buffer_type_t buft); // tensor alignment
+        size_t                (*get_alloc_size)  (wsp_ggml_backend_buffer_type_t buft, struct wsp_ggml_tensor * tensor); // data size needed to allocate the tensor, including padding
+        bool                  (*supports_backend)(wsp_ggml_backend_buffer_type_t buft, wsp_ggml_backend_t backend); // check if the buffer type is usable by the backend
+    };
+
+    struct wsp_ggml_backend_buffer_type {
+        struct wsp_ggml_backend_buffer_type_i  iface;
+        wsp_ggml_backend_buffer_type_context_t context;
+    };
+
+    // buffer
+    typedef void * wsp_ggml_backend_buffer_context_t;
+
+    struct wsp_ggml_backend_buffer_i {
+        void     (*free_buffer)(wsp_ggml_backend_buffer_t buffer);
+        //void     (*reset)      (wsp_ggml_backend_buffer_t buffer); // reset any internal state due to tensor initialization, such as tensor extras
+        void *   (*get_base)   (wsp_ggml_backend_buffer_t buffer);
+        void     (*init_tensor)(wsp_ggml_backend_buffer_t buffer, struct wsp_ggml_tensor * tensor);
+        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) copy tensor between different buffer-type, allow for single-copy tranfers
+        void (*cpy_tensor_from)(wsp_ggml_backend_buffer_t buffer, struct wsp_ggml_tensor * src, struct wsp_ggml_tensor * dst);
+        void (*cpy_tensor_to)  (wsp_ggml_backend_buffer_t buffer, struct wsp_ggml_tensor * src, struct wsp_ggml_tensor * dst);
+    };
+
+    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;
+        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,
+                   wsp_ggml_backend_buffer_context_t   context,
+                   size_t                          size);
+
+
+    //
+    // Backend
+    //
+
+    typedef void * wsp_ggml_backend_context_t;
+
+    struct wsp_ggml_backend_i {
+        const char * (*get_name)(wsp_ggml_backend_t backend);
+
+        void (*free)(wsp_ggml_backend_t backend);
+
+        // buffer allocation
+        wsp_ggml_backend_buffer_type_t (*get_default_buffer_type)(wsp_ggml_backend_t backend);
+
+        // (optional) asynchroneous tensor data access
+        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);
+
+        // (optional) asynchroneous tensor copy
+        void (*cpy_tensor_from_async)(wsp_ggml_backend_t backend, struct wsp_ggml_tensor * src, struct wsp_ggml_tensor * dst);
+        void (*cpy_tensor_to_async)  (wsp_ggml_backend_t backend, struct wsp_ggml_tensor * src, struct wsp_ggml_tensor * dst);
+
+        void (*synchronize)     (wsp_ggml_backend_t backend);
+
+        // compute graph with a plan
+        wsp_ggml_backend_graph_plan_t (*graph_plan_create) (wsp_ggml_backend_t backend, struct wsp_ggml_cgraph * cgraph);
+        void                      (*graph_plan_free)   (wsp_ggml_backend_t backend, wsp_ggml_backend_graph_plan_t plan);
+        void                      (*graph_plan_compute)(wsp_ggml_backend_t backend, wsp_ggml_backend_graph_plan_t plan);
+
+        // compute graph without a plan
+        void (*graph_compute)(wsp_ggml_backend_t backend, struct wsp_ggml_cgraph * cgraph);
+
+        // check if the backend supports an operation
+        bool (*supports_op)(wsp_ggml_backend_t backend, const struct wsp_ggml_tensor * op);
+    };
+
+    struct wsp_ggml_backend {
+        struct wsp_ggml_backend_i iface;
+
+        wsp_ggml_backend_context_t context;
+    };
+
+
+    //
+    // Backend registry
+    //
+
+    typedef wsp_ggml_backend_t (*wsp_ggml_backend_init_fn)(const char * params, void * user_data);
+
+    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);
+
+#ifdef  __cplusplus
+}
+#endif