tensor-compiler/runtime.py at mainline · a-rahimi/tensor-compiler · GitHub

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
from typing import Dict, List, Sequence
import dataclasses as dc

import opcodes


@dc.dataclass
class Interpreter:
    """An interpreter for a opcodes defined in opcodes.py."""

    registers: Dict[str, int] = dc.field(
        default_factory=lambda: {f"r{r}": 0 for r in range(8)}
    )
    memory: List[int] = dc.field(default_factory=lambda: [0] * 1024)
    program_counter: int = 0
    top_of_heap: int = 0
    status: int = -1

    def step(self, op: opcodes.OpCode):
        print("executing", op)
        op.validate(self.registers)

        # Execute opcodes that modify the program counter.
        if isinstance(op, opcodes.Halt):
            self.status = self.registers[op.status]
            self.program_counter = -1
            return
        elif isinstance(op, opcodes.Jump):
            self.program_counter = op.target_program_counter
            return
        elif isinstance(op, opcodes.JumpIf):
            if self.registers[op.condition]:
                self.program_counter = op.target_program_counter
            else:
                self.program_counter += 1
            return

        # Execute opcodes that don't modify the program counter.
        if isinstance(op, opcodes.Load):
            self.registers[op.dst_address] = self.memory[self.registers[op.src_address]]
        elif isinstance(op, opcodes.LoadImmediate):
            self.registers[op.dst] = op.value
        elif isinstance(op, opcodes.Store):
            self.memory[self.registers[op.dst_address]] = self.registers[op.src]
        elif isinstance(op, opcodes.Move):
            self.registers[op.dst] = self.registers[op.src]
        elif isinstance(op, opcodes.Add):
            self.registers[op.dst] = self.registers[op.src1] + self.registers[op.src2]
        elif isinstance(op, opcodes.Multiply):
            self.registers[op.dst] = self.registers[op.src1] * self.registers[op.src2]
        elif isinstance(op, opcodes.Subtract):
            self.registers[op.dst] = self.registers[op.src1] - self.registers[op.src2]
        elif isinstance(op, opcodes.LessThan):
            self.registers[op.dst] = int(
                self.registers[op.src1] < self.registers[op.src2]
            )
        elif isinstance(op, opcodes.LessThanEqual):
            self.registers[op.dst] = int(
                self.registers[op.src1] <= self.registers[op.src2]
            )
        elif isinstance(op, opcodes.GreaterThan):
            self.registers[op.dst] = int(
                self.registers[op.src1] > self.registers[op.src2]
            )
        elif isinstance(op, opcodes.GreaterThanEqual):
            self.registers[op.dst] = int(
                self.registers[op.src1] >= self.registers[op.src2]
            )
        elif isinstance(op, opcodes.Negate):
            self.registers[op.dst] = int(not self.registers[op.src])
        elif isinstance(op, opcodes.AllocateMemory):
            self.registers[op.address] = self.top_of_heap
            self.top_of_heap += self.registers[op.size]
        elif isinstance(op, opcodes.CopyMemory):
            dst = self.registers[op.dst_address]
            src = self.registers[op.src_address]
            size = self.registers[op.size]
            self.memory[dst : dst + size] = self.memory[src : src + size]
        elif isinstance(op, opcodes.CopyLiteralArray):
            dst = self.registers[op.dst_address]
            self.memory[dst : dst + len(op.src)] = op.src

        self.program_counter += 1

    def execute(self, program: Sequence[opcodes.OpCode]):
        while 0 <= self.program_counter < len(program):
            self.step(program[self.program_counter])