Merge pull request #504 from JuliaGPU/tb/reclaim

API for manually reclaiming memory

Author: maleadt

.gitlab-ci.yml

@@ -111,7 +111,7 @@ flux:
   tags:
     - nvidia
   script:
-    - export FLUX="$HOME/.julia/dev/Flux"
+    - export FLUX=".julia/dev/Flux"
     - julia -e 'using Pkg;
                 Pkg.develop("Flux");'
     - julia --project -e 'using Pkg;

Manifest.toml

@@ -50,12 +50,6 @@ git-tree-sha1 = "ed2c4abadf84c53d9e58510b5fc48912c2336fbb"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
 version = "2.2.0"
 
-[[Crayons]]
-deps = ["Test"]
-git-tree-sha1 = "f621b8ef51fd2004c7cf157ea47f027fdeac5523"
-uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
-version = "4.0.0"
-
 [[DataStructures]]
 deps = ["InteractiveUtils", "OrderedCollections"]
 git-tree-sha1 = "0809951a1774dc724da22d26e4289bbaab77809a"
@@ -176,10 +170,10 @@ deps = ["Distributed", "InteractiveUtils", "Logging", "Random"]
 uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [[TimerOutputs]]
-deps = ["Crayons", "Printf", "Test", "Unicode"]
-git-tree-sha1 = "b80671c06f8f8bae08c55d67b5ce292c5ae2660c"
+deps = ["Printf"]
+git-tree-sha1 = "311765af81bbb48d7bad01fb016d9c328c6ede03"
 uuid = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
-version = "0.5.0"
+version = "0.5.3"
 
 [[Tokenize]]
 git-tree-sha1 = "dfcdbbfb2d0370716c815cbd6f8a364efb6f42cf"

src/indexing.jl

@@ -192,7 +192,7 @@ function Base.findfirst(vals::CuArray, xs::CuArray)
 
         # blocks to cover the remaining dimensions
         dev = CUDAdrv.device(kernel.fun.mod.ctx)
-        max_other_blocks = attribute(dev, CUDAdrv.MAX_GRID_DIM_Y)
+        max_other_blocks = attribute(dev, CUDAdrv.DEVICE_ATTRIBUTE_MAX_GRID_DIM_Y)
         blocks_other = (min(length(Rother), max_other_blocks),
                         cld(length(Rother), max_other_blocks))
 

src/memory.jl

@@ -86,7 +86,7 @@ function actual_free(ptr::CuPtr{Nothing})
 end
 
 
-## memory pool
+## memory pools
 
 const pool_to = TimerOutput()
 
@@ -96,11 +96,11 @@ end
 
 pool_timings() = (show(pool_to; allocations=false, sortby=:name); println())
 
-# API:
+# pool API:
 # - init()
-# - deinit()
 # - alloc(sz)::CuPtr{Nothing}
 # - free(::CuPtr{Nothing})
+# - reclaim(nb::Int=typemax(Int))::Int
 # - used_memory()
 # - cached_memory()
 
@@ -111,8 +111,22 @@ include("memory/dummy.jl")
 
 const pool = Ref{Module}(BinnedPool)
 
+
+## interface
+
+export OutOfGPUMemoryError
+
 const requested = Dict{CuPtr{Nothing},Int}()
 
+struct OutOfGPUMemoryError <: Exception
+  sz::Int
+end
+
+function Base.showerror(io::IO, err::OutOfGPUMemoryError)
+    println(io, "Out of GPU memory trying to allocate $(Base.format_bytes(err.sz))")
+    memory_status(io)
+end
+
 @inline function alloc(sz)::CuPtr{Nothing}
   # 0-byte allocations shouldn't hit the pool
   sz == 0 && return CU_NULL
@@ -121,10 +135,7 @@ const requested = Dict{CuPtr{Nothing},Int}()
     @pool_timeit "pooled alloc" ptr = pool[].alloc(sz)
   end
   if ptr === nothing
-    @error "Out of GPU memory trying to allocate $(Base.format_bytes(sz))"
-    pool[].dump()
-    memory_status()
-    throw(OutOfMemoryError())
+    throw(OutOfGPUMemoryError(sz))
   end
 
   alloc_stats.pool_nalloc += 1
@@ -151,7 +162,7 @@ end
   return
 end
 
-pool_dump() = pool[].dump()
+reclaim(sz::Int=typemax(Int)) = pool[].reclaim(sz)
 
 
 ## utilities
@@ -238,35 +249,35 @@ macro time(ex)
     end
 end
 
-function memory_status()
+function memory_status(io::IO=stdout)
   free_bytes, total_bytes = CUDAdrv.Mem.info()
   used_bytes = total_bytes - free_bytes
   used_ratio = used_bytes / total_bytes
 
-  @printf("Effective GPU memory usage: %.2f%% (%s/%s)\n",
-          100*used_ratio, Base.format_bytes(used_bytes),
-          Base.format_bytes(total_bytes))
+  @printf(io, "Effective GPU memory usage: %.2f%% (%s/%s)\n",
+              100*used_ratio, Base.format_bytes(used_bytes),
+              Base.format_bytes(total_bytes))
 
-  @printf("CuArrays GPU memory usage: %s", Base.format_bytes(usage[]))
+  @printf(io, "CuArrays GPU memory usage: %s", Base.format_bytes(usage[]))
   if usage_limit[] !== nothing
-    @printf(" (capped at %s)", Base.format_bytes(usage_limit[]))
+    @printf(io, " (capped at %s)", Base.format_bytes(usage_limit[]))
   end
-  println()
+  println(io)
 
   alloc_used_bytes = pool[].used_memory()
   alloc_cached_bytes = pool[].cached_memory()
   alloc_total_bytes = alloc_used_bytes + alloc_cached_bytes
 
-  @printf("%s usage: %s (%s allocated, %s cached)\n", nameof(pool[]),
-          Base.format_bytes(alloc_total_bytes), Base.format_bytes(alloc_used_bytes),
-          Base.format_bytes(alloc_cached_bytes))
+  @printf(io, "%s usage: %s (%s allocated, %s cached)\n", nameof(pool[]),
+              Base.format_bytes(alloc_total_bytes), Base.format_bytes(alloc_used_bytes),
+              Base.format_bytes(alloc_cached_bytes))
 
   requested_bytes = reduce(+, values(requested); init=0)
 
-  @printf("%s efficiency: %.2f%% (%s requested, %s allocated)\n", nameof(pool[]),
-          100*requested_bytes/usage[],
-          Base.format_bytes(requested_bytes),
-          Base.format_bytes(usage[]))
+  @printf(io, "%s efficiency: %.2f%% (%s requested, %s allocated)\n", nameof(pool[]),
+              100*requested_bytes/usage[],
+              Base.format_bytes(requested_bytes),
+              Base.format_bytes(usage[]))
 
   # check if the memory usage as counted by the CUDA allocator wrapper
   # matches what is reported by the pool implementation
@@ -276,6 +287,49 @@ function memory_status()
   end
 end
 
+"""
+  extalloc(f::Function; check::Function=isa(OutOfGPUMemoryError), nb::Integer=typemax(Int))
+
+Run a function `f` repeatedly until it successfully allocates the memory it needs. Only
+out-of-memory exceptions that pass `check` are considered for retry; this defaults to
+checking for the CuArrays out-of-memory exception but should be customized as to detect how
+an out-of-memory situation is reported by the function `f`. The argument `nb` indicates how
+many bytes of memory `f` requires, and serves as a hint for how much memory to reclaim
+before trying `f` again.
+"""
+function extalloc(f::Function; check::Function=ex->isa(ex,OutOfGPUMemoryError), nb::Integer=typemax(Int))
+  phase = 0
+  while true
+    phase += 1
+    return try
+      f()
+    catch ex
+      check(ex) || rethrow()
+
+      # incrementally costly reclaim of more and more memory
+      if phase == 1
+        reclaim(nb)
+      elseif phase == 2
+        GC.gc(false)
+        reclaim(nb)
+      elseif phase == 3
+        GC.gc(true)
+        reclaim(nb)
+      elseif phase == 4
+        # maybe the user lied, so try reclaiming all memory
+        GC.gc(true)
+        reclaim()
+      else
+        # give up
+        rethrow()
+      end
+
+      # try again
+      continue
+    end
+  end
+end
+
 
 ## init
 

src/memory/binned.jl

@@ -17,6 +17,8 @@ module BinnedPool
 #                       or just use unified memory for all allocations.
 # - per-device pools
 
+# TODO: move the management thread one level up, to be shared by all allocators
+
 using ..CuArrays
 using ..CuArrays: @pool_timeit
 
@@ -132,7 +134,7 @@ function scan()
 end
 
 # reclaim unused buffers
-function reclaim(full::Bool=false, target_bytes::Int=typemax(Int))
+function reclaim(target_bytes::Int=typemax(Int); full::Bool=true)
   # find inactive buffers
   @pool_timeit "scan" begin
     pools_inactive = Vector{Int}(undef, length(pools_avail)) # pid => buffers that can be freed
@@ -161,6 +163,7 @@ function reclaim(full::Bool=false, target_bytes::Int=typemax(Int))
 
   # reclaim buffers (in reverse, to discard largest buffers first)
   @pool_timeit "reclaim" begin
+    freed = 0
     for pid in reverse(eachindex(pools_inactive))
       bytes = poolsize(pid)
       avail = pools_avail[pid]
@@ -172,13 +175,14 @@ function reclaim(full::Bool=false, target_bytes::Int=typemax(Int))
 
         actual_free(block)
 
-        target_bytes -= bytes
-        target_bytes <= 0 && return true
+        freed += bytes
+        if freed >= target_bytes
+          return freed
+        end
       end
     end
+    return freed
   end
-
-  return false
 end
 
 
@@ -208,7 +212,7 @@ function pool_alloc(bytes, pid=-1)
   #       would require proper block splitting + compaction to be any efficient.
 
   @pool_timeit "3. reclaim unused" begin
-    reclaim(true, bytes)
+    reclaim(bytes)
   end
 
   @pool_timeit "4. try alloc" begin
@@ -226,7 +230,7 @@ function pool_alloc(bytes, pid=-1)
   end
 
   @pool_timeit "6. reclaim unused" begin
-    reclaim(true, bytes)
+    reclaim(bytes)
   end
 
   @pool_timeit "7. try alloc" begin
@@ -236,7 +240,7 @@ function pool_alloc(bytes, pid=-1)
   end
 
   @pool_timeit "8. reclaim everything" begin
-    reclaim(true)
+    reclaim()
   end
 
   @pool_timeit "9. try alloc" begin
@@ -266,7 +270,7 @@ function init()
             delay = min(delay*2, MAX_DELAY)
           end
 
-          reclaim()
+          reclaim(full=false)
         end
 
         sleep(delay)
@@ -275,13 +279,6 @@ function init()
   end
 end
 
-function deinit()
-    @assert sum(length, pools_used) == 0 "Cannot deinitialize memory pool with outstanding allocations"
-    reclam(fulll)
-
-    return
-end
-
 function alloc(bytes)
   # only manage small allocations in the pool
   if bytes <= MAX_POOL

src/memory/dummy.jl

@@ -8,8 +8,6 @@ using CUDAdrv
 
 init() = return
 
-deinit() = @assert isempty(allocated) "Cannot deinitialize memory pool with outstanding allocations"
-
 const allocated = Dict{CuPtr{Nothing},Int}()
 
 function alloc(sz)
@@ -42,6 +40,8 @@ function free(ptr)
     return
 end
 
+reclaim(target_bytes::Int=typemax(Int)) = return 0
+
 used_memory() = isempty(allocated) ? 0 : sum(sizeof, values(allocated))
 
 cached_memory() = 0

src/memory/simple.jl

@@ -60,14 +60,13 @@ function scan(sz)
     return
 end
 
-function reclaim(sz)
+function reclaim(sz::Int=typemax(Int))
     freed = 0
     while freed < sz && !isempty(available)
         block = pop!(available)
-        actual_free(block)
         freed += sizeof(block)
+        actual_free(block)
     end
-
     return freed
 end
 
@@ -110,17 +109,6 @@ end
 
 init() = return
 
-function deinit()
-    @assert isempty(allocated) "Cannot deinitialize memory pool with outstanding allocations"
-
-    for block in available
-        actual_free(block)
-    end
-    empty!(available)
-
-    return
-end
-
 function alloc(sz)
     block = pool_alloc(sz)
     if block !== nothing

src/memory/split.jl

@@ -402,23 +402,6 @@ end
 
 init() = return
 
-function deinit()
-    @assert isempty(allocated) "Cannot deinitialize memory pool with outstanding allocations"
-
-    repopulate(freed)
-    incremental_compact!(Set(freed))
-    empty!(freed)
-
-    for available in (available_small, available_large, available_huge)
-        while !isempty(available)
-            block = pop!(available)
-            actual_free(block)
-        end
-    end
-
-    return
-end
-
 function alloc(sz)
     block = pool_alloc(sz)
     if block !== nothing
@@ -440,6 +423,15 @@ function free(ptr)
     return
 end
 
+function reclaim(sz::Int=typemax(Int))
+    freed = 0
+    for available in (available_huge, available_large, available_small)
+        freed >= sz && break
+        freed += reclaim!(available, sz-freed)
+    end
+    return freed
+end
+
 used_memory() = mapreduce(sizeof, +, values(allocated); init=0)
 
 cached_memory() = mapreduce(sizeof, +, union(available_small, available_large, available_huge); init=0)