Parameterize CPU circuits for N-bit data and address widths

Build changes:
- add_fetch_load_store_buffers: N-bit data bus support
- add_stack_ops: N-bit data, M-bit stack pointer
- add_conditional_jumps: 8 jump types with N-bit addresses
- add_status_flags: Z/N/C/V flags for N-bit ALU results
- add_rol_ror_nbits: Parameterized rotation circuits
- update_manifest: Version 4.0 with data_bits field
- Stack ops moved from cmd_alu to cmd_memory

Eval changes:
- get_manifest: Extract data_bits/addr_bits from model
- BatchedFitnessEvaluator reads manifest for N-bit support
- _test_conditional_jump uses self.addr_bits
- _test_stack_ops uses self.addr_bits for SP width

Tested configurations:
- 8-bit data / 10-bit addr (small memory)
- 32-bit data / 10-bit addr (small memory)

build.py

@@ -255,10 +255,19 @@ def add_memory_write_cells(tensors: Dict[str, torch.Tensor], mem_bytes: int) ->
     tensors["memory.write.or.bias"] = or_bias
 
 
-def add_fetch_load_store_buffers(tensors: Dict[str, torch.Tensor], addr_bits: int) -> None:
-    for bit in range(16):
+def add_fetch_load_store_buffers(tensors: Dict[str, torch.Tensor], data_bits: int, addr_bits: int) -> None:
+    """Add control buffers for fetch, load, store operations.
+
+    Args:
+        data_bits: Width of data bus (8/16/32)
+        addr_bits: Width of address bus (determines instruction register width)
+    """
+    # Instruction register width = opcode (8) + operands (depends on arch)
+    # For simplicity, IR width = max(16, addr_bits) to hold jump targets
+    ir_bits = max(16, addr_bits)
+    for bit in range(ir_bits):
         add_gate(tensors, f"control.fetch.ir.bit{bit}", [1.0], [-1.0])
-    for bit in range(8):
+    for bit in range(data_bits):
         add_gate(tensors, f"control.load.bit{bit}", [1.0], [-1.0])
         add_gate(tensors, f"control.store.bit{bit}", [1.0], [-1.0])
     for bit in range(addr_bits):
@@ -555,116 +564,197 @@ def add_neg(tensors: Dict[str, torch.Tensor]) -> None:
 
 
 def add_rol_ror(tensors: Dict[str, torch.Tensor]) -> None:
-    """Add ROL and ROR circuits (rotate left/right).
+    """Add 8-bit ROL and ROR circuits (legacy wrapper)."""
+    add_rol_ror_nbits(tensors, 8)
+
 
-    ROL: out[i] = in[i+1] for i<7, out[7] = in[0] (MSB wraps to LSB)
-    ROR: out[0] = in[7], out[i] = in[i-1] for i>0 (LSB wraps to MSB)
+def add_rol_ror_nbits(tensors: Dict[str, torch.Tensor], bits: int) -> None:
+    """Add N-bit ROL and ROR circuits (rotate left/right).
 
-    Identity gates with circular wiring.
+    ROL: out[i] = in[i+1] for i<N-1, out[N-1] = in[0] (MSB wraps to LSB)
+    ROR: out[0] = in[N-1], out[i] = in[i-1] for i>0 (LSB wraps to MSB)
+
+    Args:
+        bits: Data width (8, 16, 32, etc.)
     """
     # ROL: rotate left (toward MSB)
-    for bit in range(8):
-        src = (bit + 1) % 8  # Circular: bit 7 gets bit 0
-        add_gate(tensors, f"alu.alu8bit.rol.bit{bit}", [2.0], [-1.0])
+    for bit in range(bits):
+        add_gate(tensors, f"alu.alu{bits}bit.rol.bit{bit}", [2.0], [-1.0])
 
     # ROR: rotate right (toward LSB)
-    for bit in range(8):
-        src = (bit - 1) % 8  # Circular: bit 0 gets bit 7
-        add_gate(tensors, f"alu.alu8bit.ror.bit{bit}", [2.0], [-1.0])
+    for bit in range(bits):
+        add_gate(tensors, f"alu.alu{bits}bit.ror.bit{bit}", [2.0], [-1.0])
 
 
-def add_stack_ops(tensors: Dict[str, torch.Tensor]) -> None:
+def add_stack_ops(tensors: Dict[str, torch.Tensor], data_bits: int, addr_bits: int) -> None:
     """Add RET, PUSH, POP circuit components.
 
     These are higher-level operations that use memory read/write.
     We create the control logic gates.
 
+    Args:
+        data_bits: Width of data to push/pop (8/16/32)
+        addr_bits: Width of stack pointer and return addresses
+
     RET: Pop return address from stack, jump to it
     PUSH: Decrement SP, write value to [SP]
     POP: Read value from [SP], increment SP
     """
-    # SP decrement for PUSH (16-bit)
-    for bit in range(16):
+    # SP decrement for PUSH (addr_bits wide)
+    for bit in range(addr_bits):
         add_gate(tensors, f"control.push.sp_dec.bit{bit}.xor.layer1.or", [1.0, 1.0], [-1.0])
         add_gate(tensors, f"control.push.sp_dec.bit{bit}.xor.layer1.nand", [-1.0, -1.0], [1.0])
         add_gate(tensors, f"control.push.sp_dec.bit{bit}.xor.layer2", [1.0, 1.0], [-2.0])
         add_gate(tensors, f"control.push.sp_dec.bit{bit}.borrow", [1.0, 1.0], [-2.0])
 
-    # SP increment for POP (16-bit)
-    for bit in range(16):
+    # SP increment for POP (addr_bits wide)
+    for bit in range(addr_bits):
         add_gate(tensors, f"control.pop.sp_inc.bit{bit}.xor.layer1.or", [1.0, 1.0], [-1.0])
         add_gate(tensors, f"control.pop.sp_inc.bit{bit}.xor.layer1.nand", [-1.0, -1.0], [1.0])
         add_gate(tensors, f"control.pop.sp_inc.bit{bit}.xor.layer2", [1.0, 1.0], [-2.0])
         add_gate(tensors, f"control.pop.sp_inc.bit{bit}.carry", [1.0, 1.0], [-2.0])
 
-    # RET uses POP twice (for 16-bit address) then jumps
-    # Buffer gates for return address
-    for bit in range(16):
+    # Data buffers for PUSH (data_bits wide)
+    for bit in range(data_bits):
+        add_gate(tensors, f"control.push.data.bit{bit}", [2.0], [-1.0])
+
+    # Data buffers for POP (data_bits wide)
+    for bit in range(data_bits):
+        add_gate(tensors, f"control.pop.data.bit{bit}", [2.0], [-1.0])
+
+    # RET: Buffer gates for return address (addr_bits wide)
+    for bit in range(addr_bits):
         add_gate(tensors, f"control.ret.addr.bit{bit}", [2.0], [-1.0])
 
 
+def add_conditional_jumps(tensors: Dict[str, torch.Tensor], addr_bits: int) -> None:
+    """Add conditional jump circuits (JZ, JNZ, JC, JNC, JP, JN, JV, JNV).
+
+    Each conditional jump is a 2:1 MUX per address bit:
+    - If flag is set: output = target_bit
+    - If flag is clear: output = pc_bit
+
+    Structure per bit:
+    - not_sel: NOT(flag)
+    - and_a: pc_bit AND NOT(flag)
+    - and_b: target_bit AND flag
+    - or: and_a OR and_b
+
+    Args:
+        addr_bits: Width of program counter / jump target
+    """
+    jump_types = ['jz', 'jnz', 'jc', 'jnc', 'jp', 'jn', 'jv', 'jnv']
+
+    for jmp in jump_types:
+        for bit in range(addr_bits):
+            prefix = f"control.{jmp}.bit{bit}"
+            # NOT sel (invert flag)
+            add_gate(tensors, f"{prefix}.not_sel", [-1.0], [0.0])
+            # AND a: pc_bit AND NOT(flag)
+            add_gate(tensors, f"{prefix}.and_a", [1.0, 1.0], [-2.0])
+            # AND b: target_bit AND flag
+            add_gate(tensors, f"{prefix}.and_b", [1.0, 1.0], [-2.0])
+            # OR: combine
+            add_gate(tensors, f"{prefix}.or", [1.0, 1.0], [-1.0])
+
+
+def add_status_flags(tensors: Dict[str, torch.Tensor], data_bits: int) -> None:
+    """Add status flag computation circuits (Z, N, C, V).
+
+    Args:
+        data_bits: Width of ALU data (8/16/32)
+
+    Flags:
+    - Z (Zero): NOR of all result bits (1 if result == 0)
+    - N (Negative): Copy of MSB (sign bit)
+    - C (Carry): Carry out from adder (external input)
+    - V (Overflow): XOR of carry into and out of MSB (signed overflow)
+    """
+    # Z flag: NOR of all bits (result == 0)
+    # Single threshold gate: fires if sum of all bits < 1
+    add_gate(tensors, "flags.zero", [-1.0] * data_bits, [0.0])
+
+    # N flag: Buffer for MSB (sign bit)
+    add_gate(tensors, "flags.negative", [2.0], [-1.0])
+
+    # C flag: Buffer for carry out (input from adder)
+    add_gate(tensors, "flags.carry", [2.0], [-1.0])
+
+    # V flag: XOR of carry_in_msb and carry_out_msb
+    # Two-layer XOR: (A OR B) AND (A NAND B)
+    add_gate(tensors, "flags.overflow.layer1.or", [1.0, 1.0], [-1.0])
+    add_gate(tensors, "flags.overflow.layer1.nand", [-1.0, -1.0], [1.0])
+    add_gate(tensors, "flags.overflow.layer2", [1.0, 1.0], [-2.0])
+
+
 def add_barrel_shifter(tensors: Dict[str, torch.Tensor]) -> None:
-    """Add barrel shifter circuit.
+    """Add 8-bit barrel shifter circuit (legacy wrapper)."""
+    add_barrel_shifter_nbits(tensors, 8)
 
-    Shifts input by 0-7 positions based on 3-bit shift amount.
+
+def add_barrel_shifter_nbits(tensors: Dict[str, torch.Tensor], bits: int) -> None:
+    """Add N-bit barrel shifter circuit.
+
+    Shifts input by 0 to (bits-1) positions based on ceil(log2(bits))-bit shift amount.
     Uses layers of 2:1 muxes controlled by shift amount bits.
 
-    Layer 0: shift by 0 or 1 (controlled by shift[2], LSB)
-    Layer 1: shift by 0 or 2 (controlled by shift[1])
-    Layer 2: shift by 0 or 4 (controlled by shift[0], MSB)
+    Args:
+        bits: Data width (8, 16, 32, etc.)
     """
-    # 3 layers of muxes, 8 bits each
-    for layer in range(3):
-        shift_amount = 1 << (2 - layer)  # 4, 2, 1 for layers 0, 1, 2
-        for bit in range(8):
+    import math
+    num_layers = max(1, math.ceil(math.log2(bits)))
+
+    for layer in range(num_layers):
+        shift_amount = 1 << (num_layers - 1 - layer)
+        for bit in range(bits):
+            prefix = f"combinational.barrelshifter{bits}.layer{layer}.bit{bit}"
             # 2:1 mux: if sel then shifted else original
-            # NOT for inverting select
-            add_gate(tensors, f"combinational.barrelshifter.layer{layer}.bit{bit}.not_sel", [-1.0], [0.0])
-            # AND gates
-            add_gate(tensors, f"combinational.barrelshifter.layer{layer}.bit{bit}.and_a", [1.0, 1.0], [-2.0])
-            add_gate(tensors, f"combinational.barrelshifter.layer{layer}.bit{bit}.and_b", [1.0, 1.0], [-2.0])
-            # OR gate
-            add_gate(tensors, f"combinational.barrelshifter.layer{layer}.bit{bit}.or", [1.0, 1.0], [-1.0])
+            add_gate(tensors, f"{prefix}.not_sel", [-1.0], [0.0])
+            add_gate(tensors, f"{prefix}.and_a", [1.0, 1.0], [-2.0])
+            add_gate(tensors, f"{prefix}.and_b", [1.0, 1.0], [-2.0])
+            add_gate(tensors, f"{prefix}.or", [1.0, 1.0], [-1.0])
 
 
 def add_priority_encoder(tensors: Dict[str, torch.Tensor]) -> None:
-    """Add priority encoder circuit.
+    """Add 8-bit priority encoder circuit (legacy wrapper)."""
+    add_priority_encoder_nbits(tensors, 8)
+
 
-    Finds the position of the highest set bit (0-7).
-    Output is 3-bit index + valid flag.
+def add_priority_encoder_nbits(tensors: Dict[str, torch.Tensor], bits: int) -> None:
+    """Add N-bit priority encoder circuit.
 
-    Uses cascaded comparisons: check bit 7 first, then 6, etc.
+    Finds the position of the highest set bit (0 to bits-1).
+    Output is ceil(log2(bits))-bit index + valid flag.
+
+    Args:
+        bits: Input width (8, 16, 32, etc.)
     """
-    # Check each bit position (8 OR gates to detect any bit set at or above position)
-    for pos in range(8):
-        # OR of bits pos through 7
-        num_inputs = 8 - pos
+    import math
+    out_bits = max(1, math.ceil(math.log2(bits)))
+    prefix = f"combinational.priorityencoder{bits}"
+
+    # Check each bit position (OR gates to detect any bit set at or above position)
+    for pos in range(bits):
+        num_inputs = bits - pos
         weights = [1.0] * num_inputs
-        add_gate(tensors, f"combinational.priorityencoder.any_ge{pos}",
-                weights, [-1.0])
+        add_gate(tensors, f"{prefix}.any_ge{pos}", weights, [-1.0])
 
     # Priority logic: pos N is highest if bit N is set AND no higher bit is set
-    for pos in range(8):
-        # bit[pos] AND NOT(any bit > pos)
-        add_gate(tensors, f"combinational.priorityencoder.is_highest{pos}.not_higher", [-1.0], [0.0])
-        add_gate(tensors, f"combinational.priorityencoder.is_highest{pos}.and", [1.0, 1.0], [-2.0])
-
-    # Encode position to 3-bit output
-    # out[0] (LSB): positions 1,3,5,7
-    # out[1]: positions 2,3,6,7
-    # out[2] (MSB): positions 4,5,6,7
-    for out_bit in range(3):
+    for pos in range(bits):
+        add_gate(tensors, f"{prefix}.is_highest{pos}.not_higher", [-1.0], [0.0])
+        add_gate(tensors, f"{prefix}.is_highest{pos}.and", [1.0, 1.0], [-2.0])
+
+    # Encode position to output bits
+    for out_bit in range(out_bits):
         weights = []
-        for pos in range(8):
+        for pos in range(bits):
             if (pos >> out_bit) & 1:
                 weights.append(1.0)
         if weights:
-            add_gate(tensors, f"combinational.priorityencoder.out{out_bit}",
-                    weights, [-1.0])
+            add_gate(tensors, f"{prefix}.out{out_bit}", weights, [-1.0])
 
     # Valid flag: any bit set
-    add_gate(tensors, f"combinational.priorityencoder.valid",
-            [1.0] * 8, [-1.0])
+    add_gate(tensors, f"{prefix}.valid", [1.0] * bits, [-1.0])
 
 
 def add_comparators(tensors: Dict[str, torch.Tensor]) -> None:
@@ -908,10 +998,19 @@ def add_neg_nbits(tensors: Dict[str, torch.Tensor], bits: int) -> None:
         add_gate(tensors, f"alu.alu{bits}bit.neg.inc.bit{bit}.carry", [1.0, 1.0], [-2.0])
 
 
-def update_manifest(tensors: Dict[str, torch.Tensor], addr_bits: int, mem_bytes: int) -> None:
+def update_manifest(tensors: Dict[str, torch.Tensor], data_bits: int, addr_bits: int, mem_bytes: int) -> None:
+    """Update manifest metadata tensors.
+
+    Args:
+        data_bits: ALU/register width (8/16/32)
+        addr_bits: Address bus width (determines memory size)
+        mem_bytes: Memory size in bytes (2^addr_bits)
+    """
+    tensors["manifest.data_bits"] = torch.tensor([float(data_bits)], dtype=torch.float32)
+    tensors["manifest.addr_bits"] = torch.tensor([float(addr_bits)], dtype=torch.float32)
     tensors["manifest.memory_bytes"] = torch.tensor([float(mem_bytes)], dtype=torch.float32)
     tensors["manifest.pc_width"] = torch.tensor([float(addr_bits)], dtype=torch.float32)
-    tensors["manifest.version"] = torch.tensor([3.0], dtype=torch.float32)
+    tensors["manifest.version"] = torch.tensor([4.0], dtype=torch.float32)  # Bump version for N-bit support
 
 
 def write_manifest(path: Path, tensors: Dict[str, torch.Tensor]) -> None:
@@ -2028,6 +2127,10 @@ def cmd_memory(args) -> None:
     drop_prefixes(tensors, [
         "memory.addr_decode.", "memory.read.", "memory.write.",
         "control.fetch.ir.", "control.load.", "control.store.", "control.mem_addr.",
+        "control.push.", "control.pop.", "control.ret.",
+        "control.jz.", "control.jnz.", "control.jc.", "control.jnc.",
+        "control.jp.", "control.jn.", "control.jv.", "control.jnv.",
+        "flags.",
     ])
     print(f"  Now {len(tensors)} tensors")
 
@@ -2040,16 +2143,42 @@ def cmd_memory(args) -> None:
 
         print("\nGenerating buffer gates...")
         try:
-            add_fetch_load_store_buffers(tensors, addr_bits)
-            print("  Added fetch/load/store/mem_addr buffers")
+            add_fetch_load_store_buffers(tensors, args.bits, addr_bits)
+            print(f"  Added fetch/load/store/mem_addr buffers ({args.bits}-bit data, {addr_bits}-bit addr)")
         except ValueError as e:
             print(f"  Buffers already exist: {e}")
+
+        print("\nGenerating stack operation circuits...")
+        try:
+            add_stack_ops(tensors, args.bits, addr_bits)
+            sp_gates = addr_bits * 4 * 2  # SP inc/dec gates
+            data_gates = args.bits * 2  # PUSH/POP data buffers
+            ret_gates = addr_bits  # RET address buffers
+            total_gates = sp_gates + data_gates + ret_gates
+            print(f"  Added PUSH/POP/RET ({total_gates} gates: {args.bits}-bit data, {addr_bits}-bit SP)")
+        except ValueError as e:
+            print(f"  Stack ops already exist: {e}")
+
+        print("\nGenerating conditional jump circuits...")
+        try:
+            add_conditional_jumps(tensors, addr_bits)
+            jump_gates = 8 * addr_bits * 4  # 8 jump types × addr_bits × 4 gates each
+            print(f"  Added JZ/JNZ/JC/JNC/JP/JN/JV/JNV ({jump_gates} gates: {addr_bits}-bit addresses)")
+        except ValueError as e:
+            print(f"  Conditional jumps already exist: {e}")
+
+        print("\nGenerating status flag circuits...")
+        try:
+            add_status_flags(tensors, args.bits)
+            print(f"  Added Z/N/C/V flags ({args.bits}-bit aware)")
+        except ValueError as e:
+            print(f"  Status flags already exist: {e}")
     else:
         print("\nSkipping memory circuits (addr_bits=0, pure ALU mode)")
 
     print("\nUpdating manifest...")
-    update_manifest(tensors, addr_bits, mem_bytes)
-    print(f"  memory_bytes={mem_bytes:,}, pc_width={addr_bits}")
+    update_manifest(tensors, args.bits, addr_bits, mem_bytes)
+    print(f"  data_bits={args.bits}, addr_bits={addr_bits}, memory_bytes={mem_bytes:,}")
 
     if args.apply:
         print(f"\nSaving: {args.model}")
@@ -2113,7 +2242,6 @@ def cmd_alu(args) -> None:
         "arithmetic.greaterthan8bit.", "arithmetic.lessthan8bit.",
         "arithmetic.greaterorequal8bit.", "arithmetic.lessorequal8bit.",
         "arithmetic.equality8bit.", "arithmetic.add3_8bit.", "arithmetic.expr_add_mul.", "arithmetic.expr_paren.",
-        "control.push.", "control.pop.", "control.ret.",
         "combinational.barrelshifter.", "combinational.priorityencoder.",
     ]
 
@@ -2164,12 +2292,6 @@ def cmd_alu(args) -> None:
         print("  Added ROL (8 gates), ROR (8 gates)")
     except ValueError as e:
         print(f"  ROL/ROR already exist: {e}")
-    print("\nGenerating stack operation circuits...")
-    try:
-        add_stack_ops(tensors)
-        print("  Added PUSH/POP/RET (144 gates)")
-    except ValueError as e:
-        print(f"  Stack ops already exist: {e}")
     print("\nGenerating barrel shifter...")
     try:
         add_barrel_shifter(tensors)

eval.py

@@ -67,6 +67,21 @@ def load_metadata(path: str = MODEL_PATH) -> Dict:
         return {'signal_registry': {}}
 
 
+def get_manifest(tensors: Dict[str, torch.Tensor]) -> Dict[str, int]:
+    """Extract manifest values from tensors.
+
+    Returns dict with data_bits, addr_bits, memory_bytes, version.
+    Defaults to 8-bit data, 16-bit addr for legacy models.
+    """
+    return {
+        'data_bits': int(tensors.get('manifest.data_bits', torch.tensor([8.0])).item()),
+        'addr_bits': int(tensors.get('manifest.addr_bits',
+                         tensors.get('manifest.pc_width', torch.tensor([16.0]))).item()),
+        'memory_bytes': int(tensors.get('manifest.memory_bytes', torch.tensor([65536.0])).item()),
+        'version': float(tensors.get('manifest.version', torch.tensor([1.0])).item()),
+    }
+
+
 def create_population(
     base_tensors: Dict[str, torch.Tensor],
     pop_size: int,
@@ -889,7 +904,7 @@ class BatchedFitnessEvaluator:
     Tests all circuits comprehensively.
     """
 
-    def __init__(self, device: str = 'cuda', model_path: str = MODEL_PATH):
+    def __init__(self, device: str = 'cuda', model_path: str = MODEL_PATH, tensors: Dict[str, torch.Tensor] = None):
         self.device = device
         self.model_path = model_path
         self.metadata = load_metadata(model_path)
@@ -897,6 +912,16 @@ class BatchedFitnessEvaluator:
         self.results: List[CircuitResult] = []
         self.category_scores: Dict[str, Tuple[float, int]] = {}
         self.total_tests = 0
+
+        # Get manifest for N-bit support
+        if tensors is not None:
+            self.manifest = get_manifest(tensors)
+        else:
+            base_tensors = load_model(model_path)
+            self.manifest = get_manifest(base_tensors)
+        self.data_bits = self.manifest['data_bits']
+        self.addr_bits = self.manifest['addr_bits']
+
         self._setup_tests()
 
     def _setup_tests(self):
@@ -2897,7 +2922,7 @@ class BatchedFitnessEvaluator:
     # =========================================================================
 
     def _test_conditional_jump(self, pop: Dict, name: str, debug: bool) -> Tuple[torch.Tensor, int]:
-        """Test conditional jump circuit."""
+        """Test conditional jump circuit (N-bit address aware)."""
         pop_size = next(iter(pop.values())).shape[0]
         prefix = f'control.{name}'
 
@@ -2911,7 +2936,7 @@ class BatchedFitnessEvaluator:
         scores = torch.zeros(pop_size, device=self.device)
         total = 0
 
-        for bit in range(8):
+        for bit in range(self.addr_bits):
             bit_prefix = f'{prefix}.bit{bit}'
             try:
                 # NOT sel
@@ -2979,27 +3004,34 @@ class BatchedFitnessEvaluator:
         return scores, total
 
     def _test_stack_ops(self, pop: Dict, debug: bool) -> Tuple[torch.Tensor, int]:
-        """Test PUSH/POP/RET stack operation circuits."""
+        """Test PUSH/POP/RET stack operation circuits (N-bit address aware)."""
         pop_size = next(iter(pop.values())).shape[0]
         scores = torch.zeros(pop_size, device=self.device)
         total = 0
+        addr_bits = self.addr_bits
+        addr_mask = (1 << addr_bits) - 1
 
         if debug:
-            print("\n=== STACK OPERATIONS ===")
+            print(f"\n=== STACK OPERATIONS ({addr_bits}-bit SP) ===")
 
-        # Test PUSH SP decrement (16-bit, borrow chain)
+        # Test PUSH SP decrement (addr_bits wide, borrow chain)
         try:
-            sp_tests = [0x0000, 0x0001, 0x0100, 0x8000, 0xFFFF, 0x1234]
+            # Generate test values appropriate for addr_bits
+            sp_tests = [0, 1, addr_mask // 2, addr_mask]
+            if addr_bits >= 8:
+                sp_tests.append(0x100 & addr_mask)
+            if addr_bits >= 12:
+                sp_tests.append(0x1234 & addr_mask)
             op_scores = torch.zeros(pop_size, device=self.device)
             op_total = 0
 
             for sp_val in sp_tests:
-                expected_val = (sp_val - 1) & 0xFFFF
-                sp_bits = [float((sp_val >> (15 - i)) & 1) for i in range(16)]
+                expected_val = (sp_val - 1) & addr_mask
+                sp_bits = [float((sp_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)]
 
                 borrow = 1.0
                 out_bits = []
-                for bit in range(15, -1, -1):  # LSB to MSB
+                for bit in range(addr_bits - 1, -1, -1):  # LSB to MSB
                     prefix = f'control.push.sp_dec.bit{bit}'
 
                     w_or = pop[f'{prefix}.xor.layer1.or.weight'].view(pop_size, 2)
@@ -3024,11 +3056,11 @@ class BatchedFitnessEvaluator:
                     borrow = heaviside((borrow_inp * w_borrow).sum(-1) + b_borrow)[0].item()
 
                 out = torch.stack(out_bits, dim=-1)
-                expected = torch.tensor([((expected_val >> (15 - i)) & 1) for i in range(16)],
+                expected = torch.tensor([((expected_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)],
                                        device=self.device, dtype=torch.float32)
                 correct = (out == expected.unsqueeze(0)).float().sum(1)
                 op_scores += correct
-                op_total += 16
+                op_total += addr_bits
 
             scores += op_scores
             total += op_total
@@ -3040,18 +3072,18 @@ class BatchedFitnessEvaluator:
             if debug:
                 print(f"  control.push.sp_dec: SKIP ({e})")
 
-        # Test POP SP increment (16-bit, carry chain)
+        # Test POP SP increment (addr_bits wide, carry chain)
         try:
             op_scores = torch.zeros(pop_size, device=self.device)
             op_total = 0
 
             for sp_val in sp_tests:
-                expected_val = (sp_val + 1) & 0xFFFF
-                sp_bits = [float((sp_val >> (15 - i)) & 1) for i in range(16)]
+                expected_val = (sp_val + 1) & addr_mask
+                sp_bits = [float((sp_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)]
 
                 carry = 1.0
                 out_bits = []
-                for bit in range(15, -1, -1):  # LSB to MSB
+                for bit in range(addr_bits - 1, -1, -1):  # LSB to MSB
                     prefix = f'control.pop.sp_inc.bit{bit}'
 
                     w_or = pop[f'{prefix}.xor.layer1.or.weight'].view(pop_size, 2)
@@ -3074,11 +3106,11 @@ class BatchedFitnessEvaluator:
                     carry = heaviside((inp * w_carry).sum(-1) + b_carry)[0].item()
 
                 out = torch.stack(out_bits, dim=-1)
-                expected = torch.tensor([((expected_val >> (15 - i)) & 1) for i in range(16)],
+                expected = torch.tensor([((expected_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)],
                                        device=self.device, dtype=torch.float32)
                 correct = (out == expected.unsqueeze(0)).float().sum(1)
                 op_scores += correct
-                op_total += 16
+                op_total += addr_bits
 
             scores += op_scores
             total += op_total
@@ -3090,27 +3122,29 @@ class BatchedFitnessEvaluator:
             if debug:
                 print(f"  control.pop.sp_inc: SKIP ({e})")
 
-        # Test RET address buffer (16 identity gates)
+        # Test RET address buffer (addr_bits identity gates)
         try:
             op_scores = torch.zeros(pop_size, device=self.device)
             op_total = 0
 
-            addr_tests = [0x0000, 0xFFFF, 0x1234, 0x8000, 0x00FF]
-            for addr_val in addr_tests:
-                addr_bits = torch.tensor([float((addr_val >> (15 - i)) & 1) for i in range(16)],
+            ret_tests = [0, addr_mask, addr_mask // 2, 1]
+            if addr_bits >= 12:
+                ret_tests.append(0x1234 & addr_mask)
+            for addr_val in ret_tests:
+                ret_bits_tensor = torch.tensor([float((addr_val >> (addr_bits - 1 - i)) & 1) for i in range(addr_bits)],
                                         device=self.device, dtype=torch.float32)
 
                 out_bits = []
-                for bit in range(16):
+                for bit in range(addr_bits):
                     w = pop[f'control.ret.addr.bit{bit}.weight'].view(pop_size)
                     b = pop[f'control.ret.addr.bit{bit}.bias'].view(pop_size)
-                    out = heaviside(addr_bits[bit] * w + b)
+                    out = heaviside(ret_bits_tensor[bit] * w + b)
                     out_bits.append(out)
 
                 out = torch.stack(out_bits, dim=-1)
-                correct = (out == addr_bits.unsqueeze(0)).float().sum(1)
+                correct = (out == ret_bits_tensor.unsqueeze(0)).float().sum(1)
                 op_scores += correct
-                op_total += 16
+                op_total += addr_bits
 
             scores += op_scores
             total += op_total

neural_alu32.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8a292e8d1dc5b29fd84d25d0333599a9946849e456aeb30b7519156dc150a623
-size 4985016
+oid sha256:b5a0f6cdfb4ba0ebdfc863f43e5f8fd4f41626c0fd4e7258a0a581a117a79d97
+size 5031612

neural_computer.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:eaabeed4fa50c13129fe4f83f6a8f31b6ccd41de12e83c62448460881373fc3e
-size 34838348
+oid sha256:08a39c4758f6e5236f84d231be7f2d54364099309a89cf484d607a6544194d20
+size 2591660