Simplify tensor allocation logic.

This commit is contained in:
Will Findley 2024-01-08 16:22:05 -06:00
parent 1fc2f265ff
commit a522f7e90f

View file

@ -93,7 +93,7 @@ void ggml_tallocr_alloc(ggml_tallocr_t alloc, struct ggml_tensor * tensor) {
// find the best fitting free block besides the last block // find the best fitting free block besides the last block
int best_fit_block = -1; int best_fit_block = -1;
size_t best_fit_size = SIZE_MAX; size_t best_fit_size = SIZE_MAX;
for (int i = 0; i < alloc->n_free_blocks - 1; i++) { for (int i = 0; i < alloc->n_free_blocks; i++) {
struct free_block * block = &alloc->free_blocks[i]; struct free_block * block = &alloc->free_blocks[i];
max_avail = MAX(max_avail, block->size); max_avail = MAX(max_avail, block->size);
if (block->size >= size && block->size <= best_fit_size) { if (block->size >= size && block->size <= best_fit_size) {
@ -105,17 +105,10 @@ void ggml_tallocr_alloc(ggml_tallocr_t alloc, struct ggml_tensor * tensor) {
AT_PRINTF("block %d\n", best_fit_block); AT_PRINTF("block %d\n", best_fit_block);
if (best_fit_block == -1) { if (best_fit_block == -1) {
// the last block is our last resort fprintf(stderr, "%s: not enough space in the buffer (needed %zu, largest block available %zu)\n",
struct free_block * block = &alloc->free_blocks[alloc->n_free_blocks - 1]; __func__, size, max_avail);
max_avail = MAX(max_avail, block->size); GGML_ASSERT(!"not enough space in the buffer");
if (block->size >= size) { return;
best_fit_block = alloc->n_free_blocks - 1;
} else {
fprintf(stderr, "%s: not enough space in the buffer (needed %zu, largest block available %zu)\n",
__func__, size, max_avail);
GGML_ASSERT(!"not enough space in the buffer");
return;
}
} }
struct free_block * block = &alloc->free_blocks[best_fit_block]; struct free_block * block = &alloc->free_blocks[best_fit_block];
void * addr = block->addr; void * addr = block->addr;