Merge pull request #23 from LuxDL/ap/nested_structures

Handle nested array structures nicely
LuxDL · Jan 12, 2024 · 0fe0576 · 0fe0576 · avik-pal · Jan 12, 2024
2 parents 0c761ba + aa4d84c
commit 0fe0576
Show file tree

Hide file tree

Showing 7 changed files with 68 additions and 16 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxDeviceUtils"
 uuid = "34f89e08-e1d5-43b4-8944-0b49ac560553"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.1.12"
+version = "0.1.13"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -14,36 +14,44 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
 [weakdeps]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 LuxAMDGPU = "83120cb1-ca15-4f04-bf3b-6967d2e6b60b"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 LuxDeviceUtilsFillArraysExt = "FillArrays"
+LuxDeviceUtilsGPUArraysExt = "GPUArrays"
 LuxDeviceUtilsLuxAMDGPUExt = "LuxAMDGPU"
 LuxDeviceUtilsLuxCUDAExt = "LuxCUDA"
-LuxDeviceUtilsMetalExt = "Metal"
+LuxDeviceUtilsMetalGPUArraysExt = ["GPUArrays", "Metal"]
+LuxDeviceUtilsRecursiveArrayToolsExt = "RecursiveArrayTools"
 LuxDeviceUtilsZygoteExt = "Zygote"
 
 [compat]
 Adapt = "3, 4"
 ChainRulesCore = "1"
 FillArrays = "0.13, 1"
 Functors = "0.2, 0.3, 0.4"
+GPUArrays = "9, 10"
 LuxAMDGPU = "0.1, 0.2"
 LuxCUDA = "0.2, 0.3"
 LuxCore = "0.1.4"
 Metal = "0.4, 0.5"
 Preferences = "1"
 Random = "<0.0.1, 1"
+RecursiveArrayTools = "3"
 SparseArrays = "<0.0.1, 1"
 Zygote = "0.6"
 julia = "1.9"
 
 [extras]
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
+GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 LuxAMDGPU = "83120cb1-ca15-4f04-bf3b-6967d2e6b60b"
 LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
 Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
diff --git a/ext/LuxDeviceUtilsGPUArraysExt.jl b/ext/LuxDeviceUtilsGPUArraysExt.jl
@@ -0,0 +1,8 @@
+module LuxDeviceUtilsGPUArraysExt
+
+using GPUArrays, LuxDeviceUtils, Random
+import Adapt: adapt_storage, adapt
+
+adapt_storage(::LuxCPUAdaptor, rng::GPUArrays.RNG) = Random.default_rng()
+
+end
diff --git a/ext/LuxDeviceUtilsLuxAMDGPUExt.jl b/ext/LuxDeviceUtilsLuxAMDGPUExt.jl
@@ -9,13 +9,15 @@ __init__() = reset_gpu_device!()
 LuxDeviceUtils.__is_loaded(::LuxAMDGPUDevice) = true
 LuxDeviceUtils.__is_functional(::LuxAMDGPUDevice) = LuxAMDGPU.functional()
 
+# Default RNG
+LuxDeviceUtils.default_device_rng(::LuxAMDGPUDevice) = AMDGPU.rocrand_rng()
+
 # Device Transfer
 ## To GPU
 adapt_storage(::LuxAMDGPUAdaptor, x) = roc(x)
 adapt_storage(::LuxAMDGPUAdaptor, rng::AbstractRNG) = rng
-adapt_storage(::LuxAMDGPUAdaptor, rng::Random.TaskLocalRNG) = AMDGPU.rocRAND.RNG()
+adapt_storage(::LuxAMDGPUAdaptor, rng::Random.TaskLocalRNG) = AMDGPU.rocrand_rng()
 
-## Is this a correct thing to do?
 adapt_storage(::LuxCPUAdaptor, rng::AMDGPU.rocRAND.RNG) = Random.default_rng()
 
 ## Chain Rules

diff --git a/ext/LuxDeviceUtilsLuxCUDAExt.jl b/ext/LuxDeviceUtilsLuxCUDAExt.jl
@@ -9,13 +9,15 @@ __init__() = reset_gpu_device!()
 LuxDeviceUtils.__is_loaded(::LuxCUDADevice) = true
 LuxDeviceUtils.__is_functional(::LuxCUDADevice) = LuxCUDA.functional()
 
+# Default RNG
+LuxDeviceUtils.default_device_rng(::LuxCUDADevice) = CUDA.default_rng()
+
 # Device Transfer
 ## To GPU
 adapt_storage(::LuxCUDAAdaptor, x) = cu(x)
 adapt_storage(::LuxCUDAAdaptor, rng::AbstractRNG) = rng
 adapt_storage(::LuxCUDAAdaptor, rng::Random.TaskLocalRNG) = CUDA.default_rng()
 
-## Is this a correct thing to do?
 adapt_storage(::LuxCPUAdaptor, rng::CUDA.RNG) = Random.default_rng()
 
 ## To CPU

diff --git a/ext/LuxDeviceUtilsMetalExt.jl → ext/LuxDeviceUtilsMetalGPUArraysExt.jl b/ext/LuxDeviceUtilsMetalExt.jl → ext/LuxDeviceUtilsMetalGPUArraysExt.jl
@@ -1,6 +1,6 @@
-module LuxDeviceUtilsMetalExt
+module LuxDeviceUtilsMetalGPUArraysExt
 
-using ChainRulesCore, LuxDeviceUtils, Metal, Random
+using ChainRulesCore, GPUArrays, LuxDeviceUtils, Metal, Random
 import Adapt: adapt_storage, adapt
 import ChainRulesCore as CRC
 
@@ -9,16 +9,14 @@ __init__() = reset_gpu_device!()
 LuxDeviceUtils.__is_loaded(::LuxMetalDevice) = true
 LuxDeviceUtils.__is_functional(::LuxMetalDevice) = Metal.functional()
 
-__default_rng() = Metal.GPUArrays.default_rng(MtlArray)
+# Default RNG
+LuxDeviceUtils.default_device_rng(::LuxMetalDevice) = GPUArrays.default_rng(MtlArray)
 
 # Device Transfer
 ## To GPU
 adapt_storage(::LuxMetalAdaptor, x) = mtl(x)
 adapt_storage(::LuxMetalAdaptor, rng::AbstractRNG) = rng
-adapt_storage(::LuxMetalAdaptor, rng::Random.TaskLocalRNG) = __default_rng()
-
-## Is this a correct thing to do?
-adapt_storage(::LuxCPUAdaptor, rng::Metal.GPUArrays.RNG) = Random.default_rng()
+adapt_storage(::LuxMetalAdaptor, rng::Random.TaskLocalRNG) = GPUArrays.default_rng(MtlArray)
 
 ## Chain Rules
 CRC.rrule(::Type{Array}, x::MtlArray) = Array(x), Δ -> (NoTangent(), MtlArray(Δ))

diff --git a/ext/LuxDeviceUtilsRecursiveArrayToolsExt.jl b/ext/LuxDeviceUtilsRecursiveArrayToolsExt.jl
@@ -0,0 +1,17 @@
+module LuxDeviceUtilsRecursiveArrayToolsExt
+
+using Adapt, LuxDeviceUtils, RecursiveArrayTools
+
+# We want to preserve the structure
+function Adapt.adapt_structure(to::LuxDeviceUtils.AbstractLuxDeviceAdaptor,
+        x::VectorOfArray)
+    return VectorOfArray(map(Base.Fix1(adapt, to), x.u))
+end
+
+function Adapt.adapt_structure(to::LuxDeviceUtils.AbstractLuxDeviceAdaptor,
+        x::DiffEqArray)
+    # Don't move the `time` to the GPU
+    return DiffEqArray(map(Base.Fix1(adapt, to), x.u), x.t)
+end
+
+end
diff --git a/src/LuxDeviceUtils.jl b/src/LuxDeviceUtils.jl
@@ -4,6 +4,7 @@ using ChainRulesCore, Functors, LuxCore, Preferences, Random, SparseArrays
 import Adapt: adapt, adapt_storage
 
 export gpu_backend!, supported_gpu_backends, reset_gpu_device!
+export default_device_rng
 export gpu_device, cpu_device, LuxCPUDevice, LuxCUDADevice, LuxAMDGPUDevice, LuxMetalDevice
 export LuxCPUAdaptor, LuxCUDAAdaptor, LuxAMDGPUAdaptor, LuxMetalAdaptor
 
@@ -207,6 +208,22 @@ Return a `LuxCPUDevice` object which can be used to transfer data to CPU.
 """
 @inline cpu_device() = LuxCPUDevice()
 
+"""
+    default_device_rng(::AbstractLuxDevice)
+
+Returns the default RNG for the device. This can be used to directly generate parameters
+and states on the device using
+[WeightInitializers.jl](https://github.com/LuxDL/WeightInitializers.jl).
+"""
+function default_device_rng(D::AbstractLuxDevice)
+    return error("""`default_device_rng` not implemented for $(typeof(D)). This is either because:
+
+           1. The default RNG for this device is not known / officially provided.
+           2. The trigger package for the device is not loaded.
+           """)
+end
+default_device_rng(::LuxCPUDevice) = Random.default_rng()
+
 # Dispatches for Different Data Structures
 # Abstract Array / Tuples / NamedTuples have special fast paths to facilitate type stability
 # For all other types we rely on fmap which means we lose type stability.
@@ -215,12 +232,12 @@ for (dev) in (:CPU, :CUDA, :AMDGPU, :Metal)
     ldev = Symbol("Lux$(dev)Device")
     ladaptor = Symbol("Lux$(dev)Adaptor")
     @eval begin
-        function (::$(ldev))(x::AbstractArray)
+        function (D::$(ldev))(x::AbstractArray)
             fn = Base.Fix1(adapt, $(ladaptor)())
-            return _isbitsarray(x) ? fn(x) : map(fn, x)
+            return _isbitsarray(x) ? fn(x) : map(D, x)
         end
-        (::$(ldev))(x::Tuple) = map(Base.Fix1(adapt, $(ladaptor)()), x)
-        (dev::$(ldev))(x::NamedTuple{F}) where {F} = NamedTuple{F}(dev(values(x)))
+        (D::$(ldev))(x::Tuple) = map(D, x)
+        (D::$(ldev))(x::NamedTuple{F}) where {F} = NamedTuple{F}(D(values(x)))
         function (::$(ldev))(x)
             _isleaf(x) && return adapt($(ladaptor)(), x)
             return fmap(Base.Fix1(adapt, $(ladaptor)()), x; exclude=_isleaf)