From 6d0f9b0195fc7b83ff77f8d352a3dd5179c44015 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ey=C3=BCp=20Can=20Akman?= Date: Fri, 14 Aug 2026 21:32:30 +0300 Subject: [PATCH 1/2] Fix int32 overflow in conv padded input and pad shapes --- mlx/backend/cpu/conv.cpp | 9 ++++++-- mlx/backend/metal/conv.cpp | 15 +++++++++----- mlx/ops.cpp | 5 ++++- tests/ops_tests.cpp | 42 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 63 insertions(+), 8 deletions(-) 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..3ad80b0f38 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -4491,6 +4491,48 @@ 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 conv padded input overflow") { + // A large stride can keep the output shape in range while the padded input + // (in + pad_lo + pad_hi) the backend allocates overflows int32; the backend + // must reject it rather than wrap the buffer size. + // https://github.com/ml-explore/mlx/issues/3611 + const int imax = 2147483647; + // 1D path (explicit_gemm_conv_1D_cpu). + std::vector stride = {imax}, pad_lo = {imax}, pad_hi = {imax}, + dilation = {1}; + CHECK_THROWS_AS( + eval(conv_general( + zeros({1, 8, 1}), + zeros({1, 3, 1}), + stride, + pad_lo, + pad_hi, + dilation, + dilation)), + std::overflow_error); + // ND path (explicit_gemm_conv_ND_cpu), overflow on one spatial axis only. + std::vector stride2 = {imax, 1}, pad_lo2 = {imax, 0}, + pad_hi2 = {imax, 0}, dilation2 = {1, 1}; + CHECK_THROWS_AS( + eval(conv_general( + zeros({1, 8, 8, 1}), + zeros({1, 3, 3, 1}), + stride2, + pad_lo2, + pad_hi2, + dilation2, + dilation2)), + 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); From a9a7bc1a1f469b65ce66f64c6007669616c42667 Mon Sep 17 00:00:00 2001 From: Cheng Date: Sun, 16 Aug 2026 13:42:23 +0900 Subject: [PATCH 2/2] nit --- tests/ops_tests.cpp | 34 ---------------------------------- 1 file changed, 34 deletions(-) diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 3ad80b0f38..09236f1da1 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -4499,40 +4499,6 @@ TEST_CASE("test pad shape overflow") { pad(zeros({8}), {0}, Shape{imax}, Shape{imax}), std::overflow_error); } -TEST_CASE("test conv padded input overflow") { - // A large stride can keep the output shape in range while the padded input - // (in + pad_lo + pad_hi) the backend allocates overflows int32; the backend - // must reject it rather than wrap the buffer size. - // https://github.com/ml-explore/mlx/issues/3611 - const int imax = 2147483647; - // 1D path (explicit_gemm_conv_1D_cpu). - std::vector stride = {imax}, pad_lo = {imax}, pad_hi = {imax}, - dilation = {1}; - CHECK_THROWS_AS( - eval(conv_general( - zeros({1, 8, 1}), - zeros({1, 3, 1}), - stride, - pad_lo, - pad_hi, - dilation, - dilation)), - std::overflow_error); - // ND path (explicit_gemm_conv_ND_cpu), overflow on one spatial axis only. - std::vector stride2 = {imax, 1}, pad_lo2 = {imax, 0}, - pad_hi2 = {imax, 0}, dilation2 = {1, 1}; - CHECK_THROWS_AS( - eval(conv_general( - zeros({1, 8, 8, 1}), - zeros({1, 3, 3, 1}), - stride2, - pad_lo2, - pad_hi2, - dilation2, - dilation2)), - 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);