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

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

docs/source/cuda.rst

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

test/test_cuda.py

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

torch/_tensor.py

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

torch/cuda/__init__.py

@@ -19,7 +19,7 @@ from typing import List, Optional, Tuple, Union, Any
 from ._utils import _get_device_index, _dummy_type
 from .._utils import classproperty
 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
 import torch._C