Repository navigation
nvcomp-simple #4
Description
https://github.com/TomAugspurger/cuda-streams-sample/blob/main/nvcomp-simple.py is looking at nvcomp, trying to overlap host to device memory transfers and decompressing. I've attached a (compressed) nsys profile.
One thing that stands out, which the script attempts to debug: is codec.decode blocking or non-blocking? In the "direct" case I'm just directly doing codec.encode(device_arrays) on the host thread, one per stream. I've put each codec.encode inside an nvtx annotation
with nvtx.annotate("launch-decode"):
codec.encode(device_arrays)From this screenshot, we can see that each codec.encode(device_arrays) takes ~60ms to run, and they're all run serially.
Contrast that with launching the codec.encode in a worker thread:
The codec.encode() is done from a different thread now (via a ThreadPoolExecutor). You can't see it, but there is an nvtx annotation around the future submit:
with nvtx.annotate("launch-decode"):
futures.append(pool.submit(launch_decode, codec, device_arrays))which takes only 20 μs.
I'm not sure yet out much this matters. Even with launching the decode via a thread, we don't get much overlapping H2D copying and decoding.
Activity
TomAugspurger commentedon Jul 11, 2025
The threaded case (calling
codec.decodeon separate worker threads) might be a bit more complicated... Hard to see, but I've highlighted the fourblocks here
I would expect that all of those happen very near each other. And the bulk of the time would be in the
stream.synchronize()below...TomAugspurger commentedon Jul 11, 2025
f60c327 is exploring those long
cudaFreeHost. Previously the return value ofcodec.decodewas immediately discared and so the__del__ran immediately. With that change, we keep around the result until the end of the iteration:Here's a screenshot of the "direct" (no threads) submission:
Note that
launch_decodenow takes ~5-10mscudaFreeHostcalls take long (they all seem to be 10s of μs)New nsys profile attached:
nvcomp-simple.nsys-rep.zip
akshaysubr commentedon Jul 12, 2025
@TomAugspurger These are some very nice tests. Just to confirm, the initial issues were just an artifact of the testing harness not keeping references to the ouputs or the codec object itself right? With the latest screenshot looks like things are fully non-blocking and multiple streams let us get both overlapping memcpys and decodes but also overlapping decodes with other decodes when the GPU is not yet fully saturated.
One other thing I noticed is that the numba host to device copy is using
cuMemAlloc.cuda-streams-sample/nvcomp-simple.py
Line 79 in f60c327
If this application was using CuPy for some ops (parts of the zarr pipeline) and PyTorch for some other ops and nvcomp for the decodes, all of those have their own memory allocators and creating an artificial partitioning of device memory. To help solve this, nvcomp added a custom memory allocator interface which can be stream ordered as well. If we use RMM, we should be able to use the same stream ordered memory allocator across all of these libraries and enable even better overlap opportunities. The challenge is that this is really something that probably must be done in use-space code rather than inside a library since the memory allocator choice impacts unrelated parts of user code, so not quite sure how to properly harmonize here.
TomAugspurger commentedon Jul 14, 2025
Yes, one or both of those. Previously, both the
nvcomp.Codecdevice and the output ofcodec.decodewere dropped inside theforloop. One or both of those was, I think, triggering the blocking behavior.For zarr-python, the output buffers will be consumed by the user, so I don't think we need to worry about the lifetime of them (assuming that we can ensure that the
cupy.ndarraygets is a view on the output ofcodec.decode; that's not currently true but I think we can get there). The codec is cached here. We'll need to go through the codec pipeline code to ensure that it does what we need.TomAugspurger commentedon Jul 14, 2025
Thanks for the note about the allocators. I see that cupy has memory pools for both host and device buffers (https://docs.cupy.dev/en/stable/user_guide/memory.html). I'll keep that in mind.