diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c index 8b6e6028361d0..5536d3c03293d 100644 --- a/ggml/src/ggml-alloc.c +++ b/ggml/src/ggml-alloc.c @@ -23,7 +23,7 @@ static bool ggml_is_view(const struct ggml_tensor * t) { } // ops that return true for this function must not use restrict pointers for their backend implementations -static bool ggml_op_can_inplace(enum ggml_op op) { +bool ggml_op_can_inplace(enum ggml_op op) { switch (op) { case GGML_OP_SCALE: case GGML_OP_DIAG_MASK_ZERO: @@ -105,6 +105,7 @@ struct ggml_dyn_tallocr { int n_free_blocks; struct free_block free_blocks[MAX_FREE_BLOCKS]; size_t max_size; + size_t max_chunk_size; #ifdef GGML_ALLOCATOR_DEBUG struct { @@ -114,6 +115,14 @@ struct ggml_dyn_tallocr { #endif }; +// the memory range [0, max_size) is divided into n chunks of size max_chunk_size (with the last chunk possibly being smaller). +// tensor allocations may not cross chunk boundaries. +static void ggml_dyn_tallocr_new_chunk(struct ggml_dyn_tallocr * alloc, struct free_block * block) { + size_t n_chunks = (alloc->max_size + alloc->max_chunk_size - 1) / alloc->max_chunk_size; + block->offset = n_chunks * alloc->max_chunk_size; + block->size = alloc->max_chunk_size; +} + #ifdef GGML_ALLOCATOR_DEBUG static void add_allocated_tensor(struct ggml_dyn_tallocr * alloc, size_t offset, const struct ggml_tensor * tensor) { for (int i = 0; i < 1024; i++) { @@ -140,6 +149,10 @@ static size_t ggml_dyn_tallocr_alloc(struct ggml_dyn_tallocr * alloc, size_t siz size = aligned_offset(NULL, size, alloc->alignment); AT_PRINTF("%s: allocating %s (%zu bytes) - ", __func__, tensor->name, size); + if (size > alloc->max_chunk_size) { + GGML_ABORT("allocation failed: tensor %s (%zu bytes) exceeds maximum backend buffer size (%zu bytes)\n", + tensor->name, size, alloc->max_chunk_size); + } size_t max_avail = 0; @@ -156,16 +169,17 @@ static size_t ggml_dyn_tallocr_alloc(struct ggml_dyn_tallocr * alloc, size_t siz } if (best_fit_block == -1) { - // the last block is our last resort + // the last block represents memory still available in an existing chunk struct free_block * block = &alloc->free_blocks[alloc->n_free_blocks - 1]; max_avail = MAX(max_avail, block->size); if (block->size >= size) { best_fit_block = alloc->n_free_blocks - 1; } else { - // this should never happen - GGML_LOG_ERROR("%s: not enough space in the buffer to allocate %zu bytes, largest block available %zu bytes\n", - __func__, size, max_avail); - GGML_ABORT("not enough space in the buffer"); + // not enough space in existing chunk, create a new one at the end + best_fit_block = alloc->n_free_blocks; + alloc->n_free_blocks += 1; + GGML_ASSERT(alloc->n_free_blocks < MAX_FREE_BLOCKS && "out of free blocks"); + ggml_dyn_tallocr_new_chunk(alloc, &alloc->free_blocks[alloc->n_free_blocks - 1]); } } @@ -179,9 +193,14 @@ static size_t ggml_dyn_tallocr_alloc(struct ggml_dyn_tallocr * alloc, size_t siz for (int j = best_fit_block; j < alloc->n_free_blocks; j++) { alloc->free_blocks[j] = alloc->free_blocks[j+1]; } + // if there are no remaining blocks all memory in current chunk was used up -> start the next one + if (alloc->n_free_blocks == 0) { + alloc->n_free_blocks = 1; + ggml_dyn_tallocr_new_chunk(alloc, &alloc->free_blocks[0]); + } } - AT_PRINTF("block %d, offset %zu\n", best_fit_block, offset); + AT_PRINTF("block %d, offset %zu, chunk %d\n", best_fit_block, offset, offset / alloc->max_chunk_size); #ifdef GGML_ALLOCATOR_DEBUG add_allocated_tensor(alloc, offset, tensor); @@ -229,19 +248,28 @@ static void ggml_dyn_tallocr_free_tensor(struct ggml_dyn_tallocr * alloc, size_t #ifdef GGML_ALLOCATOR_DEBUG remove_allocated_tensor(alloc, offset, tensor); #endif + size_t chunk = offset / alloc->max_chunk_size; // see if we can merge with an existing block for (int i = 0; i < alloc->n_free_blocks; i++) { struct free_block * block = &alloc->free_blocks[i]; + // can only merge with blocks within the same chunk + size_t block_chunk = block->offset / alloc->max_chunk_size; + if (chunk != block_chunk) { + continue; + } // check if ptr is at the end of the block if (block->offset + block->size == offset) { block->size += size; - // check if we can merge with the next block - if (i < alloc->n_free_blocks - 1 && block->offset + block->size == alloc->free_blocks[i+1].offset) { - block->size += alloc->free_blocks[i+1].size; - alloc->n_free_blocks--; - for (int j = i+1; j < alloc->n_free_blocks; j++) { - alloc->free_blocks[j] = alloc->free_blocks[j+1]; + // check if we can merge with the next block (within the same chunk) + if (i < alloc->n_free_blocks - 1) { + struct free_block * next = &alloc->free_blocks[i+1]; + if (block->offset + block->size == next->offset && block_chunk == (next->offset / alloc->max_chunk_size)) { + block->size += next->size; + alloc->n_free_blocks--; + for (int j = i+1; j < alloc->n_free_blocks; j++) { + alloc->free_blocks[j] = alloc->free_blocks[j+1]; + } } } return; @@ -250,12 +278,15 @@ static void ggml_dyn_tallocr_free_tensor(struct ggml_dyn_tallocr * alloc, size_t if (offset + size == block->offset) { block->offset = offset; block->size += size; - // check if we can merge with the previous block - if (i > 0 && alloc->free_blocks[i-1].offset + alloc->free_blocks[i-1].size == block->offset) { - alloc->free_blocks[i-1].size += block->size; - alloc->n_free_blocks--; - for (int j = i; j < alloc->n_free_blocks; j++) { - alloc->free_blocks[j] = alloc->free_blocks[j+1]; + // check if we can merge with the previous block (within the same chunk) + if (i > 0) { + struct free_block * prev = &alloc->free_blocks[i-1]; + if (prev->offset + prev->size == block->offset && block_chunk == (prev->offset / alloc->max_chunk_size)) { + prev->size += block->size; + alloc->n_free_blocks--; + for (int j = i; j < alloc->n_free_blocks; j++) { + alloc->free_blocks[j] = alloc->free_blocks[j+1]; + } } } return; @@ -283,9 +314,13 @@ static void ggml_dyn_tallocr_free_tensor(struct ggml_dyn_tallocr * alloc, size_t static void ggml_dyn_tallocr_reset(struct ggml_dyn_tallocr * alloc) { alloc->n_free_blocks = 1; alloc->free_blocks[0].offset = 0; - alloc->free_blocks[0].size = SIZE_MAX/2; // restrict maximum size of a measure allocator to half size_t max to avoid overflows + alloc->free_blocks[0].size = alloc->max_chunk_size; alloc->max_size = 0; + if (alloc->free_blocks[0].size == SIZE_MAX) { + alloc->free_blocks[0].size = SIZE_MAX/2; // avoid overflows + } + #ifdef GGML_ALLOCATOR_DEBUG for (int i = 0; i < 1024; i++) { alloc->allocated_tensors[i].tensor = NULL; @@ -293,14 +328,15 @@ static void ggml_dyn_tallocr_reset(struct ggml_dyn_tallocr * alloc) { #endif } -static struct ggml_dyn_tallocr * ggml_dyn_tallocr_new(size_t alignment) { +static struct ggml_dyn_tallocr * ggml_dyn_tallocr_new(size_t alignment, size_t max_buffer_size) { struct ggml_dyn_tallocr * alloc = (struct ggml_dyn_tallocr *)malloc(sizeof(struct ggml_dyn_tallocr)); *alloc = (struct ggml_dyn_tallocr) { - /*.alignment = */ alignment, - /*.n_free_blocks = */ 0, - /*.free_blocks = */ {{0}}, - /*.max_size = */ 0, + /*.alignment = */ alignment, + /*.n_free_blocks = */ 0, + /*.free_blocks = */ {{0}}, + /*.max_size = */ 0, + /*.max_chunk_size = */ max_buffer_size, #ifdef GGML_ALLOCATOR_DEBUG /*.allocated_tensors = */ {{0}}, #endif @@ -320,6 +356,95 @@ static size_t ggml_dyn_tallocr_max_size(struct ggml_dyn_tallocr * alloc) { } +// virtual buffer with contiguous memory range, split into multiple backend buffers (chunks) + +#define GGML_VBUFFER_MAX_CHUNKS 8 + +struct vbuffer { + ggml_backend_buffer_type_t buft; + ggml_backend_buffer_t chunks[GGML_VBUFFER_MAX_CHUNKS]; +}; + +static struct vbuffer * ggml_vbuffer_new(ggml_backend_buffer_type_t buft) { + struct vbuffer * buf = calloc(1, sizeof(struct vbuffer)); + buf->buft = buft; + memset(buf->chunks, 0, sizeof(buf->chunks)); + return buf; +} + +static void ggml_vbuffer_free_chunks(struct vbuffer * buf) { + for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS; ++i) { + ggml_backend_buffer_free(buf->chunks[i]); + buf->chunks[i] = NULL; + } +} + +static void ggml_vbuffer_free(struct vbuffer * buf) { + if (buf == NULL) { + return; + } + ggml_vbuffer_free_chunks(buf); + free(buf); +} + +static int ggml_vbuffer_n_chunks(struct vbuffer * buf) { + int n = 0; + while (n < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[n]) n++; + return n; +} + +static size_t ggml_vbuffer_size(struct vbuffer * buf) { + size_t size = 0; + for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[i]; ++i) { + size += ggml_backend_buffer_get_size(buf->chunks[i]); + } + return size; +} + +static int ggml_vbuffer_alloc(struct vbuffer * buf, size_t size, enum ggml_backend_buffer_usage usage) { + size_t max_chunk_size = ggml_backend_buft_get_max_size(buf->buft); + if (size > GGML_VBUFFER_MAX_CHUNKS * max_chunk_size) { + return 0; + } + + int n = 0; + // always allocate at least 1 chunk even if requested size is 0 + while (size > 0 || n == 0) { + GGML_ASSERT(n < GGML_VBUFFER_MAX_CHUNKS); + size_t chunk_size = MIN(size, max_chunk_size); + buf->chunks[n] = ggml_backend_buft_alloc_buffer(buf->buft, chunk_size); + if (buf->chunks[n] == NULL) { + ggml_vbuffer_free_chunks(buf); + return 0; + } + ggml_backend_buffer_set_usage(buf->chunks[n], usage); + + GGML_ASSERT(size >= chunk_size); + size -= chunk_size; + n += 1; + } + return n; +} + +static void ggml_vbuffer_tensor_alloc(struct vbuffer * buf, struct ggml_tensor * tensor, size_t offset) { + size_t max_chunk_size = ggml_backend_buft_get_max_size(buf->buft); + size_t chunk_index = offset / max_chunk_size; + size_t chunk_offset = offset % max_chunk_size; + GGML_ASSERT(chunk_index < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[chunk_index] != NULL); + + void * base = ggml_backend_buffer_get_base(buf->chunks[chunk_index]); + void * addr = (char *)base + chunk_offset; + ggml_backend_tensor_alloc(buf->chunks[chunk_index], tensor, addr); +} + +static void ggml_vbuffer_reset(struct vbuffer * buf) { + for (int i = 0; i < GGML_VBUFFER_MAX_CHUNKS && buf->chunks[i]; ++i) { + ggml_backend_buffer_reset(buf->chunks[i]); + } +} + + + ///////////////////////////////////// // graph allocator @@ -349,7 +474,7 @@ struct node_alloc { struct ggml_gallocr { ggml_backend_buffer_type_t * bufts; // [n_buffers] - ggml_backend_buffer_t * buffers; // [n_buffers] + struct vbuffer ** buffers; // [n_buffers] struct ggml_dyn_tallocr ** buf_tallocs; // [n_buffers] int n_buffers; @@ -370,7 +495,7 @@ ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs galloc->bufts = calloc(n_bufs, sizeof(ggml_backend_buffer_type_t)); GGML_ASSERT(galloc->bufts != NULL); - galloc->buffers = calloc(n_bufs, sizeof(ggml_backend_buffer_t)); + galloc->buffers = calloc(n_bufs, sizeof(struct vbuffer *)); GGML_ASSERT(galloc->buffers != NULL); galloc->buf_tallocs = calloc(n_bufs, sizeof(struct ggml_dyn_tallocr *)); @@ -378,7 +503,7 @@ ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs for (int i = 0; i < n_bufs; i++) { galloc->bufts[i] = bufts[i]; - galloc->buffers[i] = NULL; + galloc->buffers[i] = ggml_vbuffer_new(bufts[i]); // check if the same buffer type is used multiple times and reuse the same allocator for (int j = 0; j < i; j++) { @@ -390,7 +515,8 @@ ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs if (galloc->buf_tallocs[i] == NULL) { size_t alignment = ggml_backend_buft_get_alignment(bufts[i]); - galloc->buf_tallocs[i] = ggml_dyn_tallocr_new(alignment); + size_t max_size = ggml_backend_buft_get_max_size(bufts[i]); + galloc->buf_tallocs[i] = ggml_dyn_tallocr_new(alignment, max_size); } } galloc->n_buffers = n_bufs; @@ -418,7 +544,7 @@ void ggml_gallocr_free(ggml_gallocr_t galloc) { } } if (!freed) { - ggml_backend_buffer_free(galloc->buffers[i]); + ggml_vbuffer_free(galloc->buffers[i]); } } if (galloc->buf_tallocs != NULL) { @@ -744,22 +870,20 @@ bool ggml_gallocr_reserve_n(ggml_gallocr_t galloc, struct ggml_cgraph * graph, c } } - size_t cur_size = galloc->buffers[i] ? ggml_backend_buffer_get_size(galloc->buffers[i]) : 0; + size_t cur_size = ggml_vbuffer_size(galloc->buffers[i]); size_t new_size = ggml_dyn_tallocr_max_size(galloc->buf_tallocs[i]); // even if there are no tensors allocated in this buffer, we still need to allocate it to initialize views - if (new_size > cur_size || galloc->buffers[i] == NULL) { + if (new_size > cur_size || ggml_vbuffer_n_chunks(galloc->buffers[i]) == 0) { #ifndef NDEBUG GGML_LOG_DEBUG("%s: reallocating %s buffer from size %.02f MiB to %.02f MiB\n", __func__, ggml_backend_buft_name(galloc->bufts[i]), cur_size / 1024.0 / 1024.0, new_size / 1024.0 / 1024.0); #endif - ggml_backend_buffer_free(galloc->buffers[i]); - galloc->buffers[i] = ggml_backend_buft_alloc_buffer(galloc->bufts[i], new_size); - if (galloc->buffers[i] == NULL) { + ggml_vbuffer_free_chunks(galloc->buffers[i]); + if (!ggml_vbuffer_alloc(galloc->buffers[i], new_size, GGML_BACKEND_BUFFER_USAGE_COMPUTE)) { GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(galloc->bufts[i]), new_size); return false; } - ggml_backend_buffer_set_usage(galloc->buffers[i], GGML_BACKEND_BUFFER_USAGE_COMPUTE); } } @@ -772,7 +896,7 @@ bool ggml_gallocr_reserve(ggml_gallocr_t galloc, struct ggml_cgraph *graph) { static void ggml_gallocr_init_tensor(ggml_gallocr_t galloc, struct ggml_tensor * tensor, struct tensor_alloc * tensor_alloc) { int buffer_id = tensor_alloc->buffer_id; - assert(tensor->data || tensor->view_src || ggml_backend_buffer_get_alloc_size(galloc->buffers[buffer_id], tensor) <= tensor_alloc->size_max); + assert(tensor->data || tensor->view_src || ggml_backend_buft_get_alloc_size(galloc->bufts[buffer_id], tensor) <= tensor_alloc->size_max); if (tensor->view_src != NULL) { if (tensor->buffer == NULL) { @@ -786,10 +910,8 @@ static void ggml_gallocr_init_tensor(ggml_gallocr_t galloc, struct ggml_tensor * } else { if (tensor->data == NULL) { assert(tensor_alloc->offset != SIZE_MAX); - assert(ggml_backend_buffer_get_alloc_size(galloc->buffers[buffer_id], tensor) <= tensor_alloc->size_max); - void * base = ggml_backend_buffer_get_base(galloc->buffers[buffer_id]); - void * addr = (char *)base + tensor_alloc->offset; - ggml_backend_tensor_alloc(galloc->buffers[buffer_id], tensor, addr); + assert(ggml_backend_buft_get_alloc_size(galloc->bufts[buffer_id], tensor) <= tensor_alloc->size_max); + ggml_vbuffer_tensor_alloc(galloc->buffers[buffer_id], tensor, tensor_alloc->offset); } else { if (tensor->buffer == NULL) { // this tensor was allocated without ggml-backend @@ -874,7 +996,7 @@ bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph * graph) // reset buffers for (int i = 0; i < galloc->n_buffers; i++) { if (galloc->buffers[i] != NULL) { - ggml_backend_buffer_reset(galloc->buffers[i]); + ggml_vbuffer_reset(galloc->buffers[i]); } } @@ -917,7 +1039,7 @@ size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id) { } } - return ggml_backend_buffer_get_size(galloc->buffers[buffer_id]); + return ggml_vbuffer_size(galloc->buffers[buffer_id]); } // utils diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index 19a7adb2d101b..0fc42846f0a77 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl.h @@ -329,6 +329,10 @@ struct ggml_cgraph { // if you need the gradients, get them from the original graph struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph, int i0, int i1); +// ggml-alloc.c: true if the operation can reuse memory from its sources +GGML_API bool ggml_op_can_inplace(enum ggml_op op); + + // Memory allocation GGML_API void * ggml_aligned_malloc(size_t size); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index cd1c66ba7b476..67b183754542a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -11129,7 +11129,7 @@ static size_t ggml_backend_vk_buffer_type_get_alignment(ggml_backend_buffer_type static size_t ggml_backend_vk_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { ggml_backend_vk_buffer_type_context * ctx = (ggml_backend_vk_buffer_type_context *) buft->context; - return ctx->device->suballocation_block_size; + return ctx->device->max_memory_allocation_size; } static size_t ggml_backend_vk_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 91719577564a9..3e9e082d93c18 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -219,3 +219,6 @@ target_link_libraries(${LLAMA_TEST_NAME} PRIVATE mtmd) get_filename_component(TEST_TARGET test-c.c NAME_WE) add_executable(${TEST_TARGET} test-c.c) target_link_libraries(${TEST_TARGET} PRIVATE llama) + +llama_build_and_test(test-alloc.cpp) +target_include_directories(test-alloc PRIVATE ${PROJECT_SOURCE_DIR}/ggml/src) diff --git a/tests/test-alloc.cpp b/tests/test-alloc.cpp new file mode 100644 index 0000000000000..031a44e811037 --- /dev/null +++ b/tests/test-alloc.cpp @@ -0,0 +1,483 @@ +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +// +// dummy backend with configurable max_buffer_size, tracks allocations + +uint8_t * const alloc_base = (uint8_t *) 16; + +struct dummy_backend_context { + size_t max_buffer_size = 64; + + ggml_backend_buffer_i buffer_interface; + std::vector buffers; + + size_t allocated_total() const { + size_t n = 0; + for (ggml_backend_buffer_t buf : buffers) { + n += ggml_backend_buffer_get_size(buf); + } + return n; + } +}; + +// ggml_backend_buffer_type interface + +static const char * dummy_backend_buffer_type_get_name(ggml_backend_buffer_type_t) { + return "dummy_buffer_type"; +} + +static ggml_backend_buffer_t dummy_backend_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + dummy_backend_context * ctx = (dummy_backend_context *) buft->context; + ggml_backend_buffer_t & buffer = ctx->buffers.emplace_back(); + buffer = ggml_backend_buffer_init(buft, ctx->buffer_interface, ctx, size); + return buffer; +} + +static size_t dummy_backend_buffer_type_get_alignment(ggml_backend_buffer_type_t) { + return 8; +} + +static size_t dummy_backend_buffer_type_get_max_size(ggml_backend_buffer_type_t buft) { + dummy_backend_context * ctx = (dummy_backend_context *) buft->context; + return ctx->max_buffer_size; +} + +static bool dummy_backend_buffer_type_is_host(ggml_backend_buffer_type_t) { + return true; +} + +// ggml_backend_buffer interface + +static void dummy_backend_buffer_free_buffer(ggml_backend_buffer_t buffer) { + dummy_backend_context * ctx = (dummy_backend_context *) buffer->context; + + auto i = std::find(ctx->buffers.begin(), ctx->buffers.end(), buffer); + GGML_ASSERT(i != ctx->buffers.end()); + ctx->buffers.erase(i); +} + +static void * dummy_backend_buffer_get_base(ggml_backend_buffer_t) { + return alloc_base; +} + +static ggml_status dummy_backend_buffer_init_tensor(ggml_backend_buffer_t, ggml_tensor *) { + return GGML_STATUS_SUCCESS; +} + +static void dummy_backend_buffer_memset_tensor(ggml_backend_buffer_t, ggml_tensor *, uint8_t, size_t, size_t) {} + +static void dummy_backend_buffer_set_tensor(ggml_backend_buffer_t, ggml_tensor *, const void *, size_t, size_t) {} + +static void dummy_backend_buffer_get_tensor(ggml_backend_buffer_t, const ggml_tensor *, void *, size_t, size_t) {} + +static void dummy_backend_buffer_clear(ggml_backend_buffer_t, uint8_t) {} + +// dummy_backend (not really a full backend, just provides what gallocr needs) + +struct dummy_backend { + std::unique_ptr context; + ggml_backend_buffer_type buffer_type; +}; + +static dummy_backend dummy_backend_init(size_t max_buffer_size) { + dummy_backend b{}; + b.context = std::make_unique(); + b.context->max_buffer_size = max_buffer_size; + + b.context->buffer_interface.free_buffer = dummy_backend_buffer_free_buffer; + b.context->buffer_interface.get_base = dummy_backend_buffer_get_base; + b.context->buffer_interface.init_tensor = dummy_backend_buffer_init_tensor; + b.context->buffer_interface.memset_tensor = dummy_backend_buffer_memset_tensor; + b.context->buffer_interface.set_tensor = dummy_backend_buffer_set_tensor; + b.context->buffer_interface.get_tensor = dummy_backend_buffer_get_tensor; + b.context->buffer_interface.clear = dummy_backend_buffer_clear; + + b.buffer_type.context = b.context.get(); + b.buffer_type.iface.get_name = dummy_backend_buffer_type_get_name; + b.buffer_type.iface.alloc_buffer = dummy_backend_buffer_type_alloc_buffer; + b.buffer_type.iface.get_alignment = dummy_backend_buffer_type_get_alignment; + b.buffer_type.iface.get_max_size = dummy_backend_buffer_type_get_max_size; + b.buffer_type.iface.is_host = dummy_backend_buffer_type_is_host; + return b; +} + +// +// test utilities + +struct test_context_with_graph { + ggml_context * ctx; + ggml_cgraph * graph; + ggml_context_ptr ctx_ptr; +}; + +static test_context_with_graph make_context() { + ggml_init_params params{}; + params.mem_size = 32 * ggml_tensor_overhead() + ggml_graph_overhead(); + params.no_alloc = true; + + ggml_context * ctx = ggml_init(params); + ggml_context_ptr ctx_ptr = ggml_context_ptr(ctx); + ggml_cgraph * graph = ggml_new_graph(ctx); + return { ctx, graph, std::move(ctx_ptr) }; +} + +static ggml_tensor * make_input_1d(ggml_context * ctx, int64_t n_elements) { + ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_elements); + ggml_set_input(t); + return t; +} + +static ggml_tensor * make_input_with_size(ggml_context * ctx, size_t size_bytes) { + GGML_ASSERT(size_bytes % 4 == 0); + return make_input_1d(ctx, size_bytes / 4); +} + +static void assign_names(ggml_context * ctx, const char * prefix = "x") { + int i = 0; + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; t = ggml_get_next_tensor(ctx, t)) { + ggml_format_name(t, "%s%d", prefix, i++); + } +} + +static int get_leaf_id(ggml_cgraph * graph, const char * tensor_name) { + for (int i = 0; i < graph->n_leafs; ++i) { + if (strncmp(graph->leafs[i]->name, tensor_name, GGML_MAX_NAME) == 0) { + return i; + } + } + fprintf(stderr, "leaf not found: %s\n", tensor_name); + return -1; +} + +static int get_node_id(ggml_cgraph * graph, const char * tensor_name) { + for (int i = 0; i < graph->n_nodes; ++i) { + if (strncmp(graph->nodes[i]->name, tensor_name, GGML_MAX_NAME) == 0) { + return i; + } + } + fprintf(stderr, "node not found: %s", tensor_name); + return -1; +} + +static ggml_gallocr_ptr allocate_graph(ggml_cgraph * graph, ggml_tensor * out, ggml_backend_buffer_type_t buft) { + ggml_set_output(out); + ggml_build_forward_expand(graph, out); + + ggml_gallocr_ptr galloc = ggml_gallocr_ptr(ggml_gallocr_new(buft)); + bool result = ggml_gallocr_alloc_graph(galloc.get(), graph); + GGML_ASSERT(result); + return galloc; +} + +// +// correctness checks for result allocations + +static void check_all_allocated(ggml_cgraph * graph) { + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + ggml_tensor * t = ggml_graph_node(graph, i); + GGML_ASSERT(t->buffer != nullptr); + GGML_ASSERT(t->data != nullptr); + } +} + +static void check_max_size(ggml_context * ctx) { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t; t = ggml_get_next_tensor(ctx, t)) { + auto buft = ggml_backend_buffer_get_type(t->buffer); + size_t max_size = ggml_backend_buft_get_max_size(buft); + size_t offset = (char *) t->data - (char *) ggml_backend_buffer_get_base(t->buffer); + GGML_ASSERT(t->data >= ggml_backend_buffer_get_base(t->buffer)); + GGML_ASSERT((size_t) offset + ggml_nbytes(t) <= max_size); + } +} + +static bool can_reuse_memory(ggml_cgraph * graph, int current_i, ggml_tensor * current, ggml_tensor * other) { + if (other->flags & GGML_TENSOR_FLAG_OUTPUT) { + return false; + } + // Check if `other` is still "alive", ie. an input to any node after the `current` op + for (int i = current_i; i < ggml_graph_n_nodes(graph); ++i) { + ggml_tensor * t = ggml_graph_node(graph, i); + for (int s = 0; s < GGML_MAX_SRC; s++) { + if (t == current && ggml_op_can_inplace(t->op)) { + continue; + } + if (t->src[s] == other) { + return false; + } + if (t->src[s] && t->src[s]->view_src == other) { + return false; + } + } + } + return true; +} + +static bool memory_overlap(ggml_tensor * a, ggml_tensor * b) { + if (a->buffer != b->buffer) { + return false; + } + int64_t a0 = (int64_t) a->data; + int64_t a1 = a0 + ggml_nbytes(a); + int64_t b0 = (int64_t) b->data; + int64_t b1 = b0 + ggml_nbytes(b); + return a1 > b0 && b1 > a0; +} + +static ggml_tensor * get_view_source(ggml_tensor * t) { + while (t->view_src) { + t = t->view_src; + } + return t; +} + +static void check_no_overlap(ggml_cgraph * graph) { + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + for (int j = 0; j < i; ++j) { + ggml_tensor * t = ggml_graph_node(graph, i); + ggml_tensor * o = ggml_graph_node(graph, j); + GGML_ASSERT(t != o); + + if (get_view_source(t) == get_view_source(o)) { + continue; + } + if (memory_overlap(t, o)) { + GGML_ASSERT(can_reuse_memory(graph, i, t, o)); + } + } + } +} + +// +// test cases + +// scenario where the first backend buffer is completely exhausted and there are further +// tensors which require a second buffer +static void test_max_size_too_many_tensors() { + dummy_backend backend = dummy_backend_init(16); + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * x[7]; + x[0] = make_input_with_size(ctx, 8); + x[1] = make_input_with_size(ctx, 8); + x[2] = make_input_with_size(ctx, 8); + x[3] = ggml_mul(ctx, x[0], x[1]); + x[4] = ggml_add(ctx, x[1], x[2]); + x[5] = ggml_add(ctx, x[3], x[0]); + x[6] = ggml_add(ctx, x[4], x[5]); + assign_names(ctx); + + ggml_gallocr_ptr galloc = allocate_graph(graph, x[6], &backend.buffer_type); + check_all_allocated(graph); + check_no_overlap(graph); + check_max_size(ctx); + GGML_ASSERT(backend.context->allocated_total() <= 16 + 16); +} + +// scenario where there is some space left in the first buffer, but not enough to accomodate +// a larger tensor, so a second buffer is required +static void test_max_size_tensor_too_large() { + dummy_backend backend = dummy_backend_init(32); + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * x[3]; + x[0] = make_input_with_size(ctx, 16); // chunk 0, [0 , 16) + x[1] = make_input_with_size(ctx, 8); // chunk 0, [16, 24) + x[2] = ggml_concat(ctx, x[0], x[1], 0); // chunk 1, [0 , 24) + assign_names(ctx); + + ggml_gallocr_ptr galloc = allocate_graph(graph, x[2], &backend.buffer_type); + check_all_allocated(graph); + check_no_overlap(graph); + check_max_size(ctx); + GGML_ASSERT(backend.context->allocated_total() <= 32 + 24); +} + +// check that views don't require any extra memory +static void test_view_inplace() { + dummy_backend backend = dummy_backend_init(32); + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * x[6]; + x[0] = make_input_1d(ctx, 4); // chunk 0, [0, 16) + x[1] = ggml_reshape_2d(ctx, x[0], 2, 2); // view of x0 + x[2] = ggml_permute(ctx, x[1], 1, 0, 2, 3); // view of x0 + x[3] = ggml_view_1d(ctx, x[2], 2, 4); // view of x0 + x[4] = make_input_1d(ctx, 2); // chunk 0, [16, 24) + x[5] = ggml_add(ctx, x[3], x[4]); // reuse (inplace add) + assign_names(ctx); + + ggml_gallocr_ptr galloc = allocate_graph(graph, x[5], &backend.buffer_type); + check_all_allocated(graph); + check_no_overlap(graph); + check_max_size(ctx); + GGML_ASSERT(backend.context->allocated_total() <= 24); +} + +static void test_reuse_and_free() { + dummy_backend backend = dummy_backend_init(32); + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * x[9]; + x[0] = make_input_with_size(ctx, 24); + x[1] = make_input_with_size(ctx, 8); + x[2] = make_input_with_size(ctx, 8); + x[3] = ggml_add(ctx, x[1], x[2]); // reuse, free x2 + x[4] = ggml_pad(ctx, x[0], 2, 0, 0, 0); // alloc new buffer, free x0 + x[5] = ggml_scale(ctx, x[4], 2.0f); // alloc from free block + x[6] = ggml_add(ctx, x[4], x[5]); // reuse, free x5 + x[7] = ggml_view_1d(ctx, x[6], 2, 8); // view + x[8] = ggml_add(ctx, x[3], x[7]); // reuse + assign_names(ctx); + + ggml_gallocr_ptr galloc = allocate_graph(graph, x[8], &backend.buffer_type); + check_all_allocated(graph); + check_no_overlap(graph); + check_max_size(ctx); + GGML_ASSERT(backend.context->allocated_total() <= 32 + 32 + 32); +} + +static void test_merge_free_block(size_t max_buffer_size) { + dummy_backend backend = dummy_backend_init(max_buffer_size); + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * x[9]; + x[0] = make_input_with_size(ctx, 16); + x[1] = make_input_with_size(ctx, 16); + x[2] = make_input_with_size(ctx, 16); + x[3] = ggml_mean(ctx, x[0]); + x[4] = ggml_mean(ctx, x[1]); + x[5] = ggml_pad(ctx, x[2], 2, 0, 0, 0); + x[6] = ggml_add(ctx, x[3], x[4]); + x[7] = ggml_pad(ctx, x[6], 5, 0, 0, 0); + x[8] = ggml_add(ctx, x[5], x[7]); + assign_names(ctx); + + ggml_gallocr_ptr galloc = allocate_graph(graph, x[8], &backend.buffer_type); + check_all_allocated(graph); + check_no_overlap(graph); + check_max_size(ctx); + GGML_ASSERT(backend.context->allocated_total() <= 32 + 32 + 24); +} + +// test for allocating on multiple devices with some tensors in the graph +// allocated externally (not by gallocr). +static void test_multiple_buffer_types() { + dummy_backend backend_a = dummy_backend_init(32); + dummy_backend backend_b = dummy_backend_init(SIZE_MAX); + + auto [ctx_a, _a, ctx_a_ptr] = make_context(); + auto [ctx_b, _b, ctx_b_ptr] = make_context(); + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * a[2]; + a[0] = make_input_with_size(ctx_a, 16); + a[1] = make_input_with_size(ctx_a, 16); + assign_names(ctx_a, "a"); + + ggml_tensor * b[2]; + b[0] = make_input_with_size(ctx_b, 24); + b[1] = make_input_with_size(ctx_b, 4); + assign_names(ctx_b, "b"); + + ggml_tensor * x[9]; + x[0] = make_input_with_size(ctx, 16); + x[1] = ggml_mul(ctx, x[0], a[0]); + x[2] = ggml_pad(ctx, x[1], 2, 0, 0, 0); + x[3] = ggml_mul(ctx, x[2], b[0]); + x[4] = ggml_mean(ctx, x[3]); + x[5] = ggml_add(ctx, x[4], b[1]); + x[6] = ggml_pad(ctx, x[5], 3, 0, 0, 0); + x[7] = ggml_add(ctx, x[6], a[1]); + x[8] = ggml_scale(ctx, x[7], 2.0f); + assign_names(ctx, "x"); + + ggml_backend_buffer_ptr buf_a(ggml_backend_alloc_ctx_tensors_from_buft(ctx_a, &backend_a.buffer_type)); + ggml_backend_buffer_ptr buf_b(ggml_backend_alloc_ctx_tensors_from_buft(ctx_b, &backend_b.buffer_type)); + ggml_backend_buffer_type_t bufts[2] = { &backend_a.buffer_type, &backend_b.buffer_type }; + + // assign buffer types manually to avoid extra complexity from backend scheduler + ggml_set_output(x[8]); + ggml_build_forward_expand(graph, x[8]); + + GGML_ASSERT(graph->n_leafs == 5); + int leaf_buffer_ids[5]; + leaf_buffer_ids[get_leaf_id(graph, "a0")] = 0; + leaf_buffer_ids[get_leaf_id(graph, "a1")] = 0; + leaf_buffer_ids[get_leaf_id(graph, "b0")] = 1; + leaf_buffer_ids[get_leaf_id(graph, "b1")] = 1; + leaf_buffer_ids[get_leaf_id(graph, "x0")] = 0; + + GGML_ASSERT(graph->n_nodes == 8); + int node_buffer_ids[8]; + node_buffer_ids[get_node_id(graph, "x1")] = 0; + node_buffer_ids[get_node_id(graph, "x2")] = 0; + node_buffer_ids[get_node_id(graph, "x3")] = 1; + node_buffer_ids[get_node_id(graph, "x4")] = 1; + node_buffer_ids[get_node_id(graph, "x5")] = 1; + node_buffer_ids[get_node_id(graph, "x6")] = 1; + node_buffer_ids[get_node_id(graph, "x7")] = 0; + node_buffer_ids[get_node_id(graph, "x8")] = 0; + + ggml_gallocr_ptr galloc(ggml_gallocr_new_n(bufts, 2)); + ggml_gallocr_reserve_n(galloc.get(), graph, node_buffer_ids, leaf_buffer_ids); + ggml_gallocr_alloc_graph(galloc.get(), graph); + + check_all_allocated(graph); + check_no_overlap(graph); + check_max_size(ctx); + GGML_ASSERT(backend_a.context->allocated_total() <= 32 + 32 + 24); + GGML_ASSERT(backend_b.context->allocated_total() <= 32 + 24); +} + +static void test_buffer_size_zero() { + dummy_backend backend_a = dummy_backend_init(SIZE_MAX); + dummy_backend backend_b = dummy_backend_init(SIZE_MAX); + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * x[2]; + x[0] = make_input_with_size(ctx, 16); + x[1] = ggml_scale(ctx, x[0], 2.0f); + + ggml_set_output(x[1]); + ggml_build_forward_expand(graph, x[1]); + + int leaf_buffer_ids[1] = { 0 }; + int node_buffer_ids[1] = { 0 }; + + ggml_backend_buffer_type_t bufts[2] = { &backend_a.buffer_type, &backend_b.buffer_type }; + ggml_gallocr_ptr galloc = ggml_gallocr_ptr(ggml_gallocr_new_n(bufts, 2)); + ggml_gallocr_reserve_n(galloc.get(), graph, node_buffer_ids, leaf_buffer_ids); + ggml_gallocr_alloc_graph(galloc.get(), graph); + + check_all_allocated(graph); + GGML_ASSERT(backend_a.context->allocated_total() == 16); + GGML_ASSERT(backend_b.context->allocated_total() == 0); +} + +static void run(const char * name, void (*f)()) { + printf("%s ", name); + fflush(stdout); + f(); + printf("PASSED\n"); +} + +int main() { + run("test_max_size_too_many_tensors", test_max_size_too_many_tensors); + run("test_max_size_tensor_too_large", test_max_size_tensor_too_large); + run("test_view_inplace", test_view_inplace); + run("test_reuse_and_free", test_reuse_and_free); + run("test_merge_free_block(32)", []() { test_merge_free_block(32); }); + run("test_merge_free_block(SIZE_MAX)", []() { test_merge_free_block(SIZE_MAX); }); + run("test_multiple_buffer_types", test_multiple_buffer_types); + run("test_buffer_size_zero", test_buffer_size_zero); + return 0; +}