From 54b1181bedeb18b3176f18717e1a78626ae20155 Mon Sep 17 00:00:00 2001
From: 0cc4m <picard12@live.de>
Date: Tue, 30 Jan 2024 19:02:48 +0100
Subject: [PATCH] Fix Vulkan context shift crash

---
 ggml-vulkan.cpp | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/ggml-vulkan.cpp b/ggml-vulkan.cpp
index 1d93ec6bb..7a09fc400 100644
--- a/ggml-vulkan.cpp
+++ b/ggml-vulkan.cpp
@@ -4556,8 +4556,15 @@ GGML_CALL static bool ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml
     }
     ggml_vk_preallocate_buffers();
 
+    int last_node = cgraph->n_nodes - 1;
+
+    // If the last op in the cgraph isn't backend GPU, the command buffer doesn't get closed properly
+    while (last_node > 0 && cgraph->nodes[last_node]->backend != GGML_BACKEND_GPU) {
+        last_node -= 1;
+    }
+
     for (int i = 0; i < cgraph->n_nodes; i++) {
-        ggml_vk_build_graph(cgraph->nodes[i], i == cgraph->n_nodes - 1);
+        ggml_vk_build_graph(cgraph->nodes[i], i == last_node);
     }
 
     ggml_compute_params params = {};