JuliaGPU / Metal.jl

Metal programming in Julia
MIT License
338 stars 36 forks source link

Generalize `adapt` to allow specifying the storage mode but not the element type and/or number of dimensions #363

Open mtfishman opened 3 weeks ago

mtfishman commented 3 weeks ago

For example:

julia> using Metal

julia> using Adapt

julia> adapt(MtlArray{Float32,<:Any,Managed}, randn(2, 2))
2×2 MtlMatrix{Float32, Private}:
  0.527836  0.597629
 -0.225364  0.211664

julia> adapt(MtlArray{<:Any,<:Any,Managed}, randn(Float32, 2, 2))
2×2 Matrix{Float32}:
 -0.933958  0.784281
 -1.37475   0.566195

julia> using Pkg

julia> Pkg.status(["Metal", "Adapt"])
Status `.../Project.toml`
  [79e6a3ab] Adapt v4.0.4
  [dde4c033] Metal v1.1.0

I think there is the same issue in other GPU backends.

maleadt commented 3 weeks ago

I'm not sure that's so valuable. Why don't you just call adapt(MtlArray{eltype(x),ndims(x),Managed}, x)? The MtlArray-based adaptor shouldn't ever change the eltype.

mtfishman commented 3 weeks ago

The application is that I want to adapt a nested data structure that stores a number of GPU arrays where the GPU arrays can have different numbers of dimensions (say schematically of type Vector{Union{MtlArray{Float32,1},MtlArray{Float32,2}}}), so it isn't straightforward/possible to choose the number of dimensions ahead of time. Same thing with the element type, the GPU arrays could have different element types, i.e. Vector{Union{MtlArray{Float32,1},MtlArray{Complex{Float32},2}}}.

mtfishman commented 3 weeks ago

This could of course be done with custom adapters, but I prefer to not need to define my own adapters if it can be avoided and also give users options that "just work" that are directly from the GPU ecosystem so we don't have to support our own custom GPU functionality/interfaces.