Use NSAutoreleasePool to clean up memory. (#294)

Author: maleadt

Project.toml

@@ -38,7 +38,7 @@ KernelAbstractions = "0.9.1"
 LLVM = "6"
 LLVMDowngrader_jll = "0.1"
 ObjectFile = "0.4"
-ObjectiveC = "1.1, 2"
+ObjectiveC = "2.1"
 Reexport = "1.0"
 Requires = "1"
 SHA = "0.7"

lib/mtl/MTL.jl

@@ -3,6 +3,12 @@ module MTL
 using CEnum
 using ObjectiveC, .Foundation, .Dispatch
 
+# Metal APIs generally expect to be running under an autorelease pool.
+# In most cases, we handle this in the code calling into the MTL module,
+# however, finalizers are out of the caller's control, so we need to
+# ensure here already that they are running under an autorelease pool.
+release(obj) = @autoreleasepool unsafe=true Foundation.release(obj)
+
 
 ## source code includes
 

src/compiler/compilation.jl

@@ -107,7 +107,8 @@ function compile(@nospecialize(job::CompilerJob))
 end
 
 # link into an executable kernel
-function link(@nospecialize(job::CompilerJob), compiled; return_function=false)
+@autoreleasepool function link(@nospecialize(job::CompilerJob), compiled;
+                               return_function=false)
     dev = current_device()
     lib = MTLLibraryFromData(dev, compiled.image)
     fun = MTLFunction(lib, compiled.entry)

src/compiler/execution.jl

@@ -259,8 +259,8 @@ end
     return argument_buffer
 end
 
-function (kernel::HostKernel)(args...; groups=1, threads=1,
-                              queue=global_queue(current_device()))
+@autoreleasepool function (kernel::HostKernel)(args...; groups=1, threads=1,
+                                               queue=global_queue(current_device()))
     groups = MTLSize(groups)
     threads = MTLSize(threads)
     (groups.width>0 && groups.height>0 && groups.depth>0) ||

src/compiler/reflection.jl

@@ -36,7 +36,7 @@ function code_agx(io::IO, @nospecialize(func), @nospecialize(types),
     code_agx(io, job)
 end
 
-function code_agx(io::IO, job::MetalCompilerJob)
+@autoreleasepool function code_agx(io::IO, job::MetalCompilerJob)
     if !job.config.kernel
         error("Can only generate AGX code for kernel functions")
     end
@@ -58,7 +58,7 @@ function code_agx(io::IO, job::MetalCompilerJob)
     bin = MTLBinaryArchive(current_device(), bin_desc)
     add_functions!(bin, pipeline_desc)
 
-    code = mktempdir() do dir
+    mktempdir() do dir
         # serialize the archive to a file
         binary = joinpath(dir, "kernel.macho")
         write(binary, bin)
@@ -82,7 +82,6 @@ function code_agx(io::IO, job::MetalCompilerJob)
             first = false
         end
     end
-
 end
 
 @enum GPUMachineType::UInt32 begin

src/initialization.jl

@@ -28,7 +28,7 @@ function __init__()
         end
     end
 
-    try
+    @autoreleasepool try
         load_framework("CoreGraphics")
         ver = MTL.MTLCompileOptions().languageVersion
         @debug "Successfully loaded Metal; targeting v$ver."

src/memory.jl

@@ -69,8 +69,10 @@ function Base.unsafe_copyto!(dev::MTLDevice, dst::Ptr{T}, src::MtlPointer{T}, N:
 end
 
 # GPU -> GPU
-function Base.unsafe_copyto!(dev::MTLDevice, dst::MtlPointer{T}, src::MtlPointer{T}, N::Integer;
-                             queue::MTLCommandQueue=global_queue(dev), async::Bool=false) where T
+@autoreleasepool function Base.unsafe_copyto!(dev::MTLDevice, dst::MtlPointer{T},
+                                              src::MtlPointer{T}, N::Integer;
+                                              queue::MTLCommandQueue=global_queue(dev),
+                                              async::Bool=false) where T
     cmdbuf = MTLCommandBuffer(queue)
     MTLBlitCommandEncoder(cmdbuf) do enc
         MTL.append_copy!(enc, dst.buffer, dst.offset, src.buffer, src.offset, N * sizeof(T))
@@ -79,7 +81,8 @@ function Base.unsafe_copyto!(dev::MTLDevice, dst::MtlPointer{T}, src::MtlPointer
     async || wait_completed(cmdbuf)
 end
 
-function unsafe_fill!(dev::MTLDevice, ptr::MtlPointer{T}, value::Union{UInt8,Int8}, N::Integer) where T
+@autoreleasepool function unsafe_fill!(dev::MTLDevice, ptr::MtlPointer{T},
+                                       value::Union{UInt8,Int8}, N::Integer) where T
     cmdbuf = MTLCommandBuffer(global_queue(dev))
     MTLBlitCommandEncoder(cmdbuf) do enc
         MTL.append_fillbuffer!(enc, ptr.buffer, value, N * sizeof(T), ptr.offset)

src/pool.jl

@@ -51,14 +51,11 @@ The storage kwarg controls where the buffer is stored. Possible values are:
 Note that `Private` buffers can't be directly accessed from the CPU, therefore you cannot
 use this option if you pass a ptr to initialize the memory.
 """
-function alloc(dev::Union{MTLDevice,MTLHeap},
-               bytesize::Integer,
-               args...;
-               storage,
-               kwargs...)
+function alloc(dev::Union{MTLDevice,MTLHeap}, bytesize::Integer, args...;
+               storage, kwargs...)
 
     time = Base.@elapsed begin
-        buf = MTLBuffer(dev, bytesize, args...; storage, kwargs...)
+        buf = @autoreleasepool MTLBuffer(dev, bytesize, args...; storage, kwargs...)
     end
 
     Base.@atomic alloc_stats.alloc_count + 1
@@ -78,7 +75,7 @@ function free(buf::MTLBuffer)
     sz::Int = buf.length
 
     time = Base.@elapsed begin
-        release(buf)
+        @autoreleasepool unsafe=true release(buf)
     end
 
     Base.@atomic alloc_stats.free_count + 1

src/state.jl

@@ -11,8 +11,10 @@ effectively returns the only system GPU.
 function current_device()
     get!(task_local_storage(), :MTLDevice) do
         dev = MTLDevice(1)
-        supports_family(dev, MTL.MTLGPUFamilyApple7) || @warn """Metal.jl is only supported on M-series Macs, you may run into issues.
-                                                    See https://github.com/JuliaGPU/Metal.jl/issues/22 for more details.""" maxlog=1
+        if !supports_family(dev, MTL.MTLGPUFamilyApple7)
+            @warn """Metal.jl is only supported on M-series Macs, you may run into issues.
+                     See https://github.com/JuliaGPU/Metal.jl/issues/22 for more details.""" maxlog=1
+        end
         return dev
     end::MTLDevice
 end
@@ -33,10 +35,14 @@ Return the Metal command queue associated with the current Julia thread.
 """
 function global_queue(dev::MTLDevice)
     get!(task_local_storage(), (:MTLCommandQueue, dev)) do
-        queue = MTLCommandQueue(dev)
-        queue.label = "global_queue($(current_task()))"
-        global_queues[queue] = nothing
-        queue
+        @autoreleasepool begin
+            # NOTE: MTLCommandQueue itself is manually reference-counted,
+            #       the release pool is for resources used during its construction.
+            queue = MTLCommandQueue(dev)
+            queue.label = "global_queue($(current_task()))"
+            global_queues[queue] = nothing
+            queue
+        end
     end::MTLCommandQueue
 end
 
@@ -50,7 +56,7 @@ Create a new MTLCommandBuffer from the global command queue, commit it to the qu
 and simply wait for it to be completed. Since command buffers *should* execute in a
 First-In-First-Out manner, this synchronizes the GPU.
 """
-function synchronize(queue::MTLCommandQueue=global_queue(current_device()))
+@autoreleasepool function synchronize(queue::MTLCommandQueue=global_queue(current_device()))
     cmdbuf = MTLCommandBuffer(queue)
     commit!(cmdbuf)
     wait_completed(cmdbuf)

test/metal.jl

@@ -1,4 +1,5 @@
 @testset "MTL" begin
+@autoreleasepool begin
 
 @testset "devices" begin
 
@@ -462,3 +463,4 @@ end
 # TODO: continue adding tests
 
 end
+end