diff --git a/mlx/backend/metal/compiled.cpp b/mlx/backend/metal/compiled.cpp index cdb0a471be..e3e8712140 100644 --- a/mlx/backend/metal/compiled.cpp +++ b/mlx/backend/metal/compiled.cpp @@ -12,6 +12,20 @@ #include "mlx/utils.h" namespace mlx::core { +namespace { + +bool has_negative_strides(const std::vector& strides) { + for (const auto& stride : strides) { + for (auto value : stride) { + if (value < 0) { + return true; + } + } + } + return false; +} + +} // namespace inline void build_kernel( std::string& os, @@ -139,8 +153,10 @@ 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 +332,18 @@ 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); - } + 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, @@ -364,7 +378,8 @@ void Compiled::eval_gpu( compiled_collapse_contiguous_dims(inputs, outputs[0], is_constant_); // Whether to use large index. - bool large = compiled_use_large_index(inputs, outputs, contiguous); + bool large = compiled_use_large_index(inputs, outputs, contiguous) || + (!contiguous && has_negative_strides(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 632b34119a..88751189ab 100644 --- a/python/tests/test_compile.py +++ b/python/tests/test_compile.py @@ -5,6 +5,7 @@ import io import math import threading +import unittest from functools import partial, wraps from io import StringIO @@ -71,6 +72,25 @@ def test_compile_nonfinite_constants(self): self.assertEqual(out[0].item(), 1.0) self.assertEqual(out[1].item(), float("-inf")) + @unittest.skipIf(not mx.metal.is_available(), "Metal is not available") + def test_compile_negative_strided_slice_update_expr(self): + def add_update(x): + out = mx.zeros_like(x) + out[::-1] += 2.0 * x[::-1] + return out + + def set_update(x): + out = mx.zeros_like(x) + out[::-1] = 2.0 * x[::-1] + return out + + mx.set_default_device(mx.gpu) + x = mx.arange(6, dtype=mx.float32) + for fn in (add_update, set_update): + expected = fn(x) + actual = mx.compile(fn)(x) + self.assertEqualArray(actual, expected) + def test_compile_tuple_output_in_thread(self): @mx.compile def fun(x):