Add support for NVTX events to the integrated profiler. (#2043)

Author: maleadt

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.8.5"
 manifest_format = "2.0"
-project_hash = "56ddf5ba3ee0cbfd2fc5e239e07860c2adb9b489"
+project_hash = "62b8680a215cbfae7c33bbe0a528c176464c0f65"
 
 [[deps.AbstractFFTs]]
 deps = ["ChainRulesCore", "LinearAlgebra", "Test"]
@@ -73,6 +73,18 @@ git-tree-sha1 = "2fba81a302a7be671aefe194f0525ef231104e7f"
 uuid = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0"
 version = "0.1.8"
 
+[[deps.ColorTypes]]
+deps = ["FixedPointNumbers", "Random"]
+git-tree-sha1 = "eb7f0f8307f71fac7c606984ea5fb2817275d6e4"
+uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f"
+version = "0.11.4"
+
+[[deps.Colors]]
+deps = ["ColorTypes", "FixedPointNumbers", "Reexport"]
+git-tree-sha1 = "fc08e5930ee9a4e03f84bfb5211cb54e7769758a"
+uuid = "5ae59095-9a9b-59fe-a467-6f913c188581"
+version = "0.12.10"
+
 [[deps.Compat]]
 deps = ["Dates", "LinearAlgebra", "UUIDs"]
 git-tree-sha1 = "e460f044ca8b99be31d35fe54fc33a5c33dd8ed7"
@@ -134,6 +146,12 @@ version = "0.1.10"
 [[deps.FileWatching]]
 uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
 
+[[deps.FixedPointNumbers]]
+deps = ["Statistics"]
+git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
+uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
+version = "0.8.4"
+
 [[deps.Future]]
 deps = ["Random"]
 uuid = "9fa8497b-333b-5362-9e8d-4d0656e87820"
@@ -193,6 +211,12 @@ git-tree-sha1 = "7e5d6779a1e09a36db2a7b6cff50942a0a7d0fca"
 uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
 version = "1.5.0"
 
+[[deps.JuliaNVTXCallbacks_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "af433a10f3942e882d3c671aacb203e006a5808f"
+uuid = "9c1d0b0a-7046-5b2e-a33f-ea22f176ac7e"
+version = "0.2.1+0"
+
 [[deps.KernelAbstractions]]
 deps = ["Adapt", "Atomix", "InteractiveUtils", "LinearAlgebra", "MacroTools", "PrecompileTools", "Requires", "SparseArrays", "StaticArrays", "UUIDs", "UnsafeAtomics", "UnsafeAtomicsLLVM"]
 git-tree-sha1 = "4c5875e4c228247e1c2b087669846941fb6e0118"
@@ -280,6 +304,18 @@ version = "1.1.0"
 uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
 version = "2022.2.1"
 
+[[deps.NVTX]]
+deps = ["Colors", "JuliaNVTXCallbacks_jll", "Libdl", "NVTX_jll"]
+git-tree-sha1 = "ab760fa11c7f12df87334ea9e0dffc54af3025bc"
+uuid = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
+version = "0.3.2"
+
+[[deps.NVTX_jll]]
+deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
+git-tree-sha1 = "ce3269ed42816bf18d500c9f63418d4b0d9f5a3b"
+uuid = "e98f9f5b-d649-5603-91fd-7774390e6439"
+version = "3.1.0+2"
+
 [[deps.NetworkOptions]]
 uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
 version = "1.2.0"

Project.toml

@@ -21,6 +21,7 @@ LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
@@ -49,6 +50,7 @@ GPUArrays = "8.6"
 GPUCompiler = "0.22"
 KernelAbstractions = "0.9.2"
 LLVM = "6"
+NVTX = "0.3.2"
 Preferences = "1"
 PrettyTables = "2"
 Random123 = "1.2"

lib/cupti/wrappers.jl

@@ -135,6 +135,9 @@ function process(f, cfg::ActivityConfig)
         CUPTI_ACTIVITY_KIND_DRIVER              => CUpti_ActivityAPI,
         CUPTI_ACTIVITY_KIND_RUNTIME             => CUpti_ActivityAPI,
         CUPTI_ACTIVITY_KIND_INTERNAL_LAUNCH_API => CUpti_ActivityAPI,
+        CUPTI_ACTIVITY_KIND_NAME                => CUpti_ActivityName,
+        CUPTI_ACTIVITY_KIND_MARKER              => CUpti_ActivityMarker2,
+        CUPTI_ACTIVITY_KIND_MARKER_DATA         => CUpti_ActivityMarkerData,
     )
     # NOTE: the CUPTI version is unreliable, e.g., both CUDA 11.5 and 11.6 have CUPTI 16,
     #       yet CUpti_ActivityMemset4 is only available in CUDA 11.6.

src/CUDA.jl

@@ -34,6 +34,8 @@ import Preferences
 
 using Libdl
 
+import NVTX
+
 
 ## source code includes
 

src/initialization.jl

@@ -119,6 +119,12 @@ function __init__()
                  It is recommended to upgrade your driver, or switch to automatic installation of CUDA."""
     end
 
+    # if we're not running under an external profiler, let CUPTI handle NVTX events
+    if !NVTX.isactive()
+        ENV["NVTX_INJECTION64_PATH"] = CUDA_Runtime.libcupti
+        NVTX.activate()
+    end
+
     # finally, initialize CUDA
     try
         cuInit(0)

src/profile.jl

@@ -1,7 +1,7 @@
 # Profiler control
 
 """
-    @profile [io=stdout] [host=true] [device=true] [trace=false] [raw=false] code...
+    @profile [io=stdout] [trace=false] [raw=false] code...
     @profile external=true code...
 
 Profile the GPU execution of `code`.
@@ -11,22 +11,19 @@ There are two modes of operation, depending on whether `external` is `true` or `
 ## Integrated profiler (`external=false`, the default)
 
 In this mode, CUDA.jl will profile the execution of `code` and display the result. By
-default, both host-side and device-side activity is captured; this can be controlled with
-the `host` and `device` keyword arguments. If `trace` is `true`, a chronological trace of
-the captured activity will be generated, where the ID column can be used to match host-side
-and device-side activity; by default, only a summary will be shown. If `raw` is `true`, all
-data will always be included, even if it may not be relevant. The output will be written to
-`io`, which defaults to `stdout`.
+default, a summary of host and device-side execution will be show, including any NVTX
+events. To display a chronological trace of the captured activity instead, `trace` can be
+set to `true`. Trace output will include an ID column that can be used to match host-side
+and device-side activity. If `raw` is `true`, all data will always be included, even if it
+may not be relevant. The output will be written to `io`, which defaults to `stdout`.
 
 Slow operations will be highlighted in the output: Entries colored in yellow are among the
 slowest 25%, while entries colored in red are among the slowest 5% of all operations.
 
-!!! compat "Julia 1.9"
-    This functionality is only available on Julia 1.9 and later.
+!!! compat "Julia 1.9" This functionality is only available on Julia 1.9 and later.
 
-!!! compat "CUDA 11.2"
-    Older versions of CUDA, before 11.2, contain bugs that may prevent the
-    `CUDA.@profile` macro to work. It is recommended to use a newer runtime.
+!!! compat "CUDA 11.2" Older versions of CUDA, before 11.2, contain bugs that may prevent
+    the `CUDA.@profile` macro to work. It is recommended to use a newer runtime.
 
 ## External profilers (`external=true`)
 
@@ -187,6 +184,8 @@ function emit_integrated_profile(code, kwargs)
         # memory operations
         CUPTI.CUPTI_ACTIVITY_KIND_MEMCPY,
         CUPTI.CUPTI_ACTIVITY_KIND_MEMSET,
+        # NVTX markers
+        CUPTI.CUPTI_ACTIVITY_KIND_MARKER,
     ]
     if CUDA.runtime_version() >= v"11.2"
         # additional information for API host calls
@@ -260,6 +259,14 @@ function generate_traces(cfg)
         id      = Int[],
         details = String[],
     )
+    nvtx_trace = DataFrame(
+        id      = Int[],
+        start   = Float64[],
+        type    = Symbol[],
+        tid     = Int[],
+        name    = Union{Missing,String}[],
+        domain  = Union{Missing,String}[],
+    )
 
     # memory_kind fields are sometimes typed CUpti_ActivityMemoryKind, sometimes UInt
     as_memory_kind(x) = isa(x, CUPTI.CUpti_ActivityMemoryKind) ? x : CUPTI.CUpti_ActivityMemoryKind(x)
@@ -349,18 +356,44 @@ function generate_traces(cfg)
                                    stream=record.streamId,
                                    grid, block, registers,
                                    static_shmem, dynamic_shmem); cols=:union)
+
+        # NVTX markers
+        elseif record.kind == CUPTI.CUPTI_ACTIVITY_KIND_MARKER
+            start = record.timestamp/1e9
+            name = record.name == C_NULL ? missing : unsafe_string(record.name)
+            domain = record.domain == C_NULL ? missing : unsafe_string(record.domain)
+
+            if record.flags == CUPTI.CUPTI_ACTIVITY_FLAG_MARKER_INSTANTANEOUS
+                @assert record.objectKind == CUDA.CUPTI.CUPTI_ACTIVITY_OBJECT_THREAD
+                tid = record.objectId.pt.threadId
+                push!(nvtx_trace, (; record.id, start, tid, type=:instant, name, domain))
+            elseif record.flags == CUPTI.CUPTI_ACTIVITY_FLAG_MARKER_START
+                @assert record.objectKind == CUDA.CUPTI.CUPTI_ACTIVITY_OBJECT_THREAD
+                tid = record.objectId.pt.threadId
+                push!(nvtx_trace, (; record.id, start, tid, type=:start, name, domain))
+            elseif record.flags == CUPTI.CUPTI_ACTIVITY_FLAG_MARKER_END
+                @assert record.objectKind == CUDA.CUPTI.CUPTI_ACTIVITY_OBJECT_THREAD
+                tid = record.objectId.pt.threadId
+                push!(nvtx_trace, (; record.id, start, tid, type=:end, name, domain))
+            else
+                @error "Unexpected NVTX marker kind $(Int(record.flags)). Please file an issue."
+            end
         else
-            error("Unexpected CUPTI activity kind: $(record.kind). Please file an issue.")
+            @error "Unexpected CUPTI activity kind $(Int(record.kind)). Please file an issue."
         end
     end
 
-    return host_trace, device_trace, details
+    # merge in the details
+    host_trace = leftjoin(host_trace, details, on=:id, order=:left)
+    device_trace = leftjoin(device_trace, details, on=:id, order=:left)
+
+    return host_trace, device_trace, nvtx_trace
 end
 
 # render traces to a table
-function render_traces(host_trace, device_trace, details;
+function render_traces(host_trace, device_trace, nvtx_trace;
                        io=stdout isa Base.TTY ? IOContext(stdout, :limit => true) : stdout,
-                       host=true, device=true, trace=false, raw=false)
+                       trace=false, raw=false)
     # find the relevant part of the trace (marked by calls to 'cuCtxSynchronize')
     trace_first_sync = findfirst(host_trace.name .== "cuCtxSynchronize")
     trace_first_sync === nothing && error("Could not find the start of the profiling trace.")
@@ -400,16 +433,21 @@ function render_traces(host_trace, device_trace, details;
         df.start .-= trace_begin
         df.stop .-= trace_begin
     end
+    nvtx_trace.start .-= trace_begin
     if !raw
         # renumber event IDs from 1
-        first_id = minimum([host_trace.id; device_trace.id; details.id])
-        for df in (host_trace, device_trace, details)
+        first_id = minimum([host_trace.id; device_trace.id])
+        for df in (host_trace, device_trace)
             df.id .-= first_id - 1
         end
 
         # renumber thread IDs from 1
-        first_tid = minimum(host_trace.tid)
-        host_trace.tid .-= first_tid - 1
+        threads = unique([host_trace.tid; nvtx_trace.tid])
+        for df in (host_trace, nvtx_trace)
+            broadcast!(df.tid, df.tid) do tid
+                findfirst(isequal(tid), threads)
+            end
+        end
 
     end
 
@@ -480,7 +518,7 @@ function render_traces(host_trace, device_trace, details;
         "name"          => "Name"
     )
 
-    summary_formatter = function(v, i, j)
+    summary_formatter(df) = function(v, i, j)
         if names(df)[j] == "time_ratio"
             format_percentage(v)
         elseif names(df)[j] in ["time", "time_avg", "time_min", "time_max"]
@@ -497,11 +535,11 @@ function render_traces(host_trace, device_trace, details;
         :horizontal
     end
 
-    if host
+    # host-side activity
+    let
         # to determine the time the host was active, we should look at threads separately
         host_time = maximum(combine(groupby(host_trace, :tid), :time => sum => :time).time)
         host_ratio = host_time / trace_time
-        println(io, "\nHost-side activity: calling CUDA APIs took $(format_time(host_time)) ($(format_percentage(host_ratio)) of the trace)")
 
         # get rid of API call version suffixes
         host_trace.name = replace.(host_trace.name, r"_v\d+$" => "")
@@ -524,11 +562,28 @@ function render_traces(host_trace, device_trace, details;
             end
         end
 
-        # add in details
-        df = leftjoin(df, details, on=:id, order=:left)
+        # instantaneous NVTX markers can be added to the API trace
+        if trace
+            markers = copy(nvtx_trace[nvtx_trace.type .== :instant, :])
+            markers.id .= missing
+            markers.time .= 0.0
+            markers.details = map(markers.name, markers.domain) do name, domain
+                if name !== missing && domain !== missing
+                    "$(domain).$(name)"
+                elseif name !== missing
+                    "$name"
+                end
+            end
+            markers.name .= "NVTX marker"
+            append!(df, markers; cols=:subset)
+            sort!(df, :start)
+        end
 
+        if !isempty(df)
+            println(io, "\nHost-side activity: calling CUDA APIs took $(format_time(host_time)) ($(format_percentage(host_ratio)) of the trace)")
+        end
         if isempty(df)
-            println(io, "No host-side activity was recorded.")
+            println(io, "\nNo host-side activity was recorded.")
         elseif trace
             # determine columns to show, based on whether they contain useful information
             columns = [:id, :start, :time]
@@ -566,20 +621,23 @@ function render_traces(host_trace, device_trace, details;
             header = [summary_column_names[name] for name in names(df)]
             alignment = [i == lastindex(header) ? :l : :r for i in 1:length(header)]
             highlighters = time_highlighters(df)
-            pretty_table(io, df; header, alignment, formatters=summary_formatter, highlighters, crop)
+            pretty_table(io, df; header, alignment, formatters=summary_formatter(df), highlighters, crop)
         end
     end
 
-    if device
+    # device-side activity
+    let
         device_time = sum(device_trace.time)
         device_ratio = device_time / trace_time
-        println(io, "\nDevice-side activity: GPU was busy for $(format_time(device_time)) ($(format_percentage(device_ratio)) of the trace)")
+        if !isempty(device_trace)
+            println(io, "\nDevice-side activity: GPU was busy for $(format_time(device_time)) ($(format_percentage(device_ratio)) of the trace)")
+        end
 
         # add memory throughput information
         device_trace.throughput = device_trace.size ./ device_trace.time
 
         if isempty(device_trace)
-            println(io, "No device-side activity was recorded.")
+            println(io, "\nNo device-side activity was recorded.")
         elseif trace
             # determine columns to show, based on whether they contain useful information
             columns = [:id, :start, :time]
@@ -645,9 +703,43 @@ function render_traces(host_trace, device_trace, details;
             header = [summary_column_names[name] for name in names(df)]
             alignment = [i == lastindex(header) ? :l : :r for i in 1:length(header)]
             highlighters = time_highlighters(df)
-            pretty_table(io, df; header, alignment, formatters=summary_formatter, highlighters, crop)
+            pretty_table(io, df; header, alignment, formatters=summary_formatter(df), highlighters, crop)
         end
     end
+
+    # show NVTX ranges
+    # TODO: do we also want to repeat the host/device summary for each NVTX range?
+    #       that's what nvprof used to do, but it's a little verbose...
+    nvtx_ranges = copy(nvtx_trace[nvtx_trace.type .== :start, :])
+    nvtx_ranges = leftjoin(nvtx_ranges, nvtx_trace[nvtx_trace.type .== :end,
+                                                   [:id, :start]],
+                           on=:id, makeunique=true)
+    if !isempty(nvtx_ranges)
+        println(io, "\nNVTX ranges:")
+
+        rename!(nvtx_ranges, :start_1 => :stop)
+        nvtx_ranges.id .= missing
+        nvtx_ranges.time .= nvtx_ranges.stop .- nvtx_ranges.start
+        nvtx_ranges.name = map(nvtx_ranges.name, nvtx_ranges.domain) do name, domain
+            if name !== missing && domain !== missing
+                "$(domain).$(name)"
+            elseif name !== missing
+                "$name"
+            end
+        end
+
+        df = summarize_trace(nvtx_ranges)
+
+        columns = [:time_ratio, :time, :calls, :time_avg, :time_min, :time_max, :name]
+        df = df[:, columns]
+
+        header = [summary_column_names[name] for name in names(df)]
+        alignment = [i == lastindex(header) ? :l : :r for i in 1:length(header)]
+        highlighters = time_highlighters(df)
+        pretty_table(io, df; header, alignment, formatters=summary_formatter(df), highlighters, crop)
+    end
+
+    return
 end
 
 format_percentage(x::Number) = @sprintf("%.2f%%", x * 100)

test/Project.toml

@@ -12,6 +12,7 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"