Simplify back-end interface.

Author: maleadt

lib/JLArrays/src/JLArrays.jl

@@ -89,21 +89,15 @@ mutable struct JLArray{T, N} <: AbstractGPUArray{T, N}
         check_eltype(T)
         maxsize = prod(dims) * sizeof(T)
 
-        function _alloc_f()
+        GPUArrays.cached_alloc((JLArray, T, dims)) do
             data = Vector{UInt8}(undef, maxsize)
             ref = DataRef(data) do data
                 resize!(data, 0)
             end
             obj = new{T, N}(ref, 0, dims)
-            return finalizer(unsafe_free!, obj)
-        end
-
-        cache = GPUArrays.ALLOC_CACHE[]
-        return if cache ≡ nothing
-            _alloc_f()
-        else
-            GPUArrays.alloc!(_alloc_f, cache, (JLArray, T, dims))::JLArray{T, N}
-        end
+            finalizer(unsafe_free!, obj)
+            return obj
+        end::JLArray{T,N}
     end
 
     # low-level constructor for wrapping existing data
@@ -112,6 +106,7 @@ mutable struct JLArray{T, N} <: AbstractGPUArray{T, N}
         check_eltype(T)
         obj = new{T,N}(ref, offset, dims)
         finalizer(unsafe_free!, obj)
+        return obj
     end
 end
 

src/host/alloc_cache.jl

@@ -30,13 +30,18 @@ function get_pool!(cache::AllocCache{T}, pool::Symbol, uid::UInt64) where {T <:
     return uid_pool
 end
 
-function alloc!(alloc_f, cache::AllocCache, key)
+function cached_alloc(f, key)
+    cache = ALLOC_CACHE[]
+    if cache === nothing
+        return f()
+    end
+
     x = nothing
     uid = hash(key)
 
     busy_pool = get_pool!(cache, :busy, uid)
     free_pool = get_pool!(cache, :free, uid)
-    isempty(free_pool) && (x = alloc_f())
+    isempty(free_pool) && (x = f())
 
     while !isempty(free_pool) && x ≡ nothing
         tmp = Base.@lock cache.lock pop!(free_pool)
@@ -45,7 +50,7 @@ function alloc!(alloc_f, cache::AllocCache, key)
         x = tmp
     end
 
-    x ≡ nothing && (x = alloc_f())
+    x ≡ nothing && (x = f())
     Base.@lock cache.lock push!(busy_pool, x)
     return x
 end