[ggml / Vulkan] queue user inputs in the backend stream;

ggml/src/ggml-backend-impl.h

@@ -154,6 +154,11 @@ extern "C" {
 
         // (optional) sort/optimize the nodes in the graph
         void                      (*graph_optimize)    (ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params);
+
+        // (optional) queue a write of host data to a tensor, ordered after the work already queued on the backend
+        // unlike set_tensor_async, the data is copied before the call returns, so the caller may reuse it at once
+        // returns false if the backend cannot do this for the tensor; the caller then writes it synchronously
+        bool (*set_tensor_async_staged)(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size);
     };
 
     struct ggml_backend {

ggml/src/ggml-backend-meta.cpp

@@ -2499,6 +2499,7 @@ static const ggml_backend_i ggml_backend_meta_i = {
     /* .event_record            = */ nullptr,
     /* .event_wait              = */ nullptr,
     /* .graph_optimize          = */ nullptr,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 bool ggml_backend_is_meta(ggml_backend_t backend) {

ggml/src/ggml-backend.cpp

@@ -1816,12 +1816,17 @@ static void ggml_backend_sched_copy_input(ggml_backend_sched_t sched, struct ggm
 
     if (input->flags & GGML_TENSOR_FLAG_INPUT) {
         // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done
-        if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
-            ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
-        } else {
-            ggml_backend_synchronize(split_backend);
+        // a backend that stages the data when the write is queued orders it after its pending work, without a wait
+        const bool staged = split_backend->iface.set_tensor_async_staged != NULL && ggml_backend_buffer_is_host(input->buffer) &&
+            split_backend->iface.set_tensor_async_staged(split_backend, input_cpy, input->data, 0, ggml_nbytes(input));
+        if (!staged) {
+            if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
+                ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
+            } else {
+                ggml_backend_synchronize(split_backend);
+            }
+            ggml_backend_tensor_copy(input, input_cpy);
         }
-        ggml_backend_tensor_copy(input, input_cpy);
         return;
     }
 

ggml/src/ggml-blas/ggml-blas.cpp

@@ -276,6 +276,7 @@ static struct ggml_backend_i blas_backend_i = {
     /* .event_record            = */ NULL,
     /* .event_wait              = */ NULL,
     /* .graph_optimize          = */ NULL,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 static ggml_guid_t ggml_backend_blas_guid(void) {

ggml/src/ggml-cann/ggml-cann.cpp

@@ -2769,6 +2769,7 @@ static const ggml_backend_i ggml_backend_cann_interface = {
     /* .event_record            = */ ggml_backend_cann_event_record,
     /* .event_wait              = */ ggml_backend_cann_event_wait,
     /* .graph_optimize          = */ NULL,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 /**

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -207,6 +207,7 @@ static const struct ggml_backend_i ggml_backend_cpu_i = {
     /* .event_record            = */ NULL,
     /* .event_wait              = */ NULL,
     /* .graph_optimize          = */ NULL,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 static ggml_guid_t ggml_backend_cpu_guid(void) {

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5001,6 +5001,7 @@ static const ggml_backend_i ggml_backend_cuda_interface = {
     /* .event_record            = */ ggml_backend_cuda_event_record,
     /* .event_wait              = */ ggml_backend_cuda_event_wait,
     /* .graph_optimize          = */ ggml_backend_cuda_graph_optimize,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 static ggml_guid_t ggml_backend_cuda_guid() {

ggml/src/ggml-et/ggml-et.cpp

@@ -1613,6 +1613,7 @@ static const struct ggml_backend_i ggml_backend_et_i = {
     /* .event_record            = */ NULL,
     /* .event_wait              = */ NULL,
     /* .graph_optimize          = */ NULL,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 static const char * ggml_backend_et_device_get_name(ggml_backend_dev_t dev) {

ggml/src/ggml-hexagon/ggml-hexagon.cpp

@@ -8704,6 +8704,7 @@ static struct ggml_backend_i hexagon_backend_i = {
     /* .event_record            = */ ggml_backend_hexagon_event_record,
     /* .event_wait              = */ ggml_backend_hexagon_event_wait,
     /* .graph_optimize          = */ ggml_backend_hexagon_graph_optimize,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 static ggml_guid_t ggml_backend_hexagon_guid() {

ggml/src/ggml-metal/ggml-metal.cpp

@@ -604,6 +604,7 @@ static ggml_backend_i ggml_backend_metal_i = {
     /* .event_record            = */ ggml_backend_metal_event_record,
     /* .event_wait              = */ ggml_backend_metal_event_wait,
     /* .graph_optimize          = */ ggml_backend_metal_graph_optimize,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 static ggml_guid_t ggml_backend_metal_guid(void) {

ggml/src/ggml-opencl/ggml-opencl.cpp

@@ -9320,6 +9320,7 @@ static ggml_backend_i ggml_backend_opencl_i = {
     /* .event_record            = */ NULL,
     /* .event_wait              = */ NULL,
     /* .graph_optimize          = */ NULL,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 ggml_backend_t ggml_backend_opencl_init(void) {

ggml/src/ggml-openvino/ggml-openvino.cpp

@@ -896,6 +896,7 @@ static const ggml_backend_i ggml_backend_openvino_interface = {
     /* .event_record            = */ NULL,
     /* .event_wait              = */ NULL,
     /* .graph_optimize          = */ NULL,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 int ggml_backend_openvino_get_device_count() {

ggml/src/ggml-rpc/ggml-rpc.cpp

@@ -1235,6 +1235,7 @@ static ggml_backend_i ggml_backend_rpc_interface = {
     /* .event_record            = */ ggml_backend_rpc_event_record,
     /* .event_wait              = */ ggml_backend_rpc_event_wait,
     /* .graph_optimize          = */ NULL,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 ggml_backend_buffer_type_t ggml_backend_rpc_buffer_type(const char * endpoint, uint32_t device) {

ggml/src/ggml-sycl/ggml-sycl.cpp

@@ -6505,6 +6505,7 @@ static ggml_backend_i ggml_backend_sycl_interface = {
     /* .event_record            = */ ggml_backend_sycl_event_record,
     /* .event_wait              = */ ggml_backend_sycl_event_wait,
     /* .graph_optimize          = */ NULL,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 static ggml_guid_t ggml_backend_sycl_guid() {

ggml/src/ggml-virtgpu/ggml-backend.cpp

@@ -47,6 +47,7 @@ static ggml_backend_i ggml_backend_remoting_interface = {
     /* .event_record            = */ NULL,
     /* .event_wait              = */ NULL,
     /* .graph_optimize          = */ ggml_backend_remoting_graph_optimize,
+    /* .set_tensor_async_staged = */ NULL,
 };
 
 static ggml_guid_t ggml_backend_remoting_guid() {

ggml/src/ggml-vulkan/ggml-vulkan-types.h

@@ -712,6 +712,9 @@ struct vk_device_struct {
     bool single_queue;
     bool support_async;
     bool async_use_transfer_queue;
+    // largest host-to-device copy that goes into the compute stream, measured on first use
+    size_t compute_stream_copy_max;
+    bool compute_stream_copy_max_known {};
     bool has_internally_synchronized_queues = false;
     uint32_t subgroup_size;
     uint32_t subgroup_size_log2;
@@ -1250,6 +1253,8 @@ struct ggml_backend_vk_context {
     ggml_vk_garbage_collector gc;
     size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k, prealloc_size_add_rms_partials, prealloc_size_add_rms_partials_offset;
     vk_buffer prealloc_x, prealloc_y, prealloc_split_k, prealloc_add_rms_partials, sync_staging;
+    std::vector<vk_buffer> input_staging;
+    size_t input_staging_used {};
     vk::Fence fence, almost_ready_fence;
     bool submit_pending {};
     bool almost_ready_fence_pending {};