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
//! ANE persistent memory — experiment 3 of the silicon-duel series.
//!
//! `MLState` is a writable tensor that survives across predicts — that is a
//! hardware-resident key-value store nobody uses as one. This graph is a
//! mutable memory bank: fp16 [64,32] state, random-access read via `gather`,
//! write via `slice_update`, all inside one predict:
//!
//! m = read_state(mem) → out = gather(m, rslot)
//! m' = slice_update(m, wvec, [wslot,0]..[wslot+1,32]) → write_state(m')
//!
//! Read-before-write semantics per call; contents persist across calls
//! inside the MLState object. Verified: write N vectors, read them back
//! fp16-exact, then prove a fresh state reads zeros — the memory lives in
//! the state object, not the program.
//!
//! Usage: cargo run -p mil_infer --example mem --release
#[cfg(target_os = "macos")]
mod imp {
use half::f16;
use mil_infer::{ComputeUnits, Input, Model};
use mil_spec::DType;
use std::time::Duration;
const SLOTS: i64 = 64;
const DIM: i64 = 32;
pub fn run() -> Result<(), String> {
let dir = std::env::temp_dir().join("mil_mem");
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).map_err(|e| e.to_string())?;
let pkg = dir.join("mem.mlpackage");
mil_machines::write_memory_package(&pkg, SLOTS, DIM).map_err(|e| e.to_string())?;
let compiled =
mil_compile::compile(&pkg, &dir.join("compiled")).map_err(|e| e.to_string())?;
println!(
"compiled {} via {}",
compiled.path.display(),
compiled.backend
);
let model = Model::load(&compiled.path, ComputeUnits::CpuAndNeuralEngine)
.map_err(|e| e.to_string())?;
let state = model.new_state().map_err(|e| e.to_string())?;
let call = |wslot: i32, wvec: &[f32], rslot: i32| -> Result<Vec<f32>, String> {
let ws = wslot.to_le_bytes();
let rs = rslot.to_le_bytes();
let wv: Vec<u8> = wvec
.iter()
.flat_map(|v| f16::from_f32(*v).to_le_bytes())
.collect();
let p = model
.predict_with_state(
Some(&state),
&[
Input {
name: "wslot",
shape: &[1],
data: &ws,
dtype: DType::Int32,
},
Input {
name: "wvec",
shape: &[1, DIM],
data: &wv,
dtype: DType::Fp16,
},
Input {
name: "rslot",
shape: &[1],
data: &rs,
dtype: DType::Int32,
},
],
)
.map_err(|e| e.to_string())?;
Ok(p.outputs
.iter()
.find(|o| o.name == "out")
.map(|o| o.values())
.unwrap_or_default())
};
// write 8 distinct vectors into slots
let zeros = vec![0.0f32; DIM as usize];
let mut wrote = Vec::new();
let mut total = Duration::ZERO;
for s in 0..8i32 {
let v: Vec<f32> = (0..DIM).map(|d| s as f32 + d as f32 / 100.0).collect();
wrote.push(v.clone());
let t = std::time::Instant::now();
call(s, &v, 0)?;
total += t.elapsed();
}
// read them back — fp16-exact
let mut ok = 0;
for s in 0..8i32 {
let got = call(63, &zeros, s)?; // harmless write to scratch slot 63
if got.len() == DIM as usize
&& wrote[s as usize]
.iter()
.zip(got.iter())
.all(|(w, g)| (*w - *g).abs() < 0.02)
{
ok += 1;
}
}
println!(
"memory: wrote 8 vectors, read back {ok}/8 fp16-exact ({:.0}µs/write avg)",
total.as_secs_f64() * 1e6 / 8.0
);
// persistence is state-object-scoped: fresh state must read zeros
let state2 = model.new_state().map_err(|e| e.to_string())?;
let st2 = &state2;
let rs = 0i32.to_le_bytes();
let ws = 62i32.to_le_bytes();
let wv: Vec<u8> = zeros
.iter()
.flat_map(|v| f16::from_f32(*v).to_le_bytes())
.collect();
let p = model
.predict_with_state(
Some(st2),
&[
Input {
name: "wslot",
shape: &[1],
data: &ws,
dtype: DType::Int32,
},
Input {
name: "wvec",
shape: &[1, DIM],
data: &wv,
dtype: DType::Fp16,
},
Input {
name: "rslot",
shape: &[1],
data: &rs,
dtype: DType::Int32,
},
],
)
.map_err(|e| e.to_string())?;
let fresh = p
.outputs
.iter()
.find(|o| o.name == "out")
.map(|o| o.values())
.unwrap_or_default();
let zeros_read = fresh.iter().all(|v| v.abs() < 1e-3);
println!(
"fresh state reads zeros: {zeros_read} (memory lives in MLState, not the program)"
);
if ok == 8 && zeros_read {
println!("verdict: hardware-resident mutable KV store — proven");
Ok(())
} else {
Err("memory verification failed".into())
}
}
}
fn main() {
#[cfg(target_os = "macos")]
if let Err(e) = imp::run() {
eprintln!("mem: {e}");
std::process::exit(1);
}
#[cfg(not(target_os = "macos"))]
eprintln!("mem: macOS only");
}