Kernel function `interpolate_primary_atmospheric_state!` calls `intrinsic_vector`, which can lead to `sind` call inside GPU kernel.

[jlk9]

This was discovered through this issue: EnzymeAD/Reactant.jl#2345

The kernel function _interpolate_primary_atmospheric_state calls the Oceananigans function intrinsic_vector:

ClimaOcean.jl/src/Atmospheres/interpolate_atmospheric_state.jl

Lines 121 to 159 in 1e9823e

	@kernel function _interpolate_primary_atmospheric_state!(surface_atmos_state,
	space_fractional_indices,
	time_interpolator,
	exchange_grid,
	atmos_velocities,
	atmos_tracers,
	atmos_pressure,
	downwelling_radiation,
	prescribed_freshwater_flux,
	atmos_backend,
	atmos_time_indexing)
	i, j = @index(Global, NTuple)
	ii = space_fractional_indices.i
	jj = space_fractional_indices.j
	fi = get_fractional_index(i, j, ii)
	fj = get_fractional_index(i, j, jj)
	x_itp = FractionalIndices(fi, fj, nothing)
	t_itp = time_interpolator
	atmos_args = (x_itp, t_itp, atmos_backend, atmos_time_indexing)
	uₐ = interp_atmos_time_series(atmos_velocities.u, atmos_args...)
	vₐ = interp_atmos_time_series(atmos_velocities.v, atmos_args...)
	Tₐ = interp_atmos_time_series(atmos_tracers.T, atmos_args...)
	qₐ = interp_atmos_time_series(atmos_tracers.q, atmos_args...)
	pₐ = interp_atmos_time_series(atmos_pressure, atmos_args...)
	Qs = interp_atmos_time_series(downwelling_radiation.shortwave, atmos_args...)
	Qℓ = interp_atmos_time_series(downwelling_radiation.longwave, atmos_args...)
	# Usually precipitation
	Mh = interp_atmos_time_series(prescribed_freshwater_flux, atmos_args...)
	# Convert atmosphere velocities (usually defined on a latitude-longitude grid) to
	# the frame of reference of the native grid
	kᴺ = size(exchange_grid, 3) # index of the top ocean cell
	uₐ, vₐ = intrinsic_vector(i, j, kᴺ, exchange_grid, uₐ, vₐ)

The problem is that (at least in some situations, like the script in the linked issue) the exchange_grid passed to intrinsic_vector is an OrthogonalSphericalShellGrid. The method for intrinsic_vector in this case calls sind, cosd, and atand (some through rotation_angle) which are all not supported in GPU kernels:

https://github.com/CliMA/Oceananigans.jl/blob/ae70538358c5ddf9e38267cccbd2bbf05bf7b66a/src/Operators/vector_rotation_operators.jl#L96-L109

Thus calling _interpolate_primary_atmospheric_state with a reactant or GPU backend can error.

A fix would be to modify _interpolate_primary_atmospheric_state so it doesn't call intrinsic_vector. Alternatively that method call of intrinsic_vector can be modified to get rid of the problem trig functions.

[glwagner]

Instead of inlining your links, can you paste them? it will help reading the code

[glwagner]

for example:

ClimaOcean.jl/src/Atmospheres/interpolate_atmospheric_state.jl

Line 159 in 1e9823e

uₐ, vₐ = intrinsic_vector(i, j, kᴺ, exchange_grid, uₐ, vₐ)

[glwagner]

https://github.com/CliMA/Oceananigans.jl/blob/86df36554019c27b7b5d527e033f21646f4e2ee4/src/Operators/vector_rotation_operators.jl#L96-L110

[glwagner]

This is an Oceananigans issue right? we need to use hack_sind and hack_cosd (and a new hack_tand)

[simone-silvestri]

A fix would be to modify _interpolate_primary_atmospheric_state so it doesn't call intrinsic_vector

we need to rotate the vectors in any case otherwise the velocity difference is not computed correctly. @glwagner's solution works nicely.

[jlk9]

Instead of inlining your links, can you paste them? it will help reading the code

I did that with the link to ClimaOcean code, apparently the link to intrinsic_vector won't show code since it's a different repo.

we need to rotate the vectors in any case otherwise the velocity difference is not computed correctly. @glwagner's solution works nicely.

Agree that this makes the most sense for a fix. I can make a PR for it later today.

[simone-silvestri]

another option would be to not use atand in the rotation_angle function, but the atan function, which will then return the angle in radians and then simple sin, cos, and tan will work

[glwagner]

Instead of inlining your links, can you paste them? it will help reading the code

I did that with the link to ClimaOcean code, apparently the link to intrinsic_vector won't show code since it's a different repo.

thank you <3 ! even for external repos, the full link is better. We can read the link to see what file it points to.