backend_sched : fix assignments

ggml-ci

diff --git a/ggml-backend.c b/ggml-backend.c
index 4c2d8b0b26f18a9ce2e178ec1219b86559793dff..505dbba47625390c7a596d0c0688dab6363ba6cc 100644 (file)
--- a/ggml-backend.c
+++ b/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++;