1#include "virtgpu-forward-impl.h"
  2
  3void * apir_buffer_get_base(virtgpu * gpu, apir_buffer_context_t * buffer_context) {
  4    apir_encoder *        encoder;
  5    apir_decoder *        decoder;
  6    ApirForwardReturnCode ret;
  7
  8    REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_GET_BASE);
  9
 10    apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle);
 11
 12    REMOTE_CALL(gpu, encoder, decoder, ret);
 13
 14    uintptr_t base;
 15    apir_decode_uintptr_t(decoder, &base);
 16
 17    remote_call_finish(gpu, encoder, decoder);
 18
 19    return (void *) base;
 20}
 21
 22void apir_buffer_set_tensor(virtgpu *               gpu,
 23                            apir_buffer_context_t * buffer_context,
 24                            ggml_tensor *           tensor,
 25                            const void *            data,
 26                            size_t                  offset,
 27                            size_t                  size) {
 28    apir_encoder *        encoder;
 29    apir_decoder *        decoder;
 30    ApirForwardReturnCode ret;
 31
 32    REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_SET_TENSOR);
 33
 34    apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle);
 35    apir_encode_ggml_tensor(encoder, tensor);
 36
 37    virtgpu_shmem   temp_shmem;  // Local storage for large buffers
 38    virtgpu_shmem * shmem = &temp_shmem;
 39    bool using_shared_shmem = false;
 40
 41    if (size <= gpu->data_shmem.mmap_size) {
 42        // Lock mutex before using shared data_shmem buffer
 43        if (mtx_lock(&gpu->data_shmem_mutex) != thrd_success) {
 44            GGML_ABORT(GGML_VIRTGPU "%s: Failed to lock data_shmem mutex", __func__);
 45        }
 46        using_shared_shmem = true;
 47        shmem = &gpu->data_shmem;
 48
 49    } else if (virtgpu_shmem_create(gpu, size, shmem)) {
 50        GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the guest-host shared buffer", __func__);
 51    }
 52
 53    memcpy(shmem->mmap_ptr, data, size);
 54    apir_encode_virtgpu_shmem_res_id(encoder, shmem->res_id);
 55
 56    apir_encode_size_t(encoder, &offset);
 57    apir_encode_size_t(encoder, &size);
 58
 59    REMOTE_CALL(gpu, encoder, decoder, ret);
 60
 61    remote_call_finish(gpu, encoder, decoder);
 62
 63    // Unlock mutex before cleanup
 64    if (using_shared_shmem) {
 65        mtx_unlock(&gpu->data_shmem_mutex);
 66    } else {
 67        virtgpu_shmem_destroy(gpu, shmem);
 68    }
 69
 70    return;
 71}
 72
 73void apir_buffer_get_tensor(virtgpu *               gpu,
 74                            apir_buffer_context_t * buffer_context,
 75                            const ggml_tensor *     tensor,
 76                            void *                  data,
 77                            size_t                  offset,
 78                            size_t                  size) {
 79    apir_encoder *        encoder;
 80    apir_decoder *        decoder;
 81    ApirForwardReturnCode ret;
 82
 83    REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_GET_TENSOR);
 84
 85    apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle);
 86    apir_encode_ggml_tensor(encoder, tensor);
 87
 88    virtgpu_shmem   temp_shmem;  // Local storage for large buffers
 89    virtgpu_shmem * shmem = &temp_shmem;
 90    bool using_shared_shmem = false;
 91
 92    if (size <= gpu->data_shmem.mmap_size) {
 93        // Lock mutex before using shared data_shmem buffer
 94        if (mtx_lock(&gpu->data_shmem_mutex) != thrd_success) {
 95            GGML_ABORT(GGML_VIRTGPU "%s: Failed to lock data_shmem mutex", __func__);
 96        }
 97        using_shared_shmem = true;
 98        shmem = &gpu->data_shmem;
 99
100    } else if (virtgpu_shmem_create(gpu, size, shmem)) {
101        GGML_ABORT(GGML_VIRTGPU "%s: Couldn't allocate the guest-host shared buffer", __func__);
102    }
103
104    apir_encode_virtgpu_shmem_res_id(encoder, shmem->res_id);
105    apir_encode_size_t(encoder, &offset);
106    apir_encode_size_t(encoder, &size);
107
108    REMOTE_CALL(gpu, encoder, decoder, ret);
109
110    memcpy(data, shmem->mmap_ptr, size);
111
112    remote_call_finish(gpu, encoder, decoder);
113
114    // Unlock mutex before cleanup
115    if (using_shared_shmem) {
116        mtx_unlock(&gpu->data_shmem_mutex);
117    } else {
118        virtgpu_shmem_destroy(gpu, shmem);
119    }
120}
121
122bool apir_buffer_cpy_tensor(virtgpu *               gpu,
123                            apir_buffer_context_t * buffer_context,
124                            const ggml_tensor *     src,
125                            const ggml_tensor *     dst) {
126    apir_encoder *        encoder;
127    apir_decoder *        decoder;
128    ApirForwardReturnCode ret;
129
130    REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_CPY_TENSOR);
131
132    apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle);
133    apir_encode_ggml_tensor(encoder, src);
134    apir_encode_ggml_tensor(encoder, dst);
135
136    REMOTE_CALL(gpu, encoder, decoder, ret);
137
138    bool ret_val;
139    apir_decode_bool_t(decoder, &ret_val);
140
141    remote_call_finish(gpu, encoder, decoder);
142
143    return ret_val;
144}
145
146void apir_buffer_clear(virtgpu * gpu, apir_buffer_context_t * buffer_context, uint8_t value) {
147    apir_encoder *        encoder;
148    apir_decoder *        decoder;
149    ApirForwardReturnCode ret;
150
151    REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_CLEAR);
152
153    apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle);
154    apir_encode_uint8_t(encoder, &value);
155
156    REMOTE_CALL(gpu, encoder, decoder, ret);
157
158    remote_call_finish(gpu, encoder, decoder);
159}
160
161void apir_buffer_free_buffer(virtgpu * gpu, apir_buffer_context_t * buffer_context) {
162    apir_encoder *        encoder;
163    apir_decoder *        decoder;
164    ApirForwardReturnCode ret;
165
166    REMOTE_CALL_PREPARE(gpu, encoder, APIR_COMMAND_TYPE_BUFFER_FREE_BUFFER);
167
168    apir_encode_apir_buffer_host_handle_t(encoder, &buffer_context->host_handle);
169
170    REMOTE_CALL(gpu, encoder, decoder, ret);
171
172    remote_call_finish(gpu, encoder, decoder);
173}