Indexing error when using CUDA arrays for custom Lagrangian particles

[sebasti4nmc]

Hello, this is related to the same simulation I'm working on with @amrapallig that I was running on my previous discussion post (#4655), implementing Lagrangian particles with interpolated temp/velocity fields. I was successfully able to run the code on my CPU without any issues using StructArrays.

I am now currently trying to run my code at higher resolution using an NVIDIA A100 GPU. So far, I have been able to run my simulation with particles without any custom properties or field tracking. However, when I try to use my custom particle StructArrays, I run into errors when I try to use CudaArrays. Specifically, I get a "scalar indexing disallowed error" when I try to store CudaArrays in my StructArrays, even though I am using @allowscalar in my global scope. I have also tried directly passing a StructArray with CPU arrays to LagrangianParticles(), but then I get an error in my model constructor because the arrays aren't isbits.

Any idea how to fix this? Is there a more suitable approach/command for handling this on GPU? Here is my current code:

using Oceananigans, JLD2, Printf, StructArrays, CUDA
using CUDA: CUDA.allowscalar, @allowscalar
using Oceananigans.Units: minute, minutes,hour, hours
using Oceananigans: VelocityFields

grid = RectilinearGrid(GPU(), size=(256, 256, 160),   # Nx, Ny, Nz
                              extent=(128, 128, 80)) # Lx, Ly, Lz

u, v, w = VelocityFields(grid)

T = CenterField(grid)
S = CenterField(grid)

### Lagrangian particles
struct Particle
     x :: Float64
     y :: Float64
     z :: Float64
     u :: Float64
     v :: Float64
     w :: Float64
     T :: Float64
 end

# Function to set initial particle fields
function initialize_particle_fields(grid)
    spacing = 1
    Nparticles = Int.((grid.Lx * grid.Ly * 64) / (spacing ^ 3)) # Nx/spacing * Ny/spacing * Nz/spacing
    u₀, v₀, w₀ = [zeros(Nparticles) for i in 1:3]
    T₀ = fill(20.0, Nparticles)

    u₀, v₀, w₀, T₀ = CuArray(u₀), CuArray(v₀), CuArray(w₀), CuArray(T₀)

    return u₀, v₀, w₀, T₀
end

u₀, v₀, w₀, T₀ = initialize_particle_fields(grid)

# Function to set initial particle positions
function initialize_particles(grid)
    Δx = grid.Lx / grid.Nx
    Δy = grid.Ly / grid.Ny
    Δz = grid.Lz / grid.Nz
    spacing = 1
    ϵ = 1e-6

    x₀, y₀, z₀ = Float64[], Float64[], Float64[]

    for k in 1:spacing:64, j in 1:spacing:grid.Ly, i in 1:spacing:grid.Lx
        push!(x₀, (i - 0.5) * Δx - ϵ)
        push!(y₀, (j - 0.5) * Δy - ϵ)
        push!(z₀, -((k - 0.5) * Δz - ϵ))
    end

    x₀, y₀, z₀ = CuArray(x₀), CuArray(y₀), CuArray(z₀)

    return x₀, y₀, z₀
end

x₀, y₀, z₀ = initialize_particles(grid)

# StructArray with initial conditions
particle_array = StructArray{Particle}(x = x₀,
                                       y = y₀,
                                       z = z₀,
                                       u = u₀,
                                       v = v₀,
                                       w = w₀,
                                       T = T₀)

# Construct keyword argument for model
lagrangian_particles = LagrangianParticles(particle_array;
    tracked_fields = (; u, v, w, T))

model = NonhydrostaticModel(; grid,
                            buoyancy = SeawaterBuoyancy(; equation_of_state),
                            advection = UpwindBiased(order=5),
                            timestepper = :RungeKutta3,
                            tracers = (; T, S),
                            coriolis = FPlane(f = 1e-4),
                            closure = SmagorinskyLilly(),
                            stokes_drift = UniformStokesDrift(∂z_uˢ=∂z_uˢ),
                            boundary_conditions = (u=u_bcs, T=T_bcs, S=S_bcs),
                            forcing = (;u = u_forcing, v = u_forcing, w = u_forcing),
                            particles = lagrangian_particles
                            )

Let me know if you have any solutions.