cublas.py 471 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
# nvprof --print-gpu-trace python examples/stream/cublas.py
import cupy

x = cupy.array([1, 2, 3])
y = cupy.array([[1], [2], [3]])
expected = cupy.matmul(x, y)
cupy.cuda.Device().synchronize()

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

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