[user-streams] Support streams as contexts (pytorch#164507)

Pull Request resolved: pytorch#164507
Approved by: https://github.com/williamwen42
ghstack dependencies: pytorch#162903, pytorch#164343, pytorch#164344

Author: mlazos

test/cpp_extensions/open_registration_extension/torch_openreg/tests/test_streams.py

@@ -5,6 +5,7 @@
 
 
 class TestStream(TestCase):
+    @skipIfTorchDynamo()
     def test_stream_create(self):
         stream = torch.Stream(device="openreg")
         self.assertEqual(stream.device_index, torch.openreg.current_device())
@@ -24,6 +25,7 @@ def test_stream_create(self):
         )
         self.assertEqual(stream, stream1)
 
+    @skipIfTorchDynamo()
     def test_stream_context(self):
         with torch.Stream(device="openreg:1") as stream:
             self.assertEqual(torch.accelerator.current_stream(), stream)
@@ -40,6 +42,7 @@ def test_stream_switch(self):
         current_stream = torch.accelerator.current_stream()
         self.assertEqual(current_stream, stream2)
 
+    @skipIfTorchDynamo()
     def test_stream_synchronize(self):
         stream = torch.Stream(device="openreg:1")
         self.assertEqual(True, stream.query())
@@ -49,12 +52,14 @@ def test_stream_synchronize(self):
         stream.synchronize()
         self.assertEqual(True, stream.query())
 
+    @skipIfTorchDynamo()
     def test_stream_repr(self):
         stream = torch.Stream(device="openreg:1")
         self.assertTrue(
             "torch.Stream device_type=openreg, device_index=1" in repr(stream)
         )
 
+    @skipIfTorchDynamo()
     def test_stream_wait_stream(self):
         stream_1 = torch.Stream(device="openreg:0")
         stream_2 = torch.Stream(device="openreg:1")

torch/_dynamo/variables/streams.py

@@ -92,7 +92,6 @@ def __init__(
             target_values=target_values, initial_values=initial_values, **kwargs
         )
         self.device = device
-        self.set_stream_id = get_interface_for_device(self.device)._set_stream_by_id
 
     def enter(self, tx: "InstructionTranslator") -> "VariableTracker":
         # to stream, from stream is the order of the arguments
@@ -124,7 +123,7 @@ def _initial_stream_proxies(self) -> tuple[Proxy, Proxy]:
 
     def _target_stream_proxies(self) -> tuple[Proxy, Proxy]:
         return StreamContextVariable._extract_stream_properties(
-            self.target_values[0].as_proxy()
+            self._get_target_values()[0].as_proxy()
         )
 
     @staticmethod
@@ -152,6 +151,15 @@ def _get_current_stream(
         )
         return current_stream
 
+    def _get_target_values(self) -> list["StreamVariable"]:
+        # We need this to be overridable, since StreamVariable does
+        # not store target values (it does not require any arguments)
+        # and captures the current stream at the time of entering the context
+        return self.target_values
+
+    def supports_graph_breaks(self) -> bool:
+        return True
+
 
 class StreamVariable(StreamContextVariable):
     """Represents the device-agnostic torch.Stream class"""
@@ -168,9 +176,7 @@ def __init__(
         assert value.device.type == device.type, (
             "stream value is not equal to the passed device"
         )
-        super().__init__(
-            target_values=[self], initial_values=None, device=device, **kwargs
-        )
+        super().__init__(target_values=[], initial_values=None, device=device, **kwargs)
         self.proxy = proxy
         self.value = value
         # pyrefly: ignore [read-only]
@@ -233,18 +239,23 @@ def call_method(
         return super().call_method(tx, name, args, kwargs)
 
     def enter(self, tx: "InstructionTranslator") -> "VariableTracker":
-        # NB: Set initial values and target values when we enter
+        # NB: Set initial values when we enter
         # Don't do this at object creation, as we need to record the current stream
         # at the time the context is entered.
         self.initial_values = [
             StreamContextVariable._get_current_stream(self.device, tx)
         ]
-        self.target_values = [self]
         return super().enter(tx)
 
     def as_proxy(self) -> Proxy:
         return self.proxy
 
+    def module_name(self) -> str:
+        return "torch._C"
+
+    def fn_name(self) -> str:
+        return "Stream"
+
     def reconstruct(self, codegen: "PyCodegen") -> None:
         # If we got here, this stream is fully subsumed by the graph - this means it is
         # not an input or global
@@ -259,6 +270,9 @@ def reconstruct(self, codegen: "PyCodegen") -> None:
         name = codegen.tx.output.install_global_by_id(prefix, self.value)
         codegen.append_output(codegen.create_load_global(name, add=True))
 
+    def _get_target_values(self) -> list["StreamVariable"]:
+        return [self]
+
 
 class EventVariable(VariableTracker):
     def __init__(self, proxy: Proxy, value: torch.Event, **kwargs: Any) -> None: