implement device API

JuliaGPU · Jan 21, 2025 · 67f269f · 67f269f
1 parent 4bec614
commit 67f269f
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/Project.toml b/Project.toml
@@ -61,7 +61,7 @@ EnzymeCore = "0.8.2"
 ExprTools = "0.1"
 GPUArrays = "11.2.1"
 GPUCompiler = "0.24, 0.25, 0.26, 0.27, 1"
-KernelAbstractions = "0.9.2"
+KernelAbstractions = "0.9.32"
 LLVM = "9.1"
 LLVMLoopInfo = "1"
 LazyArtifacts = "1"

diff --git a/src/CUDAKernels.jl b/src/CUDAKernels.jl
@@ -35,9 +35,6 @@ Adapt.adapt_storage(::KA.CPU, a::CuArray) = convert(Array, a)
 ## 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)
@@ -47,6 +44,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(state.device)+1
+end
+
+function KA.device!(::CUDABackend, id)
+    device!(id-1)
+end
+
 ## kernel launch
 
 function KA.mkcontext(kernel::KA.Kernel{CUDABackend}, _ndrange, iterspace)