# Is there a GPU optimized \`Distances\` package?

**URL:** <https://discourse.julialang.org/t/is-there-a-gpu-optimized-distances-package/46434>\
**Category:** General Usage\
**Created:** [September 11, 2020, 6:53am UTC](https://discourse.julialang.org/t/is-there-a-gpu-optimized-distances-package/46434 "2020-09-11T06:53:13Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![yiyuezhuo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yiyuezhuo/32/27101_2.png) [@yiyuezhuo](https://discourse.julialang.org/u/yiyuezhuo)\
**Post date:** [September 11, 2020, 6:53am UTC](https://discourse.julialang.org/t/is-there-a-gpu-optimized-distances-package/46434/1 "2020-09-11T06:53:13Z")

</div>

I wonder if there is a package which implements `Distances.jl`’s interface with `CUDA.AbstractGPUArray` type to leverage GPU acceleration. I found those issues:

[https://github.com/JuliaStats/Distances.jl/issues/143](https://github.com/JuliaStats/Distances.jl/issues/143)

[https://github.com/JuliaStats/Distances.jl/issues/137](https://github.com/JuliaStats/Distances.jl/issues/137)

From them, it seems that it’s impossible (improper) to introduce CUDA specific type to `Distances.jl` even with `Requires`. So maybe is there a package like `DistributionsAD.jl` which extends CUDA specific methods with `Distances.jl`’s interface?

For example, I wrote following code to speed up my Haversine distance computing:

```julia
using Distances
import Distances: pairwise
using CUDA

function pairwise(dist::Haversine, x::CUDA.AbstractGPUArray{Float32, 2}, y::CUDA.AbstractGPUArray{Float32, 2})
    # x, y: (2, n) GPU Array
    
    # longtitudes
    Δλ = deg2rad.(y[1,:] .- x[1,:]')
    
    # latitudes
    φ₁ = deg2rad.(x[2,:])
    φ₂ = deg2rad.(y[2,:])
    
    Δφ = φ₂ .- φ₁'
    
    # haversine formula
    a = sin.(Δφ./2).^2 .+ cos.(φ₂) .* cos.(φ₁') .* sin.(Δλ./2) .^ 2
    
    # distance on the sphere
    return 2 .* dist.radius .* asin.(min.(sqrt.(a), fill!(similar(a), 1))) # take care of floating point errors
end

```

While I don’t consider any deep optimization, the performance boosting looks huge:

```julia
using BenchmarkTools

hd = Haversine(100)

tx2 = randn(2, 10000)
cu_tx2 = cu(tx2)

# precompile and ensure the correctness
pairwise(hd, tx2, tx2) ≈ collect(pairwise(hd, cu_tx2, cu_tx2))
# true

@benchmark pairwise(hd, tx2, tx2)
#=
BenchmarkTools.Trial: 
  memory estimate: 762.94 MiB
  allocs estimate: 3
  --------------
  minimum time: 3.525 s (0.04% GC)
  median time: 3.532 s (0.03% GC)
  mean time: 3.532 s (0.03% GC)
  maximum time: 3.539 s (0.03% GC)
  --------------
  samples: 2
  evals/sample: 1
=#

@benchmark pairwise(hd, cu_tx2, cu_tx2)
#=
BenchmarkTools.Trial: 
  memory estimate: 17.45 KiB
  allocs estimate: 505
  --------------
  minimum time: 1.789 ms (0.00% GC)
  median time: 39.941 ms (0.00% GC)
  mean time: 39.150 ms (0.90% GC)
  maximum time: 87.350 ms (2.36% GC)
  --------------
  samples: 128
  evals/sample: 1
=#

```

---

<div class="post-metadata">

**Author:** ![aplavin](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/aplavin/32/222056_2.png) [@aplavin](https://discourse.julialang.org/u/aplavin)\
**Post date:** [September 12, 2020, 12:08am UTC](https://discourse.julialang.org/t/is-there-a-gpu-optimized-distances-package/46434/2 "2020-09-12T00:08:57Z")

</div>

Does your function definition really require CUDA? I would expect that replacing `CUDA.AbstractGPUArray{Float32, 2}` with `AbstractMatrix` gives exactly the same results and performance.

---

<div class="post-metadata">

**Author:** ![yiyuezhuo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yiyuezhuo/32/27101_2.png) [@yiyuezhuo](https://discourse.julialang.org/u/yiyuezhuo)\
**Post date:** [September 12, 2020, 1:02am UTC](https://discourse.julialang.org/t/is-there-a-gpu-optimized-distances-package/46434/3 "2020-09-12T01:02:47Z")

</div>

`AbstractMatrix` will override original implementation of `Distances.jl` (the performance will be worse than original implementation on CPU I guess). But it’s possible to just `using GPUArrays` to prevent `using CUDA`, so `GPUArrays.JLArrays` can be used to test code in pure CPU mode.

---

<div class="post-metadata">

**Author:** ![yiyuezhuo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yiyuezhuo/32/27101_2.png) [@yiyuezhuo](https://discourse.julialang.org/u/yiyuezhuo)\
**Post date:** [September 12, 2020, 1:11am UTC](https://discourse.julialang.org/t/is-there-a-gpu-optimized-distances-package/46434/4 "2020-09-12T01:11:41Z")

</div>

Such a package would be useful as I have noticed someone is trying to add GPU support using similar style. If a package `DistancesGPU.jl` exists, one will just `using` it and pass a `GPUArray` to the normal `Distances.jl` function.

[https://github.com/willtebbutt/Stheno.jl/commit/f4120f7429113b642b8218a74ae5dcb854057594#diff-206b99ce211ae6ceba2a8bef4c11980fR34](https://github.com/willtebbutt/Stheno.jl/commit/f4120f7429113b642b8218a74ae5dcb854057594#diff-206b99ce211ae6ceba2a8bef4c11980fR34)

---

<div class="post-metadata">

**Author:** ![mcabbott](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mcabbott/32/6603_2.png) [@mcabbott](https://discourse.julialang.org/u/mcabbott)\
**Post date:** [September 12, 2020, 1:49am UTC](https://discourse.julialang.org/t/is-there-a-gpu-optimized-distances-package/46434/5 "2020-09-12T01:49:21Z")

</div>

It’s possible that Distances should only call its methods on things fairly close to `::Array`, and then methods like this for `::AbstractArray` could be added there.

Yours seems to be quicker than what’s in Distances on the CPU, and can be made faster. You could probably avoid some more trig by re-writing in terms of `cos(φ/2)`.

```julia
julia> function pairwise_bc2(dist::Haversine, x, y)
           @views Δλ = deg2rad.(y[1,:] .- x[1,:]')
           φ₁ = deg2rad.(@view x[2,:])
           φ₂ = deg2rad.(@view y[2,:])
           a = Δλ .= sin.((φ₂ .- φ₁')./2).^2 .+ (cos.(φ₂)) .* (cos.(φ₁))' .* sin.(Δλ./2) .^ 2
           return Δλ .= 2 .* dist.radius .* asin.(min.(sqrt.(a), 1))
       end;

julia> tx3 = randn(2, 2_000);

julia> @btime pairwise_bc1($hd, $tx3, $tx3); # as in Q
  122.366 ms (16 allocations: 152.68 MiB)

julia> @btime pairwise_bc2($hd, $tx3, $tx3);
  88.243 ms (5 allocations: 30.56 MiB)

julia> @btime Distances.pairwise($hd, $tx3, $tx3);
  166.760 ms (3 allocations: 30.52 MiB)

julia> Distances.pairwise(hd, tx3, tx3) ≈ pairwise_bc2(hd, tx3, tx3)
true

```