Fix unaligned scalar reads and oversized contiguous buffer views

mlx/array.h

@@ -3,6 +3,7 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <cstring>
 #include <functional>
 #include <memory>
 #include <utility>
@@ -565,20 +566,19 @@ T array::item() {
     throw std::invalid_argument("item can only be called on arrays of size 1.");
   }
   eval();
-  return *data<T>();
+  // Use memcpy to read unaligned data.
+  T value;
+  std::memcpy(&value, data<void>(), sizeof(T));
+  return value;
 }
 
 template <typename T>
 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<array*>(this)->eval();
-  return *data<T>();
+  return const_cast<array*>(this)->item<T>();
 }
 
 template <typename It>

mlx/backend/common/common.cpp

@@ -9,18 +9,8 @@
 
 namespace mlx::core {
 
-void AsStrided::eval(const std::vector<array>& 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;

mlx/backend/cpu/primitives.cpp

@@ -86,7 +86,14 @@ static std::pair<array, bool> compute_dynamic_offset(
 }
 
 void AsStrided::eval_cpu(const std::vector<array>& 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<array>& inputs, array& out) {
   eval(inputs, out);

mlx/backend/gpu/primitives.cpp

@@ -8,6 +8,9 @@
 
 #if defined(MLX_USE_CUDA)
 #include <nvtx3/nvtx3.hpp>
+#include "mlx/backend/cuda/device.h"
+#else
+#include "mlx/backend/metal/device.h"
 #endif
 
 #include <cassert>
@@ -22,7 +25,18 @@ namespace mlx::core {
 
 void AsStrided::eval_gpu(const std::vector<array>& 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<array>& inputs, array& out) {

mlx/ops.cpp

@@ -333,8 +333,7 @@ array as_strided(
       dtype,
       std::make_shared<AsStrided>(
           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 /* = {} */) {

mlx/primitives.h

@@ -443,7 +443,7 @@ class AsStrided : public UnaryPrimitive {
   Strides strides_;
   size_t offset_;
 
-  void eval(const std::vector<array>& inputs, array& out);
+  void eval(const array& in, array& out);
 };
 
 class BitwiseBinary : public UnaryPrimitive {

tests/array_tests.cpp

@@ -1,6 +1,7 @@
 // Copyright © 2023 Apple Inc.
 #include <cassert>
 #include <climits>
+#include <cstring>
 #include <stdexcept>
 #include <vector>
 
@@ -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<bool>());
 }
+
+TEST_CASE("test unaligned scalar item") {
+  uint64_t expected = 0x0123456789abcdef;
+  for (size_t offset = 1; offset < sizeof(expected); ++offset) {
+    std::vector<uint8_t> data(offset + sizeof(expected));
+    std::memcpy(data.data() + offset, &expected, sizeof(expected));
+    auto storage = array(data.data(), {static_cast<int>(data.size())}, uint8);
+    auto scalar = view(as_strided(storage, {8}, {1}, offset), uint64);
+    eval(scalar);
+    if (default_device() == Device::cpu) {
+      CHECK_EQ(scalar.data<void>(), storage.data<uint8_t>() + offset);
+    }
+    CHECK_EQ(scalar.item<uint64_t>(), expected);
+    CHECK_EQ(static_cast<const array&>(scalar).item<uint64_t>(), expected);
+  }
+}

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<bool>());
+
+  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<bool>());
+
+  out = jvp(fun, x, x).second;
+  CHECK(array_equal(out, fun(x)).item<bool>());
 }
 
 TEST_CASE("test jvp from vjp") {

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<bool>());
+
+  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<bool>());
+
+  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<bool>());
 }
 
 TEST_CASE("test sort and scan on empty arrays") {