backend_sched : fix assignments

ggml-ci
ggerganov · Jan 12, 2024 · 8fb376b · 8fb376b
1 parent 45d3329
commit 8fb376b
Showing 1 changed file with 20 additions and 0 deletions.
diff --git a/src/ggml-backend.c b/src/ggml-backend.c
@@ -1087,6 +1087,24 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
  }
  }
  }
+
+ // pass 2.4 expand rest down
+ {
+ ggml_tallocr_t cur_allocr = NULL;
+ for (int i = 0; i < graph->n_nodes; i++) {
+ struct ggml_tensor * node = graph->nodes[i];
+ if (ggml_is_view_op(node->op)) {
+ continue;
+ }
+ ggml_tallocr_t node_allocr = node_allocr(node);
+ if (node_allocr != NULL) {
+ cur_allocr = node_allocr;
+ } else {
+ node_allocr(node) = cur_allocr;
+ SET_CAUSE(node, "2.4");
+ }
+ }
+ }
 #ifdef DEBUG_PASS2
  fprintf(stderr, "PASS 2 ASSIGNMENTS\n"); sched_print_assignments(sched, graph);
 #endif
@@ -1146,6 +1164,8 @@ static void sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgraph * g
 
  ggml_tallocr_t node_allocr = node_allocr(node);
 
+ GGML_ASSERT(node_allocr != NULL); // all nodes should be assigned by now
+
  if (node_allocr != cur_allocr) {
  sched->splits[cur_split].i_end = i;
  cur_split++;