CUPTI: Add high-level wrappers for the callback API. (#2239)

Author: maleadt

lib/cupti/error.jl

@@ -88,7 +88,9 @@ function description(err)
     elseif err.code == CUPTI_ERROR_CUDA_COMPILER_NOT_COMPATIBLE
         "Profiling results might be incorrect for CUDA applications compiled with nvcc version older than 9.0 for devices with compute capability 6.0 and 6.1"
     elseif err.code == CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
-        "user doesn't have sufficient privileges which are required to start the profiling session"
+        """Insufficient privileges: You don't have permissions to profile GPU code.
+           Please configure your system to allow all users to profile, or run Julia with
+           elevated permissions: https://developer.nvidia.com/ERR_NVGPUCTRPERM#SolnAdminTag"""
     elseif err.code == CUPTI_ERROR_OLD_PROFILER_API_INITIALIZED
         "old profiling api's are not supported with new profiling api's"
     elseif err.code == CUPTI_ERROR_OPENACC_UNDEFINED_ROUTINE

lib/cupti/wrappers.jl

@@ -5,18 +5,98 @@ function version()
 end
 
 
+#
+# callback API
+#
+
+# multiple subscribers aren't supported, so make sure we only call CUPTI once
+const callback_lock = ReentrantLock()
+
+function callback(userdata::Ptr{Cvoid}, domain::CUpti_CallbackDomain,
+                  id::CUpti_CallbackId, data_ptr::Ptr{Cvoid})
+    cfg = Base.unsafe_pointer_to_objref(userdata)::CallbackConfig
+
+    # decode the callback data
+    datatype = if domain in (CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_CB_DOMAIN_RUNTIME_API)
+        CUpti_CallbackData
+    elseif domain == CUPTI_CB_DOMAIN_RESOURCE
+        CUpti_ResourceData
+    elseif domain == CUPTI_CB_DOMAIN_SYNCHRONIZE
+        CUpti_SynchronizeData
+    elseif domain == CUPTI_CB_DOMAIN_NVTX
+        CUpti_NvtxData
+    else
+        @warn """Unsupported callback domain: $(domain).
+                 Please file an issue, or extend the implementation of `CUPTI.callback` to handle this callback kind."""
+        return
+    end
+    data = unsafe_load(convert(Ptr{datatype}, data_ptr))
+
+    # invoke the actual user callback
+    cfg.callback(domain, id, data)
+
+    return
+end
+
+"""
+    cfg = CUPTI.CallbackConfig(callback_kinds) do domain, id, data
+        # inspect data
+    end
+
+    CUPTI.enable!(cfg) do
+        # do stuff
+    end
+"""
+mutable struct CallbackConfig
+    callback::Function
+    callback_kinds::Vector{CUpti_CallbackDomain}
+end
+
+function enable!(f::Base.Callable, cfg::CallbackConfig)
+    @lock callback_lock begin
+        callback_ptr =
+            @cfunction(callback, Cvoid,
+                       (Ptr{Cvoid}, CUpti_CallbackDomain, CUpti_CallbackId, Ptr{Cvoid}))
+
+        GC.@preserve cfg begin
+            # set-up subscriber
+            subscriber_ref = Ref{CUpti_SubscriberHandle}()
+            cuptiSubscribe(subscriber_ref, callback_ptr, Base.pointer_from_objref(cfg))
+            subscriber = subscriber_ref[]
+
+            # enable domains
+            for callback_kind in cfg.callback_kinds
+                CUPTI.cuptiEnableDomain(true, subscriber, callback_kind)
+            end
+
+            try
+                f()
+            finally
+                # disable callback kinds
+                for callback_kind in cfg.callback_kinds
+                    CUPTI.cuptiEnableDomain(false, subscriber, callback_kind)
+                end
+
+                # disable the subscriber
+                CUPTI.cuptiUnsubscribe(subscriber)
+            end
+        end
+    end
+end
+
+
 #
 # activity API
 #
 
 """
-    cfg = ActvitiyConfig(activity_kinds)
+    cfg = CUPTI.ActivityConfig(activity_kinds)
 
-    enable!(cfg)
-    # do stuff
-    disable!(cfg)
+    CUPTI.enable!(cfg) do
+        # do stuff
+    end
 
-    process(cfg) do (ctx, stream, record)
+    CUPTI.process(cfg) do ctx, stream_id, record
         # inspect record
     end
 
@@ -43,6 +123,7 @@ function allocate_buffer()
     Array{UInt8}(undef, 8 * 1024 * 1024)  # 8 MB
 end
 
+const activity_lock = ReentrantLock()
 const activity_config = Ref{Union{Nothing,ActivityConfig}}(nothing)
 
 function request_buffer(dest_ptr, sz_ptr, max_num_records_ptr)
@@ -93,41 +174,40 @@ function complete_buffer(ctx_handle, stream_id, buf_ptr, sz, valid_sz)
     return
 end
 
-function enable!(cfg::ActivityConfig)
-    activity_config[] === nothing ||
-        error("Only one profiling session can be active at a time.")
-
-    # set-up callbacks
-    request_buffer_ptr = @cfunction(request_buffer, Cvoid,
-                                    (Ptr{Ptr{UInt8}}, Ptr{Csize_t}, Ptr{Csize_t}))
-    complete_buffer_ptr = @cfunction(complete_buffer, Cvoid,
-                                     (CUDA.CUcontext, UInt32, Ptr{UInt8}, Csize_t, Csize_t))
-    cuptiActivityRegisterCallbacks(request_buffer_ptr, complete_buffer_ptr)
+function enable!(f::Base.Callable, cfg::ActivityConfig)
+    @lock activity_lock begin
+        activity_config[] = cfg
 
-    activity_config[] = cfg
+        # set-up callbacks
+        request_buffer_ptr =
+            @cfunction(request_buffer, Cvoid,
+                       (Ptr{Ptr{UInt8}}, Ptr{Csize_t}, Ptr{Csize_t}))
+        complete_buffer_ptr =
+            @cfunction(complete_buffer, Cvoid,
+                       (CUDA.CUcontext, UInt32, Ptr{UInt8}, Csize_t, Csize_t))
+        cuptiActivityRegisterCallbacks(request_buffer_ptr, complete_buffer_ptr)
 
-    # enable requested activity kinds
-    for activity_kind in cfg.activity_kinds
-        cuptiActivityEnable(activity_kind)
-    end
-end
+        activity_config[] = cfg
 
-function disable!(cfg::ActivityConfig)
-    if activity_config[] !== cfg
-        error("This profiling session is not active.")
-    end
-
-    # disable activity kinds
-    for activity_kind in cfg.activity_kinds
-        cuptiActivityDisable(activity_kind)
-    end
+        # enable requested activity kinds
+        for activity_kind in cfg.activity_kinds
+            cuptiActivityEnable(activity_kind)
+        end
 
-    # flush all activity records, even incomplete ones
-    cuptiActivityFlushAll(CUPTI_ACTIVITY_FLAG_FLUSH_FORCED)
+        try
+            f()
+        finally
+            # disable activity kinds
+            for activity_kind in cfg.activity_kinds
+                cuptiActivityDisable(activity_kind)
+            end
 
-    activity_config[] = nothing
+            # flush all activity records, even incomplete ones
+            cuptiActivityFlushAll(CUPTI_ACTIVITY_FLAG_FLUSH_FORCED)
 
-    return
+            activity_config[] = nothing
+        end
+    end
 end
 
 function process(f, cfg::ActivityConfig)

src/compiler/reflection.jl

@@ -1,33 +1,14 @@
 # code reflection entry-points
 
 using .CUPTI
-using .CUPTI: CUpti_CallbackDomain, CUpti_CallbackId, CUpti_SubscriberHandle,
-              CUpti_ResourceData, CUpti_ModuleResourceData
+using .CUPTI: CUpti_ModuleResourceData
 
 
 
 #
 # code_* replacements
 #
 
-function code_sass_callback(userdata::Ptr{Cvoid}, domain::CUpti_CallbackDomain,
-                            cbid::CUpti_CallbackId, cbdada::Ptr{Cvoid})
-    dest = Base.unsafe_pointer_to_objref(userdata)::Ref{Any}
-
-    if domain == CUPTI.CUPTI_CB_DOMAIN_RESOURCE
-        cbdada = unsafe_load(reinterpret(Ptr{CUpti_ResourceData}, cbdada))
-        if cbid == CUPTI.CUPTI_CBID_RESOURCE_MODULE_LOADED
-            resourceDescriptor =
-                unsafe_load(reinterpret(Ptr{CUpti_ModuleResourceData}, cbdada.resourceDescriptor))
-            cubin = unsafe_wrap(Vector{Cchar}, pointer(resourceDescriptor.pCubin),
-                                resourceDescriptor.cubinSize)
-            dest[] = copy(cubin)
-        end
-    end
-
-    return
-end
-
 """
     code_sass([io], f, types; raw=false)
 
@@ -49,8 +30,8 @@ function code_sass(io::IO, @nospecialize(func), @nospecialize(types); kwargs...)
     code_sass(io, job; kwargs...)
 end
 
-# multiple subscribers aren't supported, so make sure we only call CUPTI once
-const cupti_lock = ReentrantLock()
+code_sass(@nospecialize(func), @nospecialize(types); kwargs...) =
+    code_sass(stdout, func, types; kwargs...)
 
 function code_sass(io::IO, job::CompilerJob; raw::Bool=false)
     if !job.config.kernel
@@ -64,35 +45,27 @@ function code_sass(io::IO, job::CompilerJob; raw::Bool=false)
         return
     end
 
-    compiled = compile(job)
+    cfg = CUPTI.CallbackConfig([CUPTI.CUPTI_CB_DOMAIN_RESOURCE]) do domain, id, data
+        id == CUPTI.CUPTI_CBID_RESOURCE_MODULE_LOADED || return
+        resourceDescriptor =
+            unsafe_load(convert(Ptr{CUpti_ModuleResourceData}, data.resourceDescriptor))
+        cubin = unsafe_wrap(Vector{Cchar}, pointer(resourceDescriptor.pCubin),
+                            resourceDescriptor.cubinSize)
+        disassemble_cubin(io, cubin; raw)
+    end
 
-    cubin = Ref{Any}()
-    callback = @cfunction(code_sass_callback, Cvoid,
-                          (Ptr{Cvoid}, CUpti_CallbackDomain, CUpti_CallbackId, Ptr{Cvoid}))
-
-    # JIT compile and capture the generated object file
-    lock(cupti_lock) do
-        subscriber_ref = Ref{CUpti_SubscriberHandle}()
-        res = CUPTI.unsafe_cuptiSubscribe(subscriber_ref, callback, Base.pointer_from_objref(cubin))
-        if res === CUPTI.CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
-            error("""Insufficient privilege: You don't have permissions to profile GPU code, which is required for `code_sass`.
-                    Get administrative privileges or allow all users to profile: https://developer.nvidia.com/ERR_NVGPUCTRPERM#SolnAdminTag""")
-        elseif res != CUPTI.CUPTI_SUCCESS
-            throw(CUPTIError(res))
-        end
-        subscriber = subscriber_ref[]
-        try
-            CUPTI.cuptiEnableDomain(1, subscriber, CUPTI.CUPTI_CB_DOMAIN_RESOURCE)
-            link(job, compiled)
-        finally
-            CUPTI.cuptiUnsubscribe(subscriber)
-        end
+    compiled = compile(job)
+    CUPTI.enable!(cfg) do
+        link(job, compiled)
     end
 
-    # disassemble to SASS
-    isassigned(cubin) || error("No kernels compiled")
+    return
+end
+
+# disassemble a cubin to SASS
+function disassemble_cubin(io::IO, cubin::Vector{Cchar}; raw::Bool)
     mktemp() do cubin_path,cubin_io
-        write(cubin_io, cubin[])
+        write(cubin_io, cubin)
         flush(cubin_io)
 
         cmd = `$(nvdisasm()) --print-code --print-line-info $cubin_path`
@@ -111,9 +84,6 @@ function code_sass(io::IO, job::CompilerJob; raw::Bool=false)
     end
 end
 
-code_sass(@nospecialize(func), @nospecialize(types); kwargs...) =
-    code_sass(stdout, func, types; kwargs...)
-
 
 # forward the rest to GPUCompiler with an appropriate CompilerJob
 

src/profile.jl

@@ -270,18 +270,16 @@ function profile_internally(f; concurrent=true, kwargs...)
     GC.gc(false)
     GC.gc(true)
 
-    CUPTI.enable!(cfg)
+    CUPTI.enable!(cfg) do
+        # sink the initial profiler overhead into a synchronization call
+        CUDA.cuCtxSynchronize()
 
-    # sink the initial profiler overhead into a synchronization call
-    CUDA.cuCtxSynchronize()
-    try
         f()
 
         # synchronize to ensure we capture all activity
         CUDA.cuCtxSynchronize()
-    finally
-        CUPTI.disable!(cfg)
     end
+
     data = capture(cfg)
     ProfileResults(; data..., kwargs...)
 end