diff --git a/mlx/backend/cuda/conv.cpp b/mlx/backend/cuda/conv.cpp index 9e471d235d..42bd2e8331 100644 --- a/mlx/backend/cuda/conv.cpp +++ b/mlx/backend/cuda/conv.cpp @@ -9,6 +9,7 @@ #include +#include #include namespace mlx::core { @@ -31,6 +32,7 @@ struct ConvCacheKey { std::array padding_lo; std::array padding_hi; std::array dilation; + std::array input_dilation; int groups; bool flip; uint8_t input_alignment; @@ -281,6 +283,7 @@ void Convolution::eval_gpu(const std::vector& inputs, array& out_) { cache_key.pod.padding_lo = vector_key(padding_lo_); cache_key.pod.padding_hi = vector_key(padding_hi_); cache_key.pod.dilation = vector_key(kernel_dilation_); + cache_key.pod.input_dilation = vector_key(input_dilation_); cache_key.pod.groups = groups_; cache_key.pod.flip = flip_; cache_key.pod.input_alignment = get_alignment(in); @@ -320,11 +323,20 @@ void Convolution::eval_gpu(const std::vector& inputs, array& out_) { // There is no reliable way to deduce the proper cuDNN backend for the // convolution, so we make a best guess and then try. + // The backward input convolution can not express a kernel stride and the + // other two can not express input dilation, those use the fallback kernel. + auto not_one = [](int v) { return v != 1; }; + bool has_stride = + std::any_of(kernel_strides_.begin(), kernel_strides_.end(), not_one); + bool has_input_dilation = + std::any_of(input_dilation_.begin(), input_dilation_.end(), not_one); SmallVector try_backends; if (flip_) { // When weight is flipped, we assume it is backward input convolution. - try_backends.push_back(CONV_BACKWARD_INPUT); - } else { + if (!has_stride) { + try_backends.push_back(CONV_BACKWARD_INPUT); + } + } else if (!has_input_dilation) { // Otherwise it could be backward weight convolution or forward convolution, // mathematically there is no difference so we have to use heuristics. // Empirically backward convolutions have large kernel dimensions, and diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 26e7e8caa3..97ff8ae881 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -4779,22 +4779,35 @@ array conv_general( has_neg_padding |= (pd < 0); } - // Handle negative padding + // Negative padding counts positions of the dilated input. Keep the input + // elements inside the window and pad the gap left before or after them. + bool only_padding = false; if (has_neg_padding) { Shape starts(in.ndim(), 0); auto stops = in.shape(); for (int i = 0; i < spatial_dims; i++) { - if (padding_lo[i] < 0) { - starts[i + 1] = safe_cast( - starts[i + 1] - static_cast(padding_lo[i]), "conv"); - padding_lo[i] = 0; - } - - if (padding_hi[i] < 0) { - stops[i + 1] += padding_hi[i]; - padding_hi[i] = 0; + int64_t dil = input_dilation[i]; + int64_t len = (in.shape(i + 1) - 1) * dil + 1; + int64_t lo = padding_lo[i]; + int64_t hi = padding_hi[i]; + int64_t start = lo < 0 ? (safe_cast(-lo, "conv") + dil - 1) / dil : 0; + int64_t stop = hi < 0 ? std::max(len + hi + dil - 1, 0) / dil + : in.shape(i + 1); + if (stop <= start) { + // Only zeros are left, keep the padded size for the output shape + only_padding = true; + start = stop = 0; + lo = 0; + hi = len + padding_lo[i] + hi + dil - 1; + } else { + lo += start * dil; + hi += len - (stop - 1) * dil - 1; } + starts[i + 1] = safe_cast(start, "conv"); + stops[i + 1] = safe_cast(stop, "conv"); + padding_lo[i] = safe_cast(lo, "conv"); + padding_hi[i] = safe_cast(hi, "conv"); } in = slice(in, std::move(starts), std::move(stops), s); @@ -4810,6 +4823,10 @@ array conv_general( kernel_dilation, input_dilation); + if (only_padding) { + return zeros(std::move(out_shape), in.dtype(), s); + } + return array( std::move(out_shape), in.dtype(), diff --git a/mlx/primitives.cpp b/mlx/primitives.cpp index 05939862dc..ecfc1c3eab 100644 --- a/mlx/primitives.cpp +++ b/mlx/primitives.cpp @@ -1506,54 +1506,20 @@ std::vector Convolution::vjp( padding_hi[i] = safe_cast(in_size - out_size + padding_lo_[i], "conv"); } - // Check for negative padding - bool has_neg_padding = false; - for (auto& pd : padding_lo) { - has_neg_padding |= (pd < 0); - } - for (auto& pd : padding_hi) { - has_neg_padding |= (pd < 0); - } - - // Negative padding is cropped from the output below - std::vector conv_padding_lo = padding_lo; - std::vector conv_padding_hi = padding_hi; - for (int i = 0; i < conv_padding_lo.size(); ++i) { - conv_padding_lo[i] = std::max(conv_padding_lo[i], 0); - conv_padding_hi[i] = std::max(conv_padding_hi[i], 0); - } - + // conv_general crops any negative padding from the cotangent auto wt_trans = group_transpose(wt, 0, 1, -1); auto grad = conv_general( /* const array& input = */ cotan, /* const array& weight = */ wt_trans, /* std::vector stride = */ input_dilation_, - /* std::vector padding_lo = */ conv_padding_lo, - /* std::vector padding_hi = */ conv_padding_hi, + /* std::vector padding_lo = */ padding_lo, + /* std::vector padding_hi = */ padding_hi, /* std::vector kernel_dilation = */ kernel_dilation_, /* std::vector input_dilation = */ kernel_strides_, /* int groups = */ groups_, /* bool flip = */ !flip_, stream()); - // Handle negative padding - if (has_neg_padding) { - Shape starts(grad.ndim(), 0); - auto stops = grad.shape(); - - for (int i = 0; i < grad.ndim() - 2; i++) { - if (padding_lo[i] < 0) { - starts[i + 1] = safe_cast( - starts[i + 1] - static_cast(padding_lo[i]), "conv"); - } - if (padding_hi[i] < 0) { - stops[i + 1] += padding_hi[i]; - } - } - - grad = slice(grad, std::move(starts), std::move(stops), stream()); - } - grads.push_back(grad); } // Grads for weight diff --git a/python/tests/test_conv.py b/python/tests/test_conv.py index 98210dd396..cc51765379 100644 --- a/python/tests/test_conv.py +++ b/python/tests/test_conv.py @@ -1219,6 +1219,58 @@ def padded_loss(x, w, stride, padding): self.assertEqual(grad.shape, x.shape) self.assertTrue(mx.allclose(grad, expected, atol=1e-4, rtol=1e-4)) + def test_input_dilation_grad(self): + def dilate(x, dilation): + shape = [(n - 1) * d + 1 for n, d in zip(x.shape[1:-1], dilation)] + idx = (slice(None),) + tuple(slice(None, None, d) for d in dilation) + return mx.zeros([x.shape[0], *shape, x.shape[-1]]).at[idx].add(x) + + def conv_loss(x, w, stride, padding, dilation): + out = mx.conv_general( + x, w, stride=stride, padding=padding, input_dilation=dilation + ) + return mx.sum(out * out) + + def dilated_loss(x, w, stride, padding, dilation): + out = mx.conv_general( + dilate(x, dilation), w, stride=stride, padding=padding + ) + return mx.sum(out * out) + + for in_shape, w_shape, stride, padding, dilation in [ + ((1, 5, 1), (1, 1, 1), (1,), ([0], [1]), (2,)), + ((1, 5, 1), (1, 1, 1), (1,), ([1], [0]), (2,)), + ((1, 5, 1), (1, 2, 1), (2,), ([0], [0]), (2,)), + ((1, 5, 1), (1, 3, 1), (2,), ([1], [0]), (3,)), + ((1, 6, 2), (3, 2, 2), (1,), ([4], [3]), (3,)), + ((1, 3, 1), (1, 1, 1), (2,), ([2], [1]), (2,)), + ((1, 3, 1), (1, 3, 1), (3,), ([0], [0]), (2,)), + ((2, 5, 4, 3), (4, 2, 3, 3), (2, 1), ([1, 0], [0, 2]), (2, 3)), + ( + (1, 3, 4, 3, 2), + (2, 2, 2, 2, 2), + (1, 2, 1), + ([0, 1, 0], [1, 0, 2]), + (2, 2, 3), + ), + ]: + x = mx.random.normal(in_shape) + w = mx.random.normal(w_shape) + args = (stride, padding, dilation) + out = mx.conv_general( + x, w, stride=stride, padding=padding, input_dilation=dilation + ) + expected = mx.conv_general( + dilate(x, dilation), w, stride=stride, padding=padding + ) + self.assertEqual(out.shape, expected.shape) + self.assertTrue(mx.allclose(out, expected, atol=1e-4, rtol=1e-4)) + expected = mx.grad(dilated_loss, argnums=(0, 1))(x, w, *args) + grads = mx.grad(conv_loss, argnums=(0, 1))(x, w, *args) + for grad, exp in zip(grads, expected): + self.assertEqual(grad.shape, exp.shape) + self.assertTrue(mx.allclose(grad, exp, atol=1e-4, rtol=1e-4)) + @unittest.skipIf(mx.cuda.is_available() and "CI" in os.environ, "flaky in CI") def test_conv_1d_with_2d(self): x = mx.random.uniform(shape=(2, 10, 16)) diff --git a/python/tests/test_conv_transpose.py b/python/tests/test_conv_transpose.py index 60038df109..8c7567f4d5 100644 --- a/python/tests/test_conv_transpose.py +++ b/python/tests/test_conv_transpose.py @@ -88,6 +88,9 @@ def run_conv_transpose_1D( (1, 1, 1, 0), (3, 3, 1, 0), (31, 5, 5, 2), + (5, 2, 2, 2), + (6, 3, 2, 3), + (2, 1, 4, 1), ): run_conv_transpose_1D( N, C, OC, iH, kH, stride, padding, dtype=dtype @@ -221,6 +224,8 @@ def f(a, b): (1, 1, 1, 0), (3, 3, 1, 0), (31, 5, 5, 2), + (5, 2, 2, 2), + (6, 3, 2, 3), ): run_conv_transpose1D_grad( N, C, OC, iH, kH, stride, padding, dtype=dtype @@ -659,6 +664,7 @@ def run_conv_transpose_1d_output_padding( (3, 2, 2, 0, 1), (5, 3, 2, 1, 0), (7, 4, 3, 1, 2), + (3, 2, 3, 3, 1), ): run_conv_transpose_1d_output_padding( N, C, OC, iH, kH, stride, padding, output_padding, dtype=dtype