GPU issue indexing flat fields

[jagoosw]

When I try to index into a field which has nothing topology in a kernel, e.g.:

@kernel function add_one!(A)
    i, j = @index(Global, NTuple)
    A[i, j] += 1
end

then it works on CPU, but hits a dynamic function invocation on line 82 here:

Oceananigans.jl/src/Grids/grid_utils.jl

Lines 81 to 89 in 2b601f1

	@inline function Base.size(loc, topo, sz, indices=default_indices(Val(length(loc))))
	D = length(loc)
	# (it's type stable?)
	return ntuple(Val(D)) do d
	Base.@_inline_meta
	length(instantiate(loc[d]), instantiate(topo[d]), sz[d], indices[d])
	end
	end

could we rewrite this like:

@inline function Base.size(loc::NTuple{N},...) where N

instead?

[simone-silvestri]

The problem here is actually this line

Oceananigans.jl/src/Fields/abstract_field.jl

Line 54 in b67ad56

Base.size(f::AbstractField) = size(f.grid, location(f))

that gets called here

Oceananigans.jl/src/Fields/abstract_field.jl

Line 66 in b67ad56

@inline function Base.axes(f::Abstract3DField)

It is not a problem on CPUs because on CPUs we have a grid associated to a field. On GPUs we throw away the grid so this function errors. I can open a PR to fix this

[simone-silvestri]

PS this does not happen if you index with 3 indices since axes does not need to be called since we have specific methods for 3d indexing of reduced fields

[glwagner]

Before you fix this @simone-silvestri, the issue is that the indexing should be A[i, j, 1]

[simone-silvestri]

I see a lot of users try to access 2D fields with [i, j], which does make sense (this was supported in Oceananigans up to the change to axes. I think it should be included. Open to not include it though.

[glwagner]

Hmm interesting. I think that support was an accident?

[glwagner]

I guess we have to decide whether its important for people to explicitly declare the dead indices or not. Another consideration is whether this permits more general code (eg reduced code independent of direction). I haven't seen the latter though.

[glwagner]

Somehow internally in Oceananigans we have never used a 1- or 2-index form, but @simone-silvestri you're saying this appears in user code? If you have examples that will help shape this discussion.