numba-cuda 0.16.0__py3-none-any.whl → 0.18.0__py3-none-any.whl

numba_cuda/numba/cuda/mathimpl.py

@@ -1,11 +1,11 @@
 import math
 import operator
 from llvmlite import ir
-from numba.core import types, typing, cgutils, targetconfig
+from numba.core import types, typing, targetconfig
+from numba.cuda import cgutils
 from numba.core.imputils import Registry
 from numba.types import float32, float64, int64, uint64
 from numba.cuda import libdevice
-from numba import cuda
 
 registry = Registry()
 lower = registry.lower
@@ -89,94 +89,6 @@ def math_isfinite_int(context, builder, sig, args):
     return context.get_constant(types.boolean, 1)
 
 
-@lower(math.sin, types.float16)
-def fp16_sin_impl(context, builder, sig, args):
-    def fp16_sin(x):
-        return cuda.fp16.hsin(x)
-
-    return context.compile_internal(builder, fp16_sin, sig, args)
-
-
-@lower(math.cos, types.float16)
-def fp16_cos_impl(context, builder, sig, args):
-    def fp16_cos(x):
-        return cuda.fp16.hcos(x)
-
-    return context.compile_internal(builder, fp16_cos, sig, args)
-
-
-@lower(math.log, types.float16)
-def fp16_log_impl(context, builder, sig, args):
-    def fp16_log(x):
-        return cuda.fp16.hlog(x)
-
-    return context.compile_internal(builder, fp16_log, sig, args)
-
-
-@lower(math.log10, types.float16)
-def fp16_log10_impl(context, builder, sig, args):
-    def fp16_log10(x):
-        return cuda.fp16.hlog10(x)
-
-    return context.compile_internal(builder, fp16_log10, sig, args)
-
-
-@lower(math.log2, types.float16)
-def fp16_log2_impl(context, builder, sig, args):
-    def fp16_log2(x):
-        return cuda.fp16.hlog2(x)
-
-    return context.compile_internal(builder, fp16_log2, sig, args)
-
-
-@lower(math.exp, types.float16)
-def fp16_exp_impl(context, builder, sig, args):
-    def fp16_exp(x):
-        return cuda.fp16.hexp(x)
-
-    return context.compile_internal(builder, fp16_exp, sig, args)
-
-
-@lower(math.floor, types.float16)
-def fp16_floor_impl(context, builder, sig, args):
-    def fp16_floor(x):
-        return cuda.fp16.hfloor(x)
-
-    return context.compile_internal(builder, fp16_floor, sig, args)
-
-
-@lower(math.ceil, types.float16)
-def fp16_ceil_impl(context, builder, sig, args):
-    def fp16_ceil(x):
-        return cuda.fp16.hceil(x)
-
-    return context.compile_internal(builder, fp16_ceil, sig, args)
-
-
-@lower(math.sqrt, types.float16)
-def fp16_sqrt_impl(context, builder, sig, args):
-    def fp16_sqrt(x):
-        return cuda.fp16.hsqrt(x)
-
-    return context.compile_internal(builder, fp16_sqrt, sig, args)
-
-
-@lower(math.fabs, types.float16)
-def fp16_fabs_impl(context, builder, sig, args):
-    def fp16_fabs(x):
-        return cuda.fp16.habs(x)
-
-    return context.compile_internal(builder, fp16_fabs, sig, args)
-
-
-@lower(math.trunc, types.float16)
-def fp16_trunc_impl(context, builder, sig, args):
-    def fp16_trunc(x):
-        return cuda.fp16.htrunc(x)
-
-    return context.compile_internal(builder, fp16_trunc, sig, args)
-
-
 def impl_boolean(key, ty, libfunc):
     def lower_boolean_impl(context, builder, sig, args):
         libfunc_impl = context.get_function(

numba_cuda/numba/cuda/memory_management/nrt.py

@@ -143,7 +143,7 @@ class _Runtime:
             1,
             1,
             0,
-            stream.handle,
+            stream.handle.value,
             params,
             cooperative=False,
         )

numba_cuda/numba/cuda/nvvmutils.py

@@ -1,6 +1,7 @@
 import itertools
 from llvmlite import ir
-from numba.core import cgutils, targetconfig
+from numba.core import targetconfig
+from numba.cuda import cgutils
 from .cudadrv import nvvm
 
 

numba_cuda/numba/cuda/printimpl.py

@@ -1,6 +1,7 @@
 from functools import singledispatch
 from llvmlite import ir
-from numba.core import types, cgutils
+from numba.core import types
+from numba.cuda import cgutils
 from numba.core.errors import NumbaWarning
 from numba.core.imputils import Registry
 from numba.cuda import nvvmutils

numba_cuda/numba/cuda/serialize.py

@@ -0,0 +1,264 @@
+"""
+Serialization support for compiled functions.
+"""
+
+import sys
+import abc
+import io
+import copyreg
+
+
+import pickle
+from numba import cloudpickle
+from llvmlite import ir
+
+
+#
+# Pickle support
+#
+
+
+def _rebuild_reduction(cls, *args):
+    """
+    Global hook to rebuild a given class from its __reduce__ arguments.
+    """
+    return cls._rebuild(*args)
+
+
+# Keep unpickled object via `numba_unpickle` alive.
+_unpickled_memo = {}
+
+
+def _numba_unpickle(address, bytedata, hashed):
+    """Used by `numba_unpickle` from _helperlib.c
+
+    Parameters
+    ----------
+    address : int
+    bytedata : bytes
+    hashed : bytes
+
+    Returns
+    -------
+    obj : object
+        unpickled object
+    """
+    key = (address, hashed)
+    try:
+        obj = _unpickled_memo[key]
+    except KeyError:
+        _unpickled_memo[key] = obj = cloudpickle.loads(bytedata)
+    return obj
+
+
+def dumps(obj):
+    """Similar to `pickle.dumps()`. Returns the serialized object in bytes."""
+    pickler = NumbaPickler
+    with io.BytesIO() as buf:
+        p = pickler(buf, protocol=4)
+        p.dump(obj)
+        pickled = buf.getvalue()
+
+    return pickled
+
+
+def runtime_build_excinfo_struct(static_exc, exc_args):
+    exc, static_args, locinfo = cloudpickle.loads(static_exc)
+    real_args = []
+    exc_args_iter = iter(exc_args)
+    for arg in static_args:
+        if isinstance(arg, ir.Value):
+            real_args.append(next(exc_args_iter))
+        else:
+            real_args.append(arg)
+    return (exc, tuple(real_args), locinfo)
+
+
+# Alias to pickle.loads to allow `serialize.loads()`
+loads = cloudpickle.loads
+
+
+class _CustomPickled:
+    """A wrapper for objects that must be pickled with `NumbaPickler`.
+
+    Standard `pickle` will pick up the implementation registered via `copyreg`.
+    This will spawn a `NumbaPickler` instance to serialize the data.
+
+    `NumbaPickler` overrides the handling of this type so as not to spawn a
+    new pickler for the object when it is already being pickled by a
+    `NumbaPickler`.
+    """
+
+    __slots__ = "ctor", "states"
+
+    def __init__(self, ctor, states):
+        self.ctor = ctor
+        self.states = states
+
+    def _reduce(self):
+        return _CustomPickled._rebuild, (self.ctor, self.states)
+
+    @classmethod
+    def _rebuild(cls, ctor, states):
+        return cls(ctor, states)
+
+
+def _unpickle__CustomPickled(serialized):
+    """standard unpickling for `_CustomPickled`.
+
+    Uses `NumbaPickler` to load.
+    """
+    ctor, states = loads(serialized)
+    return _CustomPickled(ctor, states)
+
+
+def _pickle__CustomPickled(cp):
+    """standard pickling for `_CustomPickled`.
+
+    Uses `NumbaPickler` to dump.
+    """
+    serialized = dumps((cp.ctor, cp.states))
+    return _unpickle__CustomPickled, (serialized,)
+
+
+# Register custom pickling for the standard pickler.
+copyreg.pickle(_CustomPickled, _pickle__CustomPickled)
+
+
+def custom_reduce(cls, states):
+    """For customizing object serialization in `__reduce__`.
+
+    Object states provided here are used as keyword arguments to the
+    `._rebuild()` class method.
+
+    Parameters
+    ----------
+    states : dict
+        Dictionary of object states to be serialized.
+
+    Returns
+    -------
+    result : tuple
+        This tuple conforms to the return type requirement for `__reduce__`.
+    """
+    return custom_rebuild, (_CustomPickled(cls, states),)
+
+
+def custom_rebuild(custom_pickled):
+    """Customized object deserialization.
+
+    This function is referenced internally by `custom_reduce()`.
+    """
+    cls, states = custom_pickled.ctor, custom_pickled.states
+    return cls._rebuild(**states)
+
+
+def is_serialiable(obj):
+    """Check if *obj* can be serialized.
+
+    Parameters
+    ----------
+    obj : object
+
+    Returns
+    --------
+    can_serialize : bool
+    """
+    with io.BytesIO() as fout:
+        pickler = NumbaPickler(fout)
+        try:
+            pickler.dump(obj)
+        except pickle.PicklingError:
+            return False
+        else:
+            return True
+
+
+def _no_pickle(obj):
+    raise pickle.PicklingError(f"Pickling of {type(obj)} is unsupported")
+
+
+def disable_pickling(typ):
+    """This is called on a type to disable pickling"""
+    NumbaPickler.disabled_types.add(typ)
+    # Return `typ` to allow use as a decorator
+    return typ
+
+
+class NumbaPickler(cloudpickle.CloudPickler):
+    disabled_types = set()
+    """A set of types that pickling cannot is disabled.
+    """
+
+    def reducer_override(self, obj):
+        # Overridden to disable pickling of certain types
+        if type(obj) in self.disabled_types:
+            _no_pickle(obj)  # noreturn
+        return super().reducer_override(obj)
+
+
+def _custom_reduce__custompickled(cp):
+    return cp._reduce()
+
+
+NumbaPickler.dispatch_table[_CustomPickled] = _custom_reduce__custompickled
+
+
+class ReduceMixin(abc.ABC):
+    """A mixin class for objects that should be reduced by the NumbaPickler
+    instead of the standard pickler.
+    """
+
+    # Subclass MUST override the below methods
+
+    @abc.abstractmethod
+    def _reduce_states(self):
+        raise NotImplementedError
+
+    @abc.abstractclassmethod
+    def _rebuild(cls, **kwargs):
+        raise NotImplementedError
+
+    # Subclass can override the below methods
+
+    def _reduce_class(self):
+        return self.__class__
+
+    # Private methods
+
+    def __reduce__(self):
+        return custom_reduce(self._reduce_class(), self._reduce_states())
+
+
+class PickleCallableByPath:
+    """Wrap a callable object to be pickled by path to workaround limitation
+    in pickling due to non-pickleable objects in function non-locals.
+
+    Note:
+    - Do not use this as a decorator.
+    - Wrapped object must be a global that exist in its parent module and it
+      can be imported by `from the_module import the_object`.
+
+    Usage:
+
+    >>> def my_fn(x):
+    >>>     ...
+    >>> wrapped_fn = PickleCallableByPath(my_fn)
+    >>> # refer to `wrapped_fn` instead of `my_fn`
+    """
+
+    def __init__(self, fn):
+        self._fn = fn
+
+    def __call__(self, *args, **kwargs):
+        return self._fn(*args, **kwargs)
+
+    def __reduce__(self):
+        return type(self)._rebuild, (
+            self._fn.__module__,
+            self._fn.__name__,
+        )
+
+    @classmethod
+    def _rebuild(cls, modname, fn_path):
+        return cls(getattr(sys.modules[modname], fn_path))

numba_cuda/numba/cuda/simulator/__init__.py

@@ -33,6 +33,7 @@ del vector_types, name, svty, alias
 if config.ENABLE_CUDASIM:
     import sys
     from numba.cuda.simulator import cudadrv
+    from . import dispatcher
 
     sys.modules["numba.cuda.cudadrv"] = cudadrv
     sys.modules["numba.cuda.cudadrv.devicearray"] = cudadrv.devicearray
@@ -43,6 +44,7 @@ if config.ENABLE_CUDASIM:
     sys.modules["numba.cuda.cudadrv.drvapi"] = cudadrv.drvapi
     sys.modules["numba.cuda.cudadrv.error"] = cudadrv.error
     sys.modules["numba.cuda.cudadrv.nvvm"] = cudadrv.nvvm
+    sys.modules["numba.cuda.dispatcher"] = dispatcher
 
     from . import bf16, compiler, _internal
 

numba_cuda/numba/cuda/simulator/dispatcher.py

@@ -0,0 +1,7 @@
+class CUDADispatcher:
+    """
+    Dummy class so that consumers that try to import the real CUDADispatcher
+    do not get an import failure when running with the simulator.
+    """
+
+    ...