Automatically flush to stack new

python · Fidget-Spinner · Jun 28, 2024 · Jun 28, 2024 · Jun 29, 2024 · Jun 29, 2024
commit f4b40223bbaa764da44e9e326bf5d9e4475d997f
diff --git a/Include/internal/pycore_frame.h b/Include/internal/pycore_frame.h
@@ -131,7 +131,7 @@ static inline void _PyFrame_Copy(_PyInterpreterFrame *src, _PyInterpreterFrame *
 
 #ifdef Py_GIL_DISABLED
     PyCodeObject *co = (PyCodeObject *)dest->f_executable;
-    for (int i = src->stacktop; i < co->co_nlocalsplus + co->co_stacksize; i++) {
+    for (int i = stacktop; i < co->co_nlocalsplus + co->co_stacksize; i++) {
         dest->localsplus[i] = PyStackRef_NULL;
     }
 #endif

diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
diff --git a/Python/gc.c b/Python/gc.c
@@ -537,12 +537,10 @@ visit_decref(PyObject *op, void *parent)
 int
 _PyGC_VisitFrameStack(_PyInterpreterFrame *frame, visitproc visit, void *arg)
 {
-    /* locals */
-    _PyStackRef *locals = _PyFrame_GetLocalsArray(frame);
-    int i = 0;
+    _PyStackRef *i = frame->localsplus;
     /* locals and stack */
-    for (; i <frame->stacktop; i++) {
-        Py_VISIT(PyStackRef_AsPyObjectBorrow(locals[i]));
+    for (; i < frame->stackpointer; i++) {
+        Py_VISIT(PyStackRef_AsPyObjectBorrow(*i));
     }
     return 0;
 }

diff --git a/Python/gc_free_threading.c b/Python/gc_free_threading.c
@@ -703,8 +703,8 @@ clear_weakrefs(struct collection_state *state)
             // freed after its function object.
             PyGenObject *gen = (PyGenObject *)op;
             struct _PyInterpreterFrame *frame = &(gen->gi_iframe);
-            for (int i = 0; i < frame->stacktop; i++) {
-                gc_visit_stackref(frame->localsplus[i]);
+            for (_PyStackRef *i = frame->localsplus; i < frame->stackpointer; i++) {
+                gc_visit_stackref(*i);
             }
         }
         if (PyWeakref_Check(op)) {
@@ -939,16 +939,14 @@ visit_decref_unreachable(PyObject *op, void *data)
 int
 _PyGC_VisitFrameStack(_PyInterpreterFrame *frame, visitproc visit, void *arg)
 {
-    /* locals */
-    _PyStackRef *locals = _PyFrame_GetLocalsArray(frame);
-    int i = 0;
+    _PyStackRef *i = frame->localsplus;
     /* locals and stack */
-    for (; i <frame->stacktop; i++) {
-        if (PyStackRef_IsDeferred(locals[i]) &&
+    for (; i < frame->stackpointer; i++) {
+        if (PyStackRef_IsDeferred(*i) &&
             (visit == visit_decref || visit == visit_decref_unreachable)) {
             continue;
         }
-        Py_VISIT(PyStackRef_AsPyObjectBorrow(locals[i]));
+        Py_VISIT(PyStackRef_AsPyObjectBorrow(*i));
     }
     return 0;
 }

diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
@@ -156,6 +156,7 @@ class Uop:
     stack: StackEffect
     caches: list[CacheEntry]
     body: list[lexer.Token]
+    token_requires_flush: list[bool]  # bitmap to body
     properties: Properties
     _size: int = -1
     implicitly_created: bool = False
@@ -345,6 +346,24 @@ def analyze_caches(inputs: list[parser.InputEffect]) -> list[CacheEntry]:
     return [CacheEntry(i.name, int(i.size)) for i in caches]
 
 
+def analyze_specials_requiring_flush(name: str, op: list[lexer.Token], outputs: list[StackItem]) -> list[bool]:
+    res = []
+    for idx, token in enumerate(op):
+        if token.kind == "IDENTIFIER" and token.text == "PyStackRef_FromPyObjectNew":
+            should_flush = True
+            assert op[idx-1].kind == "EQUALS", (f"{name} Result of a specials that is flushed"
+                                                " must be written to a stack variable directly")
+            # Scan to look for the first thing it's assigned to
+            for assgn_target in reversed(op[:idx-1]):
+                out_names = [out.name for out in outputs]
+                if assgn_target.kind == "IDENTIFIER":
+                    should_flush = assgn_target.text in out_names
+                    break
+            res.append(should_flush)
+        else:
+            res.append(False)
+    return res
+
 def variable_used(node: parser.InstDef, name: str) -> bool:
     """Determine whether a variable with a given name is used in a node."""
     return any(
@@ -619,13 +638,15 @@ def compute_properties(op: parser.InstDef) -> Properties:
 
 
 def make_uop(name: str, op: parser.InstDef, inputs: list[parser.InputEffect], uops: dict[str, Uop]) -> Uop:
+    stack = analyze_stack(op)
     result = Uop(
         name=name,
         context=op.context,
         annotations=op.annotations,
-        stack=analyze_stack(op),
+        stack=stack,
         caches=analyze_caches(inputs),
         body=op.block.tokens,
+        token_requires_flush=analyze_specials_requiring_flush(name, op.block.tokens, stack.outputs),
         properties=compute_properties(op),
     )
     if effect_depends_on_oparg_1(op) and "split" in op.annotations:
@@ -636,13 +657,15 @@ def make_uop(name: str, op: parser.InstDef, inputs: list[parser.InputEffect], uo
             if properties.oparg:
                 # May not need oparg anymore
                 properties.oparg = any(token.text == "oparg" for token in op.block.tokens)
+            stack = analyze_stack(op, bit)
             rep = Uop(
                 name=name_x,
                 context=op.context,
                 annotations=op.annotations,
-                stack=analyze_stack(op, bit),
+                stack=stack,
                 caches=analyze_caches(inputs),
                 body=op.block.tokens,
+                token_requires_flush=analyze_specials_requiring_flush(name, op.block.tokens, stack.outputs),
                 properties=properties,
             )
             rep.replicates = result
@@ -658,13 +681,15 @@ def make_uop(name: str, op: parser.InstDef, inputs: list[parser.InputEffect], uo
         properties = compute_properties(op)
         properties.oparg = False
         properties.const_oparg = oparg
+        stack = analyze_stack(op)
         rep = Uop(
             name=name_x,
             context=op.context,
             annotations=op.annotations,
-            stack=analyze_stack(op),
+            stack=stack,
             caches=analyze_caches(inputs),
             body=op.block.tokens,
+            token_requires_flush=analyze_specials_requiring_flush(name, op.block.tokens, stack.outputs),
             properties=properties,
         )
         rep.replicates = result