Abstraction for using "custom" kernels to compute fields (plus an example)?

[glwagner]

In a perfect world one would be able to use AbstractOperations for all output needs. But alas, there is #1241. The discussion on #1234 makes clear that this will continue to be a serious issue. There may also be output that cannot be expressed as an abstract operation.

One work around for advanced users (which may be nice to have anyways) is to design an abstraction for KernelComputedFields that are computed based on a user-supplied kernel. I think usage could be something like (using ViscousDissipation as an example)

model = IncompressibleModel(...)

@kernel function compute_viscous_dissipation!(ϵ, grid, clock, νₑ, u, v, w)
    i, j, k = @index(Global, NTuple)

    Σˣˣ = ∂xᶜᵃᵃ(i, j, k, grid, u)
    Σʸʸ = ∂yᵃᶜᵃ(i, j, k, grid, v)
    Σᶻᶻ = ∂zᵃᵃᶜ(i, j, k, grid, w)

    Σˣʸ = (ℑxyᶜᶜᵃ(i, j, k, grid, ∂yᵃᶠᵃ, u) + ℑxyᶜᶜᵃ(i, j, k, grid, ∂xᶠᵃᵃ, v)) / 2
    Σˣᶻ = (ℑxzᶜᵃᶜ(i, j, k, grid, ∂zᵃᵃᶠ, u) + ℑxzᶜᵃᶜ(i, j, k, grid, ∂xᶠᵃᵃ, w)) / 2
    Σʸᶻ = (ℑyzᵃᶜᶜ(i, j, k, grid, ∂zᵃᵃᶠ, v) + ℑyzᵃᶜᶜ(i, j, k, grid, ∂yᵃᶠᵃ, w)) / 2

    @inbounds ϵ[i, j, k] = νₑ[i, j, k] * 2 * (Σˣˣ^2 + Σʸʸ^2 + Σᶻᶻ^2 + Σˣʸ^2 + Σˣᶻ^2 + Σʸᶻ^2)
end

νₑ = model.diffusivities.νₑ
u, v, w = model.velocities

viscous_dissipation = KernelComputedField(Cell, Cell, Cell, model, compute_viscous_dissipation;
                                          field_dependencies=(νₑ, u, v, w))

In this example viscous_dissipation is then a bonafide field with a fully-fledged compute! function that can be used in output, time-averaging, in an AveragedField, etc.

Users will need to write a kernel with KernelAbstractions language (which we'll have to document and provide examples for), but they won't have to write all the other boilerplate that goes into defining the struct and writing a compute! function (see again ViscousDissipation to get an idea of what this boiler plate looks like).

The kernel can have field_dependencies, which are passed into the KernelComputedField constructor. It also seems nice to pass clock.

[glwagner]

cc @tomchor

[tomchor]

@glwagner I modified the docstring with a small example that I think illustrates the patterns relatively in this commit. I'm assuming that's what you meant by "example", right?

I've tested it already in some simple examples and it produces correct results. Is there anything else left besides testing this with an actual complex calculation that can only be done via KernelComputedField on a GPU?

[glwagner]

@tomchor you may have mean t to make this comment on PR #1293!