-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathbackend.cpp
More file actions
151 lines (126 loc) · 4.91 KB
/
Copy pathbackend.cpp
File metadata and controls
151 lines (126 loc) · 4.91 KB
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
#include "backend.h"
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <stdexcept>
#include <strings.h>
#include <thread>
#include "ggml-cpu.h"
namespace muscriptor {
namespace {
BackendPair g_backend_pair = {nullptr, nullptr, nullptr, false};
bool g_initialized = false;
int get_physical_cores() {
int n = static_cast<int>(std::thread::hardware_concurrency());
return n > 1 ? n / 2 : 1;
}
} // namespace
BackendPair backend_init(const char* name) {
if (g_initialized) {
return g_backend_pair;
}
BackendPair bp = {nullptr, nullptr, nullptr, false};
const char* env_backend = std::getenv("GGML_BACKEND");
if (env_backend) {
if (strcasecmp(env_backend, "CPU") == 0) {
bp.backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
} else if (strcasecmp(env_backend, "Metal") == 0 ||
strcasecmp(env_backend, "GPU") == 0 ||
strcasecmp(env_backend, "CUDA") == 0) {
bp.backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_GPU, nullptr);
} else {
bp.backend = ggml_backend_init_by_name(env_backend, nullptr);
}
if (!bp.backend) {
std::fprintf(stderr, "[%s] warning: requested backend '%s' not found\n", name,
env_backend);
}
}
if (!bp.backend) {
bp.backend = ggml_backend_init_best();
}
if (!bp.backend) {
std::fprintf(stderr, "[%s] error: no backend available\n", name);
throw std::runtime_error("No GGML backend available");
}
ggml_backend_dev_t dev = ggml_backend_get_device(bp.backend);
const char* type_name = "unknown";
if (dev) {
enum ggml_backend_dev_type dev_type = ggml_backend_dev_type(dev);
// Dedicated GPU and integrated/unified-memory GPU (e.g. GB10) both count.
bp.has_gpu = (dev_type == GGML_BACKEND_DEVICE_TYPE_GPU ||
dev_type == GGML_BACKEND_DEVICE_TYPE_IGPU);
switch (dev_type) {
case GGML_BACKEND_DEVICE_TYPE_CPU: type_name = "CPU"; break;
case GGML_BACKEND_DEVICE_TYPE_GPU: type_name = "GPU"; break;
case GGML_BACKEND_DEVICE_TYPE_IGPU: type_name = "IGPU"; break;
case GGML_BACKEND_DEVICE_TYPE_ACCEL: type_name = "ACCEL"; break;
case GGML_BACKEND_DEVICE_TYPE_META: type_name = "META"; break;
}
}
std::fprintf(stderr, "[%s] primary backend: %s (%s)%s\n",
name, ggml_backend_name(bp.backend), type_name,
bp.has_gpu ? "" : " — host compute");
bp.cpu_backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
if (bp.cpu_backend) {
int n_threads = get_physical_cores();
ggml_backend_reg_t cpu_reg = ggml_backend_dev_backend_reg(
ggml_backend_get_device(bp.cpu_backend));
if (cpu_reg) {
auto set_threads = (ggml_backend_set_n_threads_t)
ggml_backend_reg_get_proc_address(cpu_reg, "ggml_backend_set_n_threads");
if (set_threads) {
set_threads(bp.cpu_backend, n_threads);
std::fprintf(stderr, "[%s] CPU backend: %d threads\n", name, n_threads);
}
}
}
if (bp.cpu_backend && bp.backend == bp.cpu_backend) {
bp.cpu_backend = nullptr;
}
g_backend_pair = bp;
g_initialized = true;
return bp;
}
ggml_backend_sched_t backend_sched_new(const BackendPair& bp, size_t graph_size) {
ggml_backend_t backends[3];
ggml_backend_buffer_type_t buffer_types[3];
int n_backends = 0;
backends[n_backends] = bp.backend;
buffer_types[n_backends] = ggml_backend_get_default_buffer_type(bp.backend);
n_backends++;
if (bp.cpu_backend && bp.cpu_backend != bp.backend) {
backends[n_backends] = bp.cpu_backend;
if (bp.has_gpu) {
ggml_backend_dev_t dev = ggml_backend_get_device(bp.backend);
ggml_backend_buffer_type_t host_buft = ggml_backend_dev_host_buffer_type(dev);
buffer_types[n_backends] = host_buft ? host_buft : ggml_backend_cpu_buffer_type();
} else {
buffer_types[n_backends] = ggml_backend_cpu_buffer_type();
}
n_backends++;
}
ggml_backend_sched_t sched = ggml_backend_sched_new(
backends, buffer_types, n_backends, graph_size, false, true);
if (!sched) {
throw std::runtime_error("Failed to create backend scheduler");
}
return sched;
}
void backend_release(ggml_backend_t backend, ggml_backend_t cpu_backend) {
g_initialized = false;
if (cpu_backend && cpu_backend != backend) {
ggml_backend_free(cpu_backend);
}
if (backend) {
ggml_backend_free(backend);
}
g_backend_pair = {nullptr, nullptr, nullptr, false};
}
BackendPair& global_backend_pair() {
if (!g_initialized) {
backend_init("global");
}
return g_backend_pair;
}
} // namespace muscriptor