File size: 7,095 Bytes
c9c3262 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 | #include "ggml-remoting.h"
#include "ggml-virtgpu.h"
#include <iostream>
#include <mutex>
void ggml_virtgpu_cleanup(virtgpu * gpu);
static virtgpu * apir_initialize() {
static virtgpu * gpu = NULL;
static std::atomic<bool> initialized = false;
if (initialized) {
// fast track
return gpu;
}
{
static std::mutex mutex;
std::lock_guard<std::mutex> lock(mutex);
if (initialized) {
// thread safe
return gpu;
}
gpu = create_virtgpu();
if (!gpu) {
initialized = true;
return NULL;
}
// Pre-fetch and cache all device information, it will not change
gpu->cached_device_info.description = apir_device_get_description(gpu);
if (!gpu->cached_device_info.description) {
GGML_ABORT(GGML_VIRTGPU "%s: failed to initialize the virtgpu device description", __func__);
}
gpu->cached_device_info.device_count = apir_device_get_count(gpu);
gpu->cached_device_info.type = apir_device_get_type(gpu);
{
// Get the remote name and create prefixed version
char * rmt_device_name = apir_device_get_name(gpu);
if (!rmt_device_name) {
GGML_ABORT(GGML_VIRTGPU "%s: failed to get the virtgpu device name", __func__);
}
size_t device_name_len = strlen(rmt_device_name) + 11; // "[virtgpu] " + null terminator
gpu->cached_device_info.name = (char *) malloc(device_name_len);
if (!gpu->cached_device_info.name) {
free(rmt_device_name);
GGML_ABORT(GGML_VIRTGPU "%s: failed to allocate memory for prefixed device name", __func__);
}
snprintf(gpu->cached_device_info.name, device_name_len, "[virtgpu] %s", rmt_device_name);
free(rmt_device_name);
}
apir_device_get_memory(gpu, &gpu->cached_device_info.memory_free, &gpu->cached_device_info.memory_total);
apir_buffer_type_host_handle_t buft_host_handle = apir_device_get_buffer_type(gpu);
gpu->cached_buffer_type.host_handle = buft_host_handle;
{
// Get the remote name and create prefixed version
char * rmt_name = apir_buffer_type_get_name(gpu, buft_host_handle);
if (!rmt_name) {
GGML_ABORT(GGML_VIRTGPU "%s: failed to get the virtgpu buffer type name", __func__);
}
size_t prefixed_len = strlen(rmt_name) + 11; // "[virtgpu] " + null terminator
gpu->cached_buffer_type.name = (char *) malloc(prefixed_len);
if (!gpu->cached_buffer_type.name) {
free(rmt_name);
GGML_ABORT(GGML_VIRTGPU "%s: failed to allocate memory for prefixed buffer type name", __func__);
}
snprintf(gpu->cached_buffer_type.name, prefixed_len, "[virtgpu] %s", rmt_name);
free(rmt_name);
}
gpu->cached_buffer_type.alignment = apir_buffer_type_get_alignment(gpu, buft_host_handle);
gpu->cached_buffer_type.max_size = apir_buffer_type_get_max_size(gpu, buft_host_handle);
initialized = true;
}
return gpu;
}
static int ggml_backend_remoting_get_device_count() {
virtgpu * gpu = apir_initialize();
if (!gpu) {
return 0;
}
return gpu->cached_device_info.device_count;
}
static size_t ggml_backend_remoting_reg_get_device_count(ggml_backend_reg_t reg) {
UNUSED(reg);
return ggml_backend_remoting_get_device_count();
}
static std::vector<ggml_backend_dev_t> devices;
ggml_backend_dev_t ggml_backend_remoting_get_device(size_t device) {
GGML_ASSERT(device < devices.size());
return devices[device];
}
static void ggml_backend_remoting_reg_init_devices(ggml_backend_reg_t reg) {
if (devices.size() > 0) {
GGML_LOG_INFO(GGML_VIRTGPU "%s: already initialized\n", __func__);
return;
}
virtgpu * gpu = apir_initialize();
if (!gpu) {
GGML_LOG_ERROR(GGML_VIRTGPU "%s: apir_initialize failed\n", __func__);
return;
}
static std::atomic<bool> initialized = false;
if (initialized) {
return; // fast track
}
{
static std::mutex mutex;
std::lock_guard<std::mutex> lock(mutex);
if (!initialized) {
for (int i = 0; i < ggml_backend_remoting_get_device_count(); i++) {
ggml_backend_remoting_device_context * ctx = new ggml_backend_remoting_device_context;
char desc[256] = "ggml-virtgpu API Remoting device";
ctx->device = i;
ctx->name = GGML_VIRTGPU_NAME + std::to_string(i);
ctx->description = desc;
ctx->gpu = gpu;
ggml_backend_dev_t dev = new ggml_backend_device{
/* .iface = */ ggml_backend_remoting_device_interface,
/* .reg = */ reg,
/* .context = */ ctx,
};
devices.push_back(dev);
}
initialized = true;
}
}
}
static ggml_backend_dev_t ggml_backend_remoting_reg_get_device(ggml_backend_reg_t reg, size_t device) {
UNUSED(reg);
return ggml_backend_remoting_get_device(device);
}
static const char * ggml_backend_remoting_reg_get_name(ggml_backend_reg_t reg) {
UNUSED(reg);
return GGML_VIRTGPU_NAME;
}
static const ggml_backend_reg_i ggml_backend_remoting_reg_i = {
/* .get_name = */ ggml_backend_remoting_reg_get_name,
/* .get_device_count = */ ggml_backend_remoting_reg_get_device_count,
/* .get_device = */ ggml_backend_remoting_reg_get_device,
/* .get_proc_address = */ NULL,
};
ggml_backend_reg_t ggml_backend_virtgpu_reg() {
virtgpu * gpu = apir_initialize();
if (!gpu) {
GGML_LOG_ERROR(GGML_VIRTGPU "%s: virtgpu_apir_initialize failed\n", __func__);
}
static ggml_backend_reg reg = {
/* .api_version = */ GGML_BACKEND_API_VERSION,
/* .iface = */ ggml_backend_remoting_reg_i,
/* .context = */ gpu,
};
static bool initialized = false;
if (initialized) {
return ®
}
initialized = true;
ggml_backend_remoting_reg_init_devices(®);
return ®
}
// public function, not exposed in the GGML interface at the moment
void ggml_virtgpu_cleanup(virtgpu * gpu) {
if (gpu->cached_device_info.name) {
free(gpu->cached_device_info.name);
gpu->cached_device_info.name = NULL;
}
if (gpu->cached_device_info.description) {
free(gpu->cached_device_info.description);
gpu->cached_device_info.description = NULL;
}
if (gpu->cached_buffer_type.name) {
free(gpu->cached_buffer_type.name);
gpu->cached_buffer_type.name = NULL;
}
mtx_destroy(&gpu->data_shmem_mutex);
}
GGML_BACKEND_DL_IMPL(ggml_backend_virtgpu_reg)
|