# Anything like Distributions.jl for GPU?

**URL:** https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147
**Category:** Statistics
**Created:** [November 4, 2018, 8:20pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147 "2018-11-04T20:20:41Z")
**Posts on this page:** 16
**Page:** 1

<div class="post-metadata">

### Author: ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)
#### Post date: [November 4, 2018, 8:20pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/1 "2018-11-04T20:20:41Z")

</div>

Is there any package supporting multivariate distributions and basic operations on them like sampling and log pdf, but working on GPU?

---

<div class="post-metadata">

### Author: ![datnamer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/datnamer/32/3471_2.png) [@datnamer](https://discourse.julialang.org/u/datnamer)
#### Post date: [November 5, 2018, 2:31pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/2 "2018-11-05T14:31:10Z")

</div>

Or, what would it take to get Distributions.jl working on the GPU?

---

<div class="post-metadata">

### Author: ![jonathanBieler](https://avatars.discourse-cdn.com/v4/letter/j/82dd89/32.png) [@jonathanBieler](https://discourse.julialang.org/u/jonathanBieler)
#### Post date: [November 5, 2018, 2:48pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/3 "2018-11-05T14:48:51Z")

</div>

The issue might be that Distributions.jl still uses some C libraries (Rmath, openspecfun). There was some discussion to port them in [Julia](https://discourse.julialang.org/t/where-to-discuss-replacement-of-rmath/9671) but I’m not sure what’s the status of that.

Sometimes I just translate scipy code for probability distributions when I need automatic differentiation.

---

<div class="post-metadata">

### Author: ![andreasnoack](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/andreasnoack/32/27_2.png) [@andreasnoack](https://discourse.julialang.org/u/andreasnoack)
#### Post date: [November 5, 2018, 3:05pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/4 "2018-11-05T15:05:20Z")

</div>

`Distributions` is indeed made with CPUs and `Array`s in mind. It would be great to modify the code to work more generally even if the support on GPUs would initially be sparser than on CPUs. As mentioned, we use `Rmath` for many of the functions in `Distributions` but quite a few of them already have fallbacks written in Julia. However, since the number types on the GPU are the same as the CPU number types, it would be a problem to dispatch to the Julia implementation. This is a general problem.

We might want to consider a split of `Distributions` into a package with all the basic API definitions, say `DistributionsBase`, and then make that a dependency of `Distributions`. In that way, it would be possible for `CuArray`s to use (at least most of) the same API as `Distributions` without depending on all of `Distributions`.

---

<div class="post-metadata">

### Author: ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)
#### Post date: [November 5, 2018, 10:23pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/5 "2018-11-05T22:23:22Z")

</div>

I considered modifying Distributions.jl to work on GPU, but it seems easier to go the other way - start with some basic GPU-compatible implementation and gradually move to existing API of Distributions.jl.

I also think only multivariate distributions since putting scalars to GPU looks quite meaningless.

---

<div class="post-metadata">

### Author: ![andreasnoack](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/andreasnoack/32/27_2.png) [@andreasnoack](https://discourse.julialang.org/u/andreasnoack)
#### Post date: [November 6, 2018, 10:17am UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/6 "2018-11-06T10:17:24Z")

</div>

> …but it seems easier to go the other way - start with some basic GPU-compatible implementation and gradually move to existing API of Distributions.jl.

True but it would be a shame to develop a separate API for GPU. Is would nice to be able to use e.g. `Normal(μ,σ)` and `Binomial(n,p)` on the GPU.

> I also think only multivariate distributions since putting scalars to GPU looks quite meaningless.

I don’t understand why you wouldn’t like to work with scalar valued distributions on the GPU, e.g. for things like

```julia
rand!(Normal(), dA)

```

to also work on the GPU.

---

<div class="post-metadata">

### Author: ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)
#### Post date: [November 6, 2018, 7:14pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/7 "2018-11-06T19:14:49Z")

</div>

> I don’t understand why you wouldn’t like to work with scalar valued distributions on the GPU, e.g. for things like […]

Ah, I thought more about generating separate numbers (e.g. something like `rand(Normal())`) when I said about univariate distributions, generating a vector of such numbers on GPU makes perfect sense to me.

But in the case of `rand!(Normal(), dA)` it seems like we don’t have a problem with dispatching since we can always dispatch on type of `dA`. Or am I missing something?

---

<div class="post-metadata">

### Author: ![andreasnoack](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/andreasnoack/32/27_2.png) [@andreasnoack](https://discourse.julialang.org/u/andreasnoack)
#### Post date: [November 6, 2018, 7:19pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/8 "2018-11-06T19:19:55Z")

</div>

I think the specific example might be fine but I suspect that something like

```julia
pdf.(Normal(), dA)

```

could be a problem because the broadcast function would end up calling e.g. `pdf(Normal(), dA[i])` and the elements would be, say `Float32`s, so `pdf` would dispatch to the wrong method.

---

<div class="post-metadata">

### Author: ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)
#### Post date: [November 6, 2018, 9:29pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/9 "2018-11-06T21:29:48Z")

</div>

Indexing of individual elements in `CuArrays` is inefficient and disabled by default if I remember correctly, so dispatching is the least of the problems I would say. On the other hand, in Julia 1.0 broadcasting is relatively easy to override.

---

<div class="post-metadata">

### Author: ![andreasnoack](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/andreasnoack/32/27_2.png) [@andreasnoack](https://discourse.julialang.org/u/andreasnoack)
#### Post date: [November 7, 2018, 10:11am UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/10 "2018-11-07T10:11:26Z")

</div>

> [@dfdx](#):
>
> Indexing of individual elements in `CuArrays` is inefficient and disabled by default if I remember correctly, so dispatching is the least of the problems I would say. On the other hand, in Julia 1.0 broadcasting is relatively easy to override.

I was probably too brief. The call to `pdf(Normal(), dA[i])` would happen on the GPU device where indexing is, of course, not disabled. This is already what CuArray broadcasting does so the problem can be illustrated with

```julia
julia> cdf.(Normal(), cu(rand(3)))
ERROR: InvalidIRError: compiling #25(CuArrays.CuKernelState, CUDAnative.CuDeviceArray{Float64,1,CUDAnative.AS.Global}, Base.Broadcast.Broadcasted{Nothing,Tupl
e{Base.OneTo{Int64}},typeof(cdf),Tuple{CuArrays.CuRefValue{Normal{Float64}},Base.Broadcast.Extruded{CUDAnative.CuDeviceArray{Float32,1,CUDAnative.AS.Global},T
uple{Bool},Tuple{Int64}}}}) resulted in invalid LLVM IR
Reason: unsupported call to the Julia runtime (call to jl_f_tuple)
Stacktrace:
 [1] erfc at /home/andreasnoack/.julia/packages/SpecialFunctions/fvheQ/src/erf.jl:8
 [2] normcdf at /home/andreasnoack/.julia/packages/StatsFuns/0W2sM/src/distrs/norm.jl:15
 [3] normcdf at /home/andreasnoack/.julia/packages/StatsFuns/0W2sM/src/distrs/norm.jl:16
 [4] cdf at /home/andreasnoack/.julia/packages/Distributions/WHjOk/src/univariates.jl:528
 [5] _broadcast_getindex_evalf at broadcast.jl:574
 [6] _broadcast_getindex at broadcast.jl:547
 [7] getindex at broadcast.jl:507
 [8] macro expansion at /home/andreasnoack/.julia/packages/GPUArrays/hzyWn/src/abstract_gpu_interface.jl:69
 [9] #25 at /home/andreasnoack/.julia/packages/GPUArrays/hzyWn/src/broadcast.jl:58
Reason: unsupported call to the Julia runtime (call to jl_type_error_rt)
Stacktrace:
 [1] erfc at /home/andreasnoack/.julia/packages/SpecialFunctions/fvheQ/src/erf.jl:8
 [2] normcdf at /home/andreasnoack/.julia/packages/StatsFuns/0W2sM/src/distrs/norm.jl:15
 [3] normcdf at /home/andreasnoack/.julia/packages/StatsFuns/0W2sM/src/distrs/norm.jl:16
 [4] cdf at /home/andreasnoack/.julia/packages/Distributions/WHjOk/src/univariates.jl:528
 [5] _broadcast_getindex_evalf at broadcast.jl:574
 [6] _broadcast_getindex at broadcast.jl:547
 [7] getindex at broadcast.jl:507
 [8] macro expansion at /home/andreasnoack/.julia/packages/GPUArrays/hzyWn/src/abstract_gpu_interface.jl:69
 [9] #25 at /home/andreasnoack/.julia/packages/GPUArrays/hzyWn/src/broadcast.jl:58

```

---

<div class="post-metadata">

### Author: ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)
#### Post date: [November 16, 2018, 10:42pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/11 "2018-11-16T22:42:38Z")

</div>

It took me some time to learn internals of Distributions.jl, it’s dependencies and theory behind them. I tried to imagine what it would take to make a single distribution - namely, `MvNormal` - working on CuArrays (or any other form of GPU arrays). By “working” I mean:

- parameters are represented as CuArrays or wrappers (e.g. `Diagonal{T, CuArray{T},1}`)
- `logpdf` or any similar method is implemented

The most acute problem that I’ve encountered is that under the hood Distributions.jl heavily uses LinearAlgebra types such as `UpperTriangular` matrix and `Cholesky` factorization. These types, however, use concrete `Array` instead of `AbstractArray`. I didn’t dive deeper, but I believe using concrete types there is intended and will be hard to change in Base. Also, Cholesky factorization for CuArrays is better to do using cuSOLVER.

Unless I miss some important detail, this means we will have to re-implement half of the LinearAlgebra types, which sounds like a terrible idea. I’ll continue thinking about it, fresh idea are welcome too.

---

<div class="post-metadata">

### Author: ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)
#### Post date: [November 16, 2018, 11:40pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/12 "2018-11-16T23:40:07Z")

</div>

It seems like I misinterpreted what I saw in REPL while inspecting types in LinearAlgebra - at least Cholesky factorization turns to work perfectly well:

```julia
A = cu(rand(3,3))
A = A*A' # make positive-definite 
cholesky(A)

```

---

<div class="post-metadata">

### Author: ![andreasnoack](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/andreasnoack/32/27_2.png) [@andreasnoack](https://discourse.julialang.org/u/andreasnoack)
#### Post date: [November 18, 2018, 7:31pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/13 "2018-11-18T19:31:23Z")

</div>

I’d argue the right place to start for this effort is not with the multivariate distributions but the univariates. The multivariate distributions are more complicated and would probably some work before they’ll work for a broader range of array types.

---

<div class="post-metadata">

### Author: ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)
#### Post date: [November 19, 2018, 9:34pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/14 "2018-11-19T21:34:00Z")

</div>

I’m not so sure about this. As I see it:

1. You can’t represent a single number on GPU, the closest you can do is a single-element `CuArray`. Even if you could, it only makes sense as a part of processing of a larger batch of data - either broadcasting or batch operations like matrix multiplication.

2. Broadcasting on univariate distribution type, e.g. `cdf.(Normal(), cu(rand(3)))`, is hard since you need to pass a custom Julia struct (like `Normal`) to GPU which is poorly documented in C and seems to be completely missing from CUDAnative. So working towards this approach means improving CUDAnative runtime itself, which is, well, quite low-level.  
An alternative approach is to overload broadcasting and rewrite such calls into e.g. `cdfnormal(mu, sigma, cu(rand(3)))`.

3. Multivariate distributions, on the other hand, can operate directly on high-level `CuArray`s. In the simplest case, you just use `CuArray`s for distribution type fields (e.g. `mu` and `sigma` for `MvNormal`) and apply generic functions to them. Perhaps, most of the work in this case goes on re-implementing CPU-only and third-party code, which sounds to me like an easier starting point than digging into internals of CUDAnative.

Perhaps you have another vision, so I’d be happy to learn about it!

---

<div class="post-metadata">

### Author: ![andreasnoack](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/andreasnoack/32/27_2.png) [@andreasnoack](https://discourse.julialang.org/u/andreasnoack)
#### Post date: [November 20, 2018, 8:44am UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/15 "2018-11-20T08:44:06Z")

</div>

> Broadcasting on univariate distribution type, e.g. `cdf.(Normal(), cu(rand(3)))`, is hard since you need to pass a custom Julia struct (like `Normal`) to GPU which is poorly documented in C and seems to be completely missing from `CUDAnative`.

I think you might be overthinking this. What I’m talking about would be something like

```julia
julia> using CuArrays

julia> struct Normal{T}
         μ::T
         σ::T
       end

julia> pdf(D::Normal{T}, x::T) where T = CuArrays.CUDAnative.exp(-((x-D.μ)/D.σ)^2/2)/(sqrt(2*T(π))*D.σ)
pdf (generic function with 1 method)

julia> X = cu(rand(10))
10-element CuArray{Float32,1}:
 0.49907735
 0.0075366762
 0.4265667
 0.044595595
 0.6838826
 0.24200504
 0.95713586
 0.8394181
 0.8209742
 0.46985078

julia> Base.Broadcast.broadcastable(D::Normal) = Ref(D)

julia> pdf.(Normal(0.3f0, 0.9f0), X)
10-element CuArray{Float32,1}:
 0.43255657
 0.42047203
 0.4389076
 0.42577493
 0.40472618
 0.44234982
 0.3395486
 0.37039313
 0.3748919
 0.43544522

```

---

<div class="post-metadata">

### Author: ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)
#### Post date: [September 8, 2019, 9:29pm UTC](https://discourse.julialang.org/t/anything-like-distributions-jl-for-gpu/17147/16 "2019-09-08T21:29:42Z")

</div>

> [@dfdx](#):
>
> Broadcasting on univariate distribution type, e.g. `cdf.(Normal(), cu(rand(3)))` , is hard since you need to pass a custom Julia struct (like `Normal` ) to GPU which is poorly documented in C and seems to be completely missing from CUDAnative. So working towards this approach means improving CUDAnative runtime itself, which is, well, quite low-level.

CUDAnative can do codegen with broadcast. You can make and return any Julia struct. You can internally use mutable structs as long as they go away. This makes pure-Julia codes “just work” on the GPU as long as they are fully inferred and return isbits types.

The issue that you will run into with CuArrays is that you need to use the RNG from CUDA if you want to do `rand.(dist)`, and without a way to pass a rng that won’t work. But the other functions should be fine.
