6346937198

Committed 29 Sep 2023 03:46AM UTC coverage: 77.107% (+0.7%) from 76.426%

Build # 6346937198

Build Type

Pull #1148

github

Committed by

web-flow

Commit Message

Merge f3b7d2678 into 49109704d

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

Run Details

1093 of 1654 branches covered (0.0%)

Branch coverage included in aggregate %.

244 of 254 new or added lines in 5 files covered. (96.06%)

1 existing line in 1 file now uncovered.

4697 of 5855 relevant lines covered (80.22%)

0.8 hits per line

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

80.34

/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,
    )
    range_datamodel = default_manager.lookup(ty_global_range)

    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
        )
        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!
46	raise TypeError("dim1 of a Range must be an int.")	×
47	_values.append(dim1)	1✔
48	if dim2:	1✔
49	if not isinstance(dim2, int):	1!
50	raise TypeError("dim2 of a Range must be an int.")	×
51	_values.append(dim2)	1✔
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	"""
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	"""
74	n = len(self)	×
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	"""
89	return len(self)	1✔
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	"""
98	return self[0]	1✔
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	"""
108	try:	1✔
109	return self[1]	1✔
110	except:	1✔
111	return -1	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	"""
121	try:	1✔
122	return self[2]	1✔
123	except:	1✔
124	return -1	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:
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:
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	"""
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	"""
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	"""
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	"""
220	return self.__str__()	×
221
222
223	@intrinsic	1✔
224	def _intrin_range_alloc(typingctx, ty_dim0, ty_dim1, ty_dim2, ty_range):	1✔
225	ty_retty = ty_range.instance_type	1✔
226	sig = ty_retty(	1✔
227	ty_dim0,
228	ty_dim1,
229	ty_dim2,
230	ty_range,
231	)
232
233	def codegen(context, builder, sig, args):	1✔
234	typ = sig.return_type	1✔
235	dim0, dim1, dim2, _ = args	1✔
236	range_struct = cgutils.create_struct_proxy(typ)(context, builder)	1✔
237	range_struct.dim0 = dim0	1✔
238
239	if not isinstance(sig.args[1], types.NoneType):	1✔
240	range_struct.dim1 = dim1	1✔
241	else:
242	range_struct.dim1 = llvmir.Constant(llvmir.types.IntType(64), -1)	1✔
243
244	if not isinstance(sig.args[2], types.NoneType):	1✔
245	range_struct.dim2 = dim2	1✔
246	else:
247	range_struct.dim2 = llvmir.Constant(llvmir.types.IntType(64), -1)	1✔
248
249	range_struct.ndim = llvmir.Constant(llvmir.types.IntType(64), typ.ndim)	1✔
250
251	return range_struct._getvalue()	1✔
252
253	return sig, codegen	1✔
254
255
256	@intrinsic	1✔
257	def _intrin_ndrange_alloc(	1✔
258	typingctx, ty_global_range, ty_local_range, ty_ndrange
259	):
260	ty_retty = ty_ndrange.instance_type	1✔
261	sig = ty_retty(	1✔
262	ty_global_range,
263	ty_local_range,
264	ty_ndrange,
265	)
266	range_datamodel = default_manager.lookup(ty_global_range)	1✔
267
268	def codegen(context, builder, sig, args):	1✔
269	typ = sig.return_type	1✔
270
271	global_range, local_range, _ = args	1✔
272	ndrange_struct = cgutils.create_struct_proxy(typ)(context, builder)	1✔
273	ndrange_struct.ndim = llvmir.Constant(	1✔
274	llvmir.types.IntType(64), typ.ndim
275	)
276	ndrange_struct.gdim0 = builder.extract_value(	1✔
277	global_range,
278	range_datamodel.get_field_position("dim0"),
279	)
280	ndrange_struct.gdim1 = builder.extract_value(	1✔
281	global_range,
282	range_datamodel.get_field_position("dim1"),
283	)
284	ndrange_struct.gdim2 = builder.extract_value(	1✔
285	global_range,
286	range_datamodel.get_field_position("dim2"),
287	)
288	ndrange_struct.ldim0 = builder.extract_value(	1✔
289	local_range,
290	range_datamodel.get_field_position("dim0"),
291	)
292	ndrange_struct.ldim1 = builder.extract_value(	1✔
293	local_range,
294	range_datamodel.get_field_position("dim1"),
295	)
296	ndrange_struct.ldim2 = builder.extract_value(	1✔
297	local_range,
298	range_datamodel.get_field_position("dim2"),
299	)
300
301	return ndrange_struct._getvalue()	1✔
302
303	return sig, codegen	1✔
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	"""
312	from numba_dpex.core.types import RangeType	1✔
313
314	ndims = 1	1✔
315	ty_optional_dims = (dim1, dim2)	1✔
316
317	# A Range should at least have the 0th dimension populated
318	if not isinstance(dim0, types.Integer):	1!
NEW 319	raise errors.TypingError(	×
320	"Expected a Range's dimension should to be an Integer value, but "
321	"encountered " + dim0.name
322	)
323
324	for ty_dim in ty_optional_dims:	1✔
325	if isinstance(ty_dim, types.Integer):	1✔
326	ndims += 1	1✔
327	elif ty_dim is not None:	1!
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
333	ret_ty = RangeType(ndims)	1✔
334
335	def impl(dim0, dim1=None, dim2=None):	1✔
NEW 336	return _intrin_range_alloc(dim0, dim1, dim2, ret_ty)	×
337
338	return impl	1✔
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	"""
347	from numba_dpex.core.exceptions import UnmatchedNumberOfRangeDimsError	1✔
348	from numba_dpex.core.types import NdRangeType, RangeType	1✔
349
350	if not isinstance(global_range, RangeType):	1!
NEW 351	raise errors.TypingError(	×
352	"Only global range values specified as a Range are "
353	"supported inside dpjit"
354	)
355
356	if not isinstance(local_range, RangeType):	1!
NEW 357	raise errors.TypingError(	×
358	"Only local range values specified as a Range are "
359	"supported inside dpjit"
360	)
361
362	if not global_range.ndim == local_range.ndim:	1!
NEW 363	raise UnmatchedNumberOfRangeDimsError(	×
364	kernel_name="",
365	global_ndims=global_range.ndim,
366	local_ndims=local_range.ndim,
367	)
368
369	ret_ty = NdRangeType(global_range.ndim)	1✔
370
371	def impl(global_range, local_range):	1✔
NEW 372	return _intrin_ndrange_alloc(global_range, local_range, ret_ty)	×
373
374	return impl	1✔

IntelPython / numba-dpex / 6346937198

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