Type conversions now use auto gpu

apaszke · apaszke · commit da5bb373e6ad · 2016-09-15T18:48:27.000-07:00
diff --git a/test/test_cuda.py b/test/test_cuda.py
@@ -280,6 +280,27 @@ def test_serialization(self):
         q_copy[1].fill_(10)
         self.assertTrue(q_copy[3], torch.cuda.IntStorage(10).fill_(10))
 
+    def test_type_conversions(self):
+        x = torch.randn(5, 5)
+        self.assertIs(type(x.float()), torch.FloatTensor)
+        self.assertIs(type(x.cuda()), torch.cuda.DoubleTensor)
+        self.assertIs(type(x.cuda().float()), torch.cuda.FloatTensor)
+        self.assertIs(type(x.cuda().float().cpu()), torch.FloatTensor)
+        self.assertIs(type(x.cuda().float().cpu().int()), torch.IntTensor)
+
+        y = x.storage()
+        self.assertIs(type(y.float()), torch.FloatStorage)
+        self.assertIs(type(y.cuda()), torch.cuda.DoubleStorage)
+        self.assertIs(type(y.cuda().float()), torch.cuda.FloatStorage)
+        self.assertIs(type(y.cuda().float().cpu()), torch.FloatStorage)
+        self.assertIs(type(y.cuda().float().cpu().int()), torch.IntStorage)
+
+    @unittest.skipIf(torch.cuda.deviceCount() < 2, "only one GPU detected")
+    def test_type_conversions_same_gpu(self):
+        x = torch.randn(5, 5).cuda(1)
+        self.assertEqual(x.int().getDevice(), 1)
+
+
 for decl in tests:
     for t in types:
         tensor = t()
diff --git a/torch/Storage.py b/torch/Storage.py
@@ -1,4 +1,5 @@
 import torch
+from ._utils import _type
 
 
 class _StorageBase(object):
@@ -35,3 +36,27 @@ def clone(self):
 
     def tolist(self):
         return [v for v in self]
+
+    def double(self):
+        return self.type(type(self).__module__ + '.DoubleStorage')
+
+    def float(self):
+        return self.type(type(self).__module__ + '.FloatStorage')
+
+    def long(self):
+        return self.type(type(self).__module__ + '.LongStorage')
+
+    def int(self):
+        return self.type(type(self).__module__ + '.IntStorage')
+
+    def short(self):
+        return self.type(type(self).__module__ + '.ShortStorage')
+
+    def char(self):
+        return self.type(type(self).__module__ + '.CharStorage')
+
+    def byte(self):
+        return self.type(type(self).__module__ + '.ByteStorage')
+
+
+_StorageBase.type = _type
diff --git a/torch/Tensor.py b/torch/Tensor.py
@@ -1,5 +1,6 @@
 import torch
 from . import TensorPrinting
+from ._utils import _type
 from functools import reduce
 from itertools import chain
 import sys
@@ -25,43 +26,29 @@ class _TensorBase(object):
     def new(self, *args, **kwargs):
         return self.__class__(*args, **kwargs)
 
-    def type(self, t=None):
-        if isinstance(t, str) or t is None:
-            current = self.__module__ + '.' + self.__class__.__name__
-            if t is None:
-                return current
-            if t == current:
-                return self
-            _, _, typename = t.partition('.')
-            return torch._import_dotted_name(t)(self.size()).copy_(self)
-        else:
-            if t == type(self):
-                return self
-            return t(self.size()).copy_(self)
-
     def typeAs(self, t):
         return self.type(t.type())
 
     def double(self):
-        return self.type(torch.DoubleTensor)
+        return self.type(type(self).__module__ + '.DoubleTensor')
 
     def float(self):
-        return self.type(torch.FloatTensor)
+        return self.type(type(self).__module__ + '.FloatTensor')
 
     def long(self):
-        return self.type(torch.LongTensor)
+        return self.type(type(self).__module__ + '.LongTensor')
 
     def int(self):
-        return self.type(torch.IntTensor)
+        return self.type(type(self).__module__ + '.IntTensor')
 
     def short(self):
-        return self.type(torch.ShortTensor)
+        return self.type(type(self).__module__ + '.ShortTensor')
 
     def char(self):
-        return self.type(torch.CharTensor)
+        return self.type(type(self).__module__ + '.CharTensor')
 
     def byte(self):
-        return self.type(torch.ByteTensor)
+        return self.type(type(self).__module__ + '.ByteTensor')
 
     def copy_(self, other):
         torch._C._tensorCopy(self, other)
@@ -267,3 +254,6 @@ def __mod__(self, other):
 
     def __neg__(self):
         return self.neg()
+
+
+_TensorBase.type = _type
diff --git a/torch/_utils.py b/torch/_utils.py
@@ -0,0 +1,17 @@
+import torch
+
+
+def _type(self, t=None):
+    if isinstance(t, str) or t is None:
+        current = self.__module__ + '.' + self.__class__.__name__
+        if t is None:
+            return current
+        if t == current:
+            return self
+        _, _, typename = t.partition('.')
+        return torch._import_dotted_name(t)(self.size()).copy_(self)
+    else:
+        if t == type(self):
+            return self
+        return t(self.size()).copy_(self)
+
diff --git a/torch/csrc/generic/StorageMethods.cpp b/torch/csrc/generic/StorageMethods.cpp
@@ -405,6 +405,15 @@ PyObject * THPStorage_(_sharedFd)(THPStorage *self)
 }
 #endif
 
+#ifdef THC_GENERIC_FILE
+PyObject * THPStorage_(getDevice)(THPStorage *self)
+{
+  HANDLE_TH_ERRORS
+  return PyLong_FromLong(THCStorage_(getDevice)(LIBRARY_STATE self->cdata));
+  END_HANDLE_TH_ERRORS
+}
+#endif
+
 PyObject * THPStorage_(_setCdata)(THPStorage *self, PyObject *new_cdata)
 {
   HANDLE_TH_ERRORS
@@ -438,6 +447,9 @@ static PyMethodDef THPStorage_(methods)[] = {
   {"_get_shared_fd", (PyCFunction)THPStorage_(_sharedFd), METH_NOARGS, NULL},
   {"_shared_decref", (PyCFunction)THPStorage_(_sharedDecref), METH_NOARGS, NULL},
   {"_shared_incref", (PyCFunction)THPStorage_(_sharedIncref), METH_NOARGS, NULL},
+#endif
+#ifdef THC_GENERIC_FILE
+  {"getDevice", (PyCFunction)THPStorage_(getDevice), METH_NOARGS, NULL},
 #endif
   {"_set_cdata", (PyCFunction)THPStorage_(_setCdata), METH_O, NULL},
   {NULL}
diff --git a/torch/cuda/__init__.py b/torch/cuda/__init__.py
@@ -24,56 +24,76 @@
             "for old NVIDIA drivers")
 
 
-from torch.Storage import _StorageBase
-from torch.Tensor import _TensorBase
+@contextlib.contextmanager
+def device(idx):
+    prev_idx = torch._C._cuda_getDevice()
+    torch._C._cuda_setDevice(idx)
+    yield
+    torch._C._cuda_setDevice(prev_idx)
+
+
+@contextlib.contextmanager
+def _dummy_ctx():
+    yield
+
+
+def deviceCount():
+    return torch._C._cuda_getDeviceCount()
+
 
 ################################################################################
 # Define Storage and Tensor classes
 ################################################################################
 
-class DoubleStorage(torch._C.CudaDoubleStorageBase, _StorageBase):
+
+from .tensor import _CudaTensorBase
+from .storage import _CudaStorageBase
+
+
+class DoubleStorage(torch._C.CudaDoubleStorageBase, _CudaStorageBase):
     pass
-class FloatStorage(torch._C.CudaFloatStorageBase, _StorageBase):
+class FloatStorage(torch._C.CudaFloatStorageBase, _CudaStorageBase):
     pass
-class LongStorage(torch._C.CudaLongStorageBase, _StorageBase):
+class LongStorage(torch._C.CudaLongStorageBase, _CudaStorageBase):
     pass
-class IntStorage(torch._C.CudaIntStorageBase, _StorageBase):
+class IntStorage(torch._C.CudaIntStorageBase, _CudaStorageBase):
     pass
-class ShortStorage(torch._C.CudaShortStorageBase, _StorageBase):
+class ShortStorage(torch._C.CudaShortStorageBase, _CudaStorageBase):
     pass
-class CharStorage(torch._C.CudaCharStorageBase, _StorageBase):
+class CharStorage(torch._C.CudaCharStorageBase, _CudaStorageBase):
     pass
-class ByteStorage(torch._C.CudaByteStorageBase, _StorageBase):
+class ByteStorage(torch._C.CudaByteStorageBase, _CudaStorageBase):
     pass
-class HalfStorage(torch._C.CudaHalfStorageBase, _StorageBase):
+class HalfStorage(torch._C.CudaHalfStorageBase, _CudaStorageBase):
     pass
 
-class DoubleTensor(torch._C.CudaDoubleTensorBase, _TensorBase):
+class DoubleTensor(torch._C.CudaDoubleTensorBase, _CudaTensorBase):
     def is_signed(self):
         return True
-class FloatTensor(torch._C.CudaFloatTensorBase, _TensorBase):
+class FloatTensor(torch._C.CudaFloatTensorBase, _CudaTensorBase):
     def is_signed(self):
         return True
-class LongTensor(torch._C.CudaLongTensorBase, _TensorBase):
+class LongTensor(torch._C.CudaLongTensorBase, _CudaTensorBase):
     def is_signed(self):
         return True
-class IntTensor(torch._C.CudaIntTensorBase, _TensorBase):
+class IntTensor(torch._C.CudaIntTensorBase, _CudaTensorBase):
     def is_signed(self):
         return True
-class ShortTensor(torch._C.CudaShortTensorBase, _TensorBase):
+class ShortTensor(torch._C.CudaShortTensorBase, _CudaTensorBase):
     def is_signed(self):
         return True
-class CharTensor(torch._C.CudaCharTensorBase, _TensorBase):
+class CharTensor(torch._C.CudaCharTensorBase, _CudaTensorBase):
     def is_signed(self):
         # TODO
         return False
-class ByteTensor(torch._C.CudaByteTensorBase, _TensorBase):
+class ByteTensor(torch._C.CudaByteTensorBase, _CudaTensorBase):
     def is_signed(self):
         return False
-class HalfTensor(torch._C.CudaHalfTensorBase, _TensorBase):
+class HalfTensor(torch._C.CudaHalfTensorBase, _CudaTensorBase):
     def is_signed(self):
         return True
 
+
 torch._storage_classes.add(DoubleStorage)
 torch._storage_classes.add(FloatStorage)
 torch._storage_classes.add(LongStorage)
@@ -90,36 +110,33 @@ def is_signed(self):
 torch._tensor_classes.add(CharTensor)
 torch._tensor_classes.add(ByteTensor)
 
-@contextlib.contextmanager
-def device(idx):
-    prev_idx = torch._C._cuda_getDevice()
-    torch._C._cuda_setDevice(idx)
-    yield
-    torch._C._cuda_setDevice(prev_idx)
 
-@contextlib.contextmanager
-def _dummy_ctx():
-    yield
-
-def _tensor_cuda(self, idx=None):
+def _cuda(self, idx=None):
     # This already is a CUDA tensor.
     # Let's check if it needs to be transfered to another GPU.
     if hasattr(self, 'getDevice'):
         target_device = idx if idx else torch._C._cuda_getDevice()
         if self.getDevice() != target_device:
             with device(target_device):
                 return type(self)(self.size()).copy_(self)
+        else:
+            return self
     else:
         ctx = device(idx) if idx else _dummy_ctx()
         with ctx:
             return self.type(getattr(torch.cuda, self.__class__.__name__))
-_TensorBase.cuda = _tensor_cuda
 
-def _tensor_cpu(self):
+
+def _cpu(self):
     return self.type(getattr(torch, self.__class__.__name__))
-_TensorBase.cpu = _tensor_cpu
 
-def deviceCount():
-    return torch._C._cuda_getDeviceCount()
+
+from ..Tensor import _TensorBase
+from ..Storage import _StorageBase
+_TensorBase.cuda = _cuda
+_TensorBase.cpu = _cpu
+_StorageBase.cuda = _cuda
+_StorageBase.cpu = _cpu
+
 
 assert torch._C._cuda_init()
diff --git a/torch/cuda/storage.py b/torch/cuda/storage.py
@@ -0,0 +1,12 @@
+from . import device, _dummy_ctx
+from ..Storage import _StorageBase
+
+
+class _CudaStorageBase(_StorageBase):
+
+    def type(self, *args, **kwargs):
+        source_device = self.getDevice()
+        ctx = device(source_device) if source_device != -1 else _dummy_ctx()
+        with ctx:
+            return super(_CudaStorageBase, self).type(*args, **kwargs)
+
diff --git a/torch/cuda/tensor.py b/torch/cuda/tensor.py
@@ -0,0 +1,12 @@
+from . import device, _dummy_ctx
+from ..Tensor import _TensorBase
+
+
+class _CudaTensorBase(_TensorBase):
+
+    def type(self, *args, **kwargs):
+        source_device = self.getDevice()
+        ctx = device(source_device) if source_device != -1 else _dummy_ctx()
+        with ctx:
+            return super(_CudaTensorBase, self).type(*args, **kwargs)
+
