Oceananigans.jl: Convert Grid and Field structs to isbitstype

[ali-ramadhan]

@peterahrens suggested removing ModelMetadata from Field structs and make the Grid a parameter, i.e. Grid{T,G} then it should be isbitstype. It's not isbits right now because T<:AbstractArray.

Then we should be able to use some abstractions in the CPU/GPU element-wise kernels as well as multiple dispatch and won't need to have δx_c2f, δx_f2c, δx_e2f, δx_f2e, etc.

Some goals for guidance:

• [ ] Prototype grid and field types that are isbitstype.
• [ ] Test that they do work with GPUifyLoops on the GPU. For this I will extend the example from https://github.com/vchuravy/GPUifyLoops.jl/pull/18
• [ ] Benchmark the simple example to see things slow down. I don't expect to see much of a difference.
• [ ] Refactor the operators and time stepping loop to use the new abstractions!
• [ ] Add tests to make sure any structure that may be passed to a CUDA kernel isbitstype. Just construct a bunch of grids, fields, and forcing functions and test that they are indeed isbitstype.
• [ ] Benchmark the model with isbitstype abstractions to make sure performance hasn't degraded.

[glwagner]

Just noticed the title! Duh.

Anyways:

abstract type Grid end 

struct RegularCartesianGrid{T<:AbstractFloat,TxC} <: Grid
    dim::Int
    # Number of grid points in (x,y,z).
    Nx::Int
    Ny::Int
    Nz::Int
    # Domain size [m].
    Lx::T
    Ly::T
    Lz::T
    # Grid spacing [m].
    Δx::T
    Δy::T
    Δz::T
    # Cell face areas [m²].
    Ax::T
    Ay::T
    Az::T
    # Volume of a cell [m³].
    V::T
    # Ranges of coordinates at the centers of the cells.
    xC::TxC
    yC::TxC
    zC::TxC
    # Ranges of coordinates at the faces of the cells.
    xF::TxC
    yF::TxC
    zF::TxC
end

function RegularCartesianGrid(T, N, L)
    @assert length(N) == 3 && length(L) == 3 "N, L must have all three dimensions to specify which dimensions are used."
    @assert all(L .> 0) "Domain lengths must be nonzero and positive!"

    # Count the number of dimensions with 1 grid point, i.e. the number of flat
    # dimensions, and use it to determine the dimension of the model.
    num_flat_dims = count(i->(i==1), N)
    dim = 3 - num_flat_dims
    @assert 1 <= dim <= 3 "Only 1D, 2D, and 3D grids are supported right now."

    Nx, Ny, Nz = N 
    Lx, Ly, Lz = L 

    dim == 2 && @assert Nx == 1 || Ny == 1 || Nz == 1 "For 2D grids, Nx, Ny, or Nz must be 1."
    dim == 3 && @assert Nx != 1 && Ny != 1 && Nz != 1 "For 3D grid, cannot have dimensions of size 1."

    Δx = Lx / Nx
    Δy = Ly / Ny
    Δz = Lz / Nz

    Ax = Δy*Δz
    Ay = Δx*Δz
    Az = Δx*Δy

    V = Δx*Δy*Δz

    xC = Δx/2:Δx:Lx
    yC = Δy/2:Δy:Ly
    zC = -Δz/2:-Δz:-Lz

    xF = 0:Δx:Lx
    yF = 0:Δy:Ly
    zF = 0:-Δz:-Lz

    RegularCartesianGrid{T,typeof(xC)}(dim, Nx, Ny, Nz, Lx, Ly, Lz, Δx, Δy, Δz, Ax, Ay, 
                                       Az, V, xC, yC, zC, xF, yF, zF) 
end

RegularCartesianGrid(metadata::ModelMetadata, N, L) = RegularCartesianGrid(metadata.float_type, N, L)

then

<somehow include above code>

julia> g = RegularCartesianGrid(Float64, (3, 3, 3), (1.0, 1.0, 1.0))
...
julia> isbitstype(typeof(g))
true

[vchuravy]

Was about to note the same thing: https://github.com/ali-ramadhan/Oceananigans.jl/blob/f7ba1ce353eaa6abc2b3c9381fa22fe53ba27fce/src/grids.jl#L33-L46 need to be typed to be isbits.

[glwagner]

Perhaps move this to v0.6?

[ali-ramadhan]

Agreed.

Useful and related comment

I was thinking of doing some prototyping and benchmarking in a sandbox by building off the example in my PR vchuravy/GPUifyLoops.jl#18.

The PR contains an example that can be extended to rely on a Grid struct, multiple FaceFields and CellField. So I'll prototype grids and fields that are isbitstype (you already helped by doing this for a grid in #59 (comment)) and test to see if they work on the GPU with GPUifyLoops.jl. If they do work and performance isn't degraded then I'll rewrite the operators to use grid and field structs.

You probably know how to do this better than me, but might be good if I rewrite the operators as they's still undocumented and do some _slightly convoluted_ stuff to avoid having to store intermediate calculations.

Right now I'm focusing on system tests and benchmarks but once @christophernhill @jm-c and I get closer to implementing the variable _Δz_ grid #47 I will work on this.

_Originally posted by @ali-ramadhan in https://github.com/climate-machine/Oceananigans.jl/issues/115#issuecomment-470782067_

[glwagner]

@jm-c and I talked today and we have an additional point to make (this is related to #115 as well).

Notice first that the operators δz_f2c and δz_e2f are identical:

@inline function δz_f2c(f, Nz, i, j, k)
    if k == Nz
        @inbounds return f[i, j, k]
    else
        @inbounds return f[i, j, k] - f[i, j, k+1]
    end
end

@inline function δz_e2f(f, Nz, i, j, k)
    if k == Nz
        @inbounds return f[i, j, k]
    else
        @inbounds return f[i, j, k] - f[i, j, k+1]
    end
end

Our solution:

abstract type Location end
struct Center end
struct Interface end

struct Field{Lx<:Location, Ly<:Location, Lz<:Location, A, G}
    data::A
    grid::G
end

we then need only two δz functions that dispatch on Lz; for example:

δz(f::Field{Lx, Ly, Lz}, i, j, k) where {Lx, Ly, Lz<:Center} = ...
δz(f::Field{Lx, Ly, Lz}, i, j, k) where {Lx, Ly, Lz<:Interface} = ...

[ali-ramadhan]

Dispatching on the location as a parameterized type looks neat, would definitely clean things up! Thanks for coding up the example.

I see how this would clean things up for δz_f2c and δz_e2f but the δz operator also needs to know onto which location to interpolate, so how would this work with δz_f2c and δz_f2e?

@inline function δz_f2c(f, Nz, i, j, k)
    if k == Nz
        @inbounds return f[i, j, k]
    else
        @inbounds return f[i, j, k] - f[i, j, k+1]
    end
end

@inline function δz_f2e(f, Nz, i, j, k)
    if k == 1
        return 0
    else
        @inbounds return f[i, j, k-1] - f[i, j, k]
    end
end

To me it looks like they both fit

δz(f::Field{Lx, Ly, Lz}, i, j, k) where {Lx, Ly, Lz<:Interface} = ...

so I'm not sure how dispatch between the two.

[glwagner]

“face to edge” delta z must refer to an operation on an x-interface or
y-interface — not a z-interface to an edge?

On Sat, Mar 9, 2019 at 8:38 AM Ali Ramadhan [e-mail]
wrote:

Dispatching on the location as a parameterized type looks neat, would
definitely clean things up! Thanks for coding up the example.

I see how this would clean things up for δz_f2c and δz_e2f but the δz
operator also needs to know onto which location to interpolate, so how
would this work with δz_f2c and δz_f2e?

@inline function δz_f2c(f, Nz, i, j, k)

if k == Nz

    @inbounds return f[i, j, k]

else

    @inbounds return f[i, j, k] - f[i, j, k+1]

end

end

@inline function δz_f2e(f, Nz, i, j, k)

if k == 1

    return 0

else

    @inbounds return f[i, j, k-1] - f[i, j, k]

end

end

To me it looks like they both fit

δz(f::Field{Lx, Ly, Lz}, i, j, k) where {Lx, Ly, Lz<:Interface} = ...

so I'm not sure how dispatch between the two.

—
You are receiving this because you were assigned.
Reply to this email directly, view it on GitHub
https://github.com/climate-machine/Oceananigans.jl/issues/59#issuecomment-471177715,
or mute the thread
https://github.com/notifications/unsubscribe-auth/AOkIBkzcm1N5EI8VXyk1UgseDS6muPk9ks5vU7lYgaJpZM4bEKew
.

[ali-ramadhan]

Ah yes I remember @jm-c's concern now! Yes, just "face" and "edge" are ambiguous.

So your proposed solution will make it explicit exactly which interface the field is being interpolated from and to, which will be nice.

[glwagner]

An “edge” is the case in which 2 of Lx, Ly, or Lz are “interface”. We just
don’t need special notation. This method of specifying locations in a cube
is simpler, but also more general.

On Sat, Mar 9, 2019 at 8:52 AM Ali Ramadhan [e-mail]
wrote:

Ah yes I remember @jm-c https://github.com/jm-c's concern now! Yes,
just "face" and "edge" are ambiguous.

So your proposed solution will make it explicit exactly which interface
the field is being interpolated from and to, which will be nice.

—
You are receiving this because you were assigned.
Reply to this email directly, view it on GitHub
https://github.com/climate-machine/Oceananigans.jl/issues/59#issuecomment-471178907,
or mute the thread
https://github.com/notifications/unsubscribe-auth/AOkIBlzM_EBoVLsfUYqw-V_v_Pe2admmks5vU7yqgaJpZM4bEKew
.

[ali-ramadhan]

I'm finally starting to work on this issue now. I should be able to refactor all the code without changing any of the numerics. Resolving this issue will also resolve #60.

I have moved the slightly separate issue of just using cell centers and cell interfaces (and explicitly stating which interface) to #146 as that will change the operators and numerics, and may be trickier to get right.

I've added some goals to the original post.

[ali-ramadhan]

Hit a slight roadblock as the Field structs are not isbitstype:

struct CellField{A<:AbstractArray,G<:Grid} <: Field
    data::A
    grid::G
end

since Arrays and CuArrays are not isbits.

However, @vchuravy points out that a CuArray is not isbits only because it contains a finalize function (essentially a destructor) which makes it not isbits on the host. Apparently CUDAnative converts the CuArray to some sort of CuDeviceArray using Adapt.jl which is isbits and can be used as a device argument.

So we'll have to figure out how to _adapt_ our Field structs to be isbits on the GPU device using Adapt.jl. It shouldn't be too hard as the struct is simple.

I should also check out the cudaconvert(args...) function in CUDAnative.

[glwagner]

Oh no!

We can still use the Grid structs though, right?

[ali-ramadhan]

Well, road block might not be the right word as I'm sure we can figure it out. So we can already use the Grid struct but would be nice to also be able to use the Field struct too. But yeah, using the Grid struct would clean up most of the clutter.