95 lines
3.4 KiB
C++
95 lines
3.4 KiB
C++
#include "ggml.h"
|
|
#include "ggml-alloc.h"
|
|
#include "ggml-backend.h"
|
|
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <vector>
|
|
|
|
static void require(bool ok) {
|
|
if (!ok) std::exit(1);
|
|
}
|
|
|
|
struct graph {
|
|
ggml_context * ctx;
|
|
ggml_tensor * input;
|
|
ggml_tensor * output;
|
|
ggml_cgraph * gf;
|
|
graph(int size) {
|
|
ctx = ggml_init({1024 * 1024, nullptr, true});
|
|
input = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, size);
|
|
ggml_set_input(input);
|
|
output = ggml_scale(ctx, input, 2.0f);
|
|
ggml_set_output(output);
|
|
gf = ggml_new_graph(ctx);
|
|
ggml_build_forward_expand(gf, output);
|
|
}
|
|
~graph() { ggml_free(ctx); }
|
|
void check(ggml_backend_t backend, float value) {
|
|
std::vector<float> data(ggml_nelements(input), value);
|
|
ggml_backend_tensor_set(input, data.data(), 0, ggml_nbytes(input));
|
|
require(ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS);
|
|
ggml_backend_tensor_get(output, data.data(), 0, ggml_nbytes(output));
|
|
for (float x : data) require(x == 2 * value);
|
|
}
|
|
};
|
|
|
|
int main() {
|
|
ggml_backend_load_all();
|
|
auto backend = ggml_backend_init_by_name("CUDA0", nullptr);
|
|
require(backend != nullptr);
|
|
auto buft = ggml_backend_get_default_buffer_type(backend);
|
|
for (int grow_dst : {0, 1}) {
|
|
auto src = ggml_gallocr_new(buft);
|
|
auto dst = ggml_gallocr_new(buft);
|
|
graph a(1024), b(1024), big(1024 * 1024);
|
|
require(ggml_gallocr_alloc_graph(src, a.gf));
|
|
require(ggml_gallocr_share_buffers(dst, src));
|
|
require(ggml_gallocr_alloc_graph(dst, b.gf));
|
|
require(a.input->buffer == b.input->buffer);
|
|
for (int i = 1; i <= 10; ++i) {
|
|
a.check(backend, i);
|
|
b.check(backend, -i);
|
|
}
|
|
require(ggml_gallocr_alloc_graph(grow_dst ? dst : src, big.gf));
|
|
require(big.input->buffer != (grow_dst ? a.input : b.input)->buffer);
|
|
big.check(backend, 0.5f);
|
|
(grow_dst ? a : b).check(backend, 4);
|
|
ggml_gallocr_free(grow_dst ? dst : src);
|
|
(grow_dst ? a : b).check(backend, 5);
|
|
ggml_gallocr_free(grow_dst ? src : dst);
|
|
}
|
|
for (int free_dst : {0, 1}) {
|
|
auto src = ggml_gallocr_new(buft);
|
|
auto dst = ggml_gallocr_new(buft);
|
|
graph a(1024), b(1024);
|
|
require(ggml_gallocr_alloc_graph(src, a.gf));
|
|
require(ggml_gallocr_share_buffers(dst, src));
|
|
require(ggml_gallocr_alloc_graph(dst, b.gf));
|
|
ggml_gallocr_free(free_dst ? dst : src);
|
|
(free_dst ? a : b).check(backend, 7);
|
|
ggml_gallocr_free(free_dst ? src : dst);
|
|
}
|
|
for (int larger_dst : {0, 1}) {
|
|
auto src = ggml_gallocr_new(buft);
|
|
auto dst = ggml_gallocr_new(buft);
|
|
{
|
|
graph a(larger_dst ? 1024 : 2048), b(larger_dst ? 2048 : 1024);
|
|
require(ggml_gallocr_alloc_graph(src, a.gf));
|
|
require(ggml_gallocr_alloc_graph(dst, b.gf));
|
|
}
|
|
require(ggml_gallocr_share_buffers(dst, src));
|
|
graph a(1024), b(1024);
|
|
require(ggml_gallocr_alloc_graph(src, a.gf));
|
|
require(ggml_gallocr_alloc_graph(dst, b.gf));
|
|
require(a.input->buffer == b.input->buffer);
|
|
a.check(backend, 9);
|
|
b.check(backend, 10);
|
|
ggml_gallocr_free(src);
|
|
b.check(backend, 11);
|
|
ggml_gallocr_free(dst);
|
|
}
|
|
ggml_backend_free(backend);
|
|
std::puts("PASS: shared buffers, alternating graphs, growth detachment, both free orders, larger-buffer selection");
|
|
}
|