Fix negative padding with input_dilation > 1 in conv_general

mlx/backend/cuda/conv.cpp

@@ -9,6 +9,7 @@
 
 #include <nvtx3/nvtx3.hpp>
 
+#include <algorithm>
 #include <cassert>
 
 namespace mlx::core {
@@ -31,6 +32,7 @@ struct ConvCacheKey {
   std::array<int, MAX_NDIM> padding_lo;
   std::array<int, MAX_NDIM> padding_hi;
   std::array<int, MAX_NDIM> dilation;
+  std::array<int, MAX_NDIM> input_dilation;
   int groups;
   bool flip;
   uint8_t input_alignment;
@@ -281,6 +283,7 @@ void Convolution::eval_gpu(const std::vector<array>& 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<array>& 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<ConvBackendType, 2> 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

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

mlx/primitives.cpp

@@ -1506,54 +1506,20 @@ std::vector<array> 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<int> conv_padding_lo = padding_lo;
-      std::vector<int> 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<int> stride = */ input_dilation_,
-          /* std::vector<int> padding_lo = */ conv_padding_lo,
-          /* std::vector<int> padding_hi = */ conv_padding_hi,
+          /* std::vector<int> padding_lo = */ padding_lo,
+          /* std::vector<int> padding_hi = */ padding_hi,
           /* std::vector<int> kernel_dilation = */ kernel_dilation_,
           /* std::vector<int> 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<int64_t>(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

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

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