perf: optimize Vulkan kernels, graph fusion, and workspace reuse (#2251)

src/core/compute_workspace.cpp

@@ -92,7 +92,8 @@ namespace sd {
     ComputeWorkspace::Measurement ComputeWorkspace::measure(
         ggml_cgraph* graph,
         const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
-        const AssignNodes& assign_nodes) {
+        const AssignNodes& assign_nodes,
+        bool cache_measurement) {
         std::vector<const ggml_tensor*> tensors;
         std::unordered_set<const ggml_tensor*> seen;
         auto visit = [&](const ggml_tensor* tensor) {
@@ -112,6 +113,52 @@ namespace sd {
                 visit(source);
             }
         }
+        std::vector<uint8_t> key;
+        if (cache_measurement) {
+            key.reserve(tensors.size() * sizeof(ggml_tensor));
+            auto append = [&](const auto& value) {
+                const auto* bytes = reinterpret_cast<const uint8_t*>(&value);
+                key.insert(key.end(), bytes, bytes + sizeof(value));
+            };
+            append(graph->n_nodes);
+            append(graph->n_leafs);
+            append(graph->order);
+            for (auto tensor : tensors) {
+                append(tensor);
+                append(tensor->type);
+                append(tensor->ne);
+                append(tensor->nb);
+                append(tensor->op);
+                append(tensor->op_params);
+                append(tensor->flags);
+                append(tensor->src);
+                append(tensor->view_src);
+                append(tensor->view_offs);
+                append(tensor->name);
+                const auto external = external_backend(tensor);
+                append(external);
+                if (external == nullptr || tensor->view_src != nullptr) {
+                    append(tensor->buffer);
+                    append(tensor->data);
+                    append(tensor->extra);
+                    if (tensor->buffer != nullptr) {
+                        const auto usage = ggml_backend_buffer_get_usage(tensor->buffer);
+                        const auto type  = ggml_backend_buffer_get_type(tensor->buffer);
+                        append(usage);
+                        append(type);
+                    }
+                }
+                const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor);
+                const int32_t uses = index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index)
+                                         ? graph->use_counts[index]
+                                         : 0;
+                append(uses);
+            }
+            // Match the full allocation metadata, including topology and external placement.
+            if (key == measurement_key_) {
+                return {measured_buffers_};
+            }
+        }
         const size_t graph_size = tensors.size() + 8;
         auto context            = ggml_init({tensors.size() * ggml_tensor_overhead() + ggml_graph_overhead_custom(graph_size, false), nullptr, true});
         if (context == nullptr) {
@@ -178,6 +225,10 @@ namespace sd {
             ggml_backend_buffer_free(entry.second);
         }
         ggml_free(context);
+        if (cache_measurement && !result.buffers.empty()) {
+            measurement_key_  = std::move(key);
+            measured_buffers_ = result.buffers;
+        }
         return result;
     }
 
@@ -212,6 +263,8 @@ namespace sd {
         if (active_) {
             return false;
         }
+        measurement_key_.clear();
+        measured_buffers_.clear();
         ggml_backend_sched_free(scheduler_);
         scheduler_          = nullptr;
         scheduler_capacity_ = 0;

src/core/compute_workspace.h

@@ -19,6 +19,8 @@ namespace sd {
         ggml_backend_sched_t scheduler_ = nullptr;
         size_t scheduler_capacity_      = 0;
         bool active_                    = false;
+        std::vector<uint8_t> measurement_key_;
+        std::vector<BackendBufferSize> measured_buffers_;
 
         ggml_backend_sched_t make_scheduler(size_t graph_size);
 
@@ -41,7 +43,8 @@ namespace sd {
         Measurement measure(
             ggml_cgraph* graph,
             const std::function<ggml_backend_t(const ggml_tensor*)>& external_backend,
-            const AssignNodes& assign_nodes);
+            const AssignNodes& assign_nodes,
+            bool cache_measurement = false);
         void synchronize() const;
         bool segment_end() noexcept;
         bool release();

src/core/ggml_extend.cpp

@@ -3,6 +3,7 @@
 #include <cmath>
 #include <stdexcept>
 #include <utility>
+#include <vector>
 
 #include "core/ggml_extend_backend.h"
 
@@ -478,21 +479,33 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
         x          = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2));
         x          = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N);
     } else {
-        // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends
-        // (e.g. Metal, see #850) do not implement. Fall back to
-        // GGML_OP_CONV_3D on those backends.
-        bool im2col_3d_supported = true;
+        auto convolution = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
         if (backend != nullptr) {
-            ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
-            im2col_3d_supported = ggml_backend_supports_op(backend, im2col);
-        }
-        if (im2col_3d_supported) {
-            x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2);
-        } else {
-            int64_t OC = w->ne[3] / IC;
-            int64_t N  = x->ne[3] / IC;
-            x          = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+            // im2col can support large buffers that its reshape/GEMM consumers cannot.
+            std::vector<ggml_tensor*> nodes{convolution};
+            bool supported = true;
+            for (size_t i = 0; i < nodes.size() && supported; ++i) {
+                auto node = nodes[i];
+                if (node == x || node == w) {
+                    continue;
+                }
+                supported = ggml_backend_supports_op(backend, node);
+                for (auto source : node->src) {
+                    if (source != nullptr) {
+                        nodes.push_back(source);
+                    }
+                }
+            }
+            if (!supported) {
+                int64_t OC    = w->ne[3] / IC;
+                int64_t N     = x->ne[3] / IC;
+                auto fallback = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+                if (ggml_backend_supports_op(backend, fallback)) {
+                    convolution = fallback;
+                }
+            }
         }
+        x = convolution;
     }
 
     if (scale != 1.f) {

src/core/ggml_runner.cpp

@@ -735,7 +735,7 @@ ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, const std:
     auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) {
         pin_multi_device_nodes(scheduler, copy, graph);
     };
-    return workspace_.measure(graph, external_backend, assign_nodes);
+    return workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device());
 }
 
 std::vector<DeviceMemoryRequest> GGMLRunner::memory_requests(