Repository: cupy/cupy

Issue:
How to leverage memory of multiple GPUs (with NCCL) using unified memory
Using `set_allocator(MemoryPool(malloc_managed).malloc)`, the "unified memory" seems to allocate/use CPU and (one) GPU memory, not memory of multiple GPUs available.

How can I use memory of multiple GPUs available on the system, possibly with NCCL?
I have tested on a system with four P100 GPUs.

My use-case is - I have large 3D arrays that don't fit into single GPU memory, and I'd like to do some operations (e.g., FFT/IFFT) on these arrays. For example,
```python
import cupy as cp

...

# A and B are large 3D numpy arrays, e.g., 4096x4096x100
arrayA = cp.array(A)
arrayB = cp.array(B)

output = cp.fft.ifftn(cp.multiply(cp.fft.fft(arrayA), arrayB)).real
```
Discussion on multi-GPU support for cuFFT and cuSOLVER
Original context: See #2644.

Some of CUDA libraries like cuFFT and cuSOLVER provide native multi-GPU support. However, there is no guarantee that multiple GPUs will outperform single GPU; on the contrary, for cuFFT I found in https://github.com/cupy/cupy/pull/2644#issuecomment-558835428 that any workload that can be handled in one GPU should ***not*** be processed using multi GPUs; the extra data transfer just kills the performance (even if NVLinks are present). I suspect the same holds for cuSOLVER.

Since CuPy operations incurs additional memory overhead, I think it is clear that multi-GPU is useful only when the workload is too large to fit in one GPU. For such excessive workloads, it's better/unavoidable to start from NumPy arrays directly, such as the "NumPy-in, NumPy-out" feature currently implemented in #2644. **This thread is for discussions regarding this feature and other possible alternative strategies.**


Discussion:
Sorry, current CuPy does not support such feature.

FYI, Can you NVLINK on your environment? If you use only PCI-ex, you can not get performance improvement.
Thanks for your reply, and sorry for my delay in getting back to you.
If I have NVLINK, then the above code should execute faster using (one) GPU and CPU memory?

Thank you
> If I have NVLINK, then the above code should execute faster using (one) GPU and CPU memory?

Maybe yes. Nvidia provides the unified memory and wide band width memory access on NVLINK environment.
Thank you for your comments!
@asi1024 Any chance you've discussed this issue further internally? Thanks.
@leofang Sorry for my late response. We are considering that unified memory may be useful for this issue. ([Related issue in stack overflow](https://stackoverflow.com/questions/57752516/how-to-use-cuda-pinned-zero-copy-memory-for-a-memory-mapped-file))
@asi1024 Did you mean ***pinned*** memory or ***unified*** memory? For the former, yes it will benefit the performance of host-device transfer, and it is possible to create a NumPy array from a pinned memory, something like 
```python
    ptr = cp.cuda.alloc_pinned_memory(itemsize*size)
    np_array = np.frombuffer(ptr, dtype, size).reshape(shape)
```
So the "NumPy-in, NumPy-out" feature will speed up further. The SO post you linked is along this line, with extra effort that the host memory is backed by disk via mmap.

For the latter, I don't think it would help. Unified memory allocated via `cudaMallocManaged` is simply for writing CPU/GPU agnostic code and making code porting easier. It does not provide the sought-after optimization: 
> An important point is that a carefully tuned CUDA program that uses streams and cudaMemcpyAsync to efficiently overlap execution with data transfers may very well perform better than a CUDA program that only uses Unified Memory.
> ...
> Unified Memory is first and foremost a productivity feature that provides a smoother on-ramp to parallel computing, without taking away any of CUDA’s features for power users.

(from https://devblogs.nvidia.com/unified-memory-in-cuda-6/). My understanding for using unified memory is that it would need lots of device synchronization (see the examples therein).
you can have a Cupy array living in unified memory, which allows it to be bigger than a single gpu memory size.
@emcastillo True, but I don’t think this is supported by cuFFT, and I suspect it’s not supported by cuSOLVER either. 

For multi-GPU cuFFT the content must live on GPUs in a specified format, which is supposed to be managed by `cufftXtMalloc` but I handled manually in #2644. If unified memory is supported there, it must be an undocumented feature.
Update: #2644 is merged, but currently the "NumPy-in, NumPy-out" feature is removed (3e2b7d0 and 705fe1a). We should revisit this in the near future. I still think it is a necessary feature for performance reasons.
One more thing:
> For multi-GPU cuFFT the content must live on GPUs in a specified format, which is supposed to be managed by cufftXtMalloc but I handled manually in #2644. If unified memory is supported there, it must be an undocumented feature.

I think I was wrong about the unified memory. From [the cuFFT doc](https://docs.nvidia.com/cuda/cufft/index.html#accessing-cufft):
> Functions in the cuFFT and cuFFTW library assume that the data is in GPU visible memory. This means any memory allocated by cudaMalloc, cudaMallocHost and cudaMallocManaged or registered with cudaHostRegister can be used as input, output or plan work area with cuFFT and cuFFTW functions. For the best performance input data, output data and plan work area should reside in device memory.

It seems data managed by the unified memory system can be used, and moreover host data pointer can be passed to cuFFT routines. But we will need to do some performance benchmarks to determine the best strategy. It is also unclear if multi-GPU cuFFT supports unified memory similarly.

It's unclear to me if the cuSOLVER can do similar things. 
Rel: #1429
@asi1024 Can I bring back part of the "NumPy-in, NumPy-out" feature (at the `cupy.cuda.cufft.Plan1d.fft()` level)? I still think it's very useful for handling massive FFT. We can mark this feature experimental (or just don't make any announcement yet until we're ready) and leave the high-level API (`cupy.fft.fft()` and friends) untouched for further discussion. 

Resolving patch:
diff --git a/cupy/fft/fft.py b/cupy/fft/fft.py
--- a/cupy/fft/fft.py
+++ b/cupy/fft/fft.py
@@ -155,6 +155,8 @@ def _fft_c2c(a, direction, norm, axes, overwrite_x, plan=None):
 
 def _fft(a, s, axes, norm, direction, value_type='C2C', overwrite_x=False,
          plan=None):
+    if isinstance(a, np.ndarray):
+        raise TypeError('The input array a must be a cupy.ndarray')
     if norm not in (None, 'ortho'):
         raise ValueError('Invalid norm value %s, should be None or "ortho".'
                          % norm)
@@ -475,6 +477,8 @@ def _exec_fftn(a, direction, value_type, norm, axes, overwrite_x,
 
 def _fftn(a, s, axes, norm, direction, value_type='C2C', order='A', plan=None,
           overwrite_x=False, out=None):
+    if isinstance(a, np.ndarray):
+        raise TypeError('The input array a must be a cupy.ndarray')
     if norm not in (None, 'ortho'):
         raise ValueError('Invalid norm value %s, should be None or "ortho".'
                          % norm)


Test patch:
diff --git a/tests/cupy_tests/cuda_tests/test_cufft.py b/tests/cupy_tests/cuda_tests/test_cufft.py
--- a/tests/cupy_tests/cuda_tests/test_cufft.py
+++ b/tests/cupy_tests/cuda_tests/test_cufft.py
@@ -1,7 +1,14 @@
 import pickle
 import unittest
 
+import numpy
+
+from cupy import testing
 from cupy.cuda import cufft
+from cupy.fft import config
+from cupy.fft.fft import _convert_fft_type
+
+from ..fft_tests.test_fft import (multi_gpu_config, _skip_multi_gpu_bug)
 
 
 class TestExceptionPicklable(unittest.TestCase):
@@ -11,3 +18,82 @@ def test(self):
         e2 = pickle.loads(pickle.dumps(e1))
         assert e1.args == e2.args
         assert str(e1) == str(e2)
+
+
+# This class tests multi-GPU Plan1d with data sitting on host.
+# Internally, we use cuFFT's data transfer API to ensure the
+# data is in order.
+
+@testing.parameterize(*testing.product({
+    'shape': [(64,), (4, 16), (128,), (8, 32)],
+}))
+@testing.multi_gpu(2)
+class TestMultiGpuPlan1dNumPy(unittest.TestCase):
+
+    @multi_gpu_config(gpu_configs=[[0, 1], [1, 0]])
+    @testing.for_complex_dtypes()
+    def test_fft(self, dtype):
+        _skip_multi_gpu_bug(self.shape, self.gpus)
+
+        a = testing.shaped_random(self.shape, numpy, dtype)
+
+        if len(self.shape) == 1:
+            batch = 1
+            nx = self.shape[0]
+        elif len(self.shape) == 2:
+            batch = self.shape[0]
+            nx = self.shape[1]
+
+        # compute via cuFFT
+        cufft_type = _convert_fft_type(a.dtype, 'C2C')
+        plan = cufft.Plan1d(nx, cufft_type, batch, devices=config._devices)
+        out_cp = numpy.empty_like(a)
+        plan.fft(a, out_cp, cufft.CUFFT_FORWARD)
+
+        out_np = numpy.fft.fft(a)
+        # np.fft.fft alway returns np.complex128
+        if dtype is numpy.complex64:
+            out_np = out_np.astype(dtype)
+
+        assert numpy.allclose(out_cp, out_np, rtol=1e-4, atol=1e-7)
+
+        # compute it again to ensure Plan1d's internal state is reset
+        plan.fft(a, out_cp, cufft.CUFFT_FORWARD)
+
+        assert numpy.allclose(out_cp, out_np, rtol=1e-4, atol=1e-7)
+
+    @multi_gpu_config(gpu_configs=[[0, 1], [1, 0]])
+    @testing.for_complex_dtypes()
+    def test_ifft(self, dtype):
+        _skip_multi_gpu_bug(self.shape, self.gpus)
+
+        a = testing.shaped_random(self.shape, numpy, dtype)
+
+        if len(self.shape) == 1:
+            batch = 1
+            nx = self.shape[0]
+        elif len(self.shape) == 2:
+            batch = self.shape[0]
+            nx = self.shape[1]
+
+        # compute via cuFFT
+        cufft_type = _convert_fft_type(a.dtype, 'C2C')
+        plan = cufft.Plan1d(nx, cufft_type, batch, devices=config._devices)
+        out_cp = numpy.empty_like(a)
+        plan.fft(a, out_cp, cufft.CUFFT_INVERSE)
+        # normalization
+        out_cp /= nx
+
+        out_np = numpy.fft.ifft(a)
+        # np.fft.fft alway returns np.complex128
+        if dtype is numpy.complex64:
+            out_np = out_np.astype(dtype)
+
+        assert numpy.allclose(out_cp, out_np, rtol=1e-4, atol=1e-7)
+
+        # compute it again to ensure Plan1d's internal state is reset
+        plan.fft(a, out_cp, cufft.CUFFT_INVERSE)
+        # normalization
+        out_cp /= nx
+
+        assert numpy.allclose(out_cp, out_np, rtol=1e-4, atol=1e-7)