Make unsafe_wrap of host memory safe, and support wrapping Arrays

[maleadt]

unsafe_wrap(ROCArray, ...) lets the GPU work directly on host memory, without copying it. This PR makes that safe to use, and adds the Array convenience method that CUDA.jl, Metal.jl and OpenCL.jl already have:

a = rand(Float32, 1024)
b = unsafe_wrap(ROCArray, a)   # no copy: the GPU accesses `a` directly
b .+= 1
AMDGPU.synchronize()
a                              # updated in place

The wrapper keeps a alive for as long as it's used, so wrapping a temporary array is fine. You can still wrap a raw pointer, in which case keeping the memory alive is up to you. The other direction, unsafe_wrap(Array, ::ROCArray), now works for arrays backed by host memory. Its docstring warns that the returned Array doesn't keep the ROCArray alive.

Why

Wrapping host memory page-locks it with hipHostRegister, but with the default own=false it was never unregistered:

x = zeros(Float32, 16)
xd = unsafe_wrap(ROCArray, pointer(x), size(x))
AMDGPU.unsafe_free!(xd)
AMDGPU.Mem.is_pinned(pointer(x))   # main: true, this PR: false

Besides leaking pinned memory, this meant a later allocation at the same address silently reused the stale registration. Wrapping the same pointer again at a larger size then crashed:

Memory access fault by GPU node-1 (Agent handle: 0x35c14a70) on address 0x7f94d0a01000. Reason: Page not present or supervisor privilege.

How

The registration is undone when the wrapper is freed, but only once outstanding GPU work on it has finished. Finalizers can't wait for the GPU: they can't switch tasks, and blocking the thread could deadlock with kernels doing host calls. So freeing a wrapper launches a host function on the stream that last used it. When the GPU reaches it, that host function wakes a Julia task, which unregisters the memory and lets go of the wrapped Array. If anything prevents the GPU from getting there, the memory stays registered and the Array stays alive, rather than being released while it may still be in use.

When you free a wrapper explicitly with unsafe_free!, the call waits for the release, so the memory can be wrapped again right away. Some cases need special handling:

• If the stream is being captured into a graph, the release waits until the capture is over, so the host function doesn't end up in the graph.
• Freeing a wrapper while another stream is being captured no longer invalidates that capture.
• Wrapping and freeing the same memory from several threads at once no longer races. HIP can even segfault when hipPointerGetAttributes and hipHostUnregister run concurrently on the same pointer, so those calls are now serialized.
• If the wrapper's stream has been handed to another task since (see Recycle the streams of finished tasks #1120, which this PR builds on), the work on it has finished, so the memory is released right away.

Testing

@gbaraldi tested an earlier version on MI300A and MI250 (thanks!), and the problems found there are fixed and covered by tests. The new tests in test/core/rocarray_base.jl fail on that version and pass now. The full test suite passes on a gfx1036 iGPU (ROCm 7.2.4). The only failure is an Int128 axpby! test marked broken that unexpectedly passes, which also happens on the commit this stack is based on.

[gbaraldi]

Tested on MI300A + MI250 (ROCm 7.2.4):

1. CI failure is the test. Inside the @testset, wrap_tracked's a = … and xd = … assign the testset's a/xd, so neither ever gets collected. local a, xd fixes it (1.12 and 1.13).
2. Freeing during a capture. unsafe_free!(xd) inside capture() do … end, where xd was last used on another stream: HIP.isdone calls hipStreamQuery, which returns 900 during a global-mode capture and invalidates it. The wrapper leaks and the stream stays broken (next sync 900, next launch 901). main is fine. Use relaxed capture mode for the release calls, or defer them while capturing.
3. Register/unregister race. Mem.unregister calls hipHostUnregister after dropping __pin_lock, so a concurrent register of the same pointer sees Host memory and doesn't track it. 2 threads × 3000 wrap/free of one Array: hipErrorInvalidValue in get_device_ptr in 5 of 6 runs. Pre-existing for own=true, but default wraps with async releases now hit it. Make the HIP calls under the lock.
4. Larger re-wrap. Re-wrapping a pointer at a larger size before its pending release has run reuses the old registration (the count > 0 path ignores the size) → hipErrorIllegalAddress. Same on main, but an explicit unsafe_free! could wait for the release (cf. CUDA.jl#3308).
5. With #1120, managed.stream may belong to another task by release time; release immediately if it was recycled.

[maleadt]

Addressed review comments, and put in a stack to handle the stream recycling.

[github-actions]

AMDGPU.jl Benchmarks

Details

Benchmark suite	Current: c5843a1	Previous: 130f06f	Ratio
amdgpu/synchronization/context/device	550 ns	550 ns	1
amdgpu/synchronization/stream/blocking	232.5 ns	225 ns	1.03
amdgpu/synchronization/stream/nonblocking	315 ns	307.5 ns	1.02
applications/bitonic_sort	1117586 ns	1119448.5 ns	1.00
applications/convolution	91151.25 ns	91583.75 ns	1.00
applications/floyd_warshall	9095680 ns	9109628.25 ns	1.00
applications/histogram	774053.75 ns	818519.5 ns	0.95
applications/prefix_sum	185982.75 ns	187865 ns	0.99
array/accumulate/Float32/1d	82771.25 ns	64661 ns	1.28
array/accumulate/Float32/dims=1	257873.5 ns	258818.75 ns	1.00
array/accumulate/Float32/dims=1L	71516 ns	103104 ns	0.69
array/accumulate/Float32/dims=2	103231.5 ns	104694 ns	0.99
array/accumulate/Float32/dims=2L	3032683.25 ns	3033111 ns	1.00
array/accumulate/Int64/1d	87916.25 ns	91849 ns	0.96
array/accumulate/Int64/dims=1	244438.5 ns	244421 ns	1.00
array/accumulate/Int64/dims=1L	111784.25 ns	112366.5 ns	0.99
array/accumulate/Int64/dims=2	103186.5 ns	107364 ns	0.96
array/accumulate/Int64/dims=2L	3391645.75 ns	3391723.75 ns	1.00
array/broadcast	29260.5 ns	29182.75 ns	1.00
array/construct	2417.5 ns	2227.5 ns	1.09
array/copy	41438.25 ns	40513 ns	1.02
array/copyto!/cpu_to_gpu	89163.75 ns	89373.75 ns	1.00
array/copyto!/gpu_to_cpu	89341.25 ns	89406.5 ns	1.00
array/copyto!/gpu_to_gpu	33670.5 ns	33853 ns	0.99
array/iteration/findall/bool	119841.75 ns	113754 ns	1.05
array/iteration/findall/int	112831.5 ns	111856.5 ns	1.01
array/iteration/findfirst/bool	150242.25 ns	152492.25 ns	0.99
array/iteration/findfirst/int	136404.25 ns	163867.25 ns	0.83
array/iteration/findmin/1d	104419 ns	106526.5 ns	0.98
array/iteration/findmin/2d	91206.25 ns	94613.75 ns	0.96
array/iteration/logical	193660.5 ns	183560.25 ns	1.06
array/iteration/scalar	309254.5 ns	297834.25 ns	1.04
array/permutedims/2d	61968.5 ns	58705.75 ns	1.06
array/permutedims/3d	57718.25 ns	56185.75 ns	1.03
array/permutedims/4d	65928.5 ns	65440.75 ns	1.01
array/random/rand/Float32	51735.75 ns	42163 ns	1.23
array/random/rand/Int64	37658.25 ns	49873.25 ns	0.76
array/random/rand!/Float32	41920.75 ns	33603 ns	1.25
array/random/rand!/Int64	38310.5 ns	37550.5 ns	1.02
array/random/randn/Float32	71836 ns	71281 ns	1.01
array/random/randn!/Float32	54053.25 ns	53618.25 ns	1.01
array/reductions/mapreduce/Float32/1d	108001.75 ns	109219 ns	0.99
array/reductions/mapreduce/Float32/dims=1	69691 ns	87476.25 ns	0.80
array/reductions/mapreduce/Float32/dims=1L	833619.25 ns	832684.5 ns	1.00
array/reductions/mapreduce/Float32/dims=2	94141.5 ns	93696.25 ns	1.00
array/reductions/mapreduce/Float32/dims=2L	135284.5 ns	135931.75 ns	1.00
array/reductions/mapreduce/Int64/1d	108506.75 ns	109336.5 ns	0.99
array/reductions/mapreduce/Int64/dims=1	88708.75 ns	87891.25 ns	1.01
array/reductions/mapreduce/Int64/dims=1L	838427 ns	838782 ns	1.00
array/reductions/mapreduce/Int64/dims=2	92963.75 ns	92433.75 ns	1.01
array/reductions/mapreduce/Int64/dims=2L	135924.5 ns	136807 ns	0.99
array/reductions/reduce/Float32/1d	108461.5 ns	84738.75 ns	1.28
array/reductions/reduce/Float32/dims=1	88063.75 ns	87951.25 ns	1.00
array/reductions/reduce/Float32/dims=1L	833584.5 ns	843852.25 ns	0.99
array/reductions/reduce/Float32/dims=2	94226.5 ns	93908.75 ns	1.00
array/reductions/reduce/Float32/dims=2L	135729.5 ns	135749.5 ns	1.00
array/reductions/reduce/Int64/1d	107906.5 ns	109031.75 ns	0.99
array/reductions/reduce/Int64/dims=1	88743.75 ns	87966.25 ns	1.01
array/reductions/reduce/Int64/dims=1L	840439.5 ns	835229.5 ns	1.01
array/reductions/reduce/Int64/dims=2	92761.25 ns	92278.75 ns	1.01
array/reductions/reduce/Int64/dims=2L	136479.5 ns	136729.5 ns	1.00
array/reverse/1d	43533.25 ns	40120.5 ns	1.09
array/reverse/1dL	68886 ns	67163.5 ns	1.03
array/reverse/1dL_inplace	54130.75 ns	53348.25 ns	1.01
array/reverse/1d_inplace	37763 ns	37115.5 ns	1.02
array/reverse/2d	46968 ns	46145.75 ns	1.02
array/reverse/2dL	90796.25 ns	90801.25 ns	1.00
array/reverse/2dL_inplace	64923.5 ns	48145.5 ns	1.35
array/reverse/2d_inplace	37193 ns	36485.75 ns	1.02
array/sorting/1d	333732.25 ns	322017 ns	1.04
gemm/tiled	1949770.25 ns	1955698.25 ns	1.00
gemm/tiled_unbounded	1985073.5 ns	1945963 ns	1.02
integration/byval/reference	39070 ns	39630 ns	0.99
integration/byval/slices=1	41091 ns	39681 ns	1.04
integration/byval/slices=2	156892 ns	156123 ns	1.00
integration/byval/slices=3	239384 ns	239254 ns	1.00
integration/volumerhs	4881199 ns	4910710 ns	0.99
kernel/indexing	28943 ns	28272.75 ns	1.02
kernel/indexing_checked	34188 ns	36060.75 ns	0.95
kernel/launch	1200 ns	1240 ns	0.97
kernel/rand	51015.75 ns	50543.25 ns	1.01
latency/import	1542452108 ns	1536649245 ns	1.00
latency/precompile	28599596712 ns	28276810633 ns	1.01
latency/ttfp	2123479767 ns	2328295960 ns	0.91
stencil/diffusion3d	1581127.75 ns	1587540.5 ns	1.00
stencil/diffusion3d_checked	1632146 ns	1618798 ns	1.01

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