VLIW Kernel Optimization

Earlier this year, Anthropic released one of their interview questions - a take home challenge to optimize a kernel (not an OS kernel, a kernel in the sense of being the core operation loop of a program) for a toy VLIW (Very Large Instruction Word) architecture. While this type of programming is not my specialty per-se, the more I mentioned this challenge to people in casual conversation, the more this challenge felt like an exceptionally good way to help people understand and visualize what is going on with these types of problems, and what each iterative improvement towards making the kernel faster is actually doing. In this two part blog series, I’d like to look at this problem from two perspectives:

1. Looking at the problem “manually”, and walking through some basic improvements to optimize it.
2. Examining and comparing how various LLMs do on the challenge.

The Challenge

The problem implements a toy VLIW SIMD machine, and has a sample kernel implemented for it which is currently (terribly) inefficient.

VLIW?

Unless you happen to have done DSP or some TPU/NPU programming (or maybe were unlucky enough to deal with Itanium…) VLIW may not be a familiar term. At its core, VLIW architectures allow assembly level control over each processing unit (arithmetic (ALU), load/store, etc.) in each instruction instead of dispatching and pipelining single instructions automatically to effeciently use all available units (as most modern CPU architectures do). What does this look like? Well a VLIW instruction might look like this (adapted from the challenge repo):

{"valu": [("*", 4, 0, 0), ("+", 8, 4, 4)], "load": [("load", 16, 17)]}

So in one instruction, we’re actually commanding two separate vector ALU units plus one load unit. On most CPU architectures of today we’d have something like:

$4 = MUL $0, $0
$8 = ADD $4, $4
$16 = LOAD $17

And the CPU hardware would be responsible for dispatching each of these sequential, independent instructions to most effectively utilize the available hardware units. Conceptually you can think of this as moving the responsibility for dispatch from hardware up into software, specifically the compiler.

Why would we want this? Well in short, it’s a difficult problem! If this dispatching can be computed before-hand (potentially with better knowledge about aliasing etc. from the source level), we could get overall efficiency gains by increasing unit utilization.

The Architecture

With that high-level primer of VLIW out of the way, there’s just a few other points worth noting about the specific toy architecture:

• We have 12 ALU units, 6 vector (SIMD) ALU units, and 2 each load and store units
• The word size is 32 bits
• Vector ALUs operate on 8 words (=256 bits) at a time
• All instructions complete in one step() - i.e. there are no instructions which take more clock cycles than others
• There is no limit to the number of those units we can dispatch in a single instruction - in theory we could load up all 22

The Kernel

The kernel as implemented essentially does multiple independent walks of a perfect binary tree, using an accumulated value to figure out which way to walk next. Each walker starts with a random value. On each iteration, the value is XOR’d with the binary tree node’s value, the result hashed, and then the low bit of the resultant used to determine whether to walk left or right down the tree. This behavior isn’t directly applicable to anything real-world, but as we’ll see it ends up with similar opportunities for optimization as LLM inference.

Finally, optimizations

With all of that prereq out of the way, let’s take a look at a trace of the kernel as originally written.

What we’re looking to optimize for here is basically slot utilization - we want the processor to be doing as much as possible on each step. Without even looking at the code, as you can see out of the gate we’re only using a single ALU (and not even a vector one at that). And if you scroll down, you’ll also see we’re thrashing the load/store units with a lot of tmp var load/stores.

Not good. 147,734 cycles for the test.

So how can we do better?

Reference Kernel

Let’s examine the main loop in the kernel. A few things pop out immediately, lining up with our observations from the visualization:

# Slightly modified from perf_takehome.pyfor round in range(rounds):    for i in range(batch_size):        i_const = self.scratch_const(i)        # idx = mem[inp_indices_p + i]        self.instrs.append(            {"alu": [("+", tmp_addr, self.scratch["inp_indices_p"], i_const)]}        )        self.instrs.append({"load": [("load", tmp_idx, tmp_addr)]})        self.instrs.append(            {"debug": [("compare", tmp_idx, (round, i, "idx"))]}        )        # val = mem[inp_values_p + i]        self.instrs.append(            {"alu": [("+", tmp_addr, self.scratch["inp_values_p"], i_const)]}        )        self.instrs.append({"load": [("load", tmp_val, tmp_addr)]})        self.instrs.append(            {"debug": [("compare", tmp_val, (round, i, "val"))]}        )        # node_val = mem[forest_values_p + idx]        self.instrs.append(            {                "alu": [                    ("+", tmp_addr, self.scratch["forest_values_p"], tmp_idx)                ]            }        )        self.instrs.append({"load": [("load", tmp_node_val, tmp_addr)]})        self.instrs.append(            {"debug": [("compare", tmp_node_val, (round, i, "node_val"))]}        )        # val = myhash(val ^ node_val)        self.instrs.append({"alu": [("^", tmp_val, tmp_val, tmp_node_val)]})        self.instrs.extend(self.build_hash(tmp_val, tmp1, tmp2, round, i))        self.instrs.append(            {"debug": [("compare", tmp_val, (round, i, "hashed_val"))]}        )        # idx = 2*idx + (1 if val % 2 == 0 else 2)        self.instrs.append({"alu": [("%", tmp1, tmp_val, two_const)]})        self.instrs.append({"alu": [("==", tmp1, tmp1, zero_const)]})        self.instrs.append(            {"flow": [("select", tmp3, tmp1, one_const, two_const)]}        )        self.instrs.append({"alu": [("*", tmp_idx, tmp_idx, two_const)]})        self.instrs.append({"alu": [("+", tmp_idx, tmp_idx, tmp3)]})        self.instrs.append(            {"debug": [("compare", tmp_idx, (round, i, "next_idx"))]}        )        # idx = 0 if idx >= n_nodes else idx        self.instrs.append(            {"alu": [("<", tmp1, tmp_idx, self.scratch["n_nodes"])]}        )        self.instrs.append(            {"flow": [("select", tmp_idx, tmp1, tmp_idx, zero_const)]}        )        self.instrs.append(            {"debug": [("compare", tmp_idx, (round, i, "wrapped_idx"))]}        )        # mem[inp_indices_p + i] = idx        self.instrs.append(            {"alu": [("+", tmp_addr, self.scratch["inp_indices_p"], i_const)]}        )        self.instrs.append({"store": [("store", tmp_addr, tmp_idx)]})        # mem[inp_values_p + i] = val        self.instrs.append(            {"alu": [("+", tmp_addr, self.scratch["inp_values_p"], i_const)]}        )        self.instrs.append({"store": [("store", tmp_addr, tmp_val)]})

Let’s go through one-by-one and see what these get us.

Making Very Large Instruction Words

With just a few changes, we can fold both the idx and val loading into a single VLIW instruction

# ...
# idx = mem[inp_indices_p + i], val = mem[inp_values_p + i]
self.instrs.append(
    {"alu": [
        ("+", tmp_addr, self.scratch["inp_indices_p"], i_const),
        ("+", tmp_addr2, self.scratch["inp_values_p"], i_const)
    ]}
)
self.instrs.append({"load": [
    ("load", tmp_idx, tmp_addr),
    ("load", tmp_val, tmp_addr2),
]})
# ...

and see our first signs of multi-unit utilization:

All of these traces are interactive - scroll to see all of the different lanes, ctrl+scroll (or pinch) to zoom, and grab to pan around.

131,350 cycles now, about a 12.5% perf increase. Not bad! But we can do a lot better.

diff

Making Even Larger Instruction Words

As stated above, conceptually this kernel is doing a bunch of independent walks over a binary tree. We like (data) independence! That means parallelization!

We’ve got 12 (non-vector) ALU units, and are currently max’ing out with 2 ALU units per “walker”, so let’s parallelize out up to 6 walkers (“lanes”) in a single batch.

That’s looking much better, and we’re also down to 37,286 cycles. Almost 4x faster than the original!

With that, all of our ALUs are working pretty hard but there’s still very noticable gaps if you zoom in which appear to be taken up by load/store operations. Can we improve there?

diff

Swapping Loop Order

You may have noticed this from the original kernel code - instead of seeing each “walker” through to its final state (thus keeping the most amount of in-progress state local vs. having to write it out to memory), we’re instead stepping all of them forward at once, resulting in all intermediate state having to go through load/stores instead of being in fast local scratch memory.

How about we swap the loop nesting so we process one thing all the way through?

Down to 28,316 cycles now! That alone shaves off almost 25% of the cycles, and makes us ~5x faster than the original.

diff

“Branchless”

There’s many, many more places for optimization in this (we’re still over 20x slower than the best numbers Anthropic has seen), but I want to highlight one last bottleneck before wrapping up this post.

You may have noticed that the flow unit is being used quite heavily - eyeballing it I’d say maybe 30-40% of cycles have flow utilization. Given we only have one of those units, that could be bottlenecking everything else. So can we remove it?

Even outside of massively parallel programming, you may be familiar with “branchless” programming: code which has no conditional statements. In traditional CPUs, branchless can lead to performance speedups because the CPU does not have to fully serialize its pipeline on the condition. Conditionals cannot be resolved (meaning the CPU doesn’t actually know what to do next) until the real computation they depend on has happened, and with CPUs holding many (possibly many dozens) of instructions in flight, this can mean the CPU stalls the pipeline (unable to fetch and load new instructions) until most if not all of the existing pipeline has finished/retired, and the result of the conditional is known.

While it’s for a different reason here (GPUs and other parallel architectures usually don’t do any type of pipelining), we still end up with a similar bottleneck in flow control but also have a similar fix: use math!

There’s 2 places we need to fix up. First,

idx = 2*idx + (1 if val % 2 == 0 else 2)

Can be replaced with

idx = 2*idx + ((val % 2) + 1)

And further, % 2 is really just & 1, so

idx = 2*idx + ((val & 1) + 1)

With this one too, we can use our VLIW powers and do the 2*idx and val&1 computations in the same cycle, which means the whole computation only takes three cycles

Cycle 1:
idx = 2*idx
tmp_addend = val & 1

Cycle 2:
tmp_addend = tmp_addend + 1

Cycle 3:
idx = idx + tmp_addend

The second conditional

idx = 0 if idx >= n_nodes else idx

Can be replaced with

idx = idx * (idx >= n_nodes)

Requiring two cycles.

20,124 cycles now, an ~30% decrease from before and ~7.3x faster than what we started with.

diff

Conclusion

Again, 20k cycles is still slow compared to what a fully optimized kernel looks like, but this post would go on for a looong time if I were to detail each and every optimization to get down to 1.5k-2k cycles.

I hope this highlights some of the principles of optimization though. A lot of what is required to get even further is just more of the same. Using all of the units available to us (we haven’t touched SIMD yet), even better instruction-level parallelism (we haven’t touched the hash function at all either!), reducing memory load/stores and keeping more in scratch.

In part 2 (coming soon) we’ll see how far modern(ish) LLMs are able to push this down, their “strategies” for doing so, and what it was like to see them run for a few hours on a hillclimbing (or well, hilldescending?) task like this.