From 725b751008bb9d11fb387646df682046fa766351 Mon Sep 17 00:00:00 2001 From: Noah Lyons Date: Fri, 19 Jun 2026 00:37:37 -0400 Subject: [PATCH 1/3] Fix compiled kernel correctness for negative-strided inputs 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), 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 #3716. Co-Authored-By: Claude Opus 4.6 (1M context) --- mlx/backend/common/compiled.cpp | 18 ++++++++++--- mlx/backend/common/compiled.h | 5 +++- mlx/backend/cuda/compiled.cpp | 5 ++-- mlx/backend/metal/compiled.cpp | 36 ++++++++++++------------- python/tests/test_compile.py | 47 +++++++++++++++++++++++++++++++++ 5 files changed, 87 insertions(+), 24 deletions(-) diff --git a/mlx/backend/common/compiled.cpp b/mlx/backend/common/compiled.cpp index aceeb1f7fd..a5c9934bd8 100644 --- a/mlx/backend/common/compiled.cpp +++ b/mlx/backend/common/compiled.cpp @@ -102,7 +102,7 @@ bool compiled_check_contiguity( } 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)) { contiguous = false; } else if (non_scalar_inputs == 0 && !shape.empty()) { contiguous = false; @@ -224,7 +224,8 @@ std::tuple> compiled_collapse_contiguous_dims( bool compiled_use_large_index( const std::vector& inputs, const std::vector& outputs, - bool contiguous) { + bool contiguous, + const std::vector& strides) { if (contiguous) { size_t max_size = 0; for (const auto& in : inputs) { @@ -236,7 +237,18 @@ bool compiled_use_large_index( for (const auto& o : outputs) { max_size = std::max(max_size, o.size()); } - return max_size > UINT32_MAX; + if (max_size > UINT32_MAX) { + return true; + } + // Check for negative strides in inputs (strides[0] is the output). + for (size_t i = 1; i < strides.size(); ++i) { + for (auto v : strides[i]) { + if (v < 0) { + return true; + } + } + } + return false; } } diff --git a/mlx/backend/common/compiled.h b/mlx/backend/common/compiled.h index 84a3460459..7a01642e77 100644 --- a/mlx/backend/common/compiled.h +++ b/mlx/backend/common/compiled.h @@ -87,9 +87,12 @@ std::tuple> compiled_collapse_contiguous_dims( const std::function& is_constant); // Return whether the kernel should use large index. +// Also returns true when any non-contiguous input has negative strides, +// since unsigned index arithmetic wraps negative stride values. bool compiled_use_large_index( const std::vector& inputs, const std::vector& outputs, - bool contiguous); + bool contiguous, + const std::vector& strides = {}); } // namespace mlx::core diff --git a/mlx/backend/cuda/compiled.cpp b/mlx/backend/cuda/compiled.cpp index 4ffb959ac2..1389038951 100644 --- a/mlx/backend/cuda/compiled.cpp +++ b/mlx/backend/cuda/compiled.cpp @@ -288,8 +288,9 @@ void Compiled::eval_gpu( auto [contiguous, shape, strides_vec] = compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_); - // Whether to use large index. - bool large = compiled_use_large_index(inputs, outputs, contiguous); + // Whether to use large index (also true for negative strides). + bool large = + compiled_use_large_index(inputs, outputs, contiguous, strides_vec); cu::KernelArgs args; // Put inputs. diff --git a/mlx/backend/metal/compiled.cpp b/mlx/backend/metal/compiled.cpp index cdb0a471be..6076136455 100644 --- a/mlx/backend/metal/compiled.cpp +++ b/mlx/backend/metal/compiled.cpp @@ -139,8 +139,8 @@ inline void build_kernel( os += fmt::format(" {0} index_{1} = ", idx_type, xname); if (ndim == 1) { int offset = i * ndim; - os += - fmt::format("elem_to_loc_1(pos.x, in_strides[{0}]);\n", offset); + os += fmt::format( + "elem_to_loc_1<{0}>(pos.x, in_strides[{1}]);\n", idx_type, offset); } else if (ndim == 2) { int offset = i * ndim; os += fmt::format( @@ -316,20 +316,20 @@ void Compiled::eval_gpu( /* dynamic_dims = */ false, /* use_big_index = */ false, /* work_per_thread = */ i > 3 ? 2 : 1); - if (i > 1) { - build_kernel( - kernel, - kernel_lib_ + "_strided_" + std::to_string(i) + "_large", - inputs_, - outputs_, - tape_, - is_constant_, - /* contiguous = */ false, - /* ndim = */ i, - /* dynamic_dims = */ false, - /* use_big_index = */ true, - /* work_per_thread = */ i > 3 ? 4 : 1); - } + // Generate int64_t index variant for all ndim, including ndim=1. + // Negative strides force large mode even for small arrays. + build_kernel( + kernel, + kernel_lib_ + "_strided_" + std::to_string(i) + "_large", + inputs_, + outputs_, + tape_, + is_constant_, + /* contiguous = */ false, + /* ndim = */ i, + /* dynamic_dims = */ false, + /* use_big_index = */ true, + /* work_per_thread = */ i > 3 ? 4 : 1); } build_kernel( kernel, @@ -363,8 +363,8 @@ void Compiled::eval_gpu( auto [contiguous, shape, strides] = compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_); - // Whether to use large index. - bool large = compiled_use_large_index(inputs, outputs, contiguous); + // Whether to use large index (also true for negative strides). + bool large = compiled_use_large_index(inputs, outputs, contiguous, strides); // Get the kernel from the lib int ndim = shape.size(); diff --git a/python/tests/test_compile.py b/python/tests/test_compile.py index 4b7643dadd..35e51600e1 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -1454,6 +1454,53 @@ def fun(x): np.asarray(out, copy=False).__array_interface__["data"][0], in_ptr ) + def test_compile_negative_strides(self): + # 1D negative stride with elementwise expression + @mx.compile + def f(x): + return 2.0 * x[::-1] + + x = mx.arange(8, dtype=mx.float32) + expected = 2.0 * x[::-1] + self.assertTrue(mx.array_equal(f(x), expected)) + + # 1D negative stride with slice update + def g_eager(x): + base = mx.zeros_like(x) + base[::-1] += 2.0 * x[::-1] + return base + + g_compiled = mx.compile(g_eager) + expected = g_eager(x) + self.assertTrue(mx.array_equal(g_compiled(x), expected)) + + # 2D negative stride + @mx.compile + def h(x): + return x[::-1] + 1.0 + + y = mx.arange(12, dtype=mx.float32).reshape(3, 4) + expected = y[::-1] + 1.0 + self.assertTrue(mx.array_equal(h(y), expected)) + + # Mixed positive and negative strides + @mx.compile + def m(x): + return x[::-1, ::2] * 3.0 + + z = mx.arange(24, dtype=mx.float32).reshape(4, 6) + expected = z[::-1, ::2] * 3.0 + self.assertTrue(mx.array_equal(m(z), expected)) + + # 4D negative stride (exercises work_per_thread > 1 path) + @mx.compile + def p(x): + return x + 1.0 + + w = mx.arange(120, dtype=mx.float32).reshape(2, 3, 4, 5) + expected = w[::-1, :, ::-1, :] + 1.0 + self.assertTrue(mx.array_equal(p(w[::-1, :, ::-1, :]), expected)) + if __name__ == "__main__": mlx_tests.MLXTestRunner() From 6155599df04c51e4b46092deb22bc296a978e529 Mon Sep 17 00:00:00 2001 From: Noah Lyons Date: Mon, 22 Jun 2026 03:41:18 -0400 Subject: [PATCH 2/3] Clarify strides vector ordering comment in compiled_use_large_index 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) --- mlx/backend/common/compiled.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/mlx/backend/common/compiled.cpp b/mlx/backend/common/compiled.cpp index a5c9934bd8..b795287395 100644 --- a/mlx/backend/common/compiled.cpp +++ b/mlx/backend/common/compiled.cpp @@ -240,7 +240,10 @@ bool compiled_use_large_index( if (max_size > UINT32_MAX) { return true; } - // Check for negative strides in inputs (strides[0] is the output). + // Force int64_t indices when any input has negative strides, since + // unsigned elem_to_loc wraps negative values. strides is ordered + // [output, input_0, input_1, ...] from compiled_collapse_contiguous_dims; + // the output is freshly allocated so its strides are always non-negative. for (size_t i = 1; i < strides.size(); ++i) { for (auto v : strides[i]) { if (v < 0) { From 5c239440428a81e1e62a53946207be81a4a28b17 Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Mon, 6 Jul 2026 11:04:58 -0700 Subject: [PATCH 3/3] Use int64 only when contiguous and negative strides --- mlx/backend/common/compiled.cpp | 52 +++++++++++++++++---------------- mlx/backend/common/compiled.h | 13 ++++----- mlx/backend/cpu/compiled.cpp | 2 +- mlx/backend/cuda/compiled.cpp | 4 +-- mlx/backend/metal/compiled.cpp | 5 ++-- 5 files changed, 39 insertions(+), 37 deletions(-) diff --git a/mlx/backend/common/compiled.cpp b/mlx/backend/common/compiled.cpp index b795287395..5173be421f 100644 --- a/mlx/backend/common/compiled.cpp +++ b/mlx/backend/common/compiled.cpp @@ -82,7 +82,16 @@ std::string get_type_string(Dtype d) { } } -bool compiled_check_contiguity( +bool has_negative_strides(const array& x) { + for (auto s : x.strides()) { + if (s < 0) { + return true; + } + } + return false; +} + +std::pair compiled_check_contiguity( const std::vector& inputs, const Shape& shape) { bool contiguous = true; @@ -90,6 +99,7 @@ bool compiled_check_contiguity( bool all_row_contig = true; bool all_col_contig = true; int non_scalar_inputs = 0; + bool negative_strides = false; for (const auto& x : inputs) { if (is_scalar(x)) { continue; @@ -99,15 +109,16 @@ bool compiled_check_contiguity( all_contig &= (x.flags().contiguous && shape_eq); all_row_contig &= (x.flags().row_contiguous && shape_eq); all_col_contig &= (x.flags().col_contiguous && shape_eq); + negative_strides |= has_negative_strides(x); } if (non_scalar_inputs > 1 && !all_row_contig && !all_col_contig) { contiguous = false; - } else if (non_scalar_inputs == 1 && !(all_row_contig || all_col_contig)) { + } else if (non_scalar_inputs == 1 && !all_contig) { contiguous = false; } else if (non_scalar_inputs == 0 && !shape.empty()) { contiguous = false; } - return contiguous; + return {contiguous, negative_strides}; } void compiled_allocate_outputs( @@ -170,14 +181,16 @@ void compiled_allocate_outputs( } } -std::tuple> compiled_collapse_contiguous_dims( +std::tuple> +compiled_collapse_contiguous_dims( const std::vector& inputs, const array& out, const std::function& is_constant) { const Shape& shape = out.shape(); - bool contiguous = compiled_check_contiguity(inputs, shape); - if (contiguous) { - return {true, shape, {}}; + auto [contiguous, negative_strides] = + compiled_check_contiguity(inputs, shape); + if (contiguous && !negative_strides) { + return {true, false, shape, {}}; } std::vector strides_vec{out.strides()}; @@ -218,14 +231,17 @@ std::tuple> compiled_collapse_contiguous_dims( } auto tup = collapse_contiguous_dims(shape, strides_vec, INT32_MAX); - return {false, std::move(std::get<0>(tup)), std::move(std::get<1>(tup))}; + return { + false, + negative_strides, + std::move(std::get<0>(tup)), + std::move(std::get<1>(tup))}; } bool compiled_use_large_index( const std::vector& inputs, const std::vector& outputs, - bool contiguous, - const std::vector& strides) { + bool contiguous) { if (contiguous) { size_t max_size = 0; for (const auto& in : inputs) { @@ -237,21 +253,7 @@ bool compiled_use_large_index( for (const auto& o : outputs) { max_size = std::max(max_size, o.size()); } - if (max_size > UINT32_MAX) { - return true; - } - // Force int64_t indices when any input has negative strides, since - // unsigned elem_to_loc wraps negative values. strides is ordered - // [output, input_0, input_1, ...] from compiled_collapse_contiguous_dims; - // the output is freshly allocated so its strides are always non-negative. - for (size_t i = 1; i < strides.size(); ++i) { - for (auto v : strides[i]) { - if (v < 0) { - return true; - } - } - } - return false; + return max_size > UINT32_MAX; } } diff --git a/mlx/backend/common/compiled.h b/mlx/backend/common/compiled.h index 7a01642e77..8c6466da03 100644 --- a/mlx/backend/common/compiled.h +++ b/mlx/backend/common/compiled.h @@ -66,8 +66,9 @@ inline bool is_scalar(const array& x) { return x.ndim() == 0; } -// Check if we can use a contiguous operation given inputs and the output shape -bool compiled_check_contiguity( +// Check if we can use a contiguous operation given inputs and the output shape. +// Also returns if any input has negative strides. +std::pair compiled_check_contiguity( const std::vector& inputs, const Shape& shape); @@ -81,18 +82,16 @@ void compiled_allocate_outputs( allocator::malloc); // Collapse contiguous dims ignoring scalars and constants. -std::tuple> compiled_collapse_contiguous_dims( +std::tuple> +compiled_collapse_contiguous_dims( const std::vector& inputs, const array& out, const std::function& is_constant); // Return whether the kernel should use large index. -// Also returns true when any non-contiguous input has negative strides, -// since unsigned index arithmetic wraps negative stride values. bool compiled_use_large_index( const std::vector& inputs, const std::vector& outputs, - bool contiguous, - const std::vector& strides = {}); + bool contiguous); } // namespace mlx::core diff --git a/mlx/backend/cpu/compiled.cpp b/mlx/backend/cpu/compiled.cpp index 2cd70b1cdb..f5b516e44b 100644 --- a/mlx/backend/cpu/compiled.cpp +++ b/mlx/backend/cpu/compiled.cpp @@ -298,7 +298,7 @@ void Compiled::eval_cpu( // Collapse contiguous dims to route to a faster kernel if possible. Also // handle all broadcasting. - auto [contiguous, shape, strides] = + auto [contiguous, negative_strides, shape, strides] = compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_); // Collect function input arguments. diff --git a/mlx/backend/cuda/compiled.cpp b/mlx/backend/cuda/compiled.cpp index 1389038951..a363cf9c31 100644 --- a/mlx/backend/cuda/compiled.cpp +++ b/mlx/backend/cuda/compiled.cpp @@ -285,12 +285,12 @@ void Compiled::eval_gpu( // Collapse contiguous dims to route to a faster kernel if possible. Also // handle all broadcasting. - auto [contiguous, shape, strides_vec] = + auto [contiguous, negative_strides, shape, strides_vec] = compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_); // Whether to use large index (also true for negative strides). bool large = - compiled_use_large_index(inputs, outputs, contiguous, strides_vec); + negative_strides || compiled_use_large_index(inputs, outputs, contiguous); cu::KernelArgs args; // Put inputs. diff --git a/mlx/backend/metal/compiled.cpp b/mlx/backend/metal/compiled.cpp index 6076136455..cda06143d6 100644 --- a/mlx/backend/metal/compiled.cpp +++ b/mlx/backend/metal/compiled.cpp @@ -360,11 +360,12 @@ void Compiled::eval_gpu( // Collapse contiguous dims to route to a faster kernel if possible. Also // handle all broadcasting. - auto [contiguous, shape, strides] = + auto [contiguous, negative_strides, shape, strides] = compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_); // Whether to use large index (also true for negative strides). - bool large = compiled_use_large_index(inputs, outputs, contiguous, strides); + bool large = + negative_strides || compiled_use_large_index(inputs, outputs, contiguous); // Get the kernel from the lib int ndim = shape.size();