numba-cuda 0.21.1__cp313-cp313-win_amd64.whl

numba_cuda/numba/cuda/vector_types.py

@@ -0,0 +1,214 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+# CUDA built-in Vector Types
+# https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#built-in-vector-types
+
+from typing import List, Tuple, Dict
+
+from numba.cuda import types
+from numba.cuda import cgutils
+from numba.cuda.datamodel import models
+from numba.cuda.core.imputils import Registry as ImplRegistry
+from numba.cuda.typing.templates import ConcreteTemplate
+from numba.cuda.typing.templates import Registry as TypingRegistry
+from numba.cuda.typing.templates import signature
+from numba.cuda import stubs
+from numba.cuda.errors import CudaLoweringError
+from numba.cuda.extending import make_attribute_wrapper, register_model
+
+
+typing_registry = TypingRegistry()
+impl_registry = ImplRegistry()
+
+register = typing_registry.register
+register_attr = typing_registry.register_attr
+register_global = typing_registry.register_global
+lower = impl_registry.lower
+
+
+class VectorType(types.Type):
+    def __init__(self, name, base_type, attr_names, user_facing_object):
+        self._base_type = base_type
+        self._attr_names = attr_names
+        self._user_facing_object = user_facing_object
+        super().__init__(name=name)
+
+    @property
+    def base_type(self):
+        return self._base_type
+
+    @property
+    def attr_names(self):
+        return self._attr_names
+
+    @property
+    def num_elements(self):
+        return len(self._attr_names)
+
+    @property
+    def user_facing_object(self):
+        return self._user_facing_object
+
+
+def make_vector_type(
+    name: str,
+    base_type: types.Type,
+    attr_names: Tuple[str, ...],
+    user_facing_object,
+) -> types.Type:
+    """Create a vector type.
+
+    Parameters
+    ----------
+    name: str
+        The name of the type.
+    base_type: numba.cuda.types.Type
+        The primitive type for each element in the vector.
+    attr_names: tuple of str
+        Name for each attribute.
+    user_facing_object: object
+        The handle to be used in cuda kernel.
+    """
+
+    class _VectorType(VectorType):
+        """Internal instantiation of VectorType."""
+
+        pass
+
+    class VectorTypeModel(models.StructModel):
+        def __init__(self, dmm, fe_type):
+            members = [(attr_name, base_type) for attr_name in attr_names]
+            super().__init__(dmm, fe_type, members)
+
+    vector_type = _VectorType(name, base_type, attr_names, user_facing_object)
+    register_model(_VectorType)(VectorTypeModel)
+    for attr_name in attr_names:
+        make_attribute_wrapper(_VectorType, attr_name, attr_name)
+
+    return vector_type
+
+
+def enable_vector_type_ctor(
+    vector_type: VectorType, overloads: List[List[types.Type]]
+):
+    """Create typing and lowering for vector type constructor.
+
+    Parameters
+    ----------
+    vector_type: VectorType
+        The type whose constructor to type and lower.
+    overloads: List of argument types
+        A list containing different overloads of the constructor. Each base type
+        in the argument list should either be primitive type or VectorType.
+    """
+    ctor = vector_type.user_facing_object
+
+    @register
+    class CtorTemplate(ConcreteTemplate):
+        key = ctor
+        cases = [signature(vector_type, *arglist) for arglist in overloads]
+
+    register_global(ctor, types.Function(CtorTemplate))
+
+    # Lowering
+
+    def make_lowering(fml_arg_list):
+        """Meta function to create a lowering for the constructor. Flattens
+        the arguments by converting vector_type into load instructions for each
+        of its attributes. Such as float2 -> float2.x, float2.y.
+        """
+
+        def lowering(context, builder, sig, actual_args):
+            # A list of elements to assign from
+            source_list = []
+            # Convert the list of argument types to a list of load IRs.
+            for argidx, fml_arg in enumerate(fml_arg_list):
+                if isinstance(fml_arg, VectorType):
+                    pxy = cgutils.create_struct_proxy(fml_arg)(
+                        context, builder, actual_args[argidx]
+                    )
+                    source_list += [
+                        getattr(pxy, attr) for attr in fml_arg.attr_names
+                    ]
+                else:
+                    # assumed primitive type
+                    source_list.append(actual_args[argidx])
+
+            if len(source_list) != vector_type.num_elements:
+                raise CudaLoweringError(
+                    f"Unmatched number of source elements ({len(source_list)}) "
+                    "and target elements ({vector_type.num_elements})."
+                )
+
+            out = cgutils.create_struct_proxy(vector_type)(context, builder)
+
+            for attr_name, source in zip(vector_type.attr_names, source_list):
+                setattr(out, attr_name, source)
+            return out._getvalue()
+
+        return lowering
+
+    for arglist in overloads:
+        lowering = make_lowering(arglist)
+        lower(ctor, *arglist)(lowering)
+
+
+vector_types: Dict[str, VectorType] = {}
+
+
+def build_constructor_overloads(base_type, vty_name, num_elements, arglists, l):
+    """
+    For a given vector type, build a list of overloads for its constructor.
+    """
+
+    # TODO: speed up with memoization
+    if num_elements == 0:
+        arglists.append(l[:])
+
+    for i in range(1, num_elements + 1):
+        if i == 1:
+            # For 1-element component, it can construct with either a
+            # primitive type or other 1-element component.
+            l.append(base_type)
+            build_constructor_overloads(
+                base_type, vty_name, num_elements - i, arglists, l
+            )
+            l.pop(-1)
+
+            l.append(vector_types[f"{vty_name[:-1]}1"])
+            build_constructor_overloads(
+                base_type, vty_name, num_elements - i, arglists, l
+            )
+            l.pop(-1)
+        else:
+            l.append(vector_types[f"{vty_name[:-1]}{i}"])
+            build_constructor_overloads(
+                base_type, vty_name, num_elements - i, arglists, l
+            )
+            l.pop(-1)
+
+
+def _initialize():
+    """
+    Construct the vector types, populate `vector_types` dictionary, and
+    enable the constructors.
+    """
+    vector_type_attribute_names = ("x", "y", "z", "w")
+    for stub in stubs._vector_type_stubs:
+        type_name = stub.__name__
+        base_type = getattr(types, type_name[:-2])
+        num_elements = int(type_name[-1])
+        attributes = vector_type_attribute_names[:num_elements]
+        vector_type = make_vector_type(type_name, base_type, attributes, stub)
+        vector_types[type_name] = vector_type
+
+    for vty in vector_types.values():
+        arglists, l = [], []
+        build_constructor_overloads(
+            vty.base_type, vty.name, vty.num_elements, arglists, l
+        )
+        enable_vector_type_ctor(vty, arglists)
+
+
+_initialize()

numba_cuda/numba/cuda/vectorizers.py

@@ -0,0 +1,260 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+from numba import cuda
+from numpy import array as np_array
+from numba.cuda import deviceufunc
+from numba.cuda.deviceufunc import (
+    UFuncMechanism,
+    GeneralizedUFunc,
+    GUFuncCallSteps,
+)
+
+
+class CUDAUFuncDispatcher(object):
+    """
+    Invoke the CUDA ufunc specialization for the given inputs.
+    """
+
+    def __init__(self, types_to_retty_kernels, pyfunc):
+        self.functions = types_to_retty_kernels
+        self.__name__ = pyfunc.__name__
+
+    def __call__(self, *args, **kws):
+        """
+        *args: numpy arrays or DeviceArrayBase (created by cuda.to_device).
+               Cannot mix the two types in one call.
+
+        **kws:
+            stream -- cuda stream; when defined, asynchronous mode is used.
+            out    -- output array. Can be a numpy array or DeviceArrayBase
+                      depending on the input arguments.  Type must match
+                      the input arguments.
+        """
+        return CUDAUFuncMechanism.call(self.functions, args, kws)
+
+    def reduce(self, arg, stream=0):
+        assert len(list(self.functions.keys())[0]) == 2, (
+            "must be a binary ufunc"
+        )
+        assert arg.ndim == 1, "must use 1d array"
+
+        n = arg.shape[0]
+        gpu_mems = []
+
+        if n == 0:
+            raise TypeError("Reduction on an empty array.")
+        elif n == 1:  # nothing to do
+            return arg[0]
+
+        # always use a stream
+        stream = stream or cuda.stream()
+        with stream.auto_synchronize():
+            # transfer memory to device if necessary
+            if cuda.cudadrv.devicearray.is_cuda_ndarray(arg):
+                mem = arg
+            else:
+                mem = cuda.to_device(arg, stream)
+                # do reduction
+            out = self.__reduce(mem, gpu_mems, stream)
+            # use a small buffer to store the result element
+            buf = np_array((1,), dtype=arg.dtype)
+            out.copy_to_host(buf, stream=stream)
+
+        return buf[0]
+
+    def __reduce(self, mem, gpu_mems, stream):
+        n = mem.shape[0]
+        if n % 2 != 0:  # odd?
+            fatcut, thincut = mem.split(n - 1)
+            # prevent freeing during async mode
+            gpu_mems.append(fatcut)
+            gpu_mems.append(thincut)
+            # execute the kernel
+            out = self.__reduce(fatcut, gpu_mems, stream)
+            gpu_mems.append(out)
+            return self(out, thincut, out=out, stream=stream)
+        else:  # even?
+            left, right = mem.split(n // 2)
+            # prevent freeing during async mode
+            gpu_mems.append(left)
+            gpu_mems.append(right)
+            # execute the kernel
+            self(left, right, out=left, stream=stream)
+            if n // 2 > 1:
+                return self.__reduce(left, gpu_mems, stream)
+            else:
+                return left
+
+
+class _CUDAGUFuncCallSteps(GUFuncCallSteps):
+    __slots__ = [
+        "_stream",
+    ]
+
+    def __init__(self, nin, nout, args, kwargs):
+        super().__init__(nin, nout, args, kwargs)
+        self._stream = kwargs.get("stream", 0)
+
+    def is_device_array(self, obj):
+        return cuda.is_cuda_array(obj)
+
+    def as_device_array(self, obj):
+        # We don't want to call as_cuda_array on objects that are already Numba
+        # device arrays, because this results in exporting the array as a
+        # Producer then importing it as a Consumer, which causes a
+        # synchronization on the array's stream (if it has one) by default.
+        # When we have a Numba device array, we can simply return it.
+        if cuda.cudadrv.devicearray.is_cuda_ndarray(obj):
+            return obj
+        return cuda.as_cuda_array(obj)
+
+    def to_device(self, hostary):
+        return cuda.to_device(hostary, stream=self._stream)
+
+    def to_host(self, devary, hostary):
+        out = devary.copy_to_host(hostary, stream=self._stream)
+        return out
+
+    def allocate_device_array(self, shape, dtype):
+        return cuda.device_array(shape=shape, dtype=dtype, stream=self._stream)
+
+    def launch_kernel(self, kernel, nelem, args):
+        kernel.forall(nelem, stream=self._stream)(*args)
+
+
+class CUDAGeneralizedUFunc(GeneralizedUFunc):
+    def __init__(self, kernelmap, engine, pyfunc):
+        self.__name__ = pyfunc.__name__
+        super().__init__(kernelmap, engine)
+
+    @property
+    def _call_steps(self):
+        return _CUDAGUFuncCallSteps
+
+    def _broadcast_scalar_input(self, ary, shape):
+        return cuda.cudadrv.devicearray.DeviceNDArray(
+            shape=shape, strides=(0,), dtype=ary.dtype, gpu_data=ary.gpu_data
+        )
+
+    def _broadcast_add_axis(self, ary, newshape):
+        newax = len(newshape) - len(ary.shape)
+        # Add 0 strides for missing dimension
+        newstrides = (0,) * newax + ary.strides
+        return cuda.cudadrv.devicearray.DeviceNDArray(
+            shape=newshape,
+            strides=newstrides,
+            dtype=ary.dtype,
+            gpu_data=ary.gpu_data,
+        )
+
+
+class CUDAUFuncMechanism(UFuncMechanism):
+    """
+    Provide CUDA specialization
+    """
+
+    DEFAULT_STREAM = 0
+
+    def launch(self, func, count, stream, args):
+        func.forall(count, stream=stream)(*args)
+
+    def is_device_array(self, obj):
+        return cuda.is_cuda_array(obj)
+
+    def as_device_array(self, obj):
+        # We don't want to call as_cuda_array on objects that are already Numba
+        # device arrays, because this results in exporting the array as a
+        # Producer then importing it as a Consumer, which causes a
+        # synchronization on the array's stream (if it has one) by default.
+        # When we have a Numba device array, we can simply return it.
+        if cuda.cudadrv.devicearray.is_cuda_ndarray(obj):
+            return obj
+        return cuda.as_cuda_array(obj)
+
+    def to_device(self, hostary, stream):
+        return cuda.to_device(hostary, stream=stream)
+
+    def to_host(self, devary, stream):
+        return devary.copy_to_host(stream=stream)
+
+    def allocate_device_array(self, shape, dtype, stream):
+        return cuda.device_array(shape=shape, dtype=dtype, stream=stream)
+
+    def broadcast_device(self, ary, shape):
+        ax_differs = [
+            ax
+            for ax in range(len(shape))
+            if ax >= ary.ndim or ary.shape[ax] != shape[ax]
+        ]
+
+        missingdim = len(shape) - len(ary.shape)
+        strides = [0] * missingdim + list(ary.strides)
+
+        for ax in ax_differs:
+            strides[ax] = 0
+
+        return cuda.cudadrv.devicearray.DeviceNDArray(
+            shape=shape, strides=strides, dtype=ary.dtype, gpu_data=ary.gpu_data
+        )
+
+
+vectorizer_stager_source = """
+def __vectorized_{name}({args}, __out__):
+    __tid__ = __cuda__.grid(1)
+    if __tid__ < __out__.shape[0]:
+        __out__[__tid__] = __core__({argitems})
+"""
+
+
+class CUDAVectorize(deviceufunc.DeviceVectorize):
+    def _compile_core(self, sig):
+        cudevfn = cuda.jit(sig, device=True, inline="always")(self.pyfunc)
+        return cudevfn, cudevfn.overloads[sig.args].signature.return_type
+
+    def _get_globals(self, corefn):
+        glbl = self.pyfunc.__globals__.copy()
+        glbl.update({"__cuda__": cuda, "__core__": corefn})
+        return glbl
+
+    def _compile_kernel(self, fnobj, sig):
+        return cuda.jit(fnobj)
+
+    def build_ufunc(self):
+        return CUDAUFuncDispatcher(self.kernelmap, self.pyfunc)
+
+    @property
+    def _kernel_template(self):
+        return vectorizer_stager_source
+
+
+# ------------------------------------------------------------------------------
+# Generalized CUDA ufuncs
+
+_gufunc_stager_source = """
+def __gufunc_{name}({args}):
+    __tid__ = __cuda__.grid(1)
+    if __tid__ < {checkedarg}:
+        __core__({argitems})
+"""
+
+
+class CUDAGUFuncVectorize(deviceufunc.DeviceGUFuncVectorize):
+    def build_ufunc(self):
+        engine = deviceufunc.GUFuncEngine(self.inputsig, self.outputsig)
+        return CUDAGeneralizedUFunc(
+            kernelmap=self.kernelmap, engine=engine, pyfunc=self.pyfunc
+        )
+
+    def _compile_kernel(self, fnobj, sig):
+        return cuda.jit(sig)(fnobj)
+
+    @property
+    def _kernel_template(self):
+        return _gufunc_stager_source
+
+    def _get_globals(self, sig):
+        corefn = cuda.jit(sig, device=True)(self.pyfunc)
+        glbls = self.py_func.__globals__.copy()
+        glbls.update({"__cuda__": cuda, "__core__": corefn})
+        return glbls

numba_cuda-0.21.1.dist-info/METADATA

@@ -0,0 +1,109 @@
+Metadata-Version: 2.4
+Name: numba-cuda
+Version: 0.21.1
+Summary: CUDA target for Numba
+Author: Anaconda Inc., NVIDIA Corporation
+License-Expression: BSD-2-Clause
+Project-URL: Homepage, https://nvidia.github.io/numba-cuda/
+Project-URL: Documentation, https://nvidia.github.io/numba-cuda/
+Project-URL: Repository, https://github.com/NVIDIA/numba-cuda
+Project-URL: License, https://github.com/NVIDIA/numba-cuda/blob/main/LICENSE
+Project-URL: Issues, https://github.com/NVIDIA/numba-cuda/issues
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: LICENSE.numba
+Requires-Dist: numba>=0.60.0
+Requires-Dist: cuda-bindings<14.0.0,>=12.9.1
+Requires-Dist: cuda-core<1.0.0,>=0.3.2
+Provides-Extra: cu12
+Requires-Dist: cuda-bindings<13.0.0,>=12.9.1; extra == "cu12"
+Requires-Dist: cuda-core<1.0.0,>=0.3.0; extra == "cu12"
+Requires-Dist: cuda-python==12.9.*; extra == "cu12"
+Requires-Dist: nvidia-cuda-nvcc-cu12; extra == "cu12"
+Requires-Dist: nvidia-cuda-runtime-cu12; extra == "cu12"
+Requires-Dist: nvidia-cuda-nvrtc-cu12; extra == "cu12"
+Requires-Dist: nvidia-nvjitlink-cu12; extra == "cu12"
+Requires-Dist: nvidia-cuda-cccl-cu12; extra == "cu12"
+Provides-Extra: cu13
+Requires-Dist: cuda-bindings==13.*; extra == "cu13"
+Requires-Dist: cuda-core<1.0.0,>=0.3.2; extra == "cu13"
+Requires-Dist: cuda-python==13.*; extra == "cu13"
+Requires-Dist: nvidia-nvvm==13.*; extra == "cu13"
+Requires-Dist: nvidia-cuda-runtime==13.*; extra == "cu13"
+Requires-Dist: nvidia-cuda-nvrtc==13.*; extra == "cu13"
+Requires-Dist: nvidia-nvjitlink==13.*; extra == "cu13"
+Requires-Dist: nvidia-cuda-cccl==13.*; extra == "cu13"
+Dynamic: license-file
+
+<div align="center"><img src="docs/source/_static/numba-green-icon-rgb.svg" width="200"/></div>
+
+# Numba CUDA Target
+
+The CUDA target for Numba. Please visit the [official
+documentation](https://nvidia.github.io/numba-cuda) to get started!
+
+
+To report issues or file feature requests, please use the [issue
+tracker](https://github.com/NVIDIA/numba-cuda/issues).
+
+To raise questions or initiate discussions, please use the [Numba Discourse
+forum](https://numba.discourse.group).
+
+## Installation with pip or conda
+
+Please refer to the [Installation documentation](https://nvidia.github.io/numba-cuda/user/installation.html#installation-with-a-python-package-manager).
+
+
+## Installation from source
+
+Install as an editable install:
+
+```
+pip install -e .
+```
+
+If you want to manage all run-time dependencies yourself, also pass the `--no-deps` flag.
+
+## Running tests
+
+Tests must be run from the `testing` folder, which contains the pytest
+configuration and code to generate binaries used during the tests. The test
+binaries need to be built on the system on which the tests are run, so that
+they are compiled for the appropriate compute capability.
+
+```
+cd testing
+# Optionally, build test binaries and point to their location for the test suite
+make -j $(nproc)
+export NUMBA_CUDA_TEST_BIN_DIR=`pwd`
+# Execute tests
+pytest -n auto -v --dist loadscope
+```
+
+Alternatively, you can use [pixi](https://pixi.sh/latest/installation/) to wrap all of that up for you:
+
+```
+# run tests against CUDA 13
+pixi run -e cu13 test -n auto -v --dist loadscope
+```
+
+
+Testing should discover the `numba.cuda` module from the `numba_cuda` package. You
+can check where `numba.cuda` files are being located by running
+
+```
+python -c "from numba import cuda; print(cuda.__file__)"
+```
+
+which will show a path like:
+
+```
+<path to numba-cuda repo>/numba_cuda/numba/cuda/__init__.py
+```
+
+## Contributing Guide
+
+Review the
+[CONTRIBUTING.md](https://github.com/NVIDIA/numba-cuda/blob/main/CONTRIBUTING.md)
+file for information on how to contribute code and issues to the project.