[user-streams] Mark stream ops as side effectful (pytorch#167152)

Pull Request resolved: pytorch#167152
Approved by: https://github.com/Lucaskabela
ghstack dependencies: pytorch#167141, pytorch#167151

Author: mlazos

test/dynamo/test_streams.py

@@ -486,6 +486,22 @@ def test_run_opcheck_wait_record(self):
             torch.accelerator.set_stream(original_stream)
             reset_user_object_tracking()
 
+    def test_is_marked_side_effectful(self):
+        self.assertIn(
+            torch.ops.streams.fork.default, torch.fx.node._side_effectful_functions
+        )
+        self.assertIn(
+            torch.ops.streams.join.default, torch.fx.node._side_effectful_functions
+        )
+        self.assertIn(
+            torch.ops.streams.wait_event.default,
+            torch.fx.node._side_effectful_functions,
+        )
+        self.assertIn(
+            torch.ops.streams.record_event.default,
+            torch.fx.node._side_effectful_functions,
+        )
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests

torch/_dynamo/variables/streams.py

@@ -5,7 +5,7 @@
 import torch
 from torch._dynamo.variables.dicts import ConstDictVariable
 from torch._dynamo.variables.lists import TupleVariable
-from torch.fx import Proxy
+from torch.fx import has_side_effect, Proxy
 
 from .. import graph_break_hints
 from ..bytecode_transformation import create_call_function
@@ -60,6 +60,9 @@ def _(
     pass
 
 
+has_side_effect(torch.ops.streams.fork.default)
+
+
 @custom_op("streams::join", mutates_args=())
 def join_stream(from_index: int, to_index: int) -> None:
     torch.accelerator.set_stream(_get_stream_by_index(to_index))
@@ -73,6 +76,9 @@ def _(
     pass
 
 
+has_side_effect(torch.ops.streams.join.default)
+
+
 @custom_op("streams::record_event", mutates_args=())
 def record_event(event_index: int, stream_index: int) -> None:
     event = _get_event_by_index(event_index)
@@ -88,6 +94,9 @@ def _(
     pass
 
 
+has_side_effect(torch.ops.streams.record_event.default)
+
+
 @custom_op("streams::wait_event", mutates_args=())
 def wait_event(event_index: int, stream_index: int) -> None:
     event = _get_event_by_index(event_index)
@@ -103,6 +112,9 @@ def _(
     pass
 
 
+has_side_effect(torch.ops.streams.wait_event.default)
+
+
 class SymbolicStreamState:
     """Track the currently entered stream if any"""