Kernel-optimization take-home scored by speedup

"""A toy VLIW machine with vector lanes. Do not edit; your score is the
number of cycles `run` reports.

A program is a list of bundles; a bundle is a list of instructions that run in
the same cycle. Every instruction in a bundle reads its inputs before any of
them writes, so an instruction never sees a value written in its own bundle.
"""

VLEN = 8
SLOTS = {"alu": 8, "valu": 4, "load": 2, "store": 1}  # per bundle
ENGINE = {
    "const": "alu", "alu": "alu",
    "vbroadcast": "valu", "valu": "valu",
    "load": "load", "vload": "load",
    "store": "store", "vstore": "store",
}  # fmt: skip
N_SCALAR, N_VECTOR = 1024, 128
MASK = 0xFFFFFFFF
OPS = {
    "+": lambda a, b: a + b, "-": lambda a, b: a - b, "*": lambda a, b: a * b,
    "^": lambda a, b: a ^ b, "&": lambda a, b: a & b, "|": lambda a, b: a | b,
    "<<": lambda a, b: a << b, ">>": lambda a, b: a >> b,
}  # fmt: skip

# The kernel: three rounds of v = v * mul + add; v ^= v >> shift (32-bit).
ROUNDS = [(31, 7, 3), (17, 11, 5), (13, 3, 7)]
N, IN, OUT = 256, 0, 256  # 256 inputs at mem[0:256], results go to mem[256:512]


def reference(values):
    out = []
    for v in values:
        for mul, add, shift in ROUNDS:
            v = (v * mul + add) & MASK
            v ^= v >> shift
        out.append(v)
    return out


class Machine:
    def __init__(self, mem):
        self.mem = list(mem)
        self.s = [0] * N_SCALAR
        self.v = [[0] * VLEN for _ in range(N_VECTOR)]
        self.cycles = 0

    def run(self, program):
        for bundle in program:
            used = {}
            for ins in bundle:
                engine = ENGINE[ins[0]]
                used[engine] = used.get(engine, 0) + 1
                if used[engine] > SLOTS[engine]:
                    raise ValueError(f"cycle {self.cycles}: more than {SLOTS[engine]} {engine} slots")
            writes = [w for ins in bundle for w in self.execute(ins)]
            targets = [(kind, idx) for kind, idx, _ in writes]
            if len(targets) != len(set(targets)):
                raise ValueError(f"cycle {self.cycles}: two writes to the same place")
            for kind, idx, value in writes:
                if kind == "s":
                    self.s[idx] = value
                elif kind == "v":
                    self.v[idx] = value
                else:
                    self.mem[idx] = value
            self.cycles += 1
        return self

    def execute(self, ins):
        """Return the writes an instruction makes, reading current state only."""
        op = ins[0]
        s, v, mem = self.s, self.v, self.mem
        if op == "const":
            return [("s", ins[1], ins[2] & MASK)]
        if op == "alu":
            _, f, dst, a, b = ins
            return [("s", dst, OPS[f](s[a], s[b]) & MASK)]
        if op == "valu":
            _, f, dst, a, b = ins
            return [("v", dst, [OPS[f](x, y) & MASK for x, y in zip(v[a], v[b])])]
        if op == "vbroadcast":
            return [("v", ins[1], [s[ins[2]]] * VLEN)]
        if op == "load":
            return [("s", ins[1], mem[s[ins[2]]])]
        if op == "vload":
            base = s[ins[2]]
            return [("v", ins[1], mem[base:base + VLEN])]
        if op == "store":
            return [("m", s[ins[1]], s[ins[2]])]
        if op == "vstore":
            base = s[ins[1]]
            return [("m", base + i, x) for i, x in enumerate(v[ins[2]])]
        raise ValueError(f"unknown instruction {ins!r}")


def run_kernel(build_kernel, values):
    """Build the program, run it on `values`, return (outputs, cycles)."""
    machine = Machine(list(values) + [0] * N).run(build_kernel())
    return machine.mem[OUT:OUT + N], machine.cycles

"""Edit this file. build_kernel() returns the program for machine.Machine."""

from machine import IN, N, OUT, ROUNDS

ADDR, VAL, TMP, K = 0, 1, 2, 3  # scalar registers


def build_kernel():
    program = []

    def emit(ins):
        program.append([ins])  # one instruction per bundle

    for i in range(N):
        emit(("const", ADDR, IN + i))
        emit(("load", VAL, ADDR))
        for mul, add, shift in ROUNDS:
            emit(("const", K, mul))
            emit(("alu", "*", VAL, VAL, K))
            emit(("const", K, add))
            emit(("alu", "+", VAL, VAL, K))
            emit(("const", K, shift))
            emit(("alu", ">>", TMP, VAL, K))
            emit(("alu", "^", VAL, VAL, TMP))
        emit(("const", ADDR, OUT + i))
        emit(("store", ADDR, VAL))
    return program