Kernel-optimization take-home scored by speedup

Speedup take-homes reward a loop: measure, find the bound, remove it, check correctness, repeat. They punish one big rewrite that is fast and wrong.

AI tools are allowed and expected. They are good at the scheduler boilerplate and at spotting register clashes. Don't let them write code you can't explain: the follow-up conversation is about your choices.

Per element, the kernel needs 1 load, 3 rounds × 4 operations (multiply, add, shift, xor) = 12 ALU operations, and 1 store, plus the addresses. The baseline also reloads every constant before each use and issues one instruction per cycle: per element, 2 (input address, load) + 3 × 7 + 2 (output address, store) = 25 cycles, × 256 = 6,400 cycles.

Lower bounds come from dividing the work on each engine by its slots per cycle:

The last row is the important insight. Once the work is vectorized, the 8 scalar ALUs are idle except for addresses, and 8 scalar operations do as much work as one vector operation over 8 lanes. Moving 40 elements to the scalar side balances the two units.

Notes on the steps:

• Step 2 is the big idea. The 12 operations of one element form a chain, but different elements are independent. A list scheduler places each instruction in the earliest bundle where its inputs are ready and its engine has a free slot. The chains of 8 or more elements then run side by side.
• Steps 3 and 4 show that vectorizing alone is not enough. One vector instruction per cycle uses 1 of 4 vector slots and none of the others. Packing is what turns 526 into 103.
• Step 5 is about the store engine. With one store slot, 32 vector stores take 32 cycles. Emitting the work stage by stage (all multiplies, then all adds, …) makes every vector finish at about the same time, and the stores queue up at the end: 124 cycles. Emitting element by element lets early elements finish and store while later ones are still computing: 103.
• Step 6 trades scalar ALU time for vector ALU time. Measured: 0 scalar elements → 103, 32 → 91, 40 → 90, 48 → 97. Past 40 the scalar ALUs become the bottleneck, as the bound table predicts.

• Reads happen before writes in a bundle. An instruction that reads r must go in a later bundle than the one that writes r. An instruction that overwrites r may share a bundle with the last reader of r, but never come before it.
• Reusing registers creates false dependencies. If every vector uses the same temporary register, the scheduler must serialize them. Give each element its own registers; the machine has 1,024 scalar and 128 vector registers, and this kernel uses about 220 and 60.
• Two writes to the same place in one bundle are an error. The simulator raises; a real machine might silently keep one.
• Correctness first, every step. Compare against reference() on random inputs and on edge values (0, 2^32 − 1, 2^31) after each change. Overflow is the classic bug: the machine wraps at 32 bits, and so must the reference.

"""Optimized kernel: SIMD over 8 lanes, a list scheduler that packs bundles,
and the idle scalar ALUs take a share of the elements."""

import itertools

from machine import ENGINE, IN, N, OUT, ROUNDS, SLOTS, VLEN


def regs(ins):
    """(registers read, registers written) as ("s"|"v", index) pairs."""
    op = ins[0]
    if op == "const":
        return [], [("s", ins[1])]
    if op == "alu":
        return [("s", ins[3]), ("s", ins[4])], [("s", ins[2])]
    if op == "valu":
        return [("v", ins[3]), ("v", ins[4])], [("v", ins[2])]
    if op == "vbroadcast":
        return [("s", ins[2])], [("v", ins[1])]
    if op in ("load", "vload"):
        return [("s", ins[2])], [("s" if op == "load" else "v", ins[1])]
    if op == "store":
        return [("s", ins[1]), ("s", ins[2])], []
    if op == "vstore":
        return [("s", ins[1]), ("v", ins[2])], []
    raise ValueError(op)


class Scheduler:
    """Place each instruction in the earliest bundle that respects data
    dependencies and slot limits. Reads happen before writes in a bundle, so a
    reader of r must come after r's writer, and a new writer of r may share a
    bundle with the last reader of r (but not come before it)."""

    def __init__(self):
        self.bundles = []
        self.written = {}  # reg -> bundle of its latest writer
        self.read = {}  # reg -> latest bundle that reads it

    def add(self, ins):
        reads, writes = regs(ins)
        engine = ENGINE[ins[0]]
        t = 0
        for r in reads:
            t = max(t, self.written.get(r, -1) + 1)  # read after write
        for w in writes:
            t = max(t, self.written.get(w, -1) + 1, self.read.get(w, 0))  # write after write / read
        while True:
            if t == len(self.bundles):
                self.bundles.append([])
            if sum(ENGINE[i[0]] == engine for i in self.bundles[t]) < SLOTS[engine]:
                break
            t += 1
        self.bundles[t].append(ins)
        for w in writes:
            self.written[w] = t
        for r in reads:
            self.read[r] = max(self.read.get(r, 0), t)


def build_kernel(scalar_elems=40):
    """The first N - scalar_elems elements run 8 lanes at a time on the vector
    unit; the last scalar_elems run on the scalar ALUs, which would otherwise
    sit idle after computing addresses."""
    assert (N - scalar_elems) % VLEN == 0
    sched = Scheduler()
    s_regs, v_regs = itertools.count(), itertools.count()  # fresh registers: no false dependencies
    sreg, vreg = s_regs.__next__, v_regs.__next__

    # Constants once: a scalar copy for the scalar path, a broadcast for vectors.
    scalar_k, vector_k = {}, {}
    for value in sorted({c for rnd in ROUNDS for c in rnd}):
        scalar_k[value] = sreg()
        sched.add(("const", scalar_k[value], value))
        vector_k[value] = vreg()
        sched.add(("vbroadcast", vector_k[value], scalar_k[value]))

    vectors = [(i, sreg(), sreg(), vreg(), vreg()) for i in range(0, N - scalar_elems, VLEN)]
    scalars = [(i, sreg(), sreg(), sreg(), sreg()) for i in range(N - scalar_elems, N)]
    lanes = [("v", *x) for x in vectors] + [("s", *x) for x in scalars]

    # Element by element: each one finishes early, so its store overlaps the
    # arithmetic of the elements still in flight.
    for kind, i, a_in, a_out, val, tmp in lanes:
        k = vector_k if kind == "v" else scalar_k
        unit = "valu" if kind == "v" else "alu"
        sched.add(("const", a_in, IN + i))
        sched.add(("vload" if kind == "v" else "load", val, a_in))
        for mul, add, shift in ROUNDS:
            sched.add((unit, "*", val, val, k[mul]))
            sched.add((unit, "+", val, val, k[add]))
            sched.add((unit, ">>", tmp, val, k[shift]))
            sched.add((unit, "^", val, val, tmp))
        sched.add(("const", a_out, OUT + i))
        sched.add(("vstore" if kind == "v" else "store", a_out, val))
    return sched.bundles

The scheduler is greedy and processes instructions in program order, which is why the order of emission matters (step 5). It is 30 lines and good enough here: 90 cycles is about 8% above the bound of 83.

Level 1 checks correctness on random and edge inputs. Level 2 asks for a 10× speedup (packing or vectorizing). Level 3 asks for both.

import random

from kernel import build_kernel
from machine import MASK, N, reference, run_kernel


def test_matches_the_reference_on_random_inputs():
    rng = random.Random(42)
    for _ in range(3):
        values = [rng.getrandbits(32) for _ in range(N)]
        out, _ = run_kernel(build_kernel, values)
        assert out == reference(values)


def test_matches_the_reference_on_edge_values():
    values = [0, 1, MASK, MASK - 1, 1 << 31, 7] * (N // 6) + [0] * (N % 6)
    out, _ = run_kernel(build_kernel, values)
    assert out == reference(values)

import random

from kernel import build_kernel
from machine import N, reference, run_kernel


def test_at_least_10x_faster_than_the_baseline():
    rng = random.Random(1)
    values = [rng.getrandbits(32) for _ in range(N)]
    out, cycles = run_kernel(build_kernel, values)
    assert out == reference(values)
    assert cycles <= 600, cycles  # baseline: 6400

import random

from kernel import build_kernel
from machine import N, reference, run_kernel


def test_vectorized_and_packed():
    rng = random.Random(3)
    values = [rng.getrandbits(32) for _ in range(N)]
    out, cycles = run_kernel(build_kernel, values)
    assert out == reference(values)
    # The vector unit alone needs (32 vectors x 12 ops + 7 broadcasts) / 4 slots = 98 cycles.
    assert cycles <= 110, cycles

The baseline passes level 1 and fails levels 2 and 3 (6,400 cycles). The kernel above passes all three at 90 cycles.

Say these in the write-up even if you run out of time:

1. Close the 90 → 83 gap. The remaining cycles are startup (constants, the first loads) and the tail (the last chains and their stores). Schedule the longest chains first (critical-path priority) instead of in program order, and start the scalar elements before the vectors, since their chains are the longest per engine slot.
2. Strength-reduce the multiply. v × 17 = (v << 4) + v. That adds operations here, so it only pays if a multiply costs more cycles than a shift. On this machine all ALU operations cost one slot, so leave it. Saying why you didn't do it shows you read the cost model.
3. Scale. For millions of elements the program would not be straight-line code. Software-pipeline a loop body (load for iteration i+2, compute i+1, store i) so each bundle stays full in steady state.

Mid-level (L4)

• Keeps the kernel exactly correct throughout, with a reference check run after every change.
• Hoists the constants and either packs bundles or vectorizes, reaching about 10×.
• Explains the read-before-write rule and why reusing registers serializes work.

Senior (L5)

• Computes the per-engine lower bounds before optimizing and uses them to pick the next step.
• Vectorizes and writes a list scheduler, reaching about 60× (around 100 cycles).
• Explains the store-engine tail and fixes it with the emission order; keeps a table of steps and cycle counts.

Staff (L6+)

• Sees the idle scalar ALUs and balances work across both units with the bound formula, reaching about 70× (90 cycles), and says how close that is to the bound.
• Discusses scheduler priorities (critical path first), software pipelining for loops, and when a strength reduction does or doesn't pay on this cost model.
• Writes a short, honest report: what bounds the result, what was tried and dropped, what comes next.