Use adapt symmetrically in CuIterator (#1769)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: mcabbott

src/iterator.jl

@@ -5,8 +5,10 @@ export CuIterator
 
 Return a `CuIterator` that can iterate through the provided `batches` via `Base.iterate`.
 
-Upon each iteration, the current `batch` is adapted to the GPU (via `map(x -> adapt(CuArray, x), batch)`)
+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.
 
 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
@@ -20,10 +22,10 @@ end
 
 function Base.iterate(c::CuIterator, state...)
     item = iterate(c.batches, state...)
-    isdefined(c, :previous) && foreach(unsafe_free!, c.previous)
+    isdefined(c, :previous) && adapt(CuIteratorFree, c.previous)
     item === nothing && return nothing
     batch, next_state = item
-    cubatch = map(x -> adapt(CuArray, x), batch)
+    cubatch = adapt(CuIterator, batch)
     c.previous = cubatch
     return cubatch, next_state
 end
@@ -35,3 +37,19 @@ Base.axes(c::CuIterator) = axes(c.batches)  # required for HasShape{N}
 Base.IteratorEltype(::Type{CuIterator{B}}) where {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)
+end
+function Adapt.adapt_storage(::Type{CuIteratorFree}, xs::AbstractArray{T}) where T
+    foreach(adapt(CuIteratorFree), xs)
+    xs
+end

test/iterator.jl

@@ -25,3 +25,18 @@ end
 @test eltype(cubatches) == eltype(batch for batch in batches) == Any
 @test Base.IteratorEltype(typeof(CuIterator(batches))) isa Base.HasEltype
 @test eltype(CuIterator(batches)) == eltype(batches)  # Vector
+
+it_nt = CuIterator((x=Float32[i,i/2], y=i) for i in 1:4)
+@test first(it_nt).x isa CuArray{Float32}
+batch1, state = iterate(it_nt)
+@test batch1.x == cu([1,1/2])
+batch2, _ = iterate(it_nt, state)
+@test batch2.x == cu([2,2/2])
+@test batch1.x.storage === nothing  # unsafe_free! has worked inside
+
+it_vec = CuIterator([[i,i/2], [i/3, i/4]] for i in 1:4)
+@test first(it_vec)[1] isa CuArray{Float64}
+
+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}