Merge pull request #5 from LuxDL/ap/simplify

Simplify Accelerator Loading + Aqua

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "LuxDeviceUtils"
 uuid = "34f89e08-e1d5-43b4-8944-0b49ac560553"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.1.1"
+version = "0.1.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

README.md

@@ -8,6 +8,7 @@
 [![Buildkite](https://badge.buildkite.com/b098d6387b2c69bd0ab684293ff66332047b219e1b8f9bb486.svg)](https://buildkite.com/julialang/luxdeviceutils-dot-jl)
 [![codecov](https://codecov.io/gh/LuxDL/LuxDeviceUtils.jl/branch/main/graph/badge.svg?token=1ZY0A2NPEM)](https://codecov.io/gh/LuxDL/LuxDeviceUtils.jl)
 [![Package Downloads](https://shields.io/endpoint?url=https://pkgs.genieframework.com/api/v1/badge/LuxDeviceUtils)](https://pkgs.genieframework.com?packages=LuxDeviceUtils)
+[![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
 
 [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
 [![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle)

docs/src/index.md

@@ -37,5 +37,11 @@ gpu_backend!
 ```@docs
 cpu_device
 gpu_device
+```
+
+### Miscellaneous
+
+```@docs
+reset_gpu_device!
 supported_gpu_backends
 ```

ext/LuxDeviceUtilsLuxAMDGPUExt.jl

@@ -4,10 +4,7 @@ using ChainRulesCore, LuxAMDGPU, LuxDeviceUtils, Random
 import Adapt: adapt_storage, adapt
 import ChainRulesCore as CRC
 
-function __init__()
-    LuxDeviceUtils.ACCELERATOR_STATE_CHANGED[] = true
-    return
-end
+__init__() = reset_gpu_device!()
 
 # Device Transfer
 ## To GPU

ext/LuxDeviceUtilsLuxCUDAExt.jl

@@ -4,10 +4,7 @@ using ChainRulesCore, LuxCUDA, LuxDeviceUtils, Random
 import Adapt: adapt_storage, adapt
 import ChainRulesCore as CRC
 
-function __init__()
-    LuxDeviceUtils.ACCELERATOR_STATE_CHANGED[] = true
-    return
-end
+__init__() = reset_gpu_device!()
 
 # Device Transfer
 ## To GPU

ext/LuxDeviceUtilsMetalExt.jl

@@ -4,10 +4,7 @@ using ChainRulesCore, LuxDeviceUtils, Metal, Random
 import Adapt: adapt_storage, adapt
 import ChainRulesCore as CRC
 
-function __init__()
-    LuxDeviceUtils.ACCELERATOR_STATE_CHANGED[] = true
-    return
-end
+__init__() = reset_gpu_device!()
 
 # Device Transfer
 ## To GPU

src/LuxDeviceUtils.jl

@@ -1,6 +1,6 @@
 module LuxDeviceUtils
 
-using Functors, LuxCore, Preferences, Random, SparseArrays
+using ChainRulesCore, Functors, LuxCore, Preferences, Random, SparseArrays
 import Adapt: adapt, adapt_storage
 import Base: PkgId, UUID
 
@@ -9,12 +9,10 @@ function __init__()
     @require_extensions
 end
 
-export gpu_backend!, supported_gpu_backends
+export gpu_backend!, supported_gpu_backends, reset_gpu_device!
 export gpu_device, cpu_device, LuxCPUDevice, LuxCUDADevice, LuxAMDGPUDevice, LuxMetalDevice
 export LuxCPUAdaptor, LuxCUDAAdaptor, LuxAMDGPUAdaptor, LuxMetalAdaptor
 
-const ACCELERATOR_STATE_CHANGED = Ref{Bool}(false)
-
 abstract type AbstractLuxDevice <: Function end
 abstract type AbstractLuxGPUDevice <: AbstractLuxDevice end
 
@@ -58,6 +56,16 @@ const GPU_DEVICES = (LuxCUDADevice(), LuxAMDGPUDevice(), LuxMetalDevice())
 
 const GPU_DEVICE = Ref{Union{Nothing, AbstractLuxDevice}}(nothing)
 
+"""
+    reset_gpu_device!()
+
+Resets the selected GPU device. This is useful when automatic GPU selection needs to be
+run again.
+"""
+function reset_gpu_device!()
+    return GPU_DEVICE[] = nothing
+end
+
 """
     supported_gpu_backends() -> Tuple{String, ...}
 
@@ -85,17 +93,14 @@ Selects GPU device based on the following criteria:
  4. If nothing works, an error is thrown.
 """
 function gpu_device(; force_gpu_usage::Bool=false)::AbstractLuxDevice
-    if !ACCELERATOR_STATE_CHANGED[]
-        if GPU_DEVICE[] !== nothing
-            force_gpu_usage &&
-                !(GPU_DEVICE[] isa AbstractLuxGPUDevice) &&
-                throw(LuxDeviceSelectionException())
-            return GPU_DEVICE[]
-        end
+    if GPU_DEVICE[] !== nothing
+        force_gpu_usage &&
+            !(GPU_DEVICE[] isa AbstractLuxGPUDevice) &&
+            throw(LuxDeviceSelectionException())
+        return GPU_DEVICE[]
     end
 
     device = _get_gpu_device(; force_gpu_usage)
-    ACCELERATOR_STATE_CHANGED[] = false
     GPU_DEVICE[] = device
 
     return device

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

test/amdgpu.jl

@@ -10,7 +10,7 @@ end
 using LuxAMDGPU
 
 @testset "Loaded Trigger Package" begin
-    @test LuxDeviceUtils.ACCELERATOR_STATE_CHANGED[]
+    @test LuxDeviceUtils.GPU_DEVICE[] === nothing
 
     if LuxAMDGPU.functional()
         @info "LuxAMDGPU is functional"
@@ -22,7 +22,7 @@ using LuxAMDGPU
         @test_throws LuxDeviceUtils.LuxDeviceSelectionException gpu_device(;
             force_gpu_usage=true)
     end
-    @test !LuxDeviceUtils.ACCELERATOR_STATE_CHANGED[]
+    @test LuxDeviceUtils.GPU_DEVICE[] !== nothing
 end
 
 using FillArrays, Zygote  # Extensions

test/cuda.jl

@@ -10,7 +10,7 @@ end
 using LuxCUDA
 
 @testset "Loaded Trigger Package" begin
-    @test LuxDeviceUtils.ACCELERATOR_STATE_CHANGED[]
+    @test LuxDeviceUtils.GPU_DEVICE[] === nothing
 
     if LuxCUDA.functional()
         @info "LuxCUDA is functional"
@@ -22,7 +22,7 @@ using LuxCUDA
         @test_throws LuxDeviceUtils.LuxDeviceSelectionException gpu_device(;
             force_gpu_usage=true)
     end
-    @test !LuxDeviceUtils.ACCELERATOR_STATE_CHANGED[]
+    @test LuxDeviceUtils.GPU_DEVICE[] !== nothing
 end
 
 using FillArrays, Zygote  # Extensions

test/metal.jl

@@ -10,7 +10,7 @@ end
 using Metal
 
 @testset "Loaded Trigger Package" begin
-    @test LuxDeviceUtils.ACCELERATOR_STATE_CHANGED[]
+    @test LuxDeviceUtils.GPU_DEVICE[] === nothing
 
     if Metal.functional()
         @info "Metal is functional"
@@ -22,7 +22,7 @@ using Metal
         @test_throws LuxDeviceUtils.LuxDeviceSelectionException gpu_device(;
             force_gpu_usage=true)
     end
-    @test !LuxDeviceUtils.ACCELERATOR_STATE_CHANGED[]
+    @test LuxDeviceUtils.GPU_DEVICE[] !== nothing
 end
 
 using FillArrays, Zygote  # Extensions

test/runtests.jl

@@ -1,4 +1,4 @@
-using SafeTestsets, Test, Pkg
+using Aqua, SafeTestsets, Test, Pkg
 using LuxCore, LuxDeviceUtils
 
 const GROUP = get(ENV, "GROUP", "CUDA")
@@ -41,4 +41,8 @@ end
             end
         end
     end
+
+    @testset "Aqua Tests" begin
+        Aqua.test_all(LuxDeviceUtils; piracy=false)
+    end
 end