Add hostcalls for device-to-host services

CUDACore/lib/cudadrv/context.jl

@@ -95,6 +95,8 @@ respect any users of the context, and might make other objects unusable.
 """
 function unsafe_destroy!(ctx::CuContext)
     if isvalid(ctx)
+        hostcall_forget!(ctx)
+        forget_exceptions!(ctx)
         cuCtxDestroy_v2(ctx)
     end
 end
@@ -226,6 +228,8 @@ in the current process. Note that this forcibly invalidates all contexts derived
 primary context, and as a result outstanding resources might become invalid.
 """
 function unsafe_reset!(pctx::CuPrimaryContext)
+    hostcall_forget!(pctx.dev)
+    forget_exceptions!(pctx.dev)
     cuDevicePrimaryCtxReset_v2(pctx.dev)
     return
 end

CUDACore/lib/cudadrv/module.jl

@@ -13,7 +13,10 @@ function checked_cuModuleLoadDataEx(_module, image, numOptions, options, optionV
     #      available, but cached by the allocator. by configuring the allocator with a
     #      release threshold, we have it actually free up that memory, but that requires
     #      synchronizing all streams to make sure pending frees are actually executed.
-    if !is_capturing()
+    #
+    #      inside a hostcall handler we must not synchronize the device: it would wait for
+    #      the kernel that is waiting for the handler.
+    if !is_capturing() && !hostcall_in_handler()
         device_synchronize()
     end
 

CUDACore/lib/cudadrv/synchronization.jl

@@ -183,7 +183,7 @@ function nonblocking_synchronize(val)
 end
 
 function device_synchronize(; blocking::Bool=false, spin::Bool=true)
-    if use_nonblocking_synchronization && !blocking
+    if use_nonblocking_synchronization && !blocking && !hostcall_in_handler()
         if spin && spinning_synchronization(isdone, legacy_stream())
             cuCtxSynchronize()
         else
@@ -199,7 +199,10 @@ function device_synchronize(; blocking::Bool=false, spin::Bool=true)
 end
 
 function synchronize(stream::CuStream=stream(); blocking::Bool=false, spin::Bool=true)
-    if use_nonblocking_synchronization && !blocking
+    # the special streams (default, legacy, per-thread) belong to the current context
+    ctx = something(stream.ctx, context())
+    # hostcall handlers run on a foreign thread that must not wait on the Julia scheduler
+    if use_nonblocking_synchronization && !blocking && !hostcall_in_handler()
         if spin && spinning_synchronization(isdone, stream)
             cuStreamSynchronize(stream)
         else
@@ -211,11 +214,11 @@ function synchronize(stream::CuStream=stream(); blocking::Bool=false, spin::Bool
         cuStreamSynchronize(stream)
     end
 
-    check_exceptions()
+    check_exceptions(ctx)
 end
 
 function synchronize(event::CuEvent; blocking::Bool=false, spin::Bool=true)
-    if use_nonblocking_synchronization && !blocking
+    if use_nonblocking_synchronization && !blocking && !hostcall_in_handler()
         if spin && spinning_synchronization(isdone, event)
             cuEventSynchronize(event)
         else
@@ -226,4 +229,6 @@ function synchronize(event::CuEvent; blocking::Bool=false, spin::Bool=true)
         maybe_collect(true)
         cuEventSynchronize(event)
     end
+
+    check_exceptions(event.ctx)
 end

CUDACore/src/CUDACore.jl

@@ -105,6 +105,7 @@ include("refpointer.jl")
 include("compiler/compilation.jl")
 include("compiler/execution.jl")
 include("compiler/exceptions.jl")
+include("hostcall.jl")
 
 # array implementation
 include("utilities.jl")

CUDACore/src/compiler/compilation.jl

@@ -201,10 +201,19 @@ mutable struct CUDACompilerResults
     entry::Union{Nothing,String}
     relocations::GPUCompiler.Relocations
 
+    # whether the kernel calls host functions, and the key types of the statically-known
+    # targets of those calls, as recovered from the compiled method instances. the wire
+    # identifiers are not stored: they are the addresses of the egal-rooted target markers,
+    # recomputed at registration in each session (the image's relocated marker slots
+    # resolve to the same addresses)
+    hostcall::Bool
+    hostcall_targets::Vector{Type}
+
     # session-local kernel handles, linear-scanned by context; usually holds a single entry
     kernels::Vector{Tuple{CuContext,CuFunction}}
 
     CUDACompilerResults() = new(nothing, nothing, GPUCompiler.Relocations(),
+                                false, Type[],
                                 Tuple{CuContext,CuFunction}[])
 end
 
@@ -379,6 +388,17 @@ device_compatible_layout(@nospecialize(T)) =
     Base.datatype_alignment(Int128) == 16 ||
     !layout_reaches(S -> device_layout(S) === :mismatch, T)
 
+# Recover the value of a type-valued dispatch key. Julia 1.14 uses `Core.TypeEgal{T}`
+# for closed type arguments, while earlier releases represent them as `Type{T}`.
+@inline function hostcall_key_type(@nospecialize(T))
+    Base.isType(T) || return nothing
+    @static if isdefined(Base, :type_parameter)
+        return Base.type_parameter(T)
+    else
+        return T.parameters[1]
+    end
+end
+
 # compile to executable machine code
 function compile(@nospecialize(job::CompilerJob))
     # lower to PTX
@@ -387,6 +407,20 @@ function compile(@nospecialize(job::CompilerJob))
         invoke_frozen(GPUCompiler.compile, :asm, job)
     end
 
+    # recover the hostcall targets: the device-side `hostcall_impl` is specialized on the
+    # key type of every statically-known call, and `meta.compiled` lists everything codegen
+    # emitted (including deferred compilation jobs)
+    hostcall = false
+    hostcall_targets = Type[]
+    for mi in keys(meta.compiled)
+        mi.def isa Method || continue
+        mi.def.module === CUDACore && mi.def.name === :hostcall_impl || continue
+        hostcall = true
+        K = hostcall_key_type(mi.specTypes.parameters[2])
+        K === nothing && continue
+        push!(hostcall_targets, K)
+    end
+
     # check if we'll need the device runtime
     undefined_fs = filter(collect(functions(meta.ir))) do f
         isdeclaration(f) && !LLVM.isintrinsic(f) &&
@@ -544,7 +578,8 @@ function compile(@nospecialize(job::CompilerJob))
         rm(ptxas_output)
     end
 
-    return (image, entry=LLVM.name(meta.entry), relocations=meta.relocations)
+    return (image, entry=LLVM.name(meta.entry), relocations=meta.relocations,
+            hostcall, hostcall_targets)
 end
 
 # link a compiled image into a session-local `CuFunction` on the active context
@@ -553,7 +588,10 @@ function link_kernel(image::Vector{UInt8}, entry::String,
     # load as an executable kernel object on the current context
     mod = try
         CuModule(image)
-    catch
+    catch err
+        # loading synchronizes the device first, which may surface an unrelated exception
+        # from an earlier kernel; only driver errors are about our image
+        err isa CuError || rethrow()
         # the driver rejected our compiled image (e.g. ERROR_NOT_SUPPORTED). dump the cubin
         # so the failure can be reported with a reproducer, mirroring how we keep the PTX
         # around when `ptxas` fails above.
@@ -597,6 +635,8 @@ function compile_or_lookup(@nospecialize(job::CompilerJob))::CUDACompilerResults
         res.image = compiled.image
         res.entry = compiled.entry
         res.relocations = compiled.relocations
+        res.hostcall = compiled.hostcall
+        res.hostcall_targets = compiled.hostcall_targets
     end
     return res
 end