Move `GPUArrays.backend` here

N5N3 commented 2 years ago

This function is the last piece we need to enable gpu broadcast for StructArray. Move it here as Adapt is more lightweight.

maleadt commented 2 years ago

Adapt is a generic package, moving GPUArrays-specific functionality here isn't its purpose. There's no reason GPUArrays.backend should live here other than it being a lightweight dependency, so I don't feel comfortable using Adapt as a container for that.

N5N3 commented 2 years ago

Sorry for the noise. I thought Adapt.jl is a good choice as it has been included as StructArrays.jl's dependency. I guess there's no plan for GPUArraysCore.jl? Maybe a glue package is a better choice.

maleadt commented 2 years ago

Is GPUArrays.jl that heavy of a dependency? It doesn't depend on much, and doesn't load much code.

N5N3 commented 2 years ago

I didn't test it. But we only want to add one definition in StructArrays.jl.

function GPUArrays.backend(x::StructArray)
    cs = components(x)
    back = GPUArrays.backend(cs[1])
    for i in 2:length(cs)
        back === Adapt.backend(cs[i]) || error("backend mismatch!")
    end
    back
end

So it seems good to not add more dependency. @piever

piever commented 2 years ago

So, I think GPUArrays is not particularly heavy, but it's also true that StructArrays is one of those low-level, low-dependency packages that should ideally depend on very little (in an ideal scenario, it should have no dependencies at all IMO).

The main goal, IIUC, is to have copyto! work out of the box even with StructArrays of GPU arrays. @N5N3 can correct me if I'm wrong, but I think the method we need to work is copyto!(dest::StructArray, bc::Broadcasted{<:AbstractGPUArrayStyle}). That requires

https://github.com/JuliaGPU/GPUArrays.jl/pull/393, and
backend(StructArray{T, N, C}) to give the correct result, based on C, the type of component arrays.

I suggested to move backend here, but I see how that can be a bad idea. OTOH, I also think it's a bit odd to have to depend on both Adapt and GPUArrays interfaces to define a wrapper type with custom broadcast. Maybe a milder alternative would be to define a "general use" function (which StructArrays would overload) in Adapt to be used as a fallback for GPUArrays.backend. I see something similar is done with Adapt.parent here. At least in principle, if we know how to adapt a type to the GPU, we should also know the backend of the adapted type.

EDIT: an alternative option would be to use the broadcast style rather than the array type to infer the backend in copyto!, but I'm not sure how feasible / clean that is.

JuliaGPU / Adapt.jl

Move `GPUArrays.backend` here #50