Merge pull request #626 from JuliaGPU/tb/at_lock

Lock memory allocator operations

src/memory.jl

@@ -3,6 +3,12 @@
 using Printf
 using TimerOutputs
 
+using Base: @lock
+
+# global lock for shared resources (alloc stats, usage limits, etc).
+# each allocator needs to lock its own resources separately too.
+const memory_lock = ReentrantLock()
+
 
 ## allocation statistics
 
@@ -56,36 +62,50 @@ function actual_alloc(bytes)::Union{Nothing,CuPtr{Nothing}}
   end
 
   # try the actual allocation
-  try
-    alloc_stats.actual_time += Base.@elapsed begin
+  time, buf = try
+    time = Base.@elapsed begin
       @timeit_debug alloc_to "alloc" buf = Mem.alloc(Mem.Device, bytes)
     end
-    @assert sizeof(buf) == bytes
+    time, buf
+  catch err
+    (isa(err, CuError) && err.code == CUDAdrv.ERROR_OUT_OF_MEMORY) || rethrow()
+    return nothing
+  end
+  @assert sizeof(buf) == bytes
+  ptr = convert(CuPtr{Nothing}, buf)
+
+  # manage state
+  @lock memory_lock begin
+    alloc_stats.actual_time += time
     alloc_stats.actual_nalloc += 1
     alloc_stats.actual_alloc += bytes
     usage[] += bytes
-    ptr = convert(CuPtr{Nothing}, buf)
     @assert !haskey(allocated, ptr)
     allocated[ptr] = buf
-    return ptr
-  catch err
-    (isa(err, CuError) && err.code == CUDAdrv.ERROR_OUT_OF_MEMORY) || rethrow()
   end
 
-  return nothing
+  return ptr
 end
 
 function actual_free(ptr::CuPtr{Nothing})
-  buf = allocated[ptr]
-  delete!(allocated, ptr)
+  # look up the buffer
+  buf = @lock memory_lock begin
+    allocated[ptr]
+  end
   bytes = sizeof(buf)
 
-  alloc_stats.actual_nfree += 1
-  alloc_stats.actual_free += bytes
-  usage[] -= bytes
+  # free the memory
+  @timeit_debug alloc_to "free" begin
+    time = Base.@elapsed Mem.free(buf)
+  end
 
-  @timeit_debug alloc_to "free"  begin
-    alloc_stats.actual_time += Base.@elapsed Mem.free(buf)
+  # manage state
+  @lock memory_lock begin
+    alloc_stats.actual_time += time
+    alloc_stats.actual_nfree += 1
+    alloc_stats.actual_free += bytes
+    usage[] -= bytes
+    delete!(allocated, ptr)
   end
 
   return
@@ -155,18 +175,20 @@ a [`OutOfGPUMemoryError`](@ref) if the allocation request cannot be satisfied.
   # 0-byte allocations shouldn't hit the pool
   sz == 0 && return CU_NULL
 
-  alloc_stats.pool_time += Base.@elapsed begin
+  time = Base.@elapsed begin
     @pool_timeit "pooled alloc" ptr = pool[].alloc(sz)
   end
-  if ptr === nothing
-    throw(OutOfGPUMemoryError(sz))
+  ptr === nothing && throw(OutOfGPUMemoryError(sz))
+
+  # manage state
+  @lock memory_lock begin
+    alloc_stats.pool_time += time
+    alloc_stats.pool_nalloc += 1
+    alloc_stats.pool_alloc += sz
+    @assert !haskey(requested, ptr)
+    requested[ptr] = sz
   end
 
-  alloc_stats.pool_nalloc += 1
-  alloc_stats.pool_alloc += sz
-  @assert !haskey(requested, ptr)
-  requested[ptr] = sz
-
   return ptr
 end
 
@@ -179,15 +201,19 @@ Releases a buffer pointed to by `ptr` to the memory pool.
   # 0-byte allocations shouldn't hit the pool
   ptr == CU_NULL && return
 
-  @assert haskey(requested, ptr)
-  sz = requested[ptr]
-  delete!(requested, ptr)
-
-  alloc_stats.pool_nfree += 1
-  alloc_stats.pool_time += Base.@elapsed begin
+  time = Base.@elapsed begin
     @pool_timeit "pooled free" pool[].free(ptr)
   end
 
+  # manage state
+  @lock memory_lock begin
+    alloc_stats.pool_time += time
+    alloc_stats.pool_nfree += 1
+    @assert haskey(requested, ptr)
+    sz = requested[ptr]
+    delete!(requested, ptr)
+  end
+
   return
 end
 

src/memory/binned.jl

@@ -24,7 +24,8 @@ using ..CuArrays: @pool_timeit
 
 using CUDAdrv
 
-const pool_lock = ReentrantLock()
+# use a macro-version of Base.lock to avoid closures
+using Base: @lock
 
 
 ## tunables
@@ -64,6 +65,8 @@ end
 
 ## infrastructure
 
+const pool_lock = ReentrantLock()
+
 const pools_used = Vector{Set{Block}}()
 const pools_avail = Vector{Vector{Block}}()
 const allocated = Dict{CuPtr{Nothing},Block}()
@@ -79,7 +82,7 @@ function create_pools(idx)
     return
   end
 
-  lock(pool_lock) do
+  @lock pool_lock begin
     while length(pool_usage) < idx
       push!(pool_usage, 1)
       push!(pool_history, initial_usage)
@@ -263,7 +266,7 @@ function init()
     delay = MIN_DELAY
     @async begin
       while true
-        @pool_timeit "background task" lock(pool_lock) do
+        @pool_timeit "background task" @lock pool_lock begin
           if scan()
             delay = MIN_DELAY
           else
@@ -289,7 +292,7 @@ function alloc(bytes)
     @inbounds used = pools_used[pid]
     @inbounds avail = pools_avail[pid]
 
-    lock(pool_lock) do
+    @lock pool_lock begin
       block = pool_alloc(alloc_bytes, pid)
 
       if block !== nothing
@@ -327,7 +330,7 @@ function free(ptr)
     @inbounds used = pools_used[pid]
     @inbounds avail = pools_avail[pid]
 
-    lock(pool_lock) do
+    @lock pool_lock begin
       # mark the buffer as available
       delete!(used, block)
       push!(avail, block)