Cache FFT handles & more aggressive maybe_collect (#646)

- More aggressive `maybe_collect`. Always trigger it when pressure is `> 0.9` and last time we called it we freed a lot of memory.
- Add `soft_memory_limit` preference (similar to `hard_memory_limit`). It controls how much memory is returned back to the OS if it is not used.
- Actually cache rocFFT handles. Significantly improves dispatch time.

Author: pxl-th

src/fft/fft.jl

@@ -1,9 +1,16 @@
 abstract type ROCFFTPlan{T, K, inplace} <: Plan{T} end
 
+Base.eltype(::ROCFFTPlan{T}) where T = T
+
+is_inplace(::ROCFFTPlan{<:Any, <:Any, inplace}) where inplace = inplace
+
 Base.unsafe_convert(::Type{rocfft_plan}, p::ROCFFTPlan) = p.handle
 
-function unsafe_free!(plan::ROCFFTPlan)
-    rocfft_plan_destroy(plan.handle)
+function AMDGPU.unsafe_free!(plan::ROCFFTPlan)
+    if plan.handle != C_NULL
+        release_plan!(plan)
+        plan.handle = C_NULL
+    end
     unsafe_free!(plan.workarea)
     rocfft_execution_info_destroy(plan.execution_info)
 end
@@ -44,7 +51,7 @@ mutable struct cROCFFTPlan{T, K, inplace, N} <: ROCFFTPlan{T, K, inplace}
             rocfft_execution_info_set_work_buffer(info, workarea, length(workarea))
         end
         p = new(handle, stream, workarea, info, size(X), sizey, xtype, region)
-        finalizer(unsafe_free!, p)
+        finalizer(AMDGPU.unsafe_free!, p)
         p
     end
 end

src/fft/wrappers.jl

@@ -13,6 +13,20 @@ function get_plan(args...)
     return handle, workarea
 end
 
+function release_plan!(plan)
+    key = (
+        AMDGPU.context(), plan.xtype, plan.sz,
+        eltype(plan), is_inplace(plan), plan.region)
+    value = (plan.handle, length(plan.workarea))
+    push!(IDLE_HANDLES, key, value) do
+        destroy_plan!(plan)
+    end
+end
+
+function destroy_plan!(plan)
+    rocfft_plan_destroy(plan.handle)
+end
+
 function create_plan(xtype::rocfft_transform_type, xdims, T, inplace, region)
     precision = (real(T) == Float64) ?
         rocfft_precision_double : rocfft_precision_single
@@ -169,9 +183,3 @@ function create_plan(xtype::rocfft_transform_type, xdims, T, inplace, region)
     rocfft_plan_get_work_buffer_size(handle_ref[], worksize_ref)
     return handle_ref[], Int(worksize_ref[])
 end
-
-function release_plan(plan)
-    push!(IDLE_HANDLES, plan) do
-        unsafe_free!(plan)
-    end
-end

src/memory.jl

@@ -67,16 +67,34 @@ end
 """
 Set a hard limit for total GPU memory allocations.
 """
-set_memory_alloc_limit!(limit::String) =
+hard_memory_limit!(limit::String) =
     @set_preferences!("hard_memory_limit" => limit)
 
+soft_memory_limit!(limit::String) =
+    @set_preferences!("soft_memory_limit" => limit)
+
 const HARD_MEMORY_LIMIT = Ref{Union{Nothing, UInt64}}(nothing)
 function hard_memory_limit()
-    l = HARD_MEMORY_LIMIT[]
-    l ≢ nothing && return l
+    hard_limit = HARD_MEMORY_LIMIT[]
+    hard_limit ≢ nothing && return hard_limit
 
-    HARD_MEMORY_LIMIT[] = parse_memory_limit(
+    hard_limit = parse_memory_limit(
         @load_preference("hard_memory_limit", "none"))
+
+    @debug "Setting hard memory limit: $(Base.format_bytes(hard_limit))"
+    HARD_MEMORY_LIMIT[] = hard_limit
+end
+
+const SOFT_MEMORY_LIMIT = Ref{Union{Nothing, UInt64}}(nothing)
+function soft_memory_limit()
+    soft_limit = SOFT_MEMORY_LIMIT[]
+    soft_limit ≢ nothing && return soft_limit
+
+    soft_limit = parse_memory_limit(
+        @load_preference("soft_memory_limit", "none"))
+
+    @debug "Setting soft memory limit: $(Base.format_bytes(soft_limit))"
+    SOFT_MEMORY_LIMIT[] = soft_limit
 end
 
 mutable struct MemoryStats
@@ -152,19 +170,23 @@ function maybe_collect(; blocking::Bool = false)
     # Tolerate 5% GC time.
     max_gc_rate = 0.05
     # If freed a lot of memory last time, double max GC rate.
-    (stats.last_freed > 0.1 * stats.size) && (max_gc_rate *= 2;)
+    freed_alot = stats.last_freed > 0.1 * stats.size
+    freed_alot && (max_gc_rate *= 2;)
     # Be more aggressive if we are going to block.
     blocking && (max_gc_rate *= 2;)
 
-    # And even more if we are at a limit.
-    pressure > 0.9 && (max_gc_rate *= 2;)
-    pressure > 0.95 && (max_gc_rate *= 2;)
+    # And even more if the pressure is high.
+    pressure > 0.5 && (max_gc_rate *= 2;)
+    pressure > 0.7 && (max_gc_rate *= 2;)
+
+    # Always free if pressure is 0.9 and we freed a lot.
+    pressure > 0.9 && (max_gc_rate *= freed_alot ? Inf : 2;)
     gc_rate > max_gc_rate && return
 
     # Call the GC.
     Base.@atomic stats.last_time = current_time
     pre_gc_live = stats.live
-    gc_time = Base.@elapsed GC.gc(pressure > 0.9 ? true : false)
+    gc_time = Base.@elapsed GC.gc(pressure > 0.7 ? true : false)
     post_gc_live = stats.live
 
     # Update stats.

src/runtime/memory/hip.jl

@@ -9,7 +9,10 @@ function pool_create(dev::HIPDevice)
             max_size = max_size != typemax(UInt64) ? max_size : UInt64(0)
 
             pool = HIP.HIPMemoryPool(dev; max_size)
-            # TODO set soft threshold?
+            # Allow pool to use up all device memory.
+            soft_limit = AMDGPU.soft_memory_limit()
+            HIP.attribute!(pool, HIP.hipMemPoolAttrReleaseThreshold, soft_limit)
+
             HIP.memory_pool!(dev, pool)
             return pool
         end