Decompose small kernel-depth 3D convs into 2D convs (#3625)

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<int>& padding,
+    const std::vector<int>& wt_strides,
+    const std::vector<int>& wt_dilation,
+    const std::vector<int>& in_dilation,
+    const int groups,
+    bool flip,
+    std::vector<array>& 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<array>& 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<int64_t>(H) * W * C,
+         static_cast<int64_t>(W) * C,
+         static_cast<int64_t>(C),
+         1},
+        {true, true, false},
+        static_cast<size_t>(OD) * H * W * C,
+        static_cast<int64_t>(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<int64_t>(KD) * KH * KW * C,
+         static_cast<int64_t>(KW) * C,
+         static_cast<int64_t>(C),
+         1},
+        {false, false, false},
+        static_cast<size_t>(O - 1) * KD * KH * KW * C +
+            static_cast<size_t>(KH) * KW * C,
+        static_cast<int64_t>(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<int64_t>(OD) * OH * OW * O,
+       static_cast<int64_t>(OH) * OW * O,
+       static_cast<int64_t>(OW) * O,
+       static_cast<int64_t>(O),
+       1},
+      {true, true, false},
+      static_cast<size_t>(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);

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()