General improvements for use in CUDAnative.jl

[maleadt]

Redesign as an attempt to address https://github.com/JuliaGPU/CUDAnative.jl/issues/121

There's been lots of duplication between CUDAnative.cudaconvert and adapt, so it makes sense to try and merge them. However, cudaconvert was designed to "convert for execution on CUDA GPUs" (ie. towards an environment/context), while adapt is more like a generalized form of convert (ie. towards a type), so there's a difference in how these APIs are used and extended.

In many cases though, adapting towards a specific type doesn't make much sense IMO. Having people do or implement adapt(CuDeviceArray, arr) seems awfully specific, and doesn't solve the problem of conversions that do not yield a CuDeviceArray. So in this PR I propose to use a AbstractAdaptor type hierarchy for deciding how to adapt structs. It should be possible to express the same things we're doing right now with adapt, and make it possible to replace cudaconvert by having a CUDA adaptor and dispatch on that.

I've also renamed adapt and adapt_ to adapt_structure/adapt_storage which seems to match the pattern that matters 99% of the time (but feel free to bikeshed on this one).

Single-blob version of the diff in this PR + usage in external libraries:

## Adapt.jl

export adapt, AbstractAdaptor

abstract type AbstractAdaptor end

# external interface
adapt(A::Type{AbstractAdaptor}, x) = adapt_structure(A, x)

# interface for libraries to implement
# two levels of priority:
# - structure, eg. array wrappers, should get preserved and get priority in dispatch
# - storage, eg. CuArray
adapt_structure(A::Type{AbstractAdaptor}, x) = adapt_storage(A, x)
adapt_storage(::Type{AbstractAdaptor}, x) = x

# Base
adapt_structure(A::Type{<:AbstractAdaptor}, xs::Tuple) = Tuple(adapt(A, x) for x in xs)
@generated adapt_structure(A::Type{<:AbstractAdaptor}, x::NamedTuple) =
    Expr(:tuple, (:($f=adapt(A, x.$f)) for f in fieldnames(x))...)

# LinearAlgebra
import LinearAlgebra: Adjoint, Transpose
adapt_structure(A::Type{AbstractAdaptor}, x::Adjoint)   = Adjoint(adapt(A, parent(x)))
adapt_structure(A::Type{AbstractAdaptor}, x::Transpose) = Transpose(adapt(A, parent(x)))

## CUDAnative

struct CUDAAdaptor <: AbstractAdaptor end

## CuArrays.jl

adapt_storage(::Type{CUDAAdaptor}, x::Array) = (println("adapt $(typeof(x)) to CuDeviceArray"); x)

## User

adapt(CUDAAdaptor, [1]')

[maleadt]

@vchuravy mentions how it would probably be better to simplify this interface using singleton instances (eg. adapt(CUDAAdaptor(), ...)) cfr. BroadcastStyle.

[MikeInnes]

Not wanting to use CuDeviceArray seems totally fine, but you could already use adapt(CUDAAdaptor(), ...) with the current design, no? I guess the main change from there is enforcing a type heirarchy, is there a particular benefit to that?

I like the idea of splitting structure/storage more clearly, but perhaps we should still define an adapt alias just for user purposes – then it's totally clear what you're actually supposed to call.

[maleadt]

I like the idea of splitting structure/storage more clearly, but perhaps we should still define an adapt alias just for user purposes – then it's totally clear what you're actually supposed to call.

Definitely, that's still there, right?

Not wanting to use CuDeviceArray seems totally fine, but you could already use adapt(CUDAAdaptor(), ...) with the current design, no? I guess the main change from there is enforcing a type heirarchy, is there a particular benefit to that?

I personally dislike fully untyped interfaces. Besides, without type constraints user code that uses the current adapt(CuArray, ...) paradigm would silently fail. It seems like a confusing interface design to me to support both the existing convert(T, x)::T-like adapt, and the "new" approach where the first argument is an adaptor.

[maleadt]

Added to CuArrays and CUDAnative (passing all tests) as a try-out: https://github.com/JuliaGPU/CUDAnative.jl/compare/tb/adapt https://github.com/JuliaGPU/CuArrays.jl/compare/tb/adapt

[MikeInnes]

Definitely, that's still there, right?

Whoops, missed that, great.

On typing, I guess that while I see an adaptor making sense for cudaconvert, I'd still be ok with writing adapt(CuArray, x); it seems to me that viewing this as a generalised conversion-with-structure is reasonable here. On top of which the original thing was more flexible in that you could adapt to CuArray{Float32} or CuArray{Float64}. Of course, you could reproduce type parameters in CUDAHostAdaptor, but that just starts to seem like code duplication to me.

[maleadt]

On top of which the original thing was more flexible in that you could adapt to CuArray{Float32} or CuArray{Float64}.

That is true. Is that a common thing to do?

Either way, I'd argue that an element-type parameter to CUDAHostAdaptor is a more sensible thing to do, because for non-CuArray types it doesn't make sense to adapt towards eg. CuArray{Float32}. Just stripping types and using adapt quite differently seems like bad API design to me.

[maleadt]

https://github.com/JuliaGPU/CuArrays.jl/commit/4b25cbdbee09562ea2ac05183acf57a79c2d55bc implements that feature now (in exactly the same way as before)

[MikeInnes]

At least for my purposes I'm happy just always going to f32, but not doing that seems to be a thing occasionally as well.

Anyway, I can get on board with this if you're sure it's the right way to go. I guess it will have to come with a round of package bounds before merging/tagging?

[maleadt]

Anyway, I can get on board with this if you're sure it's the right way to go.

I'm really not, just trying to figure something out here. Maybe you're right though, since CUDAHostAdaptor{T} looks awfully similar to CuArray{T} now. I mostly was against confounding those two since I was planning on using that adaptor in CUDAnative.jl too, but I've moved to a separate one over there. I'll play with it some more.

[maleadt]

Moved away from most of the radical redesign, keeping it simple instead. Mainly features the split into adapt_structure and adapt_storage, and a port of most of the cudaconvert functionality from CuArrays and CUDAnative. Good to go? Will need a tag before we can use it from the other packages, as well as updates over there to overload the proper functions.

[codecov-io]

Codecov Report

Merging #9 into master will not change coverage. The diff coverage is 100%.

@@          Coverage Diff          @@
##           master     #9   +/-   ##
=====================================
  Coverage     100%   100%           
=====================================
  Files           1      2    +1     
  Lines           1      6    +5     
=====================================
+ Hits            1      6    +5

Impacted Files	Coverage Δ
src/base.jl	100% <100%> (ø)
src/Adapt.jl	100% <100%> (ø)	:arrow_up:

---

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update 8bb4c0a...10a3752. Read the comment docs.

[MikeInnes]

Yes, I'm happy with this if you are. Very glad to see this being used in CUDAnative, it will simplify things a good bit on the Flux side.