/
device_init.py
59 lines (47 loc) · 2.21 KB
/
device_init.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
# Re export
from .stubs import (threadIdx, blockIdx, blockDim, gridDim, laneid,
warpsize, syncthreads, syncthreads_count, syncwarp,
syncthreads_and, syncthreads_or, shared, local,
const, grid, gridsize, atomic, shfl_sync_intrinsic,
vote_sync_intrinsic, match_any_sync, match_all_sync,
threadfence_block, threadfence_system,
threadfence, selp, popc, brev, clz, ffs, fma)
from .cudadrv.error import CudaSupportError
from numba.cuda.cudadrv.driver import (BaseCUDAMemoryManager,
HostOnlyCUDAMemoryManager,
GetIpcHandleMixin, MemoryPointer,
MappedMemory, PinnedMemory, MemoryInfo,
IpcHandle, set_memory_manager)
from numba.cuda.cudadrv.runtime import runtime
from .cudadrv import nvvm
from numba.cuda import initialize
from .errors import KernelRuntimeError
from .decorators import jit, declare_device
from .api import *
from .api import _auto_device
from .args import In, Out, InOut
from .intrinsic_wrapper import (all_sync, any_sync, eq_sync, ballot_sync,
shfl_sync, shfl_up_sync, shfl_down_sync,
shfl_xor_sync)
from .kernels import reduction
reduce = Reduce = reduction.Reduce
def is_available():
"""Returns a boolean to indicate the availability of a CUDA GPU.
This will initialize the driver if it hasn't been initialized.
"""
# whilst `driver.is_available` will init the driver itself,
# the driver initialization may raise and as a result break
# test discovery/orchestration as `cuda.is_available` is often
# used as a guard for whether to run a CUDA test, the try/except
# below is to handle this case.
driver_is_available = False
try:
driver_is_available = driver.driver.is_available
except CudaSupportError:
pass
return driver_is_available and nvvm.is_available()
def cuda_error():
"""Returns None or an exception if the CUDA driver fails to initialize.
"""
return driver.driver.initialization_error
initialize.initialize_all()