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

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) {