Review changes and test

pytorch · emcastillo · Apr 28, 2021 · Jun 17, 2021 · Jun 21, 2021 · Jun 21, 2021
commit aa107ef2b7a9afb24c3938f1c1c6dfb3056cfaa7
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -7095,8 +7095,8 @@ def compare_strides(s1, s2, div):
                 _test_helper(x, op, unary=True)
 
     @skipMeta
-    @dtypes(*get_all_dtypes())
-    def test_dlpack_conversion(self, device, dtype):
+    @dtypes(*torch.testing.get_all_dtypes())
+    def test_dlpack_capsule_conversion(self, device, dtype):
         # DLpack does not explicitly support bool
         # It does it through uint8 type
         if dtype is torch.bool:
@@ -7105,6 +7105,46 @@ def test_dlpack_conversion(self, device, dtype):
         z = from_dlpack(to_dlpack(x))
         self.assertEqual(z, x)
 
+    @skipMeta
+    @dtypes(*torch.testing.get_all_dtypes())
+    def test_dlpack_protocol_conversion(self, device, dtype):
+        # DLpack does not explicitly support bool
+        # It does it through uint8 type
+        if dtype is torch.bool:
+            return
+        x = make_tensor((5,), device, dtype, low=-9, high=9)
+        z = from_dlpack(x)
+        self.assertEqual(z, x)
+
+    @skipMeta
+    @dtypes(*torch.testing.get_all_dtypes())
+    def test_dlpack_conversion_with_streams(self, device, dtype):
+        # DLpack does not explicitly support bool
+        # It does it through uint8 type
+        if dtype is torch.bool:
+            return
+        # Create a stream where the tensor will reside
+        if device == 'cuda':
+            x = make_tensor((5,), device, dtype, low=-9, high=9)
+            stream = torch.cuda.Stream()
+            with torch.cuda.stream(stream):
+                assert stream.query()
+                z = from_dlpack(to_dlpack(x))
+                assert not stream.query()
+            assert stream.query()
+            self.assertEqual(z, x)
+
+    @skipMeta
+    @dtypes(*torch.testing.get_all_dtypes())
+    def test_dlpack_tensor_invalid_stream(self, device, dtype):
+        # DLpack does not explicitly support bool
+        # It does it through uint8 type
+        if dtype is torch.bool:
+            return
+        with self.assertRaises(TypeError):
+            x = make_tensor((5,), device, dtype, low=-9, high=9)
+            x.__dlpack__(stream=object())
+
     @onlyCUDA
     @unittest.skipIf(PYTORCH_CUDA_MEMCHECK, "is_pinned uses failure to detect pointer property")
     def test_pin_memory_from_constructor(self, device):

diff --git a/torch/_tensor.py b/torch/_tensor.py
@@ -1,4 +1,5 @@
 from collections import OrderedDict
+import enum
 import functools
 from numbers import Number
 from typing import Any, Dict, Optional, Tuple, Union
@@ -1063,18 +1064,19 @@ def __dlpack__(self, stream=None):
 
         Args:
             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
+            to a stream (CUDA or ROCm). `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 stream is not None and type(stream) is not int:
-            # currently in pytorch is not possible to create a stream
-            # from a given pointer
+            # Stream pointers in CUDA/ROCm are uniquely numbered and can
+            # be retrieved from their integer value. 
             raise TypeError('stream must be ``int`` or ``none``')
         elif stream is not None and stream != -1:
-            if self.device.type in ('cuda', 'rocm'):
+            if self.device.type == 'cuda':
                 stream = torch.cuda.streams.ExternalStream(stream)
                 # Only synchronize on different streams
                 if stream != torch.cuda.current_stream:
@@ -1083,16 +1085,22 @@ def __dlpack__(self, stream=None):
                     torch.cuda.current_stream().wait_event(event)
         return torch.utils.dlpack.to_dlpack(self)
 
-    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}
+    def __dlpack_device__(self) -> Tuple[enum.IntEnum, int]:
+        class DLPackIds(enum.IntEnum): 
+            cpu = 1
+            cpu_pinned = 3
+            cuda = 2
+            opencl = 4
+            vulkan = 7
+            rocm = 10
+
         idx = self.device.index if self.device.index is not None else 0
-        # TODO(ecastill) detect HIP or CUDA
-        # in torch rocm device is cuda too
-        return (dlpack_ids[self.device.type], idx)
+        device_type = self.device.type
+        if device_type == 'cuda' and torch.version.hip is not None:
+            device_type = 'rocm'
+        elif device_type == 'cpu' and self.is_pinned():
+            device_type = 'cpu_pinned'
+        return (DLPackIds[device_type], idx)
 
     __module__ = 'torch'
 

diff --git a/torch/utils/dlpack.py b/torch/utils/dlpack.py
@@ -1,3 +1,5 @@
+from typing import Any
+
 import torch
 
 from torch._C import _from_dlpack
@@ -15,17 +17,18 @@
 """)
 
 
-def from_dlpack(ext_tensor) -> torch.Tensor:
+def from_dlpack(ext_tensor: Any) -> torch.Tensor:
     """from_dlpack(ext_tensor) -> Tensor
 
-    Decodes a DLPack to a tensor.
-
-    Args:
-        ext_tensor: a PyCapsule object with the dltensor
+    Convers a tensor from a external library into a ``torch.Tensor``
+    by means of the ``__dlpack__`` protocol.
 
     The tensor will share the memory with the object represented
     in the dlpack.
-    Note that each dlpack can only be consumed once.
+
+    In order to keep backward compatibility, this function also admits
+    to pass a dlpack capsule object.
+    Note that each dlpack capsule can only be consumed once.
 
     Args:
         ext_tensor (object with __dlpack__ attribute or dlpack capsule):
@@ -45,5 +48,5 @@ def from_dlpack(ext_tensor) -> torch.Tensor:
             dlpack = ext_tensor.__dlpack__()
     else:
         # Old versions just call the converter
-        dlpack = tensor
-    _from_dlpack(dlpack)
+        dlpack = ext_tensor
+    return _from_dlpack(dlpack)