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- .gitattributes +36 -0
- .gitignore +4 -0
- README.md +613 -0
- build.py +0 -0
- build_all.py +214 -0
- cpu_programs.py +703 -0
- eval.py +0 -0
- eval_all.py +850 -0
- llm_integration/baseline.py +221 -0
- llm_integration/circuits.py +320 -0
- llm_integration/fitness.py +218 -0
- llm_integration/model.py +1263 -0
- llm_integration/smollm2/SMOLLM2_ARCHITECTURE.md +456 -0
- llm_integration/smollm2/analyze_smollm2.py +232 -0
- llm_integration/smollm2/smollm2_analysis.json +439 -0
- llm_integration/train.py +955 -0
- llm_integration/trained/llm.pt +3 -0
- llm_integration/trained/router.pt +3 -0
- machines.py +1719 -0
- neural_computer.safetensors +3 -0
- play.py +262 -0
- prune_weights.py +484 -0
- quantize.py +406 -0
- routing/routing.json +0 -0
- test_cpu.py +138 -0
- tests/__init__.py +0 -0
- tests/test_play.py +72 -0
- todo.md +21 -0
- variants/neural_alu16.safetensors +3 -0
- variants/neural_alu32.safetensors +3 -0
- variants/neural_alu8.safetensors +3 -0
- variants/neural_computer16.safetensors +3 -0
- variants/neural_computer16_reduced.safetensors +3 -0
- variants/neural_computer16_registers.safetensors +3 -0
- variants/neural_computer16_scratchpad.safetensors +3 -0
- variants/neural_computer16_small.safetensors +3 -0
- variants/neural_computer32.safetensors +3 -0
- variants/neural_computer32_reduced.safetensors +3 -0
- variants/neural_computer32_registers.safetensors +3 -0
- variants/neural_computer32_scratchpad.safetensors +3 -0
- variants/neural_computer32_small.safetensors +3 -0
- variants/neural_computer8.safetensors +3 -0
- variants/neural_computer8_reduced.safetensors +3 -0
- variants/neural_computer8_registers.safetensors +3 -0
- variants/neural_computer8_scratchpad.safetensors +3 -0
- variants/neural_computer8_small.safetensors +3 -0
- variants/neural_rv32.safetensors +3 -0
- 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
|