Fix elementwise unary ops giving different results for strided vs contiguous inputs (CPU) (#4163) - #4187
Closed
AxelNoun wants to merge 3 commits into
Closed
Conversation
Strided views previously fell back to scalar libm while contiguous inputs used Accelerate SIMD, causing ~1 ULP differences for the same values. Gather into a SIMD buffer so both paths share one math kernel. Co-authored-by: Cursor <cursoragent@cursor.com>
The contiguous fast path used Accelerate SIMD for full vectors and scalar libm for the leftover lanes, so the same value could differ by ~1 ULP depending on its position in the array. Route the residual through the padded SIMD kernel instead. Related to the unary/lane-index inconsistency described in ml-explore#4161 (binary ops like x**2 still need a binary.h follow-up). Co-authored-by: Cursor <cursoragent@cursor.com>
Cover float32/float64 sizes that cross SIMD widths, checking both strided/contiguous bit-identity and position-independent residuals. Co-authored-by: Cursor <cursoragent@cursor.com>
Member
|
Thanks for the PR but we plan to switch to use google/highway for SIMD and at the moment (#3019), so we don't accept non-critical changes on SIMD code. |
Author
Oh okay! So much the better, tell me if there is a referenced place where the different steps that have not yet been resolved would be marked to be able to contribute to this task? |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes #4163.
Possibly related to #4161: this PR fixes the unary lane-index / length dependence shown there ( an, �rctan, etc.). Binary cases in that issue (x2, x0.5, …) share the same Accelerate-vs-libm pattern in �inary.h and are not fixed here happy to follow up the same way in �inary.h/ ernary.h if this approach is accepted. Leaving #4161 open for the binary remainder.
Validation (no local Mac GitHub Actions macos-26 / Accelerate)
python -m unittest discover python/testsgreenMicro-benchmark (strided before vs after, float64, n=1e6, 50 iters, same runner)
The useful comparison is strided-before vs strided-after (not strided vs contiguous). Trivial ops included because gather overhead is most visible there:
absnegativetanabsnegativetanNo regression observed: SIMD+gather is faster than the previous pure-scalar strided path for both trivial and expensive ops on this runner.
Test plan
mx.tan(a[:, 0])differs frommx.tan(contiguous_copy)on float64 arrays #4163 repro script: 296/300 on main → 0/300 with fix (CI links above)mx.tan(and**,arctan) on float64 give different results depending on ARRAY LENGTH and LANE INDEX #4161 (unarytan)python -m unittest discover python/testson macos-26 CPU/Acceleratepre-commit(clang-format / black) on touched files