Metal Capture improvements

- Default to capturing with MTLCaptureScope
- Test improvments
- Spinlock until capture is over in `stopCapture`

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

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)