From 6d1db72fb886e52a5ca49720007c2e747e293be3 Mon Sep 17 00:00:00 2001 From: Alessio Pollero Date: Sun, 18 Oct 2026 10:02:06 +0400 Subject: [PATCH 1/3] Fix unaligned scalar reads and oversized contiguous buffer views --- mlx/array.h | 9 +++++++-- mlx/backend/cpu/primitives.cpp | 8 +++++++- mlx/backend/gpu/primitives.cpp | 15 +++++++++++++- mlx/ops.cpp | 3 +-- tests/array_tests.cpp | 17 ++++++++++++++++ tests/autograd_tests.cpp | 8 ++++++++ tests/ops_tests.cpp | 10 ++++++++++ tests/residency_tests.cpp | 36 ++++++++++++++++++++++++++++++++++ 8 files changed, 100 insertions(+), 6 deletions(-) diff --git a/mlx/array.h b/mlx/array.h index fd71104bc2..e10807a8cd 100644 --- a/mlx/array.h +++ b/mlx/array.h @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -565,7 +566,9 @@ T array::item() { throw std::invalid_argument("item can only be called on arrays of size 1."); } eval(); - return *data(); + T value; + std::memcpy(&value, data(), sizeof(T)); + return value; } template @@ -578,7 +581,9 @@ T array::item() const { "item() const can only be called on evaled arrays"); } const_cast(this)->eval(); - return *data(); + T value; + std::memcpy(&value, data(), sizeof(T)); + return value; } template diff --git a/mlx/backend/cpu/primitives.cpp b/mlx/backend/cpu/primitives.cpp index d2f5481b99..0474d02fda 100644 --- a/mlx/backend/cpu/primitives.cpp +++ b/mlx/backend/cpu/primitives.cpp @@ -86,7 +86,13 @@ static std::pair compute_dynamic_offset( } void AsStrided::eval_cpu(const std::vector& inputs, array& out) { - eval(inputs, out); + if (inputs[0].flags().row_contiguous) { + eval(inputs, 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..af10a14132 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,17 @@ namespace mlx::core { void AsStrided::eval_gpu(const std::vector& inputs, array& out) { MLX_PROFILER_RANGE("AsStrided::eval_gpu"); - eval(inputs, out); + if (inputs[0].flags().row_contiguous) { + eval(inputs, 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/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") { diff --git a/tests/residency_tests.cpp b/tests/residency_tests.cpp index 0edac9b6dc..3ca6f5535c 100644 --- a/tests/residency_tests.cpp +++ b/tests/residency_tests.cpp @@ -1,5 +1,8 @@ // Copyright © 2026 Apple Inc. +#include + +#include #include #include "doctest/doctest.h" @@ -294,3 +297,36 @@ TEST_CASE("test the wired limit is used up to its boundary") { allocator::allocator().free(second); CHECK_EQ(residency().wired_size(), baseline); } + +TEST_CASE("test as_strided on a mapped buffer larger than INT_MAX") { + constexpr size_t length = size_t{131073} * 16384; + if (length > metal::device(Device::gpu).mtl_device()->maxBufferLength()) { + return; + } + LimitGuard guard(0, 0); + auto pointer = mmap( + nullptr, length, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0); + REQUIRE(pointer != MAP_FAILED); + auto mapping = + std::shared_ptr(pointer, [](void* data) { munmap(data, length); }); + auto buffer = allocator::make_buffer(mapping.get(), length); + REQUIRE(buffer.ptr() != nullptr); + auto data = allocator::Data(buffer, [mapping](allocator::Buffer buffer) { + allocator::release(buffer); + }); + auto storage = array(std::move(data), {131073, 16384}, uint8); + auto bytes = static_cast(mapping.get()); + for (int i = 0; i < 64; ++i) { + bytes[length - 64 + i] = i + 1; + } + for (const auto& input : {storage, copy(storage)}) { + auto active = get_active_memory(); + reset_peak_memory(); + auto result = as_strided(input, {2, 32}, {32, 1}, length - 64); + eval(result); + CHECK_LE(get_peak_memory(), active); + CHECK_EQ(result.data(), bytes + length - 64); + CHECK_EQ(result.data()[0], 1); + CHECK_EQ(result.data()[63], 64); + } +} From 4def7a95bfbc842bd62e677f8bf95b55ac0ad155 Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 19 Oct 2026 13:55:19 +0900 Subject: [PATCH 2/3] avoid test based on active memory --- tests/residency_tests.cpp | 36 ------------------------------------ 1 file changed, 36 deletions(-) diff --git a/tests/residency_tests.cpp b/tests/residency_tests.cpp index 3ca6f5535c..0edac9b6dc 100644 --- a/tests/residency_tests.cpp +++ b/tests/residency_tests.cpp @@ -1,8 +1,5 @@ // Copyright © 2026 Apple Inc. -#include - -#include #include #include "doctest/doctest.h" @@ -297,36 +294,3 @@ TEST_CASE("test the wired limit is used up to its boundary") { allocator::allocator().free(second); CHECK_EQ(residency().wired_size(), baseline); } - -TEST_CASE("test as_strided on a mapped buffer larger than INT_MAX") { - constexpr size_t length = size_t{131073} * 16384; - if (length > metal::device(Device::gpu).mtl_device()->maxBufferLength()) { - return; - } - LimitGuard guard(0, 0); - auto pointer = mmap( - nullptr, length, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANON, -1, 0); - REQUIRE(pointer != MAP_FAILED); - auto mapping = - std::shared_ptr(pointer, [](void* data) { munmap(data, length); }); - auto buffer = allocator::make_buffer(mapping.get(), length); - REQUIRE(buffer.ptr() != nullptr); - auto data = allocator::Data(buffer, [mapping](allocator::Buffer buffer) { - allocator::release(buffer); - }); - auto storage = array(std::move(data), {131073, 16384}, uint8); - auto bytes = static_cast(mapping.get()); - for (int i = 0; i < 64; ++i) { - bytes[length - 64 + i] = i + 1; - } - for (const auto& input : {storage, copy(storage)}) { - auto active = get_active_memory(); - reset_peak_memory(); - auto result = as_strided(input, {2, 32}, {32, 1}, length - 64); - eval(result); - CHECK_LE(get_peak_memory(), active); - CHECK_EQ(result.data(), bytes + length - 64); - CHECK_EQ(result.data()[0], 1); - CHECK_EQ(result.data()[63], 64); - } -} From 92e1de97ec416fbf63c1ca00d76bf26ebebbb5cb Mon Sep 17 00:00:00 2001 From: Cheng Date: Mon, 19 Oct 2026 13:58:03 +0900 Subject: [PATCH 3/3] nit --- mlx/array.h | 9 ++------- mlx/backend/common/common.cpp | 14 ++------------ mlx/backend/cpu/primitives.cpp | 5 +++-- mlx/backend/gpu/primitives.cpp | 5 +++-- mlx/primitives.h | 2 +- 5 files changed, 11 insertions(+), 24 deletions(-) diff --git a/mlx/array.h b/mlx/array.h index e10807a8cd..59ab941945 100644 --- a/mlx/array.h +++ b/mlx/array.h @@ -566,6 +566,7 @@ T array::item() { throw std::invalid_argument("item can only be called on arrays of size 1."); } eval(); + // Use memcpy to read unaligned data. T value; std::memcpy(&value, data(), sizeof(T)); return value; @@ -573,17 +574,11 @@ T array::item() { 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(); - T value; - std::memcpy(&value, data(), sizeof(T)); - return value; + 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 0474d02fda..1633e0a647 100644 --- a/mlx/backend/cpu/primitives.cpp +++ b/mlx/backend/cpu/primitives.cpp @@ -86,12 +86,13 @@ static std::pair compute_dynamic_offset( } void AsStrided::eval_cpu(const std::vector& inputs, array& out) { + assert(inputs.size() == 1); if (inputs[0].flags().row_contiguous) { - eval(inputs, out); + eval(inputs[0], out); } else { auto copy = contiguous_copy_cpu(inputs[0], stream()); cpu::get_command_encoder(stream()).add_temporary(copy); - eval({copy}, out); + eval(copy, out); } } void Broadcast::eval_cpu(const std::vector& inputs, array& out) { diff --git a/mlx/backend/gpu/primitives.cpp b/mlx/backend/gpu/primitives.cpp index af10a14132..6df6ab3e09 100644 --- a/mlx/backend/gpu/primitives.cpp +++ b/mlx/backend/gpu/primitives.cpp @@ -25,8 +25,9 @@ namespace mlx::core { void AsStrided::eval_gpu(const std::vector& inputs, array& out) { MLX_PROFILER_RANGE("AsStrided::eval_gpu"); + assert(inputs.size() == 1); if (inputs[0].flags().row_contiguous) { - eval(inputs, out); + eval(inputs[0], out); } else { auto copy = contiguous_copy_gpu(inputs[0], stream()); #if defined(MLX_USE_CUDA) @@ -34,7 +35,7 @@ void AsStrided::eval_gpu(const std::vector& inputs, array& out) { #else metal::get_command_encoder(stream()).add_temporary(copy); #endif - eval({copy}, out); + eval(copy, out); } } 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 {