feat: add automatic direct convolution selection (#2134)

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;