Zero result when using pressure derivatives in KernelComputedFields

[tomchor]

I'm encountering a weird behavior that I can't really understand. I'm trying to calculate the pressure redistribution terms for an LES using a KernelComputedField. Everything else I use a KernelComputedField in that script works. Since this is the only diagnostic that uses the pressure field I'm thinking that's the source of the problem, but I really can't understand why.

Below is a MWE:

using Oceananigans
using Oceananigans.Utils
using Oceananigans.Fields

Lx = 150; Ly = 6000; Lz = 80
topology = (Periodic, Bounded, Bounded)
grid = RegularRectilinearGrid(size=(1, 512, 8), x=(0, Lx), y=(0, Ly), z=(-Lz, 0),
                              topology=(Periodic, Bounded, Bounded))


model = IncompressibleModel(architecture = CPU(),
                            grid = grid, 
                            )
                            
w_ic(x, y, z) = 0.01*y
v_ic(x, y, z) = 0.01*x

set!(model, w=w_ic, v=v_ic)


import Oceananigans.Fields: ComputedField, KernelComputedField
using Oceananigans.AbstractOperations: @at, ∂x, ∂y, ∂z
using Oceananigans.Grids: Center, Face 
using KernelAbstractions: @index, @kernel
using Oceananigans.Operators

u, v, w = model.velocities
p = sum(model.pressures)


@kernel function pressure_distribution_z_ccc!(dwpdz_ρ, grid, w, p, ρ₀)
    i, j, k = @index(Global, NTuple)
    wp = ℑzᵃᵃᶠ(i, j, k, grid, p) * w[i, j, k] # C, C, C  → C, C, F
    @inbounds dwpdz_ρ[i, j, k] = (1/ρ₀) * ∂zᵃᵃᶜ(i, j, k, grid, wp) # C, C, F  → C, C, C
end 



dwpdz = KernelComputedField(Center, Center, Center, pressure_distribution_z_ccc!, model;
                            field_dependencies=(w, p), 
                            parameters=1024)
                            
dwpdz_2 = ComputedField((1/1024) * ∂z(w*p))

compute!(dwpdz)
compute!(dwpdz_2)

After running this I get that the result from dwpdz_2 is correct (which is a ComputedField), but the result from dwpdz (a KernelComputedField) is not. Output from the REPL:

julia> interior(dwpdz_2)
1×512×8 view(OffsetArray(::Array{Float64,3}, 0:2, 0:513, 0:9), 1:1, 1:512, 1:8) with eltype Float64:
[:, :, 1] =
 0.0262007  0.0149044  0.00899933  0.00558196  0.00350401  0.00221144  0.00139908  0.000886123  …  0.000262703  0.000419683  0.000672159  0.00108255  0.00176524  0.00295908  0.00527151

[:, :, 2] =
 0.0156273  0.0109913  0.00718346  0.00459756  0.00292272  0.00185359  0.00117448  0.00074391  …  0.000912443  0.0014431  0.00228338  0.00361579  0.00573161  0.00908374  0.0142605

[:, :, 3] =
 0.00842069  0.00650486  0.0045005  0.00296556  0.00191117  0.00121906  0.00077387  0.000490158  …  0.00125962  0.00198565  0.00312493  0.00490036  0.00762628  0.0116795  0.0173084

[:, :, 4] =
 0.00262593  0.00210762  0.00149743  0.00100023  0.000647807  0.000413185  0.000261557  0.000164957  …  0.00117422  0.00184884  0.00290523  0.00454661  0.00705759  0.0107859  0.0160343

[:, :, 5] =
 -0.00275745  -0.00220287  -0.00157002  -0.00105617  -0.000690513  -0.000445253  -0.000285225  …  0.000627442  0.000990405  0.00156691  0.00249021  0.00399015  0.00647543  0.0106509

[:, :, 6] =
 -0.00849839  -0.00655705  -0.00453929  -0.00299523  -0.00193378  -0.00123604  -0.000786409  …  -0.000469446  -0.000720018  -0.00106043  -0.00141352  -0.00138246  0.000389347

[:, :, 7] =
 -0.0155988  -0.0109644  -0.00716212  -0.00458112  -0.00291024  -0.00184427  -0.00116763  …  -0.00142967  -0.00225476  -0.00354958  -0.00556289  -0.00861397  -0.012872  -0.0169656

[:, :, 8] =
 -0.0260201  -0.0147839  -0.00890929  -0.00551279  -0.00345117  -0.00217171  -0.00136973  …  -0.00250611  -0.00396347  -0.00628301  -0.0100122  -0.0161434  -0.0267292  -0.0469493

julia> interior(dwpdz)
1×512×8 view(OffsetArray(::Array{Float64,3}, 0:2, 0:513, 0:9), 1:1, 1:512, 1:8) with eltype Float64:
[:, :, 1] =
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

[:, :, 2] =
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

[:, :, 3] =
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

[:, :, 4] =
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

[:, :, 5] =
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

[:, :, 6] =
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

[:, :, 7] =
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

[:, :, 8] =
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  …  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0

Am I missing something here? Is computation for pressure somehow different from other variables?

Notes:

• If I replace p = sum(model.pressures) for p = model.pressures.pNHS the result is the same.
• If I calculate instead just the w * p term in the kernel the result is correct! So apparently it has to do with the derivative of the pressure, not just the pressure itself.

[glwagner]

What is the type of the object p = sum(model.pressures)?

[tomchor]

It's a binary operation, which I guess is what we expect, right?

julia> p = sum(model.pressures)
BinaryOperation at (Center, Center, Center)
├── grid: RegularRectilinearGrid{Float64, Periodic, Bounded, Bounded}(Nx=1, Ny=512, Nz=8)
│   └── domain: x ∈ [0.0, 150.0], y ∈ [0.0, 6000.0], z ∈ [-80.0, 0.0]
└── tree: 
    + at (Center, Center, Center) via identity
    ├── Field located at (Center, Center, Center)
    └── Field located at (Center, Center, Center)

julia> typeof(p)
Oceananigans.AbstractOperations.BinaryOperation{Center,Center,Center,typeof(+),Field{Center,Center,Center,OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}},RegularRectilinearGrid{Float64,Periodic,Bounded,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}},NamedTuple{(:x, :y, :z),Tuple{CoordinateBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic,Nothing},BoundaryCondition{Oceananigans.BoundaryConditions.Periodic,Nothing}},CoordinateBoundaryConditions{BoundaryCondition{Flux,Nothing},BoundaryCondition{Flux,Nothing}},CoordinateBoundaryConditions{BoundaryCondition{Flux,Nothing},BoundaryCondition{Flux,Nothing}}}}},Field{Center,Center,Center,OffsetArrays.OffsetArray{Float64,3,Array{Float64,3}},RegularRectilinearGrid{Float64,Periodic,Bounded,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}},NamedTuple{(:x, :y, :z),Tuple{CoordinateBoundaryConditions{BoundaryCondition{Oceananigans.BoundaryConditions.Periodic,Nothing},BoundaryCondition{Oceananigans.BoundaryConditions.Periodic,Nothing}},CoordinateBoundaryConditions{BoundaryCondition{Flux,Nothing},BoundaryCondition{Flux,Nothing}},CoordinateBoundaryConditions{BoundaryCondition{Flux,Nothing},BoundaryCondition{Flux,Nothing}}}}},typeof(identity),typeof(identity),typeof(identity),RegularRectilinearGrid{Float64,Periodic,Bounded,Bounded,OffsetArrays.OffsetArray{Float64,1,StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}}}

[glwagner]

Pressure is different from other variables because it doesn't have a dedicated compute! function. I do actually think it would be nice to "regularize" the code in this way by making the pressure (and eddy diffusivities) into KernelComputedFields (or other kinds of special fields) so that compute! works on them in the same way it works on other ComputedFields. The reason it is different is because while we have always had pressure auxiliary variables, ComputedFields are only a relatively recent addition to the code. In fact, it hadn't even occurred to me that it would be nice if we could use compute!(hydrostatic_pressure) until recently.

More on the top of this issue, there is a small difference between the two kernels: in one the reference density (which is otherwise arbitrary?) appears to be 1027 and in the other its 1024.

But otherwise I don't see why it wouldn't work and I am a bit stumped.

[tomchor]

More on the top of this issue, there is a small difference between the two kernels: in one the reference density (which is otherwise arbitrary?) appears to be 1027 and in the other its 1024.

Sorry that was a typo. Just fixed it.

But otherwise I don't see why it wouldn't work and I am a bit stumped.

Yes, I've dedicated several hours to investigating this issue also and got nothing. The issue also is that I'm actually depending on that result for research, and it's pretty expensive to calculate that offline.

Do you have any suggested workarounds?

[glwagner]

Does the order in which you call compute! in your example matter?

You could try making p a ComputedField.

[tomchor]

Does the order in which you call compute! in your example matter?

Just tested it and no.

You could try making p a ComputedField.

I'd tried that already and the results are the same :/

Something interesting is that the code below works:

@kernel function pressure_correlation_z_ccc!(wp, grid, w, p)
    i, j, k = @index(Global, NTuple)

    @inbounds wp[i, j, k] = ℑzᵃᵃᶠ(i, j, k, grid, p) * w[i, j, k] # C, C, C  → C, C, F
end

wp = KernelComputedField(Center, Center, Center, pressure_correlation_z_ccc!, model;
                         field_dependencies=(w, p))

wp_2 = ComputedField(w*p)

compute!(wp)
compute!(wp_2)

The outputs are correct in this case. So apparently what's causing the problem is the term ∂zᵃᵃᶜ(i, j, k, grid, wp) I think. But I cannot figure out why. Why would an interpolation operation work, but not a differentiation operation?

CC'ing @ali-ramadhan to see if he has any ideas.

[glwagner]

Oh I see the problem. In this code

@kernel function pressure_distribution_z_ccc!(dwpdz_ρ, grid, w, p, ρ₀)
    i, j, k = @index(Global, NTuple)
    wp = ℑzᵃᵃᶠ(i, j, k, grid, p) * w[i, j, k] # C, C, C  → C, C, F
    @inbounds dwpdz_ρ[i, j, k] = (1/ρ₀) * ∂zᵃᵃᶜ(i, j, k, grid, wp) # C, C, F  → C, C, C
end 

the object wp is a number. When you differentiate the number wp with ∂zᵃᵃᶜ(i, j, k, grid, wp), ∂zᵃᵃᶜ returns 0.

It might work to instead pass the AbstractOperation wp = w * p into the kernel, and differentiate that.

[tomchor]

Hmm... indeed when I do it like you suggest it seems to work:

@kernel function pressure_correlation_z_ccc!(dwpdz, grid, wp, ρ₀)
    i, j, k = @index(Global, NTuple)

    @inbounds dwpdz[i, j, k] = (1/ρ₀) * ∂zᵃᵃᶜ(i, j, k, grid, wp) # C, C, F  → C, C, C
end

dwpdz_3 = KernelComputedField(Center, Center, Center, pressure_correlation_z_ccc!, model;
                         field_dependencies=(w*p), parameters=1024)

But I don't understand why the other case doesn't work. You mentioned that wp is an number, but isn't it an array?

Also I remember using this approach before in kernels with other properties and it seems to work in other cases. So why not here?

[glwagner]

This:

wp = ℑzᵃᵃᶠ(i, j, k, grid, p) * w[i, j, k]

is a number (the result of interpolating p in z and multiplying the result by w, at the index i, j, k). You can test this at the REPL by picking i, j, k.

∂zᵃᵃᶜ needs to be able to index into wp at both i, j, k and i, j, k+1 in order to calculate a derivative.

[glwagner]

Also I remember using this approach before in kernels with other properties and it seems to work in other cases. So why not here?

If we take a look at those other kernels, maybe we can get to the bottom of it?

[tomchor]

If we take a look at those other kernels, maybe we can get to the bottom of it?

Actually now that I'm checking my scripts, I have always differentiated arguments passed as a computed_dependency. I hadn't paid attention to this before but I've never differentiated a variable that I created inside the kernel, as is the case with wp. Apparently this is yet another limitation of GPU computing that I wasn't aware of.

For my application this already solves the issue, as I can just define a vertical derivative kernel and pass the abstract operation w*p. But for the sake of the general user: is there any way around this limitation?

[tomchor]

Also another thought. When re-writing the script based on your suggestion I found myself writing:

@kernel function ∂zᶜ!(result, grid, operand)
    i, j, k = @index(Global, NTuple)
    @inbounds result[i, j, k] = ∂zᵃᵃᶜ(i, j, k, grid, operand) # C, C, F  → C, C, C
end

vw = @at (Center, Center, Face) w*v
dvwdz = KernelComputedField(Center, Center, Center, ∂zᶜ!, model;
                            field_dependencies=(vw,))

After which I immediately thought: "I thought re-wrote the operation Oceananigans.AbstractOperations.∂z!"

So this to me begs the question: can these operations be generalized using KernelComputedFields to also work on GPUs?

[glwagner]

Ahh hmm... I think you can get more diagnostics to work this way. The reason complex diagnostics fail, as discussed on #1241, is because of the recursive calls to identity (for example). We have to do this with AbstractOperations because that framework is written so that users don't have to "know" where their computations occur (or end up) --- at least not necessarily. In the above code, the user (you) is manually providing location information by selecting ∂zᵃᵃᶜ; therefore we don't have to call an interpolation function (such as identity).

This can be seen in the code in the crucial getindex function defined for all AbstractOperations, such as:

Oceananigans.jl/src/AbstractOperations/derivatives.jl

Line 25 in f6d65f8

@inline Base.getindex(d::Derivative, i, j, k) = d.▶(i, j, k, d.grid, d.∂, d.arg)

I believe in this case that the function d.▶ from the above is identity, which could fail because it's also called in the calculation of the binary operation w * v?

I think we could avoid some calls to identity in AbstractOperations if we restrict their generality. But I'm not sure that's the right philosophy. Maybe we just want to fix the compilation issues instead and preserve full generality in AbstractOperations.

[glwagner]

Apparently this is yet another limitation of GPU computing that I wasn't aware of.

But for the sake of the general user: is there any way around this limitation?

What is the limitation you're referring to?

[tomchor]

I believe in this case that the function d.▶ from the above is identity, which could fail because it's also called in the calculation of the binary operation w * v?

I gotta be honest that whole discussion went right over my head, haha. I think mostly because I don't really understand what's the exact use of identities and what the functions d.▶ actually do.

Am I correct in assuming that if I specify where each calculation is performed (by using @at) I can decrease the computational complexity and maybe make ComputedFields easier to compile on GPUs?

[tomchor]

What is the limitation you're referring to?

I'm referring to the fact that if you calculate a new quantity inside the kernel and then try to differentiate it the differentiation will fail because the new quantity will be a number and not a field. At least that's what I understood is happening in the piece of code below based on your feedback:

wp = ℑzᵃᵃᶠ(i, j, k, grid, p) * w[i, j, k]

Also I'm closing the issue since my I've found a workaround to it based on your comments.

[glwagner]

Great!

Just to clear up the conversation, I think its incorrect to refer to what we've discovered as a "limitation". Instead, we found a bug in the kernel function. The "workaround" is not a work around in the usual sense --- its just one correct way to express the intent with code.

There are several other ways to write this relatively simple kernel, which may be illuminating to explore. Exploiting the fact the the derivative function can be applied to functions of the form f(i, j, k, grid, args...) (as we have designed it) we might also write

@inline wpᵃᵃᶠ(i, j, k, grid, w, p) = @inbounds ℑzᵃᵃᶠ(i, j, k, grid, p) * w[i, j, k] 

@kernel function pressure_distribution_z_ccc!(dwpdz_ρ, grid, w, p, ρ₀)
    i, j, k = @index(Global, NTuple)
    @inbounds dwpdz_ρ[i, j, k] = (1/ρ₀) * ∂zᵃᵃᶜ(i, j, k, grid, wpᵃᵃᶠ, w, p) # C, C, F  → C, C, C
end 

An even more mundane to express our intent is to manually difference the product wp, eg:

@kernel function pressure_distribution_z_ccc!(dwpdz_ρ, grid, w, p, ρ₀)
    i, j, k = @index(Global, NTuple)
    @inbounds wpᵃᵃᶠ_above = ℑzᵃᵃᶠ(i, j, k+1, grid, p) * w[i, j, k+1] 
    @inbounds wpᵃᵃᶠ_below = ℑzᵃᵃᶠ(i, j, k, grid, p) * w[i, j, k] 
    @inbounds dwpdz_ρ[i, j, k] = (1/ρ₀) * (wpᵃᵃᶠ_above - wpᵃᵃᶠ_below) / Δzᵃᵃᶜ(i, j, k, grid) # C, C, F  → C, C, C
end 

Using the differencing function is probably better (its what we do in the source) because it ensures that the indexing convection is treated correctly.