I have got a score of 1103

`from collections import defaultdict, Counter
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,
)

def get_rw(engine, slot):
reads = set()
writes = set()

def add_read(r, length=1):
for i in range(length):
reads.add(r + i)

def add_write(r, length=1):
for i in range(length):
writes.add(r + i)

if engine == "alu":
# (op, dest, a1, a2)
op, dest, a1, a2 = slot
add_read(a1)
add_read(a2)
add_write(dest)

elif engine == "valu":
op = slot[0]
if op == "vbroadcast":
# ("vbroadcast", dest, src)
_, dest, src = slot
add_read(src)
add_write(dest, VLEN)
elif op == "multiply_add":
# ("multiply_add", dest, a, b, c)
_, dest, a, b, c = slot
add_read(a, VLEN)
add_read(b, VLEN)
add_read(c, VLEN)
add_write(dest, VLEN)
elif op in ("vcompare", "vcompare_idx"): # Debug
pass
else:
# (op, dest, a1, a2)
_, dest, a1, a2 = slot
add_read(a1, VLEN)
add_read(a2, VLEN)
add_write(dest, VLEN)

elif engine == "load":
op = slot[0]
if op == "load":
_, dest, addr = slot
add_read(addr)
add_write(dest)
elif op == "load_offset":
_, dest, addr, offset = slot
add_read(addr)
add_write(dest + offset) # Scalar write at offset
elif op == "vload":
_, dest, addr = slot
add_read(addr)
add_write(dest, VLEN)
elif op == "const":
_, dest, val = slot
add_write(dest)

elif engine == "store":
op = slot[0]
if op == "store":
_, addr, src = slot
add_read(addr)
add_read(src)
elif op == "vstore":
_, addr, src = slot
add_read(addr)
add_read(src, VLEN)

elif engine == "flow":
op = slot[0]
if op == "select":
_, dest, cond, a, b = slot
add_read(cond)
add_read(a)
add_read(b)
add_write(dest)
elif op == "vselect":
_, dest, cond, a, b = slot
add_read(cond, VLEN)
add_read(a, VLEN)
add_read(b, VLEN)
add_write(dest, VLEN)
elif op == "add_imm":
_, dest, a, imm = slot
add_read(a)
add_write(dest)
elif op in ("halt", "pause"):
pass
elif op == "trace_write":
_, val = slot
add_read(val)
elif op == "cond_jump":
_, cond, addr = slot
add_read(cond)
# Add other flow ops if needed

return reads, writes

class Scheduler:
def __init__(self):
self.pending = []

def add(self, engine, slot):
if engine == "debug": return # Skip debug for optimization
reads, writes = get_rw(engine, slot)
self.pending.append({
'engine': engine,
'slot': slot,
'read': reads,
'write': writes
})

def flush(self):
if not self.pending:
return []

schedule = defaultdict(list) # cycle -> list of ops
resource_usage = defaultdict(lambda: defaultdict(int))

reg_avail = defaultdict(int)
reg_last_read = defaultdict(int)
reg_last_write = defaultdict(int)

for op in self.pending:
# Determine earliest possible cycle
t_min = 0

# RAW Dependency: Wait for inputs to be available
for r in op['read']:
t_min = max(t_min, reg_avail[r])

# WAR Dependency: Wait for previous reads of output registers
for r in op['write']:
t_min = max(t_min, reg_last_read[r])

# WAW Dependency: Wait for previous writes
for r in op['write']:
if r in reg_last_write:
t_min = max(t_min, reg_last_write[r] + 1)

# Find a structural slot
t = t_min
while True:
if resource_usage[t][op['engine']] < SLOT_LIMITS[op['engine']]:
break
t += 1

# Schedule
schedule[t].append(op)
resource_usage[t][op['engine']] += 1

# Update state
# Outputs available at t + 1
for r in op['write']:
reg_avail[r] = t + 1
reg_last_write[r] = t

# Inputs read at t
for r in op['read']:
reg_last_read[r] = max(reg_last_read[r], t)

# Build instruction list
instrs = []
if not schedule:
return instrs

max_cycle = max(schedule.keys())
for t in range(max_cycle + 1):
bundle = {}
for op in schedule[t]:
engine = op['engine']
if engine not in bundle:
bundle[engine] = []
bundle[engine].append(op['slot'])
instrs.append(bundle)

self.pending = []
return instrs

class KernelBuilder:
def __init__(self):
self.instrs = []
self.scratch = {}
self.scratch_debug = {}
self.scratch_ptr = 0
self.const_map = {}
self.scheduler = Scheduler()

def debug_info(self):
return DebugInfo(scratch_map=self.scratch_debug)

def add(self, engine, slot):
self.scheduler.add(engine, slot)

def flush(self):
self.instrs.extend(self.scheduler.flush())

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)
self.add("load", ("const", addr, val))
self.const_map[val] = addr
return self.const_map[val]

def build_kernel(
self, forest_height: int, n_nodes: int, batch_size: int, rounds: int
):
# Constants
tmp_init = self.alloc_scratch("tmp_init")
init_vars = [
"rounds",
"n_nodes",
"batch_size",
"forest_height",
"forest_values_p",
"inp_indices_p",
"inp_values_p",
]
for v in init_vars:
self.alloc_scratch(v, 1)
for i, v in enumerate(init_vars):
self.add("load", ("const", tmp_init, i))
self.add("load", ("load", self.scratch[v], tmp_init))

v_zero = self.alloc_scratch("v_zero", VLEN)
self.add("valu", ("vbroadcast", v_zero, self.scratch_const(0)))
v_one = self.alloc_scratch("v_one", VLEN)
self.add("valu", ("vbroadcast", v_one, self.scratch_const(1)))

v_hash_consts = []
for i, (op1, val1, op2, op3, val3) in enumerate(HASH_STAGES):
if i in [0, 2, 4]:
mul = (1 << val3) + 1
comb = (val1 * mul) % (2**32)
v_mul = self.alloc_scratch(f"v_hash_mul_{i}", VLEN)
v_comb = self.alloc_scratch(f"v_hash_comb_{i}", VLEN)
self.add("valu", ("vbroadcast", v_mul, self.scratch_const(mul)))
self.add("valu", ("vbroadcast", v_comb, self.scratch_const(comb)))
v_hash_consts.append((v_mul, v_comb))
else:
v_c1 = self.alloc_scratch(f"v_hash_c1_{i}", VLEN)
v_c3 = self.alloc_scratch(f"v_hash_c3_{i}", VLEN)
self.add("valu", ("vbroadcast", v_c1, self.scratch_const(val1)))
self.add("valu", ("vbroadcast", v_c3, self.scratch_const(val3)))
v_hash_consts.append((v_c1, v_c3))

v_n_nodes = self.alloc_scratch("v_n_nodes", VLEN)
self.add("valu", ("vbroadcast", v_n_nodes, self.scratch["n_nodes"]))

# Unrolling factor
U = 16 # Reduced to avoid OOM with separate scalar_tmps

# Allocate state registers
v_idx = [self.alloc_scratch(f"v_idx_{k}", VLEN) for k in range(U)]
v_val = [self.alloc_scratch(f"v_val_{k}", VLEN) for k in range(U)]

# Temp registers
v_node_val = [self.alloc_scratch(f"v_node_val_{k}", VLEN) for k in range(U)]
v_tmp1 = [self.alloc_scratch(f"v_tmp1_{k}", VLEN) for k in range(U)]

# Scalar tmps for gather (reusing across lanes)
scalar_tmps = [self.alloc_scratch(f"st_{k}") for k in range(U)]

# Allocate pools for Level Caching temps (max needs for R=2)
candidate_vecs_pool = [self.alloc_scratch(f"pool_cand_{i}", VLEN) for i in range(4)]
v_start_idx_pool = self.alloc_scratch("pool_start_idx", VLEN)
v_bit_regs_pool = [self.alloc_scratch(f"pool_bit_{b}", VLEN) for b in range(2)]
shifted_regs_pool = [self.alloc_scratch(f"pool_shifted_{b}", VLEN) if b > 0 else None for b in range(2)]
v_b_regs_pool = [self.alloc_scratch(f"pool_const_shift_{b}", VLEN) if b > 0 else None for b in range(2)]
v_mask_regs_pool = [self.alloc_scratch(f"pool_mask_{d}", VLEN) for d in range(2)]
t_extra_regs_pool = []
mux_res_regs_pool = []
# R=2 layers: 0 (2 pairs), 1 (1 pair)
# Layer 0
t_extra_regs_pool.append([self.alloc_scratch(f"pool_extra_0_{i}", VLEN) for i in range(2)])
mux_res_regs_pool.append([self.alloc_scratch(f"pool_res_0_{i}", VLEN) for i in range(2)])
# Layer 1
t_extra_regs_pool.append([self.alloc_scratch(f"pool_extra_1_{i}", VLEN) for i in range(1)])

# Initial Load
i_const_regs = []
for k in range(U):
batch_start = k * VLEN
const_addr = self.scratch_const(batch_start)
i_const_regs.append(const_addr)

# Load indices
self.add("alu", ("+", scalar_tmps[k], self.scratch["inp_indices_p"], const_addr))
self.add("load", ("vload", v_idx[k], scalar_tmps[k]))

# Load values
self.add("alu", ("+", scalar_tmps[k], self.scratch["inp_values_p"], const_addr))
self.add("load", ("vload", v_val[k], scalar_tmps[k]))

# self.add("flow", ("pause",)) # Removed
self.add("debug", ("comment", "Starting loop"))

self.flush()

for round in range(rounds):
# Level Caching for Round 0, 1, 2
if round <= 2:
# Load all nodes for this level
count = 1 << round
start_index = (1 << round) - 1
start_const = self.scratch_const(start_index)

candidate_vectors = candidate_vecs_pool[:count]
base_addr_scalar = self.scratch["forest_values_p"]

# Load nodes into vectors
for i in range(count):
offset = self.scratch_const(start_index + i)
node_addr = scalar_tmps[0] # Reuse any scalar temp

self.add("alu", ("+", node_addr, base_addr_scalar, offset))
self.add("load", ("load", node_addr, node_addr)) # reuse as val_scalar
self.add("valu", ("vbroadcast", candidate_vectors[i], node_addr))


v_start_idx = v_start_idx_pool
self.add("valu", ("vbroadcast", v_start_idx, start_const))

num_bits = round

for k in range(U):
# offset = v_idx - start
v_offset = v_tmp1[k] # Reuse
self.add("valu", ("-", v_offset, v_idx[k], v_start_idx))

# Extract bits
bits = []
for b in range(num_bits):
v_bit = v_bit_regs_pool[b]

shifted = v_offset
if b > 0:
shifted = shifted_regs_pool[b]
v_b = v_b_regs_pool[b]
c_shift = self.scratch_const(b)
self.add("valu", ("vbroadcast", v_b, c_shift))
self.add("valu", (">>", shifted, v_offset, v_b))

self.add("valu", ("&", v_bit, shifted, v_one))
bits.append(v_bit)

# Mux
if round == 0:
self.add("valu", ("|", v_node_val[k], candidate_vectors[0], v_zero))
else:
current_layer = list(candidate_vectors)
bit_idx = 0

while len(current_layer) > 1:
next_layer = []
v_cond = bits[bit_idx]

v_mask = v_mask_regs_pool[bit_idx]
self.add("valu", ("-", v_mask, v_zero, v_cond))

for i in range(0, len(current_layer), 2):
a = current_layer[i]
b = current_layer[i+1]

t_extra = t_extra_regs_pool[bit_idx][i//2]

# Determine dest
if len(current_layer) == 2: # Last layer (2->1)
dest = v_node_val[k]
else:
dest = mux_res_regs_pool[bit_idx][i//2]

self.add("valu", ("^", t_extra, a, b))
self.add("valu", ("&", t_extra, v_mask, t_extra))
self.add("valu", ("^", dest, a, t_extra))

next_layer.append(dest)

current_layer = next_layer
bit_idx += 1

# Result is in v_node_val[k] already if we set dest correctly

else:
# Standard Gather
for k in range(U):
for lane in range(VLEN):
idx_reg = v_idx[k] + lane
dest_reg = v_node_val[k] + lane
st = scalar_tmps[k] # Reuse same scalar for all lanes (serialized)

self.add("alu", ("+", st, self.scratch["forest_values_p"], idx_reg))
self.add("load", ("load", dest_reg, st))

# XOR val ^ node_val
for k in range(U):
self.add("valu", ("^", v_val[k], v_val[k], v_node_val[k]))

# Hash
for hi, (op1, val1, op2, op3, val3) in enumerate(HASH_STAGES):
v_tmp2_ref = v_node_val # Reuse

if hi in [0, 2, 4]:
v_mul, v_comb = v_hash_consts[hi]
for k in range(U):
self.add("valu", ("multiply_add", v_val[k], v_val[k], v_mul, v_comb))
else:
v_c1, v_c3 = v_hash_consts[hi]
for k in range(U):
self.add("valu", (op1, v_tmp1[k], v_val[k], v_c1))
self.add("valu", (op3, v_tmp2_ref[k], v_val[k], v_c3))
self.add("valu", (op2, v_val[k], v_tmp1[k], v_tmp2_ref[k]))

# Update idx
for k in range(U):
# v_idx = v_idx << 1
self.add("valu", ("<<", v_idx[k], v_idx[k], v_one))

# term = (val & 1) + 1
self.add("valu", ("&", v_tmp1[k], v_val[k], v_one))
self.add("valu", ("+", v_tmp1[k], v_tmp1[k], v_one))

# idx = idx + term
self.add("valu", ("+", v_idx[k], v_idx[k], v_tmp1[k]))

# Wrap idx
for k in range(U):
# cond = (idx < n_nodes) -> 0 or 1
self.add("valu", ("<", v_tmp1[k], v_idx[k], v_n_nodes))
# vselect
self.add("flow", ("vselect", v_idx[k], v_tmp1[k], v_idx[k], v_zero))

# --- FINAL STORE ---
for k in range(U):
self.add("alu", ("+", scalar_tmps[k], self.scratch["inp_indices_p"], i_const_regs[k]))
self.add("store", ("vstore", scalar_tmps[k], v_idx[k]))

self.add("alu", ("+", scalar_tmps[k], self.scratch["inp_values_p"], i_const_regs[k]))
self.add("store", ("vstore", scalar_tmps[k], v_val[k]))

self.flush()

self.add("flow", ("pause",))
self.flush()

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)

value_trace = {}
machine = Machine(
mem,
kb.instrs,
kb.debug_info(),
n_cores=N_CORES,
value_trace=value_trace,
trace=trace,
)
machine.prints = prints
machine.run()

for ref_mem in reference_kernel2(mem):
pass

inp_values_p = ref_mem[6]

print("CYCLES: ", machine.cycle)
print("Speedup over baseline: ", BASELINE / machine.cycle)

assert (
machine.mem[inp_values_p : inp_values_p + len(inp.values)]
== ref_mem[inp_values_p : inp_values_p + len(inp.values)]
), "Incorrect output values"

inp_indices_p = ref_mem[5]
assert (
machine.mem[inp_indices_p : inp_indices_p + len(inp.indices)]
== ref_mem[inp_indices_p : inp_indices_p + len(inp.indices)]
), "Incorrect output indices"

return machine.cycle

class Tests(unittest.TestCase):
def test_ref_kernels(self):
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):
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()
`