Skip to content

Commit 2f9414d

Browse files
committed
M12: Zero-copy to_binary via resource binary
to_binary now returns a BEAM resource binary aliasing the MLX Metal buffer directly, eliminating the memcpy that previously copied every tensor's bytes into a BEAM-heap binary. The resource binary pins a fresh Tensor resource so the MLX buffer survives until the binary is GC'd. from_binary zero-copy is deferred to M12.5 — MLX's allocator::Buffer stores an MTL::Buffer* on Metal, so wrapping a BEAM heap pointer is unsound without MTL::Device::newBufferWithBytesNoCopy. Also adds EMILY_ASAN=1 Makefile flag (requires sanitizer-enabled OTP) and a deferred ASan CI note (macOS SIP blocks DYLD_INSERT_LIBRARIES).
1 parent e70d5c8 commit 2f9414d

8 files changed

Lines changed: 305 additions & 35 deletions

File tree

.github/workflows/ci.yml

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,3 +71,7 @@ jobs:
7171
# DistilBERT forward pass is the canonical integration signal
7272
# that no Nx op on the transformer critical path has regressed.
7373
- run: mix test --only conformance
74+
75+
# ASan CI deferred: requires OTP built with --enable-sanitizers=address
76+
# (macOS SIP blocks DYLD_INSERT_LIBRARIES, and late-loaded libasan
77+
# fails). See Makefile and RELEASE.md for details.

Makefile

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,17 @@ else
2828
JOBS := $(shell nproc)
2929
endif
3030

31+
# Optional: AddressSanitizer build. Set EMILY_ASAN=1 to instrument the
32+
# NIF. Requires an OTP built with --enable-sanitizers=address so that
33+
# beam.smp links the ASan runtime at startup (interceptors must install
34+
# before any allocation). macOS SIP strips DYLD_INSERT_LIBRARIES from
35+
# processes launched through /bin/sh, and loading libasan late via
36+
# dlopen fails, so preloading is not an option on stock macOS+OTP.
37+
ifeq ($(EMILY_ASAN),1)
38+
CXXFLAGS += -fsanitize=address -fno-omit-frame-pointer -g -O1
39+
LDFLAGS += -fsanitize=address
40+
endif
41+
3142
MAKE_JOBS ?= $(JOBS)
3243

3344
.PHONY: all clean bench-native

PLAN.md

Lines changed: 58 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -517,40 +517,82 @@ version is a future follow-on.
517517
enabled and pass conformance; benchmark shows a measurable speedup
518518
over the M9 baseline (target ≥1.5× on M3 hardware).
519519

520-
### M12 — Zero-copy binary round-trip
520+
### M12 — Zero-copy binary round-trip (`to_binary`)
521521

522522
PLAN design decision #9 claims unified-memory zero-copy for
523523
`from_binary` / `to_binary`. The current code memcpys unconditionally
524-
(`emily_nif.cpp:57-58`, `:74-84`). M12 delivers the claim.
524+
(`emily_nif.cpp:57-58`, `:74-84`). M12 delivers the claim for
525+
`to_binary`; `from_binary` is deferred to M12.5 because MLX's Metal
526+
`allocator::Buffer` stores an `MTL::Buffer*`, not a raw CPU pointer,
527+
so wrapping a BEAM heap pointer is unsound without routing through
528+
`MTL::Device::newBufferWithBytesNoCopy`.
525529

526530
- **`to_binary`**: wrap the materialized MLX buffer pointer as a BEAM
527531
resource binary via `enif_make_resource_binary`, with the resource
528532
retaining a refcount on the MLX array so the buffer survives until
529533
the BEAM binary is GC'd. No copy; the BEAM binary aliases MLX
530534
storage directly.
531-
- **`from_binary`**: when the input binary is heap-resident and
532-
page-aligned (and the MLX array would otherwise live on the unified
533-
arena), construct the array with a deleter that releases the BEAM
534-
binary refcount instead of freeing. Fall back to the current memcpy
535-
path when alignment doesn't hold or when the source is a sub-binary.
535+
- **`from_binary`**: deferred. See M12.5.
536536
- **Stride-aware materialize**: `to_binary` currently routes through
537537
`mx::contiguous`; for already-contiguous arrays this is a no-op, but
538538
the wrap-as-resource path needs an explicit guard since aliasing a
539539
non-contiguous buffer would lie about its layout.
540540

541541
**Testing**:
542-
- Allocate a 256 MB tensor, round-trip through `to_binary` then
543-
`from_binary`, assert MLX active memory grew by ~256 MB not ~512 MB.
544-
- Soak: repeated round-trip with cache-clear, assert peak memory is
545-
bounded by the working-set size, not 2× it.
542+
- Allocate a tensor, call `to_binary`, assert MLX active memory did
543+
not grow (aliasing, not copying).
544+
- Soak: repeated `to_binary` with cache-clear, assert peak memory is
545+
bounded by the working-set size.
546546
- Correctness: the M2 property suite must still pass — this is a perf
547547
change, not a semantics change.
548-
- Refcount safety: drop the MLX-side reference, then read the BEAM
549-
binary; must not segfault. Use-after-free is the failure mode, so
550-
this milestone gates on an AddressSanitizer build in CI.
548+
- Refcount safety: drop the original tensor reference, then read the
549+
BEAM binary returned by `to_binary`; must not segfault. Use-after-
550+
free is the failure mode, so this milestone gates on an
551+
AddressSanitizer build in CI.
552+
553+
**Exit:** `to_binary` zero-copy verified by allocator stats; M2
554+
property suite green; lifecycle and soak tests verify refcount
555+
safety. AddressSanitizer CI deferred (macOS SIP prevents
556+
`DYLD_INSERT_LIBRARIES` propagation through `/bin/sh`-launched BEAM;
557+
requires a custom `--enable-sanitizers=address` OTP build).
558+
`EMILY_ASAN=1` Makefile flag ships for users with sanitizer-enabled
559+
OTP.
560+
561+
### M12.5 — `from_binary` zero-copy via MTL no-copy buffer
562+
563+
Deferred half of M12. The BEAM → MLX direction can't be done by
564+
wrapping a BEAM pointer as an `mx::allocator::Buffer` — Metal's
565+
allocator stores `MTL::Buffer*`, not a raw CPU pointer, so a wrapped
566+
heap pointer would be dereferenced as an `MTL::Buffer*` and crash on
567+
GPU dispatch. True zero-copy requires registering the BEAM memory
568+
with Metal.
569+
570+
- **NIF changes**: accept `fine::Term` instead of `ErlNifBinary` so
571+
we can `enif_make_copy` the term into a persistent `ErlNifEnv` and
572+
keep the refc binary alive. Call
573+
`MTL::Device::newBufferWithBytesNoCopy:length:options:deallocator:`
574+
to hand the BEAM pointer to Metal; the deallocator block calls
575+
`enif_free_env`. Wrap the resulting `MTL::Buffer*` as an
576+
`allocator::Buffer` and pass to the 4-arg `mx::array` constructor.
577+
- **MLX integration**: register the MTL::Buffer with MLX's
578+
residency set so command buffers keep it resident. The residency
579+
API is not in public headers — either upstream a public entry
580+
point or bypass via implementation-detail APIs.
581+
- **Build**: link the Metal framework from the NIF; add metal-cpp
582+
headers to the build.
583+
- **Pre-conditions**: heap-resident refc binary, page-aligned,
584+
page-sized. Fall back to the M12 memcpy path otherwise.
551585

552-
**Exit:** zero-copy verified by allocator stats; M2 property suite
553-
green; AddressSanitizer build clean.
586+
**Testing**:
587+
- Allocate a 256 MB page-aligned binary, `from_binary`
588+
`to_binary`, assert MLX active memory grew by ~256 MB not ~512 MB.
589+
- Weight-loading soak: simulate Bumblebee's mmap'd-weights path with
590+
realistic tensor counts, assert peak memory matches the
591+
theoretical working-set lower bound.
592+
- Refcount safety under ASan, same pattern as M12.
593+
594+
**Exit:** `from_binary` zero-copy verified by allocator stats for
595+
page-aligned inputs; PLAN decision #9 claim fully delivered.
554596

555597
### M13 — EXLA gradient conformance
556598

RELEASE.md

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,56 @@
22

33
## Added
44

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

c_src/emily_nif.cpp

Lines changed: 29 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@ FINE_RESOURCE(Tensor);
2929
// from_binary/3 — build a lazy MLX array from a BEAM binary.
3030
// Regular scheduler: MLX copies the buffer into its own storage during
3131
// construction, so this is cheap and bounded.
32+
//
33+
// BEAM→MLX zero-copy is not possible with the current allocator API:
34+
// on Metal, allocator::Buffer stores an MTL::Buffer*, so wrapping a
35+
// BEAM heap pointer would crash on GPU dispatch.
3236
fine::ResourcePtr<Tensor> from_binary(
3337
ErlNifEnv *,
3438
ErlNifBinary data,
@@ -50,9 +54,9 @@ fine::ResourcePtr<Tensor> from_binary(
5054
" got " + std::to_string(data.size));
5155
}
5256

53-
// MLX has no void*-accepting array constructor. The canonical path
54-
// is: allocate an MLX-owned buffer, memcpy into it, hand ownership
55-
// to the array with a matching deleter.
57+
// Allocate an MLX-owned buffer, memcpy into it, hand ownership to
58+
// the array with a matching deleter. See comment above for why we
59+
// don't alias the BEAM binary directly.
5660
auto buf = mx::allocator::malloc(expected);
5761
std::memcpy(buf.raw_ptr(), data.data, expected);
5862
auto deleter = [](mx::allocator::Buffer b) { mx::allocator::free(b); };
@@ -62,26 +66,33 @@ fine::ResourcePtr<Tensor> from_binary(
6266
}
6367
FINE_NIF(from_binary, 0);
6468

65-
// to_binary/1 — materialize the array and return its bytes as a binary.
66-
// Dirty CPU: eval() triggers kernel launch and waits for completion.
69+
// to_binary/1 — materialize the array and return its bytes as a BEAM
70+
// resource binary aliasing MLX storage (no memcpy). The binary pins a
71+
// Tensor resource → mx::array → MLX buffer; the buffer survives until
72+
// the BEAM binary is GC'd.
6773
//
68-
// We route through mx::contiguous() first so views with non-standard
69-
// strides (transpose, slice, swapaxes, broadcast_to) produce the
70-
// correct in-memory layout. For already-contiguous arrays MLX elides
71-
// the copy. A handful of MLX ops (notably cumulative reductions on
72-
// interior axes of some 4-D shapes) raise "Unable to safely factor
73-
// shape" here; the Backend layer routes the known cases around us.
74-
std::string to_binary(ErlNifEnv *, fine::ResourcePtr<Tensor> tensor) {
74+
// We route through mx::contiguous() so views with non-standard strides
75+
// produce the correct in-memory layout. A handful of MLX ops (notably
76+
// cumulative reductions on interior axes of some 4-D shapes) raise
77+
// "Unable to safely factor shape" here; the Backend layer routes the
78+
// known cases around us.
79+
fine::Term to_binary(ErlNifEnv *env, fine::ResourcePtr<Tensor> tensor) {
7580
auto materialized = mx::contiguous(tensor->array);
7681
mx::eval(materialized);
7782

78-
const void *src = materialized.data<void>();
79-
size_t nbytes = materialized.nbytes();
83+
// Defensive: mx::contiguous is supposed to give a row-contiguous
84+
// layout. If it ever doesn't, we'd be aliasing a strided buffer and
85+
// lying about its layout. Throw rather than silently corrupt.
86+
if (!materialized.flags().row_contiguous) {
87+
throw std::runtime_error(
88+
"to_binary: array is not row-contiguous after mx::contiguous");
89+
}
90+
91+
auto nbytes = materialized.nbytes();
92+
auto data = reinterpret_cast<const char *>(materialized.data<void>());
8093

81-
std::string out;
82-
out.resize(nbytes);
83-
std::memcpy(out.data(), src, nbytes);
84-
return out;
94+
auto pin = wrap(std::move(materialized));
95+
return fine::make_resource_binary(env, std::move(pin), data, nbytes);
8596
}
8697
FINE_NIF(to_binary, ERL_NIF_DIRTY_JOB_CPU_BOUND);
8798

test/emily_test.exs

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,5 +97,24 @@ defmodule EmilyTest do
9797

9898
assert Emily.to_binary(t) == bin
9999
end
100+
101+
test "to_binary aliased binary survives after tensor goes out of scope" do
102+
# to_binary returns a resource binary that aliases MLX storage.
103+
# If the refcount wiring is wrong, reading after the tensor is
104+
# dropped would segfault or return garbage.
105+
bin =
106+
for x <- 1..1024, into: <<>>, do: <<x * 1.0::float-32-native>>
107+
108+
out =
109+
(fn ->
110+
t = Emily.from_binary(bin, [1024], {:f, 32})
111+
Emily.to_binary(t)
112+
end).()
113+
114+
:erlang.garbage_collect()
115+
:erlang.garbage_collect(self())
116+
117+
assert out == bin
118+
end
100119
end
101120
end

test/soak/memory_test.exs

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,12 @@ defmodule Emily.Soak.MemoryTest do
1818
passes, but the test fails once the cumulative leak exceeds
1919
~2 retained tensors total — i.e. a leak rate above ~0.1 %.
2020
21+
`Native.to_binary/1` returns a resource binary aliasing MLX memory.
22+
The ProcBin is tiny (~64 B) so BEAM's binary-vheap GC doesn't fire
23+
based on the external data size; resource binaries can accumulate
24+
and pin MLX buffers. We GC every @gc_every iterations so the test
25+
reflects steady-state usage, not transient accumulation.
26+
2127
A tight tolerance matters — a generous one makes the test pass even
2228
when refcounts are broken.
2329
"""
@@ -34,6 +40,11 @@ defmodule Emily.Soak.MemoryTest do
3440
@iters 2_000
3541
@warmup 20
3642

43+
# See moduledoc: periodic GC keeps to_binary's aliased resource
44+
# binaries from accumulating. Every 10 iters ≈ 10 MB of MLX-pinned
45+
# memory at the peak, well within the tolerance window.
46+
@gc_every 10
47+
3748
# Measured drift under the full suite is ~2 MB (Metal pool keeps
3849
# ~2 tensors live past clear_cache). 4 MB tolerance absorbs that
3950
# without masking a leak of more than a couple of retained tensors;
@@ -59,7 +70,10 @@ defmodule Emily.Soak.MemoryTest do
5970
baseline = Native.get_active_memory()
6071
Native.reset_peak_memory()
6172

62-
for _ <- 1..@iters, do: workload(data)
73+
for i <- 1..@iters do
74+
workload(data)
75+
if rem(i, @gc_every) == 0, do: :erlang.garbage_collect()
76+
end
6377

6478
:erlang.garbage_collect()
6579
Native.clear_cache()

0 commit comments

Comments
 (0)