1#include "ggml-remoting.h"
  2
  3#define BUFFER_TO_GPU(name) ((ggml_backend_remoting_buffer_context *) (name)->context)->gpu
  4
  5static void * ggml_backend_remoting_buffer_get_base(ggml_backend_buffer_t buffer) {
  6    ggml_backend_remoting_buffer_context * context = (ggml_backend_remoting_buffer_context *) buffer->context;
  7    if (context->base) {
  8        return context->base;
  9    }
 10
 11    context->base = apir_buffer_get_base(BUFFER_TO_GPU(buffer), BUFFER_TO_APIR_CONTEXT(buffer));
 12
 13    return context->base;
 14}
 15
 16static void ggml_backend_remoting_buffer_set_tensor(ggml_backend_buffer_t buffer,
 17                                                    ggml_tensor *         tensor,
 18                                                    const void *          data,
 19                                                    size_t                offset,
 20                                                    size_t                size) {
 21    virtgpu * gpu = BUFFER_TO_GPU(buffer);
 22
 23    ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer);
 24    if (context->is_from_ptr) {
 25        memcpy((char *) tensor->data + offset, data, size);
 26    } else {
 27        apir_buffer_set_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), tensor, data, offset, size);
 28    }
 29
 30    return;
 31}
 32
 33static void ggml_backend_remoting_buffer_get_tensor(ggml_backend_buffer_t buffer,
 34                                                    const ggml_tensor *   tensor,
 35                                                    void *                data,
 36                                                    size_t                offset,
 37                                                    size_t                size) {
 38    virtgpu *                              gpu     = BUFFER_TO_GPU(buffer);
 39    ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer);
 40    if (context->is_from_ptr) {
 41        memcpy(data, (const char *) tensor->data + offset, size);
 42    } else {
 43        apir_buffer_get_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), tensor, data, offset, size);
 44    }
 45}
 46
 47static void ggml_backend_remoting_buffer_set_tensor_from_ptr(ggml_backend_buffer_t buffer,
 48                                                             ggml_tensor *         tensor,
 49                                                             const void *          data,
 50                                                             size_t                offset,
 51                                                             size_t                size) {
 52    UNUSED(buffer);
 53
 54    memcpy((char *) tensor->data + offset, data, size);
 55
 56    return;
 57}
 58
 59static void ggml_backend_remoting_buffer_get_tensor_from_ptr(ggml_backend_buffer_t buffer,
 60                                                             const ggml_tensor *   tensor,
 61                                                             void *                data,
 62                                                             size_t                offset,
 63                                                             size_t                size) {
 64    UNUSED(buffer);
 65
 66    memcpy(data, (const char *) tensor->data + offset, size);
 67}
 68
 69static bool ggml_backend_remoting_buffer_cpy_tensor(ggml_backend_buffer_t buffer,
 70                                                    const ggml_tensor *   src,
 71                                                    ggml_tensor *         dst) {
 72    virtgpu * gpu = BUFFER_TO_GPU(buffer);
 73
 74    bool ret = apir_buffer_cpy_tensor(gpu, BUFFER_TO_APIR_CONTEXT(buffer), src, dst);
 75
 76    return ret;
 77}
 78
 79static void ggml_backend_remoting_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) {
 80    virtgpu * gpu = BUFFER_TO_GPU(buffer);
 81
 82    apir_buffer_clear(gpu, BUFFER_TO_APIR_CONTEXT(buffer), value);
 83
 84    return;
 85}
 86
 87static void ggml_backend_remoting_buffer_free_buffer(ggml_backend_buffer_t buffer) {
 88    virtgpu * gpu = BUFFER_TO_GPU(buffer);
 89
 90    apir_buffer_free_buffer(gpu, BUFFER_TO_APIR_CONTEXT(buffer));
 91
 92    ggml_backend_remoting_buffer_context * context = BUFFER_TO_GGML_CONTEXT(buffer);
 93    free(context);
 94    buffer->context = NULL;
 95}
 96
 97const ggml_backend_buffer_i ggml_backend_remoting_buffer_interface = {
 98    /* .free_buffer     = */ ggml_backend_remoting_buffer_free_buffer,
 99    /* .get_base        = */ ggml_backend_remoting_buffer_get_base,
100    /* .init_tensor     = */ NULL,
101    /* .memset_tensor   = */ NULL,
102    /* .set_tensor      = */ ggml_backend_remoting_buffer_set_tensor,
103    /* .get_tensor      = */ ggml_backend_remoting_buffer_get_tensor,
104    /* .cpy_tensor      = */ ggml_backend_remoting_buffer_cpy_tensor,
105    /* .clear           = */ ggml_backend_remoting_buffer_clear,
106    /* .reset           = */ NULL,
107};
108
109const ggml_backend_buffer_i ggml_backend_remoting_buffer_from_ptr_interface = {
110    /* .free_buffer     = */ ggml_backend_remoting_buffer_free_buffer,
111    /* .get_base        = */ ggml_backend_remoting_buffer_get_base,
112    /* .init_tensor     = */ NULL,
113    /* .memset_tensor   = */ NULL,
114    /* .set_tensor      = */ ggml_backend_remoting_buffer_set_tensor_from_ptr,
115    /* .get_tensor      = */ ggml_backend_remoting_buffer_get_tensor_from_ptr,
116    /* .cpy_tensor      = */ ggml_backend_remoting_buffer_cpy_tensor,
117    /* .clear           = */ ggml_backend_remoting_buffer_clear,
118    /* .reset           = */ NULL,
119};