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