Adapt to JuliaLang/julia#38487.

Author: maleadt

src/pool.jl

@@ -24,39 +24,17 @@ end
 
 Base.unlock(nrl::NonReentrantLock) = unlock(nrl.rl)
 
-# the above lock is taken around code that might gc, which might reenter through finalizers.
-# avoid that by temporarily disabling finalizers running concurrently on this thread.
-enable_finalizers(on::Bool) = ccall(:jl_gc_enable_finalizers, Cvoid,
-                                    (Ptr{Cvoid}, Int32,), Core.getptls(), on)
-macro safe_lock(l, ex)
-  quote
-    temp = $(esc(l))
-    lock(temp)
-    enable_finalizers(false)
-    try
-      $(esc(ex))
-    finally
-      unlock(temp)
-      enable_finalizers(true)
-    end
-  end
-end
-
-# if we actually want to acquire these locks from a finalizer, we can't just wait on them
-# (which might cause a task switch). as the lock can only be taken by another thread that
-# should be running, and not a concurrent task we'd need to switch to, we can safely spin.
-macro safe_lock_spin(l, ex)
+# a safe way to acquire locks from finalizers, where we can't wait (which switches tasks)
+macro spinlock(l, ex)
   quote
     temp = $(esc(l))
     while !trylock(temp)
       # we can't yield here
     end
-    enable_finalizers(false) # retains compatibility with non-finalizer callers
     try
       $(esc(ex))
     finally
       unlock(temp)
-      enable_finalizers(true)
     end
   end
 end
@@ -299,7 +277,7 @@ a [`OutOfGPUMemoryError`](@ref) if the allocation request cannot be satisfied.
 
   # record the memory block
   ptr = pointer(block)
-  @safe_lock allocated_lock begin
+  @lock allocated_lock begin
       @assert !haskey(allocated[dev], ptr)
       allocated[dev][ptr] = block, 1
   end
@@ -337,7 +315,7 @@ multiple calls to `free` before this buffer is put back into the memory pool.
   dev = device()
 
   # look up the memory block
-  @safe_lock_spin allocated_lock begin
+  @spinlock allocated_lock begin
     block, refcount = allocated[dev][ptr]
     allocated[dev][ptr] = block, refcount+1
   end
@@ -365,7 +343,7 @@ Releases a buffer pointed to by `ptr` to the memory pool.
   # so perform our own error handling.
   try
     # look up the memory block, and bail out if its refcount isn't 1
-    block = @safe_lock_spin allocated_lock begin
+    block = @spinlock allocated_lock begin
         block, refcount = allocated[dev][ptr]
         if refcount == 1
           delete!(allocated[dev], ptr)
@@ -473,7 +451,7 @@ end
 
 ## utilities
 
-used_memory(dev=device()) = @safe_lock allocated_lock begin
+used_memory(dev=device()) = @lock allocated_lock begin
     mapreduce(sizeof, +, values(allocated[dev]); init=0)
 end
 

src/pool/binned.jl

@@ -85,7 +85,7 @@ end
 
 # repopulate the "available" pools from the list of freed blocks
 function pool_repopulate(dev)
-  blocks = @safe_lock freed_lock begin
+  blocks = @lock freed_lock begin
     isempty(freed[dev]) && return
     blocks = Set(freed[dev])
     empty!(freed[dev])
@@ -232,7 +232,7 @@ function pool_free(dev, block)
 
   # we don't do any work here to reduce pressure on the GC (spending time in finalizers)
   # and to simplify locking (preventing concurrent access during GC interventions)
-  @safe_lock_spin freed_lock begin
+  @spinlock freed_lock begin
     push!(freed[dev], block)
   end
 end
@@ -245,7 +245,7 @@ function pool_init()
 end
 
 function cached_memory(dev=device())
-  sz = @safe_lock freed_lock mapreduce(sizeof, +, freed[dev]; init=0)
+  sz = @lock freed_lock mapreduce(sizeof, +, freed[dev]; init=0)
   @lock pool_lock for (pid, pl) in enumerate(pools_avail[dev])
     bytes = poolsize(pid)
     sz += bytes * length(pl)

src/pool/simple.jl

@@ -105,7 +105,7 @@ end
 function pool_free(dev, block)
     # we don't do any work here to reduce pressure on the GC (spending time in finalizers)
     # and to simplify locking (preventing concurrent access during GC interventions)
-    @safe_lock_spin freed_lock begin
+    @spinlock freed_lock begin
         push!(freed[dev], block)
     end
 end
@@ -116,7 +116,7 @@ function pool_init()
 end
 
 function cached_memory(dev=device())
-    sz = @safe_lock freed_lock mapreduce(sizeof, +, freed[dev]; init=0)
+    sz = @lock freed_lock mapreduce(sizeof, +, freed[dev]; init=0)
     sz += @lock pool_lock mapreduce(sizeof, +, pool[dev]; init=0)
     return sz
 end

src/pool/split.jl

@@ -151,7 +151,7 @@ end
 
 # repopulate the pools from the list of freed blocks
 function pool_repopulate(dev)
-    blocks = @safe_lock freed_lock begin
+    blocks = @lock freed_lock begin
         isempty(freed[dev]) && return
         blocks = Set(freed[dev])
         empty!(freed[dev])
@@ -304,7 +304,7 @@ function pool_free(dev, block)
     # and to simplify locking (preventing concurrent access during GC interventions)
     block.state == ALLOCATED || error("Cannot free a $(block.state) block")
     block.state = FREED
-    @safe_lock_spin freed_lock begin
+    @spinlock freed_lock begin
         push!(freed[dev], block)
     end
 end
@@ -329,7 +329,7 @@ function pool_reclaim(dev, sz::Int=typemax(Int))
 end
 
 function cached_memory(dev=device())
-    sz = @safe_lock freed_lock mapreduce(sizeof, +, freed[dev]; init=0)
+    sz = @lock freed_lock mapreduce(sizeof, +, freed[dev]; init=0)
     @lock pool_lock for pool in (pool_small[dev], pool_large[dev], pool_huge[dev])
         sz += mapreduce(sizeof, +, pool; init=0)
     end