Merge pull request #9 from darsnack/gpu-support

Allow any backing array type and add Adapt support

Author: findmyway

.github/workflows/ci.yml

@@ -10,7 +10,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.0'
+          # - '1.0'  # incompatible with CUDA.jl
           - '1'
           - 'nightly'
         os:

Project.toml

@@ -3,11 +3,15 @@ uuid = "9de3a189-e0c0-4e15-ba3b-b14b9fb0aec1"
 authors = ["Jun Tian <[email protected]> and contributors"]
 version = "0.1.5"
 
+[deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+
 [compat]
 julia = "1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 
 [targets]
-test = ["Test"]
+test = ["CUDA", "Test"]

src/CircularArrayBuffers.jl

@@ -1,21 +1,23 @@
 module CircularArrayBuffers
 
+using Adapt
+
 export CircularArrayBuffer, CircularVectorBuffer, capacity, isfull
 
 """
-    CircularArrayBuffer{T}(sz::Integer...) -> CircularArrayBuffer{T, N}
+    CircularArrayBuffer{T}(sz::Integer...) -> CircularArrayBuffer{T, N, Array{T, N}}
 
 `CircularArrayBuffer` uses a `N`-dimension `Array` of size `sz` to serve as a buffer for
 `N-1`-dimension `Array`s of the same size.
 """
-mutable struct CircularArrayBuffer{T,N} <: AbstractArray{T,N}
-    buffer::Array{T,N}
+mutable struct CircularArrayBuffer{T,N,S<:AbstractArray{T,N}} <: AbstractArray{T,N}
+    buffer::S
     first::Int
     nframes::Int
     step_size::Int
 end
 
-const CircularVectorBuffer{T} = CircularArrayBuffer{T,1}
+const CircularVectorBuffer{T,S} = CircularArrayBuffer{T,1,S}
 
 CircularVectorBuffer{T}(n::Integer) where {T} = CircularArrayBuffer{T}(n)
 
@@ -28,12 +30,25 @@ function CircularArrayBuffer(A::AbstractArray{T,N}) where {T,N}
     CircularArrayBuffer(A, 1, size(A, N), N == 1 ? 1 : *(size(A)[1:end-1]...))
 end
 
+Adapt.adapt_structure(to, cb::CircularArrayBuffer) =
+    CircularArrayBuffer(adapt(to, cb.buffer), cb.first, cb.nframes, cb.step_size)
+
+function Base.show(io::IO, ::MIME"text/plain", cb::CircularArrayBuffer{T}) where T
+    print(io, ndims(cb) == 1 ? "CircularVectorBuffer(" : "CircularArrayBuffer(")
+    Base.showarg(io, cb.buffer, false)
+    print(io, ") with eltype $T:\n")
+    Base.print_array(io, adapt(Array, cb))
+    return nothing
+end
+
 Base.IndexStyle(::CircularArrayBuffer) = IndexLinear()
 
 Base.size(cb::CircularArrayBuffer{T,N}, i::Integer) where {T,N} = i == N ? cb.nframes : size(cb.buffer, i)
 Base.size(cb::CircularArrayBuffer{T,N}) where {T,N} = ntuple(i -> size(cb, i), N)
 Base.getindex(cb::CircularArrayBuffer{T,N}, i::Int) where {T,N} = getindex(cb.buffer, _buffer_index(cb, i))
+Base.getindex(cb::CircularArrayBuffer{T,N}, I...) where {T,N} = getindex(cb.buffer, Base.front(I)..., _buffer_frame(cb, Base.last(I)))
 Base.setindex!(cb::CircularArrayBuffer{T,N}, v, i::Int) where {T,N} = setindex!(cb.buffer, v, _buffer_index(cb, i))
+Base.setindex!(cb::CircularArrayBuffer{T,N}, v, I...) where {T,N} = setindex!(cb.buffer, v, Base.front(I)..., _buffer_frame(cb, Base.last(I)))
 
 capacity(cb::CircularArrayBuffer{T,N}) where {T,N} = size(cb.buffer, N)
 isfull(cb::CircularArrayBuffer) = cb.nframes == capacity(cb)
@@ -47,6 +62,7 @@ Base.isempty(cb::CircularArrayBuffer) = cb.nframes == 0
         ind
     end
 end
+@inline _buffer_index(cb::CircularArrayBuffer, I::AbstractVector{<:Integer}) = map(Base.Fix1(_buffer_index, cb), I)
 
 @inline function _buffer_frame(cb::CircularArrayBuffer, i::Int)
     n = capacity(cb)
@@ -58,7 +74,7 @@ end
     end
 end
 
-_buffer_frame(cb::CircularArrayBuffer, I::Vector{Int}) = map(i -> _buffer_frame(cb, i), I)
+_buffer_frame(cb::CircularArrayBuffer, I::AbstractVector{<:Integer}) = map(i -> _buffer_frame(cb, i), I)
 
 function Base.empty!(cb::CircularArrayBuffer)
     cb.nframes = 0
@@ -72,13 +88,14 @@ function Base.push!(cb::CircularArrayBuffer{T,N}, data) where {T,N}
         cb.nframes += 1
     end
     if N == 1
+        i = _buffer_frame(cb, cb.nframes)
         if ndims(data) == 0
-            cb[cb.nframes] = data[]
+            cb.buffer[i:i] .= data[]
         else
-            cb[cb.nframes] = data
+            cb.buffer[i:i] .= data
         end
     else
-        cb[ntuple(_ -> (:), N - 1)..., cb.nframes] .= data
+        cb.buffer[ntuple(_ -> (:), N - 1)..., _buffer_frame(cb, cb.nframes)] .= data
     end
     cb
 end

test/runtests.jl

@@ -1,7 +1,10 @@
 using CircularArrayBuffers
 using Test
+using Adapt
+using CUDA
+CUDA.allowscalar(false)
 
-@testset "CircularArrayBuffers.jl" begin
+@testset "CircularArrayBuffers (Array)" begin
     A = ones(2, 2)
     C = ones(Float32, 2, 2)
 
@@ -165,3 +168,178 @@ using Test
         ]
     end
 end
+
+if CUDA.functional()
+    @testset "CircularArrayBuffers (CuArray)" begin
+        A = CUDA.ones(2, 2)
+        Ac = adapt(Array, A)
+        C = CUDA.ones(Float32, 2, 2)
+
+        @testset "Adapt" begin
+            X = CircularArrayBuffer(rand(2, 3))
+            Xc = adapt(CuArray, X)
+            @test Xc isa CircularArrayBuffer{Float64,2,<:CuArray}
+            @test adapt(Array, Xc) == X
+        end
+
+        # https://github.com/JuliaReinforcementLearning/ReinforcementLearning.jl/issues/551
+        @testset "1D with 0d data" begin
+            b = adapt(CuArray, CircularArrayBuffer{Int}(3))
+            CUDA.@allowscalar push!(b, CUDA.zeros(Int, ()))
+            @test length(b) == 1
+            @test CUDA.@allowscalar b[1] == 0
+        end
+
+        @testset "1D Int" begin
+            b = adapt(CuArray, CircularArrayBuffer{Int}(3))
+
+            @test eltype(b) == Int
+            @test capacity(b) == 3
+            @test isfull(b) == false
+            @test isempty(b) == true
+            @test length(b) == 0
+            @test size(b) == (0,)
+            # element must has the exact same length with the element of buffer
+            @test_throws Exception push!(b, [1, 2])
+
+            for x in 1:3
+                push!(b, x)
+            end
+
+            @test capacity(b) == 3
+            @test isfull(b) == true
+            @test length(b) == 3
+            @test size(b) == (3,)
+            # scalar indexing is not allowed
+            @test_throws ErrorException b[1]
+            @test_throws ErrorException b[end]
+            @test CUDA.@allowscalar b[1:end] == cu([1, 2, 3])
+
+            for x in 4:5
+                push!(b, x)
+            end
+
+            @test capacity(b) == 3
+            @test length(b) == 3
+            @test size(b) == (3,)
+            @test CUDA.@allowscalar b[1:end] == [3, 4, 5]
+
+            empty!(b)
+            @test isfull(b) == false
+            @test isempty(b) == true
+            @test length(b) == 0
+            @test size(b) == (0,)
+
+            push!(b, 6)
+            @test isfull(b) == false
+            @test isempty(b) == false
+            @test length(b) == 1
+            @test size(b) == (1,)
+            @test CUDA.@allowscalar b[1] == 6
+
+            push!(b, 7)
+            push!(b, 8)
+            @test isfull(b) == true
+            @test isempty(b) == false
+            @test length(b) == 3
+            @test size(b) == (3,)
+            @test CUDA.@allowscalar b[1:3] == cu([6, 7, 8])
+
+            push!(b, 9)
+            @test isfull(b) == true
+            @test isempty(b) == false
+            @test length(b) == 3
+            @test size(b) == (3,)
+            @test CUDA.@allowscalar b[1:3] == cu([7, 8, 9])
+
+            x = CUDA.@allowscalar pop!(b)
+            @test x == 9
+            @test length(b) == 2
+            @test CUDA.@allowscalar b[1:2] == cu([7, 8])
+
+            x = CUDA.@allowscalar popfirst!(b)
+            @test x == 7
+            @test length(b) == 1
+            @test CUDA.@allowscalar b[1] == 8
+
+            x = CUDA.@allowscalar pop!(b)
+            @test x == 8
+            @test length(b) == 0
+
+            @test_throws ArgumentError pop!(b)
+            @test_throws ArgumentError popfirst!(b)
+        end
+
+        @testset "2D Float64" begin
+            b = adapt(CuArray, CircularArrayBuffer{Float64}(2, 2, 3))
+
+            @test eltype(b) == Float64
+            @test capacity(b) == 3
+            @test isfull(b) == false
+            @test length(b) == 0
+            @test size(b) == (2, 2, 0)
+
+            for x in 1:3
+                push!(b, x * A)
+            end
+
+            @test capacity(b) == 3
+            @test isfull(b) == true
+            @test length(b) == 2 * 2 * 3
+            @test size(b) == (2, 2, 3)
+            for i in 1:3
+                @test b[:, :, i] == i * A
+            end
+            @test b[:, :, end] == 3 * A
+
+            for x in 4:5
+                push!(b, x * CUDA.ones(Float64, 2, 2))
+            end
+
+            @test capacity(b) == 3
+            @test length(b) == 2 * 2 * 3
+            @test size(b) == (2, 2, 3)
+            @test b[:, :, 1] == 3 * A
+            @test b[:, :, end] == 5 * A
+
+            # doing b == ... triggers scalar indexing
+            @test CUDA.@allowscalar b == cu(reshape([c for x in 3:5 for c in x * Ac], 2, 2, 3))
+
+            push!(b, 6 * CUDA.ones(Float32, 2, 2))
+            push!(b, 7 * CUDA.ones(Int, 2, 2))
+            @test CUDA.@allowscalar b == cu(reshape([c for x in 5:7 for c in x * Ac], 2, 2, 3))
+
+            x = pop!(b)
+            @test x == 7 * CUDA.ones(Float64, 2, 2)
+            @test CUDA.@allowscalar b == cu(reshape([c for x in 5:6 for c in x * Ac], 2, 2, 2))
+        end
+
+        @testset "append!" begin
+            b = adapt(CuArray, CircularArrayBuffer{Int}(2, 3))
+            append!(b, CUDA.zeros(2))
+            append!(b, 1:4)
+            @test CUDA.@allowscalar b == cu([
+                0 1 3
+                0 2 4
+            ])
+
+
+            b = adapt(CuArray, CircularArrayBuffer{Int}(2, 3))
+            for i in 1:5
+                push!(b, CUDA.fill(i, 2))
+            end
+            empty!(b)
+            append!(b, 1:4)
+            @test CUDA.@allowscalar b == cu([
+                1 3
+                2 4
+            ])
+
+            append!(b, 5:8)
+            @test CUDA.@allowscalar b == cu([
+                3 5 7
+                4 6 8
+            ])
+        end
+    end
+end