Skip to content

Fix elementwise unary ops giving different results for strided vs contiguous inputs (CPU) (#4163) - #4187

Closed
AxelNoun wants to merge 3 commits into
ml-explore:mainfrom
AxelNoun:fix/4163-strided-unary-consistency-pr
Closed

Fix elementwise unary ops giving different results for strided vs contiguous inputs (CPU) (#4163)#4187
AxelNoun wants to merge 3 commits into
ml-explore:mainfrom
AxelNoun:fix/4163-strided-unary-consistency-pr

Conversation

@AxelNoun

@AxelNoun AxelNoun commented Aug 11, 2026

Copy link
Copy Markdown

Summary

  • On Apple Silicon, CPU unary ops used Accelerate SIMD for contiguous inputs and scalar libm for strided views (and for the contiguous residual), so the same values could differ by ~1 ULP depending on memory layout / lane index.
  • Route strided loads and the contiguous residual through the same padded SIMD kernel so results are bit-identical to the contiguous fast path.
  • Commits are split: (1) strided gather, (2) contiguous residual, (3) regression tests so the residual-only fix remains submittable if bit-identity across layouts is considered out of scope (NumPy does not guarantee it either).

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)

Run Result
Before (main) baseline 296/300 mismatches (issue repro)
After (this fix) fix 0/300, position-independence OK, full python -m unittest discover python/tests green

Micro-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:

Op (strided) cols Before (main) After (fix)
abs 2 14.14 ms 7.91 ms
negative 2 11.86 ms 3.23 ms
tan 2 20.03 ms 9.29 ms
abs 8 11.01 ms 8.37 ms
negative 8 10.52 ms 3.38 ms
tan 8 17.28 ms 9.06 ms

No regression observed: SIMD+gather is faster than the previous pure-scalar strided path for both trivial and expensive ops on this runner.

Test plan

AxelNoun and others added 3 commits August 11, 2026 19:37
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>
@zcbenz

zcbenz commented Aug 12, 2026

Copy link
Copy Markdown
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.

@zcbenz zcbenz closed this Aug 12, 2026
@AxelNoun
AxelNoun deleted the fix/4163-strided-unary-consistency-pr branch August 12, 2026 02:32
@AxelNoun

Copy link
Copy Markdown
Author

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.

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?

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.tan(a[:, 0]) differs from mx.tan(contiguous_copy) on float64 arrays

2 participants