traits.h 1.1 KB

1234567891011121314151617181920212223242526272829303132333435363738
  1. #pragma once
  2. #include "ggml-backend-impl.h"
  3. #include "ggml-cpu-impl.h"
  4. #include "ggml.h"
  5. #ifdef __cplusplus
  6. # include <vector>
  7. extern "C" {
  8. #endif
  9. // return true if op part of extra "accelerator"
  10. bool ggml_cpu_extra_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op);
  11. bool ggml_cpu_extra_work_size(int n_threads, const struct ggml_tensor * op, size_t * size);
  12. #ifdef __cplusplus
  13. }
  14. namespace ggml::cpu {
  15. // register in tensor->extra
  16. class tensor_traits {
  17. public:
  18. virtual ~tensor_traits();
  19. virtual bool work_size(int n_threads, const struct ggml_tensor * op, size_t & size) = 0;
  20. virtual bool compute_forward(struct ggml_compute_params * params, struct ggml_tensor * op) = 0;
  21. };
  22. class extra_buffer_type {
  23. public:
  24. virtual ~extra_buffer_type();
  25. virtual bool supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) = 0;
  26. virtual tensor_traits * get_tensor_traits(const struct ggml_tensor * op) = 0;
  27. };
  28. } // namespace ggml::cpu
  29. // implemented in ggml-cpu.cpp.
  30. std::vector<ggml_backend_buffer_type_t> & ggml_backend_cpu_get_extra_buffer_types();
  31. #endif