Merge pull request #22 from ausimian/m12-zero-copy-binary

M12: Zero-copy to_binary via resource binary

Author: ausimian

.github/workflows/ci.yml

@@ -71,3 +71,7 @@ jobs:
       # DistilBERT forward pass is the canonical integration signal
       # that no Nx op on the transformer critical path has regressed.
       - run: mix test --only conformance
+
+  # ASan CI deferred: requires OTP built with --enable-sanitizers=address
+  # (macOS SIP blocks DYLD_INSERT_LIBRARIES, and late-loaded libasan
+  # fails). See Makefile and RELEASE.md for details.

Makefile

@@ -28,6 +28,17 @@ else
     JOBS    := $(shell nproc)
 endif
 
+# Optional: AddressSanitizer build. Set EMILY_ASAN=1 to instrument the
+# NIF. Requires an OTP built with --enable-sanitizers=address so that
+# beam.smp links the ASan runtime at startup (interceptors must install
+# before any allocation). macOS SIP strips DYLD_INSERT_LIBRARIES from
+# processes launched through /bin/sh, and loading libasan late via
+# dlopen fails, so preloading is not an option on stock macOS+OTP.
+ifeq ($(EMILY_ASAN),1)
+    CXXFLAGS += -fsanitize=address -fno-omit-frame-pointer -g -O1
+    LDFLAGS  += -fsanitize=address
+endif
+
 MAKE_JOBS ?= $(JOBS)
 
 .PHONY: all clean bench-native

PLAN.md

@@ -517,40 +517,82 @@ version is a future follow-on.
 enabled and pass conformance; benchmark shows a measurable speedup
 over the M9 baseline (target ≥1.5× on M3 hardware).
 
-### M12 — Zero-copy binary round-trip
+### M12 — Zero-copy binary round-trip (`to_binary`)
 
 PLAN design decision #9 claims unified-memory zero-copy for
 `from_binary` / `to_binary`. The current code memcpys unconditionally
-(`emily_nif.cpp:57-58`, `:74-84`). M12 delivers the claim.
+(`emily_nif.cpp:57-58`, `:74-84`). M12 delivers the claim for
+`to_binary`; `from_binary` is deferred to M12.5 because MLX's Metal
+`allocator::Buffer` stores an `MTL::Buffer*`, not a raw CPU pointer,
+so wrapping a BEAM heap pointer is unsound without routing through
+`MTL::Device::newBufferWithBytesNoCopy`.
 
 - **`to_binary`**: wrap the materialized MLX buffer pointer as a BEAM
   resource binary via `enif_make_resource_binary`, with the resource
   retaining a refcount on the MLX array so the buffer survives until
   the BEAM binary is GC'd. No copy; the BEAM binary aliases MLX
   storage directly.
-- **`from_binary`**: when the input binary is heap-resident and
-  page-aligned (and the MLX array would otherwise live on the unified
-  arena), construct the array with a deleter that releases the BEAM
-  binary refcount instead of freeing. Fall back to the current memcpy
-  path when alignment doesn't hold or when the source is a sub-binary.
+- **`from_binary`**: deferred. See M12.5.
 - **Stride-aware materialize**: `to_binary` currently routes through
   `mx::contiguous`; for already-contiguous arrays this is a no-op, but
   the wrap-as-resource path needs an explicit guard since aliasing a
   non-contiguous buffer would lie about its layout.
 
 **Testing**:
-- Allocate a 256 MB tensor, round-trip through `to_binary` then
-  `from_binary`, assert MLX active memory grew by ~256 MB not ~512 MB.
-- Soak: repeated round-trip with cache-clear, assert peak memory is
-  bounded by the working-set size, not 2× it.
+- Allocate a tensor, call `to_binary`, assert MLX active memory did
+  not grow (aliasing, not copying).
+- Soak: repeated `to_binary` with cache-clear, assert peak memory is
+  bounded by the working-set size.
 - Correctness: the M2 property suite must still pass — this is a perf
   change, not a semantics change.
-- Refcount safety: drop the MLX-side reference, then read the BEAM
-  binary; must not segfault. Use-after-free is the failure mode, so
-  this milestone gates on an AddressSanitizer build in CI.
+- Refcount safety: drop the original tensor reference, then read the
+  BEAM binary returned by `to_binary`; must not segfault. Use-after-
+  free is the failure mode, so this milestone gates on an
+  AddressSanitizer build in CI.
+
+**Exit:** `to_binary` zero-copy verified by allocator stats; M2
+property suite green; lifecycle and soak tests verify refcount
+safety. AddressSanitizer CI deferred (macOS SIP prevents
+`DYLD_INSERT_LIBRARIES` propagation through `/bin/sh`-launched BEAM;
+requires a custom `--enable-sanitizers=address` OTP build).
+`EMILY_ASAN=1` Makefile flag ships for users with sanitizer-enabled
+OTP.
+
+### M12.5 — `from_binary` zero-copy via MTL no-copy buffer
+
+Deferred half of M12. The BEAM → MLX direction can't be done by
+wrapping a BEAM pointer as an `mx::allocator::Buffer` — Metal's
+allocator stores `MTL::Buffer*`, not a raw CPU pointer, so a wrapped
+heap pointer would be dereferenced as an `MTL::Buffer*` and crash on
+GPU dispatch. True zero-copy requires registering the BEAM memory
+with Metal.
+
+- **NIF changes**: accept `fine::Term` instead of `ErlNifBinary` so
+  we can `enif_make_copy` the term into a persistent `ErlNifEnv` and
+  keep the refc binary alive. Call
+  `MTL::Device::newBufferWithBytesNoCopy:length:options:deallocator:`
+  to hand the BEAM pointer to Metal; the deallocator block calls
+  `enif_free_env`. Wrap the resulting `MTL::Buffer*` as an
+  `allocator::Buffer` and pass to the 4-arg `mx::array` constructor.
+- **MLX integration**: register the MTL::Buffer with MLX's
+  residency set so command buffers keep it resident. The residency
+  API is not in public headers — either upstream a public entry
+  point or bypass via implementation-detail APIs.
+- **Build**: link the Metal framework from the NIF; add metal-cpp
+  headers to the build.
+- **Pre-conditions**: heap-resident refc binary, page-aligned,
+  page-sized. Fall back to the M12 memcpy path otherwise.
 
-**Exit:** zero-copy verified by allocator stats; M2 property suite
-green; AddressSanitizer build clean.
+**Testing**:
+- Allocate a 256 MB page-aligned binary, `from_binary` →
+  `to_binary`, assert MLX active memory grew by ~256 MB not ~512 MB.
+- Weight-loading soak: simulate Bumblebee's mmap'd-weights path with
+  realistic tensor counts, assert peak memory matches the
+  theoretical working-set lower bound.
+- Refcount safety under ASan, same pattern as M12.
+
+**Exit:** `from_binary` zero-copy verified by allocator stats for
+page-aligned inputs; PLAN decision #9 claim fully delivered.
 
 ### M13 — EXLA gradient conformance
 

RELEASE.md

@@ -2,6 +2,56 @@
 
 ## Added
 
+- M12 — Zero-copy `to_binary`. `Emily.to_binary/1` (and everything
+  that routes through `Nx.to_binary` on the Emily backend) now
+  returns a BEAM resource binary that aliases the MLX buffer
+  directly, instead of memcpying the bytes into a fresh BEAM binary.
+  The resource binary's lifetime pins a fresh `Tensor` resource so
+  the underlying MLX storage survives until the binary is GC'd.
+  Savings are most visible when handing large tensors back to Nx
+  (logits from inference, weight exports): one memcpy eliminated
+  per call.
+  - **NIF change** (`c_src/emily_nif.cpp`): `to_binary` now returns
+    `fine::Term` via `fine::make_resource_binary`. Defensive assert
+    on `row_contiguous` after `mx::contiguous` guards against
+    aliasing a strided buffer.
+  - **`from_binary` unchanged.** BEAM → MLX zero-copy on Metal
+    requires `MTL::Device::newBufferWithBytesNoCopy` (MLX's
+    `allocator::Buffer` stores an `MTL::Buffer*`, not a raw CPU
+    pointer — wrapping a BEAM heap pointer is unsound). Deferred
+    to M12.5; see `PLAN.md`.
+  - **Tests**: round-trip lifetime test at `test/emily_test.exs`
+    (`"to_binary aliased binary survives after tensor goes out of
+    scope"`); zero-copy memory soak at
+    `test/soak/zero_copy_roundtrip_test.exs` (asserts MLX active
+    memory and BEAM binary heap both stay flat across 200 round
+    trips). M2 property suite still green — semantics unchanged.
+  - **Build**: `EMILY_ASAN=1` env var enables an AddressSanitizer
+    build of the NIF (`Makefile`). Requires an OTP built with
+    `--enable-sanitizers=address` so beam.smp links the ASan runtime
+    at startup — macOS SIP strips `DYLD_INSERT_LIBRARIES` from
+    processes launched through `/bin/sh` (which `erl`/`elixir` use),
+    and loading libasan late (via dlopen of the NIF) fails because
+    the malloc/free interceptors must be installed before any
+    allocation. With a sanitizer-enabled OTP:
+    ```
+    EMILY_ASAN=1 mix compile --force
+    ASAN_OPTIONS=detect_leaks=0:abort_on_error=1 mix test
+    ```
+    No `DYLD_INSERT_LIBRARIES` or `ERL_FLAGS` needed; beam.smp
+    already has the runtime linked. A CI job for this is deferred
+    until the custom OTP build cost is justified; the lifecycle
+    and soak tests provide empirical refcount-safety coverage.
+  - **Behavioral change**: `to_binary` returns a resource binary
+    aliasing MLX storage. BEAM's binary-vheap GC heuristics don't
+    account for the aliased external data (the ProcBin is ~64 bytes
+    regardless of the underlying MLX buffer size). In tight loops
+    calling `to_binary`, resource binaries can accumulate without
+    triggering collection, holding MLX memory longer than expected.
+    Callers in hot paths should trigger
+    `:erlang.garbage_collect/0` periodically or ensure the binary
+    escapes to a short-lived process where GC runs naturally.
+
 - M11 — MLX fused transformer kernels. Wires MLX's handwritten
   `mx::fast::*` fused kernels (RMSNorm, LayerNorm, RoPE, scaled-dot-
   product attention) into Emily as `defn`-callable helpers, and ships

c_src/emily_nif.cpp

@@ -29,6 +29,10 @@ FINE_RESOURCE(Tensor);
 // from_binary/3 — build a lazy MLX array from a BEAM binary.
 // Regular scheduler: MLX copies the buffer into its own storage during
 // construction, so this is cheap and bounded.
+//
+// BEAM→MLX zero-copy is not possible with the current allocator API:
+// on Metal, allocator::Buffer stores an MTL::Buffer*, so wrapping a
+// BEAM heap pointer would crash on GPU dispatch.
 fine::ResourcePtr<Tensor> from_binary(
     ErlNifEnv *,
     ErlNifBinary data,
@@ -50,9 +54,9 @@ fine::ResourcePtr<Tensor> from_binary(
         " got " + std::to_string(data.size));
   }
 
-  // MLX has no void*-accepting array constructor. The canonical path
-  // is: allocate an MLX-owned buffer, memcpy into it, hand ownership
-  // to the array with a matching deleter.
+  // Allocate an MLX-owned buffer, memcpy into it, hand ownership to
+  // the array with a matching deleter. See comment above for why we
+  // don't alias the BEAM binary directly.
   auto buf = mx::allocator::malloc(expected);
   std::memcpy(buf.raw_ptr(), data.data, expected);
   auto deleter = [](mx::allocator::Buffer b) { mx::allocator::free(b); };
@@ -62,26 +66,33 @@ fine::ResourcePtr<Tensor> from_binary(
 }
 FINE_NIF(from_binary, 0);
 
-// to_binary/1 — materialize the array and return its bytes as a binary.
-// Dirty CPU: eval() triggers kernel launch and waits for completion.
+// to_binary/1 — materialize the array and return its bytes as a BEAM
+// resource binary aliasing MLX storage (no memcpy). The binary pins a
+// Tensor resource → mx::array → MLX buffer; the buffer survives until
+// the BEAM binary is GC'd.
 //
-// We route through mx::contiguous() first so views with non-standard
-// strides (transpose, slice, swapaxes, broadcast_to) produce the
-// correct in-memory layout. For already-contiguous arrays MLX elides
-// the copy. A handful of MLX ops (notably cumulative reductions on
-// interior axes of some 4-D shapes) raise "Unable to safely factor
-// shape" here; the Backend layer routes the known cases around us.
-std::string to_binary(ErlNifEnv *, fine::ResourcePtr<Tensor> tensor) {
+// We route through mx::contiguous() so views with non-standard strides
+// produce the correct in-memory layout. A handful of MLX ops (notably
+// cumulative reductions on interior axes of some 4-D shapes) raise
+// "Unable to safely factor shape" here; the Backend layer routes the
+// known cases around us.
+fine::Term to_binary(ErlNifEnv *env, fine::ResourcePtr<Tensor> tensor) {
   auto materialized = mx::contiguous(tensor->array);
   mx::eval(materialized);
 
-  const void *src = materialized.data<void>();
-  size_t nbytes = materialized.nbytes();
+  // Defensive: mx::contiguous is supposed to give a row-contiguous
+  // layout. If it ever doesn't, we'd be aliasing a strided buffer and
+  // lying about its layout. Throw rather than silently corrupt.
+  if (!materialized.flags().row_contiguous) {
+    throw std::runtime_error(
+        "to_binary: array is not row-contiguous after mx::contiguous");
+  }
+
+  auto nbytes = materialized.nbytes();
+  auto data = reinterpret_cast<const char *>(materialized.data<void>());
 
-  std::string out;
-  out.resize(nbytes);
-  std::memcpy(out.data(), src, nbytes);
-  return out;
+  auto pin = wrap(std::move(materialized));
+  return fine::make_resource_binary(env, std::move(pin), data, nbytes);
 }
 FINE_NIF(to_binary, ERL_NIF_DIRTY_JOB_CPU_BOUND);
 

test/emily_test.exs

@@ -97,5 +97,24 @@ defmodule EmilyTest do
 
       assert Emily.to_binary(t) == bin
     end
+
+    test "to_binary aliased binary survives after tensor goes out of scope" do
+      # to_binary returns a resource binary that aliases MLX storage.
+      # If the refcount wiring is wrong, reading after the tensor is
+      # dropped would segfault or return garbage.
+      bin =
+        for x <- 1..1024, into: <<>>, do: <<x * 1.0::float-32-native>>
+
+      out =
+        (fn ->
+           t = Emily.from_binary(bin, [1024], {:f, 32})
+           Emily.to_binary(t)
+         end).()
+
+      :erlang.garbage_collect()
+      :erlang.garbage_collect(self())
+
+      assert out == bin
+    end
   end
 end

test/soak/memory_test.exs

@@ -18,6 +18,12 @@ defmodule Emily.Soak.MemoryTest do
       passes, but the test fails once the cumulative leak exceeds
       ~2 retained tensors total — i.e. a leak rate above ~0.1 %.
 
+  `Native.to_binary/1` returns a resource binary aliasing MLX memory.
+  The ProcBin is tiny (~64 B) so BEAM's binary-vheap GC doesn't fire
+  based on the external data size; resource binaries can accumulate
+  and pin MLX buffers. We GC every @gc_every iterations so the test
+  reflects steady-state usage, not transient accumulation.
+
   A tight tolerance matters — a generous one makes the test pass even
   when refcounts are broken.
   """
@@ -34,6 +40,11 @@ defmodule Emily.Soak.MemoryTest do
   @iters 2_000
   @warmup 20
 
+  # See moduledoc: periodic GC keeps to_binary's aliased resource
+  # binaries from accumulating. Every 10 iters ≈ 10 MB of MLX-pinned
+  # memory at the peak, well within the tolerance window.
+  @gc_every 10
+
   # Measured drift under the full suite is ~2 MB (Metal pool keeps
   # ~2 tensors live past clear_cache). 4 MB tolerance absorbs that
   # without masking a leak of more than a couple of retained tensors;
@@ -59,7 +70,10 @@ defmodule Emily.Soak.MemoryTest do
     baseline = Native.get_active_memory()
     Native.reset_peak_memory()
 
-    for _ <- 1..@iters, do: workload(data)
+    for i <- 1..@iters do
+      workload(data)
+      if rem(i, @gc_every) == 0, do: :erlang.garbage_collect()
+    end
 
     :erlang.garbage_collect()
     Native.clear_cache()