[Do not merge] Metal 4

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

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")

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)

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
 

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

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

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