Support new functionality from KA 0.9.32 (#2774)

Co-authored-by: Valentin Churavy <[email protected]>

Author: michel2323

Project.toml

@@ -63,7 +63,7 @@ ExprTools = "0.1"
 GPUArrays = "11.2.1"
 GPUCompiler = "0.24, 0.25, 0.26, 0.27, 1"
 GPUToolbox = "0.2"
-KernelAbstractions = "0.9.2"
+KernelAbstractions = "0.9.32"
 LLVM = "9.1"
 LLVMLoopInfo = "1"
 LazyArtifacts = "1"

src/CUDAKernels.jl

@@ -38,9 +38,6 @@ Adapt.adapt_storage(::KA.CPU, a::Union{CuArray,CUSPARSE.AbstractCuSparseArray})
 ## memory operations
 
 function KA.copyto!(::CUDABackend, A, B)
-    A isa Array && CUDA.pin(A)
-    B isa Array && CUDA.pin(B)
-
     GC.@preserve A B begin
         destptr = pointer(A)
         srcptr  = pointer(B)
@@ -50,6 +47,25 @@ function KA.copyto!(::CUDABackend, A, B)
     return A
 end
 
+function KA.pagelock!(::CUDABackend, A::Array)
+    CUDA.pin(A)
+    return nothing
+end
+
+## device operations
+
+function KA.ndevices(::CUDABackend)
+    return ndevices()
+end
+
+function KA.device(::CUDABackend)
+    deviceid(CUDA.active_state().device) + 1
+end
+
+function KA.device!(::CUDABackend, id::Int32)
+    device!(id - 1)
+end
+
 ## kernel launch
 
 function KA.mkcontext(kernel::KA.Kernel{CUDABackend}, _ndrange, iterspace)