-
Notifications
You must be signed in to change notification settings - Fork 0
/
stream.py
114 lines (81 loc) · 2.48 KB
/
stream.py
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
# -*- coding: utf-8 -*-
__all__ = [
"Stream",
]
import numpy as np
# Local imports
from shared import (get_nbytes,
Shared) #Shared calls between Device and Stream
from cublas_helpers import cublas
from cufft_helpers.cufft import cufft
from cuda_helpers import (cu_memcpy_3d_async,
cu_stream_create,
cu_sync_stream)
class Stream(Shared, object):
def __init__(self, device, stream_id):
"""
CUDA stream object.
Parameters
----------
device : object
The CUDA device object.
stream_id : int
The CUDA set stream id.
Attributes
----------
id : int
Stream ID.
stream : c_void_p (cudaStream_t*)
Pointer to the CUDA Stream handle.
cublas: object
The callable cuBLAS object.
cufft : object
The callable cuFFT object.
"""
super(Stream, self).__init__()
self._device = device
self._stream = cu_stream_create()
self._id = stream_id
self._cublas = cublas(self.stream)
self._cufft = cufft(self.stream)
def malloc(self, shape, dtype=None, stream=None, fill=None):
"""
Allocates device memory.
Parameters
----------
shape : tuple
The shape of the array to allocate.
dtype : np.dtype
That data type of the array.
fill : scalar or np.ndarray, optional
Default value to set allocated array to.
stream : c_void_p, optional
CUDA stream to associate the returned object with.
Returns
-------
dev_ptr: c_void_p
Pointer to allocated device memory.
"""
stream = stream or self.stream
dtype = dtype or self.device._default_dtype
return self.device.malloc(shape, dtype, fill, stream)
def sync(self):
"""
Block the host thread until the stream has completed its task.
"""
cu_sync_stream(self.stream)
@property
def cublas(self):
return self._cublas
@property
def cufft(self):
return self._cufft
@property
def device(self):
return self._device
@property
def id(self):
return self.stream_id
@property
def stream(self):
return self._stream