Skip to content

Commit

Permalink
ggml-alloc : fix backend assignments of views (ggerganov#3982)
Browse files Browse the repository at this point in the history
  • Loading branch information
slaren authored Nov 8, 2023
1 parent 0a7c980 commit 875fb42
Showing 1 changed file with 12 additions and 9 deletions.
21 changes: 12 additions & 9 deletions ggml-alloc.c
Original file line number Diff line number Diff line change
Expand Up @@ -378,9 +378,13 @@ static bool ggml_op_can_inplace(enum ggml_op op) {
}
}

static void init_view(struct ggml_allocr * alloc, struct ggml_tensor * view) {
static void init_view(struct ggml_allocr * alloc, struct ggml_tensor * view, bool update_backend) {
assert(view->view_src != NULL && view->view_src->data != NULL);
view->backend = view->view_src->backend;

if (update_backend) {
view->backend = view->view_src->backend;
}

view->buffer = view->view_src->buffer;
view->data = (char *)view->view_src->data + view->view_offs;

Expand All @@ -394,7 +398,7 @@ static void allocate_node(struct ggml_allocr * alloc, struct ggml_tensor * node)
struct hash_node * ht = alloc->hash_table;
if (node->data == NULL) {
if (ggml_is_view(node)) {
init_view(alloc, node);
init_view(alloc, node, true);
} else {
// see if we can reuse a parent's buffer (inplace)
if (ggml_op_can_inplace(node->op)) {
Expand Down Expand Up @@ -424,15 +428,14 @@ static void allocate_node(struct ggml_allocr * alloc, struct ggml_tensor * node)
AT_PRINTF("reusing view parent %s (%s) for %s\n", parent->name, view_src->name, node->name);
node->view_src = view_src;
view_src_hn->n_views += 1;
init_view(alloc, node);
init_view(alloc, node, false);
return;
}
}
else {
} else {
AT_PRINTF("reusing parent %s for %s\n", parent->name, node->name);
node->view_src = parent;
p_hn->n_views += 1;
init_view(alloc, node);
init_view(alloc, node, false);
return;
}
}
Expand Down Expand Up @@ -463,7 +466,7 @@ size_t ggml_allocr_alloc_graph_n(
hash_get(ht, view_src)->n_views += 1;
if (node->buffer == NULL && node->data != NULL) {
// view of a pre-allocated tensor, didn't call init_view() yet
init_view(alloc, node);
init_view(alloc, node, true);
}
}

Expand All @@ -474,7 +477,7 @@ size_t ggml_allocr_alloc_graph_n(
}
hash_get(ht, parent)->n_children += 1;
if (ggml_is_view(parent) && parent->buffer == NULL && parent->data != NULL) {
init_view(alloc, parent);
init_view(alloc, parent, true);
}
}
}
Expand Down

0 comments on commit 875fb42

Please sign in to comment.