Fix and test the legacy memory pool. (#2402)

[skip julia]
[skip cuda]
[skip downstream]
[skip subpackages]

.buildkite/pipeline.yml

@@ -280,6 +280,25 @@ steps:
  if: build.message !~ /\[skip tests\]/ && build.message !~ /\[skip special\]/ && !build.pull_request.draft
  timeout_in_minutes: 30
 
+ - label: "Legacy memory allocator"
+ plugins:
+ - JuliaCI/julia#v1:
+ version: "1.10"
+ - JuliaCI/julia-test#v1:
+ test_args: "--quickfail core base"
+ - JuliaCI/julia-coverage#v1:
+ dirs:
+ - src
+ - lib
+ - examples
+ agents:
+ queue: "juliagpu"
+ cuda: "*"
+ env:
+ JULIA_CUDA_MEMORY_POOL: 'none'
+ if: build.message !~ /\[skip tests\]/ && build.message !~ /\[skip special\]/ && !build.pull_request.draft
+ timeout_in_minutes: 30
+
  - label: "CuArray with {{matrix.memory}} memory"
  plugins:
  - JuliaCI/julia#v1:

Project.toml

@@ -1,6 +1,6 @@
 name = "CUDA"
 uuid = "052768ef-5323-5732-b1bb-66c8b64840ba"
-version = "5.4.1"
+version = "5.4.2"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/memory.jl

@@ -627,28 +627,36 @@ end
 @inline function _pool_alloc(::Type{DeviceMemory}, sz)
  state = active_state()
 
- if stream_ordered(state.device)
+ mem = if stream_ordered(state.device)
  pool_mark!(state.device, true)
  pool = pool_create(state.device)
- end
 
- mem = let pool = pool # closure capture bug
  retry_reclaim(isnothing) do
  memory_limit_exceeded(sz) && return nothing
 
  # try the actual allocation
  try
- if stream_ordered(state.device)
- alloc(DeviceMemory, sz; async=true, state.stream, pool)
- else
- alloc(DeviceMemory, sz; async=false)
- end
+ alloc(DeviceMemory, sz; async=true, state.stream, pool)
+ catch err
+ isa(err, OutOfGPUMemoryError) || rethrow()
+ return nothing
+ end
+ end
+ else
+ retry_reclaim(isnothing) do
+ memory_limit_exceeded(sz) && return nothing
+
+ # try the actual allocation
+ try
+ alloc(DeviceMemory, sz; async=false)
  catch err
  isa(err, OutOfGPUMemoryError) || rethrow()
  return nothing
  end
  end
  end
+ # NOTE: the `retry_reclaim` body is duplicated to work around
+ # closure capture issues with the `pool` variable
  mem === nothing && throw(OutOfGPUMemoryError(sz))
 
  account!(memory_stats(state.device), sz)