Skip to content

Fix signed-integer overflow in convolution shape arithmetic - #3938

Merged
zcbenz merged 2 commits into
ml-explore:mainfrom
eyupcanakman:conv-shape-overflow-guards
Aug 6, 2026
Merged

Fix signed-integer overflow in convolution shape arithmetic#3938
zcbenz merged 2 commits into
ml-explore:mainfrom
eyupcanakman:conv-shape-overflow-guards

Conversation

@eyupcanakman

@eyupcanakman eyupcanakman commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Follow-up to #3611, which was closed as mostly fixed with an invitation to open PRs for whatever still hit UB. #3524 handled the Metal and CUDA sides. The shape arithmetic itself is still 32-bit.

conv_out_shape computes 1 + dilation * (dim - 1) and the padding sum in 32 bits. conv_transpose_general works out all of its padding in 32 bits before the nested conv_general validates anything, and the backward pass repeats those same formulas in Convolution::vjp and conv_weight_backward_patches. Extreme but in-range int32 parameters wrap rather than raise. On an 8x8 input with a 3x3 weight, mx.conv_general(..., kernel_dilation=2**31-1) returns a 10x10 output.

These now compute in int64_t and narrow through safe_cast, the same way #3524 and #3604 did. Convolution::vjp calls the 64-bit dilate_size helper rather than repeating the formula, which is what #2601 pulled it out for.

Most of those inputs raised nothing before. The ones that did move from std::invalid_argument to std::overflow_error, so Python raises OverflowError where it used to raise ValueError. In-range convolutions are unchanged.

Built CPU-only with -DUSE_UBSAN=ON. The added test drives 15 signed-overflow reports before the fix and none after, and 9 of its 11 assertions fail without it. Full C++ suite passes 245/245, Python 784 passed.

Two things #3611 also lists are left alone. The lcm jump-table sizes are in the Metal backend, and the dim + 1 half of its negative-padding item is in normalize_slice, reachable from mx.slice with no convolution involved.

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@zcbenz
zcbenz force-pushed the conv-shape-overflow-guards branch from e728567 to 2c3f1fe Compare August 6, 2026 07:51
@zcbenz
zcbenz merged commit 7e9c3f6 into ml-explore:main Aug 6, 2026
28 checks passed
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