Fix Metal strided-input handling: recip/sign, index_select, gather#3706
Open
v-code01 wants to merge 3 commits into
Open
Fix Metal strided-input handling: recip/sign, index_select, gather#3706v-code01 wants to merge 3 commits into
v-code01 wants to merge 3 commits into
Conversation
Author
|
cc @ivarflakstad @EricLBuehler — small dispatch gap: the strided |
candle-metal-kernels generates strided `recip` and `sign` kernels (both are in
the unary `ops!` list and both already have contiguous dispatch entries), but the
strided unary dispatch never mapped them. As a result `.recip()` and `.sign()` on
a non-contiguous tensor (e.g. a transpose) errored with "Metal strided unary
u{recip,sign} ... not implemented" while every other unary op and the CPU backend
handled it fine.
Map `urecip`/`usign` to the existing strided kernels for F32/F16/BF16, matching
their contiguous coverage. Adds a regression test comparing Metal strided recip
and sign against the CPU reference across the three dtypes.
The Metal index_select strided path passed `src_dim_size` (the size of the selected dimension) to `get_strided_index` where the source rank (`num_dims`) is expected. The strided-index loop therefore ran `src_dim_size` iterations instead of `rank`, reading `src_dims`/`src_strides` out of bounds and computing the wrong source offset -- so index_select on a non-contiguous source (e.g. a transpose) silently returned wrong values on Metal, while matching the CPU backend only when the rank happened to equal the selected-dim size (as in the existing 2x2 `index_select_strided` test). Thread the source rank through the kernel and pass it to `get_strided_index`. Adds a regression test comparing strided-source index_select against the CPU reference for shapes where rank != selected-dim size, across dim 0 and dim 1.
The Metal gather kernel indexes the source assuming a contiguous layout, but the dispatch only checked the ids for contiguity, not the source. A strided source (e.g. a transpose) therefore produced silently-wrong results, while the CPU backend errors with RequiresContiguous. Require a contiguous source, matching the CPU backend and the existing ids check. Adds a regression test that a strided-source gather now errors and that a contiguous gather still matches the CPU result.
v-code01
force-pushed
the
metal-strided-recip-sign
branch
from
7bb5015 to
02b6896
Compare
v-code01
changed the title
This was referenced Jul 10, 2026
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.
What
Found via a Metal-vs-CPU differential sweep over ops on non-contiguous (strided) inputs. Three related issues in the Metal backend's strided-input handling, each with a regression test vs the CPU reference:
recip/sign(unary)index_selectget_strided_indexnum_dims)gatherRequiresContiguous, matching the CPU-side contractThe two silent-correctness ones (
index_select,gather) are the important fixes — they returned wrong results with no error.Why one PR
All three are the same class (Metal op mishandles a strided/non-contiguous input) found by the same differential-testing pass; grouping keeps the review as one coherent correctness sweep.
Tests
metal_strided_unary_tests,metal_index_select_strided,metal_gather_contiguous— each compares Metal vs CPU on strided inputs. Verified on Apple M4 (Metal).cargo fmt --allclean; no new clippy warnings.cc @LaurentMazare @ivarflakstad @EricLBuehler