Time kernels on the device without synchronizing (prototype)

lib/KernelInterface/src/KernelInterface.jl

@@ -39,6 +39,7 @@ include("host.jl")
             # host side
             :allocate, :zeros, :ones, :copyto!, :pagelock!, :unsafe_free!,
             :synchronize, :record_event, :wait_event, :priority!,
+            :record_timestamp, :elapsed_time,
             :device, :ndevices, :device!,
             :functional, :versioninfo,
             :supports_unified, :supports_atomics, :supports_float64,

lib/KernelInterface/src/host.jl

@@ -379,3 +379,33 @@ block until it has completed. For a simple, synchronous copy, use `Base.copyto!`
     device-to-device copies.
 """
 function copyto! end
+
+"""
+    record_timestamp(backend::Backend)
+
+Enqueue a timestamp on the calling task's queue of `backend`'s active device, without
+blocking, and return a handle for [`elapsed_time`](@ref). Returns `nothing` if the backend
+doesn't support timestamps.
+
+Profilers use this to measure the device time of kernels without synchronizing the host
+with the device.
+
+!!! note
+    Backend implementations **may** implement this function, e.g. with a CUDA event that
+    records timing. The fallback returns `nothing`.
+"""
+record_timestamp(::Backend) = nothing
+
+"""
+    elapsed_time(backend::Backend, start, stop)::Int64
+
+The device time in nanoseconds between the timestamps `start` and `stop`, returned by
+[`record_timestamp`](@ref) on the same device. Waits for `stop` to be reached, as
+cooperatively as [`synchronize`](@ref) does. The result may be negative if `stop` was
+reached before `start`, e.g. when they were recorded on different queues.
+
+!!! note
+    Backend implementations **must** implement this function if they implement
+    [`record_timestamp`](@ref).
+"""
+function elapsed_time end

src/KernelAbstractions.jl

@@ -665,6 +665,7 @@ automatically when a kernel is launched.
 argconvert(k::Kernel{T}, arg) where {T} =
     error("Don't know how to convert arguments for Kernel{$T}")
 
+include("profiling.jl")
 include("backend_launch.jl")
 
 # Enzyme support
@@ -743,7 +744,6 @@ number as its compute units: `KernelAbstractions.POCL.device().max_compute_units
 """
 const CPU = POCLBackend
 
-include("profiling.jl")
 include("profiler.jl")
 include("precompile.jl")
 

src/backend_launch.jl

@@ -93,12 +93,22 @@ end
 Base.@nospecializeinfer @noinline function launch_traced(
         @nospecialize(obj::Kernel), @nospecialize(args::Tuple), @nospecialize(ndrange), @nospecialize(workgroupsize)
     )
-    id = start_launch_range(kernel_label(obj.f))
+    label = kernel_label(obj.f)
+    id = start_launch_range(label)
+    # time the kernel on the device, if a tracer wants that
+    timer = records_kernels(id) ? KernelTimer() : nothing
     try
-        launch_untraced(obj, args, ndrange, workgroupsize)
+        if timer === nothing
+            launch_untraced(obj, args, ndrange, workgroupsize)
+        else
+            # passed to `launch_kernel` in a scoped value rather than an argument, so that
+            # the launch path is inferred once for timed and other launches
+            with(() -> launch_untraced(obj, args, ndrange, workgroupsize), KERNEL_TIMER => timer)
+        end
         synchronize_launch(id, backend(obj))
     finally
         profiling_range_end(id)
+        timer === nothing || timer.issued == 0 || trace_kernel(id, label, timer)
     end
     return nothing
 end
@@ -146,11 +156,15 @@ function launch_kernel(obj::Kernel, launch, ndrange, _workgroupsize, iterspace,
     # launching through the `KI.Kernel` validates the sizes against the kernel's limits
     groups = size(blocks(iterspace))
     items = size(workitems(iterspace))
+    # timed around the launch alone, so that compilation isn't counted as device time; the
+    # timer comes from `launch_traced`
+    timer = start_kernel_timing(b)
     if launch isa NDLaunch
         call_kernel(kernel, ctx, args, groups, items)
     else
         call_kernel(kernel, ctx, args, prod(groups), prod(items))
     end
+    stop_kernel_timing(timer, b)
     return nothing
 end