Skip to content

Fix int32 overflow in conv padded input and pad shapes - #4258

Open
eyupcanakman wants to merge 2 commits into
ml-explore:mainfrom
eyupcanakman:fix/pad-conv-shape-overflow-3611
Open

Fix int32 overflow in conv padded input and pad shapes#4258
eyupcanakman wants to merge 2 commits into
ml-explore:mainfrom
eyupcanakman:fix/pad-conv-shape-overflow-3611

Conversation

@eyupcanakman

Copy link
Copy Markdown
Contributor

Follow-up to #3611 and #3938. #3938 made conv_out_shape and the backward pass compute in int64_t, but it guards the output shape. The backends build a padded input buffer of in + pad_lo + pad_hi in 32 bits, which can overflow even when the output stays in range, because a large stride shrinks the output while the padded input does not.

mx.conv_general(mx.ones((1, 8, 1)), mx.ones((1, 3, 1)), stride=3, padding=([2**31-1], [2**31-1])) returns a valid (1, 1431655767, 1) output, then wraps the padded buffer and writes out of bounds on eval. mx.pad computes the same 32-bit sum in its own shape.

The CPU explicit_gemm paths (1D and ND), the Metal Winograd path, and pad now compute these sums in int64_t and narrow through safe_cast, the same way #3938 did for the output shape. In-range convolutions and pads are unchanged.

Inputs that used to wrap now raise std::overflow_error, so Python raises OverflowError.

The Winograd change is hardening for the same class rather than a live fix. With stride one, a padding large enough to overflow the padded input also forces an output near INT32_MAX, so the output buffer allocation raises first. It is guarded anyway because the tile round-up it does was itself 32-bit.

Built CPU-only under ASan and the Metal backend with -DMLX_BUILD_METAL=ON. The added tests crash or fail on the current code and pass with the fix. The full C++ suite passes 251/251.

@zcbenz
zcbenz force-pushed the fix/pad-conv-shape-overflow-3611 branch from bd04769 to a9a7bc1 Compare August 16, 2026 04:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants