ggml-alloc.h 2.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. #pragma once
  2. #include "ggml.h"
  3. #ifdef __cplusplus
  4. extern "C" {
  5. #endif
  6. typedef struct ggml_backend_buffer_type * ggml_backend_buffer_type_t;
  7. typedef struct ggml_backend_buffer * ggml_backend_buffer_t;
  8. typedef struct ggml_backend * ggml_backend_t;
  9. // Tensor allocator
  10. typedef struct ggml_tallocr * ggml_tallocr_t;
  11. GGML_API ggml_tallocr_t ggml_tallocr_new(ggml_backend_buffer_t buffer);
  12. GGML_API void ggml_tallocr_free(ggml_tallocr_t talloc);
  13. GGML_API void ggml_tallocr_alloc(ggml_tallocr_t talloc, struct ggml_tensor * tensor);
  14. // Graph allocator
  15. /*
  16. Example usage:
  17. ggml_gallocr_t galloc = ggml_gallocr_new(ggml_bacckend_cpu_buffer_type());
  18. // optional: create a worst-case graph and reserve the buffers to avoid reallocations
  19. ggml_gallocr_reserve(galloc, build_graph(max_batch));
  20. // allocate the graph
  21. struct ggml_cgraph * graph = build_graph(batch);
  22. ggml_gallocr_alloc_graph(galloc, graph);
  23. printf("compute buffer size: %zu bytes\n", ggml_gallocr_get_buffer_size(galloc, 0));
  24. // evaluate the graph
  25. ggml_backend_graph_compute(backend, graph);
  26. */
  27. // special tensor flags for use with the graph allocator:
  28. // ggml_set_input(): all input tensors are allocated at the beginning of the graph in non-overlapping addresses
  29. // ggml_set_output(): output tensors are never freed and never overwritten
  30. typedef struct ggml_gallocr * ggml_gallocr_t;
  31. GGML_API ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft);
  32. GGML_API ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs);
  33. GGML_API void ggml_gallocr_free(ggml_gallocr_t galloc);
  34. // pre-allocate buffers from a measure graph - does not allocate or modify the graph
  35. // call with a worst-case graph to avoid buffer reallocations
  36. // not strictly required for single buffer usage: ggml_gallocr_alloc_graph will reallocate the buffers automatically if needed
  37. // returns false if the buffer allocation failed
  38. GGML_API bool ggml_gallocr_reserve(ggml_gallocr_t galloc, struct ggml_cgraph * graph);
  39. GGML_API bool ggml_gallocr_reserve_n(ggml_gallocr_t galloc, struct ggml_cgraph * graph, const int * node_buffer_ids);
  40. // automatic reallocation if the topology changes when using a single buffer
  41. // returns false if using multiple buffers and a re-allocation is needed (call ggml_gallocr_reserve_n first to set the node buffers)
  42. GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph * graph);
  43. GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id);
  44. // Utils
  45. // Create a buffer and allocate all the tensors in a ggml_context
  46. GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors_from_buft(struct ggml_context * ctx, ggml_backend_buffer_type_t buft);
  47. GGML_API struct ggml_backend_buffer * ggml_backend_alloc_ctx_tensors(struct ggml_context * ctx, ggml_backend_t backend);
  48. #ifdef __cplusplus
  49. }
  50. #endif