feat: add unsafe_free!

LuxDL · Aug 28, 2024 · 50df4fa · 50df4fa
1 parent 1c2619a
commit 50df4fa
Show file tree

Hide file tree

Showing 6 changed files with 42 additions and 0 deletions.
diff --git a/ext/MLDataDevicesAMDGPUExt.jl b/ext/MLDataDevicesAMDGPUExt.jl
@@ -64,6 +64,12 @@ function MLDataDevices.set_device!(::Type{AMDGPUDevice}, ::Nothing, rank::Intege
     return MLDataDevices.set_device!(AMDGPUDevice, id)
 end
 
+# unsafe_free!
+function Internal.unsafe_free_internal!(::Type{AMDGPUDevice}, x::AbstractArray)
+    AMDGPU.unsafe_free!(x)
+    return
+end
+
 # Device Transfer
 ## To GPU
 Adapt.adapt_storage(::AMDGPUDevice{Nothing}, x::AbstractArray) = AMDGPU.roc(x)

diff --git a/ext/MLDataDevicesCUDAExt.jl b/ext/MLDataDevicesCUDAExt.jl
@@ -42,6 +42,12 @@ function MLDataDevices.set_device!(::Type{CUDADevice}, ::Nothing, rank::Integer)
     return MLDataDevices.set_device!(CUDADevice, id)
 end
 
+# unsafe_free!
+function Internal.unsafe_free_internal!(::Type{CUDADevice}, x::AbstractArray)
+    CUDA.unsafe_free!(x)
+    return
+end
+
 # Device Transfer
 Adapt.adapt_storage(::CUDADevice{Nothing}, x::AbstractArray) = CUDA.cu(x)
 function Adapt.adapt_storage(to::CUDADevice, x::AbstractArray)

diff --git a/ext/MLDataDevicesMLUtilsExt.jl b/ext/MLDataDevicesMLUtilsExt.jl
@@ -0,0 +1,5 @@
+module MLDataDevicesMLUtilsExt
+
+using MLUtils: DataLoader
+
+end
diff --git a/ext/MLDataDevicesMetalExt.jl b/ext/MLDataDevicesMetalExt.jl
@@ -18,6 +18,12 @@ Internal.get_device(::MtlArray) = MetalDevice()
 
 Internal.get_device_type(::MtlArray) = MetalDevice
 
+# unsafe_free!
+function Internal.unsafe_free_internal!(::Type{MetalDevice}, x::AbstractArray)
+    Metal.unsafe_free!(x)
+    return
+end
+
 # Device Transfer
 ## To GPU
 Adapt.adapt_storage(::MetalDevice, x::AbstractArray) = Metal.mtl(x)

diff --git a/ext/MLDataDevicesoneAPIExt.jl b/ext/MLDataDevicesoneAPIExt.jl
@@ -29,6 +29,12 @@ Internal.get_device(::oneArray) = oneAPIDevice()
 
 Internal.get_device_type(::oneArray) = oneAPIDevice
 
+# unsafe_free!
+function Internal.unsafe_free_internal!(::Type{oneAPIDevice}, x::AbstractArray)
+    oneAPI.unsafe_free!(x)
+    return
+end
+
 # Device Transfer
 ## To GPU
 for (T1, T2) in ((Float64, Float32), (ComplexF64, ComplexF32))

diff --git a/src/internal.jl b/src/internal.jl
@@ -1,5 +1,6 @@
 module Internal
 
+using Functors: fmap
 using Preferences: load_preference
 using Random: AbstractRNG
 using UnrolledUtilities: unrolled_mapreduce
@@ -149,4 +150,16 @@ for op in (:get_device, :get_device_type)
     end
 end
 
+function unsafe_free_internal!(x::AbstractArray)
+    unsafe_free_internal!(MLDataDevices.get_device_type(x), x)
+    return
+end
+unsafe_free_internal!(::Type, x::AbstractArray) = nothing
+unsafe_free_internal!(_) = nothing
+
+function unsafe_free!(x)
+    fmap(unsafe_free_internal!, x)
+    return
+end
+
 end