/
_pinned_array.py
141 lines (108 loc) · 4.98 KB
/
_pinned_array.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
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
import numpy
from cupy import cuda
from cupy._creation.basic import _new_like_order_and_strides
from cupy._core import internal
def _update_shape(a, shape):
if shape is None and a is not None:
shape = a.shape
elif isinstance(shape, int):
shape = (shape,)
else:
shape = tuple(shape)
return shape
def empty_pinned(shape, dtype=float, order='C'):
"""Returns a new, uninitialized NumPy array with the given shape
and dtype.
This is a convenience function which is just :func:`numpy.empty`,
except that the underlying memory is pinned/pagelocked.
Args:
shape (int or tuple of ints): Dimensionalities of the array.
dtype: Data type specifier.
order ({'C', 'F'}): Row-major (C-style) or column-major
(Fortran-style) order.
Returns:
numpy.ndarray: A new array with elements not initialized.
.. seealso:: :func:`numpy.empty`
"""
shape = _update_shape(None, shape)
nbytes = internal.prod(shape) * numpy.dtype(dtype).itemsize
mem = cuda.alloc_pinned_memory(nbytes)
out = numpy.ndarray(shape, dtype=dtype, buffer=mem, order=order)
return out
def empty_like_pinned(a, dtype=None, order='K', subok=None, shape=None):
"""Returns a new, uninitialized NumPy array with the same shape and dtype
as those of the given array.
This is a convenience function which is just :func:`numpy.empty_like`,
except that the underlying memory is pinned/pagelocked.
This function currently does not support ``subok`` option.
Args:
a (numpy.ndarray or cupy.ndarray): Base array.
dtype: Data type specifier. The data type of ``a`` is used by default.
order ({'C', 'F', 'A', or 'K'}): Overrides the memory layout of the
result. ``'C'`` means C-order, ``'F'`` means F-order, ``'A'`` means
``'F'`` if ``a`` is Fortran contiguous, ``'C'`` otherwise.
``'K'`` means match the layout of ``a`` as closely as possible.
subok: Not supported yet, must be None.
shape (int or tuple of ints): Overrides the shape of the result. If
``order='K'`` and the number of dimensions is unchanged, will try
to keep order, otherwise, ``order='C'`` is implied.
Returns:
numpy.ndarray: A new array with same shape and dtype of ``a`` with
elements not initialized.
.. seealso:: :func:`numpy.empty_like`
"""
# We're kinda duplicating the code here because order='K' needs special
# treatment: strides need to be computed
if subok is not None:
raise TypeError('subok is not supported yet')
if dtype is None:
dtype = a.dtype
shape = _update_shape(a, shape)
order, strides, _ = _new_like_order_and_strides(
a, dtype, order, shape, get_memptr=False)
nbytes = internal.prod(shape) * numpy.dtype(dtype).itemsize
mem = cuda.alloc_pinned_memory(nbytes)
out = numpy.ndarray(shape, dtype=dtype, buffer=mem,
strides=strides, order=order)
return out
def zeros_pinned(shape, dtype=float, order='C'):
"""Returns a new, zero-initialized NumPy array with the given shape
and dtype.
This is a convenience function which is just :func:`numpy.zeros`,
except that the underlying memory is pinned/pagelocked.
Args:
shape (int or tuple of ints): Dimensionalities of the array.
dtype: Data type specifier.
order ({'C', 'F'}): Row-major (C-style) or column-major
(Fortran-style) order.
Returns:
numpy.ndarray: An array filled with zeros.
.. seealso:: :func:`numpy.zeros`
"""
out = empty_pinned(shape, dtype, order)
numpy.copyto(out, 0, casting='unsafe')
return out
def zeros_like_pinned(a, dtype=None, order='K', subok=None, shape=None):
"""Returns a new, zero-initialized NumPy array with the same shape and dtype
as those of the given array.
This is a convenience function which is just :func:`numpy.zeros_like`,
except that the underlying memory is pinned/pagelocked.
This function currently does not support ``subok`` option.
Args:
a (numpy.ndarray or cupy.ndarray): Base array.
dtype: Data type specifier. The dtype of ``a`` is used by default.
order ({'C', 'F', 'A', or 'K'}): Overrides the memory layout of the
result. ``'C'`` means C-order, ``'F'`` means F-order, ``'A'`` means
``'F'`` if ``a`` is Fortran contiguous, ``'C'`` otherwise.
``'K'`` means match the layout of ``a`` as closely as possible.
subok: Not supported yet, must be None.
shape (int or tuple of ints): Overrides the shape of the result. If
``order='K'`` and the number of dimensions is unchanged, will try
to keep order, otherwise, ``order='C'`` is implied.
Returns:
numpy.ndarray: An array filled with zeros.
.. seealso:: :func:`numpy.zeros_like`
"""
out = empty_like_pinned(a, dtype, order, subok, shape)
numpy.copyto(out, 0, casting='unsafe')
return out