[KA] support unified memory API

[simeonschaub]

No description provided.

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic master) to apply these changes.

Click here to view the suggested changes.

diff --git a/lib/cl/state.jl b/lib/cl/state.jl
index bed42dc..da99fe9 100644
--- a/lib/cl/state.jl
+++ b/lib/cl/state.jl
@@ -9,7 +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)
+    return delete!(task_local_storage(), :CLUnifiedMemoryBackend)
 end
 
 
@@ -164,7 +164,7 @@ struct SVMBackend <: AbstractMemoryBackend end
 struct USMBackend <: AbstractMemoryBackend end
 struct BufferBackend <: AbstractMemoryBackend end
 
-function supported_memory_backends(dev::Device; unified=false)
+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
@@ -197,7 +197,7 @@ function supported_memory_backends(dev::Device; unified=false)
     return backends
 end
 
-function default_memory_backend(dev::Device; unified=false)
+function default_memory_backend(dev::Device; unified = false)
     supported_backends = supported_memory_backends(dev; unified)
     isempty(supported_backends) && return nothing
 
@@ -234,7 +234,7 @@ end
 function unified_memory_backend()
     return get!(task_local_storage(), :CLUnifiedMemoryBackend) do
         dev = device()
-        backend = default_memory_backend(dev; unified=true)
+        backend = default_memory_backend(dev; unified = true)
         if backend === nothing
             error("Device $(dev) does not support any of the available unified memory backends")
         end
diff --git a/src/OpenCLKernels.jl b/src/OpenCLKernels.jl
index e06102c..9891360 100644
--- a/src/OpenCLKernels.jl
+++ b/src/OpenCLKernels.jl
@@ -17,7 +17,7 @@ export OpenCLBackend
 struct OpenCLBackend <: KA.GPU
 end
 
-function KA.allocate(::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where T
+function KA.allocate(::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where {T}
     if unified
         memory_backend = cl.unified_memory_backend()
         if memory_backend === cl.USMBackend()
@@ -32,7 +32,7 @@ function KA.allocate(::OpenCLBackend, ::Type{T}, dims::Tuple; unified::Bool = fa
     end
 end
 
-KA.supports_unified(::OpenCLBackend) = cl.default_memory_backend(cl.device(); unified=true) !== nothing
+KA.supports_unified(::OpenCLBackend) = cl.default_memory_backend(cl.device(); unified = true) !== nothing
 
 KA.get_backend(::CLArray) = OpenCLBackend()
 # TODO should be non-blocking

[codecov]

Codecov Report

❌ Patch coverage is 90.00000% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 80.24%. Comparing base (95039b8) to head (adca066).

Files with missing lines	Patch %	Lines
src/OpenCLKernels.jl	90.00%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master     #394       +/-   ##
===========================================
+ Coverage   31.39%   80.24%   +48.84%     
===========================================
  Files          11       12        +1     
  Lines         672      729       +57     
===========================================
+ Hits          211      585      +374     
+ Misses        461      144      -317     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[maleadt]

SharedVirtualMemory is also shared. Maybe we should do something smarter based on the back-end, a la

OpenCL.jl/src/array.jl

Lines 96 to 105 in 53f450a

	# default to non-unified memory
	function memory_type()
	if cl.memory_backend() == cl.USMBackend()
	return cl.UnifiedDeviceMemory
	elseif cl.memory_backend() == cl.SVMBackend()
	return cl.SharedVirtualMemory
	elseif cl.memory_backend() == cl.BufferBackend()
	return cl.Buffer
	end
	end

(which could also be used in the mixed broadcast implementation).

[simeonschaub]

Like this?

[maleadt]

Do we need a separate unified_memory_backend? I figured just basing the decision to use USM or SVM for unified memory based on the existing memory_backend.

[simeonschaub]

The issue I saw with that is that if we have no USM support but support for BDA, we'll always pick BufferBackend over SVMBackend, in that case we won't be able to support unified memory. If you don't think we should care about this edge case I can also simply use memory_backend though

[maleadt]

Yeah, I guess. Although it's definitely an edge case, e.g., I don't know how well we support copies from unified BDA memory to Buffers (to support the unified <-> device copies).

[simeonschaub]

Merge?