diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index ba087d3c..ebbed10b 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -1038,6 +1038,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_linear(ggml_context* ctx, } __STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, + ggml_backend_t backend, ggml_tensor* x, int lp0, int rp0, @@ -1063,7 +1064,17 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, } if (lp0 != 0 || rp0 != 0 || lp1 != 0 || rp1 != 0 || lp2 != 0 || rp2 != 0 || lp3 != 0 || rp3 != 0) { - x = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); + ggml_tensor* padded = ggml_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); + if (backend == nullptr || ggml_backend_supports_op(backend, padded)) { + x = padded; + } else { + // Some backends (e.g. Metal) only implement right-padding for + // GGML_OP_PAD (see #850): pad right by lp+rp instead, then roll + // the padding around to the left. shift < ne always holds because + // ne grew by lp+rp. + x = ggml_pad_ext(ctx, x, 0, lp0 + rp0, 0, lp1 + rp1, 0, lp2 + rp2, 0, lp3 + rp3); + x = ggml_roll(ctx, x, lp0, lp1, lp2, lp3); + } } return x; } @@ -1076,7 +1087,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_pad(ggml_context* ctx, int p3 = 0, bool circular_x = false, bool circular_y = false) { - return ggml_ext_pad_ext(ctx, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y); + return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y); } // w: [OC,IC, KH, KW] @@ -1105,7 +1116,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, } if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) { - x = ggml_ext_pad_ext(ctx, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y); + x = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y); p0 = 0; p1 = 0; } @@ -1130,6 +1141,7 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, // b: [OC,] // result: [N*OC, OD, OH, OW] __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, + ggml_backend_t backend, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, @@ -1159,7 +1171,21 @@ __STATIC_INLINE__ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2)); x = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N); } else { - x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2); + // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends + // (e.g. Metal, see #850) do not implement. Fall back to + // GGML_OP_CONV_3D on those backends. + bool im2col_3d_supported = true; + if (backend != nullptr) { + ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type); + im2col_3d_supported = ggml_backend_supports_op(backend, im2col); + } + if (im2col_3d_supported) { + x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2); + } else { + 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); + } } if (b != nullptr) { @@ -3368,7 +3394,7 @@ public: b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); } } - return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels, + 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), std::get<2>(padding), std::get<1>(padding), std::get<0>(padding), std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), diff --git a/src/model/vae/ltx_audio_vae.hpp b/src/model/vae/ltx_audio_vae.hpp index 2f8c03d1..71070d6a 100644 --- a/src/model/vae/ltx_audio_vae.hpp +++ b/src/model/vae/ltx_audio_vae.hpp @@ -214,7 +214,7 @@ namespace LTXV { auto x = ggml_reshape_3d(ctx, waveform, time, 1, channels * batch); if (left_pad > 0) { - x = ggml_pad_ext(ctx, x, static_cast(left_pad), 0, 0, 0, 0, 0, 0, 0); + x = ggml_ext_pad_ext(ctx, runner_ctx->backend, x, static_cast(left_pad), 0, 0, 0, 0, 0, 0, 0); } auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1); @@ -451,6 +451,7 @@ namespace LTXV { int pad_h = kernel_size.first - 1; int pad_w = kernel_size.second - 1; x = ggml_ext_pad_ext(ctx->ggml_ctx, + ctx->backend, x, pad_w / 2, pad_w - pad_w / 2, diff --git a/src/model/vae/tae.hpp b/src/model/vae/tae.hpp index a78e5e96..558c0544 100644 --- a/src/model/vae/tae.hpp +++ b/src/model/vae/tae.hpp @@ -408,7 +408,7 @@ public: h = conv->forward(ctx, h); for (int j = 0; j < num_blocks; j++) { auto block = std::dynamic_pointer_cast(blocks[std::to_string(index++)]); - auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0); + auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0); mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0); h = block->forward(ctx, h, mem); } @@ -479,7 +479,7 @@ public: int index = 3; for (int i = 0; i < num_layers; i++) { for (int j = 0; j < num_blocks; j++) { - auto mem = ggml_pad_ext(ctx->ggml_ctx, h, 0, 0, 0, 0, 0, 0, 1, 0); + auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0); mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0); if (is_wide) { auto block = std::dynamic_pointer_cast(blocks[std::to_string(index++)]); diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index d28b5b25..94956209 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -72,8 +72,8 @@ namespace WAN { lp2 -= (int)cache_x->ne[2]; } - x = ggml_ext_pad_ext(ctx->ggml_ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); - return ggml_ext_conv_3d(ctx->ggml_ctx, x, w, b, in_channels, + x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); + 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)); @@ -195,7 +195,7 @@ namespace WAN { 2); } if (chunk_idx == 1 && cache_x->ne[2] < 2) { // Rep - cache_x = ggml_pad_ext(ctx->ggml_ctx, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0); + cache_x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0); // aka cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device),cache_x],dim=2) } if (chunk_idx == 1) { @@ -283,7 +283,7 @@ namespace WAN { int pad_t = (factor_t - T % factor_t) % factor_t; - x = ggml_pad_ext(ctx->ggml_ctx, x, 0, 0, 0, 0, pad_t, 0, 0, 0); + x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, 0, 0, 0, 0, pad_t, 0, 0, 0); T = x->ne[2]; x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, factor_t, T / factor_t, C); // [C, T/factor_t, factor_t, H*W]