support SVM as well

Author: simeonschaub

lib/cl/state.jl

@@ -9,6 +9,7 @@ function clear_task_local_storage!()
     delete!(task_local_storage(), :CLPlatform)
     delete!(task_local_storage(), :CLQueue)
     delete!(task_local_storage(), :CLMemoryBackend)
+    delete!(task_local_storage(), :CLUnifiedMemoryBackend)
 end
 
 
@@ -163,21 +164,23 @@ struct SVMBackend <: AbstractMemoryBackend end
 struct USMBackend <: AbstractMemoryBackend end
 struct BufferBackend <: AbstractMemoryBackend end
 
-function supported_memory_backends(dev::Device)
+function supported_memory_backends(dev::Device; unified=false)
     backends = AbstractMemoryBackend[]
 
     # unified shared memory is the first choice, as it gives us separate host and device
     # memory spaces that can be directly referenced by raw pointers.
     if usm_supported(dev)
         usm_caps = usm_capabilities(dev)
-        if usm_caps.host.access && usm_caps.device.access
+        if unified
+            usm_caps.shared.access && push!(backends, USMBackend())
+        elseif usm_caps.host.access && usm_caps.device.access
             push!(backends, USMBackend())
         end
     end
 
     # plain old buffers are always supported, but we only want to use them if we have the
     # buffer device address extension, which allows us to reference them by raw pointers.
-    if bda_supported(dev)
+    if !unified && bda_supported(dev)
         push!(backends, BufferBackend())
     end
 
@@ -187,7 +190,7 @@ function supported_memory_backends(dev::Device)
         push!(backends, SVMBackend())
     end
 
-    if isempty(backends)
+    if !unified && isempty(backends)
         # as a last resort, use plain buffers without the ability to reference by pointer.
         # this severely limits compatibility, but it's better than nothing.
         push!(backends, BufferBackend())
@@ -196,8 +199,9 @@ function supported_memory_backends(dev::Device)
     return backends
 end
 
-function default_memory_backend(dev::Device)
-    supported_backends = supported_memory_backends(dev)
+function default_memory_backend(dev::Device; unified=false)
+    supported_backends = supported_memory_backends(dev; unified)
+    isempty(supported_backends) && return nothing
 
     backend_str = load_preference(OpenCL, "default_memory_backend")
     backend_str === nothing && return first(supported_backends)
@@ -211,8 +215,7 @@ function default_memory_backend(dev::Device)
     else
         error("Unknown memory backend '$backend_str' requested")
     end
-    in(backend, supported_backends) ? backend : nothing
-    backend
+    return in(backend, supported_backends) ? backend : nothing
 end
 
 function memory_backend()
@@ -230,6 +233,17 @@ function memory_backend()
     end
 end
 
+function unified_memory_backend()
+    return get!(task_local_storage(), :CLUnifiedMemoryBackend) do
+        dev = device()
+        backend = default_memory_backend(dev; unified=true)
+        if backend === nothing
+            error("Device $(dev) does not support any of the available unified memory backends")
+        end
+        backend
+    end
+end
+
 
 ## per-task queues
 

src/OpenCLKernels.jl

@@ -19,13 +19,20 @@ end
 
 function KA.allocate(::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T
     if unified
-        return CLArray{T, length(dims), cl.UnifiedSharedMemory}(undef, dims)
+        memory_backend = cl.unified_memory_backend()
+        if memory_backend === cl.USMBackend()
+            return CLArray{T, length(dims), cl.UnifiedSharedMemory}(undef, dims)
+        elseif memory_backend === cl.SVMBackend()
+            return CLArray{T, length(dims), cl.SharedVirtualMemory}(undef, dims)
+        else
+            throw(ArgumentError("Unified memory not supported"))
+        end
     else
         return CLArray{T}(undef, dims)
     end
 end
 
-KA.supports_unified(::OpenCLBackend) = cl.usm_supported(cl.device())
+KA.supports_unified(::OpenCLBackend) = cl.default_memory_backend(cl.device(); unified=true) !== nothing
 
 KA.get_backend(::CLArray) = OpenCLBackend()
 # TODO should be non-blocking