Skip to content
 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

14 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Optimizing a VLIW SIMD Tree-Traversal Kernel: 147,734 → 1,453 Cycles (101× Speedup)

The Problem

The kernel performs a parallel tree traversal on a custom VLIW SIMD processor. For each of 256 batch elements across 16 rounds, it:

  1. Reads a node value from a binary tree using the current index
  2. XORs the node value into the running value
  3. Hashes the value through a 6-stage hash function (18 scalar ops)
  4. Updates the index: idx = 2*idx + 1 + (val%2)
  5. Wraps the index back to 0 if it exceeds the tree

The target hardware has per-cycle slot limits: alu:12, valu:6, load:2, store:2, flow:1, debug:64 (free).

The Baseline: 147,734 Cycles

The original build_kernel emitted one operation per instruction bundle. Every ALU op, every load, every store each consumed an entire cycle by itself. For 256 elements × 16 rounds × ~25 ops per element, this produced ~100K+ single-op bundles, each wasting 5 of the 6 engine slots.

Phase 1: VLIW Scheduling (~147K → ~18K)

ASAP Greedy Scheduler (lines 59–211)

The build() method was extended with a VLIW-aware ASAP scheduler that packs independent operations into the same cycle, respecting:

  • RAW dependencies: A read must follow the write that produces its value
  • WAW dependencies: Two writes to the same scratch address need separate cycles
  • WAR dependencies: A write must not precede a prior read of the same address (same cycle is OK since reads see pre-cycle state)
  • Slot limits: Each engine can only execute SLOT_LIMITS[engine] slots per cycle
  • Debug ops are free: They don't increment the cycle counter, only need RAW ordering

The scheduler scans operations in emission order, assigns each to the earliest legal cycle, and assembles multi-engine bundles. This alone packed 6–12 operations per cycle where previously there was 1.

Phase 2: Vectorization (~18K → ~3K)

SIMD with VLEN=8

Instead of processing 256 elements one at a time (256 scalar scatter-loads per round), the kernel processes them in 32 chunks of 8 using vector instructions:

  • vbroadcast to replicate scalars across 8 lanes
  • valu ops (+, ^, %, <<, >>) operate on 8 elements per slot
  • vload/vstore move 8 contiguous elements per slot

This reduced the operation count by ~8× for ALU-bound work. The hash function dropped from 18 scalar ops to 18 vector ops (but now does 8 elements each).

Phase 3: Persistent Scratch & Multi-Buffering (~3K → ~2K)

Persistent Scratch (lines 319–321)

Instead of loading indices/values from memory and storing them back every round, the kernel keeps them in on-chip scratch space (p_idx[k], p_val[k]) across all 16 rounds. Memory loads happen only once at init; memory stores happen only once at the end.

32-Way Multi-Buffering (lines 326–335)

Each of the 32 chunks gets its own independent temporary buffer (bufs[k]). This eliminates WAW/WAR dependencies between chunks, allowing the VLIW scheduler to freely interleave operations from different chunks. Without this, chunk B's ALU would stall waiting for chunk A's store to release the shared temp register.

Phase 4: Hash Optimization (~2K → ~1.8K)

multiply_add Fusion (lines 259–279)

Three of the six hash stages have the pattern val = (val + const1) + (val << shift), which is algebraically val = val * (1 + 2^shift) + const1. The hardware's multiply_add instruction computes dest = a*b + c in a single valu slot, collapsing 3 valu ops into 1 for each of these 3 stages. This saved 6 valu ops per chunk per round = 6 × 32 × 16 = 3,072 valu ops total.

Phase 5: Index Update Optimization

multiply_add for Index (lines 584–588)

The index update idx = 2*idx + 1 + (val%2) was restructured into:

  • multiply_add(idx, idx, 2, 1)idx = 2*idx + 1 (1 valu)
  • %(val, 2)bit = val%2 (1 valu, runs in parallel — no data dependency)
  • +(idx, idx, bit)idx += bit (1 valu)

Total: 3 valu ops with 2 of them parallelizable, down from 4 sequential ops.

Phase 6: Emission Ordering & VLIW Packing (~1.8K → ~1.7K)

Interleaved Chunk-Round Ordering (lines 452–465)

The ASAP scheduler is greedy — it processes operations in emission order and assigns the earliest legal cycle. The emission order dramatically affects packing quality.

The optimal order is IL4-RG15: interleave groups of 4 chunks (INTERLEAVE=4), processing all 15 rounds for each group before moving to the next. This gives the scheduler a mix of operations at different pipeline stages (some chunks doing loads while others do ALU), maximizing engine utilization. This alone saved ~100 cycles over naive round-first ordering.

Phase 7: Tree Pre-Loading (~1.7K → ~1.6K)

Eliminating Scatter Loads for Early Rounds (lines 389–417)

In the first rounds, the tree indices are deterministic (starting from 0):

  • Round 0: all elements at node 0 → 1 possible value
  • Round 1: nodes 1 or 2 → 2 values
  • Round 2: nodes 3–6 → 4 values
  • Round 3: nodes 7–14 → 8 values

These 1+2+4+8 = 15 tree values are pre-loaded and broadcast into scratch vectors during initialization. At runtime, the correct value is selected using vselect (the flow engine) instead of scatter loads (the load engine, limited to 2/cycle).

Pipelined Loading (lines 409–417)

The 15 tree loads are pipelined using 2 alternating temp registers, overlapping the ALU address computation of load N+1 with the memory load of load N.

Phase 8: Binary Tree Selection (~1.6K → ~1.5K)

O(log N) vselect Cascade (lines 511–544)

For depths 2–3, selecting among 4 or 8 pre-loaded values originally used O(N) multiply_add accumulation (4 or 8 valu ops). Replaced with a binary tree cascade:

  • Depth 2 (4 values): 2 valu comparisons + 3 vselect (flow) ops
  • Depth 3 (8 values): 3 valu comparisons + 9 vselect (flow) ops

This shifts work from the valu engine (6 slots/cycle, heavily loaded by hash) to the flow engine (1 slot/cycle, mostly idle), trading a congested resource for an underutilized one.

Phase 9: Deterministic Wrap Exploitation (~1.5K → ~1.5K)

Implicit Wrap at effective_depth == forest_height (lines 594–601)

For a tree of height H, after H+1 rounds the index exceeds n_nodes = 2^(H+1)−1 and wraps to 0. This wrap is deterministic — ALL 256 elements wrap simultaneously. Instead of computing if idx >= n_nodes: idx = 0 (expensive vselect), the code simply:

  1. Skips the wrapped_idx debug check at the wrap round
  2. At the next round (effective_depth == 0), uses the idx-is-zero fast path

idx-is-zero Optimization (lines 577–581)

When effective_depth == 0, the index is known to be 0. The update simplifies from idx = 2*idx + 1 + val%2 (3 ops including multiply_add) to idx = val%2 + 1 (2 ops: one %, one +). This saves 1 valu op × 32 chunks per traversal restart.

Second-Traversal Preloading (line 488)

Since indices restart from 0 after each wrap, the tree values for depths 0–1 (already in scratch) are reused for rounds after the wrap. The can_preload condition detects both first-traversal rounds (rnd ≤ MAX_PRELOAD) and later-traversal rounds (effective_depth ≤ max_preload2 and rnd > forest_height).

Phase 10: VLIW Constant Merging (~1.5K → 1.453K)

Absorbing scratch_const into the VLIW Schedule (lines 609–626)

Throughout the kernel, scratch_const() lazily emits load const instructions. These were appended to self.instrs as single-op bundles outside the VLIW-scheduled body — each consuming a full cycle for just 1 load op.

The fix: after building the body, extract all post-pause const ops from self.instrs and prepend them to the body list before VLIW scheduling. The scheduler then packs them 2/cycle alongside other load ops, saving 51 cycles (one per const that was previously wasted).


Summary

Phase Technique Cycles Speedup
0 Baseline (1 op/cycle) 147,734 1.0×
1 VLIW ASAP scheduler ~18,000 ~8×
2 SIMD vectorization (VLEN=8) ~3,000 ~49×
3 Persistent scratch + 32-way buffering ~2,000 ~74×
4 multiply_add hash fusion ~1,800 ~82×
5 Index update optimization ~1,795 ~82×
6 IL4-RG15 emission ordering ~1,702 ~87×
7 Tree pre-loading (depths 0–3) ~1,625 ~91×
8 Binary tree vselect cascade ~1,515 ~98×
9 Deterministic wrap + idx-is-zero ~1,504 ~98×
10 VLIW constant merging 1,453 101.7×

Theoretical Floor

The workload requires 256 elements × 12 scatter-load rounds × (8 load ops / 2 per cycle) = 1,536 cycles on the load engine alone (rounds 4–10 + 2–10 in second traversal). At 1,453 cycles we are actually below this naive floor because pre-loading and the idx-is-zero path eliminate many scatter rounds entirely.

Generalization

The kernel is fully parameterized for forest_height 8–10 and rounds 8–20. All optimizations adapt automatically via cycle_len = forest_height + 1 and effective_depth = rnd % cycle_len — no round numbers or tree sizes are hardcoded.

About

Anthropic's original performance take-home, now open for you to try!

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages