-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmetal_decode.cpp
More file actions
258 lines (222 loc) · 7.7 KB
/
Copy pathmetal_decode.cpp
File metadata and controls
258 lines (222 loc) · 7.7 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
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
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
#if defined(__APPLE__) \
&& (defined(__aarch64__) || defined(__arm64__)) \
&& !defined(LLAMACPP_METAL_TEST)
#error "LLAMACPP_METAL_TEST must be enabled on macOS ARM64"
#endif
#ifdef LLAMACPP_METAL_TEST
import std;
import llamacpp;
#ifdef LLAMA_H
#error "import llamacpp leaked LLAMA_H"
#endif
#ifdef LLAMA_API
#error "import llamacpp leaked LLAMA_API"
#endif
namespace {
std::string logs;
void capture_log(enum ggml_log_level, const char * text, void *) {
if (text) {
logs += text;
std::cerr << text;
}
}
int fail(const char * message) {
std::cerr << "Metal smoke test failed: " << message << "\n";
return 1;
}
bool has_positive_buffer(const std::string & text, const std::regex & pattern) {
std::smatch match;
return std::regex_search(text, match, pattern)
&& std::stod(match[1].str()) > 0.0;
}
bool has_positive_metal_model_buffer(const std::string & text) {
static const std::regex pattern(
R"(MTL[0-9]+(?:_[A-Za-z0-9]+)? model buffer size\s*=\s*([0-9]+(?:\.[0-9]+)?) MiB)"
);
return has_positive_buffer(text, pattern);
}
bool has_positive_metal_compute_buffer(const std::string & text) {
static const std::regex pattern(
R"(MTL[0-9]+ compute buffer size\s*=\s*([0-9]+(?:\.[0-9]+)?) MiB)"
);
return has_positive_buffer(text, pattern);
}
bool has_embedded_library(ggml_backend_reg_t registry) {
auto get_features = reinterpret_cast<ggml_backend_get_features_t>(
ggml_backend_reg_get_proc_address(
registry, "ggml_backend_get_features"
)
);
if (!get_features) return false;
for (auto * feature = get_features(registry);
feature && feature->name;
++feature) {
if (std::strcmp(feature->name, "EMBED_LIBRARY") == 0
&& std::strcmp(feature->value, "1") == 0) {
return true;
}
}
return false;
}
bool run_metal_add_probe(ggml_backend_dev_t device) {
ggml_backend_t backend = ggml_backend_dev_init(device, nullptr);
if (!backend) return false;
ggml_init_params params = {};
params.mem_size = 1024 * 1024;
params.no_alloc = true;
ggml_context * context = ggml_init(params);
if (!context) {
ggml_backend_free(backend);
return false;
}
ggml_cgraph * graph = ggml_new_graph(context);
ggml_tensor * lhs = ggml_new_tensor_1d(context, GGML_TYPE_F32, 4);
ggml_tensor * rhs = ggml_new_tensor_1d(context, GGML_TYPE_F32, 4);
ggml_tensor * sum = ggml_add(context, lhs, rhs);
ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(
context, backend
);
if (!graph || !lhs || !rhs || !sum || !buffer) {
if (buffer) ggml_backend_buffer_free(buffer);
ggml_free(context);
ggml_backend_free(backend);
return false;
}
ggml_build_forward_expand(graph, sum);
const float lhs_values[] = {1.0F, -2.0F, 3.5F, 10.0F};
const float rhs_values[] = {4.0F, 5.0F, -1.5F, -3.0F};
ggml_backend_tensor_set(lhs, lhs_values, 0, sizeof(lhs_values));
ggml_backend_tensor_set(rhs, rhs_values, 0, sizeof(rhs_values));
bool passed = ggml_backend_graph_compute(backend, graph)
== GGML_STATUS_SUCCESS;
float actual[4] = {};
if (passed) {
ggml_backend_synchronize(backend);
ggml_backend_tensor_get(sum, actual, 0, sizeof(actual));
const float expected[] = {5.0F, 3.0F, 2.0F, 7.0F};
for (std::size_t index = 0; index < 4; ++index) {
if (actual[index] != expected[index]) {
passed = false;
break;
}
}
}
ggml_backend_buffer_free(buffer);
ggml_free(context);
ggml_backend_free(backend);
return passed;
}
} // namespace
int main() {
const char * model_path = std::getenv("LLAMACPP_TEST_MODEL");
if (!model_path || !*model_path) {
return fail("LLAMACPP_TEST_MODEL is not set");
}
llama_log_set(capture_log, nullptr);
llama_backend_init();
ggml_backend_reg_t metal = ggml_backend_reg_by_name("MTL");
if (!metal || ggml_backend_reg_dev_count(metal) == 0) {
llama_backend_free();
return fail("MTL registry or device is missing");
}
ggml_backend_dev_t device = ggml_backend_reg_dev_get(metal, 0);
if (!device || ggml_backend_dev_type(device) != GGML_BACKEND_DEVICE_TYPE_GPU) {
llama_backend_free();
return fail("MTL device is not a GPU");
}
if (!has_embedded_library(metal)) {
llama_backend_free();
return fail("MTL registry does not report EMBED_LIBRARY=1");
}
if (!llama_supports_gpu_offload()) {
llama_backend_free();
return fail("llama does not report GPU offload support");
}
if (!run_metal_add_probe(device)) {
llama_backend_free();
return fail("MTL F32 ADD graph did not execute correctly");
}
if (logs.find("using embedded metal library") == std::string::npos) {
llama_backend_free();
return fail("embedded Metal source path was not used by the probe");
}
logs.clear();
llama_model_params model_params = llama_model_default_params();
model_params.n_gpu_layers = std::numeric_limits<int>::max();
llama_model * model = llama_model_load_from_file(model_path, model_params);
if (!model) {
llama_backend_free();
return fail("model load failed");
}
const std::regex offload_pattern(
"offloaded ([1-9][0-9]*)/([1-9][0-9]*) layers to GPU"
);
if (!std::regex_search(logs, offload_pattern)) {
llama_model_free(model);
llama_backend_free();
return fail("positive GPU layer offload was not logged");
}
if (!has_positive_metal_model_buffer(logs)) {
llama_model_free(model);
llama_backend_free();
return fail("positive Metal model buffer was not logged");
}
llama_context_params context_params = llama_context_default_params();
context_params.n_ctx = 64;
llama_context * context = llama_init_from_model(model, context_params);
if (!context) {
llama_model_free(model);
llama_backend_free();
return fail("context creation failed");
}
if (!has_positive_metal_compute_buffer(logs)) {
llama_free(context);
llama_model_free(model);
llama_backend_free();
return fail("positive Metal compute buffer was not logged");
}
llama_token tokens[] = {1, 2, 3};
const int decode_result = llama_decode(
context,
llama_batch_get_one(tokens, sizeof(tokens) / sizeof(tokens[0]))
);
if (decode_result != 0) {
llama_free(context);
llama_model_free(model);
llama_backend_free();
return fail("decode failed");
}
llama_sampler * sampler = llama_sampler_chain_init(
llama_sampler_chain_default_params()
);
llama_sampler_chain_add(sampler, llama_sampler_init_greedy());
const llama_token sampled = llama_sampler_sample(sampler, context, -1);
std::cout << "sampled token: " << sampled << "\n";
const int vocabulary_size = llama_vocab_n_tokens(
llama_model_get_vocab(model)
);
if (sampled < 0 || sampled >= vocabulary_size) {
llama_sampler_free(sampler);
llama_free(context);
llama_model_free(model);
llama_backend_free();
return fail("sampled token is outside the vocabulary");
}
llama_sampler_free(sampler);
llama_free(context);
llama_model_free(model);
llama_backend_free();
std::cout << "LLAMACPP_METAL_TEST=PASS\n";
return 0;
}
#else
#if defined(__APPLE__) \
&& (defined(__aarch64__) || defined(__arm64__))
#error "Metal smoke test cannot skip on macOS ARM64"
#endif
import std;
int main() {
std::cout << "Metal smoke test skipped on unsupported target\n";
return 0;
}
#endif