[Python] How to speedup cupy code with frequent data transfer & writing

I have a computational code largely written based on cupy. The computation I'm running now requires high-frequency data transfer (from GPU to CPU) and writing to a h5 file.

Here is a sketch of what I'm doing now (FYI, this is done on a single GPU):

import numpy as np
import cupy as cp
import h5py

u_gpu = cp.zeros((n, n, ...))
u_cpu = np.zeros((n, n, ...))

f = h5py.File(...)
while ...:
# computation goes here
...

# data transfer
u_cpu[:] = u_gpu.get()

# write u_cpu to h5 file
f.create_dataset(..., data=u_cpu)

f.close()


I know that, by default, data transfer by .get() requires synchronization of the GPU & CPU, thus slowing down the whole process. In my tests, the code slows down by more than 20 times compared to NO data transfer & writing. Even worse, such a slowdown somehow accumulates: for example, the first 50000 steps take 30 mins (which is comparable to the time needed for NO data transfer & writing), but then the next 50000 will take a longer time, around 1 hr, and the next takes 1.5 hrs, and so on.

The potential solutions I've found so far:

1. Use pinned memory (see this; although this thread seems somewhat outdated)
2. Async data transfer; see this comment. But I'm not sure how this works and if this will overwrite data or anything due to this asynchronization
3. Buffer data (honestly I got this from ChatGPT). Save the data for x steps on GPU, and then transfer the accumulated data to CPU for writing, thus reducing the frequency of doing data transfer & writing.
4. Probably a more trivial improvement: changing from u_cpu[:] = u_gpu.get() to u_gpu.get(out=u_cpu). This speeds up the data transfer by roughly 1 time, based on my naive tests.

I'm new to GPU computation and so cupy. I'm not sure which one would work or if there are better ways to resolve my problem.

Any suggestions are welcomed. Thanks!

Continue reading...