Implement @mtlprintf using os_log

tgymnich · tgymnich · commit 0f527468b295 · 2024-09-16T13:54:02.000+02:00
diff --git a/src/Metal.jl b/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
diff --git a/src/device/intrinsics/output.jl b/src/device/intrinsics/output.jl
@@ -0,0 +1,262 @@
+export @mtlprintf
+
+@generated function promote_c_argument(arg)
+    # > When a function with a variable-length argument list is called, the variable
+    # > arguments are passed using C's old ``default argument promotions.'' These say that
+    # > types char and short int are automatically promoted to int, and type float is
+    # > automatically promoted to double. Therefore, varargs functions will never receive
+    # > arguments of type char, short int, or float.
+
+    if arg == Cchar || arg == Cshort
+        return :(Cint(arg))
+    elseif arg == Cfloat
+        return :(Cdouble(arg))
+    else
+        return :(arg)
+    end
+end
+
+"""
+    @mtlprintf("%Fmt", args...)
+
+Print a formatted string in device context on the host standard output.
+"""
+macro mtlprintf(fmt::String, args...)
+    fmt_val = Val(Symbol(fmt))
+
+    return :(_mtlprintf($fmt_val, $(map(arg -> :(promote_c_argument($arg)), esc.(args))...)))
+end
+
+@generated function _mtlprintf(::Val{fmt}, argspec...) where {fmt}
+    @dispose ctx=Context() begin
+        arg_exprs = [:( argspec[$i] ) for i in 1:length(argspec)]
+        arg_types = [argspec...]
+
+        T_void = LLVM.VoidType()
+        T_int32 = LLVM.Int32Type()
+        T_int64 = LLVM.Int64Type()
+        T_pint8 = LLVM.PointerType(LLVM.Int8Type())
+        T_pint8a2 = LLVM.PointerType(LLVM.Int8Type(), 2)
+
+        # create functions
+        param_types = LLVMType[convert(LLVMType, typ) for typ in arg_types]
+        llvm_f, llvm_ft = create_function(T_void, param_types)
+        # push!(function_attributes(llvm_f), EnumAttribute("alwaysinline", 0))
+
+        mod = LLVM.parent(llvm_f)
+
+        # generate IR
+        @dispose builder=IRBuilder() begin
+            entry = BasicBlock(llvm_f, "entry")
+            position!(builder, entry)
+
+            str = globalstring_ptr!(builder, String(fmt), addrspace=2)
+            argsize = 0
+
+            # construct and fill args buffer
+            if isempty(argspec)
+                buffer = LLVM.PointerNull(T_pint8)
+            else
+                argtypes = LLVM.StructType("os_log_args")
+                elements!(argtypes, param_types)
+
+                args = alloca!(builder, argtypes)
+                for (i, param) in enumerate(parameters(llvm_f))
+                    p = struct_gep!(builder, argtypes, args, i-1)
+                    store!(builder, param, p)
+                end
+
+                dl = datalayout(mod) 
+                argsize = sizeof(dl, argtypes)
+                buffer = bitcast!(builder, args, T_pint8)
+            end
+
+            # invoke @air.os_log and return
+            subsystem_str = LLVM.PointerNull(T_pint8a2)
+            intptr = LLVM.ConstantInt(T_int64, Int64(-1))
+            category_str = const_inttoptr(intptr, T_pint8a2)
+            log_type_default = LLVM.ConstantInt(T_int32, Int32(0))
+            log_type_debug = LLVM.ConstantInt(T_int32, Int32(2))
+            log_type_error = LLVM.ConstantInt(T_int32, Int32(16))
+            arg_size = LLVM.ConstantInt(T_int64, Int64(argsize))
+
+            # declare void @air.os_log(i8 addrspace(2)*, i8 addrspace(2)*, i32, i8 addrspace(2)*, i8*, i64) local_unnamed_addr #4
+            os_log_fty = LLVM.FunctionType(T_void, [T_pint8a2, T_pint8a2, T_int32, T_pint8a2, T_pint8, T_int64])
+            os_log = LLVM.Function(mod, "air.os_log", os_log_fty)
+            call!(builder, os_log_fty, os_log, [subsystem_str, category_str, log_type_error, str, buffer, arg_size])
+            ret!(builder)
+        end
+
+        call_function(llvm_f, Nothing, Tuple{arg_types...}, arg_exprs...)
+    end
+end
+
+
+## print-like functionality
+
+export @mtlprint, @mtlprintln
+
+# simple conversions, defining an expression and the resulting argument type. nothing fancy,
+# `@mtlprint` pretty directly maps to `@mtlprintf`; we should just support `write(::IO)`.
+const mtlprint_conversions = [
+    Float32         => (x->:(Float64($x)),          Float64),
+    Ptr{<:Any}      => (x->:(reinterpret(Int, $x)), Ptr{Cvoid}),
+    LLVMPtr{<:Any}  => (x->:(reinterpret(Int, $x)), Ptr{Cvoid}),
+    Bool            => (x->:(Int32($x)),            Int32),
+]
+
+# format specifiers
+const mtlprint_specifiers = Dict(
+    # integers
+    Int16       => "%hd",
+    Int32       => "%d",
+    Int64       => "%ld",
+    UInt16      => "%hu",
+    UInt32      => "%u",
+    UInt64      => "%lu",
+
+    # floating-point
+    Float32     => "%f",
+
+    # other
+    Cchar       => "%c",
+    Ptr{Cvoid}  => "%p",
+    Cstring     => "%s",
+)
+
+@inline @generated function _mtlprint(parts...)
+    fmt = ""
+    args = Expr[]
+
+    for i in 1:length(parts)
+        part = :(parts[$i])
+        T = parts[i]
+
+        # put literals directly in the format string
+        if T <: Val
+            fmt *= string(T.parameters[1])
+            continue
+        end
+
+        # try to convert arguments if they are not supported directly
+        if !haskey(mtlprint_specifiers, T)
+            for (Tmatch, rule) in mtlprint_conversions
+                if T <: Tmatch
+                    part = rule[1](part)
+                    T = rule[2]
+                    break
+                end
+            end
+        end
+
+        # render the argument
+        if haskey(mtlprint_specifiers, T)
+            fmt *= mtlprint_specifiers[T]
+            push!(args, part)
+        elseif T <: Tuple
+            fmt *= "("
+            for (j, U) in enumerate(T.parameters)
+                if haskey(mtlprint_specifiers, U)
+                    fmt *= mtlprint_specifiers[U]
+                    push!(args, :($part[$j]))
+                    if j < length(T.parameters)
+                        fmt *= ", "
+                    elseif length(T.parameters) == 1
+                        fmt *= ","
+                    end
+                else
+                    @error("@mtlprint does not support values of type $U")
+                end
+            end
+            fmt *= ")"
+        elseif T <: String
+            @error("@mtlprint does not support non-literal strings")
+        elseif T <: Type
+            fmt *= string(T.parameters[1])
+        else
+            @warn("@mtlprint does not support values of type $T")
+            fmt *= "$(T)(...)"
+        end
+    end
+
+    quote
+        @mtlprintf($fmt, $(args...))
+    end
+end
+
+"""
+    @mtlprint(xs...)
+    @mtlprintln(xs...)
+
+Print a textual representation of values `xs` to standard output from the GPU. The
+functionality builds on `@mtlprintf`, and is intended as a more use friendly alternative of
+that API. However, that also means there's only limited support for argument types, handling
+16/32/64 signed and unsigned integers, 32 and 64-bit floating point numbers, `Cchar`s and
+pointers. For more complex output, use `@mtlprintf` directly.
+
+Limited string interpolation is also possible:
+
+```julia
+    @mtlprint("Hello, World ", 42, "\\n")
+    @mtlprint "Hello, World \$(42)\\n"
+```
+"""
+macro mtlprint(parts...)
+    args = Union{Val,Expr,Symbol}[]
+
+    parts = [parts...]
+    while true
+        isempty(parts) && break
+
+        part = popfirst!(parts)
+
+        # handle string interpolation
+        if isa(part, Expr) && part.head == :string
+            parts = vcat(part.args, parts)
+            continue
+        end
+
+        # expose literals to the generator by using Val types
+        if isbits(part) # literal numbers, etc
+            push!(args, Val(part))
+        elseif isa(part, QuoteNode) # literal symbols
+            push!(args, Val(part.value))
+        elseif isa(part, String) # literal strings need to be interned
+            push!(args, Val(Symbol(part)))
+        else # actual values that will be passed to printf
+            push!(args, part)
+        end
+    end
+
+    quote
+        _mtlprint($(map(esc, args)...))
+    end
+end
+
+@doc (@doc @mtlprint) ->
+macro mtlprintln(parts...)
+    esc(quote
+        Metal.@mtlprint($(parts...), "\n")
+    end)
+end
+
+export @mtlshow
+
+"""
+    @mtlshow(ex)
+
+GPU analog of `Base.@show`. It comes with the same type restrictions as [`@mtlprintf`](@ref).
+
+```julia
+@mtlshow thread_position_in_grid_1d()
+```
+"""
+macro mtlshow(exs...)
+    blk = Expr(:block)
+    for ex in exs
+        push!(blk.args, :(Metal.@mtlprintln($(sprint(Base.show_unquoted,ex)*" = "),
+                                          begin local value = $(esc(ex)) end)))
+    end
+    isempty(exs) || push!(blk.args, :value)
+    blk
+end
diff --git a/src/state.jl b/src/state.jl
@@ -47,7 +47,29 @@ function global_queue(dev::MTLDevice)
         @autoreleasepool begin
             # NOTE: MTLCommandQueue itself is manually reference-counted,
             #       the release pool is for resources used during its construction.
-            queue = MTLCommandQueue(dev)
+            log_state_descriptor = MTLLogStateDescriptor()
+            log_state = MTLLogState(dev, log_state_descriptor)
+
+            function log_handler(subSystem, category, logLevel, message)
+                println(String(message))
+                return nothing
+            end
+            function wrapper(subSystem, category, logLevel, message)
+                log_handler(subSystem == nil ? nothing : NSString(subSystem), 
+                            category == nil ? nothing : NSString(category),
+                            logLevel,
+                            message == nil ? nothing : NSString(message)
+                )
+                return nothing
+            end
+
+            block = @objcblock(wrapper, Nothing, (id{NSString}, id{NSString}, MTLLogLevel, id{NSString}))
+
+            @objc [log_state::id{MTLLogState} addLogHandler:block::id{NSBlock}]::Nothing
+
+            queue_descriptor = MTLCommandQueueDescriptor()
+            queue_descriptor.logState = log_state
+            queue = MTLCommandQueue(dev, queue_descriptor)
             queue.label = "global_queue($(current_task()))"
             global_queues[queue] = nothing
             queue
