Felipe97/llama-cpp-compiled
01.1k
1#include "ggml-alloc.h"2#include "ggml-backend.h"3#include "ggml-impl.h"4#include "ggml-rpc.h"5#include "ggml.h"6 7int main(int argc, char ** argv) {8 GGML_ASSERT(argc == 3);9 ggml_backend_load_all();10 11 const char * endpoint_a = argv[1];12 const char * endpoint_b = argv[2];13 14 ggml_backend_t backend_a = ggml_backend_rpc_init(endpoint_a, 0);15 ggml_backend_t backend_b = ggml_backend_rpc_init(endpoint_b, 0);16 GGML_ASSERT(backend_a != nullptr);17 GGML_ASSERT(backend_b != nullptr);18 19 ggml_init_params params = {20 /* .mem_size = */ ggml_tensor_overhead() + ggml_graph_overhead_custom(1, false),21 /* .mem_buffer = */ nullptr,22 /* .no_alloc = */ true,23 };24 ggml_context * ctx = ggml_init(params);25 GGML_ASSERT(ctx != nullptr);26 27 ggml_tensor * tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);28 ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend_a);29 GGML_ASSERT(buffer != nullptr);30 31 // A remote pointer allocated by server A is not meaningful to server B.32 ggml_cgraph * graph = ggml_new_graph_custom(ctx, 1, false);33 graph->nodes[0] = tensor;34 graph->n_nodes = 1;35 36 GGML_ASSERT(ggml_backend_graph_compute(backend_b, graph) == GGML_STATUS_SUCCESS);37 // Wait for server B to finish the graph before the script checks its log.38 size_t free_mem;39 size_t total_mem;40 ggml_backend_rpc_get_device_memory(endpoint_b, 0, &free_mem, &total_mem);41 GGML_ASSERT(total_mem > 0);42 ggml_backend_buffer_free(buffer);43 ggml_free(ctx);44 ggml_backend_free(backend_b);45 ggml_backend_free(backend_a);46 return 0;47}48 