Support 8-byte dtypes in assignment scatter on Metal - #4328
Open
erwinzhang7 wants to merge 1 commit into
Open
Conversation
Metal scatter rejected every 8-byte output type, so mx.eye and mx.diag both failed on the GPU for int64, uint64 and complex64. The cause is that all scatters write through mlx_atomic<T>. For a type with no native Metal atomic that falls back to packing several values into one atomic<uint>, which cannot represent an 8-byte value at all, so the dispatch refused them outright. Only the reducing scatters need read-modify-write. Assignment is a plain store: scattering duplicate indices under Scatter::None is already a race whose result is whichever thread writes last, and a plain store has exactly those semantics. So the output pointer is now chosen by the op — mlx_atomic<T> for the reducing scatters, plain T for assignment — and the two type gates only reject the reducing case. Kernel names already encode the op, so the two forms cannot collide in the library cache. This fixes eye and diag as a consequence rather than as a special case, and it covers put_along_axis and every other 8-byte assignment scatter with it. test_put_along_axis asserted the old limitation and now verifies the result on both streams. test_scatter_8_byte_types covers the reported cases, checks assignment against the CPU with unique indices, and keeps the invariant that a reducing scatter still raises rather than returning something plausible. Indices are unique in the new test on purpose: assignment with a repeated index is a race by definition, so the GPU and the CPU are both correct and need not agree.
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
1 participant
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.
Proposed changes
Root cause identified: fixes #4300, supersedes the two closed workarounds, #4301 (eye) and #4309 (diag)
mx.eyeandmx.diagfail on the GPU forint64,uint64andcomplex64because Metal scatter rejects every 8-byte output type. Both build their result
with an assignment scatter, so fixing the scatter fixes them, and
put_along_axisand anything else assigning 8-byte data with it.
Why they were rejected. Every scatter writes through
device mlx_atomic<T>* out. For a type with no native Metal atomic,mlx_atomic<T>packs several values into a singleatomic<uint>, which cannotrepresent an 8-byte value, so
Scatter::eval_gpuand thescatterop refusedthem outright.
That is real for the reducing scatters and was never real for assignment.
Metal has no 8-byte atomic, so
Sum,Prod,MaxandMindo needread-modify-write. Scattering duplicate indices under
Scatter::Noneisa race and the result is whichever thread writes last; a plain store has those
semantics.
The change. The output pointer type is chosen by the op:
mlx_atomic<T>forthe reducing scatters, plain
Tfor assignment.Nonegains a non-atomicoverload,
scatter_impldeduces the pointer type,scatter_axistakes it as atemplate parameter, and the JIT signature takes it as a placeholder. Three type
gates now reject only the reducing case:
Scatter::eval_gpu,ScatterAxis::eval_gpu, and the one inops.cppcarrying// TODO, remove when scatter supports 64-bit outputs. Kernel names alreadyencode the op, so the two pointer forms cannot collide in the library cache.
CUDA and CPU are unaffected; both already handled these dtypes, which is why the
old gate was GPU-only.
complex64 + Sumkeeps its existing path.Verified
macOS 26.2, M5 Max.
mx.eyeandmx.diagforint64,uint64andcomplex64on both streams. Assignment scatter against the CPU across those three plus seven
control dtypes over four shapes. Duplicate indices land a valid candidate and
leave other slots untouched. Reducing scatter still raises for
int64anduint64.test_put_along_axisasserted the old limitation and now checks the result onboth streams.
test_scatter_8_byte_typescovers the reported ops and keeps theinvariant that a reducing scatter raises rather than returning something
plausible. Its indices are unique on purpose: assignment with a repeated index is
a race by definition, so the GPU and CPU need not agree.
Full python suite 839 passed, 11,814 subtests, 5 skipped.