File size: 3,148 Bytes
8f34e5f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """neural_matrix8 -- a fleet of processors as one batched matrix product.
The whole CPU is compiled to 108 ternary weight matrices with a Heaviside step
between them, so one clock cycle is one matrix-vector product plus a threshold.
Add a batch dimension and N processors are the same 108 matmuls: here 65,536
independent CPUs step in lockstep on the GPU. A 256-wide correctness fleet
confirms every CPU halts exactly at cycle 2n+1 for its own countdown input.
python demos/neural_matrix8_gpu_cpu_fleet.py
"""
import os, sys, time
HERE = os.path.dirname(os.path.abspath(__file__))
REPO = os.path.dirname(HERE)
sys.path.insert(0, os.path.join(REPO, "src"))
import torch
from matrix8 import MatrixMachine, state_to_vec, _mk_state, _instr
def prog_bytes():
return (_instr(0x1, 0, 1) # SUB R0, R1
+ _instr(0xD, 0, 0, 1) + [0, 0] # JNZ 0x0000 (address word)
+ _instr(0xF) # HALT
+ [0] * 8)
if __name__ == "__main__":
dev = "cuda" if torch.cuda.is_available() else "cpu"
mm = MatrixMachine.from_file(device=dev)
n_neurons = sum(int(b.numel()) for b in mm.B)
name = torch.cuda.get_device_name(0) if dev == "cuda" else "cpu"
print("neural_matrix8: 65,536 CPUs as one batched matrix product")
print("=" * 56)
print(f"loaded {len(mm.W)} ternary matrices, {n_neurons:,} threshold neurons "
f"per transition, device={dev} ({name})")
mem = prog_bytes()
# correctness fleet: one CPU per 8-bit input, all 256 at once
V = torch.stack([state_to_vec(_mk_state(mem=mem, regs=(x, 1, 0, 0)))
for x in range(256)]).to(dev)
first_halt = torch.full((256,), -1, dtype=torch.long)
step = 0
while step < 600:
halted = V[:, MatrixMachine.HALT_IDX] > 0.5
first_halt[(first_halt < 0) & halted.cpu()] = step
if bool(halted.all()):
break
V = mm.step(V); step += 1
expect = torch.tensor([2 * (x if x else 256) + 1 for x in range(256)])
ok_halt = bool((first_halt == expect).all())
ok_r0 = bool((V[:, 4:12].cpu().long().sum(1) == 0).all())
print(f"\ncorrectness fleet of 256: every CPU halts at cycle 2n+1 for its own "
f"input: {'EXACT' if ok_halt else 'MISMATCH'}; all results R0==0: {ok_r0}")
# throughput fleet: 65,536 CPUs
B = 65536
V = torch.stack([state_to_vec(_mk_state(mem=mem, regs=(x & 0xFF, 1, 0, 0)))
for x in range(256)]).repeat(256, 1).to(dev)
if dev == "cuda":
torch.cuda.synchronize()
t0, steps = time.perf_counter(), 0
while steps < 520:
V = mm.step(V); steps += 1
if steps % 64 == 0 and bool((V[:, MatrixMachine.HALT_IDX] > 0.5).all()):
break
if dev == "cuda":
torch.cuda.synchronize()
dt = time.perf_counter() - t0
instr = B * steps
print(f"\nthroughput fleet of {B:,}: {steps} transitions in {dt:.1f}s")
print(f" {instr / dt / 1e6:,.1f} million instructions/s aggregate "
f"({instr:,} instructions retired)")
print(f" {instr * n_neurons / dt / 1e9:,.1f} billion threshold-neuron "
f"evaluations/s")
|