Support GPUArrays allocations cache

[pxl-th]

Ref: JuliaGPU/GPUArrays.jl#576

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 61.04%. Comparing base (234ec4c) to head (c4c3639).
Report is 3 commits behind head on master.

Additional details and impacted files

@@             Coverage Diff             @@
##           master    #2593       +/-   ##
===========================================
- Coverage   73.64%   61.04%   -12.60%     
===========================================
  Files         157      157               
  Lines       15220    14955      -265     
===========================================
- Hits        11209     9130     -2079     
- Misses       4011     5825     +1814     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[maleadt]

It's also kinda weird that we're always attaching a finalizer to the array, even when it's put inside of cache (meaning it should get freed when the cache goes out of scope). Should we move the finalizer into cached_alloc? That starts to make it a slightly messy API, though.

[pxl-th]

It's also kinda weird that we're always attaching a finalizer to the array, even when it's put inside of cache (meaning it should get freed when the cache goes out of scope). Should we move the finalizer into cached_alloc? That starts to make it a slightly messy API, though.

But we need them anyway, so it might be fine to attach them in ctor and have a single place for this, rather than catching when they go out of cache.

[pxl-th]

Passing bytesize as a key results in errors, where arrays are of the same bytesize, but of different dims.

julia> GPUArrays.@cached cache ROCArray(zeros(Float32, 16));

julia> GPUArrays.@cached cache ROCArray(zeros(Float32, 4, 4))
ERROR: TypeError: in typeassert, expected ROCArray{Float32, 2, AMDGPU.Runtime.Mem.HIPBuffer}, got a value of type ROCArray{Float32, 1, AMDGPU.Runtime.Mem.HIPBuffer}
Stacktrace:
 [1] ROCArray
   @ ~/.julia/dev/AMDGPU/src/array.jl:9 [inlined]
 [2] ROCArray
   @ ~/.julia/dev/AMDGPU/src/array.jl:111 [inlined]
 [3] ROCArray
   @ ~/.julia/dev/AMDGPU/src/array.jl:116 [inlined]
 [4] ROCArray(A::Matrix{Float32})
   @ AMDGPU ~/.julia/dev/AMDGPU/src/array.jl:119
 [5] macro expansion
   @ ~/.julia/packages/GPUArrays/2FKRj/src/host/alloc_cache.jl:147 [inlined]
 [6] top-level scope
   @ REPL[4]:1

To fix this we can either revert to dims or reshape an array post GPUArrays.cached_alloc and only then do type assertion:

return if size(x) != dims
    reshape(x, dims)
else
    x
end::ROCArray{T, N, B}

[maleadt]

julia> GPUArrays.@cached cache ROCArray(zeros(Float32, 16));

julia> GPUArrays.@cached cache ROCArray(zeros(Float32, 4, 4))
ERROR: TypeError: in typeassert, expected ROCArray{Float32, 2, AMDGPU.Runtime.Mem.HIPBuffer}, got a value of type ROCArray{Float32, 1, AMDGPU.Runtime.Mem.HIPBuffer}

Yes, that's what I meant in #2593 (comment), and I have reverted the change before merging this PR already.

[maleadt]

That said, it would be nice if the cache could fulfill such requests. I considered storing the DataRef instead, which tracks a refcount (so we could check if it has been unsafe_freed!d). However, that does make finalization more tricky: Since the cache doesn't keep the array objects alive anymore, that means array finalizers can now affect the refcount during execution of the @cached expressions. I guess we could move registration of the finalizer to cached_alloc, skipping that when the cache is active (because the cache is now responsible for freeing contained data)...