214 lines
6.7 KiB
C++
214 lines
6.7 KiB
C++
#include "ggml.h"
|
||
#include "ggml-cpu.h"
|
||
#include "ggml-alloc.h"
|
||
#include "ggml-backend.h"
|
||
|
||
#ifdef GGML_USE_CUDA
|
||
#include "ggml-cuda.h"
|
||
#endif
|
||
|
||
#ifdef GGML_USE_METAL
|
||
#include "ggml-metal.h"
|
||
#endif
|
||
|
||
#include <string.h>
|
||
#include <stdio.h>
|
||
#include <stdlib.h>
|
||
|
||
|
||
static void ggml_log_callback_default(ggml_log_level level, const char * text, void * user_data) {
|
||
(void) level;
|
||
(void) user_data;
|
||
fputs(text, stderr);
|
||
fflush(stderr);
|
||
}
|
||
|
||
struct ggml_context* make_ctx(void) {
|
||
struct ggml_init_params params = {
|
||
/*.mem_size =*/ 2 * 1024 * 1024,
|
||
/*.mem_buffer =*/ nullptr,
|
||
/*.no_alloc. =*/ false
|
||
};
|
||
return ggml_init(params);
|
||
}
|
||
|
||
void check_tensor(struct ggml_tensor * t, float * expected_t_d, int ne0, int ne1, int ne2) {
|
||
GGML_ASSERT(t->type == GGML_TYPE_F32);
|
||
GGML_ASSERT(t->ne[0] == ne0);
|
||
GGML_ASSERT(t->ne[1] == ne1);
|
||
GGML_ASSERT(t->ne[2] == ne2);
|
||
for (int i2 = 0; i2 < ne2; ++i2) {
|
||
for (int i1 = 0; i1 < ne1; ++i1) {
|
||
for (int i0 = 0; i0 < ne0; ++i0) {
|
||
float expected = *(expected_t_d + i2 * ne1 * ne0 + i1 * ne0 + i0);
|
||
float actual = ggml_get_data_f32(t)[i2 * ne1 * ne0 + i1 * ne0 + i0];
|
||
if (expected != actual) {
|
||
printf("expected %.1f, got %.1f at (%d,%d,%d)\n", expected, actual, i0, i1, i2);
|
||
}
|
||
GGML_ASSERT(expected == actual);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
void test_pad_reflect_1d(bool use_gpu) {
|
||
ggml_backend_t backend = NULL;
|
||
struct ggml_init_params params;
|
||
ggml_backend_buffer_t buffer;
|
||
struct ggml_context * ctx;
|
||
struct ggml_tallocr tallocr;
|
||
ggml_gallocr_t gallocr;
|
||
|
||
// initialize the backend
|
||
#ifdef GGML_USE_CUDA
|
||
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 (use_gpu) {
|
||
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
|
||
|
||
if (!backend) {
|
||
fprintf(stderr, "%s: using CPU backend\n", __func__);
|
||
backend = ggml_backend_cpu_init();
|
||
}
|
||
|
||
// Test cases for different padding configurations
|
||
{
|
||
params = ggml_init_params{
|
||
/*.mem_size =*/ 16*1024*1024,
|
||
/*.mem_buffer =*/ nullptr,
|
||
/*.no_alloc. =*/ true
|
||
};
|
||
|
||
ggml_log_set(ggml_log_callback_default, nullptr);
|
||
|
||
ctx = ggml_init(params);
|
||
buffer = ggml_backend_alloc_buffer(backend, 16*1024*1024);
|
||
tallocr = ggml_tallocr_new(buffer);
|
||
gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
|
||
|
||
// Create a simple 1D input tensor [1, 2, 3, 4]
|
||
struct ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 4);
|
||
float input_data[] = {1.0f, 2.0f, 3.0f, 4.0f};
|
||
ggml_tallocr_alloc(&tallocr, t);
|
||
|
||
// load data to buffer
|
||
if(ggml_backend_is_cpu(backend)) {
|
||
memcpy(t->data, input_data, ggml_nbytes(t));
|
||
} else {
|
||
ggml_backend_tensor_set(t, input_data, 0, ggml_nbytes(t));
|
||
}
|
||
|
||
// Test case 1: pad left=1, right=1
|
||
// Expected: [2, 1, 2, 3, 4, 3]
|
||
float expected_1[] = {2.0f, 1.0f, 2.0f, 3.0f, 4.0f, 3.0f};
|
||
struct ggml_tensor * out_1 = ggml_pad_reflect_1d(ctx, t, 1, 1);
|
||
|
||
// Test case 2: pad left=2, right=1
|
||
// Expected: [3, 2, 1, 2, 3, 4, 3]
|
||
float expected_2[] = {3.0f, 2.0f, 1.0f, 2.0f, 3.0f, 4.0f, 3.0f};
|
||
struct ggml_tensor * out_2 = ggml_pad_reflect_1d(ctx, t, 2, 1);
|
||
|
||
// Test case 3: pad left=1, right=2
|
||
// Expected: [2, 1, 2, 3, 4, 3, 2]
|
||
float expected_3[] = {2.0f, 1.0f, 2.0f, 3.0f, 4.0f, 3.0f, 2.0f};
|
||
struct ggml_tensor * out_3 = ggml_pad_reflect_1d(ctx, t, 1, 2);
|
||
|
||
struct ggml_cgraph * gf = ggml_new_graph(ctx);
|
||
ggml_build_forward_expand(gf, out_1);
|
||
ggml_build_forward_expand(gf, out_2);
|
||
ggml_build_forward_expand(gf, out_3);
|
||
|
||
ggml_gallocr_alloc_graph(gallocr, gf);
|
||
|
||
ggml_backend_graph_compute(backend, gf);
|
||
|
||
check_tensor(out_1, expected_1, 6, 1, 1);
|
||
check_tensor(out_2, expected_2, 7, 1, 1);
|
||
check_tensor(out_3, expected_3, 7, 1, 1);
|
||
|
||
ggml_free(ctx);
|
||
ggml_backend_buffer_free(buffer);
|
||
ggml_gallocr_free(gallocr);
|
||
}
|
||
|
||
{
|
||
params = ggml_init_params{
|
||
/*.mem_size =*/ 16*1024*1024,
|
||
/*.mem_buffer =*/ nullptr,
|
||
/*.no_alloc. =*/ true
|
||
};
|
||
|
||
ggml_log_set(ggml_log_callback_default, nullptr);
|
||
|
||
ctx = ggml_init(params);
|
||
buffer = ggml_backend_alloc_buffer(backend, 16*1024*1024);
|
||
tallocr = ggml_tallocr_new(buffer);
|
||
gallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend));
|
||
|
||
// Create a 2D input tensor (5 columns × 4 rows)
|
||
struct ggml_tensor * t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 5, 4);
|
||
float input_data[] = {
|
||
1.0f, 2.0f, 3.0f, 4.0f, 5.0f, // row 1
|
||
6.0f, 7.0f, 8.0f, 9.0f, 10.0f, // row 2
|
||
11.0f, 12.0f, 13.0f, 14.0f, 15.0f, // row 3
|
||
16.0f, 17.0f, 18.0f, 19.0f, 20.0f // row 4
|
||
};
|
||
ggml_tallocr_alloc(&tallocr, t);
|
||
|
||
// load data to buffer
|
||
if(ggml_backend_is_cpu(backend)) {
|
||
memcpy(t->data, input_data, ggml_nbytes(t));
|
||
} else {
|
||
ggml_backend_tensor_set(t, input_data, 0, ggml_nbytes(t));
|
||
}
|
||
|
||
// Test case 4: pad left=3, right=2 on a 2D tensor
|
||
// Each row should be padded independently
|
||
float expected_4[] = {
|
||
4.0f, 3.0f, 2.0f, 1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 4.0f, 3.0f, // row 1
|
||
9.0f, 8.0f, 7.0f, 6.0f, 7.0f, 8.0f, 9.0f, 10.0f, 9.0f, 8.0f, // row 2
|
||
14.0f, 13.0f, 12.0f, 11.0f, 12.0f, 13.0f, 14.0f, 15.0f, 14.0f, 13.0f, // row 3
|
||
19.0f, 18.0f, 17.0f, 16.0f, 17.0f, 18.0f, 19.0f, 20.0f, 19.0f, 18.0f // row 4
|
||
};
|
||
struct ggml_tensor * out_4 = ggml_pad_reflect_1d(ctx, t, 3, 2);
|
||
|
||
struct ggml_cgraph * gf = ggml_new_graph(ctx);
|
||
ggml_build_forward_expand(gf, out_4);
|
||
|
||
ggml_gallocr_alloc_graph(gallocr, gf);
|
||
|
||
ggml_backend_graph_compute(backend, gf);
|
||
|
||
check_tensor(out_4, expected_4, 10, 4, 1);
|
||
|
||
ggml_free(ctx);
|
||
ggml_gallocr_free(gallocr);
|
||
ggml_backend_buffer_free(buffer);
|
||
}
|
||
|
||
ggml_backend_free(backend);
|
||
}
|
||
|
||
int main(int argc, const char * argv[]) {
|
||
bool use_gpu = false;
|
||
if (argc > 1) {
|
||
use_gpu = strcmp(argv[1], "--gpu") == 0;
|
||
}
|
||
test_pad_reflect_1d(use_gpu);
|
||
return 0;
|
||
}
|