`resize_immersed_boundary` leads to scalar indexing of GPU arrays

[jackdfranklin]

The issue arises when using distributed grids on GPU, and results in an error:

Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.

A shortened version of the stack trace is:

  [5] getindex
    @ /glade/work/jfranklin/.julia/packages/GPUArrays/VNJ6B/src/host/indexing.jl:50 [inlined]
  [6] getindex
    @ /glade/work/jfranklin/.julia/packages/OffsetArrays/b7qpm/src/OffsetArrays.jl:437 [inlined]
  [7] copyto_unaliased!(deststyle::IndexCartesian, dest::SubArray{Float64, 3, CuArray{Float64, 3, CUDACore.DeviceMemory}, Tuple{UnitRange{Int64}, UnitRange{Int64}, UnitRange{Int64}}, false}, srcstyle::IndexLinear, src::OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDACore.DeviceMemory}})
    @ Base ./abstractarray.jl:1109
  [8] copyto!(dest::SubArray{Float64, 3, CuArray{Float64, 3, CUDACore.DeviceMemory}, Tuple{UnitRange{Int64}, UnitRange{Int64}, UnitRange{Int64}}, false}, src::OffsetArrays.OffsetArray{Float64, 3, CuArray{Float64, 3, CUDACore.DeviceMemory}})
    @ Base ./abstractarray.jl:1070
  [9] set_to_array!(u::Field...)
    @ Oceananigans.Fields ~/CliMA/Oceananigans.jl/src/Fields/set!.jl:131
 [10] _set!(u::Field...)
    @ Oceananigans.DistributedComputations ~/CliMA/Oceananigans.jl/src/DistributedComputations/distributed_fields.jl:58
 [11] set!(u::Field...)
    @ Oceananigans.DistributedComputations ~/CliMA/Oceananigans.jl/src/DistributedComputations/distributed_fields.jl:42
 [12] resize_immersed_boundary(ib::PartialCellBottom...)
    @ Oceananigans.DistributedComputations ~/CliMA/Oceananigans.jl/src/DistributedComputations/distributed_immersed_boundaries.jl:120

I am running the distributed benchmarks on GPUs, with the --distributed -partition "1x2x1" flags.

[jackdfranklin]

I think the solution could be to use Architectures.device_copy_to!

[glwagner]

Is it possible to do a mem copy to a non-contiguous array?

If it is not, then the only way I see to solve this is to allocate a temporary contiguous temp_dest on the GPU and then broadcast into dest from there?

[jackdfranklin]

Yeah I think that might be the only option. I can overload copyto! when the destination is a subarray on GPU in order to do that?

[glwagner]

Yes but don't directly overload copyto! because that's type piracy. But we can define our own internal copyto method for it

[jackdfranklin]

Maybe there is a better way...if both arrays are on the same device can I broadcast them automatically?