From 3db19dc3c2921292879573b9b12a7113989b5235 Mon Sep 17 00:00:00 2001 From: Vladimir Iglovikov Date: Tue, 18 Aug 2026 18:51:38 +0300 Subject: [PATCH] [Metal] Optimize stride-2 ConvTranspose3d --- mlx/backend/metal/conv.cpp | 114 ++++++++++++++++++++++++++++ python/tests/test_conv_transpose.py | 2 + 2 files changed, 116 insertions(+) diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 85e137bf4f..8600f069a5 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -842,6 +842,115 @@ void small_kd_conv_3D_gpu( 0); } +// A stride-2, kernel-2 transposed convolution has exactly one valid kernel +// phase for every output coordinate. The explicit unfold path nevertheless +// materializes all eight phases and fills seven of them with zeros. Compute +// each phase as a small regular GEMM instead. This path is intentionally +// narrow: other transposed-convolution configurations retain the general +// implementation below. +bool is_stride_two_conv_transpose_3D(const MLXConvParams<3>& p) { + return p.groups == 1 && p.flip && p.str[0] == 1 && p.str[1] == 1 && + p.str[2] == 1 && p.idil[0] == 2 && p.idil[1] == 2 && p.idil[2] == 2 && + p.kdil[0] == 1 && p.kdil[1] == 1 && p.kdil[2] == 1 && p.wS[0] == 2 && + p.wS[1] == 2 && p.wS[2] == 2 && p.pad[0] == 1 && p.pad[1] == 1 && + p.pad[2] == 1 && static_cast(p.oS[0]) == 2LL * p.iS[0] && + static_cast(p.oS[1]) == 2LL * p.iS[1] && + static_cast(p.oS[2]) == 2LL * p.iS[2]; +} + +void stride_two_conv_transpose_3D_gpu( + const Stream& s, + metal::Device& d, + const array& in, + const array& wt, + array& out, + const MLXConvParams<3>& p, + std::vector& copies) { + constexpr int kernel_volume = 8; + const int C = p.C; + const int O = p.O; + + // The input and weight are contiguous by the time this helper is called. + // Every phase covers the complete input volume; the phase bit only selects + // the interleaved output coordinates and the corresponding weight slice. + for (int phase = 0; phase < kernel_volume; ++phase) { + const int pd = (phase >> 2) & 1; + const int ph = (phase >> 1) & 1; + const int pw = phase & 1; + const int D = p.iS[0]; + const int H = p.iS[1]; + const int W = p.iS[2]; + + const int M = safe_cast(static_cast(p.N) * D * H * W, "conv"); + + // The weight is [O, 2, 2, 2, C]. Present one spatial phase as a [C, O] + // matrix with the layout expected by a transposed Steel GEMM. + array wt_phase({C, O}, wt.dtype(), nullptr, {}); + array::Flags wt_flags = wt.flags(); + wt_flags.contiguous = false; + wt_flags.row_contiguous = false; + wt_flags.col_contiguous = true; + wt_phase.copy_shared_buffer( + wt, + {1, wt.strides(0)}, + wt_flags, + wt.data_size(), + static_cast(pd * 4 + ph * 2 + pw) * C); + + array in_matrix({M, C}, in.dtype(), nullptr, {}); + in_matrix.copy_shared_buffer(in, {C, 1}, in.flags(), in.data_size()); + + array phase_out({M, O}, out.dtype(), nullptr, {}); + phase_out.set_data(allocator::malloc(phase_out.nbytes())); + + std::vector gemm_copies = {in_matrix, wt_phase}; + steel_matmul( + s, + d, + /* a = */ in_matrix, + /* b = */ wt_phase, + /* out = */ phase_out, + /* M = */ M, + /* N = */ O, + /* K = */ C, + /* batch_size_out = */ 1, + /* lda = */ C, + /* ldb = */ kernel_volume * C, + /* a_transposed = */ false, + /* b_transposed = */ true, + /* copies = */ gemm_copies); + + Shape phase_out_shape{p.N, D, H, W, O}; + array phase_out_nd(phase_out_shape, out.dtype(), nullptr, {}); + phase_out_nd.copy_shared_buffer( + phase_out, + make_contiguous_strides(phase_out_shape), + phase_out.flags(), + phase_out.data_size()); + + Strides out_phase_strides = out.strides(); + out_phase_strides[1] *= 2; + out_phase_strides[2] *= 2; + out_phase_strides[3] *= 2; + array::Flags out_phase_flags = out.flags(); + out_phase_flags.contiguous = false; + out_phase_flags.row_contiguous = false; + out_phase_flags.col_contiguous = false; + array out_phase_view(phase_out_shape, out.dtype(), nullptr, {}); + out_phase_view.copy_shared_buffer( + out, + out_phase_strides, + out_phase_flags, + out.data_size(), + pd * out.strides(1) + ph * out.strides(2) + pw * out.strides(3)); + + copy_gpu_inplace(phase_out_nd, out_phase_view, CopyType::GeneralGeneral, s); + copies.push_back(phase_out); + copies.push_back(phase_out_nd); + copies.push_back(out_phase_view); + } +} + void dispatch_conv_3D_gpu( const Stream& s, metal::Device& d, @@ -872,6 +981,11 @@ void dispatch_conv_3D_gpu( auto in = ensure_row_contiguous(in_pre, d, s); auto wt = ensure_row_contiguous(wt_pre, d, s); + if (is_stride_two_conv_transpose_3D(conv_params)) { + return stride_two_conv_transpose_3D_gpu( + s, d, in, wt, out, conv_params, copies); + } + // Decompose 3D conv to per-frame 2D convs constexpr int kSmallKdLimit3D = 7; if (is_idil_one && mod16_channels && conv_params.groups == 1 && diff --git a/python/tests/test_conv_transpose.py b/python/tests/test_conv_transpose.py index e6def081a7..7d12cb16f4 100644 --- a/python/tests/test_conv_transpose.py +++ b/python/tests/test_conv_transpose.py @@ -486,6 +486,8 @@ def run_conv_transpose3D( ((1, 1, 1), (1, 1, 1), (1, 1, 1), (0, 0, 0)), ((3, 3, 3), (3, 1, 1), (1, 1, 1), (0, 0, 0)), ((15, 15, 15), (3, 3, 3), (3, 3, 3), (2, 2, 2)), + # Exercises the Metal phase-aware stride-2/kernel-2 path. + ((3, 4, 5), (2, 2, 2), (2, 2, 2), (0, 0, 0)), ): run_conv_transpose3D( N, C, O, idim, kdim, stride, padding, dtype=dtype