diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ca28c08..b48c673 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -44,11 +44,15 @@ pytest tests/ -v # run all tests ## Adding a New Optimization Pass 1. Create `scratchv/optimizer/my_pass.py`. -2. Implement a class with a `run(program) → int` method (returns number of - transformations applied). -3. Register it in `scratchv/main.py` → `run_optimizer()`. -4. Add test cases (positive: should transform; negative: should not). -5. Run `pytest` to verify. +2. Subclass `OptimizationPass` and define a stable, lowercase, kebab-case + `name` such as `"constant-folding"`. The manager uses this identifier in + ordered reports and failure diagnostics. +3. Implement `optimize(program) → int`. It must return a non-negative count + of transformations made by that invocation, not a lifetime total. +4. Register it in `create_optimization_pass_manager()` in + `scratchv/compiler.py`. +5. Add test cases (positive: should transform; negative: should not). +6. Run `pytest` to verify. ## Documentation diff --git a/benchmarks/run_benchmark.py b/benchmarks/run_benchmark.py index f473214..8e42553 100644 --- a/benchmarks/run_benchmark.py +++ b/benchmarks/run_benchmark.py @@ -76,27 +76,14 @@ def _parse_onnx(path: str) -> Program: def _optimize(program: Program, level: str) -> float: - """Run optimizations. Returns elapsed time in seconds.""" - if level == "none": - return 0.0 + """Run the selected unified pipeline and return its reported elapsed time. - t0 = time.perf_counter() - - from scratchv.optimizer.constant_folding import ConstantFolder - from scratchv.optimizer.dead_code import DeadCodeEliminator - - ConstantFolder(program).run() - DeadCodeEliminator(program).run() + The ``none`` pipeline is empty, so its report is exactly ``0.0`` seconds. + """ + from scratchv.compiler import create_optimization_pass_manager - if level == "all": - from scratchv.optimizer.peephole import IRPeepholeOptimizer - from scratchv.optimizer.muladd_fusion import MulAddFusion - from scratchv.optimizer.licm import LICM - IRPeepholeOptimizer(program).run() - MulAddFusion(program).run() - LICM(program).run() - - return time.perf_counter() - t0 + manager = create_optimization_pass_manager(level) + return manager.run(program).elapsed_seconds def _codegen_riscv(program: Program) -> tuple[str, float]: @@ -162,11 +149,8 @@ def run_benchmark(model_name: str, model_path: str, *, result.ir_inst_count, result.ir_bb_count = _count_ir(program) # 2. Optimize - if optimize_level != "none": - result.optimize_time_s = _optimize(program, optimize_level) - result.ir_opt_inst_count, _ = _count_ir(program) - else: - result.ir_opt_inst_count = result.ir_inst_count + result.optimize_time_s = _optimize(program, optimize_level) + result.ir_opt_inst_count, _ = _count_ir(program) # 3. Codegen if backend == "llvm": diff --git a/benchmarks/test_benchmark.py b/benchmarks/test_benchmark.py index 74e08c2..e591708 100644 --- a/benchmarks/test_benchmark.py +++ b/benchmarks/test_benchmark.py @@ -65,22 +65,18 @@ def test_parse_onnx(model_name: str, benchmark_models: dict[str, str]): @pytest.mark.parametrize("model_name", MODEL_PARAMS, ids=_model_id) def test_optimize(model_name: str, benchmark_models: dict[str, str]): """Parse + optimize, check IR is not empty.""" + from scratchv.compiler import create_optimization_pass_manager from scratchv.frontend.onnx_parser import ONNXParser - from scratchv.optimizer.constant_folding import ConstantFolder - from scratchv.optimizer.dead_code import DeadCodeEliminator - from scratchv.optimizer.peephole import IRPeepholeOptimizer path = benchmark_models[model_name] program = ONNXParser().parse(path) inst_before = sum(1 for f in program.functions for bb in f.blocks for _ in bb.instructions) - ConstantFolder(program).run() - DeadCodeEliminator(program).run() - IRPeepholeOptimizer(program).run() + create_optimization_pass_manager("all").run(program) inst_after = sum(1 for f in program.functions for bb in f.blocks for _ in bb.instructions) - assert inst_after >= 0, f"Optimization failed for {model_name}" + assert 0 < inst_after <= inst_before, f"Optimization failed for {model_name}" print(f"\n {model_name}: {inst_before} → {inst_after} instructions") @@ -92,9 +88,8 @@ def test_optimize(model_name: str, benchmark_models: dict[str, str]): @pytest.mark.parametrize("backend", BACKEND_PARAMS) def test_codegen_riscv(model_name: str, backend: str, benchmark_models: dict[str, str]): """Parse + codegen → RISC-V assembly, check output is non-empty.""" + from scratchv.compiler import create_optimization_pass_manager from scratchv.frontend.onnx_parser import ONNXParser - from scratchv.optimizer.constant_folding import ConstantFolder - from scratchv.optimizer.dead_code import DeadCodeEliminator from scratchv.backend.instruction_select import InstructionSelector from scratchv.backend.register_alloc import RegisterAllocator from scratchv.backend.asm_emit import AsmEmitter @@ -102,8 +97,7 @@ def test_codegen_riscv(model_name: str, backend: str, benchmark_models: dict[str path = benchmark_models[model_name] program = ONNXParser().parse(path) - ConstantFolder(program).run() - DeadCodeEliminator(program).run() + create_optimization_pass_manager("basic").run(program) selector = InstructionSelector(program) machine = selector.run() diff --git a/docs/optimization_guide.md b/docs/optimization_guide.md index 6960142..1d9eff3 100644 --- a/docs/optimization_guide.md +++ b/docs/optimization_guide.md @@ -3,6 +3,28 @@ Six beginner-friendly optimization passes for ScratchV, ordered by difficulty. Each pass can be implemented as a standalone task. +All IR passes define a stable `name` and implement +`OptimizationPass.optimize(self, program: Program) -> int`. The return value is +the non-negative number of transformations made by the current invocation. +Use the canonical factory when selecting an optimization level: + +```python +from scratchv.compiler import create_optimization_pass_manager + +manager = create_optimization_pass_manager("all") +report = manager.run(program) +print(report.total_changes) +``` + +`manager.run()` returns an immutable `OptimizationReport`. Its ordered +`executions` contain each pass name, change count, and elapsed time; +`total_changes` and `elapsed_seconds` summarize the complete pipeline. + +The main CLI accepts `--opt-level none|basic|all`; `--optimize` remains an +alias for the option name and still requires one of those three values. A bare +`--optimize` is invalid. The standalone LLVM tool has a separate numeric +`--opt-level 0|1|2|3` option. + --- ## 1. Constant Folding (⭐) @@ -16,6 +38,24 @@ b = 5 c = add(a, b) → c = 8 (replaced with load_const) ``` +### W3 typed ADD/SUB semantics + +The first ConstantFolder milestone folds only structurally valid scalar +`ADD`/`SUB` candidates whose two operands and destination have the same +`FLOAT32` or `INT32` type: + +- `FLOAT32` inputs and results are rounded through IEEE-754 binary32. NaN, + infinity, and finite operations that overflow to infinity are left intact. +- `INT32` accepts in-range Python integers and integer-valued finite floats. + Results wrap as signed 32-bit two's-complement values. +- Mixed or unsupported types, malformed payloads, non-scalars, missing + destinations, and instructions with unknown attributes or targets are left + unchanged. + +The pass performs one forward scan and reports changes from that invocation. +Existing simple `MUL`/`DIV` folding remains for compatibility; W4 will give +those operations typed semantics and add fixed-point folding. + --- ## 2. Dead Code Elimination (⭐⭐) @@ -66,21 +106,27 @@ Scans assembly for redundant patterns and removes them: **Implementation** in `scratchv/optimizer/peephole.py`: ```python class PeepholeOptimizer: - def run(self, program: Program) -> int: + def optimize(self, program: Program) -> int: + changes = 0 for func in program.functions: for block in func.blocks: - self._optimize_block(block) + changes += self._optimize_block(block) + return changes - def _optimize_block(self, block): + def _optimize_block(self, block) -> int: + changes = 0 i = 0 while i < len(block.instructions): if self._is_addi_zero(block.instructions[i]): block.instructions.pop(i) + changes += 1 continue elif self._is_jump_to_next(block, i): block.instructions.pop(i) + changes += 1 continue i += 1 + return changes ``` --- @@ -109,14 +155,19 @@ for out_y in range(H_out): **Implementation** in `scratchv/optimizer/licm.py`: ```python class LICM: - def run(self, program: Program) -> int: + def optimize(self, program: Program) -> int: + changes = 0 for func in program.functions: - self._find_loops_and_hoist(func) + changes += self._find_loops_and_hoist(func) + return changes - def _find_loops_and_hoist(self, func): + def _find_loops_and_hoist(self, func) -> int: + hoisted_count = 0 # 1. Find FOR/ENDFOR pairs # 2. Identify instructions whose operands don't change in loop # 3. Move them before the FOR instruction + # 4. Increment hoisted_count for each moved instruction + return hoisted_count ``` --- diff --git "a/docs/topics/04-IR\344\274\230\345\214\226\345\231\250\346\241\206\346\236\266.md" "b/docs/topics/04-IR\344\274\230\345\214\226\345\231\250\346\241\206\346\236\266.md" index 3b726b2..a44223c 100644 --- "a/docs/topics/04-IR\344\274\230\345\214\226\345\231\250\346\241\206\346\236\266.md" +++ "b/docs/topics/04-IR\344\274\230\345\214\226\345\231\250\346\241\206\346\236\266.md" @@ -84,11 +84,43 @@ IR 优化器框架包含 5 个 IR 层面的优化 pass,它们共享相同接 所有 pass 实现相同签名: ```python class OptimizationPass: + @property + def name(self) -> str: + """返回稳定的小写连字符名称,如 constant-folding。""" + ... + def optimize(self, program: Program) -> int: - """对 Program 原地优化,返回变更次数""" + """对 Program 原地优化,返回本次调用的非负变更次数。""" ... ``` +`name` 是报告和失败诊断中使用的稳定标识符;PassManager 按注册顺序 +执行,不会根据名称排序或去重。具体 pass 可以用类属性实现: + +```python +class ConstantFolder(OptimizationPass): + name = "constant-folding" +``` + +默认管线由唯一工厂创建,避免主编译器与 benchmark 各自维护 Pass 列表: + +```python +from scratchv.compiler import create_optimization_pass_manager + +manager = create_optimization_pass_manager("basic") +report = manager.run(program) +for execution in report.executions: + print(execution.name, execution.changes, execution.elapsed_seconds) +``` + +主命令行使用 `--opt-level none|basic|all`;旧参数名 `--optimize` 暂作参数名的 +兼容别名,仍须显式提供三个级别之一,裸 `--optimize` 不合法。 + +```bash +scratchv model.onnx --opt-level basic +scratchv model.onnx --optimize basic +``` + ### 死代码消除的回溯标记 ```python diff --git a/docs/topics/archive/optimizer_framework.md b/docs/topics/archive/optimizer_framework.md index b247d57..a2ca63d 100644 --- a/docs/topics/archive/optimizer_framework.md +++ b/docs/topics/archive/optimizer_framework.md @@ -1,5 +1,8 @@ # Topic: IR 优化器框架 (Optimizer Framework) +> **归档说明:** 本文保留为课题历史摘要;当前公共接口以 +> [`docs/topics/04-IR优化器框架.md`](../04-IR优化器框架.md) 和实际代码为准。 + > **源文件**: `scratchv/optimizer/` > **层级**: IR → IR (源和目标都是 IR Program) > **策略**: 所有 pass 原地修改 Program, 返回变更计数 @@ -98,7 +101,34 @@ if changes > 0: print(f"Folded {changes} constants") ``` -所有 pass 遵循相同接口,可在 `PassManager` 中链式调用。 +所有 pass 遵循相同接口。当前实现的默认三级管线通过唯一工厂函数创建: + +```python +from scratchv.compiler import create_optimization_pass_manager + +manager = create_optimization_pass_manager("all") +report = manager.run(program) +print(report.total_changes) +``` + +编译器入口和 benchmark 共用该工厂,当前映射为: + +- `none`:空管线 +- `basic`:`constant-folding` → `dead-code-elim` +- `all`:`constant-folding` → `dead-code-elim` → `ir-peephole` → + `muladd-fusion` → `licm` + +`manager.run()` 返回不可变的 +`OptimizationReport`;其 `executions` 按执行顺序记录 pass 名称、变更数和耗时, +`total_changes` 与 `elapsed_seconds` 提供总计。 + +主 CLI 使用 `--opt-level none|basic|all`。`--optimize` 只是参数名的兼容别名, +仍须显式提供 `none`、`basic` 或 `all`;裸 `--optimize` 不合法。 + +```bash +scratchv model.onnx --opt-level all +scratchv model.onnx --optimize all +``` ## 相关 Topic diff --git a/docs/verification.md b/docs/verification.md index 1e5356f..8188eb3 100644 --- a/docs/verification.md +++ b/docs/verification.md @@ -201,10 +201,10 @@ Add verification to your workflow: ```bash # 1. Compile with ScratchV (RISC-V backend) -scratchv model.onnx -o output.s --optimize +scratchv model.onnx -o output.s --opt-level all # 2. Compile with LLVM backend -scratchv model.onnx --backend llvm -o model.ll --optimize +scratchv model.onnx --backend llvm -o model.ll --opt-level all # 3. Verify against ONNX Runtime scratchv model.onnx --verify @@ -218,6 +218,10 @@ lli model.ll # LLVM JIT execution # (before vs after optimization) ``` +> **Note:** 核心 CLI 中 `--optimize` 只是 `--opt-level` 的参数名兼容别名, +> 仍须显式提供 `none`、`basic` 或 `all`。例如 `--optimize all` 等价于 +> `--opt-level all`;裸 `--optimize` 不合法。新文档优先使用后者。 + --- ## LLVM IR Verification diff --git a/examples/end_to_end_pipeline.py b/examples/end_to_end_pipeline.py index 593b702..77dc559 100644 --- a/examples/end_to_end_pipeline.py +++ b/examples/end_to_end_pipeline.py @@ -107,14 +107,20 @@ def demo_matmul(backend: str): parser = DSLParser() program = parser.parse(dsl_source) - # Optimize - from scratchv.optimizer.constant_folding import ConstantFolder - from scratchv.optimizer.dead_code import DeadCodeEliminator - folder = ConstantFolder(program) - folded = folder.run() - elim = DeadCodeEliminator(program) - eliminated = elim.run() - print(f" Optimizer: {folded} folded, {eliminated} eliminated") + # "basic" is the stable constant-folding -> dead-code-elim pipeline. + # run() optimizes program in place and returns an immutable report. + from scratchv.compiler import create_optimization_pass_manager + report = create_optimization_pass_manager("basic").run(program) + changes_by_name = { + execution.name: execution.changes for execution in report.executions + } + folded_changes = changes_by_name["constant-folding"] + eliminated_changes = changes_by_name["dead-code-elim"] + print( + " Optimizer: " + f"{folded_changes} constant fold(s), " + f"{eliminated_changes} dead-code elimination(s)" + ) from scratchv.backend.llvm_codegen import LLVMCodegen codegen = LLVMCodegen(program) @@ -160,18 +166,16 @@ def demo_optimized_pipeline(): print(codegen1.emit()[:400]) print("...") - # With optimization + # "all" runs the five canonical passes in registration order. + # run() optimizes program2 in place and returns an immutable report. parser2 = DSLParser() program2 = parser2.parse(dsl_source) - from scratchv.optimizer.constant_folding import ConstantFolder - from scratchv.optimizer.dead_code import DeadCodeEliminator - from scratchv.optimizer.peephole import IRPeepholeOptimizer - folder = ConstantFolder(program2) - folder.run() - elim = DeadCodeEliminator(program2) - elim.run() - peep = IRPeepholeOptimizer(program2) - peep.run() + from scratchv.compiler import create_optimization_pass_manager + report = create_optimization_pass_manager("all").run(program2) + print( + f" Optimizer: {report.total_changes} change(s) " + f"across {len(report.executions)} passes" + ) codegen2 = LLVMCodegen(program2) print("After optimization (fold + dce + peephole):") print(codegen2.emit()[:400]) diff --git a/examples/llvm_optimization_pipeline.py b/examples/llvm_optimization_pipeline.py index a8a6b6d..1912252 100644 --- a/examples/llvm_optimization_pipeline.py +++ b/examples/llvm_optimization_pipeline.py @@ -46,13 +46,19 @@ def main(): parser2 = DSLParser() program2 = parser2.parse(dsl_source) - from scratchv.optimizer.constant_folding import ConstantFolder - from scratchv.optimizer.dead_code import DeadCodeEliminator - folder = ConstantFolder(program2) - folded = folder.run() - elim = DeadCodeEliminator(program2) - eliminated = elim.run() - print(f" Folded: {folded}, Eliminated: {eliminated}") + # "basic" is the stable constant-folding -> dead-code-elim pipeline. + # run() optimizes program2 in place and returns an immutable report. + from scratchv.compiler import create_optimization_pass_manager + report = create_optimization_pass_manager("basic").run(program2) + changes_by_name = { + execution.name: execution.changes for execution in report.executions + } + folded_changes = changes_by_name["constant-folding"] + eliminated_changes = changes_by_name["dead-code-elim"] + print( + f" Constant folds: {folded_changes}, " + f"dead-code eliminations: {eliminated_changes}" + ) codegen2 = LLVMCodegen(program2) basic_ir = codegen2.emit() @@ -64,14 +70,9 @@ def main(): parser3 = DSLParser() program3 = parser3.parse(dsl_source) - folder3 = ConstantFolder(program3) - folder3.run() - elim3 = DeadCodeEliminator(program3) - elim3.run() - from scratchv.optimizer.peephole import IRPeepholeOptimizer - peep = IRPeepholeOptimizer(program3) - peeped = peep.run() - print(f" Folded+DCE+Peephole: {peeped} optimizations") + # "all" runs the five canonical passes in registration order and mutates program3. + report = create_optimization_pass_manager("all").run(program3) + print(f" Full pipeline total changes: {report.total_changes}") codegen3 = LLVMCodegen(program3) opt_ir = codegen3.emit() diff --git a/examples/onnx_llvm_verification.py b/examples/onnx_llvm_verification.py index 691ecfe..af57481 100644 --- a/examples/onnx_llvm_verification.py +++ b/examples/onnx_llvm_verification.py @@ -56,13 +56,19 @@ def main(): # Step 2: Optimize print("\n[2/5] Optimizing IR...") - from scratchv.optimizer.constant_folding import ConstantFolder - from scratchv.optimizer.dead_code import DeadCodeEliminator - folder = ConstantFolder(program) - folded = folder.run() - elim = DeadCodeEliminator(program) - eliminated = elim.run() - print(f" Folded: {folded}, Eliminated: {eliminated}") + # "basic" is the stable constant-folding -> dead-code-elim pipeline. + # run() optimizes program in place and returns an immutable report. + from scratchv.compiler import create_optimization_pass_manager + report = create_optimization_pass_manager("basic").run(program) + changes_by_name = { + execution.name: execution.changes for execution in report.executions + } + folded_changes = changes_by_name["constant-folding"] + eliminated_changes = changes_by_name["dead-code-elim"] + print( + f" Constant folds: {folded_changes}, " + f"dead-code eliminations: {eliminated_changes}" + ) # Step 3: Generate LLVM IR print("\n[3/5] Generating LLVM IR...") diff --git a/examples/verify_with_tinyfive.py b/examples/verify_with_tinyfive.py index 02a352e..5d2b91c 100644 --- a/examples/verify_with_tinyfive.py +++ b/examples/verify_with_tinyfive.py @@ -14,8 +14,7 @@ from scratchv.backend.instruction_select import InstructionSelector from scratchv.backend.register_alloc import RegisterAllocator from scratchv.backend.asm_emit import AsmEmitter -from scratchv.optimizer.constant_folding import ConstantFolder -from scratchv.optimizer.dead_code import DeadCodeEliminator +from scratchv.compiler import create_optimization_pass_manager def compile_and_count(path: str, optimize: bool = False) -> tuple[str, int]: @@ -27,11 +26,20 @@ def compile_and_count(path: str, optimize: bool = False) -> tuple[str, int]: program = parser.parse(source) if optimize: - folder = ConstantFolder(program) - folded = folder.run() - elim = DeadCodeEliminator(program) - eliminated = elim.run() - print(f" Optimization: {folded} folded, {eliminated} eliminated") + # "basic" is the stable constant-folding -> dead-code-elim pipeline. + # run() optimizes program in place and returns an immutable report. + report = create_optimization_pass_manager("basic").run(program) + changes_by_name = { + execution.name: execution.changes + for execution in report.executions + } + folded_changes = changes_by_name["constant-folding"] + eliminated_changes = changes_by_name["dead-code-elim"] + print( + " Optimization: " + f"{folded_changes} constant fold(s), " + f"{eliminated_changes} dead-code elimination(s)" + ) selector = InstructionSelector(program) instrs = selector.run() diff --git a/scratchv/compiler.py b/scratchv/compiler.py index a3484d2..960b411 100644 --- a/scratchv/compiler.py +++ b/scratchv/compiler.py @@ -1,6 +1,6 @@ """Compiler driver and pass manager for ScratchV. -Provides a ``PassManager`` that chains compiler passes together and a +Provides a ``PassManager`` that runs IR optimization passes and a ``CompilerDriver`` that orchestrates the full compilation pipeline: parse → optimise → codegen → verify → emit. @@ -21,9 +21,15 @@ import time from dataclasses import dataclass, field -from typing import Any, Optional +from typing import Any -from scratchv.pass_interface import CompilerPass, PassResult +from scratchv.ir.types import Program +from scratchv.pass_interface import ( + OptimizationPass, + OptimizationPassError, + OptimizationReport, + PassExecutionStats, +) # ═══════════════════════════════════════════════════════════════════════════════ @@ -80,81 +86,96 @@ class CompilerConfig: # ═══════════════════════════════════════════════════════════════════════════════ class PassManager: - """Manages a sequence of compiler passes and runs them in order. - - Each pass receives the output of the previous pass as its input. The - first pass receives the initial ``input_data`` provided to ``run()``. + """Register and run IR optimization passes in a deterministic order. Usage:: pm = PassManager() - pm.add(ConstantFolder(program)) - pm.add(DeadCodeEliminator(program)) - result = pm.run(program) + pm.register(ConstantFolder()) + pm.register(DeadCodeEliminator()) + report = pm.run(program) """ def __init__(self, name: str = "pipeline"): + if not isinstance(name, str): + raise TypeError("pipeline name must be a string") + if not name.strip(): + raise ValueError("pipeline name must not be empty") self._name = name - self._passes: list[CompilerPass] = [] + self._passes: list[OptimizationPass] = [] @property def name(self) -> str: return self._name @property - def passes(self) -> list[CompilerPass]: + def passes(self) -> list[OptimizationPass]: return list(self._passes) - def add(self, pass_: CompilerPass) -> "PassManager": - """Add a pass to the end of the pipeline. Returns self for chaining.""" + def register(self, pass_: OptimizationPass) -> "PassManager": + """Append an optimization pass and return ``self`` for chaining.""" + if not isinstance(pass_, OptimizationPass): + raise TypeError("registered pass must implement OptimizationPass") + if not isinstance(pass_.name, str): + raise TypeError("optimization pass name must be a string") + if not pass_.name.strip(): + raise ValueError("optimization pass name must not be empty") self._passes.append(pass_) return self - def run(self, input_data: Any) -> PassResult: - """Run all passes sequentially. + def add(self, pass_: OptimizationPass) -> "PassManager": + """Compatibility alias for :meth:`register`.""" + return self.register(pass_) - Returns the final ``PassResult``. If any pass returns ``None`` - data the pipeline stops early and returns the last result. - """ - data = input_data - total_changes = 0 - messages: list[str] = [] - all_warnings: list[str] = [] - timings: dict[str, float] = {} + def run(self, program: Program) -> OptimizationReport: + """Optimize ``program`` in place and return ordered execution stats.""" + if not isinstance(program, Program): + raise TypeError("PassManager.run() requires a Program") - for p in self._passes: + executions: list[PassExecutionStats] = [] + for index, pass_ in enumerate(self._passes): t0 = time.perf_counter() try: - result = p.run(data) + changes = pass_.optimize(program) + self._validate_change_count(changes) except Exception as exc: - return PassResult( - data=None, - changes=total_changes, - message=f"Pass '{p.name}' failed: {exc}", - warnings=all_warnings, - ) + elapsed = time.perf_counter() - t0 + completed_report = self._make_report(executions) + raise OptimizationPassError( + pass_index=index, + pass_name=pass_.name, + elapsed_seconds=elapsed, + completed_report=completed_report, + cause=exc, + ) from exc elapsed = time.perf_counter() - t0 - timings[p.name] = elapsed - - if result.data is None: - return PassResult( - data=None, - changes=total_changes, - message=f"Pipeline stopped after '{p.name}': {result.message}", - warnings=all_warnings + result.warnings, + executions.append( + PassExecutionStats( + index=index, + name=pass_.name, + changes=changes, + elapsed_seconds=elapsed, ) + ) - data = result.data - total_changes += result.changes - if result.message: - messages.append(f"[{p.name}] {result.message}") - all_warnings.extend(result.warnings) - - return PassResult( - data=data, - changes=total_changes, - message="; ".join(messages) if messages else "pipeline complete", - warnings=all_warnings, + return self._make_report(executions) + + @staticmethod + def _validate_change_count(changes: int) -> None: + if isinstance(changes, bool) or not isinstance(changes, int): + raise TypeError("optimization pass must return an integer change count") + if changes < 0: + raise ValueError("optimization pass change count must be non-negative") + + def _make_report(self, executions: list[PassExecutionStats]) -> OptimizationReport: + records = tuple(executions) + return OptimizationReport( + pipeline_name=self._name, + executions=records, + total_changes=sum(item.changes for item in records), + elapsed_seconds=sum( + (item.elapsed_seconds for item in records), 0.0 + ), ) def report(self) -> str: @@ -165,6 +186,35 @@ def report(self) -> str: return "\n".join(lines) +def create_optimization_pass_manager(level: str) -> PassManager: + """Build the canonical ``none``/``basic``/``all`` IR pipeline.""" + if not isinstance(level, str): + raise TypeError("optimization level must be a string") + if level not in {"none", "basic", "all"}: + raise ValueError("optimization level must be one of: none, basic, all") + + manager = PassManager("optimizer") + if level == "none": + return manager + + from scratchv.optimizer.constant_folding import ConstantFolder + from scratchv.optimizer.dead_code import DeadCodeEliminator + + manager.register(ConstantFolder()) + manager.register(DeadCodeEliminator()) + + if level == "all": + from scratchv.optimizer.licm import LICM + from scratchv.optimizer.muladd_fusion import MulAddFusion + from scratchv.optimizer.peephole import IRPeepholeOptimizer + + manager.register(IRPeepholeOptimizer()) + manager.register(MulAddFusion()) + manager.register(LICM()) + + return manager + + # ═══════════════════════════════════════════════════════════════════════════════ # CompileResult # ═══════════════════════════════════════════════════════════════════════════════ @@ -233,7 +283,6 @@ def compile(self, input_path: str, output_path: str | None = None, Returns: A ``CompileResult`` with output text and statistics. """ - errors: list[str] = [] warnings: list[str] = [] # Resolve output path @@ -259,9 +308,24 @@ def compile(self, input_path: str, output_path: str | None = None, # --- 3. Optimize --- opt_message = "" - if self.config.optimize_level != "none": - opt_result = self._run_optimizations(program) - opt_message = opt_result.message + try: + optimization_report = self._run_optimizations(program) + except (OptimizationPassError, TypeError, ValueError) as exc: + if isinstance(exc, OptimizationPassError): + completed_report = exc.completed_report + else: + completed_report = OptimizationReport("optimizer", (), 0, 0.0) + optimization_stats = self._optimization_stats(completed_report) + return CompileResult( + success=False, + errors=[f"Optimization error: {exc}"], + stats={ + "optimization": optimization_stats, + "opt_message": opt_message, + "cycle_report": "", + }, + ) + opt_message = self._optimization_message(optimization_report) ir_dump_after = "" if self.config.dump_ir: @@ -316,7 +380,11 @@ def compile(self, input_path: str, output_path: str | None = None, output_text=asm_text, output_path=output_path, ir_dump=ir_dump, - stats={"opt_message": opt_message, "cycle_report": cycle_report}, + stats={ + "optimization": self._optimization_stats(optimization_report), + "opt_message": opt_message, + "cycle_report": cycle_report, + }, warnings=warnings, ) @@ -361,25 +429,35 @@ def _verify_ir(self, program, warnings: list[str]) -> None: # ── Internal: optimizations ───────────────────────────────────────────── - def _run_optimizations(self, program) -> PassResult: + def _run_optimizations(self, program: Program) -> OptimizationReport: """Run all configured optimization passes.""" - from scratchv.optimizer.constant_folding import ConstantFolder - from scratchv.optimizer.dead_code import DeadCodeEliminator - - pm = PassManager("optimizer") - pm.add(_PassAdapter("constant-folding", ConstantFolder(program))) - pm.add(_PassAdapter("dead-code-elim", DeadCodeEliminator(program))) - - if self.config.optimize_level == "all": - from scratchv.optimizer.peephole import IRPeepholeOptimizer - from scratchv.optimizer.muladd_fusion import MulAddFusion - from scratchv.optimizer.licm import LICM - - pm.add(_PassAdapter("ir-peephole", IRPeepholeOptimizer(program))) - pm.add(_PassAdapter("muladd-fusion", MulAddFusion(program))) - pm.add(_PassAdapter("licm", LICM(program))) - - return pm.run(program) + manager = create_optimization_pass_manager(self.config.optimize_level) + return manager.run(program) + + def _optimization_stats(self, report: OptimizationReport) -> dict[str, Any]: + """Convert an immutable report to the CompileResult stats schema.""" + return { + "level": self.config.optimize_level, + "total_changes": report.total_changes, + "elapsed_seconds": report.elapsed_seconds, + "passes": [ + { + "index": item.index, + "name": item.name, + "changes": item.changes, + "elapsed_seconds": item.elapsed_seconds, + } + for item in report.executions + ], + } + + @staticmethod + def _optimization_message(report: OptimizationReport) -> str: + """Build the legacy human-readable optimization summary.""" + return "; ".join( + f"[{item.name}] {item.changes} change(s)" + for item in report.executions + ) # ── Internal: code generation ─────────────────────────────────────────── @@ -480,31 +558,3 @@ def _run_asm_passes(self, asm_text: str, warnings: list[str]) -> str: warnings.append(f"Instruction count: {total}") return asm_text - - -# ═══════════════════════════════════════════════════════════════════════════════ -# _PassAdapter — wraps legacy passes that don't implement CompilerPass -# ═══════════════════════════════════════════════════════════════════════════════ - -class _PassAdapter(CompilerPass): - """Adapter that wraps a legacy pass object into the ``CompilerPass`` API. - - Legacy passes are expected to have a ``run()`` method that returns an - integer (number of changes) and mutate the data in place. - """ - - def __init__(self, name: str, legacy_pass: Any): - self._name = name - self._legacy = legacy_pass - - @property - def name(self) -> str: - return self._name - - def run(self, input_data: Any) -> PassResult: - changes = self._legacy.run() - return PassResult( - data=input_data, - changes=changes, - message=f"{changes} change(s)", - ) diff --git a/scratchv/main.py b/scratchv/main.py index 25eaf33..e958f7c 100644 --- a/scratchv/main.py +++ b/scratchv/main.py @@ -38,7 +38,8 @@ def build_arg_parser() -> argparse.ArgumentParser: # ── Optimizations ─────────────────────────────────────────────────── parser.add_argument( - "--optimize", choices=["none", "basic", "all"], + "--opt-level", "--optimize", dest="optimize_level", + choices=["none", "basic", "all"], default="none", help="Optimization level (none, basic, all)", ) @@ -136,7 +137,7 @@ def args_to_config(args: argparse.Namespace) -> CompilerConfig: """Translate parsed CLI arguments to a CompilerConfig.""" return CompilerConfig( backend=args.backend, - optimize_level=args.optimize, + optimize_level=args.optimize_level, reg_alloc=args.reg_alloc, dump_ir=args.dump_ir, verify=args.verify, diff --git a/scratchv/optimizer/constant_folding.py b/scratchv/optimizer/constant_folding.py index ce3e493..1a1e20c 100644 --- a/scratchv/optimizer/constant_folding.py +++ b/scratchv/optimizer/constant_folding.py @@ -6,43 +6,53 @@ from __future__ import annotations +import math +import struct + from scratchv.ir.types import ( - OpCode, Instruction, BasicBlock, Function, Program, + BasicBlock, + DataType, + Function, + Instruction, + OpCode, + Program, ) +from scratchv.pass_interface import OptimizationPass -class ConstantFolder: +class ConstantFolder(OptimizationPass): """Fold constant expressions in an IR Program.""" - def __init__(self, program: Program): - self.program = program - self._stats = {"folded": 0} + name = "constant-folding" - def run(self) -> int: - """Run constant folding on all functions. Returns number of folds.""" - for func in self.program.functions: - self._fold_function(func) - return self._stats["folded"] + def optimize(self, program: Program) -> int: + """Fold constants in one forward scan and return this run's count.""" + return sum(self._fold_function(func) for func in program.functions) - def _fold_function(self, func: Function) -> None: - for block in func.blocks: - self._fold_block(block) + def _fold_function(self, func: Function) -> int: + return sum(self._fold_block(block) for block in func.blocks) - def _fold_block(self, block: BasicBlock) -> None: + def _fold_block(self, block: BasicBlock) -> int: + changes = 0 new_instrs: list[Instruction] = [] for instr in block.instructions: folded = self._try_fold(instr) if folded is not None: new_instrs.append(folded) + changes += 1 else: new_instrs.append(instr) block.instructions = new_instrs + return changes def _try_fold(self, instr: Instruction) -> Instruction | None: """Try to fold an instruction. Returns a replacement or None.""" if instr.opcode not in ( - OpCode.ADD, OpCode.SUB, - OpCode.MUL, OpCode.DIV): + OpCode.ADD, + OpCode.SUB, + OpCode.MUL, + OpCode.DIV, + ): return None if len(instr.operands) != 2: return None @@ -53,29 +63,109 @@ def _try_fold(self, instr: Instruction) -> Instruction | None: if lhs.const_value is None or rhs.const_value is None: return None - a, b = float(lhs.const_value), float(rhs.const_value) - result = self._compute(instr.opcode, a, b) + dest = instr.dest + if instr.opcode in (OpCode.ADD, OpCode.SUB): + if dest is None or instr.target is not None or instr.attrs: + return None + if not ( + lhs.dtype is rhs.dtype + and rhs.dtype is dest.dtype + and dest.dtype in (DataType.FLOAT32, DataType.INT32) + ): + return None + if lhs.shape != () or rhs.shape != () or dest.shape != (): + return None + + if instr.opcode in (OpCode.ADD, OpCode.SUB): + result = self._compute( + instr.opcode, + lhs.const_value, + rhs.const_value, + lhs.dtype, + ) + else: + try: + legacy_lhs = float(lhs.const_value) + legacy_rhs = float(rhs.const_value) + except (TypeError, ValueError, OverflowError): + return None + result = self._compute( + instr.opcode, + legacy_lhs, + legacy_rhs, + lhs.dtype, + ) if result is None: return None - self._stats["folded"] += 1 - dest = instr.dest - if dest is not None: - dest.is_constant = True - dest.const_value = result - - return Instruction( + replacement = Instruction( opcode=OpCode.LOAD_CONST, dest=dest, attrs={"value": result}, ) + if dest is not None: + dest.is_constant = True + dest.const_value = result + return replacement + + @staticmethod + def _quantize_float32(value: float | int) -> float | None: + """Return a finite IEEE-754 binary32 value as a Python float.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + try: + result = struct.unpack(" int | None: + """Normalize a valid INT32 scalar payload to a Python int.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + if isinstance(value, float) and ( + not math.isfinite(value) or not value.is_integer() + ): + return None + normalized = int(value) + if not -(2**31) <= normalized <= 2**31 - 1: + return None + return normalized + + @staticmethod + def _wrap_int32(value: int) -> int: + """Interpret the low 32 bits of ``value`` as signed two's complement.""" + return (value + 2**31) % 2**32 - 2**31 @staticmethod - def _compute(opcode: OpCode, a: float, b: float) -> float | None: - mapping = { - OpCode.ADD: a + b, - OpCode.SUB: a - b, - OpCode.MUL: a * b, - OpCode.DIV: a / b if b != 0 else None, - } - return mapping.get(opcode) + def _compute( + opcode: OpCode, + a: float | int, + b: float | int, + dtype: DataType, + ) -> float | int | None: + """Compute a typed scalar result, or return None when unsafe.""" + if dtype is DataType.FLOAT32 and opcode in (OpCode.ADD, OpCode.SUB): + lhs = ConstantFolder._quantize_float32(a) + rhs = ConstantFolder._quantize_float32(b) + if lhs is None or rhs is None: + return None + raw_result = lhs + rhs if opcode is OpCode.ADD else lhs - rhs + return ConstantFolder._quantize_float32(raw_result) + + if dtype is DataType.INT32 and opcode in (OpCode.ADD, OpCode.SUB): + lhs = ConstantFolder._normalize_int32(a) + rhs = ConstantFolder._normalize_int32(b) + if lhs is None or rhs is None: + return None + raw_result = lhs + rhs if opcode is OpCode.ADD else lhs - rhs + return ConstantFolder._wrap_int32(raw_result) + + # W3 preserves the reference implementation's simple MUL/DIV behavior; + # W4 will bring these operations under the typed numeric contract. + if opcode is OpCode.MUL: + return a * b + if opcode is OpCode.DIV: + return a / b if b != 0 else None + return None diff --git a/scratchv/optimizer/dead_code.py b/scratchv/optimizer/dead_code.py index 05c77ed..f2e5485 100644 --- a/scratchv/optimizer/dead_code.py +++ b/scratchv/optimizer/dead_code.py @@ -9,29 +9,26 @@ from scratchv.ir.types import ( OpCode, Instruction, BasicBlock, Function, Program, ) +from scratchv.pass_interface import OptimizationPass -class DeadCodeEliminator: +class DeadCodeEliminator(OptimizationPass): """Remove unused instructions from an IR Program.""" - def __init__(self, program: Program): - self.program = program - self._stats = {"eliminated": 0} + name = "dead-code-elim" - def run(self) -> int: + def optimize(self, program: Program) -> int: """Run dead code elimination. Returns number of eliminated instructions. """ - for func in self.program.functions: - self._eliminate_function(func) - return self._stats["eliminated"] + return sum(self._eliminate_function(func) for func in program.functions) - def _eliminate_function(self, func: Function) -> None: - for block in func.blocks: - self._eliminate_block(block) + def _eliminate_function(self, func: Function) -> int: + return sum(self._eliminate_block(block) for block in func.blocks) - def _eliminate_block(self, block: BasicBlock) -> None: + def _eliminate_block(self, block: BasicBlock) -> int: + changes = 0 # Collect all used value names used: set[str | None] = set() # Return values and branch targets are always live @@ -52,9 +49,10 @@ def _eliminate_block(self, block: BasicBlock) -> None: elif instr.dest is None or instr.dest.name in used: new_instrs.append(instr) else: - self._stats["eliminated"] += 1 + changes += 1 block.instructions = new_instrs + return changes @staticmethod def _is_side_effect(instr: Instruction) -> bool: diff --git a/scratchv/optimizer/licm.py b/scratchv/optimizer/licm.py index e332b22..4d9b254 100644 --- a/scratchv/optimizer/licm.py +++ b/scratchv/optimizer/licm.py @@ -15,30 +15,27 @@ from scratchv.ir.types import ( OpCode, Instruction, BasicBlock, Function, Program, ) +from scratchv.pass_interface import OptimizationPass -class LICM: +class LICM(OptimizationPass): """Hoist loop-invariant code out of loops.""" - def __init__(self, program: Program): - self.program = program - self._stats = {"hoisted": 0} + name = "licm" - def run(self) -> int: + def optimize(self, program: Program) -> int: """Run LICM on all functions. Returns number of hoisted instructions. """ - for func in self.program.functions: - self._process_function(func) - return self._stats["hoisted"] + return sum(self._process_function(func) for func in program.functions) - def _process_function(self, func: Function) -> None: - for block in func.blocks: - self._process_block(block) + def _process_function(self, func: Function) -> int: + return sum(self._process_block(block) for block in func.blocks) - def _process_block(self, block: BasicBlock) -> None: + def _process_block(self, block: BasicBlock) -> int: """Find FOR/ENDFOR pairs in a block and hoist invariants.""" + changes = 0 instrs = block.instructions i = 0 while i < len(instrs): @@ -84,10 +81,12 @@ def _process_block(self, block: BasicBlock) -> None: instrs.pop(idx) instrs.insert(loop_start, instr) loop_end += 1 # adjust for shift - self._stats["hoisted"] += 1 + changes += 1 i = loop_end + 1 + return changes + def _find_matching_endfor(self, instrs: list[Instruction], start: int): """Find matching ENDFOR for a FOR at given index.""" depth = 0 diff --git a/scratchv/optimizer/muladd_fusion.py b/scratchv/optimizer/muladd_fusion.py index bbcb6b8..ac79cfa 100644 --- a/scratchv/optimizer/muladd_fusion.py +++ b/scratchv/optimizer/muladd_fusion.py @@ -14,23 +14,24 @@ from __future__ import annotations from scratchv.ir.types import OpCode, Instruction, BasicBlock, Program +from scratchv.pass_interface import OptimizationPass -class MulAddFusion: +class MulAddFusion(OptimizationPass): """Combine consecutive mul+add instruction pairs.""" - def __init__(self, program: Program): - self.program = program - self._stats = {"fused": 0} + name = "muladd-fusion" - def run(self) -> int: + def optimize(self, program: Program) -> int: """Run mul-add fusion. Returns number of fusions performed.""" - for func in self.program.functions: - for block in func.blocks: - self._fuse_block(block) - return self._stats["fused"] + return sum( + self._fuse_block(block) + for func in program.functions + for block in func.blocks + ) - def _fuse_block(self, block: BasicBlock) -> None: + def _fuse_block(self, block: BasicBlock) -> int: + changes = 0 instrs = block.instructions i = 0 while i < len(instrs) - 1: @@ -54,10 +55,12 @@ def _fuse_block(self, block: BasicBlock) -> None: instrs[i] = fused # Remove the original add instrs.pop(i + 1) - self._stats["fused"] += 1 + changes += 1 i += 1 + return changes + def _matches_pattern(self, mul: Instruction, add: Instruction) -> bool: """Check if mul + add form a fusible pattern. diff --git a/scratchv/optimizer/peephole.py b/scratchv/optimizer/peephole.py index b58520e..8ec7a7f 100644 --- a/scratchv/optimizer/peephole.py +++ b/scratchv/optimizer/peephole.py @@ -10,26 +10,27 @@ from __future__ import annotations from scratchv.ir.types import OpCode, Instruction, BasicBlock, Program +from scratchv.pass_interface import OptimizationPass -class IRPeepholeOptimizer: +class IRPeepholeOptimizer(OptimizationPass): """Eliminate redundant instruction patterns in IR.""" - def __init__(self, program: Program): - self.program = program - self._stats = {"eliminated": 0} + name = "ir-peephole" - def run(self) -> int: + def optimize(self, program: Program) -> int: """Run peephole optimization. Returns number of eliminated instructions. """ - for func in self.program.functions: - for block in func.blocks: - self._optimize_block(block) - return self._stats["eliminated"] - - def _optimize_block(self, block: BasicBlock) -> None: + return sum( + self._optimize_block(block) + for func in program.functions + for block in func.blocks + ) + + def _optimize_block(self, block: BasicBlock) -> int: + changes = 0 instrs = block.instructions i = 0 while i < len(instrs): @@ -38,7 +39,7 @@ def _optimize_block(self, block: BasicBlock) -> None: # Pattern 1: addi rd, rs, 0 → delete (no-op) if self._is_addi_zero(instr): instrs.pop(i) - self._stats["eliminated"] += 1 + changes += 1 continue # Pattern 2: mul rd, rs, 1 → mv rd, rs (use add rd, rs, x0) @@ -50,7 +51,7 @@ def _optimize_block(self, block: BasicBlock) -> None: dest=dest, operands=[src, self._make_zero_operand(instr)], ) - self._stats["eliminated"] += 1 # counts as optimization + changes += 1 i += 1 continue @@ -62,18 +63,20 @@ def _optimize_block(self, block: BasicBlock) -> None: dest=dest, attrs={"value": 0}, ) - self._stats["eliminated"] += 1 + changes += 1 i += 1 continue # Pattern 4: j L followed immediately by L: if self._is_jump_to_next(instrs, i): instrs.pop(i) - self._stats["eliminated"] += 1 + changes += 1 continue i += 1 + return changes + def _is_addi_zero(self, instr: Instruction) -> bool: """Check for: add rd, rs, 0 (or add rd, rs, const where const=0).""" if instr.opcode != OpCode.ADD: diff --git a/scratchv/pass_interface.py b/scratchv/pass_interface.py index 12598ad..a94af94 100644 --- a/scratchv/pass_interface.py +++ b/scratchv/pass_interface.py @@ -1,7 +1,8 @@ -"""Unified compiler pass interface for ScratchV. +"""Unified compiler pass interfaces for ScratchV. -All optimisation, analysis, and code-generation passes implement this -protocol so they can be composed by ``PassManager`` and ``CompilerDriver``. +IR optimization passes use the strongly typed ``OptimizationPass`` contract. +The older generic ``CompilerPass`` and ``PassResult`` types remain available +for compatibility with non-optimization callers. Usage:: @@ -26,6 +27,111 @@ def run(self, input_data): from dataclasses import dataclass, field from typing import Any +from scratchv.ir.types import Program + + +class OptimizationPass(ABC): + """Abstract interface implemented by every IR optimization pass.""" + + @property + @abstractmethod + def name(self) -> str: + """Return a stable, non-empty name for reports and diagnostics.""" + ... + + @abstractmethod + def optimize(self, program: Program) -> int: + """Optimize ``program`` in place and return this run's change count.""" + ... + + +@dataclass(frozen=True) +class PassExecutionStats: + """Statistics for one execution in an optimization pipeline.""" + + index: int + name: str + changes: int + elapsed_seconds: float + + def __post_init__(self) -> None: + if isinstance(self.index, bool) or not isinstance(self.index, int): + raise TypeError("pass execution index must be an integer") + if self.index < 0: + raise ValueError("pass execution index must be non-negative") + if not isinstance(self.name, str): + raise TypeError("pass execution name must be a string") + if not self.name.strip(): + raise ValueError("pass execution name must not be empty") + if isinstance(self.changes, bool) or not isinstance(self.changes, int): + raise TypeError("pass changes must be an integer") + if self.changes < 0: + raise ValueError("pass changes must be non-negative") + if isinstance(self.elapsed_seconds, bool) or not isinstance( + self.elapsed_seconds, (int, float) + ): + raise TypeError("pass elapsed time must be numeric") + if self.elapsed_seconds < 0: + raise ValueError("pass elapsed time must be non-negative") + + +@dataclass(frozen=True) +class OptimizationReport: + """Ordered statistics produced by one optimization pipeline run.""" + + pipeline_name: str + executions: tuple[PassExecutionStats, ...] + total_changes: int + elapsed_seconds: float + + def __post_init__(self) -> None: + if not isinstance(self.pipeline_name, str): + raise TypeError("pipeline name must be a string") + if not self.pipeline_name.strip(): + raise ValueError("pipeline name must not be empty") + if not isinstance(self.executions, tuple): + raise TypeError("pipeline executions must be a tuple") + if [item.index for item in self.executions] != list( + range(len(self.executions)) + ): + raise ValueError("pass execution indexes must be consecutive") + if isinstance(self.total_changes, bool) or not isinstance( + self.total_changes, int + ): + raise TypeError("total changes must be an integer") + if self.total_changes < 0: + raise ValueError("total changes must be non-negative") + if self.total_changes != sum(item.changes for item in self.executions): + raise ValueError("total changes must equal the execution sum") + if isinstance(self.elapsed_seconds, bool) or not isinstance( + self.elapsed_seconds, (int, float) + ): + raise TypeError("pipeline elapsed time must be numeric") + if self.elapsed_seconds < 0: + raise ValueError("pipeline elapsed time must be non-negative") + + +class OptimizationPassError(RuntimeError): + """A pass failed or returned a value that violates its contract.""" + + def __init__( + self, + pass_index: int, + pass_name: str, + elapsed_seconds: float, + completed_report: OptimizationReport, + cause: Exception, + ) -> None: + self.pass_index = pass_index + self.pass_name = pass_name + self.elapsed_seconds = elapsed_seconds + self.completed_report = completed_report + self.cause = cause + + super().__init__( + f"optimization pass #{pass_index} '{pass_name}' failed: {cause}" + ) + # ═══════════════════════════════════════════════════════════════════════════════ # PassResult — uniform return value for all passes @@ -64,8 +170,8 @@ class CompilerPass(ABC): Subclasses must implement ``name`` and ``run``. ``stats`` is optional and defaults to an empty dict. - The ``run`` method accepts and returns arbitrary data — passes are - composed by ``PassManager`` which chains their inputs and outputs. + The legacy ``run`` method accepts and returns arbitrary data. IR + optimization pipelines instead use the typed ``OptimizationPass``. """ @property diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 64441ff..3bd2c4c 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -1,12 +1,468 @@ """Tests for optimizer passes.""" +from __future__ import annotations + +import math + +import pytest + from scratchv.ir.builder import IRBuilder +from scratchv.ir.types import ( + BasicBlock, + DataType, + Function, + Instruction, + OpCode, + Program, + Value, +) from scratchv.optimizer.constant_folding import ConstantFolder from scratchv.optimizer.dead_code import DeadCodeEliminator -from scratchv.ir.types import OpCode + + +def _typed_binary_program( + opcode: OpCode, + lhs_value: float | int, + rhs_value: float | int, + dtype: DataType, + *, + dest_dtype: DataType | None = None, + shape: tuple[int, ...] = (), + candidate_attrs: dict[str, object] | None = None, + target: str | None = None, +) -> tuple[Program, BasicBlock, Value, Instruction]: + """Build minimal, explicitly typed scalar IR for folding tests.""" + lhs = Value( + "lhs", + dtype=dtype, + is_constant=True, + const_value=lhs_value, + shape=shape, + ) + rhs = Value( + "rhs", + dtype=dtype, + is_constant=True, + const_value=rhs_value, + shape=shape, + ) + dest = Value("result", dtype=dest_dtype or dtype, shape=shape) + candidate = Instruction( + opcode, + dest=dest, + operands=[lhs, rhs], + attrs=candidate_attrs or {}, + target=target, + ) + user = Instruction(OpCode.RETURN, operands=[dest]) + block = BasicBlock("entry") + block.instructions = [ + Instruction(OpCode.LOAD_CONST, dest=lhs, attrs={"value": lhs_value}), + Instruction(OpCode.LOAD_CONST, dest=rhs, attrs={"value": rhs_value}), + candidate, + user, + ] + program = Program() + program.add_function(Function("test", blocks=[block])) + return program, block, dest, user class TestConstantFolder: + def test_compute_rejects_unsupported_typed_operation(self): + result = ConstantFolder._compute( + OpCode.ADD, + 1.0, + 2.0, + DataType.FLOAT64, + ) + + assert result is None + + def test_compute_rejects_unsupported_opcode(self): + result = ConstantFolder._compute( + OpCode.NEG, + 1.0, + 2.0, + DataType.FLOAT32, + ) + + assert result is None + + def test_fold_float32_add_uses_binary32_rounding(self): + program, block, _, _ = _typed_binary_program( + OpCode.ADD, + 16_777_216.0, + 1.0, + DataType.FLOAT32, + ) + + count = ConstantFolder().optimize(program) + + result = block.instructions[2].attrs["value"] + assert count == 1 + assert result == 16_777_216.0 + assert type(result) is float + + def test_fold_float32_preserves_signed_zero(self): + program, block, _, _ = _typed_binary_program( + OpCode.SUB, + -0.0, + 0.0, + DataType.FLOAT32, + ) + + count = ConstantFolder().optimize(program) + + result = block.instructions[2].attrs["value"] + assert count == 1 + assert result == 0.0 + assert math.copysign(1.0, result) == -1.0 + + @pytest.mark.parametrize( + ("opcode", "lhs", "rhs", "expected"), + [ + (OpCode.ADD, 3.0, 4.0, 7.0), + (OpCode.SUB, 3.5, 5.0, -1.5), + ], + ) + def test_fold_float32_add_and_sub(self, opcode, lhs, rhs, expected): + program, block, _, _ = _typed_binary_program( + opcode, + lhs, + rhs, + DataType.FLOAT32, + ) + + count = ConstantFolder().optimize(program) + + result = block.instructions[2].attrs["value"] + assert count == 1 + assert result == expected + assert type(result) is float + + def test_fold_int32_add_wraps_on_overflow(self): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 2_147_483_647, + 1, + DataType.INT32, + ) + + count = ConstantFolder().optimize(program) + + result = block.instructions[2].attrs["value"] + assert count == 1 + assert result == -2_147_483_648 + assert type(result) is int + assert dest.dtype is DataType.INT32 + + @pytest.mark.parametrize( + ("opcode", "lhs", "rhs", "expected"), + [ + (OpCode.ADD, 20, 22, 42), + (OpCode.SUB, 3, 5, -2), + (OpCode.SUB, -2_147_483_648, 1, 2_147_483_647), + (OpCode.ADD, 2.0, 3.0, 5), + ], + ) + def test_fold_int32_add_and_sub(self, opcode, lhs, rhs, expected): + program, block, dest, _ = _typed_binary_program( + opcode, + lhs, + rhs, + DataType.INT32, + ) + + count = ConstantFolder().optimize(program) + + result = block.instructions[2].attrs["value"] + assert count == 1 + assert result == expected + assert type(result) is int + assert dest.const_value == expected + + @pytest.mark.parametrize( + "invalid_payload", + [True, 1.5, 2_147_483_648, -2_147_483_649, math.inf], + ) + def test_skip_invalid_int32_payload(self, invalid_payload): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + invalid_payload, + 1, + DataType.INT32, + ) + candidate = block.instructions[2] + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + @pytest.mark.parametrize("invalid_payload", [math.nan, math.inf, -math.inf]) + def test_skip_non_finite_float32_payload(self, invalid_payload): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + invalid_payload, + 1.0, + DataType.FLOAT32, + ) + candidate = block.instructions[2] + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + def test_skip_float32_result_overflow(self): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 3.4028234663852886e38, + 3.4028234663852886e38, + DataType.FLOAT32, + ) + candidate = block.instructions[2] + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + def test_skip_add_when_destination_dtype_differs(self): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 20, + 22, + DataType.INT32, + dest_dtype=DataType.FLOAT32, + ) + candidate = block.instructions[2] + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + assert dest.const_value is None + + def test_skip_add_when_operand_dtypes_differ(self): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 20, + 22.0, + DataType.INT32, + ) + candidate = block.instructions[2] + candidate.operands[1].dtype = DataType.FLOAT32 + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + @pytest.mark.parametrize("dtype", [DataType.FLOAT64, DataType.INT64]) + def test_skip_unsupported_dtype(self, dtype): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 1, + 2, + dtype, + ) + candidate = block.instructions[2] + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + def test_skip_non_scalar_add(self): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 1.0, + 2.0, + DataType.FLOAT32, + shape=(2,), + ) + candidate = block.instructions[2] + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + def test_skip_add_with_unknown_instruction_semantics(self): + for extra_fields in ( + {"candidate_attrs": {"saturating": True}}, + {"target": "other"}, + ): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 1.0, + 2.0, + DataType.FLOAT32, + **extra_fields, + ) + candidate = block.instructions[2] + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + def test_skip_float32_add_with_bool_payload(self): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + True, + 1.0, + DataType.FLOAT32, + ) + candidate = block.instructions[2] + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + def test_skip_float32_add_with_non_numeric_payload(self): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 1.0, + 2.0, + DataType.FLOAT32, + ) + candidate = block.instructions[2] + candidate.operands[0].const_value = "not-a-number" # type: ignore[assignment] + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + @pytest.mark.parametrize( + "malformation", + ["zero-operands", "one-operand", "three-operands", "no-dest"], + ) + def test_skip_structurally_invalid_add(self, malformation): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 1.0, + 2.0, + DataType.FLOAT32, + ) + candidate = block.instructions[2] + if malformation == "zero-operands": + candidate.operands = [] + elif malformation == "one-operand": + candidate.operands = candidate.operands[:1] + elif malformation == "three-operands": + candidate.operands.append(candidate.operands[0]) + else: + candidate.dest = None + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + @pytest.mark.parametrize("invalid_constant", ["not-constant", "no-value"]) + def test_skip_unavailable_constant_payload(self, invalid_constant): + program, block, dest, _ = _typed_binary_program( + OpCode.ADD, + 1.0, + 2.0, + DataType.FLOAT32, + ) + candidate = block.instructions[2] + lhs = candidate.operands[0] + if invalid_constant == "not-constant": + lhs.is_constant = False + else: + lhs.const_value = None + + count = ConstantFolder().optimize(program) + + assert count == 0 + assert block.instructions[2] is candidate + assert not dest.is_constant + + def test_fold_preserves_destination_and_uses(self): + program, block, dest, user = _typed_binary_program( + OpCode.SUB, + 9.0, + 2.0, + DataType.FLOAT32, + ) + original_shape = dest.shape + + count = ConstantFolder().optimize(program) + + replacement = block.instructions[2] + assert count == 1 + assert replacement.opcode is OpCode.LOAD_CONST + assert replacement.dest is dest + assert replacement.operands == [] + assert replacement.attrs == {"value": 7.0} + assert dest.dtype is DataType.FLOAT32 + assert dest.shape == original_shape + assert dest.is_constant + assert dest.const_value == replacement.attrs["value"] + assert user.operands[0] is dest + + def test_repeated_run_reports_only_current_changes(self): + program, _, _, _ = _typed_binary_program( + OpCode.ADD, + 1.0, + 2.0, + DataType.FLOAT32, + ) + folder = ConstantFolder() + + first_count = folder.optimize(program) + second_count = folder.optimize(program) + + assert first_count == 1 + assert second_count == 0 + + def test_fold_counts_candidates_across_functions_and_blocks(self): + program = Program() + blocks = [] + for index, opcode in enumerate((OpCode.ADD, OpCode.SUB, OpCode.ADD)): + candidate_program, block, _, _ = _typed_binary_program( + opcode, + float(index + 3), + 1.0, + DataType.FLOAT32, + ) + candidate_function = candidate_program.functions[0] + if index == 0: + candidate_function.name = "first" + program.add_function(candidate_function) + elif index == 1: + block.name = "next" + program.functions[0].blocks.extend(candidate_function.blocks) + else: + candidate_function.name = "second" + program.add_function(candidate_function) + blocks.append(block) + + count = ConstantFolder().optimize(program) + + assert count == 3 + assert all( + block.instructions[2].opcode is OpCode.LOAD_CONST + for block in blocks + ) + def test_fold_add_constants(self): builder = IRBuilder() builder.new_function("test") @@ -17,8 +473,8 @@ def test_fold_add_constants(self): r = builder.add(c1, c2) builder.ret(r) - folder = ConstantFolder(builder.program) - count = folder.run() + folder = ConstantFolder() + count = folder.optimize(builder.program) assert count == 1 block = builder.program.functions[0].blocks[0] @@ -36,8 +492,8 @@ def test_fold_mul_constants(self): r = builder.mul(c1, c2) builder.ret(r) - folder = ConstantFolder(builder.program) - count = folder.run() + folder = ConstantFolder() + count = folder.optimize(builder.program) assert count == 1 # The mul (index 2) was replaced by load_const 10.0 block = builder.program.functions[0].blocks[0] @@ -53,8 +509,8 @@ def test_no_fold_with_variable(self): r = builder.add(a, c) builder.ret(r) - folder = ConstantFolder(builder.program) - count = folder.run() + folder = ConstantFolder() + count = folder.optimize(builder.program) assert count == 0 # cannot fold because 'a' is not constant @@ -70,8 +526,8 @@ def test_eliminate_unused(self): d = builder.load_const(3.0) builder.ret(d) - elim = DeadCodeEliminator(builder.program) - count = elim.run() + elim = DeadCodeEliminator() + count = elim.optimize(builder.program) assert count == 1 # the add should be eliminated block = builder.program.functions[0].blocks[0] @@ -88,6 +544,6 @@ def test_keep_used_value(self): c = builder.add(a, b) # used by ret builder.ret(c) - elim = DeadCodeEliminator(builder.program) - count = elim.run() + elim = DeadCodeEliminator() + count = elim.optimize(builder.program) assert count == 0 # nothing eliminated diff --git a/tests/test_optimizer_advanced.py b/tests/test_optimizer_advanced.py index 1d4fa5a..44c2bc9 100644 --- a/tests/test_optimizer_advanced.py +++ b/tests/test_optimizer_advanced.py @@ -19,8 +19,8 @@ def test_eliminate_addi_zero(self): builder._emit(OpCode.ADD, c, [a, zero]) builder.ret(c) - opt = IRPeepholeOptimizer(builder.program) - opt.run() + opt = IRPeepholeOptimizer() + opt.optimize(builder.program) # The addi 0 should have been removed block = builder.program.functions[0].blocks[0] assert all(i.opcode != OpCode.ADD for i in block.instructions) @@ -35,8 +35,8 @@ def test_eliminate_mul_one(self): builder._emit(OpCode.MUL, c, [a, one]) builder.ret(c) - opt = IRPeepholeOptimizer(builder.program) - count = opt.run() + opt = IRPeepholeOptimizer() + count = opt.optimize(builder.program) assert count >= 1 # MUL with 1 should be replaced block = builder.program.functions[0].blocks[0] @@ -53,8 +53,8 @@ def test_eliminate_mul_zero(self): builder._emit(OpCode.MUL, c, [a, zero]) builder.ret(c) - opt = IRPeepholeOptimizer(builder.program) - opt.run() + opt = IRPeepholeOptimizer() + opt.optimize(builder.program) block = builder.program.functions[0].blocks[0] mul_instrs = [i for i in block.instructions if i.opcode == OpCode.MUL] assert len(mul_instrs) == 0 @@ -77,8 +77,8 @@ def test_fuse_mul_add(self): result = builder.add(tmp, acc) builder.ret(result) - opt = MulAddFusion(builder.program) - count = opt.run() + opt = MulAddFusion() + count = opt.optimize(builder.program) assert count == 1 block = builder.program.functions[0].blocks[0] @@ -97,8 +97,8 @@ def test_no_fuse_without_mul(self): c = builder.add(a, b) builder.ret(c) - opt = MulAddFusion(builder.program) - count = opt.run() + opt = MulAddFusion() + count = opt.optimize(builder.program) assert count == 0 @@ -119,8 +119,8 @@ def test_hoist_invariant(self): builder.endfor() builder.ret() - opt = LICM(builder.program) - count = opt.run() + opt = LICM() + count = opt.optimize(builder.program) assert count == 1 # mul should be hoisted block = builder.program.functions[0].blocks[0] @@ -142,8 +142,8 @@ def test_no_hoist_variant(self): builder.endfor() builder.ret() - opt = LICM(builder.program) - opt.run() + opt = LICM() + opt.optimize(builder.program) # The load_const IS invariant and gets hoisted. # But the add depending on iv stays in the loop. # Verify the add remains after the FOR. diff --git a/tests/test_pass_manager.py b/tests/test_pass_manager.py new file mode 100644 index 0000000..bc52a23 --- /dev/null +++ b/tests/test_pass_manager.py @@ -0,0 +1,381 @@ +"""Tests for the unified IR optimization pass framework.""" + +from __future__ import annotations + +from dataclasses import FrozenInstanceError + +import pytest + +import scratchv.compiler as compiler_module +from scratchv.compiler import ( + CompilerConfig, + CompilerDriver, + PassManager, + create_optimization_pass_manager, +) +from scratchv.ir.builder import IRBuilder +from scratchv.ir.types import OpCode, Program +from scratchv.main import args_to_config, build_arg_parser +from scratchv.optimizer import ( + ConstantFolder, + DeadCodeEliminator, + IRPeepholeOptimizer, + LICM, + MulAddFusion, +) +from scratchv.pass_interface import ( + OptimizationPass, + OptimizationPassError, + OptimizationReport, + PassExecutionStats, +) + + +class _RecordingPass(OptimizationPass): + """Small public-contract fake used by PassManager tests.""" + + def __init__( + self, + name: str, + changes: int = 0, + calls: list[str] | None = None, + ): + self._name = name + self._changes = changes + self._calls = calls + + @property + def name(self) -> str: + return self._name + + def optimize(self, program: Program) -> int: + if self._calls is not None: + self._calls.append(self.name) + return self._changes + + +class TestOptimizationPass: + def test_base_class_is_abstract(self): + with pytest.raises(TypeError): + OptimizationPass() + + @pytest.mark.parametrize("name", ["", " ", None, 7]) + def test_register_rejects_invalid_name(self, name): + manager = PassManager() + + with pytest.raises((TypeError, ValueError)): + manager.register(_RecordingPass(name)) + + @pytest.mark.parametrize( + "pass_type", + [ + ConstantFolder, + DeadCodeEliminator, + IRPeepholeOptimizer, + MulAddFusion, + LICM, + ], + ) + def test_production_passes_share_the_interface_and_local_counts( + self, pass_type + ): + pass_ = pass_type() + program = Program() + + assert isinstance(pass_, OptimizationPass) + assert pass_.name + assert pass_.optimize(program) == 0 + assert pass_.optimize(program) == 0 + + +class TestPassManager: + def test_runs_in_registration_order_and_reports_each_execution(self): + calls: list[str] = [] + program = Program() + manager = PassManager("ordered") + manager.register(_RecordingPass("first", 2, calls)) + manager.add(_RecordingPass("second", 3, calls)) + + report = manager.run(program) + + assert calls == ["first", "second"] + assert report.pipeline_name == "ordered" + assert report.total_changes == 5 + assert report.elapsed_seconds >= 0 + assert [execution.index for execution in report.executions] == [0, 1] + assert [execution.name for execution in report.executions] == [ + "first", + "second", + ] + assert [execution.changes for execution in report.executions] == [2, 3] + assert all( + execution.elapsed_seconds >= 0 for execution in report.executions + ) + + def test_allows_duplicate_pass_names(self): + manager = PassManager() + manager.register(_RecordingPass("repeat", 1)) + manager.register(_RecordingPass("repeat", 2)) + + report = manager.run(Program()) + + assert [(item.index, item.name) for item in report.executions] == [ + (0, "repeat"), + (1, "repeat"), + ] + assert report.total_changes == 3 + + def test_empty_pipeline_returns_an_empty_report(self): + report = PassManager("empty").run(Program()) + + assert report == OptimizationReport("empty", (), 0, 0.0) + + def test_reports_are_immutable(self): + execution = PassExecutionStats(0, "pass", 0, 0.0) + + with pytest.raises(FrozenInstanceError): + execution.changes = 1 + + @pytest.mark.parametrize("result", [-1, True, 1.5, "1", None]) + def test_rejects_invalid_change_counts(self, result): + manager = PassManager().register(_RecordingPass("invalid", result)) + + with pytest.raises(OptimizationPassError) as caught: + manager.run(Program()) + + assert caught.value.pass_index == 0 + assert caught.value.pass_name == "invalid" + assert caught.value.completed_report.executions == () + assert caught.value.__cause__ is caught.value.cause + + def test_wraps_failure_and_stops_before_later_passes(self): + calls: list[str] = [] + + class _FailingPass(_RecordingPass): + def optimize(self, program: Program) -> int: + calls.append(self.name) + raise RuntimeError("boom") + + manager = PassManager("failure") + manager.register(_RecordingPass("done", 4, calls)) + manager.register(_FailingPass("broken")) + manager.register(_RecordingPass("never", 1, calls)) + + with pytest.raises(OptimizationPassError) as caught: + manager.run(Program()) + + error = caught.value + assert calls == ["done", "broken"] + assert error.pass_index == 1 + assert error.pass_name == "broken" + assert error.elapsed_seconds >= 0 + assert error.completed_report.total_changes == 4 + assert [item.name for item in error.completed_report.executions] == ["done"] + assert isinstance(error.cause, RuntimeError) + assert error.__cause__ is error.cause + + def test_rejects_non_program_before_running_passes(self): + calls: list[str] = [] + manager = PassManager().register(_RecordingPass("pass", calls=calls)) + + with pytest.raises(TypeError): + manager.run(object()) + + assert calls == [] + + def test_does_not_catch_keyboard_interrupt(self): + class _InterruptingPass(_RecordingPass): + def optimize(self, program: Program) -> int: + raise KeyboardInterrupt + + manager = PassManager().register(_InterruptingPass("interrupt")) + + with pytest.raises(KeyboardInterrupt): + manager.run(Program()) + + +class TestOptimizationLevels: + @pytest.mark.parametrize( + ("level", "names"), + [ + ("none", []), + ("basic", ["constant-folding", "dead-code-elim"]), + ( + "all", + [ + "constant-folding", + "dead-code-elim", + "ir-peephole", + "muladd-fusion", + "licm", + ], + ), + ], + ) + def test_factory_builds_the_documented_pipeline(self, level, names): + manager = create_optimization_pass_manager(level) + + assert [pass_.name for pass_ in manager.passes] == names + + @pytest.mark.parametrize("level", ["", "Basic", "fast", None, 1]) + def test_factory_rejects_unknown_level(self, level): + with pytest.raises((TypeError, ValueError)): + create_optimization_pass_manager(level) + + def test_none_pipeline_is_noop_with_zero_elapsed_time(self): + from benchmarks.run_benchmark import _optimize + + builder = IRBuilder() + builder.new_function("test") + block = builder.new_block("entry") + result = builder.add( + builder.make_value(name="lhs"), + builder.make_value(name="rhs"), + ) + builder.ret(result) + instructions_before = tuple(block.instructions) + + report = create_optimization_pass_manager("none").run(builder.program) + + assert report == OptimizationReport("optimizer", (), 0, 0.0) + assert type(report.elapsed_seconds) is float + assert tuple(block.instructions) == instructions_before + assert tuple(map(id, block.instructions)) == tuple( + map(id, instructions_before) + ) + + benchmark_elapsed = _optimize(builder.program, "none") + assert benchmark_elapsed == 0.0 + assert type(benchmark_elapsed) is float + assert tuple(block.instructions) == instructions_before + + def test_basic_pipeline_reports_typed_constant_folding(self): + builder = IRBuilder() + builder.new_function("test") + block = builder.new_block("entry") + lhs = builder.load_const(16_777_216.0) + rhs = builder.load_const(1.0) + result = builder.add(lhs, rhs) + builder.ret(result) + + report = create_optimization_pass_manager("basic").run( + builder.program + ) + + assert report.executions[0].name == "constant-folding" + assert report.executions[0].changes == 1 + folded = next( + instruction + for instruction in block.instructions + if instruction.dest is result + ) + assert folded.opcode is OpCode.LOAD_CONST + assert folded.attrs["value"] == 16_777_216.0 + + +class TestOptimizationCli: + @pytest.mark.parametrize("flag", ["--opt-level", "--optimize"]) + @pytest.mark.parametrize("level", ["none", "basic", "all"]) + def test_accepts_canonical_and_compatibility_flags(self, flag, level): + args = build_arg_parser().parse_args(["model.onnx", flag, level]) + + assert args.optimize_level == level + assert args_to_config(args).optimize_level == level + + def test_defaults_to_none(self): + args = build_arg_parser().parse_args(["model.onnx"]) + + assert args.optimize_level == "none" + + @pytest.mark.parametrize("flag", ["--opt-level", "--optimize"]) + def test_rejects_unknown_level(self, flag): + with pytest.raises(SystemExit): + build_arg_parser().parse_args(["model.onnx", flag, "fast"]) + + +class _ProgramDriver(CompilerDriver): + """Compiler driver test seam that avoids parser and backend dependencies.""" + + def __init__(self, config: CompilerConfig): + super().__init__(config) + self.codegen_called = False + + def _parse(self, input_path: str, dsl_source: str | None = None) -> Program: + return Program() + + def _generate_code(self, program: Program) -> str: + self.codegen_called = True + return "generated" + + +class TestCompilerOptimizationIntegration: + def test_exposes_structured_optimization_stats(self, tmp_path): + output_path = tmp_path / "output.s" + driver = _ProgramDriver(CompilerConfig(optimize_level="basic")) + + result = driver.compile("input.dsl", str(output_path)) + + assert result.success + assert result.stats["optimization"]["level"] == "basic" + assert result.stats["optimization"]["total_changes"] == 0 + assert [ + item["name"] for item in result.stats["optimization"]["passes"] + ] == ["constant-folding", "dead-code-elim"] + assert result.stats["opt_message"] + + def test_invalid_driver_level_stops_before_codegen_and_output(self, tmp_path): + output_path = tmp_path / "must-not-exist.s" + driver = _ProgramDriver(CompilerConfig(optimize_level="fast")) + + result = driver.compile("input.dsl", str(output_path)) + + assert not result.success + assert "optimization level" in result.errors[0] + assert not driver.codegen_called + assert not output_path.exists() + + @pytest.mark.parametrize("level", ["none", "basic", "all"]) + def test_driver_and_benchmark_use_the_same_factory( + self, monkeypatch, level + ): + from benchmarks.run_benchmark import _optimize + + requested_levels: list[str] = [] + real_factory = compiler_module.create_optimization_pass_manager + + def recording_factory(requested_level: str) -> PassManager: + requested_levels.append(requested_level) + return real_factory(requested_level) + + monkeypatch.setattr( + compiler_module, + "create_optimization_pass_manager", + recording_factory, + ) + driver = _ProgramDriver(CompilerConfig(optimize_level=level)) + + driver._run_optimizations(Program()) + _optimize(Program(), level) + + assert requested_levels == [level, level] + + def test_pass_failure_stops_codegen_and_output(self, monkeypatch, tmp_path): + class _FailingPass(_RecordingPass): + def optimize(self, program: Program) -> int: + raise RuntimeError("optimizer failed") + + manager = PassManager("optimizer").register(_FailingPass("broken")) + monkeypatch.setattr( + compiler_module, + "create_optimization_pass_manager", + lambda level: manager, + ) + output_path = tmp_path / "must-not-exist.s" + driver = _ProgramDriver(CompilerConfig(optimize_level="all")) + + result = driver.compile("input.dsl", str(output_path)) + + assert not result.success + assert "broken" in result.errors[0] + assert not driver.codegen_called + assert not output_path.exists()