Converting entire custom structure from cpu to gpu and viceversa

Specific Domains GPU
,
1

Hi!

I am looking for a good way to convert a custom structure of Arrays to CuArrays and viceversa. I suspect the answer is related to them and to this topic, but I am not well versed in the usage of convert and adapt so I don’t really understand what I should do.

In my code, I have the following custom structure where A is either Array or CuArray

struct MyStruct{A}
    x::A
    y::A
end

I first create a GPU version of it

n = 3
a_gpu = CUDA.rand(n)
b_gpu = CUDA.rand(n)
MS_gpu = MyStruct(a_gpu, b_gpu)

and then perform operations with this on the GPU (thank you for CUDA.jl, it’s amazing!).

But then I need to perform operations with scalar indexing on these objects, so my understanding is that it is more efficient to “bring” the structure to the CPU for that. What I currently do is the following.

First I create a CPU version

a_cpu = zeros(n)
b_cpu = zeros(n)
MS_cpu = MyStruct(a_cpu, b_cpu)

Then I created the following function (my structures are various and larger, so it made sense to have a somewhat automated way to do this)

function convert_to_cpu!(struct_gpu::T_gpu, struct_cpu::T_cpu) where {T_gpu<:MyStruct, T_cpu<:MyStruct}
    names = fieldnames(T_gpu)
    types_cpu = fieldtypes(T_cpu)

    for i in 1:lastindex(names)
        getfield(struct_cpu, names[i]) .= convert(types_cpu[i], getfield(struct_gpu, names[i]))
    end
    return nothing
end

and use it to copy the GPU stuff into the CPU structure

convert_to_cpu!(MS_gpu, MS_cpu)

I do get the desired final result, which is to have a structure of normal Arrays on the CPU with the values I calculated on the GPU, but I am not sure this is the proper way to do it.

PS: I marked the topic as GPU domain but I suppose this question has a much more general answer, apologies for that!

2

You should be able to use Adapt.jl for this, it should be sufficient to define an overload of Adapt.adapt_structure for you custom struct as described in the README and docstring of Adapt.adapt, or use the macro Adapt.@adapt_structure MyStruct which will do the same thing automatically. Then you can use adapt(CuArray, MS_cpu) or adapt(Array, MS_gpu) to transfer your object back and forth between CPU and GPU.

3

Thank you, it works!

I had seen the documentation of Adapt.adapt but had never understood that to had to be replaced by a type and would just work out of the box.

Indeed, I defined

function Adapt.adapt_structure(to, MS::MyStruct)
    x = adapt(to, MS.x)
    y = adapt(to, MS.y)
    return MyStruct(x, y)
end

and now the commands adapt(CuArray, MS_cpu) and adapt(Array, MS_gpu) do precisely what I was after.

1 Like