Implement @mtlprintf using os_log

Author: tgymnich

lib/mtl/MTL.jl

@@ -34,6 +34,7 @@ include("events.jl")
 include("fences.jl")
 include("heap.jl")
 include("buffer.jl")
+include("log_state.jl")
 include("command_queue.jl")
 include("command_buf.jl")
 include("compute_pipeline.jl")

lib/mtl/command_buf.jl

@@ -28,6 +28,7 @@ export MTLCommandBufferDescriptor
 @objcproperties MTLCommandBufferDescriptor begin
     @autoproperty retainedReferences::Bool setter=setRetainedReferences
     @autoproperty errorOptions::MTLCommandBufferErrorOption setter=setErrorOptions
+    @autoproperty logState::id{MTLLogState} setter=setLogState
 end
 
 function MTLCommandBufferDescriptor()

lib/mtl/command_queue.jl

@@ -1,3 +1,21 @@
+export MTLCommandQueueDescriptor
+
+@objcwrapper immutable=false MTLCommandQueueDescriptor <: NSObject
+
+@objcproperties MTLCommandQueueDescriptor begin
+    @autoproperty maxCommandBufferCount::NSUInteger
+    @autoproperty logState::id{MTLLogState} setter=setLogState
+end
+
+function MTLCommandQueueDescriptor()
+    handle = @objc [MTLCommandQueueDescriptor alloc]::id{MTLCommandQueueDescriptor}
+    obj = MTLCommandQueueDescriptor(handle)
+    finalizer(release, obj)
+    @objc [obj::id{MTLCommandQueueDescriptor} init]::id{MTLCommandQueueDescriptor}
+    return obj
+end
+
+
 export MTLCommandQueue
 
 @objcwrapper immutable=false MTLCommandQueue <: NSObject
@@ -13,3 +31,10 @@ function MTLCommandQueue(dev::MTLDevice)
     finalizer(release, obj)
     return obj
 end
+
+function MTLCommandQueue(dev::MTLDevice, descriptor::MTLCommandQueueDescriptor)
+    handle = @objc [dev::id{MTLDevice} newCommandQueueWithDescriptor:descriptor::id{MTLCommandQueueDescriptor}]::id{MTLCommandQueue}
+    obj = MTLCommandQueue(handle)
+    finalizer(release, obj)
+    return obj
+end

lib/mtl/log_state.jl

@@ -0,0 +1,40 @@
+export MTLLogLevel
+
+@cenum MTLLogLevel::NSInteger begin
+    MTLLogLevelUndefined = 0
+    MTLLogLevelDebug = 1
+    MTLLogLevelInfo = 2
+    MTLLogLevelNotice = 3
+    MTLLogLevelError = 4
+    MTLLogLevelFault = 5
+end
+
+export MTLLogStateDescriptor
+
+@objcwrapper immutable=false MTLLogStateDescriptor <: NSObject
+
+@objcproperties MTLLogStateDescriptor begin
+    @autoproperty level::MTLLogLevel setter=setLevel
+    @autoproperty bufferSize::NSInteger setter=setBufferSize
+end
+
+function MTLLogStateDescriptor()
+    handle = @objc [MTLLogStateDescriptor alloc]::id{MTLLogStateDescriptor}
+    obj = MTLLogStateDescriptor(handle)
+    finalizer(release, obj)
+    @objc [obj::id{MTLLogStateDescriptor} init]::id{MTLLogStateDescriptor}
+    return obj
+end
+
+
+export MTLLogState
+
+@objcwrapper MTLLogState <: NSObject
+
+function MTLLogState(dev::MTLDevice, descriptor::MTLLogStateDescriptor)
+    err = Ref{id{NSError}}(nil)
+    handle = @objc [dev::id{MTLDevice} newLogStateWithDescriptor:descriptor::id{MTLLogStateDescriptor} 
+                                       error:err::Ptr{id{NSError}}]::id{MTLLogState}
+    err[] == nil || throw(NSError(err[]))
+    MTLLogState(handle)
+end

src/Metal.jl

@@ -36,6 +36,7 @@ include("device/intrinsics/memory.jl")
 include("device/intrinsics/simd.jl")
 include("device/intrinsics/version.jl")
 include("device/intrinsics/atomics.jl")
+include("device/intrinsics/output.jl")
 include("device/quirks.jl")
 
 # array essentials

src/compiler/compilation.jl

@@ -97,9 +97,9 @@ function compile(@nospecialize(job::CompilerJob))
 
     @signpost_interval log=log_compiler() "Generate LLVM IR" begin
         # TODO: on 1.9, this actually creates a context. cache those.
-        ir, entry = JuliaContext() do ctx
+        ir, entry, loggingEnabled = JuliaContext() do ctx
             mod, meta = GPUCompiler.compile(:llvm, job)
-            string(mod), LLVM.name(meta.entry)
+            string(mod), LLVM.name(meta.entry), haskey(functions(mod), "air.os_log")
         end
     end
 
@@ -160,7 +160,7 @@ function compile(@nospecialize(job::CompilerJob))
         end
     end
 
-    return (; image, entry)
+    return (; image, entry, loggingEnabled)
 end
 
 # link into an executable kernel
@@ -192,5 +192,5 @@ end
 
     # most of the time, we don't need the function object,
     # so don't keep it alive unconditionally in GPUCompiler's caches
-    pipeline_state, return_function ? fun : nothing
+    pipeline_state, return_function ? fun : nothing, compiled.loggingEnabled
 end

src/compiler/execution.jl

@@ -161,6 +161,7 @@ mtlconvert(arg, cce=nothing) = adapt(Adaptor(cce), arg)
 struct HostKernel{F,TT}
     f::F
     pipeline::MTLComputePipelineState
+    loggingEnabled::Bool
 end
 
 const mtlfunction_lock = ReentrantLock()
@@ -182,15 +183,15 @@ function mtlfunction(f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT}
         cache = compiler_cache(dev)
         source = methodinstance(F, tt)
         config = compiler_config(dev; name, kwargs...)::MetalCompilerConfig
-        pipeline, _ = GPUCompiler.cached_compilation(cache, source, config, compile, link)
+        pipeline, _, loggingEnabled = GPUCompiler.cached_compilation(cache, source, config, compile, link)
 
         # create a callable object that captures the function instance. we don't need to think
         # about world age here, as GPUCompiler already does and will return a different object
         h = hash(pipeline, hash(f, hash(tt)))
         kernel = get(_kernel_instances, h, nothing)
         if kernel === nothing
             # create the kernel state object
-            kernel = HostKernel{F,tt}(f, pipeline)
+            kernel = HostKernel{F,tt}(f, pipeline, loggingEnabled)
             _kernel_instances[h] = kernel
         end
         return kernel::HostKernel{F,tt}
@@ -271,7 +272,34 @@ end
     (threads.width * threads.height * threads.depth) > kernel.pipeline.maxTotalThreadsPerThreadgroup &&
         throw(ArgumentError("Number of threads in group ($(threads.width * threads.height * threads.depth)) should not exceed $(kernel.pipeline.maxTotalThreadsPerThreadgroup)"))
 
-    cmdbuf = MTLCommandBuffer(queue)
+    cmdbuf = if kernel.loggingEnabled
+        if macos_version() < v"15"
+            @error "Logging is only supported on macOS 15 or higher"
+        end
+
+        if MTLCaptureManager().isCapturing
+            @error "Logging is not supported while GPU frame capturing"
+        end
+
+        log_state_descriptor = MTLLogStateDescriptor()
+        log_state_descriptor.level = MTL.MTLLogLevelDebug
+        log_state = MTLLogState(device(), log_state_descriptor)
+
+        function log_handler(subSystem, category, logLevel, message)
+            print(String(NSString(message)))
+            return nothing
+        end
+
+        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
+        MTLCommandBuffer(queue, cmdbuf_descriptor)
+    else
+        MTLCommandBuffer(queue)
+    end
+
     cmdbuf.label = "MTLCommandBuffer($(nameof(kernel.f)))"
     cce = MTLComputeCommandEncoder(cmdbuf)
     argument_buffers = try