diff --git a/lib/level-zero/memory.jl b/lib/level-zero/memory.jl index 66cb3fa9..6febca3a 100644 --- a/lib/level-zero/memory.jl +++ b/lib/level-zero/memory.jl @@ -104,6 +104,36 @@ function host_alloc(ctx::ZeContext, bytesize::Integer, alignment::Integer=1; return HostBuffer(ptr_ref[], bytesize, ctx) end +# from the ZE_extension_external_memmap_sysmem extension, which our headers predate +struct ze_external_memmap_sysmem_ext_desc_t + stype::ze_structure_type_t + pNext::Ptr{Cvoid} + pSystemMemory::Ptr{Cvoid} + size::UInt64 +end +const ZE_STRUCTURE_TYPE_EXTERNAL_MEMMAP_SYSMEM_EXT_DESC = + reinterpret(ze_structure_type_t, UInt32(0x00020037)) + +""" + host_memmap(ctx, ptr, bytesize) -> HostBuffer + +Create a host allocation backed by existing system memory, using the +`ZE_extension_external_memmap_sysmem` extension. Both `ptr` and `bytesize` must be +page-aligned. Freeing the resulting buffer removes the mapping, but does not free the +memory itself. +""" +function host_memmap(ctx::ZeContext, ptr::Ptr, bytesize::Integer) + sysmem_ref = Ref(ze_external_memmap_sysmem_ext_desc_t( + ZE_STRUCTURE_TYPE_EXTERNAL_MEMMAP_SYSMEM_EXT_DESC, C_NULL, ptr, bytesize)) + ptr_ref = Ref{Ptr{Cvoid}}() + GC.@preserve sysmem_ref begin + pNext = Base.unsafe_convert(Ptr{Cvoid}, sysmem_ref) + desc_ref = Ref(ze_host_mem_alloc_desc_t(; pNext)) + zeMemAllocHost(ctx, desc_ref, bytesize, 1, ptr_ref) + end + return HostBuffer(ptr_ref[], bytesize, ctx) +end + Base.pointer(buf::HostBuffer) = buf.ptr Base.sizeof(buf::HostBuffer) = buf.bytesize context(buf::HostBuffer) = buf.context diff --git a/src/array.jl b/src/array.jl index c0a183b6..36af573d 100644 --- a/src/array.jl +++ b/src/array.jl @@ -370,14 +370,19 @@ end ## indexing -# Host-accessible arrays can be indexed from CPU, bypassing GPUArrays restrictions +# Host-accessible arrays can be indexed from CPU, bypassing GPUArrays restrictions. +# Wait for queued work first, e.g., the kernel computing the result of a reduction. This +# only synchronizes the current task's stream; work submitted by other tasks, or to an +# explicitly created queue, needs to be synchronized explicitly. function Base.getindex(x::oneArray{<:Any, <:Any, <:Union{oneL0.HostBuffer, oneL0.SharedBuffer}}, I::Int) @boundscheck checkbounds(x, I) + synchronize(global_stream(context(x), device())) return unsafe_load(pointer(x, I; type = oneL0.HostBuffer)) end function Base.setindex!(x::oneArray{<:Any, <:Any, <:Union{oneL0.HostBuffer, oneL0.SharedBuffer}}, v, I::Int) @boundscheck checkbounds(x, I) + synchronize(global_stream(context(x), device())) return unsafe_store!(pointer(x, I; type = oneL0.HostBuffer), v) end @@ -599,18 +604,290 @@ Base.unsafe_convert(::Type{ZePtr{T}}, A::PermutedDimsArray) where {T} = ## unsafe_wrap """ - unsafe_wrap(Array, arr::oneArray{_,_,oneL0.SharedBuffer}) + unsafe_wrap(Array, arr::oneArray{_,_,<:Union{oneL0.SharedBuffer,oneL0.HostBuffer}}) -Wrap a Julia `Array` around the buffer that backs a `oneArray`. This is only possible if the -GPU array is backed by a shared buffer, i.e. if it was created with `oneArray{T}(undef, ...)`. +Wrap a Julia `Array` around the buffer that backs a `oneArray`, without copying. This is +only possible if the GPU array is backed by memory that is accessible from the host, i.e., +a shared buffer (as created by `oneArray{T}(undef, ...)`), a host buffer, or host memory +that was itself wrapped using `unsafe_wrap(oneArray, ...)`. + +!!! warning + + The returned `Array` does **not** keep `arr` alive. The caller has to keep a reference + to `arr` for as long as the `Array`, or anything derived from it, is used; otherwise + the `Array` may end up referring to freed memory. Device operations execute + asynchronously, so call `synchronize()` before accessing the returned array after + using `arr` on the device. """ -function Base.unsafe_wrap(::Type{Array}, arr::oneArray{T,N,oneL0.SharedBuffer}) where {T,N} +function Base.unsafe_wrap(::Type{Array}, + arr::oneArray{T,N,<:Union{oneL0.SharedBuffer,oneL0.HostBuffer}}) where {T,N} # TODO: can we make this more convenient by increasing the buffer's refcount and using # a finalizer on the Array? does that work when taking views etc of the Array? ptr = reinterpret(Ptr{T}, pointer(arr)) unsafe_wrap(Array, ptr, size(arr)) end +""" + unsafe_wrap(oneArray, a::Array) + unsafe_wrap(oneArray, ptr::Ptr{T}, dims) + unsafe_wrap(oneArray{T,N,oneL0.HostBuffer}, ...) + +Wrap a `oneArray` around host memory, without copying, so that it can be used on the +device, e.g., in kernels or broadcasts. Changes made through the `oneArray` are visible in +the original array, and vice versa. The resulting array is backed by a host buffer. + +This requires a driver that supports mapping system memory (the +`ZE_extension_external_memmap_sysmem` extension); otherwise an `ArgumentError` is thrown, +and `oneArray(a)` can be used to copy the data instead. + +When wrapping an `Array`, the returned `oneArray` keeps it alive. When wrapping a pointer, +the caller has to make sure the memory stays valid for as long as the `oneArray` is used. +In both cases, the memory must not be freed or reallocated while it is wrapped (e.g., by +calling `resize!` on the original array), and resizing the wrapper is not supported. + +Device operations execute asynchronously, so call `synchronize()` before accessing the +original memory on the host. + +!!! warning + + Level Zero maps host memory in whole pages, and does not support overlapping mappings. + Memory within pages that are already mapped for another wrapper can be wrapped as long + as it lies entirely within that mapping. Memory whose pages only partially overlap an + existing mapping cannot be wrapped, and results in an `ArgumentError`, until the other + wrappers are freed. As a result, whether wrapping succeeds can depend on the order in + which neighboring arrays are wrapped: wrapping a large array first allows wrapping + smaller arrays inside of it, but not the other way around. + +```julia +a = rand(Float32, 1024) +b = unsafe_wrap(oneArray, a) +b .= sin.(b) # executes on the device, updating `a` +synchronize() +``` +""" +unsafe_wrap(::Type{<:oneArray}, ::Any, ::Any...) + +# Level Zero can only map whole pages of system memory, and does not support overlapping +# mappings (e.g., of two small arrays that share a page). Each mapping covers a range of +# pages, and is shared by all wrappers of memory within that range. Memory that only +# partially overlaps an existing mapping cannot be wrapped: that would require replacing +# the mapping, which the driver cannot do while it is in use. Mappings are never replaced +# or freed while wrappers use them. +mutable struct SystemMapping + const ctx::ZeContext + const lo::UInt + const hi::UInt + const buf::oneL0.HostBuffer + # devices the mapping has been made resident on + const resident::Vector{ZeDevice} + # owners of memory within the mapping, which have to outlive it. they are kept until + # the mapping is released, not just until their own wrapper is freed, because other + # wrappers may keep the mapping alive. + const owners::Base.IdSet{Any} + refcount::Int + # set when freeing the mapping failed; its pages then remain occupied + broken::Bool +end + +const system_mappings = SystemMapping[] +const system_mappings_lock = ReentrantLock() + +page_range(ptr::Ptr, bytesize::Integer) = let pagesize = UInt(ccall(:getpagesize, Cint, ())) + lo = UInt(ptr) & ~(pagesize - 1) + hi = Base.checked_add(UInt(ptr), UInt(bytesize), pagesize - 1) & ~(pagesize - 1) + lo, hi +end + +function map_system_memory(ctx::ZeContext, dev::ZeDevice, ptr::Ptr, bytesize::Integer, owner) + lo, hi = page_range(ptr, bytesize) + @lock system_mappings_lock begin + for m in system_mappings + m.ctx == ctx && m.lo < hi && lo < m.hi || continue + if m.broken || !(m.lo <= lo && hi <= m.hi) + throw(ArgumentError("""Cannot wrap host memory at $ptr: its pages ($(repr(lo))-$(repr(hi))) partially overlap memory that is already wrapped (pages $(repr(m.lo))-$(repr(m.hi))). + Level Zero maps host memory in whole pages and does not support overlapping mappings. + Free the other wrapper first, wrap memory in a different order, or use `oneArray(a)` to copy the data instead.""")) + end + if !(dev in m.resident) + make_resident(ctx, dev, m.buf) + push!(m.resident, dev) + end + m.refcount += 1 + owner === nothing || push!(m.owners, owner) + return m + end + + buf = oneL0.host_memmap(ctx, Ptr{Cvoid}(lo), hi - lo) + try + # like any host allocation, the mapping has to be resident to be usable by kernels + make_resident(ctx, dev, buf) + catch + release(buf) + rethrow() + end + owners = Base.IdSet{Any}() + owner === nothing || push!(owners, owner) + m = SystemMapping(ctx, lo, hi, buf, [dev], owners, 1, false) + push!(system_mappings, m) + return m + end +end + +# a wrapper's claim on a mapping. records are linked into a queue when they are freed from a +# finalizer, which keeps the mapping (and the owners it roots) alive until it is released. +mutable struct SystemWrapper + const mapping::SystemMapping + next::Union{Nothing,SystemWrapper} +end + +function unmap_system_memory(m::SystemMapping) + @lock system_mappings_lock begin + m.refcount -= 1 + m.refcount == 0 || return + try + # this waits for outstanding work using the mapping + release(m.buf) + catch err + # keep the mapping registered, which keeps its pages occupied and its owners + # alive, rather than risking mapping them again or freeing memory that the + # device may still access + m.broken = true + @error "Failed to release a mapping of host memory; its memory will be leaked" exception=(err, catch_backtrace()) + return + end + filter!(x -> x !== m, system_mappings) + end + return +end + +# finalizers must not block or call into the driver, so wrappers that are freed from a +# finalizer are queued (without allocating or yielding while holding the lock), and +# released by a task that is woken up through an async condition. +const pending_wrappers = Ref{Union{Nothing,SystemWrapper}}(nothing) +const pending_lock = Threads.SpinLock() +const release_condition = Ref{Base.AsyncCondition}() +const release_task_lock = ReentrantLock() + +function start_release_task() + isassigned(release_condition) && return + @lock release_task_lock begin + isassigned(release_condition) && return + cond = Base.AsyncCondition() + errormonitor(@async while true + wait(cond) + process_pending_wrappers() + end) + release_condition[] = cond + end + return +end + +function process_pending_wrappers() + lock(pending_lock) + w = pending_wrappers[] + pending_wrappers[] = nothing + unlock(pending_lock) + while w !== nothing + next = w.next + w.next = nothing + unmap_system_memory(w.mapping) + w = next + end +end + +function free_system_wrapper(w::SystemWrapper) + if ccall(:jl_gc_is_in_finalizer, Int8, ()) == 0 + unmap_system_memory(w.mapping) + else + lock(pending_lock) + w.next = pending_wrappers[] + pending_wrappers[] = w + unlock(pending_lock) + ccall(:uv_async_send, Cint, (Ptr{Cvoid},), release_condition[].handle) + end + return +end + +system_memmap_supported(drv::ZeDriver) = + haskey(oneL0.extension_properties(drv), "ZE_extension_external_memmap_sysmem") + +# `owner` is kept alive for as long as the wrapper +function wrap_system_memory(::Type{oneArray{T,N,B}}, ptr::Ptr{T}, dims::NTuple{N,Int}, + owner=nothing) where {T,N,B} + B == oneL0.HostBuffer || + throw(ArgumentError("Cannot wrap host memory as $B; use oneL0.HostBuffer")) + check_eltype(T) + isbitstype(T) || throw(ArgumentError("Can only unsafe_wrap a pointer to a bits type")) + all(>=(0), dims) || throw(ArgumentError("Invalid dimensions $dims")) + bytesize = Base.checked_mul(foldl(Base.checked_mul, dims; init=1), sizeof(T)) + if bytesize > 0 && ptr == C_NULL + throw(ArgumentError("Cannot wrap a NULL pointer")) + end + if !iszero(UInt(ptr) % Base.datatype_alignment(T)) + throw(ArgumentError("Pointer $ptr is not sufficiently aligned for elements of type $T")) + end + if Base.Checked.add_with_overflow(UInt(ptr), UInt(bytesize) + UInt(ccall(:getpagesize, Cint, ())))[2] + throw(ArgumentError("Memory range of $bytesize bytes at $ptr exceeds the address space")) + end + if !system_memmap_supported(driver()) + throw(ArgumentError("""The Level Zero driver does not support mapping host memory, which is required to wrap it as a oneArray. + Use `oneArray(a)` to copy the data instead.""")) + end + + ctx = context() + buf = oneL0.HostBuffer(Ptr{Cvoid}(ptr), bytesize, ctx) + bytesize == 0 && return oneArray{T,N}(DataRef(Returns(nothing), buf), dims) + + start_release_task() + m = map_system_memory(ctx, device(), ptr, bytesize, owner) + arr, data = try + w = SystemWrapper(m, nothing) + data = DataRef(buf) do _ + free_system_wrapper(w) + end + oneArray{T,N}(data, dims), data + catch + # give up our claim on the mapping (which roots the owner until it is released) + unmap_system_memory(m) + rethrow() + end + # the array holds its own reference + unsafe_free!(data) + return arr +end +wrap_system_memory(::Type{oneArray{T,N}}, ptr::Ptr{T}, dims::NTuple{N,Int}, owner=nothing) where {T,N} = + wrap_system_memory(oneArray{T,N,oneL0.HostBuffer}, ptr, dims, owner) + +Base.unsafe_wrap(::Union{Type{oneArray},Type{oneArray{T}},Type{oneArray{T,N}}}, + ptr::Ptr{T}, dims::NTuple{N,Int}) where {T,N} = + wrap_system_memory(oneArray{T,N}, ptr, dims) +Base.unsafe_wrap(::Type{oneArray{T,N,B}}, ptr::Ptr{T}, dims::NTuple{N,Int}) where {T,N,B} = + wrap_system_memory(oneArray{T,N,B}, ptr, dims) + +# integer size input +Base.unsafe_wrap(::Union{Type{oneArray},Type{oneArray{T}},Type{oneArray{T,1}}}, + ptr::Ptr{T}, dim::Integer) where {T} = + unsafe_wrap(oneArray{T,1}, ptr, (Int(dim),)) +Base.unsafe_wrap(::Type{oneArray{T,1,B}}, ptr::Ptr{T}, dim::Integer) where {T,B} = + unsafe_wrap(oneArray{T,1,B}, ptr, (Int(dim),)) + +# array input: keep the array alive for as long as the wrapper +Base.unsafe_wrap(::Union{Type{oneArray},Type{oneArray{T}},Type{oneArray{T,N}}}, + a::Array{T,N}) where {T,N} = + wrap_system_memory(oneArray{T,N}, pointer(a), size(a), a) +Base.unsafe_wrap(::Type{oneArray{T,N,B}}, a::Array{T,N}) where {T,N,B} = + wrap_system_memory(oneArray{T,N,B}, pointer(a), size(a), a) + +# whether an array wraps host memory using `unsafe_wrap` +function is_system(a::oneArray) + buf = a.data[] + buf isa oneL0.HostBuffer && sizeof(buf) > 0 || return false + lo, hi = page_range(pointer(buf), sizeof(buf)) + @lock system_mappings_lock begin + any(m -> m.ctx == context(a) && m.lo <= lo && hi <= m.hi, system_mappings) + end +end + ## resizing @@ -622,6 +899,9 @@ the first `n` elements will be retained. If `n` is larger, the new elements are guaranteed to be initialized. """ function Base.resize!(a::oneVector{T}, n::Integer) where {T} + # resizing would detach the array from the host memory it wraps + is_system(a) && throw(ArgumentError("Cannot resize a oneArray that wraps host memory")) + # TODO: add additional space to allow for quicker resizing maxsize = n * sizeof(T) bufsize = if isbitstype(T) diff --git a/test/array.jl b/test/array.jl index 5d92d759..6ae0c085 100644 --- a/test/array.jl +++ b/test/array.jl @@ -101,6 +101,179 @@ end @test Array(a) == [100, 42] oneAPI.@sync copyto!(a, 2, [200], 1, 1) @test b == [100, 200] + + # the same works for arrays backed by host memory + c = oneVector{Int,oneL0.HostBuffer}([1, 2]) + d = unsafe_wrap(Array, c) + @test d == [1, 2] + d[1] = 100 + @test Array(c) == [100, 2] +end + +# arrays that occupy whole pages, so that wrapping them cannot be rejected because their +# pages partially overlap memory that is already wrapped +function page_aligned_array(T, n) + pagesize = ccall(:getpagesize, Cint, ()) + ref = Ref{Ptr{Cvoid}}() + bytes = cld(n * sizeof(T), pagesize) * pagesize + @assert ccall(:posix_memalign, Cint, (Ptr{Ptr{Cvoid}}, Csize_t, Csize_t), ref, pagesize, bytes) == 0 + return unsafe_wrap(Array, Ptr{T}(ref[]), n; own=true) +end + +# wrap an array and memory inside of it, free the outer wrapper, and return the inner one +@noinline function wrap_inner() + big = page_aligned_array(Float32, 4096) + outer = unsafe_wrap(oneArray, big) + inner = GC.@preserve big unsafe_wrap(oneArray, pointer(big) + 4096, 16) + oneAPI.unsafe_free!(outer) + return inner, WeakRef(big) +end + +# create wrappers that are unreachable once this returns +@noinline function wrap_garbage(n) + for _ in 1:n + unsafe_wrap(oneArray, page_aligned_array(Float32, 4096)) .+= 1 + end +end + +function wrap_kernel(a) + i = get_global_id() + @inbounds a[i] = i + return +end + +@testset "wrapping host memory" begin + if !oneAPI.system_memmap_supported(oneAPI.driver()) + @test_throws ArgumentError unsafe_wrap(oneArray, Float32[1]) + else + a = Float32[1, 2, 3, 4] + b = unsafe_wrap(oneArray, a) + @test b isa oneVector{Float32, oneL0.HostBuffer} + @test size(b) == size(a) + @test UInt(pointer(b)) == UInt(pointer(a)) + + # changes are visible in both directions + b .+= 1 + synchronize() + @test a == [2, 3, 4, 5] + a[1] = 10 + c = oneArray{Float32, 1, oneL0.DeviceBuffer}(undef, 4) + c .= b # read on the device, not via the host + @test Array(c) == [10, 3, 4, 5] + @test pointer(unsafe_wrap(Array, b)) == pointer(a) + + for AT in [oneArray, oneArray{Float32}, oneArray{Float32, 1}, + oneArray{Float32, 1, oneL0.HostBuffer}], + f in [x -> unsafe_wrap(AT, pointer(x), length(x)), + x -> unsafe_wrap(AT, pointer(x), size(x)), + x -> unsafe_wrap(AT, x)] + d = f(a) + @test d isa oneVector{Float32, oneL0.HostBuffer} + @test Array(d) == a + end + let m = rand(Float32, 3, 4) + d = unsafe_wrap(oneArray, m) + @test d isa oneMatrix{Float32} + d .*= 2 + synchronize() + @test Array(d) == m + end + @test isempty(Array(unsafe_wrap(oneArray, Float32[]))) + + # large arrays, and copies to and from regular device arrays + let x = copyto!(page_aligned_array(Float32, 10^6), rand(Float32, 10^6)) + ref = x .* 2 + d = unsafe_wrap(oneArray, x) + d .*= 2 + synchronize() + @test x == ref + e = oneArray(ref .+ 1) + copyto!(d, e) + synchronize() + @test x == ref .+ 1 + end + + fill!(b, 42) + synchronize() + @test all(==(42), a) + @oneapi items=length(b) wrap_kernel(b) + synchronize() + @test a == [1, 2, 3, 4] + + # memory within pages that are already mapped can be wrapped too + nmappings = length(oneAPI.system_mappings) + let big = page_aligned_array(Float32, 4096) + wbig = unsafe_wrap(oneArray, big) + GC.@preserve big begin + inner = unsafe_wrap(oneArray, pointer(big) + 4096, 16) + wbig .= 1 + inner .= 2 + synchronize() + @test big[1] == 1 && big[1025] == 2 + # partially overlapping memory cannot be wrapped + @test_throws ArgumentError unsafe_wrap(oneArray, pointer(big) + 4 * 4096 - 16, 1024) + oneAPI.unsafe_free!(inner) + end + oneAPI.unsafe_free!(wbig) + end + # small arrays share pages, which are mapped once + # a mapping keeps the memory it covers alive, even after the wrapper that created it + # is freed while other wrappers still use the mapping + inner, big_ref = wrap_inner() + for _ in 1:3 + GC.gc(true) + end + @test big_ref.value !== nothing + inner .= 5 + synchronize() + @test Array(inner) == fill(5, 16) + oneAPI.unsafe_free!(inner) + @test length(oneAPI.system_mappings) == nmappings + + xs = [Float32[i, i] for i in 1:8] + ws = [unsafe_wrap(oneArray, x) for x in xs] + push!(ws, unsafe_wrap(oneArray, xs[1])) # wrapping the same memory twice + for w in ws + w .+= 1 + end + synchronize() + @test xs[1] == [3, 3] + @test xs[2] == [3, 3] + foreach(oneAPI.unsafe_free!, ws) + @test length(oneAPI.system_mappings) == nmappings + + # wrappers freed by the GC are released asynchronously + wrap_garbage(10) + synchronize() + for _ in 1:100 + GC.gc(true) + length(oneAPI.system_mappings) == nmappings && break + sleep(0.1) + end + @test length(oneAPI.system_mappings) == nmappings + + # the wrapper keeps the array alive + d = unsafe_wrap(oneArray, fill!(page_aligned_array(Float32, 1024), 1)) + GC.gc(true) + @test sum(d) == 1024 + + @test_throws ArgumentError resize!(b, 5) + @test_throws ArgumentError unsafe_wrap(oneVector{Float32, oneL0.DeviceBuffer}, a) + @test_throws ArgumentError unsafe_wrap(oneArray, Ptr{Float32}(C_NULL), 1) + @test_throws ArgumentError unsafe_wrap(oneArray, pointer(a), (-1,)) + @test_throws ArgumentError unsafe_wrap(oneArray, Ptr{Float32}(typemax(UInt) - 15), 16) + GC.@preserve a begin + @test_throws ArgumentError unsafe_wrap(oneArray, Ptr{Float32}(pointer(a) + 1), 1) + end + end +end + +@testset "reductions of host-accessible arrays" begin + for B in (oneL0.SharedBuffer, oneL0.HostBuffer) + a = oneArray{Float32, 1, B}(fill(1.0f0, 1024)) + @test sum(a) == 1024 + @test maximum(a) == 1 + end end # https://github.com/JuliaGPU/CUDA.jl/issues/2191