Use ExternalStream

pytorch · emcastillo · Apr 28, 2021 · Jun 17, 2021 · Jun 21, 2021 · Jun 21, 2021
commit 56af6f329ada96307528a0ccf12c6766d796edf9
diff --git a/torch/_tensor.py b/torch/_tensor.py
@@ -1062,19 +1062,32 @@ def __dlpack__(self, stream=None):
         of the consumer library.
 
         Args:
-            stream (integer or None): Object that represents a stream and provides
-                a `synchronize` method. Optional.
+            stream (integer or None): A Python integer representing a pointer
+            to a stream. `stream` is provided by the consumer to the producer
+            to instruct the producer to ensure that operations can safely be
+            performed on the array. The pointer must be a positive integer or
+            -1 . If stream is -1 , the value may be used by the consumer to
+            signal "producer must not perform any synchronization. Optional.
         """
-        if isinstance(stream, torch.cuda.Stream) or hasattr(stream, 'synchronize'):
-            stream.synchronize()
-        elif stream is not None and type(stream) is int:
+        if stream is not None and type(stream) is not int:
             # currently in pytorch is not possible to create a stream
             # from a given pointer
-            raise TypeError('Can\'t create a stream from an integer in PyTorch')
-
+            raise TypeError('stream must be ``int`` or ``none``')
+        elif stream is not None and stream != -1:
+            if self.device.type in ('cuda', 'rocm'):
+                stream = torch.cuda.streams.ExternalStream(stream)
+                # Only synchronize on different streams
+                if stream != torch.cuda.current_stream:
+                    event = torch.cuda.Event()
+                    event.record(stream)
+                    torch.cuda.current_stream().wait_event(event)
-                    event.record(stream)
-                    torch.cuda.current_stream().wait_event(event)
+                    event.record(torch.cuda.current_stream())
+                    stream.wait_event(event)
-                    event.record(stream)
-                    torch.cuda.current_stream().wait_event(event)
+                    event.record(torch.cuda.current_stream())
+                    stream.wait_event(event)
         return torch.utils.dlpack.to_dlpack(self)
 
-    def __dlpack_device__(self) -> Tuple[Int, Int]:
+    def __dlpack_device__(self) -> Tuple[int, int]:
+        # TODO(ecastill)
+        # Add support for the following devices
+        # CPU = 1 CPU_PINNED = 3 OPENCL = 4 VULKAN = 7
+        # METAL = 8 VPI = 9
         dlpack_ids = {'cpu': 1, 'cuda': 2, 'rocm': 10}
         idx = self.device.index if self.device.index is not None else 0
         # TODO(ecastill) detect HIP or CUDA

diff --git a/torch/utils/dlpack.py b/torch/utils/dlpack.py
@@ -16,12 +16,12 @@
 
 
 def from_dlpack(ext_tensor) -> torch.Tensor:
-    """from_dlpack(dlpack) -> Tensor
+    """from_dlpack(ext_tensor) -> Tensor
 
     Decodes a DLPack to a tensor.
 
     Args:
-        dlpack: a PyCapsule object with the dltensor
+        ext_tensor: a PyCapsule object with the dltensor
 
     The tensor will share the memory with the object represented
     in the dlpack.
@@ -32,14 +32,17 @@ def from_dlpack(ext_tensor) -> torch.Tensor:
             The tensor from an external library that will be converted
             to a PyTorch one.
     """
-    if hasattr(dlpack, '__dlpack__'):
-        device = dlpack.__dlpack_device__()
+    if hasattr(ext_tensor, '__dlpack__'):
+        device = ext_tensor.__dlpack_device__()
         # device is either CUDA or ROCm, we need to pass the current
         # stream
         if device[0] in (2, 10):
-            stream = torch.cuda.stream.current_stream('cuda:{}'.format(device[1]))
-            # Should we pass an id? or the producer can accept a stream object
+            stream = torch.cuda.current_stream('cuda:{}'.format(device[1]))
+            # cuda_stream is the pointer to the stream and it is a public
+            # attribute, but it is not documented
             dlpack = ext_tensor.__dlpack__(stream=stream.cuda_stream)
+        else:
+            dlpack = ext_tensor.__dlpack__()
     else:
         # Old versions just call the converter
         dlpack = tensor