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

sonolus/backend/finalize.py

@@ -1,12 +1,12 @@
-from sonolus.backend.flow import BasicBlock, traverse_cfg_preorder
 from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet
 from sonolus.backend.node import ConstantNode, EngineNode, FunctionNode
 from sonolus.backend.ops import Op
+from sonolus.backend.optimize.flow import BasicBlock, traverse_cfg_reverse_postorder
 from sonolus.backend.place import BlockPlace
 
 
 def cfg_to_engine_node(entry: BasicBlock):
-    block_indexes = {block: i for i, block in enumerate(traverse_cfg_preorder(entry))}
+    block_indexes = {block: i for i, block in enumerate(traverse_cfg_reverse_postorder(entry))}
     block_statements = []
     for block in block_indexes:
         statements = []
@@ -28,16 +28,24 @@ def cfg_to_engine_node(entry: BasicBlock):
                         ],
                     )
                 )
-            case dict() as tgt:
+            case dict() as targets:
                 args = [ir_to_engine_node(block.test)]
                 default = len(block_indexes)
-                for cond, target in tgt.items():
-                    if cond is None:
-                        default = block_indexes[target]
-                    args.append(ConstantNode(value=cond))
-                    args.append(ConstantNode(value=block_indexes[target]))
-                args.append(ConstantNode(value=default))
-                statements.append(FunctionNode(Op.SwitchWithDefault, args))
+                conds = [cond for cond in targets if cond is not None]
+                if min(conds) == 0 and max(conds) == len(conds) - 1 and all(int(cond) == cond for cond in conds):
+                    args.extend(ConstantNode(value=block_indexes[targets[cond]]) for cond in range(len(conds)))
+                    if None in targets:
+                        default = block_indexes[targets[None]]
+                    args.append(ConstantNode(value=default))
+                    statements.append(FunctionNode(Op.SwitchIntegerWithDefault, args))
+                else:
+                    for cond, target in targets.items():
+                        if cond is None:
+                            default = block_indexes[target]
+                        args.append(ConstantNode(value=cond))
+                        args.append(ConstantNode(value=block_indexes[target]))
+                    args.append(ConstantNode(value=default))
+                    statements.append(FunctionNode(Op.SwitchWithDefault, args))
         block_statements.append(FunctionNode(Op.Execute, statements))
     block_statements.append(ConstantNode(value=0))
     return FunctionNode(Op.Block, [FunctionNode(Op.JumpLoop, block_statements)])

sonolus/backend/interpret.py

@@ -62,16 +62,16 @@ class Interpreter:
                 return self.run(default)
             case Op.SwitchInteger:
                 test, *branches = args
-                test_result = int(self.run(test))
-                if 0 <= test_result < len(branches):
+                test_result = self.run(test)
+                if 0 <= test_result < len(branches) and int(test_result) == test_result:
                     return self.run(branches[test_result])
                 else:
                     return 0.0
             case Op.SwitchIntegerWithDefault:
                 test, *branches, default = args
-                test_result = int(self.run(test))
-                if 0 <= test_result < len(branches):
-                    return self.run(branches[test_result])
+                test_result = self.run(test)
+                if 0 <= test_result < len(branches) and int(test_result) == test_result:
+                    return self.run(branches[int(test_result)])
                 else:
                     return self.run(default)
             case Op.While:
@@ -242,10 +242,10 @@ class Interpreter:
                 return math.radians(self.run(args[0]))
             case Op.Random:
                 lo, hi = (self.run(arg) for arg in args)
-                return lo + (hi - lo) * random.random()
+                return random.uniform(lo, hi)
             case Op.RandomInteger:
                 lo, hi = (self.ensure_int(self.run(arg)) for arg in args)
-                return random.randint(lo, hi)
+                return random.randrange(lo, hi)
             case Op.Rem:
                 return self.reduce_args(args, math.remainder)
             case Op.Remap:

sonolus/backend/ir.py

@@ -41,6 +41,12 @@ class IRPureInstr:
     def __str__(self):
         return f"{self.op.name}({', '.join(map(str, self.args))})"
 
+    def __eq__(self, other):
+        return isinstance(other, IRPureInstr) and self.op == other.op and self.args == other.args
+
+    def __hash__(self):
+        return hash((self.op, tuple(self.args)))
+
 
 class IRInstr:
     op: Op
@@ -56,6 +62,12 @@ class IRInstr:
     def __str__(self):
         return f"{self.op.name}({', '.join(map(str, self.args))})"
 
+    def __eq__(self, other):
+        return isinstance(other, IRInstr) and self.op == other.op and self.args == other.args
+
+    def __hash__(self):
+        return hash((self.op, tuple(self.args)))
+
 
 class IRGet:
     place: Place
@@ -69,6 +81,12 @@ class IRGet:
     def __str__(self):
         return f"{self.place}"
 
+    def __eq__(self, other):
+        return isinstance(other, IRGet) and self.place == other.place
+
+    def __hash__(self):
+        return hash(self.place)
+
 
 class IRSet:
     place: Place
@@ -89,3 +107,9 @@ class IRSet:
                 return f"{self.place} := {self.value}"
             case _:
                 raise TypeError(f"Invalid place: {self.place}")
+
+    def __eq__(self, other):
+        return isinstance(other, IRSet) and self.place == other.place and self.value == other.value
+
+    def __hash__(self):
+        return hash((self.place, self.value))

sonolus/backend/{allocate.py → optimize/allocate.py}

@@ -1,7 +1,7 @@
-from sonolus.backend.flow import BasicBlock, traverse_cfg_preorder
 from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet
-from sonolus.backend.liveness import LivenessAnalysis, get_live
-from sonolus.backend.passes import CompilerPass
+from sonolus.backend.optimize.flow import BasicBlock, traverse_cfg_preorder
+from sonolus.backend.optimize.liveness import LivenessAnalysis, get_live
+from sonolus.backend.optimize.passes import CompilerPass
 from sonolus.backend.place import BlockPlace, TempBlock
 
 TEMP_SIZE = 4096
@@ -86,6 +86,7 @@ class Allocate(CompilerPass):
                     if block.size == 0:
                         offset = -1
                     else:
+                        assert block in mapping, f"Missing mapping for {block}, dead variable not removed?"
                         offset = mapping[block] + offset
                     return BlockPlace(block=10000, index=self.update_stmt(index, mapping), offset=offset)
                 return BlockPlace(

sonolus/backend/{constant_evaluation.py → optimize/constant_evaluation.py}

@@ -4,11 +4,11 @@ import math
 import operator
 from typing import ClassVar
 
-import sonolus.script.math as smath
-from sonolus.backend.flow import BasicBlock, FlowEdge, traverse_cfg_preorder
+import sonolus.script.internal.math_impls as smath
 from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet, IRStmt
 from sonolus.backend.ops import Op
-from sonolus.backend.passes import CompilerPass
+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
 
 
@@ -68,7 +68,7 @@ class SparseConditionalConstantPropagation(CompilerPass):
         executable_edges: set[FlowEdge] = set()
 
         # BasicBlock key means the block's test
-        values: dict[SSAPlace | BasicBlock, Value] = {}
+        values: dict[SSAPlace | BasicBlock, Value] = {SSAPlace("err", 0): UNDEF}
         defs: dict[SSAPlace | BasicBlock, IRStmt | dict[FlowEdge, SSAPlace]] = {}
         places_to_blocks: dict[SSAPlace, BasicBlock] = {}
         reachable_blocks: set[BasicBlock] = set()
@@ -164,9 +164,9 @@ class SparseConditionalConstantPropagation(CompilerPass):
                             flow_worklist.update(p.outgoing)
                             reachable_blocks.update(e.dst for e in p.outgoing)
                         else:
-                            taken_edge = next(edge for edge in p.outgoing if edge.cond == new_test_value) or next(
-                                (edge for edge in p.outgoing if edge.cond is None), None
-                            )
+                            taken_edge = next(
+                                (edge for edge in p.outgoing if edge.cond == new_test_value), None
+                            ) or next((edge for edge in p.outgoing if edge.cond is None), None)
                             if taken_edge:
                                 flow_worklist.add(taken_edge)
                                 reachable_blocks.add(taken_edge.dst)

sonolus/backend/{coalesce.py → optimize/copy_coalesce.py}

@@ -1,7 +1,7 @@
-from sonolus.backend.flow import BasicBlock, traverse_cfg_preorder
 from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet
-from sonolus.backend.liveness import LivenessAnalysis, get_live
-from sonolus.backend.passes import CompilerPass
+from sonolus.backend.optimize.flow import BasicBlock, traverse_cfg_preorder
+from sonolus.backend.optimize.liveness import LivenessAnalysis, get_live
+from sonolus.backend.optimize.passes import CompilerPass
 from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
 
 

sonolus/backend/optimize/dead_code.py

@@ -0,0 +1,185 @@
+from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet, IRStmt
+from sonolus.backend.optimize.flow import BasicBlock, traverse_cfg_preorder
+from sonolus.backend.optimize.liveness import HasLiveness, LivenessAnalysis, get_live, get_live_phi_targets
+from sonolus.backend.optimize.passes import CompilerPass
+from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
+
+
+class UnreachableCodeElimination(CompilerPass):
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        original_blocks = [*traverse_cfg_preorder(entry)]
+        worklist = {entry}
+        visited = set()
+        while worklist:
+            block = worklist.pop()
+            if block in visited:
+                continue
+            visited.add(block)
+            match block.test:
+                case IRConst(value=value):
+                    block.test = IRConst(0)
+                    taken_edge = next(
+                        (edge for edge in block.outgoing if edge.cond == value),
+                        None,
+                    ) or next((edge for edge in block.outgoing if edge.cond is None), None)
+                    assert not block.outgoing or taken_edge
+                    for edge in [*block.outgoing]:
+                        if edge is not taken_edge:
+                            edge.dst.incoming.remove(edge)
+                            block.outgoing.remove(edge)
+                    if taken_edge:
+                        taken_edge.cond = None
+                        block.outgoing.add(taken_edge)
+                        worklist.add(taken_edge.dst)
+                case _:
+                    worklist.update(edge.dst for edge in block.outgoing)
+        for block in original_blocks:
+            if block not in visited:
+                for edge in block.outgoing:
+                    edge.dst.incoming.remove(edge)
+            else:
+                for args in block.phis.values():
+                    for src_block in [*args]:
+                        if src_block not in visited:
+                            args.pop(src_block)
+        return entry
+
+
+class DeadCodeElimination(CompilerPass):
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        uses = set()
+        defs = {}
+        for block in traverse_cfg_preorder(entry):
+            for statement in block.statements:
+                self.handle_statement(statement, uses, defs)
+            for target, args in block.phis.items():
+                if target not in defs:
+                    defs[target] = []
+                defs[target].append(tuple(args.values()))
+            self.update_uses(block.test, uses)
+
+        queue = [*uses]
+        while queue:
+            val = queue.pop()
+            if val not in defs:
+                continue
+            for stmt in defs[val]:
+                if isinstance(stmt, tuple):
+                    stmt_uses = stmt
+                else:
+                    stmt_uses = self.update_uses(stmt, set())
+                for use in stmt_uses:
+                    if use not in uses:
+                        uses.add(use)
+                        queue.append(use)
+
+        for block in traverse_cfg_preorder(entry):
+            live_stmts = []
+            for statement in block.statements:
+                match statement:
+                    case IRSet(place=place, value=value):
+                        is_live = not (
+                            (isinstance(place, SSAPlace) and place not in uses)
+                            or (
+                                isinstance(place, BlockPlace)
+                                and isinstance(place.block, TempBlock)
+                                and place.block not in uses
+                            )
+                            or (isinstance(value, IRGet) and place == value.place)
+                        )
+                        if is_live:
+                            live_stmts.append(statement)
+                        elif isinstance(value, IRInstr) and value.op.side_effects:
+                            live_stmts.append(value)
+                    case other:
+                        live_stmts.append(other)
+            block.statements = live_stmts
+            block.phis = {place: phi for place, phi in block.phis.items() if place in uses}
+        return entry
+
+    def handle_statement(
+        self,
+        stmt: IRStmt | BlockPlace | SSAPlace | TempBlock | int,
+        uses: set[HasLiveness],
+        defs: dict[HasLiveness, list[IRStmt | tuple[HasLiveness]]],
+    ):
+        if isinstance(stmt, IRSet):
+            place = stmt.place
+            value = stmt.value
+            if isinstance(place, SSAPlace):
+                if place not in defs:
+                    defs[place] = []
+                defs[place].append(stmt)
+                if isinstance(value, IRInstr) and value.op.side_effects:
+                    self.update_uses(value, uses)
+            elif isinstance(place, BlockPlace) and isinstance(place.block, TempBlock):
+                if place.block not in defs:
+                    defs[place.block] = []
+                defs[place.block].append(stmt)
+                if isinstance(value, IRInstr) and value.op.side_effects:
+                    self.update_uses(value, uses)
+            else:
+                self.update_uses(place, uses)
+                self.update_uses(value, uses)
+        else:
+            self.update_uses(stmt, uses)
+
+    def update_uses(
+        self, stmt: IRStmt | BlockPlace | SSAPlace | TempBlock | int, uses: set[HasLiveness]
+    ) -> set[HasLiveness]:
+        match stmt:
+            case IRPureInstr(op=_, args=args) | IRInstr(op=_, args=args):
+                for arg in args:
+                    self.update_uses(arg, uses)
+            case IRGet(place=place):
+                self.update_uses(place, uses)
+            case IRSet(place=place, value=value):
+                if isinstance(place, BlockPlace):
+                    if not isinstance(place.block, TempBlock):
+                        self.update_uses(place.block, uses)
+                    self.update_uses(place.index, uses)
+                self.update_uses(value, uses)
+            case IRConst() | int():
+                pass
+            case BlockPlace(block=block, index=index, offset=_):
+                self.update_uses(block, uses)
+                self.update_uses(index, uses)
+            case TempBlock() | SSAPlace():
+                uses.add(stmt)
+            case _:
+                raise TypeError(f"Unexpected statement type: {type(stmt)}")
+        return uses
+
+
+class AdvancedDeadCodeElimination(CompilerPass):
+    """Slower than regular DeadCodeElimination but can handle cases like definitions after the last use and so on."""
+
+    def requires(self) -> set[CompilerPass]:
+        return {LivenessAnalysis()}
+
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        for block in traverse_cfg_preorder(entry):
+            live_stmts = []
+            for statement in block.statements:
+                live = get_live(statement)
+                match statement:
+                    case IRSet(place=place, value=value):
+                        is_live = not (
+                            (isinstance(place, SSAPlace) and place not in live)
+                            or (
+                                isinstance(place, BlockPlace)
+                                and isinstance(place.block, TempBlock)
+                                and place.block not in live
+                            )
+                            or (isinstance(value, IRGet) and place == value.place)
+                        )
+                        if is_live:
+                            live_stmts.append(statement)
+                        elif isinstance(value, IRInstr) and value.op.side_effects:
+                            live_stmts.append(value)
+                            value.live = live
+                    case other:
+                        live_stmts.append(other)
+            block.statements = live_stmts
+            block.phis = {place: phi for place, phi in block.phis.items() if place in get_live_phi_targets(block)}
+        return entry

sonolus/backend/{dominance.py → optimize/dominance.py}

@@ -1,23 +1,8 @@
-from sonolus.backend.flow import (
+from sonolus.backend.optimize.flow import (
     BasicBlock,
     traverse_cfg_reverse_postorder,
 )
-from sonolus.backend.passes import CompilerPass
-
-# traverse_cfg_preorder(entry: BasicBlock) -> Iterator[BasicBlock]
-# traverse_cfg_postorder(entry: BasicBlock) -> Iterator[BasicBlock]
-# traverse_cfg_reverse_postorder(entry: BasicBlock) -> Iterator[BasicBlock]
-
-# class BasicBlock:
-#     phis: dict[SSAPlace, dict[Self, SSAPlace]]
-#     statements: list[IRStmt]
-#     test: IRExpr
-#     incoming: set[FlowEdge]
-#     outgoing: set[FlowEdge]
-
-# class FlowEdge:
-#     src: "BasicBlock"
-#     dst: "BasicBlock"
+from sonolus.backend.optimize.passes import CompilerPass
 
 
 class DominanceFrontiers(CompilerPass):

sonolus/backend/{flow.py → optimize/flow.py}

@@ -87,11 +87,10 @@ def cfg_to_mermaid(entry: BasicBlock):
         else:
             return "{}"
 
-    block_indexes = {block: i for i, block in enumerate(traverse_cfg_preorder(entry))}
+    block_indexes = {block: i for i, block in enumerate(traverse_cfg_reverse_postorder(entry))}
 
     lines = ["Entry([Entry]) --> 0"]
-    for block in traverse_cfg_preorder(entry):
-        index = block_indexes[block]
+    for block, index in block_indexes.items():
         lines.append(f"{index}[{pre(fmt([f'#{index}', *(
             f"{dst} := phi({", ".join(f"{block_indexes.get(src_block, "<dead>")}: {src_place}"
                                       for src_block, src_place
@@ -106,7 +105,7 @@ def cfg_to_mermaid(entry: BasicBlock):
             case {None: target, **other} if not other:
                 lines.append(f"{index} --> {block_indexes[target]}")
             case {0: f_branch, None: t_branch, **other} if not other:
-                lines.append(f"{index}_{{{pre(fmt([block.test]))}}}")
+                lines.append(f"{index}_{{{{{pre(fmt([block.test]))}}}}}")
                 lines.append(f"{index} --> {index}_")
                 lines.append(f"{index}_ --> |true| {block_indexes[t_branch]}")
                 lines.append(f"{index}_ --> |false| {block_indexes[f_branch]}")
@@ -114,7 +113,9 @@ def cfg_to_mermaid(entry: BasicBlock):
                 lines.append(f"{index}_{{{{{pre(fmt([block.test]))}}}}}")
                 lines.append(f"{index} --> {index}_")
                 for cond, target in tgt.items():
-                    lines.append(f"{index}_ --> |{pre(fmt([cond or "default"]))}| {block_indexes[target]}")
+                    lines.append(
+                        f"{index}_ --> |{pre(fmt([cond if cond is not None else "default"]))}| {block_indexes[target]}"
+                    )
     lines.append("Exit([Exit])")
 
     body = textwrap.indent("\n".join(lines), "    ")

sonolus/backend/{inlining.py → optimize/inlining.py}

@@ -1,6 +1,6 @@
-from sonolus.backend.flow import BasicBlock, traverse_cfg_preorder
 from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet, IRStmt
-from sonolus.backend.passes import CompilerPass
+from sonolus.backend.optimize.flow import BasicBlock, traverse_cfg_preorder
+from sonolus.backend.optimize.passes import CompilerPass
 from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
 
 
@@ -11,16 +11,7 @@ class InlineVars(CompilerPass):
 
         for block in traverse_cfg_preorder(entry):
             for stmt in block.statements:
-                if (
-                    isinstance(stmt, IRSet)
-                    and isinstance(stmt.place, SSAPlace)
-                    and isinstance(stmt.value, IRGet)
-                    and isinstance(stmt.value.place, SSAPlace)
-                ):
-                    # This is effectively an alias
-                    pass
-                else:
-                    self.count_uses(stmt, use_counts)
+                self.count_uses(stmt, use_counts)
                 if isinstance(stmt, IRSet) and isinstance(stmt.place, SSAPlace):
                     definitions[stmt.place] = stmt.value
             self.count_uses(block.test, use_counts)
@@ -50,11 +41,7 @@ class InlineVars(CompilerPass):
                 defn = self.substitute(defn, subs)
             definitions[p] = defn
 
-        valid = {
-            p
-            for p, count in use_counts.items()
-            if (p in definitions and count == 1 and self.is_inlinable(definitions[p]))
-        }
+        valid = {p for p in definitions if self.is_inlinable(definitions[p]) and use_counts.get(p, 0) <= 1}
 
         for block in traverse_cfg_preorder(entry):
             new_statements = []