Reland "Add torch.utils.device_mode" (pytorch#91796)

Original PR pytorch#91525

Signed-off-by: Edward Z. Yang <ezyangfb.com>
Pull Request resolved: pytorch#91796
Approved by: https://github.com/albanD

Author: ezyang

docs/source/tensor_attributes.rst

@@ -177,6 +177,20 @@ Via a string and device ordinal:
     >>> torch.device('cpu', 0)
     device(type='cpu', index=0)
 
+The device object can also be used as a context manager to change the default
+device tensors are allocated on:
+
+::
+
+    >>> with torch.device('cuda:1'):
+    ...     r = torch.randn(2, 3)
+    >>> r.device
+    device(type='cuda', index=1)
+
+This context manager has no effect if a factory function is passed an explicit,
+non-None device argument.  To globally change the default device, see also
+:func:`torch.set_default_device`.
+
 .. note::
    The :class:`torch.device` argument in functions can generally be substituted with a string.
    This allows for fast prototyping of code.

docs/source/torch.rst

@@ -17,6 +17,7 @@ Tensors
     is_nonzero
     set_default_dtype
     get_default_dtype
+    set_default_device
     set_default_tensor_type
     numel
     set_printoptions

test/test_autograd.py

@@ -122,22 +122,36 @@ def test_grad_mode_class_decoration(self):
         with self.assertWarnsRegex(UserWarning, "Decorating classes is deprecated"):
             @torch.no_grad()
             class Foo():
-                pass
+                def __init__(self):
+                    assert not torch.is_grad_enabled()
+
+                def foo(self):
+                    # Not applied to methods
+                    assert torch.is_grad_enabled()
+
+            # Show that we can actually construct the class
+            foo = Foo()
+            foo.foo()
 
         # Decorating functions or methods is fine though
         with warnings.catch_warnings(record=True) as w:
             @torch.no_grad()
             def foo():
-                pass
+                assert not torch.is_grad_enabled()
+
+            foo()
 
             class Foo2():
                 @torch.no_grad()
                 def __init__(self):
-                    pass
+                    assert not torch.is_grad_enabled()
 
                 @torch.no_grad()
                 def foo(self):
-                    pass
+                    assert not torch.is_grad_enabled()
+
+            foo2 = Foo2()
+            foo2.foo()
 
         self.assertEqual(len(w), 0)
 

test/test_overrides.py

@@ -760,7 +760,7 @@ def __getattr__(self, name):
         val = getattr(self._data, name)
 
         # If it's a method
-        if callable(val):
+        if not isinstance(val, torch.device) and callable(val):
             c = getattr(type(self._data), name)
             # Don't append self to args if classmethod/staticmethod
             if c is val:

test/test_utils.py

@@ -14,8 +14,17 @@
 import torch.nn as nn
 import torch.utils.data
 from torch.utils.data import DataLoader
+from torch.testing._internal.common_device_type import (
+    ops,
+    onlyCPU,
+    instantiate_device_type_tests,
+)
+from torch.testing._internal.common_methods_invocations import op_db
 import torch.cuda
+from torch.utils._pytree import tree_any, tree_all_only
 from torch.utils.checkpoint import checkpoint, checkpoint_sequential
+from torch import set_default_device
+from torch.utils._device import set_device
 import torch.utils.cpp_extension
 from torch.autograd._functions.utils import check_onnx_broadcast
 from torch.onnx.symbolic_opset9 import _prepare_onnx_paddings
@@ -796,6 +805,74 @@ def test_external_module_register(self):
             torch._register_device_module('xpu', DummyXPUModule)
 
 
+class TestDeviceUtils(TestCase):
+    def test_basic(self):
+        with torch.device('meta') as dev:
+            x = torch.empty(3, 3)
+        self.assertEqual(x.device.type, 'meta')
+        self.assertEqual(dev, torch.device('meta'))
+
+    def test_decorator(self):
+        @set_device('meta')
+        def f():
+            return torch.empty(3, 3)
+        self.assertEqual(f().device.type, 'meta')
+
+    def test_decorator_generator(self):
+        @set_device('meta')
+        def f():
+            yield torch.empty(3, 3)
+            yield torch.empty(3, 3)
+        r1, r2 = list(f())
+        self.assertEqual(r1.device.type, 'meta')
+        self.assertEqual(r2.device.type, 'meta')
+
+
+    def test_nn_module(self):
+        with torch.device('meta'):
+            m = nn.Linear(40, 50)
+        self.assertEqual(m.weight.device.type, 'meta')
+
+    def test_set_default_device(self):
+        try:
+            set_default_device('meta')
+            r = torch.empty(2, 2)
+        finally:
+            set_default_device(None)
+
+        self.assertEqual(r.device.type, 'meta')
+
+    @onlyCPU
+    @ops(op_db)
+    def test_device_mode_ops(self, device, dtype, op):
+        func = op.get_op()
+        samples = op.sample_inputs(device, dtype, requires_grad=False)
+        for sample in samples:
+            # Only test samples which don't have Tensor inputs.  However,
+            # we don't test the factory property on OpInfo as it is very,
+            # very incomplete
+            if tree_any(
+                lambda x: isinstance(x, torch.Tensor),
+                (sample.input, sample.args, sample.kwargs)
+            ):
+                continue
+            # Many OpInfos will explicitly pass in a device.  DeviceContext
+            # will respect device if it is explicitly specified.  To test
+            # DeviceContext, we have to remove the device kwarg in this case.
+            # NB: Can't pass None to sample_inputs, the function can't
+            # handle it.
+            kwargs = sample.kwargs.copy()
+            kwargs.pop('device', None)
+            with torch.device('meta'):
+                r = func(sample.input, *sample.args, **kwargs)
+            self.assertTrue(
+                tree_all_only(torch.Tensor, lambda x: x.device.type == 'meta', r)
+            )
+
+
+instantiate_device_type_tests(TestDeviceUtils, globals())
+
+
 class TestCppExtensionUtils(TestCase):
     def test_cpp_compiler_is_ok(self):
         self.assertTrue(torch.utils.cpp_extension.check_compiler_ok_for_platform('c++'))

torch/_C/__init__.pyi.in

@@ -45,6 +45,13 @@ class device:
     @overload
     def __init__(self, type: str, index: _int) -> None: ...
 
+    # Uncomment if we ever make torch.device a decorator
+    # def __call__(self, func: T) -> T: ...
+
+    def __enter__(self) -> "device": ...
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None: ...
+
     def __reduce__(self) -> Tuple[Any, ...]: ...  # THPDevice_reduce
 
 # Defined in torch/csrc/Stream.cpp

torch/__init__.py

@@ -35,6 +35,7 @@
 
 __all__ = [
     'typename', 'is_tensor', 'is_storage', 'set_default_tensor_type',
+    'set_default_device',
     'set_rng_state', 'get_rng_state', 'manual_seed', 'initial_seed', 'seed',
     'save', 'load', 'set_printoptions', 'chunk', 'split', 'stack', 'matmul',
     'no_grad', 'enable_grad', 'rand', 'randn', 'inference_mode',
@@ -444,6 +445,49 @@ def is_storage(obj):
     return type(obj) in _storage_classes
 
 
+_GLOBAL_DEVICE_CONTEXT = None
+
+def set_default_device(device):
+    """Sets the default ``torch.Tensor`` to be allocated on ``device``.  This
+    does not affect factory function calls which are called with an explicit
+    ``device`` argument.  Factory calls will be performed as if they
+    were passed ``device`` as an argument.
+
+    To only temporarily change the default device instead of setting it
+    globally, use ``with torch.device(device):`` instead.
+
+    The default device is initially ``cpu``.  If you set the default tensor
+    device to another device (e.g., ``cuda``) without a device index, tensors
+    will be allocated on whatever the current device for the device type,
+    even after :func:`torch.cuda.set_device` is called.
+
+    Args:
+        device (device or string): the device to set as default
+
+    Example::
+
+        >>> # xdoctest: +SKIP("requires cuda, changes global state")
+        >>> torch.tensor([1.2, 3]).device
+        device(type='cpu')
+        >>> torch.set_default_device('cuda')  # current device is 0
+        >>> torch.tensor([1.2, 3]).device
+        device(type='cuda', index=0)
+        >>> torch.set_default_device('cuda:1')
+        >>> torch.tensor([1.2, 3]).device
+        device(type='cuda', index=1)
+
+    """
+    global _GLOBAL_DEVICE_CONTEXT
+    if _GLOBAL_DEVICE_CONTEXT is not None:
+        _GLOBAL_DEVICE_CONTEXT.__exit__(None, None, None)
+    if device is None:
+        _GLOBAL_DEVICE_CONTEXT = None
+        return
+    from torch.utils._device import DeviceContext
+    _GLOBAL_DEVICE_CONTEXT = DeviceContext(device)
+    _GLOBAL_DEVICE_CONTEXT.__enter__()
+
+
 def set_default_tensor_type(t):
     r"""Sets the default ``torch.Tensor`` type to floating point tensor type
     ``t``. This type will also be used as default floating point type for

torch/autograd/grad_mode.py

@@ -1,96 +1,11 @@
-import sys
 import torch
-import functools
-import inspect
-import warnings
-from typing import Any, Callable, TypeVar, cast
+from typing import Any
+
+from torch.utils._contextlib import _DecoratorContextManager
 
 __all__ = ['no_grad', 'enable_grad', 'set_grad_enabled',
            'inference_mode', 'set_multithreading_enabled']
 
-
-# Used for annotating the decorator usage of 'no_grad' and 'enable_grad'.
-# See https://mypy.readthedocs.io/en/latest/generics.html#declaring-decorators
-FuncType = Callable[..., Any]
-F = TypeVar('F', bound=FuncType)
-
-
-class _DecoratorContextManager:
-    """Allow a context manager to be used as a decorator"""
-
-    def __call__(self, func: F) -> F:
-        if inspect.isclass(func):
-            warnings.warn("Decorating classes is deprecated and will be disabled in "
-                          "future versions. You should only decorate functions or methods. "
-                          "To preserve the current behavior of class decoration, you can "
-                          "directly decorate the `__init__` method and nothing else.")
-
-        if inspect.isgeneratorfunction(func):
-            return self._wrap_generator(func)
-
-        @functools.wraps(func)
-        def decorate_context(*args, **kwargs):
-            with self.clone():
-                return func(*args, **kwargs)
-        return cast(F, decorate_context)
-
-    def _wrap_generator(self, func):
-        """Wrap each generator invocation with the context manager"""
-        @functools.wraps(func)
-        def generator_context(*args, **kwargs):
-            gen = func(*args, **kwargs)
-
-            # Generators are suspended and unsuspended at `yield`, hence we
-            # make sure the grad mode is properly set every time the execution
-            # flow returns into the wrapped generator and restored when it
-            # returns through our `yield` to our caller (see PR #49017).
-            try:
-                # Issuing `None` to a generator fires it up
-                with self.clone():
-                    response = gen.send(None)
-
-                while True:
-                    try:
-                        # Forward the response to our caller and get its next request
-                        request = yield response
-
-                    except GeneratorExit:
-                        # Inform the still active generator about its imminent closure
-                        with self.clone():
-                            gen.close()
-                        raise
-
-                    except BaseException:
-                        # Propagate the exception thrown at us by the caller
-                        with self.clone():
-                            response = gen.throw(*sys.exc_info())
-
-                    else:
-                        # Pass the last request to the generator and get its response
-                        with self.clone():
-                            response = gen.send(request)
-
-            # We let the exceptions raised above by the generator's `.throw` or
-            # `.send` methods bubble up to our caller, except for StopIteration
-            except StopIteration as e:
-                # The generator informed us that it is done: take whatever its
-                # returned value (if any) was and indicate that we're done too
-                # by returning it (see docs for python's return-statement).
-                return e.value
-
-        return generator_context
-
-    def __enter__(self) -> None:
-        raise NotImplementedError
-
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-        raise NotImplementedError
-
-    def clone(self):
-        # override this method if your children class takes __init__ parameters
-        return self.__class__()
-
-
 class no_grad(_DecoratorContextManager):
     r"""Context-manager that disabled gradient calculation.