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__