feat(optimizer): add loop unrolling pass (topic 10)

Implement LoopUnroller that expands FOR/ENDFOR loops with constant trip
counts by replicating the loop body. Two modes:

- Full unroll: trip_count <= unroll_factor removes FOR/ENDFOR and inlines
  the body with induction variable replaced by compile-time constants.
- Partial unroll: trip_count > unroll_factor replicates the body
  unroll_factor times per iteration with offset ADDs, plus a fully-unrolled
  remainder for leftover iterations.

Changes:
- Add scratchv/optimizer/loop_unroll.py (LoopUnroller class)
- Export LoopUnroller in scratchv/optimizer/__init__.py
- Register loop-unroll pass in compiler.py after LICM
- Add tests/test_loop_unroll.py with 13 test cases covering:
  full unroll, partial unroll, nested loops, edge cases

Author: watney1024

scratchv/compiler.py

@@ -374,10 +374,12 @@ def _run_optimizations(self, program) -> PassResult:
             from scratchv.optimizer.peephole import IRPeepholeOptimizer
             from scratchv.optimizer.muladd_fusion import MulAddFusion
             from scratchv.optimizer.licm import LICM
+            from scratchv.optimizer.loop_unroll import LoopUnroller
 
             pm.add(_PassAdapter("ir-peephole", IRPeepholeOptimizer(program)))
             pm.add(_PassAdapter("muladd-fusion", MulAddFusion(program)))
             pm.add(_PassAdapter("licm", LICM(program)))
+            pm.add(_PassAdapter("loop-unroll", LoopUnroller(program)))
 
         return pm.run(program)
 

scratchv/optimizer/__init__.py

@@ -3,11 +3,13 @@
 from .peephole import IRPeepholeOptimizer
 from .muladd_fusion import MulAddFusion
 from .licm import LICM
+from .loop_unroll import LoopUnroller
 
 __all__ = [
     "ConstantFolder",
     "DeadCodeEliminator",
     "IRPeepholeOptimizer",
     "MulAddFusion",
     "LICM",
+    "LoopUnroller",
 ]

scratchv/optimizer/loop_unroll.py

@@ -0,0 +1,317 @@
+"""Loop Unrolling optimization pass.
+
+Expands FOR/ENDFOR loops with constant trip counts by replicating the
+loop body.  Two modes:
+
+  * **Full unroll** – trip_count <= unroll_factor: the FOR/ENDFOR wrapper
+    is removed entirely and the body is inlined ``trip_count`` times with
+    the induction variable replaced by compile-time constants.
+
+  * **Partial unroll** – trip_count > unroll_factor: the loop is kept but
+    each iteration now executes ``unroll_factor`` copies of the original
+    body.  A remainder loop handles the leftover iterations.
+
+Only loops whose ``start``, ``end`` and ``step`` attributes are all
+compile-time integers are eligible.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from scratchv.ir.types import (
+    DataType,
+    OpCode,
+    Value,
+    Instruction,
+    BasicBlock,
+    Program,
+)
+
+
+class LoopUnroller:
+    """Unroll FOR/ENDFOR loops with constant trip counts."""
+
+    def __init__(self, program: Program, unroll_factor: int = 4):
+        self.program = program
+        self.unroll_factor = unroll_factor
+        self._changes = 0
+        self._name_counter = 0
+
+    def run(self) -> int:
+        """Run loop unrolling on all functions.
+
+        Returns the number of loops unrolled.
+        """
+        for func in self.program.functions:
+            for block in func.blocks:
+                self._process_block(block)
+        return self._changes
+
+    # ── helpers ────────────────────────────────────────────────────────
+
+    def _fresh_name(self, prefix: str = "lu") -> str:
+        self._name_counter += 1
+        return f"{prefix}_{self._name_counter}"
+
+    def _fresh_value(self, orig: Value) -> Value:
+        """Create a fresh Value with the same type/shape but a new name."""
+        return Value(
+            name=self._fresh_name(orig.name),
+            dtype=orig.dtype,
+            is_constant=orig.is_constant,
+            const_value=orig.const_value,
+            shape=orig.shape,
+        )
+
+    # ── core logic ─────────────────────────────────────────────────────
+
+    def _process_block(self, block: BasicBlock) -> None:
+        """Scan *block* for FOR/ENDFOR pairs and unroll them."""
+        instrs = block.instructions
+        i = 0
+        while i < len(instrs):
+            if instrs[i].opcode != OpCode.FOR:
+                i += 1
+                continue
+
+            for_idx = i
+            endfor_idx = self._find_matching_endfor(instrs, for_idx)
+            if endfor_idx is None:
+                i += 1
+                continue
+
+            for_instr = instrs[for_idx]
+            start = for_instr.attrs.get("start")
+            end = for_instr.attrs.get("end")
+            step = for_instr.attrs.get("step", 1)
+
+            if not self._is_const_int(start, end, step):
+                i += 1
+                continue
+
+            trip_count = self._compute_trip_count(start, end, step)
+            if trip_count <= 0:
+                i += 1
+                continue
+
+            body = instrs[for_idx + 1 : endfor_idx]
+            if not body:
+                # empty loop – just remove
+                del instrs[for_idx : endfor_idx + 1]
+                self._changes += 1
+                i = for_idx
+                continue
+
+            iv = for_instr.dest
+            factor = self.unroll_factor
+
+            if trip_count <= factor:
+                self._full_unroll(instrs, for_idx, endfor_idx, body, iv,
+                                  start, step, trip_count)
+                self._changes += 1
+                i = 0  # restart: unrolled body may contain new FORs
+            else:
+                self._partial_unroll(instrs, for_idx, endfor_idx,
+                                     body, iv, start, end, step,
+                                     trip_count, factor)
+                self._changes += 1
+                # Don't restart: the loop itself is still eligible but
+                # was already partially unrolled.  Advance past it.
+                i = for_idx + 1
+
+    # ── full unroll ────────────────────────────────────────────────────
+
+    def _full_unroll(
+        self,
+        instrs: list[Instruction],
+        for_idx: int,
+        endfor_idx: int,
+        body: list[Instruction],
+        iv: Optional[Value],
+        start: int,
+        step: int,
+        trip_count: int,
+    ) -> None:
+        """Remove FOR/ENDFOR and inline body *trip_count* times."""
+        replacements: list[Instruction] = []
+        for k in range(trip_count):
+            iv_const = start + k * step
+            for orig_instr in body:
+                cloned = self._clone_instr(orig_instr)
+                if iv is not None:
+                    self._subst_iv(cloned, iv, iv_const)
+                replacements.append(cloned)
+
+        instrs[for_idx : endfor_idx + 1] = replacements
+
+    # ── partial unroll ─────────────────────────────────────────────────
+
+    def _partial_unroll(
+        self,
+        instrs: list[Instruction],
+        for_idx: int,
+        endfor_idx: int,
+        body: list[Instruction],
+        iv: Optional[Value],
+        start: int,
+        end: int,
+        step: int,
+        trip_count: int,
+        factor: int,
+    ) -> int:
+        """Partial unroll: replicate body *factor* times per iteration.
+
+        Returns the number of instructions inserted (for caller to update
+        its index).
+        """
+        full_iters = trip_count // factor
+        remainder = trip_count % factor
+
+        # Build the unrolled body (factor copies)
+        unrolled_body: list[Instruction] = []
+        for k in range(factor):
+            for orig_instr in body:
+                cloned = self._clone_instr(orig_instr)
+                # Replace iv with (iv + k*step) – but only for k>0
+                # k==0 keeps the original iv reference
+                if iv is not None and k > 0:
+                    self._subst_iv_offset(cloned, iv, k * step, unrolled_body)
+                unrolled_body.append(cloned)
+
+        # Adjust FOR bounds: new_end = start + full_iters * factor * step
+        new_end = start + full_iters * factor * step
+        instrs[for_idx].attrs["end"] = new_end
+
+        # Replace old body with unrolled body
+        instrs[for_idx + 1 : endfor_idx] = unrolled_body
+        new_endfor_idx = for_idx + 1 + len(unrolled_body)
+        # Ensure ENDFOR is still there
+        if new_endfor_idx < len(instrs) and instrs[new_endfor_idx].opcode == OpCode.ENDFOR:
+            pass  # already there
+        else:
+            instrs.insert(new_endfor_idx, Instruction(opcode=OpCode.ENDFOR))
+
+        insert_point = new_endfor_idx + 1
+
+        # Remainder: fully unroll the leftover iterations
+        if remainder > 0:
+            remainder_instrs: list[Instruction] = []
+            base_iv_val = new_end
+            for k in range(remainder):
+                iv_const = base_iv_val + k * step
+                for orig_instr in body:
+                    cloned = self._clone_instr(orig_instr)
+                    if iv is not None:
+                        self._subst_iv(cloned, iv, iv_const)
+                    remainder_instrs.append(cloned)
+            instrs[insert_point:insert_point] = remainder_instrs
+            insert_point += len(remainder_instrs)
+
+        return insert_point - for_idx
+
+    # ── instruction cloning ────────────────────────────────────────────
+
+    def _clone_instr(self, instr: Instruction) -> Instruction:
+        """Deep-clone an instruction, giving dest a fresh name."""
+        new_dest = self._fresh_value(instr.dest) if instr.dest else None
+        # Operands keep their names (we fixup iv separately)
+        new_operands = list(instr.operands)
+        return Instruction(
+            opcode=instr.opcode,
+            dest=new_dest,
+            operands=new_operands,
+            attrs=dict(instr.attrs),
+            target=instr.target,
+        )
+
+    # ── induction variable substitution ────────────────────────────────
+
+    def _subst_iv(self, instr: Instruction, iv: Value, const_val: int) -> None:
+        """Replace all references to *iv* in *instr* with a constant."""
+        if instr.dest and instr.dest.name == iv.name:
+            instr.dest = Value(
+                name=instr.dest.name,
+                dtype=instr.dest.dtype,
+                is_constant=True,
+                const_value=const_val,
+                shape=instr.dest.shape,
+            )
+        new_ops: list[Value] = []
+        for op in instr.operands:
+            if op.name == iv.name:
+                new_ops.append(Value(
+                    name=op.name,
+                    dtype=op.dtype,
+                    is_constant=True,
+                    const_value=const_val,
+                    shape=op.shape,
+                ))
+            else:
+                new_ops.append(op)
+        instr.operands = new_ops
+
+    def _subst_iv_offset(
+        self,
+        instr: Instruction,
+        iv: Value,
+        offset: int,
+        prefix_instrs: list[Instruction],
+    ) -> None:
+        """Replace references to *iv* with ``(iv + offset)``.
+
+        For partial unrolling, we cannot replace iv with a constant because
+        the loop variable is still live.  Instead we insert a fresh
+        ``ADD iv, offset`` before this instruction and rewrite the operand
+        reference to the new value.
+        """
+        iv_offset_val = Value(
+            name=self._fresh_name(f"{iv.name}_off"),
+            dtype=iv.dtype,
+        )
+
+        # Emit: iv_offset_val = ADD iv, <offset_const>
+        offset_const = Value(
+            name=self._fresh_name("off"),
+            dtype=DataType.INT32,
+            is_constant=True,
+            const_value=offset,
+        )
+        prefix_instrs.append(Instruction(
+            opcode=OpCode.ADD,
+            dest=iv_offset_val,
+            operands=[iv, offset_const],
+        ))
+
+        # Rewrite operands
+        new_ops: list[Value] = []
+        for op in instr.operands:
+            if op.name == iv.name:
+                new_ops.append(iv_offset_val)
+            else:
+                new_ops.append(op)
+        instr.operands = new_ops
+
+    # ── utilities ──────────────────────────────────────────────────────
+
+    @staticmethod
+    def _find_matching_endfor(instrs: list[Instruction], start: int) -> Optional[int]:
+        depth = 0
+        for i in range(start, len(instrs)):
+            if instrs[i].opcode == OpCode.FOR:
+                depth += 1
+            elif instrs[i].opcode == OpCode.ENDFOR:
+                depth -= 1
+                if depth == 0:
+                    return i
+        return None
+
+    @staticmethod
+    def _is_const_int(*vals) -> bool:
+        return all(isinstance(v, int) for v in vals)
+
+    @staticmethod
+    def _compute_trip_count(start: int, end: int, step: int) -> int:
+        if step <= 0 or end <= start:
+            return 0
+        return (end - start + step - 1) // step  # ceiling division