Merge pull request #745 from JuliaGPU/tb/alloc_opt

Use the default memory pool.

Author: maleadt

lib/cudadrv/pool.jl

@@ -1,30 +1,30 @@
 # Stream-orderdered memory allocator
 
-export CuMemoryPool, trim
+export CuMemoryPool, default_memory_pool, memory_pool, memory_pool!, trim
 
 mutable struct CuMemoryPool
     handle::CUmemoryPool
     ctx::CuContext
+end
 
-    function CuMemoryPool(dev::CuDevice)
-        props = Ref(CUmemPoolProps(
-            CU_MEM_ALLOCATION_TYPE_PINNED,
-            CU_MEM_HANDLE_TYPE_NONE,
-            CUmemLocation(
-                CU_MEM_LOCATION_TYPE_DEVICE,
-                deviceid(dev)
-            ),
-            C_NULL,
-            ntuple(i->Cuchar(0), 64)
-        ))
-        handle_ref = Ref{CUmemoryPool}()
-        cuMemPoolCreate(handle_ref, props)
-
-        ctx = CuCurrentContext()
-        obj = new(handle_ref[], ctx)
-        finalizer(unsafe_destroy!, obj)
-        return obj
-    end
+function CuMemoryPool(dev::CuDevice)
+    props = Ref(CUmemPoolProps(
+        CU_MEM_ALLOCATION_TYPE_PINNED,
+        CU_MEM_HANDLE_TYPE_NONE,
+        CUmemLocation(
+            CU_MEM_LOCATION_TYPE_DEVICE,
+            deviceid(dev)
+        ),
+        C_NULL,
+        ntuple(i->Cuchar(0), 64)
+    ))
+    handle_ref = Ref{CUmemoryPool}()
+    cuMemPoolCreate(handle_ref, props)
+
+    ctx = CuCurrentContext()
+    obj = CuMemoryPool(handle_ref[], ctx)
+    finalizer(unsafe_destroy!, obj)
+    return obj
 end
 
 function unsafe_destroy!(pool::CuMemoryPool)
@@ -38,4 +38,22 @@ Base.unsafe_convert(::Type{CUmemoryPool}, pool::CuMemoryPool) = pool.handle
 Base.:(==)(a::CuMemoryPool, b::CuMemoryPool) = a.handle == b.handle
 Base.hash(pool::CuMemoryPool, h::UInt) = hash(pool.handle, h)
 
+function default_memory_pool(dev::CuDevice)
+    handle_ref = Ref{CUmemoryPool}()
+    cuDeviceGetDefaultMemPool(handle_ref, dev)
+
+    ctx = CuCurrentContext()
+    CuMemoryPool(handle_ref[], ctx)
+end
+
+function memory_pool(dev::CuDevice)
+    handle_ref = Ref{CUmemoryPool}()
+    cuDeviceGetMemPool(handle_ref, dev)
+
+    ctx = CuCurrentContext()
+    CuMemoryPool(handle_ref[], ctx)
+end
+
+memory_pool!(dev::CuDevice, pool::CuMemoryPool) = cuDeviceSetMemPool(dev, pool)
+
 trim(pool::CuMemoryPool, bytes_to_keep::Integer=0) = cuMemPoolTrimTo(pool, bytes_to_keep)

src/pool.jl

@@ -127,11 +127,7 @@ function actual_alloc(dev::CuDevice, bytes::Integer, last_resort::Bool=false)
     try
       time = Base.@elapsed begin
         @timeit_debug alloc_to "alloc" begin
-          buf = Mem.alloc(Mem.Device, bytes; async=true,
-                          pool = (async_alloc[] ? pool() : nothing))
-          # we only need a memory pool when we'll be using the async allocator.
-          # this avoids a needless warning when running under cuda-memcheck,
-          # which doesn't support the stream-ordered memory allocator.
+          buf = Mem.alloc(Mem.Device, bytes; async=true)
         end
       end
 
@@ -419,7 +415,8 @@ macro retry_reclaim(isfailed, ex)
       elseif phase == 4 && async_alloc[]
         # this phase is unique to retry_reclaim, as regular allocations come from the pool
         # so are assumed to never need to trim its contents.
-        trim(pool())
+        pool = memory_pool(device())
+        trim(pool)
       end
     end
     ret
@@ -652,6 +649,16 @@ function __init_pool__()
   TimerOutputs.reset_timer!(alloc_to)
   TimerOutputs.reset_timer!(PoolUtils.to)
 
+  if isdebug(:init, CUDA)
+    TimerOutputs.enable_debug_timings(CUDA)
+    atexit() do
+        println("Memory pool timings:")
+        pool_timings()
+        println("Allocator timings:")
+        alloc_timings()
+    end
+  end
+
   if isinteractive()
     @async @pooled pool_cleanup()
   end

src/pool/utils.jl

@@ -103,7 +103,7 @@ export @pool_timeit
 const to = TimerOutput()
 
 macro pool_timeit(args...)
-    TimerOutputs.timer_expr(CUDA, true, :($CUDA.to), args...)
+    TimerOutputs.timer_expr(CUDA, true, :($PoolUtils.to), args...)
 end
 
 end

src/state.jl

@@ -561,18 +561,3 @@ function stream!(f::Function, s::CuStream)
         set_library_streams(old_s)
     end
 end
-
-
-## memory pools
-
-const memory_pools = Dict{CuContext,CuMemoryPool}()
-
-function pool()
-    if CUDA.version() < v"11.2"
-        return nothing
-    end
-
-    return get!(memory_pools, context()) do
-        CuMemoryPool(device())
-    end
-end

test/cudadrv/pool.jl

@@ -0,0 +1,14 @@
+let
+    dev = device()
+
+    pool = memory_pool(dev)
+
+    pool2 = CuMemoryPool(dev)
+    @test pool2 != pool
+    memory_pool!(dev, pool2)
+    @test pool2 == memory_pool(dev)
+    @test pool2 != default_memory_pool(dev)
+
+    memory_pool!(dev, pool)
+    @test pool == memory_pool(dev)
+end