Fix compiled kernel correctness for negative-strided inputs#3720
Conversation
lyonsno
force-pushed
the
cc/fix-compile-negative-stride
branch
from
db1bdb0 to
ac40f52
Compare
| if (non_scalar_inputs > 1 && !all_row_contig && !all_col_contig) { | ||
| contiguous = false; | ||
| } else if (non_scalar_inputs == 1 && !all_contig) { | ||
| } else if (non_scalar_inputs == 1 && !(all_row_contig || all_col_contig)) { |
There was a problem hiding this comment.
Why the change here? Note that all_contig != (all_row_contig || all_col_contig)
There was a problem hiding this comment.
A negative-strided array (e.g. x[::-1]) has flags().contiguous == true (densely packed, no gaps) but flags().row_contiguous == false. The old check let these through to the contiguous compiled kernel, which iterates linearly from the data pointer — but for a reversed view the data pointer is at the last element, so the linear iteration reads the wrong values. This aligns the single-input path with the multi-input check on line 103 (which already uses row_contiguous / col_contiguous).
There was a problem hiding this comment.
Yeah I know that but for a single contiguous input with positive strides we don't need to use the strides at all which is why that check was like that.
Basically now doing gelu on a transposed array will be slower than before for no reason.
angeloskath
left a comment
angeloskath
left a comment
There was a problem hiding this comment.
Thanks for the fix!
Did a little change, feel free to take a look, but basically gelu(x.transpose(0, 2, 1)) will now route to the appropriate contiguous kernel.
angeloskath
force-pushed
the
cc/fix-compile-negative-stride
branch
from
2069dbf to
7bb86ee
Compare
The compiled (fused) kernel path produced wrong results when inputs had negative strides (e.g. x[::-1]). Two issues: 1. compiled_check_contiguity used the broad contiguous flag for single inputs, which is true for negative-strided arrays (no data gaps). Changed to require row_contiguous or col_contiguous, matching the multi-input path. 2. Metal/CUDA strided compiled kernels used unsigned index arithmetic (elem_to_loc_1<uint>), wrapping negative strides. Force int64_t indices when any input has negative strides. Also generate the _large (int64_t) strided kernel variant for ndim=1. The CPU compiled path uses signed pointer arithmetic and only needed the contiguity check fix. Fixes ml-explore#3716. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
The strides vector from compiled_collapse_contiguous_dims is ordered [output, input_0, input_1, ...]. Expand the comment to make this clear and explain why only input entries need the negative-stride check. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
angeloskath
force-pushed
the
cc/fix-compile-negative-stride
branch
from
7bb86ee to
5c23944
Compare
3 participants
Fixes #3716.
This is an alternative to #3717 that fixes the compiled-kernel path in place instead of materializing negative-strided inputs before the generated kernel runs.
mx.compileproduces incorrect results for negative-strided inputs such asx[::-1], while eager execution is correct.This patch fixes two compiled-kernel path issues:
_largeMetal variant is now always generated (matching ndim > 1 behavior) since negative strides force int64_t indexing even for 1D arrays.Added regression coverage for 1D reverse, slice update, 2D reverse, mixed positive/negative strides, and a 4D negative-stride case.
Tests:
python -m pytest python/tests/test_compile.py -q # 59 passed