fix: benchmark Vulkan direct convolution in auto mode (#2141)

docs/backend.md

@@ -228,17 +228,20 @@ sd-cli -m model.safetensors -p "a cat" --conv-direct off,diffusion=on
 
 - `on` requests direct 2D and 3D convolution.
 - `off` uses the existing im2col-based convolution paths.
-- `auto` enables direct convolution on compute backends with implicit GEMM
+- `auto` first checks support for a representative F16-weight/F32-input direct
-  implementations, without running a benchmark. A metadata-only probe first
+  convolution graph. Vulkan then benchmarks direct convolution against im2col
-  checks support for a representative F16-weight/F32-input direct convolution
+  using two small shapes per dimension, with warmup, output validation, and five
-  graph. Results and decision logs are cached per compute backend and dimension
+  interleaved measurements per shape. Direct convolution must improve the sum of
-  in the context, as with flash attention. Vulkan enables both 2D and 3D.
+  median times by at least 5%; a failed probe defaults to `off`.
+  Other backends retain their implicit GEMM checks without benchmarking:
   CUDA 2D requires a compiled Turing-or-newer MMA architecture; HIP 2D requires
   the supported CDNA MFMA or RDNA WMMA architectures described below for flash
   attention. MUSA 2D has a SIMT implementation. CUDA, HIP, and MUSA 3D also have
   a SIMT implementation, so they do not require MMA. Other backends default to
-  `off`. Each dimension is selected independently, and all compute devices in
+  `off`. Results and decision logs are cached per compute backend and dimension
-  a module must qualify. Parameter-storage devices do not affect this choice.
+  in the context, as with flash attention. Each dimension is selected independently,
+  and all compute devices in a module must qualify. Parameter-storage devices do
+  not affect this choice.
 
 These are backend-level defaults. GGML still chooses the actual kernel based on
 weight type, shape, and hardware; enabling direct convolution does not guarantee

docs/backend_probe.md

@@ -5,7 +5,8 @@ policy stays separate: `src/core/probe/flash_attn.cpp`, for example, defines att
 shapes and its CPU/CUDA/HIP/Vulkan shortcuts, then calls `sd_run_backend_probe` for unknown
 configurations. `src/core/probe/conv_direct.cpp` uses metadata-only probe graphs
 to check direct 2D/3D convolution support, then enables known implicit GEMM
-paths. It does not allocate device tensors, execute kernels, or run measurements.
+paths on non-Vulkan backends. Vulkan compares direct convolution with im2col
+using the shared measurement machinery.
 
 For each representative shape, create a `BackendProbeCase`:
 

docs/performance.md

@@ -5,8 +5,13 @@ eg.:
  - flux 768x768 ~600mb
  - SD2 768x768 ~1400mb
 
-For most backends, it slows things down, but for cuda it generally speeds it up too.
+`--flash-attn on|off|auto` controls flash attention. Its effect on speed depends
-At the moment, it is only supported for some models and some backends (like cpu, cuda/rocm, metal).
+on GPU architecture, backend, driver, model, and tensor shapes. Automatic
+selection uses backend rules and, for some paths, small probes; it does not
+guarantee the fastest end-to-end generation. Compare `on` and `off` with the
+same model and generation settings over repeated runs, separating startup/probe
+time from steady-state generation time. See [backend selection](backend.md)
+for automatic selection rules and per-module overrides.
 
 Watch for:
 ```
@@ -21,6 +26,19 @@ CPU fallback. It excludes weights and cache buffers. Within a runner lifecycle,
 the summary is printed only on the first graph or when backend capacities or the
 segment count change.
 
+## Tune direct convolution.
+
+`--conv-direct on|off|auto` selects direct convolution or the im2col-based
+path. Performance depends on GPU architecture, backend, driver, model, and
+tensor shapes, so enabling direct convolution can either improve or reduce speed.
+
+Automatic selection uses small benchmarks on Vulkan and implicit GEMM checks
+on other backends. It does not guarantee the fastest end-to-end generation.
+Compare `on` and `off` with the same model and generation settings over
+repeated runs, keeping flash attention unchanged and separating startup/probe
+time from steady-state generation time. See [backend selection](backend.md)
+for automatic selection rules and per-module overrides.
+
 ## Run conditional and unconditional CFG in one batched UNet forward.
 
 For UNet models, the conditional and unconditional guidance branches are

src/core/probe/conv_direct.cpp

@@ -1,26 +1,52 @@
 #include "core/probe/conv_direct.h"
 
+#include <cmath>
 #include <exception>
+#include <utility>
+#include <vector>
 
 #include "core/ggml_extend_backend.h"
 #include "core/probe/backend_probe.h"
 #include "core/util.h"
 
-static void build_conv_direct_probe(BackendProbeGraph& graph, bool conv3d) {
+static void build_conv_direct_probe(BackendProbeGraph& graph, bool conv3d, bool direct, int size = 32, int channels = 64) {
-    constexpr int channels = 64;
     if (conv3d) {
-        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, 16, 16, 4, channels);
+        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, size / 2, size / 2, 4, channels);
         auto weight  = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F16, 3, 3, 3, channels * channels);
         graph.inputs = {input, weight};
-        graph.output = ggml_conv_3d_direct(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1, 1, 1, 1, channels, 1, channels);
+        graph.output = direct ? ggml_conv_3d_direct(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1, 1, 1, 1, channels, 1, channels)
+                              : ggml_conv_3d(graph.ctx, weight, input, channels, 1, 1, 1, 1, 1, 1, 1, 1, 1);
     } else {
-        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, 32, 32, channels, 1);
+        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, size, size, channels, 1);
         auto weight  = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F16, 3, 3, channels, channels);
         graph.inputs = {input, weight};
-        graph.output = ggml_conv_2d_direct(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1);
+        graph.output = direct ? ggml_conv_2d_direct(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1)
+                              : ggml_conv_2d(graph.ctx, weight, input, 1, 1, 1, 1, 1, 1);
     }
 }
 
+static bool initialize_conv_probe(BackendProbeGraph& graph, bool conv3d) {
+    uint32_t state = 42;
+    for (auto* tensor : graph.inputs) {
+        const size_t count = static_cast<size_t>(ggml_nelements(tensor));
+        const float scale  = tensor->type == GGML_TYPE_F16 ? 1.f / std::sqrt(float(graph.inputs[0]->ne[conv3d ? 3 : 2] * (conv3d ? 27 : 9))) : 1.f;
+        std::vector<ggml_fp16_t> values(count);
+        for (auto& value : values) {
+            state = state * 1664525u + 1013904223u;
+            value = ggml_fp32_to_fp16((static_cast<float>(state >> 8) / 16777216.f - 0.5f) * scale);
+        }
+        if (tensor->type == GGML_TYPE_F16) {
+            ggml_backend_tensor_set(tensor, values.data(), 0, count * sizeof(ggml_fp16_t));
+        } else {
+            // im2col rounds inputs to F16; use the same input precision for both paths.
+            std::vector<float> input(count);
+            ggml_fp16_to_fp32_row(values.data(), input.data(), count);
+            ggml_backend_tensor_set(tensor, input.data(), 0, count * sizeof(float));
+        }
+    }
+    return true;
+}
+
 bool sd_probe_conv_direct(ggml_backend_t backend, bool conv3d) {
     if (backend == nullptr) {
         return false;
@@ -33,11 +59,30 @@ bool sd_probe_conv_direct(ggml_backend_t backend, bool conv3d) {
             LOG_INFO("conv%dd direct auto: %s -> off (probe graph unavailable)", dims, name);
             return false;
         }
-        build_conv_direct_probe(capability, conv3d);
+        build_conv_direct_probe(capability, conv3d, true);
         if (!capability.supported()) {
             LOG_INFO("conv%dd direct auto: %s -> off (unsupported probe graph)", dims, name);
             return false;
         }
+        if (sd_backend_is(backend, "Vulkan")) {
+            std::vector<BackendProbeCase> cases;
+            for (const auto& shape : {std::pair<int, int>{32, 64}, {16, 128}}) {
+                BackendProbeCase test;
+                test.build = [conv3d, shape](BackendProbeGraph& graph, bool direct) {
+                    build_conv_direct_probe(graph, conv3d, direct, shape.first, shape.second);
+                };
+                test.initialize = [conv3d](BackendProbeGraph& graph) { return initialize_conv_probe(graph, conv3d); };
+                cases.push_back(std::move(test));
+            }
+            const auto result = sd_run_backend_probe(backend, 0, cases);
+            if (!result.completed) {
+                LOG_WARN("conv%dd direct auto: %s -> off (%s)", dims, name, result.error.c_str());
+                return false;
+            }
+            LOG_INFO("conv%dd direct auto: %s -> %s (probe: im2col %.3f ms, direct %.3f ms)",
+                     dims, name, result.candidate_faster ? "on" : "off", result.reference_ms, result.candidate_ms);
+            return result.candidate_faster;
+        }
         if (sd_backend_supports_conv_implicit_gemm(backend, conv3d)) {
             LOG_INFO("conv%dd direct auto: %s -> on (implicit GEMM)", dims, name);
             return true;