Repository: cupy/cupy

Issue:
Offer a `cupy.cuda.get_allocator` , and a pinned allocator that can associate with a particular device. Current workaround allows 110x speed over Pytorch CPU pinned tensors
I am filing these requests on the suggestion of this stackoverflow answer

https://stackoverflow.com/questions/57752516/how-to-use-cuda-pinned-zero-copy-memory-for-a-memory-mapped-file

Which is an answer to how to pinn cupy tensors to CPU memory. This is important, because pinned Cupy tensors with dlpack offer the fastest CPU<->GPU transfer for Pytorch cuda variables. 

I have made a library around this method, which can achieve speeds up to 110x than Pytorch pinned cpu tensors. 

https://github.com/Santosh-Gupta/SpeedTorch

Here is the direct issue suggestions from the stackoverflow answer

> It appears to me that currently, cupy doesn't offer a pinned allocator that can be used in place of the usual device memory allocator, i.e. could be used as the backing for cupy.ndarray. If this is important to you, you might consider filing a cupy issue.

> after doing what you need with pinned memory (allocations) you should probably revert the cupy allocator to its default value. Unfortunately, unlike cupy.cuda.set_allocator, I did not find a corresponding cupy.cuda.get_allocator, which strikes me as a deficiency in cupy, something that also seems worthy of filing a cupy issue to me. 

> Ordinary device memory associated with cupy's default device memory allocator, has an association with a particular device. pinned memory need not have such an association, however our trivial replacement of BaseMemory with a lookalike class means that we are suggesting to cupy that this "device" memory, like all other ordinary device memory, has a specific device association. In a single device setting such as yours, this distinction is meaningless. However, this isn't suitable for robust multi-device use of pinned memory. For that, again the suggestion would be a more robust change to cupy, perhaps by filing an issue.





Discussion:
Probably related to #1899.
I think that the main reason behind the original stackoverflow post might be this
https://github.com/cupy/cupy/blob/master/cupy/core/core.pyx#L1945

Even though the array is memory-mapped, it is unraveled loading it into the pinned host memory and then the object is explicitly transferred to the device.

The zero-copy case can only be addressed by creating a new memory pool as the reply in SO proposes. However, as the array is residing in a pinned memory pool for fast data transfers, that memory should be directly accessible from the device, but there is an explicit cuda allocation on the device and a copy afterward as generally, we want the arrays to live in the device memory.

The allocator changing trick what it does is to allocate unified memory in the host instead of allocating it into the device, so all the data lives in the host memory and is transferred to the device only when it is accessed.
Some details on how it's speeding up transfer processes in certain cases for Pytorch

https://discuss.pytorch.org/t/introducing-speedtorch-4x-speed-cpu-gpu-transfer-110x-gpu-cpu-transfer/56147

Resolving patch:
diff --git a/cupy/cuda/__init__.py b/cupy/cuda/__init__.py
--- a/cupy/cuda/__init__.py
+++ b/cupy/cuda/__init__.py
@@ -97,6 +97,7 @@ def get_cuda_path():
 from cupy.cuda.memory import MemoryPointer  # NOQA
 from cupy.cuda.memory import MemoryPool  # NOQA
 from cupy.cuda.memory import set_allocator  # NOQA
+from cupy.cuda.memory import get_allocator  # NOQA
 from cupy.cuda.memory import UnownedMemory  # NOQA
 from cupy.cuda.memory_hook import MemoryHook  # NOQA
 from cupy.cuda.pinned_memory import alloc_pinned_memory  # NOQA


Test patch:
diff --git a/tests/cupy_tests/cuda_tests/test_memory.py b/tests/cupy_tests/cuda_tests/test_memory.py
--- a/tests/cupy_tests/cuda_tests/test_memory.py
+++ b/tests/cupy_tests/cuda_tests/test_memory.py
@@ -626,6 +626,9 @@ def test_set_allocator(self):
             self.assertEqual(1024, arr.data.mem.size)
             self.assertEqual(1024, self.pool.used_bytes())
 
+    def test_get_allocator(self):
+        assert memory.get_allocator() == self.pool.malloc
+
     @unittest.skipUnless(sys.version_info[0] >= 3,
                          'Only for Python3 or higher')
     def test_reuse_between_thread(self):