Adding device objects for selecting GPU backends (and defaulting to CPU if none exists).

ext/FluxAMDGPUExt/FluxAMDGPUExt.jl

@@ -17,6 +17,9 @@ const MIOPENFloat = AMDGPU.MIOpen.MIOPENFloat
 # Set to boolean on the first call to check_use_amdgpu
 const USE_AMDGPU = Ref{Union{Nothing, Bool}}(nothing)
 
+Flux.isavailable(device::Flux.FluxAMDDevice) = true
+Flux.isfunctional(device::Flux.FluxAMDDevice) = AMDGPU.functional()
+
 function check_use_amdgpu()
     if !isnothing(USE_AMDGPU[])
         return

ext/FluxCUDAExt/FluxCUDAExt.jl

@@ -14,6 +14,9 @@ import Adapt: adapt_storage
 
 const USE_CUDA = Ref{Union{Nothing, Bool}}(nothing)
 
+Flux.isavailable(device::Flux.FluxCUDADevice) = true
+Flux.isfunctional(device::Flux.FluxCUDADevice) = CUDA.functional()
+
 function check_use_cuda()
     if !isnothing(USE_CUDA[])
         return

ext/FluxMetalExt/FluxMetalExt.jl

@@ -12,6 +12,9 @@ using Zygote
 
 const USE_METAL = Ref{Union{Nothing, Bool}}(nothing)
 
+Flux.isavailable(device::Flux.FluxMetalDevice) = true
+Flux.isfunctional(device::Flux.FluxMetalDevice) = Metal.functional()
+
 function check_use_metal()
     isnothing(USE_METAL[]) || return
 

src/functor.jl

@@ -187,7 +187,7 @@ _isbitsarray(x) = false
 _isleaf(::AbstractRNG) = true
 _isleaf(x) = _isbitsarray(x) || Functors.isleaf(x)
 
-const GPU_BACKENDS = ("CUDA", "AMD", "Metal")
+const GPU_BACKENDS = ("CUDA", "AMD", "Metal", "CPU")
 const GPU_BACKEND = @load_preference("gpu_backend", "CUDA")
 
 function gpu_backend!(backend::String)
@@ -249,6 +249,8 @@ function gpu(x)
         gpu(FluxAMDAdaptor(), x)
     elseif GPU_BACKEND == "Metal"
         gpu(FluxMetalAdaptor(), x)
+    elseif GPU_BACKEND == "CPU"
+        cpu(x)
     else
         error("""
         Unsupported GPU backend: $GPU_BACKEND.
@@ -444,3 +446,205 @@ function gpu(d::MLUtils.DataLoader)
     d.rng,
   )
 end
+
+"""
+    Flux.AbstractDevice <: Function
+
+An abstract type representing `device` objects for different GPU backends. The currently supported backends are `"CUDA"`, `"AMD"`, `"Metal"` and `"CPU"`; the `"CPU"` backend is the fallback case when no GPU is available.
+"""
+abstract type AbstractDevice <: Function end
+
+"""
+    Flux.FluxCPUDevice <: Flux.AbstractDevice
+
+A type representing `device` objects for the `"CPU"` backend for Flux. This is the fallback case when no GPU is available to Flux.
+"""
+Base.@kwdef struct FluxCPUDevice <: AbstractDevice
+	name::String = "CPU"
+end
+
+"""
+    Flux.FluxCUDADevice <: Flux.AbstractDevice
+
+A type representing `device` objects for the `"CUDA"` backend for Flux.
+"""
+Base.@kwdef struct FluxCUDADevice <: AbstractDevice
+    name::String = "CUDA"
+end
+
+"""
+    Flux.FluxAMDDevice <: Flux.AbstractDevice
+
+A type representing `device` objects for the `"AMD"` backend for Flux.
+"""
+Base.@kwdef struct FluxAMDDevice <: AbstractDevice
+    name::String = "AMD"
+end
+
+"""
+    Flux.FluxMetalDevice <: Flux.AbstractDevice
+
+A type representing `device` objects for the `"Metal"` backend for Flux.
+"""
+Base.@kwdef struct FluxMetalDevice <: AbstractDevice
+    name::String = "Metal"
+end
+
+(::FluxCPUDevice)(x) = cpu(x)
+(::FluxCUDADevice)(x) = gpu(FluxCUDAAdaptor(), x)
+(::FluxAMDDevice)(x) = gpu(FluxAMDAdaptor(), x)
+(::FluxMetalDevice)(x) = gpu(FluxMetalAdaptor(), x)
+
+# Applying device to DataLoader
+function _apply_to_dataloader(device::T, d::MLUtils.DataLoader) where {T <: AbstractDevice}
+    MLUtils.DataLoader(MLUtils.mapobs(device, d.data),
+        d.batchsize,
+        d.buffer,
+        d.partial,
+        d.shuffle,
+        d.parallel,
+        d.collate,
+        d.rng,
+    )
+end
+
+(device::FluxCPUDevice)(d::MLUtils.DataLoader) = _apply_to_dataloader(device, d)
+(device::FluxCUDADevice)(d::MLUtils.DataLoader) = _apply_to_dataloader(device, d)
+(device::FluxAMDDevice)(d::MLUtils.DataLoader) = _apply_to_dataloader(device, d)
+(device::FluxMetalDevice)(d::MLUtils.DataLoader) = _apply_to_dataloader(device, d)
+function _get_device_name(t::T) where {T <: AbstractDevice}
+    return hasfield(T, :name) ? t.name : ""
+end
+
+## check device availability; more definitions in corresponding extensions
+isavailable(device::AbstractDevice) = false
+isfunctional(device::AbstractDevice) = false
+
+isavailable(device::FluxCPUDevice) = true
+isfunctional(device::FluxCPUDevice) = true
+
+# below order is important
+const DEVICES = (FluxCUDADevice(), FluxAMDDevice(), FluxMetalDevice(), FluxCPUDevice())
+
+"""
+    Flux.supported_devices()
+
+Get all supported backends for Flux, in order of preference.
+
+# Example
+
+```jldoctest
+julia> using Flux;
+
+julia> Flux.supported_devices()
+("CUDA", "AMD", "Metal", "CPU")
+```
+"""
+supported_devices() = map(_get_device_name, DEVICES)
+
+"""
+    Flux.get_device()::AbstractDevice
+
+Returns a `device` object for the most appropriate backend for the current Julia session. 
+
+First, the function checks whether a backend preference has been set via the `gpu_backend!` function. If so, an attempt is made to load this backend. If the corresponding trigger package has been loaded and the backend is functional, a `device` corresponding to the given backend is loaded. Otherwise, the backend is chosen automatically. To update the backend preference, use [gpu_backend!](@ref).
+
+If there is no preference, then for each of the `"CUDA"`, `"AMD"`, `"Metal"` and `"CPU"` backends in the given order, this function checks whether the given backend has been loaded via the corresponding trigger package, and whether the backend is functional. If so, the `device` corresponding to the backend is returned. If no GPU backend is available, a `Flux.FluxCPUDevice` is returned.
+
+# Examples
+For the example given below, the backend preference was set to `"AMD"` via the [`gpu_backend!`](@ref) function.
+
+```julia-repl
+julia> using Flux;
+
+julia> model = Dense(2 => 3)
+Dense(2 => 3)       # 9 parameters
+
+julia> device = Flux.get_device()       # this will just load the CPU device
+[ Info: Using backend set in preferences: AMD.
+┌ Warning: Trying to use backend AMD but package AMDGPU [21141c5a-9bdb-4563-92ae-f87d6854732e] is not loaded.
+│ Please load the package and call this function again to respect the preferences backend.
+└ @ Flux ~/fluxml/Flux.jl/src/functor.jl:496
+[ Info: Running automatic device selection...
+(::Flux.FluxCPUDevice) (generic function with 1 method)
+
+julia> model = model |> device
+Dense(2 => 3)       # 9 parameters
+
+julia> model.weight
+3×2 Matrix{Float32}:
+ -0.304362  -0.700477
+ -0.861201   0.67825
+ -0.176017   0.234188
+```
+
+Here is the same example, but using `"CUDA"`:
+
+```julia-repl
+julia> using Flux, CUDA;
+
+julia> model = Dense(2 => 3)
+Dense(2 => 3)       # 9 parameters
+
+julia> device = Flux.get_device()
+[ Info: Using backend set in preferences: AMD.
+┌ Warning: Trying to use backend AMD but package AMDGPU [21141c5a-9bdb-4563-92ae-f87d6854732e] is not loaded.
+│ Please load the package and call this function again to respect the preferences backend.
+└ @ Flux ~/fluxml/Flux.jl/src/functor.jl:496
+[ Info: Running automatic device selection...
+(::Flux.FluxCUDADevice) (generic function with 1 method)
+
+julia> model = model |> device
+Dense(2 => 3)       # 9 parameters
+
+julia> model.weight
+3×2 CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}:
+  0.820013   0.527131
+ -0.915589   0.549048
+  0.290744  -0.0592499
+```
+"""
+function get_device()::AbstractDevice
+    backend = @load_preference("gpu_backend", nothing) 
+    if backend !== nothing
+        allowed_backends = supported_devices()
+        idx = findfirst(isequal(backend), allowed_backends)
+        if backend ∉  allowed_backends
+            @warn """
+                `gpu_backend` preference is set to $backend, which is not allowed.
+                Defaulting to automatic device selection.
+            """ maxlog=1
+        else
+            @info "Using backend set in preferences: $backend."
+            device = DEVICES[idx] 
+
+            if !isavailable(device)
+                @warn """
+                Trying to use backend $(_get_device_name(device)) but it's trigger package is not loaded.
+                Please load the package and call this function again to respect the preferences backend.
+                """ maxlog=1
+            else 
+                if isfunctional(device)
+                    @info "Using backend: $(_get_device_name(device))"
+                    return device
+                else
+                    @warn "Backend: $(_get_device_name(device)) from the set preferences is not functional. Defaulting to autmatic device selection." maxlog=1
+                end
+            end
+        end
+    end
+
+    @info "Running automatic device selection..."
+    for device in DEVICES 
+        if isavailable(device)
+            @info "Trying backend: $(_get_device_name(device))."
+            if isfunctional(device)
+                @debug "Using backend: $(_get_device_name(device))."
+                return device
+            end
+            @info "Backend: $(_get_device_name(device)) is not functional."
+        else
+            @info "Trigger package for backend ($(_get_device_name(device))) is not loaded."
+        end
+    end
+end