User memory allocation & non-blocking execution

[mariogeiger]

I’d like to request two related enhancements to cuFINUFFT:

1. User-provided GPU memory allocation:
   Allow passing in user-managed GPU workspaces instead of performing internal cudaMalloc/cudaFree calls. This would enable seamless integration with frameworks like PyTorch or JAX, which maintain their own GPU memory pools and expect full control over allocation and deallocation.

2. Non-blocking, asynchronous execution:
   Support fully asynchronous launches that avoid implicit CPU synchronizations (e.g., from hidden memory allocations or stream synchronizations). Frameworks like PyTorch and JAX rely on overlapping GPU execution with CPU-side scheduling — allowing the CPU to stay ahead and queue work — to minimize Python and other host-side overheads. Blocking behavior prevents these frameworks from efficiently pipelining GPU workloads.

[ahbarnett]

Thanks. We already have
https://github.com/flatironinstitute/jax-finufft
https://github.com/flatironinstitute/pytorch-finufft
so we'd be interested in knowing if/how your requests would change/improve them.

I am not the GPU expert here, so I'll loop in @DiamonDinoia @lu1and10 @lgarrison @blackwer @janden to see what our team has to say.

[blackwer]

Not the current maintainer, so just adding my 2 cents from my experience with various things in the cufinufft ecosystem.

1. Allow passing in user-managed GPU workspaces instead of performing internal cudaMalloc/cudaFree calls. This would enable seamless integration with frameworks like PyTorch or JAX, which maintain their own GPU memory pools and expect full control over allocation and deallocation.

cufinufft's memory footprint is quite small compared to most other things going on in GPU world, so I'd need to see a real use case scenario where this is proving to be a problem. A lot of the memory would be coming from the torch vectors being passed in, which would be handled by torch. The engineering burden would be quite large I imagine, and it would be an increased maintenance burden. I'm not saying no (nor do I have the right to :D), but I would need to see something concrete before saying that this would be worth it.

2. Non-blocking, asynchronous execution:
   Support fully asynchronous launches that avoid implicit CPU synchronizations (e.g., from hidden memory allocations or stream synchronizations). Frameworks like PyTorch and JAX rely on overlapping GPU execution with CPU-side scheduling — allowing the CPU to stay ahead and queue work — to minimize Python and other host-side overheads. Blocking behavior prevents these frameworks from efficiently pipelining GPU workloads.

Are you looking for something more than our existing stream API? There are only one or two stream syncs afaik (during setpts and plan creation -- execution should be sync free, iirc). https://github.com/flatironinstitute/finufft/blob/master/python/cufinufft/examples/example3d2many_async_cupy.py for our basic example of abusing stream concurrency. Once again, I'm not saying no (or that this can't be improved), but I'd want a concrete example of our current algorithm causing problems before saying it'd be worth it.

All that said, a completely non-blocking, memory source agnostic implementation of everything would be a nice have, but there are a lot of structural elements of cufinufft that I think would need considerable overhauling to make it work. Some of the algorithms would need some serious rethinking to make properly fully async as well, I think. The code is mostly written in more "classic" cuda with a lot of manual memory handling. This makes certain improvements like external memory management likely a nightmare to implement (getting async in the first place was a monumental amount of work partly because of this).

[lu1and10]

There are only one or two stream syncs afaik (during setpts and plan creation -- execution should be sync free, iirc).

I second this, I double checked the cufinufft execute, it seems after @blackwer and @DiamonDinoia cleaned up the cuda code, execute should be sync free.