Kernel-optimization take-home scored by speedup

Question: Kernel-optimization take-home scored by speedup

Companies: Anthropic

Difficulty: Hard

Problem statement

A two-hour take-home that replaces the coding screen. AI tools are allowed. You get a simulator for a small accelerator and a working but slow kernel; your score is the speedup in simulated cycles over the baseline, and the result must stay exactly correct.

The machine runs bundles of instructions, one bundle per cycle, with a fixed number of slots per engine (8 scalar ALU, 4 vector ALU over 8 lanes, 2 loads, 1 store). All instructions in a bundle read before any of them write. The kernel applies three rounds of a multiply, add, shift and xor hash to 256 32-bit values. The baseline issues one instruction per cycle and takes 6,400 cycles.

Edit only kernel.py. build_kernel() does not see the input values. Explain what bounds your cycle count and what you would try next.

"""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

How to spend the two hours

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.

Read the machine: the cost model

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.

Step by step

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.

The rules that cause wrong answers

• 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.

The optimized kernel

"""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.

Tests

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.

What to try next

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.

Interviewer expectations

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.