numba-cuda 0.17.0__py3-none-any.whl → 0.18.1__py3-none-any.whl

numba_cuda/numba/cuda/core/callconv.py

@@ -0,0 +1,338 @@
+from numba.core import types
+from numba.cuda import cgutils
+from collections import namedtuple
+
+from llvmlite import ir
+
+int32_t = ir.IntType(32)
+int64_t = ir.IntType(64)
+errcode_t = int32_t
+
+
+Status = namedtuple(
+    "Status",
+    (
+        "code",
+        # If the function returned ok (a value or None)
+        "is_ok",
+        # If the function returned None
+        "is_none",
+        # If the function errored out (== not is_ok)
+        "is_error",
+        # If the generator exited with StopIteration
+        "is_stop_iteration",
+        # If the function errored with an already set exception
+        "is_python_exc",
+        # If the function errored with a user exception
+        "is_user_exc",
+        # The pointer to the exception info structure (for user
+        # exceptions)
+        "excinfoptr",
+    ),
+)
+
+
+def _const_int(code):
+    return ir.Constant(errcode_t, code)
+
+
+RETCODE_OK = _const_int(0)
+RETCODE_EXC = _const_int(-1)
+RETCODE_NONE = _const_int(-2)
+# StopIteration
+RETCODE_STOPIT = _const_int(-3)
+
+FIRST_USEREXC = 1
+
+RETCODE_USEREXC = _const_int(FIRST_USEREXC)
+
+
+class BaseCallConv(object):
+    def __init__(self, context):
+        self.context = context
+
+    def return_optional_value(self, builder, retty, valty, value):
+        if valty == types.none:
+            # Value is none
+            self.return_native_none(builder)
+
+        elif retty == valty:
+            # Value is an optional, need a runtime switch
+            optval = self.context.make_helper(builder, retty, value=value)
+
+            validbit = cgutils.as_bool_bit(builder, optval.valid)
+            with builder.if_then(validbit):
+                retval = self.context.get_return_value(
+                    builder, retty.type, optval.data
+                )
+                self.return_value(builder, retval)
+
+            self.return_native_none(builder)
+
+        elif not isinstance(valty, types.Optional):
+            # Value is not an optional, need a cast
+            if valty != retty.type:
+                value = self.context.cast(
+                    builder, value, fromty=valty, toty=retty.type
+                )
+            retval = self.context.get_return_value(builder, retty.type, value)
+            self.return_value(builder, retval)
+
+        else:
+            raise NotImplementedError(
+                "returning {0} for {1}".format(valty, retty)
+            )
+
+    def return_native_none(self, builder):
+        self._return_errcode_raw(builder, RETCODE_NONE)
+
+    def return_exc(self, builder):
+        self._return_errcode_raw(builder, RETCODE_EXC)
+
+    def return_stop_iteration(self, builder):
+        self._return_errcode_raw(builder, RETCODE_STOPIT)
+
+    def get_return_type(self, ty):
+        """
+        Get the actual type of the return argument for Numba type *ty*.
+        """
+        restype = self.context.data_model_manager[ty].get_return_type()
+        return restype.as_pointer()
+
+    def init_call_helper(self, builder):
+        """
+        Initialize and return a call helper object for the given builder.
+        """
+        ch = self._make_call_helper(builder)
+        builder.__call_helper = ch
+        return ch
+
+    def _get_call_helper(self, builder):
+        return builder.__call_helper
+
+    def unpack_exception(self, builder, pyapi, status):
+        return pyapi.unserialize(status.excinfoptr)
+
+    def raise_error(self, builder, pyapi, status):
+        """
+        Given a non-ok *status*, raise the corresponding Python exception.
+        """
+        bbend = builder.function.append_basic_block()
+
+        with builder.if_then(status.is_user_exc):
+            # Unserialize user exception.
+            # Make sure another error may not interfere.
+            pyapi.err_clear()
+            exc = self.unpack_exception(builder, pyapi, status)
+            with cgutils.if_likely(builder, cgutils.is_not_null(builder, exc)):
+                pyapi.raise_object(exc)  # steals ref
+            builder.branch(bbend)
+
+        with builder.if_then(status.is_stop_iteration):
+            pyapi.err_set_none("PyExc_StopIteration")
+            builder.branch(bbend)
+
+        with builder.if_then(status.is_python_exc):
+            # Error already raised => nothing to do
+            builder.branch(bbend)
+
+        pyapi.err_set_string(
+            "PyExc_SystemError", "unknown error when calling native function"
+        )
+        builder.branch(bbend)
+
+        builder.position_at_end(bbend)
+
+    def decode_arguments(self, builder, argtypes, func):
+        """
+        Get the decoded (unpacked) Python arguments with *argtypes*
+        from LLVM function *func*.  A tuple of LLVM values is returned.
+        """
+        raw_args = self.get_arguments(func)
+        arginfo = self._get_arg_packer(argtypes)
+        return arginfo.from_arguments(builder, raw_args)
+
+    def _get_arg_packer(self, argtypes):
+        """
+        Get an argument packer for the given argument types.
+        """
+        return self.context.get_arg_packer(argtypes)
+
+
+class MinimalCallConv(BaseCallConv):
+    """
+    A minimal calling convention, suitable for e.g. GPU targets.
+    The implemented function signature is:
+
+        retcode_t (<Python return type>*, ... <Python arguments>)
+
+    The return code will be one of the RETCODE_* constants or a
+    function-specific user exception id (>= RETCODE_USEREXC).
+
+    Caller is responsible for allocating a slot for the return value
+    (passed as a pointer in the first argument).
+    """
+
+    def _make_call_helper(self, builder):
+        return _MinimalCallHelper()
+
+    def return_value(self, builder, retval):
+        retptr = builder.function.args[0]
+        assert retval.type == retptr.type.pointee, (
+            str(retval.type),
+            str(retptr.type.pointee),
+        )
+        builder.store(retval, retptr)
+        self._return_errcode_raw(builder, RETCODE_OK)
+
+    def return_user_exc(
+        self, builder, exc, exc_args=None, loc=None, func_name=None
+    ):
+        if exc is not None and not issubclass(exc, BaseException):
+            raise TypeError(
+                "exc should be None or exception class, got %r" % (exc,)
+            )
+        if exc_args is not None and not isinstance(exc_args, tuple):
+            raise TypeError(
+                "exc_args should be None or tuple, got %r" % (exc_args,)
+            )
+
+        # Build excinfo struct
+        if loc is not None:
+            fname = loc._raw_function_name()
+            if fname is None:
+                # could be exec(<string>) or REPL, try func_name
+                fname = func_name
+
+            locinfo = (fname, loc.filename, loc.line)
+            if None in locinfo:
+                locinfo = None
+        else:
+            locinfo = None
+
+        call_helper = self._get_call_helper(builder)
+        exc_id = call_helper._add_exception(exc, exc_args, locinfo)
+        self._return_errcode_raw(builder, _const_int(exc_id))
+
+    def return_status_propagate(self, builder, status):
+        self._return_errcode_raw(builder, status.code)
+
+    def _return_errcode_raw(self, builder, code):
+        if isinstance(code, int):
+            code = _const_int(code)
+        builder.ret(code)
+
+    def _get_return_status(self, builder, code):
+        """
+        Given a return *code*, get a Status instance.
+        """
+        norm = builder.icmp_signed("==", code, RETCODE_OK)
+        none = builder.icmp_signed("==", code, RETCODE_NONE)
+        ok = builder.or_(norm, none)
+        err = builder.not_(ok)
+        exc = builder.icmp_signed("==", code, RETCODE_EXC)
+        is_stop_iteration = builder.icmp_signed("==", code, RETCODE_STOPIT)
+        is_user_exc = builder.icmp_signed(">=", code, RETCODE_USEREXC)
+
+        status = Status(
+            code=code,
+            is_ok=ok,
+            is_error=err,
+            is_python_exc=exc,
+            is_none=none,
+            is_user_exc=is_user_exc,
+            is_stop_iteration=is_stop_iteration,
+            excinfoptr=None,
+        )
+        return status
+
+    def get_function_type(self, restype, argtypes):
+        """
+        Get the implemented Function type for *restype* and *argtypes*.
+        """
+        arginfo = self._get_arg_packer(argtypes)
+        argtypes = list(arginfo.argument_types)
+        resptr = self.get_return_type(restype)
+        fnty = ir.FunctionType(errcode_t, [resptr] + argtypes)
+        return fnty
+
+    def decorate_function(self, fn, args, fe_argtypes, noalias=False):
+        """
+        Set names and attributes of function arguments.
+        """
+        assert not noalias
+        arginfo = self._get_arg_packer(fe_argtypes)
+        arginfo.assign_names(self.get_arguments(fn), ["arg." + a for a in args])
+        fn.args[0].name = ".ret"
+
+    def get_arguments(self, func):
+        """
+        Get the Python-level arguments of LLVM *func*.
+        """
+        return func.args[1:]
+
+    def call_function(self, builder, callee, resty, argtys, args):
+        """
+        Call the Numba-compiled *callee*.
+        """
+        retty = callee.args[0].type.pointee
+        retvaltmp = cgutils.alloca_once(builder, retty)
+        # initialize return value
+        builder.store(cgutils.get_null_value(retty), retvaltmp)
+
+        arginfo = self._get_arg_packer(argtys)
+        args = arginfo.as_arguments(builder, args)
+        realargs = [retvaltmp] + list(args)
+        code = builder.call(callee, realargs)
+        status = self._get_return_status(builder, code)
+        retval = builder.load(retvaltmp)
+        out = self.context.get_returned_value(builder, resty, retval)
+        return status, out
+
+
+class _MinimalCallHelper(object):
+    """
+    A call helper object for the "minimal" calling convention.
+    User exceptions are represented as integer codes and stored in
+    a mapping for retrieval from the caller.
+    """
+
+    def __init__(self):
+        self.exceptions = {}
+
+    def _add_exception(self, exc, exc_args, locinfo):
+        """
+        Add a new user exception to this helper. Returns an integer that can be
+        used to refer to the added exception in future.
+
+        Parameters
+        ----------
+        exc :
+            exception type
+        exc_args : None or tuple
+            exception args
+        locinfo : tuple
+            location information
+        """
+        exc_id = len(self.exceptions) + FIRST_USEREXC
+        self.exceptions[exc_id] = exc, exc_args, locinfo
+        return exc_id
+
+    def get_exception(self, exc_id):
+        """
+        Get information about a user exception. Returns a tuple of
+        (exception type, exception args, location information).
+
+        Parameters
+        ----------
+        id : integer
+            The ID of the exception to look up
+        """
+        try:
+            return self.exceptions[exc_id]
+        except KeyError:
+            msg = "unknown error %d in native function" % exc_id
+            exc = SystemError
+            exc_args = (msg,)
+            locinfo = None
+            return exc, exc_args, locinfo

numba_cuda/numba/cuda/core/codegen.py

@@ -0,0 +1,168 @@
+from abc import abstractmethod, ABCMeta
+from numba.misc.llvm_pass_timings import PassTimingsCollection
+
+
+class CodeLibrary(metaclass=ABCMeta):
+    """
+    An interface for bundling LLVM code together and compiling it.
+    It is tied to a *codegen* instance (e.g. JITCUDACodegen) that will
+    determine how the LLVM code is transformed and linked together.
+    """
+
+    _finalized = False
+    _object_caching_enabled = False
+    _disable_inspection = False
+
+    def __init__(self, codegen: "Codegen", name: str):
+        self._codegen = codegen
+        self._name = name
+        ptc_name = f"{self.__class__.__name__}({self._name!r})"
+        self._recorded_timings = PassTimingsCollection(ptc_name)
+        # Track names of the dynamic globals
+        self._dynamic_globals = []
+
+    @property
+    def has_dynamic_globals(self):
+        self._ensure_finalized()
+        return len(self._dynamic_globals) > 0
+
+    @property
+    def recorded_timings(self):
+        return self._recorded_timings
+
+    @property
+    def codegen(self):
+        """
+        The codegen object owning this library.
+        """
+        return self._codegen
+
+    @property
+    def name(self):
+        return self._name
+
+    def __repr__(self):
+        return "<Library %r at 0x%x>" % (self.name, id(self))
+
+    def _raise_if_finalized(self):
+        if self._finalized:
+            raise RuntimeError(
+                "operation impossible on finalized object %r" % (self,)
+            )
+
+    def _ensure_finalized(self):
+        if not self._finalized:
+            self.finalize()
+
+    def create_ir_module(self, name):
+        """
+        Create an LLVM IR module for use by this library.
+        """
+        self._raise_if_finalized()
+        ir_module = self._codegen._create_empty_module(name)
+        return ir_module
+
+    @abstractmethod
+    def add_linking_library(self, library):
+        """
+        Add a library for linking into this library, without losing
+        the original library.
+        """
+
+    @abstractmethod
+    def add_ir_module(self, ir_module):
+        """
+        Add an LLVM IR module's contents to this library.
+        """
+
+    @abstractmethod
+    def finalize(self):
+        """
+        Finalize the library.  After this call, nothing can be added anymore.
+        Finalization involves various stages of code optimization and
+        linking.
+        """
+
+    @abstractmethod
+    def get_function(self, name):
+        """
+        Return the function named ``name``.
+        """
+
+    @abstractmethod
+    def get_llvm_str(self):
+        """
+        Get the human-readable form of the LLVM module.
+        """
+
+    @abstractmethod
+    def get_asm_str(self):
+        """
+        Get the human-readable assembly.
+        """
+
+    #
+    # Object cache hooks and serialization
+    #
+
+    def enable_object_caching(self):
+        self._object_caching_enabled = True
+        self._compiled_object = None
+        self._compiled = False
+
+    def _get_compiled_object(self):
+        if not self._object_caching_enabled:
+            raise ValueError("object caching not enabled in %s" % (self,))
+        if self._compiled_object is None:
+            raise RuntimeError("no compiled object yet for %s" % (self,))
+        return self._compiled_object
+
+    def _set_compiled_object(self, value):
+        if not self._object_caching_enabled:
+            raise ValueError("object caching not enabled in %s" % (self,))
+        if self._compiled:
+            raise ValueError("library already compiled: %s" % (self,))
+        self._compiled_object = value
+        self._disable_inspection = True
+
+
+class Codegen(metaclass=ABCMeta):
+    """
+    Base Codegen class. It is expected that subclasses set the class attribute
+    ``_library_class``, indicating the CodeLibrary class for the target.
+
+    Subclasses should also initialize:
+
+    ``self._data_layout``: the data layout for the target.
+    ``self._target_data``: the binding layer ``TargetData`` for the target.
+    """
+
+    @abstractmethod
+    def _create_empty_module(self, name):
+        """
+        Create a new empty module suitable for the target.
+        """
+
+    @abstractmethod
+    def _add_module(self, module):
+        """
+        Add a module to the execution engine. Ownership of the module is
+        transferred to the engine.
+        """
+
+    @property
+    def target_data(self):
+        """
+        The LLVM "target data" object for this codegen instance.
+        """
+        return self._target_data
+
+    def create_library(self, name, **kwargs):
+        """
+        Create a :class:`CodeLibrary` object for use with this codegen
+        instance.
+        """
+        return self._library_class(self, name, **kwargs)
+
+    def unserialize_library(self, serialized):
+        return self._library_class._unserialize(self, serialized)

numba_cuda/numba/cuda/core/compiler.py

@@ -0,0 +1,205 @@
+from numba.core.tracing import event
+
+from numba.core import callconv, bytecode, config, errors
+from numba.core.errors import CompilerError
+from numba.parfors.parfor import ParforDiagnostics
+
+from numba.core.untyped_passes import ExtractByteCode, FixupArgs
+from numba.core.targetconfig import ConfigStack
+
+
+class _CompileStatus(object):
+    """
+    Describes the state of compilation. Used like a C record.
+    """
+
+    __slots__ = ["fail_reason", "can_fallback"]
+
+    def __init__(self, can_fallback):
+        self.fail_reason = None
+        self.can_fallback = can_fallback
+
+    def __repr__(self):
+        vals = []
+        for k in self.__slots__:
+            vals.append("{k}={v}".format(k=k, v=getattr(self, k)))
+        return ", ".join(vals)
+
+
+class StateDict(dict):
+    """
+    A dictionary that has an overloaded getattr and setattr to permit getting
+    and setting key/values through the use of attributes.
+    """
+
+    def __getattr__(self, attr):
+        try:
+            return self[attr]
+        except KeyError:
+            raise AttributeError(attr)
+
+    def __setattr__(self, attr, value):
+        self[attr] = value
+
+
+class _EarlyPipelineCompletion(Exception):
+    """
+    Raised to indicate that a pipeline has completed early
+    """
+
+    def __init__(self, result):
+        self.result = result
+
+
+def _make_subtarget(targetctx, flags):
+    """
+    Make a new target context from the given target context and flags.
+    """
+    subtargetoptions = {}
+    if flags.debuginfo:
+        subtargetoptions["enable_debuginfo"] = True
+    if flags.boundscheck:
+        subtargetoptions["enable_boundscheck"] = True
+    if flags.nrt:
+        subtargetoptions["enable_nrt"] = True
+    if flags.auto_parallel:
+        subtargetoptions["auto_parallel"] = flags.auto_parallel
+    if flags.fastmath:
+        subtargetoptions["fastmath"] = flags.fastmath
+    error_model = callconv.create_error_model(flags.error_model, targetctx)
+    subtargetoptions["error_model"] = error_model
+
+    return targetctx.subtarget(**subtargetoptions)
+
+
+class CompilerBase(object):
+    """
+    Stores and manages states for the compiler
+    """
+
+    def __init__(
+        self, typingctx, targetctx, library, args, return_type, flags, locals
+    ):
+        # Make sure the environment is reloaded
+        config.reload_config()
+        typingctx.refresh()
+        targetctx.refresh()
+
+        self.state = StateDict()
+
+        self.state.typingctx = typingctx
+        self.state.targetctx = _make_subtarget(targetctx, flags)
+        self.state.library = library
+        self.state.args = args
+        self.state.return_type = return_type
+        self.state.flags = flags
+        self.state.locals = locals
+
+        # Results of various steps of the compilation pipeline
+        self.state.bc = None
+        self.state.func_id = None
+        self.state.func_ir = None
+        self.state.lifted = None
+        self.state.lifted_from = None
+        self.state.typemap = None
+        self.state.calltypes = None
+        self.state.type_annotation = None
+        # holds arbitrary inter-pipeline stage meta data
+        self.state.metadata = {}
+        self.state.reload_init = []
+        # hold this for e.g. with_lifting, null out on exit
+        self.state.pipeline = self
+
+        # parfor diagnostics info, add to metadata
+        self.state.parfor_diagnostics = ParforDiagnostics()
+        self.state.metadata["parfor_diagnostics"] = (
+            self.state.parfor_diagnostics
+        )
+        self.state.metadata["parfors"] = {}
+
+        self.state.status = _CompileStatus(
+            can_fallback=self.state.flags.enable_pyobject
+        )
+
+    def compile_extra(self, func):
+        self.state.func_id = bytecode.FunctionIdentity.from_function(func)
+        ExtractByteCode().run_pass(self.state)
+
+        self.state.lifted = ()
+        self.state.lifted_from = None
+        return self._compile_bytecode()
+
+    def compile_ir(self, func_ir, lifted=(), lifted_from=None):
+        self.state.func_id = func_ir.func_id
+        self.state.lifted = lifted
+        self.state.lifted_from = lifted_from
+        self.state.func_ir = func_ir
+        self.state.nargs = self.state.func_ir.arg_count
+
+        FixupArgs().run_pass(self.state)
+        return self._compile_ir()
+
+    def define_pipelines(self):
+        """Child classes override this to customize the pipelines in use."""
+        raise NotImplementedError()
+
+    def _compile_core(self):
+        """
+        Populate and run compiler pipeline
+        """
+        with ConfigStack().enter(self.state.flags.copy()):
+            pms = self.define_pipelines()
+            for pm in pms:
+                pipeline_name = pm.pipeline_name
+                func_name = "%s.%s" % (
+                    self.state.func_id.modname,
+                    self.state.func_id.func_qualname,
+                )
+
+                event("Pipeline: %s for %s" % (pipeline_name, func_name))
+                self.state.metadata["pipeline_times"] = {
+                    pipeline_name: pm.exec_times
+                }
+                is_final_pipeline = pm == pms[-1]
+                res = None
+                try:
+                    pm.run(self.state)
+                    if self.state.cr is not None:
+                        break
+                except _EarlyPipelineCompletion as e:
+                    res = e.result
+                    break
+                except Exception as e:
+                    if not isinstance(e, errors.NumbaError):
+                        raise e
+                    self.state.status.fail_reason = e
+                    if is_final_pipeline:
+                        raise e
+            else:
+                raise CompilerError("All available pipelines exhausted")
+
+            # Pipeline is done, remove self reference to release refs to user
+            # code
+            self.state.pipeline = None
+
+            # organise a return
+            if res is not None:
+                # Early pipeline completion
+                return res
+            else:
+                assert self.state.cr is not None
+                return self.state.cr
+
+    def _compile_bytecode(self):
+        """
+        Populate and run pipeline for bytecode input
+        """
+        assert self.state.func_ir is None
+        return self._compile_core()
+
+    def _compile_ir(self):
+        """
+        Populate and run pipeline for IR input
+        """
+        assert self.state.func_ir is not None
+        return self._compile_core()