Add more tests

Author: pxl-th

Project.toml

@@ -57,4 +57,4 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Documenter", "IterTools", "LinearAlgebra", "FillArrays", "ComponentArrays"]
+test = ["AMDGPU", "Test", "Documenter", "IterTools", "LinearAlgebra", "FillArrays", "ComponentArrays"]

ext/AMDGPUExt/AMDGPUExt.jl

@@ -1,20 +1,42 @@
 module AMDGPUExt
 
+import ChainRulesCore
+import Flux
+import Flux: FluxCPUAdaptor, _amd, _isleaf, adapt_storage, fmap
+
 using AMDGPU
 using Adapt
 using Random
 using Zygote
-import ChainRulesCore
-import Functors: fmap
-import Flux
-import Flux: FluxCPUAdaptor, adapt_storage, _isleaf, _amd
 
-const use_amdgpu = Ref{Bool}(false)
+const USE_AMDGPU = Ref{Union{Nothing, Bool}}(nothing)
+
+function check_use_amdgpu()
+    isnothing(USE_AMDGPU[]) || return
+
+    USE_AMDGPU[] = AMDGPU.functional()
+    if USE_AMDGPU[]
+        if !AMDGPU.functional(:MIOpen)
+            @warn "MIOpen is not functional in AMDGPU.jl, some functionality will not be available."
+        end
+    else
+        @info """
+        The AMDGPU function is being called but the AMDGPU is not functional.
+        Defaulting back to the CPU. (No action is required if you want to run on the CPU).
+        """ maxlog=1
+    end
+    return
+end
+ChainRulesCore.@non_differentiable check_use_amdgpu()
 
 include("functor.jl")
 
 function __init__()
-    Flux.amdgpu_loaded[] = true
+    Flux.AMDGPU_LOADED[] = true
 end
 
+# TODO
+# fail early if input to the model is not on the device (e.g. on the host)
+# otherwise we get very cryptic errors & segfaults at the rocBLAS level
+
 end

ext/AMDGPUExt/functor.jl

@@ -1,13 +1,20 @@
-struct FluxAMDGPUAdaptor end
+struct FluxAMDAdaptor end
 
-adapt_storage(::FluxAMDGPUAdaptor, x) = ROCArray(x)
-adapt_storage(::FluxAMDGPUAdaptor, x::Zygote.FillArrays.AbstractFill) =
+# Convert Float64 to Float32, but preserve Float16.
+adapt_storage(::FluxAMDAdaptor, x::T) where T <: AbstractArray =
+    isbits(x) ? x : ROCArray(x)
+adapt_storage(::FluxAMDAdaptor, x::AbstractArray{T, N}) where {T <: AbstractFloat, N} =
+    isbits(x) ? x : ROCArray{Float32, N}(x)
+adapt_storage(::FluxAMDAdaptor, x::AbstractArray{Float16, N}) where N =
+    isbits(x) ? x : ROCArray{Float16, N}(x)
+
+adapt_storage(::FluxAMDAdaptor, x::Zygote.FillArrays.AbstractFill) =
     ROCArray(collect(x))
-adapt_storage(::FluxAMDGPUAdaptor, x::Zygote.OneElement) = ROCArray(collect(x))
-adapt_storage(::FluxAMDGPUAdaptor, x::Random.TaskLocalRNG) =
+adapt_storage(::FluxAMDAdaptor, x::Zygote.OneElement) = ROCArray(collect(x))
+adapt_storage(::FluxAMDAdaptor, x::Random.TaskLocalRNG) =
     AMDGPU.rocRAND.default_rng()
-adapt_storage(::FluxAMDGPUAdaptor, x::AMDGPU.rocRAND.RNG) = x
-adapt_storage(::FluxAMDGPUAdaptor, x::AbstractRNG) = error("""
+adapt_storage(::FluxAMDAdaptor, x::AMDGPU.rocRAND.RNG) = x
+adapt_storage(::FluxAMDAdaptor, x::AbstractRNG) = error("""
     Cannot map RNG of type $(typeof(x)) to AMDGPU.
     AMDGPU execution only supports Random.default_rng().""")
 
@@ -24,28 +31,12 @@ function ChainRulesCore.rrule(
 )
     adapt_storage(to, x), dx -> (
         NoTangent(), NoTangent(),
-        adapt_storage(FluxAMDGPUAdaptor(), unthunk(dx)))
+        adapt_storage(FluxAMDAdaptor(), unthunk(dx)))
 end
 
 function _amd(x)
     check_use_amdgpu()
-    use_amdgpu[] ? fmap(x -> Adapt.adapt(FluxAMDGPUAdaptor(), x)) : x
-end
-
-function check_use_amdgpu()
-    use_amdgpu[] === nothing || return
-
-    use_amdgpu[] = AMDGPU.functional()
-    if use_amdgpu[]
-        if !AMDGPU.functional(:MIOpen)
-            @warn "MIOpen is not functional in AMDGPU.jl, some functionality will not be available."
-        end
-    else
-        @info """
-        The AMDGPU function is being called but the AMDGPU is not functional.
-        Defaulting back to the CPU. (No action is required if you want to run on the CPU).
-        """ maxlog=1
-    end
-    return
+    USE_AMDGPU[] ?
+        fmap(x -> Adapt.adapt(FluxAMDAdaptor(), x), x; exclude=_isleaf) :
+        x
 end
-ChainRulesCore.@non_differentiable check_use_amdgpu()

src/functor.jl

@@ -282,10 +282,10 @@ trainable(c::Cholesky) = ()
 
 # AMDGPU extension.
 
-const amdgpu_loaded = Ref{Bool}(false)
+const AMDGPU_LOADED = Ref{Bool}(false)
 
 function amd(x)
-    if amdgpu_loaded[]
+    if AMDGPU_LOADED[]
         return _amd(x)
     else
         @info """

test/amd/basic.jl

@@ -0,0 +1,64 @@
+@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 gradient(x -> sum(Flux.amd(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
+    @test sum(m(x)) ≈ 0f0
+    gs = gradient(m -> sum(m(x)), m)
+    @test isnothing(gs[1].bias)
+end
+
+@testset "Chain of Dense layers" begin
+    m = Chain(Dense(10, 5, tanh), Dense(5, 2), softmax) |> f32
+    x = rand(Float32, 10, 10)
+    amdgputest(m, x)
+end
+
+@testset "Cross-correlation" begin
+    m = CrossCor((2, 2), 3 => 4) |> f32
+    x = rand(Float32, 10, 10, 3, 2)
+    amdgputest(m, x; atol=1f-3)
+end
+
+@testset "Restructure" begin
+    m = Dense(1, 1) |> Flux.amd
+    θ, m̂ = Flux.destructure(m)
+    foo(x) = sum(re(p)(x))
+
+    x = Flux.amd(rand(Float32, 1))
+    @test gradient(x -> sum(m̂(θ)(x)), x)[1] isa ROCArray{Float32}
+end
+
+@testset "Flux.amd(x) on structured arrays" begin
+    g1 = Zygote.OneElement(1, (2, 3), axes(ones(4, 5)))
+    @test Flux.amd(g1) isa ROCMatrix{Int64}
+    g2 = Zygote.Fill(1f0, 2)
+    @test Flux.amd(g2) isa ROCArray{Float32, 1}
+    g3 = transpose(Float32[1 2; 3 4])
+    @test parent(Flux.amd(g3)) isa ROCMatrix{Float32}
+end
+
+@testset "Flux.onecold gpu" begin
+    y = Flux.onehotbatch(ones(3), 1:10) |> Flux.amd
+    l = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j']
+    @test Flux.onecold(y) isa ROCArray
+    @test y[3, :] isa ROCArray
+    @test Flux.onecold(y, l) == ['a', 'a', 'a']
+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
+#     @test gradient(x -> sum(x * y), x)[1] isa ROCArray
+# end

test/amd/runtests.jl

@@ -0,0 +1,7 @@
+include("utils.jl")
+
+AMDGPU.allowscalar(false)
+
+@testset "Basic" begin
+    include("basic.jl")
+end

test/amd/utils.jl

@@ -0,0 +1,47 @@
+function amdgputest(model, xs...; checkgrad=true, atol=1e-6, kws...)
+    cpu_model = model
+    gpu_model = Flux.amd(model)
+
+    cpu_in = xs
+    gpu_in = Flux.amd.(xs)
+
+    cpu_out = cpu_model(cpu_in...)
+    gpu_out = gpu_model(gpu_in...)
+    @test collect(cpu_out) ≈ collect(gpu_out) atol=atol
+
+    if checkgrad
+        cpu_grad = gradient(m -> sum(m(cpu_in...)), cpu_model)
+        gpu_grad = gradient(m -> sum(m(gpu_in...)), gpu_model)
+        amd_check_grad(gpu_grad, cpu_grad; atol)
+    end
+end
+
+function amd_check_grad(g_gpu, g_cpu; atol)
+  @show g_gpu g_cpu
+  @test false
+end
+
+amd_check_grad(g_gpu::Base.RefValue, g_cpu::Base.RefValue, atol) =
+    amd_check_grad(g_gpu[], g_cpu[]; atol)
+amd_check_grad(g_gpu::Nothing, g_cpu::Nothing; atol) =
+    @test true
+amd_check_grad(g_gpu::Float32, g_cpu::Float32; atol) =
+    @test g_cpu ≈ g_gpu atol=atol
+amd_check_grad(g_gpu::ROCArray{Float32}, g_cpu::Array{Float32}; atol) =
+    @test g_cpu ≈ collect(g_gpu) atol=atol
+amd_check_grad(
+    g_gpu::ROCArray{Float32}, g_cpu::Zygote.FillArrays.AbstractFill; atol,
+) = @test collect(g_cpu) ≈ collect(g_gpu) atol=atol
+
+function amd_check_grad(g_gpu::Tuple, g_cpu::Tuple; atol)
+    for (v1, v2) in zip(g_gpu, g_cpu)
+        amd_check_grad(v1, v2; atol)
+    end
+end
+
+function amd_check_grad(g_gpu::NamedTuple, g_cpu::NamedTuple; atol)
+    for ((k1, v1), (k2, v2)) in zip(pairs(g_gpu), pairs(g_cpu))
+        @test k1 == k2
+        amd_check_grad(v1, v2; atol)
+    end
+end

test/cuda/cuda.jl

@@ -91,7 +91,7 @@ end
     struct SimpleBits
       field::Int32
     end
-    
+
     @test gpu((;a=ones(1))).a isa CuVector{Float32}
     @test gpu((;a=['a', 'b', 'c'])).a isa CuVector{Char}
     @test gpu((;a=[SimpleBits(1)])).a isa CuVector{SimpleBits}
@@ -167,14 +167,14 @@ end
   @test parent(gpu(g3)) isa CuArray
 
 
-  #Issue #2116  
+  #Issue #2116
   struct A2116
     x::Int
     y::Int
   end
   x = [A2116(1,1), A2116(2,2)]
-  xgpu = gpu(x) 
+  xgpu = gpu(x)
   @test xgpu isa CuVector{A2116}
-  @test cpu(xgpu) isa Vector{A2116} 
+  @test cpu(xgpu) isa Vector{A2116}
   @test cpu(gpu([CartesianIndex(1)])) isa Vector{CartesianIndex{1}}
 end

test/runtests.jl

@@ -60,4 +60,19 @@ Random.seed!(0)
       doctest(Flux)
     end
   end
+
+    if get(ENV, "FLUX_TEST_AMDGPU", "false") == "true"
+        using AMDGPU
+        AMDGPU.versioninfo()
+        if AMDGPU.functional() && AMDGPU.functional(:MIOpen)
+            @show AMDGPU.MIOpen.version()
+            @testset "AMDGPU" begin
+                include("amd/runtests.jl")
+            end
+        else
+            @info "AMDGPU.jl package is not functional. Skipping AMDGPU tests."
+        end
+    else
+        @info "Skipping AMDGPU tests, set FLUX_TEST_CUDA=true to run them."
+    end
 end