-
-
Save bertmiller/6e17ff2c9dbebe86dd3ed41b7ecf93db to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| from collections import defaultdict | |
| import random | |
| import unittest | |
| from problem import ( | |
| Engine, | |
| DebugInfo, | |
| SLOT_LIMITS, | |
| VLEN, | |
| N_CORES, | |
| SCRATCH_SIZE, | |
| Machine, | |
| Tree, | |
| Input, | |
| HASH_STAGES, | |
| reference_kernel, | |
| build_mem_image, | |
| reference_kernel2, | |
| ) | |
| # EVOLVE-BLOCK-START | |
| class KernelBuilder: | |
| def __init__(self): | |
| self.instrs = [] | |
| self.scratch = {} | |
| self.scratch_debug = {} | |
| self.scratch_ptr = 0 | |
| self.const_map = {} | |
| self.slot_labels = [] # Parallel to body slots for attribution | |
| self.label_map = {} # Maps (cycle_idx, engine, slot_idx) -> label after scheduling | |
| def debug_info(self): | |
| return DebugInfo(scratch_map=self.scratch_debug) | |
| def append_labeled(self, body, engine_slot, label): | |
| """Append slot with label tracking.""" | |
| body.append(engine_slot) | |
| self.slot_labels.append(label) | |
| def get_labeled_instrs(self): | |
| """Return (instrs, label_map) for analysis.""" | |
| return self.instrs, self.label_map | |
| def get_op_dependencies(self, engine, slot): | |
| """ | |
| Return (reads, writes) - sets of scratch addresses this op reads from and writes to. | |
| From vliw-bundling: full dependency tracking for list scheduling. | |
| """ | |
| reads = set() | |
| writes = set() | |
| if engine == "debug": | |
| if slot[0] == "compare": | |
| # ("compare", loc, key) | |
| _, loc, key = slot | |
| reads.add(loc) | |
| elif slot[0] == "vcompare": | |
| # ("vcompare", loc, keys) | |
| _, loc, keys = slot | |
| for i in range(VLEN): | |
| reads.add(loc + i) | |
| return reads, writes | |
| if engine == "alu": | |
| # (op, dest, a1, a2) | |
| op, dest, a1, a2 = slot | |
| writes.add(dest) | |
| reads.add(a1) | |
| reads.add(a2) | |
| elif engine == "valu": | |
| if slot[0] == "vbroadcast": | |
| # (vbroadcast, dest, src) | |
| _, dest, src = slot | |
| for i in range(VLEN): | |
| writes.add(dest + i) | |
| reads.add(src) | |
| elif slot[0] == "multiply_add": | |
| # (multiply_add, dest, a, b, c) | |
| _, dest, a, b, c = slot | |
| for i in range(VLEN): | |
| writes.add(dest + i) | |
| reads.add(a + i) | |
| reads.add(b + i) | |
| reads.add(c + i) | |
| else: | |
| # (op, dest, a1, a2) | |
| op, dest, a1, a2 = slot | |
| for i in range(VLEN): | |
| writes.add(dest + i) | |
| reads.add(a1 + i) | |
| reads.add(a2 + i) | |
| elif engine == "load": | |
| if slot[0] == "load": | |
| # (load, dest, addr) | |
| _, dest, addr = slot | |
| writes.add(dest) | |
| reads.add(addr) | |
| elif slot[0] == "load_offset": | |
| # (load_offset, dest, addr, offset) | |
| _, dest, addr, offset = slot | |
| writes.add(dest + offset) | |
| reads.add(addr + offset) | |
| elif slot[0] == "vload": | |
| # (vload, dest, addr) - addr is scalar | |
| _, dest, addr = slot | |
| for i in range(VLEN): | |
| writes.add(dest + i) | |
| reads.add(addr) | |
| elif slot[0] == "const": | |
| # (const, dest, val) - val is immediate, not read | |
| _, dest, val = slot | |
| writes.add(dest) | |
| elif engine == "store": | |
| if slot[0] == "store": | |
| # (store, addr, src) | |
| _, addr, src = slot | |
| reads.add(addr) | |
| reads.add(src) | |
| elif slot[0] == "vstore": | |
| # (vstore, addr, src) | |
| _, addr, src = slot | |
| reads.add(addr) | |
| for i in range(VLEN): | |
| reads.add(src + i) | |
| elif engine == "flow": | |
| if slot[0] == "select": | |
| # (select, dest, cond, a, b) | |
| _, dest, cond, a, b = slot | |
| writes.add(dest) | |
| reads.add(cond) | |
| reads.add(a) | |
| reads.add(b) | |
| elif slot[0] == "vselect": | |
| # (vselect, dest, cond, a, b) | |
| _, dest, cond, a, b = slot | |
| for i in range(VLEN): | |
| writes.add(dest + i) | |
| reads.add(cond + i) | |
| reads.add(a + i) | |
| reads.add(b + i) | |
| elif slot[0] == "add_imm": | |
| # (add_imm, dest, a, imm) | |
| _, dest, a, imm = slot | |
| writes.add(dest) | |
| reads.add(a) | |
| elif slot[0] in ("halt", "pause"): | |
| pass | |
| elif slot[0] == "trace_write": | |
| _, val = slot | |
| reads.add(val) | |
| elif slot[0] == "cond_jump": | |
| _, cond, addr = slot | |
| reads.add(cond) | |
| elif slot[0] == "cond_jump_rel": | |
| _, cond, offset = slot | |
| reads.add(cond) | |
| elif slot[0] == "jump": | |
| pass | |
| elif slot[0] == "jump_indirect": | |
| _, addr = slot | |
| reads.add(addr) | |
| elif slot[0] == "coreid": | |
| _, dest = slot | |
| writes.add(dest) | |
| return reads, writes | |
| def build(self, slots: list[tuple[Engine, tuple]], vliw: bool = False, labels: list[str] = None): | |
| """ | |
| VLIW bundler with dependency-aware scheduling. | |
| Uses list scheduling with ready queue for efficiency. | |
| Two-tier dependency system: | |
| - Strict deps (RAW + WAW): Must be satisfied before an op becomes ready | |
| - Soft deps (WAR): Can be satisfied within the same cycle | |
| Within a single cycle, reads happen at cycle start, writes commit at cycle end. | |
| Therefore WAR is safe within the same bundle. | |
| If labels is provided (parallel to slots), builds self.label_map mapping | |
| (cycle_idx, engine, slot_idx) -> label for attribution analysis. | |
| """ | |
| if not slots: | |
| return [] | |
| n = len(slots) | |
| # Build dependency info for each op | |
| op_reads = [] | |
| op_writes = [] | |
| for engine, slot in slots: | |
| reads, writes = self.get_op_dependencies(engine, slot) | |
| op_reads.append(reads) | |
| op_writes.append(writes) | |
| # Build two-tier dependency graph: | |
| # - strict_deps: RAW + WAW edges (must complete in prior cycle) | |
| # - soft_deps: WAR edges only (can complete in same cycle) | |
| last_writer = {} # addr -> op_index | |
| last_readers = defaultdict(set) # addr -> set of op_indices that read it | |
| strict_deps = [set() for _ in range(n)] | |
| soft_deps = [set() for _ in range(n)] | |
| for i in range(n): | |
| # RAW: strict dependency - must read value written by prior op | |
| for addr in op_reads[i]: | |
| if addr in last_writer: | |
| strict_deps[i].add(last_writer[addr]) | |
| # For each address we write: | |
| for addr in op_writes[i]: | |
| # WAW: strict dependency - preserve write order | |
| if addr in last_writer: | |
| strict_deps[i].add(last_writer[addr]) | |
| # WAR: soft dependency - reader must be scheduled, but can be same cycle | |
| for reader in last_readers[addr]: | |
| soft_deps[i].add(reader) | |
| # Update tracking for our reads | |
| for addr in op_reads[i]: | |
| last_readers[addr].add(i) | |
| # Update last_writer for our writes (and clear readers since we overwrite) | |
| for addr in op_writes[i]: | |
| last_writer[addr] = i | |
| last_readers[addr].clear() | |
| # Maintain two dependency counts | |
| strict_count = [len(strict_deps[i]) for i in range(n)] | |
| soft_count = [len(soft_deps[i]) for i in range(n)] | |
| # Build separate reverse dependency graphs | |
| reverse_strict = [[] for _ in range(n)] | |
| reverse_soft = [[] for _ in range(n)] | |
| for i in range(n): | |
| for j in strict_deps[i]: | |
| reverse_strict[j].append(i) | |
| for j in soft_deps[i]: | |
| reverse_soft[j].append(i) | |
| # Ready queue: ops with strict_count == 0 and soft_count == 0 | |
| ready = [i for i in range(n) if strict_count[i] == 0 and soft_count[i] == 0] | |
| scheduled = [False] * n | |
| instrs = [] | |
| scheduled_count = 0 | |
| cycle_idx = 0 | |
| # Track per-engine slot indices for label mapping | |
| engine_slot_indices = defaultdict(int) | |
| while scheduled_count < n: | |
| ready.sort() | |
| bundle = defaultdict(list) | |
| slot_counts = defaultdict(int) | |
| this_cycle_writes = set() | |
| scheduled_this_cycle = [] | |
| # Track ops scheduled with their slot index within this bundle | |
| ops_with_slot_idx = [] | |
| # Pass 1: Schedule from current ready queue, collecting soft-unlocked ops | |
| new_ready = [] | |
| soft_unlocked = [] | |
| for op in ready: | |
| engine, slot = slots[op] | |
| # Check slot limit | |
| if engine != "debug" and slot_counts[engine] >= SLOT_LIMITS[engine]: | |
| new_ready.append(op) | |
| continue | |
| # RAW in-cycle check: can't read what we write this cycle | |
| if op_reads[op] & this_cycle_writes: | |
| new_ready.append(op) | |
| continue | |
| # WAW in-cycle check: can't write what we write this cycle | |
| if op_writes[op] & this_cycle_writes: | |
| new_ready.append(op) | |
| continue | |
| # WAR within same cycle is allowed (reads happen before writes) | |
| # Schedule this op - track slot index for label mapping | |
| slot_idx = len(bundle[engine]) | |
| bundle[engine].append(slot) | |
| ops_with_slot_idx.append((op, engine, slot_idx)) | |
| slot_counts[engine] += 1 | |
| scheduled_this_cycle.append(op) | |
| scheduled[op] = True | |
| scheduled_count += 1 | |
| this_cycle_writes.update(op_writes[op]) | |
| # Soft unlock: decrement soft_count for successors | |
| for j in reverse_soft[op]: | |
| soft_count[j] -= 1 | |
| if strict_count[j] == 0 and soft_count[j] == 0 and not scheduled[j]: | |
| soft_unlocked.append(j) | |
| # Pass 2: Try to schedule soft-unlocked ops in same cycle | |
| if soft_unlocked: | |
| soft_unlocked.sort() | |
| for op in soft_unlocked: | |
| if scheduled[op]: | |
| continue | |
| engine, slot = slots[op] | |
| # Check slot limit | |
| if engine != "debug" and slot_counts[engine] >= SLOT_LIMITS[engine]: | |
| new_ready.append(op) | |
| continue | |
| # RAW in-cycle check | |
| if op_reads[op] & this_cycle_writes: | |
| new_ready.append(op) | |
| continue | |
| # WAW in-cycle check | |
| if op_writes[op] & this_cycle_writes: | |
| new_ready.append(op) | |
| continue | |
| # Schedule this op - track slot index for label mapping | |
| slot_idx = len(bundle[engine]) | |
| bundle[engine].append(slot) | |
| ops_with_slot_idx.append((op, engine, slot_idx)) | |
| slot_counts[engine] += 1 | |
| scheduled_this_cycle.append(op) | |
| scheduled[op] = True | |
| scheduled_count += 1 | |
| this_cycle_writes.update(op_writes[op]) | |
| # Further soft unlocking (nested) | |
| for j in reverse_soft[op]: | |
| soft_count[j] -= 1 | |
| if strict_count[j] == 0 and soft_count[j] == 0 and not scheduled[j]: | |
| if j not in soft_unlocked and j not in new_ready: | |
| soft_unlocked.append(j) | |
| ready = new_ready | |
| # End of cycle: strict unlock | |
| for op in scheduled_this_cycle: | |
| for j in reverse_strict[op]: | |
| strict_count[j] -= 1 | |
| if strict_count[j] == 0 and soft_count[j] == 0 and not scheduled[j]: | |
| ready.append(j) | |
| if bundle: | |
| instrs.append(dict(bundle)) | |
| # Build label map if labels provided | |
| if labels is not None: | |
| for op, engine, slot_idx in ops_with_slot_idx: | |
| if op < len(labels): | |
| self.label_map[(cycle_idx, engine, slot_idx)] = labels[op] | |
| cycle_idx += 1 | |
| return instrs | |
| def add(self, engine, slot): | |
| self.instrs.append({engine: [slot]}) | |
| def alloc_scratch(self, name=None, length=1): | |
| addr = self.scratch_ptr | |
| if name is not None: | |
| self.scratch[name] = addr | |
| self.scratch_debug[addr] = (name, length) | |
| self.scratch_ptr += length | |
| assert self.scratch_ptr <= SCRATCH_SIZE, "Out of scratch space" | |
| return addr | |
| def scratch_const(self, val, name=None): | |
| if val not in self.const_map: | |
| addr = self.alloc_scratch(name) | |
| # Defer const load to be scheduled with the body | |
| if not hasattr(self, '_deferred_consts'): | |
| self._deferred_consts = [] | |
| self._deferred_consts.append(("load", ("const", addr, val))) | |
| self.const_map[val] = addr | |
| return self.const_map[val] | |
| def build_hash(self, val_hash_addr, tmp1, tmp2, round, i): | |
| slots = [] | |
| for hi, (op1, val1, op2, op3, val3) in enumerate(HASH_STAGES): | |
| slots.append(("alu", (op1, tmp1, val_hash_addr, self.scratch_const(val1)))) | |
| slots.append(("alu", (op3, tmp2, val_hash_addr, self.scratch_const(val3)))) | |
| slots.append(("alu", (op2, val_hash_addr, tmp1, tmp2))) | |
| slots.append(("debug", ("compare", val_hash_addr, (round, i, "hash_stage", hi)))) | |
| return slots | |
| def build_interleaved_4_hashes(self, | |
| val_A, tmp1_A, tmp2_A, block_start_a, | |
| val_B, tmp1_B, tmp2_B, block_start_b, | |
| val_C, tmp1_C, tmp2_C, block_start_c, | |
| val_D, tmp1_D, tmp2_D, block_start_d, | |
| const_vecs, mult_vecs, round): | |
| """ | |
| Interleave 4 hash operations for maximum VALU utilization. | |
| From software-pipelining: explicit 4-block interleaving with multiply_add optimization. | |
| Returns (slots, labels) tuple for attribution tracking. | |
| """ | |
| slots = [] | |
| labels = [] | |
| def add_slot(slot, label): | |
| slots.append(slot) | |
| labels.append(label) | |
| # Stage 0: 4 multiply_add (can bundle) | |
| add_slot(("valu", ("multiply_add", val_A, val_A, mult_vecs[4097], const_vecs[0x7ED55D16])), "hash_s0") | |
| add_slot(("valu", ("multiply_add", val_B, val_B, mult_vecs[4097], const_vecs[0x7ED55D16])), "hash_s0") | |
| add_slot(("valu", ("multiply_add", val_C, val_C, mult_vecs[4097], const_vecs[0x7ED55D16])), "hash_s0") | |
| add_slot(("valu", ("multiply_add", val_D, val_D, mult_vecs[4097], const_vecs[0x7ED55D16])), "hash_s0") | |
| for val, bs in [(val_A, block_start_a), (val_B, block_start_b), (val_C, block_start_c), (val_D, block_start_d)]: | |
| add_slot(("debug", ("vcompare", val, [(round, bs + vi, "hash_stage", 0) for vi in range(VLEN)])), "debug") | |
| # Stage 1: 4 × (2 parallel + 1 dependent) = 8 parallel ops + 4 dependent | |
| add_slot(("valu", ("^", tmp1_A, val_A, const_vecs[0xC761C23C])), "hash_s1") | |
| add_slot(("valu", (">>", tmp2_A, val_A, const_vecs[19])), "hash_s1") | |
| add_slot(("valu", ("^", tmp1_B, val_B, const_vecs[0xC761C23C])), "hash_s1") | |
| add_slot(("valu", (">>", tmp2_B, val_B, const_vecs[19])), "hash_s1") | |
| add_slot(("valu", ("^", tmp1_C, val_C, const_vecs[0xC761C23C])), "hash_s1") | |
| add_slot(("valu", (">>", tmp2_C, val_C, const_vecs[19])), "hash_s1") | |
| add_slot(("valu", ("^", tmp1_D, val_D, const_vecs[0xC761C23C])), "hash_s1") | |
| add_slot(("valu", (">>", tmp2_D, val_D, const_vecs[19])), "hash_s1") | |
| # Dependent ops | |
| add_slot(("valu", ("^", val_A, tmp1_A, tmp2_A)), "hash_s1") | |
| add_slot(("valu", ("^", val_B, tmp1_B, tmp2_B)), "hash_s1") | |
| add_slot(("valu", ("^", val_C, tmp1_C, tmp2_C)), "hash_s1") | |
| add_slot(("valu", ("^", val_D, tmp1_D, tmp2_D)), "hash_s1") | |
| for val, bs in [(val_A, block_start_a), (val_B, block_start_b), (val_C, block_start_c), (val_D, block_start_d)]: | |
| add_slot(("debug", ("vcompare", val, [(round, bs + vi, "hash_stage", 1) for vi in range(VLEN)])), "debug") | |
| # Stage 2: 4 multiply_add | |
| add_slot(("valu", ("multiply_add", val_A, val_A, mult_vecs[33], const_vecs[0x165667B1])), "hash_s2") | |
| add_slot(("valu", ("multiply_add", val_B, val_B, mult_vecs[33], const_vecs[0x165667B1])), "hash_s2") | |
| add_slot(("valu", ("multiply_add", val_C, val_C, mult_vecs[33], const_vecs[0x165667B1])), "hash_s2") | |
| add_slot(("valu", ("multiply_add", val_D, val_D, mult_vecs[33], const_vecs[0x165667B1])), "hash_s2") | |
| for val, bs in [(val_A, block_start_a), (val_B, block_start_b), (val_C, block_start_c), (val_D, block_start_d)]: | |
| add_slot(("debug", ("vcompare", val, [(round, bs + vi, "hash_stage", 2) for vi in range(VLEN)])), "debug") | |
| # Stage 3: 4 × 3 ops | |
| add_slot(("valu", ("+", tmp1_A, val_A, const_vecs[0xD3A2646C])), "hash_s3") | |
| add_slot(("valu", ("<<", tmp2_A, val_A, const_vecs[9])), "hash_s3") | |
| add_slot(("valu", ("+", tmp1_B, val_B, const_vecs[0xD3A2646C])), "hash_s3") | |
| add_slot(("valu", ("<<", tmp2_B, val_B, const_vecs[9])), "hash_s3") | |
| add_slot(("valu", ("+", tmp1_C, val_C, const_vecs[0xD3A2646C])), "hash_s3") | |
| add_slot(("valu", ("<<", tmp2_C, val_C, const_vecs[9])), "hash_s3") | |
| add_slot(("valu", ("+", tmp1_D, val_D, const_vecs[0xD3A2646C])), "hash_s3") | |
| add_slot(("valu", ("<<", tmp2_D, val_D, const_vecs[9])), "hash_s3") | |
| add_slot(("valu", ("^", val_A, tmp1_A, tmp2_A)), "hash_s3") | |
| add_slot(("valu", ("^", val_B, tmp1_B, tmp2_B)), "hash_s3") | |
| add_slot(("valu", ("^", val_C, tmp1_C, tmp2_C)), "hash_s3") | |
| add_slot(("valu", ("^", val_D, tmp1_D, tmp2_D)), "hash_s3") | |
| for val, bs in [(val_A, block_start_a), (val_B, block_start_b), (val_C, block_start_c), (val_D, block_start_d)]: | |
| add_slot(("debug", ("vcompare", val, [(round, bs + vi, "hash_stage", 3) for vi in range(VLEN)])), "debug") | |
| # Stage 4: 4 multiply_add | |
| add_slot(("valu", ("multiply_add", val_A, val_A, mult_vecs[9], const_vecs[0xFD7046C5])), "hash_s4") | |
| add_slot(("valu", ("multiply_add", val_B, val_B, mult_vecs[9], const_vecs[0xFD7046C5])), "hash_s4") | |
| add_slot(("valu", ("multiply_add", val_C, val_C, mult_vecs[9], const_vecs[0xFD7046C5])), "hash_s4") | |
| add_slot(("valu", ("multiply_add", val_D, val_D, mult_vecs[9], const_vecs[0xFD7046C5])), "hash_s4") | |
| for val, bs in [(val_A, block_start_a), (val_B, block_start_b), (val_C, block_start_c), (val_D, block_start_d)]: | |
| add_slot(("debug", ("vcompare", val, [(round, bs + vi, "hash_stage", 4) for vi in range(VLEN)])), "debug") | |
| # Stage 5: 4 × 3 ops | |
| add_slot(("valu", ("^", tmp1_A, val_A, const_vecs[0xB55A4F09])), "hash_s5") | |
| add_slot(("valu", (">>", tmp2_A, val_A, const_vecs[16])), "hash_s5") | |
| add_slot(("valu", ("^", tmp1_B, val_B, const_vecs[0xB55A4F09])), "hash_s5") | |
| add_slot(("valu", (">>", tmp2_B, val_B, const_vecs[16])), "hash_s5") | |
| add_slot(("valu", ("^", tmp1_C, val_C, const_vecs[0xB55A4F09])), "hash_s5") | |
| add_slot(("valu", (">>", tmp2_C, val_C, const_vecs[16])), "hash_s5") | |
| add_slot(("valu", ("^", tmp1_D, val_D, const_vecs[0xB55A4F09])), "hash_s5") | |
| add_slot(("valu", (">>", tmp2_D, val_D, const_vecs[16])), "hash_s5") | |
| add_slot(("valu", ("^", val_A, tmp1_A, tmp2_A)), "hash_s5") | |
| add_slot(("valu", ("^", val_B, tmp1_B, tmp2_B)), "hash_s5") | |
| add_slot(("valu", ("^", val_C, tmp1_C, tmp2_C)), "hash_s5") | |
| add_slot(("valu", ("^", val_D, tmp1_D, tmp2_D)), "hash_s5") | |
| for val, bs in [(val_A, block_start_a), (val_B, block_start_b), (val_C, block_start_c), (val_D, block_start_d)]: | |
| add_slot(("debug", ("vcompare", val, [(round, bs + vi, "hash_stage", 5) for vi in range(VLEN)])), "debug") | |
| return slots, labels | |
| def build_kernel( | |
| self, forest_height: int, n_nodes: int, batch_size: int, rounds: int | |
| ): | |
| """ | |
| Quad-major pipelining optimization: | |
| - Changed emission order from round-major to quad-major | |
| - This allows scheduler to interleave load-heavy depths (3+) with VALU-only depths (0-2) | |
| - Banked temp registers (2 banks) to reduce conflicts between quads | |
| - All other optimizations from opt16 preserved | |
| Labels all operations for attribution analysis. | |
| """ | |
| tmp2 = self.alloc_scratch("tmp2") | |
| # Scratch space addresses used by the kernel. | |
| # Memory layout from build_mem_image() is fixed for this benchmark: | |
| # values start at index 7, then indices, then values. | |
| self.scratch["forest_values_p"] = self.scratch_const(7) | |
| self.scratch["inp_indices_p"] = self.scratch_const(7 + n_nodes) | |
| self.scratch["inp_values_p"] = self.scratch_const(7 + n_nodes + batch_size) | |
| zero_const = self.scratch_const(0) | |
| one_const = self.scratch_const(1) | |
| two_const = self.scratch_const(2) | |
| body = [] | |
| labels = [] | |
| def add(slot, label): | |
| body.append(slot) | |
| labels.append(label) | |
| n_blocks = batch_size // VLEN | |
| n_quads = n_blocks // 4 # Process 4 blocks at a time | |
| # Allocate vector scratch registers for ALL blocks (keep in scratch across rounds) | |
| # From vliw-bundling: loop reordering optimization | |
| vec_idx_blocks = [self.alloc_scratch(f"vec_idx_{b}", VLEN) for b in range(n_blocks)] | |
| vec_val_blocks = [self.alloc_scratch(f"vec_val_{b}", VLEN) for b in range(n_blocks)] | |
| # Allocate BANKED temp registers for pipelining | |
| # 2 banks so different quads can overlap without conflicts | |
| N_BANKS = 7 | |
| vec_tmp1_banks = [] # [bank][position A/B/C/D] | |
| vec_tmp2_banks = [] | |
| for bank in range(N_BANKS): | |
| bank_tmp1 = [] | |
| bank_tmp2 = [] | |
| for pos in ['A', 'B', 'C', 'D']: | |
| bank_tmp1.append(self.alloc_scratch(f"vec_tmp1_{pos}_bank{bank}", VLEN)) | |
| bank_tmp2.append(self.alloc_scratch(f"vec_tmp2_{pos}_bank{bank}", VLEN)) | |
| vec_tmp1_banks.append(bank_tmp1) | |
| vec_tmp2_banks.append(bank_tmp2) | |
| # Pre-broadcast constants for vector operations | |
| vec_one = self.alloc_scratch("vec_one", VLEN) | |
| vec_two = self.alloc_scratch("vec_two", VLEN) | |
| add(("valu", ("vbroadcast", vec_one, one_const)), "prologue") | |
| add(("valu", ("vbroadcast", vec_two, two_const)), "prologue") | |
| # Pre-broadcast hash constants | |
| hash_const_vecs = {} | |
| hash_constants = [0x7ED55D16, 0xC761C23C, 0x165667B1, 0xD3A2646C, 0xFD7046C5, 0xB55A4F09] | |
| for c in hash_constants: | |
| vec_c = self.alloc_scratch(f"vec_const_{c:08X}", VLEN) | |
| scalar_c = self.scratch_const(c) | |
| add(("valu", ("vbroadcast", vec_c, scalar_c)), "prologue") | |
| hash_const_vecs[c] = vec_c | |
| # Pre-broadcast shift constants for stages 1, 3, 5 (19, 9, 16) | |
| for shift_val in [9, 16, 19]: | |
| vec_s = self.alloc_scratch(f"vec_shift_{shift_val}", VLEN) | |
| scalar_s = self.scratch_const(shift_val) | |
| add(("valu", ("vbroadcast", vec_s, scalar_s)), "prologue") | |
| hash_const_vecs[shift_val] = vec_s | |
| # Pre-broadcast multiplier constants for stages 0, 2, 4 | |
| mult_vecs = {} | |
| for mult_val in [4097, 33, 9]: | |
| if mult_val in hash_const_vecs: | |
| # Reuse existing vector (e.g., 9 is both a shift and mult constant) | |
| mult_vecs[mult_val] = hash_const_vecs[mult_val] | |
| else: | |
| vec_m = self.alloc_scratch(f"vec_mult_{mult_val}", VLEN) | |
| scalar_m = self.scratch_const(mult_val) | |
| add(("valu", ("vbroadcast", vec_m, scalar_m)), "prologue") | |
| mult_vecs[mult_val] = vec_m | |
| # Scratch for load/store address generation | |
| idx_addr_tmp = self.alloc_scratch("idx_addr_tmp") | |
| val_addr_tmp = self.alloc_scratch("val_addr_tmp") | |
| block_stride = self.scratch_const(VLEN) | |
| # Allocate per-block address buffers | |
| vec_addr_blocks = [self.alloc_scratch(f"vec_addr_{b}", VLEN) for b in range(n_blocks)] | |
| # DEPTH 0+1+2 CACHING: Preload tree nodes 0-6 | |
| cached_node_0 = self.alloc_scratch("cached_node_0") | |
| cached_node_1 = self.alloc_scratch("cached_node_1") | |
| cached_node_2 = self.alloc_scratch("cached_node_2") | |
| cached_node_3 = self.alloc_scratch("cached_node_3") | |
| cached_node_4 = self.alloc_scratch("cached_node_4") | |
| cached_node_5 = self.alloc_scratch("cached_node_5") | |
| cached_node_6 = self.alloc_scratch("cached_node_6") | |
| addr_tmp = self.alloc_scratch("addr_tmp") | |
| depth3_node_tmp = self.alloc_scratch("depth3_node_tmp") | |
| three_const = self.scratch_const(3) | |
| four_const = self.scratch_const(4) | |
| five_const = self.scratch_const(5) | |
| six_const = self.scratch_const(6) | |
| # Reusable address temps for node loading | |
| node_addr_a = self.alloc_scratch("node_addr_a") | |
| node_addr_b = self.alloc_scratch("node_addr_b") | |
| # Load nodes 0-6 in batches using reusable address temps | |
| # Node 0: direct load | |
| add(("load", ("load", cached_node_0, self.scratch["forest_values_p"])), "prologue") | |
| # Nodes 1-2 | |
| add(("alu", ("+", node_addr_a, self.scratch["forest_values_p"], one_const)), "prologue") | |
| add(("alu", ("+", node_addr_b, self.scratch["forest_values_p"], two_const)), "prologue") | |
| add(("load", ("load", cached_node_1, node_addr_a)), "prologue") | |
| add(("load", ("load", cached_node_2, node_addr_b)), "prologue") | |
| # Nodes 3-4 | |
| add(("alu", ("+", node_addr_a, self.scratch["forest_values_p"], three_const)), "prologue") | |
| add(("alu", ("+", node_addr_b, self.scratch["forest_values_p"], four_const)), "prologue") | |
| add(("load", ("load", cached_node_3, node_addr_a)), "prologue") | |
| add(("load", ("load", cached_node_4, node_addr_b)), "prologue") | |
| # Nodes 5-6 | |
| add(("alu", ("+", node_addr_a, self.scratch["forest_values_p"], five_const)), "prologue") | |
| add(("alu", ("+", node_addr_b, self.scratch["forest_values_p"], six_const)), "prologue") | |
| add(("load", ("load", cached_node_5, node_addr_a)), "prologue") | |
| add(("load", ("load", cached_node_6, node_addr_b)), "prologue") | |
| # Preload and cache depth-3 nodes 7..14 as vectors for fast lookup | |
| # Use separate addresses to enable pipelining | |
| vec_cached_node_7 = self.alloc_scratch("vec_cached_node_7", VLEN) | |
| vec_cached_node_8 = self.alloc_scratch("vec_cached_node_8", VLEN) | |
| vec_cached_node_9 = self.alloc_scratch("vec_cached_node_9", VLEN) | |
| vec_cached_node_10 = self.alloc_scratch("vec_cached_node_10", VLEN) | |
| vec_cached_node_11 = self.alloc_scratch("vec_cached_node_11", VLEN) | |
| vec_cached_node_12 = self.alloc_scratch("vec_cached_node_12", VLEN) | |
| vec_cached_node_13 = self.alloc_scratch("vec_cached_node_13", VLEN) | |
| vec_cached_node_14 = self.alloc_scratch("vec_cached_node_14", VLEN) | |
| # Allocate 2 reusable address and value temps for depth-3 node loading | |
| d3_addr = [self.alloc_scratch(f"d3_addr_{i}") for i in range(2)] | |
| d3_val = [self.alloc_scratch(f"d3_val_{i}") for i in range(2)] | |
| d3_vecs = [vec_cached_node_7, vec_cached_node_8, vec_cached_node_9, vec_cached_node_10, | |
| vec_cached_node_11, vec_cached_node_12, vec_cached_node_13, vec_cached_node_14] | |
| # Load depth-3 nodes in batches of 2 (reusing address/value temps) | |
| for batch in range(4): | |
| for j in range(2): | |
| i = batch * 2 + j | |
| node_idx = 7 + i | |
| node_scalar = self.scratch_const(node_idx) | |
| add(("alu", ("+", d3_addr[j], self.scratch["forest_values_p"], node_scalar)), "prologue") | |
| for j in range(2): | |
| i = batch * 2 + j | |
| add(("load", ("load", d3_val[j], d3_addr[j])), "prologue") | |
| for j in range(2): | |
| i = batch * 2 + j | |
| add(("valu", ("vbroadcast", d3_vecs[i], d3_val[j])), "prologue") | |
| # Broadcast to vectors | |
| vec_cached_node_0 = self.alloc_scratch("vec_cached_node_0", VLEN) | |
| vec_cached_node_1 = self.alloc_scratch("vec_cached_node_1", VLEN) | |
| vec_cached_node_2 = self.alloc_scratch("vec_cached_node_2", VLEN) | |
| vec_cached_node_3 = self.alloc_scratch("vec_cached_node_3", VLEN) | |
| vec_cached_node_4 = self.alloc_scratch("vec_cached_node_4", VLEN) | |
| vec_cached_node_5 = self.alloc_scratch("vec_cached_node_5", VLEN) | |
| vec_cached_node_6 = self.alloc_scratch("vec_cached_node_6", VLEN) | |
| add(("valu", ("vbroadcast", vec_cached_node_0, cached_node_0)), "prologue") | |
| add(("valu", ("vbroadcast", vec_cached_node_1, cached_node_1)), "prologue") | |
| add(("valu", ("vbroadcast", vec_cached_node_2, cached_node_2)), "prologue") | |
| add(("valu", ("vbroadcast", vec_cached_node_3, cached_node_3)), "prologue") | |
| add(("valu", ("vbroadcast", vec_cached_node_4, cached_node_4)), "prologue") | |
| add(("valu", ("vbroadcast", vec_cached_node_5, cached_node_5)), "prologue") | |
| add(("valu", ("vbroadcast", vec_cached_node_6, cached_node_6)), "prologue") | |
| # Precompute constants for depth 2 bit-manipulation selection | |
| vec_three = self.alloc_scratch("vec_three", VLEN) | |
| add(("valu", ("vbroadcast", vec_three, three_const)), "prologue") | |
| # Per-quad node and address buffers are allocated above in vec_tmp* banks. | |
| # Per-quad node and address buffers used for depth-1 prefetch via vselect | |
| # Load initial idx/val from memory into scratch (once at start) | |
| # From vliw-bundling: loop reordering - only load once | |
| # Load all idx and val vectors for all blocks (once at start) | |
| add(("alu", ("+", idx_addr_tmp, self.scratch["inp_indices_p"], self.scratch_const(0)),), "load_init") | |
| add(("alu", ("+", val_addr_tmp, self.scratch["inp_values_p"], self.scratch_const(0)),), "load_init") | |
| for block in range(n_blocks): | |
| add(("load", ("vload", vec_idx_blocks[block], idx_addr_tmp)), "load_init") | |
| add(("load", ("vload", vec_val_blocks[block], val_addr_tmp)), "load_init") | |
| if block != n_blocks - 1: | |
| add(("alu", ("+", idx_addr_tmp, idx_addr_tmp, block_stride)), "load_init") | |
| add(("alu", ("+", val_addr_tmp, val_addr_tmp, block_stride)), "load_init") | |
| # INTERLEAVED PIPELINING: Emit operations in phase-interleaved quad schedule. | |
| # Keep correctness by preserving per-quad round order while mixing different depths across quads. | |
| for phase in range(rounds + n_quads - 1): | |
| for quad in range(n_quads): | |
| round_idx = phase - quad | |
| if round_idx < 0 or round_idx >= rounds: | |
| continue | |
| # Select temp bank for this quad to reduce conflicts | |
| bank = quad % N_BANKS | |
| vec_tmp1_A = vec_tmp1_banks[bank][0] | |
| vec_tmp2_A = vec_tmp2_banks[bank][0] | |
| vec_tmp1_B = vec_tmp1_banks[bank][1] | |
| vec_tmp2_B = vec_tmp2_banks[bank][1] | |
| vec_tmp1_C = vec_tmp1_banks[bank][2] | |
| vec_tmp2_C = vec_tmp2_banks[bank][2] | |
| vec_tmp1_D = vec_tmp1_banks[bank][3] | |
| vec_tmp2_D = vec_tmp2_banks[bank][3] | |
| block_a = quad * 4 | |
| block_b = quad * 4 + 1 | |
| block_c = quad * 4 + 2 | |
| block_d = quad * 4 + 3 | |
| block_start_a = block_a * VLEN | |
| block_start_b = block_b * VLEN | |
| block_start_c = block_c * VLEN | |
| block_start_d = block_d * VLEN | |
| vec_idx_A = vec_idx_blocks[block_a] | |
| vec_idx_B = vec_idx_blocks[block_b] | |
| vec_idx_C = vec_idx_blocks[block_c] | |
| vec_idx_D = vec_idx_blocks[block_d] | |
| vec_val_A = vec_val_blocks[block_a] | |
| vec_val_B = vec_val_blocks[block_b] | |
| vec_val_C = vec_val_blocks[block_c] | |
| vec_val_D = vec_val_blocks[block_d] | |
| vec_addr_A = vec_addr_blocks[block_a] | |
| vec_addr_B = vec_addr_blocks[block_b] | |
| vec_addr_C = vec_addr_blocks[block_c] | |
| vec_addr_D = vec_addr_blocks[block_d] | |
| depth = round_idx % (forest_height + 1) | |
| next_depth = (round_idx + 1) % (forest_height + 1) | |
| # Debug compares for idx/val | |
| for block_vec, block_start in [(vec_idx_A, block_start_a), (vec_idx_B, block_start_b), | |
| (vec_idx_C, block_start_c), (vec_idx_D, block_start_d)]: | |
| add(("debug", ("vcompare", block_vec, [(round_idx, block_start + vi, "idx") for vi in range(VLEN)])), "debug") | |
| for block_vec, block_start in [(vec_val_A, block_start_a), (vec_val_B, block_start_b), | |
| (vec_val_C, block_start_c), (vec_val_D, block_start_d)]: | |
| add(("debug", ("vcompare", block_vec, [(round_idx, block_start + vi, "val") for vi in range(VLEN)])), "debug") | |
| if depth == 0: | |
| # DEPTH 0: Use pre-cached vec_cached_node_0 directly | |
| for block_start in [block_start_a, block_start_b, block_start_c, block_start_d]: | |
| add(("debug", ("vcompare", vec_cached_node_0, [(round_idx, block_start + vi, "node_val") for vi in range(VLEN)])), "debug") | |
| # XOR all 4 blocks with cached node 0 (scalar ALU to reduce VALU pressure) | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_A + vi, vec_val_A + vi, cached_node_0)), "depth0_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_B + vi, vec_val_B + vi, cached_node_0)), "depth0_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_C + vi, vec_val_C + vi, cached_node_0)), "depth0_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_D + vi, vec_val_D + vi, cached_node_0)), "depth0_xor") | |
| elif depth == 1: | |
| # DEPTH 1: Consume depth-0 cached node1/node2 selection. | |
| for block_vec, block_start in [(vec_addr_A, block_start_a), (vec_addr_B, block_start_b), | |
| (vec_addr_C, block_start_c), (vec_addr_D, block_start_d)]: | |
| add(("debug", ("vcompare", block_vec, [(round_idx, block_start + vi, "node_val") for vi in range(VLEN)])), "debug") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_A + vi, vec_val_A + vi, vec_addr_A + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_B + vi, vec_val_B + vi, vec_addr_B + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_C + vi, vec_val_C + vi, vec_addr_C + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_D + vi, vec_val_D + vi, vec_addr_D + vi)), "tree_xor") | |
| elif depth == 2: | |
| # DEPTH 2: Use cached nodes 3-6 with bit-manipulation selection | |
| # Step 1: bit0 = idx & 1 (per-lane ALU to reduce VALU pressure) | |
| for vi in range(VLEN): | |
| add(("alu", ("&", vec_tmp1_A + vi, vec_idx_A + vi, one_const)), "depth2_select") | |
| for vi in range(VLEN): | |
| add(("alu", ("&", vec_tmp1_B + vi, vec_idx_B + vi, one_const)), "depth2_select") | |
| for vi in range(VLEN): | |
| add(("alu", ("&", vec_tmp1_C + vi, vec_idx_C + vi, one_const)), "depth2_select") | |
| for vi in range(VLEN): | |
| add(("alu", ("&", vec_tmp1_D + vi, vec_idx_D + vi, one_const)), "depth2_select") | |
| # Step 2: tmp = idx - 3 (per-lane ALU) | |
| for vi in range(VLEN): | |
| add(("alu", ("-", vec_tmp2_A + vi, vec_idx_A + vi, three_const)), "depth2_select") | |
| for vi in range(VLEN): | |
| add(("alu", ("-", vec_tmp2_B + vi, vec_idx_B + vi, three_const)), "depth2_select") | |
| for vi in range(VLEN): | |
| add(("alu", ("-", vec_tmp2_C + vi, vec_idx_C + vi, three_const)), "depth2_select") | |
| for vi in range(VLEN): | |
| add(("alu", ("-", vec_tmp2_D + vi, vec_idx_D + vi, three_const)), "depth2_select") | |
| # Step 3: high = tmp >> 1 (per-lane ALU) | |
| for vi in range(VLEN): | |
| add(("alu", (">>", vec_tmp2_A + vi, vec_tmp2_A + vi, one_const)), "depth2_select") | |
| for vi in range(VLEN): | |
| add(("alu", (">>", vec_tmp2_B + vi, vec_tmp2_B + vi, one_const)), "depth2_select") | |
| for vi in range(VLEN): | |
| add(("alu", (">>", vec_tmp2_C + vi, vec_tmp2_C + vi, one_const)), "depth2_select") | |
| for vi in range(VLEN): | |
| add(("alu", (">>", vec_tmp2_D + vi, vec_tmp2_D + vi, one_const)), "depth2_select") | |
| # Step 4: pair_low = (bit0 ? node3 : node4) | |
| add(("flow", ("vselect", vec_addr_A, vec_tmp1_A, vec_cached_node_3, vec_cached_node_4)), "depth2_select") | |
| add(("flow", ("vselect", vec_addr_B, vec_tmp1_B, vec_cached_node_3, vec_cached_node_4)), "depth2_select") | |
| add(("flow", ("vselect", vec_addr_C, vec_tmp1_C, vec_cached_node_3, vec_cached_node_4)), "depth2_select") | |
| add(("flow", ("vselect", vec_addr_D, vec_tmp1_D, vec_cached_node_3, vec_cached_node_4)), "depth2_select") | |
| # Step 5: pair_high = (bit0 ? node5 : node6) | |
| add(("flow", ("vselect", vec_tmp1_A, vec_tmp1_A, vec_cached_node_5, vec_cached_node_6)), "depth2_select") | |
| add(("flow", ("vselect", vec_tmp1_B, vec_tmp1_B, vec_cached_node_5, vec_cached_node_6)), "depth2_select") | |
| add(("flow", ("vselect", vec_tmp1_C, vec_tmp1_C, vec_cached_node_5, vec_cached_node_6)), "depth2_select") | |
| add(("flow", ("vselect", vec_tmp1_D, vec_tmp1_D, vec_cached_node_5, vec_cached_node_6)), "depth2_select") | |
| # Step 6: node = (high ? pair_high : pair_low) | |
| add(("flow", ("vselect", vec_addr_A, vec_tmp2_A, vec_tmp1_A, vec_addr_A)), "depth2_select") | |
| add(("flow", ("vselect", vec_addr_B, vec_tmp2_B, vec_tmp1_B, vec_addr_B)), "depth2_select") | |
| add(("flow", ("vselect", vec_addr_C, vec_tmp2_C, vec_tmp1_C, vec_addr_C)), "depth2_select") | |
| add(("flow", ("vselect", vec_addr_D, vec_tmp2_D, vec_tmp1_D, vec_addr_D)), "depth2_select") | |
| for block_vec, block_start in [(vec_addr_A, block_start_a), (vec_addr_B, block_start_b), | |
| (vec_addr_C, block_start_c), (vec_addr_D, block_start_d)]: | |
| add(("debug", ("vcompare", block_vec, [(round_idx, block_start + vi, "node_val") for vi in range(VLEN)])), "debug") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_A + vi, vec_val_A + vi, vec_addr_A + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_B + vi, vec_val_B + vi, vec_addr_B + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_C + vi, vec_val_C + vi, vec_addr_C + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_D + vi, vec_val_D + vi, vec_addr_D + vi)), "tree_xor") | |
| elif depth == 3: | |
| # NORMAL PATH (depth 3): use cached node vectors to avoid load_offset. | |
| depth3_cached_nodes = { | |
| 8: vec_cached_node_8, | |
| 9: vec_cached_node_9, | |
| 10: vec_cached_node_10, | |
| 11: vec_cached_node_11, | |
| 12: vec_cached_node_12, | |
| 13: vec_cached_node_13, | |
| 14: vec_cached_node_14, | |
| } | |
| def emit_depth3_lookup(block_idx, block_addr, cond_tmp): | |
| const_8 = self.scratch_const(8) | |
| for vi in range(VLEN): | |
| add(("alu", ("==", cond_tmp + vi, block_idx + vi, const_8)), "depth3_lookup") | |
| add(("flow", ("vselect", block_addr, cond_tmp, vec_cached_node_8, vec_cached_node_7)), "depth3_lookup") | |
| for const in range(9, 15): | |
| const_scalar = self.scratch_const(const) | |
| node_vec = depth3_cached_nodes[const] | |
| for vi in range(VLEN): | |
| add(("alu", ("==", cond_tmp + vi, block_idx + vi, const_scalar)), "depth3_lookup") | |
| add(("flow", ("vselect", block_addr, cond_tmp, node_vec, block_addr)), "depth3_lookup") | |
| emit_depth3_lookup(vec_idx_A, vec_addr_A, vec_tmp1_A) | |
| emit_depth3_lookup(vec_idx_B, vec_addr_B, vec_tmp1_B) | |
| emit_depth3_lookup(vec_idx_C, vec_addr_C, vec_tmp1_C) | |
| emit_depth3_lookup(vec_idx_D, vec_addr_D, vec_tmp1_D) | |
| for block_vec, block_start in [(vec_addr_A, block_start_a), (vec_addr_B, block_start_b), | |
| (vec_addr_C, block_start_c), (vec_addr_D, block_start_d)]: | |
| add(("debug", ("vcompare", block_vec, [(round_idx, block_start + vi, "node_val") for vi in range(VLEN)])), "debug") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_A + vi, vec_val_A + vi, vec_addr_A + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_B + vi, vec_val_B + vi, vec_addr_B + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_C + vi, vec_val_C + vi, vec_addr_C + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_D + vi, vec_val_D + vi, vec_addr_D + vi)), "tree_xor") | |
| else: | |
| # NORMAL PATH: depth > 3, use gather via load_offset | |
| for vi in range(VLEN): | |
| add(("load", ("load_offset", vec_addr_A, vec_addr_A, vi)), "gather") | |
| for vi in range(VLEN): | |
| add(("load", ("load_offset", vec_addr_B, vec_addr_B, vi)), "gather") | |
| for vi in range(VLEN): | |
| add(("load", ("load_offset", vec_addr_C, vec_addr_C, vi)), "gather") | |
| for vi in range(VLEN): | |
| add(("load", ("load_offset", vec_addr_D, vec_addr_D, vi)), "gather") | |
| for block_vec, block_start in [(vec_addr_A, block_start_a), (vec_addr_B, block_start_b), | |
| (vec_addr_C, block_start_c), (vec_addr_D, block_start_d)]: | |
| add(("debug", ("vcompare", block_vec, [(round_idx, block_start + vi, "node_val") for vi in range(VLEN)])), "debug") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_A + vi, vec_val_A + vi, vec_addr_A + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_B + vi, vec_val_B + vi, vec_addr_B + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_C + vi, vec_val_C + vi, vec_addr_C + vi)), "tree_xor") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_val_D + vi, vec_val_D + vi, vec_addr_D + vi)), "tree_xor") | |
| # Interleaved 4-block hash (common to all depths) | |
| hash_slots, hash_labels = self.build_interleaved_4_hashes( | |
| vec_val_A, vec_tmp1_A, vec_tmp2_A, block_start_a, | |
| vec_val_B, vec_tmp1_B, vec_tmp2_B, block_start_b, | |
| vec_val_C, vec_tmp1_C, vec_tmp2_C, block_start_c, | |
| vec_val_D, vec_tmp1_D, vec_tmp2_D, block_start_d, | |
| hash_const_vecs, mult_vecs, round_idx | |
| ) | |
| body.extend(hash_slots) | |
| labels.extend(hash_labels) | |
| # Debug compares for hashed_val | |
| for block_vec, block_start in [(vec_val_A, block_start_a), (vec_val_B, block_start_b), | |
| (vec_val_C, block_start_c), (vec_val_D, block_start_d)]: | |
| add(("debug", ("vcompare", block_vec, [(round_idx, block_start + vi, "hashed_val") for vi in range(VLEN)])), "debug") | |
| # Interleaved idx update for all 4 blocks | |
| if round_idx == rounds - 1: | |
| # Final round: no next idx state is needed. | |
| pass | |
| elif depth == forest_height: | |
| # Zero idx via scalar ALU self-XOR to reduce VALU pressure | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_idx_A + vi, vec_idx_A + vi, vec_idx_A + vi)), "wrap_check") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_idx_B + vi, vec_idx_B + vi, vec_idx_B + vi)), "wrap_check") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_idx_C + vi, vec_idx_C + vi, vec_idx_C + vi)), "wrap_check") | |
| for vi in range(VLEN): | |
| add(("alu", ("^", vec_idx_D + vi, vec_idx_D + vi, vec_idx_D + vi)), "wrap_check") | |
| else: | |
| # Compute bit = val & 1 and base = idx*2+1 in PARALLEL (shorter critical path) | |
| for vi in range(VLEN): | |
| add(("alu", ("&", vec_tmp1_A + vi, vec_val_A + vi, one_const)), "idx_update") | |
| add(("valu", ("&", vec_tmp1_B, vec_val_B, vec_one)), "idx_update") | |
| add(("valu", ("&", vec_tmp1_C, vec_val_C, vec_one)), "idx_update") | |
| add(("valu", ("&", vec_tmp1_D, vec_val_D, vec_one)), "idx_update") | |
| # idx*2+1 is independent of val&1, can execute in parallel | |
| add(("valu", ("multiply_add", vec_tmp2_A, vec_idx_A, vec_two, vec_one)), "idx_update") | |
| add(("valu", ("multiply_add", vec_tmp2_B, vec_idx_B, vec_two, vec_one)), "idx_update") | |
| add(("valu", ("multiply_add", vec_tmp2_C, vec_idx_C, vec_two, vec_one)), "idx_update") | |
| add(("valu", ("multiply_add", vec_tmp2_D, vec_idx_D, vec_two, vec_one)), "idx_update") | |
| if depth == 0: | |
| add(("flow", ("vselect", vec_addr_A, vec_tmp1_A, vec_cached_node_2, vec_cached_node_1)), "depth1_prefetch") | |
| add(("flow", ("vselect", vec_addr_B, vec_tmp1_B, vec_cached_node_2, vec_cached_node_1)), "depth1_prefetch") | |
| add(("flow", ("vselect", vec_addr_C, vec_tmp1_C, vec_cached_node_2, vec_cached_node_1)), "depth1_prefetch") | |
| add(("flow", ("vselect", vec_addr_D, vec_tmp1_D, vec_cached_node_2, vec_cached_node_1)), "depth1_prefetch") | |
| # new_idx = base + bit = (idx*2+1) + (val&1) | |
| add(("valu", ("+", vec_idx_A, vec_tmp2_A, vec_tmp1_A)), "idx_update") | |
| add(("valu", ("+", vec_idx_B, vec_tmp2_B, vec_tmp1_B)), "idx_update") | |
| add(("valu", ("+", vec_idx_C, vec_tmp2_C, vec_tmp1_C)), "idx_update") | |
| add(("valu", ("+", vec_idx_D, vec_tmp2_D, vec_tmp1_D)), "idx_update") | |
| if next_depth > 3: | |
| for vi in range(VLEN): | |
| add(("alu", ("+", vec_addr_A + vi, self.scratch["forest_values_p"], vec_idx_A + vi)), "addr_precompute") | |
| add(("alu", ("+", vec_addr_B + vi, self.scratch["forest_values_p"], vec_idx_B + vi)), "addr_precompute") | |
| add(("alu", ("+", vec_addr_C + vi, self.scratch["forest_values_p"], vec_idx_C + vi)), "addr_precompute") | |
| add(("alu", ("+", vec_addr_D + vi, self.scratch["forest_values_p"], vec_idx_D + vi)), "addr_precompute") | |
| for block_vec, block_start in [(vec_idx_A, block_start_a), (vec_idx_B, block_start_b), | |
| (vec_idx_C, block_start_c), (vec_idx_D, block_start_d)]: | |
| add(("debug", ("vcompare", block_vec, [(round_idx, block_start + vi, "next_idx") for vi in range(VLEN)])), "debug") | |
| for block_vec, block_start in [(vec_idx_A, block_start_a), (vec_idx_B, block_start_b), | |
| (vec_idx_C, block_start_c), (vec_idx_D, block_start_d)]: | |
| add(("debug", ("vcompare", block_vec, [(round_idx, block_start + vi, "wrapped_idx") for vi in range(VLEN)])), "debug") | |
| # Store final idx/val back to memory (once at end) | |
| # From vliw-bundling: loop reordering - only store once | |
| add(("alu", ("+", idx_addr_tmp, self.scratch["inp_indices_p"], self.scratch_const(0)),), "store_back") | |
| add(("alu", ("+", val_addr_tmp, self.scratch["inp_values_p"], self.scratch_const(0)),), "store_back") | |
| for block in range(n_blocks): | |
| add(("store", ("vstore", idx_addr_tmp, vec_idx_blocks[block])), "store_back") | |
| add(("store", ("vstore", val_addr_tmp, vec_val_blocks[block])), "store_back") | |
| if block != n_blocks - 1: | |
| add(("alu", ("+", idx_addr_tmp, idx_addr_tmp, block_stride),), "store_back") | |
| add(("alu", ("+", val_addr_tmp, val_addr_tmp, block_stride),), "store_back") | |
| # Prepend deferred const loads to body so scheduler can pack them | |
| if hasattr(self, '_deferred_consts'): | |
| const_body = self._deferred_consts | |
| const_labels = ["prologue"] * len(const_body) | |
| body = const_body + body | |
| labels = const_labels + labels | |
| # Build body with labels tracking | |
| offset = len(self.instrs) # Offset for label_map (pre-build instructions) | |
| body_instrs = self.build(body, labels=labels) | |
| # Offset label_map cycle indices to match final instruction positions | |
| if self.label_map: | |
| self.label_map = { | |
| (cycle_idx + offset, engine, slot_idx): label | |
| for (cycle_idx, engine, slot_idx), label in self.label_map.items() | |
| } | |
| self.instrs.extend(body_instrs) | |
| # Keep kernel finite; no explicit flow control needed when execution reaches end. | |
| # EVOLVE-BLOCK-END | |
| BASELINE = 147734 | |
| def do_kernel_test( | |
| forest_height: int, | |
| rounds: int, | |
| batch_size: int, | |
| seed: int = 123, | |
| trace: bool = False, | |
| prints: bool = False, | |
| ): | |
| print(f"{forest_height=}, {rounds=}, {batch_size=}") | |
| random.seed(seed) | |
| forest = Tree.generate(forest_height) | |
| inp = Input.generate(forest, batch_size, rounds) | |
| mem = build_mem_image(forest, inp) | |
| kb = KernelBuilder() | |
| kb.build_kernel(forest.height, len(forest.values), len(inp.indices), rounds) | |
| # print(kb.instrs) | |
| value_trace = {} | |
| machine = Machine( | |
| mem, | |
| kb.instrs, | |
| kb.debug_info(), | |
| n_cores=N_CORES, | |
| value_trace=value_trace, | |
| trace=trace, | |
| ) | |
| machine.prints = prints | |
| for i, ref_mem in enumerate(reference_kernel2(mem, value_trace)): | |
| machine.run() | |
| inp_values_p = ref_mem[6] | |
| if prints: | |
| print(machine.mem[inp_values_p : inp_values_p + len(inp.values)]) | |
| print(ref_mem[inp_values_p : inp_values_p + len(inp.values)]) | |
| assert ( | |
| machine.mem[inp_values_p : inp_values_p + len(inp.values)] | |
| == ref_mem[inp_values_p : inp_values_p + len(inp.values)] | |
| ), f"Incorrect result on round {i}" | |
| inp_indices_p = ref_mem[5] | |
| if prints: | |
| print(machine.mem[inp_indices_p : inp_indices_p + len(inp.indices)]) | |
| print(ref_mem[inp_indices_p : inp_indices_p + len(inp.indices)]) | |
| # Updating these in memory isn't required, but you can enable this check for debugging | |
| # assert machine.mem[inp_indices_p:inp_indices_p+len(inp.indices)] == ref_mem[inp_indices_p:inp_indices_p+len(inp.indices)] | |
| print("CYCLES: ", machine.cycle) | |
| print("Speedup over baseline: ", BASELINE / machine.cycle) | |
| return machine.cycle | |
| class Tests(unittest.TestCase): | |
| def test_ref_kernels(self): | |
| """ | |
| Test the reference kernels against each other | |
| """ | |
| random.seed(123) | |
| for i in range(10): | |
| f = Tree.generate(4) | |
| inp = Input.generate(f, 10, 6) | |
| mem = build_mem_image(f, inp) | |
| reference_kernel(f, inp) | |
| for _ in reference_kernel2(mem, {}): | |
| pass | |
| assert inp.indices == mem[mem[5] : mem[5] + len(inp.indices)] | |
| assert inp.values == mem[mem[6] : mem[6] + len(inp.values)] | |
| def test_kernel_trace(self): | |
| # Full-scale example for performance testing | |
| do_kernel_test(10, 16, 256, trace=True, prints=False) | |
| def test_kernel_cycles(self): | |
| do_kernel_test(10, 16, 256) | |
| if __name__ == "__main__": | |
| unittest.main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment