OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Document torch.cuda.ExternalStream, torch.cuda.caching_allocator_alloc and torch.cuda.caching_allocator_delete (#70126)

Browse Source 
Summary:
Fixes https://github.com/pytorch/pytorch/issues/67414. Fixes https://github.com/pytorch/pytorch/issues/70117.

cc brianjo mruberry ngimel

Pull Request resolved: https://github.com/pytorch/pytorch/pull/70126

Reviewed By: mruberry

Differential Revision: D33542910

Pulled By: ngimel

fbshipit-source-id: 4b870f4dceca6ee4cc8fba58819f1cb18ac9f857
This commit is contained in:
Leo Fang 2022-01-12 15:43:07 -08:00 committed by Facebook GitHub Bot
parent ad803936d1
commit 67941c8a94
4 changed files with 7 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
docs/source/cuda.rst
View File

@ -71,6 +71,7 @@ Streams and events
:nosignatures: :nosignatures:
Stream Stream
ExternalStream
Event Event
Graphs (beta) Graphs (beta)
@ -106,6 +107,8 @@ Memory management
max_memory_cached max_memory_cached
reset_max_memory_cached reset_max_memory_cached
reset_peak_memory_stats reset_peak_memory_stats
caching_allocator_alloc
caching_allocator_delete
.. FIXME The following doesn't seem to exist. Is it supposed to? .. FIXME The following doesn't seem to exist. Is it supposed to?
https://github.com/pytorch/pytorch/issues/27785 https://github.com/pytorch/pytorch/issues/27785
.. autofunction:: reset_max_memory_reserved .. autofunction:: reset_max_memory_reserved

4
test/test_cuda.py
View File

@ -1336,7 +1336,7 @@ class TestCuda(TestCase):
def test_external_streams(self): def test_external_streams(self):
device = torch.cuda.device(0) device = torch.cuda.device(0)
with self._get_external_stream(device) as stream_v: with self._get_external_stream(device) as stream_v:
ext_stream = torch.cuda.streams.ExternalStream(stream_v) ext_stream = torch.cuda.ExternalStream(stream_v)
self.assertEqual(stream_v, ext_stream.cuda_stream) self.assertEqual(stream_v, ext_stream.cuda_stream)
self.assertEqual(ext_stream.device.index, device.idx) self.assertEqual(ext_stream.device.index, device.idx)
@ -1345,7 +1345,7 @@ class TestCuda(TestCase):
def test_external_streams_multi_device(self): def test_external_streams_multi_device(self):
device = torch.cuda.device(1) device = torch.cuda.device(1)
with self._get_external_stream(device) as stream_v: with self._get_external_stream(device) as stream_v:
ext_stream = torch.cuda.streams.ExternalStream( ext_stream = torch.cuda.ExternalStream(
stream_v, device=device) stream_v, device=device)
self.assertEqual(stream_v, ext_stream.cuda_stream) self.assertEqual(stream_v, ext_stream.cuda_stream)
self.assertEqual(ext_stream.device.index, device.idx) self.assertEqual(ext_stream.device.index, device.idx)

2
torch/_tensor.py
View File

@ -1169,7 +1169,7 @@ class Tensor(torch._C._TensorBase):
raise TypeError('stream must be ``int`` or ``none``') raise TypeError('stream must be ``int`` or ``none``')
elif stream is not None and stream != -1: elif stream is not None and stream != -1:
if self.device.type == 'cuda': if self.device.type == 'cuda':
stream = torch.cuda.streams.ExternalStream(stream) stream = torch.cuda.ExternalStream(stream)
# Only synchronize on different streams # Only synchronize on different streams
if stream != torch.cuda.current_stream: if stream != torch.cuda.current_stream:
event = torch.cuda.Event() event = torch.cuda.Event()

2
torch/cuda/__init__.py
View File

@ -19,7 +19,7 @@ from typing import List, Optional, Tuple, Union, Any
from ._utils import _get_device_index, _dummy_type from ._utils import _get_device_index, _dummy_type
from .._utils import classproperty from .._utils import classproperty
from .graphs import CUDAGraph, graph_pool_handle, graph, make_graphed_callables from .graphs import CUDAGraph, graph_pool_handle, graph, make_graphed_callables
from .streams import Stream, Event from .streams import ExternalStream, Stream, Event
from .. import device as _device from .. import device as _device
import torch._C import torch._C