diff --git a/asm_templates/normalization_asm.py b/asm_templates/normalization_asm.py index 01addff..6b0d7e7 100644 --- a/asm_templates/normalization_asm.py +++ b/asm_templates/normalization_asm.py @@ -10,6 +10,7 @@ def rms_norm_asm( vlen: int, batch_size: int, hidden_dim: int, + unroll: bool = True, ) -> str: """ Generate assembly code for RMS normalization. @@ -17,6 +18,9 @@ def rms_norm_asm( 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) @@ -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" @@ -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" @@ -78,6 +95,7 @@ def layer_norm_asm( vlen: int, batch_size: int, hidden_dim: int, + unroll: bool = True, ) -> str: """ Generate assembly code for layer normalization. @@ -85,6 +103,9 @@ def layer_norm_asm( 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) @@ -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" @@ -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" diff --git a/asm_templates/rope_asm.py b/asm_templates/rope_asm.py index 23c3d1d..7f81b9b 100644 --- a/asm_templates/rope_asm.py +++ b/asm_templates/rope_asm.py @@ -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: @@ -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) @@ -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" diff --git a/assembler/assembly_to_binary.py b/assembler/assembly_to_binary.py index 01336fd..4ea6746 100644 --- a/assembler/assembly_to_binary.py +++ b/assembler/assembly_to_binary.py @@ -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): """ @@ -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)) @@ -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 diff --git a/assembler/parser.py b/assembler/parser.py index 69aeae4..d9dd875 100644 --- a/assembler/parser.py +++ b/assembler/parser.py @@ -92,6 +92,32 @@ def __repr__(self): return f"Instruction(opcode='{self.opcode}', rd='{self.rd}', rs1='{self.rs1}', rs2='{self.rs2}', rstride = '{self.rstride}', funct1={self.funct1}, funct2={self.funct2}, imm={self.imm}, rflag={self.rflag})" +_REG_PREFIXES = ("gp", "f", "a") +# Hoisted to module scope: these were previously re-created for every line of the +# .asm (millions of times for large programs), which dominated sim_env re-parse time. +vector_masked_unary_or_reduction_ops = frozenset({"V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"}) +vector_masked_binary_ops = frozenset({"V_ADD_VV", "V_ADD_VF", "V_MUL_VV", "V_SUB_VV", "V_MUL_VF"}) + + +def _parse_operand(operand): + """Parse a register (gp/f/a prefix, decimal index) or integer operand; None if neither. + + Operands are already whitespace-stripped by the caller (the line-132 split), so this + does no stripping. Hoisted out of the per-line loop, where it was redefined every line. + """ + if operand.endswith(";"): + operand = operand[:-1] + if operand.startswith("gp"): + return int(operand[2:]) # decimal, not hex + elif operand.startswith(("f", "a")): + return int(operand[1:]) # decimal, not hex + else: + try: + return int(operand) + except ValueError: + return None + + def parse_asm_file(file_path: str) -> list[Instruction]: """ Parse an ASM file into a list of Instruction objects. @@ -109,28 +135,28 @@ def parse_asm_file(file_path: str) -> list[Instruction]: """ instructions = [] - vector_masked_unary_or_reduction_ops = {"V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"} - vector_masked_binary_ops = {"V_ADD_VV", "V_ADD_VF", "V_MUL_VV", "V_SUB_VV", "V_MUL_VF"} - with open(file_path) as file: for line in file: - # Remove comments and strip whitespace - # Handle both // and ; style comments - if line.startswith("//") or line.strip().startswith(";"): - continue - line = line.split("//")[0] # Remove // comments - line = line.split(";")[0] # Remove ; comments + # Strip once, then skip blanks and whole-line comments with cheap char checks + # (most lines are neither). Comments are // or ; style. line = line.strip() - if not line: + if not line or line[0] == ";" or line.startswith("//"): continue - - # Split the opcode and operands - parts = line.split() - if len(parts) < 2 or ";" in parts[0]: + # Remove inline // and ; comments only when present. + c = line.find("//") + if c != -1: + line = line[:c] + c = line.find(";") + if c != -1: + line = line[:c] + + # Split opcode off the operand string on the first whitespace run; this avoids + # the old split()+" ".join() round-trip. Lines with no operands are skipped. + head = line.split(None, 1) + if len(head) < 2: continue # Invalid line - opcode = parts[0] - operands = [part.strip() for part in " ".join(parts[1:]).split(",")] - # print(f"Parsing instruction: {line}", "operand length:", len(operands), "operands:", operands) + opcode, rest = head + operands = [part.strip() for part in rest.split(",")] # Decode based on number of operands, case-structure by length rd = None @@ -141,34 +167,17 @@ def parse_asm_file(file_path: str) -> list[Instruction]: funct2 = None imm = None - # Helper to parse a register or int operand - def parse_reg_or_int(operand): - operand = operand.strip() - if operand.endswith(";"): - operand = operand[:-1] - if operand.startswith("gp"): - return int(operand[2:]) # decimal, not hex - elif operand.startswith("f"): - return int(operand[1:]) # decimal, not hex - elif operand.startswith("a"): - return int(operand[1:]) # decimal, not hex - else: - try: - return int(operand) - except ValueError: - return None - if len(operands) == 1: operand_0 = operands[0] - rd = parse_reg_or_int(operand_0) + rd = _parse_operand(operand_0) elif len(operands) == 2: operand_0 = operands[0] operand_1 = operands[1] - rd = parse_reg_or_int(operand_0) + rd = _parse_operand(operand_0) # rs1 is a register, imm is a number # Heuristics: if it looks like a reg, it's rs1; else, it's imm - if operand_1.strip().startswith(("gp", "f", "a")): - rs1 = parse_reg_or_int(operand_1) + if operand_1.startswith(("gp", "f", "a")): + rs1 = _parse_operand(operand_1) else: try: imm = int(operand_1) @@ -176,18 +185,18 @@ def parse_reg_or_int(operand): imm = None elif len(operands) == 3: operand_0, operand_1, operand_2 = operands - rd = parse_reg_or_int(operand_0) + rd = _parse_operand(operand_0) # If looks like register, rs1; else, imm - if operand_1.strip().startswith(("gp", "f", "a")): - rs1 = parse_reg_or_int(operand_1) + if operand_1.startswith(("gp", "f", "a")): + rs1 = _parse_operand(operand_1) else: try: imm = int(operand_1) except ValueError: imm = None # If it looks like register, rs2; else, imm (overwrites imm if rs1 not present) - if operand_2.strip().startswith(("gp", "f", "a")): - rs2 = parse_reg_or_int(operand_2) + if operand_2.startswith(("gp", "f", "a")): + rs2 = _parse_operand(operand_2) else: try: imm = int(operand_2) @@ -201,16 +210,16 @@ def parse_reg_or_int(operand): rstride = 0 elif len(operands) == 4: operand_0, operand_1, operand_2, operand_3 = operands - rd = parse_reg_or_int(operand_0) - if operand_1.strip().startswith(("gp", "f", "a")): - rs1 = parse_reg_or_int(operand_1) + rd = _parse_operand(operand_0) + if operand_1.startswith(("gp", "f", "a")): + rs1 = _parse_operand(operand_1) else: try: imm = int(operand_1) except ValueError: imm = None - if operand_2.strip().startswith(("gp", "f", "a")): - rs2 = parse_reg_or_int(operand_2) + if operand_2.startswith(("gp", "f", "a")): + rs2 = _parse_operand(operand_2) else: try: imm = int(operand_2) @@ -223,16 +232,16 @@ def parse_reg_or_int(operand): rstride = None elif len(operands) == 5: operand_0, operand_1, operand_2, operand_3, operand_4 = operands - rd = parse_reg_or_int(operand_0) - if operand_1.strip().startswith(("gp", "f", "a")): - rs1 = parse_reg_or_int(operand_1) + rd = _parse_operand(operand_0) + if operand_1.startswith(("gp", "f", "a")): + rs1 = _parse_operand(operand_1) else: try: imm = int(operand_1) except ValueError: imm = None - if operand_2.strip().startswith(("gp", "f", "a")): - rs2 = parse_reg_or_int(operand_2) + if operand_2.startswith(("gp", "f", "a")): + rs2 = _parse_operand(operand_2) else: try: imm = int(operand_2) @@ -251,16 +260,16 @@ def parse_reg_or_int(operand): funct1 = funct1_raw # fallback, if not int, keep as string elif len(operands) == 6: operand_0, operand_1, operand_2, operand_3, operand_4, operand_5 = operands - rd = parse_reg_or_int(operand_0) - if operand_1.strip().startswith(("gp", "f", "a")): - rs1 = parse_reg_or_int(operand_1) + rd = _parse_operand(operand_0) + if operand_1.startswith(("gp", "f", "a")): + rs1 = _parse_operand(operand_1) else: try: imm = int(operand_1) except ValueError: imm = None - if operand_2.strip().startswith(("gp", "f", "a")): - rs2 = parse_reg_or_int(operand_2) + if operand_2.startswith(("gp", "f", "a")): + rs2 = _parse_operand(operand_2) else: try: imm = int(operand_2) diff --git a/aten/plena/compiler.py b/aten/plena/compiler.py index 80ee9c5..e0ab4cf 100644 --- a/aten/plena/compiler.py +++ b/aten/plena/compiler.py @@ -90,9 +90,9 @@ def __init__( PLENA_SETTINGS_TOML / plena_settings.toml. unroll_loops: If True, unroll sub-projection and attention helper loops at ASM-gen time to eliminate C_LOOP_START/END overhead. - Overridden by the ATEN_UNROLL env var ("1"=True, "0"=False). + Overridden by the ATEN_OPS_UNROLL env var ("1"=True, "0"=False). """ - _env_unroll = os.environ.get("ATEN_UNROLL", "") + _env_unroll = os.environ.get("ATEN_OPS_UNROLL", "") if _env_unroll == "1": unroll_loops = True elif _env_unroll == "0": diff --git a/aten/plena/isa_compiler.py b/aten/plena/isa_compiler.py index 2c22da9..3a6c8d8 100644 --- a/aten/plena/isa_compiler.py +++ b/aten/plena/isa_compiler.py @@ -284,6 +284,7 @@ def normalize( vlen=vlen, batch_size=batch_size, hidden_dim=hidden_dim, + unroll=self._unroll, ) else: isa_code += layer_norm_asm( @@ -295,6 +296,7 @@ def normalize( vlen=vlen, batch_size=batch_size, hidden_dim=hidden_dim, + unroll=self._unroll, ) return self._emit(isa_code) @@ -332,7 +334,9 @@ def rope( if head_dim % vlen != 0: raise ValueError(f"head_dim ({head_dim}) must be divisible by vlen ({vlen}) for rope") - gp_regs = self.register_allocator.allocate_gp(5) + # Rolled RoPE needs one extra GP register (the C_LOOP counter); the unrolled + # path keeps its original 5-register allocation so its output is byte-identical. + gp_regs = self.register_allocator.allocate_gp(5 if self._unroll else 6) scratch_name = f"__rope_scratch__{x_name}__{len(self.generated_code)}" scratch_addr = self.vram_allocator.allocate(vlen, name=scratch_name) @@ -348,6 +352,7 @@ def rope( vlen=vlen, seq_len=seq_len, head_dim=head_dim, + unroll=self._unroll, ) return self._emit(isa_code) finally: diff --git a/aten/plena/isa_emit.py b/aten/plena/isa_emit.py index adad6be..d15afe6 100644 --- a/aten/plena/isa_emit.py +++ b/aten/plena/isa_emit.py @@ -17,6 +17,24 @@ def _reg(self) -> RegisterAllocator: """Shorthand for self.register_allocator (used by FPVar ISA helpers).""" return self.register_allocator + # ------------------------------------------------------------------ + # Generated ISA buffer + # + # Backed by a list of rendered chunks rather than one growing string: + # ``self.generated_code += rendered`` per instruction is O(n) per call + # (it copies the whole buffer), i.e. O(n^2) overall, which runs away when + # the instruction count is large (e.g. mlen=16 vision attention tiles into + # 4 col-blocks). Appending to a list is amortised O(1); the getter joins + # on read, producing a byte-identical string. Callers still see a ``str``. + # ------------------------------------------------------------------ + @property + def generated_code(self) -> str: + return "".join(getattr(self, "_code_chunks", ())) + + @generated_code.setter + def generated_code(self, value: str) -> None: + self._code_chunks = [value] if value else [] + @property def _unroll(self) -> bool: """Shorthand for self.unroll_loops.""" @@ -25,7 +43,7 @@ def _unroll(self) -> bool: def _emit(self, isa_code: AsmInput) -> str: """Append ISA text to the output buffer and return it.""" rendered = render_asm(isa_code) - self.generated_code += rendered + self._code_chunks.append(rendered) return rendered def emit(self, isa_code: AsmInput) -> str: