Add support for signposts (#300)

Author: maleadt

docs/src/profiling.md

@@ -46,6 +46,24 @@ julia> Metal.@profile @metal threads=length(c) vadd(a, b, c);
 [ Info: System trace saved to julia_3.trace; open the resulting trace in Instruments
 ```
 
+It is possible to augment the trace with additional information by using signposts: Similar
+to NVTX markers and ranges in CUDA.jl, signpost intervals and events can be used to add
+respectively time intervals and points of interest to the trace. This can be done by using
+the signpost functionality from ObjectiveC.jl:
+
+```julia
+using ObjectiveC, .OS
+
+@signpost_interval "My Interval" begin
+    # code to profile
+    @signpost_event "My Event"
+end
+```
+
+For more information, e.g. how to pass additional messages to the signposts, or how to
+use a custom logger, consult the ObjectiveC.jl documentation, or the docstrings of the
+`@signpost_interval` and `@signpost_event` macros.
+
 ## Frame capture
 
 For more details on specific operations, you can use Metal's frame capture feature to

src/Metal.jl

@@ -11,7 +11,7 @@ using Python_jll
 using ObjectFile
 using ExprTools: splitdef, combinedef
 using Artifacts
-using ObjectiveC, .CoreFoundation, .Foundation, .Dispatch
+using ObjectiveC, .CoreFoundation, .Foundation, .Dispatch, .OS
 
 if !isdefined(Base, :get_extension)
     using Requires: @require

src/compiler/compilation.jl

@@ -55,52 +55,59 @@ end
 
 # compile to executable machine code
 function compile(@nospecialize(job::CompilerJob))
-    # TODO: on 1.9, this actually creates a context. cache those.
-    ir, entry = JuliaContext() do ctx
-        mod, meta = GPUCompiler.compile(:llvm, job)
-        string(mod), LLVM.name(meta.entry)
-    end
-
-    # generate AIR
-    air = let
-        input = Pipe()
-        output = Pipe()
+    @signpost_event log=log_compiler() "Compile" "Job=$job"
 
-        cmd = `$(LLVMDowngrader_jll.llvm_as()) --bitcode-version=5.0 -o -`
-        proc = run(pipeline(cmd, stdout=output, stderr=stderr, stdin=input); wait=false)
-        close(output.in)
-
-        writer = @async begin
-            write(input, ir)
-            close(input)
+    @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
+            mod, meta = GPUCompiler.compile(:llvm, job)
+            string(mod), LLVM.name(meta.entry)
         end
-        reader = @async read(output)
-
-        wait(proc)
-        if !success(proc)
-            file = tempname(cleanup=false) * ".ll"
-            write(file, ir)
-            error("""Compilation to AIR failed; see above for details.
-                     If you think this is a bug, please file an issue and attach $(file)""")
+    end
+
+    @signpost_interval log=log_compiler() "Downgrade to AIR" begin
+        # generate AIR
+        air = let
+            input = Pipe()
+            output = Pipe()
+
+            cmd = `$(LLVMDowngrader_jll.llvm_as()) --bitcode-version=5.0 -o -`
+            proc = run(pipeline(cmd, stdout=output, stderr=stderr, stdin=input); wait=false)
+            close(output.in)
+
+            writer = @async begin
+                write(input, ir)
+                close(input)
+            end
+            reader = @async read(output)
+
+            wait(proc)
+            if !success(proc)
+                file = tempname(cleanup=false) * ".ll"
+                write(file, ir)
+                error("""Compilation to AIR failed; see above for details.
+                        If you think this is a bug, please file an issue and attach $(file)""")
+            end
+            fetch(reader)
         end
-        fetch(reader)
     end
 
-    # create a Metal library
-    image = try
-        metallib_fun = MetalLibFunction(; name=entry, air_module=air,
-                                          air_version=job.config.target.air,
-                                          metal_version=job.config.target.metal)
-        metallib = MetalLib(; functions = [metallib_fun])
-
-        image_stream = IOBuffer()
-        write(image_stream, metallib)
-        take!(image_stream)
-    catch err
-        file = tempname(cleanup=false) * ".air"
-        write(file, air)
-        error("""Compilation to Metal library failed; see below for details.
-                 If you think this is a bug, please file an issue and attach $(file)""")
+    @signpost_interval log=log_compiler() "Create Metal library" begin
+        image = try
+            metallib_fun = MetalLibFunction(; name=entry, air_module=air,
+                                            air_version=job.config.target.air,
+                                            metal_version=job.config.target.metal)
+            metallib = MetalLib(; functions = [metallib_fun])
+
+            image_stream = IOBuffer()
+            write(image_stream, metallib)
+            take!(image_stream)
+        catch err
+            file = tempname(cleanup=false) * ".air"
+            write(file, air)
+            error("""Compilation to Metal library failed; see below for details.
+                    If you think this is a bug, please file an issue and attach $(file)""")
+        end
     end
 
     return (; image, entry)
@@ -109,20 +116,24 @@ end
 # link into an executable kernel
 @autoreleasepool function link(@nospecialize(job::CompilerJob), compiled;
                                return_function=false)
-    dev = current_device()
-    lib = MTLLibraryFromData(dev, compiled.image)
-    fun = MTLFunction(lib, compiled.entry)
-    pipeline_state = try
-        MTLComputePipelineState(dev, fun)
-    catch err
-        isa(err, NSError) || rethrow()
-
-        # the back-end compiler likely failed
-        # XXX: check more accurately? the error domain doesn't help much here
-        file = tempname(cleanup=false) * ".metallib"
-        write(file, compiled.image)
-        error("""Compilation to native code failed; see below for details.
-                 If you think this is a bug, please file an issue and attach $(file)""")
+    @signpost_event log=log_compiler() "Link" "Job=$job"
+
+    @signpost_interval log=log_compiler() "Instantiate compute pipeline" begin
+        dev = current_device()
+        lib = MTLLibraryFromData(dev, compiled.image)
+        fun = MTLFunction(lib, compiled.entry)
+        pipeline_state = try
+            MTLComputePipelineState(dev, fun)
+        catch err
+            isa(err, NSError) || rethrow()
+
+            # the back-end compiler likely failed
+            # XXX: check more accurately? the error domain doesn't help much here
+            file = tempname(cleanup=false) * ".metallib"
+            write(file, compiled.image)
+            error("""Compilation to native code failed; see below for details.
+                    If you think this is a bug, please file an issue and attach $(file)""")
+        end
     end
 
     # most of the time, we don't need the function object,

src/pool.jl

@@ -51,15 +51,15 @@ The storage kwarg controls where the buffer is stored. Possible values are:
 Note that `Private` buffers can't be directly accessed from the CPU, therefore you cannot
 use this option if you pass a ptr to initialize the memory.
 """
-function alloc(dev::Union{MTLDevice,MTLHeap}, bytesize::Integer, args...;
-               storage, kwargs...)
+function alloc(dev::Union{MTLDevice,MTLHeap}, sz::Integer, args...; storage, kwargs...)
+    @signpost_event log=log_array() "Allocate" "Size=$(Base.format_bytes(sz))"
 
     time = Base.@elapsed begin
-        buf = @autoreleasepool MTLBuffer(dev, bytesize, args...; storage, kwargs...)
+        buf = @autoreleasepool MTLBuffer(dev, sz, args...; storage, kwargs...)
     end
 
     Base.@atomic alloc_stats.alloc_count + 1
-    Base.@atomic alloc_stats.alloc_bytes + bytesize
+    Base.@atomic alloc_stats.alloc_bytes + sz
     Base.@atomic alloc_stats.total_time + time
 
     return buf
@@ -73,6 +73,7 @@ This does not protect against double-freeing of the same buffer!
 """
 function free(buf::MTLBuffer)
     sz::Int = buf.length
+    @signpost_event log=log_array() "Free" "Size=$(Base.format_bytes(sz))"
 
     time = Base.@elapsed begin
         @autoreleasepool unsafe=true release(buf)

src/state.jl

@@ -1,5 +1,10 @@
 export current_device, device!, global_queue, synchronize, device_synchronize
 
+log_compiler()          = OSLog("org.juliagpu.metal", "Compiler")
+log_compiler(args...)   = log_compiler()(args...)
+log_array()             = OSLog("org.juliagpu.metal", "Array")
+log_array(args...)      = log_array()(args...)
+
 """
     current_device()::MTLDevice
 

src/utilities.jl

@@ -161,6 +161,7 @@ function profiled(f)
     notification_name = "julia.metal.profile"
     folder = profile_dir()
     instruments = [
+        # relevant instruments taken from `xcrun xctrace list instruments`
         "GPU",
 
         # CPU
@@ -169,6 +170,8 @@ function profiled(f)
         "Metal Application",
         "Metal GPU Counters",
         "Metal Resource Events",
+
+        "os_signpost",
     ]
     cmd = `xctrace record`
     for instrument in instruments