1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
//! ADR-040 §0.21c-track2 — regression guard for the concurrent-encoder
//! ordering ROOT FIX (encoder retain, mlx-native 80a58de).
//!
//! ## The bug this guards against
//!
//! mlx-native drives a SINGLE persistent `MTLDispatchTypeConcurrent` compute
//! encoder (`encoder.rs::get_or_create_encoder`). metal-rs returns that encoder
//! as a +0 autoreleased object (gfx-rs/metal-rs#128). Before the fix, the
//! engine held it through a borrowed `*const ComputeCommandEncoderRef`
//! (`encoder.rs:757 active_encoder`) WITHOUT taking ownership. An autorelease-
//! pool drain between two dependent dispatches could drop the ownership state
//! Metal's barrier ordering keys on, so `memoryBarrierWithScope:Buffers` no
//! longer reliably ordered a slow producer (e.g. the Q6_K mvN lm_head matmul)
//! before its in-place consumer (softcap). The faster nr2 producer won the race
//! by timing; mvN lost it → ~1/3 of runtime-compile runs diverged. No crash
//! because the command buffer also references the encoder (no UAF), so the
//! symptom was silent data-race non-determinism, NOT a panic.
//!
//! llama.cpp uses the IDENTICAL primitive (concurrent encoder +
//! `[encoder memoryBarrierWithScope:MTLBarrierScopeBuffers]`,
//! ggml-metal-device.m:512) and works — because it RETAINS the encoder
//! (`[res->obj retain]`, released at end). The fix matches that: retain in
//! `get_or_create_encoder`, balanced release in `end_active_encoder`, and a
//! leak-safe `reset_command_buffer` that routes through `end_active_encoder`
//! (review Finder B, mlx-native 3e7b030).
//!
//! ## Why this is a STRUCTURAL guard + a functional round-trip, not a
//! timing-race reproduction
//!
//! A faithful timing-race reproduction in a unit test is not reliable: the
//! race window depends on a real multi-stage decode workload (the slow Q6_K
//! mvN lm_head over vocab=262144 vs an in-place softcap) plus an autorelease
//! drain landing in exactly the wrong place. An isolated single-session writer
//! →barrier→reader with a per-iteration full sync does NOT open that window —
//! an earlier attempt at a pure timing test passed WITH and WITHOUT the fix and
//! was therefore deleted (it gave false confidence). The codebase convention
//! for "the fix is real but the failure can't be force-injected from a unit
//! test" is a source-structural regression guard plus a functional GPU round-
//! trip of the fixed path (see `encoder_session_lifecycle.rs` Test 6). This
//! file follows that convention:
//!
//! 1. `retain_root_fix_source_invariants` — reads `encoder.rs` and asserts the
//! three load-bearing edits are present (retain, balanced release, reset
//! routed through end_active_encoder). A revert of any one fails CI. THIS is
//! the durable regression catch.
//! 2. `writer_barrier_reader_orders_through_retained_encoder` — a cheap
//! functional SMOKE of the writer→barrier→reader path on the concurrent
//! encoder. HONEST LIMIT: it does NOT reproduce the race (verified — it
//! passes even with the retain reverted; an isolated add→mul edge with a
//! per-iteration full-sync does not open the window, same as the earlier
//! deleted timing test). It guards the happy path only. The real
//! fail-direction discriminator is (1) above plus the production model-
//! parity oracle `slot_aware_n8_per_slot_parity_vs_serial` under
//! `MLX_PRECOMPILED_METALLIB=0` (which ran the actual slow Q6_K mvN lm_head
//! vs in-place softcap and reproduced ~1/3 without the fix).
//!
//! ## CI recommendation
//!
//! The durable regression catch for the encoder-ordering fix is the structural
//! guard here PLUS running the hf2q model-parity test on the runtime-source-
//! compile path: `MLX_PRECOMPILED_METALLIB=0 ... slot_aware_n8_per_slot_parity\
//! _vs_serial`. That path is the only one that exercised the real race; the
//! -O3 metallib masks it. (Note: that parity test is currently confounded by a
//! separate, pre-existing batched-flash-attention non-determinism, §0.19 /
//! task #19 — kernel-independent — so it is not yet a clean 0-flake gate.)
#![allow(clippy::expect_used, clippy::unwrap_used)]
use mlx_native::{DType, GraphExecutor, KernelRegistry, MlxDevice};
/// Structural regression guard — the THREE load-bearing edits of the retain
/// root fix (80a58de) + the leak-safe reset (3e7b030) must remain in
/// `encoder.rs`. This runs without a GPU and is the durable CI catch: any
/// revert that drops the retain, the balancing release, or the reset routing
/// fails here.
#[test]
fn retain_root_fix_source_invariants() {
let manifest_dir = env!("CARGO_MANIFEST_DIR");
let src_path = std::path::Path::new(manifest_dir)
.join("src")
.join("encoder.rs");
let src = std::fs::read_to_string(&src_path)
.unwrap_or_else(|e| panic!("read {}: {e}", src_path.display()));
// Match against ACTIVE CODE ONLY: drop any line whose first non-whitespace
// characters are `//`. A text-only `contains` would be fooled by a
// commented-out `// ... msg_send![encoder, retain]` revert (verified: a
// comment-out evades a raw contains check). Stripping line comments makes a
// comment-out revert fail this guard, not just a deletion.
let active: String = src
.lines()
.filter(|l| !l.trim_start().starts_with("//"))
.collect::<Vec<_>>()
.join("\n");
let active_norm: String = active.split_whitespace().collect::<Vec<_>>().join(" ");
// (1) The concurrent compute encoder is RETAINED when created, as ACTIVE
// code. This is the root fix: metal-rs returns it +0 autoreleased; we take
// ownership so an autorelease drain cannot drop the state Metal's barrier
// ordering keys on. Catches both deletion AND comment-out reverts.
assert!(
active_norm.contains("msg_send![encoder, retain]"),
"encoder.rs::get_or_create_encoder MUST retain the concurrent compute \
encoder as ACTIVE code (ADR-040 §0.21c-track2 root fix 80a58de) — not \
deleted, not commented out. Without it, an autorelease-pool drain drops \
the encoder ownership state and memoryBarrierWithScope no longer \
reliably orders dependent dispatches."
);
// (2) The retain is BALANCED by a release in end_active_encoder (no leak),
// also as active code.
assert!(
active_norm.contains("msg_send![enc, release]"),
"encoder.rs::end_active_encoder MUST release the encoder it retained \
in get_or_create_encoder as ACTIVE code (balance the +1; ADR-040 \
§0.21c-track2). A missing/commented release leaks one \
MTLComputeCommandEncoder per encoder cycle."
);
// (3) reset_command_buffer is leak-safe: it routes through
// end_active_encoder (which ends + releases if non-null, no-op if null)
// instead of a bare `active_encoder = null` that would leak the +1 retain
// and leave a zombie encoder (review Finder B, 3e7b030). Checked on active
// lines of the fn body.
let reset_body = active
.split("fn reset_command_buffer")
.nth(1)
.expect("encoder.rs must define reset_command_buffer (active code)");
let reset_head: String = reset_body.lines().take(40).collect::<Vec<_>>().join("\n");
assert!(
reset_head.contains("self.end_active_encoder()"),
"encoder.rs::reset_command_buffer MUST route through \
self.end_active_encoder() as ACTIVE code (review Finder B, 3e7b030) so \
the retained encoder is released on reset. A bare `active_encoder = \
null` guarded only by debug_assert leaks the +1 retain in release builds."
);
}
/// Functional SMOKE test of the fixed writer→barrier→reader path.
///
/// Writer (`tmp = a + b`) → `barrier()` → reader (`out = tmp * c`) on the
/// persistent concurrent compute encoder, ×16, asserting `out == (a+b)*c`.
///
/// HONEST SCOPE — this is NOT a fail-direction repro of the ownership bug.
/// Verified: with the retain temporarily reverted (commented out), this test
/// still PASSES. An isolated elementwise add→mul edge with a per-iteration
/// `finish()` full-sync does not open the race window — the fast producer wins
/// by timing regardless of retain, exactly as the earlier deleted timing test
/// did. So this test guards only that the happy path stays correct; it does
/// NOT discriminate the fix. The DISCRIMINATING guard is
/// `retain_root_fix_source_invariants` above (structural, catches a
/// deleted-or-commented retain) plus the production model-parity oracle
/// (`hf2q slot_aware_n8_per_slot_parity_vs_serial` under
/// `MLX_PRECOMPILED_METALLIB=0`, which exercised the real slow-producer race).
/// Kept as a cheap smoke + documentation of the path, not as the regression
/// gate — see the module docstring.
#[test]
fn writer_barrier_reader_orders_through_retained_encoder() {
let device = MlxDevice::new().expect("MlxDevice::new");
let mut registry = KernelRegistry::new();
let n = 4096usize;
let byte_len = n * std::mem::size_of::<f32>();
// Deterministic, non-degenerate inputs (no zeros so a stale tmp read can't
// accidentally match the correct product).
let a_data: Vec<f32> = (0..n).map(|i| 1.0 + (i % 97) as f32 * 0.5).collect();
let b_data: Vec<f32> = (0..n).map(|i| 2.0 + (i % 53) as f32 * 0.25).collect();
let c_data: Vec<f32> = (0..n).map(|i| 0.5 + (i % 31) as f32 * 0.125).collect();
let expected: Vec<f32> = (0..n).map(|i| (a_data[i] + b_data[i]) * c_data[i]).collect();
let mut a_buf = device.alloc_buffer(byte_len, DType::F32, vec![n]).expect("a");
let mut b_buf = device.alloc_buffer(byte_len, DType::F32, vec![n]).expect("b");
let mut c_buf = device.alloc_buffer(byte_len, DType::F32, vec![n]).expect("c");
a_buf.as_mut_slice::<f32>().expect("a").copy_from_slice(&a_data);
b_buf.as_mut_slice::<f32>().expect("b").copy_from_slice(&b_data);
c_buf.as_mut_slice::<f32>().expect("c").copy_from_slice(&c_data);
let tmp = device.alloc_buffer(byte_len, DType::F32, vec![n]).expect("tmp");
let out = device.alloc_buffer(byte_len, DType::F32, vec![n]).expect("out");
let executor = GraphExecutor::new(MlxDevice::new().expect("device2"));
// Repeat the writer→barrier→reader edge several times. Each iteration opens
// the persistent concurrent encoder, encodes the dependent pair across a
// drained autorelease pool, and full-syncs. Under the pre-fix unretained
// encoder this is the shape that lost barrier ordering; under the fix every
// iteration must produce the byte-correct product.
for iter in 0..16 {
{
let mut session = executor.begin().expect("begin");
// Writer: tmp = a + b
session
.elementwise_add(
&mut registry,
device.metal_device(),
&a_buf,
&b_buf,
&tmp,
n,
DType::F32,
)
.expect("graph add (writer)");
// RAW dependency: tmp written by add, read by mul. The barrier must
// order the reader after the writer on the concurrent encoder. This
// is the same primitive (`memoryBarrierWithScope:Buffers` on a
// concurrent encoder) that lost ordering pre-fix when the encoder
// was held unretained across the per-`begin()` command-buffer cycle.
session.barrier();
// Reader: out = tmp * c
session
.elementwise_mul(
&mut registry,
device.metal_device(),
&tmp,
&c_buf,
&out,
n,
DType::F32,
)
.expect("graph mul (reader)");
session.finish().expect("finish");
}
let result: Vec<f32> = out.as_slice::<f32>().expect("read out").to_vec();
for i in 0..n {
let diff = (result[i] - expected[i]).abs();
assert!(
diff < 1e-4,
"iter {iter}, element {i}: reader observed a stale `tmp` — the \
barrier did not order the writer before the reader on the \
concurrent encoder (ordering regression; ADR-040 §0.21c-track2 \
retain root fix). expected={}, got={}, diff={diff}",
expected[i],
result[i]
);
}
}
}