Possible elegant solution for compiling kernels with fields as arguments

[glwagner]

One solution for compiling kernels with fields as arguments is to define

Adapt.adapt_structure(to, field::AbstractField) = data(field)

and

Adapt.adapt_structure(to, fields::NamedTuple{S, NTuple{N, <:AbstractField}}) where {S, N} =
    datatuple(fields)

This approach will automatically unwrap fields when they are passed to GPU kernels (I think).

The disadvantage of this approach is that the code is a little bit harder to interpret, because if you miss the definition of this function, you might be confused why field.boundary_conditions and field.grid were not accessible from inside GPU kernels.

On the other hand, this only affects the lowest-level kernels, and GPU programmers probably know that an adapt_structure method must be defined somewhere for an exotic object like Field to be passed into a kernel. If we document what we are doing clearly, we may solve this problem.

Doing this means we would no longer have to unwrap fields manually prior to passing them to GPU kernels.

Are there any other issues that I'm not seeing?

[ali-ramadhan]

Hmmm, do we know why we can't adapt the full Field to work inside GPU kernels? I think you said you've tried it but ran into some issues a while back?

The definition is

struct Field{X, Y, Z, A, G, B} <: AbstractField{X, Y, Z, A, G}
                   data :: A
                   grid :: G
    boundary_conditions :: B

where the grid has been adapted, data is usually an offset array which has also been adapted, and I thought individual boundary conditions have been adapted as well so feels like it should be easy to adapt the full Field.

Note: boundary_conditions is a named tuple of FieldBoundaryConditions so maybe we just need to adapt FieldBoundaryConditions, CoordinateBoundaryConditions? Hmmm but they're really named tuples so maybe they're already adapted.

If we can't adapt Field then yeah the adapt rules you've suggested sound pretty good and would simplify the time stepping a lot. The kernels already deal with data tuples so not much would even have to change I think.

If we can adapt Field though, then we should get all the benefits of simpler time stepping code without the confusion of a non-vanilla adapt rule.

X-Ref: #298

[glwagner]

I'm not sure why my attempt failed. It's worth documenting an attempt.

[sandreza]

If you are looking to pass structs with CuArray members into kernels, I think that you might be looking for something like https://github.com/sandreza/Learning/blob/master/nla/gmres.jl#L50
(thanks to @vchuravy for showing me how to use adapt!)
It should work recursively so that composite structs will function as desired as long as each individual struct is adapted

[glwagner]

@sandreza thanks! We utilize this functionality for many of our objects, eg:

https://github.com/climate-machine/Oceananigans.jl/blob/9ef95e7bef2db1dc9ac04af78664418b0caaf99b/src/AbstractOperations/binary_operations.jl#L144

For some reason, during an undocumented attempt to apply this logic to fields back in October, we were unsuccessful to get code to work on the GPU.

Back then, the field consisted of an OffsetArray wrapped around a CuArray, and a grid. Both of those objects can be adapted to GPU kernels, so it should have worked, I think.

So I'm not 100% sure why our attempt to use adapt_structure failed for fields, while working for other objects. Any insight appreciated...

[vchuravy]

Any insight appreciated...

Without knowing what you tried to do, not really.

[glwagner]

Yes, that's not surprising. I regret not being more careful.

[glwagner]

I've documented an attempt to adapt Field to the GPU at #746 .