How do I run computations on data in a checkpoint file on a GPU?

[kenflat2]

Hi all,

I have a series of checkpoints from my simulation and I would like to perform a large calculation using them (specifically, I want to compute a buoyancy variance budget). My computation requires derivatives, cumulative integrals and averages over time and hence is quite expensive. When I load my checkpoint files, by default the fields are over a grid on the CPU. I would like to change this to the GPU in hopes that parallelization will accelerate the computation, but unlike loading a FieldTimeSeries, there doesn't appear to be a way to directly load a checkpoint onto a GPU grid.

Here is the code I use to load checkpoints:

function checkpoint_fields(file)
    file = jldopen(file)
    grid = file["HydrostaticFreeSurfaceModel"]["grid"]

    u = XFaceField(grid)
    v = YFaceField(grid)
    T = CenterField(grid)
    
    Hx, Hy, Hz = halo_size(grid)
    for (var, name) in zip((u, v, T), ("u", "v", "T"))
        set!(var, file["HydrostaticFreeSurfaceModel"][name * "/data"][Hx+1:end-Hx, Hy+1:end-Hy, Hz+1:end-Hz])
    end

    return (; u, v, T)
end

Ideally, I would like to cast the grid to be on the GPU then load each field onto it. One solution is to recreate the grid exactly as it was defined in my simulation while specifying the GPU() architecture and replace grid in the above function with that. But this seems cumbersome because we have direct access to the grid through file["HydrostaticFreeSurfaceModel"]["grid"].

Any help would be appreciated, thanks!

[glwagner]

hmm possibly you want something like

cpu_grid = file["HydrostaticFreeSurfaceModel"]["grid"]
grid = Oceananigans.Architectures.on_architecture(GPU(), cpu_grid)

let me know if that works