Issue with stochastic forcing implementation on GPU

[mncrowe]

Hi All,

I'm having an issue with running the singlelayerqg_betaforced.jl example with dev = GPU(). It seems to be that calcF! calls are failing though it's unclear to me what the issue is (error output below). Removing calcF=calcF! from the prob definition allows the script to run and the unforced example runs with no issues.

Is this some new bug introduced by an update to one of the packages or perhaps an issue with my installation?

Matt

out.log

[navidcy]

thanks!

if you change

@. Fh = sqrt(forcing_spectrum) * cis(2π * random_uniform(T)) / sqrt(clock.dt)

to

Fh .= sqrt.(forcing_spectrum) .* cis.(2π * random_uniform(T, size(forcing_spectrum)...)) ./ sqrt(clock.dt)

it should work. However, it allocates an array and I'm trying to find a better way to do this and will open a PR.

[glwagner]

The problem is to generate random numbers without allocating data? Is rand! supported by CUDA?

[glwagner]

I think it is:

help?> CUDA.rand!
  rand!([rng=default_rng()], A, [S=eltype(A)])

  Populate the array A with random values. If S is specified (S can be a type or a collection, cf. rand for details), the values are picked
  randomly from S. This is equivalent to copyto!(A, rand(rng, S, size(A))) but without allocating a new array.

  Examples
  ≡≡≡≡≡≡≡≡

  julia> rng = MersenneTwister(1234);

  julia> rand!(rng, zeros(5))
  5-element Vector{Float64}:
   0.5908446386657102
   0.7667970365022592
   0.5662374165061859
   0.4600853424625171
   0.7940257103317943

[glwagner]

Not sure what distribution you're going for but basically you need to preallocate an array:

sz = size(forcing_spectrum)
ϵ = CuArray(rand(T, sz...))

# then later, within calcF!
rand!(ϵ)
@. Fh = sqrt(forcing_spectrum) * cis(2π * ϵ) / sqrt(clock.dt)

bit of a sketch but its something like that

[navidcy]

Which version of CUDA you using? Because I don't get this for CUDA v5.2.

[glwagner]

5.1.2

[navidcy]

I found that Random.rand! will work both on Arrays and CuArrays provided you give

Random.rand!(Random.default_rng(), A)

or

Random.rand!(CUDA.default_rng(), A)

[glwagner]

(BaroclinicAdjustments) pkg> add CUDA@5.2
   Resolving package versions...
  No Changes to `~/Projects/AwesomeOceananigans/BaroclinicAdjustments/Project.toml`
  No Changes to `~/Projects/AwesomeOceananigans/BaroclinicAdjustments/Manifest.toml`

(BaroclinicAdjustments) pkg> st
Status `~/Projects/AwesomeOceananigans/BaroclinicAdjustments/Project.toml`
  [052768ef] CUDA v5.2.0
  [e9467ef8] GLMakie v0.9.9
⌃ [9e8cae18] Oceananigans v0.90.9
  [de0858da] Printf
Info Packages marked with ⌃ have new versions available and may be upgradable.

julia> using CUDA

julia> CUDA.rand!
rand! (generic function with 90 methods)

help?> CUDA.rand!
  rand!([rng=default_rng()], A, [S=eltype(A)])

  Populate the array A with random values. If S is specified (S can be a type or a collection, cf. rand for details), the values are picked
  randomly from S. This is equivalent to copyto!(A, rand(rng, S, size(A))) but without allocating a new array.

  Examples
  ≡≡≡≡≡≡≡≡

  julia> rng = MersenneTwister(1234);

  julia> rand!(rng, zeros(5))
  5-element Vector{Float64}:
   0.5908446386657102
   0.7667970365022592
   0.5662374165061859
   0.4600853424625171
   0.7940257103317943

[glwagner]

oh so it does work

[navidcy]

OK, then that's the cleanest solution! rand! it is!

[navidcy]

OK, for the record, the best solution is to change the calcF! and bit above to:

random_normal! = dev==CPU() ? Random.randn! :
                 dev==GPU() ? CUDA.randn! :
                 error("dev must be CPU() or GPU()")

function calcF!(Fh, sol, t, clock, vars, params, grid)
  random_normal!(Fh)
  @. Fh *= sqrt(forcing_spectrum) / sqrt(clock.dt)
end

[glwagner]

OK, for the record, the best solution is to change the calcF! and bit above to:

random_normal! = dev==CPU() ? Random.randn! :
                 dev==GPU() ? CUDA.randn! :
                 error("dev must be CPU() or GPU()")

function calcF!(Fh, sol, t, clock, vars, params, grid)
  random_normal!(Fh)
  @. Fh *= sqrt(forcing_spectrum) / sqrt(clock.dt)
end

Do you need the if-statement? Is it really true that randn! does not dispatch on array type?

[navidcy]

Do you need the if-statement? Is it really true that randn! does not dispatch on array type?

You are correct!

julia> using CUDA, Random, BenchmarkTools

julia> A = CUDA.zeros(3, 4)
3×4 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
 0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0
 0.0  0.0  0.0  0.0

julia> @btime CUDA.@sync Random.randn!($A);
  16.554 μs (10 allocations: 352 bytes)

julia> @btime CUDA.@sync CUDA.randn!($A);
  17.563 μs (10 allocations: 352 bytes)

[glwagner]

Here's how you see it directly:

julia> Random.randn! === CUDA.randn!
true