forked from JuliaGPU/CUDA.jl
-
Notifications
You must be signed in to change notification settings - Fork 0
/
initialization.jl
148 lines (117 loc) · 3.07 KB
/
initialization.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
@test has_cuda(true)
@test has_cuda_gpu(true)
# the API shouldn't have been initialized
@test CuCurrentContext() == nothing
@test CuCurrentDevice() == nothing
task_cb = Any[nothing for tid in 1:Threads.nthreads()]
CUDA.attaskswitch() do
task_cb[Threads.threadid()] = current_task()
end
device_switch_cb = Any[nothing for tid in 1:Threads.nthreads()]
CUDA.atdeviceswitch() do
device_switch_cb[Threads.threadid()] = (dev=device(), ctx=context())
end
device_reset_cb = Any[nothing for tid in 1:Threads.nthreads()]
CUDA.atdevicereset() do dev
device_reset_cb[Threads.threadid()] = dev
end
function reset_cb()
fill!(task_cb, nothing)
fill!(device_switch_cb, nothing)
fill!(device_reset_cb, nothing)
end
# now cause initialization
ctx = context()
dev = device()
@test CuCurrentContext() == ctx
@test CuCurrentDevice() == dev
@test task_cb[1] == current_task()
@test device_switch_cb[1].ctx == ctx
@test device_switch_cb[1].dev == dev
reset_cb()
# ... on a different task
task = @async begin
context()
end
@test ctx == fetch(task)
@test task_cb[1] == task
@test device_switch_cb[1].ctx == ctx
@test device_switch_cb[1].dev == dev
reset_cb()
# ... back to the main task
ctx = context()
dev = device()
@test task_cb[1] == current_task()
@test device_switch_cb[1] == nothing
device!(CuDevice(0))
device!(CuDevice(0)) do
nothing
end
context!(ctx)
context!(ctx) do
nothing
end
@test_throws ErrorException device!(0, CUDA.CU_CTX_SCHED_YIELD)
reset_cb()
device_reset!()
@test device_reset_cb[1] == CuDevice(0)
reset_cb()
device!(0, CUDA.CU_CTX_SCHED_YIELD)
@test task_cb[1] == nothing
@test device_switch_cb[1].dev == CuDevice(0)
# reset on a different task
let ctx = context()
@test CUDA.isvalid(ctx)
@test ctx == fetch(@async context())
@sync @async device_reset!()
@test CUDA.isvalid(context())
@test ctx != context()
end
# test the device selection functionality
if length(devices()) > 1
device!(0)
device!(1) do
@test device() == CuDevice(1)
end
@test device() == CuDevice(0)
device!(1)
@test device() == CuDevice(1)
end
# test that each task can work with devices independently from other tasks
if length(devices()) > 1
device!(0)
@test device() == CuDevice(0)
task = @async begin
device!(1)
@test device() == CuDevice(1)
end
fetch(task)
@test device() == CuDevice(0)
# reset on a task
task = @async begin
device!(1)
device_reset!()
end
fetch(task)
@test device() == CuDevice(0)
# tasks on multiple threads
Threads.@threads for d in 0:1
for x in 1:100 # give threads a chance to trample over each other
device!(d)
yield()
@test device() == CuDevice(d)
yield()
sleep(rand(0.001:0.001:0.01))
device!(1-d)
yield()
@test device() == CuDevice(1-d)
yield()
end
end
@test device() == CuDevice(0)
end
@test deviceid() >= 0
@test deviceid(CuDevice(0)) == 0
if length(devices()) > 1
@test deviceid(CuDevice(1)) == 1
end