Kernel calls for AMD diffusivity calculations

[tomchor]

Pasted from a slack discussion:

Hi, everyone. Can someone please give me a little push in the right direction? I'm trying to understand the calls to calculate_tracer_diffusivity but I can't make sense of them.
The relevant code is here:

function calculate_diffusivities!(K, arch, grid, closure::AnisotropicMinimumDissipation, buoyancy, U, C)
    workgroup, worksize = work_layout(grid, :xyz)
    barrier = Event(device(arch))
    viscosity_kernel! = calculate_viscosity!(device(arch), workgroup, worksize)
    diffusivity_kernel! = calculate_tracer_diffusivity!(device(arch), workgroup, worksize)
    viscosity_event = viscosity_kernel!(K.νₑ, grid, closure, buoyancy, U, C, dependencies=barrier)
    events = [viscosity_event]
    for (tracer_index, κₑ) in enumerate(K.κₑ)
        @inbounds c = C[tracer_index]
        event = diffusivity_kernel!(κₑ, grid, closure, c, Val(tracer_index), U, dependencies=barrier)
        push!(events, event)
    end
    wait(device(arch), MultiEvent(Tuple(events)))
    return nothing
end
@kernel function calculate_viscosity!(νₑ, grid, closure::AnisotropicMinimumDissipation, buoyancy, U, C)
    i, j, k = @index(Global, NTuple)
    @inbounds νₑ[i, j, k] = νᶜᶜᶜ(i, j, k, grid, closure, buoyancy, U, C)
end
@kernel function calculate_tracer_diffusivity!(κₑ, grid, closure, c, tracer_index, U)
    i, j, k = @index(Global, NTuple)
    @inbounds κₑ[i, j, k] = κᶜᶜᶜ(i, j, k, grid, closure, c, tracer_index, U)
end

So it seems that calculate_tracer_diffusivity! is defined with the signature κₑ, grid, closure, c, tracer_index, U, but is called with the signature device(arch), workgroup, worksize.
I can see that the diffusivity_kernel! is somehow what actually performs the "correct" signature call in the line event = diffusivity_kernel!(κₑ, grid, closure, c, Val(tracer_index), U, dependencies=barrier), but I honestly can't understand why that's the case or even how that works.

Could someone please shed some light?

CC @glwagner

[tomchor]

The answer for this, as per @glwagner's comments, is because that's the syntax for KernelAbstractions, which we use to compute GPU kernels:

From @glwagner :

The kernel is created here:
diffusivity_kernel! = calculate_tracer_diffusivity!(device(arch), workgroup, worksize)
and launched here:
event = diffusivity_kernel!(κₑ, grid, closure, c, Val(tracer_index), U, dependencies=barrier)
the diffiusivity calculations do not depend on one another, so they are all launch simultaneously. At the end, we wait for all of them to finish here:
wait(device(arch), MultiEvent(Tuple(events)))

Noticeably apparently we don't use lauch! in this case, even though apparently we should?

[tomchor]

Here's a question: wouldn't it be better if this were a KernelComputedField? (Or maybe even a KernelFunctionOperation?)

[glwagner]

Here's a link to the Oceananigans source code relevant to this discussion:

Oceananigans.jl/src/TurbulenceClosures/turbulence_closure_implementations/anisotropic_minimum_dissipation.jl

Lines 190 to 221 in 472bdca

	function calculate_diffusivities!(K, arch, grid, closure::AnisotropicMinimumDissipation, buoyancy, U, C)
	workgroup, worksize = work_layout(grid, :xyz)
	barrier = Event(device(arch))
	viscosity_kernel! = calculate_viscosity!(device(arch), workgroup, worksize)
	diffusivity_kernel! = calculate_tracer_diffusivity!(device(arch), workgroup, worksize)
	viscosity_event = viscosity_kernel!(K.νₑ, grid, closure, buoyancy, U, C, dependencies=barrier)
	events = [viscosity_event]
	for (tracer_index, κₑ) in enumerate(K.κₑ)
	@inbounds c = C[tracer_index]
	event = diffusivity_kernel!(κₑ, grid, closure, c, Val(tracer_index), U, dependencies=barrier)
	push!(events, event)
	end
	wait(device(arch), MultiEvent(Tuple(events)))
	return nothing
	end
	@kernel function calculate_viscosity!(νₑ, grid, closure::AnisotropicMinimumDissipation, buoyancy, U, C)
	i, j, k = @index(Global, NTuple)
	@inbounds νₑ[i, j, k] = νᶜᶜᶜ(i, j, k, grid, closure, buoyancy, U, C)
	end
	@kernel function calculate_tracer_diffusivity!(κₑ, grid, closure, c, tracer_index, U)
	i, j, k = @index(Global, NTuple)
	@inbounds κₑ[i, j, k] = κᶜᶜᶜ(i, j, k, grid, closure, c, tracer_index, U)
	end

[glwagner]

Here's a question: wouldn't it be better if this were a KernelComputedField? (Or maybe even a KernelFunctionOperation?)

I think reusing more features that have been developed for user applications internally is a great idea in general. More code reuse, more understandable internals. I have thought about this a bit for diffusivities and pressure fields, which are both auxiliary fields to a model and some flavor of "ComputedField".

I think we would need to expand on the user interface to KernelComputedField for this purpose. There are a few reasons. First, the user interface to ComputedField does not supply a way to combine kernel computations for more than one field. Second, we need to develop a user interface for propagating arguments from compute! into the kernel. The reason is that we have to pass in both closure and buoyancy into this kernel. KernelComputedField supports parameters right now. However, I think its important that these arguments are explicitly passed from the model (rather than being bound to the diffusivities at model construction time) --- because models are mutable, which means that users can change the closure or buoyancy model on the fly (provided the new object has the same type as the old one).

Expanding the computed field infrastructure to support simultaneous computation and propagated arguments would be really cool I think and a great addition to the code. Once that's done we could consider refactoring internal code, which would also be awesome.

[navidcy]

Should we convert this to Discussions?

[tomchor]

Sorry I just realized opened this in the wrong section! Yes, feel free to move it to a discussion if that's better