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

numba_cuda/numba/cuda/vector_types.py

@@ -0,0 +1,209 @@
+# 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 import types
+from numba.core import cgutils
+from numba.core.extending import make_attribute_wrapper, models, register_model
+from numba.core.imputils import Registry as ImplRegistry
+from numba.core.typing.templates import ConcreteTemplate
+from numba.core.typing.templates import Registry as TypingRegistry
+from numba.core.typing.templates import signature
+from numba.cuda import stubs
+from numba.cuda.errors import CudaLoweringError
+
+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.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,252 @@
+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=True)(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.0.12.dist-info/LICENSE

@@ -0,0 +1,25 @@
+Copyright (c) 2012, Anaconda, Inc.
+Copyright (c) 2024, NVIDIA CORPORATION.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are
+met:
+
+Redistributions of source code must retain the above copyright notice,
+this list of conditions and the following disclaimer.
+
+Redistributions in binary form must reproduce the above copyright
+notice, this list of conditions and the following disclaimer in the
+documentation and/or other materials provided with the distribution.
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

numba_cuda-0.0.12.dist-info/METADATA

@@ -0,0 +1,68 @@
+Metadata-Version: 2.1
+Name: numba-cuda
+Version: 0.0.12
+Summary: CUDA target for Numba
+Author: Anaconda Inc., NVIDIA Corporation
+License: BSD 2-clause
+Project-URL: Homepage, https://github.com/rapidsai/numba-cuda
+Project-URL: Documentation, https://github.com/rapidsai/numba-cuda/blob/main/README.md
+Project-URL: Repository, https://github.com/rapidsai/numba-cuda
+Project-URL: License, https://github.com/rapidsai/numba-cuda/blob/main/LICENSE
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+
+# Numba CUDA Target
+
+An out-of-tree CUDA target for Numba.
+
+This contains an entire copy of Numba's CUDA target (the `numba.cuda` module),
+and a mechanism to ensure the code from this module (`numba_cuda.numba.cuda`) is
+used as the `numba.cuda` module instead of the code from the `numba` package.
+
+This is presently in an early state and is published for testing and feedback.
+
+## Building / testing
+
+Install as an editable install:
+
+```
+pip install -e .
+```
+
+Running tests:
+
+```
+python -m numba.runtests numba.cuda.tests
+```
+
+This 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
+```
+
+## Branching strategy
+
+Presently the `main` branch is being used to target the exact behavior of the
+built-in CUDA target. New feature development and bug fixes should be applied to
+`develop`. Once the `main` branch is widely tested and confirmed to work well as
+a drop-in replacement for the built-in `numba.cuda`, the `develop` branch will
+be merged in and new feature development will proceed on `main`.
+
+### Current PR targets
+
+- PRs related to replacing the built-in CUDA target's features should target
+  `main`.
+- PRs adding new features and bug fixes should target `develop`.
+
+### Future PR targets
+
+- In future, all PRs should target the `main` branch.