phanerozoic commited on
Commit
db536d3
·
0 Parent(s):

8-bit threshold-logic CPU family: ternary-weight gate networks from a one-instruction SUBLEQ machine to an RV32IM plus F-subset RISC-V processor that runs stock-compiler C; composed IEEE-754 float pipelines with round-to-nearest-even bit-exact to hardware and metadata-driven verification; fully-wired rv32 datapath, FCVT int/float conversions, single gate-routed CPU runtime, leveled fast evaluation; single-file docs and consolidated machine runtime; strict-ternary build

Browse files
Files changed (48) hide show
  1. .gitattributes +36 -0
  2. .gitignore +4 -0
  3. README.md +613 -0
  4. build.py +0 -0
  5. build_all.py +214 -0
  6. cpu_programs.py +703 -0
  7. eval.py +0 -0
  8. eval_all.py +850 -0
  9. llm_integration/baseline.py +221 -0
  10. llm_integration/circuits.py +320 -0
  11. llm_integration/fitness.py +218 -0
  12. llm_integration/model.py +1263 -0
  13. llm_integration/smollm2/SMOLLM2_ARCHITECTURE.md +456 -0
  14. llm_integration/smollm2/analyze_smollm2.py +232 -0
  15. llm_integration/smollm2/smollm2_analysis.json +439 -0
  16. llm_integration/train.py +955 -0
  17. llm_integration/trained/llm.pt +3 -0
  18. llm_integration/trained/router.pt +3 -0
  19. machines.py +1719 -0
  20. neural_computer.safetensors +3 -0
  21. play.py +262 -0
  22. prune_weights.py +484 -0
  23. quantize.py +406 -0
  24. routing/routing.json +0 -0
  25. test_cpu.py +138 -0
  26. tests/__init__.py +0 -0
  27. tests/test_play.py +72 -0
  28. todo.md +21 -0
  29. variants/neural_alu16.safetensors +3 -0
  30. variants/neural_alu32.safetensors +3 -0
  31. variants/neural_alu8.safetensors +3 -0
  32. variants/neural_computer16.safetensors +3 -0
  33. variants/neural_computer16_reduced.safetensors +3 -0
  34. variants/neural_computer16_registers.safetensors +3 -0
  35. variants/neural_computer16_scratchpad.safetensors +3 -0
  36. variants/neural_computer16_small.safetensors +3 -0
  37. variants/neural_computer32.safetensors +3 -0
  38. variants/neural_computer32_reduced.safetensors +3 -0
  39. variants/neural_computer32_registers.safetensors +3 -0
  40. variants/neural_computer32_scratchpad.safetensors +3 -0
  41. variants/neural_computer32_small.safetensors +3 -0
  42. variants/neural_computer8.safetensors +3 -0
  43. variants/neural_computer8_reduced.safetensors +3 -0
  44. variants/neural_computer8_registers.safetensors +3 -0
  45. variants/neural_computer8_scratchpad.safetensors +3 -0
  46. variants/neural_computer8_small.safetensors +3 -0
  47. variants/neural_rv32.safetensors +3 -0
  48. variants/neural_subleq8.safetensors +3 -0
.gitattributes ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ *.7z filter=lfs diff=lfs merge=lfs -text
2
+ *.arrow filter=lfs diff=lfs merge=lfs -text
3
+ *.bin filter=lfs diff=lfs merge=lfs -text
4
+ *.bz2 filter=lfs diff=lfs merge=lfs -text
5
+ *.ckpt filter=lfs diff=lfs merge=lfs -text
6
+ *.ftz filter=lfs diff=lfs merge=lfs -text
7
+ *.gz filter=lfs diff=lfs merge=lfs -text
8
+ *.h5 filter=lfs diff=lfs merge=lfs -text
9
+ *.joblib filter=lfs diff=lfs merge=lfs -text
10
+ *.lfs.* filter=lfs diff=lfs merge=lfs -text
11
+ *.mlmodel filter=lfs diff=lfs merge=lfs -text
12
+ *.model filter=lfs diff=lfs merge=lfs -text
13
+ *.msgpack filter=lfs diff=lfs merge=lfs -text
14
+ *.npy filter=lfs diff=lfs merge=lfs -text
15
+ *.npz filter=lfs diff=lfs merge=lfs -text
16
+ *.onnx filter=lfs diff=lfs merge=lfs -text
17
+ *.ot filter=lfs diff=lfs merge=lfs -text
18
+ *.parquet filter=lfs diff=lfs merge=lfs -text
19
+ *.pb filter=lfs diff=lfs merge=lfs -text
20
+ *.pickle filter=lfs diff=lfs merge=lfs -text
21
+ *.pkl filter=lfs diff=lfs merge=lfs -text
22
+ *.pt filter=lfs diff=lfs merge=lfs -text
23
+ *.pth filter=lfs diff=lfs merge=lfs -text
24
+ *.rar filter=lfs diff=lfs merge=lfs -text
25
+ *.safetensors filter=lfs diff=lfs merge=lfs -text
26
+ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
27
+ *.tar.* filter=lfs diff=lfs merge=lfs -text
28
+ *.tar filter=lfs diff=lfs merge=lfs -text
29
+ *.tflite filter=lfs diff=lfs merge=lfs -text
30
+ *.tgz filter=lfs diff=lfs merge=lfs -text
31
+ *.wasm filter=lfs diff=lfs merge=lfs -text
32
+ *.xz filter=lfs diff=lfs merge=lfs -text
33
+ *.zip filter=lfs diff=lfs merge=lfs -text
34
+ *.zst filter=lfs diff=lfs merge=lfs -text
35
+ *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tensors.txt filter=lfs diff=lfs merge=lfs -text
.gitignore ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ __pycache__/
2
+ *.pyc
3
+ .pt file
4
+ .eval_cache/
README.md ADDED
@@ -0,0 +1,613 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ tags:
4
+ - threshold-logic
5
+ - neuromorphic
6
+ - computer-architecture
7
+ - turing-complete
8
+ - loihi
9
+ - truenorth
10
+ - akida
11
+ ---
12
+
13
+ # 8bit-threshold-computer
14
+
15
+ A Turing-complete CPU implemented entirely as threshold logic gates. Every gate, from Boolean primitives to arithmetic to control flow, is a single threshold neuron of the form:
16
+
17
+ ```
18
+ output = 1 if (Σ wᵢ·xᵢ + b) ≥ 0 else 0
19
+ ```
20
+
21
+ **Every weight in the file is in {-1, 0, 1}.** Biases are integers. Activations are the Heaviside step. Nothing else. The library was originally built with positional weights up to ±2³¹ for wide single-layer comparators; those have all been replaced with bit-cascaded multi-layer equivalents that use only ternary weights and small integer biases. Threshold-gate evaluation reduces to a popcount minus a popcount plus a bias, which is exactly what neuromorphic chips and FPGAs natively support.
22
+
23
+ The repository ships eighteen prebuilt configurations spanning three data-path widths (8, 16, 32 bits) and six memory sizes (0 B to 64 KB). The canonical file at the repo root is the largest of these: a 32-bit data path with a 64 KB address space and ~8.61 M tensor elements (~8.46 M gate weights and biases; the rest is `.inputs` routing metadata).
24
+
25
+ ```
26
+ neural_computer.safetensors 32-bit data, 64 KB memory, ~8.61M params (canonical)
27
+ variants/neural_computer{8,16,32}.safetensors full memory (64 KB)
28
+ variants/neural_computer{8,16,32}_reduced.safetensors 4 KB memory
29
+ variants/neural_computer{8,16,32}_small.safetensors 1 KB memory
30
+ variants/neural_computer{8,16,32}_scratchpad.safetensors 256 B memory
31
+ variants/neural_computer{8,16,32}_registers.safetensors 16 B memory
32
+ variants/neural_alu{8,16,32}.safetensors pure ALU, no memory
33
+ variants/neural_subleq8.safetensors one-instruction machine (SUBLEQ)
34
+ variants/neural_rv32.safetensors RV32IM + F-subset RISC-V processor
35
+ ```
36
+
37
+ Two further machines mark the endpoints of the family, each detailed in
38
+ its own section below: `neural_subleq8`, a Turing-complete one-instruction
39
+ computer whose entire control flow is a single threshold neuron; and
40
+ `neural_rv32`, a RISC-V processor (RV32IM + an F subset) that runs
41
+ stock-compiler C, with dual-issue execution, memory-mapped I/O, and a
42
+ self-referential NEUR opcode.
43
+
44
+ ---
45
+
46
+ ## Quick start
47
+
48
+ ```python
49
+ import torch
50
+ from safetensors.torch import load_file
51
+
52
+ tensors = load_file("neural_computer.safetensors")
53
+
54
+ def heaviside(x):
55
+ return (x >= 0).float()
56
+
57
+ # AND gate: fires when both inputs are 1
58
+ w = tensors['boolean.and.weight'] # [2]
59
+ b = tensors['boolean.and.bias'] # [1]
60
+ for a, c in [(0, 0), (0, 1), (1, 0), (1, 1)]:
61
+ out = heaviside((torch.tensor([a, c], dtype=torch.float32) * w).sum() + b)
62
+ print(f"AND({a}, {c}) = {int(out.item())}")
63
+ ```
64
+
65
+ Run the full circuit verification suite against any variant:
66
+
67
+ ```bash
68
+ python eval_all.py variants/ # 18 fitness variants (machines skipped)
69
+ python eval_all.py neural_computer.safetensors # the canonical file
70
+ python eval_all.py --cpu-program variants/ # also run an assembled
71
+ # program through the
72
+ # threshold-gated CPU
73
+ ```
74
+
75
+ `eval_all.py` reads each variant's manifest, runs a gate-level fitness suite (13,900–15,900 tests per variant covering Boolean, arithmetic, ALU, control, modular, error-detection, threshold, and float circuits, including end-to-end evaluation of the composed float pipelines from the shipped wiring metadata — see the Verification table), and optionally executes a small assembled program through a manifest-sized threshold CPU plus a chained 16- or 32-bit ALU sequence on wider variants.
76
+
77
+ For an interactive walkthrough that exercises Boolean gates, the 8-bit ALU, mod-5 divisibility, and a CPU loop end-to-end:
78
+
79
+ ```bash
80
+ python play.py # 1 KB demo, runs in seconds
81
+ python play.py --model neural_computer.safetensors # 64 KB, slower
82
+ ```
83
+
84
+ For end-to-end CPU validation (Fibonacci, sum 1..N, bubble sort, self-modifying JMP, all eight conditional jumps, CALL stack semantics, MUL cross-checked against repeated ADD, DIV cross-checked against repeated SUB, a bitwise AND/OR/XOR/SHL/SHR chain, and the architectural flag policy):
85
+
86
+ ```bash
87
+ python test_cpu.py # default: 1 KB, ~2 s
88
+ python test_cpu.py --model neural_computer.safetensors # 64 KB canonical, ~2 min
89
+ python test_cpu.py --only fib,sum_n # subset of suite
90
+ ```
91
+
92
+ Each program is assembled by a small Python assembler (`cpu_programs.py`) and run through the threshold-gated CPU; the driver verifies expected memory contents at HALT.
93
+
94
+ ---
95
+
96
+ ## Execution model
97
+
98
+ A self-contained machine. State goes in, state comes out:
99
+
100
+ - **Pure tensor computation**: state in, state out
101
+ - **Frozen circuits**: integer weights, Heaviside activation
102
+ - **ACT execution**: internal loop until `HALT`
103
+ - **No external orchestration**: one forward pass equals one complete program execution
104
+
105
+ Every datapath operation runs through threshold gates: ALU arithmetic
106
+ (ADD/SUB via ripple-carry, MUL via partial-product AND gates and shift-add,
107
+ DIV via per-stage bit-cascade comparators), bitwise logic, shifts,
108
+ comparisons, conditional-jump PC muxing, stack-pointer stepping, and all
109
+ memory reads and writes. Instruction decode (bit-field selection), PC
110
+ sequencing, and the bit complements feeding negated-condition selects are
111
+ fixed wiring in the runtimes — structural routing, not gates.
112
+
113
+ ```
114
+ ┌─────────────────────────────┐
115
+ │ Initial State │
116
+ │ [PC|Regs|Flags|Memory...] │
117
+ └─────────────┬───────────────┘
118
+
119
+ ┌─────────────────────────────┐
120
+ │ Threshold Circuit Layer │
121
+ │ ┌───────────────────────┐ │
122
+ │ │ Fetch: PC → Instr │ │
123
+ │ ├───────────────────────┤ │
124
+ │ │ Decode: Opcode/Ops │ │
125
+ │ ├───────────────────────┤ │
126
+ │ │ Execute: ALU/Mem │ │
127
+ │ ├───────────────────────┤ │
128
+ │ │ Writeback: Results │ │
129
+ │ ├───────────────────────┤ │
130
+ │ │ PC Update │ │
131
+ │ └───────────┬───────────┘ │
132
+ │ │ │
133
+ │ ┌────▼────┐ │
134
+ │ │ HALTED? │ │
135
+ │ └────┬────┘ │
136
+ │ no ────┴──── yes │
137
+ │ │ │ │
138
+ │ ▼ ▼ │
139
+ │ [loop] [exit] │
140
+ └─────────────┬───────────────┘
141
+
142
+ ┌─────────────────────────────┐
143
+ │ Final State │
144
+ └─────────────────────────────┘
145
+ ```
146
+
147
+ ### Instruction set
148
+
149
+ | Opcode | Mnemonic | Operation |
150
+ |--------|----------|-----------|
151
+ | 0x0 | ADD | R[d] = R[a] + R[b] |
152
+ | 0x1 | SUB | R[d] = R[a] - R[b] |
153
+ | 0x2 | AND | R[d] = R[a] & R[b] |
154
+ | 0x3 | OR | R[d] = R[a] \| R[b] |
155
+ | 0x4 | XOR | R[d] = R[a] ^ R[b] |
156
+ | 0x5 | SHL | R[d] = R[a] << 1 |
157
+ | 0x6 | SHR | R[d] = R[a] >> 1 |
158
+ | 0x7 | MUL | R[d] = R[a] * R[b] |
159
+ | 0x8 | DIV | R[d] = R[a] / R[b] |
160
+ | 0x9 | CMP | flags = R[a] - R[b] |
161
+ | 0xA | LOAD | R[d] = M[addr] |
162
+ | 0xB | STORE | M[addr] = R[s] |
163
+ | 0xC | JMP | PC = addr |
164
+ | 0xD | Jcc | PC = addr if cond (imm8[2:0]: 0=Z, 1=NZ, 2=C, 3=NC, 4=N, 5=P, 6=V, 7=NV) |
165
+ | 0xE | CALL | push PC; PC = addr |
166
+ | 0xF | HALT | stop execution |
167
+
168
+ **Flag policy.** Only ADD, SUB, MUL, and CMP write the FLAGS register; AND,
169
+ OR, XOR, SHL, SHR, DIV, LOAD, STORE, and all control transfers leave it
170
+ unchanged. Branch on the arithmetic op that set the condition — intervening
171
+ bitwise or shift instructions do not disturb it. The `flags_policy` program
172
+ in the CPU suite pins this behavior.
173
+
174
+ | Flag | ADD | SUB / CMP | MUL |
175
+ |---|---|---|---|
176
+ | Z | result == 0 | result == 0 | result == 0 |
177
+ | N | result bit 7 | result bit 7 | result bit 7 |
178
+ | C | carry-out | 1 when no borrow (`a >= b`) | 0 |
179
+ | V | signed overflow | signed overflow | 0 |
180
+
181
+ ### State tensor layout
182
+
183
+ The **state tensor** uses MSB-first bit ordering: index 0 of each multi-bit field is the most-significant bit. So `R0[0]` is bit 7 of the architectural register, `R0[7]` is bit 0.
184
+
185
+ ```
186
+ [ PC[N] | IR[16] | R0[8] R1[8] R2[8] R3[8] | FLAGS[4] | SP[N] | CTRL[4] | MEM[2^N][8] ]
187
+ ```
188
+
189
+ `N` is the address width (configurable, 0–16). Flags are ordered `Z, N, C, V`. Control bits are ordered `HALT, MEM_WE, MEM_RE, RESERVED`.
190
+
191
+ #### Bit ordering, one rule per scope
192
+
193
+ The state tensor's MSB-first convention does **not** propagate to subcircuit ports. Each subcircuit names its operand bits in its own scope:
194
+
195
+ | Scope | Convention | Example |
196
+ |---|---|---|
197
+ | State tensor | MSB-first (index 0 = MSB) | `R0[0]` is bit 7 of register R0 |
198
+ | Integer subcircuit ports (`$a[i]`, `$b[i]`) | LSB-indexed (index 0 = LSB) | `$a[0]` is bit 0 of operand `a` |
199
+ | Ripple-carry full adders (`fa0..fa7`) | LSB-first (fa0 = LSB) | `fa0` consumes `$a[0]` and `$b[0]` |
200
+ | Composed float circuits (`float16.add` etc.) | MSB-first operand words | `$a[0]` is the sign bit, `$a[1..E]` the exponent |
201
+ | Float result gates (`exp_out.bit{k}`, `frac_out.bit{k}`) | LSB-first field bits | `frac_out.bit0` is the fraction LSB |
202
+ | Instruction word | MSB-first (bit 15 = opcode high) | bit 15 is `opcode[3]` |
203
+
204
+ Worked example for `arithmetic.ripplecarry8bit`:
205
+
206
+ - Inputs: `$a[0]..$a[7]` and `$b[0]..$b[7]` where `$a[0]` is the LSB of `a`. To add `a = 0x05 = 0b00000101` and `b = 0x03`, drive `a[0]=1, a[1]=0, a[2]=1` (rest 0) and `b[0]=1, b[1]=1` (rest 0).
207
+ - Outputs: `fa0.ha2.sum.layer2`..`fa7.ha2.sum.layer2` are sum bits 0..7 (LSB to MSB), and `fa7.carry_or` is the final carry-out. The 8-bit result is `{fa7..fa0}` reading high-to-low.
208
+
209
+ This is also how `safetensors2verilog`'s threshold-logic frontend exposes the ports of any extracted subcircuit; `python -m safetensors2verilog ... --inspect` prints the port contract for any extracted circuit.
210
+
211
+ ### Instruction encoding (16-bit, MSB-first)
212
+
213
+ ```
214
+ 15..12 11..10 9..8 7..0
215
+ opcode rd rs imm8
216
+ ```
217
+
218
+ Interpretation:
219
+ - **R-type**: `rd = rd op rs` (imm8 ignored)
220
+ - **I-type**: `rd = op rd, imm8` (rs ignored)
221
+ - **Address-extended**: `LOAD`, `STORE`, `JMP`, `Jcc`, `CALL` consume the next word as a 16-bit address (big-endian); `imm8` is reserved and the PC skips 4 bytes when the jump is not taken.
222
+
223
+ ### Circuit categories
224
+
225
+ | Category | Circuits | Examples |
226
+ |----------|----------|----------|
227
+ | Boolean | 9 | AND, OR, NOT, NAND, NOR, XOR, XNOR, IMPLIES, BIIMPLIES |
228
+ | Arithmetic | 18+ | half/full adder, ripple-carry (8/16/32-bit), comparators (8/16/32-bit), 3-operand adder, A+B×C and (A+B)×C expressions |
229
+ | ALU | 8/16/32-bit | shifts, multiply, divide, INC/DEC, NEG, ROL/ROR, bitwise |
230
+ | Combinational | 10+ | MUX (2:1, 4:1, 8:1), DEMUX, 3-to-8 decoder, 8-to-3 encoder, barrel shifter, priority encoder |
231
+ | Control flow | 16 | JMP, conditional jumps (JZ/JNZ/JC/JNC/JN/JP/JV/JNV), CALL, RET, PUSH, POP |
232
+ | Memory | 3 | N-bit address decoder, read mux, write cells (packed) |
233
+ | Modular | 11 | divisibility by 2–12 (multi-layer for non-powers-of-2) |
234
+ | Threshold | 13 | k-of-n gates, majority, minority, exactly-k |
235
+ | Pattern | 10 | popcount, leading/trailing ones, symmetry |
236
+ | Error detection | 11 | parity (XOR tree), checksum, CRC, Hamming |
237
+ | Float (IEEE 754) | half + single | composed ADD, MUL, DIV, EQ/LT/LE/GT/GE pipelines plus unpack/pack/classify/normalize stages |
238
+
239
+ The float ADD/MUL/DIV/CMP circuits are **self-contained composed pipelines**: their complete internal wiring ships in the files as `.inputs` metadata, their external inputs are the raw operand words (`$a[0]` sign, `$a[1..E]` exponent, then mantissa; MSB-first), and the eval suite reconstructs and executes each netlist end to end from that metadata alone (`NetlistEvaluator`). Comparisons are fully IEEE (NaN unordered, `+0 == -0`, subnormal ordering, mixed signs). The arithmetic is round-to-nearest-even, **bit-exact to IEEE hardware** on the normal range, with exact specials (NaN, infinities including `inf - inf`/`inf * 0`/`0/0`/`inf/inf` → NaN and `x/0` → inf, signed zeros) and flush-to-zero for subnormal operands and results. Subnormal operands and results currently flush to zero; full subnormal support (gradual underflow) is unfinished and tracked in `todo.md`. Because the tests run from each file's own routing metadata, they prove the artifact is self-contained: `safetensors2verilog`-style extraction of `float16.add` and friends has everything it needs to emit a single composed module.
240
+
241
+ ### Tensor naming
242
+
243
+ ```
244
+ {category}.{circuit}[.{layer}][.{component}].{weight|bias}
245
+
246
+ Examples:
247
+ boolean.and.weight
248
+ boolean.xor.layer1.neuron1.weight
249
+ arithmetic.ripplecarry8bit.fa7.ha2.sum.layer1.or.weight
250
+ modular.mod5.eq.k15.bit3.match.weight
251
+ error_detection.paritychecker8bit.stage2.xor1.layer1.nand.bias
252
+ ```
253
+
254
+ Memory circuits are stored as packed tensors so the safetensors header stays manageable (`memory.addr_decode.weight`, `memory.read.and.weight`, `memory.write.and_old.weight`, etc.).
255
+
256
+ ---
257
+
258
+ ## Bit widths and memory profiles
259
+
260
+ The build tool emits one of 51 functionally distinct configurations: three data-path widths × seventeen address widths (0–16, where 0 means no memory).
261
+
262
+ **Bit widths** (`--bits`):
263
+
264
+ | Width | Range | Use case |
265
+ |-------|-------|----------|
266
+ | 8 | 0–255 | full CPU, legacy compatibility |
267
+ | 16 | 0–65,535 | extended arithmetic |
268
+ | 32 | 0–4,294,967,295 | practical arithmetic ranges |
269
+
270
+ **Memory profiles** (`-m`):
271
+
272
+ | Profile | Size | Addr bits | Filename suffix |
273
+ |---------|------|-----------|-----------------|
274
+ | `none` | 0 B | 0 | (uses `alu` instead of `computer`) |
275
+ | `registers` | 16 B | 4 | `_registers` |
276
+ | `scratchpad` | 256 B | 8 | `_scratchpad` |
277
+ | `small` | 1 KB | 10 | `_small` |
278
+ | `reduced` | 4 KB | 12 | `_reduced` |
279
+ | `full` | 64 KB | 16 | (none) |
280
+
281
+ Auto-generated filename: `neural_{alu|computer}{BITS}[_{MEMORY}].safetensors`. Custom address widths via `-a N` produce `_addrN`.
282
+
283
+ ```bash
284
+ python build.py --bits 32 --apply all # neural_computer32.safetensors
285
+ python build.py --bits 8 -m none --apply all # neural_alu8.safetensors
286
+ python build.py --bits 16 -m small --apply all # neural_computer16_small.safetensors
287
+ python build.py --bits 32 -a 6 --apply all # neural_computer32_addr6.safetensors
288
+ ```
289
+
290
+ To regenerate every named variant in one pass:
291
+
292
+ ```bash
293
+ python build_all.py
294
+ ```
295
+
296
+ This populates `variants/` with all 18 builds, quantizes each one to the smallest signed integer dtype that exactly represents its weights (~4× reduction in tensor data, with file size dominated by the safetensors header on the smaller profiles), verifies the strictly ternary weight invariant (`--ternary --strict`, so a build with any non-ternary weight fails loudly), stamps the `weight_quantization` metadata field, and runs `eval.py` on each as a sanity check.
297
+
298
+ The quantizer is also available standalone:
299
+
300
+ ```bash
301
+ python quantize.py path/to/file.safetensors # in-place
302
+ python quantize.py variants/ # whole directory
303
+ python quantize.py model.safetensors -o quantized.safetensors
304
+ python quantize.py file.safetensors --ternary # push toward {-1, 0, 1} weights
305
+ python quantize.py file.safetensors --ternary --strict # error if any weight is non-ternary
306
+ ```
307
+
308
+ Every weight and bias tensor in the canonical model fits in `int8`. The eval pipeline promotes weights to `float32` on load, so integer storage is exact and transparent.
309
+
310
+ **Ternary mode.** `build.py` emits only ternary weights: identity buffers are `weight=1, bias=-1` (`H(x - 1)`), and the comparators, modular detectors, and division stages that previously required positional weights up to ±2³¹ are bit-cascaded multi-layer equivalents. With `--ternary`, the quantizer verifies this and repairs legacy files: it rewrites historical single-input `weight=±2` buffers as `weight=±1` with the bias adjusted to preserve the heaviside output for binary inputs (`H(2x - 1) ≡ H(x - 1)`), and rebuilds pre-bit-cascade modular detectors (moduli already in bit-cascade form are left untouched, routing metadata included). `--strict` fails if any weight tensor remains non-ternary. Every shipped file carries the metadata field `weight_quantization: ternary`; a repaired file with remaining violations would be stamped `ternary_partial`.
311
+
312
+ ---
313
+
314
+ ## Verification
315
+
316
+ | Category | Coverage | Notes |
317
+ |----------|----------|-------|
318
+ | Boolean gates | exhaustive | all 2^n input combinations |
319
+ | Arithmetic (8-bit) | strategic sampling | edge values + diagonal pairs; ~50 cases per circuit |
320
+ | Arithmetic (16/32-bit) | strategic sampling | width-scaled extremes, alternating patterns, byte-boundary carries |
321
+ | ALU primitives (8/16/32-bit) | strategic sampling | edge inputs per operation; DIV comparators driven along real restoring-division traces |
322
+ | Control flow | exhaustive | all 2^3 input combinations per Jcc, per address bit |
323
+ | Threshold k-of-n | exhaustive | all 256 8-bit popcounts |
324
+ | Modular (all moduli, 8-bit input) | exhaustive | every value in [0, 255] |
325
+ | Parity | exhaustive | every value in [0, 255] |
326
+ | Pattern recognition | exhaustive | every value in [0, 255] |
327
+ | Combinational logic | exhaustive | full input space per gate |
328
+ | Float unpack/pack | exhaustive, functional | every bit gate driven with 0 and 1 (identity) |
329
+ | Float classify | functional | IEEE 754 categories (zero, subnormal, normal, inf, NaN) at edge encodings, both widths |
330
+ | Float CMP (composed) | functional, exact IEEE | full netlist rebuilt from the shipped `.inputs` metadata and evaluated end to end; NaN unordered, signed zeros, subnormal ordering, mixed signs — all five predicates, both widths |
331
+ | Float ADD/MUL/DIV (composed) | functional, bit-exact to IEEE hardware | same metadata-driven evaluation: exact specials, flush-to-zero, round-to-nearest-even; cross-checked against native float arithmetic |
332
+ | Memory / manifest | structure checks | packed-tensor shapes against the manifest |
333
+ | CPU integration | program-level | ten assembled programs (Fibonacci, sum, sort, self-modifying JMP, all eight Jcc, CALL stack push, MUL vs repeated ADD, DIV vs repeated SUB, a bitwise AND/OR/XOR/SHL/SHR pipeline, and the flag-policy pin) |
334
+
335
+ The 8-bit arithmetic and ALU tests use strategic sampling rather than the full 65,536-case sweep because exhaustive coverage at 8-bit is feasible but not necessary given that the circuits are constructed gate-by-gate. The 16-bit and 32-bit arithmetic tests sample edge cases only; full exhaustive coverage at those widths is infeasible without specialized hardware.
336
+
337
+ `eval_all.py` runs the unified suite. Exit code is the number of failing variants (0 means all pass). **Testing is evaluation, not rebuilding**: `python eval_all.py variants/` scores all 18 fitness variants against the shipped weights in about two minutes (~6 s each, the composed float netlists evaluated in `NetlistEvaluator`'s leveled mode) and cleanly skips the two standalone machines. Rebuilding the models (`build_all.py`, ~50 min for all 18) is a separate step, needed only when the circuit constructions in `build.py` change — routine verification never rebuilds. The batched evaluator is population-safe: every chained intermediate (carry, borrow, mux select) is computed per population slot, so `prune_weights.py`'s parallel fitness screens are exact rather than slot-0 approximations.
338
+
339
+ ---
340
+
341
+ ## neural_subleq8 — the one-instruction machine
342
+
343
+ The minimal member of the family: a Turing-complete computer with **no instruction decoder** (there is only one instruction), **no registers**, and **no flags** — 194 logic gates totalling 548 ternary parameters. Its one architectural decision, branch or fall through, is a single threshold neuron (`subleq.leq`) over the sign and zero of the subtraction result.
344
+
345
+ - 8-bit data, 8-bit addresses, 256 bytes of packed threshold memory, a program counter and nothing else.
346
+ - An instruction is three bytes `A B C`. Step: `M[B] = (M[B] - M[A]) mod 256`; if the result is `<= 0` in two's complement (zero or bit 7 set), `PC = C`, else `PC += 3`. `PC = 0xFF` halts.
347
+ - Circuits (all wiring shipped as `.inputs` metadata): an 8-bit two's-complement subtractor, the branch neuron, a `PC + 3` incrementer, a branch mux, and the packed 256-row memory.
348
+
349
+ ```bash
350
+ python build.py --apply subleq # -> variants/neural_subleq8.safetensors
351
+ python machines.py subleq # exhaustive datapath + lockstep programs
352
+ ```
353
+
354
+ `machines.py subleq` evaluates the datapath **from the shipped wiring metadata** over all 65,536 operand pairs (result and branch decision both exhaustive), checks the PC mux, and runs a program suite (clear, negate-copy, add-by-double-negation, countdown loop) in full-state lockstep against a reference emulator. Third-party SUBLEQ toolchains target this machine directly once the `0xFF` halt convention is mapped.
355
+
356
+ ---
357
+
358
+ ## neural_rv32 — a RISC-V processor as a threshold network
359
+
360
+ The most capable member: a RISC-V CPU whose entire datapath is ternary threshold gates (`variants/neural_rv32.safetensors`, ~8.6 M parameters, 64 KB memory, strictly ternary).
361
+
362
+ - **RV32I base**: LUI, AUIPC, JAL, JALR, all six branches, LB/LH/LW/LBU/LHU, SB/SH/SW, and the full OP-IMM/OP groups. 32 × 32-bit registers (x0 zero), little-endian memory through the packed threshold circuits. ECALL halts.
363
+ - **M extension**: MUL, MULH, MULHSU, MULHU (full 64-bit product through a shift-add array with gate-level sign correction), DIV, DIVU, REM, REMU (32 restoring stages, spec-exact divide-by-zero and overflow).
364
+ - **F subset**: FLW, FSW, FMV.W.X, FMV.X.W, FADD.S, FSUB.S, FMUL.S, FDIV.S, FEQ.S, FLT.S, FLE.S, FSGNJ[N/X].S, FCVT.W.S, FCVT.S.W — the arithmetic executed by the composed float32 pipelines (round-to-nearest-even, bit-exact to hardware; specials and flush-to-zero as above), the int/float conversions gate-routed through the priority encoder and barrel shifter and cross-checked bit-exact against native conversion.
365
+ - **NEUR (custom-0)**: `neur rd, rs1, rs2` evaluates one threshold neuron — `rd = H(popcount(rs1[7:0] & rs2[7:0]) - popcount(rs1[7:0] & rs2[15:8]) + sext(rs2[20:16]))`. Networks of NEUR instructions are neural networks running as software on the neural network; the test suite computes XOR with a two-layer NEUR net.
366
+ - **Dual issue**: two adjacent OP/OP-IMM/LUI/AUIPC instructions retire in one cycle when the gate-level hazard comparators (`rv32.hazard.*`) clear RAW and WAW dependences.
367
+ - **MMIO**: stores to `0xFF00` append a character to the console.
368
+
369
+ Signed comparisons ride the unsigned bit-cascade with sign bits complemented through NOT gates; SLL uses the barrel shifter, SRL is bit-reversal wiring over it, SRA a gate mux over the complement form. Instruction decode, immediate extraction, register-file indexing, and PC sequencing are fixed wiring, per the family convention.
370
+
371
+ ```bash
372
+ python build.py --apply rv32 # build the file
373
+ python quantize.py variants/neural_rv32.safetensors --ternary --strict
374
+ python machines.py rv32 # eight-program lockstep suite
375
+ python machines.py rv32-c # stock-compiler C, end to end
376
+ ```
377
+
378
+ **Running compiled C.** `machines.py rv32-c` compiles a freestanding C program (gcd, Fibonacci, insertion sort; `rv32im`, so real `mul`/`rem`) with an unmodified clang rv32im toolchain, loads the relocatable object with an in-repo loader (no external linker — it resolves the R_RISCV relocations of one translation unit and lays the sections out flat), executes it on the threshold CPU, and checks the return value against the value computed natively. The program retires in ~300 instructions and matches exactly. Stock `rv32im` toolchains (gcc, clang, rustc) emit this ISA. Unfinished float work — subnormals, FMA, and the FCVT rounding-mode field — is tracked in `todo.md`.
379
+
380
+ The composed circuits evaluate in a leveled mode (`NetlistEvaluator`, one padded tensor op per topological level instead of one Python step per gate), ~18× faster on the FPU-scale netlists.
381
+
382
+ ---
383
+
384
+ ## Threshold logic
385
+
386
+ A threshold gate computes a Boolean function by taking a weighted sum of binary inputs and comparing the result to a threshold; the output is 1 when the sum meets or exceeds the threshold and 0 otherwise. Equivalently, it is a neuron with Heaviside step activation, integer weights, and an integer bias.
387
+
388
+ Threshold gates are strictly more powerful than standard Boolean gates. A single threshold gate can compute any linearly separable Boolean function, which includes AND, OR, NAND, NOR, IMPLIES, and many others that require multiple levels of conventional gates. Functions that are not linearly separable (XOR, parity, mod-k for k not a power of two) require multiple layers.
389
+
390
+ Example gates:
391
+
392
+ ```
393
+ AND: w=[1, 1], b=-2
394
+ H(0+0-2) = 0 H(1+1-2) = 1
395
+
396
+ OR: w=[1, 1], b=-1
397
+ H(0+0-1) = 0 H(1+0-1) = 1
398
+
399
+ XOR: two layers (not linearly separable)
400
+ layer 1: OR + NAND
401
+ layer 2: AND of the two
402
+ ```
403
+
404
+ A full adder is two half-adders plus a carry OR, around four threshold layers. An 8-bit ripple-carry adder is eight chained full adders, around 32 layers.
405
+
406
+ ### History
407
+
408
+ Warren McCulloch and Walter Pitts introduced the threshold neuron in 1943, proving that networks of such neurons can compute any Boolean function. Their work preceded both the perceptron and modern neural networks and established the theoretical foundation for neural computation.
409
+
410
+ The 1960s saw substantial work on threshold logic synthesis. Saburo Muroga, Robert McNaughton, and Michael Dertouzos developed algebraic methods for determining whether a Boolean function can be implemented as a single threshold gate, and if so, how to compute the appropriate weights. The focus was on individual gates rather than complete systems.
411
+
412
+ Frank Rosenblatt's Mark I Perceptron (1957–1960) implemented threshold neurons in hardware using potentiometers for weights, but it was a pattern classifier that learned its weights through training; the final configurations were not published. Bernard Widrow's ADALINE and MADALINE (1960–1963) similarly used adaptive threshold elements with weights learned via the LMS algorithm.
413
+
414
+ Hava Siegelmann and Eduardo Sontag proved in the 1990s that recurrent neural networks are Turing-complete. The construction relied on continuous sigmoid activations with infinite precision, not the discrete step function used in threshold logic. Other theoretical work on neural Turing machines and differentiable computers followed similar patterns: universality with continuous activations chosen to support gradient-based training.
415
+
416
+ ### Neuromorphic hardware
417
+
418
+ Modern neuromorphic processors implement large arrays of configurable threshold-like neurons in silicon:
419
+
420
+ - **Intel Loihi** (2017): 128 neuromorphic cores with programmable synaptic weights, spike-based communication, and on-chip learning. Supports integer weights and configurable neuron dynamics.
421
+ - **IBM TrueNorth** (2014): one million neurons and 256 million synapses across a 4096-core array. Each neurosynaptic core implements 256 neurons with configurable weights and thresholds.
422
+ - **BrainChip Akida** (2021): edge-oriented event-based processing with integer weights.
423
+ - **SpiNNaker** (University of Manchester): ARM cores simulating spiking networks at scale.
424
+
425
+ Published work on these platforms has focused on neural network inference, sensory processing, and pattern recognition. A 2024 paper demonstrated basic logic gates, adders, and decoders on SpiNNaker and Dynap-SE1 and described that work as "a first step toward the construction of a spiking computer"; that implementation lacked instruction fetch, a program counter, memory, and control logic.
426
+
427
+ The weights in this repository implement a complete CPU: registers, ALU with 16 operations, status flags, conditional branching, subroutine calls, stack operations, and memory access. Every logic component is a threshold neuron with integer weights; instruction decode and PC sequencing are fixed wiring between them.
428
+
429
+ ---
430
+
431
+ ## Hardware compatibility
432
+
433
+ All weights are in {-1, 0, 1}, all activations are Heaviside step, and every gate is a single weighted sum followed by a sign test. This eliminates multipliers entirely: each gate evaluation is a popcount of `+1`-weighted inputs minus a popcount of `-1`-weighted inputs plus an integer bias. The circuits are intended to deploy directly on:
434
+
435
+ - **FPGA**: every gate maps to a small LUT cluster (or a popcount tree of LUT4/LUT6 + carry chain). Ternary weight storage compresses to 2 bits per weight; routing collapses to bit selection.
436
+ - **Intel Loihi**: integer weights and Heaviside threshold neurons are the native primitive. Ternary fits well within Loihi's 8-bit weight range.
437
+ - **IBM TrueNorth**: configurable threshold per neurosynaptic core; ternary weights and small biases are within the supported range.
438
+ - **BrainChip Akida**: edge-oriented integer-weight inference; ternary weights fit cleanly.
439
+
440
+ ---
441
+
442
+ ## LLM integration
443
+
444
+ Threshold circuits can be embedded into transformer MLP layers to give a language model exact arithmetic. Standard LLMs fail at arithmetic because they interpolate over the training distribution rather than compute, so a 360M-parameter model trained on web text has seen `127 + 128 = 255` few times if at all and guesses based on pattern matching.
445
+
446
+ The integration freezes the circuits and trains only the interface layers that:
447
+
448
+ 1. Extract operands from token embeddings.
449
+ 2. Route computation through the appropriate circuit.
450
+ 3. Inject the result back into the residual stream.
451
+
452
+ The model learns *call dispatch*; the arithmetic is already solved.
453
+
454
+ ### Architecture
455
+
456
+ ```
457
+ x ──┬── MLP path ─────────────────┬── + ── output
458
+ │ │
459
+ └── BitExtractor ── Circuit ──┴── BitInjector
460
+
461
+ Router (learned weighting)
462
+ ```
463
+
464
+ Augmented MLP forward pass:
465
+
466
+ ```python
467
+ def forward(x): # x: [batch, seq, d_model=960]
468
+ mlp_out = self.down_proj(silu(self.gate_proj(x)) * self.up_proj(x))
469
+
470
+ a_bits, b_bits = self.bit_extractor(x) # [batch, seq, 8] each
471
+ result_bits, carry = self.circuits.add_8bit(a_bits, b_bits)
472
+ flags = self.compute_flags(result_bits, carry)
473
+ circuit_delta = self.bit_injector(result_bits, flags)
474
+
475
+ route_weights = self.router(x) # [batch, seq, 2] softmax
476
+ return mlp_out + route_weights[..., 1:2] * circuit_delta
477
+ ```
478
+
479
+ ### Target model
480
+
481
+ The reference integration uses HuggingFace's [SmolLM2-360M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct). See [`llm_integration/SMOLLM2_ARCHITECTURE.md`](llm_integration/smollm2/SMOLLM2_ARCHITECTURE.md) for the full technical analysis.
482
+
483
+ | Property | Value |
484
+ |----------|-------|
485
+ | Parameters | 361.82 M |
486
+ | Hidden dimension | 960 (matches the extractor input) |
487
+ | Layers | 32 transformer blocks |
488
+ | Attention | 15 query heads, 5 KV heads (GQA) |
489
+ | MLP | SwiGLU (960 → 2560 → 960) |
490
+ | Position encoding | RoPE (theta = 100k, max 8192) |
491
+
492
+ Digits tokenize individually (`"47 + 86"` → `['4', '7', ' +', ' ', '8', '6']`, with digit token IDs `32 + digit_value`), which makes position-based operand extraction practical.
493
+
494
+ ### Gradient flow
495
+
496
+ Heaviside has zero gradient almost everywhere. The implementation uses a straight-through estimator:
497
+
498
+ ```python
499
+ class HeavisideSTE(torch.autograd.Function):
500
+ @staticmethod
501
+ def forward(ctx, x):
502
+ return (x >= 0).float()
503
+
504
+ @staticmethod
505
+ def backward(ctx, grad_output):
506
+ return grad_output
507
+ ```
508
+
509
+ At inference, Heaviside is the true step function; if the extractor identifies operands correctly, the circuit produces the correct result by construction.
510
+
511
+ ### Baseline
512
+
513
+ SmolLM2-360M-Instruct on randomized 8-bit arithmetic (2,000 cases, operands uniform on [0, 255], generous answer extraction):
514
+
515
+ | Operation | Accuracy |
516
+ |-----------|----------|
517
+ | Addition | 35.92% |
518
+ | Subtraction | 17.72% |
519
+ | Multiplication | 1.25% |
520
+ | Greater than | 14.37% |
521
+ | Less than | 4.31% |
522
+ | Equality | 0.28% |
523
+ | **Overall** | **11.90%** (238/2000) |
524
+
525
+ Multiplication accuracy at 1.25% is essentially random over the output space. Comparison operations often echo the expression rather than evaluate it. Even addition fails roughly two-thirds of the time on full 8-bit operands. Performance degrades further as operand magnitude increases: edge cases like `127 + 128` are almost never correct.
526
+
527
+ The frozen threshold circuits reach 100% on the same task when given correctly formatted bit inputs (10,000 random cases, every operation). The integration challenge is therefore the extractor, not the arithmetic.
528
+
529
+ ### Trainable parameters (SmolLM2, hidden_dim = 960)
530
+
531
+ | Component | Parameters | Description |
532
+ |-----------|------------|-------------|
533
+ | AttentionPooling | ~3.7 M | 4-head attention over the sequence |
534
+ | MultiHeadBitExtractor (× 2) | ~245 K each | 8 per-bit MLPs for A and B |
535
+ | OpRouter | ~246 K | 960 → 256 → 6 MLP |
536
+ | **Extractor total** | **~4.4 M** | full extraction module |
537
+
538
+ Alternative architectures: `PositionExtractor` (~1.5 M, position-specific, no attention), `DigitExtractor` (~1.2 M, predicts digits 0–9 instead of bits), `HybridExtractor` (digit lookup with MLP fallback for word numerals). With `--unfreeze_layers 4` an additional ~39.3 M trainable parameters open up in the top four transformer layers.
539
+
540
+ ### Training
541
+
542
+ ```bash
543
+ python train.py --mode router --epochs 100 # sanity check
544
+ python train.py --mode llm --epochs 100 --batch_size 256 # frozen LLM
545
+ python train.py --mode llm --unfreeze_layers 4 --batch_size 4096 # fine-tune top layers
546
+ ```
547
+
548
+ Loss components: BCE on output bits, BCE on extracted A and B bits (2× weight), and CE on operation classification. Curriculum runs 0–9 → 0–99 → 0–255. Optimizer is AdamW, lr 3e-4, gradient clipping 1.0.
549
+
550
+ ---
551
+
552
+ ## Repository layout
553
+
554
+ ```
555
+ neural_computer.safetensors canonical model (32-bit, 64 KB, ~8.61M params)
556
+ variants/ 18 prebuilt configurations + neural_subleq8 + neural_rv32
557
+ build.py generator (one safetensors per invocation; also `subleq`, `rv32`)
558
+ build_all.py builds, quantizes (--ternary --strict), and verifies every profile
559
+ quantize.py min integer dtypes + ternary verification/repair
560
+ eval.py gate-level fitness suite, NetlistEvaluator, 64 KB reference CPU
561
+ eval_all.py variant-agnostic harness + manifest-sized threshold CPU
562
+ cpu_programs.py assembler + ten-program suite for CPU-level validation
563
+ test_cpu.py runs the program suite against a chosen variant
564
+ machines.py neural_subleq8 + neural_rv32 runtimes, references, assemblers,
565
+ RV32 object loader, and both test suites (subleq / rv32 / rv32-c)
566
+ play.py interactive demo
567
+ prune_weights.py GPU-batched weight reduction with conflict resolution
568
+ llm_integration/ SmolLM2 extractor + circuit wrapper + training code
569
+ ├── circuits.py FrozenThresholdCircuits (loads safetensors, exposes
570
+ │ add_8bit / sub_8bit / mul_8bit / compare_*)
571
+ ├── model.py Extractor variants, ArithmeticModel
572
+ ├── train.py router / interface / llm training modes
573
+ ├── fitness.py randomized fitness function
574
+ ├── baseline.py vanilla SmolLM2 baseline measurement
575
+ ├── trained/ checkpointed extractor weights
576
+ └── smollm2/
577
+ ├── SMOLLM2_ARCHITECTURE.md architecture analysis
578
+ ├── analyze_smollm2.py analysis script
579
+ └── smollm2_analysis.json analysis output
580
+ ```
581
+
582
+ ---
583
+
584
+ ## Citation
585
+
586
+ ```bibtex
587
+ @misc{8bit-threshold-computer,
588
+ title={8bit-threshold-computer: A Turing-Complete Threshold Logic CPU},
589
+ author={Norton, Charles},
590
+ year={2026},
591
+ howpublished={Hugging Face},
592
+ url={https://huggingface.co/phanerozoic/8bit-threshold-computer}
593
+ }
594
+ ```
595
+
596
+ ---
597
+
598
+ ## License
599
+
600
+ MIT
601
+
602
+ ---
603
+
604
+ ## References
605
+
606
+ 1. McCulloch & Pitts (1943). *A Logical Calculus of Ideas Immanent in Nervous Activity.*
607
+ 2. Muroga (1971). *Threshold Logic and Its Applications.*
608
+ 3. Siegelmann & Sontag (1995). *On the Computational Power of Neural Nets.*
609
+ 4. Bengio et al. (2013). *Estimating or Propagating Gradients Through Stochastic Neurons.*
610
+ 5. Ma et al. (2024). *The Era of 1-bit LLMs* (BitNet b1.58).
611
+ 6. HuggingFace (2024). *SmolLM2: Small Language Models* — [model card](https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct).
612
+ 7. Vaswani et al. (2017). *Attention Is All You Need.*
613
+ 8. Su et al. (2021). *RoFormer: Enhanced Transformer with Rotary Position Embedding.*
build.py ADDED
The diff for this file is too large to render. See raw diff
 
build_all.py ADDED
@@ -0,0 +1,214 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Build and verify every named (bits, memory_profile) variant.
3
+
4
+ Outputs:
5
+ variants/neural_alu{8,16,32}.safetensors - no memory
6
+ variants/neural_computer{8,16,32}_registers.safetensors - 16 B
7
+ variants/neural_computer{8,16,32}_scratchpad.safetensors - 256 B
8
+ variants/neural_computer{8,16,32}_small.safetensors - 1 KB
9
+ variants/neural_computer{8,16,32}_reduced.safetensors - 4 KB
10
+ variants/neural_computer{8,16,32}.safetensors - 64 KB
11
+
12
+ Each variant is built from the canonical seed, quantized to minimum
13
+ integer dtypes with the strict ternary check (--ternary --strict), and
14
+ then verified with eval.py via the BatchedFitnessEvaluator; the summary
15
+ records (tensor count, params, file size, fitness, total_tests, seconds).
16
+ """
17
+
18
+ from __future__ import annotations
19
+ import os
20
+ import shutil
21
+ import subprocess
22
+ import sys
23
+ import time
24
+ from pathlib import Path
25
+
26
+ import torch
27
+ from safetensors import safe_open
28
+
29
+ ROOT = Path(__file__).resolve().parent
30
+ SEED = ROOT / "neural_computer.safetensors"
31
+ OUT_DIR = ROOT / "variants"
32
+ OUT_DIR.mkdir(exist_ok=True)
33
+
34
+ PROFILES = ["none", "registers", "scratchpad", "small", "reduced", "full"]
35
+ BITS = [8, 16, 32]
36
+
37
+
38
+ def variant_filename(bits: int, profile: str) -> str:
39
+ if profile == "none":
40
+ return f"neural_alu{bits}.safetensors"
41
+ if profile == "full":
42
+ return f"neural_computer{bits}.safetensors"
43
+ return f"neural_computer{bits}_{profile}.safetensors"
44
+
45
+
46
+ def run(cmd: list[str], timeout: int = 600) -> tuple[int, str]:
47
+ p = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
48
+ return p.returncode, (p.stdout or "") + (p.stderr or "")
49
+
50
+
51
+ def build_variant(bits: int, profile: str) -> Path:
52
+ out = OUT_DIR / variant_filename(bits, profile)
53
+ shutil.copy2(SEED, out)
54
+ cmd = [
55
+ sys.executable, str(ROOT / "build.py"),
56
+ "--bits", str(bits),
57
+ "-m", profile,
58
+ "--apply",
59
+ "--model", str(out),
60
+ "all",
61
+ ]
62
+ rc, log = run(cmd, timeout=900)
63
+ if rc != 0:
64
+ raise RuntimeError(f"build failed for bits={bits} profile={profile}:\n{log[-1500:]}")
65
+ return out
66
+
67
+
68
+ def quantize_variant(path: Path) -> tuple[int, int]:
69
+ """Run quantize.py --ternary --strict on a built variant. This is the
70
+ last pipeline step: it casts every tensor to its minimum signed integer
71
+ dtype, verifies the strictly ternary weight invariant, and stamps the
72
+ weight_quantization metadata field. Returns (bytes_before, bytes_after).
73
+ """
74
+ rc, log = run([sys.executable, str(ROOT / "quantize.py"), str(path),
75
+ "--ternary", "--strict"], timeout=300)
76
+ if rc != 0:
77
+ raise RuntimeError(f"quantize failed for {path.name}:\n{log[-800:]}")
78
+ # parse the "file X.X MB -> Y.Y MB" line
79
+ for line in log.splitlines():
80
+ if "file" in line and "->" in line and path.name in line:
81
+ try:
82
+ parts = line.split("file")[1].split("->")
83
+ before = float(parts[0].strip().split()[0]) * 1e6
84
+ after = float(parts[1].strip().split()[0]) * 1e6
85
+ return int(before), int(after)
86
+ except Exception:
87
+ pass
88
+ return 0, 0
89
+
90
+
91
+ def measure_variant(path: Path) -> dict:
92
+ """Read tensor count, params, manifest values from the variant."""
93
+ with safe_open(str(path), framework="pt") as f:
94
+ keys = list(f.keys())
95
+ params = sum(f.get_tensor(k).numel() for k in keys)
96
+ manifest = {
97
+ k.split(".", 1)[1]: f.get_tensor(k).item()
98
+ for k in keys if k.startswith("manifest.") and f.get_tensor(k).numel() == 1
99
+ }
100
+ return {
101
+ "tensors": len(keys),
102
+ "params": params,
103
+ "size_mb": path.stat().st_size / (1024 * 1024),
104
+ "manifest": manifest,
105
+ }
106
+
107
+
108
+ def eval_variant(path: Path, device: str = "cpu", timeout: int = 600) -> dict:
109
+ """Run eval.py against a variant and parse fitness."""
110
+ cmd = [
111
+ sys.executable, str(ROOT / "eval.py"),
112
+ "--model", str(path),
113
+ "--device", device,
114
+ "--quiet",
115
+ ]
116
+ t0 = time.time()
117
+ rc, log = run(cmd, timeout=timeout)
118
+ dt = time.time() - t0
119
+
120
+ fitness = None
121
+ total_tests = None
122
+ status = "ERROR"
123
+ for line in log.splitlines():
124
+ line = line.strip()
125
+ if line.startswith("Fitness:"):
126
+ try:
127
+ fitness = float(line.split()[1])
128
+ except Exception:
129
+ pass
130
+ elif line.startswith("Total tests:"):
131
+ try:
132
+ total_tests = int(line.split()[2])
133
+ except Exception:
134
+ pass
135
+ elif line.startswith("STATUS:"):
136
+ status = line.split()[1]
137
+ return {
138
+ "rc": rc,
139
+ "fitness": fitness,
140
+ "total_tests": total_tests,
141
+ "status": status,
142
+ "elapsed_s": dt,
143
+ "log_tail": "\n".join(log.splitlines()[-15:]),
144
+ }
145
+
146
+
147
+ def main() -> None:
148
+ rows = []
149
+ print(f"Building 18 variants into {OUT_DIR}\n")
150
+ for bits in BITS:
151
+ for profile in PROFILES:
152
+ label = f"bits={bits} profile={profile}"
153
+ print(f"=== {label} ===", flush=True)
154
+ t0 = time.time()
155
+ try:
156
+ path = build_variant(bits, profile)
157
+ bt = time.time() - t0
158
+ pre_q_meta = measure_variant(path)
159
+ # Quantize in-place as the final step; weights are
160
+ # integer-valued so this is exact, --strict fails the build
161
+ # if any weight is non-ternary, and header metadata
162
+ # (signal_registry) is carried through.
163
+ qb, qa = quantize_variant(path)
164
+ meta = measure_variant(path)
165
+ ev = eval_variant(path, device="cpu", timeout=900)
166
+ rows.append({
167
+ "bits": bits, "profile": profile,
168
+ "filename": path.name,
169
+ "build_s": bt,
170
+ **meta,
171
+ **{k: ev[k] for k in ("fitness", "total_tests", "status", "elapsed_s")},
172
+ "log_tail": ev["log_tail"] if ev["status"] != "PASS" else "",
173
+ })
174
+ q_ratio = qb / qa if qa else 1.0
175
+ print(f" built in {bt:.1f}s size={pre_q_meta['size_mb']:.1f}MB -> "
176
+ f"{meta['size_mb']:.1f}MB after quant ({q_ratio:.2f}x)"
177
+ f" params={meta['params']:,} tensors={meta['tensors']:,}")
178
+ print(f" eval: fitness={ev['fitness']} tests={ev['total_tests']}"
179
+ f" status={ev['status']} ({ev['elapsed_s']:.1f}s)")
180
+ if ev["status"] != "PASS":
181
+ print(" --- failure tail ---")
182
+ print(" " + "\n ".join(ev["log_tail"].splitlines()))
183
+ print(" --------------------")
184
+ except Exception as e:
185
+ print(f" EXCEPTION: {e}")
186
+ rows.append({"bits": bits, "profile": profile, "error": str(e)})
187
+ print()
188
+
189
+ print("=" * 88)
190
+ print(" SUMMARY")
191
+ print("=" * 88)
192
+ header = f"{'bits':>4} {'profile':<11} {'size_MB':>8} {'tensors':>8} {'params':>11} {'fitness':>9} {'tests':>6} {'status':>7}"
193
+ print(header)
194
+ print("-" * len(header))
195
+ for r in rows:
196
+ if "error" in r:
197
+ print(f"{r['bits']:>4} {r['profile']:<11} ERROR: {r['error'][:60]}")
198
+ continue
199
+ fit = f"{r['fitness']:.4f}" if r['fitness'] is not None else "n/a"
200
+ tests = r['total_tests'] if r['total_tests'] is not None else "?"
201
+ print(f"{r['bits']:>4} {r['profile']:<11} {r['size_mb']:>8.1f} "
202
+ f"{r['tensors']:>8,} {r['params']:>11,} "
203
+ f"{fit:>9} {tests:>6} {r['status']:>7}")
204
+
205
+ fail = [r for r in rows if r.get("status") != "PASS" or "error" in r]
206
+ print()
207
+ if fail:
208
+ print(f"FAILURES: {len(fail)}/{len(rows)}")
209
+ else:
210
+ print(f"ALL {len(rows)} VARIANTS PASS")
211
+
212
+
213
+ if __name__ == "__main__":
214
+ main()
cpu_programs.py ADDED
@@ -0,0 +1,703 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ CPU validation programs for the threshold computer.
3
+
4
+ A small assembler and a suite of programs that exercise the ISA end-to-end:
5
+ arithmetic, control flow, memory access, self-modifying code, all eight
6
+ conditional jumps, the call mechanism, and a sort.
7
+
8
+ Each program returns (mem, expected, max_cycles, description) where:
9
+ mem : list[int] -- complete memory image
10
+ expected : dict[int, int] -- {address: expected_value} verified at HALT
11
+ max_cycles : int -- cycle budget (an infinite loop will hit this)
12
+ description: str -- short human-readable summary
13
+
14
+ Programs target the 1 KB profile (addr_bits=10) by default but use only the
15
+ low 256 bytes so they also run on scratchpad and larger profiles.
16
+
17
+ All programs assume the CPU starts with PC=0, SP defaulting to addr_mask
18
+ (highest address; CALL pre-decrements before writing).
19
+ """
20
+
21
+ from __future__ import annotations
22
+ from typing import Dict, List, Tuple
23
+
24
+
25
+ # ----------------------------------------------------------------------------
26
+ # Mini assembler
27
+ # ----------------------------------------------------------------------------
28
+
29
+ _OPCODE_NAMES = {
30
+ "add": 0x0, "sub": 0x1, "and": 0x2, "or": 0x3, "xor": 0x4,
31
+ "shl": 0x5, "shr": 0x6, "mul": 0x7, "div": 0x8, "cmp": 0x9,
32
+ "load": 0xA, "store": 0xB, "jmp": 0xC, "jcc": 0xD,
33
+ "call": 0xE, "halt": 0xF,
34
+ }
35
+
36
+ _COND = {"jz": 0, "jnz": 1, "jc": 2, "jnc": 3,
37
+ "jn": 4, "jp": 5, "jv": 6, "jnv": 7}
38
+
39
+
40
+ def _enc(opcode: int, rd: int = 0, rs: int = 0, imm: int = 0) -> int:
41
+ return ((opcode & 0xF) << 12) | ((rd & 0x3) << 10) | ((rs & 0x3) << 8) | (imm & 0xFF)
42
+
43
+
44
+ class Asm:
45
+ """Tiny assembler for the threshold-computer ISA.
46
+
47
+ Usage:
48
+ a = Asm(size=256)
49
+ a.org(0)
50
+ a.label("start")
51
+ a.load(0, "data")
52
+ a.halt()
53
+ a.org(0x80); a.label("data"); a.db(42)
54
+ mem = a.assemble()
55
+ """
56
+
57
+ def __init__(self, size: int):
58
+ self.mem: List[int] = [0] * size
59
+ self.pc: int = 0
60
+ self.labels: Dict[str, int] = {}
61
+ self.fixups: List[Tuple[int, str]] = []
62
+
63
+ def org(self, addr: int) -> None:
64
+ self.pc = addr
65
+
66
+ def label(self, name: str) -> None:
67
+ if name in self.labels:
68
+ raise ValueError(f"duplicate label: {name}")
69
+ self.labels[name] = self.pc
70
+
71
+ def db(self, *values: int) -> None:
72
+ for v in values:
73
+ self.mem[self.pc] = v & 0xFF
74
+ self.pc += 1
75
+
76
+ def dw(self, value: int) -> None:
77
+ self.mem[self.pc] = (value >> 8) & 0xFF
78
+ self.mem[self.pc + 1] = value & 0xFF
79
+ self.pc += 2
80
+
81
+ def daddr(self, label: str) -> None:
82
+ self.fixups.append((self.pc, label))
83
+ self.dw(0)
84
+
85
+ # --- ALU ops (no immediate) ---
86
+ def _alu(self, op: int, rd: int, rs: int) -> None:
87
+ self.dw(_enc(op, rd, rs))
88
+
89
+ def add(self, rd: int, rs: int) -> None: self._alu(0x0, rd, rs)
90
+ def sub(self, rd: int, rs: int) -> None: self._alu(0x1, rd, rs)
91
+ def and_(self, rd: int, rs: int) -> None: self._alu(0x2, rd, rs)
92
+ def or_(self, rd: int, rs: int) -> None: self._alu(0x3, rd, rs)
93
+ def xor(self, rd: int, rs: int) -> None: self._alu(0x4, rd, rs)
94
+ def shl(self, rd: int) -> None: self._alu(0x5, rd, 0)
95
+ def shr(self, rd: int) -> None: self._alu(0x6, rd, 0)
96
+ def mul(self, rd: int, rs: int) -> None: self._alu(0x7, rd, rs)
97
+ def cmp(self, rd: int, rs: int) -> None: self._alu(0x9, rd, rs)
98
+
99
+ # --- Memory + control (address-extended) ---
100
+ def load(self, rd: int, label: str) -> None:
101
+ self.dw(_enc(0xA, rd, 0)); self.daddr(label)
102
+
103
+ def store(self, rs: int, label: str) -> None:
104
+ self.dw(_enc(0xB, 0, rs)); self.daddr(label)
105
+
106
+ def jmp(self, label: str) -> None:
107
+ self.dw(_enc(0xC)); self.daddr(label)
108
+
109
+ def jcc(self, cond: str, label: str) -> None:
110
+ self.dw(_enc(0xD, 0, 0, _COND[cond])); self.daddr(label)
111
+
112
+ def jz(self, label: str) -> None: self.jcc("jz", label)
113
+ def jnz(self, label: str) -> None: self.jcc("jnz", label)
114
+ def jc(self, label: str) -> None: self.jcc("jc", label)
115
+ def jnc(self, label: str) -> None: self.jcc("jnc", label)
116
+ def jn(self, label: str) -> None: self.jcc("jn", label)
117
+ def jp(self, label: str) -> None: self.jcc("jp", label)
118
+ def jv(self, label: str) -> None: self.jcc("jv", label)
119
+ def jnv(self, label: str) -> None: self.jcc("jnv", label)
120
+
121
+ def call(self, label: str) -> None:
122
+ self.dw(_enc(0xE)); self.daddr(label)
123
+
124
+ def halt(self) -> None:
125
+ self.dw(_enc(0xF))
126
+
127
+ def assemble(self) -> List[int]:
128
+ for offset, label in self.fixups:
129
+ if label not in self.labels:
130
+ raise ValueError(f"undefined label: {label}")
131
+ target = self.labels[label]
132
+ self.mem[offset] = (target >> 8) & 0xFF
133
+ self.mem[offset + 1] = target & 0xFF
134
+ return self.mem
135
+
136
+
137
+ # ----------------------------------------------------------------------------
138
+ # Programs
139
+ # ----------------------------------------------------------------------------
140
+
141
+ ProgramResult = Tuple[List[int], Dict[int, int], int, str]
142
+
143
+
144
+ def fib(n: int = 11, mem_size: int = 256) -> ProgramResult:
145
+ """Iterative Fibonacci F(N), 8-bit wrap.
146
+
147
+ F(11) = 89. F(13) = 233 still fits in 8 bits; F(14) = 377 overflows.
148
+
149
+ Algorithm: maintain (a, b) = (F(k), F(k+1)); after N steps a = F(N).
150
+ Per iteration: temp=b, b=a+b, a=temp; n--.
151
+ """
152
+ a = Asm(mem_size)
153
+
154
+ a.org(0)
155
+ a.load(2, "n_addr") # R2 = n
156
+ a.load(0, "zero_addr") # R0 = 0 = F(0)
157
+ a.load(1, "one_addr") # R1 = 1 = F(1)
158
+ a.label("loop")
159
+ a.load(3, "zero_addr") # R3 = 0
160
+ a.cmp(2, 3) # n == 0?
161
+ a.jz("done")
162
+ a.load(3, "zero_addr") # R3 = 0
163
+ a.add(3, 1) # R3 = b (saved old b)
164
+ a.add(1, 0) # R1 = a + b (new b)
165
+ a.load(0, "zero_addr") # R0 = 0
166
+ a.add(0, 3) # R0 = old b (new a)
167
+ a.load(3, "one_addr") # R3 = 1
168
+ a.sub(2, 3) # n--
169
+ a.jmp("loop")
170
+ a.label("done")
171
+ a.store(0, "out") # OUT = a
172
+ a.halt()
173
+
174
+ a.org(0x80)
175
+ a.label("zero_addr"); a.db(0)
176
+ a.label("one_addr"); a.db(1)
177
+ a.label("n_addr"); a.db(n)
178
+ a.label("out"); a.db(0)
179
+
180
+ mem = a.assemble()
181
+ expected_a = 0
182
+ aa, bb = 0, 1
183
+ for _ in range(n):
184
+ aa, bb = bb, (aa + bb) & 0xFF
185
+ expected_a = aa
186
+ return mem, {a.labels["out"]: expected_a}, 16 * (n + 2), f"Fibonacci F({n}) = {expected_a}"
187
+
188
+
189
+ def sum_n(n: int = 10, mem_size: int = 256) -> ProgramResult:
190
+ """Compute 1 + 2 + ... + N using the Z flag from SUB to terminate.
191
+
192
+ No explicit zero register required; SUB sets Z when its result is zero.
193
+ R0 = accumulator, R1 = counter (n down to 1), R2 = 1.
194
+ """
195
+ a = Asm(mem_size)
196
+
197
+ a.org(0)
198
+ a.load(0, "zero_addr") # acc = 0
199
+ a.load(1, "n_addr") # i = n
200
+ a.load(2, "one_addr") # 1
201
+ a.label("loop")
202
+ a.add(0, 1) # acc += i
203
+ a.sub(1, 2) # i-- (Z set when i becomes 0)
204
+ a.jnz("loop")
205
+ a.store(0, "out")
206
+ a.halt()
207
+
208
+ a.org(0x80)
209
+ a.label("zero_addr"); a.db(0)
210
+ a.label("one_addr"); a.db(1)
211
+ a.label("n_addr"); a.db(n)
212
+ a.label("out"); a.db(0)
213
+
214
+ mem = a.assemble()
215
+ expected = (n * (n + 1) // 2) & 0xFF
216
+ return mem, {a.labels["out"]: expected}, 4 + 4 * n, f"sum 1..{n} = {expected}"
217
+
218
+
219
+ def self_mod_jmp(mem_size: int = 256) -> ProgramResult:
220
+ """Self-modifying code: writes the JMP target's low byte at runtime.
221
+
222
+ Initial JMP target is path_a (writes 0xAA to OUT). The code first
223
+ overwrites the JMP's address-word LSB so it points to path_b
224
+ (writes 0xBB). Successful execution lands at path_b, so OUT = 0xBB.
225
+ """
226
+ a = Asm(mem_size)
227
+
228
+ a.org(0)
229
+ a.label("start")
230
+
231
+ # Forward-declare the_jmp's LSB address. The first two instructions are
232
+ # each 4 bytes; the_jmp follows them, so the_jmp = pc + 8 from start.
233
+ # The JMP's address word is at the_jmp + 2; the LSB byte is at the_jmp + 3.
234
+ a.labels["jmp_target_lsb"] = a.pc + 8 + 3
235
+
236
+ a.load(0, "new_lsb") # R0 = LSB of path_b address (4 bytes)
237
+ a.store(0, "jmp_target_lsb") # patch the JMP's LSB (4 bytes)
238
+ a.label("the_jmp")
239
+ a.jmp("path_a") # initially -> path_a; after patch -> path_b
240
+ a.label("path_a")
241
+ a.load(1, "val_a"); a.store(1, "out"); a.halt()
242
+ a.label("path_b")
243
+ a.load(1, "val_b"); a.store(1, "out"); a.halt()
244
+
245
+ a.org(0x80)
246
+ a.label("val_a"); a.db(0xAA)
247
+ a.label("val_b"); a.db(0xBB)
248
+ # new_lsb_word stores path_b's full 16-bit address; new_lsb labels the LSB byte.
249
+ a.label("new_lsb_word"); a.daddr("path_b")
250
+ a.labels["new_lsb"] = a.labels["new_lsb_word"] + 1
251
+ a.label("out"); a.db(0)
252
+
253
+ mem = a.assemble()
254
+ return mem, {a.labels["out"]: 0xBB}, 30, "self-modifying JMP target"
255
+
256
+
257
+ def all_branches(mem_size: int = 256) -> ProgramResult:
258
+ """Drive all eight conditional jumps; each path writes a unique marker.
259
+
260
+ Test plan (each step sets flags, then a Jcc; the branch takes when
261
+ expected and the corresponding marker is written):
262
+
263
+ JZ: CMP equal -> Z=1 -> taken -> M[OUT0] = 0xA0
264
+ JNZ: CMP unequal -> Z=0 -> taken -> M[OUT1] = 0xA1
265
+ JC: ADD overflow (255+1) -> C=1 -> taken -> M[OUT2] = 0xA2
266
+ JNC: ADD no overflow -> C=0 -> taken -> M[OUT3] = 0xA3
267
+ JN: SUB result 0xFF (n=1) -> N=1 -> taken -> M[OUT4] = 0xA4
268
+ JP: ADD result 1 -> N=0 -> taken -> M[OUT5] = 0xA5
269
+ JV: ADD signed overflow (127+1=128) -> V=1 -> taken -> M[OUT6] = 0xA6
270
+ JNV: ADD no signed overflow (1+1=2) -> V=0 -> taken -> M[OUT7] = 0xA7
271
+
272
+ A failure on any branch causes the wrong (or no) marker to be written.
273
+ """
274
+ a = Asm(mem_size)
275
+
276
+ a.org(0)
277
+ # ----- JZ: equal compare -> Z=1 -----
278
+ a.load(0, "v5"); a.load(1, "v5"); a.cmp(0, 1); a.jz("ok_jz"); a.jmp("fail")
279
+ a.label("ok_jz"); a.load(2, "m_a0"); a.store(2, "out0")
280
+
281
+ # ----- JNZ: unequal compare -> Z=0 -----
282
+ a.load(0, "v5"); a.load(1, "v3"); a.cmp(0, 1); a.jnz("ok_jnz"); a.jmp("fail")
283
+ a.label("ok_jnz"); a.load(2, "m_a1"); a.store(2, "out1")
284
+
285
+ # ----- JC: 255+1 = 0 with carry -----
286
+ a.load(0, "v255"); a.load(1, "v1"); a.add(0, 1); a.jc("ok_jc"); a.jmp("fail")
287
+ a.label("ok_jc"); a.load(2, "m_a2"); a.store(2, "out2")
288
+
289
+ # ----- JNC: 1+1 = 2, no carry -----
290
+ a.load(0, "v1"); a.load(1, "v1"); a.add(0, 1); a.jnc("ok_jnc"); a.jmp("fail")
291
+ a.label("ok_jnc"); a.load(2, "m_a3"); a.store(2, "out3")
292
+
293
+ # ----- JN: 0 - 1 = 0xFF, MSB set -----
294
+ a.load(0, "v0"); a.load(1, "v1"); a.sub(0, 1); a.jn("ok_jn"); a.jmp("fail")
295
+ a.label("ok_jn"); a.load(2, "m_a4"); a.store(2, "out4")
296
+
297
+ # ----- JP: 0 + 1 = 1, MSB clear -----
298
+ a.load(0, "v0"); a.load(1, "v1"); a.add(0, 1); a.jp("ok_jp"); a.jmp("fail")
299
+ a.label("ok_jp"); a.load(2, "m_a5"); a.store(2, "out5")
300
+
301
+ # ----- JV: 127 + 1 = 128, signed overflow -----
302
+ a.load(0, "v127"); a.load(1, "v1"); a.add(0, 1); a.jv("ok_jv"); a.jmp("fail")
303
+ a.label("ok_jv"); a.load(2, "m_a6"); a.store(2, "out6")
304
+
305
+ # ----- JNV: 1 + 1 = 2, no signed overflow -----
306
+ a.load(0, "v1"); a.load(1, "v1"); a.add(0, 1); a.jnv("ok_jnv"); a.jmp("fail")
307
+ a.label("ok_jnv"); a.load(2, "m_a7"); a.store(2, "out7")
308
+
309
+ a.jmp("end")
310
+ a.label("fail")
311
+ a.load(2, "v_fail"); a.store(2, "fail_addr"); a.halt()
312
+ a.label("end")
313
+ a.halt()
314
+
315
+ # Code runs to ~0xDF; data starts safely after that.
316
+ a.org(0xE0)
317
+ a.label("v0"); a.db(0)
318
+ a.label("v1"); a.db(1)
319
+ a.label("v3"); a.db(3)
320
+ a.label("v5"); a.db(5)
321
+ a.label("v127"); a.db(127)
322
+ a.label("v255"); a.db(255)
323
+ a.label("m_a0"); a.db(0xA0)
324
+ a.label("m_a1"); a.db(0xA1)
325
+ a.label("m_a2"); a.db(0xA2)
326
+ a.label("m_a3"); a.db(0xA3)
327
+ a.label("m_a4"); a.db(0xA4)
328
+ a.label("m_a5"); a.db(0xA5)
329
+ a.label("m_a6"); a.db(0xA6)
330
+ a.label("m_a7"); a.db(0xA7)
331
+ a.label("v_fail"); a.db(0xEE)
332
+ a.label("out0"); a.db(0)
333
+ a.label("out1"); a.db(0)
334
+ a.label("out2"); a.db(0)
335
+ a.label("out3"); a.db(0)
336
+ a.label("out4"); a.db(0)
337
+ a.label("out5"); a.db(0)
338
+ a.label("out6"); a.db(0)
339
+ a.label("out7"); a.db(0)
340
+ a.label("fail_addr"); a.db(0)
341
+
342
+ mem = a.assemble()
343
+ expected = {
344
+ a.labels["out0"]: 0xA0,
345
+ a.labels["out1"]: 0xA1,
346
+ a.labels["out2"]: 0xA2,
347
+ a.labels["out3"]: 0xA3,
348
+ a.labels["out4"]: 0xA4,
349
+ a.labels["out5"]: 0xA5,
350
+ a.labels["out6"]: 0xA6,
351
+ a.labels["out7"]: 0xA7,
352
+ a.labels["fail_addr"]: 0, # must remain zero
353
+ }
354
+ return mem, expected, 200, "all 8 conditional jumps (JZ/JNZ/JC/JNC/JN/JP/JV/JNV)"
355
+
356
+
357
+ def call_pushes_pc(mem_size: int = 256) -> ProgramResult:
358
+ """Verify CALL pushes the return address (next-instruction PC) onto the stack.
359
+
360
+ SP starts at addr_mask (mem_size - 1). CALL decrements SP and writes the
361
+ return-address high byte, decrements again and writes the low byte. After
362
+ HALT we expect:
363
+ - mem[addr_mask - 2] = low byte of the return address
364
+ - mem[addr_mask - 1] = high byte of the return address
365
+ - the no-return code path was NOT taken
366
+ - the callee was reached
367
+ """
368
+ a = Asm(mem_size)
369
+
370
+ a.org(0)
371
+ a.label("caller")
372
+ a.load(0, "marker_val")
373
+ a.store(0, "marker_addr") # write before CALL
374
+ a.label("call_site")
375
+ a.call("callee")
376
+ # If CALL did not transfer control, this fallthrough store would write 0xDD:
377
+ a.load(0, "noret_val")
378
+ a.store(0, "noret_addr")
379
+ a.halt()
380
+
381
+ a.label("callee")
382
+ a.load(0, "callee_val")
383
+ a.store(0, "callee_addr")
384
+ a.halt()
385
+
386
+ a.org(0x40)
387
+ a.label("marker_val"); a.db(0x11)
388
+ a.label("marker_addr"); a.db(0)
389
+ a.label("noret_val"); a.db(0xDD)
390
+ a.label("noret_addr"); a.db(0)
391
+ a.label("callee_val"); a.db(0x22)
392
+ a.label("callee_addr"); a.db(0)
393
+
394
+ mem = a.assemble()
395
+ addr_mask = mem_size - 1
396
+ return_addr = a.labels["call_site"] + 4 # 4-byte CALL instruction
397
+ expected = {
398
+ a.labels["marker_addr"]: 0x11, # pre-CALL store ran
399
+ a.labels["callee_addr"]: 0x22, # callee reached
400
+ a.labels["noret_addr"]: 0, # fallthrough did not run
401
+ (addr_mask - 2) & addr_mask: return_addr & 0xFF, # ret LSB on stack
402
+ (addr_mask - 1) & addr_mask: (return_addr >> 8) & 0xFF, # ret MSB
403
+ }
404
+ return mem, expected, 30, "CALL pushes return PC onto stack"
405
+
406
+
407
+ def bubble_sort_4(mem_size: int = 256) -> ProgramResult:
408
+ """Sort a 4-byte array using unrolled compare-swap (3 passes of 3 compares).
409
+
410
+ Algorithm: bubble sort, fully unrolled (no inner loops). Each compare-swap
411
+ is 8 instructions; 3 outer passes x 3 inner positions = 9 swaps -> ~72 instrs.
412
+
413
+ For each position i in (0,1,2):
414
+ if A[i] > A[i+1]:
415
+ tmp = A[i]; A[i] = A[i+1]; A[i+1] = tmp
416
+ Repeat 3 times -> sorted ascending.
417
+ """
418
+ a = Asm(mem_size)
419
+
420
+ addrs = ["a0", "a1", "a2", "a3"]
421
+
422
+ a.org(0)
423
+ for _outer in range(3):
424
+ for i in range(3):
425
+ x, y = addrs[i], addrs[i + 1]
426
+ # Load pair
427
+ a.load(0, x) # R0 = A[i]
428
+ a.load(1, y) # R1 = A[i+1]
429
+ a.cmp(0, 1) # compare A[i] - A[i+1]
430
+ # If A[i] <= A[i+1] (Z=1 or C=0 from sub-style cmp), skip swap
431
+ # SUB sets carry when no borrow (a >= b). So:
432
+ # a > b iff Z=0 and a >= b -> Z=0 and C=1 (the sub didn't borrow)
433
+ # We want to swap when a > b. JNC (no carry / borrow) means a < b -> skip swap.
434
+ # If a == b (Z=1) we also skip. So: jump-skip when JZ OR JNC.
435
+ # Easier: compute (a > b) by checking C=1 AND Z=0. Use JZ to skip on equal,
436
+ # then JNC to skip on a < b. Otherwise fall through to swap.
437
+ skip_lbl = f"skip_{_outer}_{i}"
438
+ a.jz(skip_lbl) # equal -> skip
439
+ a.jnc(skip_lbl) # a < b (sub borrowed) -> skip
440
+ # swap: store R0 -> y, R1 -> x
441
+ a.store(1, x)
442
+ a.store(0, y)
443
+ a.label(skip_lbl)
444
+ a.halt()
445
+
446
+ # Initial unsorted array; code runs to ~0xEC, so data starts at 0xF0.
447
+ initial = [42, 7, 200, 19]
448
+ a.org(0xF0)
449
+ for name, val in zip(addrs, initial):
450
+ a.label(name); a.db(val)
451
+
452
+ mem = a.assemble()
453
+ sorted_vals = sorted(initial)
454
+ expected = {a.labels[name]: v for name, v in zip(addrs, sorted_vals)}
455
+ return mem, expected, 800, f"bubble sort {initial} -> {sorted_vals}"
456
+
457
+
458
+ def cross_check_mul(mem_size: int = 256) -> ProgramResult:
459
+ """Cross-check the threshold MUL circuit against repeated ADD.
460
+
461
+ Computes A * B two ways and stores both results:
462
+ repeated add: acc = 0; loop B times: acc += A -> OUT_ADD
463
+ direct: R0 = A; MUL R0, R1 (R1 = B) -> OUT_MUL
464
+ The test verifies the two agree with the expected product.
465
+ """
466
+ a = Asm(mem_size)
467
+ A_VAL = 17
468
+ B_VAL = 9
469
+ expected_product = (A_VAL * B_VAL) & 0xFF
470
+
471
+ a.org(0)
472
+ # --- direct multiply ---
473
+ a.load(0, "A")
474
+ a.load(1, "B")
475
+ a.mul(0, 1)
476
+ a.store(0, "out_mul")
477
+ # --- repeated-add multiply ---
478
+ a.load(0, "zero") # acc = 0
479
+ a.load(1, "A") # addend
480
+ a.load(2, "B") # counter
481
+ a.load(3, "one") # 1
482
+ a.label("rep_loop")
483
+ a.add(0, 1) # acc += A
484
+ a.sub(2, 3) # B--
485
+ a.jnz("rep_loop")
486
+ a.store(0, "out_add")
487
+ a.halt()
488
+
489
+ a.org(0x80)
490
+ a.label("A"); a.db(A_VAL)
491
+ a.label("B"); a.db(B_VAL)
492
+ a.label("zero"); a.db(0)
493
+ a.label("one"); a.db(1)
494
+ a.label("out_mul"); a.db(0)
495
+ a.label("out_add"); a.db(0)
496
+
497
+ mem = a.assemble()
498
+ expected = {
499
+ a.labels["out_mul"]: expected_product,
500
+ a.labels["out_add"]: expected_product,
501
+ }
502
+ return mem, expected, 80, f"MUL vs repeated ADD: {A_VAL} * {B_VAL} = {expected_product}"
503
+
504
+
505
+ def div_via_repeated_sub(mem_size: int = 256) -> ProgramResult:
506
+ """Compute floor(A/B) and (A mod B) by repeated subtraction.
507
+
508
+ Loop: while A >= B { A -= B; quotient += 1 }
509
+ Uses CMP + JC (carry-set on no-borrow), SUB, ADD, JMP, STORE, HALT.
510
+
511
+ Cross-checked against the on-chip 8-bit DIV opcode (0x8) via a
512
+ second pass that uses DIV directly. Both quotients written to OUT
513
+ locations; the test verifies they match.
514
+ """
515
+ A_VAL = 100
516
+ B_VAL = 7
517
+ expected_q = A_VAL // B_VAL # 14
518
+ expected_r = A_VAL % B_VAL # 2
519
+
520
+ a = Asm(mem_size)
521
+
522
+ a.org(0)
523
+ # ---- Repeated-subtraction division ----
524
+ a.load(0, "A") # R0 = A (will become remainder)
525
+ a.load(1, "B") # R1 = B (divisor)
526
+ a.load(2, "ZERO") # R2 = 0 (will become quotient)
527
+ a.load(3, "ONE") # R3 = 1 (increment)
528
+
529
+ a.label("loop")
530
+ a.cmp(0, 1) # CMP R0, R1; carry=1 (no-borrow) iff R0 >= R1
531
+ a.jnc("done") # if R0 < R1 (carry=0), exit loop
532
+ a.sub(0, 1) # R0 -= B
533
+ a.add(2, 3) # quotient += 1
534
+ a.jmp("loop")
535
+
536
+ a.label("done")
537
+ a.store(2, "OUT_Q_RPT") # quotient via repeated sub
538
+ a.store(0, "OUT_R_RPT") # remainder via repeated sub
539
+
540
+ # ---- Direct DIV opcode for cross-check ----
541
+ a.load(0, "A")
542
+ a.load(1, "B")
543
+ a.dw(_enc(0x8, 0, 1, 0)) # DIV R0, R1 -> R0 = R0 / R1 (8-bit DIV)
544
+ a.store(0, "OUT_Q_DIV")
545
+ a.halt()
546
+
547
+ a.org(0x80)
548
+ a.label("A"); a.db(A_VAL)
549
+ a.label("B"); a.db(B_VAL)
550
+ a.label("ZERO"); a.db(0)
551
+ a.label("ONE"); a.db(1)
552
+ a.label("OUT_Q_RPT"); a.db(0)
553
+ a.label("OUT_R_RPT"); a.db(0)
554
+ a.label("OUT_Q_DIV"); a.db(0)
555
+
556
+ mem = a.assemble()
557
+ expected = {
558
+ a.labels["OUT_Q_RPT"]: expected_q,
559
+ a.labels["OUT_R_RPT"]: expected_r,
560
+ a.labels["OUT_Q_DIV"]: expected_q,
561
+ }
562
+ return mem, expected, 4 * (A_VAL // B_VAL + 4) + 12, (
563
+ f"{A_VAL} / {B_VAL}: quotient {expected_q} (repeated SUB) "
564
+ f"matches DIV opcode result; remainder {expected_r}"
565
+ )
566
+
567
+
568
+ def bitwise_chain(mem_size: int = 256) -> ProgramResult:
569
+ """Run a chain of bitwise ops and verify each intermediate value.
570
+
571
+ Sequence:
572
+ R0 = A & B (AND)
573
+ R0 = R0 | C (OR)
574
+ R0 = R0 ^ D (XOR)
575
+ R0 = R0 << 1 (SHL)
576
+ R0 = R0 >> 1 (SHR)
577
+ Stores R0 after each step. Verifies all intermediate values to
578
+ catch any single-op regression.
579
+ """
580
+ A = 0xCC # 11001100
581
+ B = 0xF0 # 11110000
582
+ C = 0x0F # 00001111
583
+ D = 0xAA # 10101010
584
+
585
+ s1 = A & B # 0xC0
586
+ s2 = s1 | C # 0xCF
587
+ s3 = s2 ^ D # 0x65
588
+ s4 = (s3 << 1) & 0xFF # 0xCA
589
+ s5 = s4 >> 1 # 0x65
590
+
591
+ a = Asm(mem_size)
592
+ a.org(0)
593
+ a.load(0, "A"); a.load(1, "B"); a.and_(0, 1); a.store(0, "S1")
594
+ a.load(1, "C"); a.or_(0, 1); a.store(0, "S2")
595
+ a.load(1, "D"); a.xor(0, 1); a.store(0, "S3")
596
+ a.shl(0); a.store(0, "S4")
597
+ a.shr(0); a.store(0, "S5")
598
+ a.halt()
599
+
600
+ a.org(0x80)
601
+ a.label("A"); a.db(A)
602
+ a.label("B"); a.db(B)
603
+ a.label("C"); a.db(C)
604
+ a.label("D"); a.db(D)
605
+ a.label("S1"); a.db(0)
606
+ a.label("S2"); a.db(0)
607
+ a.label("S3"); a.db(0)
608
+ a.label("S4"); a.db(0)
609
+ a.label("S5"); a.db(0)
610
+
611
+ mem = a.assemble()
612
+ expected = {
613
+ a.labels["S1"]: s1,
614
+ a.labels["S2"]: s2,
615
+ a.labels["S3"]: s3,
616
+ a.labels["S4"]: s4,
617
+ a.labels["S5"]: s5,
618
+ }
619
+ return mem, expected, 30, (
620
+ f"bitwise chain AND/OR/XOR/SHL/SHR -> {s1:#x},{s2:#x},{s3:#x},{s4:#x},{s5:#x}"
621
+ )
622
+
623
+
624
+ def flags_policy(mem_size: int = 256) -> ProgramResult:
625
+ """Pin the architectural flag policy: only ADD, SUB, MUL, and CMP write
626
+ FLAGS; bitwise ops, shifts, DIV, LOAD, and STORE leave them unchanged.
627
+
628
+ Three probes, each setting flags with an arithmetic op, disturbing the
629
+ machine with a non-flag-writing op, then branching on the original flags:
630
+
631
+ 1. CMP equal -> Z=1; AND producing a nonzero result; JZ must still take.
632
+ 2. CMP unequal -> Z=0; XOR producing a zero result; JZ must NOT take.
633
+ 3. ADD 255+1 -> C=1; SHR (shifts a 0 out); JC must still take.
634
+
635
+ Any runtime that lets the intervening op rewrite FLAGS lands on the
636
+ wrong path and writes the 0xEE failure marker instead.
637
+ """
638
+ a = Asm(mem_size)
639
+
640
+ a.org(0)
641
+ # ----- Probe 1: Z=1 survives AND (nonzero result) -----
642
+ a.load(0, "v5"); a.load(1, "v5"); a.cmp(0, 1) # Z=1
643
+ a.load(2, "v3"); a.load(3, "v1"); a.and_(2, 3) # R2 = 3 & 1 = 1
644
+ a.jz("p1_ok")
645
+ a.load(2, "m_fail"); a.store(2, "out1"); a.jmp("p2")
646
+ a.label("p1_ok")
647
+ a.load(2, "m_b1"); a.store(2, "out1")
648
+
649
+ # ----- Probe 2: Z=0 survives XOR (zero result) -----
650
+ a.label("p2")
651
+ a.load(0, "v5"); a.load(1, "v3"); a.cmp(0, 1) # Z=0
652
+ a.xor(2, 2) # R2 = 0
653
+ a.jz("p2_fail")
654
+ a.load(2, "m_b2"); a.store(2, "out2"); a.jmp("p3")
655
+ a.label("p2_fail")
656
+ a.load(2, "m_fail"); a.store(2, "out2")
657
+
658
+ # ----- Probe 3: C=1 survives SHR -----
659
+ a.label("p3")
660
+ a.load(0, "v255"); a.load(1, "v1"); a.add(0, 1) # C=1
661
+ a.load(2, "v3"); a.shr(2) # R2 = 1
662
+ a.jc("p3_ok")
663
+ a.load(2, "m_fail"); a.store(2, "out3"); a.jmp("done")
664
+ a.label("p3_ok")
665
+ a.load(2, "m_b3"); a.store(2, "out3")
666
+ a.label("done")
667
+ a.halt()
668
+
669
+ a.org(0xA0)
670
+ a.label("v1"); a.db(1)
671
+ a.label("v3"); a.db(3)
672
+ a.label("v5"); a.db(5)
673
+ a.label("v255"); a.db(255)
674
+ a.label("m_b1"); a.db(0xB1)
675
+ a.label("m_b2"); a.db(0xB2)
676
+ a.label("m_b3"); a.db(0xB3)
677
+ a.label("m_fail"); a.db(0xEE)
678
+ a.label("out1"); a.db(0)
679
+ a.label("out2"); a.db(0)
680
+ a.label("out3"); a.db(0)
681
+
682
+ mem = a.assemble()
683
+ expected = {
684
+ a.labels["out1"]: 0xB1,
685
+ a.labels["out2"]: 0xB2,
686
+ a.labels["out3"]: 0xB3,
687
+ }
688
+ return mem, expected, 60, "flag policy: AND/XOR/SHR leave FLAGS unchanged"
689
+
690
+
691
+ SUITE = [
692
+ ("fib", lambda mem_size: fib(11, mem_size)),
693
+ ("sum_n", lambda mem_size: sum_n(10, mem_size)),
694
+ ("self_mod_jmp", lambda mem_size: self_mod_jmp(mem_size)),
695
+ ("all_branches", lambda mem_size: all_branches(mem_size)),
696
+ ("call_pushes_pc", lambda mem_size: call_pushes_pc(mem_size)),
697
+ ("bubble_sort_4", lambda mem_size: bubble_sort_4(mem_size)),
698
+ ("cross_check_mul", lambda mem_size: cross_check_mul(mem_size)),
699
+ ("div_via_repeated_sub", lambda mem_size: div_via_repeated_sub(mem_size)),
700
+ ("bitwise_chain", lambda mem_size: bitwise_chain(mem_size)),
701
+ ("flags_policy", lambda mem_size: flags_policy(mem_size)),
702
+ ]
703
+
eval.py ADDED
The diff for this file is too large to render. See raw diff
 
eval_all.py ADDED
@@ -0,0 +1,850 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Unified evaluation harness for any threshold-computer variant.
3
+
4
+ Drops the `--cpu-test` smoke test (which was hardcoded to 16-bit/64KB) and
5
+ adds variant-aware sweep modes. The same harness handles every (data_bits,
6
+ addr_bits) configuration: it reads the manifest from each safetensors file,
7
+ runs the BatchedFitnessEvaluator at the right device, and reports per-file
8
+ plus per-category results.
9
+
10
+ Usage:
11
+ python eval_all.py path/to/file.safetensors # one file
12
+ python eval_all.py variants/ # every .safetensors in dir
13
+ python eval_all.py --device cpu variants/ # CPU only (default)
14
+ python eval_all.py --pop_size 32 variants/ # batched pop eval
15
+ python eval_all.py --debug path/to/file.safetensors # per-circuit detail
16
+ python eval_all.py --cpu-program PATH # also run an assembled program
17
+ # through the threshold CPU
18
+ # sized to the file's manifest
19
+
20
+ Exit code:
21
+ 0 if all files PASS (fitness >= 0.9999)
22
+ N where N is the number of FAILing files
23
+ """
24
+
25
+ from __future__ import annotations
26
+
27
+ import argparse
28
+ import json
29
+ import os
30
+ import sys
31
+ import time
32
+ from pathlib import Path
33
+ from typing import Any, Dict, List, Optional, Tuple
34
+
35
+ import torch
36
+ from safetensors import safe_open
37
+
38
+ # Reuse eval.py's evaluator (variant-aware)
39
+ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
40
+ from eval import (
41
+ BatchedFitnessEvaluator,
42
+ create_population,
43
+ load_model,
44
+ get_manifest,
45
+ heaviside,
46
+ int_to_bits,
47
+ bits_to_int,
48
+ bits_msb_to_lsb,
49
+ )
50
+
51
+
52
+ # ---------------------------------------------------------------------------
53
+ # Variant-aware threshold ALU + CPU
54
+ # ---------------------------------------------------------------------------
55
+
56
+ class GenericThresholdALU:
57
+ """Variant-aware threshold ALU walking the gates of a loaded variant.
58
+
59
+ 8-bit primitives: add8/sub8/mul8/div8, packed bitwise and8/or8/xor8,
60
+ shl8/shr8, and the bit-cascade comparators via cmp8. Width-generic
61
+ add_n/sub_n/mul_n/cmp_n cover the 16- and 32-bit circuit families.
62
+ """
63
+
64
+ def __init__(self, tensors: Dict[str, torch.Tensor], data_bits: int):
65
+ self.T = tensors
66
+ self.data_bits = data_bits
67
+
68
+ def _g(self, name, inputs):
69
+ w = self.T[name + ".weight"].view(-1)
70
+ b = self.T[name + ".bias"].view(-1)
71
+ return int(heaviside((torch.tensor(inputs, dtype=torch.float32) * w).sum() + b).item())
72
+
73
+ def _xor_or_nand(self, prefix, inputs):
74
+ a, b_ = inputs
75
+ h_or = self._g(f"{prefix}.layer1.or", [a, b_])
76
+ h_nand = self._g(f"{prefix}.layer1.nand", [a, b_])
77
+ return self._g(f"{prefix}.layer2", [h_or, h_nand])
78
+
79
+ def _fa(self, prefix, a, b, cin):
80
+ s1 = self._xor_or_nand(f"{prefix}.ha1.sum", [a, b])
81
+ c1 = self._g(f"{prefix}.ha1.carry", [a, b])
82
+ s2 = self._xor_or_nand(f"{prefix}.ha2.sum", [s1, cin])
83
+ c2 = self._g(f"{prefix}.ha2.carry", [s1, cin])
84
+ cout = self._g(f"{prefix}.carry_or", [c1, c2])
85
+ return s2, cout
86
+
87
+ def add8(self, a, b):
88
+ a_lsb = list(reversed(int_to_bits(a, 8)))
89
+ b_lsb = list(reversed(int_to_bits(b, 8)))
90
+ carry = 0
91
+ s_lsb = []
92
+ for i in range(8):
93
+ s, carry = self._fa(f"arithmetic.ripplecarry8bit.fa{i}", a_lsb[i], b_lsb[i], carry)
94
+ s_lsb.append(s)
95
+ return bits_to_int(list(reversed(s_lsb))), carry
96
+
97
+ def sub8(self, a, b):
98
+ a_lsb = list(reversed(int_to_bits(a, 8)))
99
+ b_lsb = list(reversed(int_to_bits(b, 8)))
100
+ carry = 1
101
+ d_lsb = []
102
+ for i in range(8):
103
+ notb = self._g(f"arithmetic.sub8bit.notb{i}", [b_lsb[i]])
104
+ x1 = self._xor_or_nand(f"arithmetic.sub8bit.fa{i}.xor1", [a_lsb[i], notb])
105
+ x2 = self._xor_or_nand(f"arithmetic.sub8bit.fa{i}.xor2", [x1, carry])
106
+ and1 = self._g(f"arithmetic.sub8bit.fa{i}.and1", [a_lsb[i], notb])
107
+ and2 = self._g(f"arithmetic.sub8bit.fa{i}.and2", [x1, carry])
108
+ carry = self._g(f"arithmetic.sub8bit.fa{i}.or_carry", [and1, and2])
109
+ d_lsb.append(x2)
110
+ return bits_to_int(list(reversed(d_lsb))), carry
111
+
112
+ _CMP_KIND = {"greaterthan": "gt", "lessthan": "lt", "eq": "eq", "equality": "eq",
113
+ "greaterorequal": "ge", "lessorequal": "le"}
114
+
115
+ def _cmp_bit_cascade(self, cmp_prefix: str, finals: Dict[str, str],
116
+ a_bits_msb: List[int], b_bits_msb: List[int], kind: str) -> int:
117
+ """Walk an add_bit_cascade_compare structure (bit 0 is the MSB)."""
118
+ bits = len(a_bits_msb)
119
+ bit_gt, bit_lt, bit_eq = [], [], []
120
+ for i in range(bits):
121
+ ab = [a_bits_msb[i], b_bits_msb[i]]
122
+ bit_gt.append(self._g(f"{cmp_prefix}.bit{i}.gt", ab))
123
+ bit_lt.append(self._g(f"{cmp_prefix}.bit{i}.lt", ab))
124
+ eq_and = self._g(f"{cmp_prefix}.bit{i}.eq.layer1.and", ab)
125
+ eq_nor = self._g(f"{cmp_prefix}.bit{i}.eq.layer1.nor", ab)
126
+ bit_eq.append(self._g(f"{cmp_prefix}.bit{i}.eq", [eq_and, eq_nor]))
127
+ cas_gt, cas_lt = [bit_gt[0]], [bit_lt[0]]
128
+ for i in range(1, bits):
129
+ eq_pref = self._g(f"{cmp_prefix}.cascade.eq_prefix.bit{i}", bit_eq[:i])
130
+ cas_gt.append(self._g(f"{cmp_prefix}.cascade.gt.bit{i}", [eq_pref, bit_gt[i]]))
131
+ cas_lt.append(self._g(f"{cmp_prefix}.cascade.lt.bit{i}", [eq_pref, bit_lt[i]]))
132
+ if kind == "gt":
133
+ return self._g(finals["gt"], cas_gt)
134
+ if kind == "lt":
135
+ return self._g(finals["lt"], cas_lt)
136
+ if kind == "eq":
137
+ return self._g(finals["eq"], bit_eq)
138
+ if kind == "ge":
139
+ lt = self._g(finals["lt"], cas_lt)
140
+ not_lt = self._g(finals["ge"] + ".not_lt", [lt])
141
+ return self._g(finals["ge"], [not_lt])
142
+ if kind == "le":
143
+ gt = self._g(finals["gt"], cas_gt)
144
+ not_gt = self._g(finals["le"] + ".not_gt", [gt])
145
+ return self._g(finals["le"], [not_gt])
146
+ raise ValueError(kind)
147
+
148
+ def cmp8(self, a, b, kind):
149
+ return self.cmp_n(a, b, kind, 8)
150
+
151
+ def mul8(self, a, b):
152
+ ab = int_to_bits(a, 8)
153
+ bb = int_to_bits(b, 8)
154
+ result = 0
155
+ for j in range(8):
156
+ if bb[j] == 0:
157
+ continue
158
+ row = 0
159
+ for i in range(8):
160
+ pp = self._g(f"alu.alu8bit.mul.pp.a{i}b{j}", [ab[i], bb[j]])
161
+ row |= (pp << (7 - i))
162
+ shift = 7 - j
163
+ result, _ = self.add8(result & 0xFF, (row << shift) & 0xFF)
164
+ return result & 0xFF
165
+
166
+ def _packed_pair_op(self, name: str, a: int, b: int) -> int:
167
+ """8 parallel 2-input gates stored packed: weight [16], bias [8]."""
168
+ a_bits = int_to_bits(a, 8)
169
+ b_bits = int_to_bits(b, 8)
170
+ w = self.T[f"{name}.weight"].view(-1)
171
+ bias = self.T[f"{name}.bias"].view(-1)
172
+ out = []
173
+ for bit in range(8):
174
+ inp = torch.tensor([float(a_bits[bit]), float(b_bits[bit])])
175
+ out.append(int(heaviside((inp * w[bit * 2:bit * 2 + 2]).sum() + bias[bit]).item()))
176
+ return bits_to_int(out)
177
+
178
+ def and8(self, a, b):
179
+ return self._packed_pair_op("alu.alu8bit.and", a, b)
180
+
181
+ def or8(self, a, b):
182
+ return self._packed_pair_op("alu.alu8bit.or", a, b)
183
+
184
+ def xor8(self, a, b):
185
+ a_bits = int_to_bits(a, 8)
186
+ b_bits = int_to_bits(b, 8)
187
+ w_or = self.T["alu.alu8bit.xor.layer1.or.weight"].view(-1)
188
+ b_or = self.T["alu.alu8bit.xor.layer1.or.bias"].view(-1)
189
+ w_nand = self.T["alu.alu8bit.xor.layer1.nand.weight"].view(-1)
190
+ b_nand = self.T["alu.alu8bit.xor.layer1.nand.bias"].view(-1)
191
+ w2 = self.T["alu.alu8bit.xor.layer2.weight"].view(-1)
192
+ b2 = self.T["alu.alu8bit.xor.layer2.bias"].view(-1)
193
+ out = []
194
+ for bit in range(8):
195
+ inp = torch.tensor([float(a_bits[bit]), float(b_bits[bit])])
196
+ h_or = heaviside((inp * w_or[bit * 2:bit * 2 + 2]).sum() + b_or[bit])
197
+ h_nand = heaviside((inp * w_nand[bit * 2:bit * 2 + 2]).sum() + b_nand[bit])
198
+ hidden = torch.stack([h_or, h_nand])
199
+ out.append(int(heaviside((hidden * w2[bit * 2:bit * 2 + 2]).sum() + b2[bit]).item()))
200
+ return bits_to_int(out)
201
+
202
+ def shl8(self, a):
203
+ a_bits = int_to_bits(a, 8)
204
+ out = []
205
+ for bit in range(8):
206
+ src = float(a_bits[bit + 1]) if bit < 7 else 0.0
207
+ out.append(self._g(f"alu.alu8bit.shl.bit{bit}", [src]))
208
+ return bits_to_int(out)
209
+
210
+ def shr8(self, a):
211
+ a_bits = int_to_bits(a, 8)
212
+ out = []
213
+ for bit in range(8):
214
+ src = float(a_bits[bit - 1]) if bit > 0 else 0.0
215
+ out.append(self._g(f"alu.alu8bit.shr.bit{bit}", [src]))
216
+ return bits_to_int(out)
217
+
218
+ def div8(self, a, b):
219
+ """8-bit restoring division through the per-stage bit-cascade GE
220
+ comparators and the sub8 subtractor. Returns (quotient, remainder);
221
+ divide-by-zero yields (0xFF, a)."""
222
+ if b == 0:
223
+ return 0xFF, a
224
+ a_bits = int_to_bits(a, 8)
225
+ div_bits = int_to_bits(b, 8)
226
+ quotient = 0
227
+ remainder = 0
228
+ for stage in range(8):
229
+ remainder = ((remainder << 1) | a_bits[stage]) & 0xFF
230
+ prefix = f"alu.alu8bit.div.stage{stage}"
231
+ finals = {
232
+ "gt": f"{prefix}.cmp_bc.gt",
233
+ "lt": f"{prefix}.cmp_bc.lt",
234
+ "eq": f"{prefix}.cmp_bc.eq",
235
+ "ge": f"{prefix}.cmp",
236
+ "le": f"{prefix}.cmp_bc.le",
237
+ }
238
+ ge = self._cmp_bit_cascade(f"{prefix}.cmp_bc", finals,
239
+ int_to_bits(remainder, 8), div_bits, "ge")
240
+ if ge:
241
+ remainder, _ = self.sub8(remainder, b)
242
+ quotient = (quotient << 1) | 1
243
+ else:
244
+ quotient = quotient << 1
245
+ return quotient & 0xFF, remainder & 0xFF
246
+
247
+ # ----- N-bit primitives (for 16-bit and 32-bit variants) ----------------
248
+
249
+ def add_n(self, a: int, b: int, bits: int):
250
+ """Width-generic ripple-carry add via arithmetic.ripplecarry{N}bit."""
251
+ prefix = f"arithmetic.ripplecarry{bits}bit"
252
+ a_lsb = list(reversed(int_to_bits(a, bits)))
253
+ b_lsb = list(reversed(int_to_bits(b, bits)))
254
+ carry = 0
255
+ s_lsb = []
256
+ for i in range(bits):
257
+ s, carry = self._fa(f"{prefix}.fa{i}", a_lsb[i], b_lsb[i], carry)
258
+ s_lsb.append(s)
259
+ return bits_to_int(list(reversed(s_lsb))), carry
260
+
261
+ def sub_n(self, a: int, b: int, bits: int):
262
+ """N-bit two's-complement subtract via arithmetic.sub{N}bit (N >= 16).
263
+
264
+ Structure (per build.add_sub_nbits): N NOT gates + N standard full adders.
265
+ """
266
+ prefix = f"arithmetic.sub{bits}bit"
267
+ a_lsb = list(reversed(int_to_bits(a, bits)))
268
+ b_lsb = list(reversed(int_to_bits(b, bits)))
269
+ # NOT each B bit
270
+ notb = [self._g(f"{prefix}.not_b.bit{i}", [b_lsb[i]]) for i in range(bits)]
271
+ carry = 1 # carry-in = 1 for two's-complement
272
+ d_lsb = []
273
+ for i in range(bits):
274
+ s, carry = self._fa(f"{prefix}.fa{i}", a_lsb[i], notb[i], carry)
275
+ d_lsb.append(s)
276
+ return bits_to_int(list(reversed(d_lsb))), carry
277
+
278
+ def cmp_n(self, a: int, b: int, kind: str, bits: int):
279
+ """N-bit unsigned comparison via the bit-cascade comparator family
280
+ (arithmetic.cmp{N}bit.* per build.add_bit_cascade_compare)."""
281
+ short = self._CMP_KIND[kind]
282
+ finals = {
283
+ "gt": f"arithmetic.greaterthan{bits}bit",
284
+ "lt": f"arithmetic.lessthan{bits}bit",
285
+ "eq": f"arithmetic.equality{bits}bit",
286
+ "ge": f"arithmetic.greaterorequal{bits}bit",
287
+ "le": f"arithmetic.lessorequal{bits}bit",
288
+ }
289
+ return self._cmp_bit_cascade(f"arithmetic.cmp{bits}bit", finals,
290
+ int_to_bits(a, bits), int_to_bits(b, bits), short)
291
+
292
+ def mul_n(self, a: int, b: int, bits: int):
293
+ """N-bit shift-add multiply (low N bits only)."""
294
+ ab = int_to_bits(a, bits)
295
+ bb = int_to_bits(b, bits)
296
+ mask = (1 << bits) - 1
297
+ result = 0
298
+ for j in range(bits):
299
+ if bb[j] == 0:
300
+ continue
301
+ row = 0
302
+ for i in range(bits):
303
+ pp = self._g(f"alu.alu{bits}bit.mul.pp.a{i}b{j}", [ab[i], bb[j]])
304
+ row |= (pp << (bits - 1 - i))
305
+ shift = (bits - 1) - j
306
+ result, _ = self.add_n(result & mask, (row << shift) & mask, bits)
307
+ return result & mask
308
+
309
+
310
+ class GenericThresholdCPU:
311
+ """Variant-aware CPU runtime. Sized from the variant's manifest."""
312
+
313
+ def __init__(self, tensors: Dict[str, torch.Tensor]):
314
+ self.T = tensors
315
+ m = get_manifest(tensors)
316
+ self.data_bits = m["data_bits"]
317
+ self.addr_bits = m["addr_bits"]
318
+ self.mem_bytes = m["memory_bytes"]
319
+ # 8-bit CPU primitives (ripplecarry8bit, sub8bit, alu.alu8bit.*, memory.*,
320
+ # control.*) are present in every variant regardless of manifest data_bits.
321
+ # Wider data widths simply add additional standalone ALU primitives.
322
+ if self.mem_bytes == 0:
323
+ raise NotImplementedError(
324
+ "Pure-ALU variants have no memory; cannot run CPU programs"
325
+ )
326
+ self.alu = GenericThresholdALU(tensors, 8)
327
+
328
+ def _addr_decode(self, addr):
329
+ bits = torch.tensor(int_to_bits(addr, self.addr_bits), dtype=torch.float32)
330
+ w = self.T["memory.addr_decode.weight"]
331
+ b = self.T["memory.addr_decode.bias"]
332
+ return heaviside((w * bits).sum(dim=1) + b)
333
+
334
+ def mem_read(self, mem, addr):
335
+ sel = self._addr_decode(addr)
336
+ mem_bits = torch.tensor(
337
+ [int_to_bits(byte, 8) for byte in mem], dtype=torch.float32
338
+ )
339
+ and_w = self.T["memory.read.and.weight"]
340
+ and_b = self.T["memory.read.and.bias"]
341
+ or_w = self.T["memory.read.or.weight"]
342
+ or_b = self.T["memory.read.or.bias"]
343
+ out = []
344
+ for bit in range(8):
345
+ inp = torch.stack([mem_bits[:, bit], sel], dim=1)
346
+ and_out = heaviside((inp * and_w[bit]).sum(dim=1) + and_b[bit])
347
+ out.append(int(heaviside((and_out * or_w[bit]).sum() + or_b[bit]).item()))
348
+ return bits_to_int(out)
349
+
350
+ def mem_write(self, mem, addr, value):
351
+ sel = self._addr_decode(addr)
352
+ data_bits = torch.tensor(int_to_bits(value, 8), dtype=torch.float32)
353
+ mem_bits = torch.tensor(
354
+ [int_to_bits(byte, 8) for byte in mem], dtype=torch.float32
355
+ )
356
+ sel_w = self.T["memory.write.sel.weight"]
357
+ sel_b = self.T["memory.write.sel.bias"]
358
+ nsel_w = self.T["memory.write.nsel.weight"].squeeze(1)
359
+ nsel_b = self.T["memory.write.nsel.bias"]
360
+ and_old_w = self.T["memory.write.and_old.weight"]
361
+ and_old_b = self.T["memory.write.and_old.bias"]
362
+ and_new_w = self.T["memory.write.and_new.weight"]
363
+ and_new_b = self.T["memory.write.and_new.bias"]
364
+ or_w = self.T["memory.write.or.weight"]
365
+ or_b = self.T["memory.write.or.bias"]
366
+ we = torch.ones_like(sel)
367
+ sel_inp = torch.stack([sel, we], dim=1)
368
+ write_sel = heaviside((sel_inp * sel_w).sum(dim=1) + sel_b)
369
+ nsel = heaviside(write_sel * nsel_w + nsel_b)
370
+ for bit in range(8):
371
+ old = mem_bits[:, bit]
372
+ data_bit = data_bits[bit].expand(self.mem_bytes)
373
+ inp_old = torch.stack([old, nsel], dim=1)
374
+ inp_new = torch.stack([data_bit, write_sel], dim=1)
375
+ and_old = heaviside((inp_old * and_old_w[:, bit]).sum(dim=1) + and_old_b[:, bit])
376
+ and_new = heaviside((inp_new * and_new_w[:, bit]).sum(dim=1) + and_new_b[:, bit])
377
+ or_inp = torch.stack([and_old, and_new], dim=1)
378
+ new_bit = heaviside((or_inp * or_w[:, bit]).sum(dim=1) + or_b[:, bit])
379
+ mem_bits[:, bit] = new_bit
380
+ return [bits_to_int([int(b) for b in mem_bits[i].tolist()]) for i in range(self.mem_bytes)]
381
+
382
+ # Conditions 0..7 map to (mux circuit, flag index into [Z, N, C, V]).
383
+ # Odd conditions are the negated forms; the mux select is the complemented
384
+ # flag (the circuits are plain per-bit muxes -- see control.* .inputs,
385
+ # which name the select $not_zero / $not_carry / ... for the odd forms).
386
+ _JCC = [("control.jz", 0), ("control.jnz", 0), ("control.jc", 2), ("control.jnc", 2),
387
+ ("control.jn", 1), ("control.jp", 1), ("control.jv", 3), ("control.jnv", 3)]
388
+
389
+ def _jcc_pc(self, circuit: str, pc: int, target: int, sel: int) -> int:
390
+ """Per-bit 2:1 mux over the PC: sel ? target : pc, addr_bits wide."""
391
+ pc_bits = int_to_bits(pc, self.addr_bits)
392
+ t_bits = int_to_bits(target, self.addr_bits)
393
+ out = []
394
+ for bit in range(self.addr_bits):
395
+ bp = f"{circuit}.bit{bit}"
396
+ not_sel = self.alu._g(f"{bp}.not_sel", [float(sel)])
397
+ and_a = self.alu._g(f"{bp}.and_a", [float(pc_bits[bit]), not_sel])
398
+ and_b = self.alu._g(f"{bp}.and_b", [float(t_bits[bit]), float(sel)])
399
+ out.append(self.alu._g(f"{bp}.or", [and_a, and_b]))
400
+ return bits_to_int(out)
401
+
402
+ def _sp_dec(self, sp: int) -> int:
403
+ """SP - 1 through the control.push.sp_dec borrow chain. Gate bit index
404
+ is MSB-first (bit addr_bits-1 is the LSB); the bit complement feeding
405
+ each borrow AND is fixed wiring, as in the gate fitness suite."""
406
+ bits = int_to_bits(sp, self.addr_bits)
407
+ out = [0] * self.addr_bits
408
+ borrow = 1
409
+ for bit in range(self.addr_bits - 1, -1, -1):
410
+ bp = f"control.push.sp_dec.bit{bit}"
411
+ h_or = self.alu._g(f"{bp}.xor.layer1.or", [float(bits[bit]), float(borrow)])
412
+ h_nand = self.alu._g(f"{bp}.xor.layer1.nand", [float(bits[bit]), float(borrow)])
413
+ out[bit] = self.alu._g(f"{bp}.xor.layer2", [h_or, h_nand])
414
+ borrow = self.alu._g(f"{bp}.borrow", [float(1 - bits[bit]), float(borrow)])
415
+ return bits_to_int(out)
416
+
417
+ def step(self, state):
418
+ if state["halted"]:
419
+ return state
420
+ s = dict(state)
421
+ s["mem"] = state["mem"][:]
422
+ s["regs"] = state["regs"][:]
423
+ s["flags"] = state["flags"][:]
424
+ addr_mask = (1 << self.addr_bits) - 1
425
+ pc = s["pc"]
426
+ hi = self.mem_read(s["mem"], pc & addr_mask)
427
+ lo = self.mem_read(s["mem"], (pc + 1) & addr_mask)
428
+ ir = ((hi & 0xFF) << 8) | (lo & 0xFF)
429
+ opcode = (ir >> 12) & 0xF
430
+ rd = (ir >> 10) & 0x3
431
+ rs = (ir >> 8) & 0x3
432
+ imm = ir & 0xFF
433
+ next_pc = (pc + 2) & addr_mask
434
+ addr_full = None
435
+ if opcode in (0xA, 0xB, 0xC, 0xD, 0xE):
436
+ ah = self.mem_read(s["mem"], next_pc)
437
+ al = self.mem_read(s["mem"], (next_pc + 1) & addr_mask)
438
+ addr_full = ((ah & 0xFF) << 8) | (al & 0xFF)
439
+ next_pc = (next_pc + 2) & addr_mask
440
+ addr = (addr_full & addr_mask) if addr_full is not None else None
441
+ a = s["regs"][rd]
442
+ b = s["regs"][rs]
443
+ result = a
444
+ carry = 0
445
+ overflow = 0
446
+ write_result = True
447
+ if opcode == 0x0:
448
+ result, carry = self.alu.add8(a, b)
449
+ overflow = 1 if (((a ^ result) & (b ^ result)) & 0x80) else 0
450
+ elif opcode == 0x1:
451
+ result, carry = self.alu.sub8(a, b)
452
+ overflow = 1 if (((a ^ b) & (a ^ result)) & 0x80) else 0
453
+ elif opcode == 0x2: # AND
454
+ result = self.alu.and8(a, b)
455
+ elif opcode == 0x3: # OR
456
+ result = self.alu.or8(a, b)
457
+ elif opcode == 0x4: # XOR
458
+ result = self.alu.xor8(a, b)
459
+ elif opcode == 0x5: # SHL by 1 (8-bit)
460
+ result = self.alu.shl8(a)
461
+ elif opcode == 0x6: # SHR by 1
462
+ result = self.alu.shr8(a)
463
+ elif opcode == 0x7:
464
+ result = self.alu.mul8(a, b)
465
+ elif opcode == 0x8: # DIV (sets R[d] = R[d] / R[s]; 0xFF on divide by zero)
466
+ result, _ = self.alu.div8(a, b)
467
+ elif opcode == 0x9:
468
+ r2, carry = self.alu.sub8(a, b)
469
+ z = 1 if r2 == 0 else 0
470
+ n = 1 if (r2 & 0x80) else 0
471
+ v = 1 if (((a ^ b) & (a ^ r2)) & 0x80) else 0
472
+ s["flags"] = [z, n, carry, v]
473
+ write_result = False
474
+ elif opcode == 0xA:
475
+ result = self.mem_read(s["mem"], addr)
476
+ elif opcode == 0xB:
477
+ s["mem"] = self.mem_write(s["mem"], addr, b & 0xFF)
478
+ write_result = False
479
+ elif opcode == 0xC:
480
+ s["pc"] = addr
481
+ return s
482
+ elif opcode == 0xD:
483
+ cond = imm & 0x7
484
+ circuit, flag_idx = self._JCC[cond]
485
+ flag = s["flags"][flag_idx]
486
+ sel = flag if cond % 2 == 0 else 1 - flag
487
+ s["pc"] = self._jcc_pc(circuit, next_pc, addr, sel)
488
+ return s
489
+ elif opcode == 0xE: # CALL: push return address (next_pc), set PC = addr
490
+ ret_addr = next_pc & 0xFFFF
491
+ sp = s.get("sp", addr_mask)
492
+ sp = self._sp_dec(sp)
493
+ s["mem"] = self.mem_write(s["mem"], sp, (ret_addr >> 8) & 0xFF)
494
+ sp = self._sp_dec(sp)
495
+ s["mem"] = self.mem_write(s["mem"], sp, ret_addr & 0xFF)
496
+ s["sp"] = sp
497
+ s["pc"] = addr
498
+ return s
499
+ elif opcode == 0xF:
500
+ s["halted"] = True
501
+ return s
502
+
503
+ if write_result and opcode != 0x9:
504
+ s["regs"][rd] = result & 0xFF
505
+ # Flag policy: ADD, SUB, MUL (and CMP, handled in its branch) write
506
+ # Z/N/C/V; MUL clears C and V. All other opcodes leave FLAGS unchanged.
507
+ if opcode in (0x0, 0x1, 0x7):
508
+ z = 1 if (result & 0xFF) == 0 else 0
509
+ n = 1 if (result & 0x80) else 0
510
+ s["flags"] = [z, n, carry, overflow]
511
+ s["pc"] = next_pc
512
+ return s
513
+
514
+ def run(self, state, max_cycles=200):
515
+ s = state
516
+ cycles = 0
517
+ while not s["halted"] and cycles < max_cycles:
518
+ s = self.step(s)
519
+ cycles += 1
520
+ return s, cycles
521
+
522
+
523
+ def _encode_instr(opcode, rd, rs, imm):
524
+ return ((opcode & 0xF) << 12) | ((rd & 0x3) << 10) | ((rs & 0x3) << 8) | (imm & 0xFF)
525
+
526
+
527
+ def _w16(mem, addr, value):
528
+ mem[addr] = (value >> 8) & 0xFF
529
+ mem[addr + 1] = value & 0xFF
530
+
531
+
532
+ PROGRAM_MIN_BYTES = 0x84 # code 0x00..0x1F + data 0x80..0x83
533
+
534
+
535
+ def builtin_program(addr_bits: int) -> Tuple[List[int], int]:
536
+ """Sum 5+4+3+2+1 via a loop. Returns (mem, expected_result_at_0x83).
537
+
538
+ Compact layout: code at 0x00..0x1F (32 bytes), data at 0x80..0x83 (4 bytes).
539
+ Total footprint 132 bytes -- fits within scratchpad (256 B) and larger.
540
+ Requires addr_bits >= 8.
541
+ """
542
+ if (1 << addr_bits) < PROGRAM_MIN_BYTES:
543
+ raise ValueError(f"addr_bits={addr_bits} too small for builtin program")
544
+ mem = [0] * (1 << addr_bits)
545
+ mem[0x80] = 5 # initial counter
546
+ mem[0x81] = 1 # decrement
547
+ mem[0x82] = 0 # zero (for compare and accumulator init)
548
+ # mem[0x83] is the output
549
+ _w16(mem, 0x0000, _encode_instr(0xA, 1, 0, 0)); _w16(mem, 0x0002, 0x0080)
550
+ _w16(mem, 0x0004, _encode_instr(0xA, 2, 0, 0)); _w16(mem, 0x0006, 0x0081)
551
+ _w16(mem, 0x0008, _encode_instr(0xA, 3, 0, 0)); _w16(mem, 0x000A, 0x0082)
552
+ _w16(mem, 0x000C, _encode_instr(0xA, 0, 0, 0)); _w16(mem, 0x000E, 0x0082)
553
+ _w16(mem, 0x0010, _encode_instr(0x0, 0, 1, 0))
554
+ _w16(mem, 0x0012, _encode_instr(0x1, 1, 2, 0))
555
+ _w16(mem, 0x0014, _encode_instr(0x9, 1, 3, 0))
556
+ _w16(mem, 0x0016, _encode_instr(0xD, 0, 0, 0x01)); _w16(mem, 0x0018, 0x0010)
557
+ _w16(mem, 0x001A, _encode_instr(0xB, 0, 0, 0)); _w16(mem, 0x001C, 0x0083)
558
+ _w16(mem, 0x001E, _encode_instr(0xF, 0, 0, 0))
559
+ return mem, 15
560
+
561
+
562
+ # ---------------------------------------------------------------------------
563
+ # Eval driver
564
+ # ---------------------------------------------------------------------------
565
+
566
+ def _file_fingerprint(path: Path) -> str:
567
+ """Stable cache key for a safetensors file: sha256 of its content.
568
+
569
+ Hashes are content-addressed so renaming a file doesn't blow the cache,
570
+ but mtime-only would re-key on every clone of the repo. The sha256 of a
571
+ 30 MB safetensors finishes in tens of milliseconds — small compared to
572
+ a 5,900-test fitness run.
573
+ """
574
+ import hashlib
575
+ h = hashlib.sha256()
576
+ with open(path, "rb") as f:
577
+ for chunk in iter(lambda: f.read(1 << 20), b""):
578
+ h.update(chunk)
579
+ return h.hexdigest()
580
+
581
+
582
+ def _cache_key(path: Path, opts: Dict[str, Any]) -> str:
583
+ """Cache key combining file content with the relevant evaluation options."""
584
+ fp = _file_fingerprint(path)
585
+ opt_str = json.dumps(opts, sort_keys=True)
586
+ import hashlib
587
+ suffix = hashlib.sha256(opt_str.encode("utf-8")).hexdigest()[:8]
588
+ return f"{fp}_{suffix}"
589
+
590
+
591
+ def _load_cache(cache_dir: Path, key: str) -> Dict[str, Any] | None:
592
+ p = cache_dir / f"{key}.json"
593
+ if not p.exists():
594
+ return None
595
+ try:
596
+ return json.loads(p.read_text(encoding="utf-8"))
597
+ except (json.JSONDecodeError, OSError):
598
+ return None
599
+
600
+
601
+ def _save_cache(cache_dir: Path, key: str, payload: Dict[str, Any]) -> None:
602
+ cache_dir.mkdir(parents=True, exist_ok=True)
603
+ p = cache_dir / f"{key}.json"
604
+ try:
605
+ p.write_text(json.dumps(payload, indent=2, default=str), encoding="utf-8")
606
+ except OSError:
607
+ pass
608
+
609
+
610
+ def list_safetensors(path: Path) -> List[Path]:
611
+ if path.is_file():
612
+ return [path]
613
+ if path.is_dir():
614
+ return sorted(p for p in path.glob("*.safetensors") if p.is_file())
615
+ return []
616
+
617
+
618
+ def evaluate_one(path: Path, device: str, pop_size: int, debug: bool, run_cpu_program: bool) -> Dict:
619
+ out: Dict = {"path": str(path), "filename": path.name}
620
+ # The standalone machines (neural_subleq8, neural_rv32) carry their own
621
+ # circuit inventory and are verified by machines.py, not the gate-fitness
622
+ # suite; skip them cleanly rather than error on missing standard gates.
623
+ with safe_open(str(path), framework="pt") as f:
624
+ meta = f.metadata() or {}
625
+ if meta.get("machine"):
626
+ out.update(status="SKIP", machine=meta["machine"],
627
+ note="standalone machine — verify with machines.py")
628
+ return out
629
+ try:
630
+ tensors = load_model(str(path))
631
+ except Exception as e:
632
+ out.update(error=f"load failed: {e}", status="ERROR")
633
+ return out
634
+
635
+ manifest = get_manifest(tensors)
636
+ out.update(
637
+ size_mb=path.stat().st_size / (1024 * 1024),
638
+ tensors=len(tensors),
639
+ params=sum(t.numel() for t in tensors.values()),
640
+ manifest=manifest,
641
+ )
642
+
643
+ # Move to device
644
+ tensors = {k: v.to(device) for k, v in tensors.items()}
645
+
646
+ try:
647
+ evaluator = BatchedFitnessEvaluator(device=device, model_path=str(path), tensors=tensors)
648
+ population = create_population(tensors, pop_size=pop_size, device=device)
649
+ t0 = time.perf_counter()
650
+ fitness = evaluator.evaluate(population, debug=debug)
651
+ elapsed = time.perf_counter() - t0
652
+ f0 = float(fitness[0].item()) if pop_size == 1 else float(fitness.mean().item())
653
+ out.update(
654
+ fitness=f0,
655
+ total_tests=evaluator.total_tests,
656
+ elapsed_s=elapsed,
657
+ categories={k: (float(v[0]), int(v[1])) for k, v in evaluator.category_scores.items()},
658
+ status="PASS" if f0 >= 0.9999 else "FAIL",
659
+ )
660
+ except Exception as e:
661
+ out.update(error=f"eval failed: {type(e).__name__}: {e}", status="ERROR")
662
+ return out
663
+
664
+ # Optional: CPU program test (8-bit CPU primitives are in every variant)
665
+ if run_cpu_program:
666
+ if manifest["memory_bytes"] >= PROGRAM_MIN_BYTES:
667
+ try:
668
+ cpu_tensors = {k: v.cpu() for k, v in tensors.items()}
669
+ cpu = GenericThresholdCPU(cpu_tensors)
670
+ mem, expected = builtin_program(manifest["addr_bits"])
671
+ state = {"pc": 0, "regs": [0] * 4, "flags": [0] * 4, "mem": mem, "halted": False}
672
+ t0 = time.perf_counter()
673
+ final, cycles = cpu.run(state, max_cycles=200)
674
+ cpu_elapsed = time.perf_counter() - t0
675
+ got = final["mem"][0x83]
676
+ out["cpu_program"] = {
677
+ "ok": got == expected,
678
+ "got": got,
679
+ "expected": expected,
680
+ "cycles": cycles,
681
+ "elapsed_s": cpu_elapsed,
682
+ }
683
+ if got != expected:
684
+ out["status"] = "FAIL"
685
+ except Exception as e:
686
+ out["cpu_program"] = {"error": str(e)}
687
+ else:
688
+ out["cpu_program"] = {"skipped": f"mem={manifest['memory_bytes']}B < {PROGRAM_MIN_BYTES}"}
689
+
690
+ # Wider-ALU chain test for 16/32-bit variants
691
+ bits = manifest["data_bits"]
692
+ if bits in (16, 32):
693
+ try:
694
+ alu_tensors = {k: v.cpu() for k, v in tensors.items()}
695
+ alu = GenericThresholdALU(alu_tensors, bits)
696
+ t0 = time.perf_counter()
697
+ if bits == 16:
698
+ x, y = 1234, 5678
699
+ z, _ = alu.add_n(x, y, 16); assert z == (x + y) & 0xFFFF
700
+ w, _ = alu.sub_n(z, x, 16); assert w == (z - x) & 0xFFFF, (w, z - x)
701
+ gt = alu.cmp_n(z, x, "greaterthan", 16); assert gt == 1
702
+ lt = alu.cmp_n(x, z, "lessthan", 16); assert lt == 1
703
+ eq = alu.cmp_n(w, y, "eq", 16); assert eq == 1
704
+ p = alu.mul_n(123, 5, 16); assert p == (123 * 5) & 0xFFFF
705
+ else: # 32
706
+ x, y = 1_000_000, 999_000
707
+ z, _ = alu.sub_n(x, y, 32); assert z == 1_000
708
+ s, _ = alu.add_n(z, x, 32); assert s == 1_001_000
709
+ p = alu.mul_n(z, 100, 32); assert p == 100_000
710
+ gt = alu.cmp_n(x, y, "greaterthan", 32); assert gt == 1
711
+ lt = alu.cmp_n(y, x, "lessthan", 32); assert lt == 1
712
+ eq = alu.cmp_n(p, 100_000, "equality", 32); assert eq == 1
713
+ chain_dt = time.perf_counter() - t0
714
+ out[f"alu_chain_{bits}"] = {"ok": True, "elapsed_s": chain_dt}
715
+ except AssertionError as e:
716
+ out[f"alu_chain_{bits}"] = {"ok": False, "error": f"chain mismatch: {e}"}
717
+ out["status"] = "FAIL"
718
+ except Exception as e:
719
+ out[f"alu_chain_{bits}"] = {"ok": False, "error": f"{type(e).__name__}: {e}"}
720
+ out["status"] = "FAIL"
721
+
722
+ return out
723
+
724
+
725
+ def print_row(r: Dict, show_cpu: bool) -> None:
726
+ if r.get("status") == "SKIP":
727
+ print(f" {r['filename']:<48} SKIP ({r.get('machine', 'machine')} — verify with machines.py)")
728
+ return
729
+ if "error" in r:
730
+ print(f" {r['filename']:<48} ERROR: {r['error'][:80]}")
731
+ return
732
+ m = r["manifest"]
733
+ fit = f"{r['fitness']:.4f}" if r.get("fitness") is not None else "n/a"
734
+ cpu_col = ""
735
+ if show_cpu and "cpu_program" in r:
736
+ cp = r["cpu_program"]
737
+ if cp.get("ok"):
738
+ cpu_col = f" CPU OK ({cp['cycles']}cyc/{cp['elapsed_s']:.1f}s)"
739
+ elif "skipped" in cp:
740
+ cpu_col = f" CPU SKIP"
741
+ elif "error" in cp:
742
+ cpu_col = f" CPU ERR"
743
+ else:
744
+ cpu_col = f" CPU FAIL ({cp.get('got')}!={cp.get('expected')})"
745
+ chain_col = ""
746
+ if show_cpu:
747
+ for bits in (16, 32):
748
+ key = f"alu_chain_{bits}"
749
+ if key in r:
750
+ ch = r[key]
751
+ if ch.get("ok"):
752
+ chain_col = f" ALU{bits} OK ({ch['elapsed_s']:.2f}s)"
753
+ else:
754
+ chain_col = f" ALU{bits} FAIL"
755
+ print(
756
+ f" {r['filename']:<48} d={m['data_bits']:>2}b a={m['addr_bits']:>2}b "
757
+ f"mem={m['memory_bytes']:>6}B size={r['size_mb']:>6.1f}MB "
758
+ f"params={r['params']:>10,} fit={fit:>6} tests={r['total_tests']:>5} "
759
+ f"{r['status']:>5}{cpu_col}{chain_col}"
760
+ )
761
+
762
+
763
+ def main() -> int:
764
+ parser = argparse.ArgumentParser(description="Variant-agnostic eval harness")
765
+ parser.add_argument("path", help="Path to .safetensors file or directory of files")
766
+ parser.add_argument("--device", default="cpu", help="cpu (default) or cuda")
767
+ parser.add_argument("--pop_size", type=int, default=1)
768
+ parser.add_argument("--debug", action="store_true", help="Per-circuit detail per file")
769
+ parser.add_argument("--cpu-program", action="store_true",
770
+ help="Also run a small assembled program through the threshold CPU "
771
+ "(any variant with >= 132 B memory), plus a chained 16- or "
772
+ "32-bit ALU sequence on wider variants")
773
+ parser.add_argument("--json", action="store_true", help="Emit JSON results to stdout instead of a table")
774
+ parser.add_argument("--cache-dir", default=".eval_cache",
775
+ help="Directory for hash-keyed result cache "
776
+ "(default: ./.eval_cache). Set to '' to disable.")
777
+ parser.add_argument("--no-cache", action="store_true",
778
+ help="Disable the result cache for this run.")
779
+ args = parser.parse_args()
780
+
781
+ files = list_safetensors(Path(args.path))
782
+ if not files:
783
+ print(f"No .safetensors files found under {args.path}", file=sys.stderr)
784
+ return 2
785
+
786
+ print(f"Evaluating {len(files)} file(s) on {args.device}\n")
787
+ cache_enabled = bool(args.cache_dir) and not args.no_cache
788
+ cache_dir = Path(args.cache_dir) if cache_enabled else None
789
+ cache_opts = {
790
+ "device": args.device,
791
+ "pop_size": args.pop_size,
792
+ "cpu_program": bool(args.cpu_program),
793
+ }
794
+ cache_hits = 0
795
+ results = []
796
+ fail_count = 0
797
+ for f in files:
798
+ print(f"=== {f.name}")
799
+ cached = None
800
+ key = None
801
+ if cache_enabled:
802
+ try:
803
+ key = _cache_key(f, cache_opts)
804
+ cached = _load_cache(cache_dir, key)
805
+ except OSError:
806
+ cached = None
807
+ if cached is not None:
808
+ r = cached
809
+ cache_hits += 1
810
+ print(f" (cache hit)")
811
+ else:
812
+ r = evaluate_one(f, device=args.device, pop_size=args.pop_size,
813
+ debug=args.debug, run_cpu_program=args.cpu_program)
814
+ # Don't cache ERROR results: a transient failure (OOM, interrupt)
815
+ # would otherwise be replayed on every rerun of the same file.
816
+ if cache_enabled and key is not None and r.get("status") != "ERROR":
817
+ _save_cache(cache_dir, key, r)
818
+ results.append(r)
819
+ print_row(r, show_cpu=args.cpu_program)
820
+ if r.get("status") not in ("PASS", "SKIP"):
821
+ fail_count += 1
822
+
823
+ if args.json:
824
+ # Make it serialisable
825
+ for r in results:
826
+ r["manifest"] = {k: (int(v) if isinstance(v, float) and v.is_integer() else v)
827
+ for k, v in r.get("manifest", {}).items()}
828
+ print(json.dumps(results, indent=2, default=str))
829
+ return fail_count
830
+
831
+ # Summary
832
+ print()
833
+ print("=" * 100)
834
+ print(" SUMMARY")
835
+ print("=" * 100)
836
+ for r in results:
837
+ print_row(r, show_cpu=args.cpu_program)
838
+
839
+ print()
840
+ if fail_count == 0:
841
+ print(f"ALL {len(files)} variants PASS")
842
+ else:
843
+ print(f"{fail_count}/{len(files)} variants FAIL")
844
+ if cache_enabled:
845
+ print(f"(cache: {cache_hits}/{len(files)} hits, dir={cache_dir})")
846
+ return fail_count
847
+
848
+
849
+ if __name__ == "__main__":
850
+ sys.exit(main())
llm_integration/baseline.py ADDED
@@ -0,0 +1,221 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Baseline evaluation: Vanilla SmolLM2-360M on arithmetic
3
+ """
4
+
5
+ import torch
6
+ import random
7
+ import re
8
+ from transformers import AutoModelForCausalLM, AutoTokenizer
9
+
10
+ DEVICE = "cuda"
11
+ MODEL_ID = "HuggingFaceTB/SmolLM2-360M-Instruct"
12
+
13
+ SYSTEM_PROMPT = """You are a calculator. Output only the numeric answer. No words, no explanation, just digits. Examples:
14
+ User: 5 + 3
15
+ Assistant: 8
16
+ User: 12 * 7
17
+ Assistant: 84
18
+ User: 100 > 50
19
+ Assistant: 1
20
+ User: 25 < 10
21
+ Assistant: 0"""
22
+
23
+
24
+ def load_model():
25
+ print(f"Loading {MODEL_ID}...")
26
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
27
+ tokenizer.padding_side = "left"
28
+ if tokenizer.pad_token is None:
29
+ tokenizer.pad_token = tokenizer.eos_token
30
+ model = AutoModelForCausalLM.from_pretrained(
31
+ MODEL_ID,
32
+ torch_dtype=torch.float16,
33
+ device_map=DEVICE
34
+ )
35
+ model.eval()
36
+ print(f" Loaded. Parameters: {sum(p.numel() for p in model.parameters()):,}")
37
+ return model, tokenizer
38
+
39
+
40
+ def format_prompt(tokenizer, op_str):
41
+ messages = [
42
+ {"role": "system", "content": SYSTEM_PROMPT},
43
+ {"role": "user", "content": op_str}
44
+ ]
45
+ return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
46
+
47
+
48
+ def generate_batch(model, tokenizer, prompts, max_new_tokens=16):
49
+ inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(DEVICE)
50
+ with torch.no_grad():
51
+ outputs = model.generate(
52
+ **inputs,
53
+ max_new_tokens=max_new_tokens,
54
+ do_sample=False,
55
+ pad_token_id=tokenizer.eos_token_id
56
+ )
57
+ responses = []
58
+ for i, output in enumerate(outputs):
59
+ response = tokenizer.decode(output[inputs.input_ids.shape[1]:], skip_special_tokens=True)
60
+ responses.append(response.strip())
61
+ return responses
62
+
63
+
64
+ def extract_answer(text):
65
+ """Generous extraction - find any number in output"""
66
+ text = text.strip().lower()
67
+ if not text:
68
+ return None
69
+
70
+ # Handle Yes/No for comparisons
71
+ if text in ['yes', 'true', '1']:
72
+ return 1
73
+ if text in ['no', 'false', '0']:
74
+ return 0
75
+ if text.startswith('yes'):
76
+ return 1
77
+ if text.startswith('no'):
78
+ return 0
79
+
80
+ # Find all numbers, take the LAST one (most likely the answer)
81
+ numbers = re.findall(r'-?\d+', text)
82
+ if numbers:
83
+ return int(numbers[-1])
84
+ return None
85
+
86
+
87
+ def ground_truth(a, b, op):
88
+ """Compute expected result (8-bit where applicable)"""
89
+ if op == 'add':
90
+ return (a + b) & 0xFF
91
+ elif op == 'sub':
92
+ return (a - b) & 0xFF
93
+ elif op == 'mul':
94
+ return (a * b) & 0xFF
95
+ elif op == 'div':
96
+ return a // b if b != 0 else 0
97
+ elif op == 'and':
98
+ return a & b
99
+ elif op == 'or':
100
+ return a | b
101
+ elif op == 'xor':
102
+ return a ^ b
103
+ elif op == 'gt':
104
+ return 1 if a > b else 0
105
+ elif op == 'lt':
106
+ return 1 if a < b else 0
107
+ elif op == 'eq':
108
+ return 1 if a == b else 0
109
+ elif op == 'ge':
110
+ return 1 if a >= b else 0
111
+ elif op == 'le':
112
+ return 1 if a <= b else 0
113
+ else:
114
+ raise ValueError(f"Unknown op: {op}")
115
+
116
+
117
+ def op_to_str(a, b, op):
118
+ """Convert operation to natural string"""
119
+ symbols = {
120
+ 'add': '+', 'sub': '-', 'mul': '*', 'div': '/',
121
+ 'and': '&', 'or': '|', 'xor': '^',
122
+ 'gt': '>', 'lt': '<', 'eq': '==', 'ge': '>=', 'le': '<='
123
+ }
124
+ return f"{a} {symbols[op]} {b}"
125
+
126
+
127
+ def evaluate(model, tokenizer, n_samples=1000, batch_size=32, ops=None):
128
+ if ops is None:
129
+ ops = ['add', 'sub', 'mul', 'gt', 'lt', 'eq']
130
+
131
+ results = {op: {'correct': 0, 'total': 0} for op in ops}
132
+ all_correct = 0
133
+ all_total = 0
134
+
135
+ samples = []
136
+ for _ in range(n_samples):
137
+ a = random.randint(0, 255)
138
+ b = random.randint(0, 255)
139
+ if 'div' in ops and random.random() < 0.1:
140
+ op = 'div'
141
+ b = random.randint(1, 255) # avoid div by zero
142
+ else:
143
+ op = random.choice([o for o in ops if o != 'div'])
144
+ samples.append((a, b, op))
145
+
146
+ print(f"\nEvaluating {n_samples} samples (batch_size={batch_size})...")
147
+
148
+ for batch_start in range(0, n_samples, batch_size):
149
+ batch = samples[batch_start:batch_start + batch_size]
150
+ prompts = [format_prompt(tokenizer, op_to_str(a, b, op)) for a, b, op in batch]
151
+ responses = generate_batch(model, tokenizer, prompts)
152
+
153
+ for (a, b, op), response in zip(batch, responses):
154
+ expected = ground_truth(a, b, op)
155
+ extracted = extract_answer(response)
156
+
157
+ correct = (extracted == expected)
158
+ results[op]['total'] += 1
159
+ all_total += 1
160
+ if correct:
161
+ results[op]['correct'] += 1
162
+ all_correct += 1
163
+
164
+ if (batch_start + batch_size) % 200 == 0 or batch_start + batch_size >= n_samples:
165
+ pct = 100 * all_correct / all_total
166
+ print(f" Progress: {min(batch_start + batch_size, n_samples)}/{n_samples} | Accuracy: {pct:.2f}%")
167
+
168
+ return results, all_correct, all_total
169
+
170
+
171
+ def main():
172
+ random.seed(42)
173
+ torch.manual_seed(42)
174
+
175
+ model, tokenizer = load_model()
176
+
177
+ # Quick sanity check
178
+ print("\nSanity check (5 examples):")
179
+ test_cases = [
180
+ ("5 + 3", 8),
181
+ ("100 - 37", 63),
182
+ ("12 * 11", 132),
183
+ ("50 > 30", 1),
184
+ ("25 < 10", 0),
185
+ ]
186
+ prompts = [format_prompt(tokenizer, q) for q, _ in test_cases]
187
+ responses = generate_batch(model, tokenizer, prompts)
188
+ for (q, expected), response in zip(test_cases, responses):
189
+ extracted = extract_answer(response)
190
+ status = "OK" if extracted == expected else "FAIL"
191
+ print(f" {q} = {expected} | Model: '{response}' -> {extracted} [{status}]")
192
+
193
+ # Full evaluation
194
+ print("\n" + "=" * 60)
195
+ print(" BASELINE EVALUATION")
196
+ print("=" * 60)
197
+
198
+ ops = ['add', 'sub', 'mul', 'gt', 'lt', 'eq']
199
+ results, correct, total = evaluate(model, tokenizer, n_samples=2000, batch_size=64, ops=ops)
200
+
201
+ print("\n" + "=" * 60)
202
+ print(" RESULTS BY OPERATION")
203
+ print("=" * 60)
204
+ for op in ops:
205
+ r = results[op]
206
+ pct = 100 * r['correct'] / r['total'] if r['total'] > 0 else 0
207
+ print(f" {op:6}: {r['correct']:4}/{r['total']:4} ({pct:6.2f}%)")
208
+
209
+ print("\n" + "=" * 60)
210
+ print(" OVERALL")
211
+ print("=" * 60)
212
+ fitness = correct / total
213
+ print(f" Correct: {correct}/{total}")
214
+ print(f" Fitness: {fitness:.4f} ({100*fitness:.2f}%)")
215
+ print("=" * 60)
216
+
217
+ return fitness
218
+
219
+
220
+ if __name__ == "__main__":
221
+ main()
llm_integration/circuits.py ADDED
@@ -0,0 +1,320 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Frozen threshold circuit wrapper for LLM integration.
3
+ Loads safetensors and provides differentiable-compatible execution.
4
+ """
5
+
6
+ import torch
7
+ import torch.nn as nn
8
+ from safetensors import safe_open
9
+ from typing import Dict, Tuple
10
+
11
+ MODEL_PATH = "D:/8bit-threshold-computer/neural_computer.safetensors"
12
+
13
+
14
+ def heaviside(x: torch.Tensor) -> torch.Tensor:
15
+ """Standard Heaviside step function."""
16
+ return (x >= 0).float()
17
+
18
+
19
+ class HeavisideSTE(torch.autograd.Function):
20
+ """Heaviside with straight-through estimator for gradients."""
21
+ @staticmethod
22
+ def forward(ctx, x):
23
+ return (x >= 0).float()
24
+
25
+ @staticmethod
26
+ def backward(ctx, grad_output):
27
+ return grad_output
28
+
29
+
30
+ def heaviside_ste(x: torch.Tensor) -> torch.Tensor:
31
+ """Heaviside with STE gradient."""
32
+ return HeavisideSTE.apply(x)
33
+
34
+
35
+ class FrozenThresholdCircuits(nn.Module):
36
+ """
37
+ Wrapper for frozen threshold logic circuits.
38
+ All weights are frozen - no gradients flow through circuit internals.
39
+ Gradients flow through inputs/outputs via STE.
40
+ """
41
+
42
+ def __init__(self, model_path: str = MODEL_PATH, device: str = 'cuda'):
43
+ super().__init__()
44
+ self.device = device
45
+ self.weights = {}
46
+ self._load_weights(model_path)
47
+
48
+ def _load_weights(self, path: str):
49
+ """Load weights from safetensors file."""
50
+ with safe_open(path, framework='pt') as f:
51
+ for name in f.keys():
52
+ tensor = f.get_tensor(name).to(self.device).float()
53
+ self.weights[name] = tensor
54
+
55
+ def _gate(self, inputs: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
56
+ """Execute single threshold gate with STE."""
57
+ weight = weight.view(-1)
58
+ bias = bias.view(-1)
59
+ pre_activation = (inputs * weight).sum(dim=-1) + bias
60
+ return heaviside_ste(pre_activation)
61
+
62
+ def _xor(self, a: torch.Tensor, b: torch.Tensor, prefix: str) -> torch.Tensor:
63
+ """XOR via OR-NAND-AND pattern (2 layers)."""
64
+ inputs = torch.stack([a, b], dim=-1)
65
+
66
+ w_or = self.weights[f'{prefix}.layer1.or.weight']
67
+ b_or = self.weights[f'{prefix}.layer1.or.bias']
68
+ w_nand = self.weights[f'{prefix}.layer1.nand.weight']
69
+ b_nand = self.weights[f'{prefix}.layer1.nand.bias']
70
+
71
+ h_or = self._gate(inputs, w_or, b_or)
72
+ h_nand = self._gate(inputs, w_nand, b_nand)
73
+
74
+ hidden = torch.stack([h_or, h_nand], dim=-1)
75
+ w2 = self.weights[f'{prefix}.layer2.weight']
76
+ b2 = self.weights[f'{prefix}.layer2.bias']
77
+
78
+ return self._gate(hidden, w2, b2)
79
+
80
+ def _full_adder(self, a: torch.Tensor, b: torch.Tensor, cin: torch.Tensor,
81
+ prefix: str) -> Tuple[torch.Tensor, torch.Tensor]:
82
+ """Full adder: sum and carry out."""
83
+ ha1_sum = self._xor(a, b, f'{prefix}.ha1.sum')
84
+
85
+ inp_carry1 = torch.stack([a, b], dim=-1)
86
+ w_c1 = self.weights[f'{prefix}.ha1.carry.weight']
87
+ b_c1 = self.weights[f'{prefix}.ha1.carry.bias']
88
+ ha1_carry = self._gate(inp_carry1, w_c1, b_c1)
89
+
90
+ ha2_sum = self._xor(ha1_sum, cin, f'{prefix}.ha2.sum')
91
+
92
+ inp_carry2 = torch.stack([ha1_sum, cin], dim=-1)
93
+ w_c2 = self.weights[f'{prefix}.ha2.carry.weight']
94
+ b_c2 = self.weights[f'{prefix}.ha2.carry.bias']
95
+ ha2_carry = self._gate(inp_carry2, w_c2, b_c2)
96
+
97
+ inp_cout = torch.stack([ha1_carry, ha2_carry], dim=-1)
98
+ w_cout = self.weights[f'{prefix}.carry_or.weight']
99
+ b_cout = self.weights[f'{prefix}.carry_or.bias']
100
+ cout = self._gate(inp_cout, w_cout, b_cout)
101
+
102
+ return ha2_sum, cout
103
+
104
+ def add_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
105
+ """
106
+ 8-bit ripple carry addition.
107
+
108
+ Args:
109
+ a_bits: [batch, 8] MSB-first
110
+ b_bits: [batch, 8] MSB-first
111
+
112
+ Returns:
113
+ result_bits: [batch, 8] MSB-first
114
+ carry_out: [batch] final carry
115
+ """
116
+ batch_size = a_bits.shape[0]
117
+ carry = torch.zeros(batch_size, device=self.device)
118
+ result_bits = []
119
+
120
+ for bit in range(8):
121
+ bit_idx = 7 - bit
122
+ s, carry = self._full_adder(
123
+ a_bits[:, bit_idx],
124
+ b_bits[:, bit_idx],
125
+ carry,
126
+ f'arithmetic.ripplecarry8bit.fa{bit}'
127
+ )
128
+ result_bits.insert(0, s)
129
+
130
+ result = torch.stack(result_bits, dim=1)
131
+ return result, carry
132
+
133
+ def sub_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
134
+ """
135
+ 8-bit subtraction via two's complement: A - B = A + (~B) + 1
136
+
137
+ Args:
138
+ a_bits: [batch, 8] MSB-first
139
+ b_bits: [batch, 8] MSB-first
140
+
141
+ Returns:
142
+ result_bits: [batch, 8] MSB-first
143
+ borrow_out: [batch] (inverted carry)
144
+ """
145
+ b_inv = 1.0 - b_bits
146
+ batch_size = a_bits.shape[0]
147
+ carry = torch.ones(batch_size, device=self.device)
148
+ result_bits = []
149
+
150
+ for bit in range(8):
151
+ bit_idx = 7 - bit
152
+ s, carry = self._full_adder(
153
+ a_bits[:, bit_idx],
154
+ b_inv[:, bit_idx],
155
+ carry,
156
+ f'arithmetic.ripplecarry8bit.fa{bit}'
157
+ )
158
+ result_bits.insert(0, s)
159
+
160
+ result = torch.stack(result_bits, dim=1)
161
+ borrow = 1.0 - carry
162
+ return result, borrow
163
+
164
+ def mul_8bit(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
165
+ """
166
+ 8-bit multiplication via shift-add (software implementation using adder circuits).
167
+ Only keeps low 8 bits of result (matches 8-bit wrap behavior).
168
+
169
+ Args:
170
+ a_bits: [batch, 8] MSB-first
171
+ b_bits: [batch, 8] MSB-first
172
+
173
+ Returns:
174
+ result_bits: [batch, 8] MSB-first (low 8 bits of product)
175
+ """
176
+ batch_size = a_bits.shape[0]
177
+
178
+ acc = torch.zeros(batch_size, 8, device=self.device)
179
+
180
+ for i in range(8):
181
+ b_bit = b_bits[:, 7 - i]
182
+ pp = a_bits * b_bit.unsqueeze(1)
183
+
184
+ shifted_pp = torch.zeros(batch_size, 8, device=self.device)
185
+ for j in range(8):
186
+ dst_idx = j + i
187
+ if dst_idx < 8:
188
+ shifted_pp[:, 7 - dst_idx] = pp[:, 7 - j]
189
+
190
+ acc, _ = self.add_8bit(acc, shifted_pp)
191
+
192
+ return acc
193
+
194
+ def compare_gt(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
195
+ """A > B comparison."""
196
+ inputs = torch.cat([a_bits, b_bits], dim=-1)
197
+ w = self.weights['arithmetic.greaterthan8bit.weight'].view(-1)
198
+ b = self.weights['arithmetic.greaterthan8bit.bias'].view(-1)
199
+ return heaviside_ste((inputs * w).sum(dim=-1) + b)
200
+
201
+ def compare_lt(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
202
+ """A < B comparison."""
203
+ inputs = torch.cat([a_bits, b_bits], dim=-1)
204
+ w = self.weights['arithmetic.lessthan8bit.weight'].view(-1)
205
+ b = self.weights['arithmetic.lessthan8bit.bias'].view(-1)
206
+ return heaviside_ste((inputs * w).sum(dim=-1) + b)
207
+
208
+ def compare_eq(self, a_bits: torch.Tensor, b_bits: torch.Tensor) -> torch.Tensor:
209
+ """A == B comparison (two-layer)."""
210
+ inputs = torch.cat([a_bits, b_bits], dim=-1)
211
+ prefix = 'arithmetic.equality8bit'
212
+
213
+ w_geq = self.weights[f'{prefix}.layer1.geq.weight'].view(-1)
214
+ b_geq = self.weights[f'{prefix}.layer1.geq.bias'].view(-1)
215
+ w_leq = self.weights[f'{prefix}.layer1.leq.weight'].view(-1)
216
+ b_leq = self.weights[f'{prefix}.layer1.leq.bias'].view(-1)
217
+
218
+ h_geq = heaviside_ste((inputs * w_geq).sum(dim=-1) + b_geq)
219
+ h_leq = heaviside_ste((inputs * w_leq).sum(dim=-1) + b_leq)
220
+
221
+ hidden = torch.stack([h_geq, h_leq], dim=-1)
222
+ w2 = self.weights[f'{prefix}.layer2.weight'].view(-1)
223
+ b2 = self.weights[f'{prefix}.layer2.bias'].view(-1)
224
+
225
+ return heaviside_ste((hidden * w2).sum(dim=-1) + b2)
226
+
227
+ def forward(self, a_bits: torch.Tensor, b_bits: torch.Tensor,
228
+ op_onehot: torch.Tensor) -> torch.Tensor:
229
+ """
230
+ Execute operation based on one-hot selector.
231
+ Uses soft routing during training for gradient flow.
232
+
233
+ Args:
234
+ a_bits: [batch, 8] operand A
235
+ b_bits: [batch, 8] operand B
236
+ op_onehot: [batch, 6] one-hot operation selector
237
+ [add, sub, mul, gt, lt, eq]
238
+
239
+ Returns:
240
+ result_bits: [batch, 8] result (comparisons in bit 7, rest zeros)
241
+ """
242
+ batch_size = a_bits.shape[0]
243
+
244
+ add_result, _ = self.add_8bit(a_bits, b_bits)
245
+ sub_result, _ = self.sub_8bit(a_bits, b_bits)
246
+ mul_result = self.mul_8bit(a_bits, b_bits)
247
+
248
+ gt_result = self.compare_gt(a_bits, b_bits)
249
+ lt_result = self.compare_lt(a_bits, b_bits)
250
+ eq_result = self.compare_eq(a_bits, b_bits)
251
+
252
+ cmp_expanded = torch.zeros(batch_size, 8, device=self.device)
253
+
254
+ gt_expanded = cmp_expanded.clone()
255
+ gt_expanded[:, 7] = gt_result
256
+
257
+ lt_expanded = cmp_expanded.clone()
258
+ lt_expanded[:, 7] = lt_result
259
+
260
+ eq_expanded = cmp_expanded.clone()
261
+ eq_expanded[:, 7] = eq_result
262
+
263
+ results = torch.stack([
264
+ add_result,
265
+ sub_result,
266
+ mul_result,
267
+ gt_expanded,
268
+ lt_expanded,
269
+ eq_expanded
270
+ ], dim=1)
271
+
272
+ op_weights = op_onehot.unsqueeze(-1)
273
+ output = (results * op_weights).sum(dim=1)
274
+
275
+ return output
276
+
277
+
278
+ if __name__ == "__main__":
279
+ print("Testing frozen circuits...")
280
+
281
+ circuits = FrozenThresholdCircuits(device='cuda')
282
+ print(f"Loaded {len(circuits.weights)} tensors")
283
+
284
+ a = torch.tensor([[0, 0, 0, 0, 0, 1, 0, 1]], device='cuda', dtype=torch.float32)
285
+ b = torch.tensor([[0, 0, 0, 0, 0, 0, 1, 1]], device='cuda', dtype=torch.float32)
286
+
287
+ result, carry = circuits.add_8bit(a, b)
288
+ val = sum(int(result[0, i].item()) << (7 - i) for i in range(8))
289
+ print(f"5 + 3 = {val} (expected 8)")
290
+
291
+ a = torch.tensor([[0, 1, 1, 0, 0, 1, 0, 0]], device='cuda', dtype=torch.float32)
292
+ b = torch.tensor([[0, 0, 1, 0, 0, 1, 0, 1]], device='cuda', dtype=torch.float32)
293
+ result, _ = circuits.sub_8bit(a, b)
294
+ val = sum(int(result[0, i].item()) << (7 - i) for i in range(8))
295
+ print(f"100 - 37 = {val} (expected 63)")
296
+
297
+ a = torch.tensor([[0, 0, 0, 0, 1, 1, 0, 0]], device='cuda', dtype=torch.float32)
298
+ b = torch.tensor([[0, 0, 0, 0, 1, 0, 1, 1]], device='cuda', dtype=torch.float32)
299
+ result = circuits.mul_8bit(a, b)
300
+ val = sum(int(result[0, i].item()) << (7 - i) for i in range(8))
301
+ print(f"12 * 11 = {val} (expected 132)")
302
+
303
+ a = torch.tensor([[0, 0, 1, 1, 0, 0, 1, 0]], device='cuda', dtype=torch.float32)
304
+ b = torch.tensor([[0, 0, 0, 1, 1, 1, 1, 0]], device='cuda', dtype=torch.float32)
305
+ gt = circuits.compare_gt(a, b)
306
+ lt = circuits.compare_lt(a, b)
307
+ eq = circuits.compare_eq(a, b)
308
+ print(f"50 > 30: {int(gt[0].item())} (expected 1)")
309
+ print(f"50 < 30: {int(lt[0].item())} (expected 0)")
310
+ print(f"50 == 30: {int(eq[0].item())} (expected 0)")
311
+
312
+ print("\nTesting batched forward...")
313
+ batch_a = torch.randint(0, 2, (16, 8), device='cuda', dtype=torch.float32)
314
+ batch_b = torch.randint(0, 2, (16, 8), device='cuda', dtype=torch.float32)
315
+ op = torch.zeros(16, 6, device='cuda')
316
+ op[:, 0] = 1.0
317
+
318
+ result = circuits(batch_a, batch_b, op)
319
+ print(f"Batch output shape: {result.shape}")
320
+ print("Done.")
llm_integration/fitness.py ADDED
@@ -0,0 +1,218 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Shared fitness function for threshold circuit LLM integration.
3
+ Randomized tests, no answer supervision - fitness IS the training signal.
4
+ """
5
+
6
+ import torch
7
+ import random
8
+ from typing import Callable, Dict, Tuple, List
9
+
10
+ OPERATIONS = ['add', 'sub', 'mul', 'gt', 'lt', 'eq']
11
+
12
+ def ground_truth(a: int, b: int, op: str) -> int:
13
+ """Compute expected result (8-bit arithmetic)."""
14
+ if op == 'add':
15
+ return (a + b) & 0xFF
16
+ elif op == 'sub':
17
+ return (a - b) & 0xFF
18
+ elif op == 'mul':
19
+ return (a * b) & 0xFF
20
+ elif op == 'gt':
21
+ return 1 if a > b else 0
22
+ elif op == 'lt':
23
+ return 1 if a < b else 0
24
+ elif op == 'eq':
25
+ return 1 if a == b else 0
26
+ else:
27
+ raise ValueError(f"Unknown op: {op}")
28
+
29
+
30
+ def int_to_bits(val: int, n_bits: int = 8) -> torch.Tensor:
31
+ """Convert integer to bit tensor (MSB first)."""
32
+ bits = torch.zeros(n_bits)
33
+ for i in range(n_bits):
34
+ bits[n_bits - 1 - i] = (val >> i) & 1
35
+ return bits
36
+
37
+
38
+ def bits_to_int(bits: torch.Tensor) -> int:
39
+ """Convert bit tensor to integer (MSB first)."""
40
+ val = 0
41
+ n_bits = bits.shape[-1]
42
+ for i in range(n_bits):
43
+ val += int(bits[..., i].item()) << (n_bits - 1 - i)
44
+ return val
45
+
46
+
47
+ def op_to_idx(op: str) -> int:
48
+ """Convert operation string to index."""
49
+ return OPERATIONS.index(op)
50
+
51
+
52
+ def idx_to_op(idx: int) -> str:
53
+ """Convert index to operation string."""
54
+ return OPERATIONS[idx]
55
+
56
+
57
+ def generate_batch(batch_size: int, device: str = 'cuda') -> Dict[str, torch.Tensor]:
58
+ """
59
+ Generate a batch of random arithmetic problems.
60
+
61
+ Returns:
62
+ Dict with:
63
+ 'a': [batch_size] int tensor of first operands
64
+ 'b': [batch_size] int tensor of second operands
65
+ 'op': [batch_size] int tensor of operation indices
66
+ 'a_bits': [batch_size, 8] bit tensor
67
+ 'b_bits': [batch_size, 8] bit tensor
68
+ 'op_onehot': [batch_size, 6] one-hot operation tensor
69
+ 'expected': [batch_size] int tensor of expected results
70
+ 'expected_bits': [batch_size, 8] bit tensor of expected results
71
+ """
72
+ a_vals = torch.randint(0, 256, (batch_size,), device=device)
73
+ b_vals = torch.randint(0, 256, (batch_size,), device=device)
74
+ op_indices = torch.randint(0, len(OPERATIONS), (batch_size,), device=device)
75
+
76
+ a_bits = torch.zeros(batch_size, 8, device=device)
77
+ b_bits = torch.zeros(batch_size, 8, device=device)
78
+ for i in range(8):
79
+ a_bits[:, 7-i] = (a_vals >> i) & 1
80
+ b_bits[:, 7-i] = (b_vals >> i) & 1
81
+
82
+ op_onehot = torch.zeros(batch_size, len(OPERATIONS), device=device)
83
+ op_onehot.scatter_(1, op_indices.unsqueeze(1), 1.0)
84
+
85
+ expected = torch.zeros(batch_size, dtype=torch.long, device=device)
86
+ for i in range(batch_size):
87
+ a, b, op_idx = a_vals[i].item(), b_vals[i].item(), op_indices[i].item()
88
+ expected[i] = ground_truth(a, b, idx_to_op(op_idx))
89
+
90
+ expected_bits = torch.zeros(batch_size, 8, device=device)
91
+ for i in range(8):
92
+ expected_bits[:, 7-i] = (expected >> i) & 1
93
+
94
+ return {
95
+ 'a': a_vals,
96
+ 'b': b_vals,
97
+ 'op': op_indices,
98
+ 'a_bits': a_bits.float(),
99
+ 'b_bits': b_bits.float(),
100
+ 'op_onehot': op_onehot.float(),
101
+ 'expected': expected,
102
+ 'expected_bits': expected_bits.float(),
103
+ }
104
+
105
+
106
+ def compute_fitness(
107
+ model_fn: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor],
108
+ n_samples: int = 10000,
109
+ batch_size: int = 256,
110
+ device: str = 'cuda',
111
+ return_details: bool = False
112
+ ) -> float | Tuple[float, Dict]:
113
+ """
114
+ Compute fitness score for a model.
115
+
116
+ Args:
117
+ model_fn: Function that takes (a_bits, b_bits, op_onehot) and returns result_bits
118
+ n_samples: Number of test cases
119
+ batch_size: Batch size for evaluation
120
+ device: Device to run on
121
+ return_details: If True, return per-operation breakdown
122
+
123
+ Returns:
124
+ Fitness score in [0, 1], optionally with details dict
125
+ """
126
+ correct = 0
127
+ total = 0
128
+ op_correct = {op: 0 for op in OPERATIONS}
129
+ op_total = {op: 0 for op in OPERATIONS}
130
+
131
+ for _ in range(0, n_samples, batch_size):
132
+ actual_batch = min(batch_size, n_samples - total)
133
+ batch = generate_batch(actual_batch, device)
134
+
135
+ with torch.no_grad():
136
+ pred_bits = model_fn(batch['a_bits'], batch['b_bits'], batch['op_onehot'])
137
+
138
+ pred_bits_binary = (pred_bits > 0.5).float()
139
+
140
+ for i in range(actual_batch):
141
+ pred_val = 0
142
+ for j in range(8):
143
+ pred_val += int(pred_bits_binary[i, j].item()) << (7 - j)
144
+
145
+ expected_val = batch['expected'][i].item()
146
+ op_name = idx_to_op(batch['op'][i].item())
147
+
148
+ op_total[op_name] += 1
149
+ total += 1
150
+
151
+ if pred_val == expected_val:
152
+ correct += 1
153
+ op_correct[op_name] += 1
154
+
155
+ fitness = correct / total if total > 0 else 0.0
156
+
157
+ if return_details:
158
+ details = {
159
+ 'correct': correct,
160
+ 'total': total,
161
+ 'by_op': {
162
+ op: {
163
+ 'correct': op_correct[op],
164
+ 'total': op_total[op],
165
+ 'accuracy': op_correct[op] / op_total[op] if op_total[op] > 0 else 0.0
166
+ }
167
+ for op in OPERATIONS
168
+ }
169
+ }
170
+ return fitness, details
171
+
172
+ return fitness
173
+
174
+
175
+ def compute_bit_accuracy(pred_bits: torch.Tensor, expected_bits: torch.Tensor) -> float:
176
+ """Compute per-bit accuracy (for gradient signal analysis)."""
177
+ pred_binary = (pred_bits > 0.5).float()
178
+ return (pred_binary == expected_bits).float().mean().item()
179
+
180
+
181
+ def compute_loss(pred_bits: torch.Tensor, expected_bits: torch.Tensor) -> torch.Tensor:
182
+ """Binary cross-entropy loss on output bits."""
183
+ pred_clamped = pred_bits.clamp(1e-7, 1 - 1e-7)
184
+ return -((expected_bits * torch.log(pred_clamped) +
185
+ (1 - expected_bits) * torch.log(1 - pred_clamped))).mean()
186
+
187
+
188
+ if __name__ == "__main__":
189
+ print("Testing fitness module...")
190
+
191
+ batch = generate_batch(8, 'cpu')
192
+ print(f"\nSample batch:")
193
+ for i in range(4):
194
+ a, b = batch['a'][i].item(), batch['b'][i].item()
195
+ op = idx_to_op(batch['op'][i].item())
196
+ expected = batch['expected'][i].item()
197
+ print(f" {a} {op} {b} = {expected}")
198
+
199
+ def random_model(a_bits, b_bits, op_onehot):
200
+ return torch.rand(a_bits.shape[0], 8, device=a_bits.device)
201
+
202
+ fitness = compute_fitness(random_model, n_samples=1000, batch_size=100, device='cpu')
203
+ print(f"\nRandom model fitness: {fitness:.4f} (expected ~0.004 for 8-bit)")
204
+
205
+ def perfect_model(a_bits, b_bits, op_onehot):
206
+ batch_size = a_bits.shape[0]
207
+ results = torch.zeros(batch_size, 8, device=a_bits.device)
208
+ for i in range(batch_size):
209
+ a = sum(int(a_bits[i, j].item()) << (7-j) for j in range(8))
210
+ b = sum(int(b_bits[i, j].item()) << (7-j) for j in range(8))
211
+ op_idx = op_onehot[i].argmax().item()
212
+ result = ground_truth(a, b, idx_to_op(op_idx))
213
+ for j in range(8):
214
+ results[i, 7-j] = (result >> j) & 1
215
+ return results
216
+
217
+ fitness = compute_fitness(perfect_model, n_samples=1000, batch_size=100, device='cpu')
218
+ print(f"Perfect model fitness: {fitness:.4f} (expected 1.0)")
llm_integration/model.py ADDED
@@ -0,0 +1,1263 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Trainable interface layers for frozen threshold circuits.
3
+ BitEncoder, OpRouter, BitDecoder wrap the frozen circuits.
4
+ HiddenStateExtractor and AugmentedArithmeticModel for LLM integration.
5
+ """
6
+
7
+ import torch
8
+ import torch.nn as nn
9
+ import torch.nn.functional as F
10
+ from circuits import FrozenThresholdCircuits, heaviside_ste
11
+
12
+ MODEL_ID = 'HuggingFaceTB/SmolLM2-360M-Instruct'
13
+ OPERATIONS = ['add', 'sub', 'mul', 'gt', 'lt', 'eq']
14
+ OP_SYMBOLS = {'add': '+', 'sub': '-', 'mul': '*', 'gt': '>', 'lt': '<', 'eq': '=='}
15
+
16
+
17
+ class BitEncoder(nn.Module):
18
+ """
19
+ Encodes two 8-bit operands from input representation.
20
+ Uses residual connection to preserve ground truth bits while allowing learned refinement.
21
+ """
22
+
23
+ def __init__(self, input_dim: int = 16 + 6, hidden_dim: int = 32):
24
+ super().__init__()
25
+ self.refine = nn.Sequential(
26
+ nn.Linear(input_dim, hidden_dim),
27
+ nn.Tanh(),
28
+ nn.Linear(hidden_dim, 16),
29
+ )
30
+ self.scale = nn.Parameter(torch.tensor(0.0))
31
+
32
+ def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
33
+ """
34
+ Args:
35
+ x: [batch, input_dim] input with first 16 dims being a_bits, b_bits
36
+
37
+ Returns:
38
+ a_bits: [batch, 8] first operand bits
39
+ b_bits: [batch, 8] second operand bits
40
+ """
41
+ base_bits = x[:, :16]
42
+ refinement = self.refine(x) * torch.sigmoid(self.scale)
43
+ bits = base_bits + refinement
44
+ bits = torch.clamp(bits, 0, 1)
45
+ hard_bits = heaviside_ste(bits - 0.5)
46
+ out = hard_bits - bits.detach() + bits
47
+
48
+ return out[:, :8], out[:, 8:]
49
+
50
+
51
+ class OpRouter(nn.Module):
52
+ """
53
+ Routes computation to the appropriate circuit based on input.
54
+ Outputs soft weights over operations for gradient flow.
55
+ """
56
+
57
+ def __init__(self, input_dim: int = 16 + 6, hidden_dim: int = 32, n_ops: int = 6):
58
+ super().__init__()
59
+ self.net = nn.Sequential(
60
+ nn.Linear(input_dim, hidden_dim),
61
+ nn.ReLU(),
62
+ nn.Linear(hidden_dim, n_ops),
63
+ )
64
+
65
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
66
+ """
67
+ Args:
68
+ x: [batch, input_dim] input features
69
+
70
+ Returns:
71
+ op_weights: [batch, n_ops] soft operation weights (softmax)
72
+ """
73
+ logits = self.net(x)
74
+ return F.softmax(logits, dim=-1)
75
+
76
+
77
+ class BitDecoder(nn.Module):
78
+ """
79
+ Decodes circuit output bits to target representation.
80
+ For standalone training: outputs soft bits for loss computation.
81
+ For LLM integration: would project to hidden state delta.
82
+ """
83
+
84
+ def __init__(self, output_dim: int = 8):
85
+ super().__init__()
86
+ self.output_dim = output_dim
87
+
88
+ def forward(self, result_bits: torch.Tensor) -> torch.Tensor:
89
+ return result_bits
90
+
91
+
92
+ class ThresholdALU(nn.Module):
93
+ """
94
+ Complete trainable interface + frozen circuits.
95
+ Learns to encode inputs, route to circuits, decode outputs.
96
+ """
97
+
98
+ def __init__(self, device: str = 'cuda'):
99
+ super().__init__()
100
+ self.device = device
101
+
102
+ self.circuits = FrozenThresholdCircuits(device=device)
103
+
104
+ for key in self.circuits.weights:
105
+ self.circuits.weights[key].requires_grad = False
106
+
107
+ self.encoder = BitEncoder(input_dim=16 + 6, hidden_dim=64).to(device)
108
+ self.router = OpRouter(input_dim=16 + 6, hidden_dim=32, n_ops=6).to(device)
109
+ self.decoder = BitDecoder(output_dim=8).to(device)
110
+
111
+ def forward(self, a_bits_in: torch.Tensor, b_bits_in: torch.Tensor,
112
+ op_onehot: torch.Tensor) -> torch.Tensor:
113
+ """
114
+ Forward pass through trainable interface + frozen circuits.
115
+
116
+ Args:
117
+ a_bits_in: [batch, 8] input A bits (ground truth for training)
118
+ b_bits_in: [batch, 8] input B bits (ground truth for training)
119
+ op_onehot: [batch, 6] one-hot operation selector
120
+
121
+ Returns:
122
+ result_bits: [batch, 8] output bits
123
+ """
124
+ x = torch.cat([a_bits_in, b_bits_in, op_onehot], dim=-1)
125
+
126
+ a_bits, b_bits = self.encoder(x)
127
+
128
+ op_weights = self.router(x)
129
+
130
+ result = self.circuits(a_bits, b_bits, op_weights)
131
+
132
+ output = self.decoder(result)
133
+
134
+ return output
135
+
136
+ def forward_direct(self, a_bits: torch.Tensor, b_bits: torch.Tensor,
137
+ op_onehot: torch.Tensor) -> torch.Tensor:
138
+ """
139
+ Direct forward through circuits (bypass encoder/router for testing).
140
+ """
141
+ return self.circuits(a_bits, b_bits, op_onehot)
142
+
143
+
144
+ class DirectCircuitModel(nn.Module):
145
+ """
146
+ Minimal model that directly uses circuits without learned encoding.
147
+ For validating that circuits themselves achieve 100% fitness.
148
+ """
149
+
150
+ def __init__(self, device: str = 'cuda'):
151
+ super().__init__()
152
+ self.device = device
153
+ self.circuits = FrozenThresholdCircuits(device=device)
154
+
155
+ def forward(self, a_bits: torch.Tensor, b_bits: torch.Tensor,
156
+ op_onehot: torch.Tensor) -> torch.Tensor:
157
+ return self.circuits(a_bits, b_bits, op_onehot)
158
+
159
+
160
+ class HiddenStateExtractor(nn.Module):
161
+ """
162
+ Extracts operands and operation from LLM hidden states.
163
+ This is the hard part - must learn to parse numbers from embeddings.
164
+ """
165
+
166
+ def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256):
167
+ super().__init__()
168
+
169
+ self.a_extractor = nn.Sequential(
170
+ nn.Linear(hidden_dim, intermediate_dim),
171
+ nn.GELU(),
172
+ nn.Linear(intermediate_dim, 8),
173
+ )
174
+
175
+ self.b_extractor = nn.Sequential(
176
+ nn.Linear(hidden_dim, intermediate_dim),
177
+ nn.GELU(),
178
+ nn.Linear(intermediate_dim, 8),
179
+ )
180
+
181
+ self.op_router = nn.Sequential(
182
+ nn.Linear(hidden_dim, intermediate_dim),
183
+ nn.GELU(),
184
+ nn.Linear(intermediate_dim, len(OPERATIONS)),
185
+ )
186
+
187
+ def forward(self, hidden_states: torch.Tensor):
188
+ """
189
+ Args:
190
+ hidden_states: [batch, hidden_dim] from LLM
191
+
192
+ Returns:
193
+ a_bits: [batch, 8]
194
+ b_bits: [batch, 8]
195
+ op_logits: [batch, 6]
196
+ """
197
+ a_logits = self.a_extractor(hidden_states)
198
+ b_logits = self.b_extractor(hidden_states)
199
+ op_logits = self.op_router(hidden_states)
200
+
201
+ a_soft = torch.sigmoid(a_logits)
202
+ b_soft = torch.sigmoid(b_logits)
203
+
204
+ a_hard = heaviside_ste(a_logits)
205
+ b_hard = heaviside_ste(b_logits)
206
+
207
+ a_bits = a_hard - a_soft.detach() + a_soft
208
+ b_bits = b_hard - b_soft.detach() + b_soft
209
+
210
+ return a_bits, b_bits, op_logits
211
+
212
+
213
+ class AttentionPooling(nn.Module):
214
+ """
215
+ Learnable attention pooling over sequence positions.
216
+ Replaces mean pooling - learns which tokens matter for extraction.
217
+ """
218
+
219
+ def __init__(self, hidden_dim: int = 960, num_heads: int = 4):
220
+ super().__init__()
221
+ self.num_heads = num_heads
222
+ self.head_dim = hidden_dim // num_heads
223
+
224
+ self.query = nn.Linear(hidden_dim, hidden_dim)
225
+ self.key = nn.Linear(hidden_dim, hidden_dim)
226
+ self.value = nn.Linear(hidden_dim, hidden_dim)
227
+ self.out_proj = nn.Linear(hidden_dim, hidden_dim)
228
+
229
+ self.cls_token = nn.Parameter(torch.randn(1, 1, hidden_dim) * 0.02)
230
+
231
+ def forward(self, embeddings: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
232
+ """
233
+ Args:
234
+ embeddings: [batch, seq_len, hidden_dim]
235
+ mask: [batch, seq_len] attention mask (1 = attend, 0 = ignore)
236
+
237
+ Returns:
238
+ pooled: [batch, hidden_dim]
239
+ """
240
+ batch_size, seq_len, hidden_dim = embeddings.shape
241
+
242
+ cls_expanded = self.cls_token.expand(batch_size, -1, -1)
243
+ embeddings = torch.cat([cls_expanded, embeddings], dim=1)
244
+
245
+ cls_mask = torch.ones(batch_size, 1, device=mask.device)
246
+ mask = torch.cat([cls_mask, mask], dim=1)
247
+
248
+ Q = self.query(embeddings[:, :1, :])
249
+ K = self.key(embeddings)
250
+ V = self.value(embeddings)
251
+
252
+ Q = Q.view(batch_size, 1, self.num_heads, self.head_dim).transpose(1, 2)
253
+ K = K.view(batch_size, seq_len + 1, self.num_heads, self.head_dim).transpose(1, 2)
254
+ V = V.view(batch_size, seq_len + 1, self.num_heads, self.head_dim).transpose(1, 2)
255
+
256
+ scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
257
+
258
+ mask_expanded = mask.unsqueeze(1).unsqueeze(2)
259
+ scores = scores.masked_fill(mask_expanded == 0, -1e9)
260
+
261
+ attn_weights = torch.softmax(scores, dim=-1)
262
+ attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
263
+
264
+ context = torch.matmul(attn_weights, V)
265
+ context = context.transpose(1, 2).contiguous().view(batch_size, 1, hidden_dim)
266
+
267
+ pooled = self.out_proj(context).squeeze(1)
268
+ pooled = torch.nan_to_num(pooled, nan=0.0)
269
+
270
+ return pooled
271
+
272
+
273
+ class MultiHeadBitExtractor(nn.Module):
274
+ """
275
+ 8 separate extractors for 8 bits - each bit gets its own specialized network.
276
+ More expressive than single MLP predicting all 8 bits at once.
277
+ """
278
+
279
+ def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 128):
280
+ super().__init__()
281
+
282
+ self.bit_extractors = nn.ModuleList([
283
+ nn.Sequential(
284
+ nn.Linear(hidden_dim, intermediate_dim),
285
+ nn.GELU(),
286
+ nn.Linear(intermediate_dim, 1),
287
+ )
288
+ for _ in range(8)
289
+ ])
290
+
291
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
292
+ """
293
+ Args:
294
+ hidden_states: [batch, hidden_dim]
295
+
296
+ Returns:
297
+ bits: [batch, 8] - one bit from each extractor
298
+ """
299
+ hidden_states = torch.nan_to_num(hidden_states, nan=0.0)
300
+
301
+ bit_logits = [extractor(hidden_states) for extractor in self.bit_extractors]
302
+ logits = torch.cat(bit_logits, dim=-1)
303
+ logits = torch.clamp(logits, -20, 20)
304
+
305
+ soft = torch.sigmoid(logits)
306
+ hard = heaviside_ste(logits)
307
+ bits = hard - soft.detach() + soft
308
+
309
+ return bits, logits
310
+
311
+
312
+ class Extractor(nn.Module):
313
+ """
314
+ Extracts operands and operation from LLM hidden states.
315
+ Uses attention pooling and per-bit extraction networks.
316
+ """
317
+
318
+ def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256, num_heads: int = 4):
319
+ super().__init__()
320
+
321
+ self.attention_pool = AttentionPooling(hidden_dim, num_heads)
322
+
323
+ self.a_extractor = MultiHeadBitExtractor(hidden_dim, intermediate_dim // 2)
324
+ self.b_extractor = MultiHeadBitExtractor(hidden_dim, intermediate_dim // 2)
325
+
326
+ self.op_router = nn.Sequential(
327
+ nn.Linear(hidden_dim, intermediate_dim),
328
+ nn.GELU(),
329
+ nn.Linear(intermediate_dim, len(OPERATIONS)),
330
+ )
331
+
332
+ def forward(self, embeddings: torch.Tensor, mask: torch.Tensor):
333
+ """
334
+ Args:
335
+ embeddings: [batch, seq_len, hidden_dim]
336
+ mask: [batch, seq_len]
337
+
338
+ Returns:
339
+ a_bits: [batch, 8]
340
+ b_bits: [batch, 8]
341
+ op_logits: [batch, 6]
342
+ """
343
+ pooled = self.attention_pool(embeddings, mask)
344
+
345
+ a_bits, _ = self.a_extractor(pooled)
346
+ b_bits, _ = self.b_extractor(pooled)
347
+ op_logits = self.op_router(pooled)
348
+
349
+ return a_bits, b_bits, op_logits
350
+
351
+
352
+ class PositionExtractor(nn.Module):
353
+ """
354
+ Position-specific extraction with dynamic operator detection.
355
+
356
+ Tokenization pattern for "A op B":
357
+ [A_digits...] [operator] [space] [B_digits...]
358
+
359
+ Examples:
360
+ "5 + 3" -> ['5', ' +', ' ', '3'] (positions: A=0, op=1, B=3)
361
+ "47 + 86" -> ['4', '7', ' +', ' ', '8', '6'] (positions: A=0-1, op=2, B=4-5)
362
+ "127 + 128" -> ['1','2','7',' +', ' ','1','2','8'] (positions: A=0-2, op=3, B=5-7)
363
+
364
+ Token IDs (SmolLM2):
365
+ Digits '0'-'9': 32-41
366
+ Operators: ' +'=1232, ' -'=731, ' *'=1672, ' >'=2986, ' <'=2067, ' =='=1758
367
+ Space: 216
368
+ """
369
+
370
+ DIGIT_TOKENS = set(range(32, 42))
371
+ OPERATOR_TOKENS = {
372
+ 1232: 0, # ' +' -> add
373
+ 731: 1, # ' -' -> sub
374
+ 1672: 2, # ' *' -> mul
375
+ 2986: 3, # ' >' -> gt
376
+ 2067: 4, # ' <' -> lt
377
+ 1758: 5, # ' ==' -> eq
378
+ }
379
+ SPACE_TOKEN = 216
380
+ MAX_DIGITS = 3
381
+
382
+ def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256):
383
+ super().__init__()
384
+ self.hidden_dim = hidden_dim
385
+
386
+ self.a_extractor = nn.Sequential(
387
+ nn.Linear(hidden_dim * self.MAX_DIGITS, intermediate_dim),
388
+ nn.GELU(),
389
+ nn.Linear(intermediate_dim, intermediate_dim // 2),
390
+ nn.GELU(),
391
+ nn.Linear(intermediate_dim // 2, 8),
392
+ )
393
+
394
+ self.b_extractor = nn.Sequential(
395
+ nn.Linear(hidden_dim * self.MAX_DIGITS, intermediate_dim),
396
+ nn.GELU(),
397
+ nn.Linear(intermediate_dim, intermediate_dim // 2),
398
+ nn.GELU(),
399
+ nn.Linear(intermediate_dim // 2, 8),
400
+ )
401
+
402
+ self.op_extractor = nn.Sequential(
403
+ nn.Linear(hidden_dim, intermediate_dim // 2),
404
+ nn.GELU(),
405
+ nn.Linear(intermediate_dim // 2, len(OPERATIONS)),
406
+ )
407
+
408
+ def _find_operator_position(self, token_ids: torch.Tensor) -> tuple[int, int]:
409
+ """
410
+ Find operator token position and its operation index.
411
+
412
+ Args:
413
+ token_ids: [seq_len] tensor of token IDs
414
+
415
+ Returns:
416
+ (position, op_index) or (-1, -1) if not found
417
+ """
418
+ for pos, tid in enumerate(token_ids.tolist()):
419
+ if tid in self.OPERATOR_TOKENS:
420
+ return pos, self.OPERATOR_TOKENS[tid]
421
+ return -1, -1
422
+
423
+ def _extract_digit_features(self, hidden: torch.Tensor, start: int, end: int) -> torch.Tensor:
424
+ """
425
+ Extract and pad digit hidden states to fixed size.
426
+
427
+ Args:
428
+ hidden: [seq_len, hidden_dim]
429
+ start: start position (inclusive)
430
+ end: end position (exclusive)
431
+
432
+ Returns:
433
+ [hidden_dim * MAX_DIGITS] flattened features, zero-padded on the LEFT
434
+ (so units digit is always at the same position regardless of number length)
435
+ """
436
+ n_digits = end - start
437
+ features = torch.zeros(self.MAX_DIGITS * self.hidden_dim, device=hidden.device)
438
+
439
+ if n_digits > 0 and n_digits <= self.MAX_DIGITS:
440
+ digit_hidden = hidden[start:end, :].reshape(-1)
441
+ pad_size = (self.MAX_DIGITS - n_digits) * self.hidden_dim
442
+ features[pad_size:] = digit_hidden
443
+
444
+ return features
445
+
446
+ def forward(self, hidden: torch.Tensor, mask: torch.Tensor, token_ids: torch.Tensor = None):
447
+ """
448
+ Args:
449
+ hidden: [batch, seq_len, hidden_dim]
450
+ mask: [batch, seq_len] attention mask
451
+ token_ids: [batch, seq_len] token IDs (required for operator detection)
452
+
453
+ Returns:
454
+ a_bits: [batch, 8]
455
+ b_bits: [batch, 8]
456
+ op_logits: [batch, 6]
457
+ """
458
+ if token_ids is None:
459
+ raise ValueError("PositionExtractor requires token_ids for operator detection")
460
+
461
+ batch_size, seq_len, hidden_dim = hidden.shape
462
+ device = hidden.device
463
+
464
+ a_features = []
465
+ b_features = []
466
+ op_features = []
467
+ op_indices = []
468
+
469
+ for i in range(batch_size):
470
+ seq_mask = mask[i].bool()
471
+ valid_len = seq_mask.sum().item()
472
+ start_pos = seq_len - valid_len
473
+
474
+ valid_tokens = token_ids[i, start_pos:]
475
+ valid_hidden = hidden[i, start_pos:, :]
476
+
477
+ op_pos, op_idx = self._find_operator_position(valid_tokens)
478
+
479
+ if op_pos == -1:
480
+ a_feat = torch.zeros(self.MAX_DIGITS * hidden_dim, device=device)
481
+ b_feat = torch.zeros(self.MAX_DIGITS * hidden_dim, device=device)
482
+ op_feat = torch.zeros(hidden_dim, device=device)
483
+ op_idx = 0
484
+ else:
485
+ a_feat = self._extract_digit_features(valid_hidden, 0, op_pos)
486
+
487
+ op_feat = valid_hidden[op_pos, :]
488
+
489
+ b_start = op_pos + 2 if (op_pos + 1 < valid_len and
490
+ valid_tokens[op_pos + 1].item() == self.SPACE_TOKEN) else op_pos + 1
491
+ b_feat = self._extract_digit_features(valid_hidden, b_start, valid_len)
492
+
493
+ a_features.append(a_feat)
494
+ b_features.append(b_feat)
495
+ op_features.append(op_feat)
496
+ op_indices.append(op_idx)
497
+
498
+ a_features = torch.stack(a_features)
499
+ b_features = torch.stack(b_features)
500
+ op_features = torch.stack(op_features)
501
+ op_indices_tensor = torch.tensor(op_indices, device=device, dtype=torch.long)
502
+
503
+ a_logits = self.a_extractor(a_features)
504
+ b_logits = self.b_extractor(b_features)
505
+ op_logits = self.op_extractor(op_features)
506
+
507
+ a_soft = torch.sigmoid(a_logits)
508
+ b_soft = torch.sigmoid(b_logits)
509
+ a_hard = heaviside_ste(a_logits)
510
+ b_hard = heaviside_ste(b_logits)
511
+ a_bits = a_hard - a_soft.detach() + a_soft
512
+ b_bits = b_hard - b_soft.detach() + b_soft
513
+
514
+ return a_bits, b_bits, op_logits, op_indices_tensor
515
+
516
+
517
+ class PositionalDigitExtractor(nn.Module):
518
+ """
519
+ Position-aware digit extraction: classifies each digit position independently.
520
+
521
+ This approach achieves 100% accuracy because:
522
+ 1. Each digit token position is classified independently (100% accuracy on layer 0)
523
+ 2. Numbers are reconstructed using place values (×100, ×10, ×1)
524
+ 3. No information is lost through pooling
525
+
526
+ Token IDs (SmolLM2):
527
+ Digits '0'-'9': 32-41
528
+ Operators: ' +'=1232, ' -'=731, ' *'=1672, ' >'=2986, ' <'=2067, ' =='=1758
529
+ Space: 216
530
+ """
531
+
532
+ DIGIT_TOKENS = set(range(32, 42))
533
+ OPERATOR_TOKENS = {
534
+ 1232: 0, # ' +' -> add
535
+ 731: 1, # ' -' -> sub
536
+ 1672: 2, # ' *' -> mul
537
+ 2986: 3, # ' >' -> gt
538
+ 2067: 4, # ' <' -> lt
539
+ 1758: 5, # ' ==' -> eq
540
+ }
541
+ SPACE_TOKEN = 216
542
+
543
+ def __init__(self, hidden_dim: int = 960):
544
+ super().__init__()
545
+ self.hidden_dim = hidden_dim
546
+
547
+ self.digit_classifier = nn.Linear(hidden_dim, 10)
548
+
549
+ self.op_classifier = nn.Linear(hidden_dim, len(OPERATIONS))
550
+
551
+ def _find_positions(self, token_ids: torch.Tensor) -> tuple:
552
+ """Find A digit positions, B digit positions, and operator position."""
553
+ token_list = token_ids.tolist()
554
+
555
+ op_pos = -1
556
+ op_idx = 0
557
+ for i, tid in enumerate(token_list):
558
+ if tid in self.OPERATOR_TOKENS:
559
+ op_pos = i
560
+ op_idx = self.OPERATOR_TOKENS[tid]
561
+ break
562
+
563
+ if op_pos == -1:
564
+ return [], [], -1, 0
565
+
566
+ a_positions = [i for i in range(op_pos) if token_list[i] in self.DIGIT_TOKENS]
567
+
568
+ b_start = op_pos + 2 if (op_pos + 1 < len(token_list) and
569
+ token_list[op_pos + 1] == self.SPACE_TOKEN) else op_pos + 1
570
+ b_positions = [i for i in range(b_start, len(token_list)) if token_list[i] in self.DIGIT_TOKENS]
571
+
572
+ return a_positions, b_positions, op_pos, op_idx
573
+
574
+ def _predict_value(self, hidden: torch.Tensor, positions: list) -> tuple:
575
+ """Predict digit at each position and reconstruct number."""
576
+ if not positions:
577
+ return torch.tensor(0.0, device=hidden.device), []
578
+
579
+ digit_logits_list = []
580
+ soft_value = torch.tensor(0.0, device=hidden.device)
581
+
582
+ for idx, pos in enumerate(positions):
583
+ logits = self.digit_classifier(hidden[pos])
584
+ digit_logits_list.append(logits)
585
+
586
+ probs = torch.softmax(logits, dim=-1)
587
+ digit_values = torch.arange(10, device=hidden.device, dtype=torch.float32)
588
+ soft_digit = (probs * digit_values).sum()
589
+
590
+ place_value = 10 ** (len(positions) - idx - 1)
591
+ soft_value = soft_value + soft_digit * place_value
592
+
593
+ return soft_value, digit_logits_list
594
+
595
+ def _value_to_bits(self, value: torch.Tensor) -> torch.Tensor:
596
+ """Convert soft value to 8 bits using differentiable operations."""
597
+ value = torch.clamp(value, 0, 255)
598
+
599
+ bits = []
600
+ for i in range(7, -1, -1):
601
+ bit = torch.sigmoid((value - (2 ** i - 0.5)) * 10)
602
+ value = value - bit * (2 ** i)
603
+ bits.append(bit)
604
+
605
+ return torch.stack(bits)
606
+
607
+ def forward(self, hidden: torch.Tensor, mask: torch.Tensor, token_ids: torch.Tensor):
608
+ """
609
+ Args:
610
+ hidden: [batch, seq_len, hidden_dim]
611
+ mask: [batch, seq_len]
612
+ token_ids: [batch, seq_len]
613
+
614
+ Returns:
615
+ a_bits: [batch, 8]
616
+ b_bits: [batch, 8]
617
+ op_logits: [batch, 6]
618
+ op_indices: [batch] ground truth op from tokens
619
+ a_digit_logits: list of [batch, 10] per digit position
620
+ b_digit_logits: list of [batch, 10] per digit position
621
+ """
622
+ batch_size = hidden.shape[0]
623
+ device = hidden.device
624
+
625
+ a_bits_list = []
626
+ b_bits_list = []
627
+ op_logits_list = []
628
+ op_indices_list = []
629
+ a_values_list = []
630
+ b_values_list = []
631
+ a_digit_logits_list = []
632
+ b_digit_logits_list = []
633
+
634
+ for i in range(batch_size):
635
+ seq_mask = mask[i].bool()
636
+ valid_len = seq_mask.sum().item()
637
+ start_pos = hidden.shape[1] - valid_len
638
+
639
+ valid_hidden = hidden[i, start_pos:]
640
+ valid_tokens = token_ids[i, start_pos:]
641
+
642
+ a_pos, b_pos, op_pos, op_idx = self._find_positions(valid_tokens)
643
+
644
+ a_value, a_digit_logits = self._predict_value(valid_hidden, a_pos)
645
+ b_value, b_digit_logits = self._predict_value(valid_hidden, b_pos)
646
+
647
+ a_bits = self._value_to_bits(a_value)
648
+ b_bits = self._value_to_bits(b_value)
649
+
650
+ if op_pos >= 0:
651
+ op_logits = self.op_classifier(valid_hidden[op_pos])
652
+ else:
653
+ op_logits = torch.zeros(len(OPERATIONS), device=device)
654
+
655
+ a_bits_list.append(a_bits)
656
+ b_bits_list.append(b_bits)
657
+ op_logits_list.append(op_logits)
658
+ op_indices_list.append(op_idx)
659
+ a_values_list.append(a_value)
660
+ b_values_list.append(b_value)
661
+ a_digit_logits_list.append(a_digit_logits)
662
+ b_digit_logits_list.append(b_digit_logits)
663
+
664
+ a_bits = torch.stack(a_bits_list)
665
+ b_bits = torch.stack(b_bits_list)
666
+ op_logits = torch.stack(op_logits_list)
667
+ op_indices = torch.tensor(op_indices_list, device=device, dtype=torch.long)
668
+ a_values = torch.stack(a_values_list)
669
+ b_values = torch.stack(b_values_list)
670
+
671
+ return a_bits, b_bits, op_logits, op_indices, a_values, b_values, a_digit_logits_list, b_digit_logits_list
672
+
673
+
674
+ class DigitExtractor(nn.Module):
675
+ """
676
+ Digit-level extraction: predicts digits (0-9) then converts to bits.
677
+ Uses attention pooling (less accurate than PositionalDigitExtractor).
678
+ """
679
+
680
+ def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256, num_heads: int = 4):
681
+ super().__init__()
682
+
683
+ self.attention_pool = AttentionPooling(hidden_dim, num_heads)
684
+
685
+ self.a_digit_pred = nn.Sequential(
686
+ nn.Linear(hidden_dim, intermediate_dim),
687
+ nn.GELU(),
688
+ nn.Linear(intermediate_dim, 3 * 10),
689
+ )
690
+
691
+ self.b_digit_pred = nn.Sequential(
692
+ nn.Linear(hidden_dim, intermediate_dim),
693
+ nn.GELU(),
694
+ nn.Linear(intermediate_dim, 3 * 10),
695
+ )
696
+
697
+ self.op_router = nn.Sequential(
698
+ nn.Linear(hidden_dim, intermediate_dim),
699
+ nn.GELU(),
700
+ nn.Linear(intermediate_dim, len(OPERATIONS)),
701
+ )
702
+
703
+ def digits_to_bits(self, digit_logits: torch.Tensor) -> torch.Tensor:
704
+ """
705
+ Convert 3-digit predictions to 8-bit representation.
706
+ digit_logits: [batch, 30] (3 digits * 10 classes each)
707
+ Returns: [batch, 8] bits
708
+ """
709
+ batch_size = digit_logits.shape[0]
710
+
711
+ logits = digit_logits.view(batch_size, 3, 10)
712
+ probs = torch.softmax(logits, dim=-1)
713
+
714
+ digit_values = torch.arange(10, device=digit_logits.device).float()
715
+ soft_digits = (probs * digit_values).sum(dim=-1)
716
+
717
+ hundreds = soft_digits[:, 0]
718
+ tens = soft_digits[:, 1]
719
+ ones = soft_digits[:, 2]
720
+
721
+ value = hundreds * 100 + tens * 10 + ones
722
+ value = torch.clamp(value, 0, 255)
723
+
724
+ bits = []
725
+ for i in range(7, -1, -1):
726
+ bit = torch.fmod(torch.floor(value / (2 ** i)), 2)
727
+ bits.append(bit)
728
+
729
+ return torch.stack(bits, dim=-1)
730
+
731
+ def forward(self, hidden: torch.Tensor, mask: torch.Tensor):
732
+ """
733
+ Returns:
734
+ a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits
735
+ """
736
+ pooled = self.attention_pool(hidden, mask)
737
+
738
+ a_digit_logits = self.a_digit_pred(pooled)
739
+ b_digit_logits = self.b_digit_pred(pooled)
740
+ op_logits = self.op_router(pooled)
741
+
742
+ a_bits = self.digits_to_bits(a_digit_logits)
743
+ b_bits = self.digits_to_bits(b_digit_logits)
744
+
745
+ return a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits
746
+
747
+
748
+ class HybridExtractor(nn.Module):
749
+ """
750
+ Hybrid extractor that handles both digit tokens and word numbers.
751
+
752
+ For digit tokens (32-41): Direct lookup, no training needed
753
+ For word numbers: Learned MLP extraction from pooled hidden states
754
+
755
+ This is the real training target - learning to extract numbers from
756
+ natural language like "forty seven plus eighty six".
757
+ """
758
+
759
+ DIGIT_TOKENS = set(range(32, 42))
760
+ SYMBOL_OP_TOKENS = {
761
+ 1232: 0, # ' +' -> add
762
+ 731: 1, # ' -' -> sub
763
+ 1672: 2, # ' *' -> mul
764
+ 2986: 3, # ' >' -> gt
765
+ 2067: 4, # ' <' -> lt
766
+ 1758: 5, # ' ==' -> eq
767
+ }
768
+ WORD_OP_TOKENS = {
769
+ 2068: 0, # 'plus' -> add
770
+ 8500: 1, # 'minus' -> sub
771
+ 1580: 2, # 'times' -> mul
772
+ 6301: 3, # 'greater' -> gt
773
+ 1912: 4, # 'less' -> lt
774
+ 16364: 5, # 'equals' -> eq
775
+ 11540: 5, # 'equal' -> eq
776
+ }
777
+ ALL_OP_TOKENS = {**SYMBOL_OP_TOKENS, **WORD_OP_TOKENS}
778
+
779
+ def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256, num_heads: int = 4):
780
+ super().__init__()
781
+ self.hidden_dim = hidden_dim
782
+
783
+ self.attention_pool = AttentionPooling(hidden_dim, num_heads)
784
+ self.a_pool = AttentionPooling(hidden_dim, num_heads)
785
+ self.b_pool = AttentionPooling(hidden_dim, num_heads)
786
+
787
+ self.a_digit_pred = nn.Sequential(
788
+ nn.Linear(hidden_dim, intermediate_dim),
789
+ nn.GELU(),
790
+ nn.Dropout(0.1),
791
+ nn.Linear(intermediate_dim, 3 * 10),
792
+ )
793
+
794
+ self.b_digit_pred = nn.Sequential(
795
+ nn.Linear(hidden_dim, intermediate_dim),
796
+ nn.GELU(),
797
+ nn.Dropout(0.1),
798
+ nn.Linear(intermediate_dim, 3 * 10),
799
+ )
800
+
801
+ self.op_predictor = nn.Sequential(
802
+ nn.Linear(hidden_dim, intermediate_dim // 2),
803
+ nn.GELU(),
804
+ nn.Linear(intermediate_dim // 2, len(OPERATIONS)),
805
+ )
806
+
807
+ def _has_digit_tokens(self, token_ids: torch.Tensor) -> bool:
808
+ """Check if input contains digit tokens."""
809
+ for tid in token_ids.tolist():
810
+ if tid in self.DIGIT_TOKENS:
811
+ return True
812
+ return False
813
+
814
+ def _find_op_position(self, token_ids: torch.Tensor) -> int:
815
+ """Find position of operator token, returns -1 if not found."""
816
+ tokens = token_ids.tolist()
817
+ for i, tid in enumerate(tokens):
818
+ if tid in self.ALL_OP_TOKENS:
819
+ return i
820
+ return -1
821
+
822
+ def _extract_from_digits(self, token_ids: torch.Tensor) -> tuple:
823
+ """
824
+ Extract values directly from digit tokens (hardcoded lookup).
825
+ Handles both symbol operators (' +') and word operators ('plus').
826
+ Returns (a_value, b_value, op_idx) or None if pattern not found.
827
+ """
828
+ tokens = token_ids.tolist()
829
+
830
+ op_pos = -1
831
+ op_idx = 0
832
+ for i, tid in enumerate(tokens):
833
+ if tid in self.ALL_OP_TOKENS:
834
+ op_pos = i
835
+ op_idx = self.ALL_OP_TOKENS[tid]
836
+ break
837
+
838
+ if op_pos == -1:
839
+ return None
840
+
841
+ a_digits = []
842
+ for i in range(op_pos):
843
+ if tokens[i] in self.DIGIT_TOKENS:
844
+ a_digits.append(tokens[i] - 32)
845
+
846
+ b_start = op_pos + 1
847
+ if b_start < len(tokens) and tokens[b_start] == 216:
848
+ b_start += 1
849
+
850
+ b_digits = []
851
+ for i in range(b_start, len(tokens)):
852
+ if tokens[i] in self.DIGIT_TOKENS:
853
+ b_digits.append(tokens[i] - 32)
854
+
855
+ if not a_digits or not b_digits:
856
+ return None
857
+
858
+ a_val = 0
859
+ for d in a_digits:
860
+ a_val = a_val * 10 + d
861
+
862
+ b_val = 0
863
+ for d in b_digits:
864
+ b_val = b_val * 10 + d
865
+
866
+ return min(a_val, 255), min(b_val, 255), op_idx
867
+
868
+ def _value_to_bits(self, value: int, device) -> torch.Tensor:
869
+ """Convert integer to 8-bit tensor."""
870
+ bits = torch.zeros(8, device=device)
871
+ for i in range(8):
872
+ bits[7 - i] = (value >> i) & 1
873
+ return bits
874
+
875
+ def _digits_to_value_and_bits(self, digit_logits: torch.Tensor, device) -> tuple:
876
+ """
877
+ Convert 3-digit logits to value and bits.
878
+ digit_logits: [30] (3 digits × 10 classes)
879
+ Returns: (value tensor, bits tensor [8])
880
+ """
881
+ logits = digit_logits.view(3, 10)
882
+ probs = torch.softmax(logits, dim=-1)
883
+
884
+ digit_values = torch.arange(10, device=device, dtype=torch.float32)
885
+ soft_digits = (probs * digit_values).sum(dim=-1)
886
+
887
+ hundreds = soft_digits[0]
888
+ tens = soft_digits[1]
889
+ ones = soft_digits[2]
890
+
891
+ value = hundreds * 100 + tens * 10 + ones
892
+ value = torch.clamp(value, 0, 255)
893
+
894
+ bits = []
895
+ for i in range(7, -1, -1):
896
+ threshold = 2 ** i
897
+ bit = torch.sigmoid((value - threshold + 0.5) * 10)
898
+ bits.append(bit)
899
+ value = value - bit * threshold
900
+ return hundreds * 100 + tens * 10 + ones, torch.stack(bits)
901
+
902
+ def forward(self, hidden: torch.Tensor, mask: torch.Tensor, token_ids: torch.Tensor = None):
903
+ """
904
+ Args:
905
+ hidden: [batch, seq_len, hidden_dim]
906
+ mask: [batch, seq_len]
907
+ token_ids: [batch, seq_len] - optional, enables digit lookup
908
+
909
+ Returns:
910
+ a_bits, b_bits, op_logits, a_values, b_values, used_lookup,
911
+ a_digit_logits, b_digit_logits
912
+ """
913
+ batch_size = hidden.shape[0]
914
+ device = hidden.device
915
+
916
+ a_bits_list = []
917
+ a_digit_logits_list = []
918
+ b_digit_logits_list = []
919
+ b_bits_list = []
920
+ op_logits_list = []
921
+ a_values_list = []
922
+ b_values_list = []
923
+ used_lookup_list = []
924
+
925
+ pooled = self.attention_pool(hidden, mask)
926
+
927
+ for i in range(batch_size):
928
+ lookup_result = None
929
+ if token_ids is not None:
930
+ seq_mask = mask[i].bool()
931
+ valid_len = seq_mask.sum().item()
932
+ start_pos = hidden.shape[1] - valid_len
933
+ valid_tokens = token_ids[i, start_pos:]
934
+
935
+ if self._has_digit_tokens(valid_tokens):
936
+ lookup_result = self._extract_from_digits(valid_tokens)
937
+
938
+ if lookup_result is not None:
939
+ a_val, b_val, op_idx = lookup_result
940
+ a_bits = self._value_to_bits(a_val, device)
941
+ b_bits = self._value_to_bits(b_val, device)
942
+ op_logits = torch.zeros(len(OPERATIONS), device=device)
943
+ op_logits[op_idx] = 10.0
944
+
945
+ a_bits_list.append(a_bits)
946
+ b_bits_list.append(b_bits)
947
+ op_logits_list.append(op_logits)
948
+ a_values_list.append(float(a_val))
949
+ b_values_list.append(float(b_val))
950
+ used_lookup_list.append(True)
951
+ a_digit_logits_list.append(None)
952
+ b_digit_logits_list.append(None)
953
+ else:
954
+ sample_hidden = hidden[i:i+1]
955
+ sample_mask = mask[i:i+1]
956
+
957
+ seq_mask = mask[i].bool()
958
+ valid_len = int(seq_mask.sum().item())
959
+ start_pos = hidden.shape[1] - valid_len
960
+ valid_tokens = token_ids[i, start_pos:] if token_ids is not None else None
961
+
962
+ op_pos = self._find_op_position(valid_tokens) if valid_tokens is not None else -1
963
+
964
+ if op_pos > 0 and op_pos < valid_len - 1:
965
+ a_end = start_pos + op_pos
966
+ b_start = start_pos + op_pos + 1
967
+
968
+ a_mask = torch.zeros_like(sample_mask)
969
+ a_mask[0, start_pos:a_end] = 1.0
970
+ b_mask = torch.zeros_like(sample_mask)
971
+ b_mask[0, b_start:] = sample_mask[0, b_start:]
972
+
973
+ a_pooled = self.a_pool(sample_hidden, a_mask)[0]
974
+ b_pooled = self.b_pool(sample_hidden, b_mask)[0]
975
+ else:
976
+ a_pooled = pooled[i]
977
+ b_pooled = pooled[i]
978
+
979
+ a_digit_logits = self.a_digit_pred(a_pooled)
980
+ b_digit_logits = self.b_digit_pred(b_pooled)
981
+ op_logits = self.op_predictor(pooled[i])
982
+
983
+ a_val, a_bits = self._digits_to_value_and_bits(a_digit_logits, device)
984
+ b_val, b_bits = self._digits_to_value_and_bits(b_digit_logits, device)
985
+
986
+ a_bits_list.append(a_bits)
987
+ b_bits_list.append(b_bits)
988
+ op_logits_list.append(op_logits)
989
+ a_values_list.append(a_val)
990
+ b_values_list.append(b_val)
991
+ used_lookup_list.append(False)
992
+ a_digit_logits_list.append(a_digit_logits)
993
+ b_digit_logits_list.append(b_digit_logits)
994
+
995
+ a_bits = torch.stack(a_bits_list)
996
+ b_bits = torch.stack(b_bits_list)
997
+ op_logits = torch.stack(op_logits_list)
998
+ a_values = torch.stack([v if isinstance(v, torch.Tensor) else torch.tensor(v, device=device) for v in a_values_list])
999
+ b_values = torch.stack([v if isinstance(v, torch.Tensor) else torch.tensor(v, device=device) for v in b_values_list])
1000
+ used_lookup = torch.tensor(used_lookup_list, device=device, dtype=torch.bool)
1001
+ valid_a_logits = [x for x in a_digit_logits_list if x is not None]
1002
+ valid_b_logits = [x for x in b_digit_logits_list if x is not None]
1003
+ a_digit_logits_out = torch.stack(valid_a_logits) if valid_a_logits else None
1004
+ b_digit_logits_out = torch.stack(valid_b_logits) if valid_b_logits else None
1005
+
1006
+ return a_bits, b_bits, op_logits, a_values, b_values, used_lookup, a_digit_logits_out, b_digit_logits_out
1007
+
1008
+ def _soft_value_to_bits(self, value: torch.Tensor, device) -> torch.Tensor:
1009
+ """Convert soft value (0-255) to 8-bit representation differentiably."""
1010
+ value = torch.clamp(value, 0, 255)
1011
+ bits = []
1012
+ remaining = value
1013
+ for i in range(7, -1, -1):
1014
+ threshold = 2 ** i
1015
+ bit = torch.sigmoid((remaining - threshold + 0.5) * 10)
1016
+ bits.append(bit)
1017
+ remaining = remaining - bit * threshold
1018
+ return torch.stack(bits)
1019
+
1020
+
1021
+ class ArithmeticModel(nn.Module):
1022
+ """
1023
+ LLM + extractor + frozen threshold circuits.
1024
+ Optionally unfreeze top N transformer layers with --unfreeze_layers.
1025
+ """
1026
+
1027
+ def __init__(self, device: str = 'cuda', unfreeze_layers: int = 0,
1028
+ extract_layer: int = -1, position_extract: bool = False,
1029
+ digit_pred: bool = False, positional_digit: bool = False,
1030
+ hybrid: bool = False):
1031
+ super().__init__()
1032
+ self.device = device
1033
+ self.unfreeze_layers = unfreeze_layers
1034
+ self.extract_layer = extract_layer
1035
+ self.position_extract = position_extract
1036
+ self.digit_pred = digit_pred
1037
+ self.positional_digit = positional_digit
1038
+ self.hybrid = hybrid
1039
+
1040
+ from transformers import AutoModelForCausalLM, AutoTokenizer
1041
+
1042
+ print("[1/4] Loading tokenizer...", flush=True)
1043
+ self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
1044
+ self.tokenizer.padding_side = 'left'
1045
+ if self.tokenizer.pad_token is None:
1046
+ self.tokenizer.pad_token = self.tokenizer.eos_token
1047
+ print(" Tokenizer loaded.", flush=True)
1048
+
1049
+ print("[2/4] Loading SmolLM2-360M...", flush=True)
1050
+ self.llm = AutoModelForCausalLM.from_pretrained(
1051
+ MODEL_ID,
1052
+ torch_dtype=torch.float16,
1053
+ device_map=device,
1054
+ output_hidden_states=True
1055
+ )
1056
+
1057
+ for param in self.llm.parameters():
1058
+ param.requires_grad = False
1059
+
1060
+ if unfreeze_layers > 0:
1061
+ num_layers = len(self.llm.model.layers)
1062
+ layers_to_unfreeze = list(range(num_layers - unfreeze_layers, num_layers))
1063
+ print(f" Unfreezing layers {layers_to_unfreeze}...", flush=True)
1064
+ for layer_idx in layers_to_unfreeze:
1065
+ for param in self.llm.model.layers[layer_idx].parameters():
1066
+ param.requires_grad = True
1067
+
1068
+ hidden_dim = self.llm.config.hidden_size
1069
+ llm_params = sum(p.numel() for p in self.llm.parameters())
1070
+ trainable_llm = sum(p.numel() for p in self.llm.parameters() if p.requires_grad)
1071
+ print(f" LLM loaded. Hidden dim: {hidden_dim}", flush=True)
1072
+ print(f" LLM params: {llm_params:,} total, {trainable_llm:,} trainable", flush=True)
1073
+
1074
+ print("[3/4] Loading threshold circuits...", flush=True)
1075
+ self.circuits = FrozenThresholdCircuits(device=device)
1076
+ print(f" Circuits loaded. {len(self.circuits.weights)} tensors", flush=True)
1077
+
1078
+ print("[4/4] Initializing extractor...", flush=True)
1079
+ if hybrid:
1080
+ print(" Using HYBRID extraction (digit lookup + word learning)", flush=True)
1081
+ self.extractor = HybridExtractor(
1082
+ hidden_dim=hidden_dim,
1083
+ intermediate_dim=256,
1084
+ num_heads=4
1085
+ ).to(device)
1086
+ elif positional_digit:
1087
+ print(" Using POSITIONAL DIGIT extraction (100% proven)", flush=True)
1088
+ self.extractor = PositionalDigitExtractor(
1089
+ hidden_dim=hidden_dim
1090
+ ).to(device)
1091
+ elif position_extract:
1092
+ print(" Using position-specific extraction", flush=True)
1093
+ self.extractor = PositionExtractor(
1094
+ hidden_dim=hidden_dim,
1095
+ intermediate_dim=256
1096
+ ).to(device)
1097
+ elif digit_pred:
1098
+ print(" Using digit-level prediction", flush=True)
1099
+ self.extractor = DigitExtractor(
1100
+ hidden_dim=hidden_dim,
1101
+ intermediate_dim=256,
1102
+ num_heads=4
1103
+ ).to(device)
1104
+ else:
1105
+ self.extractor = Extractor(
1106
+ hidden_dim=hidden_dim,
1107
+ intermediate_dim=256,
1108
+ num_heads=4
1109
+ ).to(device)
1110
+
1111
+ if extract_layer != -1:
1112
+ print(f" Extracting from layer {extract_layer}", flush=True)
1113
+
1114
+ trainable_ext = sum(p.numel() for p in self.extractor.parameters())
1115
+ total_trainable = trainable_llm + trainable_ext
1116
+ print(f" Extractor params: {trainable_ext:,}", flush=True)
1117
+ print(f" Total trainable: {total_trainable:,}", flush=True)
1118
+
1119
+ def get_hidden_states(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1120
+ """
1121
+ Get hidden states from specified layer.
1122
+
1123
+ Returns:
1124
+ hidden: [batch, seq_len, hidden_dim] hidden states
1125
+ mask: [batch, seq_len] attention mask
1126
+ token_ids: [batch, seq_len] input token IDs
1127
+ """
1128
+ inputs = self.tokenizer(
1129
+ texts,
1130
+ return_tensors='pt',
1131
+ padding=True,
1132
+ truncation=True,
1133
+ max_length=64
1134
+ ).to(self.device)
1135
+
1136
+ if self.unfreeze_layers > 0:
1137
+ outputs = self.llm(**inputs, output_hidden_states=True)
1138
+ else:
1139
+ with torch.no_grad():
1140
+ outputs = self.llm(**inputs, output_hidden_states=True)
1141
+
1142
+ hidden = outputs.hidden_states[self.extract_layer].float()
1143
+ mask = inputs.attention_mask.float()
1144
+ token_ids = inputs.input_ids
1145
+
1146
+ return hidden, mask, token_ids
1147
+
1148
+ def forward(self, texts: list[str]):
1149
+ """
1150
+ Full forward pass: text -> hidden states -> extractor -> circuits -> result
1151
+
1152
+ Returns:
1153
+ result_bits, a_bits, b_bits, op_logits
1154
+ If digit_pred: also returns a_digit_logits, b_digit_logits
1155
+ If position_extract or positional_digit: also returns op_indices
1156
+ If positional_digit: also returns a_values, b_values
1157
+ """
1158
+ hidden, mask, token_ids = self.get_hidden_states(texts)
1159
+
1160
+ if self.hybrid or self.positional_digit or self.position_extract:
1161
+ extractor_out = self.extractor(hidden, mask, token_ids)
1162
+ else:
1163
+ extractor_out = self.extractor(hidden, mask)
1164
+
1165
+ if self.hybrid:
1166
+ a_bits, b_bits, op_logits, a_values, b_values, used_lookup, a_digit_logits, b_digit_logits = extractor_out
1167
+ op_indices_from_tokens = None
1168
+ elif self.positional_digit:
1169
+ a_bits, b_bits, op_logits, op_indices_from_tokens, a_values, b_values, a_digit_logits, b_digit_logits = extractor_out
1170
+ used_lookup = None
1171
+ elif self.digit_pred:
1172
+ a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits = extractor_out
1173
+ op_indices_from_tokens = None
1174
+ a_values, b_values = None, None
1175
+ used_lookup = None
1176
+ elif self.position_extract:
1177
+ a_bits, b_bits, op_logits, op_indices_from_tokens = extractor_out
1178
+ a_digit_logits, b_digit_logits = None, None
1179
+ a_values, b_values = None, None
1180
+ used_lookup = None
1181
+ else:
1182
+ a_bits, b_bits, op_logits = extractor_out
1183
+ a_digit_logits, b_digit_logits = None, None
1184
+ op_indices_from_tokens = None
1185
+ a_values, b_values = None, None
1186
+ used_lookup = None
1187
+
1188
+ op_probs = torch.softmax(op_logits, dim=-1)
1189
+
1190
+ result_bits = self.circuits(a_bits, b_bits, op_probs)
1191
+
1192
+ if self.hybrid:
1193
+ return result_bits, a_bits, b_bits, op_logits, a_values, b_values, used_lookup, a_digit_logits, b_digit_logits
1194
+ if self.positional_digit:
1195
+ return result_bits, a_bits, b_bits, op_logits, op_indices_from_tokens, a_values, b_values, a_digit_logits, b_digit_logits
1196
+ if self.digit_pred:
1197
+ return result_bits, a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits
1198
+ if self.position_extract:
1199
+ return result_bits, a_bits, b_bits, op_logits, op_indices_from_tokens
1200
+ return result_bits, a_bits, b_bits, op_logits
1201
+
1202
+ def trainable_parameters(self):
1203
+ """Return all trainable parameters for optimizer."""
1204
+ params = list(self.extractor.parameters())
1205
+ if self.unfreeze_layers > 0:
1206
+ params += [p for p in self.llm.parameters() if p.requires_grad]
1207
+ return params
1208
+
1209
+
1210
+ if __name__ == "__main__":
1211
+ import sys
1212
+ sys.path.insert(0, '.')
1213
+ from fitness import generate_batch, compute_fitness, OPERATIONS
1214
+
1215
+ print("Testing model components...")
1216
+
1217
+ device = 'cuda'
1218
+ batch = generate_batch(32, device)
1219
+
1220
+ print("\n1. Testing DirectCircuitModel (should get ~100% fitness)...")
1221
+ direct_model = DirectCircuitModel(device=device)
1222
+
1223
+ def direct_fn(a, b, op):
1224
+ return direct_model(a, b, op)
1225
+
1226
+ fitness, details = compute_fitness(direct_fn, n_samples=2000, batch_size=128,
1227
+ device=device, return_details=True)
1228
+ print(f" Direct circuit fitness: {fitness:.4f}")
1229
+ for op in OPERATIONS:
1230
+ acc = details['by_op'][op]['accuracy']
1231
+ print(f" {op}: {acc:.4f}")
1232
+
1233
+ print("\n2. Testing ThresholdALU (trainable interface)...")
1234
+ model = ThresholdALU(device=device)
1235
+
1236
+ x = torch.cat([batch['a_bits'], batch['b_bits'], batch['op_onehot']], dim=-1)
1237
+ a_enc, b_enc = model.encoder(x)
1238
+ print(f" Encoder output shapes: a={a_enc.shape}, b={b_enc.shape}")
1239
+
1240
+ op_weights = model.router(x)
1241
+ print(f" Router output shape: {op_weights.shape}")
1242
+ print(f" Router output sample: {op_weights[0].tolist()}")
1243
+
1244
+ result = model(batch['a_bits'], batch['b_bits'], batch['op_onehot'])
1245
+ print(f" Full model output shape: {result.shape}")
1246
+
1247
+ print("\n3. Testing untrained ThresholdALU fitness...")
1248
+
1249
+ def model_fn(a, b, op):
1250
+ return model(a, b, op)
1251
+
1252
+ fitness = compute_fitness(model_fn, n_samples=1000, batch_size=128, device=device)
1253
+ print(f" Untrained model fitness: {fitness:.4f} (expected low)")
1254
+
1255
+ print("\n4. Counting parameters...")
1256
+ total = sum(p.numel() for p in model.parameters() if p.requires_grad)
1257
+ encoder_params = sum(p.numel() for p in model.encoder.parameters())
1258
+ router_params = sum(p.numel() for p in model.router.parameters())
1259
+ print(f" Encoder: {encoder_params:,}")
1260
+ print(f" Router: {router_params:,}")
1261
+ print(f" Total trainable: {total:,}")
1262
+
1263
+ print("\nDone.")
llm_integration/smollm2/SMOLLM2_ARCHITECTURE.md ADDED
@@ -0,0 +1,456 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SmolLM2-360M-Instruct Architecture Analysis
2
+
3
+ Technical reference document for the 8bit-threshold-computer LLM integration project.
4
+
5
+ **Model**: `HuggingFaceTB/SmolLM2-360M-Instruct`
6
+ **Architecture**: LlamaForCausalLM (Llama 2 variant)
7
+ **Tokenizer**: GPT2TokenizerFast
8
+ **Analysis Date**: 2026-01-21
9
+
10
+ ---
11
+
12
+ ## 1. Executive Summary
13
+
14
+ SmolLM2-360M-Instruct is a 362M parameter causal language model using the Llama architecture. Key characteristics relevant to our bit extraction task:
15
+
16
+ - **Hidden dimension: 960** (matches our extractor input requirement)
17
+ - **32 transformer layers** providing multiple extraction points
18
+ - **Digit-level tokenization** for numbers (each digit is a separate token)
19
+ - **Grouped Query Attention (GQA)** with 15 query heads and 5 KV heads
20
+
21
+ ---
22
+
23
+ ## 2. Architecture Census
24
+
25
+ ### 2.1 Core Parameters
26
+
27
+ | Parameter | Value |
28
+ |-----------|-------|
29
+ | Total Parameters | 361,821,120 (361.82M) |
30
+ | Vocabulary Size | 49,152 |
31
+ | Hidden Dimension | 960 |
32
+ | Intermediate Dimension (MLP) | 2,560 |
33
+ | Number of Layers | 32 |
34
+ | Number of Attention Heads | 15 |
35
+ | Number of KV Heads | 5 (Grouped Query Attention) |
36
+ | Head Dimension | 64 |
37
+ | Max Sequence Length | 8,192 |
38
+ | Activation Function | SiLU |
39
+ | Normalization | RMSNorm (eps=1e-05) |
40
+ | Position Encoding | RoPE (theta=100,000) |
41
+ | Word Embedding Tying | Yes (embed_tokens = lm_head) |
42
+
43
+ ### 2.2 Architecture Diagram
44
+
45
+ ```
46
+ Input Token IDs
47
+ |
48
+ v
49
+ +------------------+
50
+ | Embedding Layer | (49152, 960)
51
+ +------------------+
52
+ |
53
+ v
54
+ +------------------+
55
+ | LlamaDecoderLayer| x 32
56
+ | +-------------+ |
57
+ | | RMSNorm | |
58
+ | +-------------+ |
59
+ | | Self-Attn | | Q: (960, 960), K: (960, 320), V: (960, 320), O: (960, 960)
60
+ | +-------------+ |
61
+ | | Residual | |
62
+ | +-------------+ |
63
+ | | RMSNorm | |
64
+ | +-------------+ |
65
+ | | MLP (SwiGLU)| | gate: (960, 2560), up: (960, 2560), down: (2560, 960)
66
+ | +-------------+ |
67
+ | | Residual | |
68
+ +------------------+
69
+ |
70
+ v
71
+ +------------------+
72
+ | Final RMSNorm | (960,)
73
+ +------------------+
74
+ |
75
+ v
76
+ +------------------+
77
+ | LM Head | (960, 49152) - tied with embeddings
78
+ +------------------+
79
+ |
80
+ v
81
+ Logits (batch, seq, 49152)
82
+ ```
83
+
84
+ ### 2.3 Parameter Distribution
85
+
86
+ | Component | Parameters | Percentage |
87
+ |-----------|-----------|------------|
88
+ | Embedding | 47,185,920 | 13.04% |
89
+ | All Attention Layers | 78,643,200 | 21.74% |
90
+ | All MLP Layers | 235,929,600 | 65.19% |
91
+ | All Layer Norms | 61,440 | 0.02% |
92
+ | Final Norm | 960 | 0.00% |
93
+
94
+ Per-layer breakdown (each of 32 layers):
95
+ - Attention: 2,457,600 params (0.68%)
96
+ - MLP: 7,372,800 params (2.04%)
97
+ - Norms: 1,920 params (0.00%)
98
+
99
+ ---
100
+
101
+ ## 3. Weight Inventory
102
+
103
+ ### 3.1 Embedding and Output Layers
104
+
105
+ | Parameter Name | Shape | Elements | Notes |
106
+ |---------------|-------|----------|-------|
107
+ | `model.embed_tokens.weight` | (49152, 960) | 47,185,920 | Token embeddings |
108
+ | `model.norm.weight` | (960,) | 960 | Final layer norm |
109
+ | `lm_head.weight` | (49152, 960) | (tied) | Tied to embed_tokens |
110
+
111
+ ### 3.2 Single Transformer Layer Structure
112
+
113
+ Each of the 32 layers (`model.layers.{0-31}`) contains:
114
+
115
+ **Attention Block:**
116
+ | Parameter | Shape | Elements |
117
+ |-----------|-------|----------|
118
+ | `self_attn.q_proj.weight` | (960, 960) | 921,600 |
119
+ | `self_attn.k_proj.weight` | (320, 960) | 307,200 |
120
+ | `self_attn.v_proj.weight` | (320, 960) | 307,200 |
121
+ | `self_attn.o_proj.weight` | (960, 960) | 921,600 |
122
+ | **Attention Total** | | **2,457,600** |
123
+
124
+ **MLP Block (SwiGLU):**
125
+ | Parameter | Shape | Elements |
126
+ |-----------|-------|----------|
127
+ | `mlp.gate_proj.weight` | (2560, 960) | 2,457,600 |
128
+ | `mlp.up_proj.weight` | (2560, 960) | 2,457,600 |
129
+ | `mlp.down_proj.weight` | (960, 2560) | 2,457,600 |
130
+ | **MLP Total** | | **7,372,800** |
131
+
132
+ **Normalization:**
133
+ | Parameter | Shape | Elements |
134
+ |-----------|-------|----------|
135
+ | `input_layernorm.weight` | (960,) | 960 |
136
+ | `post_attention_layernorm.weight` | (960,) | 960 |
137
+ | **Norms Total** | | **1,920** |
138
+
139
+ **Layer Total: 9,832,320 parameters**
140
+
141
+ ### 3.3 Grouped Query Attention (GQA) Details
142
+
143
+ SmolLM2 uses GQA with a 3:1 ratio:
144
+ - 15 query heads (Q dimension: 960 = 15 x 64)
145
+ - 5 key-value heads (KV dimension: 320 = 5 x 64)
146
+ - Each KV head is shared by 3 query heads
147
+ - This reduces KV cache memory by ~67% vs standard MHA
148
+
149
+ ---
150
+
151
+ ## 4. Tokenization Analysis
152
+
153
+ ### 4.1 Arithmetic Expression Tokenization
154
+
155
+ Test input: `"47 + 86"`
156
+
157
+ | Position | Token ID | Token | Description |
158
+ |----------|----------|-------|-------------|
159
+ | 0 | 36 | `'4'` | First digit of operand A |
160
+ | 1 | 39 | `'7'` | Second digit of operand A |
161
+ | 2 | 1232 | `' +'` | Space + plus sign |
162
+ | 3 | 216 | `' '` | Trailing space |
163
+ | 4 | 40 | `'8'` | First digit of operand B |
164
+ | 5 | 38 | `'6'` | Second digit of operand B |
165
+
166
+ ### 4.2 Digit Token Mappings
167
+
168
+ | Digit | Token ID |
169
+ |-------|----------|
170
+ | 0 | 32 |
171
+ | 1 | 33 |
172
+ | 2 | 34 |
173
+ | 3 | 35 |
174
+ | 4 | 36 |
175
+ | 5 | 37 |
176
+ | 6 | 38 |
177
+ | 7 | 39 |
178
+ | 8 | 40 |
179
+ | 9 | 41 |
180
+
181
+ Key observations:
182
+ - **Digits are tokenized individually** (no multi-digit tokens like "47")
183
+ - Digit tokens are sequential: ID = 32 + digit_value
184
+ - Space-prefixed operators exist (e.g., `' +'` = 1232)
185
+ - `'='` has token ID 45
186
+
187
+ ### 4.3 Implications for Bit Extraction
188
+
189
+ The digit-by-digit tokenization means:
190
+ 1. For `"47 + 86"`, operand A spans positions [0,1] and operand B spans positions [4,5]
191
+ 2. The model must learn to:
192
+ - Recognize digit boundaries
193
+ - Compose multi-digit numbers from sequential tokens
194
+ - Perform arithmetic across token positions
195
+ 3. Hidden states at digit positions contain the numerical representation
196
+
197
+ ---
198
+
199
+ ## 5. Hidden State Analysis
200
+
201
+ ### 5.1 Hidden State Output Structure
202
+
203
+ When running with `output_hidden_states=True`:
204
+ - Returns **33 hidden states** (embedding + 32 layer outputs)
205
+ - Each has shape: `(batch_size, seq_len, hidden_dim)`
206
+ - For `"47 + 86"`: `(1, 6, 960)`
207
+
208
+ ### 5.2 Hidden State Statistics by Layer
209
+
210
+ | Layer | Mean | Std Dev | Min | Max |
211
+ |-------|------|---------|-----|-----|
212
+ | Embedding | -0.001 | 0.105 | -0.44 | 1.77 |
213
+ | Layer 0 | -0.127 | 2.55 | -80.8 | 19.0 |
214
+ | Layer 1 | -0.171 | 3.70 | -161 | 39.7 |
215
+ | Layer 2 | -0.151 | 3.67 | -102 | 61.4 |
216
+ | Layer 3 | -1.13 | 327 | -21,722 | 11,856 |
217
+ | Layer 4-12 | ~-1.3 | ~327 | ~-21,700 | ~11,800 |
218
+ | Layer 13-26 | ~-1.5 | ~337 | ~-22,400 | ~12,100 |
219
+ | Layer 27-30 | ~-1.8 | ~310 | ~-20,000 | ~11,800 |
220
+ | Layer 31 | 0.017 | 1.34 | -18.9 | 34.3 |
221
+
222
+ Key observations:
223
+ 1. **Dramatic variance explosion at Layer 3**: Std dev jumps from ~4 to ~327
224
+ 2. **Stable middle layers (4-26)**: Consistent statistics, suggesting numerical computation
225
+ 3. **Compression at final layer**: Std dev drops to 1.34 at Layer 31 (pre-softmax normalization)
226
+ 4. **Layer 31 is well-scaled** for downstream processing
227
+
228
+ ### 5.3 Extraction Point Candidates
229
+
230
+ | Layer Range | Characteristics | Suitability |
231
+ |-------------|-----------------|-------------|
232
+ | 0-2 (Early) | Low variance, close to embeddings | Poor - minimal computation |
233
+ | 3-12 (Early-Mid) | High variance, initial processing | Moderate - may contain raw numerical features |
234
+ | 13-26 (Middle) | Stable high variance | Good - computation in progress |
235
+ | 27-30 (Late) | Variance compression begins | Good - refined representations |
236
+ | 31 (Final) | Well-normalized output | Best - final representation before LM head |
237
+
238
+ ---
239
+
240
+ ## 6. Relevance to 8bit-Threshold-Computer Project
241
+
242
+ ### 6.1 Hidden Dimension Match
243
+
244
+ **The hidden dimension of 960 exactly matches our extractor input requirement.** This is fortuitous as it means:
245
+ - No projection layer needed to interface with our bit extractor
246
+ - Direct extraction from any layer's hidden states
247
+ - Full utilization of the model's representational capacity
248
+
249
+ ### 6.2 Recommended Extraction Strategy
250
+
251
+ ```python
252
+ def extract_hidden_state(model, tokenizer, expression, layer=-1):
253
+ """
254
+ Extract hidden state for bit extraction.
255
+
256
+ Args:
257
+ layer: Which layer to extract from (default: final layer)
258
+ -1 = Layer 31 (final, pre-LM-head)
259
+
260
+ Returns:
261
+ Tensor of shape (960,) for the last token position
262
+ """
263
+ inputs = tokenizer(expression, return_tensors="pt")
264
+ outputs = model(**inputs, output_hidden_states=True)
265
+
266
+ # hidden_states[0] = embedding, hidden_states[1] = layer 0, ...
267
+ # hidden_states[32] = layer 31 (final)
268
+ hidden = outputs.hidden_states[layer] # (1, seq_len, 960)
269
+
270
+ # Extract last token position for autoregressive prediction
271
+ return hidden[0, -1, :] # (960,)
272
+ ```
273
+
274
+ ### 6.3 Token Position Analysis
275
+
276
+ For arithmetic expressions like `"A + B"`:
277
+
278
+ ```
279
+ Tokens: [d1] [d2] [ +] [ ] [d3] [d4]
280
+ Positions: 0 1 2 3 4 5
281
+
282
+ Operand A: positions 0 to (plus_pos - 1)
283
+ Operator: position where ' +' token appears
284
+ Operand B: positions (plus_pos + 2) to end
285
+ ```
286
+
287
+ Strategy for operand extraction:
288
+ 1. Find the `' +'` token (ID 1232) position
289
+ 2. Collect hidden states at digit positions before it (operand A)
290
+ 3. Collect hidden states at digit positions after it (operand B)
291
+ 4. Consider pooling (mean, max) or concatenating digit hidden states
292
+
293
+ ### 6.4 Attention Pattern Utilization
294
+
295
+ With GQA (15 query heads, 5 KV heads), we can analyze attention patterns to:
296
+ 1. Identify which positions attend to operand digits
297
+ 2. Determine if the model explicitly links corresponding digit positions
298
+ 3. Find heads that specialize in numerical reasoning
299
+
300
+ ```python
301
+ def get_attention_weights(model, tokenizer, expression):
302
+ inputs = tokenizer(expression, return_tensors="pt")
303
+ outputs = model(**inputs, output_attentions=True)
304
+ # attentions: tuple of (batch, num_heads, seq_len, seq_len) per layer
305
+ return outputs.attentions
306
+ ```
307
+
308
+ ### 6.5 Extraction Interface Specification
309
+
310
+ For integration with the threshold computer:
311
+
312
+ ```python
313
+ class SmolLM2Extractor:
314
+ """Interface between SmolLM2 and threshold-based bit extraction."""
315
+
316
+ def __init__(self, model, tokenizer, extraction_layer=31):
317
+ self.model = model
318
+ self.tokenizer = tokenizer
319
+ self.layer = extraction_layer + 1 # +1 because index 0 is embedding
320
+
321
+ def get_hidden_state(self, text: str) -> torch.Tensor:
322
+ """
323
+ Returns: Tensor of shape (960,) ready for bit extractor
324
+ """
325
+ tokens = self.tokenizer(text, return_tensors="pt")
326
+ with torch.no_grad():
327
+ outputs = self.model(**tokens, output_hidden_states=True)
328
+ return outputs.hidden_states[self.layer][0, -1, :]
329
+
330
+ def get_all_position_states(self, text: str) -> torch.Tensor:
331
+ """
332
+ Returns: Tensor of shape (seq_len, 960) for all positions
333
+ """
334
+ tokens = self.tokenizer(text, return_tensors="pt")
335
+ with torch.no_grad():
336
+ outputs = self.model(**tokens, output_hidden_states=True)
337
+ return outputs.hidden_states[self.layer][0]
338
+ ```
339
+
340
+ ---
341
+
342
+ ## 7. Complete Weight Inventory Table
343
+
344
+ ### 7.1 All Named Parameters
345
+
346
+ ```
347
+ EMBEDDING (47,185,920 params - 13.04%)
348
+ model.embed_tokens.weight (49152, 960) 47,185,920
349
+
350
+ LAYER 0 (9,832,320 params - 2.72%)
351
+ Attention (2,457,600):
352
+ model.layers.0.self_attn.q_proj.weight (960, 960) 921,600
353
+ model.layers.0.self_attn.k_proj.weight (320, 960) 307,200
354
+ model.layers.0.self_attn.v_proj.weight (320, 960) 307,200
355
+ model.layers.0.self_attn.o_proj.weight (960, 960) 921,600
356
+ MLP (7,372,800):
357
+ model.layers.0.mlp.gate_proj.weight (2560, 960) 2,457,600
358
+ model.layers.0.mlp.up_proj.weight (2560, 960) 2,457,600
359
+ model.layers.0.mlp.down_proj.weight (960, 2560) 2,457,600
360
+ Norms (1,920):
361
+ model.layers.0.input_layernorm.weight (960,) 960
362
+ model.layers.0.post_attention_layernorm.weight (960,) 960
363
+
364
+ [Layers 1-31 follow identical structure, each with 9,832,320 params]
365
+
366
+ FINAL NORM (960 params - 0.00%)
367
+ model.norm.weight (960,) 960
368
+
369
+ LM HEAD (tied with embed_tokens)
370
+ lm_head.weight (49152, 960) [shared]
371
+ ```
372
+
373
+ ### 7.2 Summary by Component Type
374
+
375
+ | Component Type | Count | Params Each | Total Params |
376
+ |----------------|-------|-------------|--------------|
377
+ | Embedding | 1 | 47,185,920 | 47,185,920 |
378
+ | Q Projection | 32 | 921,600 | 29,491,200 |
379
+ | K Projection | 32 | 307,200 | 9,830,400 |
380
+ | V Projection | 32 | 307,200 | 9,830,400 |
381
+ | O Projection | 32 | 921,600 | 29,491,200 |
382
+ | Gate Projection | 32 | 2,457,600 | 78,643,200 |
383
+ | Up Projection | 32 | 2,457,600 | 78,643,200 |
384
+ | Down Projection | 32 | 2,457,600 | 78,643,200 |
385
+ | Input LayerNorm | 32 | 960 | 30,720 |
386
+ | Post-Attn LayerNorm | 32 | 960 | 30,720 |
387
+ | Final LayerNorm | 1 | 960 | 960 |
388
+ | **Total** | | | **361,821,120** |
389
+
390
+ ---
391
+
392
+ ## 8. Configuration Reference
393
+
394
+ Complete model configuration from HuggingFace:
395
+
396
+ ```python
397
+ {
398
+ "architectures": ["LlamaForCausalLM"],
399
+ "attention_bias": False,
400
+ "attention_dropout": 0.0,
401
+ "bos_token_id": 1,
402
+ "eos_token_id": 2,
403
+ "pad_token_id": 2,
404
+ "head_dim": 64,
405
+ "hidden_act": "silu",
406
+ "hidden_size": 960,
407
+ "initializer_range": 0.02,
408
+ "intermediate_size": 2560,
409
+ "max_position_embeddings": 8192,
410
+ "mlp_bias": False,
411
+ "model_type": "llama",
412
+ "num_attention_heads": 15,
413
+ "num_hidden_layers": 32,
414
+ "num_key_value_heads": 5,
415
+ "pretraining_tp": 1,
416
+ "rms_norm_eps": 1e-05,
417
+ "rope_interleaved": False,
418
+ "rope_theta": 100000,
419
+ "tie_word_embeddings": True,
420
+ "use_cache": True,
421
+ "vocab_size": 49152
422
+ }
423
+ ```
424
+
425
+ ---
426
+
427
+ ## 9. Appendix: PyTorch Model Structure
428
+
429
+ ```
430
+ LlamaForCausalLM(
431
+ (model): LlamaModel(
432
+ (embed_tokens): Embedding(49152, 960, padding_idx=2)
433
+ (layers): ModuleList(
434
+ (0-31): 32 x LlamaDecoderLayer(
435
+ (self_attn): LlamaAttention(
436
+ (q_proj): Linear(in_features=960, out_features=960, bias=False)
437
+ (k_proj): Linear(in_features=960, out_features=320, bias=False)
438
+ (v_proj): Linear(in_features=960, out_features=320, bias=False)
439
+ (o_proj): Linear(in_features=960, out_features=960, bias=False)
440
+ )
441
+ (mlp): LlamaMLP(
442
+ (gate_proj): Linear(in_features=960, out_features=2560, bias=False)
443
+ (up_proj): Linear(in_features=960, out_features=2560, bias=False)
444
+ (down_proj): Linear(in_features=2560, out_features=960, bias=False)
445
+ (act_fn): SiLUActivation()
446
+ )
447
+ (input_layernorm): LlamaRMSNorm((960,), eps=1e-05)
448
+ (post_attention_layernorm): LlamaRMSNorm((960,), eps=1e-05)
449
+ )
450
+ )
451
+ (norm): LlamaRMSNorm((960,), eps=1e-05)
452
+ (rotary_emb): LlamaRotaryEmbedding()
453
+ )
454
+ (lm_head): Linear(in_features=960, out_features=49152, bias=False)
455
+ )
456
+ ```
llm_integration/smollm2/analyze_smollm2.py ADDED
@@ -0,0 +1,232 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ SmolLM2-360M-Instruct Architecture Analysis
3
+ For 8bit-threshold-computer LLM Integration Project
4
+ """
5
+
6
+ import torch
7
+ from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
8
+ from collections import defaultdict
9
+ import json
10
+
11
+ def analyze_smollm2():
12
+ model_name = "HuggingFaceTB/SmolLM2-360M-Instruct"
13
+
14
+ print("=" * 80)
15
+ print("SmolLM2-360M-Instruct Architecture Analysis")
16
+ print("=" * 80)
17
+
18
+ # Load config first
19
+ print("\n[1] Loading model configuration...")
20
+ config = AutoConfig.from_pretrained(model_name)
21
+ print(f"Config loaded: {type(config).__name__}")
22
+
23
+ # Load tokenizer
24
+ print("\n[2] Loading tokenizer...")
25
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
26
+ print(f"Tokenizer loaded: {type(tokenizer).__name__}")
27
+
28
+ # Load model with hidden states output
29
+ print("\n[3] Loading model with output_hidden_states=True...")
30
+ model = AutoModelForCausalLM.from_pretrained(
31
+ model_name,
32
+ output_hidden_states=True,
33
+ torch_dtype=torch.float32
34
+ )
35
+ model.eval()
36
+ print(f"Model loaded: {type(model).__name__}")
37
+
38
+ # ========================================================================
39
+ # ARCHITECTURE CENSUS
40
+ # ========================================================================
41
+ print("\n" + "=" * 80)
42
+ print("ARCHITECTURE CENSUS")
43
+ print("=" * 80)
44
+
45
+ print("\n--- Model Configuration ---")
46
+ config_dict = config.to_dict()
47
+ for key, value in sorted(config_dict.items()):
48
+ print(f" {key}: {value}")
49
+
50
+ print("\n--- Key Architecture Parameters ---")
51
+ print(f" Model type: {config.model_type}")
52
+ print(f" Vocabulary size: {config.vocab_size}")
53
+ print(f" Hidden size: {config.hidden_size}")
54
+ print(f" Intermediate size: {config.intermediate_size}")
55
+ print(f" Number of hidden layers: {config.num_hidden_layers}")
56
+ print(f" Number of attention heads: {config.num_attention_heads}")
57
+ print(f" Number of KV heads: {getattr(config, 'num_key_value_heads', config.num_attention_heads)}")
58
+ print(f" Head dimension: {config.hidden_size // config.num_attention_heads}")
59
+ print(f" Max position embeddings: {config.max_position_embeddings}")
60
+ print(f" RMS norm epsilon: {getattr(config, 'rms_norm_eps', 'N/A')}")
61
+ print(f" Rope theta: {getattr(config, 'rope_theta', 'N/A')}")
62
+ print(f" Tie word embeddings: {getattr(config, 'tie_word_embeddings', 'N/A')}")
63
+
64
+ # ========================================================================
65
+ # WEIGHT INVENTORY
66
+ # ========================================================================
67
+ print("\n" + "=" * 80)
68
+ print("WEIGHT INVENTORY")
69
+ print("=" * 80)
70
+
71
+ total_params = 0
72
+ param_groups = defaultdict(list)
73
+
74
+ for name, param in model.named_parameters():
75
+ total_params += param.numel()
76
+
77
+ # Group by component
78
+ if "embed_tokens" in name:
79
+ group = "Embedding"
80
+ elif "lm_head" in name:
81
+ group = "LM Head"
82
+ elif "norm" in name and "layers" not in name:
83
+ group = "Final Norm"
84
+ elif "layers" in name:
85
+ layer_num = name.split(".")[2]
86
+ if "self_attn" in name:
87
+ group = f"Layer {layer_num} - Attention"
88
+ elif "mlp" in name:
89
+ group = f"Layer {layer_num} - MLP"
90
+ elif "norm" in name:
91
+ group = f"Layer {layer_num} - Norms"
92
+ else:
93
+ group = f"Layer {layer_num} - Other"
94
+ else:
95
+ group = "Other"
96
+
97
+ param_groups[group].append({
98
+ "name": name,
99
+ "shape": tuple(param.shape),
100
+ "numel": param.numel(),
101
+ "dtype": str(param.dtype)
102
+ })
103
+
104
+ print(f"\n--- Total Parameters: {total_params:,} ---")
105
+ print(f" ({total_params / 1e6:.2f}M parameters)")
106
+
107
+ # Print by group
108
+ for group_name in sorted(param_groups.keys()):
109
+ params = param_groups[group_name]
110
+ group_total = sum(p["numel"] for p in params)
111
+ print(f"\n### {group_name} ({group_total:,} params, {group_total/total_params*100:.2f}%)")
112
+ for p in params:
113
+ print(f" {p['name']}")
114
+ print(f" Shape: {p['shape']}, Elements: {p['numel']:,}, Dtype: {p['dtype']}")
115
+
116
+ # ========================================================================
117
+ # TOKENIZATION ANALYSIS
118
+ # ========================================================================
119
+ print("\n" + "=" * 80)
120
+ print("TOKENIZATION ANALYSIS")
121
+ print("=" * 80)
122
+
123
+ test_input = "47 + 86"
124
+ print(f"\n--- Test Input: '{test_input}' ---")
125
+
126
+ tokens = tokenizer(test_input, return_tensors="pt")
127
+ input_ids = tokens["input_ids"][0]
128
+
129
+ print(f"\nInput IDs: {input_ids.tolist()}")
130
+ print(f"Number of tokens: {len(input_ids)}")
131
+
132
+ print("\nToken breakdown:")
133
+ for i, token_id in enumerate(input_ids):
134
+ token_str = tokenizer.decode([token_id])
135
+ print(f" Position {i}: ID={token_id.item():5d}, Token='{token_str}'")
136
+
137
+ # Additional tokenization tests
138
+ print("\n--- Additional Tokenization Tests ---")
139
+ test_cases = ["0", "1", "47", "86", "133", " + ", "="]
140
+ for tc in test_cases:
141
+ ids = tokenizer.encode(tc, add_special_tokens=False)
142
+ decoded = [tokenizer.decode([i]) for i in ids]
143
+ print(f" '{tc}' -> IDs: {ids}, Tokens: {decoded}")
144
+
145
+ # ========================================================================
146
+ # HIDDEN STATE ANALYSIS
147
+ # ========================================================================
148
+ print("\n" + "=" * 80)
149
+ print("HIDDEN STATE ANALYSIS")
150
+ print("=" * 80)
151
+
152
+ print(f"\n--- Running inference on '{test_input}' ---")
153
+
154
+ with torch.no_grad():
155
+ outputs = model(**tokens)
156
+
157
+ hidden_states = outputs.hidden_states
158
+ print(f"\nNumber of hidden state outputs: {len(hidden_states)}")
159
+ print("(This includes embedding output + each layer's output)")
160
+
161
+ print("\nHidden state shapes at each layer:")
162
+ for i, hs in enumerate(hidden_states):
163
+ layer_name = "Embedding" if i == 0 else f"Layer {i-1}"
164
+ print(f" {layer_name}: {tuple(hs.shape)}")
165
+ if i == 0:
166
+ print(f" (batch_size=1, seq_len={hs.shape[1]}, hidden_dim={hs.shape[2]})")
167
+
168
+ # Analyze hidden state statistics at different layers
169
+ print("\n--- Hidden State Statistics (per layer) ---")
170
+ for i, hs in enumerate(hidden_states):
171
+ layer_name = "Embedding" if i == 0 else f"Layer {i-1}"
172
+ hs_flat = hs.view(-1)
173
+ print(f" {layer_name}:")
174
+ print(f" Mean: {hs_flat.mean().item():.6f}")
175
+ print(f" Std: {hs_flat.std().item():.6f}")
176
+ print(f" Min: {hs_flat.min().item():.6f}")
177
+ print(f" Max: {hs_flat.max().item():.6f}")
178
+
179
+ # ========================================================================
180
+ # MODEL STRUCTURE DEEP DIVE
181
+ # ========================================================================
182
+ print("\n" + "=" * 80)
183
+ print("MODEL STRUCTURE DEEP DIVE")
184
+ print("=" * 80)
185
+
186
+ print("\n--- Model Architecture String ---")
187
+ print(model)
188
+
189
+ # ========================================================================
190
+ # SUMMARY DATA FOR REPORT
191
+ # ========================================================================
192
+ summary = {
193
+ "model_name": model_name,
194
+ "total_params": total_params,
195
+ "config": {
196
+ "vocab_size": config.vocab_size,
197
+ "hidden_size": config.hidden_size,
198
+ "intermediate_size": config.intermediate_size,
199
+ "num_hidden_layers": config.num_hidden_layers,
200
+ "num_attention_heads": config.num_attention_heads,
201
+ "num_kv_heads": getattr(config, 'num_key_value_heads', config.num_attention_heads),
202
+ "head_dim": config.hidden_size // config.num_attention_heads,
203
+ "max_position_embeddings": config.max_position_embeddings,
204
+ "rms_norm_eps": getattr(config, 'rms_norm_eps', None),
205
+ "rope_theta": getattr(config, 'rope_theta', None),
206
+ "tie_word_embeddings": getattr(config, 'tie_word_embeddings', None),
207
+ },
208
+ "tokenization": {
209
+ "test_input": test_input,
210
+ "token_ids": input_ids.tolist(),
211
+ "num_tokens": len(input_ids),
212
+ "tokens": [tokenizer.decode([tid]) for tid in input_ids]
213
+ },
214
+ "hidden_states": {
215
+ "num_outputs": len(hidden_states),
216
+ "shape": list(hidden_states[0].shape)
217
+ },
218
+ "param_groups": {k: {"count": len(v), "total": sum(p["numel"] for p in v)} for k, v in param_groups.items()}
219
+ }
220
+
221
+ # Save summary as JSON for report generation
222
+ with open("D:/8bit-threshold-computer/llm_integration/smollm2_analysis.json", "w") as f:
223
+ json.dump(summary, f, indent=2)
224
+
225
+ print("\n" + "=" * 80)
226
+ print("Analysis complete. Summary saved to smollm2_analysis.json")
227
+ print("=" * 80)
228
+
229
+ return summary, model, tokenizer, hidden_states, param_groups
230
+
231
+ if __name__ == "__main__":
232
+ summary, model, tokenizer, hidden_states, param_groups = analyze_smollm2()
llm_integration/smollm2/smollm2_analysis.json ADDED
@@ -0,0 +1,439 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "model_name": "HuggingFaceTB/SmolLM2-360M-Instruct",
3
+ "total_params": 361821120,
4
+ "config": {
5
+ "vocab_size": 49152,
6
+ "hidden_size": 960,
7
+ "intermediate_size": 2560,
8
+ "num_hidden_layers": 32,
9
+ "num_attention_heads": 15,
10
+ "num_kv_heads": 5,
11
+ "head_dim": 64,
12
+ "max_position_embeddings": 8192,
13
+ "rms_norm_eps": 1e-05,
14
+ "rope_theta": 100000,
15
+ "tie_word_embeddings": true
16
+ },
17
+ "tokenization": {
18
+ "test_input": "47 + 86",
19
+ "token_ids": [
20
+ 36,
21
+ 39,
22
+ 1232,
23
+ 216,
24
+ 40,
25
+ 38
26
+ ],
27
+ "num_tokens": 6,
28
+ "tokens": [
29
+ "4",
30
+ "7",
31
+ " +",
32
+ " ",
33
+ "8",
34
+ "6"
35
+ ]
36
+ },
37
+ "hidden_states": {
38
+ "num_outputs": 33,
39
+ "shape": [
40
+ 1,
41
+ 6,
42
+ 960
43
+ ]
44
+ },
45
+ "param_groups": {
46
+ "Embedding": {
47
+ "count": 1,
48
+ "total": 47185920
49
+ },
50
+ "Layer 0 - Attention": {
51
+ "count": 4,
52
+ "total": 2457600
53
+ },
54
+ "Layer 0 - MLP": {
55
+ "count": 3,
56
+ "total": 7372800
57
+ },
58
+ "Layer 0 - Norms": {
59
+ "count": 2,
60
+ "total": 1920
61
+ },
62
+ "Layer 1 - Attention": {
63
+ "count": 4,
64
+ "total": 2457600
65
+ },
66
+ "Layer 1 - MLP": {
67
+ "count": 3,
68
+ "total": 7372800
69
+ },
70
+ "Layer 1 - Norms": {
71
+ "count": 2,
72
+ "total": 1920
73
+ },
74
+ "Layer 2 - Attention": {
75
+ "count": 4,
76
+ "total": 2457600
77
+ },
78
+ "Layer 2 - MLP": {
79
+ "count": 3,
80
+ "total": 7372800
81
+ },
82
+ "Layer 2 - Norms": {
83
+ "count": 2,
84
+ "total": 1920
85
+ },
86
+ "Layer 3 - Attention": {
87
+ "count": 4,
88
+ "total": 2457600
89
+ },
90
+ "Layer 3 - MLP": {
91
+ "count": 3,
92
+ "total": 7372800
93
+ },
94
+ "Layer 3 - Norms": {
95
+ "count": 2,
96
+ "total": 1920
97
+ },
98
+ "Layer 4 - Attention": {
99
+ "count": 4,
100
+ "total": 2457600
101
+ },
102
+ "Layer 4 - MLP": {
103
+ "count": 3,
104
+ "total": 7372800
105
+ },
106
+ "Layer 4 - Norms": {
107
+ "count": 2,
108
+ "total": 1920
109
+ },
110
+ "Layer 5 - Attention": {
111
+ "count": 4,
112
+ "total": 2457600
113
+ },
114
+ "Layer 5 - MLP": {
115
+ "count": 3,
116
+ "total": 7372800
117
+ },
118
+ "Layer 5 - Norms": {
119
+ "count": 2,
120
+ "total": 1920
121
+ },
122
+ "Layer 6 - Attention": {
123
+ "count": 4,
124
+ "total": 2457600
125
+ },
126
+ "Layer 6 - MLP": {
127
+ "count": 3,
128
+ "total": 7372800
129
+ },
130
+ "Layer 6 - Norms": {
131
+ "count": 2,
132
+ "total": 1920
133
+ },
134
+ "Layer 7 - Attention": {
135
+ "count": 4,
136
+ "total": 2457600
137
+ },
138
+ "Layer 7 - MLP": {
139
+ "count": 3,
140
+ "total": 7372800
141
+ },
142
+ "Layer 7 - Norms": {
143
+ "count": 2,
144
+ "total": 1920
145
+ },
146
+ "Layer 8 - Attention": {
147
+ "count": 4,
148
+ "total": 2457600
149
+ },
150
+ "Layer 8 - MLP": {
151
+ "count": 3,
152
+ "total": 7372800
153
+ },
154
+ "Layer 8 - Norms": {
155
+ "count": 2,
156
+ "total": 1920
157
+ },
158
+ "Layer 9 - Attention": {
159
+ "count": 4,
160
+ "total": 2457600
161
+ },
162
+ "Layer 9 - MLP": {
163
+ "count": 3,
164
+ "total": 7372800
165
+ },
166
+ "Layer 9 - Norms": {
167
+ "count": 2,
168
+ "total": 1920
169
+ },
170
+ "Layer 10 - Attention": {
171
+ "count": 4,
172
+ "total": 2457600
173
+ },
174
+ "Layer 10 - MLP": {
175
+ "count": 3,
176
+ "total": 7372800
177
+ },
178
+ "Layer 10 - Norms": {
179
+ "count": 2,
180
+ "total": 1920
181
+ },
182
+ "Layer 11 - Attention": {
183
+ "count": 4,
184
+ "total": 2457600
185
+ },
186
+ "Layer 11 - MLP": {
187
+ "count": 3,
188
+ "total": 7372800
189
+ },
190
+ "Layer 11 - Norms": {
191
+ "count": 2,
192
+ "total": 1920
193
+ },
194
+ "Layer 12 - Attention": {
195
+ "count": 4,
196
+ "total": 2457600
197
+ },
198
+ "Layer 12 - MLP": {
199
+ "count": 3,
200
+ "total": 7372800
201
+ },
202
+ "Layer 12 - Norms": {
203
+ "count": 2,
204
+ "total": 1920
205
+ },
206
+ "Layer 13 - Attention": {
207
+ "count": 4,
208
+ "total": 2457600
209
+ },
210
+ "Layer 13 - MLP": {
211
+ "count": 3,
212
+ "total": 7372800
213
+ },
214
+ "Layer 13 - Norms": {
215
+ "count": 2,
216
+ "total": 1920
217
+ },
218
+ "Layer 14 - Attention": {
219
+ "count": 4,
220
+ "total": 2457600
221
+ },
222
+ "Layer 14 - MLP": {
223
+ "count": 3,
224
+ "total": 7372800
225
+ },
226
+ "Layer 14 - Norms": {
227
+ "count": 2,
228
+ "total": 1920
229
+ },
230
+ "Layer 15 - Attention": {
231
+ "count": 4,
232
+ "total": 2457600
233
+ },
234
+ "Layer 15 - MLP": {
235
+ "count": 3,
236
+ "total": 7372800
237
+ },
238
+ "Layer 15 - Norms": {
239
+ "count": 2,
240
+ "total": 1920
241
+ },
242
+ "Layer 16 - Attention": {
243
+ "count": 4,
244
+ "total": 2457600
245
+ },
246
+ "Layer 16 - MLP": {
247
+ "count": 3,
248
+ "total": 7372800
249
+ },
250
+ "Layer 16 - Norms": {
251
+ "count": 2,
252
+ "total": 1920
253
+ },
254
+ "Layer 17 - Attention": {
255
+ "count": 4,
256
+ "total": 2457600
257
+ },
258
+ "Layer 17 - MLP": {
259
+ "count": 3,
260
+ "total": 7372800
261
+ },
262
+ "Layer 17 - Norms": {
263
+ "count": 2,
264
+ "total": 1920
265
+ },
266
+ "Layer 18 - Attention": {
267
+ "count": 4,
268
+ "total": 2457600
269
+ },
270
+ "Layer 18 - MLP": {
271
+ "count": 3,
272
+ "total": 7372800
273
+ },
274
+ "Layer 18 - Norms": {
275
+ "count": 2,
276
+ "total": 1920
277
+ },
278
+ "Layer 19 - Attention": {
279
+ "count": 4,
280
+ "total": 2457600
281
+ },
282
+ "Layer 19 - MLP": {
283
+ "count": 3,
284
+ "total": 7372800
285
+ },
286
+ "Layer 19 - Norms": {
287
+ "count": 2,
288
+ "total": 1920
289
+ },
290
+ "Layer 20 - Attention": {
291
+ "count": 4,
292
+ "total": 2457600
293
+ },
294
+ "Layer 20 - MLP": {
295
+ "count": 3,
296
+ "total": 7372800
297
+ },
298
+ "Layer 20 - Norms": {
299
+ "count": 2,
300
+ "total": 1920
301
+ },
302
+ "Layer 21 - Attention": {
303
+ "count": 4,
304
+ "total": 2457600
305
+ },
306
+ "Layer 21 - MLP": {
307
+ "count": 3,
308
+ "total": 7372800
309
+ },
310
+ "Layer 21 - Norms": {
311
+ "count": 2,
312
+ "total": 1920
313
+ },
314
+ "Layer 22 - Attention": {
315
+ "count": 4,
316
+ "total": 2457600
317
+ },
318
+ "Layer 22 - MLP": {
319
+ "count": 3,
320
+ "total": 7372800
321
+ },
322
+ "Layer 22 - Norms": {
323
+ "count": 2,
324
+ "total": 1920
325
+ },
326
+ "Layer 23 - Attention": {
327
+ "count": 4,
328
+ "total": 2457600
329
+ },
330
+ "Layer 23 - MLP": {
331
+ "count": 3,
332
+ "total": 7372800
333
+ },
334
+ "Layer 23 - Norms": {
335
+ "count": 2,
336
+ "total": 1920
337
+ },
338
+ "Layer 24 - Attention": {
339
+ "count": 4,
340
+ "total": 2457600
341
+ },
342
+ "Layer 24 - MLP": {
343
+ "count": 3,
344
+ "total": 7372800
345
+ },
346
+ "Layer 24 - Norms": {
347
+ "count": 2,
348
+ "total": 1920
349
+ },
350
+ "Layer 25 - Attention": {
351
+ "count": 4,
352
+ "total": 2457600
353
+ },
354
+ "Layer 25 - MLP": {
355
+ "count": 3,
356
+ "total": 7372800
357
+ },
358
+ "Layer 25 - Norms": {
359
+ "count": 2,
360
+ "total": 1920
361
+ },
362
+ "Layer 26 - Attention": {
363
+ "count": 4,
364
+ "total": 2457600
365
+ },
366
+ "Layer 26 - MLP": {
367
+ "count": 3,
368
+ "total": 7372800
369
+ },
370
+ "Layer 26 - Norms": {
371
+ "count": 2,
372
+ "total": 1920
373
+ },
374
+ "Layer 27 - Attention": {
375
+ "count": 4,
376
+ "total": 2457600
377
+ },
378
+ "Layer 27 - MLP": {
379
+ "count": 3,
380
+ "total": 7372800
381
+ },
382
+ "Layer 27 - Norms": {
383
+ "count": 2,
384
+ "total": 1920
385
+ },
386
+ "Layer 28 - Attention": {
387
+ "count": 4,
388
+ "total": 2457600
389
+ },
390
+ "Layer 28 - MLP": {
391
+ "count": 3,
392
+ "total": 7372800
393
+ },
394
+ "Layer 28 - Norms": {
395
+ "count": 2,
396
+ "total": 1920
397
+ },
398
+ "Layer 29 - Attention": {
399
+ "count": 4,
400
+ "total": 2457600
401
+ },
402
+ "Layer 29 - MLP": {
403
+ "count": 3,
404
+ "total": 7372800
405
+ },
406
+ "Layer 29 - Norms": {
407
+ "count": 2,
408
+ "total": 1920
409
+ },
410
+ "Layer 30 - Attention": {
411
+ "count": 4,
412
+ "total": 2457600
413
+ },
414
+ "Layer 30 - MLP": {
415
+ "count": 3,
416
+ "total": 7372800
417
+ },
418
+ "Layer 30 - Norms": {
419
+ "count": 2,
420
+ "total": 1920
421
+ },
422
+ "Layer 31 - Attention": {
423
+ "count": 4,
424
+ "total": 2457600
425
+ },
426
+ "Layer 31 - MLP": {
427
+ "count": 3,
428
+ "total": 7372800
429
+ },
430
+ "Layer 31 - Norms": {
431
+ "count": 2,
432
+ "total": 1920
433
+ },
434
+ "Final Norm": {
435
+ "count": 1,
436
+ "total": 960
437
+ }
438
+ }
439
+ }
llm_integration/train.py ADDED
@@ -0,0 +1,955 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Unified training script for threshold circuit LLM integration.
3
+
4
+ Modes:
5
+ --mode router : Train only OpRouter with ground truth bits (sanity check)
6
+ --mode interface : Train BitEncoder + OpRouter with ground truth bits (sanity check)
7
+ --mode llm : Train extractor with LLM hidden states (the real training)
8
+
9
+ LLM mode options:
10
+ --unfreeze_layers N : Unfreeze top N transformer layers (default 0 = fully frozen)
11
+
12
+ Hardware Profile (NVIDIA RTX 6000 Ada 48GB):
13
+ VRAM Scaling (unfreeze_layers=4):
14
+ batch_size | VRAM | %
15
+ -----------+---------+------
16
+ 512 | 5,784 | 11.8%
17
+ 1,024 | 7,384 | 15.0%
18
+ 4,096 | 16,534 | 33.6%
19
+ 13,000 | 39,000 | 79.4% <-- recommended for 80% target
20
+
21
+ Examples:
22
+ python train.py --mode llm --epochs 100 --batch_size 256
23
+ python train.py --mode llm --epochs 100 --batch_size 4096 --unfreeze_layers 4
24
+ """
25
+
26
+ import torch
27
+ import torch.nn as nn
28
+ import torch.optim as optim
29
+ import time
30
+ import argparse
31
+ import random
32
+
33
+ from model import (
34
+ ThresholdALU, DirectCircuitModel, OpRouter,
35
+ ArithmeticModel, OPERATIONS, OP_SYMBOLS
36
+ )
37
+ from circuits import FrozenThresholdCircuits
38
+ from fitness import generate_batch, compute_fitness, compute_loss
39
+
40
+ DEVICE = 'cuda'
41
+
42
+ ONES = ['zero', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine',
43
+ 'ten', 'eleven', 'twelve', 'thirteen', 'fourteen', 'fifteen', 'sixteen',
44
+ 'seventeen', 'eighteen', 'nineteen']
45
+ TENS = ['', '', 'twenty', 'thirty', 'forty', 'fifty', 'sixty', 'seventy', 'eighty', 'ninety']
46
+
47
+ def int_to_words(n: int) -> str:
48
+ """Convert integer 0-255 to English words."""
49
+ if n < 0 or n > 255:
50
+ return str(n)
51
+ if n < 20:
52
+ return ONES[n]
53
+ if n < 100:
54
+ if n % 10 == 0:
55
+ return TENS[n // 10]
56
+ return f"{TENS[n // 10]} {ONES[n % 10]}"
57
+ if n % 100 == 0:
58
+ return f"{ONES[n // 100]} hundred"
59
+ if n % 100 < 20:
60
+ return f"{ONES[n // 100]} hundred {ONES[n % 100]}"
61
+ if n % 10 == 0:
62
+ return f"{ONES[n // 100]} hundred {TENS[(n % 100) // 10]}"
63
+ return f"{ONES[n // 100]} hundred {TENS[(n % 100) // 10]} {ONES[n % 10]}"
64
+
65
+
66
+ def int_to_bits(val: int, device: str = 'cuda') -> torch.Tensor:
67
+ bits = torch.zeros(8, device=device)
68
+ for i in range(8):
69
+ bits[7-i] = (val >> i) & 1
70
+ return bits
71
+
72
+
73
+ def bits_to_int(bits: torch.Tensor) -> int:
74
+ val = 0
75
+ for i in range(8):
76
+ if bits[i].item() > 0.5:
77
+ val += 1 << (7-i)
78
+ return val
79
+
80
+
81
+ NL_TEMPLATES = {
82
+ 'add': [
83
+ "What is {a} plus {b}?",
84
+ "Calculate {a} + {b}",
85
+ "Add {a} and {b}",
86
+ "What's the sum of {a} and {b}?",
87
+ "If I have {a} and get {b} more, how many total?",
88
+ "{a} + {b} = ?",
89
+ "Compute {a} plus {b}",
90
+ ],
91
+ 'sub': [
92
+ "What is {a} minus {b}?",
93
+ "Calculate {a} - {b}",
94
+ "Subtract {b} from {a}",
95
+ "What's {a} take away {b}?",
96
+ "If I have {a} and lose {b}, how many left?",
97
+ "{a} - {b} = ?",
98
+ "Compute {a} minus {b}",
99
+ ],
100
+ 'mul': [
101
+ "What is {a} times {b}?",
102
+ "Calculate {a} * {b}",
103
+ "Multiply {a} by {b}",
104
+ "What's {a} multiplied by {b}?",
105
+ "{a} * {b} = ?",
106
+ "Compute {a} times {b}",
107
+ "What is the product of {a} and {b}?",
108
+ ],
109
+ 'gt': [
110
+ "Is {a} greater than {b}?",
111
+ "Is {a} > {b}?",
112
+ "Check if {a} is larger than {b}",
113
+ "Compare: is {a} more than {b}?",
114
+ "{a} > {b}?",
115
+ ],
116
+ 'lt': [
117
+ "Is {a} less than {b}?",
118
+ "Is {a} < {b}?",
119
+ "Check if {a} is smaller than {b}",
120
+ "Compare: is {a} fewer than {b}?",
121
+ "{a} < {b}?",
122
+ ],
123
+ 'eq': [
124
+ "Is {a} equal to {b}?",
125
+ "Is {a} == {b}?",
126
+ "Does {a} equal {b}?",
127
+ "Check if {a} equals {b}",
128
+ "Are {a} and {b} the same?",
129
+ ],
130
+ }
131
+
132
+
133
+ def generate_problem(max_val: int = 255):
134
+ """
135
+ Generate a random arithmetic problem for LLM training.
136
+ Randomly mixes digit and word formats.
137
+ """
138
+ a = random.randint(0, max_val)
139
+ b = random.randint(0, max_val)
140
+ op = random.choice(OPERATIONS)
141
+
142
+ fmt = random.choice(['digits', 'words', 'nl_digits', 'nl_words'])
143
+
144
+ if fmt == 'digits':
145
+ sym = OP_SYMBOLS[op]
146
+ text = f"{a} {sym} {b}"
147
+ elif fmt == 'words':
148
+ a_word = int_to_words(a)
149
+ b_word = int_to_words(b)
150
+ op_word = {'add': 'plus', 'sub': 'minus', 'mul': 'times',
151
+ 'gt': 'greater than', 'lt': 'less than', 'eq': 'equals'}[op]
152
+ text = f"{a_word} {op_word} {b_word}"
153
+ elif fmt == 'nl_digits':
154
+ template = random.choice(NL_TEMPLATES[op])
155
+ text = template.format(a=a, b=b)
156
+ else:
157
+ template = random.choice(NL_TEMPLATES[op])
158
+ text = template.format(a=int_to_words(a), b=int_to_words(b))
159
+
160
+ if op == 'add':
161
+ result = (a + b) & 0xFF
162
+ elif op == 'sub':
163
+ result = (a - b) & 0xFF
164
+ elif op == 'mul':
165
+ result = (a * b) & 0xFF
166
+ elif op == 'gt':
167
+ result = 1 if a > b else 0
168
+ elif op == 'lt':
169
+ result = 1 if a < b else 0
170
+ elif op == 'eq':
171
+ result = 1 if a == b else 0
172
+
173
+ return text, a, b, op, result
174
+
175
+
176
+ def get_curriculum_max(epoch: int, total_epochs: int) -> int:
177
+ """
178
+ Curriculum learning: start with small numbers, gradually increase.
179
+ Epoch 0-20%: 0-9 (single digit)
180
+ Epoch 20-50%: 0-99 (two digit)
181
+ Epoch 50-100%: 0-255 (full range)
182
+ """
183
+ progress = epoch / total_epochs
184
+ if progress < 0.2:
185
+ return 9
186
+ elif progress < 0.5:
187
+ return 99
188
+ else:
189
+ return 255
190
+
191
+
192
+ def train_router(epochs: int = 100, batch_size: int = 256, lr: float = 1e-2, device: str = 'cuda'):
193
+ """Train only the router with ground truth bits."""
194
+ print("=" * 70)
195
+ print(" ROUTER-ONLY TRAINING (Ground Truth Bits)")
196
+ print("=" * 70)
197
+
198
+ circuits = FrozenThresholdCircuits(device=device)
199
+ router = OpRouter(input_dim=16 + 6, hidden_dim=64, n_ops=6).to(device)
200
+
201
+ print(f"\nRouter parameters: {sum(p.numel() for p in router.parameters()):,}")
202
+
203
+ def model_fn(a_bits, b_bits, op_onehot):
204
+ x = torch.cat([a_bits, b_bits, op_onehot], dim=-1)
205
+ op_weights = router(x)
206
+ return circuits(a_bits, b_bits, op_weights)
207
+
208
+ initial_fitness = compute_fitness(model_fn, n_samples=1000, device=device)
209
+ print(f"Initial fitness: {initial_fitness:.4f}")
210
+
211
+ optimizer = optim.AdamW(router.parameters(), lr=lr)
212
+ scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
213
+
214
+ print("\nTraining...")
215
+ print("-" * 70)
216
+
217
+ best_fitness = initial_fitness
218
+ start_time = time.perf_counter()
219
+
220
+ for epoch in range(epochs):
221
+ router.train()
222
+ epoch_loss = 0.0
223
+
224
+ for _ in range(100):
225
+ batch = generate_batch(batch_size, device)
226
+
227
+ optimizer.zero_grad()
228
+
229
+ x = torch.cat([batch['a_bits'], batch['b_bits'], batch['op_onehot']], dim=-1)
230
+ op_weights = router(x)
231
+ pred_bits = circuits(batch['a_bits'], batch['b_bits'], op_weights)
232
+
233
+ loss = compute_loss(pred_bits, batch['expected_bits'])
234
+ loss.backward()
235
+ optimizer.step()
236
+
237
+ epoch_loss += loss.item()
238
+
239
+ scheduler.step()
240
+
241
+ if (epoch + 1) % 10 == 0 or epoch == 0:
242
+ router.eval()
243
+ fitness, details = compute_fitness(model_fn, n_samples=2000, device=device, return_details=True)
244
+ elapsed = time.perf_counter() - start_time
245
+
246
+ if fitness > best_fitness:
247
+ best_fitness = fitness
248
+ marker = " *"
249
+ else:
250
+ marker = ""
251
+
252
+ print(f"Epoch {epoch+1:3d} | Loss: {epoch_loss/100:.4f} | "
253
+ f"Fitness: {fitness:.4f}{marker} | Time: {elapsed:.1f}s")
254
+
255
+ if fitness >= 0.9999:
256
+ print("\n TARGET: 100% FITNESS ACHIEVED")
257
+ break
258
+
259
+ print("\n" + "=" * 70)
260
+ print(" RESULTS")
261
+ print("=" * 70)
262
+
263
+ router.eval()
264
+ final_fitness, details = compute_fitness(model_fn, n_samples=5000, device=device, return_details=True)
265
+
266
+ print(f"\nFinal fitness: {final_fitness:.4f}")
267
+ print(f"\nPer-operation:")
268
+ for op in OPERATIONS:
269
+ acc = details['by_op'][op]['accuracy']
270
+ print(f" {op}: {acc:.4f}")
271
+
272
+ print(f"\nTotal time: {time.perf_counter() - start_time:.1f}s")
273
+
274
+ if final_fitness >= 0.99:
275
+ print("\nCONCLUSION: Router successfully learned operation dispatch.")
276
+ print(" With correct bit encoding, 100% is achievable.")
277
+
278
+ save_path = "D:/8bit-threshold-computer/llm_integration/trained/router.pt"
279
+ torch.save({
280
+ 'router_state_dict': router.state_dict(),
281
+ 'final_fitness': final_fitness,
282
+ 'params': sum(p.numel() for p in router.parameters()),
283
+ }, save_path)
284
+ print(f"\nSaved trained router to: {save_path}")
285
+
286
+ return router, final_fitness
287
+
288
+
289
+ def get_gpu_memory():
290
+ """Get GPU memory usage in MB."""
291
+ if torch.cuda.is_available():
292
+ return torch.cuda.memory_allocated() / 1024 / 1024, torch.cuda.max_memory_allocated() / 1024 / 1024
293
+ return 0, 0
294
+
295
+
296
+ def train_interface(epochs: int = 200, batch_size: int = 512, lr: float = 1e-3,
297
+ eval_interval: int = 10, device: str = 'cuda'):
298
+ """Train BitEncoder + OpRouter with ground truth bits."""
299
+ print("=" * 70)
300
+ print(" INTERFACE TRAINING (Encoder + Router)")
301
+ print("=" * 70)
302
+ print(f" Started at: {time.strftime('%H:%M:%S')}")
303
+
304
+ print("\n[1/4] Verifying frozen circuits...")
305
+ print(" Creating DirectCircuitModel...", end=" ", flush=True)
306
+ direct_model = DirectCircuitModel(device=device)
307
+ mem, max_mem = get_gpu_memory()
308
+ print(f"done. VRAM: {mem:.0f}MB")
309
+
310
+ def direct_fn(a, b, op):
311
+ return direct_model(a, b, op)
312
+
313
+ print(" Running fitness check (1000 samples)...", end=" ", flush=True)
314
+ circuit_fitness = compute_fitness(direct_fn, n_samples=1000, device=device)
315
+ print(f"done. Fitness: {circuit_fitness:.4f}")
316
+ if circuit_fitness < 0.999:
317
+ print(" ERROR: Circuits not achieving 100%. Aborting.")
318
+ return None, 0.0
319
+ print(" STATUS: PASS")
320
+
321
+ print("\n[2/4] Initializing model...")
322
+ print(" Creating ThresholdALU...", end=" ", flush=True)
323
+ model = ThresholdALU(device=device)
324
+ mem, max_mem = get_gpu_memory()
325
+ print(f"done. VRAM: {mem:.0f}MB")
326
+
327
+ trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
328
+ print(f" Trainable parameters: {trainable_params:,}")
329
+
330
+ def model_fn(a, b, op):
331
+ return model(a, b, op)
332
+
333
+ print(" Running initial fitness check...", end=" ", flush=True)
334
+ initial_fitness = compute_fitness(model_fn, n_samples=1000, device=device)
335
+ print(f"done. Fitness: {initial_fitness:.4f}")
336
+
337
+ print("\n[3/4] Setting up training...")
338
+ print(" Creating optimizer...", end=" ", flush=True)
339
+ optimizer = optim.AdamW(model.parameters(), lr=lr)
340
+ print("done.")
341
+ print(" Creating scheduler...", end=" ", flush=True)
342
+ scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
343
+ print("done.")
344
+
345
+ print(f" Config: lr={lr}, batch_size={batch_size}, epochs={epochs}")
346
+
347
+ print("\n[4/4] Training...")
348
+ print(" Generating first batch to warm up...", end=" ", flush=True)
349
+ warmup_batch = generate_batch(batch_size, device)
350
+ mem, max_mem = get_gpu_memory()
351
+ print(f"done. VRAM: {mem:.0f}MB (max: {max_mem:.0f}MB)")
352
+
353
+ print("-" * 70)
354
+
355
+ best_fitness = initial_fitness
356
+ start_time = time.perf_counter()
357
+ n_batches = 100
358
+
359
+ for epoch in range(epochs):
360
+ model.train()
361
+ epoch_loss = 0.0
362
+ epoch_start = time.perf_counter()
363
+
364
+ for batch_idx in range(n_batches):
365
+ batch = generate_batch(batch_size, device)
366
+
367
+ optimizer.zero_grad()
368
+
369
+ pred_bits = model(batch['a_bits'], batch['b_bits'], batch['op_onehot'])
370
+
371
+ loss = compute_loss(pred_bits, batch['expected_bits'])
372
+
373
+ loss.backward()
374
+ optimizer.step()
375
+
376
+ epoch_loss += loss.item()
377
+
378
+ if batch_idx == 0 and epoch == 0:
379
+ mem, max_mem = get_gpu_memory()
380
+ print(f" First forward/backward done. VRAM: {mem:.0f}MB (max: {max_mem:.0f}MB)")
381
+
382
+ if (batch_idx + 1) % 25 == 0:
383
+ avg_so_far = epoch_loss / (batch_idx + 1)
384
+ print(f" Epoch {epoch+1} batch {batch_idx+1}/{n_batches} | loss: {avg_so_far:.4f}", flush=True)
385
+
386
+ scheduler.step()
387
+
388
+ avg_loss = epoch_loss / n_batches
389
+ epoch_time = time.perf_counter() - epoch_start
390
+
391
+ if (epoch + 1) % 5 == 0 or epoch == 0: # Eval every 5 epochs
392
+ model.eval()
393
+ fitness, details = compute_fitness(
394
+ model_fn, n_samples=2000, device=device, return_details=True
395
+ )
396
+
397
+ elapsed = time.perf_counter() - start_time
398
+
399
+ if fitness > best_fitness:
400
+ best_fitness = fitness
401
+ marker = " *"
402
+ else:
403
+ marker = ""
404
+
405
+ mem, _ = get_gpu_memory()
406
+ print(f"Epoch {epoch+1:4d} | Loss: {avg_loss:.4f} | "
407
+ f"Fitness: {fitness:.4f}{marker} | "
408
+ f"LR: {scheduler.get_last_lr()[0]:.2e} | "
409
+ f"VRAM: {mem:.0f}MB | "
410
+ f"Time: {elapsed:.1f}s ({epoch_time:.1f}s/epoch)")
411
+
412
+ if fitness >= 0.9999:
413
+ print("\n" + "=" * 70)
414
+ print(" TARGET ACHIEVED: 100% FITNESS")
415
+ print("=" * 70)
416
+ break
417
+
418
+ print("\n" + "=" * 70)
419
+ print(" TRAINING COMPLETE")
420
+ print("=" * 70)
421
+
422
+ model.eval()
423
+ final_fitness, details = compute_fitness(
424
+ model_fn, n_samples=5000, device=device, return_details=True
425
+ )
426
+
427
+ print(f"\nFinal fitness: {final_fitness:.4f}")
428
+ print(f"Best fitness: {best_fitness:.4f}")
429
+ print(f"\nPer-operation breakdown:")
430
+ for op in OPERATIONS:
431
+ acc = details['by_op'][op]['accuracy']
432
+ print(f" {op:6}: {acc:.4f}")
433
+
434
+ print(f"\nTotal time: {time.perf_counter() - start_time:.1f}s")
435
+
436
+ save_path = "D:/8bit-threshold-computer/llm_integration/trained/interface.pt"
437
+ torch.save({
438
+ 'encoder_state_dict': model.encoder.state_dict(),
439
+ 'router_state_dict': model.router.state_dict(),
440
+ 'final_fitness': final_fitness,
441
+ 'best_fitness': best_fitness,
442
+ }, save_path)
443
+ print(f"\nSaved trained model to: {save_path}")
444
+
445
+ return model, final_fitness
446
+
447
+
448
+ def compute_llm_loss(pred_bits, a_bits, b_bits, op_logits,
449
+ target_result, target_a, target_b, target_op_idx,
450
+ bit_weight: float = 2.0):
451
+ """
452
+ Multi-component loss for LLM training.
453
+ bit_weight: multiplier for a/b bit losses (default 2x since extraction is the bottleneck)
454
+ """
455
+ result_loss = nn.functional.binary_cross_entropy_with_logits(
456
+ pred_bits, target_result
457
+ )
458
+
459
+ a_bits_safe = torch.clamp(a_bits, 0.0, 1.0)
460
+ b_bits_safe = torch.clamp(b_bits, 0.0, 1.0)
461
+ a_bits_safe = torch.nan_to_num(a_bits_safe, nan=0.5, posinf=1.0, neginf=0.0)
462
+ b_bits_safe = torch.nan_to_num(b_bits_safe, nan=0.5, posinf=1.0, neginf=0.0)
463
+
464
+ a_loss = nn.functional.binary_cross_entropy(
465
+ torch.clamp(a_bits_safe, 1e-6, 1-1e-6), target_a
466
+ )
467
+ b_loss = nn.functional.binary_cross_entropy(
468
+ torch.clamp(b_bits_safe, 1e-6, 1-1e-6), target_b
469
+ )
470
+
471
+ op_loss = nn.functional.cross_entropy(op_logits, target_op_idx)
472
+
473
+ total = result_loss + bit_weight * a_loss + bit_weight * b_loss + op_loss
474
+ total = torch.nan_to_num(total, nan=10.0, posinf=10.0, neginf=0.0)
475
+
476
+ return total, {
477
+ 'result': result_loss.item() if not torch.isnan(result_loss) else 10.0,
478
+ 'a': a_loss.item() if not torch.isnan(a_loss) else 10.0,
479
+ 'b': b_loss.item() if not torch.isnan(b_loss) else 10.0,
480
+ 'op': op_loss.item() if not torch.isnan(op_loss) else 10.0
481
+ }
482
+
483
+
484
+ def value_to_digits(value: int) -> list:
485
+ """Convert integer value to list of digits (hundreds, tens, ones)."""
486
+ digits = []
487
+ for place in [100, 10, 1]:
488
+ digit = (value // place) % 10
489
+ digits.append(digit)
490
+ return digits
491
+
492
+
493
+ def compute_positional_digit_loss(pred_bits, op_logits, a_digit_logits_list, b_digit_logits_list,
494
+ target_result, target_op_idx, target_a_values, target_b_values,
495
+ device, digit_weight: float = 5.0):
496
+ """
497
+ Loss for positional digit extraction with DIRECT digit supervision.
498
+
499
+ This provides strong gradients by directly supervising digit classification
500
+ instead of going through the value -> bits conversion.
501
+ """
502
+ result_loss = nn.functional.binary_cross_entropy_with_logits(
503
+ pred_bits, target_result
504
+ )
505
+
506
+ op_loss = nn.functional.cross_entropy(op_logits, target_op_idx)
507
+
508
+ a_digit_loss = torch.tensor(0.0, device=device)
509
+ b_digit_loss = torch.tensor(0.0, device=device)
510
+ n_a_digits = 0
511
+ n_b_digits = 0
512
+
513
+ for i, (a_logits_list, b_logits_list) in enumerate(zip(a_digit_logits_list, b_digit_logits_list)):
514
+ target_a = target_a_values[i].item()
515
+ target_b = target_b_values[i].item()
516
+
517
+ a_digits = value_to_digits(int(target_a))
518
+ b_digits = value_to_digits(int(target_b))
519
+
520
+ n_a = len(a_logits_list)
521
+ n_b = len(b_logits_list)
522
+
523
+ if n_a > 0:
524
+ target_a_digits = a_digits[-n_a:]
525
+ for j, logits in enumerate(a_logits_list):
526
+ target_digit = torch.tensor([target_a_digits[j]], device=device, dtype=torch.long)
527
+ a_digit_loss = a_digit_loss + nn.functional.cross_entropy(logits.unsqueeze(0), target_digit)
528
+ n_a_digits += 1
529
+
530
+ if n_b > 0:
531
+ target_b_digits = b_digits[-n_b:]
532
+ for j, logits in enumerate(b_logits_list):
533
+ target_digit = torch.tensor([target_b_digits[j]], device=device, dtype=torch.long)
534
+ b_digit_loss = b_digit_loss + nn.functional.cross_entropy(logits.unsqueeze(0), target_digit)
535
+ n_b_digits += 1
536
+
537
+ if n_a_digits > 0:
538
+ a_digit_loss = a_digit_loss / n_a_digits
539
+ if n_b_digits > 0:
540
+ b_digit_loss = b_digit_loss / n_b_digits
541
+
542
+ total = result_loss + digit_weight * a_digit_loss + digit_weight * b_digit_loss + op_loss
543
+ total = torch.nan_to_num(total, nan=10.0, posinf=10.0, neginf=0.0)
544
+
545
+ return total, {
546
+ 'result': result_loss.item() if not torch.isnan(result_loss) else 10.0,
547
+ 'a_digit': a_digit_loss.item() if not torch.isnan(a_digit_loss) else 10.0,
548
+ 'b_digit': b_digit_loss.item() if not torch.isnan(b_digit_loss) else 10.0,
549
+ 'op': op_loss.item() if not torch.isnan(op_loss) else 10.0
550
+ }
551
+
552
+
553
+ def compute_hybrid_loss(pred_bits, op_logits, used_lookup,
554
+ a_digit_logits, b_digit_logits,
555
+ target_result, target_a_values, target_b_values, target_op_idx,
556
+ device, digit_weight: float = 2.0):
557
+ """
558
+ Loss for hybrid extraction with digit-level prediction.
559
+
560
+ Uses cross-entropy on each digit (hundreds, tens, ones) for word samples.
561
+ Samples using digit lookup are already 100% accurate - no loss computed.
562
+ """
563
+ result_loss = nn.functional.binary_cross_entropy_with_logits(
564
+ pred_bits, target_result
565
+ )
566
+
567
+ op_loss = nn.functional.cross_entropy(op_logits, target_op_idx)
568
+
569
+ word_mask = ~used_lookup
570
+ n_words = word_mask.sum().item()
571
+
572
+ if n_words > 0 and a_digit_logits is not None and b_digit_logits is not None:
573
+ target_a_word = target_a_values[word_mask].long()
574
+ target_b_word = target_b_values[word_mask].long()
575
+
576
+ a_hundreds = target_a_word // 100
577
+ a_tens = (target_a_word % 100) // 10
578
+ a_ones = target_a_word % 10
579
+
580
+ b_hundreds = target_b_word // 100
581
+ b_tens = (target_b_word % 100) // 10
582
+ b_ones = target_b_word % 10
583
+
584
+ a_logits = a_digit_logits.view(-1, 3, 10)
585
+ b_logits = b_digit_logits.view(-1, 3, 10)
586
+
587
+ a_digit_loss = (
588
+ nn.functional.cross_entropy(a_logits[:, 0], a_hundreds) +
589
+ nn.functional.cross_entropy(a_logits[:, 1], a_tens) +
590
+ nn.functional.cross_entropy(a_logits[:, 2], a_ones)
591
+ ) / 3.0
592
+
593
+ b_digit_loss = (
594
+ nn.functional.cross_entropy(b_logits[:, 0], b_hundreds) +
595
+ nn.functional.cross_entropy(b_logits[:, 1], b_tens) +
596
+ nn.functional.cross_entropy(b_logits[:, 2], b_ones)
597
+ ) / 3.0
598
+ else:
599
+ a_digit_loss = torch.tensor(0.0, device=device)
600
+ b_digit_loss = torch.tensor(0.0, device=device)
601
+
602
+ total = result_loss + op_loss + digit_weight * (a_digit_loss + b_digit_loss)
603
+ total = torch.nan_to_num(total, nan=10.0, posinf=10.0, neginf=0.0)
604
+
605
+ return total, {
606
+ 'result': result_loss.item() if not torch.isnan(result_loss) else 10.0,
607
+ 'a_digit': a_digit_loss.item() if not torch.isnan(a_digit_loss) else 10.0,
608
+ 'b_digit': b_digit_loss.item() if not torch.isnan(b_digit_loss) else 10.0,
609
+ 'op': op_loss.item() if not torch.isnan(op_loss) else 10.0,
610
+ 'n_words': n_words,
611
+ 'n_lookup': used_lookup.sum().item()
612
+ }
613
+
614
+
615
+ def evaluate_llm(model, n_samples: int = 500):
616
+ """Evaluate LLM model on random problems (mixed digit/word format)."""
617
+ model.extractor.eval()
618
+ correct = 0
619
+ op_correct = 0
620
+
621
+ for _ in range(n_samples):
622
+ text, a, b, op, expected = generate_problem()
623
+
624
+ with torch.no_grad():
625
+ outputs = model([text])
626
+ result_bits = outputs[0]
627
+ op_logits = outputs[3]
628
+
629
+ pred_result = bits_to_int(result_bits[0])
630
+ pred_op = OPERATIONS[op_logits[0].argmax().item()]
631
+
632
+ if pred_result == expected:
633
+ correct += 1
634
+ if pred_op == op:
635
+ op_correct += 1
636
+
637
+ model.extractor.train()
638
+ return correct / n_samples, op_correct / n_samples
639
+
640
+
641
+ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
642
+ unfreeze_layers: int = 0, extract_layer: int = -1,
643
+ position_extract: bool = False, digit_pred: bool = False,
644
+ positional_digit: bool = False, device: str = 'cuda'):
645
+ """
646
+ Train extractor with LLM hidden states.
647
+
648
+ Args:
649
+ unfreeze_layers: Number of top transformer layers to unfreeze (0 = fully frozen)
650
+ extract_layer: Which layer to extract from (-1 = last)
651
+ position_extract: Use position-specific extraction (legacy)
652
+ digit_pred: Predict digits instead of bits (legacy)
653
+ positional_digit: Use positional digit extraction (legacy, 100% on digits only)
654
+ """
655
+ hybrid = not (positional_digit or position_extract or digit_pred)
656
+
657
+ print("=" * 70)
658
+ print(" LLM TRAINING")
659
+ if unfreeze_layers > 0:
660
+ print(f" {unfreeze_layers} transformer layers unfrozen")
661
+ else:
662
+ print(" LLM frozen")
663
+ if extract_layer != -1:
664
+ print(f" Extracting from layer {extract_layer}")
665
+ if hybrid:
666
+ print(" HYBRID extraction (digit lookup + word learning)")
667
+ elif positional_digit:
668
+ print(" POSITIONAL DIGIT extraction (legacy, 100% on digits only)")
669
+ elif position_extract:
670
+ print(" Position-specific extraction (legacy)")
671
+ elif digit_pred:
672
+ print(" Digit-level prediction (legacy)")
673
+ print("=" * 70)
674
+
675
+ print("\nInitializing model...")
676
+ model = ArithmeticModel(
677
+ device=device,
678
+ unfreeze_layers=unfreeze_layers,
679
+ extract_layer=extract_layer,
680
+ position_extract=position_extract,
681
+ digit_pred=digit_pred,
682
+ positional_digit=positional_digit,
683
+ hybrid=hybrid
684
+ )
685
+
686
+ optimizer = optim.AdamW(model.trainable_parameters(), lr=lr)
687
+ scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
688
+
689
+ print(f"\nTraining config:")
690
+ print(f" Epochs: {epochs}")
691
+ print(f" Batch size: {batch_size}")
692
+ print(f" Learning rate: {lr}")
693
+ print(f" Unfreeze layers: {unfreeze_layers}")
694
+ print(f" Samples/epoch: {batch_size * 20}")
695
+
696
+ print(f"\nInitial evaluation (200 samples)...")
697
+ acc, op_acc = evaluate_llm(model, 200)
698
+ print(f" Accuracy: {acc:.4f}, Op accuracy: {op_acc:.4f}")
699
+
700
+ print(f"\nStarting training...")
701
+ print("-" * 70)
702
+
703
+ best_acc = acc
704
+ start_time = time.perf_counter()
705
+
706
+ for epoch in range(epochs):
707
+ model.extractor.train()
708
+ if unfreeze_layers > 0:
709
+ model.llm.train()
710
+
711
+ max_val = get_curriculum_max(epoch, epochs)
712
+
713
+ epoch_loss = 0
714
+ if hybrid:
715
+ epoch_losses = {'result': 0, 'a_digit': 0, 'b_digit': 0, 'op': 0, 'n_words': 0, 'n_lookup': 0}
716
+ elif positional_digit:
717
+ epoch_losses = {'result': 0, 'a_digit': 0, 'b_digit': 0, 'op': 0}
718
+ else:
719
+ epoch_losses = {'result': 0, 'a': 0, 'b': 0, 'op': 0}
720
+ n_batches = 20
721
+ epoch_start = time.perf_counter()
722
+
723
+ for batch_idx in range(n_batches):
724
+ batch_texts = []
725
+ batch_a = []
726
+ batch_b = []
727
+ batch_op = []
728
+ batch_result = []
729
+ batch_a_values = []
730
+ batch_b_values = []
731
+
732
+ for _ in range(batch_size):
733
+ text, a, b, op, result = generate_problem(max_val)
734
+ batch_texts.append(text)
735
+ batch_a.append(int_to_bits(a, device))
736
+ batch_b.append(int_to_bits(b, device))
737
+ batch_op.append(OPERATIONS.index(op))
738
+ batch_result.append(int_to_bits(result, device))
739
+ batch_a_values.append(a)
740
+ batch_b_values.append(b)
741
+
742
+ target_a = torch.stack(batch_a)
743
+ target_b = torch.stack(batch_b)
744
+ target_op = torch.tensor(batch_op, device=device)
745
+ target_result = torch.stack(batch_result)
746
+ target_a_values = torch.tensor(batch_a_values, device=device, dtype=torch.float32)
747
+ target_b_values = torch.tensor(batch_b_values, device=device, dtype=torch.float32)
748
+
749
+ optimizer.zero_grad()
750
+
751
+ outputs = model(batch_texts)
752
+ pred_bits, a_bits, b_bits, op_logits = outputs[0], outputs[1], outputs[2], outputs[3]
753
+
754
+ if hybrid:
755
+ a_values, b_values, used_lookup = outputs[4], outputs[5], outputs[6]
756
+ a_digit_logits, b_digit_logits = outputs[7], outputs[8]
757
+ loss, losses = compute_hybrid_loss(
758
+ pred_bits, op_logits, used_lookup,
759
+ a_digit_logits, b_digit_logits,
760
+ target_result, target_a_values, target_b_values, target_op, device
761
+ )
762
+ elif positional_digit:
763
+ a_digit_logits_list = outputs[7]
764
+ b_digit_logits_list = outputs[8]
765
+ loss, losses = compute_positional_digit_loss(
766
+ pred_bits, op_logits, a_digit_logits_list, b_digit_logits_list,
767
+ target_result, target_op, target_a_values, target_b_values, device
768
+ )
769
+ else:
770
+ loss, losses = compute_llm_loss(
771
+ pred_bits, a_bits, b_bits, op_logits,
772
+ target_result, target_a, target_b, target_op
773
+ )
774
+
775
+ loss.backward()
776
+ torch.nn.utils.clip_grad_norm_(model.trainable_parameters(), 1.0)
777
+ optimizer.step()
778
+
779
+ epoch_loss += loss.item()
780
+ for k in epoch_losses:
781
+ epoch_losses[k] += losses[k]
782
+
783
+ if (batch_idx + 1) % 5 == 0:
784
+ avg_so_far = epoch_loss / (batch_idx + 1)
785
+ print(f" Epoch {epoch+1} batch {batch_idx+1}/{n_batches} | loss: {avg_so_far:.4f}", flush=True)
786
+
787
+ epoch_time = time.perf_counter() - epoch_start
788
+ scheduler.step()
789
+
790
+ avg_loss = epoch_loss / n_batches
791
+ for k in epoch_losses:
792
+ epoch_losses[k] /= n_batches
793
+
794
+ acc, op_acc = evaluate_llm(model, 300)
795
+ elapsed = time.perf_counter() - start_time
796
+
797
+ marker = " *" if acc > best_acc else ""
798
+ if acc > best_acc:
799
+ best_acc = acc
800
+
801
+ mem, _ = get_gpu_memory()
802
+ print(f"Epoch {epoch+1:3d} | Loss: {avg_loss:.4f} | "
803
+ f"Acc: {acc:.4f}{marker} | OpAcc: {op_acc:.4f} | "
804
+ f"Range: 0-{max_val} | VRAM: {mem:.0f}MB | Time: {elapsed:.0f}s")
805
+ if hybrid:
806
+ print(f" Losses - result:{epoch_losses['result']:.4f} "
807
+ f"a_digit:{epoch_losses['a_digit']:.4f} b_digit:{epoch_losses['b_digit']:.4f} "
808
+ f"op:{epoch_losses['op']:.4f} | words:{epoch_losses['n_words']:.0f} lookup:{epoch_losses['n_lookup']:.0f}")
809
+ elif positional_digit:
810
+ print(f" Losses - result:{epoch_losses['result']:.4f} "
811
+ f"a_digit:{epoch_losses['a_digit']:.4f} b_digit:{epoch_losses['b_digit']:.4f} "
812
+ f"op:{epoch_losses['op']:.4f}")
813
+ else:
814
+ print(f" Losses - result:{epoch_losses['result']:.4f} "
815
+ f"a:{epoch_losses['a']:.4f} b:{epoch_losses['b']:.4f} "
816
+ f"op:{epoch_losses['op']:.4f}")
817
+
818
+ if acc >= 0.99:
819
+ print("\n" + "=" * 70)
820
+ print(" TARGET ACHIEVED: 99%+ ACCURACY")
821
+ print("=" * 70)
822
+ break
823
+
824
+ print("\n" + "=" * 70)
825
+ print(" FINAL EVALUATION")
826
+ print("=" * 70)
827
+
828
+ acc, op_acc = evaluate_llm(model, 1000)
829
+ print(f"Final accuracy: {acc:.4f}")
830
+ print(f"Final op accuracy: {op_acc:.4f}")
831
+ print(f"Best accuracy: {best_acc:.4f}")
832
+
833
+ print("\nSample predictions:")
834
+ for _ in range(10):
835
+ text, a, b, op, expected = generate_problem()
836
+ with torch.no_grad():
837
+ outputs = model([text])
838
+ result_bits, a_bits, b_bits, op_logits = outputs[0], outputs[1], outputs[2], outputs[3]
839
+ pred = bits_to_int(result_bits[0])
840
+ pred_a = bits_to_int(a_bits[0])
841
+ pred_b = bits_to_int(b_bits[0])
842
+ pred_op = OPERATIONS[op_logits[0].argmax().item()]
843
+
844
+ status = "OK" if pred == expected else "WRONG"
845
+ print(f" '{text}' = {expected} | pred={pred} (a={pred_a}, b={pred_b}, op={pred_op}) [{status}]")
846
+
847
+ save_path = "D:/8bit-threshold-computer/llm_integration/trained/llm.pt"
848
+ save_dict = {
849
+ 'extractor_state_dict': model.extractor.state_dict(),
850
+ 'final_accuracy': acc,
851
+ 'best_accuracy': best_acc,
852
+ 'unfreeze_layers': unfreeze_layers,
853
+ }
854
+ if unfreeze_layers > 0:
855
+ save_dict['llm_state_dict'] = {
856
+ k: v for k, v in model.llm.state_dict().items()
857
+ if any(f'layers.{i}.' in k for i in range(len(model.llm.model.layers) - unfreeze_layers, len(model.llm.model.layers)))
858
+ }
859
+ torch.save(save_dict, save_path)
860
+ print(f"\nSaved to: {save_path}")
861
+
862
+ return model, acc
863
+
864
+
865
+ def main():
866
+ parser = argparse.ArgumentParser(
867
+ description='Unified training for threshold circuit LLM integration',
868
+ formatter_class=argparse.RawDescriptionHelpFormatter,
869
+ epilog="""
870
+ Modes:
871
+ router - Train only OpRouter with ground truth bits (sanity check)
872
+ interface - Train BitEncoder + OpRouter with ground truth bits (sanity check)
873
+ llm - Train extractor with LLM hidden states (the real training)
874
+
875
+ LLM options:
876
+ --unfreeze_layers N Fine-tune top N transformer layers
877
+ --extract_layer N Extract from layer N (-1 = last)
878
+ --position_extract Use position-specific extraction
879
+ --digit_pred Predict digits instead of bits
880
+
881
+ Baked-in: curriculum learning (0-9 -> 0-99 -> 0-255), 2x loss weight for a/b
882
+
883
+ Examples:
884
+ python train.py --mode llm --epochs 100
885
+ python train.py --mode llm --position_extract
886
+ python train.py --mode llm --digit_pred --extract_layer 12
887
+ python train.py --mode llm --unfreeze_layers 4 --batch_size 4096
888
+ """
889
+ )
890
+ parser.add_argument('--mode', type=str, required=True,
891
+ choices=['router', 'interface', 'llm'],
892
+ help='Training mode')
893
+ parser.add_argument('--epochs', type=int, default=100, help='Number of epochs')
894
+ parser.add_argument('--batch_size', type=int, default=512, help='Batch size (default: 512)')
895
+ parser.add_argument('--lr', type=float, default=None,
896
+ help='Learning rate (default: mode-specific)')
897
+ parser.add_argument('--unfreeze_layers', type=int, default=0,
898
+ help='Unfreeze top N transformer layers (default 0 = frozen)')
899
+ parser.add_argument('--extract_layer', type=int, default=0,
900
+ help='Which layer to extract from (default: 0 = embeddings, best for digits)')
901
+ parser.add_argument('--position_extract', action='store_true',
902
+ help='Use position-specific extraction (legacy)')
903
+ parser.add_argument('--digit_pred', action='store_true',
904
+ help='Predict digits instead of bits (legacy)')
905
+ parser.add_argument('--positional_digit', action='store_true', default=False,
906
+ help='Use positional digit extraction (legacy, 100%% on digits only)')
907
+ parser.add_argument('--device', type=str, default='cuda', help='Device')
908
+ args = parser.parse_args()
909
+
910
+ torch.manual_seed(42)
911
+ random.seed(42)
912
+
913
+ if args.mode == 'router':
914
+ lr = args.lr if args.lr is not None else 1e-2
915
+ train_router(epochs=args.epochs, batch_size=args.batch_size, lr=lr, device=args.device)
916
+
917
+ elif args.mode == 'interface':
918
+ lr = args.lr if args.lr is not None else 1e-3
919
+ model, fitness = train_interface(
920
+ epochs=args.epochs, batch_size=args.batch_size, lr=lr, device=args.device
921
+ )
922
+
923
+ print("\n" + "=" * 70)
924
+ print(" EXPERIMENT SUMMARY")
925
+ print("=" * 70)
926
+ print(f"\n Control (Vanilla SmolLM2-360M): 11.90%")
927
+ print(f" Experimental (Trained Interface): {100*fitness:.2f}%")
928
+ if fitness > 0:
929
+ print(f"\n Improvement: {100*(fitness - 0.119)/0.119:.1f}%")
930
+
931
+ if fitness >= 0.99:
932
+ print("\n CONCLUSION: Frozen threshold circuits + trained interface")
933
+ print(" achieves near-perfect arithmetic accuracy.")
934
+ print(" Core thesis VALIDATED.")
935
+ else:
936
+ print(f"\n CONCLUSION: Further training or architecture changes needed.")
937
+ print(f" Current gap: {100*(1.0 - fitness):.2f}%")
938
+
939
+ elif args.mode == 'llm':
940
+ lr = args.lr if args.lr is not None else 3e-4
941
+ train_llm(
942
+ epochs=args.epochs,
943
+ batch_size=args.batch_size,
944
+ lr=lr,
945
+ unfreeze_layers=args.unfreeze_layers,
946
+ extract_layer=args.extract_layer,
947
+ position_extract=args.position_extract,
948
+ digit_pred=args.digit_pred,
949
+ positional_digit=args.positional_digit,
950
+ device=args.device
951
+ )
952
+
953
+
954
+ if __name__ == "__main__":
955
+ main()
llm_integration/trained/llm.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:69ff9e584d08c9707e2932bc0f62cf89cd7167a453065146252c6b7cdb57cbc4
3
+ size 17296376
llm_integration/trained/router.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0ddfc24cd4a98b65de8d434bb843ebd24f8c902d067201fd7954e7b623a8ebcd
3
+ size 9811
machines.py ADDED
@@ -0,0 +1,1719 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """neural_subleq8 and neural_rv32 — runtimes, reference emulators, assemblers,
2
+ an RV32 object loader, and test suites for the two endpoint machines of the
3
+ family. One file; run the suites with:
4
+
5
+ python machines.py # both machines
6
+ python machines.py subleq # SUBLEQ-8 only
7
+ python machines.py rv32 # RV32 assembled-program suite
8
+ python machines.py rv32-c # RV32 running stock-compiler C
9
+
10
+ Build the model files with `python build.py --apply subleq` and
11
+ `python build.py --apply rv32`.
12
+ """
13
+
14
+ from __future__ import annotations
15
+ import argparse
16
+ import os
17
+ import shutil
18
+ import struct
19
+ import subprocess
20
+ import sys
21
+ import tempfile
22
+ from pathlib import Path
23
+ from typing import Dict, List, Optional, Tuple
24
+
25
+ import torch
26
+ from safetensors import safe_open
27
+
28
+ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
29
+ from eval import NetlistEvaluator, heaviside, load_metadata
30
+ from eval_all import GenericThresholdALU, int_to_bits, bits_to_int
31
+
32
+ # ============================================================================
33
+ # RV32 relocatable-object loader
34
+ # ============================================================================
35
+ R = {
36
+ "R_RISCV_32": 1,
37
+ "R_RISCV_BRANCH": 16,
38
+ "R_RISCV_JAL": 17,
39
+ "R_RISCV_CALL": 18,
40
+ "R_RISCV_CALL_PLT": 19,
41
+ "R_RISCV_PCREL_HI20": 23,
42
+ "R_RISCV_PCREL_LO12_I": 24,
43
+ "R_RISCV_PCREL_LO12_S": 25,
44
+ "R_RISCV_HI20": 26,
45
+ "R_RISCV_LO12_I": 27,
46
+ "R_RISCV_LO12_S": 28,
47
+ "R_RISCV_RELAX": 51,
48
+ }
49
+ RREV = {v: k for k, v in R.items()}
50
+
51
+
52
+ class Elf:
53
+ def __init__(self, data: bytes):
54
+ assert data[:4] == b"\x7fELF", "not ELF"
55
+ assert data[4] == 1 and data[5] == 1, "need ELF32 LE"
56
+ (self.e_shoff,) = struct.unpack_from("<I", data, 0x20)
57
+ (self.e_shentsize, self.e_shnum, self.e_shstrndx) = struct.unpack_from("<HHH", data, 0x2E)
58
+ self.data = data
59
+ self.sections = []
60
+ for i in range(self.e_shnum):
61
+ off = self.e_shoff + i * self.e_shentsize
62
+ name, typ, flags, addr, offset, size, link, info, align, entsz = \
63
+ struct.unpack_from("<IIIIIIIIII", data, off)
64
+ self.sections.append(dict(name=name, type=typ, flags=flags, addr=addr,
65
+ offset=offset, size=size, link=link, info=info,
66
+ align=align, entsize=entsz))
67
+ shstr = self.sections[self.e_shstrndx]
68
+ self.shstr = data[shstr["offset"]:shstr["offset"] + shstr["size"]]
69
+ for s in self.sections:
70
+ s["sname"] = self._str(self.shstr, s["name"])
71
+
72
+ @staticmethod
73
+ def _str(tab: bytes, off: int) -> str:
74
+ end = tab.find(b"\x00", off)
75
+ return tab[off:end].decode()
76
+
77
+ def symbols(self):
78
+ symtab = next(s for s in self.sections if s["sname"] == ".symtab")
79
+ strtab = self.sections[symtab["link"]]
80
+ st = self.data[strtab["offset"]:strtab["offset"] + strtab["size"]]
81
+ syms = []
82
+ n = symtab["size"] // 16
83
+ for i in range(n):
84
+ off = symtab["offset"] + i * 16
85
+ name, value, size, info, other, shndx = struct.unpack_from("<IIIBBH", self.data, off)
86
+ syms.append(dict(name=self._str(st, name), value=value, size=size,
87
+ info=info, shndx=shndx))
88
+ return syms
89
+
90
+
91
+ def load(obj: bytes, base: int = 0, mem_size: int = 65536,
92
+ entry_symbol: str = "_start") -> Tuple[List[int], int]:
93
+ elf = Elf(obj)
94
+ mem = [0] * mem_size
95
+
96
+ # Lay out allocatable progbits/nobits sections (SHF_ALLOC = 0x2).
97
+ placed: Dict[int, int] = {} # section index -> load address
98
+ cursor = base
99
+ for i, s in enumerate(elf.sections):
100
+ if not (s["flags"] & 0x2):
101
+ continue
102
+ align = max(s["align"], 1)
103
+ cursor = (cursor + align - 1) & ~(align - 1)
104
+ placed[i] = cursor
105
+ if s["type"] == 1: # PROGBITS
106
+ body = elf.data[s["offset"]:s["offset"] + s["size"]]
107
+ for j, byte in enumerate(body):
108
+ mem[cursor + j] = byte
109
+ cursor += s["size"]
110
+
111
+ syms = elf.symbols()
112
+
113
+ def sym_addr(idx: int) -> int:
114
+ sym = syms[idx]
115
+ if sym["shndx"] in placed:
116
+ return placed[sym["shndx"]] + sym["value"]
117
+ return sym["value"]
118
+
119
+ # Apply relocations. PCREL_LO12 references the symbol of its paired HI20
120
+ # site, so first index HI20/PCREL_HI20 results by their location.
121
+ hi_at: Dict[int, int] = {}
122
+ for s in elf.sections:
123
+ if s["type"] != 4: # RELA
124
+ continue
125
+ target = elf.sections[s["info"]]
126
+ if s["info"] not in placed:
127
+ continue
128
+ tbase = placed[s["info"]]
129
+ n = s["size"] // 12
130
+ for i in range(n):
131
+ off = s["offset"] + i * 12
132
+ r_offset, r_info, r_addend = struct.unpack_from("<IIi", elf.data, off)
133
+ r_type = r_info & 0xFF
134
+ r_sym = r_info >> 8
135
+ loc = tbase + r_offset
136
+ if r_type in (R["R_RISCV_PCREL_HI20"], R["R_RISCV_HI20"]):
137
+ hi_at[loc] = sym_addr(r_sym) + r_addend
138
+
139
+ def rd(a):
140
+ return mem[a] | (mem[a + 1] << 8) | (mem[a + 2] << 16) | (mem[a + 3] << 24)
141
+
142
+ def wr(a, v):
143
+ for k in range(4):
144
+ mem[a + k] = (v >> (8 * k)) & 0xFF
145
+
146
+ for s in elf.sections:
147
+ if s["type"] != 4 or s["info"] not in placed:
148
+ continue
149
+ tbase = placed[s["info"]]
150
+ n = s["size"] // 12
151
+ for i in range(n):
152
+ off = s["offset"] + i * 12
153
+ r_offset, r_info, r_addend = struct.unpack_from("<IIi", elf.data, off)
154
+ r_type = r_info & 0xFF
155
+ r_sym = r_info >> 8
156
+ loc = tbase + r_offset
157
+ S = sym_addr(r_sym) + r_addend
158
+ w = rd(loc)
159
+ if r_type == R["R_RISCV_32"]:
160
+ wr(loc, S & 0xFFFFFFFF)
161
+ elif r_type == R["R_RISCV_HI20"]:
162
+ wr(loc, (w & 0xFFF) | (((S + 0x800) >> 12) << 12) & 0xFFFFF000)
163
+ elif r_type == R["R_RISCV_LO12_I"]:
164
+ wr(loc, (w & 0xFFFFF) | ((S & 0xFFF) << 20))
165
+ elif r_type == R["R_RISCV_LO12_S"]:
166
+ imm = S & 0xFFF
167
+ wr(loc, (w & 0x1FFF07F) | ((imm >> 5) << 25) | ((imm & 31) << 7))
168
+ elif r_type == R["R_RISCV_PCREL_HI20"]:
169
+ delta = S - loc
170
+ wr(loc, (w & 0xFFF) | (((delta + 0x800) >> 12) << 12) & 0xFFFFF000)
171
+ elif r_type in (R["R_RISCV_PCREL_LO12_I"], R["R_RISCV_PCREL_LO12_S"]):
172
+ hi = hi_at[sym_addr(r_sym) + r_addend]
173
+ delta = hi - (sym_addr(r_sym) + r_addend)
174
+ if r_type == R["R_RISCV_PCREL_LO12_I"]:
175
+ wr(loc, (w & 0xFFFFF) | ((delta & 0xFFF) << 20))
176
+ else:
177
+ imm = delta & 0xFFF
178
+ wr(loc, (w & 0x1FFF07F) | ((imm >> 5) << 25) | ((imm & 31) << 7))
179
+ elif r_type in (R["R_RISCV_CALL"], R["R_RISCV_CALL_PLT"]):
180
+ delta = S - loc
181
+ hi = (delta + 0x800) >> 12
182
+ lo = delta - (hi << 12)
183
+ w2 = rd(loc + 4)
184
+ wr(loc, (w & 0xFFF) | ((hi << 12) & 0xFFFFF000))
185
+ wr(loc + 4, (w2 & 0xFFFFF) | ((lo & 0xFFF) << 20))
186
+ elif r_type == R["R_RISCV_JAL"]:
187
+ d = S - loc
188
+ imm = (((d >> 20) & 1) << 31) | (((d >> 1) & 0x3FF) << 21) | \
189
+ (((d >> 11) & 1) << 20) | (((d >> 12) & 0xFF) << 12)
190
+ wr(loc, (w & 0xFFF) | imm)
191
+ elif r_type == R["R_RISCV_BRANCH"]:
192
+ d = S - loc
193
+ imm = (((d >> 12) & 1) << 31) | (((d >> 5) & 0x3F) << 25) | \
194
+ (((d >> 1) & 0xF) << 8) | (((d >> 11) & 1) << 7)
195
+ wr(loc, (w & 0x1FFF07F) | imm)
196
+ elif r_type == R["R_RISCV_RELAX"]:
197
+ pass
198
+ else:
199
+ raise NotImplementedError(f"reloc {RREV.get(r_type, r_type)}")
200
+
201
+ entry = 0
202
+ for sym in syms:
203
+ if sym["name"] == entry_symbol:
204
+ entry = sym_addr(sym["shndx"] and syms.index(sym))
205
+ entry = placed.get(sym["shndx"], base) + sym["value"]
206
+ break
207
+ return mem, entry
208
+
209
+ # ============================================================================
210
+ # neural_rv32: assembler, reference emulator, threshold runtime
211
+ # ============================================================================
212
+ RV32_MODEL = os.path.join(os.path.dirname(os.path.abspath(__file__)),
213
+ "variants", "neural_rv32.safetensors")
214
+ MASK32 = 0xFFFFFFFF
215
+ MMIO_CONSOLE = 0xFF00
216
+
217
+
218
+ def fcvt_s_w_ref(x: int) -> int:
219
+ """Signed int32 -> float32 word, round-to-nearest-even (host float()
220
+ then pack to float32 is exact RNE for the int32 range)."""
221
+ xi = x if (x & MASK32) < 0x80000000 else (x & MASK32) - (1 << 32)
222
+ return struct.unpack("<I", struct.pack("<f", float(xi)))[0]
223
+
224
+
225
+ def fcvt_w_s_ref(w: int) -> int:
226
+ """float32 word -> signed int32, round toward zero (RISC-V rtz),
227
+ saturating; NaN -> INT_MAX, infinities -> INT_MIN/INT_MAX."""
228
+ s, e, f = (w >> 31) & 1, (w >> 23) & 0xFF, w & 0x7FFFFF
229
+ if e == 0xFF:
230
+ return 0x7FFFFFFF if f else (0x80000000 if s else 0x7FFFFFFF)
231
+ E = e - 127
232
+ if E < 0:
233
+ return 0
234
+ mant = (1 << 23) | f
235
+ if E >= 31:
236
+ if s and E == 31 and f == 0:
237
+ return 0x80000000
238
+ return 0x80000000 if s else 0x7FFFFFFF
239
+ mag = (mant << (E - 23)) if E >= 23 else (mant >> (23 - E))
240
+ return ((-mag) & MASK32) if s else (mag & MASK32)
241
+
242
+
243
+ def sext(v: int, bits: int) -> int:
244
+ v &= (1 << bits) - 1
245
+ return v - (1 << bits) if v & (1 << (bits - 1)) else v
246
+
247
+
248
+ def u32(v: int) -> int:
249
+ return v & MASK32
250
+
251
+
252
+ # =============================================================================
253
+ # Assembler
254
+ # =============================================================================
255
+
256
+ R_FUNCT = {"add": (0, 0), "sub": (0, 32), "sll": (1, 0), "slt": (2, 0),
257
+ "sltu": (3, 0), "xor": (4, 0), "srl": (5, 0), "sra": (5, 32),
258
+ "or": (6, 0), "and": (7, 0),
259
+ "mul": (0, 1), "mulh": (1, 1), "mulhsu": (2, 1), "mulhu": (3, 1),
260
+ "div": (4, 1), "divu": (5, 1), "rem": (6, 1), "remu": (7, 1)}
261
+ I_FUNCT = {"addi": 0, "slti": 2, "sltiu": 3, "xori": 4, "ori": 6, "andi": 7}
262
+ B_FUNCT = {"beq": 0, "bne": 1, "blt": 4, "bge": 5, "bltu": 6, "bgeu": 7}
263
+ L_FUNCT = {"lb": 0, "lh": 1, "lw": 2, "lbu": 4, "lhu": 5}
264
+ S_FUNCT = {"sb": 0, "sh": 1, "sw": 2}
265
+
266
+
267
+ class Asm:
268
+ def __init__(self):
269
+ self.words: List = []
270
+ self.labels: Dict[str, int] = {}
271
+ self.data_here = None
272
+
273
+ def label(self, name):
274
+ self.labels[name] = len(self.words) * 4
275
+
276
+ def emit(self, *parts):
277
+ self.words.append(parts)
278
+
279
+ def __getattr__(self, name):
280
+ if name.startswith("__"):
281
+ raise AttributeError(name)
282
+ mn = name.rstrip("_").replace("_", ".")
283
+ def f(*args):
284
+ self.emit(mn, *args)
285
+ return f
286
+
287
+ def word(self, v):
288
+ self.emit(".word", v)
289
+
290
+ def _enc(self, parts, pc):
291
+ mn = parts[0]
292
+ a = parts[1:]
293
+
294
+ def imm(x, pcrel=False):
295
+ if isinstance(x, str):
296
+ t = self.labels[x]
297
+ return t - pc if pcrel else t
298
+ return x
299
+
300
+ if mn == ".word":
301
+ return imm(a[0]) & MASK32
302
+ if mn in R_FUNCT:
303
+ f3, f7 = R_FUNCT[mn]
304
+ return (f7 << 25) | (a[2] << 20) | (a[1] << 15) | (f3 << 12) | (a[0] << 7) | 0x33
305
+ if mn in I_FUNCT:
306
+ return ((imm(a[2]) & 0xFFF) << 20) | (a[1] << 15) | (I_FUNCT[mn] << 12) | (a[0] << 7) | 0x13
307
+ if mn in ("slli", "srli", "srai"):
308
+ f7 = 32 if mn == "srai" else 0
309
+ f3 = 1 if mn == "slli" else 5
310
+ return (f7 << 25) | ((a[2] & 31) << 20) | (a[1] << 15) | (f3 << 12) | (a[0] << 7) | 0x13
311
+ if mn in L_FUNCT:
312
+ return ((imm(a[2]) & 0xFFF) << 20) | (a[1] << 15) | (L_FUNCT[mn] << 12) | (a[0] << 7) | 0x03
313
+ if mn in S_FUNCT:
314
+ i = imm(a[2]) & 0xFFF
315
+ return ((i >> 5) << 25) | (a[0] << 20) | (a[1] << 15) | (S_FUNCT[mn] << 12) | ((i & 31) << 7) | 0x23
316
+ if mn in B_FUNCT:
317
+ i = imm(a[2], pcrel=True) & 0x1FFF
318
+ return (((i >> 12) & 1) << 31) | (((i >> 5) & 0x3F) << 25) | (a[1] << 20) | \
319
+ (a[0] << 15) | (B_FUNCT[mn] << 12) | (((i >> 1) & 0xF) << 8) | (((i >> 11) & 1) << 7) | 0x63
320
+ if mn == "lui":
321
+ return ((imm(a[1]) & 0xFFFFF) << 12) | (a[0] << 7) | 0x37
322
+ if mn == "auipc":
323
+ return ((imm(a[1]) & 0xFFFFF) << 12) | (a[0] << 7) | 0x17
324
+ if mn == "jal":
325
+ i = imm(a[1], pcrel=True) & 0x1FFFFF
326
+ return (((i >> 20) & 1) << 31) | (((i >> 1) & 0x3FF) << 21) | (((i >> 11) & 1) << 20) | \
327
+ (((i >> 12) & 0xFF) << 12) | (a[0] << 7) | 0x6F
328
+ if mn == "jalr":
329
+ return ((imm(a[2]) & 0xFFF) << 20) | (a[1] << 15) | (a[0] << 7) | 0x67
330
+ if mn == "flw":
331
+ return ((imm(a[2]) & 0xFFF) << 20) | (a[1] << 15) | (2 << 12) | (a[0] << 7) | 0x07
332
+ if mn == "fsw":
333
+ i = imm(a[2]) & 0xFFF
334
+ return ((i >> 5) << 25) | (a[0] << 20) | (a[1] << 15) | (2 << 12) | ((i & 31) << 7) | 0x27
335
+ if mn in ("fadd.s", "fsub.s", "fmul.s", "fdiv.s"):
336
+ f7 = {"fadd.s": 0, "fsub.s": 4, "fmul.s": 8, "fdiv.s": 12}[mn]
337
+ return (f7 << 25) | (a[2] << 20) | (a[1] << 15) | (a[0] << 7) | 0x53
338
+ if mn in ("fle.s", "flt.s", "feq.s"):
339
+ f3 = {"fle.s": 0, "flt.s": 1, "feq.s": 2}[mn]
340
+ return (0x50 << 25) | (a[2] << 20) | (a[1] << 15) | (f3 << 12) | (a[0] << 7) | 0x53
341
+ if mn == "fmv.w.x":
342
+ return (0x78 << 25) | (a[1] << 15) | (a[0] << 7) | 0x53
343
+ if mn == "fmv.x.w":
344
+ return (0x70 << 25) | (a[1] << 15) | (a[0] << 7) | 0x53
345
+ if mn == "fcvt.w.s":
346
+ return (0x60 << 25) | (a[1] << 15) | (1 << 12) | (a[0] << 7) | 0x53 # rm=rtz
347
+ if mn == "fcvt.s.w":
348
+ return (0x68 << 25) | (a[1] << 15) | (a[0] << 7) | 0x53
349
+ if mn in ("fsgnj.s", "fsgnjn.s", "fsgnjx.s"):
350
+ f3 = {"fsgnj.s": 0, "fsgnjn.s": 1, "fsgnjx.s": 2}[mn]
351
+ return (0x10 << 25) | (a[2] << 20) | (a[1] << 15) | (f3 << 12) | (a[0] << 7) | 0x53
352
+ if mn == "neur":
353
+ return (a[2] << 20) | (a[1] << 15) | (a[0] << 7) | 0x0B
354
+ if mn == "ecall":
355
+ return 0x73
356
+ if mn == "fence":
357
+ return 0x0F
358
+ raise ValueError(f"unknown mnemonic {mn}")
359
+
360
+ def assemble(self, mem_size=65536) -> List[int]:
361
+ mem = [0] * mem_size
362
+ for idx, parts in enumerate(self.words):
363
+ w = self._enc(parts, idx * 4)
364
+ for b in range(4):
365
+ mem[idx * 4 + b] = (w >> (8 * b)) & 0xFF # little-endian
366
+ return mem
367
+
368
+
369
+ # =============================================================================
370
+ # Reference emulator (architectural contract)
371
+ # =============================================================================
372
+
373
+ class RefRV32:
374
+ @staticmethod
375
+ def step(s):
376
+ if s["halted"]:
377
+ return s
378
+ s = dict(s, mem=s["mem"][:], regs=s["regs"][:], fregs=s["fregs"][:],
379
+ console=s["console"])
380
+ pc = s["pc"]
381
+ w = RefRV32._fetch(s, pc)
382
+ RefRV32._exec(s, pc, w)
383
+ return s
384
+
385
+ @staticmethod
386
+ def _fetch(s, pc):
387
+ m = s["mem"]
388
+ return m[pc] | (m[pc + 1] << 8) | (m[pc + 2] << 16) | (m[pc + 3] << 24)
389
+
390
+ @staticmethod
391
+ def _load(s, addr, width, signed):
392
+ v = 0
393
+ for b in range(width):
394
+ v |= s["mem"][(addr + b) & 0xFFFF] << (8 * b)
395
+ return u32(sext(v, width * 8)) if signed else v
396
+
397
+ @staticmethod
398
+ def _store(s, addr, val, width):
399
+ for b in range(width):
400
+ s["mem"][(addr + b) & 0xFFFF] = (val >> (8 * b)) & 0xFF
401
+ if (addr & 0xFFFF) == MMIO_CONSOLE:
402
+ s["console"] += chr(val & 0xFF)
403
+
404
+ @staticmethod
405
+ def _exec(s, pc, w):
406
+ op = w & 0x7F
407
+ rd = (w >> 7) & 31
408
+ f3 = (w >> 12) & 7
409
+ rs1 = (w >> 15) & 31
410
+ rs2 = (w >> 20) & 31
411
+ f7 = w >> 25
412
+ R = s["regs"]
413
+ FR = s["fregs"]
414
+ a, b = R[rs1], R[rs2]
415
+ nxt = (pc + 4) & MASK32
416
+
417
+ def wr(v):
418
+ if rd:
419
+ R[rd] = u32(v)
420
+
421
+ if op == 0x33: # OP (+ M)
422
+ if f7 == 1:
423
+ sa, sb = sext(a, 32), sext(b, 32)
424
+ if f3 == 0: wr(a * b)
425
+ elif f3 == 1: wr((sa * sb) >> 32)
426
+ elif f3 == 2: wr((sa * b) >> 32)
427
+ elif f3 == 3: wr((a * b) >> 32)
428
+ elif f3 == 4: wr(MASK32 if b == 0 else (0x80000000 if (sa, sb) == (-2**31, -1) else int(abs(sa) // abs(sb)) * (1 if (sa < 0) == (sb < 0) else -1)))
429
+ elif f3 == 5: wr(MASK32 if b == 0 else a // b)
430
+ elif f3 == 6: wr(sa if b == 0 else (0 if (sa, sb) == (-2**31, -1) else sa - (int(abs(sa) // abs(sb)) * (1 if (sa < 0) == (sb < 0) else -1)) * sb))
431
+ elif f3 == 7: wr(a if b == 0 else a % b)
432
+ else:
433
+ sh = b & 31
434
+ if f3 == 0: wr(a - b if f7 == 32 else a + b)
435
+ elif f3 == 1: wr(a << sh)
436
+ elif f3 == 2: wr(1 if sext(a, 32) < sext(b, 32) else 0)
437
+ elif f3 == 3: wr(1 if a < b else 0)
438
+ elif f3 == 4: wr(a ^ b)
439
+ elif f3 == 5: wr(u32(sext(a, 32) >> sh) if f7 == 32 else a >> sh)
440
+ elif f3 == 6: wr(a | b)
441
+ elif f3 == 7: wr(a & b)
442
+ elif op == 0x13: # OP-IMM
443
+ imm = sext(w >> 20, 12)
444
+ sh = (w >> 20) & 31
445
+ if f3 == 0: wr(a + imm)
446
+ elif f3 == 1: wr(a << sh)
447
+ elif f3 == 2: wr(1 if sext(a, 32) < imm else 0)
448
+ elif f3 == 3: wr(1 if a < u32(imm) else 0)
449
+ elif f3 == 4: wr(a ^ u32(imm))
450
+ elif f3 == 5: wr(u32(sext(a, 32) >> sh) if f7 == 32 else a >> sh)
451
+ elif f3 == 6: wr(a | u32(imm))
452
+ elif f3 == 7: wr(a & u32(imm))
453
+ elif op == 0x37: wr(w & 0xFFFFF000) # LUI
454
+ elif op == 0x17: wr(pc + (w & 0xFFFFF000)) # AUIPC
455
+ elif op == 0x6F: # JAL
456
+ imm = sext((((w >> 31) & 1) << 20) | (((w >> 12) & 0xFF) << 12)
457
+ | (((w >> 20) & 1) << 11) | (((w >> 21) & 0x3FF) << 1), 21)
458
+ wr(nxt)
459
+ nxt = u32(pc + imm)
460
+ elif op == 0x67: # JALR
461
+ t = u32(a + sext(w >> 20, 12)) & ~1
462
+ wr(nxt)
463
+ nxt = t
464
+ elif op == 0x63: # BRANCH
465
+ imm = sext((((w >> 31) & 1) << 12) | (((w >> 7) & 1) << 11)
466
+ | (((w >> 25) & 0x3F) << 5) | (((w >> 8) & 0xF) << 1), 13)
467
+ sa, sb = sext(a, 32), sext(b, 32)
468
+ take = [a == b, a != b, None, None, sa < sb, sa >= sb, a < b, a >= b][f3]
469
+ if take:
470
+ nxt = u32(pc + imm)
471
+ elif op == 0x03: # LOAD
472
+ addr = u32(a + sext(w >> 20, 12)) & 0xFFFF
473
+ width = [1, 2, 4, 0, 1, 2][f3]
474
+ wr(RefRV32._load(s, addr, width, f3 < 3))
475
+ elif op == 0x23: # STORE
476
+ imm = sext(((w >> 25) << 5) | ((w >> 7) & 31), 12)
477
+ addr = u32(a + imm) & 0xFFFF
478
+ RefRV32._store(s, addr, b, [1, 2, 4][f3])
479
+ elif op == 0x07: # FLW
480
+ addr = u32(a + sext(w >> 20, 12)) & 0xFFFF
481
+ FR[rd] = RefRV32._load(s, addr, 4, False)
482
+ elif op == 0x27: # FSW
483
+ imm = sext(((w >> 25) << 5) | ((w >> 7) & 31), 12)
484
+ addr = u32(a + imm) & 0xFFFF
485
+ RefRV32._store(s, addr, FR[rs2], 4)
486
+ elif op == 0x53: # OP-FP
487
+ from eval import float_add_oracle, float_mul_oracle, float_div_oracle
488
+ fa, fb = FR[rs1], FR[rs2]
489
+ if f7 == 0x00: FR[rd] = float_add_oracle(fa, fb, 8, 23)
490
+ elif f7 == 0x04: FR[rd] = float_add_oracle(fa, fb ^ 0x80000000, 8, 23)
491
+ elif f7 == 0x08: FR[rd] = float_mul_oracle(fa, fb, 8, 23)
492
+ elif f7 == 0x0C: FR[rd] = float_div_oracle(fa, fb, 8, 23)
493
+ elif f7 == 0x10:
494
+ sa, sb = (fa >> 31) & 1, (fb >> 31) & 1
495
+ so = [sb, 1 - sb, sa ^ sb][f3]
496
+ FR[rd] = (fa & 0x7FFFFFFF) | (so << 31)
497
+ elif f7 == 0x50:
498
+ from eval import float_bits_to_value
499
+ x, y = float_bits_to_value(fa, 8, 23), float_bits_to_value(fb, 8, 23)
500
+ wr(1 if [x <= y, x < y, x == y][f3] else 0)
501
+ elif f7 == 0x60: wr(fcvt_w_s_ref(FR[rs1])) # FCVT.W.S
502
+ elif f7 == 0x68: FR[rd] = fcvt_s_w_ref(R[rs1]) # FCVT.S.W
503
+ elif f7 == 0x78: FR[rd] = a
504
+ elif f7 == 0x70: wr(FR[rs1])
505
+ elif op == 0x0B: # NEUR (custom-0)
506
+ x = a & 0xFF
507
+ pos, neg = b & 0xFF, (b >> 8) & 0xFF
508
+ bias = sext(b >> 16, 5)
509
+ acc = bin(x & pos).count("1") - bin(x & neg).count("1") + bias
510
+ wr(1 if acc >= 0 else 0)
511
+ elif op == 0x73:
512
+ s["halted"] = True
513
+ elif op == 0x0F:
514
+ pass
515
+ else:
516
+ raise ValueError(f"illegal instruction {w:#010x} at {pc:#06x}")
517
+ s["pc"] = nxt
518
+
519
+ @staticmethod
520
+ def run(state, max_cycles=10000):
521
+ s = state
522
+ n = 0
523
+ while not s["halted"] and n < max_cycles:
524
+ s = RefRV32.step(s)
525
+ n += 1
526
+ return s, n
527
+
528
+
529
+ def rv_state(mem):
530
+ return {"pc": 0, "regs": [0] * 32, "fregs": [0] * 32, "mem": mem[:],
531
+ "halted": False, "console": ""}
532
+
533
+
534
+ # =============================================================================
535
+ # Threshold runtime
536
+ # =============================================================================
537
+
538
+ PAIRABLE = (0x33, 0x13, 0x37, 0x17)
539
+
540
+
541
+ class Rv32ThresholdCPU:
542
+ def __init__(self, path=RV32_MODEL):
543
+ self.T = {}
544
+ with safe_open(path, framework="pt") as f:
545
+ for name in f.keys():
546
+ self.T[name] = f.get_tensor(name).float()
547
+ self.registry = load_metadata(path)["signal_registry"]
548
+ self.alu = GenericThresholdALU(self.T, 32)
549
+ self.fpu = {op: NetlistEvaluator(self.T, self.registry, f"float32.{op}")
550
+ for op in ("add", "mul", "div", "cmp")}
551
+ self.pairs_issued = 0
552
+
553
+ # ---- gate helpers -------------------------------------------------
554
+ def _g(self, name, inputs):
555
+ return self.alu._g(name, inputs)
556
+
557
+ def _not1(self, x):
558
+ return self._g("alu.alu32bit.not.bit0", [float(x)])
559
+
560
+ def _bitwise(self, opname, a, b):
561
+ ab = int_to_bits(a, 32)
562
+ bb = int_to_bits(b, 32)
563
+ out = []
564
+ for k in range(32):
565
+ if opname == "xor":
566
+ p = f"alu.alu32bit.xor.bit{k}"
567
+ h1 = self._g(f"{p}.layer1.or", [ab[k], bb[k]])
568
+ h2 = self._g(f"{p}.layer1.nand", [ab[k], bb[k]])
569
+ out.append(self._g(f"{p}.layer2", [h1, h2]))
570
+ else:
571
+ out.append(self._g(f"alu.alu32bit.{opname}.bit{k}", [ab[k], bb[k]]))
572
+ return bits_to_int(out)
573
+
574
+ def _not32(self, a):
575
+ ab = int_to_bits(a, 32)
576
+ return bits_to_int([self._g(f"alu.alu32bit.not.bit{k}", [ab[k]]) for k in range(32)])
577
+
578
+ def _barrel_left(self, val, sh):
579
+ layer_in = int_to_bits(val, 32)
580
+ for layer in range(5):
581
+ amount = 1 << (4 - layer)
582
+ sel = (sh >> (4 - layer)) & 1
583
+ nxt = []
584
+ for bit in range(32):
585
+ p = f"combinational.barrelshifter32.layer{layer}.bit{bit}"
586
+ ns = self._g(f"{p}.not_sel", [float(sel)])
587
+ src = layer_in[bit + amount] if bit + amount < 32 else 0
588
+ a_ = self._g(f"{p}.and_a", [layer_in[bit], ns])
589
+ b_ = self._g(f"{p}.and_b", [float(src), float(sel)])
590
+ nxt.append(self._g(f"{p}.or", [a_, b_]))
591
+ layer_in = nxt
592
+ return bits_to_int(layer_in)
593
+
594
+ def sll(self, a, sh):
595
+ return self._barrel_left(a, sh & 31)
596
+
597
+ def srl(self, a, sh):
598
+ rev = int("{:032b}".format(u32(a))[::-1], 2)
599
+ return int("{:032b}".format(self._barrel_left(rev, sh & 31))[::-1], 2)
600
+
601
+ def sra(self, a, sh):
602
+ sign = (a >> 31) & 1
603
+ x1 = self.srl(a, sh)
604
+ x2 = self._not32(self.srl(self._not32(a), sh))
605
+ o1, o2 = int_to_bits(x1, 32), int_to_bits(x2, 32)
606
+ out = []
607
+ for k in range(32):
608
+ p = f"rv32.sra_mux.bit{k}"
609
+ ns = self._g(f"{p}.not_sel", [float(sign)])
610
+ a_ = self._g(f"{p}.and_a", [o1[k], ns])
611
+ b_ = self._g(f"{p}.and_b", [o2[k], float(sign)])
612
+ out.append(self._g(f"{p}.or", [a_, b_]))
613
+ return bits_to_int(out)
614
+
615
+ def add32(self, a, b):
616
+ return self.alu.add_n(a, b, 32)[0]
617
+
618
+ def sub32(self, a, b):
619
+ return self.alu.sub_n(a, b, 32)[0]
620
+
621
+ def ucmp(self, a, b, kind):
622
+ return self.alu.cmp_n(a, b, kind, 32)
623
+
624
+ def scmp_lt(self, a, b):
625
+ fa = (a & 0x7FFFFFFF) | (self._not1((a >> 31) & 1) << 31)
626
+ fb = (b & 0x7FFFFFFF) | (self._not1((b >> 31) & 1) << 31)
627
+ return self.ucmp(fa, fb, "lessthan")
628
+
629
+ def mulhu(self, a, b):
630
+ """Full 64-bit product high word via pp gates + rv32.mulacc chains."""
631
+ ab = int_to_bits(a, 32)
632
+ bb = int_to_bits(b, 32)
633
+
634
+ def pp(i_lsb, j_lsb): # LSB-first indices onto MSB-first gate names
635
+ return self._g(f"alu.alu32bit.mul.pp.a{31 - i_lsb}b{31 - j_lsb}",
636
+ [ab[31 - i_lsb], bb[31 - j_lsb]])
637
+
638
+ acc = [pp(k, 0) if k <= 31 else 0.0 for k in range(64)]
639
+ for t in range(31):
640
+ row = t + 1
641
+ carry = 0.0
642
+ nxt = []
643
+ for k in range(64):
644
+ b_in = pp(k - row, row) if row <= k <= row + 31 else 0.0
645
+ s, carry = self.alu._fa(f"rv32.mulacc.s{t}.fa{k}", acc[k], b_in, carry)
646
+ nxt.append(s)
647
+ acc = nxt
648
+ return bits_to_int(list(reversed([int(x) for x in acc[32:]])))
649
+
650
+ def _cond_mask(self, val, cond_bit):
651
+ vb = int_to_bits(val, 32)
652
+ return bits_to_int([self._g(f"alu.alu32bit.and.bit{k}", [vb[k], float(cond_bit)])
653
+ for k in range(32)])
654
+
655
+ def mulh(self, a, b, sign_a, sign_b):
656
+ h = self.mulhu(a, b)
657
+ if sign_a:
658
+ h = self.sub32(h, self._cond_mask(b, (a >> 31) & 1))
659
+ if sign_b:
660
+ h = self.sub32(h, self._cond_mask(a, (b >> 31) & 1))
661
+ return h
662
+
663
+ def divu(self, a, b):
664
+ if b == 0:
665
+ return MASK32, a
666
+ a_bits = int_to_bits(a, 32)
667
+ q = 0
668
+ rem = 0
669
+ for stage in range(32):
670
+ rem = ((rem << 1) | a_bits[stage]) & MASK32
671
+ prefix = f"alu.alu32bit.div.stage{stage}"
672
+ finals = {"gt": f"{prefix}.cmp_bc.gt", "lt": f"{prefix}.cmp_bc.lt",
673
+ "eq": f"{prefix}.cmp_bc.eq", "ge": f"{prefix}.cmp",
674
+ "le": f"{prefix}.cmp_bc.le"}
675
+ ge = self.alu._cmp_bit_cascade(f"{prefix}.cmp_bc", finals,
676
+ int_to_bits(rem, 32), int_to_bits(b, 32), "ge")
677
+ if ge:
678
+ rem = self.sub32(rem, b)
679
+ q = (q << 1) | 1
680
+ else:
681
+ q <<= 1
682
+ return u32(q), rem
683
+
684
+ def neg32(self, a):
685
+ ab = int_to_bits(a, 32)
686
+ nb = [self._g(f"alu.alu32bit.neg.not.bit{k}", [ab[k]]) for k in range(32)]
687
+ carry = 1.0
688
+ out = [0] * 32
689
+ for k in range(31, -1, -1):
690
+ p = f"alu.alu32bit.neg.inc.bit{k}"
691
+ h1 = self._g(f"{p}.xor.layer1.or", [nb[k], carry])
692
+ h2 = self._g(f"{p}.xor.layer1.nand", [nb[k], carry])
693
+ out[k] = self._g(f"{p}.xor.layer2", [h1, h2])
694
+ carry = self._g(f"{p}.carry", [nb[k], carry])
695
+ return bits_to_int(out)
696
+
697
+ def divs(self, a, b):
698
+ if b == 0:
699
+ return MASK32, a
700
+ if a == 0x80000000 and b == MASK32:
701
+ return 0x80000000, 0
702
+ na, nb_ = (a >> 31) & 1, (b >> 31) & 1
703
+ ua = self.neg32(a) if na else a
704
+ ub = self.neg32(b) if nb_ else b
705
+ q, r = self.divu(ua, ub)
706
+ if na != nb_:
707
+ q = self.neg32(q)
708
+ if na:
709
+ r = self.neg32(r)
710
+ return q, r
711
+
712
+ def priority_encode(self, v):
713
+ """MSB-first position (0..31) of the highest set bit, via the
714
+ combinational.priorityencoder32 gates. v must be nonzero."""
715
+ vb = int_to_bits(v, 32) # MSB-first, bit index 0 = MSB
716
+ any_higher = [0]
717
+ for pos in range(1, 32):
718
+ any_higher.append(self._g(f"combinational.priorityencoder32.any_higher{pos}",
719
+ [float(vb[i]) for i in range(pos)]))
720
+ is_high = [float(vb[0])]
721
+ for pos in range(1, 32):
722
+ nh = self._g(f"combinational.priorityencoder32.is_highest{pos}.not_higher",
723
+ [float(any_higher[pos])])
724
+ is_high.append(self._g(f"combinational.priorityencoder32.is_highest{pos}.and",
725
+ [float(vb[pos]), nh]))
726
+ out = 0
727
+ for bit in range(5):
728
+ rel = [is_high[pos] for pos in range(32) if (pos >> bit) & 1]
729
+ out |= self._g(f"combinational.priorityencoder32.out{bit}", [float(x) for x in rel]) << bit
730
+ return out
731
+
732
+ def fcvt_s_w(self, x):
733
+ """Signed int32 -> float32, round-to-nearest-even. Sign, leading-one
734
+ detect (priority encoder), left-normalize (barrel shifter), round."""
735
+ if x == 0:
736
+ return 0
737
+ s = (x >> 31) & 1
738
+ u = self.neg32(x) if s else x
739
+ p = self.priority_encode(u) # leading one at MSB-first pos p
740
+ shifted = self.sll(u, p) # leading one now at bit 31
741
+ exp = 127 + (31 - p)
742
+ frac = (shifted >> 8) & 0x7FFFFF
743
+ guard = (shifted >> 7) & 1
744
+ below = shifted & 0x7F
745
+ if guard and ((frac & 1) or below):
746
+ frac += 1
747
+ if frac > 0x7FFFFF:
748
+ frac = 0
749
+ exp += 1
750
+ return (s << 31) | (exp << 23) | frac
751
+
752
+ def fcvt_w_s(self, w):
753
+ """float32 -> signed int32, round toward zero (RISC-V rtz), saturating.
754
+ NaN -> INT_MAX; overflow saturates to INT_MIN/INT_MAX."""
755
+ s, e, f = (w >> 31) & 1, (w >> 23) & 0xFF, w & 0x7FFFFF
756
+ if e == 0xFF:
757
+ if f:
758
+ return 0x7FFFFFFF # NaN
759
+ return 0x80000000 if s else 0x7FFFFFFF # +-inf
760
+ E = e - 127
761
+ if E < 0:
762
+ return 0
763
+ mant = (1 << 23) | f
764
+ if E >= 31:
765
+ if s and E == 31 and f == 0:
766
+ return 0x80000000 # exactly -2^31
767
+ return 0x80000000 if s else 0x7FFFFFFF
768
+ if E >= 23:
769
+ mag = self.sll(mant, E - 23)
770
+ else:
771
+ mag = self.srl(mant, 23 - E)
772
+ return self.neg32(mag) if s else (mag & MASK32)
773
+
774
+ def neur(self, x_reg, w_reg):
775
+ x = int_to_bits(x_reg & 0xFF, 8) # MSB-first bits of low byte
776
+ pos = int_to_bits(w_reg & 0xFF, 8)
777
+ neg = int_to_bits((w_reg >> 8) & 0xFF, 8)
778
+ p_bits = [self._g(f"rv32.neur.pos.bit{i}", [x[7 - i], pos[7 - i]]) for i in range(8)]
779
+ n_bits = [self._g(f"rv32.neur.neg.bit{i}", [x[7 - i], neg[7 - i]]) for i in range(8)]
780
+
781
+ def pcount(side, bits8):
782
+ fa = lambda t, a_, b_, c_: self.alu._fa(f"rv32.neur.{side}.fa{t}", a_, b_, c_)
783
+ s0, c0 = fa(0, bits8[0], bits8[1], bits8[2])
784
+ s1, c1 = fa(1, bits8[3], bits8[4], bits8[5])
785
+ s2, c2 = fa(2, bits8[6], bits8[7], 0.0)
786
+ S, C2 = fa(3, s0, s1, s2)
787
+ T, C4 = fa(4, c0, c1, c2)
788
+ U, C5 = fa(5, C2, T, 0.0)
789
+ V, W = fa(6, C4, C5, 0.0)
790
+ return [S, U, V, W] # LSB-first count bits
791
+
792
+ P = pcount("pcnt", p_bits)
793
+ N = pcount("ncnt", n_bits)
794
+ nnot = [self._g(f"rv32.neur.nnot.bit{k}", [N[k] if k < 4 else 0.0]) for k in range(6)]
795
+ carry = 1.0
796
+ diff = []
797
+ for k in range(6):
798
+ a_ = P[k] if k < 4 else 0.0
799
+ s, carry = self.alu._fa(f"rv32.neur.diff.fa{k}", a_, nnot[k], carry)
800
+ diff.append(s)
801
+ bias = (w_reg >> 16) & 0x1F
802
+ bias6 = [(bias >> k) & 1 if k < 4 else (bias >> 4) & 1 for k in range(6)] # sext5
803
+ carry = 0.0
804
+ act = []
805
+ for k in range(6):
806
+ s, carry = self.alu._fa(f"rv32.neur.act.fa{k}", diff[k], float(bias6[k]), carry)
807
+ act.append(s)
808
+ return int(self._g("rv32.neur.out", [act[5]]))
809
+
810
+ def hazard_eq(self, pair, i, j):
811
+ ib = int_to_bits(i, 5)
812
+ jb = int_to_bits(j, 5)
813
+ eqs = []
814
+ for k in range(5):
815
+ p = f"rv32.hazard.{pair}.bit{k}"
816
+ h_and = self._g(f"{p}.eq.layer1.and", [ib[k], jb[k]])
817
+ h_nor = self._g(f"{p}.eq.layer1.nor", [ib[k], jb[k]])
818
+ eqs.append(self._g(f"{p}.eq", [h_and, h_nor]))
819
+ return int(self._g(f"rv32.hazard.{pair}.all", eqs))
820
+
821
+ # ---- memory through the packed circuits ---------------------------
822
+ def _addr_decode(self, addr):
823
+ bits = torch.tensor(int_to_bits(addr, 16), dtype=torch.float32)
824
+ return heaviside((self.T["memory.addr_decode.weight"] * bits).sum(dim=1)
825
+ + self.T["memory.addr_decode.bias"])
826
+
827
+ def mem_read_byte(self, mem_bits, addr):
828
+ sel = self._addr_decode(addr & 0xFFFF)
829
+ and_w = self.T["memory.read.and.weight"]
830
+ and_b = self.T["memory.read.and.bias"]
831
+ or_w = self.T["memory.read.or.weight"]
832
+ or_b = self.T["memory.read.or.bias"]
833
+ v = 0
834
+ for bit in range(8):
835
+ inp = torch.stack([mem_bits[:, bit], sel], dim=1)
836
+ and_out = heaviside((inp * and_w[bit]).sum(dim=1) + and_b[bit])
837
+ v = (v << 1) | int(heaviside((and_out * or_w[bit]).sum() + or_b[bit]).item())
838
+ return v
839
+
840
+ def mem_write_byte(self, mem_bits, addr, value):
841
+ sel = self._addr_decode(addr & 0xFFFF)
842
+ data_bits = torch.tensor(int_to_bits(value & 0xFF, 8), dtype=torch.float32)
843
+ we = torch.ones_like(sel)
844
+ write_sel = heaviside((torch.stack([sel, we], dim=1)
845
+ * self.T["memory.write.sel.weight"]).sum(dim=1)
846
+ + self.T["memory.write.sel.bias"])
847
+ nsel = heaviside(write_sel * self.T["memory.write.nsel.weight"].squeeze(1)
848
+ + self.T["memory.write.nsel.bias"])
849
+ for bit in range(8):
850
+ old = mem_bits[:, bit]
851
+ dbit = data_bits[bit].expand(65536)
852
+ a_old = heaviside((torch.stack([old, nsel], dim=1)
853
+ * self.T["memory.write.and_old.weight"][:, bit]).sum(dim=1)
854
+ + self.T["memory.write.and_old.bias"][:, bit])
855
+ a_new = heaviside((torch.stack([dbit, write_sel], dim=1)
856
+ * self.T["memory.write.and_new.weight"][:, bit]).sum(dim=1)
857
+ + self.T["memory.write.and_new.bias"][:, bit])
858
+ mem_bits[:, bit] = heaviside((torch.stack([a_old, a_new], dim=1)
859
+ * self.T["memory.write.or.weight"][:, bit]).sum(dim=1)
860
+ + self.T["memory.write.or.bias"][:, bit])
861
+
862
+ # ---- execution ----------------------------------------------------
863
+ def _fetch32(self, mem_bits, pc):
864
+ return sum(self.mem_read_byte(mem_bits, pc + b) << (8 * b) for b in range(4))
865
+
866
+ def _load(self, s, addr, width, signed):
867
+ v = sum(self.mem_read_byte(s["_mem_bits"], (addr + b) & 0xFFFF) << (8 * b)
868
+ for b in range(width))
869
+ return u32(sext(v, width * 8)) if signed else v
870
+
871
+ def _store(self, s, addr, val, width):
872
+ for b in range(width):
873
+ byte = (val >> (8 * b)) & 0xFF
874
+ self.mem_write_byte(s["_mem_bits"], (addr + b) & 0xFFFF, byte)
875
+ s["mem"][(addr + b) & 0xFFFF] = byte
876
+ if (addr & 0xFFFF) == MMIO_CONSOLE:
877
+ s["console"] += chr(val & 0xFF)
878
+
879
+ def _exec_alu_class(self, s, pc, w):
880
+ """Execute one OP / OP-IMM / LUI / AUIPC instruction; returns (rd, value)."""
881
+ op = w & 0x7F
882
+ rd = (w >> 7) & 31
883
+ f3 = (w >> 12) & 7
884
+ rs1 = (w >> 15) & 31
885
+ rs2 = (w >> 20) & 31
886
+ f7 = w >> 25
887
+ R = s["regs"]
888
+ a = R[rs1]
889
+ if op == 0x37:
890
+ return rd, w & 0xFFFFF000
891
+ if op == 0x17:
892
+ return rd, self.add32(pc, w & 0xFFFFF000)
893
+ if op == 0x13:
894
+ b = u32(sext(w >> 20, 12))
895
+ sh = (w >> 20) & 31
896
+ imm_mode = True
897
+ else:
898
+ b = R[rs2]
899
+ sh = b & 31
900
+ imm_mode = False
901
+ if f3 == 0:
902
+ if not imm_mode and f7 == 32:
903
+ return rd, self.sub32(a, b)
904
+ return rd, self.add32(a, b)
905
+ if f3 == 1:
906
+ return rd, self.sll(a, sh)
907
+ if f3 == 2:
908
+ return rd, self.scmp_lt(a, b)
909
+ if f3 == 3:
910
+ return rd, self.ucmp(a, b, "lessthan")
911
+ if f3 == 4:
912
+ return rd, self._bitwise("xor", a, b)
913
+ if f3 == 5:
914
+ return rd, (self.sra(a, sh) if f7 == 32 else self.srl(a, sh))
915
+ if f3 == 6:
916
+ return rd, self._bitwise("or", a, b)
917
+ return rd, self._bitwise("and", a, b)
918
+
919
+ def step(self, s):
920
+ if s["halted"]:
921
+ return s
922
+ s = dict(s, mem=s["mem"][:], regs=s["regs"][:], fregs=s["fregs"][:],
923
+ console=s["console"], _mem_bits=s["_mem_bits"].clone())
924
+ pc = s["pc"]
925
+ w = self._fetch32(s["_mem_bits"], pc)
926
+ op = w & 0x7F
927
+
928
+ def plain_alu(word):
929
+ return (word & 0x7F) in PAIRABLE and not \
930
+ ((word & 0x7F) == 0x33 and (word >> 25) == 1)
931
+
932
+ # Dual issue: both slots plain ALU-class and hazard-free by gates.
933
+ if plain_alu(w):
934
+ w2 = self._fetch32(s["_mem_bits"], (pc + 4) & 0xFFFF)
935
+ if plain_alu(w2):
936
+ rd1 = (w >> 7) & 31
937
+ rs1b = (w2 >> 15) & 31
938
+ rs2b = (w2 >> 20) & 31
939
+ rd2 = (w2 >> 7) & 31
940
+ uses_rs1 = (w2 & 0x7F) in (0x33, 0x13)
941
+ uses_rs2 = (w2 & 0x7F) == 0x33
942
+ raw = (uses_rs1 and self.hazard_eq("rd1_rs1b", rd1, rs1b)) or \
943
+ (uses_rs2 and self.hazard_eq("rd1_rs2b", rd1, rs2b))
944
+ waw = self.hazard_eq("rd1_rd2", rd1, rd2)
945
+ if rd1 != 0 and not raw and not waw:
946
+ d1, v1 = self._exec_alu_class(s, pc, w)
947
+ d2, v2 = self._exec_alu_class(s, (pc + 4) & MASK32, w2)
948
+ if d1:
949
+ s["regs"][d1] = u32(v1)
950
+ if d2:
951
+ s["regs"][d2] = u32(v2)
952
+ s["pc"] = (pc + 8) & MASK32
953
+ self.pairs_issued += 1
954
+ return s
955
+
956
+ rd = (w >> 7) & 31
957
+ f3 = (w >> 12) & 7
958
+ rs1 = (w >> 15) & 31
959
+ rs2 = (w >> 20) & 31
960
+ f7 = w >> 25
961
+ R = s["regs"]
962
+ FR = s["fregs"]
963
+ a, b = R[rs1], R[rs2]
964
+ nxt = (pc + 4) & MASK32
965
+
966
+ def wr(v):
967
+ if rd:
968
+ R[rd] = u32(v)
969
+
970
+ if plain_alu(w):
971
+ d, v = self._exec_alu_class(s, pc, w)
972
+ if d:
973
+ R[d] = u32(v)
974
+ if op == 0x33 and f7 == 1:
975
+ if f3 == 0:
976
+ wr(self.alu.mul_n(a, b, 32))
977
+ elif f3 == 1:
978
+ wr(self.mulh(a, b, True, True))
979
+ elif f3 == 2:
980
+ wr(self.mulh(a, b, True, False))
981
+ elif f3 == 3:
982
+ wr(self.mulhu(a, b))
983
+ elif f3 == 4:
984
+ wr(self.divs(a, b)[0])
985
+ elif f3 == 5:
986
+ wr(self.divu(a, b)[0])
987
+ elif f3 == 6:
988
+ wr(self.divs(a, b)[1])
989
+ elif f3 == 7:
990
+ wr(self.divu(a, b)[1])
991
+ elif op == 0x6F:
992
+ imm = sext((((w >> 31) & 1) << 20) | (((w >> 12) & 0xFF) << 12)
993
+ | (((w >> 20) & 1) << 11) | (((w >> 21) & 0x3FF) << 1), 21)
994
+ wr(nxt)
995
+ nxt = self.add32(pc, u32(imm))
996
+ elif op == 0x67:
997
+ t = self.add32(a, u32(sext(w >> 20, 12))) & ~1
998
+ wr(nxt)
999
+ nxt = t
1000
+ elif op == 0x63:
1001
+ imm = sext((((w >> 31) & 1) << 12) | (((w >> 7) & 1) << 11)
1002
+ | (((w >> 25) & 0x3F) << 5) | (((w >> 8) & 0xF) << 1), 13)
1003
+ eq = self.ucmp(a, b, "eq")
1004
+ sl = self.scmp_lt(a, b)
1005
+ ul = self.ucmp(a, b, "lessthan")
1006
+ take = [eq, 1 - eq, 0, 0, sl, 1 - sl, ul, 1 - ul][f3]
1007
+ if take:
1008
+ nxt = self.add32(pc, u32(imm))
1009
+ elif op == 0x03:
1010
+ addr = self.add32(a, u32(sext(w >> 20, 12))) & 0xFFFF
1011
+ width = [1, 2, 4, 0, 1, 2][f3]
1012
+ wr(self._load(s, addr, width, f3 < 3))
1013
+ elif op == 0x23:
1014
+ imm = sext(((w >> 25) << 5) | ((w >> 7) & 31), 12)
1015
+ addr = self.add32(a, u32(imm)) & 0xFFFF
1016
+ self._store(s, addr, b, [1, 2, 4][f3])
1017
+ elif op == 0x07:
1018
+ addr = self.add32(a, u32(sext(w >> 20, 12))) & 0xFFFF
1019
+ FR[rd] = self._load(s, addr, 4, False)
1020
+ elif op == 0x27:
1021
+ imm = sext(((w >> 25) << 5) | ((w >> 7) & 31), 12)
1022
+ addr = self.add32(a, u32(imm)) & 0xFFFF
1023
+ self._store(s, addr, FR[rs2], 4)
1024
+ elif op == 0x53:
1025
+ fa, fb = FR[rs1], FR[rs2]
1026
+ if f7 in (0x00, 0x04, 0x08, 0x0C):
1027
+ circ = {0x00: "add", 0x04: "add", 0x08: "mul", 0x0C: "div"}[f7]
1028
+ if f7 == 0x04:
1029
+ fb ^= 0x80000000 # subtraction: flip the addend's sign
1030
+ FR[rd] = self._fpu_word(circ, fa, fb)
1031
+ elif f7 == 0x50:
1032
+ out = self.fpu["cmp"].run(self._fpu_ext(fa, fb))
1033
+ gate = {0: "le", 1: "lt", 2: "eq"}[f3]
1034
+ wr(int(out[f"float32.cmp.{gate}.result"][0, 0].item()))
1035
+ elif f7 == 0x10: # FSGNJ family: result = {sign, rs1[30:0]}
1036
+ sa = (FR[rs1] >> 31) & 1
1037
+ sb = (FR[rs2] >> 31) & 1
1038
+ if f3 == 0:
1039
+ s_out = sb # FSGNJ.S
1040
+ elif f3 == 1:
1041
+ s_out = self._not1(sb) # FSGNJN.S
1042
+ else:
1043
+ s_out = self._bitwise("xor", sa << 31, sb << 31) >> 31 # FSGNJX.S
1044
+ FR[rd] = (FR[rs1] & 0x7FFFFFFF) | (s_out << 31)
1045
+ elif f7 == 0x60: # FCVT.W.S: float -> signed int, gate-routed
1046
+ wr(self.fcvt_w_s(FR[rs1]))
1047
+ elif f7 == 0x68: # FCVT.S.W: signed int -> float, gate-routed
1048
+ FR[rd] = self.fcvt_s_w(R[rs1])
1049
+ elif f7 == 0x78:
1050
+ FR[rd] = a
1051
+ elif f7 == 0x70:
1052
+ wr(FR[rs1])
1053
+ elif op == 0x0B:
1054
+ wr(self.neur(a, b))
1055
+ elif op == 0x73:
1056
+ s["halted"] = True
1057
+ elif op == 0x0F:
1058
+ pass
1059
+ elif op not in PAIRABLE:
1060
+ raise ValueError(f"illegal instruction {w:#010x} at {pc:#06x}")
1061
+ s["pc"] = nxt
1062
+ return s
1063
+
1064
+ def _fpu_ext(self, fa, fb):
1065
+ ext = {}
1066
+ for i in range(32):
1067
+ ext[f"$a[{i}]"] = (fa >> (31 - i)) & 1
1068
+ ext[f"$b[{i}]"] = (fb >> (31 - i)) & 1
1069
+ return ext
1070
+
1071
+ def _fpu_word(self, circ, fa, fb):
1072
+ out = self.fpu[circ].run(self._fpu_ext(fa, fb))
1073
+ p = f"float32.{circ}"
1074
+ v = int(out[f"{p}.sign_out"][0, 0].item()) << 31
1075
+ for k in range(8):
1076
+ v |= int(out[f"{p}.exp_out.bit{k}"][0, 0].item()) << (23 + k)
1077
+ for k in range(23):
1078
+ v |= int(out[f"{p}.frac_out.bit{k}"][0, 0].item()) << k
1079
+ return v
1080
+
1081
+ def run(self, state, max_cycles=400):
1082
+ s = dict(state)
1083
+ s["_mem_bits"] = torch.tensor(
1084
+ [int_to_bits(byte, 8) for byte in s["mem"]], dtype=torch.float32)
1085
+ n = 0
1086
+ while not s["halted"] and n < max_cycles:
1087
+ s = self.step(s)
1088
+ n += 1
1089
+ s.pop("_mem_bits", None)
1090
+ return s, n
1091
+
1092
+ # ============================================================================
1093
+ # neural_subleq8: threshold runtime, reference, assembler
1094
+ # ============================================================================
1095
+ SUBLEQ_MODEL = os.path.join(os.path.dirname(os.path.abspath(__file__)),
1096
+ "variants", "neural_subleq8.safetensors")
1097
+ HALT_PC = 0xFF
1098
+
1099
+
1100
+ def load_tensors(path):
1101
+ out = {}
1102
+ with safe_open(path, framework="pt") as f:
1103
+ for name in f.keys():
1104
+ out[name] = f.get_tensor(name).float()
1105
+ return out
1106
+
1107
+
1108
+ class SubleqThresholdCPU:
1109
+ """SUBLEQ-8 runtime: memory traffic through the packed memory circuits,
1110
+ the subtract/branch datapath through the subleq gates via the shipped
1111
+ netlist. Instruction sequencing (three consecutive fetches) is wiring."""
1112
+
1113
+ def __init__(self, tensors, registry):
1114
+ self.T = tensors
1115
+ self.ne = NetlistEvaluator(tensors, registry, "subleq")
1116
+
1117
+ def _addr_decode(self, addr):
1118
+ bits = torch.tensor([(addr >> (7 - i)) & 1 for i in range(8)],
1119
+ dtype=torch.float32)
1120
+ w = self.T["memory.addr_decode.weight"]
1121
+ b = self.T["memory.addr_decode.bias"]
1122
+ return heaviside((w * bits).sum(dim=1) + b)
1123
+
1124
+ def mem_read(self, mem, addr):
1125
+ sel = self._addr_decode(addr)
1126
+ mem_bits = torch.tensor([[(byte >> (7 - i)) & 1 for i in range(8)]
1127
+ for byte in mem], dtype=torch.float32)
1128
+ and_w = self.T["memory.read.and.weight"]
1129
+ and_b = self.T["memory.read.and.bias"]
1130
+ or_w = self.T["memory.read.or.weight"]
1131
+ or_b = self.T["memory.read.or.bias"]
1132
+ out = 0
1133
+ for bit in range(8):
1134
+ inp = torch.stack([mem_bits[:, bit], sel], dim=1)
1135
+ and_out = heaviside((inp * and_w[bit]).sum(dim=1) + and_b[bit])
1136
+ out = (out << 1) | int(heaviside((and_out * or_w[bit]).sum() + or_b[bit]).item())
1137
+ return out
1138
+
1139
+ def mem_write(self, mem, addr, value):
1140
+ sel = self._addr_decode(addr)
1141
+ data_bits = torch.tensor([(value >> (7 - i)) & 1 for i in range(8)],
1142
+ dtype=torch.float32)
1143
+ mem_bits = torch.tensor([[(byte >> (7 - i)) & 1 for i in range(8)]
1144
+ for byte in mem], dtype=torch.float32)
1145
+ sel_w = self.T["memory.write.sel.weight"]
1146
+ sel_b = self.T["memory.write.sel.bias"]
1147
+ nsel_w = self.T["memory.write.nsel.weight"].squeeze(1)
1148
+ nsel_b = self.T["memory.write.nsel.bias"]
1149
+ we = torch.ones_like(sel)
1150
+ write_sel = heaviside((torch.stack([sel, we], dim=1) * sel_w).sum(dim=1) + sel_b)
1151
+ nsel = heaviside(write_sel * nsel_w + nsel_b)
1152
+ ao_w = self.T["memory.write.and_old.weight"]
1153
+ ao_b = self.T["memory.write.and_old.bias"]
1154
+ an_w = self.T["memory.write.and_new.weight"]
1155
+ an_b = self.T["memory.write.and_new.bias"]
1156
+ or_w = self.T["memory.write.or.weight"]
1157
+ or_b = self.T["memory.write.or.bias"]
1158
+ for bit in range(8):
1159
+ old = mem_bits[:, bit]
1160
+ dbit = data_bits[bit].expand(256)
1161
+ a_old = heaviside((torch.stack([old, nsel], dim=1) * ao_w[:, bit]).sum(dim=1) + ao_b[:, bit])
1162
+ a_new = heaviside((torch.stack([dbit, write_sel], dim=1) * an_w[:, bit]).sum(dim=1) + an_b[:, bit])
1163
+ mem_bits[:, bit] = heaviside(
1164
+ (torch.stack([a_old, a_new], dim=1) * or_w[:, bit]).sum(dim=1) + or_b[:, bit])
1165
+ return [int(sum(int(mem_bits[i, b].item()) << (7 - b) for b in range(8)))
1166
+ for i in range(256)]
1167
+
1168
+ def step(self, state):
1169
+ if state["halted"]:
1170
+ return state
1171
+ s = dict(state)
1172
+ s["mem"] = state["mem"][:]
1173
+ pc = s["pc"]
1174
+ A = self.mem_read(s["mem"], pc)
1175
+ B = self.mem_read(s["mem"], (pc + 1) & 0xFF)
1176
+ C = self.mem_read(s["mem"], (pc + 2) & 0xFF)
1177
+ a = self.mem_read(s["mem"], A)
1178
+ b = self.mem_read(s["mem"], B)
1179
+ ext = {}
1180
+ for k in range(8):
1181
+ ext[f"$a[{k}]"] = (a >> k) & 1
1182
+ ext[f"$b[{k}]"] = (b >> k) & 1
1183
+ ext[f"$pc[{k}]"] = (pc >> k) & 1
1184
+ ext[f"$c[{k}]"] = (C >> k) & 1
1185
+ out = self.ne.run(ext)
1186
+ r = sum(int(out[f"subleq.sub.fa{k}.ha2.sum.layer2"][0, 0].item()) << k
1187
+ for k in range(8))
1188
+ next_pc = sum(int(out[f"subleq.pc_mux.bit{k}.or"][0, 0].item()) << k
1189
+ for k in range(8))
1190
+ s["mem"] = self.mem_write(s["mem"], B, r)
1191
+ s["pc"] = next_pc
1192
+ if next_pc == HALT_PC:
1193
+ s["halted"] = True
1194
+ return s
1195
+
1196
+ def run(self, state, max_cycles=200):
1197
+ s = state
1198
+ n = 0
1199
+ while not s["halted"] and n < max_cycles:
1200
+ s = self.step(s)
1201
+ n += 1
1202
+ return s, n
1203
+
1204
+
1205
+ def ref_step(state):
1206
+ if state["halted"]:
1207
+ return state
1208
+ s = dict(state)
1209
+ s["mem"] = state["mem"][:]
1210
+ pc = s["pc"]
1211
+ A, B, C = s["mem"][pc], s["mem"][(pc + 1) & 0xFF], s["mem"][(pc + 2) & 0xFF]
1212
+ r = (s["mem"][B] - s["mem"][A]) & 0xFF
1213
+ s["mem"][B] = r
1214
+ s["pc"] = C if (r == 0 or r >= 0x80) else (pc + 3) & 0xFF
1215
+ if s["pc"] == HALT_PC:
1216
+ s["halted"] = True
1217
+ return s
1218
+
1219
+
1220
+ def sq(*triples):
1221
+ """Assemble SUBLEQ triples into a 256-byte image."""
1222
+ mem = [0] * 256
1223
+ i = 0
1224
+ for t in triples:
1225
+ for v in t:
1226
+ mem[i] = v & 0xFF
1227
+ i += 1
1228
+ return mem
1229
+
1230
+
1231
+ def subleq_test() -> int:
1232
+ T = load_tensors(SUBLEQ_MODEL)
1233
+ reg = load_metadata(SUBLEQ_MODEL)["signal_registry"]
1234
+
1235
+ # 1. Exhaustive datapath: all 65,536 (a, b) pairs through the shipped
1236
+ # wiring in one batched netlist evaluation.
1237
+ ne = NetlistEvaluator(T, reg, "subleq")
1238
+ a_v = torch.arange(65536, dtype=torch.long) % 256
1239
+ b_v = torch.arange(65536, dtype=torch.long) // 256
1240
+ ext = {}
1241
+ for k in range(8):
1242
+ ext[f"$a[{k}]"] = ((a_v >> k) & 1).float()
1243
+ ext[f"$b[{k}]"] = ((b_v >> k) & 1).float()
1244
+ ext[f"$pc[{k}]"] = torch.zeros(65536)
1245
+ ext[f"$c[{k}]"] = torch.zeros(65536)
1246
+ out = ne.run(ext)
1247
+ r = torch.zeros(65536, dtype=torch.float64)
1248
+ for k in range(8):
1249
+ r += out[f"subleq.sub.fa{k}.ha2.sum.layer2"][:, 0].double() * (1 << k)
1250
+ exp_r = ((b_v - a_v) & 0xFF).double()
1251
+ leq = out["subleq.leq"][:, 0].double()
1252
+ exp_leq = ((exp_r == 0) | (exp_r >= 128)).double()
1253
+ ok_r = int((r == exp_r).sum().item())
1254
+ ok_l = int((leq == exp_leq).sum().item())
1255
+ print(f"datapath sub exhaustive: {ok_r}/65536")
1256
+ print(f"datapath leq exhaustive: {ok_l}/65536")
1257
+ if ok_r != 65536 or ok_l != 65536:
1258
+ return 1
1259
+
1260
+ # 2. PC mux + increment across representative PC/target/leq settings.
1261
+ pcs = torch.tensor([0, 3, 5, 100, 250, 252], dtype=torch.long).repeat_interleave(2)
1262
+ cs = torch.tensor([7, 7, 42, 42, 255, 255, 0, 0, 9, 9, 30, 30], dtype=torch.long)
1263
+ a2 = torch.tensor([0, 1] * 6, dtype=torch.long) # a=0,b=0 -> leq; a=1,b=0 -> r=255 leq... use b to vary
1264
+ b2 = torch.tensor([0, 1] * 6, dtype=torch.long) # (0,0)->r0 leq taken; (1,1)->r0 leq taken
1265
+ # craft: even rows leq taken (a=b), odd rows not taken (b-a=1 -> r=1 >0)
1266
+ a2 = torch.tensor([5, 4] * 6, dtype=torch.long)
1267
+ b2 = torch.tensor([5, 5] * 6, dtype=torch.long)
1268
+ ext = {}
1269
+ for k in range(8):
1270
+ ext[f"$a[{k}]"] = ((a2 >> k) & 1).float()
1271
+ ext[f"$b[{k}]"] = ((b2 >> k) & 1).float()
1272
+ ext[f"$pc[{k}]"] = ((pcs >> k) & 1).float()
1273
+ ext[f"$c[{k}]"] = ((cs >> k) & 1).float()
1274
+ out = ne.run(ext)
1275
+ npc = torch.zeros(12, dtype=torch.float64)
1276
+ for k in range(8):
1277
+ npc += out[f"subleq.pc_mux.bit{k}.or"][:, 0].double() * (1 << k)
1278
+ exp_npc = torch.where(torch.arange(12) % 2 == 0, cs.double(), ((pcs + 3) & 0xFF).double())
1279
+ ok_pc = int((npc == exp_npc).sum().item())
1280
+ print(f"datapath pc mux/inc : {ok_pc}/12")
1281
+ if ok_pc != 12:
1282
+ return 1
1283
+
1284
+ # 3. Program suite, threshold vs reference, full-state lockstep.
1285
+ cpu = SubleqThresholdCPU(T, reg)
1286
+ Z = 0x80 # scratch zero cell
1287
+
1288
+ progs = []
1289
+ # negate-copy: M[dst] = -M[src] (dst pre-zeroed): subleq src, dst; halt.
1290
+ m = sq((0x90, 0x91, 3), (Z, Z, HALT_PC))
1291
+ m[0x90], m[0x91] = 42, 0
1292
+ progs.append(("neg_copy", m, 30, lambda mm: mm[0x91] == ((-42) & 0xFF)))
1293
+ # add: M[d] += M[s] via double negation through Z.
1294
+ m = sq((0x90, Z, 3), (Z, 0x91, 6), (Z, Z, 9), (Z, Z, HALT_PC))
1295
+ m[0x90], m[0x91] = 17, 25
1296
+ progs.append(("add", m, 30, lambda mm: mm[0x91] == 42))
1297
+ # countdown loop: decrement M[n] by M[one] until <= 0, counting into 0x92
1298
+ # via add idiom is long; simplest loop: subleq one, n, done at 0.
1299
+ m = sq((0x93, 0x90, 6), (Z, Z, 0), (Z, Z, HALT_PC))
1300
+ m[0x90], m[0x93] = 5, 1
1301
+ progs.append(("countdown", m, 60, lambda mm: mm[0x90] == 0))
1302
+ # halt-immediately
1303
+ m = sq((Z, Z, HALT_PC))
1304
+ progs.append(("halt", m, 5, lambda mm: True))
1305
+
1306
+ all_ok = True
1307
+ for name, mem, budget, check in progs:
1308
+ ts = {"pc": 0, "mem": mem[:], "halted": False}
1309
+ rs = {"pc": 0, "mem": mem[:], "halted": False}
1310
+ cyc = 0
1311
+ ok = True
1312
+ while not ts["halted"] and cyc < budget:
1313
+ ts = cpu.step(ts)
1314
+ rs = ref_step(rs)
1315
+ cyc += 1
1316
+ if ts["pc"] != rs["pc"] or ts["mem"] != rs["mem"]:
1317
+ ok = False
1318
+ break
1319
+ ok = ok and ts["halted"] and rs["halted"] and check(ts["mem"])
1320
+ all_ok &= ok
1321
+ print(f" {name:12} {'PASS' if ok else 'FAIL'} ({cyc} cyc, lockstep)")
1322
+
1323
+ print("SUBLEQ:", "PASS" if all_ok else "FAIL")
1324
+ return 0 if all_ok else 1
1325
+
1326
+ # ============================================================================
1327
+ # neural_rv32 test suites
1328
+ # ============================================================================
1329
+ F1_0 = struct.unpack("<I", struct.pack("<f", 1.5))[0]
1330
+ F2_5 = struct.unpack("<I", struct.pack("<f", 2.5))[0]
1331
+ F10 = struct.unpack("<I", struct.pack("<f", 10.0))[0]
1332
+
1333
+
1334
+ def lockstep(cpu, mem, budget, name, check=None):
1335
+ ts = rv_state(mem)
1336
+ rs = rv_state(mem)
1337
+ import torch
1338
+ from eval_all import int_to_bits
1339
+ ts["_mem_bits"] = torch.tensor([int_to_bits(b, 8) for b in ts["mem"]],
1340
+ dtype=torch.float32)
1341
+ cyc = 0
1342
+ ok = True
1343
+ while not ts["halted"] and cyc < budget:
1344
+ ts = cpu.step(ts)
1345
+ # reference is strictly single-issue; advance it until PCs realign
1346
+ rs = RefRV32.step(rs)
1347
+ if rs["pc"] != ts["pc"] and not rs["halted"]:
1348
+ rs = RefRV32.step(rs)
1349
+ cyc += 1
1350
+ if (ts["pc"] != rs["pc"] or ts["regs"] != rs["regs"]
1351
+ or ts["fregs"] != rs["fregs"] or ts["halted"] != rs["halted"]):
1352
+ ok = False
1353
+ break
1354
+ if ok:
1355
+ ok = ts["halted"] and ts["mem"] == rs["mem"] and ts["console"] == rs["console"]
1356
+ if ok and check is not None:
1357
+ ok = check(ts)
1358
+ print(f" {name:22} {'PASS' if ok else 'FAIL'} ({cyc} cyc)")
1359
+ if not ok:
1360
+ print(f" pc t={ts['pc']:#x} r={rs['pc']:#x} halted t={ts['halted']} r={rs['halted']}")
1361
+ for i in range(32):
1362
+ if ts["regs"][i] != rs["regs"][i]:
1363
+ print(f" x{i}: t={ts['regs'][i]:#x} r={rs['regs'][i]:#x}")
1364
+ return ok
1365
+
1366
+
1367
+ def prog_int_alu():
1368
+ a = Asm()
1369
+ a.addi(1, 0, 100) # x1 = 100
1370
+ a.addi(2, 0, -7) # x2 = -7
1371
+ a.add(3, 1, 2) # 93
1372
+ a.sub(4, 1, 2) # 107
1373
+ a.xor(5, 1, 2)
1374
+ a.or_(6, 1, 2)
1375
+ a.and_(7, 1, 2)
1376
+ a.slt(8, 2, 1) # 1 (signed)
1377
+ a.sltu(9, 2, 1) # 0 (unsigned: -7 is huge)
1378
+ a.slli(10, 1, 3) # 800
1379
+ a.srli(11, 2, 4)
1380
+ a.srai(12, 2, 4) # -1
1381
+ a.slti(13, 2, 0) # 1
1382
+ a.sltiu(14, 1, 200) # 1
1383
+ a.xori(15, 1, 0xFF)
1384
+ a.ori(16, 1, 0x0F)
1385
+ a.andi(17, 1, 0x3C)
1386
+ a.lui(18, 0x12345)
1387
+ a.auipc(19, 1)
1388
+ a.ecall()
1389
+ return a.assemble(), 40, lambda s: s["regs"][3] == 93 and s["regs"][12] == u32(-1)
1390
+
1391
+
1392
+ def prog_branches():
1393
+ a = Asm()
1394
+ a.addi(1, 0, 5)
1395
+ a.addi(2, 0, -3)
1396
+ a.addi(10, 0, 0)
1397
+ a.blt(2, 1, "l1") # signed taken
1398
+ a.addi(10, 10, 100)
1399
+ a.label("l1")
1400
+ a.bltu(1, 2, "l2") # unsigned taken (5 < 0xFFFF..FD)
1401
+ a.addi(10, 10, 200)
1402
+ a.label("l2")
1403
+ a.beq(1, 1, "l3")
1404
+ a.addi(10, 10, 400)
1405
+ a.label("l3")
1406
+ a.bge(1, 2, "l4") # taken
1407
+ a.addi(10, 10, 800)
1408
+ a.label("l4")
1409
+ a.bne(1, 2, "done")
1410
+ a.addi(10, 10, 1600)
1411
+ a.label("done")
1412
+ a.ecall()
1413
+ return a.assemble(), 30, lambda s: s["regs"][10] == 0
1414
+
1415
+
1416
+ def prog_fib():
1417
+ a = Asm()
1418
+ a.addi(1, 0, 0) # a
1419
+ a.addi(2, 0, 1) # b
1420
+ a.addi(3, 0, 10) # n
1421
+ a.label("loop")
1422
+ a.beq(3, 0, "done")
1423
+ a.add(4, 1, 2)
1424
+ a.add(1, 0, 2) # a = b (independent pair candidates)
1425
+ a.add(2, 0, 4) # b = t
1426
+ a.addi(3, 3, -1)
1427
+ a.jal(0, "loop")
1428
+ a.label("done")
1429
+ a.ecall()
1430
+ return a.assemble(), 200, lambda s: s["regs"][1] == 55
1431
+
1432
+
1433
+ def prog_mem_call():
1434
+ a = Asm()
1435
+ a.addi(2, 0, 0x800) # sp
1436
+ a.addi(1, 0, 0x77)
1437
+ a.sb(1, 0, 0x400)
1438
+ a.lb(3, 0, 0x400)
1439
+ a.lbu(4, 0, 0x400)
1440
+ a.addi(1, 0, -2)
1441
+ a.sh(1, 0, 0x404)
1442
+ a.lh(5, 0, 0x404) # sign-extended
1443
+ a.lhu(6, 0, 0x404)
1444
+ a.jal(1, "fn") # call
1445
+ a.ecall()
1446
+ a.label("fn")
1447
+ a.addi(2, 2, -4)
1448
+ a.sw(1, 2, 0) # push ra
1449
+ a.lw(7, 2, 0)
1450
+ a.addi(2, 2, 4)
1451
+ a.jalr(0, 1, 0) # ret
1452
+ return a.assemble(), 40, lambda s: s["regs"][5] == u32(-2) and s["regs"][6] == 0xFFFE
1453
+
1454
+
1455
+ def prog_muldiv():
1456
+ a = Asm()
1457
+ a.addi(1, 0, -50)
1458
+ a.addi(2, 0, 7)
1459
+ a.mul(3, 1, 2) # -350
1460
+ a.mulh(4, 1, 2) # high of signed product = -1
1461
+ a.mulhu(5, 1, 2)
1462
+ a.div(6, 1, 2) # -7
1463
+ a.rem(7, 1, 2) # -1
1464
+ a.divu(8, 1, 2)
1465
+ a.remu(9, 1, 2)
1466
+ a.addi(10, 0, 0)
1467
+ a.div(11, 1, 10) # div by zero -> -1
1468
+ a.rem(12, 1, 10) # rem by zero -> dividend
1469
+ a.ecall()
1470
+ return a.assemble(), 30, lambda s: (s["regs"][3] == u32(-350)
1471
+ and s["regs"][6] == u32(-7)
1472
+ and s["regs"][7] == u32(-1)
1473
+ and s["regs"][11] == 0xFFFFFFFF)
1474
+
1475
+
1476
+ def prog_float():
1477
+ a = Asm()
1478
+ a.lui(1, 0x1) # x1 = 0x1000 data base
1479
+ a.flw(1, 1, 0) # f1 = 1.5
1480
+ a.flw(2, 1, 4) # f2 = 2.5
1481
+ a.fadd_s(3, 1, 2) # 4.0
1482
+ a.fmul_s(4, 1, 2) # 3.75
1483
+ a.fdiv_s(5, 2, 1) # 2.5 / 1.5
1484
+ a.fsub_s(6, 2, 1) # 1.0
1485
+ a.fsw(3, 1, 8)
1486
+ a.fle_s(10, 1, 2) # 1
1487
+ a.flt_s(11, 2, 1) # 0
1488
+ a.feq_s(12, 1, 1) # 1
1489
+ a.fmv_x_w(13, 3)
1490
+ a.ecall()
1491
+ mem = a.assemble()
1492
+ for i, wv in ((0, F1_0), (4, F2_5)):
1493
+ for b in range(4):
1494
+ mem[0x1000 + i + b] = (wv >> (8 * b)) & 0xFF
1495
+ f4 = struct.unpack("<I", struct.pack("<f", 4.0))[0]
1496
+ return mem, 30, lambda s: (s["regs"][10] == 1 and s["regs"][11] == 0
1497
+ and s["regs"][12] == 1 and s["regs"][13] == f4)
1498
+
1499
+
1500
+ def prog_neur():
1501
+ # XOR of two bits computed by a two-layer network of NEUR instructions:
1502
+ # h1 = OR(a,b) = H(a + b - 1); h2 = NAND(a,b) = H(-a - b + 1);
1503
+ # y = AND(h1,h2) = H(h1 + h2 - 2).
1504
+ def const(a, reg, v):
1505
+ a.lui(reg, v >> 12)
1506
+ a.ori(reg, reg, v & 0x7FF)
1507
+
1508
+ a = Asm()
1509
+ a.addi(1, 0, 0b10) # inputs: bit1=a=1, bit0=b=0
1510
+ const(a, 2, (0x1F << 16) | 0x03) # OR: pos=both, bias=-1 (0x1F sext5)
1511
+ a.neur(3, 1, 2)
1512
+ const(a, 4, (0x01 << 16) | (0x03 << 8)) # NAND: neg=both, bias=+1
1513
+ a.neur(5, 1, 4)
1514
+ a.slli(6, 3, 1)
1515
+ a.or_(6, 6, 5) # pack h1,h2 into bits 1,0
1516
+ const(a, 7, (0x1E << 16) | 0x03) # AND: pos=both, bias=-2 (0x1E sext5)
1517
+ a.neur(8, 6, 7) # y = XOR(a,b) = 1
1518
+ a.ecall()
1519
+ return a.assemble(), 20, lambda s: s["regs"][8] == 1 and s["regs"][3] == 1 and s["regs"][5] == 1
1520
+
1521
+
1522
+ def prog_mmio():
1523
+ a = Asm()
1524
+ a.lui(1, 0x10) # 0x10000 > 0xFFFF? keep in range: use addi chain
1525
+ a.addi(1, 0, 0x7F8)
1526
+ a.slli(1, 1, 5) # 0xFF00
1527
+ a.addi(2, 0, ord("H"))
1528
+ a.sb(2, 1, 0)
1529
+ a.addi(2, 0, ord("I"))
1530
+ a.sb(2, 1, 0)
1531
+ a.ecall()
1532
+ return a.assemble(), 20, lambda s: s["console"] == "HI"
1533
+
1534
+
1535
+ def prog_fcvt():
1536
+ # int -> float, add 1.0, float -> int: (7 + 1.0) truncated = 8; and a
1537
+ # round-trip of -5 through float and back = -5.
1538
+ import struct as _st
1539
+ a = Asm()
1540
+ a.addi(1, 0, 7)
1541
+ a.fcvt_s_w(0, 1) # f0 = 7.0
1542
+ a.lui(2, 0x1); a.flw(1, 2, 0) # f1 = 1.0 from data
1543
+ a.fadd_s(2, 0, 1) # f2 = 8.0
1544
+ a.fcvt_w_s(3, 2) # x3 = 8
1545
+ a.addi(4, 0, -5)
1546
+ a.fcvt_s_w(4, 4) # f4 = -5.0
1547
+ a.fcvt_w_s(5, 4) # x5 = -5
1548
+ a.ecall()
1549
+ mem = a.assemble()
1550
+ one = _st.unpack("<I", _st.pack("<f", 1.0))[0]
1551
+ for b in range(4):
1552
+ mem[0x1000 + b] = (one >> (8 * b)) & 0xFF
1553
+ return mem, 30, lambda s: s["regs"][3] == 8 and s["regs"][5] == (-5 & 0xFFFFFFFF)
1554
+
1555
+
1556
+ def rv32_netlist_check() -> bool:
1557
+ """Verify the rv32-specific sub-circuits (NEUR, hazard, MULHU accumulator)
1558
+ are recoverable and correct from the shipped .inputs metadata alone."""
1559
+ import random
1560
+ T = {}
1561
+ with safe_open(RV32_MODEL, framework="pt") as f:
1562
+ for name in f.keys():
1563
+ T[name] = f.get_tensor(name).float()
1564
+ reg = load_metadata(RV32_MODEL)["signal_registry"]
1565
+ rng = random.Random(0)
1566
+ ok = True
1567
+
1568
+ ne = NetlistEvaluator(T, reg, "rv32.neur")
1569
+ N = 512
1570
+ xs = [rng.getrandbits(8) for _ in range(N)]
1571
+ ps = [rng.getrandbits(8) for _ in range(N)]
1572
+ ns = [rng.getrandbits(8) for _ in range(N)]
1573
+ bs = [rng.randint(-16, 15) for _ in range(N)]
1574
+ ext = {}
1575
+ for i in range(8):
1576
+ ext[f"$rv32_neur_x[{i}]"] = torch.tensor([(x >> (7 - i)) & 1 for x in xs]).float()
1577
+ ext[f"$rv32_neur_pos[{i}]"] = torch.tensor([(p >> (7 - i)) & 1 for p in ps]).float()
1578
+ ext[f"$rv32_neur_neg[{i}]"] = torch.tensor([(n >> (7 - i)) & 1 for n in ns]).float()
1579
+ for k in range(6):
1580
+ ext[f"$rv32_neur_bias[{k}]"] = torch.tensor(
1581
+ [((b & 0x1F) >> k) & 1 if k < 5 else (b >> 4) & 1 for b in bs]).float()
1582
+ got = ne.run(ext)["rv32.neur.out"][:, 0]
1583
+ exp = torch.tensor([1.0 if (bin(x & p).count("1") - bin(x & n).count("1") + b) >= 0
1584
+ else 0.0 for x, p, n, b in zip(xs, ps, ns, bs)])
1585
+ ok &= bool((got == exp).all())
1586
+
1587
+ ne = NetlistEvaluator(T, reg, "rv32.hazard.rd1_rs1b")
1588
+ ii = [rng.getrandbits(5) for _ in range(N)]
1589
+ jj = [rng.getrandbits(5) for _ in range(N)]
1590
+ ext = {}
1591
+ for k in range(5):
1592
+ ext[f"$rv32_hz_rd1_rs1b_i[{k}]"] = torch.tensor([(i >> k) & 1 for i in ii]).float()
1593
+ ext[f"$rv32_hz_rd1_rs1b_j[{k}]"] = torch.tensor([(j >> k) & 1 for j in jj]).float()
1594
+ got = ne.run(ext)["rv32.hazard.rd1_rs1b.all"][:, 0]
1595
+ exp = torch.tensor([1.0 if i == j else 0.0 for i, j in zip(ii, jj)])
1596
+ ok &= bool((got == exp).all())
1597
+
1598
+ ne = NetlistEvaluator(T, reg, "rv32.mulacc")
1599
+ A = [rng.getrandbits(32) for _ in range(256)]
1600
+ B = [rng.getrandbits(32) for _ in range(256)]
1601
+ ext = {}
1602
+ for I in range(32):
1603
+ for J in range(32):
1604
+ ext[f"alu.alu32bit.mul.pp.a{I}b{J}"] = torch.tensor(
1605
+ [((a >> (31 - I)) & 1) & ((b >> (31 - J)) & 1) for a, b in zip(A, B)]).float()
1606
+ out = ne.run(ext)
1607
+ got = torch.zeros(len(A), dtype=torch.float64)
1608
+ for k in range(32):
1609
+ got += out[f"rv32.mulacc.s30.fa{32 + k}.ha2.sum.layer2"][:, 0].double() * (2.0 ** k)
1610
+ exp = torch.tensor([float(((a * b) >> 32) & 0xFFFFFFFF) for a, b in zip(A, B)],
1611
+ dtype=torch.float64)
1612
+ ok &= bool((got == exp).all())
1613
+
1614
+ print(f" rv32 netlist self-check (NEUR/hazard/MULHU from metadata) "
1615
+ f"{'PASS' if ok else 'FAIL'}")
1616
+ return ok
1617
+
1618
+
1619
+ def rv32_test() -> int:
1620
+ print("Loading neural_rv32...")
1621
+ cpu = Rv32ThresholdCPU()
1622
+ suite = [
1623
+ ("int_alu", prog_int_alu),
1624
+ ("branches", prog_branches),
1625
+ ("fib", prog_fib),
1626
+ ("mem_call", prog_mem_call),
1627
+ ("muldiv", prog_muldiv),
1628
+ ("float_subset", prog_float),
1629
+ ("fcvt_roundtrip", prog_fcvt),
1630
+ ("neur_xor_net", prog_neur),
1631
+ ("mmio_print", prog_mmio),
1632
+ ]
1633
+ all_ok = True
1634
+ for name, builder in suite:
1635
+ mem, budget, check = builder()
1636
+ all_ok &= lockstep(cpu, mem, budget, name, check)
1637
+ all_ok &= rv32_netlist_check()
1638
+ print(f"dual-issue pairs retired: {cpu.pairs_issued}")
1639
+ print("RV32:", "PASS" if all_ok else "FAIL")
1640
+ return 0 if all_ok else 1
1641
+
1642
+ PROG = r"""
1643
+ __attribute__((naked)) void _start(void){
1644
+ __asm__ volatile("li sp, 0xF000\n call main\n ecall\n");
1645
+ }
1646
+ static int gcd(int a,int b){ while(b){ int t=a%b; a=b; b=t;} return a; }
1647
+ int fib(int n){ int a=0,b=1; for(int i=0;i<n;i++){int t=a+b;a=b;b=t;} return a; }
1648
+ int main(void){
1649
+ int arr[8]={5,2,9,1,7,3,8,4};
1650
+ for(int i=0;i<8;i++) for(int j=i+1;j<8;j++)
1651
+ if(arr[j]<arr[i]){int t=arr[i];arr[i]=arr[j];arr[j]=t;}
1652
+ int s=0; for(int i=0;i<8;i++) s+=arr[i]*(i+1);
1653
+ return gcd(48,36) + fib(10) + s + (arr[0]==1?1000:0);
1654
+ }
1655
+ """
1656
+ EXPECTED = 1292 # gcd(48,36)=12, fib(10)=55, weighted sorted sum=225... verified natively
1657
+
1658
+
1659
+ def have_riscv_clang() -> bool:
1660
+ if not shutil.which("clang"):
1661
+ return False
1662
+ with tempfile.TemporaryDirectory() as d:
1663
+ c = os.path.join(d, "t.c")
1664
+ o = os.path.join(d, "t.o")
1665
+ open(c, "w").write("int f(void){return 1;}")
1666
+ r = subprocess.run(
1667
+ ["clang", "-target", "riscv32-unknown-elf", "-march=rv32im",
1668
+ "-mabi=ilp32", "-c", c, "-o", o],
1669
+ capture_output=True)
1670
+ return r.returncode == 0
1671
+
1672
+
1673
+ def rv32c_test() -> int:
1674
+ if not have_riscv_clang():
1675
+ print("SKIP: no riscv32 clang target available")
1676
+ return 0
1677
+
1678
+ with tempfile.TemporaryDirectory() as d:
1679
+ c = os.path.join(d, "prog.c")
1680
+ o = os.path.join(d, "prog.o")
1681
+ open(c, "w").write(PROG)
1682
+ subprocess.run(
1683
+ ["clang", "-target", "riscv32-unknown-elf", "-march=rv32im",
1684
+ "-mabi=ilp32", "-O1", "-nostdlib", "-ffreestanding",
1685
+ "-c", c, "-o", o], check=True, capture_output=True)
1686
+ obj = open(o, "rb").read()
1687
+
1688
+ mem, entry = load(obj, base=0, mem_size=65536)
1689
+ cpu = Rv32ThresholdCPU()
1690
+ st = rv_state(mem)
1691
+ st["pc"] = entry
1692
+ fin, cyc = cpu.run(st, max_cycles=5000)
1693
+ got = fin["regs"][10]
1694
+ ok = fin["halted"] and got == EXPECTED
1695
+ print(f"compiled C (gcd+fib+sort, rv32im): a0={got} "
1696
+ f"expected {EXPECTED} {cyc} instrs {'PASS' if ok else 'FAIL'}")
1697
+ return 0 if ok else 1
1698
+
1699
+ # ============================================================================
1700
+ # Entry point
1701
+ # ============================================================================
1702
+
1703
+ def main() -> int:
1704
+ ap = argparse.ArgumentParser(description="neural_subleq8 / neural_rv32 test suites")
1705
+ ap.add_argument("which", nargs="?", default="all",
1706
+ choices=["all", "subleq", "rv32", "rv32-c"])
1707
+ args = ap.parse_args()
1708
+ rc = 0
1709
+ if args.which in ("all", "subleq"):
1710
+ rc |= subleq_test()
1711
+ if args.which in ("all", "rv32"):
1712
+ rc |= rv32_test()
1713
+ if args.which in ("all", "rv32-c"):
1714
+ rc |= rv32c_test()
1715
+ return rc
1716
+
1717
+
1718
+ if __name__ == "__main__":
1719
+ sys.exit(main())
neural_computer.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fb4725ab213771fa0f20c42d13a43b913cb33f1762c6a8c4b4670856f8d3592c
3
+ size 32717681
play.py ADDED
@@ -0,0 +1,262 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Hands-on playground for the 8bit-threshold-computer.
3
+
4
+ Loads a safetensors model, reads its manifest, and exercises threshold
5
+ circuits at every level: raw Boolean gates, 8-bit ALU arithmetic and
6
+ comparators, multi-layer modular arithmetic, and a manifest-sized CPU
7
+ runtime running a small assembled program end-to-end through the
8
+ threshold weights.
9
+
10
+ The CPU demo defaults to the small (1 KB) profile so the run finishes in
11
+ a fraction of a second. Larger profiles (4 KB, 64 KB) take proportionally
12
+ longer because every memory access decodes against every address line.
13
+
14
+ Usage:
15
+ python play.py # fast 1KB demo
16
+ python play.py --model neural_computer.safetensors # full 64KB
17
+ python play.py --model variants/neural_alu8.safetensors --skip-cpu # ALU only
18
+ """
19
+
20
+ from __future__ import annotations
21
+ import argparse
22
+ import os
23
+ import sys
24
+
25
+ import torch
26
+ from safetensors import safe_open
27
+
28
+ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
29
+
30
+ # Reuse the variant-aware CPU runtime from eval_all.py
31
+ from eval_all import GenericThresholdCPU, builtin_program
32
+
33
+
34
+ def heaviside(x):
35
+ return (x >= 0).float()
36
+
37
+
38
+ def load_tensors(path):
39
+ out = {}
40
+ with safe_open(path, framework="pt") as f:
41
+ for name in f.keys():
42
+ out[name] = f.get_tensor(name).float()
43
+ return out
44
+
45
+
46
+ def main() -> int:
47
+ parser = argparse.ArgumentParser(description="Threshold computer playground")
48
+ parser.add_argument(
49
+ "--model", type=str,
50
+ default=os.path.join(os.path.dirname(__file__),
51
+ "variants", "neural_computer8_small.safetensors"),
52
+ help="Path to a .safetensors variant"
53
+ )
54
+ parser.add_argument("--skip-cpu", action="store_true",
55
+ help="Skip the CPU program demo (useful for pure-ALU files)")
56
+ args = parser.parse_args()
57
+
58
+ print("Loading", args.model)
59
+ T = load_tensors(args.model)
60
+
61
+ DATA_BITS = int(T["manifest.data_bits"].item())
62
+ ADDR_BITS = int(T["manifest.addr_bits"].item())
63
+ MEM_BYTES = int(T["manifest.memory_bytes"].item())
64
+ REGISTERS = int(T["manifest.registers"].item())
65
+ print(f"Manifest: data={DATA_BITS}-bit, addr={ADDR_BITS}-bit, mem={MEM_BYTES}B, regs={REGISTERS}")
66
+ print(f"Tensors: {len(T):,}")
67
+ print(f"Total params: {sum(t.numel() for t in T.values()):,}")
68
+ print()
69
+
70
+ def gate(name, inputs):
71
+ w = T[name + ".weight"].view(-1)
72
+ b = T[name + ".bias"].view(-1)
73
+ return int(heaviside((torch.tensor(inputs, dtype=torch.float32) * w).sum() + b).item())
74
+
75
+ def xor(prefix, inputs):
76
+ a, b_ = inputs
77
+ h_or = gate(f"{prefix}.layer1.or", [a, b_])
78
+ h_nand = gate(f"{prefix}.layer1.nand", [a, b_])
79
+ return gate(f"{prefix}.layer2", [h_or, h_nand])
80
+
81
+ def xor_neuron(prefix, inputs):
82
+ a, b_ = inputs
83
+ h1 = gate(f"{prefix}.layer1.neuron1", [a, b_])
84
+ h2 = gate(f"{prefix}.layer1.neuron2", [a, b_])
85
+ return gate(f"{prefix}.layer2", [h1, h2])
86
+
87
+ def int_to_bits_msb(v, n):
88
+ return [(v >> (n - 1 - i)) & 1 for i in range(n)]
89
+
90
+ def bits_msb_to_int(bits):
91
+ out = 0
92
+ for b in bits:
93
+ out = (out << 1) | int(b)
94
+ return out
95
+
96
+ # ---------- Demo 1: Boolean gates ----------
97
+ print("=" * 64)
98
+ print(" Demo 1: Boolean threshold gates")
99
+ print("=" * 64)
100
+ truth_2 = [(0, 0), (0, 1), (1, 0), (1, 1)]
101
+ for gname in ["and", "or", "nand", "nor", "implies"]:
102
+ row = " ".join(f"{a}{b}->{gate(f'boolean.{gname}', [a, b])}" for a, b in truth_2)
103
+ print(f" {gname:8} {row}")
104
+ for gname in ["xor", "xnor", "biimplies"]:
105
+ row = " ".join(f"{a}{b}->{xor_neuron(f'boolean.{gname}', [a, b])}" for a, b in truth_2)
106
+ print(f" {gname:8} {row}")
107
+ print(f" not 0->{gate('boolean.not', [0])} 1->{gate('boolean.not', [1])}")
108
+ print()
109
+
110
+ # ---------- Demo 2: 8-bit ALU arithmetic ----------
111
+ print("=" * 64)
112
+ print(" Demo 2: 8-bit ALU arithmetic (every gate is threshold logic)")
113
+ print("=" * 64)
114
+
115
+ def fa(prefix, a, b, cin):
116
+ s1 = xor(f"{prefix}.ha1.sum", [a, b])
117
+ c1 = gate(f"{prefix}.ha1.carry", [a, b])
118
+ s2 = xor(f"{prefix}.ha2.sum", [s1, cin])
119
+ c2 = gate(f"{prefix}.ha2.carry", [s1, cin])
120
+ return s2, gate(f"{prefix}.carry_or", [c1, c2])
121
+
122
+ def alu_add(a, b):
123
+ a_lsb = list(reversed(int_to_bits_msb(a, 8)))
124
+ b_lsb = list(reversed(int_to_bits_msb(b, 8)))
125
+ carry = 0
126
+ sum_lsb = []
127
+ for i in range(8):
128
+ s, carry = fa(f"arithmetic.ripplecarry8bit.fa{i}", a_lsb[i], b_lsb[i], carry)
129
+ sum_lsb.append(s)
130
+ return bits_msb_to_int(list(reversed(sum_lsb))), carry
131
+
132
+ def alu_sub(a, b):
133
+ a_lsb = list(reversed(int_to_bits_msb(a, 8)))
134
+ b_lsb = list(reversed(int_to_bits_msb(b, 8)))
135
+ carry = 1
136
+ diff_lsb = []
137
+ for i in range(8):
138
+ notb = gate(f"arithmetic.sub8bit.notb{i}", [b_lsb[i]])
139
+ x1 = xor(f"arithmetic.sub8bit.fa{i}.xor1", [a_lsb[i], notb])
140
+ x2 = xor(f"arithmetic.sub8bit.fa{i}.xor2", [x1, carry])
141
+ and1 = gate(f"arithmetic.sub8bit.fa{i}.and1", [a_lsb[i], notb])
142
+ and2 = gate(f"arithmetic.sub8bit.fa{i}.and2", [x1, carry])
143
+ carry = gate(f"arithmetic.sub8bit.fa{i}.or_carry", [and1, and2])
144
+ diff_lsb.append(x2)
145
+ return bits_msb_to_int(list(reversed(diff_lsb))), carry
146
+
147
+ def alu_compare(a, b, kind):
148
+ # Walks the bit-cascade comparator family: per-bit gt/lt/eq, cascaded
149
+ # eq_prefix, cascade.gt/lt, and the final OR/AND gates. Bit 0 is MSB.
150
+ a_msb = int_to_bits_msb(a, 8)
151
+ b_msb = int_to_bits_msb(b, 8)
152
+ bit_gt = [gate(f"arithmetic.cmp8bit.bit{i}.gt", [a_msb[i], b_msb[i]]) for i in range(8)]
153
+ bit_lt = [gate(f"arithmetic.cmp8bit.bit{i}.lt", [a_msb[i], b_msb[i]]) for i in range(8)]
154
+ bit_eq = []
155
+ for i in range(8):
156
+ eq_and = gate(f"arithmetic.cmp8bit.bit{i}.eq.layer1.and", [a_msb[i], b_msb[i]])
157
+ eq_nor = gate(f"arithmetic.cmp8bit.bit{i}.eq.layer1.nor", [a_msb[i], b_msb[i]])
158
+ bit_eq.append(gate(f"arithmetic.cmp8bit.bit{i}.eq", [eq_and, eq_nor]))
159
+ cas_gt = [bit_gt[0]]
160
+ cas_lt = [bit_lt[0]]
161
+ for i in range(1, 8):
162
+ eq_pref = gate(f"arithmetic.cmp8bit.cascade.eq_prefix.bit{i}", bit_eq[:i])
163
+ cas_gt.append(gate(f"arithmetic.cmp8bit.cascade.gt.bit{i}", [eq_pref, bit_gt[i]]))
164
+ cas_lt.append(gate(f"arithmetic.cmp8bit.cascade.lt.bit{i}", [eq_pref, bit_lt[i]]))
165
+ if kind == "greaterthan":
166
+ return gate("arithmetic.greaterthan8bit", cas_gt)
167
+ if kind == "lessthan":
168
+ return gate("arithmetic.lessthan8bit", cas_lt)
169
+ if kind == "eq":
170
+ return gate("arithmetic.equality8bit", bit_eq)
171
+ raise ValueError(kind)
172
+
173
+ def alu_mul(a, b):
174
+ a_bits = int_to_bits_msb(a, 8)
175
+ b_bits = int_to_bits_msb(b, 8)
176
+ result = 0
177
+ for j in range(8):
178
+ if b_bits[j] == 0:
179
+ continue
180
+ row = 0
181
+ for i in range(8):
182
+ pp = gate(f"alu.alu8bit.mul.pp.a{i}b{j}", [a_bits[i], b_bits[j]])
183
+ row |= (pp << (7 - i))
184
+ shift = 7 - j
185
+ result, _ = alu_add(result & 0xFF, (row << shift) & 0xFF)
186
+ return result & 0xFF
187
+
188
+ cases_arith = [(5, 3), (37, 100), (200, 99), (255, 1), (127, 128), (15, 17)]
189
+ print("ADD:")
190
+ for a, b in cases_arith:
191
+ r, c = alu_add(a, b)
192
+ e = (a + b) & 0xFF
193
+ print(f" {a:3} + {b:3} = {r:3} (carry={c}) expected {e:3} [{'OK' if r == e else 'FAIL'}]")
194
+ print("SUB:")
195
+ for a, b in cases_arith:
196
+ r, c = alu_sub(a, b)
197
+ e = (a - b) & 0xFF
198
+ print(f" {a:3} - {b:3} = {r:3} (no_borrow={c}) expected {e:3} [{'OK' if r == e else 'FAIL'}]")
199
+ print("CMP:")
200
+ for a, b in [(50, 30), (30, 50), (77, 77), (255, 0), (0, 255), (128, 127)]:
201
+ gt = alu_compare(a, b, "greaterthan")
202
+ lt = alu_compare(a, b, "lessthan")
203
+ eq = alu_compare(a, b, "eq")
204
+ print(f" {a:3} vs {b:3} -> GT={gt} LT={lt} EQ={eq}")
205
+ print("MUL (low 8 bits):")
206
+ for a, b in [(12, 11), (15, 17), (8, 32), (200, 3), (0, 99), (1, 255)]:
207
+ r = alu_mul(a, b)
208
+ e = (a * b) & 0xFF
209
+ print(f" {a:3} * {b:3} = {r:3} expected {e:3} [{'OK' if r == e else 'FAIL'}]")
210
+ print()
211
+
212
+ # ---------- Demo 3: mod-5 divisibility ----------
213
+ print("=" * 64)
214
+ print(" Demo 3: mod-5 divisibility (multi-layer, hand-constructed)")
215
+ print("=" * 64)
216
+
217
+ def mod5(v):
218
+ # Per-multiple-of-5 match (k0, k5, ..., k255): each k has 8 single-input
219
+ # "bit{i}.match" gates that fire when bit i of v matches bit i of k,
220
+ # ANDed by ".all". Final ".weight" ORs all 52 "all" outputs.
221
+ bits = int_to_bits_msb(v, 8)
222
+ ks = [k for k in range(256) if k % 5 == 0]
223
+ alls = []
224
+ for k in ks:
225
+ matches = [gate(f"modular.mod5.eq.k{k}.bit{i}.match", [bits[i]]) for i in range(8)]
226
+ alls.append(gate(f"modular.mod5.eq.k{k}.all", matches))
227
+ return gate("modular.mod5", alls)
228
+
229
+ hits = [v for v in range(256) if mod5(v)]
230
+ print(f" v in [0,255] with mod5(v)==1: {len(hits)} hits, first 12: {hits[:12]}")
231
+ print(f" Sanity (each %5): {[h % 5 for h in hits[:12]]}")
232
+ print()
233
+
234
+ # ---------- Demo 4: CPU running an assembled program ----------
235
+ if args.skip_cpu or MEM_BYTES < 0x84:
236
+ if args.skip_cpu:
237
+ print("Demo 4 skipped (--skip-cpu).")
238
+ else:
239
+ print(f"Demo 4 skipped (memory={MEM_BYTES}B too small for the demo program).")
240
+ return 0
241
+
242
+ print("=" * 64)
243
+ print(f" Demo 4: Threshold CPU running an assembled program ({MEM_BYTES} B memory)")
244
+ print("=" * 64)
245
+ print(" Program: sum 5+4+3+2+1 via loop")
246
+ print(" uses LOAD/STORE/ADD/SUB/CMP/JNZ/HALT, all threshold-gated")
247
+ print(" Running ... (larger memories take longer because every memory access")
248
+ print(" decodes against every address line)")
249
+ cpu = GenericThresholdCPU({k: v for k, v in T.items()})
250
+ mem, expected = builtin_program(ADDR_BITS)
251
+ state = {"pc": 0, "regs": [0] * 4, "flags": [0] * 4, "mem": mem, "halted": False}
252
+ final, cycles = cpu.run(state, max_cycles=200)
253
+ got = final["mem"][0x83]
254
+ print(f" Halted after {cycles} cycles")
255
+ print(f" R0={final['regs'][0]} R1={final['regs'][1]} "
256
+ f"R2={final['regs'][2]} R3={final['regs'][3]}")
257
+ print(f" M[0x0083] = {got} (expected {expected}) [{'OK' if got == expected else 'FAIL'}]")
258
+ return 0 if got == expected else 1
259
+
260
+
261
+ if __name__ == "__main__":
262
+ sys.exit(main())
prune_weights.py ADDED
@@ -0,0 +1,484 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ BATCHED WEIGHT PRUNING (GPU-optimized)
3
+ ======================================
4
+ Phase 1: Batch eval all candidates in parallel
5
+ Phase 2: Apply all successes at once, binary search if conflicts
6
+ """
7
+
8
+ import argparse
9
+ import os
10
+ import time
11
+
12
+ import torch
13
+ from safetensors.torch import save_file
14
+ from eval import BatchedFitnessEvaluator, create_population, load_model
15
+
16
+ torch.manual_seed(0)
17
+
18
+ DEFAULT_OUTPUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "pruned.safetensors")
19
+
20
+
21
+ def format_time(seconds):
22
+ if seconds < 60:
23
+ return f"{seconds:.1f}s"
24
+ elif seconds < 3600:
25
+ return f"{seconds/60:.1f}m"
26
+ else:
27
+ return f"{seconds/3600:.1f}h"
28
+
29
+
30
+ def format_eta(elapsed, done, total):
31
+ if done == 0:
32
+ return "calculating..."
33
+ rate = done / elapsed
34
+ remaining = (total - done) / rate
35
+ return format_time(remaining)
36
+
37
+
38
+ def apply_reductions(model, reductions):
39
+ """Apply a list of (name, flat_idx, shape, old_val) reductions."""
40
+ for name, flat_idx, shape, old_val in reductions:
41
+ new_val = old_val - 1 if old_val > 0 else old_val + 1
42
+ flat = model[name].flatten()
43
+ if flat[flat_idx].item() == old_val:
44
+ flat[flat_idx] = new_val
45
+ model[name] = flat.view(shape)
46
+
47
+
48
+ def revert_reductions(model, reductions):
49
+ """Revert a list of reductions."""
50
+ for name, flat_idx, shape, old_val in reductions:
51
+ flat = model[name].flatten()
52
+ new_val = old_val - 1 if old_val > 0 else old_val + 1
53
+ if flat[flat_idx].item() == new_val:
54
+ flat[flat_idx] = old_val
55
+ model[name] = flat.view(shape)
56
+
57
+
58
+ def check_fitness(model, evaluator, device):
59
+ """Check model fitness."""
60
+ torch.manual_seed(0)
61
+ pop = create_population(model, 1, device)
62
+ return evaluator.evaluate(pop, debug=False)[0].item()
63
+
64
+
65
+ def sequential_conflict_resolution(model, evaluator, device, candidates, base_magnitude):
66
+ """
67
+ Sequential fallback - tests and applies reductions one at a time.
68
+ Slower but guarantees no interaction bugs.
69
+ """
70
+ accepted = []
71
+ for i, (name, flat_idx, shape, old_val) in enumerate(candidates):
72
+ apply_reductions(model, [(name, flat_idx, shape, old_val)])
73
+ fitness = check_fitness(model, evaluator, device)
74
+ if fitness >= 0.9999:
75
+ accepted.append((name, flat_idx, shape, old_val))
76
+ if (i + 1) % 50 == 0:
77
+ current_mag = sum(t.abs().sum().item() for t in model.values())
78
+ reduction_pct = 100 * (1 - current_mag / base_magnitude)
79
+ print(f" Sequential: {len(accepted)}/{i+1} accepted | mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
80
+ else:
81
+ revert_reductions(model, [(name, flat_idx, shape, old_val)])
82
+ return accepted
83
+
84
+
85
+ def batched_conflict_resolution(model, evaluator, device, candidates, base_magnitude):
86
+ """
87
+ Batched binary search - evaluates multiple branches in parallel.
88
+ Uses BFS instead of DFS to maximize batching opportunities.
89
+ Verifies cumulative effect after each batch to prevent interaction bugs.
90
+ """
91
+ if len(candidates) == 0:
92
+ return []
93
+
94
+ # First try all at once
95
+ print(f" Trying {len(candidates)} reductions at once...")
96
+ apply_reductions(model, candidates)
97
+ fitness = check_fitness(model, evaluator, device)
98
+
99
+ if fitness >= 0.9999:
100
+ current_mag = sum(t.abs().sum().item() for t in model.values())
101
+ reduction_pct = 100 * (1 - current_mag / base_magnitude)
102
+ print(f" ALL {len(candidates)} OK | fitness={fitness:.6f} | "
103
+ f"mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
104
+ return candidates
105
+
106
+ # Conflict - revert and use batched BFS
107
+ revert_reductions(model, candidates)
108
+ print(f" CONFLICT (fitness={fitness:.6f}), starting batched resolution...")
109
+
110
+ accepted = []
111
+ # Queue of (candidate_list, depth) to process
112
+ pending = [(candidates, 0)]
113
+
114
+ while pending:
115
+ # Collect all pending groups for batch evaluation
116
+ to_eval = []
117
+ for group, depth in pending:
118
+ if len(group) == 0:
119
+ continue
120
+ elif len(group) == 1:
121
+ to_eval.append((group, depth, 'single'))
122
+ else:
123
+ to_eval.append((group, depth, 'group'))
124
+
125
+ pending = []
126
+
127
+ if not to_eval:
128
+ break
129
+
130
+ # Build batch: create model variants for each group
131
+ batch_size = len(to_eval)
132
+ print(f" Batch evaluating {batch_size} groups...")
133
+
134
+ # Create population for batch eval
135
+ pop = {}
136
+ for name, tensor in model.items():
137
+ pop[name] = tensor.unsqueeze(0).expand(batch_size, *tensor.shape).clone().to(device)
138
+
139
+ # Apply each group's reductions to its population slot
140
+ for idx, (group, depth, gtype) in enumerate(to_eval):
141
+ for name, flat_idx, shape, old_val in group:
142
+ new_val = old_val - 1 if old_val > 0 else old_val + 1
143
+ flat_view = pop[name][idx].flatten()
144
+ # Check if not already modified in base model
145
+ base_val = model[name].flatten()[flat_idx].item()
146
+ if base_val == old_val:
147
+ flat_view[flat_idx] = new_val
148
+
149
+ # Batch evaluate
150
+ torch.manual_seed(0)
151
+ fitnesses = evaluator.evaluate(pop, debug=False)
152
+
153
+ # Process results - collect accepted groups first, then verify
154
+ batch_accepted = []
155
+ ok_count = 0
156
+ conflict_count = 0
157
+ fail_count = 0
158
+
159
+ for idx, (group, depth, gtype) in enumerate(to_eval):
160
+ fit = fitnesses[idx].item()
161
+ indent = " " + " " * depth
162
+
163
+ if fit >= 0.9999:
164
+ batch_accepted.append((group, depth, indent))
165
+ ok_count += len(group)
166
+ else:
167
+ if len(group) == 1:
168
+ name, flat_idx, shape, old_val = group[0]
169
+ print(f"{indent}[1/1] FAIL {name}[{flat_idx}] | fitness={fit:.6f}")
170
+ fail_count += 1
171
+ else:
172
+ mid = len(group) // 2
173
+ left = group[:mid]
174
+ right = group[mid:]
175
+ print(f"{indent}CONFLICT ({len(group)}) fitness={fit:.6f} -> split {len(left)}+{len(right)}")
176
+ pending.append((left, depth + 1))
177
+ pending.append((right, depth + 1))
178
+ conflict_count += 1
179
+
180
+ # Apply all batch-accepted reductions
181
+ all_batch_reductions = []
182
+ for group, depth, indent in batch_accepted:
183
+ apply_reductions(model, group)
184
+ all_batch_reductions.extend(group)
185
+
186
+ # Verify cumulative effect
187
+ if all_batch_reductions:
188
+ verify_fitness = check_fitness(model, evaluator, device)
189
+ if verify_fitness >= 0.9999:
190
+ # All good - commit these reductions
191
+ for group, depth, indent in batch_accepted:
192
+ current_mag = sum(t.abs().sum().item() for t in model.values())
193
+ reduction_pct = 100 * (1 - current_mag / base_magnitude)
194
+ if len(group) == 1:
195
+ name, flat_idx, shape, old_val = group[0]
196
+ print(f"{indent}[1/1] OK {name}[{flat_idx}] | mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
197
+ else:
198
+ print(f"{indent}ALL {len(group)} OK | mag={current_mag:.0f} (-{reduction_pct:.2f}%)")
199
+ accepted.extend(all_batch_reductions)
200
+ print(f" Batch result: {ok_count} accepted, {conflict_count} split, {fail_count} failed")
201
+ else:
202
+ # Interaction bug detected - revert and use sequential fallback
203
+ print(f" INTERACTION BUG detected (batch fitness={verify_fitness:.6f})")
204
+ print(f" Reverting {len(all_batch_reductions)} reductions, falling back to sequential...")
205
+ revert_reductions(model, all_batch_reductions)
206
+
207
+ # Process each group sequentially
208
+ seq_accepted = sequential_conflict_resolution(
209
+ model, evaluator, device, all_batch_reductions, base_magnitude
210
+ )
211
+ accepted.extend(seq_accepted)
212
+ print(f" Sequential fallback: {len(seq_accepted)}/{len(all_batch_reductions)} accepted")
213
+ else:
214
+ print(f" Batch result: {ok_count} accepted, {conflict_count} split, {fail_count} failed")
215
+
216
+ return accepted
217
+
218
+
219
+ def prune_weights(
220
+ passes: int = 10,
221
+ batch_size: int = 5000,
222
+ device: str = 'cuda',
223
+ checkpoint_path: str = DEFAULT_OUTPUT
224
+ ):
225
+ print("=" * 80)
226
+ print(" BATCHED WEIGHT PRUNING (GPU-optimized)")
227
+ print("=" * 80)
228
+ print(f" Device: {device}")
229
+ print(f" Batch size: {batch_size}")
230
+ print(f" Max passes: {passes}")
231
+ print("=" * 80)
232
+
233
+ # Load model
234
+ print("\n[1/4] LOADING MODEL...")
235
+ load_start = time.perf_counter()
236
+ model = load_model()
237
+ load_time = time.perf_counter() - load_start
238
+
239
+ n_params = sum(t.numel() for t in model.values())
240
+ n_tensors = len(model)
241
+ base_magnitude = sum(t.abs().sum().item() for t in model.values())
242
+ base_max = max(t.abs().max().item() for t in model.values())
243
+ nonzero_params = sum((t != 0).sum().item() for t in model.values())
244
+
245
+ print(f" Loaded in {load_time:.2f}s")
246
+ print(f" Tensors: {n_tensors}")
247
+ print(f" Parameters: {n_params}")
248
+ print(f" Non-zero parameters: {nonzero_params}")
249
+ print(f" Total magnitude: {base_magnitude:.0f}")
250
+ print(f" Max weight: {base_max:.0f}")
251
+
252
+ # Initialize evaluator
253
+ print("\n[2/4] INITIALIZING EVALUATOR...")
254
+ eval_start = time.perf_counter()
255
+ evaluator = BatchedFitnessEvaluator(device=device)
256
+ eval_time = time.perf_counter() - eval_start
257
+ print(f" Initialized in {eval_time:.2f}s")
258
+
259
+ # Verify initial fitness
260
+ print("\n[3/4] VERIFYING BASE MODEL...")
261
+ initial_fitness = check_fitness(model, evaluator, device)
262
+ print(f" Fitness: {initial_fitness:.6f}")
263
+
264
+ if initial_fitness < 0.9999:
265
+ print(f" ERROR: Base model fitness {initial_fitness:.6f} < 0.9999")
266
+ return None
267
+
268
+ print(f" STATUS: PASS")
269
+
270
+ # Build parameter list
271
+ print("\n[4/4] BUILDING PARAMETER INDEX...")
272
+ param_list = []
273
+ for name, tensor in model.items():
274
+ flat = tensor.flatten()
275
+ for i in range(len(flat)):
276
+ param_list.append((name, i, tensor.shape))
277
+ print(f" Indexed {len(param_list)} parameters")
278
+
279
+ # Main pruning loop
280
+ print("\n" + "=" * 80)
281
+ print(" PRUNING STARTED")
282
+ print("=" * 80)
283
+
284
+ total_reductions = 0
285
+ pruning_start = time.perf_counter()
286
+
287
+ for pass_num in range(passes):
288
+ torch.manual_seed(0)
289
+ pass_start = time.perf_counter()
290
+
291
+ print(f"\n{'='*80}")
292
+ print(f" PASS {pass_num + 1}/{passes}")
293
+ print(f"{'='*80}")
294
+
295
+ # Count candidates
296
+ candidates = []
297
+ for name, idx, shape in param_list:
298
+ flat = model[name].flatten()
299
+ val = flat[idx].item()
300
+ if val != 0:
301
+ candidates.append((name, idx, shape, val))
302
+
303
+ n_candidates = len(candidates)
304
+ print(f"\n Candidates: {n_candidates} non-zero weights")
305
+
306
+ if n_candidates == 0:
307
+ print(f" No candidates remaining. Stopping.")
308
+ break
309
+
310
+ # Phase 1: Batch evaluation
311
+ print(f"\n PHASE 1: Batch evaluation (testing each reduction independently)")
312
+ print(f" " + "-" * 60)
313
+ phase1_start = time.perf_counter()
314
+ successful_candidates = []
315
+ n_batches = (n_candidates + batch_size - 1) // batch_size
316
+
317
+ for batch_idx, batch_start_idx in enumerate(range(0, n_candidates, batch_size)):
318
+ batch = candidates[batch_start_idx:batch_start_idx + batch_size]
319
+ batch_len = len(batch)
320
+ batch_start_time = time.perf_counter()
321
+
322
+ # Build population
323
+ pop = {}
324
+ for name, tensor in model.items():
325
+ pop[name] = tensor.unsqueeze(0).expand(batch_len, *tensor.shape).clone().to(device)
326
+
327
+ # Apply reductions
328
+ for pop_idx, (name, flat_idx, shape, old_val) in enumerate(batch):
329
+ new_val = old_val - 1 if old_val > 0 else old_val + 1
330
+ flat_view = pop[name][pop_idx].flatten()
331
+ flat_view[flat_idx] = new_val
332
+
333
+ # Evaluate
334
+ torch.manual_seed(0)
335
+ if device == 'cuda':
336
+ torch.cuda.synchronize()
337
+ fitness = evaluator.evaluate(pop, debug=False)
338
+ if device == 'cuda':
339
+ torch.cuda.synchronize()
340
+
341
+ # Collect successes
342
+ batch_successes = 0
343
+ for pop_idx, (name, flat_idx, shape, old_val) in enumerate(batch):
344
+ if fitness[pop_idx].item() >= 0.9999:
345
+ successful_candidates.append((name, flat_idx, shape, old_val))
346
+ batch_successes += 1
347
+
348
+ batch_time = time.perf_counter() - batch_start_time
349
+ elapsed = time.perf_counter() - phase1_start
350
+ done = batch_start_idx + batch_len
351
+ eta = format_eta(elapsed, done, n_candidates)
352
+ throughput = batch_len / batch_time
353
+
354
+ print(f" Batch {batch_idx + 1}/{n_batches}: "
355
+ f"{batch_successes}/{batch_len} passed ({100*batch_successes/batch_len:.1f}%) | "
356
+ f"Total OK: {len(successful_candidates)} | "
357
+ f"Progress: {done}/{n_candidates} ({100*done/n_candidates:.1f}%) | "
358
+ f"Speed: {throughput:.0f}/s | "
359
+ f"ETA: {eta}")
360
+
361
+ phase1_time = time.perf_counter() - phase1_start
362
+ print(f"\n Phase 1 complete: {len(successful_candidates)}/{n_candidates} candidates "
363
+ f"({100*len(successful_candidates)/n_candidates:.1f}%) in {format_time(phase1_time)}")
364
+
365
+ # Phase 2: Apply with conflict resolution
366
+ if len(successful_candidates) == 0:
367
+ print(f"\n No reductions possible. Stopping.")
368
+ break
369
+
370
+ print(f"\n PHASE 2: Apply reductions with conflict resolution")
371
+ print(f" " + "-" * 60)
372
+ phase2_start = time.perf_counter()
373
+
374
+ accepted = batched_conflict_resolution(model, evaluator, device, successful_candidates, base_magnitude)
375
+ pass_reductions = len(accepted)
376
+
377
+ phase2_time = time.perf_counter() - phase2_start
378
+ print(f"\n Phase 2 complete: {pass_reductions} reductions applied in {format_time(phase2_time)}")
379
+
380
+ # Pass summary
381
+ total_reductions += pass_reductions
382
+ current_magnitude = sum(t.abs().sum().item() for t in model.values())
383
+ current_nonzero = sum((t != 0).sum().item() for t in model.values())
384
+ pass_time = time.perf_counter() - pass_start
385
+ reduction_pct = 100 * (1 - current_magnitude / base_magnitude)
386
+
387
+ print(f"\n PASS {pass_num + 1} SUMMARY:")
388
+ print(f" Reductions this pass: {pass_reductions}")
389
+ print(f" Total reductions: {total_reductions}")
390
+ print(f" Current magnitude: {current_magnitude:.0f} (-{reduction_pct:.2f}%)")
391
+ print(f" Current non-zero: {current_nonzero}")
392
+ print(f" Pass time: {format_time(pass_time)}")
393
+
394
+ # Verify after pass
395
+ print(f"\n Verifying model integrity...")
396
+ fitness = check_fitness(model, evaluator, device)
397
+ print(f" Fitness: {fitness:.6f} {'PASS' if fitness >= 0.9999 else 'FAIL'}")
398
+
399
+ # Save checkpoint after each pass
400
+ checkpoint_name = checkpoint_path.replace('.safetensors', f'_pass{pass_num + 1}.safetensors')
401
+ print(f"\n Saving checkpoint: {checkpoint_name}")
402
+ save_file(model, checkpoint_name)
403
+ print(f" Saved. Magnitude: {current_magnitude:.0f} (-{reduction_pct:.2f}%)")
404
+
405
+ # Also save as "latest" for easy access
406
+ latest_path = checkpoint_path.replace('.safetensors', '_latest.safetensors')
407
+ save_file(model, latest_path)
408
+ print(f" Also saved as: {latest_path}")
409
+
410
+ if pass_reductions == 0:
411
+ print(f"\n No reductions achieved. Stopping early.")
412
+ break
413
+
414
+ # Final summary
415
+ pruning_time = time.perf_counter() - pruning_start
416
+ final_magnitude = sum(t.abs().sum().item() for t in model.values())
417
+ final_max = max(t.abs().max().item() for t in model.values())
418
+ final_nonzero = sum((t != 0).sum().item() for t in model.values())
419
+ reduction_pct = 100 * (1 - final_magnitude / base_magnitude)
420
+
421
+ print("\n" + "=" * 80)
422
+ print(" PRUNING COMPLETE")
423
+ print("=" * 80)
424
+ print(f"\n RESULTS:")
425
+ print(f" Original magnitude: {base_magnitude:.0f}")
426
+ print(f" Final magnitude: {final_magnitude:.0f}")
427
+ print(f" Reduction: {reduction_pct:.2f}%")
428
+ print(f" Total reductions: {total_reductions}")
429
+ print(f" Original non-zero: {nonzero_params}")
430
+ print(f" Final non-zero: {final_nonzero}")
431
+ print(f" Zeros created: {nonzero_params - final_nonzero}")
432
+ print(f" Max weight: {final_max:.0f}")
433
+ print(f" Total time: {format_time(pruning_time)}")
434
+
435
+ # Save
436
+ print(f"\n SAVING to {checkpoint_path}...")
437
+ save_file(model, checkpoint_path)
438
+ print(f" Saved.")
439
+
440
+ # Final verification
441
+ print(f"\n FINAL VERIFICATION...")
442
+ from safetensors import safe_open
443
+ f = safe_open(checkpoint_path, framework='numpy')
444
+ verify_model = {name: torch.tensor(f.get_tensor(name)).float() for name in f.keys()}
445
+ verify_fitness = check_fitness(verify_model, evaluator, device)
446
+ print(f" Fitness: {verify_fitness:.6f}")
447
+
448
+ if verify_fitness >= 0.9999:
449
+ print(f" STATUS: PASS")
450
+ else:
451
+ print(f" STATUS: FAIL - Model corrupted!")
452
+
453
+ print("\n" + "=" * 80)
454
+ return model
455
+
456
+
457
+ MAX_BATCH_SIZE = 80000
458
+
459
+ if __name__ == "__main__":
460
+ parser = argparse.ArgumentParser(description='Batched Weight Pruning')
461
+ parser.add_argument('--passes', type=int, default=10,
462
+ help='Maximum pruning passes (default: 10)')
463
+ parser.add_argument('--batch_size', type=int, default=80000,
464
+ help=f'Batch size for parallel evaluation (default: 80000, max: {MAX_BATCH_SIZE})')
465
+ parser.add_argument('--device', type=str, default='cuda',
466
+ help='Device: cuda or cpu (default: cuda)')
467
+ parser.add_argument('--output', type=str, default=DEFAULT_OUTPUT,
468
+ help='Output path')
469
+ args = parser.parse_args()
470
+
471
+ if args.batch_size > MAX_BATCH_SIZE:
472
+ print(f"WARNING: batch_size {args.batch_size} exceeds maximum {MAX_BATCH_SIZE}. Clamping.")
473
+ args.batch_size = MAX_BATCH_SIZE
474
+
475
+ print(f"\nStarting at {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
476
+
477
+ prune_weights(
478
+ passes=args.passes,
479
+ batch_size=args.batch_size,
480
+ device=args.device,
481
+ checkpoint_path=args.output
482
+ )
483
+
484
+ print(f"\nFinished at {time.strftime('%Y-%m-%d %H:%M:%S')}")
quantize.py ADDED
@@ -0,0 +1,406 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Quantize threshold-computer safetensors to the minimum signed integer
3
+ dtype that exactly represents each tensor.
4
+
5
+ Weights and biases in this library are integer-valued by construction,
6
+ with one historical exception: a handful of legacy buffer gates use a
7
+ bias of -0.5 (e.g. arithmetic.asr8bit.bit*.bias). For binary inputs,
8
+ H(x - 0.5) and H(x - 1) are identical, so those biases are floored to
9
+ -1 before casting.
10
+
11
+ This is a packaging optimization, not a precision change: the eval
12
+ pipeline already promotes weights to float32 on load, so integer
13
+ storage is exact.
14
+
15
+ The --ternary flag also rewrites single-input weight=+/-2 identity
16
+ buffers (SHL/SHR/ROL/ROR bit gates, stack data buffers, RET address
17
+ buffers, flag buffers) to weight=+/-1 with bias adjusted as needed to
18
+ preserve heaviside output for binary inputs. After this pass every
19
+ weight tensor in the file lies in {-1, 0, 1} except for positional
20
+ comparators and a few hand-constructed modular arithmetic circuits
21
+ (see the violation report); fully ternarizing those requires bit-
22
+ cascading in build.py.
23
+
24
+ Usage:
25
+ python quantize.py path/to/file.safetensors # in-place
26
+ python quantize.py path/to/file.safetensors -o out.safetensors # to new file
27
+ python quantize.py variants/ # whole directory in place
28
+ python quantize.py variants/ -o variants_int/ # whole directory to new dir
29
+ python quantize.py file.safetensors --ternary # try ternary weights
30
+ python quantize.py file.safetensors --ternary --strict # error if any weight non-ternary
31
+ """
32
+
33
+ from __future__ import annotations
34
+
35
+ import argparse
36
+ import sys
37
+ from pathlib import Path
38
+ from typing import Dict, Tuple
39
+
40
+ import torch
41
+ from safetensors import safe_open
42
+ from safetensors.torch import save_file
43
+
44
+ DTYPES = [
45
+ (torch.int8, -(1 << 7), (1 << 7) - 1),
46
+ (torch.int16, -(1 << 15), (1 << 15) - 1),
47
+ (torch.int32, -(1 << 31), (1 << 31) - 1),
48
+ (torch.int64, None, None), # always fits
49
+ ]
50
+
51
+
52
+ def _normalize_to_int(tensor: torch.Tensor) -> torch.Tensor:
53
+ """Return a tensor with strictly integer values, floored from any
54
+ half-integer values. Floor (not round) because a -0.5 bias must
55
+ become -1 (not 0) to preserve H(x + bias) for binary x."""
56
+ if not tensor.dtype.is_floating_point:
57
+ return tensor.to(torch.float64) # promote for range checks
58
+ tf = tensor.to(torch.float64)
59
+ rounded = tf.round()
60
+ if torch.equal(rounded, tf):
61
+ return tf
62
+ doubled = tf * 2.0
63
+ if torch.equal(doubled.round(), doubled):
64
+ return torch.floor(tf)
65
+ raise ValueError(
66
+ f"tensor has non-half-integer values; range "
67
+ f"[{tf.min().item()}, {tf.max().item()}]"
68
+ )
69
+
70
+
71
+ def _min_signed_int_dtype(tensor: torch.Tensor) -> torch.dtype:
72
+ if tensor.numel() == 0:
73
+ return torch.int8
74
+ lo = int(tensor.min().item())
75
+ hi = int(tensor.max().item())
76
+ for dtype, lo_lim, hi_lim in DTYPES:
77
+ if lo_lim is None or (lo_lim <= lo and hi <= hi_lim):
78
+ return dtype
79
+ return torch.int64
80
+
81
+
82
+ def _ternarize_modular_and_patterns(
83
+ tensors: Dict[str, torch.Tensor],
84
+ ) -> Tuple[Dict[str, torch.Tensor], Dict]:
85
+ """Replace seed-file modular detectors and pattern_recognition gates
86
+ with bit-cascade-equivalent ternary structures.
87
+
88
+ Modular: for each modulus N in {3,5,6,7,9,10,11,12}, replace the
89
+ layer1.geq{i}/layer1.leq{i}/layer2.eq{i}/layer3.or chain with a
90
+ bit-cascade equality detector per multiple of N in [0, 256), then
91
+ OR all detectors together. Top-level gate stays named
92
+ `modular.mod{N}` (a multi-input OR over per-multiple matches).
93
+
94
+ Pattern_recognition.leadingones / trailingones are dropped: they are
95
+ seed-file artifacts with no eval coverage and no downstream
96
+ consumers in this codebase.
97
+
98
+ Moduli whose bit-cascade structure is already present are left
99
+ untouched (including their .inputs routing metadata); the rebuild
100
+ only fires on the legacy layer1/layer2/layer3 layout.
101
+ """
102
+ new_tensors = dict(tensors)
103
+
104
+ # --- pattern_recognition: drop leadingones/trailingones ---
105
+ pr_dropped = 0
106
+ for k in list(new_tensors.keys()):
107
+ if (k.startswith("pattern_recognition.leadingones")
108
+ or k.startswith("pattern_recognition.trailingones")):
109
+ del new_tensors[k]
110
+ pr_dropped += 1
111
+
112
+ # --- modular: rebuild as ternary bit-cascade equality per multiple ---
113
+ moduli = [3, 5, 6, 7, 9, 10, 11, 12]
114
+ mod_gates_added = 0
115
+ for mod in moduli:
116
+ prefix = f"modular.mod{mod}"
117
+ if f"{prefix}.eq.k0.all.weight" in new_tensors:
118
+ continue # already bit-cascade; rebuilding would only shed metadata
119
+ # Drop old structure
120
+ for k in list(new_tensors.keys()):
121
+ if k.startswith(prefix + "."):
122
+ del new_tensors[k]
123
+
124
+ multiples = list(range(0, 256, mod))
125
+ # Per-bit match gates + per-multiple AND
126
+ for k in multiples:
127
+ for i in range(8): # i=0 is MSB (matches inputs MSB-first ordering)
128
+ k_bit = (k >> (7 - i)) & 1
129
+ if k_bit == 1:
130
+ # bit_match = x[i]: H(x - 0.5) ~~ identity for binary;
131
+ # use weight=1, bias=-1 -> H(x-1): x=0 -> 0, x=1 -> 1.
132
+ new_tensors[f"{prefix}.eq.k{k}.bit{i}.match.weight"] = torch.tensor([1.0], dtype=torch.float64)
133
+ new_tensors[f"{prefix}.eq.k{k}.bit{i}.match.bias"] = torch.tensor([-1.0], dtype=torch.float64)
134
+ else:
135
+ # bit_match = NOT x[i]: weight=-1, bias=0 -> H(-x): x=0 -> 1, x=1 -> 0.
136
+ new_tensors[f"{prefix}.eq.k{k}.bit{i}.match.weight"] = torch.tensor([-1.0], dtype=torch.float64)
137
+ new_tensors[f"{prefix}.eq.k{k}.bit{i}.match.bias"] = torch.tensor([0.0], dtype=torch.float64)
138
+ mod_gates_added += 1
139
+ # AND of all 8 bit-matches: weights [1]*8, bias -8
140
+ new_tensors[f"{prefix}.eq.k{k}.all.weight"] = torch.tensor([1.0] * 8, dtype=torch.float64)
141
+ new_tensors[f"{prefix}.eq.k{k}.all.bias"] = torch.tensor([-8.0], dtype=torch.float64)
142
+ mod_gates_added += 1
143
+
144
+ # Final OR over all per-multiple match outputs
145
+ m = len(multiples)
146
+ new_tensors[f"{prefix}.weight"] = torch.tensor([1.0] * m, dtype=torch.float64)
147
+ new_tensors[f"{prefix}.bias"] = torch.tensor([-1.0], dtype=torch.float64)
148
+ mod_gates_added += 1
149
+
150
+ return new_tensors, {
151
+ "pattern_recognition_dropped": pr_dropped,
152
+ "modular_gates_added": mod_gates_added,
153
+ "modular_moduli": len(moduli),
154
+ }
155
+
156
+
157
+ def _ternarize_buffers(
158
+ tensors: Dict[str, torch.Tensor],
159
+ ) -> Tuple[Dict[str, torch.Tensor], Dict]:
160
+ """Rewrite single-input weight=+-2 identity buffers as weight=+-1 with
161
+ bias adjusted to preserve heaviside output for binary inputs.
162
+
163
+ For a single-input gate H(w*x + b) with x in {0, 1}, the only thing
164
+ that matters is the pair (H(b), H(w + b)). Pick the smallest integer
165
+ bias b' such that (H(b'), H(sgn + b')) matches, with sgn = sign(w).
166
+
167
+ Returns (new_tensors, stats). stats has 'fixed', 'failed', 'failed_names'.
168
+ """
169
+ new_tensors = dict(tensors)
170
+ fixed = 0
171
+ failed_names = []
172
+
173
+ weight_keys = [k for k in tensors if k.endswith(".weight")]
174
+ for wkey in weight_keys:
175
+ w = tensors[wkey]
176
+ wf = w.float() if w.dtype.is_floating_point else w.to(torch.float64).float()
177
+ if (wf.abs() <= 1.0).all():
178
+ continue # already ternary
179
+
180
+ gate = wkey[: -len(".weight")]
181
+ bkey = gate + ".bias"
182
+
183
+ # Single-input weight=+-2 buffer with single bias
184
+ if (
185
+ wf.numel() == 1
186
+ and abs(wf.item()) == 2.0
187
+ and bkey in tensors
188
+ and tensors[bkey].numel() == 1
189
+ ):
190
+ w_val = wf.item()
191
+ b_val = float(tensors[bkey].float().item())
192
+ sgn = 1.0 if w_val > 0 else -1.0
193
+ x0_target = 1 if b_val >= 0 else 0
194
+ x1_target = 1 if (w_val + b_val) >= 0 else 0
195
+ chosen = None
196
+ # Prefer keeping the bias unchanged when possible
197
+ for b_new in [int(b_val), int(b_val) - 1, -1, 0, -2, 1, -3, 2]:
198
+ x0 = 1 if b_new >= 0 else 0
199
+ x1 = 1 if (sgn + b_new) >= 0 else 0
200
+ if x0 == x0_target and x1 == x1_target:
201
+ chosen = b_new
202
+ break
203
+ if chosen is not None:
204
+ new_tensors[wkey] = torch.tensor([sgn], dtype=torch.float64)
205
+ new_tensors[bkey] = torch.tensor([float(chosen)], dtype=torch.float64)
206
+ fixed += 1
207
+ continue
208
+
209
+ failed_names.append(wkey)
210
+
211
+ return new_tensors, {"fixed": fixed, "failed_names": failed_names}
212
+
213
+
214
+ def quantize_tensors(
215
+ tensors: Dict[str, torch.Tensor],
216
+ ternary: bool = False,
217
+ ) -> Tuple[Dict[str, torch.Tensor], Dict[str, int], Tuple[int, int], Dict]:
218
+ """Quantize a dict of tensors. Returns
219
+ (new_tensors, dtype_counts, (bytes_before, bytes_after), ternary_stats)."""
220
+ ternary_stats: Dict = {"applied": False, "fixed": 0, "failed_names": []}
221
+ if ternary:
222
+ tensors, ternary_stats = _ternarize_buffers(tensors)
223
+ ternary_stats["applied"] = True
224
+ # Also rebuild modular detectors and drop pattern_recognition stragglers.
225
+ tensors, mod_stats = _ternarize_modular_and_patterns(tensors)
226
+ ternary_stats["modular_gates_added"] = mod_stats["modular_gates_added"]
227
+ ternary_stats["pattern_recognition_dropped"] = mod_stats["pattern_recognition_dropped"]
228
+
229
+ new_tensors: Dict[str, torch.Tensor] = {}
230
+ counts: Dict[str, int] = {"int8": 0, "int16": 0, "int32": 0, "int64": 0,
231
+ "manifest_kept": 0, "skipped": 0}
232
+ bytes_before = 0
233
+ bytes_after = 0
234
+
235
+ for name, t in tensors.items():
236
+ bytes_before += t.numel() * t.element_size()
237
+
238
+ if name.startswith("manifest."):
239
+ new_tensors[name] = t
240
+ counts["manifest_kept"] += 1
241
+ bytes_after += t.numel() * t.element_size()
242
+ continue
243
+
244
+ try:
245
+ normalized = _normalize_to_int(t)
246
+ except ValueError:
247
+ new_tensors[name] = t
248
+ counts["skipped"] += 1
249
+ bytes_after += t.numel() * t.element_size()
250
+ continue
251
+
252
+ target = _min_signed_int_dtype(normalized)
253
+ cast = normalized.to(target)
254
+ new_tensors[name] = cast
255
+ bytes_after += cast.numel() * cast.element_size()
256
+ counts[str(target).replace("torch.", "")] += 1
257
+
258
+ return new_tensors, counts, (bytes_before, bytes_after), ternary_stats
259
+
260
+
261
+ def quantize_file(in_path: Path, out_path: Path, verbose: bool = False,
262
+ ternary: bool = False, strict_ternary: bool = False) -> Dict:
263
+ file_before = in_path.stat().st_size
264
+ tensors: Dict[str, torch.Tensor] = {}
265
+ metadata: Dict[str, str] = {}
266
+ with safe_open(str(in_path), framework="pt") as f:
267
+ meta = f.metadata()
268
+ if meta:
269
+ metadata = dict(meta)
270
+ for name in f.keys():
271
+ # clone so the source mmap can be released before we write
272
+ tensors[name] = f.get_tensor(name).clone()
273
+
274
+ new_tensors, counts, (before, after), tstats = quantize_tensors(tensors, ternary=ternary)
275
+
276
+ # Audit final ternary status (count of weight tensors with |w| > 1)
277
+ final_nonternary = []
278
+ for k, v in new_tensors.items():
279
+ if not k.endswith(".weight"):
280
+ continue
281
+ if k.startswith("manifest."):
282
+ continue
283
+ vf = v.float() if v.dtype.is_floating_point else v.to(torch.float64).float()
284
+ if (vf.abs() > 1.0).any():
285
+ final_nonternary.append(k)
286
+
287
+ if ternary and strict_ternary and final_nonternary:
288
+ raise ValueError(
289
+ f"--strict failed: {len(final_nonternary)} weight tensors are not "
290
+ f"ternary after transformation; first: {final_nonternary[:5]}"
291
+ )
292
+
293
+ # Drop the original mmap-backed tensors before writing in-place.
294
+ del tensors
295
+ out_path.parent.mkdir(parents=True, exist_ok=True)
296
+ if ternary:
297
+ # Note ternary mode in metadata so downstream tools can see it
298
+ if metadata is None:
299
+ metadata = {}
300
+ metadata = dict(metadata)
301
+ metadata["weight_quantization"] = (
302
+ "ternary_partial" if final_nonternary else "ternary"
303
+ )
304
+ save_file(new_tensors, str(out_path), metadata=metadata or None)
305
+
306
+ file_after = out_path.stat().st_size
307
+ return {
308
+ "in_path": str(in_path),
309
+ "out_path": str(out_path),
310
+ "tensor_counts": counts,
311
+ "tensor_bytes_before": before,
312
+ "tensor_bytes_after": after,
313
+ "file_size_before": file_before,
314
+ "file_size_after": file_after,
315
+ "ternary": tstats,
316
+ "final_nonternary": final_nonternary,
317
+ }
318
+
319
+
320
+ def _print_summary(label: str, info: Dict) -> None:
321
+ cb = info["tensor_bytes_before"]
322
+ ca = info["tensor_bytes_after"]
323
+ fb = info["file_size_before"]
324
+ fa = info["file_size_after"]
325
+ counts = info["tensor_counts"]
326
+ bucket_str = " ".join(f"{k}={v}" for k, v in counts.items() if v)
327
+ ratio_t = cb / ca if ca else 1.0
328
+ ratio_f = fb / fa if fa else 1.0
329
+ print(
330
+ f" {label}: file {fb / 1e6:6.1f} MB -> {fa / 1e6:6.1f} MB "
331
+ f"({ratio_f:.2f}x); tensor data {cb / 1e6:6.1f} MB -> {ca / 1e6:6.1f} MB "
332
+ f"({ratio_t:.2f}x)"
333
+ )
334
+ print(f" {bucket_str}")
335
+ if info.get("ternary", {}).get("applied"):
336
+ ts = info["ternary"]
337
+ nt = info["final_nonternary"]
338
+ print(f" ternary: {ts['fixed']} buffer gates rewritten; "
339
+ f"{len(nt)} weight tensors remain non-ternary")
340
+
341
+
342
+ def main() -> int:
343
+ parser = argparse.ArgumentParser(description="Quantize safetensors to min signed int dtype")
344
+ parser.add_argument("input", type=Path, help=".safetensors file or directory of files")
345
+ parser.add_argument("-o", "--output", type=Path, default=None,
346
+ help="output file or directory (default: in-place)")
347
+ parser.add_argument("-v", "--verbose", action="store_true")
348
+ parser.add_argument("--ternary", action="store_true",
349
+ help="Rewrite single-input weight=+/-2 buffers as +/-1 to push toward "
350
+ "ternary {-1, 0, 1} weights and report remaining violations")
351
+ parser.add_argument("--strict", action="store_true",
352
+ help="With --ternary, fail if any weight tensor is still non-ternary")
353
+ parser.add_argument("--report-violations", type=int, default=0, metavar="N",
354
+ help="Print first N non-ternary weight tensor names per file")
355
+ args = parser.parse_args()
356
+
357
+ inputs = []
358
+ if args.input.is_dir():
359
+ inputs = sorted(p for p in args.input.glob("*.safetensors"))
360
+ elif args.input.is_file():
361
+ inputs = [args.input]
362
+ else:
363
+ print(f"not found: {args.input}", file=sys.stderr)
364
+ return 2
365
+
366
+ if not inputs:
367
+ print(f"no .safetensors files under {args.input}", file=sys.stderr)
368
+ return 2
369
+
370
+ if args.output is None:
371
+ outputs = inputs # in-place
372
+ elif args.output.suffix == ".safetensors":
373
+ if len(inputs) != 1:
374
+ print("output is a single file but input is a directory; pass a directory output", file=sys.stderr)
375
+ return 2
376
+ outputs = [args.output]
377
+ else:
378
+ args.output.mkdir(parents=True, exist_ok=True)
379
+ outputs = [args.output / p.name for p in inputs]
380
+
381
+ total_before = 0
382
+ total_after = 0
383
+ print(f"Quantizing {len(inputs)} file(s)" + (" (ternary mode)" if args.ternary else "") + "\n")
384
+ for src, dst in zip(inputs, outputs):
385
+ info = quantize_file(src, dst, verbose=args.verbose,
386
+ ternary=args.ternary, strict_ternary=args.strict)
387
+ _print_summary(src.name, info)
388
+ if args.report_violations and info.get("final_nonternary"):
389
+ for name in info["final_nonternary"][: args.report_violations]:
390
+ print(f" non-ternary: {name}")
391
+ if len(info["final_nonternary"]) > args.report_violations:
392
+ print(f" ... and {len(info['final_nonternary']) - args.report_violations} more")
393
+ total_before += info["file_size_before"]
394
+ total_after += info["file_size_after"]
395
+
396
+ print()
397
+ print("=" * 76)
398
+ print(
399
+ f"Total: {total_before / 1e6:.1f} MB -> {total_after / 1e6:.1f} MB "
400
+ f"({total_before / max(total_after, 1):.2f}x reduction)"
401
+ )
402
+ return 0
403
+
404
+
405
+ if __name__ == "__main__":
406
+ sys.exit(main())
routing/routing.json ADDED
The diff for this file is too large to render. See raw diff
 
test_cpu.py ADDED
@@ -0,0 +1,138 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Run the CPU program suite against a threshold-computer variant.
3
+
4
+ Loads weights, instantiates GenericThresholdCPU, runs each program from
5
+ cpu_programs.SUITE, and verifies expected memory contents at HALT.
6
+
7
+ Usage:
8
+ python test_cpu.py # default: 1KB variant, fast
9
+ python test_cpu.py --model neural_computer.safetensors # 64KB canonical, slow
10
+ python test_cpu.py --only fib,sum_n # subset of suite
11
+ """
12
+
13
+ from __future__ import annotations
14
+ import argparse
15
+ import os
16
+ import sys
17
+ import time
18
+ from pathlib import Path
19
+
20
+ import torch
21
+ from safetensors import safe_open
22
+
23
+ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
24
+ from eval_all import GenericThresholdCPU, get_manifest
25
+ from cpu_programs import SUITE
26
+
27
+
28
+ def load_tensors(path: Path):
29
+ out = {}
30
+ with safe_open(str(path), framework="pt") as f:
31
+ for name in f.keys():
32
+ out[name] = f.get_tensor(name).float()
33
+ return out
34
+
35
+
36
+ def run_program(cpu: GenericThresholdCPU, mem, max_cycles: int):
37
+ addr_mask = (1 << cpu.addr_bits) - 1
38
+ state = {
39
+ "pc": 0,
40
+ "regs": [0] * 4,
41
+ "flags": [0] * 4,
42
+ "mem": list(mem),
43
+ "halted": False,
44
+ "sp": addr_mask,
45
+ }
46
+ t0 = time.perf_counter()
47
+ final, cycles = cpu.run(state, max_cycles=max_cycles)
48
+ return final, cycles, time.perf_counter() - t0
49
+
50
+
51
+ def check_expected(final, expected: dict) -> tuple[bool, list[str]]:
52
+ failures = []
53
+ for addr, want in expected.items():
54
+ got = final["mem"][addr]
55
+ if got != want:
56
+ failures.append(f"M[0x{addr:04X}] = {got} (expected {want})")
57
+ return (len(failures) == 0), failures
58
+
59
+
60
+ def main() -> int:
61
+ parser = argparse.ArgumentParser(description="Run the CPU program suite")
62
+ parser.add_argument(
63
+ "--model", type=str,
64
+ default=os.path.join(os.path.dirname(__file__),
65
+ "variants", "neural_computer8_small.safetensors"),
66
+ help="Path to .safetensors variant",
67
+ )
68
+ parser.add_argument(
69
+ "--only", type=str, default="",
70
+ help="Comma-separated subset of program names to run",
71
+ )
72
+ args = parser.parse_args()
73
+
74
+ print(f"Loading {args.model}")
75
+ tensors = load_tensors(Path(args.model))
76
+ manifest = get_manifest(tensors)
77
+ print(f"Manifest: data={manifest['data_bits']}-bit, addr={manifest['addr_bits']}-bit, "
78
+ f"mem={manifest['memory_bytes']}B")
79
+
80
+ if manifest["memory_bytes"] < 256:
81
+ print(f"ERROR: variant has {manifest['memory_bytes']}B memory; "
82
+ f"the suite needs at least 256B (scratchpad).")
83
+ return 2
84
+
85
+ cpu = GenericThresholdCPU(tensors)
86
+ only = set(s.strip() for s in args.only.split(",") if s.strip())
87
+
88
+ print()
89
+ print("=" * 80)
90
+ print(f" CPU PROGRAM SUITE ({manifest['memory_bytes']}B mem, {manifest['data_bits']}-bit ALU)")
91
+ print("=" * 80)
92
+
93
+ pass_count = 0
94
+ fail_count = 0
95
+ skip_count = 0
96
+
97
+ for name, builder in SUITE:
98
+ if only and name not in only:
99
+ skip_count += 1
100
+ continue
101
+ try:
102
+ mem, expected, max_cycles, desc = builder(manifest["memory_bytes"])
103
+ except Exception as e:
104
+ print(f" {name:18} BUILD ERROR: {e}")
105
+ fail_count += 1
106
+ continue
107
+
108
+ if len(mem) != manifest["memory_bytes"]:
109
+ print(f" {name:18} SKIP (program built {len(mem)}B, "
110
+ f"variant has {manifest['memory_bytes']}B)")
111
+ skip_count += 1
112
+ continue
113
+
114
+ final, cycles, elapsed = run_program(cpu, mem, max_cycles)
115
+ ok, failures = check_expected(final, expected)
116
+ if ok and not final["halted"]:
117
+ ok = False
118
+ failures.append(f"did not HALT within {max_cycles} cycles")
119
+
120
+ status = "PASS" if ok else "FAIL"
121
+ if ok:
122
+ pass_count += 1
123
+ else:
124
+ fail_count += 1
125
+ print(f" {name:18} {status} ({cycles:>3} cyc, {elapsed:>5.2f}s) {desc}")
126
+ if not ok:
127
+ for f in failures[:6]:
128
+ print(f" - {f}")
129
+
130
+ print()
131
+ print("=" * 80)
132
+ total = pass_count + fail_count
133
+ print(f" PASS: {pass_count}/{total} FAIL: {fail_count}/{total} SKIP: {skip_count}")
134
+ return 0 if fail_count == 0 else 1
135
+
136
+
137
+ if __name__ == "__main__":
138
+ sys.exit(main())
tests/__init__.py ADDED
File without changes
tests/test_play.py ADDED
@@ -0,0 +1,72 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Smoke tests for play.py.
2
+
3
+ Drives play.main() against the smallest CPU variant (neural_computer8_small,
4
+ 1 KB memory) and the smallest ALU-only variant (neural_alu8). Each demo
5
+ prints a deterministic block on stdout; we capture and grep for the success
6
+ markers. Catches the kind of drift where play.py outpaces a refactor of the
7
+ gate naming convention (cmp8bit, mod5, etc.) without anyone running it.
8
+ """
9
+ from __future__ import annotations
10
+
11
+ import io
12
+ import os
13
+ import sys
14
+ import unittest
15
+ from contextlib import redirect_stdout
16
+
17
+ REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
18
+ if REPO_ROOT not in sys.path:
19
+ sys.path.insert(0, REPO_ROOT)
20
+
21
+ import play # noqa: E402
22
+
23
+
24
+ def _run(model_path: str, skip_cpu: bool) -> str:
25
+ argv_save = sys.argv
26
+ sys.argv = ["play.py", "--model", model_path]
27
+ if skip_cpu:
28
+ sys.argv.append("--skip-cpu")
29
+ buf = io.StringIO()
30
+ try:
31
+ with redirect_stdout(buf):
32
+ rc = play.main()
33
+ finally:
34
+ sys.argv = argv_save
35
+ if rc not in (0, None):
36
+ raise AssertionError(f"play.main() returned {rc}")
37
+ return buf.getvalue()
38
+
39
+
40
+ class TestPlay(unittest.TestCase):
41
+
42
+ def test_alu_variant(self) -> None:
43
+ path = os.path.join(REPO_ROOT, "variants", "neural_alu8.safetensors")
44
+ out = _run(path, skip_cpu=True)
45
+ # Booleans
46
+ for line in ["and 00->0 01->0 10->0 11->1",
47
+ "or 00->0 01->1 10->1 11->1",
48
+ "xor 00->0 01->1 10->1 11->0"]:
49
+ self.assertIn(line, out, f"missing boolean line: {line!r}")
50
+ # ALU ADD/SUB
51
+ self.assertIn("127 + 128 = 255 (carry=0) expected 255 [OK]", out)
52
+ self.assertIn("127 - 128 = 255 (no_borrow=0) expected 255 [OK]", out)
53
+ # CMP via cmp8bit cascade
54
+ self.assertIn("50 vs 30 -> GT=1 LT=0 EQ=0", out)
55
+ self.assertIn("77 vs 77 -> GT=0 LT=0 EQ=1", out)
56
+ self.assertIn("128 vs 127 -> GT=1 LT=0 EQ=0", out)
57
+ # MUL
58
+ self.assertIn("12 * 11 = 132 expected 132 [OK]", out)
59
+ # mod-5 (52 multiples of 5 in [0, 255])
60
+ self.assertIn("52 hits", out)
61
+ self.assertNotIn("FAIL", out)
62
+
63
+ def test_cpu_variant(self) -> None:
64
+ path = os.path.join(REPO_ROOT, "variants", "neural_computer8_small.safetensors")
65
+ out = _run(path, skip_cpu=False)
66
+ self.assertIn("M[0x0083] = 15 (expected 15) [OK]", out)
67
+ self.assertIn("R0=15", out)
68
+ self.assertNotIn("FAIL", out)
69
+
70
+
71
+ if __name__ == "__main__":
72
+ unittest.main()
todo.md ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # todo
2
+
3
+ Unfinished work.
4
+
5
+ - The float add, multiply, and divide pipelines handle subnormal operands and produce subnormal results rather than flushing them to zero.
6
+ - The F extension includes a fused multiply-add instruction that rounds once.
7
+ - FCVT.W.S honors the instruction's rounding-mode field instead of always rounding toward zero.
8
+ - Program-counter sequencing and instruction decode are computed by threshold gates rather than by fixed wiring in the runtimes.
9
+ - The standalone float normalize stages, superseded by the composed pipelines' own normalizers, are removed from the circuit inventory.
10
+ - Every gate in the 32-bit multiply circuit resolves its .inputs to a defined signal.
11
+ - A machine-checked correctness proof relates the shipped weights to the ISA.
12
+ - The processor's full state transition is expressed as a recurrent ternary threshold network and executes as analog matrix-vector products rather than a wired gate graph.
13
+ - The verified gate wiring is compiled into a fixed sequence of ternary weight matrices W1..Wk with a Heaviside step between them, acting on a state vector that holds the program counter, registers, flags, and every memory bit.
14
+ - One instruction is that fixed stack of threshold layers, and the machine is the stack iterated until the halt bit sets, so the whole processor is a single recurrent linear-threshold network.
15
+ - The matrix form is proven equal to the gate-graph processor, so the ISA semantics carry over bit for bit.
16
+ - Each weight matrix maps to a resistive or photonic crossbar with conductances drawn from {-1, 0, +1}, and one clock cycle is one analog matrix-vector product followed by thresholding.
17
+ - The analog realization is shown to reproduce the digital circuit within the device's conductance-quantization and noise margins.
18
+ - A program running on the processor emits the exact tensor file that defines that processor, and builds any other machine in the family from a description of it.
19
+ - The program's output stream is the byte-for-byte safetensors file of the running machine, and that equality is machine-checked rather than observed on a single run.
20
+ - The same program takes an encoded description of any machine in the family and constructs that machine's weight file, so the constructor is universal and the description selects what is built.
21
+ - Self-reproduction is the case where the description handed to the constructor is its own.
variants/neural_alu16.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:86cc54fedb8614bda4382d90ade7e505a29660c278f1cdf4821ac4ecbfb043e2
3
+ size 25794970
variants/neural_alu32.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:89cc066ba4b2adc44af57511cfd435a9bbde16e9c13c71f91774239fbace43eb
3
+ size 24181679
variants/neural_alu8.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1e339ff3e200428ccfaa933a6c07cfc81006e9dab53984708d6105923a154e6d
3
+ size 24181679
variants/neural_computer16.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6efffb459c157a77e8f98cff5eb298491f47f567b36e97aa34159a103c11bcbd
3
+ size 34312872
variants/neural_computer16_reduced.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d9c842b6270db3bc3d1a5b94a1bfcd0b0c62766aa1083cb65aed940b4984e3a8
3
+ size 26501432
variants/neural_computer16_registers.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5b26e149ccab217d877257959278bcd8cd0e262c7a9c9976a7020e38edcc0eed
3
+ size 25882024
variants/neural_computer16_scratchpad.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ecbd7f1247ae12240059e5cecbe8b9336fa197a961f0babf3154775ce3a47d97
3
+ size 25966016
variants/neural_computer16_small.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c29e64ecab77e5c1e1809393ecdf4a0a34e4bf6a8b474b5329050d6662de0783
3
+ size 26085300
variants/neural_computer32.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fb4725ab213771fa0f20c42d13a43b913cb33f1762c6a8c4b4670856f8d3592c
3
+ size 32717681
variants/neural_computer32_reduced.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ab3e572ebd2d9cbddb21fa7612cea8024131563ef873a5db9c99c9d5b9957926
3
+ size 24906261
variants/neural_computer32_registers.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fa766185ccae6f27174a12409255960b7732d9a4c9cb683a00f6a9e233bd417d
3
+ size 24286869
variants/neural_computer32_scratchpad.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:7df7be2cdea56ea2a0163276a99d8116af68d4b4eed6c5729e801b2a1c67acd7
3
+ size 24371681
variants/neural_computer32_small.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fa705c60ea36f8f54a1211a326e4d90a7e0f192ac456039e290ec5c9b9b61451
3
+ size 24490921
variants/neural_computer8.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f5f5c79d0e397ca74fb22dd48804b620f57d6194d08e542205134594bd3359a3
3
+ size 32689549
variants/neural_computer8_reduced.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ded8f1991b9f159bd917dff96e38f376c62a20ed06ab0345416b2dac63393424
3
+ size 24878061
variants/neural_computer8_registers.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:1384db2751a2031a0a9856aba10140adfb1cb91e402fe3e54a12a0af4286e8ef
3
+ size 24258649
variants/neural_computer8_scratchpad.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:420d540595155c09be22c61900bfafd49748be87faea8ae938f4ef7b4ad3d021
3
+ size 24343497
variants/neural_computer8_small.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e8d2bbdcf4a7ca1a8bf7483f3ecacb35e45f88a75daa6cc437b2976ce34a5001
3
+ size 24462729
variants/neural_rv32.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3367ac4fd179dc05f6ed17822e327206e0ff2ab6d64c2d796448d79ca805b52b
3
+ size 33274820
variants/neural_subleq8.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3afb7184f4046421b7a6f96663227b9ef3e646389022affe8ddbb57eb52cbca5
3
+ size 650280