feat: add automatic direct convolution selection (#2134)

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__

src/model_manager.cpp

@@ -1762,7 +1762,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity(
             }
             // A wrapped report may reflect real pressure; reclaim before accepting the tracked estimate.
             LOG_VERBOSE("ignoring impossible free memory report on %s (%.2f MB free / %.2f MB total), using tracked residency",
-                        ggml_backend_name(backend), free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0));
+                           ggml_backend_name(backend), free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0));
             free_bytes                  = total_bytes;
             result.requires_reclamation = true;
         }

src/pipeline/diffusion_engine.cpp

@@ -18,6 +18,7 @@
 #include "core/ggml_runner.h"
 #include "core/ggml_tensor_utils.h"
 #include "core/layer_split_partition.h"
+#include "core/probe/conv_direct.h"
 #include "core/probe/flash_attn.h"
 #include "model.h"
 
@@ -951,7 +952,15 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) {
         LOG_ERROR("invalid flash_attn_mode '%s': %s", flash_mode.c_str(), flash_error.c_str());
         return false;
     }
+    const std::string conv_mode = SAFE_STR(sd_ctx_params->conv_direct_mode);
+    std::string conv_error;
+    if (!sd_parse_conv_direct_assignment(conv_mode, &conv_direct_assignment, &conv_error)) {
+        LOG_ERROR("invalid conv_direct_mode '%s': %s", conv_mode.c_str(), conv_error.c_str());
+        return false;
+    }
     flash_attn_by_backend.clear();
+    conv2d_direct_by_backend.clear();
+    conv3d_direct_by_backend.clear();
     conditioning_cache_->set_capacity(static_cast<size_t>(sd_ctx_params->conditioning_cache_size));
     auto configuration = std::make_unique<ModelConfig>(*sd_ctx_params);
     n_threads          = sd_ctx_params->n_threads;
@@ -1224,6 +1233,34 @@ bool StableDiffusionGGML::flash_attention_for(SDBackendModule module) {
     return enabled;
 }
 
+bool StableDiffusionGGML::conv_direct_for(SDBackendModule module, bool conv3d) {
+    const auto& params     = config_->params;
+    const std::string mode = conv_direct_assignment.get(module);
+    if (mode == "on" || mode == "off") {
+        return mode == "on";
+    }
+    if (strlen(SAFE_STR(params.conv_direct_mode)) == 0 &&
+        (module == SDBackendModule::VAE ? params.vae_conv_direct : params.diffusion_conv_direct)) {
+        return true;
+    }
+    auto& by_backend = conv3d ? conv3d_direct_by_backend : conv2d_direct_by_backend;
+    auto backends    = backend_manager.runtime_backends(module);
+    bool enabled     = !backends.empty();
+    bool any_enabled = false;
+    for (auto backend : backends) {
+        auto it = by_backend.find(backend);
+        if (it == by_backend.end()) {
+            it = by_backend.emplace(backend, sd_probe_conv_direct(backend, conv3d)).first;
+        }
+        enabled     = enabled && it->second;
+        any_enabled = any_enabled || it->second;
+    }
+    if (any_enabled && !enabled) {
+        LOG_INFO("conv%dd direct auto: %s -> off (compute devices disagree)", conv3d ? 3 : 2, sd_backend_module_name(module));
+    }
+    return enabled;
+}
+
 bool StableDiffusionGGML::validate_and_load_runners() {
     const auto* sd_ctx_params   = &config_->params;
     const bool use_tae          = uses_tae();
@@ -1261,6 +1298,23 @@ bool StableDiffusionGGML::validate_and_load_runners() {
     if (sd_ctx_params->sage_attn && !set_sage_attention_enabled(true)) {
         return false;
     }
+    const auto set_conv_direct = [this](SDBackendModule module, const std::vector<GGMLRunner*>& runners) {
+        const bool conv2d = conv_direct_for(module, false);
+        const bool conv3d = conv_direct_for(module, true);
+        for (auto runner : runners) {
+            if (runner != nullptr) {
+                runner->set_conv2d_direct_enabled(conv2d);
+                runner->set_conv3d_direct_enabled(conv3d);
+            }
+        }
+    };
+    set_conv_direct(SDBackendModule::DIFFUSION, {diffusion_model.get(), high_noise_diffusion_model.get()});
+    if (first_stage_model || preview_vae || audio_vae_model) {
+        set_conv_direct(SDBackendModule::VAE, {first_stage_model.get(), preview_vae.get(), audio_vae_model.get()});
+    }
+    if (control_net) {
+        set_conv_direct(SDBackendModule::CONTROL_NET, {control_net.get()});
+    }
     LOG_VERBOSE("validating model metadata");
 
     std::set<std::string> ignore_tensors;

src/pipeline/diffusion_engine.h

@@ -51,6 +51,9 @@ public:
     SDBackendManager backend_manager;
     std::unordered_map<ggml_backend_t, bool> flash_attn_by_backend;
     SDBackendAssignment flash_attn_assignment;
+    SDBackendAssignment conv_direct_assignment;
+    std::unordered_map<ggml_backend_t, bool> conv2d_direct_by_backend;
+    std::unordered_map<ggml_backend_t, bool> conv3d_direct_by_backend;
 
     SDVersion version;
     bool external_vae_is_invalid = false;
@@ -134,7 +137,7 @@ public:
                                 &sd_ctx_params_t::clip_g_path, &sd_ctx_params_t::clip_vision_path,
                                 &sd_ctx_params_t::t5xxl_path, &sd_ctx_params_t::llm_path,
                                 &sd_ctx_params_t::llm_vision_path, &sd_ctx_params_t::diffusion_model_path,
-                                &sd_ctx_params_t::tokenizer, &sd_ctx_params_t::flash_attn_mode,
+                                &sd_ctx_params_t::tokenizer, &sd_ctx_params_t::flash_attn_mode, &sd_ctx_params_t::conv_direct_mode,
                                 &sd_ctx_params_t::high_noise_diffusion_model_path, &sd_ctx_params_t::uncond_diffusion_model_path,
                                 &sd_ctx_params_t::embeddings_connectors_path, &sd_ctx_params_t::vae_path,
                                 &sd_ctx_params_t::audio_vae_path, &sd_ctx_params_t::taesd_path,
@@ -338,6 +341,7 @@ public:
 
     bool validate_and_load_runners();
     bool flash_attention_for(SDBackendModule module);
+    bool conv_direct_for(SDBackendModule module, bool conv3d);
 
     bool build_denoiser();
 

src/pipeline/model_builders.cpp

@@ -478,10 +478,6 @@ namespace sd::model_builders {
                                                                  "model.diffusion_model",
                                                                  version,
                                                                  weight_manager);
-            if (sd_ctx_params->diffusion_conv_direct) {
-                LOG_INFO("Using Conv2d direct in the diffusion model");
-                result.diffusion->set_conv2d_direct_enabled(true);
-            }
         }
 
         if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && result.clip_vision == nullptr) {
@@ -668,15 +664,6 @@ namespace sd::model_builders {
             }
         }
 
-        if (sd_ctx_params->vae_conv_direct) {
-            LOG_INFO("Using Conv2d/Conv3d direct in the vae model");
-            result.vae->set_conv2d_direct_enabled(true);
-            result.vae->set_conv3d_direct_enabled(true);
-            if (result.preview) {
-                result.preview->set_conv2d_direct_enabled(true);
-                result.preview->set_conv3d_direct_enabled(true);
-            }
-        }
         if (result.vae) {
             result.vae->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
         }
@@ -703,10 +690,6 @@ namespace sd::model_builders {
                                                         version,
                                                         "",
                                                         weight_manager);
-        if (sd_ctx_params->diffusion_conv_direct) {
-            LOG_INFO("Using Conv2d direct in the control net");
-            control_net->set_conv2d_direct_enabled(true);
-        }
         control_net->set_scale_overrides(sd_ctx_params->linear_scale, sd_ctx_params->attn_scale);
         runner = std::move(control_net);
         return true;

src/stable-diffusion.cpp

@@ -339,6 +339,7 @@ void sd_ctx_params_init(sd_ctx_params_t* sd_ctx_params) {
     sd_ctx_params->eager_load                = false;
     sd_ctx_params->enable_mmap               = false;
     sd_ctx_params->flash_attn_mode           = nullptr;
+    sd_ctx_params->conv_direct_mode          = nullptr;
     sd_ctx_params->diffusion_flash_attn      = false;
     sd_ctx_params->sage_attn                 = false;
     sd_ctx_params->linear_scale              = 0.f;
@@ -397,6 +398,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              "auto_fit: %s\n"
              "flash_attn: %s\n"
              "flash_attn_mode: %s\n"
+             "conv_direct_mode: %s\n"
              "diffusion_flash_attn: %s\n"
              "sage_attn: %s\n"
              "linear_scale: %g\n"
@@ -439,6 +441,7 @@ char* sd_ctx_params_to_str(const sd_ctx_params_t* sd_ctx_params) {
              BOOL_STR(sd_ctx_params->auto_fit),
              BOOL_STR(sd_ctx_params->flash_attn),
              sd_ctx_params->flash_attn_mode && sd_ctx_params->flash_attn_mode[0] ? sd_ctx_params->flash_attn_mode : "auto (legacy flags honored)",
+             sd_ctx_params->conv_direct_mode && sd_ctx_params->conv_direct_mode[0] ? sd_ctx_params->conv_direct_mode : "auto (legacy flags honored)",
              BOOL_STR(sd_ctx_params->diffusion_flash_attn),
              BOOL_STR(sd_ctx_params->sage_attn),
              sd_ctx_params->linear_scale,