Port storage handling + wrapper materialization from CUDA.jl. (#468)

Author: maleadt

lib/JLArrays/src/JLArrays.jl

@@ -5,7 +5,7 @@
 
 module JLArrays
 
-export JLArray, jl
+export JLArray, JLVector, JLMatrix, jl
 
 using GPUArrays
 
@@ -86,18 +86,27 @@ end
 # array type
 
 struct JLDeviceArray{T, N} <: AbstractDeviceArray{T, N}
-    data::Array{T, N}
+    data::Vector{UInt8}
+    offset::Int
     dims::Dims{N}
-
-    function JLDeviceArray{T,N}(data::Array{T, N}, dims::Dims{N}) where {T,N}
-        new(data, dims)
-    end
 end
 
+Base.elsize(::Type{<:JLDeviceArray{T}}) where {T} = sizeof(T)
+
 Base.size(x::JLDeviceArray) = x.dims
+Base.sizeof(x::JLDeviceArray) = Base.elsize(x) * length(x)
+
+Base.unsafe_convert(::Type{Ptr{T}}, x::JLDeviceArray{T}) where {T} =
+    Base.unsafe_convert(Ptr{T}, x.data) + x.offset*Base.elsize(x)
+
+# conversion of untyped data to a typed Array
+function typed_data(x::JLDeviceArray{T}) where {T}
+    unsafe_wrap(Array, pointer(x), x.dims)
+end
+
+@inline Base.getindex(A::JLDeviceArray, index::Integer) = getindex(typed_data(A), index)
+@inline Base.setindex!(A::JLDeviceArray, x, index::Integer) = setindex!(typed_data(A), x, index)
 
-@inline Base.getindex(A::JLDeviceArray, index::Integer) = getindex(A.data, index)
-@inline Base.setindex!(A::JLDeviceArray, x, index::Integer) = setindex!(A.data, x, index)
 
 # indexing
 
@@ -139,23 +148,60 @@ end
 # Host abstractions
 #
 
-struct JLArray{T, N} <: AbstractGPUArray{T, N}
-    data::Array{T, N}
+function check_eltype(T)
+  if !Base.allocatedinline(T)
+    explanation = explain_allocatedinline(T)
+    error("""
+      JLArray only supports element types that are allocated inline.
+      $explanation""")
+  end
+end
+
+mutable struct JLArray{T, N} <: AbstractGPUArray{T, N}
+    data::DataRef{Vector{UInt8}}
+
+    offset::Int        # offset of the data in the buffer, in number of elements
+
     dims::Dims{N}
 
-    function JLArray{T,N}(data::Array{T, N}, dims::Dims{N}) where {T,N}
-        isbitstype(T) || error("JLArray only supports bits types")
-        # when supporting isbits-union types, use `Base.allocatedinline` here.
-        new(data, dims)
+    # allocating constructor
+    function JLArray{T,N}(::UndefInitializer, dims::Dims{N}) where {T,N}
+        check_eltype(T)
+        maxsize = prod(dims) * sizeof(T)
+        data = Vector{UInt8}(undef, maxsize)
+        ref = DataRef(data)
+        obj = new{T,N}(ref, 0, dims)
+        finalizer(unsafe_free!, obj)
     end
+
+    # low-level constructor for wrapping existing data
+    function JLArray{T,N}(ref::DataRef{Vector{UInt8}}, dims::Dims{N};
+                          offset::Int=0) where {T,N}
+        check_eltype(T)
+        obj = new{T,N}(ref, offset, dims)
+        finalizer(unsafe_free!, obj)
+    end
+end
+
+unsafe_free!(a::JLArray) = GPUArrays.unsafe_free!(a.data)
+
+# conversion of untyped data to a typed Array
+function typed_data(x::JLArray{T}) where {T}
+    unsafe_wrap(Array, pointer(x), x.dims)
+end
+
+function GPUArrays.derive(::Type{T}, N::Int, a::JLArray, dims::Dims, offset::Int) where {T}
+    ref = copy(a.data)
+    offset = (a.offset * Base.elsize(a)) ÷ sizeof(T) + offset
+    JLArray{T,N}(ref, dims; offset)
 end
 
 
-## constructors
+## convenience constructors
 
-# type and dimensionality specified, accepting dims as tuples of Ints
-JLArray{T,N}(::UndefInitializer, dims::Dims{N}) where {T,N} =
-  JLArray{T,N}(Array{T, N}(undef, dims), dims)
+const JLVector{T} = JLArray{T,1}
+const JLMatrix{T} = JLArray{T,2}
+const JLVecOrMat{T} = Union{JLVector{T},JLMatrix{T}}
 
 # type and dimensionality specified, accepting dims as series of Ints
 JLArray{T,N}(::UndefInitializer, dims::Integer...) where {T,N} = JLArray{T,N}(undef, dims)
@@ -172,7 +218,10 @@ Base.similar(a::JLArray{T,N}) where {T,N} = JLArray{T,N}(undef, size(a))
 Base.similar(a::JLArray{T}, dims::Base.Dims{N}) where {T,N} = JLArray{T,N}(undef, dims)
 Base.similar(a::JLArray, ::Type{T}, dims::Base.Dims{N}) where {T,N} = JLArray{T,N}(undef, dims)
 
-Base.copy(a::JLArray{T,N}) where {T,N} = JLArray{T,N}(copy(a.data), size(a))
+function Base.copy(a::JLArray{T,N}) where {T,N}
+    b = similar(a)
+    @inbounds copyto!(b, a)
+end
 
 
 ## derived types
@@ -181,31 +230,26 @@ export DenseJLArray, DenseJLVector, DenseJLMatrix, DenseJLVecOrMat,
        StridedJLArray, StridedJLVector, StridedJLMatrix, StridedJLVecOrMat,
        AnyJLArray, AnyJLVector, AnyJLMatrix, AnyJLVecOrMat
 
-ContiguousSubJLArray{T,N,A<:JLArray} = Base.FastContiguousSubArray{T,N,A}
-
 # dense arrays: stored contiguously in memory
-DenseReinterpretJLArray{T,N,A<:Union{JLArray,ContiguousSubJLArray}} =
-    Base.ReinterpretArray{T,N,S,A} where S
-DenseReshapedJLArray{T,N,A<:Union{JLArray,ContiguousSubJLArray,DenseReinterpretJLArray}} =
-    Base.ReshapedArray{T,N,A}
-DenseSubJLArray{T,N,A<:Union{JLArray,DenseReshapedJLArray,DenseReinterpretJLArray}} =
-    Base.FastContiguousSubArray{T,N,A}
-DenseJLArray{T,N} = Union{JLArray{T,N}, DenseSubJLArray{T,N}, DenseReshapedJLArray{T,N},
-                          DenseReinterpretJLArray{T,N}}
+DenseJLArray{T,N} = JLArray{T,N}
 DenseJLVector{T} = DenseJLArray{T,1}
 DenseJLMatrix{T} = DenseJLArray{T,2}
 DenseJLVecOrMat{T} = Union{DenseJLVector{T}, DenseJLMatrix{T}}
 
 # strided arrays
-StridedSubJLArray{T,N,A<:Union{JLArray,DenseReshapedJLArray,DenseReinterpretJLArray},
-                  I<:Tuple{Vararg{Union{Base.RangeIndex, Base.ReshapedUnitRange,
-                                        Base.AbstractCartesianIndex}}}} = SubArray{T,N,A,I}
-StridedJLArray{T,N} = Union{JLArray{T,N}, StridedSubJLArray{T,N}, DenseReshapedJLArray{T,N},
-                            DenseReinterpretJLArray{T,N}}
+StridedSubJLArray{T,N,I<:Tuple{Vararg{Union{Base.RangeIndex, Base.ReshapedUnitRange,
+                                            Base.AbstractCartesianIndex}}}} =
+  SubArray{T,N,<:JLArray,I}
+StridedJLArray{T,N} = Union{JLArray{T,N}, StridedSubJLArray{T,N}}
 StridedJLVector{T} = StridedJLArray{T,1}
 StridedJLMatrix{T} = StridedJLArray{T,2}
 StridedJLVecOrMat{T} = Union{StridedJLVector{T}, StridedJLMatrix{T}}
 
+Base.pointer(x::StridedJLArray{T}) where {T} = Base.unsafe_convert(Ptr{T}, x)
+@inline function Base.pointer(x::StridedJLArray{T}, i::Integer) where T
+    Base.unsafe_convert(Ptr{T}, x) + Base._memory_offset(x, i)
+end
+
 # anything that's (secretly) backed by a JLArray
 AnyJLArray{T,N} = Union{JLArray{T,N}, WrappedArray{T,N,JLArray,JLArray{T,N}}}
 AnyJLVector{T} = AnyJLArray{T,1}
@@ -221,13 +265,16 @@ Base.size(x::JLArray) = x.dims
 Base.sizeof(x::JLArray) = Base.elsize(x) * length(x)
 
 Base.unsafe_convert(::Type{Ptr{T}}, x::JLArray{T}) where {T} =
-    Base.unsafe_convert(Ptr{T}, x.data)
+    Base.unsafe_convert(Ptr{T}, x.data[]) + x.offset*Base.elsize(x)
 
 
 ## interop with Julia arrays
 
-JLArray{T,N}(x::AbstractArray{<:Any,N}) where {T,N} =
-    JLArray{T,N}(convert(Array{T}, x), size(x))
+function JLArray{T,N}(xs::AbstractArray{<:Any,N}) where {T,N}
+    A = JLArray{T,N}(undef, size(xs))
+    copyto!(A, convert(Array{T}, xs))
+    return A
+end
 
 # underspecified constructors
 JLArray{T}(xs::AbstractArray{S,N}) where {T,N,S} = JLArray{T,N}(xs)
@@ -345,14 +392,15 @@ end
 GPUArrays.backend(::Type{<:JLArray}) = JLBackend()
 
 Adapt.adapt_storage(::Adaptor, x::JLArray{T,N}) where {T,N} =
-  JLDeviceArray{T,N}(x.data, x.dims)
+  JLDeviceArray{T,N}(x.data[], x.offset, x.dims)
 
 function GPUArrays.mapreducedim!(f, op, R::AnyJLArray, A::Union{AbstractArray,Broadcast.Broadcasted};
                                  init=nothing)
     if init !== nothing
         fill!(R, init)
     end
-    @allowscalar Base.reducedim!(op, R.data, map(f, A))
+    @allowscalar Base.reducedim!(op, typed_data(R), map(f, A))
+    R
 end
 
 end

src/host/abstractarray.jl

@@ -1,6 +1,98 @@
 # core definition of the AbstractGPUArray type
 
 
+# storage handling
+
+export DataRef
+
+# DataRef provides a helper class to manage the storage of an array.
+#
+# There's multiple reasons we don't just put the data directly in a GPUArray struct:
+# - to share data between multiple arrays, e.g., to create views;
+# - to be able to early-free data and release GC pressure.
+#
+# To support this, wrap the data in a DataRef instead, and use it with the following methods:
+# - `ref[]`: get the data;
+# - `copy(ref)`: create a new reference, increasing the reference count;
+# - `unsafe_free!(ref)`: decrease the reference count, and free the data if it reaches 0.
+#
+# The contained RefCounted struct should not be used directly.
+
+# shared, reference-counted state.
+mutable struct RefCounted{D}
+  obj::D
+  finalizer
+  count::Threads.Atomic{Int}
+end
+
+function retain(rc::RefCounted)
+    if rc.count[] == 0
+        throw(ArgumentError("Attempt to retain freed data."))
+    end
+    Threads.atomic_add!(rc.count, 1)
+    return
+end
+
+function release(rc::RefCounted, args...)
+    if rc.count[] == 0
+        throw(ArgumentError("Attempt to release freed data."))
+    end
+    refcount = Threads.atomic_add!(rc.count, -1)
+    if refcount == 1 && rc.finalizer !== nothing
+        rc.finalizer(rc.obj, args...)
+    end
+    return
+end
+
+function Base.getindex(rc::RefCounted)
+    if rc.count[] == 0
+        throw(ArgumentError("Attempt to use freed data."))
+    end
+    rc.obj
+end
+
+# per-object state, with a flag to indicate whether the object has been freed.
+# this is to support multiple calls to `unsafe_free!` on the same object,
+# while only lowering the referene count of the underlying data once.
+mutable struct DataRef{D}
+    rc::RefCounted{D}
+    freed::Bool
+end
+
+function DataRef(finalizer, data::D) where {D}
+    rc = RefCounted{D}(data, finalizer, Threads.Atomic{Int}(1))
+    DataRef{D}(rc, false)
+end
+DataRef(data; kwargs...) = DataRef(nothing, data; kwargs...)
+
+function Base.getindex(ref::DataRef)
+    if ref.freed
+        throw(ArgumentError("Attempt to use a freed reference."))
+    end
+    ref.rc[]
+end
+
+function Base.copy(ref::DataRef{D}) where {D}
+    if ref.freed
+        throw(ArgumentError("Attempt to copy a freed reference."))
+    end
+    retain(ref.rc)
+    return DataRef{D}(ref.rc, false)
+end
+
+function unsafe_free!(ref::DataRef, args...)
+    if ref.freed
+        # multiple frees *of the same object* are allowed.
+        # we should only ever call `release` once per object, though,
+        # as multiple releases of the underlying data is not allowed.
+        return
+    end
+    release(ref.rc, args...)
+    ref.freed = true
+    return
+end
+
+
 # input/output
 
 ## serialization