Types of parameters that can be passed to forcing function on GPU

[raphaelouillon]

Hi all,

I am interested in a simulation where a buoyancy source moves in space and time and the position of the source is interpolated from a time and position array. It appears that Interpolations.jl doesn't work on GPU, so I tried instead to pass the time and position arrays as parameters to the forcing function, but that return an error, as the arrays I am passing to the function are CPU arrays of type Vector{Float64} and not GPU compatible. Is there a way I can pass arrays to the forcing function?

[glwagner]

On Nvidia GPUs we use CuArray rather than Array. You'd need to use CuArray with Interpolations.jl too on the GPU.

[raphaelouillon]

Thanks @glwagner, that's what I have been trying to do following an example on JuliaGPU (see code snippet below)

struct Interpolate{A}
    xs::A
    ys::A
end

function (itp::Interpolate)(x)
    i = searchsortedfirst(itp.xs, x)
    i = clamp(i, firstindex(itp.ys), lastindex(itp.ys))
    @inbounds itp.ys[i]
end

function Adapt.adapt_structure(to, itp::Interpolate)
    xs = Adapt.adapt_structure(to, itp.xs)
    ys = Adapt.adapt_structure(to, itp.ys)
    Interpolate(xs, ys)
end

itp_x = Interpolate(CuArray(ctime), CuArray(cE))
itp_y = Interpolate(CuArray(ctime), CuArray(cN))

@inline function buoyancy_source(x, y, z, t, p)

	R = p.R
        δᴸ = p.δᴸ
	s = p.s

	xc = itp_x(t)#xc = 50*cos(0.01*t)#
	yc = itp_y(t)
	zc = 2.0 #zc = zcolf(xc,yc)

	d = sqrt((x - xc)^2 + (y - yc)^2+ (z - zc)^2)

    return - 0.5*s* (1 - erf((d - R) / δᴸ))#  sᵦ(t) * 0.5 * (1 - erf((d - R) / δᴸ)) # 0.5 * (1 - erf((d - R) / δᴸ))#-
end

but this returns an "unsupported dynamic function invocation" error

ERROR: LoadError: InvalidIRError: compiling kernel gpu_calculate_Gc!(Cassette.Context{nametype(CUDACtx), Nothing, Nothing, KernelAbstractions.var"##PassType#291", Nothing, Cassette.DisableHooks}, typeof(Oceananigans.Models.NonhydrostaticModels.gpu_calculate_Gc!), KernelAbstractions.CompilerMetadata{KernelAbstractions.NDIteration.StaticSize{(64, 64, 16)}, KernelAbstractions.NDIteration.DynamicCheck, Nothing, Nothing, KernelAbstractions.NDIteration.NDRange{3, KernelAbstractions.NDIteration.StaticSize{(4, 4, 16)}, KernelAbstractions.NDIteration.StaticSize{(16, 16, 1)}, Nothing, Nothing}}, OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}, RectilinearGrid{Float64, Oceananigans.Grids.Periodic, Oceananigans.Grids.Periodic, Bounded, Float64, Float64, Float64, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, OffsetArrays.OffsetVector{Float64, StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}}, Nothing}, Val{1}, WENO5{Float64, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, true}, ScalarDiffusivity{Oceananigans.TurbulenceClosures.ExplicitTimeDiscretization, Oceananigans.TurbulenceClosures.ThreeDimensionalFormulation, Float64, NamedTuple{(:b,), Tuple{Float64}}}, BoundaryCondition{Oceananigans.BoundaryConditions.Flux, Nothing}, Buoyancy{BuoyancyTracer, Oceananigans.Grids.ZDirection}, NamedTuple{(:velocities, :tracers), Tuple{NamedTuple{(:u, :v, :w), Tuple{Oceananigans.Fields.ZeroField{Int64, 3}, Oceananigans.Fields.ZeroField{Int64, 3}, Oceananigans.Fields.ZeroField{Int64, 3}}}, NamedTuple{(:b,), Tuple{Oceananigans.Fields.ZeroField{Int64, 3}}}}}, NamedTuple{(:u, :v, :w), Tuple{OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}, OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}, OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}}}, NamedTuple{(:b,), Tuple{OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}}}, Nothing, Oceananigans.Forcings.ContinuousForcing{Center, Center, Center, NamedTuple{(:δᴸ, :R, :s), Tuple{Float64, Float64, Float64}}, typeof(buoyancy_source), Nothing, Tuple{}, Tuple{}}, NamedTuple{(:time, :iteration, :stage), Tuple{Float64, Int64, Int64}}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to overdub)
Stacktrace:
 [1] overdub
   @ ~/Documents/MIT/simulations/oceananigans/MI2_moving_source/moving_source_MI2.jl:112
 [2] overdub
   @ ~/.julia/packages/Oceananigans/rtj1r/src/Forcings/continuous_forcing.jl:124
 [3] overdub
   @ ~/.julia/packages/Oceananigans/rtj1r/src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl:226
 [4] macro expansion
   @ ~/.julia/packages/Oceananigans/rtj1r/src/Models/NonhydrostaticModels/calculate_nonhydrostatic_tendencies.jl:179
 [5] overdub
   @ ~/.julia/packages/KernelAbstractions/3ZHln/src/macros.jl:80
 [6] overdub
   @ ~/.julia/packages/Cassette/34vIw/src/overdub.jl:0

[raphaelouillon]

I am also not using Interpolations.jl as I have read that it is not compatible with CUDA.

[glwagner]

Oh, interesting!

I can't say this will help for sure, but it might make sense to start your implementation using discrete_form=true. This will be a little more work, because you'll have to use node to get x, y, z and t = clock.time but may help with compilation.

there are also some interpolation utilities here:

https://github.com/CliMA/Oceananigans.jl/blob/main/src/Fields/interpolate.jl

that may be helpful

I didn't know that searchsortedfirst and clamp work on the GPU. Is there an implicit assumption about boundary conditions built-in here?

[glwagner]

If I may ask, can you describe the problem you're trying to solve in more general terms? I wonder if we can implement it without interpolation.

[raphaelouillon]

@glwagner I seem to have gotten this to work, but I think it might be highly inefficient. The goal is to model a moving source of buoyancy over complex topography. The topography is interpolated from a bathymetry file of some region of the pacific ocean. This is relatively slowly varying and could easily be represented by some polynomial. The position and intensity of the source are interpolated from a time series from actual field trials that took place in the pacific ocean some time ago. I don't think there is an alternative to interpolation here. This is what I am currently running for the source

@inline function buoyancy_source(x, y, z, t, p)

	R = p.R
    δᴸ = p.δᴸ
	ctime = p.ctime
	cE = p.cE
	cN = p.cN
	mflux = p.mflux
	depth = p.depth
	easting = p.easting
	northing = p.northing

	i = searchsortedfirst(ctime, t)
    i = clamp(i, firstindex(cE), lastindex(cE))
	xc = @inbounds cE[i]
	yc = @inbounds cN[i]
	mdot = @inbounds mflux[i]

	i = searchsortedfirst(easting, xc)
	i = clamp(i, firstindex(easting), lastindex(easting))
	j = searchsortedfirst(northing, yc)
	j = clamp(j, firstindex(northing), lastindex(northing))
	zc = 2.0 + depth[i,j]
	d = sqrt((x - xc)^2 + (y - yc)^2+ (z - zc)^2)

    return - 0.5*mdot* (1 - erf((d - R) / δᴸ))#  sᵦ(t) * 0.5 * (1 - erf((d - R) / δᴸ)) # 0.5 * (1 - erf((d - R) / δᴸ))#-
end


b_forcing = Forcing(buoyancy_source,
                    discrete_form = false,
					parameters = (δᴸ=0.05, R=2.0, s=s,ctime = CuArray(ctime), cE =CuArray(cE), cN = CuArray(cN), mflux = CuArray(mflux), northing = CuArray(northing), easting = CuArray(easting), depth = CuArray(depth)))

[glwagner]

I think you do need to interpolate in time, but you may not need to interpolate in space. Can't you generate a coarse time-series of the source field as a preprocessing step, and then linearly interpolate between the saved time points to produce a smoothly varying buoyancy source in time, on the grid that the simulation is using?

Interpolating in time will be easy because you don't have to fiddle with indices.

[glwagner]

We use a similar method to apply boundary conditions from observed fluxes, for example. Here's a case where we have an annual cycle of monthly-averaged momentum, temperature, and salinity fluxes generated from observations. We thus have 12 arrays for each flux field representing the monthly-averaged spatial distribution of the fluxes, and linearly interpolate in time between the snapshots to generate a smoothly varying time-series of fluxes at every space index (i, j in this case, since its a boundary condition) and time:

https://github.com/CliMA/Oceananigans.jl/blob/main/validation/near_global_lat_lon/cyclic_interpolate_utils.jl

Those utilities also make the interpolation "cyclic", ie repeating in time, which you don't need (so your case may be a little simpler, no need for unsafe_trunc).

[glwagner]

For bathymetry, it's almost always better to use an array rather than a function. You can use Interpolations there, because if we are given a bathymetry function, we will precalculate the bathymetry on the CPU and then later transfer it to the GPU:

Oceananigans.jl/src/ImmersedBoundaries/grid_fitted_immersed_boundaries.jl

Lines 43 to 51 in bc3e6c0

	function ImmersedBoundaryGrid(grid, ib::GridFittedBottom)
	arch = grid.architecture
	bottom_field = Field{Center, Center, Nothing}(grid)
	set!(bottom_field, ib.bottom_height)
	fill_halo_regions!(bottom_field)
	offset_bottom_array = dropdims(bottom_field.data, dims=3)
	new_ib = GridFittedBottom(offset_bottom_array)
	return ImmersedBoundaryGrid(grid, new_ib)
	end