diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 0fe6b69615..25ab123731 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -5,6 +5,7 @@ #include "mlx/backend/gpu/copy.h" #include "mlx/backend/gpu/slicing.h" +#include "mlx/backend/metal/binary.h" #include "mlx/backend/metal/device.h" #include "mlx/backend/metal/kernels.h" #include "mlx/backend/metal/kernels/defines.h" @@ -676,6 +677,124 @@ void pad_and_slice_conv_3D_gpu( intermediate, intermediate.strides(), {0}, intermediate.data_size()); } +// Forward declaration; conv_2D_gpu is defined later in this file. +void conv_2D_gpu( + const Stream& s, + metal::Device& d, + const array& in_pre, + const array& wt_pre, + array& out, + const std::vector& padding, + const std::vector& wt_strides, + const std::vector& wt_dilation, + const std::vector& in_dilation, + const int groups, + bool flip, + std::vector& copies); + +// Small kernel-depth 3D conv via per-depth-tap 2D convs (#3625). +// +// The 3D implicit-gemm path has no Winograd / 3x3-specialized kernel, so a +// (small KD) 3D conv is 2-5x slower than decomposing it into KD 2D convs, each +// of which hits the tuned 2D dispatch (Winograd for 3x3 stride-1). For each +// depth tap kd we run a 2D conv over the OD output frames (a zero-copy strided +// view of the input at depth offset kd) with the weight's depth slice, and +// accumulate. +// +// Preconditions (enforced by the dispatch guard): input dilation 1, depth +// stride and depth kernel-dilation 1, no depth padding, groups == 1, N == 1, +// mod16 channels, and KD small. Other cases fall through to the implicit gemm. +void small_kd_conv_3D_gpu( + const Stream& s, + metal::Device& d, + const array& in, + const array& wt, + array& out, + const MLXConvParams<3>& conv_params, + std::vector& copies) { + const int H = conv_params.iS[1]; + const int W = conv_params.iS[2]; + const int C = conv_params.C; + const int O = conv_params.O; + const int KD = conv_params.wS[0]; + const int KH = conv_params.wS[1]; + const int KW = conv_params.wS[2]; + const int OD = conv_params.oS[0]; + const int OH = conv_params.oS[1]; + const int OW = conv_params.oS[2]; + + // Accumulate the KD per-depth-tap 2D convs into `acc`, then repoint `out`. + array acc({OD, OH, OW, O}, out.dtype(), nullptr, {}); + for (int kd = 0; kd < KD; ++kd) { + // Input frames for this depth tap: [OD, H, W, C] view of `in` at depth kd. + // With depth stride/dilation 1 and no depth padding these OD frames are the + // contiguous slab in[kd : kd + OD], so a plain strided view suffices. + array in_2d({OD, H, W, C}, in.dtype(), nullptr, {}); + in_2d.copy_shared_buffer( + in, + {static_cast(H) * W * C, + static_cast(W) * C, + static_cast(C), + 1}, + {true, true, false}, + static_cast(OD) * H * W * C, + static_cast(kd) * H * W * C); + + // Weight depth slice wt[:, kd] -> [O, KH, KW, C], strided over O. + array wt_2d({O, KH, KW, C}, wt.dtype(), nullptr, {}); + wt_2d.copy_shared_buffer( + wt, + {static_cast(KD) * KH * KW * C, + static_cast(KW) * C, + static_cast(C), + 1}, + {false, false, false}, + static_cast(O - 1) * KD * KH * KW * C + + static_cast(KH) * KW * C, + static_cast(kd) * KH * KW * C); + + // 2D conv into a fresh output; conv_2D_gpu allocates and dispatches + // (Winograd etc.). Spatial params come from dims 1,2 of the 3D params. + array conv_out({OD, OH, OW, O}, out.dtype(), nullptr, {}); + conv_2D_gpu( + s, + d, + in_2d, + wt_2d, + conv_out, + {conv_params.pad[1], conv_params.pad[2]}, + {conv_params.str[1], conv_params.str[2]}, + {conv_params.kdil[1], conv_params.kdil[2]}, + {conv_params.idil[1], conv_params.idil[2]}, + /* groups = */ 1, + /* flip = */ conv_params.flip, + copies); + + if (kd == 0) { + // First tap owns the accumulator buffer. + acc = conv_out; + } else { + // Elementwise in-place accumulate (safe: add is per-element). + binary_op_gpu_inplace({acc, conv_out}, acc, "Add", s); + copies.push_back(conv_out); + } + } + + // Repoint the [1, OD, OH, OW, O] output at the contiguous [OD, OH, OW, O] + // accumulator buffer (same element count). No temporary needed for `acc` + // since `out` shares its buffer. + out.copy_shared_buffer( + acc, + {static_cast(OD) * OH * OW * O, + static_cast(OH) * OW * O, + static_cast(OW) * O, + static_cast(O), + 1}, + {true, true, false}, + static_cast(OD) * OH * OW * O, + 0); +} + void dispatch_conv_3D_gpu( const Stream& s, metal::Device& d, @@ -706,6 +825,19 @@ void dispatch_conv_3D_gpu( auto in = ensure_row_contiguous(in_pre, d, s); auto wt = ensure_row_contiguous(wt_pre, d, s); + // Small kernel-depth 3D conv: decompose into KD 2D convs, which hit the tuned + // 2D path (Winograd for 3x3 stride-1) that the 3D implicit gemm lacks. Only + // valid for depth stride/dilation 1, no depth padding, groups == 1, N == 1. + // (#3625) + constexpr int kSmallKdLimit3D = 7; + bool small_kd_ok = is_idil_one && mod16_channels && conv_params.groups == 1 && + conv_params.N == 1 && conv_params.wS[0] <= kSmallKdLimit3D && + conv_params.str[0] == 1 && conv_params.kdil[0] == 1 && + conv_params.pad[0] == 0; + if (small_kd_ok) { + return small_kd_conv_3D_gpu(s, d, in, wt, out, conv_params, copies); + } + // Perform the implicit gemm if (is_idil_one && mod16_channels) { return implicit_gemm_conv_3D_gpu(s, d, in, wt, out, conv_params); diff --git a/python/tests/test_conv.py b/python/tests/test_conv.py index f05fa2aaa9..dbbb2a023d 100644 --- a/python/tests/test_conv.py +++ b/python/tests/test_conv.py @@ -14,7 +14,7 @@ import torch.nn.functional as F has_torch = True -except ImportError as e: +except ImportError: has_torch = False @@ -309,9 +309,11 @@ def run_conv2D( lambda x: mx.array(x).astype(mx_dtype), (in_np, wt_np) ) in_pt, wt_pt = map( - lambda x: torch.from_numpy(x.transpose(0, 3, 1, 2)) - .to("cpu") - .to(torch_dtype), + lambda x: ( + torch.from_numpy(x.transpose(0, 3, 1, 2)) + .to("cpu") + .to(torch_dtype) + ), (in_np, wt_np), ) @@ -1054,7 +1056,6 @@ def test_repeated_conv(self): @unittest.skipIf(not has_torch, "requires Torch") def test_torch_conv_depthwise(self): - # fmt: off shapes = ( # N, H, W, C kH, kW, O, strides, padding, groups @@ -1194,6 +1195,47 @@ def test_conv2d_large_filter_small_channels(self): y_hat = mx.conv2d(x, w, (1, 1), (1, 1)) self.assertTrue(mx.allclose(y, y_hat, rtol=1e-3, atol=1e-3)) + def test_conv_3D_small_kd_decomposition(self): + # Exercises the small kernel-depth 3D -> KD x 2D decomposition (#3625): + # N=1, small KD, depth stride/dilation 1, no depth padding, mod16 channels. + # Validated against the CPU reference, which uses a different code path. + for T, H, W, Cin, Cout, kd, kh, kw in [ + (5, 16, 16, 32, 32, 3, 3, 3), # canonical 3x3x3 (2D hits Winograd) + (4, 12, 10, 16, 48, 3, 3, 3), # Cout != Cin + (6, 14, 14, 32, 32, 1, 3, 3), # KD = 1 + (5, 12, 12, 16, 16, 5, 1, 1), # larger KD, 1x1 spatial + (4, 10, 10, 32, 16, 2, 3, 3), # KD = 2 + ]: + x = mx.random.normal((1, T, H, W, Cin)) + w = mx.random.normal((Cout, kd, kh, kw, Cin)) + y_gpu = mx.conv_general(x, w, stride=(1, 1, 1)) + y_cpu = mx.conv_general(x, w, stride=(1, 1, 1), stream=mx.cpu) + mx.eval(y_gpu, y_cpu) + self.assertTrue( + mx.allclose(y_gpu, y_cpu, rtol=1e-4, atol=1e-4), + f"3D small-kd mismatch T{T} H{H} W{W} C{Cin}->{Cout} k{kd}{kh}{kw}", + ) + + def test_conv_3D_small_kd_fallback_cases(self): + # Cases that must NOT take the fast path (guarded out) but stay correct. + for kwargs in [ + dict(stride=(2, 1, 1)), # depth stride > 1 + dict(stride=(1, 1, 1), padding=(1, 0, 0)), # depth padding + ]: + x = mx.random.normal((1, 6, 12, 12, 32)) + w = mx.random.normal((32, 3, 3, 3, 32)) + y_gpu = mx.conv_general(x, w, **kwargs) + y_cpu = mx.conv_general(x, w, stream=mx.cpu, **kwargs) + mx.eval(y_gpu, y_cpu) + self.assertTrue(mx.allclose(y_gpu, y_cpu, rtol=1e-4, atol=1e-4)) + # non-mod16 channels (falls to pad-and-slice / implicit gemm) + x = mx.random.normal((1, 5, 12, 12, 24)) + w = mx.random.normal((24, 3, 3, 3, 24)) + y_gpu = mx.conv_general(x, w, stride=(1, 1, 1)) + y_cpu = mx.conv_general(x, w, stride=(1, 1, 1), stream=mx.cpu) + mx.eval(y_gpu, y_cpu) + self.assertTrue(mx.allclose(y_gpu, y_cpu, rtol=1e-4, atol=1e-4)) + if __name__ == "__main__": mlx_tests.MLXTestRunner()