sonolus.py 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

sonolus/backend/optimize/passes.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections import deque
+
+from sonolus.backend.optimize.flow import BasicBlock
+
+
+class CompilerPass(ABC):
+    def requires(self) -> set[CompilerPass]:
+        return set()
+
+    def preserves(self) -> set[CompilerPass] | None:
+        return None
+
+    def destroys(self) -> set[CompilerPass] | None:
+        return None
+
+    def applies(self) -> set[CompilerPass]:
+        return {self}
+
+    def exists_after(self, passes: set[CompilerPass]) -> set[CompilerPass]:
+        preserved = self.preserves()
+        destroyed = self.destroys()
+        if destroyed is None and preserved is None:
+            return self.applies()
+        if preserved is not None:
+            passes = {p for p in passes if p in preserved}
+        if destroyed is not None:
+            passes = {p for p in passes if p not in destroyed}
+        return passes | self.applies()
+
+    @abstractmethod
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        pass
+
+
+def run_passes(entry: BasicBlock, passes: list[CompilerPass]) -> BasicBlock:
+    active_passes = set()
+    queue = deque(passes)
+    while queue:
+        if len(queue) > 99:
+            raise RuntimeError("Likely unsatisfiable pass requirements")
+        current_pass = queue.popleft()
+        missing_requirements = current_pass.requires() - active_passes
+        if missing_requirements:
+            queue.appendleft(current_pass)
+            queue.extendleft(missing_requirements)
+            continue
+        entry = current_pass.run(entry)
+        active_passes = current_pass.exists_after(active_passes)
+    return entry

sonolus/backend/optimize/simplify.py

@@ -0,0 +1,191 @@
+from sonolus.backend.ir import IRConst, IRGet, IRPureInstr, IRSet
+from sonolus.backend.ops import Op
+from sonolus.backend.optimize.flow import BasicBlock, traverse_cfg_preorder
+from sonolus.backend.optimize.passes import CompilerPass
+
+
+class CoalesceFlow(CompilerPass):
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        queue = [entry]
+        processed = set()
+        while queue:
+            block = queue.pop()
+            if block in processed:
+                continue
+            processed.add(block)
+            for edge in block.outgoing:
+                while True:
+                    dst = edge.dst
+                    if dst.phis or dst.statements or len(dst.outgoing) != 1 or dst is block:
+                        break
+                    next_dst = next(iter(dst.outgoing)).dst
+                    if next_dst.phis:
+                        break
+                    dst.incoming.remove(edge)
+                    if not dst.incoming:
+                        for dst_edge in dst.outgoing:
+                            dst_edge.dst.incoming.remove(dst_edge)
+                        processed.add(dst)
+                    edge.dst = next_dst
+                    next_dst.incoming.add(edge)
+                    if dst is edge.dst:
+                        break
+            default_edge = next((edge for edge in block.outgoing if edge.cond is None), None)
+            if default_edge is not None:
+                for edge in [*block.outgoing]:
+                    if edge is default_edge:
+                        continue
+                    if edge.dst is default_edge.dst:
+                        block.outgoing.remove(edge)
+                        edge.dst.incoming.remove(edge)
+            if len(block.outgoing) != 1:
+                queue.extend(edge.dst for edge in block.outgoing)
+                continue
+            next_block = next(iter(block.outgoing)).dst
+            if len(next_block.incoming) != 1:
+                queue.append(next_block)
+                if not block.statements and not block.phis and not next_block.phis:
+                    for edge in block.incoming:
+                        edge.dst = next_block
+                        next_block.incoming.add(edge)
+                    for edge in block.outgoing:  # There should be exactly one
+                        next_block.incoming.remove(edge)
+                    if block is entry:
+                        entry = next_block
+                continue
+            for p, args in next_block.phis.items():
+                if block not in args:
+                    continue
+                block.statements.append(IRSet(p, IRGet(args[block])))
+            block.statements.extend(next_block.statements)
+            block.test = next_block.test
+            block.outgoing = next_block.outgoing
+            for edge in block.outgoing:
+                edge.src = block
+                dst = edge.dst
+                for args in dst.phis.values():
+                    if next_block in args:
+                        args[block] = args.pop(next_block)
+            processed.add(next_block)
+            queue.extend(edge.dst for edge in block.outgoing)
+            processed.remove(block)
+            queue.append(block)
+        return entry
+
+
+class RewriteToSwitch(CompilerPass):
+    """Rewrite if-else chains to switch statements.
+
+    Note that this needs inlining (and dead code elimination) to be run first to really do anything useful.
+    """
+
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        self.ifs_to_switch(entry)
+        self.combine_blocks(entry)
+        self.remove_unreachable(entry)
+        return entry
+
+    def ifs_to_switch(self, entry: BasicBlock):
+        for block in traverse_cfg_preorder(entry):
+            if len(block.outgoing) != 2 or {edge.cond for edge in block.outgoing} != {None, 0}:
+                continue
+            test = block.test
+            if not isinstance(test, IRPureInstr) or test.op != Op.Equal:
+                continue
+            assert len(test.args) == 2
+            if isinstance(test.args[0], IRConst):
+                const, other = test.args
+            elif isinstance(test.args[1], IRConst):
+                other, const = test.args
+            else:
+                continue
+            block.test = other
+            for edge in block.outgoing:
+                if edge.cond is None:
+                    edge.cond = const.value
+                else:
+                    edge.cond = None
+
+    def combine_blocks(self, entry: BasicBlock):
+        queue = [entry]
+        processed = set()
+        while queue:
+            block = queue.pop()
+            if block in processed:
+                continue
+            processed.add(block)
+            queue.extend(edge.dst for edge in block.outgoing)
+
+            default_edge = next((edge for edge in block.outgoing if edge.cond is None), None)
+            if default_edge is None:
+                continue
+
+            next_block = default_edge.dst
+            if (
+                len(next_block.incoming) > 1
+                or next_block.statements
+                or next_block.phis
+                or block.test != next_block.test
+                or block is next_block
+                or next_block is entry
+            ):
+                continue
+
+            outgoing_by_cond = {edge.cond: edge for edge in block.outgoing}
+            assert len(outgoing_by_cond) == len(block.outgoing)
+            outgoing_by_cond.pop(None)
+            for edge in next_block.outgoing:
+                if edge.cond in outgoing_by_cond:
+                    # This edge is unreachable since an equivalent edge would have been taken
+                    edge.dst.incoming.remove(edge)
+                    continue
+                outgoing_by_cond[edge.cond] = edge
+                edge.src = block
+                for args in edge.dst.phis.values():
+                    if next_block in args:
+                        args[block] = args.pop(next_block)
+            block.outgoing = set(outgoing_by_cond.values())
+            processed.add(next_block)
+            queue.append(block)
+            processed.remove(block)
+
+    def remove_unreachable(self, entry: BasicBlock):
+        reachable = {*traverse_cfg_preorder(entry)}
+        for block in traverse_cfg_preorder(entry):
+            block.incoming = {edge for edge in block.incoming if edge.src in reachable}
+            block.outgoing = {edge for edge in block.outgoing if edge.dst in reachable}
+
+
+class NormalizeSwitch(CompilerPass):
+    """Normalize branches like cond -> case a, case a + b, case a + 2b to ((cond - a) / b) -> case 0, case 1, case 2."""
+
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        for block in traverse_cfg_preorder(entry):
+            cases = {edge.cond for edge in block.outgoing}
+            if len(cases) <= 2:
+                continue
+            assert None in cases, "Non-terminal blocks should always have a default edge"
+            cases.remove(None)
+            offset, stride = self.get_offset_stride(cases)
+            if offset is None or (offset == 0 and stride == 1):
+                continue
+            for edge in block.outgoing:
+                if edge.cond is None:
+                    continue
+                edge.cond = (edge.cond - offset) // stride
+            if offset != 0:
+                block.test = IRPureInstr(Op.Subtract, [block.test, IRConst(offset)])
+            if stride != 1:
+                block.test = IRPureInstr(Op.Divide, [block.test, IRConst(stride)])
+        return entry
+
+    def get_offset_stride(self, cases: set[int]) -> tuple[int | None, int | None]:
+        cases = sorted(cases)
+        offset = cases[0]
+        stride = cases[1] - offset
+        if int(offset) != offset or int(stride) != stride:
+            return None, None
+        for i, case in enumerate(cases[2:], 2):
+            if case != offset + i * stride:
+                return None, None
+        return offset, stride

sonolus/backend/optimize/ssa.py

@@ -0,0 +1,200 @@
+from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet, IRStmt
+from sonolus.backend.optimize.dominance import DominanceFrontiers, get_df, get_dom_children
+from sonolus.backend.optimize.flow import BasicBlock, FlowEdge, traverse_cfg_preorder
+from sonolus.backend.optimize.passes import CompilerPass
+from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
+
+
+class ToSSA(CompilerPass):
+    def requires(self) -> set[CompilerPass]:
+        return {DominanceFrontiers()}
+
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        defs = self.defs_to_blocks(entry)
+        self.insert_phis(defs)
+        self.rename(entry, defs, {var: [] for var in defs}, {})
+        self.remove_placeholder_phis(entry)
+        return entry
+
+    def rename(
+        self,
+        block: BasicBlock,
+        defs: dict[TempBlock, set[BasicBlock]],
+        ssa_places: dict[TempBlock, list[SSAPlace]],
+        used: dict[str, int],
+    ):
+        to_pop = []
+        for var, args in [*block.phis.items()]:
+            if isinstance(var, SSAPlace):
+                continue
+            ssa_places[var].append(self.get_new_ssa_place(var.name, used))
+            to_pop.append(var)
+            block.phis[ssa_places[var][-1]] = args
+        block.statements = [self.rename_stmt(stmt, ssa_places, used, to_pop) for stmt in block.statements]
+        for edge in block.outgoing:
+            dst = edge.dst
+            for var, args in dst.phis.items():
+                if isinstance(var, SSAPlace):
+                    continue
+                if ssa_places[var]:
+                    args[block] = ssa_places[var][-1]
+        block.test = self.rename_stmt(block.test, ssa_places, used, to_pop)
+        for dom_child in get_dom_children(block):
+            self.rename(dom_child, defs, ssa_places, used)
+        for var in to_pop:
+            ssa_places[var].pop()
+
+    def remove_placeholder_phis(self, entry: BasicBlock):
+        for block in traverse_cfg_preorder(entry):
+            block.phis = {var: args for var, args in block.phis.items() if isinstance(var, SSAPlace)}
+
+    def rename_stmt(
+        self, stmt: IRStmt, ssa_places: dict[TempBlock, list[SSAPlace]], used: dict[str, int], to_pop: list[SSAPlace]
+    ):
+        match stmt:
+            case IRConst():
+                return stmt
+            case IRPureInstr(op=op, args=args):
+                return IRPureInstr(op=op, args=[self.rename_stmt(arg, ssa_places, used, to_pop) for arg in args])
+            case IRInstr(op=op, args=args):
+                return IRInstr(op=op, args=[self.rename_stmt(arg, ssa_places, used, to_pop) for arg in args])
+            case IRGet(place=place):
+                return IRGet(place=self.rename_stmt(place, ssa_places, used, to_pop))
+            case IRSet(place=place, value=value):
+                value = self.rename_stmt(value, ssa_places, used, to_pop)
+                if isinstance(place, BlockPlace) and isinstance(place.block, TempBlock) and place.block.size == 1:
+                    ssa_places[place.block].append(self.get_new_ssa_place(place.block.name, used))
+                    to_pop.append(place.block)
+                place = self.rename_stmt(place, ssa_places, used, to_pop)
+                return IRSet(place=place, value=value)
+            case SSAPlace():
+                return stmt
+            case TempBlock() if stmt.size == 1:
+                if stmt not in ssa_places or not ssa_places[stmt]:
+                    # This is an access to a definitely undefined variable
+                    # But it might not be reachable in reality, so we should allow it
+                    # Maybe there should be an error if this still happens after optimization,
+                    # but recovering the location of the error in the original code is hard.
+                    # This can happen in places like matching a VarArray[Num, 1] which was just created.
+                    # IR generation won't immediately fold a check that size > 0 to false, so here we
+                    # might see an access to uninitialized memory even though it's not reachable in reality.
+                    return SSAPlace("err", 0)
+                return ssa_places[stmt][-1]
+            case TempBlock():
+                return stmt
+            case int():
+                return stmt
+            case BlockPlace(block=block, index=index, offset=offset):
+                if isinstance(block, TempBlock) and block.size == 1:
+                    return self.rename_stmt(block, ssa_places, used, to_pop)
+                return BlockPlace(
+                    block=self.rename_stmt(block, ssa_places, used, to_pop),
+                    index=self.rename_stmt(index, ssa_places, used, to_pop),
+                    offset=self.rename_stmt(offset, ssa_places, used, to_pop),
+                )
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")
+
+    def insert_phis(self, defs: dict[TempBlock, set[BasicBlock]]):
+        for var, blocks in defs.items():
+            df = self.get_iterated_df(blocks)
+            for block in df:
+                block.phis[var] = {}
+
+    def defs_to_blocks(self, entry: BasicBlock) -> dict[TempBlock, set[BasicBlock]]:
+        result = {}
+        for block in traverse_cfg_preorder(entry):
+            for stmt in block.statements:
+                def_block = self.get_stmt_def(stmt)
+                if def_block is not None:
+                    result.setdefault(def_block, set()).add(block)
+        return result
+
+    def get_stmt_def(self, stmt: IRStmt) -> TempBlock:
+        if (
+            isinstance(stmt, IRSet)
+            and isinstance(stmt.place, BlockPlace)
+            and isinstance(stmt.place.block, TempBlock)
+            and stmt.place.block.size == 1
+        ):
+            return stmt.place.block
+        return None
+
+    def get_iterated_df(self, blocks: set[BasicBlock]) -> set[BasicBlock]:
+        df = set()
+        worklist = set(blocks)
+        while worklist:
+            block = worklist.pop()
+            new_df = get_df(block) - df
+            if new_df:
+                df.update(new_df)
+                worklist.update(new_df)
+        return df
+
+    def get_new_ssa_place(self, name: str, used: dict[str, int]) -> SSAPlace:
+        if name not in used:
+            used[name] = 0
+        used[name] += 1
+        return SSAPlace(name, used[name])
+
+
+class FromSSA(CompilerPass):
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        for block in [*traverse_cfg_preorder(entry)]:
+            self.process_block(block)
+        return entry
+
+    def process_block(self, block: BasicBlock):
+        incoming = [*block.incoming]
+        block.incoming.clear()
+        for edge in incoming:
+            between_block = BasicBlock()
+            edge.dst = between_block
+            between_block.incoming.add(edge)
+            next_edge = FlowEdge(between_block, block, None)
+            block.incoming.add(next_edge)
+            between_block.outgoing.add(next_edge)
+            for args in block.phis.values():
+                if edge.src in args:
+                    args[between_block] = args.pop(edge.src)
+        for var, args in block.phis.items():
+            for src, arg in args.items():
+                src.statements.append(
+                    IRSet(place=self.place_from_ssa_place(var), value=IRGet(place=self.place_from_ssa_place(arg)))
+                )
+        block.phis = {}
+        block.statements = [self.process_stmt(stmt) for stmt in block.statements]
+        block.test = self.process_stmt(block.test)
+
+    def process_stmt(self, stmt: IRStmt):
+        match stmt:
+            case IRConst():
+                return stmt
+            case IRPureInstr(op=op, args=args):
+                return IRPureInstr(op=op, args=[self.process_stmt(arg) for arg in args])
+            case IRInstr(op=op, args=args):
+                return IRInstr(op=op, args=[self.process_stmt(arg) for arg in args])
+            case IRGet(place=place):
+                return IRGet(place=self.process_stmt(place))
+            case IRSet(place=place, value=value):
+                return IRSet(place=self.process_stmt(place), value=self.process_stmt(value))
+            case SSAPlace():
+                return self.place_from_ssa_place(stmt)
+            case TempBlock():
+                return stmt
+            case int():
+                return stmt
+            case BlockPlace(block=block, index=index, offset=offset):
+                return BlockPlace(
+                    block=self.process_stmt(block),
+                    index=self.process_stmt(index),
+                    offset=self.process_stmt(offset),
+                )
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")
+
+    def temp_block_from_ssa_place(self, ssa_place: SSAPlace) -> TempBlock:
+        return TempBlock(f"{ssa_place.name}.{ssa_place.num}")
+
+    def place_from_ssa_place(self, ssa_place: SSAPlace) -> BlockPlace:
+        return BlockPlace(block=self.temp_block_from_ssa_place(ssa_place), index=0, offset=0)

sonolus/backend/place.py

@@ -1,5 +1,5 @@
 from collections.abc import Iterator
-from typing import Self
+from typing import NamedTuple, Self
 
 from sonolus.backend.blocks import Block
 
@@ -8,13 +8,9 @@ type BlockValue = Block | int | TempBlock | Place
 type IndexValue = int | Place
 
 
-class TempBlock:
+class TempBlock(NamedTuple):
     name: str
-    size: int
-
-    def __init__(self, name: str, size: int = 1):
-        self.name = name
-        self.size = size
+    size: int = 1
 
     def __repr__(self):
         return f"TempBlock(name={self.name!r}, size={self.size!r})"
@@ -33,19 +29,14 @@ class TempBlock:
         return isinstance(other, TempBlock) and self.name == other.name and self.size == other.size
 
     def __hash__(self):
-        return hash((self.name, self.size))
+        return hash(self.name)  # Typically will be unique by name alone
 
 
-class BlockPlace:
+class BlockPlace(NamedTuple):
     block: BlockValue
-    index: IndexValue
+    index: IndexValue = 0
     offset: int = 0
 
-    def __init__(self, block: BlockValue, index: IndexValue = 0, offset: int = 0):
-        self.block = block
-        self.index = index
-        self.offset = offset
-
     def __repr__(self):
         return f"BlockPlace(block={self.block!r}, index={self.index!r}, offset={self.offset!r})"
 
@@ -59,24 +50,25 @@ class BlockPlace:
         else:
             return f"{self.block}[{self.index} + {self.offset}]"
 
+    def add_offset(self, offset: int) -> Self:
+        return BlockPlace(self.block, self.index, self.offset + offset)
+
     def __eq__(self, other):
-        return isinstance(other, BlockPlace) and self.block == other.block and self.index == other.index
+        return (
+            isinstance(other, BlockPlace)
+            and self.block == other.block
+            and self.index == other.index
+            and self.offset == other.offset
+        )
 
     def __hash__(self):
-        return hash((self.block, self.index))
+        return hash((self.block, self.index, self.offset))
 
-    def add_offset(self, offset: int) -> Self:
-        return BlockPlace(self.block, self.index, self.offset + offset)
 
-
-class SSAPlace:
+class SSAPlace(NamedTuple):
     name: str
     num: int
 
-    def __init__(self, name: str, num: int):
-        self.name = name
-        self.num = num
-
     def __repr__(self):
         return f"SSAPlace(name={self.name!r}, num={self.num!r})"
 

sonolus/backend/utils.py

@@ -1,48 +1,58 @@
-# ruff: noqa: N802
-import ast
-import inspect
-from collections.abc import Callable
-from pathlib import Path
-
-
-def get_function(fn: Callable) -> tuple[str, ast.FunctionDef]:
-    # This preserves both line number and column number in the returned node
-    source_file = inspect.getsourcefile(fn)
-    _, start_line = inspect.getsourcelines(fn)
-    base_tree = ast.parse(Path(source_file).read_text(encoding="utf-8"))
-    return source_file, find_function(base_tree, start_line)
-
-
-class FindFunction(ast.NodeVisitor):
-    def __init__(self, line):
-        self.line = line
-        self.node: ast.FunctionDef | None = None
-
-    def visit_FunctionDef(self, node: ast.FunctionDef):
-        if node.lineno == self.line or (
-            node.decorator_list and (node.decorator_list[-1].end_lineno <= self.line <= node.lineno)
-        ):
-            self.node = node
-        else:
-            self.generic_visit(node)
-
-
-def find_function(tree: ast.Module, line: int):
-    visitor = FindFunction(line)
-    visitor.visit(tree)
-    return visitor.node
-
-
-class ScanWrites(ast.NodeVisitor):
-    def __init__(self):
-        self.writes = []
-
-    def visit_Name(self, node):
-        if isinstance(node.ctx, ast.Store | ast.Delete):
-            self.writes.append(node.id)
-
-
-def scan_writes(node: ast.AST) -> set[str]:
-    visitor = ScanWrites()
-    visitor.visit(node)
-    return set(visitor.writes)
+# ruff: noqa: N802
+import ast
+import inspect
+from collections.abc import Callable
+from functools import cache
+from pathlib import Path
+
+
+@cache
+def get_function(fn: Callable) -> tuple[str, ast.FunctionDef]:
+    # This preserves both line number and column number in the returned node
+    source_file = inspect.getsourcefile(fn)
+    _, start_line = inspect.getsourcelines(fn)
+    base_tree = ast.parse(Path(source_file).read_text(encoding="utf-8"))
+    return source_file, find_function(base_tree, start_line)
+
+
+class FindFunction(ast.NodeVisitor):
+    def __init__(self, line):
+        self.line = line
+        self.node: ast.FunctionDef | None = None
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        if node.lineno == self.line or (
+            node.decorator_list and (node.decorator_list[-1].end_lineno <= self.line <= node.lineno)
+        ):
+            self.node = node
+        else:
+            self.generic_visit(node)
+
+    def visit_Lambda(self, node: ast.Lambda):
+        if node.lineno == self.line:
+            if self.node is not None:
+                raise ValueError("Multiple functions defined on the same line are not supported")
+            self.node = node
+        else:
+            self.generic_visit(node)
+
+
+def find_function(tree: ast.Module, line: int):
+    visitor = FindFunction(line)
+    visitor.visit(tree)
+    return visitor.node
+
+
+class ScanWrites(ast.NodeVisitor):
+    def __init__(self):
+        self.writes = []
+
+    def visit_Name(self, node):
+        if isinstance(node.ctx, ast.Store | ast.Delete):
+            self.writes.append(node.id)
+
+
+def scan_writes(node: ast.AST) -> set[str]:
+    visitor = ScanWrites()
+    visitor.visit(node)
+    return set(visitor.writes)