Merge pull request #141 from exanauts/ms/active_kernels

Add support for active_kernels

Author: jpsamaroo

Project.toml

@@ -1,7 +1,7 @@
 name = "AMDGPU"
 uuid = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 authors = ["Julian P Samaroo <jpsamaroo@jpsamaroo.me>"]
-version = "0.2.11"
+version = "0.2.12"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

deps/loaddeps.jl

@@ -10,7 +10,7 @@ try
     include("ext.jl")
 catch err
     if !isfile(joinpath(@__DIR__, "ext.jl"))
-        @warn "Didn't find $ext, please build AMDGPU.jl"
+        @warn "Didn't find deps/ext.jl, please build AMDGPU.jl"
         @eval const hsa_configured = false
         @eval const hip_configured = false
         @eval const device_libs_configured = false

docs/src/queues_signals.md

@@ -12,6 +12,12 @@ q = AMDGPU.HSAQueue(agent)
 @roc queue=q kernel(...)
 ```
 
+If you want to query which kernels are currently executing on a given queue,
+calling `AMDGPU.active_kernels(queue)` will return a `Vector{HSAStatusSignal}`,
+which can be inspected to determine how many (and which) kernels are executing
+by comparing the signals returned from `@roc`. You can also omit the `queue`
+argument, which will then check the default queue.
+
 If a kernel ever gets "stuck" and locks up the GPU (noticeable with 100% GPU
 usage in `rocm-smi`), you can kill the kernel and all other kernels in the
 queue with `kill_queue!(queue)`. This can be "safely" done to the default

src/AMDGPU.jl

@@ -31,6 +31,8 @@ import .HSA: Agent, Queue, Executable, Status, Signal
 
 struct Adaptor end
 
+const RT_LOCK = Threads.SpinLock()
+
 include("extras.jl")
 include("error.jl")
 include("agent.jl")
@@ -48,6 +50,7 @@ if get(ENV, "AMDGPUNATIVE_OPENCL", "") != ""
 end
 =#
 include("runtime.jl")
+include("statussignal.jl")
 include("sync.jl")
 
 # Device sources must load _before_ the compiler infrastructure
@@ -60,6 +63,7 @@ include(joinpath("device", "runtime.jl"))
 include(joinpath("device", "llvm.jl"))
 include(joinpath("device", "globals.jl"))
 
+include("query.jl")
 include("compiler.jl")
 include("execution_utils.jl")
 include("execution.jl")

src/exceptions.jl

@@ -6,7 +6,8 @@ struct KernelException <: Exception
     dev::RuntimeDevice
     exstr::Union{String,Nothing}
 end
-KernelException(dev) = KernelException(dev, nothing)
+KernelException(dev::RuntimeDevice) = KernelException(dev::RuntimeDevice, nothing)
+KernelException(agent::HSAAgent, exstr=nothing) = KernelException(RuntimeDevice(agent), exstr)
 
 function Base.showerror(io::IO, err::KernelException)
     print(io, "KernelException: exception(s) thrown during kernel execution on device $(err.dev.device)")

src/execution.jl

@@ -192,7 +192,7 @@ macro roc(ex...)
                     local $kernel = $rocfunction($f, $kernel_tt; $(compiler_kwargs...))
                     foreach($wait!, ($(var_exprs...),))
                     if $launch
-                        local $signal = $create_event($kernel.mod.exe)
+                        local $signal = $create_event($kernel.mod.exe; $(call_kwargs...))
                         $kernel($kernel_args...; signal=$signal, $(call_kwargs...))
                         foreach(x->$mark!(x, $signal), ($(var_exprs...),))
                         $signal

src/execution_utils.jl

@@ -204,6 +204,12 @@ unpreserve!(ev::RuntimeEvent) = unpreserve!(ev.event)
             kern = create_kernel(get_device(queue), f.mod.exe, f.entry, args)
 
             # launch kernel
+            lock($RT_LOCK)
+            try
+                push!($_active_kernels[queue.queue], signal.event)
+            finally
+                unlock($RT_LOCK)
+            end
             launch_kernel(queue, kern, signal;
                           groupsize=groupsize, gridsize=gridsize)
 

src/query.jl

@@ -0,0 +1,13 @@
+const _active_kernels = IdDict{HSAQueue,Vector{AMDGPU.HSAStatusSignal}}()
+
+"""
+    active_kernels(queue) -> Vector{AMDGPU.HSAStatusSignal}
+
+Returns the set of actively-executing kernels on `queue`.
+"""
+function active_kernels(queue::HSAQueue=get_default_queue())
+    lock(RT_LOCK) do
+        copy(_active_kernels[queue])
+    end
+end
+active_kernels(queue::RuntimeQueue) = active_kernels(queue.queue)

src/queue.jl

@@ -21,11 +21,18 @@ function HSAQueue(agent::HSAAgent)
                      C_NULL, C_NULL, typemax(UInt32), typemax(UInt32),
                      queue.queue) |> check
 
+    lock(RT_LOCK) do
+        _active_kernels[queue] = Vector{AMDGPU.RuntimeEvent{AMDGPU.HSAStatusSignal}}()
+    end
+
     hsaref!()
     finalizer(queue) do queue
         if queue.active
             HSA.queue_destroy(queue.queue[]) |> check
         end
+        lock(RT_LOCK) do
+            delete!(_active_kernels, queue)
+        end
         hsaunref!()
     end
     return queue

src/runtime.jl

@@ -19,63 +19,8 @@ default_isa(device::RuntimeDevice{HSAAgent}) =
 struct RuntimeEvent{E}
     event::E
 end
-create_event(exe) = RuntimeEvent(create_event(RUNTIME[], exe))
-Base.wait(event::RuntimeEvent, exe) = wait(event.event, exe)
-
-"Tracks the completion and status of a kernel's execution."
-struct HSAStatusSignal
-    signal::HSASignal
-    exe::HSAExecutable
-end
-create_event(::typeof(HSA_rt), exe) = HSAStatusSignal(HSASignal(), exe.exe)
-Base.wait(event::HSAStatusSignal; kwargs...) = wait(RuntimeEvent(event); kwargs...)
-function Base.wait(event::RuntimeEvent{HSAStatusSignal}; check_exceptions=true, cleanup=true, kwargs...)
-    wait(event.event.signal; kwargs...) # wait for completion signal
-    unpreserve!(event) # allow kernel-associated objects to be freed
-    exe = event.event.exe::HSAExecutable{Mem.Buffer}
-    mod = EXE_TO_MODULE_MAP[exe].value
-    agent = exe.agent
-    ex = nothing
-    signal_handle = (event.event.signal.signal[]::HSA.Signal).handle
-    if haskey(exe.globals, :__global_exception_flag)
-        # Check if any wavefront for this kernel threw an exception
-        ex_flag = get_global(exe, :__global_exception_flag)
-        ex_flag_ptr = Base.unsafe_convert(Ptr{Int64}, ex_flag)
-        ex_flag_value = Base.unsafe_load(ex_flag_ptr)
-        if ex_flag_value != 0
-            ex_strings = String[]
-            if check_exceptions && haskey(exe.globals, :__global_exception_ring)
-                # Check for and collect any exceptions, and clear their slots
-                ex_ring = get_global(exe, :__global_exception_ring)
-                ex_ring_ptr_ptr = Base.unsafe_convert(Ptr{Ptr{ExceptionEntry}}, ex_ring)
-                ex_ring_ptr = unsafe_load(ex_ring_ptr_ptr)
-                while (ex_ring_value = unsafe_load(ex_ring_ptr)).kern != 1
-                    if ex_ring_value.kern == signal_handle
-                        push!(ex_strings, unsafe_string(reinterpret(Ptr{UInt8}, ex_ring_value.ptr)))
-                        # FIXME: Write rest of entry first, then CAS 0 to kern field
-                        unsafe_store!(ex_ring_ptr, ExceptionEntry(UInt64(0), LLVMPtr{UInt8,1}(0)))
-                    end
-                    ex_ring_ptr += sizeof(ExceptionEntry)
-                end
-            end
-            unique!(ex_strings)
-            ex = KernelException(RuntimeDevice(agent), isempty(ex_strings) ? nothing : join(ex_strings, '\n'))
-        end
-    end
-    if cleanup
-        # Clean-up malloc'd data
-        for idx in length(mod.metadata):-1:1
-            metadata_value = mod.metadata[idx]
-            if metadata_value.kern == signal_handle
-                @debug "Cleaning up data: $metadata_value"
-                Mem.free(metadata_value.buf)
-                deleteat!(mod.metadata, idx)
-            end
-        end
-    end
-    ex !== nothing && throw(ex)
-end
-
+create_event(exe; kwargs...) = RuntimeEvent(create_event(RUNTIME[], exe; kwargs...))
+Base.wait(event::RuntimeEvent) = wait(event.event)
 
 struct RuntimeExecutable{E}
     exe::E
@@ -100,6 +45,9 @@ end
 get_global(exe::RuntimeExecutable, sym::Symbol) =
     get_global(exe.exe, sym)
 
+create_event(::typeof(HSA_rt), exe::RuntimeExecutable{<:HSAExecutable}; queue=default_queue(), kwargs...) =
+    HSAStatusSignal(HSASignal(), exe.exe, queue.queue)
+
 struct RuntimeKernel{K}
     kernel::K
 end
@@ -116,6 +64,4 @@ function launch_kernel(::typeof(HSA_rt), queue, kern, event;
                        workgroup_size=groupsize, grid_size=gridsize)
 end
 barrier_and!(queue, events::Vector{<:RuntimeEvent}) = barrier_and!(queue, map(x->x.event,events))
-barrier_and!(queue, signals::Vector{HSAStatusSignal}) = barrier_and!(queue, map(x->x.signal,signals))
 barrier_or!(queue, events::Vector{<:RuntimeEvent}) = barrier_or!(queue, map(x->x.event,events))
-barrier_or!(queue, signals::Vector{HSAStatusSignal}) = barrier_or!(queue, map(x->x.signal,signals))

src/statussignal.jl

@@ -0,0 +1,73 @@
+"Tracks the completion and status of a kernel's execution."
+mutable struct HSAStatusSignal
+    signal::HSASignal
+    exe::HSAExecutable
+    queue::HSAQueue
+    done::Base.Event
+    exception::Union{Exception,Nothing}
+    function HSAStatusSignal(signal::HSASignal, exe::HSAExecutable, queue::HSAQueue; kwargs...)
+        signal = new(signal, exe, queue, Base.Event(), nothing)
+        @async _wait(signal; kwargs...) # the real waiter
+        signal
+    end
+end
+
+function Base.wait(signal::HSAStatusSignal)
+    wait(signal.done)
+    ex = signal.exception
+    if ex !== nothing
+        throw(ex)
+    end
+end
+function _wait(signal::HSAStatusSignal; check_exceptions=true, cleanup=true, kwargs...)
+    wait(signal.signal; kwargs...) # wait for completion signal
+    unpreserve!(signal) # allow kernel-associated objects to be freed
+    exe = signal.exe::HSAExecutable{Mem.Buffer}
+    mod = EXE_TO_MODULE_MAP[exe].value
+    agent = exe.agent
+    ex = nothing
+    signal_handle = (signal.signal.signal[]::HSA.Signal).handle
+    if haskey(exe.globals, :__global_exception_flag)
+        # Check if any wavefront for this kernel threw an exception
+        ex_flag = get_global(exe, :__global_exception_flag)
+        ex_flag_ptr = Base.unsafe_convert(Ptr{Int64}, ex_flag)
+        ex_flag_value = Base.unsafe_load(ex_flag_ptr)
+        if ex_flag_value != 0
+            ex_strings = String[]
+            if check_exceptions && haskey(exe.globals, :__global_exception_ring)
+                # Check for and collect any exceptions, and clear their slots
+                ex_ring = get_global(exe, :__global_exception_ring)
+                ex_ring_ptr_ptr = Base.unsafe_convert(Ptr{Ptr{ExceptionEntry}}, ex_ring)
+                ex_ring_ptr = unsafe_load(ex_ring_ptr_ptr)
+                while (ex_ring_value = unsafe_load(ex_ring_ptr)).kern != 1
+                    if ex_ring_value.kern == signal_handle
+                        push!(ex_strings, unsafe_string(reinterpret(Ptr{UInt8}, ex_ring_value.ptr)))
+                        # FIXME: Write rest of entry first, then CAS 0 to kern field
+                        unsafe_store!(ex_ring_ptr, ExceptionEntry(UInt64(0), LLVMPtr{UInt8,1}(0)))
+                    end
+                    ex_ring_ptr += sizeof(ExceptionEntry)
+                end
+            end
+            unique!(ex_strings)
+            ex = KernelException(agent, isempty(ex_strings) ? nothing : join(ex_strings, '\n'))
+        end
+    end
+    if cleanup
+        # Clean-up malloc'd data
+        for idx in length(mod.metadata):-1:1
+            metadata_value = mod.metadata[idx]
+            if metadata_value.kern == signal_handle
+                @debug "Cleaning up data: $metadata_value"
+                Mem.free(metadata_value.buf)
+                deleteat!(mod.metadata, idx)
+            end
+        end
+    end
+    signal.exception = ex
+    lock(RT_LOCK) do
+        deleteat!(_active_kernels[signal.queue], findall(x->x==signal, _active_kernels[signal.queue]))
+    end
+    notify(signal.done)
+end
+barrier_and!(queue, signals::Vector{HSAStatusSignal}) = barrier_and!(queue, map(x->x.signal,signals))
+barrier_or!(queue, signals::Vector{HSAStatusSignal}) = barrier_or!(queue, map(x->x.signal,signals))

test/device/exceptions.jl

@@ -12,7 +12,7 @@ catch err
     @test err isa AMDGPU.KernelException
     if err isa AMDGPU.KernelException
         @test err.exstr !== nothing
-        @test occursin("BoundsError", err.exstr)
+        @test occursin("julia_throw_boundserror", err.exstr)
     end
 end
 

test/device/queries.jl

@@ -0,0 +1,18 @@
+@testset "Active kernels" begin
+    function kernel(sig)
+        hostcall!(sig)
+        nothing
+    end
+
+    wait_ev = Base.Event()
+    hc = HostCall(Nothing, Tuple{}) do
+        wait(wait_ev)
+    end
+
+    sig = @roc kernel(hc)
+    @test sig.event in AMDGPU.active_kernels()
+    @test !sig.event.done.set
+    notify(wait_ev)
+    wait(sig)
+    @test !(sig.event in AMDGPU.active_kernels())
+end

test/runtests.jl

@@ -55,6 +55,7 @@ if AMDGPU.configured
             include("device/execution_control.jl")
             include("device/exceptions.jl")
             include("device/deps.jl")
+            include("device/queries.jl")
         end
         @testset "ROCArray" begin
             @testset "GPUArrays test suite" begin