feat: add automatic direct convolution selection

docs/backend.md

@@ -215,6 +215,47 @@ If a VAE decode fails, decoding retries with spatial tiling even when `--auto-fi
 is off; supported video decoders try temporal tiling first and can then add
 spatial tiling. Spatial retries use half-size tiles along each latent dimension.
 
+## Direct convolution (`--conv-direct on|off|auto`)
+
+`--conv-direct` defaults to `auto` and accepts module assignments for
+`diffusion`, `vae` (including preview VAE), and `controlnet`, using the same
+aliases and precedence rules as `--flash-attn`:
+
+```shell
+sd-cli -m model.safetensors -p "a cat" --conv-direct auto,vae=off
+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
+  implementations, without running a benchmark. A metadata-only probe first
+  checks support for a representative F16-weight/F32-input direct convolution
+  graph. Results and decision logs are cached per compute backend and dimension
+  in the context, as with flash attention. Vulkan enables both 2D and 3D.
+  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
+  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
+that every operation uses implicit GEMM. The CUDA-family implicit GEMM paths
+require F16 weights and can use ordinary direct kernels for other cases.
+This option does not change the separately created image upscaler context.
+
+Repeating `--conv-direct` replaces the previous assignment string. Existing
+`--diffusion-conv-direct` (diffusion and ControlNet) and `--vae-conv-direct`
+flags are deprecated and emit a warning, but remain supported when
+`--conv-direct` is omitted. Use `--conv-direct auto,diffusion=on,controlnet=on`
+and `--conv-direct auto,vae=on`, respectively. An explicit
+`--conv-direct` takes precedence over both legacy flags regardless of order.
+Library callers can set the appended `sd_ctx_params_t::conv_direct_mode` field
+to the same string; null or empty selects `auto` while honoring legacy true
+flags. The context copies the string.
+
 ## Flash attention (`--flash-attn on|off|auto`)
 
 `--flash-attn` defaults to `auto`. Selection happens while creating the model

docs/backend_probe.md

@@ -3,8 +3,9 @@
 `src/core/probe/backend_probe.h` provides the shared measurement machinery. Operator
 policy stays separate: `src/core/probe/flash_attn.cpp`, for example, defines attention
 shapes and its CPU/CUDA/HIP/Vulkan shortcuts, then calls `sd_run_backend_probe` for unknown
-configurations. A future convolution probe can use the same runner without
-changing the attention implementation.
+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.
 
 For each representative shape, create a `BackendProbeCase`:
 
@@ -40,7 +41,9 @@ cache by both operation and backend); an attention result must not select a
 convolution implementation. If shapes or settings vary within one context,
 include them in the corresponding cache key as well.
 
-No convolution policy or CLI option is added by this infrastructure.
+Convolution results are cached separately for 2D and 3D, once per compute backend
+in the context, just like attention results. Each probe logs its backend, decision,
+and reason; mixed-device modules log when their compute devices disagree.
 
 Module assignment parsing remains in `src/core/ggml_extend_backend.cpp`, next
 to backend assignment parsing. The probe directory does not parse CLI/API options.

examples/cli/README.md

@@ -17,6 +17,11 @@ for probing behavior.
 Use module assignments such as `--flash-attn auto,te=off,vae=on` to override
 individual modules.
 
+`--conv-direct on|off|auto` also defaults to `auto`, enabling direct convolution
+on backends with implicit GEMM implementations. Use `--conv-direct auto,vae=off`
+for module overrides; supported modules are `diffusion`, `vae`, and `controlnet`.
+See [Direct convolution](../../docs/backend.md#direct-convolution---conv-direct-onoffauto).
+
 For direct image repair or automatic post-generation YOLOv8 detection followed by cropped inpainting, see
 [ADetailer](../../docs/adetailer.md).
 

examples/common/common.cpp

@@ -636,14 +636,6 @@ ArgOptions SDContextParams::get_options() {
          "--sage-attn",
          "use native CUDA SageAttention in the diffusion model, with flash/default attention fallback",
          true, &sage_attn},
-        {"",
-         "--diffusion-conv-direct",
-         "use ggml_conv2d_direct in the diffusion model",
-         true, &diffusion_conv_direct},
-        {"",
-         "--vae-conv-direct",
-         "use direct 2D and 3D convolutions in the vae model",
-         true, &vae_conv_direct},
     };
 
     auto on_type_arg = [&](int argc, const char** argv, int index) {
@@ -729,6 +721,33 @@ ArgOptions SDContextParams::get_options() {
          [this](int argc, const char** argv, int index) {
              return parse_scale_override(argc, argv, index, attn_scale);
          }},
+        {"",
+         "--diffusion-conv-direct",
+         "deprecated; use --conv-direct auto,diffusion=on,controlnet=on",
+         [this](int, const char**, int) {
+             LOG_WARN("--diffusion-conv-direct is deprecated; use --conv-direct auto,diffusion=on,controlnet=on instead.");
+             diffusion_conv_direct = true;
+             return 0;
+         }},
+        {"",
+         "--vae-conv-direct",
+         "deprecated; use --conv-direct auto,vae=on",
+         [this](int, const char**, int) {
+             LOG_WARN("--vae-conv-direct is deprecated; use --conv-direct auto,vae=on instead.");
+             vae_conv_direct = true;
+             return 0;
+         }},
+        {"",
+         "--conv-direct",
+         "on|off|auto or diffusion/vae/controlnet assignments (default: auto), e.g. auto,vae=off. Overrides legacy convolution flags",
+         [this](int argc, const char** argv, int index) {
+             if (++index >= argc || argv[index][0] == '\0') {
+                 LOG_ERROR("--conv-direct requires on, off, auto, or module assignments");
+                 return -1;
+             }
+             conv_direct_mode = argv[index];
+             return 1;
+         }},
         {"",
          "--flash-attn",
          "on|off|auto or module assignments (default: auto), e.g. auto,te=off,vae=on. Last repeated option wins",
@@ -984,6 +1003,7 @@ std::string SDContextParams::to_string() const {
         << "  clip_on_cpu: " << (clip_on_cpu ? "true" : "false") << ",\n"
         << "  vae_on_cpu: " << (vae_on_cpu ? "true" : "false") << ",\n"
         << "  flash_attn_mode: " << flash_attn_mode << ",\n"
+        << "  conv_direct_mode: " << (conv_direct_mode.empty() ? "auto (legacy flags honored)" : conv_direct_mode) << ",\n"
         << "  sage_attn: " << (sage_attn ? "true" : "false") << ",\n"
         << "  linear_scale: " << linear_scale << ",\n"
         << "  attn_scale: " << attn_scale << ",\n"
@@ -1042,6 +1062,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) {
     sd_ctx_params.lora_apply_mode                 = lora_apply_mode;
     sd_ctx_params.enable_mmap                     = enable_mmap;
     sd_ctx_params.flash_attn_mode                 = flash_attn_mode.c_str();
+    sd_ctx_params.conv_direct_mode                = conv_direct_mode.c_str();
     sd_ctx_params.sage_attn                       = sage_attn;
     sd_ctx_params.linear_scale                    = linear_scale;
     sd_ctx_params.attn_scale                      = attn_scale;

examples/common/common.h

@@ -171,9 +171,10 @@ struct SDContextParams {
     bool clip_on_cpu            = false;
     bool vae_on_cpu             = false;
     std::string flash_attn_mode = "auto";
-    bool sage_attn              = false;
-    bool diffusion_conv_direct  = false;
-    bool vae_conv_direct        = false;
+    std::string conv_direct_mode;  // Empty selects auto while honoring legacy convolution flags.
+    bool sage_attn             = false;
+    bool diffusion_conv_direct = false;
+    bool vae_conv_direct       = false;
 
     prediction_t prediction           = PREDICTION_COUNT;
     lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO;

include/stable-diffusion.h

@@ -252,6 +252,7 @@ typedef struct {
     bool sage_attn;
     int conditioning_cache_size;  // Maximum cached conditioning entries per context; 0 disables caching (default: 4)
     const char* flash_attn_mode;  // on/off/auto or module assignments (e.g. "auto,te=off"); NULL/empty honors legacy flags.
+    const char* conv_direct_mode;  // on/off/auto or diffusion/vae/controlnet assignments; NULL/empty honors legacy flags.
 } sd_ctx_params_t;
 
 typedef struct {

src/core/ggml_extend_backend.cpp

@@ -457,23 +457,23 @@ ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t dev
     return buffer;
 }
 
-bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
+static int cuda_compiled_arch(ggml_backend_t backend) {
 #ifdef SD_USE_CUDA
     if (!sd_backend_is(backend, "CUDA")) {
-        return false;
+        return 0;
     }
     auto dev = ggml_backend_get_device(backend);
     if (dev == nullptr) {
-        return false;
+        return 0;
     }
     static std::mutex mutex;
-    static std::unordered_map<ggml_backend_dev_t, bool> cache;
+    static std::unordered_map<ggml_backend_dev_t, int> cache;
     std::lock_guard<std::mutex> lock(mutex);
     auto it = cache.find(dev);
     if (it != cache.end()) {
         return it->second;
     }
-    const bool supported = [&]() {
+    const int arch = [&]() {
         ggml_backend_dev_props props{};
         ggml_backend_dev_get_props(dev, &props);
         CUdevice device;
@@ -482,13 +482,13 @@ bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
             cuDeviceGetByPCIBusId(&device, props.device_id) != CUDA_SUCCESS ||
             cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device) != CUDA_SUCCESS ||
             cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device) != CUDA_SUCCESS) {
-            return false;
+            return 0;
         }
         auto reg          = ggml_backend_dev_backend_reg(dev);
         auto get_features = reinterpret_cast<ggml_backend_get_features_t>(
             ggml_backend_reg_get_proc_address(reg, "ggml_backend_get_features"));
         if (get_features == nullptr) {
-            return false;
+            return 0;
         }
         // Match ggml's highest compiled architecture for this device, including PTX fallback.
         const int cc      = 100 * major + 10 * minor;
@@ -511,16 +511,21 @@ bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
                 arch = end;
             }
         }
-        return compiled_arch == 700 || compiled_arch >= 750;
+        return compiled_arch;
     }();
-    cache.emplace(dev, supported);
-    return supported;
+    cache.emplace(dev, arch);
+    return arch;
 #else
     (void)backend;
-    return false;
+    return 0;
 #endif
 }
 
+bool sd_backend_supports_cuda_mma(ggml_backend_t backend) {
+    const int arch = cuda_compiled_arch(backend);
+    return arch == 700 || arch >= 750;
+}
+
 bool sd_backend_supports_hip_mma(ggml_backend_t backend) {
 #ifdef SD_USE_HIP
     if (!sd_backend_is(backend, "ROCm")) {
@@ -564,6 +569,20 @@ bool sd_backend_supports_hip_mma(ggml_backend_t backend) {
 #endif
 }
 
+bool sd_backend_supports_conv_implicit_gemm(ggml_backend_t backend, bool conv3d) {
+    if (backend == nullptr) {
+        return false;
+    }
+    if (sd_backend_is(backend, "Vulkan") || sd_backend_is(backend, "MUSA")) {
+        return true;
+    }
+    // CUDA/HIP Conv3d also implements SIMT implicit GEMM; Conv2d requires MMA.
+    if (conv3d) {
+        return sd_backend_is(backend, "CUDA") || sd_backend_is(backend, "ROCm");
+    }
+    return cuda_compiled_arch(backend) >= 750 || sd_backend_supports_hip_mma(backend);
+}
+
 ggml_backend_t sd_backend_cpu_init() {
     ggml_backend_load_all_once();
     return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
@@ -782,7 +801,8 @@ bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* a
     return true;
 }
 
-bool sd_parse_flash_attention_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
+static bool parse_operation_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error, bool convolution) {
+    const std::string operation = convolution ? "direct convolution" : "flash attention";
     if (assignment == nullptr) {
         return false;
     }
@@ -800,7 +820,7 @@ bool sd_parse_flash_attention_assignment(const std::string& spec, SDBackendAssig
                 return true;
             }
             if (error != nullptr) {
-                *error = "invalid flash attention mode '" + mode + "': expected on, off, or auto";
+                *error = "invalid " + operation + " mode '" + mode + "': expected on, off, or auto";
             }
             return false;
         };
@@ -813,15 +833,19 @@ bool sd_parse_flash_attention_assignment(const std::string& spec, SDBackendAssig
         for (const auto& setting : entry.module_names) {
             switch (setting.first) {
                 case SDBackendModule::DIFFUSION:
-                case SDBackendModule::TE:
-                case SDBackendModule::CLIP_VISION:
                 case SDBackendModule::VAE:
                 case SDBackendModule::CONTROL_NET:
-                case SDBackendModule::AUDIO_ENCODER:
                     break;
+                case SDBackendModule::TE:
+                case SDBackendModule::CLIP_VISION:
+                case SDBackendModule::AUDIO_ENCODER:
+                    if (!convolution) {
+                        break;
+                    }
+                    [[fallthrough]];
                 default:
                     if (error != nullptr) {
-                        *error = std::string("flash attention is not configurable for module '") + sd_backend_module_name(setting.first) + "'";
+                        *error = operation + " is not configurable for module '" + sd_backend_module_name(setting.first) + "'";
                     }
                     return false;
             }
@@ -835,6 +859,14 @@ bool sd_parse_flash_attention_assignment(const std::string& spec, SDBackendAssig
     return true;
 }
 
+bool sd_parse_flash_attention_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
+    return parse_operation_assignment(spec, assignment, error, false);
+}
+
+bool sd_parse_conv_direct_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error) {
+    return parse_operation_assignment(spec, assignment, error, true);
+}
+
 bool SDBackendAssignment::empty() const {
     return default_name.empty() && module_names.empty();
 }

src/core/ggml_extend_backend.h

@@ -89,6 +89,7 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name);
 bool sd_backend_is_cpu(ggml_backend_t backend);
 bool sd_backend_supports_cuda_mma(ggml_backend_t backend);
 bool sd_backend_supports_hip_mma(ggml_backend_t backend);
+bool sd_backend_supports_conv_implicit_gemm(ggml_backend_t backend, bool conv3d);
 ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
                                                           void* ptr,
                                                           size_t size,
@@ -102,6 +103,7 @@ ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
 std::string sd_backend_resolve_name(const std::string& name);
 bool sd_parse_backend_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error);
 bool sd_parse_flash_attention_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error);
+bool sd_parse_conv_direct_assignment(const std::string& spec, SDBackendAssignment* assignment, std::string* error);
 const char* sd_backend_module_name(SDBackendModule module);
 void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value);
 bool add_rpc_devices(const std::string& servers);

src/core/probe/conv_direct.cpp

@@ -0,0 +1,51 @@
+#include "core/probe/conv_direct.h"
+
+#include <exception>
+
+#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) {
+    constexpr int channels = 64;
+    if (conv3d) {
+        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, 16, 16, 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);
+    } else {
+        auto input   = ggml_new_tensor_4d(graph.ctx, GGML_TYPE_F32, 32, 32, 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);
+    }
+}
+
+bool sd_probe_conv_direct(ggml_backend_t backend, bool conv3d) {
+    if (backend == nullptr) {
+        return false;
+    }
+    const char* name = ggml_backend_name(backend);
+    const int dims   = conv3d ? 3 : 2;
+    try {
+        BackendProbeGraph capability(backend);
+        if (capability.ctx == nullptr) {
+            LOG_INFO("conv%dd direct auto: %s -> off (probe graph unavailable)", dims, name);
+            return false;
+        }
+        build_conv_direct_probe(capability, conv3d);
+        if (!capability.supported()) {
+            LOG_INFO("conv%dd direct auto: %s -> off (unsupported probe graph)", dims, name);
+            return false;
+        }
+        if (sd_backend_supports_conv_implicit_gemm(backend, conv3d)) {
+            LOG_INFO("conv%dd direct auto: %s -> on (implicit GEMM)", dims, name);
+            return true;
+        }
+        LOG_INFO("conv%dd direct auto: %s -> off (no implicit GEMM path)", dims, name);
+        return false;
+    } catch (const std::exception& error) {
+        LOG_WARN("conv%dd direct auto: %s -> off (probe failed: %s)", dims, name, error.what());
+        return false;
+    }
+}

src/core/probe/conv_direct.h

@@ -0,0 +1,8 @@
+#ifndef __SD_CORE_PROBE_CONV_DIRECT_H__
+#define __SD_CORE_PROBE_CONV_DIRECT_H__
+
+#include "ggml-backend.h"
+
+bool sd_probe_conv_direct(ggml_backend_t backend, bool conv3d);
+
+#endif  // __SD_CORE_PROBE_CONV_DIRECT_H__