| """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] |
| |
| 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 |
| |
| 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) |
| |
| 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]) |
|
|
