5681978645

# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
#
# SPDX-License-Identifier: Apache-2.0

import re
from functools import cached_property

import dpnp
from llvmlite import binding as ll
from llvmlite import ir as llvmir
from numba import typeof
from numba.core import cgutils, types, typing, utils
from numba.core.base import BaseContext
from numba.core.callconv import MinimalCallConv
from numba.core.registry import cpu_target
from numba.core.target_extension import GPU, target_registry
from numba.core.types import Array as NpArrayType

from numba_dpex.core.datamodel.models import _init_data_model_manager
from numba_dpex.core.exceptions import UnsupportedKernelArgumentError
from numba_dpex.core.typeconv import to_usm_ndarray
from numba_dpex.core.types import DpnpNdArray, USMNdArray
from numba_dpex.core.utils import get_info_from_suai
from numba_dpex.utils import address_space, calling_conv

from .. import codegen

CC_SPIR_KERNEL = "spir_kernel"
CC_SPIR_FUNC = "spir_func"
VALID_CHARS = re.compile(r"[^a-z0-9]", re.I)
LINK_ATOMIC = 111
LLVM_SPIRV_ARGS = 112


class DpexKernelTypingContext(typing.BaseContext):
    """Custom typing context to support kernel compilation.

    The customized typing context provides two features required to compile
    Python functions decorated by the kernel decorator: An overridden
    :func:`resolve_argument_type` that changes all ``npytypes.Array`` to
    :class:`numba_depx.core.types.Array`. An overridden
    :func:`load_additional_registries` that registers OpenCL math and other
    functions to the typing context.

    """

    def resolve_argument_type(self, val):
        """Return the Numba type of a Python value used as a function argument.

        Overrides the implementation of ``numba.core.typing.BaseContext`` to
        handle the special case of ``numba.core.types.npytypes.Array``. Whenever
        a NumPy ndarray argument is encountered as an argument to a ``kernel``
        function, it is converted to a ``numba_dpex.core.types.Array`` type.

        Args:
            val : A Python value that is passed as an argument to a ``kernel``
                  function.

        Returns: The Numba type corresponding to the Python value.

        Raises:
            ValueError: If the type of the Python value is not supported.

        """
        try:
            numba_type = typeof(val)

            if isinstance(numba_type, NpArrayType) and not isinstance(
                numba_type, USMNdArray
            ):
                raise UnsupportedKernelArgumentError(
                    type=str(type(val)), value=val
                )

            # A cast from DpnpNdArray type to USMNdArray is needed for all
            # arguments of DpnpNdArray type. Although, DpnpNdArray derives from
            # USMNdArray the two types use different data models. USMNdArray
            # uses the numba_dpex.core.datamodel.models.ArrayModel data model
            # that defines all CPointer type members in the GLOBAL address
            # space. The DpnpNdArray uses Numba's default ArrayModel that does
            # not define pointers in any specific address space. For OpenCL HD
            # Graphics devices, defining a kernel function (spir_kernel calling
            # convention) with pointer arguments that have no address space
            # qualifier causes a run time crash. By casting the argument type
            # for parfor arguments from DpnpNdArray type to the USMNdArray type
            # the generated kernel always has an address space qualifier,
            # avoiding the issue on OpenCL HD graphics devices.
            if isinstance(numba_type, DpnpNdArray):
                return USMNdArray(
                    ndim=numba_type.ndim,
                    layout=numba_type.layout,
                    dtype=numba_type.dtype,
                    usm_type=numba_type.usm_type,
                    queue=numba_type.queue,
                )

        except ValueError:
            # When an array-like kernel argument is not recognized by
            # numba-dpex, this additional check sees if the array-like object
            # implements the __sycl_usm_array_interface__ protocol. For such
            # cases, we treat the object as an UsmNdArray type.
            try:
                suai_attrs = get_info_from_suai(val)
                return to_usm_ndarray(suai_attrs)
            except Exception:
                raise UnsupportedKernelArgumentError(
                    type=str(type(val)), value=val
                )

        return super().resolve_argument_type(val)

    def load_additional_registries(self):
        """Register the OpenCL API and math and other functions."""
        from numba.core.typing import cmathdecl, npydecl

        from ...ocl import mathdecl, ocldecl

        self.install_registry(ocldecl.registry)
        self.install_registry(mathdecl.registry)
        self.install_registry(cmathdecl.registry)
        self.install_registry(npydecl.registry)


class SyclDevice(GPU):
    """Mark the hardware target as SYCL Device."""

    pass


DPEX_KERNEL_TARGET_NAME = "SyclDevice"

target_registry[DPEX_KERNEL_TARGET_NAME] = SyclDevice


class DpexKernelTargetContext(BaseContext):
    """A target context inheriting Numba's ``BaseContext`` that is customized
    for generating SYCL kernels.

    A customized target context for generating SPIR-V kernels. The class defines
    helper functions to generates SPIR-V kernels as LLVM IR using the required
    calling conventions and metadata. The class also registers OpenCL math and
    API functions, helper functions for inserting LLVM address
    space cast instructions, and other functionalities used by the compiler
    to generate SPIR-V kernels.

    """

    implement_powi_as_math_call = True

    def _gen_arg_addrspace_md(self, fn):
        """Generate kernel_arg_addr_space metadata."""
        mod = fn.module
        fnty = fn.type.pointee
        codes = []

        for a in fnty.args:
            if cgutils.is_pointer(a):
                codes.append(address_space.GLOBAL)
            else:
                codes.append(address_space.PRIVATE)

        consts = [llvmir.Constant(llvmir.IntType(32), x) for x in codes]
        name = llvmir.MetaDataString(mod, "kernel_arg_addr_space")
        return mod.add_metadata([name] + consts)

    def _gen_arg_type(self, fn):
        """Generate kernel_arg_type metadata."""
        mod = fn.module
        fnty = fn.type.pointee
        consts = [llvmir.MetaDataString(mod, str(a)) for a in fnty.args]
        name = llvmir.MetaDataString(mod, "kernel_arg_type")
        return mod.add_metadata([name] + consts)

    def _gen_arg_type_qual(self, fn):
        """Generate kernel_arg_type_qual metadata."""
        mod = fn.module
        fnty = fn.type.pointee
        consts = [llvmir.MetaDataString(mod, "") for _ in fnty.args]
        name = llvmir.MetaDataString(mod, "kernel_arg_type_qual")
        return mod.add_metadata([name] + consts)

    def _gen_arg_base_type(self, fn):
        """Generate kernel_arg_base_type metadata."""
        mod = fn.module
        fnty = fn.type.pointee
        consts = [llvmir.MetaDataString(mod, str(a)) for a in fnty.args]
        name = llvmir.MetaDataString(mod, "kernel_arg_base_type")
        return mod.add_metadata([name] + consts)

    def _finalize_wrapper_module(self, fn):
        """Add metadata and calling convention to the wrapper function.

        The helper function adds function metadata to the wrapper function and
        also module level metadata to the LLVM module containing the wrapper.
        We also make sure the wrapper function has ``spir_kernel`` calling
        convention, without which the function cannot be used as a kernel.

        Args:
            fn: LLVM function representing the "kernel" wrapper function.

        """
        mod = fn.module
        # Set norecurse
        fn.attributes.add("norecurse")
        # Set SPIR kernel calling convention
        fn.calling_convention = CC_SPIR_KERNEL

        # Mark kernels
        ocl_kernels = cgutils.get_or_insert_named_metadata(
            mod, "opencl.kernels"
        )
        ocl_kernels.add(
            mod.add_metadata(
                [
                    fn,
                    self._gen_arg_addrspace_md(fn),
                    self._gen_arg_type(fn),
                    self._gen_arg_type_qual(fn),
                    self._gen_arg_base_type(fn),
                ],
            )
        )

        # Other metadata
        others = [
            "opencl.used.extensions",
            "opencl.used.optional.core.features",
            "opencl.compiler.options",
        ]

        for name in others:
            nmd = cgutils.get_or_insert_named_metadata(mod, name)
            if not nmd.operands:
                mod.add_metadata([])

    def _generate_kernel_wrapper(self, func, argtypes):
        module = func.module
        arginfo = self.get_arg_packer(argtypes)
        wrapperfnty = llvmir.FunctionType(
            llvmir.VoidType(), arginfo.argument_types
        )
        wrapper_module = self.create_module("dpex.kernel.wrapper")
        wrappername = "dpexPy_{name}".format(name=func.name)
        argtys = list(arginfo.argument_types)
        fnty = llvmir.FunctionType(
            llvmir.IntType(32),
            [self.call_conv.get_return_type(types.pyobject)] + argtys,
        )
        func = llvmir.Function(wrapper_module, fnty, name=func.name)
        func.calling_convention = CC_SPIR_FUNC
        wrapper = llvmir.Function(wrapper_module, wrapperfnty, name=wrappername)
        builder = llvmir.IRBuilder(wrapper.append_basic_block(""))

        callargs = arginfo.from_arguments(builder, wrapper.args)

        # XXX handle error status
        status, _ = self.call_conv.call_function(
            builder, func, types.void, argtypes, callargs
        )
        builder.ret_void()

        self._finalize_wrapper_module(wrapper)

        # Link the spir_func module to the wrapper module
        module.link_in(ll.parse_assembly(str(wrapper_module)))
        # Make sure the spir_func has internal linkage to be inlinable.
        func.linkage = "internal"
        wrapper = module.get_function(wrapper.name)
        module.get_function(func.name).linkage = "internal"
        return wrapper

    def __init__(self, typingctx, target=DPEX_KERNEL_TARGET_NAME):
        super().__init__(typingctx, target)

    def init(self):
        self._internal_codegen = codegen.JITSPIRVCodegen("numba_dpex.jit")
        self._target_data = ll.create_target_data(
            codegen.SPIR_DATA_LAYOUT[utils.MACHINE_BITS]
        )
        # Override data model manager to SPIR model
        import numba.cpython.unicode

        self.data_model_manager = _init_data_model_manager()
        self.extra_compile_options = dict()

        import copy

        from numba.np.ufunc_db import _lazy_init_db

        _lazy_init_db()
        from numba.np.ufunc_db import _ufunc_db as ufunc_db

        self.ufunc_db = copy.deepcopy(ufunc_db)
        self.cpu_context = cpu_target.target_context

    # Overrides
    def create_module(self, name):
        return self._internal_codegen._create_empty_module(name)

    def replace_dpnp_ufunc_with_ocl_intrinsics(self):
        from numba_dpex.ocl.mathimpl import lower_ocl_impl, sig_mapper

        ufuncs = [
            ("fabs", dpnp.fabs),
            ("exp", dpnp.exp),
            ("log", dpnp.log),
            ("log10", dpnp.log10),
            ("expm1", dpnp.expm1),
            ("log1p", dpnp.log1p),
            ("sqrt", dpnp.sqrt),
            ("sin", dpnp.sin),
            ("cos", dpnp.cos),
            ("tan", dpnp.tan),
            ("asin", dpnp.arcsin),
            ("acos", dpnp.arccos),
            ("atan", dpnp.arctan),
            ("atan2", dpnp.arctan2),
            ("sinh", dpnp.sinh),
            ("cosh", dpnp.cosh),
            ("tanh", dpnp.tanh),
            ("asinh", dpnp.arcsinh),
            ("acosh", dpnp.arccosh),
            ("atanh", dpnp.arctanh),
            ("floor", dpnp.floor),
            ("ceil", dpnp.ceil),
            ("trunc", dpnp.trunc),
            ("hypot", dpnp.hypot),
            ("exp2", dpnp.exp2),
            ("log2", dpnp.log2),
        ]

        for name, ufunc in ufuncs:
            for sig in self.ufunc_db[ufunc].keys():
                if (
                    sig in sig_mapper
                    and (name, sig_mapper[sig]) in lower_ocl_impl
                ):
                    self.ufunc_db[ufunc][sig] = lower_ocl_impl[
                        (name, sig_mapper[sig])
                    ]

    def load_additional_registries(self):
        """Register OpenCL functions into numba_depx's target context.

        To make sure we are calling supported OpenCL math functions, we replace
        the dpnp functions that default to NUMBA's NumPy ufunc with OpenCL
        intrinsics that are equivalent to those functions. The replacement is
        done after the OpenCL functions have been registered into the
        target context.

        """
        from numba_dpex.dpnp_iface import dpnpimpl

        from ... import printimpl
        from ...ocl import mathimpl, oclimpl

        self.insert_func_defn(oclimpl.registry.functions)
        self.insert_func_defn(mathimpl.registry.functions)
        self.insert_func_defn(dpnpimpl.registry.functions)
        self.install_registry(printimpl.registry)
        # Replace dpnp math functions with their OpenCL versions.
        self.replace_dpnp_ufunc_with_ocl_intrinsics()

    @cached_property
    def call_conv(self):
        return DpexCallConv(self)

    def codegen(self):
        return self._internal_codegen

    @property
    def target_data(self):
        return self._target_data

    def mangler(self, name, argtypes, abi_tags=(), uid=None):
        def repl(m):
            ch = m.group(0)
            return "_%X_" % ord(ch)

        qualified = name + "." + ".".join(str(a) for a in argtypes)
        mangled = VALID_CHARS.sub(repl, qualified)
        return "dpex_py_devfn_" + mangled

    def prepare_ocl_kernel(self, func, argtypes):
        module = func.module
        func.linkage = "linkonce_odr"
        module.data_layout = codegen.SPIR_DATA_LAYOUT[self.address_size]
        wrapper = self._generate_kernel_wrapper(func, argtypes)
        return wrapper

    def mark_ocl_device(self, func):
        # Adapt to SPIR
        func.calling_convention = CC_SPIR_FUNC
        func.linkage = "linkonce_odr"
        return func

    def declare_function(self, module, fndesc):
        """Create the LLVM function from a ``numba_dpex.kernel`` decorated
        function.

        Args:
            module (llvmlite.ir.Module) : The LLVM module into which
                the kernel function will be inserted.
            fndesc (numba.core.funcdesc.PythonFunctionDescriptor) : The
                signature of the function.

        Returns:
            llvmlite.ir.values.Function: The reference to the LLVM Function
                that was inserted into the module.

        """
        fnty = self.call_conv.get_function_type(fndesc.restype, fndesc.argtypes)
        fn = cgutils.get_or_insert_function(
            module, fnty, name=fndesc.mangled_name
        )
        if not self.enable_debuginfo:
            fn.attributes.add("alwaysinline")
        ret = super(DpexKernelTargetContext, self).declare_function(
            module, fndesc
        )
        ret.calling_convention = calling_conv.CC_SPIR_FUNC
        return ret

    def insert_const_string(self, mod, string):
        """Create a global string from the passed in string argument and return
        a void* in the GENERIC address space pointing to that string.

        Args:
            mod: LLVM module where the global string value is to be inserted.
            string: A Python string that will be converted to a global constant
            string and inserted into the module.

        Returns: A LLVM Constant pointing to the global string value inserted
        into the module.

        """
        text = cgutils.make_bytearray(string.encode("utf-8") + b"\x00")

        name = "$".join(["__conststring__", self.mangler(string, ["str"])])

        # Try to reuse existing global
        gv = mod.globals.get(name)
        if gv is None:
            # Not defined yet
            gv = cgutils.add_global_variable(
                mod, text.type, name=name, addrspace=address_space.GENERIC
            )
            gv.linkage = "internal"
            gv.global_constant = True
            gv.initializer = text

        # Cast to a i8* pointer
        charty = gv.type.pointee.element
        return gv.bitcast(charty.as_pointer(address_space.GENERIC))

    def addrspacecast(self, builder, src, addrspace):
        """Insert an LLVM addressspace cast instruction into the module.

        FIXME: Move this function into utils.

        """
        ptras = llvmir.PointerType(src.type.pointee, addrspace=addrspace)
        return builder.addrspacecast(src, ptras)

    # Overrides
    def get_ufunc_info(self, ufunc_key):
        return self.ufunc_db[ufunc_key]


class DpexCallConv(MinimalCallConv):
    """Custom calling convention class used by numba-dpex.

    numba_dpex's calling convention derives from
    :class:`numba.core.callconv import MinimalCallConv`. The
    :class:`DpexCallConv` overriddes :func:`call_function`.

    """

    def call_function(self, builder, callee, resty, argtys, args, env=None):
        """Call the Numba-compiled *callee*."""
        assert env is None
        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.context.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

1	# SPDX-FileCopyrightText: 2020 - 2023 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import re	1✔
6	from functools import cached_property	1✔
7
8	import dpnp	1✔
9	from llvmlite import binding as ll	1✔
10	from llvmlite import ir as llvmir	1✔
11	from numba import typeof	1✔
12	from numba.core import cgutils, types, typing, utils	1✔
13	from numba.core.base import BaseContext	1✔
14	from numba.core.callconv import MinimalCallConv	1✔
15	from numba.core.registry import cpu_target	1✔
16	from numba.core.target_extension import GPU, target_registry	1✔
17	from numba.core.types import Array as NpArrayType	1✔
18
19	from numba_dpex.core.datamodel.models import _init_data_model_manager	1✔
20	from numba_dpex.core.exceptions import UnsupportedKernelArgumentError	1✔
21	from numba_dpex.core.typeconv import to_usm_ndarray	1✔
22	from numba_dpex.core.types import DpnpNdArray, USMNdArray	1✔
23	from numba_dpex.core.utils import get_info_from_suai	1✔
24	from numba_dpex.utils import address_space, calling_conv	1✔
25
26	from .. import codegen	1✔
27
28	CC_SPIR_KERNEL = "spir_kernel"	1✔
29	CC_SPIR_FUNC = "spir_func"	1✔
30	VALID_CHARS = re.compile(r"[^a-z0-9]", re.I)	1✔
31	LINK_ATOMIC = 111	1✔
32	LLVM_SPIRV_ARGS = 112	1✔
33
34
35	class DpexKernelTypingContext(typing.BaseContext):	1✔
36	"""Custom typing context to support kernel compilation.
37
38	The customized typing context provides two features required to compile
39	Python functions decorated by the kernel decorator: An overridden
40	:func:`resolve_argument_type` that changes all ``npytypes.Array`` to
41	:class:`numba_depx.core.types.Array`. An overridden
42	:func:`load_additional_registries` that registers OpenCL math and other
43	functions to the typing context.
44
45	"""
46
47	def resolve_argument_type(self, val):	1✔
48	"""Return the Numba type of a Python value used as a function argument.
49
50	Overrides the implementation of ``numba.core.typing.BaseContext`` to
51	handle the special case of ``numba.core.types.npytypes.Array``. Whenever
52	a NumPy ndarray argument is encountered as an argument to a ``kernel``
53	function, it is converted to a ``numba_dpex.core.types.Array`` type.
54
55	Args:
56	val : A Python value that is passed as an argument to a ``kernel``
57	function.
58
59	Returns: The Numba type corresponding to the Python value.
60
61	Raises:
62	ValueError: If the type of the Python value is not supported.
63
64	"""
65	try:	1✔
66	numba_type = typeof(val)	1✔
67
68	if isinstance(numba_type, NpArrayType) and not isinstance(	1✔
69	numba_type, USMNdArray
70	):
71	raise UnsupportedKernelArgumentError(	1✔
72	type=str(type(val)), value=val
73	)
74
75	# A cast from DpnpNdArray type to USMNdArray is needed for all
76	# arguments of DpnpNdArray type. Although, DpnpNdArray derives from
77	# USMNdArray the two types use different data models. USMNdArray
78	# uses the numba_dpex.core.datamodel.models.ArrayModel data model
79	# that defines all CPointer type members in the GLOBAL address
80	# space. The DpnpNdArray uses Numba's default ArrayModel that does
81	# not define pointers in any specific address space. For OpenCL HD
82	# Graphics devices, defining a kernel function (spir_kernel calling
83	# convention) with pointer arguments that have no address space
84	# qualifier causes a run time crash. By casting the argument type
85	# for parfor arguments from DpnpNdArray type to the USMNdArray type
86	# the generated kernel always has an address space qualifier,
87	# avoiding the issue on OpenCL HD graphics devices.
88	if isinstance(numba_type, DpnpNdArray):	1✔
89	return USMNdArray(	1✔
90	ndim=numba_type.ndim,
91	layout=numba_type.layout,
92	dtype=numba_type.dtype,
93	usm_type=numba_type.usm_type,
94	queue=numba_type.queue,
95	)
96
97	except ValueError:	1✔
98	# When an array-like kernel argument is not recognized by
99	# numba-dpex, this additional check sees if the array-like object
100	# implements the __sycl_usm_array_interface__ protocol. For such
101	# cases, we treat the object as an UsmNdArray type.
102	try:	1✔
103	suai_attrs = get_info_from_suai(val)	1✔
104	return to_usm_ndarray(suai_attrs)	1✔
105	except Exception:	1✔
106	raise UnsupportedKernelArgumentError(	1✔
107	type=str(type(val)), value=val
108	)
109
110	return super().resolve_argument_type(val)	1✔
111
112	def load_additional_registries(self):	1✔
113	"""Register the OpenCL API and math and other functions."""
114	from numba.core.typing import cmathdecl, npydecl	1✔
115
116	from ...ocl import mathdecl, ocldecl	1✔
117
118	self.install_registry(ocldecl.registry)	1✔
119	self.install_registry(mathdecl.registry)	1✔
120	self.install_registry(cmathdecl.registry)	1✔
121	self.install_registry(npydecl.registry)	1✔
122
123
124	class SyclDevice(GPU):	1✔
125	"""Mark the hardware target as SYCL Device."""
126
127	pass	1✔
128
129
130	DPEX_KERNEL_TARGET_NAME = "SyclDevice"	1✔
131
132	target_registry[DPEX_KERNEL_TARGET_NAME] = SyclDevice	1✔
133
134
135	class DpexKernelTargetContext(BaseContext):	1✔
136	"""A target context inheriting Numba's ``BaseContext`` that is customized
137	for generating SYCL kernels.
138
139	A customized target context for generating SPIR-V kernels. The class defines
140	helper functions to generates SPIR-V kernels as LLVM IR using the required
141	calling conventions and metadata. The class also registers OpenCL math and
142	API functions, helper functions for inserting LLVM address
143	space cast instructions, and other functionalities used by the compiler
144	to generate SPIR-V kernels.
145
146	"""
147
148	implement_powi_as_math_call = True	1✔
149
150	def _gen_arg_addrspace_md(self, fn):	1✔
151	"""Generate kernel_arg_addr_space metadata."""
152	mod = fn.module	1✔
153	fnty = fn.type.pointee	1✔
154	codes = []	1✔
155
156	for a in fnty.args:	1✔
157	if cgutils.is_pointer(a):	1✔
158	codes.append(address_space.GLOBAL)	1✔
159	else:
160	codes.append(address_space.PRIVATE)	1✔
161
162	consts = [llvmir.Constant(llvmir.IntType(32), x) for x in codes]	1✔
163	name = llvmir.MetaDataString(mod, "kernel_arg_addr_space")	1✔
164	return mod.add_metadata([name] + consts)	1✔
165
166	def _gen_arg_type(self, fn):	1✔
167	"""Generate kernel_arg_type metadata."""
168	mod = fn.module	1✔
169	fnty = fn.type.pointee	1✔
170	consts = [llvmir.MetaDataString(mod, str(a)) for a in fnty.args]	1✔
171	name = llvmir.MetaDataString(mod, "kernel_arg_type")	1✔
172	return mod.add_metadata([name] + consts)	1✔
173
174	def _gen_arg_type_qual(self, fn):	1✔
175	"""Generate kernel_arg_type_qual metadata."""
176	mod = fn.module	1✔
177	fnty = fn.type.pointee	1✔
178	consts = [llvmir.MetaDataString(mod, "") for _ in fnty.args]	1✔
179	name = llvmir.MetaDataString(mod, "kernel_arg_type_qual")	1✔
180	return mod.add_metadata([name] + consts)	1✔
181
182	def _gen_arg_base_type(self, fn):	1✔
183	"""Generate kernel_arg_base_type metadata."""
184	mod = fn.module	1✔
185	fnty = fn.type.pointee	1✔
186	consts = [llvmir.MetaDataString(mod, str(a)) for a in fnty.args]	1✔
187	name = llvmir.MetaDataString(mod, "kernel_arg_base_type")	1✔
188	return mod.add_metadata([name] + consts)	1✔
189
190	def _finalize_wrapper_module(self, fn):	1✔
191	"""Add metadata and calling convention to the wrapper function.
192
193	The helper function adds function metadata to the wrapper function and
194	also module level metadata to the LLVM module containing the wrapper.
195	We also make sure the wrapper function has ``spir_kernel`` calling
196	convention, without which the function cannot be used as a kernel.
197
198	Args:
199	fn: LLVM function representing the "kernel" wrapper function.
200
201	"""
202	mod = fn.module	1✔
203	# Set norecurse
204	fn.attributes.add("norecurse")	1✔
205	# Set SPIR kernel calling convention
206	fn.calling_convention = CC_SPIR_KERNEL	1✔
207
208	# Mark kernels
209	ocl_kernels = cgutils.get_or_insert_named_metadata(	1✔
210	mod, "opencl.kernels"
211	)
212	ocl_kernels.add(	1✔
213	mod.add_metadata(
214	[
215	fn,
216	self._gen_arg_addrspace_md(fn),
217	self._gen_arg_type(fn),
218	self._gen_arg_type_qual(fn),
219	self._gen_arg_base_type(fn),
220	],
221	)
222	)
223
224	# Other metadata
225	others = [	1✔
226	"opencl.used.extensions",
227	"opencl.used.optional.core.features",
228	"opencl.compiler.options",
229	]
230
231	for name in others:	1✔
232	nmd = cgutils.get_or_insert_named_metadata(mod, name)	1✔
233	if not nmd.operands:	1!
234	mod.add_metadata([])	1✔
235
236	def _generate_kernel_wrapper(self, func, argtypes):	1✔
237	module = func.module	1✔
238	arginfo = self.get_arg_packer(argtypes)	1✔
239	wrapperfnty = llvmir.FunctionType(	1✔
240	llvmir.VoidType(), arginfo.argument_types
241	)
242	wrapper_module = self.create_module("dpex.kernel.wrapper")	1✔
243	wrappername = "dpexPy_{name}".format(name=func.name)	1✔
244	argtys = list(arginfo.argument_types)	1✔
245	fnty = llvmir.FunctionType(	1✔
246	llvmir.IntType(32),
247	[self.call_conv.get_return_type(types.pyobject)] + argtys,
248	)
249	func = llvmir.Function(wrapper_module, fnty, name=func.name)	1✔
250	func.calling_convention = CC_SPIR_FUNC	1✔
251	wrapper = llvmir.Function(wrapper_module, wrapperfnty, name=wrappername)	1✔
252	builder = llvmir.IRBuilder(wrapper.append_basic_block(""))	1✔
253
254	callargs = arginfo.from_arguments(builder, wrapper.args)	1✔
255
256	# XXX handle error status
257	status, _ = self.call_conv.call_function(	1✔
258	builder, func, types.void, argtypes, callargs
259	)
260	builder.ret_void()	1✔
261
262	self._finalize_wrapper_module(wrapper)	1✔
263
264	# Link the spir_func module to the wrapper module
265	module.link_in(ll.parse_assembly(str(wrapper_module)))	1✔
266	# Make sure the spir_func has internal linkage to be inlinable.
267	func.linkage = "internal"	1✔
268	wrapper = module.get_function(wrapper.name)	1✔
269	module.get_function(func.name).linkage = "internal"	1✔
270	return wrapper	1✔
271
272	def __init__(self, typingctx, target=DPEX_KERNEL_TARGET_NAME):	1✔
273	super().__init__(typingctx, target)	1✔
274
275	def init(self):	1✔
276	self._internal_codegen = codegen.JITSPIRVCodegen("numba_dpex.jit")	1✔
277	self._target_data = ll.create_target_data(	1✔
278	codegen.SPIR_DATA_LAYOUT[utils.MACHINE_BITS]
279	)
280	# Override data model manager to SPIR model
281	import numba.cpython.unicode	1✔
282
283	self.data_model_manager = _init_data_model_manager()	1✔
284	self.extra_compile_options = dict()	1✔
285
286	import copy	1✔
287
288	from numba.np.ufunc_db import _lazy_init_db	1✔
289
290	_lazy_init_db()	1✔
291	from numba.np.ufunc_db import _ufunc_db as ufunc_db	1✔
292
293	self.ufunc_db = copy.deepcopy(ufunc_db)	1✔
294	self.cpu_context = cpu_target.target_context	1✔
295
296	# Overrides
297	def create_module(self, name):	1✔
298	return self._internal_codegen._create_empty_module(name)	1✔
299
300	def replace_dpnp_ufunc_with_ocl_intrinsics(self):	1✔
301	from numba_dpex.ocl.mathimpl import lower_ocl_impl, sig_mapper	1✔
302
303	ufuncs = [	1✔
304	("fabs", dpnp.fabs),
305	("exp", dpnp.exp),
306	("log", dpnp.log),
307	("log10", dpnp.log10),
308	("expm1", dpnp.expm1),
309	("log1p", dpnp.log1p),
310	("sqrt", dpnp.sqrt),
311	("sin", dpnp.sin),
312	("cos", dpnp.cos),
313	("tan", dpnp.tan),
314	("asin", dpnp.arcsin),
315	("acos", dpnp.arccos),
316	("atan", dpnp.arctan),
317	("atan2", dpnp.arctan2),
318	("sinh", dpnp.sinh),
319	("cosh", dpnp.cosh),
320	("tanh", dpnp.tanh),
321	("asinh", dpnp.arcsinh),
322	("acosh", dpnp.arccosh),
323	("atanh", dpnp.arctanh),
324	("floor", dpnp.floor),
325	("ceil", dpnp.ceil),
326	("trunc", dpnp.trunc),
327	("hypot", dpnp.hypot),
328	("exp2", dpnp.exp2),
329	("log2", dpnp.log2),
330	]
331
332	for name, ufunc in ufuncs:	1✔
333	for sig in self.ufunc_db[ufunc].keys():	1✔
334	if (	1!
335	sig in sig_mapper
336	and (name, sig_mapper[sig]) in lower_ocl_impl
337	):
338	self.ufunc_db[ufunc][sig] = lower_ocl_impl[	1✔
339	(name, sig_mapper[sig])
340	]
341
342	def load_additional_registries(self):	1✔
343	"""Register OpenCL functions into numba_depx's target context.
344
345	To make sure we are calling supported OpenCL math functions, we replace
346	the dpnp functions that default to NUMBA's NumPy ufunc with OpenCL
347	intrinsics that are equivalent to those functions. The replacement is
348	done after the OpenCL functions have been registered into the
349	target context.
350
351	"""
352	from numba_dpex.dpnp_iface import dpnpimpl	1✔
353
354	from ... import printimpl	1✔
355	from ...ocl import mathimpl, oclimpl	1✔
356
357	self.insert_func_defn(oclimpl.registry.functions)	1✔
358	self.insert_func_defn(mathimpl.registry.functions)	1✔
359	self.insert_func_defn(dpnpimpl.registry.functions)	1✔
360	self.install_registry(printimpl.registry)	1✔
361	# Replace dpnp math functions with their OpenCL versions.
362	self.replace_dpnp_ufunc_with_ocl_intrinsics()	1✔
363
364	@cached_property	1✔
365	def call_conv(self):	1✔
366	return DpexCallConv(self)	1✔
367
368	def codegen(self):	1✔
369	return self._internal_codegen	1✔
370
371	@property	1✔
372	def target_data(self):	1✔
373	return self._target_data	1✔
374
375	def mangler(self, name, argtypes, abi_tags=(), uid=None):	1✔
376	def repl(m):	1✔
377	ch = m.group(0)	1✔
378	return "_%X_" % ord(ch)	1✔
379
380	qualified = name + "." + ".".join(str(a) for a in argtypes)	1✔
381	mangled = VALID_CHARS.sub(repl, qualified)	1✔
382	return "dpex_py_devfn_" + mangled	1✔
383
384	def prepare_ocl_kernel(self, func, argtypes):	1✔
385	module = func.module	1✔
386	func.linkage = "linkonce_odr"	1✔
387	module.data_layout = codegen.SPIR_DATA_LAYOUT[self.address_size]	1✔
388	wrapper = self._generate_kernel_wrapper(func, argtypes)	1✔
389	return wrapper	1✔
390
391	def mark_ocl_device(self, func):	1✔
392	# Adapt to SPIR
393	func.calling_convention = CC_SPIR_FUNC	1✔
394	func.linkage = "linkonce_odr"	1✔
395	return func	1✔
396
397	def declare_function(self, module, fndesc):	1✔
398	"""Create the LLVM function from a ``numba_dpex.kernel`` decorated
399	function.
400
401	Args:
402	module (llvmlite.ir.Module) : The LLVM module into which
403	the kernel function will be inserted.
404	fndesc (numba.core.funcdesc.PythonFunctionDescriptor) : The
405	signature of the function.
406
407	Returns:
408	llvmlite.ir.values.Function: The reference to the LLVM Function
409	that was inserted into the module.
410
411	"""
412	fnty = self.call_conv.get_function_type(fndesc.restype, fndesc.argtypes)	1✔
413	fn = cgutils.get_or_insert_function(	1✔
414	module, fnty, name=fndesc.mangled_name
415	)
416	if not self.enable_debuginfo:	1✔
417	fn.attributes.add("alwaysinline")	1✔
418	ret = super(DpexKernelTargetContext, self).declare_function(	1✔
419	module, fndesc
420	)
421	ret.calling_convention = calling_conv.CC_SPIR_FUNC	1✔
422	return ret	1✔
423
424	def insert_const_string(self, mod, string):	1✔
425	"""Create a global string from the passed in string argument and return
426	a void* in the GENERIC address space pointing to that string.
427
428	Args:
429	mod: LLVM module where the global string value is to be inserted.
430	string: A Python string that will be converted to a global constant
431	string and inserted into the module.
432
433	Returns: A LLVM Constant pointing to the global string value inserted
434	into the module.
435
436	"""
437	text = cgutils.make_bytearray(string.encode("utf-8") + b"\x00")	1✔
438
439	name = "$".join(["__conststring__", self.mangler(string, ["str"])])	1✔
440
441	# Try to reuse existing global
442	gv = mod.globals.get(name)	1✔
443	if gv is None:	1!
444	# Not defined yet
445	gv = cgutils.add_global_variable(	1✔
446	mod, text.type, name=name, addrspace=address_space.GENERIC
447	)
448	gv.linkage = "internal"	1✔
449	gv.global_constant = True	1✔
450	gv.initializer = text	1✔
451
452	# Cast to a i8* pointer
453	charty = gv.type.pointee.element	1✔
454	return gv.bitcast(charty.as_pointer(address_space.GENERIC))	1✔
455
456	def addrspacecast(self, builder, src, addrspace):	1✔
457	"""Insert an LLVM addressspace cast instruction into the module.
458
459	FIXME: Move this function into utils.
460
461	"""
UNCOV 462	ptras = llvmir.PointerType(src.type.pointee, addrspace=addrspace)	×
UNCOV 463	return builder.addrspacecast(src, ptras)	×
464
465	# Overrides
466	def get_ufunc_info(self, ufunc_key):	1✔
467	return self.ufunc_db[ufunc_key]	1✔
468
469
470	class DpexCallConv(MinimalCallConv):	1✔
471	"""Custom calling convention class used by numba-dpex.
472
473	numba_dpex's calling convention derives from
474	:class:`numba.core.callconv import MinimalCallConv`. The
475	:class:`DpexCallConv` overriddes :func:`call_function`.
476
477	"""
478
479	def call_function(self, builder, callee, resty, argtys, args, env=None):	1✔
480	"""Call the Numba-compiled callee."""
481	assert env is None	1✔
482	retty = callee.args[0].type.pointee	1✔
483	retvaltmp = cgutils.alloca_once(builder, retty)	1✔
484	# initialize return value
485	builder.store(cgutils.get_null_value(retty), retvaltmp)	1✔
486	arginfo = self.context.get_arg_packer(argtys)	1✔
487	args = arginfo.as_arguments(builder, args)	1✔
488	realargs = [retvaltmp] + list(args)	1✔
489	code = builder.call(callee, realargs)	1✔
490	status = self._get_return_status(builder, code)	1✔
491	retval = builder.load(retvaltmp)	1✔
492	out = self.context.get_returned_value(builder, resty, retval)	1✔
493	return status, out	1✔

IntelPython / numba-dpex / 5681978645

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous