Select diag's diagonal out of a broadcast for 8-byte types - #4309
Closed
erwinzhang7 wants to merge 1 commit into
Closed
Select diag's diagonal out of a broadcast for 8-byte types#4309erwinzhang7 wants to merge 1 commit into
erwinzhang7 wants to merge 1 commit into
Conversation
diag places a 1-D input on the diagonal with a scatter, and the Metal scatter has no 8-byte output path, so int64, uint64 and complex64 fail on the GPU and work on the CPU. eye has the same problem and can build in a narrower type and cast, because the only values it writes are 0 and 1. That is not available here: these are the caller's values, and an int64 past float32's exact range comes back changed. 2^53 + 1 becomes 2^53, and 2^62 - 1 becomes 2^62. So the diagonal is selected out of a broadcast instead, which never leaves the dtype. Only the types that currently throw take that path, since it allocates an n x n mask alongside the n x n output and the scatter does not.
Member
|
Since this is not an important bug I prefer fixing the root issue rather than adding a workaround. |
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
2 participants
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.
Fixes #4308.
diagplaces a 1-D input on the diagonal with ascatter, andScatter::eval_gpurefusesany 8-byte output that is not a
complex64sum, soint64,uint64andcomplex64failon the GPU and work on the CPU.
eyereached the same guard and #4301 fixes it by building in a narrower type and casting.That works there because the only values
eyewrites are 0 and 1, exact in every type itcan return. It does not work here: these are the caller's values, and float32 has 24 bits
of mantissa.
So the 8-byte types select the diagonal out of a broadcast rather than scattering into it,
which never leaves the dtype. Only those three take that path: the mask is an n x n
allocation the scatter does not need, and there is no reason to spend it on the types that
already work.
Verified
M5 Max, cpu and gpu streams. 180 cases over int64, uint64, complex64, int32, float32 and
uint8, offsets -3 to 2, and lengths 1, 3 and 6, values checked against numpy: all pass.
2^53 + 1,2^62 - 1and2^64 - 1round trip exactly, which is the case a cast wouldhave changed. The empty input and the 2-D extraction path are unchanged.
test_ops,test_autogradandtest_arraypass. Fork CI green on all 22 jobs, includingmetal and jit on macOS 14, 15 and 26.2, and the cuda 12.6, 12.9 and 13.0 legs.