More alloc cache improvements #583
Conversation
It's not safe to concurrently access a dict while it may be mutated.
|
Thanks! I'll take a look at it this evening |
|
|
||
| if !isempty(free_pool) | ||
| ref = Base.@lock cache.lock pop!(free_pool) | ||
| @assert !ref.freed |
There was a problem hiding this comment.
This essentially forbids user calling CUDA.unsafe_free!(x) on an array inside GPUArrays.@cached.
I imagined GPUArrays.@cached to be like a toggle, such that with and without it the code can be used without any modifications.
Even if putting aside users calling unsafe_free! directly I think it will introduce problems with things like with_workspace.
Why not cycle in a loop until we find a non-freed array as it was before?
There was a problem hiding this comment.
Ah, good catch.
Why not cycle in a loop until we find a non-freed array as it was before?
Because only the reference is freed, not the underlying data. I should find a way to still reuse that (but we can't simply toggle freed because then the GC may free it again).
There was a problem hiding this comment.
As a way to mitigate this, I've added a cached field to DataRef.
When arrays are created within @cached their DataRef will have cached set to true.
When cached=true, calling unsafe_free! on an array is a no-op.
When freeing the cache via unsafe_free! (or in a finalizer) it sets all DataRef cached field to false allowing all of the arrays to be freed.
All copies of DataRef (copies happen during creation of views) are marked as not cached.
Because AllocCache only stores DataRef that allocate, marking copies as cached (or transferring the property from the original DataRef) would prevent the refcount from going to 0 and will result in a memory leak.
There was a problem hiding this comment.
I've tested on GaussianSplatting.jl and HiFiGAN.jl and it works fine.
Let me know what you think of this approach.
There was a problem hiding this comment.
Hmm, I'm not a fan of coupling DataRef to the cache (which to me conceptually is a lower-level abstraction). I'll try to come up with a different approach Monday, or else go with your suggested one.
This allows calling `unsafe_free!` on arrays inside @cached. When doing so, `cached` arrays are not freed until they are no longer maintained by AllocCache (which is done via `unsafe_free!` on the cache or in a finalizer).
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
The manual refcount fiddling is a bit dangerous (it should be compatible with manual
unsafe_free!s and the GC), but I think is correct. @pxl-th Maybe test with one of your apps using this? Ideally, first adapt your back-end to key on bufsize instead of dims and remove T.