From 984cfd8125d57ed8e1f49c5f5ae7c929215b2a0e Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 23 Jul 2024 19:46:57 -0700 Subject: [PATCH 1/3] refactor!: rename round 2 to `MLDataDevices` --- Project.toml | 26 ++++++------- README.md | 12 +++--- ext/DeviceUtilsMetalExt.jl | 27 ------------- ext/DeviceUtilsReverseDiffExt.jl | 17 --------- ...AMDGPUExt.jl => MLDataDevicesAMDGPUExt.jl} | 32 ++++++++-------- ...tilsCUDAExt.jl => MLDataDevicesCUDAExt.jl} | 32 ++++++++-------- ...ysExt.jl => MLDataDevicesFillArraysExt.jl} | 2 +- ...aysExt.jl => MLDataDevicesGPUArraysExt.jl} | 2 +- ext/MLDataDevicesMetalExt.jl | 27 +++++++++++++ ...=> MLDataDevicesRecursiveArrayToolsExt.jl} | 6 +-- ext/MLDataDevicesReverseDiffExt.jl | 17 +++++++++ ...Ext.jl => MLDataDevicesSparseArraysExt.jl} | 2 +- ...ackerExt.jl => MLDataDevicesTrackerExt.jl} | 10 ++--- ...ZygoteExt.jl => MLDataDevicesZygoteExt.jl} | 2 +- ...lscuDNNExt.jl => MLDataDevicescuDNNExt.jl} | 6 +-- ...oneAPIExt.jl => MLDataDevicesoneAPIExt.jl} | 12 +++--- src/{DeviceUtils.jl => MLDataDevices.jl} | 6 +-- test/amdgpu_tests.jl | 34 ++++++++--------- test/cuda_tests.jl | 38 +++++++++---------- test/metal_tests.jl | 28 +++++++------- test/misc_tests.jl | 12 +++--- test/oneapi_tests.jl | 28 +++++++------- test/qa_tests.jl | 18 ++++----- test/runtests.jl | 2 +- 24 files changed, 199 insertions(+), 199 deletions(-) delete mode 100644 ext/DeviceUtilsMetalExt.jl delete mode 100644 ext/DeviceUtilsReverseDiffExt.jl rename ext/{DeviceUtilsAMDGPUExt.jl => MLDataDevicesAMDGPUExt.jl} (63%) rename ext/{DeviceUtilsCUDAExt.jl => MLDataDevicesCUDAExt.jl} (65%) rename ext/{DeviceUtilsFillArraysExt.jl => MLDataDevicesFillArraysExt.jl} (79%) rename ext/{DeviceUtilsGPUArraysExt.jl => MLDataDevicesGPUArraysExt.jl} (85%) create mode 100644 ext/MLDataDevicesMetalExt.jl rename ext/{DeviceUtilsRecursiveArrayToolsExt.jl => MLDataDevicesRecursiveArrayToolsExt.jl} (74%) create mode 100644 ext/MLDataDevicesReverseDiffExt.jl rename ext/{DeviceUtilsSparseArraysExt.jl => MLDataDevicesSparseArraysExt.jl} (83%) rename ext/{DeviceUtilsTrackerExt.jl => MLDataDevicesTrackerExt.jl} (59%) rename ext/{DeviceUtilsZygoteExt.jl => MLDataDevicesZygoteExt.jl} (82%) rename ext/{DeviceUtilscuDNNExt.jl => MLDataDevicescuDNNExt.jl} (77%) rename ext/{DeviceUtilsoneAPIExt.jl => MLDataDevicesoneAPIExt.jl} (71%) rename src/{DeviceUtils.jl => MLDataDevices.jl} (99%) diff --git a/Project.toml b/Project.toml index ab06f0f..d015883 100644 --- a/Project.toml +++ b/Project.toml @@ -1,4 +1,4 @@ -name = "DeviceUtils" +name = "MLDataDevices" uuid = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40" authors = ["Avik Pal and contributors"] version = "1.0.0" @@ -26,18 +26,18 @@ cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] -DeviceUtilsAMDGPUExt = "AMDGPU" -DeviceUtilsCUDAExt = "CUDA" -DeviceUtilsFillArraysExt = "FillArrays" -DeviceUtilsGPUArraysExt = "GPUArrays" -DeviceUtilsMetalExt = ["GPUArrays", "Metal"] -DeviceUtilsRecursiveArrayToolsExt = "RecursiveArrayTools" -DeviceUtilsReverseDiffExt = "ReverseDiff" -DeviceUtilsSparseArraysExt = "SparseArrays" -DeviceUtilsTrackerExt = "Tracker" -DeviceUtilsZygoteExt = "Zygote" -DeviceUtilscuDNNExt = ["CUDA", "cuDNN"] -DeviceUtilsoneAPIExt = ["GPUArrays", "oneAPI"] +MLDataDevicesAMDGPUExt = "AMDGPU" +MLDataDevicesCUDAExt = "CUDA" +MLDataDevicesFillArraysExt = "FillArrays" +MLDataDevicesGPUArraysExt = "GPUArrays" +MLDataDevicesMetalExt = ["GPUArrays", "Metal"] +MLDataDevicesRecursiveArrayToolsExt = "RecursiveArrayTools" +MLDataDevicesReverseDiffExt = "ReverseDiff" +MLDataDevicesSparseArraysExt = "SparseArrays" +MLDataDevicesTrackerExt = "Tracker" +MLDataDevicesZygoteExt = "Zygote" +MLDataDevicescuDNNExt = ["CUDA", "cuDNN"] +MLDataDevicesoneAPIExt = ["GPUArrays", "oneAPI"] [compat] AMDGPU = "0.9.6" diff --git a/README.md b/README.md index 5e4ab35..b580383 100644 --- a/README.md +++ b/README.md @@ -1,18 +1,18 @@ -# DeviceUtils +# MLDataDevices [![Join the chat at https://julialang.zulipchat.com #machine-learning](https://img.shields.io/static/v1?label=Zulip&message=chat&color=9558b2&labelColor=389826)](https://julialang.zulipchat.com/#narrow/stream/machine-learning) [![Latest Docs](https://img.shields.io/badge/docs-latest-blue.svg)](https://lux.csail.mit.edu/dev/api/Accelerator_Support/LuxDeviceUtils) [![Stable Docs](https://img.shields.io/badge/docs-stable-blue.svg)](https://lux.csail.mit.edu/stable/api/Accelerator_Support/LuxDeviceUtils) -[![CI](https://github.com/LuxDL/DeviceUtils.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/LuxDL/DeviceUtils.jl/actions/workflows/CI.yml) -[![Buildkite](https://badge.buildkite.com/b098d6387b2c69bd0ab684293ff66332047b219e1b8f9bb486.svg?branch=main)](https://buildkite.com/julialang/DeviceUtils-dot-jl) -[![codecov](https://codecov.io/gh/LuxDL/DeviceUtils.jl/branch/main/graph/badge.svg?token=1ZY0A2NPEM)](https://codecov.io/gh/LuxDL/DeviceUtils.jl) +[![CI](https://github.com/LuxDL/MLDataDevices.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/LuxDL/MLDataDevices.jl/actions/workflows/CI.yml) +[![Buildkite](https://badge.buildkite.com/b098d6387b2c69bd0ab684293ff66332047b219e1b8f9bb486.svg?branch=main)](https://buildkite.com/julialang/MLDataDevices-dot-jl) +[![codecov](https://codecov.io/gh/LuxDL/MLDataDevices.jl/branch/main/graph/badge.svg?token=1ZY0A2NPEM)](https://codecov.io/gh/LuxDL/MLDataDevices.jl) [![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl) [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac) [![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle) -`DeviceUtils.jl` is a lightweight package defining rules for transferring data across +`MLDataDevices.jl` is a lightweight package defining rules for transferring data across devices. It is used in deep learning frameworks such as [Lux.jl](https://lux.csail.mit.edu/). Currently we provide support for the following backends: @@ -24,6 +24,6 @@ Currently we provide support for the following backends: ## Updating to v1.0 - * Package was renamed from `LuxDeviceUtils.jl` to `DeviceUtils.jl`. + * Package was renamed from `LuxDeviceUtils.jl` to `MLDataDevices.jl`. * `Lux(***)Device` has been renamed to `(***)Device`. * `Lux(***)Adaptor` objects have been removed. Use `(***)Device` objects instead. diff --git a/ext/DeviceUtilsMetalExt.jl b/ext/DeviceUtilsMetalExt.jl deleted file mode 100644 index 75f605b..0000000 --- a/ext/DeviceUtilsMetalExt.jl +++ /dev/null @@ -1,27 +0,0 @@ -module DeviceUtilsMetalExt - -using Adapt: Adapt -using GPUArrays: GPUArrays -using DeviceUtils: DeviceUtils, MetalDevice, reset_gpu_device! -using Metal: Metal, MtlArray - -__init__() = reset_gpu_device!() - -DeviceUtils.loaded(::Union{MetalDevice, Type{<:MetalDevice}}) = true -function DeviceUtils.functional(::Union{MetalDevice, Type{<:MetalDevice}}) - return Metal.functional() -end - -# Default RNG -DeviceUtils.default_device_rng(::MetalDevice) = GPUArrays.default_rng(MtlArray) - -# Query Device from Array -DeviceUtils._get_device(::MtlArray) = MetalDevice() - -DeviceUtils._get_device_type(::MtlArray) = MetalDevice - -# Device Transfer -## To GPU -Adapt.adapt_storage(::MetalDevice, x::AbstractArray) = Metal.mtl(x) - -end diff --git a/ext/DeviceUtilsReverseDiffExt.jl b/ext/DeviceUtilsReverseDiffExt.jl deleted file mode 100644 index d54fd35..0000000 --- a/ext/DeviceUtilsReverseDiffExt.jl +++ /dev/null @@ -1,17 +0,0 @@ -module DeviceUtilsReverseDiffExt - -using DeviceUtils: DeviceUtils -using ReverseDiff: ReverseDiff - -for op in (:_get_device, :_get_device_type) - @eval begin - function DeviceUtils.$op(x::ReverseDiff.TrackedArray) - return DeviceUtils.$op(ReverseDiff.value(x)) - end - function DeviceUtils.$op(x::AbstractArray{<:ReverseDiff.TrackedReal}) - return DeviceUtils.$op(ReverseDiff.value.(x)) - end - end -end - -end diff --git a/ext/DeviceUtilsAMDGPUExt.jl b/ext/MLDataDevicesAMDGPUExt.jl similarity index 63% rename from ext/DeviceUtilsAMDGPUExt.jl rename to ext/MLDataDevicesAMDGPUExt.jl index ab89c04..5b008f1 100644 --- a/ext/DeviceUtilsAMDGPUExt.jl +++ b/ext/MLDataDevicesAMDGPUExt.jl @@ -2,7 +2,7 @@ module DeviceUtilsAMDGPUExt using Adapt: Adapt using AMDGPU: AMDGPU -using DeviceUtils: DeviceUtils, AMDGPUDevice, CPUDevice, reset_gpu_device! +using MLDataDevices: MLDataDevices, AMDGPUDevice, CPUDevice, reset_gpu_device! using Random: Random __init__() = reset_gpu_device!() @@ -21,16 +21,16 @@ function _check_use_amdgpu!() return end -DeviceUtils.loaded(::Union{AMDGPUDevice, <:Type{AMDGPUDevice}}) = true -function DeviceUtils.functional(::Union{AMDGPUDevice, <:Type{AMDGPUDevice}})::Bool +MLDataDevices.loaded(::Union{AMDGPUDevice, <:Type{AMDGPUDevice}}) = true +function MLDataDevices.functional(::Union{AMDGPUDevice, <:Type{AMDGPUDevice}})::Bool _check_use_amdgpu!() return USE_AMD_GPU[] end -function DeviceUtils._with_device(::Type{AMDGPUDevice}, ::Nothing) +function MLDataDevices._with_device(::Type{AMDGPUDevice}, ::Nothing) return AMDGPUDevice(nothing) end -function DeviceUtils._with_device(::Type{AMDGPUDevice}, id::Integer) +function MLDataDevices._with_device(::Type{AMDGPUDevice}, id::Integer) id > length(AMDGPU.devices()) && throw(ArgumentError("id = $id > length(AMDGPU.devices()) = $(length(AMDGPU.devices()))")) old_dev = AMDGPU.device() @@ -40,30 +40,30 @@ function DeviceUtils._with_device(::Type{AMDGPUDevice}, id::Integer) return device end -DeviceUtils._get_device_id(dev::AMDGPUDevice) = AMDGPU.device_id(dev.device) +MLDataDevices._get_device_id(dev::AMDGPUDevice) = AMDGPU.device_id(dev.device) # Default RNG -DeviceUtils.default_device_rng(::AMDGPUDevice) = AMDGPU.rocrand_rng() +MLDataDevices.default_device_rng(::AMDGPUDevice) = AMDGPU.rocrand_rng() # Query Device from Array -function DeviceUtils._get_device(x::AMDGPU.AnyROCArray) +function MLDataDevices._get_device(x::AMDGPU.AnyROCArray) parent_x = parent(x) parent_x === x && return AMDGPUDevice(AMDGPU.device(x)) - return DeviceUtils._get_device(parent_x) + return MLDataDevices._get_device(parent_x) end -DeviceUtils._get_device_type(::AMDGPU.AnyROCArray) = AMDGPUDevice +MLDataDevices._get_device_type(::AMDGPU.AnyROCArray) = AMDGPUDevice # Set Device -function DeviceUtils.set_device!(::Type{AMDGPUDevice}, dev::AMDGPU.HIPDevice) +function MLDataDevices.set_device!(::Type{AMDGPUDevice}, dev::AMDGPU.HIPDevice) return AMDGPU.device!(dev) end -function DeviceUtils.set_device!(::Type{AMDGPUDevice}, id::Integer) - return DeviceUtils.set_device!(AMDGPUDevice, AMDGPU.devices()[id]) +function MLDataDevices.set_device!(::Type{AMDGPUDevice}, id::Integer) + return MLDataDevices.set_device!(AMDGPUDevice, AMDGPU.devices()[id]) end -function DeviceUtils.set_device!(::Type{AMDGPUDevice}, ::Nothing, rank::Integer) +function MLDataDevices.set_device!(::Type{AMDGPUDevice}, ::Nothing, rank::Integer) id = mod1(rank + 1, length(AMDGPU.devices())) - return DeviceUtils.set_device!(AMDGPUDevice, id) + return MLDataDevices.set_device!(AMDGPUDevice, id) end # Device Transfer @@ -71,7 +71,7 @@ end Adapt.adapt_storage(::AMDGPUDevice{Nothing}, x::AbstractArray) = AMDGPU.roc(x) function Adapt.adapt_storage(to::AMDGPUDevice, x::AbstractArray) old_dev = AMDGPU.device() # remember the current device - dev = DeviceUtils.get_device(x) + dev = MLDataDevices.get_device(x) if !(dev isa AMDGPUDevice) AMDGPU.device!(to.device) x_new = AMDGPU.roc(x) diff --git a/ext/DeviceUtilsCUDAExt.jl b/ext/MLDataDevicesCUDAExt.jl similarity index 65% rename from ext/DeviceUtilsCUDAExt.jl rename to ext/MLDataDevicesCUDAExt.jl index f035a0c..a353b42 100644 --- a/ext/DeviceUtilsCUDAExt.jl +++ b/ext/MLDataDevicesCUDAExt.jl @@ -3,10 +3,10 @@ module DeviceUtilsCUDAExt using Adapt: Adapt using CUDA: CUDA using CUDA.CUSPARSE: AbstractCuSparseMatrix, AbstractCuSparseVector -using DeviceUtils: DeviceUtils, CUDADevice, CPUDevice +using MLDataDevices: MLDataDevices, CUDADevice, CPUDevice using Random: Random -function DeviceUtils._with_device(::Type{CUDADevice}, id::Integer) +function MLDataDevices._with_device(::Type{CUDADevice}, id::Integer) id > length(CUDA.devices()) && throw(ArgumentError("id = $id > length(CUDA.devices()) = $(length(CUDA.devices()))")) old_dev = CUDA.device() @@ -16,47 +16,47 @@ function DeviceUtils._with_device(::Type{CUDADevice}, id::Integer) return device end -function DeviceUtils._with_device(::Type{CUDADevice}, ::Nothing) +function MLDataDevices._with_device(::Type{CUDADevice}, ::Nothing) return CUDADevice(nothing) end -DeviceUtils._get_device_id(dev::CUDADevice) = CUDA.deviceid(dev.device) + 1 +MLDataDevices._get_device_id(dev::CUDADevice) = CUDA.deviceid(dev.device) + 1 # Default RNG -DeviceUtils.default_device_rng(::CUDADevice) = CUDA.default_rng() +MLDataDevices.default_device_rng(::CUDADevice) = CUDA.default_rng() # Query Device from Array -function DeviceUtils._get_device(x::CUDA.AnyCuArray) +function MLDataDevices._get_device(x::CUDA.AnyCuArray) parent_x = parent(x) parent_x === x && return CUDADevice(CUDA.device(x)) - return DeviceUtils.get_device(parent_x) + return MLDataDevices.get_device(parent_x) end -function DeviceUtils._get_device(x::CUDA.CUSPARSE.AbstractCuSparseArray) +function MLDataDevices._get_device(x::CUDA.CUSPARSE.AbstractCuSparseArray) return CUDADevice(CUDA.device(x.nzVal)) end -function DeviceUtils._get_device_type(::Union{ +function MLDataDevices._get_device_type(::Union{ <:CUDA.AnyCuArray, <:CUDA.CUSPARSE.AbstractCuSparseArray}) return CUDADevice end # Set Device -function DeviceUtils.set_device!(::Type{CUDADevice}, dev::CUDA.CuDevice) +function MLDataDevices.set_device!(::Type{CUDADevice}, dev::CUDA.CuDevice) return CUDA.device!(dev) end -function DeviceUtils.set_device!(::Type{CUDADevice}, id::Integer) - return DeviceUtils.set_device!(CUDADevice, collect(CUDA.devices())[id]) +function MLDataDevices.set_device!(::Type{CUDADevice}, id::Integer) + return MLDataDevices.set_device!(CUDADevice, collect(CUDA.devices())[id]) end -function DeviceUtils.set_device!(::Type{CUDADevice}, ::Nothing, rank::Integer) +function MLDataDevices.set_device!(::Type{CUDADevice}, ::Nothing, rank::Integer) id = mod1(rank + 1, length(CUDA.devices())) - return DeviceUtils.set_device!(CUDADevice, id) + return MLDataDevices.set_device!(CUDADevice, id) end # Device Transfer Adapt.adapt_storage(::CUDADevice{Nothing}, x::AbstractArray) = CUDA.cu(x) function Adapt.adapt_storage(to::CUDADevice, x::AbstractArray) old_dev = CUDA.device() # remember the current device - dev = DeviceUtils.get_device(x) + dev = MLDataDevices.get_device(x) if !(dev isa CUDADevice) CUDA.device!(to.device) x_new = CUDA.cu(x) @@ -84,7 +84,7 @@ Adapt.adapt_storage(::CPUDevice, rng::CUDA.RNG) = Random.default_rng() end else @warn "CUDA.CUSPARSE seems to have removed SparseArrays as a dependency. Please open \ - an issue in DeviceUtils.jl repository." + an issue in MLDataDevices.jl repository." end end diff --git a/ext/DeviceUtilsFillArraysExt.jl b/ext/MLDataDevicesFillArraysExt.jl similarity index 79% rename from ext/DeviceUtilsFillArraysExt.jl rename to ext/MLDataDevicesFillArraysExt.jl index 25a9d61..36a5d6f 100644 --- a/ext/DeviceUtilsFillArraysExt.jl +++ b/ext/MLDataDevicesFillArraysExt.jl @@ -2,7 +2,7 @@ module DeviceUtilsFillArraysExt using Adapt: Adapt using FillArrays: FillArrays, AbstractFill -using DeviceUtils: DeviceUtils, CPUDevice, AbstractDevice +using MLDataDevices: MLDataDevices, CPUDevice, AbstractDevice Adapt.adapt_structure(::CPUDevice, x::AbstractFill) = x Adapt.adapt_structure(to::AbstractDevice, x::AbstractFill) = Adapt.adapt(to, collect(x)) diff --git a/ext/DeviceUtilsGPUArraysExt.jl b/ext/MLDataDevicesGPUArraysExt.jl similarity index 85% rename from ext/DeviceUtilsGPUArraysExt.jl rename to ext/MLDataDevicesGPUArraysExt.jl index 304b3f0..328222a 100644 --- a/ext/DeviceUtilsGPUArraysExt.jl +++ b/ext/MLDataDevicesGPUArraysExt.jl @@ -2,7 +2,7 @@ module DeviceUtilsGPUArraysExt using Adapt: Adapt using GPUArrays: GPUArrays -using DeviceUtils: CPUDevice +using MLDataDevices: CPUDevice using Random: Random Adapt.adapt_storage(::CPUDevice, rng::GPUArrays.RNG) = Random.default_rng() diff --git a/ext/MLDataDevicesMetalExt.jl b/ext/MLDataDevicesMetalExt.jl new file mode 100644 index 0000000..f82d55c --- /dev/null +++ b/ext/MLDataDevicesMetalExt.jl @@ -0,0 +1,27 @@ +module DeviceUtilsMetalExt + +using Adapt: Adapt +using GPUArrays: GPUArrays +using MLDataDevices: MLDataDevices, MetalDevice, reset_gpu_device! +using Metal: Metal, MtlArray + +__init__() = reset_gpu_device!() + +MLDataDevices.loaded(::Union{MetalDevice, Type{<:MetalDevice}}) = true +function MLDataDevices.functional(::Union{MetalDevice, Type{<:MetalDevice}}) + return Metal.functional() +end + +# Default RNG +MLDataDevices.default_device_rng(::MetalDevice) = GPUArrays.default_rng(MtlArray) + +# Query Device from Array +MLDataDevices._get_device(::MtlArray) = MetalDevice() + +MLDataDevices._get_device_type(::MtlArray) = MetalDevice + +# Device Transfer +## To GPU +Adapt.adapt_storage(::MetalDevice, x::AbstractArray) = Metal.mtl(x) + +end diff --git a/ext/DeviceUtilsRecursiveArrayToolsExt.jl b/ext/MLDataDevicesRecursiveArrayToolsExt.jl similarity index 74% rename from ext/DeviceUtilsRecursiveArrayToolsExt.jl rename to ext/MLDataDevicesRecursiveArrayToolsExt.jl index abbe2a7..cc006ba 100644 --- a/ext/DeviceUtilsRecursiveArrayToolsExt.jl +++ b/ext/MLDataDevicesRecursiveArrayToolsExt.jl @@ -1,7 +1,7 @@ module DeviceUtilsRecursiveArrayToolsExt using Adapt: Adapt, adapt -using DeviceUtils: DeviceUtils, AbstractDevice +using MLDataDevices: MLDataDevices, AbstractDevice using RecursiveArrayTools: VectorOfArray, DiffEqArray # We want to preserve the structure @@ -15,9 +15,9 @@ function Adapt.adapt_structure(to::AbstractDevice, x::DiffEqArray) end for op in (:_get_device, :_get_device_type) - @eval function DeviceUtils.$op(x::Union{VectorOfArray, DiffEqArray}) + @eval function MLDataDevices.$op(x::Union{VectorOfArray, DiffEqArray}) length(x.u) == 0 && return $(op == :_get_device ? nothing : Nothing) - return mapreduce(DeviceUtils.$op, DeviceUtils.__combine_devices, x.u) + return mapreduce(MLDataDevices.$op, MLDataDevices.__combine_devices, x.u) end end diff --git a/ext/MLDataDevicesReverseDiffExt.jl b/ext/MLDataDevicesReverseDiffExt.jl new file mode 100644 index 0000000..14915d9 --- /dev/null +++ b/ext/MLDataDevicesReverseDiffExt.jl @@ -0,0 +1,17 @@ +module DeviceUtilsReverseDiffExt + +using MLDataDevices: MLDataDevices +using ReverseDiff: ReverseDiff + +for op in (:_get_device, :_get_device_type) + @eval begin + function MLDataDevices.$op(x::ReverseDiff.TrackedArray) + return MLDataDevices.$op(ReverseDiff.value(x)) + end + function MLDataDevices.$op(x::AbstractArray{<:ReverseDiff.TrackedReal}) + return MLDataDevices.$op(ReverseDiff.value.(x)) + end + end +end + +end diff --git a/ext/DeviceUtilsSparseArraysExt.jl b/ext/MLDataDevicesSparseArraysExt.jl similarity index 83% rename from ext/DeviceUtilsSparseArraysExt.jl rename to ext/MLDataDevicesSparseArraysExt.jl index 6c3c15d..1851872 100644 --- a/ext/DeviceUtilsSparseArraysExt.jl +++ b/ext/MLDataDevicesSparseArraysExt.jl @@ -1,7 +1,7 @@ module DeviceUtilsSparseArraysExt using Adapt: Adapt -using DeviceUtils: CPUDevice +using MLDataDevices: CPUDevice using SparseArrays: AbstractSparseArray Adapt.adapt_storage(::CPUDevice, x::AbstractSparseArray) = x diff --git a/ext/DeviceUtilsTrackerExt.jl b/ext/MLDataDevicesTrackerExt.jl similarity index 59% rename from ext/DeviceUtilsTrackerExt.jl rename to ext/MLDataDevicesTrackerExt.jl index 0854d62..a30da57 100644 --- a/ext/DeviceUtilsTrackerExt.jl +++ b/ext/MLDataDevicesTrackerExt.jl @@ -1,19 +1,19 @@ module DeviceUtilsTrackerExt using Adapt: Adapt -using DeviceUtils: DeviceUtils, AMDGPUDevice, CUDADevice, MetalDevice, oneAPIDevice +using MLDataDevices: MLDataDevices, AMDGPUDevice, CUDADevice, MetalDevice, oneAPIDevice using Tracker: Tracker for op in (:_get_device, :_get_device_type) @eval begin - DeviceUtils.$op(x::Tracker.TrackedArray) = DeviceUtils.$op(Tracker.data(x)) - function DeviceUtils.$op(x::AbstractArray{<:Tracker.TrackedReal}) - return DeviceUtils.$op(Tracker.data.(x)) + MLDataDevices.$op(x::Tracker.TrackedArray) = MLDataDevices.$op(Tracker.data(x)) + function MLDataDevices.$op(x::AbstractArray{<:Tracker.TrackedReal}) + return MLDataDevices.$op(Tracker.data.(x)) end end end -DeviceUtils.__special_aos(::AbstractArray{<:Tracker.TrackedReal}) = true +MLDataDevices.__special_aos(::AbstractArray{<:Tracker.TrackedReal}) = true for T in (AMDGPUDevice, AMDGPUDevice{Nothing}, CUDADevice, CUDADevice{Nothing}, MetalDevice, oneAPIDevice) diff --git a/ext/DeviceUtilsZygoteExt.jl b/ext/MLDataDevicesZygoteExt.jl similarity index 82% rename from ext/DeviceUtilsZygoteExt.jl rename to ext/MLDataDevicesZygoteExt.jl index 5b7e6b0..7c4c202 100644 --- a/ext/DeviceUtilsZygoteExt.jl +++ b/ext/MLDataDevicesZygoteExt.jl @@ -1,7 +1,7 @@ module DeviceUtilsZygoteExt using Adapt: Adapt -using DeviceUtils: AbstractDevice, CPUDevice +using MLDataDevices: AbstractDevice, CPUDevice using Zygote: OneElement Adapt.adapt_structure(::CPUDevice, x::OneElement) = x diff --git a/ext/DeviceUtilscuDNNExt.jl b/ext/MLDataDevicescuDNNExt.jl similarity index 77% rename from ext/DeviceUtilscuDNNExt.jl rename to ext/MLDataDevicescuDNNExt.jl index c87cfaf..308cc7f 100644 --- a/ext/DeviceUtilscuDNNExt.jl +++ b/ext/MLDataDevicescuDNNExt.jl @@ -2,7 +2,7 @@ module DeviceUtilscuDNNExt using CUDA: CUDA using cuDNN: cuDNN -using DeviceUtils: DeviceUtils, CUDADevice, reset_gpu_device! +using MLDataDevices: MLDataDevices, CUDADevice, reset_gpu_device! __init__() = reset_gpu_device!() @@ -26,9 +26,9 @@ function _check_use_cuda!() return end -DeviceUtils.loaded(::Union{CUDADevice, Type{<:CUDADevice}}) = true +MLDataDevices.loaded(::Union{CUDADevice, Type{<:CUDADevice}}) = true -function DeviceUtils.functional(::Union{CUDADevice, Type{<:CUDADevice}})::Bool +function MLDataDevices.functional(::Union{CUDADevice, Type{<:CUDADevice}})::Bool _check_use_cuda!() return USE_CUDA_GPU[] end diff --git a/ext/DeviceUtilsoneAPIExt.jl b/ext/MLDataDevicesoneAPIExt.jl similarity index 71% rename from ext/DeviceUtilsoneAPIExt.jl rename to ext/MLDataDevicesoneAPIExt.jl index 24ef8c4..68db94e 100644 --- a/ext/DeviceUtilsoneAPIExt.jl +++ b/ext/MLDataDevicesoneAPIExt.jl @@ -2,7 +2,7 @@ module DeviceUtilsoneAPIExt using Adapt: Adapt using GPUArrays: GPUArrays -using DeviceUtils: DeviceUtils, oneAPIDevice, reset_gpu_device! +using MLDataDevices: MLDataDevices, oneAPIDevice, reset_gpu_device! using oneAPI: oneAPI, oneArray, oneL0 const SUPPORTS_FP64 = Dict{oneL0.ZeDevice, Bool}() @@ -16,18 +16,18 @@ function __init__() end end -DeviceUtils.loaded(::Union{oneAPIDevice, Type{<:oneAPIDevice}}) = true -function DeviceUtils.functional(::Union{oneAPIDevice, Type{<:oneAPIDevice}}) +MLDataDevices.loaded(::Union{oneAPIDevice, Type{<:oneAPIDevice}}) = true +function MLDataDevices.functional(::Union{oneAPIDevice, Type{<:oneAPIDevice}}) return oneAPI.functional() end # Default RNG -DeviceUtils.default_device_rng(::oneAPIDevice) = GPUArrays.default_rng(oneArray) +MLDataDevices.default_device_rng(::oneAPIDevice) = GPUArrays.default_rng(oneArray) # Query Device from Array -DeviceUtils._get_device(::oneArray) = oneAPIDevice() +MLDataDevices._get_device(::oneArray) = oneAPIDevice() -DeviceUtils._get_device_type(::oneArray) = oneAPIDevice +MLDataDevices._get_device_type(::oneArray) = oneAPIDevice # Device Transfer ## To GPU diff --git a/src/DeviceUtils.jl b/src/MLDataDevices.jl similarity index 99% rename from src/DeviceUtils.jl rename to src/MLDataDevices.jl index da8b23b..556bfab 100644 --- a/src/DeviceUtils.jl +++ b/src/MLDataDevices.jl @@ -1,4 +1,4 @@ -module DeviceUtils +module MLDataDevices using Adapt: Adapt using ChainRulesCore: ChainRulesCore, NoTangent @@ -25,7 +25,7 @@ abstract type AbstractGPUDevice <: AbstractDevice end Checks if the device is functional. This is used to determine if the device can be used for computation. Note that even if the backend is loaded (as checked via -[`DeviceUtils.loaded`](@ref)), the device may not be functional. +[`MLDataDevices.loaded`](@ref)), the device may not be functional. Note that while this function is not exported, it is considered part of the public API. """ @@ -108,7 +108,7 @@ Return a tuple of supported GPU backends. !!! warning This is not the list of functional backends on the system, but rather backends which - `DeviceUtils.jl` supports. + `MLDataDevices.jl` supports. """ @inline supported_gpu_backends() = map(_get_device_name, GPU_DEVICES) diff --git a/test/amdgpu_tests.jl b/test/amdgpu_tests.jl index f7c4dac..3d8bf57 100644 --- a/test/amdgpu_tests.jl +++ b/test/amdgpu_tests.jl @@ -1,31 +1,31 @@ -using DeviceUtils, Random, Test +using MLDataDevices, Random, Test using ArrayInterface: parameterless_type @testset "CPU Fallback" begin - @test !DeviceUtils.functional(AMDGPUDevice) + @test !MLDataDevices.functional(AMDGPUDevice) @test cpu_device() isa CPUDevice @test gpu_device() isa CPUDevice - @test_throws DeviceUtils.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) @test_throws Exception default_device_rng(AMDGPUDevice(nothing)) - @test_logs (:warn, "`AMDGPU.jl` hasn't been loaded. Ignoring the device setting.") DeviceUtils.set_device!( + @test_logs (:warn, "`AMDGPU.jl` hasn't been loaded. Ignoring the device setting.") MLDataDevices.set_device!( AMDGPUDevice, nothing, 1) end using AMDGPU @testset "Loaded Trigger Package" begin - @test DeviceUtils.GPU_DEVICE[] === nothing + @test MLDataDevices.GPU_DEVICE[] === nothing - if DeviceUtils.functional(AMDGPUDevice) + if MLDataDevices.functional(AMDGPUDevice) @info "AMDGPU is functional" @test gpu_device() isa AMDGPUDevice @test gpu_device(; force_gpu_usage=true) isa AMDGPUDevice else @info "AMDGPU is NOT functional" @test gpu_device() isa CPUDevice - @test_throws DeviceUtils.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) end - @test DeviceUtils.GPU_DEVICE[] !== nothing + @test MLDataDevices.GPU_DEVICE[] !== nothing end using FillArrays, Zygote # Extensions @@ -38,8 +38,8 @@ using FillArrays, Zygote # Extensions one_elem=Zygote.OneElement(2.0f0, (2, 3), (1:3, 1:4)), farray=Fill(1.0f0, (2, 3))) device = gpu_device() - aType = DeviceUtils.functional(AMDGPUDevice) ? ROCArray : Array - rngType = DeviceUtils.functional(AMDGPUDevice) ? AMDGPU.rocRAND.RNG : Random.AbstractRNG + aType = MLDataDevices.functional(AMDGPUDevice) ? ROCArray : Array + rngType = MLDataDevices.functional(AMDGPUDevice) ? AMDGPU.rocRAND.RNG : Random.AbstractRNG ps_xpu = ps |> device @test get_device(ps_xpu) isa AMDGPUDevice @@ -57,7 +57,7 @@ using FillArrays, Zygote # Extensions @test ps_xpu.rng_default isa rngType @test ps_xpu.rng == ps.rng - if DeviceUtils.functional(AMDGPUDevice) + if MLDataDevices.functional(AMDGPUDevice) @test ps_xpu.one_elem isa ROCArray @test ps_xpu.farray isa ROCArray else @@ -83,7 +83,7 @@ using FillArrays, Zygote # Extensions @test ps_cpu.rng_default isa Random.TaskLocalRNG @test ps_cpu.rng == ps.rng - if DeviceUtils.functional(AMDGPUDevice) + if MLDataDevices.functional(AMDGPUDevice) @test ps_cpu.one_elem isa Array @test ps_cpu.farray isa Array else @@ -100,7 +100,7 @@ using FillArrays, Zygote # Extensions @test get_device(x_dev) isa parameterless_type(typeof(dev)) @test get_device_type(x_dev) <: parameterless_type(typeof(dev)) - if DeviceUtils.functional(AMDGPUDevice) + if MLDataDevices.functional(AMDGPUDevice) dev2 = gpu_device(length(AMDGPU.devices())) x_dev2 = x_dev |> dev2 @test get_device(x_dev2) isa typeof(dev2) @@ -117,7 +117,7 @@ using FillArrays, Zygote # Extensions end @testset "Wrapped Arrays" begin - if DeviceUtils.functional(AMDGPUDevice) + if MLDataDevices.functional(AMDGPUDevice) x = rand(10, 10) |> AMDGPUDevice() @test get_device(x) isa AMDGPUDevice @test get_device_type(x) <: AMDGPUDevice @@ -128,7 +128,7 @@ end end @testset "Multiple Devices AMDGPU" begin - if DeviceUtils.functional(AMDGPUDevice) + if MLDataDevices.functional(AMDGPUDevice) ps = (; weight=rand(Float32, 10), bias=rand(Float32, 10)) ps_cpu = deepcopy(ps) cdev = cpu_device() @@ -153,9 +153,9 @@ end end @testset "setdevice!" begin - if DeviceUtils.functional(AMDGPUDevice) + if MLDataDevices.functional(AMDGPUDevice) for i in 1:10 - @test_nowarn DeviceUtils.set_device!(AMDGPUDevice, nothing, i) + @test_nowarn MLDataDevices.set_device!(AMDGPUDevice, nothing, i) end end end diff --git a/test/cuda_tests.jl b/test/cuda_tests.jl index 0d08ffa..9465b99 100644 --- a/test/cuda_tests.jl +++ b/test/cuda_tests.jl @@ -1,31 +1,31 @@ -using DeviceUtils, Random, Functors, Test +using MLDataDevices, Random, Functors, Test using ArrayInterface: parameterless_type @testset "CPU Fallback" begin - @test !DeviceUtils.functional(CUDADevice) + @test !MLDataDevices.functional(CUDADevice) @test cpu_device() isa CPUDevice @test gpu_device() isa CPUDevice - @test_throws DeviceUtils.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) @test_throws Exception default_device_rng(CUDADevice(nothing)) - @test_logs (:warn, "`CUDA.jl` hasn't been loaded. Ignoring the device setting.") DeviceUtils.set_device!( + @test_logs (:warn, "`CUDA.jl` hasn't been loaded. Ignoring the device setting.") MLDataDevices.set_device!( CUDADevice, nothing, 1) end using LuxCUDA @testset "Loaded Trigger Package" begin - @test DeviceUtils.GPU_DEVICE[] === nothing + @test MLDataDevices.GPU_DEVICE[] === nothing - if DeviceUtils.functional(CUDADevice) + if MLDataDevices.functional(CUDADevice) @info "LuxCUDA is functional" @test gpu_device() isa CUDADevice @test gpu_device(; force_gpu_usage=true) isa CUDADevice else @info "LuxCUDA is NOT functional" @test gpu_device() isa CPUDevice - @test_throws DeviceUtils.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) end - @test DeviceUtils.GPU_DEVICE[] !== nothing + @test MLDataDevices.GPU_DEVICE[] !== nothing end using FillArrays, Zygote # Extensions @@ -38,8 +38,8 @@ using FillArrays, Zygote # Extensions one_elem=Zygote.OneElement(2.0f0, (2, 3), (1:3, 1:4)), farray=Fill(1.0f0, (2, 3))) device = gpu_device() - aType = DeviceUtils.functional(CUDADevice) ? CuArray : Array - rngType = DeviceUtils.functional(CUDADevice) ? CUDA.RNG : Random.AbstractRNG + aType = MLDataDevices.functional(CUDADevice) ? CuArray : Array + rngType = MLDataDevices.functional(CUDADevice) ? CUDA.RNG : Random.AbstractRNG ps_xpu = ps |> device @test get_device(ps_xpu) isa CUDADevice @@ -57,7 +57,7 @@ using FillArrays, Zygote # Extensions @test ps_xpu.rng_default isa rngType @test ps_xpu.rng == ps.rng - if DeviceUtils.functional(CUDADevice) + if MLDataDevices.functional(CUDADevice) @test ps_xpu.one_elem isa CuArray @test ps_xpu.farray isa CuArray else @@ -83,7 +83,7 @@ using FillArrays, Zygote # Extensions @test ps_cpu.rng_default isa Random.TaskLocalRNG @test ps_cpu.rng == ps.rng - if DeviceUtils.functional(CUDADevice) + if MLDataDevices.functional(CUDADevice) @test ps_cpu.one_elem isa Array @test ps_cpu.farray isa Array else @@ -101,7 +101,7 @@ using FillArrays, Zygote # Extensions @test get_device(data) isa CPUDevice @test get_device_type(data) <: CPUDevice data_dev = data |> device - if DeviceUtils.functional(CUDADevice) + if MLDataDevices.functional(CUDADevice) @test get_device(data_dev) isa CUDADevice @test get_device_type(data_dev) <: CUDADevice else @@ -123,7 +123,7 @@ using FillArrays, Zygote # Extensions @test get_device(x_dev) isa parameterless_type(typeof(dev)) @test get_device_type(x_dev) <: parameterless_type(typeof(dev)) - if DeviceUtils.functional(CUDADevice) + if MLDataDevices.functional(CUDADevice) dev2 = gpu_device(length(CUDA.devices())) x_dev2 = x_dev |> dev2 @test get_device(x_dev2) isa typeof(dev2) @@ -143,7 +143,7 @@ using FillArrays, Zygote # Extensions end @testset "Wrapped Arrays" begin - if DeviceUtils.functional(CUDADevice) + if MLDataDevices.functional(CUDADevice) x = rand(10, 10) |> CUDADevice() @test get_device(x) isa CUDADevice @test get_device_type(x) <: CUDADevice @@ -154,7 +154,7 @@ end end @testset "Multiple Devices CUDA" begin - if DeviceUtils.functional(CUDADevice) + if MLDataDevices.functional(CUDADevice) ps = (; weight=rand(Float32, 10), bias=rand(Float32, 10)) ps_cpu = deepcopy(ps) cdev = cpu_device() @@ -181,7 +181,7 @@ end using SparseArrays @testset "CUDA Sparse Arrays" begin - if DeviceUtils.functional(CUDADevice) + if MLDataDevices.functional(CUDADevice) ps = (; weight=sprand(Float32, 10, 10, 0.1), bias=sprand(Float32, 10, 0.1)) ps_cpu = deepcopy(ps) cdev = cpu_device() @@ -206,9 +206,9 @@ using SparseArrays end @testset "setdevice!" begin - if DeviceUtils.functional(CUDADevice) + if MLDataDevices.functional(CUDADevice) for i in 1:10 - @test_nowarn DeviceUtils.set_device!(CUDADevice, nothing, i) + @test_nowarn MLDataDevices.set_device!(CUDADevice, nothing, i) end end end diff --git a/test/metal_tests.jl b/test/metal_tests.jl index 2d89a43..1e25c53 100644 --- a/test/metal_tests.jl +++ b/test/metal_tests.jl @@ -1,29 +1,29 @@ -using DeviceUtils, Random, Test +using MLDataDevices, Random, Test using ArrayInterface: parameterless_type @testset "CPU Fallback" begin - @test !DeviceUtils.functional(MetalDevice) + @test !MLDataDevices.functional(MetalDevice) @test cpu_device() isa CPUDevice @test gpu_device() isa CPUDevice - @test_throws DeviceUtils.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) @test_throws Exception default_device_rng(MetalDevice()) end using Metal @testset "Loaded Trigger Package" begin - @test DeviceUtils.GPU_DEVICE[] === nothing + @test MLDataDevices.GPU_DEVICE[] === nothing - if DeviceUtils.functional(MetalDevice) + if MLDataDevices.functional(MetalDevice) @info "Metal is functional" @test gpu_device() isa MetalDevice @test gpu_device(; force_gpu_usage=true) isa MetalDevice else @info "Metal is NOT functional" @test gpu_device() isa MetalDevice - @test_throws DeviceUtils.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) end - @test DeviceUtils.GPU_DEVICE[] !== nothing + @test MLDataDevices.GPU_DEVICE[] !== nothing end using FillArrays, Zygote # Extensions @@ -36,8 +36,8 @@ using FillArrays, Zygote # Extensions one_elem=Zygote.OneElement(2.0f0, (2, 3), (1:3, 1:4)), farray=Fill(1.0f0, (2, 3))) device = gpu_device() - aType = DeviceUtils.functional(MetalDevice) ? MtlArray : Array - rngType = DeviceUtils.functional(MetalDevice) ? Metal.GPUArrays.RNG : Random.AbstractRNG + aType = MLDataDevices.functional(MetalDevice) ? MtlArray : Array + rngType = MLDataDevices.functional(MetalDevice) ? Metal.GPUArrays.RNG : Random.AbstractRNG ps_xpu = ps |> device @test get_device(ps_xpu) isa MetalDevice @@ -55,7 +55,7 @@ using FillArrays, Zygote # Extensions @test ps_xpu.rng_default isa rngType @test ps_xpu.rng == ps.rng - if DeviceUtils.functional(MetalDevice) + if MLDataDevices.functional(MetalDevice) @test ps_xpu.one_elem isa MtlArray @test ps_xpu.farray isa MtlArray else @@ -81,7 +81,7 @@ using FillArrays, Zygote # Extensions @test ps_cpu.rng_default isa Random.TaskLocalRNG @test ps_cpu.rng == ps.rng - if DeviceUtils.functional(MetalDevice) + if MLDataDevices.functional(MetalDevice) @test ps_cpu.one_elem isa Array @test ps_cpu.farray isa Array else @@ -106,7 +106,7 @@ using FillArrays, Zygote # Extensions end @testset "Wrapper Arrays" begin - if DeviceUtils.functional(MetalDevice) + if MLDataDevices.functional(MetalDevice) x = rand(Float32, 10, 10) |> MetalDevice() @test get_device(x) isa MetalDevice @test get_device_type(x) <: MetalDevice @@ -117,9 +117,9 @@ end end @testset "setdevice!" begin - if DeviceUtils.functional(MetalDevice) + if MLDataDevices.functional(MetalDevice) @test_logs (:warn, - "Support for Multi Device Metal hasn't been implemented yet. Ignoring the device setting.") DeviceUtils.set_device!( + "Support for Multi Device Metal hasn't been implemented yet. Ignoring the device setting.") MLDataDevices.set_device!( MetalDevice, nothing, 1) end end diff --git a/test/misc_tests.jl b/test/misc_tests.jl index 653c1f2..e3f3ed8 100644 --- a/test/misc_tests.jl +++ b/test/misc_tests.jl @@ -1,10 +1,10 @@ -using Adapt, DeviceUtils, ComponentArrays, Random +using Adapt, MLDataDevices, ComponentArrays, Random using ArrayInterface: parameterless_type using ChainRulesTestUtils: test_rrule using ReverseDiff, Tracker, ForwardDiff using SparseArrays, FillArrays, Zygote, RecursiveArrayTools -@testset "https://github.com/LuxDL/DeviceUtils.jl/issues/10 patch" begin +@testset "https://github.com/LuxDL/MLDataDevices.jl/issues/10 patch" begin dev = CPUDevice() ps = (; weight=randn(10, 1), bias=randn(1)) @@ -95,7 +95,7 @@ end @testset "CPU setdevice!" begin @test_logs (:warn, - "Setting device for `CPUDevice` doesn't make sense. Ignoring the device setting.") DeviceUtils.set_device!( + "Setting device for `CPUDevice` doesn't make sense. Ignoring the device setting.") MLDataDevices.set_device!( CPUDevice, nothing, 1) end @@ -116,8 +116,8 @@ end end @testset "loaded and functional" begin - @test DeviceUtils.loaded(CPUDevice) - @test DeviceUtils.functional(CPUDevice) + @test MLDataDevices.loaded(CPUDevice) + @test MLDataDevices.functional(CPUDevice) end @testset "writing to preferences" begin @@ -127,7 +127,7 @@ end for backend in (:CUDA, :AMDGPU, :oneAPI, :Metal, AMDGPUDevice(), CUDADevice(), MetalDevice(), oneAPIDevice()) backend_name = backend isa Symbol ? string(backend) : - DeviceUtils._get_device_name(backend) + MLDataDevices._get_device_name(backend) @test_logs (:info, "GPU backend has been set to $(backend_name). Restart Julia to use the new backend.") gpu_backend!(backend) end diff --git a/test/oneapi_tests.jl b/test/oneapi_tests.jl index 638836e..25b1ed3 100644 --- a/test/oneapi_tests.jl +++ b/test/oneapi_tests.jl @@ -1,29 +1,29 @@ -using DeviceUtils, Random, Test +using MLDataDevices, Random, Test using ArrayInterface: parameterless_type @testset "CPU Fallback" begin - @test !DeviceUtils.functional(oneAPIDevice) + @test !MLDataDevices.functional(oneAPIDevice) @test cpu_device() isa CPUDevice @test gpu_device() isa CPUDevice - @test_throws DeviceUtils.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) @test_throws Exception default_device_rng(oneAPIDevice()) end using oneAPI @testset "Loaded Trigger Package" begin - @test DeviceUtils.GPU_DEVICE[] === nothing + @test MLDataDevices.GPU_DEVICE[] === nothing - if DeviceUtils.functional(oneAPIDevice) + if MLDataDevices.functional(oneAPIDevice) @info "oneAPI is functional" @test gpu_device() isa oneAPIDevice @test gpu_device(; force_gpu_usage=true) isa oneAPIDevice else @info "oneAPI is NOT functional" @test gpu_device() isa oneAPIDevice - @test_throws DeviceUtils.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) end - @test DeviceUtils.GPU_DEVICE[] !== nothing + @test MLDataDevices.GPU_DEVICE[] !== nothing end using FillArrays, Zygote # Extensions @@ -36,8 +36,8 @@ using FillArrays, Zygote # Extensions one_elem=Zygote.OneElement(2.0f0, (2, 3), (1:3, 1:4)), farray=Fill(1.0f0, (2, 3))) device = gpu_device() - aType = DeviceUtils.functional(oneAPIDevice) ? oneArray : Array - rngType = DeviceUtils.functional(oneAPIDevice) ? oneAPI.GPUArrays.RNG : + aType = MLDataDevices.functional(oneAPIDevice) ? oneArray : Array + rngType = MLDataDevices.functional(oneAPIDevice) ? oneAPI.GPUArrays.RNG : Random.AbstractRNG ps_xpu = ps |> device @@ -56,7 +56,7 @@ using FillArrays, Zygote # Extensions @test ps_xpu.rng_default isa rngType @test ps_xpu.rng == ps.rng - if DeviceUtils.functional(oneAPIDevice) + if MLDataDevices.functional(oneAPIDevice) @test ps_xpu.one_elem isa oneArray @test ps_xpu.farray isa oneArray else @@ -82,7 +82,7 @@ using FillArrays, Zygote # Extensions @test ps_cpu.rng_default isa Random.TaskLocalRNG @test ps_cpu.rng == ps.rng - if DeviceUtils.functional(oneAPIDevice) + if MLDataDevices.functional(oneAPIDevice) @test ps_cpu.one_elem isa Array @test ps_cpu.farray isa Array else @@ -107,7 +107,7 @@ using FillArrays, Zygote # Extensions end @testset "Wrapper Arrays" begin - if DeviceUtils.functional(oneAPIDevice) + if MLDataDevices.functional(oneAPIDevice) x = rand(10, 10) |> oneAPIDevice() @test get_device(x) isa oneAPIDevice @test get_device_type(x) <: oneAPIDevice @@ -118,9 +118,9 @@ end end @testset "setdevice!" begin - if DeviceUtils.functional(oneAPIDevice) + if MLDataDevices.functional(oneAPIDevice) @test_logs (:warn, - "Support for Multi Device oneAPI hasn't been implemented yet. Ignoring the device setting.") DeviceUtils.set_device!( + "Support for Multi Device oneAPI hasn't been implemented yet. Ignoring the device setting.") MLDataDevices.set_device!( oneAPIDevice, nothing, 1) end end diff --git a/test/qa_tests.jl b/test/qa_tests.jl index b08a873..965e818 100644 --- a/test/qa_tests.jl +++ b/test/qa_tests.jl @@ -1,17 +1,17 @@ -using Aqua, ExplicitImports, DeviceUtils, Test +using Aqua, ExplicitImports, MLDataDevices, Test @testset "Aqua Tests" begin - Aqua.test_all(DeviceUtils) + Aqua.test_all(MLDataDevices) end import FillArrays, RecursiveArrayTools, SparseArrays, Zygote @testset "Explicit Imports" begin - @test check_no_implicit_imports(DeviceUtils) === nothing - @test check_no_stale_explicit_imports(DeviceUtils) === nothing - @test check_no_self_qualified_accesses(DeviceUtils) === nothing - @test check_all_explicit_imports_via_owners(DeviceUtils) === nothing - @test check_all_qualified_accesses_via_owners(DeviceUtils) === nothing - @test_broken check_all_explicit_imports_are_public(DeviceUtils) === nothing # mostly upstream problems - @test_broken check_all_qualified_accesses_are_public(DeviceUtils) === nothing # mostly upstream problem + @test check_no_implicit_imports(MLDataDevices) === nothing + @test check_no_stale_explicit_imports(MLDataDevices) === nothing + @test check_no_self_qualified_accesses(MLDataDevices) === nothing + @test check_all_explicit_imports_via_owners(MLDataDevices) === nothing + @test check_all_qualified_accesses_via_owners(MLDataDevices) === nothing + @test_broken check_all_explicit_imports_are_public(MLDataDevices) === nothing # mostly upstream problems + @test_broken check_all_qualified_accesses_are_public(MLDataDevices) === nothing # mostly upstream problem end diff --git a/test/runtests.jl b/test/runtests.jl index 8448f4b..b9fb136 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -18,7 +18,7 @@ if !isempty(EXTRA_PKGS) Pkg.instantiate() end -@testset "DeviceUtils Tests" begin +@testset "MLDataDevices Tests" begin file_names = BACKEND_GROUP == "all" ? ["cuda_tests.jl", "amdgpu_tests.jl", "metal_tests.jl", "oneapi_tests.jl"] : (BACKEND_GROUP == "cpu" ? [] : [BACKEND_GROUP * "_tests.jl"]) From a6e46d030aaf69ba2f37b4b93189fd0ae4147270 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 23 Jul 2024 19:49:35 -0700 Subject: [PATCH 2/3] chore: apply formatting --- test/amdgpu_tests.jl | 6 ++++-- test/cuda_tests.jl | 3 ++- test/metal_tests.jl | 6 ++++-- test/oneapi_tests.jl | 3 ++- 4 files changed, 12 insertions(+), 6 deletions(-) diff --git a/test/amdgpu_tests.jl b/test/amdgpu_tests.jl index 3d8bf57..0338031 100644 --- a/test/amdgpu_tests.jl +++ b/test/amdgpu_tests.jl @@ -23,7 +23,8 @@ using AMDGPU else @info "AMDGPU is NOT functional" @test gpu_device() isa CPUDevice - @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; + force_gpu_usage=true) end @test MLDataDevices.GPU_DEVICE[] !== nothing end @@ -39,7 +40,8 @@ using FillArrays, Zygote # Extensions device = gpu_device() aType = MLDataDevices.functional(AMDGPUDevice) ? ROCArray : Array - rngType = MLDataDevices.functional(AMDGPUDevice) ? AMDGPU.rocRAND.RNG : Random.AbstractRNG + rngType = MLDataDevices.functional(AMDGPUDevice) ? AMDGPU.rocRAND.RNG : + Random.AbstractRNG ps_xpu = ps |> device @test get_device(ps_xpu) isa AMDGPUDevice diff --git a/test/cuda_tests.jl b/test/cuda_tests.jl index 9465b99..7804183 100644 --- a/test/cuda_tests.jl +++ b/test/cuda_tests.jl @@ -23,7 +23,8 @@ using LuxCUDA else @info "LuxCUDA is NOT functional" @test gpu_device() isa CPUDevice - @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; + force_gpu_usage=true) end @test MLDataDevices.GPU_DEVICE[] !== nothing end diff --git a/test/metal_tests.jl b/test/metal_tests.jl index 1e25c53..3bf98ec 100644 --- a/test/metal_tests.jl +++ b/test/metal_tests.jl @@ -21,7 +21,8 @@ using Metal else @info "Metal is NOT functional" @test gpu_device() isa MetalDevice - @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; + force_gpu_usage=true) end @test MLDataDevices.GPU_DEVICE[] !== nothing end @@ -37,7 +38,8 @@ using FillArrays, Zygote # Extensions device = gpu_device() aType = MLDataDevices.functional(MetalDevice) ? MtlArray : Array - rngType = MLDataDevices.functional(MetalDevice) ? Metal.GPUArrays.RNG : Random.AbstractRNG + rngType = MLDataDevices.functional(MetalDevice) ? Metal.GPUArrays.RNG : + Random.AbstractRNG ps_xpu = ps |> device @test get_device(ps_xpu) isa MetalDevice diff --git a/test/oneapi_tests.jl b/test/oneapi_tests.jl index 25b1ed3..a9f25cf 100644 --- a/test/oneapi_tests.jl +++ b/test/oneapi_tests.jl @@ -21,7 +21,8 @@ using oneAPI else @info "oneAPI is NOT functional" @test gpu_device() isa oneAPIDevice - @test_throws MLDataDevices.DeviceSelectionException gpu_device(; force_gpu_usage=true) + @test_throws MLDataDevices.DeviceSelectionException gpu_device(; + force_gpu_usage=true) end @test MLDataDevices.GPU_DEVICE[] !== nothing end From 611c28eb50cf5736e855d397d3040877391605e4 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Tue, 23 Jul 2024 19:53:12 -0700 Subject: [PATCH 3/3] fix: change names --- ext/MLDataDevicesAMDGPUExt.jl | 2 +- ext/MLDataDevicesCUDAExt.jl | 2 +- ext/MLDataDevicesFillArraysExt.jl | 2 +- ext/MLDataDevicesGPUArraysExt.jl | 2 +- ext/MLDataDevicesMetalExt.jl | 2 +- ext/MLDataDevicesRecursiveArrayToolsExt.jl | 2 +- ext/MLDataDevicesReverseDiffExt.jl | 2 +- ext/MLDataDevicesSparseArraysExt.jl | 2 +- ext/MLDataDevicesTrackerExt.jl | 2 +- ext/MLDataDevicesZygoteExt.jl | 2 +- ext/MLDataDevicescuDNNExt.jl | 2 +- ext/MLDataDevicesoneAPIExt.jl | 2 +- 12 files changed, 12 insertions(+), 12 deletions(-) diff --git a/ext/MLDataDevicesAMDGPUExt.jl b/ext/MLDataDevicesAMDGPUExt.jl index 5b008f1..7769b84 100644 --- a/ext/MLDataDevicesAMDGPUExt.jl +++ b/ext/MLDataDevicesAMDGPUExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsAMDGPUExt +module MLDataDevicesAMDGPUExt using Adapt: Adapt using AMDGPU: AMDGPU diff --git a/ext/MLDataDevicesCUDAExt.jl b/ext/MLDataDevicesCUDAExt.jl index a353b42..6362f80 100644 --- a/ext/MLDataDevicesCUDAExt.jl +++ b/ext/MLDataDevicesCUDAExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsCUDAExt +module MLDataDevicesCUDAExt using Adapt: Adapt using CUDA: CUDA diff --git a/ext/MLDataDevicesFillArraysExt.jl b/ext/MLDataDevicesFillArraysExt.jl index 36a5d6f..5a88241 100644 --- a/ext/MLDataDevicesFillArraysExt.jl +++ b/ext/MLDataDevicesFillArraysExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsFillArraysExt +module MLDataDevicesFillArraysExt using Adapt: Adapt using FillArrays: FillArrays, AbstractFill diff --git a/ext/MLDataDevicesGPUArraysExt.jl b/ext/MLDataDevicesGPUArraysExt.jl index 328222a..daf7eb3 100644 --- a/ext/MLDataDevicesGPUArraysExt.jl +++ b/ext/MLDataDevicesGPUArraysExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsGPUArraysExt +module MLDataDevicesGPUArraysExt using Adapt: Adapt using GPUArrays: GPUArrays diff --git a/ext/MLDataDevicesMetalExt.jl b/ext/MLDataDevicesMetalExt.jl index f82d55c..1c81689 100644 --- a/ext/MLDataDevicesMetalExt.jl +++ b/ext/MLDataDevicesMetalExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsMetalExt +module MLDataDevicesMetalExt using Adapt: Adapt using GPUArrays: GPUArrays diff --git a/ext/MLDataDevicesRecursiveArrayToolsExt.jl b/ext/MLDataDevicesRecursiveArrayToolsExt.jl index cc006ba..4277150 100644 --- a/ext/MLDataDevicesRecursiveArrayToolsExt.jl +++ b/ext/MLDataDevicesRecursiveArrayToolsExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsRecursiveArrayToolsExt +module MLDataDevicesRecursiveArrayToolsExt using Adapt: Adapt, adapt using MLDataDevices: MLDataDevices, AbstractDevice diff --git a/ext/MLDataDevicesReverseDiffExt.jl b/ext/MLDataDevicesReverseDiffExt.jl index 14915d9..9e6553e 100644 --- a/ext/MLDataDevicesReverseDiffExt.jl +++ b/ext/MLDataDevicesReverseDiffExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsReverseDiffExt +module MLDataDevicesReverseDiffExt using MLDataDevices: MLDataDevices using ReverseDiff: ReverseDiff diff --git a/ext/MLDataDevicesSparseArraysExt.jl b/ext/MLDataDevicesSparseArraysExt.jl index 1851872..a52871f 100644 --- a/ext/MLDataDevicesSparseArraysExt.jl +++ b/ext/MLDataDevicesSparseArraysExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsSparseArraysExt +module MLDataDevicesSparseArraysExt using Adapt: Adapt using MLDataDevices: CPUDevice diff --git a/ext/MLDataDevicesTrackerExt.jl b/ext/MLDataDevicesTrackerExt.jl index a30da57..49ef3ea 100644 --- a/ext/MLDataDevicesTrackerExt.jl +++ b/ext/MLDataDevicesTrackerExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsTrackerExt +module MLDataDevicesTrackerExt using Adapt: Adapt using MLDataDevices: MLDataDevices, AMDGPUDevice, CUDADevice, MetalDevice, oneAPIDevice diff --git a/ext/MLDataDevicesZygoteExt.jl b/ext/MLDataDevicesZygoteExt.jl index 7c4c202..1b705c5 100644 --- a/ext/MLDataDevicesZygoteExt.jl +++ b/ext/MLDataDevicesZygoteExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsZygoteExt +module MLDataDevicesZygoteExt using Adapt: Adapt using MLDataDevices: AbstractDevice, CPUDevice diff --git a/ext/MLDataDevicescuDNNExt.jl b/ext/MLDataDevicescuDNNExt.jl index 308cc7f..a332c7a 100644 --- a/ext/MLDataDevicescuDNNExt.jl +++ b/ext/MLDataDevicescuDNNExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilscuDNNExt +module MLDataDevicescuDNNExt using CUDA: CUDA using cuDNN: cuDNN diff --git a/ext/MLDataDevicesoneAPIExt.jl b/ext/MLDataDevicesoneAPIExt.jl index 68db94e..ebffa02 100644 --- a/ext/MLDataDevicesoneAPIExt.jl +++ b/ext/MLDataDevicesoneAPIExt.jl @@ -1,4 +1,4 @@ -module DeviceUtilsoneAPIExt +module MLDataDevicesoneAPIExt using Adapt: Adapt using GPUArrays: GPUArrays