This repository was archived by the owner on Nov 4, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 4
/
Copy pathamdgpu_tests.jl
192 lines (166 loc) · 6.56 KB
/
amdgpu_tests.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
using MLDataDevices, Random, Test
using ArrayInterface: parameterless_type
@testset "CPU Fallback" begin
@test !MLDataDevices.functional(AMDGPUDevice)
@test cpu_device() isa CPUDevice
@test gpu_device() isa CPUDevice
@test_throws MLDataDevices.Internal.DeviceSelectionException gpu_device(; force=true)
@test_throws Exception default_device_rng(AMDGPUDevice(nothing))
@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 MLDataDevices.GPU_DEVICE[] === nothing
if MLDataDevices.functional(AMDGPUDevice)
@info "AMDGPU is functional"
@test gpu_device() isa AMDGPUDevice
@test gpu_device(; force=true) isa AMDGPUDevice
else
@info "AMDGPU is NOT functional"
@test gpu_device() isa CPUDevice
@test_throws MLDataDevices.Internal.DeviceSelectionException gpu_device(;
force=true)
end
@test MLDataDevices.GPU_DEVICE[] !== nothing
end
using FillArrays, Zygote # Extensions
@testset "Data Transfer" begin
ps = (a=(c=zeros(10, 1), d=1), b=ones(10, 1), e=:c,
d="string", mixed=[2.0f0, 3.0, ones(2, 3)], # mixed array types
range=1:10,
rng_default=Random.default_rng(), rng=MersenneTwister(),
one_elem=Zygote.OneElement(2.0f0, (2, 3), (1:3, 1:4)), farray=Fill(1.0f0, (2, 3)))
device = gpu_device()
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
@test get_device_type(ps_xpu) <: AMDGPUDevice
@test ps_xpu.a.c isa aType
@test ps_xpu.b isa aType
@test ps_xpu.a.d == ps.a.d
@test ps_xpu.mixed isa Vector
@test ps_xpu.mixed[1] isa Float32
@test ps_xpu.mixed[2] isa Float64
@test ps_xpu.mixed[3] isa aType
@test ps_xpu.range isa aType
@test ps_xpu.e == ps.e
@test ps_xpu.d == ps.d
@test ps_xpu.rng_default isa rngType
@test get_device(ps_xpu.rng_default) isa AMDGPUDevice
@test get_device_type(ps_xpu.rng_default) <: AMDGPUDevice
@test ps_xpu.rng == ps.rng
@test get_device(ps_xpu.rng) === nothing
@test get_device_type(ps_xpu.rng) <: Nothing
if MLDataDevices.functional(AMDGPUDevice)
@test ps_xpu.one_elem isa ROCArray
@test ps_xpu.farray isa ROCArray
else
@test ps_xpu.one_elem isa Zygote.OneElement
@test ps_xpu.farray isa Fill
end
ps_cpu = ps_xpu |> cpu_device()
@test get_device(ps_cpu) isa CPUDevice
@test get_device_type(ps_cpu) <: CPUDevice
@test ps_cpu.a.c isa Array
@test ps_cpu.b isa Array
@test ps_cpu.a.c == ps.a.c
@test ps_cpu.b == ps.b
@test ps_cpu.a.d == ps.a.d
@test ps_cpu.mixed isa Vector
@test ps_cpu.mixed[1] isa Float32
@test ps_cpu.mixed[2] isa Float64
@test ps_cpu.mixed[3] isa Array
@test ps_cpu.range isa Array
@test ps_cpu.e == ps.e
@test ps_cpu.d == ps.d
@test ps_cpu.rng_default isa Random.TaskLocalRNG
@test get_device(ps_cpu.rng_default) === nothing
@test get_device_type(ps_cpu.rng_default) <: Nothing
@test ps_cpu.rng == ps.rng
@test get_device(ps_cpu.rng) === nothing
@test get_device_type(ps_cpu.rng) <: Nothing
if MLDataDevices.functional(AMDGPUDevice)
@test ps_cpu.one_elem isa Array
@test ps_cpu.farray isa Array
else
@test ps_cpu.one_elem isa Zygote.OneElement
@test ps_cpu.farray isa Fill
end
ps_mixed = (; a=rand(2), b=device(rand(2)))
@test_throws ArgumentError get_device(ps_mixed)
dev = gpu_device()
x = rand(Float32, 10, 2)
x_dev = x |> dev
@test get_device(x_dev) isa parameterless_type(typeof(dev))
@test get_device_type(x_dev) <: parameterless_type(typeof(dev))
if MLDataDevices.functional(AMDGPUDevice)
dev2 = gpu_device(length(AMDGPU.devices()))
x_dev2 = x_dev |> dev2
@test get_device(x_dev2) isa typeof(dev2)
@test get_device_type(x_dev2) <: parameterless_type(typeof(dev2))
end
@testset "get_device_type compile constant" begin
x = rand(10, 10) |> device
ps = (; weight=x, bias=x, d=(x, x))
return_val(x) = Val(get_device_type(x)) # If it is a compile time constant then type inference will work
@test @inferred(return_val(ps)) isa Val{parameterless_type(typeof(device))}
end
end
@testset "Functions" begin
if MLDataDevices.functional(AMDGPUDevice)
@test get_device(tanh) isa MLDataDevices.UnknownDevice
@test get_device_type(tanh) <: MLDataDevices.UnknownDevice
f(x, y) = () -> (x, x .^ 2, y)
ff = f([1, 2, 3], 1)
@test get_device(ff) isa CPUDevice
@test get_device_type(ff) <: CPUDevice
ff_xpu = ff |> AMDGPUDevice()
@test get_device(ff_xpu) isa AMDGPUDevice
@test get_device_type(ff_xpu) <: AMDGPUDevice
ff_cpu = ff_xpu |> cpu_device()
@test get_device(ff_cpu) isa CPUDevice
@test get_device_type(ff_cpu) <: CPUDevice
end
end
@testset "Wrapped Arrays" begin
if MLDataDevices.functional(AMDGPUDevice)
x = rand(10, 10) |> AMDGPUDevice()
@test get_device(x) isa AMDGPUDevice
@test get_device_type(x) <: AMDGPUDevice
x_view = view(x, 1:5, 1:5)
@test get_device(x_view) isa AMDGPUDevice
@test get_device_type(x_view) <: AMDGPUDevice
end
end
@testset "Multiple Devices AMDGPU" begin
if MLDataDevices.functional(AMDGPUDevice)
ps = (; weight=rand(Float32, 10), bias=rand(Float32, 10))
ps_cpu = deepcopy(ps)
cdev = cpu_device()
for idx in 1:length(AMDGPU.devices())
amdgpu_device = gpu_device(idx)
@test typeof(amdgpu_device.device) <: AMDGPU.HIPDevice
@test AMDGPU.device_id(amdgpu_device.device) == idx
ps = ps |> amdgpu_device
@test ps.weight isa ROCArray
@test ps.bias isa ROCArray
@test AMDGPU.device_id(AMDGPU.device(ps.weight)) == idx
@test AMDGPU.device_id(AMDGPU.device(ps.bias)) == idx
@test isequal(cdev(ps.weight), ps_cpu.weight)
@test isequal(cdev(ps.bias), ps_cpu.bias)
end
ps = ps |> cdev
@test ps.weight isa Array
@test ps.bias isa Array
end
end
@testset "setdevice!" begin
if MLDataDevices.functional(AMDGPUDevice)
for i in 1:10
@test_nowarn MLDataDevices.set_device!(AMDGPUDevice, nothing, i)
end
end
end