Make it possible to customize the CuIterator adaptor. (#2297)

Author: maleadt

src/iterator.jl

@@ -1,55 +1,69 @@
 export CuIterator
 
 """
-    CuIterator(batches)
+    CuIterator([to], batches)
 
-Return a `CuIterator` that can iterate through the provided `batches` via `Base.iterate`.
+Create a `CuIterator` that iterates through the provided `batches` via `iterate`. Upon each
+iteration, the current `batch` is copied to the GPU, and the previous iteration is marked as
+freeable from GPU memory (via `unsafe_free!`).
 
-Upon each iteration, the current `batch` is copied to the GPU,
-and the previous iteration is marked as freeable from GPU memory (via `unsafe_free!`).
-Both of these use `adapt`, so that each `batch` can be an array, an array of arrays,
-or a more complex object such as a nested set of NamedTuples, which is explored recursively.
+The conversion to GPU memory is done recursively, using Adapt.jl, so that each batch can be
+an array, an array of arrays, or more complex iterable objects. To customize the conversion,
+an adaptor can be specified as the first argument, e.g., to change the element type:
 
-This abstraction is useful for batching data into GPU memory in a manner that
-allows old iterations to potentially be freed (or marked as reusable) earlier
-than they otherwise would via CuArray's internal polling mechanism.
+```julia
+julia> first(CuIterator([[1.]]))
+1-element CuArray{Float64, 1, CUDA.Mem.DeviceBuffer}:
+ 1.0
+
+julia> first(CuIterator(CuArray{Float32}, [[1.]]))
+1-element CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}:
+ 1.0
+```
+
+This abstraction is useful for batching data into GPU memory in a manner that allows old
+iterations to potentially be freed (or marked as reusable) earlier than they otherwise would
+via `CuArray`'s internal polling mechanism.
 """
-mutable struct CuIterator{B}
+mutable struct CuIterator{T,B}
+    to::T
     batches::B
     previous::Any
-    CuIterator(batches) = new{typeof(batches)}(batches)
+
+    CuIterator(batches) = CuIterator(nothing, batches)
+    CuIterator(to, batches) = new{typeof(to),typeof(batches)}(to, batches)
 end
 
 function Base.iterate(c::CuIterator, state...)
     item = iterate(c.batches, state...)
-    isdefined(c, :previous) && adapt(CuIteratorFree, c.previous)
+    isdefined(c, :previous) && adapt(CuIteratorFree(), c.previous)
     item === nothing && return nothing
     batch, next_state = item
-    cubatch = adapt(CuIterator, batch)
+    cubatch = adapt(c, batch)
     c.previous = cubatch
     return cubatch, next_state
 end
 
-Base.IteratorSize(::Type{CuIterator{B}}) where {B} = Base.IteratorSize(B)
+Base.IteratorSize(::Type{CuIterator{T,B}}) where {T,B} = Base.IteratorSize(B)
 Base.length(c::CuIterator) = length(c.batches)  # required for HasLength
 Base.axes(c::CuIterator) = axes(c.batches)  # required for HasShape{N}
 
-Base.IteratorEltype(::Type{CuIterator{B}}) where {B} = Base.IteratorEltype(B)
+Base.IteratorEltype(::Type{CuIterator{T,B}}) where {T,B} = Base.IteratorEltype(B)
 Base.eltype(c::CuIterator) = eltype(c.batches)  # required for HasEltype
 
-# This struct exists to control adapt for clean-up-afterwards step:
-struct CuIteratorFree end
-Adapt.adapt_storage(::Type{CuIteratorFree}, x::CuArray) = unsafe_free!(x)
-
-# We re-purpose struct CuIterator for the matching transfer-before-use step,
-# mostly fall back to adapt(CuArray, x) which recurses into Tuples etc:
-Adapt.adapt_storage(::Type{<:CuIterator}, x) = adapt(CuArray, x)
-
-# But unlike adapt(CuArray, x), returse into arrays of arrays:
-function Adapt.adapt_storage(::Type{<:CuIterator}, xs::AbstractArray{T}) where T
-    isbitstype(T) ? adapt(CuArray, xs) : map(adapt(CuArray), xs)
+# adaptor for uploading
+Adapt.adapt_storage(c::CuIterator, x) = adapt(something(c.to, CuArray), x)
+## unlike adapt(CuArray, x), recurse into arrays of arrays
+function Adapt.adapt_storage(c::CuIterator, xs::AbstractArray{T}) where T
+    to = something(c.to, CuArray)
+    isbitstype(T) ? adapt(to, xs) : map(to, xs)
 end
-function Adapt.adapt_storage(::Type{CuIteratorFree}, xs::AbstractArray{T}) where T
-    foreach(adapt(CuIteratorFree), xs)
+
+# adaptor for clean-up
+struct CuIteratorFree end
+Adapt.adapt_storage(::CuIteratorFree, x::CuArray) = unsafe_free!(x)
+function Adapt.adapt_storage(::CuIteratorFree, xs::AbstractArray{T}) where T
+    foreach(adapt(CuIteratorFree()), xs)
     xs
 end
+

test/base/iterator.jl

@@ -37,6 +37,12 @@ batch2, _ = iterate(it_nt, state)
 it_vec = CuIterator([[i,i/2], [i/3, i/4]] for i in 1:4)
 @test first(it_vec)[1] isa CuArray{Float64}
 
+# test element type conversion using a custom adaptor
+it_float64 = CuIterator([[1.0]])
+@test first(it_float64) isa CuArray{Float64}
+it_float32 = CuIterator(CuArray{Float32}, [[1.0]])
+@test first(it_float32) isa CuArray{Float32}
+
 using StaticArrays: SVector, SA
 it_static = CuIterator([SA[i,i/2], SA[i/3, i/4]] for i in 1:4)
 @test first(it_static) isa CuArray{<:SVector}