forked from ggerganov/ggml
-
Notifications
You must be signed in to change notification settings - Fork 0
/
test-arange.cpp
99 lines (76 loc) · 2.56 KB
/
test-arange.cpp
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
#include "ggml/ggml.h"
#include "ggml/ggml-alloc.h"
#include "ggml/ggml-backend.h"
#ifdef GGML_USE_CUBLAS
#include "ggml-cuda.h"
#endif
#ifdef GGML_USE_METAL
#include "ggml-metal.h"
#endif
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
int main(int /*argc*/, const char** /*argv*/) {
{
bool use_gpu = true;
GGML_UNUSED(use_gpu);
ggml_backend_t backend = NULL;
//ggml_backend_buffer_t buffer;
#ifdef GGML_USE_CUBLAS
if (use_gpu) {
fprintf(stderr, "%s: using CUDA backend\n", __func__);
backend = ggml_backend_cuda_init(0);
if (!backend) {
fprintf(stderr, "%s: ggml_backend_cuda_init() failed\n", __func__);
}
}
#endif
#ifdef GGML_USE_METAL
if (!backend) {
fprintf(stderr, "%s: using Metal backend\n", __func__);
backend = ggml_backend_metal_init();
if (!backend) {
fprintf(stderr, "%s: ggml_backend_metal_init() failed\n", __func__);
}
}
#endif
const int num_tensors = 2;
struct ggml_init_params params = {
/*.mem_size =*/ ggml_tensor_overhead() * num_tensors + 2 * 1024 * 1024,
/*.mem_size =*/ NULL,
/*.mem_size =*/ true,
};
if (!backend) {
// fallback to CPU backend
backend = ggml_backend_cpu_init();
}
// create context
struct ggml_context* ctx = ggml_init(params);
struct ggml_tensor * t = ggml_arange(ctx, 0, 3, 1);
GGML_ASSERT(t->ne[0] == 3);
ggml_gallocr_t galloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
struct ggml_cgraph * graph = ggml_new_graph(ctx);
ggml_build_forward_expand(graph, t);
// allocate tensors
ggml_gallocr_alloc_graph(galloc, graph);
int n_threads = 4;
if (ggml_backend_is_cpu(backend)) {
ggml_backend_cpu_set_n_threads(backend, n_threads);
}
ggml_backend_graph_compute(backend, graph);
float * output = new float[ggml_nelements(t)];
ggml_backend_tensor_get(t, output, 0, ggml_nbytes(t));
for (int i = 0; i < t->ne[0]; i++) {
printf("%.2f ", output[i]);
}
printf("\n");
GGML_ASSERT(output[0] == 0);
GGML_ASSERT(output[1] == 1);
GGML_ASSERT(output[2] == 2);
delete[] output;
ggml_free(ctx);
ggml_gallocr_free(galloc);
ggml_backend_free(backend);
}
return 0;
}