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()
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]

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.

[maleadt]

Looks good! Thanks for keeping up with my review requests.

CI failures look related though.

[github-actions]

Runic suggested the following formatting changes.

[github-actions]

Runic suggested the following formatting changes.

[maleadt]

Pushed a simplification to the back-end interface to avoid having to reach into the scoped value. The only cost is having to allocate the key tuple unconditionally, but I think that should be fine.

[pxl-th]

Looks good!

[maleadt]

The show method is broken:

julia> cache = GPUArrays.AllocCache(CuArray)
GPUArrays.AllocCache{CuArray}(Error showing value of type GPUArrays.AllocCache{CuArray}:
ERROR: TypeError: in typeassert, expected Int64, got a value of type UInt64

Would probably be useful if it showed some stats.

[pxl-th]

The show method is broken:

julia> cache = GPUArrays.AllocCache(CuArray)
GPUArrays.AllocCache{CuArray}(Error showing value of type GPUArrays.AllocCache{CuArray}:
ERROR: TypeError: in typeassert, expected Int64, got a value of type UInt64

Would probably be useful if it showed some stats.

Fixed.

[github-actions]

Runic suggested the following formatting changes.

[maleadt]

Alright, I think we arrived at something great. Thanks for the PR! Let's merge and tag.

[maleadt]

Just saw this:

      From worker 2:	error in running finalizer: ErrorException("Invalidating allocations cache that's currently in use. Invalidating inside `@cached` is not allowed.")
      From worker 2:	error at ./error.jl:35
      From worker 2:	macro expansion at /home/tim/Julia/pkg/GPUArrays/src/host/alloc_cache.jl:75 [inlined]
      From worker 2:	macro expansion at ./lock.jl:273 [inlined]
      From worker 2:	unsafe_free! at /home/tim/Julia/pkg/GPUArrays/src/host/alloc_cache.jl:73
      From worker 2:	unknown function (ip: 0x7385147b5df2)
      From worker 2:	run_finalizer at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/gc.c:299
      From worker 2:	jl_gc_run_finalizers_in_list at /cache/build/tester-amdci5-12/julialang/julia-release-1-dot-11/src/gc.c:391

@cached should probably root the cache until it completes?

[pxl-th]

Oh... right