Merge MTLSharedEvent synchronization and Metal capture improvements #690

Author: christiangnrd

lib/mtl/capture.jl

@@ -12,6 +12,15 @@ Use [`beginScope()`](@ref) and [`endScope()`](@ref) to set the boundaries for a
 """
 MTLCaptureScope
 
+function MTLCaptureScope(queue::MTLDevice, manager=MTLCaptureManager())
+    handle = @objc [manager::id{MTLCaptureManager} newCaptureScopeWithDevice:queue::id{MTLDevice}]::id{MTLCaptureScope}
+    MTLCaptureScope(handle)
+end
+function MTLCaptureScope(queue::MTLCommandQueue, manager=MTLCaptureManager())
+    handle = @objc [manager::id{MTLCaptureManager} newCaptureScopeWithCommandQueue:queue::id{MTLCommandQueue}]::id{MTLCaptureScope}
+    MTLCaptureScope(handle)
+end
+
 # @objcwrapper MTLCaptureScope <: NSObject
 
 """
@@ -59,7 +68,7 @@ function MTLCaptureDescriptor()
 end
 
 # TODO: Add capture state
-function MTLCaptureDescriptor(obj::Union{MTLDevice,MTLCommandQueue, MTLCaptureScope},
+function MTLCaptureDescriptor(obj::Union{MTLDevice, MTLCommandQueue, MTLCaptureScope},
                               destination::MTLCaptureDestination;
                               folder::String=nothing)
     desc = MTLCaptureDescriptor()
@@ -110,7 +119,7 @@ end
 
 Start GPU frame capture using the default capture object and specifying capture descriptor parameters directly.
 """
-function startCapture(obj::Union{MTLDevice,MTLCommandQueue, MTLCaptureScope},
+function startCapture(obj::Union{MTLDevice, MTLCommandQueue, MTLCaptureScope},
                       destination::MTLCaptureDestination=MTLCaptureDestinationGPUTraceDocument;
                       folder::String=nothing)
     if destination == MTLCaptureDestinationGPUTraceDocument && folder === nothing
@@ -141,8 +150,17 @@ Stop GPU frame capture.
 """
 function stopCapture(manager::MTLCaptureManager=MTLCaptureManager())
     @objc [manager::id{MTLCaptureManager} stopCapture]::Nothing
+    # Spinlock until capture is fully complete. Otherwise, `stopCapture`
+    #  sometimes returns while the capture manager is still capturing,
+    #  causing test failures and potentially corrupted captures.
+    while manager.isCapturing
+    end
 end
 
+"""
+    supports_destination(manager::MTLCaptureManager, destination::MTLCaptureDestination)
+Checks if a given capture destination is supported.
+"""
 function supports_destination(manager::MTLCaptureManager, destination::MTLCaptureDestination)
     @objc [manager::id{MTLCaptureManager} supportsDestination:destination::MTLCaptureDestination]::Bool
 end

lib/mtl/events.jl

@@ -29,6 +29,10 @@ function MTLSharedEvent(dev::MTLDevice)
     return obj
 end
 
+function waitUntilSignaledValue(ev::MTLSharedEvent, value, timeoutMS=typemax(UInt64))
+    @objc [ev::id{MTLSharedEvent} waitUntilSignaledValue:value::UInt64
+                        timeoutMS:timeoutMS::UInt64]::Bool
+end
 
 ## shared event handle
 

src/state.jl

@@ -55,6 +55,17 @@ function global_queue(dev::MTLDevice)
     end::MTLCommandQueue
 end
 
+"""
+    queue_event(queue::MTLCommandQueue)::MTLSharedEvent
+
+Return the `MTLSharedEvent` used to synchronize a queue
+"""
+function queue_event(queue::MTLCommandQueue)
+    get!(task_local_storage(), (:MTLSharedEvent, queue)) do
+        MTLSharedEvent(queue.device)
+    end::MTLSharedEvent
+end
+
 # TODO: Increase performance (currently ~15us)
 """
     synchronize(queue)
@@ -66,9 +77,13 @@ and simply wait for it to be completed. Since command buffers *should* execute i
 First-In-First-Out manner, this synchronizes the GPU.
 """
 @autoreleasepool function synchronize(queue::MTLCommandQueue=global_queue(device()))
+    ev = queue_event(queue)
+    val = ev.signaledValue + 1
     cmdbuf = MTLCommandBuffer(queue)
+    MTL.encode_signal!(cmdbuf, ev, val)
     commit!(cmdbuf)
-    wait_completed(cmdbuf)
+    MTL.waitUntilSignaledValue(ev, val)
+    return
 end
 
 """

src/utilities.jl

@@ -137,10 +137,13 @@ function captured(f; dest=MTL.MTLCaptureDestinationGPUTraceDocument,
     end
 
     folder = capture_dir()
-    startCapture(object, dest; folder)
+    scope = MTLCaptureScope(object)
+    startCapture(scope, dest; folder)
     try
-       f()
-       synchronize()
+        beginScope(scope)
+        f()
+        endScope(scope)
+        synchronize()
     finally
         @info "GPU frame capture saved to $folder; open the resulting trace in Xcode"
         stopCapture()

test/capturing.jl

@@ -7,8 +7,6 @@ else
 # Verify Metal capture is enabled via environment variable
 @test capturing
 
-@testset "capturing" begin
-
 mktempdir() do tmpdir
 cd(tmpdir) do
 
@@ -38,6 +36,10 @@ desc.captureObject = dev
 desc.destination = MTL.MTLCaptureDestinationGPUTraceDocument
 @test desc.destination == MTL.MTLCaptureDestinationGPUTraceDocument
 
+# Capture Manager supports destination
+@test !supports_destination(manager, MTL.MTLCaptureDestinationDeveloperTools)
+@test supports_destination(manager, MTL.MTLCaptureDestinationGPUTraceDocument)
+
 # Output URL
 @test desc.outputURL === nothing
 path = joinpath(tmpdir, "test.gputrace")
@@ -67,7 +69,7 @@ bufferA = MtlArray{Float32,1,SharedStorage}(undef, tuple(4))
 @test manager.isCapturing == false
 startCapture(manager, desc)
 @test manager.isCapturing
-@test_throws ErrorException startCapture(manager, desc)
+@test_throws "Capture manager is already capturing." startCapture(manager, desc)
 Metal.@sync @metal threads=4 tester(bufferA)
 stopCapture(manager)
 @test manager.isCapturing == false
@@ -85,5 +87,4 @@ end
 end # cd(tmpdir) do
 end # mktempdir() do tmpdir
 
-end # @testset "capturing" begin
 end # if shader_validation (else branch)