Add gpu backend switch mechanism

Author: pxl-th

LocalPreferences.toml

@@ -0,0 +1,2 @@
+[Flux]
+gpu_backend = "AMD"

NEWS.md

@@ -2,6 +2,9 @@
 
 ## v0.13.13
 * Added `f16` which changes precision to `Float16`, recursively.
+* Initial support for AMDGPU via extension mechanism.
+* Add `gpu_backend` preference to select GPU backend using `LocalPreference.toml`.
+* Add `Flux.gpu_backend!` method to switch between GPU backends.
 
 ## v0.13.12
 * CUDA.jl 4.0 compatibility.

Project.toml

@@ -14,6 +14,7 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NNlibCUDA = "a00861dc-f156-4864-bf3c-e6376f28a68d"
 OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -30,8 +31,8 @@ AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 AMDGPUExt = "AMDGPU"
 
 [compat]
-Adapt = "3.0"
 AMDGPU = "0.4.8"
+Adapt = "3.0"
 CUDA = "3, 4"
 ChainRulesCore = "1.12"
 Functors = "0.3, 0.4"
@@ -57,4 +58,4 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AMDGPU", "Test", "Documenter", "IterTools", "LinearAlgebra", "FillArrays", "ComponentArrays"]
+test = ["Test", "Documenter", "IterTools", "LinearAlgebra", "FillArrays", "ComponentArrays"]

ext/AMDGPUExt/AMDGPUExt.jl

@@ -3,7 +3,7 @@ module AMDGPUExt
 import ChainRulesCore
 import ChainRulesCore: NoTangent
 import Flux
-import Flux: FluxCPUAdaptor, _amd, _isleaf, adapt_storage, fmap
+import Flux: FluxCPUAdaptor, FluxAMDAdaptor, _amd, _isleaf, adapt_storage, fmap
 
 using AMDGPU
 using Adapt

ext/AMDGPUExt/functor.jl

@@ -1,5 +1,3 @@
-struct FluxAMDAdaptor end
-
 # Convert Float64 to Float32, but preserve Float16.
 adapt_storage(::FluxAMDAdaptor, x::T) where T <: AbstractArray =
     isbits(x) ? x : ROCArray(x)

src/Flux.jl

@@ -1,6 +1,7 @@
 module Flux
 
 using Base: tail
+using Preferences
 using LinearAlgebra, Statistics, Random  # standard lib
 using MacroTools, Reexport, ProgressLogging, SpecialFunctions
 using MacroTools: @forward
@@ -72,4 +73,22 @@ include("deprecations.jl")
 
 include("cuda/cuda.jl")
 
+const GPU_BACKENDS = Dict(
+    "CUDA" => FluxCUDAAdaptor(),
+    "AMD" => FluxAMDAdaptor())
+
+const GPU_BACKEND = Ref{Union{FluxCUDAAdaptor, FluxAMDAdaptor}}(
+    GPU_BACKENDS[@load_preference("gpu_backend", "CUDA")])
+
+function gpu_backend!(backend::String)
+    backend in keys(GPU_BACKENDS) || throw(ArgumentError("""
+    Unsupported GPU backend: $backend.
+    Supported backends are: $(keys(GPU_BACKENDS)).
+    """))
+
+    @set_preferences!("gpu_backend" => backend)
+    GPU_BACKEND[] = GPU_BACKENDS[@load_preference("gpu_backend")]
+    return
+end
+
 end # module

src/functor.jl

@@ -209,6 +209,10 @@ CUDA.CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}
 ```
 """
 function gpu(x)
+    gpu(GPU_BACKEND[], x)
+end
+
+function gpu(::FluxCUDAAdaptor, x)
   check_use_cuda()
   use_cuda[] ? fmap(x -> Adapt.adapt(FluxCUDAAdaptor(), x), x; exclude = _isleaf) : x
 end
@@ -282,9 +286,11 @@ trainable(c::Cholesky) = ()
 
 # AMDGPU extension.
 
+struct FluxAMDAdaptor end
+
 const AMDGPU_LOADED = Ref{Bool}(false)
 
-function amd(x)
+function gpu(::FluxAMDAdaptor, x)
     if AMDGPU_LOADED[]
         return _amd(x)
     else

test/amd/basic.jl

@@ -1,18 +1,18 @@
 @test Flux.AMDGPU_LOADED[]
 
 @testset "Basic GPU movement" begin
-    @test Flux.amd(rand(Float64, 16)) isa ROCArray{Float32, 1}
-    @test Flux.amd(rand(Float64, 16, 16)) isa ROCArray{Float32, 2}
-    @test Flux.amd(rand(Float32, 16, 16)) isa ROCArray{Float32, 2}
-    @test Flux.amd(rand(Float16, 16, 16, 16)) isa ROCArray{Float16, 3}
+    @test Flux.gpu(rand(Float64, 16)) isa ROCArray{Float32, 1}
+    @test Flux.gpu(rand(Float64, 16, 16)) isa ROCArray{Float32, 2}
+    @test Flux.gpu(rand(Float32, 16, 16)) isa ROCArray{Float32, 2}
+    @test Flux.gpu(rand(Float16, 16, 16, 16)) isa ROCArray{Float16, 3}
 
-    @test gradient(x -> sum(Flux.amd(x)), rand(Float32, 4, 4)) isa Tuple
+    @test gradient(x -> sum(Flux.gpu(x)), rand(Float32, 4, 4)) isa Tuple
     @test gradient(x -> sum(cpu(x)), AMDGPU.rand(Float32, 4, 4)) isa Tuple
 end
 
 @testset "Dense no bias" begin
-    m = Dense(3 => 4; bias=false) |> Flux.amd
-    x = zeros(Float32, 3, 4) |> Flux.amd
+    m = Dense(3 => 4; bias=false) |> Flux.gpu
+    x = zeros(Float32, 3, 4) |> Flux.gpu
     @test sum(m(x)) ≈ 0f0
     gs = gradient(m -> sum(m(x)), m)
     @test isnothing(gs[1].bias)
@@ -25,15 +25,15 @@ end
 end
 
 @testset "Convolution" begin
-    for nd in (1, 2, 3)
+    for nd in 1:3
         m = Conv(tuple(fill(2, nd)...), 3 => 4) |> f32
         x = rand(Float32, fill(10, nd)..., 3, 5)
 
         # Ensure outputs are the same.
         amdgputest(m, x; atol=1f-3, checkgrad=false)
 
         # Gradients are flipped as well.
-        md, xd = Flux.amd.((m, x))
+        md, xd = Flux.gpu.((m, x))
         gs = gradient(m -> sum(m(x)), m)
         gsd = gradient(m -> sum(m(xd)), md)
 
@@ -53,25 +53,25 @@ end
 end
 
 @testset "Restructure" begin
-    m = Dense(1, 1) |> Flux.amd
+    m = Dense(1, 1) |> Flux.gpu
     θ, m̂ = Flux.destructure(m)
     foo(x) = sum(re(p)(x))
 
-    x = Flux.amd(rand(Float32, 1))
+    x = Flux.gpu(rand(Float32, 1))
     @test gradient(x -> sum(m̂(θ)(x)), x)[1] isa ROCArray{Float32}
 end
 
-@testset "Flux.amd(x) on structured arrays" begin
+@testset "Flux.gpu(x) on structured arrays" begin
     g1 = Zygote.OneElement(1, (2, 3), axes(ones(4, 5)))
-    @test Flux.amd(g1) isa ROCMatrix{Int64}
+    @test Flux.gpu(g1) isa ROCMatrix{Int64}
     g2 = Zygote.Fill(1f0, 2)
-    @test Flux.amd(g2) isa ROCArray{Float32, 1}
+    @test Flux.gpu(g2) isa ROCArray{Float32, 1}
     g3 = transpose(Float32[1 2; 3 4])
-    @test parent(Flux.amd(g3)) isa ROCMatrix{Float32}
+    @test parent(Flux.gpu(g3)) isa ROCMatrix{Float32}
 end
 
 @testset "Flux.onecold gpu" begin
-    y = Flux.onehotbatch(ones(3), 1:10) |> Flux.amd
+    y = Flux.onehotbatch(ones(3), 1:10) |> Flux.gpu
     l = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']
     @test Flux.onecold(y) isa ROCArray
     @test y[3, :] isa ROCArray
@@ -80,13 +80,15 @@ end
 
 @testset "Batchnorm" begin
     bn = BatchNorm(3, σ)
-    x = rand(Float32, 16, 16, 3, 4)
-    amdgputest(bn, x; atol=1f-3)
+    for nd in 1:3
+        x = rand(Float32, fill(16, nd - 1)..., 3, 4)
+        amdgputest(bn, x; atol=1f-3)
+    end
 end
 
 # FIXME scalar indexing. Needs NNlib.scatter?
 # @testset "Flux.onehot gpu" begin
-#     y = Flux.onehotbatch(ones(3), 1:2) |> Flux.amd
-#     x = rand(3, 2) |> Flux.amd
+#     y = Flux.onehotbatch(ones(3), 1:2) |> Flux.gpu
+#     x = rand(3, 2) |> Flux.gpu
 #     @test gradient(x -> sum(x * y), x)[1] isa ROCArray
 # end

test/amd/runtests.jl

@@ -1,3 +1,5 @@
+Flux.gpu_backend!("AMD")
+
 include("utils.jl")
 
 AMDGPU.allowscalar(false)

test/amd/utils.jl

@@ -1,9 +1,9 @@
 function amdgputest(model, xs...; checkgrad=true, atol=1e-6)
     cpu_model = model
-    gpu_model = Flux.amd(model)
+    gpu_model = Flux.gpu(model)
 
     cpu_in = xs
-    gpu_in = Flux.amd.(xs)
+    gpu_in = Flux.gpu.(xs)
 
     cpu_out = cpu_model(cpu_in...)
     gpu_out = gpu_model(gpu_in...)