Keep stream and rocFFT plan destructors out of task-local state

[luraess]

Fixes the intermittent hip_rocarray/multi_gpu failure on the CSCS MI300 runners:

Test Failed at test/hip_rocarray/multi_gpu.jl:15
  Expression: id1 == 1
   Evaluated: 2 == 1

followed by d1 != d2, s1 != s2, id1 != id2 and d1 == AMDGPU.device(). Seen on CI.

Cause

The HIPStream finalizer switched context with AMDGPU.context!, the task-local variant. Finalizers run on whichever task triggers GC. If that task has no AMDGPU state yet, context! creates it on the stream's device and returns nothing as the previous context, so the finally never switches back. The task stays on the stream's device for good.

ParallelTestRunner reuses workers and runs each file in a fresh task that starts with GC.gc(true), so this carries from one file to the next:

1. device/launch switches to device 2 and launches, which leaves a device-2 stream in its task-local state. It restores device 1 correctly.
2. The next file on that worker starts in a fresh task. Its first GC runs that stream's finalizer, which pins the task to device 2.
3. That file does its GPU work on device 2 and leaves its own device-2 stream behind, which pins the next file the same way.

This continues until multi_gpu lands on the worker and finds itself on device 2 at line 15. In the last 237 CSCS jobs, all 3 line-15 failures had device/launch earlier on the same worker. Of the 176 passing runs, only one did, and in that one gpuarrays/sparse (which likely creates no stream) broke the chain.

Fix

Use HIP.context! in the two destructors that can run from a finalizer and still went through AMDGPU.context!:

• the HIPStream finalizer (src/hip/stream.jl)
• the rocFFT plan destroy closure, which runs on cache eviction from the plan finalizer (src/fft/wrappers.jl)

HIP.context! saves and restores the thread's current HIP context and never touches task-local state, which is what a finalizer needs. A finalizer can't yield, so the save and restore happen on the same OS thread. Both bodies are plain C calls that read no task-local state.

AMDGPU.context! came in with #552 for its ctx.valid check, back when HIPContext had a finalizer that invalidated it. #605 removed that finalizer, so the check can no longer fail. e321d7d already made this switch for pool frees, and #1096 makes it for library_state.

Related

• rocFFT: capture plan context at creation time, not finalization #1092 fixes the AMDGPU.context() read in release_plan!, which also runs from the plan finalizer. This PR leaves that line alone, and the two merge cleanly.
• library_state: key handles on the creation context, and call the destructor that was passed #1096 fixes the library_state destructor. Also merges cleanly.
• Follow-up for src/tls.jl itself, plus an audit for other sites: Task-local state follow-ups: prepare_state in finalizers, context! restore, device! global default #1105

[github-actions]

AMDGPU.jl Benchmarks

Details

Benchmark suite	Current: de6ecfe	Previous: 8a66547	Ratio
amdgpu/synchronization/context/device	550 ns	550 ns	1
amdgpu/synchronization/stream/blocking	225 ns	222.5 ns	1.01
amdgpu/synchronization/stream/nonblocking	307.5 ns	307.5 ns	1
applications/bitonic_sort	1160559.5 ns	1160103.75 ns	1.00
applications/convolution	72606.25 ns	72376 ns	1.00
applications/floyd_warshall	9107993.75 ns	9100856.25 ns	1.00
applications/histogram	818982 ns	798998.75 ns	1.03
applications/prefix_sum	188765.25 ns	190535 ns	0.99
array/accumulate/Float32/1d	69846 ns	79346 ns	0.88
array/accumulate/Float32/dims=1	272979 ns	271536.25 ns	1.01
array/accumulate/Float32/dims=1L	81821.25 ns	72536 ns	1.13
array/accumulate/Float32/dims=2	74911 ns	74041 ns	1.01
array/accumulate/Float32/dims=2L	3023800.25 ns	2834326.75 ns	1.07
array/accumulate/Int64/1d	78746.25 ns	79098.5 ns	1.00
array/accumulate/Int64/dims=1	248038.75 ns	248393.25 ns	1.00
array/accumulate/Int64/dims=1L	88898.75 ns	89128.75 ns	1.00
array/accumulate/Int64/dims=2	84888.75 ns	76041 ns	1.12
array/accumulate/Int64/dims=2L	3389240.5 ns	3139866 ns	1.08
array/broadcast	54420.75 ns	50443.25 ns	1.08
array/construct	2147.5 ns	2247.5 ns	0.96
array/copy	40625.5 ns	39278 ns	1.03
array/copyto!/cpu_to_gpu	86068.75 ns	85676.25 ns	1.00
array/copyto!/gpu_to_cpu	85558.75 ns	85591.25 ns	1.00
array/copyto!/gpu_to_gpu	58565.75 ns	58288.25 ns	1.00
array/iteration/findall/bool	129164.5 ns	125954.25 ns	1.03
array/iteration/findall/int	142407 ns	146617 ns	0.97
array/iteration/findfirst/bool	155652.25 ns	155292.25 ns	1.00
array/iteration/findfirst/int	191105.25 ns	189260 ns	1.01
array/iteration/findmin/1d	104296.5 ns	103258.75 ns	1.01
array/iteration/findmin/2d	89819 ns	90081.25 ns	1.00
array/iteration/logical	226928.5 ns	226565.5 ns	1.00
array/iteration/scalar	303434.5 ns	295931.5 ns	1.03
array/permutedims/2d	71163.75 ns	70736 ns	1.01
array/permutedims/3d	69243.5 ns	69253.5 ns	1.00
array/permutedims/4d	75503.5 ns	50278.25 ns	1.50
array/random/rand/Float32	46035.75 ns	40065.75 ns	1.15
array/random/rand/Int64	55463.25 ns	48863 ns	1.14
array/random/rand!/Float32	44345.5 ns	65541 ns	0.68
array/random/rand!/Int64	73833.5 ns	69623.25 ns	1.06
array/random/randn/Float32	68468.75 ns	71161 ns	0.96
array/random/randn!/Float32	82438.75 ns	58721 ns	1.40
array/reductions/mapreduce/Float32/1d	81983.75 ns	81621 ns	1.00
array/reductions/mapreduce/Float32/dims=1	73201.25 ns	56660.75 ns	1.29
array/reductions/mapreduce/Float32/dims=1L	835414.75 ns	844514.25 ns	0.99
array/reductions/mapreduce/Float32/dims=2	74458.75 ns	74073.5 ns	1.01
array/reductions/mapreduce/Float32/dims=2L	137144.5 ns	136524.5 ns	1.00
array/reductions/mapreduce/Int64/1d	81958.75 ns	81663.5 ns	1.00
array/reductions/mapreduce/Int64/dims=1	73968.75 ns	73933.5 ns	1.00
array/reductions/mapreduce/Int64/dims=1L	835562.25 ns	839174.25 ns	1.00
array/reductions/mapreduce/Int64/dims=2	74486 ns	75486.25 ns	0.99
array/reductions/mapreduce/Int64/dims=2L	138224.5 ns	137611.75 ns	1.00
array/reductions/reduce/Float32/1d	82116.25 ns	81541.25 ns	1.01
array/reductions/reduce/Float32/dims=1	72806.25 ns	70296 ns	1.04
array/reductions/reduce/Float32/dims=1L	838372.5 ns	836984.25 ns	1.00
array/reductions/reduce/Float32/dims=2	73791 ns	75123.5 ns	0.98
array/reductions/reduce/Float32/dims=2L	136502 ns	136774.25 ns	1.00
array/reductions/reduce/Int64/1d	81778.5 ns	81368.75 ns	1.01
array/reductions/reduce/Int64/dims=1	73628.5 ns	74831 ns	0.98
array/reductions/reduce/Int64/dims=1L	837297.25 ns	844554 ns	0.99
array/reductions/reduce/Int64/dims=2	74928.75 ns	75411.25 ns	0.99
array/reductions/reduce/Int64/dims=2L	137529.75 ns	137047 ns	1.00
array/reverse/1d	44028.25 ns	43458 ns	1.01
array/reverse/1dL	68703.75 ns	70606 ns	0.97
array/reverse/1dL_inplace	56671 ns	79671 ns	0.71
array/reverse/1d_inplace	38233 ns	60390.75 ns	0.63
array/reverse/2d	47975.5 ns	47625.75 ns	1.01
array/reverse/2dL	74533.5 ns	57993.5 ns	1.29
array/reverse/2dL_inplace	70986.25 ns	79096 ns	0.90
array/reverse/2d_inplace	52488.25 ns	61191 ns	0.86
array/sorting/1d	327442.5 ns	326449.5 ns	1.00
gemm/tiled	1982016.5 ns	1961659.75 ns	1.01
gemm/tiled_unbounded	1943141 ns	1960334.75 ns	0.99
integration/byval/reference	39990 ns	39821 ns	1.00
integration/byval/slices=1	40791 ns	40830 ns	1.00
integration/byval/slices=2	159433 ns	130842 ns	1.22
integration/byval/slices=3	241634 ns	241884 ns	1.00
integration/volumerhs	4891712 ns	4900148 ns	1.00
kernel/indexing	30742.75 ns	41880.5 ns	0.73
kernel/indexing_checked	49660.75 ns	37075.5 ns	1.34
kernel/launch	1162.5 ns	1130 ns	1.03
kernel/rand	69038.5 ns	70343.5 ns	0.98
latency/import	1451232106 ns	1456646355 ns	1.00
latency/precompile	22683230230 ns	22642954217 ns	1.00
latency/ttfp	2240398517 ns	2251455325 ns	1.00
stencil/diffusion3d	1614636.25 ns	1611230 ns	1.00
stencil/diffusion3d_checked	1655797 ns	1655470.5 ns	1.00

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