6334344536

Committed 28 Sep 2023 03:58AM UTC coverage: 74.89% (-2.1%) from 76.941%

Build # 6334344536

Build Type

Pull #1148

github

Committed by

web-flow

Commit Message

Merge cbffa1dcb into b85d37926

Pull Request Pull Request #1148: Adds Range and NdRange as supported types in numba_dpex.dpjit.

Run Details

1079 of 1653 branches covered (0.0%)

Branch coverage included in aggregate %.

55 of 243 new or added lines in 5 files covered. (22.63%)

25 existing lines in 3 files now uncovered.

4534 of 5842 relevant lines covered (77.61%)

0.78 hits per line

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

36.52

/numba_dpex/core/kernel_interface/indexers.py

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

from collections.abc import Iterable

from llvmlite import ir as llvmir
from numba.core import cgutils, errors, types
from numba.core.datamodel import default_manager
from numba.extending import intrinsic, overload


class Range(tuple):
    """A data structure to encapsulate a single kernel launch parameter.

    The range is an abstraction that describes the number of elements
    in each dimension of buffers and index spaces. It can contain
    1, 2, or 3 numbers, depending on the dimensionality of the
    object it describes.

    This is just a wrapper class on top of a 3-tuple. The kernel launch
    parameter is consisted of three int's. This class basically mimics
    the behavior of `sycl::range`.
    """

    def __new__(cls, dim0, dim1=None, dim2=None):
        """Constructs a 1, 2, or 3 dimensional range.

        Args:
            dim0 (int): The range of the first dimension.
            dim1 (int, optional): The range of second dimension.
                                    Defaults to None.
            dim2 (int, optional): The range of the third dimension.
                                    Defaults to None.

        Raises:
            TypeError: If dim0 is not an int.
            TypeError: If dim1 is not an int.
            TypeError: If dim2 is not an int.
        """
        if not isinstance(dim0, int):
            raise TypeError("dim0 of a Range must be an int.")
        _values = [dim0]
        if dim1:
            if not isinstance(dim1, int):
                raise TypeError("dim1 of a Range must be an int.")
            _values.append(dim1)
            if dim2:
                if not isinstance(dim2, int):
                    raise TypeError("dim2 of a Range must be an int.")
                _values.append(dim2)
        return super(Range, cls).__new__(cls, tuple(_values))

    def get(self, index):
        """Returns the range of a single dimension.

        Args:
            index (int): The index of the dimension, i.e. [0,2]

        Returns:
            int: The range of the dimension indexed by `index`.
        """
        return self[index]

    def size(self):
        """Returns the size of a range.

        Returns the size of a range by multiplying
        the range of the individual dimensions.

        Returns:
            int: The size of a range.
        """
        n = len(self)
        if n > 2:
            return self[0] * self[1] * self[2]
        elif n > 1:
            return self[0] * self[1]
        else:
            return self[0]

    @property
    def ndim(self) -> int:
        """Returns the rank of a Range object.

        Returns:
            int: Number of dimensions in the Range object
        """
        return len(self)

    @property
    def dim0(self) -> int:
        """Return the extent of the first dimension for the Range object.

        Returns:
            int: Extent of first dimension for the Range object
        """
        return self[0]

    @property
    def dim1(self) -> int:
        """Return the extent of the second dimension for the Range object.

        Returns:
            int: Extent of second dimension for the Range object or -1 for 1D
            Range
        """
        try:
            return self[1]
        except:
            return -1

    @property
    def dim2(self) -> int:
        """Return the extent of the second dimension for the Range object.

        Returns:
            int: Extent of second dimension for the Range object or -1 for 1D or
            2D Range
        """
        try:
            return self[2]
        except:
            return -1


class NdRange:
    """A class to encapsulate all kernel launch parameters.

    The NdRange defines the index space for a work group as well as
    the global index space. It is passed to parallel_for to execute
    a kernel on a set of work items.

    This class basically contains two Range object, one for the global_range
    and the other for the local_range. The global_range parameter contains
    the global index space and the local_range parameter contains the index
    space of a work group. This class mimics the behavior of `sycl::nd_range`
    class.
    """

    def __init__(self, global_size, local_size):
        """Constructor for NdRange class.

        Args:
            global_size (Range or tuple of int's): The values for
                the global_range.
            local_size (Range or tuple of int's, optional): The values for
                the local_range. Defaults to None.
        """
        if isinstance(global_size, Range):
            self._global_range = global_size
        elif isinstance(global_size, Iterable):
            self._global_range = Range(*global_size)
        else:
            raise TypeError(
                "Unknown argument type for NdRange global_size, "
                + "must be of either type Range or Iterable of int's."
            )

        if isinstance(local_size, Range):
            self._local_range = local_size
        elif isinstance(local_size, Iterable):
            self._local_range = Range(*local_size)
        else:
            raise TypeError(
                "Unknown argument type for NdRange local_size, "
                + "must be of either type Range or Iterable of int's."
            )

    @property
    def global_range(self):
        """Accessor for global_range.

        Returns:
            Range: The `global_range` `Range` object.
        """
        return self._global_range

    @property
    def local_range(self):
        """Accessor for local_range.

        Returns:
            Range: The `local_range` `Range` object.
        """
        return self._local_range

    def get_global_range(self):
        """Returns a Range defining the index space.

        Returns:
            Range: A `Range` object defining the index space.
        """
        return self._global_range

    def get_local_range(self):
        """Returns a Range defining the index space of a work group.

        Returns:
            Range: A `Range` object to specify index space of a work group.
        """
        return self._local_range

    def __str__(self):
        """str() function for NdRange class.

        Returns:
            str: str representation for NdRange class.
        """
        return (
            "(" + str(self._global_range) + ", " + str(self._local_range) + ")"
        )

    def __repr__(self):
        """repr() function for NdRange class.

        Returns:
            str: str representation for NdRange class.
        """
        return self.__str__()


@intrinsic
def _intrin_range_alloc(typingctx, ty_dim0, ty_dim1, ty_dim2, ty_range):
    ty_retty = ty_range.instance_type
    sig = ty_retty(
        ty_dim0,
        ty_dim1,
        ty_dim2,
        ty_range,
    )

    def codegen(context, builder, sig, args):
        typ = sig.return_type
        dim0, dim1, dim2, _ = args
        range_struct = cgutils.create_struct_proxy(typ)(context, builder)
        range_struct.dim0 = dim0

        if not isinstance(sig.args[1], types.NoneType):
            range_struct.dim1 = dim1
        else:
            range_struct.dim1 = llvmir.Constant(llvmir.types.IntType(64), -1)

        if not isinstance(sig.args[2], types.NoneType):
            range_struct.dim2 = dim2
        else:
            range_struct.dim2 = llvmir.Constant(llvmir.types.IntType(64), -1)

        range_struct.ndim = llvmir.Constant(llvmir.types.IntType(64), typ.ndim)

        return range_struct._getvalue()

    return sig, codegen


@intrinsic
def _intrin_ndrange_alloc(
    typingctx, ty_global_range, ty_local_range, ty_ndrange
):
    ty_retty = ty_ndrange.instance_type
    sig = ty_retty(
        ty_global_range,
        ty_local_range,
        ty_ndrange,
    )

    def codegen(context, builder, sig, args):
        typ = sig.return_type

        global_range, local_range, _ = args
        ndrange_struct = cgutils.create_struct_proxy(typ)(context, builder)
        ndrange_struct.ndim = llvmir.Constant(
            llvmir.types.IntType(64), typ.ndim
        )
        range_datamodel = default_manager.lookup(sig.args[0])
        ndrange_struct.gdim0 = builder.extract_value(
            global_range,
            range_datamodel.get_field_position("dim0"),
        )
        ndrange_struct.gdim1 = builder.extract_value(
            global_range,
            range_datamodel.get_field_position("dim1"),
        )
        ndrange_struct.gdim2 = builder.extract_value(
            global_range,
            range_datamodel.get_field_position("dim2"),
        )
        ndrange_struct.ldim0 = builder.extract_value(
            local_range,
            range_datamodel.get_field_position("dim0"),
        )
        ndrange_struct.ldim1 = builder.extract_value(
            local_range,
            range_datamodel.get_field_position("dim1"),
        )
        ndrange_struct.ldim2 = builder.extract_value(
            local_range,
            range_datamodel.get_field_position("dim2"),
        )

        return ndrange_struct._getvalue()

    return sig, codegen


@overload(Range)
def _ol_range_init(dim0, dim1=None, dim2=None):
    """Numba overload of the Range constructor to make it usable inside an
    njit and dpjit decorated function.

    """
    from numba_dpex.core.types import RangeType

    ndims = 1
    ty_optional_dims = (dim1, dim2)

    # A Range should at least have the 0th dimension populated
    if not isinstance(dim0, types.Integer):
        raise errors.TypingError(
            "Expected a Range's dimension should to be an Integer value, but "
            "encountered " + dim0.name
        )

    for ty_dim in ty_optional_dims:
        if isinstance(ty_dim, types.Integer):
            ndims += 1
        elif ty_dim is not None:
            raise errors.TypingError(
                "Expected a Range's dimension to be an Integer value, "
                f"but {type(ty_dim)} was provided."
            )

    ret_ty = RangeType(ndims)

    def impl(dim0, dim1=None, dim2=None):
        return _intrin_range_alloc(dim0, dim1, dim2, ret_ty)

    return impl


@overload(NdRange)
def _ol_ndrange_init(global_range, local_range):
    """Numba overload of the NdRange constructor to make it usable inside an
    njit and dpjit decorated function.

    """
    from numba_dpex.core.exceptions import UnmatchedNumberOfRangeDimsError
    from numba_dpex.core.types import NdRangeType, RangeType

    if not isinstance(global_range, RangeType):
        raise errors.TypingError(
            "Only global range values specified as a Range are "
            "supported inside dpjit"
        )

    if not isinstance(local_range, RangeType):
        raise errors.TypingError(
            "Only local range values specified as a Range are "
            "supported inside dpjit"
        )

    if not global_range.ndim == local_range.ndim:
        raise UnmatchedNumberOfRangeDimsError(
            kernel_name="",
            global_ndims=global_range.ndim,
            local_ndims=local_range.ndim,
        )

    ret_ty = NdRangeType(global_range.ndim)

    def impl(global_range, local_range):
        return _intrin_ndrange_alloc(global_range, local_range, ret_ty)

    return impl

1	# SPDX-FileCopyrightText: 2023 Intel Corporation
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	from collections.abc import Iterable	1✔
6
7	from llvmlite import ir as llvmir	1✔
8	from numba.core import cgutils, errors, types	1✔
9	from numba.core.datamodel import default_manager	1✔
10	from numba.extending import intrinsic, overload	1✔
11
12
13	class Range(tuple):	1✔
14	"""A data structure to encapsulate a single kernel launch parameter.
15
16	The range is an abstraction that describes the number of elements
17	in each dimension of buffers and index spaces. It can contain
18	1, 2, or 3 numbers, depending on the dimensionality of the
19	object it describes.
20
21	This is just a wrapper class on top of a 3-tuple. The kernel launch
22	parameter is consisted of three int's. This class basically mimics
23	the behavior of `sycl::range`.
24	"""
25
26	def __new__(cls, dim0, dim1=None, dim2=None):	1✔
27	"""Constructs a 1, 2, or 3 dimensional range.
28
29	Args:
30	dim0 (int): The range of the first dimension.
31	dim1 (int, optional): The range of second dimension.
32	Defaults to None.
33	dim2 (int, optional): The range of the third dimension.
34	Defaults to None.
35
36	Raises:
37	TypeError: If dim0 is not an int.
38	TypeError: If dim1 is not an int.
39	TypeError: If dim2 is not an int.
40	"""
41	if not isinstance(dim0, int):	1✔
42	raise TypeError("dim0 of a Range must be an int.")	1✔
43	_values = [dim0]	1✔
44	if dim1:	1✔
45	if not isinstance(dim1, int):	1!
UNCOV 46	raise TypeError("dim1 of a Range must be an int.")	×
47	_values.append(dim1)	1✔
48	if dim2:	1!
UNCOV 49	if not isinstance(dim2, int):	×
UNCOV 50	raise TypeError("dim2 of a Range must be an int.")	×
51	_values.append(dim2)	×
52	return super(Range, cls).__new__(cls, tuple(_values))	1✔
53
54	def get(self, index):	1✔
55	"""Returns the range of a single dimension.
56
57	Args:
58	index (int): The index of the dimension, i.e. [0,2]
59
60	Returns:
61	int: The range of the dimension indexed by `index`.
62	"""
UNCOV 63	return self[index]	×
64
65	def size(self):	1✔
66	"""Returns the size of a range.
67
68	Returns the size of a range by multiplying
69	the range of the individual dimensions.
70
71	Returns:
72	int: The size of a range.
73	"""
UNCOV 74	n = len(self)	×
UNCOV 75	if n > 2:	×
76	return self[0] * self[1] * self[2]	×
77	elif n > 1:	×
78	return self[0] * self[1]	×
79	else:
80	return self[0]	×
81
82	@property	1✔
83	def ndim(self) -> int:	1✔
84	"""Returns the rank of a Range object.
85
86	Returns:
87	int: Number of dimensions in the Range object
88	"""
NEW 89	return len(self)	×
90
91	@property	1✔
92	def dim0(self) -> int:	1✔
93	"""Return the extent of the first dimension for the Range object.
94
95	Returns:
96	int: Extent of first dimension for the Range object
97	"""
NEW 98	return self[0]	×
99
100	@property	1✔
101	def dim1(self) -> int:	1✔
102	"""Return the extent of the second dimension for the Range object.
103
104	Returns:
105	int: Extent of second dimension for the Range object or -1 for 1D
106	Range
107	"""
NEW 108	try:	×
NEW 109	return self[1]	×
NEW 110	except:	×
NEW 111	return -1	×
112
113	@property	1✔
114	def dim2(self) -> int:	1✔
115	"""Return the extent of the second dimension for the Range object.
116
117	Returns:
118	int: Extent of second dimension for the Range object or -1 for 1D or
119	2D Range
120	"""
NEW 121	try:	×
NEW 122	return self[2]	×
NEW 123	except:	×
NEW 124	return -1	×
125
126
127	class NdRange:	1✔
128	"""A class to encapsulate all kernel launch parameters.
129
130	The NdRange defines the index space for a work group as well as
131	the global index space. It is passed to parallel_for to execute
132	a kernel on a set of work items.
133
134	This class basically contains two Range object, one for the global_range
135	and the other for the local_range. The global_range parameter contains
136	the global index space and the local_range parameter contains the index
137	space of a work group. This class mimics the behavior of `sycl::nd_range`
138	class.
139	"""
140
141	def __init__(self, global_size, local_size):	1✔
142	"""Constructor for NdRange class.
143
144	Args:
145	global_size (Range or tuple of int's): The values for
146	the global_range.
147	local_size (Range or tuple of int's, optional): The values for
148	the local_range. Defaults to None.
149	"""
150	if isinstance(global_size, Range):	1✔
151	self._global_range = global_size	1✔
152	elif isinstance(global_size, Iterable):	1!
153	self._global_range = Range(*global_size)	1✔
154	else:
UNCOV 155	raise TypeError(	×
156	"Unknown argument type for NdRange global_size, "
157	+ "must be of either type Range or Iterable of int's."
158	)
159
160	if isinstance(local_size, Range):	1✔
161	self._local_range = local_size	1✔
162	elif isinstance(local_size, Iterable):	1!
163	self._local_range = Range(*local_size)	1✔
164	else:
UNCOV 165	raise TypeError(	×
166	"Unknown argument type for NdRange local_size, "
167	+ "must be of either type Range or Iterable of int's."
168	)
169
170	@property	1✔
171	def global_range(self):	1✔
172	"""Accessor for global_range.
173
174	Returns:
175	Range: The `global_range` `Range` object.
176	"""
177	return self._global_range	1✔
178
179	@property	1✔
180	def local_range(self):	1✔
181	"""Accessor for local_range.
182
183	Returns:
184	Range: The `local_range` `Range` object.
185	"""
186	return self._local_range	1✔
187
188	def get_global_range(self):	1✔
189	"""Returns a Range defining the index space.
190
191	Returns:
192	Range: A `Range` object defining the index space.
193	"""
UNCOV 194	return self._global_range	×
195
196	def get_local_range(self):	1✔
197	"""Returns a Range defining the index space of a work group.
198
199	Returns:
200	Range: A `Range` object to specify index space of a work group.
201	"""
UNCOV 202	return self._local_range	×
203
204	def __str__(self):	1✔
205	"""str() function for NdRange class.
206
207	Returns:
208	str: str representation for NdRange class.
209	"""
UNCOV 210	return (	×
211	"(" + str(self._global_range) + ", " + str(self._local_range) + ")"
212	)
213
214	def __repr__(self):	1✔
215	"""repr() function for NdRange class.
216
217	Returns:
218	str: str representation for NdRange class.
219	"""
UNCOV 220	return self.__str__()	×
221
222
223	@intrinsic	1✔
224	def _intrin_range_alloc(typingctx, ty_dim0, ty_dim1, ty_dim2, ty_range):	1✔
NEW 225	ty_retty = ty_range.instance_type	×
NEW 226	sig = ty_retty(	×
227	ty_dim0,
228	ty_dim1,
229	ty_dim2,
230	ty_range,
231	)
232
NEW 233	def codegen(context, builder, sig, args):	×
NEW 234	typ = sig.return_type	×
NEW 235	dim0, dim1, dim2, _ = args	×
NEW 236	range_struct = cgutils.create_struct_proxy(typ)(context, builder)	×
NEW 237	range_struct.dim0 = dim0	×
238
NEW 239	if not isinstance(sig.args[1], types.NoneType):	×
NEW 240	range_struct.dim1 = dim1	×
241	else:
NEW 242	range_struct.dim1 = llvmir.Constant(llvmir.types.IntType(64), -1)	×
243
NEW 244	if not isinstance(sig.args[2], types.NoneType):	×
NEW 245	range_struct.dim2 = dim2	×
246	else:
NEW 247	range_struct.dim2 = llvmir.Constant(llvmir.types.IntType(64), -1)	×
248
NEW 249	range_struct.ndim = llvmir.Constant(llvmir.types.IntType(64), typ.ndim)	×
250
NEW 251	return range_struct._getvalue()	×
252
NEW 253	return sig, codegen	×
254
255
256	@intrinsic	1✔
257	def _intrin_ndrange_alloc(	1✔
258	typingctx, ty_global_range, ty_local_range, ty_ndrange
259	):
NEW 260	ty_retty = ty_ndrange.instance_type	×
NEW 261	sig = ty_retty(	×
262	ty_global_range,
263	ty_local_range,
264	ty_ndrange,
265	)
266
NEW 267	def codegen(context, builder, sig, args):	×
NEW 268	typ = sig.return_type	×
269
NEW 270	global_range, local_range, _ = args	×
NEW 271	ndrange_struct = cgutils.create_struct_proxy(typ)(context, builder)	×
NEW 272	ndrange_struct.ndim = llvmir.Constant(	×
273	llvmir.types.IntType(64), typ.ndim
274	)
NEW 275	range_datamodel = default_manager.lookup(sig.args[0])	×
NEW 276	ndrange_struct.gdim0 = builder.extract_value(	×
277	global_range,
278	range_datamodel.get_field_position("dim0"),
279	)
NEW 280	ndrange_struct.gdim1 = builder.extract_value(	×
281	global_range,
282	range_datamodel.get_field_position("dim1"),
283	)
NEW 284	ndrange_struct.gdim2 = builder.extract_value(	×
285	global_range,
286	range_datamodel.get_field_position("dim2"),
287	)
NEW 288	ndrange_struct.ldim0 = builder.extract_value(	×
289	local_range,
290	range_datamodel.get_field_position("dim0"),
291	)
NEW 292	ndrange_struct.ldim1 = builder.extract_value(	×
293	local_range,
294	range_datamodel.get_field_position("dim1"),
295	)
NEW 296	ndrange_struct.ldim2 = builder.extract_value(	×
297	local_range,
298	range_datamodel.get_field_position("dim2"),
299	)
300
NEW 301	return ndrange_struct._getvalue()	×
302
NEW 303	return sig, codegen	×
304
305
306	@overload(Range)	1✔
307	def _ol_range_init(dim0, dim1=None, dim2=None):	1✔
308	"""Numba overload of the Range constructor to make it usable inside an
309	njit and dpjit decorated function.
310
311	"""
NEW 312	from numba_dpex.core.types import RangeType	×
313
NEW 314	ndims = 1	×
NEW 315	ty_optional_dims = (dim1, dim2)	×
316
317	# A Range should at least have the 0th dimension populated
NEW 318	if not isinstance(dim0, types.Integer):	×
NEW 319	raise errors.TypingError(	×
320	"Expected a Range's dimension should to be an Integer value, but "
321	"encountered " + dim0.name
322	)
323
NEW 324	for ty_dim in ty_optional_dims:	×
NEW 325	if isinstance(ty_dim, types.Integer):	×
NEW 326	ndims += 1	×
NEW 327	elif ty_dim is not None:	×
NEW 328	raise errors.TypingError(	×
329	"Expected a Range's dimension to be an Integer value, "
330	f"but {type(ty_dim)} was provided."
331	)
332
NEW 333	ret_ty = RangeType(ndims)	×
334
NEW 335	def impl(dim0, dim1=None, dim2=None):	×
NEW 336	return _intrin_range_alloc(dim0, dim1, dim2, ret_ty)	×
337
NEW 338	return impl	×
339
340
341	@overload(NdRange)	1✔
342	def _ol_ndrange_init(global_range, local_range):	1✔
343	"""Numba overload of the NdRange constructor to make it usable inside an
344	njit and dpjit decorated function.
345
346	"""
NEW 347	from numba_dpex.core.exceptions import UnmatchedNumberOfRangeDimsError	×
NEW 348	from numba_dpex.core.types import NdRangeType, RangeType	×
349
NEW 350	if not isinstance(global_range, RangeType):	×
NEW 351	raise errors.TypingError(	×
352	"Only global range values specified as a Range are "
353	"supported inside dpjit"
354	)
355
NEW 356	if not isinstance(local_range, RangeType):	×
NEW 357	raise errors.TypingError(	×
358	"Only local range values specified as a Range are "
359	"supported inside dpjit"
360	)
361
NEW 362	if not global_range.ndim == local_range.ndim:	×
NEW 363	raise UnmatchedNumberOfRangeDimsError(	×
364	kernel_name="",
365	global_ndims=global_range.ndim,
366	local_ndims=local_range.ndim,
367	)
368
NEW 369	ret_ty = NdRangeType(global_range.ndim)	×
370
NEW 371	def impl(global_range, local_range):	×
NEW 372	return _intrin_ndrange_alloc(global_range, local_range, ret_ty)	×
373
NEW 374	return impl	×

IntelPython / numba-dpex / 6334344536

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