[maleadt]

CUDA.jl gives each Julia task its own stream. That keeps concurrently running tasks independent, but a finished task's stream stays alive until the GC collects the task. Each stream holds on to about half a MiB of device memory, and the stream-ordered allocator gets slower with every stream that has used it, so applications that spawn many short GPU tasks accumulate thousands of streams and slow down allocation for every task.

On an RTX 5080 (driver 610), 20,000 tasks that each compute sum(x .+ 1f0) on a small array took 17–35 s and reached 11,700 live streams (6.5 GiB of device memory). A producer/consumer workload with 20,000 iterations took 10.0 s with up to 15,300 streams and 9.2 GiB of device memory. The same work in a single task takes 0.71 s and 0.3 s respectively.

This PR keeps the streams of tasks in a pool per context, priority and flags, and hands a stream to the next task that needs one once its owner has finished and its GPU work has completed. Tasks running at the same time never share a stream, and up to 32 idle streams are kept around; idle streams beyond that are dropped the next time the pool is searched. With this, the two workloads above take 0.95 s with 2 streams (~350 MiB) and 0.45 s with 8 streams (~450 MiB).

Handing a stream to another task bumps a generation counter on the stream. Memory remembers the generation it was last used in, so after its stream has been recycled it knows that work has finished: using it doesn't wait for the new owner's work, and freeing it happens on the context's disposal stream instead of the recycled one, which the new owner may be capturing.

Because streams are now selected by the pool, this PR also adds task-level stream priorities. CUDA fixes a stream's priority when it is created, while Julia code schedules work in tasks, so priority! selects a stream of the requested priority for the current task and orders it after the task's previous stream:

CUDA.priority!(:high)            # subsequent GPU work in this task
# launch latency-sensitive work
CUDA.priority!(:normal)

CUDA.priority!(:high) do
    # work in this block uses high priority
end                              # restores the previous stream

priority! also accepts an integer from priority_range(), and KernelAbstractions.priority! uses the same mechanism (it used to create a new stream on every call). Selecting a priority before the task first uses the GPU creates only the requested stream, and switching back reuses the task's earlier stream. As in CUDA itself, priority is a scheduling hint for pending work, not preemption. Child tasks start at normal priority, and changing priority during graph capture is rejected. Using an array after a switch may wait on the CPU for work on its previous stream. If the block passed to priority! throws, the previous stream is selected again without ordering it after the block's work, so that the block's error is what gets reported.

The trade-off is that a stream returned by stream() belongs to the current task: code that keeps using it after the task has finished may share it with another task. Code that needs a stream beyond a task's lifetime should create one with CuStream(); the multitasking docs and NEWS say so. Looking for an idle stream scans the pool under a global lock, which is linear in the number of live tasks, but only happens when a task first uses the GPU or switches priority. The 32-stream retention limit is a fixed policy rather than a preference, since it only bounds idle streams.

Two small fixes come along: capture no longer ends the capture in progress when starting a nested one fails, and the docstring of priority(::CuStream) shows the right signature.

Tested on an RTX 5080 with the core/initialization, core/resources, core/cudadrv and core/kernelabstractions tests, including new tests that check that concurrent tasks get distinct streams, that the stream count stays bounded across 1,024 tasks, and that memory last used on a recycled stream is usable and freed without waiting for the stream's new owner.

[github-actions]

CUDA.jl Benchmarks

Details

Benchmark suite	Current: 0bc3bde	Previous: ae637e5	Ratio
array/accumulate/Float32/1d	105361 ns	105001 ns	1.00
array/accumulate/Float32/dims=1	77333 ns	77262 ns	1.00
array/accumulate/Float32/dims=1L	1684889 ns	1684940 ns	1.00
array/accumulate/Float32/dims=2	147917 ns	148043 ns	1.00
array/accumulate/Float32/dims=2L	685614 ns	684338 ns	1.00
array/accumulate/Int64/1d	124101 ns	123503 ns	1.00
array/accumulate/Int64/dims=1	82085 ns	81902 ns	1.00
array/accumulate/Int64/dims=1L	1802775 ns	1803636 ns	1.00
array/accumulate/Int64/dims=2	163401 ns	163105 ns	1.00
array/accumulate/Int64/dims=2L	989125 ns	988275 ns	1.00
array/broadcast	16980 ns	16981 ns	1.00
array/broadcast launch	7627.75 ns	7497.25 ns	1.02
array/construct	901.7826086956521 ns	886.76 ns	1.02
array/copy	16023 ns	16325 ns	0.98
array/copyto!/cpu_to_gpu	208881 ns	206249 ns	1.01
array/copyto!/gpu_to_cpu	241363 ns	240560 ns	1.00
array/copyto!/gpu_to_gpu	9045 ns	8967 ns	1.01
array/iteration/findall/bool	139880 ns	135815 ns	1.03
array/iteration/findall/int	151051 ns	145400 ns	1.04
array/iteration/findfirst/bool	75002 ns	71077 ns	1.06
array/iteration/findfirst/int	77749 ns	73309 ns	1.06
array/iteration/findmin/1d	63437 ns	59845 ns	1.06
array/iteration/findmin/2d	104162 ns	103127 ns	1.01
array/iteration/logical	196738 ns	187405 ns	1.05
array/iteration/scalar	57787 ns	59115 ns	0.98
array/permutedims/2d	49498 ns	49364 ns	1.00
array/permutedims/3d	50767 ns	50795 ns	1.00
array/permutedims/4d	52150 ns	52271 ns	1.00
array/random/rand/Float32	10906 ns	11109 ns	0.98
array/random/rand/Int64	18359 ns	18411 ns	1.00
array/random/rand!/Float32	8165 ns	8183 ns	1.00
array/random/rand!/Int64	16020 ns	15765 ns	1.02
array/random/randn/Float32	34628 ns	34501 ns	1.00
array/random/randn!/Float32	24592 ns	24403 ns	1.01
array/reductions/mapreduce/Float32/1d	32717 ns	30573 ns	1.07
array/reductions/mapreduce/Float32/dims=1	40104 ns	40213 ns	1.00
array/reductions/mapreduce/Float32/dims=1L	53680 ns	53356 ns	1.01
array/reductions/mapreduce/Float32/dims=2	59423 ns	59358 ns	1.00
array/reductions/mapreduce/Float32/dims=2L	71677 ns	71210 ns	1.01
array/reductions/mapreduce/Int64/1d	39961 ns	37469 ns	1.07
array/reductions/mapreduce/Int64/dims=1	43212 ns	42915 ns	1.01
array/reductions/mapreduce/Int64/dims=1L	89306 ns	89249 ns	1.00
array/reductions/mapreduce/Int64/dims=2	61442 ns	61370 ns	1.00
array/reductions/mapreduce/Int64/dims=2L	87745 ns	87543 ns	1.00
array/reductions/reduce/Float32/1d	32733 ns	30615 ns	1.07
array/reductions/reduce/Float32/dims=1	39974 ns	40264 ns	0.99
array/reductions/reduce/Float32/dims=1L	52870 ns	52907 ns	1.00
array/reductions/reduce/Float32/dims=2	59248 ns	59462 ns	1.00
array/reductions/reduce/Float32/dims=2L	71688 ns	71728 ns	1.00
array/reductions/reduce/Int64/1d	40149 ns	37570 ns	1.07
array/reductions/reduce/Int64/dims=1	43162 ns	42862 ns	1.01
array/reductions/reduce/Int64/dims=1L	89265 ns	89319 ns	1.00
array/reductions/reduce/Int64/dims=2	61490 ns	61437 ns	1.00
array/reductions/reduce/Int64/dims=2L	87817 ns	86880 ns	1.01
array/reverse/1d	16141 ns	16128 ns	1.00
array/reverse/1dL	70069 ns	70115 ns	1.00
array/reverse/1dL_inplace	67794 ns	67602 ns	1.00
array/reverse/1d_inplace	9343 ns	9293 ns	1.01
array/reverse/2d	21623 ns	21118 ns	1.02
array/reverse/2dL	74984 ns	74528 ns	1.01
array/reverse/2dL_inplace	67781 ns	67778 ns	1.00
array/reverse/2d_inplace	10810 ns	10676 ns	1.01
array/sorting/1d	2841671 ns	2846981 ns	1.00
array/sorting/2d	1086564 ns	1088944 ns	1.00
array/sorting/by	3393893 ns	3400388 ns	1.00
cuda/synchronization/context/auto	7920 ns	6766.6 ns	1.17
cuda/synchronization/context/blocking	873.0181818181818 ns	793.7692307692307 ns	1.10
cuda/synchronization/context/nonblocking	7856.333333333333 ns	6805.4 ns	1.15
cuda/synchronization/stream/auto	720.2700729927008 ns	691.0939597315436 ns	1.04
cuda/synchronization/stream/blocking	929.4375 ns	868.875 ns	1.07
cuda/synchronization/stream/nonblocking	8333 ns	7338.5 ns	1.14
integration/byval/reference	148185 ns	147882 ns	1.00
integration/byval/slices=1	149102 ns	149033 ns	1.00
integration/byval/slices=2	291152 ns	291277 ns	1.00
integration/byval/slices=3	433919 ns	434160 ns	1.00
integration/cudadevrt	105461 ns	105116 ns	1.00
integration/volumerhs	9205220 ns	9206594 ns	1.00
kernel/indexing	11769 ns	11855 ns	0.99
kernel/indexing_checked	12333 ns	12209 ns	1.01
kernel/launch	2506.5555555555557 ns	2433 ns	1.03
kernel/occupancy	978 ns	1008.8235294117648 ns	0.97
kernel/rand	15074 ns	14899 ns	1.01
latency/import	4314904060 ns	4305717656 ns	1.00
latency/precompile	5082642505 ns	5095420673 ns	1.00
latency/ttfp	4782318460 ns	4793974717 ns	1.00

This comment was automatically generated by workflow using github-action-benchmark.

[codecov]

Codecov Report

❌ Patch coverage is 95.71429% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 84.89%. Comparing base (0151118) to head (0bc3bde).

Files with missing lines	Patch %	Lines
CUDACore/lib/cudadrv/state.jl	95.37%	5 Missing ⚠️
CUDACore/src/resources/streams.jl	75.00%	1 Missing ⚠️

Additional details and impacted files

@@                   Coverage Diff                   @@
##           tb/deferred-release    #3317      +/-   ##
=======================================================
- Coverage                84.93%   84.89%   -0.05%     
=======================================================
  Files                      193      193              
  Lines                    19126    19234     +108     
=======================================================
+ Hits                     16245    16329      +84     
- Misses                    2881     2905      +24     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[maleadt]

While working on this, I noticed AMDGPU.jl has a priority! function to change the priority of the current task. We didn't have that, instead requiring people to manually configure the underlying stream. With streams being recycled, that became a little iffy, so I had Sol port the task priority capability from AMDGPU.jl.