numba-cuda 0.22.0__cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

numba_cuda/numba/cuda/serialize.py

@@ -0,0 +1,267 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Serialization support for compiled functions.
+"""
+
+import sys
+import abc
+import io
+import copyreg
+
+
+import pickle
+from numba.cuda 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

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+import sys
+
+from .api import *
+from .vector_types import vector_types
+from .reduction import Reduce
+from .cudadrv.devicearray import (
+    device_array,
+    device_array_like,
+    pinned,
+    pinned_array,
+    pinned_array_like,
+    mapped_array,
+    to_device,
+    auto_device,
+)
+from .cudadrv import devicearray
+from .cudadrv.devices import require_context, gpus
+from .cudadrv.devices import get_context as current_context
+from .cudadrv.runtime import runtime
+from .cudadrv.linkable_code import LinkableCode
+from numba.cuda.core import config
+
+reduce = Reduce
+
+# Register simulated vector types as module level variables
+for name, svty in vector_types.items():
+    setattr(sys.modules[__name__], name, svty)
+    for alias in svty.aliases:
+        setattr(sys.modules[__name__], alias, svty)
+del vector_types, name, svty, alias
+
+# Ensure that any user code attempting to import cudadrv etc. gets the
+# simulator's version and not the real version if the simulator is enabled.
+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
+    sys.modules["numba.cuda.cudadrv.devices"] = cudadrv.devices
+    sys.modules["numba.cuda.cudadrv.driver"] = cudadrv.driver
+    sys.modules["numba.cuda.cudadrv.linkable_code"] = cudadrv.linkable_code
+    sys.modules["numba.cuda.cudadrv.runtime"] = cudadrv.runtime
+    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
+
+    sys.modules["numba.cuda.bf16"] = bf16
+    sys.modules["numba.cuda.compiler"] = compiler
+    sys.modules["numba.cuda._internal"] = _internal
+    sys.modules["numba.cuda._internal.cuda_bf16"] = _internal.cuda_bf16
+
+    from numba.cuda.simulator import memory_management
+
+    sys.modules["numba.cuda.memory_management"] = memory_management
+    sys.modules["numba.cuda.memory_management.nrt"] = memory_management.nrt

numba_cuda/numba/cuda/simulator/_internal/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.simulator._internal import cuda_bf16  # noqa: F401

numba_cuda/numba/cuda/simulator/_internal/cuda_bf16.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause

numba_cuda/numba/cuda/simulator/api.py

@@ -0,0 +1,179 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Contains CUDA API functions
+"""
+
+# Imports here bring together parts of the API from other modules, so some of
+# them appear unused.
+from contextlib import contextmanager
+
+from .cudadrv.devices import require_context, reset, gpus  # noqa: F401
+from .cudadrv.linkable_code import (
+    PTXSource,  # noqa: F401
+    CUSource,  # noqa: F401
+    Cubin,  # noqa: F401
+    Fatbin,  # noqa: F401
+    Archive,  # noqa: F401
+    Object,  # noqa: F401
+    LTOIR,  # noqa: F401
+)  # noqa: F401
+from .kernel import FakeCUDAKernel
+from numba.cuda.core import config
+from numba.cuda.core.sigutils import is_signature
+from ..args import In, Out, InOut  # noqa: F401
+
+
+def select_device(dev=0):
+    assert dev == 0, "Only a single device supported by the simulator"
+
+
+def is_float16_supported():
+    return True
+
+
+def is_bfloat16_supported():
+    return False
+
+
+class stream(object):
+    """
+    The stream API is supported in the simulator - however, all execution
+    occurs synchronously, so synchronization requires no operation.
+    """
+
+    @contextmanager
+    def auto_synchronize(self):
+        yield
+
+    def synchronize(self):
+        pass
+
+
+# Default stream APIs. Since execution from the perspective of the host is
+# synchronous in the simulator, these can be the same as the stream class.
+default_stream = stream
+legacy_default_stream = stream
+per_thread_default_stream = stream
+
+
+# There is no way to use external streams with the simulator. Since the
+# implementation is not really using streams, we can't meaningfully interact
+# with external ones.
+def external_stream(ptr):
+    raise RuntimeError("External streams are unsupported in the simulator")
+
+
+def synchronize():
+    pass
+
+
+def close():
+    gpus.closed = True
+
+
+def declare_device(*args, **kwargs):
+    pass
+
+
+def detect():
+    print("Found 1 CUDA devices")
+    print("id %d    %20s %40s" % (0, "SIMULATOR", "[SUPPORTED]"))
+    print("%40s: 5.0" % "compute capability")
+
+
+def list_devices():
+    return gpus
+
+
+def get_current_device():
+    return gpus[0].device
+
+
+# Events
+
+
+class Event(object):
+    """
+    The simulator supports the event API, but they do not record timing info,
+    and all simulation is synchronous. Execution time is not recorded.
+    """
+
+    def record(self, stream=0):
+        pass
+
+    def wait(self, stream=0):
+        pass
+
+    def synchronize(self):
+        pass
+
+    def elapsed_time(self, event):
+        """This is here to preserve the API; the output is meaningless."""
+        return -1.0
+
+
+event = Event
+
+
+def jit(
+    func_or_sig=None,
+    device=False,
+    debug=None,
+    argtypes=None,
+    inline=False,
+    restype=None,
+    fastmath=False,
+    link=None,
+    boundscheck=None,
+    opt=None,
+    cache=None,
+):
+    # Here for API compatibility
+    if boundscheck:
+        raise NotImplementedError("bounds checking is not supported for CUDA")
+
+    if link is not None:
+        raise NotImplementedError("Cannot link PTX in the simulator")
+
+    debug = config.CUDA_DEBUGINFO_DEFAULT if debug is None else debug
+
+    # Check for first argument specifying types - in that case the
+    # decorator is not being passed a function
+    if (
+        func_or_sig is None
+        or is_signature(func_or_sig)
+        or isinstance(func_or_sig, list)
+    ):
+
+        def jitwrapper(fn):
+            return FakeCUDAKernel(
+                fn, device=device, fastmath=fastmath, debug=debug
+            )
+
+        return jitwrapper
+    return FakeCUDAKernel(func_or_sig, device=device, debug=debug)
+
+
+@contextmanager
+def defer_cleanup():
+    # No effect for simulator
+    yield
+
+
+def is_supported_version():
+    return True
+
+
+# cache hints operations
+ldca = None
+ldcg = None
+ldcs = None
+ldlu = None
+ldcv = None
+
+stcg = None
+stcs = None
+stwb = None
+stwt = None

numba_cuda/numba/cuda/simulator/bf16.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+bfloat16 = None

numba_cuda/numba/cuda/simulator/compiler.py

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+The compiler is not implemented in the simulator. This module provides a stub
+to allow tests to import successfully.
+"""
+
+compile = None
+compile_for_current_device = None
+compile_ptx = None
+compile_ptx_for_current_device = None
+declare_device_function = None
+compile_all = None
+
+
+def run_frontend(func):
+    pass
+
+
+class DefaultPassBuilder(object):
+    @staticmethod
+    def define_nopython_lowering_pipeline(state, name="nopython_lowering"):
+        pass
+
+    @staticmethod
+    def define_typed_pipeline(state, name="typed"):
+        pass
+
+
+class CompilerBase:
+    def __init__(
+        self, typingctx, targetctx, library, args, return_type, flags, locals
+    ):
+        pass
+
+
+PassManager = None

numba_cuda/numba/cuda/simulator/cudadrv/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba.cuda.simulator.cudadrv import (
+    devicearray,
+    devices,
+    driver,
+    drvapi,
+    error,
+    nvvm,
+)