From efaf445794f5b805163f1fcec94cee66e02867d5 Mon Sep 17 00:00:00 2001 From: Seven Gao <799889633@qq.com> Date: Mon, 14 Sep 2026 01:48:44 +0800 Subject: [PATCH 1/6] feat(topic16): implement real NN op lowering and fix invalid LLVM IR --- scratchv/backend/llvm_codegen.py | 1442 +++++++++++++++++++++------- tests/test_llvm_codegen_topic16.py | 689 +++++++++++++ 2 files changed, 1810 insertions(+), 321 deletions(-) create mode 100644 tests/test_llvm_codegen_topic16.py diff --git a/scratchv/backend/llvm_codegen.py b/scratchv/backend/llvm_codegen.py index be90352..1b97919 100644 --- a/scratchv/backend/llvm_codegen.py +++ b/scratchv/backend/llvm_codegen.py @@ -2,10 +2,18 @@ Produces human-readable .ll files suitable for ``llc``, ``opt``, or ``lli``. No external dependencies beyond Python — the output is standard LLVM IR. + +The generator is single-pass and memory-based: control flow uses +``alloca``/``load``/``store`` plus explicit branches, and every SSA name and +label is uniquified by :class:`SSANamer` so the emitted module always +assembles with ``llvm-as``. """ from __future__ import annotations +import struct +from dataclasses import dataclass + from scratchv.ir.types import ( OpCode, DataType, Value, Instruction, BasicBlock, Function, Program, ) @@ -23,21 +31,99 @@ _LLVM_I32 = "i32" _LLVM_I64 = "i64" +_FLOAT_TYPES = ("float", "double") +_INT_TYPES = ("i32", "i64") + +_CMP_PRED = { + "==": ("oeq", "eq"), + "!=": ("one", "ne"), + "<": ("olt", "slt"), + "<=": ("ole", "sle"), + ">": ("ogt", "sgt"), + ">=": ("oge", "sge"), +} + +_PROLOGUE_MARKER = ";__SCRATCHV_PROLOGUE__" + + +class LLVMCodegenError(Exception): + """Raised for unlowerable IR (missing operand, unmatched endfor, ...).""" + + +class SSANamer: + """Generates unique SSA register and label names per function.""" + + def __init__(self) -> None: + self._reg_n = 0 + self._label_n = 0 + self._defs: set[str] = set() + + @staticmethod + def sanitize(name: str) -> str: + """Map an IR name to a legal LLVM identifier body.""" + s = "".join( + ch if (ch.isalnum() or ch in "_.") else "_" for ch in str(name) + ) + if not s or s[0].isdigit(): + s = "v_" + s + return s + + def fresh(self, hint: str = "r") -> str: + """Return a fresh register name such as ``%hint_3``.""" + self._reg_n += 1 + return f"%{self.sanitize(hint)}_{self._reg_n}" + + def fresh_label(self, hint: str = "bb") -> str: + """Return a fresh label name such as ``hint_7`` (no ``%``).""" + self._label_n += 1 + return f"{self.sanitize(hint)}_{self._label_n}" + + def register_definition(self, reg: str) -> None: + if reg in self._defs: + raise LLVMCodegenError(f"duplicate SSA definition: {reg}") + self._defs.add(reg) + + def registered(self, reg: str) -> bool: + return reg in self._defs + + +@dataclass +class LoopContext: + """State of one structured ``for`` loop (canonical form).""" + + ir_name: str | None + ptr: str + value: str + header: str + body: str + exit: str + limit: int + step: int = 1 + class LLVMCodegen: """Translate ScratchV IR Program to LLVM IR text (.ll).""" - def __init__(self, program: Program): + def __init__(self, program: Program, target_triple: str | None = None): self.program = program + self.target_triple = target_triple self._lines: list[str] = [] - self._indent = 0 - # IR value name -> LLVM register + # IR value name -> current LLVM reference (register or pointer) self._named_values: dict[str, str] = {} + # IR value name -> memory slot (memory-based scalar variables) + self._slots: dict[str, str] = {} + # LLVM reference -> LLVM type string + self._ref_types: dict[str, str] = {} # function name -> return type self._func_type: dict[str, str] = {} - self._block_counter = 0 - self._loop_context: dict | None = None + # allocas hoisted into the function entry block + self._prologue: list[str] = [] + self._defined_labels: set[str] = set() + self._terminated: bool = True + self._loop_stack: list[LoopContext] = [] self._current_func: str | None = None + self._ptr_names: set[str] = set() + self._namer = SSANamer() # ------------------------------------------------------------------ # Public API @@ -48,10 +134,10 @@ def emit(self) -> str: self._lines = [] self._p("; LLVM IR generated by ScratchV") self._p('; ModuleID = "scratchv_module"') - self._p("target triple = \"riscv64-unknown-elf\"") + if self.target_triple is not None: + self._p(f'target triple = "{self.target_triple}"') self._p("") - # Declare external helpers self._emit_externals() for func in self.program.functions: @@ -76,17 +162,16 @@ def _emit_externals(self) -> None: self._p("declare void @print_f32(float) nounwind") self._p("") + def _p(self, line: str = "") -> None: + self._lines.append(line) + # ------------------------------------------------------------------ # Functions # ------------------------------------------------------------------ @staticmethod def _infer_function_params(func: Function) -> None: - """Add undefined external value references as function params. - - Handles DSL-parsed programs where free variables (e.g. 'a', 'b' - in "y = add(a, b)") are referenced but not declared as params. - """ + """Add undefined external value references as function params.""" defined: set[str] = {p.name for p in func.params} referenced: set[str] = set() @@ -101,169 +186,486 @@ def _infer_function_params(func: Function) -> None: existing_param_names = {p.name for p in func.params} for name in referenced - defined: if name not in existing_param_names: - # Find the value from the program's globals or create a new one - from scratchv.ir.types import Value, DataType - val = Value(name=name, dtype=DataType.FLOAT32) - func.params.append(val) + func.params.append(Value(name=name, dtype=DataType.FLOAT32)) def _emit_function(self, func: Function) -> None: self._current_func = func.name - self._named_values.clear() - self._block_counter = 0 + self._named_values = {} + self._slots = {} + self._ref_types = {} + self._prologue = [] + self._terminated = False + self._defined_labels = set() + self._loop_stack = [] + self._namer = SSANamer() + self._ptr_names = { + instr.dest.name + for block in func.blocks + for instr in block.instructions + if instr.opcode is OpCode.ALLOCA and instr.dest is not None + } - # Auto-detect undefined external variable references as params self._infer_function_params(func) - # Build param list - params = [] + used_names: set[str] = set() + params: list[str] = [] for p in func.params: - llvm_ty = _llvm_type(p.dtype) - params.append(f"{llvm_ty} %{p.name}") + ty = self._value_type(p) + safe = self._unique_param_name(p.name, used_names) + params.append(f"{ty} %{safe}") + self._bind(p.name, f"%{safe}", ty) - # Determine return type - ret_ty = "void" - if func.returns: - ret_ty = _llvm_type(func.returns[0].dtype) - else: - # Scan blocks for return instructions to infer return type - for block in func.blocks: - for instr in block.instructions: - if instr.opcode == OpCode.RETURN and instr.operands: - ret_ty = _llvm_type(instr.operands[0].dtype) - break - if ret_ty != "void": - break - - self._p(f"define {ret_ty} @{func.name}({', '.join(params)}) {{") - - # Map function params to named values - for p in func.params: - self._named_values[p.name] = f"%{p.name}" + ret_ty = self._function_return_type(func) + params_str = ", ".join(params) - self._indent = 1 + self._p(f"define {ret_ty} @{func.name}({params_str}) {{") + define_idx = len(self._lines) - 1 + marker = len(self._lines) + self._lines.append(_PROLOGUE_MARKER) - # Emit each basic block for block in func.blocks: self._emit_block(block) - self._indent = 0 + if not self._terminated: + self._p(" unreachable") + self._terminated = True + + actual_ret = self._first_ret_type(define_idx) + if actual_ret is not None and actual_ret != ret_ty: + self._lines[define_idx] = ( + f"define {actual_ret} @{func.name}({params_str}) {{" + ) + + self._lines[marker:marker + 1] = self._prologue self._p("}") self._p("") + def _first_ret_type(self, start_idx: int) -> str | None: + for line in self._lines[start_idx + 1:]: + stripped = line.strip() + if stripped.startswith("ret "): + parts = stripped.split() + if len(parts) >= 2: + return parts[1] + return "void" + return None + + @staticmethod + def _unique_param_name(name: str, used: set[str]) -> str: + safe = SSANamer.sanitize(name) + candidate = safe + n = 1 + while candidate in used: + n += 1 + candidate = f"{safe}_{n}" + used.add(candidate) + return candidate + + def _function_return_type(self, func: Function) -> str: + if func.returns: + return self._value_type(func.returns[0]) + for block in func.blocks: + for instr in block.instructions: + if instr.opcode is OpCode.RETURN and instr.operands: + return self._value_type(instr.operands[0]) + return "void" + # ------------------------------------------------------------------ - # Basic blocks + # Basic blocks and control-flow state machine # ------------------------------------------------------------------ def _emit_block(self, block: BasicBlock) -> None: - if block.name == "entry" and self._is_first_block(): - self._p(f"; --- {block.name} ---") - else: - self._p(f"{block.name}:") - self._p(f"; --- {block.name} ---") - + self._start_block(block.name) + self._p(f"; --- {block.name} ---") for instr in block.instructions: self._emit_instruction(instr) - def _is_first_block(self) -> bool: - """Check if we're in the first block.""" - return True + def _start_block(self, label: str) -> None: + """Close the current block (fallthrough) and open ``label``.""" + if not label: + label = self._namer.fresh_label("bb") + if label in self._defined_labels: + label = self._namer.fresh_label(label) + while label in self._defined_labels: + label = self._namer.fresh_label(label) + if not self._terminated: + self._p(f" br label %{label}") + self._p(f"{label}:") + self._defined_labels.add(label) + self._terminated = False + + def _terminate(self, line: str) -> None: + if self._terminated: + raise LLVMCodegenError("terminator emitted twice in one block") + self._p(f" {line}") + self._terminated = True # ------------------------------------------------------------------ # Instructions # ------------------------------------------------------------------ def _emit_instruction(self, instr: Instruction) -> None: + if self._terminated: + self._start_block(self._namer.fresh_label("dead")) handler = getattr(self, f"_emit_{instr.opcode.value}", None) if handler is None: - ops = ' '.join(str(v.name) for v in instr.operands) - self._p(f" ; UNSUPPORTED: {instr.opcode.value} {ops}") + handler = self._emit_unsupported + handler(instr) + self._flush_dest(instr) + + def _flush_dest(self, instr: Instruction) -> None: + """Store a just-defined scalar value into its memory slot.""" + if instr.dest is None: + return + name = instr.dest.name + reg = self._named_values.get(name) + slot = self._slots.get(name) + if reg is None or slot is None or self._terminated: + return + ty = self._ref_types.get(reg) + if ty is None or ty.endswith("*"): + return + self._p(f" store {ty} {reg}, {ty}* {slot}") + + def _emit_unsupported(self, instr: Instruction) -> None: + self._p(f" ; fallback: {instr.opcode.value} not lowered") + if instr.dest is None: + return + base = _llvm_type(instr.dest.dtype) + if self._is_pointer_value(instr.dest): + count = 1 + for dim in instr.dest.shape: + count *= max(1, int(dim)) + self._dest_buffer(instr, count, base) + return + dst = self._dest(instr) + if base in _FLOAT_TYPES: + self._p(f" {dst} = fadd {base} 0.0, 0.0") else: - handler(instr) + self._p(f" {dst} = add {base} 0, 0") + self._ref_types[dst] = base + + # ------------------------------------------------------------------ + # Naming / value binding helpers + # ------------------------------------------------------------------ + + def _fresh(self, hint: str = "r") -> str: + return self._namer.fresh(hint) + + def _fresh_label(self, hint: str = "bb") -> str: + return self._namer.fresh_label(hint) + + def _bind(self, name: str, ref: str, llvm_ty: str) -> None: + self._named_values[name] = ref + self._ref_types[ref] = llvm_ty def _dest(self, instr: Instruction) -> str: """Get or create an LLVM register for the instruction's destination.""" if instr.dest is None: return "" - reg = self._fresh(instr.dest.name) - self._named_values[instr.dest.name] = reg + name = instr.dest.name + if name in self._named_values: + return self._named_values[name] + reg = self._fresh(name) + ty = self._value_type(instr.dest) + self._bind(name, reg, ty) + self._ensure_slot(name, ty) return reg + def _ensure_slot(self, name: str, ty: str) -> str | None: + """Allocate (once) a memory slot for a scalar IR variable.""" + if ty.endswith("*"): + return None + slot = self._slots.get(name) + if slot is None: + slot = self._alloc_slot(ty, 1, f"{name}_mem") + self._slots[name] = slot + return slot + + def _bind_dest(self, instr: Instruction, ref: str, ty: str) -> None: + """Bind an instruction destination and give it a scalar slot.""" + if instr.dest is None: + return + self._bind(instr.dest.name, ref, ty) + self._ensure_slot(instr.dest.name, ty) + def _op(self, instr: Instruction, idx: int) -> str: """Resolve operand idx to an LLVM value reference.""" if idx >= len(instr.operands): - return "" - op = instr.operands[idx] - return self._value_ref(op) + raise LLVMCodegenError( + f"missing operand {idx} for {instr.opcode.value}" + ) + return self._value_ref(instr.operands[idx]) - def _value_ref(self, val) -> str: - """Get LLVM reference for a Value.""" + def _value_ref(self, val: Value) -> str: + """Get LLVM reference for a Value, materializing if necessary.""" if val.name in self._named_values: + slot = self._slots.get(val.name) + if slot is not None: + ty = self._ref_types.get(self._named_values[val.name]) + if ty is None or not ty.endswith("*"): + return self._load(ty or _llvm_type(val.dtype), + slot, f"{val.name}_ld") return self._named_values[val.name] if val.is_constant and val.const_value is not None: - return str(_llvm_const(val)) - reg = self._fresh(val.name) - self._named_values[val.name] = reg + return _llvm_const(val) + base = _llvm_type(val.dtype) + if self._is_pointer_value(val): + count = 1 + for dim in val.shape: + count *= max(1, int(dim)) + ptr = self._alloc_slot(base, count, "tbuf") + self._bind(val.name, ptr, base + "*") + return ptr + reg = self._materialize_const(0, base, "undef") + self._bind(val.name, reg, base) return reg - def _fresh(self, hint: str) -> str: - """Create a fresh SSA register name.""" - safe = hint.replace(".", "_").replace("-", "_") - self._block_counter += 1 - return f"%{safe}_{self._block_counter}" + def _value_type(self, val: Value) -> str: + base = _llvm_type(val.dtype) + if self._is_pointer_value(val): + return base + "*" + ref = self._named_values.get(val.name) + if ref is not None and self._ref_types.get(ref, "").endswith("*"): + return self._ref_types[ref] + return base - def _p(self, line: str = "") -> None: - if line and not line.startswith(";"): - indent = " " * self._indent + def _is_pointer_value(self, val: Value) -> bool: + return bool(val.shape) or val.name in self._ptr_names + + def _type_of_operand(self, instr: Instruction, idx: int) -> str: + if idx >= len(instr.operands): + return "float" + val = instr.operands[idx] + if (val.is_constant and val.const_value is not None + and val.name not in self._named_values): + return _llvm_type(val.dtype) + ref = self._op(instr, idx) + return self._ref_types.get(ref, _llvm_type(val.dtype)) + + def _infer_type(self, instr: Instruction) -> str: + if instr.dest is not None: + return _llvm_type(instr.dest.dtype) + for op in instr.operands: + return _llvm_type(op.dtype) + return "float" + + # ------------------------------------------------------------------ + # Memory helpers + # ------------------------------------------------------------------ + + def _alloc_slot(self, elem_ty: str, count: int = 1, + hint: str = "slot") -> str: + """Allocate stack storage in the entry prologue; return pointer ref.""" + reg = self._fresh(hint) + self._prologue.append( + f" {reg} = alloca {elem_ty}, i32 {max(1, int(count))}" + ) + self._ref_types[reg] = elem_ty + "*" + return reg + + def _dest_buffer(self, instr: Instruction, count: int, + elem_ty: str) -> str: + """Return a destination buffer pointer, binding ``instr.dest``.""" + count = max(1, int(count)) + if instr.dest is None: + return self._alloc_slot(elem_ty, count, "tmp") + name = instr.dest.name + ref = self._named_values.get(name) + if ref is not None and self._ref_types.get(ref, "").endswith("*"): + return ref + ptr = self._alloc_slot(elem_ty, count, "buf") + self._bind(name, ptr, elem_ty + "*") + return ptr + + def _ptr_of(self, instr: Instruction, idx: int, ty: str = "float") -> str: + """Use operand idx as a pointer; scalar operands are spilled.""" + if idx >= len(instr.operands): + raise LLVMCodegenError( + f"missing pointer operand {idx} for {instr.opcode.value}" + ) + val = instr.operands[idx] + ref = self._value_ref(val) + rty = self._ref_types.get(ref, "") + if rty.endswith("*"): + return ref + ptr = self._alloc_slot(ty, 1, "spin") + if not ref.startswith("%") and val.is_constant: + ref = _llvm_const_val(val.const_value, ty) + self._p(f" store {ty} {ref}, {ty}* {ptr}") + return ptr + + def _ptr_elem_ty(self, ptr: str, fallback: str) -> str: + """Element type of a known pointer operand, else ``fallback``.""" + pty = self._ref_types.get(ptr, "") + if pty.endswith("*") and pty[:-1] in _FLOAT_TYPES: + return pty[:-1] + return fallback + + def _materialize_const(self, value: float | int, ty: str, + hint: str) -> str: + reg = self._fresh(hint) + self._namer.register_definition(reg) + if ty in _FLOAT_TYPES: + self._p(f" {reg} = fadd {ty} {_float_literal(float(value))}, 0.0") else: - indent = "" - self._lines.append(f"{indent}{line}") + self._p(f" {reg} = add {ty} 0, {int(value)}") + self._ref_types[reg] = ty + return reg + + def _coerce_operand(self, ref: str, from_ty: str, + to_ty: str, hint: str) -> str: + if from_ty == to_ty: + return ref + if from_ty.endswith("*") or to_ty.endswith("*"): + raise LLVMCodegenError(f"no coercion {from_ty} -> {to_ty}") + reg = self._fresh(hint) + self._namer.register_definition(reg) + if from_ty in _INT_TYPES and to_ty in _FLOAT_TYPES: + self._p(f" {reg} = sitofp {from_ty} {ref} to {to_ty}") + elif from_ty in _FLOAT_TYPES and to_ty in _INT_TYPES: + self._p(f" {reg} = fptosi {from_ty} {ref} to {to_ty}") + elif from_ty in _FLOAT_TYPES and to_ty in _FLOAT_TYPES: + op = "fpext" if to_ty == "double" else "fptrunc" + self._p(f" {reg} = {op} {from_ty} {ref} to {to_ty}") + elif from_ty in _INT_TYPES and to_ty in _INT_TYPES: + op = "sext" if to_ty == "i64" else "trunc" + self._p(f" {reg} = {op} {from_ty} {ref} to {to_ty}") + else: + raise LLVMCodegenError(f"no coercion {from_ty} -> {to_ty}") + self._ref_types[reg] = to_ty + return reg + + def _operand_for(self, instr: Instruction, idx: int, + target_ty: str, hint: str) -> str: + val = instr.operands[idx] + if (val.is_constant and val.const_value is not None + and val.name not in self._named_values): + return _llvm_const_val(val.const_value, target_ty) + ref = self._op(instr, idx) + from_ty = self._ref_types.get(ref, _llvm_type(val.dtype)) + if from_ty.endswith("*") and not target_ty.endswith("*"): + elem_ty = from_ty[:-1] + ref = self._load(elem_ty, ref, f"{hint}_ld") + from_ty = elem_ty + return self._coerce_operand(ref, from_ty, target_ty, hint) + + def _gep(self, ty: str, base: str, idx, hint: str) -> str: + reg = self._fresh(hint) + self._namer.register_definition(reg) + self._p(f" {reg} = getelementptr {ty}, {ty}* {base}, i32 {idx}") + self._ref_types[reg] = ty + "*" + return reg + + def _load(self, ty: str, ptr: str, hint: str) -> str: + reg = self._fresh(hint) + self._namer.register_definition(reg) + self._p(f" {reg} = load {ty}, {ty}* {ptr}") + self._ref_types[reg] = ty + return reg + + def _bin(self, op: str, ty: str, lhs, rhs, hint: str) -> str: + reg = self._fresh(hint) + self._namer.register_definition(reg) + self._p(f" {reg} = {op} {ty} {lhs}, {rhs}") + self._ref_types[reg] = ty + return reg + + def _call_math(self, name: str, ty: str, arg: str, hint: str) -> str: + fn = name + ("f" if ty == "float" else "") + reg = self._fresh(hint) + self._namer.register_definition(reg) + self._p(f" {reg} = call {ty} @{fn}({ty} {arg})") + self._ref_types[reg] = ty + return reg + + def _iadd(self, lhs, rhs, hint: str): + if isinstance(lhs, int) and isinstance(rhs, int): + return lhs + rhs + reg = self._fresh(hint) + self._namer.register_definition(reg) + self._p(f" {reg} = add i32 {lhs}, {rhs}") + self._ref_types[reg] = "i32" + return reg + + def _isub(self, lhs, rhs, hint: str): + if isinstance(lhs, int) and isinstance(rhs, int): + return lhs - rhs + reg = self._fresh(hint) + self._namer.register_definition(reg) + self._p(f" {reg} = sub i32 {lhs}, {rhs}") + self._ref_types[reg] = "i32" + return reg + + def _imul(self, lhs, rhs, hint: str): + if isinstance(lhs, int) and isinstance(rhs, int): + return lhs * rhs + reg = self._fresh(hint) + self._namer.register_definition(reg) + self._p(f" {reg} = mul i32 {lhs}, {rhs}") + self._ref_types[reg] = "i32" + return reg + + def _icmp(self, pred: str, lhs, rhs, hint: str) -> str: + reg = self._fresh(hint) + self._namer.register_definition(reg) + self._p(f" {reg} = icmp {pred} i32 {lhs}, {rhs}") + self._ref_types[reg] = "i1" + return reg + + def _bin_i1(self, op: str, lhs: str, rhs: str, hint: str) -> str: + reg = self._fresh(hint) + self._namer.register_definition(reg) + self._p(f" {reg} = {op} i1 {lhs}, {rhs}") + self._ref_types[reg] = "i1" + return reg # ------------------------------------------------------------------ # Arithmetic # ------------------------------------------------------------------ - def _emit_add(self, instr: Instruction) -> None: - dst = self._dest(instr) - lhs = self._op(instr, 0) - rhs = self._op(instr, 1) + def _emit_binary(self, instr: Instruction, fop: str, iop: str) -> None: ty = self._infer_type(instr) - self._p(f" {dst} = fadd {ty} {lhs}, {rhs}") + lhs = self._operand_for(instr, 0, ty, "cvt_l") + rhs = self._operand_for(instr, 1, ty, "cvt_r") + dst = self._dest(instr) + op = fop if ty in _FLOAT_TYPES else iop + self._p(f" {dst} = {op} {ty} {lhs}, {rhs}") + + def _emit_add(self, instr: Instruction) -> None: + self._emit_binary(instr, "fadd", "add") def _emit_sub(self, instr: Instruction) -> None: - dst = self._dest(instr) - lhs = self._op(instr, 0) - rhs = self._op(instr, 1) - ty = self._infer_type(instr) - self._p(f" {dst} = fsub {ty} {lhs}, {rhs}") + self._emit_binary(instr, "fsub", "sub") def _emit_mul(self, instr: Instruction) -> None: - dst = self._dest(instr) - lhs = self._op(instr, 0) - rhs = self._op(instr, 1) - ty = self._infer_type(instr) - self._p(f" {dst} = fmul {ty} {lhs}, {rhs}") + self._emit_binary(instr, "fmul", "mul") def _emit_div(self, instr: Instruction) -> None: - dst = self._dest(instr) - lhs = self._op(instr, 0) - rhs = self._op(instr, 1) - ty = self._infer_type(instr) - self._p(f" {dst} = fdiv {ty} {lhs}, {rhs}") + self._emit_binary(instr, "fdiv", "sdiv") def _emit_neg(self, instr: Instruction) -> None: - dst = self._dest(instr) - src = self._op(instr, 0) ty = self._infer_type(instr) - self._p(f" {dst} = fneg {ty} {src}") + src = self._operand_for(instr, 0, ty, "cvt") + dst = self._dest(instr) + if ty in _FLOAT_TYPES: + self._p(f" {dst} = fneg {ty} {src}") + else: + self._p(f" {dst} = sub {ty} 0, {src}") def _emit_exp(self, instr: Instruction) -> None: - dst = self._dest(instr) - src = self._op(instr, 0) ty = self._infer_type(instr) - if ty == "double": - self._p(f" {dst} = call double @exp(double {src})") - else: - self._p(f" {dst} = call float @expf(float {src})") + if ty in _FLOAT_TYPES: + src = self._operand_for(instr, 0, ty, "cvt") + dst = self._dest(instr) + fn = "exp" if ty == "double" else "expf" + self._p(f" {dst} = call {ty} @{fn}({ty} {src})") + return + src = self._operand_for(instr, 0, "float", "cvt") + exp_reg = self._call_math("exp", "float", src, "expf") + dst = self._dest(instr) + self._p(f" {dst} = fptosi float {exp_reg} to {ty}") + self._ref_types[dst] = ty # ------------------------------------------------------------------ # Constants & memory @@ -272,279 +674,656 @@ def _emit_exp(self, instr: Instruction) -> None: def _emit_load_const(self, instr: Instruction) -> None: dst = self._dest(instr) raw_val = instr.attrs.get("value", 0) - assert isinstance(raw_val, (int, float)) - val: float | int = raw_val ty = _llvm_type(instr.dest.dtype) if instr.dest else "float" - self._p(f" {dst} = fadd {ty} {_llvm_const_val(val, ty)}, 0.0") + if ty in _FLOAT_TYPES: + lit = _float_literal(float(raw_val)) + self._p(f" {dst} = fadd {ty} 0.0, {lit}") + else: + self._p(f" {dst} = add {ty} 0, {int(raw_val)}") def _emit_load(self, instr: Instruction) -> None: - dst = self._dest(instr) + ty = _llvm_type(instr.dest.dtype) if instr.dest else "float" ptr = self._op(instr, 0) - ty = self._infer_type(instr) - ptr_ty = f"{ty}*" - self._p(f" {dst} = load {ty}, {ptr_ty} {ptr}") + pty = self._ref_types.get(ptr, "") + if pty.endswith("*"): + ty = pty[:-1] + else: + ptr = self._ptr_of(instr, 0, ty) + dst = self._dest(instr) + self._p(f" {dst} = load {ty}, {ty}* {ptr}") + self._ref_types[dst] = ty def _emit_store(self, instr: Instruction) -> None: - val = self._op(instr, 1) ptr = self._op(instr, 0) - ty = self._infer_type(instr) - ptr_ty = f"{ty}*" - self._p(f" store {ty} {val}, {ptr_ty} {ptr}") + pty = self._ref_types.get(ptr, "") + if not pty.endswith("*"): + raise LLVMCodegenError("store target is not a pointer") + elem_ty = pty[:-1] + val = self._operand_for(instr, 1, elem_ty, "stv") + self._p(f" store {elem_ty} {val}, {elem_ty}* {ptr}") def _emit_alloca(self, instr: Instruction) -> None: dst = self._dest(instr) size = instr.attrs.get("size", 4) - ty = self._infer_type(instr) - self._p(f" {dst} = alloca {ty}, i32 {size}") + try: + size = int(size) + except (TypeError, ValueError): + size = 1 + ty = _llvm_type(instr.dest.dtype) if instr.dest else "float" + self._p(f" {dst} = alloca {ty}, i32 {max(1, size)}") + self._ref_types[dst] = ty + "*" # ------------------------------------------------------------------ # Control flow # ------------------------------------------------------------------ def _emit_for(self, instr: Instruction) -> None: - self._dest(instr) # register loop variable name - start = instr.attrs.get("start", 0) - end = instr.attrs.get("end", 0) - - header = self._fresh_block("loop_header") - body = self._fresh_block("loop_body") - exit = self._fresh_block("loop_exit") - - # Initialize induction variable - start_reg = self._fresh("iv_start") - self._p(f" {start_reg} = add i32 0, {start}") - # Actually, we need an alloca or phi for the IV - iv_alloca = self._fresh("iv_ptr") - self._p(f" {iv_alloca} = alloca i32, i32 1") - self._p(f" store i32 {start_reg}, i32* {iv_alloca}") - self._p(f" br label %{body}") - - self._p(f"{header}:") - # Load IV, compare - loaded = self._fresh("iv_val") - self._p(f" {loaded} = load i32, i32* {iv_alloca}") - cond = self._fresh("cond") - self._p(f" {cond} = icmp slt i32 {loaded}, {end}") - self._p(f" br i1 {cond}, label %{body}, label %{exit}") - - self._p(f"{body}:") - - self._loop_context = { - "iv_alloca": iv_alloca, - "end": end, - "header": header, - "exit": exit, - } + limit = max(0, self._as_int(instr.attrs.get("end", 0), 0)) + start = self._as_int(instr.attrs.get("start", 0), 0) + step = self._as_int(instr.attrs.get("step", 1), 1) or 1 + ir_name = instr.dest.name if instr.dest is not None else None + self._loop_open(limit, ir_name, "loop_i", start, step) def _emit_endfor(self, instr: Instruction) -> None: - ctx = self._loop_context - if ctx is None: - self._p(" ; ERROR: endfor without matching for") - return - - iv_alloca = ctx["iv_alloca"] - header = ctx["header"] - - # Load, increment, store - loaded = self._fresh("iv_val") - self._p(f" {loaded} = load i32, i32* {iv_alloca}") - inc = self._fresh("iv_next") - self._p(f" {inc} = add i32 {loaded}, 1") - self._p(f" store i32 {inc}, i32* {iv_alloca}") - self._p(f" br label %{header}") - - # Exit label - self._p(f"{ctx['exit']}:") + if not self._loop_stack: + raise LLVMCodegenError("endfor without matching for") + self._loop_close(self._loop_stack[-1]) + + def _loop_open(self, limit: int, ir_name: str | None, hint: str, + start: int = 0, step: int = 1) -> LoopContext: + limit = max(0, int(limit)) + start = int(start) + step = int(step) or 1 + + ptr = self._alloc_slot("i32", 1, f"{hint}_ptr") + self._p(f" store i32 {start}, i32* {ptr}") + header = self._fresh_label(f"{hint}_hdr") + body = self._fresh_label(f"{hint}_bdy") + exit_ = self._fresh_label(f"{hint}_ext") + + self._start_block(header) + iv = self._fresh(f"{hint}_ld") + self._namer.register_definition(iv) + self._p(f" {iv} = load i32, i32* {ptr}") + cond = self._fresh(f"{hint}_cond") + self._namer.register_definition(cond) + self._p(f" {cond} = icmp slt i32 {iv}, {limit}") + self._terminate(f"br i1 {cond}, label %{body}, label %{exit_}") + self._start_block(body) + + ctx = LoopContext(ir_name, ptr, iv, header, body, exit_, limit, step) + self._loop_stack.append(ctx) + if ir_name is not None: + self._bind(ir_name, iv, "i32") + return ctx + + def _loop_close(self, ctx: LoopContext) -> None: + if not self._loop_stack or self._loop_stack[-1] is not ctx: + raise LLVMCodegenError("endfor without matching for") + self._loop_stack.pop() + if not self._terminated: + cur = self._fresh("iv_cur") + self._namer.register_definition(cur) + self._p(f" {cur} = load i32, i32* {ctx.ptr}") + nxt = self._fresh("iv_nxt") + self._namer.register_definition(nxt) + self._p(f" {nxt} = add i32 {cur}, {ctx.step}") + self._p(f" store i32 {nxt}, i32* {ctx.ptr}") + self._terminate(f"br label %{ctx.header}") + self._start_block(ctx.exit) def _emit_br(self, instr: Instruction) -> None: - target = instr.target or "" - self._p(f" br label %{target}") + target = (instr.target or "").strip() + if not target: + self._terminate("unreachable") + else: + self._terminate(f"br label %{target}") def _emit_br_if(self, instr: Instruction) -> None: - cond_op = self._op(instr, 0) if instr.operands else "" - targets = (instr.target or ",").split(",") - true_t = targets[0].strip() if len(targets) > 0 else "" - false_t = targets[1].strip() if len(targets) > 1 else "" + targets = (instr.target or "").split(",") + true_t = targets[0].strip() if targets else "" + false_t = targets[1].strip() if len(targets) > 1 else true_t + if not false_t: + false_t = true_t + if not true_t: + self._terminate("unreachable") + return - if cond_op: - self._p(f" br i1 {cond_op}, label %{true_t}, label %{false_t}") + cmp_op = instr.attrs.get("cmp_op") + if cmp_op and len(instr.operands) >= 2: + key = str(cmp_op) + if key not in _CMP_PRED: + raise LLVMCodegenError(f"unknown comparison operator: {key}") + ty = self._infer_type(instr) + lhs = self._operand_for(instr, 0, ty, "cmpl") + rhs = self._operand_for(instr, 1, ty, "cmpr") + fpred, ipred = _CMP_PRED[key] + is_float = ty in _FLOAT_TYPES + kind = "fcmp" if is_float else "icmp" + pred = fpred if is_float else ipred + cond = self._fresh(f"brc_{kind}") + self._namer.register_definition(cond) + self._p(f" {cond} = {kind} {pred} {ty} {lhs}, {rhs}") + self._ref_types[cond] = "i1" + elif instr.operands: + cond = self._op(instr, 0) + cty = self._ref_types.get(cond) or self._type_of_operand(instr, 0) + if cty.endswith("*"): + elem_ty = cty[:-1] + cond = self._load(elem_ty, cond, "brc_ld") + cty = elem_ty + if cty == "i32": + cond = self._icmp("ne", cond, 0, "brc") + elif cty in _FLOAT_TYPES: + cmp_reg = self._fresh("brc") + self._namer.register_definition(cmp_reg) + self._p(f" {cmp_reg} = fcmp one {cty} {cond}, 0.0") + self._ref_types[cmp_reg] = "i1" + cond = cmp_reg + elif cty == "i64": + cmp_reg = self._fresh("brc") + self._namer.register_definition(cmp_reg) + self._p(f" {cmp_reg} = icmp ne i64 {cond}, 0") + self._ref_types[cmp_reg] = "i1" + cond = cmp_reg else: - self._p(f" br label %{true_t}") + self._terminate("unreachable") + return + + self._terminate(f"br i1 {cond}, label %{true_t}, label %{false_t}") def _emit_return(self, instr: Instruction) -> None: if instr.operands: - val = self._op(instr, 0) - ty = self._infer_type(instr) - self._p(f" ret {ty} {val}") + val = instr.operands[0] + ref = self._value_ref(val) + ty = self._ref_types.get(ref) + if ty is None: + ty = self._value_type(val) + self._terminate(f"ret {ty} {ref}") else: - self._p(" ret void") + self._terminate("ret void") def _emit_label(self, instr: Instruction) -> None: """IR labels become LLVM block labels.""" if instr.target: - self._p(f"{instr.target}:") + self._start_block(instr.target) # ------------------------------------------------------------------ - # Neural-network ops (implemented as inline LLVM IR) + # Neural-network ops # ------------------------------------------------------------------ - def _emit_relu(self, instr: Instruction) -> None: - """ReLU(x) = select x > 0 ? x : 0.0""" - dst = self._dest(instr) - src = self._op(instr, 0) - ty = self._infer_type(instr) - zero = "0.0" - if ty == "double": - zero = "0.0" - cmp = self._fresh("cmp") - self._p(f" {cmp} = fcmp ogt {ty} {src}, {zero}") - self._p(f" {dst} = select i1 {cmp}, {ty} {src}, {ty} {zero}") - - def _emit_gelu(self, instr: Instruction) -> None: - """ - GELU(x) = x * 0.5 * (1.0 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) - All computed inline using LLVM IR. - """ - dst = self._dest(instr) - x = self._op(instr, 0) - ty = self._infer_type(instr) - - sqrt_2pi = "0.7978845608028654" - coeff = "0.044715" - half = "0.5" - one = "1.0" - - # x^3 - x3 = self._fresh("x3") - self._p(f" {x3} = fmul {ty} {x}, {x}") - self._p(f" {x3} = fmul {ty} {x3}, {x}") + def _finalize_from_acc(self, instr: Instruction, + acc: str, ty: str) -> None: + """Result scalar lives in ``acc``; publish it to ``instr.dest``.""" + if instr.dest is None: + return + value = self._load(ty, acc, "res") + if self._is_pointer_value(instr.dest): + count = 1 + for dim in instr.dest.shape: + count *= max(1, int(dim)) + buf = self._dest_buffer(instr, count, ty) + self._p(f" store {ty} {value}, {ty}* {buf}") + else: + self._bind_dest(instr, value, ty) - # inner = coeff * x^3 + x - inner = self._fresh("inner") - self._p(f" {inner} = fmul {ty} {coeff}, {x3}") - self._p(f" {inner} = fadd {ty} {inner}, {x}") + def _finalize_from_buffer(self, instr: Instruction, + buf: str, ty: str) -> None: + """Result tensor lives in ``buf``; load scalar if dest is scalar.""" + if instr.dest is None: + return + if not self._is_pointer_value(instr.dest): + value = self._load(ty, buf, "res") + self._bind_dest(instr, value, ty) + + def _dim_of(self, instr: Instruction, keys: tuple[str, ...], + default: int = 1, operand: int | None = None, + axis: int | None = None) -> int: + for key in keys: + value = instr.attrs.get(key) + if isinstance(value, (int, float)) and int(value) > 0: + return int(value) + if operand is not None and operand < len(instr.operands): + shape = instr.operands[operand].shape + if shape: + idx = axis if axis is not None else 0 + if -len(shape) <= idx < len(shape): + dim = int(shape[idx]) + if dim > 0: + return dim + return default - # inner *= sqrt(2/pi) - self._p(f" {inner} = fmul {ty} {inner}, {sqrt_2pi}") + @staticmethod + def _as_int(value, default: int = 0) -> int: + if isinstance(value, Value): + if value.const_value is not None: + return int(value.const_value) + return default + try: + return int(value) + except (TypeError, ValueError): + return default - # tanh - if ty == "double": - tanh_reg = self._fresh("tanh") - self._p(f" {tanh_reg} = call double @tanh(double {inner})") + def _emit_relu(self, instr: Instruction) -> None: + ty = self._infer_type(instr) + src = self._operand_for(instr, 0, ty, "cvt") + dst = self._dest(instr) + cmp_reg = self._fresh("cmp") + self._namer.register_definition(cmp_reg) + if ty in _FLOAT_TYPES: + self._p(f" {cmp_reg} = fcmp ogt {ty} {src}, 0.0") + self._p(f" {dst} = select i1 {cmp_reg}, {ty} {src}, {ty} 0.0") else: - tanh_reg = self._fresh("tanh") - self._p(f" {tanh_reg} = call float @tanhf(float {inner})") - - # 1 + tanh - plus_one = self._fresh("plus_one") - self._p(f" {plus_one} = fadd {ty} {one}, {tanh_reg}") + self._p(f" {cmp_reg} = icmp sgt {ty} {src}, 0") + self._p(f" {dst} = select i1 {cmp_reg}, {ty} {src}, {ty} 0") - # x * 0.5 - half_x = self._fresh("half_x") - self._p(f" {half_x} = fmul {ty} {x}, {half}") + def _emit_gelu(self, instr: Instruction) -> None: + """GELU(x) = 0.5*x*(1 + tanh(sqrt(2/pi)*(x + 0.044715*x^3))).""" + ty = self._infer_type(instr) + x = self._operand_for(instr, 0, ty, "cvt") + if instr.dest is None: + return + t1 = self._bin("fmul", ty, x, x, "gelu_t1") + x3 = self._bin("fmul", ty, t1, x, "gelu_x3") + in1 = self._bin("fmul", ty, _float_literal(0.044715), x3, "gelu_in1") + in2 = self._bin("fadd", ty, in1, x, "gelu_in2") + inner = self._bin( + "fmul", ty, in2, _float_literal(0.7978845608028654), "gelu_inner" + ) + tanh_reg = self._call_math("tanh", ty, inner, "gelu_tanh") + p1 = self._bin("fadd", ty, "1.0", tanh_reg, "gelu_p1") + hx = self._bin("fmul", ty, x, "0.5", "gelu_hx") + out = self._bin("fmul", ty, hx, p1, "gelu_out") + self._bind_dest(instr, out, ty) - # result - self._p(f" {dst} = fmul {ty} {half_x}, {plus_one}") + def _emit_sigmoid(self, instr: Instruction) -> None: + """Sigmoid(x) = 1 / (1 + exp(-x)).""" + ty = self._infer_type(instr) + x = self._operand_for(instr, 0, ty, "cvt") + if instr.dest is None: + return + neg = self._fresh("sig_neg") + self._namer.register_definition(neg) + self._p(f" {neg} = fneg {ty} {x}") + self._ref_types[neg] = ty + exp_reg = self._call_math("exp", ty, neg, "sig_exp") + den = self._bin("fadd", ty, "1.0", exp_reg, "sig_den") + out = self._bin("fdiv", ty, "1.0", den, "sig_out") + self._bind_dest(instr, out, ty) def _emit_softmax(self, instr: Instruction) -> None: - """ - Softmax: for a vector of N elements: - max_val = max(x) - sum = sum(exp(x[i] - max_val)) - result[i] = exp(x[i] - max_val) / sum - Implemented as a loop. - """ - dst = self._dest(instr) - src = self._op(instr, 0) + """Numerically stable softmax over the last axis (three passes).""" ty = self._infer_type(instr) - - self._p(" ; softmax: TODO full vector implementation required") - self._p(" ; placeholder: return exp(x) / sum(exp(x))") - # For now, call external softmax helper - if ty == "double": - self._p(f" {dst} = call double @exp(double {src})") - else: - self._p(f" {dst} = call float @expf(float {src})") + x = self._ptr_of(instr, 0, ty) + ty = self._ptr_elem_ty(x, ty) + n = self._dim_of( + instr, ("length", "len", "n"), 1, operand=0, axis=-1 + ) + out = self._dest_buffer(instr, n, ty) + + max_ptr = self._alloc_slot(ty, 1, "sm_m") + self._p(f" store {ty} {_float_literal(-3.4e38)}, {ty}* {max_ptr}") + ctx = self._loop_open(n, None, "sm_max_i") + xv = self._load(ty, self._gep(ty, x, ctx.value, "sm_xp1"), "sm_xv1") + mv = self._load(ty, max_ptr, "sm_mv1") + cmp_reg = self._fresh("sm_cmp1") + self._namer.register_definition(cmp_reg) + self._p(f" {cmp_reg} = fcmp ogt {ty} {xv}, {mv}") + new_max = self._fresh("sm_new1") + self._namer.register_definition(new_max) + self._p(f" {new_max} = select i1 {cmp_reg}, {ty} {xv}, {ty} {mv}") + self._p(f" store {ty} {new_max}, {ty}* {max_ptr}") + self._loop_close(ctx) + + sum_ptr = self._alloc_slot(ty, 1, "sm_s") + self._p(f" store {ty} 0.0, {ty}* {sum_ptr}") + ctx = self._loop_open(n, None, "sm_sum_i") + xv = self._load(ty, self._gep(ty, x, ctx.value, "sm_xp2"), "sm_xv2") + mv = self._load(ty, max_ptr, "sm_mv2") + diff = self._bin("fsub", ty, xv, mv, "sm_diff2") + exp_reg = self._call_math("exp", ty, diff, "sm_exp2") + sv = self._load(ty, sum_ptr, "sm_sv2") + new_sum = self._bin("fadd", ty, sv, exp_reg, "sm_new2") + self._p(f" store {ty} {new_sum}, {ty}* {sum_ptr}") + self._loop_close(ctx) + + ctx = self._loop_open(n, None, "sm_div_i") + xv = self._load(ty, self._gep(ty, x, ctx.value, "sm_xp3"), "sm_xv3") + mv = self._load(ty, max_ptr, "sm_mv3") + diff = self._bin("fsub", ty, xv, mv, "sm_diff3") + exp_reg = self._call_math("exp", ty, diff, "sm_exp3") + sv = self._load(ty, sum_ptr, "sm_sv3") + quotient = self._bin("fdiv", ty, exp_reg, sv, "sm_q3") + op = self._gep(ty, out, ctx.value, "sm_op3") + self._p(f" store {ty} {quotient}, {ty}* {op}") + self._loop_close(ctx) + + self._finalize_from_buffer(instr, out, ty) def _emit_maxpool(self, instr: Instruction) -> None: - dst = self._dest(instr) - src = self._op(instr, 0) - self._p(" ; maxpool: passthrough (requires full tensor support)") - self._p(f" {dst} = fadd {self._infer_type(instr)} {src}, 0.0") - - def _emit_matmul(self, instr: Instruction) -> None: - """Matrix multiplication: C[m,n] = A[m,k] @ B[k,n] - NOTE: Full tensor MatMul requires multi-dimensional arrays. - For scalar test cases, we compute a simple dot product approximation. - """ - dst = self._dest(instr) - a = self._op(instr, 0) - b = self._op(instr, 1) + """2D max pooling (NCHW layout, no padding).""" ty = self._infer_type(instr) - self._p(f" ; matmul: A[{a}], B[{b}] - requires multi-dim support") - self._p(f" {dst} = fmul {ty} {a}, {b}") + x = self._ptr_of(instr, 0, ty) + ty = self._ptr_elem_ty(x, ty) + shape = tuple(instr.operands[0].shape) + if len(shape) >= 3: + channels, height, width = shape[-3], shape[-2], shape[-1] + else: + channels = height = width = 1 + channels = max(1, int(channels)) + height = max(1, int(height)) + width = max(1, int(width)) + kernel = max(1, self._dim_of(instr, ("kernel", "kernel_shape"), 2)) + stride = max(1, self._dim_of(instr, ("stride", "strides"), 2)) + out_h = max(0, (height - kernel) // stride + 1) + out_w = max(0, (width - kernel) // stride + 1) + out = self._dest_buffer(instr, channels * out_h * out_w, ty) + + c_ctx = self._loop_open(channels, None, "mp_c") + h_ctx = self._loop_open(out_h, None, "mp_oh") + w_ctx = self._loop_open(out_w, None, "mp_ow") + max_ptr = self._alloc_slot(ty, 1, "mp_max") + self._p(f" store {ty} {_float_literal(-3.4e38)}, {ty}* {max_ptr}") + + kh_ctx = self._loop_open(kernel, None, "mp_kh") + ih = self._iadd( + self._imul(h_ctx.value, stride, "mp_ih0"), + kh_ctx.value, "mp_ih", + ) + in_h = self._imul(ih, width, "mp_in_h") + kw_ctx = self._loop_open(kernel, None, "mp_kw") + iw = self._iadd( + self._imul(w_ctx.value, stride, "mp_iw0"), + kw_ctx.value, "mp_iw", + ) + offset = self._iadd( + self._iadd( + self._imul(c_ctx.value, height * width, "mp_coff"), + in_h, "mp_hoff", + ), + iw, "mp_off", + ) + xv = self._load(ty, self._gep(ty, x, offset, "mp_xp"), "mp_xv") + mv = self._load(ty, max_ptr, "mp_mv") + cmp_reg = self._fresh("mp_cmp") + self._namer.register_definition(cmp_reg) + self._p(f" {cmp_reg} = fcmp ogt {ty} {xv}, {mv}") + new_max = self._fresh("mp_new") + self._namer.register_definition(new_max) + self._p(f" {new_max} = select i1 {cmp_reg}, {ty} {xv}, {ty} {mv}") + self._p(f" store {ty} {new_max}, {ty}* {max_ptr}") + self._loop_close(kw_ctx) + self._loop_close(kh_ctx) + + out_offset = self._iadd( + self._iadd( + self._imul(c_ctx.value, out_h * out_w, "mp_ooff0"), + self._imul(h_ctx.value, out_w, "mp_ooff1"), + "mp_ooff1", + ), + w_ctx.value, "mp_ooff", + ) + result = self._load(ty, max_ptr, "mp_res") + self._p( + f" store {ty} {result}, {ty}* " + f"{self._gep(ty, out, out_offset, 'mp_op')}" + ) + self._loop_close(w_ctx) + self._loop_close(h_ctx) + self._loop_close(c_ctx) + self._finalize_from_buffer(instr, out, ty) def _emit_dot(self, instr: Instruction) -> None: - """Dot product: sum(a[i] * b[i]) for i in 0..len-1""" - dst = self._dest(instr) - a = self._op(instr, 0) - b = self._op(instr, 1) - length = instr.attrs.get("length", 1) + """Dot product: sum(a[i] * b[i]).""" ty = self._infer_type(instr) - self._p(f" ; dot product len={length} - scalar approximation") - self._p(f" {dst} = fmul {ty} {a}, {b}") + a = self._ptr_of(instr, 0, ty) + ty = self._ptr_elem_ty(a, ty) + b = self._ptr_of(instr, 1, ty) + length = max(1, self._dim_of( + instr, ("length", "len", "n"), 1, operand=0, axis=-1 + )) + acc = self._alloc_slot(ty, 1, "dot_acc") + self._p(f" store {ty} 0.0, {ty}* {acc}") + ctx = self._loop_open(length, None, "dot_i") + av = self._load(ty, self._gep(ty, a, ctx.value, "dot_ap"), "dot_av") + bv = self._load(ty, self._gep(ty, b, ctx.value, "dot_bp"), "dot_bv") + prod = self._bin("fmul", ty, av, bv, "dot_pr") + old = self._load(ty, acc, "dot_old") + new_acc = self._bin("fadd", ty, old, prod, "dot_new") + self._p(f" store {ty} {new_acc}, {ty}* {acc}") + self._loop_close(ctx) + self._finalize_from_acc(instr, acc, ty) - def _emit_conv(self, instr: Instruction) -> None: - dst = self._dest(instr) - out_c = instr.attrs.get("out_channels", 1) - ksize = instr.attrs.get("kernel_size", 3) - stride = instr.attrs.get("stride", 1) + def _emit_matmul(self, instr: Instruction) -> None: + """Matrix multiply: C[m,n] = A[m,k] x B[k,n] (row major).""" ty = self._infer_type(instr) - self._p(f" ; conv: out_c={out_c} k={ksize} s={stride}") - self._p(f" {dst} = fadd {ty} 0.0, 0.0 ; conv placeholder") + a = self._ptr_of(instr, 0, ty) + ty = self._ptr_elem_ty(a, ty) + b = self._ptr_of(instr, 1, ty) + m = self._dim_of(instr, ("m", "rows"), 1, operand=0, axis=0) + k = self._dim_of(instr, ("k", "inner"), 1, operand=0, axis=1) + n = self._dim_of(instr, ("n", "cols"), 1, operand=1, axis=1) + m, k, n = max(1, m), max(1, k), max(1, n) + c = self._dest_buffer(instr, m * n, ty) + acc = self._alloc_slot(ty, 1, "mm_acc") + + i_ctx = self._loop_open(m, None, "mm_i") + j_ctx = self._loop_open(n, None, "mm_j") + self._p(f" store {ty} 0.0, {ty}* {acc}") + k_ctx = self._loop_open(k, None, "mm_k") + a_off = self._iadd( + self._imul(i_ctx.value, k, "mm_aoff0"), k_ctx.value, "mm_aoff" + ) + av = self._load(ty, self._gep(ty, a, a_off, "mm_ap"), "mm_av") + b_off = self._iadd( + self._imul(k_ctx.value, n, "mm_boff0"), j_ctx.value, "mm_boff" + ) + bv = self._load(ty, self._gep(ty, b, b_off, "mm_bp"), "mm_bv") + prod = self._bin("fmul", ty, av, bv, "mm_pr") + old = self._load(ty, acc, "mm_old") + new_acc = self._bin("fadd", ty, old, prod, "mm_new") + self._p(f" store {ty} {new_acc}, {ty}* {acc}") + self._loop_close(k_ctx) + + c_off = self._iadd( + self._imul(i_ctx.value, n, "mm_coff0"), j_ctx.value, "mm_coff" + ) + result = self._load(ty, acc, "mm_res") + c_ptr = self._gep(ty, c, c_off, "mm_cp") + self._p(f" store {ty} {result}, {ty}* {c_ptr}") + self._loop_close(j_ctx) + self._loop_close(i_ctx) + self._finalize_from_buffer(instr, c, ty) def _emit_gemm(self, instr: Instruction) -> None: - dst = self._dest(instr) - a = self._op(instr, 0) - b = self._op(instr, 1) + """General matrix multiply: C = A x W (+ bias), optional trans_b.""" ty = self._infer_type(instr) - self._p(f" ; gemm: {a} @ {b}") - self._p(f" {dst} = fmul {ty} {a}, {b} ; gemm placeholder") + a = self._ptr_of(instr, 0, ty) + ty = self._ptr_elem_ty(a, ty) + w = self._ptr_of(instr, 1, ty) + bias = self._ptr_of(instr, 2, ty) if len(instr.operands) >= 3 else None + if instr.attrs.get("trans_a"): + raise LLVMCodegenError("gemm trans_a is not supported") + + trans_b = bool(instr.attrs.get( + "trans_b", instr.attrs.get("transB", False) + )) + m = self._dim_of(instr, ("M", "m", "rows"), 1, operand=0, axis=0) + k = self._dim_of(instr, ("K", "k", "inner"), 1, operand=0, axis=1) + if trans_b: + n = self._dim_of(instr, ("N", "n", "cols"), 1, operand=1, axis=0) + else: + n = self._dim_of(instr, ("N", "n", "cols"), 1, operand=1, axis=1) + m, k, n = max(1, m), max(1, k), max(1, n) + c = self._dest_buffer(instr, m * n, ty) + acc = self._alloc_slot(ty, 1, "gemm_acc") + + i_ctx = self._loop_open(m, None, "gemm_i") + j_ctx = self._loop_open(n, None, "gemm_j") + if bias is not None: + bias_val = self._load( + ty, self._gep(ty, bias, j_ctx.value, "gemm_bp"), "gemm_bias" + ) + else: + bias_val = _float_literal(0.0) + self._p(f" store {ty} {bias_val}, {ty}* {acc}") + k_ctx = self._loop_open(k, None, "gemm_k") + a_off = self._iadd( + self._imul(i_ctx.value, k, "gemm_aoff0"), k_ctx.value, "gemm_aoff" + ) + av = self._load(ty, self._gep(ty, a, a_off, "gemm_ap"), "gemm_av") + if trans_b: + w_off = self._iadd( + self._imul(j_ctx.value, k, "gemm_woff0"), + k_ctx.value, "gemm_woff", + ) + else: + w_off = self._iadd( + self._imul(k_ctx.value, n, "gemm_woff0"), + j_ctx.value, "gemm_woff", + ) + wv = self._load(ty, self._gep(ty, w, w_off, "gemm_wp"), "gemm_wv") + prod = self._bin("fmul", ty, av, wv, "gemm_pr") + old = self._load(ty, acc, "gemm_old") + new_acc = self._bin("fadd", ty, old, prod, "gemm_new") + self._p(f" store {ty} {new_acc}, {ty}* {acc}") + self._loop_close(k_ctx) + + c_off = self._iadd( + self._imul(i_ctx.value, n, "gemm_coff0"), j_ctx.value, "gemm_coff" + ) + result = self._load(ty, acc, "gemm_res") + c_ptr = self._gep(ty, c, c_off, "gemm_cp") + self._p(f" store {ty} {result}, {ty}* {c_ptr}") + self._loop_close(j_ctx) + self._loop_close(i_ctx) + self._finalize_from_buffer(instr, c, ty) - def _emit_sigmoid(self, instr: Instruction) -> None: - dst = self._dest(instr) - src = self._op(instr, 0) + def _emit_conv(self, instr: Instruction) -> None: + """2D convolution (NCHW layout) with stride/padding/bias.""" ty = self._infer_type(instr) - if ty == "double": - self._p(f" {dst} = call double @exp(double {src})") - self._p(f" {dst} = fadd double 1.0, {dst}") - self._p(f" {dst} = fdiv double 1.0, {dst}") + x = self._ptr_of(instr, 0, ty) + ty = self._ptr_elem_ty(x, ty) + w = self._ptr_of(instr, 1, ty) + bias = self._ptr_of(instr, 2, ty) if len(instr.operands) >= 3 else None + + x_shape = tuple(instr.operands[0].shape) + if len(x_shape) >= 3: + cin, height, width = x_shape[-3], x_shape[-2], x_shape[-1] + else: + cin = height = width = 1 + w_shape = tuple(instr.operands[1].shape) + if len(w_shape) >= 3: + cout, cin, kernel = w_shape[0], w_shape[1], w_shape[2] else: - self._p(f" {dst} = call float @expf(float {src})") - self._p(f" {dst} = fadd float 1.0, {dst}") - self._p(f" {dst} = fdiv float 1.0, {dst}") + cout = self._dim_of(instr, ("out_channels", "cout"), 1) + kernel = self._dim_of( + instr, ("kernel_size", "kernel_shape"), 3 + ) + cin = max(1, int(cin)) + height = max(1, int(height)) + width = max(1, int(width)) + cout = max(1, int(cout)) + kernel = max(1, int(kernel)) + stride = max(1, self._dim_of(instr, ("stride", "strides"), 1)) + padding = max(0, self._as_int( + instr.attrs.get("padding", instr.attrs.get("pads", 0)), 0 + )) + out_h = max(0, (height + 2 * padding - kernel) // stride + 1) + out_w = max(0, (width + 2 * padding - kernel) // stride + 1) + out = self._dest_buffer(instr, cout * out_h * out_w, ty) + acc = self._alloc_slot(ty, 1, "conv_acc") + + oc_ctx = self._loop_open(cout, None, "conv_oc") + out_oc = self._imul(oc_ctx.value, out_h * out_w, "conv_out_oc") + w_oc = self._imul(oc_ctx.value, cin * kernel * kernel, "conv_w_oc") + + oh_ctx = self._loop_open(out_h, None, "conv_oh") + oh_s = self._imul(oh_ctx.value, stride, "conv_ohS") + ow_ctx = self._loop_open(out_w, None, "conv_ow") + ow_s = self._imul(ow_ctx.value, stride, "conv_owS") + if bias is not None: + bias_val = self._load( + ty, self._gep(ty, bias, oc_ctx.value, "conv_bp"), "conv_bias" + ) + else: + bias_val = _float_literal(0.0) + self._p(f" store {ty} {bias_val}, {ty}* {acc}") + + ic_ctx = self._loop_open(cin, None, "conv_ic") + in_ic = self._imul(ic_ctx.value, height * width, "conv_in_ic") + w_ic = self._imul(ic_ctx.value, kernel * kernel, "conv_w_ic") + + kh_ctx = self._loop_open(kernel, None, "conv_kh") + ih = self._iadd(oh_s, kh_ctx.value, "conv_ih0") + if padding: + ih = self._isub(ih, padding, "conv_ih") + ok_h_lo = self._icmp("sge", ih, 0, "conv_okh_lo") + ok_h_hi = self._icmp("slt", ih, height, "conv_okh_hi") + ok_h = self._bin_i1("and", ok_h_lo, ok_h_hi, "conv_okh") + in_h = self._imul(ih, width, "conv_in_h") + w_kh = self._imul(kh_ctx.value, kernel, "conv_w_kh") + + kw_ctx = self._loop_open(kernel, None, "conv_kw") + iw = self._iadd(ow_s, kw_ctx.value, "conv_iw0") + if padding: + iw = self._isub(iw, padding, "conv_iw") + ok_w_lo = self._icmp("sge", iw, 0, "conv_okw_lo") + ok_w_hi = self._icmp("slt", iw, width, "conv_okw_hi") + ok_w = self._bin_i1("and", ok_w_lo, ok_w_hi, "conv_okw") + ok = self._bin_i1("and", ok_h, ok_w, "conv_ok") + mac_label = self._fresh_label("conv_mac") + skip_label = self._fresh_label("conv_skip") + self._terminate( + f"br i1 {ok}, label %{mac_label}, label %{skip_label}" + ) + self._start_block(mac_label) + x_off = self._iadd( + self._iadd(in_ic, in_h, "conv_xoff0"), iw, "conv_xoff" + ) + xv = self._load(ty, self._gep(ty, x, x_off, "conv_xp"), "conv_xv") + w_off = self._iadd( + self._iadd(w_oc, w_ic, "conv_woff0"), + self._iadd(w_kh, kw_ctx.value, "conv_woff1"), + "conv_woff", + ) + wv = self._load(ty, self._gep(ty, w, w_off, "conv_wp"), "conv_wv") + prod = self._bin("fmul", ty, xv, wv, "conv_pr") + old = self._load(ty, acc, "conv_acc_v") + new_acc = self._bin("fadd", ty, old, prod, "conv_acc_n") + self._p(f" store {ty} {new_acc}, {ty}* {acc}") + self._start_block(skip_label) + self._loop_close(kw_ctx) + self._loop_close(kh_ctx) + self._loop_close(ic_ctx) + + out_off = self._iadd( + self._iadd( + out_oc, self._imul(oh_ctx.value, out_w, "conv_oh_off"), + "conv_ooff0", + ), + ow_ctx.value, "conv_ooff", + ) + result = self._load(ty, acc, "conv_res") + out_ptr = self._gep(ty, out, out_off, "conv_op") + self._p(f" store {ty} {result}, {ty}* {out_ptr}") + self._loop_close(ow_ctx) + self._loop_close(oh_ctx) + self._loop_close(oc_ctx) + self._finalize_from_buffer(instr, out, ty) def _emit_reshape(self, instr: Instruction) -> None: - dst = self._dest(instr) - src = self._op(instr, 0) - ty = self._infer_type(instr) - self._p(" ; reshape: passthrough") - self._p(f" {dst} = fadd {ty} {src}, 0.0 ; reshape identity") - - # ------------------------------------------------------------------ - # Helpers - # ------------------------------------------------------------------ - - def _infer_type(self, instr: Instruction) -> str: - if instr.dest is not None: - return _llvm_type(instr.dest.dtype) - for op in instr.operands: - return _llvm_type(op.dtype) - return "float" - - def _fresh_block(self, hint: str = "block") -> str: - self._block_counter += 1 - return f"{hint}_{self._block_counter}" + """Reshape is a pointer/register alias; no instruction is emitted.""" + if instr.dest is None: + return + if not instr.operands: + raise LLVMCodegenError("reshape without operand") + val = instr.operands[0] + ref = self._value_ref(val) + ty = self._ref_types.get(ref) + if ty is None: + ty = self._value_type(val) + self._bind_dest(instr, ref, ty) # Module-level helpers -------------------------------------------------------- @@ -558,12 +1337,33 @@ def _llvm_const(val: Value) -> str: """Format an IR Value as an LLVM constant.""" if val.const_value is not None: cv = val.const_value - assert isinstance(cv, (float, int)) return _llvm_const_val(cv, _llvm_type(val.dtype)) return "0.0" +def _float_to_llvm_hex(value: float) -> str: + """Encode a Python float as an LLVM double-precision hex literal. + + The value is first rounded to float32 so it is exactly representable + in either ``float`` or ``double`` IR types. + """ + f32 = struct.unpack(" str: + """Short form for 0.0/1.0/-1.0, hex otherwise (LLVM 10 strict rules).""" + if value == 0.0: + return "0.0" + if value == 1.0: + return "1.0" + if value == -1.0: + return "-1.0" + return _float_to_llvm_hex(value) + + def _llvm_const_val(value: float | int, ty: str) -> str: - if ty in ("float", "double"): - return f"{float(value):e}" + if ty in _FLOAT_TYPES: + return _float_literal(float(value)) return str(int(value)) diff --git a/tests/test_llvm_codegen_topic16.py b/tests/test_llvm_codegen_topic16.py new file mode 100644 index 0000000..6edcc6d --- /dev/null +++ b/tests/test_llvm_codegen_topic16.py @@ -0,0 +1,689 @@ +"""Topic 16 LLVM codegen correctness tests. + +Covers SSA uniqueness, canonical loop CFG, constant/type legality, real +lowering of the 8 NN ops, plus ``llvm-as``/``lli`` integration (skipped when +the LLVM tools are not installed). +""" + +from __future__ import annotations + +import re +import shutil +import subprocess +import textwrap +from pathlib import Path + +import pytest + +from scratchv.backend.llvm_codegen import ( + LLVMCodegen, LLVMCodegenError, SSANamer, +) +from scratchv.frontend.dsl_extended import ExtendedDSLParser +from scratchv.frontend.dsl_parser import DSLParser +from scratchv.ir.builder import IRBuilder +from scratchv.ir.types import DataType + +LLVM_AS = shutil.which("llvm-as") +LLI = shutil.which("lli") +requires_asm = pytest.mark.skipif( + LLVM_AS is None, reason="llvm-as not installed" +) +requires_lli = pytest.mark.skipif( + LLI is None, reason="lli not installed" +) + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +CNN_MODEL = PROJECT_ROOT / "models" / "graph" / "cnn.onnx" + +_OPS = ( + "dot", "matmul", "gemm", "maxpool", "conv", "softmax", "gelu", "sigmoid", +) + +_LOOP_HEADERS = { + "dot": 1, + "matmul": 3, + "gemm": 3, + "maxpool": 5, + "conv": 6, + "softmax": 3, +} + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _defs(ir: str) -> list[str]: + return re.findall(r"^\s*(%[A-Za-z0-9_.]+)\s*=", ir, re.M) + + +def _labels(ir: str) -> list[str]: + return re.findall(r"^([A-Za-z0-9_.]+):", ir, re.M) + + +def _assemble(ir: str, tmp_path) -> subprocess.CompletedProcess: + ll = tmp_path / "m.ll" + ll.write_text(ir) + return subprocess.run( + [LLVM_AS, str(ll), "-o", str(tmp_path / "m.bc")], + capture_output=True, text=True, + ) + + +def _run(ir: str, tmp_path) -> subprocess.CompletedProcess: + asm = _assemble(ir, tmp_path) + assert asm.returncode == 0, asm.stderr + return subprocess.run( + [LLI, str(tmp_path / "m.bc")], capture_output=True, text=True, + ) + + +def _param(builder, func, name, shape=None): + val = builder.make_value(name) + if shape: + val.shape = tuple(shape) + func.params.append(val) + return val + + +def _build_gelu_sigmoid(): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + x = _param(b, f, "x") + b.ret(b.sigmoid(b.gelu(x))) + return b.program + + +def _build_dot(length: int = 4, shape=(4,)): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + a = _param(b, f, "a", shape) + c = _param(b, f, "b", shape) + b.ret(b.dot(a, c, length)) + return b.program + + +def _build_matmul(m: int = 2, n: int = 2, k: int = 2): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + a = _param(b, f, "a", (m, k)) + c = _param(b, f, "b", (k, n)) + res = b.matmul(a, c, m, n, k) + res.shape = (m, n) + b.ret(res) + return b.program + + +def _build_matmul_scalar(): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + a = _param(b, f, "a") + c = _param(b, f, "b") + b.ret(b.matmul(a, c, 1, 1, 1)) + return b.program + + +def _build_gemm(m: int = 2, n: int = 2, k: int = 2, trans_b: bool = False): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + a = _param(b, f, "a", (m, k)) + w_shape = (n, k) if trans_b else (k, n) + w = _param(b, f, "w", w_shape) + bias = _param(b, f, "bias", (n,)) + res = b.gemm(a, w, bias, trans_b=trans_b) + res.shape = (m, n) + b.ret(res) + return b.program + + +def _build_gemm_scalar(): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + a = _param(b, f, "a") + w = _param(b, f, "w") + bias = _param(b, f, "bias") + b.ret(b.gemm(a, w, bias)) + return b.program + + +def _build_maxpool(kernel: int = 2, stride: int = 1, shape=(1, 2, 2)): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + x = _param(b, f, "x", shape) + b.ret(b.maxpool(x, kernel, stride)) + return b.program + + +def _build_conv(out_channels=2, kernel=2, stride=1, padding=0, + shape=(1, 1, 4, 4)): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + x = _param(b, f, "x", shape) + cin = shape[-3] + h = shape[-2] + k = kernel + out_h = (h + 2 * padding - k) // stride + 1 + w = _param(b, f, "w", (out_channels, cin, kernel, kernel)) + bias = _param(b, f, "bias", (out_channels,)) + res = b.conv(x, w, bias, out_channels, kernel, stride, padding) + res.shape = (1, out_channels, out_h, out_h) + b.ret(res) + return b.program + + +def _build_conv_scalar(): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + x = _param(b, f, "x", (1, 1, 1)) + w = _param(b, f, "w", (1, 1, 1)) + bias = _param(b, f, "bias", (1,)) + b.ret(b.conv(x, w, bias, 1, 1, 1, 0)) + return b.program + + +def _build_softmax(n: int): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + x = _param(b, f, "x", (n,)) + res = b.softmax(x) + res.shape = (n,) + b.ret(res) + return b.program + + +def _build_nested_for_accumulator(): + b = IRBuilder() + b.new_function("kernel") + b.new_block("entry") + acc = b.alloca(1) + b.store(acc, b.load_const(0.0)) + outer = b.for_loop(0, 3) + b.for_loop(0, 3) + current = b.load(acc) + b.store(acc, b.add(current, outer)) + b.endfor() + b.endfor() + b.ret(b.load(acc)) + return b.program + + +# --------------------------------------------------------------------------- +# Structure tests (no external tools) +# --------------------------------------------------------------------------- + +def test_gelu_sigmoid_ssa_definitions_are_unique(): + ir = LLVMCodegen(_build_gelu_sigmoid()).emit() + defs = _defs(ir) + assert len(defs) == len(set(defs)), "duplicate SSA definitions" + assert "call float @tanhf" in ir + assert "call float @expf" in ir + assert "fneg float" in ir + assert "fdiv float 1.0" in ir + + +def test_float_constants_use_hex_encoding(): + ir = LLVMCodegen(_build_gelu_sigmoid()).emit() + assert "0x" in ir + assert not re.search(r"\d\.\d+[eE][+-]?\d+", ir), "exponent literal in IR" + + +def test_int_constants_do_not_use_float_ops(): + b = IRBuilder() + b.new_function("kernel") + b.new_block("entry") + const = b.load_const(3, dtype=DataType.INT32) + b.ret(const) + ir = LLVMCodegen(b.program).emit() + assert "add i32 0, 3" in ir + assert "fadd i32" not in ir + + +def test_nested_for_labels_unique_and_canonical_cfg(): + program = DSLParser().parse( + "for i = 0, 3\n" + "for j = 0, 3\n" + "s = add(s, i)\n" + "endfor\n" + "endfor\n" + "return s" + ) + ir = LLVMCodegen(program).emit() + labels = _labels(ir) + assert len(labels) == len(set(labels)), "duplicate block label" + headers = [label for label in labels if "_hdr" in label] + bodies = [label for label in labels if "_bdy" in label] + exits = [label for label in labels if "_ext" in label] + assert len(headers) == 2 + assert len(bodies) == 2 + assert len(exits) == 2 + # preheader must branch to the header (not the body) + first_branch = ir.index("br label %") + first_header = min(ir.index(f"{label}:") for label in headers) + assert first_branch < first_header + assert "icmp slt i32" in ir + assert "br i1" in ir + assert "sitofp i32" in ir # mixed float/i32 arithmetic is coerced + + +def _op_program(name: str): + if name == "dot": + return _build_dot() + if name == "matmul": + return _build_matmul() + if name == "gemm": + return _build_gemm() + if name == "maxpool": + return _build_maxpool() + if name == "conv": + return _build_conv() + if name == "softmax": + return _build_softmax(2) + if name == "gelu": + return _build_gelu_sigmoid() + if name == "sigmoid": + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + x = _param(b, f, "x") + b.ret(b.sigmoid(x)) + return b.program + raise AssertionError(name) + + +@pytest.mark.parametrize("name", _OPS) +def test_no_placeholder_lowering_text(name): + ir = LLVMCodegen(_op_program(name)).emit() + lowered = ir.lower() + assert "placeholder" not in lowered + assert "passthrough" not in lowered + assert "unsupported" not in lowered + assert "requires" not in lowered + + +@pytest.mark.parametrize("name", sorted(_LOOP_HEADERS)) +def test_tensor_ops_emit_geps_macs_and_loops(name): + ir = LLVMCodegen(_op_program(name)).emit() + assert "getelementptr" in ir + if name in ("dot", "matmul", "gemm", "conv"): + assert "fmul" in ir + assert "fadd" in ir + else: + assert "fcmp" in ir + assert "icmp slt i32" in ir + assert "br i1" in ir + header_count = len([x for x in _labels(ir) if "_hdr" in x]) + assert header_count == _LOOP_HEADERS[name] + + +def test_conv_and_maxpool_have_full_loop_nesting(): + conv_ir = LLVMCodegen(_build_conv()).emit() + pool_ir = LLVMCodegen(_build_maxpool()).emit() + assert len([x for x in _labels(conv_ir) if "_hdr" in x]) == 6 + assert len([x for x in _labels(pool_ir) if "_hdr" in x]) == 5 + assert "conv_mac" in conv_ir + assert "conv_skip" in conv_ir + + +def test_softmax_has_three_passes(): + ir = LLVMCodegen(_build_softmax(2)).emit() + headers = [x for x in _labels(ir) if "_hdr" in x] + assert any("sm_max" in x for x in headers) + assert any("sm_sum" in x for x in headers) + assert any("sm_div" in x for x in headers) + assert "fcmp ogt" in ir + assert "call float @expf" in ir + + +def test_target_triple_is_configurable(): + program = _build_dot() + default_ir = LLVMCodegen(program).emit() + assert "target triple" not in default_ir + riscv_ir = LLVMCodegen(program, "riscv64-unknown-elf").emit() + assert 'target triple = "riscv64-unknown-elf"' in riscv_ir + + +def test_ssanamer_sanitizes_and_tracks_definitions(): + assert SSANamer.sanitize("foo.bar") == "foo.bar" + assert SSANamer.sanitize("1bad") == "v_1bad" + assert SSANamer.sanitize("a/b") == "a_b" + assert SSANamer.sanitize("") == "v_" + namer = SSANamer() + first = namer.fresh("x") + second = namer.fresh("x") + assert first != second + namer.register_definition(first) + assert namer.registered(first) + with pytest.raises(LLVMCodegenError): + namer.register_definition(first) + + +def test_unmatched_endfor_raises(): + b = IRBuilder() + b.new_function("kernel") + b.new_block("entry") + b.endfor() + with pytest.raises(LLVMCodegenError): + LLVMCodegen(b.program).emit() + + +# --------------------------------------------------------------------------- +# llvm-as integration +# --------------------------------------------------------------------------- + +@requires_asm +@pytest.mark.parametrize("name", _OPS) +def test_asm_all_op_programs(name, tmp_path): + result = _assemble(LLVMCodegen(_op_program(name)).emit(), tmp_path) + assert result.returncode == 0, result.stderr + + +@requires_asm +def test_asm_dsl_nested_for(tmp_path): + program = DSLParser().parse( + "for i = 0, 3\n" + "for j = 0, 3\n" + "s = add(s, i)\n" + "endfor\n" + "endfor\n" + "return s" + ) + result = _assemble(LLVMCodegen(program).emit(), tmp_path) + assert result.returncode == 0, result.stderr + + +@requires_asm +def test_asm_extended_if_while(tmp_path): + source = ( + "i = add(i, 1.0)\n" + "while (i < 10):\n" + "i = add(i, 1.0)\n" + "endwhile\n" + "if (i == 10):\n" + "i = add(i, 1.0)\n" + "else:\n" + "i = sub(i, 1.0)\n" + "endif\n" + "return i" + ) + program = ExtendedDSLParser().parse(source) + result = _assemble(LLVMCodegen(program).emit(), tmp_path) + assert result.returncode == 0, result.stderr + ir = LLVMCodegen(program).emit() + assert "fcmp oeq float" in ir + assert "fcmp olt float" in ir + + +@requires_asm +def test_asm_onnx_cnn(tmp_path): + if not CNN_MODEL.exists(): + pytest.skip("cnn.onnx not found") + try: + from scratchv.frontend.onnx_parser import ONNXParser + except ImportError: # pragma: no cover + pytest.skip("onnx package not installed") + program = ONNXParser().parse(str(CNN_MODEL)) + ir = LLVMCodegen(program).emit() + result = _assemble(ir, tmp_path) + assert result.returncode == 0, result.stderr + + +@requires_asm +def test_asm_int_arithmetic_uses_integer_ops(tmp_path): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + x = _param(b, f, "x") + y = _param(b, f, "y") + x.dtype = DataType.INT32 + y.dtype = DataType.INT32 + product = b.mul(x, y) + product.dtype = DataType.INT32 + total = b.add(product, x) + total.dtype = DataType.INT32 + b.ret(total) + ir = LLVMCodegen(b.program).emit() + assert "mul i32" in ir + assert "add i32" in ir + assert "fmul i32" not in ir + assert "fadd i32" not in ir + result = _assemble(ir, tmp_path) + assert result.returncode == 0, result.stderr + + +@requires_asm +def test_asm_double_dtype_softmax(tmp_path): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + x = _param(b, f, "x", (3,)) + x.dtype = DataType.FLOAT64 + res = b.softmax(x) + res.shape = (3,) + b.ret(res) + ir = LLVMCodegen(b.program).emit() + assert "call double @exp(" in ir + assert "define double* @kernel(double* %x)" in ir + result = _assemble(ir, tmp_path) + assert result.returncode == 0, result.stderr + + +@requires_asm +def test_asm_tensor_operand_degenerates_to_first_element(tmp_path): + b = IRBuilder() + f = b.new_function("kernel") + b.new_block("entry") + a = _param(b, f, "a", (4,)) + scalar = _param(b, f, "s") + b.ret(b.add(a, scalar)) + ir = LLVMCodegen(b.program).emit() + assert "load float, float* %a" in ir + result = _assemble(ir, tmp_path) + assert result.returncode == 0, result.stderr + + +# --------------------------------------------------------------------------- +# lli numerical tests (harness: append @main calling @kernel) +# --------------------------------------------------------------------------- + +def _scalar_main(call_expr: str, expected: str) -> str: + return textwrap.dedent(f"""\ + define i32 @main() {{ + entry: + %r = {call_expr} + %ok = fcmp oeq float %r, {expected} + %rc = select i1 %ok, i32 0, i32 1 + ret i32 %rc + }} + """) + + +def _store_array(reg: str, values) -> str: + lines = [ + f" %{reg} = alloca float, i32 {len(values)}", + f" store float {values[0]}, float* %{reg}", + ] + for i in range(1, len(values)): + lines.append( + f" %{reg}{i} = getelementptr float, float* %{reg}, i32 {i}" + ) + lines.append(f" store float {values[i]}, float* %{reg}{i}") + return "\n".join(lines) + + +def _check_numeric(program, main_ir: str, tmp_path): + ir = LLVMCodegen(program).emit() + "\n" + main_ir + result = _run(ir, tmp_path) + assert result.returncode == 0, result.stderr + + +@requires_lli +def test_lli_gelu_sigmoid_zero(tmp_path): + _check_numeric( + _build_gelu_sigmoid(), + _scalar_main("call float @kernel(float 0.0)", "0.5"), + tmp_path, + ) + + +@requires_lli +def test_lli_nested_for_accumulator(tmp_path): + _check_numeric( + _build_nested_for_accumulator(), + _scalar_main("call float @kernel()", "9.0"), + tmp_path, + ) + + +@requires_lli +def test_lli_dot(tmp_path): + body = "\n".join([ + _store_array("a", ["1.0", "2.0", "3.0", "4.0"]), + _store_array("b", ["1.0", "1.0", "1.0", "1.0"]), + ]) + main = textwrap.dedent("""\ + define i32 @main() { + entry: + %s + %%r = call float @kernel(float* %%a, float* %%b) + %%ok = fcmp oeq float %%r, 10.0 + %%rc = select i1 %%ok, i32 0, i32 1 + ret i32 %%rc + } + """ % body) + _check_numeric(_build_dot(), main, tmp_path) + + +@requires_lli +def test_lli_matmul_scalar(tmp_path): + _check_numeric( + _build_matmul_scalar(), + _scalar_main("call float @kernel(float 2.0, float 3.0)", "6.0"), + tmp_path, + ) + + +@requires_lli +def test_lli_matmul_2x2(tmp_path): + body = "\n".join([ + _store_array("a", ["1.0", "2.0", "3.0", "4.0"]), + _store_array("b", ["1.0", "0.0", "0.0", "1.0"]), + ]) + main = textwrap.dedent("""\ + define i32 @main() { + entry: + %s + %%out = call float* @kernel(float* %%a, float* %%b) + %%o0 = load float, float* %%out + %%op1 = getelementptr float, float* %%out, i32 1 + %%o1 = load float, float* %%op1 + %%op2 = getelementptr float, float* %%out, i32 2 + %%o2 = load float, float* %%op2 + %%op3 = getelementptr float, float* %%out, i32 3 + %%o3 = load float, float* %%op3 + %%c0 = fcmp oeq float %%o0, 1.0 + %%c1 = fcmp oeq float %%o1, 2.0 + %%c2 = fcmp oeq float %%o2, 3.0 + %%c3 = fcmp oeq float %%o3, 4.0 + %%t1 = and i1 %%c0, %%c1 + %%t2 = and i1 %%c2, %%c3 + %%all = and i1 %%t1, %%t2 + %%rc = select i1 %%all, i32 0, i32 1 + ret i32 %%rc + } + """ % body) + _check_numeric(_build_matmul(), main, tmp_path) + + +@requires_lli +def test_lli_gemm_scalar_with_bias(tmp_path): + _check_numeric( + _build_gemm_scalar(), + _scalar_main( + "call float @kernel(float 2.0, float 3.0, float 0.5)", "6.5" + ), + tmp_path, + ) + + +@requires_lli +def test_lli_maxpool_2x2(tmp_path): + body = _store_array("x", ["1.0", "2.0", "3.0", "4.0"]) + main = textwrap.dedent("""\ + define i32 @main() { + entry: + %s + %%r = call float @kernel(float* %%x) + %%ok = fcmp oeq float %%r, 4.0 + %%rc = select i1 %%ok, i32 0, i32 1 + ret i32 %%rc + } + """ % body) + _check_numeric(_build_maxpool(), main, tmp_path) + + +@requires_lli +def test_lli_conv_1x1x1(tmp_path): + body = "\n".join([ + _store_array("x", ["3.0"]), + _store_array("w", ["2.0"]), + _store_array("bias", ["1.0"]), + ]) + main = textwrap.dedent("""\ + define i32 @main() { + entry: + %s + %%r = call float @kernel(float* %%x, float* %%w, float* %%bias) + %%ok = fcmp oeq float %%r, 7.0 + %%rc = select i1 %%ok, i32 0, i32 1 + ret i32 %%rc + } + """ % body) + _check_numeric(_build_conv_scalar(), main, tmp_path) + + +@requires_lli +def test_lli_softmax_pair(tmp_path): + body = _store_array("x", ["0.0", "0.0"]) + main = textwrap.dedent("""\ + define i32 @main() { + entry: + %s + %%out = call float* @kernel(float* %%x) + %%o0 = load float, float* %%out + %%op1 = getelementptr float, float* %%out, i32 1 + %%o1 = load float, float* %%op1 + %%c0 = fcmp oeq float %%o0, 0.5 + %%c1 = fcmp oeq float %%o1, 0.5 + %%both = and i1 %%c0, %%c1 + %%rc = select i1 %%both, i32 0, i32 1 + ret i32 %%rc + } + """ % body) + _check_numeric(_build_softmax(2), main, tmp_path) + + +@requires_lli +def test_lli_softmax_single(tmp_path): + body = _store_array("x", ["7.0"]) + main = textwrap.dedent("""\ + define i32 @main() { + entry: + %s + %%out = call float* @kernel(float* %%x) + %%o0 = load float, float* %%out + %%ok = fcmp oeq float %%o0, 1.0 + %%rc = select i1 %%ok, i32 0, i32 1 + ret i32 %%rc + } + """ % body) + _check_numeric(_build_softmax(1), main, tmp_path) From 33b8366b4615f82c07dde21dc3d8caac1e3084e2 Mon Sep 17 00:00:00 2001 From: Seven Gao <799889633@qq.com> Date: Mon, 14 Sep 2026 21:21:36 +0800 Subject: [PATCH 2/6] docs(topic16): add design and development documents --- ...00\345\217\221\346\226\207\346\241\243.md" | 709 ++++++++++++++++++ ...76\350\256\241\346\226\207\346\241\243.md" | 597 +++++++++++++++ 2 files changed, 1306 insertions(+) create mode 100644 "docs/topics/16-LLVM\344\273\243\347\240\201\347\224\237\346\210\220-\345\274\200\345\217\221\346\226\207\346\241\243.md" create mode 100644 "docs/topics/16-LLVM\344\273\243\347\240\201\347\224\237\346\210\220-\350\256\276\350\256\241\346\226\207\346\241\243.md" diff --git "a/docs/topics/16-LLVM\344\273\243\347\240\201\347\224\237\346\210\220-\345\274\200\345\217\221\346\226\207\346\241\243.md" "b/docs/topics/16-LLVM\344\273\243\347\240\201\347\224\237\346\210\220-\345\274\200\345\217\221\346\226\207\346\241\243.md" new file mode 100644 index 0000000..0b24915 --- /dev/null +++ "b/docs/topics/16-LLVM\344\273\243\347\240\201\347\224\237\346\210\220-\345\274\200\345\217\221\346\226\207\346\241\243.md" @@ -0,0 +1,709 @@ +# 课题16 LLVM 代码生成后端（库路径）开发文档 + +> 文档版本：v1.0 +> 编写日期：2026-09-14 +> 涉及模块：`scratchv/backend/llvm_codegen.py`、`tests/test_llvm_codegen.py`、`tests/test_llvm_codegen_llvm_tools.py`（新增） +> 配套文档：《设计文档.md》（同目录）——本文件是其实现指南，二者若冲突以设计文档为准 +> 前置环境：Python 3.11+（项目 venv）、`/usr/bin/llvm-as`（LLVM 10，已确认可用）、可选 `/usr/bin/lli`、`/usr/bin/opt` + +--- + +## 一、实施范围与接口契约 + +### 1.1 范围 + +**做**： + +- 修复 `llvm_codegen.py` 的 SSA 唯一性、类型/常量合法性、`_emit_for/_emit_endfor` CFG、`_emit_br_if` 条件类型； +- 为 conv/gemm/matmul/dot/maxpool/softmax 实现真实循环 + GEP + MAC；gelu/sigmoid 改为单定义展开； +- 目标 triple 可配置（构造函数参数，默认省略）； +- 扩充/新增测试，含 `llvm-as`、`lli` 集成（缺失时 skip）。 + +**不做**： + +- 不动 `scratchv/standalone/onnx_to_llvm_standalone.py` 及任何 standalone 文件； +- 不动 `scratchv/compiler.py`、`scratchv/main.py`、`scratchv/ir/*`； +- 不引入 llvmlite 或任何新依赖； +- 不做 mem2reg/循环展开/向量化等 IR→IR 优化（只要求正确可汇编）。 + +### 1.2 接口契约（精确名称） + +#### 1.2.1 模块级常量与函数（`scratchv/backend/llvm_codegen.py`） + +| 名称 | 签名 | 说明 | +|------|------|------| +| `_TYPE_MAP` | `dict[DataType, str]` | dtype→LLVM 基类型（保持） | +| `_LLVM_FLOAT` / `_LLVM_DOUBLE` / `_LLVM_I32` / `_LLVM_I64` | `str` 常量 | `"float"`/`"double"`/`"i32"`/`"i64"`（保持） | +| `_llvm_type` | `(dtype: DataType) -> str` | dtype→基类型（保持） | +| `_float_to_llvm_hex` | `(value: float) -> str` | **新增**，复制 standalone 算法：float32 舍入→double 位型→`0x%016X` | +| `_float_literal` | `(value: float) -> str` | **新增**，`0.0/1.0/-1.0` 短写，否则 `_float_to_llvm_hex` | +| `_llvm_const_val` | `(value: float \| int, ty: str) -> str` | 重写：浮点走 `_float_literal`，整型十进制 | +| `_llvm_const` | `(val: Value) -> str` | 按 `val.dtype` 调用 `_llvm_const_val`（保持签名） | +| `_is_pointer_value` | `(val: Value) -> bool` | **新增**，`bool(val.shape)` | + +#### 1.2.2 异常 + +```python +class LLVMCodegenError(Exception): + """Raised for unlowerable IR (missing operand, unmatched endfor, ...).""" +``` + +#### 1.2.3 `SSANamer`（新增类） + +```python +class SSANamer: + @staticmethod + def sanitize(name: str) -> str: ... + def fresh(self, hint: str = "r") -> str: ... # -> "%hint_" + def fresh_label(self, hint: str = "bb") -> str: ... # -> "hint_" + def register_definition(self, reg: str) -> None: ... # 重复注册即抛错 + def registered(self, reg: str) -> bool: ... +``` + +#### 1.2.4 `LoopContext`（新增 dataclass） + +```python +@dataclass +class LoopContext: + ir_name: str | None # IR 循环变量名（DSL for）；张量算子内部循环为 None + ptr: str # "%_ptr_": alloca i32* + value: str # "%_ld_": header 中 load 出的 i32 值 + header: str # 标签名（无 %） + body: str + exit: str + limit: int # 循环上界（i32） + step: int = 1 # 步长（i32） +``` + +#### 1.2.5 `LLVMCodegen` 公开 API（兼容保持不变） + +```python +class LLVMCodegen: + def __init__(self, program: Program, target_triple: str | None = None) -> None: ... + def emit(self) -> str: ... + def save(self, path: str) -> None: ... +``` + +- 位置参数 `program` 不变，`LLVMCodegen(program)` 全项目兼容（`compiler.py:390`、examples、benchmarks）； +- `target_triple=None` ⇒ 不输出 `target triple` 行；传字符串则原样输出。 + +#### 1.2.6 `LLVMCodegen` 内部 helper（实现契约，供 review/测试引用） + +| 名称 | 签名 | 职责 | +|------|------|------| +| `_fresh` | `(hint: str = "r") -> str` | 委托 `SSANamer.fresh` | +| `_fresh_label` | `(hint: str = "bb") -> str` | 委托 `SSANamer.fresh_label` | +| `_value_ref` | `(val: Value) -> str` | 常量内联；否则查/建 SSA 引用 | +| `_value_type` | `(val: Value) -> str` | 标量/指针的 LLVM 类型 | +| `_bind` | `(name: str, ref: str, llvm_ty: str) -> None` | 写 `_named_values` + `_ref_types` | +| `_dest` | `(instr: Instruction) -> str` | 幂等分配 dst 引用 | +| `_dest_buffer` | `(instr: Instruction, count: int, elem_ty: str) -> str` | dst 缓冲：ALLOCA 复用 / 按 count 分配 | +| `_op` | `(instr: Instruction, idx: int) -> str` | 第 idx 操作数引用；缺失抛 `LLVMCodegenError` | +| `_ptr_of` | `(instr: Instruction, idx: int, ty: str = "float") -> str` | 操作数当指针用；标量 spill 到 `alloca` | +| `_alloc_slot` | `(elem_ty: str, count: int = 1, hint: str = "slot") -> str` | 入口 prologue alloca，返回指针 SSA 名 | +| `_materialize_const` | `(value: float \| int, ty: str, hint: str) -> str` | 常量实体化为 SSA 值 | +| `_coerce_operand` | `(ref: str, from_ty: str, to_ty: str, hint: str) -> str` | 混型算术转换：`sitofp`/`fptosi`；同型原样返回 | +| `_emit_binary` | `(instr, op: str) -> None` | add/sub/mul/div/fadd... 统一发射（内部先 `_coerce_operand`） | +| `_start_block` | `(label: str) -> None` | 结束当前块（必要时补 `br`）并打开新块 | +| `_terminate` | `(line: str) -> None` | 发射终止指令并置位 | +| `_ensure_terminator` | `() -> None` | 当前块无终止符时补 `br` 到合成续块 | +| `_loop_open` | `(limit: int, ir_name: str \| None, hint: str, start: int = 0, step: int = 1) -> LoopContext` | 循环规范形的前半；`ir_name` 非空时绑定 IR 循环变量 | +| `_loop_close` | `(ctx: LoopContext) -> None` | 循环规范形的后半 | +| `_dim_of` | `(instr: Instruction, keys: tuple[str, ...], default: int = 1, operand: int \| None = None, axis: int \| None = None) -> int` | 从 attrs/shape 取维度，键名兼容 | + +#### 1.2.7 CLI（不变） + +```bash +scratchv model.onnx --backend llvm -o out.ll # 既有入口，行为=默认目标无关 IR +python -c "from scratchv.backend.llvm_codegen import LLVMCodegen; \ + open('o.ll','w').write(LLVMCodegen(p, 'riscv64-unknown-elf').emit())" # 显式 triple +``` + +不新增命令行参数；triple 覆盖只走 Python API（范围约束）。 + +--- + +## 二、通用机制实现方案 + +### 2.1 张量表示与 slot 分配 + +**判定**：`_is_pointer_value(val)` 为真 ⇔ `val.shape` 非空；此外 `ALLOCA` 指令的 dest 显式绑定指针类型。 + +**指针类型表**（`_value_type`）： + +```python +def _value_type(self, val) -> str: + base = _llvm_type(val.dtype) + if _is_pointer_value(val): + return base + "*" + ref = self._named_values.get(val.name) + if ref is not None and self._ref_types.get(ref, "").endswith("*"): + return self._ref_types[ref] + return base +``` + +**入口 prologue**：`_alloc_slot()` 把 `%p = alloca , i32 ` 追加到 `self._prologue: list[str]`；`_emit_function` 在 `define ... {` 之后、第一个 `_emit_block` 之前输出 prologue（这些指令自动属于 entry 块）。这样循环内的 alloca 不会随迭代增长栈帧（`lli` 数值测试必需）。 + +**标量 spill**（`_ptr_of`）： + +```python +def _ptr_of(self, instr, idx, ty="float"): + val = instr.operands[idx] + ref = self._op(instr, idx) + if self._ref_types.get(ref, "").endswith("*"): + return ref + p = self._alloc_slot(ty, 1, "spin") # 退化 1 元素张量 + self._p(f" store {ty} {ref}, {ty}* {p}") + return p +``` + +**结果缓冲**（`_dest_buffer`）：`ALLOCA` dest 直接返回其指针；`dest.shape` 非空返回 `_alloc_slot(elem_ty, prod(shape))` 并把 dest 名绑定该指针；否则分配 1 元素缓冲，算子执行完由调用方 `load` 出标量（见 3.2 通用尾巴）。 + +### 2.2 SSA 命名器 + +```python +class SSANamer: + def __init__(self): + self._reg_n = 0 + self._label_n = 0 + self._defs: set[str] = set() + + @staticmethod + def sanitize(name: str) -> str: + s = "".join(ch if ch.isalnum() or ch == "_" else "_" for ch in name) + if not s or s[0].isdigit(): + s = "v_" + s + return s + + def fresh(self, hint: str = "r") -> str: + self._reg_n += 1 + return f"%{self.sanitize(hint)}_{self._reg_n}" + + def fresh_label(self, hint: str = "bb") -> str: + self._label_n += 1 + return f"{self.sanitize(hint)}_{self._label_n}" + + def register_definition(self, reg: str) -> None: + if reg in self._defs: + raise LLVMCodegenError(f"duplicate SSA definition: {reg}") + self._defs.add(reg) + + def registered(self, reg: str) -> bool: + return reg in self._defs +``` + +要点： + +- `_dest()` 幂等：若 `instr.dest.name` 已在 `_named_values` 中，直接返回旧引用，**不**调用 `_fresh`； +- 所有 `= ...` 发射点（含 `_materialize_const`、`_emit_*` 内部中间值）必须用 `_fresh` 并通过 `register_definition` 登记； +- 参数名不登记（参数定义不在函数体内），使用前先 `_bind`； +- `gelu/sigmoid` 等展开算子禁止 `_dest` 复用作为多行左值。 + +### 2.3 类型处理与常量 + +**分派规则**： + +| 指令 | 类型来源 | 生成 | +|------|----------|------| +| `fadd/fsub/fmul/fdiv/fneg/fcmp` | `_infer_type(instr)`（dst 优先，其次首操作数） | 浮点指令 | +| `load` | `dest.dtype`（含指针判定） | `%d = load , * %p` | +| `store` | `operands[1].dtype`（值类型） | `store %v, * %p` | +| `load_const` | `dest.dtype` | 浮点：`fadd , 0.0`；整型：`add 0, ` | +| `alloca` | `dest.dtype`，`attrs["size"]` 默认 4 | `%d = alloca , i32 ` | +| `for` | 计数固定 `i32` | 见第五章 | +| `return` | `_ref_types[ref]` 优先，其次操作数类型 | `ret float* %buf` 等 | +| 函数返回类型 | 首个 `RETURN` 操作数的 `_value_type` | 与 `ret` 一致 | + +**常量编码**（P5/P6 修复）： + +```python +def _float_to_llvm_hex(value: float) -> str: + f32 = struct.unpack(" str: + if value == 0.0: return "0.0" + if value == 1.0: return "1.0" + if value == -1.0: return "-1.0" + return _float_to_llvm_hex(value) + +def _llvm_const_val(value, ty): + if ty in ("float", "double"): + assert isinstance(value, (int, float)) + return _float_literal(float(value)) + assert isinstance(value, int) or float(value).is_integer() + return str(int(value)) +``` + +**不变量**：浮点指令的参数串中不出现十进制指数形式；整型指令的参数串中不出现小数点。建议在 `_emit_binary`/`_emit_load_const` 后加开发期断言（`if ty in ("float","double"): assert "." in lit or "0x" in lit`）。 + +**混型算术协调**（DSL `for` 循环体常见 `s = add(s, i)`，`s: float`、`i: i32`）： + +```python +def _coerce_operand(self, ref: str, from_ty: str, to_ty: str, hint: str) -> str: + if from_ty == to_ty: + return ref + r = self._fresh(hint) + self.namer.register_definition(r) + if from_ty == "i32" and to_ty in ("float", "double"): + self._p(f" {r} = sitofp i32 {ref} to {to_ty}") + elif from_ty in ("float", "double") and to_ty == "i32": + self._p(f" {r} = fptosi {from_ty} {ref} to i32") + else: + raise LLVMCodegenError(f"no coercion {from_ty} -> {to_ty}") + return r + +def _emit_binary(self, instr, fop: str, iop: str | None = None): + ty = self._infer_type(instr) # 结果类型 + lhs = self._coerce_operand(self._op(instr, 0), + self._type_of_operand(instr, 0), ty, "cvt") + rhs = self._coerce_operand(self._op(instr, 1), + self._type_of_operand(instr, 1), ty, "cvt") + dst = self._dest(instr) + op = fop if ty in ("float", "double") else iop + self._p(f" {dst} = {op} {ty} {lhs}, {rhs}") +``` + +`_type_of_operand` 返回操作数引用的 LLVM 类型（常量按 `val.dtype`）；`ty` 优先取 `dest.dtype`，dest 缺失时取浮点操作数类型。 + +### 2.4 控制流状态机 + +新增字段：`_terminated: bool`、`_defined_labels: set[str]`、`_prologue: list[str]`、`_loop_stack: list[LoopContext]`。 + +```python +def _start_block(self, label: str) -> None: + if label in self._defined_labels: + label = self._fresh_label(label) # 绝不重复 + if not self._terminated: + self._p(f" br label %{label}") # 显式落空转跳转 + self._p(f"{label}:") + self._defined_labels.add(label) + self._terminated = False + +def _terminate(self, line: str) -> None: + if self._terminated: + raise LLVMCodegenError("terminator emitted twice in one block") + self._p(f" {line}") + self._terminated = True + +def _ensure_terminator(self) -> None: + if not self._terminated: + self._p(f" br label %{self._fresh_label('cont')}") + self._terminated = True +``` + +- `_emit_block` 改为：首块不打印标签（entry 隐式），其余块 `_start_block(block.name)`；块注释保留； +- `_emit_instruction` 若在 `_terminated` 状态下收到非终止指令，先 `_start_block(self._fresh_label("dead"))`，保证不产生“终止符后接指令”的非法文本； +- `_emit_br/_emit_br_if/_emit_return` 全部改调 `_terminate`。 + +### 2.5 `_emit_for` / `_emit_endfor` 修复方案 + +**修复目标**：preheader 跳 header 而非 body；循环上下文栈化；IV 有定义；嵌套标签不重复不悬空（修复 P3/P4）。 + +```python +def _emit_for(self, instr): + self._dest(instr) # 占位登记，值随后绑定 + limit = int(instr.attrs.get("end", 0)) + start = int(instr.attrs.get("start", 0)) + step = int(instr.attrs.get("step", 1)) + ctx = self._loop_open(limit, instr.dest.name, "loop_i", start, step) + self._loop_stack.append(ctx) + +def _loop_open(self, limit, ir_name, hint, start=0, step=1): + ptr = self._alloc_slot("i32", 1, "iv_ptr") + self._p(f" store i32 {start}, i32* {ptr}") + header = self._fresh_label(f"{hint}_hdr") + body = self._fresh_label(f"{hint}_bdy") + exit_ = self._fresh_label(f"{hint}_ext") + self._start_block(header) # 自动补 br label %header + iv = self._fresh(f"{hint}_ld") + self._p(f" {iv} = load i32, i32* {ptr}") + self.namer.register_definition(iv) + cond = self._fresh(f"{hint}_cond") + self._p(f" {cond} = icmp slt i32 {iv}, {limit}") + self._terminate(f"br i1 {cond}, label %{body}, label %{exit_}") + self._start_block(body) + if ir_name is not None: # 仅 DSL for 需要绑定 IR 循环变量 + self._bind(ir_name, iv, "i32") # header load 支配 body+exit + return LoopContext(ir_name, ptr, iv, header, body, exit_, limit, step) + +def _loop_close(self, ctx): + if not self._loop_stack or self._loop_stack[-1] is not ctx: + raise LLVMCodegenError("endfor without matching for") + self._loop_stack.pop() + if not self._terminated: # body 以 ret 结束时为死块，跳过回边 + cur = self._fresh("iv_cur") + self._p(f" {cur} = load i32, i32* {ctx.ptr}") + self.namer.register_definition(cur) + nxt = self._fresh("iv_nxt") + self._p(f" {nxt} = add i32 {cur}, {ctx.step}") + self.namer.register_definition(nxt) + self._p(f" store i32 {nxt}, i32* {ctx.ptr}") + self._terminate(f"br label %{ctx.header}") + self._start_block(ctx.exit) + +def _emit_endfor(self, instr): + if not self._loop_stack: + raise LLVMCodegenError("endfor without matching for") + self._loop_close(self._loop_stack[-1]) +``` + +**边界**： + +- 循环变量在 `endfor` 之后的引用：绑定的是 header 的 load，仍支配 exit，值等于终止时的 IV（合法且语义可解释）； +- 空循环体：body 与 header 同名结构仍合法；`_ensure_terminator` 保证 body 有回边； +- `for` 出现在已终止块之后：`_start_block(header)` 会先补 `br`，但这时其实应已由 `_emit_instruction` 开了 `dead` 合成块，无需额外处理。 + +### 2.6 `_emit_br_if` 修复方案（P7） + +```python +_CMP_PRED = {"==": ("oeq", "eq"), "!=": ("one", "ne"), + "<": ("olt", "slt"), "<=": ("ole", "sle"), + ">": ("ogt", "sgt"), ">=": ("oge", "sge")} + +def _emit_br_if(self, instr): + targets = (instr.target or ",").split(",") + true_t = targets[0].strip() + false_t = targets[1].strip() if len(targets) > 1 else true_t + cmp_op = instr.attrs.get("cmp_op") + if cmp_op and len(instr.operands) >= 2: + lhs, rhs = self._op(instr, 0), self._op(instr, 1) + ty = self._infer_type(instr) + fpred, ipred = self._CMP_PRED[str(cmp_op)] + pred = fpred if ty in ("float", "double") else ipred + kind = "fcmp" if ty in ("float", "double") else "icmp" + cond = self._fresh("brc") + self._p(f" {cond} = {kind} {pred} {ty} {lhs}, {rhs}") + self.namer.register_definition(cond) + else: + cond = self._op(instr, 0) # 已是 i1 + self._terminate(f"br i1 {cond}, label %{true_t}, label %{false_t}") +``` + +映射覆盖 `ExtendedDSLParser._parse_condition` 的六种运算符；`IRBuilder.br_if` 的 `operands=[cond]` 形态不受影响。 + +--- + +## 三、逐算子实现方案 + +统一约定：`acc`/`max`/`sum` 等所有 alloca 走 `_alloc_slot`（入口 prologue）；循环走 `_loop_open/_loop_close`；`hint` 见设计文档 2.3 表；结果写入 `_dest_buffer` 得到的缓冲，标量 dest 在算子末尾 `load` 首元素。 + +### 3.1 dot + +```python +def _emit_dot(self, instr): + ty = self._infer_type(instr) # "float"/"double" + a = self._ptr_of(instr, 0, ty) + b = self._ptr_of(instr, 1, ty) + n = int(instr.attrs.get("length", instr.attrs.get("len", 1))) + acc = self._alloc_slot(ty, 1, "dot_acc") + self._p(f" store {ty} 0.0, {ty}* {acc}") + ctx = self._loop_open(n, None, "dot_i") + ap = self._gep(ty, a, ctx.value, "dot_ap") + av = self._load(ty, ap, "dot_av") + bp = self._gep(ty, b, ctx.value, "dot_bp") + bv = self._load(ty, bp, "dot_bv") + pr = self._bin("fmul", ty, av, bv, "dot_pr") + old = self._load(ty, acc, "dot_old") + nw = self._bin("fadd", ty, old, pr, "dot_new") + self._p(f" store {ty} {nw}, {ty}* {acc}") + self._loop_close(ctx) + self._finish_scalar_result(instr, acc, ty) # 见 3.2 +``` + +需要的微 helper（也供其余算子复用）： + +```python +def _gep(self, ty, base, idx, hint): # %r = getelementptr ty, ty* base, i32 idx +def _load(self, ty, ptr, hint): # %r = load ty, ty* ptr +def _bin(self, op, ty, lhs, rhs, hint): # %r = op ty lhs, rhs +``` + +### 3.2 结果收尾与维度解析 + +```python +def _finish_scalar_result(self, instr, buf_ptr, ty): + if _is_pointer_value(instr.dest): # 张量 dest：绑定缓冲指针 + self._bind(instr.dest.name, buf_ptr, ty + "*") + else: # 标量 dest：load 首元素 + r = self._load(ty, buf_ptr, "res") + self._bind(instr.dest.name, r, ty) + +def _dim_of(self, instr, keys, default=1, operand=None, axis=None): + for k in keys: + v = instr.attrs.get(k) + if isinstance(v, (int, float)) and int(v) > 0: + return int(v) + if operand is not None and operand < len(instr.operands): + shape = instr.operands[operand].shape + if shape: + idx = axis if axis is not None else 0 + if len(shape) > abs(idx): + return int(shape[idx]) + return default +``` + +### 3.3 matmul + +```python +def _emit_matmul(self, instr): + ty = self._infer_type(instr) + a = self._ptr_of(instr, 0, ty) + b = self._ptr_of(instr, 1, ty) + m = self._dim_of(instr, ("m", "rows"), 1, operand=0, axis=0) + k = self._dim_of(instr, ("k", "inner"), 1, operand=0, axis=1) + n = self._dim_of(instr, ("n", "cols"), 1, operand=1, axis=1) + c = self._dest_buffer(instr, m * n, ty) + acc = self._alloc_slot(ty, 1, "mm_acc") + ci = self._loop_open(m, None, "mm_i") + cj = self._loop_open(n, None, "mm_j") + self._p(f" store {ty} 0.0, {ty}* {acc}") + ck = self._loop_open(k, None, "mm_k") + # --- innermost body (i32 offsets) --- + # aoff = add (mul ci.value, k), ck.value + # av = load(gep(a, aoff)) + # boff = add (mul ck.value, n), cj.value + # bv = load(gep(b, boff)) + # pr = fmul av, bv ; acc = fadd load(acc), pr ; store acc + self._loop_close(ck) + # coff = add (mul ci.value, n), cj.value ; store load(acc) -> gep(c, coff) + self._loop_close(cj) + self._loop_close(ci) + self._finish_scalar_result(instr, c, ty) +``` + +（文档中的缩进仅表嵌套层次；实现时按 `_loop_open/_loop_close` 顺序配对。） + +### 3.4 gemm + +维度：`M=_dim_of(attrs ("M",), A.shape[0])`、`K=A.shape[1]`；`trans_b = bool(attrs.get("trans_b", attrs.get("transB", False)))`；`N = W.shape[0] if trans_b else W.shape[1]`，退化 1。循环体： + +```python +acc = bias[j] # load(gep(bias, cj.value)) +for kk: + av = A[i*K + kk] + woff = j*K + kk if trans_b else kk*N + j + wv = W[woff] + acc += av*wv +C[i*N + j] = acc +``` + +`bias` 为标量时 `_ptr_of` spill；`attrs` 中 `trans_a` 暂不支持（如为真，抛 `LLVMCodegenError` 并注释说明）。 + +### 3.5 conv + +维度解析优先级：`attrs` → `operand.shape` → 退化值。 + +```python +x = self._ptr_of(instr, 0, ty); w = self._ptr_of(instr, 1, ty); bias = self._ptr_of(instr, 2, ty) +xs = instr.operands[0].shape # 取后三维 (C,H,W) +cin = xs[-3] if len(xs) >= 3 else 1 +h = xs[-2] if len(xs) >= 3 else 1 +ww = xs[-1] if len(xs) >= 3 else 1 +ws = instr.operands[1].shape +cout = ws[0] if len(ws) >= 4 else _dim_of(instr, ("out_channels",), 1) +k = ws[2] if len(ws) >= 4 else _dim_of(instr, ("kernel_size", "kernel_shape"), 3) +s = _dim_of(instr, ("stride", "strides"), 1) +p = _dim_of(instr, ("padding", "pads"), 0) +ho = (h + 2*p - k)//s + 1; wo = (ww + 2*p - k)//s + 1 +# 若 len(instr.operands) < 3（无 bias），bias 用 0.0 常量替代（_materialize_const），不调用 _ptr_of +``` + +循环体（6 层，oc/oh/ow/ic/kh/kw）：每层用 `_loop_open`；kh 内计算 `ih`、`ok_h`；kw 内计算 `iw`、`ok`，并用 `br i1 ok, label %mac, label %skip` 包住 MAC；`mac`/`skip` 用 `_fresh_label("conv_mac"/"conv_skip")` 并 `_start_block`。MAC 地址：输入 `ic*H*W + ih*W + iw`，权重 `oc*Cin*K*K + ic*K*K + kh*K + kw`；kw 退出后写 `out[oc*Ho*Wo + oh*Wo + ow] = acc`。 + +### 3.6 maxpool + +`C/H/W` 来自 `operands[0].shape`（取后三维，退化 1）；`k=_dim_of(("kernel","kernel_shape"),2)`、`s=_dim_of(("stride","strides"),2)`；`ho=(h-k)//s+1`、`wo=(w-k)//s+1`。5 层循环；每 (c,oh,ow) 初始化 `m=-3.4e38`（`_float_literal`），内两层 `fcmp ogt + select` 更新，循环结束写回。 + +### 3.7 softmax + +`n = _dim_of(instr, ("length", "n"), 1, operand=0, axis=-1)`；输出缓冲 `_dest_buffer(instr, n, ty)`。 + +- pass1（标签 `sm_max`）：`m=-3.4e38`；`m = select(fcmp ogt x[i], m, x[i], m)`； +- pass2（`sm_sum`）：`s += expf(x[i] - m)`；`fsub` 后 `call @expf`（f64 用 `@exp`）； +- pass3（`sm_div`）：`out[i] = expf(x[i] - m) / s`； +- 三趟各自独立 `_loop_open/_loop_close`；结果收尾同 3.2。 + +### 3.8 gelu + +按设计文档 2.5.7 的九行展开实现，全部 `_fresh` + `register_definition`；常量 `0.044715`、`0.7978845608028654` 经 `_float_literal` 内联；`float` 调 `@tanhf`，`double` 调 `@tanh`；最后一行 `_dest` 只定义一次并 `_finish_scalar_result` 绑定。 + +### 3.9 sigmoid + +按设计文档 2.5.8 的四行展开实现；`float` 调 `@expf`，`double` 调 `@exp`。 + +### 3.10 reshape / relu / exp / 算术 + +- `reshape`：同 dtype 直通，`_bind(dest.name, _value_ref(src), _value_type(src))`，不发指令（当前 `fadd x, 0.0` 保留也可，但要求整型安全：按类型选 `fadd`/`add`）； +- `relu`：`fcmp ogt + select`（保持）； +- `exp`：`call @expf/@exp`（保持）； +- `add/sub/mul/div/neg`：`_emit_binary` 统一按 `_infer_type` 分派浮点/整型指令（整型用 `add/sub/mul/sdiv`），避免 int 走 `fadd`。 + +--- + +## 四、测试文件与用例 + +### 4.1 文件清单 + +| 文件 | 内容 | +|------|------| +| `tests/test_llvm_codegen.py`（扩充） | 纯 Python 单元断言：SSA 唯一、标签唯一、常量写法、8 算子结构关键词 | +| `tests/test_llvm_codegen_llvm_tools.py`（新增） | `llvm-as` 可汇编测试 + `lli` 数值测试；工具缺失 skip | + +### 4.2 `tests/test_llvm_codegen.py` 扩充用例（无外部工具依赖） + +```python +def _name_defs(ir: str) -> list[str]: + return re.findall(r"^\s*(%[A-Za-z0-9_.]+)\s*=", ir, re.M) + +def test_no_duplicate_ssa_gelu_sigmoid(...): # set(defs) 数量 == len(defs) +def test_for_labels_unique(...): # 每个 label 形如 ^\s*(\w+):$ 唯一 +def test_float_const_is_hex(...): # 无 "e-0" 指数常量；含 0x +def test_int_const_not_float_op(...): # 无 "fadd i32" +def test_all_ops_emit_loops(...): # 8 算子每个含 "getelementptr"/"icmp slt" +``` + +### 4.3 `llvm-as` 集成测试（环境缺失 skip） + +```python +import shutil, subprocess, pytest + +LLVM_AS = shutil.which("llvm-as") +LLI = shutil.which("lli") +requires_asm = pytest.mark.skipif(LLVM_AS is None, reason="llvm-as not installed") +requires_lli = pytest.mark.skipif(LLI is None, reason="lli not installed") + +def _assemble(ir: str, tmp_path): + ll = tmp_path / "m.ll"; ll.write_text(ir) + return subprocess.run([LLVM_AS, str(ll), "-o", str(tmp_path / "m.bc")], + capture_output=True, text=True) +``` + +用例矩阵（对应设计文档第三章）： + +| 用例 | 输入 | 断言 | +|------|------|------| +| `test_asm_gelu_sigmoid` | IRBuilder 展开 | rc==0；无 `multiple definition`；定义次数唯一 | +| `test_asm_for_nested` | DSLParser 双重 for | rc==0；`loop_i_hdr` 先于 body；标签计数为 1 | +| `test_asm_tensor_ops` | dot/matmul/gemm/maxpool/conv 五程序 | rc==0；GEP/MAC/循环结构 | +| `test_asm_softmax` | shape=(2,) 与 (1,) | rc==0；三趟循环标签齐全 | +| `test_asm_onnx_cnn` | `models/graph/cnn.onnx`（存在则跑） | rc==0；无占位注释 | + +### 4.4 `lli` 数值测试（harness 拼接法） + +测试端构造程序：被测函数命名 `kernel`，参数用 `IRBuilder.make_value` 后设 `shape` 得到 `float*`；生成后拼接手写 `@main`： + +```llvm +define i32 @main() { + %a = alloca float, i32 4 + ; store 常量数组 ... + %r = call float @kernel(float* %a, float* %b) + %ten = sitofp i32 10 to float ; 期望值 10.0，避免十进制浮点字面量 + %ok = fcmp oeq float %r, %ten + %rc = select i1 %ok, i32 0, i32 1 + ret i32 %rc +} +``` + +然后用 `LLI m.bc` 执行并断言 `returncode == 0`。用例与期望值： + +| 算子 | 输入 | 期望 | +|------|------|------| +| gelu→sigmoid | x=0.0 | 0.5 | +| 嵌套 for | i,j∈[0,3)，`s+=i` | 9.0 | +| dot | a=[1,2,3,4], b=[1,1,1,1] | 10.0 | +| matmul 1×1 | a=2, b=3 | 6.0 | +| gemm 1×1 | a=2, w=3, bias=0.5 | 6.5 | +| maxpool 2×2 | x=[1,2,3,4], K=2,S=1 | 4.0 | +| conv 1×1×1 | x=3, w=2, bias=1, K=1,S=1,P=0 | 7.0 | +| softmax N=2 | x=[0,0] | out=[0.5,0.5] | +| softmax N=1 | x=[7] | out=[1.0] | + +注意：这些期望值全部可被 float32 精确表示，用 `fcmp oeq` 判定不引入容差问题。若 `lli` 存在但符号解析失败（`@expf` 未注册等），把对应激活类用例标记为 `xfail(strict=False)` 并在测试注释中说明；张量类用例不依赖 libm，必须真实执行。 + +### 4.5 回归清单 + +```bash +python -m pytest tests/test_llvm_codegen.py tests/test_llvm_codegen_llvm_tools.py -v +python -m pytest tests/ -q # 全量，无新增失败 +llvm-as /tmp/out.ll -o /dev/null # 手工冒烟（cnn.onnx 输出） +``` + +--- + +## 五、验收标准 + +1. **可汇编（硬门槛）**：8 个算子的样例 IR、DSL `for/if/while` 样例、`models/graph/cnn.onnx` 经 `LLVMCodegen.emit()` 后 `llvm-as` 全部退出码 0；输出中不存在 `; UNSUPPORTED`、`; placeholder`、`; passthrough`。 +2. **结构断言**： + - 每个函数内 SSA 定义名唯一（正则计数 == set 大小）； + - 每个标签唯一定义，且 preheader 分支目标为 header； + - 每个基本块以 `br`/`ret` 结尾； + - 张量算子至少含 `getelementptr , *`、`fmul`+`fadd`、`icmp slt i32`、`br i1`；循环层数符合：conv 6、maxpool 5、gemm/matmul 3、dot 1、softmax 3 趟。 +3. **数值断言**：4.4 表格全部通过（`lli` 环境存在时必须真实执行；缺失则 skip，不得以 skip 充当通过）。 +4. **类型断言**：浮点常量全部为 `0x` 十六进制或 `0.0/1.0/-1.0`；整型常量不出现在浮点指令中。 +5. **兼容性**：`tests/test_llvm_codegen.py` 原有用例全绿；`benchmarks/test_benchmark.py::test_codegen_llvm` 通过；`LLVMCodegen(program)` 单参数调用可用。 +6. **范围纪律**：`git diff --stat` 仅包含 `scratchv/backend/llvm_codegen.py` 与两个测试文件。 + +--- + +## 六、风险与回退 + +| 风险 | 概率/影响 | 缓释 | 回退 | +|------|-----------|------|------| +| 标量 DSL 调用张量算子只有退化 1 元素语义（如 `dot(a,b,len:8)` 中 a/b 是标量） | 高/中 | 文档明确“degenerate 1-element tensor”；测试用 IRBuilder 显式 `shape` 构造真张量；不做语义欺骗 | 如评审要求完整向量语义，需扩 IR（超范围，另立课题） | +| ONNX 路径参数变 `float*`、返回 `float*` 改变模块形态 | 中/中 | `benchmarks` 仅断言非空；无其他测试依赖具体签名 | 保持标量签名 + 内部 spill，放弃指针优化（改回点在 `_value_type`） | +| `_start_block` 自动补 `br` 改变 IR 文本，既有结构断言用例误判 | 中/低 | 同步更新 `tests/test_llvm_codegen.py` 中 `": " in ir` 类弱断言 | 保留旧的 `_emit_block` 分支作为开关（不推荐） | +| `lli` 数值测试受 libm 符号/平台影响 | 中/低 | 张量用例零 libm；激活用例 `xfail(strict=False)` | 数值验证降级为“结构 + 常量折叠冒烟” | +| LLVM 版本差异（十进制浮点、attr 语法） | 低/高 | 全部浮点常量走十六进制；外部声明只用 Llvm 10 已有语法 | 以 `/usr/bin/llvm-as` 实测为准，修正编码 | +| 重构范围蔓延到 optimizer/backend 其他文件 | 低/高 | 实施顺序按第七节，单文件提交；scope 检查用 `git diff --stat` | 按文件粒度 `git checkout --` 回退非目标文件 | + +整体回退策略：本课题所有改动集中在 `llvm_codegen.py` 与测试；如集成失败，`git revert` 对应提交即可恢复旧行为（占位实现仍可汇编出错误数值的 IR，不影响 RISC-V 主路径）。 + +--- + +## 七、实施顺序（建议提交粒度） + +1. **命名器 + 类型/常量**：`SSANamer`、`_float_literal`、`_llvm_const_val`、`_emit_load_const`、`_value_type`/`_ref_types`；跑通 `gelu/sigmoid` 的 `llvm-as`（修 P1/P2/P5/P6）。 +2. **控制流状态机**：`_start_block`/`_terminate`/`_loop_open`/`_loop_close`、`_emit_br_if` cmp_op；跑通 `for`（含嵌套）与 `if/while` 的 `llvm-as`（修 P3/P4/P7）。 +3. **基础设施**：`_alloc_slot` prologue、`_ptr_of`、`_dest_buffer`、`_gep/_load/_bin`。 +4. **张量算子**：dot → matmul → gemm → maxpool → conv → softmax，逐个补 `llvm-as` 与 `lli` 用例。 +5. **triple 可选化** + 文档注释清理（删除 placeholder 文本）。 +6. **测试收口**：全量 pytest、`llvm-as` ONNX 冒烟、验收清单逐条勾选。 + +每步结束运行：`python -m pytest tests/test_llvm_codegen.py -q && llvm-as <样例> -o /dev/null`；完成后按项目 Harness 规则执行 self-review 与验证（本任务仅编写文档，不改仓库、不做 git 操作）。 + +--- + +## 实现结果（2026-09-14 集成） + +> **集成 commit**：`bd0db49`（`feat(topic16): implement real NN op lowering and fix invalid LLVM IR`） +> **集成位置**：`Seven_big_summary` 上第 4 个 topic commit（顺序 06 → 07 → 09 → **16** → 17 → …） +> **集成后全量**：`PYTHONPATH=. python3.11 -m pytest tests/ -q` → **1011 passed / 13 xfailed / 20 xpassed / 0 failed** + +### 实现文件与要点 + +| 文件 | 要点 | +|------|------| +| `scratchv/backend/llvm_codegen.py` | 重写（+1121 / −321）：SSA 命名、类型/常量、控制流状态机、8 个 NN 算子真实 lowering | +| `tests/test_llvm_codegen_topic16.py` | 47 个新测试（含 `lli` 数值断言；外部工具缺失时 skip） | + +### 测试数字 + +| 口径 | 结果 | +|------|------| +| 定向（`tests/test_llvm_codegen_topic16.py`） | 47 passed（含 lli 数值） | +| 分支全量（cherry-pick 前） | 612 passed | +| 集成后全量 | 1011 passed / 13 xfailed / 20 xpassed / 0 failed | + +### 与本文档的偏差 / 未完成项 + +- 标量变量统一使用 alloca + load/store（不再区分寄存器直出形态）。 +- DSL 嵌套 `for` 的数值断言只做结构 + 可汇编，未做 `lli` 数值。 +- 多元素 initializer 在无 `shape` 信息时退化为 1 元素。 +- `gemm trans_a=True` 明确抛错（不支持）。 + +### 已知限制 + +- `llvm-as` / `lli` 环境缺失时相关用例 skip，不得以 skip 充当通过。 +- ONNX 路径 `float*` 签名与返回类型变化仍按文档风险表处理；无其他调用方依赖具体签名。 diff --git "a/docs/topics/16-LLVM\344\273\243\347\240\201\347\224\237\346\210\220-\350\256\276\350\256\241\346\226\207\346\241\243.md" "b/docs/topics/16-LLVM\344\273\243\347\240\201\347\224\237\346\210\220-\350\256\276\350\256\241\346\226\207\346\241\243.md" new file mode 100644 index 0000000..a61d377 --- /dev/null +++ "b/docs/topics/16-LLVM\344\273\243\347\240\201\347\224\237\346\210\220-\350\256\276\350\256\241\346\226\207\346\241\243.md" @@ -0,0 +1,597 @@ +# 课题16 LLVM 代码生成后端（库路径）技术设计文档 + +> 文档版本：v1.0 +> 编写日期：2026-09-14 +> 涉及模块：`scratchv/backend/llvm_codegen.py`（库路径 LLVM IR 生成器）、`tests/test_llvm_codegen*.py` +> 功能范围：库路径 `--backend llvm` 的正确性修复与算子补全——SSA 唯一命名、类型/常量合法化、`_emit_for` CFG 修复、控制流 `br_if` 修复、目标 triple 可配置、8 个 NN 算子（conv/gemm/matmul/dot/maxpool/softmax/gelu/sigmoid）的真实循环生成；验收工具 `llvm-as` +> 状态说明：本文档描述**目标实现**。当前 `llvm_codegen.py` 存在下述已实测缺陷，文档给出修复设计；实现前代码未修改。 + +--- + +## 一、功能介绍 + +### 1.1 功能概述 + +LLVM 代码生成后端负责把 ScratchV IR（`Program`/`Function`/`BasicBlock`/`Instruction`）翻译为 LLVM IR 文本（`.ll`），供 `llvm-as`、`opt`、`llc`、`lli` 使用。项目有两条 LLVM 路径： + +| | 路径 A（库路径，本课题） | 路径 B（Standalone，不修改） | +|---|---|---| +| 文件 | `scratchv/backend/llvm_codegen.py` | `scratchv/standalone/onnx_to_llvm_standalone.py` | +| 输入 | ScratchV IR `Program` | ONNXModel（手工解析） | +| 算子状态 | conv/softmax/dot 等为占位实现，数值错误 | float32 完整循环生成 | +| 目标 | 正确、可汇编、结构可断言 | 完整可执行 CNN | + +#### 当前缺陷清单（已在 LLVM 10 / `llvm-as` 下实测复现） + +| 编号 | 缺陷 | 触发输入 | 实测结果 | +|------|------|----------|----------| +| P1 | GELU 重复定义同名 SSA | `y = gelu(x)` | `multiple definition of local value named 'x3_2'` | +| P2 | Sigmoid 重复定义同名 SSA | `sigmoid` 算子 | `multiple definition of local value named 'v_1_1'` | +| P3 | for 循环变量从未定义 | `for i = 0, 3` + 循环体引用 `i` | `use of undefined value '%v_1_1'` | +| P4 | 嵌套 for 标签重复/悬空 | 双重 `for` | 标签 `loop_exit_12` 重复定义且 `%loop_exit_4` 未定义；`_loop_context` 单槽被内层覆盖 | +| P5 | 整数常量走浮点指令 | `load_const(3, INT32)` | `%v = fadd i32 3, 0.0` → `floating point constant invalid for type` | +| P6 | 十进制浮点字面量不可表示 | ONNX initializer 如 `-6.987718e-02` | LLVM 10 要求 float 十进制常量**精确可表示**；`fadd float -6.987718e-02, 0.0` 报错 | +| P7 | `br_if` 条件类型错误 | `ExtendedDSLParser` 的 `if/while` | `br i1 %a`（`%a: float`）→ `'%a' defined with type 'float' but expected 'i1'`；`cmp_op` 属性被忽略 | +| P8 | 6 个张量算子占位、数值错误 | conv/gemm/matmul/dot/maxpool/softmax | softmax 只发 `expf(x)`；conv 发 `fadd 0.0, 0.0`；maxpool/reshape 直通；matmul/dot/gemm 发标量乘——**无循环、无 GEP、无 MAC** | +| P9 | 目标 triple 硬编码 | 模块头 | `target triple = "riscv64-unknown-elf"` 对所有用途写死，宿主 `lli`/交叉目标不可配置 | + +#### 修复后期望能力 + +- 任意由 `DSLParser`/`ExtendedDSLParser`/`ONNXParser` 产生的 IR，经 `LLVMCodegen.emit()` 输出后 `llvm-as` 全部通过。 +- 8 个算子的 IR 具有与 standalone 路径一致的循环嵌套结构、GEP 地址计算、浮点 MAC，可被 `lli`/`opt` 进一步消费。 +- 每个 SSA 名字全函数唯一、每个标签唯一定义、每个基本块以终止指令结束。 + +### 1.2 设计目标 + +- **正确可汇编**：`llvm-as file.ll -o /dev/null` 零错误是硬门槛，优先于性能与可读性。 +- **结构可对齐**：循环结构、索引公式、GEP 形式、浮点常量编码与 `onnx_to_llvm_standalone.py`（下称 standalone）保持一致，便于两条路径对比。 +- **嵌套安全**：for/if/while 任意深度嵌套，标签与 SSA 名无冲突。 +- **数值可验证**：小规模张量场景可用 `lli` 执行并断言数值（整数与 0.5 等可精确表示值）。 +- **改动收敛**：只改库路径与其测试；standalone 不动；不引入 ScratchV IR→LLVM 的高级优化（不做 mem2reg、循环展开、向量化）。 + +--- + +## 二、设计规范 + +### 2.1 总体架构 + +`LLVMCodegen` 采用**单趟文本生成**：遍历 `Program.functions`，逐函数遍历 `BasicBlock`，逐指令分派到 `_emit_`。新增/重构三类基础设施： + +1. **命名器**（`SSANamer`）：寄存器与标签分开计数，保证唯一性。 +2. **控制流状态机**：显式跟踪“当前块是否已终止”，所有标签经由 `_start_block()` 打开。 +3. **张量存储约定**：IR `Value` 携带 `shape`，非空即视为指针（`float*` 等）；标量参与张量算子时溢写（spill）到 `alloca`，退化为 1 元素张量。 + +不改变 `emit()`/`save()` 的调用契约：`scratchv/compiler.py:388-390` 的 `LLVMCodegen(program).emit()` 继续工作。 + +### 2.2 类型规则 + +| IR dtype | 标量 LLVM 类型 | 指针 LLVM 类型 | 常量写法 | 允许的指令族 | +|----------|----------------|----------------|----------|--------------| +| `FLOAT32` | `float` | `float*` | `0.0`/`1.0`/`-1.0` 或 64 位十六进制（见下） | `fadd/fsub/fmul/fdiv/fneg/fcmp/call @expf @tanhf` | +| `FLOAT64` | `double` | `double*` | 同上 | 同上（`@exp @tanh`） | +| `INT32` | `i32` | `i32*` | 十进制整数（如 `3`、`-1`） | `add/sub/mul/sdiv/srem/icmp` | +| `INT64` | `i64` | `i64*` | 十进制整数 | 同上 | + +约束规则： + +- **不隐式转换**：二元指令两侧类型必须一致；索引、循环计数、维度一律 `i32`。 +- **混型算术显式转换**：当操作数类型与结果类型不一致时插入转换指令——int→float 用 `sitofp v to `，float→int 用 `fptosi v to `（结果类型为整型时）；常量直接按结果类型格式化，不生成转换。例：`s = add(s, i)`（`s: float`、`i: i32`）⇒ `%if = sitofp i32 %i to float` + `%r = fadd float %s, %if`。该规则覆盖 DSL `for` 循环体把循环变量（i32）与 float 变量混用的常见写法。 +- **张量指针判定**：`Value.shape != ()` 或该值由 `OpCode.ALLOCA` 定义 ⇒ 类型为“元素类型 + `*`”。函数参数、返回值、算子操作数/结果统一遵循此规则。 +- **整数常量禁止浮点算子**：`load_const` 目标为整型时，生成 `%r = add i32 0, `（或直接把常量作为使用点立即数），绝不生成 `fadd i32 ...`。 +- **浮点常量必须精确可表示**：LLVM IR 十进制浮点常量必须能被目标类型精确表示（LLVM 10 校验严格）。统一使用 standalone 的编码算法： + - 先经 `struct.pack("_`，`n` 在**整个函数内单调递增、永不复用**；函数内与函数间均不依赖 LLVM 自动改名。 +- 寄存器计数器与标签计数器**分离**（现有实现共用 `_block_counter`，是 P1/P2 的根因之一）。 +- 合法字符集：`[A-Za-z0-9_.]`，首字符必须为字母/`_`；`sanitize()` 将非法字符替换为 `_`，空串/数字开头加前缀 `v_`。 +- **一次定义**：同一 SSA 名在函数内只允许一次 `= ...` 定义（参数行除外）。 + - `_dest(instr)` 幂等：同一 IR `Value` 重复出现时返回已绑定引用，不再生成新定义； + - 多指令展开的算子（gelu/sigmoid）**每一步都用 fresh 中间名**，最终结果单独用 fresh 名，禁止“复用 dst 名写三行”（P2 根因）。 +- **循环变量**：IR 循环变量 `Value` 不直接作为 SSA 定义，而是在 header 中 `load` 一次得到 `%_ld_`，该 load 及其结果支配 body 与 exit；循环体与循环后的引用统一绑定到它（P3 根因修复）。 +- 标签名由 `SSANamer.fresh_label(hint)` 产生（无 `%` 前缀），进入 `defined_labels` 集合；重复定义时追加后缀，绝不输出两个同名标签。 +- 中间名 hint 约定（增强可读性）： + +| 算子 | 寄存器 hint | 标签 hint | +|------|-------------|-----------| +| gelu | `gelu_t1, gelu_x3, gelu_inner, gelu_tanh, gelu_p1, gelu_hx` | — | +| sigmoid | `sig_neg, sig_exp, sig_den, sig_out` | — | +| dot | `dot_i, dot_acc, dot_prod` | `dot_i` | +| matmul | `mm_i, mm_j, mm_k, mm_acc` | `mm_i/mm_j/mm_k` | +| gemm | `gemm_i, gemm_j, gemm_k, gemm_acc` | `gemm_i/j/k` | +| conv | `conv_oc/oh/ow/ic/kh/kw, conv_acc, conv_mac, conv_skip` | `conv_oc/.../conv_kw` | +| maxpool | `mp_c/oh/ow/kh/kw, mp_max` | `mp_c/.../mp_kw` | +| softmax | `sm_i1/i2/i3, sm_max, sm_sum, sm_e` | `sm_max/sm_sum/sm_div` | + +### 2.4 控制流合法性规则 + +LLVM 基本块要求：**每个块恰好一条终止指令（`ret`/`br`/`br i1`）且在块尾**；标签唯一定义；所有使用被定义支配。对应实现规则： + +1. **块状态机**：字段 `self._terminated: bool`。`_start_block(label)` 的语义是“结束当前块并打开新块”： + - 若 `not self._terminated`，先发射 `br label %