Allow direct host-interactions with SVM-backed arrays (#336)

Co-authored-by: Tim Besard <[email protected]>

Author: VarLad

.github/workflows/Test.yml

@@ -15,7 +15,7 @@ jobs:
   test:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ matrix.memory_backend }} - PoCL ${{ matrix.pocl }}
     runs-on: ${{ matrix.os }}
-    timeout-minutes: 180
+    timeout-minutes: 60
     permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created
       actions: write
       contents: read

lib/cl/memory/svm.jl

@@ -45,7 +45,7 @@ Base.show(io::IO, mem::SharedVirtualMemory) =
     @printf(io, "SharedVirtualMemory(%s at %p)", Base.format_bytes(sizeof(mem)), Int(pointer(mem)))
 
 Base.convert(::Type{Ptr{T}}, mem::SharedVirtualMemory) where {T} =
-    convert(Ptr{T}, pointer(mem))
+    convert(Ptr{T}, reinterpret(Ptr{Cvoid}, pointer(mem)))
 
 Base.convert(::Type{CLPtr{T}}, mem::SharedVirtualMemory) where {T} =
     reinterpret(CLPtr{T}, pointer(mem))

src/array.jl

@@ -1,4 +1,5 @@
-export CLArray, CLVector, CLMatrix, CLVecOrMat, is_device, is_shared, is_host
+export CLArray, CLVector, CLMatrix, CLVecOrMat,
+       device_accessible, host_accessible
 
 
 ## array type
@@ -176,11 +177,9 @@ memtype(x::CLArray) = memtype(typeof(x))
 memtype(::Type{<:CLArray{<:Any, <:Any, M}}) where {M} = @isdefined(M) ? M : Any
 
 # can we read this array from the device (i.e. derive a CLPtr)?
-is_device(a::CLArray) =
+device_accessible(a::CLArray) =
     memtype(a) in (cl.UnifiedDeviceMemory, cl.UnifiedSharedMemory, cl.SharedVirtualMemory, cl.Buffer)
-is_shared(a::CLArray) =
-    memtype(a) in (cl.UnifiedSharedMemory, cl.SharedVirtualMemory)
-is_host(a::CLArray) =
+host_accessible(a::CLArray) =
     memtype(a) in (cl.UnifiedHostMemory, cl.UnifiedSharedMemory, cl.SharedVirtualMemory)
 
 
@@ -272,12 +271,12 @@ Base.convert(::Type{T}, x::T) where {T <: CLArray} = x
 
 ## indexing
 
-function Base.getindex(x::CLArray{<:Any, <:Any, <:Union{cl.UnifiedHostMemory, cl.UnifiedSharedMemory}}, I::Int)
+function Base.getindex(x::CLArray{<:Any, <:Any, <:Union{cl.UnifiedHostMemory, cl.UnifiedSharedMemory, cl.SharedVirtualMemory}}, I::Int)
     @boundscheck checkbounds(x, I)
     return GC.@preserve x unsafe_load(host_pointer(x, I))
 end
 
-function Base.setindex!(x::CLArray{<:Any, <:Any, <:Union{cl.UnifiedHostMemory, cl.UnifiedSharedMemory}}, v, I::Int)
+function Base.setindex!(x::CLArray{<:Any, <:Any, <:Union{cl.UnifiedHostMemory, cl.UnifiedSharedMemory, cl.SharedVirtualMemory}}, v, I::Int)
     @boundscheck checkbounds(x, I)
     return GC.@preserve x unsafe_store!(host_pointer(x, I), v)
 end
@@ -286,14 +285,14 @@ end
 ## interop with libraries
 
 function Base.unsafe_convert(::Type{Ptr{T}}, x::CLArray{T}) where {T}
-    if !is_host(x)
+    if !host_accessible(x)
         throw(ArgumentError("cannot take the CPU address of a $(typeof(x))"))
     end
     return convert(Ptr{T}, x.data[]) + x.offset * Base.elsize(x)
 end
 
 function Base.unsafe_convert(::Type{CLPtr{T}}, x::CLArray{T}) where {T}
-    if !is_device(x)
+    if !device_accessible(x)
         throw(ArgumentError("cannot take the device address of a $(typeof(x))"))
     end
     return convert(CLPtr{T}, x.data[]) + x.offset * Base.elsize(x)
@@ -485,16 +484,14 @@ Base.unsafe_convert(::Type{CLPtr{T}}, A::PermutedDimsArray) where {T} =
 ## unsafe_wrap
 
 """
-    unsafe_wrap(Array, arr::CLArray{_,_,cl.UnifiedSharedMemory})
+    unsafe_wrap(Array, arr::CLArray)
 
 Wrap a Julia `Array` around the buffer that backs a `CLArray`. This is only possible if the
-GPU array is backed by a shared buffer, i.e. if it was created with `CLArray{T}(undef, ...)`.
+GPU array is backed by host memory, such as unified (host or shared) memory, or shared
+virtual memory.
 """
-function Base.unsafe_wrap(::Type{Array}, arr::CLArray{T, N, cl.UnifiedSharedMemory}) 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))
-    return unsafe_wrap(Array, ptr, size(arr))
+function Base.unsafe_wrap(::Type{Array}, arr::CLArray{T, N}) where {T, N}
+    return unsafe_wrap(Array, host_pointer(arr), size(arr))
 end
 
 

src/memory.jl

@@ -15,10 +15,15 @@ mutable struct Managed{M}
     # whether there are outstanding operations that haven't been synchronized
     dirty::Bool
 
-    function Managed(mem::cl.AbstractMemory; queue = cl.queue(), dirty = true)
+    # who is currently using the memory
+    user::Symbol
+
+    function Managed(mem::cl.AbstractMemory; queue = cl.queue(), dirty = true, user = :device)
         # NOTE: memory starts as dirty, because stream-ordered allocations are only
         #       guaranteed to be physically allocated at a synchronization event.
-        return new{typeof(mem)}(mem, queue, dirty)
+        # NOTE: memory also starts as device-owned, because we need to map it as soon as
+        #       the host accesses it.
+        return new{typeof(mem)}(mem, queue, dirty, user)
     end
 end
 
@@ -38,7 +43,7 @@ function maybe_synchronize(managed::Managed)
     return nothing
 end
 
-function Base.convert(typ::Union{Type{<:CLPtr}, Type{cl.Buffer}}, managed::Managed)
+function Base.convert(typ::Union{Type{<:CLPtr}, Type{cl.Buffer}}, managed::Managed{M}) where {M}
     # let null pointers pass through as-is
     # XXX: does not work for buffers
     ptr = convert(typ, managed.mem)
@@ -52,6 +57,13 @@ function Base.convert(typ::Union{Type{<:CLPtr}, Type{cl.Buffer}}, managed::Manag
         managed.queue = cl.queue()
     end
 
+    # coarse-grained SVM needs to be unmapped when accessing it back from the device
+    # TODO: support fine-grained SVM
+    if M == cl.SharedVirtualMemory && managed.user == :host
+        cl.enqueue_svm_unmap(pointer(managed.mem))
+        managed.user = :device
+    end
+
     managed.dirty = true
     return ptr
 end
@@ -74,6 +86,14 @@ function Base.convert(typ::Type{<:Ptr}, managed::Managed{M}) where {M}
 
     # make sure any work on the memory has finished.
     maybe_synchronize(managed)
+
+    # coarse-grained SVM needs to be mapped when initially accessing it from the host
+    # TODO: support fine-grained SVM
+    if M == cl.SharedVirtualMemory && managed.user != :host
+        cl.enqueue_svm_map(pointer(managed.mem), sizeof(managed.mem), :rw; blocking = true)
+        managed.user = :host
+    end
+
     return ptr
 end
 
@@ -162,6 +182,10 @@ function free(managed::Managed)
         end
 
         if mem isa cl.SharedVirtualMemory
+            if managed.user == :host
+                # if the coarse-grained SVM buffer is mapped on the host, unmap it first.
+                cl.enqueue_svm_unmap(pointer(mem))
+            end
             cl.svm_free(mem)
         elseif mem isa cl.UnifiedMemory
             cl.usm_free(mem)

test/Project.toml

@@ -8,6 +8,7 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 REPL = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

test/array.jl

@@ -9,8 +9,10 @@ import Adapt
         @test Base.elsize(xs) == sizeof(Int)
         @test CLArray{Int, 2}(xs) === xs
 
-        @test_throws ArgumentError Base.unsafe_convert(Ptr{Int}, xs)
-        @test_throws ArgumentError Base.unsafe_convert(Ptr{Float32}, xs)
+        if !host_accessible(xs)
+            @test_throws ArgumentError Base.unsafe_convert(Ptr{Int}, xs)
+            @test_throws ArgumentError Base.unsafe_convert(Ptr{Float32}, xs)
+        end
 
         @test collect(OpenCL.zeros(Float32, 2, 2)) == zeros(Float32, 2, 2)
         @test collect(OpenCL.ones(Float32, 2, 2)) == ones(Float32, 2, 2)

test/pointer.jl

@@ -37,17 +37,17 @@ end
 
 
 @testset "GPU or CPU integration" begin
-
     a = [1]
     ccall(:clock, Nothing, (Ptr{Int},), a)
     @test_throws Exception ccall(:clock, Nothing, (CLPtr{Int},), a)
     ccall(:clock, Nothing, (PtrOrCLPtr{Int},), a)
 
     b = CLArray{eltype(a), ndims(a)}(undef, size(a))
     ccall(:clock, Nothing, (CLPtr{Int},), b)
-    @test_throws Exception ccall(:clock, Nothing, (Ptr{Int},), b)
+    if !host_accessible(b)
+        @test_throws Exception ccall(:clock, Nothing, (Ptr{Int},), b)
+    end
     ccall(:clock, Nothing, (PtrOrCLPtr{Int},), b)
-
 end
 
 

test/runtests.jl

@@ -3,6 +3,7 @@ using Dates
 import REPL
 using Printf: @sprintf
 using Base.Filesystem: path_separator
+using Preferences
 
 # parse some command-line arguments
 function extract_flag!(args, flag, default=nothing)
@@ -110,7 +111,18 @@ if !isempty(optlike_args)
     error("Unknown test options `$(join(optlike_args, " "))` (try `--help` for usage instructions)")
 end
 ## the remaining args filter tests
-if !isempty(ARGS)
+if isempty(ARGS)
+  # default to running all tests, except:
+  filter!(tests) do test
+    if load_preference(OpenCL, "default_memory_backend") == "svm" &&
+       test == "gpuarrays/indexing scalar"
+        # GPUArrays' scalar indexing tests assume that indexing is not supported
+        return false
+    end
+
+    return true
+  end
+else
   filter!(tests) do test
     any(arg->startswith(test, arg), ARGS)
   end