Finish optimization algorithm

ex.py

@@ -1,24 +1,67 @@
+from time import time
+
+from gbso.cpu.state import CPUState
+from gbso.program.test_case import TestCase
 from gbso.cpu.insn import *
-from gbso.cpu.regs import R8, R16
+from gbso.cpu.regs import R8
 
-from gbso.optimize import optimize
+from gbso.optimize import cost, optimize
 from gbso.program.program import Program
 
-rLCDC = 0xFF40
-rIE = 0xFFFF
-
-IEF_VBLANK = 1 << 0
-
 prgm = Program(
     insns=[
-        LD_R_N8(R8.A, 12),
-        LD_RR_NN(R16.HL, rLCDC),
-        RES_N_HL(7),
-        RES_N_HL(2),
-        LD_RR_NN(R16.HL, rIE),
-        LD_HL_N(IEF_VBLANK),
+        INC_R(R8.A),
+        INC_R(R8.A),
+        INC_R(R8.A),
+        INC_R(R8.A),
     ],
 )
 
-optimized_prgm = optimize(prgm)
+outputs = [R8.A]
+test_cases = [TestCase()]
+max_size = 4
+num_iters = 1_000_000  # 10_000_000
+
+initial_cost = cost(prgm, test_cases, outputs, prgm)
+initial_cycles = prgm.perf()
+
+prgm.display()
+
+print(f"Cost: {initial_cost}")
+print(f"Cycles: {initial_cycles}")
+
+start_time = time()
+
+optimized_prgm = optimize(
+    prgm,
+    max_size=max_size,
+    test_cases=test_cases,
+    outputs=outputs,
+    num_iters=num_iters,
+)
+
+end_time = time()
+
 optimized_prgm.display()
+
+optimized_cost = cost(prgm, test_cases, outputs, optimized_prgm)
+optimized_cycles = optimized_prgm.perf()
+
+print(f"Cost: {optimized_cost}")
+print(f"Cycles: {optimized_cycles}")
+
+
+optimized_cost = cost(prgm, test_cases, outputs, optimized_prgm)
+optimized_cycles = optimized_prgm.perf()
+
+print(f"Cost: {optimized_cost}")
+print(f"Cycles: {optimized_cycles}")
+
+
+optimized_cost = cost(prgm, test_cases, outputs, optimized_prgm)
+optimized_cycles = optimized_prgm.perf()
+
+print(f"Cost: {optimized_cost}")
+print(f"Cycles: {optimized_cycles}")
+
+print(f"Took {round(end_time - start_time, 3)} seconds")

gbso/cpu/state.py

@@ -11,3 +11,11 @@ class CPUState:
     sp: int = 0
     reg8: Dict[R8, int] = field(default_factory=lambda: defaultdict(lambda: 0))
     memory: Dict[int, int] = field(default_factory=lambda: defaultdict(lambda: 0))
+
+    def copy(self) -> "CPUState":
+        return CPUState(
+            carry=self.carry,
+            sp=self.sp,
+            reg8=self.reg8.copy(),
+            memory=self.memory.copy(),
+        )

gbso/optimize.py

@@ -1,3 +1,9 @@
+from math import exp, log
+from random import random
+from typing import List
+
+from gbso.program.test_case import Output, TestCase, eq_on_testcase
+from gbso.program.mutate import mutate_program
 from gbso.program.program import Program
 
 EPSILON = 0.00001
@@ -12,8 +18,23 @@ DEFAULT_PROB_INSN = 0.25
 DEFAULT_PROB_INSN_UNUSED = 0.1
 
 
+def cost(orig_prgm, test_cases, outputs, prgm) -> int:
+    c = prgm.perf() - orig_prgm.perf()
+    # print(f"init cost: {c}")
+
+    for test_case in test_cases:
+        c += eq_on_testcase(orig_prgm, prgm, test_case, outputs)
+        # print(f"cost after testcase: {c}")
+
+    return c
+
+
 def optimize(
     prgm: Program,
+    max_size: int,
+    test_cases: List[TestCase],
+    outputs: List[Output],
+    beta: int = 0.75,  # How far away in cost you are allowed to search
     num_iters: int = DEFAULT_NUM_ITERS,
     prob_opcode: float = DEFAULT_PROB_OPCODE,
     prob_operand: float = DEFAULT_PROB_OPERAND,
@@ -21,7 +42,19 @@ def optimize(
     prob_insn: float = DEFAULT_PROB_INSN,
     prob_insn_unused: float = DEFAULT_PROB_INSN_UNUSED,
 ) -> Program:
-    prob_sum = sum([prob_opcode, prob_operand, prob_swap, prob_insn])
-    assert abs(1 - prob_sum) < EPSILON
+    padded_prgm = prgm.pad(max_size)
+
+    last_prgm = padded_prgm
+    last_cost = cost(padded_prgm, test_cases, outputs, last_prgm)
+
+    for _ in range(num_iters):
+        candidate_prgm = mutate_program(
+            last_prgm, prob_opcode, prob_operand, prob_swap, prob_insn, prob_insn_unused
+        )
+        candidate_cost = cost(padded_prgm, test_cases, outputs, candidate_prgm)
+
+        if candidate_cost < last_cost - log(random()) / beta:
+            last_prgm = candidate_prgm
+            last_cost = candidate_cost
 
-    return prgm
+    return last_prgm

gbso/program/mutate.py

@@ -13,6 +13,7 @@ class Mutation(Enum):
     OPERAND = "OPERAND"
     SWAP = "SWAP"
     INSTRUCTION = "INSTRUCTION"
+    NONE = "NONE"
 
 
 def mutate_program(
@@ -23,12 +24,14 @@ def mutate_program(
     prob_insn: float,
     prob_insn_unused: float,
 ) -> Program:
+    prob_none = max(0, 1 - prob_opcode - prob_operand - prob_swap - prob_insn)
     mutation = sample_probs(
         [
             (Mutation.OPCODE, prob_opcode),
             (Mutation.OPERAND, prob_operand),
             (Mutation.SWAP, prob_swap),
             (Mutation.INSTRUCTION, prob_insn),
+            (Mutation.NONE, prob_none),
         ]
     )
     new_prgm = prgm
@@ -61,6 +64,9 @@ def mutate_operand(prgm: Program) -> Program:
     insn_index = randrange(len(prgm.insns))
     insn = prgm.insns[insn_index]
 
+    if len(insn.signature()) == 0:
+        return prgm
+
     operand_index = randrange(len(insn.signature()))
     operand_type = insn.signature()[operand_index]
     operand_name = list(vars(insn).keys())[operand_index]
@@ -135,7 +141,7 @@ def create_random_operand(op: Operand) -> Union[int, R8, R16]:
     elif op == Operand.IMM8:
         value = randrange(256)
     elif op == Operand.IMM16:
-        value = randrange(65536)
+        value = randrange(256)
     elif op == Operand.SIMM8:
         value = randint(-128, 127)
 

gbso/program/program.py

@@ -1,15 +1,15 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 
 from typing import List, Optional
 
 from gbso.cpu.cpu import CPU
 from gbso.cpu.state import CPUState
-from gbso.cpu.insn import Insn
+from gbso.cpu.insn import Insn, UNUSED
 
 
 @dataclass
 class Program:
-    insns: List[Insn]
+    insns: List[Insn] = field(default_factory=list)
 
     def execute(self, init_state: Optional[CPUState] = None) -> CPU:
         cpu = CPU(state=init_state)
@@ -24,3 +24,11 @@ class Program:
         print("=" * 10)
         for insn in self.insns:
             print(insn.pretty())
+        print("=" * 10)
+
+    def pad(self, max_bytes: int) -> "Program":
+        if len(self.insns) > max_bytes:
+            raise ValueError("larger than amount to pad to")
+
+        insns = self.insns.copy() + [UNUSED()] * (max_bytes - len(self.insns))
+        return Program(insns=insns)

gbso/program/test_case.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field, replace
 from typing import List, Union
 
 from gbso.cpu.cpu import CPU
@@ -14,7 +14,7 @@ Output = Union[R8, int]
 @dataclass
 class TestCase:
     __test__ = False
-    state: CPUState
+    state: CPUState = field(default_factory=CPUState)
 
 
 # TODO: Allow p_cpu to be computed AOT
@@ -22,8 +22,8 @@ class TestCase:
 def eq_on_testcase(
     p: Program, q: Program, case: TestCase, outputs: List[Output]
 ) -> int:
-    p_cpu = p.execute(case.state)
-    q_cpu = q.execute(case.state)
+    p_cpu = p.execute(case.state.copy())
+    q_cpu = q.execute(case.state.copy())
     return sum([eq_on_output(o, p_cpu, q_cpu) for o in outputs])
 
 
@@ -44,4 +44,4 @@ def eq_8bit(x: int, y: int) -> int:
         mask = 1 << i
         if x & mask != y & mask:
             delta += 1
-    return 0
+    return delta