Progress? towards synchronization

Author: christiangnrd

lib/mtl/command_alloc4.jl

@@ -1,3 +1,20 @@
+
+export MTL4CommandAllocatorDescriptor
+
+function MTL4CommandAllocatorDescriptor()
+    handle = @objc [MTL4CommandAllocatorDescriptor new]::id{MTL4CommandAllocatorDescriptor}
+    obj = MTL4CommandAllocatorDescriptor(handle)
+    finalizer(release, obj)
+    return obj
+end
+function MTL4CommandAllocatorDescriptor(label)
+    desc = MTL4CommandAllocatorDescriptor()
+    desc.label = label
+    return desc
+end
+
+
+
 #
 # command allocator
 #
@@ -13,6 +30,20 @@ function MTL4CommandAllocator(device::MTLDevice)
     return obj
 end
 
+function MTL4CommandAllocator(dev::MTLDevice, descriptor::MTL4CommandAllocatorDescriptor)
+    err = Ref{id{NSError}}(nil)
+    handle = @objc [dev::id{MTLDevice} newCommandAllocatorWithDescriptor:descriptor::id{MTL4CommandAllocatorDescriptor}
+                                        error:err::Ptr{id{NSError}}]::id{MTL4CommandAllocator}
+    obj = MTL4CommandAllocator(handle)
+    finalizer(release, obj)
+    return obj
+end
+
+function MTL4CommandAllocator(dev::MTLDevice, label)
+    desc = MTL4CommandAllocatorDescriptor(label)
+    return MTL4CommandAllocator(dev, desc)
+end
+
 function allocatedSize(alloc::MTL4CommandAllocator)::UInt64
     @objc [alloc::id{MTL4CommandAllocator} allocatedSize]::UInt64
 end

lib/mtl/command_buf4.jl

@@ -6,23 +6,19 @@ export MTL4CommandBuffer, commit!, beginCommandBufferWithAllocator!, endCommandB
 
 # @objcwrapper immutable=false MTL4CommandBuffer <: NSObject
 
-function MTL4CommandBuffer(device::MTLDevice)
+function MTL4CommandBuffer(device::MTLDevice, label=nothing)
     handle = @objc [device::id{MTLDevice} newCommandBuffer]::id{MTL4CommandBuffer}
-    return MTL4CommandBuffer(handle)
+    buf = MTL4CommandBuffer(handle)
+    if !isnothing(label)
+        buf.label = label
+    end
+    return buf
 end
 
-function MTL4CommandBuffer(f::Base.Callable, device::MTLDevice; queue::MTL4CommandQueue=MTL4CommandQueue(device), allocator::MTL4CommandAllocator=MTL4CommandAllocator(device))
-    cmdbuf = MTL4CommandBuffer(device)
-
-    beginCommandBufferWithAllocator!(cmdbuf, allocator)
+function MTL4CommandBuffer(f::Base.Callable, device::MTLDevice, label=nothing; queue::MTL4CommandQueue=MTL4CommandQueue(device), allocator::MTL4CommandAllocator=MTL4CommandAllocator(device))
+    cmdbuf = MTL4CommandBuffer(device, label)
 
-    try
-        ret = f(cmdbuf)
-        return ret
-    finally
-        endCommandBuffer!(cmdbuf)
-        commit!(queue, cmdbuf)
-    end
+    commit!(f, cmdbuf, queue, allocator)
 end
 
 function beginCommandBufferWithAllocator!(cmdbuf::MTL4CommandBuffer, allocator::MTL4CommandAllocator, options::Union{Nothing, MTL4CommandBufferOptions} = nothing)
@@ -49,3 +45,15 @@ function commit!(cmdqueue::MTL4CommandQueue, cmdbuf::MTL4CommandBuffer, options:
                                     count:1::NSUInteger
                                     options:options::id{MTL4CommitOptions}]::Nothing
 end
+
+function commit!(f::Base.Callable, cmdbuf::MTL4CommandBuffer, queue::MTL4CommandQueue, allocator::MTL4CommandAllocator)
+    beginCommandBufferWithAllocator!(cmdbuf, allocator)
+
+    try
+        ret = f(cmdbuf)
+        return ret
+    finally
+        endCommandBuffer!(cmdbuf)
+        commit!(queue, cmdbuf)
+    end
+end

lib/mtl/command_queue4.jl

@@ -36,6 +36,21 @@ end
 
 
 
+export MTL4CommandQueueDescriptor
+
+function MTL4CommandQueueDescriptor()
+    handle = @objc [MTL4CommandQueueDescriptor new]::id{MTL4CommandQueueDescriptor}
+    obj = MTL4CommandQueueDescriptor(handle)
+    finalizer(release, obj)
+    return obj
+end
+function MTL4CommandQueueDescriptor(label)
+    desc = MTL4CommandQueueDescriptor()
+    desc.label = label
+    return desc
+end
+
+
 export MTL4CommandQueue
 
 # @objcwrapper immutable=false MTL4CommandQueue <: NSObject
@@ -46,3 +61,12 @@ function MTL4CommandQueue(dev::MTLDevice)
     finalizer(release, obj)
     return obj
 end
+
+function MTL4CommandQueue(dev::MTLDevice, descriptor::MTL4CommandQueueDescriptor)
+    err = Ref{id{NSError}}(nil)
+    handle = @objc [dev::id{MTLDevice} newMTL4CommandQueueWithDescriptor:descriptor::id{MTL4CommandQueueDescriptor}
+                                        error:err::Ptr{id{NSError}}]::id{MTL4CommandQueue}
+    obj = MTL4CommandQueue(handle)
+    finalizer(release, obj)
+    return obj
+end

lib/mtl/libmtl.jl

@@ -2882,7 +2882,7 @@ end
 
 @objcwrapper immutable = false availability = macos(v"26.0.0") MTL4CommitOptions <: NSObject
 
-@objcwrapper immutable = true availability = macos(v"26.0.0") MTL4CommandQueueDescriptor <: NSObject
+@objcwrapper immutable = false availability = macos(v"26.0.0") MTL4CommandQueueDescriptor <: NSObject
 
 @objcproperties MTL4CommandQueueDescriptor begin
     @autoproperty label::id{NSString} setter = setLabel
@@ -3327,7 +3327,7 @@ end
     @autoproperty functionType::MTLFunctionType
 end
 
-@objcwrapper immutable = true availability = macos(v"26.0.0") MTL4CommandAllocatorDescriptor <: NSObject
+@objcwrapper immutable = false availability = macos(v"26.0.0") MTL4CommandAllocatorDescriptor <: NSObject
 
 @objcproperties MTL4CommandAllocatorDescriptor begin
     @autoproperty label::id{NSString} setter = setLabel

res/wrap/libmtl.toml

@@ -51,6 +51,9 @@ immutable=false
 [api.MTLCommandQueue]
 immutable=false
 
+[api.MTL4CommandQueueDescriptor]
+immutable=false
+
 [api.MTL4CommandQueue]
 immutable=false
 
@@ -68,6 +71,9 @@ immutable=false
 [api.MTL4ComputeCommandEncoder]
 immutable=false
 
+[api.MTL4CommandAllocatorDescriptor]
+immutable=false
+
 [api.MTL4CommandAllocator]
 immutable=false
 

src/compiler/execution.jl

@@ -281,9 +281,7 @@ end
     if use_mtl4
         allocator = MTL4CommandAllocator(device())
         cmdbuf = MTL4CommandBuffer(device())
-
-        # TODO: Initialize with descriptor to set label
-        # cmdbuf.label = "MTL4CommandBuffer($(nameof(kernel.f)))"
+        cmdbuf.label = "MTL4CommandBuffer($(nameof(kernel.f)))"
 
         beginCommandBufferWithAllocator!(cmdbuf, allocator)
         cce = MTL4ComputeCommandEncoder(cmdbuf)

src/state.jl

@@ -36,7 +36,7 @@ Sets the Metal GPU device associated with the current Julia task.
 device!(dev::MTLDevice) = task_local_storage(:MTLDevice, dev)
 
 const global_queues = WeakKeyDict{MTLCommandQueue,Nothing}()
-const global_queues4 = WeakKeyDict{MTL4CommandQueue,Nothing}()
+const global_queues4 = WeakKeyDict{MTL4CommandQueue,Tuple{MTL4CommandBuffer, MTL4CommandAllocator}}()
 
 """
     global_queue(dev::MTLDevice)::MTLCommandQueue
@@ -64,13 +64,14 @@ Return the Metal 4 command queue associated with the current Julia thread.
 function global_queue4(dev::MTLDevice)
     get!(task_local_storage(), (:MTL4CommandQueue, dev)) do
         @autoreleasepool begin
+            desc = MTL4CommandQueueDescriptor("global_queue4($(current_task()))")
+
             # NOTE: MTL4CommandQueue itself is manually reference-counted,
             #       the release pool is for resources used during its construction.
-            queue = MTL4CommandQueue(dev)
+            queue = MTL4CommandQueue(dev, desc)
+
 
-            # # TODO: Initialize with descriptor to set label
-            # queue.label = "global_queue4($(current_task()))"
-            global_queues4[queue] = nothing
+            global_queues4[queue] = (MTL4CommandBuffer(dev, "sync_buffer($(current_task()))"), MTL4CommandAllocator(dev, "sync_allocater($(current_task()))"))
             queue
         end
     end::MTL4CommandQueue
@@ -86,10 +87,36 @@ Create a new MTLCommandBuffer from the global command queue, commit it to the qu
 and simply wait for it to be completed. Since command buffers *should* execute in a
 First-In-First-Out manner, this synchronizes the GPU.
 """
-@autoreleasepool function synchronize(queue::MTLCommandQueue=global_queue(device()))
+@autoreleasepool function synchronize(queue::MTLCommandQueue)
     cmdbuf = MTLCommandBuffer(queue)
     commit!(cmdbuf)
     wait_completed(cmdbuf)
+    return
+end
+@autoreleasepool function synchronize(queue::MTL4CommandQueue)
+    cmdbuf, allocator = get(global_queues4, queue) do
+        @info "hello"
+        dev = queue.device
+        MTL4CommandBuffer(dev), MTL4CommandAllocator(dev)
+    end
+
+    cmdbuf = commit!(cmdbuf, queue, allocator) do cmdbuf
+        MTL4ComputeCommandEncoder(identity, cmdbuf, #=sync=#true)
+    end
+    return
+end
+function synchronize()
+    dev = device()
+    tlskeys = keys(task_local_storage())
+    # hasmtl3key = (:MTLCommandQueue, dev) in tlskeys
+    # hasmtl4key = use_metal4() && (:MTL4CommandQueue, dev) in tlskeys
+    if (:MTLCommandQueue, dev) in tlskeys
+        synchronize(global_queue(dev))
+    end
+    if use_metal4() && (:MTL4CommandQueue, dev) in tlskeys
+        synchronize(global_queue4(dev))
+    end
+    return
 end
 
 """
@@ -101,4 +128,7 @@ function device_synchronize()
     for queue in keys(global_queues)
         synchronize(queue)
     end
+    for queue in keys(global_queues4)
+        synchronize(queue)
+    end
 end