diff --git a/mlx/backend/cpu/conv.cpp b/mlx/backend/cpu/conv.cpp index 70b5f270f0..17bc6cb9ff 100644 --- a/mlx/backend/cpu/conv.cpp +++ b/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(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(iDim[i]) + padding_lo[i] + padding_hi[i], "conv"); } padded_shape.back() = C; array in_padded(padded_shape, conv_dtype, nullptr, {}); diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 5c32b3d963..fc95feff2a 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -898,15 +898,20 @@ void winograd_conv_2D_gpu( array& out, const MLXConvParams<2>& conv_params, std::vector& 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(conv_params.iS[0]) + + 2 * static_cast(conv_params.pad[0]); + int64_t pad_w = static_cast(conv_params.iS[1]) + + 2 * static_cast(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 diff --git a/mlx/ops.cpp b/mlx/ops.cpp index d154bd3d19..1b7c0a21b4 100644 --- a/mlx/ops.cpp +++ b/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(out_shape[ax]) + low_pad_size[i] + + high_pad_size[i], + "pad"); } if (mode == "constant") { diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 3da0a2950b..09236f1da1 100644 --- a/tests/ops_tests.cpp +++ b/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);