[user-streams] Add record/wait ops (pytorch#167151)

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

Author: mlazos

test/dynamo/test_streams.py

@@ -441,7 +441,7 @@ def forward(self, tangents_1: "f32[2, 2]", tangents_2: "f32[2, 2]"):
         )
 
     @requires_cuda
-    def test_run_opcheck(self):
+    def test_run_opcheck_fork_join(self):
         from torch._dynamo.variables.streams import fork_stream, join_stream
         from torch.library import opcheck
 
@@ -462,6 +462,30 @@ def test_run_opcheck(self):
             torch.accelerator.set_stream(original_stream)
             reset_user_object_tracking()
 
+    @requires_cuda
+    def test_run_opcheck_wait_record(self):
+        from torch._dynamo.variables.streams import record_event, wait_event
+        from torch.library import opcheck
+
+        original_stream = torch.accelerator.current_stream()
+        try:
+            s0 = torch.Stream()
+            s1 = torch.Stream()
+            e0 = torch.Event()
+            e1 = torch.Event()
+            store_user_object_weakrefs(s0, s1, e0, e1)
+
+            sample_inputs = [
+                (2, 0),
+                (3, 1),
+            ]
+            for args in sample_inputs:
+                opcheck(wait_event, args)
+                opcheck(record_event, args)
+        finally:
+            torch.accelerator.set_stream(original_stream)
+            reset_user_object_tracking()
+
 
 if __name__ == "__main__":
     from torch._dynamo.test_case import run_tests

torch/_dynamo/variables/streams.py

@@ -28,16 +28,28 @@
 Tensor = torch.Tensor
 
 
+def _get_stream_by_index(index: int) -> torch.Stream:
+    stream = get_external_object_by_index(index)
+    assert isinstance(stream, torch.Stream), (
+        f"Fork/join stream expected a stream object at index {index}"
+    )
+    return stream
+
+
+def _get_event_by_index(index: int) -> torch.Event:
+    event = get_external_object_by_index(index)
+    assert isinstance(event, torch.Event), (
+        f"Record/wait event expected an event object at index {index}"
+    )
+    return event
+
+
 @custom_op("streams::fork", mutates_args=())
 def fork_stream(
     from_index: int,  # kept to make stream transitions clearer
     to_index: int,
 ) -> None:
-    stream = get_external_object_by_index(to_index)
-    assert isinstance(stream, torch.Stream), (
-        f"fork_stream expects a stream object at index {to_index}"
-    )
-    torch.accelerator.set_stream(stream)
+    torch.accelerator.set_stream(_get_stream_by_index(to_index))
 
 
 @fork_stream.register_fake
@@ -50,11 +62,7 @@ def _(
 
 @custom_op("streams::join", mutates_args=())
 def join_stream(from_index: int, to_index: int) -> None:
-    stream = get_external_object_by_index(to_index)
-    assert isinstance(stream, torch.Stream), (
-        f"join_stream expects a stream object at index {to_index}"
-    )
-    torch.accelerator.set_stream(stream)
+    torch.accelerator.set_stream(_get_stream_by_index(to_index))
 
 
 @join_stream.register_fake
@@ -65,6 +73,36 @@ def _(
     pass
 
 
+@custom_op("streams::record_event", mutates_args=())
+def record_event(event_index: int, stream_index: int) -> None:
+    event = _get_event_by_index(event_index)
+    stream = _get_stream_by_index(stream_index)
+    stream.record_event(event)
+
+
+@record_event.register_fake
+def _(
+    event_index: int,
+    stream_index: int,
+) -> None:
+    pass
+
+
+@custom_op("streams::wait_event", mutates_args=())
+def wait_event(event_index: int, stream_index: int) -> None:
+    event = _get_event_by_index(event_index)
+    stream = _get_stream_by_index(stream_index)
+    stream.wait_event(event)
+
+
+@wait_event.register_fake
+def _(
+    event_index: int,
+    stream_index: int,
+) -> None:
+    pass
+
+
 class SymbolicStreamState:
     """Track the currently entered stream if any"""