perf: accelerate VAE direct 3D convolutions (#1996)

examples/common/common.cpp

@@ -624,7 +624,7 @@ ArgOptions SDContextParams::get_options() {
          true, &diffusion_conv_direct},
         {"",
          "--vae-conv-direct",
-         "use ggml_conv2d_direct in the vae model",
+         "use direct 2D and 3D convolutions in the vae model",
          true, &vae_conv_direct},
     };
 

ggml

@@ -1 +1 @@
-Subproject commit 1e22ec0d04b43afa69963e4b4ea6683535d54d79
+Subproject commit c6632cd905401abc58b6f5cdd52d228aa7ca1b88

src/core/ggml_extend.cpp

@@ -452,8 +452,13 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d0,
                               int d1,
                               int d2,
-                              bool force_prec_f32) {
+                              bool force_prec_f32,
-    if (force_prec_f32) {
+                              bool direct) {
+    if (direct) {
+        int64_t OC = w->ne[3] / IC;
+        int64_t N  = x->ne[3] / IC;
+        x          = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC);
+    } else if (force_prec_f32) {
         ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type);
 
         int64_t OC = w->ne[3] / IC;

src/core/ggml_extend.h

@@ -153,7 +153,8 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx,
                               int d0              = 1,
                               int d1              = 1,
                               int d2              = 1,
-                              bool force_prec_f32 = false);
+                              bool force_prec_f32 = false,
+                              bool direct         = false);
 
 // w: [OC，IC, KD, 1 * 1]
 // x: [N, IC, ID, IH*IW]

src/core/ggml_runner.cpp

@@ -530,6 +530,7 @@ GGMLRunnerContext GGMLRunner::get_context() {
     runner_ctx.linear_scale          = linear_scale;
     runner_ctx.attn_scale            = attn_scale;
     runner_ctx.conv2d_direct_enabled = conv2d_direct_enabled;
+    runner_ctx.conv3d_direct_enabled = conv3d_direct_enabled;
     runner_ctx.circular_x_enabled    = circular_x_enabled;
     runner_ctx.circular_y_enabled    = circular_y_enabled;
     runner_ctx.weight_adapter        = weight_adapter;

src/core/ggml_runner.h

@@ -71,6 +71,7 @@ struct GGMLRunnerContext {
     float linear_scale                                               = 0.f;
     float attn_scale                                                 = 0.f;
     bool conv2d_direct_enabled                                       = false;
+    bool conv3d_direct_enabled                                       = false;
     bool circular_x_enabled                                          = false;
     bool circular_y_enabled                                          = false;
     ggml_tensor* ip_context                                          = nullptr;
@@ -178,6 +179,7 @@ protected:
     float linear_scale         = 0.f;
     float attn_scale           = 0.f;
     bool conv2d_direct_enabled = false;
+    bool conv3d_direct_enabled = false;
     bool circular_x_enabled    = false;
     bool circular_y_enabled    = false;
 
@@ -346,6 +348,10 @@ public:
         conv2d_direct_enabled = enabled;
     }
 
+    void set_conv3d_direct_enabled(bool enabled) {
+        conv3d_direct_enabled = enabled;
+    }
+
     void set_circular_axes(bool circular_x, bool circular_y) {
         circular_x_enabled = circular_x;
         circular_y_enabled = circular_y;

src/model/common/ggml_block.hpp

@@ -728,7 +728,7 @@ public:
                                 std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                 std::get<2>(padding), std::get<1>(padding), std::get<0>(padding),
                                 std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
-                                force_prec_f32);
+                                force_prec_f32, ctx->conv3d_direct_enabled);
     }
 };
 

src/model/vae/wan_vae.hpp

@@ -78,7 +78,8 @@ namespace WAN {
             return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels,
                                     std::get<2>(stride), std::get<1>(stride), std::get<0>(stride),
                                     0, 0, 0,
-                                    std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation));
+                                    std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation),
+                                    false, ctx->conv3d_direct_enabled);
         }
     };
 

src/pipeline/model_builders.cpp

@@ -584,10 +584,12 @@ namespace sd::model_builders {
         }
 
         if (sd_ctx_params->vae_conv_direct) {
-            LOG_INFO("Using Conv2d direct in the vae model");
+            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) {