Add caching allocator interface

[pxl-th]

Since Julia's GC is not aware of GPU memory, in scenarios with lots of allocations we end up in either OOM situations or in excessively high memory usage.
Even though the program may require only fraction of it.

To help with GPU memory utilizaton in a program with repeating blocks of code, we can wrap those regions in a scope that will utilize caching allocator every time the program enters this scope during the execution.

For example, this is especially useful when training models, where you compute loss, gradients w.r.t. loss and perform in-place parameter update of the model.

cache = GPUArrays.AllocCache(CuArray)
model = ...
for i in 1:1000
    GPUArrays.@enable cache begin
        loss, grads = ...
        update!(optimizer, model, grads)
    end
end

Caching is done on: (ArrayType, current device, eltype, dims[, buffer type]).

Example

In the following example we apply caching allocator at every iteration of the for-loop.
Every iteration requires 8 GiB of gpu memory, without caching allocator
GC wouldn't be able to free arrays in time resulting in higher memory usage.
With caching allocator, memory usage stays at exactly 8 GiB.

After the loop, we free all cached memory if there's any.
Alternatively, it will be freed automatically when cache is collected by GC.

cache = GPUArrays.AllocCache(CuArray)
n = 1024^3
CUDA.@sync for i in 1:1000
    GPUArrays.@enable cache begin
        sin.(CUDA.rand(Float32, n))
    end
end
GPUArrays.unsafe_free!(cache)

Performance impact

Executing GaussianSplatting.jl benchmark (1k training iterations) on RX 7900XTX:

	Without caching allocator	With caching allocator
GPU memory utilization
Time	59.656476 seconds	46.365646 seconds

TODO

• Support for 1.10.
• Support bulk-freeing instead of caching.
• Add PR description.
• Documentation.
• Tests.

[maleadt]

Could you add some high-level design description to the PR?

As I mentioned on Slack, CUDA already has a caching allocator, so I'm not sure if for those back-ends this shouldn't boil down to basically batch-calling unsafe_free! at the end of each iteration, instead of actively caching arrays. Would be good to compare performance, if possible.

[pxl-th]

Could you add some high-level design description to the PR?

As I mentioned on Slack, CUDA already has a caching allocator, so I'm not sure if for those back-ends this shouldn't boil down to basically batch-calling unsafe_free! at the end of each iteration, instead of actively caching arrays. Would be good to compare performance, if possible.

Yeah, I'm planning to add both detailed PR description and documentation.
And to allow batch-freeing instead of caching the arrays (which can be just an option in the caching allocator).

[pxl-th]

@maleadt, I've updated the PR.
Also, I've added tests but they are not enabled right now, because no backend currently has the implementation merged (including JLArrays, because tests use released version of it).
However, they pass locally on my machines.

Let me know what you think.

[pxl-th]

One difference I've found between Julia 1.10 and Julia 1.11:

• Julia 1.10:

julia> GPUArrays.AllocCache.@enable CuArray :loop begin
           x1 = CuArray(rand(Float32, 1))
       end
1-element CuArray{Float32, 1, CUDA.DeviceMemory}:
 0.680597

julia> x1
ERROR: UndefVarError: `x1` not defined

• Julia 1.11:

julia> GPUArrays.AllocCache.@enable CuArray :loop begin
           x1 = CuArray(rand(Float32, 1))
       end
1-element CuArray{Float32, 1, CUDA.DeviceMemory}:
 0.7703809

julia> x1
1-element CuArray{Float32, 1, CUDA.DeviceMemory}:
 0.7703809

Not sure where is it coming from.

[github-actions]

Runic suggested the following formatting changes.

Done.

[maleadt]

Why is this needed, now that you switched to the array type? Isn't all information there for the caller to construct an appropriate allocator cache?

[pxl-th]

This is for the internal implementation to retrieve the actual cache for @enable.
E.g. when CUDA calls alloc! we retrieve its allocator cache based on its array type.
Otherwise the user would have to pass the cache itself to the macro, no?

[maleadt]

I mean that we can get rid of the alias, and replace calls to cache_allocator by AllocCache.PerDeviceCacheAllocator(AT). Just trying to minimize the interface to be implemented by back-ends.

[pxl-th]

I'm not opposed, just a bit confused.
AllocCache.PerDeviceCacheAllocator(AT) is call to ctor, when with cache_allocator we retrieve an instance.
In CUDA.jl we define a global variable that we then retrieve with cache_allocator.
How are we going to access this variable with AllocCache.PerDeviceCacheAllocator(AT) if its a call to ctor?

Or is the point just to rename cache_allocator to AllocCache.PerDeviceCacheAllocator? But then it is ambiguous because ctor has the same method signature.

[maleadt]

Is it expected for users to need this? Why not have them wrap code in multiple @enable blocks?

[pxl-th]

Yes, because the cache caches arrays based on their dims, there may be a situation where dims change (e.g. different batch size or the number of parameters of the model change) you need to invalidate the cache, because with new dims the old ones won't be retrieved.

E.g. with GaussianSplatting, where I enable cache for the training step.
But at some point the number of parameters of the model changes so we need to invalidate the cache, because old dims are not used anymore.

Why not have them wrap code in multiple @enable blocks?

IIUC, you mean something like this?

GPUArrays.AllocCache.@enable CuArray :train_step begin
    # some code
end

# some code outside of caching

GPUArrays.AllocCache.@enable CuArray :train_step begin
    # some code
end

But then again when you either no longer need the cache (done training) or the dims change you need to somehow invalidate it.

[maleadt]

Oh, so the cache persists until the user calls invalidate!? I somehow missed that. It seems like a dangerous design to me; if you forget to invalidate! on any path outside of the @enable, memory will leak?

[pxl-th]

Yes, although I think this is fine tradeoff (technically you can always call invalidate and free the memory). As a last resort we could invalidate the cache in alloc/retry mechanism, registering a hook. Similar to how we do with fft handle cache

[maleadt]

One difference I've found between Julia 1.10 and Julia 1.11

Hmm, that seems problematic. Macros should not introduce scope:

❯ jl +1.10
julia> @time begin
       x1 = []
       end
  0.000002 seconds (1 allocation: 48 bytes)
Any[]

julia> x1
Any[]

[pxl-th]

ScopedValues.jl on Julia 1.10 introduce a scope:

julia> using ScopedValues

julia> x = ScopedValue(1)
ScopedValue{Int64}(1)

julia> @with x => 2 begin
           x2 = x[]
           x3 = 1
       end
1

julia> x2
ERROR: UndefVarError: `x2` not defined

[maleadt]

Another fundamental question (sorry for stretching this out): Why do you even care about the array type in the @enable interface? Wouldn't it be better if the user didn't have to worry about this? The cache can be sharded internally so that back-ends can invalidate! only their portion (e.g. from within a retry_reclaim hook), but I don't see why it has to be provided.

Maybe the cache name should be optional as well. It could default to something derived from the current task's name, so that's it's really convenient to do:

AllocCache.@enable begin
  for i in epocs
    ...
  end
end
AllocCache.invalidate!()

Just spitballing here, you probably have a better view regarding it based on your experiments with it already.

---

Seeing the above written out, I wonder if a wholly different API wouldn't be much more idiomatic, reifing the now implicit stuff like the name of the cache:

cache = AllocCache()
cache() do
  for i in epocs
    ...
  end
end
empty!(cache)

A finalizer could then also empty the cache, avoiding the risk of leaking memory if you forget to empty! it.

[github-actions]

Runic suggested the following formatting changes.

[github-actions]

Runic suggested the following formatting changes.

[github-actions]

Runic suggested the following formatting changes.

[pxl-th]

Seeing the above written out, I wonder if a wholly different API wouldn't be much more idiomatic, reifing the now implicit stuff like the name of the cache:

@maleadt , I've updated the implementation based on this, see examples in PR description for TL;DR.
It will now also free all memory in the finalizer.

[github-actions]

Runic suggested the following formatting changes.

[github-actions]

Suggested change

	function AllocCache(::Type{T}) where {T <: AbstractGPUArray}
	cache = new{T}(
	ReentrantLock(),
	Dict{UInt64, Vector{T}}(),
	Dict{UInt64, Vector{T}}()
	)
	function AllocCache(::Type{T}) where {T <: AbstractGPUArray}
	cache = new{T}(
	ReentrantLock(),
	Dict{UInt64, Vector{T}}(),
	Dict{UInt64, Vector{T}}()
	)
	return finalizer(unsafe_free!, cache)
	end