[mypyc] Borrow final attributes more aggressively (#21702)

Final instance attributes of native classes can't be rebound after
initialization, so we can borrow them more aggressively than regular
attributes, as long as the base object is kept alive for the duration of
the borrowing. This allows skipping some incref/decref operations.

Many workloads spend a significant fraction of CPU on incref/decref, and
these operations are quite a bit more expensive on free-threaded builds,
so the potential performance impact is significant especially on
free-threaded builds.

For example, the attribute `x` can only be safely borrowed if it's
Final, since otherwise `foo` could assign to the attribute and free the
old value:
```py
def func(o: C) -> None:
    foo(o.x)
```

There are some subtleties in the implementation. Here are the main
things that required extra care:
* We don't allow borrowed values to escape from conditionally executed
code paths (e.g. conditional expressions, comprehensions).
* If a local variable can be modified with an assignment expression, we
restrict borrowing based on that variable.
* We can only borrow for a longer duration if the attribute value
doesn't depend on a subexpression with a smaller borrow scope.
* Lambda expressions generate a complete separate expression and
borrowing scope.
 
I used coding agent assist, especially for tests, but created the
implementation is short, individually reviewed increments.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: JukkaL

.pre-commit-config.yaml

@@ -26,7 +26,7 @@ repos:
     hooks:
       - id: codespell
         args:
-          - --ignore-words-list=HAX,Nam,ccompiler,ot,statics,whet,zar
+          - --ignore-words-list=HAX,Nam,ccompiler,keep-alives,ot,statics,whet,zar
         exclude: ^(mypy/test/|mypy/typeshed/|mypyc/test-data/|test-data/).+$
   - repo: https://github.com/rhysd/actionlint
     rev: v1.7.7

mypyc/common.py

@@ -68,6 +68,14 @@
 BITMAP_TYPE: Final = "uint32_t"
 BITMAP_BITS: Final = 32
 
+# Constant for keeping a (often borrowed) op alive for a short time (a few subexpressions,
+# no arbitrary computation or memory allocations), until flush_keep_alives() is called.
+KEEP_ALIVE_SHORT_LIVED: Final = 0
+# Keep value alive longer, up to until the current top-level expression has been fully
+# evaluated. This allows keeping alive values across function calls and arbitrary
+# computation. Note that some expressions (e.g. lambdas), restrict the scope of borrowing.
+KEEP_ALIVE_WHOLE_EXPRESSION: Final = 1
+
 # Runtime C library files that are always included (some ops may bring
 # extra dependencies via mypyc.ir.deps.SourceDep or mypyc.ir.deps.HeaderDep)
 RUNTIME_C_FILES: Final = [

mypyc/ir/ops.py

@@ -31,7 +31,7 @@ class to enable the new behavior. Sometimes adding a new abstract
 
 from mypy_extensions import trait
 
-from mypyc.common import PROPSET_PREFIX
+from mypyc.common import KEEP_ALIVE_SHORT_LIVED, PROPSET_PREFIX
 from mypyc.ir.deps import Dependency
 from mypyc.ir.rtypes import (
     RArray,
@@ -898,6 +898,7 @@ def __init__(
         *,
         borrow: bool = False,
         allow_error_value: bool = False,
+        borrow_scope: int = KEEP_ALIVE_SHORT_LIVED,
     ) -> None:
         super().__init__(line)
         self.obj = obj
@@ -912,6 +913,8 @@ def __init__(
         elif attr_type.error_overlap:
             self.error_kind = ERR_MAGIC_OVERLAPPING
         self.is_borrowed = borrow and attr_type.is_refcounted
+        # How long a borrowed result of this op stays valid (a KEEP_ALIVE_* constant).
+        self.borrow_scope = borrow_scope
 
     def sources(self) -> list[Value]:
         return [self.obj]

mypyc/irbuild/builder.py

@@ -20,12 +20,17 @@
     TYPE_VAR_KIND,
     TYPE_VAR_TUPLE_KIND,
     ArgKind,
+    AssignmentExpr,
+    AwaitExpr,
     CallExpr,
     Decorator,
+    DictionaryComprehension,
     Expression,
     FuncDef,
+    GeneratorExpr,
     IndexExpr,
     IntExpr,
+    LambdaExpr,
     Lvalue,
     MemberExpr,
     MypyFile,
@@ -41,7 +46,10 @@
     TypeInfo,
     TypeParam,
     Var,
+    YieldExpr,
+    YieldFromExpr,
 )
+from mypy.traverser import TraverserVisitor
 from mypy.types import (
     AnyType,
     DeletedType,
@@ -63,6 +71,8 @@
     EXT_SUFFIX,
     GENERATOR_ATTRIBUTE_PREFIX,
     IS_FREE_THREADED,
+    KEEP_ALIVE_SHORT_LIVED,
+    KEEP_ALIVE_WHOLE_EXPRESSION,
     MODULE_PREFIX,
     SELF_NAME,
     TEMP_ATTR_NAME,
@@ -80,6 +90,7 @@
     BasicBlock,
     Branch,
     Call,
+    Cast,
     ComparisonOp,
     GetAttr,
     InitStatic,
@@ -284,6 +295,20 @@ def __init__(
         self.imports: dict[str, None] = {}
 
         self.can_borrow = False
+        self.expression_depth = 0
+        # Symbols (local vars) reassigned via a walrus expression within the current
+        # top-level expression. Used to avoid borrowing an attribute over the whole
+        # expression when the borrow root could be rebound (and thus freed) partway.
+        self.reassigned_in_expr: set[SymbolNode] = set()
+        # Whether the current top-level expression contains a suspension point
+        # (await, yield or yield from). A whole-expression borrow can't span such a
+        # point, since the borrowed value (and its root) live in registers that are
+        # not spilled into the generator environment across the suspend.
+        self.expr_has_suspend = False
+        # Saved expression state for enclosing functions (see enter()/leave()).
+        self.expression_depth_stack: list[int] = []
+        self.reassigned_in_expr_stack: list[set[SymbolNode]] = []
+        self.expr_has_suspend_stack: list[bool] = []
 
         # When set, load_globals_dict uses this module instead of self.module_name.
         # Used by generate_attr_defaults_init for cross-module inherited defaults.
@@ -315,6 +340,10 @@ def accept(self, node: Statement | Expression, *, can_borrow: bool = False) -> V
         """
         with self.catch_errors(node.line):
             if isinstance(node, Expression):
+                self.expression_depth += 1
+                if self.expression_depth == 1:
+                    self.reassigned_in_expr = find_walrus_targets(node)
+                    self.expr_has_suspend = expr_has_suspend(node)
                 old_can_borrow = self.can_borrow
                 self.can_borrow = can_borrow
                 try:
@@ -329,6 +358,11 @@ def accept(self, node: Statement | Expression, *, can_borrow: bool = False) -> V
                 self.can_borrow = old_can_borrow
                 if not can_borrow:
                     self.flush_keep_alives(node.line)
+                self.expression_depth -= 1
+                if self.expression_depth == 0:
+                    self.flush_keep_alives(node.line, scope=KEEP_ALIVE_WHOLE_EXPRESSION)
+                    self.reassigned_in_expr = set()
+                    self.expr_has_suspend = False
                 return res
             else:
                 try:
@@ -337,8 +371,8 @@ def accept(self, node: Statement | Expression, *, can_borrow: bool = False) -> V
                     pass
                 return None
 
-    def flush_keep_alives(self, line: int) -> None:
-        self.builder.flush_keep_alives(line)
+    def flush_keep_alives(self, line: int, *, scope: int = KEEP_ALIVE_SHORT_LIVED) -> None:
+        self.builder.flush_keep_alives(line, scope=scope)
 
     # Pass through methods for the most common low-level builder ops, for convenience.
 
@@ -1192,7 +1226,9 @@ def is_synthetic_type(self, typ: TypeInfo) -> bool:
         return typ.is_named_tuple or typ.is_newtype or typ.typeddict_type is not None
 
     def get_final_ref(self, expr: MemberExpr) -> tuple[str, Var, bool] | None:
-        """Check if `expr` is a final attribute.
+        """Check if `expr` is a final class or module attribute.
+
+        Return False for instance attributes.
 
         This needs to be done differently for class and module attributes to
         correctly determine fully qualified name. Return a tuple that consists of
@@ -1341,6 +1377,15 @@ def enter(self, fn_info: FuncInfo | str = "", *, ret_type: RType = none_rprimiti
         self.fn_info = fn_info
         self.fn_infos.append(self.fn_info)
         self.ret_types.append(ret_type)
+        # A function body is its own top-level expression context, even when the
+        # function (e.g. a lambda) is being generated in the middle of an outer
+        # expression. Save the outer expression state and start fresh.
+        self.expression_depth_stack.append(self.expression_depth)
+        self.reassigned_in_expr_stack.append(self.reassigned_in_expr)
+        self.expr_has_suspend_stack.append(self.expr_has_suspend)
+        self.expression_depth = 0
+        self.reassigned_in_expr = set()
+        self.expr_has_suspend = False
         if fn_info.is_generator:
             self.nonlocal_control.append(GeneratorNonlocalControl())
         else:
@@ -1354,6 +1399,9 @@ def leave(self) -> tuple[list[Register], list[RuntimeArg], list[BasicBlock], RTy
         ret_type = self.ret_types.pop()
         fn_info = self.fn_infos.pop()
         self.nonlocal_control.pop()
+        self.expression_depth = self.expression_depth_stack.pop()
+        self.reassigned_in_expr = self.reassigned_in_expr_stack.pop()
+        self.expr_has_suspend = self.expr_has_suspend_stack.pop()
         self.builder = self.builders[-1]
         self.fn_info = self.fn_infos[-1]
         return builder.args, runtime_args, builder.blocks, ret_type, fn_info
@@ -1389,6 +1437,27 @@ def enter_scope(self, fn_info: FuncInfo) -> Iterator[None]:
             self.builder = self.builders[-1]
             self.fn_info = self.fn_infos[-1]
 
+    @contextmanager
+    def enter_borrow_scope(self, line: int) -> Iterator[None]:
+        """Enter new borrow scope from which borrows can't leak to outer expressions.
+
+        This is a borrow region (see LowLevelIRBuilder.borrow_region) that also
+        resets the per-expression borrowing heuristic state, since the body forms
+        its own top-level expression context (e.g. a comprehension iteration or a
+        lambda body).
+        """
+        old_expression_depth = self.expression_depth
+        old_reassigned_in_expr = self.reassigned_in_expr
+        old_expr_has_suspend = self.expr_has_suspend
+        self.expression_depth = 0
+        try:
+            with self.builder.borrow_region(line):
+                yield
+        finally:
+            self.expression_depth = old_expression_depth
+            self.reassigned_in_expr = old_reassigned_in_expr
+            self.expr_has_suspend = old_expr_has_suspend
+
     @contextmanager
     def enter_method(
         self,
@@ -1581,6 +1650,32 @@ def is_final_native_attr_ref(self, expr: MemberExpr) -> bool:
                 return expr.name in ir.final_attributes
         return False
 
+    def root_is_reassigned(self, v: Value) -> bool:
+        """Is the root local variable a borrow chain 'v' reads from reassigned this expression?
+
+        A whole-expression borrow of an attribute keeps the borrow root alive only
+        via the register holding it. If that register belongs to a local variable
+        that is rebound (via a walrus assignment) during the same top-level
+        expression, the old value may be freed while the borrow is still live.
+        """
+        if not self.reassigned_in_expr:
+            return False
+        # Peel borrowed links back to the root value the chain reads from.
+        while True:
+            if isinstance(v, GetAttr) and v.is_borrowed:
+                v = v.obj
+            elif isinstance(v, Cast) and v.is_borrowed:
+                v = v.src
+            else:
+                break
+        if not isinstance(v, Register):
+            return False
+        for symbol in self.reassigned_in_expr:
+            target = self.symtables[-1].get(symbol)
+            if isinstance(target, AssignmentTargetRegister) and target.register is v:
+                return True
+        return False
+
     def mark_block_unreachable(self) -> None:
         """Mark statements in the innermost block being processed as unreachable.
 
@@ -1695,6 +1790,78 @@ def get_call_target_fullname(ref: RefExpr) -> str:
     return ref.fullname
 
 
+class WalrusTargetCollector(TraverserVisitor):
+    """Collect the symbols assigned to by walrus expressions in a subtree."""
+
+    def __init__(self) -> None:
+        self.targets: set[SymbolNode] = set()
+
+    def visit_assignment_expr(self, o: AssignmentExpr) -> None:
+        if o.target.node is not None:
+            self.targets.add(o.target.node)
+        super().visit_assignment_expr(o)
+
+    def visit_lambda_expr(self, o: LambdaExpr) -> None:
+        # A lambda body forms its own expression context, so don't descend into it.
+        pass
+
+
+def find_walrus_targets(expr: Expression) -> set[SymbolNode]:
+    """Return the symbols reassigned via a walrus expression within 'expr'.
+
+    Walrus (':=') is the only way to rebind a variable in the middle of evaluating
+    an expression, so this is the complete set of in-expression reassignments.
+    """
+    collector = WalrusTargetCollector()
+    expr.accept(collector)
+    return collector.targets
+
+
+class SuspendDetector(TraverserVisitor):
+    """Detect await/yield/yield from expressions in a subtree."""
+
+    def __init__(self) -> None:
+        self.found = False
+
+    def visit_await_expr(self, o: AwaitExpr) -> None:
+        self.found = True
+
+    def visit_yield_expr(self, o: YieldExpr) -> None:
+        self.found = True
+
+    def visit_yield_from_expr(self, o: YieldFromExpr) -> None:
+        self.found = True
+
+    def visit_generator_expr(self, o: GeneratorExpr) -> None:
+        # An 'async for' clause suspends via an implicit await on __anext__ that
+        # isn't represented as an AwaitExpr node in the AST (list/set comprehensions
+        # delegate to a GeneratorExpr, so they are covered here too).
+        if any(o.is_async):
+            self.found = True
+        super().visit_generator_expr(o)
+
+    def visit_dictionary_comprehension(self, o: DictionaryComprehension) -> None:
+        if any(o.is_async):
+            self.found = True
+        super().visit_dictionary_comprehension(o)
+
+    def visit_lambda_expr(self, o: LambdaExpr) -> None:
+        # A lambda body forms its own function (and suspension) context.
+        pass
+
+
+def expr_has_suspend(expr: Expression) -> bool:
+    """Does evaluating 'expr' involve a suspension point (await/yield/yield from)?
+
+    A whole-expression borrow can't safely span a suspension point, since the
+    borrowed value and its borrow root are held in registers that aren't spilled
+    into the generator environment across the suspend.
+    """
+    detector = SuspendDetector()
+    expr.accept(detector)
+    return detector.found
+
+
 def create_type_params(
     builder: IRBuilder, typing_mod: Value, type_args: list[TypeParam], line: int
 ) -> list[Value]: