Free GPU memory eagerly

ext/CUDAExt.jl

@@ -56,6 +56,11 @@ function Dagger.aliasing(x::CuArray{T}) where T
     return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
 end
 
+function Dagger.unsafe_free!(x::CuArray)
+    CUDA.unsafe_free!(x)
+    return
+end
+
 Dagger.memory_spaces(proc::CuArrayDeviceProc) = Set([CUDAVRAMMemorySpace(proc.owner, proc.device, proc.device_uuid)])
 Dagger.processors(space::CUDAVRAMMemorySpace) = Set([CuArrayDeviceProc(space.owner, space.device, space.device_uuid)])
 
@@ -284,11 +289,14 @@ end
 
 # Adapt BLAS/LAPACK functions
 import LinearAlgebra: BLAS, LAPACK
+_keep_blas_functions = Set(["iamax"])
 for lib in [BLAS, LAPACK]
     for name in names(lib; all=true)
         name == nameof(lib) && continue
         startswith(string(name), '#') && continue
-        endswith(string(name), '!') || continue
+        if !endswith(string(name), '!') && !any(endswith(string(name), func) for func in _keep_blas_functions)
+            continue
+        end
 
         for culib in [CUBLAS, CUSOLVER]
             if name in names(culib; all=true)
@@ -300,6 +308,37 @@ for lib in [BLAS, LAPACK]
     end
 end
 
+# Adapt RefValue
+Dagger.move(from_proc::CPUProc, to_proc::CuArrayDeviceProc, x::Base.RefValue) =
+    Dagger.GPURef(Dagger.move(from_proc, to_proc, x[]), only(Dagger.memory_spaces(to_proc)))
+Dagger.move(from_proc::CuArrayDeviceProc, to_proc::CPUProc, x::Dagger.GPURef{T,CUDAVRAMMemorySpace} where T) =
+    Ref(Dagger.move(from_proc, to_proc, x[]))
+function Dagger.move!(dep_mod, to_space::CPURAMMemorySpace, from_space::CUDAVRAMMemorySpace, to::Base.RefValue, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::CUDAVRAMMemorySpace, from_space::CPURAMMemorySpace, to::Dagger.GPURef, from::Base.RefValue)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::CUDAVRAMMemorySpace, from_space::CUDAVRAMMemorySpace, to::Dagger.GPURef, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+
+# Adapt HaloArray
 CuArray(H::Dagger.HaloArray) = convert(CuArray, H)
 Base.convert(::Type{C}, H::Dagger.HaloArray) where {C<:CuArray} =
     Dagger.HaloArray(C(H.center),

ext/IntelExt.jl

@@ -55,6 +55,11 @@ function Dagger.aliasing(x::oneArray{T}) where T
     return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
 end
 
+function Dagger.unsafe_free!(x::oneArray)
+    oneAPI.unsafe_free!(x)
+    return
+end
+
 Dagger.memory_spaces(proc::oneArrayDeviceProc) = Set([IntelVRAMMemorySpace(proc.owner, proc.device_id)])
 Dagger.processors(space::IntelVRAMMemorySpace) = Set([oneArrayDeviceProc(space.owner, space.device_id)])
 
@@ -278,6 +283,37 @@ function Dagger.execute!(proc::oneArrayDeviceProc, f, args...; kwargs...)
     end
 end
 
+# Adapt RefValue
+Dagger.move(from_proc::CPUProc, to_proc::oneArrayDeviceProc, x::Base.RefValue) =
+    Dagger.GPURef(Dagger.move(from_proc, to_proc, x[]), only(Dagger.memory_spaces(to_proc)))
+Dagger.move(from_proc::oneArrayDeviceProc, to_proc::CPUProc, x::Dagger.GPURef{T,IntelVRAMMemorySpace} where T) =
+    Ref(Dagger.move(from_proc, to_proc, x[]))
+function Dagger.move!(dep_mod, to_space::CPURAMMemorySpace, from_space::IntelVRAMMemorySpace, to::Base.RefValue, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::IntelVRAMMemorySpace, from_space::CPURAMMemorySpace, to::Dagger.GPURef, from::Base.RefValue)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::IntelVRAMMemorySpace, from_space::IntelVRAMMemorySpace, to::Dagger.GPURef, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+
+# Adapt HaloArray
 oneArray(H::Dagger.HaloArray) = convert(oneArray, H)
 Base.convert(::Type{C}, H::Dagger.HaloArray) where {C<:oneArray} =
     Dagger.HaloArray(C(H.center),

ext/MetalExt.jl

@@ -52,6 +52,11 @@ function Dagger.aliasing(x::MtlArray{T}) where T
     return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
 end
 
+function Dagger.unsafe_free!(x::MtlArray)
+    Metal.unsafe_free!(x)
+    return
+end
+
 Dagger.memory_spaces(proc::MtlArrayDeviceProc) = Set([MetalVRAMMemorySpace(proc.owner, proc.device_id)])
 Dagger.processors(space::MetalVRAMMemorySpace) = Set([MtlArrayDeviceProc(space.owner, space.device_id)])
 
@@ -284,6 +289,37 @@ function Dagger.execute!(proc::MtlArrayDeviceProc, f, args...; kwargs...)
     end
 end
 
+# Adapt RefValue
+Dagger.move(from_proc::CPUProc, to_proc::MtlArrayDeviceProc, x::Base.RefValue) =
+    Dagger.GPURef(Dagger.move(from_proc, to_proc, x[]), only(Dagger.memory_spaces(to_proc)))
+Dagger.move(from_proc::MtlArrayDeviceProc, to_proc::CPUProc, x::Dagger.GPURef{T,MetalVRAMMemorySpace} where T) =
+    Ref(Dagger.move(from_proc, to_proc, x[]))
+function Dagger.move!(dep_mod, to_space::CPURAMMemorySpace, from_space::MetalVRAMMemorySpace, to::Base.RefValue, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::MetalVRAMMemorySpace, from_space::CPURAMMemorySpace, to::Dagger.GPURef, from::Base.RefValue)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::MetalVRAMMemorySpace, from_space::MetalVRAMMemorySpace, to::Dagger.GPURef, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+
+# Adapt HaloArray
 MtlArray(H::Dagger.HaloArray) = convert(MtlArray, H)
 Base.convert(::Type{C}, H::Dagger.HaloArray) where {C<:MtlArray} =
     Dagger.HaloArray(C(H.center),

ext/OpenCLExt.jl

@@ -52,6 +52,11 @@ function Dagger.aliasing(x::CLArray{T}) where T
     return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
 end
 
+function Dagger.unsafe_free!(x::CLArray)
+    cl.unsafe_free!(x)
+    return
+end
+
 Dagger.memory_spaces(proc::CLArrayDeviceProc) = Set([CLMemorySpace(proc.owner, proc.device)])
 Dagger.processors(space::CLMemorySpace) = Set([CLArrayDeviceProc(space.owner, space.device)])
 
@@ -251,6 +256,37 @@ function Dagger.execute!(proc::CLArrayDeviceProc, f, args...; kwargs...)
     end
 end
 
+# Adapt RefValue
+Dagger.move(from_proc::CPUProc, to_proc::CLArrayDeviceProc, x::Base.RefValue) =
+    Dagger.GPURef(Dagger.move(from_proc, to_proc, x[]), only(Dagger.memory_spaces(to_proc)))
+Dagger.move(from_proc::CLArrayDeviceProc, to_proc::CPUProc, x::Dagger.GPURef{T,CLMemorySpace} where T) =
+    Ref(Dagger.move(from_proc, to_proc, x[]))
+function Dagger.move!(dep_mod, to_space::CPURAMMemorySpace, from_space::CLMemorySpace, to::Base.RefValue, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::CLMemorySpace, from_space::CPURAMMemorySpace, to::Dagger.GPURef, from::Base.RefValue)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::CLMemorySpace, from_space::CLMemorySpace, to::Dagger.GPURef, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+
+# Adapt HaloArray
 CLArray(H::Dagger.HaloArray) = convert(CLArray, H)
 Base.convert(::Type{C}, H::Dagger.HaloArray) where {C<:CLArray} =
     Dagger.HaloArray(C(H.center),

ext/ROCExt.jl

@@ -47,6 +47,11 @@ function Dagger.aliasing(x::ROCArray{T}) where T
     return Dagger.ContiguousAliasing(Dagger.MemorySpan{S}(rptr, sizeof(T)*length(x)))
 end
 
+function Dagger.unsafe_free!(x::ROCArray)
+    AMDGPU.unsafe_free!(x)
+    return
+end
+
 Dagger.memory_spaces(proc::ROCArrayDeviceProc) = Set([ROCVRAMMemorySpace(proc.owner, proc.device_id)])
 Dagger.processors(space::ROCVRAMMemorySpace) = Set([ROCArrayDeviceProc(space.owner, space.device_id)])
 
@@ -231,24 +236,6 @@ Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Function) = x
 Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Chunk{T}) where {T<:Function} =
     Dagger.move(from_proc, to_proc, fetch(x))
 
-# Adapt BLAS/LAPACK functions
-import LinearAlgebra: BLAS, LAPACK
-for lib in [BLAS, LAPACK]
-    for name in names(lib; all=true)
-        name == nameof(lib) && continue
-        startswith(string(name), '#') && continue
-        endswith(string(name), '!') || continue
-
-        for roclib in [rocBLAS, rocSOLVER]
-            if name in names(roclib; all=true)
-                fn = getproperty(lib, name)
-                rocfn = getproperty(roclib, name)
-                @eval Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, ::$(typeof(fn))) = $rocfn
-            end
-        end
-    end
-end
-
 # Task execution
 function Dagger.execute!(proc::ROCArrayDeviceProc, f, args...; kwargs...)
     @nospecialize f args kwargs
@@ -270,6 +257,58 @@ function Dagger.execute!(proc::ROCArrayDeviceProc, f, args...; kwargs...)
     end
 end
 
+# Adapt BLAS/LAPACK functions
+import LinearAlgebra: BLAS, LAPACK
+_keep_blas_functions = Set(["iamax"])
+for lib in [BLAS, LAPACK]
+    for name in names(lib; all=true)
+        name == nameof(lib) && continue
+        startswith(string(name), '#') && continue
+        if !endswith(string(name), '!') && !any(endswith(string(name), func) for func in _keep_blas_functions)
+            continue
+        end
+
+        for roclib in [rocBLAS, rocSOLVER]
+            if name in names(roclib; all=true)
+                fn = getproperty(lib, name)
+                rocfn = getproperty(roclib, name)
+                @eval Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, ::$(typeof(fn))) = $rocfn
+            end
+        end
+    end
+end
+
+# Adapt RefValue
+Dagger.move(from_proc::CPUProc, to_proc::ROCArrayDeviceProc, x::Base.RefValue) =
+    Dagger.GPURef(Dagger.move(from_proc, to_proc, x[]), only(Dagger.memory_spaces(to_proc)))
+Dagger.move(from_proc::ROCArrayDeviceProc, to_proc::CPUProc, x::Dagger.GPURef{T,ROCVRAMMemorySpace} where T) =
+    Ref(Dagger.move(from_proc, to_proc, x[]))
+function Dagger.move!(dep_mod, to_space::CPURAMMemorySpace, from_space::ROCVRAMMemorySpace, to::Base.RefValue, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::ROCVRAMMemorySpace, from_space::CPURAMMemorySpace, to::Dagger.GPURef, from::Base.RefValue)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+function Dagger.move!(dep_mod, to_space::ROCVRAMMemorySpace, from_space::ROCVRAMMemorySpace, to::Dagger.GPURef, from::Dagger.GPURef)
+    if Dagger.type_may_alias(typeof(from[]))
+        Dagger.move!(dep_mod, to_space, from_space, to[], from[])
+    else
+        to[] = dep_mod(from[])
+    end
+    return
+end
+
+# Adapt HaloArray
 ROCArray(H::Dagger.HaloArray) = convert(ROCArray, H)
 Base.convert(::Type{C}, H::Dagger.HaloArray) where {C<:ROCArray} =
     Dagger.HaloArray(C(H.center),

src/array/alloc.jl

@@ -206,3 +206,12 @@ function Base.view(A::AbstractArray{T,N}, p::Blocks{N}) where {T,N}
 end
 Base.view(A::AbstractArray, ::AutoBlocks) =
     view(A, auto_blocks(size(A)))
+
+function unsafe_free!(A::DArray)
+    spawn_datadeps() do
+        for chunk in A.chunks
+            scope = UnionScope(map(ExactScope, collect(processors(memory_space(chunk)))))
+            Dagger.@spawn scope=scope unsafe_free!(chunk)
+        end
+    end
+end

src/array/copy.jl

@@ -18,6 +18,18 @@ function copy_buffered(f, args...)
     for (buf_arg, arg) in zip(buffered_args, real_args)
         copyto!(arg, buf_arg)
     end
+
+    # Free the buffers
+    foreach(unsafe_free!, buffered_args)
+
+    # If the result is one of the buffered args, return the corresponding
+    # original arg instead (since we've already copied data back to it,
+    # and the buffer has been freed)
+    result_idx = findfirst(buf_arg -> buf_arg === result, buffered_args)
+    if result_idx !== nothing
+        return real_args[result_idx]
+    end
+
     return result
 end
 function allocate_copy_buffer(part::Blocks{N}, A::DArray{T,N}) where {T,N}