diff --git a/lib/mps/command_buf.jl b/lib/mps/command_buf.jl index d6f47d80c..1d899e71e 100644 --- a/lib/mps/command_buf.jl +++ b/lib/mps/command_buf.jl @@ -16,6 +16,15 @@ function MPSCommandBuffer(commandQueue) @objc [MPSCommandBuffer commandBufferFromCommandQueue:commandQueue::id{MTLCommandQueue}]::MPSCommandBuffer end +# Deriving from a batched queue orders the buffer after the Metal 4 work committed so far +# with a GPU-side wait, instead of the host-side wait that converting the queue implies. +# The commit hooks see the underlying `MTLCommandBuffer`, so that is what gets registered. +function MPSCommandBuffer(bq::Metal.BatchedCommandQueue) + cmdbuf = MPSCommandBuffer(bq.queue) + Metal.order_after_batch!(bq, cmdbuf, Metal.derived_key(cmdbuf.commandBuffer)) + return cmdbuf +end + function MPSCommandBuffer(f::Base.Callable, queueOrBuf) cmdbuf = MPSCommandBuffer(queueOrBuf) commitAndContinue!(f, cmdbuf) @@ -40,6 +49,8 @@ function commitAndContinue!(cmdbuf::MPSCommandBuffer) submitted = cmdbuf.commandBuffer hook = MTL.submit_hook[] hook === nothing || hook(submitted) + hook = MTL.commit_hook[] + hook === nothing || hook(submitted) @objc [cmdbuf::id{MPSCommandBuffer} commitAndContinue]::Nothing MTL.record_committed!(submitted, pointer(submitted.commandQueue)) end diff --git a/lib/mtl/MTL.jl b/lib/mtl/MTL.jl index e6317cf39..d3fc65015 100644 --- a/lib/mtl/MTL.jl +++ b/lib/mtl/MTL.jl @@ -41,11 +41,18 @@ include("buffer.jl") include("log_state.jl") include("residency_set.jl") include("command_queue.jl") +include("command_queue4.jl") include("command_buf.jl") +include("command_buf4.jl") include("compute_pipeline.jl") +include("compute_pipeline4.jl") include("command_enc.jl") +include("command_alloc4.jl") +include("command_enc4.jl") +include("arg_table.jl") include("command_enc/blit.jl") include("command_enc/compute.jl") +include("command_enc/compute4.jl") include("binary_archive.jl") include("capture.jl") include("texture.jl") diff --git a/lib/mtl/arg_table.jl b/lib/mtl/arg_table.jl new file mode 100644 index 000000000..94265e587 --- /dev/null +++ b/lib/mtl/arg_table.jl @@ -0,0 +1,102 @@ +# +# argument table descriptor +# + +export MTL4ArgumentTable, MTL4ArgumentTableDescriptor + +# @objcwrapper managed = true MTL4ArgumentTableDescriptor <: NSObject + +function MTL4ArgumentTableDescriptor() + return @objc [MTL4ArgumentTableDescriptor new]::MTL4ArgumentTableDescriptor +end + + +# +# argument table +# + +# @objcwrapper managed = true MTL4ArgumentTable <: NSObject + +""" + MTL4ArgumentTable(device, desc::MTL4ArgumentTableDescriptor) + +Create an argument table: the Metal 4 replacement for the per-encoder `setBuffer:`/ +`setBytes:` binding state. Bindings are captured when a command is encoded, so a single +table can be rebound between dispatches on the same encoder. +""" +function MTL4ArgumentTable(device::MTLDevice, desc::MTL4ArgumentTableDescriptor) + err = Ref{id{NSError}}(nil) + argtab = @objc [device::id{MTLDevice} newArgumentTableWithDescriptor:desc::id{MTL4ArgumentTableDescriptor} + error:err::Ptr{id{NSError}}]::Union{Nothing,MTL4ArgumentTable} + argtab === nothing && throw_error(err[]) + return argtab +end + +function MTL4ArgumentTable(device::MTLDevice; buffers::Integer=0, textures::Integer=0, + samplers::Integer=0, label=nothing) + desc = MTL4ArgumentTableDescriptor() + desc.maxBufferBindCount = buffers + desc.maxTextureBindCount = textures + desc.maxSamplerStateBindCount = samplers + label === nothing || (desc.label = label) + return MTL4ArgumentTable(device, desc) +end + +# NOTE: binding indices are 1-based here, like `set_buffer!`/`set_bytes!` on a Metal 3 +# compute command encoder; they are lowered to Metal's 0-based `[[buffer(n)]]` slots. + +""" + set_address!(argtab::MTL4ArgumentTable, address, index) + +Bind a raw GPU virtual address to the 1-based buffer binding `index`. +""" +function set_address!(argtab::MTL4ArgumentTable, address::Integer, index::Integer) + @objc [argtab::id{MTL4ArgumentTable} setAddress:UInt64(address)::MTLGPUAddress + atIndex:(index-1)::NSUInteger]::Nothing +end + +""" + set_address!(argtab::MTL4ArgumentTable, address, stride, index) + +Bind a raw GPU virtual address with an explicit attribute stride. The table must have been +created with `supportAttributeStrides` enabled. +""" +function set_address!(argtab::MTL4ArgumentTable, address::Integer, stride::Integer, + index::Integer) + @objc [argtab::id{MTL4ArgumentTable} setAddress:UInt64(address)::MTLGPUAddress + attributeStride:stride::NSUInteger + atIndex:(index-1)::NSUInteger]::Nothing +end + +""" + set_buffer!(argtab::MTL4ArgumentTable, buf::MTLBuffer, offset, index) + +Bind `buf` (at a byte `offset`) to the 1-based buffer binding `index`. + +Unlike Metal 3's `setBuffer:offset:atIndex:`, this binds the buffer's GPU address and does +*not* make it resident: the buffer must be covered by a residency set that is used by the +command buffer, or attached to the queue. +""" +function set_buffer!(argtab::MTL4ArgumentTable, buf::MTLBuffer, offset::Integer, + index::Integer) + @inline set_address!(argtab, UInt64(buf.gpuAddress) + offset, index) +end + +function set_resource!(argtab::MTL4ArgumentTable, resource::MTLResourceID, index::Integer) + @objc [argtab::id{MTL4ArgumentTable} setResource:resource::MTLResourceID + atBufferIndex:(index-1)::NSUInteger]::Nothing +end + +function set_texture!(argtab::MTL4ArgumentTable, texture::MTLResourceID, index::Integer) + @objc [argtab::id{MTL4ArgumentTable} setTexture:texture::MTLResourceID + atIndex:(index-1)::NSUInteger]::Nothing +end +set_texture!(argtab::MTL4ArgumentTable, texture::MTLTexture, index::Integer) = + set_texture!(argtab, texture.gpuResourceID, index) + +function set_sampler_state!(argtab::MTL4ArgumentTable, sampler::MTLResourceID, index::Integer) + @objc [argtab::id{MTL4ArgumentTable} setSamplerState:sampler::MTLResourceID + atIndex:(index-1)::NSUInteger]::Nothing +end +set_sampler_state!(argtab::MTL4ArgumentTable, sampler::MTLSamplerState, index::Integer) = + set_sampler_state!(argtab, sampler.gpuResourceID, index) diff --git a/lib/mtl/capture.jl b/lib/mtl/capture.jl index 18e258e94..2d8fcc741 100644 --- a/lib/mtl/capture.jl +++ b/lib/mtl/capture.jl @@ -18,6 +18,9 @@ end function MTLCaptureScope(queue::MTLCommandQueue, manager=MTLCaptureManager()) @objc [manager::id{MTLCaptureManager} newCaptureScopeWithCommandQueue:queue::id{MTLCommandQueue}]::MTLCaptureScope end +function MTLCaptureScope(queue::MTL4CommandQueue, manager=MTLCaptureManager()) + @objc [manager::id{MTLCaptureManager} newCaptureScopeWithMTL4CommandQueue:queue::id{MTL4CommandQueue}]::MTLCaptureScope +end # @objcwrapper MTLCaptureScope <: NSObject diff --git a/lib/mtl/command_alloc4.jl b/lib/mtl/command_alloc4.jl new file mode 100644 index 000000000..2f7788ea7 --- /dev/null +++ b/lib/mtl/command_alloc4.jl @@ -0,0 +1,63 @@ +# +# command allocator descriptor +# + +export MTL4CommandAllocatorDescriptor + +# @objcwrapper managed = true MTL4CommandAllocatorDescriptor <: NSObject + +function MTL4CommandAllocatorDescriptor() + return @objc [MTL4CommandAllocatorDescriptor new]::MTL4CommandAllocatorDescriptor +end + +function MTL4CommandAllocatorDescriptor(label) + desc = MTL4CommandAllocatorDescriptor() + desc.label = label + return desc +end + + +# +# command allocator +# + +export MTL4CommandAllocator + +# @objcwrapper managed = true MTL4CommandAllocator <: NSObject + +function MTL4CommandAllocator(dev::MTLDevice) + return @objc [dev::id{MTLDevice} newCommandAllocator]::MTL4CommandAllocator +end + +function MTL4CommandAllocator(dev::MTLDevice, desc::MTL4CommandAllocatorDescriptor) + err = Ref{id{NSError}}(nil) + alloc = @objc [dev::id{MTLDevice} newCommandAllocatorWithDescriptor:desc::id{MTL4CommandAllocatorDescriptor} + error:err::Ptr{id{NSError}}]::Union{Nothing,MTL4CommandAllocator} + alloc === nothing && throw_error(err[]) + return alloc +end + +function MTL4CommandAllocator(dev::MTLDevice, label::Union{String,NSString}) + desc = MTL4CommandAllocatorDescriptor(label) + return MTL4CommandAllocator(dev, desc) +end + +""" + allocatedSize(alloc::MTL4CommandAllocator)::UInt64 + +The amount of memory, in bytes, that `alloc` currently holds for command storage. +""" +function allocatedSize(alloc::MTL4CommandAllocator) + @objc [alloc::id{MTL4CommandAllocator} allocatedSize]::UInt64 +end + +""" + reset!(alloc::MTL4CommandAllocator) + +Return `alloc`'s memory to it for reuse. This invalidates the commands of every command +buffer that was encoded with `alloc`, so it is only legal once the GPU has finished +executing all of them. +""" +function reset!(alloc::MTL4CommandAllocator) + @objc [alloc::id{MTL4CommandAllocator} reset]::Nothing +end diff --git a/lib/mtl/command_buf.jl b/lib/mtl/command_buf.jl index 8323eaf76..7eb77625e 100644 --- a/lib/mtl/command_buf.jl +++ b/lib/mtl/command_buf.jl @@ -203,6 +203,11 @@ const profile_hook = Ref{Any}(nothing) # enqueued or committed. const submit_hook = Ref{Any}(nothing) +# optional commit hook. when set, it is invoked just before a command buffer is handed to +# Metal, i.e. with all of its encoders closed but before it becomes immutable. unlike +# `submit_hook` this is the point at which additional commands may still be encoded. +const commit_hook = Ref{Any}(nothing) + # optional profiling data for operation metadata (e.g. kernel dimensions, copy sizes). const profile_metadata = Ref{Any}(nothing) @@ -215,7 +220,7 @@ end ProfileCollector() = ProfileCollector(ReentrantLock(), IdDict{Any,Vector{Any}}(), Tuple{String,Any}[]) -@inline function note_operation!(collector::ProfileCollector, cmdbuf::MTLCommandBufferLike, op) +@inline function note_operation!(collector::ProfileCollector, cmdbuf, op) @lock collector.lock begin ops = get(collector.metadata, cmdbuf, nothing) if ops === nothing @@ -255,6 +260,8 @@ end function commit_with_queue_key!(cmdbuf::MTLCommandBufferLike, key::id{MTLCommandQueue}) cmdbuf.status in [MTLCommandBufferStatusCompleted, MTLCommandBufferStatusCommitted] && error("Cannot commit an already committed/completed command buffer") + hook = commit_hook[] + hook === nothing || hook(cmdbuf) @objc [cmdbuf::id{MTLCommandBuffer} commit]::Nothing # Record every submission for error accounting. The most recent buffer remains # the queue tail used by synchronization, while older completed buffers are diff --git a/lib/mtl/command_buf4.jl b/lib/mtl/command_buf4.jl new file mode 100644 index 000000000..e0d8554f8 --- /dev/null +++ b/lib/mtl/command_buf4.jl @@ -0,0 +1,104 @@ +# +# command buffer options +# + +export MTL4CommandBufferOptions + +# @objcwrapper managed = true MTL4CommandBufferOptions <: NSObject + +function MTL4CommandBufferOptions() + return @objc [MTL4CommandBufferOptions new]::MTL4CommandBufferOptions +end + + +# +# command buffer +# + +export MTL4CommandBuffer, beginCommandBuffer!, endCommandBuffer! + +# @objcwrapper managed = true MTL4CommandBuffer <: NSObject + +""" + MTL4CommandBuffer(dev::MTLDevice) + +Create a Metal 4 command buffer. Unlike a Metal 3 `MTLCommandBuffer`, this object is not +tied to a queue and carries no storage of its own: it is a reusable encoding cursor that +writes into the [`MTL4CommandAllocator`](@ref) passed to [`beginCommandBuffer!`](@ref), and +it can be re-used for another encoding pass as soon as it has been committed. +""" +function MTL4CommandBuffer(dev::MTLDevice) + return @objc [dev::id{MTLDevice} newCommandBuffer]::MTL4CommandBuffer +end + +function MTL4CommandBuffer(dev::MTLDevice, label::Union{String,NSString}) + cmdbuf = MTL4CommandBuffer(dev) + cmdbuf.label = label + return cmdbuf +end + +""" + beginCommandBuffer!(cmdbuf::MTL4CommandBuffer, allocator::MTL4CommandAllocator, [options]) + +Open `cmdbuf` for encoding, storing its commands in `allocator`. +""" +function beginCommandBuffer!(cmdbuf::MTL4CommandBuffer, allocator::MTL4CommandAllocator) + @objc [cmdbuf::id{MTL4CommandBuffer} beginCommandBufferWithAllocator:allocator::id{MTL4CommandAllocator}]::Nothing +end + +function beginCommandBuffer!(cmdbuf::MTL4CommandBuffer, allocator::MTL4CommandAllocator, + options::MTL4CommandBufferOptions) + @objc [cmdbuf::id{MTL4CommandBuffer} beginCommandBufferWithAllocator:allocator::id{MTL4CommandAllocator} + options:options::id{MTL4CommandBufferOptions}]::Nothing +end + +""" + endCommandBuffer!(cmdbuf::MTL4CommandBuffer) + +Close `cmdbuf` for encoding, making it eligible for `commit!`. +""" +function endCommandBuffer!(cmdbuf::MTL4CommandBuffer) + @objc [cmdbuf::id{MTL4CommandBuffer} endCommandBuffer]::Nothing +end + +""" + MTL4CommandBuffer(f, dev, allocator; queue, options) + +Open a command buffer, apply `f` to it, then end and commit it on `queue`, returning `f`'s +value. +""" +function MTL4CommandBuffer(f::Base.Callable, dev::MTLDevice, + allocator::MTL4CommandAllocator=MTL4CommandAllocator(dev); + queue::MTL4CommandQueue=MTL4CommandQueue(dev), + options::Union{Nothing,MTL4CommitOptions}=nothing) + cmdbuf = MTL4CommandBuffer(dev) + beginCommandBuffer!(cmdbuf, allocator) + ret = try + f(cmdbuf) + finally + endCommandBuffer!(cmdbuf) + end + if options === nothing + commit!(queue, cmdbuf) + else + commit!(queue, cmdbuf, options) + end + return ret +end + +function use_residency_set!(cmdbuf::MTL4CommandBuffer, resset::MTLResidencySet) + @objc [cmdbuf::id{MTL4CommandBuffer} useResidencySet:resset::id{MTLResidencySet}]::Nothing +end + +function use_residency_sets!(cmdbuf::MTL4CommandBuffer, ressets, count) + @objc [cmdbuf::id{MTL4CommandBuffer} useResidencySets:ressets::Ptr{id{MTLResidencySet}} + count:count::NSUInteger]::Nothing +end + +function push_debug_group!(cmdbuf::MTL4CommandBuffer, name::Union{String,NSString}) + @objc [cmdbuf::id{MTL4CommandBuffer} pushDebugGroup:name::id{NSString}]::Nothing +end + +function pop_debug_group!(cmdbuf::MTL4CommandBuffer) + @objc [cmdbuf::id{MTL4CommandBuffer} popDebugGroup]::Nothing +end diff --git a/lib/mtl/command_enc/compute4.jl b/lib/mtl/command_enc/compute4.jl new file mode 100644 index 000000000..a8f67c0b9 --- /dev/null +++ b/lib/mtl/command_enc/compute4.jl @@ -0,0 +1,101 @@ +export MTL4ComputeCommandEncoder +export set_argument_table!, set_threadgroup_memory_length! + +# @objcwrapper managed = true MTL4ComputeCommandEncoder <: MTL4CommandEncoder + +function MTL4ComputeCommandEncoder(cmdbuf::MTL4CommandBuffer) + return @objc [cmdbuf::id{MTL4CommandBuffer} computeCommandEncoder]::MTL4ComputeCommandEncoder +end + +function MTL4ComputeCommandEncoder(f::Base.Callable, cmdbuf::MTL4CommandBuffer) + encoder = MTL4ComputeCommandEncoder(cmdbuf) + try + f(encoder) + finally + close(encoder) + end +end + +## pipeline and argument state + +function set_function!(cce::MTL4ComputeCommandEncoder, pipeline::MTLComputePipelineState) + @objc [cce::id{MTL4ComputeCommandEncoder} setComputePipelineState:pipeline::id{MTLComputePipelineState}]::Nothing +end + +function set_argument_table!(cce::MTL4ComputeCommandEncoder, argtab::MTL4ArgumentTable) + @objc [cce::id{MTL4ComputeCommandEncoder} setArgumentTable:argtab::id{MTL4ArgumentTable}]::Nothing +end + +function set_threadgroup_memory_length!(cce::MTL4ComputeCommandEncoder, length::Integer, + index::Integer) + @objc [cce::id{MTL4ComputeCommandEncoder} setThreadgroupMemoryLength:length::NSUInteger + atIndex:(index-1)::NSUInteger]::Nothing +end + +""" + stages(cce::MTL4ComputeCommandEncoder)::MTLStages + +The set of pipeline stages this encoder can encode work for; the natural argument for the +`after`/`before` parameters of the barrier functions. +""" +function stages(cce::MTL4ComputeCommandEncoder) + @objc [cce::id{MTL4ComputeCommandEncoder} stages]::MTLStages +end + +## dispatch + +function dispatchThreadgroups!(cce::MTL4ComputeCommandEncoder, threadgroupsPerGrid::MTLSize, + threadsPerThreadgroup::MTLSize) + @objc [cce::id{MTL4ComputeCommandEncoder} dispatchThreadgroups:threadgroupsPerGrid::MTLSize + threadsPerThreadgroup:threadsPerThreadgroup::MTLSize]::Nothing +end + +function dispatchThreads!(cce::MTL4ComputeCommandEncoder, threadsPerGrid::MTLSize, + threadsPerThreadgroup::MTLSize) + @objc [cce::id{MTL4ComputeCommandEncoder} dispatchThreads:threadsPerGrid::MTLSize + threadsPerThreadgroup:threadsPerThreadgroup::MTLSize]::Nothing +end + +function append_current_function!(cce::MTL4ComputeCommandEncoder, threadgroupsPerGrid, + threadsPerThreadgroup) + dispatchThreadgroups!(cce, threadgroupsPerGrid, threadsPerThreadgroup) +end + +## copy and fill +# +# Metal 4 folds the Metal 3 blit encoder's buffer operations into the compute encoder. + +function append_copy!(cce::MTL4ComputeCommandEncoder, dst::MTLBuffer, doff, + src::MTLBuffer, soff, len) + @objc [cce::id{MTL4ComputeCommandEncoder} copyFromBuffer:src::id{MTLBuffer} + sourceOffset:soff::NSUInteger + toBuffer:dst::id{MTLBuffer} + destinationOffset:doff::NSUInteger + size:len::NSUInteger]::Nothing +end + +for T in (UInt8, Int8) + @eval function append_fillbuffer!(cce::MTL4ComputeCommandEncoder, buf::MTLBuffer, + value::$T, bytesize, offset=0) + range = NSRange(offset, bytesize) + @objc [cce::id{MTL4ComputeCommandEncoder} fillBuffer:buf::id{MTLBuffer} + range:range::NSRange + value:value::$T]::Nothing + end +end + +function append_fillbuffer!(cce::MTL4ComputeCommandEncoder, buf::MTLBuffer, range::NSRange, + value::UInt8) + @objc [cce::id{MTL4ComputeCommandEncoder} fillBuffer:buf::id{MTLBuffer} + range:range::NSRange + value:value::UInt8]::Nothing +end + +## residency + +function use!(cce::MTL4ComputeCommandEncoder, bufs::Vector{MTLBuffer}, + mode::MTLResourceUsage=ReadWriteUsage) + @objc [cce::id{MTL4ComputeCommandEncoder} useResources:bufs::id{MTLBuffer} + count:length(bufs)::NSUInteger + usage:mode::MTLResourceUsage]::Nothing +end diff --git a/lib/mtl/command_enc4.jl b/lib/mtl/command_enc4.jl new file mode 100644 index 000000000..6d61e4fcc --- /dev/null +++ b/lib/mtl/command_enc4.jl @@ -0,0 +1,95 @@ +export endEncoding!, updateFence!, waitForFence! +export barrierAfterEncoderStages!, barrierAfterQueueStages!, barrierAfterStages! + +# @objcwrapper managed = true MTL4CommandEncoder <: NSObject + +""" + updateFence!(encoder, fence, [afterEncoderStages]) + +Update `fence` once the given stages of the work encoded so far in `encoder` complete. +""" +function updateFence!(encoder::MTL4CommandEncoderLike, fence::MTLFence, + afterEncoderStages::MTLStages=MTLStageAll) + @objc [encoder::id{MTL4CommandEncoder} updateFence:fence::id{MTLFence} + afterEncoderStages:afterEncoderStages::MTLStages]::Nothing +end + +""" + waitForFence!(encoder, fence, [beforeEncoderStages]) + +Block the given stages of subsequently encoded work in `encoder` until `fence` is updated. +""" +function waitForFence!(encoder::MTL4CommandEncoderLike, fence::MTLFence, + beforeEncoderStages::MTLStages=MTLStageAll) + @objc [encoder::id{MTL4CommandEncoder} waitForFence:fence::id{MTLFence} + beforeEncoderStages:beforeEncoderStages::MTLStages]::Nothing +end + +""" + barrierAfterEncoderStages!(encoder, after, before, [visibilityOptions]) + +Order work *within* `encoder`: commands encoded after this barrier do not start their +`before` stages until the `after` stages of all previously encoded commands complete. + +Metal 4 command encoders execute their commands concurrently by default, so this is what +gives consecutive dispatches in one encoder the serial, program-order semantics that Metal +3's compute encoder provided implicitly. +""" +function barrierAfterEncoderStages!(encoder::MTL4CommandEncoderLike, + afterEncoderStages::MTLStages=MTLStageAll, + beforeEncoderStages::MTLStages=MTLStageAll, + visibilityOptions::MTL4VisibilityOptions=MTL4VisibilityOptionDevice) + @objc [encoder::id{MTL4CommandEncoder} barrierAfterEncoderStages:afterEncoderStages::MTLStages + beforeEncoderStages:beforeEncoderStages::MTLStages + visibilityOptions:visibilityOptions::MTL4VisibilityOptions]::Nothing +end + +""" + barrierAfterQueueStages!(encoder, after, before, [visibilityOptions]) + +Block the `before` stages of work encoded after this barrier until the `after` stages of +all work previously committed to the queue complete. +""" +function barrierAfterQueueStages!(encoder::MTL4CommandEncoderLike, + afterQueueStages::MTLStages=MTLStageAll, + beforeStages::MTLStages=MTLStageAll, + visibilityOptions::MTL4VisibilityOptions=MTL4VisibilityOptionDevice) + @objc [encoder::id{MTL4CommandEncoder} barrierAfterQueueStages:afterQueueStages::MTLStages + beforeStages:beforeStages::MTLStages + visibilityOptions:visibilityOptions::MTL4VisibilityOptions]::Nothing +end + +""" + barrierAfterStages!(encoder, after, beforeQueue, [visibilityOptions]) + +Block subsequently committed queue work until the `after` stages of the work encoded so far +in `encoder` complete. +""" +function barrierAfterStages!(encoder::MTL4CommandEncoderLike, + afterStages::MTLStages=MTLStageAll, + beforeQueueStages::MTLStages=MTLStageAll, + visibilityOptions::MTL4VisibilityOptions=MTL4VisibilityOptionDevice) + @objc [encoder::id{MTL4CommandEncoder} barrierAfterStages:afterStages::MTLStages + beforeQueueStages:beforeQueueStages::MTLStages + visibilityOptions:visibilityOptions::MTL4VisibilityOptions]::Nothing +end + +function push_debug_group!(encoder::MTL4CommandEncoderLike, name::Union{String,NSString}) + @objc [encoder::id{MTL4CommandEncoder} pushDebugGroup:name::id{NSString}]::Nothing +end + +function pop_debug_group!(encoder::MTL4CommandEncoderLike) + @objc [encoder::id{MTL4CommandEncoder} popDebugGroup]::Nothing +end + +endEncoding!(ce::MTL4CommandEncoderLike) = + @objc [ce::id{MTL4CommandEncoder} endEncoding]::Nothing + +function Base.close(ce::MTL4CommandEncoderLike) + try + endEncoding!(ce) + finally + release(ce) + end + return nothing +end diff --git a/lib/mtl/command_queue4.jl b/lib/mtl/command_queue4.jl new file mode 100644 index 000000000..7f3781896 --- /dev/null +++ b/lib/mtl/command_queue4.jl @@ -0,0 +1,174 @@ +# +# command queue descriptor +# + +export MTL4CommandQueueDescriptor + +# @objcwrapper managed = true MTL4CommandQueueDescriptor <: NSObject + +function MTL4CommandQueueDescriptor() + return @objc [MTL4CommandQueueDescriptor new]::MTL4CommandQueueDescriptor +end + +function MTL4CommandQueueDescriptor(label) + desc = MTL4CommandQueueDescriptor() + desc.label = label + return desc +end + + +# +# commit options +# + +export MTL4CommitOptions + +# @objcwrapper managed = true MTL4CommitOptions <: NSObject + +function MTL4CommitOptions() + return @objc [MTL4CommitOptions new]::MTL4CommitOptions +end + +function MTL4CommitOptions(f::Base.Callable) + options = MTL4CommitOptions() + add_feedback_handler!(f, options) + return options +end + +function _commit_feedback_callback(f) + # convert the incoming pointer, and discard any return value + function wrapper(ptr) + try + f(ptr == nil ? nothing : MTL4CommitFeedback(ptr)) + catch err + # we might be on an unmanaged thread here, so display the error + # (otherwise it may get lost, or worse, crash Julia) + @error "Commit feedback handler encountered an error: " * sprint(showerror, err) + end + return + end + @objcblock(wrapper, Nothing, (id{MTL4CommitFeedback},)) +end + +""" + add_feedback_handler!(options::MTL4CommitOptions) do feedback + ... + return + end + +Register a block that Metal invokes once the GPU finishes executing the command buffers +committed with `options`. Unlike Metal 3's `addCompletedHandler:`, the handler receives an +`MTL4CommitFeedback` rather than the command buffer, since Metal 4 command buffers are +recycled as soon as they are committed. +""" +function add_feedback_handler!(f::Base.Callable, options::MTL4CommitOptions) + block = _commit_feedback_callback(f) + @objc [options::id{MTL4CommitOptions} addFeedbackHandler:block::id{NSBlock}]::Nothing +end + +""" + add_feedback_handler!(options::MTL4CommitOptions, cond::Base.AsyncCondition) + +Signal `cond` once the GPU finishes executing the command buffers committed with `options`, +without running Julia code on Metal's feedback thread. +""" +function add_feedback_handler!(options::MTL4CommitOptions, cond::Base.AsyncCondition) + block = @objcasyncblock(cond) + @objc [options::id{MTL4CommitOptions} addFeedbackHandler:block::id{NSBlock}]::Nothing +end + + +# +# command queue +# + +export MTL4CommandQueue + +# @objcwrapper managed = true MTL4CommandQueue <: NSObject + +function MTL4CommandQueue(dev::MTLDevice) + return @objc [dev::id{MTLDevice} newMTL4CommandQueue]::MTL4CommandQueue +end + +function MTL4CommandQueue(dev::MTLDevice, desc::MTL4CommandQueueDescriptor) + err = Ref{id{NSError}}(nil) + queue = @objc [dev::id{MTLDevice} newMTL4CommandQueueWithDescriptor:desc::id{MTL4CommandQueueDescriptor} + error:err::Ptr{id{NSError}}]::Union{Nothing,MTL4CommandQueue} + queue === nothing && throw_error(err[]) + return queue +end + +function add_residency_set!(queue::MTL4CommandQueue, resset::MTLResidencySet) + @objc [queue::id{MTL4CommandQueue} addResidencySet:resset::id{MTLResidencySet}]::Nothing +end + +function remove_residency_set!(queue::MTL4CommandQueue, resset::MTLResidencySet) + @objc [queue::id{MTL4CommandQueue} removeResidencySet:resset::id{MTLResidencySet}]::Nothing +end + +""" + commit!(queue::MTL4CommandQueue, cmdbuf::MTL4CommandBuffer, [options]) + +Submit `cmdbuf` for execution on `queue`. The command buffer must have been ended with +[`endCommandBuffer!`](@ref); it may be reused for encoding as soon as this returns, but its +allocator may only be reset once the GPU signals completion (see `MTL4CommitOptions`). +""" +function commit!(queue::MTL4CommandQueue, cmdbuf::MTL4CommandBuffer) + ref = Ref(pointer(cmdbuf)) + GC.@preserve cmdbuf begin + @objc [queue::id{MTL4CommandQueue} commit:ref::Ptr{id{MTL4CommandBuffer}} + count:1::NSUInteger]::Nothing + end + return +end + +function commit!(queue::MTL4CommandQueue, cmdbuf::MTL4CommandBuffer, + options::MTL4CommitOptions) + ref = Ref(pointer(cmdbuf)) + GC.@preserve cmdbuf begin + @objc [queue::id{MTL4CommandQueue} commit:ref::Ptr{id{MTL4CommandBuffer}} + count:1::NSUInteger + options:options::id{MTL4CommitOptions}]::Nothing + end + return +end + +function commit!(queue::MTL4CommandQueue, cmdbufs::Vector{MTL4CommandBuffer}) + ptrs = map(pointer, cmdbufs) + GC.@preserve cmdbufs ptrs begin + @objc [queue::id{MTL4CommandQueue} commit:ptrs::Ptr{id{MTL4CommandBuffer}} + count:length(ptrs)::NSUInteger]::Nothing + end + return +end + +function commit!(queue::MTL4CommandQueue, cmdbufs::Vector{MTL4CommandBuffer}, + options::MTL4CommitOptions) + ptrs = map(pointer, cmdbufs) + GC.@preserve cmdbufs ptrs begin + @objc [queue::id{MTL4CommandQueue} commit:ptrs::Ptr{id{MTL4CommandBuffer}} + count:length(ptrs)::NSUInteger + options:options::id{MTL4CommitOptions}]::Nothing + end + return +end + +""" + signal_event!(queue::MTL4CommandQueue, ev::MTLEvent, val::Integer) + +Enqueue an update of `ev` to `val` once all previously committed work on `queue` completes. +""" +function signal_event!(queue::MTL4CommandQueue, ev::MTLEventLike, val::Integer) + @objc [queue::id{MTL4CommandQueue} signalEvent:ev::id{MTLEvent} + value:val::UInt64]::Nothing +end + +""" + wait_event!(queue::MTL4CommandQueue, ev::MTLEvent, val::Integer) + +Block subsequently committed work on `queue` until `ev` reaches `val`. +""" +function wait_event!(queue::MTL4CommandQueue, ev::MTLEventLike, val::Integer) + @objc [queue::id{MTL4CommandQueue} waitForEvent:ev::id{MTLEvent} + value:val::UInt64]::Nothing +end diff --git a/lib/mtl/compute_pipeline4.jl b/lib/mtl/compute_pipeline4.jl new file mode 100644 index 000000000..fa0123dce --- /dev/null +++ b/lib/mtl/compute_pipeline4.jl @@ -0,0 +1,22 @@ +# +# compute pipeline descriptor +# + +export MTL4ComputePipelineDescriptor, MTL4PipelineOptions + +# @objcwrapper managed = true MTL4PipelineOptions <: NSObject + +function MTL4PipelineOptions() + return @objc [MTL4PipelineOptions new]::MTL4PipelineOptions +end + +# @objcwrapper managed = true MTL4ComputePipelineDescriptor <: MTL4PipelineDescriptor + +function MTL4ComputePipelineDescriptor() + return @objc [MTL4ComputePipelineDescriptor new]::MTL4ComputePipelineDescriptor +end + +# NOTE: Metal.jl builds its `MTLComputePipelineState`s through the Metal 3 device API (see +# `link_pipeline`), because that is what supports the binary archives used to cache +# compiled kernels across sessions. Metal 4 command encoders accept those pipeline +# states unchanged, so the `MTL4Compiler` path is not needed to run kernels. diff --git a/src/command_batching.jl b/src/command_batching.jl index fb4db00f7..a737101b3 100644 --- a/src/command_batching.jl +++ b/src/command_batching.jl @@ -1,5 +1,39 @@ +# Metal 4 command batching. +# +# Metal.jl submits its own GPU work -- kernel launches, buffer copies and fills -- through +# Metal 4. A `BatchedCommandQueue` owns an `MTL4CommandQueue` and encodes operations into a +# single open `MTL4CommandBuffer`, backed by a recycled `MTL4CommandAllocator` and bound +# through a shared `MTL4ArgumentTable`, committing lazily. +# +# Three Metal 4 differences drive the design: +# +# * command buffers carry no completion state, so a committed batch is tracked by an +# `Mtl4Submission` that a commit-feedback handler marks completed (and annotates with +# GPU timings and any error); +# * commands in an encoder run concurrently unless separated by a barrier, so every +# operation but the first in a batch is preceded by one, reproducing the serial +# semantics of Metal 3's compute encoder; +# * binding a buffer no longer makes it resident, so every buffer an operation touches is +# added to the batch's residency set. +# +# The queue also keeps an `MTLCommandQueue` for interoperability -- MPS, MPSGraph, and any +# command buffer derived from the queue by user code. The two are ordered against each +# other with a pair of shared events: a flushed Metal 4 batch signals `event`, which is +# also what tracks its completion and what a Metal 3 command buffer derived from the queue +# waits on; that buffer in turn signals `event3` when committed, which the next Metal 4 +# batch waits on. + @enum EncoderKind NoEncoder ComputeEncoder BlitEncoder +# The stage mask used for the barriers that reimpose program order. `MTL4CommandEncoder`'s +# own `stages` property reports 0 on current drivers, so ask for a full barrier rather than +# narrowing it to the dispatch and blit stages a compute encoder actually uses. +const BARRIER_STAGES = MTL.MTLStageAll + +# Label of every batched command buffer (only applied when `label_resources()`), and the +# name a submission reports under in profiles and errors either way. +const BATCH_LABEL = "MTL4CommandBuffer(batched queue)" + # Tunables for command batching (see `BatchedCommandQueue`), read once per # process. Numeric tuning knobs can also be set for one run via the matching # JULIA_METAL_* env var (read once at startup), which is convenient for @@ -39,24 +73,122 @@ command_batching_inflight() = @memoize begin @load_preference("command_batching_inflight", 3)) end::Int -# Completion handlers run on libdispatch worker threads, where compiling or -# running Julia code can overflow the small foreign stack. Keep command buffers -# alive and drain Julia roots from normal managed threads instead. -struct PendingCommand - cmdbuf::MTL.MTLCommandBufferLike + +## by-value argument storage + +# Metal 4 argument tables bind GPU addresses only; there is no `setBytes:`. Arguments that +# are passed by value are therefore bump-allocated out of a shared-storage scratch buffer, +# whose address is bound instead. A scratch buffer belongs to the batch that filled it and +# is recycled once that batch completes. + +const ARGUMENT_SCRATCH_SIZE = 1024 * 1024 +const ARGUMENT_ALIGNMENT = 64 + +mutable struct ArgumentScratch + const buf::MTLBuffer + const ptr::Ptr{UInt8} + const addr::UInt64 + const capacity::Int + offset::Int +end + +function ArgumentScratch(dev::MTLDevice, capacity::Integer=ARGUMENT_SCRATCH_SIZE) + buf = @autoreleasepool MTLBuffer(dev, capacity; storage=SharedStorage) + @label! buf "argument scratch" + ArgumentScratch(buf, convert(Ptr{UInt8}, MTL.contents(buf)), + UInt64(buf.gpuAddress), Int(capacity), 0) +end + +# Reserve `nbytes` of aligned storage, returning the CPU pointer to write through and the +# GPU address to bind, or `nothing` when this scratch buffer is full. +@inline function reserve!(scratch::ArgumentScratch, nbytes::Int) + offset = (scratch.offset + ARGUMENT_ALIGNMENT - 1) & ~(ARGUMENT_ALIGNMENT - 1) + offset + nbytes > scratch.capacity && return nothing + scratch.offset = offset + nbytes + return (scratch.ptr + offset, scratch.addr + offset) +end + +reset!(scratch::ArgumentScratch) = (scratch.offset = 0; scratch) + + +## committed batches + +# Metal 4 recycles command buffers on commit and reports completion out-of-band, so the +# per-batch state that Metal 3 would hang off an `MTLCommandBuffer` lives here instead. +# `completed` is set by the commit-feedback handler *after* the error and timings have been +# recorded, so observing it is enough to safely read the rest. +mutable struct Mtl4Submission + const label::String + const cmdbuf::MTL4CommandBuffer + const allocator::MTL4CommandAllocator + const resset::MTLResidencySet + const scratch::Vector{ArgumentScratch} + # value the queue's ordering event reaches once the GPU is done with this batch + const event::MTLSharedEvent + const seq::UInt64 + # whether a commit-feedback handler was registered for this batch (see `commit_options`) + const expect_feedback::Bool roots::Vector{Any} + Base.@atomic completed::Bool + Base.@atomic gpu_start::Float64 + Base.@atomic gpu_end::Float64 + Base.@atomic error::Union{Nothing,MTL.CommandBufferErrorInfo} +end + +# The ordering event is the authoritative completion signal: unlike the commit-feedback +# handler it needs no callback into Julia, which a precompilation worker cannot service. +# `diagnosed` additionally reports whether the handler has run and filled in the timings +# and error, which callers wait for on a bounded basis only. +is_completed(sub::Mtl4Submission) = + (Base.@atomic sub.completed) || sub.event.signaledValue >= sub.seq + +diagnosed(sub::Mtl4Submission) = !sub.expect_feedback || (Base.@atomic sub.completed) + +gpu_time_range(sub::Mtl4Submission) = + diagnosed(sub) ? ((Base.@atomic sub.gpu_start), (Base.@atomic sub.gpu_end)) : nothing + +function gpu_time_range(cmdbuf::MTL.MTLCommandBufferLike) + cmdbuf.status == MTL.MTLCommandBufferStatusCompleted || return nothing + return (cmdbuf.GPUStartTime, cmdbuf.GPUEndTime) +end + +# Runs on one of Metal's feedback threads: keep it to Objective-C property reads and field +# stores, and never touch the Julia roots (`drain_cleanups!` releases those from the owning +# task instead). +# +# Metal invokes the handler on a libdispatch worker, which a precompilation worker cannot +# adopt without hanging image serialization, so `flush!` registers no handler there. +# Completion is tracked by the ordering event either way; only the timings and error are +# lost. +function commit_options(sub::Mtl4Submission, queue_label::Union{Nothing,String}) + return MTL.MTL4CommitOptions() do feedback + if feedback !== nothing + Base.@atomic sub.gpu_start = feedback.GPUStartTime + Base.@atomic sub.gpu_end = feedback.GPUEndTime + err = feedback.error + if err !== nothing + Base.@atomic sub.error = MTL.CommandBufferErrorInfo( + String(err.domain), Int(err.code), String(err.localizedDescription), + sub.label, queue_label) + end + end + Base.@atomic sub.completed = true + return + end end + """ BatchedCommandQueue A command queue that batches GPU work to amortize Metal's per-command-buffer -submission latency. Kernel launches (`@metal`) and GPU-side blit operations -(`copyto!`, `fill!`) are encoded into a single open command buffer and committed +submission latency. Kernel launches (`@metal`) and GPU-side copy and fill operations +(`copyto!`, `fill!`) are encoded into a single open Metal 4 command buffer and committed lazily, instead of one command buffer per operation. -It wraps, and is a drop-in for, the `MTLCommandQueue` it batches: properties that -aren't its own (e.g. `label`) forward to the underlying queue. +It wraps, and is a drop-in for, an `MTLCommandQueue`: properties that aren't its own +(e.g. `label`) forward to that queue, and command buffers derived from it interoperate with +the batched Metal 4 work in program order. The open batch is committed ("flushed") when any of the following happens: @@ -64,16 +196,16 @@ The open batch is committed ("flushed") when any of the following happens: * a command buffer derived from the queue is enqueued or committed; * command batching is disabled with the `command_batching = false` preference; * the batch reaches `command_batching_ops()` operations or - `command_batching_bytes()` of blit traffic; + `command_batching_bytes()` of copy traffic; * a GPU profiler is attached, in which case batching is disabled (each operation gets its own command buffer) so per-operation GPU timing is preserved; * an immediate submission is requested via `@metal ... submit=true`, or `Metal.flush!` is called explicitly. -Program order is preserved across flushes: command buffers execute in commit order -and dispatches within an encoder run serially. At most `command_batching_inflight()` -command buffers are kept in flight; further submissions block until the GPU drains -one. Obtain the current task's batched queue with [`global_queue`](@ref). +Program order is preserved across flushes: command buffers execute in commit order, and +operations within a batch are separated by encoder barriers. At most +`command_batching_inflight()` command buffers are kept in flight; further submissions block +until the GPU drains one. Obtain the current task's batched queue with [`global_queue`](@ref). `BatchedCommandQueue`s are task-local and mutated lock-free by their owning task. Sharing a raw `MTLCommandQueue` across tasks is unsupported. [`device_synchronize`](@ref) @@ -81,24 +213,93 @@ may flush batches owned by other tasks after those tasks have yielded or complet which supports `@async` work and the REPL synchronization hook. """ mutable struct BatchedCommandQueue + # Metal 3 queue, kept for interoperability and as this queue's identity queue::MTLCommandQueue + # Metal 4 queue, carrying all of Metal.jl's own work + queue4::MTL4CommandQueue device::MTLDevice - cmdbuf::Union{Nothing,MTLCommandBuffer} - encoder::Union{Nothing,MTLComputeCommandEncoder,MTLBlitCommandEncoder} + + # open batch + cmdbuf::Union{Nothing,MTL4CommandBuffer} + allocator::Union{Nothing,MTL4CommandAllocator} + encoder::Union{Nothing,MTL4ComputeCommandEncoder} + resset::Union{Nothing,MTLResidencySet} + scratch::Vector{ArgumentScratch} kind::EncoderKind + needs_barrier::Bool roots::Vector{Any} last_pipeline::Union{Nothing,MTLComputePipelineState} nops::Int nbytes::Int pending_ops::Vector{Any} - cleanups::Vector{PendingCommand} -end + + # argument binding + argtable::MTL4ArgumentTable + + # submitted batches awaiting completion, in commit order + cleanups::Vector{Mtl4Submission} + errors::Union{Nothing,Vector{MTL.CommandBufferErrorInfo}} + # the queue's label, as reported in `CommandBufferError`s; fixed at construction + queue_label::Union{Nothing,String} + + # completion of, and ordering against, batched Metal 4 work. `order` is bumped once per + # flush and is what an `Mtl4Submission`'s `seq` refers to, so nothing else may signal + # this event. + event::MTLSharedEvent + order::UInt64 + pending4::UInt64 + # Metal 3 command buffers derived from this queue and not yet committed, keyed by + # pointer, with the `pending4` value each was made to wait on at derivation + derived::Dict{UInt,UInt64} + + # ordering of batched Metal 4 work after Metal 3 work derived from `queue` + event3::MTLSharedEvent + order3::UInt64 + pending3::UInt64 + + # recycled Metal 4 objects, returned by completed submissions + free_cmdbufs::Vector{MTL4CommandBuffer} + free_allocators::Vector{MTL4CommandAllocator} + free_ressets::Vector{MTLResidencySet} + free_scratch::Vector{ArgumentScratch} +end + +# The largest number of buffer bindings a kernel launch can use. Metal 4 caps argument +# tables at 31 buffer bindings, which matches the Metal 3 buffer-argument limit that +# `mtlfunction` already compiles against. +const MAX_BUFFER_BINDINGS = 31 function BatchedCommandQueue(queue::MTLCommandQueue) dev = queue.device - can_use_residency_sets(dev) && install_queue_residency!(queue, dev) - BatchedCommandQueue(queue, dev, nothing, nothing, NoEncoder, - Any[], nothing, 0, 0, Any[], PendingCommand[]) + @autoreleasepool begin + desc = MTL4CommandQueueDescriptor() + queue_label = let label = queue.label + label === nothing ? nothing : String(label) + end + queue_label === nothing || (desc.label = queue_label) + queue4 = MTL4CommandQueue(dev, desc) + + argtable = MTL4ArgumentTable(dev; buffers=MAX_BUFFER_BINDINGS) + @label! argtable "batched queue arguments" + + event = MTLSharedEvent(dev) + @label! event "batched queue completion" + event3 = MTLSharedEvent(dev) + @label! event3 "batched queue interop ordering" + + bq = BatchedCommandQueue(queue, queue4, dev, + nothing, nothing, nothing, nothing, ArgumentScratch[], + NoEncoder, false, + Any[], nothing, 0, 0, Any[], + argtable, + Mtl4Submission[], nothing, queue_label, + event, UInt64(0), UInt64(0), Dict{UInt,UInt64}(), + event3, UInt64(0), UInt64(0), + MTL4CommandBuffer[], MTL4CommandAllocator[], + MTLResidencySet[], ArgumentScratch[]) + install_queue_residency!(bq) + return bq + end end # Properties that aren't our own fields (e.g. `label`) forward to the wrapped @@ -116,6 +317,15 @@ end const batched_queues = IdDict{BatchedCommandQueue,Nothing}() const batched_queues_lock = ReentrantLock() +# Metal 3 command buffers (MPS work, user-derived buffers) still report completion +# themselves; their Julia roots are released by `drain_cleanups!` alongside Metal 4 batches. +struct PendingCommand + cmdbuf::MTL.MTLCommandBufferLike + roots::Vector{Any} +end + +const pending_commands = IdDict{BatchedCommandQueue,Vector{PendingCommand}}() + @inline batched_queue_key(queue::MTLCommandQueue) = (:BatchedCommandQueue, pointer(queue)) @@ -126,12 +336,16 @@ function register_queue!(bq::BatchedCommandQueue) return end +# A queue with unreported errors stays registered so `device_synchronize` on another +# task still surfaces them. +queue_is_idle(bq::BatchedCommandQueue) = + bq.cmdbuf === nothing && isempty(bq.cleanups) && !haskey(pending_commands, bq) && + bq.errors === nothing + function unregister_queue_if_idle!(bq::BatchedCommandQueue) - (bq.cmdbuf === nothing && isempty(bq.cleanups)) || return + queue_is_idle(bq) || return Base.@lock batched_queues_lock begin - if bq.cmdbuf === nothing && isempty(bq.cleanups) - delete!(batched_queues, bq) - end + queue_is_idle(bq) && delete!(batched_queues, bq) end return end @@ -157,22 +371,52 @@ raw_queue(queue::MTLCommandQueue) = queue profiling_command_buffers() = MTL.profile_hook[] !== nothing || MTL.profile_metadata[] !== nothing +# A batched queue stands for both of the Metal queues it owns: the Metal 3 one it wraps, +# and the Metal 4 one it submits its own work to (which is what frame capture sees). +Base.:(==)(bq::BatchedCommandQueue, obj::NSObject) = + UInt(pointer(obj)) in (UInt(pointer(bq.queue)), UInt(pointer(bq.queue4))) +Base.:(==)(obj::NSObject, bq::BatchedCommandQueue) = bq == obj +Base.:(==)(bq::BatchedCommandQueue, queue::MTLCommandQueue) = bq.queue == queue +Base.:(==)(queue::MTLCommandQueue, bq::BatchedCommandQueue) = queue == bq.queue + + +## interoperability with Metal 3 command buffers +# +# A command buffer derived from the queue is ordered after the batched Metal 4 work in +# two steps. At derivation it gets a GPU-side wait on the batch flushed at that moment +# (`order_after_batch!`). Batches opened afterwards are only flushed when the buffer +# commits, by which point its commands are already encoded and a wait could no longer +# precede them -- so the commit hook falls back to a host-side wait for those. Buffers +# derived from the raw `MTLCommandQueue` are unknown to the hook and always take the +# host-side path; the constructors below (and `MPSCommandBuffer(::BatchedCommandQueue)`) +# are the fast path and cover Metal.jl's own MPS and MPSGraph use. + function Base.cconvert(::Type{<:id{MTLCommandQueue}}, bq::BatchedCommandQueue) flush!(bq) + wait_submissions!(bq) return bq.queue end -Base.:(==)(bq::BatchedCommandQueue, queue::MTLCommandQueue) = bq.queue == queue -Base.:(==)(queue::MTLCommandQueue, bq::BatchedCommandQueue) = queue == bq.queue -Base.:(==)(bq::BatchedCommandQueue, obj::NSObject) = bq.queue == obj -Base.:(==)(obj::NSObject, bq::BatchedCommandQueue) = obj == bq.queue +# The key the commit hooks will see `cmdbuf` under: they are handed the underlying +# `MTLCommandBuffer`, so wrappers like `MPSCommandBuffer` must register that one. +derived_key(cmdbuf::MTL.MTLCommandBufferLike) = UInt(pointer(cmdbuf)) + +# Order `cmdbuf`, which was just derived from this queue, after the batched Metal 4 work +# committed so far, and remember what it waited on for the commit hook. +function order_after_batch!(bq::BatchedCommandQueue, cmdbuf::MTL.MTLCommandBufferLike, + key::UInt=derived_key(cmdbuf)) + flush!(bq) + bq.pending4 == 0 || MTL.encode_wait!(cmdbuf, bq.event, bq.pending4) + bq.derived[key] = bq.pending4 + return cmdbuf +end function MTL.MTLCommandBuffer(bq::BatchedCommandQueue) - return MTLCommandBuffer(bq.queue) + return order_after_batch!(bq, MTLCommandBuffer(bq.queue)) end function MTL.MTLCommandBuffer(bq::BatchedCommandQueue, desc::MTLCommandBufferDescriptor) - return MTLCommandBuffer(bq.queue, desc) + return order_after_batch!(bq, MTLCommandBuffer(bq.queue, desc)) end function MTL.MTLCommandBuffer(f::Base.Callable, bq::BatchedCommandQueue, @@ -182,20 +426,49 @@ function MTL.MTLCommandBuffer(f::Base.Callable, bq::BatchedCommandQueue, return cmdbuf end +# `MTL.submit_hook`: a Metal 3 command buffer on this queue is about to be enqueued or +# committed, so the open Metal 4 batch has to go out first. function flush_open_batch(cmdbuf) - queue = cmdbuf.commandQueue - bq = get(task_local_storage(), batched_queue_key(queue), nothing) + bq = lookup_batched_queue(cmdbuf) bq === nothing || flush!(bq) return end +# `MTL.commit_hook`: a Metal 3 command buffer is about to be committed, with all of its +# encoders closed. Make sure every batch it should follow has actually finished (see the +# section comment), then have it signal the ordering event so the next Metal 4 batch can +# wait on it. +function order_metal4_after(cmdbuf) + bq = lookup_batched_queue(cmdbuf) + bq === nothing && return + + flush!(bq) + waited = pop!(bq.derived, derived_key(cmdbuf), UInt64(0)) + bq.pending4 > waited && wait_submissions!(bq) + + value = (bq.order3 += 1) + MTL.encode_signal!(cmdbuf, bq.event3, value) + bq.pending3 = value + return +end + +function lookup_batched_queue(cmdbuf) + queue = cmdbuf.commandQueue + queue === nothing && return nothing + return get(task_local_storage(), batched_queue_key(queue), nothing) +end + +# Xcode's capture tooling does not accept a scope created from an `MTL4CommandQueue` +# (`startCaptureWithDescriptor:` sends it `traceStream`, which such a scope does not +# implement, aborting the process), so scope-based capture falls back to a device-wide +# scope. That still covers this queue's work, since the device only drives one GPU. function MTL.MTLCaptureScope(bq::BatchedCommandQueue, manager=MTLCaptureManager()) flush!(bq) - return MTLCaptureScope(bq.queue, manager) + return MTLCaptureScope(bq.device, manager) end function Base.setproperty!(desc::MTLCaptureDescriptor, name::Symbol, bq::BatchedCommandQueue) - name === :captureObject && return setproperty!(desc, name, bq.queue) + name === :captureObject && return setproperty!(desc, name, bq.queue4) return invoke(Base.setproperty!, Tuple{MTLCaptureDescriptor, Symbol, Any}, desc, name, bq) end @@ -203,25 +476,54 @@ end function MTL.MTLCaptureDescriptor(bq::BatchedCommandQueue, destination::MTL.MTLCaptureDestination; folder::String=nothing) - return MTLCaptureDescriptor(bq.queue, destination; folder) + return MTLCaptureDescriptor(bq.queue4, destination; folder) end function MTL.startCapture(bq::BatchedCommandQueue, destination::MTL.MTLCaptureDestination=MTL.MTLCaptureDestinationGPUTraceDocument; folder::String=nothing) flush!(bq) - return MTL.startCapture(bq.queue, destination; folder) + return MTL.startCapture(bq.queue4, destination; folder) end -function ensure_cmdbuf!(bq::BatchedCommandQueue) + +## the open batch + +function ensure_cmdbuf!(bq::BatchedCommandQueue, + options::Union{Nothing,MTL.MTL4CommandBufferOptions}=nothing) cmdbuf = bq.cmdbuf - if cmdbuf === nothing - cmdbuf = MTLCommandBuffer(bq.queue) - @label! cmdbuf "MTLCommandBuffer(batched queue)" - bq.cmdbuf = cmdbuf - register_queue!(bq) + cmdbuf === nothing || return cmdbuf::MTL4CommandBuffer + + # order this batch after any Metal 3 work committed on the interop queue + if bq.pending3 != 0 + MTL.wait_event!(bq.queue4, bq.event3, bq.pending3) + bq.pending3 = UInt64(0) end - return cmdbuf::MTLCommandBuffer + + cmdbuf = isempty(bq.free_cmdbufs) ? MTL4CommandBuffer(bq.device) : pop!(bq.free_cmdbufs) + @label! cmdbuf BATCH_LABEL + allocator = isempty(bq.free_allocators) ? MTL4CommandAllocator(bq.device) : + pop!(bq.free_allocators) + resset = if isempty(bq.free_ressets) + desc = MTLResidencySetDescriptor() + desc.initialCapacity = 16 + @label! desc "batched queue residency" + MTLResidencySet(bq.device, desc) + else + pop!(bq.free_ressets) + end + + if options === nothing + MTL.beginCommandBuffer!(cmdbuf, allocator) + else + MTL.beginCommandBuffer!(cmdbuf, allocator, options) + end + bq.cmdbuf = cmdbuf + bq.allocator = allocator + bq.resset = resset + bq.needs_barrier = false + register_queue!(bq) + return cmdbuf end function end_encoder!(bq::BatchedCommandQueue) @@ -231,35 +533,56 @@ function end_encoder!(bq::BatchedCommandQueue) bq.encoder = nothing bq.kind = NoEncoder bq.last_pipeline = nothing + bq.needs_barrier = false close(enc) return end -function compute_encoder(bq::BatchedCommandQueue) - if bq.kind == ComputeEncoder - return bq.encoder::MTLComputeCommandEncoder +# Metal 4 encodes dispatches and buffer copies through the same encoder, so the two +# accessors below hand out the same one, and `kind` only records what it was last used for. +function batch_encoder!(bq::BatchedCommandQueue, kind::EncoderKind) + enc = bq.encoder + if enc === nothing + enc = MTL4ComputeCommandEncoder(ensure_cmdbuf!(bq)) + MTL.set_argument_table!(enc, bq.argtable) + bq.encoder = enc + # Metal 4 makes no ordering promise between command buffers either, so fence this + # batch against everything already committed to the queue. + MTL.barrierAfterQueueStages!(enc, BARRIER_STAGES, BARRIER_STAGES, + MTL.MTL4VisibilityOptionDevice) + bq.needs_barrier = false + elseif bq.needs_barrier + # Metal 4 encoders run their commands concurrently; keep Metal.jl's program-order + # guarantee by fencing every operation against the ones already encoded. + MTL.barrierAfterEncoderStages!(enc, BARRIER_STAGES, BARRIER_STAGES, + MTL.MTL4VisibilityOptionDevice) + bq.needs_barrier = false end - - end_encoder!(bq) - enc = MTLComputeCommandEncoder(ensure_cmdbuf!(bq)) - bq.encoder = enc - bq.kind = ComputeEncoder - return enc + bq.kind = kind + return enc::MTL4ComputeCommandEncoder end -function blit_encoder(bq::BatchedCommandQueue) - if bq.kind == BlitEncoder - return bq.encoder::MTLBlitCommandEncoder - end +compute_encoder(bq::BatchedCommandQueue) = batch_encoder!(bq, ComputeEncoder) +blit_encoder(bq::BatchedCommandQueue) = batch_encoder!(bq, BlitEncoder) +# Start a batch of this operation's own, with command-buffer options that cannot be shared +# with unrelated work (currently only the `MTLLogState` of a logging-enabled kernel). +function begin_batch!(bq::BatchedCommandQueue; + options::Union{Nothing,MTL.MTL4CommandBufferOptions}=nothing) + flush!(bq) + ensure_cmdbuf!(bq, options) + return compute_encoder(bq) +end + +# Throw away a partially encoded batch, keeping the queue usable. +function abort_batch!(bq::BatchedCommandQueue) end_encoder!(bq) - enc = MTLBlitCommandEncoder(ensure_cmdbuf!(bq)) - bq.encoder = enc - bq.kind = BlitEncoder - return enc + cmdbuf = bq.cmdbuf + cmdbuf === nothing || discard_open_cmdbuf!(bq, cmdbuf) + return end -function set_pipeline!(bq::BatchedCommandQueue, cce::MTLComputeCommandEncoder, +function set_pipeline!(bq::BatchedCommandQueue, cce::MTL4ComputeCommandEncoder, pipeline::MTLComputePipelineState) if bq.last_pipeline !== pipeline MTL.set_function!(cce, pipeline) @@ -268,6 +591,48 @@ function set_pipeline!(bq::BatchedCommandQueue, cce::MTLComputeCommandEncoder, return end +# Declare `buf` resident for the open batch. Metal 4 argument tables bind bare GPU +# addresses, so nothing is made resident implicitly. +@inline function make_resident!(bq::BatchedCommandQueue, buf::MTLBuffer) + bq.resset === nothing && ensure_cmdbuf!(bq) + MTL.add_allocation!(bq.resset::MTLResidencySet, buf) + return +end + +@inline function bind_buffer!(bq::BatchedCommandQueue, buf::MTLBuffer, offset::Integer, + index::Integer) + make_resident!(bq, buf) + MTL.set_buffer!(bq.argtable, buf, offset, index) + return +end + +# Bind an `isbits` value by copying it into the batch's argument scratch buffer, Metal 4 +# having no equivalent of `setBytes:`. +@inline function bind_bytes!(bq::BatchedCommandQueue, ptr::Ptr{Cvoid}, nbytes::Int, + index::Integer) + addr = reserve_arguments!(bq, nbytes) + unsafe_copyto!(reinterpret(Ptr{UInt8}, addr[1]), reinterpret(Ptr{UInt8}, ptr), nbytes) + MTL.set_address!(bq.argtable, addr[2], index) + return +end + +function reserve_arguments!(bq::BatchedCommandQueue, nbytes::Int) + nbytes <= ARGUMENT_SCRATCH_SIZE || + throw(ArgumentError("Kernel argument of $nbytes bytes exceeds the $(ARGUMENT_SCRATCH_SIZE)-byte argument buffer")) + + scratch = isempty(bq.scratch) ? nothing : last(bq.scratch) + if scratch !== nothing + res = reserve!(scratch, nbytes) + res === nothing || return res + end + + scratch = isempty(bq.free_scratch) ? ArgumentScratch(bq.device) : + reset!(pop!(bq.free_scratch)) + push!(bq.scratch, scratch) + make_resident!(bq, scratch.buf) + return reserve!(scratch, nbytes)::Tuple{Ptr{UInt8},UInt64} +end + function note_operation!(bq::BatchedCommandQueue, op) MTL.profile_metadata[] === nothing && return push!(bq.pending_ops, op) @@ -279,21 +644,32 @@ function record_operation!(bq::BatchedCommandQueue, roots...; bytes::Integer=0, op === nothing || note_operation!(bq, op) bq.nops += 1 bq.nbytes += bytes + bq.needs_barrier = true return end -function register_operations!(bq::BatchedCommandQueue, cmdbuf) +function register_operations!(bq::BatchedCommandQueue, sub) md = MTL.profile_metadata[] md === nothing && return for op in bq.pending_ops - MTL.note_operation!(md, cmdbuf, op) + MTL.note_operation!(md, sub, op) end return end + +## completion tracking + +function defer_cleanup!(bq::BatchedCommandQueue, sub::Mtl4Submission) + push!(bq.cleanups, sub) + register_queue!(bq) + return +end + function defer_cleanup!(bq::BatchedCommandQueue, cmdbuf::MTL.MTLCommandBufferLike, roots::Vector{Any}) - push!(bq.cleanups, PendingCommand(cmdbuf, roots)) + push!(get!(() -> PendingCommand[], pending_commands, bq), + PendingCommand(cmdbuf, roots)) register_queue!(bq) return end @@ -301,21 +677,72 @@ end defer_cleanup!(queue, cmdbuf::MTL.MTLCommandBufferLike, roots::Vector{Any}) = defer_cleanup!(batched_queue(queue), cmdbuf, roots) +function recycle!(bq::BatchedCommandQueue, sub::Mtl4Submission) + empty!(sub.roots) + + err = Base.@atomic sub.error + if err !== nothing + errors = bq.errors + if errors === nothing + bq.errors = MTL.CommandBufferErrorInfo[err] + else + push!(errors, err) + end + end + + # Resetting the allocator invalidates the batch's commands, which is only legal now + # that the GPU is done with them. + MTL.reset!(sub.allocator) + push!(bq.free_allocators, sub.allocator) + push!(bq.free_cmdbufs, sub.cmdbuf) + + MTL.remove_all_allocations!(sub.resset) + MTL.commit!(sub.resset) + push!(bq.free_ressets, sub.resset) + + for scratch in sub.scratch + push!(bq.free_scratch, reset!(scratch)) + end + return +end + +# A submission is only recycled once the GPU is done with it *and* its commit-feedback +# handler has recorded the timings and error; recycling on completion alone would let a +# late handler write its error into an already-forgotten object. Forcing waits for both. +recyclable(sub::Mtl4Submission) = is_completed(sub) && diagnosed(sub) + function drain_cleanups!(bq::BatchedCommandQueue; force::Bool=false) n = 0 - for cleanup in bq.cleanups - if !(force || cleanup.cmdbuf.status >= MTL.MTLCommandBufferStatusCompleted) + for sub in bq.cleanups + if force + wait_diagnosed!(sub) + elseif !recyclable(sub) break end n += 1 end - n == 0 && return - - completed = bq.cleanups[1:n] - deleteat!(bq.cleanups, 1:n) + if n > 0 + completed = bq.cleanups[1:n] + deleteat!(bq.cleanups, 1:n) + for sub in completed + recycle!(bq, sub) + end + end - for cleanup in completed - empty!(cleanup.roots) + pending = get(pending_commands, bq, nothing) + if pending !== nothing && !isempty(pending) + m = 0 + for cleanup in pending + (force || cleanup.cmdbuf.status >= MTL.MTLCommandBufferStatusCompleted) || break + m += 1 + end + if m > 0 + for cleanup in view(pending, 1:m) + empty!(cleanup.roots) + end + deleteat!(pending, 1:m) + end + isempty(pending) && delete!(pending_commands, bq) end unregister_queue_if_idle!(bq) @@ -331,14 +758,67 @@ function drain_cleanups!(queue; force::Bool=false) return end -function pending_cleanup_count(bq::BatchedCommandQueue) - return length(bq.cleanups) +pending_cleanup_count(bq::BatchedCommandQueue) = length(bq.cleanups) + +# Metal 4 has no pollable command-buffer status, so completion is observed through the +# queue's ordering event; the commit-feedback handler is then given a bounded chance to +# record the batch's timings and error before the submission is recycled. +# +# Like `wait_cmdbuf!`, this polls from the Julia scheduler unless the +# `nonblocking_synchronization` preference is off or a precompilation worker is running, +# in which case it parks the thread in Metal's blocking event wait instead. +function wait_submission!(sub::Mtl4Submission) + if !is_completed(sub) + precompiling = ccall(:jl_generating_output, Cint, ()) != 0 + if use_nonblocking_synchronization && !precompiling + spins = 0 + while spins < 256 + if spins < 32 + ccall(:jl_cpu_pause, Cvoid, ()) + ccall(:jl_gc_safepoint, Cvoid, ()) + else + yield() + end + is_completed(sub) && break + spins += 1 + end + + while !is_completed(sub) + yield() + end + else + MTL.waitUntilSignaledValue(sub.event, sub.seq) + end + end + + spins = 0 + while !diagnosed(sub) && spins < 1024 + ccall(:jl_cpu_pause, Cvoid, ()) + ccall(:jl_gc_safepoint, Cvoid, ()) + spins += 1 + end + return +end + +# Wait until the commit-feedback handler for `sub` has run. Metal delivers a command +# buffer's `MTLLogState` blocks on the same feedback path, so this is what flushes the +# output of a logging-enabled kernel. +function wait_diagnosed!(sub::Mtl4Submission) + wait_submission!(sub) + while !diagnosed(sub) + yield() + end + return +end + +function wait_submissions!(bq::BatchedCommandQueue) + isempty(bq.cleanups) || wait_submission!(last(bq.cleanups)) + return end function wait_oldest_cleanup!(bq::BatchedCommandQueue) isempty(bq.cleanups) && return - cmdbuf = first(bq.cleanups).cmdbuf - wait_cmdbuf!(cmdbuf) + wait_diagnosed!(first(bq.cleanups)) drain_cleanups!(bq) return end @@ -351,33 +831,103 @@ function limit_inflight!(bq::BatchedCommandQueue) return end +function take_errors!(bq::BatchedCommandQueue) + errors = bq.errors + bq.errors = nothing + return errors +end + + +## flushing + function reset_open_cmdbuf!(bq::BatchedCommandQueue, cmdbuf) bq.cmdbuf = nothing + bq.allocator = nothing + bq.resset = nothing + bq.scratch = ArgumentScratch[] bq.roots = Any[] bq.pending_ops = Any[] bq.nops = 0 bq.nbytes = 0 + bq.needs_barrier = false unregister_queue_if_idle!(bq) return end -function discard_open_cmdbuf!(bq::BatchedCommandQueue, cmdbuf) +# Abandon the open batch without committing it. The allocator was never submitted, so it +# can be reset and recycled right away. +function discard_open_cmdbuf!(bq::BatchedCommandQueue, cmdbuf::MTL4CommandBuffer) + allocator = bq.allocator + resset = bq.resset + scratch = bq.scratch reset_open_cmdbuf!(bq, cmdbuf) - release(cmdbuf) + + MTL.endCommandBuffer!(cmdbuf) + push!(bq.free_cmdbufs, cmdbuf) + if allocator !== nothing + MTL.reset!(allocator) + push!(bq.free_allocators, allocator) + end + if resset !== nothing + MTL.remove_all_allocations!(resset) + MTL.commit!(resset) + push!(bq.free_ressets, resset) + end + for s in scratch + push!(bq.free_scratch, reset!(s)) + end return end -function flush!(bq::BatchedCommandQueue) +# Commit the open batch, returning its `Mtl4Submission` (or `nothing` if there was none). +# The submission may already have been drained by `limit_inflight!` on return; waiting on +# it stays valid either way, since recycling does not touch its completion state. +function flush_batch!(bq::BatchedCommandQueue) cmdbuf = bq.cmdbuf - cmdbuf === nothing && return + cmdbuf === nothing && return nothing end_encoder!(bq) - register_operations!(bq, cmdbuf) + + allocator = bq.allocator::MTL4CommandAllocator + resset = bq.resset::MTLResidencySet + scratch = bq.scratch roots = bq.roots - MTL.commit_with_queue_key!(cmdbuf, pointer(bq.queue)) - defer_cleanup!(bq, cmdbuf, roots) + + MTL.commit!(resset) + MTL.use_residency_set!(cmdbuf, resset) + MTL.endCommandBuffer!(cmdbuf) + + value = bq.order + 1 + feedback = ccall(:jl_generating_output, Cint, ()) == 0 + sub = Mtl4Submission(BATCH_LABEL, cmdbuf, allocator, resset, scratch, bq.event, value, + feedback, roots, false, 0.0, 0.0, nothing) + + register_operations!(bq, sub) reset_open_cmdbuf!(bq, cmdbuf) + + if feedback + MTL.commit!(bq.queue4, cmdbuf, commit_options(sub, bq.queue_label)) + else + MTL.commit!(bq.queue4, cmdbuf) + end + + # signal completion, which both tracks this batch and lets Metal 3 work derived from + # this queue order itself after it + MTL.signal_event!(bq.queue4, bq.event, value) + bq.order = value + bq.pending4 = value + + defer_cleanup!(bq, sub) + + hook = MTL.profile_hook[] + hook === nothing || hook(sub) + limit_inflight!(bq) + return sub +end + +function flush!(bq::BatchedCommandQueue) + flush_batch!(bq) return end diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 4c20c046c..2551e57a5 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -105,15 +105,14 @@ end ## argument conversion struct Adaptor - # the current command encoder, if any. - cce::Union{Nothing,MTLComputeCommandEncoder} + # the queue the arguments are being encoded on, if any. + bq::Union{Nothing,BatchedCommandQueue} end # convert Metal buffers to their GPU address function Adapt.adapt_storage(to::Adaptor, buf::MTLBuffer) - if to.cce !== nothing - MTL.use!(to.cce, buf, MTL.ReadWriteUsage) - end + # only the address is encoded, so the buffer has to be declared resident explicitly + to.bq === nothing || make_resident!(to.bq, buf) reinterpret(Core.LLVMPtr{Nothing,AS.Device}, buf.gpuAddress) end function Adapt.adapt_storage(to::Adaptor, ptr::MtlPtr{T}) where {T} @@ -148,7 +147,7 @@ Adapt.adapt_structure(to::Adaptor, Broadcast.Broadcasted{Style}((x...) -> T(x...), adapt(to, bc.args), bc.axes) """ - mtlconvert(x, [cce]) + mtlconvert(x, [bq]) This function is called for every argument to be passed to a kernel, allowing it to be converted to a GPU-friendly format. By default, the function does nothing and returns the @@ -157,7 +156,7 @@ input object `x` as-is. Do not add methods to this function, but instead extend the underlying Adapt.jl package and register methods for the the `Metal.Adaptor` type. """ -mtlconvert(arg, cce=nothing) = adapt(Adaptor(cce), arg) +mtlconvert(arg, bq=nothing) = adapt(Adaptor(bq), arg) ## host-side kernel API @@ -170,7 +169,6 @@ struct HostKernel{F,TT} maxthreads::Int tgmem::Int exec_width::Int - use_residency_sets::Bool # this session's relocation words, or `nothing` for a relocation-free kernel. The buffer # keeps the storage alive; `launch` passes its address in the `KernelState` and declares # it resident, since only the address (not the buffer) is encoded. @@ -251,7 +249,6 @@ function mtlfunction(f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT} Int(pipeline[].maxTotalThreadsPerThreadgroup), Int(pipeline[].staticThreadgroupMemoryLength), Int(pipeline[].threadExecutionWidth), - can_use_residency_sets(dev), reloc_table) end::HostKernel{F,tt} end @@ -288,7 +285,11 @@ const kernel_instances = Dict{UInt, Any}() ## kernel launching and argument encoding -@inline @generated function encode_arguments!(cce, kernel, args::Vararg{Any,N}) where {N} +# Bind a kernel's arguments into the queue's argument table. Every binding is a GPU +# address in Metal 4, so buffers are declared resident and by-value arguments are staged +# through the batch's argument scratch buffer. +@inline @generated function encode_arguments!(bq::BatchedCommandQueue, kernel, + args::Vararg{Any,N}) where {N} ex = quote end # the arguments passed into this function have not been `mtlconvert`ed, because we need @@ -301,16 +302,16 @@ const kernel_instances = Dict{UInt, Any}() argex = :(args[$argidx]) if argtyp <: MTLBuffer # top-level buffers are passed as a pointer-valued argument - push!(ex.args, :(set_buffer!(cce, $argex, 0, $idx))) + push!(ex.args, :(bind_buffer!(bq, $argex, 0, $idx))) elseif argtyp <: MtlPtr # the same as a buffer, but with an offset - push!(ex.args, :(set_buffer!(cce, $argex.buffer, $argex.offset, $idx))) + push!(ex.args, :(bind_buffer!(bq, $argex.buffer, $argex.offset, $idx))) elseif isghosttype(argtyp) || Core.Compiler.isconstType(argtyp) continue else - # everything else is passed by reference, copied into Metal's transient buffer + # everything else is passed by reference, copied into the argument scratch append!(ex.args, (quote - set_argument!(cce, mtlconvert($(argex), cce), $idx) + set_argument!(bq, mtlconvert($(argex), bq), $idx) end).args) end idx += 1 @@ -321,7 +322,7 @@ const kernel_instances = Dict{UInt, Any}() ex end -@inline function set_argument!(cce::MTLComputeCommandEncoder, arg, idx::Integer) +@inline function set_argument!(bq::BatchedCommandQueue, arg, idx::Integer) argtyp = typeof(arg) # A non-isbits argument has no fields the kernel could read — compilation would have @@ -336,7 +337,7 @@ end ref = Base.RefValue(arg) GC.@preserve ref begin ptr = Base.unsafe_convert(Ptr{argtyp}, ref) - set_bytes!(cce, reinterpret(Ptr{Cvoid}, ptr), sizeof(argtyp), idx) + bind_bytes!(bq, reinterpret(Ptr{Cvoid}, ptr), sizeof(argtyp), idx) end return end @@ -366,13 +367,14 @@ function kernel_operation(@nospecialize(kernel::HostKernel), gs::MTLSize, ts::MT tgmem = kernel.tgmem, maxthreads = kernel.maxthreads) end +# Logging-enabled kernels get a command buffer of their own, so that the `MTLLogState` +# attached to it can be drained by waiting on that specific submission. function launch_logging!(@nospecialize(kernel::HostKernel), gs::MTLSize, ts::MTLSize, bq::BatchedCommandQueue, @nospecialize(args::Tuple), - kernel_state, buf, exc) + kernel_state) flush!(bq) - queue = bq.queue - if is_virtual(queue.device) + if is_virtual(bq.device) # `MTLLogState` needs a residency set, which the paravirtualized GPU driver # cannot create (failing with `MTLLogStateErrorDomain` code 2). Bail out here # with a clear host error instead of surfacing that opaque `NSError`. @@ -384,7 +386,7 @@ function launch_logging!(@nospecialize(kernel::HostKernel), gs::MTLSize, ts::MTL log_state_descriptor = MTLLogStateDescriptor() log_state_descriptor.level = MTL.MTLLogLevelDebug - log_state = MTLLogState(queue.device, log_state_descriptor) + log_state = MTLLogState(bq.device, log_state_descriptor) function log_handler(subSystem, category, logLevel, message) Core.print(String(NSString(message))) @@ -394,33 +396,26 @@ function launch_logging!(@nospecialize(kernel::HostKernel), gs::MTLSize, ts::MTL block = @objcblock(log_handler, Nothing, (id{NSString}, id{NSString}, NSInteger, id{NSString})) @objc [log_state::id{MTLLogState} addLogHandler:block::id{NSBlock}]::Nothing - cmdbuf_descriptor = MTLCommandBufferDescriptor() - cmdbuf_descriptor.logState = log_state - cmdbuf = MTLCommandBuffer(queue, cmdbuf_descriptor) - @label! cmdbuf "MTLCommandBuffer($(nameof(kernel.f)))" - let md = MTL.profile_metadata[] - md === nothing || MTL.note_operation!(md, cmdbuf, kernel_operation(kernel, gs, ts)) - end + options = MTL.MTL4CommandBufferOptions() + options.logState = log_state - cce = MTLComputeCommandEncoder(cmdbuf) + cce = begin_batch!(bq; options) try - MTL.set_function!(cce, kernel.pipeline) - if !kernel.use_residency_sets - MTL.use!(cce, buf, MTL.ReadWriteUsage) - MTL.use!(cce, exc, MTL.ReadWriteUsage) - end + set_pipeline!(bq, cce, kernel.pipeline) let reloc = kernel.reloc_table - reloc === nothing || MTL.use!(cce, reloc, MTL.ReadUsage) + reloc === nothing || make_resident!(bq, reloc) end - encode_arguments_nospec!(cce, kernel, kernel_state, kernel.f, args) + encode_arguments_nospec!(bq, kernel, kernel_state, kernel.f, args) MTL.append_current_function!(cce, gs, ts) - finally - close(cce) + catch + abort_batch!(bq) + rethrow() end - commit!(cmdbuf, queue) - defer_cleanup!(bq, cmdbuf, Any[kernel.f, args]) - track_logging_cmdbuf!(queue, cmdbuf) + op = MTL.profile_metadata[] === nothing ? nothing : kernel_operation(kernel, gs, ts) + record_operation!(bq, kernel.f, args; op=op) + sub = flush_batch!(bq) + sub === nothing || track_logging_submission!(bq, sub) return end @@ -453,9 +448,9 @@ function launch(@nospecialize(kernel::HostKernel), gs::MTLSize, ts::MTLSize, tgmem > 32768 && throw(ArgumentError("Total used threadgroupMemoryLength($tgmem) must be <= 32768 bytes.")) - buf, buf_addr = malloc_buffer_and_gpu_address(dev) + _, buf_addr = malloc_buffer_and_gpu_address(dev) buf_ptr = reinterpret(Core.LLVMPtr{UInt8, AS.Device}, buf_addr) - exc, exc_addr = exception_info_buffer_and_gpu_address(dev) + _, exc_addr = exception_info_buffer_and_gpu_address(dev) exc_ptr = reinterpret(Core.LLVMPtr{UInt8, AS.Device}, exc_addr) reloc = kernel.reloc_table reloc_ptr = reinterpret(Core.LLVMPtr{UInt64, AS.Device}, @@ -464,7 +459,7 @@ function launch(@nospecialize(kernel::HostKernel), gs::MTLSize, ts::MTLSize, if kernel.loggingEnabled precompiling && return - launch_logging!(kernel, gs, ts, bq, args, kernel_state, buf, exc) + launch_logging!(kernel, gs, ts, bq, args, kernel_state) return end @@ -473,17 +468,12 @@ function launch(@nospecialize(kernel::HostKernel), gs::MTLSize, ts::MTLSize, set_pipeline!(bq, cce, pipeline) # The kernel state holds GPU addresses to per-device scratch buffers (malloc bump - # allocator, exception mailbox) that aren't otherwise bound to the encoder. Declare - # them so Metal Shader Validation tracks the accesses instead of dropping them. - if !kernel.use_residency_sets - MTL.use!(cce, buf, MTL.ReadWriteUsage) - MTL.use!(cce, exc, MTL.ReadWriteUsage) - end - # The relocation table is per-kernel, so it cannot join the queue's residency set - # (which only holds the per-device scratch buffers): declare it every launch. - reloc === nothing || MTL.use!(cce, reloc, MTL.ReadUsage) + # allocator, exception mailbox) that aren't bound through the argument table. Those + # live in the queue's own residency set (see `install_queue_residency!`). The + # relocation table is per-kernel, so it joins the batch's set instead. + reloc === nothing || make_resident!(bq, reloc) - encode_arguments_nospec!(cce, kernel, kernel_state, f, args) + encode_arguments_nospec!(bq, kernel, kernel_state, f, args) MTL.append_current_function!(cce, gs, ts) catch # The failing launch has not been recorded yet. Keep any earlier @@ -514,8 +504,8 @@ function launch(@nospecialize(kernel::HostKernel), gs::MTLSize, ts::MTLSize, end # force specialization on f and args, but not on the kernel -@inline encode_arguments_nospec!(cce, @nospecialize(kernel), kernel_state, f, args::Tuple) = - encode_arguments!(cce, kernel, kernel_state, f, args...) +@inline encode_arguments_nospec!(bq, @nospecialize(kernel), kernel_state, f, args::Tuple) = + encode_arguments!(bq, kernel, kernel_state, f, args...) ## Intra-warp Helpers diff --git a/src/initialization.jl b/src/initialization.jl index 41759fc61..69e6d3019 100644 --- a/src/initialization.jl +++ b/src/initialization.jl @@ -52,6 +52,7 @@ function __init__() _shader_validation_enabled[] = get(ENV, "MTL_SHADER_VALIDATION", "0") != "0" MTL.submit_hook[] = flush_open_batch + MTL.commit_hook[] = order_metal4_after if !Sys.isapple() || Sys.ARCH != :aarch64 @error "Metal.jl is only supported on Apple Silicon" @@ -90,9 +91,7 @@ function __init__() has_active_batched_queues() || return try @autoreleasepool for bq in active_batched_queues() - cmdbuf = bq.cmdbuf - end_encoder!(bq) - cmdbuf === nothing || reset_open_cmdbuf!(bq, cmdbuf) + abort_batch!(bq) end catch err @error "Failed to close open batched command queues at exit" exception=(err, catch_backtrace()) diff --git a/src/memory.jl b/src/memory.jl index b1b8981b7..b924f2350 100644 --- a/src/memory.jl +++ b/src/memory.jl @@ -99,6 +99,8 @@ end chunk_size = 2^31 bq = batched_queue(queue) enc = blit_encoder(bq) + make_resident!(bq, dst.buffer) + make_resident!(bq, src.buffer) offset = 0 while nbytes > 0 @@ -132,6 +134,7 @@ end nbytes = N * sizeof(T) bq = batched_queue(queue) enc = blit_encoder(bq) + make_resident!(bq, dst.buffer) append_fillbuffer!(enc, dst.buffer, value, nbytes, dst.offset) op = MTL.profile_metadata[] === nothing ? nothing : @@ -147,4 +150,4 @@ end return dst end -# TODO: Implement generic fill since mtBlitCommandEncoderFillBuffer is limiting +# TODO: Implement generic fill since Metal's buffer fill is limited to byte patterns diff --git a/src/precompile.jl b/src/precompile.jl index 82d1d2308..e3b4ae422 100644 --- a/src/precompile.jl +++ b/src/precompile.jl @@ -67,6 +67,8 @@ using PrecompileTools: @setup_workload, @compile_workload empty!(queue_residency_sets) Base.@lock memory_pressure_stats_lock empty!(_memory_pressure_stats) empty!(device_malloc_bufs) + empty!(logging_submissions) + empty!(pending_commands) empty!(MTL.last_committed_per_queue) empty!(MTL.submission_state_per_queue) empty!(device_exception_info) diff --git a/src/profiling.jl b/src/profiling.jl index 1e2672df2..619bd0ed7 100644 --- a/src/profiling.jl +++ b/src/profiling.jl @@ -2,7 +2,7 @@ module Profiling -import ..Metal: MTL, synchronize, device_synchronize +import ..Metal: MTL, synchronize, device_synchronize, gpu_time_range import ObjectiveC @@ -21,7 +21,7 @@ max_thread_id() = isdefined(Threads, :maxthreadid) ? Threads.maxthreadid() : Thr # function clean_label(name::String) - m = match(r"^MTLCommandBuffer\((.*)\)$", name) + m = match(r"^MTL4?CommandBuffer\((.*)\)$", name) m === nothing ? name : String(m.captures[1]) end @@ -178,15 +178,14 @@ function profile_internally(@nospecialize(f); trace::Bool=false, raw::Bool=false records end for (opname, cmdbuf) in records - if cmdbuf.status == MTL.MTLCommandBufferStatusCompleted - t0 = cmdbuf.GPUStartTime - t1 = cmdbuf.GPUEndTime - if t1 > t0 - push!(name, opname) - push!(start, t0) - push!(stop, t1) - push!(ops, get(collector.metadata, cmdbuf, Any[])) - end + times = gpu_time_range(cmdbuf) + times === nothing && continue + t0, t1 = times + if t1 > t0 + push!(name, opname) + push!(start, t0) + push!(stop, t1) + push!(ops, get(collector.metadata, cmdbuf, Any[])) end end diff --git a/src/state.jl b/src/state.jl index 4f9d44279..34a43f0ec 100644 --- a/src/state.jl +++ b/src/state.jl @@ -59,8 +59,8 @@ end Return the [`BatchedCommandQueue`](@ref) associated with the current Julia task. -This is a *batched* queue: kernel launches and blit operations accumulate into a -single command buffer and are submitted lazily, rather than one command buffer per +This is a *batched* Metal 4 queue: kernel launches, copies and fills accumulate into a +single `MTL4CommandBuffer` and are submitted lazily, rather than one command buffer per operation. It is a drop-in for a raw `MTLCommandQueue` — using it as one (e.g. to derive a command buffer, or for MPS) preserves program order by draining pending batches when command buffers are enqueued or committed. Call [`synchronize`](@ref) @@ -82,48 +82,46 @@ function global_queue(dev::MTLDevice) end::BatchedCommandQueue end -# tracks the most recently launched logging-enabled cmdbuf per queue, so that +# tracks the most recently committed logging-enabled batch per queue, so that # `synchronize` can wait on it and thereby drain its `addLogHandler:` blocks # (Metal dispatches log delivery asynchronously and offers no flush primitive; -# `waitUntilCompleted` on the specific cmdbuf is what processes its pending blocks). -const logging_cmdbufs = IdDict{MTLCommandQueue,MTLCommandBuffer}() -const logging_cmdbufs_lock = ReentrantLock() +# waiting for the specific submission is what processes its pending blocks). +const logging_submissions = IdDict{Any,Any}() +const logging_submissions_lock = ReentrantLock() -function track_logging_cmdbuf!(queue::MTLCommandQueue, cmdbuf::MTLCommandBuffer) - Base.@lock logging_cmdbufs_lock begin - logging_cmdbufs[queue] = cmdbuf +function track_logging_submission!(bq, sub) + Base.@lock logging_submissions_lock begin + logging_submissions[bq] = sub end return end -function drain_logging_cmdbufs!(queue::MTLCommandQueue) - cmdbuf = Base.@lock logging_cmdbufs_lock begin - prev = get(logging_cmdbufs, queue, nothing) - delete!(logging_cmdbufs, queue) +function drain_logging_submissions!(bq) + sub = Base.@lock logging_submissions_lock begin + prev = get(logging_submissions, bq, nothing) + delete!(logging_submissions, bq) prev end - if cmdbuf !== nothing - MTL.wait_completed(cmdbuf) - end + sub === nothing || wait_diagnosed!(sub) return end ## scratch-buffer residency - -# Fast residency path; collapse this to `true` when virtual devices support residency sets -function can_use_residency_sets(dev::MTLDevice) - @memoize key=pointer(dev)::id{MTLDevice} begin - !is_virtual(dev) - end::Bool -end +# +# Metal 4 has no implicit residency: a buffer whose GPU address a kernel dereferences must +# be covered by a residency set. The per-device scratch buffers (the `malloc` bump +# allocator and the exception mailbox) are referenced by every launch through the kernel +# state, so they get a residency set of their own that is attached to the queue once. const queue_residency_sets = Dict{UInt,MTLResidencySet}() const queue_residency_sets_lock = ReentrantLock() -command_queue_key(queue::MTLCommandQueue) = UInt(pointer(queue)) +command_queue_key(queue::MTL4CommandQueue) = UInt(pointer(queue)) -function install_queue_residency!(queue::MTLCommandQueue, dev::MTLDevice) +function install_queue_residency!(bq) + queue = bq.queue4 + dev = bq.device key = command_queue_key(queue) Base.@lock queue_residency_sets_lock begin cached_resset = get(queue_residency_sets, key, nothing) diff --git a/src/synchronization.jl b/src/synchronization.jl index 7c307cf88..2886c4c45 100644 --- a/src/synchronization.jl +++ b/src/synchronization.jl @@ -80,6 +80,8 @@ function command_buffer_errors(state::Union{Nothing,MTL.QueueSubmissionState}) return MTL.finish_submissions!(state) end +command_buffer_errors(errors::Vector{MTL.CommandBufferErrorInfo}) = errors + function command_buffer_errors(states::AbstractVector{MTL.QueueSubmissionState}) errors = nothing for state in states @@ -94,8 +96,17 @@ function command_buffer_errors(states::AbstractVector{MTL.QueueSubmissionState}) return errors end -function check_synchronization_errors(states) - errors = command_buffer_errors(states) +function check_synchronization_errors(states...) + errors = nothing + for state in states + state_errors = command_buffer_errors(state) + state_errors === nothing && continue + if errors === nothing + errors = state_errors + else + append!(errors, state_errors) + end + end kernel_error = try check_exceptions() @@ -129,23 +140,26 @@ Wait for currently committed GPU work on `queue` to finish. bq = batched_queue(queue) flush!(bq) queue = bq.queue - maybe_collect(queue.device; will_block=true) + maybe_collect(bq.device; will_block=true) # flush any pending log handlers from logging-enabled kernels on this queue # (Metal delivers logs asynchronously; `wait_completed` on the specific cmdbuf # is what processes its `addLogHandler:` blocks) - drain_logging_cmdbufs!(queue) + drain_logging_submissions!(bq) + # Metal 3 command buffers derived from this queue (MPS, user code) last, submissions = MTL.take_queue_submissions(queue) - # Handles the already-completed fast path internally. last === nothing || wait_cmdbuf!(last) + # batched Metal 4 work + wait_submissions!(bq) + drain_cleanups!(bq; force=true) # Surface Metal runtime failures and device-side Julia exceptions together, # after cleanup has released all Julia roots held by completed work. - check_synchronization_errors(submissions) + check_synchronization_errors(submissions, take_errors!(bq)) return end @@ -166,10 +180,8 @@ function device_synchronize() flush_batched_queues!() maybe_collect(device(); will_block=true) - queues = active_global_queues() - append!(queues, active_batched_queues()) - for queue in unique!(queues) - drain_logging_cmdbufs!(raw_queue(queue)) + for bq in active_batched_queues() + drain_logging_submissions!(bq) end cmdbufs, submissions = MTL.take_all_submissions() @@ -184,10 +196,15 @@ function device_synchronize() end end + errors = nothing for bq in active_batched_queues() + wait_submissions!(bq) drain_cleanups!(bq; force=true) + bq_errors = take_errors!(bq) + bq_errors === nothing && continue + errors = errors === nothing ? bq_errors : append!(errors, bq_errors) end - check_synchronization_errors(submissions) + check_synchronization_errors(submissions, errors) return end diff --git a/test/execution.jl b/test/execution.jl index 411ae917a..4ef8ac453 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -200,10 +200,21 @@ struct EncodeFailure end Adapt.adapt_structure(to::Metal.Adaptor, x::EncodeFailure) = - to.cce === nothing ? x : error("intentional encode failure") + to.bq === nothing ? x : error("intentional encode failure") failed_encode_kernel(x) = return +# burns a few milliseconds of GPU time before writing, so that ordering against it is +# observable rather than accidental +function slow_write_kernel(A, val) + acc = val + for i in 1:2_000_000 + acc = (acc * Int32(1103515245) + Int32(12345)) & Int32(0x7fffffff) + end + A[1] = acc == Int32(12345) ? Int32(0) : val + return +end + function failed_batch_increment_kernel(A) A[1] += Int32(1) return @@ -260,12 +271,11 @@ end D = MtlArray(UInt8[0]) @metal threads=1 queue=queue write_kernel(D, UInt8(1)) + # deriving a Metal 3 command buffer from the queue flushes the open Metal 4 batch, so + # that the new buffer can be ordered after it with a GPU-side wait cmdbuf = MTL.MTLCommandBuffer(queue) - if !Metal.command_batching() || Metal.profiling_command_buffers() - @test queue.cmdbuf === nothing - elseif Metal.command_batching_ops() > 1 - @test queue.nops == 1 - end + @test queue.cmdbuf === nothing + @test queue.nops == 0 @metal threads=1 queue=queue write_kernel(D, UInt8(2)) MTL.MTLBlitCommandEncoder(cmdbuf) do enc buf = Base.unsafe_convert(MTL.MTLBuffer, D) @@ -274,6 +284,32 @@ end MTL.commit!(cmdbuf) synchronize(queue) @test Array(D) == UInt8[3] + + # a batch opened *after* the command buffer was derived cannot be waited on from the + # GPU side; make it slow enough that a missing host-side wait would show + E = MtlArray(Int32[0]) + cmdbuf = MTL.MTLCommandBuffer(queue) + @metal threads=1 queue=queue slow_write_kernel(E, Int32(1)) + MTL.MTLBlitCommandEncoder(cmdbuf) do enc + buf = Base.unsafe_convert(MTL.MTLBuffer, E) + MTL.append_fillbuffer!(enc, buf, UInt8(2), sizeof(E), E.offset) + end + MTL.commit!(cmdbuf) + synchronize(queue) + @test Array(E) == Int32[0x02020202] + + # the same through an MPS command buffer, which has its own constructor (the blit + # encoder wrapper only takes the underlying `MTLCommandBuffer`) + F = MtlArray(Int32[0]) + cmdbuf = MPS.MPSCommandBuffer(queue) + @metal threads=1 queue=queue slow_write_kernel(F, Int32(1)) + MTL.MTLBlitCommandEncoder(cmdbuf.commandBuffer) do enc + buf = Base.unsafe_convert(MTL.MTLBuffer, F) + MTL.append_fillbuffer!(enc, buf, UInt8(2), sizeof(F), F.offset) + end + MTL.commit!(cmdbuf) + synchronize(queue) + @test Array(F) == Int32[0x02020202] end function inject_command_buffer_error!(queue, info) diff --git a/test/mtl/metal4.jl b/test/mtl/metal4.jl new file mode 100644 index 000000000..0ed8eaa9e --- /dev/null +++ b/test/mtl/metal4.jl @@ -0,0 +1,246 @@ +using .MTL + +# a trivial kernel to exercise the Metal 4 dispatch path without going through the +# Julia compiler +const add1_source = """ +#include +using namespace metal; +kernel void add1(device uint* out [[buffer(0)]], + device const uint* inp [[buffer(1)]], + uint gid [[thread_position_in_grid]]) { + out[gid] = inp[gid] + 1; +} +""" + +@autoreleasepool begin + +dev = MTLDevice(1) + +@testset "command queue" begin + +queue = MTL4CommandQueue(dev) +@test queue isa MTL4CommandQueue +@test queue.device == dev + +desc = MTL4CommandQueueDescriptor() +desc.label = "test queue" +@test String(desc.label) == "test queue" +queue2 = MTL4CommandQueue(dev, desc) +@test String(queue2.label) == "test queue" + +@test MTL4CommandQueueDescriptor("labelled").label isa NSString + +ev = MTLSharedEvent(dev) +MTL.signal_event!(queue, ev, 1) +@test MTL.waitUntilSignaledValue(ev, 1, 5000) +@test ev.signaledValue >= 1 +# a queue-level wait on an already-signalled value must not block later work +MTL.wait_event!(queue, ev, 1) + +end + +@testset "command allocator" begin + +alloc = MTL4CommandAllocator(dev) +@test alloc isa MTL4CommandAllocator +@test alloc.device == dev +@test MTL.allocatedSize(alloc) isa Integer +MTL.reset!(alloc) + +desc = MTL4CommandAllocatorDescriptor("test allocator") +@test String(desc.label) == "test allocator" +alloc2 = MTL4CommandAllocator(dev, desc) +@test String(alloc2.label) == "test allocator" +@test String(MTL4CommandAllocator(dev, "inline label").label) == "inline label" + +end + +@testset "command buffer" begin + +alloc = MTL4CommandAllocator(dev) +cmdbuf = MTL4CommandBuffer(dev) +@test cmdbuf isa MTL4CommandBuffer +@test cmdbuf.device == dev + +cmdbuf.label = "test buffer" +@test String(cmdbuf.label) == "test buffer" + +# a Metal 4 command buffer is a reusable encoding cursor: it can be begun, ended and +# begun again without being recreated +MTL.beginCommandBuffer!(cmdbuf, alloc) +MTL.endCommandBuffer!(cmdbuf) +MTL.beginCommandBuffer!(cmdbuf, alloc) +MTL.endCommandBuffer!(cmdbuf) + +@test String(MTL4CommandBuffer(dev, "labelled").label) == "labelled" + +end + +@testset "argument table" begin + +desc = MTL4ArgumentTableDescriptor() +desc.maxBufferBindCount = 4 +desc.label = "test table" +@test desc.maxBufferBindCount == 4 + +argtab = MTL4ArgumentTable(dev, desc) +@test argtab isa MTL4ArgumentTable +@test argtab.device == dev + +argtab2 = MTL4ArgumentTable(dev; buffers=8, label="kwargs table") +@test String(argtab2.label) == "kwargs table" + +buf = MTLBuffer(dev, 64; storage=Metal.SharedStorage) +MTL.set_buffer!(argtab2, buf, 0, 1) +MTL.set_buffer!(argtab2, buf, 16, 2) +MTL.set_address!(argtab2, UInt64(buf.gpuAddress), 3) + +end + +@testset "compute encoder" begin + +pipeline = MTLComputePipelineState(dev, MTLFunction(MTLLibrary(dev, add1_source), "add1")) + +N = 16 +bufs = [MTLBuffer(dev, N*sizeof(UInt32); storage=Metal.SharedStorage) for _ in 1:3] +A, B, C = bufs +pA = convert(Ptr{UInt32}, MTL.contents(A)) +pC = convert(Ptr{UInt32}, MTL.contents(C)) +for i in 1:N + unsafe_store!(pA, UInt32(i), i) + unsafe_store!(pC, UInt32(0), i) +end + +rdesc = MTLResidencySetDescriptor() +rdesc.initialCapacity = 3 +resset = MTLResidencySet(dev, rdesc) +foreach(b -> MTL.add_allocation!(resset, b), bufs) +MTL.commit!(resset) + +queue = MTL4CommandQueue(dev) +alloc = MTL4CommandAllocator(dev) +cmdbuf = MTL4CommandBuffer(dev) +argtab = MTL4ArgumentTable(dev; buffers=4) + +MTL.beginCommandBuffer!(cmdbuf, alloc) +MTL.use_residency_set!(cmdbuf, resset) +enc = MTL4ComputeCommandEncoder(cmdbuf) +@test enc isa MTL4ComputeCommandEncoder +@test MTL.stages(enc) isa MTL.MTLStages + +MTL.set_function!(enc, pipeline) +MTL.set_argument_table!(enc, argtab) + +# B = A + 1 +MTL.set_buffer!(argtab, B, 0, 1) +MTL.set_buffer!(argtab, A, 0, 2) +MTL.dispatchThreads!(enc, MTLSize(N,1,1), MTLSize(N,1,1)) + +# commands in a Metal 4 encoder run concurrently unless a barrier separates them +MTL.barrierAfterEncoderStages!(enc) + +# C = B + 1, rebinding the *same* argument table: bindings are captured per dispatch +MTL.set_buffer!(argtab, C, 0, 1) +MTL.set_buffer!(argtab, B, 0, 2) +MTL.dispatchThreads!(enc, MTLSize(N,1,1), MTLSize(N,1,1)) + +close(enc) +MTL.endCommandBuffer!(cmdbuf) + +feedback = Ref{Any}(nothing) +done = Threads.Atomic{Bool}(false) +options = MTL4CommitOptions() do fb + feedback[] = (fb.GPUStartTime, fb.GPUEndTime, fb.error) + done[] = true + return +end +MTL.commit!(queue, cmdbuf, options) + +ev = MTLSharedEvent(dev) +MTL.signal_event!(queue, ev, 1) +@test MTL.waitUntilSignaledValue(ev, 1, 10_000) + +@test [unsafe_load(pC, i) for i in 1:N] == UInt32[i+2 for i in 1:N] + +t0 = time() +while !done[] && time() - t0 < 5 + sleep(0.01) +end +@test done[] +start, stop, err = feedback[] +@test err === nothing +@test stop >= start + +# the allocator can only be reset once the GPU is done with the commands it holds +MTL.reset!(alloc) + +end + +@testset "copy and fill" begin + +N = 64 +src = MTLBuffer(dev, N; storage=Metal.SharedStorage) +dst = MTLBuffer(dev, N; storage=Metal.SharedStorage) +psrc = convert(Ptr{UInt8}, MTL.contents(src)) +pdst = convert(Ptr{UInt8}, MTL.contents(dst)) +for i in 1:N + unsafe_store!(psrc, UInt8(i), i) + unsafe_store!(pdst, UInt8(0), i) +end + +rdesc = MTLResidencySetDescriptor() +resset = MTLResidencySet(dev, rdesc) +MTL.add_allocation!(resset, src) +MTL.add_allocation!(resset, dst) +MTL.commit!(resset) + +queue = MTL4CommandQueue(dev) +alloc = MTL4CommandAllocator(dev) +cmdbuf = MTL4CommandBuffer(dev) + +MTL.beginCommandBuffer!(cmdbuf, alloc) +MTL.use_residency_set!(cmdbuf, resset) +MTL4ComputeCommandEncoder(cmdbuf) do enc + MTL.append_copy!(enc, dst, 0, src, 0, 32) + MTL.barrierAfterEncoderStages!(enc) + MTL.append_fillbuffer!(enc, dst, UInt8(0xff), 32, 32) +end +MTL.endCommandBuffer!(cmdbuf) +MTL.commit!(queue, cmdbuf) + +ev = MTLSharedEvent(dev) +MTL.signal_event!(queue, ev, 1) +@test MTL.waitUntilSignaledValue(ev, 1, 10_000) + +@test [unsafe_load(pdst, i) for i in 1:32] == UInt8[i for i in 1:32] +@test all([unsafe_load(pdst, i) for i in 33:64] .== 0xff) + +end + +@testset "capture scope" begin + +queue = MTL4CommandQueue(dev) +scope = MTLCaptureScope(queue) +@test scope isa MTLCaptureScope +@test scope.mtl4CommandQueue == queue + +end + +@testset "batched queue" begin + +bq = Metal.global_queue(device()) +@test bq.queue4 isa MTL4CommandQueue +@test bq.argtable isa MTL4ArgumentTable +@test bq.event isa MTLSharedEvent +@test bq.device == device() + +# a batched queue is still a drop-in for the Metal 3 queue it wraps +@test bq == bq.queue +cmdbuf = MTLCommandBuffer(bq) +@test cmdbuf isa MTLCommandBuffer +MTL.commit!(cmdbuf) +Metal.synchronize() + +end + +end