[Experimental] Add memory recording which can then be bulk-freed (#698)

Author: pxl-th

Project.toml

@@ -1,7 +1,7 @@
 name = "AMDGPU"
 uuid = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 authors = ["Julian P Samaroo <jpsamaroo@jpsamaroo.me>", "Valentin Churavy <v.churavy@gmail.com>", "Anton Smirnov <tonysmn97@gmail.com>"]
-version = "1.1.0"
+version = "1.1.1"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

src/AMDGPU.jl

@@ -114,6 +114,7 @@ include("tls.jl")
 include("highlevel.jl")
 include("reflection.jl")
 include("array.jl")
+include("memory_record.jl")
 include("conversions.jl")
 include("broadcast.jl")
 include("exception_handler.jl")

src/array.jl

@@ -8,8 +8,10 @@ mutable struct ROCArray{T, N, B} <: AbstractGPUArray{T, N}
     ) where {T, N, B <: Mem.AbstractAMDBuffer}
         @assert isbitstype(T) "ROCArray only supports bits types"
         data = DataRef(pool_free, pool_alloc(B, prod(dims) * sizeof(T)))
-        xs = new{T, N, B}(data, dims, 0)
-        return finalizer(unsafe_free!, xs)
+        x = new{T, N, B}(data, dims, 0)
+        x = finalizer(unsafe_free!, x)
+        RECORD_MEMORY[] && record!(x)
+        return x
     end
 
     function ROCArray{T, N}(

src/memory_record.jl

@@ -0,0 +1,48 @@
+# NOTE: EXPERIMENTAL API.
+
+const MemoryRecords = LockedObject(Dict{UInt64, ROCArray}())
+
+const RECORD_MEMORY::Ref{Bool} = Ref(false)
+
+function record_memory!(rec::Bool; free::Bool = true, sync::Bool = false)
+    RECORD_MEMORY[] = rec
+    if !rec
+        free && free_records!(; sync)
+    end
+    return
+end
+
+record_memory() = RECORD_MEMORY[]
+
+function record!(x)
+    Base.lock(records -> records[_hash(x)] = x, MemoryRecords)
+    return
+end
+
+function free_records!(; sync::Bool = false)
+    Base.lock(MemoryRecords) do records
+        for (k, x) in records
+            unsafe_free!(x)
+        end
+        empty!(records)
+    end
+    sync && AMDGPU.synchronize()
+    return
+end
+
+function remove_record!(x)
+    record_memory() || return
+
+    k = _hash(x)
+    Base.lock(MemoryRecords) do records
+        if k in records.keys
+            pop!(records, k)
+        end
+    end
+    return
+end
+
+_hash(x::ROCArray) =
+    Base.hash(x.buf.rc.obj.mem.ptr,
+        Base.hash(x.offset,
+            Base.hash(x.dims)))