switch to NVTX when available

Author: vchuravy

src/KernelAbstractions.jl

@@ -430,6 +430,15 @@ end
     end
 end
 
+###
+# Extras
+# - LoopInfo
+# - Timeline
+###
+
+include("extras/extras.jl")
+import .Extras.Timeline
+
 ###
 # Backends/Implementation
 ###
@@ -442,11 +451,4 @@ include("backends/cpu.jl")
 @init @require CUDAnative="be33ccc6-a3ff-5ff2-a52e-74243cff1e17" begin
     include("backends/cuda.jl")
 end
-
-###
-# Extras
-# - LoopInfo
-###
-
-include("extras/extras.jl")
 end #module

src/backends/cpu.jl

@@ -8,8 +8,11 @@ end
 
 function Event(f, args...; dependencies=nothing, progress=nothing)
     T = Threads.@spawn begin
-        wait(MultiEvent(dependencies), progress)
-        f(args...)
+        Timeline.range "Event($(nameof(f))" begin
+            wait(MultiEvent(dependencies), progress)
+            Timeline.mark("Event($(nameof(f))) waiting done")
+            f(args...)
+        end
     end
     return CPUEvent(T)
 end
@@ -79,13 +82,15 @@ function __run(obj, ndrange, iterspace, args, ::Val{dynamic}) where dynamic
         Nthreads = N
         len, rem = 1, 0
     end
+    Timeline.@range string(nameof(obj.f)) begin
     if Nthreads == 1
         __thread_run(1, len, rem, obj, ndrange, iterspace, args, Val(dynamic))
     else
         @sync for tid in 1:Nthreads
             Threads.@spawn __thread_run(tid, len, rem, obj, ndrange, iterspace, args, Val(dynamic))
         end
     end
+    end # Timeline
     return nothing
 end
 

src/extras/extras.jl

@@ -6,5 +6,6 @@ export @unroll
 
 include("timeline.jl")
 using .Timeline
+export Timeline
 
 end # module

src/extras/timeline.jl

@@ -1,5 +1,8 @@
 module Timeline
 
+using Requires
+export @range, mark
+
 module NVTXT
     const LOG_FILE=Ref{IOStream}()
     const SHOULD_LOG=Ref{Bool}(false)
@@ -91,7 +94,42 @@ module NVTXT
     end
 end # NVTXT
 
-function range_push end
-function mark end
+_mark(msg) = NVTXT.mark(msg)
+_push_range(msg) = NVTXT.push_range(msg)
+_pop_range() = NVTXT.pop_range()
+_start_range(msg) = NVTXT.start_range(msg)
+_end_range(r) = NVTXT.end_range(r)
+
+@init @require CUDAnative="be33ccc6-a3ff-5ff2-a52e-74243cff1e17" begin
+    # replace implementations
+    import CUDAnative.NVTX
+
+    _mark(msg) = NVTX.mark(msg)
+    _push_range(msg) = NVTX.push_range(msg)
+    _pop_range() = NVTX.pop_range()
+    _start_range(msg) = NVTX.start_range(msg)
+    _end_range(r) = NVTX.end_range(r)
+end
+
+import Base: invokelatest
+mark(msg) = invokelatest(_mark, msg)
+push_range(msg) = invokelatest(_push_range, msg)
+pop_range() = invokelatest(_pop_range)
+start_range(msg) = invokelatest(_start_range, msg)
+end_range(r) = invokelatest(_end_range, r)
+
+"""
+    @range "msg" ex
+Create a new range and execute `ex`. The range is popped automatically afterwards.
+See also: [`range`](@ref)
+"""
+macro range(msg, ex)
+    quote
+        local range = $start_range($(esc(msg)))
+        local ret = $(esc(ex))
+        $end_range(range)
+        ret
+    end
+end
 
-end
+end