Fix int32 overflow in conv padded input and pad shapes

mlx/backend/cpu/conv.cpp

@@ -814,7 +814,11 @@ void explicit_gemm_conv_1D_cpu(
   auto& encoder = cpu::get_command_encoder(stream);
 
   // Pad input
-  Shape padded_shape = {N, iH + padding_lo[0] + padding_hi[0], C};
+  Shape padded_shape = {
+      N,
+      safe_cast(
+          static_cast<int64_t>(iH) + padding_lo[0] + padding_hi[0], "conv"),
+      C};
   array in_padded(padded_shape, conv_dtype, nullptr, {});
 
   // Fill with zeros
@@ -961,7 +965,8 @@ void explicit_gemm_conv_ND_cpu(
   Shape padded_shape(in.shape().size());
   padded_shape.front() = N;
   for (size_t i = 0; i < iDim.size(); i++) {
-    padded_shape[i + 1] = iDim[i] + padding_lo[i] + padding_hi[i];
+    padded_shape[i + 1] = safe_cast(
+        static_cast<int64_t>(iDim[i]) + padding_lo[i] + padding_hi[i], "conv");
   }
   padded_shape.back() = C;
   array in_padded(padded_shape, conv_dtype, nullptr, {});

mlx/backend/metal/conv.cpp

@@ -898,15 +898,20 @@ void winograd_conv_2D_gpu(
     array& out,
     const MLXConvParams<2>& conv_params,
     std::vector<array>& copies_w) {
+  // Round the padded spatial dims up to the Winograd tile in int64 so the
+  // rounding cannot overflow int32 just below the limit.
+  int64_t pad_h = static_cast<int64_t>(conv_params.iS[0]) +
+      2 * static_cast<int64_t>(conv_params.pad[0]);
+  int64_t pad_w = static_cast<int64_t>(conv_params.iS[1]) +
+      2 * static_cast<int64_t>(conv_params.pad[1]);
+  pad_h = 6 * ((pad_h - 2 + 5) / 6) + 2;
+  pad_w = 6 * ((pad_w - 2 + 5) / 6) + 2;
   Shape padded_shape = {
       conv_params.N,
-      conv_params.iS[0] + 2 * conv_params.pad[0],
-      conv_params.iS[1] + 2 * conv_params.pad[1],
+      safe_cast(pad_h, "conv"),
+      safe_cast(pad_w, "conv"),
       conv_params.C};
 
-  padded_shape[1] = 6 * ((padded_shape[1] - 2 + 5) / 6) + 2;
-  padded_shape[2] = 6 * ((padded_shape[2] - 2 + 5) / 6) + 2;
-
   array in_padded(std::move(padded_shape), in.dtype(), nullptr, {});
 
   // Fill with zeros

mlx/ops.cpp

@@ -1633,7 +1633,10 @@ array pad(
     }
 
     auto ax = axes[i] < 0 ? a.ndim() + axes[i] : axes[i];
-    out_shape[ax] += low_pad_size[i] + high_pad_size[i];
+    out_shape[ax] = safe_cast(
+        static_cast<int64_t>(out_shape[ax]) + low_pad_size[i] +
+            high_pad_size[i],
+        "pad");
   }
 
   if (mode == "constant") {

tests/ops_tests.cpp

@@ -4491,6 +4491,14 @@ TEST_CASE("test conv shape overflow") {
       Shape{1, 8, 8, 1});
 }
 
+TEST_CASE("test pad shape overflow") {
+  // A padding sum that overflows int32 is rejected, not wrapped.
+  // https://github.com/ml-explore/mlx/issues/3611
+  const int imax = 2147483647;
+  CHECK_THROWS_AS(
+      pad(zeros({8}), {0}, Shape{imax}, Shape{imax}), std::overflow_error);
+}
+
 TEST_CASE("test fp8 conversion") {
   for (auto t : {float32, float16, bfloat16}) {
     array in({-1.125, -1.0, 0.0, 1.0, 1.125, 4.5, 448.0}, t);