Interpolating fields between grids on a GPU?

[sdbrenner]

Summary and context

Is there a good method for interpolating fields between grids on a GPU? (or a smarter approach than this type of interpolation?)

I am trying to run two-way-coupled simulations with Oceananigans, and the Subzero.jl sea ice discrete element model. The ice model responds to ocean currents and produces surface stress fields that affect the ocean, so those fields must be passed back and forth (also see #3892). I've realized that there is an offset in the grid definitions between the two models---in the sea ice model, the currents/stresses/etc are defined at effectively the 'corners' of the Oceananigans gridcells (see Caltech-OCTO/Subzero.jl#30).

For now, I've tried to interpolate the ocean currents onto the sea ice grid using interpolate from Oceananigans.Fields. When running on CPU, this seems to work, but on GPU it produces errors (the message in the errors seems to include a work-around, but also suggests avoiding it). I've seen that other interpolation functions in Julia (Interpolations.jl) also aren't GPU compatible.

Are there other steps I need to take, or an approach that would be better suited (perhaps somehow avoiding the interpolation issue)? Additionally, are there any suggestions for using built-in functions for interpolating stresses (τx,τy) back onto the Oceananigans grid?

Note, this seems to be related to #2565, but it's not clear if there was a resolution to that question (and one of the responses includes a broken link)

Minimal(-ish) working example

This example reproduces the error:

using Oceananigans
using Oceananigans.Fields
using CUDA

x = y = z = (0, 1)
nxy = 10
grid = RectilinearGrid(GPU(); size=(nxy, nxy, 1), x, y, z)
model = NonhydrostaticModel(; grid )

xi = yi = range(x[1],x[2],nxy+1) # equivalent Subzero grid
ui = [interpolate((x,y,0.5),model.velocities.u) for x in xi, y in yi] # Ocean u velocity at Subzero grid locations

Error

ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/GMsgk/src/GPUArraysCore.jl:155
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/GMsgk/src/GPUArraysCore.jl:128
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/GMsgk/src/GPUArraysCore.jl:116
  [5] getindex
    @ ~/.julia/packages/GPUArrays/qt4ax/src/host/indexing.jl:50 [inlined]
  [6] scalar_getindex
    @ ~/.julia/packages/GPUArrays/qt4ax/src/host/indexing.jl:36 [inlined]
  [7] _getindex
    @ ~/.julia/packages/GPUArrays/qt4ax/src/host/indexing.jl:19 [inlined]
  [8] getindex
    @ ~/.julia/packages/GPUArrays/qt4ax/src/host/indexing.jl:17 [inlined]
  [9] getindex
    @ ~/.julia/packages/OffsetArrays/hwmnB/src/OffsetArrays.jl:424 [inlined]
 [10] getindex
    @ ~/.julia/packages/Oceananigans/s1DfC/src/Fields/field.jl:401 [inlined]
 [11] _interpolate
    @ ~/.julia/packages/Oceananigans/s1DfC/src/Fields/interpolate.jl:295 [inlined]
 [12] interpolate(at_node::Tuple{…}, from_field::Field{…}, from_loc::Tuple{…}, from_grid::RectilinearGrid{…})
    @ Oceananigans.Fields ~/.julia/packages/Oceananigans/s1DfC/src/Fields/interpolate.jl:245
 [13] interpolate(to_node::Tuple{…}, from_field::Field{…})
    @ Oceananigans.Fields ~/.julia/packages/Oceananigans/s1DfC/src/Fields/interpolate.jl:312
 [14] (::var"#1#2")(::Tuple{Float64, Float64})
    @ Main ./none:0
 [15] iterate
    @ ./generator.jl:47 [inlined]
 [16] collect(itr::Base.Generator{Base.Iterators.ProductIterator{Tuple{StepRangeLen{…}, StepRangeLen{…}}}, var"#1#2"})
    @ Base ./array.jl:834
 [17] top-level scope
    @ REPL[12]:1
Some type information was truncated. Use `show(err)` to see complete types.

[glwagner]

This line:

ui = [interpolate((x,y,0.5),model.velocities.u) for x in xi, y in yi] # Ocean u velocity at Subzero grid locations

creates an array on the CPU. I suspect that's not what you want.

If you have a CuArray of locations corresponding to the locations of subzero floes, you probably have to write a GPU kernel in order to interpolate the velocity field to those locations efficiently. There are many ways to do this. You can maybe try to figure out how to use broadcasting:

@inline subzero_interp(x, y, u) = interpolate((x, y, zero(x)), u)
ui = subzero_interp.(xi, yi, Ref(u))

That works if xi, yi, ui are all arrays of the same kind. Another way is to use https://github.com/JuliaGPU/KernelAbstractions.jl.

PS don't forget spaces after commas and also be careful about hard coding 0.5, which is Float64.

A good example of using KernelAbstractions can be found in Oceananigans' own interpolate functionality:

Oceananigans.jl/src/Fields/interpolate.jl

Lines 328 to 372 in 37d0e29

	@kernel function _interpolate!(to_field, to_grid, to_location,
	from_field, from_grid, from_location)
	i, j, k = @index(Global, NTuple)
	to_node = _node(i, j, k, to_grid, to_location...)
	to_node = flatten_node(to_node...)
	@inbounds to_field[i, j, k] = interpolate(to_node, from_field, from_location, from_grid)
	end
	"""
	interpolate!(to_field::Field, from_field::AbstractField)
	Interpolate `from_field` `to_field` and then fill the halo regions of `to_field`.
	"""
	function interpolate!(to_field::Field, from_field::AbstractField)
	to_grid = to_field.grid
	from_grid = from_field.grid
	to_arch = architecture(to_field)
	from_arch = architecture(from_field)
	# In case architectures are `Distributed` we
	# verify that the fields are on the same child architecture
	to_arch = child_architecture(to_arch)
	from_arch = child_architecture(from_arch)
	if !isnothing(from_arch) && to_arch != from_arch
	msg = "Cannot interpolate! because from_field is on $from_arch while to_field is on $to_arch."
	throw(ArgumentError(msg))
	end
	# Make locations
	from_location = Tuple(L() for L in location(from_field))
	to_location = Tuple(L() for L in location(to_field))
	launch!(to_arch, to_grid, size(to_field),
	_interpolate!, to_field, to_grid, to_location,
	from_field, from_grid, from_location)
	fill_halo_regions!(to_field)
	return to_field
	end

[sdbrenner]

Thanks, as always, for your patience.
I should have also included for clarification that while I'll plan to run Oceananigans on a GPU, Subzero doesn't include GPU support, so the ui field that I'll pass to the model should not be a CuArray.

In any case, I still receive the same error when trying the inline function/broadcasting:

using Oceananigans
using Oceananigans.Fields
using CUDA

x = y = z = (0, 1)
nxy = 10
grid = RectilinearGrid(GPU(); size=(nxy, nxy, 1), x, y, z)
model = NonhydrostaticModel(; grid )

xi = yi = range(x[1],x[2],nxy+1)
xig = repeat(xi , 1, nxy+1 ) # grid
yig = repeat(yi', nxy+1, 1 ) # grid
zo = znodes(grid,Center()) 

@inline subzero_interp(x, y, u) = interpolate((x, y, zo[end]), u)

u = model.velocities.u 
ui = zeros(size(xig)) 
ui .= subzero_interp.(xig, yig, Ref(u))

(if I try instead to make xig, yig, and ui all CuArray's, I get a different, much longer, set of errors; lots of "unsupported dynamic function invocation")

I'll try to understand more about KernalAbstractions and see if I can solve my issues that way.

[glwagner]

You have to interpolate on the GPU, and then transfer to the CPU. This means that xi, yi, and ui must all be on the GPU when you interpolate. If you need the data on the CPU, then you will have to transfer it from GPU to CPU after interpolation, eg:

ui_cpu = zeros(size(xig))
copyto!(ui_cpu, ui_gpu)

to tranfser the data between devices.

In terms of fixing the dynamic function invocation, this means that broadcasting fails. I think you'll have to write a kernel. It will look something like this

const f = Face()
const c = Center()
const uo_loc = (f, c, c)

@kernel function interpolate_xy!(ui, xi, yi, uo, uo_grid)
      p = @index(Global, Linear) # particle index
      to_node = (xi[p], yi[p], zero(ocean_grid)) # (x, y, z=0)
      @inbounds ui[p] = interpolate(to_node, uo, uo_loc, uo_grid)
end

check out the source code I linked to for more details.