char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Kevin Pougetggml-virtgpu: improve the reliability of the code (#19846)ffaafde16

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