[user-streams] Stash graph created objects in keep_alive list for backwards (pytorch#167705)

Pull Request resolved: pytorch#167705
Approved by: https://github.com/williamwen42

Author: mlazos

torch/_dynamo/graph_bytecode_inputs.py

@@ -25,6 +25,11 @@ def has_user_objects() -> bool:
     return bool(index_to_bytecode_constructor)
 
 
+def stash_graph_created_object(obj: Any) -> Any:
+    keep_alive.append(obj)
+    return obj
+
+
 def get_external_object_by_index(index: int) -> Any:
     assert index in index_to_external_object_weakref, (
         "Index not registered in index_to_user_object_weakref"

torch/_dynamo/output_graph.py

@@ -2324,12 +2324,14 @@ def specialized_dispatch(*args: Any, **kwargs: Any) -> Any:
                         "store_user_object_weakrefs",
                     )
                 )
+
                 tmp_vars = []
                 for constructor in index_to_bytecode_constructor.values():
                     constructor(cg)
                     var_name = (
                         self.new_var()
-                    )  # keep alive any temp objects for the rest of the frame
+                    )  # keep alive any user objects for the rest of the frame
+                    # TODO: we could omit this for objects we create but shouldn't be too much overhead for now
                     cg.store(var_name)
                     tmp_vars.append(var_name)
 

torch/_dynamo/variables/streams.py

@@ -52,10 +52,21 @@ def new_stream(*args: tuple[Any], **kwargs: Any) -> int:
     )
 
 
+def _codegen_current_stream(device: torch.device, cg: "PyCodegen") -> None:
+    cg.add_push_null(
+        lambda: cg.load_import_from(
+            torch._dynamo.graph_bytecode_inputs.__name__,  # type: ignore[implicit-imports]
+            "stash_graph_created_object",
+        )
+    )
+    cg(CurrentStreamSource(device))
+    cg.extend_output(create_call_function(1, False))
+
+
 def get_current_stream(device: torch.device) -> int:
-    stream = torch.accelerator.current_stream()
+    stream = torch.accelerator.current_stream(device)
     return register_graph_created_object(
-        stream, lambda _, cg: cg(CurrentStreamSource(device))
+        stream, lambda _, cg: _codegen_current_stream(device, cg)
     )
 
 
@@ -362,6 +373,12 @@ def make_construct_in_graph_stream_fn(
         args: TupleVariable, kwargs: ConstDictVariable
     ) -> Callable[[int, "PyCodegen"], None]:
         def fn(index: int, codegen: "PyCodegen") -> None:
+            codegen.add_push_null(
+                lambda: codegen.load_import_from(
+                    torch._dynamo.graph_bytecode_inputs.__name__,  # type: ignore[implicit-imports]
+                    "stash_graph_created_object",
+                )
+            )
             codegen.add_push_null(
                 lambda: codegen.load_import_from(
                     torch._dynamo.utils.__name__, "build_stream"
@@ -370,6 +387,7 @@ def fn(index: int, codegen: "PyCodegen") -> None:
             codegen(args)
             codegen(kwargs)
             codegen.extend_output(create_call_function(2, False))
+            codegen.extend_output(create_call_function(1, False))
 
         return fn
 
@@ -473,6 +491,12 @@ def make_construct_in_graph_event_fn(
         args: TupleVariable, kwargs: ConstDictVariable
     ) -> Callable[[int, "PyCodegen"], None]:
         def fn(index: int, codegen: "PyCodegen") -> None:
+            codegen.add_push_null(
+                lambda: codegen.load_import_from(
+                    torch._dynamo.graph_bytecode_inputs.__name__,  # type: ignore[implicit-imports]
+                    "stash_graph_created_object",
+                )
+            )
             codegen.add_push_null(
                 lambda: codegen.load_import_from(
                     torch._dynamo.utils.__name__, "build_event"
@@ -481,6 +505,7 @@ def fn(index: int, codegen: "PyCodegen") -> None:
             codegen(args)
             codegen(kwargs)
             codegen.extend_output(create_call_function(2, False))
+            codegen.extend_output(create_call_function(1, False))
 
         return fn