Allow assignments of new refs to input as well as output variables.

python · markshannon · Oct 7, 2024 · Aug 7, 2024 · Aug 8, 2024 · Aug 8, 2024
commit 9f2f3bb79493666b76fb15c76974cd962aa401ab
@@ -4092,10 +4092,10 @@ dummy_func(
                 PyObject *self = ((PyMethodObject *)callable_o)->im_self;
                 args[0] = PyStackRef_FromPyObjectNew(self);
                 PyObject *method = ((PyMethodObject *)callable_o)->im_func;
-                args[-1] = PyStackRef_FromPyObjectNew(method);
-                PyStackRef_CLOSE(callable);
+                PyStackRef temp = callable;
+                callable = PyStackRef_FromPyObjectNew(method);
+                PyStackRef_CLOSE(temp);
                 callable_o = method;
-                callable = args[-1];
             }
             int positional_args = total_args - (int)PyTuple_GET_SIZE(kwnames_o);
             // Check if the call can be inlined or not

@@ -436,9 +436,18 @@ def analyze_deferred_refs(node: parser.InstDef) -> dict[lexer.Token, str | None]
             )
 
         name = lhs[0].text
-        if not any(var.name == name for var in node.outputs):
+        match = False
+        for var in node.inputs:
+            if var.name == name:
+                match = True
+                break
+        for var in node.outputs:
+            if var.name == name:
+                match = True
+                break
+        if not match:
             raise analysis_error(
-                f"PyStackRef_FromPyObjectNew() must be assigned to an output, not '{name}'",
+                f"PyStackRef_FromPyObjectNew() must be assigned to an input or output, not '{name}'",
                 tkn,
             )
 

@@ -192,7 +192,7 @@ def write_uop(uop: Uop, emitter: Emitter, stack: Stack) -> None:
             stack.push(peeks.pop())
         emitter.emit(stack.define_output_arrays(uop.stack.outputs))
         code_list, storage = Storage.for_uop(stack, uop, locals)
-        of code in code_list:
+        for code in code_list:
             emitter.emit(code)
         for cache in uop.caches:
             if cache.name != "unused":