File size: 7,075 Bytes
e64c5d2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | """Address semantics, working-memory isolation, and neural interface checks."""
import pytest
import torch
from lightnest import rv32
from lightnest.register_patch import RegisterConfig, RegisterPatchNetwork, encode_batch, load_register
def code(words):
return b''.join(w.to_bytes(4,'little') for w in words)
def pair():
words=[rv32.lw('x1',0,'x0'),rv32.lw('x2',4,'x0'),rv32.add('x3','x1','x2')]
words += [rv32.addi('x0','x0',31+i) for i in range(8)]
q=len(words)*4+3;words += [rv32.add('x4','x1','x2')]
words += [rv32.addi('x0','x0',127+i) for i in range(8)]
return [code(words+[rv32.addi('x10',r,0),rv32.EBREAK]) for r in ('x3','x4')],(11,q)
def test_modes_have_identical_initial_tensors_and_state_size_only_changes():
models=[]
for mode in ('no_memory','register32','register128'):
torch.manual_seed(109)
models.append(RegisterPatchNetwork(RegisterConfig(mode=mode)))
assert {sum(p.numel() for p in m.parameters()) for m in models}=={158134}
for name,value in models[0].state_dict().items():
for model in models[1:]:
torch.testing.assert_close(value,model.state_dict()[name],atol=0.,rtol=0.)
assert [m.config.history_slots for m in models]==[1,1,4]
def test_addresses_zero_register_and_history_shift_follow_read_before_write():
torch.set_num_threads(2);torch.manual_seed(19)
model=RegisterPatchNetwork(RegisterConfig(mode='register128')).eval()
raw=code([rv32.lw('x3',0,'x0'),rv32.addi('x3','x3',1),rv32.addi('x3','x0',9),
rv32.add('x0','x3','x3'),rv32.addi('x10','x3',0),rv32.EBREAK])
encoded=encode_batch(['change add to sub'],[raw],model.config)
assert encoded[2].tolist()==[[[3,0,0,1],[3,3,0,1],[3,0,0,1],[0,3,3,0],[10,3,0,1],[0,10,0,0]]]
with torch.inference_mode():
_,_,trace=model(*encoded,return_trace=True)
bank=trace['register_state'];instructions=trace['instructions']
assert torch.count_nonzero(bank[:,0])==0
assert not trace['register_present'][:,0].any()
for slot,instruction in enumerate((2,1,0)):
torch.testing.assert_close(bank[0,3,slot],instructions[0,instruction],atol=0.,rtol=0.)
assert trace['register_present'][0,3].tolist()==[True,True,True,False]
# At the second instruction the new value has not been written yet.
assert trace['history_attention'][0,1,0].tolist()==[1.,0.,0.,0.]
assert float(trace['history_attention'][0,4,0,1:].sum())>0
torch.testing.assert_close(trace['output_context'][0],instructions[0,4],atol=0.,rtol=0.)
def test_no_memory_cannot_read_far_return_but_register_memory_can():
torch.manual_seed(107)
memory=RegisterPatchNetwork(RegisterConfig()).eval()
control=RegisterPatchNetwork(RegisterConfig(mode='no_memory')).eval()
control.load_state_dict(memory.state_dict())
binaries,(p,q)=pair()
encoded=encode_batch(['change the returned add to sub']*2,binaries,memory.config)
with torch.inference_mode():
gs,_,gt=memory(*encoded,return_trace=True)
ls,_,lt=control(*encoded,return_trace=True)
assert torch.count_nonzero(lt['register_state'])==0
assert torch.count_nonzero(lt['history_attention'])==0
assert not lt['register_present'].any()
assert float(((ls[0,p]-ls[0,q])-(ls[1,p]-ls[1,q])).abs())<1e-7
# This is only a structural connection check, not a trained accuracy claim.
assert float(((gs[0,p]-gs[0,q])-(gs[1,p]-gs[1,q])).abs())>1e-6
assert not torch.equal(gt['output_context'][0],gt['output_context'][1])
def test_extra_history_does_not_change_single_write_program():
torch.manual_seed(217)
one=RegisterPatchNetwork(RegisterConfig(mode='register32')).eval()
four=RegisterPatchNetwork(RegisterConfig(mode='register128')).eval()
four.load_state_dict(one.state_dict())
raw=code([rv32.lw('x1',0,'x0'),rv32.lw('x2',4,'x0'),rv32.add('x3','x1','x2'),
rv32.addi('x10','x3',0),rv32.EBREAK])
with torch.inference_mode():
a=one(*encode_batch(['change add to sub'],[raw],one.config))
b=four(*encode_batch(['change add to sub'],[raw],four.config))
for x,y in zip(a,b):
torch.testing.assert_close(x,y,atol=2e-6,rtol=2e-5)
def test_dead_overwrite_cannot_leak_into_register32_return_state():
torch.manual_seed(401)
model=RegisterPatchNetwork(RegisterConfig(mode='register32')).eval()
binaries=[code([rv32.addi('x5','x0',first),rv32.addi('x5','x0',7),
rv32.addi('x10','x5',0),rv32.EBREAK]) for first in (1,2)]
with torch.inference_mode():
_,_,trace=model(*encode_batch(['do not change']*2,binaries,model.config),return_trace=True)
torch.testing.assert_close(trace['instructions'][0,1:],trace['instructions'][1,1:],atol=1e-6,rtol=1e-5)
torch.testing.assert_close(trace['output_context'][0],trace['output_context'][1],atol=1e-6,rtol=1e-5)
# A true read/modify/write must retain the allowed source dependency.
reads_old=[code([rv32.addi('x5','x0',first),rv32.addi('x5','x5',7),
rv32.addi('x10','x5',0),rv32.EBREAK]) for first in (1,2)]
with torch.inference_mode():
_,_,used=model(*encode_batch(['do not change']*2,reads_old,model.config),return_trace=True)
assert not torch.equal(used['output_context'][0],used['output_context'][1])
def test_padding_preserves_final_registers_and_each_forward_resets_memory():
torch.manual_seed(317)
model=RegisterPatchNetwork(RegisterConfig(mode='register128',width=32)).eval()
short=code([rv32.addi('x10','x0',9),rv32.EBREAK]);long=pair()[0][0]
with torch.inference_mode():
a=model(*encode_batch(['no changes'],[short],model.config),return_trace=True)
b=model(*encode_batch(['no changes','change the returned add to sub'],[short,long],model.config),return_trace=True)
model(*encode_batch(['change add to sub'],[long],model.config))
c=model(*encode_batch(['no changes'],[short],model.config),return_trace=True)
for x,y in zip(a[:2],b[:2]):
torch.testing.assert_close(x[0],y[0],atol=2e-6,rtol=2e-5)
torch.testing.assert_close(a[2]['register_state'][0],b[2]['register_state'][0],atol=2e-6,rtol=2e-5)
for x,y in zip(a[:2],c[:2]):
torch.testing.assert_close(x,y,atol=0.,rtol=0.)
def test_checkpoint_roundtrip_and_architecture_rejection(tmp_path):
model=RegisterPatchNetwork(RegisterConfig(mode='register128',width=32))
path=tmp_path/'model.pt';model.save(path);restored=load_register(path)
assert restored.config==model.config
for x,y in zip(model.parameters(),restored.parameters()):
torch.testing.assert_close(x,y,atol=0.,rtol=0.)
saved=torch.load(path,weights_only=True);saved['architecture']='rv32-definition-use-patch-v1';torch.save(saved,path)
with pytest.raises(ValueError,match='architecture'):
load_register(path)
@pytest.mark.parametrize('binary',[b'',b'\x13',code([rv32.beq('x1','x2',4)]),code([rv32.EBREAK,rv32.EBREAK])])
def test_unsupported_program_is_explicit(binary):
with pytest.raises(ValueError):
encode_batch(['change add to sub'],[binary])
|