Specifying deterministic forcing on GPUs

[chowland]

Hi there,

I'm trying to set up a simple Kolmogorov flow forced by a single sinusoidal mode using GeophysicalFlows. I have a working example using CPU, but I'm struggling to get it working on GPU. My calcF! function that I pass to the Problem is as follows:

forcing_mode = 4
function setF!(Fh, sol, t, clock, vars, params, grid)
  Fh .= zeros(eltype(grid), size(sol))
  CUDA.@allowscalar Fh[1,forcing_mode] = -forcing_mode*nx*ny
  return nothing
end

My inexperience with GPUs shows here, however, with CUDA complaining about GPU compilation of the kernel.
Here's a minimal example of the issue:

using GeophysicalFlows

dev = GPU()

n = 128
L = 2π
ν = 1e-2
nν= 1
μ = 0
nμ= 0
dt = 5e-4

function setF!(Fh, sol, t, clock, vars, params, grid)
    Fh = @. zeros(eltype(grid), size(sol))
    return nothing
end

prob = TwoDNavierStokes.Problem(dev; nx=n, Lx=L, ν, nν, μ, nμ, dt, stepper="ETDRK4",
    calcF=setF!, stochastic=false)

sol, clock, vars, params, grid = prob.sol, prob.clock, prob.vars, prob.params, prob.grid

setF!(vars.Fh, sol, 0.0, clock, vars, params, grid)

And the backtrace that appears afterwards:

ERROR: LoadError: MethodError: no method matching length(::TwoDGrid{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, CUDA.CUFFT.cCuFFTPlan{ComplexF64, -1, false, 2}, CUDA.CUFFT.rCuFFTPlan{Float64, -1, false, 2}, UnitRange{Int64}, GPU})
Closest candidates are:
  length(::Union{Base.KeySet, Base.ValueIterator}) at abstractdict.jl:58
  length(::CUDA.CUSPARSE.CuSparseMatrix) at ~/.julia/packages/CUDA/DfvRa/lib/cusparse/array.jl:223
  length(::Union{JLD2.Group, JLD2.JLDFile}) at ~/.julia/packages/JLD2/1YVED/src/JLD2.jl:466
  ...
Stacktrace:
 [1] _similar_shape(itr::TwoDGrid{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, CUDA.CUFFT.cCuFFTPlan{ComplexF64, -1, false, 2}, CUDA.CUFFT.rCuFFTPlan{Float64, -1, false, 2}, UnitRange{Int64}, GPU}, #unused#::Base.HasLength)
   @ Base ./array.jl:663
 [2] _collect(cont::UnitRange{Int64}, itr::TwoDGrid{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, CUDA.CUFFT.cCuFFTPlan{ComplexF64, -1, false, 2}, CUDA.CUFFT.rCuFFTPlan{Float64, -1, false, 2}, UnitRange{Int64}, GPU}, #unused#::Base.HasEltype, isz::Base.HasLength)
   @ Base ./array.jl:718
 [3] collect(itr::TwoDGrid{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, CUDA.CUFFT.cCuFFTPlan{ComplexF64, -1, false, 2}, CUDA.CUFFT.rCuFFTPlan{Float64, -1, false, 2}, UnitRange{Int64}, GPU})
   @ Base ./array.jl:712
 [4] broadcastable(x::TwoDGrid{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, CUDA.CUFFT.cCuFFTPlan{ComplexF64, -1, false, 2}, CUDA.CUFFT.rCuFFTPlan{Float64, -1, false, 2}, UnitRange{Int64}, GPU})
   @ Base.Broadcast ./broadcast.jl:704
 [5] broadcasted
   @ ./broadcast.jl:1296 [inlined]
 [6] setF!(Fh::CUDA.CuArray{ComplexF64, 2, CUDA.Mem.DeviceBuffer}, sol::CUDA.CuArray{ComplexF64, 2, CUDA.Mem.DeviceBuffer}, t::Float64, clock::FourierFlows.Clock{Float64}, vars::GeophysicalFlows.TwoDNavierStokes.Vars{CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{ComplexF64, 2, CUDA.Mem.DeviceBuffer}, CUDA.CuArray{ComplexF64, 2, CUDA.Mem.DeviceBuffer}, Nothing}, params::GeophysicalFlows.TwoDNavierStokes.Params{Float64}, grid::TwoDGrid{Float64, CUDA.CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, CUDA.CUFFT.cCuFFTPlan{ComplexF64, -1, false, 2}, CUDA.CUFFT.rCuFFTPlan{Float64, -1, false, 2}, UnitRange{Int64}, GPU})
   @ Main /gpfs/work3/0/pof/howland/kolmogorov/mwe.jl:14
 [7] top-level scope
   @ /gpfs/work3/0/pof/howland/kolmogorov/mwe.jl:23
in expression starting at /gpfs/work3/0/pof/howland/kolmogorov/mwe.jl:23

If you can help me understand how to properly write such forcing functions for GeophysicalFlows on GPUs, I would be happy to add a note to the docs to help others get started with this.
Thanks!

[navidcy]

There are a few problems with what you do here. First

zeros(eltype(grid), size(sol))

is an array on the CPU. Also, vars.Fh is complex-valued but the array you create in setF! is real-valued.

(There is another "problem" that what you are doing up there is not that efficient since you are allocating a whole array every time you call setF!.)

This

function setF!(Fh, sol, t, clock, vars, params, grid)
  @. Fh = 0
  return nothing
end

is device-agnostic and thus will work. But perhaps just zeroing out the forcing is not what you really aiming for but you were just playing around?

If you need to create an array, there is device_array:

help?> device_array
search: device_array

  device_array(device::Device)
  device_array(device::Device, T, dim)


  Return the proper array type according to the device, i.e., Array for CPU and CuArray for GPU.

which will let you create the appropriate array on the device you are on, e.g.,

device_array(dev)(zeros(eltype(grid), size(sol))

For example, this would also work:

julia> function setF!(Fh, sol, t, clock, vars, params, grid)
           Fh = device_array(dev)(zeros(complex(eltype(grid)), size(sol)))
           return nothing
       end
setF! (generic function with 1 method)

julia> setF!(vars.Fh, sol, 0.0, clock, vars, params, grid)

Btw, another thing I didn't quite understand is what the @. you have in Fh = @. ... was aiming to achieve.

[chowland]

Thanks Navid! This is super helpful. I was trying to stick close to the existing example of stochastic forcing for 2D Navier-Stokes, and now I see where the difference between CPU and GPU arrays comes from there. Zeroing the forcing array was just to get a minimal example of the same error, but the actual forcing I want to use is at the top of my post, with a single non-zero forcing mode. Our GPU cluster is down for the day, but hopefully I'll have something that can work tomorrow.
And yes, that final @. was just me trying different things in a bit of desperation - it doesn't make a lot of sense.