Argument splatting fails on GPU for too many arguments

[jagoosw]

I've been implementing PISCES in OceanBioME which has 24 tracers and 12 auxiliary fields and been struggling with an issue getting it to run on GPU. I couldn't understand why I kept getting Reason: unsupported call to an unknown function (call to jl_f__apply_iterate) errors. I finally worked out that it was coming from this line:

Oceananigans.jl/src/Biogeochemistry.jl

Line 120 in 10e9fb0

return bgc(val_tracer_name, x, y, z, clock.time, fields_ijk...)

where fields_ijk gets splatted.

In this specific instance I can solve this problem by redefining biogeochemical_transitions like:

@inline function biogeochemical_transition(i, j, k, grid, bgc::AbstractContinuousFormBiogeochemistry, val_tracer_name, clock, fields)
    tracer_names_to_extract = required_biogeochemical_tracers(bgc)
    auxiliary_names_to_extract = required_biogeochemical_auxiliary_fields(bgc)

    tracer_fields_ijk = extract_biogeochemical_fields(i, j, k, grid, fields, tracer_names_to_extract)
    auxiliary_fields_ijk = extract_biogeochemical_fields(i, j, k, grid, fields, auxiliary_names_to_extract)

    x = xnode(i, j, k, grid, Center(), Center(), Center())
    y = ynode(i, j, k, grid, Center(), Center(), Center())
    z = znode(i, j, k, grid, Center(), Center(), Center())

    return bgc(val_tracer_name, x, y, z, clock.time, tracer_fields_ijk..., auxiliary_fields_ijk...)
end

I gather that there are different implementations for _apply_iterate depending on the array length in C, and I assume that is the true root of the issue. Since in Oceananigans this seems like the only instance where this kind of error is likely, would it be okay for me to make this change in the source code? I can't think of a way to fix this otherwise.

[glwagner]

What's the current definition of biogeochemical_transition?

[glwagner]

Also though I would suggest using the discrete form for a heavy lifting model. It'll increase the options for development too, you are more restricted with the continuous form which is more for rapid prototyping than production-level

[jagoosw]

It just had all of the fields together:

Oceananigans.jl/src/Biogeochemistry.jl

Lines 108 to 121 in 10e9fb0

	@inline function biogeochemical_transition(i, j, k, grid, bgc::AbstractContinuousFormBiogeochemistry,
	val_tracer_name, clock, fields)
	names_to_extract = tuple(required_biogeochemical_tracers(bgc)...,
	required_biogeochemical_auxiliary_fields(bgc)...)
	fields_ijk = extract_biogeochemical_fields(i, j, k, grid, fields, names_to_extract)
	x = xnode(i, j, k, grid, Center(), Center(), Center())
	y = ynode(i, j, k, grid, Center(), Center(), Center())
	z = znode(i, j, k, grid, Center(), Center(), Center())
	return bgc(val_tracer_name, x, y, z, clock.time, fields_ijk...)
	end

The vision is to convert it to discrete form eventually but we had some undergrads working on it to get it working first.

[glwagner]

What if you try defining a new function to encourage the compiler to inline this:

so rather than calling bgc, call

return splat_bgc(bgc, val_tracer_name, x, y, z, clock.time, fields_ijk)

and then

@inline splat_bgc(bgc, name, x, y, z, t, fields_tuple) = bgc(name, x, y, z, t, fields_tuple...)

or maybe

@inline splat_bgc(bgc, name, x, y, z, t, fields_tuple::NTuple{N}) where N = bgc(name, x, y, z, t, fields_tuple...)

sometime the type annotations encourage the compiler to try harder / specialize where it may have given up previously.

We've run into inference issues with splatting before in other situations. I think your fix may work only for this specific case but may not generalize to even bigger models. If we can find a more general fix that'd be best, but it might be hard.

I also wonder if there is an issue specializing on the bgc type but not sure.

Another strategy could be to take the pattern you've demonstrated @jagoosw but generalize it, ie simply split fields_ijk into two equal parts, something like

function splat_by_division(bgc, name, x, y, z, t, fields::NTuple{N}) where N
    M = floor(Int, N/2)
    return bgc(name, x, y, z, t, fields[1:M]..., fields[M+1:N]...)
end

This is a bit nicer for future developers, because it is clear how to continue to generalize this pattern if the issue crops up again.

[jagoosw]

I can try these suggestions.

Another strategy could be to take the pattern you've demonstrated @jagoosw but generalize it, ie simply split fields_ijk into two equal parts, something like

function splat_by_division(bgc, name, x, y, z, t, fields::NTuple{N}) where N
M = convert(N/2, Int)
return bgc(name, x, y, z, t, fields[1:M]..., fields[M+1:N]...)
end

This is a bit nicer for future developers, because it is clear how to continue to generalize this pattern if the issue crops up again.

I'm not sure that this would work on GPU though right because I tried to do something similar to this first and got an error about building 1D arrays from the fields[1:M] bit, although I did try and do something recursive where it splatted the first 15 and then recursively called the function if there were more left so maybe a simpler way like this would compile.

[glwagner]

I gather that there are different implementations for _apply_iterate depending on the array length in C, and I assume that is the true root of the issue.

To respond to this, I believe we want the compiler to inline the splatting, so so that we don't even reach apply_iterate.

[glwagner]

With what you have written, we lose the velocity field. Are we sure that no bgc models would ever need to depend on the velocity field? I'm not so sure. It's plausible they might need to depend on strain, for example.

[jagoosw]

I see, maybe your first suggestion will work then although it is quite a long chain to infer how long that splat is going to be.

[jagoosw]

How do we lose velocities? I agree that its plausible they might be needed, but considering it now I don't think we support it at the moment because validate_biogeochemistry would fail if you put velocity components in required_biogeochemical_tracers because they wouldn't be in tracers. But we could fix that by checking that required_biogeochemical_tracers are in any of fields(model), and this change wouldn't affect that.

[glwagner]

Oh I see, I was mistaken. The original code also does not include velocities.

[glwagner]

In that case then I am fine with

@inline function biogeochemical_transition(i, j, k, grid, bgc::AbstractContinuousFormBiogeochemistry, val_tracer_name, clock, fields)
    tracer_names_to_extract = required_biogeochemical_tracers(bgc)
    auxiliary_names_to_extract = required_biogeochemical_auxiliary_fields(bgc)

    tracer_fields_ijk = extract_biogeochemical_fields(i, j, k, grid, fields, tracer_names_to_extract)
    auxiliary_fields_ijk = extract_biogeochemical_fields(i, j, k, grid, fields, auxiliary_names_to_extract)

    x = xnode(i, j, k, grid, Center(), Center(), Center())
    y = ynode(i, j, k, grid, Center(), Center(), Center())
    z = znode(i, j, k, grid, Center(), Center(), Center())

    return bgc(val_tracer_name, x, y, z, clock.time, tracer_fields_ijk..., auxiliary_fields_ijk...)
end

(I thought that this change was removing the velocity fields.) Now I understand it is identical but simply eliminates one splatting which may help compilation.

[glwagner]

Also with continuous form its not possible to use things like strain/shear anyways. It seems discrete form is required for that (which doesn't have the compilation problems that the continuous form has).