numba-cuda 0.0.0__py3-none-any.whl → 0.0.12__py3-none-any.whl

numba_cuda/numba/cuda/cudamath.py

@@ -0,0 +1,140 @@
+import math
+from numba.core import types
+from numba.core.typing.templates import ConcreteTemplate, signature, Registry
+
+
+registry = Registry()
+infer_global = registry.register_global
+
+
+@infer_global(math.acos)
+@infer_global(math.acosh)
+@infer_global(math.asin)
+@infer_global(math.asinh)
+@infer_global(math.atan)
+@infer_global(math.atanh)
+@infer_global(math.cosh)
+@infer_global(math.degrees)
+@infer_global(math.erf)
+@infer_global(math.erfc)
+@infer_global(math.expm1)
+@infer_global(math.gamma)
+@infer_global(math.lgamma)
+@infer_global(math.log1p)
+@infer_global(math.radians)
+@infer_global(math.sinh)
+@infer_global(math.tanh)
+@infer_global(math.tan)
+class Math_unary(ConcreteTemplate):
+    cases = [
+        signature(types.float64, types.int64),
+        signature(types.float64, types.uint64),
+        signature(types.float32, types.float32),
+        signature(types.float64, types.float64),
+    ]
+
+
+@infer_global(math.sin)
+@infer_global(math.cos)
+@infer_global(math.ceil)
+@infer_global(math.floor)
+@infer_global(math.sqrt)
+@infer_global(math.log)
+@infer_global(math.log2)
+@infer_global(math.log10)
+@infer_global(math.exp)
+@infer_global(math.fabs)
+@infer_global(math.trunc)
+class Math_unary_with_fp16(ConcreteTemplate):
+    cases = [
+        signature(types.float64, types.int64),
+        signature(types.float64, types.uint64),
+        signature(types.float32, types.float32),
+        signature(types.float64, types.float64),
+        signature(types.float16, types.float16),
+    ]
+
+
+@infer_global(math.atan2)
+class Math_atan2(ConcreteTemplate):
+    key = math.atan2
+    cases = [
+        signature(types.float64, types.int64, types.int64),
+        signature(types.float64, types.uint64, types.uint64),
+        signature(types.float32, types.float32, types.float32),
+        signature(types.float64, types.float64, types.float64),
+    ]
+
+
+@infer_global(math.hypot)
+class Math_hypot(ConcreteTemplate):
+    key = math.hypot
+    cases = [
+        signature(types.float64, types.int64, types.int64),
+        signature(types.float64, types.uint64, types.uint64),
+        signature(types.float32, types.float32, types.float32),
+        signature(types.float64, types.float64, types.float64),
+    ]
+
+
+@infer_global(math.copysign)
+@infer_global(math.fmod)
+class Math_binary(ConcreteTemplate):
+    cases = [
+        signature(types.float32, types.float32, types.float32),
+        signature(types.float64, types.float64, types.float64),
+    ]
+
+
+@infer_global(math.remainder)
+class Math_remainder(ConcreteTemplate):
+    cases = [
+        signature(types.float32, types.float32, types.float32),
+        signature(types.float64, types.float64, types.float64),
+    ]
+
+
+@infer_global(math.pow)
+class Math_pow(ConcreteTemplate):
+    cases = [
+        signature(types.float32, types.float32, types.float32),
+        signature(types.float64, types.float64, types.float64),
+        signature(types.float32, types.float32, types.int32),
+        signature(types.float64, types.float64, types.int32),
+    ]
+
+
+@infer_global(math.frexp)
+class Math_frexp(ConcreteTemplate):
+    cases = [
+        signature(types.Tuple([types.float32, types.int32]), types.float32),
+        signature(types.Tuple([types.float64, types.int32]), types.float64),
+    ]
+
+
+@infer_global(math.ldexp)
+class Math_ldexp(ConcreteTemplate):
+    cases = [
+        signature(types.float32, types.float32, types.int32),
+        signature(types.float64, types.float64, types.int32),
+    ]
+
+
+@infer_global(math.isinf)
+@infer_global(math.isnan)
+@infer_global(math.isfinite)
+class Math_isnan(ConcreteTemplate):
+    cases = [
+        signature(types.boolean, types.int64),
+        signature(types.boolean, types.uint64),
+        signature(types.boolean, types.float32),
+        signature(types.boolean, types.float64),
+    ]
+
+
+@infer_global(math.modf)
+class Math_modf(ConcreteTemplate):
+    cases = [
+        signature(types.UniTuple(types.float64, 2), types.float64),
+        signature(types.UniTuple(types.float32, 2), types.float32)
+    ]

numba_cuda/numba/cuda/decorators.py

@@ -0,0 +1,189 @@
+from warnings import warn
+from numba.core import types, config, sigutils
+from numba.core.errors import DeprecationError, NumbaInvalidConfigWarning
+from numba.cuda.compiler import declare_device_function
+from numba.cuda.dispatcher import CUDADispatcher
+from numba.cuda.simulator.kernel import FakeCUDAKernel
+
+
+_msg_deprecated_signature_arg = ("Deprecated keyword argument `{0}`. "
+                                 "Signatures should be passed as the first "
+                                 "positional argument.")
+
+
+def jit(func_or_sig=None, device=False, inline=False, link=[], debug=None,
+        opt=True, lineinfo=False, cache=False, **kws):
+    """
+    JIT compile a Python function for CUDA GPUs.
+
+    :param func_or_sig: A function to JIT compile, or *signatures* of a
+       function to compile. If a function is supplied, then a
+       :class:`Dispatcher <numba.cuda.dispatcher.CUDADispatcher>` is returned.
+       Otherwise, ``func_or_sig`` may be a signature or a list of signatures,
+       and a function is returned. The returned function accepts another
+       function, which it will compile and then return a :class:`Dispatcher
+       <numba.cuda.dispatcher.CUDADispatcher>`. See :ref:`jit-decorator` for
+       more information about passing signatures.
+
+       .. note:: A kernel cannot have any return value.
+    :param device: Indicates whether this is a device function.
+    :type device: bool
+    :param link: A list of files containing PTX or CUDA C/C++ source to link
+       with the function
+    :type link: list
+    :param debug: If True, check for exceptions thrown when executing the
+       kernel. Since this degrades performance, this should only be used for
+       debugging purposes. If set to True, then ``opt`` should be set to False.
+       Defaults to False.  (The default value can be overridden by setting
+       environment variable ``NUMBA_CUDA_DEBUGINFO=1``.)
+    :param fastmath: When True, enables fastmath optimizations as outlined in
+       the :ref:`CUDA Fast Math documentation <cuda-fast-math>`.
+    :param max_registers: Request that the kernel is limited to using at most
+       this number of registers per thread. The limit may not be respected if
+       the ABI requires a greater number of registers than that requested.
+       Useful for increasing occupancy.
+    :param opt: Whether to compile from LLVM IR to PTX with optimization
+                enabled. When ``True``, ``-opt=3`` is passed to NVVM. When
+                ``False``, ``-opt=0`` is passed to NVVM. Defaults to ``True``.
+    :type opt: bool
+    :param lineinfo: If True, generate a line mapping between source code and
+       assembly code. This enables inspection of the source code in NVIDIA
+       profiling tools and correlation with program counter sampling.
+    :type lineinfo: bool
+    :param cache: If True, enables the file-based cache for this function.
+    :type cache: bool
+    """
+
+    if link and config.ENABLE_CUDASIM:
+        raise NotImplementedError('Cannot link PTX in the simulator')
+
+    if kws.get('boundscheck'):
+        raise NotImplementedError("bounds checking is not supported for CUDA")
+
+    if kws.get('argtypes') is not None:
+        msg = _msg_deprecated_signature_arg.format('argtypes')
+        raise DeprecationError(msg)
+    if kws.get('restype') is not None:
+        msg = _msg_deprecated_signature_arg.format('restype')
+        raise DeprecationError(msg)
+    if kws.get('bind') is not None:
+        msg = _msg_deprecated_signature_arg.format('bind')
+        raise DeprecationError(msg)
+
+    debug = config.CUDA_DEBUGINFO_DEFAULT if debug is None else debug
+    fastmath = kws.get('fastmath', False)
+    extensions = kws.get('extensions', [])
+
+    if debug and opt:
+        msg = ("debug=True with opt=True (the default) "
+               "is not supported by CUDA. This may result in a crash"
+               " - set debug=False or opt=False.")
+        warn(NumbaInvalidConfigWarning(msg))
+
+    if debug and lineinfo:
+        msg = ("debug and lineinfo are mutually exclusive. Use debug to get "
+               "full debug info (this disables some optimizations), or "
+               "lineinfo for line info only with code generation unaffected.")
+        warn(NumbaInvalidConfigWarning(msg))
+
+    if device and kws.get('link'):
+        raise ValueError("link keyword invalid for device function")
+
+    if sigutils.is_signature(func_or_sig):
+        signatures = [func_or_sig]
+        specialized = True
+    elif isinstance(func_or_sig, list):
+        signatures = func_or_sig
+        specialized = False
+    else:
+        signatures = None
+
+    if signatures is not None:
+        if config.ENABLE_CUDASIM:
+            def jitwrapper(func):
+                return FakeCUDAKernel(func, device=device, fastmath=fastmath)
+            return jitwrapper
+
+        def _jit(func):
+            targetoptions = kws.copy()
+            targetoptions['debug'] = debug
+            targetoptions['lineinfo'] = lineinfo
+            targetoptions['link'] = link
+            targetoptions['opt'] = opt
+            targetoptions['fastmath'] = fastmath
+            targetoptions['device'] = device
+            targetoptions['extensions'] = extensions
+
+            disp = CUDADispatcher(func, targetoptions=targetoptions)
+
+            if cache:
+                disp.enable_caching()
+
+            for sig in signatures:
+                argtypes, restype = sigutils.normalize_signature(sig)
+
+                if restype and not device and restype != types.void:
+                    raise TypeError("CUDA kernel must have void return type.")
+
+                if device:
+                    from numba.core import typeinfer
+                    with typeinfer.register_dispatcher(disp):
+                        disp.compile_device(argtypes, restype)
+                else:
+                    disp.compile(argtypes)
+
+            disp._specialized = specialized
+            disp.disable_compile()
+
+            return disp
+
+        return _jit
+    else:
+        if func_or_sig is None:
+            if config.ENABLE_CUDASIM:
+                def autojitwrapper(func):
+                    return FakeCUDAKernel(func, device=device,
+                                          fastmath=fastmath)
+            else:
+                def autojitwrapper(func):
+                    return jit(func, device=device, debug=debug, opt=opt,
+                               lineinfo=lineinfo, link=link, cache=cache, **kws)
+
+            return autojitwrapper
+        # func_or_sig is a function
+        else:
+            if config.ENABLE_CUDASIM:
+                return FakeCUDAKernel(func_or_sig, device=device,
+                                      fastmath=fastmath)
+            else:
+                targetoptions = kws.copy()
+                targetoptions['debug'] = debug
+                targetoptions['lineinfo'] = lineinfo
+                targetoptions['opt'] = opt
+                targetoptions['link'] = link
+                targetoptions['fastmath'] = fastmath
+                targetoptions['device'] = device
+                targetoptions['extensions'] = extensions
+                disp = CUDADispatcher(func_or_sig, targetoptions=targetoptions)
+
+                if cache:
+                    disp.enable_caching()
+
+                return disp
+
+
+def declare_device(name, sig):
+    """
+    Declare the signature of a foreign function. Returns a descriptor that can
+    be used to call the function from a Python kernel.
+
+    :param name: The name of the foreign function.
+    :type name: str
+    :param sig: The Numba signature of the function.
+    """
+    argtypes, restype = sigutils.normalize_signature(sig)
+    if restype is None:
+        msg = 'Return type must be provided for device declarations'
+        raise TypeError(msg)
+
+    return declare_device_function(name, restype, argtypes)

numba_cuda/numba/cuda/descriptor.py

@@ -0,0 +1,33 @@
+from numba.core.descriptors import TargetDescriptor
+from numba.core.options import TargetOptions
+from .target import CUDATargetContext, CUDATypingContext
+
+
+class CUDATargetOptions(TargetOptions):
+    pass
+
+
+class CUDATarget(TargetDescriptor):
+    def __init__(self, name):
+        self.options = CUDATargetOptions
+        # The typing and target contexts are initialized only when needed -
+        # this prevents an attempt to load CUDA libraries at import time on
+        # systems that might not have them present.
+        self._typingctx = None
+        self._targetctx = None
+        super().__init__(name)
+
+    @property
+    def typing_context(self):
+        if self._typingctx is None:
+            self._typingctx = CUDATypingContext()
+        return self._typingctx
+
+    @property
+    def target_context(self):
+        if self._targetctx is None:
+            self._targetctx = CUDATargetContext(self._typingctx)
+        return self._targetctx
+
+
+cuda_target = CUDATarget('cuda')

numba_cuda/numba/cuda/device_init.py

@@ -0,0 +1,89 @@
+# Re export
+import sys
+from numba.cuda import cg
+from .stubs import (threadIdx, blockIdx, blockDim, gridDim, laneid, warpsize,
+                    syncwarp, shared, local, const, atomic,
+                    shfl_sync_intrinsic, vote_sync_intrinsic, match_any_sync,
+                    match_all_sync, threadfence_block, threadfence_system,
+                    threadfence, selp, popc, brev, clz, ffs, fma, cbrt,
+                    activemask, lanemask_lt, nanosleep, fp16,
+                    _vector_type_stubs)
+from .intrinsics import (grid, gridsize, syncthreads, syncthreads_and,
+                         syncthreads_count, syncthreads_or)
+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
+
+# Expose vector type constructors and aliases as module level attributes.
+for vector_type_stub in _vector_type_stubs:
+    setattr(sys.modules[__name__], vector_type_stub.__name__, vector_type_stub)
+    for alias in vector_type_stub.aliases:
+        setattr(sys.modules[__name__], alias, vector_type_stub)
+del vector_type_stub, _vector_type_stubs
+
+
+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 is_supported_version():
+    """Returns True if the CUDA Runtime is a supported version.
+
+    Unsupported versions (e.g. newer versions than those known to Numba)
+    may still work; this function provides a facility to check whether the
+    current Numba version is tested and known to work with the current
+    runtime version. If the current version is unsupported, the caller can
+    decide how to act. Options include:
+
+    - Continuing silently,
+    - Emitting a warning,
+    - Generating an error or otherwise preventing the use of CUDA.
+    """
+
+    return runtime.is_supported_version()
+
+
+def cuda_error():
+    """Returns None if there was no error initializing the CUDA driver.
+    If there was an error initializing the driver, a string describing the
+    error is returned.
+    """
+    return driver.driver.initialization_error
+
+
+initialize.initialize_all()