Saving diagnostics to NetCDFOutputWriter via KernelFunctionOperation

[rosecormier]

I (with @francispoulin) am running simulations whose initial conditions consist of a basic-state velocity plus a small perturbation velocity. I'm interested in tracking the kinetic energy associated with the perturbation velocity (i.e., difference between total velocity and basic-state velocity). To do this on a GPU, I am trying to build a KernelFunctionOperation analogous to those used in Oceanostics, but I'm having difficulty setting it up and not sure exactly how to pass it to the output writer.

When I run the following script, I get a long error:

ERROR: LoadError: InvalidIRError: compiling MethodInstance for Oceananigans.AbstractOperations.gpu__compute!(::KernelAbstractions.CompilerMetadata{Oceananigans.Utils.OffsetStaticSize{(1:1200, 1:1200, 1:12)}, KernelAbstractions.NDIteration.DynamicCheck, Nothing, Nothing, KernelAbstractions.NDIteration.NDRange{3, KernelAbstractions.NDIteration.StaticSize{(75, 75, 12)}, KernelAbstractions.NDIteration.StaticSize{(16, 16, 1)}, Tuple{Int64, Int64, Int64}, Oceananigans.Utils.KernelOffsets{Tuple{Int64, Int64, Int64}}}}, ::OffsetArrays.OffsetArray{Float64, 3, CuDeviceArray{Float64, 3, 1}}, ::KernelFunctionOperation{Center, Center, Center, RectilinearGrid{Float64, Bounded, Bounded, 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}, Float64, typeof(pKE_ccc!), 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}}}}) resulted in invalid LLVM IR
Reason: unsupported dynamic function invocation (call to pKE_ccc!)

I'm not familiar with this error, and not sure whether I've set up my KernelFunctionOperation incorrectly? I've read that I can pass KernelFunctionOperations to output writers, but I'm also not sure whether I'm passing it to the writer correctly. Any direction would be appreciated!

using Adapt, CUDA, NCDatasets, Oceananigans, Oceanostics, Printf, Random
using Dates: canonicalize, format, now
using KernelAbstractions: @index, @kernel
using Oceananigans.Architectures
using Oceananigans.BoundaryConditions
using Oceananigans.Coriolis
using Oceananigans.Fields
using Oceananigans.Units
using Oceananigans.Utils
using OffsetArrays: no_offset_view

#Numbers of gridpoints
const Nx = 1200
const Ny = 1200
const Nz = 12

#Lengths of axes
const Lx = 2.5e3 * kilometer
const Ly = 2.5e3 * kilometer
const Lz = 1 * kilometer

#Latitude (deg. N)
const lat = 74.0

#f-plane and Coriolis frequency
fPlane  = FPlane(latitude = lat)
const f = fPlane.f

#Characteristic scales for basic state
const σr = 250 * kilometer
const U   = 1.5e-1 * (meter/second)
const N²₀ = 3e-3 * (second^(-2))

#Time-stepping parameters
const Δt      = 10 * second
const tf      = 3 * hour
const Δt_save = 1 * hour

#Max. relative magnitude of initial u-perturbations
const max_u′ = 1e-8

grid = RectilinearGrid(GPU(),
		       topology = (Bounded, Bounded, Bounded),
                       size = (Nx, Ny, Nz), 
                       x = (-Lx/2, Lx/2), 
                       y = (-Ly/2, Ly/2), 
                       z = (-Lz, 0.0),
		       halo = (3, 3, 3))

b̄ = (x, y, z) -> N²₀ * z
ū = (x, y, z) -> (sqrt(2)*U*y/σr) * exp((1/2) - (x^2 + y^2)/(σr^2))
v̄ = (x, y, z) -> -(sqrt(2)*U*x/σr) * exp((1/2) - (x^2 + y^2)/(σr^2))

b̄z_top = (x, y, t) -> N²₀
b̄z_bot = (x, y, t) -> N²₀
b̄_BCs  = FieldBoundaryConditions(top    = GradientBoundaryCondition(b̄z_top),
         			 bottom = GradientBoundaryCondition(b̄z_bot))

model = NonhydrostaticModel(; 
                            grid = grid, 
                            timestepper = :RungeKutta3,
                            advection = UpwindBiasedFifthOrder(), 
                            coriolis = fPlane,
                            tracers = (:b),
                            buoyancy = BuoyancyTracer(),
			    boundary_conditions = (; b = b̄_BCs,))

set!(model, u = ū, v = v̄)
fill_halo_regions!(model.velocities, model.tracers.b)
set!(model, b = b̄)

datetimenow   = format(now(), "yymmdd-HHMMSS")

#Store basic-state velocities
const Ux_Field = XFaceField(model.grid; data = model.velocities.u.data)
const Uy_Field = YFaceField(model.grid; data = model.velocities.v.data)
const Uz_Field = ZFaceField(model.grid; data = model.velocities.w.data)
const Ux = adapt(CuArray, Ux_Field)
const Uy = adapt(CuArray, Uy_Field)
const Uz = adapt(CuArray, Uz_Field)

@kernel function pKE_ccc!(pKE, grid, u, v, w)
   i, j, k = @index(Global, NTuple)
   @inbounds pKE[i, j, k] = (
                              ℑxᶜᵃᵃ(i, j, k, grid, ψ′², u, Ux) +
                              ℑyᵃᶜᵃ(i, j, k, grid, ψ′², v, Uy) +
                              ℑzᵃᵃᶜ(i, j, k, grid, ψ′², w, Uz)
                             ) / 2
end

pKE_op = KernelFunctionOperation{Center, Center, Center}(pKE_ccc!, 
							 grid, 
							 model.velocities.u, 
							 model.velocities.v, 
							 model.velocities.w)

#Perturb velocity components to trigger instability
@inline u_perturbed(x, y, z) = (ū(x, y, z) + (2 * (rand() - 0.5)) * (max_u′ / sqrt(2)))
@inline v_perturbed(x, y, z) = (v̄(x, y, z) + (2 * (rand() - 0.5)) * (max_u′ / sqrt(2)))
set!(model, u = u_perturbed, v = v_perturbed)

simulation = Simulation(model, Δt = Δt, stop_time = tf)

energy_diagnostics = (; pKE = pKE_op)
energyfilepath = joinpath("./Output", "energetics_$(datetimenow).nc")
mkpath(dirname(energyfilepath)) #Make path if nonexistent

energy_writer = NetCDFOutputWriter(model, 
				   energy_diagnostics,
                                   with_halos = true,
				   filename = energyfilepath,
				   schedule = TimeInterval(Δt_save),
				   file_splitting = FileSizeLimit(30GiB))

simulation.output_writers[:energy_writer] = energy_writer

run!(simulation)

[glwagner]

There are a few things to fix here:

1. you do not want to build Ux with the data from model.velocities.u. This means that the field is identical to u, which is not what you want. You want Ux to have it's own data; which is achieved with Ux = XFaceField(model.grid).

2. Instead of declaring the background fields as constants here:

#Store basic-state velocities
const Ux_Field = XFaceField(model.grid; data = model.velocities.u.data)
const Uy_Field = YFaceField(model.grid; data = model.velocities.v.data)
const Uz_Field = ZFaceField(model.grid; data = model.velocities.w.data)
const Ux = adapt(CuArray, Ux_Field)
const Uy = adapt(CuArray, Uy_Field)
const Uz = adapt(CuArray, Uz_Field)

you should pass the fields into the kernel, like this:

@kernel function pKE_ccc!(pKE, grid, u, v, w, Ux, Uy, Uz)
   i, j, k = @index(Global, NTuple)
   @inbounds pKE[i, j, k] = (
                              ℑxᶜᵃᵃ(i, j, k, grid, ψ′², u, Ux) +
                              ℑyᵃᶜᵃ(i, j, k, grid, ψ′², v, Uy) +
                              ℑzᵃᵃᶜ(i, j, k, grid, ψ′², w, Uz)
                             ) / 2
end

pKE_op = KernelFunctionOperation{Center, Center, Center}(pKE_ccc!, grid, model.velocities...., Ux, Uy, Uz)

Note that b = (x, y, z) -> N²₀ * z is the same as writing b(x, y, z) = N²₀ * z. The former defines an anonymous function and then attaches b to it, which is a little convoluted. The latter is more clear.

[rosecormier]

Thanks for your quick response! And apologies for the convoluted definitions of b̄, etc. -- I copied those from their definitions within a function elsewhere.

I do want Ux to have the same data as u has after the first call to set!. I thought that by declaring Ux as a constant, it would be unaffected by subsequent changes to u, but am I misunderstanding that?

[glwagner]

No need to apologize about anything!

Right, the const is not achieving that. Probably the best solution is to call set!(Ux, model.velocities.u) after you set the initial condition.

Another option is to use set!(Ux, ū) to fill in the values of Ux with the function ū. This has a possible pitfall --- when the initial condition is set using set!(model, ...), there is a step where the divergent part of u is projected out. Therefore using set!(Ux, model.velocities.u) could produce a slightly different (and possibly more accurate) result than set!(Ux, ū).

But you don't have to take my word for that. You can test it pretty easily.

[rosecormier]

Thanks, I will test both options once I get this running! As of now, both approaches are still giving me the same InvalidIRError, citing unsupported dynamic function invocation (call to pKE_ccc!) as the issue. My naive understanding of this error is that it can be resolved by specifying argument types in the function definition, but I've tried specifying types in my kernel function pKE_ccc! and am getting the same error message still.

[glwagner]

My naive understanding of this error is that it can be resolved by specifying argument types in the function definition

Ah no, this will not solve your problem. More generally, specifying types is only meaningful for multiple dispatch.

The issue is that you are trying to call a function that cannot be compiled on the GPU. This is often due to the presence of global variables inside your function, whose type cannot be inferred. In your original code, there was another issue --- you were referencing CuArray as global variables. While CuArray does represent data on the GPU, it cannot be used directly inside a kernel. All arrays that need to be involved in kernel functions should be passed via parameters.

Another reason you might see this error is because the function does not compile at all. So, the very first thing to do is to try to run this on the CPU. If your code works on CPU but not on GPU, then you have narrowed down the problem a bit and you can start reducing your function bit by bit until it compiles and you have identified the issue.