Initial implementation of AMDGPU extension

Author: pxl-th

Project.toml

@@ -23,8 +23,15 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
+[weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+
+[extensions]
+AMDGPUExt = "AMDGPU"
+
 [compat]
 Adapt = "3.0"
+AMDGPU = "0.4.8"
 CUDA = "3, 4"
 ChainRulesCore = "1.12"
 Functors = "0.3, 0.4"

ext/AMDGPUExt/AMDGPUExt.jl

@@ -0,0 +1,20 @@
+module AMDGPUExt
+
+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)
+
+include("functor.jl")
+
+function __init__()
+    Flux.amdgpu_loaded[] = true
+end
+
+end

ext/AMDGPUExt/functor.jl

@@ -0,0 +1,51 @@
+struct FluxAMDGPUAdaptor end
+
+adapt_storage(::FluxAMDGPUAdaptor, x) = ROCArray(x)
+adapt_storage(::FluxAMDGPUAdaptor, x::Zygote.FillArrays.AbstractFill) =
+    ROCArray(collect(x))
+adapt_storage(::FluxAMDGPUAdaptor, x::Zygote.OneElement) = ROCArray(collect(x))
+adapt_storage(::FluxAMDGPUAdaptor, x::Random.TaskLocalRNG) =
+    AMDGPU.rocRAND.default_rng()
+adapt_storage(::FluxAMDGPUAdaptor, x::AMDGPU.rocRAND.RNG) = x
+adapt_storage(::FluxAMDGPUAdaptor, x::AbstractRNG) = error("""
+    Cannot map RNG of type $(typeof(x)) to AMDGPU.
+    AMDGPU execution only supports Random.default_rng().""")
+
+# TODO adaptor for Conv
+
+adapt_storage(::FluxCPUAdaptor, x::AMDGPU.rocRAND.RNG) = Random.default_rng()
+
+function ChainRulesCore.rrule(::Type{Array}, x::ROCArray)
+    Array(x), dx -> (NoTangent(), ROCArray(unthunk(dx)))
+end
+
+function ChainRulesCore.rrule(
+    ::typeof(Adapt.adapt_storage), to::FluxCPUAdaptor, x::AMDGPU.AnyROCArray,
+)
+    adapt_storage(to, x), dx -> (
+        NoTangent(), NoTangent(),
+        adapt_storage(FluxAMDGPUAdaptor(), 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
+end
+ChainRulesCore.@non_differentiable check_use_amdgpu()

src/functor.jl

@@ -280,3 +280,19 @@ f16(m) = _paramtype(Float16, m)
 @functor Cholesky
 trainable(c::Cholesky) = ()
 
+# AMDGPU extension.
+
+const amdgpu_loaded = Ref{Bool}(false)
+
+function amd(x)
+    if amdgpu_loaded[]
+        return _amd(x)
+    else
+        @info """
+        The AMDGPU functionality is being called via `Flux.amd` but
+        `AMDGPU` must be loaded to access it.
+        """ maxlog=1
+    end
+end
+
+function _amd end