diff --git a/mlx/array.h b/mlx/array.h index fd71104bc2..59ab941945 100644 --- a/mlx/array.h +++ b/mlx/array.h @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -565,20 +566,19 @@ T array::item() { throw std::invalid_argument("item can only be called on arrays of size 1."); } eval(); - return *data(); + // Use memcpy to read unaligned data. + T value; + std::memcpy(&value, data(), sizeof(T)); + return value; } template T array::item() const { - if (size() != 1) { - throw std::invalid_argument("item can only be called on arrays of size 1."); - } if (status() == Status::unscheduled) { throw std::invalid_argument( "item() const can only be called on evaled arrays"); } - const_cast(this)->eval(); - return *data(); + return const_cast(this)->item(); } template diff --git a/mlx/backend/common/common.cpp b/mlx/backend/common/common.cpp index dbea6d16da..1e0692973c 100644 --- a/mlx/backend/common/common.cpp +++ b/mlx/backend/common/common.cpp @@ -9,18 +9,8 @@ namespace mlx::core { -void AsStrided::eval(const std::vector& inputs, array& out) { - assert(inputs.size() == 1); - - auto& in = inputs[0]; - - if (!in.flags().row_contiguous) { - // Just ensuring that inputs[0] came from the ops which would ensure the - // input is row contiguous. - throw std::runtime_error( - "AsStrided must be used with row contiguous arrays only."); - } - +void AsStrided::eval(const array& in, array& out) { + assert(in.flags().row_contiguous); auto [_, row_contiguous, col_contiguous] = check_contiguity(shape_, strides_); int64_t l = 0, h = 0; diff --git a/mlx/backend/cpu/primitives.cpp b/mlx/backend/cpu/primitives.cpp index d2f5481b99..1633e0a647 100644 --- a/mlx/backend/cpu/primitives.cpp +++ b/mlx/backend/cpu/primitives.cpp @@ -86,7 +86,14 @@ static std::pair compute_dynamic_offset( } void AsStrided::eval_cpu(const std::vector& inputs, array& out) { - eval(inputs, out); + assert(inputs.size() == 1); + if (inputs[0].flags().row_contiguous) { + eval(inputs[0], out); + } else { + auto copy = contiguous_copy_cpu(inputs[0], stream()); + cpu::get_command_encoder(stream()).add_temporary(copy); + eval(copy, out); + } } void Broadcast::eval_cpu(const std::vector& inputs, array& out) { eval(inputs, out); diff --git a/mlx/backend/gpu/primitives.cpp b/mlx/backend/gpu/primitives.cpp index bf17029ff6..6df6ab3e09 100644 --- a/mlx/backend/gpu/primitives.cpp +++ b/mlx/backend/gpu/primitives.cpp @@ -8,6 +8,9 @@ #if defined(MLX_USE_CUDA) #include +#include "mlx/backend/cuda/device.h" +#else +#include "mlx/backend/metal/device.h" #endif #include @@ -22,7 +25,18 @@ namespace mlx::core { void AsStrided::eval_gpu(const std::vector& inputs, array& out) { MLX_PROFILER_RANGE("AsStrided::eval_gpu"); - eval(inputs, out); + assert(inputs.size() == 1); + if (inputs[0].flags().row_contiguous) { + eval(inputs[0], out); + } else { + auto copy = contiguous_copy_gpu(inputs[0], stream()); +#if defined(MLX_USE_CUDA) + cu::get_command_encoder(stream()).add_temporary(copy); +#else + metal::get_command_encoder(stream()).add_temporary(copy); +#endif + eval(copy, out); + } } void AsType::eval_gpu(const std::vector& inputs, array& out) { diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 01ad2778f8..a34de0a9cb 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -333,8 +333,7 @@ array as_strided( dtype, std::make_shared( to_stream(s), std::move(shape), std::move(strides), offset), - // Force the input array to be contiguous. - {flatten(std::move(a), s)}); + {std::move(a)}); } array copy(array a, StreamOrDevice s /* = {} */) { diff --git a/mlx/primitives.h b/mlx/primitives.h index 3ec8be22f9..68995b7503 100644 --- a/mlx/primitives.h +++ b/mlx/primitives.h @@ -443,7 +443,7 @@ class AsStrided : public UnaryPrimitive { Strides strides_; size_t offset_; - void eval(const std::vector& inputs, array& out); + void eval(const array& in, array& out); }; class BitwiseBinary : public UnaryPrimitive { diff --git a/tests/array_tests.cpp b/tests/array_tests.cpp index 7adad04a2b..f126ce3276 100644 --- a/tests/array_tests.cpp +++ b/tests/array_tests.cpp @@ -1,6 +1,7 @@ // Copyright © 2023 Apple Inc. #include #include +#include #include #include @@ -761,3 +762,19 @@ TEST_CASE("test assigning a sibling over the last reference keeps the cycle") { REQUIRE_EQ(last.siblings().size(), 1); CHECK(array_equal(last, array({2})).item()); } + +TEST_CASE("test unaligned scalar item") { + uint64_t expected = 0x0123456789abcdef; + for (size_t offset = 1; offset < sizeof(expected); ++offset) { + std::vector data(offset + sizeof(expected)); + std::memcpy(data.data() + offset, &expected, sizeof(expected)); + auto storage = array(data.data(), {static_cast(data.size())}, uint8); + auto scalar = view(as_strided(storage, {8}, {1}, offset), uint64); + eval(scalar); + if (default_device() == Device::cpu) { + CHECK_EQ(scalar.data(), storage.data() + offset); + } + CHECK_EQ(scalar.item(), expected); + CHECK_EQ(static_cast(scalar).item(), expected); + } +} diff --git a/tests/autograd_tests.cpp b/tests/autograd_tests.cpp index 1ff203e1e9..c1e3d20d44 100644 --- a/tests/autograd_tests.cpp +++ b/tests/autograd_tests.cpp @@ -1111,6 +1111,14 @@ TEST_CASE("test as_strided grads") { out = vjp(fun, x, reshape(astype(arange(9), x.dtype()), {3, 3})).second; expected = array({0, 0, 0, 9, 12, 15, 0, 0, 0, 0, 0}); CHECK(array_equal(out, expected).item()); + + x = transpose(reshape(arange(12, float32), {3, 4})); + out = vjp(fun, x, ones(shape)).second; + expected = reshape(array({0, 0, 0, 3, 3, 3, 0, 0, 0, 0, 0, 0}), x.shape()); + CHECK(array_equal(out, expected).item()); + + out = jvp(fun, x, x).second; + CHECK(array_equal(out, fun(x)).item()); } TEST_CASE("test jvp from vjp") { diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index cda9523f37..da0f9e2873 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -3012,6 +3012,16 @@ TEST_CASE("test as_strided op") { y = as_strided(x, {3, 3}, {2, 1}, 1); expected = array({5, 1, 6, 6, 2, 7, 7, 3, 8}, {3, 3}); CHECK(array_equal(y, expected).item()); + + x = slice(arange(12), {10}, {0}, {-2}); + y = as_strided(x, {2, 2}, {2, 1}, 1); + expected = array({8, 6, 4, 2}, {2, 2}); + CHECK(array_equal(y, expected).item()); + + x = broadcast_to(array({1, 2, 3}), {2, 3}); + y = as_strided(x, {2, 2}, {2, 1}, 1); + expected = array({2, 3, 1, 2}, {2, 2}); + CHECK(array_equal(y, expected).item()); } TEST_CASE("test sort and scan on empty arrays") {