GPU support for PoissionRandom

[sivasathyaseeelan]

Is your feature request related to a problem? Please describe.

Currently PoissionRandom is not supported inside kernel or is lacks GPU support

Additional context

https://github.com/SciML/JumpProcesses.jl/pull/490/files/7fe4771406476c03d9408a2188641b1c4ef88dd5#diff-b6d1e0f8515b90cd025c2200da0dbb441ae7688506cc560ac6c4080cde7d6c7b

[ChrisRackauckas]

Does using PoissonRandom.count_rand not work?

[sivasathyaseeelan]

Does using PoissonRandom.count_rand not work?

No. Any function which takes rng::AbstractRNG as argument is not work

[devmotion]

Just use CUDA.rand_poisson/CUDA.rand_poisson!?

[isaacsas]

Similarly there is https://rocm.docs.amd.com/projects/rocRAND/en/develop/doxygen/html/group__rocrandhost.html#gace088e83d118e5cf683a82573bbe5bd7 so it seems like at least those two could perhaps be wrapped with appropriate dispatches.

[ChrisRackauckas]

Just use CUDA.rand_poisson/CUDA.rand_poisson!?

No, it doesn't work right off the bat. There's work to be done and making things simpler is how we will make progress.

Two things here. First is that we will need to change the algorithm a bit to work in a GPU kernel. Some things are hardcoded to Float64 in ad_rand https://github.com/SciML/PoissonRandom.jl/blob/master/src/PoissonRandom.jl#L76-L90 so that's not going to compile with things like Metal. @oscardssmith was going to look into that. But secondly, the KernelAbstractions dispatches require rand is not called with an rng. So the first step here is to just see if

function count_rand(λ)
    n = 0
    c = randexp(rng)
    while c < λ
        n += 1
        c += randexp(rng)
    end
    return n
end

works as a GPU kernel, which has the highest change of working because of its simplicity.

So step one is getting count_rand to work, since that should be easy. Step 2 is to then get ad_rand to work, which will likely take a bit more time and a different polynomial approximation in Float32 which @oscardssmith needs to make. Step 3 is then

Similarly there is https://rocm.docs.amd.com/projects/rocRAND/en/develop/doxygen/html/group__rocrandhost.html#gace088e83d118e5cf683a82573bbe5bd7 so it seems like at least those two could perhaps be wrapped with appropriate dispatches.

to make it so that if we are in the GPU context, we just swap to the GPU kernels provided by the host libraries for Poisson random numbers. CUDA also has a specialized one https://docs.nvidia.com/cuda/curand/device-api-overview.html#poisson-api-example and we should use it there too. But Metal and Intel OpenAPI don't, so having the fallback is always necessary.

So let's take it one step at a time, first let's get the simple thing working and then build up.

[devmotion]

Actually I don't think any of this should be done in PoissonRandom but the package should just be deleted: #6

Some things are hardcoded to Float64

I fixed this in Distributions 6 years ago: JuliaStats/Distributions.jl#1021

[ChrisRackauckas]

https://github.com/JuliaStats/Distributions.jl/blob/abb151c42eb67b4994e21da15e24c65f1cc3dd9f/src/samplers/poisson.jl#L67 but it has the rngs, can those be removed so we can use it for gpu kernels?

[isaacsas]

It would be helpful if Distributions updated the docs to explain and show the direct samplers. Currently they seem like purely internal code that is not exported, and lacks both docstrings and mention in the official documentation.

[isaacsas]

@ChrisRackauckas is it not possible to create a dummy RNG that just dispatches the basic functions to call the versions that don’t take rngs?

[devmotion]

https://github.com/JuliaStats/Distributions.jl/blob/abb151c42eb67b4994e21da15e24c65f1cc3dd9f/src/samplers/poisson.jl#L67 but it has the rngs, can those be removed so we can use it for gpu kernels?

Sure, it has to have the RNGs - as for every rand implementation (also in Random) that's the only way to ensure reproducibility and have control over the stream of random numbers. AFAICT the way to go is not to remove rng arguments in rand calls (which seems undesirable apart from GPU specifics anyway and generally inconsistent with the interface of rand) but to use a custom dummy struct NoRNG <: Random.AbstractRNG end that dispatches rand(rng::NoRNG, ::Type{T}, dims...) calls to rand(T, dims...).

Edit: Ah, sorry, I completely missed #51 (comment) 😅 Yes, that's exactly what I'd suggest.

[devmotion]

It would be helpful if Distributions updated the docs to explain and show the direct samplers. Currently they seem like purely internal code that is not exported, and lacks both docstrings and mention in the official documentation.

The API is Distributions.sampler which is documented, see https://juliastats.org/Distributions.jl/stable/extends/#Distributions.sampler-Tuple{Distribution} and the interface for univariate etc. distributions below. The details of the types are internal/non-public.

[isaacsas]

Sorry for my confusion, then what is the most performant, allocation-free method to sample a Poisson RV using Distributions.jl that has similar performance to this package? I thought that thread you previously linked said to explicitly use rand(rng, Distributions.PoissonADSampler(λ)), which is using an internal type if I've understood correctly?

[devmotion]

No, I didn't suggest using the internal samplers directly - I just had to use them directly to be able to benchmark the implementations. Users should just use rand([rng, ]Poisson(lambda)[, dims...) (or the in-place version). Only in case you'd like to cache the sampler for separate rand calls you'd use sampler (without caring about the implementation details).

[oscardssmith]

Does anyone remember why this isn't using LogExpFunctions.jl for log1pmx? Was this just separated too early? If so, I would propose adding a dependency and removing log1pmx from this package. If that is acceptable, JuliaStats/LogExpFunctions.jl#94 will add the Float32 log1pmx. If not, I can duplicate this PR once merged.