sonolus.py 0.3.3__py3-none-any.whl → 0.4.0__py3-none-any.whl

sonolus/backend/optimize/inlining.py

@@ -1,11 +1,29 @@
+from sonolus.backend.blocks import BlockData
 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.passes import CompilerPass
+from sonolus.backend.optimize.passes import CompilerPass, OptimizerConfig
 from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
 
+RUNTIME_CONSTANT_BLOCKS = {
+    "RuntimeEnvironment",
+    "RuntimeUI",
+    "RuntimeUIConfiguration",
+    "LevelData",
+    "LevelOption",
+    "LevelBucket",
+    "LevelScore",
+    "LevelLife",
+    "EngineRom",
+    "ArchetypeLife",
+    "RuntimeCanvas",
+    "PreviewData",
+    "PreviewOption",
+    "TutorialData",
+}
+
 
 class InlineVars(CompilerPass):
-    def run(self, entry: BasicBlock) -> BasicBlock:
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
         use_counts: dict[SSAPlace, int] = {}
         definitions: dict[SSAPlace, IRStmt] = {}
 
@@ -15,49 +33,97 @@ class InlineVars(CompilerPass):
                 if isinstance(stmt, IRSet) and isinstance(stmt.place, SSAPlace):
                     definitions[stmt.place] = stmt.value
             self.count_uses(block.test, use_counts)
+            for tgt, args in block.phis.items():
+                for arg in args.values():
+                    self.count_uses(arg, use_counts)
+                if len(args) == 1:
+                    arg = next(iter(args.values()))
+                    definitions[tgt] = IRGet(place=arg)
 
+        for defn in definitions.values():
+            if isinstance(defn, IRGet) and isinstance(defn.place, SSAPlace):
+                use_counts[defn.place] -= 1
+
+        canonical_definitions: dict[SSAPlace, IRStmt] = {}
         for p, defn in definitions.items():
+            canonical_definitions[p] = defn
+            # Update the definition if it's a Get from another SSAPlace until we reach a definition that is not a Get
+            while defn and isinstance(defn, IRGet) and isinstance(defn.place, SSAPlace):
+                canonical_definitions[p] = defn
+                defn = definitions.get(defn.place, None)  # Can be None if it's a phi
+            canonical_defn = canonical_definitions[p]
+            if (
+                use_counts.get(p, 0) > 0
+                and isinstance(canonical_defn, IRGet)
+                and isinstance(canonical_defn.place, SSAPlace)
+            ):
+                use_counts[canonical_defn.place] = use_counts.get(canonical_defn.place, 0) + 1
+
+        for p, defn in canonical_definitions.items():
+            if isinstance(defn, IRGet) and isinstance(defn.place, SSAPlace):
+                inner_p = defn.place
+                inner_defn = canonical_definitions.get(inner_p)
+                if (
+                    inner_defn
+                    and self.is_inlinable(inner_defn, config.callback)
+                    and (use_counts.get(inner_p, 0) <= 1 or self.is_free_to_inline(inner_defn, config.callback))
+                ):
+                    canonical_definitions[p] = inner_defn
+
+        inlined_definitions = {**canonical_definitions}
+        for p, defn in canonical_definitions.items():
             while True:
-                if isinstance(defn, IRGet) and isinstance(defn.place, SSAPlace) and defn.place in definitions:
-                    inside_defn = definitions[defn.place]
-                    if not self.is_inlinable(inside_defn):
-                        break
-                    defn = inside_defn
-                    continue
                 inlinable_uses = self.get_inlinable_uses(defn, set())
                 subs = {}
                 for inside_p in inlinable_uses:
-                    if inside_p not in definitions:
+                    if inside_p not in canonical_definitions:
                         continue
-                    inside_defn = definitions[inside_p]
-                    if not self.is_inlinable(inside_defn):
+                    inside_defn = canonical_definitions[inside_p]
+                    if not self.is_inlinable(inside_defn, config.callback):
                         continue
-                    if (isinstance(inside_defn, IRGet) and isinstance(inside_defn.place, SSAPlace)) or use_counts[
-                        inside_p
-                    ] == 1:
+                    if (
+                        (isinstance(inside_defn, IRGet) and isinstance(inside_defn.place, SSAPlace))
+                        or use_counts[inside_p] == 1
+                        or self.is_free_to_inline(inside_defn, config.callback)
+                    ):
                         subs[inside_p] = inside_defn
                 if not subs:
                     break
                 defn = self.substitute(defn, subs)
-            definitions[p] = defn
+            inlined_definitions[p] = defn
 
-        valid = {p for p in definitions if self.is_inlinable(definitions[p]) and use_counts.get(p, 0) <= 1}
+        valid = {
+            p
+            for p in inlined_definitions
+            if self.is_inlinable(inlined_definitions[p], config.callback)
+            and (use_counts.get(p, 0) <= 1 or self.is_free_to_inline(inlined_definitions[p], config.callback))
+        }
 
         for block in traverse_cfg_preorder(entry):
             new_statements = []
             for stmt in [*block.statements, block.test]:
-                inlinable_uses = self.get_inlinable_uses(stmt, set())
-                subs = {}
-                for p in inlinable_uses:
-                    if p not in valid:
-                        continue
-                    definition = definitions[p]
-                    subs[p] = definition
-
-                if subs:
-                    new_statements.append(self.substitute(stmt, subs))
-                else:
+                if (
+                    isinstance(stmt, IRSet)
+                    and isinstance(stmt.place, SSAPlace)
+                    and isinstance(stmt.value, IRGet)
+                    and isinstance(stmt.value.place, SSAPlace)
+                ):
+                    # Don't bother inlining a direct alias since it can get optimized away later and
+                    # reordering can reduce optimality since we don't have many other code motion optimizations.
                     new_statements.append(stmt)
+                    continue
+                while True:
+                    inlinable_uses = self.get_inlinable_uses(stmt, set())
+                    subs = {}
+                    for p in inlinable_uses:
+                        if p in valid:
+                            subs[p] = inlined_definitions[p]
+
+                    if subs:
+                        stmt = self.substitute(stmt, subs)
+                    else:
+                        new_statements.append(stmt)
+                        break
 
             block.statements = new_statements[:-1]
             block.test = new_statements[-1]
@@ -123,14 +189,55 @@ class InlineVars(CompilerPass):
                 raise TypeError(f"Unexpected statement: {stmt}")
         return uses
 
-    def is_inlinable(self, stmt):
+    def is_inlinable(self, stmt, callback: str):
+        match stmt:
+            case IRConst():
+                return True
+            case IRInstr(op=op, args=args) | IRPureInstr(op=op, args=args):
+                return not op.side_effects and op.pure and all(self.is_inlinable(arg, callback) for arg in args)
+            case IRGet():
+                return isinstance(stmt.place, SSAPlace) or (
+                    isinstance(stmt.place, BlockPlace)
+                    and isinstance(stmt.place.block, BlockData)
+                    and callback not in stmt.place.block.writable
+                    and isinstance(stmt.place.index, int | SSAPlace)
+                )
+            case IRSet():
+                return False
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")
+
+    def is_free_to_inline(self, stmt: IRStmt, callback: str) -> bool:
+        match stmt:
+            case IRConst():
+                return True
+            case IRInstr() | IRPureInstr():
+                return self.is_runtime_constant(stmt, callback)
+            case IRGet():
+                return isinstance(stmt.place, SSAPlace) or (
+                    isinstance(stmt.place, BlockPlace)
+                    and isinstance(stmt.place.block, float | int)
+                    and isinstance(stmt.place.index, float | int)
+                )
+            case IRSet():
+                return False
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")
+
+    def is_runtime_constant(self, stmt: IRStmt, callback: str) -> bool:
         match stmt:
             case IRConst():
                 return True
             case IRInstr(op=op, args=args) | IRPureInstr(op=op, args=args):
-                return not op.side_effects and op.pure and all(self.is_inlinable(arg) for arg in args)
+                return not op.side_effects and op.pure and all(self.is_runtime_constant(arg, callback) for arg in args)
             case IRGet():
-                return isinstance(stmt.place, SSAPlace)
+                return (
+                    isinstance(stmt.place, BlockPlace)
+                    and isinstance(stmt.place.block, BlockData)
+                    and callback not in stmt.place.block.writable
+                    and stmt.place.block.name in RUNTIME_CONSTANT_BLOCKS
+                    and isinstance(stmt.place.index, int | SSAPlace)
+                )
             case IRSet():
                 return False
             case _:

sonolus/backend/optimize/liveness.py

@@ -2,7 +2,7 @@ from collections import deque
 
 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.passes import CompilerPass
+from sonolus.backend.optimize.passes import CompilerPass, OptimizerConfig
 from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
 
 type HasLiveness = SSAPlace | TempBlock
@@ -12,7 +12,7 @@ class LivenessAnalysis(CompilerPass):
     def destroys(self) -> set[CompilerPass]:
         return set()
 
-    def run(self, entry: BasicBlock) -> BasicBlock:
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
         self.preprocess(entry)
         self.process(entry)
         return entry

sonolus/backend/optimize/optimize.py

@@ -7,7 +7,13 @@ from sonolus.backend.optimize.dead_code import (
     UnreachableCodeElimination,
 )
 from sonolus.backend.optimize.inlining import InlineVars
-from sonolus.backend.optimize.simplify import CoalesceFlow, NormalizeSwitch, RewriteToSwitch
+from sonolus.backend.optimize.simplify import (
+    CoalesceFlow,
+    CoalesceSmallConditionalBlocks,
+    NormalizeSwitch,
+    RemoveRedundantArguments,
+    RewriteToSwitch,
+)
 from sonolus.backend.optimize.ssa import FromSSA, ToSSA
 
 MINIMAL_PASSES = (
@@ -27,6 +33,7 @@ STANDARD_PASSES = (
     CoalesceFlow(),
     UnreachableCodeElimination(),
     DeadCodeElimination(),
+    CoalesceSmallConditionalBlocks(),
     ToSSA(),
     SparseConditionalConstantPropagation(),
     UnreachableCodeElimination(),
@@ -34,7 +41,14 @@ STANDARD_PASSES = (
     CoalesceFlow(),
     InlineVars(),
     DeadCodeElimination(),
+    InlineVars(),
+    CoalesceFlow(),
+    SparseConditionalConstantPropagation(),
+    RemoveRedundantArguments(),
+    DeadCodeElimination(),
+    CoalesceFlow(),
     RewriteToSwitch(),
+    InlineVars(),
     FromSSA(),
     CoalesceFlow(),
     CopyCoalesce(),

sonolus/backend/optimize/passes.py

@@ -3,10 +3,18 @@ from __future__ import annotations
 from abc import ABC, abstractmethod
 from collections import deque
 from collections.abc import Sequence
+from dataclasses import dataclass
 
+from sonolus.backend.mode import Mode
 from sonolus.backend.optimize.flow import BasicBlock
 
 
+@dataclass
+class OptimizerConfig:
+    mode: Mode | None = None
+    callback: str | None = None
+
+
 class CompilerPass(ABC):
     def requires(self) -> set[CompilerPass]:
         return set()
@@ -32,11 +40,11 @@ class CompilerPass(ABC):
         return passes | self.applies()
 
     @abstractmethod
-    def run(self, entry: BasicBlock) -> BasicBlock:
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
         pass
 
 
-def run_passes(entry: BasicBlock, passes: Sequence[CompilerPass]) -> BasicBlock:
+def run_passes(entry: BasicBlock, passes: Sequence[CompilerPass], config: OptimizerConfig) -> BasicBlock:
     active_passes = set()
     queue = deque(passes)
     while queue:
@@ -48,6 +56,6 @@ def run_passes(entry: BasicBlock, passes: Sequence[CompilerPass]) -> BasicBlock:
             queue.appendleft(current_pass)
             queue.extendleft(missing_requirements)
             continue
-        entry = current_pass.run(entry)
+        entry = current_pass.run(entry, config)
         active_passes = current_pass.exists_after(active_passes)
     return entry

sonolus/backend/optimize/simplify.py

@@ -1,11 +1,12 @@
-from sonolus.backend.ir import IRConst, IRGet, IRPureInstr, IRSet
+from sonolus.backend.ir import IRConst, IRGet, IRInstr, 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
+from sonolus.backend.optimize.flow import BasicBlock, FlowEdge, traverse_cfg_preorder
+from sonolus.backend.optimize.passes import CompilerPass, OptimizerConfig
+from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
 
 
 class CoalesceFlow(CompilerPass):
-    def run(self, entry: BasicBlock) -> BasicBlock:
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
         queue = [entry]
         processed = set()
         while queue:
@@ -16,7 +17,7 @@ class CoalesceFlow(CompilerPass):
             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:
+                    if dst.phis or dst.statements or len(dst.outgoing) != 1 or dst is block or dst is entry:
                         break
                     next_dst = next(iter(dst.outgoing)).dst
                     if next_dst.phis:
@@ -42,6 +43,8 @@ class CoalesceFlow(CompilerPass):
                 queue.extend(edge.dst for edge in block.outgoing)
                 continue
             next_block = next(iter(block.outgoing)).dst
+            if next_block is block or next_block is entry:
+                continue
             if len(next_block.incoming) != 1:
                 queue.append(next_block)
                 if not block.statements and not block.phis and not next_block.phis:
@@ -55,7 +58,8 @@ class CoalesceFlow(CompilerPass):
                 continue
             for p, args in next_block.phis.items():
                 if block not in args:
-                    continue
+                    # This is the only predecessor to the block, so it must be a phi argument
+                    raise ValueError("Missing phi argument")
                 block.statements.append(IRSet(p, IRGet(args[block])))
             block.statements.extend(next_block.statements)
             block.test = next_block.test
@@ -73,13 +77,48 @@ class CoalesceFlow(CompilerPass):
         return entry
 
 
+class CoalesceSmallConditionalBlocks(CompilerPass):
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
+        queue = [entry]
+        processed = set()
+        while queue:
+            block = queue.pop()
+            if block.phis:
+                raise RuntimeError("SSA form is not supported in this pass")
+            if block in processed:
+                continue
+            processed.add(block)
+            while len(block.outgoing) == 1:
+                next_edge = next(iter(block.outgoing))
+                next_block = next_edge.dst
+                if len(next_block.statements) <= 1:
+                    next_block.incoming.remove(next_edge)
+                    block.test = next_block.test
+                    block.outgoing = {FlowEdge(src=block, dst=edge.dst, cond=edge.cond) for edge in next_block.outgoing}
+                    block.statements.extend(next_block.statements)
+                    for edge in block.outgoing:
+                        edge.dst.incoming.add(edge)
+                else:
+                    break
+            queue.extend(
+                edge.dst
+                for edge in sorted(block.outgoing, key=lambda e: (e.cond is not None, e.cond))
+                if edge.dst not in processed
+            )
+
+        reachable_blocks = set(traverse_cfg_preorder(entry))
+        for block in traverse_cfg_preorder(entry):
+            block.incoming = {edge for edge in block.incoming if edge.src in reachable_blocks}
+        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:
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
         self.ifs_to_switch(entry)
         self.combine_blocks(entry)
         self.remove_unreachable(entry)
@@ -159,7 +198,7 @@ class RewriteToSwitch(CompilerPass):
 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:
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
         for block in traverse_cfg_preorder(entry):
             cases = {edge.cond for edge in block.outgoing}
             if len(cases) <= 2:
@@ -189,3 +228,93 @@ class NormalizeSwitch(CompilerPass):
             if case != offset + i * stride:
                 return None, None
         return offset, stride
+
+
+class RemoveRedundantArguments(CompilerPass):
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
+        for block in traverse_cfg_preorder(entry):
+            block.statements = [self.update_statement(stmt) for stmt in block.statements]
+            block.test = self.update_statement(block.test)
+        return entry
+
+    def update_statement(self, stmt):
+        match stmt:
+            case IRPureInstr() | IRPureInstr():
+                op = stmt.op
+                args = stmt.args
+                match op:
+                    case Op.Add:
+                        args = [arg for arg in args if not (isinstance(arg, IRConst) and arg.value == 0)]
+                        if len(args) == 1:
+                            return args[0]
+                    case Op.Subtract:
+                        args = [
+                            args[0],
+                            *(arg for arg in args[1:] if not (isinstance(arg, IRConst) and arg.value == 0)),
+                        ]
+                        if len(args) == 1:
+                            return args[0]
+                    case Op.Multiply:
+                        args = [arg for arg in args if not (isinstance(arg, IRConst) and arg.value == 1)]
+                        if len(args) == 1:
+                            return args[0]
+                    case Op.Divide:
+                        args = [
+                            args[0],
+                            *(arg for arg in args[1:] if not (isinstance(arg, IRConst) and arg.value == 1)),
+                        ]
+                        if len(args) == 1:
+                            return args[0]
+                return type(stmt)(op=op, args=[self.update_statement(arg) for arg in args])
+            case IRSet(place=place, value=value):
+                return IRSet(place=place, value=self.update_statement(value))
+            case _:
+                return stmt
+
+
+class RenumberVars(CompilerPass):
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
+        numbers = {}
+        for block in traverse_cfg_preorder(entry):
+            for stmt in block.statements:
+                if (
+                    isinstance(stmt, IRSet)
+                    and isinstance(stmt.place, BlockPlace)
+                    and isinstance(stmt.place.block, TempBlock)
+                    and stmt.place.block.size == 1
+                    and stmt.place.block not in numbers
+                ):
+                    numbers[stmt.place.block] = len(numbers) + 1
+        for block in traverse_cfg_preorder(entry):
+            block.statements = [self.update_statement(stmt, numbers) for stmt in block.statements]
+            block.test = self.update_statement(block.test, numbers)
+        return entry
+
+    def update_statement(self, stmt, numbers: dict[TempBlock, int]):
+        match stmt:
+            case IRConst():
+                return stmt
+            case IRPureInstr(op=op, args=args):
+                return IRPureInstr(op=op, args=[self.update_statement(arg, numbers) for arg in args])
+            case IRInstr(op=op, args=args):
+                return IRInstr(op=op, args=[self.update_statement(arg, numbers) for arg in args])
+            case IRGet(place=place):
+                return IRGet(place=self.update_statement(place, numbers))
+            case IRSet(place=SSAPlace() as place, value=value):
+                return IRSet(place=place, value=self.update_statement(value, numbers))
+            case IRSet(place=place, value=value):
+                return IRSet(place=self.update_statement(place, numbers), value=self.update_statement(value, numbers))
+            case BlockPlace(block=block, index=index, offset=offset):
+                return BlockPlace(
+                    block=self.update_statement(block, numbers),
+                    index=self.update_statement(index, numbers),
+                    offset=offset,
+                )
+            case SSAPlace():
+                return stmt
+            case TempBlock() as b if b in numbers:
+                return TempBlock(f"v{numbers[b]}", size=1)
+            case int() | float() | TempBlock():
+                return stmt
+            case _:
+                raise TypeError(f"Unexpected statement: {stmt}")

sonolus/backend/optimize/ssa.py

@@ -1,7 +1,7 @@
 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.optimize.passes import CompilerPass, OptimizerConfig
 from sonolus.backend.place import BlockPlace, SSAPlace, TempBlock
 
 
@@ -9,7 +9,7 @@ class ToSSA(CompilerPass):
     def requires(self) -> set[CompilerPass]:
         return {DominanceFrontiers()}
 
-    def run(self, entry: BasicBlock) -> BasicBlock:
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
         defs = self.defs_to_blocks(entry)
         self.insert_phis(defs)
         self.rename(entry, defs, {var: [] for var in defs}, {})
@@ -139,15 +139,15 @@ class ToSSA(CompilerPass):
 
 
 class FromSSA(CompilerPass):
-    def run(self, entry: BasicBlock) -> BasicBlock:
+    def run(self, entry: BasicBlock, config: OptimizerConfig) -> BasicBlock:
         for block in [*traverse_cfg_preorder(entry)]:
             self.process_block(block)
         return entry
 
     def process_block(self, block: BasicBlock):
-        incoming = [*block.incoming]
+        orig_incoming = [*block.incoming]
         block.incoming.clear()
-        for edge in incoming:
+        for edge in orig_incoming:
             between_block = BasicBlock()
             edge.dst = between_block
             between_block.incoming.add(edge)
@@ -156,12 +156,46 @@ class FromSSA(CompilerPass):
             between_block.outgoing.add(next_edge)
             for args in block.phis.values():
                 if edge.src in args:
-                    args[between_block] = args.pop(edge.src)
+                    args[between_block] = args[edge.src]
+        for edge in orig_incoming:
+            # Multiple edges with different conditions can connect two of the same blocks,
+            # so we need to remove the old phi arguments in a pass at the end.
+            for args in block.phis.values():
+                if edge.src in args:
+                    del args[edge.src]
+        incoming_blocks = {edge.src for edge in block.incoming}
+        args_by_src = {}
+        for args in block.phis.values():
+            for src, arg in args.items():
+                if src not in args_by_src:
+                    args_by_src[src] = set()
+                args_by_src[src].add(arg)
         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)))
-                )
+                if src not in incoming_blocks:
+                    # Edges may have been rewritten so a phi refers to a block that is no longer directly connected.
+                    continue
+                if var in args_by_src[src]:
+                    # Make an extra copy first of values that may be overwritten by another assignment.
+                    src.statements.append(
+                        IRSet(
+                            place=self.place_from_ssa_place(var, "*"), value=IRGet(place=self.place_from_ssa_place(arg))
+                        )
+                    )
+        for var, args in block.phis.items():
+            for src, arg in args.items():
+                if src not in incoming_blocks:
+                    continue
+                if var in args_by_src[src]:
+                    src.statements.append(
+                        IRSet(
+                            place=self.place_from_ssa_place(var), value=IRGet(place=self.place_from_ssa_place(var, "*"))
+                        )
+                    )
+                else:
+                    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)
@@ -193,8 +227,8 @@ class FromSSA(CompilerPass):
             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 temp_block_from_ssa_place(self, ssa_place: SSAPlace, suffix: str = "") -> TempBlock:
+        return TempBlock(f"{ssa_place.name}.{ssa_place.num}{suffix}")
 
-    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)
+    def place_from_ssa_place(self, ssa_place: SSAPlace, suffix: str = "") -> BlockPlace:
+        return BlockPlace(block=self.temp_block_from_ssa_place(ssa_place, suffix), index=0, offset=0)

sonolus/backend/place.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
+
 from collections.abc import Iterator
-from typing import Self
 
 from sonolus.backend.blocks import Block
 
@@ -22,10 +23,10 @@ class TempBlock:
     def __str__(self):
         return f"{self.name}"
 
-    def __getitem__(self, item) -> "BlockPlace":
+    def __getitem__(self, item) -> BlockPlace:
         return BlockPlace(self, item)
 
-    def __iter__(self) -> "Iterator[BlockPlace]":
+    def __iter__(self) -> Iterator[BlockPlace]:
         for i in range(self.size):
             yield self[i]
 
@@ -78,7 +79,7 @@ class BlockPlace:
         else:
             return f"{self.block}[{self.index} + {self.offset}]"
 
-    def add_offset(self, offset: int) -> Self:
+    def add_offset(self, offset: int) -> BlockPlace:
         return BlockPlace(self.block, self.index, self.offset + offset)
 
     def __eq__(self, other):

sonolus/backend/utils.py

@@ -61,3 +61,27 @@ def scan_writes(node: ast.AST) -> set[str]:
     visitor = ScanWrites()
     visitor.visit(node)
     return set(visitor.writes)
+
+
+class HasDirectYield(ast.NodeVisitor):
+    def __init__(self):
+        self.started = False
+        self.has_yield = False
+
+    def visit_Yield(self, node: ast.Yield):
+        self.has_yield = True
+
+    def visit_YieldFrom(self, node: ast.YieldFrom):
+        self.has_yield = True
+
+    def visit_FunctionDef(self, node: ast.FunctionDef):
+        if self.started:
+            return
+        self.started = True
+        self.generic_visit(node)
+
+
+def has_yield(node: ast.AST) -> bool:
+    visitor = HasDirectYield()
+    visitor.visit(node)
+    return visitor.has_yield