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

sonolus/backend/optimize/constant_evaluation.py

@@ -0,0 +1,374 @@
+# ruff: noqa: PLR1702
+import functools
+import math
+import operator
+from typing import ClassVar
+
+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.optimize.flow import BasicBlock, FlowEdge, traverse_cfg_preorder
+from sonolus.backend.optimize.passes import CompilerPass
+from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
+
+
+class Undefined:
+    pass
+
+
+class NotAConstant:
+    pass
+
+
+UNDEF = Undefined()
+NAC = NotAConstant()
+
+
+type Value = float | Undefined | NotAConstant
+
+
+class SparseConditionalConstantPropagation(CompilerPass):
+    SUPPORTED_OPS: ClassVar[set[Op]] = {
+        Op.Equal,
+        Op.NotEqual,
+        Op.Greater,
+        Op.GreaterOr,
+        Op.Less,
+        Op.LessOr,
+        Op.Not,
+        Op.And,
+        Op.Or,
+        Op.Negate,
+        Op.Add,
+        Op.Subtract,
+        Op.Multiply,
+        Op.Divide,
+        Op.Power,
+        Op.Log,
+        Op.Ceil,
+        Op.Floor,
+        Op.Round,
+        Op.Frac,
+        Op.Mod,
+        Op.Rem,
+        Op.Sin,
+        Op.Cos,
+        Op.Tan,
+        Op.Sinh,
+        Op.Cosh,
+        Op.Tanh,
+        Op.Arcsin,
+        Op.Arccos,
+        Op.Arctan,
+        Op.Arctan2,
+    }
+
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        ssa_edges: dict[SSAPlace, set[SSAPlace | BasicBlock]] = {}
+        executable_edges: set[FlowEdge] = set()
+
+        # BasicBlock key means the block's test
+        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()
+
+        for block in traverse_cfg_preorder(entry):
+            incoming_by_src = {edge.src: edge for edge in block.incoming}
+            for p, args in block.phis.items():
+                if not isinstance(p, SSAPlace):
+                    continue
+                defs[p] = {incoming_by_src[b]: v for b, v in args.items()}
+                values[p] = UNDEF
+                for arg in args.values():
+                    ssa_edges.setdefault(arg, set()).add(p)
+            for stmt in block.statements:
+                if isinstance(stmt, IRSet) and isinstance(stmt.place, SSAPlace):
+                    defs[stmt.place] = stmt.value
+                    places_to_blocks[stmt.place] = block
+                    values[stmt.place] = UNDEF
+                    for dep in self.get_dependencies(stmt.value, set()):
+                        ssa_edges.setdefault(dep, set()).add(stmt.place)
+            defs[block] = block.test
+            values[block] = UNDEF
+            for dep in self.get_dependencies(block.test, set()):
+                ssa_edges.setdefault(dep, set()).add(block)
+
+        def visit_phi(p):
+            arg_values = [values[v] if b in executable_edges else UNDEF for b, v in defs[p].items()]
+            distinct_defined_arg_values = {arg for arg in arg_values if arg is not UNDEF}
+            value = values[p]
+            if len(distinct_defined_arg_values) == 1:
+                new_value = distinct_defined_arg_values.pop()
+            elif len(distinct_defined_arg_values) > 1:
+                new_value = NAC
+            else:
+                new_value = UNDEF
+            if new_value != value:
+                values[p] = new_value
+                ssa_worklist.update(ssa_edges.get(p, set()))
+
+        flow_worklist: set[FlowEdge] = {FlowEdge(entry, entry, None)}
+        ssa_worklist: set[SSAPlace | BasicBlock] = set()
+        while flow_worklist or ssa_worklist:
+            while flow_worklist:
+                edge = flow_worklist.pop()
+                if edge in executable_edges:
+                    continue
+                executable_edges.add(edge)
+                block: BasicBlock = edge.dst
+                for p in block.phis:
+                    visit_phi(p)
+                is_first_visit = sum(edge in executable_edges for edge in block.incoming) <= 1
+                if is_first_visit:
+                    for stmt in block.statements:
+                        if not (isinstance(stmt, IRSet) and isinstance(stmt.place, SSAPlace)):
+                            continue
+                        value = values[stmt.place]
+                        new_value = self.evaluate_stmt(stmt.value, values)
+                        if new_value != value:
+                            values[stmt.place] = new_value
+                            ssa_worklist.update(ssa_edges.get(stmt.place, set()))
+                    test_value = values[block]
+                    new_test_value = self.evaluate_stmt(block.test, values)
+                    if new_test_value != test_value:
+                        assert new_test_value is not UNDEF
+                        values[block] = new_test_value
+                        if new_test_value is NAC:
+                            flow_worklist.update(block.outgoing)
+                            reachable_blocks.update(e.dst for e in block.outgoing)
+                        else:
+                            taken_edge = next(
+                                (edge for edge in block.outgoing if edge.cond == new_test_value), None
+                            ) or next((edge for edge in block.outgoing if edge.cond is None), None)
+                            if taken_edge:
+                                flow_worklist.add(taken_edge)
+                                reachable_blocks.add(taken_edge.dst)
+                    elif len(block.outgoing) == 1 and next(iter(block.outgoing)).cond is None:
+                        flow_worklist.update(block.outgoing)
+                        reachable_blocks.update(e.dst for e in block.outgoing)
+            while ssa_worklist:
+                p = ssa_worklist.pop()
+                defn = defs[p]
+                if isinstance(defn, dict):
+                    # This is a phi
+                    visit_phi(p)
+                elif isinstance(p, BasicBlock):
+                    # This is the block's test
+                    test_value = values[p]
+                    new_test_value = self.evaluate_stmt(defn, values)
+                    if new_test_value != test_value:
+                        assert new_test_value is not UNDEF
+                        values[p] = new_test_value
+                        if new_test_value is NAC:
+                            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), 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)
+                else:
+                    # This is a regular SSA assignment
+                    if places_to_blocks[p] not in reachable_blocks:
+                        continue
+                    value = values[p]
+                    new_value = self.evaluate_stmt(defn, values)
+                    if new_value != value:
+                        values[p] = new_value
+                        ssa_worklist.update(ssa_edges.get(p, set()))
+
+        for block in traverse_cfg_preorder(entry):
+            block.statements = [self.substitute_constants(stmt, values) for stmt in block.statements]
+            block.test = self.substitute_constants(block.test, values)
+
+        return entry
+
+    def get_dependencies(self, stmt, dependencies: set[SSAPlace]):
+        match stmt:
+            case IRConst():
+                pass
+            case IRPureInstr(op=_, args=args) | IRInstr(op=_, args=args):
+                for arg in args:
+                    self.get_dependencies(arg, dependencies)
+            case IRGet(place=SSAPlace() as place):
+                dependencies.add(place)
+            case IRGet(place=BlockPlace() as place):
+                self.get_dependencies(place.block, dependencies)
+            case BlockPlace(block=block, index=index, offset=_):
+                self.get_dependencies(block, dependencies)
+                self.get_dependencies(index, dependencies)
+            case SSAPlace():
+                dependencies.add(stmt)
+            case int() | float() | TempBlock():
+                pass
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")
+        return dependencies
+
+    def substitute_constants(self, stmt, values: dict[SSAPlace, Value]):
+        match stmt:
+            case IRConst():
+                return stmt
+            case IRPureInstr(op=op, args=args):
+                return IRPureInstr(op=op, args=[self.substitute_constants(arg, values) for arg in args])
+            case IRInstr(op=op, args=args):
+                return IRInstr(op=op, args=[self.substitute_constants(arg, values) for arg in args])
+            case IRGet(place=SSAPlace() as place):
+                value = values[place]
+                if isinstance(value, int | float):
+                    return IRConst(value)
+                return stmt
+            case IRGet(place=place):
+                return IRGet(place=self.substitute_constants(place, values))
+            case IRSet(place=SSAPlace() as place, value=value):
+                return IRSet(place=place, value=self.substitute_constants(value, values))
+            case IRSet(place=place, value=value):
+                return IRSet(
+                    place=self.substitute_constants(place, values), value=self.substitute_constants(value, values)
+                )
+            case BlockPlace(block=block, index=index, offset=offset):
+                return BlockPlace(
+                    block=self.substitute_constants(block, values),
+                    index=self.substitute_constants(index, values),
+                    offset=offset,
+                )
+            case SSAPlace():
+                value = values[stmt]
+                if isinstance(value, int | float):
+                    return IRConst(value)
+                return stmt
+            case int() | float() | TempBlock():
+                return stmt
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")
+
+    def evaluate_stmt(self, stmt, values: dict[SSAPlace, Value]) -> Value:
+        match stmt:
+            case IRConst(value=value):
+                return value
+            case IRPureInstr(op=op, args=args) | IRInstr(op=op, args=args):
+                if op not in self.SUPPORTED_OPS:
+                    return NAC
+                args = [self.evaluate_stmt(arg, values) for arg in args]
+                match op:
+                    case Op.And:
+                        if any(arg == 0 for arg in args):
+                            return 0
+                    case Op.Or:
+                        if any(arg == 1 for arg in args):
+                            return 1
+                    case Op.Multiply:
+                        if any(arg == 0 for arg in args):
+                            return 0
+                if any(arg is NAC for arg in args):
+                    return NAC
+                if any(arg is UNDEF for arg in args):
+                    return UNDEF
+                match op:
+                    case Op.Equal:
+                        assert len(args) == 2
+                        return args[0] == args[1]
+                    case Op.NotEqual:
+                        assert len(args) == 2
+                        return args[0] != args[1]
+                    case Op.Greater:
+                        assert len(args) == 2
+                        return args[0] > args[1]
+                    case Op.GreaterOr:
+                        assert len(args) == 2
+                        return args[0] >= args[1]
+                    case Op.Less:
+                        assert len(args) == 2
+                        return args[0] < args[1]
+                    case Op.LessOr:
+                        assert len(args) == 2
+                        return args[0] <= args[1]
+                    case Op.Not:
+                        assert len(args) == 1
+                        return int(not args[0])
+                    case Op.And:
+                        return all(args)
+                    case Op.Or:
+                        return any(args)
+                    case Op.Negate:
+                        assert len(args) == 1
+                        return -args[0]
+                    case Op.Add:
+                        return sum(args)
+                    case Op.Subtract:
+                        if len(args) == 0:
+                            return 0
+                        return args[0] - sum(args[1:])
+                    case Op.Multiply:
+                        if len(args) == 0:
+                            return 1
+                        return functools.reduce(operator.mul, args, 1)
+                    case Op.Divide:
+                        if len(args) == 0:
+                            return 1
+                        return args[0] / functools.reduce(operator.mul, args[1:], 1)
+                    case Op.Power:
+                        if len(args) == 0:
+                            return 1
+                        return functools.reduce(operator.pow, args)
+                    case Op.Log:
+                        assert len(args) == 2
+                        return math.log(args[0], args[1])
+                    case Op.Ceil:
+                        assert len(args) == 1
+                        return math.ceil(args[0])
+                    case Op.Floor:
+                        assert len(args) == 1
+                        return math.floor(args[0])
+                    case Op.Round:
+                        assert len(args) == 1
+                        # This is round half to even in both Python and Sonolus
+                        return round(args[0])
+                    case Op.Frac:
+                        assert len(args) == 1
+                        return smath.frac(args[0])
+                    case Op.Mod:
+                        assert len(args) == 2
+                        return args[0] % args[1]
+                    case Op.Rem:
+                        assert len(args) == 2
+                        return smath.remainder(args[0], args[1])
+                    case Op.Sin:
+                        assert len(args) == 1
+                        return math.sin(args[0])
+                    case Op.Cos:
+                        assert len(args) == 1
+                        return math.cos(args[0])
+                    case Op.Tan:
+                        assert len(args) == 1
+                        return math.tan(args[0])
+                    case Op.Sinh:
+                        assert len(args) == 1
+                        return math.sinh(args[0])
+                    case Op.Cosh:
+                        assert len(args) == 1
+                        return math.cosh(args[0])
+                    case Op.Tanh:
+                        assert len(args) == 1
+                        return math.tanh(args[0])
+                    case Op.Arcsin:
+                        assert len(args) == 1
+                        return math.asin(args[0])
+                    case Op.Arccos:
+                        assert len(args) == 1
+                        return math.acos(args[0])
+                    case Op.Arctan:
+                        assert len(args) == 1
+                        return math.atan(args[0])
+                    case Op.Arctan2:
+                        assert len(args) == 2
+                        return math.atan2(args[0], args[1])
+            case IRGet(place=SSAPlace() as place):
+                return values[place]
+            case IRGet():
+                return NAC
+            case IRSet() | _:
+                raise TypeError(f"Unexpected statement: {stmt}")

sonolus/backend/optimize/copy_coalesce.py

@@ -0,0 +1,85 @@
+from sonolus.backend.ir import IRConst, IRGet, IRInstr, IRPureInstr, IRSet
+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
+
+
+class CopyCoalesce(CompilerPass):
+    def requires(self) -> set[CompilerPass]:
+        return {LivenessAnalysis()}
+
+    def run(self, entry: BasicBlock) -> BasicBlock:
+        mapping = self.get_mapping(entry)
+        for block in traverse_cfg_preorder(entry):
+            block.statements = [self.apply_to_stmt(stmt, mapping) for stmt in block.statements]
+            block.test = self.apply_to_stmt(block.test, mapping)
+        return entry
+
+    def apply_to_stmt(self, stmt, mapping: dict[TempBlock, TempBlock]):
+        match stmt:
+            case IRConst():
+                return stmt
+            case IRSet(place=place, value=value):
+                return IRSet(self.apply_to_stmt(place, mapping), self.apply_to_stmt(value, mapping))
+            case IRGet(place=place):
+                return IRGet(self.apply_to_stmt(place, mapping))
+            case IRPureInstr(op=op, args=args):
+                return IRPureInstr(op, [self.apply_to_stmt(arg, mapping) for arg in args])
+            case IRInstr(op=op, args=args):
+                return IRInstr(op, [self.apply_to_stmt(arg, mapping) for arg in args])
+            case BlockPlace(block=block, index=index, offset=offset):
+                return BlockPlace(self.apply_to_stmt(block, mapping), self.apply_to_stmt(index, mapping), offset)
+            case TempBlock():
+                return mapping.get(stmt, stmt)
+            case SSAPlace() | int() | float():
+                return stmt
+
+    def get_mapping(self, entry: BasicBlock) -> dict[TempBlock, TempBlock]:
+        interference = self.get_interference(entry)
+        copies = self.get_copies(entry)
+
+        mapping = {}
+
+        for target, sources in copies.items():
+            for source in sources:
+                if source in mapping:
+                    continue
+                if source in interference.get(target, set()):
+                    continue
+                mapping[source] = mapping.get(target, target)
+                combined_interference = interference.get(target, set()) | interference.get(source, set())
+                interference[source] = combined_interference
+                interference[target] = combined_interference
+
+        return mapping
+
+    def get_interference(self, entry: BasicBlock) -> dict[TempBlock, set[TempBlock]]:
+        result = {}
+        for block in traverse_cfg_preorder(entry):
+            for stmt in [*block.statements, block.test]:
+                live = {p for p in get_live(stmt) if isinstance(p, TempBlock) and p.size == 1}
+                for place in live:
+                    result.setdefault(place, set()).update(live - {place})
+        return result
+
+    def get_copies(self, entry: BasicBlock) -> dict[TempBlock, set[TempBlock]]:
+        result = {}
+        for block in traverse_cfg_preorder(entry):
+            for stmt in block.statements:
+                if (
+                    not isinstance(stmt, IRSet)
+                    or not isinstance(stmt.place, BlockPlace)
+                    or not isinstance(stmt.place.block, TempBlock)
+                    or stmt.place.block.size != 1
+                    or not isinstance(stmt.value, IRGet)
+                    or not isinstance(stmt.value.place, BlockPlace)
+                    or not isinstance(stmt.value.place.block, TempBlock)
+                    or stmt.value.place.block.size != 1
+                ):
+                    continue
+                target = stmt.place.block
+                source = stmt.value.place.block
+                result.setdefault(target, set()).add(source)
+                result.setdefault(source, set()).add(target)
+        return result

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