5433582462

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

from numba.core import datamodel, types
from numba.core.datamodel.models import PrimitiveModel, StructModel
from numba.core.extending import register_model

from numba_dpex.core.exceptions import UnreachableError
from numba_dpex.utils import address_space

from ..types import Array, DpctlSyclQueue, DpnpNdArray, USMNdArray


class GenericPointerModel(PrimitiveModel):
    def __init__(self, dmm, fe_type):
        adrsp = (
            fe_type.addrspace
            if fe_type.addrspace is not None
            else address_space.GLOBAL
        )
        be_type = dmm.lookup(fe_type.dtype).get_data_type().as_pointer(adrsp)
        super(GenericPointerModel, self).__init__(dmm, fe_type, be_type)


class USMArrayModel(StructModel):
    """A data model to represent a Dpex's array types in LLVM IR.

    Dpex's ArrayModel is based on Numba's ArrayModel for NumPy arrays. The
    dpex model adds an extra address space attribute to all pointer members
    in the array.
    """

    def __init__(self, dmm, fe_type):
        ndim = fe_type.ndim
        members = [
            (
                "meminfo",
                types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace),
            ),
            (
                "parent",
                types.CPointer(types.pyobject, addrspace=fe_type.addrspace),
            ),
            ("nitems", types.intp),
            ("itemsize", types.intp),
            (
                "data",
                types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace),
            ),
            (
                "sycl_queue",
                types.CPointer(types.void, addrspace=fe_type.addrspace),
            ),
            ("shape", types.UniTuple(types.intp, ndim)),
            ("strides", types.UniTuple(types.intp, ndim)),
        ]
        super(USMArrayModel, self).__init__(dmm, fe_type, members)


class DpnpNdArrayModel(StructModel):
    """Data model for the DpnpNdArray type.

    The data model for DpnpNdArray is similar to numb's ArrayModel used for
    the numba.types.Array type, with the additional field ``sycl_queue`. The
    `sycl_queue` attribute stores the pointer to the C++ sycl::queue object
    that was used to allocate memory for numba-dpex's native representation
    for an Python object inferred as a DpnpNdArray.
    """

    def __init__(self, dmm, fe_type):
        ndim = fe_type.ndim
        members = [
            ("meminfo", types.MemInfoPointer(fe_type.dtype)),
            ("parent", types.pyobject),
            ("nitems", types.intp),
            ("itemsize", types.intp),
            ("data", types.CPointer(fe_type.dtype)),
            ("sycl_queue", types.voidptr),
            ("shape", types.UniTuple(types.intp, ndim)),
            ("strides", types.UniTuple(types.intp, ndim)),
        ]
        super(DpnpNdArrayModel, self).__init__(dmm, fe_type, members)

    @property
    def flattened_field_count(self):
        """Return the number of fields in an instance of a DpnpNdArrayModel."""
        flattened_member_count = 0
        members = self._members
        for member in members:
            if isinstance(member, types.UniTuple):
                flattened_member_count += member.count
            elif isinstance(
                member,
                (
                    types.scalars.Integer,
                    types.misc.PyObject,
                    types.misc.RawPointer,
                    types.misc.CPointer,
                    types.misc.MemInfoPointer,
                ),
            ):
                flattened_member_count += 1
            else:
                print(member, type(member))
                raise UnreachableError

        return flattened_member_count


class SyclQueueModel(StructModel):
    """Represents the native data model for a dpctl.SyclQueue PyObject.

    Numba-dpex uses a C struct as defined in
    numba_dpex/core/runtime._queuestruct.h to store the required attributes for
    a ``dpctl.SyclQueue`` Python object.

        - ``queue_ref``: An opaque C pointer to an actual SYCL queue C++ object.
        - ``parent``: A PyObject* that stores a reference back to the original
                      ``dpctl.SyclQueue`` PyObject if the native struct is
                      created by unboxing the PyObject.
    """

    def __init__(self, dmm, fe_type):
        members = [
            (
                "parent",
                types.CPointer(types.int8),
            ),
            (
                "queue_ref",
                types.CPointer(types.int8),
            ),
        ]
        super(SyclQueueModel, self).__init__(dmm, fe_type, members)


def _init_data_model_manager():
    dmm = datamodel.default_manager.copy()
    dmm.register(types.CPointer, GenericPointerModel)
    dmm.register(Array, USMArrayModel)
    return dmm


dpex_data_model_manager = _init_data_model_manager()

# XXX A kernel function has the spir_kernel ABI and requires pointers to have an
# address space attribute. For this reason, the UsmNdArray type uses dpex's
# ArrayModel where the pointers are address space casted to have a SYCL-specific
# address space value. The DpnpNdArray type can be used inside djit functions
# as host function calls arguments, such as dpnp library calls. The DpnpNdArray
# needs to use Numba's array model as its data model. Thus, from a Numba typing
# perspective dpnp.ndarrays cannot be directly passed to a kernel. To get
# around the limitation, the DpexKernelTypingContext does not resolve the type
# of dpnp.array args to a kernel as DpnpNdArray type objects, but uses the
# ``to_usm_ndarray`` utility function to convert them into a UsmNdArray type
# object.

# Register the USMNdArray type with the dpex ArrayModel
register_model(USMNdArray)(USMArrayModel)
dpex_data_model_manager.register(USMNdArray, USMArrayModel)

# Register the DpnpNdArray type with the Numba ArrayModel
register_model(DpnpNdArray)(DpnpNdArrayModel)
dpex_data_model_manager.register(DpnpNdArray, DpnpNdArrayModel)

# Register the DpctlSyclQueue type
register_model(DpctlSyclQueue)(SyclQueueModel)
dpex_data_model_manager.register(DpctlSyclQueue, SyclQueueModel)

1	# SPDX-FileCopyrightText: 2022 - 2023 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	from numba.core import datamodel, types	1✔
6	from numba.core.datamodel.models import PrimitiveModel, StructModel	1✔
7	from numba.core.extending import register_model	1✔
8
9	from numba_dpex.core.exceptions import UnreachableError	1✔
10	from numba_dpex.utils import address_space	1✔
11
12	from ..types import Array, DpctlSyclQueue, DpnpNdArray, USMNdArray	1✔
13
14
15	class GenericPointerModel(PrimitiveModel):	1✔
16	def __init__(self, dmm, fe_type):	1✔
17	adrsp = (	1✔
18	fe_type.addrspace
19	if fe_type.addrspace is not None
20	else address_space.GLOBAL
21	)
22	be_type = dmm.lookup(fe_type.dtype).get_data_type().as_pointer(adrsp)	1✔
23	super(GenericPointerModel, self).__init__(dmm, fe_type, be_type)	1✔
24
25
26	class USMArrayModel(StructModel):	1✔
27	"""A data model to represent a Dpex's array types in LLVM IR.
28
29	Dpex's ArrayModel is based on Numba's ArrayModel for NumPy arrays. The
30	dpex model adds an extra address space attribute to all pointer members
31	in the array.
32	"""
33
34	def __init__(self, dmm, fe_type):	1✔
35	ndim = fe_type.ndim	1✔
36	members = [	1✔
37	(
38	"meminfo",
39	types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace),
40	),
41	(
42	"parent",
43	types.CPointer(types.pyobject, addrspace=fe_type.addrspace),
44	),
45	("nitems", types.intp),
46	("itemsize", types.intp),
47	(
48	"data",
49	types.CPointer(fe_type.dtype, addrspace=fe_type.addrspace),
50	),
51	(
52	"sycl_queue",
53	types.CPointer(types.void, addrspace=fe_type.addrspace),
54	),
55	("shape", types.UniTuple(types.intp, ndim)),
56	("strides", types.UniTuple(types.intp, ndim)),
57	]
58	super(USMArrayModel, self).__init__(dmm, fe_type, members)	1✔
59
60
61	class DpnpNdArrayModel(StructModel):	1✔
62	"""Data model for the DpnpNdArray type.
63
64	The data model for DpnpNdArray is similar to numb's ArrayModel used for
65	the numba.types.Array type, with the additional field ``sycl_queue`. The
66	`sycl_queue` attribute stores the pointer to the C++ sycl::queue object
67	that was used to allocate memory for numba-dpex's native representation
68	for an Python object inferred as a DpnpNdArray.
69	"""
70
71	def __init__(self, dmm, fe_type):	1✔
72	ndim = fe_type.ndim	1✔
73	members = [	1✔
74	("meminfo", types.MemInfoPointer(fe_type.dtype)),
75	("parent", types.pyobject),
76	("nitems", types.intp),
77	("itemsize", types.intp),
78	("data", types.CPointer(fe_type.dtype)),
79	("sycl_queue", types.voidptr),
80	("shape", types.UniTuple(types.intp, ndim)),
81	("strides", types.UniTuple(types.intp, ndim)),
82	]
83	super(DpnpNdArrayModel, self).__init__(dmm, fe_type, members)	1✔
84
85	@property	1✔
86	def flattened_field_count(self):	1✔
87	"""Return the number of fields in an instance of a DpnpNdArrayModel."""
88	flattened_member_count = 0	1✔
89	members = self._members	1✔
90	for member in members:	1✔
91	if isinstance(member, types.UniTuple):	1✔
92	flattened_member_count += member.count	1✔
93	elif isinstance(	1!
94	member,
95	(
96	types.scalars.Integer,
97	types.misc.PyObject,
98	types.misc.RawPointer,
99	types.misc.CPointer,
100	types.misc.MemInfoPointer,
101	),
102	):
103	flattened_member_count += 1	1✔
104	else:
NEW 105	print(member, type(member))	×
NEW 106	raise UnreachableError	×
107
108	return flattened_member_count	1✔
109
110
111	class SyclQueueModel(StructModel):	1✔
112	"""Represents the native data model for a dpctl.SyclQueue PyObject.
113
114	Numba-dpex uses a C struct as defined in
115	numba_dpex/core/runtime._queuestruct.h to store the required attributes for
116	a ``dpctl.SyclQueue`` Python object.
117
118	- ``queue_ref``: An opaque C pointer to an actual SYCL queue C++ object.
119	- ``parent``: A PyObject* that stores a reference back to the original
120	``dpctl.SyclQueue`` PyObject if the native struct is
121	created by unboxing the PyObject.
122	"""
123
124	def __init__(self, dmm, fe_type):	1✔
125	members = [	1✔
126	(
127	"parent",
128	types.CPointer(types.int8),
129	),
130	(
131	"queue_ref",
132	types.CPointer(types.int8),
133	),
134	]
135	super(SyclQueueModel, self).__init__(dmm, fe_type, members)	1✔
136
137
138	def _init_data_model_manager():	1✔
139	dmm = datamodel.default_manager.copy()	1✔
140	dmm.register(types.CPointer, GenericPointerModel)	1✔
141	dmm.register(Array, USMArrayModel)	1✔
142	return dmm	1✔
143
144
145	dpex_data_model_manager = _init_data_model_manager()	1✔
146
147	# XXX A kernel function has the spir_kernel ABI and requires pointers to have an
148	# address space attribute. For this reason, the UsmNdArray type uses dpex's
149	# ArrayModel where the pointers are address space casted to have a SYCL-specific
150	# address space value. The DpnpNdArray type can be used inside djit functions
151	# as host function calls arguments, such as dpnp library calls. The DpnpNdArray
152	# needs to use Numba's array model as its data model. Thus, from a Numba typing
153	# perspective dpnp.ndarrays cannot be directly passed to a kernel. To get
154	# around the limitation, the DpexKernelTypingContext does not resolve the type
155	# of dpnp.array args to a kernel as DpnpNdArray type objects, but uses the
156	# ``to_usm_ndarray`` utility function to convert them into a UsmNdArray type
157	# object.
158
159	# Register the USMNdArray type with the dpex ArrayModel
160	register_model(USMNdArray)(USMArrayModel)	1✔
161	dpex_data_model_manager.register(USMNdArray, USMArrayModel)	1✔
162
163	# Register the DpnpNdArray type with the Numba ArrayModel
164	register_model(DpnpNdArray)(DpnpNdArrayModel)	1✔
165	dpex_data_model_manager.register(DpnpNdArray, DpnpNdArrayModel)	1✔
166
167	# Register the DpctlSyclQueue type
168	register_model(DpctlSyclQueue)(SyclQueueModel)	1✔
169	dpex_data_model_manager.register(DpctlSyclQueue, SyclQueueModel)	1✔

IntelPython / numba-dpex / 5433582462

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