Skip to content

perf(codegen): O(n²)→O(n) ISA emit (fixes mlen=16 vision DNF), faster sim_env re-parse, roll RoPE+norm, rename ATEN_UNROLL→ATEN_OPS_UNROLL #56

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
booth-algo merged 3 commits into from
Jun 2, 2026
Merged

perf(codegen): O(n²)→O(n) ISA emit (fixes mlen=16 vision DNF), faster sim_env re-parse, roll RoPE+norm, rename ATEN_UNROLL→ATEN_OPS_UNROLL #56

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a82ce9c
perf(isa): accumulate generated ISA in a list buffer (O(n^2)->O(n)); …
booth-algo Jun 2, 2026
ef2eb2b
perf(assembler): speed up ASM re-parse in sim_env (hoist parse helper…
booth-algo Jun 2, 2026
47fbc78
perf(codegen): roll RoPE + normalization on the unroll flag; rename e…
booth-algo Jun 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 56 additions & 20 deletions asm_templates/normalization_asm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,13 +10,17 @@ def rms_norm_asm(
vlen: int,
batch_size: int,
hidden_dim: int,
unroll: bool = True,
) -> str:
"""
Generate assembly code for RMS normalization.
"""
act_addr = alive_registers[0]
scratchpad_addr = alive_registers[1]
stats_addr = alive_registers[2]
# Rolled path uses the spare 4th register (already allocated by normalize()) as the
# C_LOOP counter. Only accessed when rolled, so unrolled callers may pass 3 registers.
loop_addr = alive_registers[3] if not unroll else None

generated_code = "; RMS Norm generation \n"
generated_code += _load_large_int(scratchpad_addr, scratchpad_base_address)
Expand All @@ -35,13 +39,20 @@ def rms_norm_asm(
generated_code += _load_large_int(stats_addr, activation_base_address + vlen * batch)

# First loop: compute sum of squares using stats_addr
for i in range(hidden_dim // vlen):
# Compute square of the activation vector and summation
if unroll:
for i in range(hidden_dim // vlen):
# Compute square of the activation vector and summation
generated_code += f"V_MUL_VV gp{scratchpad_addr}, gp{stats_addr}, gp{stats_addr}, 0 \n"
generated_code += f"V_RED_SUM f2, gp{scratchpad_addr} \n"

# Move stats pointer to next vector
generated_code += f"S_ADDI_INT gp{stats_addr}, gp{stats_addr}, {vlen * batch_size} \n"
else:
generated_code += f"C_LOOP_START gp{loop_addr}, {hidden_dim // vlen} \n"
generated_code += f"V_MUL_VV gp{scratchpad_addr}, gp{stats_addr}, gp{stats_addr}, 0 \n"
generated_code += f"V_RED_SUM f2, gp{scratchpad_addr} \n"

# Move stats pointer to next vector
generated_code += f"S_ADDI_INT gp{stats_addr}, gp{stats_addr}, {vlen * batch_size} \n"
generated_code += f"C_LOOP_END gp{loop_addr} \n"

# Taking the avg
generated_code += "S_MUL_FP f2, f2, f3 \n"
Expand All @@ -56,12 +67,18 @@ def rms_norm_asm(
generated_code += "S_RECI_FP f2, f2 \n"

# Second loop: normalize using act_addr
for i in range(hidden_dim // vlen):
# Normalize the activation vector
if unroll:
for i in range(hidden_dim // vlen):
# Normalize the activation vector
generated_code += f"V_MUL_VF gp{act_addr}, gp{act_addr}, f2, 0 \n"

# Move to next vector
generated_code += f"S_ADDI_INT gp{act_addr}, gp{act_addr}, {vlen * batch_size} \n"
else:
generated_code += f"C_LOOP_START gp{loop_addr}, {hidden_dim // vlen} \n"
generated_code += f"V_MUL_VF gp{act_addr}, gp{act_addr}, f2, 0 \n"

# Move to next vector
generated_code += f"S_ADDI_INT gp{act_addr}, gp{act_addr}, {vlen * batch_size} \n"
generated_code += f"C_LOOP_END gp{loop_addr} \n"

# Reset accumulator for next batch
generated_code += "S_ADD_FP f2, f0, f0 \n"
Expand All @@ -78,13 +95,17 @@ def layer_norm_asm(
vlen: int,
batch_size: int,
hidden_dim: int,
unroll: bool = True,
) -> str:
"""
Generate assembly code for layer normalization.
"""
act_addr = alive_registers[0]
scratchpad_addr = alive_registers[1]
stats_addr = alive_registers[2]
# Rolled path uses the spare 4th register (already allocated by normalize()) as the
# C_LOOP counter. Only accessed when rolled, so unrolled callers may pass 3 registers.
loop_addr = alive_registers[3] if not unroll else None

generated_code = "; Layer Norm generation \n"
generated_code += _load_large_int(scratchpad_addr, scratchpad_base_address)
Expand All @@ -102,16 +123,24 @@ def layer_norm_asm(
generated_code += _load_large_int(stats_addr, activation_base_address + vlen * batch)

# First loop: compute sum(x) and sum(x^2) using stats_addr
for i in range(hidden_dim // vlen):
# sum(x)
if unroll:
for i in range(hidden_dim // vlen):
# sum(x)
generated_code += f"V_RED_SUM f2, gp{stats_addr} \n"

# sum(x^2)
generated_code += f"V_MUL_VV gp{scratchpad_addr}, gp{stats_addr}, gp{stats_addr}, 0 \n"
generated_code += f"V_RED_SUM f3, gp{scratchpad_addr} \n"

# Move stats pointer to next vector
generated_code += f"S_ADDI_INT gp{stats_addr}, gp{stats_addr}, {vlen * batch_size} \n"
else:
generated_code += f"C_LOOP_START gp{loop_addr}, {hidden_dim // vlen} \n"
generated_code += f"V_RED_SUM f2, gp{stats_addr} \n"

# sum(x^2)
generated_code += f"V_MUL_VV gp{scratchpad_addr}, gp{stats_addr}, gp{stats_addr}, 0 \n"
generated_code += f"V_RED_SUM f3, gp{scratchpad_addr} \n"

# Move stats pointer to next vector
generated_code += f"S_ADDI_INT gp{stats_addr}, gp{stats_addr}, {vlen * batch_size} \n"
generated_code += f"C_LOOP_END gp{loop_addr} \n"

# f2 = sum(x) * (1/hidden_dim) = mean(x)
generated_code += "S_MUL_FP f2, f2, f4 \n"
Expand All @@ -135,15 +164,22 @@ def layer_norm_asm(
generated_code += "S_RECI_FP f5, f5 \n"

# Second loop: normalize using act_addr (still at batch start)
for i in range(hidden_dim // vlen):
# normalized = (x - mean) * (1/std)
# Store (x - mean) in scratchpad first
if unroll:
for i in range(hidden_dim // vlen):
# normalized = (x - mean) * (1/std)
# Store (x - mean) in scratchpad first
generated_code += f"V_SUB_VF gp{scratchpad_addr}, gp{act_addr}, f2, 0, 0 \n"
# Then multiply by 1/std and write back to activation
generated_code += f"V_MUL_VF gp{act_addr}, gp{scratchpad_addr}, f5, 0 \n"

# Move to next vector
generated_code += f"S_ADDI_INT gp{act_addr}, gp{act_addr}, {vlen * batch_size} \n"
else:
generated_code += f"C_LOOP_START gp{loop_addr}, {hidden_dim // vlen} \n"
generated_code += f"V_SUB_VF gp{scratchpad_addr}, gp{act_addr}, f2, 0, 0 \n"
# Then multiply by 1/std and write back to activation
generated_code += f"V_MUL_VF gp{act_addr}, gp{scratchpad_addr}, f5, 0 \n"

# Move to next vector
generated_code += f"S_ADDI_INT gp{act_addr}, gp{act_addr}, {vlen * batch_size} \n"
generated_code += f"C_LOOP_END gp{loop_addr} \n"

# Reset accumulators for next batch
generated_code += "S_ADD_FP f2, f0, f0 \n"
Expand Down
45 changes: 33 additions & 12 deletions asm_templates/rope_asm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ def rope_asm(
vlen: int,
seq_len: int,
head_dim: int,
unroll: bool = True,
) -> str:
"""
Generate assembly for Rotary Position Embedding (RoPE) applied in-place:
Expand All @@ -22,7 +23,8 @@ def rope_asm(
Chunk j, position i is at: base + j * seq_len * vlen + i * vlen

Args:
alive_registers: 5 GP registers [x_addr, xrot_addr, cos_addr, sin_addr, scratch_addr]
alive_registers: GP registers [x_addr, xrot_addr, cos_addr, sin_addr, scratch_addr]
(5 for the unrolled path; a 6th [loop_addr] is required for the rolled path)
x_base_address: VRAM base of x (shape: seq_len × head_dim)
x_rot_base_address: VRAM base of rotate_half(x), preloaded from HBM
cos_base_address: VRAM base of cos values (seq_len × head_dim)
Expand All @@ -45,16 +47,35 @@ def rope_asm(
lines = ["; RoPE: x = x * cos + rotate_half(x) * sin (in-place)"]
lines.extend(_load_large_int_list(scratch_addr, scratchpad_base_address))

for j in range(num_chunks):
chunk_base = j * seq_len * vlen
for i in range(seq_len):
addr = chunk_base + i * vlen
lines.extend(_load_large_int_list(x_addr, x_base_address + addr))
lines.extend(_load_large_int_list(xrot_addr, x_rot_base_address + addr))
lines.extend(_load_large_int_list(cos_addr, cos_base_address + addr))
lines.extend(_load_large_int_list(sin_addr, sin_base_address + addr))
lines.append(f"V_MUL_VV gp{scratch_addr}, gp{xrot_addr}, gp{sin_addr}, 0 ")
lines.append(f"V_MUL_VV gp{x_addr}, gp{x_addr}, gp{cos_addr}, 0 ")
lines.append(f"V_ADD_VV gp{x_addr}, gp{x_addr}, gp{scratch_addr}, 0 ")
if unroll:
for j in range(num_chunks):
chunk_base = j * seq_len * vlen
for i in range(seq_len):
addr = chunk_base + i * vlen
lines.extend(_load_large_int_list(x_addr, x_base_address + addr))
lines.extend(_load_large_int_list(xrot_addr, x_rot_base_address + addr))
lines.extend(_load_large_int_list(cos_addr, cos_base_address + addr))
lines.extend(_load_large_int_list(sin_addr, sin_base_address + addr))
lines.append(f"V_MUL_VV gp{scratch_addr}, gp{xrot_addr}, gp{sin_addr}, 0 ")
lines.append(f"V_MUL_VV gp{x_addr}, gp{x_addr}, gp{cos_addr}, 0 ")
lines.append(f"V_ADD_VV gp{x_addr}, gp{x_addr}, gp{scratch_addr}, 0 ")
else:
# Rolled: addr = (j*seq_len + i)*vlen is a single linear progression, so the
# whole double loop collapses to ONE hardware loop of count num_chunks*seq_len
# with a +vlen stride on each of the four operand pointers (scratch is fixed).
loop_addr = alive_registers[5]
lines.extend(_load_large_int_list(x_addr, x_base_address))
lines.extend(_load_large_int_list(xrot_addr, x_rot_base_address))
lines.extend(_load_large_int_list(cos_addr, cos_base_address))
lines.extend(_load_large_int_list(sin_addr, sin_base_address))
lines.append(f"C_LOOP_START gp{loop_addr}, {num_chunks * seq_len}")
lines.append(f"V_MUL_VV gp{scratch_addr}, gp{xrot_addr}, gp{sin_addr}, 0 ")
lines.append(f"V_MUL_VV gp{x_addr}, gp{x_addr}, gp{cos_addr}, 0 ")
lines.append(f"V_ADD_VV gp{x_addr}, gp{x_addr}, gp{scratch_addr}, 0 ")
lines.append(f"S_ADDI_INT gp{x_addr}, gp{x_addr}, {vlen} ")
lines.append(f"S_ADDI_INT gp{xrot_addr}, gp{xrot_addr}, {vlen} ")
lines.append(f"S_ADDI_INT gp{cos_addr}, gp{cos_addr}, {vlen} ")
lines.append(f"S_ADDI_INT gp{sin_addr}, gp{sin_addr}, {vlen} ")
lines.append(f"C_LOOP_END gp{loop_addr}")

return "\n".join(lines) + "\n"
119 changes: 57 additions & 62 deletions assembler/assembly_to_binary.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,53 @@
from .parser import load_isa_definitions, parse_asm_file


# Opcode groups for binary encoding. Module-level frozensets so they are not rebuilt
# (and scanned linearly) on every _convert_to_binary call — that runs once per emitted
# instruction (millions of times for large programs). Membership is O(1) and identical
# to the previous `opcode in [ ... ]` list literals.
_RMASK_VECTOR_OPS = frozenset(
{"V_ADD_VV", "V_ADD_VF", "V_MUL_VV", "V_SUB_VV", "V_MUL_VF", "V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"}
)
_IMM_RS1_RD_OPS = frozenset(
{
"S_ADDI_INT",
"M_MM_WO",
"S_LD_FP",
"S_ST_FP",
"S_LD_INT",
"S_ST_INT",
"S_MAP_V_FP",
"V_RED_MAX",
"V_RECI_V",
"V_EXP_V",
}
)
_IMM_RD_OPS = frozenset({"S_LUI_INT", "M_MV_WO", "M_BMM_WO", "M_BMV_WO"})
_RS1_RD_OPS = frozenset({"S_MV_FP", "S_RECI_FP", "S_EXP_FP", "S_SQRT_FP", "V_EXP_V", "V_RED_SUM"})
_RD_ONLY_OPS = frozenset({"C_SET_SCALE_REG", "C_SET_STRIDE_REG", "C_SET_V_MASK_REG", "C_LOOP_END"})
_FUNCT_RSTRIDE_OPS = frozenset({"H_PREFETCH_M", "H_PREFETCH_V", "H_STORE_V", "V_SUB_VF"})
_RS2_RS1_RD_OPS = frozenset(
{
"S_ADD_INT",
"S_ADD_FP",
"S_SUB_INT",
"S_SUB_FP",
"S_MUL_INT",
"S_MUL_FP",
"S_MAX_FP",
"M_MM",
"M_MV",
"M_BMM",
"M_BMV",
"M_TMM",
"M_TMV",
"M_BTMM",
"M_BTMV",
"C_SET_ADDR_REG",
}
)


class AssemblyToBinary:
def __init__(self, isa_definition_file: str, config_file: str):
"""
Expand Down Expand Up @@ -40,47 +87,25 @@ def _convert_to_binary(self, instruction):
binary_instruction = 0
ow = self.operands_width
opw = self.opcode_width
vector_ops_with_rmask = {
"V_ADD_VV",
"V_ADD_VF",
"V_MUL_VV",
"V_SUB_VV",
"V_MUL_VF",
"V_EXP_V",
"V_RECI_V",
"V_RED_SUM",
"V_RED_MAX",
}

if instruction.opcode in vector_ops_with_rmask and rmask is None:
if instruction.opcode in _RMASK_VECTOR_OPS and rmask is None:
# Treat omitted rmask deterministically as "mask disabled" instead of crashing on None << ...
rmask = 0

if instruction.opcode in [
"S_ADDI_INT",
"M_MM_WO",
"S_LD_FP",
"S_ST_FP",
"S_LD_INT",
"S_ST_INT",
"S_MAP_V_FP",
"V_RED_MAX",
"V_RECI_V",
"V_EXP_V",
]:
if instruction.opcode in _IMM_RS1_RD_OPS:
binary_instruction = (imm << (opw + 2 * ow)) + (rs1 << (opw + ow)) + (rd << opw) + opcode
elif instruction.opcode in ["S_LUI_INT", "M_MV_WO", "M_BMM_WO", "M_BMV_WO"]:
elif instruction.opcode in _IMM_RD_OPS:
binary_instruction = (imm << (opw + ow)) + (rd << opw) + opcode
elif instruction.opcode in ["S_MV_FP", "S_RECI_FP", "S_EXP_FP", "S_SQRT_FP", "V_EXP_V", "V_RED_SUM"]:
elif instruction.opcode in _RS1_RD_OPS:
binary_instruction = (rs1 << (opw + ow)) + (rd << opw) + opcode
elif instruction.opcode in ["C_BREAK"]:
elif instruction.opcode == "C_BREAK":
binary_instruction = opcode
elif instruction.opcode in ["C_SET_SCALE_REG", "C_SET_STRIDE_REG", "C_SET_V_MASK_REG", "C_LOOP_END"]:
elif instruction.opcode in _RD_ONLY_OPS:
binary_instruction = (rd << opw) + opcode
elif instruction.opcode in ["C_LOOP_START"]:
elif instruction.opcode == "C_LOOP_START":
# C_LOOP_START rd, imm - uses 22-bit immediate like S_LUI_INT
binary_instruction = (imm << (opw + ow)) + (rd << opw) + opcode
elif instruction.opcode in ["H_PREFETCH_M", "H_PREFETCH_V", "H_STORE_V", "V_SUB_VF"]:
elif instruction.opcode in _FUNCT_RSTRIDE_OPS:
binary_instruction = (
(funct1 << (opw + 4 * ow))
+ (rstride << (opw + 3 * ow))
Expand All @@ -89,41 +114,11 @@ def _convert_to_binary(self, instruction):
+ (rd << opw)
+ opcode
)
elif instruction.opcode in [
"V_ADD_VV",
"V_ADD_VF",
"V_MUL_VV",
"V_SUB_VV",
"V_MUL_VF",
"V_EXP_V",
"V_RECI_V",
"V_RED_SUM",
"V_RED_MAX",
]:
elif instruction.opcode in _RMASK_VECTOR_OPS:
binary_instruction = (
(rmask << (opw + 3 * ow)) + (rs2 << (opw + 2 * ow)) + (rs1 << (opw + ow)) + (rd << opw) + opcode
)
elif instruction.opcode in [
# Scalar arithmetic (rd, rs1, rs2) — no rmask
"S_ADD_INT",
"S_ADD_FP",
"S_SUB_INT",
"S_SUB_FP",
"S_MUL_INT",
"S_MUL_FP",
"S_MAX_FP",
# Matrix ops without write-out (rd, rs1, rs2)
"M_MM",
"M_MV",
"M_BMM",
"M_BMV",
"M_TMM",
"M_TMV",
"M_BTMM",
"M_BTMV",
# CSR: addr reg destination + 2 GP sources (a{N}, gp{X}, gp{Y} → rd, rs1, rs2)
"C_SET_ADDR_REG",
]:
elif instruction.opcode in _RS2_RS1_RD_OPS:
binary_instruction = (rs2 << (opw + 2 * ow)) + (rs1 << (opw + ow)) + (rd << opw) + opcode
else:
binary_instruction = (rs2 << (opw + 2 * ow)) + (rs1 << (opw + ow)) + (rd << opw) + opcode
Expand Down
Loading
Loading