-
Notifications
You must be signed in to change notification settings - Fork 0
/
shared.py
110 lines (86 loc) · 2.87 KB
/
shared.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
# -*- coding: utf-8 -*-
__all__ = [
"get_nbytes",
"Shared",
]
from functools import reduce
from operator import mul
import numpy as np
from ctypes import cast, c_void_p
import warnings
# Local imports
from cuda_helpers import (cu_create_channel_char,
cu_create_channel_short,
cu_create_channel_float,
cu_malloc,
cu_malloc_3d,
cu_malloc_managed)
from dev_ptr import Device_Ptr
from uni_ptr import Unified_Ptr
from shared_utils import Mapping
def get_nbytes(arr, nbytes=None):
if type(arr) in [list or tuple]:
nbytes = nbytes or np.array(arr).nbytes
else:
nbytes = nbytes or arr.nbytes
return nbytes
class Shared(Mapping, object):
def __init__(self):
"""
This object holds methods that are callable by both the
Device and stream Objects.
Notes
-----
Calling these methods from the Stream object may cause
thread host blocking behavior and break the asynchronous
stream operation.
"""
super(Mapping, self).__init__()
def create_channel(self, dtype, components=1, unsigned=False):
"""
Create channel information used by a CUDA texture.
Parameters
----------
dtype : int
dtype identifier.
components : int, optional
Number of components to create in the channel.
unsigned : bool, optional
Unsigned flag.
Returns
-------
channelDesc : ctypes Structure
CUDA channel object that is passed into the CUDA 2/3d
memcpy functions.
"""
return {
0: cu_create_channel_char(components, unsigned),
1: cu_create_channel_short(components, unsigned),
2: cu_create_channel_float(components, unsigned),
}.get(dtype, 0)
def malloc_3d(self, channel, extent, layered=False):
"""
Allocates device memory as a CUDA array.
Parameters
----------
nbytes : int
Size to allocate in bytes.
extent : list, 1d np.ndarray, or tuple
The extent or dimensions of the array.
Format: (nx,ny,nz,...)
layered : bool, optional
Layered array flag.
Returns
-------
dev_ptr: c_void_p
Pointer to allocated device memory.
Notes
-----
Setting the layered flag to True will turn a 3D
array into a 2D layered array.
"""
if type(extent) in [list, tuple]:
extent = np.array(extent, dtype='i4')
dev_ptr = cu_malloc_3d(channel, extent, layered)
dev_ptr = cast(dev_ptr, c_void_p)
return dev_ptr