Allow direct devices as well

LuxDL · Mar 27, 2024 · 0faa572 · 0faa572
1 parent f0bc7eb
commit 0faa572
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 3 deletions.
diff --git a/ext/LuxDeviceUtilsAMDGPUExt.jl b/ext/LuxDeviceUtilsAMDGPUExt.jl
@@ -27,6 +27,14 @@ LuxDeviceUtils.default_device_rng(::LuxAMDGPUDevice) = AMDGPU.rocrand_rng()
 LuxDeviceUtils.get_device(x::AMDGPU.AnyROCArray) = LuxAMDGPUDevice(AMDGPU.device(x))
 
 # Set Device
+function LuxDeviceUtils.set_device!(::Type{LuxAMDGPUDevice}, dev::AMDGPU.HIPDevice)
+    if !AMDGPU.functional()
+        @warn "AMDGPU is not functional."
+        return
+    end
+    AMDGPU.device!(dev)
+    return
+end
 function LuxDeviceUtils.set_device!(::Type{LuxAMDGPUDevice}, id::Int)
     if !AMDGPU.functional()
         @warn "AMDGPU is not functional."

diff --git a/ext/LuxDeviceUtilsCUDAExt.jl b/ext/LuxDeviceUtilsCUDAExt.jl
@@ -28,6 +28,14 @@ LuxDeviceUtils.default_device_rng(::LuxCUDADevice) = CUDA.default_rng()
 LuxDeviceUtils.get_device(x::CUDA.AnyCuArray) = LuxCUDADevice(CUDA.device(x))
 
 # Set Device
+function LuxDeviceUtils.set_device!(::Type{LuxCUDADevice}, dev::CUDA.CuDevice)
+    if !CUDA.functional()
+        @warn "CUDA is not functional."
+        return
+    end
+    CUDA.device!(dev)
+    return
+end
 function LuxDeviceUtils.set_device!(::Type{LuxCUDADevice}, id::Int)
     if !CUDA.functional()
         @warn "CUDA is not functional."

diff --git a/src/LuxDeviceUtils.jl b/src/LuxDeviceUtils.jl
@@ -355,18 +355,20 @@ const SET_DEVICE_DANGER = """
 """
 
 """
-    set_device!(T::Type{<:AbstractLuxDevice}, id::Int)
+    set_device!(T::Type{<:AbstractLuxDevice}, dev_or_id)
 
 $SET_DEVICE_DOCS
 
 ## Arguments
 
   - `T::Type{<:AbstractLuxDevice}`: The device type to set.
-  - `id::Int`: The device id to set. This is `1`-indexed.
+  - `dev_or_id`: Can be the device from the corresponding package. For example for CUDA it
+    can be a `CuDevice`. If it is an integer, it is the device id to set. This is
+    `1`-indexed.
 
 $SET_DEVICE_DANGER
 """
-function set_device!(::Type{T}, id::Int) where {T <: AbstractLuxDevice}
+function set_device!(::Type{T}, dev_or_id) where {T <: AbstractLuxDevice}
     T === LuxCUDADevice &&
         @warn "`CUDA.jl` hasn't been loaded. Ignoring the device setting." maxlog=1
     T === LuxAMDGPUDevice &&