cupy_event.py 766 Bytes
Newer Older
root's avatar
root committed
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
# nvprof --print-gpu-trace python examples/stream/cupy_event.py
import cupy

x = cupy.array([1, 2, 3])

start_event = cupy.cuda.stream.Event()
stop_event = cupy.cuda.stream.Event()


def _norm_with_elapsed_time(x):
    start_event.record()
    y = cupy.linalg.norm(x)
    stop_event.record()
    stop_event.synchronize()
    print(cupy.cuda.get_elapsed_time(start_event, stop_event))
    return y


expected = _norm_with_elapsed_time(x)
cupy.cuda.Device().synchronize()

stream = cupy.cuda.stream.Stream()
with stream:
    y = _norm_with_elapsed_time(x)
stream.synchronize()
cupy.testing.assert_array_equal(y, expected)

stream = cupy.cuda.stream.Stream()
stream.use()
y = _norm_with_elapsed_time(x)
stream.synchronize()
cupy.testing.assert_array_equal(y, expected)