Thanks to visit codestin.com
Credit goes to github.com

Skip to content

[Metal] Fix 32-bit integer overflow in conv3d unfold kernel#3143

Merged
angeloskath merged 2 commits intoml-explore:mainfrom
kellen-sun:fix-conv3d-overflow
Feb 19, 2026
Merged

[Metal] Fix 32-bit integer overflow in conv3d unfold kernel#3143
angeloskath merged 2 commits intoml-explore:mainfrom
kellen-sun:fix-conv3d-overflow

Conversation

@kellen-sun
Copy link
Contributor

@kellen-sun kellen-sun commented Feb 18, 2026

Proposed changes

Fixes #3138
Bug: Large 3D convolutions ($M \times K > 2^{32}$) were causing silent data corruption due to 32-bit integer overflow in the naive_unfold_Nd kernel index calculation.
The Fix: Explicitly cast gid's to size_t before multiplying by filter_size to ensure 64-bit arithmetic for the output offset.
Verification: Verified locally (macOS 26.2, MLX 0.30.7, Apple M2) with a script very similar to the one described in the issue. Except using float16, and T=2500, H=64, W=64, C=16, to be just on the edge of 2^32, and my system's memory limit for a single buffer (some 9.5GB).
Result before change: M*K= 4,423,680,000 > 2^32 max_diff=152.884216 corr=0.956153 [FAIL]
Result after change: M*K= 4,423,680,000 > 2^32 max_diff=0.031334 corr=1.000000 [OK]

Didn't add a new test for this because it would need to allocate at least 2^32 * sizeof(type) to reproduce.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

Copy link
Member

@angeloskath angeloskath left a comment

Choose a reason for hiding this comment

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

Thanks!

@angeloskath angeloskath merged commit d4c8106 into ml-explore:main Feb 19, 2026
16 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.

[BUG] mx.conv3d produces silently wrong results when the unfolded matrix has M*K > 2^32 elements

2 participants