diff --git a/mlx/backend/common/compiled.cpp b/mlx/backend/common/compiled.cpp index aceeb1f7fd..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,6 +109,7 @@ 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; @@ -107,7 +118,7 @@ bool compiled_check_contiguity( } 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,7 +231,11 @@ 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( diff --git a/mlx/backend/common/compiled.h b/mlx/backend/common/compiled.h index 84a3460459..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,7 +82,8 @@ 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); 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 4ffb959ac2..a363cf9c31 100644 --- a/mlx/backend/cuda/compiled.cpp +++ b/mlx/backend/cuda/compiled.cpp @@ -285,11 +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. - bool large = compiled_use_large_index(inputs, outputs, contiguous); + // Whether to use large index (also true for negative strides). + bool large = + 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 cdb0a471be..cda06143d6 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, @@ -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. - bool large = compiled_use_large_index(inputs, outputs, contiguous); + // Whether to use large index (also true for negative strides). + bool large = + negative_strides || compiled_use_large_index(inputs, outputs, contiguous); // 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()