RFC: Read-only Metal storage export view for evaluated arrays

[robbiemu]

I am working on MetalGraph, a small graph-oriented runtime experiment around Metal resources. One interop path that would be useful for projects like this is the ability to borrow the backing Metal storage of an evaluated MLX array for read-only external Metal use.

I am not asking for broad mutable sharing, hidden synchronization, or a general-purpose “zero-copy” promise. The narrower request is a supported C++/Metal API that can expose a retained read-only view of an evaluated array’s backing Metal storage, when MLX can safely provide one.

The intended contract would be:

MLX produces/evaluates an array.
The caller obtains a retained Metal storage view.
External Metal code may read from that storage.
External writes back into MLX-owned storage are not supported by this API.

This avoids the hardest synchronization and mutation problem. The MLX Metal allocator appears to use shared storage with untracked hazard tracking, so external runtimes should not assume implicit cross-command-queue hazard management. See allocator.cpp. MLX also documents lazy evaluation and explicit evaluation/synchronization concepts, so the API should be clear about whether it fails, waits, or materializes before export. See the MLX lazy evaluation docs.

The proposed API would guarantee only:

storage identity
byte offset
byte span
Metal device identity
retained lifetime for the exported view
read-only external Metal use

It would not guarantee safe external writes into MLX-owned storage.

Proposed API shape

The API could add a small C++ class that privately retains enough MLX ownership to keep the backing storage alive for the lifetime of the view. The retained owner could be the source array itself, or a smaller internal owner if maintainers prefer that.

Conceptually, the view would expose facts already tracked by mlx::core::array, such as:

MTL::Buffer*
MTL::Device*
byte offset
logical byte length
storage span byte length
dtype
shape
strides
contiguity flags

array::data_size() already appears to represent the underlying data span in elements and may differ from the logical array size for broadcasted or irregularly-strided arrays. See array.h.

One possible API shape:

namespace mlx::core::metal {

enum class StorageExportStatus {
  success,
  metal_not_available,
  array_not_available,
  no_storage,
  unsupported_backend,
  unsupported_layout,
  invalid_offset,
};

struct StorageExportOptions {
  bool wait_until_available = false;
  bool require_row_contiguous = true;
};

class MetalStorageView {
 public:
  MTL::Buffer* buffer() const;
  MTL::Device* device() const;

  size_t byte_offset() const;
  size_t logical_nbytes() const;
  size_t storage_span_nbytes() const;

  Dtype dtype() const;
  const Shape& shape() const;
  const Strides& strides() const;

  bool contiguous() const;
  bool row_contiguous() const;
  bool col_contiguous() const;
};

MLX_API StorageExportStatus try_export_metal_storage(
    const array& a,
    MetalStorageView* out,
    const StorageExportOptions& options = {});

MLX_API MetalStorageView export_metal_storage(
    const array& a,
    const StorageExportOptions& options = {});

} // namespace mlx::core::metal

The exact naming and placement are not important to me. I am mainly trying to describe the contract.

Lazy evaluation behavior

The export API should not silently materialize lazy MLX expressions by default.

If the source mlx::core::array is not currently available/materialized, try_export_metal_storage(...) should return something like:

array_not_available

An option such as wait_until_available = true could allow the caller to explicitly request MLX-side waiting/materialization before export, using whatever mechanism maintainers consider appropriate.

Defaulting to “fail unless already available” would avoid surprising evaluation side effects and keep storage export separate from computation scheduling.

Layout behavior

The API could require row-contiguous layout by default.

If an array is not row-contiguous, the export call could return something like:

unsupported_layout

That kind of explicit failure is useful for external runtimes, because layout support is exactly where accidental “zero-copy” claims can become misleading.

The API could still report shape and stride metadata so that future versions, or more advanced callers, can intentionally support additional layouts later.

Possible PR scope

A PR might touch something like:

mlx/backend/metal/storage_view.h          new public header
mlx/backend/metal/storage_view.cpp        new implementation
mlx/backend/metal/CMakeLists.txt          add source/header if needed
docs/src/dev/extensions.rst               add “Borrowing Metal storage” section
tests/backend/metal/test_storage_view.cpp new tests, or equivalent test location

Possibly:

mlx/backend/metal/metal.h

if maintainers prefer exporting the declaration from an existing Metal backend umbrella header.

I would not add Python bindings in the PR. A C++/Metal API would already be enough for extension authors and would keep the initial scope smaller.

Non-goals

This RFC is not proposing:

Python raw pointer APIs
external writes into MLX-owned storage
automatic cross-runtime synchronization
support for arbitrary strided layouts by default
a dependency on MetalGraph
changes to MLX evaluation semantics
fallback copies hidden behind an export call

The goal is only to provide a supported, conservative storage export contract for advanced Metal interop.

Question for maintainers

Would this kind of read-only MetalStorageView API fit MLX’s design?

If so, I would be interested in preparing a small PR limited to that scope: evaluated arrays only, read-only external use, explicit failure for unavailable or unsupported layouts, retained lifetime through the view object, and documentation that synchronization remains caller-managed.

[robbiemu]

in case you are interested in MetalGraph, I have a couple of blog posts on it (the latter related to this RFC in part):

https://selfenrichment.hashnode.dev/the-honest-limits-reached-with-metalgraph
https://selfenrichment.hashnode.dev/revisiting-metalgraph-what-i-learned-about-mlx-replay-and-sidecar-performance