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

sonolus/backend/ssa.py

@@ -0,0 +1,187 @@
+from sonolus.backend.dominance import DominanceFrontiers, get_df, get_dom_children
+from sonolus.backend.flow import BasicBlock, FlowEdge, traverse_cfg_preorder
+from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet, IRStmt
+from sonolus.backend.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],
+    ):
+        original_ssa_place_lens = {var: len(ssa_places[var]) for var in defs}
+        for var, args in [*block.phis.items()]:
+            if isinstance(var, SSAPlace):
+                continue
+            ssa_places[var].append(self.get_new_ssa_place(var.name, used))
+            block.phis[ssa_places[var][-1]] = args
+        block.statements = [self.rename_stmt(stmt, ssa_places, used) 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)
+        for dom_child in get_dom_children(block):
+            self.rename(dom_child, defs, ssa_places, used)
+        for var, length in original_ssa_place_lens.items():
+            ssa_places[var] = ssa_places[var][:length]
+
+    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]):
+        match stmt:
+            case IRConst():
+                return stmt
+            case IRPureInstr(op=op, args=args):
+                return IRPureInstr(op=op, args=[self.rename_stmt(arg, ssa_places, used) for arg in args])
+            case IRInstr(op=op, args=args):
+                return IRInstr(op=op, args=[self.rename_stmt(arg, ssa_places, used) for arg in args])
+            case IRGet(place=place):
+                return IRGet(place=self.rename_stmt(place, ssa_places, used))
+            case IRSet(place=place, value=value):
+                value = self.rename_stmt(value, ssa_places, used)
+                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))
+                place = self.rename_stmt(place, ssa_places, used)
+                return IRSet(place=place, value=value)
+            case SSAPlace():
+                return stmt
+            case TempBlock() if stmt.size == 1:
+                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)
+                return BlockPlace(
+                    block=self.rename_stmt(block, ssa_places, used),
+                    index=self.rename_stmt(index, ssa_places, used),
+                    offset=self.rename_stmt(offset, ssa_places, used),
+                )
+            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/utils.py

@@ -1,48 +1,48 @@
-# 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 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)