Flush kernel output when synchronizing a stream that polling finds done

CUDACore/lib/cudadrv/synchronization.jl

@@ -43,7 +43,7 @@ function synchronize_object(obj::SyncObject; blocking::Bool, spin::Bool)
     # synchronization doesn't poll, as a query costs about as much as synchronizing an idle
     # stream.
     res = if nonblocking && pollable && spin && poll(obj)
-        obj isa CuEvent ? synchronize_completed(obj) : SUCCESS
+        synchronize_completed(obj)
     else
         # if we're about to wait, now may be a good time for a GC pause
         maybe_collect(true)
@@ -56,7 +56,7 @@ function synchronize_object(obj::SyncObject; blocking::Bool, spin::Bool)
             ctx = obj isa CuContext ? obj : obj.ctx === nothing ? context() : obj.ctx
             res = cooperative_wait(worker_synchronize, (obj, ctx);
                                    isdone = pollable ? worker_isdone : nothing, spin)
-            if res === nothing && obj isa CuEvent
+            if res === nothing && pollable
                 synchronize_completed(obj)
             else
                 something(res, SUCCESS)::CUresult
@@ -84,6 +84,29 @@ end
 synchronize_completed(event::CuEvent) =
     relaxed_capture_mode(() -> unchecked_synchronize(event))::CUresult
 
+# polling a stream doesn't flush the output of its kernels (from `printf`), as synchronizing
+# does. synchronizing it anyway would wait for work submitted since, so flush otherwise.
+synchronize_completed(stream::CuStream) =
+    flush_output(stream.ctx === nothing ? context() : stream.ctx)
+
+# synchronizing an event that was never recorded flushes kernel output without waiting for
+# any work, so each context keeps one around that is never recorded. that's also allowed
+# while another thread is capturing, once relaxing the capture mode.
+function flush_output(ctx::CuContext)
+    event = @memoize key=ctx::CuContext begin
+        context!(ctx) do
+            handle = Ref{CUevent}()
+            cuEventCreate(handle, EVENT_DISABLE_TIMING)
+            handle[]
+        end
+    end::CUevent
+    if active_captures[] > 0
+        relaxed_capture_mode(() -> unchecked_cuEventSynchronize(event))::CUresult
+    else
+        unchecked_cuEventSynchronize(event)
+    end
+end
+
 function device_synchronize(; blocking::Bool=false, spin::Bool=true)
     synchronize_object(context(); blocking, spin)
     check_exceptions()

CUDACore/src/compiler/exceptions.jl

@@ -33,14 +33,9 @@ function check_exceptions()
             # restore the structure
             unsafe_store!(exception_info, ExceptionInfo_st())
 
-            # flush the output of the kernel, which querying for completion doesn't do (as
-            # non-blocking synchronization does). synchronizing an event that was never
-            # recorded doesn't wait for anything, also not for unrelated work. that's also
-            # allowed while capturing, once relaxing the capture mode.
-            context!(ctx) do
-                event = CuEvent(EVENT_DISABLE_TIMING)
-                relaxed_capture_mode(() -> cuEventSynchronize(event))
-            end
+            # flush the kernel's diagnostic output before reporting the exception
+            res = flush_output(ctx)
+            res == SUCCESS || throw_api_error(res)
 
             # throw host-side
             dev = device(ctx)

NEWS.md

@@ -83,6 +83,10 @@ are listed as subsections of the minor release they belong to.
   words.
 - The output of a kernel that throws an exception is flushed before reporting
   the exception.
+- `synchronize()` flushes the output of kernels (from `@cuprintf` and friends)
+  again. Since v6.4.2, when the stream had already finished its work, which is
+  common for short kernels, the output could only appear after, e.g., a later
+  kernel launch or blocking synchronization.
 - Memory that failed to be pinned, e.g. because it was registered already
   using `CUDA.register`, or that has been unpinned, isn't recorded as pinned
   anymore. Re-pinning a resized array doesn't lead to unpinning it twice.

test/core/device/intrinsics/output.jl

@@ -154,3 +154,21 @@ end
     @test occursin("pointer(b) = ", out)
     @test occursin("= 0", out)  # 0x... on Linux, 0... on Windows
 end
+
+@testset "flushed by synchronize" begin
+    # `@grab_output` often finds the stream busy, synchronizing it through the driver.
+    # make sure the stream is done, so that polling does.
+    # (compile ahead of time, as loading the module flushes too)
+    kernel() = (@cuprintln("Hello, World"); nothing)
+    @cuda launch=false kernel()
+    out = mktemp() do fname, fout
+        redirect_stdout(fout) do
+            @cuda kernel()
+            while !CUDA.isdone(stream()) end
+            synchronize()
+        end
+        close(fout)
+        read(fname, String)
+    end
+    @test out == "Hello, World$endline"
+end

test/helpers.jl

@@ -21,8 +21,8 @@ macro grab_output(ex)
                 redirect_stdout(fout) do
                     ret = $(esc(ex))
 
-                    # NOTE: CUDA requires a 'proper' sync to flush its printf buffer
-                    synchronize(context())
+                    # (synchronizing also flushes the output of kernels)
+                    synchronize()
                 end
             end
             ret, read(fname, String)